fixing .reset() to give state, making .probabilities/probability() av…

python/pecos-rslib/pecos_rslib.pyi

@@ -711,31 +711,61 @@ class num:
 # =============================================================================
 # Quantum Simulators
 # =============================================================================
+class TableauWrapper:
+    """Wrapper for accessing stabilizer/destabilizer tableaus from simulators."""
+
+    def __init__(self, sim: object, *, is_stab: bool) -> None: ...
+    def print_tableau(self, *, verbose: bool = False) -> list[str]: ...
+    @property
+    def col_x(self) -> list[list[int]]: ...
+    @property
+    def col_z(self) -> list[list[int]]: ...
+    @property
+    def row_x(self) -> list[list[int]]: ...
+    @property
+    def row_z(self) -> list[list[int]]: ...
+
+class GateBindingsDict:
+    """Special dict that delegates gate lookups to run_gate()."""
+
+    def __init__(self, sim: object) -> None: ...
+    def __getitem__(self, key: str) -> object: ...
+    def __setitem__(self, key: str, value: object) -> None: ...
+    def __contains__(self, key: str) -> bool: ...
+    def get(self, key: str, default: object | None = None) -> object: ...
+    def __len__(self) -> int: ...
+    def keys(self) -> list[str]: ...
+
 class SparseSim:
     """Sparse stabilizer simulator."""
 
     def __init__(self, num_qubits: int) -> None: ...
+    def reset(self) -> SparseSim: ...
     @property
     def num_qubits(self) -> int: ...
     @property
-    def stabs(self) -> object: ...
+    def stabs(self) -> TableauWrapper: ...
+    @property
+    def destabs(self) -> TableauWrapper: ...
     @property
-    def destabs(self) -> object: ...
+    def gens(self) -> tuple[TableauWrapper, TableauWrapper]: ...
     @property
-    def gens(self) -> tuple[object, object]: ...
+    def bindings(self) -> GateBindingsDict: ...
     def __repr__(self) -> str: ...
 
 class SparseSimCpp:
     """C++ sparse simulator bindings."""
 
     def __init__(self, num_qubits: int) -> None: ...
+    def reset(self) -> SparseSimCpp: ...
     @property
     def num_qubits(self) -> int: ...
 
 class StateVec:
     """Rust state vector simulator."""
 
     def __init__(self, num_qubits: int) -> None: ...
+    def reset(self) -> StateVec: ...
     @property
     def num_qubits(self) -> int: ...
     @property
@@ -749,6 +779,7 @@ class Qulacs:
     """Rust Qulacs state vector simulator."""
 
     def __init__(self, num_qubits: int, *, seed: int | None = None) -> None: ...
+    def reset(self) -> Qulacs: ...
     @property
     def num_qubits(self) -> int: ...
     @property
@@ -765,13 +796,15 @@ class QuestStateVec:
     """QuEST state vector simulator."""
 
     def __init__(self, num_qubits: int) -> None: ...
+    def reset(self) -> QuestStateVec: ...
     @property
     def num_qubits(self) -> int: ...
 
 class QuestDensityMatrix:
     """QuEST density matrix simulator."""
 
     def __init__(self, num_qubits: int) -> None: ...
+    def reset(self) -> QuestDensityMatrix: ...
     @property
     def num_qubits(self) -> int: ...
 

python/pecos-rslib/src/cpp_sparse_sim_bindings.rs

@@ -38,8 +38,9 @@ impl PySparseSimCpp {
         self.inner.set_seed(seed);
     }
 
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     fn __repr__(&self) -> String {

python/pecos-rslib/src/quest_bindings.rs

@@ -50,8 +50,9 @@ impl QuestStateVec {
     }
 
     /// Resets the quantum state to the all-zero state
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     /// Prepares a computational basis state
@@ -444,8 +445,9 @@ impl QuestDensityMatrix {
     }
 
     /// Resets the quantum state to the all-zero state
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     /// Prepares a computational basis state

python/pecos-rslib/src/qulacs_bindings.rs

@@ -135,8 +135,9 @@ impl PyQulacs {
     }
 
     /// Resets the quantum state to the all-zero state
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     /// Executes a single-qubit gate based on the provided symbol and location

python/pecos-rslib/src/simulator_utils.rs

@@ -24,6 +24,14 @@ use std::collections::HashMap;
 
 use crate::sparse_stab_bindings::adjust_tableau_string;
 
+/// Raw generators data: `(col_x, col_z, row_x, row_z)`.
+pub type GensData = (
+    Vec<Vec<usize>>,
+    Vec<Vec<usize>>,
+    Vec<Vec<usize>>,
+    Vec<Vec<usize>>,
+);
+
 /// Special dict that delegates all gate lookups to Rust's `run_gate()`.
 ///
 /// This provides backwards compatibility for code that accesses sim.bindings[`gate_name`].
@@ -178,6 +186,33 @@ impl TableauWrapper {
 
         Ok(lines)
     }
+
+    /// Helper to get raw gens data from the simulator.
+    fn get_gens_data(&self, py: Python<'_>) -> PyResult<GensData> {
+        self.sim
+            .call_method1(py, "_gens_data", (self.is_stab,))?
+            .extract(py)
+    }
+
+    #[getter]
+    fn col_x(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
+        Ok(self.get_gens_data(py)?.0)
+    }
+
+    #[getter]
+    fn col_z(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
+        Ok(self.get_gens_data(py)?.1)
+    }
+
+    #[getter]
+    fn row_x(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
+        Ok(self.get_gens_data(py)?.2)
+    }
+
+    #[getter]
+    fn row_z(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
+        Ok(self.get_gens_data(py)?.3)
+    }
 }
 
 /// Register the simulator utils module

python/pecos-rslib/src/sparse_sim.rs

@@ -29,8 +29,9 @@ impl SparseSim {
         }
     }
 
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     fn __repr__(&self) -> String {
@@ -458,6 +459,24 @@ impl SparseSim {
         Ok(())
     }
 
+    /// Returns the raw gens data (`col_x`, `col_z`, `row_x`, `row_z`) for stabs or destabs.
+    fn _gens_data(&self, is_stab: bool) -> crate::simulator_utils::GensData {
+        let gens = if is_stab {
+            self.inner.stabs()
+        } else {
+            self.inner.destabs()
+        };
+        let to_vecs = |sets: &[VecSet<usize>]| -> Vec<Vec<usize>> {
+            sets.iter().map(|s| s.elements().to_vec()).collect()
+        };
+        (
+            to_vecs(&gens.col_x),
+            to_vecs(&gens.col_z),
+            to_vecs(&gens.row_x),
+            to_vecs(&gens.row_z),
+        )
+    }
+
     #[getter]
     fn bindings(slf: PyRef<'_, Self>) -> PyResult<crate::simulator_utils::GateBindingsDict> {
         // Create a Rust GateBindingsDict directly

python/pecos-rslib/src/sparse_stab_bindings.rs

@@ -29,8 +29,9 @@ impl PySparseSim {
         }
     }
 
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     #[getter]
@@ -613,6 +614,24 @@ impl PySparseSim {
         })
     }
 
+    /// Returns the raw gens data (`col_x`, `col_z`, `row_x`, `row_z`) for stabs or destabs.
+    fn _gens_data(&self, is_stab: bool) -> crate::simulator_utils::GensData {
+        let gens = if is_stab {
+            self.inner.stabs()
+        } else {
+            self.inner.destabs()
+        };
+        let to_vecs = |sets: &[VecSet<usize>]| -> Vec<Vec<usize>> {
+            sets.iter().map(|s| s.elements().to_vec()).collect()
+        };
+        (
+            to_vecs(&gens.col_x),
+            to_vecs(&gens.col_z),
+            to_vecs(&gens.row_x),
+            to_vecs(&gens.row_z),
+        )
+    }
+
     #[getter]
     fn bindings(slf: PyRef<'_, Self>) -> PyResult<crate::simulator_utils::GateBindingsDict> {
         // Create a Rust GateBindingsDict directly

python/pecos-rslib/src/state_vec_bindings.rs

@@ -42,8 +42,9 @@ impl PyStateVec {
     }
 
     /// Resets the quantum state to the all-zero state
-    fn reset(&mut self) {
-        self.inner.reset();
+    fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
+        slf.inner.reset();
+        slf
     }
 
     /// Executes a single-qubit gate based on the provided symbol and location

python/quantum-pecos/src/pecos/simulators/qulacs/state.py

@@ -70,3 +70,12 @@ def vector(self) -> Array:
             [complex(real, imag) for real, imag in complex_tuples],
             dtype="complex",
         )
+
+    @property
+    def probabilities(self) -> list[float]:
+        """Get the probability distribution over all basis states.
+
+        Returns:
+            List of probabilities for each computational basis state.
+        """
+        return self.qulacs_state.probabilities

python/quantum-pecos/src/pecos/simulators/statevec/state.py

@@ -56,6 +56,26 @@ def vector(self) -> Array:  # noqa: F821 - Array is a forward reference
         """
         return self.backend.vector_big_endian()
 
+    @property
+    def probabilities(self) -> Array:  # noqa: F821 - Array is a forward reference
+        """Get the probability distribution over all basis states.
+
+        Returns:
+            Array of probabilities for each computational basis state.
+        """
+        return self.backend.probabilities
+
+    def probability(self, basis_state: int) -> float:
+        """Get the probability of a specific computational basis state.
+
+        Args:
+            basis_state: The index of the basis state.
+
+        Returns:
+            The probability of measuring the given basis state.
+        """
+        return self.backend.probability(basis_state)
+
     def reset(self) -> StateVec:
         """Resets the quantum state to the all-zero state."""
         self.backend.reset()