{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0606569a",
   "metadata": {},
   "source": [
    "# Canonical phase estimation\n",
    "**Download Notebook** - {nb-download}`canonical_phase_estimation.ipynb`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b817628a",
   "metadata": {},
   "source": [
    "- 'Canonical' phase estimation reads out the bitstring representation of eigenphases of a unitary matrix by repeated application of the controlled unitary onto a phase register. \n",
    "- It relies on the use of the inverse Quantum Fourier Transform to transform a cosine wave in the time-domain into a delta function in the frequency (phase) domain.\n",
    "- To prepare the signal, binary powers (U^2^1, U^2^2, etc) of the unitary are used to express a phase kickback, which is entangled with the system register. \n",
    "- The controlled unitary here is defined with a guppy function, with the system and ancilla registers and unitary power as its arguments. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6add312c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang.decorator import guppy\n",
    "from guppylang.std.builtins import array, output, nat\n",
    "from guppylang.std.quantum import (\n",
    "    measure_array,\n",
    "    qubit,\n",
    "    discard_array,\n",
    "    ry,\n",
    "    crz,\n",
    "    collect_measurements\n",
    ")\n",
    "from guppylang.std.builtins import comptime\n",
    "from guppylang.std.angles import angle, pi\n",
    "from guppyalgos.utils import qarray\n",
    "\n",
    "from guppyalgos.algorithms.phase_estimation import qpe\n",
    "\n",
    "from guppyalgos.primitives.state_preparation import uniform_state\n",
    "\n",
    "from guppyalgos.utils import binary_fraction\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd931433",
   "metadata": {},
   "source": [
    "Here we prepare a toy system with Rz as the Hamiltonian. The powers of the controlled-Rz encode phi onto the phase register. To mirror the more general `UnitaryRegs` pattern, we wrap the system register in a small `RzUnitaryRegs` struct and let the oracle access `unitary_regs.system`. Linear combinations of the $\\ket{0}$ and $\\ket{1}$ eigenstates of this unitary can be prepared on the system register with an Ry gate. Their eigenphases are -phi and phi respectively. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6f1236e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_ancilla = 4\n",
    "n_system = 1\n",
    "phi = 0.75\n",
    "\n",
    "\n",
    "@guppy.struct\n",
    "class RzUnitaryRegs[n_q_s: nat]:\n",
    "    \"\"\"Wrap the system register for the unitary application.\n",
    "\n",
    "    Using a struct here keeps the QPE interface compatible with more complex\n",
    "    unitaries that may need multiple registers, not just a single system array.\n",
    "\n",
    "    \"\"\"\n",
    "    system: array[qubit, n_q_s]\n",
    "\n",
    "@guppy\n",
    "def power_oracle[n_q_s: nat](ctrl: qubit, unitary_regs: RzUnitaryRegs[n_q_s], power: int) -> None:\n",
    "    \"\"\"Use CRz gate to implement controlled e^-iπHt/2 for H = 0.5 * Z.\n",
    "\n",
    "    Args:\n",
    "        ctrl: Control qubit.\n",
    "        unitary_regs: Registers acted on by the powered unitary.\n",
    "        power: Exponent power.\n",
    "\n",
    "    Returns:\n",
    "        None\n",
    "\n",
    "    Note: This oracle is parametrerized by the power.\n",
    "\n",
    "    \"\"\"\n",
    "    crz(ctrl, unitary_regs.system[0], -2 * pi * phi * power)\n",
    "\n",
    "\n",
    "@guppy\n",
    "def state_preparation[n_q_s: nat](system_register: array[qubit, n_q_s]) -> None:\n",
    "    \"\"\"Uses Ry to create arbitrary linear combination of the simple oracle eigenstates.\"\"\"\n",
    "    # a value of 1 yields the |0> state with eigenphase -phi\n",
    "    ry(system_register[0], angle(0))\n",
    "\n",
    "\n",
    "# prebuild guppy functions for generic arrays\n",
    "ancilla_prep_function = uniform_state(2**n_ancilla)\n",
    "\n",
    "@guppy\n",
    "def abstract_canonical_qpe() -> None:\n",
    "    \"\"\"Abstract canonical QPE algorithm using functions for state prep and oracle.\"\"\"\n",
    "    # The unitary application can use a richer regs object than a bare qubit array.\n",
    "    unitary_regs = RzUnitaryRegs(qarray(n_system))\n",
    "    phase_reg = qarray(n_ancilla)\n",
    "    state_preparation(unitary_regs.system)\n",
    "    ancilla_prep_function(phase_reg)\n",
    "    qpe(phase_reg, unitary_regs, power_oracle)\n",
    "\n",
    "    output(\"qpe_bitstring\", collect_measurements(measure_array(phase_reg)))\n",
    "    discard_array(unitary_regs.system)\n",
    "\n",
    "\n",
    "#abstract_canonical_qpe.compile_function()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f53e12f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({'0110': 500})\n",
      "Target phase: 0.75\n",
      "Measured bitstring: 0110, counts: 500,              phase: 0.75,                  empirical probability: 1.0\n"
     ]
    }
   ],
   "source": [
    "n_shots = 500\n",
    "\n",
    "sim_result = (\n",
    "    abstract_canonical_qpe.emulator(n_qubits=n_ancilla + 1)\n",
    "    .with_seed(5)\n",
    "    .with_shots(n_shots)\n",
    "    .run()\n",
    ")\n",
    "result_counter = sim_result.register_counts()[\"qpe_bitstring\"]\n",
    "print(result_counter)\n",
    "\n",
    "print(\"Target phase:\", phi)\n",
    "for key, value in result_counter.items():\n",
    "    print(\n",
    "        f\"Measured bitstring: {key}, counts: {value},\\\n",
    "              phase: {binary_fraction([int(bit) for bit in key])},\\\n",
    "                  empirical probability: {value / n_shots}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44890287",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({'0110': 500})\n",
      "Target phase: 0.75\n",
      "Measured bitstring: 0110, counts: 500,              phase: 0.75,                  empirical probability: 1.0\n"
     ]
    }
   ],
   "source": [
    "# If the unitary only acts on a single qubit array, you can still pass that\n",
    "# array directly into qpe without defining a new regs struct.\n",
    "# however, due to typing, we do have to redefine the unitary\n",
    "\n",
    "@guppy\n",
    "def power_oracle_array[n_q_s: nat](ctrl: qubit, system_register: array[qubit, n_q_s], power: int) -> None:\n",
    "    \"\"\"Bare-array version of the same controlled Rz oracle.\"\"\"\n",
    "    crz(ctrl, system_register[0], -2 * pi * phi * power)\n",
    "\n",
    "\n",
    "@guppy\n",
    "def abstract_canonical_qpe_single_array() -> None:\n",
    "    \"\"\"Canonical QPE also works when the unitary registers are just an array.\"\"\"\n",
    "    system_register = qarray(n_system)\n",
    "    phase_reg = qarray(n_ancilla)\n",
    "    state_preparation(system_register)\n",
    "    ancilla_prep_function(phase_reg)\n",
    "    qpe(phase_reg, system_register, power_oracle_array)\n",
    "\n",
    "    output(\"qpe_bitstring_single_array\", collect_measurements(measure_array(phase_reg)))\n",
    "    discard_array(system_register)\n",
    "\n",
    "\n",
    "n_shots = 500\n",
    "\n",
    "sim_result = (\n",
    "    abstract_canonical_qpe_single_array.emulator(n_qubits=n_ancilla + 1)\n",
    "    .with_seed(5)\n",
    "    .with_shots(n_shots)\n",
    "    .run()\n",
    ")\n",
    "result_counter = sim_result.register_counts()[\"qpe_bitstring_single_array\"]\n",
    "print(result_counter)\n",
    "\n",
    "print(\"Target phase:\", phi)\n",
    "for key, value in result_counter.items():\n",
    "    print(\n",
    "        f\"Measured bitstring: {key}, counts: {value},\\\n",
    "              phase: {binary_fraction([int(bit) for bit in key])},\\\n",
    "                  empirical probability: {value / n_shots}\"\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "guppyalgos (3.13.3.final.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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}