Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support adjoint gradients in cylindrical coordinate #1749

Merged
merged 25 commits into from
Sep 3, 2021
1 change: 1 addition & 0 deletions python/Makefile.am
Original file line number Diff line number Diff line change
Expand Up @@ -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 \
Expand Down
38 changes: 28 additions & 10 deletions python/adjoint/optimization_problem.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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)]

Expand Down Expand Up @@ -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))
Expand All @@ -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()
Expand All @@ -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,
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions python/meep.i
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand Down Expand Up @@ -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;
Expand Down
2 changes: 2 additions & 0 deletions python/simulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -1289,6 +1289,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):
Expand Down
151 changes: 151 additions & 0 deletions python/tests/test_adjoint_cyl.py
Original file line number Diff line number Diff line change
@@ -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()
45 changes: 42 additions & 3 deletions src/meepgeom.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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++) {

Expand All @@ -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];
Expand All @@ -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];
Expand All @@ -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];
Expand Down
2 changes: 1 addition & 1 deletion src/meepgeom.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 ************************************/
Expand Down
Loading