From e2739e4a44000441f58cd61c326f426d3faf63b4 Mon Sep 17 00:00:00 2001 From: mochen4 Date: Fri, 3 Sep 2021 09:55:54 -0400 Subject: [PATCH] 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 --- python/Makefile.am | 1 + python/adjoint/optimization_problem.py | 38 +++++-- python/meep.i | 4 +- python/simulation.py | 2 + python/tests/test_adjoint_cyl.py | 151 +++++++++++++++++++++++++ src/meepgeom.cpp | 45 +++++++- src/meepgeom.hpp | 2 +- src/near2far.cpp | 5 + 8 files changed, 232 insertions(+), 16 deletions(-) create mode 100644 python/tests/test_adjoint_cyl.py diff --git a/python/Makefile.am b/python/Makefile.am index 1870a70ef..6d8a63795 100644 --- a/python/Makefile.am +++ b/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 \ diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index bb938cf45..fe5873d52 100644 --- a/python/adjoint/optimization_problem.py +++ b/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 diff --git a/python/meep.i b/python/meep.i index 57ffc0003..53301a345 100644 --- a/python/meep.i +++ b/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; diff --git a/python/simulation.py b/python/simulation.py index aa7ad787f..893d1b8e8 100644 --- a/python/simulation.py +++ b/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): diff --git a/python/tests/test_adjoint_cyl.py b/python/tests/test_adjoint_cyl.py new file mode 100644 index 000000000..da3e7833a --- /dev/null +++ b/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() diff --git a/src/meepgeom.cpp b/src/meepgeom.cpp index e503e8137..4a913d534 100644 --- a/src/meepgeom.cpp +++ b/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 fields_a, void material_grids_addgradient(double *v, size_t ng, std::complex *fields_a, std::complex *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 *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 *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 *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 *fiel } } } +} static void find_array_min_max(int n, const double *data, double &min_val, double &max_val) { min_val = data[0]; diff --git a/src/meepgeom.hpp b/src/meepgeom.hpp index 990c537f1..a8558c076 100644 --- a/src/meepgeom.hpp +++ b/src/meepgeom.hpp @@ -228,7 +228,7 @@ void material_grids_addgradient_point(double *v, std::complex fields_a, void material_grids_addgradient(double *v, size_t ng, std::complex *fields_a, std::complex *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 ************************************/ diff --git a/src/near2far.cpp b/src/near2far.cpp index 943584bf0..7be72038b 100644 --- a/src/near2far.cpp +++ b/src/near2far.cpp @@ -696,6 +696,11 @@ std::vector dft_near2far::near_sourcedata(const vec &x_0, dou if (is_electric(temp_struct.near_fd_comp)) EH0 *= -1; EH0 /= f->S.multiplicity(ix0); + if (x_0.dim == Dcyl){ + EH0 /= xs.r(); + //Somehow, a factor of (2pi)r for r of the near2far region was double counted. + //2pi is canceled out by the integral over design region, where an extra factor of 2pi*r' (r' of the design region) is needed. See meepgeom.cpp + } temp_struct.amp_arr.push_back(EH0); } }