witwin-maxwell 0.3.0__py3-none-win_amd64.whl

witwin/maxwell/compiler/monitors.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+from ..monitors import (
+    DiffractionMonitor,
+    DipoleEmissionMonitor,
+    FieldTimeMonitor,
+    FinitePlaneMonitor,
+    FluxMonitor,
+    FluxTimeMonitor,
+    MediumMonitor,
+    ModeMonitor,
+    PermittivityMonitor,
+    PlaneMonitor,
+    PointMonitor,
+    normalize_component,
+    normalize_axis,
+)
+from ..sources import PointDipole
+
+_POLARIZATION_AXES = (("x", "Ex"), ("y", "Ey"), ("z", "Ez"))
+
+
+def _find_scene_dipole(scene, source_name: str) -> PointDipole:
+    for source in getattr(scene, "sources", ()):
+        if isinstance(source, PointDipole) and getattr(source, "name", None) == source_name:
+            return source
+    raise ValueError(
+        f"DipoleEmissionMonitor references PointDipole {source_name!r}, "
+        "which is not present in the scene."
+    )
+
+
+def _point_observer_record(name, position, *, component: str) -> dict[str, object]:
+    return {
+        "name": str(name),
+        "kind": "point",
+        "position": tuple(float(value) for value in position),
+        "component": normalize_component(component),
+    }
+
+
+def _plane_observer_record(name, axis, position, *, component: str) -> dict[str, object]:
+    return {
+        "name": str(name),
+        "kind": "plane",
+        "axis": normalize_axis(axis),
+        "position": float(position),
+        "component": normalize_component(component),
+    }
+
+
+def _internal_observer_name(monitor_name: str, component: str) -> str:
+    return f"{monitor_name}::{normalize_component(component)}"
+
+
+def compile_fdtd_observers(scene):
+    compiled = []
+    monitors = scene.resolved_monitors() if hasattr(scene, "resolved_monitors") else scene.monitors
+    for monitor in monitors:
+        if isinstance(monitor, (FieldTimeMonitor, FluxTimeMonitor, PermittivityMonitor, MediumMonitor)):
+            # Time-domain monitors are handled by compile_fdtd_time_observers and
+            # material monitors are resolved from compiled material tensors at result
+            # time; neither must become a running-DFT spectral observer.
+            continue
+        if isinstance(monitor, DipoleEmissionMonitor):
+            dipole = _find_scene_dipole(scene, monitor.source_name)
+            polarization = tuple(float(value) for value in dipole.polarization)
+            position = tuple(float(value) for value in dipole.position)
+            active = tuple(
+                component
+                for (axis, component), weight in zip(_POLARIZATION_AXES, polarization)
+                if weight != 0.0
+            )
+            monitor_fields = tuple(normalize_component(component) for component in active)
+            for component in monitor_fields:
+                observer = _point_observer_record(
+                    _internal_observer_name(monitor.name, component),
+                    position,
+                    component=component,
+                )
+                observer["monitor_name"] = monitor.name
+                observer["monitor_fields"] = monitor_fields
+                observer["monitor_frequencies"] = monitor.frequencies
+                observer["monitor_type"] = "dipole_emission"
+                observer["dipole_polarization"] = polarization
+                observer["dipole_position"] = position
+                observer["source_name"] = monitor.source_name
+                compiled.append(observer)
+            continue
+        normalized_fields = tuple(normalize_component(field) for field in monitor.fields)
+        if isinstance(monitor, PointMonitor):
+            for field in normalized_fields:
+                observer = _point_observer_record(
+                    _internal_observer_name(monitor.name, field),
+                    monitor.position,
+                    component=field,
+                )
+                observer["monitor_name"] = monitor.name
+                observer["monitor_fields"] = normalized_fields
+                compiled.append(observer)
+            continue
+        if isinstance(monitor, (PlaneMonitor, FinitePlaneMonitor, FluxMonitor, ModeMonitor, DiffractionMonitor)):
+            monitor_plane_position = monitor.position if isinstance(monitor, PlaneMonitor) else monitor.plane_position
+            for field in normalized_fields:
+                observer = _plane_observer_record(
+                    _internal_observer_name(monitor.name, field),
+                    axis=monitor.axis,
+                    position=monitor_plane_position,
+                    component=field,
+                )
+                observer["monitor_name"] = monitor.name
+                observer["monitor_fields"] = normalized_fields
+                observer["monitor_frequencies"] = monitor.frequencies
+                observer["compute_flux"] = bool(getattr(monitor, "compute_flux", False))
+                observer["normal_direction"] = getattr(monitor, "normal_direction", "+")
+                if isinstance(monitor, ModeMonitor):
+                    observer["monitor_type"] = "mode"
+                    observer["mode_spec"] = monitor.mode_spec()
+                elif isinstance(monitor, DiffractionMonitor):
+                    observer["monitor_type"] = "diffraction"
+                    observer["diffraction_spec"] = monitor.diffraction_spec()
+                compiled.append(observer)
+            continue
+        raise ValueError(f"Unsupported monitor type: {type(monitor).__name__}")
+    return compiled
+
+
+def compile_fdtd_time_observers(scene):
+    compiled = []
+    monitors = scene.resolved_monitors() if hasattr(scene, "resolved_monitors") else scene.monitors
+    for monitor in monitors:
+        if isinstance(monitor, FieldTimeMonitor):
+            record = {
+                "name": monitor.name,
+                "kind": "field_time",
+                "start": int(monitor.start),
+                "stop": None if monitor.stop is None else int(monitor.stop),
+                "interval": int(monitor.interval),
+                "region_kind": monitor.region_kind,
+                "components": tuple(normalize_component(field).capitalize() for field in monitor.components),
+                "position": tuple(float(value) for value in monitor.position),
+                "size": tuple(float(value) for value in monitor.size),
+            }
+            if monitor.region_kind == "plane":
+                record["axis"] = normalize_axis(monitor.axis)
+                record["plane_position"] = float(monitor.plane_position)
+            compiled.append(record)
+            continue
+        if isinstance(monitor, FluxTimeMonitor):
+            compiled.append(
+                {
+                    "name": monitor.name,
+                    "kind": "flux_time",
+                    "start": int(monitor.start),
+                    "stop": None if monitor.stop is None else int(monitor.stop),
+                    "interval": int(monitor.interval),
+                    "axis": normalize_axis(monitor.axis),
+                    "position": float(monitor.position),
+                    "fields": tuple(normalize_component(field).capitalize() for field in monitor.fields),
+                    "normal_direction": monitor.normal_direction,
+                }
+            )
+            continue
+    return compiled

witwin/maxwell/compiler/sources.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import numpy as np
+
+from ..sources import (
+    AstigmaticGaussianBeam,
+    CustomCurrentSource,
+    CustomFieldSource,
+    GaussianBeam,
+    ModeSource,
+    PlaneWave,
+    PointDipole,
+    TFSF,
+    UniformCurrentSource,
+    compile_source_time,
+)
+
+
+def _compile_injection(injection) -> dict[str, object]:
+    if isinstance(injection, TFSF):
+        if injection.mode == "slab":
+            return {
+                "kind": "tfsf",
+                "mode": "slab",
+                "axis": injection.axis,
+                "axis_bounds": injection.axis_bounds,
+            }
+        return {
+            "kind": "tfsf",
+            "mode": "box",
+            "bounds": injection.bounds,
+        }
+    return {"kind": str(injection).lower()}
+
+
+def _compile_fdfd_point_dipole(source: PointDipole, *, default_frequency: float) -> dict:
+    source_time = compile_source_time(source.source_time, default_frequency=default_frequency)
+    if source_time["kind"] != "cw":
+        raise ValueError("FDFD sources currently support CW source_time only.")
+    phase = np.exp(1j * float(source_time["phase"]))
+    amplitude = float(source_time["amplitude"])
+    polarization = tuple(amplitude * phase * component for component in source.polarization)
+    return {
+        "name": source.name,
+        "position": source.position,
+        "width": source.width,
+        "profile": source.profile,
+        "polarization": polarization,
+    }
+
+
+def compile_fdfd_sources(scene, *, default_frequency: float):
+    compiled = []
+    for source in scene.sources:
+        if not isinstance(source, PointDipole):
+            raise ValueError(f"Unsupported FDFD source type: {type(source).__name__}")
+        compiled.append(_compile_fdfd_point_dipole(source, default_frequency=default_frequency))
+    return compiled
+
+
+def _compile_point_dipole(source: PointDipole, *, default_frequency: float) -> dict:
+    return {
+        "kind": "point_dipole",
+        "name": source.name,
+        "position": source.position,
+        "width": source.width,
+        "profile": source.profile,
+        "polarization": source.polarization,
+        "source_time": compile_source_time(source.source_time, default_frequency=default_frequency),
+    }
+
+
+def _reject_custom_source_time(source_time: dict, *, source_kind: str) -> dict:
+    if source_time["kind"] == "custom":
+        raise ValueError(
+            f"{source_kind} does not support CustomSourceTime; it requires the native time-shifted "
+            "injection kernel. Use PointDipole for arbitrary custom waveforms."
+        )
+    return source_time
+
+
+def _compile_plane_wave(source: PlaneWave, *, default_frequency: float) -> dict:
+    return {
+        "kind": "plane_wave",
+        "name": source.name,
+        "direction": source.direction,
+        "polarization": source.polarization,
+        "injection": _compile_injection(source.injection),
+        "injection_axis": source.injection_axis,
+        "source_time": _reject_custom_source_time(
+            compile_source_time(source.source_time, default_frequency=default_frequency),
+            source_kind="PlaneWave",
+        ),
+    }
+
+
+def _compile_gaussian_beam(source: GaussianBeam, *, default_frequency: float) -> dict:
+    return {
+        "kind": "gaussian_beam",
+        "name": source.name,
+        "direction": source.direction,
+        "polarization": source.polarization,
+        "beam_waist": source.beam_waist,
+        "focus": source.focus,
+        "injection": _compile_injection(source.injection),
+        "injection_axis": source.injection_axis,
+        "source_time": _reject_custom_source_time(
+            compile_source_time(source.source_time, default_frequency=default_frequency),
+            source_kind="GaussianBeam",
+        ),
+    }
+
+
+def _compile_astigmatic_gaussian_beam(source: AstigmaticGaussianBeam, *, default_frequency: float) -> dict:
+    return {
+        "kind": "astigmatic_gaussian_beam",
+        "name": source.name,
+        "direction": source.direction,
+        "polarization": source.polarization,
+        "beam_waist_u": source.beam_waist[0],
+        "beam_waist_v": source.beam_waist[1],
+        "focus": source.focus,
+        "focus_u": source.focus_u,
+        "focus_v": source.focus_v,
+        "injection": _compile_injection(source.injection),
+        "injection_axis": source.injection_axis,
+        "source_time": _reject_custom_source_time(
+            compile_source_time(source.source_time, default_frequency=default_frequency),
+            source_kind="AstigmaticGaussianBeam",
+        ),
+    }
+
+
+def _compile_uniform_current_source(source: UniformCurrentSource, *, default_frequency: float) -> dict:
+    return {
+        "kind": "uniform_current",
+        "name": source.name,
+        "center": source.center,
+        "size": source.size,
+        "polarization": source.polarization,
+        "source_time": compile_source_time(source.source_time, default_frequency=default_frequency),
+    }
+
+
+def _compile_custom_current_source(source: CustomCurrentSource, *, default_frequency: float) -> dict:
+    return {
+        "kind": "custom_current",
+        "name": source.name,
+        "dataset": source.current_dataset,
+        "source_time": compile_source_time(source.source_time, default_frequency=default_frequency),
+    }
+
+
+def _compile_custom_field_source(source: CustomFieldSource, *, default_frequency: float) -> dict:
+    return {
+        "kind": "custom_field",
+        "name": source.name,
+        "dataset": source.field_dataset,
+        "normal_axis": source.normal_axis,
+        "source_time": compile_source_time(source.source_time, default_frequency=default_frequency),
+    }
+
+
+def _compile_mode_source(source: ModeSource, *, default_frequency: float) -> dict:
+    # Broadband injection: the guided profile and effective index are solved once
+    # at the waveform's center frequency and driven by the temporal envelope, so a
+    # GaussianPulse/Ricker mode source excites the whole band from one run. Only
+    # CustomSourceTime is rejected because the native time-shifted surface kernel
+    # evaluates the analytic CW/Gaussian/Ricker forms on-device and has no table path.
+    source_time = _reject_custom_source_time(
+        compile_source_time(source.source_time, default_frequency=default_frequency),
+        source_kind="ModeSource",
+    )
+
+    direction_sign = 1 if source.direction == "+" else -1
+    direction_vector = {
+        "x": (float(direction_sign), 0.0, 0.0),
+        "y": (0.0, float(direction_sign), 0.0),
+        "z": (0.0, 0.0, float(direction_sign)),
+    }[source.normal_axis]
+    return {
+        "kind": "mode_source",
+        "name": source.name,
+        "position": source.position,
+        "size": source.size,
+        "mode_index": int(source.mode_index),
+        "direction": source.direction,
+        "direction_sign": int(direction_sign),
+        "direction_vector": direction_vector,
+        "normal_axis": source.normal_axis,
+        "polarization": source.polarization,
+        "polarization_axis": source.polarization_axis,
+        "bend_radius": source.bend_radius,
+        "bend_axis": source.bend_axis,
+        "injection": {"kind": "soft"},
+        "source_time": source_time,
+    }
+
+
+def compile_fdtd_sources(scene, *, default_frequency: float):
+    compiled = []
+    sources = scene.resolved_sources() if hasattr(scene, "resolved_sources") else scene.sources
+    for source in sources:
+        if isinstance(source, PointDipole):
+            compiled.append(_compile_point_dipole(source, default_frequency=default_frequency))
+        elif isinstance(source, PlaneWave):
+            compiled.append(_compile_plane_wave(source, default_frequency=default_frequency))
+        elif isinstance(source, AstigmaticGaussianBeam):
+            compiled.append(_compile_astigmatic_gaussian_beam(source, default_frequency=default_frequency))
+        elif isinstance(source, GaussianBeam):
+            compiled.append(_compile_gaussian_beam(source, default_frequency=default_frequency))
+        elif isinstance(source, ModeSource):
+            compiled.append(_compile_mode_source(source, default_frequency=default_frequency))
+        elif isinstance(source, UniformCurrentSource):
+            compiled.append(_compile_uniform_current_source(source, default_frequency=default_frequency))
+        elif isinstance(source, CustomCurrentSource):
+            compiled.append(_compile_custom_current_source(source, default_frequency=default_frequency))
+        elif isinstance(source, CustomFieldSource):
+            compiled.append(_compile_custom_field_source(source, default_frequency=default_frequency))
+        else:
+            raise ValueError(f"Unsupported FDTD source type: {type(source).__name__}")
+    return compiled

witwin/maxwell/fdfd/__init__.py

@@ -0,0 +1,3 @@
+from .solver import FDFD
+
+__all__ = ["FDFD"]

witwin/maxwell/fdfd/adjoint/__init__.py

@@ -0,0 +1,3 @@
+from .bridge import run_fdfd_with_gradient_bridge
+
+__all__ = ["run_fdfd_with_gradient_bridge"]

witwin/maxwell/fdfd/adjoint/bridge.py

@@ -0,0 +1,180 @@
+"""Differentiable FDFD via the adjoint method.
+
+For the linear system A(eps) x = b with real loss L(x), the adjoint solve
+A^T lam = conj(dL/dx*) yields material gradients without differentiating
+through the solver iterations. The permittivity enters A only on the
+diagonal (k0^2 * eps_face), so
+
+    dL/deps_face = -2 k0^2 Re[lam * x]   (per Yee face)
+
+which is pulled back to trainable scene inputs through the shared material
+compiler graph, exactly like the FDTD gradient bridge.
+"""
+
+from __future__ import annotations
+
+import cupy as cp
+import torch
+
+from ...adjoint_inputs import material_dependent_inputs, scene_trainable_material_tensors
+from ...fdtd.material_pullback import pullback_material_input_gradients
+from ...result import Result
+
+# The FDFD matrix uses relative permittivity directly, so the pullback's
+# absolute-permittivity scaling is the identity.
+_RELATIVE_EPS_SCALE = 1.0
+
+
+def _unsupported_fdfd_adjoint_medium(scene):
+    for structure in getattr(scene, "structures", ()):
+        material = getattr(structure, "material", None)
+        if material is None:
+            continue
+        if getattr(material, "is_medium2d", False):
+            return "FDFD adjoint does not support 2D sheet (Medium2D) media yet."
+        if getattr(material, "is_anisotropic", False):
+            return "FDFD adjoint does not support anisotropic media yet."
+        if getattr(material, "is_dispersive", False):
+            return "FDFD adjoint does not support dispersive media yet."
+    return None
+
+
+class _FDFDGradientBridge:
+    def __init__(self, simulation):
+        self.simulation = simulation
+        self.base_scene = simulation.scene
+        self.material_inputs = self._resolve_material_inputs()
+        if not self.material_inputs:
+            raise NotImplementedError(
+                "FDFD backward currently supports trainable scene inputs that contribute to "
+                "prepared-scene material tensors."
+            )
+        self.solver = None
+        self.solver_stats = None
+
+    def _material_graph_scene(self):
+        if self.simulation.scene_module is not None:
+            return self.simulation.scene_module.to_scene()
+        self.simulation._refresh_scene()
+        return self.simulation.scene
+
+    def _candidate_material_inputs(self) -> tuple[torch.Tensor, ...]:
+        if self.simulation.scene_module is not None:
+            return tuple(
+                parameter
+                for parameter in self.simulation.scene_module.parameters()
+                if parameter.requires_grad
+            )
+        return scene_trainable_material_tensors(self.base_scene)
+
+    def _resolve_material_inputs(self) -> tuple[torch.Tensor, ...]:
+        scene = self._material_graph_scene()
+        return material_dependent_inputs(scene, self._candidate_material_inputs())
+
+    def forward(self, _material_inputs):
+        solver = self.simulation._build_fdfd_solver()
+        unsupported = _unsupported_fdfd_adjoint_medium(solver.scene)
+        if unsupported is not None:
+            raise NotImplementedError(unsupported)
+
+        solver_cfg = self.simulation.config.solver
+        solver.solve(
+            max_iter=solver_cfg.max_iter,
+            tol=solver_cfg.tol,
+            restart=solver_cfg.restart,
+        )
+        if solver.E_field is None:
+            raise RuntimeError("FDFD solve did not produce any field data.")
+
+        self.solver = solver
+        self.solver_stats = {
+            "solver": {
+                "type": solver_cfg.solver_type,
+                "max_iter": solver_cfg.max_iter,
+                "tol": solver_cfg.tol,
+                "restart": solver_cfg.restart,
+                "preconditioner": solver_cfg.preconditioner,
+                "precision": solver_cfg.precision,
+            },
+            "converged": getattr(solver, "converged", None),
+            "solver_info": getattr(solver, "solver_info", None),
+            "final_residual": getattr(solver, "final_residual", None),
+        }
+        return solver.E_field
+
+    def backward(self, material_inputs, grad_outputs):
+        solver = self.solver
+        scene = solver.scene
+        solver_cfg = self.simulation.config.solver
+
+        resolved = tuple(
+            torch.zeros_like(field) if grad is None else grad
+            for field, grad in zip(solver.E_field, grad_outputs)
+        )
+        # For real loss L, PyTorch hands us g = 2 dL/dx* (real-pair gradient
+        # packed as complex). Solving A^T lam = conj(g) gives lam = 2 A^-T
+        # dL/dx, and dL/deps_face = -2 k0^2 Re[(lam/2) x] = -k0^2 Re[lam x].
+        grad_x = torch.cat([grad.reshape(-1) for grad in resolved]).to(torch.complex64)
+        rhs = cp.from_dlpack(grad_x.conj().resolve_conj().contiguous()).astype(cp.complex64, copy=False)
+        lam = solver.solve_adjoint(
+            rhs,
+            max_iter=solver_cfg.max_iter,
+            tol=solver_cfg.tol,
+            restart=solver_cfg.restart,
+        )
+
+        x = cp.from_dlpack(
+            torch.cat([field.reshape(-1) for field in solver.E_field]).contiguous()
+        )
+        grad_eps_flat = (-solver.k0**2) * (lam * x).real.astype(cp.float32, copy=False)
+
+        N_ex, N_ey = scene.N_ex, scene.N_ey
+        grad_eps_ex = torch.from_dlpack(
+            grad_eps_flat[:N_ex].reshape(scene.Nx_ex, scene.Ny_ex, scene.Nz_ex)
+        )
+        grad_eps_ey = torch.from_dlpack(
+            grad_eps_flat[N_ex:N_ex + N_ey].reshape(scene.Nx_ey, scene.Ny_ey, scene.Nz_ey)
+        )
+        grad_eps_ez = torch.from_dlpack(
+            grad_eps_flat[N_ex + N_ey:].reshape(scene.Nx_ez, scene.Ny_ez, scene.Nz_ez)
+        )
+
+        with torch.enable_grad():
+            graph_scene = self._material_graph_scene()
+            return pullback_material_input_gradients(
+                graph_scene,
+                inputs=material_inputs,
+                grad_eps_ex=grad_eps_ex,
+                grad_eps_ey=grad_eps_ey,
+                grad_eps_ez=grad_eps_ez,
+                eps0=_RELATIVE_EPS_SCALE,
+            )
+
+
+class _FDFDMaterialGradientFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, bridge, *material_inputs):
+        ctx.bridge = bridge
+        ctx.material_inputs = tuple(material_inputs)
+        return bridge.forward(tuple(material_inputs))
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        gradients = ctx.bridge.backward(ctx.material_inputs, grad_outputs)
+        return (None, *gradients)
+
+
+def run_fdfd_with_gradient_bridge(simulation) -> Result:
+    bridge = _FDFDGradientBridge(simulation)
+    Ex, Ey, Ez = _FDFDMaterialGradientFunction.apply(bridge, *bridge.material_inputs)
+    return Result(
+        method="fdfd",
+        scene=simulation.scene,
+        prepared_scene=bridge.solver.scene,
+        frequency=simulation.frequency,
+        frequencies=simulation.frequencies,
+        solver=bridge.solver,
+        fields={"EX": Ex, "EY": Ey, "EZ": Ez},
+        metadata=simulation.metadata,
+        solver_stats=bridge.solver_stats,
+    )