quantum-key-distro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1>Quantum Key Distribution (BB84) Visualisation and Basic Simulation</h1>\n",
    "<hr>\n",
    "<i>Written by Jamie</i><br>\n",
    "<p style=\"margin-bottom:1.5cm;\">This script is to visualise and simulate, on a basic level, how the quantum key distribution protocol <a href=\"https://arxiv.org/abs/2003.06557\">BB84</a> and error correction protocol <a href=\"https://link.springer.com/chapter/10.1007/3-540-48285-7_35\">cascade</a> are used to create, distribute and error correct a private encryption key.</p>\n",
    "\n",
    "<h3>Introduction</h3>\n",
    "<hr>\n",
    "<p style=\"margin-bottom:1.5cm;\">The basic premise is that two people, conventionally called Alice and Bob, want to communicate a secret securely. To do this they use the BB84 protocol to generate a private encryption key that can be used to encrypt their non-quantum secret. More information can be found in <a href=\"https://arxiv.org/abs/2003.06557\">this research paper</a>.</p>\n",
    "\n",
    "<h3>Import packages</h3>\n",
    "<hr>\n",
    "<p>Import python packages: <a href=\"https://pypi.org/project/random2/\">random</a> and the <a href=\"https://pypi.org/project/tabulate/\">tabulate</a> module from the library that shares its name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from tabulate import tabulate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Properties of Transaction</h3>\n",
    "<hr> \n",
    "<p>\n",
    "    Prompt the user for numerical inputs for:\n",
    "    <ul>\n",
    "        <li>The probability of a photon being lost in transmission.</li>\n",
    "        <li>The length in bits of the random string being generated and sent by the sender \"Alice\".</li>\n",
    "        <li>Probability of an incorrect bit being received by the receiver \"Bob\" due to noise in the system, also known as the error rate.</li>\n",
    "    </ul>\n",
    "    And run input validation on these user inputs to ensure that they are in the correct format so the program doesn't break later on. If the input validation fails, ask the user for the input again, providing the reason why the input was not sufficient.\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Chance of photon loss on transmission: 5%\n",
      "Number of random bits being sent: 16\n",
      "Probability of an incorrect bit being received due to noise (error rate): 10%\n"
     ]
    }
   ],
   "source": [
    "## Probability of photon being lost in transmission\n",
    "lose_in_tx_chance=input(\"Probability of losing a photon in transmission [%]: \")\n",
    "valid=False\n",
    "while valid == False: # Input validation\n",
    "    try:\n",
    "        lose_in_tx_chance=int(lose_in_tx_chance)\n",
    "    except:\n",
    "        lose_in_tx_chance=input(\"-----\\nError: Value must be a number\\nProbability of losing a photon in transmission [%]: \")\n",
    "    else:\n",
    "        if lose_in_tx_chance > 100 or lose_in_tx_chance < 0:\n",
    "            lose_in_tx_chance=input(\"-----\\nError: Value must be between 0 and 100\\nProbability of losing a photon in transmission [%]: \")\n",
    "        else:\n",
    "            valid=True\n",
    "print (\"Chance of photon loss on transmission: \"+str(lose_in_tx_chance)+\"%\") # Print value for debugging or later reference\n",
    "\n",
    "## Number of random bits being sent\n",
    "no_bits=input(\"Number of random bits being sent: \")\n",
    "valid=False\n",
    "while valid == False: # Input validation\n",
    "    try:\n",
    "        no_bits=int(no_bits)\n",
    "    except:\n",
    "        no_bits=input(\"-----\\nError: Value must be a number\\nNumber of random bits being sent: \")\n",
    "    else:\n",
    "        if no_bits <= 0:\n",
    "            no_bits=input(\"-----\\nError: Value must be above 0\\nNumber of random bits being sent: \")\n",
    "        else:\n",
    "            valid=True\n",
    "print (\"Number of random bits being sent: \"+str(no_bits)) # Print value for debugging or later reference\n",
    "\n",
    "## Probability of an incorrect bit being received due to noise (error rate)\n",
    "error_rate=input(\"Probability of an incorrect bit being received due to noise (error rate) [%]: \")\n",
    "valid=False\n",
    "while valid == False: # Input validation\n",
    "    try:\n",
    "        error_rate=int(error_rate)\n",
    "    except:\n",
    "        error_rate=input(\"-----\\nError: Value must be a number\\nProbability of an incorrect bit being received due to noise (error rate) [%]: \")\n",
    "    else:\n",
    "        if error_rate > 100 or error_rate < 0:\n",
    "            error_rate=input(\"-----\\nError: Value must be between 0 and 100\\nProbability of an incorrect bit being received due to noise (error rate) [%]: \")\n",
    "        else:\n",
    "            valid=True\n",
    "print (\"Probability of an incorrect bit being received due to noise (error rate): \"+str(error_rate)+\"%\") # Print value for debugging or later reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Define functions</h3>\n",
    "<hr>\n",
    "<p>To reduce code sprawl and to make life easier for programming, a number of functions are defined that can be recalled by the program later.<br>\n",
    "These include:\n",
    "<ul>\n",
    "    <li>A function to convert strings to arrays which is used to convert the randomly generated strings of 0s and 1s or Ds and Rs into arrays that are easy to preform calculations on.</li>\n",
    "    <li>A function to convert arrays to strings which is helpful in debugging/user output to print out easier-to-read strings of values in arrays, without the square brackets and commas.</li>\n",
    "    <li>A function to generate a random array of 1s and 0s that is as long as the integer inputted into the function (used for making Alice's random sending bits).</li>\n",
    "    <li>A function to generate a random array of Ds and Rs that is as long as the integer inputted into the function (used for making Alice and Bob's encoding and decoding bases).</li>\n",
    "    <li>A function that returns true or false with a probability of the integer that is inputted into it (used for making a percentage of the photons get lost in transmission and for making a percentage of the received bits contain errors).</li>\n",
    "</ul>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "def string_to_array(string): # Used for user display purposes to convert each character in a string to entries in an array\n",
    "    array=[]\n",
    "    for i in string: # iterate through each character in the string\n",
    "        array.append(i) # append each character to the array\n",
    "    return array\n",
    "\n",
    "def array_to_string(array): # Used for user display purposes to convert each entry in an array to characters in a string\n",
    "    string=\"\"\n",
    "    for i in array: # iterate through each entry in the array\n",
    "        string = string + i # append each entry to the string\n",
    "    return string\n",
    "\n",
    "def generate_random_bits(length): # Used to generate a random array of 1s and 0s of a given length\n",
    "    binary_string = ''.join(random.choice('01') for _ in range(length)) # generate a random string of 1s and 0s of a given length\n",
    "    return string_to_array(binary_string) # Convert the string to an array and return it\n",
    "\n",
    "def generate_random_bases(length): # Used to generate a random array of Ds and Rs of a given length\n",
    "    bases_string = ''.join(random.choice('DR') for _ in range(length)) # generate a random string of Ds and Rs of a given length\n",
    "    return string_to_array(bases_string) # Convert the string to an array and return it\n",
    "\n",
    "def test_percentage(prob): # Essentially a coin flip with a given probability, has the given probability of returning true \n",
    "    # generate a random float between 0 and 1\n",
    "    if random.random() < (prob/100): # if the float is less than the probability as a decimal, return true\n",
    "        return True\n",
    "    else: # otherwise, return false\n",
    "        return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Generate Alice & Bob's Properties</h3>\n",
    "<hr>\n",
    "<p>\n",
    "    Alice:\n",
    "    <ul>\n",
    "        <li>Generate the string of random bits Alice will send to Bob.</li>\n",
    "        <li>Generate the random bases Alice will use to encode the random bits.</li>\n",
    "        <li>Generate the encoded photon states that are being sent to Bob</li>\n",
    "    </ul>\n",
    "    Bob:\n",
    "    <ul>\n",
    "        <li>Generate the random bases Bob will use to decode the photon states received from Alice</li>\n",
    "    </ul>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Alice's properties\n",
    "alice_bits=generate_random_bits(no_bits) # Generate the random bits that Alice will send\n",
    "sending_bases=generate_random_bases(no_bits) # Generate the random bases that Alice will use to encode her bits\n",
    "# Generate the photons states that Alice will send based on her random bits and bases\n",
    "sending_photons=[]\n",
    "for i in range(len(sending_bases)): # iterate through each bit/base pair\n",
    "    if sending_bases[i] == \"D\" and alice_bits[i] == \"0\": # if the base is D and the bit is 0, append a \"/\" state to the photons array\n",
    "        sending_photons.append(\"/\")\n",
    "    elif sending_bases[i] == \"D\" and alice_bits[i] == \"1\": # if the base is D and the bit is 1, append a \"\\\" state to the photons array\n",
    "        sending_photons.append(\"\\\\\") # The \"\\\" character is escaped in python, so we need to use \"\\\\\" to get a single \"\\\" character\n",
    "    elif sending_bases[i] == \"R\" and alice_bits[i] == \"0\": # if the base is R and the bit is 0, append a \"-\" state to the photons array\n",
    "        sending_photons.append(\"-\")\n",
    "    elif sending_bases[i] == \"R\" and alice_bits[i] == \"1\": # if the base is R and the bit is 1, append a \"|\" state to the photons array\n",
    "        sending_photons.append(\"|\")\n",
    "    else:\n",
    "        print(\"Error: D, R or 0, 1 are expected\") # This should never happen, but is here just in case for debugging\n",
    "\n",
    "## Bob's properties\n",
    "receiving_bases=generate_random_bases(no_bits) # Generate the random bases that Bob will use to decode Alice's bits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Decode Received Photon States</h3>\n",
    "<hr>\n",
    "<p>Using Bob's random string of bases, the received photon states are decoded. The probability of photons being lost in transmission is applied, then if Bob's base is correct then the received bit is \"written down\" correctly and if Bob's base is incorrect then there is a 50:50 chance that the bit is \"written down\" incorrectly.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "received_bits=[]\n",
    "for i in range(len(sending_bases)): # iterate through each received photon state & decoding base pair\n",
    "    if test_percentage(lose_in_tx_chance) == True: # if the photon is lost in transmission, append a space to the received bits array\n",
    "        received_bits.append(\" \")\n",
    "    else:\n",
    "        # Incoming: /  (0)\n",
    "        if sending_photons[i] == \"/\":\n",
    "            if receiving_bases[i] == \"D\": # if the photon is received on the correct base, append the bit to the received bits array\n",
    "                received_bits.append(\"0\")\n",
    "            elif receiving_bases[i] == \"R\":\n",
    "                if test_percentage(50) == True: # if the photon is received on the wrong base, there is a 50% chance of it being decoded incorrectly\n",
    "                    received_bits.append(\"0\")\n",
    "                else:\n",
    "                    received_bits.append(\"1\")\n",
    "            else:\n",
    "                print(\"Error (photon = /):\",receiving_bases[i],\"is not an expected state [D,R] of the random receiving bases\") # This should never happen, but is here just in case for debugging\n",
    "\n",
    "        # Incoming: \\  (1)\n",
    "        elif sending_photons[i] == \"\\\\\":\n",
    "            if receiving_bases[i] == \"D\": # if the photon is received on the correct base, append the bit to the received bits array\n",
    "                received_bits.append(\"1\")\n",
    "            elif receiving_bases[i] == \"R\":\n",
    "                if test_percentage(50) == True: # if the photon is received on the wrong base, there is a 50% chance of it being decoded incorrectly\n",
    "                    received_bits.append(\"0\")\n",
    "                else:\n",
    "                    received_bits.append(\"1\")\n",
    "            else:\n",
    "                print(\"Error (photon = \\):\",receiving_bases[i],\"is not an expected state [D,R] of the random receiving bases\") # This should never happen, but is here just in case for debugging\n",
    "\n",
    "        # Incoming: -  (0)\n",
    "        elif sending_photons[i] == \"-\":\n",
    "            if receiving_bases[i] == \"R\": # if the photon is received on the correct base, append the bit to the received bits array\n",
    "                received_bits.append(\"0\")\n",
    "            elif receiving_bases[i] == \"D\":\n",
    "                if test_percentage(50) == True: # if the photon is received on the wrong base, there is a 50% chance of it being decoded incorrectly\n",
    "                    received_bits.append(\"0\")\n",
    "                else:\n",
    "                    received_bits.append(\"1\")\n",
    "            else:\n",
    "                print(\"Error (photon = -):\",receiving_bases[i],\"is not an expected state [D,R] of the random receiving bases\") # This should never happen, but is here just in case for debugging\n",
    "\n",
    "        # Incoming: |  (1)\n",
    "        elif sending_photons[i] == \"|\":\n",
    "            if receiving_bases[i] == \"R\": # if the photon is received on the correct base, append the bit to the received bits array\n",
    "                received_bits.append(\"1\")\n",
    "            elif receiving_bases[i] == \"D\":\n",
    "                if test_percentage(50) == True: # if the photon is received on the wrong base, there is a 50% chance of it being decoded incorrectly\n",
    "                    received_bits.append(\"0\")\n",
    "                else:\n",
    "                    received_bits.append(\"1\")\n",
    "            else:\n",
    "                print(\"Error (photon = |):\",receiving_bases[i],\"is not an expected state [D,R] of the random receiving bases\") # This should never happen, but is here just in case for debugging\n",
    "\n",
    "        else:\n",
    "            print(\"Error:\",sending_photons[i],\"is not an expected state [-,|,/,\\\\] of of transmission photons\") # This should never happen, but is here just in case for debugging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Public Discussion</h3>\n",
    "<hr>\n",
    "<p>During the public discussion stage, Bob and Alice compare bases and when they both have the same base, both can note down that the associated bit is correct. This allows them to check which bits were received correctly without releasing the bits themselves to the open public.<br>\n",
    "From these correctly received bits a private key can be produced.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Key: 0011010\n",
      "Key Length (bits): 7\n"
     ]
    }
   ],
   "source": [
    "public_dis=[]\n",
    "for i in range(len(sending_bases)): # iterate through each encoding/decoding base pair\n",
    "    if received_bits[i] != \" \" and sending_bases[i] == receiving_bases[i]: # if the received wasn't lost in transmission and the bases match, append \"OK\" to the public discussion array\n",
    "        public_dis.append(\"OK\")\n",
    "    else:\n",
    "        public_dis.append(\" \") # otherwise, append a space to the public discussion array\n",
    "if len(sending_bases) != len(public_dis):\n",
    "    print(\"Error: the public discussion array is either too long or too short\") # This should never happen, but is here just in case for debugging\n",
    "\n",
    "# Key generation\n",
    "key_bit=[] # For table display purposes (includes spaces)\n",
    "key_array=[] # Array of just the key bits, no spaces\n",
    "for i in range(len(public_dis)):\n",
    "    if public_dis[i] == \"OK\":\n",
    "        key_bit.append(received_bits[i])\n",
    "        key_array.append(received_bits[i])\n",
    "    else:\n",
    "        key_bit.append(\" \")\n",
    "if len(public_dis) != len(key_bit): # This should never happen, but is here just in case for debugging\n",
    "    print(\"Error: the public discussion array is either too long or too short\")\n",
    "\n",
    "# Turn key array into string\n",
    "key=array_to_string(key_array)\n",
    "print(\"Key:\",key)\n",
    "print(\"Key Length (bits):\",len(key)) # Print length of key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Display Each Stage in a Table</h3>\n",
    "<hr>\n",
    "<p>A label is inserted at the start of each array so that it is labelled in the table then each stage of BB84 (represented by the data from that stage in an array) is combined into a 2D array and then displayed as a table for easy viewing of how the protocol produced that key in a stage-by-stage manner.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[4mQuantum Key Distribution - BB84 Protocol\u001b[0m\u001b[0m\n",
      "┌─────────────────────┬───┬───┬────┬────┬────┬───┬───┬───┬───┬────┬───┬────┬───┬────┬───┬────┐\n",
      "│ Alice's bits        │ 0 │ 1 │ 0  │ 0  │ 1  │ 0 │ 1 │ 0 │ 1 │ 1  │ 1 │ 0  │ 1 │ 1  │ 0 │ 0  │\n",
      "├─────────────────────┼───┼───┼────┼────┼────┼───┼───┼───┼───┼────┼───┼────┼───┼────┼───┼────┤\n",
      "│ Sending bases       │ D │ R │ R  │ R  │ D  │ D │ R │ D │ D │ D  │ R │ D  │ R │ R  │ R │ D  │\n",
      "├─────────────────────┼───┼───┼────┼────┼────┼───┼───┼───┼───┼────┼───┼────┼───┼────┼───┼────┤\n",
      "│ Photon state        │ / │ | │ -  │ -  │ \\  │ / │ | │ / │ \\ │ \\  │ | │ /  │ | │ |  │ - │ /  │\n",
      "├─────────────────────┼───┼───┼────┼────┼────┼───┼───┼───┼───┼────┼───┼────┼───┼────┼───┼────┤\n",
      "│ Receiving bases     │ R │ D │ R  │ R  │ D  │ R │ D │ R │ R │ D  │ D │ D  │ D │ R  │ D │ D  │\n",
      "├─────────────────────┼───┼───┼────┼────┼────┼───┼───┼───┼───┼────┼───┼────┼───┼────┼───┼────┤\n",
      "│ Bob's received bits │ 1 │ 1 │ 0  │ 0  │ 1  │ 1 │ 1 │ 1 │ 0 │ 1  │ 0 │ 0  │ 1 │ 1  │ 0 │ 0  │\n",
      "├─────────────────────┼───┼───┼────┼────┼────┼───┼───┼───┼───┼────┼───┼────┼───┼────┼───┼────┤\n",
      "│ Public discussion   │   │   │ OK │ OK │ OK │   │   │   │   │ OK │   │ OK │   │ OK │   │ OK │\n",
      "├─────────────────────┼───┼───┼────┼────┼────┼───┼───┼───┼───┼────┼───┼────┼───┼────┼───┼────┤\n",
      "│ Key produced        │   │   │ 0  │ 0  │ 1  │   │   │   │   │ 1  │   │ 0  │   │ 1  │   │ 0  │\n",
      "└─────────────────────┴───┴───┴────┴────┴────┴───┴───┴───┴───┴────┴───┴────┴───┴────┴───┴────┘\n"
     ]
    }
   ],
   "source": [
    "print(\"\\033[1m\"+ \"\\033[4m\" + \"Quantum Key Distribution - BB84 Protocol\" + \"\\033[0m\" + \"\\033[0m\") # bold and underline table title\n",
    "# Add label to each array for tabulation\n",
    "alice_bits.insert(0, \"Alice's bits\")\n",
    "sending_bases.insert(0, \"Sending bases\")\n",
    "sending_photons.insert(0, \"Photon state\")\n",
    "receiving_bases.insert(0, \"Receiving bases\")\n",
    "received_bits.insert(0, \"Bob's received bits\")\n",
    "public_dis.insert(0, \"Public discussion\")\n",
    "key_bit.insert(0, \"Key produced\")\n",
    "data = [alice_bits,sending_bases,sending_photons,receiving_bases,received_bits,public_dis,key_bit]\n",
    "print(tabulate(data, tablefmt=\"simple_grid\")) # Print table\n",
    "# Remove labels from each array so they can be used for further calculations\n",
    "for array in data:\n",
    "    array.pop(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Conclusion</h3>\n",
    "<hr>\n",
    "<p>This implementation is not perfect and currently in the works is the inclusion of error correction via the cascade protocol.<br>\n",
    "See the file named \"<i>error-correction.ipynb</i>\" for the current version of this code.</p>"
   ]
  }
 ],
 "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.11.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}