example/qkd.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum Key Distribution\n",
    "\n",
    "Quantum key distribution is the process of distributing cryptographic keys between parties using quantum methods. Due to the unique properties of quantum information compared to classical, the security of a key can be guarunteed (as any unwelcomed measurement would change the state of quantum information transmitted).\n",
    "\n",
    "In this file, we see the use of SeQUeNCe to simulate quantum key distribution between two adjacent nodes. The first example demonstrates key distribution alone (using the BB84 protocol), while the second example demonstrates additional error correction with the cascade protocol. The network topology, including hardware components, is shown below:\n",
    "\n",
    "<img src=\"./notebook_images/QKD_topo.png\" width=\"500\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Only BB84\n",
    "\n",
    "### Import\n",
    "\n",
    "We must first import the necessary tools from SeQUeNCe to run our simulations.\n",
    "\n",
    "- `Timeline` is the main simulation tool, providing an interface for the discrete-event simulation kernel.\n",
    "- `QKDNode` provides a ready-to-use quantum node for quantum key distribution, including necessary hardware and protocol implementations.\n",
    "- `QuantumChannel` and `ClassicalChannel` are communication links between quantum nodes, providing models of optical fibers.\n",
    "- The `pair_bb84_protocols` function is used to explicitly pair 2 node instances for key distribution, and establishes one node as the sender \"Alice\" and one as the receiver \"Bob\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interact\n",
    "from matplotlib import pyplot as plt\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sequence.kernel.timeline import Timeline\n",
    "from sequence.topology.node import QKDNode\n",
    "from sequence.components.optical_channel import QuantumChannel, ClassicalChannel\n",
    "from sequence.qkd.BB84 import pair_bb84_protocols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Control and Collecting Metrics\n",
    "\n",
    "Several elements of SeQUeNCe automatically collect simple metrics. This includes the BB84 protocol implementation, which collects key error rates, throughput, and latency. For custom or more advanced metrics, custom code may need to be written and applied. See the documentation for a list of metrics provided by default for each simulation tool.\n",
    "\n",
    "Here, we create a `KeyManager` class to collect a custom metric (in this case, simply collect all of the generated keys and their generation time) and to provide an interface for the BB84 Protocol. To achieve this, we use the `push` and `pop` functions provided by the protocol stack on QKD nodes. `push` is used to send information down the stack (from the key manager to BB84 in this example) while `pop` is used to send information upwards (from BB84 to the key manager). Different protocols may use these interfaces for different data but only BB84 is shown in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KeyManager():\n",
    "    def __init__(self, timeline, keysize, num_keys):\n",
    "        self.timeline = timeline\n",
    "        self.lower_protocols = []\n",
    "        self.keysize = keysize\n",
    "        self.num_keys = num_keys\n",
    "        self.keys = []\n",
    "        self.times = []\n",
    "        \n",
    "    def send_request(self):\n",
    "        for p in self.lower_protocols:\n",
    "            p.push(self.keysize, self.num_keys) # interface for BB84 to generate key\n",
    "            \n",
    "    def pop(self, info): # interface for BB84 to return generated keys\n",
    "        self.keys.append(info)\n",
    "        self.times.append(self.timeline.now() * 1e-9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building the Simulation\n",
    "\n",
    "We are now ready to build the simulation itself. This example follows the usual process to ensure that all tools function properly:\n",
    "\n",
    "1. Create the timeline for the simulation\n",
    "2. Create the simulated network topology (here this is done explicitly, but this may also be handled by functions of the `Topology` class under `sequence.topology.topology`)\n",
    "3. Instantiate custom protocols and ensure all protocols are set up (paired) properly (if necessary)\n",
    "4. Initialize and run the simulation\n",
    "5. Collect and display the desired metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(sim_time, keysize):\n",
    "    \"\"\"\n",
    "    sim_time: duration of simulation time (ms)\n",
    "    keysize: size of generated secure key (bits)\n",
    "    \"\"\"\n",
    "    # begin by defining the simulation timeline with the correct simulation time\n",
    "    tl = Timeline(sim_time * 1e9)\n",
    "    \n",
    "    # Here, we create nodes for the network (QKD nodes for key distribution)\n",
    "    # stack_size=1 indicates that only the BB84 protocol should be included\n",
    "    n1 = QKDNode(\"n1\", tl, stack_size=1)\n",
    "    n2 = QKDNode(\"n2\", tl, stack_size=1)\n",
    "    n1.set_seed(0)\n",
    "    n2.set_seed(1)\n",
    "    pair_bb84_protocols(n1.protocol_stack[0], n2.protocol_stack[0])\n",
    "    \n",
    "    # connect the nodes and set parameters for the fibers\n",
    "    # note that channels are one-way\n",
    "    # construct a classical communication channel\n",
    "    # (with arguments for the channel name, timeline, and length (in m))\n",
    "    cc0 = ClassicalChannel(\"cc_n1_n2\", tl, distance=1e3)\n",
    "    cc1 = ClassicalChannel(\"cc_n2_n1\", tl, distance=1e3)\n",
    "    cc0.set_ends(n1, n2.name)\n",
    "    cc1.set_ends(n2, n1.name)\n",
    "    # construct a quantum communication channel\n",
    "    # (with arguments for the channel name, timeline, attenuation (in dB/km), and distance (in m))\n",
    "    qc0 = QuantumChannel(\"qc_n1_n2\", tl, attenuation=1e-5, distance=1e3,\n",
    "                         polarization_fidelity=0.97)\n",
    "    qc1 = QuantumChannel(\"qc_n2_n1\", tl, attenuation=1e-5, distance=1e3,\n",
    "                         polarization_fidelity=0.97)\n",
    "    qc0.set_ends(n1, n2.name)\n",
    "    qc1.set_ends(n2, n1.name)\n",
    "    \n",
    "    # instantiate our written keysize protocol\n",
    "    km1 = KeyManager(tl, keysize, 25)\n",
    "    km1.lower_protocols.append(n1.protocol_stack[0])\n",
    "    n1.protocol_stack[0].upper_protocols.append(km1)\n",
    "    km2 = KeyManager(tl, keysize, 25)\n",
    "    km2.lower_protocols.append(n2.protocol_stack[0])\n",
    "    n2.protocol_stack[0].upper_protocols.append(km2)\n",
    "    \n",
    "    # start simulation and record timing\n",
    "    tl.init()\n",
    "    km1.send_request()\n",
    "    tick = time.time()\n",
    "    tl.run()\n",
    "    print(\"execution time %.2f sec\" % (time.time() - tick))\n",
    "    \n",
    "    # display our collected metrics\n",
    "    plt.plot(km1.times, range(1, len(km1.keys) + 1), marker=\"o\")\n",
    "    plt.xlabel(\"Simulation time (ms)\")\n",
    "    plt.ylabel(\"Number of Completed Keys\")\n",
    "    plt.show()\n",
    "    \n",
    "    print(\"key error rates:\")\n",
    "    for i, e in enumerate(n1.protocol_stack[0].error_rates):\n",
    "        print(\"\\tkey {}:\\t{}%\".format(i + 1, e * 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the Simulation\n",
    "\n",
    "All that is left is to run the simulation with user input. (maximum execution time: ~5 sec)\n",
    "\n",
    "Parameters:\n",
    "\n",
    "    sim_time: duration of simulation time (ms)\n",
    "    keysize: size of generated secure key (bits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28d22d2059344083b9c46aef7f89a6fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=500, description='sim_time', max=1000, min=100, step=100), Dropdown(desc…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.test(sim_time, keysize)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create and run the simulation\n",
    "interactive_plot = interact(test, sim_time=(100, 1000, 100), keysize=[128, 256, 512])\n",
    "interactive_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to the imperfect polarization fidelity specified for the optical fiber, we observe that most (if not all) of the completed keys have errors that render them unusable. For this reason, error correction protocols (such as cascade, which is included in SeQUeNCe and shown in the next example) must also be used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Adding Cascade\n",
    "\n",
    "This example is simular to the first example, with slight alterations to allow for\n",
    "\n",
    "- Instatiation of the cascade error correction protocol on the qkd nodes\n",
    "- Differences in the cascade `push`/`pop` interface compared to BB84\n",
    "\n",
    "while the network topology remains unchanged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sequence.qkd.cascade import pair_cascade_protocols\n",
    "\n",
    "class KeyManager():\n",
    "    def __init__(self, timeline, keysize, num_keys):\n",
    "        self.timeline = timeline\n",
    "        self.lower_protocols = []\n",
    "        self.keysize = keysize\n",
    "        self.num_keys = num_keys\n",
    "        self.keys = []\n",
    "        self.times = []\n",
    "        \n",
    "    def send_request(self):\n",
    "        for p in self.lower_protocols:\n",
    "            p.push(self.keysize, self.num_keys) # interface for cascade to generate keys\n",
    "            \n",
    "    def pop(self, key): # interface for cascade to return generated keys\n",
    "        self.keys.append(key)\n",
    "        self.times.append(self.timeline.now() * 1e-9)\n",
    "        \n",
    "def test(sim_time, keysize):\n",
    "    \"\"\"\n",
    "    sim_time: duration of simulation time (ms)\n",
    "    keysize: size of generated secure key (bits)\n",
    "    \"\"\"\n",
    "    # begin by defining the simulation timeline with the correct simulation time\n",
    "    tl = Timeline(sim_time * 1e9)\n",
    "    \n",
    "    # Here, we create nodes for the network (QKD nodes for key distribution)\n",
    "    n1 = QKDNode(\"n1\", tl)\n",
    "    n2 = QKDNode(\"n2\", tl)\n",
    "    n1.set_seed(0)\n",
    "    n2.set_seed(1)\n",
    "    pair_bb84_protocols(n1.protocol_stack[0], n2.protocol_stack[0])\n",
    "    pair_cascade_protocols(n1.protocol_stack[1], n2.protocol_stack[1])\n",
    "    \n",
    "    # connect the nodes and set parameters for the fibers\n",
    "    cc0 = ClassicalChannel(\"cc_n1_n2\", tl, distance=1e3)\n",
    "    cc1 = ClassicalChannel(\"cc_n2_n1\", tl, distance=1e3)\n",
    "    cc0.set_ends(n1, n2.name)\n",
    "    cc1.set_ends(n2, n1.name)\n",
    "    qc0 = QuantumChannel(\"qc_n1_n2\", tl, attenuation=1e-5, distance=1e3,\n",
    "                         polarization_fidelity=0.97)\n",
    "    qc1 = QuantumChannel(\"qc_n2_n1\", tl, attenuation=1e-5, distance=1e3,\n",
    "                         polarization_fidelity=0.97)\n",
    "    qc0.set_ends(n1, n2.name)\n",
    "    qc1.set_ends(n2, n1.name)\n",
    "    \n",
    "    # instantiate our written keysize protocol\n",
    "    km1 = KeyManager(tl, keysize, 10)\n",
    "    km1.lower_protocols.append(n1.protocol_stack[1])\n",
    "    n1.protocol_stack[1].upper_protocols.append(km1)\n",
    "    km2 = KeyManager(tl, keysize, 10)\n",
    "    km2.lower_protocols.append(n2.protocol_stack[1])\n",
    "    n2.protocol_stack[1].upper_protocols.append(km2)\n",
    "    \n",
    "    # start simulation and record timing\n",
    "    tl.init()\n",
    "    km1.send_request()\n",
    "    tick = time.time()\n",
    "    tl.run()\n",
    "    print(\"execution time %.2f sec\" % (time.time() - tick))\n",
    "    \n",
    "    # display our collected metrics\n",
    "    plt.plot(km1.times, range(1, len(km1.keys) + 1), marker=\"o\")\n",
    "    plt.xlabel(\"Simulation time (ms)\")\n",
    "    plt.ylabel(\"Number of Completed Keys\")\n",
    "    plt.show()\n",
    "    \n",
    "    error_rates = []\n",
    "    for i, key in enumerate(km1.keys):\n",
    "        counter = 0\n",
    "        diff = key ^ km2.keys[i]\n",
    "        for j in range(km1.keysize):\n",
    "            counter += (diff >> j) & 1\n",
    "        error_rates.append(counter)\n",
    "\n",
    "    print(\"key error rates:\")\n",
    "    for i, e in enumerate(error_rates):\n",
    "        print(\"\\tkey {}:\\t{}%\".format(i + 1, e * 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the Simulation\n",
    "\n",
    "We can now run the cascade simulation with user input. Note that the extra steps required by the cascade protocol may cause the simulation to run much longer than the example with only BB84.\n",
    "\n",
    "Parameters:\n",
    "\n",
    "    sim_time: duration of simulation time (ms)\n",
    "    keysize: size of generated secure key (bits)\n",
    "    \n",
    "The maximum execution time (`sim_time=1000`, `keysize=512`) is around 60 seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0703f7c810843e7b4e4c3d286b20381",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=500, description='sim_time', max=1000, min=100, step=100), Dropdown(desc…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.test(sim_time, keysize)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create and run the simulation\n",
    "interactive_plot = interact(test, sim_time=(100, 1000, 100), keysize=[128, 256, 512])\n",
    "interactive_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "\n",
    "The implementation of the cascade protocol found within SeQUeNCe relies on the creation of a large sequence of bits, from which exerpts are used to create individual keys. Due to this behavior, keys are generated in large numbers in regularly spaced \"batches\". Also note that after application of error correction, the error rates for all keys are now at 0%."
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}