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.

Sign up for GitHub

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

Adjoint phase adjustments #1515

Merged
merged 9 commits into from
Mar 3, 2021
Merged

Adjoint phase adjustments #1515

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
5a1ec38
Fix dft fields memory leak
smartalecH Jun 26, 2020
0bcc35c
zero dim fix
smartalecH Jun 26, 2020
d9fe9ed
more cleanup
smartalecH Jun 26, 2020
4170d35
Merge branch 'master' of https://github.com/smartalecH/meep
smartalecH Jul 6, 2020
59e04fd
rebase
smartalecH Jul 9, 2020
4378c79
init
Mar 1, 2021
bb257c7
fix tests
Mar 2, 2021
0be340d
Merge branch 'master' of https://github.com/smartalecH/meep
Mar 2, 2021
32663d8
revert
Mar 2, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 46 additions & 12 deletions python/adjoint/objective.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ def place_adjoint_source(self,dJ):
if self.frequencies.size == 1:
# Single frequency simulations. We need to drive it with a time profile.
amp = da_dE * dJ * scale # final scale factor
src = self.time_src
src=self.time_src
else:
# multi frequency simulations
scale = da_dE * dJ * scale
Expand All @@ -92,10 +92,8 @@ def place_adjoint_source(self,dJ):
return self.source

def __call__(self):
# We just need a workable time profile, so just grab the first available time profile and use that.
self.time_src = self.sim.sources[0].src

# Eigenmode data
self.time_src = create_time_profile(self)
direction = mp.NO_DIRECTION if self.kpoint_func else mp.AUTOMATIC
ob = self.sim.get_eigenmode_coefficients(self.monitor,[self.mode],direction=direction,kpoint_func=self.kpoint_func,**self.EigenMode_kwargs)
self.eval = np.squeeze(ob.alpha[:,:,self.forward]) # record eigenmode coefficients for scaling
Expand Down Expand Up @@ -144,23 +142,39 @@ def place_adjoint_source(self,dJ):
self.sources += [mp.Source(src, component=self.component, amplitude=amp[xi, yi, zi],
center=mp.Vector3(self.dg.x[xi], self.dg.y[yi], self.dg.z[zi]))]
else:
'''The adjoint solver requires the objective function
to be scalar valued with regard to objective arguments
and position, but the function may be vector valued
with regard to frequency. In this case, the Jacobian
will be of the form [F,F,...] where F is the number of
frequencies. Because of linearity, we can sum across the
second frequency dimension to calculate a frequency
scale factor for each point (rather than a scale vector).
'''
dJ = np.sum(dJ,axis=1) # sum along first dimension bc Jacobian is always diag
dJ_4d = np.array([dJ[f].copy().reshape(x_dim, y_dim, z_dim) for f in range(self.num_freq)])
if self.component in [mp.Hx, mp.Hy, mp.Hz]:
dJ_4d = -dJ_4d
for zi in range(z_dim):
for yi in range(y_dim):
for xi in range(x_dim):
final_scale = -dJ_4d[:,xi,yi,zi] * scale
src = FilteredSource(self.time_src.frequency,self.frequencies,final_scale,dt)
self.sources += [mp.Source(src, component=self.component, amplitude=1,
center=mp.Vector3(self.dg.x[xi], self.dg.y[yi], self.dg.z[zi]))]
'''We only need to add a current source if the
jacobian is nonzero for all frequencies at
that particular point. Otherwise, the fitting
algorithm is going to fail.
'''
if not np.all((dJ_4d[:,xi,yi,zi] == 0)):
final_scale = -dJ_4d[:,xi,yi,zi] * scale
src = FilteredSource(self.time_src.frequency,self.frequencies,final_scale,dt)
self.sources += [mp.Source(src, component=self.component, amplitude=1,
center=mp.Vector3(self.dg.x[xi], self.dg.y[yi], self.dg.z[zi]))]

return self.sources

def __call__(self):
self.time_src = self.sim.sources[0].src
self.dg = Grid(*self.sim.get_array_metadata(dft_cell=self.monitor))
self.eval = np.array([self.sim.get_dft_array(self.monitor, self.component, i) for i in range(self.num_freq)])
self.time_src = create_time_profile(self)
return self.eval

def get_evaluation(self):
Expand Down Expand Up @@ -212,7 +226,7 @@ def place_adjoint_source(self,dJ):
return self.sources

def __call__(self):
self.time_src = self.sim.sources[0].src
self.time_src = create_time_profile(self)
self.eval = np.array([self.sim.get_farfield(self.monitor, far_pt) for far_pt in self.far_pts]).reshape((self.nfar_pts, self.num_freq, 6))
return self.eval

Expand All @@ -222,6 +236,19 @@ def get_evaluation(self):
except AttributeError:
raise RuntimeError("You must first run a forward simulation.")

def create_time_profile(obj_quantity,fwidth_frac=0.1):
'''
For single frequency objective functions, we should
generate a guassian pulse with a reasonable bandwidth
centered at said frequency.

TODO:
The user may specify a scalar valued objective
function across multiple frequencies (e.g. MSE)
in which case we should check that all the frequencies
fit in the specified bandwidth.
'''
return mp.GaussianSource(np.mean(obj_quantity.frequencies),fwidth=fwidth_frac*np.mean(obj_quantity.frequencies))

def adj_src_scale(obj_quantity, dt, include_resolution=True):
# -------------------------------------- #
Expand All @@ -246,8 +273,15 @@ def adj_src_scale(obj_quantity, dt, include_resolution=True):
# an ugly way to calcuate the scaled dtft of the forward source
y = np.array([src.swigobj.current(t,dt) for t in np.arange(0,T,dt)]) # time domain signal
fwd_dtft = np.matmul(np.exp(1j*2*np.pi*obj_quantity.frequencies[:,np.newaxis]*np.arange(y.size)*dt), y)*dt/np.sqrt(2*np.pi) # dtft

# we need to compensate for the phase added by the time envelope at our freq of interest

# Interestingly, the real parts of the DTFT and fourier transform match, but the imaginary parts are very different...
#fwd_dtft = src.fourier_transform(src.frequency)

'''
For some reason, there seems to be an additional phase
factor at the center frequency that needs to be applied
to *all* frequencies...
'''
src_center_dtft = np.matmul(np.exp(1j*2*np.pi*np.array([src.frequency])[:,np.newaxis]*np.arange(y.size)*dt), y)*dt/np.sqrt(2*np.pi)
adj_src_phase = np.exp(1j*np.angle(src_center_dtft))

Expand Down
172 changes: 103 additions & 69 deletions python/tests/adjoint_solver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from autograd import tensor_jacobian_product
import unittest
from enum import Enum
mp.quiet(True)

MonitorObject = Enum('MonitorObject', 'EIGENMODE DFT')

Expand Down Expand Up @@ -44,15 +45,15 @@
size=mp.Vector3(mp.inf,w,mp.inf))]

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


def forward_simulation(design_params,mon_type):
def forward_simulation(design_params,mon_type,frequencies=None):
matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
mp.air,
silicon,
Expand All @@ -70,32 +71,33 @@ def forward_simulation(design_params,mon_type):
boundary_layers=boundary_layers,
sources=sources,
geometry=geometry)
if not frequencies:
frequencies = [fcen]

if mon_type.name == 'EIGENMODE':
mode = sim.add_mode_monitor(fcen,
0,
1,
mode = sim.add_mode_monitor(frequencies,
mp.ModeRegion(center=mp.Vector3(0.5*sxy-dpml),size=mp.Vector3(0,sxy,0)),
yee_grid=True)

elif mon_type.name == 'DFT':
mode = sim.add_dft_fields([mp.Ez],
fcen,
0,
1,
center=mp.Vector3(0.5*sxy-dpml),
size=mp.Vector3(0,sxy),
frequencies,
center=mp.Vector3(1.25),
size=mp.Vector3(0.25,1,0),
yee_grid=False)

sim.run(until_after_sources=20)
sim.run(until_after_sources=50)

if mon_type.name == 'EIGENMODE':
coeff = sim.get_eigenmode_coefficients(mode,[1],eig_parity).alpha[0,0,0]
coeff = sim.get_eigenmode_coefficients(mode,[1],eig_parity).alpha[0,:,0]
S12 = abs(coeff)**2

elif mon_type.name == 'DFT':
Ez_dft = sim.get_dft_array(mode, mp.Ez, 0)
Ez2 = abs(Ez_dft[63])**2
Ez2 = []
for f in range(len(frequencies)):
Ez_dft = sim.get_dft_array(mode, mp.Ez, f)
Ez2.append(abs(Ez_dft[4,10])**2)
Ez2 = np.array(Ez2)

sim.reset_meep()

Expand All @@ -105,7 +107,7 @@ def forward_simulation(design_params,mon_type):
return Ez2


def adjoint_solver(design_params, mon_type):
def adjoint_solver(design_params, mon_type, frequencies=None):
matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
mp.air,
silicon,
Expand All @@ -126,6 +128,8 @@ def adjoint_solver(design_params, mon_type):
boundary_layers=boundary_layers,
sources=sources,
geometry=geometry)
if not frequencies:
frequencies = [fcen]

if mon_type.name == 'EIGENMODE':
obj_list = [mpa.EigenmodeCoefficient(sim,
Expand All @@ -137,19 +141,19 @@ def J(mode_mon):

elif mon_type.name == 'DFT':
obj_list = [mpa.FourierFields(sim,
mp.Volume(center=mp.Vector3(0.5*sxy-dpml),
size=mp.Vector3(0,sxy,0)),
mp.Volume(center=mp.Vector3(1.25),
size=mp.Vector3(0.25,1,0)),
mp.Ez)]

def J(mode_mon):
return npa.abs(mode_mon[0,63])**2
return npa.abs(mode_mon[:,4,10])**2

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

f, dJ_du = opt([design_params])
Expand All @@ -170,64 +174,94 @@ def mapping(x,filter_radius,eta,beta):
class TestAdjointSolver(unittest.TestCase):

def test_adjoint_solver_DFT_fields(self):
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.DFT)

## compute unperturbed Ez2
Ez2_unperturbed = forward_simulation(p, MonitorObject.DFT)

print("Ez2:, {:.6f}, {:.6f}".format(adjsol_obj,Ez2_unperturbed))
self.assertAlmostEqual(adjsol_obj,Ez2_unperturbed,places=2)

## compute perturbed Ez2
Ez2_perturbed = forward_simulation(p+dp, MonitorObject.DFT)
print("*** TESTING DFT ADJOINT FEATURES ***")
## 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)

## compare objective results
print("|Ez|^2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,S12_unperturbed))
np.testing.assert_array_almost_equal(adjsol_obj,S12_unperturbed,decimal=3)

print("directional_derivative:, {:.6f}, {:.6f}".format(np.dot(dp,adjsol_grad),Ez2_perturbed-Ez2_unperturbed))
self.assertAlmostEqual(np.dot(dp,adjsol_grad),Ez2_perturbed-Ez2_unperturbed,places=5)
## compute perturbed S12
S12_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
print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
np.testing.assert_array_almost_equal(adj_scale,fd_grad,decimal=5)


def test_adjoint_solver_eigenmode(self):
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.EIGENMODE)

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

print("S12:, {:.6f}, {:.6f}".format(adjsol_obj,S12_unperturbed))
self.assertAlmostEqual(adjsol_obj,S12_unperturbed,places=3)

## compute perturbed S12
S12_perturbed = forward_simulation(p+dp, MonitorObject.EIGENMODE)
print("*** TESTING EIGENMODE ADJOINT FEATURES ***")
## 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.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))
np.testing.assert_array_almost_equal(adjsol_obj,S12_unperturbed,decimal=3)

print("directional_derivative:, {:.6f}, {:.6f}".format(np.dot(dp,adjsol_grad),S12_perturbed-S12_unperturbed))
self.assertAlmostEqual(np.dot(dp,adjsol_grad),S12_perturbed-S12_unperturbed,places=5)
## 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))
np.testing.assert_array_almost_equal(adj_scale,fd_grad,decimal=5)


def test_gradient_backpropagation(self):
## filter/thresholding parameters
filter_radius = 0.21985
eta = 0.49093
beta = 4.0698

mapped_p = mapping(p,filter_radius,eta,beta)

## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(mapped_p, MonitorObject.EIGENMODE)

## backpropagate the gradient
bp_adjsol_grad = tensor_jacobian_product(mapping,0)(p,filter_radius,eta,beta,adjsol_grad)

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

print("S12:, {:.6f}, {:.6f}".format(adjsol_obj,S12_unperturbed))
self.assertAlmostEqual(adjsol_obj,S12_unperturbed,places=3)

## compute perturbed S12
S12_perturbed = forward_simulation(mapping(p+dp,filter_radius,eta,beta), MonitorObject.EIGENMODE)

print("directional_derivative:, {:.6f}, {:.6f}".format(np.dot(dp,bp_adjsol_grad),S12_perturbed-S12_unperturbed))
self.assertAlmostEqual(np.dot(dp,bp_adjsol_grad),S12_perturbed-S12_unperturbed,places=5)
print("*** TESTING BACKPROP FEATURES ***")
for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## filter/thresholding parameters
filter_radius = 0.21985
eta = 0.49093
beta = 4.0698

mapped_p = mapping(p,filter_radius,eta,beta)

## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(mapped_p, MonitorObject.EIGENMODE,frequencies)

## backpropagate the gradient
if len(frequencies) > 1:
bp_adjsol_grad = np.zeros(adjsol_grad.shape)
for i in range(len(frequencies)):
bp_adjsol_grad[:,i] = tensor_jacobian_product(mapping,0)(p,filter_radius,eta,beta,adjsol_grad[:,i])
else:
bp_adjsol_grad = tensor_jacobian_product(mapping,0)(p,filter_radius,eta,beta,adjsol_grad)

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

## compare objective results
print("S12 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,S12_unperturbed))
np.testing.assert_array_almost_equal(adjsol_obj,S12_unperturbed,decimal=3)

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

if bp_adjsol_grad.ndim < 2:
bp_adjsol_grad = np.expand_dims(bp_adjsol_grad,axis=1)
adj_scale = (dp[None,:]@bp_adjsol_grad).flatten()
fd_grad = S12_perturbed-S12_unperturbed
print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
np.testing.assert_array_almost_equal(adj_scale,fd_grad,decimal=5)


if __name__ == '__main__':
Expand Down