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

crates/ppvm-python-native/src/interface.rs

@@ -6,6 +6,7 @@ use ppvm_runtime::prelude::*;
 use ppvm_runtime::strategy::{
     CoefficientThreshold, CombinedStrategy, MaxLossWeight, MaxPauliWeight,
 };
+use ppvm_runtime::sum::PreserveConfig;
 use pyo3::prelude::*;
 
 macro_rules! create_interface_loss_methods {
@@ -59,21 +60,40 @@ 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 = Vec::<String>::new(), preserve_threshold = 0.0_f64, preserve_weight_lambda = 0.0_f64))]
+            #[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: Vec<String>,
+                preserve_threshold: f64,
+                preserve_weight_lambda: f64,
             ) -> Self {
                 let _ = max_loss_weight; // unused in non-loss variants
                 let strategy = create_strategy!($loss, min_abs_coeff, max_pauli_weight, max_loss_weight);
+                // Build the optional preserve-aware truncation policy. When
+                // `preserve` is empty the configured `strategy` is used as
+                // before; when non-empty the listed Pauli strings are kept
+                // regardless of coefficient and the rest are dropped below
+                // `preserve_threshold * exp(preserve_weight_lambda * weight)`.
+                let preserve_config = if preserve.is_empty() {
+                    None
+                } else {
+                    Some(PreserveConfig::from_strings(
+                        preserve.into_iter(),
+                        preserve_threshold,
+                        preserve_weight_lambda,
+                    ))
+                };
                 let mut ps = PauliSum::<$type>::builder()
                     .n_qubits(n_qubits)
                     .strategy(strategy)
                     .capacity(n_qubits)
+                    .maybe_preserve(preserve_config)
                     .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

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::config::Config;
+use crate::sum::preserve::PreserveConfig;
 use crate::traits::*;
 
 /// A sparse formal sum `Σ cᵢ Pᵢ` of Pauli strings.
@@ -42,6 +43,12 @@ pub struct PauliSum<T: Config> {
     n_qubits: usize,
     capacity: usize,
     strategy: T::Strategy,
+    /// Optional observable-aware truncation policy. When `Some(...)`,
+    /// [`PauliSum::truncate`] uses this instead of `strategy`: the
+    /// chosen Pauli strings are never dropped, and the rest are
+    /// truncated with a weight-biased threshold. See
+    /// [`PreserveConfig`].
+    preserve: Option<PreserveConfig<T::PauliWordType>>,
 }
 
 #[bon::bon]
@@ -51,6 +58,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 [`PreserveConfig`] for observable-aware truncation; when set
+    ///   it supersedes `strategy` inside [`truncate`](Self::truncate).
     #[builder]
     pub fn new(
         /// number of qubits
@@ -61,6 +70,8 @@ impl<T: Config> PauliSum<T> {
         /// capacity of the internal maps, default is strategy.capacity(n_qubits)
         #[builder(default = strategy.capacity(n_qubits))]
         capacity: usize,
+        /// optional observable-aware truncation policy
+        preserve: Option<PreserveConfig<T::PauliWordType>>,
     ) -> Self {
         Self {
             map: (
@@ -71,6 +82,7 @@ impl<T: Config> PauliSum<T> {
             n_qubits,
             capacity,
             strategy,
+            preserve,
         }
     }
 }
@@ -240,10 +252,36 @@ impl<T: Config> PauliSum<T> {
     }
 
     /// Apply the configured truncation [`Strategy`](crate::traits::Strategy)
-    /// to the primary map, dropping entries that fall outside its policy.
+    /// — or, if a [`PreserveConfig`] was supplied at construction, use
+    /// the preserve-aware policy instead — to the primary map, dropping
+    /// entries that fall outside the active policy.
     pub fn truncate(&mut self) {
-        let strategy = self.strategy;
-        strategy.truncate(self.data_mut());
+        if self.preserve.is_some() {
+            // Capture the keep-set directly as `HashSet<W>` so per-key
+            // lookup is O(1) with no string conversion. `W: Send + Sync`
+            // is part of the `PauliWordTrait` bundle, so the closure
+            // satisfies `ACMap::retain`'s `Sync + Send` requirement.
+            let preserve = self.preserve.as_ref().unwrap();
+            let keep = preserve.keep.clone();
+            let base = preserve.base_threshold;
+            let lambda = preserve.weight_lambda;
+            let data = self.data_mut();
+            data.retain(|k, v| {
+                if keep.contains(k) {
+                    return true;
+                }
+                let eff = base * (lambda * k.weight() as f64).exp();
+                !v.cutoff(eff)
+            });
+        } else {
+            let strategy = self.strategy;
+            strategy.truncate(self.data_mut());
+        }
+    }
+
+    /// Read-only access to the active [`PreserveConfig`], if any.
+    pub fn preserve(&self) -> Option<&PreserveConfig<T::PauliWordType>> {
+        self.preserve.as_ref()
     }
 }
 

crates/ppvm-runtime/src/sum/mod.rs

@@ -6,6 +6,7 @@ mod data;
 mod display;
 mod noise;
 mod ops;
+mod preserve;
 mod proj;
 mod rot1;
 mod rot2;
@@ -16,3 +17,4 @@ mod approx;
 
 pub use data::PauliSum;
 pub use ops::impl_op_mul_assign_coefficient;
+pub use preserve::PreserveConfig;