add sparse method on pytorch backend

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Add `tc.AnalogCircuit` for digital-analog hybrid simulation.
 
+- Add sparse matrix related methods for pytorch backend.
+
 ## v1.4.0
 
 ### Added

tensorcircuit/backends/pytorch_backend.py

@@ -9,6 +9,7 @@
 from operator import mul
 from functools import reduce, partial
 
+from scipy.sparse import coo_matrix
 import tensornetwork
 from tensornetwork.backends.pytorch import pytorch_backend
 from .abstract_backend import ExtendedBackend
@@ -302,6 +303,9 @@ def kron(self, a: Tensor, b: Tensor) -> Tensor:
         return torchlib.kron(a, b)
 
     def numpy(self, a: Tensor) -> Tensor:
+        if self.is_sparse(a):
+            a = a.coalesce()
+            return coo_matrix((a.values().numpy(), a.indices().numpy()), shape=a.shape)
         a = a.cpu()
         if a.is_conj():
             return a.resolve_conj().numpy()
@@ -381,6 +385,9 @@ def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
     def sort(self, a: Tensor, axis: int = -1) -> Tensor:
         return torchlib.sort(a, dim=axis).values
 
+    def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return torchlib.argsort(a, dim=axis)
+
     def all(self, tensor: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
         """
         Corresponds to torch.all.
@@ -467,6 +474,39 @@ def where(
     def reverse(self, a: Tensor) -> Tensor:
         return torchlib.flip(a, dims=(-1,))
 
+    def coo_sparse_matrix(
+        self, indices: Tensor, values: Tensor, shape: Tensor
+    ) -> Tensor:
+        # Convert COO format to PyTorch sparse tensor
+        indices = self.convert_to_tensor(indices)
+        return torchlib.sparse_coo_tensor(self.transpose(indices), values, shape)
+
+    def sparse_dense_matmul(
+        self,
+        sp_a: Tensor,
+        b: Tensor,
+    ) -> Tensor:
+        # Matrix multiplication between sparse and dense tensor
+        return torchlib.sparse.mm(sp_a, b)
+
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        try:
+            # Convert COO to CSR format if supported
+            return coo.to_sparse_csr()
+        except AttributeError as e:
+            if not strict:
+                return coo
+            else:
+                raise e
+
+    def to_dense(self, sp_a: Tensor) -> Tensor:
+        # Convert sparse tensor to dense
+        return sp_a.to_dense()
+
+    def is_sparse(self, a: Tensor) -> bool:
+        # Check if tensor is sparse
+        return a.is_sparse or a.is_sparse_csr  # type: ignore
+
     def tree_map(self, f: Callable[..., Any], *pytrees: Any) -> Any:
         # torch native tree_map not support multiple pytree args
         # return torchlib.utils._pytree.tree_map(f, *pytrees)

tests/test_backends.py

@@ -61,7 +61,7 @@ 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")])
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 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]))
@@ -583,7 +583,7 @@ def test_arg_cmp(backend):
     )
 
 
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_argsort(backend):
     # Test basic argsort functionality
     a = tc.array_to_tensor(np.array([3, 1, 2]), dtype="float32")
@@ -987,7 +987,7 @@ def grad(A, x):
         np.testing.assert_allclose(n_grad, a_grad, atol=1e-3)
 
 
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_sparse_methods(backend):
     values = tc.backend.convert_to_tensor(np.array([1.0, 2.0]))
     values = tc.backend.cast(values, "complex64")