witwin-maxwell 0.3.0__py3-none-win_amd64.whl

witwin/maxwell/fdtd/boundary/runtime.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import torch
+
+from .common import BOUNDARY_KIND_TO_CODE, BOUNDARY_PEC, initialize_complex_fields
+from .bloch import resolve_bloch_phase_factors, resolve_bloch_wavevector
+from .cpml import (
+    initialize_absorber_state,
+    initialize_cpml_state,
+    initialize_neutral_boundary_state,
+    initialize_simple_pml_state,
+    initialize_stable_pml_state,
+)
+
+
+def _configure_face_boundary_codes(solver):
+    symmetry = getattr(solver.scene, "symmetry", (None, None, None))
+    boundary = solver.scene.boundary
+    face_codes = {
+        axis: [
+            BOUNDARY_KIND_TO_CODE[boundary.face_kind(axis, "low")],
+            BOUNDARY_KIND_TO_CODE[boundary.face_kind(axis, "high")],
+        ]
+        for axis in ("x", "y", "z")
+    }
+    for axis, entry in zip(("x", "y", "z"), symmetry):
+        if entry is None:
+            continue
+        mode, face = entry
+        if boundary.face_kind(axis, face) not in {"none", "pml"}:
+            raise ValueError(
+                f"Scene.symmetry on the {axis}-{face} face requires "
+                "BoundarySpec.none() or BoundarySpec.pml(...)."
+            )
+        face_index = 0 if face == "low" else 1
+        face_codes[axis][face_index] = BOUNDARY_KIND_TO_CODE[mode.lower()]
+
+    _validate_symmetry_source_placement(solver.scene, symmetry)
+
+    for axis in ("x", "y", "z"):
+        low_code, high_code = (int(face_codes[axis][0]), int(face_codes[axis][1]))
+        setattr(solver, f"boundary_{axis}_low_code", low_code)
+        setattr(solver, f"boundary_{axis}_high_code", high_code)
+
+    solver.has_pec_faces = any(
+        code == BOUNDARY_PEC
+        for codes in face_codes.values()
+        for code in codes
+    )
+    solver.has_pml_faces = any(
+        code == BOUNDARY_KIND_TO_CODE["pml"]
+        for codes in face_codes.values()
+        for code in codes
+    )
+    solver.has_periodic_axes = tuple(
+        axis
+        for axis in ("x", "y", "z")
+        if boundary.axis_kind(axis) == "periodic"
+    )
+    solver.has_bloch_axes = tuple(
+        axis
+        for axis in ("x", "y", "z")
+        if boundary.axis_kind(axis) == "bloch"
+    )
+    solver.mur_faces = tuple(
+        (axis, side)
+        for axis in ("x", "y", "z")
+        for side in ("low", "high")
+        if boundary.face_kind(axis, side) == "mur"
+    )
+    solver.has_mur_faces = bool(solver.mur_faces)
+
+
+_AXIS_INDEX = {"x": 0, "y": 1, "z": 2}
+
+
+def _source_center(source):
+    """Return a point-like source center coordinate tuple, or ``None`` if the
+    source spans the domain (e.g. plane waves) and has no single position."""
+    for attr in ("position", "center"):
+        value = getattr(source, attr, None)
+        if value is not None:
+            return tuple(float(v) for v in value)
+    return None
+
+
+def _validate_symmetry_source_placement(scene, symmetry):
+    """Reject point-like sources located in the folded-away half of the domain.
+
+    A source may sit on the symmetry plane or anywhere in the kept half; a source
+    beyond the plane lives in the removed half and cannot be represented on the
+    folded grid, so we raise clearly instead of silently dropping it.
+    """
+    if all(entry is None for entry in symmetry):
+        return
+    resolver = getattr(scene, "resolved_sources", None)
+    sources = resolver() if callable(resolver) else getattr(scene, "sources", ())
+    bounds = scene.domain.bounds
+    tol = 1e-9
+    for source in sources:
+        center = _source_center(source)
+        if center is None:
+            continue
+        for axis, entry in zip(("x", "y", "z"), symmetry):
+            if entry is None:
+                continue
+            _mode, face = entry
+            axis_idx = _AXIS_INDEX[axis]
+            low, high = float(bounds[axis_idx][0]), float(bounds[axis_idx][1])
+            coord = center[axis_idx]
+            span = max(high - low, tol)
+            plane = low if face == "low" else high
+            if face == "low" and coord < plane - tol * span:
+                raise ValueError(
+                    f"Source at {axis}={coord:g} lies in the folded-away half below the "
+                    f"{axis}-low symmetry plane at {plane:g}; place it on the plane or in "
+                    "the kept half of the domain."
+                )
+            if face == "high" and coord > plane + tol * span:
+                raise ValueError(
+                    f"Source at {axis}={coord:g} lies in the folded-away half above the "
+                    f"{axis}-high symmetry plane at {plane:g}; place it on the plane or in "
+                    "the kept half of the domain."
+                )
+
+
+def initialize_boundary_state(solver):
+    solver.boundary_kind = solver.scene.boundary.kind
+    solver.boundary_code = (
+        BOUNDARY_KIND_TO_CODE[solver.boundary_kind]
+        if solver.boundary_kind in BOUNDARY_KIND_TO_CODE
+        else None
+    )
+    _configure_face_boundary_codes(solver)
+
+    solver.complex_fields_enabled = solver.scene.boundary.uses_kind("bloch")
+    if solver.complex_fields_enabled:
+        initialize_complex_fields(solver)
+
+    solver.resolved_bloch_wavevector = resolve_bloch_wavevector(solver)
+    phases = resolve_bloch_phase_factors(solver)
+    solver.boundary_phase_cos = tuple(float(phase.real) for phase in phases)
+    solver.boundary_phase_sin = tuple(float(phase.imag) for phase in phases)
+    solver.active_absorber_type = solver.absorber_type if solver.scene.boundary.uses_kind("pml") else "none"
+    solver.uses_cpml = solver.active_absorber_type in ("cpml", "stablepml")
+
+    if solver.scene.boundary.uses_kind("pml"):
+        if solver.absorber_type == "cpml":
+            initialize_cpml_state(solver)
+        elif solver.absorber_type == "stablepml":
+            initialize_stable_pml_state(solver)
+        elif solver.absorber_type == "pml":
+            initialize_simple_pml_state(solver)
+        elif solver.absorber_type == "absorber":
+            initialize_absorber_state(solver)
+        else:
+            raise ValueError(
+                "PML boundaries require absorber_type to be 'cpml', 'stablepml', 'pml', or 'absorber'."
+            )
+    else:
+        initialize_neutral_boundary_state(solver)
+
+    initialize_mur_state(solver)
+
+
+# Tangential electric-field components on each outer face and the Yee grid axis
+# they are indexed along for the normal direction of that face.
+_MUR_FACE_COMPONENTS = {
+    "x": (("Ey", 0), ("Ez", 0)),
+    "y": (("Ex", 1), ("Ez", 1)),
+    "z": (("Ex", 2), ("Ey", 2)),
+}
+
+
+def initialize_mur_state(solver):
+    solver._mur_state = []
+    if not getattr(solver, "has_mur_faces", False):
+        return
+
+    scene = solver.scene
+    axis_primal = {"x": scene.dx_primal64, "y": scene.dy_primal64, "z": scene.dz_primal64}
+    speed_dt = solver.c * solver.dt
+    for axis, side in solver.mur_faces:
+        # Mur advection uses the local boundary-cell spacing of each face. The
+        # coefficient depends only on c, dt, and that spacing, so it is constant
+        # across time steps and precomputed once here.
+        primal = axis_primal[axis]
+        delta = float(primal[0] if side == "low" else primal[-1])
+        coef = (speed_dt - delta) / (speed_dt + delta)
+        for field_name, field_axis in _MUR_FACE_COMPONENTS[axis]:
+            field = getattr(solver, field_name)
+            size = int(field.shape[field_axis])
+            boundary_index = 0 if side == "low" else size - 1
+            adjacent_index = 1 if side == "low" else size - 2
+            # Persistent boundary / first-interior plane buffers, updated in place
+            # by the Mur kernel each step (no per-step allocation). They start at
+            # zero, matching the initial field state.
+            plane_shape = tuple(field.shape[:field_axis] + field.shape[field_axis + 1:])
+            solver._mur_state.append(
+                {
+                    "field": field_name,
+                    "axis": int(field_axis),
+                    "boundary_index": int(boundary_index),
+                    "adjacent_index": int(adjacent_index),
+                    "coef": float(coef),
+                    "prev_boundary": torch.zeros(plane_shape, device=solver.device, dtype=field.dtype),
+                    "prev_adjacent": torch.zeros(plane_shape, device=solver.device, dtype=field.dtype),
+                }
+            )

witwin/maxwell/fdtd/checkpoint.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from .boundary import expand_cpml_memory_tensor
+
+_FIELD_STATE_NAMES = ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz")
+_COMPLEX_FIELD_STATE_NAMES = ("Ex_imag", "Ey_imag", "Ez_imag", "Hx_imag", "Hy_imag", "Hz_imag")
+_CPML_STATE_NAMES = (
+    "psi_ex_y",
+    "psi_ex_z",
+    "psi_ey_x",
+    "psi_ey_z",
+    "psi_ez_x",
+    "psi_ez_y",
+    "psi_hx_y",
+    "psi_hx_z",
+    "psi_hy_x",
+    "psi_hy_z",
+    "psi_hz_x",
+    "psi_hz_y",
+)
+_CPML_COMPLEX_STATE_NAMES = tuple(f"{name}_imag" for name in _CPML_STATE_NAMES)
+_TFSF_AUXILIARY_STATE_NAMES = (
+    "tfsf_aux_electric",
+    "tfsf_aux_magnetic",
+)
+_DISPERSIVE_STATE_TENSORS = {
+    "debye": ("polarization", "current"),
+    "drude": ("current",),
+    "lorentz": ("polarization", "current"),
+}
+_CHECKPOINT_SCHEMA_VERSION = 1
+
+
+def dispersive_state_name(component_name: str, model_name: str, index: int, tensor_name: str) -> str:
+    return f"{model_name}_{component_name}_{int(index)}_{tensor_name}"
+
+
+def iter_dispersive_state_specs(solver):
+    if not getattr(solver, "dispersive_enabled", False):
+        return
+    templates = getattr(solver, "_dispersive_templates", {})
+    for component_name in ("Ex", "Ey", "Ez"):
+        component_templates = templates.get(component_name, {})
+        for model_name in ("debye", "drude", "lorentz"):
+            tensor_names = _DISPERSIVE_STATE_TENSORS[model_name]
+            for index, entry in enumerate(component_templates.get(model_name, ())):
+                yield component_name, model_name, index, tensor_names, entry
+
+
+def iter_magnetic_dispersive_state_specs(solver):
+    if not getattr(solver, "magnetic_dispersive_enabled", False):
+        return
+    templates = getattr(solver, "_magnetic_dispersive_templates", {})
+    for component_name in ("Hx", "Hy", "Hz"):
+        component_templates = templates.get(component_name, {})
+        for model_name in ("debye", "drude", "lorentz"):
+            tensor_names = _DISPERSIVE_STATE_TENSORS[model_name]
+            for index, entry in enumerate(component_templates.get(model_name, ())):
+                yield component_name, model_name, index, tensor_names, entry
+
+
+@dataclass(frozen=True)
+class FDTDCheckpointSchema:
+    version: int
+    field_names: tuple[str, ...]
+    complex_field_names: tuple[str, ...]
+    cpml_state_names: tuple[str, ...]
+    tfsf_auxiliary_state_names: tuple[str, ...]
+    dispersive_state_names: tuple[str, ...]
+    magnetic_dispersive_state_names: tuple[str, ...] = ()
+
+    @property
+    def state_names(self) -> tuple[str, ...]:
+        return (
+            self.field_names
+            + self.complex_field_names
+            + self.cpml_state_names
+            + self.tfsf_auxiliary_state_names
+            + self.dispersive_state_names
+            + self.magnetic_dispersive_state_names
+        )
+
+
+def checkpoint_schema(solver) -> FDTDCheckpointSchema:
+    complex_field_names = _COMPLEX_FIELD_STATE_NAMES if bool(getattr(solver, "complex_fields_enabled", False)) else ()
+    cpml_state_names = ()
+    if getattr(solver, "uses_cpml", False):
+        cpml_state_names = _CPML_STATE_NAMES
+        if bool(getattr(solver, "complex_fields_enabled", False)):
+            cpml_state_names = cpml_state_names + _CPML_COMPLEX_STATE_NAMES
+
+    tfsf_auxiliary_state_names = ()
+    if getattr(solver, "tfsf_enabled", False):
+        tfsf_state = getattr(solver, "_tfsf_state", None)
+        if tfsf_state is not None and tfsf_state.get("auxiliary_grid") is not None:
+            tfsf_auxiliary_state_names = _TFSF_AUXILIARY_STATE_NAMES
+
+    dispersive_state_names = []
+    for component_name, model_name, index, tensor_names, _entry in iter_dispersive_state_specs(solver) or ():
+        for tensor_name in tensor_names:
+            dispersive_state_names.append(dispersive_state_name(component_name, model_name, index, tensor_name))
+    if bool(getattr(solver, "complex_fields_enabled", False)):
+        # Bloch (complex-field) runs keep an imaginary ADE replica per electric
+        # pole; the reverse replay differentiates through it, so it must be part
+        # of the frozen checkpoint layout too.
+        for component_name, model_name, index, tensor_names, _entry in iter_dispersive_state_specs(solver) or ():
+            for tensor_name in tensor_names:
+                dispersive_state_names.append(
+                    dispersive_state_name(component_name, model_name, index, tensor_name) + "_imag"
+                )
+
+    magnetic_dispersive_state_names = []
+    for component_name, model_name, index, tensor_names, _entry in iter_magnetic_dispersive_state_specs(solver) or ():
+        for tensor_name in tensor_names:
+            magnetic_dispersive_state_names.append(
+                dispersive_state_name(component_name, model_name, index, tensor_name)
+            )
+    if bool(getattr(solver, "complex_fields_enabled", False)):
+        magnetic_dispersive_state_names.extend(
+            name + "_imag" for name in tuple(magnetic_dispersive_state_names)
+        )
+
+    return FDTDCheckpointSchema(
+        version=_CHECKPOINT_SCHEMA_VERSION,
+        field_names=_FIELD_STATE_NAMES,
+        complex_field_names=tuple(complex_field_names),
+        cpml_state_names=tuple(cpml_state_names),
+        tfsf_auxiliary_state_names=tuple(tfsf_auxiliary_state_names),
+        dispersive_state_names=tuple(dispersive_state_names),
+        magnetic_dispersive_state_names=tuple(magnetic_dispersive_state_names),
+    )
+
+
+def checkpoint_state_names(solver) -> tuple[str, ...]:
+    return checkpoint_schema(solver).state_names
+
+
+@dataclass(frozen=True)
+class FDTDCheckpointState:
+    step: int
+    schema: FDTDCheckpointSchema
+    tensors: dict[str, torch.Tensor]
+
+
+def validate_checkpoint_state(
+    state: FDTDCheckpointState,
+    *,
+    expected_schema: FDTDCheckpointSchema | None = None,
+) -> FDTDCheckpointSchema:
+    state_names = tuple(state.tensors.keys())
+    if state_names != state.schema.state_names:
+        raise RuntimeError(
+            "Checkpoint tensor layout drifted from the frozen adjoint state schema. "
+            f"Expected {state.schema.state_names}, got {state_names}."
+        )
+    if expected_schema is not None and state.schema != expected_schema:
+        raise RuntimeError(
+            "Checkpoint schema mismatch detected during adjoint replay. "
+            f"Expected version {expected_schema.version} with names {expected_schema.state_names}, "
+            f"got version {state.schema.version} with names {state.schema.state_names}."
+        )
+    return state.schema
+
+
+def capture_checkpoint_state(solver, step: int) -> FDTDCheckpointState:
+    schema = checkpoint_schema(solver)
+    tensors = {
+        name: getattr(solver, name).detach().clone()
+        for name in _FIELD_STATE_NAMES
+    }
+    if bool(getattr(solver, "complex_fields_enabled", False)):
+        tensors.update(
+            {
+                name: getattr(solver, name).detach().clone()
+                for name in _COMPLEX_FIELD_STATE_NAMES
+            }
+        )
+    if getattr(solver, "uses_cpml", False):
+        tensors.update(
+            {
+                name: expand_cpml_memory_tensor(solver, name).detach().clone()
+                for name in schema.cpml_state_names
+            }
+        )
+    if getattr(solver, "tfsf_enabled", False):
+        tfsf_state = getattr(solver, "_tfsf_state", None)
+        auxiliary_grid = None if tfsf_state is None else tfsf_state.get("auxiliary_grid")
+        if auxiliary_grid is not None:
+            tensors["tfsf_aux_electric"] = auxiliary_grid.electric.detach().clone()
+            tensors["tfsf_aux_magnetic"] = auxiliary_grid.magnetic.detach().clone()
+    for component_name, model_name, index, tensor_names, entry in iter_dispersive_state_specs(solver) or ():
+        for tensor_name in tensor_names:
+            tensors[dispersive_state_name(component_name, model_name, index, tensor_name)] = (
+                entry[tensor_name].detach().clone()
+            )
+    if bool(getattr(solver, "complex_fields_enabled", False)):
+        for component_name, model_name, index, tensor_names, entry in iter_dispersive_state_specs(solver) or ():
+            for tensor_name in tensor_names:
+                tensors[dispersive_state_name(component_name, model_name, index, tensor_name) + "_imag"] = (
+                    entry[f"{tensor_name}_imag"].detach().clone()
+                )
+    for component_name, model_name, index, tensor_names, entry in iter_magnetic_dispersive_state_specs(solver) or ():
+        for tensor_name in tensor_names:
+            tensors[dispersive_state_name(component_name, model_name, index, tensor_name)] = (
+                entry[tensor_name].detach().clone()
+            )
+    if bool(getattr(solver, "complex_fields_enabled", False)):
+        for component_name, model_name, index, tensor_names, entry in iter_magnetic_dispersive_state_specs(solver) or ():
+            for tensor_name in tensor_names:
+                tensors[dispersive_state_name(component_name, model_name, index, tensor_name) + "_imag"] = (
+                    entry[f"{tensor_name}_imag"].detach().clone()
+                )
+    state = FDTDCheckpointState(step=int(step), schema=schema, tensors=tensors)
+    validate_checkpoint_state(state)
+    return state
+
+
+def clone_checkpoint_tensors(state: FDTDCheckpointState) -> dict[str, torch.Tensor]:
+    validate_checkpoint_state(state)
+    return {
+        name: tensor.detach().clone()
+        for name, tensor in state.tensors.items()
+    }

witwin/maxwell/fdtd/coords.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import numpy as np
+
+from ..monitors import normalize_component
+from ..scene import prepare_scene
+
+
+def component_coords(scene, component: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    resolved_scene = prepare_scene(scene)
+    component_name = normalize_component(component)
+    x, y, z = resolved_scene.x_nodes64, resolved_scene.y_nodes64, resolved_scene.z_nodes64
+    x_half, y_half, z_half = resolved_scene.x_half64, resolved_scene.y_half64, resolved_scene.z_half64
+
+    if component_name == "ex":
+        return x_half, y, z
+    if component_name == "ey":
+        return x, y_half, z
+    if component_name == "ez":
+        return x, y, z_half
+    if component_name == "hx":
+        return x, y_half, z_half
+    if component_name == "hy":
+        return x_half, y, z_half
+    if component_name == "hz":
+        return x_half, y_half, z
+    raise ValueError(f"Unsupported field component: {component!r}")
+
+
+def centered_cell_coords(scene) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    resolved_scene = prepare_scene(scene)
+    return resolved_scene.x_half64, resolved_scene.y_half64, resolved_scene.z_half64

witwin/maxwell/fdtd/cuda/__init__.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from .backend import build_info, debug_linear_indices, get_compiled_extension, get_native_fdtd_module, is_available
+
+__all__ = [
+    "build_info",
+    "debug_linear_indices",
+    "get_compiled_extension",
+    "get_native_fdtd_module",
+    "is_available",
+]