Add other solvers, umfpack and pypardiso

setup.py

@@ -115,6 +115,13 @@
     "jax": [
         "jax[cpu]",
     ],
+    "umfpack": [
+        "swig",
+        "scikit-umfpack",
+    ],
+    "pardiso": [
+        "pypardiso",
+    ],
 }
 
 CLASSIFIERS = [

tdgl/finite_volume/operators.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Tuple, Union
+from typing import Callable, Literal, Tuple, Union
 
 import numpy as np
 import scipy.sparse as sp
@@ -205,7 +205,7 @@ def __init__(
         self.divergence: Union[sp.spmatrix, None] = None
         self.mu_laplacian: Union[sp.spmatrix, None] = None
         self.mu_boundary_laplacian: Union[sp.spmatrix, None] = None
-        self.mu_laplacian_lu: Union[sp.linalg.SuperLU, None] = None
+        self.mu_laplacian_factorized: Union[Callable, None] = None
         self.psi_gradient: Union[sp.spmatrix, None] = None
         self.psi_laplacian: Union[sp.spmatrix, None] = None
         self.link_exponents: Union[np.ndarray, None] = None
@@ -221,14 +221,18 @@ def __init__(
             [edge_mesh.edges[:, 1], edge_mesh.edges[:, 0]]
         )
 
-    def build_operators(self) -> None:
+    def build_operators(self, solver: Literal["superlu", "umfpack", "pardiso"]) -> None:
         """Construct the vector potential-independent operators."""
         mesh = self.mesh
         self.mu_laplacian, _ = build_laplacian(mesh, weights=self.laplacian_weights)
-        self.mu_laplacian_lu = sp.linalg.splu(self.mu_laplacian)
         self.mu_boundary_laplacian = build_neumann_boundary_laplacian(mesh)
         self.mu_gradient = build_gradient(mesh, weights=self.gradient_weights)
         self.divergence = build_divergence(mesh)
+        if solver == "pardiso":
+            self.mu_laplacian_factorized = None
+        else:
+            sp.linalg.use_solver(useUmfpack=(solver == "umfpack"))
+            self.mu_laplacian_factorized = sp.linalg.factorized(self.mu_laplacian)
 
     def set_link_exponents(self, link_exponents: np.ndarray) -> None:
         """Set the link variables and construct the covarient gradient

tdgl/solver/options.py

@@ -1,6 +1,6 @@
 import os
 from dataclasses import dataclass
-from typing import Union
+from typing import Literal, Union
 
 
 class SolverOptionsError(ValueError):
@@ -24,6 +24,10 @@ class SolverOptions:
             given solve iteration before giving up.
         adaptive_time_step_multiplier: The factor by which to multiple the time
             step ``dt`` for each adaptive solve retry.
+        sparse_solver: One of "superlu", "umfpack", or "pardiso". "umfpack" requires
+            suitesparse, which can be installed via conda, and scikit-umfpack, which
+            can be installed via pip. "pardiso" requires pypardiso, which can be
+            installed via pip.
         terminal_psi: Fixed value for the order parameter in current terminals.
         field_units: The units for magnetic fields.
         current_units: The units for currents.
@@ -53,6 +57,7 @@ class SolverOptions:
     adaptive_window: int = 10
     max_solve_retries: int = 10
     adaptive_time_step_multiplier: float = 0.25
+    sparse_solver: Literal["superlu", "umfpack", "pardiso"] = "superlu"
     terminal_psi: Union[float, complex, None] = 0.0
     save_every: int = 100
     progress_interval: int = 0
@@ -70,6 +75,12 @@ class SolverOptions:
     def validate(self) -> None:
         if self.dt_init > self.dt_max:
             raise SolverOptionsError("dt_init must be less than or equal to dt_max.")
+        solver = self.sparse_solver = self.sparse_solver.lower()
+        valid_solvers = {"superlu", "umfpack", "pardiso"}
+        if solver not in valid_solvers:
+            raise ValueError(
+                f"sparse solver must be one of {valid_solvers!r}, got {solver},"
+            )
         if self.terminal_psi is not None and not (0 <= abs(self.terminal_psi) <= 1):
             raise SolverOptionsError(
                 "terminal_psi must be None or have absolute value in [0, 1]"

tdgl/solver/solve.py

@@ -319,12 +319,18 @@ def current_func(t):
         fixed_sites=normal_boundary_index,
         fix_psi=(terminal_psi is not None),
     )
-    operators.build_operators()
+    operators.build_operators(solver=options.sparse_solver)
     operators.set_link_exponents(vector_potential)
     divergence = operators.divergence
     mu_boundary_laplacian = operators.mu_boundary_laplacian
-    mu_laplacian_lu = operators.mu_laplacian_lu
+    mu_laplacian_factorized = operators.mu_laplacian_factorized
     mu_gradient = operators.mu_gradient
+    use_pardiso = options.sparse_solver == "pardiso"
+    if use_pardiso:
+        assert mu_laplacian_factorized is None
+        import pypardiso
+
+        mu_laplacian = operators.mu_laplacian
     # Initialize the order parameter and electric potential
     psi_init = np.ones(len(mesh.sites), dtype=np.complex128)
     if terminal_psi is not None:
@@ -427,8 +433,11 @@ def update(
             )
             # Compute the supercurrent, scalar potential, and normal current
             supercurrent = get_supercurrent(psi, operators.psi_gradient, edges)
-            lhs = (divergence @ supercurrent) - (mu_boundary_laplacian @ mu_boundary)
-            mu = mu_laplacian_lu.solve(lhs)
+            rhs = (divergence @ supercurrent) - (mu_boundary_laplacian @ mu_boundary)
+            if use_pardiso:
+                mu = pypardiso.spsolve(mu_laplacian, rhs)
+            else:
+                mu = mu_laplacian_factorized(rhs)
             normal_current = -(mu_gradient @ mu)
 
             if not options.include_screening:

tdgl/version.py

@@ -1,2 +1,2 @@
-__version_info__ = (0, 4, 0)
+__version_info__ = (0, 5, 0)
 __version__ = ".".join(map(str, __version_info__))