jit the paulistrsum

Author: refraction-ray

examples/hamiltonian_building.py

@@ -9,14 +9,12 @@
 import quimb
 import scipy
 import tensorflow as tf
-import torch
 
 import tensorcircuit as tc
 
 tc.set_dtype("complex128")
 
 nwires = 20
-# tc with numpy backend outperforms quimb in large nwires
 
 print("--------------------")
 
@@ -51,6 +49,14 @@
 time1 = time.perf_counter()
 
 print("tc (jax) time: ", time1 - time0)
+
+time0 = time.perf_counter()
+h12 = tc.quantum.heisenberg_hamiltonian(
+    g, hzz=1, hxx=0, hyy=0, hz=0, hx=-1, hy=0, sparse=True
+)
+time1 = time.perf_counter()
+
+print("tc (jax) time (after jit): ", time1 - time0)
 print("--------------------")
 
 # 1.3 tc approach for TFIM (tensorflow backend)
@@ -67,6 +73,15 @@
 time1 = time.perf_counter()
 
 print("tc (tensorflow) time: ", time1 - time0)
+
+time0 = time.perf_counter()
+h13 = tc.quantum.heisenberg_hamiltonian(
+    g, hzz=1, hxx=0, hyy=0, hz=0, hx=-1, hy=0, sparse=True
+)
+time1 = time.perf_counter()
+
+print("tc (tensorflow) time (after jit): ", time1 - time0)
+
 print("--------------------")
 
 

requirements/requirements-extra.txt

@@ -4,7 +4,7 @@ qiskit-aer<1.0
 qiskit-nature
 mitiq
 cirq
-torch
+torch==2.2.2
 # jupyter
 mthree==1.1.0
 openfermion

tensorcircuit/quantum.py

@@ -1330,12 +1330,21 @@ def PauliStringSum2COO(
     if weight is None:
         weight = [1.0 for _ in range(nterms)]
     weight = num_to_tensor(weight)
+    ls = num_to_tensor(ls)
     # rsparse = get_backend("numpy").coo_sparse_matrix(
     #     indices=np.array([[0, 0]], dtype=np.int64),
     #     values=np.array([0.0], dtype=getattr(np, dtypestr)),
     #     shape=(s, s),
     # )
-    rsparses = [backend.numpy(PauliString2COO(ls[i], weight[i])) for i in range(nterms)]  # type: ignore
+    global PauliString2COO_jit
+    if backend.name not in PauliString2COO_jit:
+        PauliString2COO_jit[backend.name] = backend.jit(
+            PauliString2COO, jit_compile=True
+        )
+    rsparses = [
+        backend.numpy(PauliString2COO_jit[backend.name](ls[i], weight[i]))  # type: ignore
+        for i in range(nterms)
+    ]
     rsparse = _dc_sum(rsparses)
     # auto transformed into csr format!!
 
@@ -1372,6 +1381,7 @@ def PauliStringSum2COO_tf(
 ) -> Tensor:
     """
     Generate tensorflow sparse matrix from Pauli string sum
+    [deprecated]
 
     :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`
@@ -1416,26 +1426,37 @@ def PauliString2COO(l: Sequence[int], weight: Optional[float] = None) -> Tensor:
     :rtype: Tensor
     """
     n = len(l)
+    l = backend.cast(l, dtype="int64")
     one = num_to_tensor(0b1, dtype="int64")
     idx_x = num_to_tensor(0b0, dtype="int64")
     idx_y = num_to_tensor(0b0, dtype="int64")
     idx_z = num_to_tensor(0b0, dtype="int64")
     i = num_to_tensor(0, dtype="int64")
-    for j in l:
-        # i, j from enumerate is python, non jittable when cond using tensor
-        if j == 1:  # xi
-            idx_x += backend.left_shift(one, n - i - 1)
-        elif j == 2:  # yi
-            idx_y += backend.left_shift(one, n - i - 1)
-        elif j == 3:  # zi
-            idx_z += backend.left_shift(one, n - i - 1)
+    # for j in l:
+    for j in range(n):
+        oh = backend.onehot(l[j], 4)
+        s = backend.left_shift(one, n - i - 1)
+        oh = backend.cast(oh, dtype="int64")
+        idx_x += oh[1] * s
+        idx_y += oh[2] * s
+        idx_z += oh[3] * s
+
+        # if j == 1:  # xi
+        #     idx_x += backend.left_shift(one, n - i - 1)
+        # elif j == 2:  # yi
+        #     idx_y += backend.left_shift(one, n - i - 1)
+        # elif j == 3:  # zi
+        #     idx_z += backend.left_shift(one, n - i - 1)
         i += 1
 
     if weight is None:
         weight = num_to_tensor(1.0, dtype="complex64")
     return ps2coo_core(idx_x, idx_y, idx_z, weight, n)
 
 
+PauliString2COO_jit = {"numpy": PauliString2COO}
+
+
 def ps2coo_core(
     idx_x: Tensor, idx_y: Tensor, idx_z: Tensor, weight: Tensor, nqubits: int
 ) -> Tuple[Tensor, Tensor]: