{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1d65e218",
   "metadata": {},
   "source": [
    "# Statevector testing utilities\n",
    "\n",
    "**Download Notebook** - {nb-download}`statevector_testing.ipynb`\n",
    "\n",
    "Simulation lets you check quantum amplitudes, complete operations, and individual measurement branches. These examples introduce the helpers for each task.\n",
    "\n",
    "- **Inspect a state:** compare its amplitudes and relative phases.\n",
    "- **Extract an operation:** check its action on every computational-basis input.\n",
    "- **Post-select a branch:** inspect the state or operation conditioned on chosen outcomes.\n",
    "- **Replay measurements:** exercise success, failure, and correction paths predictably.\n",
    "\n",
    "Run the cells in order; later examples reuse earlier imports and definitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cache0581",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang.std.quantum import (\n",
    "    discard_array,\n",
    "    discard,\n",
    "    x,\n",
    "    ry, \n",
    "    cx,\n",
    "    qubit,\n",
    "    h\n",
    ")\n",
    "from guppylang.std.angles import angle\n",
    "from guppylang import guppy\n",
    "from guppylang.std.debug import state_output\n",
    "from guppylang.std.builtins import comptime, array, nat\n",
    "from guppyalgos.utils import qarray\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23d0e6d7",
   "metadata": {},
   "source": [
    "## Inspect a statevector\n",
    "\n",
    "Use `get_statevector` to inspect the complex amplitudes of a prepared state:\n",
    "\n",
    "$$\n",
    "|\\psi\\rangle = \\sum_{x=0}^{2^n-1} a_x |x\\rangle,\n",
    "\\qquad\n",
    "P(x)=|a_x|^2.\n",
    "$$\n",
    "\n",
    "- Record the state with `state_output(\"result_state\", state)` before discarding the qubits.\n",
    "- The tag must be `\"result_state\"`; this is a simulation snapshot, not a measurement.\n",
    "- The returned NumPy array contains the amplitudes $a_x$, including their relative phases.\n",
    "\n",
    "The example below applies the same Y rotation to each of four qubits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "da0d32b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_qubits = 4\n",
    "@guppy\n",
    "def main_sv() -> None:\n",
    "    state = qarray(n_qubits)\n",
    "\n",
    "    for i in range(n_qubits):\n",
    "        ry(state[i], angle(3.14/2))\n",
    "\n",
    "    state_output(\"result_state\", state)\n",
    "    discard_array(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "190eb63a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppyalgos.tests.helpers import get_statevector, get_unitary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d9f8eff1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.37096824-0.00000000e+00j, -0.2972016 +1.81983495e-17j,\n",
       "       -0.2972016 +1.81983495e-17j,  0.23810338-2.91592544e-17j,\n",
       "       -0.2972016 +1.81983495e-17j,  0.23810338-2.91592544e-17j,\n",
       "        0.23810338-2.91592544e-17j, -0.19075678+3.50414518e-17j,\n",
       "       -0.2972016 +1.81983495e-17j,  0.23810338-2.91592544e-17j,\n",
       "        0.23810338-2.91592544e-17j, -0.19075678+3.50414518e-17j,\n",
       "        0.23810338-2.91592544e-17j, -0.19075678+3.50414518e-17j,\n",
       "       -0.19075678+3.50414518e-17j,  0.152825  -3.74313289e-17j])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_statevector(main_sv, n_qubits)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5df08c9a",
   "metadata": {},
   "source": [
    "## Extract a unitary matrix\n",
    "\n",
    "Use `get_unitary` to check a whole operation. Each column is the output for one computational-basis input:\n",
    "\n",
    "$$\n",
    "U_{yx} = \\langle y|U|x\\rangle.\n",
    "$$\n",
    "\n",
    "- Supply a circuit that accepts one qubit array and leaves it available to its caller.\n",
    "- The helper prepares the basis inputs and records the outputs for you; no `state_output` call is needed.\n",
    "- Compare complex matrix entries when relative phases matter. Measurement probabilities alone cannot detect every phase error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "af072aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def main_unitary[n_s: nat](state: array[qubit, n_s]) -> None:\n",
    "\n",
    "    for i in range(n_qubits):\n",
    "        ry(state[i], angle(3.14/2))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d088dd51",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.37096824+0.00000000e+00j,  0.2972016 +2.27152533e-17j,\n",
       "         0.2972016 +2.27152533e-17j,  0.23810338+3.63966991e-17j,\n",
       "         0.2972016 +2.27152533e-17j,  0.23810338+3.63966991e-17j,\n",
       "         0.23810338+3.63966991e-17j,  0.19075678+4.37388816e-17j,\n",
       "         0.2972016 +2.27152533e-17j,  0.23810338+3.63966991e-17j,\n",
       "         0.23810338+3.63966991e-17j,  0.19075678+4.37388816e-17j,\n",
       "         0.23810338+3.63966991e-17j,  0.19075678+4.37388816e-17j,\n",
       "         0.19075678+4.37388816e-17j,  0.152825  +4.67219358e-17j],\n",
       "       [-0.2972016 +1.81983495e-17j,  0.37096824-4.09136029e-17j,\n",
       "        -0.23810338-3.61872232e-18j,  0.2972016 -1.00627234e-17j,\n",
       "        -0.23810338-3.61872232e-18j,  0.2972016 -1.00627234e-17j,\n",
       "        -0.19075678-1.74787705e-17j,  0.23810338+1.01365879e-17j,\n",
       "        -0.23810338-3.61872232e-18j,  0.2972016 -1.00627234e-17j,\n",
       "        -0.19075678-1.74787705e-17j,  0.23810338+1.01365879e-17j,\n",
       "        -0.19075678-1.74787705e-17j,  0.23810338+1.01365879e-17j,\n",
       "        -0.152825  -2.56836196e-17j,  0.19075678+2.27005655e-17j],\n",
       "       [-0.2972016 +1.81983495e-17j, -0.23810338-3.61872232e-18j,\n",
       "         0.37096824-4.09136029e-17j,  0.2972016 -1.00627234e-17j,\n",
       "        -0.23810338-3.61872232e-18j, -0.19075678-1.74787705e-17j,\n",
       "         0.2972016 -1.00627234e-17j,  0.23810338+1.01365879e-17j,\n",
       "        -0.23810338-3.61872232e-18j, -0.19075678-1.74787705e-17j,\n",
       "         0.2972016 -1.00627234e-17j,  0.23810338+1.01365879e-17j,\n",
       "        -0.19075678-1.74787705e-17j, -0.152825  -2.56836196e-17j,\n",
       "         0.23810338+1.01365879e-17j,  0.19075678+2.27005655e-17j],\n",
       "       [ 0.23810338-2.91592544e-17j, -0.2972016 +5.09763263e-17j,\n",
       "        -0.2972016 +5.09763263e-17j,  0.37096824-8.18272057e-17j,\n",
       "         0.19075678-8.78134068e-18j, -0.23810338+2.26413888e-17j,\n",
       "        -0.23810338+2.26413888e-17j,  0.2972016 -4.28407001e-17j,\n",
       "         0.19075678-8.78134068e-18j, -0.23810338+2.26413888e-17j,\n",
       "        -0.23810338+2.26413888e-17j,  0.2972016 -4.28407001e-17j,\n",
       "         0.152825  +4.64530345e-18j, -0.19075678+3.55954569e-18j,\n",
       "        -0.19075678+3.55954569e-18j,  0.23810338-1.61235232e-17j],\n",
       "       [-0.2972016 +1.81983495e-17j, -0.23810338-3.61872232e-18j,\n",
       "        -0.23810338-3.61872232e-18j, -0.19075678-1.74787705e-17j,\n",
       "         0.37096824-4.09136029e-17j,  0.2972016 -1.00627234e-17j,\n",
       "         0.2972016 -1.00627234e-17j,  0.23810338+1.01365879e-17j,\n",
       "        -0.23810338-3.61872232e-18j, -0.19075678-1.74787705e-17j,\n",
       "        -0.19075678-1.74787705e-17j, -0.152825  -2.56836196e-17j,\n",
       "         0.2972016 -1.00627234e-17j,  0.23810338+1.01365879e-17j,\n",
       "         0.23810338+1.01365879e-17j,  0.19075678+2.27005655e-17j],\n",
       "       [ 0.23810338-2.91592544e-17j, -0.2972016 +5.09763263e-17j,\n",
       "         0.19075678-8.78134068e-18j, -0.23810338+2.26413888e-17j,\n",
       "        -0.2972016 +5.09763263e-17j,  0.37096824-8.18272057e-17j,\n",
       "        -0.23810338+2.26413888e-17j,  0.2972016 -4.28407001e-17j,\n",
       "         0.19075678-8.78134068e-18j, -0.23810338+2.26413888e-17j,\n",
       "         0.152825  +4.64530345e-18j, -0.19075678+3.55954569e-18j,\n",
       "        -0.23810338+2.26413888e-17j,  0.2972016 -4.28407001e-17j,\n",
       "        -0.19075678+3.55954569e-18j,  0.23810338-1.61235232e-17j],\n",
       "       [ 0.23810338-2.91592544e-17j,  0.19075678-8.78134068e-18j,\n",
       "        -0.2972016 +5.09763263e-17j, -0.23810338+2.26413888e-17j,\n",
       "        -0.2972016 +5.09763263e-17j, -0.23810338+2.26413888e-17j,\n",
       "         0.37096824-8.18272057e-17j,  0.2972016 -4.28407001e-17j,\n",
       "         0.19075678-8.78134068e-18j,  0.152825  +4.64530345e-18j,\n",
       "        -0.23810338+2.26413888e-17j, -0.19075678+3.55954569e-18j,\n",
       "        -0.23810338+2.26413888e-17j, -0.19075678+3.55954569e-18j,\n",
       "         0.2972016 -4.28407001e-17j,  0.23810338-1.61235232e-17j],\n",
       "       [-0.19075678+3.50414518e-17j,  0.23810338-5.54193656e-17j,\n",
       "         0.23810338-5.54193656e-17j, -0.2972016 +8.37543030e-17j,\n",
       "         0.23810338-5.54193656e-17j, -0.2972016 +8.37543030e-17j,\n",
       "        -0.2972016 +8.37543030e-17j,  0.37096824-1.22740809e-16j,\n",
       "        -0.152825  +1.63930127e-17j,  0.19075678-2.98196568e-17j,\n",
       "         0.19075678-2.98196568e-17j, -0.23810338+4.89015000e-17j,\n",
       "         0.19075678-2.98196568e-17j, -0.23810338+4.89015000e-17j,\n",
       "        -0.23810338+4.89015000e-17j,  0.2972016 -7.56186769e-17j],\n",
       "       [-0.2972016 +1.81983495e-17j, -0.23810338-3.61872232e-18j,\n",
       "        -0.23810338-3.61872232e-18j, -0.19075678-1.74787705e-17j,\n",
       "        -0.23810338-3.61872232e-18j, -0.19075678-1.74787705e-17j,\n",
       "        -0.19075678-1.74787705e-17j, -0.152825  -2.56836196e-17j,\n",
       "         0.37096824-4.09136029e-17j,  0.2972016 -1.00627234e-17j,\n",
       "         0.2972016 -1.00627234e-17j,  0.23810338+1.01365879e-17j,\n",
       "         0.2972016 -1.00627234e-17j,  0.23810338+1.01365879e-17j,\n",
       "         0.23810338+1.01365879e-17j,  0.19075678+2.27005655e-17j],\n",
       "       [ 0.23810338-2.91592544e-17j, -0.2972016 +5.09763263e-17j,\n",
       "         0.19075678-8.78134068e-18j, -0.23810338+2.26413888e-17j,\n",
       "         0.19075678-8.78134068e-18j, -0.23810338+2.26413888e-17j,\n",
       "         0.152825  +4.64530345e-18j, -0.19075678+3.55954569e-18j,\n",
       "        -0.2972016 +5.09763263e-17j,  0.37096824-8.18272057e-17j,\n",
       "        -0.23810338+2.26413888e-17j,  0.2972016 -4.28407001e-17j,\n",
       "        -0.23810338+2.26413888e-17j,  0.2972016 -4.28407001e-17j,\n",
       "        -0.19075678+3.55954569e-18j,  0.23810338-1.61235232e-17j],\n",
       "       [ 0.23810338-2.91592544e-17j,  0.19075678-8.78134068e-18j,\n",
       "        -0.2972016 +5.09763263e-17j, -0.23810338+2.26413888e-17j,\n",
       "         0.19075678-8.78134068e-18j,  0.152825  +4.64530345e-18j,\n",
       "        -0.23810338+2.26413888e-17j, -0.19075678+3.55954569e-18j,\n",
       "        -0.2972016 +5.09763263e-17j, -0.23810338+2.26413888e-17j,\n",
       "         0.37096824-8.18272057e-17j,  0.2972016 -4.28407001e-17j,\n",
       "        -0.23810338+2.26413888e-17j, -0.19075678+3.55954569e-18j,\n",
       "         0.2972016 -4.28407001e-17j,  0.23810338-1.61235232e-17j],\n",
       "       [-0.19075678+3.50414518e-17j,  0.23810338-5.54193656e-17j,\n",
       "         0.23810338-5.54193656e-17j, -0.2972016 +8.37543030e-17j,\n",
       "        -0.152825  +1.63930127e-17j,  0.19075678-2.98196568e-17j,\n",
       "         0.19075678-2.98196568e-17j, -0.23810338+4.89015000e-17j,\n",
       "         0.23810338-5.54193656e-17j, -0.2972016 +8.37543030e-17j,\n",
       "        -0.2972016 +8.37543030e-17j,  0.37096824-1.22740809e-16j,\n",
       "         0.19075678-2.98196568e-17j, -0.23810338+4.89015000e-17j,\n",
       "        -0.23810338+4.89015000e-17j,  0.2972016 -7.56186769e-17j],\n",
       "       [ 0.23810338-2.91592544e-17j,  0.19075678-8.78134068e-18j,\n",
       "         0.19075678-8.78134068e-18j,  0.152825  +4.64530345e-18j,\n",
       "        -0.2972016 +5.09763263e-17j, -0.23810338+2.26413888e-17j,\n",
       "        -0.23810338+2.26413888e-17j, -0.19075678+3.55954569e-18j,\n",
       "        -0.2972016 +5.09763263e-17j, -0.23810338+2.26413888e-17j,\n",
       "        -0.23810338+2.26413888e-17j, -0.19075678+3.55954569e-18j,\n",
       "         0.37096824-8.18272057e-17j,  0.2972016 -4.28407001e-17j,\n",
       "         0.2972016 -4.28407001e-17j,  0.23810338-1.61235232e-17j],\n",
       "       [-0.19075678+3.50414518e-17j,  0.23810338-5.54193656e-17j,\n",
       "        -0.152825  +1.63930127e-17j,  0.19075678-2.98196568e-17j,\n",
       "         0.23810338-5.54193656e-17j, -0.2972016 +8.37543030e-17j,\n",
       "         0.19075678-2.98196568e-17j, -0.23810338+4.89015000e-17j,\n",
       "         0.23810338-5.54193656e-17j, -0.2972016 +8.37543030e-17j,\n",
       "         0.19075678-2.98196568e-17j, -0.23810338+4.89015000e-17j,\n",
       "        -0.2972016 +8.37543030e-17j,  0.37096824-1.22740809e-16j,\n",
       "        -0.23810338+4.89015000e-17j,  0.2972016 -7.56186769e-17j],\n",
       "       [-0.19075678+3.50414518e-17j, -0.152825  +1.63930127e-17j,\n",
       "         0.23810338-5.54193656e-17j,  0.19075678-2.98196568e-17j,\n",
       "         0.23810338-5.54193656e-17j,  0.19075678-2.98196568e-17j,\n",
       "        -0.2972016 +8.37543030e-17j, -0.23810338+4.89015000e-17j,\n",
       "         0.23810338-5.54193656e-17j,  0.19075678-2.98196568e-17j,\n",
       "        -0.2972016 +8.37543030e-17j, -0.23810338+4.89015000e-17j,\n",
       "        -0.2972016 +8.37543030e-17j, -0.23810338+4.89015000e-17j,\n",
       "         0.37096824-1.22740809e-16j,  0.2972016 -7.56186769e-17j],\n",
       "       [ 0.152825  -3.74313289e-17j, -0.19075678+5.60797680e-17j,\n",
       "        -0.19075678+5.60797680e-17j,  0.23810338-8.16794767e-17j,\n",
       "        -0.19075678+5.60797680e-17j,  0.23810338-8.16794767e-17j,\n",
       "         0.23810338-8.16794767e-17j, -0.2972016 +1.16532280e-16j,\n",
       "        -0.19075678+5.60797680e-17j,  0.23810338-8.16794767e-17j,\n",
       "         0.23810338-8.16794767e-17j, -0.2972016 +1.16532280e-16j,\n",
       "         0.23810338-8.16794767e-17j, -0.2972016 +1.16532280e-16j,\n",
       "        -0.2972016 +1.16532280e-16j,  0.37096824-1.63654411e-16j]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_unitary(main_unitary, n_qubits)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4bc9beb",
   "metadata": {},
   "source": [
    "### Allow for internal work qubits\n",
    "\n",
    "- Pass `n_extra_qubits` when the circuit allocates work qubits internally.\n",
    "- Return those work qubits to $|0\\rangle$ before discarding them, so they are disentangled from the logical register.\n",
    "\n",
    "The two CX gates below uncompute the work qubit. The remaining logical operation is X:\n",
    "\n",
    "$$\n",
    "X=\\begin{pmatrix}0&1\\\\1&0\\end{pmatrix}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e20a24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import no_type_check\n",
    "\n",
    "@guppy\n",
    "@no_type_check\n",
    "def unitary_with_work_qubit(state: array[qubit, 1]) -> None:\n",
    "    work = qubit()\n",
    "    cx(state[0], work)\n",
    "    cx(state[0], work)\n",
    "    x(state[0])\n",
    "    discard(work)\n",
    "\n",
    "unitary = get_unitary(\n",
    "    unitary_with_work_qubit,\n",
    "    n_qubits=1,\n",
    "    n_extra_qubits=1,\n",
    ")\n",
    "\n",
    "expected = np.array([[0, 1], [1, 0]])\n",
    "assert np.allclose(unitary / unitary[0, 1], expected)\n",
    "unitary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b0208c5",
   "metadata": {},
   "source": [
    "### Compare magnitudes when phase information is unavailable\n",
    "\n",
    "`get_unitary_assumed_phase` also accepts `n_extra_qubits`, but does not retain phase information. Compare `np.abs(...)` for this helper, as shown below. This checks matrix-entry magnitudes; it does not verify relative phases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83c19b72",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppyalgos.tests.helpers import get_unitary_assumed_phase\n",
    "\n",
    "unitary_assumed_phase = get_unitary_assumed_phase(\n",
    "    unitary_with_work_qubit,\n",
    "    n_qubits=1,\n",
    "    n_extra_qubits=1,\n",
    ")\n",
    "\n",
    "assert np.allclose(np.abs(unitary_assumed_phase), expected)\n",
    "unitary_assumed_phase"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6e1f9a3",
   "metadata": {},
   "source": [
    "## Post-select a state\n",
    "\n",
    "Post-selection keeps the branch associated with chosen measurement outcomes. For an ancilla outcome $b$, the remaining system state is\n",
    "\n",
    "$$\n",
    "|\\widetilde{\\psi}_b\\rangle\n",
    "= (\\langle b|_{\\mathrm{ancilla}}\\otimes I_{\\mathrm{system}})\n",
    "|\\Psi\\rangle.\n",
    "$$\n",
    "\n",
    "Its squared norm gives the outcome probability. When that probability is nonzero, normalization gives the conditional state:\n",
    "\n",
    "$$\n",
    "p_b=\\lVert\\widetilde{\\psi}_b\\rVert^2,\n",
    "\\qquad\n",
    "|\\psi_b\\rangle=\\frac{|\\widetilde{\\psi}_b\\rangle}{\\sqrt{p_b}}.\n",
    "$$\n",
    "\n",
    "- Tag the registers with `state_output` and use those names in the selection dictionary.\n",
    "- `False` selects $0$; `True` selects $1$.\n",
    "- `get_statevector_projected` returns the state of the unprojected qubits.\n",
    "- Use `renormalize=True` for the conditional state, or `False` to preserve the branch weight.\n",
    "\n",
    "This example selects $|000\\rangle$ on the ancilla register tagged `\"projection\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1b3c9daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_state_qubits = 2\n",
    "n_ancilla_qubits = 3\n",
    "n_total_qubits = n_state_qubits + n_ancilla_qubits\n",
    "@guppy\n",
    "def main_sv_project() -> None:\n",
    "    state = qarray(n_state_qubits)\n",
    "    ancilla = qarray(n_ancilla_qubits)\n",
    "\n",
    "    for i in range(n_state_qubits):\n",
    "        ry(state[i], angle(3.14/2))\n",
    "\n",
    "    for i in range(n_ancilla_qubits):\n",
    "        ry(ancilla[i], angle(3.14/2))\n",
    "\n",
    "    for i in range(n_ancilla_qubits):\n",
    "        for j in range(n_state_qubits):\n",
    "            cx(state[j], ancilla[i])\n",
    "\n",
    "    # projection qubits to be measured\n",
    "    state_output(\"projection\", ancilla)\n",
    "    state_output(\"system\", state)\n",
    "    \n",
    "    discard_array(state)\n",
    "    discard_array(ancilla)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fbd1e8bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.7556292 -5.07091231e-16j, 0.31129088+1.81104011e-17j,\n",
       "       0.31129088+5.43312033e-17j, 0.48499534+2.17324813e-16j])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from guppyalgos.tests.helpers import get_statevector_projected\n",
    "\n",
    "post_selection = {\"projection\": [False for _ in range(n_ancilla_qubits)]} # postselect on all ancillas being 0\n",
    "get_statevector_projected(main_sv_project, n_total_qubits, post_selection, renormalize=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1736b34",
   "metadata": {},
   "source": [
    "### Select outcomes on several registers\n",
    "\n",
    "A selection dictionary can include several named registers. The helper applies their projections in dictionary order, allowing you to test a joint success condition, such as the ancilla and phase registers both being zero in a QSVT routine.\n",
    "\n",
    "Record disjoint registers at the same point in the circuit, with no computation between their snapshots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6e97ccf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_state_qubits = 2\n",
    "n_phase_qubits = 3\n",
    "n_total_qubits = n_state_qubits + n_phase_qubits + 1\n",
    "state_preparation_angle = np.pi / 4\n",
    "\n",
    "@guppy\n",
    "def main() -> None:\n",
    "    phase_qreg = qarray(n_phase_qubits)\n",
    "    state_qreg = qarray(n_state_qubits)\n",
    "    ancilla_qreg = qarray(1)\n",
    "\n",
    "    # Non-eigenstate input in computational basis.\n",
    "    ry(state_qreg[0], angle(state_preparation_angle))\n",
    "    x(state_qreg[1])\n",
    "\n",
    "    for i in range(n_phase_qubits):\n",
    "        h(phase_qreg[i])\n",
    "\n",
    "    h(ancilla_qreg[0])\n",
    "    for i in range(n_phase_qubits):\n",
    "\n",
    "        cx(phase_qreg[i], ancilla_qreg[0])\n",
    "\n",
    "    for i in range(n_state_qubits):\n",
    "        cx(state_qreg[i], ancilla_qreg[0])  \n",
    "\n",
    "    h(ancilla_qreg[0])\n",
    "\n",
    "    state_output(\"phase\", phase_qreg)\n",
    "    state_output(\"ancilla\", ancilla_qreg)\n",
    "    state_output(\"system\", state_qreg)\n",
    "    discard_array(phase_qreg)\n",
    "    discard_array(state_qreg)\n",
    "    discard_array(ancilla_qreg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a74377d",
   "metadata": {},
   "source": [
    "Select $|0\\rangle$ on `\"ancilla\"` and $|000\\rangle$ on `\"phase\"`. The returned statevector describes the remaining system register."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "37293273",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2.02524568e-17-2.17207347e-32j,  5.77861354e-17-7.47527359e-33j,\n",
       "       -8.70182863e-17-3.30747719e-01j, -3.04173840e-16-9.43719209e-01j])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from guppyalgos.tests.helpers import get_statevector_projected, statevector_projected_selene\n",
    "import numpy as np\n",
    "\n",
    "postselect_dict = {\"ancilla\": [False], \"phase\": [False, False, False]}\n",
    "\n",
    "get_statevector_projected(main, n_total_qubits, postselect_dict)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74743953",
   "metadata": {},
   "source": [
    "### Keep a state object for further inspection\n",
    "\n",
    "Use `statevector_projected_selene` when you need a `SeleneQuestState` rather than a NumPy array. The `returned_specified_qubits` argument marks the qubit IDs of interest for subsequent projections or inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f6d2d66e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SubQuestState(state=array([ 2.02524568e-17-2.17207347e-32j,  5.77861354e-17-7.47527359e-33j,\n",
       "       -8.70182863e-17-3.30747719e-01j, -3.04173840e-16-9.43719209e-01j]), total_qubits=2, specified_qubits=[3, 4])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "statevector_projected_selene(main, n_total_qubits, postselect_dict, returned_specified_qubits=[3,4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ad31369",
   "metadata": {},
   "source": [
    "## Extract a post-selected operation\n",
    "\n",
    "Use `get_unitary_projected` to extract the operation associated with chosen initial and final ancilla states:\n",
    "\n",
    "$$\n",
    "A_{b,a}\n",
    "= (\\langle b|\\otimes I)\\,U\\,(|a\\rangle\\otimes I).\n",
    "$$\n",
    "\n",
    "This block need not be unitary. Its scale matters: for a normalized input $|\\psi\\rangle$, the probability of outcome $b$ is $\\lVert A_{b,a}|\\psi\\rangle\\rVert^2$.\n",
    "\n",
    "- Put projected registers first in the circuit signature, followed by the system register.\n",
    "- The selection dictionary's key order determines the projected-register argument order.\n",
    "- The helper records those registers for post-selection.\n",
    "- Use matching keys in the pre-selection and post-selection dictionaries; omit pre-selection to initialize all projected registers to zero.\n",
    "\n",
    "Below, the ancillas start in $|111\\rangle$ and are projected onto $|000\\rangle$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f2a47198",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def main_unitary_projected[n_a: nat, n_s: nat](ancilla: array[qubit, n_a], state: array[qubit, n_s]) -> None:\n",
    "    for i in range(len(state)):\n",
    "        ry(state[i], angle(3.14/2))\n",
    "\n",
    "    for i in range(len(ancilla)):\n",
    "        ry(ancilla[i], angle(3.14/2))\n",
    "\n",
    "    for i in range(len(state)):\n",
    "        for j in range(len(ancilla)):\n",
    "            cx(state[i], ancilla[j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aae33b36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 6.93889390e-17+0.14887239j,  1.42420307e-17+0.11926928j,\n",
       "         1.42420307e-17+0.11926928j,  3.94574074e-17+0.09555271j],\n",
       "       [-1.80411242e-16-0.23194517j,  1.28749977e-16+0.28951489j,\n",
       "        -6.68895356e-17-0.18582312j,  1.39427393e-16+0.23194517j],\n",
       "       [-1.66533454e-16-0.23194517j, -6.68895356e-17-0.18582312j,\n",
       "         2.67527855e-16+0.28951489j,  1.53305181e-16+0.23194517j],\n",
       "       [-6.93889390e-18+0.09555271j, -7.30313681e-18-0.11926928j,\n",
       "        -7.30313681e-18-0.11926928j, -3.72795415e-17+0.14887239j]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from guppyalgos.tests.helpers import get_unitary_projected\n",
    "\n",
    "post_selection = {\"ancilla\": [False for _ in range(n_ancilla_qubits)]}\n",
    "pre_selection = {\"ancilla\": [True for _ in range(n_ancilla_qubits)]}\n",
    "get_unitary_projected(main_unitary_projected, n_state_qubits, post_selection, pre_selection)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c8bcc40",
   "metadata": {},
   "source": [
    "### Allow for internal work qubits\n",
    "\n",
    "Pass `n_extra_qubits` for work qubits allocated inside the circuit, beyond its explicit register arguments. The helper reserves that additional emulator capacity while extracting the logical operation.\n",
    "\n",
    "As with `get_unitary`, internal work qubits must return to $|0\\rangle$ before being discarded. The example below leaves the projected register unchanged and applies X to the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cde9fd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import no_type_check\n",
    "from guppyalgos.tests.helpers import get_unitary_projected\n",
    "\n",
    "@guppy\n",
    "@no_type_check\n",
    "def projected_unitary_with_work_qubit(\n",
    "    projection: array[qubit, 1],\n",
    "    state: array[qubit, 1],\n",
    ") -> None:\n",
    "    work = qubit()\n",
    "    cx(state[0], work)\n",
    "    cx(state[0], work)\n",
    "    x(state[0])\n",
    "    discard(work)\n",
    "\n",
    "projected = get_unitary_projected(\n",
    "    projected_unitary_with_work_qubit,\n",
    "    n_state_qubits=1,\n",
    "    post_select_dict={\"projection\": [False]},\n",
    "    n_extra_qubits=1,\n",
    ")\n",
    "\n",
    "expected = np.array([[0, 1], [1, 0]])\n",
    "assert np.allclose(projected / projected[0, 1], expected)\n",
    "projected"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3448344d",
   "metadata": {},
   "source": [
    "## Inspect only selected registers\n",
    "\n",
    "Use `get_total_state_on_only_specified_registers` to inspect a collection of tagged registers while excluding unused ancillas.\n",
    "\n",
    "- Identify the registers by their `state_output` tags.\n",
    "- Record the snapshots together, with no computation between them.\n",
    "- The selected registers must be disentangled from everything excluded. They may be entangled with each other.\n",
    "- The helper returns a state object and a mapping from register tags to qubit IDs.\n",
    "\n",
    "Here, the control and target registers contain six qubits altogether, so their state has $2^6=64$ amplitudes. The separate ancilla is excluded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1b57ece8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n"
     ]
    }
   ],
   "source": [
    "from guppyalgos.utils import transversal\n",
    "from selene_sim import Quest\n",
    "from guppyalgos.tests.helpers import get_total_state_on_only_specified_registers\n",
    "\n",
    "@guppy\n",
    "def main_get_total_state_on_only_specified_registers() -> None:\n",
    "    control_reg = qarray(3)\n",
    "    target_reg = qarray(3)\n",
    "    ancilla = qubit()\n",
    "    x(control_reg[1])\n",
    "    transversal(cx, control_reg, target_reg)\n",
    "    # state results for desired registers\n",
    "    state_output(\"control\", control_reg)\n",
    "    state_output(\"target\", target_reg)\n",
    "    discard_array(control_reg)\n",
    "    discard_array(target_reg)\n",
    "    discard(ancilla)\n",
    "\n",
    "res = main_get_total_state_on_only_specified_registers.emulator(7).run()\n",
    "states = Quest.extract_states_dict(res.results[0].entries)\n",
    "\n",
    "# returns the state and a dictionary mapping from the tags to the qubit ids for each register\n",
    "non_ancilla_state, spec_qubit_dict = get_total_state_on_only_specified_registers(\n",
    "    states, [\"target\", \"control\"]\n",
    ")\n",
    "control_qubits = spec_qubit_dict[\"control\"]\n",
    "target_qubits = spec_qubit_dict[\"target\"]\n",
    "\n",
    "# length of non_ancilla_state will exclude the ancilla, so be 2**6=64,\n",
    "print(len(non_ancilla_state.state)) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "855dc9a3",
   "metadata": {},
   "source": [
    "## Replay measurement outcomes\n",
    "\n",
    "Selene's `QuantumReplay` lets you specify measurement outcomes in advance. This makes tests of mid-circuit measurements, feedforward, and retries reproducible.\n",
    "\n",
    "The example below uses rotation resource states until a success flag is measured. It checks two paths:\n",
    "\n",
    "- `[True]`: success on the first attempt.\n",
    "- `[False, False, True]`: two failures followed by success.\n",
    "\n",
    "After each failure, the correction angle doubles; attempt $k$ uses $2^{k-1}\\theta$. Comparing the final rotation across paths checks that the corrections produce the intended result.\n",
    "\n",
    "Requested outcomes must have non-negligible probability in the simulated state. Replay selects a possible branch rather than assigning an arbitrary quantum state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c16bd85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.10000000000000005\n",
      "0.10000000000000002\n"
     ]
    }
   ],
   "source": [
    "from selene_sim import QuantumReplay\n",
    "from guppylang.std.quantum import measure, h, rz, s, sdg\n",
    "import numpy as np\n",
    "\n",
    "@guppy\n",
    "def repeat_until_success_ry_example() -> None:\n",
    "    \"\"\"Perform repeat-until-success Ry rotation.\"\"\"\n",
    "    q = qubit()\n",
    "    sdg(q)\n",
    "    h(q)\n",
    "    theta = angle(0.1)\n",
    "    attempts = 0\n",
    "    while True:\n",
    "        attempts += 1\n",
    "        a = qubit()\n",
    "        h(a)\n",
    "        rz(a, theta * (2 ** (attempts - 1)))\n",
    "        cx(q, a)\n",
    "        if not measure(a).read():\n",
    "            continue\n",
    "        break\n",
    "    h(q)\n",
    "    s(q)\n",
    "    state_output(\"result_state\",q)\n",
    "    discard(q)\n",
    "\n",
    "# set up the replay plugin with desired measurements and Quest as a backend\n",
    "n_shots = 2\n",
    "rus_replay_sim = QuantumReplay(simulator=Quest(), measurements=[[True], [False, False, True]])\n",
    "\n",
    "\n",
    "em_result = (\n",
    "        repeat_until_success_ry_example.emulator(2).with_simulator(rus_replay_sim).with_shots(n_shots).run()\n",
    "    )\n",
    "# print the rotation angle applied in each case\n",
    "for shot_result in em_result.results:\n",
    "    states = Quest.extract_states_dict(shot_result)\n",
    "    sv = states[\"result_state\"].state\n",
    "    res = abs(sv[3]) ** 2\n",
    "    theta = 2 * np.arcsin(np.sqrt(res)) / np.pi\n",
    "    print(theta)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.13.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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}