Skip to main content

oxicuda_quantum/algorithms/
amplitude_estimation.rs

1//! Quantum Amplitude Estimation (QAE).
2//!
3//! Reference: Brassard, Høyer, Mosca, Tapp, *"Quantum Amplitude Amplification
4//! and Estimation"*, Contemporary Mathematics 305 (2002), 53–74
5//! (arXiv:quant-ph/0005055).
6//!
7//! # Problem
8//!
9//! Let `A` be a state-preparation unitary on `m` qubits whose last qubit (the
10//! *flag*) marks "good" vs "bad" outcomes:
11//!
12//! ```text
13//! A|0…0⟩ = √(1−a) |ψ_0⟩|0⟩  +  √a |ψ_1⟩|1⟩,
14//! ```
15//!
16//! where `a ∈ [0, 1]` is the probability of measuring the flag in `|1⟩`. QAE
17//! estimates `a` to a precision set by the number of *counting* qubits, with a
18//! quadratically better query complexity than classical Monte-Carlo sampling.
19//!
20//! # The amplitude operator `Q`
21//!
22//! Amplitude amplification iterates the Grover-like operator
23//!
24//! ```text
25//! Q = − A · S_0 · A† · S_χ,
26//! ```
27//!
28//! where
29//!
30//! * `S_χ` is the reflection that **flips the sign of the good subspace** — here a
31//!   `Z` on the flag qubit: `S_χ = I − 2|·⟩⟨·|_{flag=1}` acting as
32//!   `|x⟩|1⟩ ↦ −|x⟩|1⟩`.
33//! * `S_0` is the reflection about `|0…0⟩`: `S_0 = I − 2|0…0⟩⟨0…0|`.
34//!
35//! Restricted to the two-dimensional span of `|ψ_good⟩` and `|ψ_bad⟩`, `Q` is a
36//! rotation by `2θ` where `a = sin²θ`. Its eigenvalues are `e^{±2iθ}`, so phase
37//! estimation of `Q` reads `θ` (equivalently `ϕ = θ/π`) and recovers
38//! `a = sin²(πϕ)`.
39//!
40//! # Algorithm (this implementation)
41//!
42//! 1. Build the `2^m × 2^m` dense matrix of `A`, of `Q`, and of `A†` from `A`
43//!    applied to the computational basis (the simulator gives us `A`'s columns
44//!    exactly).
45//! 2. Lay out `n` counting qubits as the low bits `[0, n)` and the `m`-qubit
46//!    state-prep register as `[n, n+m)`. Hadamard the counting register, apply
47//!    `A` to the prep register, then the controlled-`Q^{2^k}` ladder, then the
48//!    inverse QFT on the counting register — i.e. **standard QPE on `Q`**,
49//!    reusing this crate's QFT.
50//! 3. Read the most probable counting integer `y`, set `ϕ = y / 2^n`, and return
51//!    `a = sin²(πϕ)`.
52//!
53//! Because the eigenvalues come in the conjugate pair `e^{±2iθ}`, the readout
54//! peaks at both `ϕ` and `1−ϕ`, which map to the **same** `a = sin²(πϕ)`; the
55//! routine folds them together.
56//!
57//! # Numerics
58//!
59//! The dense `A`/`Q` matrices and all reflections are assembled in `f32`
60//! [`Complex`] consistent with the [`StateVector`]; the controlled-power ladder
61//! squares `Q` densely. The estimator is exact up to the `f32` state vector and
62//! the counting-grid resolution.
63
64use crate::error::{QuantumError, QuantumResult};
65use crate::fourier::qft::qft_inverse_inplace;
66use crate::gates::hadamard::gate_h;
67use crate::statevec::apply_1q::apply_1q_inplace;
68use crate::statevec::state::StateVector;
69use num_complex::Complex;
70
71type Complex32 = Complex<f32>;
72
73/// A state-preparation unitary `A` on `m` qubits, supplied as the closure that
74/// applies `A` to a [`StateVector`]. The last qubit (`m − 1`) is the *flag*.
75///
76/// The dense matrix of `A` is extracted once (by applying `A` to each
77/// computational basis state) and reused to build `A†` and the amplitude
78/// operator `Q`.
79pub struct StatePreparation {
80    /// Number of qubits the preparation acts on (system + flag).
81    m: usize,
82    /// Row-major `dim×dim` matrix of `A` (`dim = 2^m`): column `j` is `A|j⟩`.
83    matrix: Vec<Complex32>,
84}
85
86impl StatePreparation {
87    /// Build from a closure that applies `A` in place to an `m`-qubit state.
88    ///
89    /// The closure is invoked `2^m` times (once per basis column) to extract the
90    /// dense matrix; it must implement a genuine unitary or downstream unitarity
91    /// checks will fail.
92    ///
93    /// # Errors
94    /// * [`QuantumError::InvalidQubitCount`] if `m == 0` or `m > 12` (the dense
95    ///   `2^m × 2^m` matrix would be impractical beyond that).
96    /// * Propagates any error from the supplied `apply` closure.
97    pub fn from_apply<F>(m: usize, mut apply: F) -> QuantumResult<Self>
98    where
99        F: FnMut(&mut StateVector) -> QuantumResult<()>,
100    {
101        if m == 0 || m > 12 {
102            return Err(QuantumError::InvalidQubitCount { n: m });
103        }
104        let dim = 1usize << m;
105        let mut matrix = vec![Complex32::new(0.0, 0.0); dim * dim];
106        for j in 0..dim {
107            // Basis state |j⟩.
108            let mut amps = vec![Complex32::new(0.0, 0.0); dim];
109            amps[j] = Complex32::new(1.0, 0.0);
110            let mut sv = StateVector { amps, n_qubits: m };
111            apply(&mut sv)?;
112            // Column j of A is A|j⟩.
113            for (r, a) in sv.amps.iter().enumerate() {
114                matrix[r * dim + j] = *a;
115            }
116        }
117        Ok(Self { m, matrix })
118    }
119
120    /// Number of qubits `m` (system + flag).
121    #[must_use]
122    pub fn m(&self) -> usize {
123        self.m
124    }
125
126    /// The "good" probability `a = ⟨0…0| A† Π_1 A |0…0⟩`, computed exactly from
127    /// the dense matrix: it is the total `|amplitude|²` of `A|0…0⟩` over basis
128    /// states whose flag bit (qubit `m − 1`) is set.
129    #[must_use]
130    pub fn good_probability(&self) -> f64 {
131        let dim = 1usize << self.m;
132        let flag_mask = 1usize << (self.m - 1);
133        // A|0…0⟩ is column 0 of the matrix: entry [r][0] = matrix[r * dim].
134        let mut p = 0.0_f64;
135        for r in 0..dim {
136            if r & flag_mask != 0 {
137                p += self.matrix[r * dim].norm_sqr() as f64;
138            }
139        }
140        p
141    }
142}
143
144/// Outcome of [`amplitude_estimation`].
145#[derive(Debug, Clone)]
146pub struct AmplitudeEstimationResult {
147    /// Estimated good probability `a = sin²(πϕ) ∈ [0, 1]`.
148    pub a: f64,
149    /// Estimated angle `θ = πϕ ∈ [0, π/2]` (folded into the principal branch).
150    pub theta: f64,
151    /// The folded counting fraction `ϕ ∈ [0, 1/2]` driving the estimate.
152    pub phi: f64,
153    /// Winning counting-register integer `y` (before folding).
154    pub integer: usize,
155    /// Probability mass on the winning integer.
156    pub probability: f64,
157    /// Number of counting qubits used.
158    pub n_counting: usize,
159}
160
161/// Row-major dense `dim×dim` matrix product `c = a · b`.
162fn matmul(a: &[Complex32], b: &[Complex32], dim: usize) -> Vec<Complex32> {
163    let mut c = vec![Complex32::new(0.0, 0.0); dim * dim];
164    for i in 0..dim {
165        for k in 0..dim {
166            let aik = a[i * dim + k];
167            if aik.re == 0.0 && aik.im == 0.0 {
168                continue;
169            }
170            for j in 0..dim {
171                c[i * dim + j] += aik * b[k * dim + j];
172            }
173        }
174    }
175    c
176}
177
178/// Conjugate-transpose (adjoint) of a row-major `dim×dim` matrix.
179fn dagger(a: &[Complex32], dim: usize) -> Vec<Complex32> {
180    let mut out = vec![Complex32::new(0.0, 0.0); dim * dim];
181    for i in 0..dim {
182        for j in 0..dim {
183            out[i * dim + j] = a[j * dim + i].conj();
184        }
185    }
186    out
187}
188
189/// Build `S_χ`: a diagonal reflection that flips the sign of every state whose
190/// flag qubit (`m − 1`) is set. Returned as a dense `dim×dim` matrix.
191fn reflection_chi(m: usize) -> Vec<Complex32> {
192    let dim = 1usize << m;
193    let flag_mask = 1usize << (m - 1);
194    let mut s = vec![Complex32::new(0.0, 0.0); dim * dim];
195    for i in 0..dim {
196        let sign = if i & flag_mask != 0 { -1.0 } else { 1.0 };
197        s[i * dim + i] = Complex32::new(sign, 0.0);
198    }
199    s
200}
201
202/// Build `S_0 = I − 2|0…0⟩⟨0…0|`: a diagonal reflection that flips the sign of
203/// the all-zeros state only. Returned as a dense `dim×dim` matrix.
204fn reflection_zero(m: usize) -> Vec<Complex32> {
205    let dim = 1usize << m;
206    let mut s = vec![Complex32::new(0.0, 0.0); dim * dim];
207    for i in 0..dim {
208        let sign = if i == 0 { -1.0 } else { 1.0 };
209        s[i * dim + i] = Complex32::new(sign, 0.0);
210    }
211    s
212}
213
214/// Assemble the amplitude operator `Q = − A · S_0 · A† · S_χ` (dense, row-major).
215fn amplitude_operator(prep: &StatePreparation) -> Vec<Complex32> {
216    let dim = 1usize << prep.m;
217    let a = &prep.matrix;
218    let a_dag = dagger(a, dim);
219    let s_chi = reflection_chi(prep.m);
220    let s_zero = reflection_zero(prep.m);
221
222    // Q = −(A · (S_0 · (A† · S_χ))).
223    let t1 = matmul(&a_dag, &s_chi, dim); // A† S_χ
224    let t2 = matmul(&s_zero, &t1, dim); // S_0 A† S_χ
225    let t3 = matmul(a, &t2, dim); // A S_0 A† S_χ
226    let mut q = t3;
227    for z in &mut q {
228        *z = -*z;
229    }
230    q
231}
232
233/// Apply a dense `dim×dim` matrix `u` to the contiguous low-bit *prep* block
234/// `[base_qubit, base_qubit + m)` of `sv`, conditioned on `ctrl = 1`.
235///
236/// The prep qubits are contiguous starting at `base_qubit`, so within each
237/// outer configuration the `dim = 2^m` prep amplitudes are gathered, multiplied
238/// by `u`, and scattered back.
239fn apply_controlled_dense(
240    sv: &mut StateVector,
241    ctrl: usize,
242    base_qubit: usize,
243    m: usize,
244    u: &[Complex32],
245) -> QuantumResult<()> {
246    if ctrl >= sv.n_qubits {
247        return Err(QuantumError::QubitIndexOutOfRange {
248            index: ctrl,
249            n_qubits: sv.n_qubits,
250        });
251    }
252    let dim = 1usize << m;
253    if u.len() != dim * dim {
254        return Err(QuantumError::DimensionMismatch {
255            expected: dim * dim,
256            got: u.len(),
257        });
258    }
259    let ctrl_mask = 1usize << ctrl;
260    let prep_mask = (dim - 1) << base_qubit;
261    let total = sv.amps.len();
262
263    let mut base = 0usize;
264    while base < total {
265        // Block representative: control set, prep bits all zero.
266        if (base & ctrl_mask) != 0 && (base & prep_mask) == 0 {
267            let mut input = vec![Complex32::new(0.0, 0.0); dim];
268            for (s, slot) in input.iter_mut().enumerate() {
269                *slot = sv.amps[base | (s << base_qubit)];
270            }
271            for r in 0..dim {
272                let mut acc = Complex32::new(0.0, 0.0);
273                for (col, &xv) in input.iter().enumerate() {
274                    acc += u[r * dim + col] * xv;
275                }
276                sv.amps[base | (r << base_qubit)] = acc;
277            }
278        }
279        base += 1;
280    }
281    Ok(())
282}
283
284/// Apply a dense `dim×dim` matrix `u` to the contiguous low-bit prep block
285/// `[base_qubit, base_qubit + m)` of `sv`, unconditionally.
286fn apply_dense(
287    sv: &mut StateVector,
288    base_qubit: usize,
289    m: usize,
290    u: &[Complex32],
291) -> QuantumResult<()> {
292    let dim = 1usize << m;
293    if u.len() != dim * dim {
294        return Err(QuantumError::DimensionMismatch {
295            expected: dim * dim,
296            got: u.len(),
297        });
298    }
299    let prep_mask = (dim - 1) << base_qubit;
300    let total = sv.amps.len();
301    let mut base = 0usize;
302    while base < total {
303        if base & prep_mask == 0 {
304            let mut input = vec![Complex32::new(0.0, 0.0); dim];
305            for (s, slot) in input.iter_mut().enumerate() {
306                *slot = sv.amps[base | (s << base_qubit)];
307            }
308            for r in 0..dim {
309                let mut acc = Complex32::new(0.0, 0.0);
310                for (col, &xv) in input.iter().enumerate() {
311                    acc += u[r * dim + col] * xv;
312                }
313                sv.amps[base | (r << base_qubit)] = acc;
314            }
315        }
316        base += 1;
317    }
318    Ok(())
319}
320
321/// Run Quantum Amplitude Estimation with `n_counting` counting qubits.
322///
323/// Returns the estimated good probability `a` together with the intermediate
324/// angle/phase quantities. The estimate's resolution is set by `n_counting`:
325/// `a` lands on the grid `sin²(π y / 2^n)`.
326///
327/// # Errors
328/// * [`QuantumError::InvalidParameter`] if `n_counting == 0`.
329/// * [`QuantumError::InvalidQubitCount`] if the total register
330///   `n_counting + m` exceeds the simulator's 30-qubit limit.
331pub fn amplitude_estimation(
332    prep: &StatePreparation,
333    n_counting: usize,
334) -> QuantumResult<AmplitudeEstimationResult> {
335    if n_counting == 0 {
336        return Err(QuantumError::InvalidParameter {
337            name: "n_counting must be ≥ 1".into(),
338        });
339    }
340    let m = prep.m;
341    let n_qubits = n_counting + m;
342    if n_qubits > 30 {
343        return Err(QuantumError::InvalidQubitCount { n: n_qubits });
344    }
345
346    // Layout (little-endian): counting = [0, n_counting); prep = [n_counting,
347    // n_counting + m).
348    let count_qubits: Vec<usize> = (0..n_counting).collect();
349    let base = n_counting;
350
351    let mut sv = StateVector::new_zero_state(n_qubits)?;
352
353    // Hadamard the counting register.
354    for &q in &count_qubits {
355        apply_1q_inplace(&mut sv, q, &gate_h())?;
356    }
357
358    // Prepare A|0…0⟩ on the prep register.
359    apply_dense(&mut sv, base, m, &prep.matrix)?;
360
361    // Controlled-Q^{2^k} ladder via repeated dense squaring.
362    let dim = 1usize << m;
363    let mut q_pow = amplitude_operator(prep);
364    for &cq in &count_qubits {
365        apply_controlled_dense(&mut sv, cq, base, m, &q_pow)?;
366        q_pow = matmul(&q_pow, &q_pow, dim);
367    }
368
369    // Inverse QFT on the counting register (little-endian).
370    qft_inverse_inplace(&mut sv, &count_qubits)?;
371
372    // Marginal readout over the counting register.
373    let n_outcomes = 1usize << n_counting;
374    let mut probs = vec![0.0_f64; n_outcomes];
375    for (i, a) in sv.amps.iter().enumerate() {
376        let mut c = 0usize;
377        for (k, &q) in count_qubits.iter().enumerate() {
378            c |= ((i >> q) & 1) << k;
379        }
380        probs[c] += a.norm_sqr() as f64;
381    }
382
383    let mut integer = 0usize;
384    let mut best = f64::NEG_INFINITY;
385    for (c, &p) in probs.iter().enumerate() {
386        if p > best {
387            best = p;
388            integer = c;
389        }
390    }
391
392    // ϕ = y / 2^n; fold the conjugate eigenvalue branch (ϕ and 1−ϕ give the same
393    // a) into [0, 1/2].
394    let n_levels = n_outcomes as f64;
395    let raw_phi = integer as f64 / n_levels;
396    let phi = if raw_phi > 0.5 {
397        1.0 - raw_phi
398    } else {
399        raw_phi
400    };
401    let theta = std::f64::consts::PI * phi;
402    let a = theta.sin().powi(2);
403
404    Ok(AmplitudeEstimationResult {
405        a,
406        theta,
407        phi,
408        integer,
409        probability: best,
410        n_counting,
411    })
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417    use crate::gates::parametric::gate_ry;
418    use std::f64::consts::PI;
419
420    /// A single-qubit state-prep `A = R_y(2θ)` so that
421    /// `A|0⟩ = cos θ |0⟩ + sin θ |1⟩`, hence the flag (the only qubit) is `|1⟩`
422    /// with probability `a = sin²θ`.
423    fn ry_prep(theta: f64) -> StatePreparation {
424        StatePreparation::from_apply(1, move |sv| {
425            apply_1q_inplace(sv, 0, &gate_ry((2.0 * theta) as f32))
426        })
427        .expect("Ry state preparation with valid single-qubit operator")
428    }
429
430    /// Conjugate-transpose check helper: ‖U U† − I‖_∞.
431    fn unitarity_residual(u: &[Complex32], dim: usize) -> f32 {
432        let ud = dagger(u, dim);
433        let prod = matmul(u, &ud, dim);
434        let mut worst = 0.0_f32;
435        for i in 0..dim {
436            for j in 0..dim {
437                let expect = if i == j {
438                    Complex32::new(1.0, 0.0)
439                } else {
440                    Complex32::new(0.0, 0.0)
441                };
442                worst = worst.max((prod[i * dim + j] - expect).norm());
443            }
444        }
445        worst
446    }
447
448    // (a) a ∈ {1/4, 1/2, 3/4} land near QPE grid points → estimate ≈ a.
449    #[test]
450    fn estimates_quarter_half_three_quarter() {
451        // a = sin²θ. θ = π·ϕ with ϕ = y/2^n. For a = 1/2, θ = π/4 ⇒ ϕ = 1/4,
452        // a grid point at n ≥ 2. For a = 1/4, θ = π/6 ⇒ ϕ = 1/6 (≈ near 2/12);
453        // pick n large enough that the nearest grid point is close. For a = 3/4,
454        // θ = π/3 ⇒ ϕ = 1/3.
455        let n = 6usize;
456        for &a_true in &[0.25_f64, 0.5, 0.75] {
457            let theta = a_true.sqrt().asin();
458            let prep = ry_prep(theta);
459            let res = amplitude_estimation(&prep, n)
460                .expect("amplitude estimation succeeds for valid preparation");
461            assert!(
462                (res.a - a_true).abs() < 0.05,
463                "a_true={a_true}: estimate {} (ϕ={}, y={})",
464                res.a,
465                res.phi,
466                res.integer
467            );
468        }
469    }
470
471    // a = 1/2 is an EXACT grid point (ϕ = 1/4) and must be recovered tightly.
472    #[test]
473    fn exact_grid_point_half() {
474        let prep = ry_prep((0.5_f64).sqrt().asin()); // θ = π/4
475        let res = amplitude_estimation(&prep, 4)
476            .expect("amplitude estimation with 4 counting qubits succeeds");
477        assert!((res.a - 0.5).abs() < 1e-3, "a={}", res.a);
478        assert!((res.phi - 0.25).abs() < 1e-6, "phi={}", res.phi);
479    }
480
481    // (b) a = 0 → 0 and a = 1 → 1.
482    #[test]
483    fn extremes_zero_and_one() {
484        // a = 0: θ = 0, A = I (R_y(0)). Good probability 0.
485        let prep0 = ry_prep(0.0);
486        assert!(prep0.good_probability().abs() < 1e-7);
487        let res0 = amplitude_estimation(&prep0, 5)
488            .expect("amplitude estimation for zero-amplitude preparation");
489        assert!(res0.a.abs() < 1e-6, "a(0)={}", res0.a);
490
491        // a = 1: θ = π/2, A|0⟩ = |1⟩. Good probability 1; ϕ = 1/2 grid point.
492        let prep1 = ry_prep(PI / 2.0);
493        assert!((prep1.good_probability() - 1.0).abs() < 1e-6);
494        let res1 = amplitude_estimation(&prep1, 5)
495            .expect("amplitude estimation for unit-amplitude preparation");
496        assert!((res1.a - 1.0).abs() < 1e-6, "a(1)={}", res1.a);
497    }
498
499    // (c) Q is unitary: Q Q† = I to 1e-5.
500    #[test]
501    fn amplitude_operator_is_unitary() {
502        for &theta in &[0.1_f64, PI / 6.0, PI / 4.0, PI / 3.0, 1.2, PI / 2.0] {
503            let prep = ry_prep(theta);
504            let q = amplitude_operator(&prep);
505            let dim = 1usize << prep.m;
506            let res = unitarity_residual(&q, dim);
507            assert!(res < 1e-5, "θ={theta}: ‖QQ†−I‖={res}");
508        }
509        // Also a 2-qubit prep (system + flag) to exercise the dense path.
510        let prep2 = StatePreparation::from_apply(2, |sv| {
511            apply_1q_inplace(sv, 0, &gate_ry(0.7))?;
512            // Make qubit 1 (the flag) depend on qubit 0 via a CNOT-like effect.
513            apply_1q_inplace(sv, 1, &gate_ry(1.1))
514        })
515        .expect("2-qubit state preparation with two Ry rotations succeeds");
516        let q2 = amplitude_operator(&prep2);
517        let r2 = unitarity_residual(&q2, 4);
518        assert!(r2 < 1e-5, "2-qubit ‖QQ†−I‖={r2}");
519    }
520
521    // (d) The eigenphase relation θ = arcsin(√a) holds: Q's eigenvalues are
522    //     e^{±2iθ}, so the QPE-recovered θ matches arcsin(√a).
523    #[test]
524    fn eigenphase_relation_holds() {
525        let n = 8usize;
526        for &a_true in &[0.25_f64, 0.5, 0.75] {
527            let theta_true = a_true.sqrt().asin();
528            let prep = ry_prep(theta_true);
529            let res = amplitude_estimation(&prep, n)
530                .expect("amplitude estimation succeeds for eigenphase test");
531            // Recovered θ = π·ϕ should match arcsin(√a_true) closely.
532            assert!(
533                (res.theta - theta_true).abs() < 0.05,
534                "a={a_true}: θ_est={} vs arcsin√a={theta_true}",
535                res.theta
536            );
537        }
538    }
539
540    // (e) More counting qubits → strictly finer resolution (error shrinks) for a
541    //     non-grid value of a.
542    #[test]
543    fn more_qubits_finer_resolution() {
544        // a with θ = 1.0 rad (not a dyadic multiple of π): ϕ = 1/π ≈ 0.3183,
545        // never exactly on the grid, so the error is resolution-limited and must
546        // shrink monotonically (loosely) as n grows.
547        let theta = 1.0_f64;
548        let a_true = theta.sin().powi(2);
549        let mut prev_err = f64::INFINITY;
550        for n in [3usize, 5, 7, 9] {
551            let prep = ry_prep(theta);
552            let res = amplitude_estimation(&prep, n)
553                .expect("amplitude estimation with increasing counting qubits succeeds");
554            let err = (res.a - a_true).abs();
555            assert!(
556                err <= prev_err + 1e-9,
557                "n={n}: error {err} did not shrink (prev {prev_err})"
558            );
559            prev_err = err;
560        }
561        // And the finest resolution is genuinely accurate.
562        let res_fine = amplitude_estimation(&ry_prep(theta), 10)
563            .expect("fine-resolution amplitude estimation with 10 counting qubits");
564        assert!((res_fine.a - a_true).abs() < 0.01, "fine a={}", res_fine.a);
565    }
566
567    // (f) S_0 and S_χ reflections are correct on small states (unit tests).
568    #[test]
569    fn reflections_are_correct() {
570        // S_χ on 1 qubit flips |1⟩ only: diag(1, -1) = Z.
571        let s_chi = reflection_chi(1);
572        assert!((s_chi[0] - Complex32::new(1.0, 0.0)).norm() < 1e-7); // [0][0]
573        assert!(s_chi[1].norm() < 1e-7); // [0][1]
574        assert!(s_chi[2].norm() < 1e-7); // [1][0]
575        assert!((s_chi[3] - Complex32::new(-1.0, 0.0)).norm() < 1e-7); // [1][1]
576
577        // S_0 on 2 qubits flips |00⟩ only.
578        let s0 = reflection_zero(2);
579        assert!((s0[0] - Complex32::new(-1.0, 0.0)).norm() < 1e-7); // [0][0] = -1
580        for i in 1..4 {
581            assert!(
582                (s0[i * 4 + i] - Complex32::new(1.0, 0.0)).norm() < 1e-7,
583                "S0 diag[{i}]"
584            );
585        }
586
587        // S_χ on 2 qubits (flag = qubit 1) flips states with bit 1 set: indices
588        // 2 (|10⟩) and 3 (|11⟩).
589        let s_chi2 = reflection_chi(2);
590        let diag: Vec<f32> = (0..4).map(|i| s_chi2[i * 4 + i].re).collect();
591        assert_eq!(diag, vec![1.0, 1.0, -1.0, -1.0], "S_χ diag={diag:?}");
592    }
593
594    // good_probability matches the |R_y| amplitude formula a = sin²θ exactly.
595    #[test]
596    fn good_probability_matches_formula() {
597        for &theta in &[0.0_f64, 0.3, PI / 6.0, PI / 4.0, 1.0, PI / 2.0] {
598            let prep = ry_prep(theta);
599            let want = theta.sin().powi(2);
600            assert!(
601                (prep.good_probability() - want).abs() < 1e-6,
602                "θ={theta}: good_prob={} vs sin²θ={want}",
603                prep.good_probability()
604            );
605        }
606    }
607
608    // Validation: n_counting = 0 and oversized registers error out.
609    #[test]
610    fn rejects_invalid_inputs() {
611        let prep = ry_prep(0.5);
612        assert!(amplitude_estimation(&prep, 0).is_err());
613        // m = 0 prep is rejected at construction.
614        assert!(StatePreparation::from_apply(0, |_| Ok(())).is_err());
615    }
616}