Add sparse solvers to xobjects

examples/solve_sparse_system.py

@@ -0,0 +1,79 @@
+# copyright ################################# #
+# This file is part of the Xobjects Package.  #
+# Copyright (c) CERN, 2021.                   #
+# ########################################### #
+
+import xobjects as xo
+import scipy.sparse as sp
+import numpy as np
+from xobjects.sparse._test_helpers import rel_residual
+
+'''
+The goal of this example is to provide a short user guide for the xo.sparse module.
+
+The xo.sparse module can be used to solve sparse linear systems of equations:
+                                    A*x = b
+where A is a sparse matrix.
+
+Currently this module only supports CPU and Cupy contexts. This module contains a 
+variety of solvers for different contexts, with consistent APIs. The intended use
+is to reuse the same LHS for many solves, so the solvers work as follows:
+
+solver(A) # Performs decomposition/factorization
+solver.solve(b) # Solves Ax = b using precomputed factors
+
+For optimal performance across backends b should be a column-major (F Contiguous)
+array or vector.
+
+The intended interface for this module is:
+
+xo.sparse.factorized_sparse_solver()
+
+The function includes detailed documentation for usage, but in short, it returns
+the best performing solver based on the context and available modules. If the 
+context is not explicitly defined, it is inferred based on the input matrix.
+
+This is how modules that build upon this functionality within xsuite should interact
+with the xo.sparse module, so that cross-platform compatibility is guaranteed.
+
+For development and convenience purposes xo.sparse provides the:
+xo.sparse.solvers module
+
+which provides the following:
+xo.sparse.solvers.CPU.
+    - scipysplu : Alias for scipy SuperLU 
+    - KLUSuperLU : Alias for PyKLU
+xo.sparse.solvers.CPU.
+    - cuDSS : nvmath.sparse.advanced.DirectSolver Wrapper with a SuperLU-like interface
+    - cachedSpSM : Rewrite of cupy's SuperLU to cache the SpSM analysis step
+                   offering massive speedups compared to cupy splu when the only
+                   available backend is SpSM
+    - cupysplu : Alias for scipy SuperLU 
+'''
+
+# Example: solve small matrix system
+n = 5
+# Create matrix
+main = 2.0 + np.abs(np.random.standard_normal(n))
+lower = np.random.standard_normal(n-1)
+upper = np.random.standard_normal(n-1)
+A = sp.diags(
+    diagonals=[lower, main, upper],
+    offsets=[-1, 0, 1],
+    format="csc"
+)
+
+# Create solver:
+print("Solver selection process: ")
+solver = xo.sparse.factorized_sparse_solver(A, verbose = True)
+
+# Generate random vector to solve:
+b = np.random.standard_normal(n)
+
+# Solve system
+x = solver.solve(b)
+
+# Calculate relative residual to assess solver:
+res = rel_residual(A,x,b)
+print("Relative residual of solution ", res)
+print("Residual should be small (<10^-12)")

tests/test_sparse_solvers.py

@@ -0,0 +1,199 @@
+# copyright ################################# #
+# This file is part of the Xobjects Package.  #
+# Copyright (c) CERN, 2021.                   #
+# ########################################### #
+
+import numpy as np
+import scipy.sparse as sp
+import xobjects as xo
+from xobjects.test_helpers import fix_random_seed
+from xobjects.context import ModuleNotAvailableError
+import warnings
+import pytest
+
+'''
+The following tests rely on computing the relative residual of the solution
+The relative residual can be defined as:
+                               || A * x - b ||
+                        η = ---------------------
+                                    ||b||
+
+Typically, the expected value for this quantity is:
+* Ideally: 1e-12 - 1e-14
+* Ill-conditioned systems: 1e-9 - 1e-10
+
+In this module, we evaluate the residual as follows:
+* We compare the residual of the reference solver (scipy) with ABS_TOL
+* We compare the residual of the KLU solver with ABS_TOL
+* We ensure that the residual of the KLU Solver is within TOLERANCE_FACTOR
+  of the reference solver. 
+
+For reference, the machine precision for FP64 is ~2.2e-16 (PRECISION)
+
+Note: The completely random sparse system is more prone to failing the tests
+due to numerical noise, often requires looser tolerances. Still worth including
+but if testing larger systems, could potentially be omitted.
+'''
+
+# ---- Helper functions ----
+def issymmetric(A, tol=0):
+    if A.shape[0] != A.shape[1]:
+        return False
+    diff = A - A.T
+    if tol == 0:
+        return diff.nnz == 0
+    else:
+        # tolerance-based check
+        return abs(diff).max() <= tol
+
+
+def assert_residual_ok(res_ref, res_solver,
+                       abs_tol=1e-12,
+                       factor=10):
+    """
+    Check that our solver's residual is both:
+      - absolutely small enough (abs_tol),
+      - not catastrophically worse than the reference (factor * res_ref).
+    """
+    # sanity: reference solver itself should be good
+    assert res_ref < abs_tol, f"Reference residual too large: {res_ref}"
+
+    # absolute bound
+    assert res_solver < abs_tol, (
+        f"Residual {res_solver} exceeds absolute tolerance {abs_tol}"
+    )
+
+    # relative bound vs reference
+    assert res_solver <= factor * res_ref, (
+        f"Residual {res_solver} not within factor {factor} of "
+        f"reference residual {res_ref}"
+    )
+
+# ---- Tests ----
+cpu_tests = [
+        ("scipySLU",  xo.ContextCpu()),
+        ("PyKLU",     xo.ContextCpu()),
+        ]
+
+cupy_tests = []
+try:
+    cupy_tests = [
+        ("cuDSS",     xo.ContextCupy()),
+        ("CachedSLU", xo.ContextCupy()),
+        ("cupySLU",   xo.ContextCupy()),
+    ]
+except ModuleNotAvailableError:
+    warnings.warn("!!!ContextCupy unavailable. "
+                  "Skipping tests for Cupy Solvers!!!")
+
+PARAMS = cpu_tests + cupy_tests
+
+SPARSE_SYSTEM_SIZE = 2000 # (n,n) matrix
+NUM_BATCHES = 10
+PRECISION = np.finfo(float).eps
+ABS_TOL = 1e-12
+TOLERANCE_FACTOR = 2
+
+def batch_vectors_as_matrix(vector_list):
+    return np.asfortranarray(np.moveaxis(np.array(vector_list),0,-1))
+
+@fix_random_seed(1337)
+def make_random_sparse_system(n, nbatches, density=0.01):
+    A = sp.random(
+        n, n,
+        density=density,
+        format="csc",
+        random_state=np.random,
+        data_rvs=np.random.standard_normal
+    )
+    # Make it nonsingular & better conditioned:
+    # Add something on the diagonal so pivots aren't tiny/zero
+    A = A + sp.eye(n, format="csc") * 5.0  # tweak factor as you like
+    b_array = []
+    if nbatches == 0:
+        b = np.random.standard_normal(n)
+        b_array.append(b)
+    else:
+        for i in range(nbatches):
+            b = np.cos(2*i/(nbatches-1)*np.pi) * np.random.standard_normal(n)
+            b_array.append(b)
+        b = batch_vectors_as_matrix(b_array)
+    solver = sp.linalg.splu(A)
+    x = solver.solve(b)
+    return (A, b, x, b_array)
+
+@fix_random_seed(1337)
+def make_tridiagonal_system(n, nbatches):
+    main = 2.0 + np.abs(np.random.standard_normal(n))
+    lower = np.random.standard_normal(n-1)
+    upper = np.random.standard_normal(n-1)
+    A = sp.diags(
+        diagonals=[lower, main, upper],
+        offsets=[-1, 0, 1],
+        format="csc"
+    )
+    b_array = []
+    if nbatches == 0:
+        b = np.random.standard_normal(n)
+        b_array.append(b)
+    else:
+        for i in range(nbatches):
+            b = np.cos(2*i/(nbatches-1)*np.pi) * np.random.standard_normal(n)
+            b_array.append(b)
+        b = batch_vectors_as_matrix(b_array)
+    solver = sp.linalg.splu(A)
+    x = solver.solve(b)
+    return (A, b, x, b_array)
+
+random_system = make_random_sparse_system(SPARSE_SYSTEM_SIZE, 0)
+tridiag_system = make_tridiagonal_system(SPARSE_SYSTEM_SIZE, 0)
+
+@pytest.mark.parametrize("test_solver,test_context", PARAMS)
+@pytest.mark.parametrize("sparse_system", [random_system, tridiag_system])
+def test_vector_solve(test_solver, test_context, sparse_system):
+    A_sp, b_sp, x_sp, _ = sparse_system
+    assert not issymmetric(A_sp)
+
+    scipy_residual = xo.sparse.rel_residual(A_sp,x_sp,b_sp)
+
+    if "Cpu" in str(test_context):
+        A = test_context.splike_lib.sparse.csc_matrix(A_sp)
+    if "Cupy" in str(test_context):
+        A = test_context.splike_lib.sparse.csr_matrix(A_sp)
+    b_test = test_context.nparray_to_context_array(b_sp)
+    b = test_context.nplike_lib.asfortranarray(b_test)
+    solver = xo.sparse.factorized_sparse_solver(A, 
+                                                force_solver = test_solver, 
+                                                context = test_context
+                                                )
+    x = solver.solve(b)
+
+    solver_residual = xo.sparse.rel_residual(A,x,b)
+    assert_residual_ok(scipy_residual,solver_residual, 
+                        abs_tol = ABS_TOL, factor = TOLERANCE_FACTOR)
+
+random_system = make_random_sparse_system(SPARSE_SYSTEM_SIZE, NUM_BATCHES)
+tridiag_system = make_tridiagonal_system(SPARSE_SYSTEM_SIZE, NUM_BATCHES)
+
+@pytest.mark.parametrize("test_solver,test_context", PARAMS)
+@pytest.mark.parametrize("sparse_system", [random_system, tridiag_system])
+def test_batched_solve(test_solver, test_context, sparse_system):
+    A_sp, b_sp, x_sp, _ = sparse_system
+    assert not issymmetric(A_sp)
+    scipy_residual = xo.sparse.rel_residual(A_sp,x_sp,b_sp)
+    if "Cpu" in str(test_context):
+        A = test_context.splike_lib.sparse.csc_matrix(A_sp)
+    if "Cupy" in str(test_context):
+        A = test_context.splike_lib.sparse.csr_matrix(A_sp)
+    b_test = test_context.nparray_to_context_array(b_sp)
+    b = test_context.nplike_lib.asfortranarray(b_test)
+    solver = xo.sparse.factorized_sparse_solver(A, 
+                                                n_batches = NUM_BATCHES,
+                                                force_solver = test_solver, 
+                                                context = test_context
+                                                )
+    x = solver.solve(b)
+
+    solver_residual = xo.sparse.rel_residual(A,x,b)
+    assert_residual_ok(scipy_residual,solver_residual, 
+                        abs_tol = ABS_TOL, factor = TOLERANCE_FACTOR)

xobjects/__init__.py

@@ -19,6 +19,7 @@
 from .string import String
 from .struct import Struct, Field
 from .ref import Ref, UnionRef
+from . import sparse
 
 from .context_cpu import ContextCpu
 from .context_pyopencl import ContextPyopencl

xobjects/context.py

@@ -203,6 +203,10 @@ def __init__(self, message="Module not available"):
     def __getattr__(self, attr):
         raise NameError(self.message)
 
+class ModuleNotAvailableError(ModuleNotFoundError):
+    """Raised when an optional dependency is missing."""
+    pass
+
 
 class XContext(ABC):
     minimum_alignment = 1

xobjects/context_cupy.py

@@ -18,6 +18,7 @@
     classes_from_kernels,
     sort_classes,
     sources_from_classes,
+    ModuleNotAvailableError,
 )
 from .linkedarray import BaseLinkedArray
 from .specialize_source import specialize_source
@@ -398,6 +399,11 @@ def __init__(
         default_shared_mem_size_bytes=0,
         device=None,
     ):
+        if not _enabled:
+            raise ModuleNotAvailableError(
+                "cupy is not installed. " "ContextCupy is not available!"
+                )
+
         if device is not None:
             cupy.cuda.Device(device).use()
 

xobjects/context_pyopencl.py

@@ -18,6 +18,7 @@
     classes_from_kernels,
     sort_classes,
     sources_from_classes,
+    ModuleNotAvailableError,
 )
 from .linkedarray import BaseLinkedArray
 from .specialize_source import specialize_source
@@ -127,6 +128,11 @@ def __init__(
 
         """
 
+        if not _enabled:
+            raise ModuleNotAvailableError(
+                "pyopencl is not installed. ContextPyopencl is not available!"
+            )
+
         super().__init__()
 
         # TODO assume one device only

xobjects/sparse/__init__.py

@@ -0,0 +1,3 @@
+from ._sparse import factorized_sparse_solver, rel_residual
+from . import solvers
+__all__ = ["factorized_sparse_solver","solvers", "rel_residual"]