typing and formatting updates

meanas/eigensolvers.py

@@ -2,16 +2,18 @@
 Solvers for eigenvalue / eigenvector problems
 """
 from typing import Tuple, Callable, Optional, Union
-import numpy                          # type: ignore
-from numpy.linalg import norm         # type: ignore
+import numpy
+from numpy.typing import NDArray, ArrayLike
+from numpy.linalg import norm
 from scipy import sparse              # type: ignore
 import scipy.sparse.linalg as spalg   # type: ignore
 
 
-def power_iteration(operator: sparse.spmatrix,
-                    guess_vector: Optional[numpy.ndarray] = None,
-                    iterations: int = 20,
-                    ) -> Tuple[complex, numpy.ndarray]:
+def power_iteration(
+        operator: sparse.spmatrix,
+        guess_vector: Optional[NDArray[numpy.float64]] = None,
+        iterations: int = 20,
+        ) -> Tuple[complex, NDArray[numpy.float64]]:
     """
     Use power iteration to estimate the dominant eigenvector of a matrix.
 
@@ -37,12 +39,13 @@ def power_iteration(operator: sparse.spmatrix,
     return lm_eigval, v
 
 
-def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOperator],
-                                guess_vector: numpy.ndarray,
-                                iterations: int = 40,
-                                tolerance: float = 1e-13,
-                                solver: Optional[Callable[..., numpy.ndarray]] = None,
-                                ) -> Tuple[complex, numpy.ndarray]:
+def rayleigh_quotient_iteration(
+        operator: Union[sparse.spmatrix, spalg.LinearOperator],
+        guess_vector: NDArray[numpy.float64],
+        iterations: int = 40,
+        tolerance: float = 1e-13,
+        solver: Optional[Callable[..., NDArray[numpy.float64]]] = None,
+        ) -> Tuple[complex, NDArray[numpy.float64]]:
     """
     Use Rayleigh quotient iteration to refine an eigenvector guess.
 
@@ -69,11 +72,13 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe
             solver = spalg.spsolve
     except TypeError:
         def shift(eigval: float) -> spalg.LinearOperator:
-            return spalg.LinearOperator(shape=operator.shape,
-                                        dtype=operator.dtype,
-                                        matvec=lambda v: eigval * v)
+            return spalg.LinearOperator(
+                    shape=operator.shape,
+                    dtype=operator.dtype,
+                    matvec=lambda v: eigval * v,
+                    )
         if solver is None:
-            def solver(A: spalg.LinearOperator, b: numpy.ndarray) -> numpy.ndarray:
+            def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.float64]:
                 return spalg.bicgstab(A, b)[0]
     assert(solver is not None)
 
@@ -90,10 +95,11 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe
     return eigval, v
 
 
-def signed_eigensolve(operator: Union[sparse.spmatrix, spalg.LinearOperator],
-                      how_many: int,
-                      negative: bool = False,
-                      ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+def signed_eigensolve(
+        operator: Union[sparse.spmatrix, spalg.LinearOperator],
+        how_many: int,
+        negative: bool = False,
+        ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
     """
     Find the largest-magnitude positive-only (or negative-only) eigenvalues and
      eigenvectors of the provided matrix.

meanas/fdfd/bloch.py

@@ -80,7 +80,7 @@ This module contains functions for generating and solving the
 
 '''
 
-from typing import Tuple, Callable, Any, List, Optional, cast
+from typing import Tuple, Callable, Any, List, Optional, cast, Union
 import logging
 import numpy
 from numpy import pi, real, trace
@@ -433,11 +433,10 @@ def find_k(
         `(k, actual_frequency)`
         The found k-vector and its frequency.
     """
-
     direction = numpy.array(direction) / norm(direction)
 
     def get_f(k0_mag: float, band: int = 0) -> float:
-        k0 = direction * k0_mag
+        k0 = direction * k0_mag                         # type: ignore
         n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
         f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
         if solve_callback:
@@ -482,6 +481,8 @@ def eigsolve(
         `(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the
         vector `eigenvectors[i, :]`
     """
+    k0 = numpy.array(k0, copy=False)
+
     h_size = 2 * epsilon[0].size
 
     kmag = norm(G_matrix @ k0)
@@ -497,9 +498,9 @@ def eigsolve(
 
     y_shape = (h_size, num_modes)
 
-    prev_E = 0
-    d_scale = 1
-    prev_traceGtKG = 0
+    prev_E = 0.0
+    d_scale = 1.0
+    prev_traceGtKG = 0.0
     #prev_theta = 0.5
     D = numpy.zeros(shape=y_shape, dtype=complex)
 
@@ -545,7 +546,7 @@ def eigsolve(
 
         if prev_traceGtKG == 0 or i % reset_iters == 0:
             logger.info('CG reset')
-            gamma = 0
+            gamma = 0.0
         else:
             gamma = traceGtKG / prev_traceGtKG
 
@@ -695,7 +696,10 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to
         return x, fx, dfx
 '''
 
-def _rtrace_AtB(A: NDArray[numpy.float64], B: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
+def _rtrace_AtB(
+        A: NDArray[numpy.float64],
+        B: Union[NDArray[numpy.float64], float],
+        ) -> float:
     return real(numpy.sum(A.conj() * B))
 
 def _symmetrize(A: NDArray[numpy.float64]) -> NDArray[numpy.float64]:

meanas/fdfd/farfield.py

@@ -2,19 +2,20 @@
 Functions for performing near-to-farfield transformation (and the reverse).
 """
 from typing import Dict, List, Any
-import numpy            # type: ignore
-from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift     # type: ignore
-from numpy import pi    # type: ignore
+import numpy
+from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
+from numpy import pi
 
 from ..fdmath import fdfield_t
 
 
-def near_to_farfield(E_near: fdfield_t,
-                     H_near: fdfield_t,
-                     dx: float,
-                     dy: float,
-                     padded_size: List[int] = None
-                     ) -> Dict[str, Any]:
+def near_to_farfield(
+        E_near: fdfield_t,
+        H_near: fdfield_t,
+        dx: float,
+        dy: float,
+        padded_size: List[int] = None
+        ) -> Dict[str, Any]:
     """
     Compute the farfield, i.e. the distribution of the fields after propagation
       through several wavelengths of uniform medium.
@@ -120,12 +121,13 @@ def near_to_farfield(E_near: fdfield_t,
     return outputs
 
 
-def far_to_nearfield(E_far: fdfield_t,
-                     H_far: fdfield_t,
-                     dkx: float,
-                     dky: float,
-                     padded_size: List[int] = None
-                     ) -> Dict[str, Any]:
+def far_to_nearfield(
+        E_far: fdfield_t,
+        H_far: fdfield_t,
+        dkx: float,
+        dky: float,
+        padded_size: List[int] = None
+        ) -> Dict[str, Any]:
     """
     Compute the farfield, i.e. the distribution of the fields after propagation
       through several wavelengths of uniform medium.

meanas/fdfd/functional.py

@@ -5,8 +5,8 @@ Functional versions of many FDFD operators. These can be useful for performing
 The functions generated here expect `fdfield_t` inputs with shape (3, X, Y, Z),
 e.g. E = [E_x, E_y, E_z] where each component has shape (X, Y, Z)
 """
-from typing import Callable, Tuple
-import numpy        # type: ignore
+from typing import Callable, Tuple, Optional
+import numpy
 
 from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
 from ..fdmath.functional import curl_forward, curl_back
@@ -15,11 +15,12 @@ from ..fdmath.functional import curl_forward, curl_back
 __author__ = 'Jan Petykiewicz'
 
 
-def e_full(omega: complex,
-           dxes: dx_lists_t,
-           epsilon: fdfield_t,
-           mu: fdfield_t = None
-           ) -> fdfield_updater_t:
+def e_full(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: fdfield_t,
+        mu: fdfield_t = None
+        ) -> fdfield_updater_t:
     """
     Wave operator for use with E-field. See `operators.e_full` for details.
 
@@ -50,11 +51,12 @@ def e_full(omega: complex,
         return op_mu
 
 
-def eh_full(omega: complex,
-            dxes: dx_lists_t,
-            epsilon: fdfield_t,
-            mu: fdfield_t = None
-            ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]:
+def eh_full(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: fdfield_t,
+        mu: fdfield_t = None
+        ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]:
     """
     Wave operator for full (both E and H) field representation.
     See `operators.eh_full`.
@@ -86,9 +88,10 @@ def eh_full(omega: complex,
         return op_mu
 
 
-def e2h(omega: complex,
+def e2h(
+        omega: complex,
         dxes: dx_lists_t,
-        mu: fdfield_t = None,
+        mu: Optional[fdfield_t] = None,
         ) -> fdfield_updater_t:
     """
     Utility operator for converting the `E` field into the `H` field.
@@ -117,9 +120,10 @@ def e2h(omega: complex,
         return e2h_mu
 
 
-def m2j(omega: complex,
+def m2j(
+        omega: complex,
         dxes: dx_lists_t,
-        mu: fdfield_t = None,
+        mu: Optional[fdfield_t] = None,
         ) -> fdfield_updater_t:
     """
     Utility operator for converting magnetic current `M` distribution
@@ -151,12 +155,13 @@ def m2j(omega: complex,
         return m2j_mu
 
 
-def e_tfsf_source(TF_region: fdfield_t,
-                  omega: complex,
-                  dxes: dx_lists_t,
-                  epsilon: fdfield_t,
-                  mu: fdfield_t = None,
-                  ) -> fdfield_updater_t:
+def e_tfsf_source(
+        TF_region: fdfield_t,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: fdfield_t,
+        mu: Optional[fdfield_t] = None,
+        ) -> fdfield_updater_t:
     """
     Operator that turns an E-field distribution into a total-field/scattered-field
     (TFSF) source.

meanas/fdfd/operators.py

@@ -28,7 +28,7 @@ The following operators are included:
 """
 
 from typing import Tuple, Optional
-import numpy                        # type: ignore
+import numpy
 import scipy.sparse as sparse       # type: ignore
 
 from ..fdmath import vec, dx_lists_t, vfdfield_t
@@ -38,13 +38,14 @@ from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl
 __author__ = 'Jan Petykiewicz'
 
 
-def e_full(omega: complex,
-           dxes: dx_lists_t,
-           epsilon: vfdfield_t,
-           mu: Optional[vfdfield_t] = None,
-           pec: Optional[vfdfield_t] = None,
-           pmc: Optional[vfdfield_t] = None,
-           ) -> sparse.spmatrix:
+def e_full(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        pec: Optional[vfdfield_t] = None,
+        pmc: Optional[vfdfield_t] = None,
+        ) -> sparse.spmatrix:
     """
     Wave operator
      $$ \\nabla \\times (\\frac{1}{\\mu} \\nabla \\times) - \\Omega^2 \\epsilon $$
@@ -96,8 +97,9 @@ def e_full(omega: complex,
     return op
 
 
-def e_full_preconditioners(dxes: dx_lists_t
-                           ) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
+def e_full_preconditioners(
+        dxes: dx_lists_t,
+        ) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
     """
     Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator.
 
@@ -122,13 +124,14 @@ def e_full_preconditioners(dxes: dx_lists_t
     return P_left, P_right
 
 
-def h_full(omega: complex,
-           dxes: dx_lists_t,
-           epsilon: vfdfield_t,
-           mu: Optional[vfdfield_t] = None,
-           pec: Optional[vfdfield_t] = None,
-           pmc: Optional[vfdfield_t] = None,
-           ) -> sparse.spmatrix:
+def h_full(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        pec: Optional[vfdfield_t] = None,
+        pmc: Optional[vfdfield_t] = None,
+        ) -> sparse.spmatrix:
     """
     Wave operator
      $$ \\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu $$
@@ -178,13 +181,14 @@ def h_full(omega: complex,
     return A
 
 
-def eh_full(omega: complex,
-            dxes: dx_lists_t,
-            epsilon: vfdfield_t,
-            mu: Optional[vfdfield_t] = None,
-            pec: Optional[vfdfield_t] = None,
-            pmc: Optional[vfdfield_t] = None
-            ) -> sparse.spmatrix:
+def eh_full(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        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
      equations without cancelling out either E or H. The operator is
@@ -247,7 +251,8 @@ def eh_full(omega: complex,
     return A
 
 
-def e2h(omega: complex,
+def e2h(
+        omega: complex,
         dxes: dx_lists_t,
         mu: Optional[vfdfield_t] = None,
         pmc: Optional[vfdfield_t] = None,
@@ -278,9 +283,10 @@ def e2h(omega: complex,
     return op
 
 
-def m2j(omega: complex,
+def m2j(
+        omega: complex,
         dxes: dx_lists_t,
-        mu: Optional[vfdfield_t] = None
+        mu: Optional[vfdfield_t] = None,
         ) -> sparse.spmatrix:
     """
     Operator for converting a magnetic current M into an electric current J.
@@ -357,12 +363,13 @@ def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     return P
 
 
-def e_tfsf_source(TF_region: vfdfield_t,
-                  omega: complex,
-                  dxes: dx_lists_t,
-                  epsilon: vfdfield_t,
-                  mu: Optional[vfdfield_t] = None,
-                  ) -> sparse.spmatrix:
+def e_tfsf_source(
+        TF_region: vfdfield_t,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        ) -> sparse.spmatrix:
     """
     Operator that turns a desired E-field distribution into a
      total-field/scattered-field (TFSF) source.
@@ -387,13 +394,14 @@ def e_tfsf_source(TF_region: vfdfield_t,
     return (A @ Q - Q @ A) / (-1j * omega)
 
 
-def e_boundary_source(mask: vfdfield_t,
-                      omega: complex,
-                      dxes: dx_lists_t,
-                      epsilon: vfdfield_t,
-                      mu: Optional[vfdfield_t] = None,
-                      periodic_mask_edges: bool = False,
-                      ) -> sparse.spmatrix:
+def e_boundary_source(
+        mask: vfdfield_t,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        periodic_mask_edges: bool = False,
+        ) -> sparse.spmatrix:
     """
     Operator that turns an E-field distrubtion into a current (J) distribution
       along the edges (external and internal) of the provided mask. This is just an

meanas/fdfd/scpml.py

@@ -3,7 +3,9 @@ Functions for creating stretched coordinate perfectly matched layer (PML) absorb
 """
 
 from typing import Sequence, Union, Callable, Optional, List
-import numpy            # type: ignore
+
+import numpy
+from numpy.typing import ArrayLike, NDArray
 
 
 __author__ = 'Jan Petykiewicz'
@@ -13,9 +15,10 @@ s_function_t = Callable[[float], float]
 """Typedef for s-functions, see `prepare_s_function()`"""
 
 
-def prepare_s_function(ln_R: float = -16,
-                       m: float = 4
-                       ) -> s_function_t:
+def prepare_s_function(
+        ln_R: float = -16,
+        m: float = 4
+        ) -> s_function_t:
     """
     Create an s_function to pass to the SCPML functions. This is used when you would like to
     customize the PML parameters.
@@ -29,18 +32,19 @@ def prepare_s_function(ln_R: float = -16,
         of the cell width; needs to be divided by `sqrt(epilon_effective) * real(omega))`
         before use.
     """
-    def s_factor(distance: numpy.ndarray) -> numpy.ndarray:
+    def s_factor(distance: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
         s_max = (m + 1) * ln_R / 2  # / 2 because we assume periodic boundaries
         return s_max * (distance ** m)
     return s_factor
 
 
-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: Optional[s_function_t] = None,
-                       ) -> List[List[numpy.ndarray]]:
+def uniform_grid_scpml(
+        shape: ArrayLike,    # ints
+        thicknesses: ArrayLike,  # ints
+        omega: float,
+        epsilon_effective: float = 1.0,
+        s_function: Optional[s_function_t] = None,
+        ) -> List[List[NDArray[numpy.float64]]]:
     """
     Create dx arrays for a uniform grid with a cell width of 1 and a pml.
 
@@ -67,7 +71,11 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
         s_function = prepare_s_function()
 
     # Normalized distance to nearest boundary
-    def ll(u: numpy.ndarray, n: numpy.ndarray, t: numpy.ndarray) -> numpy.ndarray:
+    def ll(
+            u: NDArray[numpy.float64],
+            n: NDArray[numpy.float64],
+            t: NDArray[numpy.float64],
+            ) -> NDArray[numpy.float64]:
         return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t
 
     dx_a = [numpy.array(numpy.inf)] * 3
@@ -88,14 +96,15 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
     return [dx_a, dx_b]
 
 
-def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
-                       axis: int,
-                       polarity: int,
-                       omega: float,
-                       epsilon_effective: float = 1.0,
-                       thickness: int = 10,
-                       s_function: Optional[s_function_t] = None,
-                       ) -> List[List[numpy.ndarray]]:
+def stretch_with_scpml(
+        dxes: List[List[NDArray[numpy.float64]]],
+        axis: int,
+        polarity: int,
+        omega: float,
+        epsilon_effective: float = 1.0,
+        thickness: int = 10,
+        s_function: Optional[s_function_t] = None,
+        ) -> List[List[NDArray[numpy.float64]]]:
     """
         Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis.
 
@@ -132,7 +141,7 @@ def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
         bound = pos[thickness]
         d = bound - pos[0]
 
-        def l_d(x: numpy.ndarray) -> numpy.ndarray:
+        def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
             return (bound - x) / (bound - pos[0])
 
         slc = slice(thickness)
@@ -142,7 +151,7 @@ def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
         bound = pos[-thickness - 1]
         d = pos[-1] - bound
 
-        def l_d(x: numpy.ndarray) -> numpy.ndarray:
+        def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
             return (x - bound) / (pos[-1] - bound)
 
         if thickness == 0:

meanas/fdfd/solvers.py

@@ -2,11 +2,12 @@
 Solvers and solver interface for FDFD problems.
 """
 
-from typing import Callable, Dict, Any
+from typing import Callable, Dict, Any, Optional
 import logging
 
-import numpy                        # type: ignore
-from numpy.linalg import norm       # type: ignore
+import numpy
+from numpy.typing import ArrayLike, NDArray
+from numpy.linalg import norm
 import scipy.sparse.linalg          # type: ignore
 
 from ..fdmath import dx_lists_t, vfdfield_t
@@ -16,10 +17,11 @@ from . import operators
 logger = logging.getLogger(__name__)
 
 
-def _scipy_qmr(A: scipy.sparse.csr_matrix,
-               b: numpy.ndarray,
-               **kwargs: Any,
-               ) -> numpy.ndarray:
+def _scipy_qmr(
+        A: scipy.sparse.csr_matrix,
+        b: ArrayLike,
+        **kwargs: Any,
+        ) -> NDArray[numpy.float64]:
     """
     Wrapper for scipy.sparse.linalg.qmr
 
@@ -37,14 +39,14 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix,
     '''
     ii = 0
 
-    def log_residual(xk: numpy.ndarray) -> None:
+    def log_residual(xk: ArrayLike) -> None:
         nonlocal ii
         ii += 1
         if ii % 100 == 0:
             logger.info('Solver residual at iteration {} : {}'.format(ii, norm(A @ xk - b)))
 
     if 'callback' in kwargs:
-        def augmented_callback(xk: numpy.ndarray) -> None:
+        def augmented_callback(xk: ArrayLike) -> None:
             log_residual(xk)
             kwargs['callback'](xk)
 
@@ -60,17 +62,18 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix,
     return x
 
 
-def generic(omega: complex,
-            dxes: dx_lists_t,
-            J: vfdfield_t,
-            epsilon: vfdfield_t,
-            mu: vfdfield_t = None,
-            pec: vfdfield_t = None,
-            pmc: vfdfield_t = None,
-            adjoint: bool = False,
-            matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr,
-            matrix_solver_opts: Dict[str, Any] = None,
-            ) -> vfdfield_t:
+def generic(
+        omega: complex,
+        dxes: dx_lists_t,
+        J: vfdfield_t,
+        epsilon: vfdfield_t,
+        mu: vfdfield_t = None,
+        pec: vfdfield_t = None,
+        pmc: vfdfield_t = None,
+        adjoint: bool = False,
+        matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
+        matrix_solver_opts: Optional[Dict[str, Any]] = None,
+        ) -> vfdfield_t:
     """
     Conjugate gradient FDFD solver using CSR sparse matrices.
 
@@ -90,8 +93,8 @@ def generic(omega: complex,
         adjoint: If true, solves the adjoint problem.
         matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`,
                 where `A`: `scipy.sparse.csr_matrix`;
-                      `b`: `numpy.ndarray`;
-                      `x`: `numpy.ndarray`;
+                      `b`: `ArrayLike`;
+                      `x`: `ArrayLike`;
                 Default is a wrapped version of `scipy.sparse.linalg.qmr()`
                  which doesn't return convergence info and logs the residual
                  every 100 iterations.

meanas/fdfd/waveguide_2d.py

@@ -179,8 +179,9 @@ to account for numerical dispersion if the result is introduced into a space wit
 # TODO update module docs
 
 from typing import List, Tuple, Optional, Any
-import numpy                        # type: ignore
-from numpy.linalg import norm       # type: ignore
+import numpy
+from numpy.typing import NDArray, ArrayLike
+from numpy.linalg import norm
 import scipy.sparse as sparse       # type: ignore
 
 from ..fdmath.operators import deriv_forward, deriv_back, cross
@@ -191,11 +192,12 @@ from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 __author__ = 'Jan Petykiewicz'
 
 
-def operator_e(omega: complex,
-               dxes: dx_lists_t,
-               epsilon: vfdfield_t,
-               mu: Optional[vfdfield_t] = None,
-               ) -> sparse.spmatrix:
+def operator_e(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        ) -> sparse.spmatrix:
     """
     Waveguide operator of the form
 
@@ -257,11 +259,12 @@ def operator_e(omega: complex,
     return op
 
 
-def operator_h(omega: complex,
-               dxes: dx_lists_t,
-               epsilon: vfdfield_t,
-               mu: Optional[vfdfield_t] = None,
-               ) -> sparse.spmatrix:
+def operator_h(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        ) -> sparse.spmatrix:
     """
     Waveguide operator of the form
 
@@ -324,14 +327,15 @@ def operator_h(omega: complex,
     return op
 
 
-def normalized_fields_e(e_xy: numpy.ndarray,
-                        wavenumber: complex,
-                        omega: complex,
-                        dxes: dx_lists_t,
-                        epsilon: vfdfield_t,
-                        mu: Optional[vfdfield_t] = None,
-                        prop_phase: float = 0,
-                        ) -> Tuple[vfdfield_t, vfdfield_t]:
+def normalized_fields_e(
+        e_xy: ArrayLike,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        prop_phase: float = 0,
+        ) -> Tuple[vfdfield_t, vfdfield_t]:
     """
     Given a vector `e_xy` containing the vectorized E_x and E_y fields,
      returns normalized, vectorized E and H fields for the system.
@@ -358,14 +362,15 @@ def normalized_fields_e(e_xy: numpy.ndarray,
     return e_norm, h_norm
 
 
-def normalized_fields_h(h_xy: numpy.ndarray,
-                        wavenumber: complex,
-                        omega: complex,
-                        dxes: dx_lists_t,
-                        epsilon: vfdfield_t,
-                        mu: Optional[vfdfield_t] = None,
-                        prop_phase: float = 0,
-                        ) -> Tuple[vfdfield_t, vfdfield_t]:
+def normalized_fields_h(
+        h_xy: ArrayLike,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        prop_phase: float = 0,
+        ) -> Tuple[vfdfield_t, vfdfield_t]:
     """
     Given a vector `h_xy` containing the vectorized H_x and H_y fields,
      returns normalized, vectorized E and H fields for the system.
@@ -392,14 +397,15 @@ def normalized_fields_h(h_xy: numpy.ndarray,
     return e_norm, h_norm
 
 
-def _normalized_fields(e: numpy.ndarray,
-                       h: numpy.ndarray,
-                       omega: complex,
-                       dxes: dx_lists_t,
-                       epsilon: vfdfield_t,
-                       mu: Optional[vfdfield_t] = None,
-                       prop_phase: float = 0,
-                       ) -> Tuple[vfdfield_t, vfdfield_t]:
+def _normalized_fields(
+        e: ArrayLike,
+        h: ArrayLike,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None,
+        prop_phase: float = 0,
+        ) -> Tuple[vfdfield_t, vfdfield_t]:
     # TODO documentation
     shape = [s.size for s in dxes[0]]
     dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)]
@@ -434,12 +440,13 @@ def _normalized_fields(e: numpy.ndarray,
     return e, h
 
 
-def exy2h(wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfdfield_t,
-          mu: Optional[vfdfield_t] = None
-          ) -> sparse.spmatrix:
+def exy2h(
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None
+        ) -> sparse.spmatrix:
     """
     Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
      into a vectorized H containing all three H components
@@ -459,12 +466,13 @@ def exy2h(wavenumber: complex,
     return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)
 
 
-def hxy2e(wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfdfield_t,
-          mu: Optional[vfdfield_t] = None
-          ) -> sparse.spmatrix:
+def hxy2e(
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None
+        ) -> sparse.spmatrix:
     """
     Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
      into a vectorized E containing all three E components
@@ -484,10 +492,11 @@ def hxy2e(wavenumber: complex,
     return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)
 
 
-def hxy2h(wavenumber: complex,
-          dxes: dx_lists_t,
-          mu: Optional[vfdfield_t] = None
-          ) -> sparse.spmatrix:
+def hxy2h(
+        wavenumber: complex,
+        dxes: dx_lists_t,
+        mu: Optional[vfdfield_t] = None
+        ) -> sparse.spmatrix:
     """
     Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
      into a vectorized H containing all three H components
@@ -517,10 +526,11 @@ def hxy2h(wavenumber: complex,
     return op
 
 
-def exy2e(wavenumber: complex,
-          dxes: dx_lists_t,
-          epsilon: vfdfield_t,
-          ) -> sparse.spmatrix:
+def exy2e(
+        wavenumber: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        ) -> sparse.spmatrix:
     """
     Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
      into a vectorized E containing all three E components
@@ -550,7 +560,8 @@ def exy2e(wavenumber: complex,
     return op
 
 
-def e2h(wavenumber: complex,
+def e2h(
+        wavenumber: complex,
         omega: complex,
         dxes: dx_lists_t,
         mu: Optional[vfdfield_t] = None
@@ -574,7 +585,8 @@ def e2h(wavenumber: complex,
     return op
 
 
-def h2e(wavenumber: complex,
+def h2e(
+        wavenumber: complex,
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t
@@ -636,13 +648,14 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
     return cross([Dbx, Dby, Bz])
 
 
-def h_err(h: vfdfield_t,
-          wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfdfield_t,
-          mu: vfdfield_t = None
-          ) -> float:
+def h_err(
+        h: vfdfield_t,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: Optional[vfdfield_t] = None
+        ) -> float:
     """
     Calculates the relative error in the H field
 
@@ -670,13 +683,14 @@ def h_err(h: vfdfield_t,
     return norm(op) / norm(h)
 
 
-def e_err(e: vfdfield_t,
-          wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfdfield_t,
-          mu: vfdfield_t = None
-          ) -> float:
+def e_err(
+        e: vfdfield_t,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: vfdfield_t = Optional[None]
+        ) -> float:
     """
     Calculates the relative error in the E field
 
@@ -703,13 +717,14 @@ def e_err(e: vfdfield_t,
     return norm(op) / norm(e)
 
 
-def solve_modes(mode_numbers: List[int],
-                omega: complex,
-                dxes: dx_lists_t,
-                epsilon: vfdfield_t,
-                mu: vfdfield_t = None,
-                mode_margin: int = 2,
-                ) -> Tuple[numpy.ndarray, List[complex]]:
+def solve_modes(
+        mode_numbers: List[int],
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: vfdfield_t = None,
+        mode_margin: int = 2,
+        ) -> Tuple[NDArray[numpy.float64], List[complex]]:
     """
     Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
 
@@ -752,10 +767,11 @@ def solve_modes(mode_numbers: List[int],
     return e_xys, wavenumbers
 
 
-def solve_mode(mode_number: int,
-               *args: Any,
-               **kwargs: Any,
-               ) -> Tuple[vfdfield_t, complex]:
+def solve_mode(
+        mode_number: int,
+        *args: Any,
+        **kwargs: Any,
+        ) -> Tuple[vfdfield_t, complex]:
     """
     Wrapper around `solve_modes()` that solves for a single mode.
 

meanas/fdfd/waveguide_3d.py

@@ -5,21 +5,23 @@ 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, Optional, Sequence, Union, Any
-import numpy                    # type: ignore
+import numpy
+from numpy.typing import NDArray
 
 from ..fdmath import vec, unvec, dx_lists_t, fdfield_t
 from . import operators, waveguide_2d
 
 
-def solve_mode(mode_number: int,
-               omega: complex,
-               dxes: dx_lists_t,
-               axis: int,
-               polarity: int,
-               slices: Sequence[slice],
-               epsilon: fdfield_t,
-               mu: Optional[fdfield_t] = None,
-               ) -> Dict[str, Union[complex, numpy.ndarray]]:
+def solve_mode(
+        mode_number: int,
+        omega: complex,
+        dxes: dx_lists_t,
+        axis: int,
+        polarity: int,
+        slices: Sequence[slice],
+        epsilon: fdfield_t,
+        mu: Optional[fdfield_t] = None,
+        ) -> Dict[str, Union[complex, NDArray[numpy.float_]]]:
     """
     Given a 3D grid, selects a slice from the grid and attempts to
      solve for an eigenmode propagating through that slice.
@@ -36,7 +38,13 @@ def solve_mode(mode_number: int,
         mu: Magnetic permeability (default 1 everywhere)
 
     Returns:
-        `{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}`
+        ```
+        {
+            'E': List[NDArray[numpy.float_]],
+            'H': List[NDArray[numpy.float_]],
+            'wavenumber': complex,
+        }
+        ```
     """
     if mu is None:
         mu = numpy.ones_like(epsilon)
@@ -97,16 +105,17 @@ def solve_mode(mode_number: int,
     return results
 
 
-def compute_source(E: fdfield_t,
-                   wavenumber: complex,
-                   omega: complex,
-                   dxes: dx_lists_t,
-                   axis: int,
-                   polarity: int,
-                   slices: Sequence[slice],
-                   epsilon: fdfield_t,
-                   mu: Optional[fdfield_t] = None,
-                   ) -> fdfield_t:
+def compute_source(
+        E: fdfield_t,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        axis: int,
+        polarity: int,
+        slices: Sequence[slice],
+        epsilon: fdfield_t,
+        mu: Optional[fdfield_t] = None,
+        ) -> fdfield_t:
     """
     Given an eigenmode obtained by `solve_mode`, returns the current source distribution
     necessary to position a unidirectional source at the slice location.
@@ -142,18 +151,21 @@ def compute_source(E: fdfield_t,
     return J
 
 
-def compute_overlap_e(E: fdfield_t,
-                      wavenumber: complex,
-                      dxes: dx_lists_t,
-                      axis: int,
-                      polarity: int,
-                      slices: Sequence[slice],
-                      ) -> fdfield_t:                 # TODO DOCS
+def compute_overlap_e(
+        E: fdfield_t,
+        wavenumber: complex,
+        dxes: dx_lists_t,
+        axis: int,
+        polarity: int,
+        slices: Sequence[slice],
+        ) -> fdfield_t:                 # TODO DOCS
     """
     Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the
     mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn)
     [assumes reflection symmetry].
 
+    TODO: add reference
+
     Args:
         E: E-field of the mode
         H: H-field of the mode (advanced by half of a Yee cell from E)
@@ -187,13 +199,14 @@ def compute_overlap_e(E: fdfield_t,
     return Etgt
 
 
-def expand_e(E: fdfield_t,
-             wavenumber: complex,
-             dxes: dx_lists_t,
-             axis: int,
-             polarity: int,
-             slices: Sequence[slice],
-             ) -> fdfield_t:
+def expand_e(
+        E: fdfield_t,
+        wavenumber: complex,
+        dxes: dx_lists_t,
+        axis: int,
+        polarity: int,
+        slices: Sequence[slice],
+        ) -> fdfield_t:
     """
     Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D
     slice where the mode was calculated to the entire domain (along the propagation

meanas/fdfd/waveguide_cyl.py

@@ -9,7 +9,7 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
 # TODO update module docs
 
 from typing import Dict, Union
-import numpy                        # type: ignore
+import numpy
 import scipy.sparse as sparse       # type: ignore
 
 from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t
@@ -17,11 +17,12 @@ from ..fdmath.operators import deriv_forward, deriv_back
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
 
-def cylindrical_operator(omega: complex,
-                         dxes: dx_lists_t,
-                         epsilon: vfdfield_t,
-                         r0: float,
-                         ) -> sparse.spmatrix:
+def cylindrical_operator(
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        r0: float,
+        ) -> sparse.spmatrix:
     """
     Cylindrical coordinate waveguide operator of the form
 
@@ -78,12 +79,13 @@ def cylindrical_operator(omega: complex,
     return op
 
 
-def solve_mode(mode_number: int,
-               omega: complex,
-               dxes: dx_lists_t,
-               epsilon: vfdfield_t,
-               r0: float,
-               ) -> Dict[str, Union[complex, fdfield_t]]:
+def solve_mode(
+        mode_number: int,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        r0: float,
+        ) -> Dict[str, Union[complex, fdfield_t]]:
     """
     TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@@ -99,7 +101,13 @@ def solve_mode(mode_number: int,
             r within the simulation domain.
 
     Returns:
-        `{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}`
+        ```
+        {
+            'E': List[NDArray[numpy.float_]],
+            'H': List[NDArray[numpy.float_]],
+            'wavenumber': complex,
+        }
+        ```
     """
 
     '''

meanas/fdmath/functional.py

@@ -5,13 +5,15 @@ Basic discrete calculus etc.
 """
 from typing import Sequence, Tuple, Optional, Callable
 
-import numpy            # type: ignore
+import numpy
+from numpy.typing import NDArray
 
 from .types import fdfield_t, fdfield_updater_t
 
 
-def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None
-                  ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
+def deriv_forward(
+        dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
+        ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
     """
     Utility operators for taking discretized derivatives (backward variant).
 
@@ -33,8 +35,9 @@ def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None
     return derivs
 
 
-def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None
-               ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
+def deriv_back(
+        dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
+        ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
     """
     Utility operators for taking discretized derivatives (forward variant).
 
@@ -56,7 +59,9 @@ def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None
     return derivs
 
 
-def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t:
+def curl_forward(
+        dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
+        ) -> fdfield_updater_t:
     """
     Curl operator for use with the E field.
 
@@ -83,7 +88,9 @@ def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_upda
     return ce_fun
 
 
-def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t:
+def curl_back(
+        dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
+        ) -> fdfield_updater_t:
     """
     Create a function which takes the backward curl of a field.
 
@@ -110,7 +117,9 @@ def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater
     return ch_fun
 
 
-def curl_forward_parts(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> Callable:
+def curl_forward_parts(
+        dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
+        ) -> Callable:
     Dx, Dy, Dz = deriv_forward(dx_e)
 
     def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:
@@ -121,7 +130,9 @@ def curl_forward_parts(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> Callab
     return mkparts_fwd
 
 
-def curl_back_parts(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> Callable:
+def curl_back_parts(
+        dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
+        ) -> Callable:
     Dx, Dy, Dz = deriv_back(dx_e)
 
     def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:

meanas/fdmath/operators.py

@@ -4,13 +4,18 @@ Matrix operators for finite difference simulations
 Basic discrete calculus etc.
 """
 from typing import Sequence, List
-import numpy                    # type: ignore
+import numpy
+from numpy.typing import NDArray
 import scipy.sparse as sparse   # type: ignore
 
 from .types import vfdfield_t
 
 
-def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> sparse.spmatrix:
+def shift_circ(
+        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 +51,11 @@ def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> spar
     return d
 
 
-def shift_with_mirror(axis: int, shape: Sequence[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.
@@ -67,7 +76,7 @@ def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1)
         raise Exception('Shift ({}) is too large for axis {} of size {}'.format(
                         shift_distance, axis, shape[axis]))
 
-    def mirrored_range(n: int, s: int) -> numpy.ndarray:
+    def mirrored_range(n: int, s: int) -> NDArray[numpy.int_]:
         v = numpy.arange(n) + s
         v = numpy.where(v >= n, 2 * n - v - 1, v)
         v = numpy.where(v < 0, - 1 - v, v)
@@ -87,7 +96,9 @@ def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1)
     return d
 
 
-def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
+def deriv_forward(
+        dx_e: Sequence[NDArray[numpy.float_]],
+        ) -> List[sparse.spmatrix]:
     """
     Utility operators for taking discretized derivatives (forward variant).
 
@@ -112,7 +123,9 @@ def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
     return Ds
 
 
-def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
+def deriv_back(
+        dx_h: Sequence[NDArray[numpy.float_]],
+        ) -> List[sparse.spmatrix]:
     """
     Utility operators for taking discretized derivatives (backward variant).
 
@@ -137,7 +150,9 @@ def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
     return Ds
 
 
-def cross(B: Sequence[sparse.spmatrix]) -> sparse.spmatrix:
+def cross(
+        B: Sequence[sparse.spmatrix],
+        ) -> sparse.spmatrix:
     """
     Cross product operator
 
@@ -203,7 +218,9 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
     return avg_forward(axis, shape).T
 
 
-def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix:
+def curl_forward(
+        dx_e: Sequence[NDArray[numpy.float_]],
+        ) -> sparse.spmatrix:
     """
     Curl operator for use with the E field.
 
@@ -217,7 +234,9 @@ def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix:
     return cross(deriv_forward(dx_e))
 
 
-def curl_back(dx_h: Sequence[numpy.ndarray]) -> sparse.spmatrix:
+def curl_back(
+        dx_h: Sequence[NDArray[numpy.float_]],
+        ) -> sparse.spmatrix:
     """
     Curl operator for use with the H field.
 

meanas/fdmath/types.py

@@ -2,31 +2,20 @@
 Types shared across multiple submodules
 """
 from typing import Sequence, Callable, MutableSequence
-import numpy            # type: ignore
+import numpy
+from numpy.typing import NDArray
 
 
 # Field types
-# TODO: figure out a better way to set the docstrings without creating actual subclasses?
-#   Probably not a big issue since they're only used for type hinting
-class fdfield_t(numpy.ndarray):
-    """
-    Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)
+fdfield_t = NDArray[numpy.float_]
+"""Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
 
-    This is actually is just an unaltered `numpy.ndarray`
-    """
-    pass
-
-class vfdfield_t(numpy.ndarray):
-    """
-    Linearized vector field (single vector of length 3*X*Y*Z)
-
-    This is actually just an unaltered `numpy.ndarray`
-    """
-    pass
+vfdfield_t = NDArray[numpy.float_]
+"""Linearized vector field (single vector of length 3*X*Y*Z)"""
 
 
-dx_lists_t = Sequence[Sequence[numpy.ndarray]]
-'''
+dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]]
+"""
  'dxes' datastructure which contains grid cell width information in the following format:
 
      [[[dx_e[0], dx_e[1], ...], [dy_e[0], ...], [dz_e[0], ...]],
@@ -34,15 +23,11 @@ dx_lists_t = Sequence[Sequence[numpy.ndarray]]
 
    where `dx_e[0]` is the x-width of the `x=0` cells, as used when calculating dE/dx,
    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`
-'''
+dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]]
+"""Mutable version of `dx_lists_t`"""
 
 
 fdfield_updater_t = Callable[..., fdfield_t]
-'''
- Convenience type for functions which take and return an fdfield_t
-'''
+"""Convenience type for functions which take and return an fdfield_t"""

meanas/fdmath/vectorization.py

@@ -5,7 +5,8 @@ Vectorized versions of the field use row-major (ie., C-style) ordering.
 """
 
 from typing import Optional, overload, Union, List
-import numpy                # type: ignore
+import numpy
+from numpy.typing import ArrayLike
 
 from .types import fdfield_t, vfdfield_t
 
@@ -15,10 +16,10 @@ def vec(f: None) -> None:
     pass
 
 @overload
-def vec(f: Union[fdfield_t, List[numpy.ndarray]]) -> vfdfield_t:
+def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t:
     pass
 
-def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield_t]:
+def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]:
     """
     Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
 
@@ -37,14 +38,14 @@ def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield
 
 
 @overload
-def unvec(v: None, shape: numpy.ndarray) -> None:
+def unvec(v: None, shape: ArrayLike) -> None:
     pass
 
 @overload
-def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t:
+def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t:
     pass
 
-def unvec(v: Optional[vfdfield_t], shape: numpy.ndarray) -> Optional[fdfield_t]:
+def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> 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

meanas/fdtd/boundaries.py

@@ -9,9 +9,10 @@ from typing import Tuple, Any, List
 from ..fdmath import fdfield_t, fdfield_updater_t
 
 
-def conducting_boundary(direction: int,
-                        polarity: int
-                        ) -> Tuple[fdfield_updater_t, fdfield_updater_t]:
+def conducting_boundary(
+        direction: int,
+        polarity: int
+        ) -> Tuple[fdfield_updater_t, fdfield_updater_t]:
     dirs = [0, 1, 2]
     if direction not in dirs:
         raise Exception('Invalid direction: {}'.format(direction))

meanas/fdtd/energy.py

@@ -1,5 +1,5 @@
 from typing import Optional, Union
-import numpy        # type: ignore
+import numpy
 
 from ..fdmath import dx_lists_t, fdfield_t
 from ..fdmath.functional import deriv_back
@@ -8,10 +8,11 @@ from ..fdmath.functional import deriv_back
 # TODO documentation
 
 
-def poynting(e: fdfield_t,
-             h: fdfield_t,
-             dxes: Optional[dx_lists_t] = None,
-             ) -> fdfield_t:
+def poynting(
+        e: fdfield_t,
+        h: fdfield_t,
+        dxes: Optional[dx_lists_t] = None,
+        ) -> fdfield_t:
     """
     Calculate the poynting vector `S` ($S$).
 
@@ -87,12 +88,13 @@ def poynting(e: fdfield_t,
     return s
 
 
-def poynting_divergence(s: Optional[fdfield_t] = None,
-                        *,
-                        e: Optional[fdfield_t] = None,
-                        h: Optional[fdfield_t] = None,
-                        dxes: Optional[dx_lists_t] = None,
-                        ) -> fdfield_t:
+def poynting_divergence(
+        s: Optional[fdfield_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.
 
@@ -124,13 +126,14 @@ def poynting_divergence(s: Optional[fdfield_t] = None,
     return ds
 
 
-def energy_hstep(e0: fdfield_t,
-                 h1: fdfield_t,
-                 e2: fdfield_t,
-                 epsilon: Optional[fdfield_t] = None,
-                 mu: Optional[fdfield_t] = None,
-                 dxes: Optional[dx_lists_t] = None,
-                 ) -> fdfield_t:
+def energy_hstep(
+        e0: fdfield_t,
+        h1: fdfield_t,
+        e2: fdfield_t,
+        epsilon: Optional[fdfield_t] = None,
+        mu: Optional[fdfield_t] = None,
+        dxes: Optional[dx_lists_t] = None,
+        ) -> fdfield_t:
     """
     Calculate energy `U` at the time of the provided H-field `h1`.
 
@@ -151,13 +154,14 @@ def energy_hstep(e0: fdfield_t,
     return u
 
 
-def energy_estep(h0: fdfield_t,
-                 e1: fdfield_t,
-                 h2: fdfield_t,
-                 epsilon: Optional[fdfield_t] = None,
-                 mu: Optional[fdfield_t] = None,
-                 dxes: Optional[dx_lists_t] = None,
-                 ) -> fdfield_t:
+def energy_estep(
+        h0: fdfield_t,
+        e1: fdfield_t,
+        h2: fdfield_t,
+        epsilon: Optional[fdfield_t] = None,
+        mu: Optional[fdfield_t] = None,
+        dxes: Optional[dx_lists_t] = None,
+        ) -> fdfield_t:
     """
     Calculate energy `U` at the time of the provided E-field `e1`.
 
@@ -178,15 +182,16 @@ def energy_estep(h0: fdfield_t,
     return u
 
 
-def delta_energy_h2e(dt: float,
-                     e0: fdfield_t,
-                     h1: fdfield_t,
-                     e2: fdfield_t,
-                     h3: fdfield_t,
-                     epsilon: Optional[fdfield_t] = None,
-                     mu: Optional[fdfield_t] = None,
-                     dxes: Optional[dx_lists_t] = None,
-                     ) -> fdfield_t:
+def delta_energy_h2e(
+        dt: float,
+        e0: fdfield_t,
+        h1: fdfield_t,
+        e2: fdfield_t,
+        h3: fdfield_t,
+        epsilon: Optional[fdfield_t] = None,
+        mu: Optional[fdfield_t] = None,
+        dxes: Optional[dx_lists_t] = None,
+        ) -> fdfield_t:
     """
     Change in energy during the half-step from `h1` to `e2`.
 
@@ -210,15 +215,16 @@ def delta_energy_h2e(dt: float,
     return du
 
 
-def delta_energy_e2h(dt: float,
-                     h0: fdfield_t,
-                     e1: fdfield_t,
-                     h2: fdfield_t,
-                     e3: fdfield_t,
-                     epsilon: Optional[fdfield_t] = None,
-                     mu: Optional[fdfield_t] = None,
-                     dxes: Optional[dx_lists_t] = None,
-                     ) -> fdfield_t:
+def delta_energy_e2h(
+        dt: float,
+        h0: fdfield_t,
+        e1: fdfield_t,
+        h2: fdfield_t,
+        e3: fdfield_t,
+        epsilon: Optional[fdfield_t] = None,
+        mu: Optional[fdfield_t] = None,
+        dxes: Optional[dx_lists_t] = None,
+        ) -> fdfield_t:
     """
     Change in energy during the half-step from `e1` to `h2`.
 
@@ -242,10 +248,11 @@ def delta_energy_e2h(dt: float,
     return du
 
 
-def delta_energy_j(j0: fdfield_t,
-                   e1: fdfield_t,
-                   dxes: Optional[dx_lists_t] = None,
-                   ) -> fdfield_t:
+def delta_energy_j(
+        j0: fdfield_t,
+        e1: fdfield_t,
+        dxes: Optional[dx_lists_t] = None,
+        ) -> fdfield_t:
     """
     Calculate
 
@@ -264,12 +271,13 @@ def delta_energy_j(j0: fdfield_t,
     return du
 
 
-def dxmul(ee: fdfield_t,
-          hh: fdfield_t,
-          epsilon: Optional[Union[fdfield_t, float]] = None,
-          mu: Optional[Union[fdfield_t, float]] = None,
-          dxes: Optional[dx_lists_t] = None
-          ) -> fdfield_t:
+def dxmul(
+        ee: fdfield_t,
+        hh: fdfield_t,
+        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
     if mu is None:

meanas/fdtd/pml.py

@@ -8,7 +8,8 @@ PML implementations
 # TODO retest pmls!
 
 from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional
-import numpy            # type: ignore
+import numpy
+from typing import NDArray
 
 from ..fdmath import fdfield_t, dx_lists_t
 from ..fdmath.functional import deriv_forward, deriv_back
@@ -61,7 +62,7 @@ def cpml_params(
     expand_slice_l[axis] = slice(None)
     expand_slice = tuple(expand_slice_l)
 
-    def par(x: numpy.ndarray) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+    def par(x: NDArray[numpy.float64]) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
         scaling = (x / thickness) ** m
         sigma = scaling * sigma_max
         kappa = 1 + scaling * (kappa_max - 1)

meanas/test/conftest.py

@@ -4,7 +4,8 @@ Test fixtures
 
 """
 from typing import Tuple, Iterable, List, Any
-import numpy        # type: ignore
+import numpy
+from numpy.typing import NDArray, ArrayLike
 import pytest       # type: ignore
 
 from .utils import PRNG
@@ -34,11 +35,12 @@ def epsilon_fg(request: FixtureRequest) -> Iterable[float]:
 
 
 @pytest.fixture(scope='module', params=['center', '000', 'random'])
-def epsilon(request: FixtureRequest,
-            shape: Tuple[int, ...],
-            epsilon_bg: float,
-            epsilon_fg: float,
-            ) -> Iterable[numpy.ndarray]:
+def epsilon(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        epsilon_bg: float,
+        epsilon_fg: float,
+        ) -> Iterable[NDArray[numpy.float64]]:
     is3d = (numpy.array(shape) == 1).sum() == 0
     if is3d:
         if request.param == '000':
@@ -72,10 +74,11 @@ def dx(request: FixtureRequest) -> Iterable[float]:
 
 
 @pytest.fixture(scope='module', params=['uniform', 'centerbig'])
-def dxes(request: FixtureRequest,
-         shape: Tuple[int, ...],
-         dx: float,
-         ) -> Iterable[List[List[numpy.ndarray]]]:
+def dxes(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        dx: float,
+        ) -> Iterable[List[List[NDArray[numpy.float64]]]]:
     if request.param == 'uniform':
         dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
     elif request.param == 'centerbig':

meanas/test/test_fdfd.py

@@ -1,7 +1,8 @@
 from typing import List, Tuple, Iterable, Optional
 import dataclasses
 import pytest       # type: ignore
-import numpy        # type: ignore
+import numpy
+from numpy.typing import NDArray, ArrayLike
 #from numpy.testing import assert_allclose, assert_array_equal
 
 from .. import fdfd
@@ -59,12 +60,12 @@ def omega(request: FixtureRequest) -> Iterable[float]:
 
 
 @pytest.fixture(params=[None])
-def pec(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
+def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
     yield request.param
 
 
 @pytest.fixture(params=[None])
-def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
+def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
     yield request.param
 
 
@@ -75,10 +76,11 @@ def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
 
 
 @pytest.fixture(params=['diag'])        # 'center'
-def j_distribution(request: FixtureRequest,
-                   shape: Tuple[int, ...],
-                   j_mag: float,
-                   ) -> Iterable[numpy.ndarray]:
+def j_distribution(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        j_mag: float,
+        ) -> Iterable[NDArray[numpy.float64]]:
     j = numpy.zeros(shape, dtype=complex)
     center_mask = numpy.zeros(shape, dtype=bool)
     center_mask[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = True
@@ -94,24 +96,25 @@ def j_distribution(request: FixtureRequest,
 @dataclasses.dataclass()
 class FDResult:
     shape: Tuple[int, ...]
-    dxes: List[List[numpy.ndarray]]
-    epsilon: numpy.ndarray
+    dxes: List[List[NDArray[numpy.float64]]]
+    epsilon: NDArray[numpy.float64]
     omega: complex
-    j: numpy.ndarray
-    e: numpy.ndarray
-    pmc: numpy.ndarray
-    pec: numpy.ndarray
+    j: NDArray[numpy.float64]
+    e: NDArray[numpy.float64]
+    pmc: Optional[NDArray[numpy.float64]]
+    pec: Optional[NDArray[numpy.float64]]
 
 
 @pytest.fixture()
-def sim(request: FixtureRequest,
+def sim(
+        request: FixtureRequest,
         shape: Tuple[int, ...],
-        epsilon: numpy.ndarray,
-        dxes: List[List[numpy.ndarray]],
-        j_distribution: numpy.ndarray,
+        epsilon: NDArray[numpy.float64],
+        dxes: List[List[NDArray[numpy.float64]]],
+        j_distribution: NDArray[numpy.float64],
         omega: float,
-        pec: Optional[numpy.ndarray],
-        pmc: Optional[numpy.ndarray],
+        pec: Optional[NDArray[numpy.float64]],
+        pmc: Optional[NDArray[numpy.float64]],
         ) -> FDResult:
     """
     Build simulation from parts

meanas/test/test_fdfd_pml.py

@@ -1,7 +1,8 @@
 from typing import Optional, Tuple, Iterable, List
 import pytest       # type: ignore
-import numpy        # type: ignore
-from numpy.testing import assert_allclose   # type: ignore
+import numpy
+from numpy.typing import NDArray, ArrayLike
+from numpy.testing import assert_allclose
 
 from .. import fdfd
 from ..fdmath import vec, unvec, dx_lists_mut
@@ -48,12 +49,12 @@ def omega(request: FixtureRequest) -> Iterable[float]:
 
 
 @pytest.fixture(params=[None])
-def pec(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
+def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
     yield request.param
 
 
 @pytest.fixture(params=[None])
-def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
+def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
     yield request.param
 
 
@@ -70,13 +71,14 @@ def src_polarity(request: FixtureRequest) -> Iterable[int]:
 
 
 @pytest.fixture()
-def j_distribution(request: FixtureRequest,
-                   shape: Tuple[int, ...],
-                   epsilon: numpy.ndarray,
-                   dxes: dx_lists_mut,
-                   omega: float,
-                   src_polarity: int,
-                   ) -> Iterable[numpy.ndarray]:
+def j_distribution(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        epsilon: NDArray[numpy.float64],
+        dxes: dx_lists_mut,
+        omega: float,
+        src_polarity: int,
+        ) -> Iterable[NDArray[numpy.float64]]:
     j = numpy.zeros(shape, dtype=complex)
 
     dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0]    # Propagation axis
@@ -108,47 +110,60 @@ def j_distribution(request: FixtureRequest,
 
 
 @pytest.fixture()
-def epsilon(request: FixtureRequest,
-            shape: Tuple[int, ...],
-            epsilon_bg: float,
-            epsilon_fg: float,
-            ) -> Iterable[numpy.ndarray]:
+def epsilon(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        epsilon_bg: float,
+        epsilon_fg: float,
+        ) -> Iterable[NDArray[numpy.float64]]:
     epsilon = numpy.full(shape, epsilon_fg, dtype=float)
     yield epsilon
 
 
 @pytest.fixture(params=['uniform'])
-def dxes(request: FixtureRequest,
-         shape: Tuple[int, ...],
-         dx: float,
-         omega: float,
-         epsilon_fg: float,
-         ) -> Iterable[List[List[numpy.ndarray]]]:
+def dxes(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        dx: float,
+        omega: float,
+        epsilon_fg: float,
+        ) -> Iterable[List[List[NDArray[numpy.float64]]]]:
     if request.param == 'uniform':
         dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
     dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0]    # Propagation axis
     for axis in (dim,):
         for polarity in (-1, 1):
-            dxes = fdfd.scpml.stretch_with_scpml(dxes, axis=axis, polarity=polarity,
-                                                 omega=omega, epsilon_effective=epsilon_fg,
-                                                 thickness=10)
+            dxes = fdfd.scpml.stretch_with_scpml(
+                dxes,
+                axis=axis,
+                polarity=polarity,
+                omega=omega,
+                epsilon_effective=epsilon_fg,
+                thickness=10,
+                )
     yield dxes
 
 
 @pytest.fixture()
-def sim(request: FixtureRequest,
+def sim(
+        request: FixtureRequest,
         shape: Tuple[int, ...],
-        epsilon: numpy.ndarray,
+        epsilon: NDArray[numpy.float64],
         dxes: dx_lists_mut,
-        j_distribution: numpy.ndarray,
+        j_distribution: NDArray[numpy.float64],
         omega: float,
-        pec: Optional[numpy.ndarray],
-        pmc: Optional[numpy.ndarray],
+        pec: Optional[NDArray[numpy.float64]],
+        pmc: Optional[NDArray[numpy.float64]],
         ) -> FDResult:
     j_vec = vec(j_distribution)
     eps_vec = vec(epsilon)
-    e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec,
-                                 matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11})
+    e_vec = fdfd.solvers.generic(
+        J=j_vec,
+        omega=omega,
+        dxes=dxes,
+        epsilon=eps_vec,
+        matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11},
+        )
     e = unvec(e_vec, shape[1:])
 
     sim = FDResult(

meanas/test/test_fdtd.py

@@ -1,8 +1,9 @@
-from typing import List, Tuple, Iterable
+from typing import List, Tuple, Iterable, Any, Dict
 import dataclasses
 import pytest       # type: ignore
-import numpy        # type: ignore
-#from numpy.testing import assert_allclose, assert_array_equal       # type: ignore
+import numpy
+from numpy.typing import NDArray, ArrayLike
+#from numpy.testing import assert_allclose, assert_array_equal
 
 from .. import fdtd
 from .utils import assert_close, assert_fields_close, PRNG
@@ -32,8 +33,10 @@ def test_initial_energy(sim: 'TDResult') -> None:
 
     dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0)
     u0 = (j0 * j0.conj() / sim.epsilon * dV).sum(axis=0)
-    args = {'dxes': sim.dxes,
-            'epsilon': sim.epsilon}
+    args: Dict[str, Any] = {
+        'dxes': sim.dxes,
+        'epsilon': sim.epsilon,
+        }
 
     # Make sure initial energy and E dot J are correct
     energy0 = fdtd.energy_estep(h0=h0, e1=e0, h2=h1, **args)
@@ -49,8 +52,10 @@ def test_energy_conservation(sim: 'TDResult') -> None:
     e0 = sim.es[0]
     j0 = sim.js[0]
     u = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes).sum()
-    args = {'dxes': sim.dxes,
-            'epsilon': sim.epsilon}
+    args: Dict[str, Any] = {
+        'dxes': sim.dxes,
+        'epsilon': sim.epsilon,
+        }
 
     for ii in range(1, 8):
         u_hstep = fdtd.energy_hstep(e0=sim.es[ii - 1], h1=sim.hs[ii], e2=sim.es[ii],     **args)
@@ -65,8 +70,10 @@ def test_energy_conservation(sim: 'TDResult') -> None:
 
 
 def test_poynting_divergence(sim: 'TDResult') -> None:
-    args = {'dxes': sim.dxes,
-            'epsilon': sim.epsilon}
+    args: Dict[str, Any] = {
+        'dxes': sim.dxes,
+        'epsilon': sim.epsilon,
+        }
 
     u_eprev = None
     for ii in range(1, 8):
@@ -96,8 +103,10 @@ def test_poynting_planes(sim: 'TDResult') -> None:
     if mask.sum() > 1:
         pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum()))
 
-    args = {'dxes': sim.dxes,
-            'epsilon': sim.epsilon}
+    args: Dict[str, Any] = {
+        'dxes': sim.dxes,
+        'epsilon': sim.epsilon,
+        }
 
     mx = numpy.roll(mask, -1, axis=0)
     my = numpy.roll(mask, -1, axis=1)
@@ -149,13 +158,13 @@ def dt(request: FixtureRequest) -> Iterable[float]:
 class TDResult:
     shape: Tuple[int, ...]
     dt: float
-    dxes: List[List[numpy.ndarray]]
-    epsilon: numpy.ndarray
-    j_distribution: numpy.ndarray
+    dxes: List[List[NDArray[numpy.float64]]]
+    epsilon: NDArray[numpy.float64]
+    j_distribution: NDArray[numpy.float64]
     j_steps: Tuple[int, ...]
-    es: List[numpy.ndarray] = dataclasses.field(default_factory=list)
-    hs: List[numpy.ndarray] = dataclasses.field(default_factory=list)
-    js: List[numpy.ndarray] = dataclasses.field(default_factory=list)
+    es: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
+    hs: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
+    js: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
 
 
 @pytest.fixture(params=[(0, 4, 8)])  # (0,)
@@ -164,10 +173,11 @@ def j_steps(request: FixtureRequest) -> Iterable[Tuple[int, ...]]:
 
 
 @pytest.fixture(params=['center', 'random'])
-def j_distribution(request: FixtureRequest,
-                   shape: Tuple[int, ...],
-                   j_mag: float,
-                   ) -> Iterable[numpy.ndarray]:
+def j_distribution(
+        request: FixtureRequest,
+        shape: Tuple[int, ...],
+        j_mag: float,
+        ) -> Iterable[NDArray[numpy.float64]]:
     j = numpy.zeros(shape)
     if request.param == 'center':
         j[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = j_mag
@@ -179,12 +189,13 @@ def j_distribution(request: FixtureRequest,
 
 
 @pytest.fixture()
-def sim(request: FixtureRequest,
+def sim(
+        request: FixtureRequest,
         shape: Tuple[int, ...],
-        epsilon: numpy.ndarray,
-        dxes: List[List[numpy.ndarray]],
+        epsilon: NDArray[numpy.float64],
+        dxes: List[List[NDArray[numpy.float64]]],
         dt: float,
-        j_distribution: numpy.ndarray,
+        j_distribution: NDArray[numpy.float64],
         j_steps: Tuple[int, ...],
         ) -> TDResult:
     is3d = (numpy.array(shape) == 1).sum() == 0

meanas/test/utils.py

@@ -1,24 +1,27 @@
 from typing import Any
-import numpy        # type: ignore
+import numpy
+from typing import ArrayLike
 
 
 PRNG = numpy.random.RandomState(12345)
 
 
-def assert_fields_close(x: numpy.ndarray,
-                        y: numpy.ndarray,
-                        *args: Any,
-                        **kwargs: Any,
-                        ) -> None:
+def assert_fields_close(
+        x: ArrayLike,
+        y: ArrayLike,
+        *args: Any,
+        **kwargs: Any,
+        ) -> None:
     numpy.testing.assert_allclose(
         x, y, verbose=False,
         err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1),
                                                        numpy.rollaxis(y, -1)), *args, **kwargs)
 
-def assert_close(x: numpy.ndarray,
-                 y: numpy.ndarray,
-                 *args: Any,
-                 **kwargs: Any,
-                 ) -> None:
+def assert_close(
+        x: ArrayLike,
+        y: ArrayLike,
+        *args: Any,
+        **kwargs: Any,
+        ) -> None:
     numpy.testing.assert_allclose(x, y, *args, **kwargs)