Skip to content

Commit

Permalink
fix gradients
Browse files Browse the repository at this point in the history
  • Loading branch information
@smartalecH
smartalecH committed Jul 9, 2020
1 parent 101e7f4 commit 37ac6e1
Show file tree
Hide file tree
Showing 7 changed files with 114 additions and 106 deletions.
5 changes: 3 additions & 2 deletions python/adjoint/objective.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,9 @@ def place_adjoint_source(self,dJ,dt):
dV = 1/self.sim.resolution * 1/self.sim.resolution
else:
dV = 1/self.sim.resolution * 1/self.sim.resolution * 1/self.sim.resolution
da_dE = 0.5*(dV * self.cscale)
scale = da_dE * dJ * 1j * 2 * np.pi * self.frequencies / np.array([self.time_src.fourier_transform(f) for f in self.frequencies]) # final scale factor
da_dE = 0.5 * self.cscale
iomega = ((1.0 - np.exp(-1j * (2 * np.pi * self.frequencies) * dt)) * (1.0 / dt))
scale = da_dE * dV * dJ * iomega / np.array([self.time_src.fourier_transform(f) for f in self.frequencies]) # final scale factor
if self.frequencies.size == 1:
# Single frequency simulations. We need to drive it with a time profile.
src = self.time_src
Expand Down
69 changes: 37 additions & 32 deletions python/adjoint/optimization_problem.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from collections import namedtuple

Grid = namedtuple('Grid', ['x', 'y', 'z', 'w'])
YeeDims = namedtuple('YeeDims', ['Ex','Ey','Ez'])

class DesignRegion(object):
def __init__(self,design_parameters,volume=None, size=None, center=mp.Vector3()):
Expand All @@ -16,9 +17,11 @@ def __init__(self,design_parameters,volume=None, size=None, center=mp.Vector3())
def update_design_parameters(self,design_parameters):
self.design_parameters.update_parameters(design_parameters)
def get_gradient(self,fields_a,fields_f,frequencies,geom_list,f):
# sanitize the input
if (fields_a.ndim < 5) or (fields_f.ndim < 5):
raise ValueError("Fields arrays must have 5 dimensions (x,y,z,frequency,component)")
for c in range(3):
fields_a[c] = fields_a[c].flatten(order='C')
fields_f[c] = fields_f[c].flatten(order='C')
fields_a = np.concatenate(fields_a)
fields_f = np.concatenate(fields_f)
num_freqs = np.array(frequencies).size

grad = np.zeros((self.num_design_params*num_freqs,)) # preallocate
Expand Down Expand Up @@ -176,17 +179,20 @@ def prepare_forward_run(self):
self.forward_monitors.append(m.register_monitors(self.frequencies))

# register design region
self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=False) for dr in self.design_regions]
self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=True) for dr in self.design_regions]

# store design region voxel parameters
self.design_grids = [Grid(*self.sim.get_array_metadata(dft_cell=drm)) for drm in self.design_region_monitors]
self.design_grids = []
for drm in self.design_region_monitors:
s = [self.sim.get_array_slice_dimensions(c,vol=drm.where) for c in [mp.Ex,mp.Ey,mp.Ez]]
self.design_grids += [YeeDims(*s)]

def forward_run(self):
# set up monitors
self.prepare_forward_run()

# Forward run
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, False, self.minimum_run_time))

# record objective quantities from user specified monitors
self.results_list = []
Expand All @@ -199,11 +205,11 @@ def forward_run(self):
self.f0 = self.f0[0]

# Store forward fields for each set of design variables in array (x,y,z,field_components,frequencies)
self.d_E = [np.zeros((len(dg.x),len(dg.y),len(dg.z),self.nf,3),dtype=np.complex128) for dg in self.design_grids]
self.d_E = [[np.zeros((c[0],c[1],c[2],self.nf),dtype=np.complex128) for c in dg] for dg in self.design_grids]
for nb, dgm in enumerate(self.design_region_monitors):
for f in range(self.nf):
for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
self.d_E[nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
for f in range(self.nf):
self.d_E[nb][ic][:,:,:,f] = atleast_3d(self.sim.get_dft_array(dgm,c,f))

# store objective function evaluation in memory
self.f_bank.append(self.f0)
Expand All @@ -227,17 +233,17 @@ def adjoint_run(self,objective_idx=0):

# register design flux
# TODO use yee grid directly
self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=False) for dr in self.design_regions]
self.design_region_monitors = [self.sim.add_dft_fields([mp.Ex,mp.Ey,mp.Ez],self.frequencies,where=dr.volume,yee_grid=True) for dr in self.design_regions]

# Adjoint run
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.design_region_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.design_region_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, True, self.minimum_run_time))

# Store adjoint fields for each design set of design variables in array (x,y,z,field_components,frequencies)
self.a_E.append([np.zeros((len(dg.x),len(dg.y),len(dg.z),self.nf,3),dtype=np.complex128) for dg in self.design_grids])
self.a_E.append([[np.zeros((c[0],c[1],c[2],self.nf),dtype=np.complex128) for c in dg] for dg in self.design_grids])
for nb, dgm in enumerate(self.design_region_monitors):
for f in range(self.nf):
for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
self.a_E[objective_idx][nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
for ic, c in enumerate([mp.Ex,mp.Ey,mp.Ez]):
for f in range(self.nf):
self.a_E[objective_idx][nb][ic][:,:,:,f] = atleast_3d(self.sim.get_dft_array(dgm,c,f))

# update optimizer's state
self.current_state = "ADJ"
Expand Down Expand Up @@ -309,7 +315,7 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
for m in self.objective_arguments:
self.forward_monitors.append(m.register_monitors(self.frequencies))

self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, False, self.minimum_run_time))

# record final objective function value
results_list = []
Expand All @@ -332,7 +338,7 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
self.forward_monitors.append(m.register_monitors(self.frequencies))

# add monitor used to track dft convergence
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, self.minimum_run_time))
self.sim.run(until_after_sources=stop_when_dft_decayed(self.sim, self.forward_monitors, self.decay_dt, self.decay_fields, self.fcen_idx, self.decay_by, False, self.minimum_run_time))

# record final objective function value
results_list = []
Expand Down Expand Up @@ -380,7 +386,7 @@ def plot2D(self,init_opt=False, **kwargs):

self.sim.plot2D(**kwargs)

def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, minimum_run_time=0):
def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, yee_grid=False, minimum_run_time=0):
'''Step function that monitors the relative change in DFT fields for a list of monitors.
mon ............. a list of monitors
Expand All @@ -389,18 +395,17 @@ def stop_when_dft_decayed(simob, mon, dt, c, fcen_idx, decay_by, minimum_run_tim
'''

# get monitor dft output array dimensions
x = []
y = []
z = []
dims = []
for m in mon:
xi,yi,zi,wi = simob.get_array_metadata(dft_cell=m)
x.append(len(xi))
y.append(len(yi))
z.append(len(zi))
ci_dims = []
for ci in c:
comp = ci if yee_grid else mp.Dielectric
ci_dims += [simob.get_array_slice_dimensions(comp,vol=m.where)]
dims.append(ci_dims)

# Record data in closure so that we can persitently edit
closure = {
'previous_fields': [np.ones((x[mi],y[mi],z[mi],len(c)),dtype=np.complex128) for mi, m in enumerate(mon)],
'previous_fields': [[np.ones(di,dtype=np.complex128) for di in d] for d in dims],
't0': 0
}

Expand All @@ -413,17 +418,17 @@ def _stop(sim):

# Pull all the relevant frequency and spatial dft points
relative_change = []
current_fields = [np.zeros((x[mi],y[mi],z[mi],len(c)), dtype=np.complex128) for mi in range(len(mon))]
current_fields = [[np.zeros(di,dtype=np.complex128) for di in d] for d in dims]
for mi, m in enumerate(mon):
for ic, cc in enumerate(c):
if isinstance(m,mp.DftFlux):
current_fields[mi][:,:,:,ic] = atleast_3d(mp.get_fluxes(m)[fcen_idx])
current_fields[mi][ic][:,:,:] = atleast_3d(mp.get_fluxes(m)[fcen_idx])
elif isinstance(m,mp.DftFields):
current_fields[mi][:,:,:,ic] = atleast_3d(sim.get_dft_array(m, cc, fcen_idx))
current_fields[mi][ic][:,:,:] = atleast_3d(sim.get_dft_array(m, cc, fcen_idx))
else:
raise TypeError("Monitor of type {} not supported".format(type(m)))
relative_change_raw = np.abs(previous_fields[mi] - current_fields[mi]) / np.abs(previous_fields[mi])
relative_change.append(np.mean(relative_change_raw.flatten())) # average across space and frequency
relative_change_raw = np.abs(previous_fields[mi][ic] - current_fields[mi][ic]) / np.abs(previous_fields[mi][ic])
relative_change.append(np.mean(relative_change_raw.flatten())) # average across space and frequency
relative_change = np.mean(relative_change) # average across monitors
closure['previous_fields'] = current_fields
closure['t0'] = sim.round_time()
Expand Down
Loading

0 comments on commit 37ac6e1

Please sign in to comment.