diff --git a/docs/source/07_api.rst b/docs/source/07_api.rst index ffb2922c..55e90083 100644 --- a/docs/source/07_api.rst +++ b/docs/source/07_api.rst @@ -7,6 +7,7 @@ API fdtdx.apply_params fdtdx.ArrayContainer fdtdx.autoinit + fdtdx.BaseModeOverlapDetector fdtdx.BinaryMedianFilterModule fdtdx.BlochBoundary fdtdx.boundary_objects_from_config @@ -30,6 +31,7 @@ API fdtdx.compute_pole_coefficients_tensor fdtdx.compute_poynting_flux fdtdx.ConnectHolesAndStructures + fdtdx.CustomModeOverlapDetector fdtdx.CustomTimeSignalProfile fdtdx.Cylinder fdtdx.Detector @@ -60,6 +62,8 @@ API fdtdx.frozen_field fdtdx.frozen_private_field fdtdx.full_backward + fdtdx.gaussian_mode_function + fdtdx.GaussianModeOverlapDetector fdtdx.GaussianPlaneSource fdtdx.GaussianPulseProfile fdtdx.GaussianSmoothing2D diff --git a/src/fdtdx/__init__.py b/src/fdtdx/__init__.py index 71eb4e4e..577a9c33 100644 --- a/src/fdtdx/__init__.py +++ b/src/fdtdx/__init__.py @@ -68,7 +68,13 @@ FieldProjectionCartesianDetector, FieldProjectionKSpaceDetector, ) -from fdtdx.objects.detectors.mode import ModeOverlapDetector +from fdtdx.objects.detectors.mode import ( + BaseModeOverlapDetector, + CustomModeOverlapDetector, + GaussianModeOverlapDetector, + ModeOverlapDetector, + gaussian_mode_function, +) from fdtdx.objects.detectors.phasor import PhasorDetector from fdtdx.objects.detectors.poynting_flux import ( ClosedSurfacePhasorPoyntingFluxDetector, @@ -163,6 +169,7 @@ __all__ = [ "ArrayContainer", + "BaseModeOverlapDetector", "BinaryMedianFilterModule", "BlochBoundary", "BoundaryConfig", @@ -173,6 +180,7 @@ "ClosestIndex", "Color", "ConnectHolesAndStructures", + "CustomModeOverlapDetector", "CustomTimeSignalProfile", "Cylinder", "Detector", @@ -193,6 +201,7 @@ "GDSLayerObject", "GDSLayerSpec", "GDSPortSpec", + "GaussianModeOverlapDetector", "GaussianPlaneSource", "GaussianPulseProfile", "GaussianSmoothing2D", @@ -281,6 +290,7 @@ "frozen_field", "frozen_private_field", "full_backward", + "gaussian_mode_function", "gds_layer_stack", "gds_layer_stack_from_component", "import_from_json", diff --git a/src/fdtdx/core/misc.py b/src/fdtdx/core/misc.py index 1719c35b..967d696b 100644 --- a/src/fdtdx/core/misc.py +++ b/src/fdtdx/core/misc.py @@ -6,8 +6,9 @@ import jax.numpy as jnp import numpy as np +from fdtdx.core.axis import get_oriented_transverse_axes from fdtdx.core.jax.pytrees import TreeClass, autoinit, frozen_field -from fdtdx.core.linalg import get_orthogonal_vector +from fdtdx.core.linalg import get_orthogonal_vector, get_wave_vector_raw, rotate_vector from fdtdx.materials import Material @@ -504,6 +505,54 @@ def normalize_polarization_for_source( return e_pol, h_pol +def tilted_polarization_vectors( + direction: Literal["+", "-"], + propagation_axis: int, + *, + fixed_E_polarization_vector: tuple[float, float, float] | None = None, + fixed_H_polarization_vector: tuple[float, float, float] | None = None, + azimuth_radians: float | jax.Array = 0.0, + elevation_radians: float | jax.Array = 0.0, + dtype: jnp.dtype = jnp.float32, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Resolve the E/H polarization unit vectors and wave vector of a plane wave. + + Shared by the linearly-polarized plane sources and the analytic Gaussian + mode-overlap detector so both derive polarization and off-normal propagation + identically. The raw E/H polarization comes from + :func:`normalize_polarization_for_source`, the wave vector from + :func:`~fdtdx.core.linalg.get_wave_vector_raw`, and all three are rotated by the + azimuth/elevation angles (radians) about the ``(horizontal, vertical, propagation)`` + axis triple via :func:`~fdtdx.core.linalg.rotate_vector` (a zero angle is the + identity, so normal incidence is unaffected). + + Args: + direction: ``"+"`` (forward) or ``"-"`` (backward) along ``propagation_axis``. + propagation_axis: Physical propagation axis (0=x, 1=y, 2=z). + fixed_E_polarization_vector: Explicit E polarization 3-vector, or ``None``. + fixed_H_polarization_vector: Explicit H polarization 3-vector, or ``None``. + azimuth_radians: Tilt around the vertical axis, in radians. + elevation_radians: Tilt around the horizontal axis, in radians. + dtype: Float dtype used to build the vectors. + + Returns: + ``(e_pol, h_pol, wave_vector)`` unit vectors, each of shape ``(3,)``. + """ + e_pol, h_pol = normalize_polarization_for_source( + direction=direction, + propagation_axis=propagation_axis, + fixed_E_polarization_vector=fixed_E_polarization_vector, + fixed_H_polarization_vector=fixed_H_polarization_vector, + dtype=dtype, + ) + wave_vector = get_wave_vector_raw(direction=direction, propagation_axis=propagation_axis, dtype=dtype) + axes_tpl = (*get_oriented_transverse_axes(propagation_axis), propagation_axis) + e_pol = rotate_vector(e_pol, azimuth_radians, elevation_radians, axes_tpl) + h_pol = rotate_vector(h_pol, azimuth_radians, elevation_radians, axes_tpl) + wave_vector = rotate_vector(wave_vector, azimuth_radians, elevation_radians, axes_tpl) + return e_pol, h_pol, wave_vector + + @overload def expand_to_3x3(arr: None) -> None: ... diff --git a/src/fdtdx/objects/detectors/mode.py b/src/fdtdx/objects/detectors/mode.py index 73734391..f6b6739a 100644 --- a/src/fdtdx/objects/detectors/mode.py +++ b/src/fdtdx/objects/detectors/mode.py @@ -1,3 +1,6 @@ +import warnings +from abc import ABC, abstractmethod +from collections.abc import Callable from typing import Literal, Self, Sequence import jax @@ -6,60 +9,203 @@ 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 +from fdtdx.core.wavelength import WaveCharacter from fdtdx.dispersion import effective_complex_inv_permittivity from fdtdx.objects.detectors.detector import DetectorState from fdtdx.objects.detectors.phasor import PhasorDetector from fdtdx.typing import SliceTuple3D -@autoinit -class ModeOverlapDetector(PhasorDetector): +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. """ - 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 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)``. + 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. """ - #: 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() + 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, + ) - #: 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) + return _mode_function - #: 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) +@autoinit +class BaseModeOverlapDetector(PhasorDetector, ABC): + """Abstract base for mode-overlap detectors. - #: 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) + 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`). - #: 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 when the detector plane's waveguide lies on a - #: symmetry plane of a reduced (half/quarter) domain so the reference mode is solved with the - #: same wall the FDTD uses (e.g. ``(0, 1)`` for PEC at y=0 and PMC at the z Si-mid plane). - #: Must match the corresponding ModePlaneSource for a consistent overlap. - symmetry: tuple[int, int] = frozen_field(default=(0, 0)) + 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( @@ -77,6 +223,7 @@ class ModeOverlapDetector(PhasorDetector): @property def propagation_axis(self) -> int: + """Physical axis normal to the detector plane, i.e. the singleton grid axis.""" if sum([a == 1 for a in self.grid_shape]) != 1: raise Exception(f"Invalid ModeOverlapDetector shape: {self.grid_shape}") return self.grid_shape.index(1) @@ -87,17 +234,12 @@ def place_on_grid( 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, ) - 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})" - ) grid = self._config.resolved_grid if grid is not None: weights = grid.face_area(axis=self.propagation_axis, slice_tuple=self.grid_slice_tuple) @@ -112,37 +254,61 @@ def _face_area_weights(self) -> jax.Array: """Return detector-plane face areas for mode-overlap integration.""" return self._cached_face_area_weights - def _transverse_edge_coordinates(self) -> tuple[jax.Array, jax.Array] | None: - """Return physical transverse edge coordinates for the mode solver. + def _plane_coordinates(self) -> tuple[jax.Array, jax.Array, jax.Array]: + """Return ``(X, Y, Z)`` cell-center coordinate meshgrids for the detector plane. - 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. + 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 - if grid is None: - return None - - transverse_edges = [] + axis_centers: list[jax.Array] = [] 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. + 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 - ``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. + @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`. """ - 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() + raise NotImplementedError def apply( self, @@ -155,6 +321,7 @@ def apply( 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: @@ -195,19 +362,10 @@ def apply( electric_conductivity=sigma_slice, conductivity_spacing=conductivity_spacing, ) - mode_E, mode_H, mode_neff = compute_mode( - frequency=wc.get_frequency(), - inv_permittivities=inv_eps_i, - 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(), + 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) @@ -274,7 +432,7 @@ def compute_overlap( 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 ModeOverlapDetector before calling compute_mode_overlap!") + raise Exception("Need to call apply on the mode-overlap detector before calling compute_overlap!") overlaps = [ self.compute_overlap_to_mode( state=state, @@ -285,3 +443,334 @@ def compute_overlap( for i in range(len(self.wave_characters)) ] return jnp.stack(overlaps, axis=0) + + +@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 when the detector plane's waveguide lies on a + #: symmetry plane of a reduced (half/quarter) domain so the reference mode is solved with the + #: same wall the FDTD uses (e.g. ``(0, 1)`` for PEC at y=0 and PMC at the z Si-mid plane). + #: Must match the corresponding ModePlaneSource for a consistent overlap. + 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.""" + 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(), + ) + + +@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 diff --git a/src/fdtdx/objects/sources/linear_polarization.py b/src/fdtdx/objects/sources/linear_polarization.py index 1ac6e7b0..1e159841 100644 --- a/src/fdtdx/objects/sources/linear_polarization.py +++ b/src/fdtdx/objects/sources/linear_polarization.py @@ -7,8 +7,7 @@ from fdtdx.core.grid import calculate_time_offset_yee from fdtdx.core.jax.pytrees import autoinit, frozen_field -from fdtdx.core.linalg import get_wave_vector_raw, rotate_vector -from fdtdx.core.misc import linear_interpolated_indexing, normalize_polarization_for_source +from fdtdx.core.misc import linear_interpolated_indexing, tilted_polarization_vectors from fdtdx.core.physics.metrics import compute_energy from fdtdx.dispersion import effective_inv_permittivity from fdtdx.objects.sources.tfsf import TFSFPlaneSource, _build_dispersive_H_filter, _source_impedance @@ -161,27 +160,19 @@ def apply( c4=c4_slice, ) - # determine E/H polarization - e_pol_raw, h_pol_raw = normalize_polarization_for_source( + center, azimuth, elevation = self._get_random_parts(key) + + # determine E/H polarization and the (tilted) wave vector — shared with the + # analytic Gaussian mode-overlap detector via tilted_polarization_vectors. + e_pol, h_pol, wave_vector = tilted_polarization_vectors( direction=self.direction, propagation_axis=self.propagation_axis, fixed_E_polarization_vector=self.fixed_E_polarization_vector, fixed_H_polarization_vector=self.fixed_H_polarization_vector, + azimuth_radians=azimuth, + elevation_radians=elevation, dtype=self._config.dtype, ) - wave_vector_raw = get_wave_vector_raw( - direction=self.direction, - propagation_axis=self.propagation_axis, - dtype=self._config.dtype, - ) - - center, azimuth, elevation = self._get_random_parts(key) - - # tilt polarizations - axes_tpl = (self.horizontal_axis, self.vertical_axis, self.propagation_axis) - wave_vector = rotate_vector(wave_vector_raw, azimuth, elevation, axes_tpl) - e_pol = rotate_vector(e_pol_raw, azimuth, elevation, axes_tpl) - h_pol = rotate_vector(h_pol_raw, azimuth, elevation, axes_tpl) # update is amplitude multiplied by polarization amplitude_raw = self._get_amplitude_raw(center)[None, ...] diff --git a/src/fdtdx/utils/sparams.py b/src/fdtdx/utils/sparams.py index 80fcf726..2225de73 100644 --- a/src/fdtdx/utils/sparams.py +++ b/src/fdtdx/utils/sparams.py @@ -17,7 +17,7 @@ from fdtdx.fdtd.wrapper import run_fdtd from fdtdx.materials import Material from fdtdx.objects.boundaries.initialization import BoundaryConfig, boundary_objects_from_config -from fdtdx.objects.detectors.mode import ModeOverlapDetector +from fdtdx.objects.detectors.mode import BaseModeOverlapDetector, ModeOverlapDetector from fdtdx.objects.sources.mode import ModePlaneSource from fdtdx.objects.static_material.polygon import ExtrudedPolygon from fdtdx.objects.static_material.static import SimulationVolume @@ -310,12 +310,12 @@ def calculate_sparam( input_det_state = final_arrays.detector_states[input_norm_name] input_det = objects[input_norm_name] - assert isinstance(input_det, ModeOverlapDetector) + assert isinstance(input_det, BaseModeOverlapDetector) input_overlap = input_det.compute_overlap(input_det_state) result: dict[tuple[str, str], jax.Array] = {} for obj in objects.object_list: - if isinstance(obj, ModeOverlapDetector): + if isinstance(obj, BaseModeOverlapDetector): state = final_arrays.detector_states[obj.name] raw_overlap = obj.compute_overlap(state) result[(obj.name, input_port_name)] = raw_overlap / input_overlap @@ -378,7 +378,7 @@ def determine_input_norm_detector_name(name_part: str, objects: ObjectContainer) exact_matches = [] results = [] for obj in objects.object_list: - if isinstance(obj, ModeOverlapDetector): + if isinstance(obj, BaseModeOverlapDetector): if obj.name == exact_name: exact_matches.append(obj.name) if name_part in obj.name: diff --git a/tests/simulation/physics/detectors/__init__.py b/tests/simulation/physics/detectors/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/simulation/physics/detectors/test_mode_overlap.py b/tests/simulation/physics/detectors/test_mode_overlap.py new file mode 100644 index 00000000..db848b33 --- /dev/null +++ b/tests/simulation/physics/detectors/test_mode_overlap.py @@ -0,0 +1,221 @@ +"""Physics simulation tests: GaussianModeOverlapDetector against a real Gaussian beam. + +A ``GaussianPlaneSource`` launches a beam into vacuum and a bank of +``GaussianModeOverlapDetector`` s on one downstream plane measure it. Every assertion is a +*ratio* between detectors sharing that plane, so the mode normalization and the overlap +integral's own scaling cancel and what is left is physics: the measured coupling is +compared against the closed-form overlap of the two Gaussians involved. + +The beam is deliberately wide (``w = 1.18 um``, Rayleigh range ``4.4 um``) so the detector +plane sits at ``0.05`` Rayleigh ranges — near-collimated, so a flat-phase reference is the +matched one — and so the beam's own angular spread stays well below the tilt under test. +Mismatched references are all *narrower* than the beam, which keeps them fully inside the +detector plane and the analytic formula free of truncation bias. + +``GaussianPlaneSource`` parameterizes its profile as ``exp(-r^2 / (2 (radius std)^2))`` with +a hard aperture at ``radius``, so the ``1/e`` amplitude radius the detector wants is +``radius * std * sqrt(2)`` — see ``_BEAM_RADIUS``. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import fdtdx +from fdtdx.constants import c as c0 + +_WAVELENGTH = 1e-6 +_RESOLUTION = 50e-9 +_PML_CELLS = 10 +_DOMAIN_XY = 8e-6 +_DOMAIN_Z = 4e-6 +_SIM_TIME = 120e-15 + +_SOURCE_Z = _PML_CELLS + 2 +_DET_Z = _SOURCE_Z + 4 # 200 nm downstream + +# Source aperture radius, and the 1/e amplitude radius of the beam it actually launches. +_SOURCE_RADIUS = 2.5e-6 +_BEAM_RADIUS = _SOURCE_RADIUS * (1 / 3) * np.sqrt(2.0) # std defaults to 1/3 + +_WAVENUMBER = 2.0 * np.pi / _WAVELENGTH +_RAYLEIGH_RANGE = np.pi * _BEAM_RADIUS**2 / _WAVELENGTH +_PROPAGATION_DISTANCE = (_DET_Z - _SOURCE_Z) * _RESOLUTION + +_NARROW_FACTORS = (0.6, 0.4) +_TILT_DEGREES = 30.0 +_DIVERGENCE_DEGREES = 30.0 + +_DT_APPROX = 0.99 * _RESOLUTION / (c0 * np.sqrt(3)) +_STEPS_PER_PERIOD = round(_WAVELENGTH / (c0 * _DT_APPROX)) + + +def _radius_overlap(w_field: float, w_ref: float) -> float: + """Analytic power coupling between two collimated circular Gaussians.""" + return (2.0 * w_field * w_ref / (w_field**2 + w_ref**2)) ** 2 + + +def _curvature_overlap(w: float, divergence_degrees: float) -> float: + """Analytic power coupling between a flat-phase Gaussian and a curved one of equal width. + + From ``eta = 4 / (w^4 [(2/w^2)^2 + c^2])`` with ``c = k tan(angle) / (2 w)`` the + quadratic phase coefficient; the envelopes are identical, so only the phase mismatch + survives. + """ + c = _WAVENUMBER * np.tan(np.deg2rad(divergence_degrees)) / (2.0 * w) + return 1.0 / (1.0 + c**2 * w**4 / 4.0) + + +def _detector(name: str, wave, **kwargs) -> fdtdx.GaussianModeOverlapDetector: + """A GaussianModeOverlapDetector on the shared measurement plane.""" + params = { + "name": name, + "partial_grid_shape": (None, None, 1), + "wave_characters": (wave,), + "mode_radius": _BEAM_RADIUS, + "direction": "+", + "fixed_E_polarization_vector": (1.0, 0.0, 0.0), + } + params.update(kwargs) + return fdtdx.GaussianModeOverlapDetector(**params) + + +def _build(): + """Assemble the free-space domain, the source and the detector bank.""" + config = fdtdx.SimulationConfig( + grid=fdtdx.UniformGrid(spacing=_RESOLUTION), + time=_SIM_TIME, + dtype=jnp.float32, + ) + objects, constraints = [], [] + + volume = fdtdx.SimulationVolume(partial_real_shape=(_DOMAIN_XY, _DOMAIN_XY, _DOMAIN_Z)) + objects.append(volume) + + bound_cfg = fdtdx.BoundaryConfig.from_uniform_bound(thickness=_PML_CELLS) + bound_dict, c_list = fdtdx.boundary_objects_from_config(bound_cfg, volume) + constraints.extend(c_list) + objects.extend(bound_dict.values()) + + wave = fdtdx.WaveCharacter(wavelength=_WAVELENGTH) + source = fdtdx.GaussianPlaneSource( + partial_grid_shape=(None, None, 1), + wave_character=wave, + direction="+", + fixed_E_polarization_vector=(1, 0, 0), + radius=_SOURCE_RADIUS, + ) + constraints.extend( + [ + source.same_size(volume, axes=(0, 1)), + source.place_at_center(volume, axes=(0, 1)), + source.set_grid_coordinates(axes=(2,), sides=("-",), coordinates=(_SOURCE_Z,)), + ] + ) + objects.append(source) + + detectors = [ + _detector("matched", wave), + _detector("backward", wave, direction="-"), + _detector("tilted", wave, azimuth_angle=_TILT_DEGREES), + _detector("diverging", wave, divergence_angle=_DIVERGENCE_DEGREES), + ] + detectors += [_detector(f"narrow_{factor}", wave, mode_radius=factor * _BEAM_RADIUS) for factor in _NARROW_FACTORS] + for det in detectors: + constraints.extend( + [ + det.same_size(volume, axes=(0, 1)), + det.place_at_center(volume, axes=(0, 1)), + det.set_grid_coordinates(axes=(2,), sides=("-",), coordinates=(_DET_Z,)), + ] + ) + objects.extend(detectors) + + return objects, constraints, config + + +@pytest.fixture(scope="module") +def overlap_powers(): + """Run the simulation once and return ``{detector name: |overlap|^2}``.""" + objects, constraints, config = _build() + key = jax.random.PRNGKey(0) + obj_container, arrays, params, config, _ = fdtdx.place_objects( + object_list=objects, + config=config, + constraints=constraints, + key=key, + ) + arrays, obj_container, _ = fdtdx.apply_params(arrays, obj_container, params, key) + _, arrays = fdtdx.run_fdtd(arrays=arrays, objects=obj_container, config=config, key=key) + + powers = {} + for det in obj_container.detectors: + if isinstance(det, fdtdx.BaseModeOverlapDetector): + overlap = det.compute_overlap(arrays.detector_states[det.name]) + powers[det.name] = float(jnp.abs(overlap[0]) ** 2) + print({name: f"{value:.4e}" for name, value in powers.items()}) + return powers + + +def test_beam_is_collimated_at_the_detector_plane(): + """Guard the premise the other tests rest on: the launched beam is near its waist.""" + assert _PROPAGATION_DISTANCE < 0.1 * _RAYLEIGH_RANGE + assert _STEPS_PER_PERIOD > 10 # enough temporal sampling for the phasor accumulation + + +def test_mismatched_radius_follows_analytic_overlap(overlap_powers): + """Coupling into a wrong-radius reference matches the two-Gaussian formula. + + This is what shows the overlap integral computes a coupling *efficiency* rather than + merely a large number wherever the field is: the predicted ratios follow from the beam + radii alone and are independent of every normalization in play. + """ + matched = overlap_powers["matched"] + for factor in _NARROW_FACTORS: + expected = _radius_overlap(_BEAM_RADIUS, factor * _BEAM_RADIUS) + measured = overlap_powers[f"narrow_{factor}"] / matched + assert measured == pytest.approx(expected, rel=0.03), ( + f"radius x{factor}: measured ratio {measured:.4f} vs analytic {expected:.4f}" + ) + + +def test_divergence_follows_analytic_overlap(overlap_powers): + """A curved reference loses exactly the predicted amount against a collimated beam. + + Fails if the curvature term is dropped (ratio -> 1) or carries the wrong magnitude. The + residual few percent is the launched beam's own ~0.7 deg of curvature at this plane, + which the analytic estimate treats as exactly flat. + """ + expected = _curvature_overlap(_BEAM_RADIUS, _DIVERGENCE_DEGREES) + measured = overlap_powers["diverging"] / overlap_powers["matched"] + assert measured == pytest.approx(expected, rel=0.08), ( + f"divergence {_DIVERGENCE_DEGREES} deg: measured ratio {measured:.4f} vs analytic {expected:.4f}" + ) + + +def test_tilt_suppresses_coupling(overlap_powers): + """A tilted reference barely couples to a normal-incidence beam. + + The transverse phase ramp is what suppresses it — a 1D estimate gives + ``exp(-(k sin(theta) w)^2 / 4) ~ 0.03`` here — so this fails if the ramp is dropped. + """ + ratio = overlap_powers["tilted"] / overlap_powers["matched"] + assert ratio < 0.05, f"tilted/matched = {ratio:.3e}; the phase ramp is not taking effect" + + +def test_backward_reference_reads_zero(overlap_powers): + """A forward beam carries no backward modal power — direction selectivity. + + Would be ~1 if the overlap merely measured field magnitude on the plane. + """ + ratio = overlap_powers["backward"] / overlap_powers["matched"] + assert ratio < 1e-3, f"backward/forward = {ratio:.3e} is not negligible" + + +def test_matched_reference_is_the_best(overlap_powers): + """No mismatched reference out-couples the matched one.""" + matched = overlap_powers["matched"] + others = [name for name in overlap_powers if name != "matched"] + for name in others: + assert overlap_powers[name] < matched, f"{name} out-coupled the matched reference" diff --git a/tests/unit/objects/detectors/test_mode_custom.py b/tests/unit/objects/detectors/test_mode_custom.py new file mode 100644 index 00000000..48ee90af --- /dev/null +++ b/tests/unit/objects/detectors/test_mode_custom.py @@ -0,0 +1,610 @@ +"""Tests for objects/detectors/mode.py - arbitrary / analytic reference modes. + +Covers the abstract :class:`BaseModeOverlapDetector` seam and the two new subclasses +that accept a user-provided mode instead of the waveguide mode solver: +:class:`CustomModeOverlapDetector` and :class:`GaussianModeOverlapDetector`, plus the +:func:`gaussian_mode_fields` / :func:`gaussian_mode_function` helpers. +""" + +import inspect + +import jax.numpy as jnp +import numpy as np +import pytest + +from fdtdx.core.physics.metrics import compute_poynting_flux +from fdtdx.core.wavelength import WaveCharacter +from fdtdx.dispersion import LorentzPole, compute_pole_coefficients, effective_inv_permittivity +from fdtdx.objects.detectors.mode import ( + BaseModeOverlapDetector, + CustomModeOverlapDetector, + GaussianModeOverlapDetector, + gaussian_mode_fields, + gaussian_mode_function, +) + + +@pytest.fixture +def single_frequency(): + """Single optical frequency wave character.""" + return [WaveCharacter(wavelength=1e-6)] + + +@pytest.fixture +def two_frequencies(): + """Two frequency wave characters.""" + return [WaveCharacter(wavelength=1e-6), WaveCharacter(wavelength=1.5e-6)] + + +def _constant_mode_function(*, coordinates, frequency, propagation_axis, inv_permittivity): + """A trivial mode_function: uniform E along axis 0, H along axis 1.""" + del frequency, propagation_axis, inv_permittivity + amp = jnp.ones(coordinates[0].shape, dtype=jnp.float32) + mode_E = jnp.zeros((3, *amp.shape), dtype=jnp.float32).at[0].set(amp) + mode_H = jnp.zeros((3, *amp.shape), dtype=jnp.float32).at[1].set(amp) + return mode_E, mode_H + + +def _weighted_flux(det, mode_E, mode_H): + """Integrated Poynting flux of a stored reference mode over the detector plane.""" + flux = compute_poynting_flux(mode_E.astype(jnp.complex64), mode_H.astype(jnp.complex64)) + s_axis = 0.5 * jnp.real(flux[det.propagation_axis]) + return float(jnp.abs(jnp.sum(s_axis * det._face_area_weights()))) + + +class TestBaseIsAbstract: + """BaseModeOverlapDetector must not be directly instantiable.""" + + def test_isabstract(self): + """The ABC is not concrete.""" + assert inspect.isabstract(BaseModeOverlapDetector) + + def test_instantiation_raises(self, single_frequency): + """Constructing the ABC directly fails.""" + with pytest.raises(TypeError): + BaseModeOverlapDetector(wave_characters=single_frequency) + + +class TestCustomModeOverlapDetector: + """CustomModeOverlapDetector evaluates a user-provided mode_function.""" + + def test_roundtrip_shapes(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """apply() stores per-frequency mode fields of the plane's shape.""" + det = CustomModeOverlapDetector( + wave_characters=single_frequency, + mode_function=_constant_mode_function, + normalize=False, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + + # one frequency, 3 components, plane shape (8, 8, 1) + assert det._mode_E.shape == (1, 3, 8, 8, 1) + assert det._mode_H.shape == (1, 3, 8, 8, 1) + + state = det.init_state() + state = det.update( + jnp.array(0), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32) * 0.5, + state, + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + 1.0, + ) + overlap = det.compute_overlap(state) + assert overlap.shape == (1,) + assert jnp.iscomplexobj(overlap) + + def test_callable_survives_placement(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """A Python callable stored as a frozen field survives pytree placement/apply.""" + det = CustomModeOverlapDetector( + wave_characters=single_frequency, + mode_function=lambda *, coordinates, frequency, propagation_axis, inv_permittivity: ( + jnp.zeros((3, *coordinates[0].shape), dtype=jnp.float32).at[0].set(1.0), + jnp.zeros((3, *coordinates[0].shape), dtype=jnp.float32).at[1].set(1.0), + ), + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + assert det._mode_E.shape == (1, 3, 8, 8, 1) + + def test_compute_overlap_requires_apply(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """compute_overlap() without apply() is an error, not a silent zero.""" + det = CustomModeOverlapDetector( + wave_characters=single_frequency, + mode_function=_constant_mode_function, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + state = det.init_state() + with pytest.raises(Exception, match="before calling compute_overlap"): + det.compute_overlap(state) + + def test_multi_frequency_stacks(self, simulation_config, plane_grid_slice, random_key, two_frequencies): + """Multiple wave characters stack along the leading axis.""" + det = CustomModeOverlapDetector( + wave_characters=two_frequencies, + mode_function=_constant_mode_function, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + assert det._mode_E.shape == (2, 3, 8, 8, 1) + + state = det.init_state() + state = det.update( + jnp.array(0), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32) * 0.5, + state, + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + 1.0, + ) + assert det.compute_overlap(state).shape == (2,) + + def test_normalize_gives_unit_flux(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """normalize=True renormalizes the mode to unit through-plane flux.""" + det = CustomModeOverlapDetector( + wave_characters=single_frequency, + mode_function=_constant_mode_function, + normalize=True, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + assert _weighted_flux(det, det._mode_E[0], det._mode_H[0]) == pytest.approx(1.0, rel=1e-4) + + def test_no_normalize_keeps_raw_amplitude(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """normalize=False leaves the supplied amplitude alone.""" + det = CustomModeOverlapDetector( + wave_characters=single_frequency, + mode_function=_constant_mode_function, + normalize=False, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + # raw E·H amplitude over the plane is far from unity (uniform mode, tiny cell areas) + assert _weighted_flux(det, det._mode_E[0], det._mode_H[0]) != pytest.approx(1.0, rel=1e-2) + + +class TestGaussianModeOverlapDetector: + """GaussianModeOverlapDetector builds an analytic Gaussian reference mode.""" + + def test_roundtrip(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """The Gaussian detector places, applies and overlaps end to end.""" + det = GaussianModeOverlapDetector( + wave_characters=single_frequency, + mode_radius=3e-7, + direction="+", + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + assert det._mode_E.shape == (1, 3, 8, 8, 1) + + state = det.init_state() + state = det.update( + jnp.array(0), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32) * 0.5, + state, + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + 1.0, + ) + overlap = det.compute_overlap(state) + assert overlap.shape == (1,) + assert jnp.all(jnp.isfinite(jnp.abs(overlap))) + + def test_only_expected_components_nonzero(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """Default polarization: E along the first transverse axis, H along the second.""" + det = GaussianModeOverlapDetector( + wave_characters=single_frequency, + mode_radius=3e-7, + direction="+", + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + mode_E, mode_H = det._mode_E[0], det._mode_H[0] + # propagation axis is 2 (singleton); transverse axes are (0, 1) + assert det.propagation_axis == 2 + assert jnp.any(mode_E[0] != 0) and not jnp.any(mode_E[1] != 0) and not jnp.any(mode_E[2] != 0) + assert jnp.any(mode_H[1] != 0) and not jnp.any(mode_H[0] != 0) and not jnp.any(mode_H[2] != 0) + + def test_matches_gaussian_mode_function(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """The detector class and the standalone factory produce the same mode.""" + det_cls = GaussianModeOverlapDetector( + wave_characters=single_frequency, + mode_radius=3e-7, + direction="+", + normalize=False, + ) + det_cls = det_cls.place_on_grid(plane_grid_slice, simulation_config, random_key) + det_cls = det_cls.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + + det_fn = CustomModeOverlapDetector( + wave_characters=single_frequency, + mode_function=gaussian_mode_function(radius=3e-7, direction="+"), + normalize=False, + ) + det_fn = det_fn.place_on_grid(plane_grid_slice, simulation_config, random_key) + det_fn = det_fn.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + + np.testing.assert_allclose(np.array(det_cls._mode_E), np.array(det_fn._mode_E), rtol=1e-6) + np.testing.assert_allclose(np.array(det_cls._mode_H), np.array(det_fn._mode_H), rtol=1e-6) + + +class TestDispersionHandling: + """Custom/Gaussian detectors pick up ε(ω_c) via the shared apply() correction.""" + + @staticmethod + def _dispersive_coeffs(det, pole): + """Single-pole Lorentz ADE coefficients for a test medium.""" + dt = det._config.time_step_duration + c1, c2, c3, c4 = compute_pole_coefficients((pole,), dt) + + def mk(a): # broadcast a per-pole scalar to (num_poles, 1, Nx, Ny, Nz) + """Build a placed detector for the given mode function.""" + return jnp.full((1, 1, 8, 8, 8), float(a[0]), dtype=jnp.float32) + + return mk(c1), mk(c2), mk(c3), mk(c4), dt + + def test_gaussian_and_custom_use_effective_permittivity(self, simulation_config, plane_grid_slice, random_key): + """Both analytic detectors see eps(omega_c), not eps_infinity.""" + wc = WaveCharacter(wavelength=1e-6) + omega = 2.0 * np.pi * wc.get_frequency() + eps_inf = 2.0 + inv_eps = jnp.full((3, 8, 8, 8), 1.0 / eps_inf, dtype=jnp.float32) + pole = LorentzPole(resonance_frequency=3.0 * wc.get_frequency(), damping=1e12, delta_epsilon=1.5) + + det = GaussianModeOverlapDetector(wave_characters=[wc], mode_radius=3e-7, direction="+") + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + c1, c2, c3, c4, dt = self._dispersive_coeffs(det, pole) + + # Independently compute the expected effective index on the detector plane slice, + # using the very same helper apply() calls internally. + sl = (slice(None), slice(0, 8), slice(0, 8), slice(0, 1)) + csl = (slice(None), slice(None), slice(0, 8), slice(0, 8), slice(0, 1)) + inv_eps_eff = effective_inv_permittivity(inv_eps[sl], c1[csl], c2[csl], c3[csl], omega, dt, c4=c4[csl]) + n_eff_expected = float(jnp.sqrt(jnp.mean(1.0 / inv_eps_eff))) + assert abs(n_eff_expected - eps_inf**0.5) > 1e-3 # the pole genuinely shifts the index + + # Gaussian: without dispersion args -> ε∞; with them -> ε(ω_c). + n_inf = det.apply(random_key, inv_eps, 1.0)._mode_neff[0] + n_disp = det.apply( + random_key, inv_eps, 1.0, dispersive_c1=c1, dispersive_c2=c2, dispersive_c3=c3, dispersive_c4=c4 + )._mode_neff[0] + assert float(jnp.real(n_inf)) == pytest.approx(eps_inf**0.5, rel=1e-5) + assert float(jnp.real(n_disp)) == pytest.approx(n_eff_expected, rel=1e-4) + + # Custom detector reports the same effective index (its callable saw ε(ω_c)). + cdet = CustomModeOverlapDetector(wave_characters=[wc], mode_function=_constant_mode_function) + cdet = cdet.place_on_grid(plane_grid_slice, simulation_config, random_key) + cdet = cdet.apply( + random_key, inv_eps, 1.0, dispersive_c1=c1, dispersive_c2=c2, dispersive_c3=c3, dispersive_c4=c4 + ) + assert float(jnp.real(cdet._mode_neff[0])) == pytest.approx(n_eff_expected, rel=1e-4) + + +class TestGaussianModeFields: + """gaussian_mode_fields orientation and Poynting-direction physics.""" + + @pytest.mark.parametrize("propagation_axis", [0, 1, 2]) + @pytest.mark.parametrize("direction", ["+", "-"]) + def test_orientation_and_poynting_sign(self, propagation_axis, direction): + # Build a small transverse plane (singleton on the propagation axis). + """E/H sit on the right axes and the flux follows `direction`.""" + line = jnp.linspace(-1e-7, 1e-7, 6) + singleton = jnp.array([0.0]) + axis_coords = [line, line, line] + axis_coords[propagation_axis] = singleton + coordinates = jnp.meshgrid(*axis_coords, indexing="ij") + + mode_E, mode_H = gaussian_mode_fields( + coordinates, + propagation_axis, + radius=5e-7, + direction=direction, + wavelength=1e-6, + refractive_index=2.0, + ) + + transverse = [a for a in range(3) if a != propagation_axis] + pol_axis, h_axis = transverse[0], transverse[1] + + # E only along the polarization axis, H only along the other transverse axis. + for comp in range(3): + if comp == pol_axis: + assert jnp.any(mode_E[comp] != 0) + else: + assert not jnp.any(mode_E[comp] != 0) + if comp == h_axis: + assert jnp.any(mode_H[comp] != 0) + else: + assert not jnp.any(mode_H[comp] != 0) + + # |H| = n |E| ratio for the analytic plane wave (n = 2.0 here). + np.testing.assert_allclose( + np.abs(np.array(mode_H[h_axis])), 2.0 * np.abs(np.array(mode_E[pol_axis])), rtol=1e-5 + ) + + # Poynting vector along the propagation axis has the sign of the direction. + flux = compute_poynting_flux(mode_E.astype(jnp.complex64), mode_H.astype(jnp.complex64)) + s_total = float(jnp.real(jnp.sum(flux[propagation_axis]))) + if direction == "+": + assert s_total > 0 + else: + assert s_total < 0 + + def test_polarization_axis_override(self): + # propagation axis 2 -> transverse (0, 1); force E along axis 1 instead of default 0. + """polarization_axis picks which transverse axis E points along.""" + line = jnp.linspace(-1e-7, 1e-7, 6) + coordinates = jnp.meshgrid(line, line, jnp.array([0.0]), indexing="ij") + mode_E, mode_H = gaussian_mode_fields( + coordinates, 2, radius=5e-7, direction="+", polarization_axis=1, wavelength=1e-6 + ) + assert jnp.any(mode_E[1] != 0) and not jnp.any(mode_E[0] != 0) + assert jnp.any(mode_H[0] != 0) and not jnp.any(mode_H[1] != 0) + + def test_invalid_polarization_axis_raises(self): + """A longitudinal polarization_axis is rejected.""" + line = jnp.linspace(-1e-7, 1e-7, 6) + coordinates = jnp.meshgrid(line, line, jnp.array([0.0]), indexing="ij") + with pytest.raises(ValueError, match="transverse"): + gaussian_mode_fields(coordinates, 2, radius=5e-7, direction="+", polarization_axis=2, wavelength=1e-6) + + +class TestGaussianModeGeneralization: + """Arbitrary polarization vector and off-normal propagation angle (like the source).""" + + @staticmethod + def _plane_coords(): + """A small square plane normal to axis 2.""" + line = jnp.linspace(-1e-7, 1e-7, 6) + return jnp.meshgrid(line, line, jnp.array([0.0]), indexing="ij") + + def test_fixed_E_polarization_vector(self): + """E follows the (normalized) fixed polarization vector; here diagonal in-plane.""" + mode_E, mode_H = gaussian_mode_fields( + self._plane_coords(), + 2, + radius=5e-7, + direction="+", + fixed_E_polarization_vector=(1.0, 1.0, 0.0), + wavelength=1e-6, + ) + # E equally along axes 0 and 1, nothing along the propagation axis. + np.testing.assert_allclose(np.array(mode_E[0]), np.array(mode_E[1]), rtol=1e-6) + assert not jnp.any(mode_E[2] != 0) + # H is orthogonal (cross(k, E)) -> also diagonal in plane, still |H| = n|E| (n=1). + assert jnp.any(mode_H[0] != 0) and jnp.any(mode_H[1] != 0) + + def test_polarization_axis_and_vector_conflict_raises(self): + """polarization_axis and an explicit vector are mutually exclusive.""" + with pytest.raises(ValueError, match="not both"): + gaussian_mode_fields( + self._plane_coords(), + 2, + radius=5e-7, + direction="+", + polarization_axis=0, + fixed_E_polarization_vector=(1.0, 0.0, 0.0), + wavelength=1e-6, + ) + + def test_normal_incidence_is_real(self): + """No tilt -> the reference mode stays real-valued (regression).""" + mode_E, _ = gaussian_mode_fields(self._plane_coords(), 2, radius=5e-7, direction="+", wavelength=1e-6) + assert not jnp.iscomplexobj(mode_E) + + def test_tilt_adds_phase_ramp_and_longitudinal_component(self): + """A tilted beam becomes complex, gains a propagation-axis E component, and ramps.""" + mode_E, _ = gaussian_mode_fields( + self._plane_coords(), + 2, + radius=5e-7, + direction="+", + azimuth_angle=25.0, + elevation_angle=15.0, + refractive_index=1.5, + wavelength=1e-6, + ) + assert jnp.iscomplexobj(mode_E) + # elevation tilt gives the E field a component along the (former) plane normal + assert jnp.any(jnp.abs(mode_E[2]) > 1e-6) + # transverse phase ramp: the phase is not constant across the plane + nonzero = mode_E[0][jnp.abs(mode_E[0]) > 0] + assert float(jnp.ptp(jnp.angle(nonzero))) > 0.1 + + @pytest.mark.parametrize("direction", ["+", "-"]) + def test_tilt_preserves_propagation_sign(self, direction): + """A tilted mode still carries flux in its nominal direction.""" + mode_E, mode_H = gaussian_mode_fields( + self._plane_coords(), + 2, + radius=5e-7, + direction=direction, + azimuth_angle=20.0, + refractive_index=1.0, + wavelength=1e-6, + ) + s_total = float(jnp.real(jnp.sum(compute_poynting_flux(mode_E, mode_H)[2]))) + assert (s_total > 0) if direction == "+" else (s_total < 0) + + def test_detector_with_tilt_and_polarization( + self, simulation_config, plane_grid_slice, random_key, single_frequency + ): + """GaussianModeOverlapDetector exposes the same polarization/angle knobs.""" + det = GaussianModeOverlapDetector( + wave_characters=single_frequency, + mode_radius=3e-7, + direction="+", + fixed_E_polarization_vector=(0.0, 1.0, 0.0), + azimuth_angle=15.0, + elevation_angle=5.0, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + assert jnp.iscomplexobj(det._mode_E) + + state = det.init_state() + state = det.update( + jnp.array(0), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + jnp.ones((3, 8, 8, 8), dtype=jnp.float32) * 0.5, + state, + jnp.ones((3, 8, 8, 8), dtype=jnp.float32), + 1.0, + ) + overlap = det.compute_overlap(state) + assert overlap.shape == (1,) + assert jnp.all(jnp.isfinite(jnp.abs(overlap))) + + +class TestGaussianRadiusValidation: + """A non-positive radius is rejected before it can produce NaNs.""" + + @staticmethod + def _plane_coords(): + """A small square plane normal to axis 2.""" + line = jnp.linspace(-1e-7, 1e-7, 6) + return jnp.meshgrid(line, line, jnp.array([0.0]), indexing="ij") + + @pytest.mark.parametrize("radius", [0.0, -1e-7]) + def test_non_positive_radius_raises(self, radius): + """A zero or negative radius is rejected.""" + with pytest.raises(ValueError, match="radius must be positive"): + gaussian_mode_fields(self._plane_coords(), 2, radius=radius, direction="+", wavelength=1e-6) + + @pytest.mark.parametrize("radius", [0.0, -1e-7]) + def test_detector_rejects_non_positive_mode_radius( + self, radius, simulation_config, plane_grid_slice, random_key, single_frequency + ): + """The detector validates mode_radius at placement.""" + det = GaussianModeOverlapDetector( + wave_characters=single_frequency, + mode_radius=radius, + direction="+", + ) + with pytest.raises(ValueError, match="mode_radius must be positive"): + det.place_on_grid(plane_grid_slice, simulation_config, random_key) + + +class TestTiltedEnvelopeProjection: + """Under tilt the envelope must use the beam's own cross-section, not the plane's. + + A collimated beam of circular cross-section cut by a tilted plane leaves an elliptical + footprint, stretched by ``1/cos(theta)`` along the tilt direction — the same projection + ``GaussianPlaneSource`` applies. Measuring that footprint is therefore the check that + source and detector describe the same beam. + """ + + RADIUS = 5e-7 + + def _amplitude_along(self, axis, azimuth): + """|mode_E| sampled along one transverse axis of a plane normal to axis 2.""" + line = jnp.linspace(-2e-6, 2e-6, 201) + zero = jnp.array([0.0]) + coords = [zero, zero, zero] + coords[axis] = line + mode_E, _ = gaussian_mode_fields( + jnp.meshgrid(*coords, indexing="ij"), + 2, + radius=self.RADIUS, + direction="+", + azimuth_angle=azimuth, + refractive_index=1.0, + wavelength=1e-6, + ) + return np.asarray(line), np.abs(np.asarray(mode_E)).max(axis=0).ravel() + + @staticmethod + def _fitted_radius(line, amplitude): + """Recover the 1/e amplitude radius from the sampled profile. + + Normalized by the peak: under tilt the polarization vector spreads over components, + so the largest component is below 1 at the beam center. + """ + positive = line >= 0 + normalized = amplitude[positive] / amplitude.max() + return float(np.interp(np.exp(-1.0), normalized[::-1], line[positive][::-1])) + + @pytest.mark.parametrize("azimuth", [20.0, 35.0]) + def test_footprint_stretches_by_one_over_cos(self, azimuth): + """Along the tilt direction the footprint widens by exactly 1/cos(theta).""" + line, amplitude = self._amplitude_along(0, azimuth) + expected = self.RADIUS / np.cos(np.deg2rad(azimuth)) + assert self._fitted_radius(line, amplitude) == pytest.approx(expected, rel=1e-3) + + @pytest.mark.parametrize("azimuth", [20.0, 35.0]) + def test_footprint_unchanged_across_the_tilt(self, azimuth): + """The axis orthogonal to the tilt keeps the untilted radius.""" + line, amplitude = self._amplitude_along(1, azimuth) + assert self._fitted_radius(line, amplitude) == pytest.approx(self.RADIUS, rel=1e-3) + + def test_untilted_footprint_is_circular(self): + """Regression: with no tilt the projection is the identity.""" + for axis in (0, 1): + line, amplitude = self._amplitude_along(axis, 0.0) + assert self._fitted_radius(line, amplitude) == pytest.approx(self.RADIUS, rel=1e-3) + + +class TestDivergenceAngle: + """Wavefront curvature at the detector plane.""" + + RADIUS = 1e-6 + WAVELENGTH = 1e-6 + WAVENUMBER = 2 * np.pi / WAVELENGTH + LINE = np.linspace(-1e-6, 1e-6, 41) + + def _fields(self, divergence_angle): + """Gaussian fields on a fixed square plane at the given divergence angle.""" + line = jnp.asarray(self.LINE) + coords = jnp.meshgrid(line, line, jnp.array([0.0]), indexing="ij") + return gaussian_mode_fields( + coords, + 2, + radius=self.RADIUS, + direction="+", + divergence_angle=divergence_angle, + refractive_index=1.0, + wavelength=self.WAVELENGTH, + ) + + def test_zero_angle_is_the_collimated_mode(self): + """Regression: the default reproduces the real, flat-phase profile.""" + flat, _ = self._fields(0.0) + assert not jnp.iscomplexobj(flat) + + def test_curvature_is_quadratic_with_the_right_sign(self): + """Phase follows k r^2 / 2R, increasing outward (diverging) for a positive angle.""" + angle = 10.0 + mode_E, _ = self._fields(angle) + assert jnp.iscomplexobj(mode_E) + + centre = len(self.LINE) // 2 + phase = np.unwrap(np.angle(np.asarray(mode_E[0])[:, centre, 0])) + inv_curvature_radius = np.tan(np.deg2rad(angle)) / self.RADIUS + expected = 0.5 * self.WAVENUMBER * inv_curvature_radius * self.LINE**2 + np.testing.assert_allclose(phase - phase[centre], expected, atol=1e-6) + + def test_converging_is_the_conjugate_of_diverging(self): + """A negative angle mirrors the phase, leaving the envelope untouched.""" + diverging, _ = self._fields(12.0) + converging, _ = self._fields(-12.0) + np.testing.assert_allclose(np.asarray(converging), np.conj(np.asarray(diverging)), rtol=1e-6) + + def test_envelope_is_independent_of_curvature(self): + """Curvature is pure phase — |E| must match the collimated profile.""" + flat, _ = self._fields(0.0) + curved, _ = self._fields(15.0) + np.testing.assert_allclose(np.abs(np.asarray(curved)), np.asarray(flat), rtol=1e-6) + + def test_detector_exposes_divergence_angle(self, simulation_config, plane_grid_slice, random_key, single_frequency): + """A diverging reference mode reaches the stored fields through the detector.""" + det = GaussianModeOverlapDetector( + wave_characters=single_frequency, + mode_radius=1e-7, + direction="+", + divergence_angle=8.0, + ) + det = det.place_on_grid(plane_grid_slice, simulation_config, random_key) + det = det.apply(random_key, jnp.ones((3, 8, 8, 8), dtype=jnp.float32), 1.0) + assert jnp.iscomplexobj(det._mode_E) + assert jnp.all(jnp.isfinite(jnp.abs(det._mode_E)))