diff --git a/python/pecos-rslib/pecos_rslib.pyi b/python/pecos-rslib/pecos_rslib.pyi index 64ac489c3..0a445402a 100644 --- a/python/pecos-rslib/pecos_rslib.pyi +++ b/python/pecos-rslib/pecos_rslib.pyi @@ -711,24 +711,53 @@ 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: ... @@ -736,6 +765,7 @@ 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,6 +796,7 @@ class QuestStateVec: """QuEST state vector simulator.""" def __init__(self, num_qubits: int) -> None: ... + def reset(self) -> QuestStateVec: ... @property def num_qubits(self) -> int: ... @@ -772,6 +804,7 @@ class QuestDensityMatrix: """QuEST density matrix simulator.""" def __init__(self, num_qubits: int) -> None: ... + def reset(self) -> QuestDensityMatrix: ... @property def num_qubits(self) -> int: ... diff --git a/python/pecos-rslib/src/cpp_sparse_sim_bindings.rs b/python/pecos-rslib/src/cpp_sparse_sim_bindings.rs index 484955baa..2423c99db 100644 --- a/python/pecos-rslib/src/cpp_sparse_sim_bindings.rs +++ b/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 { diff --git a/python/pecos-rslib/src/quest_bindings.rs b/python/pecos-rslib/src/quest_bindings.rs index e3ccb6d65..065f184c1 100644 --- a/python/pecos-rslib/src/quest_bindings.rs +++ b/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 diff --git a/python/pecos-rslib/src/qulacs_bindings.rs b/python/pecos-rslib/src/qulacs_bindings.rs index be4ca6e74..fc1231b19 100644 --- a/python/pecos-rslib/src/qulacs_bindings.rs +++ b/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 diff --git a/python/pecos-rslib/src/simulator_utils.rs b/python/pecos-rslib/src/simulator_utils.rs index a3fb14db4..75a3e7b00 100644 --- a/python/pecos-rslib/src/simulator_utils.rs +++ b/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>, + Vec>, + Vec>, +); + /// 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 { + self.sim + .call_method1(py, "_gens_data", (self.is_stab,))? + .extract(py) + } + + #[getter] + fn col_x(&self, py: Python<'_>) -> PyResult>> { + Ok(self.get_gens_data(py)?.0) + } + + #[getter] + fn col_z(&self, py: Python<'_>) -> PyResult>> { + Ok(self.get_gens_data(py)?.1) + } + + #[getter] + fn row_x(&self, py: Python<'_>) -> PyResult>> { + Ok(self.get_gens_data(py)?.2) + } + + #[getter] + fn row_z(&self, py: Python<'_>) -> PyResult>> { + Ok(self.get_gens_data(py)?.3) + } } /// Register the simulator utils module diff --git a/python/pecos-rslib/src/sparse_sim.rs b/python/pecos-rslib/src/sparse_sim.rs index 066daeca0..446c62de1 100644 --- a/python/pecos-rslib/src/sparse_sim.rs +++ b/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]| -> Vec> { + 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 { // Create a Rust GateBindingsDict directly diff --git a/python/pecos-rslib/src/sparse_stab_bindings.rs b/python/pecos-rslib/src/sparse_stab_bindings.rs index aa443adc3..c49d542f9 100644 --- a/python/pecos-rslib/src/sparse_stab_bindings.rs +++ b/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]| -> Vec> { + 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 { // Create a Rust GateBindingsDict directly diff --git a/python/pecos-rslib/src/state_vec_bindings.rs b/python/pecos-rslib/src/state_vec_bindings.rs index c3787bceb..5c9c615bc 100644 --- a/python/pecos-rslib/src/state_vec_bindings.rs +++ b/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 diff --git a/python/quantum-pecos/src/pecos/simulators/qulacs/state.py b/python/quantum-pecos/src/pecos/simulators/qulacs/state.py index 620b505c8..52c02953e 100644 --- a/python/quantum-pecos/src/pecos/simulators/qulacs/state.py +++ b/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 diff --git a/python/quantum-pecos/src/pecos/simulators/statevec/state.py b/python/quantum-pecos/src/pecos/simulators/statevec/state.py index 1cdcd7cda..4ce43b55a 100644 --- a/python/quantum-pecos/src/pecos/simulators/statevec/state.py +++ b/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() diff --git a/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_stab_sims/test_tableau_properties.py b/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_stab_sims/test_tableau_properties.py new file mode 100644 index 000000000..5416a0954 --- /dev/null +++ b/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_stab_sims/test_tableau_properties.py @@ -0,0 +1,108 @@ +# Copyright 2026 The PECOS Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with +# the License.You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. + +"""Tests for TableauWrapper col_x, col_z, row_x, row_z properties.""" + +from pecos.simulators import SparseSim + + +def test_initial_state_tableau_properties() -> None: + """Test tableau properties for the initial |0...0> state. + + For n qubits in |0...0>, stabilizers are Z_i for each qubit i, + so col_z should have qubit i in generator i, and col_x should be empty. + """ + n = 3 + state = SparseSim(n) + + stabs = state.stabs + col_x = stabs.col_x + col_z = stabs.col_z + + # Each stabilizer is Z_i, so col_z[i] should contain [i] and col_x[i] should be empty + for i in range(n): + assert col_x[i] == [], f"col_x[{i}] should be empty for |0> state" + assert col_z[i] == [i], f"col_z[{i}] should be [{i}] for |0> state" + + +def test_bell_state_tableau_properties() -> None: + """Test tableau properties after creating a Bell-like entangled state.""" + state = SparseSim(2) + state.run_gate("H", {0}) + state.run_gate("CX", {(0, 1)}) + + stabs = state.stabs + col_x = stabs.col_x + col_z = stabs.col_z + + # After H on qubit 0 and CX(0,1), stabilizers are XX and ZZ + # Verify we get non-empty data back with the right shape + assert len(col_x) == 2 + assert len(col_z) == 2 + + +def test_tableau_row_properties() -> None: + """Test that row_x and row_z are accessible and have correct dimensions.""" + n = 3 + state = SparseSim(n) + + stabs = state.stabs + row_x = stabs.row_x + row_z = stabs.row_z + + assert len(row_x) == n + assert len(row_z) == n + + +def test_destabs_tableau_properties() -> None: + """Test that destabilizer tableau properties are accessible.""" + n = 3 + state = SparseSim(n) + + destabs = state.destabs + col_x = destabs.col_x + col_z = destabs.col_z + row_x = destabs.row_x + row_z = destabs.row_z + + assert len(col_x) == n + assert len(col_z) == n + assert len(row_x) == n + assert len(row_z) == n + + # For initial |0> state, destabilizers are X_i + for i in range(n): + assert col_x[i] == [i], f"destab col_x[{i}] should be [{i}] for |0> state" + assert col_z[i] == [], f"destab col_z[{i}] should be empty for |0> state" + + +def test_gens_property() -> None: + """Test that the gens property returns (stabs, destabs) tuple.""" + state = SparseSim(2) + stabs_wrapper, destabs_wrapper = state.gens + + # Should be able to access col_x on both + assert len(stabs_wrapper.col_x) == 2 + assert len(destabs_wrapper.col_x) == 2 + + +def test_tableau_after_gate() -> None: + """Test that tableau properties update correctly after applying gates.""" + state = SparseSim(1) + + # Initial |0>: stab is Z + assert state.stabs.col_x == [[]] + assert state.stabs.col_z == [[0]] + + # After H: stab is X + state.run_gate("H", {0}) + assert state.stabs.col_x == [[0]] + assert state.stabs.col_z == [[]] diff --git a/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_statevec.py b/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_statevec.py index d88cbe9fe..70eda4796 100644 --- a/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_statevec.py +++ b/python/quantum-pecos/tests/pecos/integration/state_sim_tests/test_statevec.py @@ -398,6 +398,69 @@ def test_all_gate_circ(simulator: str) -> None: check_measurement(simulator, qc) +def test_statevec_probabilities_initial_state() -> None: + """Test that initial |0...0> state has probability 1 for basis state 0.""" + sim = StateVec(3) + probs = sim.probabilities + assert pc.isclose(probs[0], 1.0, rtol=0.0, atol=1e-10) + for i in range(1, 8): + assert pc.isclose(probs[i], 0.0, rtol=0.0, atol=1e-10) + + +def test_statevec_probabilities_bell_state() -> None: + """Test probabilities for a Bell state (|00> + |11>)/sqrt(2).""" + sim = StateVec(2) + sim.run_gate("H", {0}) + sim.run_gate("CX", {(0, 1)}) + + probs = sim.probabilities + assert pc.isclose(probs[0], 0.5, rtol=0.0, atol=1e-10) + assert pc.isclose(probs[1], 0.0, rtol=0.0, atol=1e-10) + assert pc.isclose(probs[2], 0.0, rtol=0.0, atol=1e-10) + assert pc.isclose(probs[3], 0.5, rtol=0.0, atol=1e-10) + + +def test_statevec_probability_single() -> None: + """Test probability() for individual basis states.""" + sim = StateVec(2) + sim.run_gate("H", {0}) + sim.run_gate("CX", {(0, 1)}) + + assert pc.isclose(sim.probability(0), 0.5, rtol=0.0, atol=1e-10) + assert pc.isclose(sim.probability(1), 0.0, rtol=0.0, atol=1e-10) + assert pc.isclose(sim.probability(2), 0.0, rtol=0.0, atol=1e-10) + assert pc.isclose(sim.probability(3), 0.5, rtol=0.0, atol=1e-10) + + +def test_statevec_probabilities_sum_to_one() -> None: + """Test that probabilities sum to 1 after arbitrary gates.""" + sim = StateVec(3) + sim.run_gate("H", {0}) + sim.run_gate("CX", {(0, 1)}) + sim.run_gate("H", {2}) + + total = pc.sum(sim.probabilities) + assert pc.isclose(total, 1.0, rtol=0.0, atol=1e-10) + + +def test_qulacs_probabilities() -> None: + """Test that Qulacs probabilities match StateVec probabilities.""" + check_dependencies("Qulacs") + + sv = StateVec(2) + sv.run_gate("H", {0}) + sv.run_gate("CX", {(0, 1)}) + + ql = Qulacs(2) + ql.run_gate("H", {0}) + ql.run_gate("CX", {(0, 1)}) + + sv_probs = sv.probabilities + ql_probs = ql.probabilities + for i in range(4): + assert pc.isclose(sv_probs[i], ql_probs[i], rtol=0.0, atol=1e-10) + + @pytest.mark.parametrize( "simulator", [