add csr sparse support and accelerate krylov/chebyshev

Author: refraction-ray

CHANGELOG.md

@@ -10,6 +10,8 @@
 
 - Add `jaxy_scan` for numpy and jax backends, contrasting to the tf style original backend scan method.
 
+- Add `sparse_csr_from_coo` method for numpy and jax backends to convert COO format to CSR format, the latter is more efficient for `sparse_dense_matmul`.
+
 - Add `krylov_evol` method for krylov evolution.
 
 - Add `chebyshev_evol` method for chebyshev polynomial evolution.

tensorcircuit/backends/abstract_backend.py

@@ -1266,6 +1266,26 @@ def sparse_dense_matmul(
             "Backend '{}' has not implemented `sparse_dense_matmul`.".format(self.name)
         )
 
+    def sparse_csr_from_coo(self: Any, coo: Tensor, strict: bool = False) -> Tensor:
+        """
+        transform a coo matrix to a csr matrix
+
+        :param coo: a coo matrix
+        :type coo: Tensor
+        :param strict: whether to enforce the transform, defaults to False,
+            corresponding to return the coo matrix if there is no implementation for specific backend.
+        :type strict: bool, optional
+        :return: a csr matrix
+        :rtype: Tensor
+        """
+        if strict:
+            raise NotImplementedError(
+                "Backend '{}' has not implemented `sparse_csr_from_coo`.".format(
+                    self.name
+                )
+            )
+        return coo
+
     def to_dense(self: Any, sp_a: Tensor) -> Tensor:
         """
         Convert a sparse matrix to dense tensor.

tensorcircuit/backends/jax_backend.py

@@ -649,11 +649,20 @@ def sparse_dense_matmul(
     ) -> Tensor:
         return sp_a @ b
 
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        try:
+            return sparse.BCSR.from_bcoo(coo)  # type: ignore
+        except AttributeError as e:
+            if not strict:
+                return coo
+            else:
+                raise e
+
     def to_dense(self, sp_a: Tensor) -> Tensor:
         return sp_a.todense()
 
     def is_sparse(self, a: Tensor) -> bool:
-        return isinstance(a, sparse.BCOO)  # type: ignore
+        return isinstance(a, sparse.JAXSparse)  # type: ignore
 
     def device(self, a: Tensor) -> str:
         (dev,) = a.devices()

tensorcircuit/backends/numpy_backend.py

@@ -333,6 +333,9 @@ def sparse_dense_matmul(
     ) -> Tensor:
         return sp_a @ b
 
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        return coo.tocsr()
+
     def to_dense(self, sp_a: Tensor) -> Tensor:
         return sp_a.todense()
 

tensorcircuit/timeevol.py

@@ -32,6 +32,8 @@ def lanczos_iteration_scan(
     :rtype: Tuple[Tensor, Tensor]
     """
     state_size = backend.shape_tuple(initial_vector)[0]
+    if backend.is_sparse(hamiltonian):
+        hamiltonian = backend.sparse_csr_from_coo(hamiltonian)
 
     # Main scan body for the outer loop (iterating j)
     def lanczos_step(carry: Tuple[Any, ...], j: int) -> Tuple[Any, ...]:
@@ -196,6 +198,9 @@ def lanczos_iteration(
     # Add first basis vector
     basis_vectors.append(vector)
 
+    if backend.is_sparse(hamiltonian):
+        hamiltonian = backend.sparse_csr_from_coo(hamiltonian)
+
     # Lanczos iteration (fixed number of iterations for JIT compatibility)
     for j in range(subspace_dimension):
         # Calculate H|v_j>
@@ -633,6 +638,9 @@ def chebyshev_evol(
     b = (E_max + E_min) / 2.0
     tau = a * t  # Rescaled time parameter
 
+    if backend.is_sparse(hamiltonian):
+        hamiltonian = backend.sparse_csr_from_coo(hamiltonian)
+
     def apply_h_norm(psi: Any) -> Any:
         """Applies the normalized Hamiltonian to a state."""
         return ((hamiltonian @ psi) - b * psi) / a

tests/test_backends.py

@@ -61,6 +61,27 @@ def f(x):
     np.testing.assert_allclose(f(a), np.ones([2]), atol=1e-5)
 
 
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_sparse_csr_from_coo(backend):
+    # Create a sparse matrix in COO format
+    values = tc.backend.convert_to_tensor(np.array([1.0, 2.0, 3.0]))
+    values = tc.backend.cast(values, "complex64")
+    indices = tc.backend.convert_to_tensor(np.array([[0, 0], [1, 1], [2, 3]]))
+    indices = tc.backend.cast(indices, "int64")
+    coo_matrix = tc.backend.coo_sparse_matrix(indices, values, shape=[4, 4])
+
+    # Convert COO to CSR
+    csr_matrix = tc.backend.sparse_csr_from_coo(coo_matrix)
+
+    # Check that the result is still recognized as sparse
+    assert tc.backend.is_sparse(csr_matrix) is True
+
+    # Check that the conversion preserves values by comparing dense representations
+    coo_dense = tc.backend.to_dense(coo_matrix)
+    csr_dense = tc.backend.to_dense(csr_matrix)
+    np.testing.assert_allclose(coo_dense, csr_dense, atol=1e-5)
+
+
 def test_sparse_tensor_matmul_monkey_patch(tfb):
     """
     Test the monkey-patched __matmul__ method for tf.SparseTensor.