fixup! add sensitivity calculation

examples/fdfd.py

@@ -46,24 +46,20 @@ def test0(solver=generic_solver):
     # #### Create the grid, mask, and draw the device ####
     grid = gridlock.Grid(edge_coords)
     epsilon = grid.allocate(n_air**2, dtype=numpy.float32)
-    grid.draw_cylinder(
+    grid.draw_cylinder(epsilon,
-        epsilon,
+                       surface_normal=2,
-        surface_normal=2,
+                       center=center,
-        center=center,
+                       radius=max(radii),
-        radius=max(radii),
+                       thickness=th,
-        thickness=th,
+                       eps=n_ring**2,
-        foreground=n_ring**2,
+                       num_points=24)
-        num_points=24,
+    grid.draw_cylinder(epsilon,
-        )
+                       surface_normal=2,
-    grid.draw_cylinder(
+                       center=center,
-        epsilon,
+                       radius=min(radii),
-        surface_normal=2,
+                       thickness=th*1.1,
-        center=center,
+                       eps=n_air ** 2,
-        radius=min(radii),
+                       num_points=24)
-        thickness=th*1.1,
-        foreground=n_air ** 2,
-        num_points=24,
-        )
 
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (0, 1, 2):
@@ -75,9 +71,9 @@ def test0(solver=generic_solver):
     J[1][15, grid.shape[1]//2, grid.shape[2]//2] = 1
 
 
-    #
+    '''
-    # Solve!
+    Solve!
-    #
+    '''
     sim_args = {
         'omega': omega,
         'dxes': dxes,
@@ -91,9 +87,9 @@ def test0(solver=generic_solver):
 
     E = unvec(x, grid.shape)
 
-    #
+    '''
-    # Plot results
+    Plot results
-    #
+    '''
     pyplot.figure()
     pyplot.pcolor(numpy.real(E[1][:, :, grid.shape[2]//2]), cmap='seismic')
     pyplot.axis('equal')
@@ -126,7 +122,7 @@ def test1(solver=generic_solver):
     # #### Create the grid and draw the device ####
     grid = gridlock.Grid(edge_coords)
     epsilon = grid.allocate(n_air**2, dtype=numpy.float32)
-    grid.draw_cuboid(epsilon, center=center, dimensions=[8e3, w, th], foreground=n_wg**2)
+    grid.draw_cuboid(epsilon, center=center, dimensions=[8e3, w, th], eps=n_wg**2)
 
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (0, 1, 2):
@@ -173,9 +169,9 @@ def test1(solver=generic_solver):
 #    pcolor((numpy.abs(J3).sum(axis=2).sum(axis=0) > 0).astype(float).T)
     pyplot.show(block=True)
 
-    #
+    '''
-    # Solve!
+    Solve!
-    #
+    '''
     sim_args = {
         'omega': omega,
         'dxes': dxes,
@@ -192,9 +188,9 @@ def test1(solver=generic_solver):
 
     E = unvec(x, grid.shape)
 
-    #
+    '''
-    # Plot results
+    Plot results
-    #
+    '''
     center = grid.pos2ind([0, 0, 0], None).astype(int)
     pyplot.figure()
     pyplot.subplot(2, 2, 1)
@@ -236,7 +232,7 @@ def test1(solver=generic_solver):
     pyplot.grid(alpha=0.6)
     pyplot.title('Overlap with mode')
     pyplot.show()
-    print('Average overlap with mode:', sum(q[8:32])/len(q[8:32]))
+    print('Average overlap with mode:', sum(q)/len(q))
 
 
 def module_available(name):

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)
@@ -799,52 +799,3 @@ def _rtrace_AtB(
 def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
     return (A + A.conj().T) * 0.5
 
-
-
-def inner_product(eL, hL, eR, hR) -> complex:
-    # assumes x-axis propagation
-
-    assert numpy.array_equal(eR.shape, hR.shape)
-    assert numpy.array_equal(eL.shape, hL.shape)
-    assert numpy.array_equal(eR.shape, eL.shape)
-
-    # Cross product, times 2 since it's <p | n>, then divide by 4. # TODO might want to abs() this?
-    norm2R = (eR[1] * hR[2] - eR[2] * hR[1]).sum() / 2
-    norm2L = (eL[1] * hL[2] - eL[2] * hL[1]).sum() / 2
-
-    # eRxhR_x = numpy.cross(eR.reshape(3, -1), hR.reshape(3, -1), axis=0).reshape(eR.shape)[0] / normR
-    # logger.info(f'power {eRxhR_x.sum() / 2})
-
-    eR /= numpy.sqrt(norm2R)
-    hR /= numpy.sqrt(norm2R)
-    eL /= numpy.sqrt(norm2L)
-    hL /= numpy.sqrt(norm2L)
-
-    # (eR x hL)[0] and (eL x hR)[0]
-    eRxhL_x = eR[1] * hL[2] - eR[2] - hL[1]
-    eLxhR_x = eL[1] * hR[2] - eL[2] - hR[1]
-
-    #return 1j *  (eRxhL_x - eLxhR_x).sum() / numpy.sqrt(norm2R * norm2L)
-    #return (eRxhL_x.sum() - eLxhR_x.sum()) / numpy.sqrt(norm2R * norm2L)
-    return eRxhL_x.sum() - eLxhR_x.sum()
-
-
-def trq(eI, hI, eO, hO) -> tuple[complex, complex]:
-    pp = inner_product(eO,  hO, eI,  hI)
-    pn = inner_product(eO,  hO, eI, -hI)
-    np = inner_product(eO, -hO, eI,  hI)
-    nn = inner_product(eO, -hO, eI, -hI)
-
-    assert pp == -nn
-    assert pn == -np
-
-    logger.info(f'''
-        {pp=:4g} {pn=:4g}
-        {nn=:4g} {np=:4g}
-            {nn * pp / pn=:4g}    {-np=:4g}
-        ''')
-
-    r = -pp / pn    # -<Pp|Bp>/<Pn/Bp> = -(-pp) / (-pn)
-    t = (np - nn * pp / pn) / 4
-
-    return t, r

meanas/fdfd/eme.py

@@ -1,68 +0,0 @@
-import numpy
-
-from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t
-from .waveguide_2d import inner_product
-
-
-def get_tr(ehL, wavenumbers_L, ehR, wavenumbers_R, dxes: dx_lists_t):
-    nL = len(wavenumbers_L)
-    nR = len(wavenumbers_R)
-    A12 = numpy.zeros((nL, nR), dtype=complex)
-    A21 = numpy.zeros((nL, nR), dtype=complex)
-    B11 = numpy.zeros((nL,), dtype=complex)
-    for ll in range(nL):
-        eL, hL = ehL[ll]
-        B11[ll] = inner_product(eL, hL, dxes=dxes, conj_h=False)
-        for rr in range(nR):
-            eR, hR = ehR[rr]
-            A12[ll, rr] = inner_product(eL, hR, dxes=dxes, conj_h=False)    # TODO optimize loop?
-            A21[ll, rr] = inner_product(eR, hL, dxes=dxes, conj_h=False)
-
-    # tt0 = 2 * numpy.linalg.pinv(A21 + numpy.conj(A12))
-    tt0, _resid, _rank, _sing = numpy.linalg.lstsq(A21 + A12, numpy.diag(2 * B11), rcond=None)
-
-    U, st, V = numpy.linalg.svd(tt0)
-    gain = st > 1
-    st[gain] = 1 / st[gain]
-    tt = U @ numpy.diag(st) @ V
-
-    # rr = 0.5 * (A21 - numpy.conj(A12)) @ tt
-    rr = numpy.diag(0.5 / B11) @ (A21 - A12) @ tt
-
-    return tt, rr
-
-
-def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs):
-    t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs)
-    t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs)
-    t21i = numpy.linalg.pinv(t21)
-    A = t12 - r21 @ t21i @ r12
-    B = r21 @ t21i
-    C = -t21i @ r12
-    D = t21i
-    return sparse.block_array(((A, B), (C, D)))
-
-
-def get_s(
-        eL_xys,
-        wavenumbers_L,
-        eR_xys,
-        wavenumbers_R,
-        force_nogain: bool = False,
-        force_reciprocal: bool = False,
-        **kwargs):
-    t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs)
-    t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs)
-
-    ss = numpy.block([[r12, t12],
-                      [t21, r21]])
-
-    if force_nogain:
-        # force S @ S.H diagonal
-        U, sing, V = numpy.linalg.svd(ss)
-        ss = numpy.diag(sing) @ U @ V
-
-    if force_reciprocal:
-        ss = 0.5 * (ss + ss.T)
-
-    return ss

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
@@ -40,7 +40,7 @@ __author__ = 'Jan Petykiewicz'
 def e_full(
         omega: complex,
         dxes: dx_lists_t,
-        epsilon: vfdfield_t | vcfdfield_t,
+        epsilon: vfdfield_t,
         mu: vfdfield_t | None = None,
         pec: vfdfield_t | None = None,
         pmc: vfdfield_t | None = None,
@@ -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
@@ -35,16 +34,16 @@ def _scipy_qmr(
         Guess for solution (returned even if didn't converge)
     """
 
-    #
+    '''
-    #Report on our progress
+    Report on our progress
-    #
+    '''
     ii = 0
 
     def log_residual(xk: ArrayLike) -> None:
         nonlocal ii
         ii += 1
         if ii % 100 == 0:
-            cur_norm = norm(A @ xk - b) / norm(b)
+            cur_norm = norm(A @ xk - b)
             logger.info(f'Solver residual at iteration {ii} : {cur_norm}')
 
     if 'callback' in kwargs:
@@ -56,9 +55,10 @@ def _scipy_qmr(
     else:
         kwargs['callback'] = log_residual
 
-    #
+    '''
-    # Run the actual solve
+    Run the actual solve
-    #
+    '''
+
     x, _ = scipy.sparse.linalg.qmr(A, b, **kwargs)
     return x
 
@@ -68,14 +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: Optional[vfdfield_t] = None,
-        pec: vfdfield_t | None = None,
+        pmc: Optional[vfdfield_t] = None,
-        pmc: vfdfield_t | None = 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,
-        E_guess: vcfdfield_t | None = None,
         ) -> vcfdfield_t:
     """
     Conjugate gradient FDFD solver using CSR sparse matrices.
@@ -102,8 +100,6 @@ def generic(
                  which doesn't return convergence info and logs the residual
                  every 100 iterations.
         matrix_solver_opts: Passed as kwargs to `matrix_solver(...)`
-        E_guess: Guess at the solution E-field. `matrix_solver` must accept an
-                `x0` argument with the same purpose.
 
     Returns:
         E-field which solves the system.
@@ -124,13 +120,6 @@ def generic(
         A = Pl @ A0 @ Pr
         b = Pl @ b0
 
-    if E_guess is not None:
-        if adjoint:
-            x0 = Pr.H @ E_guess
-        else:
-            x0 = Pl @ E_guess
-        matrix_solver_opts['x0'] = x0
-
     x = matrix_solver(A.tocsr(), b, **matrix_solver_opts)
 
     if adjoint:

meanas/fdfd/waveguide_2d.py

@@ -18,8 +18,8 @@ $$
 \begin{aligned}
 \nabla \times \vec{E}(x, y, z) &= -\imath \omega \mu \vec{H} \\
 \nabla \times \vec{H}(x, y, z) &=  \imath \omega \epsilon \vec{E} \\
-\vec{E}(x,y,z) &= (\vec{E}_t(x, y) + E_z(x, y)\vec{z}) e^{-\imath \beta z} \\
+\vec{E}(x,y,z) &= (\vec{E}_t(x, y) + E_z(x, y)\vec{z}) e^{-\gamma z} \\
-\vec{H}(x,y,z) &= (\vec{H}_t(x, y) + H_z(x, y)\vec{z}) e^{-\imath \beta z} \\
+\vec{H}(x,y,z) &= (\vec{H}_t(x, y) + H_z(x, y)\vec{z}) e^{-\gamma z} \\
 \end{aligned}
 $$
 
@@ -40,57 +40,56 @@ Substituting in our expressions for $\vec{E}$, $\vec{H}$ and discretizing:
 
 $$
 \begin{aligned}
--\imath \omega \mu_{xx} H_x &= \tilde{\partial}_y E_z + \imath \beta E_y \\
+-\imath \omega \mu_{xx} H_x &= \tilde{\partial}_y E_z + \gamma E_y \\
--\imath \omega \mu_{yy} H_y &= -\imath \beta E_x - \tilde{\partial}_x E_z \\
+-\imath \omega \mu_{yy} H_y &= -\gamma E_x - \tilde{\partial}_x E_z \\
 -\imath \omega \mu_{zz} H_z &= \tilde{\partial}_x E_y - \tilde{\partial}_y E_x \\
-\imath \omega \epsilon_{xx} E_x &= \hat{\partial}_y H_z + \imath \beta H_y \\
+\imath \omega \epsilon_{xx} E_x &= \hat{\partial}_y H_z + \gamma H_y \\
-\imath \omega \epsilon_{yy} E_y &= -\imath \beta H_x - \hat{\partial}_x H_z \\
+\imath \omega \epsilon_{yy} E_y &= -\gamma H_x - \hat{\partial}_x H_z \\
 \imath \omega \epsilon_{zz} E_z &= \hat{\partial}_x H_y - \hat{\partial}_y H_x \\
 \end{aligned}
 $$
 
 Rewrite the last three equations as
-
 $$
 \begin{aligned}
-\imath \beta H_y &=  \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z \\
+\gamma H_y &=  \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z \\
-\imath \beta H_x &= -\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z \\
+\gamma H_x &= -\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z \\
 \imath \omega E_z &= \frac{1}{\epsilon_{zz}} \hat{\partial}_x H_y - \frac{1}{\epsilon_{zz}} \hat{\partial}_y H_x \\
 \end{aligned}
 $$
 
-Now apply $\imath \beta \tilde{\partial}_x$ to the last equation,
+Now apply $\gamma \tilde{\partial}_x$ to the last equation,
-then substitute in for $\imath \beta H_x$ and $\imath \beta H_y$:
+then substitute in for $\gamma H_x$ and $\gamma H_y$:
 
 $$
 \begin{aligned}
-\imath \beta \tilde{\partial}_x \imath \omega E_z &= \imath \beta \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x H_y
+\gamma \tilde{\partial}_x \imath \omega E_z &= \gamma \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x H_y
-                                                   - \imath \beta \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y H_x \\
+                                             - \gamma \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y H_x \\
         &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x ( \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z)
          - \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (-\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z)  \\
         &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x ( \imath \omega \epsilon_{xx} E_x)
          - \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (-\imath \omega \epsilon_{yy} E_y)  \\
-\imath \beta \tilde{\partial}_x E_z &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
+\gamma \tilde{\partial}_x E_z &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
-                                     + \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y) \\
+                               + \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y) \\
 \end{aligned}
 $$
 
-With a similar approach (but using $\imath \beta \tilde{\partial}_y$ instead), we can get
+With a similar approach (but using $\gamma \tilde{\partial}_y$ instead), we can get
 
 $$
 \begin{aligned}
-\imath \beta \tilde{\partial}_y E_z &= \tilde{\partial}_y \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
+\gamma \tilde{\partial}_y E_z &= \tilde{\partial}_y \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
-                                     + \tilde{\partial}_y \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y) \\
+                               + \tilde{\partial}_y \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y) \\
 \end{aligned}
 $$
 
-We can combine this equation for $\imath \beta \tilde{\partial}_y E_z$ with
+We can combine this equation for $\gamma \tilde{\partial}_y E_z$ with
 the unused $\imath \omega \mu_{xx} H_x$ and $\imath \omega \mu_{yy} H_y$ equations to get
 
 $$
 \begin{aligned}
--\imath \omega \mu_{xx} \imath \beta H_x &=  -\beta^2 E_y + \imath \beta \tilde{\partial}_y E_z \\
+-\imath \omega \mu_{xx} \gamma H_x &=  \gamma^2 E_y + \gamma \tilde{\partial}_y E_z \\
--\imath \omega \mu_{xx} \imath \beta H_x &=  -\beta^2 E_y + \tilde{\partial}_y (
+-\imath \omega \mu_{xx} \gamma H_x &=  \gamma^2 E_y + \tilde{\partial}_y (
                                       \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
                                     + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
                                     )\\
@@ -101,24 +100,25 @@ and
 
 $$
 \begin{aligned}
--\imath \omega \mu_{yy} \imath \beta H_y &= \beta^2 E_x - \imath \beta \tilde{\partial}_x E_z \\
+-\imath \omega \mu_{yy} \gamma H_y &= -\gamma^2 E_x - \gamma \tilde{\partial}_x E_z \\
--\imath \omega \mu_{yy} \imath \beta H_y &= \beta^2 E_x - \tilde{\partial}_x (
+-\imath \omega \mu_{yy} \gamma H_y &= -\gamma^2 E_x - \tilde{\partial}_x (
                                       \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
                                     + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
                                     )\\
 \end{aligned}
 $$
 
-However, based on our rewritten equation for $\imath \beta H_x$ and the so-far unused
+However, based on our rewritten equation for $\gamma H_x$ and the so-far unused
 equation for $\imath \omega \mu_{zz} H_z$ we can also write
 
 $$
 \begin{aligned}
--\imath \omega \mu_{xx} (\imath \beta H_x) &= -\imath \omega \mu_{xx} (-\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z) \\
+-\imath \omega \mu_{xx} (\gamma H_x) &= -\imath \omega \mu_{xx} (-\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z) \\
-                 &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y + \imath \omega \mu_{xx} \hat{\partial}_x (
+                                     &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y
-                         \frac{1}{-\imath \omega \mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x)) \\
+                                        +\imath \omega \mu_{xx} \hat{\partial}_x (
-                 &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y
+                                           \frac{1}{-\imath \omega \mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x)) \\
-                         -\mu_{xx} \hat{\partial}_x \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
+                                     &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y
+                                        -\mu_{xx} \hat{\partial}_x \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
 \end{aligned}
 $$
 
@@ -126,7 +126,7 @@ and, similarly,
 
 $$
 \begin{aligned}
--\imath \omega \mu_{yy} (\imath \beta H_y) &= \omega^2 \mu_{yy} \epsilon_{xx} E_x
+-\imath \omega \mu_{yy} (\gamma H_y) &= \omega^2 \mu_{yy} \epsilon_{xx} E_x
                                            +\mu_{yy} \hat{\partial}_y \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
 \end{aligned}
 $$
@@ -135,12 +135,12 @@ By combining both pairs of expressions, we get
 
 $$
 \begin{aligned}
-\beta^2 E_x - \tilde{\partial}_x (
+-\gamma^2 E_x - \tilde{\partial}_x (
     \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
   + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
     ) &= \omega^2 \mu_{yy} \epsilon_{xx} E_x
         +\mu_{yy} \hat{\partial}_y \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
--\beta^2 E_y + \tilde{\partial}_y (
+\gamma^2 E_y + \tilde{\partial}_y (
     \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
   + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
     ) &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y
@@ -165,24 +165,24 @@ $$
                     E_y \end{bmatrix}
 $$
 
-In the literature, $\beta$ is usually used to denote the lossless/real part of the propagation constant,
+where $\gamma = \imath\beta$. In the literature, $\beta$ is usually used to denote
-but in `meanas` it is allowed to be complex.
+the lossless/real part of the propagation constant, but in `meanas` it is allowed to
+be complex.
 
 An equivalent eigenvalue problem can be formed using the $H_x$ and $H_y$ fields, if those are more convenient.
 
-Note that $E_z$ was never discretized, so $\beta$ will need adjustment to account for numerical dispersion
+Note that $E_z$ was never discretized, so $\gamma$ and $\beta$ will need adjustment
-if the result is introduced into a space with a discretized z-axis.
+to account for numerical dispersion if the result is introduced into a space with a discretized z-axis.
 
 
 """
 # TODO update module docs
 
 from typing import Any
-from collections.abc import Sequence
 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
@@ -413,13 +413,18 @@ def _normalized_fields(
     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)]
 
+    E = unvec(e, shape)
+    H = unvec(h, shape)
+
     # Find time-averaged Sz and normalize to it
     # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting
     phase = numpy.exp(-1j * -prop_phase / 2)
-    Sz_tavg = inner_product(e, h, dxes=dxes, prop_phase=prop_phase, conj_h=True).real
+    Sz_a = E[0] * numpy.conj(H[1] * phase) * dxes_real[0][1] * dxes_real[1][0]
+    Sz_b = E[1] * numpy.conj(H[0] * phase) * dxes_real[0][0] * dxes_real[1][1]
+    Sz_tavg = numpy.real(Sz_a.sum() - Sz_b.sum()) * 0.5       # 0.5 since E, H are assumed to be peak (not RMS) amplitudes
     assert Sz_tavg > 0, f'Found a mode propagating in the wrong direction! {Sz_tavg=}'
 
-    energy = numpy.real(epsilon * e.conj() * e)
+    energy = epsilon * e.conj() * e
 
     norm_amplitude = 1 / numpy.sqrt(Sz_tavg)
     norm_angle = -numpy.angle(e[energy.argmax()])       # Will randomly add a negative sign when mode is symmetric
@@ -429,7 +434,6 @@ def _normalized_fields(
     sign = numpy.sign(E_weighted[:,
                                  :max(shape[0] // 2, 1),
                                  :max(shape[1] // 2, 1)].real.sum())
-    assert sign != 0
 
     norm_factor = sign * norm_amplitude * numpy.exp(1j * norm_angle)
 
@@ -530,37 +534,10 @@ def exy2e(
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         ) -> sparse.spmatrix:
-    r"""
+    """
     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
 
-    From the operator derivation (see module docs), we have
-
-    $$
-    \imath \omega \epsilon_{zz} E_z = \hat{\partial}_x H_y - \hat{\partial}_y H_x \\
-    $$
-
-    as well as the intermediate equations
-
-    $$
-    \begin{aligned}
-    \imath \beta H_y &=  \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z \\
-    \imath \beta H_x &= -\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z \\
-    \end{aligned}
-    $$
-
-    Combining these, we get
-
-    $$
-    \begin{aligned}
-    E_z &= \frac{1}{- \omega \beta \epsilon_{zz}} ((
-             \hat{\partial}_y \hat{\partial}_x H_z
-            -\hat{\partial}_x \hat{\partial}_y H_z)
-          + \imath \omega (\hat{\partial}_x \epsilon_{xx} E_x + \hat{\partial}_y \epsilon{yy} E_y))
-        &= \frac{1}{\imath \beta \epsilon_{zz}} (\hat{\partial}_x \epsilon_{xx} E_x + \hat{\partial}_y \epsilon{yy} E_y)
-    \end{aligned}
-    $$
-
     Args:
         wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
                     It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
@@ -847,7 +824,7 @@ def sensitivity(
 
 
 def solve_modes(
-        mode_numbers: Sequence[int],
+        mode_numbers: list[int],
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
@@ -868,38 +845,32 @@ def solve_modes(
             ability to find the correct mode. Default 2.
 
     Returns:
-        e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
+        e_xys: list of vfdfield_t specifying fields
         wavenumbers: list of wavenumbers
     """
 
-    #
+    '''
-    # Solve for the largest-magnitude eigenvalue of the real operator
+    Solve for the largest-magnitude eigenvalue of the real operator
-    #
+    '''
     dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
     mu_real = None if mu is None else numpy.real(mu)
     A_r = operator_e(numpy.real(omega), dxes_real, numpy.real(epsilon), mu_real)
 
     eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
-    keep_inds = -(numpy.array(mode_numbers) + 1)
+    e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 1)]
-    e_xys = eigvecs[:, keep_inds].T
-    eigvals = eigvals[keep_inds]
 
-    #
+    '''
-    # Now solve for the eigenvector of the full operator, using the real operator's
+    Now solve for the eigenvector of the full operator, using the real operator's
-    #  eigenvector as an initial guess for Rayleigh quotient iteration.
+     eigenvector as an initial guess for Rayleigh quotient iteration.
-    #
+    '''
     A = operator_e(omega, dxes, epsilon, mu)
     for nn in range(len(mode_numbers)):
-        eigvals[nn], e_xys[nn, :] = rayleigh_quotient_iteration(A, e_xys[nn, :])
+        eigvals[nn], e_xys[:, nn] = rayleigh_quotient_iteration(A, e_xys[:, nn])
 
     # Calculate the wave-vector (force the real part to be positive)
     wavenumbers = numpy.sqrt(eigvals)
     wavenumbers *= numpy.sign(numpy.real(wavenumbers))
 
-    order = wavenumbers.argsort()[::-1]
-    e_xys = e_xys[order]
-    wavenumbers = wavenumbers[order]
-
     return e_xys, wavenumbers
 
 
@@ -921,38 +892,4 @@ def solve_mode(
     """
     kwargs['mode_numbers'] = [mode_number]
     e_xys, wavenumbers = solve_modes(*args, **kwargs)
-    return e_xys[0], wavenumbers[0]
+    return e_xys[:, 0], wavenumbers[0]
-
-
-def inner_product(    # TODO documentation
-        e1: vcfdfield_t,
-        h2: vcfdfield_t,
-        dxes: dx_lists_t,
-        prop_phase: float = 0,
-        conj_h: bool = False,
-        trapezoid: bool = False,
-        ) -> complex:
-
-    shape = [s.size for s in dxes[0]]
-
-    # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting
-    phase = numpy.exp(-1j * -prop_phase / 2)
-
-    E1 = unvec(e1, shape)
-    H2 = unvec(h2, shape) * phase
-
-    if conj_h:
-        H2 = numpy.conj(H2)
-
-    # Find time-averaged Sz and normalize to it
-    dxes_real = [[numpy.real(dxyz) for dxyz in dxeh] for dxeh in dxes]
-    if trapezoid:
-        Sz_a = numpy.trapezoid(numpy.trapezoid(E1[0] * H2[1], numpy.cumsum(dxes_real[0][1])), numpy.cumsum(dxes_real[1][0]))
-        Sz_b = numpy.trapezoid(numpy.trapezoid(E1[1] * H2[0], numpy.cumsum(dxes_real[0][0])), numpy.cumsum(dxes_real[1][1]))
-    else:
-        Sz_a = E1[0] * H2[1] * dxes_real[1][0][:, None] * dxes_real[0][1][None, :]
-        Sz_b = E1[1] * H2[0] * dxes_real[0][0][:, None] * dxes_real[1][1][None, :]
-    Sz = 0.5 * (Sz_a.sum() - Sz_b.sum())
-    return Sz
-
-

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,
         }
         ```
@@ -53,9 +51,9 @@ def solve_mode(
 
     slices = tuple(slices)
 
-    #
+    '''
-    # Solve the 2D problem in the specified plane
+    Solve the 2D problem in the specified plane
-    #
+    '''
     # Define rotation to set z as propagation direction
     order = numpy.roll(range(3), 2 - axis)
     reverse_order = numpy.roll(range(3), axis - 2)
@@ -73,10 +71,9 @@ def solve_mode(
     }
     e_xy, wavenumber_2d = waveguide_2d.solve_mode(mode_number, **args_2d)
 
-    #
+    '''
-    # Apply corrections and expand to 3D
+    Apply corrections and expand to 3D
-    #
+    '''
-
     # Correct wavenumber to account for numerical dispersion.
     wavenumber = 2 / dx_prop * numpy.arcsin(wavenumber_2d * dx_prop / 2)
 

meanas/fdfd/waveguide_cyl.py

@@ -1,102 +1,34 @@
-r"""
+"""
 Operators and helper functions for cylindrical waveguides with unchanging cross-section.
 
-Waveguide operator is derived according to 10.1364/OL.33.001848.
+WORK IN PROGRESS, CURRENTLY BROKEN
-The curl equations in the complex coordinate system become
 
-$$
+As the z-dependence is known, all the functions in this file assume a 2D grid
-\begin{aligned}
--\imath \omega \mu_{xx} H_x &= \tilde{\partial}_y E_z + \imath \beta frac{E_y}{\tilde{t}_x} \\
--\imath \omega \mu_{yy} H_y &= -\imath \beta E_x - \frac{1}{\hat{t}_x} \tilde{\partial}_x \tilde{t}_x E_z \\
--\imath \omega \mu_{zz} H_z &= \tilde{\partial}_x E_y - \tilde{\partial}_y E_x \\
-\imath \omega \epsilon_{xx} E_x &= \hat{\partial}_y H_z + \imath \beta \frac{H_y}{\hat{T}} \\
-\imath \omega \epsilon_{yy} E_y &= -\imath \beta H_x - \{1}{\tilde{t}_x} \hat{\partial}_x \hat{t}_x} H_z \\
-\imath \omega \epsilon_{zz} E_z &= \hat{\partial}_x H_y - \hat{\partial}_y H_x \\
-\end{aligned}
-$$
-
-where $t_x = 1 + \frac{\Delta_{x, m}}{R_0}$ is the grid spacing adjusted by the nominal radius $R0$.
-
-Rewrite the last three equations as
-
-$$
-\begin{aligned}
-\imath \beta H_y &=  \imath \omega \hat{t}_x \epsilon_{xx} E_x - \hat{t}_x \hat{\partial}_y H_z \\
-\imath \beta H_x &= -\imath \omega \hat{t}_x \epsilon_{yy} E_y - \hat{t}_x \hat{\partial}_x H_z \\
-\imath \omega E_z &= \frac{1}{\epsilon_{zz}} \hat{\partial}_x H_y - \frac{1}{\epsilon_{zz}} \hat{\partial}_y H_x \\
-\end{aligned}
-$$
-
-The derivation then follows the same steps as the straight waveguide, leading to the eigenvalue problem
-
-$$
-\beta^2 \begin{bmatrix} E_x \\
-                        E_y \end{bmatrix} =
-    (\omega^2 \begin{bmatrix} T_b T_b \mu_{yy} \epsilon_{xx} & 0 \\
-                                                            0 & T_a T_a \mu_{xx} \epsilon_{yy} \end{bmatrix} +
-              \begin{bmatrix} -T_b \mu_{yy} \hat{\partial}_y \\
-                               T_a \mu_{xx} \hat{\partial}_x \end{bmatrix} T_b \mu_{zz}^{-1}
-              \begin{bmatrix} -\tilde{\partial}_y & \tilde{\partial}_x \end{bmatrix} +
-      \begin{bmatrix} \tilde{\partial}_x \\
-                      \tilde{\partial}_y \end{bmatrix} T_a \epsilon_{zz}^{-1}
-                 \begin{bmatrix} \hat{\partial}_x T_b \epsilon_{xx} & \hat{\partial}_y T_a \epsilon_{yy} \end{bmatrix})
-    \begin{bmatrix} E_x \\
-                    E_y \end{bmatrix}
-$$
-
-which resembles the straight waveguide eigenproblem with additonal $T_a$ and $T_b$ terms. These
-are diagonal matrices containing the $t_x$ values:
-
-$$
-\begin{aligned}
-T_a &=  1 + \frac{\Delta_{x, m               }}{R_0}
-T_b &=  1 + \frac{\Delta_{x, m + \frac{1}{2} }}{R_0}
-\end{aligned}
-
-
-TODO: consider 10.1364/OE.20.021583 for an alternate approach
-$$
-
-As in the straight waveguide case, all the functions in this file assume a 2D grid
  (i.e. `dxes = [[[dr_e_0, dx_e_1, ...], [dy_e_0, ...]], [[dr_h_0, ...], [dy_h_0, ...]]]`).
 """
-from typing import Any, cast
+# TODO update module docs
-from collections.abc import Sequence
-import logging
 
 import numpy
-from numpy.typing import NDArray, ArrayLike
+import scipy.sparse as sparse       # type: ignore
-from scipy import sparse
 
-from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_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
-from . import waveguide_2d
-
-logger = logging.getLogger(__name__)
 
 
 def cylindrical_operator(
-        omega: float,
+        omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
-        rmin: float,
+        r0: float,
         ) -> sparse.spmatrix:
-    r"""
+    """
     Cylindrical coordinate waveguide operator of the form
 
-    $$
+    (NOTE: See 10.1364/OL.33.001848)
-        (\omega^2 \begin{bmatrix} T_b T_b \mu_{yy} \epsilon_{xx} & 0 \\
+    TODO: consider 10.1364/OE.20.021583
-                                                                0 & T_a T_a \mu_{xx} \epsilon_{yy} \end{bmatrix} +
+
-                  \begin{bmatrix} -T_b \mu_{yy} \hat{\partial}_y \\
+    TODO
-                                   T_a \mu_{xx} \hat{\partial}_x \end{bmatrix} T_b \mu_{zz}^{-1}
-                  \begin{bmatrix} -\tilde{\partial}_y & \tilde{\partial}_x \end{bmatrix} +
-          \begin{bmatrix} \tilde{\partial}_x \\
-                          \tilde{\partial}_y \end{bmatrix} T_a \epsilon_{zz}^{-1}
-                     \begin{bmatrix} \hat{\partial}_x T_b \epsilon_{xx} & \hat{\partial}_y T_a \epsilon_{yy} \end{bmatrix})
-        \begin{bmatrix} E_x \\
-                        E_y \end{bmatrix}
-    $$
 
     for use with a field vector of the form `[E_r, E_y]`.
 
@@ -106,13 +38,12 @@ def cylindrical_operator(
     which can then be solved for the eigenmodes of the system
     (an `exp(-i * wavenumber * theta)` theta-dependence is assumed for the fields).
 
-    (NOTE: See module docs and 10.1364/OL.33.001848)
-
     Args:
         omega: The angular frequency of the system
         dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
         epsilon: Vectorized dielectric constant grid
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
+        r0: Radius of curvature for the simulation. This should be the minimum value of
+            r within the simulation domain.
 
     Returns:
         Sparse matrix representation of the operator
@@ -121,34 +52,46 @@ def cylindrical_operator(
     Dfx, Dfy = deriv_forward(dxes[0])
     Dbx, Dby = deriv_back(dxes[1])
 
-    Ta, Tb = dxes2T(dxes=dxes, rmin=rmin)
+    rx = r0 + numpy.cumsum(dxes[0][0])
+    ry = r0 + dxes[0][0] / 2.0 + numpy.cumsum(dxes[1][0])
+    tx = rx / r0
+    ty = ry / r0
+
+    Tx = sparse.diags(vec(tx[:, None].repeat(dxes[0][1].size, axis=1)))
+    Ty = sparse.diags(vec(ty[:, None].repeat(dxes[1][1].size, axis=1)))
 
     eps_parts = numpy.split(epsilon, 3)
-    eps_x = sparse.diags_array(eps_parts[0])
+    eps_x = sparse.diags(eps_parts[0])
-    eps_y = sparse.diags_array(eps_parts[1])
+    eps_y = sparse.diags(eps_parts[1])
-    eps_z_inv = sparse.diags_array(1 / eps_parts[2])
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    pa = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dbx, Dby))
+    pb = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dby, Dbx))
+    a0 = Ty @ eps_x + omega**-2 * Dby @ Ty @ Dfy
+    a1 = Tx @ eps_y + omega**-2 * Dbx @ Ty @ Dfx
+    b0 = Dbx @ Ty @ Dfy
+    b1 = Dby @ Ty @ Dfx
 
-    omega2 = omega * omega
     diag = sparse.block_diag
 
-    sq0 = omega2 * diag((Tb @ Tb @ eps_x,
+    omega2 = omega * omega
-                         Ta @ Ta @ eps_y))
+
-    lin0 = sparse.vstack((-Tb @ Dby, Ta @ Dbx)) @ Tb @ sparse.hstack((-Dfy, Dfx))
+    op = (
-    lin1 = sparse.vstack((Dfx, Dfy)) @ Ta @ eps_z_inv @ sparse.hstack((Dbx @ Tb @ eps_x,
+        (omega2 * diag((Tx, Ty)) + pa) @ diag((a0, a1))
-                                                                       Dby @ Ta @ eps_y))
+        - (sparse.bmat(((None, Ty), (Tx, None))) + pb / omega2) @ diag((b0, b1))
-    op = sq0 + lin0 + lin1
+        )
     return op
 
 
-def solve_modes(
+def solve_mode(
-        mode_numbers: Sequence[int],
+        mode_number: int,
-        omega: float,
+        omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
-        rmin: float,
+        r0: float,
-        mode_margin: int = 2,
+        ) -> dict[str, complex | cfdfield_t]:
-        ) -> tuple[vcfdfield_t, NDArray[numpy.complex128]]:
     """
+    TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
      of the bent waveguide with the specified mode number.
 
@@ -158,345 +101,48 @@ def solve_modes(
         dxes: Grid parameters [dx_e, dx_h] as described in meanas.fdmath.types.
               The first coordinate is assumed to be r, the second is y.
         epsilon: Dielectric constant
-        rmin: Radius of curvature for the simulation. This should be the minimum value of
+        r0: Radius of curvature for the simulation. This should be the minimum value of
-               r within the simulation domain.
+            r within the simulation domain.
 
     Returns:
-        e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
+        ```
-        angular_wavenumbers: list of wavenumbers in 1/rad units.
+        {
+            'E': list[NDArray[numpy.complex_]],
+            'H': list[NDArray[numpy.complex_]],
+            'wavenumber': complex,
+        }
+        ```
     """
 
-    #
+    '''
-    # Solve for the largest-magnitude eigenvalue of the real operator
+    Solve for the largest-magnitude eigenvalue of the real operator
-    #
+    '''
     dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
 
-    A_r = cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), rmin=rmin)
+    A_r = cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0)
-    eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
+    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + 3)
-    keep_inds = -(numpy.array(mode_numbers) + 1)
+    e_xy = eigvecs[:, -(mode_number + 1)]
-    e_xys = eigvecs[:, keep_inds].T
-    eigvals = eigvals[keep_inds]
 
-    #
+    '''
-    # Now solve for the eigenvector of the full operator, using the real operator's
+    Now solve for the eigenvector of the full operator, using the real operator's
-    #  eigenvector as an initial guess for Rayleigh quotient iteration.
+     eigenvector as an initial guess for Rayleigh quotient iteration.
-    #
+    '''
-    A = cylindrical_operator(omega, dxes, epsilon, rmin=rmin)
+    A = cylindrical_operator(omega, dxes, epsilon, r0)
-    for nn in range(len(mode_numbers)):
+    eigval, e_xy = rayleigh_quotient_iteration(A, e_xy)
-        eigvals[nn], e_xys[nn, :] = rayleigh_quotient_iteration(A, e_xys[nn, :])
 
     # Calculate the wave-vector (force the real part to be positive)
-    wavenumbers = numpy.sqrt(eigvals)
+    wavenumber = numpy.sqrt(eigval)
-    wavenumbers *= numpy.sign(numpy.real(wavenumbers))
+    wavenumber *= numpy.sign(numpy.real(wavenumber))
 
-    # Wavenumbers assume the mode is at rmin, which is unlikely
+    # TODO: Perform correction on wavenumber to account for numerical dispersion.
-    # Instead, return the wavenumber in inverse radians
-    angular_wavenumbers = wavenumbers * cast(complex, rmin)
 
-    order = angular_wavenumbers.argsort()[::-1]
+    shape = [d.size for d in dxes[0]]
-    e_xys = e_xys[order]
+    e_xy = numpy.hstack((e_xy, numpy.zeros(shape[0] * shape[1])))
-    angular_wavenumbers = angular_wavenumbers[order]
+    fields = {
+        'wavenumber': wavenumber,
+        'E': unvec(e_xy, shape),
+        # 'E': unvec(e, shape),
+        # 'H': unvec(h, shape),
+    }
 
-    return e_xys, angular_wavenumbers
+    return fields
-
-
-def solve_mode(
-        mode_number: int,
-        *args: Any,
-        **kwargs: Any,
-        ) -> tuple[vcfdfield_t, complex]:
-    """
-    Wrapper around `solve_modes()` that solves for a single mode.
-
-    Args:
-       mode_number: 0-indexed mode number to solve for
-       *args: passed to `solve_modes()`
-       **kwargs: passed to `solve_modes()`
-
-    Returns:
-        (e_xy, angular_wavenumber)
-    """
-    kwargs['mode_numbers'] = [mode_number]
-    e_xys, angular_wavenumbers = solve_modes(*args, **kwargs)
-    return e_xys[0], angular_wavenumbers[0]
-
-
-def linear_wavenumbers(
-        e_xys: vcfdfield_t,
-        angular_wavenumbers: ArrayLike,
-        epsilon: vfdfield_t,
-        dxes: dx_lists_t,
-        rmin: float,
-        ) -> NDArray[numpy.complex128]:
-    """
-    Calculate linear wavenumbers (1/distance) based on angular wavenumbers (1/rad)
-      and the mode's energy distribution.
-
-    Args:
-        e_xys: Vectorized mode fields with shape (num_modes, 2 * x *y)
-        angular_wavenumbers: Wavenumbers assuming fields have theta-dependence of
-            `exp(-i * angular_wavenumber * theta)`. They should satisfy
-            `operator_e() @ e_xy == (angular_wavenumber / rmin) ** 2 * e_xy`
-        epsilon: Vectorized dielectric constant grid with shape (3, x, y)
-        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
-
-    Returns:
-        NDArray containing the calculated linear (1/distance) wavenumbers
-    """
-    angular_wavenumbers = numpy.asarray(angular_wavenumbers)
-    mode_radii = numpy.empty_like(angular_wavenumbers, dtype=float)
-
-    wavenumbers = numpy.empty_like(angular_wavenumbers)
-    shape2d = (len(dxes[0][0]), len(dxes[0][1]))
-    epsilon2d = unvec(epsilon, shape2d)[:2]
-    grid_radii = rmin + numpy.cumsum(dxes[0][0])
-    for ii in range(angular_wavenumbers.size):
-        efield = unvec(e_xys[ii], shape2d, 2)
-        energy = numpy.real((efield * efield.conj()) * epsilon2d)
-        energy_vs_x = energy.sum(axis=(0, 2))
-        mode_radii[ii] = (grid_radii * energy_vs_x).sum() / energy_vs_x.sum()
-
-    logger.info(f'{mode_radii=}')
-    lin_wavenumbers = angular_wavenumbers / mode_radii
-    return lin_wavenumbers
-
-
-def exy2h(
-        angular_wavenumber: complex,
-        omega: float,
-        dxes: dx_lists_t,
-        rmin: float,
-        epsilon: vfdfield_t,
-        mu: vfdfield_t | None = 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
-
-    Args:
-        angular_wavenumber: Wavenumber assuming fields have theta-dependence of
-            `exp(-i * angular_wavenumber * theta)`. It should satisfy
-            `operator_e() @ e_xy == (angular_wavenumber / rmin) ** 2 * e_xy`
-        omega: The angular frequency of the system
-        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
-        epsilon: Vectorized dielectric constant grid
-        mu: Vectorized magnetic permeability grid (default 1 everywhere)
-
-    Returns:
-        Sparse matrix representing the operator.
-    """
-    e2hop = e2h(angular_wavenumber=angular_wavenumber, omega=omega, dxes=dxes, rmin=rmin, mu=mu)
-    return e2hop @ exy2e(angular_wavenumber=angular_wavenumber, omega=omega, dxes=dxes, rmin=rmin, epsilon=epsilon)
-
-
-def exy2e(
-        angular_wavenumber: complex,
-        omega: float,
-        dxes: dx_lists_t,
-        rmin: float,
-        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
-
-    Unlike the straight waveguide case, the H_z components do not cancel and must be calculated
-    from E_x and E_y in order to then calculate E_z.
-
-    Args:
-        angular_wavenumber: Wavenumber assuming fields have theta-dependence of
-            `exp(-i * angular_wavenumber * theta)`. It should satisfy
-            `operator_e() @ e_xy == (angular_wavenumber / rmin) ** 2 * e_xy`
-        omega: The angular frequency of the system
-        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
-        epsilon: Vectorized dielectric constant grid
-
-    Returns:
-        Sparse matrix representing the operator.
-    """
-    Dfx, Dfy = deriv_forward(dxes[0])
-    Dbx, Dby = deriv_back(dxes[1])
-    wavenumber = angular_wavenumber / rmin
-
-    Ta, Tb = dxes2T(dxes=dxes, rmin=rmin)
-    Tai = sparse.diags_array(1 / Ta.diagonal())
-    Tbi = sparse.diags_array(1 / Tb.diagonal())
-
-    epsilon_parts = numpy.split(epsilon, 3)
-    epsilon_x, epsilon_y = (sparse.diags_array(epsi) for epsi in epsilon_parts[:2])
-    epsilon_z_inv = sparse.diags_array(1 / epsilon_parts[2])
-
-    n_pts = dxes[0][0].size * dxes[0][1].size
-    zeros = sparse.coo_array((n_pts, n_pts))
-    keep_x = sparse.block_array([[sparse.eye_array(n_pts), None], [None, zeros]])
-    keep_y = sparse.block_array([[zeros, None], [None, sparse.eye_array(n_pts)]])
-
-    mu_z = numpy.ones(n_pts)
-    mu_z_inv = sparse.diags_array(1 / mu_z)
-    exy2hz = 1 / (-1j * omega) * mu_z_inv @ sparse.hstack((Dfy, -Dfx))
-    hxy2ez = 1 / (1j * omega) * epsilon_z_inv @ sparse.hstack((Dby, -Dbx))
-
-    exy2hy = Tb / (1j * wavenumber) @ (-1j * omega * sparse.hstack((epsilon_x, zeros)) - Dby @ exy2hz)
-    exy2hx = Tb / (1j * wavenumber) @ ( 1j * omega * sparse.hstack((zeros, epsilon_y)) - Tai @ Dbx @ Tb @ exy2hz)
-
-    exy2ez = hxy2ez @ sparse.vstack((exy2hx, exy2hy))
-
-    op = sparse.vstack((sparse.eye_array(2 * n_pts),
-                        exy2ez))
-    return op
-
-
-def e2h(
-        angular_wavenumber: complex,
-        omega: float,
-        dxes: dx_lists_t,
-        rmin: float,
-        mu: vfdfield_t | None = None
-        ) -> sparse.spmatrix:
-    r"""
-    Returns an operator which, when applied to a vectorized E eigenfield, produces
-     the vectorized H eigenfield.
-
-    This operator is created directly from the initial coordinate-transformed equations:
-    $$
-    \begin{aligned}
-    \imath \omega \epsilon_{xx} E_x &= \hat{\partial}_y H_z + \imath \beta \frac{H_y}{\hat{T}} \\
-    \imath \omega \epsilon_{yy} E_y &= -\imath \beta H_x - \{1}{\tilde{t}_x} \hat{\partial}_x \hat{t}_x} H_z \\
-    \imath \omega \epsilon_{zz} E_z &= \hat{\partial}_x H_y - \hat{\partial}_y H_x \\
-    \end{aligned}
-    $$
-
-    Args:
-        angular_wavenumber: Wavenumber assuming fields have theta-dependence of
-            `exp(-i * angular_wavenumber * theta)`. It should satisfy
-            `operator_e() @ e_xy == (angular_wavenumber / rmin) ** 2 * e_xy`
-        omega: The angular frequency of the system
-        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
-        mu: Vectorized magnetic permeability grid (default 1 everywhere)
-
-    Returns:
-        Sparse matrix representation of the operator.
-    """
-    Dfx, Dfy = deriv_forward(dxes[0])
-    Ta, Tb = dxes2T(dxes=dxes, rmin=rmin)
-    Tai = sparse.diags_array(1 / Ta.diagonal())
-    Tbi = sparse.diags_array(1 / Tb.diagonal())
-
-    jB = 1j * angular_wavenumber / rmin
-    op = sparse.block_array([[    None, -jB * Tai,           -Dfy],
-                             [jB * Tbi,      None, Tbi @ Dfx @ Ta],
-                             [     Dfy,      -Dfx,           None]]) / (-1j * omega)
-    if mu is not None:
-        op = sparse.diags_array(1 / mu) @ op
-    return op
-
-
-def dxes2T(
-        dxes: dx_lists_t,
-        rmin: float,
-        ) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
-    r"""
-    Returns the $T_a$ and $T_b$ diagonal matrices which are used to apply the cylindrical
-      coordinate transformation in various operators.
-
-    Args:
-        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
-
-    Returns:
-        Sparse matrix representations of the operators Ta and Tb
-    """
-    ra = rmin + numpy.cumsum(dxes[0][0])                      # Radius at Ey points
-    rb = rmin + dxes[0][0] / 2.0 + numpy.cumsum(dxes[1][0])   # Radius at Ex points
-    ta = ra / rmin
-    tb = rb / rmin
-
-    Ta = sparse.diags_array(vec(ta[:, None].repeat(dxes[0][1].size, axis=1)))
-    Tb = sparse.diags_array(vec(tb[:, None].repeat(dxes[1][1].size, axis=1)))
-    return Ta, Tb
-
-
-def normalized_fields_e(
-        e_xy: ArrayLike,
-        angular_wavenumber: complex,
-        omega: float,
-        dxes: dx_lists_t,
-        rmin: float,
-        epsilon: vfdfield_t,
-        mu: vfdfield_t | None = None,
-        prop_phase: float = 0,
-        ) -> 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.
-
-    Args:
-        e_xy: Vector containing E_x and E_y fields
-        angular_wavenumber: Wavenumber assuming fields have theta-dependence of
-            `exp(-i * angular_wavenumber * theta)`. It should satisfy
-            `operator_e() @ e_xy == (angular_wavenumber / rmin) ** 2 * e_xy`
-        omega: The angular frequency of the system
-        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
-        rmin: Radius at the left edge of the simulation domain (at minimum 'x')
-        epsilon: Vectorized dielectric constant grid
-        mu: Vectorized magnetic permeability grid (default 1 everywhere)
-        prop_phase: Phase shift `(dz * corrected_wavenumber)` over 1 cell in propagation direction.
-                    Default 0 (continuous propagation direction, i.e. dz->0).
-
-    Returns:
-        `(e, h)`, where each field is vectorized, normalized,
-        and contains all three vector components.
-    """
-    e = exy2e(angular_wavenumber=angular_wavenumber, omega=omega, dxes=dxes, rmin=rmin, epsilon=epsilon) @ e_xy
-    h = exy2h(angular_wavenumber=angular_wavenumber, omega=omega, dxes=dxes, rmin=rmin, epsilon=epsilon, mu=mu) @ e_xy
-    e_norm, h_norm = _normalized_fields(e=e, h=h, omega=omega, dxes=dxes, rmin=rmin, epsilon=epsilon,
-                                        mu=mu, prop_phase=prop_phase)
-    return e_norm, h_norm
-
-
-def _normalized_fields(
-        e: vcfdfield_t,
-        h: vcfdfield_t,
-        omega: complex,
-        dxes: dx_lists_t,
-        rmin: float,
-        epsilon: vfdfield_t,
-        mu: vfdfield_t | None = None,
-        prop_phase: float = 0,
-        ) -> tuple[vcfdfield_t, vcfdfield_t]:
-    h *= -1
-    # TODO documentation for normalized_fields
-    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)]
-
-    # Find time-averaged Sz and normalize to it
-    # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting
-    phase = numpy.exp(-1j * -prop_phase / 2)
-    Sz_tavg = waveguide_2d.inner_product(e, h, dxes=dxes, prop_phase=prop_phase, conj_h=True).real   # Note, using linear poynting vector
-    assert Sz_tavg > 0, f'Found a mode propagating in the wrong direction! {Sz_tavg=}'
-
-    energy = numpy.real(epsilon * e.conj() * e)
-
-    norm_amplitude = 1 / numpy.sqrt(Sz_tavg)
-    norm_angle = -numpy.angle(e[energy.argmax()])       # Will randomly add a negative sign when mode is symmetric
-
-    # Try to break symmetry to assign a consistent sign [experimental]
-    E_weighted = unvec(e * energy * numpy.exp(1j * norm_angle), shape)
-    sign = numpy.sign(E_weighted[:,
-                                 :max(shape[0] // 2, 1),
-                                 :max(shape[1] // 2, 1)].real.sum())
-    assert sign != 0
-
-    norm_factor = sign * norm_amplitude * numpy.exp(1j * norm_angle)
-
-    print('\nAAA\n', waveguide_2d.inner_product(e, h, dxes, prop_phase=prop_phase))
-    e *= norm_factor
-    h *= norm_factor
-    print(f'{sign=} {norm_amplitude=} {norm_angle=} {prop_phase=}')
-    print(waveguide_2d.inner_product(e, h, dxes, prop_phase=prop_phase))
-
-    return e, h

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 | complexfloating]] | 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 | complexfloating]] | 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 | complexfloating]] | 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 | complexfloating]] | 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 | complexfloating]] | 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 | complexfloating]] | 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, complexfloating
+import scipy.sparse as sparse   # type: ignore
-from scipy import sparse
 
 from .types import vfdfield_t
 
@@ -34,8 +33,8 @@ def shift_circ(
     if axis not in range(len(shape)):
         raise Exception(f'Invalid direction: {axis}, shape is {shape}')
 
-    shifts = [abs(shift_distance) if a == axis else 0 for a in range(len(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)
@@ -82,8 +81,8 @@ def shift_with_mirror(
         v = numpy.where(v < 0, - 1 - v, v)
         return v
 
-    shifts = [shift_distance if a == axis else 0 for a in range(len(shape))]
+    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 | complexfloating]],
+        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 | complexfloating]],
+        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 | complexfloating]],
+        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 | complexfloating]],
+        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 | complexfloating]]]
+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 | complexfloating]]]
    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 | complexfloating]]]
+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
 
@@ -28,16 +27,14 @@ def vec(f: cfdfield_t) -> vcfdfield_t:
 def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t:
     pass
 
-def vec(
+def vec(f: fdfield_t | cfdfield_t | ArrayLike | None) -> vfdfield_t | vcfdfield_t | None:
-        f: fdfield_t | cfdfield_t | ArrayLike | None,
-        ) -> vfdfield_t | vcfdfield_t | None:
     """
-    Create a 1D ndarray from a vector field which spans a 1-3D region.
+    Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
 
     Returns `None` if called with `f=None`.
 
     Args:
-        f: A vector field, e.g. `[f_x, f_y, f_z]` where each `f_` component is a 1- to
+        f: A vector field, `[f_x, f_y, f_z]` where each `f_` component is a 1- to
            3-D ndarray (`f_*` should all be the same size). Doesn't fail with `f=None`.
 
     Returns:
@@ -49,38 +46,33 @@ def vec(
 
 
 @overload
-def unvec(v: None, shape: Sequence[int], nvdim: int = 3) -> None:
+def unvec(v: None, shape: Sequence[int]) -> None:
     pass
 
 @overload
-def unvec(v: vfdfield_t, shape: Sequence[int], nvdim: int = 3) -> fdfield_t:
+def unvec(v: vfdfield_t, shape: Sequence[int]) -> fdfield_t:
     pass
 
 @overload
-def unvec(v: vcfdfield_t, shape: Sequence[int], nvdim: int = 3) -> cfdfield_t:
+def unvec(v: vcfdfield_t, shape: Sequence[int]) -> cfdfield_t:
     pass
 
-def unvec(
+def unvec(v: vfdfield_t | vcfdfield_t | None, shape: Sequence[int]) -> fdfield_t | cfdfield_t | None:
-        v: vfdfield_t | vcfdfield_t | None,
-        shape: Sequence[int],
-        nvdim: int = 3,
-        ) -> fdfield_t | cfdfield_t | None:
     """
-    Perform the inverse of vec(): take a 1D ndarray and output an `nvdim`-component field
+    Perform the inverse of vec(): take a 1D ndarray and output a 3D field
-     of form e.g. `[f_x, f_y, f_z]` (`nvdim=3`) where each of `f_*` is a len(shape)-dimensional
+     of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional
      ndarray.
 
     Returns `None` if called with `v=None`.
 
     Args:
-        v: 1D ndarray representing a vector field of shape shape (or None)
+        v: 1D ndarray representing a 3D vector field of shape shape (or None)
         shape: shape of the vector field
-        nvdim: Number of components in each vector
 
     Returns:
         `[f_x, f_y, f_z]` where each `f_` is a `len(shape)` dimensional ndarray (or `None`)
     """
     if v is None:
         return None
-    return v.reshape((nvdim, *shape), order='C')
+    return v.reshape((3, *shape), order='C')
 

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