import warnings
from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Literal, Self, Sequence
import jax
import jax.numpy as jnp
import numpy as np
from fdtdx import constants
from fdtdx.config import SimulationConfig
from fdtdx.constants import c
from fdtdx.core.axis import get_transverse_axes
from fdtdx.core.jax.pytrees import autoinit, frozen_field, private_field
from fdtdx.core.misc import tilted_polarization_vectors
from fdtdx.core.null import Null
from fdtdx.core.physics.metrics import normalize_by_poynting_flux
from fdtdx.core.physics.modes import compute_mode, compute_mode_symmetry_reduced
from fdtdx.core.wavelength import WaveCharacter
from fdtdx.dispersion import effective_complex_inv_permittivity
from fdtdx.objects.detectors.detector import Detector, DetectorState
from fdtdx.objects.detectors.phasor import PhasorDetector
from fdtdx.typing import SliceTuple3D
def gaussian_mode_fields(
coordinates: Sequence[jax.Array],
propagation_axis: int,
*,
radius: float,
direction: Literal["+", "-"],
polarization_axis: int | None = None,
fixed_E_polarization_vector: tuple[float, float, float] | None = None,
fixed_H_polarization_vector: tuple[float, float, float] | None = None,
azimuth_angle: float = 0.0,
elevation_angle: float = 0.0,
divergence_angle: float = 0.0,
center: tuple[float, float] = (0.0, 0.0),
wavelength: float | jax.Array,
refractive_index: float | jax.Array = 1.0,
dtype: jnp.dtype = jnp.float32,
) -> tuple[jax.Array, jax.Array]:
"""Build the transverse profile of a Gaussian beam on a detector / source plane.
The beam's cross-section at a single plane — no ``z`` dependence, waist evolution or
Gouy phase. ``radius`` is the spot size there and ``divergence_angle`` the wavefront
curvature; paraxially this is ``exp(i k r^2 / 2q)`` with
``1/q = tan(divergence_angle)/radius + 2i/(k radius^2)`` and the medium wavenumber
``k = 2 pi n / wavelength``.
Polarization and propagation direction follow :class:`~fdtdx.GaussianPlaneSource`
exactly, via :func:`~fdtdx.core.misc.tilted_polarization_vectors`. Phase is zero at the
beam center; without tilt or divergence the mode is real. Assumes ``mu_r = 1`` —
``|H| = n |E|`` and ``k`` derive from the permittivity alone.
Args:
coordinates: ``(X, Y, Z)`` cell-center meshgrids of shape ``grid_shape`` (singleton
on ``propagation_axis``), center-origin like the grid.
propagation_axis: Physical propagation axis (0=x, 1=y, 2=z).
radius: Positive ``1/e`` amplitude radius at this plane, in metres, measured in the
beam's own cross-section.
direction: ``"+"`` (forward) or ``"-"`` (backward) along ``propagation_axis``.
polarization_axis: Transverse axis the E field points along. Mutually exclusive
with ``fixed_E/H_polarization_vector``. Defaults to the first transverse axis.
fixed_E_polarization_vector: Explicit E polarization 3-vector (mirrors the source).
fixed_H_polarization_vector: Explicit H polarization 3-vector (mirrors the source).
azimuth_angle: Propagation tilt around the vertical axis, in degrees.
elevation_angle: Propagation tilt around the horizontal axis, in degrees.
divergence_angle: Wavefront cone half-angle at this plane, in degrees, taken at the
``1/e`` radius: ``tan(angle) = radius / R``. ``0.0`` is flat phase (collimated,
or a plane at the waist); positive diverges, negative converges. Far from the
waist this equals the far-field divergence ``lambda / (pi w_0 n)``.
center: Transverse center offset ``(off_t0, off_t1)`` in metres, ascending axes.
wavelength: Vacuum wavelength in metres, setting the tilt ramp and the curvature.
refractive_index: Local medium index, for the ``|H| = n |E|`` ratio and ``k``.
dtype: Float dtype used to build polarization/rotation vectors.
Returns:
``(mode_E, mode_H)``, each ``(3, *grid_shape)``; complex under tilt or divergence.
Raises:
ValueError: If ``radius`` is not positive, or ``polarization_axis`` conflicts with
an explicit polarization vector / is not transverse.
"""
if radius <= 0:
raise ValueError(f"radius must be positive, got {radius}")
if polarization_axis is not None and (
fixed_E_polarization_vector is not None or fixed_H_polarization_vector is not None
):
raise ValueError("Specify either polarization_axis or fixed_E/H_polarization_vector, not both")
if fixed_E_polarization_vector is None and fixed_H_polarization_vector is None:
pol_axis = get_transverse_axes(propagation_axis)[0] if polarization_axis is None else polarization_axis
if pol_axis == propagation_axis:
raise ValueError(
f"polarization_axis ({pol_axis}) must be transverse to the propagation axis ({propagation_axis})"
)
e_vec = [0.0, 0.0, 0.0]
e_vec[pol_axis] = 1.0
fixed_E_polarization_vector = (e_vec[0], e_vec[1], e_vec[2])
# E/H polarization unit vectors and the (tilted) wave vector — same derivation as the
# plane sources (degrees -> radians; a zero angle is the identity rotation).
e_pol, h_pol, wave_vector = tilted_polarization_vectors(
direction=direction,
propagation_axis=propagation_axis,
fixed_E_polarization_vector=fixed_E_polarization_vector,
fixed_H_polarization_vector=fixed_H_polarization_vector,
azimuth_radians=jnp.asarray(np.deg2rad(azimuth_angle), dtype=dtype),
elevation_radians=jnp.asarray(np.deg2rad(elevation_angle), dtype=dtype),
dtype=dtype,
)
is_tilted = azimuth_angle != 0.0 or elevation_angle != 0.0
is_curved = divergence_angle != 0.0
t0, t1 = get_transverse_axes(propagation_axis)
transverse_0 = coordinates[t0] - center[0]
transverse_1 = coordinates[t1] - center[1]
# Radius measured in the beam's own cross-section, i.e. projected onto the plane
# orthogonal to the wave vector (same footprint as GaussianPlaneSource). Since the
# propagation-axis coordinate is zero here, |r_perp|^2 = |r|^2 - (r . k_hat)^2.
along_k = wave_vector[t0] * transverse_0 + wave_vector[t1] * transverse_1
r_perp_squared = transverse_0**2 + transverse_1**2
if is_tilted:
r_perp_squared = r_perp_squared - along_k**2
amplitude = jnp.exp(-r_perp_squared / (radius**2))
if is_tilted or is_curved:
wavenumber = 2.0 * jnp.pi * refractive_index / wavelength
phase_arg = jnp.zeros_like(amplitude)
if is_tilted:
phase_arg = phase_arg + wavenumber * along_k # tilted wavefront ramp
if is_curved:
inv_curvature_radius = float(np.tan(np.deg2rad(divergence_angle))) / radius
phase_arg = phase_arg + 0.5 * wavenumber * inv_curvature_radius * r_perp_squared
amplitude = amplitude * jnp.exp(1j * phase_arg)
mode_E = amplitude[None, ...] * e_pol[:, None, None, None]
mode_H = amplitude[None, ...] * (refractive_index * h_pol)[:, None, None, None]
return mode_E, mode_H
def gaussian_mode_function(
*,
radius: float,
direction: Literal["+", "-"],
polarization_axis: int | None = None,
fixed_E_polarization_vector: tuple[float, float, float] | None = None,
fixed_H_polarization_vector: tuple[float, float, float] | None = None,
azimuth_angle: float = 0.0,
elevation_angle: float = 0.0,
divergence_angle: float = 0.0,
center: tuple[float, float] = (0.0, 0.0),
) -> Callable[..., tuple[jax.Array, jax.Array]]:
"""Return a ``mode_function`` for :class:`CustomModeOverlapDetector` (analytic Gaussian).
The returned callable derives the local refractive index from the permittivity slice and
delegates to :func:`gaussian_mode_fields`; see that function for the arguments,
including the ``mu_r = 1`` assumption.
"""
def _mode_function(
*,
coordinates: Sequence[jax.Array],
frequency: float,
propagation_axis: int,
inv_permittivity: jax.Array,
) -> tuple[jax.Array, jax.Array]:
"""Evaluate the Gaussian at one frequency, taking the index from the plane."""
refractive_index = jnp.sqrt(jnp.mean(1.0 / inv_permittivity))
return gaussian_mode_fields(
coordinates,
propagation_axis,
radius=radius,
direction=direction,
polarization_axis=polarization_axis,
fixed_E_polarization_vector=fixed_E_polarization_vector,
fixed_H_polarization_vector=fixed_H_polarization_vector,
azimuth_angle=azimuth_angle,
elevation_angle=elevation_angle,
divergence_angle=divergence_angle,
center=center,
wavelength=c / frequency,
refractive_index=refractive_index,
dtype=inv_permittivity.dtype,
)
return _mode_function
@autoinit
class BaseModeOverlapDetector(PhasorDetector, ABC):
"""Abstract base for mode-overlap detectors.
Owns everything that is independent of *how* the reference mode is produced: the
stored reference mode fields (``_mode_E`` / ``_mode_H`` of shape
``(num_freqs, 3, *spatial)``), the detector-plane face-area weights, and the overlap
integral itself (:meth:`compute_overlap` / :meth:`compute_overlap_to_mode`).
Subclasses only implement :meth:`_compute_mode_fields`, which returns the reference
mode for a single frequency on the detector plane. :class:`ModeOverlapDetector` solves
it with the waveguide mode solver; :class:`CustomModeOverlapDetector` /
:class:`GaussianModeOverlapDetector` evaluate a user-supplied / analytic mode instead.
``compute_overlap()`` returns a complex array of shape ``(num_freqs,)``, where
``num_freqs = len(wave_characters)``.
"""
#: Cannot be specified here since the detector needs all components.
components: Sequence[Literal["Ex", "Ey", "Ez", "Hx", "Hy", "Hz"]] = frozen_field(
default=("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"),
init=False, # in this detector, we always want all components. Do not give user a choice
)
#: Cannot be specified here since plotting a single scalar is useless.
plot: bool = frozen_field(default=False, init=False) # single scalar is useless for plotting
_mode_E: jax.Array = private_field()
_mode_H: jax.Array = private_field()
_mode_neff: jax.Array = private_field() # not required for detection, used for inspection
_cached_face_area_weights: jax.Array = private_field()
#: Explicit propagation axis for an extruded 2-D simulation. It is only
#: needed when both the detector normal and an invariant domain axis have
#: one cell, making shape-based inference ambiguous.
fixed_propagation_axis: int | None = frozen_field(default=None)
@property
def propagation_axis(self) -> int:
"""Physical axis normal to the detector plane, i.e. the singleton grid axis."""
if self.fixed_propagation_axis is not None:
if self.fixed_propagation_axis not in (0, 1, 2):
raise ValueError("fixed_propagation_axis must be 0, 1, or 2")
if self.grid_shape[self.fixed_propagation_axis] != 1:
raise ValueError(
"fixed_propagation_axis must select a one-cell detector dimension, "
f"got axis {self.fixed_propagation_axis} for shape {self.grid_shape}"
)
return self.fixed_propagation_axis
if sum([a == 1 for a in self.grid_shape]) != 1:
raise ValueError(
"Invalid ModeOverlapDetector shape: propagation axis is ambiguous; "
"set fixed_propagation_axis "
f"for shape {self.grid_shape}"
)
return self.grid_shape.index(1)
def place_on_grid(
self: Self,
grid_slice_tuple: SliceTuple3D,
config: SimulationConfig,
key: jax.Array,
) -> Self:
"""Place the detector and cache the plane's face-area weights."""
self = super().place_on_grid(
grid_slice_tuple=grid_slice_tuple,
config=config,
key=key,
)
grid = self._config.resolved_grid
if grid is not None:
weights = grid.face_area(axis=self.propagation_axis, slice_tuple=self.grid_slice_tuple)
else:
spacing = self._config.uniform_spacing()
weights = jnp.ones(self.grid_shape, dtype=jnp.float32) * spacing * spacing
weights = weights / weights.mean() # consistent with scaling in compute_mode
self = self.aset("_cached_face_area_weights", weights, create_new_ok=True)
return self
def _face_area_weights(self) -> jax.Array:
"""Return detector-plane face areas for mode-overlap integration."""
return self._cached_face_area_weights
def _plane_coordinates(self) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Return ``(X, Y, Z)`` cell-center coordinate meshgrids for the detector plane.
Each array has shape ``grid_shape`` (singleton on the propagation axis). Used by
subclasses that evaluate an analytic mode profile on the actual placed grid. When
the grid is resolved the coordinates follow the center-origin convention (#363);
the uniform fallback (unresolved policy, only hit in lightweight tests) is
corner-relative.
"""
grid = self._config.resolved_grid
axis_centers: list[jax.Array] = []
for axis in range(3):
lower, upper = self.grid_slice_tuple[axis]
if grid is not None:
centers = jnp.asarray(grid.centers(axis))[lower:upper]
else:
spacing = self._config.uniform_spacing()
centers = (jnp.arange(lower, upper) + 0.5) * spacing
axis_centers.append(centers)
x_coords, y_coords, z_coords = jnp.meshgrid(*axis_centers, indexing="ij")
return x_coords, y_coords, z_coords
@staticmethod
def _as_real_inv_permittivity(inv_permittivity_slice: jax.Array) -> jax.Array:
"""Reduce a (possibly complex) effective inverse permittivity to ``1/Re(eps)``.
:meth:`apply` hands subclasses the full complex ``1/eps(omega_c)`` so the waveguide
mode solver can model material loss. Analytic reference modes only use the real
index (loss is integrated by the ADE / conductivity updates, not by the reference
profile), so they reduce to ``1/Re(eps)`` — the same value
:func:`~fdtdx.dispersion.effective_inv_permittivity` returns. Real input is
returned unchanged.
"""
if jnp.iscomplexobj(inv_permittivity_slice):
return 1.0 / jnp.real(1.0 / inv_permittivity_slice)
return inv_permittivity_slice
@abstractmethod
def _compute_mode_fields(
self,
*,
wave_character: WaveCharacter,
inv_permittivity_slice: jax.Array,
inv_permeability_slice: jax.Array | float,
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Produce the reference mode ``(mode_E, mode_H, mode_neff)`` for one frequency.
``mode_E`` / ``mode_H`` must have shape ``(3, *grid_shape)`` (singleton on the
propagation axis); ``mode_neff`` is a complex/real scalar used only for
inspection. ``inv_permittivity_slice`` is already restricted to the detector plane
and, in a dispersive or conductive medium, corrected to the full complex effective
permittivity at the frequency's carrier (see :meth:`apply`). Subclasses that only
support a real index reduce it via :meth:`_as_real_inv_permittivity`.
"""
raise NotImplementedError
def apply(
self,
key: jax.Array,
inv_permittivities: jax.Array,
inv_permeabilities: jax.Array | float,
dispersive_c1: jax.Array | None = None,
dispersive_c2: jax.Array | None = None,
dispersive_c3: jax.Array | None = None,
electric_conductivity: jax.Array | None = None,
dispersive_c4: jax.Array | None = None,
) -> Self:
"""Precompute and store the reference mode fields for every wave character."""
del key
inv_permittivity_slice = inv_permittivities[:, *self.grid_slice]
if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0:
inv_permeability_slice = inv_permeabilities[:, *self.grid_slice]
else:
inv_permeability_slice = inv_permeabilities
c1_slice = c2_slice = c3_slice = c4_slice = None
if dispersive_c1 is not None and dispersive_c2 is not None and dispersive_c3 is not None:
c1_slice = dispersive_c1[:, :, *self.grid_slice]
c2_slice = dispersive_c2[:, :, *self.grid_slice]
c3_slice = dispersive_c3[:, :, *self.grid_slice]
c4_slice = None if dispersive_c4 is None else dispersive_c4[:, :, *self.grid_slice]
# The reference mode is solved against the FULL complex epsilon at each
# carrier frequency (eps_inf + chi(omega) + i*sigma/(eps0*omega)), so the
# overlap basis and effective index reflect material loss. With no loss
# the complex helper is skipped and the real lossless path is preserved.
sigma_slice = None if electric_conductivity is None else electric_conductivity[:, *self.grid_slice]
conductivity_spacing = (
None if sigma_slice is None else constants.c * self._config.time_step_duration / self._config.courant_number
)
all_mode_Es: list[jax.Array] = []
all_mode_Hs: list[jax.Array] = []
all_mode_neffs: list[jax.Array] = []
for wc in self.wave_characters:
inv_eps_i = inv_permittivity_slice
if c1_slice is not None or sigma_slice is not None:
inv_eps_i = effective_complex_inv_permittivity(
inv_eps=inv_permittivity_slice,
omega=2.0 * np.pi * wc.get_frequency(),
dt=self._config.time_step_duration,
c1=c1_slice,
c2=c2_slice,
c3=c3_slice,
c4=c4_slice,
electric_conductivity=sigma_slice,
conductivity_spacing=conductivity_spacing,
)
mode_E, mode_H, mode_neff = self._compute_mode_fields(
wave_character=wc,
inv_permittivity_slice=inv_eps_i,
inv_permeability_slice=inv_permeability_slice,
)
all_mode_Es.append(mode_E)
all_mode_Hs.append(mode_H)
all_mode_neffs.append(mode_neff)
self = self.aset("_mode_E", jnp.stack(all_mode_Es, axis=0), create_new_ok=True)
self = self.aset("_mode_H", jnp.stack(all_mode_Hs, axis=0), create_new_ok=True)
self = self.aset("_mode_neff", jnp.stack(all_mode_neffs, axis=0), create_new_ok=True)
return self
def compute_overlap_to_mode(
self,
state: DetectorState,
mode_E: jax.Array,
mode_H: jax.Array,
wave_character_index: int = 0,
) -> jax.Array:
"""Compute the overlap integral of *one* mode against phasors at ``wave_character_index``.
Args:
state: Detector state holding the phasor array of shape
``(1, num_freqs, 6, *spatial)``.
mode_E: Electric mode field of shape ``(3, *spatial)``.
mode_H: Magnetic mode field of shape ``(3, *spatial)``.
wave_character_index: Index into the phasor frequency axis to use.
Returns:
Complex scalar overlap coefficient.
"""
# shape (time step, num_freqs, num_components, *spatial)
# time steps is always 1 and num_components always 6
phasors = state["phasor"]
phasors_E, phasors_H = phasors[0, wave_character_index, :3], phasors[0, wave_character_index, 3:]
E_cross_H_star_sim = jnp.cross(
mode_E,
jnp.conj(phasors_H),
axis=0,
)[self.propagation_axis]
E_star_cross_H_sim = jnp.cross(
jnp.conj(phasors_E),
mode_H,
axis=0,
)[self.propagation_axis]
integrand = E_cross_H_star_sim + E_star_cross_H_sim
integrand = integrand * self._face_area_weights()
alpha_coeff = jnp.sum(integrand)
# in pulsed mode return unscaled coefficient
if self.scaling_mode != "pulse":
alpha_coeff = alpha_coeff / 4.0
return alpha_coeff
def compute_overlap(
self,
state: DetectorState,
) -> jax.Array:
"""Compute mode overlaps at every frequency in ``wave_characters``.
Returns:
Complex array of shape ``(num_freqs,)``.
"""
if isinstance(self._mode_E, Null) or isinstance(self._mode_H, Null):
raise Exception("Need to call apply on the mode-overlap detector before calling compute_overlap!")
overlaps = [
self.compute_overlap_to_mode(
state=state,
mode_E=self._mode_E[i],
mode_H=self._mode_H[i],
wave_character_index=i,
)
for i in range(len(self.wave_characters))
]
return jnp.stack(overlaps, axis=0)
[docs]
@autoinit
class ModeOverlapDetector(BaseModeOverlapDetector):
"""
Detector for measuring the overlap of a waveguide mode with the simulation fields.
This detector computes the overlap integral at every frequency in ``wave_characters``,
enabling broadband frequency-domain analysis of the electromagnetic fields.
The reference mode is obtained from the waveguide mode solver (``compute_mode``). For a
user-supplied or analytic reference mode (e.g. a Gaussian beam) use
:class:`CustomModeOverlapDetector` or :class:`GaussianModeOverlapDetector` instead;
both share the same overlap machinery via :class:`BaseModeOverlapDetector`.
The mode overlap is calculated by integrating the cross product of the mode fields
with the simulation fields over a cross-sectional plane. This is useful for
analyzing waveguide coupling efficiency, transmission coefficients, and modal
decomposition of electromagnetic fields.
``compute_overlap()`` returns a complex array of shape ``(num_freqs,)``, where
``num_freqs = len(wave_characters)``.
"""
#: Direction of mode propagation, either "+" (forward) or "-" (backward).
#: Determines which direction along the waveguide axis the mode is assumed to propagate.
direction: Literal["+", "-"] = frozen_field()
#: Index of the waveguide mode to use for overlap calculation.
#: Defaults to 0 (fundamental mode). Higher indices correspond to higher-order modes.
mode_index: int = frozen_field(default=0)
#: Optional polarization filter for the mode calculation.
#: Can be "te" (transverse electric), "tm" (transverse magnetic), or None (no filtering).
#: When specified, only modes of the given polarization type are considered. Defaults to None.
filter_pol: Literal["te", "tm"] | None = frozen_field(default=None)
#: Bend radius of the waveguide in meters. When set, the mode solver accounts for the conformal
#: transformation introduced by the bend. Must be set together with bend_axis. Defaults to None
#: (straight waveguide).
bend_radius: float | None = frozen_field(default=None)
#: Physical axis index (0=x, 1=y, 2=z) pointing from the waveguide center toward the center of
#: curvature. Must differ from the propagation axis. Required when bend_radius is set.
bend_axis: int | None = frozen_field(default=None)
#: Symmetry-plane condition at the min edge of each transverse axis (the two non-propagation
#: physical axes, in increasing-index order): ``0`` = PEC mirror (electric wall, the default),
#: ``1`` = PMC mirror (magnetic wall). Set this only for a **hand-built** half/quarter domain
#: (your own PEC/PMC boundary at the min edge), where it asks the mode solver for its own
#: symmetric solve; it must then match the corresponding ModePlaneSource. It is ignored when
#: ``config.symmetry`` performs the reduction: there the reference mode is solved on the mirrored
#: full cross-section and restricted, exactly like the mode source, so the overlap stays
#: consistent automatically.
symmetry: tuple[int, int] = frozen_field(default=(0, 0))
def place_on_grid(
self: Self,
grid_slice_tuple: SliceTuple3D,
config: SimulationConfig,
key: jax.Array,
) -> Self:
"""Place the detector and validate the bend arguments."""
self = super().place_on_grid(
grid_slice_tuple=grid_slice_tuple,
config=config,
key=key,
)
if (self.bend_radius is None) != (self.bend_axis is None):
raise ValueError("bend_radius and bend_axis must both be set or both be None")
if self.bend_axis is not None and self.bend_axis == self.propagation_axis:
raise ValueError(
f"bend_axis ({self.bend_axis}) must differ from the propagation axis ({self.propagation_axis})"
)
return self
def _transverse_edge_coordinates(self) -> tuple[jax.Array, jax.Array] | None:
"""Return physical transverse edge coordinates for the mode solver.
Tidy3D can solve modes on rectilinear non-uniform grids when supplied
with edge-coordinate arrays. Returning ``None`` keeps the uniform scalar
spacing path for legacy configurations and older tests.
"""
grid = self._config.resolved_grid
if grid is None:
return None
transverse_edges = []
for axis in range(3):
if axis == self.propagation_axis:
continue
lower, upper = self.grid_slice_tuple[axis]
transverse_edges.append(grid.edges(axis)[lower : upper + 1])
e0, e1 = transverse_edges
return e0, e1
def _mode_solver_resolution(self) -> float:
"""Return scalar resolution only for legacy uniform mode-solver setup.
``compute_mode`` ignores this value when explicit transverse coordinates
are supplied. For non-uniform grids we pass a harmless finite value so
the compatibility argument does not force a uniform-grid check.
"""
if self._config.has_nonuniform_grid:
assert self._config.resolved_grid is not None
return self._config.resolved_grid.min_spacing
return self._config.uniform_spacing()
def _compute_mode_fields(
self,
*,
wave_character: WaveCharacter,
inv_permittivity_slice: jax.Array,
inv_permeability_slice: jax.Array | float,
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Solve the reference mode with the waveguide mode solver."""
mirrored_axes = self.symmetry_mirror_axes(exclude_axis=self.propagation_axis)
if mirrored_axes:
# Symmetry-reduced cross-section: mirror it back to the full one, solve there and
# restrict — the same route ModePlaneSource takes, so source and reference mode agree.
return compute_mode_symmetry_reduced(
mirrored_axes=mirrored_axes,
walls={axis: self._config.symmetry[axis] for axis in mirrored_axes},
frequency=wave_character.get_frequency(),
inv_permittivities=inv_permittivity_slice,
inv_permeabilities=inv_permeability_slice,
resolution=self._mode_solver_resolution(),
direction=self.direction,
mode_index=self.mode_index,
filter_pol=self.filter_pol,
dtype=self._config.dtype,
bend_radius=self.bend_radius,
bend_axis=self.bend_axis,
transverse_coords=self._transverse_edge_coordinates(),
object_name=self.name,
fixed_propagation_axis=self.propagation_axis,
)
return compute_mode(
frequency=wave_character.get_frequency(),
inv_permittivities=inv_permittivity_slice,
inv_permeabilities=inv_permeability_slice,
resolution=self._mode_solver_resolution(),
direction=self.direction,
mode_index=self.mode_index,
filter_pol=self.filter_pol,
dtype=self._config.dtype,
bend_radius=self.bend_radius,
bend_axis=self.bend_axis,
symmetry=self.symmetry,
transverse_coords=self._transverse_edge_coordinates(),
fixed_propagation_axis=self.propagation_axis,
)
@autoinit
class TimeDomainModeOverlapDetector(ModeOverlapDetector):
"""Record one directional waveguide-mode coordinate versus time.
The detector solves the same fixed reference mode as
:class:`ModeOverlapDetector`, but applies the reciprocal power-overlap
integral directly to every real-valued FDTD sample instead of first
accumulating a frequency-domain phasor. Its ``fields`` output has shape
``(recorded_time_steps, 1)`` and can therefore replace a reduced one-point
:class:`~fdtdx.FieldDetector` in temporal objectives without retaining an
entire transverse field plane.
Exactly one ``wave_character`` is required. Its frequency only defines
the fixed transverse reference mode; no carrier is mixed into the stored
trace. The reference eigenvector is phase-aligned by its largest electric
component and projected to a real quadrature, which is appropriate for a
lossless mode used with real time-domain fields. Use
:class:`ModeOverlapDetector` when a complex, frequency-domain coefficient
is required (for example in a lossy guide).
"""
dtype: jnp.dtype = frozen_field(default=jnp.float32)
#: ``"directional"`` uses the reciprocal E/H power overlap and selects
#: one propagation direction. ``"electric"`` records the electric-mode
#: inner product, making the detector exactly reciprocal to an impressed
#: electric :class:`~fdtdx.ModeProfileCurrentSource` on the same plane.
overlap_kind: Literal["directional", "electric"] = frozen_field(default="directional")
def __post_init__(self) -> None:
if self.dtype not in (jnp.float32, jnp.float64):
raise ValueError("TimeDomainModeOverlapDetector requires a real floating dtype")
if len(self.wave_characters) != 1:
raise ValueError("TimeDomainModeOverlapDetector requires exactly one wave_character")
if self.overlap_kind not in ("directional", "electric"):
raise ValueError("overlap_kind must be 'directional' or 'electric'")
def _calculate_on_list(self) -> list[bool]:
"""Use the ordinary time-domain switch, not phasor DFT thinning."""
return Detector._calculate_on_list(self)
def _num_latent_time_steps(self) -> int:
return Detector._num_latent_time_steps(self)
def _shape_dtype_single_time_step(self) -> dict[str, jax.ShapeDtypeStruct]:
return {"fields": jax.ShapeDtypeStruct(shape=(1,), dtype=self.dtype)}
def update(
self,
time_step: jax.Array,
E: jax.Array,
H: jax.Array,
state: DetectorState,
inv_permittivity: jax.Array,
inv_permeability: jax.Array | float,
) -> DetectorState:
del inv_permittivity, inv_permeability
if isinstance(self._mode_E, Null) or isinstance(self._mode_H, Null):
raise RuntimeError("apply() must initialize the reference mode before recording")
mode_E_complex = self._mode_E[0]
mode_H_complex = self._mode_H[0]
flat_e = mode_E_complex.reshape(-1)
dominant = flat_e[jnp.argmax(jnp.abs(flat_e))]
rotation = jnp.exp(-1j * jnp.angle(dominant))
mode_E = jnp.real(mode_E_complex * rotation)
mode_H = jnp.real(mode_H_complex * rotation)
if self.overlap_kind == "electric":
overlap = jnp.sum(jnp.sum(mode_E * E, axis=0) * self._face_area_weights())
else:
# Lorentz-reciprocal directional overlap. For an exact forward
# copy of the reference mode the two terms add; for its backward
# partner the magnetic-field sign reversal makes them cancel.
mode_e_cross_h = jnp.cross(mode_E, H, axis=0)[self.propagation_axis]
e_cross_mode_h = jnp.cross(E, mode_H, axis=0)[self.propagation_axis]
overlap = 0.5 * jnp.sum((mode_e_cross_h + e_cross_mode_h) * self._face_area_weights())
arr_idx = self._time_step_to_arr_idx[time_step]
fields = state["fields"].at[arr_idx, 0].set(overlap.astype(self.dtype))
return {"fields": fields}
@autoinit
class CustomModeOverlapDetector(BaseModeOverlapDetector):
"""Mode-overlap detector using a user-provided reference mode.
Instead of solving the mode with the waveguide mode solver, the reference mode is
produced by ``mode_function`` — a callable evaluated on the detector plane during
:meth:`apply`. This enables overlap against an arbitrary mode (e.g. an analytic
Gaussian beam, a fiber mode, or a mode imported from another tool) and avoids the
tidy3d mode-solver dependency.
``mode_function`` is called once per frequency with keyword arguments::
mode_function(
coordinates, # (X, Y, Z) cell-center meshgrids, each (3==None) shape grid_shape
frequency, # float, Hz
propagation_axis, # int, 0/1/2
inv_permittivity, # effective eps slice on the plane (n_comp, *grid_shape)
) -> (mode_E, mode_H) # each (3, *grid_shape)
and must return the E and H fields in FDTDX's :math:`\\eta_0`-normalized convention.
Use :func:`gaussian_mode_function` for a ready-made analytic Gaussian, or
:class:`GaussianModeOverlapDetector` for the same as a configurable detector class.
Dispersive media are handled by the shared :meth:`BaseModeOverlapDetector.apply`: in a
dispersive simulation ``inv_permittivity`` is the *effective* inverse permittivity
:math:`1/\\mathrm{Re}(\\varepsilon_\\infty + \\chi(\\omega_c))` at the wave character's
carrier frequency (same correction as ``ModeOverlapDetector`` / ``ModePlaneSource``),
not :math:`\\varepsilon_\\infty`. A frequency-aware ``mode_function`` therefore sees the
true medium index per frequency.
"""
#: Callable producing the reference ``(mode_E, mode_H)`` on the detector plane.
#: See the class docstring for the exact keyword-argument signature.
mode_function: Callable[..., tuple[jax.Array, jax.Array]] = frozen_field()
#: Whether to renormalize the provided mode to unit Poynting flux over the detector
#: plane (matching what the mode solver does). Defaults to True.
normalize: bool = frozen_field(default=True)
def _compute_mode_fields(
self,
*,
wave_character: WaveCharacter,
inv_permittivity_slice: jax.Array,
inv_permeability_slice: jax.Array | float,
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Evaluate the user-supplied ``mode_function`` on the detector plane."""
del inv_permeability_slice
coordinates = self._plane_coordinates()
inv_permittivity_slice = self._as_real_inv_permittivity(inv_permittivity_slice)
mode_E, mode_H = self.mode_function(
coordinates=coordinates,
frequency=wave_character.get_frequency(),
propagation_axis=self.propagation_axis,
inv_permittivity=inv_permittivity_slice,
)
# Nominal effective index for inspection only (mean medium index on the plane).
mode_neff = jnp.sqrt(jnp.mean(1.0 / inv_permittivity_slice))
if self.normalize:
mode_E, mode_H = normalize_by_poynting_flux(
mode_E,
mode_H,
axis=self.propagation_axis,
area_weights=self._face_area_weights(),
)
return mode_E, mode_H, mode_neff
@autoinit
class GaussianModeOverlapDetector(BaseModeOverlapDetector):
"""Mode-overlap detector using an analytic Gaussian profile as the reference mode.
Wrapper around :func:`gaussian_mode_fields`. Everything is stated *at the detector
plane*: ``mode_radius`` is the spot size there and ``divergence_angle`` the wavefront
curvature there, so a collimated beam (the default), a tilted beam, and a diverging or
converging one are all expressible without a mode solver. The reference is a transverse
cross-section, not a propagating beam. Polarization and propagation angle are configured
exactly like :class:`~fdtdx.GaussianPlaneSource`. Assumes ``mu_r = 1``.
Dispersion-aware: ``n`` (and the wavenumber ``k = n ω/c``) come from the effective
permittivity :math:`\\mathrm{Re}(\\varepsilon_\\infty + \\chi(\\omega_c))` at each wave
character's carrier frequency via :meth:`BaseModeOverlapDetector.apply`, not
:math:`\\varepsilon_\\infty`.
"""
#: Gaussian ``1/e`` amplitude radius *at the detector plane*, in metres. Must be
#: positive. This is the spot size on this plane, not the waist of the beam.
mode_radius: float = frozen_field()
#: Direction of propagation, "+" (forward) or "-" (backward) along the plane normal.
direction: Literal["+", "-"] = frozen_field()
#: Convenience polarization selector — transverse axis the E field points along.
#: Mutually exclusive with ``fixed_E/H_polarization_vector``. Defaults to the first
#: transverse axis (ascending index order).
polarization_axis: int | None = frozen_field(default=None)
#: Explicit electric polarization 3-vector (mirrors ``GaussianPlaneSource``).
fixed_E_polarization_vector: tuple[float, float, float] | None = frozen_field(default=None)
#: Explicit magnetic polarization 3-vector (mirrors ``GaussianPlaneSource``).
fixed_H_polarization_vector: tuple[float, float, float] | None = frozen_field(default=None)
#: Propagation tilt around the vertical axis, in degrees (off-normal incidence).
azimuth_angle: float = frozen_field(default=0.0)
#: Propagation tilt around the horizontal axis, in degrees (off-normal incidence).
elevation_angle: float = frozen_field(default=0.0)
#: Wavefront cone half-angle at the detector plane, in degrees, taken at the ``1/e``
#: radius: ``tan(angle) = mode_radius / R``. ``0.0`` (default) is a flat phase front —
#: a collimated beam, or a plane at the beam waist. Positive diverges, negative
#: converges. Far from the waist this equals the far-field divergence
#: ``lambda / (pi w_0 n)``. Curves the wavefront; unlike ``azimuth_angle`` /
#: ``elevation_angle`` it does not tilt the propagation direction.
divergence_angle: float = frozen_field(default=0.0)
#: Transverse center offset ``(off_t0, off_t1)`` in metres for the two transverse axes
#: (ascending index order), in the same physical coordinates as the grid.
center: tuple[float, float] = frozen_field(default=(0.0, 0.0))
#: Whether to renormalize the mode to unit Poynting flux over the detector plane.
normalize: bool = frozen_field(default=True)
def place_on_grid(
self: Self,
grid_slice_tuple: SliceTuple3D,
config: SimulationConfig,
key: jax.Array,
) -> Self:
"""Place the detector, validate ``mode_radius`` and warn on a truncated beam."""
self = super().place_on_grid(grid_slice_tuple=grid_slice_tuple, config=config, key=key)
if self.mode_radius <= 0:
raise ValueError(f"mode_radius must be positive, got {self.mode_radius}")
self._warn_if_truncated()
return self
def _warn_if_truncated(self) -> None:
"""Warn if the detector plane clips the reference mode above ~1% amplitude.
The overlap only integrates over the detector's own footprint, so a beam wider than
the plane is silently under-reported. A clean measurement needs a transverse
half-extent of ``>~ 3 x mode_radius`` (edge amplitude ``~0.01%``). Assumes the beam
is centered, so this is the best-case truncation amplitude.
"""
grid = self._config.resolved_grid
half_extents = []
for axis in get_transverse_axes(self.propagation_axis):
lower, upper = self.grid_slice_tuple[axis]
if grid is not None:
extent = grid.axis_extent(axis, (lower, upper))
else:
extent = (upper - lower) * self._config.uniform_spacing()
half_extents.append(0.5 * extent)
edge_distance = min(half_extents)
truncation_amplitude = float(np.exp(-((edge_distance / self.mode_radius) ** 2)))
if truncation_amplitude > 0.01:
warnings.warn(
f"GaussianModeOverlapDetector '{self.name}': the detector plane truncates the reference "
f"mode at {truncation_amplitude * 100:.1f}% amplitude (nearest edge at "
f"{edge_distance / self.mode_radius:.2f} x mode_radius), so the overlap under-reports "
f"the coupled power. Enlarge the detector to >~ 3 x mode_radius transversely.",
UserWarning,
stacklevel=2,
)
def _compute_mode_fields(
self,
*,
wave_character: WaveCharacter,
inv_permittivity_slice: jax.Array,
inv_permeability_slice: jax.Array | float,
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Evaluate the analytic Gaussian profile on the detector plane."""
del inv_permeability_slice
coordinates = self._plane_coordinates()
inv_permittivity_slice = self._as_real_inv_permittivity(inv_permittivity_slice)
refractive_index = jnp.sqrt(jnp.mean(1.0 / inv_permittivity_slice))
mode_E, mode_H = gaussian_mode_fields(
coordinates,
self.propagation_axis,
radius=self.mode_radius,
direction=self.direction,
polarization_axis=self.polarization_axis,
fixed_E_polarization_vector=self.fixed_E_polarization_vector,
fixed_H_polarization_vector=self.fixed_H_polarization_vector,
azimuth_angle=self.azimuth_angle,
elevation_angle=self.elevation_angle,
divergence_angle=self.divergence_angle,
center=self.center,
wavelength=wave_character.get_wavelength(),
refractive_index=refractive_index,
dtype=self._config.dtype,
)
if self.normalize:
mode_E, mode_H = normalize_by_poynting_flux(
mode_E,
mode_H,
axis=self.propagation_axis,
area_weights=self._face_area_weights(),
)
return mode_E, mode_H, refractive_index
