Major typing updates

Split field types to differentiate between complex and purely-real

Fix lots of numpy-related stuff

meanas/eigensolvers.py

@@ -11,9 +11,9 @@ import scipy.sparse.linalg as spalg   # type: ignore
 
 def power_iteration(
         operator: sparse.spmatrix,
-        guess_vector: Optional[NDArray[numpy.float64]] = None,
+        guess_vector: Optional[NDArray[numpy.complex128]] = None,
         iterations: int = 20,
-        ) -> Tuple[complex, NDArray[numpy.float64]]:
+        ) -> Tuple[complex, NDArray[numpy.complex128]]:
     """
     Use power iteration to estimate the dominant eigenvector of a matrix.
 
@@ -26,7 +26,7 @@ def power_iteration(
         (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
     """
     if numpy.any(numpy.equal(guess_vector, None)):
-        v = numpy.random.rand(operator.shape[0])
+        v = numpy.random.rand(operator.shape[0]) + 1j * numpy.random.rand(operator.shape[0])
     else:
         v = guess_vector
 
@@ -41,11 +41,11 @@ def power_iteration(
 
 def rayleigh_quotient_iteration(
         operator: Union[sparse.spmatrix, spalg.LinearOperator],
-        guess_vector: NDArray[numpy.float64],
+        guess_vector: NDArray[numpy.complex128],
         iterations: int = 40,
         tolerance: float = 1e-13,
-        solver: Optional[Callable[..., NDArray[numpy.float64]]] = None,
-        ) -> Tuple[complex, NDArray[numpy.float64]]:
+        solver: Optional[Callable[..., NDArray[numpy.complex128]]] = None,
+        ) -> Tuple[complex, NDArray[numpy.complex128]]:
     """
     Use Rayleigh quotient iteration to refine an eigenvector guess.
 
@@ -78,7 +78,7 @@ def rayleigh_quotient_iteration(
                     matvec=lambda v: eigval * v,
                     )
         if solver is None:
-            def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.float64]:
+            def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.complex128]:
                 return spalg.bicgstab(A, b)[0]
     assert(solver is not None)
 
@@ -99,7 +99,7 @@ def signed_eigensolve(
         operator: Union[sparse.spmatrix, spalg.LinearOperator],
         how_many: int,
         negative: bool = False,
-        ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
+        ) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
     """
     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, Union
+from typing import Tuple, Callable, Any, List, Optional, cast, Union, Sequence
 import logging
 import numpy
 from numpy import pi, real, trace
@@ -91,7 +91,8 @@ import scipy.optimize                   # type: ignore
 from scipy.linalg import norm           # type: ignore
 import scipy.sparse.linalg as spalg     # type: ignore
 
-from ..fdmath import fdfield_t
+from ..fdmath import fdfield_t, cfdfield_t
+
 
 logger = logging.getLogger(__name__)
 
@@ -110,10 +111,10 @@ try:
         'planner_effort': 'FFTW_PATIENT',
         }
 
-    def fftn(*args: Any, **kwargs: Any) -> NDArray[numpy.float64]:
+    def fftn(*args: Any, **kwargs: Any) -> NDArray[numpy.complex128]:
         return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args)
 
-    def ifftn(*args: Any, **kwargs: Any) -> NDArray[numpy.float64]:
+    def ifftn(*args: Any, **kwargs: Any) -> NDArray[numpy.complex128]:
         return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)
 
 except ImportError:
@@ -124,7 +125,7 @@ except ImportError:
 def generate_kmn(
         k0: ArrayLike,
         G_matrix: ArrayLike,
-        shape: ArrayLike,
+        shape: Sequence[int],
         ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
     """
     Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.
@@ -142,7 +143,7 @@ def generate_kmn(
     k0 = numpy.array(k0)
 
     Gi_grids = numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij')
-    Gi = numpy.stack(Gi_grids, axis=3)
+    Gi = numpy.moveaxis(Gi_grids, 0, -1)
 
     k_G = k0[None, None, None, :] - Gi
     k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2)
@@ -169,7 +170,7 @@ def maxwell_operator(
         G_matrix: ArrayLike,
         epsilon: fdfield_t,
         mu: Optional[fdfield_t] = None
-        ) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]:
+        ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
     """
     Generate the Maxwell operator
 
@@ -198,11 +199,11 @@ def maxwell_operator(
     shape = epsilon[0].shape + (1,)
     k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
-    epsilon = numpy.stack(epsilon, axis=3)
+    epsilon = numpy.moveaxis(epsilon, 0, -1)
     if mu is not None:
-        mu = numpy.stack(mu, axis=3)
+        mu = numpy.moveaxis(mu, 0, -1)
 
-    def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
+    def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
         """
         Maxwell operator for Bloch eigenmode simulation.
 
@@ -251,7 +252,7 @@ def hmn_2_exyz(
         k0: ArrayLike,
         G_matrix: ArrayLike,
         epsilon: fdfield_t,
-        ) -> Callable[[NDArray[numpy.float64]], fdfield_t]:
+        ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
     """
     Generate an operator which converts a vectorized spatial-frequency-space
      `h_mn` into an E-field distribution, i.e.
@@ -272,11 +273,11 @@ def hmn_2_exyz(
         Function for converting `h_mn` into `E_xyz`
     """
     shape = epsilon[0].shape + (1,)
-    epsilon = numpy.stack(epsilon, axis=3)
+    epsilon = numpy.moveaxis(epsilon, 0, -1)
 
     k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
-    def operator(h: NDArray[numpy.float64]) -> fdfield_t:
+    def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
         hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
         d_xyz = (n * hin_m
                - m * hin_n) * k_mag
@@ -291,7 +292,7 @@ def hmn_2_hxyz(
         k0: ArrayLike,
         G_matrix: ArrayLike,
         epsilon: fdfield_t
-        ) -> Callable[[NDArray[numpy.float64]], fdfield_t]:
+        ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
     """
     Generate an operator which converts a vectorized spatial-frequency-space
      `h_mn` into an H-field distribution, i.e.
@@ -314,7 +315,7 @@ def hmn_2_hxyz(
     shape = epsilon[0].shape + (1,)
     _k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
-    def operator(h: NDArray[numpy.float64]) -> fdfield_t:
+    def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
         hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
         h_xyz = (m * hin_m
                + n * hin_n)
@@ -328,7 +329,7 @@ def inverse_maxwell_operator_approx(
         G_matrix: ArrayLike,
         epsilon: fdfield_t,
         mu: Optional[fdfield_t] = None,
-        ) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]:
+        ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
     """
     Generate an approximate inverse of the Maxwell operator,
 
@@ -350,14 +351,13 @@ def inverse_maxwell_operator_approx(
         Function which applies the approximate inverse of the maxwell operator to `h_mn`.
     """
     shape = epsilon[0].shape + (1,)
-    epsilon = numpy.stack(epsilon, axis=3)
+    epsilon = numpy.moveaxis(epsilon, 0, -1)
+    if mu is not None:
+        mu = numpy.moveaxis(mu, 0, -1)
 
     k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
-    if mu is not None:
-        mu = numpy.stack(mu, axis=3)
-
-    def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
+    def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
         """
         Approximate inverse Maxwell operator for Bloch eigenmode simulation.
 
@@ -462,7 +462,7 @@ def eigsolve(
         tolerance: float = 1e-20,
         max_iters: int = 10000,
         reset_iters: int = 100,
-        ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
+        ) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
     """
     Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
      k0 of the specified structure.
@@ -697,11 +697,11 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to
 '''
 
 def _rtrace_AtB(
-        A: NDArray[numpy.float64],
-        B: Union[NDArray[numpy.float64], float],
+        A: NDArray[numpy.complex128],
+        B: Union[NDArray[numpy.complex128], float],
         ) -> float:
     return real(numpy.sum(A.conj() * B))
 
-def _symmetrize(A: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
+def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
     return (A + A.conj().T) * 0.5
 

meanas/fdfd/farfield.py

@@ -6,12 +6,12 @@ import numpy
 from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
 from numpy import pi
 
-from ..fdmath import fdfield_t
+from ..fdmath import cfdfield_t
 
 
 def near_to_farfield(
-        E_near: fdfield_t,
-        H_near: fdfield_t,
+        E_near: cfdfield_t,
+        H_near: cfdfield_t,
         dx: float,
         dy: float,
         padded_size: List[int] = None
@@ -122,8 +122,8 @@ def near_to_farfield(
 
 
 def far_to_nearfield(
-        E_far: fdfield_t,
-        H_far: fdfield_t,
+        E_far: cfdfield_t,
+        H_far: cfdfield_t,
         dkx: float,
         dky: float,
         padded_size: List[int] = None

meanas/fdfd/functional.py

@@ -2,13 +2,13 @@
 Functional versions of many FDFD operators. These can be useful for performing
  FDFD calculations without needing to construct large matrices in memory.
 
-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)
+The functions generated here expect `cfdfield_t` inputs with shape (3, X, Y, Z),
+e.g. E = [E_x, E_y, E_z] where each (complex) component has shape (X, Y, Z)
 """
 from typing import Callable, Tuple, Optional
 import numpy
 
-from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
+from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t
 from ..fdmath.functional import curl_forward, curl_back
 
 
@@ -19,8 +19,8 @@ def e_full(
         omega: complex,
         dxes: dx_lists_t,
         epsilon: fdfield_t,
-        mu: fdfield_t = None
-        ) -> fdfield_updater_t:
+        mu: Optional[fdfield_t] = None
+        ) -> cfdfield_updater_t:
     """
     Wave operator for use with E-field. See `operators.e_full` for details.
 
@@ -37,13 +37,13 @@ def e_full(
     ch = curl_back(dxes[1])
     ce = curl_forward(dxes[0])
 
-    def op_1(e: fdfield_t) -> fdfield_t:
+    def op_1(e: cfdfield_t) -> cfdfield_t:
         curls = ch(ce(e))
-        return curls - omega ** 2 * epsilon * e         # type: ignore      # issues with numpy/mypy
+        return curls - omega ** 2 * epsilon * e
 
-    def op_mu(e: fdfield_t) -> fdfield_t:
-        curls = ch(mu * ce(e))
-        return curls - omega ** 2 * epsilon * e         # type: ignore      # issues with numpy/mypy
+    def op_mu(e: cfdfield_t) -> cfdfield_t:
+        curls = ch(mu * ce(e))          # type: ignore   # mu = None ok because we don't return the function
+        return curls - omega ** 2 * epsilon * e
 
     if numpy.any(numpy.equal(mu, None)):
         return op_1
@@ -56,7 +56,7 @@ def eh_full(
         dxes: dx_lists_t,
         epsilon: fdfield_t,
         mu: fdfield_t = None
-        ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]:
+        ) -> Callable[[cfdfield_t, cfdfield_t], Tuple[cfdfield_t, cfdfield_t]]:
     """
     Wave operator for full (both E and H) field representation.
     See `operators.eh_full`.
@@ -74,13 +74,13 @@ def eh_full(
     ch = curl_back(dxes[1])
     ce = curl_forward(dxes[0])
 
-    def op_1(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]:
+    def op_1(e: cfdfield_t, h: cfdfield_t) -> Tuple[cfdfield_t, cfdfield_t]:
         return (ch(h) - 1j * omega * epsilon * e,
-                ce(e) + 1j * omega * h)                 # type: ignore    # issues with numpy/mypy
+                ce(e) + 1j * omega * h)
 
-    def op_mu(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]:
+    def op_mu(e: cfdfield_t, h: cfdfield_t) -> Tuple[cfdfield_t, cfdfield_t]:
         return (ch(h) - 1j * omega * epsilon * e,
-                ce(e) + 1j * omega * mu * h)            # type: ignore    # issues with numpy/mypy
+                ce(e) + 1j * omega * mu * h)            # type: ignore   # mu=None ok
 
     if numpy.any(numpy.equal(mu, None)):
         return op_1
@@ -92,7 +92,7 @@ def e2h(
         omega: complex,
         dxes: dx_lists_t,
         mu: Optional[fdfield_t] = None,
-        ) -> fdfield_updater_t:
+        ) -> cfdfield_updater_t:
     """
     Utility operator for converting the `E` field into the `H` field.
     For use with `e_full` -- assumes that there is no magnetic current `M`.
@@ -108,11 +108,11 @@ def e2h(
     """
     ce = curl_forward(dxes[0])
 
-    def e2h_1_1(e: fdfield_t) -> fdfield_t:
-        return ce(e) / (-1j * omega)            # type: ignore    # issues with numpy/mypy
+    def e2h_1_1(e: cfdfield_t) -> cfdfield_t:
+        return ce(e) / (-1j * omega)
 
-    def e2h_mu(e: fdfield_t) -> fdfield_t:
-        return ce(e) / (-1j * omega * mu)       # type: ignore    # issues with numpy/mypy
+    def e2h_mu(e: cfdfield_t) -> cfdfield_t:
+        return ce(e) / (-1j * omega * mu)       # type: ignore   # mu=None ok
 
     if numpy.any(numpy.equal(mu, None)):
         return e2h_1_1
@@ -124,7 +124,7 @@ def m2j(
         omega: complex,
         dxes: dx_lists_t,
         mu: Optional[fdfield_t] = None,
-        ) -> fdfield_updater_t:
+        ) -> cfdfield_updater_t:
     """
     Utility operator for converting magnetic current `M` distribution
     into equivalent electric current distribution `J`.
@@ -141,13 +141,13 @@ def m2j(
     """
     ch = curl_back(dxes[1])
 
-    def m2j_mu(m: fdfield_t) -> fdfield_t:
-        J = ch(m / mu) / (-1j * omega)
-        return J                          # type: ignore    # issues with numpy/mypy
+    def m2j_mu(m: cfdfield_t) -> cfdfield_t:
+        J = ch(m / mu) / (-1j * omega)          # type: ignore  # mu=None ok
+        return J
 
-    def m2j_1(m: fdfield_t) -> fdfield_t:
+    def m2j_1(m: cfdfield_t) -> cfdfield_t:
         J = ch(m) / (-1j * omega)
-        return J                          # type: ignore    # issues with numpy/mypy
+        return J
 
     if numpy.any(numpy.equal(mu, None)):
         return m2j_1
@@ -161,7 +161,7 @@ def e_tfsf_source(
         dxes: dx_lists_t,
         epsilon: fdfield_t,
         mu: Optional[fdfield_t] = None,
-        ) -> fdfield_updater_t:
+        ) -> cfdfield_updater_t:
     """
     Operator that turns an E-field distribution into a total-field/scattered-field
     (TFSF) source.
@@ -182,13 +182,13 @@ def e_tfsf_source(
     # TODO documentation
     A = e_full(omega, dxes, epsilon, mu)
 
-    def op(e: fdfield_t) -> fdfield_t:
+    def op(e: cfdfield_t) -> cfdfield_t:
         neg_iwj = A(TF_region * e) - TF_region * A(e)
         return neg_iwj / (-1j * omega)
     return op
 
 
-def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdfield_t]:
+def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield_t, cfdfield_t], cfdfield_t]:
     """
     Generates a function that takes the single-frequency `E` and `H` fields
     and calculates the cross product `E` x `H` = $E \\times H$ as required
@@ -210,7 +210,7 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdf
     Returns:
         Function `f` that returns E x H as required for the poynting vector.
     """
-    def exh(e: fdfield_t, h: fdfield_t) -> fdfield_t:
+    def exh(e: cfdfield_t, h: cfdfield_t) -> cfdfield_t:
         s = numpy.empty_like(e)
         ex = e[0] * dxes[0][0][:, None, None]
         ey = e[1] * dxes[0][1][None, :, None]

meanas/fdfd/operators.py

@@ -31,7 +31,7 @@ from typing import Tuple, Optional
 import numpy
 import scipy.sparse as sparse       # type: ignore
 
-from ..fdmath import vec, dx_lists_t, vfdfield_t
+from ..fdmath import vec, dx_lists_t, vfdfield_t, vcfdfield_t
 from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl_back
 
 
@@ -91,7 +91,7 @@ def e_full(
     if numpy.any(numpy.equal(mu, None)):
         m_div = sparse.eye(epsilon.size)
     else:
-        m_div = sparse.diags(1 / mu)                # type: ignore  # checked mu is not None
+        m_div = sparse.diags(1 / mu)
 
     op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe
     return op
@@ -275,7 +275,7 @@ def e2h(
     op = curl_forward(dxes[0]) / (-1j * omega)
 
     if not numpy.any(numpy.equal(mu, None)):
-        op = sparse.diags(1 / mu) @ op              # type: ignore  # checked mu is not None
+        op = sparse.diags(1 / mu) @ op
 
     if not numpy.any(numpy.equal(pmc, None)):
         op = sparse.diags(numpy.where(pmc, 0, 1)) @ op
@@ -303,12 +303,12 @@ def m2j(
     op = curl_back(dxes[1]) / (1j * omega)
 
     if not numpy.any(numpy.equal(mu, None)):
-        op = op @ sparse.diags(1 / mu)            # type: ignore  # checked mu is not None
+        op = op @ sparse.diags(1 / mu)
 
     return op
 
 
-def poynting_e_cross(e: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
+def poynting_e_cross(e: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     """
     Operator for computing the Poynting vector, containing the
     (E x) portion of the Poynting vector.
@@ -336,7 +336,7 @@ def poynting_e_cross(e: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     return P
 
 
-def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
+def poynting_h_cross(h: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     """
     Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector.
 

meanas/fdfd/scpml.py

@@ -11,7 +11,7 @@ from numpy.typing import ArrayLike, NDArray
 __author__ = 'Jan Petykiewicz'
 
 
-s_function_t = Callable[[float], float]
+s_function_t = Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]
 """Typedef for s-functions, see `prepare_s_function()`"""
 
 
@@ -39,8 +39,8 @@ def prepare_s_function(
 
 
 def uniform_grid_scpml(
-        shape: ArrayLike,    # ints
-        thicknesses: ArrayLike,  # ints
+        shape: Sequence[int],
+        thicknesses: Sequence[int],
         omega: float,
         epsilon_effective: float = 1.0,
         s_function: Optional[s_function_t] = None,
@@ -70,12 +70,11 @@ def uniform_grid_scpml(
     if s_function is None:
         s_function = prepare_s_function()
 
+    shape = tuple(shape)
+    thicknesses = tuple(thicknesses)
+
     # Normalized distance to nearest boundary
-    def ll(
-            u: NDArray[numpy.float64],
-            n: NDArray[numpy.float64],
-            t: NDArray[numpy.float64],
-            ) -> NDArray[numpy.float64]:
+    def ll(u: NDArray[numpy.float64], n: int, t: int) -> NDArray[numpy.float64]:
         return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t
 
     dx_a = [numpy.array(numpy.inf)] * 3

meanas/fdfd/solvers.py

@@ -10,7 +10,7 @@ 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
+from ..fdmath import dx_lists_t, vfdfield_t, vcfdfield_t
 from . import operators
 
 
@@ -65,7 +65,7 @@ def _scipy_qmr(
 def generic(
         omega: complex,
         dxes: dx_lists_t,
-        J: vfdfield_t,
+        J: vcfdfield_t,
         epsilon: vfdfield_t,
         mu: vfdfield_t = None,
         pec: vfdfield_t = None,
@@ -73,7 +73,7 @@ def generic(
         adjoint: bool = False,
         matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
         matrix_solver_opts: Optional[Dict[str, Any]] = None,
-        ) -> vfdfield_t:
+        ) -> vcfdfield_t:
     """
     Conjugate gradient FDFD solver using CSR sparse matrices.
 

meanas/fdfd/waveguide_2d.py

@@ -185,7 +185,7 @@ from numpy.linalg import norm
 import scipy.sparse as sparse       # type: ignore
 
 from ..fdmath.operators import deriv_forward, deriv_back, cross
-from ..fdmath import unvec, dx_lists_t, vfdfield_t
+from ..fdmath import unvec, dx_lists_t, vfdfield_t, vcfdfield_t
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
 
@@ -335,7 +335,7 @@ def normalized_fields_e(
         epsilon: vfdfield_t,
         mu: Optional[vfdfield_t] = None,
         prop_phase: float = 0,
-        ) -> Tuple[vfdfield_t, vfdfield_t]:
+        ) -> Tuple[vcfdfield_t, vcfdfield_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.
@@ -370,7 +370,7 @@ def normalized_fields_h(
         epsilon: vfdfield_t,
         mu: Optional[vfdfield_t] = None,
         prop_phase: float = 0,
-        ) -> Tuple[vfdfield_t, vfdfield_t]:
+        ) -> Tuple[vcfdfield_t, vcfdfield_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.
@@ -398,14 +398,14 @@ def normalized_fields_h(
 
 
 def _normalized_fields(
-        e: ArrayLike,
-        h: ArrayLike,
+        e: vcfdfield_t,
+        h: vcfdfield_t,
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         mu: Optional[vfdfield_t] = None,
         prop_phase: float = 0,
-        ) -> Tuple[vfdfield_t, vfdfield_t]:
+        ) -> Tuple[vcfdfield_t, vcfdfield_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)]
@@ -581,7 +581,7 @@ def e2h(
     """
     op = curl_e(wavenumber, dxes) / (-1j * omega)
     if not numpy.any(numpy.equal(mu, None)):
-        op = sparse.diags(1 / mu) @ op          # type: ignore   # checked that mu is not None
+        op = sparse.diags(1 / mu) @ op
     return op
 
 
@@ -649,7 +649,7 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
 
 
 def h_err(
-        h: vfdfield_t,
+        h: vcfdfield_t,
         wavenumber: complex,
         omega: complex,
         dxes: dx_lists_t,
@@ -684,12 +684,12 @@ def h_err(
 
 
 def e_err(
-        e: vfdfield_t,
+        e: vcfdfield_t,
         wavenumber: complex,
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
-        mu: vfdfield_t = Optional[None]
+        mu: Optional[vfdfield_t] = None,
         ) -> float:
     """
     Calculates the relative error in the E field
@@ -711,7 +711,7 @@ def e_err(
     if numpy.any(numpy.equal(mu, None)):
         op = ch @ ce @ e - omega ** 2 * (epsilon * e)
     else:
-        mu_inv = sparse.diags(1 / mu)          # type: ignore   # checked that mu is not None
+        mu_inv = sparse.diags(1 / mu)
         op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
 
     return norm(op) / norm(e)
@@ -724,7 +724,7 @@ def solve_modes(
         epsilon: vfdfield_t,
         mu: vfdfield_t = None,
         mode_margin: int = 2,
-        ) -> Tuple[NDArray[numpy.float64], List[complex]]:
+        ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.complex128]]:
     """
     Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
 
@@ -771,7 +771,7 @@ def solve_mode(
         mode_number: int,
         *args: Any,
         **kwargs: Any,
-        ) -> Tuple[vfdfield_t, complex]:
+        ) -> Tuple[vcfdfield_t, complex]:
     """
     Wrapper around `solve_modes()` that solves for a single mode.
 

meanas/fdfd/waveguide_3d.py

@@ -8,7 +8,7 @@ from typing import Dict, Optional, Sequence, Union, Any
 import numpy
 from numpy.typing import NDArray
 
-from ..fdmath import vec, unvec, dx_lists_t, fdfield_t
+from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, cfdfield_t
 from . import operators, waveguide_2d
 
 
@@ -106,7 +106,7 @@ def solve_mode(
 
 
 def compute_source(
-        E: fdfield_t,
+        E: cfdfield_t,
         wavenumber: complex,
         omega: complex,
         dxes: dx_lists_t,
@@ -115,7 +115,7 @@ def compute_source(
         slices: Sequence[slice],
         epsilon: fdfield_t,
         mu: Optional[fdfield_t] = None,
-        ) -> fdfield_t:
+        ) -> cfdfield_t:
     """
     Given an eigenmode obtained by `solve_mode`, returns the current source distribution
     necessary to position a unidirectional source at the slice location.
@@ -152,13 +152,13 @@ def compute_source(
 
 
 def compute_overlap_e(
-        E: fdfield_t,
+        E: cfdfield_t,
         wavenumber: complex,
         dxes: dx_lists_t,
         axis: int,
         polarity: int,
         slices: Sequence[slice],
-        ) -> fdfield_t:                 # TODO DOCS
+        ) -> cfdfield_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)
@@ -200,13 +200,13 @@ def compute_overlap_e(
 
 
 def expand_e(
-        E: fdfield_t,
+        E: cfdfield_t,
         wavenumber: complex,
         dxes: dx_lists_t,
         axis: int,
         polarity: int,
         slices: Sequence[slice],
-        ) -> fdfield_t:
+        ) -> cfdfield_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

@@ -12,7 +12,7 @@ from typing import Dict, Union
 import numpy
 import scipy.sparse as sparse       # type: ignore
 
-from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t
+from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t, cfdfield_t
 from ..fdmath.operators import deriv_forward, deriv_back
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
@@ -85,7 +85,7 @@ def solve_mode(
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         r0: float,
-        ) -> Dict[str, Union[complex, fdfield_t]]:
+        ) -> Dict[str, Union[complex, cfdfield_t]]:
     """
     TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@@ -103,8 +103,8 @@ def solve_mode(
     Returns:
         ```
         {
-            'E': List[NDArray[numpy.float_]],
-            'H': List[NDArray[numpy.float_]],
+            'E': List[NDArray[numpy.complex_]],
+            'H': List[NDArray[numpy.complex_]],
             'wavenumber': complex,
         }
         ```

meanas/fdmath/__init__.py

@@ -741,7 +741,8 @@ 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, dx_lists_mut, fdfield_updater_t
+from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t, dx_lists_t, dx_lists_mut
+from .types import fdfield_updater_t, cfdfield_updater_t
 from .vectorization import vec, unvec
 from . import operators, functional, types, vectorization
 

meanas/fdmath/functional.py

@@ -24,7 +24,7 @@ def deriv_forward(
     Returns:
         List of functions for taking forward derivatives along each axis.
     """
-    if dx_e:
+    if dx_e is not 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, :])
@@ -48,7 +48,7 @@ def deriv_back(
     Returns:
         List of functions for taking forward derivatives along each axis.
     """
-    if dx_h:
+    if dx_h is not 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, :])
@@ -122,7 +122,7 @@ def curl_forward_parts(
         ) -> Callable:
     Dx, Dy, Dz = deriv_forward(dx_e)
 
-    def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:
+    def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, fdfield_t], ...]:
         return ((-Dz(e[1]),  Dy(e[2])),
                 ( Dz(e[0]), -Dx(e[2])),
                 (-Dy(e[0]),  Dx(e[1])))
@@ -135,7 +135,7 @@ def curl_back_parts(
         ) -> Callable:
     Dx, Dy, Dz = deriv_back(dx_h)
 
-    def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:
+    def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, fdfield_t], ...]:
         return ((-Dz(h[1]),  Dy(h[2])),
                 ( Dz(h[0]), -Dx(h[2])),
                 (-Dy(h[0]),  Dx(h[1])))

meanas/fdmath/types.py

@@ -13,6 +13,12 @@ fdfield_t = NDArray[numpy.float_]
 vfdfield_t = NDArray[numpy.float_]
 """Linearized vector field (single vector of length 3*X*Y*Z)"""
 
+cfdfield_t = NDArray[numpy.complex_]
+"""Complex vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
+
+vcfdfield_t = NDArray[numpy.complex_]
+"""Linearized complex vector field (single vector of length 3*X*Y*Z)"""
+
 
 dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]]
 """
@@ -31,3 +37,6 @@ dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]]
 
 fdfield_updater_t = Callable[..., fdfield_t]
 """Convenience type for functions which take and return an fdfield_t"""
+
+cfdfield_updater_t = Callable[..., cfdfield_t]
+"""Convenience type for functions which take and return an cfdfield_t"""

meanas/fdmath/vectorization.py

@@ -4,11 +4,11 @@ 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 Optional, overload, Union, List
+from typing import Optional, overload, Union, Sequence
 import numpy
 from numpy.typing import ArrayLike
 
-from .types import fdfield_t, vfdfield_t
+from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t
 
 
 @overload
@@ -16,10 +16,18 @@ def vec(f: None) -> None:
     pass
 
 @overload
-def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t:
+def vec(f: fdfield_t) -> vfdfield_t:
     pass
 
-def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]:
+@overload
+def vec(f: cfdfield_t) -> vcfdfield_t:
+    pass
+
+@overload
+def vec(f: ArrayLike) -> Union[vfdfield_t, vcfdfield_t]:
+    pass
+
+def vec(f: Union[fdfield_t, cfdfield_t, ArrayLike, None]) -> Union[vfdfield_t, vcfdfield_t, None]:
     """
     Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
 
@@ -38,14 +46,18 @@ def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]:
 
 
 @overload
-def unvec(v: None, shape: ArrayLike) -> None:
+def unvec(v: None, shape: Sequence[int]) -> None:
     pass
 
 @overload
-def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t:
+def unvec(v: vfdfield_t, shape: Sequence[int]) -> fdfield_t:
     pass
 
-def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> Optional[fdfield_t]:
+@overload
+def unvec(v: vcfdfield_t, shape: Sequence[int]) -> cfdfield_t:
+    pass
+
+def unvec(v: Union[vfdfield_t, vcfdfield_t, None], shape: Sequence[int]) -> Union[fdfield_t, cfdfield_t, None]:
     """
     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
@@ -62,5 +74,5 @@ def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> Optional[fdfield_t]:
     """
     if numpy.any(numpy.equal(v, None)):
         return None
-    return v.reshape((3, *shape), order='C')            # type: ignore  # already check v is not None
+    return v.reshape((3, *shape), order='C')
 

meanas/fdtd/pml.py

@@ -8,8 +8,9 @@ PML implementations
 # TODO retest pmls!
 
 from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional
+from copy import deepcopy
 import numpy
-from typing import NDArray
+from numpy.typing import NDArray, DTypeLike
 
 from ..fdmath import fdfield_t, dx_lists_t
 from ..fdmath.functional import deriv_forward, deriv_back
@@ -97,34 +98,38 @@ def updates_with_cpml(
          dxes: dx_lists_t,
          epsilon: fdfield_t,
          *,
-         dtype: numpy.dtype = numpy.float32,
-         ) -> Tuple[Callable[[fdfield_t, fdfield_t], None],
-                    Callable[[fdfield_t, fdfield_t], None]]:
+         dtype: DTypeLike = numpy.float32,
+         ) -> Tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None],
+                    Callable[[fdfield_t, fdfield_t, fdfield_t], None]]:
 
     Dfx, Dfy, Dfz = deriv_forward(dxes[1])
     Dbx, Dby, Dbz = deriv_back(dxes[1])
 
-    psi_E = [[None, None], [None, None], [None, None]]
-    psi_H = [[None, None], [None, None], [None, None]]
-    params_E = [[None, None], [None, None], [None, None]]
-    params_H = [[None, None], [None, None], [None, None]]
+
+    psi_E: List[List[Tuple[Any, Any]]] = [[(None, None) for _ in range(2)] for _ in range(3)]
+    psi_H: List[List[Tuple[Any, Any]]] = deepcopy(psi_E)
+    params_E: List[List[Tuple[Any, Any, Any, Any]]] = [[(None, None, None, None) for _ in range(2)] for _ in range(3)]
+    params_H: List[List[Tuple[Any, Any, Any, Any]]] = deepcopy(params_E)
 
     for axis in range(3):
         for pp, polarity in enumerate((-1, 1)):
-            if cpml_params[axis][pp] is None:
+            cpml_param = cpml_params[axis][pp]
+            if cpml_param is None:
                 psi_E[axis][pp] = (None, None)
                 psi_H[axis][pp] = (None, None)
                 continue
 
-            cpml_param = cpml_params[axis][pp]
-
             region = cpml_param['region']
             region_shape = epsilon[0][region].shape
 
-            psi_E[axis][pp] = (numpy.zeros(region_shape, dtype=dtype),
-                                    numpy.zeros(region_shape, dtype=dtype))
-            psi_H[axis][pp] = (numpy.zeros(region_shape, dtype=dtype),
-                                    numpy.zeros(region_shape, dtype=dtype))
+            psi_E[axis][pp] = (
+                numpy.zeros(region_shape, dtype=dtype),
+                numpy.zeros(region_shape, dtype=dtype),
+                )
+            psi_H[axis][pp] = (
+                numpy.zeros(region_shape, dtype=dtype),
+                numpy.zeros(region_shape, dtype=dtype),
+                )
             params_E[axis][pp] = cpml_param['param_e'] + (region,)
             params_H[axis][pp] = cpml_param['param_h'] + (region,)
 
@@ -132,7 +137,11 @@ def updates_with_cpml(
     pE = numpy.empty_like(epsilon, dtype=dtype)
     pH = numpy.empty_like(epsilon, dtype=dtype)
 
-    def update_E(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t) -> None:
+    def update_E(
+            e: fdfield_t,
+            h: fdfield_t,
+            epsilon: fdfield_t,
+            ) -> None:
         dyHx = Dby(h[0])
         dzHx = Dbz(h[0])
         dxHy = Dbx(h[1])
@@ -175,7 +184,11 @@ def updates_with_cpml(
         e[2] += dt / epsilon[2] * (dxHy - dyHx + pE[2])
 
 
-    def update_H(e: fdfield_t, h: fdfield_t, mu: fdfield_t = (1, 1, 1)) -> None:
+    def update_H(
+            e: fdfield_t,
+            h: fdfield_t,
+            mu: fdfield_t = numpy.ones(3),
+            ) -> None:
         dyEx = Dfy(e[0])
         dzEx = Dfz(e[0])
         dxEy = Dfx(e[1])

meanas/test/test_fdfd.py

@@ -99,8 +99,8 @@ class FDResult:
     dxes: List[List[NDArray[numpy.float64]]]
     epsilon: NDArray[numpy.float64]
     omega: complex
-    j: NDArray[numpy.float64]
-    e: NDArray[numpy.float64]
+    j: NDArray[numpy.complex128]
+    e: NDArray[numpy.complex128]
     pmc: Optional[NDArray[numpy.float64]]
     pec: Optional[NDArray[numpy.float64]]
 
@@ -111,7 +111,7 @@ def sim(
         shape: Tuple[int, ...],
         epsilon: NDArray[numpy.float64],
         dxes: List[List[NDArray[numpy.float64]]],
-        j_distribution: NDArray[numpy.float64],
+        j_distribution: NDArray[numpy.complex128],
         omega: float,
         pec: Optional[NDArray[numpy.float64]],
         pmc: Optional[NDArray[numpy.float64]],
@@ -134,8 +134,13 @@ def sim(
 
     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_fdfd_pml.py

@@ -78,7 +78,7 @@ def j_distribution(
         dxes: dx_lists_mut,
         omega: float,
         src_polarity: int,
-        ) -> Iterable[NDArray[numpy.float64]]:
+        ) -> Iterable[NDArray[numpy.complex128]]:
     j = numpy.zeros(shape, dtype=complex)
 
     dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0]    # Propagation axis
@@ -150,7 +150,7 @@ def sim(
         shape: Tuple[int, ...],
         epsilon: NDArray[numpy.float64],
         dxes: dx_lists_mut,
-        j_distribution: NDArray[numpy.float64],
+        j_distribution: NDArray[numpy.complex128],
         omega: float,
         pec: Optional[NDArray[numpy.float64]],
         pmc: Optional[NDArray[numpy.float64]],

meanas/test/utils.py

@@ -1,14 +1,15 @@
 from typing import Any
+
 import numpy
-from typing import ArrayLike
+from numpy.typing import ArrayLike, NDArray
 
 
 PRNG = numpy.random.RandomState(12345)
 
 
 def assert_fields_close(
-        x: ArrayLike,
-        y: ArrayLike,
+        x: NDArray,
+        y: NDArray,
         *args: Any,
         **kwargs: Any,
         ) -> None:
@@ -18,8 +19,8 @@ def assert_fields_close(
                                                        numpy.rollaxis(y, -1)), *args, **kwargs)
 
 def assert_close(
-        x: ArrayLike,
-        y: ArrayLike,
+        x: NDArray,
+        y: NDArray,
         *args: Any,
         **kwargs: Any,
         ) -> None: