support adjoint gradients in cylindrical coordinate (#1749)

* cyl

* remove print

* fix

* bug

* update

* update

* cyl

* update

* update

* cyl

* fix

* update

* update

* cyl

* update

* update

* rebase fix

* rebase

* rebase fix

* minor fix

* clean up

* cyl test

* typo

* fix

Co-authored-by: Mo Chen <mochen@Mos-MacBook-Pro.local>

python/Makefile.am

@@ -38,6 +38,7 @@ TESTS =                                   \
     $(TEST_DIR)/test_3rd_harm_1d.py            \
     $(TEST_DIR)/test_absorber_1d.py            \
     $(TEST_DIR)/test_adjoint_solver.py         \
+		$(TEST_DIR)/test_adjoint_cyl.py            \
     $(TEST_DIR)/test_adjoint_jax.py            \
     $(TEST_DIR)/test_antenna_radiation.py      \
     $(TEST_DIR)/test_array_metadata.py         \

python/adjoint/optimization_problem.py

@@ -40,8 +40,8 @@ def get_gradient(self, sim, fields_a, fields_f, frequencies):
         f = sim.fields
         vol = sim._fit_volume_to_simulation(self.volume)
         # compute the gradient
-        mp._get_gradient(grad, fields_a, fields_f, vol, np.array(frequencies),
-                         geom_list, f)
+        sim_is_cylindrical = (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical
+        mp._get_gradient(grad,fields_a,fields_f,vol,np.array(frequencies),geom_list,f, sim_is_cylindrical)
 
         return np.squeeze(grad).T
 
@@ -78,6 +78,9 @@ def __init__(
     ):
 
         self.sim = simulation
+        self.components = [mp.Ex,mp.Ey,mp.Ez]
+        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+            self.components = [mp.Er,mp.Ep,mp.Ez]
 
         if isinstance(objective_functions, list):
             self.objective_functions = objective_functions
@@ -213,7 +216,7 @@ def prepare_forward_run(self):
         # register design region
         self.design_region_monitors = [
             self.sim.add_dft_fields(
-                [mp.Ex, mp.Ey, mp.Ez],
+                self.components,
                 self.frequencies,
                 where=dr.volume,
                 yee_grid=True,
@@ -223,10 +226,14 @@ def prepare_forward_run(self):
 
         # store design region voxel parameters
         self.design_grids = []
+        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+            YeeDims = namedtuple('YeeDims', ['Er','Ep','Ez'])
+        else:
+            YeeDims = namedtuple('YeeDims', ['Ex','Ey','Ez'])
         for drm in self.design_region_monitors:
             s = [
                 self.sim.get_array_slice_dimensions(c, vol=drm.where)[0]
-                for c in [mp.Ex, mp.Ey, mp.Ez]
+                for c in self.components
             ]
             self.design_grids += [YeeDims(*s)]
 
@@ -263,7 +270,7 @@ def forward_run(self):
             for c in dg
         ] for dg in self.design_grids]
         for nb, dgm in enumerate(self.design_region_monitors):
-            for ic, c in enumerate([mp.Ex, mp.Ey, mp.Ez]):
+            for ic, c in enumerate(self.components):
                 for f in range(self.nf):
                     self.d_E[nb][ic][f, :, :, :] = atleast_3d(
                         self.sim.get_dft_array(dgm, c, f))
@@ -290,7 +297,12 @@ def prepare_adjoint_run(self):
     def adjoint_run(self):
         # set up adjoint sources and monitors
         self.prepare_adjoint_run()
-        if self.sim.k_point: self.sim.k_point *= -1
+
+        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+            self.sim.m = -self.sim.m
+
+        if self.sim.k_point:
+            self.sim.k_point *= -1
         for ar in range(len(self.objective_functions)):
             # Reset the fields
             self.sim.reset_meep()
@@ -301,7 +313,7 @@ def adjoint_run(self):
             # register design flux
             self.design_region_monitors = [
                 self.sim.add_dft_fields(
-                    [mp.Ex, mp.Ey, mp.Ez],
+                    self.components,
                     self.frequencies,
                     where=dr.volume,
                     yee_grid=True,
@@ -328,10 +340,16 @@ def adjoint_run(self):
                 for c in dg
             ] for dg in self.design_grids])
             for nb, dgm in enumerate(self.design_region_monitors):
-                for ic, c in enumerate([mp.Ex, mp.Ey, mp.Ez]):
+                for ic, c in enumerate(self.components):
                     for f in range(self.nf):
-                        self.a_E[ar][nb][ic][f, :, :, :] = atleast_3d(
-                            self.sim.get_dft_array(dgm, c, f))
+                        if (self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL):
+                            # Addtional factor of 2 for cyldrical coordinate
+                            self.a_E[ar][nb][ic][f, :, :, :] = 2 * atleast_3d(self.sim.get_dft_array(dgm, c, f))
+                        else:
+                            self.a_E[ar][nb][ic][f, :, :, :] = atleast_3d(self.sim.get_dft_array(dgm, c, f))
+
+        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+            self.sim.m = -self.sim.m
 
         # update optimizer's state
         if self.sim.k_point: self.sim.k_point *= -1

python/meep.i

@@ -806,7 +806,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
 //--------------------------------------------------
 
 %inline %{
-void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObject *grid_volume, PyObject *frequencies, PyObject *py_geom_list, PyObject *f) {
+void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObject *grid_volume, PyObject *frequencies, PyObject *py_geom_list, PyObject *f, bool sim_is_cylindrical) {
     // clean the gradient array
     PyArrayObject *pao_grad = (PyArrayObject *)grad;
     if (!PyArray_Check(pao_grad)) meep::abort("grad parameter must be numpy array.");
@@ -860,7 +860,7 @@ void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObj
     meep::fields* f_c = (meep::fields *)f_v;
 
     // calculate the gradient
-    meep_geom::material_grids_addgradient(grad_c,ng,fields_a_c,fields_f_c,frequencies_c,nf,scalegrad,*where_vol,geometry_tree,f_c);
+    meep_geom::material_grids_addgradient(grad_c,ng,fields_a_c,fields_f_c,frequencies_c,nf,scalegrad,*where_vol,geometry_tree,f_c, sim_is_cylindrical);
 
     destroy_geom_box_tree(geometry_tree);
     delete l;

python/simulation.py

@@ -1295,6 +1295,8 @@ def use_2d(self, k):
                 return 2
             else:
                 return 3
+        elif self.dimensions == 2 and self.is_cylindrical:
+            return mp.CYLINDRICAL
         return self.dimensions
 
     def _get_valid_material_frequencies(self):

python/tests/test_adjoint_cyl.py

@@ -0,0 +1,151 @@
+import meep as mp
+try:
+    import meep.adjoint as mpa
+except:
+    import adjoint as mpa
+import numpy as np
+from autograd import numpy as npa
+from autograd import tensor_jacobian_product
+import unittest
+from enum import Enum
+from utils import ApproxComparisonTestCase
+
+np.random.seed(2)
+resolution = 20
+dimensions = mp.CYLINDRICAL
+m=0
+Si = mp.Medium(index=3.4)
+SiO2 = mp.Medium(index=1.44)
+
+sr = 6
+sz = 6
+cell_size = mp.Vector3(sr, 0, sz)
+dpml = 1.0
+boundary_layers = [mp.PML(thickness=dpml)]
+
+design_region_resolution = int(2*resolution)
+design_r = 4.8
+design_z = 2
+Nx = int(design_region_resolution*design_r)
+Nz = int(design_region_resolution*design_z)
+
+fcen = 1/1.55
+width = 0.2
+fwidth = width * fcen
+source_center  = [0.1+design_r/2,0,(sz/2-dpml+design_z/2)/2]
+source_size    = mp.Vector3(design_r,0,0)
+src = mp.GaussianSource(frequency=fcen,fwidth=fwidth)
+source = [mp.Source(src,component=mp.Er,
+                    center=source_center,
+                   size=source_size)]
+
+## random design region
+p = np.random.rand(Nx*Nz)
+## random epsilon perturbation for design region
+deps = 1e-5
+dp = deps*np.random.rand(Nx*Nz)
+
+
+def forward_simulation(design_params):
+    matgrid = mp.MaterialGrid(mp.Vector3(Nx,0,Nz),
+                              SiO2,
+                              Si,
+                              weights=design_params.reshape(Nx,1,Nz))
+
+    geometry = [mp.Block(center=mp.Vector3(0.1+design_r/2,0,0),
+                                 size=mp.Vector3(design_r,0,design_z),
+                                 material=matgrid)]
+
+    sim = mp.Simulation(resolution=resolution,
+                        cell_size=cell_size,
+                        boundary_layers=boundary_layers,
+                        sources=source,
+                        geometry=geometry,
+                        dimensions=dimensions,
+                        m=m)
+
+    frequencies = [fcen]
+    far_x = [mp.Vector3(5,0,20)]
+    mode = sim.add_near2far(
+        frequencies,
+        mp.Near2FarRegion(center=mp.Vector3(0.1+design_r/2,0 ,(sz/2-dpml+design_z/2)/2),size=mp.Vector3(design_r,0,0), weight=+1),
+        decimation_factor=10
+    )
+
+    sim.run(until_after_sources=1200)
+    Er = sim.get_farfield(mode, far_x[0])
+    sim.reset_meep()
+
+    return abs(Er[0])**2
+
+
+def adjoint_solver(design_params):
+
+    design_variables = mp.MaterialGrid(mp.Vector3(Nx,0,Nz),SiO2,Si)
+    design_region = mpa.DesignRegion(design_variables,volume=mp.Volume(center=mp.Vector3(0.1+design_r/2,0,0), size=mp.Vector3(design_r, 0,design_z)))
+    geometry = [mp.Block(center=design_region.center, size=design_region.size, material=design_variables)]
+
+    sim = mp.Simulation(cell_size=cell_size,
+        boundary_layers=boundary_layers,
+        geometry=geometry,
+        sources=source,
+        resolution=resolution,
+        dimensions=dimensions,
+        m=m)
+
+    far_x = [mp.Vector3(5,0,20)]
+    NearRegions = [mp.Near2FarRegion(center=mp.Vector3(0.1+design_r/2,0 ,(sz/2-dpml+design_z/2)/2),size=mp.Vector3(design_r,0,0), weight=+1)]
+    FarFields = mpa.Near2FarFields(sim, NearRegions ,far_x, decimation_factor=5)
+    ob_list = [FarFields]
+
+    def J(alpha):
+        return npa.abs(alpha[0,0,0])**2
+
+    opt = mpa.OptimizationProblem(
+        simulation=sim,
+        objective_functions=J,
+        objective_arguments=ob_list,
+        design_regions=[design_region],
+        fcen=fcen,
+        df = 0,
+        nf = 1,
+        decay_fields=[mp.Er],
+        decay_by=1e-4,
+        maximum_run_time=1200)
+
+    f, dJ_du = opt([design_params])
+    sim.reset_meep()
+
+    return f, dJ_du
+
+
+class TestAdjointSolver(ApproxComparisonTestCase):
+
+    def test_adjoint_solver_cyl_n2f_fields(self):
+        print("*** TESTING CYLINDRICAL Near2Far ADJOINT FEATURES ***")
+        adjsol_obj, adjsol_grad = adjoint_solver(p)
+
+        ## compute unperturbed S12
+        S12_unperturbed = forward_simulation(p)
+
+        ## compare objective results
+        print("|Er|^2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,S12_unperturbed))
+        self.assertClose(adjsol_obj,S12_unperturbed,epsilon=1e-3)
+
+        ## compute perturbed S12
+        S12_perturbed = forward_simulation(p+dp)
+
+        ## compare gradients
+        if adjsol_grad.ndim < 2:
+            adjsol_grad = np.expand_dims(adjsol_grad,axis=1)
+        adj_scale = (dp[None,:]@adjsol_grad).flatten()
+        fd_grad = S12_perturbed-S12_unperturbed
+        print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
+        tol = 0.1 if mp.is_single_precision() else 0.01
+        self.assertClose(adj_scale,fd_grad,epsilon=tol)
+
+
+
+
+if __name__ == '__main__':
+    unittest.main()

src/meepgeom.cpp

@@ -2579,7 +2579,9 @@ double get_material_gradient(
   if ((mm->E_susceptibilities.size() == 0) && mm->D_conductivity_diag.x == 0 &&
       mm->D_conductivity_diag.y == 0 && mm->D_conductivity_diag.z == 0){
         switch (field_dir){
+          case meep::Er:
           case meep::Ex: return (m2->epsilon_diag.x - m1->epsilon_diag.x) * (fields_a * fields_f).real();
+          case meep::Ep:
           case meep::Ey: return (m2->epsilon_diag.y - m1->epsilon_diag.y) * (fields_a * fields_f).real();
           case meep::Ez: return (m2->epsilon_diag.z - m1->epsilon_diag.z) * (fields_a * fields_f).real();
           default: meep::abort("Invalid field component specified in gradient.");
@@ -2602,9 +2604,9 @@ double get_material_gradient(
     dA_du[i] = (dA_du_1[i] - dA_du_0[i]) / (2 * du);
 
   int dir_idx = 0;
-  if (field_dir == meep::Ex)
+  if (field_dir == meep::Ex || field_dir == meep::Er)
     dir_idx = 0;
-  else if (field_dir == meep::Ey)
+  else if (field_dir == meep::Ey || field_dir == meep::Ep)
     dir_idx = 1;
   else if (field_dir == meep::Ez)
     dir_idx = 2;
@@ -2739,7 +2741,7 @@ void material_grids_addgradient_point(double *v, std::complex<double> fields_a,
 void material_grids_addgradient(double *v, size_t ng, std::complex<double> *fields_a,
                                 std::complex<double> *fields_f, double *frequencies,
                                 size_t nf, double scalegrad, const meep::volume &where,
-                                geom_box_tree geometry_tree, meep::fields *f) {
+                                geom_box_tree geometry_tree, meep::fields *f, bool sim_is_cylindrical) {
   int n2, n3, n4;
   double s[3][3], cen[3][3], c1, c2, c3, s1, s2, s3;
   vector3 p;
@@ -2748,6 +2750,10 @@ void material_grids_addgradient(double *v, size_t ng, std::complex<double> *fiel
   meep::direction dirs[3];
   meep::vec min_max_loc[2] = {meep::vec(0,0,0),meep::vec(0,0,0)}; // extremal points in subgrid
   meep::component field_dir[3] = {meep::Ex, meep::Ey, meep::Ez};
+  if (sim_is_cylindrical) {
+    field_dir[0] = meep::Er;
+    field_dir[1] = meep::Ep;
+  }
   size_t dims[9] = {1, 1, 1, 1, 1, 1, 1, 1, 1};
   for (int c = 0; c < 3; c++) {
 
@@ -2758,6 +2764,16 @@ void material_grids_addgradient(double *v, size_t ng, std::complex<double> *fiel
     double max_c_array[3] = {max_corner.x, max_corner.y, max_corner.z};
     vector3 min_corner = vec_to_vector3(min_max_loc[0]);
     double min_c_array[3] = {min_corner.x, min_corner.y, min_corner.z};
+
+    if (sim_is_cylindrical){
+      max_c_array[0] = max_corner.z;
+      max_c_array[1] = max_corner.x;
+      max_c_array[2] = max_corner.y;
+      min_c_array[0] = min_corner.z;
+      min_c_array[1] = min_corner.x;
+      min_c_array[2] = min_corner.y;
+    }
+
     for (int ci = 0; ci < 3; ci++) {
       s[c][ci] = (max_c_array[ci] - min_c_array[ci]) == 0 ? 0 : (max_c_array[ci] - min_c_array[ci]) / (dims[3 * c + ci] - 1);
       cen[c][ci] = dims[3 * c + ci] <= 1 ? 0 : min_c_array[ci];
@@ -2767,6 +2783,28 @@ void material_grids_addgradient(double *v, size_t ng, std::complex<double> *fiel
   // Loop over component, x, y, z, and frequency dimensions
   // TODO speed up with MPI (if needed)
   int xyz_index = 0;
+  if (sim_is_cylindrical){
+    for (int c = 0; c < 3; c++) {             // component
+      n2 = dims[c * 3]; n3 = dims[c * 3 + 1]; n4 = dims[c * 3 + 2];
+      c1 = cen[c][0]; c2 = cen[c][1]; c3 = cen[c][2];
+      s1 = s[c][0]; s2 = s[c][1]; s3 = s[c][2];
+
+      for (int i1 = 0; i1 < nf; ++i1) {       // freq
+        for (int i2 = 0; i2 < n2; ++i2) {     // z
+          for (int i4 = 0; i4 < n4; ++i4) {//y
+            for (int i3 = 0; i3 < n3; ++i3) {//x
+              p.z = i2 * s1 + c1; p.x = i3 * s2 + c2; p.y = i4 * s3 + c3;
+              material_grids_addgradient_point(v+ ng*i1, fields_a[xyz_index]*p.x, fields_f[xyz_index], field_dir[c], p,
+                                               scalegrad, frequencies[i1], geometry_tree);
+              //p.x is the (2pi)r' factor from integrating in cyldrical coordinate;
+              //2pi is canceled out by a overcouted factor of 2pi*r of the Near2FarRegion; See near2far.cpp
+              xyz_index++;
+            }
+          }
+        }
+      }
+    }
+  } else {
   for (int c = 0; c < 3; c++) {             // component
     n2 = dims[c * 3]; n3 = dims[c * 3 + 1]; n4 = dims[c * 3 + 2];
     c1 = cen[c][0]; c2 = cen[c][1]; c3 = cen[c][2];
@@ -2786,6 +2824,7 @@ void material_grids_addgradient(double *v, size_t ng, std::complex<double> *fiel
     }
   }
 }
+}
 
 static void find_array_min_max(int n, const double *data, double &min_val, double &max_val) {
   min_val = data[0];

src/meepgeom.hpp

@@ -228,7 +228,7 @@ void material_grids_addgradient_point(double *v, std::complex<double> fields_a,
 void material_grids_addgradient(double *v, size_t ng, std::complex<double> *fields_a,
                                 std::complex<double> *fields_f, double *frequencies,
                                 size_t nf, double scalegrad, const meep::volume &where,
-                                geom_box_tree geometry_tree, meep::fields *f);
+                                geom_box_tree geometry_tree, meep::fields *f, bool sim_is_cylindrical);
 
 /***************************************************************/
 /* routines in GDSIIgeom.cc ************************************/