Skip to content

Commit 5a3fd4d

Browse files
ciaranraciaranra
authored
fixing .reset() to give state, making .probabilities/probability() available (#250)
1 parent c291282 commit 5a3fd4d

12 files changed

Lines changed: 328 additions & 17 deletions

File tree

‎python/pecos-rslib/pecos_rslib.pyi‎

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -711,31 +711,61 @@ class num:
711711
# =============================================================================
712712
# Quantum Simulators
713713
# =============================================================================
714+
class TableauWrapper:
715+
"""Wrapper for accessing stabilizer/destabilizer tableaus from simulators."""
716+
717+
def __init__(self, sim: object, *, is_stab: bool) -> None: ...
718+
def print_tableau(self, *, verbose: bool = False) -> list[str]: ...
719+
@property
720+
def col_x(self) -> list[list[int]]: ...
721+
@property
722+
def col_z(self) -> list[list[int]]: ...
723+
@property
724+
def row_x(self) -> list[list[int]]: ...
725+
@property
726+
def row_z(self) -> list[list[int]]: ...
727+
728+
class GateBindingsDict:
729+
"""Special dict that delegates gate lookups to run_gate()."""
730+
731+
def __init__(self, sim: object) -> None: ...
732+
def __getitem__(self, key: str) -> object: ...
733+
def __setitem__(self, key: str, value: object) -> None: ...
734+
def __contains__(self, key: str) -> bool: ...
735+
def get(self, key: str, default: object | None = None) -> object: ...
736+
def __len__(self) -> int: ...
737+
def keys(self) -> list[str]: ...
738+
714739
class SparseSim:
715740
"""Sparse stabilizer simulator."""
716741

717742
def __init__(self, num_qubits: int) -> None: ...
743+
def reset(self) -> SparseSim: ...
718744
@property
719745
def num_qubits(self) -> int: ...
720746
@property
721-
def stabs(self) -> object: ...
747+
def stabs(self) -> TableauWrapper: ...
748+
@property
749+
def destabs(self) -> TableauWrapper: ...
722750
@property
723-
def destabs(self) -> object: ...
751+
def gens(self) -> tuple[TableauWrapper, TableauWrapper]: ...
724752
@property
725-
def gens(self) -> tuple[object, object]: ...
753+
def bindings(self) -> GateBindingsDict: ...
726754
def __repr__(self) -> str: ...
727755

728756
class SparseSimCpp:
729757
"""C++ sparse simulator bindings."""
730758

731759
def __init__(self, num_qubits: int) -> None: ...
760+
def reset(self) -> SparseSimCpp: ...
732761
@property
733762
def num_qubits(self) -> int: ...
734763

735764
class StateVec:
736765
"""Rust state vector simulator."""
737766

738767
def __init__(self, num_qubits: int) -> None: ...
768+
def reset(self) -> StateVec: ...
739769
@property
740770
def num_qubits(self) -> int: ...
741771
@property
@@ -749,6 +779,7 @@ class Qulacs:
749779
"""Rust Qulacs state vector simulator."""
750780

751781
def __init__(self, num_qubits: int, *, seed: int | None = None) -> None: ...
782+
def reset(self) -> Qulacs: ...
752783
@property
753784
def num_qubits(self) -> int: ...
754785
@property
@@ -765,13 +796,15 @@ class QuestStateVec:
765796
"""QuEST state vector simulator."""
766797

767798
def __init__(self, num_qubits: int) -> None: ...
799+
def reset(self) -> QuestStateVec: ...
768800
@property
769801
def num_qubits(self) -> int: ...
770802

771803
class QuestDensityMatrix:
772804
"""QuEST density matrix simulator."""
773805

774806
def __init__(self, num_qubits: int) -> None: ...
807+
def reset(self) -> QuestDensityMatrix: ...
775808
@property
776809
def num_qubits(self) -> int: ...
777810

‎python/pecos-rslib/src/cpp_sparse_sim_bindings.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@ impl PySparseSimCpp {
3838
self.inner.set_seed(seed);
3939
}
4040

41-
fn reset(&mut self) {
42-
self.inner.reset();
41+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
42+
slf.inner.reset();
43+
slf
4344
}
4445

4546
fn __repr__(&self) -> String {

‎python/pecos-rslib/src/quest_bindings.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,9 @@ impl QuestStateVec {
5050
}
5151

5252
/// Resets the quantum state to the all-zero state
53-
fn reset(&mut self) {
54-
self.inner.reset();
53+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
54+
slf.inner.reset();
55+
slf
5556
}
5657

5758
/// Prepares a computational basis state
@@ -444,8 +445,9 @@ impl QuestDensityMatrix {
444445
}
445446

446447
/// Resets the quantum state to the all-zero state
447-
fn reset(&mut self) {
448-
self.inner.reset();
448+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
449+
slf.inner.reset();
450+
slf
449451
}
450452

451453
/// Prepares a computational basis state

‎python/pecos-rslib/src/qulacs_bindings.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,8 +135,9 @@ impl PyQulacs {
135135
}
136136

137137
/// Resets the quantum state to the all-zero state
138-
fn reset(&mut self) {
139-
self.inner.reset();
138+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
139+
slf.inner.reset();
140+
slf
140141
}
141142

142143
/// Executes a single-qubit gate based on the provided symbol and location

‎python/pecos-rslib/src/simulator_utils.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,14 @@ use std::collections::HashMap;
2424

2525
use crate::sparse_stab_bindings::adjust_tableau_string;
2626

27+
/// Raw generators data: `(col_x, col_z, row_x, row_z)`.
28+
pub type GensData = (
29+
Vec<Vec<usize>>,
30+
Vec<Vec<usize>>,
31+
Vec<Vec<usize>>,
32+
Vec<Vec<usize>>,
33+
);
34+
2735
/// Special dict that delegates all gate lookups to Rust's `run_gate()`.
2836
///
2937
/// This provides backwards compatibility for code that accesses sim.bindings[`gate_name`].
@@ -178,6 +186,33 @@ impl TableauWrapper {
178186

179187
Ok(lines)
180188
}
189+
190+
/// Helper to get raw gens data from the simulator.
191+
fn get_gens_data(&self, py: Python<'_>) -> PyResult<GensData> {
192+
self.sim
193+
.call_method1(py, "_gens_data", (self.is_stab,))?
194+
.extract(py)
195+
}
196+
197+
#[getter]
198+
fn col_x(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
199+
Ok(self.get_gens_data(py)?.0)
200+
}
201+
202+
#[getter]
203+
fn col_z(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
204+
Ok(self.get_gens_data(py)?.1)
205+
}
206+
207+
#[getter]
208+
fn row_x(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
209+
Ok(self.get_gens_data(py)?.2)
210+
}
211+
212+
#[getter]
213+
fn row_z(&self, py: Python<'_>) -> PyResult<Vec<Vec<usize>>> {
214+
Ok(self.get_gens_data(py)?.3)
215+
}
181216
}
182217

183218
/// Register the simulator utils module

‎python/pecos-rslib/src/sparse_sim.rs‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ impl SparseSim {
2929
}
3030
}
3131

32-
fn reset(&mut self) {
33-
self.inner.reset();
32+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
33+
slf.inner.reset();
34+
slf
3435
}
3536

3637
fn __repr__(&self) -> String {
@@ -458,6 +459,24 @@ impl SparseSim {
458459
Ok(())
459460
}
460461

462+
/// Returns the raw gens data (`col_x`, `col_z`, `row_x`, `row_z`) for stabs or destabs.
463+
fn _gens_data(&self, is_stab: bool) -> crate::simulator_utils::GensData {
464+
let gens = if is_stab {
465+
self.inner.stabs()
466+
} else {
467+
self.inner.destabs()
468+
};
469+
let to_vecs = |sets: &[VecSet<usize>]| -> Vec<Vec<usize>> {
470+
sets.iter().map(|s| s.elements().to_vec()).collect()
471+
};
472+
(
473+
to_vecs(&gens.col_x),
474+
to_vecs(&gens.col_z),
475+
to_vecs(&gens.row_x),
476+
to_vecs(&gens.row_z),
477+
)
478+
}
479+
461480
#[getter]
462481
fn bindings(slf: PyRef<'_, Self>) -> PyResult<crate::simulator_utils::GateBindingsDict> {
463482
// Create a Rust GateBindingsDict directly

‎python/pecos-rslib/src/sparse_stab_bindings.rs‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ impl PySparseSim {
2929
}
3030
}
3131

32-
fn reset(&mut self) {
33-
self.inner.reset();
32+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
33+
slf.inner.reset();
34+
slf
3435
}
3536

3637
#[getter]
@@ -613,6 +614,24 @@ impl PySparseSim {
613614
})
614615
}
615616

617+
/// Returns the raw gens data (`col_x`, `col_z`, `row_x`, `row_z`) for stabs or destabs.
618+
fn _gens_data(&self, is_stab: bool) -> crate::simulator_utils::GensData {
619+
let gens = if is_stab {
620+
self.inner.stabs()
621+
} else {
622+
self.inner.destabs()
623+
};
624+
let to_vecs = |sets: &[VecSet<usize>]| -> Vec<Vec<usize>> {
625+
sets.iter().map(|s| s.elements().to_vec()).collect()
626+
};
627+
(
628+
to_vecs(&gens.col_x),
629+
to_vecs(&gens.col_z),
630+
to_vecs(&gens.row_x),
631+
to_vecs(&gens.row_z),
632+
)
633+
}
634+
616635
#[getter]
617636
fn bindings(slf: PyRef<'_, Self>) -> PyResult<crate::simulator_utils::GateBindingsDict> {
618637
// Create a Rust GateBindingsDict directly

‎python/pecos-rslib/src/state_vec_bindings.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,9 @@ impl PyStateVec {
4242
}
4343

4444
/// Resets the quantum state to the all-zero state
45-
fn reset(&mut self) {
46-
self.inner.reset();
45+
fn reset(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
46+
slf.inner.reset();
47+
slf
4748
}
4849

4950
/// Executes a single-qubit gate based on the provided symbol and location

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,3 +70,12 @@ def vector(self) -> Array:
7070
[complex(real, imag) for real, imag in complex_tuples],
7171
dtype="complex",
7272
)
73+
74+
@property
75+
def probabilities(self) -> list[float]:
76+
"""Get the probability distribution over all basis states.
77+
78+
Returns:
79+
List of probabilities for each computational basis state.
80+
"""
81+
return self.qulacs_state.probabilities

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,26 @@ def vector(self) -> Array: # noqa: F821 - Array is a forward reference
5656
"""
5757
return self.backend.vector_big_endian()
5858

59+
@property
60+
def probabilities(self) -> Array: # noqa: F821 - Array is a forward reference
61+
"""Get the probability distribution over all basis states.
62+
63+
Returns:
64+
Array of probabilities for each computational basis state.
65+
"""
66+
return self.backend.probabilities
67+
68+
def probability(self, basis_state: int) -> float:
69+
"""Get the probability of a specific computational basis state.
70+
71+
Args:
72+
basis_state: The index of the basis state.
73+
74+
Returns:
75+
The probability of measuring the given basis state.
76+
"""
77+
return self.backend.probability(basis_state)
78+
5979
def reset(self) -> StateVec:
6080
"""Resets the quantum state to the all-zero state."""
6181
self.backend.reset()

0 commit comments

Comments
 (0)