import math
from typing import Literal
import jax
import jax.numpy as jnp
from loguru import logger
from fdtdx import constants
from fdtdx.core.grid import QuasiUniformGrid, RectilinearGrid, UniformGrid
from fdtdx.core.jax.pytrees import TreeClass, autoinit, field, frozen_field
from fdtdx.interfaces.recorder import Recorder
from fdtdx.typing import BackendOption
ExecutionMode = Literal["auto", "jax", "cuda"]
MaterialDiscretizationMethod = Literal["yee", "staircase"]
SubpixelCsgScope = Literal["direct", "shared", "all"]
@autoinit
class MaterialDiscretizationSpec(TreeClass):
"""Policy for mapping a physical material raster to Maxwell coefficients.
``"yee"`` (the default) treats scalar/diagonal values as primal-cell
averages and integrates them onto the actual staggered Ex/Ey/Ez dual
faces. The map is fixed, linear before inversion, and JAX
differentiable, so forward and adjoint simulations see exactly the same
material operator. ``"staircase"`` retains the historical behavior of
copying one cell value to every electric component and exists only for
reproducing legacy discrete-grid results.
"""
method: MaterialDiscretizationMethod = frozen_field(default="yee")
#: Retain the immutable unregistered cell-average material raster for a
#: caller that updates a manually managed physical region. Device scenes
#: enable this cache automatically. Static-only scenes leave it disabled
#: by default to avoid unnecessary memory.
retain_cell_average_baseline: bool = frozen_field(default=False)
#: Fixed number of sub-cell samples per Cartesian axis used by the
#: placement-order-aware static CSG compositor. The sample lattice is part
#: of the simulation configuration and never changes with material/design
#: values, so repeated forward/adjoint evaluations see the same map.
subpixel_samples_per_axis: int = frozen_field(default=8)
#: ``"direct"`` preserves the historical sequential Farjadpour map.
#: ``"shared"`` preserves the direct Farjadpour coefficient for an
#: isolated interface and invokes priority CSG only where two smoothed
#: interfaces cut one cell. ``"all"`` routes every smoothed cut cell
#: through the fixed priority sampler. Every choice is fixed before a solve
#: and therefore preserves a stationary forward/adjoint material map.
subpixel_csg_scope: SubpixelCsgScope = frozen_field(default="shared")
def __post_init__(self):
if self.method not in ("yee", "staircase"):
raise ValueError(
"material discretization method must be 'yee' or 'staircase', "
f"got {self.method!r}"
)
if (
not isinstance(self.subpixel_samples_per_axis, int)
or isinstance(self.subpixel_samples_per_axis, bool)
or self.subpixel_samples_per_axis < 1
):
raise ValueError("subpixel_samples_per_axis must be a positive integer")
if self.subpixel_csg_scope not in ("direct", "shared", "all"):
raise ValueError(
"subpixel_csg_scope must be 'direct', 'shared', or 'all'"
)
[docs]
@autoinit
class GradientConfig(TreeClass):
"""Configuration for gradient computation in simulations.
This class handles settings for automatic differentiation, supporting either
invertible differentiation with a recorder or checkpointing-based differentiation.
"""
#: Method for gradient computation. ``"auto"`` is a post-placement policy:
#: call :func:`fdtdx.configure_gradient_backend` once after object placement
#: to select an exact reversible solve when its PML tape fits, or the lean
#: checkpointed solve otherwise. Explicit methods retain their historical
#: behavior. ``"adjoint"`` is the exact two-solve
#: discrete adjoint for linear, lossless dielectric designs. It records the
#: electric update drive only inside ``design_region`` and avoids both the
#: full-field reverse reconstruction and the PML interface tape used by
#: ``"reversible"``.
method: Literal["auto", "reversible", "checkpointed", "adjoint"] = frozen_field(default="reversible")
#: Optional recorder for invertible differentiation. Needs to be provided for reversible autodiff. Defaults to None
recorder: Recorder | None = field(default=None)
#: Optional fixed number of checkpoints for checkpointing-based
#: differentiation. ``None`` selects the automatic, memory-aware planner:
#: it first chooses the recomputation-optimal online schedule for the
#: trajectory length and then caps it to the memory budget below. A fixed
#: integer remains available for reproducible benchmarking and expert
#: overrides.
num_checkpoints: int | None = frozen_field(default=None)
#: Fraction of total accelerator memory that the automatic checkpoint
#: planner may devote to copies of the evolving loop state. This is used
#: only when ``num_checkpoints`` and ``checkpoint_memory_budget_bytes`` are
#: both ``None``. The count depends only on array shapes, trajectory length,
#: and device capacity, so it remains fixed as design values evolve.
checkpoint_memory_fraction: float = frozen_field(default=0.25)
#: Optional absolute checkpoint-memory budget in bytes. This overrides
#: ``checkpoint_memory_fraction`` while retaining automatic selection of the
#: checkpoint count. It is useful on shared accelerators, but is not needed
#: when moving a campaign between geometries on the same device.
checkpoint_memory_budget_bytes: int | None = frozen_field(default=None)
#: Fraction of the device memory capacity reported by the JAX allocator
#: that an automatically selected exact reversible PML tape may occupy. The
#: remaining memory is reserved for live fields, detector state,
#: compilation, and reverse workspace. This policy is evaluated from
#: placed shapes and never from material values.
automatic_reversible_memory_fraction: float = frozen_field(default=0.42)
#: Number of interior full-field checkpoints for the ``"reversible"`` method.
#: The reversible backward pass reconstructs the field state by running the simulation in
#: reverse; for lossy materials this reverse reconstruction can accumulate numerical
#: error over the full trajectory. Setting this to ``k - 1`` partitions the run into ``k`` slices
#: and stores a full-field checkpoint at each interior slice boundary during the forward pass. The
#: backward pass then resets the reverse reconstruction to the exact checkpoint at every boundary,
#: bounding the reconstruction drift to a single slice (``~time_steps_total / k`` steps) at the
#: cost of O(k) field memory. The default ``0`` reproduces the classic single full reverse pass
#: (no interior checkpoints; only the final field, which is available for free, is used). Ignored
#: by the ``"checkpointed"`` method. Must not exceed ``time_steps_total - 1``.
num_checkpoints_reversible: int = frozen_field(default=0)
#: Spatial bounds of the differentiable inverse-permittivity region for the
#: ``"adjoint"`` method, expressed as ``((x0, x1), (y0, y1), (z0, z1))``.
#: Gradients outside this region are intentionally zero.
design_region: tuple[tuple[int, int], tuple[int, int], tuple[int, int]] | None = frozen_field(default=None)
#: Interpret ``design_region`` as primal material cells when the
#: component-aware Yee map is active. The compact adjoint tape then
#: automatically includes the one-cell positive-side halo affected by
#: dual-face registration. Set this to ``False`` only when the bounds
#: already describe Maxwell coefficients rather than material cells.
design_region_is_cell_average: bool = frozen_field(default=True)
def __post_init__(self):
if self.method not in ("auto", "reversible", "checkpointed", "adjoint"):
raise ValueError(f"unknown gradient method {self.method!r}")
if self.method == "reversible" and self.recorder is None:
raise Exception("Need Recorder in gradient config to compute reversible gradients")
if self.method == "adjoint" and self.design_region is None:
raise ValueError("Need design_region in gradient config to compute adjoint gradients")
if self.num_checkpoints is not None and (
not isinstance(self.num_checkpoints, int)
or isinstance(self.num_checkpoints, bool)
or self.num_checkpoints < 1
):
raise ValueError("num_checkpoints must be a positive integer or None for automatic planning")
if not 0.0 < self.checkpoint_memory_fraction <= 1.0:
raise ValueError("checkpoint_memory_fraction must be in (0, 1]")
if not 0.0 < self.automatic_reversible_memory_fraction < 1.0:
raise ValueError("automatic_reversible_memory_fraction must be in (0, 1)")
if self.checkpoint_memory_budget_bytes is not None and (
not isinstance(self.checkpoint_memory_budget_bytes, int)
or isinstance(self.checkpoint_memory_budget_bytes, bool)
or self.checkpoint_memory_budget_bytes < 1
):
raise ValueError("checkpoint_memory_budget_bytes must be a positive integer or None")
if self.num_checkpoints_reversible < 0:
raise Exception("num_checkpoints_reversible must be >= 0")
if self.design_region is not None and (
len(self.design_region) != 3
or any(len(bounds) != 2 or bounds[0] < 0 or bounds[1] <= bounds[0] for bounds in self.design_region)
):
raise ValueError(
"design_region must contain three non-empty non-negative (start, stop) bounds, "
f"got {self.design_region!r}"
)
[docs]
@autoinit
class SimulationConfig(TreeClass):
"""Configuration settings for FDTD simulations.
This class contains all the parameters needed to configure and run an FDTD
simulation, including spatial and temporal discretization, hardware backend,
and gradient computation settings.
"""
#: Total simulation time in seconds.
time: float = frozen_field()
#: Spatial grid configuration.
#:
#: ``UniformGrid`` is an unresolved policy used while the final volume shape
#: is still being inferred. ``RectilinearGrid`` is the realized solver grid
#: with explicit physical edge coordinates. Placement resolves policies to
#: ``RectilinearGrid`` so compiled FDTD code has exactly one metric source.
grid: UniformGrid | QuasiUniformGrid | RectilinearGrid = field()
#: Mapping from physical cell-average materials to staggered Yee
#: constitutive coefficients. The component-aware ``"yee"`` policy is
#: the production default; select ``MaterialDiscretizationSpec(
#: method="staircase")`` only to replay a historical scalar-raster model.
material_discretization: MaterialDiscretizationSpec = frozen_field(
default=MaterialDiscretizationSpec()
)
#: Computation backend ('gpu', 'tpu', 'cpu' or 'METAL'). Defaults to "gpu".
backend: BackendOption = frozen_field(default="gpu")
#: Data type for numerical computations. Defaults to jnp.float32.
dtype: jnp.dtype = frozen_field(default=jnp.float32)
#: Whether to use complex-valued field arrays.
#: None (default): auto-detect based on boundary conditions (e.g. Bloch).
#: True: force complex fields (complex64 if dtype=float32, complex128 if dtype=float64).
#: False: force real fields (raises error if Bloch boundaries are present).
use_complex_fields: bool | None = frozen_field(default=None)
#: Safety factor for the Courant condition (default: 0.99).
courant_factor: float = frozen_field(default=0.99)
#: Per-axis mirror symmetry of the simulation, in the order (x, y, z).
#: Each entry is one of ``{-1, 0, +1}``:
#: ``0`` = no symmetry on this axis (default),
#: ``-1`` = PEC (electric-wall) mirror on the axis center plane,
#: ``+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) 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`.
#: Each symmetric axis must resolve to an **even** number of grid cells (so the domain
#: splits exactly down the middle and the unfolded result matches the full domain
#: cell-for-cell); otherwise :func:`fdtdx.place_objects` raises a ``ValueError``.
symmetry: tuple[int, int, int] = frozen_field(default=(0, 0, 0))
#: Optional configuration for gradient computation.
gradient_config: GradientConfig | None = field(default=None)
#: FDTD propagator implementation. ``auto`` uses the native CUDA kernel
#: only for a fully supported scene and otherwise falls back to JAX.
#: The portable JAX implementation remains the default until native
#: coverage includes the general solver feature set.
execution_mode: ExecutionMode = frozen_field(default="jax")
def __post_init__(self):
from jax import extend
if len(self.symmetry) != 3 or any(s not in (-1, 0, 1) for s in self.symmetry):
raise ValueError(
f"config.symmetry must be a length-3 tuple with each entry in {{-1, 0, +1}} "
f"(0=none, -1=PEC, +1=PMC), got {self.symmetry!r}"
)
if self.execution_mode not in ("auto", "jax", "cuda"):
raise ValueError(
f"config.execution_mode must be one of 'auto', 'jax', or 'cuda', got {self.execution_mode!r}"
)
current_platform = extend.backend.get_backend().platform
if current_platform == "METAL" and self.backend == "gpu":
self.backend = "METAL"
if self.backend == "METAL":
try:
jax.devices()
if __name__ == "__main__":
logger.info("METAL device found and will be used for computations")
jax.config.update("jax_platform_name", "metal")
except RuntimeError:
if __name__ == "__main__":
logger.warning("METAL initialization failed, falling back to CPU!")
self.backend = "cpu"
elif self.backend in ["gpu", "tpu"]:
try:
jax.devices(self.backend)
if __name__ == "__main__":
logger.info(f"{str.upper(self.backend)} found and will be used for computations")
jax.config.update("jax_platform_name", self.backend)
except RuntimeError:
if __name__ == "__main__":
logger.warning(f"{str.upper(self.backend)} not found, falling back to CPU!")
self.backend = "cpu"
if self.backend == "cpu":
jax.config.update("jax_platform_name", "cpu")
@property
def has_symmetry(self) -> bool:
"""Whether any axis requests mirror symmetry.
Returns:
bool: True if at least one entry of :attr:`symmetry` is nonzero, meaning the
domain will be reduced and a PEC/PMC wall placed on the symmetry plane(s).
"""
return any(s != 0 for s in self.symmetry)
@property
def courant_number(self) -> float:
"""Calculate the Courant number for the simulation.
The Courant number is a dimensionless quantity that determines stability
of the FDTD simulation. It represents the ratio of the physical propagation
speed to the numerical propagation speed.
Returns:
float: The Courant number, scaled by the courant_factor and normalized
for 3D simulations.
"""
return self.courant_factor / math.sqrt(3)
def resolve_grid(self, shape: tuple[int, int, int] | None = None) -> RectilinearGrid:
"""Return a concrete solver grid.
Args:
shape: Required when ``grid`` is an unresolved ``UniformGrid``.
Returns:
A concrete ``RectilinearGrid``.
"""
if isinstance(self.grid, RectilinearGrid):
return self.grid
if shape is None:
raise ValueError("A grid shape is required to resolve UniformGrid.")
return self.grid.resolve(shape)
@property
def resolved_grid(self) -> RectilinearGrid | None:
"""Return the concrete solver grid, or ``None`` if not yet resolved.
``UniformGrid`` has no edge arrays until the simulation shape is known.
Callers that need coordinates, areas, or volumes should use this
property and fall back to ``uniform_spacing`` when it returns ``None``.
"""
if isinstance(self.grid, RectilinearGrid):
return self.grid
return None
@property
def has_nonuniform_grid(self) -> bool:
"""Whether the realized solver grid is non-uniform."""
grid = self.resolved_grid
return grid is not None and not grid.is_uniform
def uniform_spacing(self) -> float:
"""Return the uniform grid spacing.
``UniformGrid`` can answer this before placement. ``RectilinearGrid``
answers only when all spacings are equal and raises for non-uniform
meshes, making unsupported scalar assumptions explicit.
"""
if isinstance(self.grid, UniformGrid):
return self.grid.spacing
if isinstance(self.grid, QuasiUniformGrid):
if self.grid.is_uniform:
return self.grid.dx
else:
raise ValueError(
"QuasiUniformGrid has no single uniform spacing:"
f" ({self.grid.dx}, {self.grid.dy}, {self.grid.dz} differ). "
)
return self.grid.uniform_spacing # RectilinearGrid — raises internally if non-uniform
@property
def time_step_duration(self) -> float:
"""Calculate the duration of a single time step.
The time step duration is determined by the Courant condition to ensure
numerical stability. Realized rectilinear grids use their smallest
per-axis spacings. Unresolved uniform grids use their configured scalar
spacing; unresolved quasi-uniform grids use their smallest per-axis
spacing as a conservative CFL bound.
Returns:
float: Time step duration in seconds, calculated using the Courant
condition and spatial resolution.
"""
if isinstance(self.grid, RectilinearGrid):
return self.grid.cfl_time_step(self.courant_factor)
if isinstance(self.grid, UniformGrid):
return self.courant_number * self.grid.spacing / constants.c
if isinstance(self.grid, QuasiUniformGrid):
return self.courant_number * self.grid.min_spacing / constants.c
raise NotImplementedError(f"time_step_duration is not implemented for grid type {type(self.grid).__name__}.")
@property
def time_steps_total(self) -> int:
"""Calculate the total number of time steps for the simulation.
Determines how many discrete time steps are needed to simulate the
specified total simulation time, based on the time step duration.
Returns:
int: Total number of time steps needed to reach the specified
simulation time.
"""
return round(self.time / self.time_step_duration)
@property
def max_travel_distance(self) -> float:
"""Calculate the maximum distance light can travel during the simulation.
This represents the theoretical maximum distance that light could travel
through the simulation volume, useful for determining if the simulation
time is sufficient for light to traverse the entire domain.
Returns:
float: Maximum travel distance in meters, based on the speed of light
and total simulation time.
"""
return constants.c * self.time
@property
def only_forward(self) -> bool:
"""Check if the simulation is forward-only (no gradient computation).
Forward-only simulations don't compute gradients and are used when only
the forward propagation of electromagnetic fields is needed, without
optimization.
Returns:
bool: True if no gradient configuration is specified, False otherwise.
"""
return self.gradient_config is None
@property
def invertible_optimization(self) -> bool:
"""Check if invertible optimization is enabled.
Invertible optimization uses time-reversibility of Maxwell's equations
to compute gradients with reduced memory requirements compared to
checkpointing-based methods.
Returns:
bool: True if gradient computation uses invertible differentiation
(recorder is specified), False otherwise.
"""
if self.gradient_config is None:
return False
return self.gradient_config.recorder is not None
DUMMY_SIMULATION_CONFIG = SimulationConfig(
time=-1,
grid=UniformGrid(spacing=1),
)
