random_walk_qpe.py

"""
Implementation of the adaptive random walk phase estimation algorithm from
https://arxiv.org/abs/2208.04526.

The example Hamiltonian and numbers are taken from https://arxiv.org/abs/2206.12950.
"""

import math
from collections.abc import Callable

from guppylang.decorator import guppy
from guppylang.module import GuppyModule
from guppylang.std.angles import angle
from guppylang.std.builtins import owned, py, result
from guppylang.std.quantum import discard, measure, qubit
from guppylang.std.quantum_functional import cx, h, rz, x


sqrt_e = math.sqrt(math.e)
sqrt_e_div = math.sqrt((math.e - 1) / math.e)


@guppy
def random_walk_phase_estimation(
    eigenstate: Callable[[], qubit],
    controlled_oracle: Callable[[qubit @owned, qubit @owned, float], tuple[qubit, qubit]],
    num_iters: int,
    reset_rate: int,
    mu: float,
    sigma: float,
) -> float:
    """Performs the random walk phase estimation algorithm on a single qubit for
    some Hamiltonian H.

    Arguments:
       eigenstate: Function preparing the eigenstate of e^itH
       controlled_oracle: The oracle circuit for a controlled e^itH
       num_iters: Number of iterations to run the algorithm for
       reset_rate: Reset the eigenstate every x iterations
       mu: Initial mean for the eigenvalue estimate
       sigma: Initial standard deviation for the eigenvalue estimate
    """
    tgt = eigenstate()
    i = 0
    while i < num_iters:
        aux = h(qubit())
        t = 1 / sigma
        aux = rz(h(aux), angle((sigma - mu) * t))
        aux, tgt = controlled_oracle(aux, tgt, t)
        if measure(h(aux)):
            mu += sigma / py(sqrt_e)
        else:
            mu -= sigma / py(sqrt_e)
        sigma *= py(sqrt_e_div)

        # Reset the eigenstate every few iterations to increase the fidelity of
        # the algorithm
        if i % reset_rate == 0:
            discard(tgt)
            tgt = eigenstate()
        i += 1
    discard(tgt)
    return mu


@guppy
def example_controlled_oracle(q1: qubit @owned, q2: qubit @owned, t: float) -> tuple[qubit, qubit]:
    """A controlled e^itH gate for the example Hamiltonian H = -0.5 * Z"""
    # This is just a controlled rz gate
    a = angle(-0.5 * t)
    q2 = rz(q2, a / 2)
    q1, q2 = cx(q1, q2)
    q2 = rz(q2, -a / 2)
    return cx(q1, q2)


@guppy
def example_eigenstate() -> qubit:
    """The eigenstate of e^itH for the example Hamiltonian H = -0.5 * Z"""
    # This is just |1>
    return x(qubit())


@guppy
def main() -> int:
    num_iters = 24  # To avoid underflows
    reset_rate = 8
    mu = py(sqrt_e)
    sigma = py(sqrt_e_div)
    eigenvalue = random_walk_phase_estimation(
        example_eigenstate,
        example_controlled_oracle,
        num_iters,
        reset_rate,
        mu,
        sigma,
    )
    result("eigenvalue", eigenvalue)  # Expected outcome is 0.5
    return 0


hugr = guppy.compile_module()