Resolve m-number adjoint bugs introduced by 1855 (#2194)

* add function to update m number

* update change m number too

* precommit cleanup

* naming fixes, added check to change_m(), improved tests

* fix precommit errors

Co-authored-by: Alec Hammond <alechammond@fb.com>

python/adjoint/optimization_problem.py

@@ -243,7 +243,7 @@ def adjoint_run(self):
 
         # flip the m number
         if utils._check_if_cylindrical(self.sim):
-            self.sim.m = -self.sim.m
+            self.sim.change_m(-self.sim.m)
 
         # flip the k point
         if self.sim.k_point:
@@ -279,11 +279,11 @@ def adjoint_run(self):
 
         # reset the m number
         if utils._check_if_cylindrical(self.sim):
-            self.sim.m = -self.sim.m
+            self.sim.change_m(-self.sim.m)
 
         # reset the k point
         if self.sim.k_point:
-            self.sim.k_point *= -1
+            self.sim.change_k_point(-1 * self.sim.k_point)
 
         # update optimizer's state
         self.current_state = "ADJ"

python/simulation.py

@@ -4345,6 +4345,21 @@ def change_k_point(self, k):
                         py_v3_to_vec(self.dimensions, self.k_point, self.is_cylindrical)
                     )
 
+    def change_m(self, m: float) -> None:
+        """Changes the simulation's `m` number (the angular ϕ dependence)."""
+        self.m = m
+
+        if self.fields:
+            needs_complex_fields = not (not self.m or self.m == 0)
+
+            if needs_complex_fields and self.fields.is_real:
+                self.fields = None
+                self._is_initialized = False
+                self.init_sim()
+            else:
+                if self.m is not None:
+                    self.fields.change_m(m)
+
     def change_sources(self, new_sources):
         """
         Change the list of sources in `Simulation.sources` to `new_sources`, and changes

python/tests/test_adjoint_cyl.py

@@ -12,11 +12,11 @@
 from autograd import numpy as npa
 from autograd import tensor_jacobian_product
 from utils import ApproxComparisonTestCase
+import parameterized
 
 rng = np.random.RandomState(2)
 resolution = 20
 dimensions = mp.CYLINDRICAL
-m = 0
 Si = mp.Medium(index=3.4)
 SiO2 = mp.Medium(index=1.44)
 
@@ -48,7 +48,7 @@
 dp = deps * rng.rand(Nr * Nz)
 
 
-def forward_simulation(design_params):
+def forward_simulation(design_params, m, far_x):
     matgrid = mp.MaterialGrid(
         mp.Vector3(Nr, 0, Nz), SiO2, Si, weights=design_params.reshape(Nr, 1, Nz)
     )
@@ -72,7 +72,7 @@ def forward_simulation(design_params):
     )
 
     frequencies = [fcen]
-    far_x = [mp.Vector3(5, 0, 20)]
+
     mode = sim.add_near2far(
         frequencies,
         mp.Near2FarRegion(
@@ -89,7 +89,7 @@ def forward_simulation(design_params):
     return abs(Er[0]) ** 2
 
 
-def adjoint_solver(design_params):
+def adjoint_solver(design_params, m, far_x):
 
     design_variables = mp.MaterialGrid(mp.Vector3(Nr, 0, Nz), SiO2, Si)
     design_region = mpa.DesignRegion(
@@ -117,7 +117,6 @@ def adjoint_solver(design_params):
         m=m,
     )
 
-    far_x = [mp.Vector3(5, 0, 20)]
     NearRegions = [
         mp.Near2FarRegion(
             center=mp.Vector3(design_r / 2, 0, (sz / 2 - dpml + design_z / 2) / 2),
@@ -149,12 +148,21 @@ def J(alpha):
 
 
 class TestAdjointSolver(ApproxComparisonTestCase):
-    def test_adjoint_solver_cyl_n2f_fields(self):
+    @parameterized.parameterized.expand(
+        [
+            (0, [mp.Vector3(5, 0, 20)]),
+            (0, [mp.Vector3(4, 0, 28)]),
+            (-1, [mp.Vector3(5, 0, 20)]),
+            (1.2, [mp.Vector3(5, 0, 20)]),
+        ]
+    )
+    def test_adjoint_solver_cyl_n2f_fields(self, m, far_x):
         print("*** TESTING CYLINDRICAL Near2Far ADJOINT FEATURES ***")
-        adjsol_obj, adjsol_grad = adjoint_solver(p)
+        print(f"Current test: m={m}, far_x={far_x}")
+        adjsol_obj, adjsol_grad = adjoint_solver(p, m, far_x)
 
         ## compute unperturbed S12
-        S12_unperturbed = forward_simulation(p)
+        S12_unperturbed = forward_simulation(p, m, far_x)
 
         ## compare objective results
         print(
@@ -164,7 +172,7 @@ def test_adjoint_solver_cyl_n2f_fields(self):
         self.assertClose(adjsol_obj, S12_unperturbed, epsilon=1e-3)
 
         ## compute perturbed S12
-        S12_perturbed = forward_simulation(p + dp)
+        S12_perturbed = forward_simulation(p + dp, m, far_x)
 
         ## compare gradients
         if adjsol_grad.ndim < 2:

src/fields.cpp

@@ -793,4 +793,19 @@ bool operator==(const comms_key &lhs, const comms_key &rhs) {
   return (lhs.ft == rhs.ft) && (lhs.phase == rhs.phase) && (lhs.pair == rhs.pair);
 }
 
+void fields::change_m(double new_m) {
+  m = new_m;
+  if ((new_m != 0) && (is_real)) {
+    meep::abort("The simulation must be reinitialized if switching to complex fields!\n");
+  }
+
+  if ((new_m == 0) && (!is_real)) { use_real_fields(); }
+
+  for (int i = 0; i < num_chunks; i++) {
+    chunks[i]->change_m(new_m);
+  }
+}
+
+void fields_chunk::change_m(double new_m) { m = new_m; }
+
 } // namespace meep

src/meep.hpp

@@ -1546,6 +1546,7 @@ class fields_chunk {
   void remove_sources();
   void remove_susceptibilities(bool shared_chunks);
   void zero_fields();
+  void change_m(double new_m);
 
   // update_eh.cpp
   bool needs_W_prev(component c) const;
@@ -1751,6 +1752,7 @@ class fields {
   void remove_fluxes();
   void reset();
   void log(const char *prefix = "");
+  void change_m(double new_m);
 
   // time.cpp
   std::vector<double> time_spent_on(time_sink sink);

src/meepgeom.cpp

@@ -2651,35 +2651,35 @@ get_material_gradient(const meep::vec &r,              // current point
   material_type md;
   geps->get_material_pt(md, r);
 
+  // get the tensor column index corresponding to the forward component
   int dir_idx = 0;
-  if (forward_c == meep::Dx || forward_c == meep::Dr)
-    dir_idx = 0;
-  else if (forward_c == meep::Dy || forward_c == meep::Dp)
-    dir_idx = 1;
-  else if (forward_c == meep::Dz)
-    dir_idx = 2;
-  else
-    meep::abort("Invalid adjoint field component");
+  switch (meep::component_direction(forward_c)) {
+    case meep::X:
+    case meep::R: dir_idx = 0; break;
+    case meep::Y:
+    case meep::P: dir_idx = 1; break;
+    case meep::Z: dir_idx = 2; break;
+    case meep::NO_DIRECTION: meep::abort("Invalid forward component!\n");
+  }
 
+  // materials are non-dispersive
   if (md->trivial) {
     const double sd = 1.0; // FIXME: make user-changable?
     meep::volume v(r);
     LOOP_OVER_DIRECTIONS(dim, d) {
       v.set_direction_min(d, r.in_direction(d) - 0.5 * gv.inva * sd);
       v.set_direction_max(d, r.in_direction(d) + 0.5 * gv.inva * sd);
     }
-    double row_1[3], row_2[3], dA_du[3];
+    double row_1[3], row_2[3];
     double orig = u[idx];
     u[idx] -= du;
     geps->eff_chi1inv_row(adjoint_c, row_1, v, geps->tol, geps->maxeval);
     u[idx] += 2 * du;
     geps->eff_chi1inv_row(adjoint_c, row_2, v, geps->tol, geps->maxeval);
     u[idx] = orig;
-
-    for (int i = 0; i < 3; i++)
-      dA_du[i] = (row_1[i] - row_2[i]) / (2 * du);
-    return dA_du[dir_idx] * fields_f;
+    return fields_f * (row_1[dir_idx] - row_2[dir_idx]) / (2 * du);
   }
+  // materials have some dispersion
   else {
     double orig = u[idx];
     std::complex<double> row_1[3], row_2[3], dA_du[3];
@@ -2689,9 +2689,8 @@ get_material_gradient(const meep::vec &r,              // current point
     eff_chi1inv_row_disp(adjoint_c, row_2, r, freq, geps);
     u[idx] = orig;
 
-    for (int i = 0; i < 3; i++)
-      dA_du[i] = (row_1[i] - row_2[i]) / (2 * du);
-    return dA_du[dir_idx] * fields_f * cond_cmp(forward_c, r, freq, geps);
+    return fields_f * (row_1[dir_idx] - row_2[dir_idx]) / (2 * du) *
+           cond_cmp(forward_c, r, freq, geps);
   }
 }
 
@@ -2720,8 +2719,8 @@ void add_interpolate_weights(double rx, double ry, double rz, double *data, int
 /* define a macro to give us data(x,y,z) on the grid,
 in row-major order (the order used by HDF5): */
 #define IDX(x, y, z) (((x)*ny + (y)) * nz + (z)) * stride
-#define D(x, y, z) (data[(((x)*ny + (y)) * nz + (z)) * stride])
-#define U(x, y, z) (udata[(((x)*ny + (y)) * nz + (z)) * stride])
+#define D(x, y, z) (data[IDX(x, y, z)])
+#define U(x, y, z) (udata[IDX(x, y, z)])
 
   u = (((U(x1, y1, z1) * (1.0 - dx) + U(x2, y1, z1) * dx) * (1.0 - dy) +
         (U(x1, y2, z1) * (1.0 - dx) + U(x2, y2, z1) * dx) * dy) *