use crate::error::{QuantumError, QuantumResult};
use crate::fourier::qft::qft_inverse_inplace;
use crate::gates::hadamard::gate_h;
use crate::statevec::apply_1q::apply_1q_inplace;
use crate::statevec::state::StateVector;
use num_complex::Complex;
type Complex32 = Complex<f32>;
pub struct StatePreparation {
m: usize,
matrix: Vec<Complex32>,
}
impl StatePreparation {
pub fn from_apply<F>(m: usize, mut apply: F) -> QuantumResult<Self>
where
F: FnMut(&mut StateVector) -> QuantumResult<()>,
{
if m == 0 || m > 12 {
return Err(QuantumError::InvalidQubitCount { n: m });
}
let dim = 1usize << m;
let mut matrix = vec![Complex32::new(0.0, 0.0); dim * dim];
for j in 0..dim {
let mut amps = vec![Complex32::new(0.0, 0.0); dim];
amps[j] = Complex32::new(1.0, 0.0);
let mut sv = StateVector { amps, n_qubits: m };
apply(&mut sv)?;
for (r, a) in sv.amps.iter().enumerate() {
matrix[r * dim + j] = *a;
}
}
Ok(Self { m, matrix })
}
#[must_use]
pub fn m(&self) -> usize {
self.m
}
#[must_use]
pub fn good_probability(&self) -> f64 {
let dim = 1usize << self.m;
let flag_mask = 1usize << (self.m - 1);
let mut p = 0.0_f64;
for r in 0..dim {
if r & flag_mask != 0 {
p += self.matrix[r * dim].norm_sqr() as f64;
}
}
p
}
}
#[derive(Debug, Clone)]
pub struct AmplitudeEstimationResult {
pub a: f64,
pub theta: f64,
pub phi: f64,
pub integer: usize,
pub probability: f64,
pub n_counting: usize,
}
fn matmul(a: &[Complex32], b: &[Complex32], dim: usize) -> Vec<Complex32> {
let mut c = vec![Complex32::new(0.0, 0.0); dim * dim];
for i in 0..dim {
for k in 0..dim {
let aik = a[i * dim + k];
if aik.re == 0.0 && aik.im == 0.0 {
continue;
}
for j in 0..dim {
c[i * dim + j] += aik * b[k * dim + j];
}
}
}
c
}
fn dagger(a: &[Complex32], dim: usize) -> Vec<Complex32> {
let mut out = vec![Complex32::new(0.0, 0.0); dim * dim];
for i in 0..dim {
for j in 0..dim {
out[i * dim + j] = a[j * dim + i].conj();
}
}
out
}
fn reflection_chi(m: usize) -> Vec<Complex32> {
let dim = 1usize << m;
let flag_mask = 1usize << (m - 1);
let mut s = vec![Complex32::new(0.0, 0.0); dim * dim];
for i in 0..dim {
let sign = if i & flag_mask != 0 { -1.0 } else { 1.0 };
s[i * dim + i] = Complex32::new(sign, 0.0);
}
s
}
fn reflection_zero(m: usize) -> Vec<Complex32> {
let dim = 1usize << m;
let mut s = vec![Complex32::new(0.0, 0.0); dim * dim];
for i in 0..dim {
let sign = if i == 0 { -1.0 } else { 1.0 };
s[i * dim + i] = Complex32::new(sign, 0.0);
}
s
}
fn amplitude_operator(prep: &StatePreparation) -> Vec<Complex32> {
let dim = 1usize << prep.m;
let a = &prep.matrix;
let a_dag = dagger(a, dim);
let s_chi = reflection_chi(prep.m);
let s_zero = reflection_zero(prep.m);
let t1 = matmul(&a_dag, &s_chi, dim); let t2 = matmul(&s_zero, &t1, dim); let t3 = matmul(a, &t2, dim); let mut q = t3;
for z in &mut q {
*z = -*z;
}
q
}
fn apply_controlled_dense(
sv: &mut StateVector,
ctrl: usize,
base_qubit: usize,
m: usize,
u: &[Complex32],
) -> QuantumResult<()> {
if ctrl >= sv.n_qubits {
return Err(QuantumError::QubitIndexOutOfRange {
index: ctrl,
n_qubits: sv.n_qubits,
});
}
let dim = 1usize << m;
if u.len() != dim * dim {
return Err(QuantumError::DimensionMismatch {
expected: dim * dim,
got: u.len(),
});
}
let ctrl_mask = 1usize << ctrl;
let prep_mask = (dim - 1) << base_qubit;
let total = sv.amps.len();
let mut base = 0usize;
while base < total {
if (base & ctrl_mask) != 0 && (base & prep_mask) == 0 {
let mut input = vec![Complex32::new(0.0, 0.0); dim];
for (s, slot) in input.iter_mut().enumerate() {
*slot = sv.amps[base | (s << base_qubit)];
}
for r in 0..dim {
let mut acc = Complex32::new(0.0, 0.0);
for (col, &xv) in input.iter().enumerate() {
acc += u[r * dim + col] * xv;
}
sv.amps[base | (r << base_qubit)] = acc;
}
}
base += 1;
}
Ok(())
}
fn apply_dense(
sv: &mut StateVector,
base_qubit: usize,
m: usize,
u: &[Complex32],
) -> QuantumResult<()> {
let dim = 1usize << m;
if u.len() != dim * dim {
return Err(QuantumError::DimensionMismatch {
expected: dim * dim,
got: u.len(),
});
}
let prep_mask = (dim - 1) << base_qubit;
let total = sv.amps.len();
let mut base = 0usize;
while base < total {
if base & prep_mask == 0 {
let mut input = vec![Complex32::new(0.0, 0.0); dim];
for (s, slot) in input.iter_mut().enumerate() {
*slot = sv.amps[base | (s << base_qubit)];
}
for r in 0..dim {
let mut acc = Complex32::new(0.0, 0.0);
for (col, &xv) in input.iter().enumerate() {
acc += u[r * dim + col] * xv;
}
sv.amps[base | (r << base_qubit)] = acc;
}
}
base += 1;
}
Ok(())
}
pub fn amplitude_estimation(
prep: &StatePreparation,
n_counting: usize,
) -> QuantumResult<AmplitudeEstimationResult> {
if n_counting == 0 {
return Err(QuantumError::InvalidParameter {
name: "n_counting must be ≥ 1".into(),
});
}
let m = prep.m;
let n_qubits = n_counting + m;
if n_qubits > 30 {
return Err(QuantumError::InvalidQubitCount { n: n_qubits });
}
let count_qubits: Vec<usize> = (0..n_counting).collect();
let base = n_counting;
let mut sv = StateVector::new_zero_state(n_qubits)?;
for &q in &count_qubits {
apply_1q_inplace(&mut sv, q, &gate_h())?;
}
apply_dense(&mut sv, base, m, &prep.matrix)?;
let dim = 1usize << m;
let mut q_pow = amplitude_operator(prep);
for &cq in &count_qubits {
apply_controlled_dense(&mut sv, cq, base, m, &q_pow)?;
q_pow = matmul(&q_pow, &q_pow, dim);
}
qft_inverse_inplace(&mut sv, &count_qubits)?;
let n_outcomes = 1usize << n_counting;
let mut probs = vec![0.0_f64; n_outcomes];
for (i, a) in sv.amps.iter().enumerate() {
let mut c = 0usize;
for (k, &q) in count_qubits.iter().enumerate() {
c |= ((i >> q) & 1) << k;
}
probs[c] += a.norm_sqr() as f64;
}
let mut integer = 0usize;
let mut best = f64::NEG_INFINITY;
for (c, &p) in probs.iter().enumerate() {
if p > best {
best = p;
integer = c;
}
}
let n_levels = n_outcomes as f64;
let raw_phi = integer as f64 / n_levels;
let phi = if raw_phi > 0.5 {
1.0 - raw_phi
} else {
raw_phi
};
let theta = std::f64::consts::PI * phi;
let a = theta.sin().powi(2);
Ok(AmplitudeEstimationResult {
a,
theta,
phi,
integer,
probability: best,
n_counting,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::gates::parametric::gate_ry;
use std::f64::consts::PI;
fn ry_prep(theta: f64) -> StatePreparation {
StatePreparation::from_apply(1, move |sv| {
apply_1q_inplace(sv, 0, &gate_ry((2.0 * theta) as f32))
})
.expect("Ry state preparation with valid single-qubit operator")
}
fn unitarity_residual(u: &[Complex32], dim: usize) -> f32 {
let ud = dagger(u, dim);
let prod = matmul(u, &ud, dim);
let mut worst = 0.0_f32;
for i in 0..dim {
for j in 0..dim {
let expect = if i == j {
Complex32::new(1.0, 0.0)
} else {
Complex32::new(0.0, 0.0)
};
worst = worst.max((prod[i * dim + j] - expect).norm());
}
}
worst
}
#[test]
fn estimates_quarter_half_three_quarter() {
let n = 6usize;
for &a_true in &[0.25_f64, 0.5, 0.75] {
let theta = a_true.sqrt().asin();
let prep = ry_prep(theta);
let res = amplitude_estimation(&prep, n)
.expect("amplitude estimation succeeds for valid preparation");
assert!(
(res.a - a_true).abs() < 0.05,
"a_true={a_true}: estimate {} (ϕ={}, y={})",
res.a,
res.phi,
res.integer
);
}
}
#[test]
fn exact_grid_point_half() {
let prep = ry_prep((0.5_f64).sqrt().asin()); let res = amplitude_estimation(&prep, 4)
.expect("amplitude estimation with 4 counting qubits succeeds");
assert!((res.a - 0.5).abs() < 1e-3, "a={}", res.a);
assert!((res.phi - 0.25).abs() < 1e-6, "phi={}", res.phi);
}
#[test]
fn extremes_zero_and_one() {
let prep0 = ry_prep(0.0);
assert!(prep0.good_probability().abs() < 1e-7);
let res0 = amplitude_estimation(&prep0, 5)
.expect("amplitude estimation for zero-amplitude preparation");
assert!(res0.a.abs() < 1e-6, "a(0)={}", res0.a);
let prep1 = ry_prep(PI / 2.0);
assert!((prep1.good_probability() - 1.0).abs() < 1e-6);
let res1 = amplitude_estimation(&prep1, 5)
.expect("amplitude estimation for unit-amplitude preparation");
assert!((res1.a - 1.0).abs() < 1e-6, "a(1)={}", res1.a);
}
#[test]
fn amplitude_operator_is_unitary() {
for &theta in &[0.1_f64, PI / 6.0, PI / 4.0, PI / 3.0, 1.2, PI / 2.0] {
let prep = ry_prep(theta);
let q = amplitude_operator(&prep);
let dim = 1usize << prep.m;
let res = unitarity_residual(&q, dim);
assert!(res < 1e-5, "θ={theta}: ‖QQ†−I‖={res}");
}
let prep2 = StatePreparation::from_apply(2, |sv| {
apply_1q_inplace(sv, 0, &gate_ry(0.7))?;
apply_1q_inplace(sv, 1, &gate_ry(1.1))
})
.expect("2-qubit state preparation with two Ry rotations succeeds");
let q2 = amplitude_operator(&prep2);
let r2 = unitarity_residual(&q2, 4);
assert!(r2 < 1e-5, "2-qubit ‖QQ†−I‖={r2}");
}
#[test]
fn eigenphase_relation_holds() {
let n = 8usize;
for &a_true in &[0.25_f64, 0.5, 0.75] {
let theta_true = a_true.sqrt().asin();
let prep = ry_prep(theta_true);
let res = amplitude_estimation(&prep, n)
.expect("amplitude estimation succeeds for eigenphase test");
assert!(
(res.theta - theta_true).abs() < 0.05,
"a={a_true}: θ_est={} vs arcsin√a={theta_true}",
res.theta
);
}
}
#[test]
fn more_qubits_finer_resolution() {
let theta = 1.0_f64;
let a_true = theta.sin().powi(2);
let mut prev_err = f64::INFINITY;
for n in [3usize, 5, 7, 9] {
let prep = ry_prep(theta);
let res = amplitude_estimation(&prep, n)
.expect("amplitude estimation with increasing counting qubits succeeds");
let err = (res.a - a_true).abs();
assert!(
err <= prev_err + 1e-9,
"n={n}: error {err} did not shrink (prev {prev_err})"
);
prev_err = err;
}
let res_fine = amplitude_estimation(&ry_prep(theta), 10)
.expect("fine-resolution amplitude estimation with 10 counting qubits");
assert!((res_fine.a - a_true).abs() < 0.01, "fine a={}", res_fine.a);
}
#[test]
fn reflections_are_correct() {
let s_chi = reflection_chi(1);
assert!((s_chi[0] - Complex32::new(1.0, 0.0)).norm() < 1e-7); assert!(s_chi[1].norm() < 1e-7); assert!(s_chi[2].norm() < 1e-7); assert!((s_chi[3] - Complex32::new(-1.0, 0.0)).norm() < 1e-7);
let s0 = reflection_zero(2);
assert!((s0[0] - Complex32::new(-1.0, 0.0)).norm() < 1e-7); for i in 1..4 {
assert!(
(s0[i * 4 + i] - Complex32::new(1.0, 0.0)).norm() < 1e-7,
"S0 diag[{i}]"
);
}
let s_chi2 = reflection_chi(2);
let diag: Vec<f32> = (0..4).map(|i| s_chi2[i * 4 + i].re).collect();
assert_eq!(diag, vec![1.0, 1.0, -1.0, -1.0], "S_χ diag={diag:?}");
}
#[test]
fn good_probability_matches_formula() {
for &theta in &[0.0_f64, 0.3, PI / 6.0, PI / 4.0, 1.0, PI / 2.0] {
let prep = ry_prep(theta);
let want = theta.sin().powi(2);
assert!(
(prep.good_probability() - want).abs() < 1e-6,
"θ={theta}: good_prob={} vs sin²θ={want}",
prep.good_probability()
);
}
}
#[test]
fn rejects_invalid_inputs() {
let prep = ry_prep(0.5);
assert!(amplitude_estimation(&prep, 0).is_err());
assert!(StatePreparation::from_apply(0, |_| Ok(())).is_err());
}
}