from pathlib import Path
from typing import Literal, Self
import jax
import jax.numpy as jnp
import numpy as np
from matplotlib import pyplot as plt
from fdtdx import constants
from fdtdx.core.axis import get_transverse_axes
from fdtdx.core.grid import calculate_time_offset_yee
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, compute_mode_symmetry_reduced
from fdtdx.dispersion import effective_complex_inv_permittivity, effective_inv_permittivity
from fdtdx.objects.sources.source import Source
from fdtdx.objects.sources.tfsf import TFSFPlaneSource, _build_dispersive_H_filter
[docs]
@autoinit
class ModePlaneSource(TFSFPlaneSource):
#: index of the mode
mode_index: int = frozen_field(default=0)
#: a literal value 'te', 'tm' to filter
filter_pol: Literal["te", "tm"] | None = frozen_field(default=None)
#: Symmetry-plane condition at the min edge of each transverse axis (the two
#: non-propagation physical axes, in increasing-index order): ``0`` = PEC
#: mirror (electric wall, the default), ``1`` = PMC mirror (magnetic wall).
#: Set this only for a **hand-built** half/quarter domain (your own PEC/PMC boundary at the min
#: edge), where it asks the mode solver for its own symmetric solve. It 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()
_inv_permeability: jax.Array | float = private_field()
_neff: jax.Array = private_field() # not required for sim, used for inspection
def _local_edge_coordinates(self) -> tuple[jax.Array, jax.Array, jax.Array] | None:
"""Return local physical edge coordinates for this source slice.
Non-uniform mode sources need edge coordinates for both Tidy3D mode
solving and Yee time offsets. Coordinates are shifted so the source
slice lower corner is at zero on each axis.
"""
grid = self._config.resolved_grid
if grid is None:
return None
local_edges = []
for axis in range(3):
lower, upper = self.grid_slice_tuple[axis]
edges = grid.edges(axis)[lower : upper + 1]
local_edges.append(edges - edges[0])
e0, e1, e2 = local_edges
return e0, e1, e2
def _transverse_edge_coordinates(self) -> tuple[jax.Array, jax.Array] | None:
"""Return local transverse edge coordinates for Tidy3D mode solving."""
local_edges = self._local_edge_coordinates()
if local_edges is None:
return None
axes = get_transverse_axes(self.propagation_axis)
return local_edges[axes[0]], local_edges[axes[1]]
def _mode_solver_resolution(self) -> float:
"""Return scalar resolution only for legacy uniform mode-solver setup.
``compute_mode`` ignores this value when explicit transverse coordinates
are provided, but the argument remains part of the compatibility API.
"""
if self._config.has_nonuniform_grid:
assert self._config.resolved_grid is not None
return self._config.resolved_grid.min_spacing
return self._config.uniform_spacing()
def _source_center_physical(self) -> jax.Array | None:
"""Return the physical source center for grid-aware Yee time offsets."""
local_edges = self._local_edge_coordinates()
if local_edges is None:
return None
center = []
for axis, edges in enumerate(local_edges):
if axis == self.propagation_axis:
center.append(jnp.asarray(0.0, dtype=self._config.dtype))
else:
center.append(0.5 * edges[-1])
return jnp.asarray(center, dtype=self._config.dtype)
def apply(
self: Self,
key: jax.Array,
inv_permittivities: jax.Array,
inv_permeabilities: jax.Array | float,
dispersive_c1: jax.Array | None = None,
dispersive_c2: jax.Array | None = None,
dispersive_c3: jax.Array | None = None,
electric_conductivity: jax.Array | None = None,
dispersive_c4: jax.Array | None = None,
) -> Self:
del key
if (
self.azimuth_angle != 0
or self.elevation_angle != 0
or self.max_angle_random_offset != 0
or self.max_vertical_offset != 0
or self.max_horizontal_offset != 0
):
raise NotImplementedError()
# inv_permittivities shape: (3, Nx, Ny, Nz) - slice with component dimension
inv_permittivity_slice = inv_permittivities[:, *self.grid_slice]
if isinstance(inv_permeabilities, jax.Array) and inv_permeabilities.ndim > 0:
# inv_permeabilities shape: (3, Nx, Ny, Nz) - slice with component dimension
inv_permeability_slice = inv_permeabilities[:, *self.grid_slice]
else:
inv_permeability_slice = inv_permeabilities
# Preserve the raw ε∞ slice before the carrier-frequency correction —
# the broadband impedance filter needs ε∞ to reconstruct ε(ω).
inv_eps_inf_slice = inv_permittivity_slice
# Frequency-correct the permittivity seen by the mode solver so that
# mode profiles computed inside a dispersive medium reflect the true
# epsilon at the carrier frequency, not epsilon_infinity.
c1_slice = c2_slice = c3_slice = c4_slice = None
if dispersive_c1 is not None and dispersive_c2 is not None and dispersive_c3 is not None:
c1_slice = dispersive_c1[:, :, *self.grid_slice]
c2_slice = dispersive_c2[:, :, *self.grid_slice]
c3_slice = dispersive_c3[:, :, *self.grid_slice]
c4_slice = None if dispersive_c4 is None else dispersive_c4[:, :, *self.grid_slice]
inv_permittivity_slice = effective_inv_permittivity(
inv_eps=inv_permittivity_slice,
c1=c1_slice,
c2=c2_slice,
c3=c3_slice,
omega=2.0 * np.pi * self.wave_character.get_frequency(),
dt=self._config.time_step_duration,
c4=c4_slice,
)
self = self.aset("_inv_permittivity", inv_permittivity_slice, create_new_ok=True)
self = self.aset("_inv_permeability", inv_permeability_slice, create_new_ok=True)
# Permittivity handed to the mode solver: the FULL complex epsilon at the
# carrier frequency (eps_inf + chi(omega) + i*sigma/(eps0*omega)), so the
# solved mode profile and effective index reflect material loss. This is
# kept separate from inv_permittivity_slice above, which stays real for the
# impedance/energy normalization — using the imaginary part there would
# double-count the absorption already integrated by the FDTD update.
sigma_slice = None if electric_conductivity is None else electric_conductivity[:, *self.grid_slice]
mode_inv_permittivity = inv_eps_inf_slice
if sigma_slice is not None or c1_slice is not None:
mode_inv_permittivity = effective_complex_inv_permittivity(
inv_eps=inv_eps_inf_slice,
omega=2.0 * np.pi * self.wave_character.get_frequency(),
dt=self._config.time_step_duration,
c1=c1_slice,
c2=c2_slice,
c3=c3_slice,
c4=c4_slice,
electric_conductivity=sigma_slice,
conductivity_spacing=(
None
if sigma_slice is None
else constants.c * self._config.time_step_duration / self._config.courant_number
),
)
# compute mode
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,
fixed_propagation_axis=self.propagation_axis,
)
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(),
fixed_propagation_axis=self.propagation_axis,
)
# 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
# the complex profile via a quadrature (cos/sin) decomposition. Lossless
# modes are projected to real (bit-identical to before). The dispersive
# path also stays real here because its broadband H-filter assumes a real
# temporal profile.
keep_complex_mode = sigma_slice is not None and c1_slice is None
if not keep_complex_mode:
mode_E, mode_H = jnp.real(mode_E), jnp.real(mode_H)
self = self.aset("_E", mode_E, create_new_ok=True)
self = self.aset("_H", mode_H, create_new_ok=True)
self = self.aset("_neff", eff_index, create_new_ok=True)
center = jnp.asarray(
[round(self.grid_shape[self.horizontal_axis]), round(self.grid_shape[self.vertical_axis])], dtype=jnp.int32
)
raw_wave_vector = get_wave_vector_raw(
direction=self.direction,
propagation_axis=self.propagation_axis,
dtype=self._config.dtype,
)
time_offset_E, time_offset_H = calculate_time_offset_yee(
center=center,
wave_vector=raw_wave_vector,
inv_permittivities=inv_permittivity_slice,
inv_permeabilities=jnp.ones_like(inv_permeability_slice),
resolution=self._mode_solver_resolution(),
time_step_duration=self._config.time_step_duration,
effective_index=jnp.real(eff_index),
coordinate_edges=self._local_edge_coordinates(),
center_physical=self._source_center_physical(),
)
self = self.aset("_time_offset_E", time_offset_E, create_new_ok=True)
self = self.aset("_time_offset_H", time_offset_H, create_new_ok=True)
# Broadband impedance correction for dispersive media. The mode solver
# above used ε(ω_c), so the resulting H profile already carries the
# correct scalar impedance at the carrier frequency. For a broadband
# pulse the medium's ε(ω) varies across the source spectrum, which
# mismatches η away from ω_c and radiates spurious reflections through
# the TFSF surface. Precompute a filtered H-side temporal profile
# whose spectrum is S(ω)·√(ε(ω)/ε(ω_c)) to bake in the frequency-
# dependent correction.
#
# Note: bulk ε(ω) is averaged uniformly over the source cells; this
# does not capture geometric modal dispersion (the fact that a
# waveguide mode's effective index also depends on frequency).
if c1_slice is not None and c2_slice is not None and c3_slice is not None:
filtered = _build_dispersive_H_filter(
temporal_profile=self.temporal_profile,
wave_character=self.wave_character,
dt=self._config.time_step_duration,
num_time_steps=self._config.time_steps_total,
c1_slice=c1_slice,
c2_slice=c2_slice,
c3_slice=c3_slice,
inv_eps_inf_slice=inv_eps_inf_slice,
dtype=self._config.dtype,
c4_slice=c4_slice,
)
self = self.aset("_temporal_H_filter", filtered, create_new_ok=True)
else:
# Reused source applied in a non-dispersive context: clear any stale
# filter from a previous dispersive apply.
self = self.aset("_temporal_H_filter", None, create_new_ok=True)
return self
def plot(self, save_path: str | Path):
if self._H is None or self._E is None:
raise Exception("Cannot plot mode without init to grid and apply params first")
energy = compute_energy(
E=self._E,
H=self._H,
inv_permittivity=self._inv_permittivity,
inv_permeability=self._inv_permeability,
)
energy_2d = energy.squeeze().T
plt.clf()
fig = plt.figure(figsize=(10, 10))
mode_cmap = "inferno"
im = plt.imshow(
energy_2d,
cmap=mode_cmap,
origin="lower",
aspect="equal",
)
plt.colorbar(im)
# Ensure the plot takes up the entire figure
plt.tight_layout(pad=0)
plt.savefig(save_path, bbox_inches="tight", pad_inches=0)
plt.close(fig)
@autoinit
class ModeProfileCurrentSource(ModePlaneSource):
"""Soft electric-current sheet shaped as a solved waveguide mode.
Unlike :class:`ModePlaneSource`, this source does not inject the magnetic
Huygens partner and therefore does not select a propagation direction. It
is an impressed electric current with transverse profile ``mode_E`` and
radiates to both sides of its plane, the distributed analogue of an
electric point dipole. ``direction`` only fixes the otherwise arbitrary
phase of the reference-mode solve.
A sheet may touch a symmetry plane normal to its mode axis when
``subcell_symmetry_centered`` is true. The kept first-cell deposit is then
weighted by one half so its mirror pair preserves the full-domain current
moment.
"""
amplitude: float = frozen_field(default=1.0)
subcell_symmetry_centered: bool = frozen_field(default=False)
def validate_placement(self, objects) -> list[str]:
# Deliberately bypass TFSFPlaneSource's prohibition on a plane normal
# to a symmetry wall: this is a soft current, not a Huygens boundary.
errors = list(Source.validate_placement(self, objects))
if self.touches_symmetry_plane(self.propagation_axis) and not self.subcell_symmetry_centered:
errors.append(
f"Mode current '{self.name}' touches the symmetry plane normal to its mode axis; "
"set subcell_symmetry_centered=True or move it off the plane."
)
return errors
@property
def _current_centering_scale(self) -> float:
if self.subcell_symmetry_centered and self.touches_symmetry_plane(self.propagation_axis):
return 0.5
return 1.0
def with_electric_current_profile(self, profile: jax.Array) -> Self:
"""Return a copy that injects ``profile`` instead of the solved mode.
The mode solve performed by :meth:`ModePlaneSource.apply` remains the
reference used to establish the source plane, normalization, and
detector reciprocity. This method only replaces the impressed
electric-current profile used by :meth:`update_E`. If ``profile`` is
assembled from differentiable coefficients inside an objective, the
checkpointed FDTD reverse pass differentiates through those
coefficients as well as through the material arrays.
Normalization and physical-support constraints deliberately belong to
the caller: different applications use different current norms and
launchable subspaces. Keeping this primitive literal prevents the
solver from silently changing source amplitude.
"""
if self._E is None:
raise ValueError("the reference mode must be applied before replacing its current profile")
values = jnp.asarray(profile)
if tuple(values.shape) != tuple(self._E.shape):
raise ValueError(
f"current profile shape {tuple(values.shape)} does not match "
f"the solved mode shape {tuple(self._E.shape)}"
)
if jnp.iscomplexobj(values):
raise ValueError("ModeProfileCurrentSource currently requires a real electric-current profile")
return self.aset("_E", values.astype(jnp.real(self._E).dtype), create_new_ok=True)
def update_E(
self,
E: jax.Array,
inv_permittivities: jax.Array,
inv_permeabilities: jax.Array | float,
time_step: jax.Array,
inverse: bool,
) -> jax.Array:
del inv_permeabilities
amplitude = self.temporal_profile.get_amplitude(
time=time_step * self._config.time_step_duration,
period=self.wave_character.get_period(),
phase_shift=self.wave_character.phase_shift,
)
current = jnp.real(self._E)
inv_eps = inv_permittivities[:, *self.grid_slice]
if inv_eps.shape[0] == 9:
inv_eps_tensor = inv_eps.reshape((3, 3, *inv_eps.shape[1:]))
response = jnp.einsum("ij...,j...->i...", inv_eps_tensor, current)
else:
response = inv_eps * current
sign = 1.0 if inverse else -1.0
scale = (
sign
* self._config.courant_number
* self.static_amplitude_factor
* self.amplitude
* self._current_centering_scale
* amplitude
)
return E.at[:, *self.grid_slice].add((scale * response).astype(E.dtype))
def update_H(
self,
H: jax.Array,
inv_permittivities: jax.Array,
inv_permeabilities: jax.Array | float,
time_step: jax.Array,
inverse: bool,
) -> jax.Array:
del inv_permittivities, inv_permeabilities, time_step, inverse
return H
