[5/13][Adjoint Module] Differentiable angular-spectrum propagation through stratified media

NEWS.md

@@ -2,6 +2,14 @@
 
 ## Meep 1.35.0 (in progress)
 
+* 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
+  medium, this handles radiation crossing a material interface, so the layers
+  above the monitor can leave the FDTD cell entirely and become differentiable
+  parameters. Supports 2D and 3D Cartesian simulations, and is usable with or
+  without the adjoint solver.
+
 * Adjoint solver: objective functions are now differentiated with a single
   vector-Jacobian product instead of a full frequency Jacobian per objective
   argument. For an objective of $M$ arguments at $F$ frequencies this reduces

doc/docs/Python_Tutorials/Adjoint_Solver.md

@@ -273,6 +273,95 @@ JAX is an optional dependency throughout. If it is not installed, JAX objective
 functions are simply not recognized and `mpa.MeepJaxWrapper` and
 `mpa.value_and_jacobian` are absent; everything else works unchanged.
 
+Propagating Through Stratified Media
+------------------------------------
+
+Meep's [near-to-far transformation](Near_to_Far_Field_Spectra.md) requires its
+surface to sit in a homogeneous medium — `dft_near2far` aborts otherwise — so a
+structure whose radiation crosses a material interface has to keep that interface
+inside the FDTD cell. A grating coupler radiating up through a cladding, across
+the chip surface, and hundreds of microns to a fiber cannot afford to.
+
+`meep.adjoint.AngularSpectrum` propagates the tangential DFT fields on a planar
+monitor through an arbitrary layer stack analytically instead. The layers above
+the monitor leave the simulation, and become differentiable parameters:
+
+```py
+stack = mpa.Stack([mpa.Layer(index=1.444, thickness=300.0),   # silica superstrate
+                   mpa.Layer(index=1.0)])                      # air, semi-infinite
+
+monitor = sim.add_dft_fields([mp.Ez, mp.Hx], frequencies, where=plane)
+sim.run(...)
+
+propagator = mpa.AngularSpectrum.from_monitor(sim, monitor, stack, plane)
+fiber = mpa.gaussian_mode(waist=5.2, tilt_deg=8.0)
+efficiency = propagator.overlap_monitor(sim, monitor, fiber, distance=305.0)
+```
+
+That path is plain NumPy in and out, so it can post-process an ordinary forward
+run without any optimization. Inside an objective function, use
+`objective_arguments()` and `take()` instead, which give the `FourierFields` the
+adjoint solver needs.
+
+Both tangential fields are required. Together they determine the up- and
+down-going plane-wave amplitudes separately, which an open near-to-far surface
+cannot do — it has no way to reject radiation heading the wrong way.
+
+### Reading the diagnostics before believing a number
+
+The transverse transform is periodic, so a field that has not decayed by the ends
+of the monitor wraps around, and it does so quietly: the result looks plausible
+and does not improve with resolution. `report()` measures the three ways a
+monitor is usually placed wrongly:
+
+```py
+propagator.report_monitor(sim, monitor)
+# {'downgoing_fraction': ...,   is something above the monitor scattering?
+#  'evanescent_fraction': ...,  is the monitor inside the near field?
+#  'edge_amplitude': ...}       has the field decayed by the monitor's ends?
+```
+
+A few percent of `edge_amplitude` is common for a grating radiating into a
+cladding and is not the beam tail — it is near-grazing radiation, which travels
+sideways rather than decaying, so widening the cell barely helps. It bounds the
+accuracy at a similar level. Power at those angles was never going to reach the
+fiber.
+
+The other constraint is the padded window. A beam that spreads over hundreds of
+microns needs `pad_factor` large enough that the spread beam still fits; too small
+and the beam wraps onto itself, which is nonsense rather than merely inaccurate.
+
+### What is and is not supported
+
+Two- and three-dimensional Cartesian simulations, with a planar monitor normal to
+a coordinate axis lying in a homogeneous region — a line in 2D, a rectangle in
+3D. Cylindrical coordinates are not supported.
+
+In 3D the transverse wavevectors form a plane, and the tangential fields are
+resolved into the s and p directions of each one. A beam is therefore specified
+by its linear polarization rather than by a single scalar component:
+`gaussian_mode(waist, polarization=...)` builds a linearly polarized beam whose s
+and p content follows the azimuth. A spectrum that were purely s at every
+wavevector would instead be azimuthally polarized, carrying a vortex at normal
+incidence, which is rarely what is wanted.
+
+One consequence worth knowing when checking gradients in 3D: the *adjoint* DFT
+takes appreciably longer to converge than the forward one, and an
+under-converged adjoint produces a gradient that is wrong by a fixed factor
+which does not shrink with the finite-difference step — so it reads like a bug
+rather than like noise. Give the adjoint run enough time, or fix the run length
+so that both runs of a finite difference cover the same interval.
+
+A single plane is complete for the half-space above it, so unlike near2far there
+is no closed surface to build: the plane plus the hemisphere at infinity already
+is one, and the up/down split discards what is heading the wrong way. To account
+for both half-spaces — the substrate as well as the superstrate — use two
+parallel monitors, each with its own stack, and flip `sign` on the lower one.
+
+[examples/adjoint_optimization/grating_coupler_asm.py](https://github.com/NanoComp/meep/blob/master/python/examples/adjoint_optimization/grating_coupler_asm.py)
+is a worked two-etch grating coupler radiating through a thick superstrate into a
+fiber, with a forward-only mode and an optimization mode.
+
 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_symmetry.py       \
     $(TEST_DIR)/test_adjoint_protocol.py       \
+    $(TEST_DIR)/test_angular_spectrum.py       \
     $(TEST_DIR)/test_adjoint_jax.py
 
 TESTS =                                        \
@@ -245,6 +246,7 @@ pkgpython_PYTHON = __init__.py $(HL_IFACE)
 
 adjointdir = $(pkgpythondir)/adjoint
 adjoint_PYTHON  = $(srcdir)/adjoint/__init__.py              \
+                  $(srcdir)/adjoint/angular_spectrum.py      \
                   $(srcdir)/adjoint/basis.py                 \
                   $(srcdir)/adjoint/objective.py             \
                   $(srcdir)/adjoint/optimization_problem.py  \

python/adjoint/__init__.py

@@ -25,5 +25,13 @@
 # so objective functions written with `jax.numpy` need no special treatment.
 try:
     from .wrapper import MeepJaxWrapper, value_and_jacobian
+    from .angular_spectrum import (
+        AngularSpectrum,
+        Layer,
+        Mode,
+        Stack,
+        TangentialFields,
+        gaussian_mode,
+    )
 except ModuleNotFoundError as _:
     pass