diff --git a/crates/ppvm-lindblad/src/lib.rs b/crates/ppvm-lindblad/src/lib.rs index 86b7c623..9e2d1879 100644 --- a/crates/ppvm-lindblad/src/lib.rs +++ b/crates/ppvm-lindblad/src/lib.rs @@ -47,6 +47,10 @@ pub(crate) mod expm; /// Matrix-free / quspin-expm-backed `exp(dt·L*)·b` engine. See module docs. pub(crate) mod mf_expm; +/// Per-step orbit-rep evolution under translation symmetry, with a +/// phase-aware complex action. See module docs. +pub mod orbit_rep; + /// Pauli-word storage chunk: `u64` on 64-bit targets, `u32` elsewhere /// (`bitvec` implements `BitStore` for `u64` only on 64-bit targets, so /// e.g. wasm32 builds use four 32-bit chunks instead of two 64-bit ones). @@ -1135,6 +1139,109 @@ mod tests { assert_eq!(out.as_slice(), &codes); } + /// Per-step orbit-rep evolution gives the SAME final orbit-rep + /// state as full-basis complex evolution followed by a single + /// projection at the end. Validates that the phase-aware complex + /// action machinery is consistent with the full-basis reference. + #[test] + fn pc_step_orbit_rep_matches_full_basis_projection() { + use ppvm_pauli_sum::symmetry::canonicalize_pauli_sum_complex; + use std::f64::consts::PI; + let n = 4usize; + let dt = 0.01f64; + let n_steps = 3usize; + let mut h_terms: Vec<(String, f64)> = Vec::new(); + for j in 0..n { + let nxt = (j + 1) % n; + for op in ["X", "Y"] { + let mut s = vec!['I'; n]; + s[j] = op.chars().next().unwrap(); + s[nxt] = op.chars().next().unwrap(); + h_terms.push((s.into_iter().collect(), 1.0)); + } + } + let spec = LindbladSpec::new(n, &h_terms, &[]).unwrap(); + let group = ppvm_pauli_sum::symmetry::TranslationGroup::chain_1d(n); + let k_mode: i32 = 1; + let k = vec![k_mode]; + + // Build the k=1 eigenstate in FULL basis form. + let basis_full: Vec = (0..n) + .map(|j| { + let mut s = vec!['I'; n]; + s[j] = 'Z'; + let (w, _) = parse_pauli_string(&s.into_iter().collect::(), n).unwrap(); + w + }) + .collect(); + let coeffs_full: Vec> = (0..n as i32) + .map(|a| { + Complex::from_polar(1.0, -2.0 * PI * (k_mode as f64) * (a as f64) / (n as f64)) + }) + .collect(); + + // ----- Full-basis path ----- + let mut bf = basis_full.clone(); + let mut cf = coeffs_full.clone(); + let protected: Vec = Vec::new(); + for _ in 0..n_steps { + // Full enrichment (tau_add = 0.0 adds every leakage string): + // for a momentum eigenstate the leakage is pure-sector, so the + // full-basis and orbit-rep paths build corresponding bases and + // the projection theorem gives an exact match. The orbit-rep + // side uses a large max_basis so its rank cap never binds. + pc_step_complex_full(&spec, &mut bf, &mut cf, dt); + } + // Project at the end. + canonicalize_pauli_sum_complex(&mut bf, &mut cf, &group, &k); + + // ----- Orbit-rep path ----- + // Initial orbit-rep form: project the full-basis input. + let mut br = basis_full.clone(); + let mut cr = coeffs_full.clone(); + canonicalize_pauli_sum_complex(&mut br, &mut cr, &group, &k); + // Evolve in orbit-rep form (max_basis large ⇒ full enrichment). + for _ in 0..n_steps { + orbit_rep::pc_step_orbit_rep( + &spec, + &mut br, + &mut cr, + dt, + &protected, + &group, + &k, + &PcStepConfig { + max_basis: 10_000_000, + ..Default::default() + }, + ) + .unwrap(); + } + + // Compare. + let mf: FxHashMap> = bf.into_iter().zip(cf).collect(); + let mr: FxHashMap> = br.into_iter().zip(cr).collect(); + assert_eq!( + mf.len(), + mr.len(), + "orbit-rep ({}) and full-basis-projected ({}) basis sizes differ", + mr.len(), + mf.len() + ); + let mut max_diff = 0.0_f64; + for (w, cm) in &mr { + let cf_val = mf + .get(w) + .copied() + .unwrap_or_else(|| panic!("rep {:?} in orbit-rep but not in full-basis", w)); + max_diff = max_diff.max((cm - cf_val).norm()); + } + assert!( + max_diff < 1e-9, + "orbit-rep diverged from full-basis: max |Δc| = {max_diff:e}" + ); + } + /// The full-space complex step at momentum k=0 must reproduce the real /// pc_step on the same trajectory exactly. #[test] diff --git a/crates/ppvm-lindblad/src/orbit_rep.rs b/crates/ppvm-lindblad/src/orbit_rep.rs new file mode 100644 index 00000000..c72e3da4 --- /dev/null +++ b/crates/ppvm-lindblad/src/orbit_rep.rs @@ -0,0 +1,463 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Per-step orbit-representative evolution under translation symmetry. +//! +//! The state lives entirely in **orbit-rep form** throughout: `basis` +//! contains only canonical translation-orbit representatives, and +//! `coeffs` are complex (one per rep). The dynamics `L*` is computed +//! with **phase-aware action** — for each output Pauli `q`, we +//! canonicalize `q` to its orbit rep `r_q` with shift counter `cnt_q`, +//! and accumulate `χ_k(g_{cnt_q}) · v · c_r` (where `v` is the matrix +//! element of `L*` between input rep `r` and output `q`). +//! +//! The orbit-rep basis is ~|G|× smaller than the full-basis +//! representation, throughout the entire evolution. +//! +//! The phase-aware action is genuinely **complex** (because of the +//! `χ_k(g)` phase factors). Rather than materialise a sparse matrix, the +//! per-column action — for each input rep, the list of `(row, χ_k·v)` +//! pairs for the in-basis outputs — is computed **once per expm call** +//! (via [`build_orbit_rep_cols`]) and then reused, CSC-style, across +//! every Krylov–Taylor matvec driving the external `quspin-expm` engine. +//! +//! ## Limitations +//! +//! - Caller is responsible for ensuring the input basis is in orbit-rep +//! form (i.e. each entry is the canonical representative of its +//! translation orbit). Use [`canonicalize_basis_to_rep`] if needed. +//! - The momentum sector `k_modes` is fixed for the duration of one +//! pc_step call. To compute a full site-resolved profile, call +//! `pc_step_orbit_rep` once per momentum mode and inverse-Fourier +//! the results. + +use crate::mf_expm::CscOp; +use crate::{Error, LindbladSpec}; +use fxhash::{FxBuildHasher, FxHashMap}; +use num::Complex; +use ppvm_pauli_sum::symmetry::TranslationGroup; +use quspin_expm::ExpmOp; +use rayon::prelude::*; + +// Word type re-exported from lib.rs. +use crate::Word; + +/// Replace each entry of `basis` with its canonical orbit +/// representative under `group`. Pure rewrite; coefficients are +/// untouched. Useful to enforce the orbit-rep invariant before calling +/// [`pc_step_orbit_rep`]. +/// +/// Does NOT deduplicate — if multiple input entries collapse to the +/// 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); + } +} + +/// Build the per-column phase-aware action of the in-basis-restricted +/// orbit-rep generator `M` at momentum sector `k_modes`. +/// +/// Returns, for each input rep `basis[c]` (column `c`), the list of +/// `(row, χ_k(g_{cnt_q}) · v_q)` pairs for every action output Pauli `q` +/// of `L*(basis[c])` whose orbit rep `r_q` is in `basis` at index `row`. +/// Outputs not in `basis` are dropped. This is the expensive part of the +/// orbit-rep dynamics (`compute_action_terms`, `canonicalize_with_shift`, +/// `character`); it is computed once and reused by the CSC-style matvec +/// in [`CscOp`]. +pub(crate) fn build_orbit_rep_cols( + spec: &LindbladSpec, + basis: &[Word], + index: &FxHashMap, + group: &TranslationGroup, + k_modes: &[i32], +) -> Vec)>> { + basis + .par_iter() + .map_init( + || { + ( + Vec::::with_capacity(spec.n_qubits()), + Vec::::with_capacity(128), + FxHashMap::>::with_capacity_and_hasher( + 128, + FxBuildHasher::default(), + ), + ) + }, + |(s1, s2, lm), r| { + let terms = spec.compute_action_terms(r, s1, s2, lm); + let mut out = Vec::with_capacity(terms.len()); + for (q, v) in terms.iter() { + let (r_q, cnt_q) = group.canonicalize_with_shift(q); + if let Some(&row) = index.get(&r_q) { + let phase = group.character(k_modes, &cnt_q); + out.push((row, phase * *v)); + } + } + out + }, + ) + .collect() +} + +/// Compute `exp(dt · M) · coeffs` for the in-basis-restricted orbit-rep +/// generator `M` at momentum sector `k_modes`, via `quspin-expm`. Returns +/// a fresh `Vec>` of length `basis.len()`. +/// +/// The expensive phase-aware action is computed ONCE here (via +/// [`build_orbit_rep_cols`]) and reused, CSC-style, across every Krylov– +/// Taylor matvec (see [`CscOp`]). One pass over the cached columns +/// extracts the diagonal shift `μ = tr(M)/n` and a valid upper bound on +/// the column 1-norm of `M − μ·I`; from `‖dt·(M−μI)‖₁` we pick the Taylor +/// partition `(m*, s)` and hand everything to +/// [`quspin_expm::ExpmOp::from_parts`] (mirroring +/// [`crate::mf_expm::expm_apply_mf`]). +pub(crate) fn expm_apply_orbit_rep_cached( + spec: &LindbladSpec, + basis: &[Word], + group: &TranslationGroup, + k_modes: &[i32], + dt: f64, + coeffs: &[Complex], +) -> Vec> { + let n = basis.len(); + if n == 0 { + return Vec::new(); + } + + let index = crate::build_basis_index(basis); + let cols = build_orbit_rep_cols(spec, basis, &index, group, k_modes); + + // One pass for the per-column `(raw, diag)` used by the `μ`/1-norm + // selection: `raw = Σ|val|` (upper bound on the absolute column sum), + // `diag = M[c,c]`. From these: `trace = Σ diag`, `μ = trace/n`, and an + // upper bound on the column 1-norm of `M − μ·I`: `raw − |diag| + |diag − μ|`. + let per_col: Vec<(f64, Complex)> = cols + .par_iter() + .enumerate() + .map(|(c, col)| { + let mut raw = 0.0_f64; + let mut diag = Complex::new(0.0, 0.0); + for &(row, val) in col.iter() { + raw += val.norm(); + if row as usize == c { + diag += val; + } + } + (raw, diag) + }) + .collect(); + + let trace: Complex = per_col.iter().map(|(_, d)| *d).sum(); + let mu = trace / n as f64; + let onenorm = per_col + .iter() + .map(|(raw, diag)| raw - diag.norm() + (diag - mu).norm()) + .fold(0.0_f64, f64::max); + + let (m_star, s) = crate::expm::select_ms(dt.abs() * onenorm); + + let mut v = coeffs.to_vec(); + let op = CscOp { + cols: &cols, + dim: n, + }; + let e = ExpmOp::from_parts( + op, + Complex::new(dt, 0.0), + mu, + s as usize, + m_star as usize, + 1e-12_f64, + ); + e.apply(ndarray::ArrayViewMut1::from(v.as_mut_slice())) + .expect("expm apply (orbit-rep cached)"); + v +} + +/// Phase-aware leakage: out-of-basis component of `L*(O_k)` where `O_k` +/// is the operator represented by `basis` (orbit reps) and `coeffs` +/// (complex coefficients in momentum sector `k_modes`). +/// +/// For each input rep `r` with coefficient `c_r`, and each output `q` +/// of `L*(r) = Σ_q v_q · q`: +/// 1. Canonicalize `q` → `(r_q, cnt_q)`. +/// 2. If `r_q` NOT in `basis` and NOT in `protected`: +/// `merged[r_q] += χ_k(g_{cnt_q}) · v_q · c_r`. +/// +/// Returns `(r_q, sum)` pairs for all candidates with nonzero sum. +/// +/// The live candidate map is capped to the *available room* +/// `room = max_basis − basis.len()` (the reps we could actually add), +/// applied during accumulation: input reps are processed in descending +/// `|c|` order and after each chunk only the `room` largest-`|sum|` +/// candidates are kept. A large `max_basis` (room ≥ all candidates) +/// disables the cap — the near-exact case. +#[allow(clippy::too_many_arguments)] +pub fn leakage_orbit_rep( + spec: &LindbladSpec, + basis: &[Word], + coeffs: &[Complex], + protected: &[Word], + group: &TranslationGroup, + k_modes: &[i32], + max_basis: usize, +) -> Result)>, Error> { + if basis.len() != coeffs.len() { + return Err(Error::LengthMismatch { + what: "basis and coeffs", + a: basis.len(), + b: coeffs.len(), + }); + } + let in_basis: FxHashMap<&Word, ()> = basis.iter().map(|w| (w, ())).collect(); + let protected_set: FxHashMap<&Word, ()> = protected.iter().map(|w| (w, ())).collect(); + + // Descending sort by |c|: process largest-magnitude contributors first + // so the running room-cap keeps the right entries. + let mut order: Vec = (0..basis.len()).collect(); + order.sort_by(|&a, &b| { + coeffs[b] + .norm() + .partial_cmp(&coeffs[a].norm()) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + const CHUNK_SIZE: usize = 4096; + let room = max_basis.saturating_sub(basis.len()); + let mut merged: FxHashMap> = FxHashMap::default(); + for chunk_indices in order.chunks(CHUNK_SIZE) { + let local: Vec)>> = chunk_indices + .par_iter() + .map_init( + || { + ( + Vec::::with_capacity(spec.n_qubits()), + Vec::::with_capacity(128), + FxHashMap::>::with_capacity_and_hasher( + 128, + FxBuildHasher::default(), + ), + ) + }, + |(s1, s2, lm), &i| { + let r = &basis[i]; + let c_r = coeffs[i]; + let terms = spec.compute_action_terms(r, s1, s2, lm); + let mut out = Vec::with_capacity(terms.len()); + for (q, v) in terms.iter() { + let (r_q, cnt_q) = group.canonicalize_with_shift(q); + if !in_basis.contains_key(&r_q) && !protected_set.contains_key(&r_q) { + let phase = group.character(k_modes, &cnt_q); + out.push((r_q, phase * *v * c_r)); + } + } + out + }, + ) + .collect(); + for v in local { + for (k, val) in v { + *merged.entry(k).or_insert(Complex::new(0.0, 0.0)) += val; + } + } + + // Room-cap: keep only the `room` largest-magnitude entries. + if merged.len() > room { + if room == 0 { + merged.clear(); + } else { + let mut mags: Vec = merged.values().map(|v| v.norm()).collect(); + let k = room.min(mags.len() - 1); + mags.select_nth_unstable_by(k, |a, b| { + b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal) + }); + let cutoff = mags[k]; + merged.retain(|_, &mut v| v.norm() >= cutoff); + } + } + } + Ok(merged.into_iter().filter(|(_, c)| c.norm() > 0.0).collect()) +} + +/// Per-step orbit-rep predictor-corrector evolution. +/// +/// All state lives in orbit-rep form throughout. Each pc step does: +/// 1. Phase-aware leakage from `(basis, coeffs)`; append the largest +/// leakage reps, up to the admission room. +/// 2. Predictor: cached-action expm ([`expm_apply_orbit_rep_cached`]). +/// 3. Phase-aware leakage from the predicted state; append further reps. +/// 4. Corrector: cached-action expm from the pre-step coefficients on +/// the doubly-enlarged basis. +/// 5. Prune `|c| < drop_tol`, then trim to the top-`max_basis` reps by +/// `|c|`; protected reps never dropped. +/// +/// `max_basis` is a hard rank cap on the live orbit-rep basis: enrichment +/// adds at most `max_basis − basis.len()` of the largest leakage reps, the +/// leakage map is capped to the same room, and the post-step basis is +/// trimmed to the top-`max_basis` by `|c|`. Pass a large value (e.g. +/// `usize::MAX`) for the near-exact, uncapped case. `drop_tol` additionally +/// prunes by magnitude. +/// +/// `basis` is assumed to contain only canonical orbit representatives. +/// If not, [`canonicalize_basis_to_rep`] should be called first. +#[allow(clippy::too_many_arguments)] +pub fn pc_step_orbit_rep( + spec: &LindbladSpec, + basis: &mut Vec, + coeffs: &mut Vec>, + dt: f64, + protected: &[Word], + group: &TranslationGroup, + k_modes: &[i32], + cfg: &crate::PcStepConfig, +) -> Result<(), Error> { + let crate::PcStepConfig { + max_basis, + admit_basis, + drop_tol, + tau_add, + .. + } = *cfg; + // Admission bound, mirroring the real-space `pc_step`: enrichment may + // grow the live basis to `admit` >= `max_basis`; the final + // `cap_basis_complex` keeps the top-`max_basis` reps by evolved |coeff| + // over the whole union (rank displacement). With `admit_basis = None` + // admission is bounded by `max_basis` itself 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); + // 1. First-hop phase-aware leakage. + let mut leak = leakage_orbit_rep(spec, basis, coeffs, protected, group, k_modes, admit)?; + if tau_add > 0.0 { + leak.retain(|(_, c)| c.norm() > tau_add); + } + add_leakage_capped_complex(basis, coeffs, leak, admit); + // 2. Predictor: cached-action expm (the phase-aware action is built + // once via `build_orbit_rep_cols`). + let coeffs_predict = expm_apply_orbit_rep_cached(spec, basis, group, k_modes, dt, coeffs); + // 3. Second-hop leakage from predicted state. + let mut leak2 = leakage_orbit_rep( + spec, + basis, + &coeffs_predict, + protected, + group, + k_modes, + admit, + )?; + drop(coeffs_predict); + if tau_add > 0.0 { + leak2.retain(|(_, c)| c.norm() > tau_add); + } + add_leakage_capped_complex(basis, coeffs, leak2, admit); + // 4. Corrector: cache-the-action expm from pre-step state (basis grew). + *coeffs = expm_apply_orbit_rep_cached(spec, basis, group, k_modes, dt, coeffs); + // 5. Prune by magnitude, then rank-cap to max_basis. + if drop_tol > 0.0 { + prune_basis_complex_local(basis, coeffs, drop_tol, protected); + } + cap_basis_complex(basis, coeffs, max_basis, protected); + Ok(()) +} + +/// Complex analogue of `crate::add_leakage_capped`: add the largest leakage +/// reps to the basis, up to the available room `room = max_basis − +/// basis.len()`, so the in-step orbit-rep basis never exceeds `max_basis`. +/// New reps get coefficient 0; the surrounding expm fills them. No +/// magnitude filter — the top-`room` by `|leakage|` are added. +fn add_leakage_capped_complex( + basis: &mut Vec, + coeffs: &mut Vec>, + mut leak: Vec<(Word, Complex)>, + max_basis: usize, +) { + let room = max_basis.saturating_sub(basis.len()); + if leak.len() > room { + if room > 0 { + leak.select_nth_unstable_by(room - 1, |a, b| { + b.1.norm() + .partial_cmp(&a.1.norm()) + .unwrap_or(std::cmp::Ordering::Equal) + }); + } + leak.truncate(room); + } + for (w, _) in leak { + basis.push(w); + coeffs.push(Complex::new(0.0, 0.0)); + } +} + +/// Complex analogue of `crate::cap_basis`: keep only the `max_basis` +/// largest-`|c|` reps (protected reps always kept), dropping the rest. +/// A `max_basis` large enough to cover the whole basis is a no-op. +fn cap_basis_complex( + basis: &mut Vec, + coeffs: &mut Vec>, + max_basis: usize, + protected: &[Word], +) { + if basis.len() <= max_basis { + return; + } + let protected_set: fxhash::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.norm()) + .collect(); + let cutoff = if slots == 0 { + f64::INFINITY + } else if slots >= mags.len() { + return; + } else { + let k = slots - 1; + mags.select_nth_unstable_by(k, |a, b| { + b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal) + }); + mags[k] + }; + let mut write = 0; + for read in 0..basis.len() { + if protected_set.contains(&basis[read]) || coeffs[read].norm() >= cutoff { + if write != read { + basis.swap(write, read); + coeffs.swap(write, read); + } + write += 1; + } + } + basis.truncate(write); + coeffs.truncate(write); +} + +/// Complex analogue of `crate::prune_basis`: drop reps with `|c| < +/// drop_tol`, never dropping `protected` reps. No-op when `drop_tol <= 0`. +fn prune_basis_complex_local( + basis: &mut Vec, + coeffs: &mut Vec>, + drop_tol: f64, + protected: &[Word], +) { + if drop_tol <= 0.0 { + return; + } + let protected_set: fxhash::FxHashSet<&Word> = protected.iter().collect(); + let mut write = 0; + for read in 0..basis.len() { + if coeffs[read].norm() >= drop_tol || protected_set.contains(&basis[read]) { + if write != read { + basis.swap(write, read); + coeffs.swap(write, read); + } + write += 1; + } + } + basis.truncate(write); + coeffs.truncate(write); +} diff --git a/crates/ppvm-python-native/src/interface.rs b/crates/ppvm-python-native/src/interface.rs index 44a47b83..54813f13 100644 --- a/crates/ppvm-python-native/src/interface.rs +++ b/crates/ppvm-python-native/src/interface.rs @@ -48,6 +48,136 @@ macro_rules! create_interface_loss_methods { }; } +macro_rules! create_interface_symmetry_methods { + // Skip loss variants: LossyPauliWord canonicalization would need + // simultaneous permutation of the loss bitmap, which we don't + // implement here. + ($name: ident, $type: ident, true) => {}; + ($name: ident, $type: ident, false) => { + #[pymethods] + impl $name { + /// Symmetry-merge this PauliSum in place: replace every + /// Pauli word by its canonical orbit representative under + /// `group`, accumulating coefficients on collision. Reduces + /// entry count by up to `|group|×` for translation-invariant + /// operators. + /// + /// See `ppvm._core.TranslationGroup` for constructors + /// (`chain_1d`, `torus_2d`, `torus_3d`, `ladder`). + /// + /// Plain real-coefficient merge (the `k=0` symmetry sector). + /// For non-trivial momentum sectors use `momentum_merge`. + pub fn symmetry_merge( + &mut self, + group: &crate::symmetry::TranslationGroup, + ) -> pyo3::PyResult<()> { + if self.inner.n_qubits() != group.core().n_qubits() { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "PauliSum has {} qubits but the TranslationGroup acts on {}", + self.inner.n_qubits(), + group.core().n_qubits(), + ))); + } + ppvm_pauli_sum::symmetry::symmetry_merge_pauli_sum(&mut self.inner, group.core()); + Ok(()) + } + + /// Phase-aware (momentum-sector) merge for a complex operator + /// carried as a *real pair*: `self` is the real part, `other` + /// the imaginary part of `O = self + i·other`. Both are + /// overwritten in place with the orbit-representative form + /// projected onto momentum sector `momentum` (one integer mode + /// per group generator; `[0,…]` is the trivial sector and + /// reduces to `symmetry_merge`). This generalizes + /// `symmetry_merge` to k != 0 while keeping real coefficients on + /// the Python side — the only place complex arithmetic appears + /// is the internal character-weighted fold, reusing the tested + /// `canonicalize_pauli_sum_complex`. + /// + /// `self` and `other` must be distinct objects with identical + /// qubit count. After a translation-covariant gate layer this + /// is exact; under a generic Trotter step it carries the same + /// O(dt^{p+1}) equivariance error as the k=0 merge. + #[pyo3(signature = (other, group, momentum))] + pub fn momentum_merge( + &mut self, + mut other: pyo3::PyRefMut<'_, Self>, + group: &crate::symmetry::TranslationGroup, + momentum: Vec, + ) -> pyo3::PyResult<()> { + let n_g = group.core().n_qubits(); + for (label, n) in [ + ("self", self.inner.n_qubits()), + ("other", other.inner.n_qubits()), + ] { + if n != n_g { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "{label} PauliSum has {n} qubits but the \ + TranslationGroup acts on {n_g}", + ))); + } + } + if momentum.len() != group.core().n_generators() { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "momentum has {} entries but the group has {} generators", + momentum.len(), + group.core().n_generators(), + ))); + } + // Gather both real components into word -> (re + i·im). + let mut combined: std::collections::HashMap< + <$type as Config>::PauliWordType, + num::Complex, + > = std::collections::HashMap::new(); + for (w, v) in self.inner.data().iter() { + combined + .entry(w.clone()) + .or_insert(num::Complex::new(0.0, 0.0)) + .re += *v; + } + for (w, v) in other.inner.data().iter() { + combined + .entry(w.clone()) + .or_insert(num::Complex::new(0.0, 0.0)) + .im += *v; + } + let mut basis = Vec::with_capacity(combined.len()); + let mut coeffs = Vec::with_capacity(combined.len()); + for (w, c) in combined { + basis.push(w); + coeffs.push(c); + } + // Character-weighted fold onto orbit reps. + // `canonicalize_pauli_sum_complex` carries a 1/|G| prefactor; + // we rescale by |G| so the merge is the *summing* projector + // (like `symmetry_merge`): idempotent on already-merged input, + // hence stable under merging after every Trotter step. + ppvm_pauli_sum::symmetry::canonicalize_pauli_sum_complex( + &mut basis, + &mut coeffs, + group.core(), + &momentum, + ); + let scale = group.core().order() as f64; + // Write the real/imag parts back into the two sums. + self.inner.data_mut().clear(); + other.inner.data_mut().clear(); + for (w, c) in basis.into_iter().zip(coeffs.into_iter()) { + let re = c.re * scale; + let im = c.im * scale; + if re != 0.0 { + self.inner += (w.clone(), re); + } + if im != 0.0 { + other.inner += (w, im); + } + } + Ok(()) + } + } + }; +} + macro_rules! create_strategy { (false, $min_abs_coeff:ident, $max_pauli_weight:ident, $_max_loss_weight:ident) => { CombinedStrategy( @@ -403,6 +533,12 @@ macro_rules! create_interface { } } + // `symmetry_merge` only makes sense on non-loss variants — the + // canonicalization permutes qubit positions and the loss + // bitmap would need a parallel permutation that we don't + // attempt here. + create_interface_symmetry_methods!($name, $type, $loss); + create_interface_loss_methods!($name, $type, $loss); }; } diff --git a/crates/ppvm-python-native/src/lib.rs b/crates/ppvm-python-native/src/lib.rs index 0bcf19d8..5ebc9532 100644 --- a/crates/ppvm-python-native/src/lib.rs +++ b/crates/ppvm-python-native/src/lib.rs @@ -16,6 +16,7 @@ pub mod interface_tableau; pub mod interface_tableau_sum; pub mod lindblad; pub mod stim_program; +pub mod symmetry; pub(crate) fn flat_pairs(targets: &[usize]) -> PyResult> { if !targets.len().is_multiple_of(2) { @@ -308,4 +309,14 @@ pub mod _core { // Lindbladian time evolution #[pymodule_export] pub use crate::lindblad::LindbladSpec; + + // Symmetry merging + #[pymodule_export] + pub use crate::symmetry::TranslationGroup; + #[pymodule_export] + pub use crate::symmetry::canonicalize_basis_arr; + #[pymodule_export] + pub use crate::symmetry::canonicalize_basis_arr_complex; + #[pymodule_export] + pub use crate::symmetry::check_momentum_sector_arr; } diff --git a/crates/ppvm-python-native/src/lindblad.rs b/crates/ppvm-python-native/src/lindblad.rs index 25d4e0b0..7c720620 100644 --- a/crates/ppvm-python-native/src/lindblad.rs +++ b/crates/ppvm-python-native/src/lindblad.rs @@ -13,11 +13,14 @@ use std::collections::HashMap; use num::Complex; -use numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2}; +use numpy::{ + Complex64, IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2, +}; use ppvm_lindblad::{JumpInput, LindbladSpec as CoreSpec, Word, codes_from_word, word_from_codes}; use pyo3::{exceptions::PyValueError, prelude::*}; type PyPauliMap<'py> = (Bound<'py, PyArray2>, Bound<'py, PyArray1>); +type PyPauliMapComplex<'py> = (Bound<'py, PyArray2>, Bound<'py, PyArray1>); type PyCoo<'py> = ( Bound<'py, PyArray1>, Bound<'py, PyArray1>, @@ -341,6 +344,120 @@ impl LindbladSpec { Ok((map, d)) } + /// Per-step orbit-rep predictor-corrector evolution under + /// translation symmetry. State lives entirely in **orbit-rep form**: + /// basis contains only canonical orbit representatives, coefficients + /// are complex. The action is phase-aware: output Paulis canonicalize + /// to their orbit rep with momentum-character weight. + /// + /// Per-step memory benefit: basis is ~|group|× smaller than the + /// full-basis representation, and the reduction persists through + /// every step. + /// + /// **Pre-condition**: every row of `basis` must be the canonical + /// orbit representative of its translation orbit under `group`. + /// Pass `canonicalize_first=True` to enforce this on entry (rewrites + /// each basis row to its canonical rep; coefficients unchanged). + /// Default `False` — the caller is trusted. + /// + /// `max_basis` is a hard rank cap on the live orbit-rep basis: + /// enrichment adds at most `max_basis − basis.len()` of the largest + /// leakage reps and the post-step basis is trimmed to the top-`max_basis` + /// by `|c|` (protected reps always kept). Pass a large value for the + /// near-exact case. `drop_tol` additionally prunes by magnitude. + #[pyo3(signature = ( + basis, coeffs, dt, max_basis, + group, momentum, + drop_tol = 0.0, + protected = None, + canonicalize_first = false, + admit_basis = None, + tau_add = None, + ))] + #[allow(clippy::too_many_arguments)] + fn pc_step_orbit_rep<'py>( + &self, + py: Python<'py>, + basis: PyReadonlyArray2<'py, u8>, + coeffs: PyReadonlyArray1<'py, Complex64>, + dt: f64, + max_basis: usize, + group: &crate::symmetry::TranslationGroup, + momentum: PyReadonlyArray1<'py, i32>, + drop_tol: f64, + protected: Option>, + canonicalize_first: bool, + admit_basis: Option, + tau_add: Option, + ) -> PyResult> { + use num::Complex; + use ppvm_lindblad::orbit_rep; + + let n_q = self.inner.n_qubits(); + let basis_view = basis.as_array(); + let mut basis_words = decode_basis(&basis_view, n_q)?; + let coeffs_slice = coeffs.as_slice()?; + if coeffs_slice.len() != basis_words.len() { + return Err(PyValueError::new_err(format!( + "coeffs has length {} but basis has {} rows", + coeffs_slice.len(), + basis_words.len() + ))); + } + let mut coeffs_vec: Vec> = coeffs_slice + .iter() + .map(|c| Complex::new(c.re, c.im)) + .collect(); + let protected_words: Vec = if let Some(ref p) = protected { + decode_basis(&p.as_array(), n_q)? + } else { + Vec::new() + }; + let k_slice = momentum.as_slice()?; + if k_slice.len() != group.core().n_generators() { + return Err(PyValueError::new_err(format!( + "momentum has {} entries but group has {} generators", + k_slice.len(), + group.core().n_generators() + ))); + } + if canonicalize_first { + orbit_rep::canonicalize_basis_to_rep(&mut basis_words, group.core()); + } + orbit_rep::pc_step_orbit_rep( + &self.inner, + &mut basis_words, + &mut coeffs_vec, + dt, + &protected_words, + group.core(), + k_slice, + &ppvm_lindblad::PcStepConfig { + max_basis, + admit_basis, + drop_tol, + tau_add, + num_threads: None, + }, + ) + .map_err(map_err)?; + + let m = basis_words.len(); + let mut out_basis = vec![0u8; m * n_q]; + for (i, w) in basis_words.iter().enumerate() { + codes_from_word(w, &mut out_basis[i * n_q..(i + 1) * n_q]); + } + let out_coeffs: Vec = coeffs_vec + .iter() + .map(|c| Complex64::new(c.re, c.im)) + .collect(); + let basis_arr = out_basis + .into_pyarray(py) + .reshape([m, n_q]) + .map_err(|e| PyValueError::new_err(format!("reshape failed: {e}")))?; + Ok((basis_arr, out_coeffs.into_pyarray(py))) + } + /// Sparse generator matrix in COO form: `(rows, cols, vals)`. fn generator<'py>( &self, diff --git a/crates/ppvm-python-native/src/symmetry.rs b/crates/ppvm-python-native/src/symmetry.rs new file mode 100644 index 00000000..88085146 --- /dev/null +++ b/crates/ppvm-python-native/src/symmetry.rs @@ -0,0 +1,329 @@ +// SPDX-FileCopyrightText: 2026 The PPVM Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Python bindings for the symmetry-merging primitive. +//! +//! Exposes: +//! - [`TranslationGroup`] PyO3 class with constructors for 1D, 2D, 3D +//! tori and multi-leg ladders, plus a generic generator-list path. +//! - [`canonicalize_basis_arr`] / [`canonicalize_basis_arr_complex`] free +//! functions that merge the numpy `(basis_arr, coeffs)` representation +//! used by `Lindbladian.pc_step_arr`. + +use num::Complex; +use numpy::{ + Complex64, IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2, +}; +use ppvm_lindblad::{codes_from_word, word_from_codes}; +use ppvm_pauli_sum::symmetry as core_sym; +use pyo3::{exceptions::PyValueError, prelude::*}; + +type PyPauliMap<'py> = (Bound<'py, PyArray2>, Bound<'py, PyArray1>); +type PyPauliMapComplex<'py> = (Bound<'py, PyArray2>, Bound<'py, PyArray1>); + +/// A finite abelian symmetry group acting on qubit positions by +/// permutations. Use this to merge translation-equivalent Pauli strings +/// in either the `Lindbladian.pc_step_arr` basis or the `PauliSum` +/// dictionary, reducing per-step memory by up to `|G|×`. +/// +/// Build via the static methods: +/// - `TranslationGroup.chain_1d(n)` — 1D chain of `n` sites with PBC. +/// - `TranslationGroup.torus_2d(lx, ly)` — 2D torus; qubit `(i, j)` at +/// index `j*lx + i`. +/// - `TranslationGroup.torus_3d(lx, ly, lz)` — 3D torus; qubit +/// `(i, j, k)` at index `k*lx*ly + j*lx + i`. +/// - `TranslationGroup.ladder(l, n_legs)` — `n_legs`-leg ladder of `l` +/// sites, translation only along chain direction; qubit `(leg, j)` at +/// index `leg*l + j`. +/// - `TranslationGroup.from_generators(n_qubits, perms, orders)` — +/// arbitrary list of generator permutations + cyclic orders. +#[pyclass(frozen)] +pub struct TranslationGroup { + pub(crate) inner: core_sym::TranslationGroup, +} + +impl TranslationGroup { + /// Accessor for the underlying [`ppvm_pauli_sum::symmetry::TranslationGroup`]. + /// Used by other crate-internal modules (e.g. the PauliSum interface + /// macro) to call into the core merging API. + pub fn core(&self) -> &core_sym::TranslationGroup { + &self.inner + } +} + +#[pymethods] +impl TranslationGroup { + #[staticmethod] + pub fn chain_1d(n: usize) -> Self { + Self { + inner: core_sym::TranslationGroup::chain_1d(n), + } + } + + #[staticmethod] + pub fn torus_2d(lx: usize, ly: usize) -> Self { + Self { + inner: core_sym::TranslationGroup::torus_2d(lx, ly), + } + } + + #[staticmethod] + pub fn torus_3d(lx: usize, ly: usize, lz: usize) -> Self { + Self { + inner: core_sym::TranslationGroup::torus_3d(lx, ly, lz), + } + } + + #[staticmethod] + pub fn ladder(l: usize, n_legs: usize) -> Self { + Self { + inner: core_sym::TranslationGroup::ladder(l, n_legs), + } + } + + #[staticmethod] + pub fn from_generators( + n_qubits: usize, + perms: Vec>, + orders: Vec, + ) -> PyResult { + if perms.len() != orders.len() { + return Err(PyValueError::new_err(format!( + "perms ({} generators) and orders ({}) must have the same length", + perms.len(), + orders.len() + ))); + } + for (g, perm) in perms.iter().enumerate() { + if perm.len() != n_qubits { + return Err(PyValueError::new_err(format!( + "generator {g}: permutation length {} != n_qubits {n_qubits}", + perm.len() + ))); + } + let mut seen = vec![false; n_qubits]; + for &p in perm { + let p = p as usize; + if p >= n_qubits { + return Err(PyValueError::new_err(format!( + "generator {g}: target {p} out of range [0, {n_qubits})" + ))); + } + if seen[p] { + return Err(PyValueError::new_err(format!( + "generator {g}: not a permutation (duplicate target {p})" + ))); + } + seen[p] = true; + } + } + Ok(Self { + inner: core_sym::TranslationGroup::from_generators(n_qubits, perms, orders), + }) + } + + /// Number of qubits this group acts on. + #[getter] + pub fn n_qubits(&self) -> usize { + self.inner.n_qubits() + } + + /// Number of generators (rank as an abelian product group). + #[getter] + pub fn n_generators(&self) -> usize { + self.inner.n_generators() + } + + /// Total group order: product of generator orders. + #[getter] + pub fn order(&self) -> usize { + self.inner.order() + } + + /// Return the canonical (lex-min) orbit representative of `pauli`. + /// `pauli` is a length-`n_qubits` uint8 array with the encoding + /// `0=I, 1=X, 2=Z, 3=Y`. Result is the same shape. + pub fn canonicalize<'py>( + &self, + py: Python<'py>, + pauli: PyReadonlyArray1<'py, u8>, + ) -> PyResult>> { + let codes = pauli.as_slice()?; + if codes.len() != self.inner.n_qubits() { + return Err(PyValueError::new_err(format!( + "pauli has length {} but group expects {} qubits", + codes.len(), + self.inner.n_qubits() + ))); + } + let w = word_from_codes(codes).map_err(|e| PyValueError::new_err(e.to_string()))?; + let canon = self.inner.canonicalize(&w); + let mut out = vec![0u8; codes.len()]; + codes_from_word(&canon, &mut out); + Ok(out.into_pyarray(py)) + } +} + +/// Phase-aware merge of a complex-coefficient `(basis_arr, coeffs)` +/// Pauli sum into orbit-rep form, projected onto momentum sector +/// `momentum`. +/// +/// `momentum` is a length-`group.n_generators` integer array of mode +/// indices; the wavenumber along generator `g` is +/// `2π · momentum[g] / group.generator_order(g)`. Use `momentum=[0, …]` +/// for the trivial (k=0) sector — equivalent to plain merging modulo +/// the 1/|G| normalization the complex merge applies. +/// +/// If the input is **not** in sector `momentum`, the projection +/// silently throws away the other components. Use +/// [`check_momentum_sector_arr`] beforehand to validate. +#[pyfunction] +pub fn canonicalize_basis_arr_complex<'py>( + py: Python<'py>, + basis: PyReadonlyArray2<'py, u8>, + coeffs: PyReadonlyArray1<'py, Complex64>, + group: &TranslationGroup, + momentum: PyReadonlyArray1<'py, i32>, +) -> PyResult> { + let basis_view = basis.as_array(); + let n_q = group.inner.n_qubits(); + if basis_view.shape().get(1).copied() != Some(n_q) { + return Err(PyValueError::new_err(format!( + "basis has {} qubits per row but group acts on {n_q}", + basis_view.shape().get(1).copied().unwrap_or(0) + ))); + } + let n = basis_view.shape()[0]; + let coeffs_slice = coeffs.as_slice()?; + if coeffs_slice.len() != n { + return Err(PyValueError::new_err(format!( + "coeffs has length {} but basis has {} rows", + coeffs_slice.len(), + n + ))); + } + let k_slice = momentum.as_slice()?; + if k_slice.len() != group.inner.n_generators() { + return Err(PyValueError::new_err(format!( + "momentum has {} entries but group has {} generators", + k_slice.len(), + group.inner.n_generators() + ))); + } + let mut basis_words = crate::lindblad::decode_basis(&basis_view, n_q)?; + let mut coeffs_vec: Vec> = coeffs_slice + .iter() + .map(|c| Complex::new(c.re, c.im)) + .collect(); + + core_sym::canonicalize_pauli_sum_complex( + &mut basis_words, + &mut coeffs_vec, + &group.inner, + k_slice, + ); + + let m = basis_words.len(); + let mut out_basis = vec![0u8; m * n_q]; + for (i, w) in basis_words.iter().enumerate() { + codes_from_word(w, &mut out_basis[i * n_q..(i + 1) * n_q]); + } + let out_coeffs: Vec = coeffs_vec + .iter() + .map(|c| Complex64::new(c.re, c.im)) + .collect(); + let basis_arr = out_basis + .into_pyarray(py) + .reshape([m, n_q]) + .map_err(|e| PyValueError::new_err(format!("reshape failed: {e}")))?; + Ok((basis_arr, out_coeffs.into_pyarray(py))) +} + +/// Verify that a `(basis_arr, complex_coeffs)` Pauli sum lies in the +/// momentum sector `momentum` under `group`. Returns `None` on pass, +/// raises a `ValueError` with diagnostic info on fail. +/// +/// `tol` is the relative tolerance on coefficient comparison; default +/// `1e-8`. +#[pyfunction] +#[pyo3(signature = (basis, coeffs, group, momentum, tol = 1e-8))] +pub fn check_momentum_sector_arr<'py>( + basis: PyReadonlyArray2<'py, u8>, + coeffs: PyReadonlyArray1<'py, Complex64>, + group: &TranslationGroup, + momentum: PyReadonlyArray1<'py, i32>, + tol: f64, +) -> PyResult<()> { + let basis_view = basis.as_array(); + let n_q = group.inner.n_qubits(); + if basis_view.shape().get(1).copied() != Some(n_q) { + return Err(PyValueError::new_err(format!( + "basis has {} qubits per row but group acts on {n_q}", + basis_view.shape().get(1).copied().unwrap_or(0) + ))); + } + let coeffs_slice = coeffs.as_slice()?; + let k_slice = momentum.as_slice()?; + let basis_words = crate::lindblad::decode_basis(&basis_view, n_q)?; + let coeffs_vec: Vec> = coeffs_slice + .iter() + .map(|c| Complex::new(c.re, c.im)) + .collect(); + core_sym::check_momentum_sector(&basis_words, &coeffs_vec, &group.inner, k_slice, tol) + .map_err(|e| PyValueError::new_err(format!("{e}"))) +} + +/// Merge a `(basis_arr, coeffs)` Pauli sum (the representation used by +/// `Lindbladian.pc_step_arr`) into orbit-representative form. +/// Each row of `basis_arr` is replaced by its canonical +/// representative; coefficients of rows collapsing to the same rep are +/// summed. +/// +/// Returns `(merged_basis_arr, merged_coeffs)`. Output length ≤ input +/// length. +/// +/// For dynamics that commute with `group` and initial states that are +/// `group`-invariant, this preserves all `group`-invariant expectation +/// values (Theorem 1 of Teng et al., arXiv:2512.12094). +#[pyfunction] +pub fn canonicalize_basis_arr<'py>( + py: Python<'py>, + basis: PyReadonlyArray2<'py, u8>, + coeffs: PyReadonlyArray1<'py, f64>, + group: &TranslationGroup, +) -> PyResult> { + let basis_view = basis.as_array(); + let n_q = group.inner.n_qubits(); + if basis_view.shape().get(1).copied() != Some(n_q) { + return Err(PyValueError::new_err(format!( + "basis has {} qubits per row but group acts on {n_q}", + basis_view.shape().get(1).copied().unwrap_or(0) + ))); + } + let n = basis_view.shape()[0]; + let coeffs_slice = coeffs.as_slice()?; + if coeffs_slice.len() != n { + return Err(PyValueError::new_err(format!( + "coeffs has length {} but basis has {} rows", + coeffs_slice.len(), + n + ))); + } + + let mut basis_words = crate::lindblad::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); + + // Re-encode. + let m = basis_words.len(); + let mut out_basis = vec![0u8; m * n_q]; + for (i, w) in basis_words.iter().enumerate() { + codes_from_word(w, &mut out_basis[i * n_q..(i + 1) * n_q]); + } + let basis_arr = out_basis + .into_pyarray(py) + .reshape([m, n_q]) + .map_err(|e| PyValueError::new_err(format!("reshape failed: {e}")))?; + Ok((basis_arr, coeffs_vec.into_pyarray(py))) +} diff --git a/ppvm-python/src/ppvm/_core.pyi b/ppvm-python/src/ppvm/_core.pyi index 4b2cfac7..b0350fc5 100644 --- a/ppvm-python/src/ppvm/_core.pyi +++ b/ppvm-python/src/ppvm/_core.pyi @@ -62,6 +62,14 @@ class _PauliSumBase: def terms(self) -> list[tuple[str, float]]: ... def weights(self) -> list[tuple[str, int]]: ... def current_max_weight(self) -> int: ... + # Only on non-loss variants (see create_interface_symmetry_methods). + def symmetry_merge(self, group: TranslationGroup) -> None: ... + def momentum_merge( + self, + other: _PauliSumBase, + group: TranslationGroup, + momentum: list[int], + ) -> None: ... class _PauliSumLossBase(_PauliSumBase): def loss_channel(self, addr0: int, p: float, truncate: bool = True) -> None: ... @@ -392,4 +400,56 @@ class LindbladSpec: admit_basis: int | None = None, tau_add: float | None = None, ) -> tuple[tuple[np.ndarray, np.ndarray], dict[str, int]]: ... + def pc_step_orbit_rep( + self, + basis: np.ndarray, + coeffs: np.ndarray, + dt: float, + max_basis: int, + group: TranslationGroup, + momentum: np.ndarray, + drop_tol: float = 0.0, + protected: np.ndarray | None = None, + canonicalize_first: bool = False, + admit_basis: int | None = None, + tau_add: float | None = None, + ) -> tuple[np.ndarray, np.ndarray]: ... def generator(self, basis: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... + +class TranslationGroup: + @staticmethod + def chain_1d(n: int) -> TranslationGroup: ... + @staticmethod + def torus_2d(lx: int, ly: int) -> TranslationGroup: ... + @staticmethod + def torus_3d(lx: int, ly: int, lz: int) -> TranslationGroup: ... + @staticmethod + def ladder(l: int, n_legs: int) -> TranslationGroup: ... + @staticmethod + def from_generators( + n_qubits: int, perms: list[list[int]], orders: list[int] + ) -> TranslationGroup: ... + @property + def n_qubits(self) -> int: ... + @property + def n_generators(self) -> int: ... + @property + def order(self) -> int: ... + def canonicalize(self, pauli: np.ndarray) -> np.ndarray: ... + +def canonicalize_basis_arr( + basis: np.ndarray, coeffs: np.ndarray, group: TranslationGroup +) -> tuple[np.ndarray, np.ndarray]: ... +def canonicalize_basis_arr_complex( + basis: np.ndarray, + coeffs: np.ndarray, + group: TranslationGroup, + momentum: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: ... +def check_momentum_sector_arr( + basis: np.ndarray, + coeffs: np.ndarray, + group: TranslationGroup, + momentum: np.ndarray, + tol: float = 1e-8, +) -> None: ... diff --git a/ppvm-python/src/ppvm/lindblad.py b/ppvm-python/src/ppvm/lindblad.py index 489cf3b2..25048112 100644 --- a/ppvm-python/src/ppvm/lindblad.py +++ b/ppvm-python/src/ppvm/lindblad.py @@ -9,7 +9,8 @@ adaptive Heisenberg-picture evolution: - ``pc_step(...)`` / ``pc_step_arr(...)``: one adaptive predictor-corrector - step ``O ← exp(dt·L*) O`` + step ``O ← exp(dt·L*) O``; ``pc_step_orbit_rep(...)`` is the + translation-symmetric (momentum-sector) variant - ``action(p)`` / ``action_arr(p)``: L*(p) for one Pauli string p - ``leakage(basis, coeffs)`` / ``leakage_arr(...)``: off-basis component of L*(Σ c_j p_j), driving basis expansion @@ -290,6 +291,63 @@ def pc_step_arr( None if tau_add is None else float(tau_add), ) + def pc_step_orbit_rep( + self, + basis_arr: np.ndarray, + coeffs: np.ndarray, + dt: float, + max_basis: int, + group, + momentum: np.ndarray, + drop_tol: float = 1e-12, + protected_arr: np.ndarray | None = None, + canonicalize_first: bool = False, + admit_basis: int | None = None, + tau_add: float | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Per-step orbit-representative pc evolution. + + State lives entirely in orbit-rep form throughout: ``basis_arr`` + contains only canonical translation-orbit representatives, + ``coeffs`` are complex, and the action is phase-aware. The basis + is ~``|group|×`` smaller than the equivalent full-basis complex + evolution, and the reduction persists across every step. + + Truncation. ``max_basis`` is a hard rank cap on the live orbit-rep + basis: enrichment adds at most ``max_basis - len(basis)`` of the + largest leakage reps, and the post-step basis is trimmed to the + top-``max_basis`` reps by ``|c|`` (``protected`` reps always kept). + Pass a large value (e.g. ``10_000_000``) for the near-exact, + uncapped case. ``drop_tol`` additionally prunes reps whose absolute + coefficient is below the threshold after the corrector. + + ``admit_basis``, when set (>= ``max_basis``), bounds the enriched + working set instead of ``max_basis``: the step may hold up to + ``admit_basis`` reps transiently and the final truncation keeps the + top-``max_basis`` by evolved ``|c|`` over the whole union — the + displacement scheme, matching the real-space ``pc_step_arr``. + + ``basis_arr`` is assumed to contain canonical reps only. Pass + ``canonicalize_first=True`` to rewrite each row to its canonical + rep on entry (coefficients unchanged). + """ + n = self.n_qubits + if protected_arr is None: + protected_arr = np.zeros((0, n), dtype=np.uint8) + return self._spec.pc_step_orbit_rep( + np.ascontiguousarray(basis_arr, dtype=np.uint8), + np.ascontiguousarray(coeffs, dtype=np.complex128), + float(dt), + int(max_basis), + group, + np.ascontiguousarray(momentum, dtype=np.int32), + float(drop_tol), + np.ascontiguousarray(protected_arr, dtype=np.uint8), + bool(canonicalize_first), + None if admit_basis is None else int(admit_basis), + None if tau_add is None else float(tau_add), + ) + def pc_step( self, basis: Sequence[str], diff --git a/ppvm-python/src/ppvm/paulisum.py b/ppvm-python/src/ppvm/paulisum.py index dca573d8..b367a724 100644 --- a/ppvm-python/src/ppvm/paulisum.py +++ b/ppvm-python/src/ppvm/paulisum.py @@ -386,6 +386,50 @@ def trace(self, pattern: str) -> float: """ return self._interface.trace(pattern) + def symmetry_merge(self, group) -> None: + """Merge entries into orbit-representative form under a translation group. + + Each Pauli word in the sum is replaced by its canonical (lex-min) + representative under the action of ``group``; coefficients of words + that collapse to the same representative are summed. Entry count + reduces by up to ``|group|×`` for translation-invariant operators. + + For a translation-invariant dynamics that you apply between + merging steps, this preserves all ``group``-invariant expectation + values (Theorem 1 of Teng et al., arXiv:2512.12094). Plain + real-coefficient merge — handles the trivial (``k=0``) momentum + sector only. + + Args: + group: A `ppvm._core.TranslationGroup` + (use ``TranslationGroup.chain_1d(n)``, ``.torus_2d``, + ``.torus_3d``, ``.ladder``, or ``.from_generators``). + """ + self._interface.symmetry_merge(group) + + def momentum_merge(self, other: "PauliSum", group, momentum) -> None: + """Phase-aware (momentum-sector) merge for a complex operator stored + as a *real pair*: ``self`` is the real part and ``other`` the + imaginary part of ``O = self + i·other``. Both are overwritten in + place with the orbit-representative form projected onto momentum + sector ``momentum``. + + Generalizes `symmetry_merge` to non-trivial momentum sectors + (``k != 0``) while keeping real coefficients on both PauliSums — the + only complex arithmetic is the internal character-weighted fold. + ``self`` and ``other`` must be distinct objects with the same qubit + count. Exact after a translation-covariant gate layer; under a + generic Trotter step it carries the same ``O(dt^{p+1})`` equivariance + error as the ``k=0`` merge. + + Args: + other: the PauliSum holding the imaginary component (modified in place). + group: a `ppvm._core.TranslationGroup`. + momentum: sequence of integer modes, one per group generator + (e.g. ``[k]`` for a 1D chain; ``[0, ...]`` is the trivial sector). + """ + self._interface.momentum_merge(other._interface, group, list(momentum)) + def amplitude_damping(self, addr0: int, gamma: float, *, truncate: bool = True): """Apply an amplitude-damping channel. diff --git a/ppvm-python/test/test_momentum_merge.py b/ppvm-python/test/test_momentum_merge.py new file mode 100644 index 00000000..0546deb1 --- /dev/null +++ b/ppvm-python/test/test_momentum_merge.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: 2026 The PPVM Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for momentum-sector (k != 0) symmetry merging of real PauliSum pairs. + +A complex operator O = O_re + i·O_im is carried as a pair of real PauliSums. +``PauliSum.momentum_merge`` folds the pair onto translation-orbit +representatives in momentum sector k, generalizing ``symmetry_merge`` (k=0). + +These checks compare against *exact* references — the projector definition, +idempotency, and exact diagonalization of the dynamics — NOT against any +other propagation scheme. +""" + +import cmath +import math + +import numpy as np +import pytest + +from ppvm import PauliSum +from ppvm._core import TranslationGroup + +# ── dense Pauli helpers (exact references) ─────────────────────────────────── +_I = np.eye(2, dtype=complex) +_X = np.array([[0, 1], [1, 0]], dtype=complex) +_Y = np.array([[0, -1j], [1j, 0]], dtype=complex) +_Z = np.array([[1, 0], [0, -1]], dtype=complex) +_P = {"I": _I, "X": _X, "Y": _Y, "Z": _Z} + + +def dense(pauli_str): + m = np.array([[1]], dtype=complex) + for ch in pauli_str: + m = np.kron(m, _P[ch]) + return m + + +def zstr(n, q): + return "".join("Z" if i == q else "I" for i in range(n)) + + +def chain_bonds(n): + return [(i, (i + 1) % n, 1.0) for i in range(n)] + + +# ── helpers shared with the k-resolved Trotter driver ──────────────────────── +def _seed_pair(n, k): + a = np.arange(n) + re = np.cos(2 * np.pi * k * a / n) + im = -np.sin(2 * np.pi * k * a / n) # e^{-2πi k a/n} = cos - i sin + Z = [zstr(n, q) for q in range(n)] + PA = PauliSum.new( + n, [(Z[q], float(re[q])) for q in range(n)], min_abs_coeff=0.0, max_pauli_weight=n + ) + PB = PauliSum.new( + n, [(Z[q], float(im[q])) for q in range(n)], min_abs_coeff=0.0, max_pauli_weight=n + ) + return PA, PB + + +def _to_complex_dict(PA, PB): + d = {} + for s, c in PA.terms: + d[s] = d.get(s, 0j) + c + for s, c in PB.terms: + d[s] = d.get(s, 0j) + 1j * c + return {s: v for s, v in d.items() if v != 0j} + + +def _ovl(sA, sB, oA, oB): + re = sA.overlap(oA) + sB.overlap(oB) + im = sA.overlap(oB) - sB.overlap(oA) + return complex(re, im) + + +# ============================================================================= +# 1. The merge is an exact sector projector: idempotent, and it leaves a +# genuine momentum-k eigenoperator unchanged. +# ============================================================================= +@pytest.mark.parametrize("k", [0, 1, 2, 3]) +def test_momentum_merge_idempotent(k): + n = 4 + g = TranslationGroup.chain_1d(n) + PA, PB = _seed_pair(n, k) # S^z_k is exactly in sector k + PA.momentum_merge(PB, g, [k]) + once = _to_complex_dict(PA, PB) + PA.momentum_merge(PB, g, [k]) # merging again must be a no-op + twice = _to_complex_dict(PA, PB) + keys = set(once) | set(twice) + assert max(abs(once.get(x, 0j) - twice.get(x, 0j)) for x in keys) < 1e-12 + + +def test_momentum_merge_projects_out_other_sectors(): + """Merging a pure sector-k operator in sector k' != k gives ~zero.""" + n = 4 + g = TranslationGroup.chain_1d(n) + PA, PB = _seed_pair(n, 1) # operator lives in k=1 + PA.momentum_merge(PB, g, [2]) # project onto k=2 + d = _to_complex_dict(PA, PB) + assert all(abs(v) < 1e-12 for v in d.values()), d + + +# ============================================================================= +# 2. End-to-end: k-resolved, symmetry-compressed Trotter reproduces the +# EXACT (dense-diagonalization) operator autocorrelator as dt -> 0. +# ============================================================================= +def _ed_autocorr(n, bonds, k, ts): + """C_k(t) = Tr[O0^dagger O(t)] / Tr[O0^dagger O0], O0 = S^z_k, exact.""" + H = np.zeros((2**n, 2**n), dtype=complex) + for i, j, J in bonds: + for q in "XY": + s = ["I"] * n + s[i] = q + s[j] = q + H += J * dense("".join(s)) + O0 = np.zeros((2**n, 2**n), dtype=complex) + for a in range(n): + O0 += cmath.exp(-2j * math.pi * k * a / n) * dense(zstr(n, a)) + E, V = np.linalg.eigh(H) + out = [] + with np.errstate(all="ignore"): # silence spurious macOS-Accelerate matmul warnings + norm = np.trace(O0.conj().T @ O0).real + for t in ts: + U = (V * np.exp(-1j * E * t)) @ V.conj().T + Ot = U.conj().T @ O0 @ U + out.append(np.trace(O0.conj().T @ Ot) / norm) + return np.array(out) + + +def _ctrotter_autocorr(n, bonds, k, dt, steps): + g = TranslationGroup.chain_1d(n) + PA, PB = _seed_pair(n, k) + PA.momentum_merge(PB, g, [k]) + refA, refB = PA.copy(), PB.copy() + C0 = _ovl(refA, refB, PA, PB) + out = [1.0 + 0j] + for _ in range(steps): + for i, j, J in bonds: # Strang: forward then reversed + PA.rxx(i, j, J * dt, truncate=False) + PA.ryy(i, j, J * dt, truncate=False) + PB.rxx(i, j, J * dt, truncate=False) + PB.ryy(i, j, J * dt, truncate=False) + for i, j, J in reversed(bonds): + PA.rxx(i, j, J * dt, truncate=False) + PA.ryy(i, j, J * dt, truncate=False) + PB.rxx(i, j, J * dt, truncate=False) + PB.ryy(i, j, J * dt, truncate=False) + PA.momentum_merge(PB, g, [k]) + out.append(_ovl(refA, refB, PA, PB) / C0) + return np.array(out) + + +@pytest.mark.parametrize("k", [0, 1, 2, 3]) +def test_k_resolved_trotter_converges_to_exact(k): + n, T = 4, 0.3 + bonds = chain_bonds(n) + # exact reference at the matching times for two step sizes + err = {} + for dt in (0.04, 0.02): + steps = round(T / dt) + ts = np.arange(steps + 1) * dt + c = _ctrotter_autocorr(n, bonds, k, dt, steps) + ed = _ed_autocorr(n, bonds, k, ts) + err[dt] = np.max(np.abs(c - ed)) + + assert abs(_ctrotter_autocorr(n, bonds, k, 0.02, 1)[0] - 1.0) < 1e-12 # C_k(0)=1 + if k == 0: + # total Z is conserved -> exact in every sector-0 step + assert err[0.02] < 1e-10 + else: + assert err[0.02] < 5e-3 # close to exact at dt=0.02 + assert err[0.02] < err[0.04] # converges toward exact as dt->0 + + +def test_compressed_matches_uncompressed_evolution(): + """Merging must not change observables beyond the O(dt^2) equivariance + error: compressed (merge each step) vs the same gates with no merge.""" + n, k, dt, steps = 4, 2, 0.02, 10 + bonds = chain_bonds(n) + g = TranslationGroup.chain_1d(n) + + # uncompressed: evolve the full real pair, project only at readout + PA, PB = _seed_pair(n, k) + rA, rB = _seed_pair(n, k) + rA.momentum_merge(rB, g, [k]) + C0 = _ovl(rA, rB, *_merged_copy(PA, PB, g, k)) + comp = _ctrotter_autocorr(n, bonds, k, dt, steps) + unc = [] + for _ in range(steps): + for i, j, J in bonds: + PA.rxx(i, j, J * dt, truncate=False) + PA.ryy(i, j, J * dt, truncate=False) + PB.rxx(i, j, J * dt, truncate=False) + PB.ryy(i, j, J * dt, truncate=False) + for i, j, J in reversed(bonds): + PA.rxx(i, j, J * dt, truncate=False) + PA.ryy(i, j, J * dt, truncate=False) + PB.rxx(i, j, J * dt, truncate=False) + PB.ryy(i, j, J * dt, truncate=False) + mA, mB = _merged_copy(PA, PB, g, k) + unc.append(_ovl(rA, rB, mA, mB) / C0) + unc = np.array([1.0 + 0j, *unc]) + assert np.max(np.abs(comp - unc)) < 5e-3 # only O(dt^2) equivariance + + +def _merged_copy(PA, PB, g, k): + a, b = PA.copy(), PB.copy() + a.momentum_merge(b, g, [k]) + return a, b