Skip to content

Commit

Permalink
shorten unit test for adjoint solver (#1790)
Browse files Browse the repository at this point in the history 
* shorten unit test for adjoint solver

* revert small gap between mode monitor and PML

* add k-point test for complex fields

* add missing k_point to adjoint solver test

* use tighter tolerances for checks
  • Loading branch information
@oskooi
oskooi authored Oct 27, 2021
1 parent dd4662f commit e8c2023
Showing 1 changed file with 176 additions and 56 deletions.
232 changes: 176 additions & 56 deletions python/tests/test_adjoint_solver.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@
cell_size = mp.Vector3(sxy,sxy,0)

dpml = 1.0
boundary_layers = [mp.PML(thickness=dpml)]
pml_xy = [mp.PML(thickness=dpml)]
pml_x = [mp.PML(thickness=dpml,direction=mp.X)]

eig_parity = mp.EVEN_Y + mp.ODD_Z

Expand All @@ -46,14 +47,20 @@

fcen = 1/1.55
df = 0.23*fcen
sources = [mp.EigenModeSource(src=mp.GaussianSource(fcen,fwidth=df),
center=mp.Vector3(-0.5*sxy+dpml+0.1,0),
size=mp.Vector3(0,sxy-2*dpml),
eig_band=1,
eig_parity=eig_parity)]
wvg_source = [mp.EigenModeSource(src=mp.GaussianSource(fcen,fwidth=df),
center=mp.Vector3(-0.5*sxy+dpml+0.1,0),
size=mp.Vector3(0,sxy-2*dpml),
eig_parity=eig_parity)]

pt_source = [mp.Source(src=mp.GaussianSource(fcen,fwidth=df),
center=mp.Vector3(-0.5*sxy+dpml,0),
size=mp.Vector3(),
component=mp.Ez)]

def forward_simulation(design_params,mon_type, frequencies=None, use_complex=False, k=False):
k_point = mp.Vector3(0.23,-0.38)


def forward_simulation(design_params, mon_type, frequencies=None):
matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
mp.air,
silicon,
Expand All @@ -67,11 +74,9 @@ def forward_simulation(design_params,mon_type, frequencies=None, use_complex=Fal

sim = mp.Simulation(resolution=resolution,
cell_size=cell_size,
boundary_layers=boundary_layers,
sources=sources,
geometry=geometry,
force_complex_fields=use_complex,
k_point=k)
boundary_layers=pml_xy,
sources=wvg_source,
geometry=geometry)

if not frequencies:
frequencies = [fcen]
Expand All @@ -82,7 +87,6 @@ def forward_simulation(design_params,mon_type, frequencies=None, use_complex=Fal
size=mp.Vector3(0,sxy-2*dpml,0)),
yee_grid=True,
eig_parity=eig_parity)

elif mon_type.name == 'DFT':
mode = sim.add_dft_fields([mp.Ez],
frequencies,
Expand All @@ -95,7 +99,6 @@ def forward_simulation(design_params,mon_type, frequencies=None, use_complex=Fal
if mon_type.name == 'EIGENMODE':
coeff = sim.get_eigenmode_coefficients(mode,[1],eig_parity).alpha[0,:,0]
S12 = np.power(np.abs(coeff),2)

elif mon_type.name == 'DFT':
Ez2 = []
for f in range(len(frequencies)):
Expand All @@ -111,15 +114,17 @@ def forward_simulation(design_params,mon_type, frequencies=None, use_complex=Fal
return Ez2


def adjoint_solver(design_params, mon_type, frequencies=None, use_complex=False, k=False):
def adjoint_solver(design_params, mon_type, frequencies=None):
matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
mp.air,
silicon,
weights=np.ones((Nx,Ny)))

matgrid_region = mpa.DesignRegion(matgrid,
volume=mp.Volume(center=mp.Vector3(),
size=mp.Vector3(design_region_size.x,design_region_size.y,0)))
size=mp.Vector3(design_region_size.x,
design_region_size.y,
0)))

matgrid_geometry = [mp.Block(center=matgrid_region.center,
size=matgrid_region.size,
Expand All @@ -129,11 +134,9 @@ def adjoint_solver(design_params, mon_type, frequencies=None, use_complex=False,

sim = mp.Simulation(resolution=resolution,
cell_size=cell_size,
boundary_layers=boundary_layers,
sources=sources,
geometry=geometry,
force_complex_fields=use_complex,
k_point=k)
boundary_layers=pml_xy,
sources=wvg_source,
geometry=geometry)

if not frequencies:
frequencies = [fcen]
Expand All @@ -147,7 +150,6 @@ def adjoint_solver(design_params, mon_type, frequencies=None, use_complex=False,

def J(mode_mon):
return npa.power(npa.abs(mode_mon),2)

elif mon_type.name == 'DFT':
obj_list = [mpa.FourierFields(sim,
mp.Volume(center=mp.Vector3(1.25),
Expand All @@ -171,6 +173,95 @@ def J(mode_mon):
return f, dJ_du


def forward_simulation_complex_fields(design_params, frequencies=None):
matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
mp.air,
silicon,
weights=design_params.reshape(Nx,Ny))

geometry = [mp.Block(center=mp.Vector3(),
size=mp.Vector3(design_region_size.x,
design_region_size.y,
0),
material=matgrid)]

sim = mp.Simulation(resolution=resolution,
cell_size=cell_size,
k_point=k_point,
boundary_layers=pml_x,
sources=pt_source,
geometry=geometry)

if not frequencies:
frequencies = [fcen]

mode = sim.add_dft_fields([mp.Ez],
frequencies,
center=mp.Vector3(0.9),
size=mp.Vector3(0.2,0.5),
yee_grid=False)

sim.run(until_after_sources=mp.stop_when_dft_decayed())

Ez2 = []
for f in range(len(frequencies)):
Ez_dft = sim.get_dft_array(mode, mp.Ez, f)
Ez2.append(np.power(np.abs(Ez_dft[3,9]),2))
Ez2 = np.array(Ez2)

sim.reset_meep()

return Ez2


def adjoint_solver_complex_fields(design_params, frequencies=None):
matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
mp.air,
silicon,
weights=np.ones((Nx,Ny)))

matgrid_region = mpa.DesignRegion(matgrid,
volume=mp.Volume(center=mp.Vector3(),
size=mp.Vector3(design_region_size.x,
design_region_size.y,
0)))

geometry = [mp.Block(center=matgrid_region.center,
size=matgrid_region.size,
material=matgrid)]

sim = mp.Simulation(resolution=resolution,
cell_size=cell_size,
k_point=k_point,
boundary_layers=pml_x,
sources=pt_source,
geometry=geometry)

if not frequencies:
frequencies = [fcen]

obj_list = [mpa.FourierFields(sim,
mp.Volume(center=mp.Vector3(0.9),
size=mp.Vector3(0.2,0.5)),
mp.Ez)]

def J(dft_mon):
return npa.power(npa.abs(dft_mon[:,3,9]),2)

opt = mpa.OptimizationProblem(
simulation=sim,
objective_functions=J,
objective_arguments=obj_list,
design_regions=[matgrid_region],
frequencies=frequencies)

f, dJ_du = opt([design_params])

sim.reset_meep()

return f, dJ_du


def mapping(x,filter_radius,eta,beta):
filtered_field = mpa.conic_filter(x,
filter_radius,
Expand All @@ -186,62 +277,64 @@ def mapping(x,filter_radius,eta,beta):
class TestAdjointSolver(ApproxComparisonTestCase):

def test_adjoint_solver_DFT_fields(self):
print("*** TESTING DFT ADJOINT FEATURES ***")
print("*** TESTING DFT ADJOINT ***")

## test the single frequency and multi frequency cases
for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.DFT, frequencies)

## compute unperturbed S12
S12_unperturbed = forward_simulation(p, MonitorObject.DFT, frequencies)
## compute unperturbed |Ez|^2
Ez2_unperturbed = forward_simulation(p, MonitorObject.DFT, frequencies)

## compare objective results
print("|Ez|^2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,S12_unperturbed))
self.assertClose(adjsol_obj,S12_unperturbed,epsilon=1e-3)
print("|Ez|^2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,Ez2_unperturbed))
self.assertClose(adjsol_obj,Ez2_unperturbed,epsilon=1e-6)

## compute perturbed S12
S12_perturbed = forward_simulation(p+dp, MonitorObject.DFT, frequencies)
## compute perturbed Ez2
Ez2_perturbed = forward_simulation(p+dp, MonitorObject.DFT, frequencies)

## 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
fd_grad = Ez2_perturbed-Ez2_unperturbed
print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
tol = 0.0461 if mp.is_single_precision() else 0.005
self.assertClose(adj_scale,fd_grad,epsilon=tol)


def test_adjoint_solver_eigenmode(self):
print("*** TESTING EIGENMODE ADJOINT FEATURES ***")
print("*** TESTING EIGENMODE ADJOINT ***")

## test the single frequency and multi frequency case
for k in [False, mp.Vector3(), mp.Vector3(0.8, 0.6)]:
for use_complex in [True, False]:
for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.EIGENMODE, frequencies, use_complex, k)

## compute unperturbed S12
S12_unperturbed = forward_simulation(p, MonitorObject.EIGENMODE, frequencies, use_complex, k)

## compare objective results
print("S12 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,S12_unperturbed))
self.assertClose(adjsol_obj,S12_unperturbed,epsilon=1e-6)

## compute perturbed S12
S12_perturbed = forward_simulation(p+dp, MonitorObject.EIGENMODE, frequencies, use_complex, k)

## 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.04 if mp.is_single_precision() else 0.01
self.assertClose(adj_scale,fd_grad,epsilon=tol)
for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.EIGENMODE, frequencies)

## compute unperturbed S12
S12_unperturbed = forward_simulation(p, MonitorObject.EIGENMODE, frequencies)

## compare objective results
print("S12 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,S12_unperturbed))
self.assertClose(adjsol_obj,S12_unperturbed,epsilon=1e-6)

## compute perturbed S12
S12_perturbed = forward_simulation(p+dp, MonitorObject.EIGENMODE, frequencies)

## 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.04 if mp.is_single_precision() else 0.01
self.assertClose(adj_scale,fd_grad,epsilon=tol)


def test_gradient_backpropagation(self):
print("*** TESTING BACKPROP FEATURES ***")
print("*** TESTING BACKPROP ***")

for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## filter/thresholding parameters
filter_radius = 0.21985
Expand Down Expand Up @@ -280,5 +373,32 @@ def test_gradient_backpropagation(self):
self.assertClose(adj_scale,fd_grad,epsilon=tol)


def test_complex_fields(self):
print("*** TESTING COMPLEX FIELDS ***")

for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver_complex_fields(p, frequencies)

## compute unperturbed |Ez|^2
Ez2_unperturbed = forward_simulation_complex_fields(p, frequencies)

## compare objective results
print("Ez2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,Ez2_unperturbed))
self.assertClose(adjsol_obj,Ez2_unperturbed,epsilon=1e-8)

## compute perturbed |Ez|^2
Ez2_perturbed = forward_simulation_complex_fields(p+dp, frequencies)

## 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 = Ez2_perturbed-Ez2_unperturbed
print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
tol = 0.005 if mp.is_single_precision() else 0.0008
self.assertClose(adj_scale,fd_grad,epsilon=tol)


if __name__ == '__main__':
unittest.main()

0 comments on commit e8c2023

Please sign in to comment.