Skip to content

Commit

Permalink
Anisotropic material gradient support (#1801)
Browse files Browse the repository at this point in the history 
* init

* get everything compiled and default passing tests

* new start

* anisotropic gradients are working

* fix multifrequency bug and other cleanup

* add support for cyl coordindates

* add rotated sapphire to adjoint solver test

* get tests passing locally

* enable dispersive material support

* some minor cleanup

* clean up comments and readability issues and fix scaling of smoothing diameter

Co-authored-by: Alec Hammond <[email protected]>
  • Loading branch information
@smartalecH
smartalecH and Alec Hammond authored Nov 3, 2021
1 parent e338e14 commit 8dd39f7
Show file tree
Hide file tree
Showing 14 changed files with 522 additions and 423 deletions.
7 changes: 4 additions & 3 deletions python/adjoint/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,18 +5,19 @@

from .objective import *

from . import utils
from .utils import DesignRegion

from .basis import BilinearInterpolationBasis

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

from .filter_source import FilteredSource

from .filters import *

from .connectivity import *

from . import utils

try:
from .wrapper import MeepJaxWrapper
except ModuleNotFoundError as _:
Expand Down
3 changes: 2 additions & 1 deletion python/adjoint/objective.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@
import numpy as np
import meep as mp
from .filter_source import FilteredSource
from .optimization_problem import Grid
from meep.simulation import py_v3_to_vec
from collections import namedtuple

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

class ObjectiveQuantity(abc.ABC):
"""A differentiable objective quantity.
Expand Down
136 changes: 28 additions & 108 deletions python/adjoint/optimization_problem.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,48 +3,7 @@
from autograd import grad, jacobian
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(),
MaterialGrid=None,
):
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
self.MaterialGrid = MaterialGrid

def update_design_parameters(self, design_parameters):
self.design_parameters.update_weights(design_parameters)

def get_gradient(self, sim, fields_a, fields_f, frequencies):
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((num_freqs, self.num_design_params)) # preallocate

geom_list = sim.geometry
f = sim.fields
vol = sim._fit_volume_to_simulation(self.volume)
# compute the gradient
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

from . import utils, DesignRegion

class OptimizationProblem(object):
"""Top-level class in the MEEP adjoint module.
Expand Down Expand Up @@ -76,9 +35,6 @@ 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 @@ -161,13 +117,13 @@ def __call__(self, rho_vector=None, need_value=True, need_gradient=True):
print("Starting forward run...")
self.forward_run()
print("Starting adjoint run...")
self.a_E = []
self.D_a = []
self.adjoint_run()
print("Calculating gradient...")
self.calculate_gradient()
elif self.current_state == "FWD":
print("Starting adjoint run...")
self.a_E = []
self.D_a = []
self.adjoint_run()
print("Calculating gradient...")
self.calculate_gradient()
Expand Down Expand Up @@ -210,28 +166,9 @@ 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(
self.components,
self.frequencies,
where=dr.volume,
yee_grid=True,
decimation_factor=self.decimation_factor,
) for dr in self.design_regions
]

# 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 self.components
]
self.design_grids += [YeeDims(*s)]
self.design_region_monitors = utils.install_design_region_monitors(
self.sim, self.design_regions, self.frequencies, self.decimation_factor
)

def forward_run(self):
# set up monitors
Expand All @@ -254,16 +191,8 @@ def forward_run(self):
if len(self.f0) == 1:
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((self.nf, c[0], c[1], c[2]), dtype=np.complex128)
for c in dg
] for dg in self.design_grids]
for nb, dgm in enumerate(self.design_region_monitors):
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))
# Store forward fields for each set of design variables in array
self.D_f = utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies)

# store objective function evaluation in memory
self.f_bank.append(self.f0)
Expand All @@ -288,11 +217,14 @@ def adjoint_run(self):
# set up adjoint sources and monitors
self.prepare_adjoint_run()

if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
# flip the m number
if utils._check_if_cylindrical(self.sim):
self.sim.m = -self.sim.m

# flip the k point
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,15 +233,9 @@ def adjoint_run(self):
self.sim.change_sources(self.adjoint_sources[ar])

# register design flux
self.design_region_monitors = [
self.sim.add_dft_fields(
self.components,
self.frequencies,
where=dr.volume,
yee_grid=True,
decimation_factor=self.decimation_factor,
) for dr in self.design_regions
]
self.design_region_monitors = utils.install_design_region_monitors(
self.sim, self.design_regions, self.frequencies, self.decimation_factor
)

# Adjoint run
self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
Expand All @@ -318,34 +244,26 @@ def adjoint_run(self):
self.maximum_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((self.nf, c[0], c[1], c[2]), dtype=np.complex128)
for c in dg
] for dg in self.design_grids])
for nb, dgm in enumerate(self.design_region_monitors):
for ic, c in enumerate(self.components):
for f in range(self.nf):
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:
# Store adjoint fields for each design set of design variables
self.D_a.append(utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies))

# reset the m number
if utils._check_if_cylindrical(self.sim):
self.sim.m = -self.sim.m

# update optimizer's state
# reset the k point
if self.sim.k_point: self.sim.k_point *= -1

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

def calculate_gradient(self):
# Iterate through all design regions and calculate gradient
self.gradient = [[
dr.get_gradient(
self.sim,
self.a_E[ar][dri],
self.d_E[dri],
self.D_a[ar][dri],
self.D_f[dri],
self.frequencies,
) for dri, dr in enumerate(self.design_regions)
] for ar in range(len(self.objective_functions))]
Expand Down Expand Up @@ -486,7 +404,7 @@ def calculate_fd_gradient(

return fd_gradient, fd_gradient_idx

def update_design(self, rho_vector):
def update_design(self, rho_vector, beta=None):
"""Update the design permittivity function.
rho_vector ....... a list of numpy arrays that maps to each design region
Expand All @@ -497,6 +415,8 @@ def update_design(self, rho_vector):
"Each vector of design variables must contain only one dimension."
)
b.update_design_parameters(rho_vector[bi])
if beta:
b.update_beta(beta)

self.sim.reset_meep()
self.current_state = "INIT"
Expand Down
90 changes: 85 additions & 5 deletions python/adjoint/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,93 @@
import meep as mp
import numpy as onp

from . import ObjectiveQuantity, DesignRegion
from . import ObjectiveQuantity

# Meep field components used to compute adjoint sensitivities
_ADJOINT_FIELD_COMPONENTS = [mp.Ex, mp.Ey, mp.Ez]
_ADJOINT_FIELD_COMPONENTS = [mp.Dx, mp.Dy, mp.Dz]
_ADJOINT_FIELD_COMPONENTS_CYL = [mp.Dr, mp.Dp, mp.Dz]

# The frequency axis in the array returned by `mp._get_gradient()`
_GRADIENT_FREQ_AXIS = 1

# The returned axis order from get_dft_array in cylindrical coordinates
_FREQ_AXIS = 0
_RHO_AXIS = 2
_PHI_AXIS = 3
_Z_AXIS = 1

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_weights(design_parameters)

def update_beta(self,beta):
self.design_parameters.beta=beta

def get_gradient(self, sim, fields_a, fields_f, frequencies):
num_freqs = onp.array(frequencies).size
shapes = []
'''We have the option to linear scale the gradients up front
using the scalegrad parameter (leftover from MPB API). Not
currently needed for any existing feature (but available for
future use)'''
scalegrad = 1
for component_idx, component in enumerate(_compute_components(sim)):
'''We need to correct for the rare cases that get_dft_array
returns a singleton element for the forward and adjoint fields.
This only occurs when we are in 2D and only working with a particular
polarization (as the other fields are never stored). For example, the
2D in-plane polarization consists of a single scalar Ez field
(technically, meep doesn't store anything for these cases, but get_dft_array
still returns 0).
Our get_gradient algorithm, however, requires we pass an array of
zeros with the proper shape as the design_region.'''
spatial_shape = sim.get_array_slice_dimensions(component, vol=self.volume)[0]
if (fields_a[component_idx][0,...].size == 1):
fields_a[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs))
fields_f[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs))
if _check_if_cylindrical(sim):
'''For some reason, get_dft_array returns the field
components in a different order than the convention used
throughout meep. So, we reorder them here so we can use
the same field macros later in our get_gradient function.
'''
fields_a[component_idx] = onp.transpose(fields_a[component_idx],(_FREQ_AXIS,_RHO_AXIS,_PHI_AXIS,_Z_AXIS))
fields_f[component_idx] = onp.transpose(fields_f[component_idx],(_FREQ_AXIS,_RHO_AXIS,_PHI_AXIS,_Z_AXIS))
shapes.append(fields_a[component_idx].shape)
fields_a[component_idx] = fields_a[component_idx].flatten(order='C')
fields_f[component_idx] = fields_f[component_idx].flatten(order='C')
shapes = onp.asarray(shapes).flatten(order='C')
fields_a = onp.concatenate(fields_a)
fields_f = onp.concatenate(fields_f)

grad = onp.zeros((num_freqs, self.num_design_params)) # preallocate
geom_list = sim.geometry
f = sim.fields
vol = sim._fit_volume_to_simulation(self.volume)

# compute the gradient
mp._get_gradient(grad,scalegrad,fields_a,fields_f,sim.gv,vol.swigobj,onp.array(frequencies),sim.geps,shapes)
return onp.squeeze(grad).T

def _check_if_cylindrical(sim):
return sim.is_cylindrical or (sim.dimensions == mp.CYLINDRICAL)

def _compute_components(sim):
return _ADJOINT_FIELD_COMPONENTS_CYL if _check_if_cylindrical(sim) else _ADJOINT_FIELD_COMPONENTS

def _make_at_least_nd(x: onp.ndarray, dims: int = 3) -> onp.ndarray:
"""Makes an array have at least the specified number of dimensions."""
Expand Down Expand Up @@ -61,15 +140,16 @@ def install_design_region_monitors(
simulation: mp.Simulation,
design_regions: List[DesignRegion],
frequencies: List[float],
decimation_factor: int = 0
) -> List[mp.DftFields]:
"""Installs DFT field monitors at the design regions of the simulation."""
design_region_monitors = [
simulation.add_dft_fields(
_ADJOINT_FIELD_COMPONENTS,
_compute_components(simulation),
frequencies,
where=design_region.volume,
yee_grid=True,
decimation_factor=0
decimation_factor=decimation_factor
) for design_region in design_regions
]
return design_region_monitors
Expand Down Expand Up @@ -120,7 +200,7 @@ def gather_design_region_fields(
fwd_fields = []
for monitor in design_region_monitors:
fields_by_component = []
for component in _ADJOINT_FIELD_COMPONENTS:
for component in _compute_components(simulation):
fields_by_freq = []
for freq_idx, _ in enumerate(frequencies):
fields = simulation.get_dft_array(monitor, component, freq_idx)
Expand Down
Loading

0 comments on commit 8dd39f7

Please sign in to comment.