[1/n][Adjoint Module] Fix MaterialGrid design-variable attribution in the adjoint gradient

python/Makefile.am

@@ -39,7 +39,9 @@ ADJOINT_TESTS =                                \
 TESTS =                                        \
     $(TEST_DIR)/test_3rd_harm_1d.py            \
     $(TEST_DIR)/test_absorber_1d.py            \
+    $(TEST_DIR)/test_adjoint_adjacent_grids.py \
     $(TEST_DIR)/test_adjoint_chunks.py         \
+    $(TEST_DIR)/test_adjoint_symmetric_grids.py \
     $(TEST_DIR)/test_antenna_radiation.py      \
     $(TEST_DIR)/test_array_metadata.py         \
     $(TEST_DIR)/test_bend_flux.py              \

python/adjoint/utils.py

@@ -73,6 +73,7 @@ def get_gradient(
             onp.array(frequencies),
             sim.geps,
             finite_difference_step,
+            getattr(self.design_parameters, "grid_id", -1),
         )
         return onp.squeeze(grad).T
 

python/geom.py

@@ -7,6 +7,7 @@
 from collections import namedtuple
 from copy import deepcopy
 import functools
+import itertools
 import math
 from numbers import Number
 import operator
@@ -585,6 +586,9 @@ def _get_epsmu(
         return np.squeeze(epsmu)
 
 
+_material_grid_counter = itertools.count()
+
+
 class MaterialGrid:
     """
     This class is used to specify materials on a rectilinear grid. A class object is passed
@@ -667,6 +671,12 @@ def __init__(
         allow you to combine any material grids that overlap in space with no intervening objects.
         """
         self.grid_size = mp.Vector3(*grid_size)
+        # A stable identity for this grid, carried into the C++ material_data so
+        # that the adjoint gradient of a design region can tell its own grid
+        # apart from a neighbouring one. Several material_data objects may be
+        # created from this single Python object (e.g. one per symmetric copy),
+        # and they all share this id.
+        self.grid_id = next(_material_grid_counter)
         self.medium1 = medium1
         self.medium2 = medium2
 

python/meep.i

@@ -871,7 +871,7 @@ void _get_gradient(PyObject *grad, double scalegrad,
                    meep::dft_fields *fields_a_0, meep::dft_fields *fields_a_1, meep::dft_fields *fields_a_2,
                    meep::dft_fields *fields_f_0, meep::dft_fields *fields_f_1, meep::dft_fields *fields_f_2,
                    meep::grid_volume *grid_volume, PyObject *frequencies,
-                   meep_geom::geom_epsilon *geps, double fd_step) {
+                   meep_geom::geom_epsilon *geps, double fd_step, long owner_grid_id) {
 
     // clean the gradient array
     PyArrayObject *pao_grad = (PyArrayObject *)grad;
@@ -896,7 +896,7 @@ void _get_gradient(PyObject *grad, double scalegrad,
     if (PyArray_DIMS(pao_grad)[0] != nf) meep::abort("Numpy grad array is allocated for %td frequencies; it should be allocated for %td.",PyArray_DIMS(pao_grad)[0],nf);
 
     // calculate the gradient
-    meep_geom::material_grids_addgradient(grad_c,ng,nf,adjoint_fields,forward_fields,frequencies_c,scalegrad,*grid_volume,geps,fd_step);
+    meep_geom::material_grids_addgradient(grad_c,ng,nf,adjoint_fields,forward_fields,frequencies_c,scalegrad,*grid_volume,geps,owner_grid_id,fd_step);
 
 }
 %}

python/tests/test_adjoint_adjacent_grids.py

@@ -0,0 +1,182 @@
+"""Each design region's gradient must come only from its own MaterialGrid.
+
+Design regions may be adjacent, separated, or overlapping; in every arrangement
+the adjoint gradient of one region must agree with a finite difference of the
+objective with respect to that region's own design variables.
+"""
+
+import unittest
+
+import numpy as np
+from autograd import numpy as npa
+
+import meep as mp
+import meep.adjoint as mpa
+
+RESOLUTION = 20
+DPML = 1.0
+CELL = 6.0
+LX, LY = 1.0, 1.0
+FCEN = 1 / 1.55
+SILICON = mp.Medium(index=3.4)
+OXIDE = mp.Medium(index=1.44)
+
+
+def _problem(grid_shapes, gap=0.0, overlap=False, grid_type="U_DEFAULT"):
+    """Design region(s) stacked in y, optionally separated by `gap`.
+
+    With `overlap`, the grids share one volume and only the first is a design
+    region.
+    """
+    if overlap:
+        boxes = [(-LY / 2, LY / 2)] * len(grid_shapes)
+    else:
+        if len(grid_shapes) == 1:
+            boxes = [(-LY / 2, LY / 2)]
+        else:
+            boxes = [(-LY / 2, -gap / 2), (gap / 2, LY / 2)]
+
+    grids = [
+        mp.MaterialGrid(
+            mp.Vector3(*s),
+            OXIDE,
+            SILICON,
+            # Non-design grids are filled at mid density: an almost-empty cell
+            # scatters too weakly for the gradient to converge at this decay_by.
+            weights=np.full(s, 0.5) if overlap else np.zeros(s),
+            do_averaging=False,
+            beta=0,
+            eta=0.5,
+            grid_type=grid_type,
+        )
+        for s in grid_shapes
+    ]
+    vols = [
+        mp.Volume(center=mp.Vector3(0, 0.5 * (a + b)), size=mp.Vector3(LX, b - a))
+        for a, b in boxes
+    ]
+    sim = mp.Simulation(
+        cell_size=mp.Vector3(CELL, CELL),
+        resolution=RESOLUTION,
+        boundary_layers=[mp.PML(DPML)],
+        geometry=[
+            mp.Block(center=v.center, size=v.size, material=g)
+            for v, g in zip(vols, grids)
+        ],
+        sources=[
+            mp.Source(
+                mp.GaussianSource(FCEN, fwidth=0.1),
+                component=mp.Ez,
+                center=mp.Vector3(-CELL / 2 + DPML + 0.3, 0),
+                size=mp.Vector3(0, CELL - 2 * DPML),
+            )
+        ],
+        default_material=OXIDE,
+        dimensions=2,
+    )
+    monitor = mp.Volume(
+        center=mp.Vector3(CELL / 2 - DPML - 0.3, 0), size=mp.Vector3(0, 2.0)
+    )
+    opt = mpa.OptimizationProblem(
+        simulation=sim,
+        objective_functions=[lambda ez: npa.sum(npa.abs(ez) ** 2, axis=1)],
+        objective_arguments=[mpa.FourierFields(sim, monitor, mp.Ez, yee_grid=False)],
+        design_regions=(
+            [mpa.DesignRegion(grids[0], volume=vols[0])]
+            if overlap
+            else [mpa.DesignRegion(g, volume=v) for g, v in zip(grids, vols)]
+        ),
+        frequencies=[FCEN],
+        decay_by=1e-9,
+    )
+    return opt
+
+
+def _directional_ratio(opt, grid_shapes, seed=0, step=3e-3, n_regions=None):
+    """adjoint / finite difference along one random direction.
+
+    Directional rather than single-pixel: perturbing one design variable moves
+    the objective by about as much as the DFT convergence floor.
+    """
+    rng = np.random.default_rng(seed)
+    if n_regions is None:
+        n_regions = len(grid_shapes)
+    weights = [rng.uniform(0.3, 0.7, int(np.prod(s))) for s in grid_shapes[:n_regions]]
+    _, gradient = opt(rho_vector=[w.copy() for w in weights])
+    if n_regions == 1:
+        gradient = [gradient]
+    gradient = [np.asarray(g).ravel() for g in gradient]
+
+    def objective(ws):
+        value, _ = opt(rho_vector=[w.copy() for w in ws], need_gradient=False)
+        return float(np.ravel(value)[0])
+
+    direction = [rng.standard_normal(w.size) for w in weights]
+    norm = np.sqrt(sum(float(np.sum(d**2)) for d in direction))
+    direction = [d / norm for d in direction]
+
+    adjoint = sum(float(np.dot(gradient[k], direction[k])) for k in range(len(weights)))
+    finite_difference = (
+        objective([weights[k] + step * direction[k] for k in range(len(weights))])
+        - objective([weights[k] - step * direction[k] for k in range(len(weights))])
+    ) / (2 * step)
+    return adjoint / finite_difference
+
+
+class TestAdjacentDesignGrids(unittest.TestCase):
+    # What the field solve reproduces across the two runs of the central
+    # difference, not what the gradient assembly contributes.
+    tol = 2e-3
+
+    def test_single_design_region(self):
+        """A single design region covering the whole area."""
+        shapes = [(8, 16)]
+        ratio = _directional_ratio(_problem(shapes), shapes)
+        self.assertAlmostEqual(ratio, 1.0, delta=self.tol)
+
+    def test_separated_design_regions(self):
+        """Two design regions separated by a pixel of background."""
+        shapes = [(8, 8), (8, 13)]
+        ratio = _directional_ratio(_problem(shapes, gap=1.0 / RESOLUTION), shapes)
+        self.assertAlmostEqual(ratio, 1.0, delta=self.tol)
+
+    def test_adjacent_design_regions(self):
+        """Two design regions sharing a boundary, with equal grid sizes."""
+        shapes = [(8, 8), (8, 8)]
+        ratio = _directional_ratio(_problem(shapes, gap=0.0), shapes)
+        self.assertAlmostEqual(ratio, 1.0, delta=self.tol)
+
+    def test_adjacent_design_regions_unequal_sizes(self):
+        """Two design regions sharing a boundary, with unequal grid sizes.
+
+        The larger neighbour's grid_size must not be used to index this
+        region's smaller gradient array.
+        """
+        shapes = [(8, 8), (8, 13)]
+        ratio = _directional_ratio(_problem(shapes, gap=0.0), shapes)
+        self.assertAlmostEqual(ratio, 1.0, delta=self.tol)
+
+    def test_overlapping_grids_are_unaffected(self):
+        """Two grids overlapping in one volume, combined with U_MEAN."""
+        shapes = [(8, 8), (8, 8)]
+        ratio = _directional_ratio(
+            _problem(shapes, overlap=True, grid_type="U_MEAN"), shapes, n_regions=1
+        )
+        self.assertAlmostEqual(ratio, 1.0, delta=self.tol)
+
+    def test_overlapping_grids_unequal_sizes(self):
+        """An overlapping grid must not supply the owner's interpolation size."""
+        # Keep the owner larger than the top grid so the regression is a wrong
+        # derivative rather than a test that relies on an out-of-bounds access.
+        shapes = [(8, 13), (8, 8)]
+        ratio = _directional_ratio(
+            _problem(shapes, overlap=True, grid_type="U_MEAN"), shapes, n_regions=1
+        )
+        # Single precision loses accuracy to cancellation in the outer
+        # objective finite difference; keep the double-precision bound unchanged.
+        tol = 5e-3 if mp.is_single_precision() else self.tol
+        self.assertAlmostEqual(ratio, 1.0, delta=tol)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/tests/test_adjoint_symmetric_grids.py

@@ -0,0 +1,142 @@
+"""A MaterialGrid shared by several transformed copies of itself.
+
+Placing one MaterialGrid in several overlapping Blocks, combined with U_MEAN, is
+how a design is constrained to a symmetry: a design variable then backs every
+copy at once. The adjoint gradient must account for all of them, so it has to
+agree with a finite difference taken with respect to that shared variable.
+"""
+
+import unittest
+
+import numpy as np
+from autograd import numpy as npa
+
+import meep as mp
+import meep.adjoint as mpa
+
+RESOLUTION = 20
+DPML = 1.0
+CELL = 6.0
+LX, LY = 1.0, 1.0
+N = 8
+FCEN = 1 / 1.55
+SILICON = mp.Medium(index=3.4)
+OXIDE = mp.Medium(index=1.44)
+
+
+def _problem(transform):
+    """One design grid, optionally overlapped with a transformed copy of itself.
+
+    `transform` is None, "rot90" or "mirror". The copy reuses the same
+    MaterialGrid object, which is what makes it the same design variable.
+    """
+    grid = mp.MaterialGrid(
+        mp.Vector3(N, N),
+        OXIDE,
+        SILICON,
+        weights=np.zeros((N, N)),
+        do_averaging=False,
+        beta=0,
+        eta=0.5,
+        grid_type="U_MEAN",
+    )
+    volume = mp.Volume(center=mp.Vector3(), size=mp.Vector3(LX, LY))
+    blocks = [mp.Block(center=volume.center, size=volume.size, material=grid)]
+    if transform == "rot90":
+        blocks.append(
+            mp.Block(
+                center=volume.center,
+                size=volume.size,
+                material=grid,
+                e1=mp.Vector3(1).rotate(mp.Vector3(z=1), np.pi / 2),
+                e2=mp.Vector3(y=1).rotate(mp.Vector3(z=1), np.pi / 2),
+            )
+        )
+    elif transform == "mirror":
+        blocks.append(
+            mp.Block(
+                center=volume.center,
+                size=volume.size,
+                material=grid,
+                e1=mp.Vector3(-1, 0, 0),
+                e2=mp.Vector3(0, 1, 0),
+            )
+        )
+
+    sim = mp.Simulation(
+        cell_size=mp.Vector3(CELL, CELL),
+        resolution=RESOLUTION,
+        boundary_layers=[mp.PML(DPML)],
+        geometry=blocks,
+        sources=[
+            mp.Source(
+                mp.GaussianSource(FCEN, fwidth=0.1),
+                component=mp.Ez,
+                center=mp.Vector3(-CELL / 2 + DPML + 0.3, 0),
+                size=mp.Vector3(0, CELL - 2 * DPML),
+            )
+        ],
+        default_material=OXIDE,
+        dimensions=2,
+    )
+    monitor = mp.Volume(
+        center=mp.Vector3(CELL / 2 - DPML - 0.3, 0), size=mp.Vector3(0, 2.0)
+    )
+    return mpa.OptimizationProblem(
+        simulation=sim,
+        objective_functions=[lambda ez: npa.sum(npa.abs(ez) ** 2, axis=1)],
+        objective_arguments=[mpa.FourierFields(sim, monitor, mp.Ez, yee_grid=False)],
+        design_regions=[mpa.DesignRegion(grid, volume=volume)],
+        frequencies=[FCEN],
+        decay_by=1e-9,
+    )
+
+
+def _directional_ratio(opt, seed=0, step=3e-3):
+    """adjoint / finite difference along one random direction.
+
+    Directional rather than single-pixel: perturbing one design variable moves
+    the objective by about as much as the DFT convergence floor.
+    """
+    rng = np.random.default_rng(seed)
+    weights = rng.uniform(0.3, 0.7, N * N)
+    _, gradient = opt(rho_vector=[weights.copy()])
+    gradient = np.asarray(gradient).ravel()
+
+    def objective(w):
+        value, _ = opt(rho_vector=[w.copy()], need_gradient=False)
+        return float(np.ravel(value)[0])
+
+    direction = rng.standard_normal(weights.size)
+    direction /= np.linalg.norm(direction)
+    adjoint = float(np.dot(gradient, direction))
+    finite_difference = (
+        objective(weights + step * direction) - objective(weights - step * direction)
+    ) / (2 * step)
+    return adjoint / finite_difference
+
+
+class TestSymmetricDesignGrids(unittest.TestCase):
+    # What the field solve reproduces across the two runs of the central
+    # difference, not what the gradient assembly contributes.
+    tol = 2e-3
+
+    def test_no_transform(self):
+        """A single block, with no overlapping copy."""
+        self.assertAlmostEqual(_directional_ratio(_problem(None)), 1.0, delta=self.tol)
+
+    def test_rotated_copy(self):
+        """Overlapped with a 90-degree rotation of itself."""
+        self.assertAlmostEqual(
+            _directional_ratio(_problem("rot90")), 1.0, delta=self.tol
+        )
+
+    def test_mirrored_copy(self):
+        """Overlapped with a reflection of itself."""
+        self.assertAlmostEqual(
+            _directional_ratio(_problem("mirror")), 1.0, delta=self.tol
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()