Add test for measurement

spinoza/examples/measurement.rs

@@ -19,7 +19,7 @@ fn h(n: usize, show_results: bool) {
     println!("{}", pretty_print_int(elapsed));
 
     let now = std::time::Instant::now();
-    measure_qubit(&mut state, n - 1, Some(1));
+    measure_qubit(&mut state, n - 1, true, Some(1));
     let elapsed = now.elapsed().as_micros();
     println!("{}", pretty_print_int(elapsed));
 }

spinoza/src/measurement.rs

@@ -1,8 +1,12 @@
-use crate::{core::State, math::modulus};
+use crate::{
+    core::State,
+    gates::{apply, Gate},
+    math::modulus,
+};
 use rand_distr::{Binomial, Distribution};
 
 /// Single qubit measurement
-pub fn measure_qubit(state: &mut State, target: usize, v: Option<u64>) {
+pub fn measure_qubit(state: &mut State, target: usize, reset: bool, v: Option<u64>) {
     let mut prob0 = 0.0;
     let mut prob1 = 0.0;
     let num_pairs = state.len() >> 1;
@@ -16,12 +20,11 @@ pub fn measure_qubit(state: &mut State, target: usize, v: Option<u64>) {
         prob1 += modulus(state.reals[s1], state.imags[s1]).powi(2);
     }
 
-    let bin = Binomial::new(1, prob1).unwrap();
-
     let val = if let Some(_v) = v {
         assert!(_v == 0 || _v == 1);
         _v
     } else {
+        let bin = Binomial::new(1, prob1).unwrap();
         bin.sample(&mut rand::thread_rng())
     };
 
@@ -45,5 +48,59 @@ pub fn measure_qubit(state: &mut State, target: usize, v: Option<u64>) {
             state.reals[s1] = 0.0;
             state.imags[s1] = 0.0;
         }
+
+        if reset {
+            apply(Gate::X, state, target);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::utils::{assert_float_closeness, gen_random_state};
+
+    #[test]
+    fn test_measure_qubit() {
+        let mut state = gen_random_state(3);
+
+        let sum = state
+            .reals
+            .iter()
+            .zip(state.imags.iter())
+            .map(|(re, im)| modulus(*re, *im).powi(2))
+            .sum();
+
+        assert_float_closeness(sum, 1.0, 0.001);
+
+        measure_qubit(&mut state, 0, true, Some(0));
+        let sum = state
+            .reals
+            .iter()
+            .zip(state.imags.iter())
+            .map(|(re, im)| modulus(*re, *im).powi(2))
+            .sum();
+
+        assert_float_closeness(sum, 1.0, 0.001);
+
+        measure_qubit(&mut state, 1, true, Some(0));
+        let sum = state
+            .reals
+            .iter()
+            .zip(state.imags.iter())
+            .map(|(re, im)| modulus(*re, *im).powi(2))
+            .sum();
+
+        assert_float_closeness(sum, 1.0, 0.001);
+
+        measure_qubit(&mut state, 2, true, Some(1));
+        let sum = state
+            .reals
+            .iter()
+            .zip(state.imags.iter())
+            .map(|(re, im)| modulus(*re, *im).powi(2))
+            .sum();
+
+        assert_float_closeness(sum, 1.0, 0.001);
     }
 }

spinoza/src/utils.rs

@@ -1,13 +1,16 @@
 //! An assortment of utility functions for visualizing, benchmarking, and testing.
 use crate::{
     core::State,
+    gates::Gate,
     math::{modulus, Float, PI},
 };
 use comfy_table::{
     presets::UTF8_FULL,
     Color::Rgb,
     {Cell, Color, Table},
 };
+use rand::distributions::Uniform;
+use rand::prelude::*;
 
 /// Formats an unsigned, 128 bit integer with commas, as a string. Used for readability
 pub fn pretty_print_int(i: u128) -> String {
@@ -180,6 +183,42 @@ pub fn assert_float_closeness(actual: Float, expected: Float, epsilon: Float) {
     assert!((actual - expected).abs() < epsilon);
 }
 
+/// Generates a random quantum state
+pub fn gen_random_state(n: usize) -> State {
+    assert!(n > 0);
+    let mut rng = rand::thread_rng();
+    let between = Uniform::from(0.0..1.0);
+    let angle_dist = Uniform::from(0.0..2.0 * PI);
+    let num_amps = 1 << n;
+
+    let mut probs: Vec<_> = (0..num_amps).map(|_| between.sample(&mut rng)).collect();
+
+    let total: Float = probs.iter().sum();
+    let total_recip = total.recip();
+
+    probs.iter_mut().for_each(|p| *p *= total_recip);
+
+    let angles = (0..num_amps).map(|_| angle_dist.sample(&mut rng));
+
+    let mut reals = Vec::with_capacity(num_amps);
+    let mut imags = Vec::with_capacity(num_amps);
+
+    probs.iter().zip(angles).for_each(|(p, a)| {
+        let p_sqrt = p.sqrt();
+        let (sin_a, cos_a) = a.sin_cos();
+        let re = p_sqrt * cos_a;
+        let im = p_sqrt * sin_a;
+        reals.push(re);
+        imags.push(im);
+    });
+
+    State {
+        reals,
+        imags,
+        n: n.try_into().unwrap(),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;