add conditional gate

Author: refraction-ray

CHANGELOG.md

@@ -10,6 +10,12 @@
 
 - add trigonometric methods on backends
 
+- add `conditional_gate` to support quantum ops based on previous measurment results
+
+### Changed
+
+- change the return information of `unitary_kraus` and `general_kraus` methods
+
 ## 0.0.220328
 
 ### Added

docker/Dockerfile

@@ -10,7 +10,9 @@ RUN wget -q -P /tmp \
   https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
     && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \
     && rm /tmp/Miniconda3-latest-Linux-x86_64.sh
+
 ENV PATH="/opt/conda/bin:$PATH"
+
 RUN conda install -y  \
     cudatoolkit=11.0 \
     pip \

tensorcircuit/circuit.py

@@ -787,13 +787,34 @@ def step_function(x: Tensor) -> Tensor:
         self.any(index, unitary=g)  # type: ignore
         return 0.0
 
+    def select_gate(self, which: Tensor, kraus: Sequence[Gate], *index: int) -> None:
+        """
+        Apply ``which``-th gate from ``kraus`` list, i.e. apply kraus[which]
+
+        :param which: Tensor of shape [] and dtype int
+        :type which: Tensor
+        :param kraus: A list of gate in the form of ``tc.gate`` or Tensor
+        :type kraus: Sequence[Gate]
+        :param index: the qubit lines the gate applied on
+        :type index: int
+        """
+        kraus = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
+        kraus = [gates.array_to_tensor(k) for k in kraus]
+        l = len(kraus)
+        r = backend.onehot(which, l)
+        r = backend.cast(r, dtype=dtypestr)
+        tensor = reduce(add, [r[i] * kraus[i] for i in range(l)])
+        self.any(*index, unitary=tensor)  # type: ignore
+
+    conditional_gate = select_gate
+
     def unitary_kraus2(
         self,
         kraus: Sequence[Gate],
         *index: int,
         prob: Optional[Sequence[float]] = None,
         status: Optional[float] = None,
-    ) -> float:
+    ) -> Tensor:
         # general impl from Monte Carlo trajectory depolarizing above
         # still jittable
         # speed is similar to ``unitary_kraus``
@@ -811,7 +832,19 @@ def unitary_kraus(
         *index: int,
         prob: Optional[Sequence[float]] = None,
         status: Optional[float] = None,
-    ) -> float:
+    ) -> Tensor:
+        """
+        Apply unitary gates in ``kraus`` randomly based on corresponding ``prob``.
+
+        :param kraus: List of ``tc.gates.Gate`` or just Tensors
+        :type kraus: Sequence[Gate]
+        :param prob: prob list with the same size as ``kraus``, defaults to None
+        :type prob: Optional[Sequence[float]], optional
+        :param status: random seed between 0 to 1, defaults to None
+        :type status: Optional[float], optional
+        :return: shape [] int dtype tensor indicates which kraus gate is applied
+        :rtype: Tensor
+        """
         # general impl from Monte Carlo trajectory depolarizing above
         # still jittable
 
@@ -835,9 +868,10 @@ def _unitary_kraus_template(
         get_gate_from_index: Optional[
             Callable[[Tensor, Sequence[Tensor]], Tensor]
         ] = None,
-    ) -> float:  # DRY
+    ) -> Tensor:  # DRY
         sites = len(index)
         kraus = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
+        kraus = [gates.array_to_tensor(k) for k in kraus]
         if prob is None:
             prob = [
                 backend.real(backend.trace(backend.adjoint(k) @ k) / k.shape[0])
@@ -867,7 +901,7 @@ def step_function(x: Tensor) -> Tensor:
         g = get_gate_from_index(r, kraus)
         g = backend.reshape(g, [2 for _ in range(sites * 2)])
         self.any(*index, unitary=g)  # type: ignore
-        return 0.0
+        return r
 
     def _general_kraus_tf(
         self,
@@ -942,7 +976,7 @@ def _general_kraus_2(
         kraus: Sequence[Gate],
         *index: int,
         status: Optional[float] = None,
-    ) -> float:
+    ) -> Tensor:
         # the graph building time is frustratingly slow, several minutes
         # though running time is in terms of ms
         # raw running time in terms of s
@@ -954,7 +988,7 @@ def _general_kraus_2(
         # layerwise jit technique can greatly boost the staging time, see in /examples/mcnoise_boost.py
         sites = len(index)
         kraus_tensor = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
-        # kraus_tensor = [k.tensor for k in kraus]
+        kraus_tensor = [gates.array_to_tensor(k) for k in kraus_tensor]
 
         # tn with hole
         newnodes, newfront = self._copy()
@@ -988,15 +1022,14 @@ def calculate_kraus_p(i: int) -> Tensor:
             for w, k in zip(prob, kraus_tensor)
         ]
 
-        self.unitary_kraus2(new_kraus, *index, prob=prob, status=status)
-        return 0.0
+        return self.unitary_kraus2(new_kraus, *index, prob=prob, status=status)
 
     def general_kraus(
         self,
         kraus: Sequence[Gate],
         *index: int,
         status: Optional[float] = None,
-    ) -> float:
+    ) -> Tensor:
         """
         Monte Carlo trajectory simulation of general Kraus channel whose Kraus operators cannot be
         amplified to unitary operators. For unitary operators composed Kraus channel, :py:meth:`unitary_kraus`

tests/test_circuit.py

@@ -708,6 +708,33 @@ def small_tn():
     np.testing.assert_allclose(small_tn(), np.zeros([2 ** n]), atol=1e-5)
 
 
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_teleportation(backend):
+    key = tc.backend.get_random_state(42)
+
+    @tc.backend.jit
+    def f(key):
+        tc.backend.set_random_state(key)
+        c = tc.Circuit(2)
+        c.H(0)
+        r = c.general_kraus(
+            [np.array([[1.0, 0], [0, 0]]), np.array([[0, 0], [0, 1]])], 0
+        )
+        c.conditional_gate(r, [tc.gates.i(), tc.gates.x()], 1)
+        return r, c.expectation([tc.gates.z(), [1]])
+
+    keys = []
+    for _ in range(6):
+        key, subkey = tc.backend.random_split(key)
+        keys.append(subkey)
+    rs = [f(k) for k in keys]
+    for r, e in rs:
+        if tc.backend.numpy(r) > 0.5:
+            np.testing.assert_allclose(e, -1, atol=1e-5)
+        else:
+            np.testing.assert_allclose(e, 1, atol=1e-5)
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
 def test_apply_mpo_gate(backend):
     gate = tc.gates.multicontrol_gate(tc.gates._x_matrix, ctrl=[1, 0])