{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "de73025f",
   "metadata": {},
   "source": [
    "# Using QDK to explore the Carbon Error Correction Code.\n",
    "\n",
    "This notebook will show how you can use the QDK with the Neutral Atom Device to explore how the Carbon error correction code runs both with and without noise.\n",
    "\n",
    "For more information on the Carbon family of codes, see [Demonstration of logical qubits and repeated error correction with better-than-physical error rates](https://arxiv.org/abs/2404.02280).\n",
    "\n",
    "First, we import the required packages and utilities, initialize a `NeutralAtomDevice` and initialize the Q# environment with the project from the current folder. This includes a Python implementation of the Carbon (C12) decoder that can turn the raw physical qubit measurement results into corrected, logical qubit measurement results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1929d94c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import decoder\n",
    "from qdk import qsharp, code\n",
    "from qdk.widgets import Histogram\n",
    "from qdk.simulation import NeutralAtomDevice, NoiseConfig\n",
    "\n",
    "# Get an instance of the NeutralAtomDevice machine\n",
    "device = NeutralAtomDevice()\n",
    "\n",
    "# Initialize Q# with the project and adaptive profile configured\n",
    "qsharp.init(target_profile=qsharp.TargetProfile.Adaptive_RIF, project_root=\".\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d36f1f",
   "metadata": {},
   "source": [
    "Next, this `%%qsharp` cell sets up the experiment we want to run, which is parametrized by the number of repetitions, which Pauli basis to prepare and measure with, and how many blocks of Carbon (C12) with ancillas should be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6177ed62",
   "metadata": {
    "vscode": {
     "languageId": "qsharp"
    }
   },
   "outputs": [],
   "source": [
    "%%qsharp\n",
    "import Std.Diagnostics.Fact;\n",
    "import Utils.TransversalCNOT;\n",
    "import C12;\n",
    "\n",
    "operation PerformTeleportExperiment(ec_repetitions : Int, basis : Pauli, num_blocks : Int) : (Result[], (Result[], Result[])[], Result[])[] {\n",
    "    Fact(num_blocks > 0, \"need at least one block\");\n",
    "    let size = (12 + 16); // 12 data + 16 ancillas per block\n",
    "    use qubits = Qubit[size * num_blocks];\n",
    "    mutable results = [];\n",
    "    for i in 0..(num_blocks - 1) {\n",
    "        let logical_block = qubits[(0+(i*size))..(11+(i*size))];\n",
    "        let ancillas = qubits[(12+(i*size))..(27+(i*size))];\n",
    "\n",
    "        results += [C12_Teleport(ec_repetitions, basis, logical_block, ancillas)];\n",
    "    }\n",
    "    return results;\n",
    "}\n",
    "\n",
    "operation C12_Teleport(ec_repetitions : Int, basis : Pauli, logical_block : Qubit[], ancillas : Qubit[]) : (Result[], (Result[], Result[])[], Result[]) {\n",
    "    Fact(basis == PauliZ or basis == PauliX or basis == PauliI, \"only PauliZ and PauliX supported\");\n",
    "    Fact(Length(logical_block) == 12, \"logical block must be 12 qubits\");\n",
    "    Fact(Length(ancillas) == 16, \"16 ancillas are required\");\n",
    "\n",
    "    // Prepare in the requested basis\n",
    "    mutable preselect = if basis == PauliX {\n",
    "        C12.PrepareXX(logical_block, ancillas[...3])\n",
    "    } else {\n",
    "        C12.PrepareZZ(logical_block, ancillas[...3])\n",
    "    };\n",
    "\n",
    "    mutable syndromes = [];\n",
    "\n",
    "    for _ in 1..(ec_repetitions) {\n",
    "        // Sequential teleport on..\n",
    "        // Prepare Z, Teleport X\n",
    "        set preselect += C12.PrepareZZ(ancillas[...11], ancillas[12...]);\n",
    "        TransversalCNOT(logical_block, ancillas[...11]);\n",
    "        ApplyToEach(H, logical_block);\n",
    "        let syndrome_x = MResetEachZ(logical_block);\n",
    "\n",
    "        // Prepare X, Teleport Z\n",
    "        set preselect += C12.PrepareXX(logical_block, ancillas[12...]);\n",
    "        TransversalCNOT(logical_block, ancillas[...11]);\n",
    "        let syndrome_z = MResetEachZ(ancillas[...11]);\n",
    "        set syndromes += [(syndrome_x, syndrome_z)];\n",
    "    }\n",
    "\n",
    "    // Final measurement\n",
    "    if basis == PauliX {\n",
    "        ApplyToEach(H, logical_block);\n",
    "    }\n",
    "    let final = MResetEachZ(logical_block);\n",
    "\n",
    "    (preselect, syndromes, final)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21031218",
   "metadata": {},
   "source": [
    "With the experiment set up, we first try running the Q# with the default, sparse state simulation and no noise. This `Histogram` should have one bar with all `[[Zero, Zero]]` results, indicating that the logical teleport was performed successfully for one Carbon (C12) block in all 100 shots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8d0c4cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = qsharp.run(code.PerformTeleportExperiment, 100, 1, qsharp.Pauli.Z, 1)\n",
    "corrected_logical_results = decoder.decode_results(results, \"Z\")\n",
    "\n",
    "# Use counter to make a simple histogram of the corrected logical results.\n",
    "Histogram(map(str, corrected_logical_results))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f2ead22",
   "metadata": {},
   "source": [
    "Next, we compile the program, this time with two Carbon (C12) blocks. To visualize and test larger programs, try increasing the number of blocks (the last property) to larger values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "313cb353",
   "metadata": {},
   "outputs": [],
   "source": [
    "qir = qsharp.compile(code.PerformTeleportExperiment, 1, qsharp.Pauli.Z, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c18dcefe",
   "metadata": {},
   "source": [
    "We can use the Neutral Atom Device we previously initialized to trace the compiled program. This will decompose the gates in the QIR and perform layout and scheduling for the device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce29c823",
   "metadata": {},
   "outputs": [],
   "source": [
    "device.show_trace(qir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fbd6f5f",
   "metadata": {},
   "source": [
    "We can also use the device to perform Clifford simulations of the experiment. Clifford simulation can scale to multiple Carbon (C12) blocks and multiple rounds of error correction. Here again, a noiseless simulation should produce a single bar in the histogram, this time with a number of logical result pairs equal to the number of blocks compiled into the program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e05e0f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = device.simulate(qir, shots=1000, type=\"clifford\")\n",
    "corrected_logical_results = decoder.decode_results(res, \"Z\")\n",
    "Histogram(map(str, corrected_logical_results))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b97ffa2c",
   "metadata": {},
   "source": [
    "Error correction is more interesting with noise configured, so now we can use the `NoiseConfig` class to specify what noise we'll use. These default values produce interesting results, showing some shots rejected as \"preselect\" and others corrected by the decoder. If we filter out those \"preselect\" results, we can see that the majority of the accepted results have been corrected into the expected `[[Zero, Zero], [Zero, Zero]]` result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0f2063e",
   "metadata": {},
   "outputs": [],
   "source": [
    "noise = NoiseConfig()\n",
    "\n",
    "def set_depolarizing_error(table, p: float):\n",
    "    \"\"\"One of three equally likely Pauli errors occurs with probability `p`.\"\"\"\n",
    "    table.x = p / 3\n",
    "    table.y = p / 3\n",
    "    table.z = p / 3\n",
    "\n",
    "# 1Q error\n",
    "set_depolarizing_error(noise.sx, 0.00225)\n",
    "\n",
    "# Error per CZ\n",
    "noise.cz.iz = 0.0105 / 3\n",
    "noise.cz.zi = 0.0105 / 3\n",
    "noise.cz.zz = 0.0105 / 3\n",
    "\n",
    "# Loss per CZ\n",
    "noise.cz.loss = 0.0027\n",
    "\n",
    "# Error for moved atoms\n",
    "noise.mov.z = 1e-3\n",
    "\n",
    "# Loss per moved atom\n",
    "noise.mov.loss = 0.0005\n",
    "\n",
    "res = device.simulate(qir, shots=10000, noise=noise, type=\"clifford\")\n",
    "corrected_logical_results = decoder.decode_results(res, \"Z\")\n",
    "display(Histogram(map(str, corrected_logical_results)))\n",
    "display(Histogram(map(str, filter(lambda res: not any(r == 'PREselect' for r in res), corrected_logical_results))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85cc8664",
   "metadata": {},
   "source": [
    "To make it easier to see the results, this further processes them to treat each Carbon (C12) block as separate rather than plot the joint behavior of the blocks. This again shows both the full results in the first histogram and the filtered, non-preselect results in the second histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35e7dc3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "flattened = [r for res in corrected_logical_results for r in res]\n",
    "display(Histogram(map(str, flattened)))\n",
    "display(Histogram(map(str, filter(lambda res: res != 'PREselect', flattened))))"
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
