fixup! add sensitivity calculation

meanas/__init__.py

@@ -11,8 +11,7 @@ __author__ = 'Jan Petykiewicz'
 
 
 try:
-    readme_path = pathlib.Path(__file__).parent / 'README.md'
+    with open(pathlib.Path(__file__).parent / 'README.md', 'r') as f:
-    with readme_path.open('r') as f:
         __doc__ = f.read()
 except Exception:
     pass

meanas/eigensolvers.py

@@ -1,12 +1,12 @@
 """
 Solvers for eigenvalue / eigenvector problems
 """
-from collections.abc import Callable
+from typing import Callable
 import numpy
 from numpy.typing import NDArray, ArrayLike
 from numpy.linalg import norm
-from scipy import sparse
+from scipy import sparse              # type: ignore
-import scipy.sparse.linalg as spalg
+import scipy.sparse.linalg as spalg   # type: ignore
 
 
 def power_iteration(
@@ -25,9 +25,8 @@ def power_iteration(
     Returns:
         (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
     """
-    rng = numpy.random.default_rng()
     if guess_vector is None:
-        v = rng.random(operator.shape[0]) + 1j * rng.random(operator.shape[0])
+        v = numpy.random.rand(operator.shape[0]) + 1j * numpy.random.rand(operator.shape[0])
     else:
         v = guess_vector
 

meanas/fdfd/__init__.py

@@ -91,12 +91,5 @@ $$
 
 
 """
-from . import (
+from . import solvers, operators, functional, scpml, waveguide_2d, waveguide_3d
-    solvers as solvers,
-    operators as operators,
-    functional as functional,
-    scpml as scpml,
-    waveguide_2d as waveguide_2d,
-    waveguide_3d as waveguide_3d,
-    )
 # from . import farfield, bloch TODO

meanas/fdfd/bloch.py

@@ -94,17 +94,16 @@ This module contains functions for generating and solving the
 
 """
 
-from typing import Any, cast
+from typing import Callable, Any, cast, Sequence
-from collections.abc import Callable, Sequence
 import logging
 import numpy
 from numpy import pi, real, trace
 from numpy.fft import fftfreq
 from numpy.typing import NDArray, ArrayLike
-import scipy
+import scipy                            # type: ignore
-import scipy.optimize
+import scipy.optimize                   # type: ignore
-from scipy.linalg import norm
+from scipy.linalg import norm           # type: ignore
-import scipy.sparse.linalg as spalg
+import scipy.sparse.linalg as spalg     # type: ignore
 
 from ..fdmath import fdfield_t, cfdfield_t
 
@@ -115,6 +114,7 @@ logger = logging.getLogger(__name__)
 try:
     import pyfftw.interfaces.numpy_fft  # type: ignore
     import pyfftw.interfaces            # type: ignore
+    import multiprocessing
     logger.info('Using pyfftw')
 
     pyfftw.interfaces.cache.enable()
@@ -155,7 +155,7 @@ def generate_kmn(
             All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
     """
     k0 = numpy.array(k0)
-    G_matrix = numpy.asarray(G_matrix)
+    G_matrix = numpy.array(G_matrix, copy=False)
 
     Gi_grids = numpy.array(numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij'))
     Gi = numpy.moveaxis(Gi_grids, 0, -1)
@@ -232,7 +232,7 @@ def maxwell_operator(
             Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)), returned
             and overwritten in-place of `h`.
         """
-        hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
 
         #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
 
@@ -303,12 +303,12 @@ def hmn_2_exyz(
     k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
     def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
-        hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
         d_xyz = (n * hin_m
                - m * hin_n) * k_mag         # noqa: E128
 
         # divide by epsilon
-        return numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)
+        return numpy.array([ei for ei in numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)])         # TODO avoid copy
 
     return operator
 
@@ -341,7 +341,7 @@ def hmn_2_hxyz(
     _k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
     def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
-        hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
         h_xyz = (m * hin_m
                + n * hin_n)     # noqa: E128
         return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])
@@ -394,7 +394,7 @@ def inverse_maxwell_operator_approx(
         Returns:
             Raveled ik x conv(eps_k, ik x conv(mu_k, h_mn))
         """
-        hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
 
         #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
 
@@ -538,7 +538,7 @@ def eigsolve(
         `(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the
         vector `eigenvectors[i, :]`
     """
-    k0 = numpy.asarray(k0)
+    k0 = numpy.array(k0, copy=False)
 
     h_size = 2 * epsilon[0].size
 
@@ -561,12 +561,11 @@ def eigsolve(
     prev_theta = 0.5
     D = numpy.zeros(shape=y_shape, dtype=complex)
 
-    rng = numpy.random.default_rng()
     Z: NDArray[numpy.complex128]
     if y0 is None:
-        Z = rng.random(y_shape) + 1j * rng.random(y_shape)
+        Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
     else:
-        Z = numpy.asarray(y0).T
+        Z = numpy.array(y0, copy=False).T
 
     while True:
         Z *= num_modes / norm(Z)
@@ -574,7 +573,7 @@ def eigsolve(
         try:
             U = numpy.linalg.inv(ZtZ)
         except numpy.linalg.LinAlgError:
-            Z = rng.random(y_shape) + 1j * rng.random(y_shape)
+            Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
             continue
 
         trace_U = real(trace(U))
@@ -647,7 +646,8 @@ def eigsolve(
 
         Qi_memo: list[float | None] = [None, None]
 
-        def Qi_func(theta: float, Qi_memo=Qi_memo, ZtZ=ZtZ, DtD=DtD, symZtD=symZtD) -> float:   # noqa: ANN001
+        def Qi_func(theta: float) -> float:
+            nonlocal Qi_memo
             if Qi_memo[0] == theta:
                 return cast(float, Qi_memo[1])
 
@@ -656,7 +656,7 @@ def eigsolve(
             Q = c * c * ZtZ + s * s * DtD + 2 * s * c * symZtD
             try:
                 Qi = numpy.linalg.inv(Q)
-            except numpy.linalg.LinAlgError as err:
+            except numpy.linalg.LinAlgError:
                 logger.info('taylor Qi')
                 # if c or s small, taylor expand
                 if c < 1e-4 * s and c != 0:
@@ -666,12 +666,12 @@ def eigsolve(
                     ZtZi = numpy.linalg.inv(ZtZ)
                     Qi = ZtZi / (c * c) - 2 * s / (c * c * c) * (ZtZi @ (ZtZi @ symZtD).conj().T)
                 else:
-                    raise Exception('Inexplicable singularity in trace_func') from err
+                    raise Exception('Inexplicable singularity in trace_func')
             Qi_memo[0] = theta
             Qi_memo[1] = cast(float, Qi)
             return cast(float, Qi)
 
-        def trace_func(theta: float, ZtAZ=ZtAZ, DtAD=DtAD, symZtAD=symZtAD) -> float:           # noqa: ANN001
+        def trace_func(theta: float) -> float:
             c = numpy.cos(theta)
             s = numpy.sin(theta)
             Qi = Qi_func(theta)
@@ -680,7 +680,7 @@ def eigsolve(
             return numpy.abs(trace)
 
         if False:
-            def trace_deriv(theta, sgn: int = sgn, ZtAZ=ZtAZ, DtAD=DtAD, symZtD=symZtD, symZtAD=symZtAD, ZtZ=ZtZ, DtD=DtD):     # noqa: ANN001
+            def trace_deriv(theta):
                 Qi = Qi_func(theta)
                 c2 = numpy.cos(2 * theta)
                 s2 = numpy.sin(2 * theta)

meanas/fdfd/farfield.py

@@ -1,8 +1,7 @@
 """
 Functions for performing near-to-farfield transformation (and the reverse).
 """
-from typing import Any, cast
+from typing import Any, Sequence, cast
-from collections.abc import Sequence
 import numpy
 from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
 from numpy import pi

meanas/fdfd/functional.py

@@ -5,7 +5,7 @@ Functional versions of many FDFD operators. These can be useful for performing
 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 collections.abc import Callable
+from typing import Callable
 import numpy
 
 from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t
@@ -47,7 +47,8 @@ def e_full(
 
     if mu is None:
         return op_1
-    return op_mu
+    else:
+        return op_mu
 
 
 def eh_full(
@@ -83,7 +84,8 @@ def eh_full(
 
     if mu is None:
         return op_1
-    return op_mu
+    else:
+        return op_mu
 
 
 def e2h(
@@ -114,7 +116,8 @@ def e2h(
 
     if mu is None:
         return e2h_1_1
-    return e2h_mu
+    else:
+        return e2h_mu
 
 
 def m2j(
@@ -148,7 +151,8 @@ def m2j(
 
     if mu is None:
         return m2j_1
-    return m2j_mu
+    else:
+        return m2j_mu
 
 
 def e_tfsf_source(

meanas/fdfd/operators.py

@@ -28,7 +28,7 @@ The following operators are included:
 """
 
 import numpy
-from scipy import sparse
+import scipy.sparse as sparse       # type: ignore
 
 from ..fdmath import vec, dx_lists_t, vfdfield_t, vcfdfield_t
 from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl_back
@@ -321,11 +321,11 @@ def poynting_e_cross(e: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     """
     shape = [len(dx) for dx in dxes[0]]
 
-    fx, fy, fz = (shift_circ(i, shape, 1) for i in range(3))
+    fx, fy, fz = [shift_circ(i, shape, 1) for i in range(3)]
 
     dxag = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
     dxbg = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
-    Ex, Ey, Ez = (ei * da for ei, da in zip(numpy.split(e, 3), dxag, strict=True))
+    Ex, Ey, Ez = [ei * da for ei, da in zip(numpy.split(e, 3), dxag)]
 
     block_diags = [[ None,     fx @ -Ez, fx @  Ey],
                    [ fy @  Ez, None,     fy @ -Ex],
@@ -349,11 +349,11 @@ def poynting_h_cross(h: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     """
     shape = [len(dx) for dx in dxes[0]]
 
-    fx, fy, fz = (shift_circ(i, shape, 1) for i in range(3))
+    fx, fy, fz = [shift_circ(i, shape, 1) for i in range(3)]
 
     dxag = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
     dxbg = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
-    Hx, Hy, Hz = (sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg, strict=True))
+    Hx, Hy, Hz = [sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg)]
 
     P = (sparse.bmat(
         [[ None,    -Hz @ fx,   Hy @ fx],

meanas/fdfd/scpml.py

@@ -2,7 +2,7 @@
 Functions for creating stretched coordinate perfectly matched layer (PML) absorbers.
 """
 
-from collections.abc import Sequence, Callable
+from typing import Sequence, Callable
 
 import numpy
 from numpy.typing import NDArray

meanas/fdfd/solvers.py

@@ -2,14 +2,13 @@
 Solvers and solver interface for FDFD problems.
 """
 
-from typing import Any
+from typing import Callable, Dict, Any, Optional
-from collections.abc import Callable
 import logging
 
 import numpy
 from numpy.typing import ArrayLike, NDArray
 from numpy.linalg import norm
-import scipy.sparse.linalg
+import scipy.sparse.linalg          # type: ignore
 
 from ..fdmath import dx_lists_t, vfdfield_t, vcfdfield_t
 from . import operators
@@ -69,12 +68,12 @@ def generic(
         dxes: dx_lists_t,
         J: vcfdfield_t,
         epsilon: vfdfield_t,
-        mu: vfdfield_t | None = None,
+        mu: Optional[vfdfield_t] = None,
-        pec: vfdfield_t | None = None,
+        pec: Optional[vfdfield_t] = None,
-        pmc: vfdfield_t | None = None,
+        pmc: Optional[vfdfield_t] = None,
         adjoint: bool = False,
         matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
-        matrix_solver_opts: dict[str, Any] | None = None,
+        matrix_solver_opts: Optional[Dict[str, Any]] = None,
         ) -> vcfdfield_t:
     """
     Conjugate gradient FDFD solver using CSR sparse matrices.

meanas/fdfd/waveguide_2d.py

@@ -182,7 +182,7 @@ from typing import Any
 import numpy
 from numpy.typing import NDArray, ArrayLike
 from numpy.linalg import norm
-from scipy import sparse
+import scipy.sparse as sparse       # type: ignore
 
 from ..fdmath.operators import deriv_forward, deriv_back, cross
 from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t

meanas/fdfd/waveguide_3d.py

@@ -4,11 +4,9 @@ Tools for working with waveguide modes in 3D domains.
 This module relies heavily on `waveguide_2d` and mostly just transforms
 its parameters into 2D equivalents and expands the results back into 3D.
 """
-from typing import Any
+from typing import Sequence, Any
-from collections.abc import Sequence
 import numpy
 from numpy.typing import NDArray
-from numpy import complexfloating
 
 from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, cfdfield_t
 from . import operators, waveguide_2d
@@ -23,7 +21,7 @@ def solve_mode(
         slices: Sequence[slice],
         epsilon: fdfield_t,
         mu: fdfield_t | None = None,
-        ) -> dict[str, complex | NDArray[complexfloating]]:
+        ) -> dict[str, 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.
@@ -42,8 +40,8 @@ def solve_mode(
     Returns:
         ```
         {
-            'E': NDArray[complexfloating],
+            'E': list[NDArray[numpy.float_]],
-            'H': NDArray[complexfloating],
+            'H': list[NDArray[numpy.float_]],
             'wavenumber': complex,
         }
         ```

meanas/fdfd/waveguide_cyl.py

@@ -9,9 +9,9 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
 # TODO update module docs
 
 import numpy
-from scipy import sparse
+import scipy.sparse as sparse       # type: ignore
 
-from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, cfdfield_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
 

meanas/fdmath/__init__.py

@@ -741,24 +741,8 @@ the true values can be multiplied back in after the simulation is complete if no
 normalized results are needed.
 """
 
-from .types import (
+from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t, dx_lists_t, dx_lists_mut
-    fdfield_t as fdfield_t,
+from .types import fdfield_updater_t, cfdfield_updater_t
-    vfdfield_t as vfdfield_t,
+from .vectorization import vec, unvec
-    cfdfield_t as cfdfield_t,
+from . import operators, functional, types, vectorization
-    vcfdfield_t as vcfdfield_t,
-    dx_lists_t as dx_lists_t,
-    dx_lists_mut as dx_lists_mut,
-    fdfield_updater_t as fdfield_updater_t,
-    cfdfield_updater_t as cfdfield_updater_t,
-    )
-from .vectorization import (
-    vec as vec,
-    unvec as unvec,
-    )
-from . import (
-    operators as operators,
-    functional as functional,
-    types as types,
-    vectorization as vectorization,
-    )
 

meanas/fdmath/functional.py

@@ -3,18 +3,16 @@ Math functions for finite difference simulations
 
 Basic discrete calculus etc.
 """
-from typing import TypeVar
+from typing import Sequence, Callable
-from collections.abc import Sequence, Callable
 
 import numpy
 from numpy.typing import NDArray
-from numpy import floating, complexfloating
 
 from .types import fdfield_t, fdfield_updater_t
 
 
 def deriv_forward(
-        dx_e: Sequence[NDArray[floating]] | None = None,
+        dx_e: Sequence[NDArray[numpy.float_]] | None = None,
         ) -> tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
     """
     Utility operators for taking discretized derivatives (backward variant).
@@ -38,7 +36,7 @@ def deriv_forward(
 
 
 def deriv_back(
-        dx_h: Sequence[NDArray[floating]] | None = None,
+        dx_h: Sequence[NDArray[numpy.float_]] | None = None,
         ) -> tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
     """
     Utility operators for taking discretized derivatives (forward variant).
@@ -61,12 +59,9 @@ def deriv_back(
     return derivs
 
 
-TT = TypeVar('TT', bound='NDArray[floating | complexfloating]')
-
-
 def curl_forward(
-        dx_e: Sequence[NDArray[floating]] | None = None,
+        dx_e: Sequence[NDArray[numpy.float_]] | None = None,
-        ) -> Callable[[TT], TT]:
+        ) -> fdfield_updater_t:
     r"""
     Curl operator for use with the E field.
 
@@ -80,7 +75,7 @@ def curl_forward(
     """
     Dx, Dy, Dz = deriv_forward(dx_e)
 
-    def ce_fun(e: TT) -> TT:
+    def ce_fun(e: fdfield_t) -> fdfield_t:
         output = numpy.empty_like(e)
         output[0] = Dy(e[2])
         output[1] = Dz(e[0])
@@ -94,8 +89,8 @@ def curl_forward(
 
 
 def curl_back(
-        dx_h: Sequence[NDArray[floating]] | None = None,
+        dx_h: Sequence[NDArray[numpy.float_]] | None = None,
-        ) -> Callable[[TT], TT]:
+        ) -> fdfield_updater_t:
     r"""
     Create a function which takes the backward curl of a field.
 
@@ -109,7 +104,7 @@ def curl_back(
     """
     Dx, Dy, Dz = deriv_back(dx_h)
 
-    def ch_fun(h: TT) -> TT:
+    def ch_fun(h: fdfield_t) -> fdfield_t:
         output = numpy.empty_like(h)
         output[0] = Dy(h[2])
         output[1] = Dz(h[0])
@@ -123,7 +118,7 @@ def curl_back(
 
 
 def curl_forward_parts(
-        dx_e: Sequence[NDArray[floating]] | None = None,
+        dx_e: Sequence[NDArray[numpy.float_]] | None = None,
         ) -> Callable:
     Dx, Dy, Dz = deriv_forward(dx_e)
 
@@ -136,7 +131,7 @@ def curl_forward_parts(
 
 
 def curl_back_parts(
-        dx_h: Sequence[NDArray[floating]] | None = None,
+        dx_h: Sequence[NDArray[numpy.float_]] | None = None,
         ) -> Callable:
     Dx, Dy, Dz = deriv_back(dx_h)
 

meanas/fdmath/operators.py

@@ -3,11 +3,10 @@ Matrix operators for finite difference simulations
 
 Basic discrete calculus etc.
 """
-from collections.abc import Sequence
+from typing import Sequence
 import numpy
 from numpy.typing import NDArray
-from numpy import floating
+import scipy.sparse as sparse   # type: ignore
-from scipy import sparse
 
 from .types import vfdfield_t
 
@@ -35,7 +34,7 @@ def shift_circ(
         raise Exception(f'Invalid direction: {axis}, shape is {shape}')
 
     shifts = [abs(shift_distance) if a == axis else 0 for a in range(3)]
-    shifted_diags = [(numpy.arange(n) + s) % n for n, s in zip(shape, shifts, strict=True)]
+    shifted_diags = [(numpy.arange(n) + s) % n for n, s in zip(shape, shifts)]
     ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
 
     n = numpy.prod(shape)
@@ -83,7 +82,7 @@ def shift_with_mirror(
         return v
 
     shifts = [shift_distance if a == axis else 0 for a in range(3)]
-    shifted_diags = [mirrored_range(n, s) for n, s in zip(shape, shifts, strict=True)]
+    shifted_diags = [mirrored_range(n, s) for n, s in zip(shape, shifts)]
     ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
 
     n = numpy.prod(shape)
@@ -97,7 +96,7 @@ def shift_with_mirror(
 
 
 def deriv_forward(
-        dx_e: Sequence[NDArray[floating]],
+        dx_e: Sequence[NDArray[numpy.float_]],
         ) -> list[sparse.spmatrix]:
     """
     Utility operators for taking discretized derivatives (forward variant).
@@ -124,7 +123,7 @@ def deriv_forward(
 
 
 def deriv_back(
-        dx_h: Sequence[NDArray[floating]],
+        dx_h: Sequence[NDArray[numpy.float_]],
         ) -> list[sparse.spmatrix]:
     """
     Utility operators for taking discretized derivatives (backward variant).
@@ -219,7 +218,7 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
 
 
 def curl_forward(
-        dx_e: Sequence[NDArray[floating]],
+        dx_e: Sequence[NDArray[numpy.float_]],
         ) -> sparse.spmatrix:
     """
     Curl operator for use with the E field.
@@ -235,7 +234,7 @@ def curl_forward(
 
 
 def curl_back(
-        dx_h: Sequence[NDArray[floating]],
+        dx_h: Sequence[NDArray[numpy.float_]],
         ) -> sparse.spmatrix:
     """
     Curl operator for use with the H field.

meanas/fdmath/types.py

@@ -1,26 +1,26 @@
 """
 Types shared across multiple submodules
 """
-from collections.abc import Sequence, Callable, MutableSequence
+from typing import Sequence, Callable, MutableSequence
+import numpy
 from numpy.typing import NDArray
-from numpy import floating, complexfloating
 
 
 # Field types
-fdfield_t = NDArray[floating]
+fdfield_t = NDArray[numpy.float_]
 """Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
 
-vfdfield_t = NDArray[floating]
+vfdfield_t = NDArray[numpy.float_]
 """Linearized vector field (single vector of length 3*X*Y*Z)"""
 
-cfdfield_t = NDArray[complexfloating]
+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[complexfloating]
+vcfdfield_t = NDArray[numpy.complex_]
 """Linearized complex vector field (single vector of length 3*X*Y*Z)"""
 
 
-dx_lists_t = Sequence[Sequence[NDArray[floating]]]
+dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]]
 """
  'dxes' datastructure which contains grid cell width information in the following format:
 
@@ -31,7 +31,7 @@ dx_lists_t = Sequence[Sequence[NDArray[floating]]]
    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[NDArray[floating]]]
+dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]]
 """Mutable version of `dx_lists_t`"""
 
 

meanas/fdmath/vectorization.py

@@ -4,8 +4,7 @@ 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 overload
+from typing import overload, Sequence
-from collections.abc import Sequence
 import numpy
 from numpy.typing import ArrayLike
 

meanas/fdtd/__init__.py

@@ -159,22 +159,8 @@ Boundary conditions
 # TODO notes about boundaries / PMLs
 """
 
-from .base import (
+from .base import maxwell_e, maxwell_h
-    maxwell_e as maxwell_e,
+from .pml import cpml_params, updates_with_cpml
-    maxwell_h as maxwell_h,
+from .energy import (poynting, poynting_divergence, energy_hstep, energy_estep,
-    )
+                     delta_energy_h2e, delta_energy_j)
-from .pml import (
+from .boundaries import conducting_boundary
-    cpml_params as cpml_params,
-    updates_with_cpml as updates_with_cpml,
-    )
-from .energy import (
-    poynting as poynting,
-    poynting_divergence as poynting_divergence,
-    energy_hstep as energy_hstep,
-    energy_estep as energy_estep,
-    delta_energy_h2e as delta_energy_h2e,
-    delta_energy_j as delta_energy_j,
-    )
-from .boundaries import (
-    conducting_boundary as conducting_boundary,
-    )

meanas/fdtd/pml.py

@@ -7,8 +7,7 @@ PML implementations
 """
 # TODO retest pmls!
 
-from typing import Any
+from typing import Callable, Sequence, Any
-from collections.abc import Callable, Sequence
 from copy import deepcopy
 import numpy
 from numpy.typing import NDArray, DTypeLike
@@ -112,7 +111,7 @@ def updates_with_cpml(
     params_H: list[list[tuple[Any, Any, Any, Any]]] = deepcopy(params_E)
 
     for axis in range(3):
-        for pp, _polarity in enumerate((-1, 1)):
+        for pp, polarity in enumerate((-1, 1)):
             cpml_param = cpml_params[axis][pp]
             if cpml_param is None:
                 psi_E[axis][pp] = (None, None)
@@ -185,7 +184,7 @@ def updates_with_cpml(
     def update_H(
             e: fdfield_t,
             h: fdfield_t,
-            mu: fdfield_t | tuple[int, int, int] = (1, 1, 1),
+            mu: fdfield_t = numpy.ones(3),
             ) -> None:
         dyEx = Dfy(e[0])
         dzEx = Dfz(e[0])

meanas/test/conftest.py

@@ -3,8 +3,7 @@
 Test fixtures
 
 """
-# ruff: noqa: ARG001
+from typing import Iterable, Any
-from typing import Any
 import numpy
 from numpy.typing import NDArray
 import pytest       # type: ignore
@@ -21,18 +20,18 @@ FixtureRequest = Any
                         (5, 5, 5),
                         # (7, 7, 7),
                        ])
-def shape(request: FixtureRequest) -> tuple[int, ...]:
+def shape(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
-    return (3, *request.param)
+    yield (3, *request.param)
 
 
 @pytest.fixture(scope='module', params=[1.0, 1.5])
-def epsilon_bg(request: FixtureRequest) -> float:
+def epsilon_bg(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(scope='module', params=[1.0, 2.5])
-def epsilon_fg(request: FixtureRequest) -> float:
+def epsilon_fg(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(scope='module', params=['center', '000', 'random'])
@@ -41,7 +40,7 @@ def epsilon(
         shape: tuple[int, ...],
         epsilon_bg: float,
         epsilon_fg: float,
-        ) -> NDArray[numpy.float64]:
+        ) -> Iterable[NDArray[numpy.float64]]:
     is3d = (numpy.array(shape) == 1).sum() == 0
     if is3d:
         if request.param == '000':
@@ -61,17 +60,17 @@ def epsilon(
                                   high=max(epsilon_bg, epsilon_fg),
                                   size=shape)
 
-    return epsilon
+    yield epsilon
 
 
 @pytest.fixture(scope='module', params=[1.0])  # 1.5
-def j_mag(request: FixtureRequest) -> float:
+def j_mag(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(scope='module', params=[1.0, 1.5])
-def dx(request: FixtureRequest) -> float:
+def dx(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(scope='module', params=['uniform', 'centerbig'])
@@ -79,7 +78,7 @@ def dxes(
         request: FixtureRequest,
         shape: tuple[int, ...],
         dx: float,
-        ) -> list[list[NDArray[numpy.float64]]]:
+        ) -> 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':
@@ -91,5 +90,5 @@ def dxes(
         dxe = [PRNG.uniform(low=1.0 * dx, high=1.1 * dx, size=s) for s in shape[1:]]
         dxh = [(d + numpy.roll(d, -1)) / 2 for d in dxe]
         dxes = [dxe, dxh]
-    return dxes
+    yield dxes
 

meanas/test/test_fdfd.py

@@ -1,4 +1,4 @@
-# ruff: noqa: ARG001
+from typing import Iterable
 import dataclasses
 import pytest       # type: ignore
 import numpy
@@ -61,24 +61,24 @@ def test_poynting_planes(sim: 'FDResult') -> None:
 # Also see conftest.py
 
 @pytest.fixture(params=[1 / 1500])
-def omega(request: FixtureRequest) -> float:
+def omega(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(params=[None])
-def pec(request: FixtureRequest) -> NDArray[numpy.float64] | None:
+def pec(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(params=[None])
-def pmc(request: FixtureRequest) -> NDArray[numpy.float64] | None:
+def pmc(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
-    return request.param
+    yield request.param
 
 
 #@pytest.fixture(scope='module',
 #                params=[(25, 5, 5)])
-#def shape(request: FixtureRequest):
+#def shape(request):
-#    return (3, *request.param)
+#    yield (3, *request.param)
 
 
 @pytest.fixture(params=['diag'])        # 'center'
@@ -86,7 +86,7 @@ def j_distribution(
         request: FixtureRequest,
         shape: tuple[int, ...],
         j_mag: float,
-        ) -> NDArray[numpy.float64]:
+        ) -> 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
@@ -96,7 +96,7 @@ def j_distribution(
     elif request.param == 'diag':
         j[numpy.roll(center_mask, [1, 1, 1], axis=(1, 2, 3))] = (1 + 1j) * j_mag
         j[numpy.roll(center_mask, [-1, -1, -1], axis=(1, 2, 3))] = (1 - 1j) * j_mag
-    return j
+    yield j
 
 
 @dataclasses.dataclass()
@@ -145,7 +145,7 @@ def sim(
         omega=omega,
         dxes=dxes,
         epsilon=eps_vec,
-        matrix_solver_opts={'atol': 1e-15, 'rtol': 1e-11},
+        matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11},
         )
     e = unvec(e_vec, shape[1:])
 

meanas/test/test_fdfd_pml.py

@@ -1,4 +1,4 @@
-# ruff: noqa: ARG001
+from typing import Iterable
 import pytest       # type: ignore
 import numpy
 from numpy.typing import NDArray
@@ -44,30 +44,30 @@ def test_pml(sim: FDResult, src_polarity: int) -> None:
 # Also see conftest.py
 
 @pytest.fixture(params=[1 / 1500])
-def omega(request: FixtureRequest) -> float:
+def omega(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(params=[None])
-def pec(request: FixtureRequest) -> NDArray[numpy.float64] | None:
+def pec(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(params=[None])
-def pmc(request: FixtureRequest) -> NDArray[numpy.float64] | None:
+def pmc(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(params=[(30, 1, 1),
                         (1, 30, 1),
                         (1, 1, 30)])
-def shape(request: FixtureRequest) -> tuple[int, int, int]:
+def shape(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
-    return (3, *request.param)
+    yield (3, *request.param)
 
 
 @pytest.fixture(params=[+1, -1])
-def src_polarity(request: FixtureRequest) -> int:
+def src_polarity(request: FixtureRequest) -> Iterable[int]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture()
@@ -78,7 +78,7 @@ def j_distribution(
         dxes: dx_lists_mut,
         omega: float,
         src_polarity: int,
-        ) -> NDArray[numpy.complex128]:
+        ) -> Iterable[NDArray[numpy.complex128]]:
     j = numpy.zeros(shape, dtype=complex)
 
     dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0]    # Propagation axis
@@ -106,7 +106,7 @@ def j_distribution(
 
     j = fdfd.waveguide_3d.compute_source(E=e, wavenumber=wavenumber_corrected, omega=omega, dxes=dxes,
                                          axis=dim, polarity=src_polarity, slices=slices, epsilon=epsilon)
-    return j
+    yield j
 
 
 @pytest.fixture()
@@ -115,9 +115,9 @@ def epsilon(
         shape: tuple[int, ...],
         epsilon_bg: float,
         epsilon_fg: float,
-        ) -> NDArray[numpy.float64]:
+        ) -> Iterable[NDArray[numpy.float64]]:
     epsilon = numpy.full(shape, epsilon_fg, dtype=float)
-    return epsilon
+    yield epsilon
 
 
 @pytest.fixture(params=['uniform'])
@@ -127,7 +127,7 @@ def dxes(
         dx: float,
         omega: float,
         epsilon_fg: float,
-        ) -> list[list[NDArray[numpy.float64]]]:
+        ) -> 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
@@ -141,7 +141,7 @@ def dxes(
                 epsilon_effective=epsilon_fg,
                 thickness=10,
                 )
-    return dxes
+    yield dxes
 
 
 @pytest.fixture()
@@ -162,7 +162,7 @@ def sim(
         omega=omega,
         dxes=dxes,
         epsilon=eps_vec,
-        matrix_solver_opts={'atol': 1e-15, 'rtol': 1e-11},
+        matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11},
         )
     e = unvec(e_vec, shape[1:])
 

meanas/test/test_fdtd.py

@@ -1,5 +1,4 @@
-# ruff: noqa: ARG001
+from typing import Iterable, Any
-from typing import Any
 import dataclasses
 import pytest       # type: ignore
 import numpy
@@ -151,8 +150,8 @@ def test_poynting_planes(sim: 'TDResult') -> None:
 
 
 @pytest.fixture(params=[0.3])
-def dt(request: FixtureRequest) -> float:
+def dt(request: FixtureRequest) -> Iterable[float]:
-    return request.param
+    yield request.param
 
 
 @dataclasses.dataclass()
@@ -169,8 +168,8 @@ class TDResult:
 
 
 @pytest.fixture(params=[(0, 4, 8)])  # (0,)
-def j_steps(request: FixtureRequest) -> tuple[int, ...]:
+def j_steps(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
-    return request.param
+    yield request.param
 
 
 @pytest.fixture(params=['center', 'random'])
@@ -178,7 +177,7 @@ def j_distribution(
         request: FixtureRequest,
         shape: tuple[int, ...],
         j_mag: float,
-        ) -> NDArray[numpy.float64]:
+        ) -> Iterable[NDArray[numpy.float64]]:
     j = numpy.zeros(shape)
     if request.param == 'center':
         j[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = j_mag
@@ -186,7 +185,7 @@ def j_distribution(
         j[:, 0, 0, 0] = j_mag
     elif request.param == 'random':
         j[:] = PRNG.uniform(low=-j_mag, high=j_mag, size=shape)
-    return j
+    yield j
 
 
 @pytest.fixture()
@@ -200,8 +199,9 @@ def sim(
         j_steps: tuple[int, ...],
         ) -> TDResult:
     is3d = (numpy.array(shape) == 1).sum() == 0
-    if is3d and dt != 0.3:
+    if is3d:
-        pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)')
+        if dt != 0.3:
+            pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)')
 
     sim = TDResult(
         shape=shape,

meanas/test/utils.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import numpy
 from numpy.typing import NDArray
 
@@ -8,25 +10,22 @@ PRNG = numpy.random.RandomState(12345)
 def assert_fields_close(
         x: NDArray,
         y: NDArray,
-        *args,
+        *args: Any,
-        **kwargs,
+        **kwargs: Any,
         ) -> None:
-    x_disp = numpy.moveaxis(x, -1, 0)
-    y_disp = numpy.moveaxis(y, -1, 0)
     numpy.testing.assert_allclose(
-        x,           # type: ignore
+        x, y, verbose=False,            # type: ignore
-        y,           # type: ignore
+        err_msg='Fields did not match:\n{}\n{}'.format(numpy.moveaxis(x, -1, 0),
+                                                       numpy.moveaxis(y, -1, 0)),
         *args,
-        verbose=False,
-        err_msg=f'Fields did not match:\n{x_disp}\n{y_disp}',
         **kwargs,
         )
 
 def assert_close(
         x: NDArray,
         y: NDArray,
-        *args,
+        *args: Any,
-        **kwargs,
+        **kwargs: Any,
         ) -> None:
     numpy.testing.assert_allclose(x, y, *args, **kwargs)
 

pyproject.toml

@@ -39,8 +39,8 @@ include = [
     ]
 dynamic = ["version"]
 dependencies = [
-    "numpy>=1.26",
+    "numpy~=1.21",
-    "scipy~=1.14",
+    "scipy",
     ]
 
 
@@ -51,48 +51,3 @@ path = "meanas/__init__.py"
 dev = ["pytest", "pdoc", "gridlock"]
 examples = ["gridlock"]
 test = ["pytest"]
-
-
-[tool.ruff]
-exclude = [
-    ".git",
-    "dist",
-    ]
-line-length = 245
-indent-width = 4
-lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-lint.select = [
-    "NPY", "E", "F", "W", "B", "ANN", "UP", "SLOT", "SIM", "LOG",
-    "C4", "ISC", "PIE", "PT", "RET", "TCH", "PTH", "INT",
-    "ARG", "PL", "R", "TRY",
-    "G010", "G101", "G201", "G202",
-    "Q002", "Q003", "Q004",
-    ]
-lint.ignore = [
-    #"ANN001",   # No annotation
-    "ANN002",   # *args
-    "ANN003",   # **kwargs
-    "ANN401",   # Any
-    "ANN101",   # self: Self
-    "SIM108",   # single-line if / else assignment
-    "RET504",   # x=y+z; return x
-    "PIE790",   # unnecessary pass
-    "ISC003",   # non-implicit string concatenation
-    "C408",     # dict(x=y) instead of {'x': y}
-    "PLR09",    # Too many xxx
-    "PLR2004",  # magic number
-    "PLC0414",  # import x as x
-    "TRY003",   # Long exception message
-    "TRY002",   # Exception()
-    ]
-
-
-[[tool.mypy.overrides]]
-module = [
-    "scipy",
-    "scipy.optimize",
-    "scipy.linalg",
-    "scipy.sparse",
-    "scipy.sparse.linalg",
-    ]
-ignore_missing_imports = true