[3/13][Adjoint Module] Fix adjoint source weighting on a symmetry plane

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_symmetry.py       \
     $(TEST_DIR)/test_adjoint_jax.py
 
 TESTS =                                        \

python/tests/test_adjoint_symmetry.py

@@ -0,0 +1,206 @@
+# Adjoint gradients when the simulation declares a mirror symmetry.
+#
+# Meep had no adjoint test that declared a symmetry at all, and the behaviour is
+# easy to mistake for a bug: with `Mirror(Y)` the design gradient comes back with
+# essentially all of its magnitude in y >= 0, which looks like a routine that
+# forgot to mirror. It is not. Meep builds the structure on the reduced grid
+# volume, so weights below the plane are never read, and a weight above it drives
+# both halves at once. A finite difference confirms both halves of that: the
+# gradient really is zero down there, and really is doubled up here.
+#
+# What was a bug is the scaling when the *objective* sits on or spans the
+# symmetry plane. `fourier_sourcedata` divided each quarter of an interpolated
+# adjoint source by the multiplicity of the voxel *centre* rather than of the
+# Yee site it was written to. Those differ exactly on a symmetry plane, where a
+# point is its own image and so has multiplicity 1 while the centre half a pixel
+# away has 2, leaving the source short there. The last two tests cover it.
+
+import unittest
+
+import numpy as np
+from autograd import numpy as npa
+
+import meep as mp
+import meep.adjoint as mpa
+
+RESOLUTION = 20
+FCEN = 1.0
+N = 16
+DESIGN = 1.6
+CELL = 6.0
+DPML = 1.0
+OXIDE = mp.Medium(index=1.44)
+SILICON = mp.Medium(index=3.48)
+FD_STEP = 1e-4
+
+
+def symmetric_weights(seed=0):
+    """A design that is its own mirror image, so both runs describe one structure."""
+    rng = np.random.default_rng(seed)
+    half = rng.uniform(0.3, 0.7, size=(N, N // 2))
+    return np.concatenate([half[:, ::-1], half], axis=1)
+
+
+def problem(weights, use_symmetry, monitor_center, monitor_size):
+    grid = mp.MaterialGrid(
+        mp.Vector3(N, N),
+        OXIDE,
+        SILICON,
+        weights=weights,
+        do_averaging=False,
+        grid_type="U_MEAN",
+    )
+    region = mpa.DesignRegion(
+        grid, volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(DESIGN, DESIGN))
+    )
+    sim = mp.Simulation(
+        cell_size=mp.Vector3(CELL, CELL),
+        resolution=RESOLUTION,
+        boundary_layers=[mp.PML(DPML)],
+        default_material=OXIDE,
+        geometry=[mp.Block(center=region.center, size=region.size, material=grid)],
+        # Ez is even about y = 0, so this is the symmetry the structure has.
+        symmetries=[mp.Mirror(mp.Y, phase=+1)] if use_symmetry else [],
+        sources=[
+            mp.Source(
+                mp.GaussianSource(FCEN, fwidth=0.2),
+                mp.Ez,
+                center=mp.Vector3(-2.0, 0),
+                size=mp.Vector3(0, 3.0),
+            )
+        ],
+        force_complex_fields=True,
+    )
+    monitor = mpa.FourierFields(
+        sim, mp.Volume(center=monitor_center, size=monitor_size), 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],
+    )
+
+
+def objective(weights, use_symmetry, center, size):
+    result = problem(weights, use_symmetry, center, size)(
+        [weights.ravel()], need_gradient=False
+    )
+    if isinstance(result, tuple):
+        result = result[0]
+    return float(np.real(np.squeeze(np.asarray(result, dtype=complex))))
+
+
+def gradient(weights, use_symmetry, center, size):
+    value, grad = problem(weights, use_symmetry, center, size)([weights.ravel()])
+    return (
+        float(np.real(np.squeeze(np.asarray(value, dtype=complex)))),
+        np.asarray(grad).reshape(N, N),
+    )
+
+
+def finite_difference(weights, use_symmetry, center, size, ix, iy):
+    plus, minus = weights.copy(), weights.copy()
+    plus[ix, iy] += FD_STEP
+    minus[ix, iy] -= FD_STEP
+    return (
+        objective(plus, use_symmetry, center, size)
+        - objective(minus, use_symmetry, center, size)
+    ) / (2 * FD_STEP)
+
+
+# A monitor well clear of the mirror plane, so nothing here depends on how the
+# plane itself is handled.
+OFF_PLANE = (mp.Vector3(2.0, 0.5), mp.Vector3(0, 0))
+# A monitor sitting exactly on it.
+ON_PLANE = (mp.Vector3(2.0, 0.0), mp.Vector3(0, 0))
+# The usual case: an extended monitor that straddles the plane.
+CROSSING = (mp.Vector3(2.0, 0.0), mp.Vector3(0, 3.0))
+
+
+class TestAdjointMirrorSymmetry(unittest.TestCase):
+    """`Mirror(Y)` declared, against a finite difference of the same setup."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.weights = symmetric_weights()
+
+    def test_forward_objective_unaffected_by_symmetry(self):
+        """Declaring the symmetry must not change the value being differentiated."""
+        w = self.weights
+        self.assertAlmostEqual(
+            objective(w, True, *CROSSING) / objective(w, False, *CROSSING),
+            1.0,
+            places=6,
+        )
+
+    def test_gradient_matches_finite_difference_off_plane(self):
+        """With the objective clear of the plane the gradient is exact."""
+        w = self.weights
+        _, g = gradient(w, True, *OFF_PLANE)
+        for ix, iy in ((8, 13), (8, 15)):
+            fd = finite_difference(w, True, *OFF_PLANE, ix, iy)
+            self.assertAlmostEqual(
+                g[ix, iy] / fd, 1.0, places=2, msg=f"node ({ix},{iy})"
+            )
+
+    def test_mirrored_half_is_inert(self):
+        """Weights below the plane are never read, so their gradient is zero.
+
+        This is the part that looks like a missing mirror operation and is not:
+        the finite difference is zero there too. Meep evaluates epsilon only on
+        the reduced grid volume, so those variables have no effect on the field.
+        """
+        w = self.weights
+        _, g = gradient(w, True, *CROSSING)
+        # Away from the plane -- close to it, a weight still reaches y >= 0
+        # through the design grid's own interpolation stencil.
+        for ix, iy in ((8, 2), (8, 3), (4, 1)):
+            fd = finite_difference(w, True, *CROSSING, ix, iy)
+            self.assertAlmostEqual(fd, 0.0, places=6, msg=f"FD at ({ix},{iy})")
+            self.assertAlmostEqual(g[ix, iy], 0.0, places=6, msg=f"grad at ({ix},{iy})")
+
+    def test_symmetric_variable_drives_both_halves(self):
+        """One weight above the plane moves the structure above and below it.
+
+        So its derivative is the sum of the two independent derivatives the same
+        structure has without the symmetry, which for a symmetric design is twice
+        either one.
+        """
+        w = self.weights
+        ix, iy = 8, 15
+        with_sym = finite_difference(w, True, *CROSSING, ix, iy)
+        without = finite_difference(w, False, *CROSSING, ix, iy)
+        self.assertAlmostEqual(with_sym / without, 2.0, places=3)
+
+    def test_gradient_matches_finite_difference_on_plane(self):
+        """An objective sitting on the symmetry plane.
+
+        This gave exactly 3/4 of the correct gradient while the adjoint source
+        was divided by the voxel centre's multiplicity: of the source's three Yee
+        rows, the one on the plane carries half the weight and was the one being
+        halved, so 1/4 of the total went missing regardless of resolution.
+        """
+        w = self.weights
+        ix, iy = 8, 15
+        _, g = gradient(w, True, *ON_PLANE)
+        fd = finite_difference(w, True, *ON_PLANE, ix, iy)
+        self.assertAlmostEqual(g[ix, iy] / fd, 1.0, places=2)
+
+    def test_gradient_matches_finite_difference_crossing_plane(self):
+        """An extended monitor spanning the plane -- the usual arrangement.
+
+        Same cause as the on-plane case, diluted over the rows the monitor
+        covers, so it used to show up as a first-order error: 4.0% at resolution
+        20, 2.9% at 30, 1.8% at 40.
+        """
+        w = self.weights
+        ix, iy = 8, 15
+        _, g = gradient(w, True, *CROSSING)
+        fd = finite_difference(w, True, *CROSSING, ix, iy)
+        self.assertAlmostEqual(g[ix, iy] / fd, 1.0, places=2)
+
+
+if __name__ == "__main__":
+    unittest.main()

src/dft.cpp

@@ -1544,12 +1544,25 @@ std::vector<struct sourcedata> dft_fields::fourier_sourcedata(const volume &wher
     direction cd = component_direction(c);
     sourcedata temp_struct = {c, idx_arr, f->fc->chunk_idx, amp_arr};
 
+    /* The directions `yee2cent_offsets` averages over, in the order it picks
+       them, so that each of the four sites below can be turned back into a grid
+       location.  `ix0` is the voxel centre and the sites are its corners, half a
+       grid step away, i.e. +/-1 in ivec units.  The entries are only read when
+       the corresponding `avg1`/`avg2` is nonzero, which is exactly when the
+       loop above filled them. */
+    direction avg_d[2] = {X, X};
+    int n_avg = 0;
+    LOOP_OVER_DIRECTIONS(f->fc->gv.dim, d) {
+      if (!f->fc->gv.iyee_shift(c).in_direction(d) && n_avg < 2) avg_d[n_avg++] = d;
+    }
+
     int position_array[3] = {0, 0, 0}; // array indicating the position of a point relative to the
                                        // minimum corner of the monitor
 
     LOOP_OVER_IVECS(f->fc->gv, f->is, f->ie, idx) {
       IVEC_LOOP_LOC(f->fc->gv, x0);
       IVEC_LOOP_ILOC(f->fc->gv, ix0);
+      const ivec ix0_local(ix0);
       x0 = f->S.transform(x0, f->sn) + rshift;
       ix0 = f->S.transform(ix0, f->sn) + f->shift;
 
@@ -1598,7 +1611,18 @@ std::vector<struct sourcedata> dft_fields::fourier_sourcedata(const volume &wher
             if (is_B(c) && f->fc->s->chi1inv[c - Bx + Hx][cd])
               EH0 /= f->fc->s->chi1inv[c - Bx + Hx][cd][idx];
 
-            EH0 /= f->S.multiplicity(ix0);
+            /* Take the multiplicity of the site this quarter is actually
+               written to, not of the voxel centre.  They differ exactly when a
+               site lands on a symmetry plane: such a point is its own image and
+               so has multiplicity 1, while the centre half a pixel away has 2.
+               Dividing it by 2 anyway leaves the adjoint source short on the
+               plane -- 3/4 of the correct gradient for an objective sitting on
+               it, and an O(dx) error for one merely spanning it. */
+            ivec isite(ix0_local);
+            if (f->avg1) isite += unit_ivec(f->fc->gv.dim, avg_d[0]) * ((j & 1) ? 1 : -1);
+            if (f->avg2) isite += unit_ivec(f->fc->gv.dim, avg_d[1]) * ((j & 2) ? 1 : -1);
+            ivec isiteS(f->S.transform(isite, f->sn) + f->shift);
+            EH0 /= f->S.multiplicity(isiteS);
             temp_struct.amp_arr.push_back(EH0);
           }
         }