Skip to content

Commit

Permalink
add docstrings to adjoint solver (#2556)
Browse files Browse the repository at this point in the history 
* add adjoint doc

* add args to adjoint doc

* fix args to adjoint doc

* minor fix

* minor fix
  • Loading branch information
@mochen4
mochen4 authored Jun 22, 2023
1 parent 0aabdcc commit 95ddacd
Show file tree
Hide file tree
Showing 2 changed files with 157 additions and 42 deletions.
139 changes: 105 additions & 34 deletions python/adjoint/objective.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,23 +147,7 @@ def _create_time_profile(self, fwidth_frac=0.1, adj_cutoff=5):


class EigenmodeCoefficient(ObjectiveQuantity):
"""A frequency-dependent eigenmode coefficient.
Attributes:
volume: the volume over which the eigenmode coefficient is calculated.
mode: the eigenmode number.
forward: whether the forward or backward mode coefficient is returned as
the result of the evaluation.
kpoint_func: an optional k-point function to use when evaluating the eigenmode
coefficient. When specified, this overrides the effect of `forward`.
kpoint_func_overlap_idx: the index of the mode coefficient to return when
specifying `kpoint_func`. When specified, this overrides the effect of
`forward` and should have a value of either 0 or 1.
subtracted_dft_fields: the DFT fields obtained using `get_flux_data` from
a previous normalization run. This is subtracted from the DFT fields
of this mode monitor in order to improve the accuracy of the
reflectance measurement (i.e., the $S_{11}$ scattering parameter).
Default is None.
"""
"""A differentiable frequency-dependent eigenmode coefficient."""

def __init__(
self,
Expand All @@ -177,8 +161,29 @@ def __init__(
subtracted_dft_fields: Optional[FluxData] = None,
**kwargs
):
"""
+ **`sim` [ `Simulation` ]** —
"""Initialize an instance of differentiable frequency-dependent eigenmode coefficient.
Args:
sim: the Meep simulation object of the problem.
volume: the volume over which the eigenmode coefficient is calculated.
mode: the eigenmode number.
forward: whether the forward or backward mode coefficient is returned as
the result of the evaluation. Default is True.
kpoint_func: an optional k-point function to use when evaluating the eigenmode
coefficient. When specified, this overrides the effect of `forward`.
kpoint_func_overlap_idx: the index of the mode coefficient to return when
specifying `kpoint_func`. When specified, this overrides the effect of
`forward` and should have a value of either 0 or 1.
decimation_factor: An integer used to specify the number of timesteps between updates of
the DFT fields. The default is 0, at which the value is automatically determined from the
Nyquist rate of the bandwidth-limited sources and the DFT monitor. It can be turned off
by setting it to 1.
subtracted_dft_fields: the DFT fields obtained using `get_flux_data` from
a previous normalization run. This is subtracted from the DFT fields
of this mode monitor in order to improve the accuracy of the
reflectance measurement (i.e., the $S_{11}$ scattering parameter).
Default is None.
eigenmode_kwargs: additional argument for EigenModeSource
"""
super().__init__(sim)
if kpoint_func_overlap_idx not in [0, 1]:
Expand Down Expand Up @@ -257,6 +262,11 @@ def place_adjoint_source(self, dJ):
return [source]

def __call__(self):
"""The values of eigenmode coefficient at each frequency
Returns:
1D array of eigenmode coefficients corresponding to each of self.frequencies
"""
if self.kpoint_func:
kpoint_func = self.kpoint_func
overlap_idx = self.kpoint_func_overlap_idx
Expand Down Expand Up @@ -284,16 +294,35 @@ def __call__(self):


class FourierFields(ObjectiveQuantity):
"""A differentiable frequency-dependent Fourier fields (dft_fields)"""

def __init__(
self,
sim: mp.Simulation,
volume: mp.Volume,
component: List[int],
component: int,
yee_grid: Optional[bool] = False,
decimation_factor: Optional[int] = 0,
subtracted_dft_fields: Optional[FluxData] = None,
):
""" """
"""Initialize an instance of differentiable Fourier fields instance.
Args:
sim: the Meep simulation object of the problem.
volume: the volume over which the eigenmode coefficient is calculated. Due to an unresolved bug,
the size must not be zero in at least one direction.
component: field component (e.g. mp.Ex, mp.Hz, etc.) of the Fourier fields
yee_grid: if True, the Fourier fields are evaluated at the corresponding Yee grid points;
otherwise, they are interpolated fields at the center of each voxel. Default is False
decimation_factor: An integer used to specify the number of timesteps between updates of
the DFT fields. The default is 0, at which the value is automatically determined from the
Nyquist rate of the bandwidth-limited sources and the DFT monitor. It can be turned off
by setting it to 1.
subtracted_dft_fields: the DFT fields obtained using `get_flux_data` from
a previous normalization run. This is subtracted from the DFT fields
of this mode monitor in order to improve the accuracy of the
reflectance measurement (i.e., the $S_{11}$ scattering parameter). Default is None.
"""
super().__init__(sim)
self.volume = sim._fit_volume_to_simulation(volume)
self.component = component
Expand Down Expand Up @@ -344,11 +373,11 @@ def place_adjoint_source(self, dJ):
scale factor for each point (rather than a scale vector).
"""

self.all_fouriersrcdata = self._monitor.swigobj.fourier_sourcedata(
all_fouriersrcdata = self._monitor.swigobj.fourier_sourcedata(
self.volume.swigobj, self.component, self.sim.fields, dJ
)

for fourier_data in self.all_fouriersrcdata:
for fourier_data in all_fouriersrcdata:
amp_arr = np.array(fourier_data.amp_arr).reshape(-1, self.num_freq)
scale = amp_arr * self._adj_src_scale(include_resolution=False)

Expand All @@ -375,6 +404,13 @@ def place_adjoint_source(self, dJ):
return sources

def __call__(self):
"""The values of Fourier Fields at each frequency
Returns:
array of Fourier Fields with dimension k+1 where k is the dimension of self.volume
The first axis corresponds to the index of frequency, and the rest k axis are for
the spatial indices of points in the monitor
"""
self._eval = np.array(
[
self.sim.get_dft_array(self._monitor, self.component, i)
Expand All @@ -385,6 +421,8 @@ def __call__(self):


class Near2FarFields(ObjectiveQuantity):
"""A differentiable near2far field transformation"""

def __init__(
self,
sim: mp.Simulation,
Expand All @@ -394,7 +432,25 @@ def __init__(
decimation_factor: Optional[int] = 0,
norm_near_fields: Optional[NearToFarData] = None,
):
""" """
"""Initialize an instance of differentiable Fourier fields instance.
Args:
sim: the Meep simulation object of the problem.
Near2FarRegions: List of mp.Near2FarRegion over which the near fields are collected
far_pts: list of far points at which fields are computed
nperiods: If nperiods > 1, sum of 2*nperiods+1 Bloch-periodic copies of near fields
is computed to approximate the lattice sum from Bloch periodic boundary condition.
Default is 1 (no sum).
decimation_factor: An integer used to specify the number of timesteps between updates of
the DFT fields. The default is 0, at which the value is automatically determined from the
Nyquist rate of the bandwidth-limited sources and the DFT monitor. It can be turned off
by setting it to 1.
norm_near_fields: the DFT fields obtained using `get_near2far_data` from
a previous normalization run. This is subtracted from the DFT fields
of this near2far monitor in order to improve the accuracy of the
reflectance measurement (i.e., the $S_{11}$ scattering parameter).
Default is None.
"""
super().__init__(sim)
self.Near2FarRegions = Near2FarRegions
self.far_pts = far_pts # list of far pts
Expand Down Expand Up @@ -461,24 +517,34 @@ def place_adjoint_source(self, dJ):
return sources

def __call__(self):
"""The values of far fields at each points at each frequency
Returns:
3D array of far fields. The first axis is the index of far field points in self.far_pts;
the second axis is the index of frequency; and the third is the index of component in
[mp.Ex(mp.Er), mp.Ey(mp.Ep), mp.Ez, mp.Hx(mp.Hr), mp.Hy(mp.Hp), mp.Hz]
"""
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


class LDOS(ObjectiveQuantity):
def __init__(
self, sim: mp.Simulation, decimation_factor: Optional[int] = 0, **kwargs
):
""" """
"""A differentiable LDOS"""

def __init__(self, sim: mp.Simulation, **kwargs):
"""Initialize a differentiable LDOS instance
Args:
sim: the Meep simulation object of the problem.
"""
super().__init__(sim)
self.decimation_factor = decimation_factor
self.srckwarg = kwargs

def register_monitors(self, frequencies):
self._frequencies = np.asarray(frequencies)
self.forward_src = self.sim.sources
self._forward_src = self.sim.sources
return

def place_adjoint_source(self, dJ):
Expand All @@ -488,7 +554,7 @@ def place_adjoint_source(self, dJ):
dJ = dJ.flatten()
sources = []
forward_f_scale = np.array(
[self.ldos_scale / self.ldos_Jdata[k] for k in range(self.num_freq)]
[self._ldos_scale / self._ldos_Jdata[k] for k in range(self.num_freq)]
)
if self._frequencies.size == 1:
amp = (dJ * self._adj_src_scale(False) * forward_f_scale)[0]
Expand All @@ -502,7 +568,7 @@ def place_adjoint_source(self, dJ):
self.sim.fields.dt,
)
amp = 1
for forward_src_i in self.forward_src:
for forward_src_i in self._forward_src:
if isinstance(forward_src_i, mp.EigenModeSource):
src_i = mp.EigenModeSource(
src,
Expand Down Expand Up @@ -530,7 +596,12 @@ def place_adjoint_source(self, dJ):
return sources

def __call__(self):
"""The values of LDOS at each frequency
Returns:
1D array of LDOS corresponding to each of self.frequencies
"""
self._eval = self.sim.ldos_data
self.ldos_scale = self.sim.ldos_scale
self.ldos_Jdata = self.sim.ldos_Jdata
self._ldos_scale = self.sim.ldos_scale
self._ldos_Jdata = self.sim.ldos_Jdata
return np.array(self._eval)
60 changes: 52 additions & 8 deletions python/adjoint/optimization_problem.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ class OptimizationProblem:
of design variables, compute the objective function value (forward
calculation) and optionally its gradient (adjoint calculation).
This is done by the __call__ method.
"""

def __init__(
Expand All @@ -40,13 +39,37 @@ def __init__(
finite_difference_step: Optional[float] = utils.FD_DEFAULT,
step_funcs: list = None,
):
"""
+ **`simulation` [ `Simulation` ]** — The corresponding Meep
`Simulation` object that describes the problem (e.g. sources,
geometry)

+ **`objective_functions` [ `list of ` ]** —
"""Initialize an instance of OptimizationProblem.
Args:
sim: the corresponding Meep `Simulation` object that describes the
problem (e.g. sources, geometry)
objective_functions: list of differentiable functions (callable objects) whose arguments are
given by objective_arguments. The functions should take all of objective_arguments
as argument, even if not all of them are used in the functions.
For example, if we are interested in functions f(A,B) and g(B,C) of quantities A, B, C, then the
objective_functions list has to be [f1, g1] where f1 = lambda A, B, C: f(A,B) and g1 = lambda A, B, C: g(B,C);
and we have to specify arguments as [A,B,C]
objective_arguments: list of ObjectiveQuantity passed as arguments of objective functions
design_regions: list of DesignRegion to be optimized
frequencies: list of frequencies of interest in the problem. If not specified, then the list of frequencies
will be created from fcen, df, and nf: a list of size nf that goes from fcen-df/2 to fcen+df/2
fcen: center frequency
df: range of frequencies, i.e. maximum frequency - minimum frequency
nf: number of frequencies
decay_by: an number used to specify the amount by which all the field components and frequencies
frequencies of every DFT object have to decay before simulation stops. Default is 1e-11.
decimation_factor: an integer used to specify the number of timesteps between updates of
updates of the DFT fields. The default is 0, at which the value is
automatically determined from the Nyquist rate of the bandwidth-limited
sources and the DFT monitor. It can be turned off by setting it to 1
minimum_run_time: a number ensures the minimum runtime for each simulation. Default is 0
maximum_run_time: a number caps the maximum runtime for each simulation
finite_difference_step: step size for calculation of finite difference gradients
step_funcs: list of step functions to be called at each timestep
"""

self.step_funcs = step_funcs if step_funcs is not None else []
self.sim = simulation

Expand Down Expand Up @@ -123,7 +146,28 @@ def __call__(
need_gradient: bool = True,
beta: float = None,
) -> Tuple[List[np.ndarray], List[List[np.ndarray]]]:
"""Evaluate value and/or gradient of objective function."""
"""Evaluate value and/or gradient of objective function.
Args:
rho_vector: lists of design variables. Each list represents design variables
of one design region. The design is updated to the specified values; functions
and gradients will then be evaluated at this configuration of design variables.
need_value: whether forward evaluations are needed. Default is True.
need_gradient: whether adjoint and gradients evaluatiosn are needed. Default is True.
beta: the strength of projection of rho_vector. Default to None.
Returns:
A tuple (f0, gradient) where:
f0 is the list of objective functions values when design variables
are set to rho_vector
gradient is a list (over objective functions) of lists (over design regions) of 2d arrays
(design variables by frequencies) of derivatives. If there is only a single objective function,
the outer 1-element list is replaced by just that element, and similarly if there is only one design region
then those 1-element list are replaced by just those elements. In addition, if there is only one frequency
then the innermost array is squeezed to a 1d array.
For example, if there is only a single objective function, a single design region, and a single frequency,
then gradient is simply a 1d array of the derivatives.
"""
if rho_vector:
self.update_design(rho_vector=rho_vector, beta=beta)

Expand Down Expand Up @@ -320,7 +364,7 @@ def calculate_gradient(self):
elif len(self.gradient[0]) == 1:
self.gradient = [
g[0] for g in self.gradient
] # multiple objective functions bu one design region
] # multiple objective functions but one design region
# Return optimizer's state to initialization
self.current_state = "INIT"

Expand Down

0 comments on commit 95ddacd

Please sign in to comment.