Fix cylindrical waveguide module

- Properly account for rmin vs r0
- Change return values to match waveguide_2d
- Change operator definition to look more like waveguide_2d

remaining TODO:
- Fix docs
- Further consolidate operators vs waveguide_2d
- Figure out E/H field conversions

meanas/fdfd/operators.py

@@ -40,7 +40,7 @@ __author__ = 'Jan Petykiewicz'
 def e_full(
         omega: complex,
         dxes: dx_lists_t,
-        epsilon: vfdfield_t,
+        epsilon: vfdfield_t | vcfdfield_t,
         mu: vfdfield_t | None = None,
         pec: vfdfield_t | None = None,
         pmc: vfdfield_t | None = None,

meanas/fdfd/solvers.py

@@ -35,9 +35,9 @@ 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:
@@ -56,10 +56,9 @@ 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
 

meanas/fdfd/waveguide_2d.py

@@ -179,6 +179,7 @@ to account for numerical dispersion if the result is introduced into a space wit
 # 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
@@ -845,13 +846,13 @@ def solve_modes(
             ability to find the correct mode. Default 2.
 
     Returns:
-        e_xys: list of vfdfield_t specifying fields
+        e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
         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)
@@ -859,10 +860,10 @@ def solve_modes(
     eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
     e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 1)]
 
-    '''
-    Now solve for the eigenvector of the full operator, using the real operator's
-     eigenvector as an initial guess for Rayleigh quotient iteration.
-    '''
+    #
+    # Now solve for the eigenvector of the full operator, using the real operator's
+    #  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])
@@ -871,7 +872,7 @@ def solve_modes(
     wavenumbers = numpy.sqrt(eigvals)
     wavenumbers *= numpy.sign(numpy.real(wavenumbers))
 
-    return e_xys, wavenumbers
+    return e_xys.T, wavenumbers
 
 
 def solve_mode(
@@ -892,4 +893,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]

meanas/fdfd/waveguide_3d.py

@@ -53,9 +53,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,9 +73,10 @@ 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

@@ -8,10 +8,14 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
 """
 # TODO update module docs
 
+from typing import Any
+from collections.abc import Sequence
+
 import numpy
+from numpy.typing import NDArray, ArrayLike
 from scipy import sparse
 
-from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, cfdfield_t
+from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t
 from ..fdmath.operators import deriv_forward, deriv_back
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
@@ -21,6 +25,7 @@ def cylindrical_operator(
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         r0: float,
+        rmin: float,
         ) -> sparse.spmatrix:
     """
     Cylindrical coordinate waveguide operator of the form
@@ -42,8 +47,8 @@ def cylindrical_operator(
         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
-        r0: Radius of curvature for the simulation. This should be the minimum value of
-            r within the simulation domain.
+        r0: Radius of curvature at x=0
+        rmin: Radius at the left edge of the simulation domain
 
     Returns:
         Sparse matrix representation of the operator
@@ -52,44 +57,52 @@ def cylindrical_operator(
     Dfx, Dfy = deriv_forward(dxes[0])
     Dbx, Dby = deriv_back(dxes[1])
 
-    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
+    ra = rmin + dxes[0][0] / 2.0 + numpy.cumsum(dxes[1][0])   # Radius at Ex points
+    rb = rmin + numpy.cumsum(dxes[0][0])                      # Radius at Ey points
+    ta = ra / r0
+    tb = rb / 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)))
+    Ta = sparse.diags(vec(ta[:, None].repeat(dxes[0][1].size, axis=1)))
+    Tb = sparse.diags(vec(tb[:, None].repeat(dxes[1][1].size, axis=1)))
 
     eps_parts = numpy.split(epsilon, 3)
     eps_x = sparse.diags(eps_parts[0])
     eps_y = sparse.diags(eps_parts[1])
     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
 
-    omega2 = omega * omega
+    sq0 = omega2 * diag((Tb @ Tb @ eps_x,
+                         Ta @ Ta @ eps_y))
+    lin0 = sparse.vstack((-Tb @ Dby, Ta @ Dbx)) @ Tb @ sparse.hstack((-Dfy, Dfx))
+    lin1 = sparse.vstack((Dfx, Dfy)) @ Ta @ eps_z_inv @ sparse.hstack((Dbx @ Tb @ eps_x,
+                                                                       Dby @ Ta @ eps_y))
+    # op = (
+    #     # E
+    #     omega * omega * mu_yx @ eps_xy
+    #     + mu_yx @ sparse.vstack((-Dby, Dbx)) @ mu_z_inv @ sparse.hstack((-Dfy, Dfx))
+    #     + sparse.vstack((Dfx, Dfy)) @ eps_z_inv @ sparse.hstack((Dbx, Dby)) @ eps_xy
 
-    op = (
-        (omega2 * diag((Tx, Ty)) + pa) @ diag((a0, a1))
-        - (sparse.bmat(((None, Ty), (Tx, None))) + pb / omega2) @ diag((b0, b1))
-        )
+    #     # H
+    #     omega * omega * eps_yx @ mu_xy
+    #     + eps_yx @ sparse.vstack((-Dfy, Dfx)) @ eps_z_inv @ sparse.hstack((-Dby, Dbx))
+    #     + sparse.vstack((Dbx, Dby)) @ mu_z_inv @ sparse.hstack((Dfx, Dfy)) @ mu_xy
+    #     )
+
+    op = sq0 + lin0 + lin1
     return op
 
 
-def solve_mode(
-        mode_number: int,
+def solve_modes(
+        mode_numbers: Sequence[int],
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         r0: float,
-        ) -> dict[str, complex | cfdfield_t]:
+        rmin: float,
+        mode_margin: int = 2,
+        ) -> tuple[vcfdfield_t, NDArray[numpy.complex64]]:
     """
     TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@@ -105,44 +118,50 @@ def solve_mode(
             r within the simulation domain.
 
     Returns:
-        ```
-        {
-            'E': list[NDArray[numpy.complex_]],
-            'H': list[NDArray[numpy.complex_]],
-            'wavenumber': complex,
-        }
-        ```
+        e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
+        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]
 
-    A_r = cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0)
-    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + 3)
-    e_xy = eigvecs[:, -(mode_number + 1)]
+    A_r = cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0=r0, rmin=rmin)
+    eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
+    e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 1)].T
 
-    '''
-    Now solve for the eigenvector of the full operator, using the real operator's
-     eigenvector as an initial guess for Rayleigh quotient iteration.
-    '''
-    A = cylindrical_operator(omega, dxes, epsilon, r0)
-    eigval, e_xy = rayleigh_quotient_iteration(A, e_xy)
+    #
+    # Now solve for the eigenvector of the full operator, using the real operator's
+    #  eigenvector as an initial guess for Rayleigh quotient iteration.
+    #
+    A = cylindrical_operator(omega, dxes, epsilon, r0=r0, rmin=rmin)
+    for nn in range(len(mode_numbers)):
+        eigvals[nn], e_xys[nn, :] = rayleigh_quotient_iteration(A, e_xys[nn, :])
 
     # Calculate the wave-vector (force the real part to be positive)
-    wavenumber = numpy.sqrt(eigval)
-    wavenumber *= numpy.sign(numpy.real(wavenumber))
+    wavenumbers = numpy.sqrt(eigvals)
+    wavenumbers *= numpy.sign(numpy.real(wavenumbers))
 
-    # TODO: Perform correction on wavenumber to account for numerical dispersion.
+    return e_xys, wavenumbers
 
-    shape = [d.size for d in dxes[0]]
-    e_xy = numpy.hstack((e_xy, numpy.zeros(shape[0] * shape[1])))
-    fields = {
-        'wavenumber': wavenumber,
-        'E': unvec(e_xy, shape),
-        # 'E': unvec(e, shape),
-        # 'H': unvec(h, shape),
-    }
 
-    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, wavenumber)
+    """
+    kwargs['mode_numbers'] = [mode_number]
+    e_xys, wavenumbers = solve_modes(*args, **kwargs)
+    return e_xys[0], wavenumbers[0]

meanas/fdmath/operators.py

@@ -34,7 +34,7 @@ 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(3)]
+    shifts = [abs(shift_distance) if a == axis else 0 for a in range(len(shape))]
     shifted_diags = [(numpy.arange(n) + s) % n for n, s in zip(shape, shifts, strict=True)]
     ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
 
@@ -82,7 +82,7 @@ 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(3)]
+    shifts = [shift_distance if a == axis else 0 for a in range(len(shape))]
     shifted_diags = [mirrored_range(n, s) for n, s in zip(shape, shifts, strict=True)]
     ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
 

meanas/fdmath/types.py

@@ -20,7 +20,7 @@ vcfdfield_t = NDArray[complexfloating]
 """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[floating | complexfloating]]]
 """
  '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[floating | complexfloating]]]
 """Mutable version of `dx_lists_t`"""