{
    "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.testing 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.testing 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": "modifier-guide",
            "metadata": {},
            "source": [
                "### Test custom unitary modifiers\n",
                "\n",
                "Use `assert_unitary_modifiers` for a circuit that accepts a single qubit array and provides custom modifier implementations. One invocation checks:\n",
                "\n",
                "- Forward: the extracted matrix is unitary and matches the supplied reference up to one overall global phase.\n",
                "- Dagger: the matrix is unitary and matches the forward adjoint.\n",
                "- Control and controlled dagger: check the full matrix including the control qubit against identity on the inactive branch and the reference or its adjoint on the active branch, preserving relative phase.\n",
                "\n",
                "The rotation below defines all four implementations explicitly. Automatically generated modifiers are the compiler's testing responsibility.\n"
            ]
        },
        {
            "cell_type": "code",
            "id": "custom-modifier-example",
            "metadata": {},
            "source": [
                "from guppylang.std.builtins import control\n",
                "from guppylang.std.quantum import rz\n",
                "from guppyalgos.testing import assert_unitary_modifiers\n",
                "\n",
                "\n",
                "@guppy.unitary\n",
                "class custom_rotation:\n",
                "    @guppy\n",
                "    def __call__(qs: array[qubit, 1]) -> None:\n",
                "        rz(qs[0], angle(0.25))\n",
                "\n",
                "    @guppy\n",
                "    def daggered(qs: array[qubit, 1]) -> None:\n",
                "        rz(qs[0], angle(-0.25))\n",
                "\n",
                "    @guppy\n",
                "    def controlled[n: nat](qs: array[qubit, 1], controls: array[qubit, n]) -> None:\n",
                "        with control(controls):\n",
                "            rz(qs[0], angle(0.25))\n",
                "\n",
                "    @guppy\n",
                "    def ctrl_daggered[n: nat](qs: array[qubit, 1], controls: array[qubit, n]) -> None:\n",
                "        with control(controls):\n",
                "            rz(qs[0], angle(-0.25))\n",
                "\n",
                "\n",
                "# angle(0.25) is a quarter half-turn: pi/4 radians.\n",
                "expected_rotation = np.diag(np.exp(1j * np.array([-np.pi / 8, np.pi / 8])))\n",
                "assert_unitary_modifiers(custom_rotation, 1, expected_unitary=expected_rotation)\n"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "id": "modifier-testing-limits",
            "metadata": {},
            "source": [
                "A successful assertion returns without output. Failures identify the mode and distinguish a nonunitary extracted matrix from an incorrect unitary operation. The supplied independent reference checks the full forward unitary up to global phase and preserves the intended relative phase in full controlled-matrix checks. The reference is required so compiler-dependent forward phases cannot determine the controlled expectation.\n",
                "\n",
                "Use this helper for small systems because full-matrix extraction scales exponentially. Pass `n_extra_qubits` for internal work qubits. Reset is compatible when those qubits are correctly uncomputed and disentangled. Measurement-based cleanup also needs separate superposition and measurement-outcome checks: basis-input runs can hide lost coherence or combine outcome-dependent phases into inconsistent column phases. Specialized tests can be added case by case.\n"
            ]
        },
        {
            "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.testing 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.testing 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.testing 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.testing 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.testing 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
}