demos/quantum-error-correction-shor-code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SQUANCH demonstration: quantum error correction with the Shor code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from squanch import *\n",
    "from scipy.stats import unitary_group\n",
    "import copy\n",
    "import numpy as np\n",
    "import matplotlib.image as image\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\newcommand{\\ket}[1]{\\lvert #1 \\rangle}$\n",
    "\n",
    "When qubits are transmitted over quantum channels, they are subject to a complex set of errors which can cause them to decohere, depolarize, or simply vanish completely. For quantum information transfer to be feasible, the information must be encoded in a error-resistant format using any of a variety of quantum error correction models. In this demonstration, we show how to use $\\texttt{SQUANCH}$’s channel and error modules to simulate quantum errors in a transmitted message, which we correct for using the [Shor Code](https://en.wikipedia.org/wiki/Quantum_error_correction#The_Shor_code). This error correction model encodes a single logical qubit into the product of 9 physical qubits and is capable of correcting for arbitrary single-qubit errors. A circuit diagram for this protocol is shown below, where $E$ represents a quantum channel which can arbitrarily corrupt a single qubit. \n",
    "\n",
    "![](../docs/source/img/shor-code-circuit.png)\n",
    "\n",
    "In this demo, we have two pairs of agents: Alice and Bob will communicate a message which is error-protected using the Shor code, and DumbAlice an DumbBob will transmit the message without error correction. Formally, for each state $\\ket{\\psi}$ to be transmitted through the channel, the following procedure is simulated:\n",
    "\n",
    "1. Alice has some state $\\ket{\\psi}=\\alpha_0\\ket{0}+\\alpha_1\\ket{1}$, which she wants to send to Bob through a noisy quantum channel. She encodes her single-qubit state in nine logical qubit as $\\ket{\\psi} \\mapsto \\alpha_0\\bigotimes_{j=1}^3\\frac{1}{\\sqrt{2}}\\left(\\ket{000}+\\ket{111}\\right) + \\alpha_1\\bigotimes_{k=1}^3\\frac{1}{\\sqrt{2}}\\left(\\ket{000}-\\ket{111}\\right)$ using the circuit diagram above.\n",
    "2. DumbAlice wants to send the same state, but she doesn't error-protect the state and transmits the unencoded state $\\ket{\\psi}\\otimes{\\ket{00\\cdots0}}$.\n",
    "3. Alice and DumbAlice send their qubits through the quantum channel $E$ to Bob and DumbBob, respectively. The channel may apply an arbitrary unitary operation to a single physical qubit in each group of nine.\n",
    "4. Bob receives Alice's qubits and decodes them using the Shor decoding circuit shown above. \n",
    "5. DumbBob expects $\\ket{\\psi}\\otimes{\\ket{00\\cdots0}}$ from DumbAlice and only measures the results of the the first qubit in each group of nine.\n",
    "\n",
    "Transmitting an image is unsuitable for this scenario due to the larger size of the Hilbert space involved compared to the previous two demonstrations. (Each $\\texttt{QSystem.state}$ for $N=9$ uses 2097264 bytes, compared to 240 bytes for $N=2$.) Instead, Alice and DumbAlice will transmit the bitwise representation of a short message encoded as $\\sigma_z$-eigenstates, and Bob and DumbBob will attempt to re-assemble the message. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantum error correction logic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shor_encode(qsys):\n",
    "    # psi is state to send, q1...q8 are ancillas from top to bottom in diagram\n",
    "    psi, q1, q2, q3, q4, q5, q6, q7, q8 = qsys.qubits\n",
    "    # Gates are enumerated left to right, top to bottom from figure\n",
    "    CNOT(psi, q3)\n",
    "    CNOT(psi, q6)\n",
    "    H(psi)\n",
    "    H(q3)\n",
    "    H(q6)\n",
    "    CNOT(psi, q1)\n",
    "    CNOT(psi, q2) \n",
    "    CNOT(q3, q4)\n",
    "    CNOT(q3, q5)\n",
    "    CNOT(q6, q7)\n",
    "    CNOT(q6, q8)\n",
    "    return psi, q1, q2, q3, q4, q5, q6, q7, q8\n",
    "\n",
    "def shor_decode(psi, q1, q2, q3, q4, q5, q6, q7, q8):\n",
    "    # same enumeration as Alice\n",
    "    CNOT(psi, q1)\n",
    "    CNOT(psi, q2)\n",
    "    TOFFOLI(q2, q1, psi)\n",
    "    CNOT(q3, q4)\n",
    "    CNOT(q3, q5)\n",
    "    TOFFOLI(q5, q4, q3)\n",
    "    CNOT(q6, q7)\n",
    "    CNOT(q6, q8)\n",
    "    TOFFOLI(q7, q8, q6) # Toffoli control qubit order doesn't matter\n",
    "    H(psi)\n",
    "    H(q3)\n",
    "    H(q6)\n",
    "    CNOT(psi, q3)\n",
    "    CNOT(psi, q6)\n",
    "    TOFFOLI(q6, q3, psi)\n",
    "    return psi # psi is now Alice's original state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent logic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Alice(Agent):\n",
    "    '''Alice sends an error-protected state to Bob'''\n",
    "    def run(self):\n",
    "        for qsys in self.qstream:\n",
    "            # send the encoded qubits to Bob \n",
    "            for qubit in shor_encode(qsys):\n",
    "                self.qsend(bob, qubit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DumbAlice(Agent):\n",
    "    '''DumbAlice sends an uncorrected state to DumbBob'''   \n",
    "    def run(self):\n",
    "        for qsys in self.qstream:\n",
    "            for qubit in qsys.qubits:\n",
    "                self.qsend(dumb_bob, qubit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bob(Agent):\n",
    "    '''Bob receives and error-corrects Alice's state'''\n",
    "    def run(self):\n",
    "        measurement_results = []\n",
    "        for _ in self.qstream:\n",
    "            # Bob receives 9 qubits representing Alice's encoded state\n",
    "            received = [self.qrecv(alice) for _ in range(9)]\n",
    "            # Decode and measure the original state\n",
    "            psi_true = shor_decode(*received)\n",
    "            measurement_results.append(psi_true.measure())\n",
    "        self.output(measurement_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DumbBob(Agent):\n",
    "    '''DumbBob gets DumbAlice's non-corrected state'''\n",
    "    def run(self):\n",
    "        measurement_results = []\n",
    "        for _ in self.qstream:\n",
    "            received = [self.qrecv(dumb_alice) for _ in range(9)]\n",
    "            psi_true = received[0]\n",
    "            measurement_results.append(psi_true.measure())\n",
    "        self.output(measurement_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantum error model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ShorError(QError):\n",
    "\n",
    "    def __init__(self, qchannel):\n",
    "        '''\n",
    "        Instatiate the error model from the parent class\n",
    "        :param QChannel qchannel: parent quantum channel\n",
    "        '''\n",
    "        QError.__init__(self, qchannel)\n",
    "        self.count = 0 \n",
    "        self.error_applied = False\n",
    "\n",
    "    def apply(self, qubit):\n",
    "        '''\n",
    "        Apply a random unitary operation to one of the qubits in a set of 9\n",
    "        :param Qubit qubit: qubit from quantum channel\n",
    "        :return: either unchanged qubit or None\n",
    "        '''\n",
    "        # reset error for each group of 9 qubits\n",
    "        if self.count == 0:\n",
    "            self.error_applied = False\n",
    "        self.count = (self.count + 1) % 9\n",
    "        # qubit could be None if combining with other error models, such as attenuation\n",
    "        if not self.error_applied and qubit is not None:\n",
    "            if np.random.rand() < 0.5: # apply the error\n",
    "                random_unitary = unitary_group.rvs(2) # pick a random U(2) matrix\n",
    "                qubit.apply(random_unitary)\n",
    "                self.error_applied = True\n",
    "        return qubit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Channel model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ShorQChannel(QChannel):\n",
    "    '''Represents a quantum channel with a Shor error applied'''\n",
    "    \n",
    "    def __init__(self, from_agent, to_agent):\n",
    "        QChannel.__init__(self, from_agent, to_agent)\n",
    "        # register the error model\n",
    "        self.errors = [ShorError(self)] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_bits(string):\n",
    "    '''Convert a string to a list of bits'''\n",
    "    result = []\n",
    "    for c in string:\n",
    "        bits = bin(ord(c))[2:]\n",
    "        bits = '00000000'[len(bits):] + bits\n",
    "        result.extend([int(b) for b in bits])\n",
    "    return result\n",
    "\n",
    "def from_bits(bits):\n",
    "    '''Convert a list of bits to a string'''\n",
    "    chars = []\n",
    "    for b in range(int(len(bits) / 8)):\n",
    "        byte = bits[b*8:(b+1)*8]\n",
    "        chars.append(chr(int(''.join([str(bit) for bit in byte]), 2)))\n",
    "    return ''.join(chars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3160f84247f1478fa529232b7d255ce9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='DumbAlice', max=456), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07ce99be55be46a7a78c061f5efcc337",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='DumbBob', max=456), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b12bcff462b4eecb4eeb5bcf00c2d7d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Alice', max=456), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ee2e71e8a594910bfd8fdd260dedb35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='Bob', max=456), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "\n",
      "DumbAlice sent:   Peter Shor once lived in Ruddock 238! But who was Airman?\n",
      "DumbBob received: øu<ar®W/>P0mnª3àyx5!ig¾Â&ÄpÛ\u001bï(r38e0V~PviK\\dab¨`á|mçj?\n",
      "Alice sent:       Peter Shor once lived in Ruddock 238! But who was Airman?\n",
      "Bob received:     Peter Shor once lived in Ruddock 238! But who was Airman?\n"
     ]
    }
   ],
   "source": [
    "# Prepare a message to send\n",
    "msg = \"Peter Shor once lived in Ruddock 238! But who was Airman?\"\n",
    "bits = to_bits(msg)\n",
    "\n",
    "# Encode the message as spin eigenstates\n",
    "qstream = QStream(9, len(bits)) # 9 qubits per encoded state\n",
    "for bit, qsystem in zip(bits, qstream):\n",
    "    if bit == 1: \n",
    "        X(qsystem.qubit(0)) \n",
    "\n",
    "# Alice and Bob will use error correction\n",
    "out = Agent.shared_output()\n",
    "alice = Alice(qstream, out)\n",
    "bob = Bob(qstream, out)\n",
    "alice.qconnect(bob, ShorQChannel)\n",
    "\n",
    "# Dumb agents won't use error correction\n",
    "qstream2 = copy.deepcopy(qstream)\n",
    "dumb_alice = DumbAlice(qstream2, out)\n",
    "dumb_bob = DumbBob(qstream2, out)\n",
    "dumb_alice.qconnect(dumb_bob, ShorQChannel)\n",
    "\n",
    "# Run everything and record results\n",
    "Simulation(dumb_alice, dumb_bob, alice, bob).run()\n",
    "print(\"DumbAlice sent:   {}\".format(msg))\n",
    "print(\"DumbBob received: {}\".format(from_bits(out[\"DumbBob\"])))\n",
    "print(\"Alice sent:       {}\".format(msg))\n",
    "print(\"Bob received:     {}\".format(from_bits(out[\"Bob\"])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}