refactor pss2coo and add from numpy method in backend

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,14 @@
 
 - Add `quoperator` method to get `QuOperator` representation of the circuit unitary
 
+- Add `coo_sparse_matrix_from_numpy` method on backend, where the scipy coo matrix is converted to sparse tensor in corresponding backend
+
+- Add sparse tensor to scipy coo matrix implementation in `numpy` method
+
+### Changed
+
+- `tc.quantum.PauliStringSum2COO`, `tc.quantum.PauliStringSum2Dense`, and `tc.quantum.heisenberg_hamiltonian` now return the tensor in current backend format if `numpy` option sets to False. (Breaking change: previously, the return are fixed in TensorFlow format)
+
 ## 0.1.0
 
 ### Added

tensorcircuit/backends/abstract_backend.py

@@ -957,6 +957,21 @@ def coo_sparse_matrix(
             "Backend '{}' has not implemented `coo`.".format(self.name)
         )
 
+    def coo_sparse_matrix_from_numpy(self: Any, a: Tensor) -> Tensor:
+        """
+        Generate the coo format sparse matrix from scipy coo sparse matrix.
+
+        :param a: Scipy coo format sparse matrix
+        :type a: Tensor
+        :return: SparseTensor in backend format
+        :rtype: Tensor
+        """
+        return self.coo_sparse_matrix(
+            indices=np.array([a.row, a.col]).T,
+            values=a.data,
+            shape=a.shape,
+        )
+
     def sparse_dense_matmul(
         self: Any,
         sp_a: Tensor,

tensorcircuit/backends/jax_backend.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
 import numpy as np
+from scipy.sparse import coo_matrix
 import tensornetwork
 from tensornetwork.backends.jax import jax_backend
 
@@ -290,6 +291,10 @@ def kron(self, a: Tensor, b: Tensor) -> Tensor:
         return jnp.kron(a, b)
 
     def numpy(self, a: Tensor) -> Tensor:
+        if self.is_sparse(a):
+            return coo_matrix(
+                (a.data, (a.indices[:, 0], a.indices[:, 1])), shape=a.shape
+            )
         return np.array(a)
 
     def i(self, dtype: Any = None) -> Tensor:

tensorcircuit/backends/tensorflow_backend.py

@@ -7,6 +7,7 @@
 from operator import mul
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
+from scipy.sparse import coo_matrix
 import tensornetwork
 from tensornetwork.backends.tensorflow import tensorflow_backend
 
@@ -326,6 +327,11 @@ def kron(self, a: Tensor, b: Tensor) -> Tensor:
         )
 
     def numpy(self, a: Tensor) -> Tensor:
+        if self.is_sparse(a):
+            return coo_matrix(
+                (a.values, (a.indices[:, 0], a.indices[:, 1])), shape=a.get_shape()
+            )
+
         return a.numpy()  # only valid in eager mode
 
     def i(self, dtype: Any = None) -> Tensor:

tensorcircuit/quantum.py

@@ -1104,7 +1104,7 @@ def quimb2qop(qb_mpo: Any) -> QuOperator:
 
 try:
     compiled_jit = partial(get_backend("tensorflow").jit, jit_compile=True)
-    # TODO(@refraction-ray): at least make the final returned sparse tensor backend agnostic?
+
     def heisenberg_hamiltonian(
         g: Graph,
         hzz: float = 1.0,
@@ -1196,7 +1196,7 @@ def heisenberg_hamiltonian(
             r = PauliStringSum2COO_numpy(ls, weight)
             if numpy:
                 return r
-            return _numpy2tf_sparse(r)
+            return backend.coo_sparse_matrix_from_numpy(r)
         return PauliStringSum2Dense(ls, weight, numpy=numpy)
 
     def PauliStringSum2Dense(
@@ -1205,50 +1205,60 @@ def PauliStringSum2Dense(
         numpy: bool = False,
     ) -> Tensor:
         """
-        Generate tensorflow dense matrix from Pauli string sum
+        Generate dense matrix from Pauli string sum
 
         :param ls: 2D Tensor, each row is for a Pauli string,
             e.g. [1, 0, 0, 3, 2] is for :math:`X_0Z_3Y_4`
         :type ls: Sequence[Sequence[int]]
         :param weight: 1D Tensor, each element corresponds the weight for each Pauli string
             defaults to None (all Pauli strings weight 1.0)
         :type weight: Optional[Sequence[float]], optional
+        :param numpy: default False. If True, return numpy coo
+            else return backend compatible sparse tensor
+        :type numpy: bool
         :return: the tensorflow dense matrix
         :rtype: Tensor
         """
         sparsem = PauliStringSum2COO_numpy(ls, weight)
         if numpy:
             return sparsem.todense()
-        sparsem = _numpy2tf_sparse(sparsem)
-        densem = get_backend("tensorflow").to_dense(sparsem)
+        sparsem = backend.coo_sparse_matrix_from_numpy(sparsem)
+        densem = backend.to_dense(sparsem)
         return densem
 
-    def _tf2numpy_sparse(a: Tensor) -> Tensor:
-        return get_backend("numpy").coo_sparse_matrix(
-            indices=a.indices,
-            values=a.values,
-            shape=a.get_shape(),
-        )
-
-    def _numpy2tf_sparse(a: Tensor) -> Tensor:
-        return get_backend("tensorflow").coo_sparse_matrix(
-            indices=np.array([a.row, a.col]).T,
-            values=a.data,
-            shape=a.shape,
-        )
+    # already implemented as backend method
+    #
+    # def _tf2numpy_sparse(a: Tensor) -> Tensor:
+    #     return get_backend("numpy").coo_sparse_matrix(
+    #         indices=a.indices,
+    #         values=a.values,
+    #         shape=a.get_shape(),
+    #     )
+
+    # def _numpy2tf_sparse(a: Tensor) -> Tensor:
+    #     return get_backend("tensorflow").coo_sparse_matrix(
+    #         indices=np.array([a.row, a.col]).T,
+    #         values=a.data,
+    #         shape=a.shape,
+    #     )
 
-    def PauliStringSum2COO_numpy(
-        ls: Sequence[Sequence[int]], weight: Optional[Sequence[float]] = None
+    def PauliStringSum2COO(
+        ls: Sequence[Sequence[int]],
+        weight: Optional[Sequence[float]] = None,
+        numpy: bool = False,
     ) -> Tensor:
         """
-        Generate scipy sparse matrix from Pauli string sum
+        Generate sparse tensor from Pauli string sum
 
         :param ls: 2D Tensor, each row is for a Pauli string,
             e.g. [1, 0, 0, 3, 2] is for :math:`X_0Z_3Y_4`
         :type ls: Sequence[Sequence[int]]
         :param weight: 1D Tensor, each element corresponds the weight for each Pauli string
             defaults to None (all Pauli strings weight 1.0)
         :type weight: Optional[Sequence[float]], optional
+        :param numpy: default False. If True, return numpy coo
+            else return backend compatible sparse tensor
+        :type numpy: bool
         :return: the scipy coo sparse matrix
         :rtype: Tensor
         """
@@ -1267,11 +1277,16 @@ def PauliStringSum2COO_numpy(
             shape=(s, s),
         )
         for i in range(nterms):
-            rsparse += _tf2numpy_sparse(PauliString2COO(ls[i], weight[i]))  # type: ignore
+            rsparse += get_backend("tensorflow").numpy(PauliString2COO(ls[i], weight[i]))  # type: ignore
             # auto transformed into csr format!!
-        return rsparse.tocoo()
+        rsparse = rsparse.tocoo()
+        if numpy:
+            return rsparse
+        return backend.coo_sparse_matrix_from_numpy(rsparse)
 
-    def PauliStringSum2COO(
+    PauliStringSum2COO_numpy = partial(PauliStringSum2COO, numpy=True)
+
+    def PauliStringSum2COO_tf(
         ls: Sequence[Sequence[int]], weight: Optional[Sequence[float]] = None
     ) -> Tensor:
         """

tests/test_backends.py

@@ -639,20 +639,28 @@ def test_sparse_methods(backend):
     indices = tc.backend.cast(indices, "int64")
     spa = tc.backend.coo_sparse_matrix(indices, values, shape=[4, 4])
     vec = tc.backend.ones([4, 1])
+    da = np.array(
+        [[1, 0, 0, 0], [0, 2, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], dtype=np.complex64
+    )
     assert tc.backend.is_sparse(spa) is True
     assert tc.backend.is_sparse(vec) is False
     np.testing.assert_allclose(
         tc.backend.to_dense(spa),
-        np.array(
-            [[1, 0, 0, 0], [0, 2, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], dtype=np.complex64
-        ),
+        da,
         atol=1e-5,
     )
     np.testing.assert_allclose(
         tc.backend.sparse_dense_matmul(spa, vec),
         np.array([[1], [2], [0], [0]], dtype=np.complex64),
         atol=1e-5,
     )
+    spa_np = tc.backend.numpy(spa)
+    np.testing.assert_allclose(spa_np.todense(), da, atol=1e-6)
+    np.testing.assert_allclose(
+        tc.backend.to_dense(tc.backend.coo_sparse_matrix_from_numpy(spa_np)),
+        da,
+        atol=1e-5,
+    )
 
 
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])

tests/test_quantum.py

@@ -373,8 +373,8 @@ def test_tn2qop(backend):
     h1 = qu_mpo.eval_matrix()
     g = tc.templates.graphs.Line1D(nwires, pbc=False)
     h2 = tc.quantum.heisenberg_hamiltonian(
-        g, hzz=0, hxx=1, hyy=0, hz=1, hx=0, hy=0, sparse=False
-    ).numpy()
+        g, hzz=0, hxx=1, hyy=0, hz=1, hx=0, hy=0, sparse=False, numpy=True
+    )
     np.testing.assert_allclose(h1, h2, atol=1e-5)
 
 
@@ -390,8 +390,8 @@ def test_qb2qop(backend):
     h1 = qu_mpo.eval_matrix()
     g = tc.templates.graphs.Line1D(nwires, pbc=True)
     h2 = tc.quantum.heisenberg_hamiltonian(
-        g, hzz=1, hxx=0, hyy=0, hz=0, hx=-1, hy=0, sparse=False
-    ).numpy()
+        g, hzz=1, hxx=0, hyy=0, hz=0, hx=-1, hy=0, sparse=False, numpy=True
+    )
     np.testing.assert_allclose(h1, h2, atol=1e-5)
 
     # in out edge order test
@@ -404,6 +404,6 @@ def test_qb2qop(backend):
     m1 = h.eval_matrix()
     g = tc.templates.graphs.Line1D(3, pbc=False)
     m2 = tc.quantum.heisenberg_hamiltonian(
-        g, hzz=0, hxx=0, hyy=0, hz=0, hy=0.5, hx=0.5, sparse=False
-    ).numpy()
+        g, hzz=0, hxx=0, hyy=0, hz=0, hy=0.5, hx=0.5, sparse=False, numpy=True
+    )
     np.testing.assert_allclose(m1, m2, atol=1e-5)