diff --git a/.claude/skills/fdtdx/SKILL.md b/.claude/skills/fdtdx/SKILL.md index 254e21dd..0597430d 100644 --- a/.claude/skills/fdtdx/SKILL.md +++ b/.claude/skills/fdtdx/SKILL.md @@ -210,7 +210,7 @@ bound_dict, constraint_list = fdtdx.boundary_objects_from_config(bound_cfg, volu ## Simulation Symmetry -Mirror-symmetry exploitation: build the **full** model, set `config.symmetry`, and fdtdx runs the reduced half/quarter/octant internally (up to 8× less memory/compute), then you unfold results back to the full domain. Implemented in `src/fdtdx/fdtd/symmetry.py`; the FDTD time loop is untouched (the symmetry plane is just an added PEC/PMC wall object). +Mirror-symmetry exploitation: build the **full** model, set `config.symmetry`, and fdtdx runs the reduced half/quarter/octant internally (up to 8× less memory/compute), then you unfold results back to the full domain. Implemented in `src/fdtdx/fdtd/symmetry.py`; the FDTD time loop is untouched (an electric plane is just an added PEC wall object, a magnetic one nothing at all — see below). **Encoding** — `symmetry: tuple[int, int, int]` on `SimulationConfig`, order `(x, y, z)`: - `0` = no symmetry on this axis @@ -219,11 +219,16 @@ Mirror-symmetry exploitation: build the **full** model, set `config.symmetry`, a Distinct from manually placing PEC/PMC via `BoundaryConfig` (that still works unchanged); `config.symmetry` is the additive auto-reduce path. +**Where each plane sits** (this is the key asymmetry between the two wall types, and it sets both the wall handling and the unfold map): +- An **electric** plane sits *on* the reduced domain's min edge — the tangential `E` samples live there and the odd symmetry makes them vanish, so it is a PEC face and gets a `PerfectElectricConductor` wall object. +- A **magnetic** plane sits *half a cell below* the min edge. Sources and materials are rasterized per cell, so the discrete problem is mirror symmetric about the tangential-`H` node one cell out, where tangential `H` vanishes — already supplied by the zero halo of the field padding. It gets **no wall object**: a `PerfectMagneticConductor` there would zero tangential `H` one cell *inside* the domain, imposing the condition half a cell off the plane (a clean first-order-wrong answer, ~4e-02 field error at 50 nm). +- Detectors touching an **electric** plane need the mirror in their co-location halo (`pad_fields_with_symmetry_mirror` in `fdtd/update.py`), or the plane row records exactly *half* the field. Magnetic planes want no halo there. + **Requirements / behavior:** - Each symmetric axis **must resolve to an even cell count** (else `place_objects` raises `ValueError`) — guarantees an exact split and cell-for-cell unfold. -- The **upper half is kept** so the plane lands at the reduced domain's min edge (matching the mode solver's "wall at min edge" convention). Objects are clipped to that half during `place_objects`; centered objects keep their upper half, objects entirely in the discarded half are dropped (with a warning). -- The min-side boundary on each symmetric axis is replaced by the PEC/PMC wall; the far (max) side keeps whatever the user set (use PML there, not periodic). -- `ModePlaneSource` / `ModeOverlapDetector` get their mode-solver `symmetry` 2-tuple **auto-derived** (PMC→1, PEC/none→0 on the two transverse axes) unless explicitly set. +- The **upper half is kept** so the plane lands at the reduced domain's min edge. Objects are clipped to that half during `place_objects`; centered objects keep their upper half, objects entirely in the discarded half are dropped (with a warning). +- The min-side boundary on each symmetric axis is dropped; the far (max) side keeps whatever the user set (use PML there, not periodic — the halo *behind* the symmetry plane is set by the mirror and never by wrapping to the far side, and `place_objects` warns if a periodic/Bloch boundary survives on a symmetric axis). +- `ModePlaneSource` / `ModeOverlapDetector` solve the mode on the **mirrored full cross-section** and restrict it to the kept half (`compute_mode_symmetry_reduced`), rather than using the mode solver's own symmetric solve — the solver samples materials on its staggered grid while FDTDX writes one cell-centred ε array per component, so a symmetric solve on the reduced cross-section shifts `neff` at first order in Δ. Their mode-solver `symmetry` 2-tuple is **not** auto-derived and is ignored (with a warning) under `config.symmetry`. - The user must place objects symmetrically about the center plane — asymmetric models are warned about but not corrected (true of every FDTD symmetry feature). **Usage:** @@ -245,13 +250,14 @@ E_full = fdtdx.unfold_fields(arrays.fields.E, config.symmetry, "E") # (3, Nx, N **Unfold helpers** (`fdtdx.unfold_fields`, `fdtdx.unfold_detector_states`, `fdtdx.unfold_source_mode`, `fdtdx.unfold_array`): - `unfold_fields(field, symmetry, field_type)` — reconstruct a full `(3, Nx, Ny, Nz)` E/H array via per-component parity mirror. The general escape hatch — derive any quantity from the full fields. +- **Mirror index map** (`mirror_pairs_on_plane` in `core/physics/symmetry.py`, the single source of truth): across an **electric** plane, components sampled *on* it (tangential `E`, normal `H`) pair as `m±j` — the plane row is its own mirror — while half-cell-offset components mirror one-to-one; across a **magnetic** plane *every* component mirrors one-to-one (plain flip), because the plane is half a cell out. Applying one convention to both axes is wrong on one of them (~8e-02 vs ~1e-02 interior error). - `unfold_detector_states(arrays, objects, config)` — pure post-processing that rebuilds each detector's full-domain output from its stored reduced output + parity (no in-loop cost, no flags). Spatial outputs are mirrored per component; `reduce_volume` sums/means are rescaled per component (even doubles/keeps, **odd vanishes**); `as_slices` energy planes are mirrored in-plane. - `unfold_source_mode(source, config)` → `(E_full, H_full)` — reconstruct the full-domain mode profile a `ModePlaneSource` *injects* (its solved-on-the-reduced-cross-section `_E`/`_H`). Unfolds only the transverse axes (the propagation axis is never a symmetry plane). Run `apply_params` first. For the fields *recorded during the run*, prefer a detector on the source plane + `unfold_detector_states`. - **Guardrails:** unfolding a non-symmetric model (`symmetry=(0,0,0)`) raises `ValueError`; `place_objects` warns that results are on the reduced domain until unfolded. - **Not unfoldable:** `DiffractiveDetector` raises `NotImplementedError` (its diffraction-order basis depends on domain size — unfold the fields and recompute instead). - **Mode-overlap S-params** are already correct on the reduced domain (source + detector share the reduced plane), so they need no unfolding. -**Mode sources are fully wired:** under symmetry, a `ModePlaneSource`'s cross-section is clipped to the reduced grid, its mode-solver `symmetry` 2-tuple is auto-derived from `config.symmetry`, and `compute_mode` solves/injects the half/quarter mode with the matching PEC/PMC wall. Use `unfold_source_mode` to inspect the reconstructed full profile. +**Mode sources are fully wired:** under symmetry, a `ModePlaneSource`'s cross-section is clipped to the reduced grid, and `compute_mode_symmetry_reduced` mirrors that cross-section back to the full one, solves there, projects onto the walls' parity subspace and restricts — reproducing the full-domain mode (`neff` to ~1e-7 at every resolution) instead of the solver's own symmetric solve. A wall type the selected mode cannot support raises. Amplitudes follow the "unit power through the plane it occupies" convention, so the reduced profile is `√(2^k)` larger than the restriction of the full-domain mode — to within a few percent at coarse resolution, because the discrete mode's flux does not split exactly evenly between the halves (first order in Δ; see the docstring for measured numbers). Use `unfold_source_mode` to inspect the reconstructed full profile. **Gradient note:** the differentiable simulation runs on the reduced domain (correct and cheaper); unfolding is a post-hoc step on the output arrays. @@ -430,3 +436,4 @@ assert jnp.all(jnp.isfinite(grads)) - **Stacking objects with mixed dispersion**: `UniformMaterialObject` always writes a full zero-padded pole-coefficient stack into its `grid_slice`, so placing a non-dispersive object over a dispersive one cleanly overwrites stale coefficients. Rely on this rather than assuming "no dispersion = leave coefficients alone". - **Symmetry results look wrong / are half-size**: with `config.symmetry` set, `run_fdtd` returns *reduced-domain* arrays — you must call `fdtdx.unfold_detector_states` / `fdtdx.unfold_fields` to get full-domain results (see Simulation Symmetry). `place_objects` warns about this. Unfolding a non-symmetric model raises. - **Symmetry needs even cells + symmetric model**: each symmetric axis must resolve to an even cell count (`place_objects` raises otherwise), and the user's full model must actually be mirror-symmetric about the center plane — asymmetric objects are only warned about. Use PML (not periodic) on the far side of a symmetric axis. +- **Symmetry wall type follows the polarization, and the two types are not mirror images of each other**: PEC where `E` is normal to the plane, PMC where it is tangential (a wrong choice is warned about, not corrected — the reduced run then faithfully simulates a field with the wrong parity). Only the electric type is a wall object; the magnetic one is a plane half a cell outside the domain carried by the zero halo. Anything that mirrors across a plane (unfold maps, mode parity projection) must branch on the wall type, never on the Yee offsets alone. diff --git a/src/fdtdx/config.py b/src/fdtdx/config.py index 9c97f796..593e3d0a 100644 --- a/src/fdtdx/config.py +++ b/src/fdtdx/config.py @@ -96,9 +96,13 @@ class SimulationConfig(TreeClass): #: ``+1`` = PMC (magnetic-wall) mirror on the axis center plane. #: When any entry is nonzero, :func:`fdtdx.place_objects` automatically reduces the #: domain to the symmetric half/quarter/octant (keeping the upper half along each - #: symmetric axis), clips every object onto that reduced grid, inserts the PEC/PMC - #: wall on the symmetry plane, and forwards the matching per-axis condition to the - #: mode solver. The FDTD then runs on the reduced domain; call + #: symmetric axis) and clips every object onto that reduced grid. An electric plane + #: lands on the reduced domain's min edge and gets a PEC wall there; a magnetic plane + #: sits half a cell below it (sources and materials are rasterized per cell), where the + #: zero field halo already is the exact mirror, so it gets no wall object. Mode sources + #: and mode-overlap detectors solve on the mirrored full cross-section and restrict, + #: rather than using the mode solver's own symmetric solve. The FDTD then runs on the + #: reduced domain; call #: :func:`fdtdx.unfold_fields` / :func:`fdtdx.unfold_detector_states` afterwards to #: reconstruct the full-domain arrays. This is additive and independent of manually #: specifying PEC/PMC as ordinary boundaries via :class:`fdtdx.BoundaryConfig`. diff --git a/src/fdtdx/core/physics/modes.py b/src/fdtdx/core/physics/modes.py index 9448aa8a..2983bc66 100644 --- a/src/fdtdx/core/physics/modes.py +++ b/src/fdtdx/core/physics/modes.py @@ -7,11 +7,19 @@ import numpy as np import tidy3d from jax.typing import ArrayLike +from loguru import logger from tidy3d.components.mode.solver import compute_modes as _compute_modes from fdtdx.core.axis import get_transverse_axes +from fdtdx.core.jax.utils import is_jax_tracer from fdtdx.core.misc import expand_to_3x3 from fdtdx.core.physics.metrics import normalize_by_poynting_flux +from fdtdx.core.physics.symmetry import ( + mirror_edge_coordinates, + mirror_material_cross_section, + project_onto_parity, + restrict_to_kept_half, +) ModeTupleType = namedtuple("ModeTupleType", ["neff", "Ex", "Ey", "Ez", "Hx", "Hy", "Hz"]) """A named tuple containing the mode fields and effective index. @@ -377,6 +385,188 @@ def mode_helper(permittivity, permeability, c0_um, c1_um): return mode_E_norm, mode_H_norm, eff_idx +def _check_parity_residual(residual: jax.Array, walls: dict[int, int], object_name: str) -> None: + """Diagnose a parity projection that removed too much of the mode. + + Skipped when ``residual`` is a JAX tracer. A mode source or mode-overlap detector that overlaps a + :class:`~fdtdx.Device` solves its mode inside :func:`fdtdx.apply_params`, which callers routinely + trace (``jax.jit`` around an optimization step), and there the residual has no value yet: + concretizing it would raise, and comparing it would fire on a tracer. The check cannot be hoisted + out of the trace either — the residual depends on the device permittivity being traced over. It is + setup guidance, not something the solve depends on, so an eager ``apply_params`` surfaces it; and + ``place_objects`` already applies (hence checks) every mode object that does not overlap a device, + which is the overwhelming majority. + + Args: + residual (jax.Array): Fraction of the mode the projection removed. + walls (dict[int, int]): Mirror axis to wall type, for the message. + object_name (str): Name used in diagnostics. + + Raises: + ValueError: If the projection removed (almost) the whole mode, i.e. the configured wall types + are incompatible with the selected mode. + """ + if is_jax_tracer(residual): + return + value = float(residual) + wall_description = ", ".join(f"{'xyz'[axis]}={'PMC' if wall == 1 else 'PEC'}" for axis, wall in walls.items()) + if value > 0.9: + raise ValueError( + f"The mode selected for '{object_name}' has (almost) none of the symmetry imposed by the " + f"walls on {wall_description}: projecting it onto the admissible parity removes " + f"{value:.1%} of the mode. The wall types do not match this mode - flip the sign of " + f"config.symmetry on those axes, pick a different mode_index/filter_pol, or run without " + f"config.symmetry." + ) + if value > 0.3: + logger.warning( + f"The mode of '{object_name}' is only approximately symmetric about the walls on " + f"{wall_description}: the parity projection removed {value:.1%} of it. Check that the " + f"structure is mirror-symmetric there and that the wall types match the mode. A coarsely " + f"resolved cross-section alone can account for a residual of a few tens of percent - the " + f"mode solver samples materials on its staggered grid, so its discrete mode is only " + f"mirror-symmetric to first order in the cell size." + ) + + +def compute_mode_symmetry_reduced( + *, + mirrored_axes: tuple[int, ...], + walls: dict[int, int], + frequency: float, + inv_permittivities: jax.Array, + inv_permeabilities: jax.Array | float, + resolution: float | None = None, + direction: Literal["+", "-"] = "+", + mode_index: int = 0, + filter_pol: Literal["te", "tm"] | None = None, + dtype: jnp.dtype = jnp.float32, + bend_radius: float | None = None, + bend_axis: int | None = None, + transverse_coords: Sequence[jax.Array] | None = None, + object_name: str = "mode object", +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Solve a mode on a symmetry-reduced cross-section by way of the full cross-section. + + A reduced simulation replaces the discarded half by the mirror image of the kept half, so the + mode it supports is the full-domain mode restricted to the kept half. Rather than asking the + mode solver for a *symmetric* solve on the reduced cross-section, this mirrors the reduced + material arrays back to the full cross-section, solves there, and restricts the result. The + detour matters: the mode solver interprets the permittivity arrays on its own staggered Yee + grid, while FDTDX rasterizes materials per cell and hands the same array to every component, so + a solve on the reduced cross-section is inconsistent with the full one at first order in the + cell size (an ``neff`` error of several percent for a high-index waveguide). Going through the + full cross-section reproduces exactly the mode the unreduced simulation would inject. + + The solved mode is then projected onto the parity subspace the walls admit (see + :func:`~fdtdx.core.physics.symmetry.project_onto_parity`), which both removes the small + non-symmetric residue of the discrete mode and detects a wall type that does not match the mode + at all. Finally the fields are renormalized to unit Poynting flux through the *reduced* plane, + keeping the convention that a mode source launches unit power through the plane it occupies. + + Both steps are limited by the same discretization the detour above avoids for ``neff``: because + the solver samples materials on its staggered grid, the discrete mode of a mirror-symmetric + cross-section is itself only symmetric to first order in the cell size, so its flux does not + split exactly evenly between the two halves. Measured for a 400x200 nm Si waveguide at 1.55 um: + 0.446 / 0.554 along the solver's first transverse axis at 25 nm (0.473 / 0.527 at 12.5 nm), and + 0.499 / 0.501 along its second. Renormalizing over the reduced plane therefore leaves the + returned profile up to ~6% (25 nm) resp. ~3% (12.5 nm) above ``sqrt(2**k)`` times the restriction + of the full-domain mode on such an axis, and the parity projection removes a few tenths of a + percent of its norm. Both vanish with refinement; the flux convention is exact by construction at + every resolution. + + Args: + mirrored_axes (tuple[int, ...]): Physical axes clipped by a symmetry plane. + walls (dict[int, int]): Mirror axis to wall type (``-1`` PEC, ``+1`` PMC). + frequency (float): Operating frequency in Hz. + inv_permittivities (jax.Array): Reduced inverse permittivity on the mode plane. + inv_permeabilities (jax.Array | float): Reduced inverse permeability on the mode plane. + resolution (float | None): Uniform grid spacing, required without ``transverse_coords``. + direction (Literal["+", "-"]): Propagation direction. + mode_index (int): Index into the sorted mode list. + filter_pol (Literal["te", "tm"] | None): Optional polarization filter. + dtype (jnp.dtype): Float dtype of the simulation. + bend_radius (float | None): Waveguide bend radius, with ``bend_axis``. + bend_axis (int | None): Physical axis pointing toward the center of curvature. + transverse_coords (Sequence[jax.Array] | None): Reduced transverse edge coordinates, or + None on a uniform grid. + object_name (str): Name used in diagnostics. + + Returns: + tuple[jax.Array, jax.Array, jax.Array]: ``(E, H, effective_index)`` on the reduced + cross-section. + + Raises: + ValueError: If a waveguide bend shares an axis with a symmetry plane, or if the parity + projection removes almost the entire mode, which means the configured wall types are + incompatible with the selected mode. The latter is only detectable where the residual is + concrete, i.e. not inside ``jax.jit`` (see :func:`_check_parity_residual`). + """ + if bend_radius is not None and bend_axis is not None and bend_axis in mirrored_axes: + raise ValueError( + f"'{object_name}' bends about the {'xyz'[bend_axis]}-axis and config.symmetry mirrors " + f"that same axis. The bend is modelled by a conformal transformation that scales the " + f"refractive index linearly across bend_axis, so the transformed cross-section is not " + f"mirror-symmetric about a plane normal to it: the mode has no definite parity there and " + f"the reduced simulation, which replaces the discarded half by the mirror of the kept " + f"half, cannot represent it. Drop config.symmetry on the {'xyz'[bend_axis]}-axis, or put " + f"the symmetry plane normal to the other transverse axis - a bend leaves that one " + f"mirror-symmetric." + ) + propagation_axis = next(a for a in range(3) if inv_permittivities.shape[1:][a] == 1) + full_inv_permittivities = mirror_material_cross_section(inv_permittivities, mirrored_axes) + if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0: + full_inv_permeabilities: jax.Array | float = mirror_material_cross_section(inv_permeabilities, mirrored_axes) + else: + full_inv_permeabilities = inv_permeabilities + + full_transverse_coords = transverse_coords + if transverse_coords is not None: + transverse_axes = get_transverse_axes(propagation_axis) + full_transverse_coords = [ + mirror_edge_coordinates(coords) if axis in mirrored_axes else coords + for axis, coords in zip(transverse_axes, transverse_coords, strict=True) + ] + + mode_E, mode_H, eff_index = compute_mode( + frequency=frequency, + inv_permittivities=full_inv_permittivities, + inv_permeabilities=full_inv_permeabilities, + resolution=resolution, + direction=direction, + mode_index=mode_index, + filter_pol=filter_pol, + dtype=dtype, + bend_radius=bend_radius, + bend_axis=bend_axis, + symmetry=(0, 0), + transverse_coords=full_transverse_coords, + ) + + mode_E, residual_E = project_onto_parity(mode_E, "E", walls) + mode_H, residual_H = project_onto_parity(mode_H, "H", walls) + _check_parity_residual(jnp.maximum(residual_E, residual_H), walls, object_name) + + mode_E = restrict_to_kept_half(mode_E, mirrored_axes) + mode_H = restrict_to_kept_half(mode_H, mirrored_axes) + + area_weights = None + if transverse_coords is not None: + widths = [jnp.diff(jnp.asarray(coords)) for coords in transverse_coords] + area_2d = widths[0][:, None] * widths[1][None, :] + weight_shape = [1, 1, 1] + for local_axis, axis in enumerate(get_transverse_axes(propagation_axis)): + weight_shape[axis] = area_2d.shape[local_axis] + area_weights = area_2d.reshape(weight_shape).astype(mode_E.real.dtype) + mode_E, mode_H = normalize_by_poynting_flux( + mode_E, + mode_H, + axis=propagation_axis, + area_weights=area_weights, + ) + return mode_E, mode_H, eff_index + + def _is_reciprocal_tensor(components: Sequence[ArrayLike], tol: float = 1e-6) -> bool: """Check whether a material tensor given as 9 row-major components is symmetric (reciprocal).""" return bool( diff --git a/src/fdtdx/core/physics/symmetry.py b/src/fdtdx/core/physics/symmetry.py new file mode 100644 index 00000000..27bcca71 --- /dev/null +++ b/src/fdtdx/core/physics/symmetry.py @@ -0,0 +1,306 @@ +"""Low-level mirror-symmetry primitives shared by sources, detectors and the unfold helpers. + +This module holds only array/parity mathematics and imports nothing from ``fdtdx.objects`` or +``fdtdx.fdtd``, so object classes can use it without an import cycle. The user-facing symmetry +machinery lives in :mod:`fdtdx.fdtd.symmetry`, which re-exports :func:`field_component_parity`. + +Convention (identical to :mod:`fdtdx.fdtd.symmetry`): a symmetry plane normal to ``axis`` sits at +the *min edge* of the reduced domain, the upper half is kept, ``-1`` denotes a PEC (electric) wall +and ``+1`` a PMC (magnetic) wall. +""" + +from typing import Literal + +import jax +import jax.numpy as jnp + + +def field_component_parity( + field_type: Literal["E", "H"], + component: int, + axis: int, + wall: int, +) -> int: + """Mirror parity of a field component across a symmetry plane. + + For a mirror plane normal to ``axis`` with the kept domain on the ``+axis`` side, this is the + sign that relates the field on the discarded side to its mirror image on the kept side. A PEC + wall forces tangential ``E`` (and normal ``H``) to zero on the plane; a PMC wall forces + tangential ``H`` (and normal ``E``) to zero. + + Args: + field_type (Literal["E", "H"]): Which field the component belongs to. + component (int): Component axis index (0=x, 1=y, 2=z). + axis (int): Axis the mirror plane is normal to (0=x, 1=y, 2=z). + wall (int): ``-1`` for a PEC (electric) wall, ``+1`` for a PMC (magnetic) wall. + + Returns: + int: ``+1`` (even / unchanged under reflection) or ``-1`` (odd / sign-flipped). + """ + normal = component == axis + if wall == -1: # PEC: tangential E = 0 + if field_type == "E": + return 1 if normal else -1 + return -1 if normal else 1 + elif wall == 1: # PMC: tangential H = 0 + if field_type == "E": + return -1 if normal else 1 + return 1 if normal else -1 + raise ValueError(f"wall must be -1 (PEC) or +1 (PMC), got {wall}") + + +def component_sits_on_plane(field_type: Literal["E", "H"], component: int, axis: int) -> bool: + """Whether a Yee component's samples lie *on* a mirror plane normal to ``axis``. + + ``E_c`` is offset by half a cell along ``c`` only, so along ``axis`` it sits at integer + positions (i.e. on the plane at index 0) for every component except ``c == axis``; ``H_c`` is + offset along the two axes other than ``c``, so the opposite holds. The distinction sets the + mirror index map: on-plane samples pair ``i`` with ``-i`` (index 0 is its own mirror), the + half-cell-offset ones pair ``i`` with ``-i-1`` (a plain flip). + + Args: + field_type (Literal["E", "H"]): Which field the component belongs to. + component (int): Component axis index (0=x, 1=y, 2=z). + axis (int): Axis the mirror plane is normal to. + + Returns: + bool: True if the component's samples include the plane itself. + """ + if field_type == "E": + return component != axis + if field_type == "H": + return component == axis + raise ValueError(f"field_type must be 'E' or 'H', got {field_type!r}") + + +def mirror_pairs_on_plane( + field_type: Literal["E", "H"], + component: int, + axis: int, + wall: int, +) -> bool: + """Whether a component's mirror map pairs its samples *about* a shared plane row. + + This is the single place that decides the mirror index map, and it depends on the wall type, not + only on the Yee offsets: + + * An **electric** plane sits on the tangential-``E`` node row at the reduced domain's min edge, + so a component sampled there (see :func:`component_sits_on_plane`) is its own mirror and the + pairs are ``m ± j``. + * A **magnetic** plane sits half a cell *below* the reduced domain. Sources and materials are + rasterized per cell, so their footprint is symmetric about the cell boundary below the min + edge — the tangential-``H`` node one cell out — and across that plane *every* component + mirrors one-to-one (the plain flip), whatever its Yee offset. + + The parity itself (:func:`field_component_parity`) is unaffected: it is a property of the wall + type, not of where the plane sits. + + Args: + field_type (Literal["E", "H"]): Which field the component belongs to. + component (int): Component axis index (0=x, 1=y, 2=z). + axis (int): Axis the mirror plane is normal to. + wall (int): ``-1`` for a PEC (electric) wall, ``+1`` for a PMC (magnetic) wall. + + Returns: + bool: True if the samples pair as ``m ± j`` about a shared row, False for a plain flip. + """ + return wall == -1 and component_sits_on_plane(field_type, component, axis) + + +def mirror_material_array( + array: jax.Array, + axis: int, + component_axis: int = 0, +) -> jax.Array: + """Reflect a material array about a mirror plane at its min edge, along a spatial axis. + + Material tensors are even under reflection except for the off-diagonal entries with exactly one + index equal to the mirror axis, which change sign (``eps_xy -> -eps_xy`` for a mirror normal to + x). Isotropic (1-component) and diagonal (3-component) layouts therefore only need the flip. + + Args: + array (jax.Array): Material array whose three trailing axes are spatial and whose + ``component_axis`` holds 1, 3 or 9 tensor components (row-major 3x3 for 9). + axis (int): Physical axis (0=x, 1=y, 2=z) the mirror plane is normal to. + component_axis (int): Axis holding the tensor components. Defaults to 0. + + Returns: + jax.Array: The reflected copy, same shape as the input. + """ + spatial_axis = array.ndim - 3 + axis + mirrored = jnp.flip(array, axis=spatial_axis) + num_components = array.shape[component_axis] + if num_components != 9: + return mirrored + signs = [ + -1.0 if ((row == axis) != (col == axis)) else 1.0 # exactly one index along the mirror axis + for row in range(3) + for col in range(3) + ] + sign_shape = [1] * array.ndim + sign_shape[component_axis] = 9 + return mirrored * jnp.asarray(signs, dtype=array.dtype).reshape(sign_shape) + + +def mirror_material_cross_section( + array: jax.Array, + axes: tuple[int, ...], + component_axis: int = 0, +) -> jax.Array: + """Rebuild the full-domain material array from a symmetry-reduced one. + + For every axis in ``axes`` the reflected copy is prepended, doubling that axis: exactly the + material distribution the reduced simulation models, since its mirror wall replaces the + discarded half by the mirror image of the kept half. + + Args: + array (jax.Array): Reduced material array (three trailing spatial axes). + axes (tuple[int, ...]): Physical axes to mirror. + component_axis (int): Axis holding the tensor components. Defaults to 0. + + Returns: + jax.Array: Full-domain material array, doubled along each mirrored axis. + """ + for axis in axes: + spatial_axis = array.ndim - 3 + axis + array = jnp.concatenate([mirror_material_array(array, axis, component_axis), array], axis=spatial_axis) + return array + + +def mirror_edge_coordinates(edges: jax.Array) -> jax.Array: + """Rebuild full-domain edge coordinates from the kept half's, mirrored about ``edges[0]``. + + Args: + edges (jax.Array): Monotone edge coordinates of the kept half, ``n + 1`` entries. + + Returns: + jax.Array: ``2n + 1`` edge coordinates, mirror-symmetric about the original first edge. + """ + widths = edges[1:] - edges[:-1] + full_widths = jnp.concatenate([jnp.flip(widths), widths]) + return jnp.concatenate([jnp.zeros((1,), dtype=edges.dtype), jnp.cumsum(full_widths)]) + (edges[0] - edges[-1]) + + +def _slice_axis(array: jax.Array, axis: int, start: int, stop: int | None = None) -> jax.Array: + """``array[..., start:stop, ...]`` along a single axis.""" + index: list[slice] = [slice(None)] * array.ndim + index[axis] = slice(start, stop) + return array[tuple(index)] + + +def mirror_about_interior_plane(array: jax.Array, axis: int, on_plane: bool) -> jax.Array: + """Mirror image of an array whose mirror plane lies at the centre of ``axis``. + + The array holds ``2n`` samples split by a plane between index ``n - 1`` and ``n``. Samples that + sit half a cell off the plane pair up as ``n - 1 - j <-> n + j`` (a plain flip); samples that sit + *on* the plane pair up as ``n - j <-> n + j``, leaving index ``n`` as its own mirror and index + ``0`` without a partner (its mirror would be one cell past the array). Index ``0`` is returned + unchanged, which is harmless because it belongs to the discarded half. + + Args: + array (jax.Array): Array to mirror along ``axis``. + axis (int): Array axis to mirror. + on_plane (bool): Whether the samples include the plane itself + (see :func:`component_sits_on_plane`). + + Returns: + jax.Array: The mirrored array, same shape as the input. + """ + flipped = jnp.flip(array, axis=axis) + if not on_plane: + return flipped + shifted = jnp.roll(flipped, shift=1, axis=axis) + # roll wraps the far edge into index 0, which has no mirror partner: restore the original value. + return jnp.concatenate([_slice_axis(array, axis, 0, 1), _slice_axis(shifted, axis, 1)], axis=axis) + + +def mirror_extend_low_side(array: jax.Array, axis: int, parity: int, on_plane: bool) -> jax.Array: + """Low-side (discarded-half) block for an array whose mirror plane is at its min edge. + + Given the ``n`` kept samples along ``axis``, returns the ``n`` samples of the mirrored half in + ascending order, so that ``concatenate([low, array])`` is the full-domain array. Samples half a + cell off the plane mirror one-to-one (a plain flip). Samples *on* the plane have index 0 as + their own mirror, so only indices ``1..n-1`` produce mirror images; the outermost missing sample + (whose partner lies past the kept half) is filled by repeating its neighbour, which is the + far-boundary cell of the reconstructed domain. + + Args: + array (jax.Array): Kept half, mirrored along ``axis``. + axis (int): Array axis to mirror. + parity (int): ``+1`` (even) or ``-1`` (odd). + on_plane (bool): Whether the samples include the plane itself. + + Returns: + jax.Array: The mirrored low-side block, same shape as ``array``. + """ + if not on_plane: + return parity * jnp.flip(array, axis=axis) + mirrored = parity * jnp.flip(_slice_axis(array, axis, 1), axis=axis) + return jnp.concatenate([_slice_axis(mirrored, axis, 0, 1), mirrored], axis=axis) + + +def project_onto_parity( + field: jax.Array, + field_type: Literal["E", "H"], + walls: dict[int, int], +) -> tuple[jax.Array, jax.Array]: + """Project a full-cross-section ``(3, Nx, Ny, Nz)`` field onto the parity subspace of the walls. + + Keeps only the part of the field the walls admit: the even part of every component a wall makes + even and the odd part of every component it makes odd. That part is exactly representable on the + reduced domain — in particular the components a wall drives to zero on the plane come out zero + there — so injecting the result does not fight the wall. The returned residual is the fraction of + the field that was removed: a few tens of percent for a coarsely resolved mode (the mode solver's + discrete mode is only mirror-symmetric to first order in the cell size), close to one when the + wall type does not match the field at all. + + Each component is mirrored with the index map of its wall type (see + :func:`mirror_pairs_on_plane` and :func:`mirror_about_interior_plane`), matching how the reduced + FDTD run treats it. + + Args: + field (jax.Array): Field array with a leading component axis and three spatial axes, spanning + the full (mirrored) cross-section. + field_type (Literal["E", "H"]): Whether ``field`` is electric or magnetic. + walls (dict[int, int]): Mapping of mirror axis to wall type (``-1`` PEC, ``+1`` PMC). + + Returns: + tuple[jax.Array, jax.Array]: The projected field and the removed fraction (relative L2 norm), + as a JAX scalar so this stays usable inside ``jax.jit`` — a mode source or mode-overlap + detector overlapping a :class:`~fdtdx.Device` solves its mode inside + :func:`fdtdx.apply_params`, which callers do trace. Inspect the value only where it is + concrete (see :func:`fdtdx.core.jax.utils.is_jax_tracer`). + """ + projected = field + for axis, wall in walls.items(): + components = [] + for component in range(3): + parity = field_component_parity(field_type, component, axis, wall) + on_plane = mirror_pairs_on_plane(field_type, component, axis, wall) + single = projected[component : component + 1] + mirrored = mirror_about_interior_plane(single, axis + 1, on_plane) + components.append(0.5 * (single + parity * mirrored)) + projected = jnp.concatenate(components, axis=0) + norm = jnp.linalg.norm(jnp.abs(field)) + residual = jnp.where(norm > 0, jnp.linalg.norm(jnp.abs(field - projected)) / norm, 0.0) + return projected, residual + + +def restrict_to_kept_half(array: jax.Array, axes: tuple[int, ...]) -> jax.Array: + """Slice the upper (kept) half off a full-domain array along each mirrored axis. + + Inverse of :func:`mirror_material_cross_section` for arrays whose three trailing axes are + spatial. + + Args: + array (jax.Array): Full-domain array (three trailing spatial axes). + axes (tuple[int, ...]): Physical axes that were mirrored. + + Returns: + jax.Array: The kept half, halved along each axis in ``axes``. + """ + index: list[slice] = [slice(None)] * array.ndim + for axis in axes: + spatial_axis = array.ndim - 3 + axis + index[spatial_axis] = slice(array.shape[spatial_axis] // 2, None) + return array[tuple(index)] diff --git a/src/fdtdx/fdtd/initialization.py b/src/fdtdx/fdtd/initialization.py index 041d7344..6e90ee21 100644 --- a/src/fdtdx/fdtd/initialization.py +++ b/src/fdtdx/fdtd/initialization.py @@ -162,8 +162,9 @@ def place_objects( # non-symmetric path is unchanged. dropped_names: set[str] = set() reduced_volume_shape = None + unreduced_slices: dict[str, Any] = {} if config.has_symmetry: - resolved_slices, dropped_names, reduced_volume_shape = reduce_resolved_slices( + resolved_slices, unreduced_slices, dropped_names, reduced_volume_shape = reduce_resolved_slices( resolved_slices=resolved_slices, object_map=object_map, config=config, @@ -193,7 +194,10 @@ def place_objects( if grid.shape != volume_shape: raise ValueError(f"Configured grid shape {grid.shape} does not match simulation volume shape {volume_shape}.") - # Step 6: Place objects on grid based on resolved slice tuples + # Step 6: Place objects on grid based on resolved slice tuples. Under symmetry each object also + # remembers its unclipped extent (in reduced coordinates) so object *contents* that depend on + # the full extent - a Gaussian beam's centre, a mode cross-section - can be derived correctly + # even though the geometry was clipped onto the kept half. placed_objects = [] for name, slice_tuple in resolved_slices.items(): if name == volume_obj.name or name in dropped_names: @@ -201,25 +205,26 @@ def place_objects( obj = object_map[name] assert key is not None key, subkey = jax.random.split(key) - placed_objects.append( - obj.place_on_grid( - grid_slice_tuple=slice_tuple, - config=config, - key=subkey, - ) + placed = obj.place_on_grid( + grid_slice_tuple=slice_tuple, + config=config, + key=subkey, ) + if name in unreduced_slices: + placed = placed.aset("_unreduced_grid_slice_tuple", unreduced_slices[name]) + placed_objects.append(placed) # Step 7: Place volume first (index 0) assert key is not None key, subkey = jax.random.split(key) - placed_objects.insert( - 0, - volume_obj.place_on_grid( - grid_slice_tuple=resolved_slices[volume_obj.name], - config=config, - key=subkey, - ), + placed_volume = volume_obj.place_on_grid( + grid_slice_tuple=resolved_slices[volume_obj.name], + config=config, + key=subkey, ) + if volume_obj.name in unreduced_slices: + placed_volume = placed_volume.aset("_unreduced_grid_slice_tuple", unreduced_slices[volume_obj.name]) + placed_objects.insert(0, placed_volume) # Step 8: Insert the PEC/PMC symmetry walls and forward the per-axis condition to mode # sources/detectors, then warn that the simulation now runs on the reduced domain. @@ -238,7 +243,9 @@ def place_objects( wall_names = [w.name for w in walls] logger.warning( f"Symmetry {config.symmetry} reduces the simulation to grid shape {reduced_volume_shape} " - f"(walls added: {wall_names}; objects dropped: {sorted(dropped_names) or 'none'}). " + f"(PEC walls added on the electric planes: {wall_names or 'none'}; magnetic planes sit " + f"half a cell out and carry their mirror in the field halo, so they get no wall object; " + f"objects dropped: {sorted(dropped_names) or 'none'}). " f"Results are on the reduced domain — call fdtdx.unfold_detector_states / " f"fdtdx.unfold_fields to reconstruct the full domain." ) diff --git a/src/fdtdx/fdtd/symmetry.py b/src/fdtdx/fdtd/symmetry.py index f7101586..b9f25f07 100644 --- a/src/fdtdx/fdtd/symmetry.py +++ b/src/fdtdx/fdtd/symmetry.py @@ -6,15 +6,31 @@ * **Reduction** (used internally by ``place_objects``): clip every resolved grid slice onto the kept upper half along each symmetric axis, drop objects that fall entirely in the discarded - half, insert a PEC/PMC wall on the symmetry plane, and forward the matching per-axis condition - to mode sources/detectors. + half, and insert a PEC wall on each electric symmetry plane. * **Unfolding** (public helpers): reconstruct full-domain field and detector arrays from the reduced simulation, applying the correct per-component mirror parity. Symmetry encoding (per axis, order ``(x, y, z)``): ``0`` = none, ``-1`` = PEC (electric wall), ``+1`` = PMC (magnetic wall). The symmetry plane is the geometric center of the axis; the upper -half is kept so the plane lands at the reduced domain's min edge — matching the mode solver's -"wall at the min edge" convention. +half is kept so the plane lands at the reduced domain's min edge. + +**Where the two wall types sit.** An electric plane sits *on* the reduced domain's min edge: the +tangential ``E`` components are sampled there and the odd symmetry makes them vanish, which is a +PEC face. A magnetic plane sits *half a cell below* it: sources and materials are rasterized per +cell, so the discrete problem is mirror symmetric about the tangential-``H`` node one cell out, +where tangential ``H`` vanishes — already supplied by the zero halo of the field padding, so no wall +object is created for it (see :func:`make_symmetry_walls`). The same distinction sets the mirror +index map used when unfolding (:func:`~fdtdx.core.physics.symmetry.mirror_pairs_on_plane`): across +an electric plane, components sampled on it pair as ``m ± j``; across a magnetic plane every +component mirrors one-to-one. + +Mode sources and mode-overlap detectors clipped by a symmetry plane do not use the mode solver's own +symmetric solve; they mirror the reduced cross-section back to the full one, solve there and +restrict (see :func:`~fdtdx.core.physics.modes.compute_mode_symmetry_reduced`). + +The parity table itself (:func:`field_component_parity`) lives in +:mod:`fdtdx.core.physics.symmetry` so that source and detector classes can use it without an +import cycle; it is re-exported here, where it is documented. """ from typing import Literal @@ -26,11 +42,15 @@ from fdtdx.config import SimulationConfig from fdtdx.core.misc import validate_symmetric_axis_cells from fdtdx.core.null import Null +from fdtdx.core.physics.symmetry import ( + field_component_parity, + mirror_extend_low_side, + mirror_pairs_on_plane, +) from fdtdx.fdtd.container import ArrayContainer, ObjectContainer from fdtdx.objects.boundaries.bloch import BlochBoundary from fdtdx.objects.boundaries.boundary import BaseBoundary from fdtdx.objects.boundaries.pec import PerfectElectricConductor -from fdtdx.objects.boundaries.pmc import PerfectMagneticConductor from fdtdx.objects.detectors.diffractive import DiffractiveDetector from fdtdx.objects.detectors.energy import EnergyDetector from fdtdx.objects.detectors.field import FieldDetector @@ -56,45 +76,6 @@ _AXIS_NAMES = ("x", "y", "z") -# ────────────────────────────────────────────────────────────────────────────── -# Parity table -# ────────────────────────────────────────────────────────────────────────────── - - -def field_component_parity( - field_type: Literal["E", "H"], - component: int, - axis: int, - wall: int, -) -> int: - """Mirror parity of a field component across a symmetry plane. - - For a mirror plane normal to ``axis`` with the kept domain on the ``+axis`` side, this is the - sign that relates the field on the discarded side to its mirror image on the kept side. A PEC - wall forces tangential ``E`` (and normal ``H``) to zero on the plane; a PMC wall forces - tangential ``H`` (and normal ``E``) to zero. - - Args: - field_type (Literal["E", "H"]): Which field the component belongs to. - component (int): Component axis index (0=x, 1=y, 2=z). - axis (int): Axis the mirror plane is normal to (0=x, 1=y, 2=z). - wall (int): ``-1`` for a PEC (electric) wall, ``+1`` for a PMC (magnetic) wall. - - Returns: - int: ``+1`` (even / unchanged under reflection) or ``-1`` (odd / sign-flipped). - """ - normal = component == axis - if wall == -1: # PEC: tangential E = 0 - if field_type == "E": - return 1 if normal else -1 - return -1 if normal else 1 - elif wall == 1: # PMC: tangential H = 0 - if field_type == "E": - return -1 if normal else 1 - return 1 if normal else -1 - raise ValueError(f"wall must be -1 (PEC) or +1 (PMC), got {wall}") - - # ────────────────────────────────────────────────────────────────────────────── # Reduction (used by place_objects) # ────────────────────────────────────────────────────────────────────────────── @@ -105,7 +86,7 @@ def reduce_resolved_slices( object_map: dict[str, SimulationObject], config: SimulationConfig, volume_name: str, -) -> tuple[dict[str, SliceTuple3D], set[str], tuple[int, int, int]]: +) -> tuple[dict[str, SliceTuple3D], dict[str, SliceTuple3D], set[str], tuple[int, int, int]]: """Clip full-domain grid slices onto the symmetric (kept upper) half along each axis. Args: @@ -115,9 +96,12 @@ def reduce_resolved_slices( volume_name (str): Name of the simulation volume object. Returns: - tuple: ``(new_slices, dropped_names, reduced_volume_shape)`` where ``new_slices`` excludes - the dropped objects, ``dropped_names`` lists objects entirely in the discarded half, and - ``reduced_volume_shape`` is the reduced volume grid shape ``(Nx', Ny', Nz')``. + tuple: ``(new_slices, unreduced_slices, dropped_names, reduced_volume_shape)`` where + ``new_slices`` excludes the dropped objects, ``unreduced_slices`` holds each surviving + object's *unclipped* slice shifted into the reduced coordinate frame (negative start on + axes where the object crossed the symmetry plane), ``dropped_names`` lists objects entirely + in the discarded half, and ``reduced_volume_shape`` is the reduced volume grid shape + ``(Nx', Ny', Nz')``. """ symmetry = config.symmetry vol_slice = resolved_slices[volume_name] @@ -144,14 +128,18 @@ def reduce_resolved_slices( ) new_slices: dict[str, SliceTuple3D] = {} + unreduced_slices: dict[str, SliceTuple3D] = {} dropped: set[str] = set() for name, sl in resolved_slices.items(): if name == volume_name: new_slices[name] = (new_vol[0], new_vol[1], new_vol[2]) + shifted = [(sl[a][0] - mid_abs[a], sl[a][1] - mid_abs[a]) if symmetry[a] != 0 else sl[a] for a in range(3)] + unreduced_slices[name] = (shifted[0], shifted[1], shifted[2]) continue obj = object_map[name] clipped: list[tuple[int, int]] = [] + unclipped: list[tuple[int, int]] = [] drop = False for a in range(3): s0, s1 = sl[a] @@ -169,8 +157,10 @@ def reduce_resolved_slices( f"the mirror of the kept half." ) clipped.append((ns0, ns1)) + unclipped.append((s0 - m, s1 - m)) else: clipped.append((s0, s1)) + unclipped.append((s0, s1)) if drop: dropped.add(name) @@ -189,8 +179,20 @@ def reduce_resolved_slices( continue new_slices[name] = (clipped[0], clipped[1], clipped[2]) + unreduced_slices[name] = (unclipped[0], unclipped[1], unclipped[2]) + + if isinstance(obj, BlochBoundary) and symmetry[obj.axis] != 0: + # The min-side boundary of a symmetric axis is dropped above; a far-side one survives. + # Its wrap padding must not reach the min-side halo, which belongs to the mirror (see + # pad_fields_for_boundaries) - say so rather than let the axis look fully periodic. + logger.warning( + f"Periodic/Bloch boundary '{name}' sits on the {_AXIS_NAMES[obj.axis]}-symmetry " + f"axis. Periodic boundaries and mirror symmetry are mutually exclusive on the same " + f"axis: the halo behind the symmetry plane is set by the mirror, not by wrapping to " + f"this boundary. Use PML (or nothing) on the far side of a symmetric axis." + ) - return new_slices, dropped, reduced_volume_shape + return new_slices, unreduced_slices, dropped, reduced_volume_shape def make_symmetry_walls( @@ -199,7 +201,20 @@ def make_symmetry_walls( key: jax.Array, existing_names: set[str], ) -> list[SimulationObject]: - """Create the PEC/PMC wall objects sitting on each symmetry plane (reduced min edge). + """Create the PEC wall objects sitting on each *electric* symmetry plane (reduced min edge). + + Only an electric (``-1``) plane gets a wall object. It sits on the tangential-``E`` node row at + the reduced domain's min edge, where the odd symmetry does make those components vanish, so the + condition is exactly a PEC face. + + A magnetic (``+1``) plane gets **no wall object**. Its plane is half a cell below the reduced + domain: sources and materials are rasterized per cell, so the discrete problem is mirror + symmetric about the tangential-``H`` node one cell below the min edge, where tangential ``H`` + vanishes — which is precisely what the zero halo of the field padding already supplies. Placing a + :class:`~fdtdx.PerfectMagneticConductor` there instead (zeroing tangential ``H`` at the first + cell, half a cell inside the domain) or filling that halo with ``-H_t[0]`` forces the mirror onto + the cell edge and displaces the plane by half a cell: a consistent discretization of the *wrong* + problem, which converges cleanly at first order and so looks plausible while being wrong. Args: config (SimulationConfig): Simulation config; ``config.symmetry`` selects the axes/types. @@ -208,12 +223,13 @@ def make_symmetry_walls( existing_names (set[str]): Names already in use, so wall names can be made unique. Returns: - list[SimulationObject]: PEC/PMC boundary objects, already placed on the reduced grid. + list[SimulationObject]: PEC boundary objects for the electric symmetry planes, already + placed on the reduced grid. Empty if every symmetric axis is magnetic. """ walls: list[SimulationObject] = [] used = set(existing_names) for a in range(3): - if config.symmetry[a] == 0: + if config.symmetry[a] != -1: continue name = f"_sym_wall_{_AXIS_NAMES[a]}" suffix = 0 @@ -227,8 +243,7 @@ def make_symmetry_walls( 1 if a == 1 else None, 1 if a == 2 else None, ) - wall_cls = PerfectElectricConductor if config.symmetry[a] == -1 else PerfectMagneticConductor - wall = wall_cls( + wall = PerfectElectricConductor( name=name, axis=a, partial_grid_shape=partial_grid_shape, @@ -244,67 +259,48 @@ def make_symmetry_walls( config=config, key=subkey, ) + wall = wall.aset("_is_symmetry_wall", True) walls.append(wall) return walls -def _derive_mode_symmetry( - propagation_axis: int, - reduced_slice: SliceTuple3D, - symmetry: tuple[int, int, int], -) -> tuple[int, int]: - """Derive the mode-solver symmetry 2-tuple for a plane perpendicular to ``propagation_axis``. - - The mode solver expects the two non-propagation axes in increasing index order, with ``0`` = - PEC and ``1`` = PMC at the min edge. A wall is imposed only when the plane sits on the - symmetry plane (its reduced slice starts at index 0 on that transverse axis). - """ - transverse = [a for a in range(3) if a != propagation_axis] - result = [] - for a in transverse: - if symmetry[a] == 1 and reduced_slice[a][0] == 0: - result.append(1) # PMC - else: - result.append(0) # PEC / electric wall (also the no-symmetry default) - return (result[0], result[1]) - - def apply_mode_symmetry( placed_objects: list[SimulationObject], config: SimulationConfig, ) -> list[SimulationObject]: - """Forward the simulation symmetry to mode sources/detectors left at their default symmetry. + """Warn about explicit mode-solver symmetry tuples that ``config.symmetry`` supersedes. - For each :class:`ModePlaneSource` / :class:`ModeOverlapDetector` whose ``symmetry`` is still - the default ``(0, 0)``, set it to the value derived from the simulation symmetry and the - object's reduced placement. Explicit user-set values are respected (a mismatch warns). + Mode objects clipped by a symmetry plane no longer ask the mode solver for a symmetric solve: + :func:`~fdtdx.core.physics.modes.compute_mode_symmetry_reduced` mirrors the reduced + cross-section back to the full one, solves there and restricts the result, which reproduces the + mode the unreduced simulation would use (the solver's own symmetric solve does not, because + FDTDX hands it one cell-centred permittivity array per component while the solver samples + materials on its staggered grid). A user-set ``symmetry`` tuple is therefore inert here; say so + rather than silently ignoring it. Args: placed_objects (list[SimulationObject]): Objects already placed on the reduced grid. config (SimulationConfig): Simulation config. Returns: - list[SimulationObject]: The objects, with mode symmetry tuples filled in where applicable. + list[SimulationObject]: The objects, unchanged. """ - result: list[SimulationObject] = [] + del config for obj in placed_objects: - if isinstance(obj, (ModePlaneSource, ModeOverlapDetector)): - shape = obj.grid_shape - if 1 not in shape: - result.append(obj) - continue - propagation_axis = shape.index(1) - derived = _derive_mode_symmetry(propagation_axis, obj._grid_slice_tuple, config.symmetry) - if obj.symmetry == (0, 0): - obj = obj.aset("symmetry", derived) - elif obj.symmetry != derived: - logger.warning( - f"Mode object '{obj.name}' has an explicit symmetry={obj.symmetry} that differs " - f"from the value {derived} implied by config.symmetry={config.symmetry}. Keeping " - f"the explicit value; ensure this is intended." - ) - result.append(obj) - return result + if not isinstance(obj, (ModePlaneSource, ModeOverlapDetector)) or obj.symmetry == (0, 0): + continue + shape = obj.grid_shape + if 1 not in shape: + continue + propagation_axis = shape.index(1) + if obj.symmetry_mirror_axes(exclude_axis=propagation_axis): + logger.warning( + f"Mode object '{obj.name}' sets symmetry={obj.symmetry} for the mode solver, but " + f"config.symmetry already reduces its cross-section. The explicit value is ignored: the " + f"mode is solved on the mirrored full cross-section and restricted to the kept half. " + f"Remove it, or drop config.symmetry on those axes and keep your own PEC/PMC boundary." + ) + return placed_objects # ────────────────────────────────────────────────────────────────────────────── @@ -331,6 +327,16 @@ def unfold_fields( concatenates the mirror image in front of the kept half, so the symmetry plane ends up at the center of the reconstructed array. + The index map is chosen per axis from its wall type (see + :func:`~fdtdx.core.physics.symmetry.mirror_pairs_on_plane`). Across an **electric** plane, which + sits on the reduced domain's min edge, a component sampled *on* it (tangential ``E``, normal + ``H``) has its first sample as its own mirror, so only the remaining ones produce images, while + the half-cell-offset components mirror one-to-one. Across a **magnetic** plane, which sits half a + cell below the min edge, every component mirrors one-to-one. In the on-plane case one + reconstructed cell per axis — the outermost of the mirrored half, whose partner lies outside the + kept domain — is filled by repeating its neighbour; it is the far-boundary cell of the + reconstructed domain (inside the absorbing layer in a typical setup). + Args: field (jax.Array): Reduced field, shape ``(3, Nx, Ny, Nz)`` (component axis first). symmetry (tuple[int, int, int]): Per-axis symmetry ``(x, y, z)``; see module docstring. @@ -346,12 +352,17 @@ def unfold_fields( for a in range(3): if symmetry[a] == 0: continue - signs = jnp.asarray( - [field_component_parity(field_type, c, a, symmetry[a]) for c in range(3)], - dtype=arr.dtype, - ).reshape((3, 1, 1, 1)) - mirror = jnp.flip(arr, axis=a + 1) * signs - arr = jnp.concatenate([mirror, arr], axis=a + 1) + components = [] + for c in range(3): + single = arr[c : c + 1] + low = mirror_extend_low_side( + single, + axis=a + 1, + parity=field_component_parity(field_type, c, a, symmetry[a]), + on_plane=mirror_pairs_on_plane(field_type, c, a, symmetry[a]), + ) + components.append(jnp.concatenate([low, single], axis=a + 1)) + arr = jnp.concatenate(components, axis=0) return arr @@ -425,11 +436,12 @@ def unfold_array( symmetry: tuple[int, int, int], spatial_axes: tuple[int, int, int], signs: dict[int, jax.Array] | None = None, + on_plane_axes: tuple[int, ...] = (), ) -> jax.Array: """Mirror-and-concatenate a spatial array along each symmetric axis. Generic building block used by :func:`unfold_detector_states`. For every axis ``a`` with - ``symmetry[a] != 0`` the array is flipped along its corresponding array axis + ``symmetry[a] != 0`` the array is mirrored along its corresponding array axis ``spatial_axes[a]`` (optionally multiplied by a broadcastable per-component sign), and the mirror image is concatenated in front of the original. @@ -439,6 +451,11 @@ def unfold_array( spatial_axes (tuple[int, int, int]): For each physical axis, the corresponding array axis. signs (dict[int, jax.Array] | None): Optional mapping physical-axis → broadcastable sign array applied to the mirror image (defaults to ``+1``). + on_plane_axes (tuple[int, ...]): Physical axes whose samples lie *on* the symmetry plane + rather than half a cell off it. There the first sample is its own mirror, so it is not + duplicated and the outermost reconstructed sample repeats its neighbour. Only an electric + plane can be on-plane at all — a magnetic one sits half a cell below the reduced domain, + so everything mirrors one-to-one across it; see :func:`_colocated_on_plane_axes`. Returns: jax.Array: The unfolded array. @@ -449,11 +466,36 @@ def unfold_array( continue ax = spatial_axes[a] sign = 1.0 if signs is None else signs.get(a, 1.0) - mirror = jnp.flip(arr, axis=ax) * sign + if a in on_plane_axes: + mirror = mirror_extend_low_side(arr, axis=ax, parity=1, on_plane=True) * sign + else: + mirror = jnp.flip(arr, axis=ax) * sign arr = jnp.concatenate([mirror, arr], axis=ax) return arr +def _colocated_on_plane_axes(detector, touched: tuple[int, int, int]) -> tuple[int, ...]: + """Which physical axes a detector's stored samples sit on, for the mirror index map. + + With ``exact_interpolation`` the six components are co-located at the ``E_z`` Yee point + ``(i, j, k+1/2)``: integer along x and y and half a cell off along z. An integer position lands on + an **electric** symmetry plane, which sits on the reduced domain's min edge, so those samples pair + as ``m ± j``. A **magnetic** plane sits half a cell below the min edge, so nothing is on-plane + there and every sample mirrors one-to-one (see + :func:`~fdtdx.core.physics.symmetry.mirror_pairs_on_plane`). Without ``exact_interpolation`` the + raw staggered components are stored, each with its own offset, and no single index map fits; those + fall back to the plain flip. + + Args: + detector: The detector whose stored output is being unfolded. + touched (tuple[int, int, int]): Per-axis wall type of the symmetry planes this detector was + clipped by (``0`` where it was not). + """ + if not getattr(detector, "exact_interpolation", False): + return () + return tuple(a for a in (0, 1) if touched[a] == -1) + + def _stored_component_spec(components) -> list[tuple[Literal["E", "H"], int]]: """Stored ``(field_type, component_axis)`` order for a Field/Phasor detector's components.""" return [_COMPONENT_SPEC[i] for i, name in enumerate(_COMPONENT_NAMES) if name in components] @@ -522,6 +564,7 @@ def _poynting_parity(component: int, axis: int, wall: int) -> int: def _unfold_poynting(detector, state: dict[str, jax.Array], touched: tuple[int, int, int]): """Unfold a PoyntingFluxDetector (scalar/vector, summed/spatial) via per-component flux parity.""" arr = state["poynting_flux"] + on_plane = _colocated_on_plane_axes(detector, touched) components = (0, 1, 2) if detector.keep_all_components else (detector.propagation_axis,) touched_axes = [a for a in range(3) if touched[a] != 0] @@ -541,14 +584,18 @@ def _unfold_poynting(detector, state: dict[str, jax.Array], touched: tuple[int, a: jnp.asarray([_poynting_parity(i, a, touched[a]) for i in components]).reshape((1, 3, 1, 1, 1)) for a in touched_axes } - return {"poynting_flux": unfold_array(arr, touched, (2, 3, 4), signs)} + return {"poynting_flux": unfold_array(arr, touched, (2, 3, 4), signs, on_plane)} # (T, nx, ny, nz): single component along the propagation axis. p = detector.propagation_axis signs = {a: jnp.asarray(float(_poynting_parity(p, a, touched[a]))) for a in touched_axes} - return {"poynting_flux": unfold_array(arr, touched, (1, 2, 3), signs)} + return {"poynting_flux": unfold_array(arr, touched, (1, 2, 3), signs, on_plane)} -def _unfold_energy_slices(state: dict[str, jax.Array], touched: tuple[int, int, int]): +def _unfold_energy_slices( + state: dict[str, jax.Array], + touched: tuple[int, int, int], + on_plane_axes: tuple[int, ...] = (), +): """Unfold an EnergyDetector(as_slices=True): mirror each 2D plane along its in-plane symmetric axes. Energy density is even, so the parity is ``+1``; the collapsed-axis mean is already the full-domain value (energy is symmetric about the plane).""" @@ -567,7 +614,12 @@ def _unfold_energy_slices(state: dict[str, jax.Array], touched: tuple[int, int, sub[phys] = touched[phys] spatial_axes[phys] = arr_axis if any(s != 0 for s in sub): - out[key] = unfold_array(arr, (sub[0], sub[1], sub[2]), (spatial_axes[0], spatial_axes[1], spatial_axes[2])) + out[key] = unfold_array( + arr, + (sub[0], sub[1], sub[2]), + (spatial_axes[0], spatial_axes[1], spatial_axes[2]), + on_plane_axes=tuple(a for a in on_plane_axes if a in phys_axes), + ) else: out[key] = arr return out @@ -610,7 +662,7 @@ def _unfold_one_detector( return {"phasor": arr * factor.astype(arr.dtype)} arr = state["phasor"] # (1, num_freqs, num_components, nx, ny, nz) signs = _component_signs(spec, touched, arr.ndim, component_axis=2) - return {"phasor": unfold_array(arr, touched, (3, 4, 5), signs)} + return {"phasor": unfold_array(arr, touched, (3, 4, 5), signs, _colocated_on_plane_axes(detector, touched))} if isinstance(detector, FieldDetector): spec = _stored_component_spec(detector.components) @@ -623,15 +675,17 @@ def _unfold_one_detector( return {"fields": arr * factor.astype(arr.dtype)} arr = state["fields"] # (T, num_components, nx, ny, nz) signs = _component_signs(spec, touched, arr.ndim, component_axis=1) - return {"fields": unfold_array(arr, touched, (2, 3, 4), signs)} + return {"fields": unfold_array(arr, touched, (2, 3, 4), signs, _colocated_on_plane_axes(detector, touched))} if isinstance(detector, EnergyDetector): if detector.as_slices: - return _unfold_energy_slices(state, touched) + return _unfold_energy_slices(state, touched, _colocated_on_plane_axes(detector, touched)) if detector.reduce_volume: # summed energy; energy density is even -> x2 per plane return {"energy": state["energy"] * (2**count)} arr = state["energy"] # (T, nx, ny, nz) - return {"energy": unfold_array(arr, touched, (1, 2, 3))} + return { + "energy": unfold_array(arr, touched, (1, 2, 3), on_plane_axes=_colocated_on_plane_axes(detector, touched)) + } if isinstance(detector, PoyntingFluxDetector): return _unfold_poynting(detector, state, touched) @@ -682,12 +736,13 @@ def unfold_detector_states( if detector is None: new_states[name] = state continue - det_slice = detector._grid_slice_tuple - # Symmetric axes whose reduced slice starts at the symmetry plane (min edge, index 0). + # Symmetric axes on which the detector was clipped because it crossed the symmetry plane. + # A detector that merely begins at the plane but lies entirely inside the kept half was not + # clipped, so its stored output is already the full-domain result on that axis. touched: tuple[int, int, int] = ( - config.symmetry[0] if (config.symmetry[0] != 0 and det_slice[0][0] == 0) else 0, - config.symmetry[1] if (config.symmetry[1] != 0 and det_slice[1][0] == 0) else 0, - config.symmetry[2] if (config.symmetry[2] != 0 and det_slice[2][0] == 0) else 0, + config.symmetry[0] if detector.straddles_symmetry_plane(0) else 0, + config.symmetry[1] if detector.straddles_symmetry_plane(1) else 0, + config.symmetry[2] if detector.straddles_symmetry_plane(2) else 0, ) count = sum(1 for a in range(3) if touched[a] != 0) if count == 0: diff --git a/src/fdtdx/fdtd/update.py b/src/fdtdx/fdtd/update.py index 81395048..8699a134 100644 --- a/src/fdtdx/fdtd/update.py +++ b/src/fdtdx/fdtd/update.py @@ -1,3 +1,5 @@ +from typing import Literal + import jax import jax.numpy as jnp @@ -5,6 +7,7 @@ from fdtdx.constants import eta0 from fdtdx.core.misc import expand_to_3x3, pad_fields from fdtdx.core.physics.curl import curl_E, curl_H, interpolate_fields +from fdtdx.core.physics.symmetry import field_component_parity, mirror_pairs_on_plane from fdtdx.core.switch import OnOffSwitch from fdtdx.fdtd.container import ArrayContainer, ObjectContainer from fdtdx.fdtd.misc import ( @@ -96,6 +99,13 @@ def pad_fields_for_boundaries( Combines wrap/constant padding with boundary-specific corrections (e.g. Bloch phase shifts) in a single call. + On a ``config.symmetry`` axis the *min-side* halo is never wrapped. Reduction drops the min-side + periodic/Bloch boundary but keeps the far-side one, so that axis still reports wrap padding — and + wrapping would fill the halo behind the symmetry plane with the far side's field, overwriting what + the mirror condition needs there (zero for a magnetic plane, the mirrored interior that + :func:`pad_fields_with_symmetry_mirror` writes for an electric one). The far-side halo keeps + whatever boundary the user placed. + Args: fields: Field array of shape (3, Nx, Ny, Nz) objects: Container with simulation objects including boundaries @@ -108,6 +118,11 @@ def pad_fields_for_boundaries( """ periodic_axes = get_wrap_padding_axes(objects) padded = pad_fields(fields, periodic_axes) + for axis in range(3): + if config.symmetry[axis] != 0 and periodic_axes[axis]: + index: list[slice] = [slice(None)] * padded.ndim + index[axis + 1] = slice(0, 1) + padded = padded.at[tuple(index)].set(0) boundaries = objects.boundary_objects if boundaries: volume_shape = objects.volume.grid_shape @@ -121,6 +136,68 @@ def pad_fields_for_boundaries( return padded +def pad_fields_with_symmetry_mirror( + fields: jax.Array, + objects: ObjectContainer, + config: SimulationConfig, + field_type: Literal["E", "H"], +) -> jax.Array: + """Pad fields, filling the halo of every *electric* ``config.symmetry`` plane with its mirror. + + This exists for the detector co-location stencil of + :func:`~fdtdx.core.physics.curl.interpolate_fields`, which is the only consumer of the min-side + halo that a symmetry plane can get wrong. That stencil takes a backward half-step average, so for + a detector touching the plane it averages the first cell against the halo — and + :func:`pad_fields_for_boundaries` puts a zero there, where the true neighbour is the mirror image + of the interior. The detector then records exactly *half* the field in the plane row. Each + component instead takes the parity-weighted value of its own mirror partner: the neighbouring + cell for components sampled half a cell off the plane, the next one in for components sampled on + it. + + The field updates do **not** use this: the halo along an axis is only read by the updates of the + components tangential to it (an update never differentiates along its own component axis), and on + an electric plane those are precisely the components the PEC wall zeroes right afterwards, so the + halo value cannot survive. + + Only **electric** symmetry planes get the mirror. A magnetic plane sits half a cell below the + reduced domain, so the zero halo already *is* its exact mirror (tangential ``H`` vanishes there); + filling it would displace the plane by half a cell. See + :func:`~fdtdx.fdtd.symmetry.make_symmetry_walls`. A user-placed PEC/PMC boundary makes no + symmetry claim about the structure behind it either, so its halo also stays as it was. + + Args: + fields: Field array of shape (3, Nx, Ny, Nz) + objects: Container with simulation objects including boundaries + config: Simulation configuration + field_type: Whether ``fields`` holds the electric or magnetic field + + Returns: + Padded fields of shape (3, Nx+2, Ny+2, Nz+2) + """ + padded = pad_fields_for_boundaries(fields, objects, config) + for boundary in objects.boundary_objects: + if not getattr(boundary, "_is_symmetry_wall", False): + continue + axis = boundary.axis + wall = config.symmetry[axis] + if wall != -1: + # Unreachable today (only electric planes get a wall object), but the mirror is wrong for + # a magnetic plane, so state the restriction here rather than inherit it by accident. + continue + for component in range(3): + parity = field_component_parity(field_type, component, axis, wall) + # Padded index 0 is one cell below the domain. Its mirror partner is the domain's first + # cell for a half-cell-offset component (padded index 1) and the second one for a + # component sampled on the plane itself (padded index 2). + source_index = 2 if mirror_pairs_on_plane(field_type, component, axis, wall) else 1 + target: list[slice] = [slice(None)] * 3 + source: list[slice] = [slice(None)] * 3 + target[axis] = slice(0, 1) + source[axis] = slice(source_index, source_index + 1) + padded = padded.at[component, *target].set(parity * padded[component, *source]) + return padded + + def get_anisotropic_averaging_widths( config: SimulationConfig, ) -> tuple[jax.Array, jax.Array, jax.Array] | None: @@ -1007,8 +1084,8 @@ def is_interior(detector: Detector) -> bool: full = None if any(d.exact_interpolation and not is_interior(d) for d in to_update): full = interpolate_fields( - E_pad=pad_fields_for_boundaries(arrays.fields.E, objects, config), - H_pad=pad_fields_for_boundaries((H_prev + arrays.fields.H) / 2, objects, config), + E_pad=pad_fields_with_symmetry_mirror(arrays.fields.E, objects, config, "E"), + H_pad=pad_fields_with_symmetry_mirror((H_prev + arrays.fields.H) / 2, objects, config, "H"), config=config, ) diff --git a/src/fdtdx/objects/boundaries/boundary.py b/src/fdtdx/objects/boundaries/boundary.py index 4e426507..d5d49837 100644 --- a/src/fdtdx/objects/boundaries/boundary.py +++ b/src/fdtdx/objects/boundaries/boundary.py @@ -3,7 +3,7 @@ import jax -from fdtdx.core.jax.pytrees import autoinit, frozen_field +from fdtdx.core.jax.pytrees import autoinit, frozen_field, frozen_private_field from fdtdx.objects.object import SimulationObject from fdtdx.typing import GridShape3D, Slice3D, SliceTuple3D @@ -23,6 +23,14 @@ class BaseBoundary(SimulationObject, ABC): #: Direction along axis ("+" or "-") direction: Literal["+", "-"] = frozen_field() + #: Whether this boundary was inserted by ``config.symmetry`` as a mirror plane rather than + #: requested by the user as an ordinary boundary. Only electric symmetry planes get a wall + #: object (see ``make_symmetry_walls``, which sets this flag), and such a wall asserts that the + #: model is mirror-symmetric about it, which lets the detector co-location stencil fill its halo + #: with the mirrored interior values instead of zeros. A user-placed PEC/PMC boundary makes no + #: such claim, so its halo stays as it was. + _is_symmetry_wall: bool = frozen_private_field(default=False) + @property @abstractmethod def descriptive_name(self) -> str: diff --git a/src/fdtdx/objects/boundaries/pmc.py b/src/fdtdx/objects/boundaries/pmc.py index bea47ff2..f37f0a02 100644 --- a/src/fdtdx/objects/boundaries/pmc.py +++ b/src/fdtdx/objects/boundaries/pmc.py @@ -61,7 +61,14 @@ def tangential_components(self) -> tuple[int, int]: @override def apply_post_H_update(self, H: jax.Array) -> jax.Array: - """Zeros tangential H components at this PMC boundary face.""" + """Zeros tangential H components at this PMC boundary face. + + ``config.symmetry`` never creates this boundary: a magnetic symmetry plane sits half a cell + below the reduced domain, where the zero field halo already is the mirror, so zeroing + tangential H at the first cell would impose the condition half a cell off the plane (see + :func:`~fdtdx.fdtd.symmetry.make_symmetry_walls`). This method therefore only ever runs for a + user-placed PMC boundary, whose behaviour is unchanged. + """ comp1, comp2 = self.tangential_components sx, sy, sz = self.grid_slice H = H.at[comp1, sx, sy, sz].set(0) diff --git a/src/fdtdx/objects/detectors/diffractive.py b/src/fdtdx/objects/detectors/diffractive.py index 0d615b81..8308c152 100644 --- a/src/fdtdx/objects/detectors/diffractive.py +++ b/src/fdtdx/objects/detectors/diffractive.py @@ -41,6 +41,27 @@ def __post_init__(self): if len(self.orders) == 0: raise ValueError("DiffractiveDetector requires at least one diffraction order.") + def validate_placement(self, objects) -> list[str]: + """Reject a diffractive detector clipped by a symmetry plane. + + The diffraction orders are defined by the plane's periodicity, which the reduction halves. + The efficiencies computed on the reduced plane therefore belong to a different order basis + than the full-domain ones and are not recoverable afterwards (see + :func:`fdtdx.unfold_detector_states`, which raises for the same reason). + """ + errors = list(super().validate_placement(objects)) + clipped = [a for a in range(3) if self.straddles_symmetry_plane(a)] + if clipped: + axis_names = ", ".join("xyz"[a] for a in clipped) + errors.append( + f"DiffractiveDetector '{self.name}' crosses the {axis_names}-symmetry plane. Its " + f"diffraction-order basis is set by the plane's period, which config.symmetry halves, so " + f"the computed efficiencies would refer to the reduced period rather than the physical " + f"one. Run without config.symmetry, or record the fields instead and recompute the " + f"orders after fdtdx.unfold_fields." + ) + return errors + @property def propagation_axis(self) -> int: """Determines the axis along which diffraction is measured. diff --git a/src/fdtdx/objects/detectors/mode.py b/src/fdtdx/objects/detectors/mode.py index f6b6739a..0a2b1601 100644 --- a/src/fdtdx/objects/detectors/mode.py +++ b/src/fdtdx/objects/detectors/mode.py @@ -15,7 +15,7 @@ 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.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 DetectorState @@ -490,10 +490,12 @@ class ModeOverlapDetector(BaseModeOverlapDetector): #: 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. + #: ``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( @@ -556,6 +558,26 @@ def _compute_mode_fields( 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, + ) return compute_mode( frequency=wave_character.get_frequency(), inv_permittivities=inv_permittivity_slice, diff --git a/src/fdtdx/objects/object.py b/src/fdtdx/objects/object.py index e99c16f6..aeb7d933 100644 --- a/src/fdtdx/objects/object.py +++ b/src/fdtdx/objects/object.py @@ -243,6 +243,12 @@ class SimulationObject(TreeClass, ABC): _grid_slice_tuple: SliceTuple3D = frozen_private_field( default=INVALID_SLICE_TUPLE_3D, ) + #: Slice this object would occupy without mirror-symmetry reduction, expressed in the + #: reduced grid's coordinates. Set by place_objects only when config.symmetry is active; + #: read via unreduced_grid_slice_tuple. + _unreduced_grid_slice_tuple: SliceTuple3D = frozen_private_field( + default=INVALID_SLICE_TUPLE_3D, + ) _config: SimulationConfig = private_field() @property @@ -251,6 +257,80 @@ def grid_slice_tuple(self) -> SliceTuple3D: raise Exception(f"Object is not yet initialized: {self}") return self._grid_slice_tuple + @property + def unreduced_grid_slice_tuple(self) -> SliceTuple3D: + """Full-domain slice of this object, in the reduced grid's coordinate frame. + + Identical to :attr:`grid_slice_tuple` unless ``config.symmetry`` clipped this object. + Where the object crossed a symmetry plane the start index is **negative**: the plane sits + at reduced index 0, so an object that straddled it extends to negative coordinates. Object + contents that depend on the object's full extent (e.g. a Gaussian beam's centre) must be + derived from this slice, not from the clipped one. + + Returns: + SliceTuple3D: Per-axis ``(start, stop)`` in reduced-grid coordinates. + """ + if self._unreduced_grid_slice_tuple == INVALID_SLICE_TUPLE_3D: + return self.grid_slice_tuple + return self._unreduced_grid_slice_tuple + + @property + def unreduced_grid_shape(self) -> GridShape3D: + """Grid shape this object would have without mirror-symmetry reduction.""" + tpl = self.unreduced_grid_slice_tuple + return ( + tpl[0][1] - tpl[0][0], + tpl[1][1] - tpl[1][0], + tpl[2][1] - tpl[2][0], + ) + + def straddles_symmetry_plane(self, axis: int) -> bool: + """Whether ``config.symmetry`` clipped this object on ``axis`` because it crossed the plane. + + This is strictly stronger than "the clipped slice starts at index 0": an object that + happens to begin exactly at the symmetry plane but lies entirely inside the kept half was + not clipped and must not be treated as mirror-extended. + + Args: + axis (int): Physical axis (0=x, 1=y, 2=z). + + Returns: + bool: True if the object extends past the symmetry plane into the discarded half. + """ + return self.unreduced_grid_slice_tuple[axis][0] < 0 + + @property + def straddled_symmetry_axes(self) -> tuple[int, ...]: + """Axes on which this object was clipped by a symmetry plane (see straddles_symmetry_plane).""" + return tuple(a for a in range(3) if self.straddles_symmetry_plane(a)) + + def symmetry_mirror_axes(self, exclude_axis: int | None = None) -> tuple[int, ...]: + """Axes that must be mirrored to reconstruct this object's full-domain extent. + + Args: + exclude_axis (int | None): Axis to skip, e.g. a plane object's propagation axis, which + is one cell thick and therefore never split by a plane. + + Returns: + tuple[int, ...]: Straddled symmetry axes, excluding ``exclude_axis``. + """ + return tuple(a for a in self.straddled_symmetry_axes if a != exclude_axis) + + def touches_symmetry_plane(self, axis: int) -> bool: + """Whether this object's placed slice begins on the symmetry plane of ``axis``. + + True both for objects clipped by the plane and for objects that merely start there; use + :meth:`straddles_symmetry_plane` to distinguish the two. + + Args: + axis (int): Physical axis (0=x, 1=y, 2=z). + + Returns: + bool: True if ``config.symmetry`` is active on ``axis`` and the object sits at its + reduced min edge, where the PEC/PMC mirror wall lives. + """ + return self._config.symmetry[axis] != 0 and self.grid_slice_tuple[axis][0] == 0 + @property def grid_slice(self) -> Slice3D: tpl = ensure_slice_tuple(self._grid_slice_tuple) diff --git a/src/fdtdx/objects/sources/dipole.py b/src/fdtdx/objects/sources/dipole.py index c1001856..40ddf6a1 100644 --- a/src/fdtdx/objects/sources/dipole.py +++ b/src/fdtdx/objects/sources/dipole.py @@ -107,6 +107,26 @@ def __post_init__(self): if self.polarization not in (0, 1, 2): raise ValueError(f"polarization must be 0, 1, or 2, got {self.polarization}") + def validate_placement(self, objects) -> list[str]: + """Reject a dipole sitting on a symmetry plane. + + The mirror plane lies on a cell *edge*, so no cell is centred on it. A dipole placed in the + first cell of the reduced domain therefore does not stand for one dipole on the plane: the + reduced simulation models it together with its mirror image, i.e. two dipoles half a cell + apart, which is not the model the user drew in the full domain. + """ + errors = list(super().validate_placement(objects)) + on_plane = [a for a in range(3) if self.touches_symmetry_plane(a)] + if on_plane: + axis_names = ", ".join("xyz"[a] for a in on_plane) + errors.append( + f"Point dipole '{self.name}' sits on the {axis_names}-symmetry plane. The plane lies on a " + f"cell edge, so a single dipole cannot be centred on it - the reduced simulation would " + f"model the dipole plus its mirror image, half a cell apart. Move the dipole off the " + f"plane, or drop the symmetry on that axis." + ) + return errors + @property def _orientation(self) -> jnp.ndarray: """Normalized orientation vector as a (3,) JAX array. diff --git a/src/fdtdx/objects/sources/linear_polarization.py b/src/fdtdx/objects/sources/linear_polarization.py index 1e159841..f7d22f79 100644 --- a/src/fdtdx/objects/sources/linear_polarization.py +++ b/src/fdtdx/objects/sources/linear_polarization.py @@ -4,11 +4,17 @@ import jax import jax.numpy as jnp import numpy as np +from loguru import logger from fdtdx.core.grid import calculate_time_offset_yee from fdtdx.core.jax.pytrees import autoinit, frozen_field -from fdtdx.core.misc import linear_interpolated_indexing, tilted_polarization_vectors +from fdtdx.core.misc import ( + linear_interpolated_indexing, + normalize_polarization_for_source, + tilted_polarization_vectors, +) from fdtdx.core.physics.metrics import compute_energy +from fdtdx.core.physics.symmetry import field_component_parity from fdtdx.dispersion import effective_inv_permittivity from fdtdx.objects.sources.tfsf import TFSFPlaneSource, _build_dispersive_H_filter, _source_impedance @@ -116,6 +122,49 @@ def _uses_physical_source_coordinates(self) -> bool: """Whether transverse source coordinates are represented in metres.""" return self._config.has_nonuniform_grid + def validate_placement(self, objects) -> list[str]: + """Warn when the polarization is incompatible with the symmetry walls it crosses. + + The injected transverse profile is mirror-*even* about every symmetry plane the source + straddles. A wall makes some field components even and others odd (see + :func:`~fdtdx.core.physics.symmetry.field_component_parity`), so if a nonzero component is odd there, + the wall drives it toward zero and the reduced simulation models a different field than the + user drew. Picking the wall type per polarization is the single most common mistake with + ``config.symmetry``, and it is silent otherwise: this reports it with the fix. + """ + errors = list(super().validate_placement(objects)) + straddled = [a for a in self.transverse_axes if self.straddles_symmetry_plane(a)] + if not straddled: + return errors + + e_pol, h_pol = normalize_polarization_for_source( + 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, + dtype=self._config.dtype, + ) + for axis in straddled: + wall = self._config.symmetry[axis] + odd: list[str] = [] + for field_type, polarization in (("E", e_pol), ("H", h_pol)): + for component in range(3): + if abs(float(polarization[component])) < 1e-6: + continue + if field_component_parity(field_type, component, axis, wall) == -1: + odd.append(f"{field_type}{'xyz'[component]}") + if odd: + wall_name = "PMC" if wall == 1 else "PEC" + other = "PEC (-1)" if wall == 1 else "PMC (+1)" + logger.warning( + f"Source '{self.name}' crosses the {'xyz'[axis]}-symmetry plane, where " + f"config.symmetry[{axis}]={wall:+d} ({wall_name}) makes {', '.join(odd)} odd — but the " + f"source injects a mirror-even profile in those components. The wall will suppress " + f"the injected field. For this polarization use {other} on that axis, or place the " + f"source so it does not cross the plane." + ) + return errors + def apply( self: Self, key: jax.Array, @@ -175,7 +224,7 @@ def apply( ) # update is amplitude multiplied by polarization - amplitude_raw = self._get_amplitude_raw(center)[None, ...] + amplitude_raw = self._get_amplitude_raw(center) # map amplitude to propagation plane. Uniform grids keep the legacy # index-space projection; non-uniform grids project physical transverse @@ -196,18 +245,25 @@ def apply( ) wh_coords = jnp.stack((w, h), axis=-1) wh_coords -= center - # basis in plane + # Orthonormal in-plane basis. u follows the horizontal axis (projected perpendicular to the + # wave vector for tilted incidence); v completes a right-handed triple with the *forward* + # propagation direction, so the untilted projection below is the identity for both + # directions. Deriving v from the signed wave vector instead would mirror the transverse + # profile about the center whenever direction == "-" (invisible for a centered, radially + # symmetric profile, wrong for every other one). h_list = [0, 0, 0] h_list[self.horizontal_axis] = 1 h_axis = jnp.asarray(h_list, dtype=self._config.dtype) u_basis = h_axis - jnp.dot(h_axis, wave_vector) * wave_vector u_basis = u_basis / jnp.linalg.norm(u_basis) - v_basis = jnp.cross(wave_vector, u_basis) + direction_sign = 1.0 if self.direction == "+" else -1.0 + v_basis = direction_sign * jnp.cross(wave_vector, u_basis) # projection def project(point): - point_list = [point[0], point[1]] - point_list.insert(self.propagation_axis, 0) + point_list = [jnp.zeros((), dtype=self._config.dtype)] * 3 + point_list[self.horizontal_axis] = point[0] + point_list[self.vertical_axis] = point[1] point = jnp.asarray(point_list, dtype=self._config.dtype) projection = point - jnp.dot(point, wave_vector) * wave_vector # Convert to plane coordinates @@ -217,19 +273,18 @@ def project(point): float_projected = jax.vmap(project)(wh_coords.reshape(-1, 2)) float_projected += center + profile_hv = self._grid_to_hv(amplitude_raw) if self._uses_physical_source_coordinates(): index_fn = jax.vmap( _linear_interpolate_rectilinear_2d, in_axes=(0, None, None, None), ) - profile_2d = jnp.take(amplitude_raw[0], 0, axis=self.propagation_axis) - interp = index_fn(float_projected, horizontal_centers, vertical_centers, profile_2d) + interp = index_fn(float_projected, horizontal_centers, vertical_centers, profile_hv) else: # interpolate floating indices in original array index_fn = jax.vmap(linear_interpolated_indexing, in_axes=(0, None)) - profile_2d = jnp.take(amplitude_raw[0], 0, axis=self.propagation_axis) - interp = index_fn(float_projected, profile_2d) - amplitude = interp.reshape(*amplitude_raw.shape) + interp = index_fn(float_projected, profile_hv) + amplitude = self._hv_to_grid(interp.reshape(profile_hv.shape))[None, ...] E = amplitude * e_pol[:, None, None, None] H = amplitude * h_pol[:, None, None, None] @@ -241,7 +296,12 @@ def project(point): inv_permittivity=inv_permittivities, inv_permeability=inv_permeabilities, ) - total_energy_root = jnp.sqrt(energy.sum()) + # Normalize by the energy of the *full-domain* source. Under config.symmetry the slice + # covers only one half/quarter of it, and the profile is mirror-symmetric about each + # symmetry plane it straddles, so the full-domain sum is the reduced sum times the + # plane multiplicity. Without this the reduced run would inject 2**(k/2) times the + # amplitude of the equivalent full-domain run. + total_energy_root = jnp.sqrt(energy.sum() * self.symmetry_profile_multiplicity) E = E / total_energy_root H = H / total_energy_root @@ -316,37 +376,64 @@ class GaussianPlaneSource(LinearlyPolarizedPlaneSource): std: float = frozen_field(default=1 / 3) # relative to radius @staticmethod - def _gauss_profile( + def _gauss_profile_2d( width: int, height: int, - axis: int, center: tuple[float, float] | jax.Array, radii: tuple[float, float], std: float, - ) -> jax.Array: # shape (*grid_shape) - # (width, height, 2) grid in the same (horizontal, vertical) order as - # ``center`` and ``radii``; an xy-indexed meshgrid swaps the - # coordinates on non-square planes and misplaces the spot. + normalization_multiplicity: int = 1, + ) -> jax.Array: # shape (width, height) + """Truncated Gaussian spot on a ``(width, height)`` transverse grid, normalized to unit sum. + + ``width``/``height``, ``center`` and ``radii`` are all in the same (horizontal, vertical) + order; an xy-indexed meshgrid would swap the coordinates on non-square planes and misplace + the spot. ``normalization_multiplicity`` divides out the number of copies of this plane in + the full domain (see ``TFSFPlaneSource.symmetry_profile_multiplicity``), so a plane clipped + by a symmetry plane still carries the amplitude of the full-domain profile it is part of. + """ grid = ( jnp.stack(jnp.meshgrid(jnp.arange(width), jnp.arange(height), indexing="ij"), axis=-1) - jnp.asarray(center) ) / jnp.asarray(radii) euc_dist = (grid**2).sum(axis=-1) mask = euc_dist < 1 - mask = jnp.expand_dims(mask, axis=axis) - exp_part = jnp.exp(-0.5 * euc_dist / std**2) - exp_part = jnp.expand_dims(exp_part, axis=axis) profile = jnp.where(mask, exp_part, 0) - profile = profile / profile.sum() + return profile / (profile.sum() * normalization_multiplicity) + + @staticmethod + def _gauss_profile( + width: int, + height: int, + axis: int, + center: tuple[float, float] | jax.Array, + radii: tuple[float, float], + std: float, + normalization_multiplicity: int = 1, + ) -> jax.Array: # shape (*grid_shape) + """:meth:`_gauss_profile_2d` with a singleton inserted at ``axis``. - return profile + ``width``/``height`` are the sizes along the two axes *other than* ``axis``, in ascending + array-axis order — which is (horizontal, vertical) for propagation along x or z but + (vertical, horizontal) for propagation along y. + """ + profile = GaussianPlaneSource._gauss_profile_2d( + width=width, + height=height, + center=center, + radii=radii, + std=std, + normalization_multiplicity=normalization_multiplicity, + ) + return jnp.expand_dims(profile, axis=axis) def _get_amplitude_raw( self, center: jax.Array, ) -> jax.Array: + multiplicity = self.symmetry_profile_multiplicity if self._config.has_nonuniform_grid: local_edges = self._local_edge_coordinates() assert local_edges is not None @@ -366,19 +453,19 @@ def _get_amplitude_raw( h_widths = horizontal_edges[1:] - horizontal_edges[:-1] v_widths = vertical_edges[1:] - vertical_edges[:-1] cell_areas = h_widths[:, None] * v_widths[None, :] - profile_2d = profile_2d / (profile_2d * cell_areas).sum() - return jnp.expand_dims(profile_2d, axis=self.propagation_axis) + profile_2d = profile_2d / ((profile_2d * cell_areas).sum() * multiplicity) + return self._hv_to_grid(profile_2d) grid_radius = self.radius / self._config.uniform_spacing() - profile = self._gauss_profile( + profile_hv = self._gauss_profile_2d( width=self.grid_shape[self.horizontal_axis], height=self.grid_shape[self.vertical_axis], - axis=self.propagation_axis, center=center, radii=(grid_radius, grid_radius), std=self.std, + normalization_multiplicity=multiplicity, ) - return profile + return self._hv_to_grid(profile_hv) @autoinit diff --git a/src/fdtdx/objects/sources/mode.py b/src/fdtdx/objects/sources/mode.py index 29ca026f..3cb1210f 100644 --- a/src/fdtdx/objects/sources/mode.py +++ b/src/fdtdx/objects/sources/mode.py @@ -12,7 +12,7 @@ from fdtdx.core.jax.pytrees import autoinit, frozen_field, private_field from fdtdx.core.linalg import get_wave_vector_raw from fdtdx.core.physics.metrics import compute_energy -from fdtdx.core.physics.modes import compute_mode +from fdtdx.core.physics.modes import compute_mode, compute_mode_symmetry_reduced from fdtdx.dispersion import effective_complex_inv_permittivity, effective_inv_permittivity from fdtdx.objects.sources.tfsf import TFSFPlaneSource, _build_dispersive_H_filter @@ -28,10 +28,12 @@ class ModePlaneSource(TFSFPlaneSource): #: 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 source plane's waveguide lies on a symmetry plane of a - #: reduced (half/quarter) domain, so the mode solver imposes the same wall - #: the FDTD uses there (e.g. ``(0, 1)`` for PEC at y=0 and PMC at the z - #: Si-mid plane of a +x-propagating quarter domain). + #: 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 is *not* needed - and is + #: ignored - when ``config.symmetry`` performs the reduction: there the mode is solved on the + #: mirrored full cross-section and restricted, which reproduces the full-domain mode exactly + #: (the mode solver's symmetric solve does not, because FDTDX rasterizes materials per cell + #: while the solver samples them on its staggered grid). symmetry: tuple[int, int] = frozen_field(default=(0, 0)) _inv_permittivity: jax.Array = private_field() @@ -171,18 +173,37 @@ def apply( ) # compute mode - mode_E, mode_H, eff_index = compute_mode( - frequency=self.wave_character.get_frequency(), - inv_permittivities=mode_inv_permittivity, - 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, - symmetry=self.symmetry, - transverse_coords=self._transverse_edge_coordinates(), - ) + mirrored_axes = self.symmetry_mirror_axes(exclude_axis=self.propagation_axis) + if mirrored_axes: + # Symmetry-reduced cross-section: solve on the mirrored full cross-section and restrict, + # so the injected mode is exactly the one the unreduced simulation would launch. + mode_E, mode_H, eff_index = compute_mode_symmetry_reduced( + mirrored_axes=mirrored_axes, + walls={axis: self._config.symmetry[axis] for axis in mirrored_axes}, + frequency=self.wave_character.get_frequency(), + inv_permittivities=mode_inv_permittivity, + 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, + transverse_coords=self._transverse_edge_coordinates(), + object_name=self.name, + ) + else: + mode_E, mode_H, eff_index = compute_mode( + frequency=self.wave_character.get_frequency(), + inv_permittivities=mode_inv_permittivity, + 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, + symmetry=self.symmetry, + transverse_coords=self._transverse_edge_coordinates(), + ) # Keep the complex modal fields when the mode was solved against a lossy # (conductivity) permittivity, so the launched source carries the # eigenmode's transverse phase — TFSFPlaneSource.update_E/update_H inject diff --git a/src/fdtdx/objects/sources/tfsf.py b/src/fdtdx/objects/sources/tfsf.py index efefcb50..65145095 100644 --- a/src/fdtdx/objects/sources/tfsf.py +++ b/src/fdtdx/objects/sources/tfsf.py @@ -527,29 +527,88 @@ def _get_azimuth_elevation( ) return azimuth_radians, elevation_radians + @property + def transverse_axes(self) -> tuple[int, int]: + """The (horizontal, vertical) axis pair of this source's plane.""" + return self.horizontal_axis, self.vertical_axis + + @property + def _transverse_axes_are_swapped(self) -> bool: + """Whether (horizontal, vertical) is *descending* in array-axis order. + + True exactly for propagation along y, where ``get_oriented_transverse_axes(1) == (2, 0)``. + Transverse profiles are built in (horizontal, vertical) order, so on that axis they must be + transposed before being placed into the source's ``(Nx, Ny, Nz)`` slice. + """ + return self.horizontal_axis > self.vertical_axis + + def _hv_to_grid(self, profile_hv: jax.Array) -> jax.Array: + """Place a ``(horizontal, vertical)`` transverse profile into this source's grid shape.""" + if self._transverse_axes_are_swapped: + profile_hv = profile_hv.T + return jnp.expand_dims(profile_hv, axis=self.propagation_axis) + + def _grid_to_hv(self, profile_grid: jax.Array) -> jax.Array: + """Extract the ``(horizontal, vertical)`` transverse profile from a grid-shaped array.""" + profile_2d = jnp.take(profile_grid, 0, axis=self.propagation_axis) + if self._transverse_axes_are_swapped: + profile_2d = profile_2d.T + return profile_2d + + @property + def symmetry_profile_multiplicity(self) -> int: + """Number of copies of this source's plane in the full domain, ``2**k``. + + ``k`` counts the transverse symmetry planes the source straddles. Amplitude + normalizations that sum over the source plane (the Gaussian profile normalization, + ``normalize_by_energy``) must divide by the *full-domain* sum, which for a profile that is + mirror-symmetric about each of those planes is this factor times the reduced-plane sum. + Without it, a symmetry-reduced run would inject ``2**(k/2)`` times the amplitude of the + equivalent full-domain run. + """ + return 2 ** sum(1 for a in self.transverse_axes if self.straddles_symmetry_plane(a)) + + def _unreduced_transverse_bounds(self, axis: int) -> tuple[float, float]: + """Lower/upper coordinate of this source's *unclipped* extent along a transverse axis. + + Returned in the same local coordinates the transverse profile is sampled in: cell indices + relative to the placed slice on a uniform grid (index ``i`` = centre of cell ``i``), metres + from the placed slice's lower edge on a non-uniform grid. Under ``config.symmetry`` an + object that crossed the symmetry plane extends to *negative* local coordinates, which is + what keeps a straddling Gaussian centred on the plane instead of on the kept quadrant. + """ + placed_start = self.grid_slice_tuple[axis][0] + full_start, full_stop = self.unreduced_grid_slice_tuple[axis] + discarded = placed_start - full_start # 0 unless a symmetry plane clipped this axis + if not self._config.has_nonuniform_grid: + return float(-discarded), float(full_stop - full_start - 1 - discarded) + grid = self._config.resolved_grid + assert grid is not None + edges = grid.edges(axis) + origin = edges[placed_start] + # The grid is validated mirror-symmetric about the plane, so the widths of the discarded + # cells equal those of the first kept cells: the lower bound mirrors edge ``discarded``. + lower = -float(edges[placed_start + discarded] - origin) + upper = float(edges[placed_start + (full_stop - full_start) - discarded] - origin) + return lower, upper + def _get_center(self, key: jax.Array) -> jax.Array: # shape(2,) """Calculate the randomized source center. + The center is the middle of the source's *unreduced* extent, so a source clipped by a + symmetry plane keeps the profile center of the full-domain source it stands for (half a + cell below the reduced min edge for a plane-symmetric source) instead of being re-centered + on the surviving part. + Uniform-grid sources use the legacy index-space center. On a non-uniform grid, tilted projections and Gaussian profiles are sampled in physical coordinates, so the returned center is measured in metres from the source-slice lower edge along the transverse axes. """ - if self._config.has_nonuniform_grid: - grid = self._config.resolved_grid - assert grid is not None - local_edges = [] - for axis in (self.horizontal_axis, self.vertical_axis): - lower, upper = self.grid_slice_tuple[axis] - edges = grid.edges(axis)[lower : upper + 1] - local_edges.append(edges - edges[0]) - center_horizontal = 0.5 * local_edges[0][-1] - center_vertical = 0.5 * local_edges[1][-1] - else: - horizontal_size = self.grid_shape[self.horizontal_axis] - vertical_size = self.grid_shape[self.vertical_axis] - center_horizontal = (horizontal_size - 1) / 2 - center_vertical = (vertical_size - 1) / 2 + h_lower, h_upper = self._unreduced_transverse_bounds(self.horizontal_axis) + v_lower, v_upper = self._unreduced_transverse_bounds(self.vertical_axis) + center_horizontal = 0.5 * (h_lower + h_upper) + center_vertical = 0.5 * (v_lower + v_upper) key, subkey = jax.random.split(key) horizontal_offset = jax.random.uniform( @@ -580,6 +639,44 @@ def _get_random_parts(self, key: jax.Array): return center, azimuth, elevation + def validate_placement(self, objects) -> list[str]: + """Reject source configurations that a symmetry-reduced domain cannot represent. + + A reduced simulation replaces the discarded half by the mirror image of the kept half, so a + source clipped by a symmetry plane must itself be mirror-symmetric about it. An off-normal + (tilted) beam and a randomly displaced beam both break that, and a plane source sitting on a + symmetry plane normal to its own propagation axis would inject into the mirror wall. + """ + errors = list(super().validate_placement(objects)) + straddled = [a for a in self.transverse_axes if self.straddles_symmetry_plane(a)] + if straddled: + axis_names = ", ".join("xyz"[a] for a in straddled) + if self.azimuth_angle != 0.0 or self.elevation_angle != 0.0 or self.max_angle_random_offset != 0.0: + errors.append( + f"Source '{self.name}' is tilted (azimuth={self.azimuth_angle}, " + f"elevation={self.elevation_angle}, max_angle_random_offset=" + f"{self.max_angle_random_offset}) and crosses the {axis_names}-symmetry plane. A " + f"tilted beam is not mirror-symmetric about that plane, so config.symmetry cannot " + f"model it. Drop the tilt, or drop the symmetry on that axis." + ) + if self.max_horizontal_offset != 0.0 or self.max_vertical_offset != 0.0: + errors.append( + f"Source '{self.name}' has a random transverse offset " + f"(max_horizontal_offset={self.max_horizontal_offset}, " + f"max_vertical_offset={self.max_vertical_offset}) and crosses the " + f"{axis_names}-symmetry plane. A displaced beam is not mirror-symmetric about that " + f"plane, so config.symmetry cannot model it. Drop the offset, or drop the symmetry " + f"on that axis." + ) + if self.touches_symmetry_plane(self.propagation_axis): + errors.append( + f"Source '{self.name}' lies on the {'xyz'[self.propagation_axis]}-symmetry plane, which is " + f"normal to its propagation axis. The PEC/PMC mirror wall sits in exactly those cells " + f"and would clobber the injected field. Move the source off the plane, or drop the " + f"symmetry on that axis." + ) + return errors + @abstractmethod def apply( self: Self, diff --git a/src/fdtdx/objects/sources/tfsf_region.py b/src/fdtdx/objects/sources/tfsf_region.py index 0d86797c..2e23124c 100644 --- a/src/fdtdx/objects/sources/tfsf_region.py +++ b/src/fdtdx/objects/sources/tfsf_region.py @@ -483,6 +483,19 @@ def validate_placement(self, objects: "ObjectContainer") -> list[str]: rejected in v1 (it would require a matching Bloch phase). """ errors: list[str] = [] + # config.symmetry clips objects onto the kept half, which for a closed TFSF box means + # silently deleting the faces that fall in the discarded half - the box would no longer + # enclose anything. Reject it instead of injecting through a partial surface. + symmetry_axes = [a for a in range(3) if self._config.symmetry[a] != 0] + clipped = [a for a in symmetry_axes if self.straddles_symmetry_plane(a) or self.touches_symmetry_plane(a)] + if clipped: + axis_names = ", ".join("xyz"[a] for a in clipped) + errors.append( + f"TFSF region '{self.name}' reaches the {axis_names}-symmetry plane. A closed TFSF box " + f"cannot be split by a mirror plane: config.symmetry would drop the faces in the " + f"discarded half, leaving an open surface. Keep the box strictly inside the kept " + f"(upper) half of each symmetric axis, or drop config.symmetry." + ) p_axis = self.propagation_axis transverse = get_oriented_transverse_axes(p_axis) periodic = tuple(self.periodic_axes) diff --git a/tests/integration/fdtd/test_symmetry_mode_source.py b/tests/integration/fdtd/test_symmetry_mode_source.py new file mode 100644 index 00000000..88c31078 --- /dev/null +++ b/tests/integration/fdtd/test_symmetry_mode_source.py @@ -0,0 +1,247 @@ +"""Mode sources and mode-overlap detectors under ``config.symmetry``. + +A symmetry-reduced run must launch the same guided mode as the full-domain run, restricted to the +kept half. Solving the mode on the reduced cross-section with the mode solver's own symmetric solve +does *not* achieve that: FDTDX rasterizes materials per cell and hands the same array to every +component, while the solver samples them on its staggered Yee grid, which breaks the discrete +mirror symmetry and shifts ``neff`` at first order in the cell size. These tests pin the +mirror-then-solve route that reproduces the full-domain mode instead. + +Only mode solves run here, no FDTD time stepping. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import fdtdx + +_SPACING = 25e-9 +_WAVELENGTH = 1.55e-6 +_TRANSVERSE_CELLS = 40 # 1 um cross-section, even so every symmetric axis splits exactly +_PROPAGATION_CELLS = 8 +_CORE_CELLS = (16, 8) # 400 x 200 nm silicon core, centered => symmetric about both center planes + + +def _build(symmetry, mode_index=0, mode_symmetry=None, detector=False): + """Centered Si waveguide with a +x mode source (and optionally a mode-overlap detector).""" + config = fdtdx.SimulationConfig( + grid=fdtdx.UniformGrid(spacing=_SPACING), + time=5e-15, + dtype=jnp.float32, + symmetry=symmetry, + ) + volume = fdtdx.SimulationVolume(partial_grid_shape=(_PROPAGATION_CELLS, _TRANSVERSE_CELLS, _TRANSVERSE_CELLS)) + objects, constraints = [volume], [] + bound_dict, boundary_constraints = fdtdx.boundary_objects_from_config( + fdtdx.BoundaryConfig.from_uniform_bound(thickness=2), volume + ) + constraints.extend(boundary_constraints) + objects.extend(bound_dict.values()) + + core = fdtdx.UniformMaterialObject( + name="core", + partial_grid_shape=(None, _CORE_CELLS[0], _CORE_CELLS[1]), + material=fdtdx.Material(permittivity=fdtdx.constants.relative_permittivity_silicon), + ) + constraints.extend([core.same_size(volume, axes=(0,)), core.place_at_center(volume, axes=(0, 1, 2))]) + objects.append(core) + + wave = fdtdx.WaveCharacter(wavelength=_WAVELENGTH) + extra = {} if mode_symmetry is None else {"symmetry": mode_symmetry} + source = fdtdx.ModePlaneSource( + name="src", + partial_grid_shape=(1, None, None), + wave_character=wave, + direction="+", + mode_index=mode_index, + **extra, + ) + constraints.extend( + [ + source.same_size(volume, axes=(1, 2)), + source.place_at_center(volume, axes=(1, 2)), + source.set_grid_coordinates(axes=(0,), sides=("-",), coordinates=(2,)), + ] + ) + objects.append(source) + + if detector: + overlap = fdtdx.ModeOverlapDetector( + name="det", + partial_grid_shape=(1, None, None), + wave_characters=(wave,), + direction="+", + mode_index=mode_index, + ) + constraints.extend( + [ + overlap.same_size(volume, axes=(1, 2)), + overlap.place_at_center(volume, axes=(1, 2)), + overlap.set_grid_coordinates(axes=(0,), sides=("-",), coordinates=(5,)), + ] + ) + objects.append(overlap) + + return objects, constraints, config + + +def _apply(symmetry, **kwargs): + objects, constraints, config = _build(symmetry, **kwargs) + key = jax.random.PRNGKey(0) + container, arrays, params, config, _ = fdtdx.place_objects( + object_list=objects, config=config, constraints=constraints, key=key + ) + arrays, container, _ = fdtdx.apply_params(arrays, container, params, key=key) + return container, config + + +def _by_name(container, name): + return next(obj for obj in container.objects if obj.name == name) + + +def _kept_half(array: np.ndarray, symmetry) -> np.ndarray: + """Upper half along every symmetric axis of a ``(3, Nx, Ny, Nz)`` array.""" + index: list[slice] = [slice(None)] * 4 + for axis in range(3): + if symmetry[axis] != 0: + index[1 + axis] = slice(array.shape[1 + axis] // 2, None) + return array[tuple(index)] + + +# PEC on y (Ey is normal there -> even) and PMC on z (Ey is tangential -> even): the walls that +# match the fundamental quasi-TE mode of this waveguide. +_HALF_Y = (0, -1, 0) +_HALF_Z = (0, 0, 1) +_QUARTER = (0, -1, 1) + + +class TestModeSourceUnderSymmetry: + @pytest.mark.parametrize( + "symmetry, label", + [(_HALF_Y, "half-y-PEC"), (_HALF_Z, "half-z-PMC"), (_QUARTER, "quarter")], + ) + def test_effective_index_matches_full_domain(self, symmetry, label): + # The reduced source must launch the *same* mode, so its effective index has to agree with + # the full-domain solve to numerical precision - not merely to within discretization error. + full = _by_name(_apply((0, 0, 0))[0], "src") + reduced = _by_name(_apply(symmetry)[0], "src") + n_full = complex(full._neff).real + n_reduced = complex(reduced._neff).real + assert abs(n_reduced - n_full) < 1e-5, f"{label}: neff {n_reduced:.6f} vs full {n_full:.6f}" + + @pytest.mark.parametrize("symmetry", [_HALF_Y, _HALF_Z, _QUARTER]) + def test_profile_matches_full_domain_kept_half(self, symmetry): + # Up to the documented normalization (unit flux through the plane the source occupies, so + # sqrt(2**k) larger than the restriction of the full-domain mode) and a global sign, the + # injected profile is the full-domain mode restricted to the kept half. It is not bit-equal: + # the parity projection drops the part of the discrete mode that the wall cannot support (a + # few percent of its norm at this resolution), so compare with a modal-fidelity measure. + # + # The amplitude tolerance is loose because the *discrete* mode is only mirror-symmetric to + # first order in the cell size: the solver samples materials on its staggered grid while + # FDTDX hands it one cell-centred permittivity array, so the solved mode's flux does not split + # exactly evenly between the halves. Measured here for the y axis: 0.446/0.554 at 25 nm and + # 0.473/0.527 at 12.5 nm, so renormalizing to unit flux over the reduced plane inflates the + # amplitude by 5.9% (25 nm) resp. 2.8% (12.5 nm) for that axis. The z axis splits to 0.2%. See + # test_magnetic_plane_projection_uses_the_plain_flip for the sharp check, which uses z only. + full = _by_name(_apply((0, 0, 0))[0], "src") + reduced = _by_name(_apply(symmetry)[0], "src") + multiplicity = 2 ** sum(1 for s in symmetry if s != 0) + + for name in ("_E", "_H"): + expected = _kept_half(np.asarray(getattr(full, name)), symmetry).ravel() + actual = np.asarray(getattr(reduced, name)).ravel() + assert actual.shape == expected.shape + scale = float(np.vdot(expected, actual) / np.vdot(expected, expected)) + assert abs(abs(scale) - np.sqrt(multiplicity)) < 0.15, f"{name}: unexpected amplitude scale {scale:.4f}" + fidelity = abs(np.vdot(expected, actual)) / (np.linalg.norm(expected) * np.linalg.norm(actual)) + assert fidelity > 0.97, f"{name}: only {fidelity:.4f} overlap with the full-domain mode" + + @pytest.mark.parametrize("symmetry", [_HALF_Y, _QUARTER]) + def test_flux_normalization_is_preserved(self, symmetry): + # The convention is unchanged by the fix: a mode source carries the same Poynting flux + # through its own plane as it would in the full domain (unit power through the plane it + # occupies), which is what keeps mode-overlap amplitudes normalized to 1 for a perfectly + # transmitted mode. The reduced fields are therefore sqrt(2**k) larger than the restriction + # of the full-domain mode, not equal to it. + def plane_flux(source): + E = np.asarray(source._E) + H = np.asarray(source._H) + return float(np.sum(E[1] * H[2] - E[2] * H[1])) # S_x over the plane + + full_flux = plane_flux(_by_name(_apply((0, 0, 0))[0], "src")) + reduced_flux = plane_flux(_by_name(_apply(symmetry)[0], "src")) + assert abs(full_flux) > 0 + ratio = reduced_flux / full_flux + assert abs(ratio - 1.0) < 2e-3, f"flux through the reduced plane is {ratio:.6f} of the full-domain one" + + def test_magnetic_plane_projection_uses_the_plain_flip(self): + # The parity projection mirrors the solved mode with the index map of the wall type. A + # magnetic plane sits half a cell below the reduced domain, so that map is the plain flip; + # using the electric m±j map here instead (the same map for both wall types) costs an order + # of magnitude in both metrics below - measured 0.9978 fidelity and a 1.0% amplitude error + # versus 0.9990 and 0.2% - because it projects onto the parity of a plane half a cell away. + full = _by_name(_apply((0, 0, 0))[0], "src") + reduced = _by_name(_apply(_HALF_Z)[0], "src") + for name in ("_E", "_H"): + expected = _kept_half(np.asarray(getattr(full, name)), _HALF_Z).ravel() + actual = np.asarray(getattr(reduced, name)).ravel() + # sqrt(2) larger than the full-domain restriction (unit power through the reduced plane). + scale = abs(np.vdot(expected, actual) / np.vdot(expected, expected)) / np.sqrt(2.0) + assert abs(scale - 1.0) < 5e-3, f"{name}: amplitude is {scale:.4f} of the expected sqrt(2)" + assert _fidelity(expected, actual) > 0.9985, f"{name}: {_fidelity(expected, actual):.6f}" + + def test_explicit_mode_solver_symmetry_is_ignored_but_kept(self): + # An explicit mode-solver symmetry tuple no longer does anything under config.symmetry (a + # warning says so). It must not change the solved mode. + with_field = _by_name(_apply(_QUARTER, mode_symmetry=(0, 1))[0], "src") + without_field = _by_name(_apply(_QUARTER)[0], "src") + assert with_field.symmetry == (0, 1) # kept verbatim, not auto-derived + assert np.allclose(np.asarray(with_field._E), np.asarray(without_field._E), atol=1e-6) + + def test_incompatible_wall_type_raises(self): + # PMC on y makes Ey (normal there) odd, which the fundamental quasi-TE mode is not: the + # parity projection would annihilate it, so this must fail loudly instead of injecting noise. + objects, constraints, config = _build((0, 1, 0)) + with pytest.raises(ValueError, match="symmetry imposed by the walls"): + key = jax.random.PRNGKey(0) + container, arrays, params, config, _ = fdtdx.place_objects( + object_list=objects, config=config, constraints=constraints, key=key + ) + fdtdx.apply_params(arrays, container, params, key=key) + + +def _fidelity(a: np.ndarray, b: np.ndarray) -> float: + """Phase- and scale-invariant modal overlap of two field arrays (1.0 = same mode).""" + a, b = a.ravel(), b.ravel() + return float(abs(np.vdot(a, b)) / (np.linalg.norm(a) * np.linalg.norm(b))) + + +class TestModeOverlapDetectorUnderSymmetry: + def test_reference_mode_matches_source_mode(self): + # Source and detector must solve the same reduced mode, otherwise every overlap (and every + # S-parameter built from it) is off by their mismatch. The detector keeps the complex mode + # (it is paired with recorded phasors) while the source injects its real part, so compare the + # real parts and the effective index. + container, _ = _apply(_QUARTER, detector=True) + source = _by_name(container, "src") + detector = _by_name(container, "det") + source_E = np.asarray(source._E) + reference_E = np.real(np.asarray(detector._mode_E[0])) + assert reference_E.shape == source_E.shape + assert _fidelity(source_E, reference_E) > 0.9999, "reference mode is not the mode being launched" + ratio = np.linalg.norm(reference_E) / np.linalg.norm(source_E) + assert abs(ratio - 1.0) < 1e-3, f"reference mode amplitude differs from the source by {ratio:.6f}" + assert abs(complex(detector._mode_neff[0]).real - complex(source._neff).real) < 1e-6 + + def test_reference_mode_matches_full_domain_kept_half(self): + full = _by_name(_apply((0, 0, 0), detector=True)[0], "det") + reduced = _by_name(_apply(_QUARTER, detector=True)[0], "det") + expected = _kept_half(np.asarray(full._mode_E[0]), _QUARTER) + actual = np.asarray(reduced._mode_E[0]) + assert actual.shape == expected.shape + assert _fidelity(expected, actual) > 0.97, "reduced reference mode is not the full-domain mode" + ratio = np.linalg.norm(actual) / np.linalg.norm(expected) + assert abs(ratio - 2.0) < 0.1, f"unexpected amplitude scale {ratio:.4f}" diff --git a/tests/integration/fdtd/test_symmetry_reduction.py b/tests/integration/fdtd/test_symmetry_reduction.py index 6f46e217..33f511a4 100644 --- a/tests/integration/fdtd/test_symmetry_reduction.py +++ b/tests/integration/fdtd/test_symmetry_reduction.py @@ -12,6 +12,7 @@ import pytest import fdtdx +from fdtdx.objects.detectors.diffractive import DiffractiveDetector def _symmetric_y_edges(n_cells, base): @@ -105,14 +106,16 @@ def test_symmetry_reduces_volume_and_inserts_walls(): assert oc.volume.grid_shape == (20, 10, 10) assert arrays.fields.E.shape == (3, 20, 10, 10) - # PEC wall on y, PMC wall on z, each a single cell on the symmetry plane (min edge). + # PEC wall on the electric plane (y), a single cell on the symmetry plane (min edge). pec = oc.pec_objects - pmc = oc.pmc_objects - assert len(pec) == 1 and len(pmc) == 1 + assert len(pec) == 1 assert pec[0].axis == 1 and pec[0].direction == "-" assert pec[0]._grid_slice_tuple[1] == (0, 1) - assert pmc[0].axis == 2 and pmc[0].direction == "-" - assert pmc[0]._grid_slice_tuple[2] == (0, 1) + + # No wall on the magnetic plane (z). It sits half a cell below the reduced domain, where the zero + # field halo already is the mirror; a PMC object would zero tangential H one cell inside instead, + # over-constraining the field half a cell off the plane. + assert len(oc.pmc_objects) == 0 def test_symmetry_clips_and_drops_objects(): @@ -132,17 +135,21 @@ def test_symmetry_clips_and_drops_objects(): assert _by_name(oc, "dropme") is None -def test_mode_symmetry_auto_derived(): +def test_mode_solver_symmetry_is_not_auto_derived(): objects, constraints, config = _build((0, -1, 1)) key = jax.random.PRNGKey(0) oc, _arrays, _params, _config, _info = fdtdx.place_objects( object_list=objects, config=config, constraints=constraints, key=key ) - # +x propagation -> transverse axes (y, z). y is PEC (-> 0), z is PMC (-> 1). + # The mode-solver symmetry option is no longer used for the config.symmetry path: the mode is + # solved on the mirrored full cross-section and restricted, which reproduces the full-domain + # mode exactly, so the field stays at its default. source = _by_name(oc, "modesrc") assert source is not None - assert source.symmetry == (0, 1) + assert source.symmetry == (0, 0) + # The source spans the whole y/z plane, so it was clipped on both transverse axes. + assert source.symmetry_mirror_axes(exclude_axis=0) == (1, 2) def test_no_symmetry_is_unchanged(): @@ -158,6 +165,72 @@ def test_no_symmetry_is_unchanged(): assert _by_name(oc, "modesrc").symmetry == (0, 0) +def _minimal_symmetric_model(symmetry=(0, -1, 0)): + """Volume + boundaries only, 20 cells per axis, symmetry plane at index 10.""" + config = fdtdx.SimulationConfig( + grid=fdtdx.UniformGrid(spacing=50e-9), time=20e-15, dtype=jnp.float32, symmetry=symmetry + ) + volume = fdtdx.SimulationVolume(partial_grid_shape=(20, 20, 20)) + bound_dict, constraints = fdtdx.boundary_objects_from_config( + fdtdx.BoundaryConfig.from_uniform_bound(thickness=4), volume + ) + return [volume, *bound_dict.values()], list(constraints), config, volume + + +class TestUnsupportedObjectsUnderSymmetry: + """Objects a mirror plane cannot represent must fail at placement with a clear message.""" + + def test_point_dipole_on_symmetry_plane_raises(self): + objects, constraints, config, _volume = _minimal_symmetric_model() + dipole = fdtdx.PointDipoleSource( + name="dip", + partial_grid_shape=(1, 1, 1), + wave_character=fdtdx.WaveCharacter(wavelength=1e-6), + polarization=0, + ) + # First cell of the kept half, i.e. right on the symmetry plane. (A dipole placed by + # place_at_center lands one cell lower, in the discarded half, and is dropped with a warning.) + constraints.append(dipole.set_grid_coordinates(axes=(0, 1, 2), sides=("-", "-", "-"), coordinates=(10, 10, 10))) + objects.append(dipole) + with pytest.raises(ValueError, match="cannot be centred on it"): + fdtdx.place_objects(object_list=objects, config=config, constraints=constraints, key=jax.random.PRNGKey(0)) + + def test_point_dipole_off_symmetry_plane_is_allowed(self): + objects, constraints, config, _volume = _minimal_symmetric_model() + dipole = fdtdx.PointDipoleSource( + name="dip", + partial_grid_shape=(1, 1, 1), + wave_character=fdtdx.WaveCharacter(wavelength=1e-6), + polarization=0, + ) + constraints.append(dipole.set_grid_coordinates(axes=(0, 1, 2), sides=("-", "-", "-"), coordinates=(10, 14, 10))) + objects.append(dipole) + container, _arrays, _params, _config, _info = fdtdx.place_objects( + object_list=objects, config=config, constraints=constraints, key=jax.random.PRNGKey(0) + ) + assert _by_name(container, "dip") is not None + + def test_diffractive_detector_across_symmetry_plane_raises(self): + objects, constraints, config, volume = _minimal_symmetric_model() + detector = DiffractiveDetector( + name="diff", + partial_grid_shape=(None, None, 1), + frequencies=(3e14,), + direction="+", + plot=False, + ) + constraints.extend( + [ + detector.same_size(volume, axes=(0, 1)), + detector.same_position(volume, axes=(0, 1)), + detector.set_grid_coordinates(axes=(2,), sides=("-",), coordinates=(10,)), + ] + ) + objects.append(detector) + with pytest.raises(ValueError, match="diffraction-order basis"): + fdtdx.place_objects(object_list=objects, config=config, constraints=constraints, key=jax.random.PRNGKey(0)) + + def test_odd_cell_count_on_symmetric_axis_raises(): # 21 cells on y (odd) cannot be split exactly down the middle -> hard error. config = fdtdx.SimulationConfig( diff --git a/tests/simulation/physics/boundaries/test_symmetry_structured_field.py b/tests/simulation/physics/boundaries/test_symmetry_structured_field.py new file mode 100644 index 00000000..a6754e7d --- /dev/null +++ b/tests/simulation/physics/boundaries/test_symmetry_structured_field.py @@ -0,0 +1,208 @@ +"""Physics test: ``config.symmetry`` reproduces the full-domain run for a *structured* field. + +``test_simulation_symmetry.py`` compares a reduced run against the full domain for a uniform plane +wave. A transversely uniform field has no transverse derivatives, which makes it blind to how the +mirror walls and the detector co-location stencil treat the cells at the symmetry plane. This test +uses a narrow source instead, so the field diffracts and every transverse derivative is exercised. + +Both mirror types are then exact, each because of where its plane sits: + +* an **electric** plane sits on the reduced domain's min edge, where the tangential ``E`` samples + live and the PEC wall zeroes them; +* a **magnetic** plane sits half a cell *below* the min edge — the source footprint is rasterized per + cell, so the discrete problem is mirror symmetric about the tangential-``H`` node one cell out, + where the zero field halo already is the mirror. It gets no wall object at all; zeroing tangential + ``H`` at the first cell (or filling that halo with its mirror) displaces the plane by half a cell, + which is a clean first-order-wrong answer: 4.3e-02 raw-field error at 50 nm, halving with the cell + size instead of vanishing. + +Measured here (50 nm / 25 nm), max deviation over the kept half relative to the full-domain peak: +raw fields ≤ 2.8e-04 / 2.1e-06, detector output ≤ 9.3e-06 / 5.6e-07. The raw fields bypass the +co-location stencil, so they isolate the walls; the detector output additionally covers the halo the +stencil reads, which is what recorded *half* the field in the plane row of an electric plane +(``5.0e-01``) before it existed. + +Unfolding is looser (~1.2e-02 / 1.7e-02 at 50 nm, first order in the cell size) and cannot be +tightened by either mirror: the co-location average is not symmetric about either candidate plane, so +reconstructing the discarded half from *co-located* samples keeps an O(cell size) residue. The raw +fields it is built from do not. + +The source is a ``UniformPlaneSource`` covering only part of the plane: uniform over its own +footprint, so this stays sensitive to the walls alone and not to a source-profile centering bug +(``tests/unit/objects/sources/test_symmetry_sources.py`` covers those). +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import fdtdx + +_SPACING = 50e-9 +_WAVELENGTH = 1e-6 # 20 cells per wavelength +_TRANSVERSE_CELLS = 60 +_PROPAGATION_CELLS = 40 +_SOURCE_CELLS = 20 # even and centered -> the footprint is a mirror-symmetric cell set +_PML_CELLS = 8 +_SOURCE_Z = 12 +_DETECTOR_Z = 26 +_SIM_TIME = 60e-15 + +# E is polarized along x and propagates in +z, so the mirror that keeps it even is PEC on x (Ex is +# normal there) and PMC on y (Ex is tangential there). +_HALF_X_PEC = (-1, 0, 0) +_HALF_Y_PMC = (0, 1, 0) +_QUARTER = (-1, 1, 0) + + +def _build(symmetry): + config = fdtdx.SimulationConfig( + grid=fdtdx.UniformGrid(spacing=_SPACING), + time=_SIM_TIME, + dtype=jnp.float32, + symmetry=symmetry, + ) + volume = fdtdx.SimulationVolume(partial_grid_shape=(_TRANSVERSE_CELLS, _TRANSVERSE_CELLS, _PROPAGATION_CELLS)) + objects, constraints = [volume], [] + bound_dict, boundary_constraints = fdtdx.boundary_objects_from_config( + fdtdx.BoundaryConfig.from_uniform_bound(thickness=_PML_CELLS), volume + ) + constraints.extend(boundary_constraints) + objects.extend(bound_dict.values()) + + wave = fdtdx.WaveCharacter(wavelength=_WAVELENGTH) + source = fdtdx.UniformPlaneSource( + name="src", + partial_grid_shape=(_SOURCE_CELLS, _SOURCE_CELLS, 1), + wave_character=wave, + direction="+", + fixed_E_polarization_vector=(1, 0, 0), + normalize_by_energy=False, + ) + constraints.extend( + [ + source.place_at_center(volume, axes=(0, 1)), + source.set_grid_coordinates(axes=(2,), sides=("-",), coordinates=(_SOURCE_Z,)), + ] + ) + objects.append(source) + + detector = fdtdx.PhasorDetector( + name="det", + partial_grid_shape=(None, None, 1), + wave_characters=(wave,), + components=("Ex", "Hy"), + reduce_volume=False, + plot=False, + ) + constraints.extend( + [ + detector.same_size(volume, axes=(0, 1)), + detector.same_position(volume, axes=(0, 1)), + detector.set_grid_coordinates(axes=(2,), sides=("-",), coordinates=(_DETECTOR_Z,)), + ] + ) + objects.append(detector) + return objects, constraints, config + + +def _run(symmetry): + objects, constraints, config = _build(symmetry) + key = jax.random.PRNGKey(0) + container, arrays, params, config, _ = fdtdx.place_objects( + object_list=objects, config=config, constraints=constraints, key=key + ) + arrays, container, _ = fdtdx.apply_params(arrays, container, params, key) + _, arrays = fdtdx.run_fdtd(arrays=arrays, objects=container, config=config, key=key) + return container, arrays, config + + +def _phasor(arrays): + # (num_components, nx, ny) + return np.asarray(arrays.detector_states["det"]["phasor"][0, 0, :, :, :, 0]) + + +def _kept_half(array, symmetry, first_spatial_axis=1): + index: list[slice] = [slice(None)] * array.ndim + for axis in (0, 1): + if symmetry[axis] != 0: + ax = first_spatial_axis + axis + index[ax] = slice(array.shape[ax] // 2, None) + return array[tuple(index)] + + +@pytest.fixture(scope="module") +def full_run(): + _container, arrays, _config = _run((0, 0, 0)) + return { + "phasor": _phasor(arrays), + "E": np.asarray(arrays.fields.E), + "H": np.asarray(arrays.fields.H), + } + + +@pytest.mark.parametrize("symmetry", [_HALF_X_PEC, _HALF_Y_PMC, _QUARTER]) +def test_reduced_raw_fields_match_full_domain_kept_half(full_run, symmetry): + """The reduced run's own field arrays must agree with the full run, plane row included. + + Reads ``arrays.fields`` directly, so this sees the mirror walls alone — no detector, no + co-location stencil. A wall condition displaced by half a cell shows up here as ~4e-02. + """ + _container, arrays, _config = _run(symmetry) + for name in ("E", "H"): + reduced = np.asarray(getattr(arrays.fields, name)) + expected = _kept_half(full_run[name], symmetry) + assert reduced.shape == expected.shape + scale = np.abs(expected).max() + assert scale > 1e-20, f"reference {name} is zero - wave not launched" + error = np.abs(reduced - expected).max() / scale + assert error < 1e-3, f"{name}: reduced run differs from the full domain by {error:.3e}" + + +@pytest.mark.parametrize("symmetry", [_HALF_X_PEC, _HALF_Y_PMC, _QUARTER]) +def test_reduced_detector_matches_full_domain_kept_half(full_run, symmetry): + """Same comparison through a detector, which adds the co-location halo at the plane. + + Without the mirror halo an electric plane records exactly half the field in its plane row, i.e. + 5e-01 here; a magnetic plane needs no halo (its plane is half a cell out, where the zero halo + already is the mirror), and filling one anyway costs ~4e-02. + """ + _container, arrays, _config = _run(symmetry) + reduced = _phasor(arrays) + expected = _kept_half(full_run["phasor"], symmetry) + assert reduced.shape == expected.shape + assert np.abs(expected).max() > 1e-20, "reference field is zero - wave not launched" + + for index, name in enumerate(("Ex", "Hy")): + scale = np.abs(expected[index]).max() + error = np.abs(reduced[index] - expected[index]).max() / scale + assert error < 1e-4, f"{name}: reduced run differs from the full domain by {error:.3e}" + + +@pytest.mark.parametrize("symmetry, tolerance", [(_HALF_X_PEC, 0.02), (_QUARTER, 0.03)]) +def test_unfolded_field_matches_full_domain(full_run, symmetry, tolerance): + """Unfolding the reduced detector output reconstructs the full-domain plane. + + Exercises the per-axis mirror index map. Across the **electric** x-plane the co-located samples + sit on the plane, so the plane row must not be duplicated and the mirrored half must not be + shifted by a cell (a plain flip leaves ~8e-02 here). Across the **magnetic** y-plane the plane is + half a cell out and the same samples mirror one-to-one, so the map is the plain flip instead (the + on-plane map leaves ~9e-02). Applying either convention to both axes fails one of them. + + The remaining ~1e-02 is the co-location average itself, which is symmetric about neither + candidate plane; it is first order in the cell size and does not affect the recorded (non- + unfolded) output, which the tests above hold to 1e-4. + """ + full_phasor = full_run["phasor"] + container, arrays, config = _run(symmetry) + unfolded = fdtdx.unfold_detector_states(arrays, container, config) + reconstructed = np.asarray(unfolded.detector_states["det"]["phasor"][0, 0, :, :, :, 0]) + assert reconstructed.shape == full_phasor.shape + + for index, name in enumerate(("Ex", "Hy")): + scale = np.abs(full_phasor[index]).max() + # The outermost reconstructed cell repeats its neighbour (its mirror partner lies outside the + # kept half); it sits inside the PML, so compare the interior. + error = np.abs(reconstructed[index][1:, 1:] - full_phasor[index][1:, 1:]).max() / scale + assert error < tolerance, f"{name}: unfolded field differs from the full domain by {error:.3e}" diff --git a/tests/unit/core/physics/test_modes.py b/tests/unit/core/physics/test_modes.py index 82265e57..c67b585d 100644 --- a/tests/unit/core/physics/test_modes.py +++ b/tests/unit/core/physics/test_modes.py @@ -1,5 +1,6 @@ from unittest.mock import patch +import jax import jax.numpy as jnp import numpy as np import pytest @@ -8,6 +9,7 @@ ModeTupleType, compute_mode, compute_mode_polarization_fraction, + compute_mode_symmetry_reduced, sort_modes, tidy3d_mode_computation_wrapper, ) @@ -946,3 +948,72 @@ def test_plane_center_is_midpoint_of_transverse_coords(self, mock_normalize, moc kwargs = mock_wrapper.call_args.kwargs expected = (0.5 * 5 * resolution / 1e-6, 0.5 * 6 * resolution / 1e-6) assert kwargs["plane_center"] == pytest.approx(expected) + + +class TestComputeModeSymmetryReduced: + """The symmetry-reduced route: mirror the cross-section, solve, project, restrict.""" + + def _make_mock_mode(self, shape): + return ModeTupleType( + neff=1.5 + 0.1j, + Ex=np.ones(shape, dtype=np.complex64), + Ey=np.ones(shape, dtype=np.complex64), + Ez=np.ones(shape, dtype=np.complex64), + Hx=np.ones(shape, dtype=np.complex64), + Hy=np.ones(shape, dtype=np.complex64), + Hz=np.ones(shape, dtype=np.complex64), + ) + + def _kwargs(self, **overrides): + # x-propagation (singleton at dim 1), 4 x 3 transverse cells. + kwargs = dict( + mirrored_axes=(2,), + walls={2: 1}, + frequency=2e14, + inv_permittivities=jnp.ones((1, 1, 4, 3)), + inv_permeabilities=1.0, + resolution=1e-8, + object_name="modesrc", + ) + kwargs.update(overrides) + return kwargs + + def test_bend_about_a_mirrored_axis_raises(self): + # The conformal bend transform scales the index linearly across bend_axis, so the mirrored + # cross-section is not symmetric about a plane normal to it and the reduced run cannot + # represent the mode. Rejected before any solve. + with pytest.raises(ValueError, match="bends about the z-axis"): + compute_mode_symmetry_reduced(**self._kwargs(bend_radius=5e-6, bend_axis=2)) + + @patch("fdtdx.core.physics.modes.tidy3d_mode_computation_wrapper") + @patch("fdtdx.core.physics.modes.normalize_by_poynting_flux") + def test_bend_about_the_other_transverse_axis_is_allowed(self, mock_normalize, mock_wrapper): + # A bend leaves the axis it does not bend about mirror-symmetric, so mirroring that one is + # consistent: the pair (mirror z, bend y) is fine, unlike (mirror z, bend z) above. + mock_wrapper.return_value = [self._make_mock_mode((4, 6))] + mock_normalize.side_effect = lambda E, H, axis, area_weights=None: (E, H) + + mode_E, mode_H, _neff = compute_mode_symmetry_reduced(**self._kwargs(bend_radius=5e-6, bend_axis=1)) + + assert mode_E.shape == (3, 1, 4, 3) # solved on the mirrored plane, restricted to the kept half + assert mode_H.shape == (3, 1, 4, 3) + + @patch("fdtdx.core.physics.modes.tidy3d_mode_computation_wrapper") + @patch("fdtdx.core.physics.modes.normalize_by_poynting_flux") + def test_survives_jit(self, mock_normalize, mock_wrapper): + # A mode source or mode-overlap detector overlapping a Device solves its mode inside + # apply_params, which callers trace. The parity residual must therefore not be concretized: + # its diagnostics are skipped under tracing instead. + mock_wrapper.return_value = [self._make_mock_mode((4, 6))] + mock_normalize.side_effect = lambda E, H, axis, area_weights=None: (E, H) + + def traced(inv_permittivities): + mode_E, _mode_H, _neff = compute_mode_symmetry_reduced( + **self._kwargs(inv_permittivities=inv_permittivities) + ) + return mode_E + + eager = traced(jnp.ones((1, 1, 4, 3))) + jitted = jax.jit(traced)(jnp.ones((1, 1, 4, 3))) + assert jitted.shape == eager.shape + assert jnp.allclose(jitted, eager) diff --git a/tests/unit/core/physics/test_symmetry.py b/tests/unit/core/physics/test_symmetry.py new file mode 100644 index 00000000..824c1fa5 --- /dev/null +++ b/tests/unit/core/physics/test_symmetry.py @@ -0,0 +1,194 @@ +"""Unit tests for the low-level mirror-symmetry primitives.""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from fdtdx.core.physics.symmetry import ( + component_sits_on_plane, + field_component_parity, + mirror_about_interior_plane, + mirror_edge_coordinates, + mirror_extend_low_side, + mirror_material_array, + mirror_material_cross_section, + mirror_pairs_on_plane, + project_onto_parity, + restrict_to_kept_half, +) + + +class TestComponentSitsOnPlane: + """Which Yee samples lie on a mirror plane, i.e. which mirror index map applies.""" + + def test_electric_components(self): + # E_c is offset half a cell along c only, so along any other axis it sits at integer + # positions - on the plane. + for axis in range(3): + for component in range(3): + assert component_sits_on_plane("E", component, axis) == (component != axis) + + def test_magnetic_components(self): + # H_c is offset along the two axes other than c, so the reverse holds. + for axis in range(3): + for component in range(3): + assert component_sits_on_plane("H", component, axis) == (component == axis) + + def test_components_that_vanish_are_exactly_the_on_plane_odd_ones(self): + # Consistency of the two tables: whatever a wall drives to zero must be sampled on the plane, + # otherwise the condition cannot be enforced pointwise (this is why a PMC wall cannot be + # implemented by zeroing tangential H at the first cell). + for wall, field_type, expect_on_plane in ((-1, "E", True), (1, "H", False)): + for axis in range(3): + for component in range(3): + if field_component_parity(field_type, component, axis, wall) == -1: + assert component_sits_on_plane(field_type, component, axis) == expect_on_plane + + def test_invalid_field_type(self): + with pytest.raises(ValueError): + component_sits_on_plane("B", 0, 0) # type: ignore[arg-type] + + +class TestMirrorPairsOnPlane: + """The mirror index map depends on the wall type, not only on the Yee offsets.""" + + def test_electric_plane_uses_the_yee_offsets(self): + for field_type in ("E", "H"): + for axis in range(3): + for component in range(3): + assert mirror_pairs_on_plane(field_type, component, axis, -1) == component_sits_on_plane( + field_type, component, axis + ) + + def test_magnetic_plane_is_always_a_plain_flip(self): + # A magnetic plane sits half a cell below the reduced domain, so nothing is sampled on it. + for field_type in ("E", "H"): + for axis in range(3): + for component in range(3): + assert mirror_pairs_on_plane(field_type, component, axis, 1) is False + + +class TestMirrorMaterialArray: + def test_isotropic_and_diagonal_only_flip(self): + for num_components in (1, 3): + arr = jnp.asarray(np.arange(num_components * 4).reshape(num_components, 4, 1, 1), dtype=jnp.float32) + out = mirror_material_array(arr, axis=0) + assert jnp.allclose(out, jnp.flip(arr, axis=1)) + + def test_full_tensor_flips_sign_of_mixed_off_diagonals(self): + # eps_xy and eps_yx change sign under a mirror normal to x; eps_yz does not. + arr = jnp.ones((9, 2, 1, 1), dtype=jnp.float32) + out = mirror_material_array(arr, axis=0) + expected_signs = [1, -1, -1, -1, 1, 1, -1, 1, 1] # row-major (xx, xy, xz, yx, ...) + for index, sign in enumerate(expected_signs): + assert jnp.allclose(out[index], sign * jnp.ones((2, 1, 1))) + + def test_cross_section_doubles_and_restrict_inverts(self): + arr = jnp.asarray(np.arange(6).reshape(1, 1, 6, 1), dtype=jnp.float32) + full = mirror_material_cross_section(arr, axes=(1,)) + assert full.shape == (1, 1, 12, 1) + assert jnp.allclose(full[:, :, 6:], arr) + assert jnp.allclose(full[:, :, :6], jnp.flip(arr, axis=2)) + assert jnp.allclose(restrict_to_kept_half(full, axes=(1,)), arr) + + +class TestMirrorEdgeCoordinates: + def test_widths_are_mirrored_and_kept_half_preserved(self): + edges = jnp.asarray([0.0, 1.0, 3.0, 6.0]) # widths 1, 2, 3 + full = mirror_edge_coordinates(edges) + assert full.shape == (7,) + widths = np.diff(np.asarray(full)) + assert np.allclose(widths, [3.0, 2.0, 1.0, 1.0, 2.0, 3.0]) + # The kept half keeps its original coordinates, and the plane is at the original first edge. + assert np.allclose(np.asarray(full)[3:], np.asarray(edges)) + + +class TestMirrorIndexMaps: + def test_interior_plane_off_plane_is_plain_flip(self): + arr = jnp.asarray([[0.0, 1.0, 2.0, 3.0]]).reshape(1, 4, 1, 1) + out = mirror_about_interior_plane(arr, axis=1, on_plane=False) + assert np.allclose(np.asarray(out).ravel(), [3.0, 2.0, 1.0, 0.0]) + + def test_interior_plane_on_plane_pairs_around_the_plane_row(self): + # 4 samples, plane between index 1 and 2 -> index 2 is on the plane and its own mirror, + # index 1 <-> 3, index 0 has no partner and is left alone. + arr = jnp.asarray([0.0, 1.0, 2.0, 3.0]).reshape(1, 4, 1, 1) + out = np.asarray(mirror_about_interior_plane(arr, axis=1, on_plane=True)).ravel() + assert out[2] == 2.0 # self-mirrored + assert out[1] == 3.0 and out[3] == 1.0 # swapped + assert out[0] == 0.0 # unpaired, untouched + + def test_low_side_off_plane_mirrors_every_sample(self): + arr = jnp.asarray([10.0, 20.0, 30.0]).reshape(1, 3, 1, 1) + low = np.asarray(mirror_extend_low_side(arr, axis=1, parity=-1, on_plane=False)).ravel() + assert np.allclose(low, [-30.0, -20.0, -10.0]) + + def test_low_side_on_plane_does_not_duplicate_the_plane_row(self): + # Index 0 is on the plane: only 1 and 2 have mirror images, and the outermost reconstructed + # sample repeats its neighbour because its partner lies outside the kept half. + arr = jnp.asarray([10.0, 20.0, 30.0]).reshape(1, 3, 1, 1) + low = np.asarray(mirror_extend_low_side(arr, axis=1, parity=1, on_plane=True)).ravel() + assert np.allclose(low, [30.0, 30.0, 20.0]) + + +class TestProjectOntoParity: + def _mode(self, values): + return jnp.asarray(values, dtype=jnp.float32).reshape(3, 1, 4, 1) + + def test_symmetric_field_is_untouched_across_a_magnetic_plane(self): + # Ex is tangential to a y-wall and even under PMC. A magnetic plane sits half a cell below the + # reduced domain, so its mirror map is the plain flip and a compatible profile is the one that + # is symmetric about the array centre: f[0] == f[3], f[1] == f[2]. + field = self._mode([[7.0, 9.0, 9.0, 7.0], [0.0] * 4, [0.0] * 4]) + projected, residual = project_onto_parity(field, "E", {1: 1}) + assert residual < 1e-6 + assert jnp.allclose(projected[0, 0, :, 0], field[0, 0, :, 0]) + + def test_symmetric_field_is_untouched_across_an_electric_plane(self): + # Ez is tangential to a y-wall and odd under PEC, and it is sampled *on* an electric plane + # (index 2 of 4), so a compatible profile vanishes there and satisfies f[1] == -f[3]. + field = self._mode([[0.0] * 4, [0.0] * 4, [0.0, -7.0, 0.0, 7.0]]) + projected, residual = project_onto_parity(field, "E", {1: -1}) + assert residual < 1e-6 + assert jnp.allclose(projected[2, 0, 1:, 0], field[2, 0, 1:, 0]) + + def test_magnetic_plane_map_differs_from_the_electric_one(self): + # The same profile cannot be compatible with both: across an electric plane the pairing is + # m±j about the plane row, across a magnetic plane it is the plain flip. + field = self._mode([[0.0, 7.0, 9.0, 7.0], [0.0] * 4, [0.0] * 4]) + _projected, residual = project_onto_parity(field, "E", {1: 1}) + assert residual > 0.1, "an m±j-symmetric profile is not plain-flip symmetric" + + def test_incompatible_field_is_annihilated(self): + # The same profile against a PEC wall, where Ex is odd: nothing survives except the unpaired + # outermost sample, so the residual is close to one. + field = self._mode([[0.0, 7.0, 9.0, 7.0], [0.0] * 4, [0.0] * 4]) + projected, residual = project_onto_parity(field, "E", {1: -1}) + assert residual > 0.9 + assert jnp.allclose(projected[0, 0, 2, 0], 0.0) # on-plane odd sample must vanish + + def test_odd_on_plane_component_is_zeroed_on_the_plane(self): + field = self._mode([[1.0, 2.0, 3.0, 4.0], [0.0] * 4, [0.0] * 4]) + projected, _ = project_onto_parity(field, "E", {1: -1}) + assert jnp.allclose(projected[0, 0, 2, 0], 0.0) + + def test_zero_field_has_zero_residual(self): + field = jnp.zeros((3, 1, 4, 1), dtype=jnp.float32) + _projected, residual = project_onto_parity(field, "H", {1: 1}) + assert residual == 0.0 + + def test_residual_is_a_jax_scalar_and_survives_tracing(self): + # A mode source or mode-overlap detector that overlaps a Device solves its mode inside + # apply_params, which callers trace (jax.jit around an optimization step). Concretizing the + # residual there would raise ConcretizationTypeError, so it must stay a JAX scalar. + field = self._mode([[1.0, 2.0, 3.0, 4.0], [0.0] * 4, [0.0] * 4]) + projected, residual = project_onto_parity(field, "E", {1: -1}) + assert isinstance(residual, jax.Array) and residual.shape == () + + def traced(f): + return project_onto_parity(f, "E", {1: -1}) + + jit_projected, jit_residual = jax.jit(traced)(field) + assert jnp.allclose(jit_projected, projected) + assert float(jit_residual) == pytest.approx(float(residual), abs=1e-6) diff --git a/tests/unit/fdtd/test_symmetry_unfold.py b/tests/unit/fdtd/test_symmetry_unfold.py index 51176e82..abbdea7a 100644 --- a/tests/unit/fdtd/test_symmetry_unfold.py +++ b/tests/unit/fdtd/test_symmetry_unfold.py @@ -71,11 +71,14 @@ def test_pec_y_even_and_odd_components(self): assert full.shape == (3, 2, 8, 2) assert jnp.allclose(full[:, :, 4:, :], half) - # Ey is normal to the PEC y-plane -> even about the center. + # Ey is normal to the y-plane, so it is sampled half a cell off it and mirrors one-to-one + # about the array center: even under PEC. assert jnp.allclose(full[1], jnp.flip(full[1], axis=1)) - # Ex, Ez are tangential -> odd about the center. - assert jnp.allclose(full[0], -jnp.flip(full[0], axis=1)) - assert jnp.allclose(full[2], -jnp.flip(full[2], axis=1)) + # Ex and Ez are tangential: sampled *on* the plane (index 4 after unfolding), so index 4 is + # its own mirror and the pairs are 4-j <-> 4+j. Both are odd under PEC. + for component in (0, 2): + for offset in (1, 2, 3): + assert jnp.allclose(full[component, :, 4 - offset, :], -full[component, :, 4 + offset, :]) def test_pmc_z_parities_for_H(self): half = self._half(seed=1) @@ -83,11 +86,37 @@ def test_pmc_z_parities_for_H(self): full = unfold_fields(half, symmetry, "H") assert full.shape == (3, 2, 4, 4) - # After indexing a component, spatial axes are (x=0, y=1, z=2); mirror about z is axis 2. - # Hz normal to PMC z-plane -> even; Hx, Hy tangential -> odd. + # A magnetic plane sits half a cell *below* the reduced domain (sources and materials are + # rasterized per cell), so across it every component mirrors one-to-one - the plain flip - + # whatever its Yee offset. After indexing a component, spatial axes are (x=0, y=1, z=2). + # Hz is normal to the plane and even under PMC. assert jnp.allclose(full[2], jnp.flip(full[2], axis=2)) + # Hx, Hy are tangential and odd under PMC. assert jnp.allclose(full[0], -jnp.flip(full[0], axis=2)) assert jnp.allclose(full[1], -jnp.flip(full[1], axis=2)) + # Plain flip means the plane row *is* duplicated across the plane: index 2 is the first kept + # cell and index 1 its mirror (the m±j map would pair 1 with 3 instead). + assert jnp.allclose(full[2, :, :, 1], full[2, :, :, 2]) + + def test_wall_type_selects_the_index_map(self): + # The same reduced field unfolds differently across an electric and a magnetic plane, for a + # component sampled on the electric plane (Ex is tangential to y). Both are even here, so only + # the index map distinguishes them. + half = self._half(nx=1, nh=4, nz=1) + pmc = unfold_fields(half, (0, 1, 0), "E") # plain flip: kept cell 0 mirrors into row 3 + pec = unfold_fields(half, (0, -1, 0), "E") # m±j: row 4 is its own mirror, cell 1 -> row 3 + assert jnp.allclose(pmc[0, :, 3, :], half[0, :, 0, :]) + assert jnp.allclose(pec[0, :, 3, :], -half[0, :, 1, :]) + assert not jnp.allclose(pmc[0, :, 3, :], pec[0, :, 3, :]) + + def test_on_plane_components_are_not_duplicated(self): + # The sample on the plane is its own mirror, so unfolding must not repeat it: the row next to + # the plane belongs to the *second* kept cell, not the first. + half = self._half(nx=1, nh=4, nz=1) + full = unfold_fields(half, (0, -1, 0), "E") + # Ex is tangential to the y-plane (odd, sampled on it). + assert jnp.allclose(full[0, :, 4, :], half[0, :, 0, :]) # plane row passed through + assert jnp.allclose(full[0, :, 3, :], -half[0, :, 1, :]) # neighbour mirrors the *next* cell def test_two_axis_shape(self): half = self._half(nx=3, nh=4, nz=5) @@ -195,8 +224,27 @@ def test_energy_as_slices_mirrors_in_plane_axes(self): assert out["XY Plane"].shape == (1, 3, 8) assert out["YZ Plane"].shape == (1, 8, 5) assert out["XZ Plane"].shape == (1, 3, 5) - # Energy is even -> the doubled axis is mirror-symmetric. - assert jnp.allclose(out["XY Plane"], jnp.flip(out["XY Plane"], axis=2)) + # Energy is even, and the co-located samples sit on the plane along y (array index 4 after + # unfolding), so the mirror pairs are 4-j <-> 4+j. + for offset in (1, 2, 3): + assert jnp.allclose(out["XY Plane"][:, :, 4 - offset], out["XY Plane"][:, :, 4 + offset]) + + def test_colocated_map_is_per_wall_type(self): + # Same detector, same data: across an electric plane the co-located samples sit *on* it (they + # are integer along x and y), so the plane row is not duplicated; across a magnetic plane they + # mirror one-to-one and it is. + det = fdtdx.EnergyDetector(name="en2", as_slices=True, plot=False) + rng = np.random.default_rng(11) + plane = jnp.asarray(rng.standard_normal((1, 3, 4)), dtype=jnp.float32) + state = { + "XY Plane": plane, + "XZ Plane": jnp.zeros((1, 3, 5), dtype=jnp.float32), + "YZ Plane": jnp.zeros((1, 4, 5), dtype=jnp.float32), + } + pec = _unfold_one_detector(det, state, (0, -1, 0), count=1)["XY Plane"] + pmc = _unfold_one_detector(det, state, (0, 1, 0), count=1)["XY Plane"] + assert jnp.allclose(pec[:, :, 3], plane[:, :, 1]) # m±j: row 4 is the plane row + assert jnp.allclose(pmc[:, :, 3], plane[:, :, 0]) # plain flip: the plane row is duplicated def test_diffractive_raises(self): from fdtdx.objects.detectors.diffractive import DiffractiveDetector diff --git a/tests/unit/fdtd/test_update.py b/tests/unit/fdtd/test_update.py index 5c01233a..69f40043 100644 --- a/tests/unit/fdtd/test_update.py +++ b/tests/unit/fdtd/test_update.py @@ -14,6 +14,7 @@ collect_interfaces, get_wrap_padding_axes, pad_fields_for_boundaries, + pad_fields_with_symmetry_mirror, update_detector_states, update_E, update_E_reverse, @@ -68,11 +69,12 @@ def _make_objects(sources=None): return obj -def _make_config(c=0.5): +def _make_config(c=0.5, symmetry=(0, 0, 0)): cfg = Mock() cfg.courant_number = c cfg.uniform_spacing.return_value = 1.0 cfg.has_nonuniform_grid = False + cfg.symmetry = symmetry return cfg @@ -166,6 +168,7 @@ def test_nonuniform_grid_passes_compatibility_spacing_to_boundaries(self): config.grid = grid config.has_nonuniform_grid = True config.resolved_grid = grid + config.symmetry = (0, 0, 0) config.uniform_spacing.side_effect = AssertionError("uniform spacing should not be required") boundary = Mock() @@ -184,6 +187,85 @@ def test_nonuniform_grid_passes_compatibility_spacing_to_boundaries(self): boundary.apply_pad_correction.assert_called_once() assert boundary.apply_pad_correction.call_args.args[2] == pytest.approx(1.0) + def test_symmetric_axis_does_not_wrap_its_min_side_halo(self): + # Reduction drops the min-side periodic boundary of a symmetric axis but keeps the far-side + # one, so the axis still reports wrap padding. Wrapping the min-side halo would fill it with + # the far side's field, overwriting what the mirror needs there; the far side is untouched. + boundary = Mock() + boundary.axis = 1 + boundary.uses_wrap_padding = True + boundary.apply_pad_correction.side_effect = lambda padded, _shape, _spacing: padded + objects = Mock() + objects.boundary_objects = [boundary] + objects.volume.grid_shape = (NX, NY, NZ) + + fields = jnp.arange(1, 3 * NX * NY * NZ + 1, dtype=jnp.float32).reshape(FIELD_SHAPE) + wrapped = pad_fields_for_boundaries(fields, objects, _make_config()) + symmetric = pad_fields_for_boundaries(fields, objects, _make_config(symmetry=(0, 1, 0))) + + assert jnp.allclose(wrapped[:, 1:-1, 0, 1:-1], fields[:, :, -1, :]) # wraps without symmetry + assert jnp.all(symmetric[:, :, 0, :] == 0.0) # min side belongs to the mirror + assert jnp.allclose(symmetric[:, 1:-1, -1, 1:-1], fields[:, :, 0, :]) # far side still wraps + + +class TestPadFieldsWithSymmetryMirror: + """The min-side halo the detector co-location stencil reads at a symmetry plane. + + Only an *electric* plane gets the mirror: it sits on the reduced domain's min edge, so the halo's + true value is the parity-weighted mirror partner (the first cell for a component sampled half a + cell off the plane, the second for one sampled on it). A magnetic plane sits half a cell below the + domain, where the zero halo already *is* the mirror, and a user-placed boundary makes no symmetry + claim at all — both keep the plain zero padding. + """ + + def _objects(self, axis=1, is_symmetry_wall=True): + boundary = Mock() + boundary.axis = axis + boundary.uses_wrap_padding = False + boundary._is_symmetry_wall = is_symmetry_wall + boundary.apply_pad_correction.side_effect = lambda padded, _shape, _spacing: padded + objects = Mock() + objects.boundary_objects = [boundary] + objects.volume.grid_shape = (NX, NY, NZ) + return objects + + def _fields(self): + # Distinct nonzero values everywhere, so a wrong source index cannot pass by coincidence. + return jnp.arange(1, 3 * NX * NY * NZ + 1, dtype=jnp.float32).reshape(FIELD_SHAPE) + + def test_electric_plane_mirrors_each_component_with_its_own_index_map(self): + fields = self._fields() + config = _make_config(symmetry=(0, -1, 0)) + + # E across a PEC y-plane: Ex, Ez are tangential (odd) and sampled *on* the plane, so their + # halo is minus the *second* cell; Ey is normal (even) and half a cell off, so it is plus the + # first cell. + padded = pad_fields_with_symmetry_mirror(fields, self._objects(), config, "E") + assert jnp.allclose(padded[0, 1:-1, 0, 1:-1], -fields[0, :, 1, :]) + assert jnp.allclose(padded[1, 1:-1, 0, 1:-1], fields[1, :, 0, :]) + assert jnp.allclose(padded[2, 1:-1, 0, 1:-1], -fields[2, :, 1, :]) + + # H is the other way round: Hy is normal (odd) and sampled on the plane, Hx and Hz are + # tangential (even) and half a cell off. + padded = pad_fields_with_symmetry_mirror(fields, self._objects(), config, "H") + assert jnp.allclose(padded[0, 1:-1, 0, 1:-1], fields[0, :, 0, :]) + assert jnp.allclose(padded[1, 1:-1, 0, 1:-1], -fields[1, :, 1, :]) + assert jnp.allclose(padded[2, 1:-1, 0, 1:-1], fields[2, :, 0, :]) + + def test_magnetic_plane_keeps_the_zero_halo(self): + fields = self._fields() + config = _make_config(symmetry=(0, 1, 0)) + for field_type in ("E", "H"): + padded = pad_fields_with_symmetry_mirror(fields, self._objects(), config, field_type) + assert jnp.all(padded[:, :, 0, :] == 0.0), field_type + + def test_user_placed_boundary_keeps_the_zero_halo(self): + fields = self._fields() + config = _make_config(symmetry=(0, -1, 0)) + objects = self._objects(is_symmetry_wall=False) + padded = pad_fields_with_symmetry_mirror(fields, objects, config, "E") + assert jnp.all(padded[:, :, 0, :] == 0.0) + # ─── TestUpdateE ────────────────────────────────────────────────────────────── diff --git a/tests/unit/objects/sources/test_symmetry_sources.py b/tests/unit/objects/sources/test_symmetry_sources.py new file mode 100644 index 00000000..fa69e454 --- /dev/null +++ b/tests/unit/objects/sources/test_symmetry_sources.py @@ -0,0 +1,315 @@ +"""Plane sources under ``config.symmetry``. + +A symmetry-reduced simulation must inject exactly the part of the full-domain source that falls +inside the kept half/quadrant — same profile, same absolute amplitude. These tests compare the +injected fields of a reduced run against the same model run without symmetry, which needs no FDTD +time stepping: ``place_objects`` + ``apply_params`` already populate ``source._E`` / ``source._H``. + +Covers both wall types (PEC and PMC), one and two symmetry axes, all three propagation axes, and +both propagation directions. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from loguru import logger + +import fdtdx +from fdtdx.core.jax.pytrees import autoinit +from fdtdx.objects.sources.linear_polarization import LinearlyPolarizedPlaneSource + +_SPACING = 100e-9 +_WAVELENGTH = 1e-6 +_TRANSVERSE_CELLS = 20 # even, so every symmetric axis splits exactly +_PROPAGATION_CELLS = 12 + + +def _axes(propagation_axis: int) -> tuple[int, int]: + """(horizontal, vertical) transverse axes, in fdtdx's right-hand cyclic order.""" + return (propagation_axis + 1) % 3, (propagation_axis + 2) % 3 + + +def _compatible_symmetry(propagation_axis: int, axes: str) -> tuple[int, int, int]: + """Symmetry tuple whose walls match an E-along-horizontal plane wave. + + E is normal to the horizontal plane (even under PEC) and tangential to the vertical plane + (even under PMC), so a mirror-symmetric beam needs PEC on the horizontal axis and PMC on the + vertical one. ``axes`` selects "h", "v" or "hv". + """ + horizontal, vertical = _axes(propagation_axis) + symmetry = [0, 0, 0] + if "h" in axes: + symmetry[horizontal] = -1 # PEC + if "v" in axes: + symmetry[vertical] = 1 # PMC + return symmetry[0], symmetry[1], symmetry[2] + + +@autoinit +class _RampSource(LinearlyPolarizedPlaneSource): + """Source whose transverse profile is unique per cell, so any mirror/transpose shows up.""" + + def _get_amplitude_raw(self, center): + del center + w = jnp.arange(self.grid_shape[self.horizontal_axis], dtype=self._config.dtype) + h = jnp.arange(self.grid_shape[self.vertical_axis], dtype=self._config.dtype) + return self._hv_to_grid(w[:, None] + 100.0 * h[None, :]) + + +def _build( + symmetry: tuple[int, int, int], + propagation_axis: int = 2, + direction: str = "+", + source_kind: str = "gaussian", + normalize_by_energy: bool = True, + transverse_cells: tuple[int, int] | None = None, + source_cells: tuple[int, int] | None = None, + source_start: tuple[int, int] | None = None, + **source_kwargs, +): + """Build a model with one plane source spanning (or centered in) the transverse plane.""" + horizontal, vertical = _axes(propagation_axis) + h_cells, v_cells = transverse_cells or (_TRANSVERSE_CELLS, _TRANSVERSE_CELLS) + shape = [0, 0, 0] + shape[propagation_axis] = _PROPAGATION_CELLS + shape[horizontal] = h_cells + shape[vertical] = v_cells + + config = fdtdx.SimulationConfig( + grid=fdtdx.UniformGrid(spacing=_SPACING), + time=5e-15, + dtype=jnp.float32, + symmetry=symmetry, + ) + volume = fdtdx.SimulationVolume(partial_grid_shape=(shape[0], shape[1], shape[2])) + objects, constraints = [volume], [] + bound_dict, boundary_constraints = fdtdx.boundary_objects_from_config( + fdtdx.BoundaryConfig.from_uniform_bound(thickness=3), volume + ) + constraints.extend(boundary_constraints) + objects.extend(bound_dict.values()) + + polarization = [0.0, 0.0, 0.0] + polarization[horizontal] = 1.0 + partial_shape: list[int | None] = [None, None, None] + partial_shape[propagation_axis] = 1 + if source_cells is not None: + partial_shape[horizontal] = source_cells[0] + partial_shape[vertical] = source_cells[1] + shared = dict( + name="src", + partial_grid_shape=(partial_shape[0], partial_shape[1], partial_shape[2]), + wave_character=fdtdx.WaveCharacter(wavelength=_WAVELENGTH), + direction=direction, + fixed_E_polarization_vector=(polarization[0], polarization[1], polarization[2]), + normalize_by_energy=normalize_by_energy, + **source_kwargs, + ) + if source_kind == "gaussian": + source = fdtdx.GaussianPlaneSource(radius=h_cells * _SPACING / 2, std=1 / 5, **shared) + elif source_kind == "uniform": + source = fdtdx.UniformPlaneSource(**shared) + else: + source = _RampSource(**shared) + + if source_start is None: + constraints.append(source.place_at_center(volume, axes=(horizontal, vertical))) + else: + # Absolute placement, used to park a small source strictly inside the kept upper half. + constraints.append( + source.set_grid_coordinates(axes=(horizontal, vertical), sides=("-", "-"), coordinates=source_start) + ) + if source_cells is None: + constraints.append(source.same_size(volume, axes=(horizontal, vertical))) + constraints.append( + source.set_grid_coordinates(axes=(propagation_axis,), sides=("-",), coordinates=(_PROPAGATION_CELLS // 2,)) + ) + objects.append(source) + return objects, constraints, config + + +def _place(symmetry, **kwargs): + """Place + apply a model and return the resulting source object.""" + objects, constraints, config = _build(symmetry, **kwargs) + key = jax.random.PRNGKey(0) + container, arrays, params, config, _ = fdtdx.place_objects( + object_list=objects, config=config, constraints=constraints, key=key + ) + arrays, container, _ = fdtdx.apply_params(arrays, container, params, key=key) + return next(obj for obj in container.objects if obj.name == "src") + + +def _kept_half(array: np.ndarray, symmetry: tuple[int, int, int], leading: int = 1) -> np.ndarray: + """Slice the upper half along every symmetric axis (``leading`` non-spatial axes in front).""" + index: list[slice] = [slice(None)] * (leading + 3) + for axis in range(3): + if symmetry[axis] != 0: + index[leading + axis] = slice(array.shape[leading + axis] // 2, None) + return array[tuple(index)] + + +class TestGaussianSourceUnderSymmetry: + """A clipped Gaussian must stay centered on the symmetry plane, not on the kept quadrant.""" + + @pytest.mark.parametrize("propagation_axis", [0, 1, 2]) + @pytest.mark.parametrize("axes", ["h", "v", "hv"]) + @pytest.mark.parametrize("direction", ["+", "-"]) + def test_matches_full_domain_kept_half(self, propagation_axis, axes, direction): + symmetry = _compatible_symmetry(propagation_axis, axes) + kwargs = dict(propagation_axis=propagation_axis, direction=direction) + full = _place((0, 0, 0), **kwargs) + reduced = _place(symmetry, **kwargs) + + for name in ("_E", "_H"): + full_field = np.asarray(getattr(full, name)) + reduced_field = np.asarray(getattr(reduced, name)) + expected = _kept_half(full_field, symmetry) + assert reduced_field.shape == expected.shape + scale = np.abs(expected).max() + assert scale > 0 + error = np.abs(reduced_field - expected).max() / scale + assert error < 1e-6, f"{name}: reduced injection differs from the full-domain half by {error:.2e}" + + @pytest.mark.parametrize("axes", ["h", "v", "hv"]) + def test_peak_sits_on_the_symmetry_corner(self, axes): + # The full-domain beam peaks at the domain center, i.e. on the symmetry plane(s), so the + # reduced profile must peak in the very first cell of every straddled axis. + symmetry = _compatible_symmetry(2, axes) + reduced = _place(symmetry) + profile = np.abs(np.asarray(reduced._E)[0, :, :, 0]) + peak = np.unravel_index(np.argmax(profile), profile.shape) + horizontal, vertical = _axes(2) + assert (peak[0] == 0) == (symmetry[horizontal] != 0) + assert (peak[1] == 0) == (symmetry[vertical] != 0) + + @pytest.mark.parametrize("normalize_by_energy", [True, False]) + def test_absolute_amplitude_matches_full_domain(self, normalize_by_energy): + # Guards the 2**k profile/energy normalization: both are sums over the source plane, which + # the reduction shrinks by the plane multiplicity. + symmetry = _compatible_symmetry(2, "hv") + full = _place((0, 0, 0), normalize_by_energy=normalize_by_energy) + reduced = _place(symmetry, normalize_by_energy=normalize_by_energy) + expected = _kept_half(np.asarray(full._E), symmetry) + ratio = np.abs(np.asarray(reduced._E)).max() / np.abs(expected).max() + assert abs(ratio - 1.0) < 1e-5, f"amplitude ratio reduced/full = {ratio:.6f}" + + def test_source_inside_kept_half_is_untouched(self): + # A source parked strictly inside the kept half is not clipped, so it keeps the full-domain + # profile centered on itself — the symmetry-aware center must not move it. + symmetry = _compatible_symmetry(2, "hv") + kwargs = dict(source_cells=(6, 6), source_start=(12, 12)) + full = _place((0, 0, 0), **kwargs) + reduced = _place(symmetry, **kwargs) + assert reduced.grid_shape == full.grid_shape + assert not any(reduced.straddles_symmetry_plane(a) for a in range(3)) + assert np.allclose(np.asarray(reduced._E), np.asarray(full._E), atol=1e-6) + + def test_source_touching_the_plane_from_above_is_not_mirror_extended(self): + # Starts exactly at the symmetry plane but lies entirely in the kept half: it was never + # clipped, so neither its center nor its normalization may assume a mirrored other half. + symmetry = _compatible_symmetry(2, "hv") + kwargs = dict(source_cells=(6, 6), source_start=(10, 10)) + full = _place((0, 0, 0), **kwargs) + reduced = _place(symmetry, **kwargs) + assert reduced.grid_slice_tuple[0][0] == 0 and reduced.grid_slice_tuple[1][0] == 0 + assert not any(reduced.straddles_symmetry_plane(a) for a in range(3)) + assert np.allclose(np.asarray(reduced._E), np.asarray(full._E), atol=1e-6) + + +class TestUniformSourceUnderSymmetry: + """UniformPlaneSource has no profile center, but its energy normalization is still a plane sum.""" + + @pytest.mark.parametrize("axes", ["h", "hv"]) + @pytest.mark.parametrize("normalize_by_energy", [True, False]) + def test_matches_full_domain_kept_half(self, axes, normalize_by_energy): + symmetry = _compatible_symmetry(2, axes) + kwargs = dict(source_kind="uniform", normalize_by_energy=normalize_by_energy) + full = _place((0, 0, 0), **kwargs) + reduced = _place(symmetry, **kwargs) + expected = _kept_half(np.asarray(full._E), symmetry) + error = np.abs(np.asarray(reduced._E) - expected).max() / np.abs(expected).max() + assert error < 1e-6, f"reduced uniform injection differs by {error:.2e}" + + +class TestTransverseProfileOrientation: + """The in-plane projection must not mirror or transpose the profile.""" + + @pytest.mark.parametrize("propagation_axis", [0, 1, 2]) + def test_backward_direction_does_not_mirror_profile(self, propagation_axis): + # A backward-propagating source carries the same transverse profile as a forward one; only + # the field orientation and the time-of-flight offsets differ. + forward = _place((0, 0, 0), propagation_axis=propagation_axis, direction="+", source_kind="ramp") + backward = _place((0, 0, 0), propagation_axis=propagation_axis, direction="-", source_kind="ramp") + + def profile(source): + field = np.asarray(source._E) + component = int(np.argmax(np.abs(field).sum(axis=(1, 2, 3)))) + values = source._grid_to_hv(jnp.asarray(field[component])) + return np.abs(np.asarray(values)) + + assert np.allclose(profile(forward), profile(backward), rtol=1e-6) + + @pytest.mark.parametrize("propagation_axis", [0, 1, 2]) + def test_non_square_plane_keeps_axis_order(self, propagation_axis): + # Propagation along y has descending oriented transverse axes; a profile built in + # (horizontal, vertical) order has to be transposed before it is placed on the grid. + source = _place( + (0, 0, 0), + propagation_axis=propagation_axis, + transverse_cells=(16, 8), + source_kind="ramp", + ) + field = np.asarray(source._E) + assert field.shape == (3, *source.grid_shape) + component = int(np.argmax(np.abs(field).sum(axis=(1, 2, 3)))) + profile = np.abs(np.asarray(source._grid_to_hv(jnp.asarray(field[component])))) + assert profile.shape == (16, 8) + # The ramp increases by 1 per horizontal cell and by 100 per vertical cell. + scale = profile[1, 0] - profile[0, 0] + assert scale > 0 + assert np.isclose((profile[0, 1] - profile[0, 0]) / scale, 100.0, rtol=1e-3) + + +class TestUnsupportedSymmetryPlacements: + """Configurations the reduction cannot represent must fail loudly, not silently.""" + + @pytest.mark.parametrize("kwargs", [{"azimuth_angle": 10.0}, {"elevation_angle": 10.0}]) + def test_tilted_source_straddling_plane_raises(self, kwargs): + objects, constraints, config = _build(_compatible_symmetry(2, "hv"), **kwargs) + with pytest.raises(ValueError, match="symmetry"): + fdtdx.place_objects(object_list=objects, config=config, constraints=constraints, key=jax.random.PRNGKey(0)) + + def test_random_offset_source_straddling_plane_raises(self): + objects, constraints, config = _build(_compatible_symmetry(2, "hv"), max_horizontal_offset=2e-7) + with pytest.raises(ValueError, match="symmetry"): + fdtdx.place_objects(object_list=objects, config=config, constraints=constraints, key=jax.random.PRNGKey(0)) + + +class TestPolarizationWallCompatibility: + """The wall type has to match the polarization; a mismatch is silent otherwise.""" + + @staticmethod + def _warnings(symmetry): + """Collect fdtdx's loguru warnings emitted while placing a model.""" + messages: list[str] = [] + handler_id = logger.add(lambda message: messages.append(message), level="WARNING") + try: + _place(symmetry) + finally: + logger.remove(handler_id) + return "".join(messages) + + def test_mismatched_polarization_warns(self): + # E along the horizontal axis is *normal* to the horizontal plane, so it is even only under + # PEC. Requesting PMC there makes it odd and the wall suppresses the injected field - the + # single most common way to get a plausible-looking but wrong symmetric simulation. + horizontal, _vertical = _axes(2) + symmetry = [0, 0, 0] + symmetry[horizontal] = 1 # PMC where PEC is required + output = self._warnings((symmetry[0], symmetry[1], symmetry[2])) + assert "makes Ex" in output and "odd" in output + assert "use PEC (-1)" in output # names the fix + + def test_matching_polarization_does_not_warn(self): + assert "mirror-even profile" not in self._warnings(_compatible_symmetry(2, "hv"))