From b7d4bbbd959e873f068b9c9e4b3628a540cee582 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Thu, 11 Jun 2020 19:26:17 -0700 Subject: [PATCH] type annotation fixup --- meanas/eigensolvers.py | 10 +++---- meanas/fdfd/operators.py | 30 ++++++++++---------- meanas/fdfd/scpml.py | 20 +++++++------- meanas/fdfd/waveguide_2d.py | 28 ++++++++++--------- meanas/fdfd/waveguide_3d.py | 28 +++++++++---------- meanas/fdfd/waveguide_cyl.py | 4 +-- meanas/fdmath/__init__.py | 2 +- meanas/fdmath/functional.py | 28 ++++++++++--------- meanas/fdmath/operators.py | 20 +++++++------- meanas/fdmath/types.py | 9 ++++-- meanas/fdmath/vectorization.py | 20 +++++++++++--- meanas/fdtd/base.py | 6 ++-- meanas/fdtd/boundaries.py | 8 +++--- meanas/fdtd/energy.py | 50 +++++++++++++++++++--------------- meanas/fdtd/pml.py | 8 +++--- 15 files changed, 149 insertions(+), 122 deletions(-) diff --git a/meanas/eigensolvers.py b/meanas/eigensolvers.py index 51aeb6e..67cecea 100644 --- a/meanas/eigensolvers.py +++ b/meanas/eigensolvers.py @@ -1,7 +1,7 @@ """ Solvers for eigenvalue / eigenvector problems """ -from typing import Tuple, List +from typing import Tuple, Callable, Optional, Union import numpy from numpy.linalg import norm from scipy import sparse @@ -9,7 +9,7 @@ import scipy.sparse.linalg as spalg def power_iteration(operator: sparse.spmatrix, - guess_vector: numpy.ndarray = None, + guess_vector: Optional[numpy.ndarray] = None, iterations: int = 20, ) -> Tuple[complex, numpy.ndarray]: """ @@ -36,11 +36,11 @@ def power_iteration(operator: sparse.spmatrix, return lm_eigval, v -def rayleigh_quotient_iteration(operator: sparse.spmatrix or spalg.LinearOperator, +def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOperator], guess_vector: numpy.ndarray, iterations: int = 40, tolerance: float = 1e-13, - solver = None, + solver: Optional[Callable[..., numpy.ndarray]] = None, ) -> Tuple[complex, numpy.ndarray]: """ Use Rayleigh quotient iteration to refine an eigenvector guess. @@ -83,7 +83,7 @@ def rayleigh_quotient_iteration(operator: sparse.spmatrix or spalg.LinearOperato return eigval, v -def signed_eigensolve(operator: sparse.spmatrix or spalg.LinearOperator, +def signed_eigensolve(operator: Union[sparse.spmatrix, spalg.LinearOperator], how_many: int, negative: bool = False, ) -> Tuple[numpy.ndarray, numpy.ndarray]: diff --git a/meanas/fdfd/operators.py b/meanas/fdfd/operators.py index 3d6374a..212ed4c 100644 --- a/meanas/fdfd/operators.py +++ b/meanas/fdfd/operators.py @@ -27,7 +27,7 @@ The following operators are included: - Cross product matrices """ -from typing import List, Tuple +from typing import Tuple, Optional import numpy import scipy.sparse as sparse @@ -41,9 +41,9 @@ __author__ = 'Jan Petykiewicz' def e_full(omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, - pec: vfdfield_t = None, - pmc: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, + pec: Optional[vfdfield_t] = None, + pmc: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Wave operator @@ -125,9 +125,9 @@ def e_full_preconditioners(dxes: dx_lists_t def h_full(omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, - pec: vfdfield_t = None, - pmc: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, + pec: Optional[vfdfield_t] = None, + pmc: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Wave operator @@ -181,9 +181,9 @@ def h_full(omega: complex, def eh_full(omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, - pec: vfdfield_t = None, - pmc: vfdfield_t = None + mu: Optional[vfdfield_t] = None, + pec: Optional[vfdfield_t] = None, + pmc: Optional[vfdfield_t] = None ) -> sparse.spmatrix: """ Wave operator for `[E, H]` field representation. This operator implements Maxwell's @@ -249,8 +249,8 @@ def eh_full(omega: complex, def e2h(omega: complex, dxes: dx_lists_t, - mu: vfdfield_t = None, - pmc: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, + pmc: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Utility operator for converting the E field into the H field. @@ -280,7 +280,7 @@ def e2h(omega: complex, def m2j(omega: complex, dxes: dx_lists_t, - mu: vfdfield_t = None + mu: Optional[vfdfield_t] = None ) -> sparse.spmatrix: """ Operator for converting a magnetic current M into an electric current J. @@ -362,7 +362,7 @@ def e_tfsf_source(TF_region: vfdfield_t, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Operator that turns a desired E-field distribution into a @@ -392,7 +392,7 @@ def e_boundary_source(mask: vfdfield_t, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, periodic_mask_edges: bool = False, ) -> sparse.spmatrix: """ diff --git a/meanas/fdfd/scpml.py b/meanas/fdfd/scpml.py index 7e7694d..099fac2 100644 --- a/meanas/fdfd/scpml.py +++ b/meanas/fdfd/scpml.py @@ -2,10 +2,10 @@ Functions for creating stretched coordinate perfectly matched layer (PML) absorbers. """ -from typing import List, Callable +from typing import Sequence, Union, Callable, Optional import numpy -from ..fdmath import dx_lists_t +from ..fdmath import dx_lists_t, dx_lists_mut __author__ = 'Jan Petykiewicz' @@ -37,12 +37,12 @@ def prepare_s_function(ln_R: float = -16, return s_factor -def uniform_grid_scpml(shape: numpy.ndarray or List[int], - thicknesses: numpy.ndarray or List[int], +def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]], + thicknesses: Union[numpy.ndarray, Sequence[int]], omega: float, epsilon_effective: float = 1.0, - s_function: s_function_t = None, - ) -> dx_lists_t: + s_function: Optional[s_function_t] = None, + ) -> dx_lists_mut: """ Create dx arrays for a uniform grid with a cell width of 1 and a pml. @@ -63,7 +63,7 @@ def uniform_grid_scpml(shape: numpy.ndarray or List[int], Default uses `prepare_s_function()` with no parameters. Returns: - Complex cell widths (dx_lists_t) as discussed in `meanas.fdmath.types`. + Complex cell widths (dx_lists_mut) as discussed in `meanas.fdmath.types`. """ if s_function is None: s_function = prepare_s_function() @@ -90,13 +90,13 @@ def uniform_grid_scpml(shape: numpy.ndarray or List[int], return [dx_a, dx_b] -def stretch_with_scpml(dxes: dx_lists_t, +def stretch_with_scpml(dxes: dx_lists_mut, axis: int, polarity: int, omega: float, epsilon_effective: float = 1.0, thickness: int = 10, - s_function: s_function_t = None, + s_function: Optional[s_function_t] = None, ) -> dx_lists_t: """ Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis. @@ -113,7 +113,7 @@ def stretch_with_scpml(dxes: dx_lists_t, of pml parameters. Default uses `prepare_s_function()` with no parameters. Returns: - Complex cell widths (dx_lists_t) as discussed in `meanas.fdmath.types`. + Complex cell widths (dx_lists_mut) as discussed in `meanas.fdmath.types`. Multiple calls to this function may be necessary if multiple absorpbing boundaries are needed. """ if s_function is None: diff --git a/meanas/fdfd/waveguide_2d.py b/meanas/fdfd/waveguide_2d.py index 4113f8d..21bb404 100644 --- a/meanas/fdfd/waveguide_2d.py +++ b/meanas/fdfd/waveguide_2d.py @@ -146,7 +146,7 @@ to account for numerical dispersion if the result is introduced into a space wit """ # TODO update module docs -from typing import List, Tuple +from typing import List, Tuple, Optional import numpy from numpy.linalg import norm import scipy.sparse as sparse @@ -163,7 +163,7 @@ __author__ = 'Jan Petykiewicz' def operator_e(omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Waveguide operator of the form @@ -229,7 +229,7 @@ def operator_e(omega: complex, def operator_h(omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Waveguide operator of the form @@ -298,7 +298,7 @@ def normalized_fields_e(e_xy: numpy.ndarray, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, prop_phase: float = 0, ) -> Tuple[vfdfield_t, vfdfield_t]: """ @@ -332,7 +332,7 @@ def normalized_fields_h(h_xy: numpy.ndarray, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, prop_phase: float = 0, ) -> Tuple[vfdfield_t, vfdfield_t]: """ @@ -366,7 +366,7 @@ def _normalized_fields(e: numpy.ndarray, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None, + mu: Optional[vfdfield_t] = None, prop_phase: float = 0, ) -> Tuple[vfdfield_t, vfdfield_t]: # TODO documentation @@ -405,7 +405,7 @@ def exy2h(wavenumber: complex, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None + mu: Optional[vfdfield_t] = None ) -> sparse.spmatrix: """ Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, @@ -430,7 +430,7 @@ def hxy2e(wavenumber: complex, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = None + mu: Optional[vfdfield_t] = None ) -> sparse.spmatrix: """ Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, @@ -453,7 +453,7 @@ def hxy2e(wavenumber: complex, def hxy2h(wavenumber: complex, dxes: dx_lists_t, - mu: vfdfield_t = None + mu: Optional[vfdfield_t] = None ) -> sparse.spmatrix: """ Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, @@ -520,7 +520,7 @@ def exy2e(wavenumber: complex, def e2h(wavenumber: complex, omega: complex, dxes: dx_lists_t, - mu: vfdfield_t = None + mu: Optional[vfdfield_t] = None ) -> sparse.spmatrix: """ Returns an operator which, when applied to a vectorized E eigenfield, produces @@ -676,7 +676,7 @@ def solve_modes(mode_numbers: List[int], epsilon: vfdfield_t, mu: vfdfield_t = None, mode_margin: int = 2, - ) -> Tuple[List[vfdfield_t], List[complex]]: + ) -> Tuple[numpy.ndarray, List[complex]]: """ Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers. @@ -691,7 +691,8 @@ def solve_modes(mode_numbers: List[int], ability to find the correct mode. Default 2. Returns: - (e_xys, wavenumbers) + e_xys: list of vfdfield_t specifying fields + wavenumbers: list of wavenumbers """ ''' @@ -733,5 +734,6 @@ def solve_mode(mode_number: int, Returns: (e_xy, wavenumber) """ - e_xys, wavenumbers = solve_modes(mode_numbers=[mode_number], *args, **kwargs) + kwargs['mode_numbers'] = [mode_number] + e_xys, wavenumbers = solve_modes(*args, **kwargs) return e_xys[:, 0], wavenumbers[0] diff --git a/meanas/fdfd/waveguide_3d.py b/meanas/fdfd/waveguide_3d.py index 02fb7fd..847986d 100644 --- a/meanas/fdfd/waveguide_3d.py +++ b/meanas/fdfd/waveguide_3d.py @@ -4,7 +4,7 @@ Tools for working with waveguide modes in 3D domains. This module relies heavily on `waveguide_2d` and mostly just transforms its parameters into 2D equivalents and expands the results back into 3D. """ -from typing import Dict, List, Tuple +from typing import Dict, List, Tuple, Optional, Sequence, Union import numpy import scipy.sparse as sparse @@ -17,10 +17,10 @@ def solve_mode(mode_number: int, dxes: dx_lists_t, axis: int, polarity: int, - slices: List[slice], + slices: Sequence[slice], epsilon: fdfield_t, - mu: fdfield_t = None, - ) -> Dict[str, complex or numpy.ndarray]: + mu: Optional[fdfield_t] = None, + ) -> Dict[str, Union[complex, numpy.ndarray]]: """ Given a 3D grid, selects a slice from the grid and attempts to solve for an eigenmode propagating through that slice. @@ -104,9 +104,9 @@ def compute_source(E: fdfield_t, dxes: dx_lists_t, axis: int, polarity: int, - slices: List[slice], + slices: Sequence[slice], epsilon: fdfield_t, - mu: fdfield_t = None, + mu: Optional[fdfield_t] = None, ) -> fdfield_t: """ Given an eigenmode obtained by `solve_mode`, returns the current source distribution @@ -148,7 +148,7 @@ def compute_overlap_e(E: fdfield_t, dxes: dx_lists_t, axis: int, polarity: int, - slices: List[slice], + slices: Sequence[slice], ) -> fdfield_t: # TODO DOCS """ Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the @@ -177,9 +177,9 @@ def compute_overlap_e(E: fdfield_t, start, stop = sorted((slices[axis].start, slices[axis].start - 2 * polarity)) - slices2 = list(slices) - slices2[axis] = slice(start, stop) - slices2 = (slice(None), *slices2) + slices2_l = list(slices) + slices2_l[axis] = slice(start, stop) + slices2 = (slice(None), *slices2_l) Etgt = numpy.zeros_like(Ee) Etgt[slices2] = Ee[slices2] @@ -193,7 +193,7 @@ def expand_e(E: fdfield_t, dxes: dx_lists_t, axis: int, polarity: int, - slices: List[slice], + slices: Sequence[slice], ) -> fdfield_t: """ Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D @@ -226,9 +226,9 @@ def expand_e(E: fdfield_t, # Expand our slice to the entire grid using the phase factors E_expanded = numpy.zeros_like(E) - slices_exp = list(slices) - slices_exp[axis] = slice(E.shape[axis + 1]) - slices_exp = (slice(None), *slices_exp) + slices_exp_l = list(slices) + slices_exp_l[axis] = slice(E.shape[axis + 1]) + slices_exp = (slice(None), *slices_exp_l) slices_in = (slice(None), *slices) diff --git a/meanas/fdfd/waveguide_cyl.py b/meanas/fdfd/waveguide_cyl.py index 0995984..8ac55bd 100644 --- a/meanas/fdfd/waveguide_cyl.py +++ b/meanas/fdfd/waveguide_cyl.py @@ -8,7 +8,7 @@ As the z-dependence is known, all the functions in this file assume a 2D grid """ # TODO update module docs -from typing import List, Tuple, Dict +from typing import List, Tuple, Dict, Union import numpy from numpy.linalg import norm import scipy.sparse as sparse @@ -85,7 +85,7 @@ def solve_mode(mode_number: int, dxes: dx_lists_t, epsilon: vfdfield_t, r0: float, - ) -> Dict[str, complex or fdfield_t]: + ) -> Dict[str, Union[complex, fdfield_t]]: """ TODO: fixup Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode diff --git a/meanas/fdmath/__init__.py b/meanas/fdmath/__init__.py index a3b1af3..4b629f9 100644 --- a/meanas/fdmath/__init__.py +++ b/meanas/fdmath/__init__.py @@ -741,7 +741,7 @@ the true values can be multiplied back in after the simulation is complete if no normalized results are needed. """ -from .types import fdfield_t, vfdfield_t, dx_lists_t, fdfield_updater_t +from .types import fdfield_t, vfdfield_t, dx_lists_t, dx_lists_mut, fdfield_updater_t from .vectorization import vec, unvec from . import operators, functional, types, vectorization diff --git a/meanas/fdmath/functional.py b/meanas/fdmath/functional.py index 7de046a..379c310 100644 --- a/meanas/fdmath/functional.py +++ b/meanas/fdmath/functional.py @@ -3,13 +3,14 @@ Math functions for finite difference simulations Basic discrete calculus etc. """ -from typing import List, Callable, Tuple, Dict +from typing import Sequence, Tuple, Dict, Optional import numpy from .types import fdfield_t, fdfield_updater_t -def deriv_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t: +def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None + ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: """ Utility operators for taking discretized derivatives (backward variant). @@ -21,17 +22,18 @@ def deriv_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t: List of functions for taking forward derivatives along each axis. """ if dx_e: - derivs = [lambda f: (numpy.roll(f, -1, axis=0) - f) / dx_e[0][:, None, None], + derivs = (lambda f: (numpy.roll(f, -1, axis=0) - f) / dx_e[0][:, None, None], lambda f: (numpy.roll(f, -1, axis=1) - f) / dx_e[1][None, :, None], - lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :]] + lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :]) else: - derivs = [lambda f: numpy.roll(f, -1, axis=0) - f, + derivs = (lambda f: numpy.roll(f, -1, axis=0) - f, lambda f: numpy.roll(f, -1, axis=1) - f, - lambda f: numpy.roll(f, -1, axis=2) - f] + lambda f: numpy.roll(f, -1, axis=2) - f) return derivs -def deriv_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t: +def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None + ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: """ Utility operators for taking discretized derivatives (forward variant). @@ -43,17 +45,17 @@ def deriv_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t: List of functions for taking forward derivatives along each axis. """ if dx_h: - derivs = [lambda f: (f - numpy.roll(f, 1, axis=0)) / dx_h[0][:, None, None], + derivs = (lambda f: (f - numpy.roll(f, 1, axis=0)) / dx_h[0][:, None, None], lambda f: (f - numpy.roll(f, 1, axis=1)) / dx_h[1][None, :, None], - lambda f: (f - numpy.roll(f, 1, axis=2)) / dx_h[2][None, None, :]] + lambda f: (f - numpy.roll(f, 1, axis=2)) / dx_h[2][None, None, :]) else: - derivs = [lambda f: f - numpy.roll(f, 1, axis=0), + derivs = (lambda f: f - numpy.roll(f, 1, axis=0), lambda f: f - numpy.roll(f, 1, axis=1), - lambda f: f - numpy.roll(f, 1, axis=2)] + lambda f: f - numpy.roll(f, 1, axis=2)) return derivs -def curl_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t: +def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t: """ Curl operator for use with the E field. @@ -80,7 +82,7 @@ def curl_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t: return ce_fun -def curl_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t: +def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t: """ Create a function which takes the backward curl of a field. diff --git a/meanas/fdmath/operators.py b/meanas/fdmath/operators.py index 3f83c8c..30a6708 100644 --- a/meanas/fdmath/operators.py +++ b/meanas/fdmath/operators.py @@ -3,14 +3,14 @@ Matrix operators for finite difference simulations Basic discrete calculus etc. """ -from typing import List, Callable, Tuple, Dict +from typing import Sequence, List, Callable, Tuple, Dict import numpy import scipy.sparse as sparse from .types import fdfield_t, vfdfield_t -def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix: +def rotation(axis: int, shape: Sequence[int], shift_distance: int=1) -> sparse.spmatrix: """ Utility operator for performing a circular shift along a specified axis by a specified number of elements. @@ -46,7 +46,7 @@ def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmat return d -def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix: +def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int=1) -> sparse.spmatrix: """ Utility operator for performing an n-element shift along a specified axis, with mirror boundary conditions applied to the cells beyond the receding edge. @@ -87,7 +87,7 @@ def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> spa return d -def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]: +def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]: """ Utility operators for taking discretized derivatives (forward variant). @@ -112,7 +112,7 @@ def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]: return Ds -def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]: +def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]: """ Utility operators for taking discretized derivatives (backward variant). @@ -137,7 +137,7 @@ def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]: return Ds -def cross(B: List[sparse.spmatrix]) -> sparse.spmatrix: +def cross(B: Sequence[sparse.spmatrix]) -> sparse.spmatrix: """ Cross product operator @@ -171,7 +171,7 @@ def vec_cross(b: vfdfield_t) -> sparse.spmatrix: return cross(B) -def avg_forward(axis: int, shape: List[int]) -> sparse.spmatrix: +def avg_forward(axis: int, shape: Sequence[int]) -> sparse.spmatrix: """ Forward average operator `(x4 = (x4 + x5) / 2)` @@ -189,7 +189,7 @@ def avg_forward(axis: int, shape: List[int]) -> sparse.spmatrix: return 0.5 * (sparse.eye(n) + rotation(axis, shape)) -def avg_back(axis: int, shape: List[int]) -> sparse.spmatrix: +def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix: """ Backward average operator `(x4 = (x4 + x3) / 2)` @@ -203,7 +203,7 @@ def avg_back(axis: int, shape: List[int]) -> sparse.spmatrix: return avg_forward(axis, shape).T -def curl_forward(dx_e: List[numpy.ndarray]) -> sparse.spmatrix: +def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix: """ Curl operator for use with the E field. @@ -217,7 +217,7 @@ def curl_forward(dx_e: List[numpy.ndarray]) -> sparse.spmatrix: return cross(deriv_forward(dx_e)) -def curl_back(dx_h: List[numpy.ndarray]) -> sparse.spmatrix: +def curl_back(dx_h: Sequence[numpy.ndarray]) -> sparse.spmatrix: """ Curl operator for use with the H field. diff --git a/meanas/fdmath/types.py b/meanas/fdmath/types.py index ea78f8f..8215fed 100644 --- a/meanas/fdmath/types.py +++ b/meanas/fdmath/types.py @@ -2,7 +2,7 @@ Types shared across multiple submodules """ import numpy -from typing import List, Callable +from typing import Sequence, Callable, MutableSequence # Field types @@ -25,7 +25,7 @@ class vfdfield_t(numpy.ndarray): pass -dx_lists_t = List[List[numpy.ndarray]] +dx_lists_t = Sequence[Sequence[numpy.ndarray]] ''' 'dxes' datastructure which contains grid cell width information in the following format: @@ -36,6 +36,11 @@ dx_lists_t = List[List[numpy.ndarray]] and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc. ''' +dx_lists_mut = MutableSequence[MutableSequence[numpy.ndarray]] +''' + Mutable version of `dx_lists_t` +''' + fdfield_updater_t = Callable[..., fdfield_t] ''' diff --git a/meanas/fdmath/vectorization.py b/meanas/fdmath/vectorization.py index 63d78ef..807fc5d 100644 --- a/meanas/fdmath/vectorization.py +++ b/meanas/fdmath/vectorization.py @@ -4,16 +4,20 @@ and a 1D array representation of that field `[f_x0, f_x1, f_x2,... f_y0,... f_z0 Vectorized versions of the field use row-major (ie., C-style) ordering. """ -from typing import List +from typing import Optional, TypeVar, overload, Union, List import numpy from .types import fdfield_t, vfdfield_t +@overload +def vec(f: None) -> None: + pass -__author__ = 'Jan Petykiewicz' +@overload +def vec(f: Union[fdfield_t, List[numpy.ndarray]]) -> vfdfield_t: + pass - -def vec(f: fdfield_t) -> vfdfield_t: +def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield_t]: """ Create a 1D ndarray from a 3D vector field which spans a 1-3D region. @@ -31,7 +35,15 @@ def vec(f: fdfield_t) -> vfdfield_t: return numpy.ravel(f, order='C') +@overload +def unvec(v: None, shape: numpy.ndarray) -> None: + pass + +@overload def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t: + pass + +def unvec(v: Optional[vfdfield_t], shape: numpy.ndarray) -> Optional[fdfield_t]: """ Perform the inverse of vec(): take a 1D ndarray and output a 3D field of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional diff --git a/meanas/fdtd/base.py b/meanas/fdtd/base.py index e16a53f..a632ca2 100644 --- a/meanas/fdtd/base.py +++ b/meanas/fdtd/base.py @@ -3,7 +3,7 @@ Basic FDTD field updates """ -from typing import List, Callable, Tuple, Dict +from typing import List, Callable, Dict, Union import numpy from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t @@ -47,7 +47,7 @@ def maxwell_e(dt: float, dxes: dx_lists_t = None) -> fdfield_updater_t: else: curl_h_fun = curl_back() - def me_fun(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t): + def me_fun(e: fdfield_t, h: fdfield_t, epsilon: Union[fdfield_t, float]) -> fdfield_t: """ Update the E-field. @@ -100,7 +100,7 @@ def maxwell_h(dt: float, dxes: dx_lists_t = None) -> fdfield_updater_t: else: curl_e_fun = curl_forward() - def mh_fun(e: fdfield_t, h: fdfield_t, mu: fdfield_t = None): + def mh_fun(e: fdfield_t, h: fdfield_t, mu: Union[fdfield_t, float, None] = None) -> fdfield_t: """ Update the H-field. diff --git a/meanas/fdtd/boundaries.py b/meanas/fdtd/boundaries.py index 10c966e..dbc1d93 100644 --- a/meanas/fdtd/boundaries.py +++ b/meanas/fdtd/boundaries.py @@ -4,7 +4,7 @@ Boundary conditions #TODO conducting boundary documentation """ -from typing import List, Callable, Tuple, Dict +from typing import Callable, Tuple, Dict, Any, List import numpy from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t @@ -20,8 +20,8 @@ def conducting_boundary(direction: int, u, v = dirs if polarity < 0: - boundary_slice = [slice(None)] * 3 - shifted1_slice = [slice(None)] * 3 + boundary_slice = [slice(None)] * 3 # type: List[Any] + shifted1_slice = [slice(None)] * 3 # type: List[Any] boundary_slice[direction] = 0 shifted1_slice[direction] = 1 @@ -42,7 +42,7 @@ def conducting_boundary(direction: int, if polarity > 0: boundary_slice = [slice(None)] * 3 shifted1_slice = [slice(None)] * 3 - shifted2_slice = [slice(None)] * 3 + shifted2_slice = [slice(None)] * 3 # type: List[Any] boundary_slice[direction] = -1 shifted1_slice[direction] = -2 shifted2_slice[direction] = -3 diff --git a/meanas/fdtd/energy.py b/meanas/fdtd/energy.py index dc5848a..443176b 100644 --- a/meanas/fdtd/energy.py +++ b/meanas/fdtd/energy.py @@ -1,5 +1,5 @@ # pylint: disable=unsupported-assignment-operation -from typing import List, Callable, Tuple, Dict +from typing import Callable, Tuple, Dict, Optional, Union import numpy from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t @@ -8,7 +8,7 @@ from ..fdmath.functional import deriv_back, deriv_forward def poynting(e: fdfield_t, h: fdfield_t, - dxes: dx_lists_t = None, + dxes: Optional[dx_lists_t] = None, ) -> fdfield_t: """ Calculate the poynting vector @@ -30,16 +30,19 @@ def poynting(e: fdfield_t, return s -def poynting_divergence(s: fdfield_t = None, +def poynting_divergence(s: Optional[fdfield_t] = None, *, - e: fdfield_t = None, - h: fdfield_t = None, - dxes: dx_lists_t = None, + e: Optional[fdfield_t] = None, + h: Optional[fdfield_t] = None, + dxes: Optional[dx_lists_t] = None, ) -> fdfield_t: """ Calculate the divergence of the poynting vector """ if s is None: + assert(e is not None) + assert(h is not None) + assert(dxes is not None) s = poynting(e, h, dxes=dxes) Dx, Dy, Dz = deriv_back() @@ -50,9 +53,9 @@ def poynting_divergence(s: fdfield_t = None, def energy_hstep(e0: fdfield_t, h1: fdfield_t, e2: fdfield_t, - epsilon: fdfield_t = None, - mu: fdfield_t = None, - dxes: dx_lists_t = None, + epsilon: Optional[fdfield_t] = None, + mu: Optional[fdfield_t] = None, + dxes: Optional[dx_lists_t] = None, ) -> fdfield_t: u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes) return u @@ -61,9 +64,9 @@ def energy_hstep(e0: fdfield_t, def energy_estep(h0: fdfield_t, e1: fdfield_t, h2: fdfield_t, - epsilon: fdfield_t = None, - mu: fdfield_t = None, - dxes: dx_lists_t = None, + epsilon: Optional[fdfield_t] = None, + mu: Optional[fdfield_t] = None, + dxes: Optional[dx_lists_t] = None, ) -> fdfield_t: u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes) return u @@ -74,9 +77,9 @@ def delta_energy_h2e(dt: float, h1: fdfield_t, e2: fdfield_t, h3: fdfield_t, - epsilon: fdfield_t = None, - mu: fdfield_t = None, - dxes: dx_lists_t = None, + epsilon: Optional[fdfield_t] = None, + mu: Optional[fdfield_t] = None, + dxes: Optional[dx_lists_t] = None, ) -> fdfield_t: """ This is just from (e2 * e2 + h3 * h1) - (h1 * h1 + e0 * e2) @@ -92,9 +95,9 @@ def delta_energy_e2h(dt: float, e1: fdfield_t, h2: fdfield_t, e3: fdfield_t, - epsilon: fdfield_t = None, - mu: fdfield_t = None, - dxes: dx_lists_t = None, + epsilon: Optional[fdfield_t] = None, + mu: Optional[fdfield_t] = None, + dxes: Optional[dx_lists_t] = None, ) -> fdfield_t: """ This is just from (h2 * h2 + e3 * e1) - (e1 * e1 + h0 * h2) @@ -105,7 +108,10 @@ def delta_energy_e2h(dt: float, return du -def delta_energy_j(j0: fdfield_t, e1: fdfield_t, dxes: dx_lists_t = None) -> fdfield_t: +def delta_energy_j(j0: fdfield_t, + e1: fdfield_t, + dxes: Optional[dx_lists_t] = None, + ) -> fdfield_t: if dxes is None: dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2)) @@ -118,9 +124,9 @@ def delta_energy_j(j0: fdfield_t, e1: fdfield_t, dxes: dx_lists_t = None) -> fdf def dxmul(ee: fdfield_t, hh: fdfield_t, - epsilon: fdfield_t = None, - mu: fdfield_t = None, - dxes: dx_lists_t = None + epsilon: Optional[Union[fdfield_t, float]] = None, + mu: Optional[Union[fdfield_t, float]] = None, + dxes: Optional[dx_lists_t] = None ) -> fdfield_t: if epsilon is None: epsilon = 1 diff --git a/meanas/fdtd/pml.py b/meanas/fdtd/pml.py index 2c4ae1e..0edd01a 100644 --- a/meanas/fdtd/pml.py +++ b/meanas/fdtd/pml.py @@ -7,7 +7,7 @@ PML implementations """ # TODO retest pmls! -from typing import List, Callable, Tuple, Dict +from typing import List, Callable, Tuple, Dict, Any import numpy from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t @@ -59,9 +59,9 @@ def cpml(direction: int, else: raise Exception('Bad polarity!') - expand_slice = [None] * 3 - expand_slice[direction] = slice(None) - expand_slice = tuple(expand_slice) + expand_slice_l: List[Any] = [None] * 3 + expand_slice_l[direction] = slice(None) + expand_slice = tuple(expand_slice_l) def par(x): scaling = (x / thickness) ** m