QROM Phase-Gradient Rotations

Download Notebook - qrom_phase_gradient_rotation.ipynb

This notebook demonstrates how to use QROMRotations together with the RotationPhaseGradient Guppy struct to implement a table of rotations whose angles are selected by a QROM index register.

Motivation

Phase-gradient rotations are useful when a circuit needs to apply many different rotations selected by an index register, for example in QROM-based state preparation and basis rotations.

The key idea is that we do not synthesize a separate arbitrary-angle rotation for every table entry. Instead, we:

  1. prepare a phase-gradient register once,

  2. use QROM to write the selected fixed-point bit string into a binary target register,

  3. use controlled addition into the phase-gradient register to kick the phase back onto an ancilla, and

  4. uncompute the QROM target again.

This keeps the rotation synthesis cost as a one-time cost in the reusable phase-gradient preparation, while each indexed rotation query only uses QROM compute/uncompute plus the phase-gradient kickback primitive.

The same overall QROM workflow also works with register-incremented rotations: the QROM still loads a little-endian fixed-point word into a data register, but the rotation box interprets that word directly as a sequence of controlled per-bit rotations instead of adding it into a phase-gradient register.

In this notebook we use a concrete 16-bit example: the QROM has 16 entries, and entry j stores the 16-bit little-endian encoding of the increment j. So the table acts like a compact library of indexed fixed-point rotations.

from math import atan2, ceil, log2
from typing import Any, no_type_check

import numpy as np
import pandas as pd
from guppylang import comptime, guppy
from guppylang.std.builtins import array
from guppylang.std.debug import state_output
from guppylang.std.quantum import discard, discard_array, h, qubit
from selene_sim import Quest

from guppyalgos.algorithms.select.qrom import qrom_unary_iteration
from guppyalgos.primitives.rotations import (
    RotationAxisZ,
    RotationPhaseGradient,
)
from guppyalgos.primitives.state_preparation.phase_gradient import Convention, phase_gradient
from guppyalgos.utils import (
    int_to_bits,
    bits_to_int,
    phase_distance_mod_2,
    qarray,
)
from guppyalgos.tests.helpers import (
    assert_allclose_ignorephase,
    extract_state_branches_in_superposition,
)

What the QROM phase-gradient rotation does

Let the QROM store little-endian fixed-point integers

\[ x^{(j)} = \sum_{k=0}^{d-1} 2^k x_k^{(j)}, \qquad x_k^{(j)} \in \{0,1\}, \]

where \(j\) is the QROM index and \(d\) is the number of phase bits. If the QROM target is initialized to \(|0^d\rangle\), then the QROM compute step prepares

\[ \operatorname{QROM}\, |j\rangle |0^d\rangle = |j\rangle |x^{(j)}\rangle. \]

Now let \(|F_d\rangle\) be the standard little-endian phase-gradient state

\[ |F_d\rangle = \frac{1}{\sqrt{2^d}} \sum_{y=0}^{2^d-1} e^{-2\pi i y / 2^d} |y\rangle. \]

Controlled addition of \(x^{(j)}\) into \(|F_d\rangle\) kicks back a positive phase onto the ancilla because RotationPhaseGradient flips the controlled addition with X gates on the target:

\[ |x^{(j)}\rangle |F_d\rangle |\psi\rangle \mapsto |x^{(j)}\rangle |F_d\rangle R_z\!\left(\theta_j\right)|\psi\rangle, \]

with half-turn rotation parameter

\[ \theta_j^{\mathrm{kick}} = \frac{2 x^{(j)}}{2^d}. \]

After the final QROM uncompute, the net indexed operation is

\[ |j\rangle |0^d\rangle |F_d\rangle |\psi\rangle \mapsto |j\rangle |0^d\rangle |F_d\rangle R_z\!\left(\theta_j\right)|\psi\rangle. \]

So on each QROM index \(j\), the circuit applies the positive kickback rotation determined by the fixed-point word stored at that index, while restoring the QROM target register back to zero.

n_phase_qubits = 4
n_index_elements = 2 ** n_phase_qubits

data_input = [
    int_to_bits(index, n_phase_qubits)
    for index in range(n_index_elements)
]
n_index_qubits = ceil(log2(len(data_input)))

def kickback_theta(bits: list[bool]) -> float:
    encoded_integer = bits_to_int(bits)
    return float(2.0 * encoded_integer / (2 ** len(bits)))


pd.DataFrame(
    {
        "index j": list(range(len(data_input))),
        "stored bits x^(j)": data_input,
        "encoded integer x^(j)": [bits_to_int(bits) for bits in data_input],
        "kickback angle 2 x^(j) / 2^d": [kickback_theta(bits) for bits in data_input],
    }
)
index j stored bits x^(j) encoded integer x^(j) kickback angle 2 x^(j) / 2^d
0 0 [False, False, False, False] 0 0.000
1 1 [True, False, False, False] 1 0.125
2 2 [False, True, False, False] 2 0.250
3 3 [True, True, False, False] 3 0.375
4 4 [False, False, True, False] 4 0.500
5 5 [True, False, True, False] 5 0.625
6 6 [False, True, True, False] 6 0.750
7 7 [True, True, True, False] 7 0.875
8 8 [False, False, False, True] 8 1.000
9 9 [True, False, False, True] 9 1.125
10 10 [False, True, False, True] 10 1.250
11 11 [True, True, False, True] 11 1.375
12 12 [False, False, True, True] 12 1.500
13 13 [True, False, True, True] 13 1.625
14 14 [False, True, True, True] 14 1.750
15 15 [True, True, True, True] 15 1.875

Testing the QROM phase-gradient rotation

def expected_ry_state(theta: float) -> np.ndarray:
    """Returns the expected state vector for a rotation around the Y axis by angle theta."""
    phase = np.pi * theta / 2.0
    return np.array([np.cos(phase), np.sin(phase)], dtype=np.complex128)

def extract_rz_theta(state: np.ndarray) -> float:
    """Extract theta from Rz(theta)|+>, up to global phase."""
    relative_phase = np.angle(state[1] / state[0])
    return float((relative_phase / np.pi) % 2.0)

Run the QROM phase-gradient rotation in superposition

from guppyalgos.primitives.rotations import QROMRotations
from guppyalgos.utils.guppy.gates import transversal

def run_phase_gradient_qrom_rotation_superposition(
    data_input: list[list[bool]],
) -> dict[str, Any]:
    n_data_qubits = len(data_input[0])
    n_index_qubits = ceil(log2(len(data_input)))
    qrom_compute = qrom_unary_iteration(data_input)
    qrom_uncompute = qrom_unary_iteration(data_input)
    fourier_state = phase_gradient(n_data_qubits, convention=Convention.Standard)


    @guppy
    @no_type_check
    def main() -> None:
        index_qreg = qarray(comptime(n_index_qubits))
        data_qreg = qarray(comptime(n_data_qubits))
        rotation_target = qubit()
        phase_state = qarray(comptime(n_data_qubits))
        fourier_state(phase_state)
        rotation = RotationPhaseGradient(
            phase_state, RotationAxisZ()
        )

        transversal(h, index_qreg)
        h(rotation_target)

        qrom_rot = QROMRotations(
            qrom_compute[array[qubit, comptime(n_data_qubits)]],
            rotation,
            qrom_uncompute[array[qubit, comptime(n_data_qubits)]],
        )

        qrom_rot.compose(index_qreg, data_qreg, rotation_target)

        state_output("index", index_qreg)
        state_output("ancilla", rotation_target)
        state_output("qrom_target", data_qreg)
        discard_array(index_qreg)
        discard_array(data_qreg)
        discard_array(qrom_rot.rotation_box.phase_gradient)
        discard(rotation_target)

    res = main.emulator(n_qubits=n_index_qubits + (3 * n_data_qubits) + 1).run()
    return Quest.extract_states_dict(res.results[0].entries)

Check that every extracted branch angle matches the stored input

For each index j, the table shows the positive kickback angle \(2x^{(j)}/2^d\) from the derivation and the extracted branch angle reported back in the same positive convention as the derivation.

expected_zero = np.zeros(2**n_phase_qubits, dtype=np.complex128)
expected_zero[0] = 1.0

states = run_phase_gradient_qrom_rotation_superposition(data_input)
qrom_target_state = states["qrom_target"].get_single_state()
assert_allclose_ignorephase(qrom_target_state, expected_zero)

# The branch helper accepts little-endian bitstrings, matching the QROM data.
index_bitstrings_le = [
    int_to_bits(index, n_index_qubits)
    for index in range(len(data_input))
]
projected_ancilla_states = extract_state_branches_in_superposition(
    states,
    "index",
    ["ancilla"],
    index_bitstrings_le,
)

rows = []
for index, bits in enumerate(data_input):
    projected_ancilla = projected_ancilla_states[tuple(index_bitstrings_le[index])]
    input_positive_kickback_angle = kickback_theta(bits)
    ancilla_state = projected_ancilla.state.state
    extracted_theta = extract_rz_theta(ancilla_state)

    rows.append(
        {
            "index j": index,
            "stored bits x^(j)": bits,
            "encoded integer x^(j)": bits_to_int(bits),
            "input positive kickback angle": input_positive_kickback_angle,
            "extracted branch angle": extracted_theta,
            "angle error mod 2": phase_distance_mod_2(extracted_theta, input_positive_kickback_angle),
        }
    )

branch_table = pd.DataFrame(rows)

branch_table
index j stored bits x^(j) encoded integer x^(j) input positive kickback angle extracted branch angle angle error mod 2
0 0 [False, False, False, False] 0 0.000 0.000 0.000000e+00
1 1 [True, False, False, False] 1 0.125 0.125 2.775558e-16
2 2 [False, True, False, False] 2 0.250 0.250 1.665335e-16
3 3 [True, True, False, False] 3 0.375 0.375 1.665335e-16
4 4 [False, False, True, False] 4 0.500 0.500 0.000000e+00
5 5 [True, False, True, False] 5 0.625 0.625 1.110223e-16
6 6 [False, True, True, False] 6 0.750 0.750 6.661338e-16
7 7 [True, True, True, False] 7 0.875 0.875 0.000000e+00
8 8 [False, False, False, True] 8 1.000 1.000 0.000000e+00
9 9 [True, False, False, True] 9 1.125 1.125 4.440892e-16
10 10 [False, True, False, True] 10 1.250 1.250 8.881784e-16
11 11 [True, True, False, True] 11 1.375 1.375 4.440892e-16
12 12 [False, False, True, True] 12 1.500 1.500 0.000000e+00
13 13 [True, False, True, True] 13 1.625 1.625 4.440892e-16
14 14 [False, True, True, True] 14 1.750 1.750 8.881784e-16
15 15 [True, True, True, True] 15 1.875 1.875 2.220446e-16