diff --git a/python/adjoint/objective.py b/python/adjoint/objective.py index 07aeea618..6a9a4fb65 100644 --- a/python/adjoint/objective.py +++ b/python/adjoint/objective.py @@ -1,6 +1,7 @@ """Handling of objective functions and objective quantities.""" import abc from collections import namedtuple +from typing import Callable, List, Optional import numpy as np from meep.simulation import py_v3_to_vec @@ -161,15 +162,18 @@ class EigenmodeCoefficient(ObjectiveQuantity): def __init__( self, - sim, - volume, - mode, - forward=True, - kpoint_func=None, - kpoint_func_overlap_idx=0, - decimation_factor=0, + sim: mp.Simulation, + volume: mp.Volume, + mode: int, + forward: Optional[bool] = True, + kpoint_func: Optional[Callable] = None, + kpoint_func_overlap_idx: Optional[int] = 0, + decimation_factor: Optional[int] = 0, **kwargs ): + """ + + **`sim` [ `Simulation` ]** — + """ super().__init__(sim) if kpoint_func_overlap_idx not in [0, 1]: raise ValueError( @@ -268,7 +272,15 @@ def __call__(self): class FourierFields(ObjectiveQuantity): - def __init__(self, sim, volume, component, yee_grid=False, decimation_factor=0): + def __init__( + self, + sim: mp.Simulation, + volume: mp.Volume, + component: List[int], + yee_grid: Optional[bool] = False, + decimation_factor: Optional[int] = 0, + ): + """ """ super().__init__(sim) self.volume = sim._fit_volume_to_simulation(volume) self.component = component @@ -354,7 +366,14 @@ def __call__(self): class Near2FarFields(ObjectiveQuantity): - def __init__(self, sim, Near2FarRegions, far_pts, decimation_factor=0): + def __init__( + self, + sim: mp.Simulation, + Near2FarRegions: List[mp.Near2FarRegion], + far_pts: List[mp.Vector3], + decimation_factor: Optional[int] = 0, + ): + """ """ super().__init__(sim) self.Near2FarRegions = Near2FarRegions self.far_pts = far_pts # list of far pts @@ -420,7 +439,10 @@ def __call__(self): class LDOS(ObjectiveQuantity): - def __init__(self, sim, decimation_factor=0, **kwargs): + def __init__( + self, sim: mp.Simulation, decimation_factor: Optional[int] = 0, **kwargs + ): + """ """ super().__init__(sim) self.decimation_factor = decimation_factor self.srckwarg = kwargs diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py index fdeb33e07..42f20b0aa 100644 --- a/python/adjoint/optimization_problem.py +++ b/python/adjoint/optimization_problem.py @@ -1,11 +1,12 @@ from collections import namedtuple +from typing import Callable, List, Union, Optional, Tuple import numpy as np from autograd import grad, jacobian import meep as mp -from . import LDOS, DesignRegion, utils +from . import LDOS, DesignRegion, utils, ObjectiveQuantity class OptimizationProblem: @@ -24,21 +25,28 @@ class OptimizationProblem: def __init__( self, - simulation, - objective_functions, - objective_arguments, - design_regions, - frequencies=None, - fcen=None, - df=None, - nf=None, - decay_by=1e-11, - decimation_factor=0, - minimum_run_time=0, - maximum_run_time=None, - finite_difference_step=utils.FD_DEFAULT, + simulation: mp.Simulation, + objective_functions: List[Callable], + objective_arguments: List[ObjectiveQuantity], + design_regions: List[DesignRegion], + frequencies: Optional[Union[float, List[float]]] = None, + fcen: Optional[float] = None, + df: Optional[float] = None, + nf: Optional[int] = None, + decay_by: Optional[float] = 1e-11, + decimation_factor: Optional[int] = 0, + minimum_run_time: Optional[float] = 0, + maximum_run_time: Optional[float] = None, + 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 ` ]** — + """ self.step_funcs = step_funcs if step_funcs is not None else [] self.sim = simulation @@ -108,7 +116,13 @@ def __init__( self.gradient = [] - def __call__(self, rho_vector=None, need_value=True, need_gradient=True, beta=None): + def __call__( + self, + rho_vector: List[List[float]] = None, + need_value: bool = True, + need_gradient: bool = True, + beta: float = None, + ) -> Tuple[List[float], List[float]]: """Evaluate value and/or gradient of objective function.""" if rho_vector: self.update_design(rho_vector=rho_vector, beta=beta) @@ -140,7 +154,7 @@ def __call__(self, rho_vector=None, need_value=True, need_gradient=True, beta=No return self.f0, self.gradient - def get_fdf_funcs(self): + def get_fdf_funcs(self) -> Tuple[Callable, Callable]: """construct callable functions for objective function value and gradient Returns @@ -306,11 +320,11 @@ def calculate_gradient(self): def calculate_fd_gradient( self, - num_gradients=1, - db=1e-4, - design_variables_idx=0, - filter=None, - ): + num_gradients: int = 1, + db: float = 1e-4, + design_variables_idx: int = 0, + filter: Callable = None, + ) -> List[float]: """ Estimate central difference gradients. @@ -448,7 +462,7 @@ def calculate_fd_gradient( return fd_gradient, fd_gradient_idx - def update_design(self, rho_vector, beta=None): + def update_design(self, rho_vector: List[float], beta: float = None) -> None: """Update the design permittivity function. rho_vector ....... a list of numpy arrays that maps to each design region @@ -465,11 +479,11 @@ def update_design(self, rho_vector, beta=None): self.sim.reset_meep() self.current_state = "INIT" - def get_objective_arguments(self): + def get_objective_arguments(self) -> List[float]: """Return list of evaluated objective arguments.""" return [m.get_evaluation() for m in self.objective_arguments] - def plot2D(self, init_opt=False, **kwargs): + def plot2D(self, init_opt=False, **kwargs) -> None: """Produce a graphical visualization of the geometry and/or fields, as appropriately autodetermined based on the current state of progress.