Repository: cai4cai/torchsparsegradutils (language: python)

Issue:
Provide BiCGSTAB with no extra dependency
This is a follow-up from #3 

Options include adapting code from:
- PyKrilov, not in pytorch by pure numpy so easily adapted https://github.com/PythonOptimizers/pykrylov/blob/master/pykrylov/bicgstab/bicgstab.py
- https://gist.github.com/bridgesign/f421f69ad4a3858430e5e235bccde8c6 (vanilla implementation of the version on wikipedia -> may not be robust enough)


Discussion / hints:

Fix patch:
diff --git a/torchsparsegradutils/utils/__init__.py b/torchsparsegradutils/utils/__init__.py
index eaba213..e798bea 100644
--- a/torchsparsegradutils/utils/__init__.py
+++ b/torchsparsegradutils/utils/__init__.py
@@ -1,4 +1,5 @@
 from .linear_cg import linear_cg, LinearCGSettings
 from .minres import minres, MINRESSettings
+from .bicgstab import bicgstab, BICGSTABSettings
 
-__all__ = ["linear_cg", "LinearCGSettings", "minres", "MINRESSettings"]
+__all__ = ["linear_cg", "LinearCGSettings", "minres", "MINRESSettings", "bicgstab", "BICGSTABSettings"]
diff --git a/torchsparsegradutils/utils/bicgstab.py b/torchsparsegradutils/utils/bicgstab.py
new file mode 100644
index 0000000..525ace5
--- /dev/null
+++ b/torchsparsegradutils/utils/bicgstab.py
@@ -0,0 +1,185 @@
+# Code imported from https://github.com/PythonOptimizers/pykrylov/blob/master/pykrylov/bicgstab/bicgstab.py
+# Modifications to fit torchsparsegradutils
+# PyKrilov is licensed under LGPL -> use at your own risks 
+
+import torch
+import warnings
+
+from typing import NamedTuple
+import types
+
+import logging
+
+# Default (null) logger.
+_null_log = logging.getLogger('bicgstab')
+_null_log.disabled = True
+
+class BICGSTABSettings(NamedTuple):
+    matvec_max: int = None  # Max. number of matrix-vector produts (2n)
+    abstol: float = 1.0e-8 # absolute stopping tolerance
+    reltol: float = 1.0e-6 # absolute stopping tolerance
+    precon: any = None # optional preconditioner
+    logger: any = _null_log # a `logging.logger` instance.
+    
+        
+def bicgstab(
+    matmul_closure,
+    rhs,
+    settings=BICGSTABSettings(),
+):
+    """
+    A pytorch implementation of the bi-conjugate gradient stabilized
+    (Bi-CGSTAB) algorithm. Bi-CGSTAB may be used to solve unsymmetric systems
+    of linear equations, i.e., systems of the form
+        A x = b
+    where the operator A is unsymmetric and nonsingular.
+
+    Bi-CGSTAB requires 2 operator-vector products, 6 dot products and 6 daxpys
+    per iteration.
+
+    In addition, if a preconditioner is supplied, it needs to solve 2
+    preconditioning systems per iteration.
+
+    The original description appears in [VdVorst92]_. This implementation is a
+    preconditioned version of that given in [Kelley]_.
+
+    Reference:
+    .. [VdVorst92] H. Van der Vorst, *Bi-CGSTAB: A Fast and Smoothly Convergent
+                   Variant of Bi-CG for the Solution of Nonsymmetric Linear
+                   Systems*, SIAM Journal on Scientific and Statistical
+                   Computing **13** (2), pp. 631--644, 1992.
+
+    
+    .. [Kelley] C. T. Kelley, *Iterative Methods for Linear and Nonlinear
+                Equations*, number 16 in *Frontiers in Applied Mathematics*,
+                SIAM, Philadelphia, 1995.
+
+    Solve a linear system with `rhs` as right-hand side by the Bi-CGSTAB
+    method. The vector `rhs` should be a Numpy array.
+
+    :keywords:
+        :guess:      Initial guess (Numpy array, default: 0)
+        :matvec_max: Max. number of matrix-vector produts (2n)
+    """
+    n = rhs.shape[0]
+    nMatvec = 0
+
+    if torch.is_tensor(matmul_closure):
+        op = matmul_closure.matmul
+    elif callable(matmul_closure):
+        op = matmul_closure
+    else:
+        raise RuntimeError("matmul_closure must be a tensor, or a callable object!")
+
+    if settings.precon is None:
+        precon = None
+    elif torch.is_tensor(settings.precon):
+        precon = settings.precon.matmul
+    elif callable(settings.precon):
+        precon = settings.precon
+    else:
+        raise RuntimeError("settings.precon must be a tensor, or a callable object!")
+
+    # Initial guess is zero unless one is supplied
+    res_device = rhs.device
+    res_dtype = rhs.dtype
+    
+    #guess_supplied = 'guess' in kwargs.keys()
+    guess_supplied = False
+    #x = kwargs.get('guess', np.zeros(n)).astype(result_type)
+    x = torch.zeros(n, dtype=res_dtype, device=res_device)
+    #matvec_max = kwargs.get('matvec_max', 2*n)
+    matvec_max = 2*n if settings.matvec_max is None else settings.matvec_max
+
+    # Initial residual is the fixed vector
+    r0 = rhs.clone()
+    if guess_supplied:
+        r0 = rhs - op( x )
+        nMatvec += 1
+
+    rho = alpha = omega = 1.0
+    rho_next = torch.dot(r0,r0)
+    residNorm = residNorm0 = torch.abs(torch.sqrt(rho_next))
+    threshold = max( settings.abstol, settings.reltol * residNorm0 )
+
+    finished = (residNorm <= threshold or nMatvec >= matvec_max)
+
+    settings.logger.info('Initial residual = %8.2e' % residNorm0)
+    settings.logger.info('Threshold = %8.2e' % threshold)
+    hdr = '%6s  %8s' % ('Matvec', 'Residual')
+    settings.logger.info(hdr)
+    settings.logger.info('-' * len(hdr))
+
+    if not finished:
+        r = r0.clone()
+        p = torch.zeros(n, dtype=res_dtype, device=res_device)
+        v = torch.zeros(n, dtype=res_dtype, device=res_device)
+
+    while not finished:
+
+        beta = rho_next/rho * alpha/omega
+        rho = rho_next
+
+        # Update p in-place
+        p *= beta
+        p -= beta * omega * v
+        p += r
+        
+        # Compute preconditioned search direction
+        if precon is not None:
+            q = precon( p )
+        else:
+            q = p
+
+        v = op( q ); nMatvec += 1
+
+        alpha = rho/torch.dot(r0, v)
+        s = r - alpha * v
+
+        # Check for CGS termination
+        residNorm = torch.linalg.norm(s)
+
+        settings.logger.info('%6d  %8.2e' % (nMatvec, residNorm))
+
+        if residNorm <= threshold:
+            x += alpha * q
+            finished = True
+            continue
+
+        if nMatvec >= matvec_max:
+            finished = True
+            continue
+
+        if precon is not None:
+            z = precon( s )
+        else:
+            z = s
+
+        t = op( z ) ; nMatvec += 1
+        omega = torch.dot(t,s)/torch.dot(t,t)
+        rho_next = -omega * torch.dot(r0,t)
+
+        # Update residual
+        r = s - omega * t
+
+        # Update solution in-place-ish. Note that 'z *= omega' alters s if
+        # precon = None. That's ok since s is no longer needed in this iter.
+        # 'q *= alpha' would alter p.
+        z *= omega
+        x += z
+        x += alpha * q
+
+        residNorm = torch.linalg.norm(r)
+
+        settings.logger.info('%6d  %8.2e' % (nMatvec, residNorm))
+
+        if residNorm <= threshold or nMatvec >= matvec_max:
+            finished = True
+            continue
+
+
+    converged = residNorm <= threshold
+    bestSolution = x
+    residNorm = residNorm
+
+    return bestSolution


Test patch:
diff --git a/tests/test_bicgstab.py b/tests/test_bicgstab.py
new file mode 100644
index 0000000..1222c4f
--- /dev/null
+++ b/tests/test_bicgstab.py
@@ -0,0 +1,58 @@
+import torch
+import unittest
+
+from torchsparsegradutils.utils import bicgstab
+
+import numpy as np
+import scipy.sparse.linalg
+
+class BICGSTABTest(unittest.TestCase):
+    """Test bicgstab implementation"""
+    
+    def setUp(self) -> None:
+        # The device can be specialised by a daughter class
+        if not hasattr(self, "device"):
+            self.device = torch.device("cpu")
+
+    def test_bicgstab(self):
+        size = 100
+        matrix_dense = torch.randn(size, size, dtype=torch.float64, device=self.device) +  10*torch.eye(size, device=self.device)
+        matrix_sparse = matrix_dense.to_sparse_csr()
+
+        rhs = torch.randn(size, dtype=torch.float64, device=self.device)
+
+        actual = torch.linalg.solve(matrix_dense, rhs)
+        #print("actual",actual)
+
+        #mn = matrix_dense.cpu().numpy()
+        #rn = rhs.cpu().numpy()
+        #solves = scipy.sparse.linalg.bicgstab(mn,rn)
+        #print("np solves",solves)
+
+        solves = bicgstab(matrix_dense, rhs=rhs)
+        #print("our solves",solves)
+        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
+        
+        solves = bicgstab(matrix_dense.matmul, rhs=rhs)
+        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
+        
+        solves = bicgstab(matrix_sparse.matmul, rhs=rhs)
+        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
+
+        solves = bicgstab(matrix_sparse, rhs=rhs)
+        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
+        
+
+
+class BICGSTABTestCUDA(BICGSTABTest):
+    """Override superclass setUp to run on GPU"""
+
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+        self.device = torch.device("cuda")
+        super().setUp()
+
+
+if __name__ == "__main__":
+    unittest.main()