diff --git a/Cargo.lock b/Cargo.lock index 526c1b310..6fd2d27aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1499,6 +1499,12 @@ version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -2190,8 +2196,11 @@ name = "ppvm-lindblad" version = "0.1.1" dependencies = [ "approx", + "bitvec", + "bytemuck", "criterion 0.7.0", "fxhash", + "hashbrown 0.17.1", "nalgebra", "ndarray", "num", diff --git a/crates/ppvm-lindblad/Cargo.toml b/crates/ppvm-lindblad/Cargo.toml index 961f48131..d5872a3a7 100644 --- a/crates/ppvm-lindblad/Cargo.toml +++ b/crates/ppvm-lindblad/Cargo.toml @@ -8,7 +8,10 @@ publish = false description = "Direct Heisenberg-picture Lindbladian evolution on an adaptive Pauli-string basis." [dependencies] +bitvec = "1.0.1" +bytemuck = { version = "1", features = ["min_const_generics"] } fxhash = "0.2.1" +hashbrown = { version = "0.17.1", default-features = false } ndarray = "0.17" num = "0.4.3" ppvm-traits = { version = "0.1.1", path = "../ppvm-traits" } diff --git a/crates/ppvm-lindblad/src/algebra.rs b/crates/ppvm-lindblad/src/algebra.rs index 5fd3f4750..95b22850b 100644 --- a/crates/ppvm-lindblad/src/algebra.rs +++ b/crates/ppvm-lindblad/src/algebra.rs @@ -13,7 +13,6 @@ use crate::word::{Chunk, W_CHUNKS, Word}; use fxhash::FxHashMap; use num::Complex; -use ppvm_traits::PauliWordTrait; /// Magnitude below which an expanded Pauli coefficient is treated as /// cancellation noise and dropped. @@ -105,7 +104,7 @@ pub(crate) fn comm_product(h: &Word, p: &Word) -> (Word /// is `ω · out` with `ω = i^phase`. #[inline(always)] pub(crate) fn pauli_mul(p: &Word, q: &Word) -> (Word, u8) { - let mut out = Word::::new(p.n_qubits()); + let mut out = Word::::default(); let mut sign_count: u32 = 0; let mut imag_count: u32 = 0; for i in 0..C { @@ -120,6 +119,5 @@ pub(crate) fn pauli_mul(p: &Word, q: &Word) -> (Word, u out.xbits.data[i] = a ^ c; out.zbits.data[i] = b ^ d; } - out.rehash(); (out, ((2 * sign_count + imag_count) & 3) as u8) } diff --git a/crates/ppvm-lindblad/src/basis.rs b/crates/ppvm-lindblad/src/basis.rs index 09e162177..2f3af1092 100644 --- a/crates/ppvm-lindblad/src/basis.rs +++ b/crates/ppvm-lindblad/src/basis.rs @@ -3,11 +3,12 @@ //! Basis-level `L*` operators: in-basis generator and off-basis leakage. -use crate::Error; +use crate::candidates::{CandidateBuffer, CandidateProposer, StreamingProposer}; use crate::sector::Sector; use crate::spec::LindbladSpec; use crate::truncate::{cap_map_to_room, order_by_desc_mag}; use crate::word::{Word, word_hash}; +use crate::{Error, ShortlistScoreProposer}; use fxhash::{FxBuildHasher, FxHashMap, FxHashSet}; use num::Complex; use rayon::prelude::*; @@ -16,6 +17,63 @@ use rayon::prelude::*; /// into the live map (and the room-cap applied) once per chunk. const CHUNK_SIZE: usize = 4096; +/// Reusable scratch for adaptive real-coefficient evolution. +/// +/// Keep one workspace per trajectory. It retains candidate lookup and source-order +/// storage, generator connections and exponential scratch. Candidate words occupy +/// the main state array's tail; streaming fallback retains its generation buffers. +#[derive(Default)] +pub struct PcStepWorkspace { + // The shortlist strategy also owns the streaming scratch used by the default path. + pub(crate) proposer: ShortlistScoreProposer, + pub(crate) generator: Option, + pub(crate) taylor: Vec, + pub(crate) coefficient_auxiliary: Vec, +} + +impl PcStepWorkspace { + /// Requested vector, shortlist-index and generator capacities retained between + /// steps. Excludes the main state, the protected-word hash set, and metadata + /// owned by a separately supplied candidate proposer. + pub fn allocation_bytes(&self) -> [(&'static str, usize); 12] { + let [terms, order, index] = self.proposer.allocation_bytes(); + [ + ("shortlist_terms", terms), + ("source_order", order), + ("shortlist_index", index), + ("candidate_words", 0), + ("candidate_rates", 0), + ( + "candidate_state", + self.proposer.streaming.indexed.state.capacity(), + ), + ( + "candidate_free", + self.proposer.streaming.indexed.free.capacity() * 4, + ), + ("candidate_rank", 0), + ( + "generation", + self.proposer + .streaming + .outputs + .iter() + .map(|v| v.capacity() * std::mem::size_of::<(Word, f64)>()) + .sum(), + ), + ( + "generator", + self.generator.as_ref().map_or(0, |p| p.allocation_bytes()), + ), + ("numeric_scratch", self.taylor.capacity() * 8), + ( + "coefficient_auxiliary", + self.coefficient_auxiliary.capacity() * 8, + ), + ] + } +} + /// Build a `word → row` map for a basis assumed to contain unique Pauli /// words; debug-asserts the uniqueness invariant. pub fn build_basis_index(basis: &[Word]) -> FxHashMap, u32> { @@ -69,61 +127,12 @@ impl LindbladSpec { b: coeffs.len(), }); } - // Hash-only membership tables: storing 8-byte `u64` keys instead - // of 48-byte Words shrinks the in-basis structure ~6×, keeping it - // in L3 (and often L2) at basis sizes where the full-Word version - // would spill to DRAM. - let in_basis: FxHashMap = basis.iter().map(|w| (word_hash(w), ())).collect(); - let protected_set: FxHashMap = - protected.iter().map(|w| (word_hash(w), ())).collect(); - - let order = order_by_desc_mag(coeffs); - let room = max_basis.saturating_sub(basis.len()); - let n_qubits = self.n_qubits(); - let mut merged: FxHashMap, f64> = FxHashMap::default(); - for chunk_indices in order.chunks(CHUNK_SIZE) { - let local: Vec, f64)>> = chunk_indices - .par_iter() - .map_init( - || { - ( - Vec::::with_capacity(n_qubits), - Vec::::with_capacity(128), - FxHashMap::, Complex>::with_capacity_and_hasher( - 128, - FxBuildHasher::default(), - ), - ) - }, - |(s1, s2, lm), &i| { - let p = &basis[i]; - let c = coeffs[i]; - let terms = self.compute_action_terms(p, s1, s2, lm); - let mut out = Vec::with_capacity(terms.len()); - for (w, v) in terms.iter() { - let h = word_hash(w); - if !in_basis.contains_key(&h) && !protected_set.contains_key(&h) { - out.push((*w, c * *v)); - } - } - out - }, - ) - .collect(); - for v in local { - for (k, val) in v { - *merged.entry(k).or_insert(0.0) += val; - } - } - cap_map_to_room(&mut merged, room); - } - // Rate-based admission: keep only candidates whose leakage rate - // exceeds `tau_add`. `tau_add = 0` admits everything except exact - // zeros. - Ok(merged - .into_iter() - .filter(|(_, c)| c.abs() > tau_add) - .collect()) + let mut words = Vec::new(); + let mut numeric = Vec::new(); + let mut buffer = CandidateBuffer::separate(basis, coeffs, &mut words, &mut numeric); + let mut proposer = StreamingProposer::default(); + proposer.propose(self, &mut buffer, protected, max_basis)?; + Ok(proposer.indexed.selected_rates(&buffer, tau_add).collect()) } /// Sparse generator matrix in COO form: returns `(row, col, val)` @@ -329,3 +338,54 @@ impl LindbladSpec { Ok(merged.into_iter().filter(|(_, c)| c.norm() > 0.0).collect()) } } + +#[cfg(test)] +mod tail_tests { + use super::*; + use crate::word_from_codes; + + #[test] + fn tail_leakage_preserves_multiple_chunk_caps_and_source_order() { + let n = 7; + let mut h = Vec::new(); + for site in 0..n - 1 { + for symbol in ['X', 'Y'] { + let mut term = vec!['I'; n]; + term[site] = symbol; + term[site + 1] = symbol; + h.push((term.into_iter().collect(), 0.3)); + } + } + let reference = LindbladSpec::<2>::new(n, &h, &[]).unwrap(); + let tail = LindbladSpec::<2>::new(n, &h, &[]).unwrap(); + let basis: Vec<_> = (0..14000) + .filter(|i| i % 3 != 0) + .map(|i| { + word_from_codes::<2>( + &(0..n) + .map(|q| ((i >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let coefficients: Vec<_> = (0..basis.len()).map(|i| (i as f64).sin()).collect(); + for room in [0, 100, 2000] { + let expected = reference + .leakage_reference_with_prune(&basis, &coefficients, &[], basis.len() + room, 0.0) + .unwrap(); + for threads in [1, 2, 6] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let actual = pool + .install(|| { + tail.leakage_with_prune(&basis, &coefficients, &[], basis.len() + room, 0.0) + }) + .unwrap(); + assert_eq!(actual, expected, "room={room}, threads={threads}"); + } + } + } +} diff --git a/crates/ppvm-lindblad/src/candidates.rs b/crates/ppvm-lindblad/src/candidates.rs new file mode 100644 index 000000000..70744966d --- /dev/null +++ b/crates/ppvm-lindblad/src/candidates.rs @@ -0,0 +1,295 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Candidate proposal using a borrowed view of the adaptive state's spare storage. + +use crate::indexed_candidates::IndexedCandidates; +use crate::truncate::order_by_desc_mag; +use crate::word::{Word, word_hash}; +use crate::{Error, LindbladSpec}; +use fxhash::FxHashMap; +use rayon::prelude::*; + +const CHUNK_SIZE: usize = 4096; +const GENERATION_BATCH: usize = 2048; + +/// Borrowed candidate storage with a read-only active prefix. +/// +/// Candidate words occupy the main word array's tail. Numeric scratch is reused +/// by the exponential after admission. Dropping an unfinished view discards its +/// tail without changing the active words or coefficients. +pub struct CandidateBuffer<'a, const C: usize = { crate::W_CHUNKS }> { + pub(crate) words: &'a mut Vec>, + pub(crate) numeric: &'a mut Vec, + pub(crate) prefix: usize, + source: Option<&'a [Word]>, + coefficients: &'a [f64], +} + +impl<'a, const C: usize> CandidateBuffer<'a, C> { + /// View the active sum and use its spare word capacity for candidates. + /// Numeric scratch is cleared while retaining its allocation. + /// + /// # Panics + /// Panics when the active word and coefficient lengths differ. + pub fn new( + words: &'a mut Vec>, + coefficients: &'a [f64], + numeric: &'a mut Vec, + ) -> Self { + assert_eq!(words.len(), coefficients.len()); + let prefix = words.len(); + numeric.clear(); + Self { + words, + numeric, + prefix, + source: None, + coefficients, + } + } + + pub(crate) fn separate( + source: &'a [Word], + coefficients: &'a [f64], + words: &'a mut Vec>, + numeric: &'a mut Vec, + ) -> Self { + assert_eq!(source.len(), coefficients.len()); + words.clear(); + numeric.clear(); + Self { + words, + numeric, + prefix: 0, + source: Some(source), + coefficients, + } + } + + /// Active strings, whose order and values remain fixed during proposal. + pub fn basis(&self) -> &[Word] { + self.source.unwrap_or(&self.words[..self.prefix]) + } + + /// Coefficients corresponding to [`Self::basis`]. + pub fn coefficients(&self) -> &[f64] { + self.coefficients + } + + /// Candidate words in the writable tail, excluding the active prefix. + pub fn candidates(&self) -> &[Word] { + &self.words[self.prefix..] + } + + /// Mutable candidate words; the active prefix is not exposed for mutation. + pub fn candidates_mut(&mut self) -> &mut [Word] { + &mut self.words[self.prefix..] + } + + /// Append a candidate word, growing the backing array only when necessary. + /// Rates and any ranking metadata are managed by the proposer. + pub fn push(&mut self, word: Word) { + if self.words.len() == self.words.capacity() { + self.words + .reserve_exact((self.words.capacity() / 4).max(32)); + } + self.words.push(word); + } + + /// Reusable numeric storage for rates, ranking and permutation scratch. + /// Its contents have no meaning after the proposal has been admitted. + pub fn numeric_scratch(&mut self) -> &mut Vec { + self.numeric + } + + pub(crate) fn scoring_parts(&mut self) -> (&[Word], &[Word], &[f64], &mut Vec) { + let (basis, candidates) = self.words.split_at(self.prefix); + ( + self.source.unwrap_or(basis), + candidates, + self.coefficients, + self.numeric, + ) + } + + pub(crate) fn finish(mut self, count: usize) { + assert!(count <= self.candidates().len()); + self.prefix += count; + self.words.truncate(self.prefix); + } +} + +impl Drop for CandidateBuffer<'_, C> { + fn drop(&mut self) { + self.words.truncate(self.prefix); + self.numeric.clear(); + } +} + +/// Strategy for proposing external strings for a real predictor-corrector step. +/// +/// Selected words must be distinct and outside the active basis and protected +/// set, with at most `maximum.saturating_sub(buffer.basis().len())` survivors. Implementations +/// may retain their own metadata between calls, but cannot retain the borrowed +/// buffer. Proposal and selection are separate to preserve per-phase timings. +/// Strategies must be sendable because a step may execute in a configured thread pool. +pub trait CandidateProposer: Send { + /// Accumulate candidate information using the supplied tail and scratch. + /// `maximum` bounds the total enlarged basis, including the active prefix. + fn propose( + &mut self, + model: &LindbladSpec, + buffer: &mut CandidateBuffer<'_, C>, + protected: &[Word], + maximum: usize, + ) -> Result<(), Error>; + + /// Compact selected words into the beginning of the candidate tail and return + /// their count. Only rates with magnitude strictly above `threshold` qualify. + /// Admission initializes the corresponding coefficients to zero. + fn select(&mut self, buffer: &mut CandidateBuffer<'_, C>, threshold: f64) -> usize; +} + +/// Chunk-capped leakage proposal with position-indexed candidate accumulation. +/// +/// Sources are processed in descending coefficient magnitude. Signed sums and +/// hashtable visitation preserve deterministic admission and tie ordering. +#[derive(Default)] +pub struct StreamingProposer { + pub(crate) indexed: IndexedCandidates, + pub(crate) outputs: Vec, f64)>>, +} + +impl CandidateProposer for StreamingProposer { + fn propose( + &mut self, + model: &LindbladSpec, + buffer: &mut CandidateBuffer<'_, C>, + protected: &[Word], + maximum: usize, + ) -> Result<(), Error> { + self.indexed.reset(); + let in_basis: FxHashMap = + buffer.basis().iter().map(|w| (word_hash(w), ())).collect(); + let protected_set: FxHashMap = + protected.iter().map(|w| (word_hash(w), ())).collect(); + let order = order_by_desc_mag(buffer.coefficients()); + let room = maximum.saturating_sub(buffer.basis().len()); + self.outputs + .resize_with(GENERATION_BATCH.div_ceil(64), Vec::new); + for chunk in order.chunks(CHUNK_SIZE) { + // Batches bound temporary words without changing the source chunk's pruning cadence. + for sources in chunk.chunks(GENERATION_BATCH) { + let basis = buffer.basis(); + let coefficients = buffer.coefficients(); + self.outputs + .par_iter_mut() + .zip(sources.par_chunks(64)) + .for_each_init( + crate::direct_expm::ActionScratch::new, + |scratch, (output, indices)| { + output.clear(); + for &index in indices { + scratch.fill(model, &basis[index], true); + for &(hash, word, value) in &scratch.terms { + if !in_basis.contains_key(&hash) + && !protected_set.contains_key(&hash) + { + output.push((word, coefficients[index] * value)); + } + } + } + }, + ); + // Serial source order preserves coherent sums and admission ties. + for output in self.outputs.iter().take(sources.len().div_ceil(64)) { + for &(word, value) in output { + self.indexed.add(buffer, word, value); + } + } + } + self.indexed.cap(buffer, room); + } + for output in &mut self.outputs { + output.clear(); + } + Ok(()) + } + + fn select(&mut self, buffer: &mut CandidateBuffer<'_, C>, threshold: f64) -> usize { + let count = self.indexed.select(buffer, threshold); + self.indexed.reset(); + count + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::truncate::cap_map_to_room; + use crate::word_from_codes; + + #[test] + fn borrowed_candidates_preserve_reference_caps_ties_and_visitation() { + let seed = word_from_codes::<2>(&[1; 6]).unwrap(); + let mut words = vec![seed]; + let mut numeric = Vec::new(); + let mut buffer = CandidateBuffer::new(&mut words, &[0.5], &mut numeric); + let mut indexed = IndexedCandidates::default(); + let mut reference = FxHashMap::default(); + for chunk in 0..12 { + for i in 0..500 { + let code = (i * 37 + chunk * 11) % 4096; + let word = word_from_codes::<2>( + &(0..6) + .map(|q| ((code >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap(); + let value = ((i + chunk) % 9) as f64 - 4.0; + *reference.entry(word).or_insert(0.0) += value; + indexed.add(&mut buffer, word, value); + } + cap_map_to_room(&mut reference, 300); + indexed.cap(&mut buffer, 300); + assert_eq!(buffer.basis(), &[seed]); + } + let expected: Vec<_> = reference + .into_iter() + .filter(|(_, value)| value.abs() > 0.0) + .collect(); + assert_eq!( + indexed.selected_rates(&buffer, 0.0).collect::>(), + expected + ); + let count = indexed.select(&mut buffer, 1.5); + buffer.finish(count); + assert_eq!( + words, + std::iter::once(seed) + .chain( + expected + .iter() + .filter(|(_, value)| value.abs() > 1.5) + .map(|(word, _)| *word) + ) + .collect::>() + ); + assert!(numeric.is_empty()); + } + + #[test] + fn nonfinite_threshold_rejects_tail_and_preserves_prefix() { + let seed = word_from_codes::<2>(&[1, 0]).unwrap(); + let candidate = word_from_codes::<2>(&[0, 2]).unwrap(); + let mut words = vec![seed]; + let mut numeric = Vec::new(); + let mut buffer = CandidateBuffer::new(&mut words, &[1.0], &mut numeric); + let mut indexed = IndexedCandidates::default(); + indexed.add(&mut buffer, candidate, 0.5); + let count = indexed.select(&mut buffer, f64::NAN); + buffer.finish(count); + assert_eq!(words, vec![seed]); + } +} diff --git a/crates/ppvm-lindblad/src/compact_index.rs b/crates/ppvm-lindblad/src/compact_index.rs new file mode 100644 index 000000000..cbe6b2fcc --- /dev/null +++ b/crates/ppvm-lindblad/src/compact_index.rs @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Collision-safe basis lookup storing positions rather than duplicate words. + +use crate::word::{Word, word_hash}; +use hashbrown::{HashTable, hash_table::Entry}; + +pub(crate) struct CompactBasisIndex<'a, const C: usize> { + basis: &'a [Word], + table: HashTable, +} + +impl<'a, const C: usize> CompactBasisIndex<'a, C> { + pub(crate) fn new(basis: &'a [Word]) -> Self { + assert!( + basis.len() < u32::MAX as usize, + "basis exceeds u32 positions" + ); + let mut table = HashTable::with_capacity(basis.len()); + for (position, word) in basis.iter().enumerate() { + match table.entry( + word_hash(word), + |&index| basis[index as usize] == *word, + |&index| word_hash(&basis[index as usize]), + ) { + Entry::Vacant(entry) => { + entry.insert(position as u32); + } + Entry::Occupied(_) => panic!("duplicate Pauli word in basis"), + } + } + Self { basis, table } + } + + /// [`Self::get`] for a word whose [`word_hash`] the caller already has. + pub(crate) fn get_hashed(&self, word: &Word, hash: u64) -> Option { + self.table + .find(hash, |&index| self.basis[index as usize] == *word) + .copied() + } + + pub(crate) fn get(&self, word: &Word) -> Option { + // Control-byte fingerprints filter probes before loading full borrowed words. + self.table + .find(word_hash(word), |&index| { + self.basis[index as usize] == *word + }) + .copied() + } +} diff --git a/crates/ppvm-lindblad/src/config.rs b/crates/ppvm-lindblad/src/config.rs index 3fe32a71d..f3d756a50 100644 --- a/crates/ppvm-lindblad/src/config.rs +++ b/crates/ppvm-lindblad/src/config.rs @@ -3,6 +3,16 @@ //! Configuration objects for the predictor-corrector stepper. +/// Built-in candidate generation for real-coefficient predictor-corrector steps. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub enum CandidateStrategy { + /// Chunk-capped signed leakage accumulation. + #[default] + Streaming, + /// Shortlist individual contributions, then score their complete signed leakage. + Shortlist, +} + /// Truncation and execution policy for a single predictor-corrector step /// ([`crate::LindbladSpec::pc_step`], [`crate::LindbladSpec::pc_step_timed`]). /// @@ -14,6 +24,9 @@ /// scheme; `tau_add` is a wall optimization at most. #[derive(Debug, Clone, Copy)] pub struct PcStepConfig { + /// Candidate generation for real-coefficient steps. Explicit proposer calls + /// use their supplied strategy instead of this field. + pub proposer: CandidateStrategy, /// Hard rank cap on the retained basis: after the corrector, only the /// top-`max_basis` strings by `|coeff|` are kept (protected words always /// survive). The primary convergence dial — verify by re-running at 2×. @@ -43,12 +56,46 @@ pub struct PcStepConfig { pub num_threads: Option, } +impl PcStepConfig { + /// Set the working-set cap to `ceil(ratio * max_basis)`. + /// + /// The ratio scales the total admission bound, including the retained + /// basis. For example, `max_basis = 1025` and `ratio = 1.5` admit at most + /// 1538 strings before truncation. Existing absolute caps are overwritten. + /// + /// # Errors + /// + /// Returns an error for nonfinite ratios, ratios below one, or caps that + /// cannot be represented as a `usize`. + pub fn with_admit_ratio(mut self, ratio: f64) -> Result { + if !ratio.is_finite() || ratio < 1.0 { + return Err(crate::Error::InvalidAdmissionRatio { ratio }); + } + // Integer multipliers preserve exact counts even beyond f64's integer precision. + if ratio.fract() == 0.0 && ratio < usize::MAX as f64 { + self.admit_basis = Some( + self.max_basis + .checked_mul(ratio as usize) + .ok_or(crate::Error::AdmissionCapOverflow)?, + ); + return Ok(self); + } + let cap = (self.max_basis as f64 * ratio).ceil(); + if !cap.is_finite() || cap as u128 > usize::MAX as u128 { + return Err(crate::Error::AdmissionCapOverflow); + } + self.admit_basis = Some(cap as usize); + Ok(self) + } +} + impl Default for PcStepConfig { /// Uncapped, unfiltered, no pruning: the near-exact reference /// configuration. Production runs should set `max_basis` (and usually /// `admit_basis ≈ 2-3×` it). fn default() -> Self { Self { + proposer: CandidateStrategy::Streaming, max_basis: usize::MAX, admit_basis: None, drop_tol: 0.0, @@ -57,3 +104,76 @@ impl Default for PcStepConfig { } } } + +#[cfg(test)] +mod tests { + use super::PcStepConfig; + + #[test] + fn fractional_admission_rounds_up() { + for (basis, ratio, expected) in [ + (1025, 1.5, 1538), + (1025, 1.234, 1265), + (1024, 3.0, 3072), + (1024, 6.0, 6144), + (1024, 1.0, 1024), + ] { + let cfg = PcStepConfig { + max_basis: basis, + ..Default::default() + } + .with_admit_ratio(ratio) + .unwrap(); + assert_eq!(cfg.admit_basis, Some(expected)); + assert_eq!(cfg.max_basis, basis); + } + } + + #[test] + fn invalid_admission_ratios_and_overflow_are_rejected() { + for ratio in [ + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, + -1.0, + 0.0, + 0.5, + f64::MAX, + ] { + assert!( + PcStepConfig { + max_basis: 1024, + ..Default::default() + } + .with_admit_ratio(ratio) + .is_err() + ); + } + assert!(PcStepConfig::default().with_admit_ratio(2.0).is_err()); + assert_eq!( + PcStepConfig::default() + .with_admit_ratio(1.0) + .unwrap() + .admit_basis, + Some(usize::MAX) + ); + } + + #[test] + fn admission_ratio_preserves_other_policy_settings() { + let cfg = PcStepConfig { + max_basis: 10, + admit_basis: Some(40), + drop_tol: 0.1, + tau_add: Some(0.2), + num_threads: Some(2), + ..Default::default() + } + .with_admit_ratio(1.5) + .unwrap(); + assert_eq!(cfg.admit_basis, Some(15)); + assert_eq!(cfg.drop_tol, 0.1); + assert_eq!(cfg.tau_add, Some(0.2)); + assert_eq!(cfg.num_threads, Some(2)); + } +} diff --git a/crates/ppvm-lindblad/src/direct_expm.rs b/crates/ppvm-lindblad/src/direct_expm.rs new file mode 100644 index 000000000..d51fbfd3a --- /dev/null +++ b/crates/ppvm-lindblad/src/direct_expm.rs @@ -0,0 +1,319 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Row-gather action for a Hamiltonian plus diagonal Pauli dissipators. + +use crate::compact_index::CompactBasisIndex; +use crate::{LindbladSpec, Word}; +use fxhash::FxHashMap; +use num::Complex; + +pub(crate) struct DirectOp<'a, const C: usize> { + pub(crate) spec: &'a LindbladSpec, + pub(crate) basis: &'a [Word], + pub(crate) index: CompactBasisIndex<'a, C>, +} + +pub(crate) struct ActionScratch { + support: Vec, + candidates: Vec, + local: FxHashMap, Complex>, + active: Vec, + /// Output words with their [`crate::word::word_hash`] and generator values. + pub(crate) terms: Vec<(u64, Word, f64)>, + pub(crate) row: Vec<(u32, f64)>, +} + +impl ActionScratch { + pub(crate) fn new() -> Self { + Self { + support: Vec::new(), + candidates: Vec::new(), + local: FxHashMap::default(), + active: Vec::new(), + terms: Vec::new(), + row: Vec::new(), + } + } + pub(crate) fn fill(&mut self, spec: &LindbladSpec, word: &Word, canonical: bool) { + if let Some(plan) = &spec.pauli_plan { + plan.fill_keyed(word, &mut self.active, &mut self.terms); + if canonical { + self.terms + .sort_unstable_by(|(ha, a, _), (hb, b, _)| ha.cmp(hb).then_with(|| a.cmp(b))); + } + return; + } + + let terms = spec.compute_action_terms( + word, + &mut self.support, + &mut self.candidates, + &mut self.local, + ); + self.terms.clear(); + self.terms.extend( + terms + .into_iter() + .map(|(word, value)| (crate::word::word_hash(&word), word, value)), + ); + } +} + +impl DirectOp<'_, C> { + pub(crate) fn gather_row( + &self, + position: usize, + scratch: &mut ActionScratch, + canonical: bool, + ) -> (f64, f64) { + let word = &self.basis[position]; + scratch.fill(self.spec, word, canonical); + scratch.row.clear(); + let mut raw = 0.0; + let mut diagonal = 0.0; + for &(hash, parent, value) in &scratch.terms { + raw += value.abs(); + if parent == *word { + diagonal = value; + } + if let Some(column) = self.index.get_hashed(&parent, hash) { + scratch + .row + .push((column, if parent == *word { value } else { -value })); + } + } + scratch.row.sort_unstable_by_key(|entry| entry.0); + (raw, diagonal) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{JumpInput, PcStepConfig, word_from_codes}; + use quspin_types::LinearOperator; + + #[test] + fn row_gather_matches_projected_signed_generator_across_threads() { + let h = vec![ + ("XXII".into(), 0.3), + ("IYYI".into(), -0.7), + ("ZIIX".into(), 0.2), + ("XXII".into(), -0.1), + ]; + let jumps = vec![JumpInput { + lincomb: vec![("IZII".into(), Complex::new(1.0, 0.0))], + rate: 0.4, + }]; + let spec = LindbladSpec::<2>::new(4, &h, &jumps).unwrap(); + let basis: Vec<_> = (0..256usize) + .filter(|i| i % 3 != 0) + .map(|i| { + word_from_codes::<2>( + &(0..4) + .map(|q| ((i >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let input: Vec = (0..basis.len()) + .map(|i| if i % 7 == 0 { 0.0 } else { (i as f64).sin() }) + .collect(); + let mut expected = vec![0.0; basis.len()]; + for (row, col, value) in spec.generator(&basis) { + if input[col] != 0.0 { + expected[row] += value * input[col]; + } + } + let (mat, _) = crate::mf_expm::build_row_cache(&spec, &basis, None); + let op = crate::mf_expm::CsrOp { mat: &mat }; + for threads in [1, 2, 6] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let mut actual = vec![0.0; basis.len()]; + pool.install(|| op.dot(true, &input, &mut actual)).unwrap(); + assert_eq!(actual, expected, "thread count {threads}"); + pool.install(|| op.dot(false, &input, &mut actual)).unwrap(); + assert_eq!(actual, expected.iter().map(|v| v + v).collect::>()); + } + } + + #[test] + fn default_exponential_preserves_nonuniform_adaptive_steps() { + let h = vec![ + ("XXI".into(), 0.3), + ("IYY".into(), 0.7), + ("ZII".into(), -0.2), + ]; + let cached = LindbladSpec::<2>::new(3, &h, &[]).unwrap(); + let direct = LindbladSpec::<2>::new(3, &h, &[]).unwrap(); + let seed = word_from_codes::<2>(&[3, 0, 0]).unwrap(); + let mut a = vec![seed]; + let mut b = a.clone(); + let mut ca = vec![1.0]; + let mut cb = ca.clone(); + let mut reused_basis = a.clone(); + let mut reused_coeffs = ca.clone(); + let mut workspace = crate::PcStepWorkspace::default(); + let tail = LindbladSpec::<2>::new(3, &h, &[]).unwrap(); + let mut tail_basis = a.clone(); + let mut tail_coeffs = ca.clone(); + let mut tail_workspace = crate::PcStepWorkspace::default(); + let cfg = PcStepConfig { + max_basis: 8, + admit_basis: Some(24), + num_threads: Some(2), + ..Default::default() + }; + for _ in 0..10 { + cached + .pc_step_reference(&mut a, &mut ca, 0.1, &[seed], &cfg) + .unwrap(); + direct.pc_step(&mut b, &mut cb, 0.1, &[seed], &cfg).unwrap(); + direct + .pc_step_with_workspace( + &mut reused_basis, + &mut reused_coeffs, + 0.1, + &[seed], + &cfg, + &mut workspace, + ) + .unwrap(); + tail.pc_step_with_workspace( + &mut tail_basis, + &mut tail_coeffs, + 0.1, + &[seed], + &cfg, + &mut tail_workspace, + ) + .unwrap(); + assert_eq!(a, tail_basis); + assert_eq!(ca, tail_coeffs); + assert_eq!(a, reused_basis); + assert_eq!(ca, reused_coeffs); + assert_eq!(a, b); + assert_eq!(ca, cb); + } + } + + #[test] + fn default_workspace_preserves_predictor_corrector_and_noise() { + let h = vec![ + ("XXI".into(), 0.3), + ("IYY".into(), -0.3), + ("ZII".into(), 0.3), + ]; + let jumps = vec![JumpInput { + lincomb: vec![("IZI".into(), Complex::new(1.0, 0.0))], + rate: 0.2, + }]; + + let cached = LindbladSpec::<2>::new(3, &h, &jumps).unwrap(); + let optimized = LindbladSpec::<2>::new(3, &h, &jumps).unwrap(); + let seed = word_from_codes::<2>(&[3, 0, 0]).unwrap(); + let mut a = vec![seed]; + let mut b = a.clone(); + let mut ca = vec![1.0]; + let mut cb = ca.clone(); + b.reserve_exact(64 - b.len()); + cb.reserve_exact(64 - cb.len()); + let basis_storage = b.as_ptr(); + let coefficient_storage = cb.as_ptr(); + for step in 0..12 { + let cfg = PcStepConfig { + max_basis: if step == 0 { usize::MAX } else { 8 + step % 3 }, + admit_basis: if step == 0 { None } else { Some(24 + step % 4) }, + num_threads: Some(2), + ..Default::default() + }; + cached + .pc_step_reference(&mut a, &mut ca, 0.1, &[seed], &cfg) + .unwrap(); + optimized + .pc_step(&mut b, &mut cb, 0.1, &[seed], &cfg) + .unwrap(); + assert_eq!(b.as_ptr(), basis_storage); + assert_eq!(cb.as_ptr(), coefficient_storage); + assert_eq!(a, b); + assert_eq!(ca, cb); + } + } + + #[test] + fn default_workspace_matches_reference_beyond_one_hundred_twenty_eight_qubits() { + let n = 129; + let term = |symbol: char, sites: &[usize]| { + let mut symbols = vec!['I'; n]; + for &site in sites { + symbols[site] = symbol; + } + symbols.into_iter().collect::() + }; + let h = vec![ + (term('X', &[0, 128]), 0.25), + (term('Y', &[64, 128]), -0.25), + (term('Z', &[0]), 0.25), + ]; + let jumps = vec![JumpInput { + lincomb: vec![(term('Z', &[64]), Complex::new(1.0, 0.0))], + rate: 0.2, + }]; + let reference = LindbladSpec::<4>::new(n, &h, &jumps).unwrap(); + let optimized = LindbladSpec::<4>::new(n, &h, &jumps).unwrap(); + let mut codes = vec![0; n]; + codes[128] = 2; + let seed = word_from_codes::<4>(&codes).unwrap(); + let mut a = vec![seed]; + let mut b = a.clone(); + let mut ca = vec![1.0]; + let mut cb = ca.clone(); + let mut workspace = crate::PcStepWorkspace::default(); + let cfg = PcStepConfig { + max_basis: 16, + admit_basis: Some(48), + num_threads: Some(2), + ..Default::default() + }; + for _ in 0..8 { + reference + .pc_step_reference(&mut a, &mut ca, 0.1, &[seed], &cfg) + .unwrap(); + optimized + .pc_step_with_workspace(&mut b, &mut cb, 0.1, &[seed], &cfg, &mut workspace) + .unwrap(); + let mut expected: Vec<_> = a.iter().copied().zip(ca.iter().copied()).collect(); + let mut actual: Vec<_> = b.iter().copied().zip(cb.iter().copied()).collect(); + expected.sort_unstable_by_key(|entry| entry.0); + actual.sort_unstable_by_key(|entry| entry.0); + assert_eq!(expected.len(), actual.len()); + for ((word_a, value_a), (word_b, value_b)) in expected.iter().zip(&actual) { + assert_eq!(word_a, word_b); + assert!((value_a - value_b).abs() < 2e-13); + } + } + } + + #[test] + fn compact_lookup_checks_words_and_wide_registers() { + let basis: Vec<_> = (0..192) + .map(|site| { + let mut codes = vec![0; 128]; + codes[site / 3] = (site % 3 + 1) as u8; + word_from_codes::<2>(&codes).unwrap() + }) + .collect(); + let index = CompactBasisIndex::new(&basis); + for (i, word) in basis.iter().enumerate() { + assert_eq!(index.get(word), Some(i as u32)); + } + assert_eq!(index.get(&word_from_codes::<2>(&[2; 128]).unwrap()), None); + let wide = vec![word_from_codes::<8>(&vec![1; 512]).unwrap()]; + assert_eq!(CompactBasisIndex::new(&wide).get(&wide[0]), Some(0)); + } +} diff --git a/crates/ppvm-lindblad/src/error.rs b/crates/ppvm-lindblad/src/error.rs index 1e10c53dd..6b11e46bf 100644 --- a/crates/ppvm-lindblad/src/error.rs +++ b/crates/ppvm-lindblad/src/error.rs @@ -46,6 +46,12 @@ pub enum Error { n: usize, m: usize, }, + /// The admission ratio must be finite and at least one. + InvalidAdmissionRatio { + ratio: f64, + }, + /// The rounded admission cap exceeds the representable allocation size. + AdmissionCapOverflow, Internal(String), } @@ -85,6 +91,13 @@ impl fmt::Display for Error { f, "kossakowski K must be Hermitian; K[{n}][{m}] ≠ conj(K[{m}][{n}])" ), + Error::InvalidAdmissionRatio { ratio } => { + write!( + f, + "admission ratio must be finite and at least 1; got {ratio}" + ) + } + Error::AdmissionCapOverflow => write!(f, "admission cap overflow"), Error::Internal(msg) => write!(f, "internal error: {msg}"), } } diff --git a/crates/ppvm-lindblad/src/indexed_candidates.rs b/crates/ppvm-lindblad/src/indexed_candidates.rs new file mode 100644 index 000000000..93b0960ed --- /dev/null +++ b/crates/ppvm-lindblad/src/indexed_candidates.rs @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Stable candidate positions and chunk-capped ranking metadata. + +use crate::candidates::CandidateBuffer; +use crate::word::{Word, word_hash}; +use hashbrown::HashTable; + +#[derive(Default)] +pub(crate) struct IndexedCandidates { + pub(crate) state: Vec, + table: HashTable, + pub(crate) free: Vec, +} + +impl IndexedCandidates { + pub(crate) fn reset(&mut self) { + self.state.clear(); + self.free.clear(); + // Table growth history determines admission order, including magnitude ties. + self.table = HashTable::new(); + } + + pub(crate) fn add( + &mut self, + buffer: &mut CandidateBuffer<'_, C>, + word: Word, + value: f64, + ) { + let hash = word_hash(&word); + if let Some(&index) = self.table.find(hash, |&index| { + buffer.words[buffer.prefix + index as usize] == word + }) { + buffer.numeric[index as usize] += value; + return; + } + let index = if let Some(index) = self.free.pop() { + buffer.words[buffer.prefix + index as usize] = word; + buffer.numeric[index as usize] = 0.0 + value; + self.state[index as usize] = 1; + index + } else { + if buffer.words.len() == buffer.words.capacity() { + buffer + .words + .reserve_exact((buffer.words.capacity() / 4).max(32)); + buffer + .numeric + .reserve_exact(buffer.words.capacity() - buffer.prefix - buffer.numeric.len()); + self.state + .reserve_exact(buffer.words.capacity() - buffer.prefix - self.state.len()); + } + let index = u32::try_from(buffer.numeric.len()).expect("candidate count exceeds u32"); + buffer.words.push(word); + buffer.numeric.push(0.0 + value); + self.state.push(1); + index + }; + self.table.insert_unique(hash, index, |&index| { + word_hash(&buffer.words[buffer.prefix + index as usize]) + }); + } + + pub(crate) fn cap(&mut self, buffer: &mut CandidateBuffer<'_, C>, room: usize) { + if self.table.len() <= room { + return; + } + if room == 0 { + buffer.words.truncate(buffer.prefix); + buffer.numeric.clear(); + self.state.clear(); + self.table.clear(); + self.free.clear(); + return; + } + let count = buffer.numeric.len(); + buffer.numeric.reserve_exact(self.table.len()); + // The rate prefix remains live while the tail holds nonnegative ranking keys. + for index in 0..count { + if self.state[index] != 0 { + buffer.numeric.push(buffer.numeric[index].abs()); + } + } + let (rates, keys) = buffer.numeric.split_at_mut(count); + let position = keys.len() - room; + let cutoff = keys + .select_nth_unstable_by_key(position, |v| v.to_bits()) + .1 + .to_bits(); + let above = keys.iter().filter(|v| v.to_bits() > cutoff).count(); + let mut ties = room - above; + for (rate, state) in rates.iter().zip(&mut self.state) { + if *state != 0 { + let magnitude = rate.abs().to_bits(); + *state = if magnitude > cutoff { + 3 + } else if magnitude == cutoff { + 2 + } else { + 1 + }; + } + } + self.table.retain(|index| { + let state = &mut self.state[*index as usize]; + if *state == 3 { + return true; + } + let keep = *state == 2 && ties > 0; + ties -= keep as usize; + if !keep { + *state = 0; + self.free.push(*index); + } + keep + }); + buffer.numeric.truncate(count); + } + + pub(crate) fn selected_rates<'a, const C: usize>( + &'a self, + buffer: &'a CandidateBuffer<'_, C>, + threshold: f64, + ) -> impl Iterator, f64)> + 'a { + self.table.iter().filter_map(move |&index| { + let index = index as usize; + (buffer.numeric[index].abs() > threshold) + .then_some((buffer.words[buffer.prefix + index], buffer.numeric[index])) + }) + } + + pub(crate) fn select( + &mut self, + buffer: &mut CandidateBuffer<'_, C>, + threshold: f64, + ) -> usize { + self.table + .retain(|index| buffer.numeric[*index as usize].abs() > threshold); + let slots = buffer.numeric.len(); + buffer.numeric.clear(); + buffer.numeric.resize(2 * slots, 0.0); + let (positions, labels) = buffer.numeric.split_at_mut(slots); + // Integer positions use lossless bit patterns in the shared numeric buffer. + // No floating arithmetic reads these temporary permutation entries. + for index in 0..slots { + let code = f64::from_bits(index as u64); + positions[index] = code; + labels[index] = code; + } + let mut count = 0; + for &original in &self.table { + if self.state[original as usize] == 0 { + continue; + } + let source = positions[original as usize].to_bits() as usize; + buffer + .words + .swap(buffer.prefix + count, buffer.prefix + source); + let displaced = labels[count].to_bits() as usize; + positions[displaced] = f64::from_bits(source as u64); + labels[source] = f64::from_bits(displaced as u64); + positions[original as usize] = f64::from_bits(count as u64); + labels[count] = f64::from_bits(original as u64); + count += 1; + } + count + } +} diff --git a/crates/ppvm-lindblad/src/lib.rs b/crates/ppvm-lindblad/src/lib.rs index 683e34751..b953abfbb 100644 --- a/crates/ppvm-lindblad/src/lib.rs +++ b/crates/ppvm-lindblad/src/lib.rs @@ -27,39 +27,56 @@ //! to `f64` at the boundary (with a debug-only check that `|Im|` is at //! FP noise). //! -//! Pauli strings are stored as [`ppvm_pauli_word::word::PauliWord`] backed by -//! a fixed array of `C` chunks (64-bit, or 32-bit on 32-bit targets) with -//! cached hashes for fast HashMap lookup. The crate is const-generic in `C` +//! Pauli strings are stored as [`Word`]: an X and a Z bit plane of `C` +//! chunks each (64-bit, or 32-bit on 32-bit targets) and nothing else, so a +//! 128-qubit word is 32 bytes. Hashes are computed from the planes when a table +//! needs them. [`GeneralWord`] is the corresponding `ppvm_pauli_word` type taken +//! by `ppvm_pauli_sum` symmetry routines. The crate is const-generic in `C` //! ([`Word`], [`LindbladSpec`]); the default is the 128-qubit width, //! and [`chunks_for`] picks the narrowest of the 128/256/512-qubit widths //! for a register. The hot-path commutator/ //! product loops bypass the higher-level word API and operate directly on //! the raw chunks for speed. +//! +//! Real predictor-corrector steps use [`StreamingProposer`] by default. +//! [`CandidateStrategy::Shortlist`] selects [`ShortlistScoreProposer`], which +//! shortlists individual contributions and scores their complete signed leakage. mod algebra; mod basis; +mod candidates; +mod compact_index; pub mod config; +mod direct_expm; pub mod error; pub(crate) mod expm; +mod indexed_candidates; mod kossakowski; +mod pauli_plan; +#[cfg(test)] +mod reference; mod scalar; pub mod sector; +mod shortlist; mod spec; mod step; mod truncate; +mod unit_expm; mod word; /// Matrix-free / quspin-expm-backed `exp(dt·L*)·b` engine. See module docs. pub(crate) mod mf_expm; -pub use basis::build_basis_index; -pub use config::PcStepConfig; +pub use basis::{PcStepWorkspace, build_basis_index}; +pub use candidates::{CandidateBuffer, CandidateProposer, StreamingProposer}; +pub use config::{CandidateStrategy, PcStepConfig}; pub use error::Error; pub use sector::{Sector, canonicalize_basis_to_rep}; +pub use shortlist::ShortlistScoreProposer; pub use spec::{JumpInput, LindbladSpec}; pub use step::PcStepTimings; pub use word::{ - CHUNK_BITS, MAX_QUBITS, MAX_SUPPORTED_QUBITS, W_CHUNKS, WIDTHS, Word, chunks_for, + CHUNK_BITS, GeneralWord, MAX_QUBITS, MAX_SUPPORTED_QUBITS, W_CHUNKS, WIDTHS, Word, chunks_for, codes_from_word, max_qubits, parse_pauli_string, word_from_codes, }; diff --git a/crates/ppvm-lindblad/src/mf_expm.rs b/crates/ppvm-lindblad/src/mf_expm.rs index cd5da646b..38217c3b0 100644 --- a/crates/ppvm-lindblad/src/mf_expm.rs +++ b/crates/ppvm-lindblad/src/mf_expm.rs @@ -7,7 +7,7 @@ //! //! The in-basis-restricted generator's per-column action is computed ONCE //! per expm call (via -//! [`build_mf_cols`] / [`build_orbit_rep_cols`]), stored row-major, and +//! [`build_row_cache`], [`build_mf_cols`] or [`build_orbit_rep_cols`]), stored row-major, and //! reused across every Krylov/Taylor matvec //! by [`CsrOp`] (a [`quspin_types::LinearOperator`]) fed to //! [`quspin_expm::ExpmOp::from_parts`]. Each matvec is then a cheap, @@ -20,12 +20,17 @@ //! [`LinearOperator::dot`] runs. //! //! `μ`, the trace, and the column 1-norm of `A − μ·I` are computed in the -//! same single action pass as the cache, and turned into an `apply` by the -//! shared [`expm_apply_cached`] tail. The `(m, s)` Taylor partition is +//! same action pass as the cache. Real steps apply the exponential in place +//! with retained work buffers; orbit steps return complex coefficients. The `(m, s)` Taylor partition is //! picked with the tolerance-matched tables in [`crate::expm`]: a relaxed //! `tol=1e-6` table when the PC prunes coarsely (`drop_tol ≥ 1e-4`), else the //! double-precision table (keeping the exact-reference test paths bit-exact). +//! +//! Hamiltonians with Hermitian Pauli jumps admit direct row construction. +//! General dissipators use column actions and stable counting-sort assembly. +//! Both store explicit coefficients in the same CSR representation. +use crate::compact_index::CompactBasisIndex; use crate::scalar::Coeff; use crate::sector::Sector; use crate::{LindbladSpec, Word, build_basis_index, expm}; @@ -241,7 +246,7 @@ where fn build_mf_cols( spec: &LindbladSpec, basis: &[Word], - index: &FxHashMap, u32>, + index: &CompactBasisIndex<'_, C>, ) -> (BlockCsr, PerCol) { build_block_csr(spec, basis.len(), |c, (s1, s2, lm), rows, vals| { let p = &basis[c]; @@ -253,7 +258,7 @@ fn build_mf_cols( if w == p { diag = *v; } - if let Some(&row) = index.get(w) { + if let Some(row) = index.get(w) { rows.push(row); vals.push(*v); } @@ -492,14 +497,212 @@ where v } -/// Compute `exp(dt · M) · coeffs` for the in-basis-restricted generator -/// `M`, matrix-free, via `quspin-expm`. Returns a fresh `Vec` of length -/// `basis.len()`. -/// -/// ONE action pass builds the CSR cache `mat` (reused across every matvec) -/// and, in the same pass, the `(raw, diag)` data the `μ`/1-norm selection -/// needs; [`expm_apply_cached`] does the rest. -pub(crate) fn expm_apply_mf( +pub(crate) enum RealCache { + Coefficients(BlockCsr), + Units(crate::unit_expm::UnitCache), +} + +impl RealCache { + pub(crate) fn allocation_bytes(&self) -> usize { + match self { + Self::Coefficients(cache) => cache.allocation_bytes(), + Self::Units(cache) => cache.allocation_bytes(), + } + } +} + +impl BlockCsr { + pub(crate) fn allocation_bytes(&self) -> usize { + self.ranges.capacity() * std::mem::size_of::>() + + self + .ranges + .iter() + .map(|r| r.offsets.capacity() * 4 + r.cols.capacity() * 4 + r.vals.capacity() * 8) + .sum::() + } +} + +pub(crate) fn real_bounds(per_col: &[(f64, f64)]) -> (f64, f64) { + let mu = per_col.iter().map(|(_, diagonal)| *diagonal).sum::() / per_col.len() as f64; + let norm = per_col + .iter() + .map(|&(raw, diagonal)| raw - diagonal.abs() + (diagonal - mu).abs()) + .fold(0.0_f64, f64::max); + (mu, norm) +} + +/// Build explicit row coefficients in reusable arrays. No coupling-magnitude restriction is imposed. +pub(crate) fn build_row_cache( + spec: &LindbladSpec, + basis: &[Word], + previous: Option>, +) -> (BlockCsr, (f64, f64)) { + let dim = basis.len(); + let shift = range_shift(dim); + let rows = 1 << shift; + let direct = crate::direct_expm::DirectOp { + spec, + basis, + index: CompactBasisIndex::new(basis), + }; + // A Hamiltonian commutator has zero diagonal for arbitrary coupling strengths. + let diagonal = spec + .pauli_plan + .as_ref() + .is_none_or(|plan| plan.has_diagonal_noise()); + let mut ranges = previous.map(|cache| cache.ranges).unwrap_or_default(); + ranges.resize_with(dim.div_ceil(rows), || CsrRange { + offsets: Vec::new(), + cols: Vec::new(), + vals: Vec::new(), + }); + let metadata: Vec<_> = ranges + .par_iter_mut() + .enumerate() + .map_init( + crate::direct_expm::ActionScratch::new, + |scratch, (block, range)| { + let first = block * rows; + let length = rows.min(dim - first); + range.cols.clear(); + range.vals.clear(); + range.offsets.clear(); + range.offsets.reserve_exact(length + 1); + range.offsets.push(0); + let mut per_col = if diagonal { + Vec::with_capacity(length) + } else { + Vec::new() + }; + let mut norm = 0.0_f64; + for row in first..first + length { + let value = direct.gather_row(row, scratch, true); + norm = norm.max(value.0); + if diagonal { + per_col.push(value); + } + let required = range.cols.len() + scratch.row.len(); + if required > range.cols.capacity() || required > range.vals.capacity() { + let capacity = required.next_power_of_two(); + range.cols.reserve_exact(capacity - range.cols.len()); + range.vals.reserve_exact(capacity - range.vals.len()); + } + for &(column, value) in &scratch.row { + range.cols.push(column); + range.vals.push(value); + } + range.offsets.push( + u32::try_from(range.cols.len()).expect("CSR range exceeds u32 entries"), + ); + } + (per_col, norm) + }, + ) + .collect(); + let mut per_col = Vec::new(); + let mut norm = 0.0_f64; + for (columns, bound) in metadata { + per_col.extend(columns); + norm = norm.max(bound); + } + let bounds = if diagonal { + real_bounds(&per_col) + } else { + (0.0, norm) + }; + ( + BlockCsr { + ranges, + range_shift: shift, + dim, + }, + bounds, + ) +} + +fn select_real_expm(t_norm: f64, drop_tol: f64) -> (u32, u32, f64) { + if drop_tol >= 1e-4 { + let (m, s) = expm::select_ms_loose(t_norm); + (m, s, 1e-6) + } else { + let (m, s) = expm::select_ms(t_norm); + (m, s, 1e-12) + } +} + +/// Apply the general real-coefficient exponential to the main coefficient vector. +pub(crate) fn expm_apply_inplace( + spec: &LindbladSpec, + basis: &[Word], + dt: f64, + coefficients: &mut [f64], + drop_tol: f64, + workspace: &mut crate::PcStepWorkspace, +) { + if basis.is_empty() { + return; + } + if let Some(plan) = &spec.pauli_plan { + let previous = match workspace.generator.take() { + Some(RealCache::Units(cache)) => Some(cache), + _ => None, + }; + let (cache, bounds) = crate::unit_expm::UnitCache::build(plan, basis, previous); + let (mu, bound) = bounds; + let (m, s, tol) = select_real_expm(dt.abs() * bound, drop_tol); + workspace.taylor.resize(2 * basis.len(), 0.0); + quspin_expm::ExpmOp::from_parts( + crate::unit_expm::UnitOp { + cache: &cache, + plan, + }, + dt, + mu, + s as usize, + m as usize, + tol, + ) + .apply_into( + ndarray::ArrayViewMut1::from(coefficients), + &mut workspace.taylor, + ) + .expect("expm apply"); + workspace.generator = Some(RealCache::Units(cache)); + return; + } + let reusable_rows = spec.supports_row_gather(); + let (mat, bounds) = if reusable_rows { + build_row_cache( + spec, + basis, + match workspace.generator.take() { + Some(RealCache::Coefficients(cache)) => Some(cache), + _ => None, + }, + ) + } else { + // General jump actions need column assembly; stale row buffers have no consumer in that build. + workspace.generator = None; + let index = CompactBasisIndex::new(basis); + let (mat, per_col) = build_mf_cols(spec, basis, &index); + (mat, real_bounds(&per_col)) + }; + let (mu, bound) = bounds; + let (m, s, tol) = select_real_expm(dt.abs() * bound, drop_tol); + workspace.taylor.resize(2 * basis.len(), 0.0); + quspin_expm::ExpmOp::from_parts(CsrOp { mat: &mat }, dt, mu, s as usize, m as usize, tol) + .apply_into( + ndarray::ArrayViewMut1::from(coefficients), + &mut workspace.taylor, + ) + .expect("expm apply"); + if reusable_rows { + workspace.generator = Some(RealCache::Coefficients(mat)); + } +} + +#[cfg(test)] +pub(crate) fn expm_apply_mf_reference( spec: &LindbladSpec, basis: &[Word], dt: f64, @@ -509,24 +712,10 @@ pub(crate) fn expm_apply_mf( if basis.is_empty() { return Vec::new(); } - let index = build_basis_index(basis); + let index = CompactBasisIndex::new(basis); let (mat, per_col) = build_mf_cols(spec, basis, &index); - - // Pick the Taylor backward-error tolerance to match the basis truncation: - // when the PC prunes coarsely (drop_tol >= 1e-4) a double-precision exp is - // ~10 orders more accurate than the state it acts on, so the relaxed - // (tol=1e-6, still >=100x tighter than the cut) table is used — it admits a - // lower-degree Taylor polynomial and cuts the SpMV count with no effect on - // the truncated result. At tight/zero drop_tol we keep double precision so - // the exact-reference paths (orbit-rep / merged) still agree bit-for-bit. - expm_apply_cached(&mat, &per_col, dt, coeffs, |t_norm| { - if drop_tol >= 1e-4 { - let (m, s) = expm::select_ms_loose(t_norm); - (m, s, 1e-6) - } else { - let (m, s) = expm::select_ms(t_norm); - (m, s, 1e-12) - } + expm_apply_cached(&mat, &per_col, dt, coeffs, |norm| { + select_real_expm(norm, drop_tol) }) } @@ -576,8 +765,8 @@ pub(crate) fn expm_apply_mf_cxvec( } let re: Vec = b.iter().map(|z| z.re).collect(); let im: Vec = b.iter().map(|z| z.im).collect(); - let re_out = expm_apply_mf(spec, basis, dt, &re, drop_tol); - let im_out = expm_apply_mf(spec, basis, dt, &im, drop_tol); + let re_out = expm_apply_mf_reference(spec, basis, dt, &re, drop_tol); + let im_out = expm_apply_mf_reference(spec, basis, dt, &im, drop_tol); re_out .into_iter() .zip(im_out) diff --git a/crates/ppvm-lindblad/src/pauli_plan.rs b/crates/ppvm-lindblad/src/pauli_plan.rs new file mode 100644 index 000000000..4afd88656 --- /dev/null +++ b/crates/ppvm-lindblad/src/pauli_plan.rs @@ -0,0 +1,246 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Shared parity tables and local phase rules for Pauli Hamiltonian actions. + +use crate::algebra::comm_product; +use crate::word::{CHUNK_BITS, Word, word_from_codes, word_hash, word_support}; +use fxhash::FxHashSet; + +struct Term { + word: Word, + sites: Option<[usize; 2]>, + signs: [i8; 16], +} + +pub(crate) struct UnitAction { + pub(crate) word: Word, + /// [`word_hash`] of `word`, computed once when the action is produced. + pub(crate) hash: u64, + pub(crate) value: f64, + pub(crate) unit: i32, +} + +pub(crate) struct PauliActionPlan { + terms: Vec>, + coefficients: Vec, + jumps: Vec, + words: usize, + tables: Vec, +} + +fn label(word: &Word, site: usize) -> usize { + let chunk = site / CHUNK_BITS; + let shift = site % CHUNK_BITS; + ((word.xbits.data[chunk] >> shift) & 1) as usize + | ((((word.zbits.data[chunk] >> shift) & 1) as usize) << 1) +} + +impl PauliActionPlan { + /// Number of Hamiltonian terms; their row-action codes are `±1..=term_count`. + pub(crate) fn term_count(&self) -> usize { + self.terms.len() + } + + pub(crate) fn has_diagonal_noise(&self) -> bool { + self.jumps.iter().any(|rate| *rate != 0.0) + } + + pub(crate) fn new(n: usize, h: &[(Word, f64)], jumps: &[(Word, f64)]) -> Option { + let mut seen = FxHashSet::default(); + if h.len() >= i32::MAX as usize || h.iter().any(|(word, _)| !seen.insert(*word)) { + return None; + } + let words = (h.len() + jumps.len()).div_ceil(64); + let mut tables = vec![0; n * 4 * words]; + for (index, (word, _)) in h.iter().chain(jumps).enumerate() { + for site in 0..n { + let generator = label(word, site); + for pauli in 1..4 { + if generator != 0 && generator != pauli { + tables[(site * 4 + pauli) * words + index / 64] |= 1u64 << (index % 64); + } + } + } + } + let terms = h + .iter() + .map(|&(word, _)| { + let mut support = Vec::new(); + word_support(&word, &mut support); + let sites = if (1..=2).contains(&support.len()) { + Some([support[0] as usize, *support.last().unwrap() as usize]) + } else { + None + }; + let mut signs = [0; 16]; + if let Some([a, b]) = sites { + for (pattern, sign) in signs.iter_mut().enumerate() { + let mut codes = vec![0; n]; + codes[a] = (pattern & 3) as u8; + if b != a { + codes[b] = (pattern >> 2) as u8; + } + *sign = comm_product(&word, &word_from_codes::(&codes).unwrap()).1 as i8; + } + } + Term { word, sites, signs } + }) + .collect(); + Some(Self { + terms, + coefficients: h.iter().map(|(_, coefficient)| *coefficient).collect(), + jumps: jumps.iter().map(|(_, rate)| *rate).collect(), + words, + tables, + }) + } + + #[cfg(test)] + pub(crate) fn fill( + &self, + word: &Word, + active: &mut Vec, + output: &mut Vec<(Word, f64)>, + ) { + output.clear(); + self.visit(word, active, |word, value, _| output.push((word, value))); + } + + /// Like [`Self::fill`], with each output word's [`word_hash`] first. + pub(crate) fn fill_keyed( + &self, + word: &Word, + active: &mut Vec, + output: &mut Vec<(u64, Word, f64)>, + ) { + output.clear(); + self.visit(word, active, |word, value, _| { + output.push((word_hash(&word), word, value)) + }); + } + + pub(crate) fn fill_units( + &self, + word: &Word, + active: &mut Vec, + output: &mut Vec>, + ) { + output.clear(); + self.visit(word, active, |word, value, unit| { + output.push(UnitAction { + hash: word_hash(&word), + word, + value, + unit, + }) + }); + } + + #[inline] + pub(crate) fn row_coefficient(&self, unit: i32) -> f64 { + let coefficient = self.coefficients[unit.unsigned_abs() as usize - 1]; + coefficient * if unit > 0 { 2.0 } else { -2.0 } + } + + // Hamiltonian codes identify the original model term and its row-action sign. + // A zero code denotes the separately accumulated Pauli-noise diagonal. + fn visit( + &self, + word: &Word, + active: &mut Vec, + mut emit: impl FnMut(Word, f64, i32), + ) { + active.clear(); + active.resize(self.words, 0); + for chunk in 0..C { + let mut support = word.xbits.data[chunk] | word.zbits.data[chunk]; + while support != 0 { + let site = chunk * CHUNK_BITS + support.trailing_zeros() as usize; + support &= support - 1; + let start = (site * 4 + label(word, site)) * self.words; + for (bits, table) in active + .iter_mut() + .zip(&self.tables[start..start + self.words]) + { + *bits ^= table; + } + } + } + let mut diagonal = 0.0; + for (block, &initial) in active.iter().enumerate() { + let mut bits = initial; + while bits != 0 { + let index = block * 64 + bits.trailing_zeros() as usize; + bits &= bits - 1; + if index >= self.terms.len() { + diagonal += -2.0 * self.jumps[index - self.terms.len()]; + continue; + } + let term = &self.terms[index]; + let (result, sign) = if let Some([a, b]) = term.sites { + let pattern = label(word, a) | if a == b { 0 } else { label(word, b) << 2 }; + let mut result = *word; + for chunk in 0..C { + result.xbits.data[chunk] ^= term.word.xbits.data[chunk]; + result.zbits.data[chunk] ^= term.word.zbits.data[chunk]; + } + (result, term.signs[pattern] as f64) + } else { + comm_product(&term.word, word) + }; + let value = self.coefficients[index] * sign; + if value != 0.0 { + let term = index as i32 + 1; + emit(result, value, if sign > 0.0 { -term } else { term }); + } + } + } + if diagonal != 0.0 { + emit(*word, diagonal, 0); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{JumpInput, LindbladSpec, parse_pauli_string}; + use num::Complex; + + #[test] + fn compiled_tables_match_every_local_pauli_and_diagonal_jump() { + let h = vec![ + ("XXII".into(), 0.3), + ("IYYI".into(), -0.7), + ("ZIIX".into(), 0.2), + ("IXYZ".into(), 0.4), + ]; + let jump = JumpInput { + lincomb: vec![("IZII".into(), Complex::new(1.0, 0.0))], + rate: 0.4, + }; + let spec = LindbladSpec::<2>::new(4, &h, &[jump]).unwrap(); + let parsed: Vec<_> = h + .iter() + .map(|(s, c)| (parse_pauli_string::<2>(s, 4).unwrap().0, *c)) + .collect(); + let jumps = [(parse_pauli_string::<2>("IZII", 4).unwrap().0, 0.4)]; + let plan = PauliActionPlan::new(4, &parsed, &jumps).unwrap(); + let mut active = Vec::new(); + let mut output = Vec::new(); + for code in 0..256 { + let word = word_from_codes::<2>( + &(0..4) + .map(|q| ((code >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap(); + plan.fill(&word, &mut active, &mut output); + let mut expected = spec.action(&word); + output.sort_unstable_by_key(|(w, _)| *w); + expected.sort_unstable_by_key(|(w, _)| *w); + assert_eq!(output, expected); + } + } +} diff --git a/crates/ppvm-lindblad/src/reference.rs b/crates/ppvm-lindblad/src/reference.rs new file mode 100644 index 000000000..8f7bb0484 --- /dev/null +++ b/crates/ppvm-lindblad/src/reference.rs @@ -0,0 +1,188 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Allocation-independent predictor/corrector reference for numerical regression tests. + +use crate::step::Phase; +use crate::truncate::{ + add_leakage_capped, cap_basis, cap_map_to_room, order_by_desc_mag, prune_basis, +}; +use crate::word::word_hash; +use crate::{Error, LindbladSpec, PcStepConfig, PcStepTimings, Word}; +use fxhash::{FxBuildHasher, FxHashMap}; +use num::Complex; +use rayon::prelude::*; +const CHUNK_SIZE: usize = 4096; + +impl LindbladSpec { + pub(crate) fn pc_step_reference( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + ) -> Result<(), Error> { + self.run_in_pool(cfg, |this| { + this.pc_step_reference_inner(basis, coeffs, dt, protected, cfg, false) + .map(|_| ()) + }) + } + fn pc_step_reference_inner( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + timed: bool, + ) -> Result { + let PcStepConfig { + max_basis, + admit_basis, + drop_tol, + tau_add, + .. + } = *cfg; + // Admission bound: enrichment may grow the live basis to `admit` + // >= `max_basis`; the final `cap_basis` then keeps the top- + // `max_basis` strings by evolved |coeff| over the whole union + // (retained + admitted) — rank displacement. With `admit_basis = + // None` admission is bounded by `max_basis` itself, `cap_basis` is + // a no-op, and membership turnover requires `drop_tol > 0`. + let admit = admit_basis.unwrap_or(max_basis).max(max_basis); + let tau_add = tau_add.unwrap_or(0.0); + let mut t = PcStepTimings::default(); + + // 1. First-hop expansion. After this, `coeffs` contains the pre-step + // coefficients followed by zeros for the newly-added leakage strings. + // We rely on `coeffs` itself as the pre-step buffer for the corrector + // — the reference exponential only borrows its input coefficients. + let p = Phase::start(timed); + let leak = self.leakage_reference_with_prune(basis, coeffs, protected, admit, tau_add)?; + p.stop(&mut t.leakage1_us); + + let p = Phase::start(timed); + add_leakage_capped(basis, coeffs, leak, admit); + p.stop(&mut t.expand1_us); + + // 2. Predictor: the reference exponential reads `coeffs` and returns a + // new owned vector with the predicted state. + let p = Phase::start(timed); + let coeffs_predict = + crate::mf_expm::expm_apply_mf_reference(self, basis, dt, coeffs, drop_tol); + p.stop(&mut t.expm1_us); + + // 3. Second-hop expansion from the predicted state. Extend `coeffs` + // with zeros for any newly-added second-hop strings so it remains a + // valid input (pre-step state) for the corrector. Once the basis is + // full, the first hop usually fills the whole admission room; with + // `room = 0` the second hop can admit nothing, so its leakage pass + // is skipped. + let n_predict = basis.len(); + if admit > n_predict { + let p = Phase::start(timed); + let leak2 = self.leakage_reference_with_prune( + basis, + &coeffs_predict, + protected, + admit, + tau_add, + )?; + p.stop(&mut t.leakage2_us); + + let p = Phase::start(timed); + add_leakage_capped(basis, coeffs, leak2, admit); + p.stop(&mut t.expand2_us); + } + + // 4. Corrector: redo from pre-step state on the doubly-enlarged basis. + // If the second hop admitted nothing, the corrector would repeat the + // predictor's computation exactly, so the predicted state is kept. + if basis.len() == n_predict { + *coeffs = coeffs_predict; + } else { + drop(coeffs_predict); + let p = Phase::start(timed); + *coeffs = crate::mf_expm::expm_apply_mf_reference(self, basis, dt, coeffs, drop_tol); + p.stop(&mut t.expm2_us); + } + + // 5. Prune basis entries below `drop_tol` (protected words never dropped). + prune_basis(basis, coeffs, drop_tol, protected); + cap_basis(basis, coeffs, max_basis, protected); + Ok(t) + } + + pub(crate) fn leakage_reference_with_prune( + &self, + basis: &[Word], + coeffs: &[f64], + protected: &[Word], + max_basis: usize, + tau_add: f64, + ) -> Result, f64)>, Error> { + if basis.len() != coeffs.len() { + return Err(Error::LengthMismatch { + what: "basis and coeffs", + a: basis.len(), + b: coeffs.len(), + }); + } + // Hash-only membership tables: storing 8-byte `u64` keys instead + // of 48-byte Words shrinks the in-basis structure ~6×, keeping it + // in L3 (and often L2) at basis sizes where the full-Word version + // would spill to DRAM. + let in_basis: FxHashMap = basis.iter().map(|w| (word_hash(w), ())).collect(); + let protected_set: FxHashMap = + protected.iter().map(|w| (word_hash(w), ())).collect(); + + let order = order_by_desc_mag(coeffs); + let room = max_basis.saturating_sub(basis.len()); + let n_qubits = self.n_qubits(); + let mut merged: FxHashMap, f64> = FxHashMap::default(); + for chunk_indices in order.chunks(CHUNK_SIZE) { + let local: Vec, f64)>> = chunk_indices + .par_iter() + .map_init( + || { + ( + Vec::::with_capacity(n_qubits), + Vec::::with_capacity(128), + FxHashMap::, Complex>::with_capacity_and_hasher( + 128, + FxBuildHasher::default(), + ), + ) + }, + |(s1, s2, lm), &i| { + let p = &basis[i]; + let c = coeffs[i]; + let terms = self.compute_action_terms(p, s1, s2, lm); + let mut out = Vec::with_capacity(terms.len()); + for (w, v) in terms.iter() { + let h = word_hash(w); + if !in_basis.contains_key(&h) && !protected_set.contains_key(&h) { + out.push((*w, c * *v)); + } + } + out + }, + ) + .collect(); + for v in local { + for (k, val) in v { + *merged.entry(k).or_insert(0.0) += val; + } + } + cap_map_to_room(&mut merged, room); + } + // Rate-based admission: keep only candidates whose leakage rate + // exceeds `tau_add`. `tau_add = 0` admits everything except exact + // zeros. + Ok(merged + .into_iter() + .filter(|(_, c)| c.abs() > tau_add) + .collect()) + } +} diff --git a/crates/ppvm-lindblad/src/sector.rs b/crates/ppvm-lindblad/src/sector.rs index 181d9359b..cb2dc7fc3 100644 --- a/crates/ppvm-lindblad/src/sector.rs +++ b/crates/ppvm-lindblad/src/sector.rs @@ -92,10 +92,15 @@ impl<'a> Sector<'a> { &self, q: &Word, ) -> Option<(Word, Complex, usize)> { - let (rep, idx, orbit_size) = self - .group - .canonicalize_in_sector_indexed(q, &self.characters)?; - Some((rep, self.characters.value(idx), orbit_size)) + let (rep, idx, orbit_size) = self.group.canonicalize_in_sector_indexed( + &q.to_pauli_word(self.group.n_qubits()), + &self.characters, + )?; + Some(( + Word::from_pauli_word(&rep), + self.characters.value(idx), + orbit_size, + )) } /// Number of **distinct** members of `w`'s translation orbit, or @@ -110,7 +115,10 @@ impl<'a> Sector<'a> { #[inline] pub fn orbit_size(&self, w: &Word) -> Option { self.group - .canonicalize_in_sector_indexed(w, &self.characters) + .canonicalize_in_sector_indexed( + &w.to_pauli_word(self.group.n_qubits()), + &self.characters, + ) .map(|(_, _, orbit_size)| orbit_size) } } @@ -124,6 +132,6 @@ impl<'a> Sector<'a> { /// same rep, both are kept (caller should run a merge afterwards). pub fn canonicalize_basis_to_rep(basis: &mut [Word], group: &TranslationGroup) { for w in basis.iter_mut() { - *w = group.canonicalize(w); + *w = Word::from_pauli_word(&group.canonicalize(&w.to_pauli_word(group.n_qubits()))); } } diff --git a/crates/ppvm-lindblad/src/shortlist.rs b/crates/ppvm-lindblad/src/shortlist.rs new file mode 100644 index 000000000..ea56c106f --- /dev/null +++ b/crates/ppvm-lindblad/src/shortlist.rs @@ -0,0 +1,799 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Bounded candidate shortlisting followed by complete signed leakage scoring. + +use crate::algebra::{anti_commutes, comm_product}; +use crate::direct_expm::ActionScratch; +use crate::word::word_hash; +use crate::{CandidateBuffer, CandidateProposer, Error, LindbladSpec, StreamingProposer, Word}; +use fxhash::FxHashSet; +use hashbrown::HashTable; +use rayon::prelude::*; + +fn source_order(order: &mut Vec, coefficients: &[f64]) { + order.clear(); + order.extend(0..coefficients.len()); + order.sort_unstable_by(|&a, &b| { + coefficients[b] + .abs() + .total_cmp(&coefficients[a].abs()) + .then_with(|| a.cmp(&b)) + }); +} + +#[derive(Default)] +struct Positions { + table: HashTable, +} + +fn positioned_word<'a, const C: usize>( + buffer: &'a CandidateBuffer<'_, C>, + position: u32, +) -> &'a Word { + let position = position as usize; + if position < buffer.basis().len() { + &buffer.basis()[position] + } else { + &buffer.candidates()[position - buffer.basis().len()] + } +} + +impl Positions { + fn reset(&mut self, buffer: &CandidateBuffer<'_, C>) { + self.table.clear(); + for (index, word) in buffer.basis().iter().enumerate() { + self.table.insert_unique( + word_hash(word), + u32::try_from(index).expect("basis exceeds u32 positions"), + |&index| word_hash(positioned_word(buffer, index)), + ); + } + } + + fn get(&self, buffer: &CandidateBuffer<'_, C>, word: &Word) -> Option { + self.table + .find(word_hash(word), |&index| { + positioned_word(buffer, index) == word + }) + .copied() + } + + fn insert_candidate( + &mut self, + buffer: &CandidateBuffer<'_, C>, + candidate: usize, + ) { + let position = u32::try_from(buffer.basis().len() + candidate) + .expect("candidate positions exceed u32"); + self.table.insert_unique( + word_hash(&buffer.candidates()[candidate]), + position, + |&index| word_hash(positioned_word(buffer, index)), + ); + } +} + +#[derive(Default)] +struct ScoreHeap { + limit: usize, + capacity: usize, + len: usize, +} + +impl ScoreHeap { + fn ensure(&mut self, numeric: &mut Vec, required: usize) { + if required <= self.capacity { + return; + } + let previous = self.capacity; + self.capacity = required + .max(previous.saturating_mul(2)) + .max(8) + .min(self.limit); + numeric.resize(2 * self.capacity, 0.0); + if previous != 0 { + numeric.copy_within(previous..2 * previous, self.capacity); + } + } + // The packed region stores heap slots in low 32 bits and inverse positions + // in high 32 bits. Floating-point arithmetic only reads the score prefix. + fn slot(&self, numeric: &[f64], rank: usize) -> usize { + numeric[self.capacity + rank].to_bits() as u32 as usize + } + fn rank(&self, numeric: &[f64], slot: usize) -> usize { + (numeric[self.capacity + slot].to_bits() >> 32) as usize + } + fn set_slot(&self, numeric: &mut [f64], rank: usize, slot: usize) { + let old = numeric[self.capacity + rank].to_bits(); + numeric[self.capacity + rank] = f64::from_bits( + (old & 0xffffffff00000000) | u32::try_from(slot).expect("heap slot exceeds u32") as u64, + ); + } + fn set_rank(&self, numeric: &mut [f64], slot: usize, rank: usize) { + let old = numeric[self.capacity + slot].to_bits(); + numeric[self.capacity + slot] = f64::from_bits( + (old & 0xffffffff) + | ((u32::try_from(rank).expect("heap rank exceeds u32") as u64) << 32), + ); + } + fn worse(numeric: &[f64], a: usize, b: usize) -> bool { + numeric[a].abs() < numeric[b].abs() || (numeric[a].abs() == numeric[b].abs() && a > b) + } + fn swap(&self, numeric: &mut [f64], a: usize, b: usize) { + let sa = self.slot(numeric, a); + let sb = self.slot(numeric, b); + self.set_slot(numeric, a, sb); + self.set_slot(numeric, b, sa); + self.set_rank(numeric, sa, b); + self.set_rank(numeric, sb, a); + } + fn insert(&mut self, numeric: &mut Vec, slot: usize, value: f64) { + self.ensure(numeric, slot + 1); + numeric[slot] = value; + let mut rank = self.len; + self.len += 1; + self.set_slot(numeric, rank, slot); + self.set_rank(numeric, slot, rank); + while rank > 0 { + let parent = (rank - 1) / 2; + if !Self::worse( + numeric, + self.slot(numeric, rank), + self.slot(numeric, parent), + ) { + break; + } + self.swap(numeric, rank, parent); + rank = parent; + } + } + fn update(&self, numeric: &mut [f64], slot: usize, value: f64) { + numeric[slot] = value; + let mut rank = self.rank(numeric, slot); + loop { + let left = 2 * rank + 1; + if left >= self.len { + break; + } + let right = left + 1; + let child = if right < self.len + && Self::worse(numeric, self.slot(numeric, right), self.slot(numeric, left)) + { + right + } else { + left + }; + if !Self::worse(numeric, self.slot(numeric, child), self.slot(numeric, rank)) { + break; + } + self.swap(numeric, rank, child); + rank = child; + } + } + fn minimum(&self, numeric: &[f64]) -> (usize, f64) { + let slot = self.slot(numeric, 0); + (slot, numeric[slot].abs()) + } +} + +#[derive(Default)] +struct Shortlist { + h: Vec<(Word, f64)>, + order: Vec, + heap: ScoreHeap, + positions: Positions, + protected: FxHashSet>, +} + +impl Shortlist { + fn prepare( + &mut self, + model: &LindbladSpec, + buffer: &CandidateBuffer<'_, C>, + protected: &[Word], + ) { + self.heap.capacity = 0; + self.heap.len = 0; + self.h.clear(); + self.h + .extend(model.hamiltonian_terms().filter(|(_, c)| *c != 0.0)); + self.h.sort_by(|a, b| b.1.abs().total_cmp(&a.1.abs())); + source_order(&mut self.order, buffer.coefficients()); + self.positions.reset(buffer); + self.protected.clear(); + self.protected.extend(protected.iter().copied()); + } + + fn build(&mut self, buffer: &mut CandidateBuffer<'_, C>, budget: usize) { + if budget == 0 || self.order.is_empty() { + return; + } + self.heap.len = 0; + self.heap.limit = budget; + for source_rank in 0..self.order.len() { + let source = self.order[source_rank]; + let coefficient = buffer.coefficients()[source]; + for &(h_word, h_coefficient) in &self.h { + let bound = 2.0 * h_coefficient.abs() * coefficient.abs(); + if self.heap.len == budget && bound <= self.heap.minimum(buffer.numeric).1 { + break; + } + if bound == 0.0 { + break; + } + if !anti_commutes(&h_word, &buffer.basis()[source]) { + continue; + } + let (word, sign) = comm_product(&h_word, &buffer.basis()[source]); + if self.protected.contains(&word) { + continue; + } + let rate = h_coefficient * sign * coefficient; + if let Some(position) = self.positions.get(buffer, &word) { + let position = position as usize; + if position < buffer.basis().len() { + continue; + } + let slot = position - buffer.basis().len(); + if rate.abs() > buffer.numeric[slot].abs() { + self.heap.update(buffer.numeric, slot, rate); + } + } else if self.heap.len < budget { + let slot = buffer.candidates().len(); + buffer.push(word); + self.heap.insert(buffer.numeric, slot, rate); + self.positions.insert_candidate(buffer, slot); + } else { + let slot = self.heap.minimum(buffer.numeric).0; + let position = u32::try_from(buffer.basis().len() + slot) + .expect("candidate position exceeds u32"); + self.positions + .table + .find_entry(word_hash(&buffer.candidates()[slot]), |&p| p == position) + .expect("heap destination absent from position table") + .remove(); + buffer.candidates_mut()[slot] = word; + self.positions.insert_candidate(buffer, slot); + self.heap.update(buffer.numeric, slot, rate); + } + } + if let Some((_, h_coefficient)) = self.h.first() + && self.heap.len == budget + && source_rank + 1 < self.order.len() + && 2.0 + * h_coefficient.abs() + * buffer.coefficients()[self.order[source_rank + 1]].abs() + <= self.heap.minimum(buffer.numeric).1 + { + break; + } + } + } +} + +fn compact_threshold(buffer: &mut CandidateBuffer<'_, C>, threshold: f64) -> usize { + let mut write = 0; + for read in 0..buffer.candidates().len() { + if buffer.numeric[read].abs() > threshold { + buffer + .words + .swap(buffer.prefix + write, buffer.prefix + read); + buffer.numeric[write] = buffer.numeric[read]; + write += 1; + } + } + write +} + +/// Shortlist external strings by individual contributions, then score signed leakage. +/// +/// Exact scores use reverse Hamiltonian actions and a basis-position lookup, +/// including all contributions and cancellations for each shortlisted word. +/// Selection is exact within the shortlist, not over all external words. +/// The default shortlist is 1.5 times the available admission slots. General +/// dissipators use [`StreamingProposer`] because their action need not be +/// skew-symmetric; Hamiltonians with Pauli noise use reverse-action scoring. +pub struct ShortlistScoreProposer { + extra_percent: usize, + room: usize, + shortlist: Shortlist, + pub(crate) streaming: StreamingProposer, + use_fallback: bool, +} +impl ShortlistScoreProposer { + pub(crate) fn allocation_bytes(&self) -> [usize; 3] { + [ + self.shortlist.h.capacity() * std::mem::size_of::<(Word, f64)>(), + self.shortlist.order.capacity() * std::mem::size_of::(), + self.shortlist.positions.table.allocation_size(), + ] + } + + /// Oversample the admission budget by `extra_percent` before exact scoring. + pub fn new(extra_percent: usize) -> Self { + Self { + extra_percent, + room: 0, + shortlist: Shortlist::default(), + streaming: StreamingProposer::default(), + use_fallback: false, + } + } +} +impl Default for ShortlistScoreProposer { + fn default() -> Self { + Self::new(50) + } +} + +fn select_largest( + buffer: &mut CandidateBuffer<'_, C>, + limit: usize, + threshold: f64, +) -> usize { + let n = buffer.candidates().len(); + if n <= limit { + return compact_threshold(buffer, threshold); + } + if limit == 0 { + return 0; + } + buffer.numeric.resize(2 * n, 0.0); + let (rates, keys) = buffer.numeric.split_at_mut(n); + for (key, rate) in keys.iter_mut().zip(rates.iter()) { + *key = rate.abs(); + } + let cutoff = *keys.select_nth_unstable_by(n - limit, f64::total_cmp).1; + let mut ties = limit - keys.iter().filter(|&&key| key > cutoff).count(); + buffer.numeric.truncate(n); + let mut write = 0; + for read in 0..n { + let magnitude = buffer.numeric[read].abs(); + let keep = magnitude > cutoff || (magnitude == cutoff && ties > 0); + if magnitude == cutoff && ties > 0 { + ties -= 1; + } + if keep && magnitude > threshold { + buffer + .words + .swap(buffer.prefix + write, buffer.prefix + read); + buffer.numeric[write] = buffer.numeric[read]; + write += 1; + } + } + write +} + +impl CandidateProposer for ShortlistScoreProposer { + fn propose( + &mut self, + model: &LindbladSpec, + buffer: &mut CandidateBuffer<'_, C>, + protected: &[Word], + maximum: usize, + ) -> Result<(), Error> { + self.room = maximum.saturating_sub(buffer.basis().len()); + self.use_fallback = !model.supports_row_gather(); + if self.use_fallback { + return self.streaming.propose(model, buffer, protected, maximum); + } + if self.room == 0 { + return Ok(()); + } + self.shortlist.prepare(model, buffer, protected); + let shortlist = self + .room + .saturating_add(self.room.saturating_mul(self.extra_percent) / 100); + self.shortlist.build(buffer, shortlist); + buffer.numeric.truncate(buffer.candidates().len()); + let (basis, words, coefficients, numeric) = buffer.scoring_parts(); + let n = basis.len(); + let positions = &self.shortlist.positions; + // The position table contains basis and shortlist entries. Only the + // original basis has nonzero coefficients in the leakage calculation. + let table = &positions.table; + numeric.par_iter_mut().zip(words.par_iter()).for_each_init( + ActionScratch::new, + |scratch, (rate, word)| { + scratch.fill(model, word, true); + scratch.row.clear(); + for &(hash, parent, value) in &scratch.terms { + if let Some(&index) = table.find(hash, |&index| { + (index as usize) < n && basis[index as usize] == parent + }) { + scratch.row.push((index, -value)); + } + } + scratch.row.sort_unstable_by_key(|entry| entry.0); + *rate = scratch + .row + .iter() + .map(|&(index, value)| coefficients[index as usize] * value) + .sum(); + }, + ); + // Selection reads neither table, and both are rebuilt by the next proposal. + // Releasing them keeps them out of the exponential's peak allocation. + self.shortlist.positions.table = HashTable::new(); + self.shortlist.order = Vec::new(); + Ok(()) + } + + fn select(&mut self, buffer: &mut CandidateBuffer<'_, C>, threshold: f64) -> usize { + if self.use_fallback { + self.streaming.select(buffer, threshold) + } else { + select_largest(buffer, self.room, threshold) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{JumpInput, parse_pauli_string, word_from_codes}; + use fxhash::FxHashMap; + use num::Complex; + + impl CandidateProposer for Shortlist { + fn propose( + &mut self, + model: &LindbladSpec, + buffer: &mut CandidateBuffer<'_, C>, + protected: &[Word], + maximum: usize, + ) -> Result<(), Error> { + let room = maximum.saturating_sub(buffer.basis().len()); + if room > 0 { + self.prepare(model, buffer, protected); + self.build(buffer, room); + } + Ok(()) + } + fn select(&mut self, buffer: &mut CandidateBuffer<'_, C>, threshold: f64) -> usize { + compact_threshold(buffer, threshold) + } + } + + fn selected>( + model: &LindbladSpec, + words: &[Word], + coefficients: &[f64], + protected: &[Word], + room: usize, + threshold: f64, + proposer: &mut P, + ) -> Vec> { + let mut words = words.to_vec(); + let mut numeric = Vec::new(); + let mut view = CandidateBuffer::new(&mut words, coefficients, &mut numeric); + proposer + .propose(model, &mut view, protected, coefficients.len() + room) + .unwrap(); + let count = proposer.select(&mut view, threshold); + let result = view.candidates()[..count].to_vec(); + assert!(count <= room); + assert_eq!( + result.iter().copied().collect::>().len(), + count + ); + assert!( + result + .iter() + .all(|word| !view.basis().contains(word) && !protected.contains(word)) + ); + result + } + + #[test] + fn greedy_matches_sorted_cartesian_product_without_merging() { + let h = vec![ + ("XII".into(), 0.31), + ("IYY".into(), -0.73), + ("ZXI".into(), 1.17), + ]; + let model = LindbladSpec::<2>::new(3, &h, &[]).unwrap(); + let words: Vec<_> = [1, 5, 9, 14, 31, 46] + .iter() + .map(|&code| { + word_from_codes( + &(0..3) + .map(|q| ((code >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let coefficients = [0.2, -0.9, 0.13, 0.7, -0.11, 0.4]; + let mut terms: Vec<_> = model.hamiltonian_terms().collect(); + terms.sort_by(|a, b| b.1.abs().total_cmp(&a.1.abs())); + let mut order = Vec::new(); + source_order(&mut order, &coefficients); + let mut pairs = Vec::new(); + for (source_rank, &source) in order.iter().enumerate() { + for (h_rank, &(term, coefficient)) in terms.iter().enumerate() { + let (word, sign) = comm_product(&term, &words[source]); + if sign != 0.0 && !words.contains(&word) { + pairs.push(( + word, + (coefficient * coefficients[source]).abs(), + source_rank, + h_rank, + )); + } + } + } + pairs.sort_by(|a, b| { + b.1.total_cmp(&a.1) + .then_with(|| a.2.cmp(&b.2)) + .then_with(|| a.3.cmp(&b.3)) + }); + let mut expected = Vec::new(); + for (word, _, _, _) in pairs { + if !expected.contains(&word) { + expected.push(word); + } + } + let mut greedy = Shortlist::default(); + for room in [0, 1, 3, 100] { + let result = selected(&model, &words, &coefficients, &[], room, 0.0, &mut greedy); + let mut actual = result.clone(); + actual.sort_unstable(); + let mut wanted = expected[..room.min(expected.len())].to_vec(); + wanted.sort_unstable(); + assert_eq!(actual, wanted); + } + } + + #[test] + fn bounded_greedy_matches_exhaustive_max_contributions_across_replacements_and_ties() { + let mut proposer = Shortlist::<2>::default(); + for trial in 0..20 { + let h: Vec<_> = (0..12) + .map(|i| { + let code = (i * 73 + trial * 11 + 1) % 256; + let word = (0..4) + .map(|q| ['I', 'X', 'Z', 'Y'][((code >> (2 * q)) & 3) as usize]) + .collect(); + ( + word, + if trial % 2 == 0 { + 1.0 + } else { + 0.2 + (i as f64 * 1.7).sin() + }, + ) + }) + .collect(); + let model = LindbladSpec::new(4, &h, &[]).unwrap(); + let words: Vec<_> = (0..32) + .map(|i| { + word_from_codes( + &(0..4) + .map(|q| (((i * 37 + trial * 7) >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let coefficients: Vec<_> = (0..32) + .map(|i| { + if trial % 2 == 0 { + if i % 3 == 0 { -1.0 } else { 1.0 } + } else { + (i as f64 * 1.13).cos() + } + }) + .collect(); + let mut maxima = FxHashMap::default(); + for (word, &coefficient) in words.iter().zip(&coefficients) { + for (h_word, h_coefficient) in model.hamiltonian_terms() { + let (output, sign) = comm_product(&h_word, word); + if sign != 0.0 && !words.contains(&output) { + let score = (sign * h_coefficient * coefficient).abs(); + let entry = maxima.entry(output).or_insert(0.0_f64); + *entry = entry.max(score); + } + } + } + let mut scores: Vec<_> = maxima.values().copied().collect(); + scores.sort_unstable_by(|a, b| b.total_cmp(a)); + for room in [1, 3, 17, 100] { + let result = selected(&model, &words, &coefficients, &[], room, 0.0, &mut proposer); + let mut actual: Vec<_> = result.iter().map(|word| maxima[word]).collect(); + actual.sort_unstable_by(|a, b| b.total_cmp(a)); + assert_eq!(actual.len(), room.min(scores.len())); + for (a, b) in actual.iter().zip(&scores) { + assert!((a - b).abs() < 1e-12); + } + } + } + } + + #[test] + fn exact_shortlist_scoring_detects_complete_cancellation() { + let model = + LindbladSpec::<2>::new(2, &[("XI".into(), 1.0), ("IX".into(), 1.0)], &[]).unwrap(); + let words: Vec<_> = ["ZY", "YZ"] + .iter() + .map(|s| parse_pauli_string(s, 2).unwrap().0) + .collect(); + assert_eq!( + selected( + &model, + &words, + &[1.0, -1.0], + &[], + 2, + 0.0, + &mut Shortlist::default() + ) + .len(), + 2 + ); + assert!( + selected( + &model, + &words, + &[1.0, -1.0], + &[], + 2, + 0.0, + &mut ShortlistScoreProposer::default() + ) + .is_empty() + ); + } + + #[test] + fn complete_shortlist_matches_full_signed_leakage_with_duplicate_h_terms() { + let model = LindbladSpec::<2>::new( + 4, + &[ + ("XXII".into(), 0.3), + ("IYYI".into(), -0.7), + ("ZIIX".into(), 0.2), + ("XXII".into(), -0.1), + ], + &[], + ) + .unwrap(); + let words: Vec<_> = (0..50) + .map(|i| { + word_from_codes( + &(0..4) + .map(|q| (((i * 37) >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let coefficients: Vec<_> = (0..50).map(|i| (i as f64 * 1.71).sin()).collect(); + let mut exact = FxHashMap::default(); + for (word, &coefficient) in words.iter().zip(&coefficients) { + for (output, value) in model.action(word) { + if !words.contains(&output) { + *exact.entry(output).or_insert(0.0_f64) += coefficient * value; + } + } + } + let mut values: Vec<_> = exact.values().map(|value| value * value).collect(); + values.sort_by(|a, b| b.total_cmp(a)); + for threads in [1, 2, 6] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let result = pool.install(|| { + selected( + &model, + &words, + &coefficients, + &[], + 5, + 0.0, + &mut ShortlistScoreProposer::new(100_000), + ) + }); + let power: f64 = result.iter().map(|word| exact[word].powi(2)).sum(); + assert!((power - values[..5].iter().sum::()).abs() < 1e-12); + } + } + + #[test] + fn shortlist_scoring_preserves_general_jump_leakage_by_fallback() { + let jump = JumpInput { + lincomb: vec![ + ("XI".into(), Complex::new(0.5, 0.0)), + ("YI".into(), Complex::new(0.0, 0.5)), + ], + rate: 0.2, + }; + let model = LindbladSpec::<2>::new(2, &[("XX".into(), 0.3)], &[jump]).unwrap(); + let words = vec![parse_pauli_string("ZI", 2).unwrap().0]; + let expected = selected( + &model, + &words, + &[1.0], + &[], + 8, + 0.0, + &mut StreamingProposer::default(), + ); + assert!(!expected.is_empty()); + assert_eq!( + expected, + selected( + &model, + &words, + &[1.0], + &[], + 8, + 0.0, + &mut ShortlistScoreProposer::default() + ) + ); + } + + #[test] + fn strategies_support_wide_words_and_protected_external_strings() { + let mut x = vec![0; 129]; + x[128] = 1; + x[0] = 1; + let mut z = vec![0; 129]; + z[128] = 2; + let h = x + .iter() + .map(|code| ['I', 'X', 'Z', 'Y'][*code as usize]) + .collect::(); + let model = LindbladSpec::<4>::new(129, &[(h, 0.3)], &[]).unwrap(); + let words = vec![word_from_codes(&z).unwrap()]; + let expected = selected( + &model, + &words, + &[1.0], + &[], + 8, + 0.0, + &mut StreamingProposer::default(), + ); + assert_eq!(expected.len(), 1); + assert_eq!( + expected, + selected( + &model, + &words, + &[1.0], + &[], + 8, + 0.0, + &mut Shortlist::default() + ) + ); + assert_eq!( + expected, + selected( + &model, + &words, + &[1.0], + &[], + 8, + 0.0, + &mut ShortlistScoreProposer::default() + ) + ); + assert!( + selected( + &model, + &words, + &[1.0], + &expected, + 8, + 0.0, + &mut Shortlist::default() + ) + .is_empty() + ); + } +} diff --git a/crates/ppvm-lindblad/src/spec.rs b/crates/ppvm-lindblad/src/spec.rs index 2a86d87ed..d092bef9f 100644 --- a/crates/ppvm-lindblad/src/spec.rs +++ b/crates/ppvm-lindblad/src/spec.rs @@ -53,6 +53,8 @@ fn candidate_terms(p_support: &[u32], index: &[Vec], scratch: &mut Vec /// costs more than the recompute (hash lookup ≳ recompute) and its several /// KB per cached word dominate memory at large basis sizes. pub struct LindbladSpec { + pub(crate) workspaces: std::sync::Mutex>>, + pub(crate) pauli_plan: Option>, n_qubits: usize, h_terms: Vec>, j_kinds: Vec>, @@ -145,7 +147,25 @@ impl LindbladSpec { }); } + let pauli_plan = if j_kinds + .iter() + .all(|j| matches!(j, JumpKind::HermitianPauli { .. })) + { + let h: Vec<_> = h_parsed.iter().map(|h| (h.word, h.coeff)).collect(); + let jumps: Vec<_> = j_kinds + .iter() + .map(|j| match j { + JumpKind::HermitianPauli { word, rate } => (*word, *rate), + _ => unreachable!(), + }) + .collect(); + crate::pauli_plan::PauliActionPlan::new(n_qubits, &h, &jumps) + } else { + None + }; Ok(Self { + workspaces: std::sync::Mutex::new(Vec::new()), + pauli_plan, n_qubits, h_terms: h_parsed, j_kinds, @@ -154,6 +174,17 @@ impl LindbladSpec { }) } + /// Whether the off-diagonal action is purely Hamiltonian and skew symmetric. + pub(crate) fn supports_row_gather(&self) -> bool { + self.j_kinds + .iter() + .all(|jump| matches!(jump, JumpKind::HermitianPauli { .. })) + } + + pub(crate) fn hamiltonian_terms(&self) -> impl Iterator, f64)> + '_ { + self.h_terms.iter().map(|term| (term.word, term.coeff)) + } + /// Add a Kossakowski-form dissipator /// `D*(O) = Σ_{n,m} K_nm ( A_n† O A_m − ½ {A_n† A_m, O} )`. /// @@ -166,6 +197,7 @@ impl LindbladSpec { ops: &[Vec<(String, Complex)>], k: &[Vec>], ) -> Result<(), Error> { + self.pauli_plan = None; for (pair, support) in kossakowski::compile(ops, k, self.n_qubits)? { let idx = self.j_kinds.len() as u32; self.j_kinds.push(JumpKind::Kossakowski(pair)); diff --git a/crates/ppvm-lindblad/src/step.rs b/crates/ppvm-lindblad/src/step.rs index debd9666f..739204295 100644 --- a/crates/ppvm-lindblad/src/step.rs +++ b/crates/ppvm-lindblad/src/step.rs @@ -5,9 +5,12 @@ use crate::sector::Sector; use crate::spec::LindbladSpec; -use crate::truncate::{add_leakage_capped, cap_basis, prune_basis}; +use crate::truncate::{add_leakage_capped, cap_basis, cap_basis_with_scratch, prune_basis}; use crate::word::Word; -use crate::{Error, PcStepConfig, mf_expm}; +use crate::{ + CandidateBuffer, CandidateProposer, CandidateStrategy, Error, PcStepConfig, PcStepWorkspace, + mf_expm, +}; use num::Complex; use std::time::Instant; @@ -36,14 +39,14 @@ impl PcStepTimings { /// Clock for one `pc_step` phase. Disarmed (`None`) on the untimed path, so /// [`LindbladSpec::pc_step`] pays no `Instant` syscalls. -struct Phase(Option); +pub(crate) struct Phase(Option); impl Phase { - fn start(timed: bool) -> Self { + pub(crate) fn start(timed: bool) -> Self { Self(timed.then(Instant::now)) } - fn stop(self, slot: &mut u64) { + pub(crate) fn stop(self, slot: &mut u64) { if let Some(t0) = self.0 { *slot = t0.elapsed().as_micros() as u64; } @@ -75,10 +78,8 @@ impl LindbladSpec { protected: &[Word], cfg: &PcStepConfig, ) -> Result<(), Error> { - self.run_in_pool(cfg, |this| { - this.pc_step_inner(basis, coeffs, dt, protected, cfg, false) - .map(|_| ()) - }) + self.pc_step_pooled(basis, coeffs, dt, protected, cfg, false) + .map(|_| ()) } /// Same as [`Self::pc_step`] but also returns a per-phase timing @@ -92,12 +93,162 @@ impl LindbladSpec { protected: &[Word], cfg: &PcStepConfig, ) -> Result { + self.pc_step_pooled(basis, coeffs, dt, protected, cfg, true) + } + + fn pc_step_pooled( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + timed: bool, + ) -> Result { + // Only checkout and return hold the lock; concurrent trajectories use separate buffers. + let mut workspace = self + .workspaces + .lock() + .unwrap_or_else(|error| error.into_inner()) + .pop() + .unwrap_or_default(); + let result = self.pc_step_execute(basis, coeffs, dt, protected, cfg, timed, &mut workspace); + self.workspaces + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push(workspace); + result + } + + /// Adaptive step using buffers owned explicitly by the caller. + /// + /// Ordinary step calls reuse an internal workspace automatically. This entry point + /// lets a caller retain the buffers independently of the model. + /// Capacity follows the active state; a rank limit does not allocate its full size. + /// Numerical ordering is identical to [`Self::pc_step_timed`]. + pub fn pc_step_with_workspace( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + workspace: &mut PcStepWorkspace, + ) -> Result { + self.pc_step_execute(basis, coeffs, dt, protected, cfg, true, workspace) + } + + #[allow(clippy::too_many_arguments)] + fn pc_step_execute( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + timed: bool, + workspace: &mut PcStepWorkspace, + ) -> Result { + let mut proposer = std::mem::take(&mut workspace.proposer); + let result = match cfg.proposer { + CandidateStrategy::Streaming => self.pc_step_execute_with_proposer( + basis, + coeffs, + dt, + protected, + cfg, + timed, + workspace, + &mut proposer.streaming, + ), + CandidateStrategy::Shortlist => self.pc_step_execute_with_proposer( + basis, + coeffs, + dt, + protected, + cfg, + timed, + workspace, + &mut proposer, + ), + }; + workspace.proposer = proposer; + result + } + + /// Adaptive step with a caller-owned candidate strategy and reusable workspace. + /// + /// The strategy receives a borrowed view of the active sum and its spare word + /// storage. The stepper initializes admitted coefficients, evolves the state, + /// and applies final truncation. Ordinary step calls use the strategy in [`PcStepConfig::proposer`]. + #[allow(clippy::too_many_arguments)] + pub fn pc_step_with_proposer>( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + workspace: &mut PcStepWorkspace, + proposer: &mut P, + ) -> Result { + self.pc_step_execute_with_proposer( + basis, coeffs, dt, protected, cfg, true, workspace, proposer, + ) + } + + #[allow(clippy::too_many_arguments)] + fn pc_step_execute_with_proposer>( + &self, + basis: &mut Vec>, + coeffs: &mut Vec, + dt: f64, + protected: &[Word], + cfg: &PcStepConfig, + timed: bool, + workspace: &mut PcStepWorkspace, + proposer: &mut P, + ) -> Result { + if basis.len() != coeffs.len() { + return Err(Error::LengthMismatch { + what: "basis and coeffs", + a: basis.len(), + b: coeffs.len(), + }); + } + let dimension = 1usize + .checked_shl((2 * self.n_qubits()) as u32) + .unwrap_or(usize::MAX); + let admit = cfg.admit_basis.unwrap_or(cfg.max_basis).max(cfg.max_basis); + { + // Reserve from the active state rather than an arbitrarily large rank limit. + let active = basis.len().min(cfg.max_basis); + let bound = active.saturating_mul(3).min(admit).min(dimension); + if let Some(required) = bound.checked_mul(2) + && workspace.taylor.capacity() < required + { + workspace + .taylor + .reserve_exact(required - workspace.taylor.len()); + } + if let Some(reserve) = active.checked_mul(4) { + let reserve = reserve.max(bound).min(dimension); + if basis.capacity() < reserve { + basis.reserve_exact(reserve.saturating_sub(basis.len())); + } + if coeffs.capacity() < reserve { + coeffs.reserve_exact(reserve.saturating_sub(coeffs.len())); + } + } + } self.run_in_pool(cfg, |this| { - this.pc_step_inner(basis, coeffs, dt, protected, cfg, true) + this.pc_step_inner( + basis, coeffs, dt, protected, cfg, timed, workspace, proposer, + ) }) } - fn run_in_pool( + pub(crate) fn run_in_pool( &self, cfg: &PcStepConfig, f: impl FnOnce(&Self) -> Result + Send, @@ -113,7 +264,8 @@ impl LindbladSpec { } } - fn pc_step_inner( + #[allow(clippy::too_many_arguments)] + fn pc_step_inner>( &self, basis: &mut Vec>, coeffs: &mut Vec, @@ -121,6 +273,8 @@ impl LindbladSpec { protected: &[Word], cfg: &PcStepConfig, timed: bool, + workspace: &mut PcStepWorkspace, + proposer: &mut P, ) -> Result { let PcStepConfig { max_basis, @@ -139,66 +293,65 @@ impl LindbladSpec { let tau_add = tau_add.unwrap_or(0.0); let mut t = PcStepTimings::default(); - // 1. First-hop expansion. After this, `coeffs` contains the pre-step - // coefficients followed by zeros for the newly-added leakage strings. - // We rely on `coeffs` itself as the pre-step buffer for the corrector - // — no `.clone()` is needed because `expm_step` only borrows it. + let original_terms = basis.len(); + + // No basis permutation occurs before final truncation, so a saved + // coefficient prefix retains the original word positions for correction. + // 1. First-hop expansion appends zero coefficients for admitted words. let p = Phase::start(timed); - let leak = self.leakage_with_prune(basis, coeffs, protected, admit, tau_add)?; + let mut buffer = CandidateBuffer::new(basis, coeffs, &mut workspace.taylor); + proposer.propose(self, &mut buffer, protected, admit)?; p.stop(&mut t.leakage1_us); let p = Phase::start(timed); - add_leakage_capped(basis, coeffs, leak, admit); + let count = proposer.select(&mut buffer, tau_add); + assert!(count <= admit.saturating_sub(buffer.basis().len())); + buffer.finish(count); + coeffs.resize(basis.len(), 0.0); p.stop(&mut t.expand1_us); - // 2. Predictor: `expm_step` reads `coeffs` immutably and returns a - // new owned vector with the predicted state. + if basis.len() < admit { + workspace.coefficient_auxiliary.clear(); + workspace + .coefficient_auxiliary + .extend_from_slice(&coeffs[..original_terms]); + } + + // The main coefficient vector is the exponential's output; the prefix backup is only for a possible corrector. let p = Phase::start(timed); - let coeffs_predict = self.expm_step(basis, dt, coeffs, drop_tol); + mf_expm::expm_apply_inplace(self, basis, dt, coeffs, drop_tol, workspace); p.stop(&mut t.expm1_us); - // 3. Second-hop expansion from the predicted state. Extend `coeffs` - // with zeros for any newly-added second-hop strings so it remains a - // valid input (pre-step state) for the corrector. Once the basis is - // full, the first hop usually fills the whole admission room; with - // `room = 0` the second hop can admit nothing, so its leakage pass - // is skipped. let n_predict = basis.len(); if admit > n_predict { let p = Phase::start(timed); - let leak2 = - self.leakage_with_prune(basis, &coeffs_predict, protected, admit, tau_add)?; + let mut buffer = CandidateBuffer::new(basis, coeffs, &mut workspace.taylor); + proposer.propose(self, &mut buffer, protected, admit)?; p.stop(&mut t.leakage2_us); - let p = Phase::start(timed); - add_leakage_capped(basis, coeffs, leak2, admit); + let count = proposer.select(&mut buffer, tau_add); + assert!(count <= admit.saturating_sub(buffer.basis().len())); + buffer.finish(count); + coeffs.resize(basis.len(), 0.0); p.stop(&mut t.expand2_us); } - // 4. Corrector: redo from pre-step state on the doubly-enlarged basis. - // If the second hop admitted nothing, the corrector would repeat the - // predictor's computation exactly, so the predicted state is kept. - if basis.len() == n_predict { - *coeffs = coeffs_predict; - } else { - drop(coeffs_predict); + if basis.len() != n_predict { let p = Phase::start(timed); - *coeffs = self.expm_step(basis, dt, coeffs, drop_tol); + coeffs.clear(); + coeffs.extend_from_slice(&workspace.coefficient_auxiliary); + coeffs.resize(basis.len(), 0.0); + mf_expm::expm_apply_inplace(self, basis, dt, coeffs, drop_tol, workspace); p.stop(&mut t.expm2_us); } // 5. Prune basis entries below `drop_tol` (protected words never dropped). prune_basis(basis, coeffs, drop_tol, protected); - cap_basis(basis, coeffs, max_basis, protected); + // The exponential's work vector is idle until the next step's proposal. + cap_basis_with_scratch(basis, coeffs, max_basis, protected, &mut workspace.taylor); Ok(t) } - /// Compute `exp(dt · M) · b` for the in-basis-restricted generator - /// `M`, matrix-free, via `quspin-expm` (see [`crate::mf_expm`]). - fn expm_step(&self, basis: &[Word], dt: f64, b: &[f64], drop_tol: f64) -> Vec { - mf_expm::expm_apply_mf(self, basis, dt, b, drop_tol) - } - /// One predictor-corrector step in **orbit-rep form** at momentum /// `sector`: the same five phases as [`Self::pc_step`], but the basis /// holds only canonical translation-orbit representatives, the @@ -297,3 +450,121 @@ impl LindbladSpec { Ok(()) } } + +#[cfg(test)] +mod default_tests { + use super::*; + use crate::{JumpInput, word_from_codes}; + + #[test] + fn default_step_matches_original_with_general_jumps_and_kossakowski() { + let h = vec![ + ("XXI".into(), 0.3), + ("IYY".into(), -0.17), + ("XXI".into(), 0.13), + ]; + let jumps = vec![JumpInput { + lincomb: vec![ + ("XII".into(), Complex::new(0.5, 0.0)), + ("YII".into(), Complex::new(0.0, 0.5)), + ], + rate: 0.4, + }]; + let mut reference = LindbladSpec::<2>::new(3, &h, &jumps).unwrap(); + let mut optimized = LindbladSpec::<2>::new(3, &h, &jumps).unwrap(); + let operators = vec![ + vec![("IXI".into(), Complex::new(1.0, 0.0))], + vec![("IYI".into(), Complex::new(1.0, 0.0))], + ]; + let k = vec![ + vec![Complex::new(0.1, 0.0), Complex::new(0.0, 0.02)], + vec![Complex::new(0.0, -0.02), Complex::new(0.2, 0.0)], + ]; + reference.add_kossakowski(&operators, &k).unwrap(); + optimized.add_kossakowski(&operators, &k).unwrap(); + let seed = word_from_codes::<2>(&[2, 0, 0]).unwrap(); + let (mut a, mut b) = (vec![seed], vec![seed]); + let (mut ca, mut cb) = (vec![1.0], vec![1.0]); + for step in 0..12 { + let cfg = PcStepConfig { + max_basis: 8 + step % 3, + admit_basis: Some(24 + step % 4), + drop_tol: 1e-7, + num_threads: Some(2), + ..Default::default() + }; + reference + .pc_step_reference(&mut a, &mut ca, 0.1, &[seed], &cfg) + .unwrap(); + optimized + .pc_step(&mut b, &mut cb, 0.1, &[seed], &cfg) + .unwrap(); + assert_eq!(a, b); + assert_eq!(ca, cb); + } + let pool = optimized.workspaces.lock().unwrap(); + assert_eq!(pool.len(), 1); + assert!(pool[0].taylor.capacity() > 0); + assert!(pool[0].generator.is_none()); + } + + #[test] + fn concurrent_default_trajectories_share_model_without_sharing_scratch() { + let spec = + LindbladSpec::<2>::new(3, &[("XXI".into(), 0.3), ("IYY".into(), -0.17)], &[]).unwrap(); + let barrier = std::sync::Barrier::new(4); + std::thread::scope(|scope| { + for trajectory in 0..4 { + let spec = &spec; + let barrier = &barrier; + scope.spawn(move || { + let seed = word_from_codes::<2>(&[trajectory % 3 + 1, 0, 0]).unwrap(); + let cfg = PcStepConfig { + max_basis: 8 + trajectory as usize, + admit_basis: Some(24), + ..Default::default() + }; + let (mut a, mut b) = (vec![seed], vec![seed]); + let (mut ca, mut cb) = (vec![1.0], vec![1.0]); + barrier.wait(); + for _ in 0..6 { + spec.pc_step_reference(&mut a, &mut ca, 0.1, &[seed], &cfg) + .unwrap(); + spec.pc_step(&mut b, &mut cb, 0.1, &[seed], &cfg).unwrap(); + assert_eq!(a, b); + assert_eq!(ca, cb); + } + }); + } + }); + let pool = spec.workspaces.lock().unwrap(); + assert!(!pool.is_empty()); + assert!(pool.len() <= 4); + } + #[test] + fn large_rank_limit_does_not_eagerly_allocate_for_a_small_state() { + let n = 129; + let mut z = vec!['I'; n]; + z[0] = 'Z'; + let spec = LindbladSpec::<4>::new(n, &[(z.into_iter().collect(), 0.37)], &[]).unwrap(); + let mut codes = vec![0; n]; + codes[0] = 2; + let seed = word_from_codes::<4>(&codes).unwrap(); + let mut basis = vec![seed]; + let mut coeffs = vec![1.0]; + let cfg = PcStepConfig { + max_basis: 1 << 30, + admit_basis: Some(3 << 30), + ..Default::default() + }; + spec.pc_step(&mut basis, &mut coeffs, 0.1, &[seed], &cfg) + .unwrap(); + assert_eq!(basis, vec![seed]); + assert_eq!(coeffs, vec![1.0]); + assert!(basis.capacity() <= 8); + assert!(coeffs.capacity() <= 8); + let pool = spec.workspaces.lock().unwrap(); + assert!(pool[0].taylor.capacity() <= 16); + assert!(pool[0].coefficient_auxiliary.capacity() <= 8); + } +} diff --git a/crates/ppvm-lindblad/src/tests.rs b/crates/ppvm-lindblad/src/tests.rs index 5c5666ba4..ac7c76037 100644 --- a/crates/ppvm-lindblad/src/tests.rs +++ b/crates/ppvm-lindblad/src/tests.rs @@ -5,6 +5,34 @@ use super::*; use fxhash::FxHashMap; use num::Complex; +/// `ppvm_pauli_sum` symmetry merges on the crate's compact words, via the general word type. +fn canonicalize_pauli_sum_complex( + basis: &mut Vec>, + coeffs: &mut Vec>, + group: &ppvm_pauli_sum::symmetry::TranslationGroup, + k: &[i32], +) { + let mut words: Vec<_> = basis + .iter() + .map(|w| w.to_pauli_word(group.n_qubits())) + .collect(); + ppvm_pauli_sum::symmetry::canonicalize_pauli_sum_complex(&mut words, coeffs, group, k); + *basis = words.iter().map(Word::from_pauli_word).collect(); +} + +fn canonicalize_pauli_sum( + basis: &mut Vec>, + coeffs: &mut Vec, + group: &ppvm_pauli_sum::symmetry::TranslationGroup, +) { + let mut words: Vec<_> = basis + .iter() + .map(|w| w.to_pauli_word(group.n_qubits())) + .collect(); + ppvm_pauli_sum::symmetry::canonicalize_pauli_sum(&mut words, coeffs, group); + *basis = words.iter().map(Word::from_pauli_word).collect(); +} + /// Test-only full-space complex predictor-corrector step, UNTRUNCATED: /// adds every nonzero leakage string (two hops) and applies the exact /// in-basis exponential to the complex coefficient vector. Reference @@ -125,8 +153,6 @@ fn assert_orbit_rep_matches_projection( dt: f64, n_steps: usize, ) { - use ppvm_pauli_sum::symmetry::canonicalize_pauli_sum_complex; - let spec = ::new(n, h_terms, &[]).unwrap(); let group = ppvm_pauli_sum::symmetry::TranslationGroup::chain_1d(n); let basis_full: Vec = seed @@ -178,7 +204,7 @@ fn assert_orbit_rep_matches_projection( let cf_val = mf .get(w) .copied() - .unwrap_or_else(|| panic!("rep {w} in orbit-rep but not in full-basis")); + .unwrap_or_else(|| panic!("rep {w:?} in orbit-rep but not in full-basis")); max_diff = max_diff.max((cm - cf_val).norm()); } assert!( @@ -326,7 +352,7 @@ fn complex_full_matches_real_at_kzero() { /// bit-identical up to FP noise. #[test] fn pc_step_matches_symmetry_merged_on_small_chain() { - use ppvm_pauli_sum::symmetry::{TranslationGroup, canonicalize_pauli_sum}; + use ppvm_pauli_sum::symmetry::TranslationGroup; let n = 4usize; let dt = 0.05f64; @@ -482,7 +508,7 @@ fn pc_step_is_bit_identical_across_thread_counts() { /// duplicate columns. The cap binds. #[test] fn pc_step_orbit_rep_is_bit_identical_across_thread_counts() { - use ppvm_pauli_sum::symmetry::{TranslationGroup, canonicalize_pauli_sum_complex}; + use ppvm_pauli_sum::symmetry::TranslationGroup; let n = 12usize; let lincomb: Vec<(String, Complex)> = (0..n) @@ -579,3 +605,30 @@ fn action_term_order_is_independent_of_scratch_capacity() { ); } } + +#[test] +fn compact_word_is_two_planes_and_hashes_and_orders_like_pauli_word() { + use std::hash::{BuildHasher, Hash, Hasher}; + assert_eq!(std::mem::size_of::>(), 32); + let fx = |h: &dyn Fn(&mut fxhash::FxHasher)| { + let mut state = fxhash::FxBuildHasher::default().build_hasher(); + h(&mut state); + state.finish() + }; + let words: Vec> = (0..400u64) + .map(|i| { + let codes: Vec = (0..100) + .map(|q| ((i.wrapping_mul(0x9e37_79b9_7f4a_7c15) >> (q % 61)) & 3) as u8) + .collect(); + word_from_codes::<2>(&codes).unwrap() + }) + .collect(); + for a in &words { + let pa = a.to_pauli_word(100); + assert_eq!(fx(&|s| a.hash(s)), fx(&|s| pa.hash(s))); + assert_eq!(Word::from_pauli_word(&pa), *a); + for b in words.iter().take(40) { + assert_eq!(a.cmp(b), pa.cmp(&b.to_pauli_word(100))); + } + } +} diff --git a/crates/ppvm-lindblad/src/truncate.rs b/crates/ppvm-lindblad/src/truncate.rs index 308e5e0a5..ae9f91a3a 100644 --- a/crates/ppvm-lindblad/src/truncate.rs +++ b/crates/ppvm-lindblad/src/truncate.rs @@ -68,6 +68,17 @@ pub(crate) fn cap_basis( coeffs: &mut Vec, max_basis: usize, protected: &[Word], +) { + cap_basis_with_scratch(basis, coeffs, max_basis, protected, &mut Vec::new()); +} + +/// [`cap_basis`] computing the magnitudes in `mags`, so an idle work buffer can hold them. +pub(crate) fn cap_basis_with_scratch( + basis: &mut Vec>, + coeffs: &mut Vec, + max_basis: usize, + protected: &[Word], + mags: &mut Vec, ) { if basis.len() <= max_basis { return; @@ -75,16 +86,18 @@ pub(crate) fn cap_basis( let protected_set: FxHashSet<&Word> = protected.iter().collect(); let n_prot = basis.iter().filter(|w| protected_set.contains(w)).count(); let slots = max_basis.saturating_sub(n_prot); - let mut mags: Vec = basis - .iter() - .zip(coeffs.iter()) - .filter(|(w, _)| !protected_set.contains(w)) - .map(|(_, c)| c.mag()) - .collect(); + mags.clear(); + mags.extend( + basis + .iter() + .zip(coeffs.iter()) + .filter(|(w, _)| !protected_set.contains(w)) + .map(|(_, c)| c.mag()), + ); if slots >= mags.len() { return; } - let mut keep = top_k_keeper(&mut mags, slots); + let mut keep = top_k_keeper(mags, slots); retain_in_place(basis, coeffs, |w, c| { protected_set.contains(w) || keep(c.mag()) }); diff --git a/crates/ppvm-lindblad/src/unit_expm.rs b/crates/ppvm-lindblad/src/unit_expm.rs new file mode 100644 index 000000000..21ebeb7bc --- /dev/null +++ b/crates/ppvm-lindblad/src/unit_expm.rs @@ -0,0 +1,686 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Coefficient-free Hamiltonian connections referencing the original model terms. +//! +//! A unit plan exists only when every dissipator is a Hermitian Pauli jump, so all +//! off-diagonal generator entries come from the Hamiltonian, with any noise on the +//! diagonal. In the orthonormal Pauli basis a Hamiltonian commutator is real +//! antisymmetric, term by term: row `i` reaches parent `j` through term `k` exactly when +//! row `j` reaches `i` through `k` with the opposite sign. Projection onto the working +//! basis keeps this, so the cache stores each connected pair once and the product +//! applies it in both directions. + +use crate::compact_index::CompactBasisIndex; +use crate::pauli_plan::{PauliActionPlan, UnitAction}; +use crate::word::Word; +use quspin_types::{LinearOperator, QuSpinError}; +use rayon::prelude::*; + +/// Signed Hamiltonian term code stored with each pair: 16 bits when the model has at +/// most `i16::MAX` terms, otherwise 32 bits. +trait Code: Copy + Default + Send + Sync { + fn pack(unit: i32) -> Self; + fn unit(self) -> i32; +} + +impl Code for i16 { + fn pack(unit: i32) -> Self { + unit as i16 + } + fn unit(self) -> i32 { + self as i32 + } +} + +impl Code for i32 { + fn pack(unit: i32) -> Self { + unit + } + fn unit(self) -> i32 { + self + } +} + +/// One stored pair: block-local row and parent positions and the signed term code +/// for the row's direction. The parent's block is given by the tile holding the entry. +#[derive(Clone, Copy, Default)] +struct Entry { + row: u16, + parent: u16, + unit: U, +} + +const _: () = assert!(std::mem::size_of::>() == 6); +const _: () = assert!(std::mem::size_of::>() == 8); + +/// Largest block-size exponent that keeps block-local positions in 16 bits. +const MAX_SHIFT: u32 = 16; + +/// Rows of one block and the pairs it owns, grouped by the other row's block. +struct Block { + /// `entries[tiles[b]..tiles[b + 1]]` pair rows of this block with rows of block `b`. + tiles: Vec, + entries: Vec>, + diagonal: Vec, +} + +impl Default for Block { + fn default() -> Self { + Self { + tiles: Vec::new(), + entries: Vec::new(), + diagonal: Vec::new(), + } + } +} + +impl Block { + fn allocation_bytes(&self) -> usize { + self.tiles.capacity() * 4 + + self.entries.capacity() * std::mem::size_of::>() + + self.diagonal.capacity() * 8 + } +} + +enum Blocks { + Narrow(Vec>), + Wide(Vec>), +} + +pub(crate) struct UnitCache { + blocks: Blocks, + shift: u32, + dimension: usize, +} + +/// The block storing pairs between blocks `a` and `b`. Alternating the owner by parity +/// gives every block about half of its pairs, which balances both product passes. +fn owner(a: usize, b: usize) -> usize { + if (a + b).is_multiple_of(2) { + a.max(b) + } else { + a.min(b) + } +} + +/// Running reduction for the exponential's shift and norm bound without per-row storage. +/// +/// [`crate::mf_expm::real_bounds`] takes `mu` as the mean diagonal and the bound as the +/// largest `raw - |d| + |d - mu|` over rows. Since `|d - mu| = max(d - mu, mu - d)`, that +/// maximum equals `max(max(raw - |d| + d) - mu, max(raw - |d| - d) + mu)`, which needs only +/// two running maxima and the diagonal sum. +#[derive(Clone, Copy)] +struct Bounds { + diagonal_sum: f64, + above: f64, + below: f64, + raw: f64, +} + +impl Default for Bounds { + fn default() -> Self { + Self { + diagonal_sum: 0.0, + above: f64::NEG_INFINITY, + below: f64::NEG_INFINITY, + raw: 0.0, + } + } +} + +impl Bounds { + fn add(&mut self, raw: f64, diagonal: f64) { + let off = raw - diagonal.abs(); + self.diagonal_sum += diagonal; + self.above = self.above.max(off + diagonal); + self.below = self.below.max(off - diagonal); + self.raw = self.raw.max(raw); + } + + fn merge(self, other: Self) -> Self { + Self { + diagonal_sum: self.diagonal_sum + other.diagonal_sum, + above: self.above.max(other.above), + below: self.below.max(other.below), + raw: self.raw.max(other.raw), + } + } + + fn finish(self, dimension: usize, noisy: bool) -> (f64, f64) { + if !noisy { + return (0.0, self.raw); + } + let mu = self.diagonal_sum / dimension as f64; + (mu, (self.above - mu).max(self.below + mu).max(0.0)) + } +} + +impl UnitCache { + pub(crate) fn allocation_bytes(&self) -> usize { + match &self.blocks { + Blocks::Narrow(b) => { + b.capacity() * std::mem::size_of::>() + + b.iter().map(Block::allocation_bytes).sum::() + } + Blocks::Wide(b) => { + b.capacity() * std::mem::size_of::>() + + b.iter().map(Block::allocation_bytes).sum::() + } + } + } + + #[cfg(test)] + fn block_count(&self) -> usize { + match &self.blocks { + Blocks::Narrow(b) => b.len(), + Blocks::Wide(b) => b.len(), + } + } + + pub(crate) fn build( + plan: &PauliActionPlan, + basis: &[Word], + previous: Option, + ) -> (Self, (f64, f64)) { + let dimension = basis.len(); + let shift = dimension + .div_ceil(256) + .max(64) + .next_power_of_two() + .trailing_zeros() + .min(MAX_SHIFT); + let index = CompactBasisIndex::new(basis); + let previous = previous.map(|cache| cache.blocks); + let (blocks, bounds) = if plan.term_count() <= i16::MAX as usize { + let blocks = match previous { + Some(Blocks::Narrow(blocks)) => blocks, + _ => Vec::new(), + }; + let (blocks, bounds) = build_blocks(plan, basis, &index, shift, blocks); + (Blocks::Narrow(blocks), bounds) + } else { + let blocks = match previous { + Some(Blocks::Wide(blocks)) => blocks, + _ => Vec::new(), + }; + let (blocks, bounds) = build_blocks(plan, basis, &index, shift, blocks); + (Blocks::Wide(blocks), bounds) + }; + ( + Self { + blocks, + shift, + dimension, + }, + bounds, + ) + } +} + +fn build_blocks( + plan: &PauliActionPlan, + basis: &[Word], + index: &CompactBasisIndex<'_, C>, + shift: u32, + mut blocks: Vec>, +) -> (Vec>, (f64, f64)) { + let dimension = basis.len(); + let rows = 1 << shift; + let count = dimension.div_ceil(rows); + let noisy = plan.has_diagonal_noise(); + blocks.resize_with(count, Block::default); + let metadata: Vec<_> = blocks + .par_iter_mut() + .enumerate() + .map_init( + || { + ( + Vec::::new(), + Vec::>::new(), + Vec::<(u32, Entry)>::new(), + Vec::::new(), + ) + }, + |(active, actions, staged, cursor), (block, storage)| { + let first = block * rows; + let length = rows.min(dimension - first); + staged.clear(); + storage.diagonal.clear(); + if noisy { + storage.diagonal.resize(length, 0.0); + } + let mut bounds = Bounds::default(); + for (row, word) in basis[first..first + length].iter().enumerate() { + plan.fill_units(word, active, actions); + // Hash order keeps successive index probes close together. + actions.sort_unstable_by(|a, b| { + a.hash.cmp(&b.hash).then_with(|| a.word.cmp(&b.word)) + }); + let mut raw = 0.0; + let mut diagonal = 0.0; + for action in actions.iter() { + raw += action.value.abs(); + if action.unit == 0 { + diagonal = action.value; + continue; + } + let Some(parent) = index.get_hashed(&action.word, action.hash) else { + continue; + }; + let parent = parent as usize; + let other = parent >> shift; + let stored = if other == block { + parent < first + row + } else { + owner(block, other) == block + }; + if stored { + staged.push(( + other as u32, + Entry { + row: row as u16, + parent: (parent - (other << shift)) as u16, + unit: U::pack(action.unit), + }, + )); + } + } + bounds.add(raw, diagonal); + if noisy { + storage.diagonal[row] = diagonal; + } + } + // Group entries by tile, keeping row order within each tile. + storage.tiles.clear(); + storage.tiles.resize(count + 1, 0); + for &(tile, _) in staged.iter() { + storage.tiles[tile as usize + 1] += 1; + } + for tile in 0..count { + storage.tiles[tile + 1] += storage.tiles[tile]; + } + // Entry storage keeps its capacity between builds and grows only when a + // build needs more, with 1/16 headroom, so capacity tracks the block's + // entry count and the allocation stays stable across exponentials. + storage.entries.clear(); + if storage.entries.capacity() < staged.len() { + storage.entries = Vec::new(); + storage + .entries + .reserve_exact(staged.len() + staged.len() / 16); + } + storage.entries.resize(staged.len(), Entry::default()); + cursor.clear(); + cursor.extend_from_slice(&storage.tiles[..count]); + for &(tile, entry) in staged.iter() { + let slot = &mut cursor[tile as usize]; + storage.entries[*slot as usize] = entry; + *slot += 1; + } + bounds + }, + ) + .collect(); + let bounds = metadata + .into_iter() + .fold(Bounds::default(), Bounds::merge) + .finish(dimension, noisy); + (blocks, bounds) +} + +pub(crate) struct UnitOp<'a, const C: usize> { + pub(crate) cache: &'a UnitCache, + pub(crate) plan: &'a PauliActionPlan, +} + +impl LinearOperator for UnitOp<'_, C> { + fn dim(&self) -> usize { + self.cache.dimension + } + fn parallel_hint(&self) -> bool { + false + } + fn dot(&self, overwrite: bool, input: &[f64], output: &mut [f64]) -> Result<(), QuSpinError> { + let shift = self.cache.shift; + match &self.cache.blocks { + Blocks::Narrow(blocks) => product(blocks, shift, self.plan, overwrite, input, output), + Blocks::Wide(blocks) => product(blocks, shift, self.plan, overwrite, input, output), + } + Ok(()) + } + fn trace(&self) -> f64 { + unreachable!("trace supplied to exponential") + } + fn onenorm(&self, _: f64) -> f64 { + unreachable!("norm supplied to exponential") + } + fn dot_transpose(&self, _: bool, _: &[f64], _: &mut [f64]) -> Result<(), QuSpinError> { + unreachable!("transpose unused by single-vector exponential") + } + fn dot_many( + &self, + _: bool, + _: ndarray::ArrayView2<'_, f64>, + _: ndarray::ArrayViewMut2<'_, f64>, + ) -> Result<(), QuSpinError> { + unreachable!("multiple vectors unused by single-vector exponential") + } + fn dot_chunk(&self, _: bool, _: &[f64], _: &mut [f64], _: usize) -> Result<(), QuSpinError> { + unreachable!("chunk products unused by single-vector exponential") + } + fn dot_transpose_chunk( + &self, + _: &[f64], + _: &[::Atomic], + _: std::ops::Range, + ) -> Result<(), QuSpinError> { + unreachable!("transpose chunks unused by single-vector exponential") + } +} + +fn product( + blocks: &[Block], + shift: u32, + plan: &PauliActionPlan, + overwrite: bool, + input: &[f64], + output: &mut [f64], +) { + let rows = 1usize << shift; + let slice = |block: usize| &input[block * rows..((block + 1) * rows).min(input.len())]; + // Rows gather from the pairs their block owns, in the stored direction. + output + .par_chunks_mut(rows) + .zip(blocks.par_iter()) + .enumerate() + .for_each(|(block, (output, storage))| { + let own = slice(block); + if storage.diagonal.is_empty() { + if overwrite { + output.fill(0.0); + } + } else { + for ((result, &d), &x) in output.iter_mut().zip(&storage.diagonal).zip(own) { + let value = if x != 0.0 { d * x } else { 0.0 }; + *result = if overwrite { value } else { *result + value }; + } + } + for other in 0..blocks.len() { + let x = slice(other); + let tile = storage.tiles[other] as usize..storage.tiles[other + 1] as usize; + for entry in &storage.entries[tile] { + let value = x[entry.parent as usize]; + if value != 0.0 { + output[entry.row as usize] += + plan.row_coefficient(entry.unit.unit()) * value; + } + } + } + }); + // Rows receive the mirrored direction of pairs owned by other blocks (and of + // their own block's pairs), with the opposite sign. + output + .par_chunks_mut(rows) + .enumerate() + .for_each(|(block, output)| { + for (other, storage) in blocks.iter().enumerate() { + let x = slice(other); + let tile = storage.tiles[block] as usize..storage.tiles[block + 1] as usize; + for entry in &storage.entries[tile] { + let value = x[entry.row as usize]; + if value != 0.0 { + output[entry.parent as usize] -= + plan.row_coefficient(entry.unit.unit()) * value; + } + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{JumpInput, LindbladSpec, PcStepConfig, word_from_codes}; + use num::Complex; + + #[test] + fn unit_products_match_general_generator_with_unequal_couplings_and_noise() { + let h = vec![ + ("XXII".into(), 0.3), + ("IYYI".into(), -0.71), + ("ZIIX".into(), 0.17), + ]; + let basis: Vec<_> = (0..256usize) + .filter(|code| code % 3 != 0) + .map(|code| { + word_from_codes::<2>( + &(0..4) + .map(|q| ((code >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let input: Vec<_> = (0..basis.len()) + .map(|i| if i % 7 == 0 { 0.0 } else { (i as f64).sin() }) + .collect(); + for rate in [0.0, 0.23] { + let jumps = vec![JumpInput { + lincomb: vec![("IZII".into(), Complex::new(1.0, 0.0))], + rate, + }]; + let spec = LindbladSpec::<2>::new(4, &h, &jumps).unwrap(); + let plan = spec.pauli_plan.as_ref().unwrap(); + let (cache, _) = UnitCache::build(plan, &basis, None); + let op = UnitOp { + cache: &cache, + plan, + }; + let mut expected = vec![0.0; basis.len()]; + for (row, column, value) in spec.generator(&basis) { + if input[column] != 0.0 { + expected[row] += value * input[column]; + } + } + for threads in [1, 2, 6] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + // Paired storage sums each row in a different order than the general + // generator, so products agree to rounding rather than bit for bit. + let close = |actual: &[f64], factor: f64| { + for (a, e) in actual.iter().zip(&expected) { + assert!( + (a - factor * e).abs() <= 1e-14 * (1.0 + e.abs()), + "{a} vs {e}" + ); + } + }; + let mut actual = vec![0.0; basis.len()]; + pool.install(|| op.dot(true, &input, &mut actual)).unwrap(); + close(&actual, 1.0); + pool.install(|| op.dot(false, &input, &mut actual)).unwrap(); + close(&actual, 2.0); + } + } + } + + #[test] + fn streaming_bounds_match_per_row_bounds() { + let rows: Vec<(f64, f64)> = (0..997) + .map(|i| { + let d = -((i * 37 % 101) as f64) * 0.013; + (d.abs() + ((i * 53 % 89) as f64) * 0.07, d) + }) + .collect(); + let expected = crate::mf_expm::real_bounds(&rows); + let merged = rows + .chunks(64) + .map(|chunk| { + let mut b = Bounds::default(); + for &(raw, d) in chunk { + b.add(raw, d); + } + b + }) + .fold(Bounds::default(), Bounds::merge) + .finish(rows.len(), true); + assert!((merged.0 - expected.0).abs() <= 1e-15); + assert!((merged.1 - expected.1).abs() <= 1e-13 * expected.1); + } + + #[test] + fn wide_term_codes_give_identical_products() { + let h = vec![ + ("XXII".into(), 0.3), + ("IYYI".into(), -0.71), + ("ZIIX".into(), 0.17), + ]; + let basis: Vec<_> = (0..256usize) + .filter(|code| code % 3 != 0) + .map(|code| { + word_from_codes::<2>( + &(0..4) + .map(|q| ((code >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let input: Vec<_> = (0..basis.len()).map(|i| (i as f64).sin()).collect(); + let jumps = vec![JumpInput { + lincomb: vec![("IZII".into(), Complex::new(1.0, 0.0))], + rate: 0.23, + }]; + let spec = LindbladSpec::<2>::new(4, &h, &jumps).unwrap(); + let plan = spec.pauli_plan.as_ref().unwrap(); + let (narrow, bounds) = UnitCache::build(plan, &basis, None); + assert!(matches!(narrow.blocks, Blocks::Narrow(_))); + let index = CompactBasisIndex::new(&basis); + let (blocks, wide_bounds) = + build_blocks::<2, i32>(plan, &basis, &index, narrow.shift, Vec::new()); + assert_eq!(bounds, wide_bounds); + let wide = UnitCache { + blocks: Blocks::Wide(blocks), + shift: narrow.shift, + dimension: narrow.dimension, + }; + let mut a = vec![0.0; basis.len()]; + let mut b = vec![0.0; basis.len()]; + UnitOp { + cache: &narrow, + plan, + } + .dot(true, &input, &mut a) + .unwrap(); + UnitOp { cache: &wide, plan } + .dot(true, &input, &mut b) + .unwrap(); + assert_eq!(a, b); + } + + #[test] + fn paired_products_cover_many_blocks_and_are_thread_count_invariant() { + // Five qubits with a three-body term and unequal couplings; the basis spans + // more than four blocks, so pairs of both owner parities and within-block + // pairs are exercised. + let h = vec![ + ("XXIII".into(), 0.31), + ("IYYII".into(), -0.52), + ("IIZXI".into(), 0.27), + ("XIYIZ".into(), 0.11), + ("IIIYY".into(), 0.9), + ]; + let basis: Vec<_> = (0..1024usize) + .filter(|code| code % 5 != 1) + .map(|code| { + word_from_codes::<2>( + &(0..5) + .map(|q| ((code >> (2 * q)) & 3) as u8) + .collect::>(), + ) + .unwrap() + }) + .collect(); + let input: Vec<_> = (0..basis.len()) + .map(|i| { + if i % 11 == 0 { + 0.0 + } else { + (0.7 * i as f64).cos() + } + }) + .collect(); + for rate in [0.0, 0.4] { + let jumps = vec![JumpInput { + lincomb: vec![("IIZII".into(), Complex::new(1.0, 0.0))], + rate, + }]; + let spec = LindbladSpec::<2>::new(5, &h, &jumps).unwrap(); + let plan = spec.pauli_plan.as_ref().unwrap(); + let (cache, _) = UnitCache::build(plan, &basis, None); + assert!(cache.block_count() > 4); + let op = UnitOp { + cache: &cache, + plan, + }; + let mut expected = vec![0.0; basis.len()]; + for (row, column, value) in spec.generator(&basis) { + expected[row] += value * input[column]; + } + let mut first: Option> = None; + for threads in [1, 3, 6] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let mut actual = vec![0.0; basis.len()]; + pool.install(|| op.dot(true, &input, &mut actual)).unwrap(); + for (a, e) in actual.iter().zip(&expected) { + assert!((a - e).abs() <= 1e-13 * (1.0 + e.abs()), "{a} vs {e}"); + } + match &first { + Some(reference) => assert_eq!(&actual, reference), + None => first = Some(actual), + } + } + } + } + + #[test] + fn adding_a_general_dissipator_replaces_a_retained_unit_cache() { + let h = vec![("XXI".into(), 0.3), ("IYY".into(), -0.17)]; + let mut reference = LindbladSpec::<2>::new(3, &h, &[]).unwrap(); + let mut optimized = LindbladSpec::<2>::new(3, &h, &[]).unwrap(); + let seed = word_from_codes::<2>(&[2, 0, 0]).unwrap(); + let (mut a, mut b) = (vec![seed], vec![seed]); + let (mut ca, mut cb) = (vec![1.0], vec![1.0]); + let cfg = PcStepConfig { + max_basis: 8, + admit_basis: Some(24), + num_threads: Some(2), + ..Default::default() + }; + for step in 0..8 { + if step == 3 { + let operators = vec![vec![ + ("XII".into(), Complex::new(0.5, 0.0)), + ("YII".into(), Complex::new(0.0, 0.5)), + ]]; + let k = vec![vec![Complex::new(0.2, 0.0)]]; + reference.add_kossakowski(&operators, &k).unwrap(); + optimized.add_kossakowski(&operators, &k).unwrap(); + } + reference + .pc_step_reference(&mut a, &mut ca, 0.1, &[seed], &cfg) + .unwrap(); + optimized + .pc_step(&mut b, &mut cb, 0.1, &[seed], &cfg) + .unwrap(); + assert_eq!(a, b); + assert_eq!(ca, cb); + } + } +} diff --git a/crates/ppvm-lindblad/src/word.rs b/crates/ppvm-lindblad/src/word.rs index d79e9fe57..fc8dd18d4 100644 --- a/crates/ppvm-lindblad/src/word.rs +++ b/crates/ppvm-lindblad/src/word.rs @@ -4,9 +4,10 @@ //! Packed Pauli-word type and the string / `u8`-label codecs. use crate::Error; +use bitvec::array::BitArray; use fxhash::FxBuildHasher; -use ppvm_pauli_word::word::PauliWord; -use ppvm_traits::PauliWordTrait; +use ppvm_traits::traits::HashFinalize; +use std::hash::{Hash, Hasher}; /// Pauli-word storage chunk: `u64` on 64-bit targets, `u32` elsewhere /// (`bitvec` implements `BitStore` for `u64` only on 64-bit targets, so @@ -51,14 +52,123 @@ pub const fn chunks_for(n_qubits: usize) -> Option { None } -/// The Pauli-word storage type used throughout this crate, `C` chunks wide. +/// The Pauli-word storage type used throughout this crate, `C` chunks wide: +/// one X and one Z bit plane, 32 bytes at the default 128-qubit width. /// /// The crate is const-generic in the chunk count so a register of any size -/// up to [`MAX_SUPPORTED_QUBITS`] gets a fixed-size word; the default -/// `C = W_CHUNKS` covers 128 qubits, byte-for-byte the historical layout. -/// The `FxBuildHasher` matches the hash used by the `FxHashMap` keys we -/// wrap with; `REHASH=true` means `set()` keeps the cached hash in sync. -pub type Word = PauliWord<[Chunk; C], FxBuildHasher, true>; +/// up to [`MAX_SUPPORTED_QUBITS`] gets a fixed-size word. The register size +/// belongs to the model and is not stored per word; bits at or above it are +/// always zero. The hash is computed from the planes on demand with the same +/// digest `ppvm_pauli_word::PauliWord` caches, so hash tables keyed by words +/// order them exactly as they would order that type. Ordering compares the X +/// plane, then the Z plane. +#[derive(Clone, Copy, Default)] +pub struct Word { + /// X-bit plane, one bit per qubit. + pub xbits: BitArray<[Chunk; C]>, + /// Z-bit plane, one bit per qubit. + pub zbits: BitArray<[Chunk; C]>, +} + +impl Word { + /// The identity word for a register of `n_qubits` (at most the width's capacity). + pub fn new(n_qubits: usize) -> Self { + debug_assert!(n_qubits <= max_qubits::()); + Self::default() + } + + #[inline(always)] + fn digest(&self) -> u64 { + FxBuildHasher::word_hash( + &[ + bytemuck::bytes_of(&self.xbits.data), + bytemuck::bytes_of(&self.zbits.data), + ], + std::mem::size_of::<[Chunk; C]>(), + ) + } +} + +/// The general `ppvm_pauli_word` word of the same width, which also stores the +/// register size and a cached hash. `ppvm_pauli_sum` symmetry routines take this type. +pub type GeneralWord = + ppvm_pauli_word::word::PauliWord<[Chunk; C], FxBuildHasher, true>; + +impl Word { + /// This word as the general `PauliWord` used by `ppvm_pauli_sum` symmetry groups. + pub fn to_pauli_word(self, n_qubits: usize) -> GeneralWord { + use ppvm_traits::PauliWordTrait; + let mut word = ppvm_pauli_word::word::PauliWord::new(n_qubits); + word.xbits = self.xbits; + word.zbits = self.zbits; + word.rehash(); + word + } + + /// The planes of a general `PauliWord`. + pub fn from_pauli_word(word: &GeneralWord) -> Self { + Self { + xbits: word.xbits, + zbits: word.zbits, + } + } +} + +impl PartialEq for Word { + #[inline(always)] + fn eq(&self, other: &Self) -> bool { + self.xbits.data == other.xbits.data && self.zbits.data == other.zbits.data + } +} + +impl Eq for Word {} + +/// Bit-lexicographic order of one plane from qubit 0 upward, as `BitArray`'s +/// `Ord` defines it, evaluated a chunk at a time. +#[inline(always)] +fn plane_cmp(a: &[Chunk; C], b: &[Chunk; C]) -> std::cmp::Ordering { + for (&x, &y) in a.iter().zip(b) { + let differ = x ^ y; + if differ != 0 { + return if (x >> differ.trailing_zeros()) & 1 == 0 { + std::cmp::Ordering::Less + } else { + std::cmp::Ordering::Greater + }; + } + } + std::cmp::Ordering::Equal +} + +impl Ord for Word { + #[inline(always)] + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + plane_cmp(&self.xbits.data, &other.xbits.data) + .then_with(|| plane_cmp(&self.zbits.data, &other.zbits.data)) + } +} + +impl PartialOrd for Word { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Hash for Word { + fn hash(&self, state: &mut H) { + state.write_u64(self.digest()); + } +} + +impl std::fmt::Debug for Word { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "Word(x={:x?}, z={:x?})", + self.xbits.data, self.zbits.data + ) + } +} /// `Err(TooManyQubits)` unless a `C`-chunk word holds `n_qubits`. pub(crate) fn check_width(n_qubits: usize) -> Result<(), Error> { @@ -89,13 +199,12 @@ pub fn word_from_codes(codes: &[u8]) -> Result, Error> { w.zbits.set(q, true); } } - w.rehash(); Ok(w) } /// Inverse of [`word_from_codes`]: write `n_qubits` Pauli labels into `out`. pub fn codes_from_word(w: &Word, out: &mut [u8]) { - debug_assert_eq!(out.len(), w.n_qubits()); + debug_assert!(out.len() <= max_qubits::()); for (q, slot) in out.iter_mut().enumerate() { let xb = w.xbits[q] as u8; let zb = w.zbits[q] as u8; @@ -138,16 +247,17 @@ pub fn parse_pauli_string( other => return Err(Error::InvalidPauliChar { c: other }), } } - w.rehash(); Ok((w, support)) } /// Compute the support (non-identity qubits) of `w`. pub(crate) fn word_support(w: &Word, out: &mut Vec) { out.clear(); - for q in 0..w.n_qubits() { - if w.xbits[q] || w.zbits[q] { - out.push(q as u32); + for chunk in 0..C { + let mut bits = w.xbits.data[chunk] | w.zbits.data[chunk]; + while bits != 0 { + out.push((chunk * CHUNK_BITS) as u32 + bits.trailing_zeros()); + bits &= bits - 1; } } } diff --git a/crates/ppvm-lindblad/tests/candidate_proposers.rs b/crates/ppvm-lindblad/tests/candidate_proposers.rs new file mode 100644 index 000000000..45181bec6 --- /dev/null +++ b/crates/ppvm-lindblad/tests/candidate_proposers.rs @@ -0,0 +1,222 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Public strategy integration and borrowed-tail lifetime checks. + +use ppvm_lindblad::{ + CandidateBuffer, CandidateProposer, Error, LindbladSpec, PcStepConfig, PcStepWorkspace, + StreamingProposer, Word, word_from_codes, +}; + +#[derive(Default)] +struct DelegatingProposer { + strategy: StreamingProposer, + calls: usize, + selections: Vec, +} + +impl CandidateProposer for DelegatingProposer { + fn propose( + &mut self, + model: &LindbladSpec, + buffer: &mut CandidateBuffer<'_>, + protected: &[Word], + maximum: usize, + ) -> Result<(), Error> { + self.calls += 1; + assert_eq!(buffer.basis().len(), buffer.coefficients().len()); + assert!(buffer.candidates().is_empty()); + self.strategy.propose(model, buffer, protected, maximum) + } + + fn select(&mut self, buffer: &mut CandidateBuffer<'_>, threshold: f64) -> usize { + let count = self.strategy.select(buffer, threshold); + self.selections.push(count); + count + } +} + +#[test] +fn explicit_strategy_matches_default_through_predictor_and_corrector() { + let model = LindbladSpec::new( + 4, + &[ + ("XXII".into(), 0.3), + ("IYYI".into(), -0.7), + ("IIXX".into(), 0.2), + ], + &[], + ) + .unwrap(); + let seed = word_from_codes(&[2, 0, 0, 0]).unwrap(); + let (mut a, mut b) = (vec![seed], vec![seed]); + let (mut ca, mut cb) = (vec![1.0], vec![1.0]); + let mut workspace = PcStepWorkspace::default(); + let mut proposer = DelegatingProposer::default(); + let cfg = PcStepConfig { + max_basis: 8, + admit_basis: Some(24), + num_threads: Some(2), + ..Default::default() + }; + let mut correctors = 0; + for _ in 0..12 { + model.pc_step(&mut a, &mut ca, 0.1, &[seed], &cfg).unwrap(); + let first_selection = proposer.selections.len(); + model + .pc_step_with_proposer( + &mut b, + &mut cb, + 0.1, + &[seed], + &cfg, + &mut workspace, + &mut proposer, + ) + .unwrap(); + correctors += usize::from( + proposer + .selections + .get(first_selection + 1) + .is_some_and(|count| *count > 0), + ); + assert_eq!(a, b); + assert_eq!(ca, cb); + } + assert!(proposer.calls > 12); + assert!(correctors > 0); +} + +struct FixedProposer { + rejected: Word, + accepted: Word, +} + +impl CandidateProposer for FixedProposer { + fn propose( + &mut self, + _: &LindbladSpec, + buffer: &mut CandidateBuffer<'_>, + _: &[Word], + _: usize, + ) -> Result<(), Error> { + buffer.push(self.rejected); + buffer.push(self.accepted); + buffer.numeric_scratch().extend([0.0, 1.0]); + Ok(()) + } + + fn select(&mut self, buffer: &mut CandidateBuffer<'_>, threshold: f64) -> usize { + if buffer.numeric_scratch()[1].abs() > threshold { + buffer.candidates_mut().swap(0, 1); + 1 + } else { + 0 + } + } +} + +#[test] +fn independent_strategy_reuses_main_tail_and_admission_initializes_zero_coefficients() { + let model = LindbladSpec::new(2, &[], &[]).unwrap(); + let seed = word_from_codes(&[2, 0]).unwrap(); + let accepted = word_from_codes(&[0, 1]).unwrap(); + let mut words = Vec::with_capacity(64); + words.push(seed); + let address = words.as_ptr(); + let mut coefficients = vec![0.75]; + let mut workspace = PcStepWorkspace::default(); + let mut proposer = FixedProposer { + rejected: word_from_codes(&[1, 0]).unwrap(), + accepted, + }; + let cfg = PcStepConfig { + max_basis: 2, + num_threads: Some(1), + ..Default::default() + }; + model + .pc_step_with_proposer( + &mut words, + &mut coefficients, + 0.0, + &[seed], + &cfg, + &mut workspace, + &mut proposer, + ) + .unwrap(); + assert_eq!(words, vec![seed, accepted]); + assert_eq!(coefficients, vec![0.75, 0.0]); + assert_eq!(words.as_ptr(), address); +} + +struct FailingProposer; + +impl CandidateProposer for FailingProposer { + fn propose( + &mut self, + _: &LindbladSpec, + buffer: &mut CandidateBuffer<'_>, + _: &[Word], + _: usize, + ) -> Result<(), Error> { + buffer.push(word_from_codes(&[1, 0]).unwrap()); + buffer.numeric_scratch().push(1.0); + Err(Error::Internal("candidate generation failed".into())) + } + fn select(&mut self, _: &mut CandidateBuffer<'_>, _: f64) -> usize { + unreachable!() + } +} + +#[test] +fn proposal_error_discards_temporary_tail_and_allows_default_step_to_resume() { + let model = LindbladSpec::new(2, &[("XX".into(), 0.3)], &[]).unwrap(); + let seed = word_from_codes(&[2, 0]).unwrap(); + let mut words = vec![seed]; + let mut coefficients = vec![0.75]; + let mut workspace = PcStepWorkspace::default(); + let cfg = PcStepConfig { + max_basis: 8, + num_threads: Some(1), + ..Default::default() + }; + assert!( + model + .pc_step_with_proposer( + &mut words, + &mut coefficients, + 0.1, + &[seed], + &cfg, + &mut workspace, + &mut FailingProposer + ) + .is_err() + ); + assert_eq!(words, vec![seed]); + assert_eq!(coefficients, vec![0.75]); + let (mut expected_words, mut expected_coefficients) = (words.clone(), coefficients.clone()); + model + .pc_step( + &mut expected_words, + &mut expected_coefficients, + 0.1, + &[seed], + &cfg, + ) + .unwrap(); + model + .pc_step_with_workspace( + &mut words, + &mut coefficients, + 0.1, + &[seed], + &cfg, + &mut workspace, + ) + .unwrap(); + assert_eq!(words, expected_words); + assert_eq!(coefficients, expected_coefficients); +} diff --git a/crates/ppvm-lindblad/tests/word_width.rs b/crates/ppvm-lindblad/tests/word_width.rs index 7d3b5ea46..81b5bcee3 100644 --- a/crates/ppvm-lindblad/tests/word_width.rs +++ b/crates/ppvm-lindblad/tests/word_width.rs @@ -92,6 +92,7 @@ fn evolve(n_total: usize, n_active: usize, steps: usize, koss: b drop_tol: 0.0, tau_add: None, num_threads: Some(1), + ..Default::default() }; for _ in 0..steps { spec.pc_step(&mut basis, &mut coeffs, 0.05, &[], &cfg) @@ -133,7 +134,7 @@ fn evolve_orbit(n: usize, k: i32, steps: usize) -> Vec<(Vec, let group = TranslationGroup::chain_1d(n); let k_modes = [k]; - let mut basis: Vec> = (0..n) + let basis: Vec> = (0..n) .map(|q| { let mut codes = vec![0u8; n]; codes[q] = 1; // X @@ -148,7 +149,9 @@ fn evolve_orbit(n: usize, k: i32, steps: usize) -> Vec<(Vec, ) }) .collect(); - canonicalize_pauli_sum_complex(&mut basis, &mut coeffs, &group, &k_modes); + let mut words: Vec<_> = basis.iter().map(|w| w.to_pauli_word(n)).collect(); + canonicalize_pauli_sum_complex(&mut words, &mut coeffs, &group, &k_modes); + let mut basis: Vec> = words.iter().map(Word::from_pauli_word).collect(); let sector = Sector::new(&group, &k_modes); let cfg = PcStepConfig { @@ -157,6 +160,7 @@ fn evolve_orbit(n: usize, k: i32, steps: usize) -> Vec<(Vec, drop_tol: 0.0, tau_add: None, num_threads: Some(1), + ..Default::default() }; for _ in 0..steps { spec.pc_step_orbit_rep(&mut basis, &mut coeffs, 0.1, &[], §or, &cfg) diff --git a/crates/ppvm-python-native/src/lindblad.rs b/crates/ppvm-python-native/src/lindblad.rs index dc209cd37..53fd27607 100644 --- a/crates/ppvm-python-native/src/lindblad.rs +++ b/crates/ppvm-python-native/src/lindblad.rs @@ -38,6 +38,17 @@ pub(crate) fn map_err(e: ppvm_lindblad::Error) -> PyErr { PyValueError::new_err(e.to_string()) } +fn candidate_strategy(name: &str) -> PyResult { + use ppvm_lindblad::CandidateStrategy::{Shortlist, Streaming}; + match name { + "streaming" => Ok(Streaming), + "shortlist" => Ok(Shortlist), + _ => Err(PyValueError::new_err( + "proposer must be streaming or shortlist", + )), + } +} + /// Reject a basis that contains the same Pauli word at two distinct rows. /// Duplicate rows would silently overwrite each other in the generator's /// row-index map and produce an incorrect sparse matrix. @@ -284,7 +295,8 @@ impl LindbladSpec { /// cap selects the top-`max_basis` strings over the whole union /// (displacement truncation). `drop_tol` prunes by magnitude after the /// step; `tau_add` filters leakage admission by inflow rate. Protected - /// words are never dropped. + /// words are never dropped. `proposer` selects streaming (the default) or + /// shortlist candidate generation. #[pyo3(signature = ( basis, coeffs, dt, max_basis, drop_tol = 0.0, @@ -292,6 +304,8 @@ impl LindbladSpec { num_threads = None, admit_basis = None, tau_add = None, + *, + proposer = "streaming", ))] #[allow(clippy::too_many_arguments)] fn pc_step<'py>( @@ -306,7 +320,9 @@ impl LindbladSpec { num_threads: Option, admit_basis: Option, tau_add: Option, + proposer: &str, ) -> PyResult> { + let proposer = candidate_strategy(proposer)?; with_spec!(&self.inner, inner, C => { let n_q = inner.n_qubits(); let basis_view = basis.as_array(); @@ -326,6 +342,7 @@ impl LindbladSpec { dt, &protected_words, &ppvm_lindblad::PcStepConfig { + proposer, max_basis, admit_basis, drop_tol, @@ -335,9 +352,9 @@ impl LindbladSpec { ) .map_err(map_err)?; - // Pack output. Basis may have grown; coeffs has the same new length. - let pairs: Vec<(Word, f64)> = basis_words.into_iter().zip(coeffs_vec).collect(); - pack_pauli_map(py, pairs, n_q) + let basis_arr = encode_basis(py, &basis_words, n_q)?; + coeffs_vec.shrink_to_fit(); + Ok((basis_arr, coeffs_vec.into_pyarray(py))) }) } @@ -350,6 +367,8 @@ impl LindbladSpec { num_threads = None, admit_basis = None, tau_add = None, + *, + proposer = "streaming", ))] #[allow(clippy::too_many_arguments)] fn pc_step_timed<'py>( @@ -364,7 +383,9 @@ impl LindbladSpec { num_threads: Option, admit_basis: Option, tau_add: Option, + proposer: &str, ) -> PyResult<(PyPauliMap<'py>, Bound<'py, pyo3::types::PyDict>)> { + let proposer = candidate_strategy(proposer)?; with_spec!(&self.inner, inner, C => { let n_q = inner.n_qubits(); let basis_view = basis.as_array(); @@ -384,6 +405,7 @@ impl LindbladSpec { dt, &protected_words, &ppvm_lindblad::PcStepConfig { + proposer, max_basis, admit_basis, drop_tol, @@ -393,8 +415,9 @@ impl LindbladSpec { ) .map_err(map_err)?; - let pairs: Vec<(Word, f64)> = basis_words.into_iter().zip(coeffs_vec).collect(); - let map = pack_pauli_map(py, pairs, n_q)?; + let basis_arr = encode_basis(py, &basis_words, n_q)?; + coeffs_vec.shrink_to_fit(); + let map = (basis_arr, coeffs_vec.into_pyarray(py)); let d = pyo3::types::PyDict::new(py); d.set_item("leakage1_us", timings.leakage1_us)?; d.set_item("expand1_us", timings.expand1_us)?; @@ -495,6 +518,7 @@ impl LindbladSpec { drop_tol, tau_add, num_threads, + ..Default::default() }, ) .map_err(map_err)?; diff --git a/crates/ppvm-python-native/src/symmetry.rs b/crates/ppvm-python-native/src/symmetry.rs index 4d8129a8a..9b8941c90 100644 --- a/crates/ppvm-python-native/src/symmetry.rs +++ b/crates/ppvm-python-native/src/symmetry.rs @@ -161,8 +161,8 @@ impl TranslationGroup { let mut out = vec![0u8; codes.len()]; with_width!(codes.len(), C => { let w = word_from_codes::(codes).map_err(|e| PyValueError::new_err(e.to_string()))?; - let canon = self.inner.canonicalize(&w); - codes_from_word(&canon, &mut out); + let canon = self.inner.canonicalize(&w.to_pauli_word(codes.len())); + codes_from_word(&ppvm_lindblad::Word::from_pauli_word(&canon), &mut out); PyResult::Ok(()) })?; Ok(out.into_pyarray(py)) @@ -205,12 +205,14 @@ pub fn canonicalize_basis_arr_complex<'py>( .map(|c| Complex::new(c.re, c.im)) .collect(); + let mut general = to_general(&basis_words, n_q); core_sym::canonicalize_pauli_sum_complex( - &mut basis_words, + &mut general, &mut coeffs_vec, &group.inner, k_slice, ); + basis_words = from_general(&general); let out_coeffs: Vec = coeffs_vec .iter() @@ -249,7 +251,7 @@ pub fn check_momentum_sector_arr<'py>( .iter() .map(|c| Complex::new(c.re, c.im)) .collect(); - core_sym::check_momentum_sector(&basis_words, &coeffs_vec, &group.inner, k_slice, tol) + core_sym::check_momentum_sector(&to_general(&basis_words, n_q), &coeffs_vec, &group.inner, k_slice, tol) .map_err(|e| PyValueError::new_err(format!("{e}"))) }) } @@ -283,9 +285,29 @@ pub fn canonicalize_basis_arr<'py>( let mut basis_words = decode_basis::(&basis_view, n_q)?; let mut coeffs_vec = coeffs_slice.to_vec(); - core_sym::canonicalize_pauli_sum(&mut basis_words, &mut coeffs_vec, &group.inner); + let mut general = to_general(&basis_words, n_q); + core_sym::canonicalize_pauli_sum(&mut general, &mut coeffs_vec, &group.inner); + basis_words = from_general(&general); let basis_arr = encode_basis(py, &basis_words, n_q)?; Ok((basis_arr, coeffs_vec.into_pyarray(py))) }) } + +/// `ppvm_pauli_sum` symmetry routines take the general word type; the Lindblad words +/// store only their bit planes. +fn to_general( + words: &[ppvm_lindblad::Word], + n_qubits: usize, +) -> Vec> { + words.iter().map(|w| w.to_pauli_word(n_qubits)).collect() +} + +fn from_general( + words: &[ppvm_lindblad::GeneralWord], +) -> Vec> { + words + .iter() + .map(ppvm_lindblad::Word::from_pauli_word) + .collect() +} diff --git a/ppvm-python/src/ppvm/_core.pyi b/ppvm-python/src/ppvm/_core.pyi index a2c435cdb..d7d09962a 100644 --- a/ppvm-python/src/ppvm/_core.pyi +++ b/ppvm-python/src/ppvm/_core.pyi @@ -396,6 +396,8 @@ class LindbladSpec: num_threads: int | None = None, admit_basis: int | None = None, tau_add: float | None = None, + *, + proposer: str = "streaming", ) -> tuple[np.ndarray, np.ndarray]: ... def pc_step_timed( self, @@ -408,6 +410,8 @@ class LindbladSpec: num_threads: int | None = None, admit_basis: int | None = None, tau_add: float | None = None, + *, + proposer: str = "streaming", ) -> tuple[tuple[np.ndarray, np.ndarray], dict[str, int]]: ... def pc_step_orbit_rep( self, diff --git a/ppvm-python/src/ppvm/lindblad.py b/ppvm-python/src/ppvm/lindblad.py index 28f25a934..e3fa17886 100644 --- a/ppvm-python/src/ppvm/lindblad.py +++ b/ppvm-python/src/ppvm/lindblad.py @@ -305,13 +305,16 @@ def pc_step_arr( num_threads: int | None = None, admit_basis: int | None = None, tau_add: float | None = None, + *, + proposer: str = "streaming", ) -> tuple[np.ndarray, np.ndarray]: """One predictor-corrector adaptive step. All work — leakage expansion, matrix-exponential step, second-hop re-expansion, corrector — runs in Rust; SciPy is not required. - The matrix-exponential action is computed matrix-free via the - external ``quspin-expm`` crate. + The matrix-exponential action uses sparse generator rows and the + external ``quspin-expm`` crate. Repeated calls on this model reuse + temporary buffers automatically. Couplings may have arbitrary strengths. Truncation. ``max_basis`` is a hard rank cap on the live basis: enrichment adds at most ``max_basis - len(basis)`` of the largest @@ -339,6 +342,9 @@ def pc_step_arr( turnover. With the default ``None``, admission is bounded by ``max_basis`` and turnover requires ``drop_tol > 0``. + ``proposer`` chooses ``"streaming"`` (default) or ``"shortlist"`` + candidate generation. + Returns ``(new_basis_arr, new_coeffs)``; the basis may have grown (or shrunk, if ``max_basis`` / ``drop_tol`` pruned entries). """ @@ -355,6 +361,7 @@ def pc_step_arr( None if num_threads is None else int(num_threads), None if admit_basis is None else int(admit_basis), None if tau_add is None else float(tau_add), + proposer=proposer, ) def pc_step_orbit_rep( @@ -438,6 +445,8 @@ def pc_step( num_threads: int | None = None, admit_basis: int | None = None, tau_add: float | None = None, + *, + proposer: str = "streaming", ) -> tuple[list[str], np.ndarray]: """String-keyed variant of `pc_step_arr`.""" n = self.n_qubits @@ -455,6 +464,7 @@ def pc_step( num_threads, admit_basis, tau_add, + proposer=proposer, ) return _codes_to_basis(new_basis_arr), new_coeffs diff --git a/ppvm-python/test/lindblad/test_proposer_options.py b/ppvm-python/test/lindblad/test_proposer_options.py new file mode 100644 index 000000000..6c1ba5247 --- /dev/null +++ b/ppvm-python/test/lindblad/test_proposer_options.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: 2026 The PPVM Authors +# SPDX-License-Identifier: Apache-2.0 +"""Candidate-strategy options across Python entry points.""" + +import numpy as np +import pytest + +from ppvm import Lindbladian + + +def model(): + return Lindbladian(4, [("XXII", 0.7), ("IYYI", 0.8), ("IIXX", 0.9), ("ZIII", 0.3)], []) + + +def evolve(op, dt=0.025, **options): + basis, coeffs = ["ZIII"], np.array([1.0]) + config = dict(max_basis=4, admit_basis=12, drop_tol=0.0, protected=basis, **options) + for _ in range(100): + basis, coeffs = op.pc_step(basis, coeffs, dt, **config) + assert "ZIII" in basis and len(basis) <= 4 + assert np.isfinite(coeffs).all() + return basis, coeffs + + +@pytest.mark.parametrize("proposer", ["streaming", "shortlist"]) +def test_proposer_options(proposer): + op = model() + ordinary = evolve(op, proposer=proposer) + if proposer == "streaming": + implicit = evolve(op) + assert implicit[0] == ordinary[0] + np.testing.assert_array_equal(implicit[1], ordinary[1]) + + +@pytest.mark.parametrize("proposer", ["streaming", "shortlist"]) +def test_array_and_timed_options_agree(proposer): + op = model() + basis = np.array([[2, 0, 0, 0]], dtype=np.uint8) + coeffs = np.array([1.0]) + options = dict(drop_tol=0.0, admit_basis=12, proposer=proposer) + a = op.pc_step_arr(basis, coeffs, 0.025, 4, protected_arr=basis, **options) + b, timings = op._spec.pc_step_timed(basis, coeffs, 0.025, 4, protected=basis, **options) + np.testing.assert_array_equal(a[0], b[0]) + np.testing.assert_array_equal(a[1], b[1]) + assert all(value >= 0 for value in timings.values()) + + +def test_invalid_proposer(): + with pytest.raises(ValueError): + evolve(model(), proposer="unknown")