From e3ec0c4032d8c9904959c5a91438010f8ef540b2 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 27 Aug 2026 09:41:35 -0700 Subject: [PATCH 1/5] adjoint: gradients with respect to source amplitudes Writing Maxwell's equations as A(rho) E = -i omega J gives dJ_obj/dJ = -i omega lambda, so the derivative with respect to a source is just the adjoint field sampled over that source's support. It needs no simulation beyond the adjoint run already being performed -- only a DFT monitor over the source region. The gathering is dft_fields::fourier_sourcegradient, written as the exact transpose of the fourier_sourcedata scatter that places adjoint sources: same loop, same points, same weights. That is what makes the array ordering automatic, since a cotangent comes back indexed exactly like the array handed in. Two conventions keep it simpler than expected -- chi1inv is a realnum, so the weights are real and no conjugation is involved, and update_dft already averages the four Yee sites that the scatter splits across, which is precisely the transpose of that split. Sources opt in with `differentiable=['currents', 'amplitude']`. The names are validated against a per-class whitelist and become the keys of the resulting gradient, so the two cannot drift apart. 'center' and 'size' are rejected with their own message: they move the grid points a source occupies rather than the amplitudes applied to them, which is outside this formulation rather than merely unimplemented. ArraySource supplies per-point amplitudes directly, indexed like get_dft_array over the same region -- the same ordering the gradient comes back in, because injection and measurement are a scatter and its transpose. Verified against central finite differences: 2.7e-7 for scalar amplitudes across frequencies, resolutions, complex amplitudes and both electric and magnetic sources, and 8.3e-8 for per-point currents. The normalization was pinned empirically rather than derived; it is adj_src_phase * dtft_forward / (dV * i omega), and _adj_src_phase is now exposed on ObjectiveQuantity so that placing an adjoint source and differentiating a source cannot drift. --- python/Makefile.am | 1 + python/adjoint/__init__.py | 2 + python/adjoint/objective.py | 50 +++-- python/adjoint/optimization_problem.py | 58 ++++++ python/adjoint/source_gradient.py | 262 +++++++++++++++++++++++++ python/adjoint/wrapper.py | 12 ++ python/meep.i | 4 +- python/simulation.py | 1 + python/source.py | 202 ++++++++++++++++++- src/dft.cpp | 85 ++++++++ src/meep.hpp | 3 + 11 files changed, 662 insertions(+), 18 deletions(-) create mode 100644 python/adjoint/source_gradient.py diff --git a/python/Makefile.am b/python/Makefile.am index e4f205fa9..6a961fbdb 100644 --- a/python/Makefile.am +++ b/python/Makefile.am @@ -250,6 +250,7 @@ adjoint_PYTHON = $(srcdir)/adjoint/__init__.py \ $(srcdir)/adjoint/filter_source.py \ $(srcdir)/adjoint/connectivity.py \ $(srcdir)/adjoint/unfilter_design.py \ + $(srcdir)/adjoint/source_gradient.py \ $(srcdir)/adjoint/wrapper.py \ $(srcdir)/adjoint/utils.py diff --git a/python/adjoint/__init__.py b/python/adjoint/__init__.py index a6a083274..16c1b3fdd 100644 --- a/python/adjoint/__init__.py +++ b/python/adjoint/__init__.py @@ -20,6 +20,8 @@ from .unfilter_design import * +from . import source_gradient + # JAX is an optional dependency; everything that needs it lives in `wrapper`. # Importing it also registers JAX as a way to differentiate objective functions, # so objective functions written with `jax.numpy` need no special treatment. diff --git a/python/adjoint/objective.py b/python/adjoint/objective.py index e51f7260a..b5dfaf1f2 100644 --- a/python/adjoint/objective.py +++ b/python/adjoint/objective.py @@ -115,6 +115,39 @@ def get_evaluation(self): "evaluation of an objective quantity." ) + def _adj_src_phase(self, src=None, y=None, dt=None): + """The unit-modulus phase `_adj_src_scale` divides out of an adjoint source. + + Placing an adjoint source and taking a gradient with respect to a + *source* amplitude are transposes of one another, so the source + gradient has to multiply this factor back in. It is exposed separately + so that the two cannot drift apart. + """ + if dt is None: + dt = self.sim.fields.dt + if src is None: + src = self._create_time_profile() + if y is None: + T = self.sim.meep_time() + y = np.array([src.swigobj.current(t, dt) for t in np.arange(0, T, dt)]) + + src_center_dtft = ( + np.matmul( + np.exp( + 1j + * 2 + * np.pi + * np.array([src.frequency])[:, np.newaxis] + * np.arange(y.size) + * dt + ), + y, + ) + * dt + / np.sqrt(2 * np.pi) + ) + return np.exp(1j * np.angle(src_center_dtft)) * self.fwidth_scale + def _adj_src_scale(self, include_resolution=True): """Calculates the scale for the adjoint sources.""" T = self.sim.meep_time() @@ -157,22 +190,7 @@ def _adj_src_scale(self, include_resolution=True): # # Note: for some reason, there seems to be an additional phase factor at # the center frequency that needs to be applied to *all* frequencies... - src_center_dtft = ( - np.matmul( - np.exp( - 1j - * 2 - * np.pi - * np.array([src.frequency])[:, np.newaxis] - * np.arange(y.size) - * dt - ), - y, - ) - * dt - / np.sqrt(2 * np.pi) - ) - adj_src_phase = np.exp(1j * np.angle(src_center_dtft)) * self.fwidth_scale + adj_src_phase = self._adj_src_phase(src=src, y=y, dt=dt) if self._frequencies.size == 1: # Single-frequency simulations. Requires a time profile. diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index 3815e0bb0..88e94ffc3 100644 --- a/python/adjoint/optimization_problem.py +++ b/python/adjoint/optimization_problem.py @@ -5,6 +5,7 @@ import meep as mp from . import LDOS, DesignRegion, utils, ObjectiveQuantity +from . import source_gradient class OptimizationProblem: @@ -146,6 +147,13 @@ def __init__( # store sources for finite difference estimations self.forward_sources = self.sim.sources + # Sources flagged with `differentiable=[...]` are a second category of + # differentiable input alongside the design regions. Their gradient is + # the adjoint field sampled over the source's own support, so it needs + # no simulation beyond the adjoint run already being performed. + self.differentiable_sources = source_gradient.differentiable_sources(self.sim) + self.source_gradient = {} + # The optimizer has three allowable states : "INIT", "FWD", and "ADJ". # INIT - The optimizer is initialized and ready to run a forward simulation # FWD - The optimizer has already run a forward simulation @@ -219,6 +227,11 @@ def __call__( f"Incorrect solver state detected: {self.current_state}" ) + if self.differentiable_sources: + # Only change the return shape when the user asked for source + # gradients; without a flagged source this is exactly as before. + return self.f0, {"design": self.gradient, **self.source_gradient} + return self.f0, self.gradient def get_fdf_funcs(self) -> Tuple[Callable, Callable]: @@ -347,6 +360,7 @@ def adjoint_run(self): self.sim.change_k_point(-1 * self.sim.k_point) self.adjoint_design_region_monitors = [] + self.adjoint_source_monitors = [] for ar in range(len(self.objective_functions)): # Reset the fields self.sim.restart_fields() @@ -364,6 +378,17 @@ def adjoint_run(self): self.decimation_factor, ) ) + + # register a monitor over each differentiable source's support; the + # adjoint field there is the gradient with respect to its currents + self.adjoint_source_monitors.append( + source_gradient.install_source_gradient_monitors( + self.sim, + self.differentiable_sources, + self.frequencies, + self.decimation_factor, + ) + ) self.sim._evaluate_dft_objects() # Adjoint run @@ -385,7 +410,40 @@ def adjoint_run(self): # update optimizer's state self.current_state = "ADJ" + def calculate_source_gradient(self): + """Gather the adjoint field over each differentiable source's support. + + Returns a dict keyed by source name (or index), each holding a dict + keyed by the parameter names the source declared in `differentiable`. + """ + if not self.differentiable_sources: + return {} + + out = {} + for ar in range(len(self.objective_functions)): + for si, src in enumerate(self.differentiable_sources): + monitor = self.adjoint_source_monitors[ar][si] + # the adjoint field carries the normalization the objective + # quantity applied when it placed the adjoint source, so the + # transpose has to use that same quantity's phase + scale = source_gradient.source_grad_scale( + self.sim, + src, + self.frequencies, + self.objective_arguments[0]._adj_src_phase(), + ) + currents = monitor.gather(self.sim, scale) + grads = source_gradient.contract(src, currents) + key = source_gradient.source_key(src, si) + if len(self.objective_functions) == 1: + out[key] = grads + else: + out.setdefault(key, []).append(grads) + return out + def calculate_gradient(self): + self.source_gradient = self.calculate_source_gradient() + # Iterate through all design regions and calculate gradient self.gradient = [ [ diff --git a/python/adjoint/source_gradient.py b/python/adjoint/source_gradient.py new file mode 100644 index 000000000..06f0dfb69 --- /dev/null +++ b/python/adjoint/source_gradient.py @@ -0,0 +1,262 @@ +"""Adjoint gradients with respect to source amplitudes. + +Meep's design gradient contracts the forward and adjoint fields inside a design +region. The source gradient is simpler: writing Maxwell's equations as +`A(rho) E = -i omega J`, we have `dE/dJ = -i omega A^-1`, so + + dJ_obj/dJ = (dJ_obj/dE) (dE/dJ) = -i omega lambda + +which is just the adjoint field sampled where the source sits. No additional +simulation is needed -- only a DFT monitor over the source's support during the +adjoint run that is already being performed. + +The gathering is done by `dft_fields::fourier_sourcegradient`, the transpose of +the `fourier_sourcedata` scatter that places adjoint sources. Writing the two +as an exact transpose pair is what makes the array ordering automatic: the +cotangent comes back indexed exactly like the array that would be handed in. + +Meep's contract stops at the discrete current amplitudes. Turning `dJ_obj/dA` +into a derivative with respect to a beam waist or a density is the caller's +job -- either JAX's, via `MeepJaxWrapper`, or the user's. +""" + +from typing import List, Optional + +import numpy as np + +import meep as mp + +# Parameters this module can currently evaluate. `source.differentiable` is +# validated against a wider list at construction time; see +# `meep.source._validate_differentiable`. +IMPLEMENTED_PARAMS = ("currents", "amplitude") + + +def differentiable_sources(sim: mp.Simulation) -> List: + """Return the sources in `sim` that were flagged for differentiation.""" + return [s for s in sim.sources if getattr(s, "differentiable", ())] + + +def source_key(source, index: int): + """The key a source's gradient appears under. + + Uses the source's `name` when it has one, and its position in + `Simulation.sources` otherwise. Names are preferred because they survive + the source list being reordered, and because they are meaningful across + processes. + """ + name = getattr(source, "name", None) + return name if name is not None else index + + +class SourceGradientMonitor: + """A DFT monitor over one differentiable source's support. + + Installed during the adjoint run. The adjoint field it records *is* the + gradient with respect to the source's currents, up to the per-frequency + scale factor supplied by the caller. + """ + + def __init__( + self, + sim: mp.Simulation, + source, + frequencies: np.ndarray, + decimation_factor: Optional[int] = 0, + ): + self.source = source + self.component = source.component + self._frequencies = np.asarray(frequencies) + + if getattr(source, "amp_func", None) is not None or getattr( + source, "amp_func_file", "" + ): + # An amp_func is evaluated inside Meep at each grid point, so the + # discrete amplitudes it produces are not visible here. The + # currents cotangent would be correct but the user has no array to + # apply it to, and `amplitude` would silently ignore the profile. + raise NotImplementedError( + "Differentiating a source defined by `amp_func` or " + "`amp_func_file` is not supported; supply the amplitudes as an " + "array (`amp_data`) or drive the source from JAX instead." + ) + + self.volume = sim._fit_volume_to_simulation( + mp.Volume(center=source.center, size=source.size) + ) + self._check_not_in_pml(sim) + self.decimation_factor = decimation_factor + self._monitor = None + + def _check_not_in_pml(self, sim: mp.Simulation) -> None: + """Refuse a source whose adjoint field would be absorbed. + + Inside a PML the adjoint field is damped, so the gradient would come + back finite, smooth, and wrong. That is worth an error rather than a + surprise. + """ + if not sim.boundary_layers: + return + thickness = max( + (getattr(bl, "thickness", 0.0) for bl in sim.boundary_layers), default=0.0 + ) + if thickness <= 0: + return + + half_cell = np.array([sim.cell_size.x, sim.cell_size.y, sim.cell_size.z]) / 2 + center = np.array([self.volume.center.x, self.volume.center.y, self.volume.center.z]) + half_size = np.array([self.volume.size.x, self.volume.size.y, self.volume.size.z]) / 2 + + for i, dim in enumerate("xyz"): + if half_cell[i] == 0: + continue # not a simulated direction + overshoot = (np.abs(center[i]) + half_size[i]) - (half_cell[i] - thickness) + if overshoot > 0: + raise ValueError( + f"Differentiable source extends {overshoot:.4g} into the PML " + f"along {dim}. The adjoint field is absorbed there, so the " + "gradient would be wrong without reporting an error. Move " + "the source away from the boundary." + ) + + def register(self, sim: mp.Simulation) -> None: + """Install the DFT monitor. Called once per adjoint run.""" + # yee_grid=True so the monitor samples exactly the points a volume + # source drives, rather than voxel centers. + self._monitor = sim.add_dft_fields( + [self.component], + self._frequencies, + where=self.volume, + yee_grid=True, + decimation_factor=self.decimation_factor, + ) + + def shape(self, sim: mp.Simulation): + """(nfreq,) + the monitor's spatial shape, with trailing 1s dropped.""" + dims = sim.fields.dft_monitor_size( + self._monitor.swigobj, self.volume.swigobj, self.component + ) + dims = [d for d in dims if d > 1] or [1] + return (len(self._frequencies), *dims) + + def gather(self, sim: mp.Simulation, scale: np.ndarray) -> np.ndarray: + """Return dJ_obj/d(currents), summed over all processes. + + `scale` is the per-frequency factor relating a current amplitude to the + adjoint field, broadcast over the spatial axes. + """ + dims = sim.fields.dft_monitor_size( + self._monitor.swigobj, self.volume.swigobj, self.component + ) + num_points = int(np.prod(dims)) + grad = np.zeros(num_points * len(self._frequencies), dtype=np.complex128) + + self._monitor.swigobj.fourier_sourcegradient( + self.volume.swigobj, self.component, sim.fields, grad + ) + + grad = grad.reshape(len(self._frequencies), num_points) + grad *= np.asarray(scale).reshape(-1, 1) + return grad.reshape(self.shape(sim)) + + +def install_source_gradient_monitors( + sim: mp.Simulation, + sources: List, + frequencies: np.ndarray, + decimation_factor: Optional[int] = 0, +) -> List[SourceGradientMonitor]: + """Install a DFT monitor over each differentiable source's support.""" + monitors = [ + SourceGradientMonitor(sim, s, frequencies, decimation_factor) for s in sources + ] + for m in monitors: + m.register(sim) + return monitors + + +def time_profile_dtft( + sim: mp.Simulation, src_time, frequencies: np.ndarray +) -> np.ndarray: + """DTFT of a source's time profile, in the convention `_adj_src_scale` uses. + + Meep relates a source's *amplitude* to the field it produces through the + Fourier transform of its time profile, so converting the adjoint field into + a derivative with respect to that amplitude requires this factor. + """ + dt = sim.fields.dt + T = sim.meep_time() + y = np.array([src_time.swigobj.current(t, dt) for t in np.arange(0, T, dt)]) + freqs = np.asarray(frequencies) + phase = np.exp(1j * 2 * np.pi * freqs[:, np.newaxis] * np.arange(y.size) * dt) + return (phase @ y) * dt / np.sqrt(2 * np.pi) + + +def source_grad_scale( + sim: mp.Simulation, + source, + frequencies: np.ndarray, + adj_src_phase: np.ndarray, +) -> np.ndarray: + """Per-frequency factor relating the adjoint field to dJ_obj/d(amplitude). + + This is the transpose of `ObjectiveQuantity._adj_src_scale`, which maps a + field cotangent onto an adjoint source amplitude. Going the other way, the + adjoint field at the source is multiplied by the *forward* source's own + transform, divided by the volume element and the discrete-time `i*omega`, + and multiplied by the same unit-modulus phase the adjoint source divided + out. + + The returned gradient follows the convention JAX and autograd use for a + real function of a complex input, namely + `g = dJ/d(Re a) - i dJ/d(Im a)`, so that it chains without adjustment. + """ + frequencies = np.asarray(frequencies) + dt = sim.fields.dt + + # discrete-time derivative, matching _adj_src_scale + iomega = (1.0 - np.exp(-1j * (2 * np.pi * frequencies) * dt)) * (1.0 / dt) + + num_dims = sim._infer_dimensions(sim.k_point) + dV = 1 / sim.resolution**num_dims + + fwd_dtft = time_profile_dtft(sim, source.src, frequencies) + + scale = np.asarray(adj_src_phase) * fwd_dtft / (dV * iomega) + + if sim.using_real_fields(): + # real fields keep only Re[J], halving the amplitude at +omega + scale *= 2 + return scale + + +def contract(source, currents_grad: np.ndarray, source_amplitudes=None) -> dict: + """Contract the currents cotangent onto the source's declared parameters. + + Every requested parameter is a linear function of the currents, so no + finite difference over Meep's source construction is involved here. The + nonlinear built-in parameters (`beam_w0` and friends) are rejected at + construction time; see `meep.source._validate_differentiable`. + """ + out = {} + for name in source.differentiable: + if name == "currents": + out[name] = currents_grad + elif name == "amplitude": + # `amplitude` scales every point identically, so its derivative is + # the sum of the currents cotangent against the unit profile. Exact, + # with no finite difference. + if source_amplitudes is None: + profile = np.ones(currents_grad.shape[1:]) + else: + profile = np.asarray(source_amplitudes) / source.amplitude + out[name] = np.sum( + currents_grad * np.conj(profile)[np.newaxis, ...], + axis=tuple(range(1, currents_grad.ndim)), + ) + else: # pragma: no cover - blocked by _validate_differentiable + raise NotImplementedError( + f"'{name}' is not implemented in this version of the source " + "adjoint; see meep.source._validate_differentiable." + ) + return out diff --git a/python/adjoint/wrapper.py b/python/adjoint/wrapper.py index d0172171e..597c80e0c 100644 --- a/python/adjoint/wrapper.py +++ b/python/adjoint/wrapper.py @@ -212,6 +212,18 @@ def __init__( self.until_after_sources = until_after_sources self.finite_difference_step = finite_difference_step + # Sources whose amplitudes are supplied from JAX are differentiated + # with respect to their currents automatically: the parameters live + # upstream in JAX, so there is nothing for Meep to name. + self.differentiable_sources = [ + s for s in sources if isinstance(s, mp.ArraySource) + ] + for s in self.differentiable_sources: + if "currents" not in getattr(s, "differentiable", ()): + s.differentiable = tuple(getattr(s, "differentiable", ())) + ( + "currents", + ) + self._simulate_fn = self._initialize_callable() def __call__( diff --git a/python/meep.i b/python/meep.i index 9c917969d..dc4ef7948 100644 --- a/python/meep.i +++ b/python/meep.i @@ -1084,7 +1084,8 @@ void _get_gradient(PyObject *grad, double scalegrad, } %apply std::complex* grid_vals { std::complex* eigfreq, std::complex* coeffs, - std::complex* dJ, std::complex* amp_arr + std::complex* dJ, std::complex* amp_arr, + std::complex* grad }; // typemaps for diffractedplanewave @@ -1776,6 +1777,7 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where with_prefix ) from .source import ( + ArraySource, ContinuousSource, CustomSource, EigenModeSource, diff --git a/python/simulation.py b/python/simulation.py index af75cf30d..0760467a2 100644 --- a/python/simulation.py +++ b/python/simulation.py @@ -24,6 +24,7 @@ import numpy as np from meep.geom import GeometricObject, Medium, Vector3, init_do_averaging from meep.source import ( + ArraySource, EigenModeSource, GaussianBeamSource, IndexedSource, diff --git a/python/source.py b/python/source.py index b88440c2a..9d17ce107 100644 --- a/python/source.py +++ b/python/source.py @@ -15,6 +15,72 @@ def check_positive(prop, val): raise ValueError(f"{prop} must be positive. Got {val}") +# Parameters every source can be differentiated with respect to. "currents" is +# the cotangent with respect to the per-point complex current amplitudes Meep +# actually applies to the Yee grid; it is the universal representation, and the +# one the JAX bridge uses. "amplitude" is exact and needs no finite difference, +# since it scales those currents linearly. +_DIFFERENTIABLE_ALWAYS = ("currents", "amplitude") + +# Parameters that move the source's support rather than change its amplitudes. +# The adjoint machinery gathers the cotangent at a *fixed* set of grid points +# (see dft_fields::fourier_sourcegradient), so a derivative here is not merely +# unimplemented -- it is outside the formulation. +_DIFFERENTIABLE_MOVES_SUPPORT = ("center", "size", "volume") + + +def _validate_differentiable(src, differentiable): + """Check and normalize a source's `differentiable` argument. + + Returns a tuple of parameter names. Raises `ValueError` for a name the + source does not have, and `NotImplementedError` for a name that is + meaningful but whose sensitivity is not yet implemented, so that the two + cases are not confused with one another. + """ + if differentiable is None: + return () + if isinstance(differentiable, str): + raise ValueError( + "`differentiable` takes a list of parameter names, not a bare " + f"string; use ['{differentiable}'] instead." + ) + + valid = tuple(_DIFFERENTIABLE_ALWAYS) + tuple( + getattr(src, "_differentiable_params", ()) + ) + implemented = set(_DIFFERENTIABLE_ALWAYS) + + names = [] + for name in differentiable: + if name in _DIFFERENTIABLE_MOVES_SUPPORT: + raise ValueError( + f"'{name}' changes which grid points the source occupies, not " + "the amplitudes it applies to them, so its derivative is not " + "defined by the adjoint formulation Meep uses for sources. " + "Parameterize the amplitudes instead, e.g. with 'currents'." + ) + if name not in valid: + raise ValueError( + f"'{name}' is not a differentiable parameter of " + f"{type(src).__name__}. Valid choices are: " + f"{', '.join(sorted(valid))}." + ) + if name not in implemented: + raise NotImplementedError( + f"'{name}' is a differentiable parameter of " + f"{type(src).__name__}, but its sensitivity requires the " + "finite-difference contraction over Meep's own source " + "construction, which is not implemented yet. Use 'currents' " + "and apply the chain rule yourself, or parameterize the " + "source from JAX." + ) + names.append(name) + + if len(set(names)) != len(names): + raise ValueError(f"`differentiable` contains duplicate names: {names}") + return tuple(names) + + class Source: """ The `Source` class is used to specify the current sources via the `Simulation.sources` @@ -49,6 +115,8 @@ def __init__( amp_func=None, amp_func_file="", amp_data=None, + differentiable=None, + name=None, ): """ Construct a `Source`. @@ -98,6 +166,22 @@ def __init__( For a 2d simulation, just pass 1 for the third dimension, e.g., `arr = np.zeros((N, M, 1), dtype=np.complex128)`. Defaults to `None`. + + **`differentiable` [`list of string`]** — Names of the parameters this + source should be differentiated with respect to by the adjoint solver + (see [Adjoint Solver](Python_Tutorials/Adjoint_Solver.md)). Every source + accepts `'currents'`, the per-point complex current amplitudes, and + `'amplitude'`; individual source classes may accept more. The names given + here are exactly the keys of the corresponding gradient, so they cannot + drift apart. `'center'` and `'size'` are rejected: they move the grid + points the source occupies rather than the amplitudes applied to them, + which the adjoint formulation for sources does not cover. Defaults to + `None`, meaning the source is not differentiated. + + + **`name` [`string`]** — An optional label used as the key for this + source's entry in the gradient returned by `OptimizationProblem`. Defaults + to `None`, in which case the source's position in `Simulation.sources` is + used instead. + As described in Section 4.2 ("Incident Fields and Equivalent Currents") in [Chapter 4](http://arxiv.org/abs/arXiv:1301.5366) ("Electromagnetic Wave Source Conditions") of the book [Advances in FDTD Computational Electrodynamics: @@ -131,6 +215,8 @@ def __init__( self.amp_func = amp_func self.amp_func_file = amp_func_file self.amp_data = amp_data + self.name = name + self.differentiable = _validate_differentiable(self, differentiable) def add_source(self, sim): where = mp.Volume( @@ -702,6 +788,11 @@ class GaussianBeam3DSource(Source): The `SourceTime` object (`Source.src`), which specifies the time dependence of the source, should normally be a narrow-band `ContinuousSource` or `GaussianSource`. (For a `CustomSource`, the beam frequency is determined by the source's `center_frequency` parameter. """ + # `beam_kdir` is deliberately absent: its length is ignored, so only its + # direction is meaningful and a component-wise derivative would report a + # spurious radial sensitivity. It needs a tangent-space projection first. + _differentiable_params = ("beam_x0", "beam_w0", "beam_E0") + def __init__( self, src, @@ -1084,15 +1175,124 @@ def add_source(self, sim): super().add_source(sim) +class ArraySource(Source): + """A volume source whose per-point complex amplitudes are given as an array. + + Ordinary `Source` objects take a single `amplitude` and, optionally, an + `amp_func` that Meep evaluates internally. This class instead takes the + amplitudes directly, one per grid point, which is what a source computed + somewhere else -- by a mode solver, a propagator, or JAX -- naturally + produces. + + The array is indexed exactly like `Simulation.get_dft_array` over the same + volume and component. That is deliberate: it is also the ordering the + adjoint solver returns `differentiable=['currents']` gradients in, so an + array can be handed in and its cotangent read back with no bookkeeping in + between. Injection and measurement share one convention because they are + implemented as a scatter and its exact transpose. + """ + + def __init__( + self, + src, + component, + amplitudes, + frequency, + center=None, + volume=None, + size=Vector3(), + differentiable=None, + name=None, + ): + """Construct an `ArraySource`. + + + **`amplitudes` [`numpy.ndarray`]** — Complex amplitude for each point + of the source region, shaped like `get_dft_array` over that region. + + **`frequency` [`number`]** — The frequency the amplitudes refer to. + """ + super().__init__( + src, + component, + center=center, + volume=volume, + size=size, + differentiable=differentiable, + name=name, + ) + self.amplitudes = np.ascontiguousarray(amplitudes, dtype=np.complex128) + self.frequency = float(frequency) + self._monitor = None + + def add_source(self, sim): + vol = sim._fit_volume_to_simulation( + mp.Volume(center=self.center, size=self.size) + ) + # The monitor is created only for its chunk decomposition and array + # ordering; its DFT storage is never read. + # A DFT object cannot be added before the field components exist, and + # components are normally allocated only once every source has been + # added. Ask for this one up front. + sim.fields.require_component(self.component) + + mon = sim.add_dft_fields( + [self.component], [self.frequency], where=vol, yee_grid=True + ) + # add_dft_fields defers construction; the scatter needs it now + sim._evaluate_dft_objects() + dims = sim.fields.dft_monitor_size(mon.swigobj, vol.swigobj, self.component) + npts = int(np.prod(dims)) + + if self.amplitudes.size != npts: + raise ValueError( + f"`amplitudes` has {self.amplitudes.size} elements but the " + f"source region holds {npts} grid points (shape {tuple(dims)})." + ) + + # Two conversions, so that one entry of `amplitudes` means exactly what + # `Source.amplitude` means for a point source at that grid point: + # fourier_sourcedata negates electric components (it was written to + # place adjoint sources), and it places a current *density*, whose + # integral over a voxel is the amplitude times dV. + num_dims = sim._infer_dimensions(sim.k_point) + dV = 1 / sim.resolution**num_dims + flat = np.ascontiguousarray( + self.amplitudes.ravel() * complex(self.amplitude) * (-1.0 / dV), + dtype=np.complex128, + ) + srcdata = mon.swigobj.fourier_sourcedata( + vol.swigobj, self.component, sim.fields, flat + ) + + sim.fields.register_src_time(self.src.swigobj) + for sd in srcdata: + amp = np.asarray(sd.amp_arr, dtype=np.complex128) + if amp.size == 0: + continue # this process owns no part of the source + sim.fields.add_srcdata(sd, self.src.swigobj, amp.size, amp, False) + + # the DFT storage is dead weight for a plane with many points + mon.remove() + + class IndexedSource(Source): """ created a source object using (SWIG-wrapped mp::srcdata*) srcdata. """ - def __init__(self, src, srcdata, amp_arr, needs_boundary_fix=False): + def __init__( + self, + src, + srcdata, + amp_arr, + needs_boundary_fix=False, + differentiable=None, + name=None, + ): self.src = src self.num_pts = len(amp_arr) self.srcdata = srcdata + self.name = name + self.differentiable = _validate_differentiable(self, differentiable) self.amp_arr = amp_arr self.needs_boundary_fix = needs_boundary_fix diff --git a/src/dft.cpp b/src/dft.cpp index 0257586e6..678ecd767 100644 --- a/src/dft.cpp +++ b/src/dft.cpp @@ -1585,4 +1585,89 @@ std::vector dft_fields::fourier_sourcedata(const volume &wher return temp; } +/* Transpose of fourier_sourcedata. + + fourier_sourcedata scatters a user-ordered array dJ onto the grid points of + this monitor, weighting each point by w(x). This routine performs the + adjoint of that operation: it gathers the DFT fields recorded by this + monitor -- normally an adjoint run -- back into the same user ordering, + applying the same w(x). The two walk the same points in the same order, so + that == to roundoff. + + Two conventions make this simpler than it first appears. The weights are + real, because chi1inv is a realnum, so no conjugation is involved. And the + 0.25 four-point split that fourier_sourcedata applies when yee_grid is false + needs no counterpart here: update_dft already averages those same four + sites, which is precisely the transpose of the split. + + `grad` must have room for freq.size() * (monitor size) elements. It is + summed across all processes before returning. */ +void dft_fields::fourier_sourcegradient(const volume &where, component c, fields &f, + std::complex *grad) { + const size_t Nfreq = freq.size(); + + ivec min_corner, max_corner; + int rank, reduced_rank; + direction dirs[3], reduced_dirs[3]; + size_t array_size, bufsz, dims[3], reduced_dims[3], reduced_stride[3], stride[3]; + dft_chunk *chunklists[1]; + chunklists[0] = chunks; + + f.get_dft_component_dims(chunklists, 1, c, min_corner, max_corner, array_size, bufsz, rank, dirs, + dims); + reduce_array_dimensions(where, rank, dims, dirs, stride, reduced_rank, reduced_dims, reduced_dirs, + reduced_stride); + size_t reduced_grid_size = reduced_dims[0] * reduced_dims[1] * reduced_dims[2]; + + std::vector > local(Nfreq * reduced_grid_size, std::complex(0, 0)); + + for (dft_chunk *fdc = chunks; fdc; fdc = fdc->next_in_dft) { + assert(Nfreq == fdc->omega.size()); + vec rshift(fdc->shift * (0.5 * fdc->fc->gv.inva)); + + component cc = component(fdc->c); + direction cd = component_direction(cc); + + int position_array[3] = {0, 0, 0}; // position of a point relative to the monitor's min corner + + LOOP_OVER_IVECS(fdc->fc->gv, fdc->is, fdc->ie, idx) { + IVEC_LOOP_LOC(fdc->fc->gv, x0); + IVEC_LOOP_ILOC(fdc->fc->gv, ix0); + size_t idx_dft = IVEC_LOOP_COUNTER; // how update_dft indexes fdc->dft + x0 = fdc->S.transform(x0, fdc->sn) + rshift; + ix0 = fdc->S.transform(ix0, fdc->sn) + fdc->shift; + + double dJ_weight = 1; // weight for linear interpolation + int nd = 0; + LOOP_OVER_DIRECTIONS(fdc->fc->gv.dim, d) { + if (where.in_direction(d) > 0) + position_array[nd++] = int((ix0.in_direction(d) - min_corner.in_direction(d)) / 2); + else + dJ_weight *= (1 - abs(x0.in_direction(d) - where.in_direction_min(d)) / + (fdc->fc->gv.inva)); // based on distances + } + + // index when the gradient is flattened to a one-dimensional array + size_t idx_1d = (position_array[0] * reduced_dims[1] + position_array[1]) * reduced_dims[2] + + position_array[2]; + + double w = dJ_weight; + if (is_electric(cc)) w *= -1; + if (is_D(cc) && fdc->fc->s->chi1inv[cc - Dx + Ex][cd]) + w /= -fdc->fc->s->chi1inv[cc - Dx + Ex][cd][idx]; + if (is_B(cc) && fdc->fc->s->chi1inv[cc - Bx + Hx][cd]) + w /= fdc->fc->s->chi1inv[cc - Bx + Hx][cd][idx]; + w /= fdc->S.multiplicity(ix0); + + for (size_t i = 0; i < Nfreq; ++i) { + std::complex EH = fdc->dft[Nfreq * idx_dft + i]; + local[reduced_grid_size * i + idx_1d] += + w * std::complex(double(EH.real()), double(EH.imag())); + } + } + } + + sum_to_all(local.data(), grad, int(Nfreq * reduced_grid_size)); +} + } // namespace meep diff --git a/src/meep.hpp b/src/meep.hpp index 911d37a4e..8a661ae32 100644 --- a/src/meep.hpp +++ b/src/meep.hpp @@ -1425,6 +1425,9 @@ class dft_fields { dft_fields(dft_chunk *chunks, const double *freq_, size_t Nfreq, const volume &where); std::vector fourier_sourcedata(const volume &where, component c, fields &f, const std::complex *dJ); + // transpose of fourier_sourcedata; see dft.cpp + void fourier_sourcegradient(const volume &where, component c, fields &f, + std::complex *grad); void scale_dfts(std::complex scale); void remove(); From 6b01582618925d8843cd374a5bb2496b7efd49bf Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 27 Aug 2026 10:08:35 -0700 Subject: [PATCH 2/5] adjoint: JAX support, tests and docs for source gradients MeepJaxWrapper takes a `sources` argument alongside `designs`, and its custom_vjp returns a cotangent for it. An ArraySource handed to the wrapper is differentiated automatically: the parameters that produced the currents live upstream in JAX, so there is nothing for Meep to name. This is the objective-side protocol run backwards -- there Meep returns monitor values and JAX owns the post-processing; here JAX supplies currents and Meep returns their cotangent. The cotangent follows the convention JAX and autograd both use for a real function of a complex input, dJ/d(Re a) - i dJ/d(Im a), so it chains with no adjustment. A convention error there would not raise, so the round trip through jax.value_and_grad is tested against a finite difference: 8.3e-8. test_source_gradient.py covers flag validation, the transpose (both the exact -dft identity for an aligned monitor and the dot-product identity against the scatter), ArraySource against an ordinary Source, and finite-difference checks of both amplitude and per-point currents gradients. Every finite-difference test pins the run length, since an adaptive stop makes a perturbed run end at a different time and turns a correct gradient into a constant-ratio failure that does not shrink with the step. --- NEWS.md | 12 + doc/docs/Python_Tutorials/Adjoint_Solver.md | 116 +++++ python/Makefile.am | 2 + python/adjoint/source_gradient.py | 8 +- python/adjoint/wrapper.py | 97 +++- python/tests/test_source_gradient.py | 462 ++++++++++++++++++++ 6 files changed, 678 insertions(+), 19 deletions(-) create mode 100644 python/tests/test_source_gradient.py diff --git a/NEWS.md b/NEWS.md index a8d273ad5..8dc2cd49e 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,6 +2,18 @@ ## Meep 1.35.0 (in progress) +* Adjoint solver: sources can now be differentiated alongside the design + regions. A source opts in with `differentiable=['amplitude']` or + `differentiable=['currents']`, and its gradient appears in the returned + dictionary under the source's `name`. It costs no extra simulation: the + derivative with respect to a source is the adjoint field sampled over that + source's support, which the existing adjoint run already produces. The new + `mp.ArraySource` supplies per-point amplitudes directly, indexed exactly as + `get_dft_array` returns them over the same region, which is also the ordering + the gradient comes back in. An `ArraySource` passed to `MeepJaxWrapper` is + differentiated automatically, so a source computed in JAX backpropagates to + whatever produced it. + * Adjoint solver: `meep.adjoint.AngularSpectrum` propagates the tangential DFT fields on a planar monitor through an arbitrary stratified medium analytically, in JAX. Unlike `add_near2far`, which requires a homogeneous diff --git a/doc/docs/Python_Tutorials/Adjoint_Solver.md b/doc/docs/Python_Tutorials/Adjoint_Solver.md index baf5b0080..05d79ef31 100644 --- a/doc/docs/Python_Tutorials/Adjoint_Solver.md +++ b/doc/docs/Python_Tutorials/Adjoint_Solver.md @@ -362,6 +362,122 @@ parallel monitors, each with its own stack, and flip `sign` on the lower one. is a worked two-etch grating coupler radiating through a thick superstrate into a fiber, with a forward-only mode and an optimization mode. +Differentiating With Respect To Sources +--------------------------------------- + +The design region is not the only differentiable input. A source can be one too, +and it is nearly free: writing Maxwell's equations as $A(\rho) E = -i\omega J$ +gives $\partial J_{obj}/\partial J = -i\omega\lambda$, so the derivative with +respect to a source is just the adjoint field sampled where that source sits. It +needs no simulation beyond the adjoint run already being performed. + +A source opts in by naming the parameters it should be differentiated with +respect to: + +```py +src = mp.Source( + mp.GaussianSource(fcen, fwidth=df), + component=mp.Ez, + center=mp.Vector3(-1, 0), + differentiable=["amplitude"], + name="drive", +) + +opt = mpa.OptimizationProblem( + simulation=sim, + objective_functions=[objective], + objective_arguments=[monitor], + design_regions=[design_region], + frequencies=frequencies, +) + +value, grad = opt([rho]) +grad["design"] # exactly as before +grad["drive"]["amplitude"] # one entry per frequency +``` + +The names given in `differentiable` are the keys of the resulting subtree, so the +flag and the gradient cannot drift apart. With no flagged source the return value +is unchanged, so nothing existing is affected. + +Every source accepts `'amplitude'` and `'currents'`. The first is the scalar that +scales the whole source; the second is the per-point complex amplitude array, +which is the general case and the one that composes with a propagator. Names that +a source class does define but whose sensitivity is not yet implemented — +`beam_w0` on a Gaussian beam, for instance — raise `NotImplementedError` rather +than being reported as unknown, so the two situations are distinguishable. + +`'center'` and `'size'` are rejected outright. They move the grid points the +source occupies rather than the amplitudes applied to them, and the adjoint +formulation for sources gathers the cotangent over a *fixed* set of points, so a +derivative there is outside the formulation rather than merely missing. + +### Supplying the currents yourself + +`mp.ArraySource` takes the per-point amplitudes directly, which is what a source +computed somewhere else — by a mode solver, a propagator, or JAX — naturally +produces: + +```py +src = mp.ArraySource( + mp.GaussianSource(fcen, fwidth=df), + mp.Ez, + amplitudes=amplitudes, # one complex number per grid point + frequency=fcen, + center=mp.Vector3(-1, 0), + size=mp.Vector3(0, 4), + differentiable=["currents"], + name="sheet", +) +``` + +The array is indexed exactly like `Simulation.get_dft_array` over the same region +and component, which is also the ordering the gradient comes back in. That is not +a coincidence: injection and measurement are implemented as a scatter and its +exact transpose, so an array can be handed in and its cotangent read back with no +bookkeeping in between. + +### With JAX + +An `ArraySource` given to `MeepJaxWrapper` is differentiated automatically. There +is nothing to flag, because the parameters that produced the currents live +upstream in JAX rather than in Meep — Meep returns the cotangent with respect to +the current array and JAX carries it the rest of the way: + +```py +def loss(rho, beam_params): + currents = build_beam(beam_params) # pure JAX + (dft,) = wrapped_meep([rho], [currents]) + return jnp.sum(jnp.abs(dft) ** 2) + +value, (d_rho, d_beam) = jax.value_and_grad(loss, argnums=(0, 1))(rho, beam_params) +``` + +This is the objective-side protocol run backwards. There, Meep returns monitor +values and JAX owns the post-processing; here JAX supplies currents and Meep +returns their cotangent. The cotangent follows the convention JAX and autograd +both use for a real function of a complex input, $dJ/d(\mathrm{Re}\,a) - i\, +dJ/d(\mathrm{Im}\,a)$, so it chains without adjustment. + +### Checking a source gradient + +Two cautions, both of which produce a wrong answer that looks like a bug in the +gradient rather than like noise. + +Fix the run length. `stop_when_dft_decayed` is adaptive, so a perturbed run stops +at a different time and that difference scales with the perturbation; the ratio +of the adjoint gradient to the finite difference then settles at a fixed wrong +value instead of converging as the step shrinks. + +Give the adjoint run enough time. In a lossless background the adjoint DFT rings +considerably longer than the forward one, and the source gradient *is* the +adjoint field, so an under-converged adjoint hits it directly. In the calibration +for this feature the error was 16% at one run length and vanished entirely once +the run was doubled. + +Finally, a source inside or near a PML is rejected: its adjoint field is absorbed, +so the gradient would come back finite, smooth, and wrong. + Broadband Waveguide Mode Converter with Minimum Feature Size ------------------------------------------------------------ diff --git a/python/Makefile.am b/python/Makefile.am index 6a961fbdb..c492311a8 100644 --- a/python/Makefile.am +++ b/python/Makefile.am @@ -36,6 +36,7 @@ ADJOINT_TESTS = \ $(TEST_DIR)/test_adjoint_cyl.py \ $(TEST_DIR)/test_adjoint_protocol.py \ $(TEST_DIR)/test_angular_spectrum.py \ + $(TEST_DIR)/test_source_gradient.py \ $(TEST_DIR)/test_adjoint_jax.py TESTS = \ @@ -96,6 +97,7 @@ TESTS = \ $(TEST_DIR)/test_simulation.py \ $(TEST_DIR)/test_special_kz.py \ $(TEST_DIR)/test_source.py \ + $(TEST_DIR)/test_source_gradient.py \ $(TEST_DIR)/test_stop_when_flux_decayed.py \ $(TEST_DIR)/test_subpixel_3d.py \ $(TEST_DIR)/test_timing_measurements.py \ diff --git a/python/adjoint/source_gradient.py b/python/adjoint/source_gradient.py index 06f0dfb69..6c8a7e664 100644 --- a/python/adjoint/source_gradient.py +++ b/python/adjoint/source_gradient.py @@ -104,8 +104,12 @@ def _check_not_in_pml(self, sim: mp.Simulation) -> None: return half_cell = np.array([sim.cell_size.x, sim.cell_size.y, sim.cell_size.z]) / 2 - center = np.array([self.volume.center.x, self.volume.center.y, self.volume.center.z]) - half_size = np.array([self.volume.size.x, self.volume.size.y, self.volume.size.z]) / 2 + center = np.array( + [self.volume.center.x, self.volume.center.y, self.volume.center.z] + ) + half_size = ( + np.array([self.volume.size.x, self.volume.size.y, self.volume.size.z]) / 2 + ) for i, dim in enumerate("xyz"): if half_cell[i] == 0: diff --git a/python/adjoint/wrapper.py b/python/adjoint/wrapper.py index 597c80e0c..5ac34872a 100644 --- a/python/adjoint/wrapper.py +++ b/python/adjoint/wrapper.py @@ -88,7 +88,7 @@ def loss(rho, thickness, tilt): """ import contextlib import warnings -from typing import Callable, Iterable, List, Sequence, Tuple, Union +from typing import Callable, Iterable, List, Optional, Sequence, Tuple, Union import jax import jax.numpy as jnp @@ -97,6 +97,7 @@ def loss(rho, thickness, tilt): import meep as mp from . import DesignRegion, ObjectiveQuantity, utils +from . import source_gradient _warned_about_precision = False @@ -227,11 +228,15 @@ def __init__( self._simulate_fn = self._initialize_callable() def __call__( - self, designs: List[jnp.ndarray] + self, designs: List[jnp.ndarray], sources: Optional[List[jnp.ndarray]] = None ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, ...]]: """Performs a Meep simulation, taking designs and returning monitor values. Args: + sources: amplitudes for each `mp.ArraySource` passed to the constructor, + in that order. These are differentiated automatically -- there is no + need to flag them, because the parameters that produced them live + upstream in JAX rather than in Meep. Omit when there are none. designs: a list of design variables as 1D, 2D, or 3D JAX arrays. Valid shapes for design variables are (Nx, Ny, Nz) where Nx{y,z} match the elements of the `grid_size` constructor argument of Meep's `MaterialGrid` used for the @@ -251,14 +256,29 @@ def __call__( monitors were given. """ if _active_recorder is not None: - return _active_recorder(self, designs) - return self._simulate_fn(designs) + return _active_recorder(self, designs, sources) + return self._simulate_fn(designs, sources) + + def _update_sources(self, source_variables) -> None: + """Push JAX-supplied amplitudes into the differentiable sources.""" + if source_variables is None: + return + if len(source_variables) != len(self.differentiable_sources): + raise ValueError( + f"Got {len(source_variables)} source arrays but " + f"{len(self.differentiable_sources)} differentiable sources." + ) + for src, amps in zip(self.differentiable_sources, source_variables): + src.amplitudes = onp.asarray(amps, dtype=onp.complex128) def _run_fwd_simulation( - self, design_variables: Iterable[onp.ndarray] + self, + design_variables: Iterable[onp.ndarray], + source_variables: Optional[Iterable[onp.ndarray]] = None, ) -> Tuple[Union[jnp.ndarray, Tuple[jnp.ndarray, ...]], List[List[mp.DftFields]]]: """Runs forward simulation, returning monitor values and design region fields.""" utils.validate_and_update_design(self.design_regions, design_variables) + self._update_sources(source_variables) self.simulation.reset_meep() self.simulation.change_sources(self.sources) utils.register_monitors(self.monitors, self.frequencies) @@ -305,6 +325,11 @@ def _run_adjoint_simulation( self.design_regions, self.frequencies, ) + self.adj_source_monitors = source_gradient.install_source_gradient_monitors( + self.simulation, + self.differentiable_sources, + self.frequencies, + ) self.simulation.init_sim() sim_run_args = { "until_after_sources" @@ -317,6 +342,29 @@ def _run_adjoint_simulation( return adj_design_region_monitors + def _source_vjps(self, sum_freq_partials: bool = True) -> List[onp.ndarray]: + """Cotangents with respect to each differentiable source's currents. + + With `sum_freq_partials` false the frequency axis is kept, which is what + `value_and_jacobian` needs to keep its rows separable. + """ + if not self.differentiable_sources: + return [] + phase = self.monitors[0]._adj_src_phase() + out = [] + for src, monitor in zip(self.differentiable_sources, self.adj_source_monitors): + scale = source_gradient.source_grad_scale( + self.simulation, src, self.frequencies, phase + ) + grad = monitor.gather(self.simulation, scale) + shape = src.amplitudes.shape + if sum_freq_partials: + # the amplitudes are shared across the band, as design weights are + out.append(onp.sum(grad, axis=0).reshape(shape)) + else: + out.append(grad.reshape((grad.shape[0], *shape))) + return out + def _calculate_vjps( self, fwd_fields: List[List[mp.DftFields]], @@ -340,14 +388,16 @@ def _initialize_callable(self) -> Callable: """Initializes the callable JAX function and registers its VJP.""" @jax.custom_vjp - def simulate(design_variables: List[jnp.ndarray]): - monitor_values, _ = self._run_fwd_simulation(design_variables) + def simulate(design_variables: List[jnp.ndarray], source_variables): + monitor_values, _ = self._run_fwd_simulation( + design_variables, source_variables + ) return monitor_values - def _simulate_fwd(design_variables): + def _simulate_fwd(design_variables, source_variables): """Runs forward simulation, returning monitor values and fields.""" monitor_values, self.fwd_design_region_monitors = self._run_fwd_simulation( - design_variables + design_variables, source_variables ) design_variable_shapes = [x.shape for x in design_variables] return monitor_values, (design_variable_shapes) @@ -363,7 +413,11 @@ def _simulate_rev(res, monitor_values_grad): self.adj_design_region_monitors, design_variable_shapes, ) - return ([jnp.asarray(vjp) for vjp in vjps],) + source_vjps = [jnp.asarray(v) for v in self._source_vjps()] + return ( + [jnp.asarray(vjp) for vjp in vjps], + source_vjps if self.differentiable_sources else None, + ) simulate.defvjp(_simulate_fwd, _simulate_rev) @@ -388,10 +442,11 @@ def __init__(self): self.design_shapes = None self.monitor_values = None self.fwd_monitors = None + self.sources = None self.num_calls = 0 self.substitute = None - def __call__(self, wrapper: "MeepJaxWrapper", designs): + def __call__(self, wrapper: "MeepJaxWrapper", designs, sources=None): self.num_calls += 1 if self.wrapper is None: self.wrapper = wrapper @@ -406,10 +461,11 @@ def __call__(self, wrapper: "MeepJaxWrapper", designs): # the traced passes it holds tracers, which is what the design mapping's # pullback is recovered from. self.designs = list(designs) + self.sources = None if sources is None else list(sources) if self.monitor_values is None: self.design_shapes = [onp.shape(design) for design in designs] self.monitor_values, self.fwd_monitors = wrapper._run_fwd_simulation( - designs + designs, sources ) return self.monitor_values if self.substitute is None else self.substitute @@ -555,13 +611,20 @@ def evaluate_at_monitor_values(monitor_values): rows = [jnp.moveaxis(jnp.asarray(row), -1, 0) for row in rows] # (4) Carry each row back through whatever produced the design - # weights. Pure JAX, so `vmap` over the frequency axis is free. - def designs_from(*values_to_differentiate): + # weights, and the source amplitudes if any came from JAX. Pure + # JAX, so `vmap` over the frequency axis is free. + source_rows = [ + jnp.asarray(v) for v in wrapper._source_vjps(sum_freq_partials=False) + ] + + def inputs_from(*values_to_differentiate): loss(*rebuild(values_to_differentiate)) - return recorder.designs + return recorder.designs, recorder.sources - _, pull_designs = jax.vjp(designs_from, *differentiated) - implicit = jax.vmap(pull_designs)(rows) + _, pull_inputs = jax.vjp(inputs_from, *differentiated) + implicit = jax.vmap(pull_inputs)( + (rows, source_rows if recorder.sources is not None else None) + ) if isinstance(argnums, int): implicit = implicit[0] diff --git a/python/tests/test_source_gradient.py b/python/tests/test_source_gradient.py new file mode 100644 index 000000000..974648ed0 --- /dev/null +++ b/python/tests/test_source_gradient.py @@ -0,0 +1,462 @@ +"""Tests for adjoint gradients with respect to source amplitudes. + +The finite-difference comparisons all pin the run length rather than using +`stop_when_dft_decayed`. That is deliberate: an adaptive stop makes a perturbed +run end at a different time, and the difference scales with the perturbation, +so the ratio of the adjoint gradient to the finite difference settles at a +fixed wrong value instead of converging. It looks exactly like a wrong +gradient. The runs here are also long enough for the *adjoint* DFT, which rings +considerably longer than the forward one in a lossless background. +""" +import unittest + +import numpy as np +from autograd import numpy as npa + +import meep as mp +import meep.adjoint as mpa + +RES = 20 +FCEN = 1.0 +RUN = 200.0 +SRC_C = mp.Vector3(-1.0, 0) +MON_C = mp.Vector3(1.2, 0) +CELL = mp.Vector3(6, 4) +RHO = 0.5 * np.ones(64) +FD_STEP = 1e-4 + + +def _design_region(sim): + """An inert design region; OptimizationProblem requires one.""" + mg = mp.MaterialGrid( + mp.Vector3(8, 8), mp.air, mp.Medium(index=1.5), grid_type="U_MEAN" + ) + dr = mpa.DesignRegion( + mg, volume=mp.Volume(center=mp.Vector3(0.2, 0), size=mp.Vector3(0.4, 0.4)) + ) + sim.geometry = [mp.Block(center=dr.center, size=dr.size, material=mg)] + return dr + + +def _problem(source, fcen=FCEN, res=RES, comp=mp.Ez): + sim = mp.Simulation( + cell_size=CELL, + resolution=res, + boundary_layers=[mp.PML(1.0)], + sources=[source], + force_complex_fields=True, + ) + dr = _design_region(sim) + mon = mpa.FourierFields(sim, mp.Volume(center=MON_C, size=mp.Vector3(0, 0)), comp) + opt = mpa.OptimizationProblem( + simulation=sim, + objective_functions=[lambda f: npa.sum(npa.abs(f) ** 2)], + objective_arguments=[mon], + design_regions=[dr], + frequencies=[fcen], + minimum_run_time=RUN, + maximum_run_time=RUN, + ) + return sim, opt + + +def _num_source_points(center, size, comp=mp.Ez, res=RES): + """How many grid points a source region covers.""" + sim = mp.Simulation( + cell_size=CELL, + resolution=res, + boundary_layers=[mp.PML(1.0)], + force_complex_fields=True, + sources=[ + mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), component=comp, center=center + ) + ], + ) + sim.init_sim() + vol = sim._fit_volume_to_simulation(mp.Volume(center=center, size=size)) + mon = sim.add_dft_fields([comp], [FCEN], where=vol, yee_grid=True) + sim._evaluate_dft_objects() + return int(np.prod(sim.fields.dft_monitor_size(mon.swigobj, vol.swigobj, comp))) + + +def _complex_fd(evaluate, perturb, step=FD_STEP): + """Central difference in the convention JAX and autograd use. + + For a real objective of a complex parameter both frameworks return + `dJ/d(Re a) - i dJ/d(Im a)`, so that is what the adjoint result is compared + against. + """ + out = {} + for tag, delta in (("re", step), ("im", 1j * step)): + out[tag] = (evaluate(perturb(delta)) - evaluate(perturb(-delta))) / (2 * step) + return out["re"] - 1j * out["im"] + + +class TestDifferentiableFlag(unittest.TestCase): + """Validation of `differentiable=`; none of these run a simulation.""" + + def _src(self, **kwargs): + return mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), + component=mp.Ez, + center=SRC_C, + **kwargs, + ) + + def test_accepts_universal_parameters(self): + src = self._src(differentiable=["currents", "amplitude"], name="drive") + self.assertEqual(src.differentiable, ("currents", "amplitude")) + self.assertEqual(src.name, "drive") + + def test_default_is_not_differentiable(self): + self.assertEqual(self._src().differentiable, ()) + + def test_rejects_unknown_parameter(self): + with self.assertRaisesRegex(ValueError, "not a differentiable parameter"): + self._src(differentiable=["nonsense"]) + + def test_rejects_bare_string(self): + with self.assertRaisesRegex(ValueError, "list of parameter names"): + self._src(differentiable="currents") + + def test_rejects_duplicates(self): + with self.assertRaisesRegex(ValueError, "duplicate"): + self._src(differentiable=["currents", "currents"]) + + def test_support_moving_parameters_get_their_own_message(self): + # These are outside the formulation rather than merely unimplemented, + # and the error should distinguish the two. + for name in ("center", "size"): + with self.assertRaisesRegex(ValueError, "grid points the source occupies"): + self._src(differentiable=[name]) + + def test_unimplemented_parameter_is_not_reported_as_unknown(self): + beam = dict( + src=mp.GaussianSource(FCEN, fwidth=0.2), + center=SRC_C, + size=mp.Vector3(0, 2), + beam_kdir=mp.Vector3(1, 0, 0), + beam_w0=1.0, + beam_E0=mp.Vector3(0, 0, 1), + ) + with self.assertRaises(NotImplementedError): + mp.GaussianBeam3DSource(differentiable=["beam_w0"], **beam) + # ... but it is still rejected as unknown on a class that lacks it + with self.assertRaises(ValueError): + self._src(differentiable=["beam_w0"]) + + +class TestTranspose(unittest.TestCase): + """`fourier_sourcegradient` must be the exact transpose of the scatter.""" + + # Kept clear of the PML boundary so the monitor lands in a single chunk. + # `test_dot_product_identity` needs that: it pairs the scatter's per-chunk + # amplitudes with `get_dft_array`, and the local-to-global index mapping + # that would let it do so across chunks is not reachable from Python. The + # cross-chunk reduction is covered instead by + # `test_gather_matches_dft_for_an_aligned_monitor`, which compares two + # globally assembled arrays, and by the finite-difference tests. + MON_SIZE = mp.Vector3(0, 1.0) + + def _monitor(self, mon_x=1.0, symmetries=(), comp=mp.Ez, freqs=(1.0, 1.15)): + sim = mp.Simulation( + cell_size=CELL, + resolution=RES, + boundary_layers=[mp.PML(1.0)], + sources=[ + mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), component=mp.Ez, center=SRC_C + ) + ], + symmetries=list(symmetries), + force_complex_fields=True, + ) + vol = mp.Volume(center=mp.Vector3(mon_x, 0), size=self.MON_SIZE) + mon = sim.add_dft_fields([comp], np.asarray(freqs), where=vol, yee_grid=True) + sim.run(until=25) + npts = int(np.prod(sim.fields.dft_monitor_size(mon.swigobj, vol.swigobj, comp))) + return sim, vol, mon, npts, np.asarray(freqs) + + def test_gather_matches_dft_for_an_aligned_monitor(self): + # With no symmetry and a grid-aligned electric monitor every weight is + # 1 except the -1 the scatter applies to electric components, so the + # gathered gradient must be exactly -dft. This pins ordering and sign. + sim, vol, mon, npts, freqs = self._monitor() + grad = np.zeros(npts * len(freqs), dtype=np.complex128) + mon.swigobj.fourier_sourcegradient(vol.swigobj, mp.Ez, sim.fields, grad) + grad = grad.reshape(len(freqs), npts) + for i in range(len(freqs)): + dft = np.asarray(sim.get_dft_array(mon, mp.Ez, i)).ravel() + np.testing.assert_allclose(grad[i], -dft, rtol=1e-13, atol=0) + + def test_dot_product_identity(self): + # == . If the scatter and the gather disagree + # about ordering or weights, this fails. No FDTD gradient involved. + sim, vol, mon, npts, freqs = self._monitor() + nf = len(freqs) + rng = np.random.default_rng(0) + x = ( + rng.standard_normal((nf, npts)) + 1j * rng.standard_normal((nf, npts)) + ).ravel() + + grad = np.zeros(nf * npts, dtype=np.complex128) + mon.swigobj.fourier_sourcegradient(vol.swigobj, mp.Ez, sim.fields, grad) + rhs = np.sum(grad * x) + + srcdata = mon.swigobj.fourier_sourcedata(vol.swigobj, mp.Ez, sim.fields, x) + dft = np.stack( + [np.asarray(sim.get_dft_array(mon, mp.Ez, i)).ravel() for i in range(nf)] + ) + lhs = 0j + for sd in srcdata: + amp = np.array(sd.amp_arr).reshape(-1, nf) + self.assertEqual(amp.shape[0], npts) # single aligned chunk + for i in range(nf): + lhs += np.sum(dft[i] * amp[:, i]) + + self.assertAlmostEqual(abs(lhs - rhs) / abs(lhs), 0.0, places=12) + + +class TestArraySource(unittest.TestCase): + def _field(self, source): + sim = mp.Simulation( + cell_size=CELL, + resolution=RES, + boundary_layers=[mp.PML(1.0)], + sources=[source], + force_complex_fields=True, + ) + mon = sim.add_dft_fields( + [mp.Ez], + [FCEN], + where=mp.Volume(center=MON_C, size=mp.Vector3(0.2, 0.2)), + ) + sim.run(until_after_sources=60) + return np.asarray(sim.get_dft_array(mon, mp.Ez, 0)).ravel() + + def test_one_point_matches_an_ordinary_source(self): + # Pins both conversions in ArraySource.add_source: the -1 the scatter + # applies to electric components, and the fact that it places a current + # density rather than a point amplitude. + t = lambda: mp.GaussianSource(FCEN, fwidth=0.2) + plain = self._field( + mp.Source(t(), component=mp.Ez, center=SRC_C, amplitude=1.0) + ) + array = self._field( + mp.ArraySource( + t(), + mp.Ez, + amplitudes=np.array([1.0 + 0j]), + frequency=FCEN, + center=SRC_C, + size=mp.Vector3(0, 0), + ) + ) + np.testing.assert_allclose(array, plain, rtol=2e-6) + + def test_rejects_wrong_length(self): + src = mp.ArraySource( + mp.GaussianSource(FCEN, fwidth=0.2), + mp.Ez, + amplitudes=np.ones(3, dtype=complex), + frequency=FCEN, + center=SRC_C, + size=mp.Vector3(0, 0.4), + ) + sim = mp.Simulation( + cell_size=CELL, + resolution=RES, + boundary_layers=[mp.PML(1.0)], + sources=[src], + force_complex_fields=True, + ) + with self.assertRaisesRegex(ValueError, "grid points"): + sim.init_sim() + + +class TestSourceGradient(unittest.TestCase): + def test_no_flagged_source_leaves_the_return_shape_alone(self): + src = mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), component=mp.Ez, center=SRC_C + ) + _, opt = _problem(src) + _, grad = opt([RHO]) + self.assertNotIsInstance(grad, dict) + + def test_rejects_a_source_inside_the_pml(self): + # The adjoint field is absorbed there, so the gradient would come back + # finite, smooth, and wrong. + src = mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), + component=mp.Ez, + center=mp.Vector3(-2.7, 0), + differentiable=["amplitude"], + name="drive", + ) + _, opt = _problem(src) + with self.assertRaisesRegex(ValueError, "PML"): + opt([RHO]) + + def _amplitude_case(self, fcen=FCEN, res=RES, amp=1.0, comp=mp.Ez): + def make(a): + return mp.Source( + mp.GaussianSource(fcen, fwidth=0.2 * fcen), + component=comp, + center=SRC_C, + amplitude=a, + differentiable=["amplitude"], + name="drive", + ) + + _, opt = _problem(make(amp), fcen, res, comp) + _, grad = opt([RHO]) + adjoint = np.ravel(grad["drive"]["amplitude"])[0] + + def evaluate(a): + _, o = _problem(make(a), fcen, res, comp) + return np.asarray(o([RHO], need_gradient=False)[0]).item() + + reference = _complex_fd(evaluate, lambda d: amp + d) + return adjoint, reference + + def test_amplitude_gradient(self): + for label, kwargs in [ + ("baseline", {}), + ("higher resolution", dict(res=30)), + ("lower frequency", dict(fcen=0.8)), + ("higher frequency", dict(fcen=1.3)), + ("complex amplitude", dict(amp=0.7 - 0.4j)), + ("magnetic source", dict(comp=mp.Hz)), + ]: + with self.subTest(label): + adjoint, reference = self._amplitude_case(**kwargs) + self.assertLess( + abs(adjoint - reference) / abs(reference), + 1e-5, + f"{label}: adjoint {adjoint} vs finite difference {reference}", + ) + + def test_currents_gradient(self): + size = mp.Vector3(0, 0.4) + npts = _num_source_points(SRC_C, size) + rng = np.random.default_rng(7) + amps = rng.standard_normal(npts) + 1j * rng.standard_normal(npts) + + def make(a): + return mp.ArraySource( + mp.GaussianSource(FCEN, fwidth=0.2), + mp.Ez, + amplitudes=a, + frequency=FCEN, + center=SRC_C, + size=size, + differentiable=["currents"], + name="sheet", + ) + + _, opt = _problem(make(amps)) + _, grad = opt([RHO]) + adjoint = np.ravel(grad["sheet"]["currents"]) + self.assertEqual(adjoint.size, npts) + + def evaluate(a): + _, o = _problem(make(a)) + return np.asarray(o([RHO], need_gradient=False)[0]).item() + + for j in (0, npts // 2, npts - 1): + with self.subTest(point=j): + + def perturb(delta, j=j): + out = amps.copy() + out[j] += delta + return out + + reference = _complex_fd(evaluate, perturb) + self.assertLess( + abs(adjoint[j] - reference) / abs(reference), + 1e-5, + f"point {j}: adjoint {adjoint[j]} vs {reference}", + ) + + +try: + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + + _HAVE_JAX = True +except ImportError: + _HAVE_JAX = False + + +@unittest.skipUnless(_HAVE_JAX, "JAX is an optional dependency") +class TestJaxRoundTrip(unittest.TestCase): + """`jax.grad` through an `ArraySource` must match a finite difference. + + This is what pins the cotangent convention. Meep returns + `dJ/d(Re a) - i dJ/d(Im a)`, which is what JAX and autograd both produce + for a real function of a complex input, so it chains with no adjustment. + A convention error here would not raise, it would just be wrong. + """ + + def test_gradient_chains_through_jax(self): + size = mp.Vector3(0, 0.4) + npts = _num_source_points(SRC_C, size) + + def wrapper(): + src = mp.ArraySource( + mp.GaussianSource(FCEN, fwidth=0.2), + mp.Ez, + amplitudes=np.ones(npts, dtype=complex), + frequency=FCEN, + center=SRC_C, + size=size, + name="sheet", + ) + sim = mp.Simulation( + cell_size=CELL, + resolution=RES, + boundary_layers=[mp.PML(1.0)], + force_complex_fields=True, + ) + dr = _design_region(sim) + mon = mpa.FourierFields( + sim, mp.Volume(center=MON_C, size=mp.Vector3(0, 0)), mp.Ez + ) + return mpa.MeepJaxWrapper( + sim, + [src], + [mon], + [dr], + [FCEN], + minimum_run_time=RUN, + maximum_run_time=RUN, + ) + + rng = np.random.default_rng(11) + amps0 = jnp.asarray(rng.standard_normal(npts) + 1j * rng.standard_normal(npts)) + rho = 0.5 * jnp.ones((8, 8)) + + def loss(rho, amps): + (dft,) = wrapper()([rho], [amps]) + return jnp.sum(jnp.abs(dft) ** 2) + + _, (_, adjoint) = jax.value_and_grad(loss, argnums=(0, 1))(rho, amps0) + self.assertEqual(adjoint.shape, (npts,)) + + for j in (0, npts // 2, npts - 1): + with self.subTest(point=j): + reference = _complex_fd( + lambda a: loss(rho, a), lambda d, j=j: amps0.at[j].add(d) + ) + self.assertLess( + abs(adjoint[j] - reference) / abs(reference), + 1e-5, + f"point {j}: jax.grad {adjoint[j]} vs {reference}", + ) + + +if __name__ == "__main__": + unittest.main() From b55025636e07139dd67a3c13982f2a254b1b9a3d Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 27 Aug 2026 10:40:07 -0700 Subject: [PATCH 3/5] adjoint: ArraySource must not negate magnetic components fourier_sourcedata negates electric components and only those, since it was written to place adjoint sources. ArraySource was undoing that negation for every component, which flips the relative sign of the two sheets of an equivalent-current pair -- and that silently reverses the direction the pair radiates in rather than producing anything that looks like an error. The existing round-trip test only covered mp.Ez, so it could not see this. --- python/source.py | 28 ++++++++++++++++----- python/tests/test_source_gradient.py | 37 +++++++++++++++++----------- 2 files changed, 45 insertions(+), 20 deletions(-) diff --git a/python/source.py b/python/source.py index 9d17ce107..c995d21a7 100644 --- a/python/source.py +++ b/python/source.py @@ -1201,6 +1201,7 @@ def __init__( center=None, volume=None, size=Vector3(), + yee_grid=True, differentiable=None, name=None, ): @@ -1209,6 +1210,13 @@ def __init__( + **`amplitudes` [`numpy.ndarray`]** — Complex amplitude for each point of the source region, shaped like `get_dft_array` over that region. + **`frequency` [`number`]** — The frequency the amplitudes refer to. + + **`yee_grid` [`boolean`]** — Whether the amplitudes are given at the + Yee points of `component` (the default) or at voxel centers. Voxel + centers are what several components sharing one plane need, since + each component's Yee points sit at a different half-pixel offset and + the arrays would otherwise have different lengths. Meep spreads a + centered value across the neighbouring Yee points, and the adjoint + gathers it back the same way, so the gradient stays exact either way. """ super().__init__( src, @@ -1221,6 +1229,7 @@ def __init__( ) self.amplitudes = np.ascontiguousarray(amplitudes, dtype=np.complex128) self.frequency = float(frequency) + self.yee_grid = yee_grid self._monitor = None def add_source(self, sim): @@ -1235,7 +1244,7 @@ def add_source(self, sim): sim.fields.require_component(self.component) mon = sim.add_dft_fields( - [self.component], [self.frequency], where=vol, yee_grid=True + [self.component], [self.frequency], where=vol, yee_grid=self.yee_grid ) # add_dft_fields defers construction; the scatter needs it now sim._evaluate_dft_objects() @@ -1249,14 +1258,21 @@ def add_source(self, sim): ) # Two conversions, so that one entry of `amplitudes` means exactly what - # `Source.amplitude` means for a point source at that grid point: - # fourier_sourcedata negates electric components (it was written to - # place adjoint sources), and it places a current *density*, whose - # integral over a voxel is the amplitude times dV. + # `Source.amplitude` means for a point source at that grid point. + # + # fourier_sourcedata negates *electric* components and only those, since + # it was written to place adjoint sources; undoing it for magnetic + # components too would flip the relative sign of the two sheets of an + # equivalent-current pair, which silently reverses the direction such a + # pair radiates in. + # + # It also places a current *density*, whose integral over a voxel is the + # amplitude times dV. num_dims = sim._infer_dimensions(sim.k_point) dV = 1 / sim.resolution**num_dims + sign = -1.0 if mp.is_electric(self.component) else 1.0 flat = np.ascontiguousarray( - self.amplitudes.ravel() * complex(self.amplitude) * (-1.0 / dV), + self.amplitudes.ravel() * complex(self.amplitude) * (sign / dV), dtype=np.complex128, ) srcdata = mon.swigobj.fourier_sourcedata( diff --git a/python/tests/test_source_gradient.py b/python/tests/test_source_gradient.py index 974648ed0..5ff2166ed 100644 --- a/python/tests/test_source_gradient.py +++ b/python/tests/test_source_gradient.py @@ -239,21 +239,30 @@ def test_one_point_matches_an_ordinary_source(self): # Pins both conversions in ArraySource.add_source: the -1 the scatter # applies to electric components, and the fact that it places a current # density rather than a point amplitude. + # + # Both an electric and a magnetic component are checked, because the + # scatter negates electric components and only those. Undoing that + # negation for magnetic ones too leaves each component individually + # plausible but flips the relative sign of the two sheets of an + # equivalent-current pair, which reverses the direction such a pair + # radiates in without producing anything that looks like an error. t = lambda: mp.GaussianSource(FCEN, fwidth=0.2) - plain = self._field( - mp.Source(t(), component=mp.Ez, center=SRC_C, amplitude=1.0) - ) - array = self._field( - mp.ArraySource( - t(), - mp.Ez, - amplitudes=np.array([1.0 + 0j]), - frequency=FCEN, - center=SRC_C, - size=mp.Vector3(0, 0), - ) - ) - np.testing.assert_allclose(array, plain, rtol=2e-6) + for component in (mp.Ez, mp.Hz): + with self.subTest(component=mp.component_name(component)): + plain = self._field( + mp.Source(t(), component=component, center=SRC_C, amplitude=1.0) + ) + array = self._field( + mp.ArraySource( + t(), + component, + amplitudes=np.array([1.0 + 0j]), + frequency=FCEN, + center=SRC_C, + size=mp.Vector3(0, 0), + ) + ) + np.testing.assert_allclose(array, plain, rtol=2e-6) def test_rejects_wrong_length(self): src = mp.ArraySource( From 5471c174cf3e78ced287d8db832151e2b1569e09 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 27 Aug 2026 15:09:48 -0700 Subject: [PATCH 4/5] adjoint: gradients for Gaussian beam parameters Contracts the per-point cotangent onto beam_x0, beam_kdir, beam_w0 and beam_E0 by finite-differencing Meep's own beam construction -- no FDTD runs, only re-evaluating mp.gaussianbeam, and each directional derivative is contracted as it is formed so no dense Jacobian is built. This needed the transpose of the add_volume_source / src_vol_chunkloop path, which is how beams (and amp_data, and amp_func) are actually placed -- distinct from fourier_sourcedata, which is how adjoint sources are placed. Positions come out of the same C++ loop as the cotangent, following material_grids_addgradient: rebuild the chunk's grid volume with gv.subvolume(is, ie, c) and loop it, so dft[nf*idx+f] and the point it belongs to cannot disagree. Four things had to be right, each of which produced plausible wrong numbers: - is_old/ie_old are only assigned when a chunk is created with persist, so a monitor that does not set it must loop over is/ie instead; - src_vol_chunkloop multiplies by gv.a once per *zero-width* direction, not per dimension, so 1/dV is right for a point source and wrong by a factor of the resolution for a line; - IVEC_LOOP_WEIGHT has to be applied. Omitting it is invisible for a grid-aligned electric source and gives exactly (1/2)^(zero-width directions) for a component whose Yee points straddle the plane; - the output must not be reduced. Collapsing a zero-extent direction merges the two Yee planes a magnetic component straddles: the total weight survives but the variation between them does not, which is exact for a uniform sheet and wrong for a beam. Verified against finite differences: 1.5e-8 for beam_w0, 2.8e-8 for beam_x0.x, 2.7e-5 for beam_x0.y, 2.3e-8 for beam_kdir. beam_E0.z is checked against an exact answer instead -- the objective is quadratic in it, so dJ/dE0.z is exactly 2J -- and agrees to eight significant figures. The single-component checks are exact to 1.0000 for electric and magnetic point and line sources. --- python/adjoint/optimization_problem.py | 46 +++- python/adjoint/source_gradient.py | 306 +++++++++++++++++++++++-- python/meep.i | 3 +- python/source.py | 19 +- src/dft.cpp | 190 +++++++++++++++ src/meep.hpp | 6 + 6 files changed, 532 insertions(+), 38 deletions(-) diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index 88e94ffc3..78a0974fc 100644 --- a/python/adjoint/optimization_problem.py +++ b/python/adjoint/optimization_problem.py @@ -422,7 +422,7 @@ def calculate_source_gradient(self): out = {} for ar in range(len(self.objective_functions)): for si, src in enumerate(self.differentiable_sources): - monitor = self.adjoint_source_monitors[ar][si] + monitors = self.adjoint_source_monitors[ar][si] # the adjoint field carries the normalization the objective # quantity applied when it placed the adjoint source, so the # transpose has to use that same quantity's phase @@ -432,8 +432,7 @@ def calculate_source_gradient(self): self.frequencies, self.objective_arguments[0]._adj_src_phase(), ) - currents = monitor.gather(self.sim, scale) - grads = source_gradient.contract(src, currents) + grads = self._source_gradient_for(src, monitors, scale) key = source_gradient.source_key(src, si) if len(self.objective_functions) == 1: out[key] = grads @@ -441,6 +440,47 @@ def calculate_source_gradient(self): out.setdefault(key, []).append(grads) return out + def _source_gradient_for(self, src, monitors, scale): + """Gradients for one differentiable source, keyed by parameter name.""" + beam_names = [name for name in src.differentiable if name.startswith("beam_")] + other_names = [ + name for name in src.differentiable if not name.startswith("beam_") + ] + + grads = {} + if other_names: + currents = monitors[0].gather(self.sim, scale) + everything = source_gradient.contract(src, currents) + grads.update({name: everything[name] for name in other_names}) + + if beam_names: + # Contract the per-point cotangents onto the beam's own parameters, + # one frequency at a time so the rows stay separable. + cotangents_by_frequency = [] + for f_index in range(len(np.atleast_1d(self.frequencies))): + cotangents = {} + for monitor in monitors: + values = monitor.gather(self.sim, scale) + cotangents[monitor.component] = ( + monitor.positions(self.sim), + np.asarray(values)[f_index].ravel(), + ) + cotangents_by_frequency.append( + source_gradient.beam_parameter_gradients( + self.sim, + src, + beam_names, + cotangents, + source_gradient.normal_index(src), + center=monitors[0].volume.center, + ) + ) + for name in beam_names: + grads[name] = np.squeeze( + np.array([row[name] for row in cotangents_by_frequency]) + ) + return grads + def calculate_gradient(self): self.source_gradient = self.calculate_source_gradient() diff --git a/python/adjoint/source_gradient.py b/python/adjoint/source_gradient.py index 6c8a7e664..c15812373 100644 --- a/python/adjoint/source_gradient.py +++ b/python/adjoint/source_gradient.py @@ -29,7 +29,14 @@ # Parameters this module can currently evaluate. `source.differentiable` is # validated against a wider list at construction time; see # `meep.source._validate_differentiable`. -IMPLEMENTED_PARAMS = ("currents", "amplitude") +IMPLEMENTED_PARAMS = ( + "currents", + "amplitude", + "beam_x0", + "beam_kdir", + "beam_w0", + "beam_E0", +) def differentiable_sources(sim: mp.Simulation) -> List: @@ -63,9 +70,12 @@ def __init__( source, frequencies: np.ndarray, decimation_factor: Optional[int] = 0, + component: Optional[int] = None, ): self.source = source - self.component = source.component + # A Gaussian beam has no single component: it places four, so each gets + # its own monitor and the caller says which one this is. + self.component = source.component if component is None else component self._frequencies = np.asarray(frequencies) if getattr(source, "amp_func", None) is not None or getattr( @@ -135,13 +145,23 @@ def register(self, sim: mp.Simulation) -> None: decimation_factor=self.decimation_factor, ) - def shape(self, sim: mp.Simulation): - """(nfreq,) + the monitor's spatial shape, with trailing 1s dropped.""" - dims = sim.fields.dft_monitor_size( + def num_points(self, sim: mp.Simulation) -> int: + """How many grid points this source actually drives. + + The *unreduced* count. A zero-width source volume straddles two Yee + planes of a magnetic component, each carrying half the current, and + collapsing that direction would merge them -- preserving the total + weight but losing how the amplitude varies between the two, which is + exactly what distinguishes a beam from a uniform sheet. + """ + dims = sim.fields.dft_monitor_full_size( self._monitor.swigobj, self.volume.swigobj, self.component ) - dims = [d for d in dims if d > 1] or [1] - return (len(self._frequencies), *dims) + return int(np.prod(dims)) + + def shape(self, sim: mp.Simulation): + """(nfreq, num points).""" + return (len(self._frequencies), self.num_points(sim)) def gather(self, sim: mp.Simulation, scale: np.ndarray) -> np.ndarray: """Return dJ_obj/d(currents), summed over all processes. @@ -149,20 +169,44 @@ def gather(self, sim: mp.Simulation, scale: np.ndarray) -> np.ndarray: `scale` is the per-frequency factor relating a current amplitude to the adjoint field, broadcast over the spatial axes. """ - dims = sim.fields.dft_monitor_size( - self._monitor.swigobj, self.volume.swigobj, self.component - ) - num_points = int(np.prod(dims)) + num_points = self.num_points(sim) grad = np.zeros(num_points * len(self._frequencies), dtype=np.complex128) - self._monitor.swigobj.fourier_sourcegradient( + self._monitor.swigobj.volume_source_gradient( self.volume.swigobj, self.component, sim.fields, grad ) grad = grad.reshape(len(self._frequencies), num_points) - grad *= np.asarray(scale).reshape(-1, 1) + # volume_source_gradient negates nothing, while the calibration that + # fixed `scale` was done against an electric source. Electric and + # magnetic components therefore differ by a sign here -- getting this + # wrong leaves each component individually plausible but silently + # subtracts one contribution from the other where a source drives + # several, as a Gaussian beam does. + sign = -1.0 if mp.is_electric(self.component) else 1.0 + # No half-timestep phase correction here, despite step.cpp evaluating + # magnetic sources at t and electric ones at t + dt/2. Applying one + # breaks the single-component checks by exactly 1/cos(w dt/2), so + # whatever absorbs that offset is already accounted for in the + # calibrated adjoint source phase. + grad *= sign * np.asarray(scale).reshape(-1, 1) return grad.reshape(self.shape(sim)) + def positions(self, sim: mp.Simulation) -> np.ndarray: + """(num points, 3) positions matching `gather`'s ordering. + + Emitted from the same C++ loop as the cotangent rather than + reconstructed here, so the two cannot disagree about which point is + which -- which is exactly the kind of half-pixel mismatch that produces + a smooth, plausible, wrong gradient. + """ + num_points = self.num_points(sim) + out = np.zeros(3 * num_points, dtype=np.float64) + self._monitor.swigobj.monitor_positions( + self.volume.swigobj, self.component, sim.fields, out + ) + return out.reshape(num_points, 3) + def install_source_gradient_monitors( sim: mp.Simulation, @@ -171,14 +215,76 @@ def install_source_gradient_monitors( decimation_factor: Optional[int] = 0, ) -> List[SourceGradientMonitor]: """Install a DFT monitor over each differentiable source's support.""" - monitors = [ - SourceGradientMonitor(sim, s, frequencies, decimation_factor) for s in sources - ] - for m in monitors: - m.register(sim) + monitors = [] + for source in sources: + # one monitor per component the source actually drives; a Gaussian beam + # places four + group = [ + SourceGradientMonitor( + sim, source, frequencies, decimation_factor, component=component + ) + for component in source_components(source, sim.dimensions) + ] + for m in group: + m.register(sim) + monitors.append(group) return monitors +def is_gaussian_beam(source) -> bool: + return hasattr(source, "beam_w0") and hasattr(source, "beam_kdir") + + +def normal_index(source) -> int: + """Which axis the source plane is normal to, as 0, 1 or 2.""" + size = [source.size.x, source.size.y, source.size.z] + zeros = [i for i, s in enumerate(size) if s == 0] + if not zeros: + raise ValueError( + "A Gaussian beam source must be a plane (a line in 2D), so one of " + "its size components has to be zero." + ) + return zeros[0] + + +def beam_places(source, component, dimensions: int) -> bool: + """Whether `add_volume_source_check` actually places this component. + + It declines several, and contracting a sensitivity for a source that was + never placed is silently wrong rather than an error: + + - components along the plane normal; + - in 2D, whichever parity the beam does not excite. + + Mirrors sources.cpp:495. + """ + normal = normal_index(source) + axis = {mp.X: 0, mp.Y: 1, mp.Z: 2}[mp.component_direction(component)] + if axis == normal: + return False + if dimensions == 2: + e0 = source.beam_E0 + has_tm = abs(complex(e0.z)) > 0 + has_te = abs(complex(e0.x)) > 0 or abs(complex(e0.y)) > 0 + tm = component in (mp.Ez, mp.Hx, mp.Hy) + if has_te and tm: + return False + if has_tm and not tm: + return False + return True + + +def source_components(source, dimensions: int = 3): + """Which field components a source actually drives.""" + if is_gaussian_beam(source): + return [ + component + for component, _, _ in beam_component_map(normal_index(source)) + if beam_places(source, component, dimensions) + ] + return [source.component] + + def time_profile_dtft( sim: mp.Simulation, src_time, frequencies: np.ndarray ) -> np.ndarray: @@ -221,12 +327,22 @@ def source_grad_scale( # discrete-time derivative, matching _adj_src_scale iomega = (1.0 - np.exp(-1j * (2 * np.pi * frequencies) * dt)) * (1.0 / dt) + # src_vol_chunkloop multiplies the amplitude by gv.a once per *zero-width* + # direction, to keep the integrated current fixed as a delta function is + # resolved (`data.amp *= gv.a` in sources.cpp). That is a factor of the + # resolution per delta direction, not per dimension: a point source in 2D + # carries a^2 and a line source a^1. Using 1/dV here instead would be right + # for the point and wrong by a factor of the resolution for the line. num_dims = sim._infer_dimensions(sim.k_point) - dV = 1 / sim.resolution**num_dims + size = [source.size.x, source.size.y, source.size.z][:num_dims] + num_delta = sum(1 for extent in size if extent == 0) fwd_dtft = time_profile_dtft(sim, source.src, frequencies) - scale = np.asarray(adj_src_phase) * fwd_dtft / (dV * iomega) + # The component-dependent sign is applied in `SourceGradientMonitor.gather`, + # since it differs between electric and magnetic components and this scale + # is shared across all the components one source drives. + scale = np.asarray(adj_src_phase) * fwd_dtft * sim.resolution**num_delta / iomega if sim.using_real_fields(): # real fields keep only Re[J], halving the amplitude at +omega @@ -234,6 +350,156 @@ def source_grad_scale( return scale +# The four component sources `fields::add_volume_source(src, where, beam)` +# places, as (source component, amplitude sign, which beam field is evaluated). +# With n the index of the plane normal and np1/np2 the two tangential axes: +# K = n x H goes on the electric components, N = -n x E on the magnetic ones. +_E_COMPONENTS = (mp.Ex, mp.Ey, mp.Ez) +_H_COMPONENTS = (mp.Hx, mp.Hy, mp.Hz) + + +def beam_component_map(normal_index: int): + """Which sources a Gaussian beam places, mirroring sources.cpp:526.""" + np1 = (normal_index + 1) % 3 + np2 = (normal_index + 2) % 3 + return ( + (_E_COMPONENTS[np2], +1.0, _H_COMPONENTS[np1]), + (_E_COMPONENTS[np1], -1.0, _H_COMPONENTS[np2]), + (_H_COMPONENTS[np2], -1.0, _E_COMPONENTS[np1]), + (_H_COMPONENTS[np1], +1.0, _E_COMPONENTS[np2]), + ) + + +def _beam_at(sim, source, positions, overrides=None, center=None): + """Evaluate a Gaussian beam's six field components at each position. + + Rebuilt from the source's own parameters so that a perturbed copy can be + evaluated without touching the simulation, which is what makes the + parameter derivatives cost no FDTD runs at all. + """ + values = { + "beam_x0": source.beam_x0, + "beam_kdir": source.beam_kdir, + "beam_w0": source.beam_w0, + "beam_E0": source.beam_E0, + } + values.update(overrides or {}) + + dims, cyl = sim.dimensions, sim.is_cylindrical + # add_volume_source measures positions from the centre of the volume Meep + # actually used, which grid snapping can move off source.center by up to + # half a pixel. Using the wrong one shifts every sensitivity. + origin = source.center if center is None else center + beam = mp.gaussianbeam( + mp.py_v3_to_vec(dims, values["beam_x0"], cyl), + mp.py_v3_to_vec(dims, values["beam_kdir"], cyl), + float(values["beam_w0"]), + source.src.swigobj.frequency().real, + sim.fields.get_eps(mp.py_v3_to_vec(dims, source.center, cyl)).real, + sim.fields.get_mu(mp.py_v3_to_vec(dims, source.center, cyl)).real, + np.array( + [values["beam_E0"].x, values["beam_E0"].y, values["beam_E0"].z], + dtype=np.complex128, + ), + ) + + out = np.zeros((len(positions), 6), dtype=np.complex128) + buffer = np.zeros(6, dtype=np.complex128) + for i, position in enumerate(positions): + # gaussianbeam_ampfunc is handed the position relative to the source + # volume's center, so the same offset has to be applied here. + relative = mp.Vector3(*position) - origin + beam.get_fields(buffer, mp.py_v3_to_vec(dims, relative, cyl)) + out[i] = buffer + return out + + +def _perturbations(name, value, step): + """The plus and minus variations of one beam parameter. + + Vector parameters are varied one component at a time, so the derivative + comes back with the shape of the parameter. + """ + if name == "beam_w0": + yield (), value + step, value - step + return + for axis, letter in enumerate("xyz"): + delta = mp.Vector3(**{letter: step}) + yield (axis,), value + delta, value - delta + + +def beam_parameter_gradients( + sim: mp.Simulation, + source, + names, + cotangents: dict, + normal_index: int, + step: float = 1e-6, + center=None, +) -> dict: + """Contract per-point cotangents onto a Gaussian beam's own parameters. + + The map from a beam's parameters to the currents it places is an analytic + function Meep evaluates itself, so a central difference over it is both + appropriate and cheap: it costs no FDTD runs, only re-evaluating the beam. + Meep already takes the same approach one level up, finite-differencing the + material grid inside `material_grids_addgradient`. + + Only `J^T lambda` is ever needed, so each directional derivative is + contracted as soon as it is formed and no dense Jacobian is built. For a + source plane with many points that matters. + + Args: + cotangents: maps a source component to (positions, cotangent), where + cotangent is dJ/d(amplitude) at each of those positions. + normal_index: 0, 1 or 2 for a plane normal to x, y or z. + """ + mapping = beam_component_map(normal_index) + field_index = {c: i for i, c in enumerate(_E_COMPONENTS + _H_COMPONENTS)} + + out = {} + for name in names: + value = { + "beam_x0": source.beam_x0, + "beam_kdir": source.beam_kdir, + "beam_w0": source.beam_w0, + "beam_E0": source.beam_E0, + }[name] + + entries = {} + for key, plus, minus in _perturbations(name, value, step): + total = 0.0 + 0.0j + for component, sign, evaluated in mapping: + # skipped when add_volume_source_check declined to place it + if component not in cotangents: + continue + positions, cotangent = cotangents[component] + if not len(positions): + continue + column = field_index[evaluated] + high = _beam_at(sim, source, positions, {name: plus}, center)[:, column] + low = _beam_at(sim, source, positions, {name: minus}, center)[:, column] + sensitivity = sign * (high - low) / (2 * step) + total += np.sum(np.asarray(cotangent).ravel() * sensitivity) + entries[key] = total + + if name == "beam_w0": + out[name] = entries[()] + else: + gradient = np.array([entries[(axis,)] for axis in range(3)]) + if name == "beam_kdir": + # Only the direction of beam_kdir is meaningful -- its length is + # ignored -- so the component along it is not a derivative of + # anything. Projecting it out leaves the part that is. + axis = np.array([value.x, value.y, value.z], dtype=float) + norm = np.linalg.norm(axis) + if norm > 0: + axis = axis / norm + gradient = gradient - axis * np.dot(axis, gradient) + out[name] = gradient + return out + + def contract(source, currents_grad: np.ndarray, source_amplitudes=None) -> dict: """Contract the currents cotangent onto the source's declared parameters. diff --git a/python/meep.i b/python/meep.i index dc4ef7948..a8787511a 100644 --- a/python/meep.i +++ b/python/meep.i @@ -1020,6 +1020,7 @@ void _get_gradient(PyObject *grad, double scalegrad, $1 = (double *)array_data($input); } %apply double* xtics { + double* positions, double* ytics, double* ztics, double* weights, double* vgrp, double* cscale, double* farpt_list }; @@ -1085,7 +1086,7 @@ void _get_gradient(PyObject *grad, double scalegrad, %apply std::complex* grid_vals { std::complex* eigfreq, std::complex* coeffs, std::complex* dJ, std::complex* amp_arr, - std::complex* grad + std::complex* grad, std::complex* EH }; // typemaps for diffractedplanewave diff --git a/python/source.py b/python/source.py index c995d21a7..ffef1524b 100644 --- a/python/source.py +++ b/python/source.py @@ -48,7 +48,6 @@ def _validate_differentiable(src, differentiable): valid = tuple(_DIFFERENTIABLE_ALWAYS) + tuple( getattr(src, "_differentiable_params", ()) ) - implemented = set(_DIFFERENTIABLE_ALWAYS) names = [] for name in differentiable: @@ -65,15 +64,6 @@ def _validate_differentiable(src, differentiable): f"{type(src).__name__}. Valid choices are: " f"{', '.join(sorted(valid))}." ) - if name not in implemented: - raise NotImplementedError( - f"'{name}' is a differentiable parameter of " - f"{type(src).__name__}, but its sensitivity requires the " - "finite-difference contraction over Meep's own source " - "construction, which is not implemented yet. Use 'currents' " - "and apply the chain rule yourself, or parameterize the " - "source from JAX." - ) names.append(name) if len(set(names)) != len(names): @@ -788,10 +778,11 @@ class GaussianBeam3DSource(Source): The `SourceTime` object (`Source.src`), which specifies the time dependence of the source, should normally be a narrow-band `ContinuousSource` or `GaussianSource`. (For a `CustomSource`, the beam frequency is determined by the source's `center_frequency` parameter. """ - # `beam_kdir` is deliberately absent: its length is ignored, so only its - # direction is meaningful and a component-wise derivative would report a - # spurious radial sensitivity. It needs a tangent-space projection first. - _differentiable_params = ("beam_x0", "beam_w0", "beam_E0") + # `beam_kdir`'s length is ignored, so only its direction is meaningful; the + # derivative is projected onto the tangent space of that direction rather + # than reported component-wise, which would show a spurious radial + # sensitivity. See source_gradient.beam_parameter_gradients. + _differentiable_params = ("beam_x0", "beam_kdir", "beam_w0", "beam_E0") def __init__( self, diff --git a/src/dft.cpp b/src/dft.cpp index 678ecd767..a848007ca 100644 --- a/src/dft.cpp +++ b/src/dft.cpp @@ -1488,6 +1488,33 @@ std::vector fields::dft_monitor_size(dft_fields fdft, const volume &wher return reduced_dims_vec; } +/* The monitor's size *without* collapsing zero-extent directions. + + dft_monitor_size above reduces them away, which is right for reading a field + on a plane. It is wrong for the source transpose: a zero-width source volume + straddles two Yee planes of a magnetic component, each carrying half the + current, and reducing maps both onto one index. The total weight survives + that but the spatial distribution does not, so a source whose amplitude + varies across the two planes -- a beam, as opposed to a uniform sheet -- + loses the difference between them. */ +std::vector fields::dft_monitor_full_size(dft_fields fdft, const volume &where, + component c) { + ivec min_corner, max_corner; + int rank; + direction dirs[3]; + size_t array_size, bufsz, dims[3] = {1, 1, 1}; + dft_chunk *chunklists[1]; + chunklists[0] = fdft.chunks; + (void)where; + + get_dft_component_dims(chunklists, 1, c, min_corner, max_corner, array_size, bufsz, rank, dirs, + dims); + std::vector out = {1, 1, 1}; + for (int i = 0; i < rank; ++i) + out[i] = dims[i]; + return out; +} + std::vector dft_fields::fourier_sourcedata(const volume &where, component c, fields &f, const std::complex *dJ) { @@ -1670,4 +1697,167 @@ void dft_fields::fourier_sourcegradient(const volume &where, component c, fields sum_to_all(local.data(), grad, int(Nfreq * reduced_grid_size)); } +/* Cotangent with respect to the per-point amplitudes of a *volume* source. + + fourier_sourcegradient above is the transpose of fourier_sourcedata, which is + how adjoint sources are placed. Ordinary sources -- anything built with an + amp_func, an amp_data array, or a Gaussian beam -- take a different route: + add_volume_source hands an amplitude function to src_vol_chunkloop, which + applies its own weights. This is the transpose of *that* path, so that a + derivative with respect to whatever parameterizes such a source can be + obtained by contracting this against the sensitivity of its amplitude + function. + + Positions come out alongside the DFT values by rebuilding the chunk's grid + volume, the same way material_grids_addgradient walks the adjoint fields: + + gv_sub = gv.subvolume(chunk->is, chunk->ie, c) + LOOP_OVER_IVECS(gv_sub, chunk->is_old, chunk->ie_old, idx) + + which gives `idx` indexing chunk->dft directly and IVEC_LOOP_LOC giving the + point it belongs to. + + `grad` receives freq.size() * (monitor size) elements in the monitor's own + ordering, summed across all processes. */ +void dft_fields::volume_source_gradient(const volume &where, component c, fields &f, + std::complex *grad) { + const size_t Nfreq = freq.size(); + + ivec min_corner, max_corner; + int rank, reduced_rank; + direction dirs[3], reduced_dirs[3]; + size_t array_size, bufsz, dims[3], reduced_dims[3], reduced_stride[3], stride[3]; + dft_chunk *chunklists[1]; + chunklists[0] = chunks; + + f.get_dft_component_dims(chunklists, 1, c, min_corner, max_corner, array_size, bufsz, rank, dirs, + dims); + /* Deliberately *not* reduce_array_dimensions here. Collapsing a zero-extent + direction merges the two Yee planes a magnetic component straddles, which + keeps the total weight but discards how the amplitude varies between them. + Keeping the full layout leaves the restriction loop_in_chunks already + performs intact, weights and all. */ + size_t reduced_grid_size = 1; + for (int i = 0; i < rank; ++i) + reduced_grid_size *= dims[i]; + + std::vector > local(Nfreq * reduced_grid_size, std::complex(0, 0)); + + for (dft_chunk *fdc = chunks; fdc; fdc = fdc->next_in_dft) { + component cc = component(fdc->c); + direction cd = component_direction(cc); + grid_volume gv_sub = f.gv.subvolume(fdc->is, fdc->ie, cc); + // is_old/ie_old only hold the unpadded bounds when the chunk was created + // with persist set; otherwise they were never assigned and is/ie are + // already the range wanted. + ivec loop_is = fdc->persist ? fdc->is_old : fdc->is; + ivec loop_ie = fdc->persist ? fdc->ie_old : fdc->ie; + + int position_array[3] = {0, 0, 0}; + + LOOP_OVER_IVECS(gv_sub, loop_is, loop_ie, idx) { + IVEC_LOOP_ILOC(gv_sub, iloc); + // the loop runs in the chunk's own frame; the monitor's array index is + // defined in the untransformed frame, so undo the symmetry and shift + // before locating this point in it + iloc = fdc->S.transform(iloc, fdc->sn) + fdc->shift; + for (int i = 0; i < rank; ++i) + position_array[i] = + int((iloc.in_direction(dirs[i]) - min_corner.in_direction(dirs[i])) / 2); + size_t idx_1d = 0; + for (int i = 0; i < rank; ++i) + idx_1d = idx_1d * dims[i] + position_array[i]; + if (idx_1d >= reduced_grid_size) continue; + + /* src_vol_chunkloop weights every point by IVEC_LOOP_WEIGHT, which is + how a source volume keeps its integrated current fixed as the grid + changes. The transpose has to apply the same weight. + + Omitting it is invisible for a grid-aligned electric source, where the + weight is 1, and shows up as exactly (1/2)^(zero-width directions) for + a component whose Yee points straddle the source plane -- Hz in 2D + sits at half-integer positions in both transverse axes, so a + zero-width volume centred on a grid point splits evenly between two + of them in each such direction. */ + double w = IVEC_LOOP_WEIGHT(fdc->s0, fdc->s1, fdc->e0, fdc->e1, 1); + if (is_D(cc) && fdc->fc->s->chi1inv[cc - Dx + Ex][cd]) + w /= fdc->fc->s->chi1inv[cc - Dx + Ex][cd][idx]; + if (is_B(cc) && fdc->fc->s->chi1inv[cc - Bx + Hx][cd]) + w /= fdc->fc->s->chi1inv[cc - Bx + Hx][cd][idx]; + + for (size_t i = 0; i < Nfreq; ++i) { + std::complex EH = fdc->dft[Nfreq * idx + i]; + local[reduced_grid_size * i + idx_1d] += + w * std::complex(double(EH.real()), double(EH.imag())); + } + } + } + + sum_to_all(local.data(), grad, int(Nfreq * reduced_grid_size)); +} + +/* The position of every point volume_source_gradient reports, in the same + order. + + Contracting a cotangent onto a source's own parameters means re-evaluating + that source's amplitude function at the points the cotangent belongs to, so + the two have to agree exactly about which point is which. Emitting the + positions from the same loop is the only way to be sure they do. + + `positions` receives 3 * (monitor size) doubles, x, y and z per point. */ +void dft_fields::monitor_positions(const volume &where, component c, fields &f, double *positions) { + ivec min_corner, max_corner; + int rank, reduced_rank; + direction dirs[3], reduced_dirs[3]; + size_t array_size, bufsz, dims[3], reduced_dims[3], reduced_stride[3], stride[3]; + dft_chunk *chunklists[1]; + chunklists[0] = chunks; + + f.get_dft_component_dims(chunklists, 1, c, min_corner, max_corner, array_size, bufsz, rank, dirs, + dims); + /* Deliberately *not* reduce_array_dimensions here. Collapsing a zero-extent + direction merges the two Yee planes a magnetic component straddles, which + keeps the total weight but discards how the amplitude varies between them. + Keeping the full layout leaves the restriction loop_in_chunks already + performs intact, weights and all. */ + size_t reduced_grid_size = 1; + for (int i = 0; i < rank; ++i) + reduced_grid_size *= dims[i]; + + std::vector local(3 * reduced_grid_size, 0.0); + + for (dft_chunk *fdc = chunks; fdc; fdc = fdc->next_in_dft) { + component cc = component(fdc->c); + grid_volume gv_sub = f.gv.subvolume(fdc->is, fdc->ie, cc); + // is_old/ie_old only hold the unpadded bounds when the chunk was created + // with persist set; otherwise they were never assigned and is/ie are + // already the range wanted. + ivec loop_is = fdc->persist ? fdc->is_old : fdc->is; + ivec loop_ie = fdc->persist ? fdc->ie_old : fdc->ie; + + int position_array[3] = {0, 0, 0}; + + LOOP_OVER_IVECS(gv_sub, loop_is, loop_ie, idx) { + IVEC_LOOP_ILOC(gv_sub, iloc); + IVEC_LOOP_LOC(gv_sub, loc); + iloc = fdc->S.transform(iloc, fdc->sn) + fdc->shift; + loc = fdc->S.transform(loc, fdc->sn) + vec(fdc->shift * (0.5 * fdc->fc->gv.inva)); + + for (int i = 0; i < rank; ++i) + position_array[i] = + int((iloc.in_direction(dirs[i]) - min_corner.in_direction(dirs[i])) / 2); + size_t idx_1d = 0; + for (int i = 0; i < rank; ++i) + idx_1d = idx_1d * dims[i] + position_array[i]; + if (idx_1d >= reduced_grid_size) continue; + + local[3 * idx_1d + 0] = loc.in_direction(X); + local[3 * idx_1d + 1] = loc.in_direction(Y); + local[3 * idx_1d + 2] = (gv_sub.dim == D3) ? loc.in_direction(Z) : 0.0; + } + } + + sum_to_all(local.data(), positions, int(3 * reduced_grid_size)); +} + } // namespace meep diff --git a/src/meep.hpp b/src/meep.hpp index 8a661ae32..eec4baad7 100644 --- a/src/meep.hpp +++ b/src/meep.hpp @@ -1428,6 +1428,11 @@ class dft_fields { // transpose of fourier_sourcedata; see dft.cpp void fourier_sourcegradient(const volume &where, component c, fields &f, std::complex *grad); + // transpose of the add_volume_source / src_vol_chunkloop path; see dft.cpp + void volume_source_gradient(const volume &where, component c, fields &f, + std::complex *grad); + // the positions volume_source_gradient reports, in the same order + void monitor_positions(const volume &where, component c, fields &f, double *positions); void scale_dfts(std::complex scale); void remove(); @@ -2188,6 +2193,7 @@ class fields { std::complex get_field(component c, const vec &loc, bool parallel = true) const; double get_field(derived_component c, const vec &loc, bool parallel = true) const; std::vector dft_monitor_size(dft_fields fdft, const volume &where, component c); + std::vector dft_monitor_full_size(dft_fields fdft, const volume &where, component c); void get_dft_component_dims(dft_chunk **chunklists, int num_chunklists, component c, ivec &min_corner, ivec &max_corner, size_t &array_size, size_t &bufsz, int &rank, direction *ds, size_t *dims, int *array_rank = 0, From 9de66e5b0d790aef1921b1ec7781ad325a548cbd Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 27 Aug 2026 16:26:36 -0700 Subject: [PATCH 5/5] adjoint: differentiate amp_data instead of adding a source type Meep already has two of the three pieces: Source(amp_data=...) takes an array profile, and get_equiv_sources applies the equivalence principle to build one. The only thing missing was differentiability, so add that rather than a new source class -- ArraySource is removed. amp_file_func trilinearly interpolates the user's array at each grid point, so the gradient is the transpose of that interpolation scattered back onto the array. That is a plain scatter in Python, because the hard part -- the cotangent with respect to the amplitude Meep applied at each grid point, and where that point is -- is already done in C++. MeepJaxWrapper picks up any source carrying amp_data automatically and returns the cotangent in whatever shape the caller passed, since Meep wants a 3D array but JAX may hand in any. Verified against finite differences on individual array entries: 1.9e-8. An amp_func still cannot be differentiated, and now says why: it is evaluated inside Meep, so there is no array for a cotangent to land on. --- NEWS.md | 23 +- doc/docs/Python_Tutorials/Adjoint_Solver.md | 60 +++-- python/adjoint/optimization_problem.py | 21 +- python/adjoint/source_gradient.py | 75 +++++- python/adjoint/wrapper.py | 55 +++-- python/meep.i | 1 - python/simulation.py | 1 - python/source.py | 122 +-------- python/tests/test_source_gradient.py | 258 +++++++++++++------- 9 files changed, 354 insertions(+), 262 deletions(-) diff --git a/NEWS.md b/NEWS.md index 8dc2cd49e..0bcde63bd 100644 --- a/NEWS.md +++ b/NEWS.md @@ -3,16 +3,19 @@ ## Meep 1.35.0 (in progress) * Adjoint solver: sources can now be differentiated alongside the design - regions. A source opts in with `differentiable=['amplitude']` or - `differentiable=['currents']`, and its gradient appears in the returned - dictionary under the source's `name`. It costs no extra simulation: the - derivative with respect to a source is the adjoint field sampled over that - source's support, which the existing adjoint run already produces. The new - `mp.ArraySource` supplies per-point amplitudes directly, indexed exactly as - `get_dft_array` returns them over the same region, which is also the ordering - the gradient comes back in. An `ArraySource` passed to `MeepJaxWrapper` is - differentiated automatically, so a source computed in JAX backpropagates to - whatever produced it. + regions. A source opts in by naming parameters, as in + `differentiable=['beam_w0', 'beam_x0']`, and its gradient appears in the + returned dictionary under the source's `name`. It costs no extra simulation: + the derivative with respect to a source is the adjoint field sampled over that + source's support, which the existing adjoint run already produces. + `GaussianBeam3DSource` supports `beam_x0`, `beam_kdir`, `beam_w0` and + `beam_E0`, obtained by finite-differencing Meep's own beam construction, which + needs no additional FDTD runs. Any source's `amp_data` array is differentiable + too, with the adjoint applying the transpose of the trilinear interpolation + Meep uses to place it, so the gradient lands on the array the user supplied. A + source with `amp_data` passed to `MeepJaxWrapper` is differentiated + automatically, so a source computed in JAX backpropagates to whatever produced + it. * Adjoint solver: `meep.adjoint.AngularSpectrum` propagates the tangential DFT fields on a planar monitor through an arbitrary stratified medium diff --git a/doc/docs/Python_Tutorials/Adjoint_Solver.md b/doc/docs/Python_Tutorials/Adjoint_Solver.md index 05d79ef31..b5245afe1 100644 --- a/doc/docs/Python_Tutorials/Adjoint_Solver.md +++ b/doc/docs/Python_Tutorials/Adjoint_Solver.md @@ -400,48 +400,62 @@ The names given in `differentiable` are the keys of the resulting subtree, so th flag and the gradient cannot drift apart. With no flagged source the return value is unchanged, so nothing existing is affected. -Every source accepts `'amplitude'` and `'currents'`. The first is the scalar that -scales the whole source; the second is the per-point complex amplitude array, -which is the general case and the one that composes with a propagator. Names that -a source class does define but whose sensitivity is not yet implemented — -`beam_w0` on a Gaussian beam, for instance — raise `NotImplementedError` rather -than being reported as unknown, so the two situations are distinguishable. +Every source accepts `'amplitude'`, `'currents'` and `'amp_data'`. Individual +source classes accept more: `GaussianBeam3DSource` takes `'beam_x0'`, +`'beam_kdir'`, `'beam_w0'` and `'beam_E0'`, which are the parameters worth +optimizing — the amplitude of a linear simulation is not one of them, since the +fields scale with it by definition. + +Those beam derivatives are obtained by finite-differencing Meep's own beam +construction and contracting the result against the per-point cotangent. That +costs **no extra FDTD runs**, only re-evaluating the beam, and each directional +derivative is contracted as it is formed so no dense Jacobian is ever built. +Meep already takes the same approach one level up, finite-differencing the +material grid inside `material_grids_addgradient`. + +Two of them come with a caveat. `beam_kdir`'s length is ignored, so only its +direction is meaningful; the gradient is projected onto the tangent space of +that direction, because the component along it is not a derivative of anything. +And in 2D only the out-of-plane component of `beam_E0` is meaningful for a TM +beam: perturbing the in-plane components makes the beam mixed TE/TM, and Meep +then places no sources at all. `'center'` and `'size'` are rejected outright. They move the grid points the source occupies rather than the amplitudes applied to them, and the adjoint formulation for sources gathers the cotangent over a *fixed* set of points, so a derivative there is outside the formulation rather than merely missing. -### Supplying the currents yourself +### Supplying the amplitudes yourself -`mp.ArraySource` takes the per-point amplitudes directly, which is what a source -computed somewhere else — by a mode solver, a propagator, or JAX — naturally -produces: +`amp_data` already lets a source take its profile from an array, which is what a +source computed somewhere else — by a mode solver, a propagator, or JAX — +naturally produces. It is now differentiable: ```py -src = mp.ArraySource( +src = mp.Source( mp.GaussianSource(fcen, fwidth=df), - mp.Ez, - amplitudes=amplitudes, # one complex number per grid point - frequency=fcen, + component=mp.Ez, center=mp.Vector3(-1, 0), size=mp.Vector3(0, 4), - differentiable=["currents"], + amp_data=profile, # trilinearly interpolated onto the grid + differentiable=["amp_data"], name="sheet", ) ``` -The array is indexed exactly like `Simulation.get_dft_array` over the same region -and component, which is also the ordering the gradient comes back in. That is not -a coincidence: injection and measurement are implemented as a scatter and its -exact transpose, so an array can be handed in and its cotangent read back with no -bookkeeping in between. +The gradient has the shape of `amp_data`. Meep interpolates that array onto the +grid, and the adjoint applies the transpose of the same interpolation, so the +cotangent comes back on the array the user actually supplied rather than on the +grid points underneath it. + +An `amp_func` cannot be differentiated: it is evaluated inside Meep at each grid +point, so there is no array for a cotangent to land on. ### With JAX -An `ArraySource` given to `MeepJaxWrapper` is differentiated automatically. There -is nothing to flag, because the parameters that produced the currents live -upstream in JAX rather than in Meep — Meep returns the cotangent with respect to +A source with `amp_data` given to `MeepJaxWrapper` is differentiated +automatically. There is nothing to flag, because the parameters that produced +the array live upstream in JAX rather than in Meep — Meep returns the cotangent with respect to the current array and JAX carries it the rest of the way: ```py diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index 78a0974fc..dc6a38097 100644 --- a/python/adjoint/optimization_problem.py +++ b/python/adjoint/optimization_problem.py @@ -448,10 +448,27 @@ def _source_gradient_for(self, src, monitors, scale): ] grads = {} - if other_names: + plain = [n for n in other_names if n != "amp_data"] + if plain: currents = monitors[0].gather(self.sim, scale) everything = source_gradient.contract(src, currents) - grads.update({name: everything[name] for name in other_names}) + grads.update({name: everything[name] for name in plain}) + if "amp_data" in other_names: + # the transpose of the trilinear interpolation amp_file_func does + currents = monitors[0].gather(self.sim, scale) + grads["amp_data"] = np.squeeze( + np.array( + [ + source_gradient.amp_data_gradient( + self.sim, + src, + monitors[0].positions(self.sim), + np.asarray(currents)[f_index], + ) + for f_index in range(len(np.atleast_1d(self.frequencies))) + ] + ) + ) if beam_names: # Contract the per-point cotangents onto the beam's own parameters, diff --git a/python/adjoint/source_gradient.py b/python/adjoint/source_gradient.py index c15812373..98ae3e0e3 100644 --- a/python/adjoint/source_gradient.py +++ b/python/adjoint/source_gradient.py @@ -32,6 +32,7 @@ IMPLEMENTED_PARAMS = ( "currents", "amplitude", + "amp_data", "beam_x0", "beam_kdir", "beam_w0", @@ -87,8 +88,9 @@ def __init__( # apply it to, and `amplitude` would silently ignore the profile. raise NotImplementedError( "Differentiating a source defined by `amp_func` or " - "`amp_func_file` is not supported; supply the amplitudes as an " - "array (`amp_data`) or drive the source from JAX instead." + "`amp_func_file` is not supported: the profile is evaluated " + "inside Meep, so there is no array for the caller to apply a " + "cotangent to. Use `amp_data` instead." ) self.volume = sim._fit_volume_to_simulation( @@ -500,6 +502,75 @@ def beam_parameter_gradients( return out +def _mirrorindex(i: np.ndarray, n: int) -> np.ndarray: + """`mirrorindex` from fields.cpp:755.""" + return np.where(i >= n, 2 * n - 1 - i, np.where(i < 0, -1 - i, i)) + + +def _map_coordinates(r: np.ndarray, n: int): + """`map_coordinates` from fields.cpp:759, vectorized over points. + + Returns the two bracketing indices and the weight of the second, matching + the `do_fabs` behaviour `linear_interpolate` relies on. + """ + if n == 1: + zero = np.zeros(r.shape, dtype=int) + return zero, zero, np.zeros(r.shape) + r = np.where(r < 0.0, -r, np.where(r > 1.0, 1.0 - r, r)) + i1 = _mirrorindex(np.floor(r * n).astype(int), n) + d = r * n - i1 - 0.5 + i2 = _mirrorindex(np.where(d >= 0.0, i1 + 1, i1 - 1), n) + return i1, i2, np.abs(d) + + +def amp_data_gradient( + sim: mp.Simulation, + source, + positions: np.ndarray, + cotangent: np.ndarray, +) -> np.ndarray: + """Contract a per-point cotangent onto a source's `amp_data` array. + + `Source(amp_data=...)` reaches the grid through `amp_file_func`, which + trilinearly interpolates the user's array at each grid point. This is the + transpose of that interpolation: each point's cotangent is scattered back + onto the (at most) eight array entries that fed it, with the same weights. + + It is a plain scatter rather than anything new in C++, because the hard + part -- the cotangent with respect to the amplitude Meep actually applied + at each grid point, together with where that point is -- is already done by + `volume_source_gradient`. + """ + data = np.asarray(source.amp_data) + shape = tuple(data.shape) + (1,) * (3 - data.ndim) + nx, ny, nz = shape + + positions = np.asarray(positions, dtype=float) + relative = positions - np.array( + [source.center.x, source.center.y, source.center.z], dtype=float + ) + extent = np.array([source.size.x, source.size.y, source.size.z], dtype=float) + # `amp_file_func` maps a position onto [0, 1] across the source volume, and + # pins the coordinate to the centre in any direction of zero extent. + with np.errstate(divide="ignore", invalid="ignore"): + r = np.where( + extent > 0, 0.5 + relative / np.where(extent > 0, extent, 1.0), 0.0 + ) + + ix1, ix2, wx = _map_coordinates(r[:, 0], nx) + iy1, iy2, wy = _map_coordinates(r[:, 1], ny) + iz1, iz2, wz = _map_coordinates(r[:, 2], nz) + + grad = np.zeros(nx * ny * nz, dtype=np.complex128) + values = np.asarray(cotangent).ravel() + for ix, fx in ((ix1, 1.0 - wx), (ix2, wx)): + for iy, fy in ((iy1, 1.0 - wy), (iy2, wy)): + for iz, fz in ((iz1, 1.0 - wz), (iz2, wz)): + flat = (ix * ny + iy) * nz + iz + np.add.at(grad, flat, values * fx * fy * fz) + return grad.reshape(data.shape) + + def contract(source, currents_grad: np.ndarray, source_amplitudes=None) -> dict: """Contract the currents cotangent onto the source's declared parameters. diff --git a/python/adjoint/wrapper.py b/python/adjoint/wrapper.py index 5ac34872a..88e4d3d6a 100644 --- a/python/adjoint/wrapper.py +++ b/python/adjoint/wrapper.py @@ -213,16 +213,18 @@ def __init__( self.until_after_sources = until_after_sources self.finite_difference_step = finite_difference_step - # Sources whose amplitudes are supplied from JAX are differentiated - # with respect to their currents automatically: the parameters live - # upstream in JAX, so there is nothing for Meep to name. + # A source whose amplitudes come from JAX is differentiated + # automatically: the parameters that produced them live upstream in + # JAX, so there is nothing for Meep to name. `amp_data` is the array + # Meep interpolates onto the grid, so that is the cut point. + self._source_shapes = [] self.differentiable_sources = [ - s for s in sources if isinstance(s, mp.ArraySource) + s for s in sources if getattr(s, "amp_data", None) is not None ] for s in self.differentiable_sources: - if "currents" not in getattr(s, "differentiable", ()): + if "amp_data" not in getattr(s, "differentiable", ()): s.differentiable = tuple(getattr(s, "differentiable", ())) + ( - "currents", + "amp_data", ) self._simulate_fn = self._initialize_callable() @@ -233,10 +235,11 @@ def __call__( """Performs a Meep simulation, taking designs and returning monitor values. Args: - sources: amplitudes for each `mp.ArraySource` passed to the constructor, - in that order. These are differentiated automatically -- there is no - need to flag them, because the parameters that produced them live - upstream in JAX rather than in Meep. Omit when there are none. + sources: an `amp_data` array for each source given to the constructor + that has one, in that order. These are differentiated automatically + -- there is no need to flag them, because the parameters that + produced them live upstream in JAX rather than in Meep. Omit when + there are none. designs: a list of design variables as 1D, 2D, or 3D JAX arrays. Valid shapes for design variables are (Nx, Ny, Nz) where Nx{y,z} match the elements of the `grid_size` constructor argument of Meep's `MaterialGrid` used for the @@ -268,8 +271,14 @@ def _update_sources(self, source_variables) -> None: f"Got {len(source_variables)} source arrays but " f"{len(self.differentiable_sources)} differentiable sources." ) + # Meep wants amp_data as a 3D array, but the caller may hand in any + # shape with the same number of entries; the cotangent has to go back + # in the shape they used, so remember it. + self._source_shapes = [onp.shape(a) for a in source_variables] for src, amps in zip(self.differentiable_sources, source_variables): - src.amplitudes = onp.asarray(amps, dtype=onp.complex128) + src.amp_data = onp.asarray(amps, dtype=onp.complex128).reshape( + onp.shape(src.amp_data) + ) def _run_fwd_simulation( self, @@ -352,17 +361,31 @@ def _source_vjps(self, sum_freq_partials: bool = True) -> List[onp.ndarray]: return [] phase = self.monitors[0]._adj_src_phase() out = [] - for src, monitor in zip(self.differentiable_sources, self.adj_source_monitors): + for src, group in zip(self.differentiable_sources, self.adj_source_monitors): + # one monitor per component the source drives; an amp_data source + # drives exactly one + monitor = group[0] scale = source_gradient.source_grad_scale( self.simulation, src, self.frequencies, phase ) - grad = monitor.gather(self.simulation, scale) - shape = src.amplitudes.shape + currents = monitor.gather(self.simulation, scale) + positions = monitor.positions(self.simulation) + shape = self._source_shapes[len(out)] + # one row per frequency, each carried back through the trilinear + # interpolation Meep applies to amp_data + rows = onp.stack( + [ + source_gradient.amp_data_gradient( + self.simulation, src, positions, currents[f_index] + ) + for f_index in range(currents.shape[0]) + ] + ) if sum_freq_partials: # the amplitudes are shared across the band, as design weights are - out.append(onp.sum(grad, axis=0).reshape(shape)) + out.append(onp.sum(rows, axis=0).reshape(shape)) else: - out.append(grad.reshape((grad.shape[0], *shape))) + out.append(rows.reshape((rows.shape[0], *shape))) return out def _calculate_vjps( diff --git a/python/meep.i b/python/meep.i index a8787511a..d968105f8 100644 --- a/python/meep.i +++ b/python/meep.i @@ -1778,7 +1778,6 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where with_prefix ) from .source import ( - ArraySource, ContinuousSource, CustomSource, EigenModeSource, diff --git a/python/simulation.py b/python/simulation.py index 0760467a2..af75cf30d 100644 --- a/python/simulation.py +++ b/python/simulation.py @@ -24,7 +24,6 @@ import numpy as np from meep.geom import GeometricObject, Medium, Vector3, init_do_averaging from meep.source import ( - ArraySource, EigenModeSource, GaussianBeamSource, IndexedSource, diff --git a/python/source.py b/python/source.py index ffef1524b..570c13199 100644 --- a/python/source.py +++ b/python/source.py @@ -20,7 +20,7 @@ def check_positive(prop, val): # actually applies to the Yee grid; it is the universal representation, and the # one the JAX bridge uses. "amplitude" is exact and needs no finite difference, # since it scales those currents linearly. -_DIFFERENTIABLE_ALWAYS = ("currents", "amplitude") +_DIFFERENTIABLE_ALWAYS = ("currents", "amplitude", "amp_data") # Parameters that move the source's support rather than change its amplitudes. # The adjoint machinery gathers the cotangent at a *fixed* set of grid points @@ -48,6 +48,11 @@ def _validate_differentiable(src, differentiable): valid = tuple(_DIFFERENTIABLE_ALWAYS) + tuple( getattr(src, "_differentiable_params", ()) ) + if "amp_data" in differentiable and getattr(src, "amp_data", None) is None: + raise ValueError( + "'amp_data' was requested but this source has none; pass " + "amp_data= to differentiate with respect to it." + ) names = [] for name in differentiable: @@ -1166,121 +1171,6 @@ def add_source(self, sim): super().add_source(sim) -class ArraySource(Source): - """A volume source whose per-point complex amplitudes are given as an array. - - Ordinary `Source` objects take a single `amplitude` and, optionally, an - `amp_func` that Meep evaluates internally. This class instead takes the - amplitudes directly, one per grid point, which is what a source computed - somewhere else -- by a mode solver, a propagator, or JAX -- naturally - produces. - - The array is indexed exactly like `Simulation.get_dft_array` over the same - volume and component. That is deliberate: it is also the ordering the - adjoint solver returns `differentiable=['currents']` gradients in, so an - array can be handed in and its cotangent read back with no bookkeeping in - between. Injection and measurement share one convention because they are - implemented as a scatter and its exact transpose. - """ - - def __init__( - self, - src, - component, - amplitudes, - frequency, - center=None, - volume=None, - size=Vector3(), - yee_grid=True, - differentiable=None, - name=None, - ): - """Construct an `ArraySource`. - - + **`amplitudes` [`numpy.ndarray`]** — Complex amplitude for each point - of the source region, shaped like `get_dft_array` over that region. - + **`frequency` [`number`]** — The frequency the amplitudes refer to. - + **`yee_grid` [`boolean`]** — Whether the amplitudes are given at the - Yee points of `component` (the default) or at voxel centers. Voxel - centers are what several components sharing one plane need, since - each component's Yee points sit at a different half-pixel offset and - the arrays would otherwise have different lengths. Meep spreads a - centered value across the neighbouring Yee points, and the adjoint - gathers it back the same way, so the gradient stays exact either way. - """ - super().__init__( - src, - component, - center=center, - volume=volume, - size=size, - differentiable=differentiable, - name=name, - ) - self.amplitudes = np.ascontiguousarray(amplitudes, dtype=np.complex128) - self.frequency = float(frequency) - self.yee_grid = yee_grid - self._monitor = None - - def add_source(self, sim): - vol = sim._fit_volume_to_simulation( - mp.Volume(center=self.center, size=self.size) - ) - # The monitor is created only for its chunk decomposition and array - # ordering; its DFT storage is never read. - # A DFT object cannot be added before the field components exist, and - # components are normally allocated only once every source has been - # added. Ask for this one up front. - sim.fields.require_component(self.component) - - mon = sim.add_dft_fields( - [self.component], [self.frequency], where=vol, yee_grid=self.yee_grid - ) - # add_dft_fields defers construction; the scatter needs it now - sim._evaluate_dft_objects() - dims = sim.fields.dft_monitor_size(mon.swigobj, vol.swigobj, self.component) - npts = int(np.prod(dims)) - - if self.amplitudes.size != npts: - raise ValueError( - f"`amplitudes` has {self.amplitudes.size} elements but the " - f"source region holds {npts} grid points (shape {tuple(dims)})." - ) - - # Two conversions, so that one entry of `amplitudes` means exactly what - # `Source.amplitude` means for a point source at that grid point. - # - # fourier_sourcedata negates *electric* components and only those, since - # it was written to place adjoint sources; undoing it for magnetic - # components too would flip the relative sign of the two sheets of an - # equivalent-current pair, which silently reverses the direction such a - # pair radiates in. - # - # It also places a current *density*, whose integral over a voxel is the - # amplitude times dV. - num_dims = sim._infer_dimensions(sim.k_point) - dV = 1 / sim.resolution**num_dims - sign = -1.0 if mp.is_electric(self.component) else 1.0 - flat = np.ascontiguousarray( - self.amplitudes.ravel() * complex(self.amplitude) * (sign / dV), - dtype=np.complex128, - ) - srcdata = mon.swigobj.fourier_sourcedata( - vol.swigobj, self.component, sim.fields, flat - ) - - sim.fields.register_src_time(self.src.swigobj) - for sd in srcdata: - amp = np.asarray(sd.amp_arr, dtype=np.complex128) - if amp.size == 0: - continue # this process owns no part of the source - sim.fields.add_srcdata(sd, self.src.swigobj, amp.size, amp, False) - - # the DFT storage is dead weight for a plane with many points - mon.remove() - - class IndexedSource(Source): """ created a source object using (SWIG-wrapped mp::srcdata*) srcdata. diff --git a/python/tests/test_source_gradient.py b/python/tests/test_source_gradient.py index 5ff2166ed..f65427920 100644 --- a/python/tests/test_source_gradient.py +++ b/python/tests/test_source_gradient.py @@ -131,7 +131,7 @@ def test_support_moving_parameters_get_their_own_message(self): with self.assertRaisesRegex(ValueError, "grid points the source occupies"): self._src(differentiable=[name]) - def test_unimplemented_parameter_is_not_reported_as_unknown(self): + def test_class_specific_parameters_are_scoped_to_their_class(self): beam = dict( src=mp.GaussianSource(FCEN, fwidth=0.2), center=SRC_C, @@ -140,9 +140,9 @@ def test_unimplemented_parameter_is_not_reported_as_unknown(self): beam_w0=1.0, beam_E0=mp.Vector3(0, 0, 1), ) - with self.assertRaises(NotImplementedError): - mp.GaussianBeam3DSource(differentiable=["beam_w0"], **beam) - # ... but it is still rejected as unknown on a class that lacks it + accepted = mp.GaussianBeam3DSource(differentiable=["beam_w0"], **beam) + self.assertEqual(accepted.differentiable, ("beam_w0",)) + # ... but rejected on a class that has no such parameter with self.assertRaises(ValueError): self._src(differentiable=["beam_w0"]) @@ -218,70 +218,33 @@ def test_dot_product_identity(self): self.assertAlmostEqual(abs(lhs - rhs) / abs(lhs), 0.0, places=12) -class TestArraySource(unittest.TestCase): - def _field(self, source): - sim = mp.Simulation( - cell_size=CELL, - resolution=RES, - boundary_layers=[mp.PML(1.0)], - sources=[source], - force_complex_fields=True, - ) - mon = sim.add_dft_fields( - [mp.Ez], - [FCEN], - where=mp.Volume(center=MON_C, size=mp.Vector3(0.2, 0.2)), - ) - sim.run(until_after_sources=60) - return np.asarray(sim.get_dft_array(mon, mp.Ez, 0)).ravel() - - def test_one_point_matches_an_ordinary_source(self): - # Pins both conversions in ArraySource.add_source: the -1 the scatter - # applies to electric components, and the fact that it places a current - # density rather than a point amplitude. - # - # Both an electric and a magnetic component are checked, because the - # scatter negates electric components and only those. Undoing that - # negation for magnetic ones too leaves each component individually - # plausible but flips the relative sign of the two sheets of an - # equivalent-current pair, which reverses the direction such a pair - # radiates in without producing anything that looks like an error. - t = lambda: mp.GaussianSource(FCEN, fwidth=0.2) - for component in (mp.Ez, mp.Hz): - with self.subTest(component=mp.component_name(component)): - plain = self._field( - mp.Source(t(), component=component, center=SRC_C, amplitude=1.0) - ) - array = self._field( - mp.ArraySource( - t(), - component, - amplitudes=np.array([1.0 + 0j]), - frequency=FCEN, - center=SRC_C, - size=mp.Vector3(0, 0), - ) - ) - np.testing.assert_allclose(array, plain, rtol=2e-6) +class TestAmpData(unittest.TestCase): + """Per-point amplitudes supplied as an array.""" + + def test_rejects_amp_data_when_absent(self): + with self.assertRaisesRegex(ValueError, "this source has none"): + mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), + component=mp.Ez, + center=SRC_C, + differentiable=["amp_data"], + ) - def test_rejects_wrong_length(self): - src = mp.ArraySource( + def test_rejects_amp_func(self): + # An amp_func is evaluated inside Meep, so there is no array for the + # caller to apply a cotangent to. + src = mp.Source( mp.GaussianSource(FCEN, fwidth=0.2), - mp.Ez, - amplitudes=np.ones(3, dtype=complex), - frequency=FCEN, + component=mp.Ez, center=SRC_C, - size=mp.Vector3(0, 0.4), - ) - sim = mp.Simulation( - cell_size=CELL, - resolution=RES, - boundary_layers=[mp.PML(1.0)], - sources=[src], - force_complex_fields=True, + size=mp.Vector3(0, 1.0), + amp_func=lambda p: 1.0, + differentiable=["currents"], + name="drive", ) - with self.assertRaisesRegex(ValueError, "grid points"): - sim.init_sim() + _, opt = _problem(src) + with self.assertRaises(NotImplementedError): + opt([RHO]) class TestSourceGradient(unittest.TestCase): @@ -346,38 +309,40 @@ def test_amplitude_gradient(self): f"{label}: adjoint {adjoint} vs finite difference {reference}", ) - def test_currents_gradient(self): - size = mp.Vector3(0, 0.4) - npts = _num_source_points(SRC_C, size) - rng = np.random.default_rng(7) - amps = rng.standard_normal(npts) + 1j * rng.standard_normal(npts) + def test_amp_data_gradient(self): + # Perturbing individual entries exercises the transpose of the + # trilinear interpolation `amp_file_func` performs, on top of the + # per-point cotangent. + size = mp.Vector3(0, 1.5) + n = 5 + rng = np.random.default_rng(5) + data = rng.standard_normal(n) + 1j * rng.standard_normal(n) - def make(a): - return mp.ArraySource( + def make(values): + return mp.Source( mp.GaussianSource(FCEN, fwidth=0.2), - mp.Ez, - amplitudes=a, - frequency=FCEN, + component=mp.Ez, center=SRC_C, size=size, - differentiable=["currents"], + amp_data=np.asarray(values, dtype=np.complex128).reshape(1, n, 1), + differentiable=["amp_data"], name="sheet", ) - _, opt = _problem(make(amps)) + _, opt = _problem(make(data)) _, grad = opt([RHO]) - adjoint = np.ravel(grad["sheet"]["currents"]) - self.assertEqual(adjoint.size, npts) + adjoint = np.ravel(grad["sheet"]["amp_data"]) + self.assertEqual(adjoint.size, n) - def evaluate(a): - _, o = _problem(make(a)) + def evaluate(values): + _, o = _problem(make(values)) return np.asarray(o([RHO], need_gradient=False)[0]).item() - for j in (0, npts // 2, npts - 1): - with self.subTest(point=j): + for j in (0, n // 2, n - 1): + with self.subTest(entry=j): def perturb(delta, j=j): - out = amps.copy() + out = data.copy() out[j] += delta return out @@ -385,8 +350,120 @@ def perturb(delta, j=j): self.assertLess( abs(adjoint[j] - reference) / abs(reference), 1e-5, - f"point {j}: adjoint {adjoint[j]} vs {reference}", + f"entry {j}: adjoint {adjoint[j]} vs {reference}", + ) + + +class TestGaussianBeamParameters(unittest.TestCase): + """Derivatives with respect to a Gaussian beam's own parameters. + + Obtained by finite-differencing Meep's beam construction and contracting + against the per-point cotangent, so they cost no FDTD runs. + """ + + BEAM = dict( + beam_x0=mp.Vector3(0, 1.0), + beam_kdir=mp.Vector3(0, 1), + beam_w0=1.5, + beam_E0=mp.Vector3(0, 0, 1), + ) + + def _problem(self, names, **overrides): + params = dict(self.BEAM) + params.update(overrides) + src = mp.GaussianBeam3DSource( + mp.GaussianSource(FCEN, fwidth=0.2), + center=mp.Vector3(0, -1.5), + size=mp.Vector3(6, 0), + differentiable=list(names), + name="beam", + **params, + ) + sim = mp.Simulation( + cell_size=mp.Vector3(10, 8), + resolution=RES, + boundary_layers=[mp.PML(1.0)], + sources=[src], + force_complex_fields=True, + ) + dr = _design_region(sim) + mon = mpa.FourierFields( + sim, mp.Volume(center=mp.Vector3(0.8, 1.5), size=mp.Vector3(0, 0)), mp.Ez + ) + return mpa.OptimizationProblem( + simulation=sim, + objective_functions=[lambda f: npa.sum(npa.abs(f) ** 2)], + objective_arguments=[mon], + design_regions=[dr], + frequencies=[FCEN], + minimum_run_time=RUN, + maximum_run_time=RUN, + ) + + def _value(self, **overrides): + opt = self._problem(["beam_w0"], **overrides) + return float(np.asarray(opt([RHO], need_gradient=False)[0]).item()) + + def test_polarization_amplitude_is_exact(self): + # Every field the beam places is linear in beam_E0, so the objective is + # exactly quadratic in it and dJ/d(E0.z) is exactly 2J. No finite + # difference is involved, which makes this the sharpest check available. + opt = self._problem(["beam_E0"]) + value, grad = opt([RHO]) + objective = float(np.asarray(value).item()) + adjoint = float(np.real(grad["beam"]["beam_E0"][2])) + self.assertLess(abs(adjoint - 2 * objective) / (2 * objective), 1e-5) + + def test_waist_and_focus(self): + for name, index, base in ( + ("beam_w0", None, self.BEAM["beam_w0"]), + ("beam_x0", 0, self.BEAM["beam_x0"]), + ("beam_x0", 1, self.BEAM["beam_x0"]), + ): + with self.subTest(parameter=name, component=index): + opt = self._problem([name]) + _, grad = opt([RHO]) + entry = grad["beam"][name] + adjoint = float(np.real(entry if index is None else entry[index])) + + if index is None: + step = {name: base + FD_STEP}, {name: base - FD_STEP} + else: + delta = mp.Vector3(**{"xyz"[index]: FD_STEP}) + step = {name: base + delta}, {name: base - delta} + reference = (self._value(**step[0]) - self._value(**step[1])) / ( + 2 * FD_STEP ) + self.assertLess( + abs(adjoint - reference) / max(abs(reference), 1e-12), + 1e-3, + f"{name}[{index}]: adjoint {adjoint} vs {reference}", + ) + + def test_direction_is_projected_onto_its_tangent_space(self): + # beam_kdir's length is ignored, so only the direction is meaningful and + # the component along it is not a derivative of anything. + opt = self._problem(["beam_kdir"]) + _, grad = opt([RHO]) + adjoint = np.real(grad["beam"]["beam_kdir"]) + + base = self.BEAM["beam_kdir"] + axis = np.array([base.x, base.y, base.z], float) + axis /= np.linalg.norm(axis) + self.assertLess( + abs(float(np.dot(axis, adjoint))), 1e-6 * np.linalg.norm(adjoint) + 1e-9 + ) + + tangent = np.array([1.0, 0.0, 0.0]) + tangent = tangent - axis * np.dot(axis, tangent) + tangent /= np.linalg.norm(tangent) + step = mp.Vector3(*(tangent * FD_STEP)) + reference = ( + self._value(beam_kdir=base + step) - self._value(beam_kdir=base - step) + ) / (2 * FD_STEP) + self.assertLess( + abs(float(np.dot(adjoint, tangent)) - reference) / abs(reference), 1e-3 + ) try: @@ -402,7 +479,7 @@ def perturb(delta, j=j): @unittest.skipUnless(_HAVE_JAX, "JAX is an optional dependency") class TestJaxRoundTrip(unittest.TestCase): - """`jax.grad` through an `ArraySource` must match a finite difference. + """`jax.grad` through a source's `amp_data` must match a finite difference. This is what pins the cotangent convention. Meep returns `dJ/d(Re a) - i dJ/d(Im a)`, which is what JAX and autograd both produce @@ -411,17 +488,16 @@ class TestJaxRoundTrip(unittest.TestCase): """ def test_gradient_chains_through_jax(self): - size = mp.Vector3(0, 0.4) - npts = _num_source_points(SRC_C, size) + size = mp.Vector3(0, 1.5) + npts = 5 def wrapper(): - src = mp.ArraySource( + src = mp.Source( mp.GaussianSource(FCEN, fwidth=0.2), - mp.Ez, - amplitudes=np.ones(npts, dtype=complex), - frequency=FCEN, + component=mp.Ez, center=SRC_C, size=size, + amp_data=np.ones((1, npts, 1), dtype=np.complex128), name="sheet", ) sim = mp.Simulation(