diff --git a/crates/ppvm-python-native/src/interface.rs b/crates/ppvm-python-native/src/interface.rs index 88afc07e3..71efcbfe5 100644 --- a/crates/ppvm-python-native/src/interface.rs +++ b/crates/ppvm-python-native/src/interface.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: 2026 The PPVM Authors // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashSet; + use paste::paste; use ppvm_runtime::prelude::*; use ppvm_runtime::strategy::{ @@ -59,21 +61,26 @@ macro_rules! create_interface { #[pymethods] impl $name { #[new] - #[pyo3(signature = (n_qubits, min_abs_coeff = 1e-10, max_pauli_weight = usize::MAX, max_loss_weight = usize::MAX, terms = Vec::::new(), coefficients = Vec::::new()))] + #[pyo3(signature = (n_qubits, min_abs_coeff = 1e-10, max_pauli_weight = usize::MAX, max_loss_weight = usize::MAX, terms = Vec::::new(), coefficients = Vec::::new(), preserve_strings = Vec::::new()))] + #[allow(clippy::too_many_arguments)] pub fn new( n_qubits: usize, min_abs_coeff: f64, max_pauli_weight: usize, max_loss_weight: usize, terms: Vec, - coefficients: Vec + coefficients: Vec, + preserve_strings: Vec, ) -> Self { let _ = max_loss_weight; // unused in non-loss variants let strategy = create_strategy!($loss, min_abs_coeff, max_pauli_weight, max_loss_weight); + let preserve_set: HashSet<_> = + preserve_strings.into_iter().map(Into::into).collect(); let mut ps = PauliSum::<$type>::builder() .n_qubits(n_qubits) .strategy(strategy) .capacity(n_qubits) + .preserve_strings(preserve_set) .build(); assert_eq!( diff --git a/crates/ppvm-runtime/src/map/dashmap.rs b/crates/ppvm-runtime/src/map/dashmap.rs index 0262f18f5..574f31aa9 100644 --- a/crates/ppvm-runtime/src/map/dashmap.rs +++ b/crates/ppvm-runtime/src/map/dashmap.rs @@ -223,9 +223,9 @@ where H: BuildHasher + Clone + Default + Sync + Send, W: PauliWordTrait, { - fn retain(&mut self, f: F) + fn retain(&mut self, mut f: F) where - F: Fn(&W, &C) -> bool + Sync + Send, + F: FnMut(&W, &C) -> bool, { Self::retain(self, |k, v| f(k, v)); } diff --git a/crates/ppvm-runtime/src/map/hashmap.rs b/crates/ppvm-runtime/src/map/hashmap.rs index ad4ecb957..bebc8069a 100644 --- a/crates/ppvm-runtime/src/map/hashmap.rs +++ b/crates/ppvm-runtime/src/map/hashmap.rs @@ -237,9 +237,9 @@ macro_rules! impl_acmap_retain { H: BuildHasher + Clone + Default, W: PauliWordTrait, { - fn retain(&mut self, f: F) + fn retain(&mut self, mut f: F) where - F: Fn(&W, &V) -> bool + Sync + Send, + F: FnMut(&W, &V) -> bool, { Self::retain(self, |k, v| f(k, v)); } @@ -351,9 +351,9 @@ mod ahash_impl { H: BuildHasher + Clone + Default, W: PauliWordTrait, { - fn retain(&mut self, f: F) + fn retain(&mut self, mut f: F) where - F: Fn(&W, &V) -> bool + Sync + Send, + F: FnMut(&W, &V) -> bool, { HashMap::retain(self, |k, v| f(k, v)); } diff --git a/crates/ppvm-runtime/src/sum/data.rs b/crates/ppvm-runtime/src/sum/data.rs index 6eaec791e..9568b9273 100644 --- a/crates/ppvm-runtime/src/sum/data.rs +++ b/crates/ppvm-runtime/src/sum/data.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: 2026 The PPVM Authors // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashSet; + use crate::config::Config; use crate::traits::*; @@ -42,6 +44,9 @@ pub struct PauliSum { n_qubits: usize, capacity: usize, strategy: T::Strategy, + /// Keep-set: strings [`PauliSum::truncate`] must always re-insert + /// after the strategy runs. Empty by default. + preserve_strings: HashSet, } #[bon::bon] @@ -51,6 +56,8 @@ impl PauliSum { /// One can optionally set /// - the strategy for truncation, initialization etc. /// - the capacity of the internal maps, default is strategy.capacity(n_qubits) + /// - a set of `preserve_strings` that [`truncate`](Self::truncate) must + /// never drop, on top of whatever the strategy decides. #[builder] pub fn new( /// number of qubits @@ -61,6 +68,9 @@ impl PauliSum { /// capacity of the internal maps, default is strategy.capacity(n_qubits) #[builder(default = strategy.capacity(n_qubits))] capacity: usize, + /// Pauli strings that truncate must always keep. Empty by default. + #[builder(default)] + preserve_strings: HashSet, ) -> Self { Self { map: ( @@ -71,6 +81,7 @@ impl PauliSum { n_qubits, capacity, strategy, + preserve_strings, } } } @@ -241,9 +252,53 @@ impl PauliSum { /// Apply the configured truncation [`Strategy`](crate::traits::Strategy) /// to the primary map, dropping entries that fall outside its policy. + /// + /// If `preserve_strings` is non-empty, any of those Pauli strings + /// that the strategy would have dropped are re-inserted afterwards + /// with their pre-truncate coefficient. The mechanism composes with + /// any [`Strategy`] (coefficient-magnitude, max-weight, combinations + /// — anything) because the strategy runs unchanged in the middle. pub fn truncate(&mut self) { + // Hot path: empty preserve set → just run the strategy. + if self.preserve_strings.is_empty() { + let strategy = self.strategy; + strategy.truncate(self.data_mut()); + return; + } + + // Snapshot the current coefficients of preserved keys. We piggy- + // back on `retain` (which always returns true here, so it's a + // pure scan) to walk `(k, v)` pairs without needing a separate + // `get`/`iter` route through the `ACMap` traits. + let preserve = self.preserve_strings.clone(); + let mut saved: Vec<(T::PauliWordType, T::Coeff)> = Vec::new(); + self.data_mut().retain(|k, v| { + if preserve.contains(k) { + saved.push((k.clone(), v.clone())); + } + true + }); + + // Run the configured strategy verbatim. let strategy = self.strategy; strategy.truncate(self.data_mut()); + + // Restore any preserved entry the strategy dropped. We use + // `add_assign` because the trait does not expose a plain + // "insert if absent"; on a missing key it inserts, which is + // what we want here (the guarded `contains_with` keeps us from + // accidentally summing into a kept entry). + let data = self.data_mut(); + for (k, v) in saved { + if !data.contains_with(&k, |_| true) { + data.add_assign(k, v); + } + } + } + + /// Read-only access to the active preserve set. + pub fn preserve_strings(&self) -> &HashSet { + &self.preserve_strings } } diff --git a/crates/ppvm-runtime/src/traits/map.rs b/crates/ppvm-runtime/src/traits/map.rs index d7fb8dcf6..ea426a4ca 100644 --- a/crates/ppvm-runtime/src/traits/map.rs +++ b/crates/ppvm-runtime/src/traits/map.rs @@ -133,7 +133,7 @@ pub trait ACMapRetain< /// Keep only entries for which `f(key, value)` returns `true`. fn retain(&mut self, f: F) where - F: Fn(&W, &V) -> bool + Sync + Send; + F: FnMut(&W, &V) -> bool; } /// Aggregate trait combining every operation a backing map must support diff --git a/crates/ppvm-runtime/tests/preserve.rs b/crates/ppvm-runtime/tests/preserve.rs new file mode 100644 index 000000000..4010369e2 --- /dev/null +++ b/crates/ppvm-runtime/tests/preserve.rs @@ -0,0 +1,114 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Integration tests for the `preserve_strings` snapshot-and-restore +//! post-filter in [`PauliSum::truncate`]: whatever the active strategy +//! decides to drop, preserved strings come back. + +use std::collections::HashSet; + +use ppvm_runtime::config::Config; +use ppvm_runtime::config::fxhash::ByteF64; +use ppvm_runtime::prelude::*; +use ppvm_runtime::strategy::CoefficientThreshold; +use ppvm_runtime::sum::PauliSum; + +type Cfg = ByteF64<1>; +type CfgThr = ByteF64<1, CoefficientThreshold>; +type PWord = ::PauliWordType; + +fn single_z(n_qubits: usize) -> HashSet { + (0..n_qubits) + .map(|i| { + let s: String = (0..n_qubits) + .map(|j| if j == i { 'Z' } else { 'I' }) + .collect(); + PWord::from(s) + }) + .collect() +} + +/// The active strategy (`CoefficientThreshold`) drops a tiny coefficient, +/// but the preserved string is re-inserted. +#[test] +fn truncate_restores_preserved_string_dropped_by_strategy() { + let mut s: PauliSum = PauliSum::builder() + .n_qubits(3) + .strategy(CoefficientThreshold(0.5)) + .preserve_strings(single_z(3)) + .build(); + // Single-Z strings get tiny coefficients (below cutoff 0.5) — must survive. + s += ("ZII", 1e-6); + s += ("IZI", 1e-6); + s += ("IIZ", 1e-6); + // A non-preserved term well below cutoff — must be dropped. + s += ("XYZ", 1e-6); + // A non-preserved term above cutoff — must survive. + s += ("XXX", 0.7); + + s.truncate(); + let kept: HashSet = s.data().keys().map(|k| k.to_string()).collect(); + assert!(kept.contains("ZII"), "preserved ZII should be kept"); + assert!(kept.contains("IZI"), "preserved IZI should be kept"); + assert!(kept.contains("IIZ"), "preserved IIZ should be kept"); + assert!( + !kept.contains("XYZ"), + "below-cutoff non-preserved XYZ should be dropped" + ); + assert!(kept.contains("XXX"), "above-cutoff XXX should be kept"); +} + +/// End-to-end conservation: `Σ Z_i` propagated through a sequence of +/// `rxx + ryy` exchange-style gates (which preserve total Z) with +/// aggressive coefficient truncation keeps every single-Z coefficient +/// at 1.0 exactly. The same setup without the preserve set would drop +/// them once their coefficients dipped below the threshold. +#[test] +fn preserve_single_z_conserves_total_z_under_aggressive_truncation() { + let n = 4; + let mut s: PauliSum = PauliSum::builder() + .n_qubits(n) + .strategy(CoefficientThreshold(0.5)) + .preserve_strings(single_z(n)) + .build(); + for j in 0..n { + let term: String = (0..n).map(|i| if i == j { 'Z' } else { 'I' }).collect(); + s += (term.as_str(), 1.0); + } + + // Apply a few rxx+ryy pairs (= XY exchange on each edge). This commutes + // with Σ Z_k, so the coefficients on Z_j should remain at 1.0. + for (a, b) in [(0, 1), (1, 2), (2, 3)] { + s.rxx(a, b, 0.37); + s.ryy(a, b, 0.37); + s.truncate(); + } + + for j in 0..n { + let term: String = (0..n).map(|i| if i == j { 'Z' } else { 'I' }).collect(); + let word: PWord = term.clone().into(); + let coeff = s.data().iter().find(|(k, _)| **k == word).map(|(_, v)| *v); + assert!( + coeff.is_some(), + "single-Z string {} must be preserved", + term + ); + assert!( + (coeff.unwrap() - 1.0).abs() < 1e-10, + "coefficient on {} should remain 1.0 (got {})", + term, + coeff.unwrap() + ); + } +} + +/// No preserve set → behaviour is identical to the bare strategy. +#[test] +fn empty_preserve_falls_back_to_strategy_unchanged() { + let n = 2; + let mut s: PauliSum = PauliSum::builder().n_qubits(n).build(); + s += ("ZI", 1.0); + s += ("XY", 1e-30); + s.truncate(); // default strategy keeps everything + assert_eq!(s.data().iter().count(), 2); +} diff --git a/ppvm-python/src/ppvm/paulisum.py b/ppvm-python/src/ppvm/paulisum.py index 4ad16248b..784501a7d 100644 --- a/ppvm-python/src/ppvm/paulisum.py +++ b/ppvm-python/src/ppvm/paulisum.py @@ -137,6 +137,7 @@ class PauliSum( min_abs_coeff: float = 1e-10 max_pauli_weight: int | None = None max_loss_weight: int | None = None + preserve_strings: Sequence[str] | None = None _interface: PauliSumInterface = field(init=False, repr=False) @@ -193,6 +194,16 @@ def _init_ppvm_interface( if self.max_loss_weight is not None: options["max_loss_weight"] = self.max_loss_weight + if self.preserve_strings: + preserve_list = list(self.preserve_strings) + for s in preserve_list: + if len(s) != n_qubits: + raise ValueError( + "All preserve strings must have length n_qubits " + f"({n_qubits}); got {len(s)}: {s!r}" + ) + options["preserve_strings"] = preserve_list + return interface( n_qubits, **options, @@ -214,6 +225,7 @@ def new( min_abs_coeff: float = 1e-10, max_pauli_weight: int | None = None, max_loss_weight: int | None = None, + preserve_strings: Sequence[str] | None = None, ) -> Self: """Create a PauliSum from one or more terms with flexible input formats. @@ -234,6 +246,8 @@ def new( Note, that this should usually be chosen to be quite low, since e.g. 10 would correspond to keeping terms that contribute if up to 10 qubits are lost simultaneously. + preserve_strings: Pauli strings (length ``n_qubits`` each) that + truncation must never drop. Empty by default. Returns: A new instance of the class this method is called on. @@ -278,6 +292,7 @@ def new( min_abs_coeff=min_abs_coeff, max_pauli_weight=max_pauli_weight, max_loss_weight=max_loss_weight, + preserve_strings=preserve_strings, ) def __str__(self) -> str: @@ -291,6 +306,7 @@ def __copy__(self) -> Self: object.__setattr__(new, "min_abs_coeff", self.min_abs_coeff) object.__setattr__(new, "max_pauli_weight", self.max_pauli_weight) object.__setattr__(new, "max_loss_weight", self.max_loss_weight) + object.__setattr__(new, "preserve_strings", self.preserve_strings) object.__setattr__(new, "_interface", self._interface.__copy__()) return new diff --git a/ppvm-python/test/test_preserve.py b/ppvm-python/test/test_preserve.py new file mode 100644 index 000000000..37dd748ff --- /dev/null +++ b/ppvm-python/test/test_preserve.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: 2026 The PPVM Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for observable-aware (preserve-set) truncation.""" + +import math + +import pytest + +from ppvm import PauliSum + + +def _single_z(n_qubits: int) -> list[str]: + return ["".join("Z" if i == j else "I" for i in range(n_qubits)) for j in range(n_qubits)] + + +# ============================================================================= +# Plumbing: when preserve_strings is set, the field round-trips and `new()` +# accepts it. +# ============================================================================= + + +def test_preserve_strings_round_trip(): + ps = PauliSum.new( + 3, + "Z1", + preserve_strings=_single_z(3), + ) + assert list(ps.preserve_strings) == ["ZII", "IZI", "IIZ"] + + +def test_preserve_strings_length_validated(): + with pytest.raises(ValueError, match="length n_qubits"): + PauliSum.new( + 3, + "Z1", + preserve_strings=["ZI"], # wrong length + ) + + +# ============================================================================= +# Behavior: a preserved string with tiny coefficient survives truncation, +# regardless of which truncation strategy is active. The preserve mechanism +# is a post-filter that re-inserts dropped preserved keys after the strategy +# runs — it composes with any strategy. +# ============================================================================= + + +def test_preserved_string_survives_coefficient_truncation(): + """`min_abs_coeff` would drop a tiny-coefficient single-Z, but the + preserve mechanism puts it back.""" + ps = PauliSum.new( + 3, + [("Z0", 1e-8), ("X0", 0.5), ("X1", 1e-8)], + min_abs_coeff=1e-3, + preserve_strings=_single_z(3), + ) + # Trigger auto-truncate via a no-op gate. + ps.rx(0, 0.0) + kept = {t for t, _ in ps.terms} + assert "ZII" in kept, "preserved tiny Z0 must survive" + assert "XII" in kept, "above-threshold XII must survive" + assert "IXI" not in kept, "below-threshold non-preserved IXI must be dropped" + + +def test_preserved_string_survives_weight_truncation(): + """`max_pauli_weight` would drop a high-weight string; if that string + happens to be in the preserve set, it's restored.""" + # Build a 4-qubit sum where one term has weight 3 (X0X1X2 → "XXXI") + # and we cap max_pauli_weight at 2. Without preserve, the weight-3 + # term is dropped. With preserve including it, it survives. + ps = PauliSum.new( + 4, + [("Z0", 1.0), ("X0X1X2", 0.7)], + max_pauli_weight=2, + preserve_strings=["XXXI"], + ) + ps.rx(0, 0.0) # no-op triggers truncate + kept = {t for t, _ in ps.terms} + assert "ZIII" in kept, "weight-1 Z0 must survive" + assert "XXXI" in kept, "weight-3 XXXI is in preserve set and must survive" + + +def test_preserved_string_survives_combined_truncation(): + """The strategy combines coefficient *and* weight cuts; preserve still + works orthogonally on top.""" + ps = PauliSum.new( + 4, + [("Z0", 1e-8), ("X0X1X2", 1e-8), ("Y0", 0.5)], + min_abs_coeff=1e-3, + max_pauli_weight=2, + preserve_strings=["ZIII", "XXXI"], # one tiny-coef, one high-weight + ) + ps.rx(0, 0.0) + kept = {t for t, _ in ps.terms} + assert "ZIII" in kept, "preserved tiny Z0 must survive coefficient cut" + assert "XXXI" in kept, "preserved weight-3 XXXI must survive weight cut" + assert "YIII" in kept, "above-threshold YIII must survive" + + +# ============================================================================= +# Transport diagnostic: <Σ_j Z_j(t) Z_i(0)>.sum() is conserved exactly with +# preserve, but drifts without. +# ============================================================================= + + +def test_total_z_conservation_with_preserve_vs_without(): + """Propagate a localized Z_i under XY exchange + Z-dephasing with + aggressive truncation. Single-Z preserve substantially reduces the + drift in `result.sum(axis=1)` versus the same run with plain + `min_abs_coeff` truncation.""" + L = 8 + i = L // 2 + threshold = 0.02 + gamma = 0.5 + dt = 0.1 + steps = 6 + noise = (1 - math.exp(-gamma * dt)) / 2 + + edges = [(a, a + 1, 0.08) for a in range(L - 1)] + z_observables = [PauliSum.new(L, f"Z{j}") for j in range(L)] + + def evolve(ps): + sums = [] + for _ in range(steps + 1): + sums.append(sum(ps.overlap(zz) for zz in z_observables)) + for q in range(L): + ps.pauli_error(q, [0.0, 0.0, noise]) + for a, b, th in reversed(edges): + ps.rxx(a, b, th) + ps.ryy(a, b, th) + for q in range(L): + ps.pauli_error(q, [0.0, 0.0, noise]) + return sums + + # Plain truncation. + ps_plain = PauliSum.new(L, f"Z{i}", min_abs_coeff=threshold, max_pauli_weight=L) + drift_plain = evolve(ps_plain) + + # Same strategy + preserve-set on top. + ps_pres = PauliSum.new( + L, + f"Z{i}", + min_abs_coeff=threshold, + max_pauli_weight=L, + preserve_strings=_single_z(L), + ) + drift_pres = evolve(ps_pres) + + plain_drift = 1.0 - drift_plain[-1] + pres_drift = abs(1.0 - drift_pres[-1]) + assert plain_drift > 1e-3, ( + f"sanity check: plain truncation should drift here " + f"(got drift={plain_drift:.2e}); try lowering threshold" + ) + assert pres_drift < plain_drift / 10, ( + f"preserve should reduce drift by >= 10x; " + f"got preserve_drift={pres_drift:.2e}, plain_drift={plain_drift:.2e}" + ) + + +# ============================================================================= +# Default behavior: when preserve_strings is None, the existing strategy +# (CoefficientThreshold + MaxPauliWeight) is used unchanged. +# ============================================================================= + + +def test_no_preserve_uses_existing_strategy(): + """Without preserve_strings, behaviour is identical to before this + change. Below-threshold strings are dropped uniformly.""" + ps = PauliSum.new(2, [("ZI", 0.5), ("XI", 1e-8)], min_abs_coeff=1e-3) + ps.rx(1, 0.0) # no-op gate triggers truncate + kept = {t for t, _ in ps.terms} + assert "ZI" in kept + assert "XI" not in kept