gam_solve/arrow_schur/

penalty_ops.rs

//! Matrix-free penalty-side `H_ββ` operators: the [`BetaPenaltyOp`] trait and
//! every concrete operator (dense, block, Kronecker, factored-frame, composite,
//! matvec-diagonal) plus the β-coupling graph used for block-Jacobi clustering.

use super::*;

#[derive(Debug, Clone)]
pub(crate) struct BetaEdge {
    pub(crate) a: usize,
    pub(crate) b: usize,
}

#[derive(Debug, Clone)]
pub(crate) struct BetaCouplingGraph {
    pub(crate) num_blocks: usize,
    pub(crate) edges: Vec<BetaEdge>,
    pub(crate) adj_start: Vec<usize>,
    pub(crate) adj_targets: Vec<usize>,
    /// Co-visibility WEIGHT parallel to `adj_targets`: `adj_weights[e]` is the
    /// number of rows in which `node` and `adj_targets[e]` both fire. This is
    /// the co-visibility count that keys the visibility-based cluster partition
    /// (Kushal & Agarwal, CVPR 2012): the reduced Schur is, up to its
    /// block-diagonal penalty, a graph Laplacian over this weighted graph, so a
    /// strong `adj_weights[e]` marks an off-block coupling a good preconditioner
    /// block must keep together. Symmetric: the `(a→b)` and `(b→a)` entries carry
    /// the same weight.
    pub(crate) adj_weights: Vec<f64>,
}

impl BetaCouplingGraph {
    /// Build the β co-firing graph over `sys.block_offsets`, through the
    /// system's ROUTED `H_tβ` convention.
    ///
    /// # Why this takes the system and not the row slabs (#2572)
    ///
    /// This used to take `&[Array2<f64>]` and every one of its five call sites
    /// passed `sys.rows.iter().map(|r| r.htbeta.clone()).collect()` — a raw read
    /// of the DENSE per-row cross block. A matrix-free system does not have one:
    /// `set_row_htbeta_operator` installs `H_tβ` as a matvec pair and the row
    /// slabs are allocated at zero columns, which is exactly what the
    /// overcomplete support-sparse lane assembles
    /// (`SaeSupportSparseTerm::assemble_arrow_schur`: `htbeta_cols = 0`, plus
    /// per-atom `block_offsets`). Subscripting a `(d, 0)` slab at a border
    /// column is `ndarray: index out of bounds`, on whichever rayon worker the
    /// escalated PCG tier happens to run on.
    ///
    /// Every OTHER reduced-Schur kernel was already taught this — see
    /// `assemble_local_schur_block` and `build_schur_scalar_inv`, whose comments
    /// say "never a raw `row.htbeta` read at the global `sys.d`: matvec-backed
    /// rows carry absent/zero-sized slabs by contract (a raw read is wrong or
    /// panics)". This was the one place on the preconditioner ladder that was
    /// not, and it is what ClusterJacobi, AdditiveSchwarz,
    /// DiagAssembledSchwarz, BlockIncompleteCholesky and the co-visibility group
    /// builder all construct first.
    ///
    /// The row's cross block is recovered EXACTLY, one row at a time, by
    /// applying the routed transpose to each unit vector of the row's own latent
    /// space: `H_βt^(i) e_c` is row `c` of `H_tβ^(i)`. That is `d_i` operator
    /// applications, and the per-column activity is accumulated in ABSOLUTE
    /// value so a column that is nonzero in two latent rows and sums to zero
    /// across them still reads as active — the predicate is "some entry is
    /// nonzero", which is what the previous element scan tested.
    ///
    /// Taking the system also removes the clone: the old signature copied every
    /// row's `(d, k)` slab, `O(n·d·k)` of transient memory on a dense system.
    ///
    /// # Cost
    ///
    /// One linear pass over the cross-block data: `Σ_i d_i · k`. The element
    /// scan it replaces exited early per block, so it usually touched only the
    /// first nonzero column of each — cheaper by roughly the average block
    /// width. That early exit needs random access to individual entries, which
    /// is precisely what a matrix-free operator does not offer, and for a dense
    /// system the new pass is one sweep over slabs the system is already holding
    /// resident (a system whose `Σ_i d_i · k` did not fit could not have
    /// allocated them). The graph is built once per preconditioner escalation.
    pub(crate) fn build_from_system(sys: &ArrowSchurSystem) -> Self {
        let num_blocks = sys.block_offsets.len();
        if num_blocks == 0 {
            return Self::empty();
        }
        let k = sys.k;
        let mut column_mass = vec![0.0_f64; k];
        let mut probe = Array1::<f64>::zeros(k);
        let mut active = Vec::<usize>::new();
        let mut unit = Array1::<f64>::zeros(sys.d.max(1));
        let mut edge_count = std::collections::BTreeMap::<(usize, usize), f64>::new();
        for (row_idx, row) in sys.rows.iter().enumerate() {
            let latent_dim = sys.row_dims[row_idx];
            column_mass.fill(0.0);
            for axis in 0..latent_dim {
                unit.fill(0.0);
                unit[axis] = 1.0;
                probe.fill(0.0);
                let unit_row = unit.slice(ndarray::s![..latent_dim]);
                sys_htbeta_accumulate_transpose(sys, row_idx, row, unit_row, &mut probe);
                for (mass, value) in column_mass.iter_mut().zip(probe.iter()) {
                    *mass += value.abs();
                }
            }
            active.clear();
            for (block, range) in sys.block_offsets.iter().enumerate() {
                if range.clone().any(|col| column_mass[col] != 0.0) {
                    active.push(block);
                }
            }
            Self::count_co_firing(&active, &mut edge_count);
        }
        Self::from_edge_counts(num_blocks, edge_count)
    }

    fn empty() -> Self {
        Self {
            num_blocks: 0,
            edges: Vec::new(),
            adj_start: vec![0],
            adj_targets: Vec::new(),
            adj_weights: Vec::new(),
        }
    }

    /// Accumulate the co-firing MULTIPLICITY per undirected pair: the number of
    /// rows where the two blocks are simultaneously active. The unweighted edge
    /// set (for `component_partition`) is its key set; the counts become the
    /// co-visibility weights that drive the bounded cluster partition.
    fn count_co_firing(
        active: &[usize],
        edge_count: &mut std::collections::BTreeMap<(usize, usize), f64>,
    ) {
        for i in 0..active.len() {
            for j in (i + 1)..active.len() {
                let key = (active[i].min(active[j]), active[i].max(active[j]));
                *edge_count.entry(key).or_insert(0.0) += 1.0;
            }
        }
    }

    fn from_edge_counts(
        num_blocks: usize,
        edge_count: std::collections::BTreeMap<(usize, usize), f64>,
    ) -> Self {
        let edges: Vec<_> = edge_count.keys().map(|&(a, b)| BetaEdge { a, b }).collect();
        let weights: Vec<f64> = edge_count.values().copied().collect();
        let mut degree = vec![0usize; num_blocks];
        for &BetaEdge { a, b } in &edges {
            degree[a] += 1;
            degree[b] += 1;
        }
        let mut adj_start = vec![0usize; num_blocks + 1];
        for block in 0..num_blocks {
            adj_start[block + 1] = adj_start[block] + degree[block];
        }
        let mut adj_targets = vec![0usize; adj_start[num_blocks]];
        let mut adj_weights = vec![0.0f64; adj_start[num_blocks]];
        let mut cursor = adj_start[..num_blocks].to_vec();
        for (&BetaEdge { a, b }, &w) in edges.iter().zip(weights.iter()) {
            adj_targets[cursor[a]] = b;
            adj_weights[cursor[a]] = w;
            cursor[a] += 1;
            adj_targets[cursor[b]] = a;
            adj_weights[cursor[b]] = w;
            cursor[b] += 1;
        }
        Self {
            num_blocks,
            edges,
            adj_start,
            adj_targets,
            adj_weights,
        }
    }

    pub(crate) fn neighbours(&self, node: usize) -> &[usize] {
        &self.adj_targets[self.adj_start[node]..self.adj_start[node + 1]]
    }

    /// Co-visibility neighbours of `node` paired with their co-firing weights.
    pub(crate) fn weighted_neighbours(
        &self,
        node: usize,
    ) -> impl Iterator<Item = (usize, f64)> + '_ {
        let lo = self.adj_start[node];
        let hi = self.adj_start[node + 1];
        self.adj_targets[lo..hi]
            .iter()
            .copied()
            .zip(self.adj_weights[lo..hi].iter().copied())
    }

    /// Add each unassigned co-visibility neighbour of `block` to `gain`,
    /// accumulating its co-firing weight toward the growing cluster.
    fn accumulate_frontier(
        &self,
        block: usize,
        assigned: &[bool],
        gain: &mut std::collections::HashMap<usize, f64>,
    ) {
        for (nbr, w) in self.weighted_neighbours(block) {
            if !assigned[nbr] {
                *gain.entry(nbr).or_insert(0.0) += w;
            }
        }
    }

    /// Bounded greedy co-visibility partition of the β-coupling graph (Kushal &
    /// Agarwal, "Visibility Based Preconditioning for Bundle Adjustment",
    /// CVPR 2012).
    ///
    /// `component_partition` groups by CONNECTED COMPONENT, which at real over-
    /// complete SAE widths collapses into a single giant component (transitive
    /// co-firing) — pushing the dense per-cluster factor past the size cap and
    /// degrading the whole cluster tier to scalar Jacobi, the scaling ceiling.
    /// This routine instead PARTITIONS the graph into clusters bounded by
    /// `max_block_cols` (total member columns) that keep the STRONGEST co-firing
    /// edges inside a cluster: seed at the lowest unassigned block, then
    /// repeatedly absorb the unassigned neighbour with the greatest accumulated
    /// co-firing weight to the current cluster, stopping when the next admissible
    /// block would push the cluster's column count past the cap. A block wider
    /// than the cap on its own becomes a singleton (the caller scalar-degrades
    /// it). Every block lands in exactly one cluster.
    ///
    /// Deterministic and independent of thread scheduling: the seed order is
    /// ascending block index and candidate ties break by ascending index, so the
    /// partition is a pure function of the graph. It is a partition of the SAME
    /// operator (the preconditioner only steers the PCG iterate), so the solve
    /// converges to the same reduced-system solution — the change is pure
    /// numerics on the CG path, REML-neutral.
    pub(crate) fn covisibility_cluster_partition(
        &self,
        block_offsets: &[Range<usize>],
        max_block_cols: usize,
    ) -> Vec<Vec<usize>> {
        let nb = self.num_blocks;
        let mut clusters: Vec<Vec<usize>> = Vec::new();
        if nb == 0 {
            return clusters;
        }
        let block_cols = |b: usize| {
            block_offsets[b]
                .end
                .saturating_sub(block_offsets[b].start)
                .max(1)
        };
        let cap = max_block_cols.max(1);
        let mut assigned = vec![false; nb];
        for seed in 0..nb {
            if assigned[seed] {
                continue;
            }
            let mut members = vec![seed];
            assigned[seed] = true;
            let mut cols = block_cols(seed);
            let mut gain = std::collections::HashMap::<usize, f64>::new();
            self.accumulate_frontier(seed, &assigned, &mut gain);
            loop {
                // Pick the admissible frontier block of greatest co-firing weight
                // to the current cluster; tie-break by ascending index for a
                // schedule-independent partition. Stale (already-assigned) map
                // entries are skipped rather than eagerly purged.
                let mut best: Option<(usize, f64)> = None;
                for (&cand, &g) in gain.iter() {
                    if assigned[cand] || cols + block_cols(cand) > cap {
                        continue;
                    }
                    let improves = match best {
                        None => true,
                        Some((bi, bg)) => g > bg || (g == bg && cand < bi),
                    };
                    if improves {
                        best = Some((cand, g));
                    }
                }
                let Some((cand, _)) = best else {
                    break;
                };
                assigned[cand] = true;
                members.push(cand);
                cols += block_cols(cand);
                gain.remove(&cand);
                self.accumulate_frontier(cand, &assigned, &mut gain);
            }
            members.sort_unstable();
            clusters.push(members);
        }
        clusters
    }

    pub(crate) fn component_partition(&self) -> Vec<Vec<usize>> {
        let mut parent: Vec<usize> = (0..self.num_blocks).collect();
        let mut rank = vec![0u8; self.num_blocks];

        fn find(parent: &mut [usize], mut x: usize) -> usize {
            while parent[x] != x {
                parent[x] = parent[parent[x]];
                x = parent[x];
            }
            x
        }

        for &BetaEdge { a, b } in &self.edges {
            let lhs = find(&mut parent, a);
            let rhs = find(&mut parent, b);
            if lhs != rhs {
                if rank[lhs] < rank[rhs] {
                    parent[lhs] = rhs;
                } else if rank[lhs] > rank[rhs] {
                    parent[rhs] = lhs;
                } else {
                    parent[rhs] = lhs;
                    rank[lhs] += 1;
                }
            }
        }

        let mut label_map = vec![usize::MAX; self.num_blocks];
        let mut parts = Vec::<Vec<usize>>::new();
        for block in 0..self.num_blocks {
            let root = find(&mut parent, block);
            let label = if label_map[root] == usize::MAX {
                label_map[root] = parts.len();
                parts.push(Vec::new());
                label_map[root]
            } else {
                label_map[root]
            };
            parts[label].push(block);
        }
        parts
    }

    pub(crate) fn expand_one_hop(&self, seed: &[usize]) -> Vec<usize> {
        let mut expanded = seed.to_vec();
        for &block in seed {
            expanded.extend_from_slice(self.neighbours(block));
        }
        expanded.sort_unstable();
        expanded.dedup();
        expanded
    }
}
// ---------------------------------------------------------------------------
// BetaPenaltyOp — matrix-free penalty-side H_ββ abstraction (#296)
// ---------------------------------------------------------------------------

/// Identifies one contiguous column block in the shared β vector for
/// block-Jacobi Schur pre-conditioning (#287).
///
/// A `BetaBlockId(i)` refers to the `i`-th range in
/// [`ArrowSchurSystem::block_offsets`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BetaBlockId(pub usize);

/// Matrix-free operator for the penalty side of `H_ββ`.
///
/// Callers must satisfy the additive convention: every method **adds** its
/// contribution to the output buffer (i.e. `y += P x`, not `y = P x`).
/// This matches the assembly pattern where multiple penalty terms are
/// accumulated into the same gradient / Hessian buffers.
pub trait BetaPenaltyOp: Send + Sync {
    /// Full dimension `K` of the β vector.
    fn dim(&self) -> usize;
    /// `y += P x` — penalty Hessian-vector product (length `K`).
    fn matvec(&self, x: &[f64], y: &mut [f64]);
    /// Penalty gradient: `out += P β`.
    fn gradient(&self, beta: &[f64], out: &mut [f64]);
    /// `diag += diag(P)` — diagonal entries used by Jacobi preconditioner.
    fn diagonal(&self, diag: &mut [f64]);
    /// Add the `b×b` dense penalty sub-block for block `id` into `out`
    /// (row-major, block size `b = offsets[id.0].len()`).
    /// Used by the block-Jacobi Schur preconditioner (#287).
    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>);
    /// Materialize the full `K×K` dense penalty matrix (needed by
    /// Direct / SqrtBA modes that form the Schur complement explicitly).
    fn to_dense(&self) -> Array2<f64>;
    /// Mix the operator's defining state into `hasher` for cache-validity
    /// fingerprinting. Must change whenever `matvec` / `to_dense` would change,
    /// so the factorization / evidence cache (`cache_matches_system`) is
    /// invalidated when the β-block content changes. Implementations hash their
    /// own compact defining data (e.g. Kronecker factors, block matrices)
    /// rather than the full `K×K` dense form, which would defeat the structured
    /// operator's storage savings.
    fn fingerprint(&self, hasher: &mut Fingerprinter);

    /// If this operator writes its `matvec` contribution into EXACTLY one
    /// contiguous output range `[start, end)` and touches no other index,
    /// return that range; otherwise `None` (the default — opaque / scattered
    /// output). This lets [`CompositePenaltyOp::matvec`] fan a leading run of
    /// mutually-disjoint operators across rayon workers, each writing its own
    /// output sub-slice with no cross-thread aliasing. The per-atom Kronecker
    /// smooth penalties (the SAE prologue's serial Amdahl ceiling at the K=32k
    /// manifold border) each cover one atom's β block, so they qualify.
    fn output_range(&self) -> Option<Range<usize>> {
        None
    }

    /// Accumulate `matvec`'s contribution into `y_local`, where `y_local[i]`
    /// aliases global output index `output_range().start + i` (so
    /// `y_local.len() == output_range().len()`), reading the FULL-length input
    /// `x`. ONLY valid when [`Self::output_range`] returns `Some`; the default
    /// panics because a `None`-range operator has no single contiguous slice to
    /// write into. Must be BIT-IDENTICAL to the corresponding indices of
    /// `matvec` (same per-index accumulation order) so the composite's parallel
    /// prefix reproduces the serial result exactly.
    fn matvec_local(&self, x: &[f64], y_local: &mut [f64]) {
        // SAFETY: hard contract guard, not a silent sentinel. A `None`-range
        // `BetaPenaltyOp` exposes no single contiguous output slice, so the
        // local matvec is undefined and MUST be routed through `matvec`. This
        // default exists only to turn that misuse into a loud, immediate
        // failure at the call site; the input/local extents are surfaced for
        // triage.
        panic!(
            "matvec_local requires output_range() == Some; a None-range \
             BetaPenaltyOp (input len {}, local output len {}) must be applied \
             through matvec",
            x.len(),
            y_local.len()
        );
    }
}

/// Dense fallback: wraps the existing `K×K` `H_ββ` accumulator.
pub struct DensePenaltyOp(pub Array2<f64>);

impl BetaPenaltyOp for DensePenaltyOp {
    fn dim(&self) -> usize {
        self.0.nrows()
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        let k = self.0.nrows();
        for a in 0..k {
            let mut acc = 0.0_f64;
            for b in 0..k {
                acc += self.0[[a, b]] * x[b];
            }
            y[a] += acc;
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        let k = self.0.nrows();
        for a in 0..k {
            let mut acc = 0.0_f64;
            for b in 0..k {
                acc += self.0[[a, b]] * beta[b];
            }
            out[a] += acc;
        }
    }

    fn diagonal(&self, diag: &mut [f64]) {
        let k = self.0.nrows().min(diag.len());
        for j in 0..k {
            diag[j] += self.0[[j, j]];
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        let range = &offsets[id.0];
        let b = range.end - range.start;
        for bi in 0..b {
            for bj in 0..b {
                out[[bi, bj]] += self.0[[range.start + bi, range.start + bj]];
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        self.0.clone()
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("dense-penalty-op-v1");
        hasher.write_f64_array2(&self.0);
    }
}

/// Coupled-carrier PSD penalty operator
///
/// ```text
///     P = Σ_{a,b} C[a,b] · v_a v_bᵀ
/// ```
///
/// with a SMALL dense symmetric coupling `C` (`ne × ne`) and `ne` carriers
/// `v_a`, each supported on a few contiguous runs of the ambient β vector
/// (in practice: the two atom decoder blocks a co-firing edge touches).
///
/// This is the FACTORED form of a sum of rank-1 terms. Diagonalizing
/// `C = Σ_r λ_r e_r e_rᵀ` and expanding `w_r = Σ_a e_r[a] v_a` gives the
/// identical operator `Σ_r λ_r w_r w_rᵀ` — and that expansion is exactly what
/// this type exists to avoid:
///
/// | | expanded `Σ_r λ_r w_r w_rᵀ` | factored (here) |
/// |---|---|---|
/// | build | `ne² · max‖v_a‖₀` | nothing to build |
/// | store | `ne · ‖∪_a supp v_a‖₀` | `ne · max‖v_a‖₀ + ne²` |
/// | matvec | `ne · ‖∪_a supp v_a‖₀` | `2·ne · max‖v_a‖₀ + ne²` |
///
/// Each `w_r` is dense over the UNION of the carriers' supports even though
/// every `v_a` touches only two blocks, so the expansion trades a factor
/// `ne · (#blocks)/2` in both storage and application for nothing — the operator
/// is the same. On the SAE separation barrier (#2731) that expansion was the
/// whole cost of the curved tier at `p=2048, charts=32`: `ne` grows with the
/// number of REALIZED co-firing chart pairs, so `ne²·‖v_a‖₀` dominated the fit.
///
/// PSD-ness is the caller's contract: `C` must be symmetric positive
/// semidefinite (the SAE barrier hands over `penalty_scale·|M|`, the matrix
/// absolute value of the exact Gauss–Newton overlap-space Hessian), which makes
/// `P` PSD for any carriers and so keeps any PSD penalty operator it is added to
/// positive definite.
///
/// The carriers are expressed in whatever coordinate space the operator runs in
/// (full-`B` on the non-frames path; factored border coordinates on the framed
/// path, where each full-`B` `v_a` is projected `v_a ← Φᵀ v_a` before
/// construction — the identity `Φᵀ(Σ C[a,b] v_a v_bᵀ)Φ = Σ C[a,b](Φᵀv_a)(Φᵀv_b)ᵀ`
/// is what makes that legal, and it is why `C` needs no transformation).
pub struct CoupledCarrierPenaltyOp {
    /// Full β dimension `K` of the coordinate space the carriers live in.
    pub k: usize,
    /// Symmetric PSD coupling `C` (`ne × ne`), with any scalar prefactor
    /// already folded in.
    pub coupling: Array2<f64>,
    /// `carriers[a]` — the runs `(global_start, values)` of `v_a`. Runs inside
    /// one carrier are disjoint and every `global_start + values.len() <= k`.
    pub carriers: Vec<Vec<(usize, Vec<f64>)>>,
}

/// Overlap of the ambient index ranges `[a_start, a_start+a_len)` and
/// `[b_start, b_start+b_len)`, as `(offset into a, offset into b, length)`.
#[inline]
fn coupled_run_overlap(
    a_start: usize,
    a_len: usize,
    b_start: usize,
    b_len: usize,
) -> Option<(usize, usize, usize)> {
    let lo = a_start.max(b_start);
    let hi = (a_start + a_len).min(b_start + b_len);
    if hi <= lo {
        None
    } else {
        Some((lo - a_start, lo - b_start, hi - lo))
    }
}

impl CoupledCarrierPenaltyOp {
    /// `out[a] = v_aᵀ x`.
    fn carrier_dots(&self, x: &[f64]) -> Vec<f64> {
        self.carriers
            .iter()
            .map(|runs| {
                let mut acc = 0.0_f64;
                for (start, values) in runs {
                    let xs = &x[*start..*start + values.len()];
                    for (value, xi) in values.iter().zip(xs) {
                        acc += value * xi;
                    }
                }
                acc
            })
            .collect()
    }

    /// `y += Σ_a coefficients[a] · v_a`.
    fn scatter_carriers(&self, coefficients: &[f64], y: &mut [f64]) {
        for (runs, &coefficient) in self.carriers.iter().zip(coefficients) {
            if coefficient == 0.0 {
                continue;
            }
            for (start, values) in runs {
                let ys = &mut y[*start..*start + values.len()];
                for (value, yi) in values.iter().zip(ys) {
                    *yi += coefficient * value;
                }
            }
        }
    }

    /// `y += P x`, shared by `matvec` and `gradient` (both are `P·vector`).
    fn apply_add(&self, x: &[f64], y: &mut [f64]) {
        let dots = self.carrier_dots(x);
        let ne = self.carriers.len();
        let mut coefficients = vec![0.0_f64; ne];
        for a in 0..ne {
            let mut acc = 0.0_f64;
            for b in 0..ne {
                acc += self.coupling[[a, b]] * dots[b];
            }
            coefficients[a] = acc;
        }
        self.scatter_carriers(&coefficients, y);
    }

    /// Visit every `(global row, global column, value)` of `C[a,b]·v_a v_bᵀ`
    /// restricted to the ambient window `[window.start, window.end)` on BOTH
    /// axes. The window is the whole space for `to_dense` and one block range
    /// for `block`; emitting through one walk keeps the two consistent.
    fn for_each_windowed_entry<F: FnMut(usize, usize, f64)>(
        &self,
        window: Range<usize>,
        mut visit: F,
    ) {
        let ne = self.carriers.len();
        let (lo, hi) = (window.start, window.end);
        let span = hi.saturating_sub(lo);
        for a in 0..ne {
            for b in 0..ne {
                let coupling = self.coupling[[a, b]];
                if coupling == 0.0 {
                    continue;
                }
                for (start_a, values_a) in &self.carriers[a] {
                    let Some((skip_a, _, len_a)) =
                        coupled_run_overlap(*start_a, values_a.len(), lo, span)
                    else {
                        continue;
                    };
                    for (start_b, values_b) in &self.carriers[b] {
                        let Some((skip_b, _, len_b)) =
                            coupled_run_overlap(*start_b, values_b.len(), lo, span)
                        else {
                            continue;
                        };
                        for i in 0..len_a {
                            let row = start_a + skip_a + i;
                            let scaled = coupling * values_a[skip_a + i];
                            if scaled == 0.0 {
                                continue;
                            }
                            for j in 0..len_b {
                                visit(row, start_b + skip_b + j, scaled * values_b[skip_b + j]);
                            }
                        }
                    }
                }
            }
        }
    }
}

impl BetaPenaltyOp for CoupledCarrierPenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        self.apply_add(x, y);
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        self.apply_add(beta, out);
    }

    fn diagonal(&self, diag: &mut [f64]) {
        // `diag(P)[i] = Σ_{a,b} C[a,b]·v_a[i]·v_b[i]` — accumulated over the
        // OVERLAP of the two carriers' runs, so the cost is the shared support
        // (each atom block, once per pair of edges touching it) rather than the
        // full `ne × ‖∪ supp‖` the expanded form pays.
        let ne = self.carriers.len();
        for a in 0..ne {
            for b in 0..ne {
                let coupling = self.coupling[[a, b]];
                if coupling == 0.0 {
                    continue;
                }
                for (start_a, values_a) in &self.carriers[a] {
                    for (start_b, values_b) in &self.carriers[b] {
                        let Some((skip_a, skip_b, len)) = coupled_run_overlap(
                            *start_a,
                            values_a.len(),
                            *start_b,
                            values_b.len(),
                        ) else {
                            continue;
                        };
                        let base = start_a + skip_a;
                        if base >= diag.len() {
                            continue;
                        }
                        let len = len.min(diag.len() - base);
                        for i in 0..len {
                            diag[base + i] += coupling * values_a[skip_a + i] * values_b[skip_b + i];
                        }
                    }
                }
            }
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        // Only the intra-block restriction reaches this block-Jacobi sub-block;
        // the cross-block coupling is invisible to the preconditioner, exactly
        // as for any operator with off-block entries.
        let range = offsets[id.0].clone();
        let start = range.start;
        self.for_each_windowed_entry(range, |row, col, value| {
            out[[row - start, col - start]] += value;
        });
    }

    fn to_dense(&self) -> Array2<f64> {
        let mut out = Array2::<f64>::zeros((self.k, self.k));
        self.for_each_windowed_entry(0..self.k, |row, col, value| {
            out[[row, col]] += value;
        });
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("coupled-carrier-penalty-op-v1");
        hasher.write_usize(self.k);
        hasher.write_f64_array2(&self.coupling);
        hasher.write_usize(self.carriers.len());
        for runs in &self.carriers {
            hasher.write_usize(runs.len());
            for (start, values) in runs {
                hasher.write_usize(*start);
                hasher.write_usize(values.len());
                for &value in values {
                    hasher.write_f64(value);
                }
            }
        }
    }
}

/// Block-local penalty operator: applies per-block penalty matrices
/// (matching `ParameterBlockSpec` boundaries) without materialising a
/// full `K×K` dense matrix.
///
/// Each entry is `(global_offset, local_matrix)` where `global_offset`
/// is the start of that block in the full β vector.
pub struct BlockPenaltyOp {
    /// Full β dimension `K`.
    pub k: usize,
    /// `(global_start, local_matrix)` for each atom/block.
    pub blocks: Vec<(usize, Array2<f64>)>,
}

impl BetaPenaltyOp for BlockPenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        for (off, local) in &self.blocks {
            let b = local.nrows();
            for i in 0..b {
                let gi = off + i;
                let mut acc = 0.0_f64;
                for j in 0..b {
                    acc += local[[i, j]] * x[off + j];
                }
                y[gi] += acc;
            }
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        for (off, local) in &self.blocks {
            let b = local.nrows();
            for i in 0..b {
                let gi = off + i;
                let mut acc = 0.0_f64;
                for j in 0..b {
                    acc += local[[i, j]] * beta[off + j];
                }
                out[gi] += acc;
            }
        }
    }

    fn diagonal(&self, diag: &mut [f64]) {
        for (off, local) in &self.blocks {
            let b = local.nrows();
            for j in 0..b {
                diag[off + j] += local[[j, j]];
            }
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        let range = &offsets[id.0];
        let b_out = range.end - range.start;
        for (off, local) in &self.blocks {
            let b = local.nrows();
            let block_end = off + b;
            if block_end <= range.start || *off >= range.end {
                continue;
            }
            for bi in 0..b_out {
                let gi = range.start + bi;
                if gi < *off || gi >= block_end {
                    continue;
                }
                let li = gi - off;
                for bj in 0..b_out {
                    let gj = range.start + bj;
                    if gj < *off || gj >= block_end {
                        continue;
                    }
                    let lj = gj - off;
                    out[[bi, bj]] += local[[li, lj]];
                }
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let mut out = Array2::<f64>::zeros((self.k, self.k));
        for (off, local) in &self.blocks {
            let b = local.nrows();
            for i in 0..b {
                for j in 0..b {
                    out[[off + i, off + j]] += local[[i, j]];
                }
            }
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("block-penalty-op-v1");
        hasher.write_usize(self.k);
        hasher.write_usize(self.blocks.len());
        for (off, local) in &self.blocks {
            hasher.write_usize(*off);
            hasher.write_f64_array2(local);
        }
    }
}

/// Kronecker-product penalty: `P = A ⊗ B` applied without materialising
/// the full `(p_a·p_b)×(p_a·p_b)` matrix.
pub struct KroneckerPenaltyOp {
    /// Left factor `A`, shape `(p_a, p_a)`.
    pub factor_a: Array2<f64>,
    /// Right factor `B`, shape `(p_b, p_b)`.
    pub factor_b: Array2<f64>,
    /// Global offset into the β vector where this block starts.
    pub global_offset: usize,
    /// Full β dimension `K`.
    pub k: usize,
}

impl BetaPenaltyOp for KroneckerPenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        let p_a = self.factor_a.nrows();
        let p_b = self.factor_b.nrows();
        let off = self.global_offset;
        // (A ⊗ B) vec(V) where V is (p_b, p_a) with Fortran/vec ordering.
        for i_a in 0..p_a {
            for i_b in 0..p_b {
                let gi = off + i_a * p_b + i_b;
                let mut acc = 0.0_f64;
                for j_a in 0..p_a {
                    let a_ij = self.factor_a[[i_a, j_a]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    for j_b in 0..p_b {
                        acc += a_ij * self.factor_b[[i_b, j_b]] * x[off + j_a * p_b + j_b];
                    }
                }
                y[gi] += acc;
            }
        }
    }

    fn output_range(&self) -> Option<Range<usize>> {
        let off = self.global_offset;
        Some(off..off + self.factor_a.nrows() * self.factor_b.nrows())
    }

    fn matvec_local(&self, x: &[f64], y_local: &mut [f64]) {
        // Byte-for-byte the `matvec` arithmetic with the output written at the
        // LOCAL index `i_a·p_b + i_b` (== global `gi - off`), so the composite
        // can apply this block into its own `y[off..off+p_a·p_b]` sub-slice in
        // parallel. Per-index accumulation order is unchanged ⇒ bit-identical.
        let p_a = self.factor_a.nrows();
        let p_b = self.factor_b.nrows();
        let off = self.global_offset;
        for i_a in 0..p_a {
            for i_b in 0..p_b {
                let li = i_a * p_b + i_b;
                let mut acc = 0.0_f64;
                for j_a in 0..p_a {
                    let a_ij = self.factor_a[[i_a, j_a]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    for j_b in 0..p_b {
                        acc += a_ij * self.factor_b[[i_b, j_b]] * x[off + j_a * p_b + j_b];
                    }
                }
                y_local[li] += acc;
            }
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        let p_a = self.factor_a.nrows();
        let p_b = self.factor_b.nrows();
        let off = self.global_offset;
        for i_a in 0..p_a {
            for i_b in 0..p_b {
                let gi = off + i_a * p_b + i_b;
                let mut acc = 0.0_f64;
                for j_a in 0..p_a {
                    let a_ij = self.factor_a[[i_a, j_a]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    for j_b in 0..p_b {
                        acc += a_ij * self.factor_b[[i_b, j_b]] * beta[off + j_a * p_b + j_b];
                    }
                }
                out[gi] += acc;
            }
        }
    }

    fn diagonal(&self, diag: &mut [f64]) {
        let p_a = self.factor_a.nrows();
        let p_b = self.factor_b.nrows();
        let off = self.global_offset;
        for i_a in 0..p_a {
            for i_b in 0..p_b {
                diag[off + i_a * p_b + i_b] +=
                    self.factor_a[[i_a, i_a]] * self.factor_b[[i_b, i_b]];
            }
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        let range = &offsets[id.0];
        let b = range.end - range.start;
        let p_a = self.factor_a.nrows();
        let p_b = self.factor_b.nrows();
        let off = self.global_offset;
        let block_end = off + p_a * p_b;
        if block_end <= range.start || off >= range.end {
            return;
        }
        for bi in 0..b {
            let gi = range.start + bi;
            if gi < off || gi >= block_end {
                continue;
            }
            let li = gi - off;
            let i_a = li / p_b;
            let i_b = li % p_b;
            for bj in 0..b {
                let gj = range.start + bj;
                if gj < off || gj >= block_end {
                    continue;
                }
                let lj = gj - off;
                let j_a = lj / p_b;
                let j_b = lj % p_b;
                out[[bi, bj]] += self.factor_a[[i_a, j_a]] * self.factor_b[[i_b, j_b]];
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let p_a = self.factor_a.nrows();
        let p_b = self.factor_b.nrows();
        let off = self.global_offset;
        let mut out = Array2::<f64>::zeros((self.k, self.k));
        for i_a in 0..p_a {
            for i_b in 0..p_b {
                let gi = off + i_a * p_b + i_b;
                for j_a in 0..p_a {
                    let a_ij = self.factor_a[[i_a, j_a]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    for j_b in 0..p_b {
                        let gj = off + j_a * p_b + j_b;
                        out[[gi, gj]] += a_ij * self.factor_b[[i_b, j_b]];
                    }
                }
            }
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("kronecker-penalty-op-v1");
        hasher.write_usize(self.global_offset);
        hasher.write_usize(self.k);
        hasher.write_f64_array2(&self.factor_a);
        hasher.write_f64_array2(&self.factor_b);
    }
}

/// Kronecker-product penalty with an identity right factor:
/// `P = A ⊗ I_p`.
///
/// This is the hot SAE smoothness case. Storing `I_p` as a dense matrix costs
/// `O(p²)` memory per atom and makes every matvec pay an unnecessary right-factor
/// loop. This operator stores only the identity dimension and keeps the same
/// layout as [`KroneckerPenaltyOp`]: local index `i_a * p + i_b`.
pub struct IdentityRightKroneckerPenaltyOp {
    /// Left factor `A`, shape `(p_a, p_a)`.
    pub factor_a: Array2<f64>,
    /// Identity right-factor dimension `p`.
    pub p: usize,
    /// Global offset into the β vector where this block starts.
    pub global_offset: usize,
    /// Full β dimension `K`.
    pub k: usize,
}

impl BetaPenaltyOp for IdentityRightKroneckerPenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        let p_a = self.factor_a.nrows();
        let p = self.p;
        let off = self.global_offset;
        for i_a in 0..p_a {
            for i_b in 0..p {
                let gi = off + i_a * p + i_b;
                let mut acc = 0.0_f64;
                for j_a in 0..p_a {
                    let a_ij = self.factor_a[[i_a, j_a]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    acc += a_ij * x[off + j_a * p + i_b];
                }
                y[gi] += acc;
            }
        }
    }

    fn output_range(&self) -> Option<Range<usize>> {
        let off = self.global_offset;
        Some(off..off + self.factor_a.nrows() * self.p)
    }

    fn matvec_local(&self, x: &[f64], y_local: &mut [f64]) {
        // Byte-for-byte the `matvec` inner arithmetic, but the output writes to
        // the LOCAL index `i_a·p + i_b` (== global `gi - off`) so the composite
        // can hand this operator its own `y[off..off+p_a·p]` sub-slice. The
        // per-index accumulation order over `j_a` is unchanged, so the result is
        // bit-identical to `matvec`.
        let p_a = self.factor_a.nrows();
        let p = self.p;
        let off = self.global_offset;
        for i_a in 0..p_a {
            for i_b in 0..p {
                let li = i_a * p + i_b;
                let mut acc = 0.0_f64;
                for j_a in 0..p_a {
                    let a_ij = self.factor_a[[i_a, j_a]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    acc += a_ij * x[off + j_a * p + i_b];
                }
                y_local[li] += acc;
            }
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        self.matvec(beta, out);
    }

    fn diagonal(&self, diag: &mut [f64]) {
        let p_a = self.factor_a.nrows();
        let p = self.p;
        let off = self.global_offset;
        for i_a in 0..p_a {
            let a_ii = self.factor_a[[i_a, i_a]];
            for i_b in 0..p {
                diag[off + i_a * p + i_b] += a_ii;
            }
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        let range = &offsets[id.0];
        let b = range.end - range.start;
        let p_a = self.factor_a.nrows();
        let p = self.p;
        let off = self.global_offset;
        let block_end = off + p_a * p;
        if block_end <= range.start || off >= range.end {
            return;
        }
        for bi in 0..b {
            let gi = range.start + bi;
            if gi < off || gi >= block_end {
                continue;
            }
            let li = gi - off;
            let i_a = li / p;
            let i_b = li % p;
            for bj in 0..b {
                let gj = range.start + bj;
                if gj < off || gj >= block_end {
                    continue;
                }
                let lj = gj - off;
                let j_a = lj / p;
                let j_b = lj % p;
                if i_b == j_b {
                    out[[bi, bj]] += self.factor_a[[i_a, j_a]];
                }
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let p_a = self.factor_a.nrows();
        let p = self.p;
        let off = self.global_offset;
        let mut out = Array2::<f64>::zeros((self.k, self.k));
        for i_a in 0..p_a {
            for j_a in 0..p_a {
                let a_ij = self.factor_a[[i_a, j_a]];
                if a_ij == 0.0 {
                    continue;
                }
                for i_b in 0..p {
                    let gi = off + i_a * p + i_b;
                    let gj = off + j_a * p + i_b;
                    out[[gi, gj]] += a_ij;
                }
            }
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("identity-right-kronecker-penalty-op-v1");
        hasher.write_usize(self.global_offset);
        hasher.write_usize(self.k);
        hasher.write_usize(self.p);
        hasher.write_f64_array2(&self.factor_a);
    }
}

/// One co-occurring atom-pair block of a block-sparse left factor `A`.
///
/// `data` is the dense `(m_i × m_j)` coupling between the basis columns of
/// atom `i` (rows, starting at left-factor offset `row_off`) and atom `j`
/// (columns, starting at `col_off`). Both offsets are in *left-factor* (`A`)
/// coordinates, i.e. `μ`-space, not β-space.
#[derive(Debug, Clone)]
pub struct SparseGBlock {
    /// Left-factor (`μ`-space) row offset = `beta_offset[atom_i] / p`.
    pub row_off: usize,
    /// Left-factor (`μ`-space) column offset = `beta_offset[atom_j] / p`.
    pub col_off: usize,
    /// Dense `(m_i × m_j)` coupling block.
    pub data: Array2<f64>,
}

/// Block-sparse Kronecker penalty `P = A ⊗ I_p` where the left factor `A`
/// (dimension `dim_a × dim_a` in `μ`-space) is stored only on its non-empty
/// co-occurring atom-pair blocks rather than as a dense `(dim_a × dim_a)`
/// matrix.
///
/// This is the sparse-atom (`K = 100K`) replacement for wrapping the dense
/// data-fit Gauss-Newton Gram `G` (`m_total × m_total`) in a
/// [`KroneckerPenaltyOp`]: with per-row active sets of size `k_active ≪ K`,
/// only the `(atom, atom')` pairs that co-occur in some row contribute a
/// non-zero `(m_i × m_j)` block, so the storage and every matvec/diagonal
/// pass cost `O(Σ_pairs m_i m_j · p)` instead of `O((m_total · p)²)`.
///
/// The β index of left-factor coordinate `μ` and output channel `oc` is
/// `μ · p + oc` (the same `μ`-major / `oc`-minor layout the dense
/// `KroneckerPenaltyOp { factor_b: I_p }` uses), so this op is a drop-in
/// structured replacement: with the full dense pair set it reproduces the
/// dense operator exactly.
pub struct SparseBlockKroneckerPenaltyOp {
    /// Right-factor identity dimension `p` (number of decoder output channels).
    pub p: usize,
    /// Left-factor dimension `dim_a` in `μ`-space (= `m_total`).
    pub dim_a: usize,
    /// Full β dimension `K = dim_a · p`.
    pub k: usize,
    /// Non-empty `(atom_i, atom_j)` coupling blocks of `A`.
    pub blocks: Vec<SparseGBlock>,
}

#[derive(Debug, Clone)]
pub struct DeviceSaeSmoothBlock {
    pub global_offset: usize,
    pub factor_a: Array2<f64>,
}

/// Frame-factored extension of [`DeviceSaePcgData`] (issue #1017/#1026,
/// frames-engaged device PCG). Present only when at least one atom is genuinely
/// frame-reduced (`ranks[k] < p`); absent (`None`) on the full-`B` path, where
/// the legacy `G ⊗ I_p` channel-identical kernel applies byte-for-byte.
///
/// On the frames path the β border is the FACTORED coordinate space `C` of width
/// `Σ_k M_k·r_k`, the data-fit β-Hessian is `G_{ij} ⊗ W_{ij}` (`W_{ij}=U_iᵀU_j`,
/// carried on `frame_blocks`), the smooth penalty is `λ S_k ⊗ I_{r_k}`
/// (`smooth_blocks`, reused — width `r_k` instead of `p`), and the per-row
/// reduced-Schur cross-block `H_tβ^(i)` is the DENSE `(q_i × border_dim)` slab
/// `row_htbeta[i]` (row-major) rather than the full-`B` factored `L_i · J_β`
/// gather (so `a_phi`/`local_jac` are unused on this path).
#[derive(Debug, Clone)]
pub struct DeviceSaeFrameData {
    /// Per-atom frame rank `r_k` (factored output width); `r_k == p` for an
    /// un-framed atom riding the identity special case.
    pub ranks: Vec<usize>,
    /// Per-atom basis size `M_k`.
    pub basis_sizes: Vec<usize>,
    /// Per-atom factored-border offset `off_C[k]` (prefix sum of `M_k·r_k`),
    /// length `n_atoms`. Atom `k`'s `C_k` block is `[off_C[k] .. +M_k·r_k)`.
    pub border_offsets: Vec<usize>,
    /// Co-occurring `(atom_i, atom_j)` data-fit blocks `g ⊗ w` (`w = U_iᵀU_j`).
    pub frame_blocks: Vec<FactoredFrameGBlock>,
    /// Right-factor width (`r_k`) of each entry of the top-level
    /// `DeviceSaePcgData::smooth_blocks`, in the SAME order. On the frames path
    /// the smooth penalty is `λ S_k ⊗ I_{r_k}` so the block at
    /// `smooth_blocks[i].global_offset` has identity width `smooth_ranks[i]`
    /// (which equals `ranks[atom]`), NOT the ambient `p`.
    pub smooth_ranks: Vec<usize>,
    /// Per-row dense cross-block `H_tβ^(i)` as a row-major `q_i × border_dim`
    /// buffer (`q_i = row_dims[i]`). Empty inner `Vec` for a 0-dim row.
    pub row_htbeta: Vec<Vec<f64>>,
}

#[derive(Debug, Clone)]
pub struct DeviceSaePcgData {
    pub p: usize,
    pub beta_dim: usize,
    // #1033 large-n: the per-row support `a_phi` and local Jacobians `local_jac`
    // are ALSO held by the host matrix-free row operator (`SaeKroneckerRows`) for
    // the lifetime of the inner solve. Storing them as `Arc<[…]>` lets the
    // assembler hand BOTH consumers the SAME backing allocation instead of a
    // second full `O(n·q·p)` clone (`device_rows = (a_phi.clone(), kron_jac.clone())`
    // was the dominant always-resident duplication on the CPU non-frames path at
    // the LLM shape p≈5120). Indexing/`.len()`/iteration are identical to `Vec`.
    pub a_phi: Arc<[Vec<(usize, f64)>]>,
    pub local_jac: Arc<[Vec<f64>]>,
    pub smooth_blocks: Vec<DeviceSaeSmoothBlock>,
    pub sparse_g_blocks: Vec<SparseGBlock>,
    /// Frame-factored metadata. `None` ⇒ legacy full-`B` `G ⊗ I_p` path
    /// (byte-identical to before this field existed). `Some` ⇒ frames-engaged
    /// path: the kernel consumes `frame.frame_blocks`/`smooth_blocks` (now
    /// rank-`r_k` wide) and `frame.row_htbeta` instead of the `⊗ I_p` gather.
    pub frame: Option<DeviceSaeFrameData>,
}

impl DeviceSaePcgData {
    /// Replace a framed current-iterate payload while retaining the nested host
    /// allocations that have stable shapes across nonlinear assemblies.
    ///
    /// Full-`B` payloads keep their freshly assembled `Arc` slices because those
    /// exact allocations are shared with the CPU row operator. Framed payloads
    /// have no such aliasing (`a_phi`/`local_jac` are empty), so `Vec::clone_from`
    /// can reuse the dominant per-row dense `H_tβ` slabs, frame-Gram blocks, and
    /// smooth blocks while replacing every numerical value.
    pub(crate) fn replace_reusing_framed_allocations(&mut self, mut data: Self) {
        if self.frame.is_none() || data.frame.is_none() {
            *self = data;
            return;
        }
        let new_frame = data.frame.take().expect("framed replacement checked above");
        let current_frame = self.frame.as_mut().expect("framed receiver checked above");

        self.p = data.p;
        self.beta_dim = data.beta_dim;
        self.a_phi = data.a_phi;
        self.local_jac = data.local_jac;
        self.smooth_blocks.clone_from(&data.smooth_blocks);
        self.sparse_g_blocks.clone_from(&data.sparse_g_blocks);
        current_frame.ranks.clone_from(&new_frame.ranks);
        current_frame.basis_sizes.clone_from(&new_frame.basis_sizes);
        current_frame
            .border_offsets
            .clone_from(&new_frame.border_offsets);
        current_frame
            .frame_blocks
            .clone_from(&new_frame.frame_blocks);
        current_frame
            .smooth_ranks
            .clone_from(&new_frame.smooth_ranks);
        current_frame.row_htbeta.clone_from(&new_frame.row_htbeta);
    }

    /// Snapshot the per-row active-atom support `a_phi` into a shared `Arc<[…]>`
    /// for the CPU residency operator ([`SaeResidentReducedSchur`]). Cloned once
    /// per CG-solve build (cost `O(Σ_i m_i)`, dwarfed by the per-row factor solves
    /// in the same build), so the resident matvec borrows the index lists without
    /// re-cloning them on every CG iteration.
    pub(crate) fn a_phi_shared(&self) -> Arc<[Vec<(usize, f64)>]> {
        // #1033: `a_phi` is already an `Arc<[…]>`; hand back a refcount bump
        // (`O(1)`) rather than re-cloning every `(idx, weight)` pair per CG build.
        Arc::clone(&self.a_phi)
    }

    /// Share the per-row local Jacobians `local_jac` with the CPU residency
    /// operator ([`SaeResidentReducedSchur`]) as an `O(1)` refcount bump. The
    /// staged row factor used to hold a verbatim row-major copy of each
    /// `local_jac[row]`; sharing the slab removes that second full `O(n·di·p)`
    /// copy with byte-for-byte identical reads (#1033).
    pub(crate) fn local_jac_shared(&self) -> Arc<[Vec<f64>]> {
        Arc::clone(&self.local_jac)
    }
}

/// #1017/#2230 residency measurement: the host→device operand bytes that
/// `flatten_device_sae_data` / `flatten_device_sae_frame_data` re-upload on EVERY
/// matrix-free PCG solve (hence on every LM ridge-ladder escalation trial),
/// broken out by category. Lets the a100 job confirm which sub-lane a real fit
/// takes — legacy sparse (`⊗ I_p` a_phi/local_jac) vs framed dense (per-row
/// `row_htbeta` at the factored border) — and size the transfer a base-resident
/// frame would remove. Pure host accounting, no device contact, so it is
/// CPU-CI testable without CUDA.
#[derive(Clone, Debug, Default)]
pub struct SaePcgOperandReport {
    /// `true` when the framed dense per-row cross path is active (frame present).
    pub framed: bool,
    /// Border width `k` (= `beta_dim`).
    pub beta_dim: usize,
    /// Per-row support `a_phi`: total `(row, weight)` pairs and their bytes.
    pub a_phi_pairs: usize,
    pub a_phi_bytes: usize,
    /// Per-row local Jacobians `local_jac`: total `f64` and their bytes.
    pub local_jac_elems: usize,
    pub local_jac_bytes: usize,
    /// Smooth penalty factors `λ S_k`: bytes.
    pub smooth_bytes: usize,
    /// Sparse `G` co-occurrence blocks (legacy `⊗ I_p`): bytes.
    pub sparse_g_bytes: usize,
    /// Framed per-row dense cross `row_htbeta`: bytes (0 on the legacy path). This
    /// is the 34MiB-vs-31GiB discriminator the #2230 design report flagged: a full
    /// factored-border dense cross balloons here, a legacy fit leaves it at 0.
    pub row_htbeta_bytes: usize,
    /// Number of rows carrying a non-empty framed `row_htbeta` slab.
    pub row_htbeta_rows: usize,
    /// Framed `G_{ij} ⊗ W_{ij}` co-occurrence factors: bytes (0 on legacy).
    pub frame_blocks_bytes: usize,
    /// Sum of the data-category bytes above — the per-solve re-upload total.
    pub total_bytes: usize,
}

impl std::fmt::Display for SaePcgOperandReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mib = |b: usize| b as f64 / (1024.0 * 1024.0);
        write!(
            f,
            "sae-pcg operand upload [{lane}]: total={total:.1}MiB k={k} \
             a_phi={pairs}pairs/{aphi:.1}MiB local_jac={lj:.1}MiB smooth={sm:.1}MiB \
             sparse_g={sg:.1}MiB row_htbeta={rh:.1}MiB({rhrows}rows) frame_blocks={fb:.1}MiB",
            lane = if self.framed {
                "framed-dense"
            } else {
                "legacy-sparse"
            },
            total = mib(self.total_bytes),
            k = self.beta_dim,
            pairs = self.a_phi_pairs,
            aphi = mib(self.a_phi_bytes),
            lj = mib(self.local_jac_bytes),
            sm = mib(self.smooth_bytes),
            sg = mib(self.sparse_g_bytes),
            rh = mib(self.row_htbeta_bytes),
            rhrows = self.row_htbeta_rows,
            fb = mib(self.frame_blocks_bytes),
        )
    }
}

impl DeviceSaePcgData {
    /// #1017/#2230 residency measurement — see [`SaePcgOperandReport`]. Sums the
    /// resident host operands by category; no device contact.
    pub fn operand_byte_report(&self) -> SaePcgOperandReport {
        const PAIR: usize = std::mem::size_of::<(usize, f64)>();
        const F64: usize = std::mem::size_of::<f64>();

        let a_phi_pairs: usize = self.a_phi.iter().map(|r| r.len()).sum();
        let local_jac_elems: usize = self.local_jac.iter().map(|r| r.len()).sum();
        let smooth_bytes = self
            .smooth_blocks
            .iter()
            .map(|b| b.factor_a.len())
            .sum::<usize>()
            * F64;
        let sparse_g_bytes = self
            .sparse_g_blocks
            .iter()
            .map(|b| b.data.len())
            .sum::<usize>()
            * F64;

        let (framed, row_htbeta_bytes, row_htbeta_rows, frame_blocks_bytes) =
            if let Some(frame) = self.frame.as_ref() {
                let rh_elems: usize = frame.row_htbeta.iter().map(|r| r.len()).sum();
                let rh_rows = frame.row_htbeta.iter().filter(|r| !r.is_empty()).count();
                let fb_elems: usize = frame
                    .frame_blocks
                    .iter()
                    .map(|b| b.g.len() + b.w.len())
                    .sum();
                (true, rh_elems * F64, rh_rows, fb_elems * F64)
            } else {
                (false, 0, 0, 0)
            };

        let a_phi_bytes = a_phi_pairs * PAIR;
        let local_jac_bytes = local_jac_elems * F64;
        SaePcgOperandReport {
            framed,
            beta_dim: self.beta_dim,
            a_phi_pairs,
            a_phi_bytes,
            local_jac_elems,
            local_jac_bytes,
            smooth_bytes,
            sparse_g_bytes,
            row_htbeta_bytes,
            row_htbeta_rows,
            frame_blocks_bytes,
            total_bytes: a_phi_bytes
                + local_jac_bytes
                + smooth_bytes
                + sparse_g_bytes
                + row_htbeta_bytes
                + frame_blocks_bytes,
        }
    }
}

impl BetaPenaltyOp for SparseBlockKroneckerPenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        let p = self.p;
        for blk in &self.blocks {
            let (m_i, m_j) = blk.data.dim();
            for li in 0..m_i {
                let gi_base = (blk.row_off + li) * p;
                for lj in 0..m_j {
                    let a_ij = blk.data[[li, lj]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    let gj_base = (blk.col_off + lj) * p;
                    for oc in 0..p {
                        y[gi_base + oc] += a_ij * x[gj_base + oc];
                    }
                }
            }
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        self.matvec(beta, out);
    }

    fn diagonal(&self, diag: &mut [f64]) {
        let p = self.p;
        for blk in &self.blocks {
            // Only on-diagonal `A` blocks (row_off == col_off) carry diagonal
            // mass; their `(li, li)` entries map to `(row_off+li)·p + oc`.
            if blk.row_off != blk.col_off {
                continue;
            }
            let (m_i, m_j) = blk.data.dim();
            let m = m_i.min(m_j);
            for li in 0..m {
                let a_ii = blk.data[[li, li]];
                let gi_base = (blk.row_off + li) * p;
                for oc in 0..p {
                    diag[gi_base + oc] += a_ii;
                }
            }
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        let range = &offsets[id.0];
        let b = range.end - range.start;
        let p = self.p;
        for blk in &self.blocks {
            let (m_i, m_j) = blk.data.dim();
            let row_start = blk.row_off * p;
            let row_end = (blk.row_off + m_i) * p;
            let col_start = blk.col_off * p;
            let col_end = (blk.col_off + m_j) * p;
            if row_end <= range.start
                || row_start >= range.end
                || col_end <= range.start
                || col_start >= range.end
            {
                continue;
            }
            for bi in 0..b {
                let gi = range.start + bi;
                if gi < row_start || gi >= row_end {
                    continue;
                }
                let li = (gi - row_start) / p;
                let oc_i = (gi - row_start) % p;
                for bj in 0..b {
                    let gj = range.start + bj;
                    if gj < col_start || gj >= col_end {
                        continue;
                    }
                    let oc_j = (gj - col_start) % p;
                    if oc_i != oc_j {
                        continue;
                    }
                    let lj = (gj - col_start) / p;
                    out[[bi, bj]] += blk.data[[li, lj]];
                }
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let p = self.p;
        let mut out = Array2::<f64>::zeros((self.k, self.k));
        for blk in &self.blocks {
            let (m_i, m_j) = blk.data.dim();
            for li in 0..m_i {
                let gi_base = (blk.row_off + li) * p;
                for lj in 0..m_j {
                    let a_ij = blk.data[[li, lj]];
                    if a_ij == 0.0 {
                        continue;
                    }
                    let gj_base = (blk.col_off + lj) * p;
                    for oc in 0..p {
                        out[[gi_base + oc, gj_base + oc]] += a_ij;
                    }
                }
            }
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("sparse-block-kronecker-penalty-op-v1");
        hasher.write_usize(self.p);
        hasher.write_usize(self.dim_a);
        hasher.write_usize(self.k);
        hasher.write_usize(self.blocks.len());
        for blk in &self.blocks {
            hasher.write_usize(blk.row_off);
            hasher.write_usize(blk.col_off);
            hasher.write_f64_array2(&blk.data);
        }
    }
}

/// One co-occurring `(atom_i, atom_j)` block of the **frame-factored** data-fit
/// Gauss–Newton β-Hessian (issue #972 / #977 T1). Carries the basis-space Gram
/// `g` (`m_i × m_j`) AND the per-pair frame output factor `w = U_iᵀ U_j`
/// (`r_i × r_j`); the contributed Hessian sub-block is the Kronecker product
/// `g ⊗ w`.
#[derive(Debug, Clone)]
pub struct FactoredFrameGBlock {
    /// Atom index of the row factor (selects rank `r_i` and β offset).
    pub atom_i: usize,
    /// Atom index of the column factor (selects rank `r_j` and β offset).
    pub atom_j: usize,
    /// Basis-space coupling `G_{ij}` (`m_i × m_j`).
    pub g: Array2<f64>,
    /// Frame output factor `U_iᵀ U_j` (`r_i × r_j`). For `i == j` with an
    /// orthonormal frame this is `I_{r_i}` (the clean within-atom `g ⊗ I_r`
    /// collapse); across atoms it is the dense principal-angle cosine matrix
    /// between the two frames.
    pub w: Array2<f64>,
}

/// Frame-factored data-fit Gauss–Newton β-Hessian operator (#972 / #977 T1):
/// the `Σ_k M_k·r_k` reduced-border analogue of [`SparseBlockKroneckerPenaltyOp`].
///
/// When every atom's decoder `B_k = C_k U_kᵀ` is profiled onto a Grassmann
/// frame `U_k ∈ St(p, r_k)`, the border carries only the shape coefficients
/// `C_k` (`M_k · r_k` entries) instead of the full `B_k` (`M_k · p`). The data
/// Gram in this reduced space is, for the isotropic likelihood,
/// `H[(i,li,a),(j,lj,b)] = G_{ij}[li,lj] · (U_iᵀ U_j)[a,b]` — within an atom the
/// orthonormal frame gives `U_iᵀU_i = I_{r_i}` and the block is the clean
/// `G ⊗ I_r` collapse; across co-active atoms the frames do not share a basis
/// so the output factor is the dense `U_iᵀU_j`.
///
/// The β layout is `μ`-major / frame-minor with a **variable** per-atom width
/// `r_k`: the index of (atom `k`, basis `li`, frame coord `a`) is
/// `offset[k] + li·r_k + a`, where `offset` is the prefix sum of `M_k · r_k`.
/// With every `r_k = p` and `U_k = I_p` this reproduces
/// [`SparseBlockKroneckerPenaltyOp`] exactly (a unit test pins the reduction),
/// so it is a strict generalization, not a separate code path.
pub struct FactoredFrameKroneckerOp {
    /// Per-atom frame rank `r_k` (the factored output width).
    pub ranks: Vec<usize>,
    /// Per-atom basis size `M_k`.
    pub basis_sizes: Vec<usize>,
    /// Per-atom β offset (prefix sum of `M_k · r_k`); `offsets[k]` is the start
    /// of atom `k`'s `C_k` block, `offsets[n_atoms]` the total dim.
    pub offsets: Vec<usize>,
    /// Total reduced β dimension `Σ_k M_k · r_k`.
    pub dim: usize,
    /// Non-empty co-occurring `(atom_i, atom_j)` blocks.
    pub blocks: Vec<FactoredFrameGBlock>,
}

impl FactoredFrameKroneckerOp {
    /// Build from per-atom ranks + basis sizes and the co-occurring blocks.
    /// Computes the β offsets (prefix sum of `M_k·r_k`) and validates that each
    /// block's `g`/`w` shapes match the atoms' `(M, r)`.
    pub fn new(
        ranks: Vec<usize>,
        basis_sizes: Vec<usize>,
        blocks: Vec<FactoredFrameGBlock>,
    ) -> Result<Self, String> {
        if ranks.len() != basis_sizes.len() {
            return Err(format!(
                "FactoredFrameKroneckerOp: {} ranks but {} basis sizes",
                ranks.len(),
                basis_sizes.len()
            ));
        }
        let n_atoms = ranks.len();
        let mut offsets = Vec::with_capacity(n_atoms + 1);
        let mut acc = 0usize;
        for k in 0..n_atoms {
            offsets.push(acc);
            acc += basis_sizes[k] * ranks[k];
        }
        offsets.push(acc);
        let dim = acc;
        for blk in &blocks {
            if blk.atom_i >= n_atoms || blk.atom_j >= n_atoms {
                return Err(format!(
                    "FactoredFrameKroneckerOp: block atom indices ({}, {}) out of range (n_atoms = {n_atoms})",
                    blk.atom_i, blk.atom_j
                ));
            }
            if blk.g.dim() != (basis_sizes[blk.atom_i], basis_sizes[blk.atom_j]) {
                return Err(format!(
                    "FactoredFrameKroneckerOp: block ({}, {}) g has shape {:?} but expected ({}, {})",
                    blk.atom_i,
                    blk.atom_j,
                    blk.g.dim(),
                    basis_sizes[blk.atom_i],
                    basis_sizes[blk.atom_j]
                ));
            }
            if blk.w.dim() != (ranks[blk.atom_i], ranks[blk.atom_j]) {
                return Err(format!(
                    "FactoredFrameKroneckerOp: block ({}, {}) w has shape {:?} but expected ({}, {})",
                    blk.atom_i,
                    blk.atom_j,
                    blk.w.dim(),
                    ranks[blk.atom_i],
                    ranks[blk.atom_j]
                ));
            }
        }
        Ok(Self {
            ranks,
            basis_sizes,
            offsets,
            dim,
            blocks,
        })
    }

}

impl BetaPenaltyOp for FactoredFrameKroneckerOp {
    fn dim(&self) -> usize {
        self.dim
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        for blk in &self.blocks {
            let r_i = self.ranks[blk.atom_i];
            let r_j = self.ranks[blk.atom_j];
            let off_i = self.offsets[blk.atom_i];
            let off_j = self.offsets[blk.atom_j];
            let (m_i, m_j) = blk.g.dim();
            for li in 0..m_i {
                let yi_base = off_i + li * r_i;
                for lj in 0..m_j {
                    let g = blk.g[[li, lj]];
                    if g == 0.0 {
                        continue;
                    }
                    let xj_base = off_j + lj * r_j;
                    // y_block[li, a] += g · Σ_b w[a, b] · x_block[lj, b]
                    for a in 0..r_i {
                        let mut acc = 0.0;
                        for b in 0..r_j {
                            acc += blk.w[[a, b]] * x[xj_base + b];
                        }
                        y[yi_base + a] += g * acc;
                    }
                }
            }
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        self.matvec(beta, out);
    }

    fn diagonal(&self, diag: &mut [f64]) {
        for blk in &self.blocks {
            // Only on-diagonal atom blocks carry diagonal mass; the entry at
            // (atom k, basis li, coord a) is g[li,li]·w[a,a].
            if blk.atom_i != blk.atom_j {
                continue;
            }
            let r = self.ranks[blk.atom_i];
            let off = self.offsets[blk.atom_i];
            let (m_i, m_j) = blk.g.dim();
            let m = m_i.min(m_j);
            for li in 0..m {
                let gii = blk.g[[li, li]];
                let base = off + li * r;
                for a in 0..r {
                    diag[base + a] += gii * blk.w[[a, a]];
                }
            }
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        // Dense sub-block over the β index range `offsets[id.0]`. Mirror the
        // global (i,a) ↔ (j,b) coupling, keeping only indices inside the range.
        let range = &offsets[id.0];
        let b_dim = range.end - range.start;
        for blk in &self.blocks {
            let r_i = self.ranks[blk.atom_i];
            let r_j = self.ranks[blk.atom_j];
            let off_i = self.offsets[blk.atom_i];
            let off_j = self.offsets[blk.atom_j];
            let (m_i, m_j) = blk.g.dim();
            for li in 0..m_i {
                for a in 0..r_i {
                    let gi = off_i + li * r_i + a;
                    if gi < range.start || gi >= range.end {
                        continue;
                    }
                    let bi = gi - range.start;
                    for lj in 0..m_j {
                        let g = blk.g[[li, lj]];
                        if g == 0.0 {
                            continue;
                        }
                        for b in 0..r_j {
                            let gj = off_j + lj * r_j + b;
                            if gj < range.start || gj >= range.end {
                                continue;
                            }
                            let bj = gj - range.start;
                            if bi < b_dim && bj < b_dim {
                                out[[bi, bj]] += g * blk.w[[a, b]];
                            }
                        }
                    }
                }
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let mut out = Array2::<f64>::zeros((self.dim, self.dim));
        for blk in &self.blocks {
            let r_i = self.ranks[blk.atom_i];
            let r_j = self.ranks[blk.atom_j];
            let off_i = self.offsets[blk.atom_i];
            let off_j = self.offsets[blk.atom_j];
            let (m_i, m_j) = blk.g.dim();
            for li in 0..m_i {
                for lj in 0..m_j {
                    let g = blk.g[[li, lj]];
                    if g == 0.0 {
                        continue;
                    }
                    for a in 0..r_i {
                        let gi = off_i + li * r_i + a;
                        for b in 0..r_j {
                            let gj = off_j + lj * r_j + b;
                            out[[gi, gj]] += g * blk.w[[a, b]];
                        }
                    }
                }
            }
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("factored-frame-kronecker-op-v1");
        hasher.write_usize(self.dim);
        for &r in &self.ranks {
            hasher.write_usize(r);
        }
        for &m in &self.basis_sizes {
            hasher.write_usize(m);
        }
        hasher.write_usize(self.blocks.len());
        for blk in &self.blocks {
            hasher.write_usize(blk.atom_i);
            hasher.write_usize(blk.atom_j);
            hasher.write_f64_array2(&blk.g);
            hasher.write_f64_array2(&blk.w);
        }
    }
}

/// Composite penalty: sum of multiple `BetaPenaltyOp` operators.
pub struct CompositePenaltyOp {
    /// Full β dimension `K`.
    pub k: usize,
    /// Component operators, each contributing additively.
    pub ops: Vec<Arc<dyn BetaPenaltyOp>>,
}

impl BetaPenaltyOp for CompositePenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        // The reduced-Schur PCG matvec applies this composite ONCE PER CG
        // ITERATION as the penalty prologue `y += (H_ββ) x`. At the K=32k
        // manifold-SAE border the composite is a leading run of per-atom
        // Kronecker smooth penalties (`λ S_k ⊗ I_{r_k}`, one per atom, over
        // DISJOINT β blocks) followed by the cross-atom data-fit op and any
        // dense analytic tail — and this whole sum ran SERIALLY while the
        // point-elimination row term already fanned across all cores, so it was
        // the prologue's Amdahl ceiling on the wide border.
        //
        // Fan the leading run of mutually-disjoint, sorted, contiguous-or-gapped
        // output-range operators across rayon workers: each writes ONLY its own
        // `y[start..end]` sub-slice (no cross-thread aliasing), then the
        // remaining (`None`-range / overlapping) operators run SERIALLY in
        // original order. Because every prefix index is touched by exactly one
        // prefix operator and all prefix work happens-before the serial tail,
        // each output index accumulates in the SAME order as the fully-serial
        // loop — the result is BIT-IDENTICAL, not merely deterministic. Stay
        // serial when already inside a rayon worker (the topology race / nested
        // matvec) to avoid oversubscription — the same guard the row loop uses.
        let mut prefix_len = 0usize;
        let mut prev_end = 0usize;
        if rayon::current_thread_index().is_none() {
            for op in &self.ops {
                match op.output_range() {
                    Some(r) if r.start >= prev_end && r.end > r.start && r.end <= y.len() => {
                        prev_end = r.end;
                        prefix_len += 1;
                    }
                    _ => break,
                }
            }
        }
        // Only worth the fan-out when there is real disjoint work: at least two
        // blocks and a covered width past the same border threshold the dense
        // prologue uses. Otherwise fall through to the plain serial sum.
        if prefix_len >= 2 && prev_end >= SCHUR_PROLOGUE_PARALLEL_K_MIN {
            use rayon::prelude::*;
            // Carve `y` into one mutable sub-slice per prefix operator, skipping
            // any gaps between ranges. Sorted, non-overlapping ranges make this
            // a single left-to-right walk of `split_at_mut`.
            let mut subslices: Vec<&mut [f64]> = Vec::with_capacity(prefix_len);
            {
                let mut consumed = 0usize;
                let mut rest: &mut [f64] = y;
                for op in &self.ops[..prefix_len] {
                    let r = op.output_range().expect("prefix op has an output range");
                    let (_, after_gap) = rest.split_at_mut(r.start - consumed);
                    let (block, tail) = after_gap.split_at_mut(r.end - r.start);
                    subslices.push(block);
                    rest = tail;
                    consumed = r.end;
                }
            }
            self.ops[..prefix_len]
                .par_iter()
                .zip(subslices.par_iter_mut())
                .for_each(|(op, y_local)| op.matvec_local(x, y_local));
            for op in &self.ops[prefix_len..] {
                op.matvec(x, y);
            }
        } else {
            for op in &self.ops {
                op.matvec(x, y);
            }
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        for op in &self.ops {
            op.gradient(beta, out);
        }
    }

    fn diagonal(&self, diag: &mut [f64]) {
        for op in &self.ops {
            op.diagonal(diag);
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        for op in &self.ops {
            op.block(id, offsets, out);
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let mut out = Array2::<f64>::zeros((self.k, self.k));
        for op in &self.ops {
            let dense = op.to_dense();
            out += &dense;
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        hasher.write_str("composite-penalty-op-v1");
        hasher.write_usize(self.k);
        hasher.write_usize(self.ops.len());
        for op in &self.ops {
            op.fingerprint(hasher);
        }
    }
}

/// Adapts a closure-based matrix-free `H_ββ` operator (from
/// [`ArrowSchurSystem::set_shared_beta_operator`]) to the `BetaPenaltyOp` trait.
///
/// `diagonal` holds the precomputed `diag(H_ββ)` supplied alongside the matvec;
/// `to_dense` falls back to probing all `K` canonical basis vectors.
pub struct MatvecDiagPenaltyOp {
    pub(crate) k: usize,
    pub(crate) matvec: SharedBetaMatvec,
    pub(crate) diagonal_vec: Array1<f64>,
}

impl MatvecDiagPenaltyOp {
    pub fn new(k: usize, matvec: SharedBetaMatvec, diagonal_vec: Array1<f64>) -> Self {
        assert_eq!(diagonal_vec.len(), k);
        Self {
            k,
            matvec,
            diagonal_vec,
        }
    }
}

impl BetaPenaltyOp for MatvecDiagPenaltyOp {
    fn dim(&self) -> usize {
        self.k
    }

    fn matvec(&self, x: &[f64], y: &mut [f64]) {
        let x_arr = Array1::from_iter(x.iter().copied());
        let mut out = Array1::<f64>::zeros(self.k);
        (self.matvec)(x_arr.view(), &mut out);
        for a in 0..self.k {
            y[a] += out[a];
        }
    }

    fn gradient(&self, beta: &[f64], out: &mut [f64]) {
        let beta_arr = Array1::from_iter(beta.iter().copied());
        let mut hb = Array1::<f64>::zeros(self.k);
        (self.matvec)(beta_arr.view(), &mut hb);
        for a in 0..self.k {
            out[a] += hb[a];
        }
    }

    fn diagonal(&self, diag: &mut [f64]) {
        for j in 0..self.k.min(diag.len()) {
            diag[j] += self.diagonal_vec[j];
        }
    }

    fn block(&self, id: BetaBlockId, offsets: &[Range<usize>], out: &mut Array2<f64>) {
        // Probe each basis vector in the block range to extract the sub-block.
        let range = &offsets[id.0];
        let b = range.end - range.start;
        let mut probe = Array1::<f64>::zeros(self.k);
        for bj in 0..b {
            probe.fill(0.0);
            probe[range.start + bj] = 1.0;
            let mut col = Array1::<f64>::zeros(self.k);
            (self.matvec)(probe.view(), &mut col);
            for bi in 0..b {
                out[[bi, bj]] += col[range.start + bi];
            }
        }
    }

    fn to_dense(&self) -> Array2<f64> {
        let k = self.k;
        let mut out = Array2::<f64>::zeros((k, k));
        let mut probe = Array1::<f64>::zeros(k);
        for j in 0..k {
            probe.fill(0.0);
            probe[j] = 1.0;
            let mut col = Array1::<f64>::zeros(k);
            (self.matvec)(probe.view(), &mut col);
            for i in 0..k {
                out[[i, j]] = col[i];
            }
        }
        out
    }

    fn fingerprint(&self, hasher: &mut Fingerprinter) {
        // The matvec closure cannot be hashed by content; the precomputed
        // diagonal is the operator's stable defining proxy (it is recomputed
        // alongside the matvec each time the operator is installed).
        hasher.write_str("matvec-diag-penalty-op-v1");
        hasher.write_usize(self.k);
        for &value in self.diagonal_vec.iter() {
            hasher.write_f64(value);
        }
    }
}