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

WIP:MaterialGrid #1242

Merged
merged 14 commits into from
Jul 4, 2020
1 change: 1 addition & 0 deletions python/Makefile.am
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ TESTS = \
$(KDOM_TEST) \
$(TEST_DIR)/ldos.py \
$(MDPYTEST) \
$(TEST_DIR)/material_grid.py \
$(TEST_DIR)/medium_evaluations.py \
$(MPBPYTEST) \
$(MODE_COEFFS_TEST) \
Expand Down
2 changes: 1 addition & 1 deletion python/adjoint/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

from .basis import BilinearInterpolationBasis

from .optimization_problem import (OptimizationProblem, Grid)
from .optimization_problem import (OptimizationProblem, Grid, DesignRegion)

from .filter_source import FilteredSource

Expand Down
12 changes: 7 additions & 5 deletions python/adjoint/objective.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,15 @@ def get_evaluation(self):
return

class EigenmodeCoefficient(ObjectiveQuantitiy):
def __init__(self,sim,volume,mode,forward=True,k0=None,**kwargs):
def __init__(self,sim,volume,mode,forward=True,kpoint_func=None,**kwargs):
'''
'''
self.sim = sim
self.volume=volume
self.mode=mode
self.forward = 0 if forward else 1
self.normal_direction = None
self.k0 = k0
self.kpoint_func = kpoint_func
self.eval = None
self.EigenMode_kwargs = kwargs
return
Expand All @@ -52,15 +52,15 @@ def place_adjoint_source(self,dJ,dt):
dJ = np.atleast_1d(dJ)
# determine starting kpoint for reverse mode eigenmode source
direction_scalar = 1 if self.forward else -1
if self.k0 is None:
if self.kpoint_func is None:
if self.normal_direction == 0:
k0 = direction_scalar * mp.Vector3(x=1)
elif self.normal_direction == 1:
k0 = direction_scalar * mp.Vector3(y=1)
elif self.normal_direction == 2:
k0 == direction_scalar * mp.Vector3(z=1)
else:
k0 = direction_scalar * self.k0
k0 = direction_scalar * self.kpoint_func(self.time_src.frequency,1)

# -------------------------------------- #
# Get scaling factor
Expand Down Expand Up @@ -94,6 +94,7 @@ def place_adjoint_source(self,dJ,dt):
direction=mp.NO_DIRECTION,
eig_kpoint=k0,
amplitude=amp,
eig_match_freq=True,
size=self.volume.size,
center=self.volume.center,
**self.EigenMode_kwargs)
Expand All @@ -105,7 +106,8 @@ def __call__(self):
self.time_src = self.sim.sources[0].src

# Eigenmode data
ob = self.sim.get_eigenmode_coefficients(self.monitor,[self.mode],**self.EigenMode_kwargs)
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
self.cscale = ob.cscale # pull scaling factor

Expand Down
64 changes: 41 additions & 23 deletions python/adjoint/optimization_problem.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,28 @@

Grid = namedtuple('Grid', ['x', 'y', 'z', 'w'])

class DesignRegion(object):
def __init__(self,design_parameters,volume=None, size=None, center=mp.Vector3()):
self.volume = volume if volume else mp.Volume(center=center,size=size)
self.size=self.volume.size
self.center=self.volume.center
self.design_parameters=design_parameters
self.num_design_params=design_parameters.num_params
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)")
num_freqs = np.array(frequencies).size

grad = np.zeros((self.num_design_params*num_freqs,)) # preallocate

# compute the gradient
mp._get_gradient(grad,fields_a,fields_f,self.volume,np.array(frequencies),geom_list,f)

return np.squeeze(grad.reshape(self.num_design_params,num_freqs,order='F'))

class OptimizationProblem(object):
"""Top-level class in the MEEP adjoint module.

Expand All @@ -24,7 +46,7 @@ def __init__(self,
simulation,
objective_functions,
objective_arguments,
design_variables,
design_regions,
frequencies=None,
fcen=None,
df=None,
Expand All @@ -44,13 +66,12 @@ def __init__(self,
self.objective_arguments = objective_arguments
self.f_bank = [] # objective function evaluation history

if isinstance(design_variables, list):
self.design_variables = design_variables
if isinstance(design_regions, list):
self.design_regions = design_regions
else:
self.design_variables = [design_variables]
self.design_regions = [design_regions]

self.num_design_params = [ni.num_design_params for ni in self.design_variables]
self.design_regions = [dr.volume for dr in self.design_variables]
self.num_design_params = [ni.num_design_params for ni in self.design_regions]
self.num_design_regions = len(self.design_regions)

# TODO typecheck frequency choices
Expand Down Expand Up @@ -119,11 +140,8 @@ def __call__(self, rho_vector=None, need_value=True, need_gradient=True):
self.calculate_gradient()
else:
raise ValueError("Incorrect solver state detected: {}".format(self.current_state))

# clean up design grid list object
dg = self.design_grids[0] if len(self.design_grids) == 1 else self.design_grids

return self.f0, self.gradient, dg
return self.f0, self.gradient

def get_fdf_funcs(self):
"""construct callable functions for objective function value and gradient
Expand Down Expand Up @@ -158,7 +176,7 @@ 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,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=False) 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]
Expand All @@ -181,11 +199,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),3,self.nf),dtype=np.complex128) for dg in self.design_grids]
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]
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][:,:,:,ic,f] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
self.d_E[nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))

# store objective function evaluation in memory
self.f_bank.append(self.f0)
Expand All @@ -209,24 +227,24 @@ 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,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=False) 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))

# 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),3,self.nf),dtype=np.complex128) for dg in self.design_grids])
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])
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][:,:,:,ic,f] = atleast_3d(self.sim.get_dft_array(dgm,c,f))
self.a_E[objective_idx][nb][:,:,:,f,ic] = atleast_3d(self.sim.get_dft_array(dgm,c,f))

# update optimizer's state
self.current_state = "ADJ"

def calculate_gradient(self):
# Iterate through all design region bases and store gradient w.r.t. permittivity
self.gradient = [[2*np.sum(np.real(self.a_E[ar][nb]*self.d_E[nb]),axis=(3)) for nb in range(self.num_design_regions)] for ar in range(len(self.objective_functions))]
# Iterate through all design regions and calculate gradient
self.gradient = [[dr.get_gradient(self.a_E[ar][dri],self.d_E[dri],self.frequencies,self.sim.geometry,self.sim.fields) for dri, dr in enumerate(self.design_regions)] for ar in range(len(self.objective_functions))]

# Cleanup list of lists
if len(self.gradient) == 1:
Expand Down Expand Up @@ -276,15 +294,15 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi
for k in fd_gradient_idx:

b0 = np.ones((self.num_design_params[design_variables_idx],))
b0[:] = (self.design_variables[design_variables_idx].rho_vector)
b0[:] = (self.design_regions[design_variables_idx].design_parameters.design_parameters)
# -------------------------------------------- #
# left function evaluation
# -------------------------------------------- #
self.sim.reset_meep()

# assign new design vector
b0[k] -= db
self.design_variables[design_variables_idx].set_rho_vector(b0)
self.design_regions[design_variables_idx].update_design_parameters(b0)

# initialize design monitors
self.forward_monitors = []
Expand All @@ -306,7 +324,7 @@ def calculate_fd_gradient(self,num_gradients=1,db=1e-4,design_variables_idx=0,fi

# assign new design vector
b0[k] += 2*db # central difference rule...
self.design_variables[design_variables_idx].set_rho_vector(b0)
self.design_regions[design_variables_idx].update_design_parameters(b0)

# initialize design monitors
self.forward_monitors = []
Expand Down Expand Up @@ -338,10 +356,10 @@ def update_design(self, rho_vector):

rho_vector ....... a list of numpy arrays that maps to each design region
"""
for bi, b in enumerate(self.design_variables):
for bi, b in enumerate(self.design_regions):
if np.array(rho_vector[bi]).ndim > 1:
raise ValueError("Each vector of design variables must contain only one dimension.")
b.set_rho_vector(rho_vector[bi])
b.update_design_parameters(rho_vector[bi])

self.sim.reset_meep()
self.current_state = "INIT"
Expand Down
91 changes: 63 additions & 28 deletions python/examples/adjoint_optimization/01-Introduction.ipynb

Large diffs are not rendered by default.

Loading