Add DumpOperation support in Q# (#1885)

This enables Q# code to call 'DumpOperation' which will output to
console/debugger as text, or Jupyter notebooks as a LaTeX matrix.

e.g. from a notebook cell 

```
%%qsharp

open Microsoft.Quantum.Diagnostics;
operation Foo(qs: Qubit[]) : Unit {
    H(qs[0]);
    CX(qs[0], qs[1]);
}

operation Main() : Unit {
    use qs = Qubit[2];
    Foo(qs);
    DumpMachine();
    Message("About to dump operation Foo");
    DumpOperation(2, Foo);
    ResetAll(qs);
}

Main()
```

Output is:

<img width="783" alt="image"
src="https://github.com/user-attachments/assets/48e94c91-d291-4a4f-9d88-ae374472641b">

---------

Co-authored-by: Stefan J. Wernli <[email protected]>
Co-authored-by: Dmitry Vasilevsky <[email protected]>
Co-authored-by: DmitryVasilevsky <[email protected]>

compiler/qsc/src/bin/qsi.rs

@@ -81,6 +81,16 @@ impl Receiver for TerminalReceiver {
         Ok(())
     }
 
+    fn matrix(&mut self, matrix: Vec<Vec<Complex64>>) -> std::result::Result<(), output::Error> {
+        println!("Matrix:");
+        for row in matrix {
+            let row = row.iter().map(|elem| format!("[{}, {}]", elem.re, elem.im));
+            println!("{}", row.collect::<Vec<_>>().join(", "));
+        }
+
+        Ok(())
+    }
+
     fn message(&mut self, msg: &str) -> Result<(), output::Error> {
         println!("{msg}");
         Ok(())

compiler/qsc/src/lib.rs

@@ -49,7 +49,10 @@ pub mod line_column {
 
 pub use qsc_eval::{
     backend::{Backend, SparseSim},
-    state::{fmt_basis_state_label, fmt_complex, format_state_id, get_latex, get_phase},
+    state::{
+        fmt_basis_state_label, fmt_complex, format_state_id, get_matrix_latex, get_phase,
+        get_state_latex,
+    },
 };
 
 pub mod linter {

compiler/qsc_eval/src/intrinsic.rs

@@ -77,6 +77,24 @@ pub(crate) fn call(
                 Err(_) => Err(Error::OutputFail(name_span)),
             }
         }
+        "DumpMatrix" => {
+            let qubits = arg.unwrap_array();
+            let qubits = qubits
+                .iter()
+                .map(|q| q.clone().unwrap_qubit().0)
+                .collect::<Vec<_>>();
+            if qubits.len() != qubits.iter().collect::<FxHashSet<_>>().len() {
+                return Err(Error::QubitUniqueness(arg_span));
+            }
+            let (state, qubit_count) = sim.capture_quantum_state();
+            let state = utils::split_state(&qubits, &state, qubit_count)
+                .map_err(|()| Error::QubitsNotSeparable(arg_span))?;
+            let matrix = utils::state_to_matrix(state, qubits.len() / 2);
+            match out.matrix(matrix) {
+                Ok(()) => Ok(Value::unit()),
+                Err(_) => Err(Error::OutputFail(name_span)),
+            }
+        }
         "PermuteLabels" => qubit_relabel(arg, arg_span, |q0, q1| sim.qubit_swap_id(q0, q1)),
         "Message" => match out.message(&arg.unwrap_string()) {
             Ok(()) => Ok(Value::unit()),

compiler/qsc_eval/src/intrinsic/tests.rs

@@ -440,6 +440,26 @@ fn dump_register_other_qubits_one_state_is_separable() {
     );
 }
 
+#[test]
+fn dump_register_other_qubits_phase_reflected_in_subset() {
+    check_intrinsic_output(
+        "",
+        indoc! {"{
+            use qs = Qubit[3];
+            H(qs[0]);
+            X(qs[2]);
+            Z(qs[2]);
+            Microsoft.Quantum.Diagnostics.DumpRegister(qs[...1]);
+            ResetAll(qs);
+        }"},
+        &expect![[r#"
+            STATE:
+            |00⟩: −0.7071+0.0000𝑖
+            |10⟩: −0.7071+0.0000𝑖
+        "#]],
+    );
+}
+
 #[test]
 fn dump_register_qubits_reorder_output() {
     check_intrinsic_output(

compiler/qsc_eval/src/intrinsic/utils.rs

@@ -5,8 +5,8 @@ use std::collections::hash_map::Entry;
 
 use num_bigint::BigUint;
 use num_complex::{Complex, Complex64};
-use num_traits::{One, Zero};
-use rustc_hash::FxHashMap;
+use num_traits::Zero;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 /// Given a state and a set of qubits, split the state into two parts: the qubits to dump and the remaining qubits.
 /// This function will return an error if the state is not separable using the provided qubit identifiers.
@@ -22,26 +22,13 @@ pub fn split_state(
     }
 
     let mut dump_state = FxHashMap::default();
-    let mut other_state = FxHashMap::default();
 
     // Compute the mask for the qubits to dump and the mask for the other qubits.
     let (dump_mask, other_mask) = compute_mask(qubit_count, qubits);
 
     // Try to split out the state for the given qubits from the whole state, detecting any entanglement
     // and returning an error if the qubits are not separable.
-    let dump_norm = collect_split_state(
-        state,
-        &dump_mask,
-        &other_mask,
-        &mut dump_state,
-        &mut other_state,
-    )?;
-
-    // If the product of the collected states is not equal to the total number of input states, then that
-    // implies some states are zero amplitude that would have to be non-zero for the state to be separable.
-    if state.len() != dump_state.len() * other_state.len() {
-        return Err(());
-    }
+    let dump_norm = collect_split_state(state, &dump_mask, &other_mask, &mut dump_state)?;
 
     let dump_norm = 1.0 / dump_norm.sqrt();
     let mut dump_state = dump_state
@@ -79,7 +66,6 @@ fn collect_split_state(
     dump_mask: &BigUint,
     other_mask: &BigUint,
     dump_state: &mut FxHashMap<BigUint, Complex64>,
-    other_state: &mut FxHashMap<BigUint, Complex64>,
 ) -> Result<f64, ()> {
     // To ensure consistent ordering, we iterate over the vector directly (returned from the simulator in a deterministic order),
     // and not the map used for arbitrary lookup.
@@ -88,12 +74,11 @@ fn collect_split_state(
     let (base_label, base_val) = state_iter.next().expect("state should never be empty");
     let dump_base_label = base_label & dump_mask;
     let other_base_label = base_label & other_mask;
-    let mut dump_norm = 1.0_f64;
+    let mut dump_norm = base_val.norm().powi(2);
+    let mut other_state = FxHashSet::default();
 
-    // Start with an amplitude of 1 in the first expected split states. This becomes the basis
-    // of the split later, but will get normalized as part of collecting the remaining states.
-    dump_state.insert(dump_base_label.clone(), Complex64::one());
-    other_state.insert(other_base_label.clone(), Complex64::one());
+    dump_state.insert(dump_base_label.clone(), *base_val);
+    other_state.insert(other_base_label.clone());
 
     for (curr_label, curr_val) in state_iter {
         let dump_label = curr_label & dump_mask;
@@ -117,23 +102,23 @@ fn collect_split_state(
         }
 
         if let Entry::Vacant(entry) = dump_state.entry(dump_label) {
-            // When capturing the amplitude for the dump state, we must divide out the amplitude for the other
-            // state, and vice-versa below.
-            let amplitude = curr_val / other_val;
+            let amplitude = *curr_val;
             let norm = amplitude.norm().powi(2);
             if !norm.is_nearly_zero() {
                 entry.insert(amplitude);
                 dump_norm += norm;
             }
         }
-        if let Entry::Vacant(entry) = other_state.entry(other_label) {
-            let amplitude = curr_val / dump_val;
-            let norm = amplitude.norm().powi(2);
-            if !norm.is_nearly_zero() {
-                entry.insert(amplitude);
-            }
+        if !(curr_val / dump_val).norm().powi(2).is_nearly_zero() {
+            other_state.insert(other_label);
         }
     }
+
+    // If the product of the collected states is not equal to the total number of input states, then that
+    // implies some states are zero amplitude that would have to be non-zero for the state to be separable.
+    if state.len() != dump_state.len() * other_state.len() {
+        return Err(());
+    }
     Ok(dump_norm)
 }
 
@@ -185,3 +170,28 @@ where
         self.re.is_nearly_zero() && self.im.is_nearly_zero()
     }
 }
+
+pub(crate) fn state_to_matrix(
+    state: Vec<(BigUint, Complex64)>,
+    qubit_count: usize,
+) -> Vec<Vec<Complex64>> {
+    let state: FxHashMap<BigUint, Complex<f64>> = state.into_iter().collect();
+    let mut matrix = Vec::new();
+    let num_entries: usize = 1 << qubit_count;
+    #[allow(clippy::cast_precision_loss)]
+    let factor = (num_entries as f64).sqrt();
+    for i in 0..num_entries {
+        let mut row = Vec::new();
+        for j in 0..num_entries {
+            let key = BigUint::from(i * num_entries + j);
+            let val = match state.get(&key) {
+                Some(val) => val * factor,
+                None => Complex::zero(),
+            };
+            row.push(val);
+        }
+        matrix.push(row);
+    }
+
+    matrix
+}

compiler/qsc_eval/src/output.rs

@@ -16,6 +16,11 @@ pub trait Receiver {
     /// This will return an error if handling the output fails.
     fn state(&mut self, state: Vec<(BigUint, Complex64)>, qubit_count: usize) -> Result<(), Error>;
 
+    /// Receive matrix output
+    /// # Errors
+    /// This will return an error if handling the output fails.
+    fn matrix(&mut self, matrix: Vec<Vec<Complex64>>) -> Result<(), Error>;
+
     /// Receive generic message output
     /// # Errors
     /// This will return an error if handling the output fails.
@@ -47,6 +52,15 @@ impl<'a> Receiver for GenericReceiver<'a> {
         Ok(())
     }
 
+    fn matrix(&mut self, matrix: Vec<Vec<Complex64>>) -> Result<(), Error> {
+        writeln!(self.writer, "MATRIX:").map_err(|_| Error)?;
+        for row in matrix {
+            let row_str = row.iter().map(fmt_complex).collect::<Vec<_>>().join(" ");
+            writeln!(self.writer, "{row_str}").map_err(|_| Error)?;
+        }
+        Ok(())
+    }
+
     fn message(&mut self, msg: &str) -> Result<(), Error> {
         writeln!(self.writer, "{msg}").map_err(|_| Error)
     }
@@ -86,6 +100,15 @@ impl<'a> Receiver for CursorReceiver<'a> {
         Ok(())
     }
 
+    fn matrix(&mut self, matrix: Vec<Vec<Complex64>>) -> Result<(), Error> {
+        writeln!(self.cursor, "MATRIX:").map_err(|_| Error)?;
+        for row in matrix {
+            let row_str = row.iter().map(fmt_complex).collect::<Vec<_>>().join(" ");
+            writeln!(self.cursor, "{row_str}").map_err(|_| Error)?;
+        }
+        Ok(())
+    }
+
     fn message(&mut self, msg: &str) -> Result<(), Error> {
         writeln!(self.cursor, "{msg}").map_err(|_| Error)
     }