wip: N-D LinearOperator and cg

benchmarks/benchmarks/sparse_linalg_solve.py

@@ -1,14 +1,16 @@
 """
 Check the speed of the conjugate gradient solver.
 """
+import inspect
+
 import numpy as np
 from numpy.testing import assert_equal
 
 from .common import Benchmark, safe_import
 
 with safe_import():
     from scipy import linalg, sparse
-    from scipy.sparse.linalg import cg, minres, gmres, tfqmr, spsolve
+    from scipy.sparse.linalg import cg, minres, gmres, tfqmr, spsolve, LinearOperator
 with safe_import():
     from scipy.sparse.linalg import lgmres
 with safe_import():
@@ -18,7 +20,11 @@
 def _create_sparse_poisson1d(n):
     # Make Gilbert Strang's favorite matrix
     # http://www-math.mit.edu/~gs/PIX/cupcakematrix.jpg
-    P1d = sparse.diags([[-1]*(n-1), [2]*n, [-1]*(n-1)], [-1, 0, 1])
+    P1d = sparse.diags_array(
+        [[-1]*(n-1), [2]*n, [-1]*(n-1)],
+        offsets=[-1, 0, 1],
+        dtype=np.float64
+    )
     assert_equal(P1d.shape, (n, n))
     return P1d
 
@@ -27,12 +33,19 @@ def _create_sparse_poisson2d(n):
     P1d = _create_sparse_poisson1d(n)
     P2d = sparse.kronsum(P1d, P1d)
     assert_equal(P2d.shape, (n*n, n*n))
-    return P2d.tocsr()
+    return sparse.csr_array(P2d)
+
+
+def _create_sparse_poisson2d_coo(n):
+    P1d = _create_sparse_poisson1d(n)
+    P2d = sparse.kronsum(P1d, P1d)
+    assert_equal(P2d.shape, (n*n, n*n))
+    return sparse.coo_array(P2d)
 
 
 class Bench(Benchmark):
     params = [
-        [4, 6, 10, 16, 25, 40, 64, 100],
+        [4, 8, 16, 32, 64, 128, 256, 512],
         ['dense', 'spsolve', 'cg', 'minres', 'gmres', 'lgmres', 'gcrotmk',
          'tfqmr']
     ]
@@ -57,6 +70,90 @@ def time_solve(self, n, solver):
             self.mapping[solver](self.P_sparse, self.b)
 
 
+class BatchedCG(Benchmark):
+    params = [
+        [2, 4, 6, 8, 16, 32, 64],
+        [1, 10, 100, 500, 1000, 5000, 10000]
+    ]
+    param_names = ['(n,n)', 'batch_size']
+
+    def setup(self, n, batch_size):
+        if n >= 32 and batch_size >= 500:
+            raise NotImplementedError()
+        if n >= 16 and batch_size > 5000:
+            raise NotImplementedError()
+        rng = np.random.default_rng(42)
+
+        self.batched = "xp" in inspect.signature(LinearOperator.__init__).parameters
+        if self.batched:
+            P_sparse = _create_sparse_poisson2d_coo(n)
+            if batch_size > 1:
+                self.P_sparse = sparse.vstack(
+                    [P_sparse] * batch_size, format="coo"
+                ).reshape(batch_size, n*n, n*n)
+                self.b = rng.standard_normal((batch_size, n*n))
+            else:
+                self.P_sparse = P_sparse
+                self.b = rng.standard_normal(n*n)
+        else:
+            self.P_sparse = _create_sparse_poisson2d(n)
+            self.b = [rng.standard_normal(n*n) for _ in range(batch_size)]
+
+    def time_solve(self, n, batch_size):
+        if self.batched:
+            cg(self.P_sparse, self.b)
+        else:
+            for i in range(batch_size):
+                cg(self.P_sparse, self.b[i])
+
+
+def _create_dense_random(n, batch_shape=None):
+    rng = np.random.default_rng(42)
+    M = rng.standard_normal((n*n, n*n))
+    reg = 1e-3
+    if batch_shape:
+        M = np.broadcast_to(M[np.newaxis, ...], (*batch_shape, n*n, n*n))
+
+    def matvec(x):
+        return np.squeeze(M.mT @ (M @ x[..., np.newaxis]), axis=-1) + reg * x
+
+    return LinearOperator(shape=M.shape, matvec=matvec, dtype=np.float64)
+
+
+class BatchedCGDense(Benchmark):
+    params = [
+        [2, 4, 8, 16, 24],
+        [1, 10, 100, 500, 1000, 5000, 10000]
+    ]
+    param_names = ['(n,n)', 'batch_size']
+
+    def setup(self, n, batch_size):
+        if n >= 24 and batch_size > 100:
+            raise NotImplementedError()
+        if n >= 16 and batch_size > 500:
+            raise NotImplementedError()
+        rng = np.random.default_rng(42)
+
+        self.batched = "xp" in inspect.signature(LinearOperator.__init__).parameters
+        if self.batched:
+            if batch_size > 1:
+                self.A = _create_dense_random(n, batch_shape=(batch_size,))
+                self.b = rng.standard_normal((batch_size, n*n))
+            else:
+                self.A = _create_dense_random(n)
+                self.b = rng.standard_normal(n*n)
+        else:
+            self.A = _create_dense_random(n)
+            self.b = [rng.standard_normal(n*n) for _ in range(batch_size)]
+
+    def time_solve(self, n, batch_size):
+        if self.batched:
+            cg(self.A, self.b)
+        else:
+            for i in range(batch_size):
+                cg(self.A, self.b[i])
+
+
 class Lgmres(Benchmark):
     params = [
         [10, 50, 100, 1000, 10000],

doc/source/_templates/autosummary/class.rst

@@ -18,12 +18,12 @@
       .. autosummary::
          :toctree:
       {% for item in all_methods %}
-         {%- if not item.startswith('_') or item in ['__call__', '__mul__', '__getitem__', '__len__', '__pow__'] %}
+         {%- if not item.startswith('_') or item in ['__call__', '__mul__', '__getitem__', '__len__', '__pow__', '__matmul__', '__truediv__', '__add__', '__rmul__', '__rmatmul__'] %}
          {{ name }}.{{ item }}
          {%- endif -%}
       {%- endfor %}
       {% for item in inherited_members %}
-         {%- if item in ['__call__', '__mul__', '__getitem__', '__len__', '__pow__'] %}
+         {%- if item in ['__call__', '__mul__', '__getitem__', '__len__', '__pow__', '__matmul__', '__truediv__', '__add__', '__rmul__', '__rmatmul__'] %}
          {{ name }}.{{ item }}
          {%- endif -%}
       {%- endfor %}

doc/source/conf.py

@@ -296,7 +296,8 @@
 # Generate plots for example sections
 numpydoc_use_plots = True
 np_docscrape.ClassDoc.extra_public_methods = [  # should match class.rst
-    '__call__', '__mul__', '__getitem__', '__len__',
+    '__call__', '__mul__', '__getitem__', '__len__', '__pow__', '__matmul__',
+    '__truediv__', '__add__', '__rmul__', '__rmatmul__'
 ]
 
 # -----------------------------------------------------------------------------