add with_prob in general_krais

refraction-ray · refraction-ray · commit ec4cdadfa207 · 2023-07-09T11:19:43.000+08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,8 @@
 
 - Add multiple GPU VQE examples using jax pmap
 
+- Add `with_prob` option to `general_kraus` so that the probability of each option can be returned together
+
 - Add benchmark example showcasing new way of implementing matrix product using vmap
 
 ## 0.10.0
diff --git a/examples/mipt_pideal.py b/examples/mipt_pideal.py
@@ -0,0 +1,111 @@
+"""
+demo example of mipt in tc style, with ideal p for each history trajectory
+p is also jittable now, change parameter p doesn't trigger recompiling
+"""
+
+from functools import partial
+import time
+import numpy as np
+from scipy import stats
+import tensorcircuit as tc
+
+K = tc.set_backend("jax")
+tc.set_dtype("complex128")
+# tf backend is slow (at least on cpu)
+
+
+def delete2(pick, plist):
+    # pick = 0, 1 : return plist[pick]/(plist[0]+plist[1])
+    # pick = 2: return 1
+    indicator = (K.sign(1.5 - pick) + 1) / 2  # 0,1 : 1, 2: 0
+    p = 0
+    p += 1 - indicator
+    p += indicator / (plist[0] + plist[1]) * (plist[0] * (1 - pick) + plist[1] * pick)
+    return p
+
+
+@partial(K.jit, static_argnums=(2, 3))
+def circuit_output(random_matrix, status, n, d, p):
+    """
+    mipt circuit
+
+    :param random_matrix: a float or complex tensor containing 4*4 random haar matrix wth size [d*n, 4, 4]
+    :type random_matrix: _type_
+    :param status: a int tensor with element in 0 or 1 or 2 (no meausrement) with size d*n
+    :type status: _type_
+    :param n: number of qubits
+    :type n: _type_
+    :param d: number of depth
+    :type d: _type_
+    :param p: measurement ratio
+    :type p: float
+    :return: output state
+    """
+    random_matrix = K.reshape(random_matrix, [d, n, 4, 4])
+    status = K.reshape(status, [d, n])
+    inputs = None
+    bs_history = []
+    prob_history = []
+    for j in range(d):
+        if inputs is None:
+            c = tc.Circuit(n)
+        else:
+            c = tc.Circuit(n, inputs=inputs)
+        for i in range(0, n, 2):
+            c.unitary(i, (i + 1) % n, unitary=random_matrix[j, i])
+        for i in range(1, n, 2):
+            c.unitary(i, (i + 1) % n, unitary=random_matrix[j, i])
+        inputs = c.state()
+        c = tc.Circuit(n, inputs=inputs)
+        for i in range(n):
+            pick, plist = c.general_kraus(
+                [
+                    K.sqrt(p) * K.convert_to_tensor(np.array([[1.0, 0], [0, 0]])),
+                    K.sqrt(p) * K.convert_to_tensor(np.array([[0, 0], [0, 1.0]])),
+                    K.sqrt(1 - p) * K.eye(2),
+                ],
+                i,
+                status=status[j, i],
+                with_prob=True,
+            )
+            bs_history.append(pick)
+            prob_history.append(delete2(pick, plist))
+            inputs = c.state()
+            c = tc.Circuit(n, inputs=inputs)
+        inputs = c.state()
+        inputs /= K.norm(inputs)
+    bs_history = K.stack(bs_history)
+    prob_history = K.stack(prob_history)
+    return inputs, bs_history, prob_history, K.sum(K.log(prob_history + 1e-11))
+
+
+@partial(K.jit, static_argnums=(2, 3))
+def cals(random_matrix, status, n, d, p):
+    state, bs_history, prob_history, prob = circuit_output(
+        random_matrix, status, n, d, p
+    )
+    rho = tc.quantum.reduced_density_matrix(state, cut=[i for i in range(n // 2)])
+    return (
+        tc.quantum.entropy(rho),
+        tc.quantum.renyi_entropy(rho, k=2),
+        bs_history,
+        prob_history,
+        prob,
+    )
+
+
+if __name__ == "__main__":
+    n = 12
+    d = 12
+    st = np.random.uniform(size=[d * n])
+    ## assume all X gate instead
+    rm = [stats.unitary_group.rvs(4) for _ in range(d * n)]
+    rm = [r / np.linalg.det(r) for r in rm]
+    rm = np.stack(rm)
+    time0 = time.time()
+    print(cals(rm, st, n, d, 0.6))
+    time1 = time.time()
+    st = np.random.uniform(size=[d * n])
+    print(cals(rm, st, n, d, 0.1))
+    time2 = time.time()
+    print(f"compiling time {time1-time0}, running time {time2-time1}")
diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
@@ -507,6 +507,7 @@ def _general_kraus_2(
         kraus: Sequence[Gate],
         *index: int,
         status: Optional[float] = None,
+        with_prob: bool = False,
         name: Optional[str] = None,
     ) -> Tensor:
         # the graph building time is frustratingly slow, several minutes
@@ -554,16 +555,20 @@ def calculate_kraus_p(i: int) -> Tensor:
             k / backend.cast(backend.sqrt(w) + eps, dtypestr)
             for w, k in zip(prob, kraus_tensor)
         ]
-
-        return self.unitary_kraus(
+        pick = self.unitary_kraus(
             new_kraus, *index, prob=prob, status=status, name=name
         )
+        if with_prob is False:
+            return pick
+        else:
+            return pick, prob
 
     def general_kraus(
         self,
         kraus: Sequence[Gate],
         *index: int,
         status: Optional[float] = None,
+        with_prob: bool = False,
         name: Optional[str] = None,
     ) -> Tensor:
         """
@@ -583,7 +588,9 @@ def general_kraus(
             when the random number will be generated automatically
         :type status: Optional[float], optional
         """
-        return self._general_kraus_2(kraus, *index, status=status, name=name)
+        return self._general_kraus_2(
+            kraus, *index, status=status, with_prob=with_prob, name=name
+        )
 
     apply_general_kraus = general_kraus
 
diff --git a/tests/test_circuit.py b/tests/test_circuit.py
@@ -1573,3 +1573,57 @@ def test_fancy_circuit_indexing(backend):
     assert c.gate_count("h") == 4
     assert c.gate_count("rzz") == 2
     assert c.gate_count("rxx") == 3
+
+
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("npb")])
+def test_general_kraus(backend):
+    c = tc.Circuit(2)
+    c.h([0, 1])
+    p = 0.5
+    status = [0.3, 0.8]
+    rs = []
+    for i in range(2):
+        rs.append(
+            c.general_kraus(
+                [
+                    np.sqrt(p) * np.array([[1.0, 0], [0, 0]]),
+                    np.sqrt(p) * np.array([[0, 0], [0, 1.0]]),
+                    np.sqrt(1 - p) * np.eye(2),
+                ],
+                i,
+                status=status[i],
+            )
+        )
+    np.testing.assert_allclose(rs[0], 1)
+    np.testing.assert_allclose(rs[1], 2)
+    np.testing.assert_allclose(c.expectation_ps(z=[0]), -1, atol=1e-5)
+    np.testing.assert_allclose(c.expectation_ps(z=[1]), 0, atol=1e-5)
+
+
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("npb")])
+def test_general_kraus_with_prob(backend):
+    c = tc.Circuit(2)
+    c.h([0, 1])
+    p = 0.5
+    status = [0.3, 0.8]
+    rs = []
+    for i in range(2):
+        rs.append(
+            c.general_kraus(
+                [
+                    np.sqrt(p) * np.array([[1.0, 0], [0, 0]]),
+                    np.sqrt(p) * np.array([[0, 0], [0, 1.0]]),
+                    np.sqrt(1 - p) * np.eye(2),
+                ],
+                i,
+                status=status[i],
+                with_prob=True,
+            )
+        )
+    np.testing.assert_allclose(rs[0][0], 1)
+    np.testing.assert_allclose(rs[1][0], 2)
+    np.testing.assert_allclose(c.expectation_ps(z=[0]), -1, atol=1e-5)
+    np.testing.assert_allclose(c.expectation_ps(z=[1]), 0, atol=1e-5)
+    np.testing.assert_allclose(rs[0][1], [0.25, 0.25, 0.5], atol=1e-5)
+    np.testing.assert_allclose(rs[1][1], [0.25, 0.25, 0.5], atol=1e-5)
+    np.testing.assert_allclose(tc.backend.norm(c.state()), 1, atol=1e-5)
