diff --git a/NEWS.md b/NEWS.md index a8d273ad5..0bcde63bd 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,6 +2,21 @@ ## Meep 1.35.0 (in progress) +* Adjoint solver: sources can now be differentiated alongside the design + 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 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..b5245afe1 100644 --- a/doc/docs/Python_Tutorials/Adjoint_Solver.md +++ b/doc/docs/Python_Tutorials/Adjoint_Solver.md @@ -362,6 +362,136 @@ 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'`, `'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 amplitudes yourself + +`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.Source( + mp.GaussianSource(fcen, fwidth=df), + component=mp.Ez, + center=mp.Vector3(-1, 0), + size=mp.Vector3(0, 4), + amp_data=profile, # trilinearly interpolated onto the grid + differentiable=["amp_data"], + name="sheet", +) +``` + +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 + +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 +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 e4f205fa9..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 \ @@ -250,6 +252,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..dc6a38097 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,97 @@ 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): + 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 + scale = source_gradient.source_grad_scale( + self.sim, + src, + self.frequencies, + self.objective_arguments[0]._adj_src_phase(), + ) + grads = self._source_gradient_for(src, monitors, scale) + 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 _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 = {} + 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 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, + # 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() + # 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..98ae3e0e3 --- /dev/null +++ b/python/adjoint/source_gradient.py @@ -0,0 +1,603 @@ +"""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", + "amp_data", + "beam_x0", + "beam_kdir", + "beam_w0", + "beam_E0", +) + + +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, + component: Optional[int] = None, + ): + self.source = source + # 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( + 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: 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( + 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 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 + ) + 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. + + `scale` is the per-frequency factor relating a current amplitude to the + adjoint field, broadcast over the spatial axes. + """ + num_points = self.num_points(sim) + grad = np.zeros(num_points * len(self._frequencies), dtype=np.complex128) + + self._monitor.swigobj.volume_source_gradient( + self.volume.swigobj, self.component, sim.fields, grad + ) + + grad = grad.reshape(len(self._frequencies), num_points) + # 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, + sources: List, + frequencies: np.ndarray, + decimation_factor: Optional[int] = 0, +) -> List[SourceGradientMonitor]: + """Install a DFT monitor over each differentiable source's support.""" + 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: + """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) + + # 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) + 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) + + # 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 + scale *= 2 + 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 _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. + + 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..88e4d3d6a 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 @@ -212,14 +213,33 @@ def __init__( self.until_after_sources = until_after_sources self.finite_difference_step = finite_difference_step + # 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 getattr(s, "amp_data", None) is not None + ] + for s in self.differentiable_sources: + if "amp_data" not in getattr(s, "differentiable", ()): + s.differentiable = tuple(getattr(s, "differentiable", ())) + ( + "amp_data", + ) + 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: 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 @@ -239,14 +259,35 @@ 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." + ) + # 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.amp_data = onp.asarray(amps, dtype=onp.complex128).reshape( + onp.shape(src.amp_data) + ) 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) @@ -293,6 +334,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" @@ -305,6 +351,43 @@ 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, 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 + ) + 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(rows, axis=0).reshape(shape)) + else: + out.append(rows.reshape((rows.shape[0], *shape))) + return out + def _calculate_vjps( self, fwd_fields: List[List[mp.DftFields]], @@ -328,14 +411,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) @@ -351,7 +436,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) @@ -376,10 +465,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 @@ -394,10 +484,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 @@ -543,13 +634,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/meep.i b/python/meep.i index 9c917969d..d968105f8 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 }; @@ -1084,7 +1085,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, std::complex* EH }; // typemaps for diffractedplanewave diff --git a/python/source.py b/python/source.py index b88440c2a..570c13199 100644 --- a/python/source.py +++ b/python/source.py @@ -15,6 +15,67 @@ 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", "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 +# (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", ()) + ) + 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: + 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))}." + ) + 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 +110,8 @@ def __init__( amp_func=None, amp_func_file="", amp_data=None, + differentiable=None, + name=None, ): """ Construct a `Source`. @@ -98,6 +161,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 +210,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 +783,12 @@ 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`'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, src, @@ -1089,10 +1176,20 @@ 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/python/tests/test_source_gradient.py b/python/tests/test_source_gradient.py new file mode 100644 index 000000000..f65427920 --- /dev/null +++ b/python/tests/test_source_gradient.py @@ -0,0 +1,547 @@ +"""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_class_specific_parameters_are_scoped_to_their_class(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), + ) + 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"]) + + +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 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_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), + component=mp.Ez, + center=SRC_C, + size=mp.Vector3(0, 1.0), + amp_func=lambda p: 1.0, + differentiable=["currents"], + name="drive", + ) + _, opt = _problem(src) + with self.assertRaises(NotImplementedError): + opt([RHO]) + + +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_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(values): + return mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), + component=mp.Ez, + center=SRC_C, + size=size, + amp_data=np.asarray(values, dtype=np.complex128).reshape(1, n, 1), + differentiable=["amp_data"], + name="sheet", + ) + + _, opt = _problem(make(data)) + _, grad = opt([RHO]) + adjoint = np.ravel(grad["sheet"]["amp_data"]) + self.assertEqual(adjoint.size, n) + + def evaluate(values): + _, o = _problem(make(values)) + return np.asarray(o([RHO], need_gradient=False)[0]).item() + + for j in (0, n // 2, n - 1): + with self.subTest(entry=j): + + def perturb(delta, j=j): + out = data.copy() + out[j] += delta + return out + + reference = _complex_fd(evaluate, perturb) + self.assertLess( + abs(adjoint[j] - reference) / abs(reference), + 1e-5, + 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: + 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 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 + 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, 1.5) + npts = 5 + + def wrapper(): + src = mp.Source( + mp.GaussianSource(FCEN, fwidth=0.2), + component=mp.Ez, + center=SRC_C, + size=size, + amp_data=np.ones((1, npts, 1), dtype=np.complex128), + 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() diff --git a/src/dft.cpp b/src/dft.cpp index 0257586e6..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) { @@ -1585,4 +1612,252 @@ 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)); +} + +/* 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 911d37a4e..eec4baad7 100644 --- a/src/meep.hpp +++ b/src/meep.hpp @@ -1425,6 +1425,14 @@ 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); + // 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(); @@ -2185,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,