[2/n][Adjoint Module] Fix source deposition across chunk boundaries

python/tests/test_source.py

@@ -125,6 +125,114 @@ def test_gaussian_fourier_transform_continuous(self):
             self.assertEqual(src.fourier_transform(f), src.swigobj.fourier_transform(f))
 
 
+class TestSourceChunkBoundary(unittest.TestCase):
+    @staticmethod
+    def _get_source(
+        cell_size,
+        geometry_center,
+        source_center,
+        source_size,
+        chunk_layout,
+        run_until=None,
+    ):
+        sim = mp.Simulation(
+            cell_size=cell_size,
+            geometry_center=geometry_center,
+            resolution=8,
+            chunk_layout=chunk_layout,
+            sources=[
+                mp.Source(
+                    mp.ContinuousSource(1.0),
+                    component=mp.Ez,
+                    center=source_center,
+                    size=source_size,
+                )
+            ],
+        )
+        try:
+            sim.init_sim()
+            source = sim.get_source(mp.Ez)
+            energy = None
+            if run_until is not None:
+                sim.run(until=run_until)
+                energy = sim.field_energy_in_box(
+                    mp.Volume(center=geometry_center, size=cell_size)
+                )
+            return source, energy
+        finally:
+            sim.reset_meep()
+
+    def test_zero_thickness_source_at_chunk_boundary(self):
+        cell_size = mp.Vector3(2.5, 2.25, 2.75)
+        geometry_center = mp.Vector3(0.0625, -0.0625, 0.0625)
+        one_chunk = mp.BinaryPartition(data=0)
+        split_at_source = mp.BinaryPartition(data=[(mp.Z, 0.0), 0, 1])
+
+        # The off-boundary case is a control that the split alone does not
+        # change deposition. The boundary case exercises interpolation onto
+        # both Ez planes adjacent to the split.
+        for source_z in (0.25, 0.0):
+            with self.subTest(source_z=source_z):
+                source_center = mp.Vector3(0.0625, -0.0625, source_z)
+                source_size = mp.Vector3(1.5, 1.25, 0)
+                expected_source, expected_energy = self._get_source(
+                    cell_size,
+                    geometry_center,
+                    source_center,
+                    source_size,
+                    one_chunk,
+                    run_until=0.25 if source_z == 0 else None,
+                )
+                actual_source, actual_energy = self._get_source(
+                    cell_size,
+                    geometry_center,
+                    source_center,
+                    source_size,
+                    split_at_source,
+                    run_until=0.25 if source_z == 0 else None,
+                )
+
+                self.assertTrue(np.any(expected_source != 0))
+                np.testing.assert_array_equal(actual_source != 0, expected_source != 0)
+                np.testing.assert_array_equal(actual_source, expected_source)
+                if expected_energy is not None:
+                    # Source introspection can include non-owned storage that
+                    # never injects into the evolved fields.
+                    rtol = 1e-5 if mp.is_single_precision() else 1e-12
+                    np.testing.assert_allclose(
+                        actual_energy, expected_energy, rtol=rtol, atol=0
+                    )
+
+    def test_point_source_at_chunk_boundary_2d(self):
+        args = (
+            mp.Vector3(2.5, 2.25),
+            mp.Vector3(0.0625, -0.0625),
+            mp.Vector3(0.0625, -0.0625),
+            mp.Vector3(),
+        )
+        expected, _ = self._get_source(*args, mp.BinaryPartition(data=0))
+        actual, _ = self._get_source(
+            *args, mp.BinaryPartition(data=[(mp.X, 0.0), 0, 1])
+        )
+
+        self.assertTrue(np.any(expected != 0))
+        np.testing.assert_array_equal(actual != 0, expected != 0)
+        np.testing.assert_array_equal(actual, expected)
+
+    def test_point_source_on_owned_chunk_boundary_2d(self):
+        # This centered-grid control guards the historical failure mode where
+        # both chunks deposited the same already-owned source sample.
+        args = (mp.Vector3(2, 2), mp.Vector3(), mp.Vector3(), mp.Vector3())
+        expected, _ = self._get_source(*args, mp.BinaryPartition(data=0))
+        actual, _ = self._get_source(
+            *args, mp.BinaryPartition(data=[(mp.X, 0.0), 0, 1])
+        )
+
+        self.assertTrue(np.any(expected != 0))
+        np.testing.assert_array_equal(actual != 0, expected != 0)
+        np.testing.assert_array_equal(actual, expected)
+
+
 class TestSourceTypemaps(unittest.TestCase):
     def setUp(self):
         def dummy_eps(v):

src/fix_boundary_sources.cpp

@@ -86,11 +86,17 @@ void fields::fix_boundary_sources() {
               src.set_amplitude(ipt, 0.0); // will no longer be needed
             }
           }
+          src.remove_zero_amplitudes();
           src.needs_boundary_fix = false;
         }
     }
   }
 
+  if (sum_to_all(boundarysources.size()) == 0) {
+    finished_working();
+    return;
+  }
+
   // we need each process's data to be contiguous
   srcpt_info_compare compare = {chunks};
   std::sort(boundarysources.begin(), boundarysources.end(), compare);

src/loop_in_chunks.cpp

@@ -418,8 +418,11 @@ void fields::loop_in_chunks(field_chunkloop chunkloop, void *chunkloop_data, con
       for (int i = 0; i < num_chunks; ++i) {
         if (!chunks[i]->is_mine()) continue;
         grid_volume gvu(chunks[i]->gv);
-        ivec _iscoS(S.transform(gvu.little_owned_corner(cS), sn));
-        ivec _iecoS(S.transform(gvu.big_owned_corner(cS), sn));
+        ivec _iscoS(S.transform(use_symmetry ? gvu.little_owned_corner(cS)
+                                             : gvu.little_corner() + gvu.iyee_shift(cS),
+                                sn));
+        ivec _iecoS(S.transform(
+            use_symmetry ? gvu.big_owned_corner(cS) : gvu.big_corner() + gvu.iyee_shift(cS), sn));
         ivec iscoS(max(user_volume.little_owned_corner(cgrid), min(_iscoS, _iecoS))),
             iecoS(max(_iscoS, _iecoS)); // fix ordering due to to transform
 
@@ -445,6 +448,18 @@ void fields::loop_in_chunks(field_chunkloop chunkloop, void *chunkloop_data, con
         }
         ivec iecS(min(ie - shifti, iecoS));
 
+        // A chunk origin may lie between points of the global component grid.
+        // For source loops (use_symmetry=false), keep the loop aligned with the
+        // global grid so that interpolation weights are independent of chunking.
+        if (!use_symmetry) {
+          LOOP_OVER_DIRECTIONS(gv.dim, d) {
+            const int parity = (iscS.in_direction(d) - (is - shifti).in_direction(d)) % 2;
+            if (parity) iscS.set_direction(d, iscS.in_direction(d) + 1);
+            if ((iecS.in_direction(d) - (is - shifti).in_direction(d)) % 2)
+              iecS.set_direction(d, iecS.in_direction(d) - 1);
+          }
+        }
+
         if (iscS <= iecS) { // non-empty intersection
           // Determine weights at chunk looping boundaries:
           ivec isc(S.transform(iscS, -sn)), iec(S.transform(iecS, -sn));

src/meep_internals.hpp

@@ -72,6 +72,8 @@ class src_vol {
   // Requirement: other.num_points() == this->num_points().
   // It is recommended to use `combinable` before calling this method.
   void add_amplitudes_from(const src_vol &other);
+  // Removes entries that no longer deposit a source after boundary relocation.
+  void remove_zero_amplitudes();
 
   const component c;       // field component the source applies to
   bool needs_boundary_fix; // whether fix_boundary_sources needs calling

src/sources.cpp

@@ -205,6 +205,17 @@ void src_vol::add_amplitudes_from(const src_vol &other) {
   }
 }
 
+void src_vol::remove_zero_amplitudes() {
+  size_t dst = 0;
+  for (size_t src = 0; src < amp.size(); ++src)
+    if (amp[src] != 0.0) {
+      index[dst] = index[src];
+      amp[dst++] = amp[src];
+    }
+  index.resize(dst);
+  amp.resize(dst);
+}
+
 /*********************************************************************/
 
 // THIS VARIANT IS FOR BACKWARDS COMPATIBILITY, and is DEPRECATED:
@@ -326,7 +337,7 @@ static void src_vol_chunkloop(fields_chunk *fc, int ichunk, component c, ivec is
     meep::abort("add_volume_source: computed wrong npts (%zd vs. %zd)", npts, idx_vol);
 
   field_type ft = is_H_or_B(c) ? B_stuff : D_stuff;
-  fc->add_source(ft, src_vol(c, data->src, std::move(index_array), std::move(amps_array)));
+  fc->add_source(ft, src_vol(c, data->src, std::move(index_array), std::move(amps_array), true));
 }
 
 void fields::add_srcdata(struct sourcedata cur_data, src_time *src, size_t n,
@@ -501,8 +512,10 @@ void fields::add_volume_source(component c, const src_time &src, const volume &w
       data.amp *= gv.a; // correct units for J delta-function amplitude
   }
   sources = src.add_to(sources, &data.src);
+  register_src_time(data.src);
   data.center = (where.get_min_corner() + where.get_max_corner()) * 0.5;
   loop_in_chunks(src_vol_chunkloop, (void *)&data, where, c, false);
+  fix_boundary_sources();
   require_component(c);
 }