From 95ddacdda9754b0c6613b2a25a1154e87ab427e9 Mon Sep 17 00:00:00 2001 From: mochen4 Date: Thu, 22 Jun 2023 16:20:49 -0400 Subject: [PATCH] add docstrings to adjoint solver (#2556) * add adjoint doc * add args to adjoint doc * fix args to adjoint doc * minor fix * minor fix --- python/adjoint/objective.py | 139 +++++++++++++++++++------ python/adjoint/optimization_problem.py | 60 +++++++++-- 2 files changed, 157 insertions(+), 42 deletions(-) diff --git a/python/adjoint/objective.py b/python/adjoint/objective.py index bfee38983..eda9f91e3 100644 --- a/python/adjoint/objective.py +++ b/python/adjoint/objective.py @@ -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, @@ -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]: @@ -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 @@ -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 @@ -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) @@ -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) @@ -385,6 +421,8 @@ def __call__(self): class Near2FarFields(ObjectiveQuantity): + """A differentiable near2far field transformation""" + def __init__( self, sim: mp.Simulation, @@ -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 @@ -461,6 +517,13 @@ 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)) @@ -468,17 +531,20 @@ def __call__(self): 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): @@ -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] @@ -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, @@ -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) diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index 245b74717..d72d89db0 100644 --- a/python/adjoint/optimization_problem.py +++ b/python/adjoint/optimization_problem.py @@ -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__( @@ -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 @@ -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) @@ -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"