tutorials/locc/QuantumTeleportation_EN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning the Quantum Teleportation Protocol\n",
    "\n",
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "Quantum teleportation is another important task that can be completed by Local Operations and Classical Communication (LOCC) protocols, which transfers quantum information between two spatially separated communication nodes (only classical communication channel allowed) with the help of entanglement. In this tutorial, we will first briefly review the original teleportation protocol and simulate it with Paddle Quantum. Then, we will go through how to learn a teleportation protocol with LOCCNet.\n",
    " \n",
    "## The original quantum teleportation protocol\n",
    "\n",
    "\n",
    "The original teleportation protocol was proposed by C. H. Bennett et al. in 1993 [1] and experimentally verified in 1997 with photonic platforms [2-3]. The workflow is illustrated in the figure below. Following the convention, this process requires 2 communication nodes or parties, namely $A$ (Alice) and $B$ (Bob). For simplicity, we only consider transferring a single-qubit quantum state $|\\psi\\rangle_C$ and this requires 3 qubits in total including the pre-shared maximally entangled state $|\\Phi^+\\rangle_{AB} = \\frac{1}{\\sqrt{2}}(|00\\rangle + |11\\rangle)$. Alice holds systems A and C，Bob holds system B. **Note: Only quantum information is transferred, not the physical qubits.** \n",
    "\n",
    "![teleportation](figures/teleportation-fig-circuit.jpg \"Figure 1: Quantum teleportation: Transferring quantum state $|\\psi\\rangle$ from Alice to receiver Bob.\")\n",
    "<div style=\"text-align:center\">Figure 1: Quantum teleportation: Transferring quantum state $|\\psi\\rangle$ from Alice to receiver Bob. </div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step I:** At the very beginning, the system state can be described as\n",
    "\n",
    "$$\n",
    "\\lvert\\varphi_{0}\\rangle \n",
    "= \\lvert\\psi\\rangle_{C}\\otimes \\lvert\\Phi^+\\rangle_{AB} \n",
    "= \\frac{1}{\\sqrt{2}}\\big[\\alpha\\lvert0\\rangle(\\lvert00\\rangle + \\lvert11\\rangle)+\\beta\\lvert1\\rangle(\\lvert00\\rangle + \\lvert11\\rangle)\\big],\n",
    "\\tag{1}\n",
    "$$\n",
    "\n",
    "where the quantum state Alice want to transmit is $|\\psi\\rangle_C = \\alpha|0\\rangle_C + \\beta|1\\rangle_C$ and the coefficients $\\alpha, \\beta \\in \\mathbb{C}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step II:** Alice applies a CNOT gate, and the resulting state is\n",
    "\n",
    "$$\n",
    "|\\varphi_1\\rangle  \n",
    "= \\frac{1}{\\sqrt{2}}\\big[\\alpha\\lvert0\\rangle(\\lvert00\\rangle + \\lvert11\\rangle)+\\beta\\lvert1\\rangle(\\lvert10\\rangle + \\lvert01\\rangle)\\big],\n",
    "\\tag{2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step III:** Alice applies a Hadamard gate, and the system state becomes $|\\varphi_2\\rangle$ \n",
    "\n",
    "$$\n",
    "|\\varphi_2\\rangle = \\frac{1}{2}\\big[\\alpha(\\lvert0\\rangle + \\lvert1\\rangle)(\\lvert00\\rangle + \\lvert11\\rangle)+\\beta(\\lvert0\\rangle - \\lvert1\\rangle)(\\lvert10\\rangle + \\lvert01\\rangle)\\big],\n",
    "\\tag{3}\n",
    "$$\n",
    "\n",
    "The above state can be rearranged to\n",
    "\n",
    "$$\n",
    "\\lvert\\varphi_{2}\\rangle = \\frac{1}{2}\\big[\\lvert00\\rangle(\\alpha\\lvert0\\rangle + \\beta\\lvert1\\rangle) + \\lvert01\\rangle(\\alpha\\lvert1\\rangle + \\beta\\lvert0\\rangle) + \\lvert10\\rangle(\\alpha\\lvert0\\rangle - \\beta\\lvert1\\rangle) + \\lvert11\\rangle(\\alpha\\lvert1\\rangle - \\beta\\lvert0\\rangle)\\big].\n",
    "\\tag{4}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step IV:** Alice measures both of her qubits in the computational basis $\\{|00\\rangle, |01\\rangle, |10\\rangle, |11\\rangle\\}$ and send the results $m_1m_2$ to Bob with a classical channel. There are 4 distinct possibilities: $m_1m_2 \\in \\{ 00, 01,10, 11\\}$. Then, Bob implements certain operations correspondingly on his qubit based on the received messages.\n",
    "\n",
    "\n",
    "- If the measurement result is $m_1m_2 = 00$, Bob's state will be $\\alpha\\lvert0\\rangle + \\beta\\lvert1\\rangle$, which is the state Alice want to transmit $\\lvert\\psi\\rangle_C$. No operations are needed and the teleportation is finished.\n",
    "- If the measurement result is $m_1m_2 = 01$, Bob's state will be $\\alpha\\lvert1\\rangle + \\beta\\lvert0\\rangle$. Bob needs to act the $X$ gate on his qubit.\n",
    "- If the measurement result is $m_1m_2 = 10$, Bob's state will be $\\alpha\\lvert0\\rangle - \\beta\\lvert1\\rangle$. Bob needs to act the $Z$ gate on his qubit.\n",
    "- If the measurement result is $m_1m_2 = 11$, Bob's state will be $\\alpha\\lvert1\\rangle - \\beta\\lvert0\\rangle$. Bob needs to act the $X$ gate followed by the $Z$ gate on his qubit.\n",
    "\n",
    "In the next section, we will go through how to simulate the teleportation protocol with Paddle Quantum."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation with Paddle Quantum\n",
    "\n",
    "First, we need to import all the dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:28:54.132971Z",
     "start_time": "2021-03-09T06:28:49.877088Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import paddle\n",
    "from paddle import matmul, trace\n",
    "import paddle_quantum\n",
    "from paddle_quantum.locc import LoccNet\n",
    "from paddle_quantum.qinfo import state_fidelity\n",
    "from paddle_quantum.state import bell_state, isotropic_state, random_state\n",
    "# Change to density matrix mode\n",
    "paddle_quantum.set_backend('density_matrix')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the quantum state, and define the quantum circuit and teleportation protocol."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:28:54.197393Z",
     "start_time": "2021-03-09T06:28:54.155480Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC, self).__init__()\n",
    "        # Add the first party Alice \n",
    "        # The first parameter 2 stands for how many qubits A holds\n",
    "        # The second parameter records the name of this party\n",
    "        self.add_new_party(2, party_name=\"Alice\")\n",
    "        # Add the second party Bob\n",
    "        # The first parameter 1 stands for how many qubits B holds\n",
    "        # The second parameter records the name of this party\n",
    "        self.add_new_party(1, party_name=\"Bob\")\n",
    "\n",
    "        # Create a bell state\n",
    "        _state = bell_state(2)\n",
    "        # _state = isotropic_state(2, 0.8)\n",
    "        # Generate random pure quantum states for teleportation\n",
    "        self.state_C = random_state(num_qubits=1, rank=1)\n",
    "\n",
    "        # Initialize the system by distributing states between Alice and Bob\n",
    "        # (\"Alice\", 0) refers to qubit C\n",
    "        # (\"Alice\", 1) refers to qubit A\n",
    "        # (\"Bob\", 0) refers to qubit B\n",
    "        # print('Pre-shared entanglement state is:\\n', _state.numpy())\n",
    "        self.set_init_state(self.state_C, [(\"Alice\", 0)])\n",
    "        self.set_init_state(_state, [(\"Alice\", 1), (\"Bob\", 0)])\n",
    "\n",
    "        # Create Alice's local operations \n",
    "        self.cirA = self.create_ansatz(\"Alice\")\n",
    "        self.cirA.cnot([0, 1])\n",
    "        self.cirA.h(0)\n",
    "        # Create Bob's local operations\n",
    "        self.cirB = [self.create_ansatz(\"Bob\") for _ in range(4)]\n",
    "        self.cirB[1].x(0)\n",
    "        self.cirB[2].z(0)\n",
    "        self.cirB[3].x(0)\n",
    "        self.cirB[3].z(0)\n",
    "\n",
    "    def teleportation(self):\n",
    "        status = self.init_status\n",
    "        # Run circuit\n",
    "        status = self.cirA(status)\n",
    "        # Alice measures both of her qubits C and A\n",
    "        status_A = self.measure(status, [(\"Alice\", 0), (\"Alice\", 1)], [\"00\", \"01\", \"10\", \"11\"])\n",
    "\n",
    "        # Record average fidelity \n",
    "        fid_list = []\n",
    "        # Bob applies different gates on his qubits depending on the measurement result of Alice\n",
    "        for i, s in enumerate(status_A):\n",
    "            # Bob's circuit\n",
    "            cirB = self.cirB[int(status_A[i].measured_result, 2)]\n",
    "            # Run circuit\n",
    "            status_B = cirB(s) \n",
    "            # Trace out the measured qubits C and A\n",
    "            # Leaving out only Bob’s qubit B\n",
    "            status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "            # Calculate the fidelity between the teleported state and the original state\n",
    "            fid = state_fidelity(self.state_C, status_fin) ** 2\n",
    "            fid_list.append(fid * status_fin.prob)\n",
    "        fid_avg = sum(fid_list)\n",
    "\n",
    "        return fid_avg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we randomly generate 200 pure quantum states and use state fidelity $F$ to benchmark the teleportation protocol, where\n",
    "\n",
    "$$\n",
    "F(\\rho,\\sigma) \\equiv \\text{tr}\\big( \\sqrt{\\sqrt{\\rho}\\sigma \\sqrt{\\rho}} \\big)^2.\n",
    "\\tag{5}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:29:03.318794Z",
     "start_time": "2021-03-09T06:28:54.202670Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Teleportation_Fidelity_Avg: 1.0 , std= 3.9433777e-07\n"
     ]
    }
   ],
   "source": [
    "SEED = 999              # Fix random seed\n",
    "num_state = 200         # Number of random states generated\n",
    "list_fid = []           # Record the fidelity\n",
    "np.random.seed(SEED)\n",
    "# Start sampling\n",
    "for idx in range(num_state):\n",
    "    list_fid.append(LOCC().teleportation().numpy())\n",
    "\n",
    "print('Teleportation_Fidelity_Avg:', np.around(sum(list_fid)[0] / len(list_fid), 4), ', std=', np.std(list_fid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** We want to point out the validity of this protocol relies on the quality of the pre-shared entanglement. Readers can change the entangled state from `bell_state(2)` to `isotropic_state(2, p)` and see how quantum noises will influence the teleportation performance. Recall the definition of isotropic states,\n",
    "\n",
    "\n",
    "$$\n",
    "\\rho_{\\text{iso}}(p) = p\\lvert\\Phi^+\\rangle \\langle\\Phi^+\\rvert + (1-p)\\frac{I}{4}, \\quad p \\in [0,1]\n",
    "\\tag{6}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning a teleportation protocol with LOCCNet\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training the protocol \n",
    "\n",
    "A general LOCC protocol can be classified by the number of classical communication rounds $r$. The original teleportation protocol is a one-round ($r=1$) protocol. For simplicity, we also restrict the communication rounds to 1. Compare to the original protocol, we use parametrized quantum circuits (PQC) to replace the fixed gates $U\\in\\{X,Z\\}$ Bob applied on his qubit with a general rotation gate $U_3$ on the Bloch sphere. \n",
    "\n",
    "$$\n",
    "U_3(\\theta, \\phi, \\varphi) =\n",
    "\\begin{bmatrix}\n",
    "\\cos(\\frac{\\theta}{2})           & -e^{i\\varphi}\\sin(\\frac{\\theta}{2})\\\\\n",
    "e^{i\\phi}\\sin(\\frac{\\theta}{2})  & e^{i(\\phi+\\varphi)} \\cos(\\frac{\\theta}{2})\n",
    "\\end{bmatrix}.\n",
    "\\tag{7}\n",
    "$$\n",
    "\n",
    "This would bring us a more powerful searching capability in finding practical LOCC protocols. Similarly, we change Alice's local operations to a more general PQC called the `universal_two_qubits()` [4]. We summarize the workflow below: \n",
    "\n",
    "1. Alice applies a 2-qubit PQC on her qubits.\n",
    "2. Alice measures both of her qubits in the computational basis and communicates with Bob through a classical channel.\n",
    "3. There are 4 possible measurement results: $m_1m_2 \\in \\{00, 01, 10, 11\\}$. Bob needs to act different operations corresponding to these measurement results. After Bob's local operations, the state on his qubit collapses to $\\lvert \\psi\\rangle_{B}$. \n",
    "4. Calculate the overlap $O$ between $\\lvert \\psi\\rangle_{B}$ and $\\lvert\\psi_C\\rangle$ (pure states). LOCCNet framework only supports density matrix formulation and hence we have to rewrite them as $\\rho_{B} = |\\psi\\rangle\\langle\\psi|_B$ and $\\rho_{C} = |\\psi\\rangle\\langle\\psi|_C$. Then, $O = \\text{Tr}(\\rho_C\\rho_{B})$. For pure states, this metric is simply the state fidelity.\n",
    "5. Set the accumulated loss function over 4 possible measurement results as $L = \\sum_{m_1m_2} \\big(1-\\text{Tr}(\\rho_C\\rho_{B})\\big)$, and use gradient-based optimization methods to update circuit parameters and hence minimize the loss function.\n",
    "6. Repeat steps 1-5 until the loss function converges to the global minimum.\n",
    "7. Generate an ensemble of arbitrary states $\\{\\lvert\\psi_C\\rangle\\}$, and benchmark the trained protocol with average state fidelity.\n",
    "\n",
    "![teleportation-LOCCNet](figures/teleportation-fig-LOCCNet.png \"Figure 2: Learning a teleportation protocol with LOCCNet.\")\n",
    "<div style=\"text-align:center\">Figure 2: Learning a teleportation protocol with LOCCNet. </div>\n",
    "\n",
    "**Note:** In order to make sure the parameters in parameterized quantum circuit is valid for all state after training, we set the training set as 4 linear independent states, which is $\\{|0\\rangle\\langle 0|,|1\\rangle\\langle 1|,|+\\rangle\\langle +|,|+\\rangle\\langle +|_y\\}$ or in the density matrix form:\n",
    "\n",
    "$$\n",
    "\\rho_0 = \\left[\\begin{array}{ccc}\n",
    "1 & 0\\\\\n",
    "0 & 0\n",
    "\\end{array}\\right], \n",
    "\\rho_1 = \\left[\\begin{array}{ccc}\n",
    "0 & 0\\\\\n",
    "0 & 1\n",
    "\\end{array}\\right], \n",
    "\\rho_2 = \\left[\\begin{array}{ccc}\n",
    "0.5 & 0.5\\\\\n",
    "0.5 & 0.5\n",
    "\\end{array}\\right], \n",
    "\\rho_3 = \\left[\\begin{array}{ccc}\n",
    "0.5 & -0.5 i\\\\\n",
    "0.5i & 0.5\n",
    "\\end{array}\\right]. \n",
    "\\tag{8}\n",
    "$$\n",
    "\n",
    "Any single qubit state can be written as a combination of the above 4 linear independent states in $\\mathcal{H}^{2\\times 2}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:29:08.878062Z",
     "start_time": "2021-03-09T06:29:08.796051Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC_Train(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC_Train, self).__init__()\n",
    "        # Add the first party Alice \n",
    "        # The first parameter 2 stands for how many qubits A holds\n",
    "        # The second parameter records the name of this party\n",
    "        self.add_new_party(2, party_name=\"Alice\")\n",
    "        # Add the second party Bob\n",
    "        # The first parameter 1 stands for how many qubits B holds\n",
    "        # The second parameter records the name of this party\n",
    "        self.add_new_party(1, party_name=\"Bob\")\n",
    "\n",
    "        # Create a bell state\n",
    "        _state = bell_state(2)\n",
    "        # _state = isotropic_state(2, 0.8)\n",
    "        # Training set: 4 linear independent states\n",
    "        _state0 = paddle_quantum.State(np.array([[1, 0], [0, 0]], dtype=np.complex64))\n",
    "        _state1 = paddle_quantum.State(np.array([[0, 0], [0, 1]], dtype=np.complex64))\n",
    "        _state2 = paddle_quantum.State(np.array([[0.5, 0.5], [0.5, 0.5]], dtype=np.complex64))\n",
    "        _state3 = paddle_quantum.State(np.array([[0.5, -0.5j], [0.5j, 0.5]], dtype=np.complex64))\n",
    "        self.init_states = [_state0, _state1, _state2, _state3]\n",
    "        # Initialize the system by distributing states between Alice and Bob\n",
    "        self.set_init_state(_state, [(\"Alice\", 1), (\"Bob\", 0)])\n",
    "        self.set_init_state(_state0, [(\"Alice\", 0)])\n",
    "\n",
    "        # Define Alice's local operations\n",
    "        self.cirA = self.create_ansatz(\"Alice\")\n",
    "        self.cirA.universal_two_qubits([0, 1])\n",
    "        # Define Bob's local operations\n",
    "        self.cirB = [self.create_ansatz(\"Bob\") for _ in range(4)]\n",
    "        for cir in self.cirB:\n",
    "            # Apply a universal single qubit gate\n",
    "            cir.u3(0)\n",
    "\n",
    "    def LOCCNet(self):\n",
    "        # Define the training process\n",
    "        loss = 0\n",
    "        temp_state = self.init_status\n",
    "        # Training\n",
    "        for init_state in self.init_states:\n",
    "            # Reset Alice's first qubit C to states in training set\n",
    "            status = self.reset_state(temp_state, init_state, [(\"Alice\", 0)])\n",
    "            # Run circuit\n",
    "            status = self.cirA(status)\n",
    "            # Obtain 4 possible measurement results\n",
    "            status_A = self.measure(status, [(\"Alice\", 0), (\"Alice\", 1)], [\"00\", \"01\", \"10\", \"11\"])\n",
    "\n",
    "            # Bob needs to apply different operation on his qubits, depending on the measurement results\n",
    "            for i, s in enumerate(status_A):\n",
    "                # Run circuit\n",
    "                status_B = self.cirB[i](s)\n",
    "                # Trace out the measured qubits C and A\n",
    "                # Leaving out only Bob’s qubit B\n",
    "                status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "                # Summing up the loss function for all possible measurement results\n",
    "                loss += 1 - paddle.real(trace(matmul(init_state.data, status_fin.data)))\n",
    "\n",
    "        return loss\n",
    "\n",
    "    # Save the optimized parameters\n",
    "    def save_module(self):\n",
    "        theta_A = self.cirA.parameters()\n",
    "        theta_B = [self.cirB[i].parameters() for i in range(4)]\n",
    "        theta = theta_A + sum(theta_B, [])\n",
    "        paddle.save(theta, 'parameters/QT_LOCCNet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:30:12.560553Z",
     "start_time": "2021-03-09T06:29:11.902232Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "itr 0: 7.721435\n",
      "itr 10: 0.5636051\n",
      "itr 20: 0.20643002\n",
      "itr 30: 0.073403895\n",
      "itr 40: 0.025113285\n",
      "itr 50: 0.008655369\n",
      "itr 60: 0.0035846233\n",
      "itr 70: 0.0014438629\n",
      "itr 80: 0.00045502186\n",
      "itr 90: 0.00018626451\n",
      "itr 100: 4.7028065e-05\n",
      "itr 110: 1.4960766e-05\n",
      "itr 120: 4.4703484e-06\n",
      "itr 130: 3.1590462e-06\n",
      "itr 140: 1.66893e-06\n"
     ]
    }
   ],
   "source": [
    "ITR = 150   # Number of iterations\n",
    "LR = 0.2    # Set up learning rate\n",
    "SEED = 999  # Fix random seed for parameters in PQC\n",
    "np.random.seed(SEED)\n",
    "paddle.seed(SEED)\n",
    "\n",
    "net = LOCC_Train()\n",
    "params = net.cirA.parameters() + sum([net.cirB[i].parameters() for i in range(4)], [])\n",
    "# Choose the Adam optimizer\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=params)\n",
    "# Optimization loop\n",
    "for itr in range(ITR):\n",
    "    # Forward propagation to calculate loss function\n",
    "    loss = net.LOCCNet()\n",
    "    # Backpropagation\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "    # Clean gradients\n",
    "    opt.clear_grad()\n",
    "    if itr % 10 == 0:\n",
    "        print(\"itr \" + str(itr) + \":\", loss.numpy()[0])\n",
    "\n",
    "# Save parameters\n",
    "net.save_module()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Benchmark\n",
    "\n",
    "If you don't want to spend time on training, you can **load the pre-trained circuit parameters** and directly test the performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-23T09:20:18.582863Z",
     "start_time": "2021-02-23T09:20:18.568348Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC_Test(LoccNet):\n",
    "    def __init__(self, theta_A, theta_B):\n",
    "        super(LOCC_Test, self).__init__()\n",
    "        self.parties = list()\n",
    "        self.add_new_party(2, party_name=\"Alice\")\n",
    "        self.add_new_party(1, party_name=\"Bob\")\n",
    "\n",
    "        _state = bell_state(2)\n",
    "        self._state0 = random_state(1)\n",
    "        self.set_init_state(_state, [(\"Alice\", 1), (\"Bob\", 0)])\n",
    "        self.set_init_state(self._state0, [(\"Alice\", 0)])\n",
    "\n",
    "        self.cirA = self.create_ansatz(\"Alice\")\n",
    "        self.cirA.universal_two_qubits(qubits_idx=[0, 1], param=theta_A)\n",
    "        self.cirB = [self.create_ansatz(\"Bob\") for _ in range(4)]\n",
    "        for i, cir in enumerate(self.cirB):\n",
    "            cir.u3(qubits_idx=0, param=theta_B[i])\n",
    "        \n",
    "\n",
    "    def benchmark(self):\n",
    "        input_state = self.init_status\n",
    "        status = self.cirA(input_state)\n",
    "        status_A = self.measure(status, [(\"Alice\", 0), (\"Alice\", 1)], [\"00\", \"01\", \"10\", \"11\"])\n",
    "        fid_list = []\n",
    "        for i, s in enumerate(status_A):\n",
    "            status_B = self.cirB[i](s)\n",
    "            status_fin = self.partial_state(status_B, [(\"Bob\", 0)])\n",
    "            fid = state_fidelity(self._state0, status_fin) ** 2\n",
    "            fid_list.append(fid * status_fin.prob)\n",
    "        fid_ave = sum(fid_list)\n",
    "        \n",
    "        return fid_ave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-02-23T09:20:33.607232Z",
     "start_time": "2021-02-23T09:20:22.038477Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LOCCNet_Fidelity_avg: 1.0 , std= 3.8644316e-07\n"
     ]
    }
   ],
   "source": [
    "# Load pre-trained circuit parameters\n",
    "para = paddle.load('parameters/QT_LOCCNet')\n",
    "np.random.seed(999)     # Fix random seed\n",
    "num_state = 200         # Number of random states generated\n",
    "list_fid = []           # Record the fidelity\n",
    "np.random.seed(SEED)\n",
    "paddle.seed(SEED)\n",
    "\n",
    "# Start sampling\n",
    "for idx in range(num_state):\n",
    "    list_fid.append(LOCC_Test(para[0], para[1:]).benchmark().numpy())\n",
    "\n",
    "print('LOCCNet_Fidelity_avg:', np.around(sum(list_fid)[0] / len(list_fid), 4), ', std=', np.std(list_fid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Based on LOCCNet we successfully learned a teleportation protocol with a noiseless pre-shared Bell state. The original teleportation protocol was designed to transfer a single-qubit quantum state. It is not clear how it can be generalize to the multi-qubit case. By comparison, LOCCNet provides the possibility of training a teleportation protocol for multi-qubit quantum states. Also, it will be an interesting question to ask how robust LOCCNet will  be against various noises. On the other hand, the teleportation protocol could be viewed as a special case of simulating the identity channel $\\mathcal{E}_I$ where Alice sends $\\psi$ and Bob receives $\\mathcal{E}_I(\\psi)$. This idea can be extended to simulate many other channels including the depolarizing channel $\\mathcal{E}_{D}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## References\n",
    "\n",
    "[1] Bennett, Charles H., et al. \"Teleporting an unknown quantum state via dual classical and Einstein-Podolsky-Rosen channels.\" [Physical Review Letters 70.13 (1993): 1895.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.70.1895)\n",
    "\n",
    "[2] Boschi, Danilo, et al. \"Experimental realization of teleporting an unknown pure quantum state via dual classical and Einstein-Podolsky-Rosen channels.\" [Physical Review Letters 80.6 (1998): 1121.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.80.1121)\n",
    "\n",
    "[3] Bouwmeester, Dik, et al. \"Experimental quantum teleportation.\" [Nature 390.6660 (1997): 575-579.](https://www.nature.com/articles/37539)\n",
    "\n",
    "[4] Vidal, Guifre, and Christopher M. Dawson. \"Universal quantum circuit for two-qubit transformations with three controlled-NOT gates.\" [Physical Review A 69.1 (2004): 010301.](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.69.010301)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}