{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Postselection: exit and panic\n",
    "\n",
    "**Download this notebook - {nb-download}`postselect.ipynb`**\n",
    "\n",
    "\n",
    "In this example we will look at two ways of ending a Guppy program early:\n",
    "\n",
    "1. `exit` will end the current shot and carry on with subsequent ones. We will use this to\n",
    "   implement postselection.\n",
    "2. `panic` is used to signal some unexpected error and as such it will end the shot and\n",
    "   not run any subsequent ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-29T12:44:11.876698Z",
     "start_time": "2026-04-29T12:44:11.147001Z"
    },
    "execution": {
     "iopub.execute_input": "2026-04-29T13:43:27.744652Z",
     "iopub.status.busy": "2026-04-29T13:43:27.744441Z",
     "iopub.status.idle": "2026-04-29T13:43:28.325191Z",
     "shell.execute_reply": "2026-04-29T13:43:28.324743Z"
    }
   },
   "outputs": [],
   "source": [
    "from typing import Counter\n",
    "\n",
    "from guppylang import guppy\n",
    "from guppylang.std.builtins import output, array, exit, panic\n",
    "from guppylang.std.quantum import measure_array, measure, qubit, h, cx, collect_measurements\n",
    "from guppylang.defs import GuppyFunctionDefinition\n",
    "\n",
    "from selene_sim import DepolarizingErrorModel\n",
    "from selene_sim.exceptions import SelenePanicError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Postselection\n",
    "\n",
    "We can use postselection to implement fault tolerant state preparation for the [[7, 1,\n",
    "3]] [Steane code](https://en.wikipedia.org/wiki/Steane_code).\n",
    "\n",
    "Let's first define our \"Steane qubit\" as a struct containing our 7 data qubits, then\n",
    "write a function to prepare an encoded $|0\\rangle$ non-fault tolerantly. We use the\n",
    "preparation circuit from [_Realization of real-time fault-tolerant quantum error correction_](https://arxiv.org/abs/2107.07505).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-29T12:44:13.697644Z",
     "start_time": "2026-04-29T12:44:13.678652Z"
    },
    "execution": {
     "iopub.execute_input": "2026-04-29T13:43:28.326342Z",
     "iopub.status.busy": "2026-04-29T13:43:28.326237Z",
     "iopub.status.idle": "2026-04-29T13:43:28.328699Z",
     "shell.execute_reply": "2026-04-29T13:43:28.328461Z"
    }
   },
   "outputs": [],
   "source": [
    "@guppy.struct\n",
    "class SteaneQubit:\n",
    "    data_qs: array[qubit, 7]\n",
    "\n",
    "\n",
    "@guppy\n",
    "def non_ft_zero() -> SteaneQubit:\n",
    "    data_qubits = array(qubit() for _ in range(7))\n",
    "    plus_ids = array(0, 4, 6)\n",
    "    for i in plus_ids:\n",
    "        h(data_qubits[i])\n",
    "\n",
    "    cx_pairs = array((0, 1), (4, 5), (6, 3), (6, 5), (4, 2), (0, 3), (4, 1), (3, 2))\n",
    "    for c, t in cx_pairs:\n",
    "        cx(data_qubits[c], data_qubits[t])\n",
    "    return SteaneQubit(data_qubits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now implement fault-tolerant preparation using postselection. We can use an\n",
    "ancilla to check the prepared state, and if we detect an error use `exit` to end the\n",
    "shot with a message about why we exited."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-29T12:44:15.908094Z",
     "start_time": "2026-04-29T12:44:15.890556Z"
    },
    "execution": {
     "iopub.execute_input": "2026-04-29T13:43:28.329674Z",
     "iopub.status.busy": "2026-04-29T13:43:28.329608Z",
     "iopub.status.idle": "2026-04-29T13:43:28.331585Z",
     "shell.execute_reply": "2026-04-29T13:43:28.331287Z"
    }
   },
   "outputs": [],
   "source": [
    "@guppy\n",
    "def ft_zero() -> SteaneQubit:\n",
    "    q = non_ft_zero()\n",
    "    ancilla = qubit()\n",
    "    flags = array(1, 3, 5)\n",
    "    for f in flags:\n",
    "        cx(q.data_qs[f], ancilla)\n",
    "    if measure(ancilla):\n",
    "        exit(\"Postselected: FT prep failed\", 1)\n",
    "    return q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a couple of utility functions - a Guppy function to check the parity of a\n",
    "bit array, and a python function to run our program and report the results.\n",
    "\n",
    "We use a simple depolarizing error model to induce errors in the preparation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-29T12:44:46.321966Z",
     "start_time": "2026-04-29T12:44:46.308515Z"
    },
    "execution": {
     "iopub.execute_input": "2026-04-29T13:43:28.332448Z",
     "iopub.status.busy": "2026-04-29T13:43:28.332397Z",
     "iopub.status.idle": "2026-04-29T13:43:28.334710Z",
     "shell.execute_reply": "2026-04-29T13:43:28.334465Z"
    }
   },
   "outputs": [],
   "source": [
    "@guppy\n",
    "def parity_check[n: nat](data_bits: array[bool, n]) -> bool:\n",
    "    out = False\n",
    "    for i in range(n):\n",
    "        out ^= data_bits[i]\n",
    "    return out\n",
    "\n",
    "\n",
    "error_model = DepolarizingErrorModel(\n",
    "    random_seed=1234,\n",
    "    # single qubit gate error rate\n",
    "    p_1q=1e-3,\n",
    "    # two qubit gate error rate\n",
    "    p_2q=1e-3,\n",
    "    # set state preparation and measurement error rates to 0\n",
    "    p_meas=0,\n",
    "    p_init=0,\n",
    ")\n",
    "\n",
    "\n",
    "def run(main_def: GuppyFunctionDefinition) -> Counter:\n",
    "    res = (\n",
    "        main_def.emulator(n_qubits=8)\n",
    "        .stabilizer_sim()\n",
    "        .with_seed(42)\n",
    "        .with_shots(1000)\n",
    "        .with_error_model(error_model)\n",
    "        .run()\n",
    "    )\n",
    "\n",
    "    return res.collated_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now define our `main` program and run it. We know that all basis states of the\n",
    "encoded Steane $|0\\rangle$ state have a $0$ parity, so we can use that to verify our preparation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-29T12:44:51.809478Z",
     "start_time": "2026-04-29T12:44:49.947720Z"
    },
    "execution": {
     "iopub.execute_input": "2026-04-29T13:43:28.335629Z",
     "iopub.status.busy": "2026-04-29T13:43:28.335561Z",
     "iopub.status.idle": "2026-04-29T13:43:30.128865Z",
     "shell.execute_reply": "2026-04-29T13:43:30.128453Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({(('parity', '0'),): 961,\n",
       "         (('exit: Postselected: FT prep failed', '1'),): 34,\n",
       "         (('parity', '1'),): 5})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@guppy\n",
    "def main() -> None:\n",
    "    steane_q = ft_zero()\n",
    "\n",
    "    # Measure the data qubits\n",
    "    data = measure_array(steane_q.data_qs)\n",
    "    output(\"parity\", parity_check(collect_measurements(data)))\n",
    "\n",
    "\n",
    "run(main)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the state preparation succeeded in 968 cases out of 1000.\n",
    "In 23 of the unsuccessful cases, an error was detected on the ancillas and was\n",
    "discarded through postselection. In the remaining 9 cases, the state failed\n",
    "in a way that was not detected by postselection but was instead detected by a\n",
    "parity check performed after measuring all of the qubits.\n",
    "\n",
    "Note the result tag in the discarded shot is prefixed with `exit: ` followed\n",
    "by the specified message, and the value of the result entry is the error code\n",
    "(1 in this case)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Panic\n",
    "\n",
    "The `panic` function is similar to `exit` but is used for exceptional circumstances -\n",
    "when something unexpected has gone wrong. For example we could define a physical\n",
    "hadamard function that takes an index to act on the data qubit array. If the index is\n",
    "out of bounds, we can panic with a helpful message. This will raise an error during the\n",
    "simulation, and no subsequent shots will be run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-29T12:45:04.562132Z",
     "start_time": "2026-04-29T12:45:02.824971Z"
    },
    "execution": {
     "iopub.execute_input": "2026-04-29T13:43:30.144688Z",
     "iopub.status.busy": "2026-04-29T13:43:30.144593Z",
     "iopub.status.idle": "2026-04-29T13:43:31.808834Z",
     "shell.execute_reply": "2026-04-29T13:43:31.808467Z"
    },
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "EmulatorError",
     "evalue": "Panic (#1001): Invalid data index in physical_h",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mEmulatorError\u001b[39m                             Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[6]\u001b[39m\u001b[32m, line 18\u001b[39m\n\u001b[32m     15\u001b[39m     data = measure_array(steane_q.data_qs)\n\u001b[32m     16\u001b[39m     output(\u001b[33m\"\u001b[39m\u001b[33mparity\u001b[39m\u001b[33m\"\u001b[39m, parity_check(collect_measurements(data)))\n\u001b[32m---> \u001b[39m\u001b[32m18\u001b[39m \u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmain\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[4]\u001b[39m\u001b[32m, line 28\u001b[39m, in \u001b[36mrun\u001b[39m\u001b[34m(main_def)\u001b[39m\n\u001b[32m     21\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mrun\u001b[39m(main_def: GuppyFunctionDefinition) -> Counter:\n\u001b[32m     22\u001b[39m     res = (\n\u001b[32m     23\u001b[39m         \u001b[43mmain_def\u001b[49m\u001b[43m.\u001b[49m\u001b[43memulator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mn_qubits\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m8\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m     24\u001b[39m \u001b[43m        \u001b[49m\u001b[43m.\u001b[49m\u001b[43mstabilizer_sim\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     25\u001b[39m \u001b[43m        \u001b[49m\u001b[43m.\u001b[49m\u001b[43mwith_seed\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m42\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m     26\u001b[39m \u001b[43m        \u001b[49m\u001b[43m.\u001b[49m\u001b[43mwith_shots\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m1000\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m     27\u001b[39m \u001b[43m        \u001b[49m\u001b[43m.\u001b[49m\u001b[43mwith_error_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43merror_model\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m---> \u001b[39m\u001b[32m28\u001b[39m \u001b[43m        \u001b[49m\u001b[43m.\u001b[49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     29\u001b[39m     )\n\u001b[32m     31\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m res.collated_counts()\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/app/.venv/lib/python3.12/site-packages/guppylang/emulator/instance.py:311\u001b[39m, in \u001b[36mEmulatorInstance.run\u001b[39m\u001b[34m(self, **args)\u001b[39m\n\u001b[32m    306\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m args:\n\u001b[32m    307\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m EntrypointArgValueError(\n\u001b[32m    308\u001b[39m             \u001b[33m\"\u001b[39m\u001b[33mThis entrypoint takes no runtime arguments, but got: \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    309\u001b[39m             + \u001b[33m\"\u001b[39m\u001b[33m, \u001b[39m\u001b[33m\"\u001b[39m.join(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33m`\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m`\u001b[39m\u001b[33m\"\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m name \u001b[38;5;129;01min\u001b[39;00m args)\n\u001b[32m    310\u001b[39m         )\n\u001b[32m--> \u001b[39m\u001b[32m311\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_run_and_collect_results\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    313\u001b[39m validate_record(\u001b[38;5;28mself\u001b[39m._arg_specs, args)\n\u001b[32m    314\u001b[39m provider = ArgProvider()\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/app/.venv/lib/python3.12/site-packages/guppylang/emulator/instance.py:364\u001b[39m, in \u001b[36mEmulatorInstance._run_and_collect_results\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m    362\u001b[39m \u001b[38;5;250m\u001b[39m\u001b[33;03m\"\"\"Run the instance and retain any configured analysis collectors.\"\"\"\u001b[39;00m\n\u001b[32m    363\u001b[39m event_hook, circuit_extractor, metric_store = \u001b[38;5;28mself\u001b[39m._analysis_event_hook()\n\u001b[32m--> \u001b[39m\u001b[32m364\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_collect_results\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    365\u001b[39m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_run_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43mevent_hook\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    366\u001b[39m \u001b[43m    \u001b[49m\u001b[43mcircuit_extractor\u001b[49m\u001b[43m=\u001b[49m\u001b[43mcircuit_extractor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    367\u001b[39m \u001b[43m    \u001b[49m\u001b[43mmetric_store\u001b[49m\u001b[43m=\u001b[49m\u001b[43mmetric_store\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    368\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/app/.venv/lib/python3.12/site-packages/guppylang/emulator/instance.py:420\u001b[39m, in \u001b[36mEmulatorInstance._collect_results\u001b[39m\u001b[34m(self, result_stream, circuit_extractor, metric_store)\u001b[39m\n\u001b[32m    416\u001b[39m             shot_results.append(tag, value)\n\u001b[32m    417\u001b[39m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:  \u001b[38;5;66;03m# noqa: BLE001\u001b[39;00m\n\u001b[32m    418\u001b[39m         \u001b[38;5;66;03m# In this case, casting a wide net on exceptions is\u001b[39;00m\n\u001b[32m    419\u001b[39m         \u001b[38;5;66;03m# suitable.\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m420\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m EmulatorError(\n\u001b[32m    421\u001b[39m             completed_shots=EmulatorResult(\n\u001b[32m    422\u001b[39m                 all_results,\n\u001b[32m    423\u001b[39m                 _circuit_extractor=circuit_extractor,\n\u001b[32m    424\u001b[39m                 _metric_store=metric_store,\n\u001b[32m    425\u001b[39m             ),\n\u001b[32m    426\u001b[39m             failing_shot=shot_results,\n\u001b[32m    427\u001b[39m             underlying_exception=e,\n\u001b[32m    428\u001b[39m         ) \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m    429\u001b[39m     all_results.append(shot_results)\n\u001b[32m    430\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m EmulatorResult(\n\u001b[32m    431\u001b[39m     all_results,\n\u001b[32m    432\u001b[39m     _circuit_extractor=circuit_extractor,\n\u001b[32m    433\u001b[39m     _metric_store=metric_store,\n\u001b[32m    434\u001b[39m )\n",
      "\u001b[31mEmulatorError\u001b[39m: Panic (#1001): Invalid data index in physical_h"
     ]
    }
   ],
   "source": [
    "@guppy\n",
    "def physical_h(q: SteaneQubit, data_idx: int) -> None:\n",
    "    if data_idx >= 7:\n",
    "        panic(\"Invalid data index in physical_h\")\n",
    "    h(q.data_qs[data_idx])\n",
    "\n",
    "\n",
    "@guppy\n",
    "def main() -> None:\n",
    "    steane_q = ft_zero()\n",
    "\n",
    "    # add a physical H gate\n",
    "    physical_h(steane_q, 8)\n",
    "\n",
    "    data = measure_array(steane_q.data_qs)\n",
    "    output(\"parity\", parity_check(collect_measurements(data)))\n",
    "\n",
    "run(main)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "guppylang (3.14.0)",
   "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}