Track PauliLindbladMap generators during parity_sample (#16072)

* feat: add PauliLindbladMap.parity_sample_generators

This adds a new method to the `PauliLindbladMap` which is identical to
its `parity_sample` except for also tracking and returning which of the
map's generators resulted in the sampled Pauli operator and sign.

* refactor: avoid code duplication of parity_sample

* refactor: clean up PauliLindbladMap.parity_sample unittests

* refactor: remove duplicate PauliLindbladMap.parity_sample

With the refactoring of this PR, the Rust-level
PauliLindbladMap.parity_sample method is no longer needed and can be
fully replaced by PauliLindbladMap.parity_sample_generators for internal
Rust API calls. Therefore, this commit removes that method and relies
on the public pyo3 API to expose the PauliLindbladMap.parity_sample
method.

* refactor: update naming

Author: mrossinek

crates/quantum_info/src/pauli_lindblad_map/pauli_lindblad_map_class.rs

@@ -11,7 +11,7 @@
 // that they have been altered from the originals.
 
 use hashbrown::HashSet;
-use numpy::{PyArray1, PyArrayMethods, ToPyArray};
+use numpy::{PyArray1, PyArray2, PyArrayMethods, ToPyArray};
 use pyo3::{
     IntoPyObjectExt, PyErr,
     exceptions::{PyTypeError, PyValueError},
@@ -386,13 +386,19 @@ impl PauliLindbladMap {
 
     /// Sample sign and Pauli operator pairs from the map.
     /// Note that here the "sign" bool is interpreted as the exponent of (-1)^b.
-    pub fn parity_sample(
+    #[allow(clippy::type_complexity)]
+    pub fn parity_sample_with_history(
         &self,
         num_samples: u64,
         seed: Option<u64>,
         scale: Option<f64>,
         local_scale: Option<Vec<f64>>,
-    ) -> (Vec<bool>, QubitSparsePauliList) {
+    ) -> (
+        Vec<bool>,
+        QubitSparsePauliList,
+        Vec<Vec<bool>>,
+        Vec<Vec<bool>>,
+    ) {
         let mut rng = match seed {
             Some(seed) => Pcg64Mcg::seed_from_u64(seed),
             None => Pcg64Mcg::try_from_rng(&mut SysRng).unwrap(),
@@ -422,13 +428,18 @@ impl PauliLindbladMap {
         };
         let mut random_signs = Vec::with_capacity(num_samples as usize);
         let mut random_paulis = QubitSparsePauliList::empty(self.num_qubits());
+        let mut pauli_history = Vec::with_capacity(num_samples as usize);
+        let mut signs_history = Vec::with_capacity(num_samples as usize);
 
         for _ in 0..num_samples {
             let mut random_sign = false;
             let mut random_pauli = QubitSparsePauli::identity(self.num_qubits());
+            let mut inner_pauli_history = vec![false; self.qubit_sparse_pauli_list.num_terms()];
+            let mut inner_signs_history = vec![false; self.qubit_sparse_pauli_list.num_terms()];
 
-            for ((probability, generator), non_negative_rate) in probabilities
+            for (((idx, probability), generator), non_negative_rate) in probabilities
                 .iter()
+                .enumerate()
                 .zip(self.qubit_sparse_pauli_list.iter())
                 .zip(non_negative_rates.iter())
             {
@@ -437,16 +448,21 @@ impl PauliLindbladMap {
                     random_pauli = random_pauli.compose(&generator.to_term()).unwrap();
                     // if rate is negative, flip random_sign
                     random_sign = random_sign == *non_negative_rate;
+                    // keep track of sampled generator
+                    inner_pauli_history[idx] = true;
+                    inner_signs_history[idx] = *non_negative_rate;
                 }
             }
 
             random_signs.push(random_sign);
             random_paulis
                 .add_qubit_sparse_pauli(random_pauli.view())
                 .unwrap();
+            pauli_history.push(inner_pauli_history);
+            signs_history.push(inner_signs_history);
         }
 
-        (random_signs, random_paulis)
+        (random_signs, random_paulis, pauli_history, signs_history)
     }
 
     /// Reduce the map to its canonical form.
@@ -1578,7 +1594,8 @@ impl PyPauliLindbladMap {
         seed: Option<u64>,
     ) -> PyResult<Bound<'py, PyTuple>> {
         let inner = self.inner.read().map_err(|_| InnerReadError)?;
-        let (signs, paulis) = py.detach(|| inner.parity_sample(num_samples, seed, None, None));
+        let (signs, paulis, _, _) =
+            py.detach(|| inner.parity_sample_with_history(num_samples, seed, None, None));
 
         let signs = PyArray1::from_vec(py, signs.iter().map(|b| !b).collect());
         let paulis = paulis.into_pyobject(py).unwrap();
@@ -1630,15 +1647,53 @@ impl PyPauliLindbladMap {
         local_scale: Option<Vec<f64>>,
     ) -> PyResult<Bound<'py, PyTuple>> {
         let inner = self.inner.read().map_err(|_| InnerReadError)?;
-        let (signs, paulis) =
-            py.detach(|| inner.parity_sample(num_samples, seed, scale, local_scale));
+        let (signs, paulis, _, _) =
+            py.detach(|| inner.parity_sample_with_history(num_samples, seed, scale, local_scale));
 
         let signs = PyArray1::from_vec(py, signs);
         let paulis = paulis.into_pyobject(py).unwrap();
 
         (signs, paulis).into_pyobject(py)
     }
 
+    /// Sample sign and Pauli operator pairs from the map, preserving the sampled generators.
+    ///
+    /// This method is identical to :meth:`parity_sample` except for also returning the information
+    /// which :meth:`generators` were actually sampled to yield the final Pauli operator and sign.
+    ///
+    ///
+    /// Args:
+    ///     num_samples (int): Number of samples to draw.
+    ///     seed (int): Random seed.
+    ///     scale (float): Scale to apply to all rates.
+    ///     local_scale (list[float]): Local scale to apply on a term-by-term basis.
+    ///
+    /// Returns:
+    ///     signs, qubit_sparse_pauli_list, pauli_history, signs_history: The boolean array of
+    ///         signs, the list of qubit sparse paulis, the two dimensional boolean array
+    ///         indicating which :meth:`generators` were sampled and the two dimensional boolean
+    ///         array indicating their signs.
+    #[pyo3(signature = (num_samples, seed=None, scale=None, local_scale=None))]
+    pub fn parity_sample_with_history<'py>(
+        &self,
+        py: Python<'py>,
+        num_samples: u64,
+        seed: Option<u64>,
+        scale: Option<f64>,
+        local_scale: Option<Vec<f64>>,
+    ) -> PyResult<Bound<'py, PyTuple>> {
+        let inner = self.inner.read().map_err(|_| InnerReadError)?;
+        let (signs, paulis, pauli_history, signs_history) =
+            py.detach(|| inner.parity_sample_with_history(num_samples, seed, scale, local_scale));
+
+        let signs = PyArray1::from_vec(py, signs);
+        let paulis = paulis.into_pyobject(py).unwrap();
+        let pauli_history = PyArray2::from_vec2(py, &pauli_history).unwrap();
+        let signs_history = PyArray2::from_vec2(py, &signs_history).unwrap();
+
+        (signs, paulis, pauli_history, signs_history).into_pyobject(py)
+    }
+
     /// For :class:`.PauliLindbladMap` instances with purely non-negative rates, sample Pauli
     /// operators from the map. If the map has negative rates, use
     /// :meth:`.PauliLindbladMap.parity_sample`.
@@ -1683,7 +1738,8 @@ impl PyPauliLindbladMap {
             }
         }
 
-        let (_, paulis) = py.detach(|| inner.parity_sample(num_samples, seed, None, None));
+        let (_, paulis, _, _) =
+            py.detach(|| inner.parity_sample_with_history(num_samples, seed, None, None));
 
         paulis.into_pyobject(py)
     }

releasenotes/notes/paulilindbladmap-parity-sample-generators-d4a439235327e1fc.yaml

@@ -0,0 +1,6 @@
+---
+features_quantum_info:
+  - Added the :meth:`.PauliLindbladMap.parity_sample_with_history` method which is
+    identical to :meth:`.PauliLindbladMap.parity_sample` except for also tracking
+    and returning the underlying history of Paulis and signs that resulted in the
+    final sampled terms.

test/python/quantum_info/test_pauli_lindblad_map.py

@@ -458,7 +458,7 @@ def test_attributes_immutable(self):
         with self.assertRaisesRegex(ValueError, "assignment destination is read-only"):
             pauli_lindblad_map.rates[0] = 1.0
 
-    def test_generators(self):
+    def test_with_history(self):
         """Test that generators() method returns the same result as
         get_qubit_sparse_pauli_list_copy()."""
         pauli_lindblad_map = PauliLindbladMap.from_list([("IIXIZ", 2), ("IIZIX", 3)])
@@ -1239,134 +1239,100 @@ def test_signed_sample(self):
         self.assertEqual(len(qubit_sparse_pauli_list), 5)
         self.assertEqual(len(signs), 5)
 
-    def test_parity_sample(self):
-
-        # test all negative rates
-        pauli_lindblad_map = PauliLindbladMap([("X", -1.0), ("Y", -0.5)])
-        probs = pauli_lindblad_map.probabilities()
+    def _assert_parity_sample_results(
+        self,
+        plm_to_sample,
+        expected_signs,
+        num_samples=10000,
+        seed=12312,
+        scale=None,
+        local_scale=None,
+        plm_reference=None,
+    ):
+        if plm_reference is None:
+            plm_reference = plm_to_sample
+
+        generators = plm_reference.generators()
+        probs = plm_reference.probabilities()
         probs_dict = {
             "I": probs[0] * probs[1],
             "X": (1 - probs[0]) * probs[1],
             "Y": probs[0] * (1 - probs[1]),
             "Z": (1 - probs[0]) * (1 - probs[1]),
         }
-        expected_signs = {"I": False, "X": True, "Y": True, "Z": False}
-
-        num_samples = 10000
-        signs, qubit_sparse_pauli_list = pauli_lindblad_map.parity_sample(num_samples, 12312)
 
+        signs, qubit_sparse_pauli_list, sampled_with_history, sampled_signs = (
+            plm_to_sample.parity_sample_with_history(
+                num_samples, seed, scale=scale, local_scale=local_scale
+            )
+        )
         counts = {"I": 0, "X": 0, "Y": 0, "Z": 0}
-        for sign, q in zip(signs, qubit_sparse_pauli_list):
+        for sign, q, _gens, _signs in zip(
+            signs, qubit_sparse_pauli_list, sampled_with_history, sampled_signs
+        ):
             for symbol in counts:
                 if q == QubitSparsePauli(symbol):
                     counts[symbol] += 1
                     self.assertEqual(expected_signs[symbol], sign)
 
-        for symbol, count in counts.items():
-            self.assertTrue(np.abs(count / num_samples - probs_dict[symbol]) < 1e-2)
+            sampled_sign = False
+            sampled_pauli = QubitSparsePauli.from_label("I")
+            for _i, (_g, _s) in enumerate(zip(_gens, _signs)):
+                if _g:
+                    sampled_sign = sampled_sign == _s
+                    sampled_pauli = sampled_pauli.compose(generators[_i])
 
-        # test all positive rates
-        pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", 0.5)])
-        probs = pauli_lindblad_map.probabilities()
-        probs_dict = {
-            "I": probs[0] * probs[1],
-            "X": (1 - probs[0]) * probs[1],
-            "Y": probs[0] * (1 - probs[1]),
-            "Z": (1 - probs[0]) * (1 - probs[1]),
-        }
-        expected_signs = {"I": False, "X": False, "Y": False, "Z": False}
+            self.assertEqual(sign, sampled_sign)
+            self.assertEqual(q, sampled_pauli)
 
-        num_samples = 10000
-        signs, qubit_sparse_pauli_list = pauli_lindblad_map.parity_sample(num_samples, 12312)
-
-        counts = {"I": 0, "X": 0, "Y": 0, "Z": 0}
-        for sign, q in zip(signs, qubit_sparse_pauli_list):
-            for symbol in counts:
-                if q == QubitSparsePauli(symbol):
-                    counts[symbol] += 1
-                    self.assertEqual(expected_signs[symbol], sign)
         for symbol, count in counts.items():
             self.assertTrue(np.abs(count / num_samples - probs_dict[symbol]) < 1e-2)
 
-        # test mix of positive and negative rates
-        pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
-        probs = pauli_lindblad_map.probabilities()
-        probs_dict = {
-            "I": probs[0] * probs[1],
-            "X": (1 - probs[0]) * probs[1],
-            "Y": probs[0] * (1 - probs[1]),
-            "Z": (1 - probs[0]) * (1 - probs[1]),
-        }
-        expected_signs = {"I": False, "X": False, "Y": True, "Z": True}
-
-        num_samples = 10000
-        signs, qubit_sparse_pauli_list = pauli_lindblad_map.parity_sample(num_samples, 12312)
+    def test_parity_sample(self):
 
-        counts = {"I": 0, "X": 0, "Y": 0, "Z": 0}
-        for sign, q in zip(signs, qubit_sparse_pauli_list):
-            for symbol in counts:
-                if q == QubitSparsePauli(symbol):
-                    counts[symbol] += 1
-                    self.assertEqual(expected_signs[symbol], sign)
+        with self.subTest("all negative rates"):
+            pauli_lindblad_map = PauliLindbladMap([("X", -1.0), ("Y", -0.5)])
+            expected_signs = {"I": False, "X": True, "Y": True, "Z": False}
 
-        for symbol, count in counts.items():
-            self.assertTrue(np.abs(count / num_samples - probs_dict[symbol]) < 1e-2)
+            self._assert_parity_sample_results(pauli_lindblad_map, expected_signs)
 
-        # test scale with mix of positive and negative rates
-        pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
-        pauli_lindblad_map_downscaled = PauliLindbladMap([("X", 0.5), ("Y", -0.25)])
-        probs = pauli_lindblad_map.probabilities()
-        probs_dict = {
-            "I": probs[0] * probs[1],
-            "X": (1 - probs[0]) * probs[1],
-            "Y": probs[0] * (1 - probs[1]),
-            "Z": (1 - probs[0]) * (1 - probs[1]),
-        }
-        expected_signs = {"I": False, "X": False, "Y": True, "Z": True}
+        with self.subTest("all positive rates"):
+            pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", 0.5)])
+            expected_signs = {"I": False, "X": False, "Y": False, "Z": False}
 
-        num_samples = 10000
-        signs, qubit_sparse_pauli_list = pauli_lindblad_map_downscaled.parity_sample(
-            num_samples, 12312, scale=2.0
-        )
+            self._assert_parity_sample_results(pauli_lindblad_map, expected_signs)
 
-        counts = {"I": 0, "X": 0, "Y": 0, "Z": 0}
-        for sign, q in zip(signs, qubit_sparse_pauli_list):
-            for symbol in counts:
-                if q == QubitSparsePauli(symbol):
-                    counts[symbol] += 1
-                    self.assertEqual(expected_signs[symbol], sign)
+        with self.subTest("mix of positive and negative rates"):
+            pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
+            expected_signs = {"I": False, "X": False, "Y": True, "Z": True}
 
-        for symbol, count in counts.items():
-            self.assertTrue(np.abs(count / num_samples - probs_dict[symbol]) < 1e-2)
+            self._assert_parity_sample_results(pauli_lindblad_map, expected_signs)
 
-        # test local_scale with mix of positive and negative rates
-        pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
-        pauli_lindblad_map_downscaled = PauliLindbladMap([("X", 1.0), ("Y", -0.25)])
-        probs = pauli_lindblad_map.probabilities()
-        probs_dict = {
-            "I": probs[0] * probs[1],
-            "X": (1 - probs[0]) * probs[1],
-            "Y": probs[0] * (1 - probs[1]),
-            "Z": (1 - probs[0]) * (1 - probs[1]),
-        }
-        expected_signs = {"I": False, "X": False, "Y": True, "Z": True}
+        with self.subTest("scale with mix of positive and negative rates"):
+            pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
+            pauli_lindblad_map_downscaled = PauliLindbladMap([("X", 0.5), ("Y", -0.25)])
+            expected_signs = {"I": False, "X": False, "Y": True, "Z": True}
 
-        num_samples = 10000
-        signs, qubit_sparse_pauli_list = pauli_lindblad_map_downscaled.parity_sample(
-            num_samples, 12312, local_scale=[1.0, 2.0]
-        )
+            self._assert_parity_sample_results(
+                pauli_lindblad_map_downscaled,
+                expected_signs,
+                scale=2.0,
+                plm_reference=pauli_lindblad_map,
+            )
 
-        counts = {"I": 0, "X": 0, "Y": 0, "Z": 0}
-        for sign, q in zip(signs, qubit_sparse_pauli_list):
-            for symbol in counts:
-                if q == QubitSparsePauli(symbol):
-                    counts[symbol] += 1
-                    self.assertEqual(expected_signs[symbol], sign)
+        with self.subTest("local_scale with mix of positive and negative rates"):
+            pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
+            pauli_lindblad_map_downscaled = PauliLindbladMap([("X", 1.0), ("Y", -0.25)])
 
-        for symbol, count in counts.items():
-            self.assertTrue(np.abs(count / num_samples - probs_dict[symbol]) < 1e-2)
+            self._assert_parity_sample_results(
+                pauli_lindblad_map_downscaled,
+                expected_signs,
+                local_scale=[1.0, 2.0],
+                plm_reference=pauli_lindblad_map,
+            )
 
-        # test callable without seed
+    def test_parity_sample_without_seed(self):
+        pauli_lindblad_map = PauliLindbladMap([("X", 1.0), ("Y", -0.5)])
         signs, qubit_sparse_pauli_list = pauli_lindblad_map.parity_sample(5)
         self.assertTrue(isinstance(qubit_sparse_pauli_list, QubitSparsePauliList))
         self.assertEqual(len(qubit_sparse_pauli_list), 5)