revise thermalrelax

Author: yutuer21

tensorcircuit/channels.py

@@ -350,7 +350,7 @@ def thermalrelaxationchannel(
     p1 = backend.cast(array_to_tensor(p1), dtype=cons.dtypestr)
 
     if method == "T1dom":
-        # git avaliable
+        # jit avaliable
         m0 = backend.convert_to_tensor(
             np.array([[1, 0], [0, 0]], dtype=cons.npdtype)
         )  # reset channel
@@ -387,7 +387,7 @@ def thermalrelaxationchannel(
         return Gkraus
 
     else:  # method == "general" or method == "T2dom":
-        # git avaliable
+        # jit avaliable
         choi = (1 - p1 * p_reset) * backend.convert_to_tensor(
             np.array(
                 [[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]],
@@ -698,11 +698,14 @@ def choi_to_kraus(
 
     >>> kraus = tc.channels.phasedampingchannel(0.2)
     >>> superop = tc.channels.kraus_to_choi(kraus)
-    >>> kraus_new = tc.channels.choi_to_kraus(superop)
+    >>> kraus_new = tc.channels.choi_to_kraus(superop, truncation_rules={"max_singular_values":3})
 
 
     :param choi: Choi matrix
     :type choi: Matrix
+    :param truncation_rules: A dictionary to restrict the calculation of kraus matrices. The restriction
+    can be the number of kraus terms, which is jitable. It can also be the truncattion error, which is not jitable.
+    :type truncation_rules: Dictionary
     :return: A list of Kraus operators
     :rtype: Sequence[Matrix]
     """
@@ -731,7 +734,7 @@ def choi_to_kraus(
             kraus.append(k)
 
     else:
-        if truncation_rules.get("max_singular_values", None) is None:
+        if truncation_rules.get("max_truncattion_err", None) is None:
             atol = 1e-5
         else:
             atol = truncation_rules["max_truncattion_err"]

tests/test_channels.py

@@ -86,29 +86,21 @@ def test_thermal(backend):
     time = 100
 
     t1 = 180
-    print(t1, t2)
     kraus = tc.channels.thermalrelaxationchannel(t1, t2, time, "general", 0.1)
     supop1 = tc.channels.kraus_to_super(kraus)
-    # print(kraus)
-    # print(supop1)
 
     kraus = tc.channels.thermalrelaxationchannel(t1, t2, time, "T1dom", 0.1)
     supop2 = tc.channels.kraus_to_super(kraus)
-    # print(kraus)
-    # print(supop2)
+
     np.testing.assert_allclose(supop1, supop2, atol=1e-5)
 
     t1 = 80
-    print(t1, t2)
     kraus = tc.channels.thermalrelaxationchannel(t1, t2, time, "general", 0.1)
     supop1 = tc.channels.kraus_to_super(kraus)
-    # print(kraus)
-    # print(supop1)
 
     kraus = tc.channels.thermalrelaxationchannel(t1, t2, time, "T2dom", 0.1)
     supop2 = tc.channels.kraus_to_super(kraus)
-    # print(kraus)
-    # print(supop2)
+
     np.testing.assert_allclose(supop1, supop2, atol=1e-5)
 
 
@@ -117,17 +109,26 @@ def test_noisecircuit(backend):
 
     # Monte carlo simulation
     def noisecircuit(X):
-        noise = tc.channels.thermalrelaxationchannel(300, 400, 1000, "T2dom", 0)
 
         n = 1
         c = tc.Circuit(n)
         c.x(0)
-        c.general_kraus(noise, 0, status=X)
+        # noise = tc.channels.thermalrelaxationchannel(300, 400, 1000, "T2dom", 0)
+        # c.general_kraus(noise, 0, status=X)
+        c.thermalrelaxation(
+            0,
+            t1=300,
+            t2=400,
+            time=1000,
+            method="T2dom",
+            excitedstatepopulation=0,
+            status=X,
+        )
+
         val = c.expectation_ps(z=[0])
         return val
 
-    noisec = noisecircuit
-    noisec_vmap = tc.backend.vmap(noisec, vectorized_argnums=0)
+    noisec_vmap = tc.backend.vmap(noisecircuit, vectorized_argnums=0)
     noisec_jit = tc.backend.jit(noisec_vmap)
 
     nmc = 1000
@@ -145,9 +146,7 @@ def noisecircuitdm():
         val = dmc.expectation_ps(z=[0])
         return val
 
-    noisec = noisecircuitdm
-    noisec_jit = tc.backend.jit(noisec)
-
+    noisec_jit = tc.backend.jit(noisecircuitdm)
     valuedm = noisec_jit()
 
     np.testing.assert_allclose(valuemc, valuedm, atol=1e-1)