EntanglementDistillation_LOCCNET_EN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Entanglement Distillation -- Protocol Design with LOCCNet\n",
    "\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 entanglement plays a vital role in quantum communication, quantum computing, and many other quantum technologies. Therefore, detecting, transmitting, and distributing quantum entanglement reliably are essential tasks if we want to build real applications in those fields. However, errors are inevitable in the real world. They could come from imperfect equipment when we create entanglement (preparation errors), or the quantum channel used to transmit entanglement is noisy, and we gradually lose the degree of entanglement as the transmission distance increases. The aim of entanglement distillation is to compensate for those losses and restore a **maximally entangled state** at the cost of many noisy entangled states. In this sense, one could also refer entanglement distillation as a purification/error-correction protocol. This process often involves two remote parties Alice and Bob such that only Local Operations and Classical Communication (LOCC) are allowed [1]. Many distillation protocols have been proposed since 1996, including the famous BBPSSW [2] and the DEJMPS protocol [3]. \n",
    "\n",
    "However, the BBPSSW and DEJMPS distillation protocols are designed for specific types of noisy states (i.e., isotropic states and Bell-diagonal states, respectively). It is nearly impossible to find a single distillation protocol that could purify all kinds of noises. Due to the complicated mathematical structure of LOCC, designing a new distillation protocol is time-consuming with paper and pencil only. LOCCNet, a machine learning framework for LOCC protocols, is designed to reduce the effort as much as possible. With LOCCNet, it will be easier to design new distillation protocols as long as we can characterize the mathematical form of noise introduced to the entanglement resources. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preliminary\n",
    "\n",
    "In the context of entanglement distillation, we usually use the **state fidelity** $F$ between the distilled state $\\rho_{out}$ and the maximally entangled Bell state $|\\Phi^+\\rangle$ to quantify the performance of a distillation protocol, where\n",
    "\n",
    "$$\n",
    "F(\\rho_{out}, \\Phi^+) \\equiv \\langle \\Phi^+|\\rho_{out}|\\Phi^+\\rangle.\n",
    "\\tag{1}\n",
    "$$\n",
    "\n",
    "**Note:** In general, LOCC distillation protocols are probabilistic, and hence we are also interested in the success rate $p_{succ}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Protocol design logic\n",
    "\n",
    "In this tutorial, we will go through an example that distills four identical **isotropic states** (or Werner state) $\\rho_{in}= \\rho_{\\text{iso}}$ into a single final state $\\rho_{out}$ with a higher state fidelity $F$ (closer to the Bell state $|\\Phi^+\\rangle$). We call this type of protocol as the $4\\rightarrow 1$ LOCC distillation class, while the original BBPSSW and DEJMPS protocols belong to the $2\\rightarrow 1$ LOCC distillation class. The isotropic state is a parametrized family of mixed states consist of $|\\Phi^+\\rangle$ and the completely mixed state (white noise) $I/4$,\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{2}\n",
    "$$\n",
    "\n",
    "In our example, we set $p=0.7$ and the input state becomes:\n",
    "\n",
    "$$\n",
    "\\rho_{in} = \\rho_{\\text{iso}}(0.7)= 0.7\\lvert\\Phi^+\\rangle \\langle\\Phi^+\\rvert + 0.075 I.\n",
    "\\tag{3}\n",
    "$$\n",
    "\n",
    "\n",
    "To fulfill the task of distillation through LOCC, we introduce two remote parties, $A$ (Alice) and $B$ (Bob). At the very beginning, they share four copies of entangled qubit pairs $\\rho_{A_0B_0}, \\rho_{A_1B_1}, \\rho_{A_2B_2}$ and $\\rho_{A_3B_3}$. Each copy is initialized as $\\rho_{in} = \\rho_{\\text{iso}}(p =0.7)$.  Alice holds four qubits $A_0, A_1, A_2, A_3$ in one place and Bob holds $B_0, B_1, B_2, B_3$ in another place. With these initial setups, Alice and Bob could choose the communication rounds $r$, which indicates how many times they would measure their subsystems and transmit classical data with each other. For simplicity, it is natural to start with $r=1$. Then, Alice and Bob can use LOCCNet to find out the correct local operations (encoded as quantum neural networks, QNNs) before communication by following the steps below:\n",
    "\n",
    "1. Design a general QNN architecture $U(\\boldsymbol\\theta)$ as shown in Figure 1, where each $R(\\theta)$ represents a general rotation operation on the Bloch sphere. They are referred as the `u3(theta, phi, lam, which_qubit)` gate in Paddle Quantum.\n",
    "2. After implementing the QNN, Alice and Bob measure all of the qubits except $A_0$ and $B_0$ in the computational basis. Then the measurement results $M = \\{m_{A_1}m_{B_1}, m_{A_2}m_{B_2}, m_{A_3}m_{B_3}\\}$ will be exchanged through a classical channel.\n",
    "3. If the measurement results of all qubit pairs are the same (each pair $m_{A_1}m_{B_1}, m_{A_2}m_{B_2}, m_{A_3}m_{B_3}$ is either 00 or 11), the distillation is successful, and the pair $A_0B_0$ is reserved as the memory qubit pair to store the purified entanglement; If the measurement results are different, the protocol fails. Discard the quantum pair $A_0B_0$. At this point, we obtain the purified quantum pair $A_0B_0$ probabilistically, and its state is denoted as $\\rho_{AB}'$.\n",
    "4. Define the accumulated loss function $L = \\sum_{m_{A_j}m_{B_j}\\in \\{00,11\\}} \\big(1- \\text{Tr}(\\rho_{tar}\\rho_{AB}')\\big)$ over the successful distillations, where $\\text{Tr}(\\rho_{tar}\\rho_{AB}')$ is the state overlap between the current state $\\rho_{AB}'$ and the target state $\\rho_{tar}=\\lvert\\Phi^+\\rangle \\langle \\Phi^+\\rvert$.\n",
    "5. Use gradient-based optimization methods to update parameters in QNN and minimize the loss function.\n",
    "6. Repeat steps 1-5 until the loss function converges.\n",
    "7. Output the purified state $\\rho_{out} = \\rho'_{A_0B_0}$.\n",
    "\n",
    "**Note:** The QNN structure used is merely an illustration. Readers are welcomed to try their own designs.\n",
    "\n",
    "\n",
    "<center><img src=\"figures/distillation-fig-LOCCNet4.png\" height=\"200\" width=\"400\"></center>\n",
    "<div style=\"text-align:center\">Figure 1: Schematic diagram of a distillation protocol designed with LOCCNet. </div>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation with Paddle Quantum\n",
    "\n",
    "First, we import all the dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:19:13.646334Z",
     "start_time": "2021-03-09T06:19:10.780432Z"
    }
   },
   "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.state import isotropic_state, bell_state\n",
    "from paddle_quantum.qinfo import logarithmic_negativity\n",
    "# Change to density matrix mode\n",
    "paddle_quantum.set_backend('density_matrix')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Secondly, we define the QNN and the loss function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:19:13.700268Z",
     "start_time": "2021-03-09T06:19:13.650127Z"
    }
   },
   "outputs": [],
   "source": [
    "class LOCC(LoccNet):\n",
    "    def __init__(self):\n",
    "        super(LOCC, self).__init__()\n",
    "\n",
    "        # Add the first party Alice \n",
    "        # The first parameter 4 stands for how many qubits A holds\n",
    "        # The second parameter records the name of this party\n",
    "        self.add_new_party(4, party_name=\"Alice\")\n",
    "        # Add the second party Bob\n",
    "        # The first parameter 4 stands for how many qubits B holds\n",
    "        # The second parameter records the name of this party\n",
    "        self.add_new_party(4, party_name=\"Bob\")\n",
    "\n",
    "        # Generate input isotropic states\n",
    "        _state = isotropic_state(2, 0.7)\n",
    "        # Distribute the pre-shared entangled states\n",
    "        self.set_init_state(_state, [[\"Alice\", 0], [\"Bob\", 0]])\n",
    "        self.set_init_state(_state, [[\"Alice\", 1], [\"Bob\", 1]])\n",
    "        self.set_init_state(_state, [[\"Alice\", 2], [\"Bob\", 2]])\n",
    "        self.set_init_state(_state, [[\"Alice\", 3], [\"Bob\", 3]])\n",
    "\n",
    "        # Create Alice's circuit\n",
    "        self.cir1 = self.create_ansatz(\"Alice\")\n",
    "        self.QNN(self.cir1)\n",
    "        # Create Bob's circuit\n",
    "        self.cir2 = self.create_ansatz(\"Bob\")\n",
    "        self.QNN(self.cir2)\n",
    "\n",
    "    def QNN(self, cir):\n",
    "        '''\n",
    "        Define the QNN illustrated in Figure 1\n",
    "        '''\n",
    "        cir.u3('full')\n",
    "        cir.cnot('cycle')\n",
    "        cir.u3('full')\n",
    "\n",
    "    def New_Protocol(self):\n",
    "        status = self.init_status\n",
    "        # Execute Alice's circuit\n",
    "        status = self.cir1(status)\n",
    "        # Execute Bob's circuit\n",
    "        status = self.cir2(status)\n",
    "\n",
    "        # Measure qubits，[\"000000\", \"000011\",\"001100\",\"110000\",\"001111\",\"111100\",\"110011\",\"111111\"] represent successful cases\n",
    "        status1 = self.measure(status, [[\"Alice\", 1], [\"Bob\", 1],[\"Alice\", 2], [\"Bob\", 2], [\"Alice\", 3], [\"Bob\", 3]], \n",
    "                               [\"000000\", \"000011\", \"001100\", \"110000\", \"001111\", \"111100\", \"110011\", \"111111\"])\n",
    "\n",
    "        # Trace out all the qubits but A_0 and B_0\n",
    "        status_fin = self.partial_state(status1, [[\"Alice\", 0], [\"Bob\", 0]])\n",
    "        target_state = bell_state(2)\n",
    "\n",
    "        # Calculate loss function\n",
    "        loss = 0\n",
    "        for idx in range(0, len(status_fin)):\n",
    "            loss += 1 - paddle.real(trace(matmul(target_state.data, status_fin[idx].data)))\n",
    "\n",
    "        return loss, status_fin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, minimize the loss function with gradient-based optimization methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-09T06:22:27.071785Z",
     "start_time": "2021-03-09T06:19:13.706482Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "itr 0: 5.7033257\n",
      "itr 10: 1.552211\n",
      "itr 20: 0.8980296\n",
      "itr 30: 0.7061223\n",
      "itr 40: 0.57139534\n",
      "itr 50: 0.5278059\n",
      "itr 60: 0.5121539\n",
      "itr 70: 0.50759834\n",
      "itr 80: 0.5052042\n",
      "itr 90: 0.50493234\n",
      "The fidelity of the input quantum state is：0.77500\n",
      "The fidelity of the purified quantum state is： 0.93690\n",
      "The probability of successful purification is：0.38654\n",
      "========================================================\n",
      "The output state is:\n",
      " [[ 4.790e-01-0.e+00j -1.000e-04-0.e+00j  1.000e-04-0.e+00j\n",
      "   4.579e-01+5.e-04j]\n",
      " [-1.000e-04+0.e+00j  2.100e-02+0.e+00j  0.000e+00+0.e+00j\n",
      "  -1.000e-04+0.e+00j]\n",
      " [ 1.000e-04+0.e+00j  0.000e+00-0.e+00j  2.100e-02+0.e+00j\n",
      "   1.000e-04+1.e-04j]\n",
      " [ 4.579e-01-5.e-04j -1.000e-04-0.e+00j  1.000e-04-1.e-04j\n",
      "   4.790e-01-0.e+00j]]\n",
      "The initial logarithmic negativity is: 0.6322681307792664\n",
      "The final logarithmic negativity is: 0.9059638381004333\n"
     ]
    }
   ],
   "source": [
    "ITR = 100    # Number of iterations\n",
    "LR = 0.2     # Learning rate\n",
    "paddle.seed(999)\n",
    "\n",
    "net = LOCC()\n",
    "# Choose Adam optimizer\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=net.cir1.parameters() + net.cir2.parameters())\n",
    "\n",
    "# Optimization loop\n",
    "for itr in range(ITR):\n",
    "    loss, status_fin = net.New_Protocol()\n",
    "    # Backpropagation\n",
    "    loss.backward()          \n",
    "    opt.minimize(loss)\n",
    "    # Clean gradients\n",
    "    opt.clear_grad()\n",
    "    # Print training result\n",
    "    if itr % 10 == 0:\n",
    "        print(\"itr \" + str(itr) + \":\", loss.numpy()[0])\n",
    "\n",
    "# Calculate input state fidelity\n",
    "fidelity_in = (3 * 0.7 + 1) / 4\n",
    "# Calculate output state fidelity\n",
    "fidelity = (len(status_fin) - loss) / len(status_fin)\n",
    "# Calculate successful rate\n",
    "suc_rate = sum([s.prob for s in status_fin])\n",
    "\n",
    "print(\"The fidelity of the input quantum state is：%.5f\" % fidelity_in)\n",
    "print(\"The fidelity of the purified quantum state is： %.5f\" % fidelity.numpy()[0])\n",
    "print(\"The probability of successful purification is：%.5f\" % suc_rate.numpy()[0])\n",
    "rho_out = status_fin[0]\n",
    "print(\"========================================================\")\n",
    "print(f\"The output state is:\\n {np.around(rho_out.data.numpy(), 4)}\")\n",
    "print(f\"The initial logarithmic negativity is: {logarithmic_negativity(isotropic_state(2, 0.7)).numpy()[0]}\")\n",
    "print(f\"The final logarithmic negativity is: {logarithmic_negativity(rho_out).numpy()[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "\n",
    "As we can see, this new distillation protocol can purify four copies of isotropic states, each with a fidelity of 0.775, into a single two-qubit state with a fidelity of 0.937, which outperforms the extended DEJMPS protocol [3] with a distillation fidelity of 0.924 under the same setting. At the same time, our framework also exhibits advantages in terms of flexibility and scalability. With the help of LOCCNet, one can try various combinations and see what's the effect of increasing the communication rounds $r$, adding non-identical noisy entanglement before distillation, and importing different noise types.\n",
    "\n",
    "LOCCNet has a wide range of applications, and distillation is merely one of them. We want to point out that the protocols trained by LOCCNet are hardware efficient since every operation is applicable in near term quantum devices, and we offer the possibility of customizing QNN architectures. Now, it's time to design and train your own LOCC protocol! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## References\n",
    "\n",
    "[1] Chitambar, Eric, et al. \"Everything you always wanted to know about LOCC (but were afraid to ask).\" [Communications in Mathematical Physics 328.1 (2014): 303-326.](https://link.springer.com/article/10.1007/s00220-014-1953-9)\n",
    "\n",
    "[2] Bennett, Charles H., et al. \"Purification of noisy entanglement and faithful teleportation via noisy channels.\" [Physical Review Letters 76.5 (1996): 722.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.76.722)\n",
    "\n",
    "[3] Deutsch, David, et al. \"Quantum privacy amplification and the security of quantum cryptography over noisy channels.\" [Physical Review Letters 77.13 (1996): 2818.](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.77.2818)"
   ]
  }
 ],
 "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": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}