diff --git a/python/simulation.py b/python/simulation.py index a8afd9e3a..0f71c51e9 100644 --- a/python/simulation.py +++ b/python/simulation.py @@ -201,12 +201,12 @@ class PML: def __init__( self, - thickness, - direction=mp.ALL, - side=mp.ALL, - R_asymptotic=1e-15, - mean_stretch=1.0, - pml_profile=DefaultPMLProfile, + thickness: float = None, + direction: int = mp.ALL, + side: int = mp.ALL, + R_asymptotic: float = 1e-15, + mean_stretch: float = 1.0, + pml_profile: Callable[[float], float] = DefaultPMLProfile, ): """ + **`thickness` [`number`]** — The spatial thickness of the PML layer which @@ -246,6 +246,9 @@ def __init__( example, one can use a cubic profile $f(u) = u^3$ by specifying `pml_profile=lambda u: u*u*u`. """ + if thickness is None: + raise ValueError("PML thickness must be specified.") + self.thickness = thickness self.direction = direction self.side = side @@ -275,8 +278,8 @@ def mean_stretch(self): return self._mean_stretch @mean_stretch.setter - def mean_stretch(self, val): - if val >= 1: + def mean_stretch(self, val: float): + if val >= 1.0: self._mean_stretch = val else: raise ValueError(f"PML.mean_stretch must be >= 1. Got {val}") @@ -333,7 +336,7 @@ class Symmetry: the axis of the rotation. """ - def __init__(self, direction, phase=1): + def __init__(self, direction: int = None, phase: complex = 1.0 + 0j): """ Construct a `Symmetry`. @@ -390,11 +393,11 @@ class Volume: def __init__( self, - center=Vector3(), - size=Vector3(), - dims=2, - is_cylindrical=False, - vertices=[], + center: Vector3Type = Vector3(), + size: Vector3Type = Vector3(), + dims: int = 2, + is_cylindrical: bool = False, + vertices: List[Vector3Type] = [], ): """ Construct a Volume. @@ -466,7 +469,7 @@ def nearly_equal(a, b, sig_fig=10): edges.append([vertices[iter1], vertices[iter2]]) return edges - def pt_in_volume(self, pt): + def pt_in_volume(self, pt: Vector3Type): xmin = self.center.x - self.size.x / 2 xmax = self.center.x + self.size.x / 2 ymin = self.center.y - self.size.y / 2 @@ -500,11 +503,11 @@ class FluxRegion: def __init__( self, - center=None, - size=Vector3(), - direction=mp.AUTOMATIC, - weight=1.0, - volume=None, + center: Vector3Type = None, + size: Vector3Type = Vector3(), + direction: int = mp.AUTOMATIC, + weight: float = 1.0, + volume: Optional[Volume] = None, ): """ Construct a `FluxRegion` object. @@ -760,7 +763,9 @@ def eps(self): def mu(self): return self.swigobj_attr("mu") - def flux(self, direction, where, resolution): + def flux( + self, direction: int = None, where: Volume = None, resolution: float = None + ): """ Given a `Volume` `where` (may be 0d, 1d, 2d, or 3d) and a `resolution` (in grid points / distance unit), compute the far fields in `where` (which may lie @@ -821,7 +826,15 @@ def chunks(self): class EigenmodeData: - def __init__(self, band_num, freq, group_velocity, k, swigobj, kdom): + def __init__( + self, + band_num, + freq: float, + group_velocity: float, + k: Vector3Type, + swigobj, + kdom: Vector3Type, + ): """Construct an `EigenmodeData`.""" self.band_num = band_num self.freq = freq @@ -883,7 +896,14 @@ class Harminv: ``` """ - def __init__(self, c, pt, fcen, df, mxbands=None): + def __init__( + self, + c: int = None, + pt: Vector3Type = None, + fcen: float = None, + df: float = None, + mxbands: Optional[int] = None, + ): """ Construct a Harminv object. @@ -1346,7 +1366,7 @@ def __init__( # to Volumes they create, the library will adjust them appropriately based # on the settings in the Simulation instance. This method must be called on # any user-defined Volume before passing it to meep via its `swigobj`. - def _fit_volume_to_simulation(self, vol): + def _fit_volume_to_simulation(self, vol: Volume) -> Volume: return Volume( vol.center, vol.size, @@ -1356,7 +1376,9 @@ def _fit_volume_to_simulation(self, vol): # Every function that takes a user volume can be specified either by a volume # (a Python Volume or a SWIG-wrapped meep::volume), or a center and a size - def _volume_from_kwargs(self, vol=None, center=None, size=None): + def _volume_from_kwargs( + self, vol: Volume = None, center: Vector3Type = None, size: Vector3Type = None + ) -> Volume: if vol: if isinstance(vol, Volume): # A pure Python Volume @@ -1374,7 +1396,7 @@ def _volume_from_kwargs(self, vol=None, center=None, size=None): else: raise ValueError("Need either a Volume, or a size and center") - def _infer_dimensions(self, k): + def _infer_dimensions(self, k: Vector3Type = None): if self.dimensions == 3: def use_2d(self, k): @@ -1437,7 +1459,7 @@ def _check_material_frequencies(self): warn_dft_fmt.format(dftf, min_freq, max_freq), RuntimeWarning ) - def _create_grid_volume(self, k): + def _create_grid_volume(self, k: Vector3Type = None): dims = self._infer_dimensions(k) if dims == 0 or dims == 1: @@ -1462,7 +1484,7 @@ def _create_grid_volume(self, k): ) return gv - def _create_symmetries(self, gv): + def _create_symmetries(self, gv) -> Symmetry: sym = mp.symmetry() # Initialize swig objects for each symmetry and combine them into one @@ -1483,7 +1505,7 @@ def _create_symmetries(self, gv): return sym - def _get_dft_volumes(self): + def _get_dft_volumes(self) -> List[Volume]: volumes = [ self._volume_from_kwargs( vol=r.where if hasattr(r, "where") else None, @@ -1496,7 +1518,7 @@ def _get_dft_volumes(self): return volumes - def _boundaries_to_vols_1d(self, boundaries): + def _boundaries_to_vols_1d(self, boundaries) -> List[Volume]: v1 = [] for bl in boundaries: @@ -1993,7 +2015,9 @@ def get_avg_chunk_communication_area(self): def get_estimated_costs(self): return [self.structure.estimated_cost(i) for i in range(mp.count_processors())] - def set_materials(self, geometry=None, default_material=None): + def set_materials( + self, geometry: List[GeometricObject] = None, default_material: Medium = None + ): """ This can be called in a step function, and is useful for changing the geometry or default material as a function of time. @@ -2060,7 +2084,7 @@ def set_materials(self, geometry=None, default_material=None): None, ) - def dump_structure(self, fname, single_parallel_file=True): + def dump_structure(self, fname: str = None, single_parallel_file: bool = True): """ Dumps the structure to the file `fname`. """ @@ -2076,7 +2100,7 @@ def dump_structure(self, fname, single_parallel_file=True): ) ) - def load_structure(self, fname, single_parallel_file=True): + def load_structure(self, fname: str = None, single_parallel_file: bool = True): """ Loads a structure from the file `fname`. """ @@ -2092,7 +2116,7 @@ def load_structure(self, fname, single_parallel_file=True): % (fname, str(single_parallel_file)) ) - def dump_fields(self, fname, single_parallel_file=True): + def dump_fields(self, fname: str = None, single_parallel_file: bool = True): """ Dumps the fields to the file `fname`. """ @@ -2106,7 +2130,7 @@ def dump_fields(self, fname, single_parallel_file=True): ) ) - def load_fields(self, fname, single_parallel_file=True): + def load_fields(self, fname: str = None, single_parallel_file: bool = True): """ Loads a fields from the file `fname`. """ @@ -2124,7 +2148,7 @@ def load_fields(self, fname, single_parallel_file=True): ) ) - def dump_chunk_layout(self, fname): + def dump_chunk_layout(self, fname: str = None): """ Dumps the chunk layout to file `fname`. """ @@ -2149,7 +2173,9 @@ def load_chunk_layout(self, br, source): ## source is either filename (string) self.structure.load_chunk_layout(source, br) - def get_load_dump_dirname(self, dirname, single_parallel_file): + def get_load_dump_dirname( + self, dirname: str = None, single_parallel_file: bool = None + ): """ Get the (possibly rank specific) dirname to dump simulation state to. """ @@ -2162,7 +2188,11 @@ def get_load_dump_dirname(self, dirname, single_parallel_file): return dump_dirname def dump( - self, dirname, dump_structure=True, dump_fields=True, single_parallel_file=True + self, + dirname: str = None, + dump_structure: bool = True, + dump_fields: bool = True, + single_parallel_file: bool = True, ): """ Dumps simulation state. @@ -2179,7 +2209,11 @@ def dump( self.dump_fields(fields_dump_filename, single_parallel_file) def load( - self, dirname, load_structure=True, load_fields=True, single_parallel_file=True + self, + dirname: str, + load_structure: bool = True, + load_fields: bool = True, + single_parallel_file: bool = True, ): """ Loads simulation state. @@ -2276,7 +2310,9 @@ def using_real_fields(self): cond5 = not (cond3 or cond4 or self.k_point == Vector3()) return not (self.force_complex_fields or cond1 or cond2 or cond5) - def initialize_field(self, cmpnt, amp_func): + def initialize_field( + self, cmpnt: int = None, amp_func: Callable[[Vector3Type], Union[float,complex]] = None + ): """ Initialize the component `c` fields using the function `func` which has a single argument, a `Vector3` giving a position and returns a complex number for the value @@ -2416,7 +2452,7 @@ def set_boundary(self, side, direction, condition): self.fields.set_boundary(side, direction, condition) - def get_field_point(self, c, pt): + def get_field_point(self, c: int = None, pt: Vector3Type = None): """ Given a `component` or `derived_component` constant `c` and a `Vector3` `pt`, returns the value of that component at that point. @@ -2424,7 +2460,7 @@ def get_field_point(self, c, pt): v3 = py_v3_to_vec(self.dimensions, pt, self.is_cylindrical) return self.fields.get_field_from_comp(c, v3) - def get_epsilon_point(self, pt, frequency=0): + def get_epsilon_point(self, pt: Vector3Type = None, frequency: float = 0.0): """ Given a frequency `frequency` and a `Vector3` `pt`, returns the average eigenvalue of the permittivity tensor at that location and frequency. If `frequency` is @@ -2434,7 +2470,7 @@ def get_epsilon_point(self, pt, frequency=0): v3 = py_v3_to_vec(self.dimensions, pt, self.is_cylindrical) return self.fields.get_eps(v3, frequency) - def get_mu_point(self, pt, frequency=0): + def get_mu_point(self, pt: Vector3Type = None, frequency: float = 0.0): """ Given a frequency `frequency` and a `Vector3` `pt`, returns the average eigenvalue of the permeability tensor at that location and frequency. If `frequency` is @@ -2444,7 +2480,13 @@ def get_mu_point(self, pt, frequency=0): v3 = py_v3_to_vec(self.dimensions, pt, self.is_cylindrical) return self.fields.get_mu(v3, frequency) - def get_epsilon_grid(self, xtics=None, ytics=None, ztics=None, frequency=0): + def get_epsilon_grid( + self, + xtics: np.ndarray = None, + ytics: np.ndarray = None, + ztics: np.ndarray = None, + frequency: float = 0.0, + ): """ Given three 1d NumPy arrays (`xtics`,`ytics`,`ztics`) which define the coordinates of a Cartesian grid anywhere within the cell volume, compute the trace of the $\\varepsilon(f)$ tensor at frequency @@ -2504,7 +2546,7 @@ def get_filename_prefix(self): "Expected a string for filename_prefix, or None for the default." ) - def use_output_directory(self, dname=""): + def use_output_directory(self, dname: str = ""): """ Output all files into a subdirectory, which is created if necessary. If the optional argument `dname` is specified, that is the name of the directory. If `dname` @@ -2625,7 +2667,7 @@ def _run_sources(self, step_funcs): """ self._run_sources_until(self, 0, step_funcs) - def run_k_point(self, t, k): + def run_k_point(self, t: float = None, k: Vector3Type = None): """ Lower level function called by `run_k_points` that runs a simulation for a single *k* point `k_point` and returns a `Harminv` instance. Useful when you need to @@ -2660,7 +2702,7 @@ def run_k_point(self, t, k): return h - def run_k_points(self, t, k_points): + def run_k_points(self, t: float = None, k_points: List[Vector3Type] = None): """ Given a list of `Vector3`, `k_points` of *k* vectors, runs a simulation for each *k* point (i.e. specifying Bloch-periodic boundary conditions) and extracts the @@ -3083,7 +3125,14 @@ def get_farfield(self, near2far, x): py_v3_to_vec(self.dimensions, x, is_cylindrical=self.is_cylindrical), ) - def get_farfields(self, near2far, resolution, where=None, center=None, size=None): + def get_farfields( + self, + near2far, + resolution: float = None, + where: Volume = None, + center: Vector3Type = None, + size: Vector3Type = None, + ): """ Like `output_farfields` but returns a dictionary of NumPy arrays instead of writing to a file. The dictionary keys are `Ex`, `Ey`, `Ez`, `Hx`, `Hy`, `Hz`. @@ -3117,7 +3166,13 @@ def get_farfields(self, near2far, resolution, where=None, center=None, size=None } def output_farfields( - self, near2far, fname, resolution, where=None, center=None, size=None + self, + near2far, + fname: str = None, + resolution: float = None, + where: Volume = None, + center: Vector3Type = None, + size: Vector3Type = None, ): """ Given an HDF5 file name `fname` (does *not* include the `.h5` suffix), a `Volume`