Adjoint typehints

python/adjoint/objective.py

@@ -1,13 +1,15 @@
 """Handling of objective functions and objective quantities."""
 import abc
 from collections import namedtuple
+from typing import Callable, List, Union, Optional
 
 import numpy as np
 from meep.simulation import py_v3_to_vec
 
 import meep as mp
 
 from .filter_source import FilteredSource
+from . import MEEP_COMPONENTS
 
 Grid = namedtuple("Grid", ["x", "y", "z", "w"])
 
@@ -161,15 +163,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[idx] = 0,
+        decimation_factor: Optional[idx] = 0,
         **kwargs
     ):
+        """
+        + **`sim` [ `Simulation` ]** —
+        """
         super().__init__(sim)
         if kpoint_func_overlap_idx not in [0, 1]:
             raise ValueError(
@@ -268,7 +273,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[MEEP_COMPONENTS],
+        yee_grid: Optional[bool] = False,
+        decimation_factor: Optional[idx] = 0,
+    ):
+        """ """
         super().__init__(sim)
         self.volume = sim._fit_volume_to_simulation(volume)
         self.component = component
@@ -354,7 +367,14 @@ def __call__(self):
 
 
 class Near2FarFields(ObjectiveQuantity):
-    def __init__(self, sim, Near2FarRegions, far_pts, decimation_factor=0):
+    def __init__(
+        self,
+        sim: mp.Simulation,
+        Near2FarRegions: mp.Near2FarRegions,
+        far_pts: List[mp.Vector3],
+        decimation_factor: Optional[idx] = 0,
+    ):
+        """ """
         super().__init__(sim)
         self.Near2FarRegions = Near2FarRegions
         self.far_pts = far_pts  # list of far pts
@@ -420,7 +440,10 @@ def __call__(self):
 
 
 class LDOS(ObjectiveQuantity):
-    def __init__(self, sim, decimation_factor=0, **kwargs):
+    def __init__(
+        self, sim: mp.Simulation, decimation_factor: Optional[idx] = 0, **kwargs
+    ):
+        """ """
         super().__init__(sim)
         self.decimation_factor = decimation_factor
         self.srckwarg = kwargs

python/adjoint/optimization_problem.py

@@ -1,4 +1,5 @@
 from collections import namedtuple
+from typing import Callable, List, Union, Optional
 
 import numpy as np
 from autograd import grad, jacobian
@@ -24,20 +25,27 @@ 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: [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,
     ):
+        """
+        + **`simulation` [ `Simulation` ]** — The corresponding Meep
+        `Simulation` object that describeds the problem (e.g. sources,
+        geometry)
+
+        + **`objective_functions` [ `list of ` ]** —
+        """
 
         self.sim = simulation
 

python/adjoint/utils.py

@@ -1,4 +1,4 @@
-from typing import Iterable, List, Tuple
+from typing import Iterable, List, Tuple, Union
 
 import numpy as onp
 
@@ -22,6 +22,11 @@
 # default finite difference step size when calculating Aᵤ
 FD_DEFAULT = 1e-3
 
+# Type hint for all components
+MEEP_COMPONENTS = Union[
+    mp.Ex, mp.Ey, mp.Ez, mp.Hx, mp.Hy, mp.Hz, mp.Dx, mp.Dy, mp.Dz, mp.Bx, mp.By, mp.Bz
+]
+
 
 class DesignRegion:
     def __init__(