meanas/fdfd/operators.py

@@ -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,

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,9 +56,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
 

meanas/fdfd/waveguide_2d.py

@@ -179,7 +179,6 @@ 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
@@ -846,13 +845,13 @@ 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)
@@ -860,10 +859,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])
@@ -872,7 +871,7 @@ def solve_modes(
     wavenumbers = numpy.sqrt(eigvals)
     wavenumbers *= numpy.sign(numpy.real(wavenumbers))
 
-    return e_xys.T, wavenumbers
+    return e_xys, wavenumbers
 
 
 def solve_mode(
@@ -893,4 +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]

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

@@ -8,14 +8,10 @@ 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, vcfdfield_t
+from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, cfdfield_t
 from ..fdmath.operators import deriv_forward, deriv_back
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
@@ -25,7 +21,6 @@ def cylindrical_operator(
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         r0: float,
-        rmin: float,
         ) -> sparse.spmatrix:
     """
     Cylindrical coordinate waveguide operator of the form
@@ -47,8 +42,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 at x=0
-        rmin: Radius at the left edge of the simulation domain
+        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
@@ -57,52 +52,44 @@ def cylindrical_operator(
     Dfx, Dfy = deriv_forward(dxes[0])
     Dbx, Dby = deriv_back(dxes[1])
 
-    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
+    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
 
-    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)))
+    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(eps_parts[0])
     eps_y = sparse.diags(eps_parts[1])
     eps_z_inv = sparse.diags(1 / eps_parts[2])
 
-    omega2 = omega * omega
+    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
+
     diag = sparse.block_diag
 
-    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
+    omega2 = omega * omega
 
-    #     # 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
+    op = (
+        (omega2 * diag((Tx, Ty)) + pa) @ diag((a0, a1))
+        - (sparse.bmat(((None, Ty), (Tx, None))) + pb / omega2) @ diag((b0, b1))
+        )
     return op
 
 
-def solve_modes(
-        mode_numbers: Sequence[int],
+def solve_mode(
+        mode_number: int,
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
         r0: float,
-        rmin: float,
-        mode_margin: int = 2,
-        ) -> tuple[vcfdfield_t, NDArray[numpy.complex64]]:
+        ) -> dict[str, complex | cfdfield_t]:
     """
     TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@@ -118,50 +105,44 @@ def solve_modes(
             r within the simulation domain.
 
     Returns:
-        e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
-        wavenumbers: list of wavenumbers
+        ```
+        {
+            '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), r0=r0, rmin=rmin)
-    eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
-    e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 1)].T
+    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)]
 
-    #
-    # 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, :])
+    '''
+    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)
 
     # Calculate the wave-vector (force the real part to be positive)
-    wavenumbers = numpy.sqrt(eigvals)
-    wavenumbers *= numpy.sign(numpy.real(wavenumbers))
+    wavenumber = numpy.sqrt(eigval)
+    wavenumber *= numpy.sign(numpy.real(wavenumber))
 
-    return e_xys, wavenumbers
+    # TODO: Perform correction on wavenumber to account for numerical dispersion.
 
+    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),
+    }
 
-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]
+    return fields

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(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)]
     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(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)]
     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 | complexfloating]]]
+dx_lists_t = Sequence[Sequence[NDArray[floating]]]
 """
  '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[floating]]]
 """Mutable version of `dx_lists_t`"""