Move eigensolver code out to separate module

fdfd_tools/eigensolvers.py

@@ -0,0 +1,95 @@
+"""
+Solvers for eigenvalue / eigenvector problems
+"""
+from typing import Tuple, List
+import numpy
+from numpy.linalg import norm
+from scipy import sparse
+import scipy.sparse.linalg as spalg
+
+
+def power_iteration(operator: sparse.spmatrix,
+                    guess_vector: numpy.ndarray = None,
+                    iterations: int = 20,
+                    ) -> Tuple[complex, numpy.ndarray]:
+    """
+    Use power iteration to estimate the dominant eigenvector of a matrix.
+
+    :param operator: Matrix to analyze.
+    :param guess_vector: Starting point for the eigenvector. Default is a randomly chosen vector.
+    :param iterations: Number of iterations to perform. Default 20.
+    :return: (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
+    """
+    if numpy.any(numpy.equal(guess_vector, None)):
+        v = numpy.random.rand(operator.shape[0])
+    else:
+        v = guess_vector
+
+    for _ in range(iterations):
+        v = operator @ v
+        v /= norm(v)
+
+    lm_eigval = v.conj() @ operator @ v
+    return lm_eigval, v
+
+
+def rayleigh_quotient_iteration(operator: sparse.spmatrix,
+                                guess_vector: numpy.ndarray,
+                                iterations: int = 40,
+                                tolerance: float = 1e-13,
+                                ) -> Tuple[complex, numpy.ndarray]:
+    """
+    Use Rayleigh quotient iteration to refine an eigenvector guess.
+
+    :param operator: Matrix to analyze.
+    :param guess_vector: Eigenvector to refine.
+    :param iterations: Maximum number of iterations to perform. Default 40.
+    :param tolerance: Stop iteration if (A - I*eigenvalue) @ v < tolerance.
+        Default 1e-13.
+    :return: (eigenvalue, eigenvector)
+    """
+    v = guess_vector
+    for _ in range(iterations):
+        eigval = v.conj() @ operator @ v
+        if norm(operator @ v - eigval * v) < tolerance:
+            break
+        v = spalg.spsolve(operator - eigval * sparse.eye(operator.shape[0]), v)
+        v /= norm(v)
+
+    return eigval, v
+
+
+def signed_eigensolve(operator: sparse.spmatrix,
+                      how_many: int,
+                      negative: bool = False,
+                      ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+    """
+    Find the largest-magnitude positive-only (or negative-only) eigenvalues and
+     eigenvectors of the provided matrix.
+
+    :param operator: Matrix to analyze.
+    :param how_many: How many eigenvalues to find.
+    :param negative: Whether to find negative-only eigenvalues.
+        Default False (positive only).
+    :return: (sorted list of eigenvalues, 2D ndarray of corresponding eigenvectors)
+        eigenvectors[:, k] corresponds to the k-th eigenvalue
+    """
+    # Use power iteration to estimate the dominant eigenvector
+    lm_eigval, _ = power_iteration(operator)
+
+    '''
+    Shift by the absolute value of the largest eigenvalue, then find a few of the
+     largest-magnitude (shifted) eigenvalues. A positive shift ensures that we find the
+     largest _positive_ eigenvalues, since any negative eigenvalues will be shifted to the
+     range 0 >= neg_eigval + abs(lm_eigval) > abs(lm_eigval)
+    '''
+    if negative:
+        shifted_operator = operator - abs(lm_eigval) * sparse.eye(operator.shape[0])
+    else:
+        shifted_operator = operator + abs(lm_eigval) * sparse.eye(operator.shape[0])
+
+    eigenvalues, eigenvectors = spalg.eigs(shifted_operator, which='LM', k=how_many, ncv=50)
+
+    k = eigenvalues.argsort()
+    return eigenvalues[k], eigenvectors[:, k]
+

fdfd_tools/waveguide.py

@@ -23,7 +23,7 @@ import numpy
 from numpy.linalg import norm
 import scipy.sparse as sparse
 
-from . import unvec, dx_lists_t, field_t, vfield_t
+from . import vec, unvec, dx_lists_t, field_t, vfield_t
 from . import operators
 
 

fdfd_tools/waveguide_mode.py

@@ -1,10 +1,10 @@
 from typing import Dict, List
 import numpy
 import scipy.sparse as sparse
-import scipy.sparse.linalg as spalg
 
 from . import vec, unvec, dx_lists_t, vfield_t, field_t
 from . import operators, waveguide, functional
+from .eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
 
 def solve_waveguide_mode_2d(mode_number: int,
@@ -12,12 +12,12 @@ def solve_waveguide_mode_2d(mode_number: int,
                             dxes: dx_lists_t,
                             epsilon: vfield_t,
                             mu: vfield_t = None,
-                            wavenumber_correction: bool = True
+                            wavenumber_correction: bool = True,
                             ) -> Dict[str, complex or field_t]:
     """
     Given a 2d region, attempts to solve for the eigenmode with the specified mode number.
 
-    :param mode_number: Number of the mode, 0-indexed
+    :param mode_number: Number of the mode, 0-indexed.
     :param omega: Angular frequency of the simulation
     :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
     :param epsilon: Dielectric constant
@@ -29,46 +29,19 @@ def solve_waveguide_mode_2d(mode_number: int,
 
     '''
     Solve for the largest-magnitude eigenvalue of the real operator
-     by using power iteration.
     '''
     dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
-
     A_r = waveguide.operator(numpy.real(omega), dxes_real, numpy.real(epsilon), numpy.real(mu))
 
-    # Use power iteration for 20 steps to estimate the dominant eigenvector
+    eigvals, eigvecs = signed_eigensolve(A_r, mode_number+3)
-    v = numpy.random.rand(A_r.shape[0])
+    v = eigvecs[:, -(mode_number + 1)]
-    for _ in range(20):
-        v = A_r @ v
-        v /= numpy.linalg.norm(v)
-
-    lm_eigval = v @ A_r @ v
-
-    '''
-    Shift by the absolute value of the largest eigenvalue, then find a few of the
-     largest-magnitude (shifted) eigenvalues. The shift ensures that we find the largest
-     _positive_ eigenvalues, since any negative eigenvalues will be shifted to the range
-     0 >= neg_eigval + abs(lm_eigval) > abs(lm_eigval)
-    '''
-    shifted_A_r = A_r + abs(lm_eigval) * sparse.eye(A_r.shape[0])
-    eigvals, eigvecs = spalg.eigs(shifted_A_r, which='LM', k=mode_number + 3, ncv=50)
-
-    # Pick the eigenvalue we want from the few we found
-    k = eigvals.argsort()[-(mode_number+1)]
-    v = eigvecs[:, k]
 
     '''
     Now solve for the eigenvector of the full operator, using the real operator's
      eigenvector as an initial guess for Rayleigh quotient iteration.
     '''
     A = waveguide.operator(omega, dxes, epsilon, mu)
-
+    v, eigval = rayleigh_quotient_iteration(A, v)
-    eigval = None
-    for _ in range(40):
-        eigval = v.conj() @ A @ v
-        if numpy.linalg.norm(A @ v - eigval * v) < 1e-13:
-            break
-        v = spalg.spsolve(A - eigval * sparse.eye(A.shape[0]), v)
-        v /= numpy.linalg.norm(v)
 
     # Calculate the wave-vector (force the real part to be positive)
     wavenumber = numpy.sqrt(eigval)