[3/n][Adjoint Module] Adjoint gradients with respect to sources

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

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
 ------------------------------------------------------------
 

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
 

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.

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.