add enable lightcone jittable example

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Add `enable_lightcone` option in circuit `expectation` method, where only gates within casual lightcone of local observable is contracted.
 
+- Add `benchmark` function into utils
+
 ### Fixed
 
 - Fixed a vital bug on circuit expectation evaluation, a wrongly transposed operator connection is fixed.

examples/lightcone_simplify.py

@@ -0,0 +1,64 @@
+"""
+comparison between expectation evaluation with/wo lightcone simplification
+"""
+
+import numpy as np
+import tensorcircuit as tc
+
+K = tc.set_backend("tensorflow")
+
+
+def brickwall_ansatz(c, params, gatename, nlayers):
+    n = c._nqubits
+    params = K.reshape(params, [nlayers, n, 2])
+    for j in range(nlayers):
+        for i in range(0, n, 2):
+            getattr(c, gatename)(i, (i + 1) % n, theta=params[j, i, 0])
+        for i in range(1, n, 2):
+            getattr(c, gatename)(i, (i + 1) % n, theta=params[j, i, 1])
+    return c
+
+
+def loss(params, n, nlayers, enable_lightcone):
+    c = tc.Circuit(n)
+    for i in range(n):
+        c.h(i)
+    c = brickwall_ansatz(c, params, "rzz", nlayers)
+    expz = K.stack(
+        [c.expectation_ps(z=[i], enable_lightcone=enable_lightcone) for i in range(n)]
+    )
+    return K.real(K.sum(expz))
+
+
+vg1 = K.jit(K.value_and_grad(loss), static_argnums=(1, 2, 3))
+
+
+def efficiency():
+    for n in range(6, 40, 4):
+        for nlayers in range(2, 6, 2):
+            print(n, nlayers)
+            print("w lightcone")
+            (v2, g2), _, _ = tc.utils.benchmark(
+                vg1, K.ones([nlayers * n * 2]), n, nlayers, True
+            )
+            if n < 16:
+                print("wo lightcone")
+                (v1, g1), _, _ = tc.utils.benchmark(
+                    vg1, K.ones([nlayers * n * 2]), n, nlayers, False
+                )
+                np.testing.assert_allclose(v1, v2, atol=1e-5)
+                np.testing.assert_allclose(g1, g2, atol=1e-5)
+
+
+## further correctness check
+def correctness(n, nlayers):
+    for _ in range(5):
+        v1, g1 = vg1(K.implicit_randn([nlayers * n * 2]), n, nlayers, False)
+        v2, g2 = vg1(K.implicit_randn([nlayers * n * 2]), n, nlayers, True)
+        np.testing.assert_allclose(v1, v2, atol=1e-5)
+        np.testing.assert_allclose(g1, g2, atol=1e-5)
+
+
+if __name__ == "__main__":
+    efficiency()
+    correctness(7, 3)

tensorcircuit/circuit.py

@@ -443,6 +443,8 @@ def apply_general_gate(
         mpo: bool = False,
         ir_dict: Optional[Dict[str, Any]] = None,
     ) -> None:
+        if name is None:
+            name = ""
         gate_dict = {
             "gate": gate,
             "index": index,

tensorcircuit/simplify.py

@@ -204,9 +204,8 @@ def _full_rank_simplify(nodes: List[Any]) -> List[Any]:
 
 def _light_cone_cancel(nodes: List[Any]) -> Tuple[List[Any], bool]:
     is_changed = False
-    for n in nodes:
-        if getattr(n, "is_dagger", None) is None:
-            break
+    for ind in range(len(nodes) // 2, 0, -1):
+        n = nodes[ind]
         if n.is_dagger is True:
             continue
         noe = len(n.shape)
@@ -216,8 +215,6 @@ def _light_cone_cancel(nodes: List[Any]) -> Tuple[List[Any], bool]:
         n1, n2 = e.node1, e.node2  # one of them is n itself
         if n1 is None or n2 is None:
             continue
-        if getattr(n2, "is_dagger", None) is None:
-            break
         if n1.is_dagger == n2.is_dagger:
             continue
         if n1.id != n2.id:
@@ -276,6 +273,9 @@ def _full_light_cone_cancel(nodes: List[Any]) -> List[Any]:
     :return: _description_
     :rtype: List[Any]
     """
+    for n in nodes:
+        if getattr(n, "is_dagger", None) is None:
+            return nodes
     nodes, is_changed = _light_cone_cancel(nodes)
     while is_changed:
         nodes, is_changed = _light_cone_cancel(nodes)

tensorcircuit/utils.py

@@ -2,9 +2,10 @@
 Helper functions
 """
 
-from typing import Any, Callable, Union, Sequence
+from typing import Any, Callable, Union, Sequence, Tuple
 from functools import wraps
 import platform
+import time
 
 
 def return_partial(
@@ -96,3 +97,33 @@ def is_m1mac() -> bool:
     if not platform.platform().startswith("macOS"):
         return False
     return True
+
+
+def benchmark(
+    f: Any, *args: Any, tries: int = 5, verbose: bool = True
+) -> Tuple[Any, float, float]:
+    """
+    benchmark jittable function with staging time and running time
+
+    :param f: _description_
+    :type f: Any
+    :param tries: _description_, defaults to 5
+    :type tries: int, optional
+    :param verbose: _description_, defaults to True
+    :type verbose: bool, optional
+    :return: _description_
+    :rtype: Tuple[Any, float, float]
+    """
+    time0 = time.time()
+    r = f(*args)
+    time1 = time.time()
+    for _ in range(tries):
+        r = f(*args)
+    time2 = time.time()
+    if tries == 0:
+        rt = 0
+    else:
+        rt = (time2 - time1) / tries  # type: ignore
+    if verbose:
+        print("staging time: ", time1 - time0, "running time: ", rt)
+    return r, time1 - time0, rt