feat(runtime): observable-aware (preserve-set) truncation

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::<String>::new(), coefficients = Vec::<f64>::new()))]
+            #[pyo3(signature = (n_qubits, min_abs_coeff = 1e-10, max_pauli_weight = usize::MAX, max_loss_weight = usize::MAX, terms = Vec::<String>::new(), coefficients = Vec::<f64>::new(), preserve_strings = Vec::<String>::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<String>,
-                coefficients: Vec<f64>
+                coefficients: Vec<f64>,
+                preserve_strings: Vec<String>,
             ) -> 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!(

crates/ppvm-runtime/src/map/dashmap.rs

@@ -223,9 +223,9 @@ where
     H: BuildHasher + Clone + Default + Sync + Send,
     W: PauliWordTrait,
 {
-    fn retain<F>(&mut self, f: F)
+    fn retain<F>(&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));
     }

crates/ppvm-runtime/src/map/hashmap.rs

@@ -237,9 +237,9 @@ macro_rules! impl_acmap_retain {
             H: BuildHasher + Clone + Default,
             W: PauliWordTrait,
         {
-            fn retain<F>(&mut self, f: F)
+            fn retain<F>(&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<F>(&mut self, f: F)
+        fn retain<F>(&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));
         }

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<T: Config> {
     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<T::PauliWordType>,
 }
 
 #[bon::bon]
@@ -51,6 +56,8 @@ impl<T: Config> PauliSum<T> {
     /// 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<T: Config> PauliSum<T> {
         /// 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<T::PauliWordType>,
     ) -> Self {
         Self {
             map: (
@@ -71,6 +81,7 @@ impl<T: Config> PauliSum<T> {
             n_qubits,
             capacity,
             strategy,
+            preserve_strings,
         }
     }
 }
@@ -241,9 +252,53 @@ impl<T: Config> PauliSum<T> {
 
     /// 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<T::PauliWordType> {
+        &self.preserve_strings
     }
 }
 

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<F>(&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

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 = <Cfg as Config>::PauliWordType;
+
+fn single_z(n_qubits: usize) -> HashSet<PWord> {
+    (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<CfgThr> = 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<String> = 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<CfgThr> = 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<Cfg> = 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);
+}

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