[7/n][Adjoint Module] Evaluate the dispersive adjoint operator at the discrete lineshape

python/Makefile.am

@@ -34,6 +34,7 @@ ADJOINT_TESTS =                                \
     $(TEST_DIR)/test_adjoint_solver.py         \
     $(TEST_DIR)/test_adjoint_utils.py          \
     $(TEST_DIR)/test_adjoint_cyl.py            \
+    $(TEST_DIR)/test_adjoint_dispersion.py     \
     $(TEST_DIR)/test_adjoint_jax.py
 
 TESTS =                                        \

python/tests/test_adjoint_dispersion.py

@@ -0,0 +1,169 @@
+"""Adjoint gradients of dispersive structures against finite differences.
+
+`get_chi1_tensor_disp` used to build the dispersive permittivity from the
+*continuum* lineshape -- `lorentzian_susceptibility::chi1` and
+`1 + i*sigma/(2*pi*f)` -- while the FDTD timesteps a discrete recurrence. The
+adjoint gradient was then the exact derivative of an operator slightly
+different from the one being simulated, and disagreed with a finite difference
+of the discrete objective at `O((freq*dt)^2)`.
+
+Measured relative disagreement at `res` 20 / 40 / 80, continuum against
+discrete:
+
+    D-conductivity   0.455%  0.107%  0.027%   ->  0.001%  0.000%  0.001%
+    Drude pole       0.232%  0.058%  0.015%   ->  0.000%  0.000%  0.000%
+
+The `4x`-per-doubling of the left-hand columns is the second-order signature;
+the right-hand columns sit at the same roundoff floor as a non-dispersive
+structure. `TOL` below is set from those numbers: loose enough to be immune to
+roundoff, tight enough that reverting the fix fails at every resolution tested.
+
+Run lengths are pinned rather than left to `stop_when_dft_decayed`: an adaptive
+stop makes a perturbed run end at a different time, and the difference scales
+with the perturbation.
+"""
+
+import unittest
+
+import numpy as np
+from autograd import numpy as npa
+
+import meep as mp
+import meep.adjoint as mpa
+
+FCEN = 1.0
+RUN = 150.0
+CELL = mp.Vector3(6, 4)
+NG = 4
+IDX = 5
+U0 = 0.3
+FD_STEP = 1e-4
+
+# Comfortably below the 0.232% that the continuum lineshape produces at the
+# coarsest resolution tested here, and ~100x above the roundoff floor the fix
+# achieves. See the measurements in the module docstring.
+TOL = 5e-4
+
+DRUDE = mp.Medium(
+    epsilon=1.0,
+    E_susceptibilities=[mp.DrudeSusceptibility(frequency=1.0, gamma=0.1, sigma=0.3)],
+)
+LORENTZ = mp.Medium(
+    epsilon=1.0,
+    E_susceptibilities=[
+        mp.LorentzianSusceptibility(frequency=1.3, gamma=0.2, sigma=0.5)
+    ],
+)
+
+
+def _problem(medium2, res, damping=0.0):
+    grid = mp.MaterialGrid(
+        mp.Vector3(NG, NG), mp.air, medium2, do_averaging=False, beta=0, damping=damping
+    )
+    region = mpa.DesignRegion(
+        grid, volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(1.0, 1.0))
+    )
+    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=mp.Vector3(-1.8, 0),
+            )
+        ],
+        geometry=[mp.Block(center=region.center, size=region.size, material=grid)],
+        eps_averaging=False,
+        force_complex_fields=True,
+    )
+    monitor = mpa.FourierFields(
+        sim, mp.Volume(center=mp.Vector3(1.8, 0), size=mp.Vector3()), mp.Ez
+    )
+    return mpa.OptimizationProblem(
+        simulation=sim,
+        objective_functions=[lambda f: npa.sum(npa.abs(f) ** 2)],
+        objective_arguments=[monitor],
+        design_regions=[region],
+        frequencies=[FCEN],
+        minimum_run_time=RUN,
+        maximum_run_time=RUN,
+    )
+
+
+class TestDispersiveGradient(unittest.TestCase):
+    def assertGradientMatchesFD(self, medium2, res=20, damping=0.0):
+        rho = U0 * np.ones(NG * NG)
+        _, grad = _problem(medium2, res, damping)([rho])
+        adjoint = float(np.real(np.atleast_1d(np.squeeze(grad))[IDX]))
+
+        def value(r):
+            return float(
+                np.asarray(
+                    _problem(medium2, res, damping)([r], need_gradient=False)[0]
+                ).real.item()
+            )
+
+        hi, lo = rho.copy(), rho.copy()
+        hi[IDX] += FD_STEP
+        lo[IDX] -= FD_STEP
+        reference = (value(hi) - value(lo)) / (2 * FD_STEP)
+
+        # Guards against the comparison passing vacuously: a structure the
+        # objective does not respond to would make any adjoint value "agree".
+        self.assertGreater(abs(reference), 1e-6, "objective does not respond to rho")
+        self.assertLess(
+            abs(adjoint - reference) / abs(reference),
+            TOL,
+            f"adjoint {adjoint} vs finite difference {reference}",
+        )
+        return adjoint, reference
+
+    def test_drude_pole(self):
+        self.assertGradientMatchesFD(DRUDE)
+
+    def test_lorentzian_pole(self):
+        # The Drude branch of chi1_discrete() drops the omega_0^2 term, so a
+        # resonant pole exercises a path the Drude test does not reach.
+        self.assertGradientMatchesFD(LORENTZ)
+
+    def test_d_conductivity(self):
+        # `damping` adds sigma = u(1-u)*damping to D_conductivity, which reaches
+        # the operator through conductivity_factor() rather than through a pole.
+        self.assertGradientMatchesFD(mp.Medium(index=1.5), damping=20.0)
+
+    def test_coarse_resolution(self):
+        # The sharpest of these. The error being removed is O((freq*dt)^2), so
+        # it is largest where dt is largest; a formula that is merely closer
+        # than the continuum one would still converge and still pass at res 20.
+        self.assertGradientMatchesFD(DRUDE, res=10)
+        self.assertGradientMatchesFD(mp.Medium(index=1.5), res=10, damping=20.0)
+
+
+class TestReportedEpsilonIsContinuum(unittest.TestCase):
+    """`chi1` describes the material; only the adjoint wants the discretization.
+
+    Guards the split: the discrete lineshape lives in `chi1_discrete`, and
+    anything reporting material properties to the user must keep using the
+    continuum one, or `get_epsilon` would start returning resolution-dependent
+    numbers for a resolution-independent material.
+    """
+
+    def test_get_epsilon_matches_the_medium_at_two_resolutions(self):
+        reference = complex(np.squeeze(DRUDE.epsilon(FCEN))[0, 0])
+        for res in (10, 40):
+            sim = mp.Simulation(
+                cell_size=mp.Vector3(2, 2),
+                resolution=res,
+                default_material=DRUDE,
+                force_complex_fields=True,
+            )
+            sim.init_sim()
+            eps = np.mean(sim.get_epsilon(frequency=FCEN))
+            self.assertAlmostEqual(eps.real, reference.real, places=5)
+            self.assertAlmostEqual(eps.imag, reference.imag, places=5)
+
+
+if __name__ == "__main__":
+    unittest.main()

src/meep.hpp

@@ -110,6 +110,20 @@ class susceptibility {
   // Returns the 1st order (linear) susceptibility (generic)
   virtual std::complex<realnum> chi1(realnum freq, realnum sigma = 1);
 
+  // Frequency response of the recurrence that update_P() actually timesteps,
+  // rather than of the continuous-time model it approximates.  Reduces to
+  // chi1() as dt -> 0, and is called with dt = 0 to mean exactly that.
+  //
+  // The adjoint solver needs this: differentiating the continuum lineshape
+  // gives the exact derivative of an operator slightly different from the one
+  // being simulated, so the gradient disagrees with a finite difference of the
+  // discrete objective at O((freq*dt)^2).  Reporting material properties to
+  // the user is the opposite case and should keep using chi1().
+  virtual std::complex<realnum> chi1_discrete(realnum freq, realnum sigma, realnum dt) {
+    (void)dt;
+    return chi1(freq, sigma);
+  }
+
   // update all of the internal polarization state given the W field
   // at the current time step, possibly the previous field W_prev, etc.
   virtual void update_P(realnum *W[NUM_FIELD_COMPONENTS][2],
@@ -252,6 +266,7 @@ class lorentzian_susceptibility : public susceptibility {
 
   // Returns the 1st order nonlinear susceptibility
   virtual std::complex<realnum> chi1(realnum freq, realnum sigma = 1);
+  virtual std::complex<realnum> chi1_discrete(realnum freq, realnum sigma, realnum dt);
 
   virtual void update_P(realnum *W[NUM_FIELD_COMPONENTS][2],
                         realnum *W_prev[NUM_FIELD_COMPONENTS][2], realnum dt, const grid_volume &gv,

src/meepgeom.cpp

@@ -740,6 +740,14 @@ geom_epsilon::geom_epsilon(const geom_epsilon &geps1) {
   extra_materials = geps1.extra_materials;
   current_pol = NULL;
 
+  // Parameters stashed for the gradient calculation.  Silently reverting these
+  // to their defaults in a copy would not fail, it would just quietly compute
+  // against a different operator than the original was built with.
+  u_p = geps1.u_p;
+  tol = geps1.tol;
+  maxeval = geps1.maxeval;
+  dt = geps1.dt;
+
   FOR_DIRECTIONS(d) FOR_SIDES(b) { cond[d][b].prof = geps1.cond[d][b].prof; }
 }
 geom_epsilon::~geom_epsilon() {
@@ -2051,6 +2059,7 @@ void set_materials_from_geom_epsilon(meep::structure *s, geom_epsilon *geps,
                                      absorber_list alist) {
 
   // store for later use in gradient calculations
+  geps->dt = s->dt;
   geps->tol = tol;
   geps->maxeval = maxeval;
 
@@ -2577,6 +2586,35 @@ void invert_tensor(std::complex<double> t_inv[9], std::complex<double> t[9]) {
 #undef minv
 }
 
+/* The factor by which D-conductivity scales the permittivity at `freq`.
+
+   Continuum: Ampere's law reads D(-i w + sigma) = curl, so factoring out the
+   time derivative gives 1 + i sigma / w.
+
+   The D update actually timestepped is
+
+     D' = ((1 - sigma dt/2) D + dcurl) / (1 + sigma dt/2)
+
+   (step_db.cpp, with condinv from structure.cpp).  Substituting
+   D ~ exp(-i w t) and dividing by exp(-i w dt/2) turns the bracket into
+
+     -(2/dt) sin(w dt/2) i + sigma cos(w dt/2)
+
+   i.e. w -> (2/dt) sin(w dt/2) and sigma -> sigma cos(w dt/2), so the factor
+   becomes 1 + i sigma (dt/2) cot(pi freq dt).  That differs from the continuum
+   form at O((freq*dt)^2); `dt = 0` selects the continuum form. */
+static std::complex<double> conductivity_factor(double sigma, double freq, double dt) {
+  if (sigma == 0) return std::complex<double>(1.0, 0.0);
+
+  double inv_omega;
+  if (dt > 0) {
+    double half_phase = meep::pi * freq * dt;
+    inv_omega = 0.5 * dt / std::tan(half_phase);
+  }
+  else { inv_omega = 1.0 / (2 * meep::pi * freq); }
+  return std::complex<double>(1.0, sigma * inv_omega);
+}
+
 void get_chi1_tensor_disp(std::complex<double> tensor[9], const meep::vec &r, double freq,
                           geom_epsilon *geps) {
   // locate the proper material
@@ -2591,15 +2629,15 @@ void get_chi1_tensor_disp(std::complex<double> tensor[9], const meep::vec &r, do
     vector3 dummy;
     dummy.x = dummy.y = dummy.z = 0.0;
     double conductivityCur = vec_to_value(mm->D_conductivity_diag, dummy, i);
-    a = std::complex<double>(1.0, conductivityCur / (2 * meep::pi * freq));
+    a = conductivity_factor(conductivityCur, freq, geps->dt);
 
     // compute lorentzian component including the instantaneous ε
     b = cvec_to_value(mm->epsilon_diag, mm->epsilon_offdiag, i);
     for (const auto &mm_susc : mm->E_susceptibilities) {
       meep::lorentzian_susceptibility sus =
           meep::lorentzian_susceptibility(mm_susc.frequency, mm_susc.gamma, mm_susc.drude);
       double sigma = vec_to_value(mm_susc.sigma_diag, mm_susc.sigma_offdiag, i);
-      b += sus.chi1(freq, sigma);
+      b += sus.chi1_discrete(freq, sigma, geps->dt);
     }
 
     // elementwise multiply
@@ -2647,13 +2685,10 @@ std::complex<double> cond_cmp(meep::component c, const meep::vec &r, double freq
   // get the row we care about
   switch (component_direction(c)) {
     case meep::X:
-    case meep::R:
-      return std::complex<double>(1.0, mm->D_conductivity_diag.x / (2 * meep::pi * freq));
+    case meep::R: return conductivity_factor(mm->D_conductivity_diag.x, freq, geps->dt);
     case meep::Y:
-    case meep::P:
-      return std::complex<double>(1.0, mm->D_conductivity_diag.y / (2 * meep::pi * freq));
-    case meep::Z:
-      return std::complex<double>(1.0, mm->D_conductivity_diag.z / (2 * meep::pi * freq));
+    case meep::P: return conductivity_factor(mm->D_conductivity_diag.y, freq, geps->dt);
+    case meep::Z: return conductivity_factor(mm->D_conductivity_diag.z, freq, geps->dt);
     case meep::NO_DIRECTION: meep::abort("Invalid adjoint field component");
   }
 }

src/meepgeom.hpp

@@ -179,6 +179,11 @@ class geom_epsilon : public meep::material_function {
   cond_profile cond[5][2]; // [direction][side]
   double tol = DEFAULT_SUBPIXEL_TOL;
   int maxeval = DEFAULT_SUBPIXEL_MAXEVAL;
+  /* Timestep of the structure these materials were installed in, so that the
+     adjoint solver can evaluate the dispersive response of the recurrence
+     actually being timestepped rather than of the continuum model.  Zero means
+     "not known", which selects the continuum form. */
+  double dt = 0;
 
   geom_epsilon(geometric_object_list g, material_type_list mlist, const meep::volume &v);
   geom_epsilon(const geom_epsilon &geps1); // copy constructor

src/susceptibility.cpp

@@ -306,6 +306,33 @@ std::complex<realnum> lorentzian_susceptibility::chi1(realnum freq, realnum sigm
   }
 }
 
+/* Frequency response of the recurrence in update_P(),
+
+     P'(1 + g dt/2) = P (2 - w0^2 dt^2) - P_prev (1 - g dt/2) + w0^2 dt^2 sigma E
+
+   with g = 2 pi gamma and w0 = 2 pi omega_0.  Substituting P ~ exp(-i w n dt)
+   and dividing through by P^n gives a denominator
+
+     w0^2 dt^2 - 4 sin^2(w dt/2) - i g dt sin(w dt)
+
+   so relative to the continuum lineshape the discretization replaces
+
+     freq^2       ->  [sin(pi freq dt) / (pi dt)]^2
+     gamma*freq   ->  gamma sin(2 pi freq dt) / (2 pi dt)
+
+   and leaves the numerator alone.  Both corrections are O((freq*dt)^2). */
+std::complex<realnum> lorentzian_susceptibility::chi1_discrete(realnum freq, realnum sigma,
+                                                               realnum dt) {
+  if (dt <= 0) return chi1(freq, sigma);
+
+  realnum freq_eff = std::sin(pi * freq * dt) / (pi * dt);
+  realnum damping = gamma * std::sin(2 * pi * freq * dt) / (2 * pi * dt);
+  realnum resonance = no_omega_0_denominator ? 0 : omega_0 * omega_0;
+
+  return sigma * omega_0 * omega_0 /
+         std::complex<realnum>(resonance - freq_eff * freq_eff, -damping);
+}
+
 void lorentzian_susceptibility::dump_params(h5file *h5f, size_t *start) {
   size_t num_params = 5;
   size_t params_dims[1] = {num_params};