from collections import namedtuple
from collections.abc import Sequence
from types import SimpleNamespace
from typing import Literal
import jax
import jax.numpy as jnp
import numpy as np
import tidy3d
from jax.typing import ArrayLike
from loguru import logger
from tidy3d.components.mode.solver import compute_modes as _compute_modes
from fdtdx.core.axis import get_transverse_axes
from fdtdx.core.jax.utils import is_jax_tracer
from fdtdx.core.misc import expand_to_3x3
from fdtdx.core.physics.metrics import normalize_by_poynting_flux
from fdtdx.core.physics.symmetry import (
mirror_edge_coordinates,
mirror_material_cross_section,
project_onto_parity,
restrict_to_kept_half,
)
ModeTupleType = namedtuple("ModeTupleType", ["neff", "Ex", "Ey", "Ez", "Hx", "Hy", "Hz"])
"""A named tuple containing the mode fields and effective index.
Attributes:
neff: Complex effective refractive index of the mode
Ex: x-component of the electric field
Ey: y-component of the electric field
Ez: z-component of the electric field
Hx: x-component of the magnetic field
Hy: y-component of the magnetic field
Hz: z-component of the magnetic field
"""
def compute_mode_polarization_fraction(
mode: ModeTupleType,
tangential_axes: tuple[int, int],
pol: Literal["te", "tm"],
) -> float:
"""Mode polarization fraction.
Args:
mode (ModeTupleType): a ModeTupleType instance
tangential_axes (tuple[int, int]): indices of transverse E-field component axes.
pol (Literal["te", "tm"]): "te" or "tm" determines which axis is 'E1'
Returns:
float: Polarization fraction between 0 and 1.
"""
E_fields = [mode.Ex, mode.Ey, mode.Ez]
E1 = E_fields[tangential_axes[0]]
E2 = E_fields[tangential_axes[1]]
if pol == "te":
numerator = np.sum(np.abs(E1) ** 2)
elif pol == "tm":
numerator = np.sum(np.abs(E2) ** 2)
else:
raise ValueError(f"pol must be 'te' or 'tm', but got {pol}")
denominator = np.sum(np.abs(E1) ** 2 + np.abs(E2) ** 2) + 1e-18
return numerator / denominator
def sort_modes(
modes: list[ModeTupleType],
filter_pol: Literal["te", "tm"] | None,
tangential_axes: tuple[int, int],
) -> list[ModeTupleType]:
"""
Sort modes by polarization.
Args:
modes (list[ModeTupleType]): list of modes.
filter_pol (Literal["te", "tm"] | None): If not none, sort by polarization specificaton.
tangential_axes (tuple[int, int]): indices of transverse E-field component axes.
Returns:
list[ModeTupleType]: sorted list of modes.
"""
if filter_pol is None:
return sorted(modes, key=lambda m: float(np.real(m.neff)), reverse=True)
def is_matching(mode):
frac = compute_mode_polarization_fraction(mode, tangential_axes, filter_pol)
return frac >= 0.5
matching = [m for m in modes if is_matching(m)]
non_matching = [m for m in modes if not is_matching(m)]
matching_sorted = sorted(matching, key=lambda m: float(np.real(m.neff)), reverse=True)
non_matching_sorted = sorted(non_matching, key=lambda m: float(np.real(m.neff)), reverse=True)
return matching_sorted + non_matching_sorted
[docs]
def compute_mode(
frequency: float,
inv_permittivities: jax.Array, # shape (nx, ny, nz)
inv_permeabilities: jax.Array | float,
resolution: float | None = None,
direction: Literal["+", "-"] = "+",
mode_index: int = 0,
filter_pol: Literal["te", "tm"] | None = None,
dtype: jnp.dtype = jnp.float32,
bend_radius: float | None = None,
bend_axis: int | None = None,
symmetry: tuple[int, int] = (0, 0),
transverse_coords: Sequence[jax.Array] | None = None,
fixed_propagation_axis: int | None = None,
) -> tuple[
jax.Array, # E
jax.Array, # H
jax.Array, # complex propagation constant
]:
"""Compute optical modes of a waveguide cross-section.
This function uses the Tidy3D mode solver to compute the optical modes of a given waveguide cross-section defined
by its permittivity distribution.
By default modes are sorted by their effective index. The mode_index argument indexes this sorted list of modes and
returns the desired mode. With filter_pol, it is also possible to only index a specific polarization.
Args:
frequency (float): Operating frequency in Hz
inv_permittivities (jax.Array): 3D array of inverse relative permittivity values
inv_permeabilities (jax.Array | float): 3D array of inverse relative permittivity values or single float for
uniform permeability distribution.
resolution (float | None): Uniform-grid spacing in metres. Required when ``transverse_coords`` is not
provided (uniform-grid path). Ignored when ``transverse_coords`` is given. Defaults to None.
direction (Literal["+", "-"]): Propagation direction, either "+" or "-".
mode_index (int, optional): Index of the mode to compute. Defaults to 0.
filter_pol (Literal["te", "tm"] | None, optional). If not None, modes are filtered by polarization.
dtype (jnp.dtype, optional): Float dtype of the simulation. Controls whether mode fields are returned
as complex64 (float32) or complex128 (float64). Defaults to jnp.float32.
bend_radius (float | None, optional): Bend radius of the waveguide in meters. Must be set together with
bend_axis. When set, the mode solver uses a conformal transformation to account for the bend. Defaults to
None (straight waveguide).
bend_axis (int | None, optional): Physical axis index (0/1/2) pointing from the waveguide toward the center
of curvature. Must differ from the propagation axis. Required when bend_radius is set. Defaults to None.
symmetry (tuple[int, int], optional): Symmetry-plane condition at the *min* edge of each transverse axis,
in the order of the two non-propagation physical axes (increasing index). ``0`` imposes a PEC mirror
(electric wall — the tidy3d default), ``1`` imposes a PMC mirror (magnetic wall). Use this when the
waveguide sits on a symmetry plane of a reduced (half/quarter) domain so the mode solver reproduces the
same boundary the FDTD uses there. For a +x-propagating TE mode on a y/z quarter domain with PEC at y=0
and PMC at the z Si-mid plane, pass ``(0, 1)``. Defaults to ``(0, 0)`` (PEC on both, i.e. no symmetry).
transverse_coords: Optional pair of physical edge-coordinate arrays, in metres, for the two axes transverse
to propagation. Each array must have one more entry than the corresponding transverse cell count.
When provided, the Tidy3D mode solver receives the non-uniform rectilinear grid directly.
JAX arrays are accepted; the numpy conversion happens inside the tidy3d callback so the function
remains compatible with ``jax.jit``.
fixed_propagation_axis: Optional physical axis normal to the mode plane.
This disambiguates an extruded 2-D plane whose normal axis and
invariant transverse axis are both one cell wide.
Returns:
Tuple[jax.Array, jax.Array, jax.Array]:
Tuple of E, H field and the effective index as complex-valued jax arrays.
"""
# Input validation
valid_permittivity = inv_permittivities.ndim == 4 and inv_permittivities.shape[0] in [1, 3, 9]
if fixed_propagation_axis is None:
valid_permittivity = valid_permittivity and sum(dim == 1 for dim in inv_permittivities.shape[1:]) == 1
else:
valid_permittivity = (
valid_permittivity
and fixed_propagation_axis in (0, 1, 2)
and inv_permittivities.shape[fixed_propagation_axis + 1] == 1
)
if not valid_permittivity:
raise ValueError(f"Invalid shape of inv_permittivities: {inv_permittivities.shape}")
if (
isinstance(inv_permeabilities, jax.Array)
and inv_permeabilities.ndim > 0
and (
not (inv_permeabilities.ndim == 4 and inv_permeabilities.shape[0] in [1, 3, 9])
or (
fixed_propagation_axis is None
and sum(dim == 1 for dim in inv_permeabilities.shape[1:]) != 1
)
or (
fixed_propagation_axis is not None
and inv_permeabilities.shape[fixed_propagation_axis + 1] != 1
)
)
):
raise ValueError(f"Invalid shape of inv_permeabilities: {inv_permeabilities.shape}")
if (bend_radius is None) != (bend_axis is None):
raise ValueError("bend_radius and bend_axis must both be set or both be None")
np_complex_dtype = np.complex128 if dtype == jnp.float64 else np.complex64
def mode_helper(permittivity, permeability, c0_um, c1_um):
coords = [np.asarray(c0_um), np.asarray(c1_um)]
# A genuinely extruded 2-D FDTD domain has one cell on one transverse
# axis. The eigensolver needs a non-degenerate 2-D cross-section, so
# solve the same uniform extrusion on four virtual cells spanning the
# identical physical extent, then average it back to one cell. Four is
# also the historical minimum used by FDTDX's compact 2-D examples,
# making source and overlap profiles invariant under this optimization.
singleton_axes = [axis for axis, size in enumerate(permittivity.shape[1:]) if size == 1]
virtual_extrusion_axis = singleton_axes[0] if singleton_axes else None
if virtual_extrusion_axis is not None:
repeat_axis = virtual_extrusion_axis + 1
permittivity = np.repeat(permittivity, 4, axis=repeat_axis)
if isinstance(permeability, np.ndarray) and permeability.ndim > 0:
permeability = np.repeat(permeability, 4, axis=repeat_axis)
coords[virtual_extrusion_axis] = np.linspace(
coords[virtual_extrusion_axis][0],
coords[virtual_extrusion_axis][-1],
5,
)
# Implicitly detect 2D mode if any transverse dimension is exactly 2
mode_2d = virtual_extrusion_axis is None and 2 in permittivity.shape[1:]
if mode_2d:
collapsed_axis = permittivity.shape[1:].index(2)
sl = (slice(None), slice(None), [0]) if collapsed_axis == 1 else (slice(None), [0], slice(None))
assert np.allclose(permittivity, permittivity[sl]), "Permittivity is not uniform across the collapsed axis!"
assert len(coords[collapsed_axis]) == 3, (
f"Assumption: Permittivity {permittivity.shape[1:]=}+1 matches ({coords[0].shape=}, {coords[1].shape})"
)
permittivity = permittivity[sl]
if isinstance(permeability, np.ndarray) and permeability.ndim > 0:
permeability = permeability[sl]
# Adjust coordinates for the collapsed dimension
coords[collapsed_axis] = coords[collapsed_axis][:2]
if bend_radius is not None:
assert bend_axis is not None
transverse_axes = get_transverse_axes(propagation_axis)
tidy3d_bend_axis = transverse_axes.index(bend_axis)
bend_radius_um = bend_radius / 1e-6
plane_center = (float(0.5 * (coords[0][0] + coords[0][-1])), float(0.5 * (coords[1][0] + coords[1][-1])))
else:
tidy3d_bend_axis = None
bend_radius_um = None
plane_center = None
modes = tidy3d_mode_computation_wrapper(
frequency=frequency,
permittivity_cross_section=permittivity,
permeability_cross_section=permeability,
coords=coords,
direction=direction,
num_modes=2 * (mode_index + 1) + 10,
bend_radius=bend_radius_um,
bend_axis=tidy3d_bend_axis,
plane_center=plane_center,
symmetry=symmetry,
)
# sort modes by polarization
# tidy3d assumes propagation in the z-direction. The tangential axes are therefore x and y.
modes = sort_modes(modes, filter_pol, (0, 1))
mode = modes[mode_index]
if propagation_axis == 0:
mode_E, mode_H = (
np.stack([mode.Ez, mode.Ex, mode.Ey], axis=0).astype(np_complex_dtype),
np.stack([mode.Hz, mode.Hx, mode.Hy], axis=0).astype(np_complex_dtype),
)
elif propagation_axis == 1:
mode_E, mode_H = (
np.stack([mode.Ex, mode.Ez, mode.Ey], axis=0).astype(np_complex_dtype),
-np.stack([mode.Hx, mode.Hz, mode.Hy], axis=0).astype(np_complex_dtype),
)
elif propagation_axis == 2:
mode_E, mode_H = (
np.stack([mode.Ex, mode.Ey, mode.Ez], axis=0).astype(np_complex_dtype),
np.stack([mode.Hx, mode.Hy, mode.Hz], axis=0).astype(np_complex_dtype),
)
else:
raise ValueError(f"Invalid propagation axis: {propagation_axis}")
if virtual_extrusion_axis is not None:
collapse_axis = virtual_extrusion_axis + 1
mode_E = np.mean(mode_E, axis=collapse_axis, keepdims=True)
mode_H = np.mean(mode_H, axis=collapse_axis, keepdims=True)
if mode_2d:
# Re-expand the collapsed dimension
mode_E = np.expand_dims(mode_E, axis=collapsed_axis + 1)
mode_E = np.repeat(mode_E, 2, axis=collapsed_axis + 1)
mode_H = np.expand_dims(mode_H, axis=collapsed_axis + 1)
mode_H = np.repeat(mode_H, 2, axis=collapsed_axis + 1)
neff = np.asarray(mode.neff).astype(np_complex_dtype)
return mode_E, mode_H, neff
# compute input to tidy3d Mode solver
if inv_permittivities.shape[0] == 9:
eps = expand_to_3x3(inv_permittivities)
# Invert the 3x3 matrix
perm = (2, 3, 4, 0, 1) # (3, 3, nx, ny, nz) -> (nx, ny, nz, 3, 3)
inv_perm = (3, 4, 0, 1, 2) # (nx, ny, nz, 3, 3) -> (3, 3, nx, ny, nz)
permittivities = (
jnp.linalg.inv(eps.transpose(perm)).transpose(inv_perm).reshape(9, *inv_permittivities.shape[1:])
)
else:
permittivities = 1 / inv_permittivities
propagation_axis = (
permittivities.shape[1:].index(1)
if fixed_propagation_axis is None
else fixed_propagation_axis
)
other_axes = [axis + 1 for axis in range(3) if axis != propagation_axis]
if transverse_coords is None:
if resolution is None:
raise ValueError("resolution is required when transverse_coords is not provided")
# Uniform grid: build concrete coordinate arrays in µm and pass as callback args.
c0_um = jnp.asarray(np.arange(permittivities.shape[other_axes[0]] + 1) * resolution / 1e-6)
c1_um = jnp.asarray(np.arange(permittivities.shape[other_axes[1]] + 1) * resolution / 1e-6)
normalization_area_weights = None
else:
if len(transverse_coords) != 2:
raise ValueError(
f"transverse_coords must contain exactly two coordinate arrays, got {len(transverse_coords)}"
)
# Shape validation uses .shape which is always concrete, even for JAX tracers.
expected_lengths = [permittivities.shape[dim] + 1 for dim in other_axes]
for axis_idx, (coord, expected_length) in enumerate(zip(transverse_coords, expected_lengths, strict=True)):
if coord.ndim != 1 or coord.shape[0] != expected_length:
raise ValueError(
f"transverse_coords[{axis_idx}] must be 1D with length {expected_length}, got {coord.shape}"
)
# Convert to µm for tidy3d; keep as JAX arrays so jax.jit can trace through.
c0_um = jnp.asarray(transverse_coords[0]) / 1e-6
c1_um = jnp.asarray(transverse_coords[1]) / 1e-6
# area_2d in m²: use jnp.diff so this works with traced JAX arrays.
area_2d = (
jnp.diff(jnp.asarray(transverse_coords[0]))[:, None] * jnp.diff(jnp.asarray(transverse_coords[1]))[None, :]
).astype(dtype)
weight_shape = [1, 1, 1]
weight_shape[other_axes[0] - 1] = area_2d.shape[0]
weight_shape[other_axes[1] - 1] = area_2d.shape[1]
normalization_area_weights = area_2d.reshape(weight_shape)
permittivity_squeezed = jnp.take(
permittivities,
indices=0,
axis=propagation_axis + 1,
)
# Rotate permittivity components to match tidy3d coordinate system
# tidy3d assumes propagation along z, so we need to map physical axes to tidy3d axes:
# - tidy3d x → first transverse axis
# - tidy3d y → second transverse axis
# - tidy3d z → propagation axis
if propagation_axis == 0:
# propagation along x: tidy3d (x,y,z) → physical (y,z,x)
perm_idx = [1, 2, 0]
perm_idx_full_anisotropy = [4, 5, 3, 7, 8, 6, 1, 2, 0]
elif propagation_axis == 1:
# propagation along y: tidy3d (x,y,z) → physical (x,z,y)
perm_idx = [0, 2, 1]
perm_idx_full_anisotropy = [0, 2, 1, 6, 8, 7, 3, 5, 4]
else: # propagation_axis == 2
# propagation along z: tidy3d (x,y,z) → physical (x,y,z)
perm_idx = [0, 1, 2]
perm_idx_full_anisotropy = [0, 1, 2, 3, 4, 5, 6, 7, 8]
# Only apply rotation if anisotropic (3 components)
if permittivity_squeezed.shape[0] == 3:
permittivity_squeezed = permittivity_squeezed[jnp.array(perm_idx), :, :]
if permittivity_squeezed.shape[0] == 9:
permittivity_squeezed = permittivity_squeezed[jnp.array(perm_idx_full_anisotropy), :, :]
jnp_complex_dtype = jnp.complex128 if dtype == jnp.float64 else jnp.complex64
result_shape_dtype = (
jnp.zeros((3, *permittivity_squeezed.shape[1:]), dtype=jnp_complex_dtype),
jnp.zeros((3, *permittivity_squeezed.shape[1:]), dtype=jnp_complex_dtype),
jnp.zeros(shape=(), dtype=jnp_complex_dtype),
)
if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0 and inv_permeabilities.shape[0] == 9:
mu = expand_to_3x3(inv_permeabilities)
# Invert the 3x3 matrix
perm = (2, 3, 4, 0, 1) # (3, 3, nx, ny, nz) -> (nx, ny, nz, 3, 3)
inv_perm = (3, 4, 0, 1, 2) # (nx, ny, nz, 3, 3) -> (3, 3, nx, ny, nz)
permeabilities = (
jnp.linalg.inv(mu.transpose(perm)).transpose(inv_perm).reshape(9, *inv_permeabilities.shape[1:])
)
else:
permeabilities = 1 / inv_permeabilities
if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0:
permeability_squeezed = jnp.take(
permeabilities,
indices=0,
axis=propagation_axis + 1,
)
# Apply same rotation to permeability if anisotropic
if permeability_squeezed.shape[0] == 3:
permeability_squeezed = permeability_squeezed[jnp.array(perm_idx), :, :]
if permeability_squeezed.shape[0] == 9:
permeability_squeezed = permeability_squeezed[jnp.array(perm_idx_full_anisotropy), :, :]
else: # float
permeability_squeezed = permeabilities
# pure callback to tidy3d is necessary to work in jitted environment.
# c0_um and c1_um are passed as explicit args so JAX materialises them to
# concrete numpy arrays before calling mode_helper, allowing np.asarray()
# inside the callback without raising TracerArrayConversionError.
mode_E_raw, mode_H_raw, eff_idx = jax.pure_callback(
mode_helper,
result_shape_dtype,
jax.lax.stop_gradient(permittivity_squeezed),
jax.lax.stop_gradient(permeability_squeezed),
jax.lax.stop_gradient(c0_um),
jax.lax.stop_gradient(c1_um),
)
mode_E = jnp.expand_dims(mode_E_raw, axis=propagation_axis + 1)
mode_H = jnp.expand_dims(mode_H_raw, axis=propagation_axis + 1)
# Tidy3D uses different scaling internally, so convert back
mode_H = mode_H * tidy3d.constants.ETA_0
mode_E_norm, mode_H_norm = normalize_by_poynting_flux(
mode_E,
mode_H,
axis=propagation_axis,
area_weights=normalization_area_weights,
)
return mode_E_norm, mode_H_norm, eff_idx
def _check_parity_residual(residual: jax.Array, walls: dict[int, int], object_name: str) -> None:
"""Diagnose a parity projection that removed too much of the mode.
Skipped when ``residual`` is a JAX tracer. A mode source or mode-overlap detector that overlaps a
:class:`~fdtdx.Device` solves its mode inside :func:`fdtdx.apply_params`, which callers routinely
trace (``jax.jit`` around an optimization step), and there the residual has no value yet:
concretizing it would raise, and comparing it would fire on a tracer. The check cannot be hoisted
out of the trace either — the residual depends on the device permittivity being traced over. It is
setup guidance, not something the solve depends on, so an eager ``apply_params`` surfaces it; and
``place_objects`` already applies (hence checks) every mode object that does not overlap a device,
which is the overwhelming majority.
Args:
residual (jax.Array): Fraction of the mode the projection removed.
walls (dict[int, int]): Mirror axis to wall type, for the message.
object_name (str): Name used in diagnostics.
Raises:
ValueError: If the projection removed (almost) the whole mode, i.e. the configured wall types
are incompatible with the selected mode.
"""
if is_jax_tracer(residual):
return
value = float(residual)
wall_description = ", ".join(f"{'xyz'[axis]}={'PMC' if wall == 1 else 'PEC'}" for axis, wall in walls.items())
if value > 0.9:
raise ValueError(
f"The mode selected for '{object_name}' has (almost) none of the symmetry imposed by the "
f"walls on {wall_description}: projecting it onto the admissible parity removes "
f"{value:.1%} of the mode. The wall types do not match this mode - flip the sign of "
f"config.symmetry on those axes, pick a different mode_index/filter_pol, or run without "
f"config.symmetry."
)
if value > 0.3:
logger.warning(
f"The mode of '{object_name}' is only approximately symmetric about the walls on "
f"{wall_description}: the parity projection removed {value:.1%} of it. Check that the "
f"structure is mirror-symmetric there and that the wall types match the mode. A coarsely "
f"resolved cross-section alone can account for a residual of a few tens of percent - the "
f"mode solver samples materials on its staggered grid, so its discrete mode is only "
f"mirror-symmetric to first order in the cell size."
)
def compute_mode_symmetry_reduced(
*,
mirrored_axes: tuple[int, ...],
walls: dict[int, int],
frequency: float,
inv_permittivities: jax.Array,
inv_permeabilities: jax.Array | float,
resolution: float | None = None,
direction: Literal["+", "-"] = "+",
mode_index: int = 0,
filter_pol: Literal["te", "tm"] | None = None,
dtype: jnp.dtype = jnp.float32,
bend_radius: float | None = None,
bend_axis: int | None = None,
transverse_coords: Sequence[jax.Array] | None = None,
object_name: str = "mode object",
fixed_propagation_axis: int | None = None,
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Solve a mode on a symmetry-reduced cross-section by way of the full cross-section.
A reduced simulation replaces the discarded half by the mirror image of the kept half, so the
mode it supports is the full-domain mode restricted to the kept half. Rather than asking the
mode solver for a *symmetric* solve on the reduced cross-section, this mirrors the reduced
material arrays back to the full cross-section, solves there, and restricts the result. The
detour matters: the mode solver interprets the permittivity arrays on its own staggered Yee
grid, while FDTDX rasterizes materials per cell and hands the same array to every component, so
a solve on the reduced cross-section is inconsistent with the full one at first order in the
cell size (an ``neff`` error of several percent for a high-index waveguide). Going through the
full cross-section reproduces exactly the mode the unreduced simulation would inject.
The solved mode is then projected onto the parity subspace the walls admit (see
:func:`~fdtdx.core.physics.symmetry.project_onto_parity`), which both removes the small
non-symmetric residue of the discrete mode and detects a wall type that does not match the mode
at all. Finally the fields are renormalized to unit Poynting flux through the *reduced* plane,
keeping the convention that a mode source launches unit power through the plane it occupies.
Both steps are limited by the same discretization the detour above avoids for ``neff``: because
the solver samples materials on its staggered grid, the discrete mode of a mirror-symmetric
cross-section is itself only symmetric to first order in the cell size, so its flux does not
split exactly evenly between the two halves. Measured for a 400x200 nm Si waveguide at 1.55 um:
0.446 / 0.554 along the solver's first transverse axis at 25 nm (0.473 / 0.527 at 12.5 nm), and
0.499 / 0.501 along its second. Renormalizing over the reduced plane therefore leaves the
returned profile up to ~6% (25 nm) resp. ~3% (12.5 nm) above ``sqrt(2**k)`` times the restriction
of the full-domain mode on such an axis, and the parity projection removes a few tenths of a
percent of its norm. Both vanish with refinement; the flux convention is exact by construction at
every resolution.
Args:
mirrored_axes (tuple[int, ...]): Physical axes clipped by a symmetry plane.
walls (dict[int, int]): Mirror axis to wall type (``-1`` PEC, ``+1`` PMC).
frequency (float): Operating frequency in Hz.
inv_permittivities (jax.Array): Reduced inverse permittivity on the mode plane.
inv_permeabilities (jax.Array | float): Reduced inverse permeability on the mode plane.
resolution (float | None): Uniform grid spacing, required without ``transverse_coords``.
direction (Literal["+", "-"]): Propagation direction.
mode_index (int): Index into the sorted mode list.
filter_pol (Literal["te", "tm"] | None): Optional polarization filter.
dtype (jnp.dtype): Float dtype of the simulation.
bend_radius (float | None): Waveguide bend radius, with ``bend_axis``.
bend_axis (int | None): Physical axis pointing toward the center of curvature.
transverse_coords (Sequence[jax.Array] | None): Reduced transverse edge coordinates, or
None on a uniform grid.
object_name (str): Name used in diagnostics.
fixed_propagation_axis (int | None): Explicit mode-plane normal for an
extruded domain with two singleton dimensions.
Returns:
tuple[jax.Array, jax.Array, jax.Array]: ``(E, H, effective_index)`` on the reduced
cross-section.
Raises:
ValueError: If a waveguide bend shares an axis with a symmetry plane, or if the parity
projection removes almost the entire mode, which means the configured wall types are
incompatible with the selected mode. The latter is only detectable where the residual is
concrete, i.e. not inside ``jax.jit`` (see :func:`_check_parity_residual`).
"""
if bend_radius is not None and bend_axis is not None and bend_axis in mirrored_axes:
raise ValueError(
f"'{object_name}' bends about the {'xyz'[bend_axis]}-axis and config.symmetry mirrors "
f"that same axis. The bend is modelled by a conformal transformation that scales the "
f"refractive index linearly across bend_axis, so the transformed cross-section is not "
f"mirror-symmetric about a plane normal to it: the mode has no definite parity there and "
f"the reduced simulation, which replaces the discarded half by the mirror of the kept "
f"half, cannot represent it. Drop config.symmetry on the {'xyz'[bend_axis]}-axis, or put "
f"the symmetry plane normal to the other transverse axis - a bend leaves that one "
f"mirror-symmetric."
)
propagation_axis = (
next(a for a in range(3) if inv_permittivities.shape[1:][a] == 1)
if fixed_propagation_axis is None
else fixed_propagation_axis
)
full_inv_permittivities = mirror_material_cross_section(inv_permittivities, mirrored_axes)
if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0:
full_inv_permeabilities: jax.Array | float = mirror_material_cross_section(inv_permeabilities, mirrored_axes)
else:
full_inv_permeabilities = inv_permeabilities
full_transverse_coords = transverse_coords
if transverse_coords is not None:
transverse_axes = get_transverse_axes(propagation_axis)
full_transverse_coords = [
mirror_edge_coordinates(coords) if axis in mirrored_axes else coords
for axis, coords in zip(transverse_axes, transverse_coords, strict=True)
]
mode_E, mode_H, eff_index = compute_mode(
frequency=frequency,
inv_permittivities=full_inv_permittivities,
inv_permeabilities=full_inv_permeabilities,
resolution=resolution,
direction=direction,
mode_index=mode_index,
filter_pol=filter_pol,
dtype=dtype,
bend_radius=bend_radius,
bend_axis=bend_axis,
symmetry=(0, 0),
transverse_coords=full_transverse_coords,
fixed_propagation_axis=propagation_axis,
)
mode_E, residual_E = project_onto_parity(mode_E, "E", walls)
mode_H, residual_H = project_onto_parity(mode_H, "H", walls)
_check_parity_residual(jnp.maximum(residual_E, residual_H), walls, object_name)
mode_E = restrict_to_kept_half(mode_E, mirrored_axes)
mode_H = restrict_to_kept_half(mode_H, mirrored_axes)
area_weights = None
if transverse_coords is not None:
widths = [jnp.diff(jnp.asarray(coords)) for coords in transverse_coords]
area_2d = widths[0][:, None] * widths[1][None, :]
weight_shape = [1, 1, 1]
for local_axis, axis in enumerate(get_transverse_axes(propagation_axis)):
weight_shape[axis] = area_2d.shape[local_axis]
area_weights = area_2d.reshape(weight_shape).astype(mode_E.real.dtype)
mode_E, mode_H = normalize_by_poynting_flux(
mode_E,
mode_H,
axis=propagation_axis,
area_weights=area_weights,
)
return mode_E, mode_H, eff_index
def _is_reciprocal_tensor(components: Sequence[ArrayLike], tol: float = 1e-6) -> bool:
"""Check whether a material tensor given as 9 row-major components is symmetric (reciprocal)."""
return bool(
np.max(np.abs(np.asarray(components[1]) - np.asarray(components[3]))) <= tol
and np.max(np.abs(np.asarray(components[2]) - np.asarray(components[6]))) <= tol
and np.max(np.abs(np.asarray(components[5]) - np.asarray(components[7]))) <= tol
)
def _canonicalize_mode_phase(fields: tuple[np.ndarray, ...], num_modes: int) -> tuple[np.ndarray, ...]:
"""Choose a deterministic global phase for each mode.
Eigenvectors have arbitrary global phase, and Tidy3D 2.9 can add a global
``±i`` to backward modes while normalizing their negative self-flux. Use
the largest transverse electric-field sample (``Ex`` or ``Ey`` in
Tidy3D's propagation-along-z frame) as a real, non-negative phase anchor.
This removes only global phase; intrinsic spatial phase in tensorial or
lossy modes is preserved.
"""
def canonicalize_one(mode_fields: tuple[np.ndarray, ...]) -> tuple[np.ndarray, ...]:
transverse_e = np.concatenate([mode_fields[0].ravel(), mode_fields[1].ravel()])
reference = transverse_e[np.argmax(np.abs(transverse_e))]
if abs(reference) == 0:
return mode_fields
phase = np.conj(reference) / abs(reference)
return tuple(field * phase for field in mode_fields)
arrays = tuple(np.asarray(field) for field in fields)
if num_modes == 1:
return canonicalize_one(arrays)
canonical_modes = [canonicalize_one(tuple(field[..., index] for field in arrays)) for index in range(num_modes)]
return tuple(np.stack([mode[field_index] for mode in canonical_modes], axis=-1) for field_index in range(6))
def tidy3d_mode_computation_wrapper(
frequency: float,
permittivity_cross_section: ArrayLike,
coords: list[np.ndarray],
direction: Literal["+", "-"],
permeability_cross_section: ArrayLike | float | None = None,
target_neff: float | None = None,
angle_theta: float = 0.0,
angle_phi: float = 0.0,
num_modes: int = 10,
precision: Literal["single", "double"] = "double",
bend_radius: float | None = None,
bend_axis: int | None = None,
plane_center: tuple[float, float] | None = None,
symmetry: tuple[int, int] = (0, 0),
) -> list[ModeTupleType]:
"""Compute optical modes of a waveguide cross-section.
This function uses the Tidy3D mode solver to compute the optical modes of a given
waveguide cross-section defined by its permittivity distribution.
Args:
frequency (float): Operating frequency in Hz
permittivity_cross_section (jax.Array): 2D array of relative permittivity values
coords (List[np.ndarray]): List of coordinate arrays [x, y] defining the grid
direction (Literal["+", "-"], optional): Propagation direction, either "+" or "-"
permeability_cross_section (jax.Array | float | None, optional): 2D array of relative permeability values.
Defauts to None.
target_neff (float | None, optional): Target effective index to search around. Defaults to None.
angle_theta (float, optional): Polar angle in radians. Defaults to 0.0.
angle_phi (float, optional): Azimuthal angle in radians. Defaults to 0.0.
num_modes (int, optional): Number of modes to compute. Defaults to 10.
precision (Literal["single", "double"], optional): Numerical precision. Defaults to "double".
bend_radius (float | None, optional): Bend radius in microns (tidy3d units). Defaults to None.
bend_axis (int | None, optional): Axis index (0 or 1) of the center of curvature in tidy3d's transverse
coordinate frame. Defaults to None.
plane_center (tuple[float, float] | None, optional): Center of the mode plane in the same units as coords.
Required by tidy3d when bend_radius is set. Defaults to None.
symmetry (tuple[int, int], optional): Per-transverse-axis symmetry condition at the min edge, forwarded to
the tidy3d mode solver. ``1`` imposes a PMC (magnetic) wall there; ``0`` (default) leaves the solver's
PEC (electric) wall. Order matches ``coords``. Defaults to ``(0, 0)``.
Notes:
tidy3d assumes propagation in z-direction. The output fields should be handled accordingly.
Returns:
List[ModeTupleType]: List of computed modes sorted by decreasing real part of
effective index. Each mode contains the field components and effective index.
"""
# see https://docs.flexcompute.com/projects/tidy3d/en/latest/_autosummary/tidy3d.ModeSpec.html#tidy3d.ModeSpec
mode_spec = SimpleNamespace(
# Note that the filter_pol argument is not used here since it does not work from tidy3d
num_modes=num_modes,
target_neff=target_neff,
num_pml=(0, 0),
angle_theta=angle_theta,
angle_phi=angle_phi,
bend_radius=bend_radius,
bend_axis=bend_axis,
precision=precision,
track_freq="central",
group_index_step=False,
)
permittivity_cross_section = jnp.asarray(permittivity_cross_section)
permittivity_cross_section = expand_to_3x3(permittivity_cross_section)
permittivity_cross_section = permittivity_cross_section.reshape(9, *permittivity_cross_section.shape[2:])
eps_cross = [
permittivity_cross_section[0],
permittivity_cross_section[1],
permittivity_cross_section[2],
permittivity_cross_section[3],
permittivity_cross_section[4],
permittivity_cross_section[5],
permittivity_cross_section[6],
permittivity_cross_section[7],
permittivity_cross_section[8],
]
mu_cross = None
if permeability_cross_section is not None:
permeability_cross_section = jnp.asarray(permeability_cross_section)
permeability_cross_section = expand_to_3x3(permeability_cross_section)
permeability_cross_section = permeability_cross_section.reshape(9, *permeability_cross_section.shape[2:])
mu_cross = [
permeability_cross_section[0],
permeability_cross_section[1],
permeability_cross_section[2],
permeability_cross_section[3],
permeability_cross_section[4],
permeability_cross_section[5],
permeability_cross_section[6],
permeability_cross_section[7],
permeability_cross_section[8],
]
if direction == "-" and (
angle_theta != 0.0
or angle_phi != 0.0
or not _is_reciprocal_tensor(eps_cross)
or (mu_cross is not None and not _is_reciprocal_tensor(mu_cross))
):
raise NotImplementedError(
"Backward ('-') modes are derived from the forward solve via the reciprocity transformation, "
"which requires symmetric material tensors and normal incidence."
)
# Solve the requested direction. A backward mode is not generally a simple
# component-sign transform of the forward mode when the material tensor
# couples longitudinal and transverse fields. Canonicalizing below removes
# Tidy3D 2.9's arbitrary backward global phase without changing the field.
EH, neffs, _ = _compute_modes(
eps_cross=eps_cross,
coords=coords,
freq=frequency,
precision=precision,
mode_spec=mode_spec,
direction=direction,
mu_cross=mu_cross,
plane_center=plane_center,
symmetry=symmetry,
)
((Ex, Ey, Ez), (Hx, Hy, Hz)) = EH.squeeze()
Ex, Ey, Ez, Hx, Hy, Hz = _canonicalize_mode_phase((Ex, Ey, Ez, Hx, Hy, Hz), num_modes)
if num_modes == 1:
modes = [
ModeTupleType(
Ex=Ex,
Ey=Ey,
Ez=Ez,
Hx=Hx,
Hy=Hy,
Hz=Hz,
neff=float(neffs.real) + 1j * float(neffs.imag),
)
for _ in range(num_modes)
]
else:
modes = [
ModeTupleType(
Ex=Ex[..., i],
Ey=Ey[..., i],
Ez=Ez[..., i],
Hx=Hx[..., i],
Hy=Hy[..., i],
Hz=Hz[..., i],
neff=neffs[i],
)
for i in range(num_modes)
]
return modes
