Fix 3D MaterialGrid subpixel smoothing: kernel normalization (#3258) and gradient consistency (#3259)

python/Makefile.am

@@ -94,6 +94,7 @@ TESTS =                                        \
     $(TEST_DIR)/test_special_kz.py             \
     $(TEST_DIR)/test_source.py                 \
     $(TEST_DIR)/test_stop_when_flux_decayed.py \
+    $(TEST_DIR)/test_subpixel_3d.py            \
     $(TEST_DIR)/test_timing_measurements.py    \
     $(TEST_DIR)/test_user_defined_material.py  \
     $(TEST_DIR)/test_verbosity_mgr.py          \

python/meep.i

@@ -1939,6 +1939,14 @@ meep_geom::geom_epsilon* _set_materials(meep::structure * s,
         meep_geom::set_materials_from_geom_epsilon(s, geps, use_anisotropic_averaging, tol,
                                              maxeval,alist);
     }
+    else {
+        /* set_materials_from_geom_epsilon() is what normally records these for the
+           gradient calculation. Skipping it leaves the geom_epsilon on its header
+           defaults (DEFAULT_SUBPIXEL_TOL/MAXEVAL), so the adjoint would smooth no
+           matter what the caller asked for. */
+        geps->tol = tol;
+        geps->maxeval = use_anisotropic_averaging ? maxeval : 0;
+    }
 
     if (meep::verbosity > 1 && !split_chunks_evenly && set_materials) {
       int num_procs = meep::count_processors();

python/simulation.py

@@ -1246,6 +1246,7 @@ def __init__(
         progress_interval: float = 4,
         subpixel_tol: float = 1e-4,
         subpixel_maxeval: int = 100000,
+        allow_3d_subpixel: bool = True,
         loop_tile_base_db: int = 0,
         loop_tile_base_eh: int = 0,
         ensure_periodicity: bool = True,
@@ -1395,6 +1396,13 @@ def __init__(
           effects and irregular
           convergence](Subpixel_Smoothing.md#what-happens-when-subpixel-smoothing-is-disabled).
 
+        + **`allow_3d_subpixel` [ `boolean` ]** — If `False`, then in a 3d simulation
+          any `MaterialGrid` with `do_averaging=True` has its level-set subpixel
+          smoothing turned off at structure-initialization time. This is a
+          convenience switch so a 3d run can drop MaterialGrid smoothing without
+          editing every grid individually; the default `True` leaves behavior
+          unchanged. Has no effect in 1d/2d/cylindrical.
+
         + **`force_complex_fields` [ `boolean` ]** — By default, Meep runs its simulations
           with purely real fields whenever possible. It uses complex fields which require
           twice the memory and computation if the `k_point` is non-zero or if `m` is
@@ -1502,6 +1510,7 @@ def __init__(
         self.eps_averaging = eps_averaging
         self.subpixel_tol = subpixel_tol
         self.subpixel_maxeval = subpixel_maxeval
+        self.allow_3d_subpixel = allow_3d_subpixel
         self.loop_tile_base_db = loop_tile_base_db
         self.loop_tile_base_eh = loop_tile_base_eh
         self.ensure_periodicity = ensure_periodicity
@@ -2002,11 +2011,45 @@ def _compute_fragment_stats(self, gv):
 
         return stats
 
+    def _iter_material_grids(self):
+        """Yield every MaterialGrid reachable from this simulation's materials."""
+        candidates = list(self.geometry or [])
+        for obj in candidates:
+            mat = getattr(obj, "material", None)
+            if isinstance(mat, mp.MaterialGrid):
+                yield mat
+        if isinstance(self.default_material, mp.MaterialGrid):
+            yield self.default_material
+        for mat in self.extra_materials or []:
+            if isinstance(mat, mp.MaterialGrid):
+                yield mat
+
+    def _apply_3d_subpixel_policy(self):
+        """Honor `allow_3d_subpixel=False` by disabling MaterialGrid smoothing.
+
+        A single simulation-level switch, so a 3d run can drop level-set subpixel
+        smoothing without editing every MaterialGrid it happens to contain. The
+        default is True, which leaves behavior exactly as it was.
+        """
+        if self.allow_3d_subpixel or self.dimensions != 3:
+            return
+        disabled = 0
+        for grid in self._iter_material_grids():
+            if grid.do_averaging:
+                grid.do_averaging = False
+                disabled += 1
+        if disabled and verbosity.meep > 0:
+            print(
+                f"allow_3d_subpixel=False: disabled do_averaging on {disabled} "
+                "MaterialGrid(s)"
+            )
+
     def _init_structure(self, k=False):
         if verbosity.meep > 0:
             print("-" * 11)
             print("Initializing structure...")
 
+        self._apply_3d_subpixel_policy()
         gv = self._create_grid_volume(k)
         sym = self._create_symmetries(gv)
         br = _create_boundary_region_from_boundary_layers(self.boundary_layers, gv)

python/tests/test_subpixel_3d.py

@@ -0,0 +1,198 @@
+"""Regression coverage for MaterialGrid subpixel smoothing in 3d.
+
+Every adjoint test upstream runs in 2d or cylindrical coordinates, which left two
+3d-only defects unnoticed:
+
+  1. get_uproj_w()'s D3 branch normalized by `4 / 3 * pi * rad^3`. `4 / 3` is
+     integer division, so the denominator was `pi * rad^3` and the smoothing
+     kernel integrated to 4/3 instead of 1, scaling epsilon everywhere the
+     material grid had a nonzero gradient.
+
+  2. set_materials_from_geom_epsilon() stored the raw `maxeval` on the
+     geom_epsilon regardless of `use_anisotropic_averaging`, while
+     structure_chunk::set_chi1inv() zeroes it when averaging is off. With
+     eps_averaging=False the forward solve therefore saw an unsmoothed operator
+     while the adjoint differentiated a smoothed one.
+"""
+
+import unittest
+
+import numpy as np
+
+import meep as mp
+
+
+# Allow finite-difference and MPI chunking noise while remaining far below the
+# 70-100% disagreement caused by the regression covered by these tests.
+FD_REL_TOL = 1e-2
+
+
+def _epsilon(dim, do_averaging, n1, n2, resolution=20, n=20):
+    """Dielectric array over a MaterialGrid carrying a linear ramp in u."""
+    design, pad = 1.0, 0.6
+    sz = 0 if dim == 2 else design + 2 * pad
+
+    ramp = np.linspace(0.25, 0.75, n)
+    weights = np.repeat(ramp[:, None], n, axis=1)[:, :, None]
+
+    grid = mp.MaterialGrid(
+        mp.Vector3(n, n, 1),
+        mp.Medium(index=n1),
+        mp.Medium(index=n2),
+        weights=weights,
+        do_averaging=do_averaging,
+        beta=0,
+    )
+    sim = mp.Simulation(
+        cell_size=mp.Vector3(design + 2 * pad, design + 2 * pad, sz),
+        resolution=resolution,
+        default_material=mp.Medium(index=n1),
+        geometry=[
+            mp.Block(
+                center=mp.Vector3(),
+                size=mp.Vector3(design, design, mp.inf if dim == 2 else design),
+                material=grid,
+            )
+        ],
+        dimensions=dim,
+        eps_averaging=True,
+    )
+    sim.init_sim()
+    extent = design * 0.8
+    return np.asarray(
+        sim.get_array(
+            component=mp.Dielectric,
+            center=mp.Vector3(),
+            size=mp.Vector3(extent, extent, 0 if dim == 2 else extent),
+        )
+    )
+
+
+class TestSubpixelKernelNormalization(unittest.TestCase):
+    """The smoothing kernel must integrate to 1 in every dimensionality.
+
+    Measuring that through eps needs care: smoothing produces an anisotropic
+    tensor whose normal component is a harmonic mean and whose transverse
+    components are arithmetic means. With real contrast those differ, and the
+    resulting (legitimate) shift swamps the normalization signal. Removing the
+    contrast collapses both means onto the same value, leaving
+
+        eps(do_averaging=True) / eps(do_averaging=False) == integral(w)
+
+    so the ratio reads the normalization off directly.
+    """
+
+    def test_kernel_normalized_2d(self):
+        ratio = _epsilon(2, True, 1.44, 1.4414) / _epsilon(2, False, 1.44, 1.4414)
+        self.assertAlmostEqual(float(np.median(ratio)), 1.0, places=6)
+
+    def test_kernel_normalized_3d(self):
+        ratio = _epsilon(3, True, 1.44, 1.4414) / _epsilon(3, False, 1.44, 1.4414)
+        # Before the fix this was 18/17 = 1.058823..., the closed form of
+        # 3 / (K + 2/K) at K = 4/3.
+        self.assertAlmostEqual(float(np.median(ratio)), 1.0, places=6)
+
+
+def _directional_fd(do_averaging, eps_averaging, resolution=12, n=6, dp=1e-3, seed=0):
+    """Return (fd, adjoint) directional derivatives along the adjoint gradient."""
+    import autograd.numpy as npa
+    import meep.adjoint as mpa
+
+    si, clad = mp.Medium(index=3.48), mp.Medium(index=1.44)
+    pml, port_pad, side_pad = 0.5, 0.8, 0.5
+    design, thickness = 1.0, 0.22
+
+    sx = 2 * pml + 2 * port_pad + design
+    sy = 2 * pml + 2 * side_pad + design
+    sz = 2 * pml + 2 * side_pad + thickness
+
+    rng = np.random.default_rng(seed)
+    weights = rng.uniform(0.2, 0.8, size=n * n)
+
+    grid = mp.MaterialGrid(
+        mp.Vector3(n, n, 1),
+        clad,
+        si,
+        weights=weights.reshape(n, n, 1),
+        do_averaging=do_averaging,
+    )
+    region = mpa.DesignRegion(
+        grid,
+        volume=mp.Volume(
+            center=mp.Vector3(), size=mp.Vector3(design, design, thickness)
+        ),
+    )
+
+    fcen = 1 / 1.55
+    port = mp.Vector3(0, sy - 2 * pml, sz - 2 * pml)
+    sim = mp.Simulation(
+        cell_size=mp.Vector3(sx, sy, sz),
+        resolution=resolution,
+        boundary_layers=[mp.PML(pml)],
+        default_material=clad,
+        geometry=[
+            mp.Block(
+                center=mp.Vector3(),
+                size=mp.Vector3(mp.inf, 0.5, thickness),
+                material=si,
+            ),
+            mp.Block(center=region.center, size=region.size, material=grid),
+        ],
+        sources=[
+            mp.EigenModeSource(
+                mp.GaussianSource(fcen, fwidth=0.1 * fcen),
+                center=mp.Vector3(-(design / 2 + port_pad / 2)),
+                size=port,
+                eig_band=1,
+            )
+        ],
+        eps_averaging=eps_averaging,
+    )
+    monitor = mpa.EigenmodeCoefficient(
+        sim,
+        mp.Volume(center=mp.Vector3(design / 2 + port_pad / 2), size=port),
+        mode=1,
+    )
+    opt = mpa.OptimizationProblem(
+        simulation=sim,
+        objective_functions=[lambda c: npa.abs(c) ** 2],
+        objective_arguments=[monitor],
+        design_regions=[region],
+        frequencies=[fcen],
+        decay_by=1e-6,
+    )
+
+    _, gradient = opt([weights], need_gradient=True)
+    gradient = np.asarray(np.real(np.squeeze(gradient)), dtype=np.float64).reshape(-1)
+
+    # Along the gradient the directional derivative is as large as it gets, which
+    # keeps the difference quotient clear of the forward solve's convergence noise.
+    direction = gradient / np.linalg.norm(gradient)
+    f_plus, _ = opt([weights + dp * direction], need_gradient=False)
+    f_minus, _ = opt([weights - dp * direction], need_gradient=False)
+    fd = (float(np.squeeze(f_plus)) - float(np.squeeze(f_minus))) / (2 * dp)
+    return fd, float(gradient @ direction)
+
+
+class TestAdjointGradient3D(unittest.TestCase):
+    def test_gradient_matches_fd_without_smoothing(self):
+        fd, adj = _directional_fd(do_averaging=False, eps_averaging=False)
+        self.assertAlmostEqual(fd / adj, 1.0, delta=FD_REL_TOL)
+
+    def test_gradient_matches_fd_with_smoothing(self):
+        fd, adj = _directional_fd(do_averaging=True, eps_averaging=True)
+        self.assertAlmostEqual(fd / adj, 1.0, delta=FD_REL_TOL)
+
+    def test_do_averaging_ignored_when_eps_averaging_off(self):
+        """do_averaging=True with eps_averaging=False must not change the gradient.
+
+        The forward solve discards MaterialGrid smoothing when eps_averaging is
+        off, so the gradient has to discard it too. Before the fix the adjoint
+        kept smoothing and this pair disagreed by ~70-100%.
+        """
+        fd, adj = _directional_fd(do_averaging=True, eps_averaging=False)
+        self.assertAlmostEqual(fd / adj, 1.0, delta=FD_REL_TOL)
+
+
+if __name__ == "__main__":
+    unittest.main()

src/meepgeom.cpp

@@ -1219,8 +1219,12 @@ static void get_uproj_w(const matgrid_volavg *mgva, double x0, double &u_proj, d
   else if (mgva->dim == meep::D2 || mgva->dim == meep::Dcyl)
     w = 2 * sqrt(mgva->rad * mgva->rad - x0 * x0) / (meep::pi * mgva->rad * mgva->rad);
   else if (mgva->dim == meep::D3)
+    // 4.0 / 3.0, not 4 / 3: the latter is integer division, which drops the
+    // denominator to pi*rad^3 instead of the sphere volume (4/3)*pi*rad^3 and
+    // leaves the kernel integrating to 4/3 rather than 1. The 1d and 2d branches
+    // above are correctly normalized, so this only ever affected 3d.
     w = meep::pi * (mgva->rad * mgva->rad - x0 * x0) /
-        (4 / 3 * meep::pi * mgva->rad * mgva->rad * mgva->rad);
+        (4.0 / 3.0 * meep::pi * mgva->rad * mgva->rad * mgva->rad);
 }
 
 #ifdef CTL_HAS_COMPLEX_INTEGRATION
@@ -2014,9 +2018,14 @@ void set_materials_from_geom_epsilon(meep::structure *s, geom_epsilon *geps,
                                      bool use_anisotropic_averaging, double tol, int maxeval,
                                      absorber_list alist) {
 
-  // store for later use in gradient calculations
+  // Store for later use in gradient calculations. These must mirror what
+  // structure_chunk::set_chi1inv() actually used, and that routine zeroes maxeval
+  // when averaging is off (see anisotropic_averaging.cpp). Passing the raw maxeval
+  // through would make get_material_gradient() finite-difference a *smoothed*
+  // operator that the forward solve never saw, so the adjoint gradient would not
+  // be the derivative of the simulation being run.
   geps->tol = tol;
-  geps->maxeval = maxeval;
+  geps->maxeval = use_anisotropic_averaging ? maxeval : 0;
 
   meep::grid_volume gv = s->gv;
   if (alist) {