Repository navigation
feat(runtime): observable-aware (preserve-set) truncation - #95
Conversation
Adds a `PreserveConfig<W>` option on `PauliSum` that tells `truncate()` to never drop a chosen set of Pauli strings, and to truncate the rest with an exponentially weight-biased threshold. Motivation: for computing transport diagnostics like `Σ_j <Z_j(t) Z_i(0)>`, the answer depends only on the projection of the propagated `Z_i(t)` onto the L single-Z strings. Standard `CoefficientThreshold` truncation drops those exact strings once their coefficients drift toward the cutoff (typically the tail of a spreading operator at far sites), causing the conserved-charge component to leak. With preserve enabled, the small chosen set of strings is always kept regardless of coefficient, and `<Σ Z>` is preserved to floating-point precision. The weight-biased threshold is the "virtual DAOE" version of the truncation tactic from Rakovszky, Pollmann & von Keyserlingk (2020): high-weight Pauli strings are truncated more aggressively (effective cutoff `base · exp(λ · weight)`), but no actual damping is applied to the dynamics. Setting `λ = 0` recovers a uniform cutoff. Rust: - New `PreserveConfig<W>` in `ppvm-runtime/src/sum/preserve.rs` with `new`, `single_z`, `from_strings`, `threshold_for_weight` helpers. - New optional `preserve` builder field on `PauliSum`; when set, `truncate()` uses the preserve-aware policy instead of `T::Strategy`. - 8 unit tests covering keep-set construction, threshold computation, end-to-end keeps below-cutoff preserved strings, weight-lambda aggressively drops high-weight terms, and total-Z conservation under aggressive truncation through XY exchange. Python: - New `preserve_strings` / `preserve_threshold` / `preserve_weight_lambda` kwargs on `PauliSum.new()` and the `PauliSum` dataclass. Defaults preserve the existing behavior. - New `preserve_single_z(n_qubits)` helper module-level function for the common `Σ Z` diagnostic case. - Exposed through the `create_interface!` macro so every Config variant picks up the new constructor parameters. - 7 tests including a side-by-side end-to-end comparison: with aggressive truncation, preserve reduces `<Σ Z>` drift by >10×. The change is purely additive: callers that don't set `preserve_strings` get exactly the previous behavior (the `T::Strategy` truncation path). The new `truncate()` does not introduce a `Sync + Send` bound on `T::PauliWordType` — the preserve closure captures a snapshot `HashSet<String>` instead, so existing benches and other generic call sites are unaffected. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Capture the keep-set as `HashSet<W>` directly so per-entry lookup is O(1) without allocating; this requires `ACMapRetain::retain`'s closure to be `FnMut` rather than `Fn + Sync + Send`, which is what every backing map (`HashMap`, `IndexMap`, `DashMap`, `AHashMap`) already delegates to anyway — the prior bound was overly restrictive and forced the preserve-aware path to capture a `HashSet<String>` and call `to_string()` on every map entry per truncate. On the user-facing reference benchmark (L=16, alpha=1, dt=0.01, steps=100, gamma=1, pbc=1, min_abs_coeff=1e-4): | config | time (s) before | time (s) after | speedup | |---------------------------|----------------:|---------------:|--------:| | exchange + preserve=1 | 1.17 | 0.45 | 2.6x | | rxx+ryy + preserve=1 | 1.85 | 0.71 | 2.6x | Bit-exact `<Σ Z>` conservation under `exchange + preserve=1` is unchanged (drift = 8.2e-15, machine ε for f64). No public-API change beyond `ACMapRetain::retain`'s relaxed bound, which is strictly more permissive — every existing caller satisfies `FnMut`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
There was a problem hiding this comment.
Pull request overview
This PR introduces observable-aware truncation for PauliSum by adding a preserve-set policy (PreserveConfig<W>) that guarantees selected Pauli strings are never dropped during truncation, while optionally applying a weight-biased cutoff (threshold * exp(λ·weight)) to all other terms. It also relaxes the Rust retain predicate bound to FnMut to avoid unnecessary constraints and enable faster preserve-mode truncation.
Changes:
- Add
PreserveConfig<W>and integrate it intoPauliSum::truncate()as an optional truncation policy that supersedes the configured strategy when present. - Expose preserve-aware truncation through the Python API (
preserve_strings,preserve_threshold,preserve_weight_lambda) and add apreserve_single_z(n_qubits)helper. - Relax
ACMapRetain::retainpredicate bounds fromFn + Sync + SendtoFnMut, and update map implementations accordingly.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| ppvm-python/test/test_preserve.py | Adds Python tests covering preserve-set behavior, transport diagnostic drift reduction, and weight-biased truncation. |
| ppvm-python/src/ppvm/paulisum.py | Adds Python plumbing for preserve-aware truncation and a preserve_single_z helper. |
| ppvm-python/src/ppvm/init.py | Re-exports preserve_single_z from the package root. |
| crates/ppvm-runtime/src/traits/map.rs | Relaxes ACMapRetain::retain closure bound to FnMut and documents the rationale. |
| crates/ppvm-runtime/src/sum/preserve.rs | Introduces PreserveConfig<W> and Rust-side tests for preserve-aware truncation semantics. |
| crates/ppvm-runtime/src/sum/mod.rs | Registers and re-exports the new preserve module/config. |
| crates/ppvm-runtime/src/sum/data.rs | Adds an optional preserve policy to PauliSum and switches truncate() behavior when configured. |
| crates/ppvm-runtime/src/map/hashmap.rs | Updates retain adapter impls to accept FnMut. |
| crates/ppvm-runtime/src/map/dashmap.rs | Updates retain adapter impls to accept FnMut. |
| crates/ppvm-python-native/src/interface.rs | Adds preserve parameters to the PyO3 constructor and wires them into the Rust PauliSum builder. |
Comments suppressed due to low confidence (1)
crates/ppvm-runtime/src/sum/data.rs:264
- The comment above the retain closure still mentions satisfying a
Sync + Sendrequirement and claimsPauliWordTraitimpliesSend + Sync, butPauliWordTraitdoes not include those bounds andACMapRetain::retainno longer requiresSync + Sendfor the predicate. Updating/removing this comment would avoid misleading future readers.
// 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();
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 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"] = preserve_list | ||
| options["preserve_threshold"] = ( | ||
| self.preserve_threshold | ||
| if self.preserve_threshold is not None | ||
| else self.min_abs_coeff | ||
| ) | ||
| options["preserve_weight_lambda"] = self.preserve_weight_lambda |
| #[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, | ||
| )) | ||
| }; |
| // 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(); |
| 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; " |
| ps = PauliSum.new( | ||
| 4, | ||
| [("X0", 0.05), ("X0X1", 0.05), ("X0X1X2", 0.05)], | ||
| preserve_strings=[], # no preserve; use only weight-biased cutoff | ||
| preserve_threshold=0.01, | ||
| preserve_weight_lambda=1.0, | ||
| ) | ||
| # We can't pass preserve_strings=[] currently — the dataclass treats | ||
| # falsy as "no preserve". So instead pass a dummy preserve set with | ||
| # something we don't have, plus set the threshold. Re-do via a | ||
| # weight-only style: | ||
| ps = PauliSum.new( | ||
| 4, | ||
| [("X0", 0.05), ("X0X1", 0.05), ("X0X1X2", 0.05)], | ||
| preserve_strings=["IIII"], # never-matching keep-set; pure weighted cutoff | ||
| preserve_threshold=0.01, | ||
| preserve_weight_lambda=1.0, | ||
| ) |
david-pl
left a comment
There was a problem hiding this comment.
Cool approach, but I think the implementation went a little overboard in trying to avoid changing anything existing to make this change purely additive.
I particularly don't like the pattern of "switching" the usual behavior based on whether you set a PreserveConfig.
I'd suggest the following, though I admit I'm also not 100% sure it's feasible:
- remove
PreserveConfig. - make preservation a new
Strategyfor truncation. - setting this via a kwarg on the python side is fine and shouldn't really look any different from how it is now.
The tricky part is, if you make Preserve a new Strategy, it doesn't compose well with the MaxWeight strategy: the latter would still prune strings you'd want to keep. To overcome this, I suggest you:
- either just change
MaxWeightto also take akeepvariable (empty by default), which will not be truncated. Really, that sounds like the simplest change to me. This would be a small change (new field + a small logic change on the truncation) and would do everything you're doing here. However, I might be missing some details. - or, you introduce a concept of mutually exclusiveness in the strategies, i.e. when creating a
CombinedStrategy, check whether these can be used together, otherwise error.
| fn retain<F>(&mut self, f: F) | ||
| fn retain<F>(&mut self, mut f: F) | ||
| where | ||
| F: Fn(&W, &C) -> bool + Sync + Send, |
There was a problem hiding this comment.
Are you sure we don't need this? We may just be missing some tests on DashMap. It can't be too bad if it compiles, however.
| /// chosen Pauli strings are never dropped, and the rest are | ||
| /// truncated with a weight-biased threshold. See | ||
| /// [`PreserveConfig`]. | ||
| preserve: Option<PreserveConfig<T::PauliWordType>>, |
There was a problem hiding this comment.
I'm not a fan of the overall implementation here, tbh. Setting this is a flag that completely alters truncation behavior. It also breaks composition with other truncation strategies. You can't use this together with CoefficientThreshold, which you now work around by making the preservation truncation basically a combination of that strategy and the coefficient cutoff (I think, I'm not sure I fully understand the exponential there).
| pub fn truncate(&mut self) { | ||
| let strategy = self.strategy; | ||
| strategy.truncate(self.data_mut()); | ||
| if self.preserve.is_some() { |
There was a problem hiding this comment.
See above. This is odd and doesn't compose with other truncation strategies anymore.
| /// [`PreserveConfig::new`] for arbitrary keep-sets, or via one of the | ||
| /// `single_z` / `from_strings` convenience constructors. | ||
| #[derive(Debug, Clone)] | ||
| pub struct PreserveConfig<W: PauliWordTrait> { |
There was a problem hiding this comment.
Why is this a PauliWord? It's a little odd, we never want to use this as an actual PauliString, right? Is it just for construction convenience? If so, maybe just implement the methods without the trait?
… any strategy
`preserve` is no longer a separate truncation policy that *replaces*
the configured `Strategy` — it is now a small post-filter on top of
whatever the strategy already does:
```
PauliSum::truncate():
saved = snapshot of preserve_strings entries
strategy.truncate(map) # unchanged
re-insert any preserved entry dropped by the strategy
```
This means preserve composes with **any** truncation strategy
(coefficient-magnitude, max-weight, the existing `CombinedStrategy`
of both, or anything else) — the strategy is run verbatim and the
preserve mechanism only kicks in for the specific Pauli strings the
caller asked to keep.
API simplifications:
- Removed `PreserveConfig<W>` struct entirely. The keep-set is just
a `HashSet<T::PauliWordType>` field on `PauliSum`, empty by
default. Builder method renamed to `.preserve_strings(set)`.
- Removed `weight_lambda` / `base_threshold` knobs and the
`PreserveConfig::threshold_for_weight` helper. The "virtual DAOE"
weight-biased threshold belonged in (and can later be added to)
the coefficient-threshold strategy itself, not piggy-backed on the
preserve mechanism.
- Python: dropped `preserve_threshold` and `preserve_weight_lambda`
kwargs on `PauliSum.new()` and the dataclass. Only
`preserve_strings: Sequence[str] | None` remains.
- The free helpers `preserve::single_z(n)` and
`preserve::from_strings(...)` now return a plain `HashSet<W>`
instead of constructing the removed `PreserveConfig`.
Tests:
- The Rust integration tests now build the `PauliSum` with an
explicit `CoefficientThreshold` strategy plus a preserve set, and
assert the preserved strings survive the strategy's drops.
- New Python tests
`test_preserved_string_survives_{coefficient,weight,combined}_truncation`
verify the post-filter composes correctly with each strategy
variant.
- Dropped the `test_weight_lambda_drops_high_weight_more_aggressively`
test (no longer applicable).
Net diff: +203 / -352 lines; the implementation gets smaller and
strictly more general.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
david-pl
left a comment
There was a problem hiding this comment.
@AlexSchuckert good stuff, the new version is much simpler, I like it.
Just some minor clean up work (I think there's leftover stuff from the previous version) and then this is good to go.
| /// Build the keep-set of all single-`Z` Pauli strings on `n_qubits` qubits — | ||
| /// `Z_0, Z_1, …, Z_{n−1}`. The natural choice when the transport | ||
| /// diagnostic is `<Σ_j Z_j(t) Z_i(0)>` (z-magnetization spread). | ||
| pub fn single_z<W: PauliWordTrait>(n_qubits: usize) -> HashSet<W> { |
There was a problem hiding this comment.
This function and from_strings is only used in tests. Are they needed anywhere? Otherwise, I'd suggest to just remove these and move the entire file into tests/ since there's no actual code anymore.
| _COMPACT_TOKEN_RE = re.compile(r"([IXYZ])(\d+)") | ||
|
|
||
|
|
||
| def preserve_single_z(n_qubits: int) -> list[str]: |
There was a problem hiding this comment.
Is this function used anywhere?
| n_qubits: usize, | ||
| capacity: usize, | ||
| strategy: T::Strategy, | ||
| /// Pauli strings that [`PauliSum::truncate`] must never drop, |
There was a problem hiding this comment.
Now there's exactly one super long docstring in here. Kind of odd 😄
- delete `single_z` / `from_strings` from `sum::preserve` (only used in tests; the one external caller in `interface.rs` is now a one-liner) - move the three integration tests to `tests/preserve.rs` and drop the now-empty `sum::preserve` module - drop the `preserve_single_z` Python helper (only used by tests; inlined as `_single_z` in `test_preserve.py`) - shorten the `preserve_strings` docstrings on the Rust field and the Python `new()` arg; remove the `:func:preserve_single_z` cross-reference Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|
should I merge this? |
| if self.max_loss_weight is not None: | ||
| options["max_loss_weight"] = self.max_loss_weight | ||
|
|
||
| if self.preserve_strings: |
| 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 | ||
| }); |
Summary
Adds a
preserve_stringskeep-set onPauliSumthat tellstruncate()to never drop a chosen set of Pauli strings, regardless of which truncation strategy is active.Motivation. For transport diagnostics like
Σ_j <Z_j(t) Z_i(0)>, the answer depends only on the projection of the propagatedZ_i(t)onto the L single-Z strings. PlainCoefficientThreshold(or any other truncation strategy) drops those exact strings once their coefficients drift toward the cutoff (typically the tail of a spreading operator at far sites), causing the conserved-charge component to leak. Withpreserve_stringsset to the single-Z list, the small chosen set is always kept regardless of coefficient, and<Σ Z>is preserved to floating-point precision.Design
preserveis not itself a truncation strategy — it composes with anyStrategyvia a snapshot-and-restore post-filter insidePauliSum::truncate:That's the whole mechanism. It composes with
CoefficientThreshold,MaxPauliWeight,CombinedStrategy(both), or anything else without modifying any strategy. The strategy runs unchanged in the middle.API
Rust:
Python:
Module-level helpers
preserve::single_z(n)andpreserve::from_strings(...)(and the Pythonpreserve_single_zre-export) build the keep-set.History of this PR
This branch had two earlier commits implementing the keep-set as
PreserveConfig<W>that replaced the configured strategy insidetruncate, and bundled aweight_lambda"virtual DAOE" knob. Review feedback said the keep-set shouldn't be a separate strategy. The third commit (41892e0) is that refactor:PreserveConfig<W>— keep-set is now justHashSet<T::PauliWordType>.weight_lambda/base_thresholdentirely.preserve_threshold/preserve_weight_lambdakwargs.PauliSum::truncateas the snapshot-and-restore post-filter above; the strategy always runs verbatim.Net effect of commit 3: +203 / −352 lines — the implementation gets smaller and strictly more general. The earlier commit (
0525523, relaxingACMapRetain::retainfromFn + Sync + SendtoFnMut) is kept because the cleanup is independently useful.Tests
Rust (5 tests in
crates/ppvm-runtime/src/sum/preserve.rs::tests):single_z_builds_correct_set,from_strings_round_trip— keep-set helpers.truncate_restores_preserved_string_dropped_by_strategy— explicitCoefficientThreshold(0.5)+ tiny-coef preserved key.preserve_single_z_conserves_total_z_under_aggressive_truncation— end-to-endΣ Zthrough XY exchange +CoefficientThreshold(0.5).empty_preserve_falls_back_to_strategy_unchanged— no preserve set → fast path.Python (8 tests in
ppvm-python/test/test_preserve.py):preserve_single_zhelper, dataclass round-trip, length validation.test_preserved_string_survives_coefficient_truncation— preserve +min_abs_coeff.test_preserved_string_survives_weight_truncation— preserve +max_pauli_weight.test_preserved_string_survives_combined_truncation— preserve + both at once.preservereduces<Σ Z>drift by ≥10× vs plain.test_no_preserve_uses_existing_strategy— default behaviour unchanged.Test plan
cargo test --workspace(83 runtime tests, all pass)uv run --project ppvm-python --group dev pytest ppvm-python/test/(110 pass, of which 8 are preserve)cargo fmt,cargo clippy --no-depsruff check ppvm-python/Notes
preserve_stringsget exactly the previous behaviour (their configured strategy, unchanged).retainbound from commit 2 is a strict relaxation — every existing caller still satisfiesFnMut.