Struct ArrowSchurSystem 

pub struct ArrowSchurSystem {

    pub rows: Vec<ArrowRowBlock>,
    pub hbb: Array2<f64>,
    pub hbb_matvec: Option<SharedBetaMatvec>,
    pub htbeta_matvec: Option<RowHtbetaMatvec>,
    pub htbeta_transpose_matvec: Option<RowHtbetaTransposeMatvec>,
    pub htbeta_dense_supplement: bool,
    pub hbb_diag: Option<Array1<f64>>,
    pub gb: Array1<f64>,
    pub d: usize,
    pub row_dims: Arc<[usize]>,
    pub row_offsets: Arc<[usize]>,
    pub k: usize,
    pub manifold_mode_fingerprint: u64,
    pub row_hessian_fingerprint: u64,
    pub analytic_row_hessian_fingerprint: u64,
    pub block_offsets: Arc<[Range<usize>]>,
    pub penalty_op: Option<Arc<dyn BetaPenaltyOp>>,
    pub device_sae_pcg: Option<Arc<DeviceSaePcgData>>,
    pub cross_row_penalties: Vec<CrossRowLatentPenalty>,
    pub row_gauge_deflation: Option<ArrowRowGaugeDeflation>,
    pub ibp_cross_row: Option<IbpCrossRowSource>,
}

Bordered (t, β) Newton system with arrow structure.

The β-block is held as a dense K × K Hessian H_ββ plus a K-length gradient g_β for direct BA modes. Large-scale inexact BA callers may additionally install a matrix-free H_ββ x operator and diagonal via ArrowSchurSystem::set_shared_beta_operator; the InexactPCG mode then avoids dense Schur formation/factorization. The t-block is a Vec<ArrowRowBlock> of length N.

Construction is the driver’s responsibility: the driver

1. evaluates Φ(t) and the radial jet ∂Φ/∂t (the latter via gam_terms::latent::LatentCoordValues::design_gradient_wrt_t);
2. forms the working-weighted Gauss–Newton blocks H_tt^(i) += (g_i β)(g_i β)^T, H_tβ^(i) += (g_i β) ⊗ Φ_i, H_ββ += Φ^T W Φ + Σ_k λ_k S_k;
3. calls ArrowSchurSystem::add_analytic_penalty_contributions to fold row-block Psi-tier analytic penalties (ARDPenalty, SparsityPenalty) into H_tt^(i) and Beta-tier penalties into H_ββ;
4. calls ArrowSchurSystem::solve to obtain (Δt, Δβ).

Fields

rows: Vec<ArrowRowBlock>

Per-row latent block (length N, each row d × d / d × K / d).

hbb: Array2<f64>

H_ββ, shape (K, K) for direct BA modes; empty when constructed by ArrowSchurSystem::new_matrix_free_shared for PCG-only use.

hbb_matvec: Option<SharedBetaMatvec>

Optional matrix-free H_ββ x operator for large BA Schur PCG.

Direct and Square-Root BA modes still require hbb; InexactPCG uses this operator when present, avoiding dense shared-block storage for SAE-manifold scale K.

htbeta_matvec: Option<RowHtbetaMatvec>

Optional row-local matrix-free multiply for H_tβ^(i) x.

When present, all inner-Schur paths route through this operator instead of indexing the per-row htbeta dense slabs: reduced_rhs_beta, schur_matvec (PCG hot loop), back-substitution, JacobiPreconditioner construction, build_dense_schur_direct, and build_dense_schur_sqrt_ba all call sys_htbeta_apply_row or sys_htbeta_materialize_row. Factor caches retain the operator for IFT/evidence consumers as before.

htbeta_transpose_matvec: Option<RowHtbetaTransposeMatvec>

Optional row-local matrix-free transpose multiply out += H_βt^(i) · v.

The sparse adjoint of Self::htbeta_matvec. When present, the reduced-Schur matvec applies H_βt^(i) directly (sparse scatter) instead of probing the forward operator against K basis vectors. This is the per-row sparse apply that lifts the O(K) column-probe in the GPU PCG and streaming Schur paths to O(m_i · p) per row. Installed in lock-step with htbeta_matvec by Self::set_row_htbeta_operator.

htbeta_dense_supplement: bool

Whether rows[*].htbeta contains a dense contribution that must be added on top of the matrix-free row operator.

hbb_diag: Option<Array1<f64>>

Optional diagonal of the matrix-free shared block, used by the Schur-Jacobi preconditioner in the Agarwal-style PCG path.

gb: Array1<f64>

g_β, shape (K,).

d: usize

Maximum per-row latent dimensionality across all rows.

For homogeneous systems (all rows have the same dim) this equals the common per-row d. For heterogeneous systems (e.g. sparse SAE rows where JumpReLU / TopK / sparsemax active sets vary per observation) this is max_i row_dims[i]. Per-row code should use row.htt.nrows() or row_dims[i]; d is an upper bound for scratch-buffer sizing.

row_dims: Arc<[usize]>

Per-row latent dimensionality: row_dims[i] == rows[i].htt.nrows().

For homogeneous systems row_dims[i] == d for all i.

row_offsets: Arc<[usize]>

Flat-buffer row offsets for the delta_t vector produced by Self::solve / solve_arrow_newton_step_core.

row_offsets[i] is the start index for row i’s slice in delta_t; row_offsets[n] is the total delta_t length. For homogeneous systems row_offsets[i] == i * d.

k: usize

β dimensionality K.

manifold_mode_fingerprint: u64

Geometry tag for the row-local latent blocks after optional Riemannian projection. Euclidean/no-op geometry uses the sentinel.

row_hessian_fingerprint: u64

Structural/value tag for row-local Hessian factors and their Schur inputs. Stale caches must be rejected when row-dependent Hessian penalties or cross-blocks change.

analytic_row_hessian_fingerprint: u64

Registry-side tag for row-dependent analytic-penalty Hessian inputs. Combined with the materialized row blocks in Self::current_row_hessian_fingerprint.

block_offsets: Arc<[Range<usize>]>

Term-block column ranges for the block-Jacobi Schur preconditioner.

Each entry r means that indices r.start..r.end belong to one coefficient block (a GAM term or a custom parameter family from ParameterBlockSpec). When populated via Self::set_block_offsets, the Jacobi preconditioner inverts the full b × b Schur block for each term instead of only its diagonal.

The default (empty slice) causes JacobiPreconditioner to fall back to pure scalar diagonal inversion, preserving the pre-#283 behaviour.

penalty_op: Option<Arc<dyn BetaPenaltyOp>>

Optional matrix-free penalty-side H_ββ operator (#296).

When set, all hot paths (schur_matvec, build_dense_schur_*, JacobiPreconditioner, quadratic-form reduction) route through this operator instead of the dense hbb accumulator, enabling BlockPenaltyOp / KroneckerPenaltyOp to skip the O(K²) dense materialisation for structured smoothness penalties.

When None, those paths fall back to wrapping hbb in a transient DensePenaltyOp — identical observable behaviour, no new allocation hot-path cost for callers that have not opted in.

device_sae_pcg: Option<Arc<DeviceSaePcgData>>

Device-uploadable SAE Kronecker data for CUDA-resident reduced PCG.

The generic matrix-free closures remain the authoritative CPU path. This descriptor is installed only when SAE assembly has a matching CUDA sparse representation for both H_tβ and H_ββ.

cross_row_penalties: Vec<CrossRowLatentPenalty>

Registered Psi-tier analytic penalties whose Hessian couples distinct latent rows (non-row-block-diagonal), captured by Self::add_analytic_penalty_contributions.

These penalties (TotalVariationPenalty, SheafConsistencyPenalty, block-orthogonality, …) produce off-row Hessian blocks ∂²P/∂t_i∂t_j (i ≠ j) that the arrow elimination — which assumes each H_tt^(i) is independent of every other row — cannot represent. Their gradient is still folded into g_t exactly like every other Psi penalty; only their curvature is held here, applied during the solve as a full-latent Hessian-vector product P_cross · Δt against the penalty’s psd_majorizer_hvp. When this vector is non-empty, [solve_arrow_newton_step_artifacts] auto-selects the matrix-free full-system PCG path (arrow block-diagonal inverse as preconditioner) instead of the exact one-shot Schur elimination. When empty, the system is purely row-block-diagonal and the exact Schur path is unchanged.

row_gauge_deflation: Option<ArrowRowGaugeDeflation>

Optional row-local gauge directions for evidence-only Faddeev-Popov deflation of an otherwise non-PD H_tt row block.

These vectors live in each row’s actual chart block, so compact SAE rows and dense rows share the same factorization path. Ordinary Newton solves ignore them; only undamped evidence factors with tolerate_ill_conditioning set may stiffen a gauge-explained row direction.

ibp_cross_row: Option<IbpCrossRowSource>

Optional exact cross-row IBP low-rank source (#1038). When set, the factorization downdates the per-row logit-slot self term and layers the exact rank-R Woodbury correction onto the evidence cache (value, log-determinant, and θ/ρ-adjoint together). None for all non-IBP systems — the row-block-diagonal arrow path is then unchanged.

Implementations

impl ArrowSchurSystem

pub fn new(n: usize, d: usize, k: usize) -> Self

Allocate an empty BA reduced-camera-system instance sized (N point/latent rows × d, K shared decoder parameters).

pub fn new_with_empty_hbb_and_htbeta_cols( n: usize, d: usize, k: usize, htbeta_cols: usize, ) -> Self

Allocate an arrow system with no dense shared H_ββ block and with per-row dense H_tβ slabs allocated at htbeta_cols columns.

pub fn new_with_hbb(n: usize, d: usize, k: usize, hbb: Array2<f64>) -> Self

Allocate an arrow system using a caller-owned dense shared-block buffer. The buffer must already have shape (k, k) and is zeroed in place before use so callers can recycle it across assemblies without changing numerics.

pub fn new_with_hbb_and_htbeta_cols( n: usize, d: usize, k: usize, hbb: Array2<f64>, htbeta_cols: usize, ) -> Self

Allocate an arrow system with a caller-owned dense shared-block buffer and per-row dense H_tβ slabs allocated at htbeta_cols columns.

pub fn new_matrix_free_shared<F>( n: usize, d: usize, k: usize, matvec: F, diag: Array1<f64>, ) -> Self

where F: for<'a> Fn(ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync + 'static,

Allocate an arrow system whose shared H_ββ block is supplied only as a matrix-free operator for large BA InexactPCG.

Direct and Square-Root BA modes require dense hbb and must not be used with this constructor. The row-local H_tβ slabs remain explicit; a future MegBA backend can replace those slab operations behind BatchedBlockSolver.

pub fn new_with_per_row_dims_empty_hbb_and_htbeta_cols( per_row_dims: Vec<usize>, k: usize, htbeta_cols: usize, ) -> Self

Allocate a heterogeneous-row arrow system with no dense shared H_ββ block and with row H_tβ slabs allocated at htbeta_cols columns.

pub fn new_with_per_row_dims_and_hbb_and_htbeta_cols( per_row_dims: Vec<usize>, k: usize, hbb: Array2<f64>, htbeta_cols: usize, ) -> Self

Allocate a heterogeneous-row system using a caller-owned dense shared block and row H_tβ slabs allocated at htbeta_cols columns.

pub fn set_row_gauge_deflation(&mut self, deflation: ArrowRowGaugeDeflation)

pub fn set_ibp_cross_row_source(&mut self, source: IbpCrossRowSource)

Register the exact cross-row IBP low-rank source (#1038). The assembly passes the per-column D-coefficients (cross_row_d) and the (global latent index, atom, z'_ik) entries built from z_jac; the factorization then carries the exact rank-R Woodbury (value + log-determinant + θ/ρ-adjoint) on the evidence cache. An empty source (r == 0 or no entries) is treated as absent so the row-block-diagonal path is unchanged.

pub fn n(&self) -> usize

Number of BA point/latent rows N.

pub fn compute_row_hessian_fingerprint(&self) -> u64

Recompute the row-system fingerprint from the currently materialized row blocks, cross-blocks, and shared-block diagonal.

pub fn current_row_hessian_fingerprint(&self) -> u64

Current effective row-system fingerprint, including the materialized row blocks and any registry metadata captured while folding analytic penalties into the system.

pub fn refresh_row_hessian_fingerprint(&mut self)

Store the current row-system fingerprint on the system.

This is intentionally explicit and expensive. Cache and evidence callers use Self::current_row_hessian_fingerprint at the point they need the value, after assembly has populated the system, instead of hashing each intermediate construction/mutation step.

pub fn set_shared_beta_operator<F>(&mut self, matvec: F, diag: Array1<f64>)

where F: for<'a> Fn(ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync + 'static,

Install a matrix-free shared-block operator for Agarwal-style inexact Schur PCG.

diag must be the diagonal of the same H_ββ operator and is used for the Schur-Jacobi preconditioner. This is the BA “large camera system” path mapped to large decoder coefficient blocks.

pub fn activate_dense_htbeta_supplement(&mut self)

Mark the dense per-row cross-block slabs as active supplements to the installed matrix-free row operator.

pub fn set_row_htbeta_operator<F, T>(&mut self, forward: F, transpose: T)

where F: for<'a> Fn(usize, ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync + 'static, T: for<'a> Fn(usize, ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync + 'static,

Install a matrix-free per-row cross-block operator and its sparse adjoint.

forward must write out = H_tβ^(row) x for out.len() == d and x.len() == K. transpose must add H_βt^(row) v into out for out.len() == K and v.len() == d (the sparse scatter adjoint).

When installed, the forward operator is used during the Newton solve (inside reduced_rhs_beta, schur_matvec, back-substitution, and JacobiPreconditioner construction) and afterwards by IFT/evidence predictors. Per-row htbeta slabs in ArrowRowBlock may be left zero-sized when this operator is installed — all inner-Schur paths route through the matvec instead of indexing the dense block. The transpose operator lets the reduced-Schur matvec apply H_βt^(row) directly (O(m_i · p)) instead of probing forward against K basis vectors.

pub fn set_block_offsets(&mut self, offsets: Arc<[Range<usize>]>)

Register term-block column ranges for the block-Jacobi Schur preconditioner.

Each Range<usize> covers the columns of one GAM term (or custom parameter family) in the shared β vector. The ranges must be non-overlapping, sorted, and their union must cover 0..k.

Call this after building the system and before Self::solve / Self::solve_with_options whenever the solver will use ArrowSolverMode::InexactPCG. Absent a call, the preconditioner falls back to scalar diagonal Jacobi (the pre-#283 behaviour).

The same plumbing is compatible with #287 (custom ParameterBlockSpec families): callers from that path simply supply ranges derived from their own block layout.

pub fn set_penalty_op(&mut self, op: Arc<dyn BetaPenaltyOp>)

Install a matrix-free penalty-side H_ββ operator (#296).

When set, all hot paths (schur_matvec, build_dense_schur_*, JacobiPreconditioner, quadratic-form reduction) route through this operator instead of the dense hbb accumulator, enabling BlockPenaltyOp / KroneckerPenaltyOp to avoid O(K²) allocation for structured smoothness penalties.

pub fn set_device_sae_pcg_data(&mut self, data: DeviceSaePcgData)

pub fn effective_penalty_op(&self) -> Arc<dyn BetaPenaltyOp>

Return the effective penalty operator: the installed penalty_op if present, otherwise a DensePenaltyOp wrapping the current hbb.

Note: when penalty_op is None, this clones hbb into a new DensePenaltyOp. Callers in hot loops should call this once and store the result, not call it per-iteration.

pub fn add_analytic_penalty_contributions( &mut self, registry: &AnalyticPenaltyRegistry, target_t: ArrayView1<'_, f64>, target_beta: ArrayView1<'_, f64>, rho_global: ArrayView1<'_, f64>, ) -> Result<(), ArrowSchurError>

Fold analytic-penalty contributions into the appropriate blocks.

BA source mapping: these are extra prior/regularization normal-equation terms before point elimination, the same place Ceres/g2o attach robust priors or gauge-fixing constraints.

Composition path. Each registered AnalyticPenaltyKind is queried for grad_target (added to g_t or g_β) and then for hessian_diag first. Diagonal penalties (ARD and the shipped sparsity kernels) are injected directly. The row-block-only Psi-tier penalties are ARDPenalty, SparsityPenalty, SoftmaxAssignmentSparsity, IBPAssignment, RowPrecisionPrior, ParametricRowPrecisionPrior, and ScadMcpPenalty. Their d × d per-row Hessian folds into rows[i].htt, so the exact arrow Schur elimination (N independent d × d row solves) represents them exactly. Dense Beta-tier penalties still fall back to hvp probes against the canonical basis vectors for β.

Cross-row Psi penalties. Penalties whose Hessian couples distinct latent rows — TotalVariationPenalty, SheafConsistencyPenalty, block-orthogonality, … — produce off-row blocks ∂²P/∂t_i∂t_j (i ≠ j) that the arrow elimination cannot store, since it assumes each H_tt^(i) is independent of every other row. These are handled without any approximation: their gradient is folded into g_t exactly as for every other Psi penalty (grad_target → g_t), and their full curvature is captured into Self::cross_row_penalties as a matrix-free operator. At solve time, K = K0 + P_cross where K0 is the block-diagonal arrow operator and P_cross · Δt = Σ_p ρ_p · psd_majorizer_hvp_p(t, Δt) is the cross-row penalty Hessian applied to the full flat latent vector. The presence of any captured cross-row penalty auto-routes Self::solve through the matrix-free full-system PCG path (the exact arrow block-diagonal inverse K0⁻¹ is the preconditioner M⁻¹); a purely row-block-diagonal system keeps the exact one-shot Schur path unchanged. No new flag is involved — the route is selected from the captured penalty set alone (magic by default).

target_t is the full flat latent-coordinate vector (row-major, N·d entries) at the current iterate; target_beta is the current β. rho is the global ρ vector restricted to each penalty’s local slice by AnalyticPenaltyRegistry::rho_layout.

pub fn apply_riemannian_latent_geometry(&mut self, latent: &LatentCoordValues)

Convert row-local Euclidean latent blocks to Riemannian tangent blocks.

This is the only arrow-Schur algebra change needed for manifold latents: g_t, H_tt, and each H_tβ column are projected to T_{t_i}M, while the shared β block and Schur structure remain untouched. Embedded constrained manifolds carry a pinned normal block so the existing ambient Cholesky factorization still works; all RHS terms live in the tangent space, so the solved update retracts cleanly.

pub fn solve( &self, ridge_t: f64, ridge_beta: f64, ) -> Result<(Array1<f64>, Array1<f64>, PcgDiagnostics), ArrowSchurError>

Schur-eliminate the per-row latent block and solve for (Δt, Δβ, diag).

This uses ArrowSolveOptions::automatic: BA dense RCS for K <= 2000, and Agarwal-style inexact Schur PCG above that size. Call ArrowSchurSystem::solve_with_options to force Square-Root BA or a specific inexact solve policy.

Returns (delta_t, delta_beta, PcgDiagnostics) with delta_t flat row-major of length N · d and delta_beta of length K. The sign convention matches solve_newton_direction_dense: the returned increments satisfy the bordered system with RHS [-g_t; -g_β], i.e. they are the negated solutions of the standard Newton-direction formulation. PcgDiagnostics is zero-valued for the Direct path and carries live counters (PCG iters, ridge escalations, residual) for InexactPCG.

ridge_t and ridge_beta are nonnegative diagonal regularizers added to the latent and β blocks respectively before factorization — used by the LM damping outer wrapper to recover from near-singular inner steps. Pass 0.0 for both to obtain the unregularized Newton direction.

pub fn solve_with_lm_escalation( &self, ridge_t: f64, ridge_beta: f64, ) -> Result<(Array1<f64>, Array1<f64>, PcgDiagnostics), ArrowSchurError>

Solve with the standard LM-style ridge escalation: if a per-row H_tt + ridge_t·I Cholesky pivot is non-PD, or the reduced Schur factor fails, geometrically grow both ridges and retry. This is the same Ceres-style proximal correction the Newton driver in run_joint_fit_arrow_schur performs around solve, lifted into the system itself so every entry point (predict OOS reconstruction, single-shot Newton refinement, …) is self-healing against the pathological per-row blocks produced by PCA-seeded latent coordinates on subset / new data — see #163 and #175.

ridge_t / ridge_beta are the caller-nominal Tikhonov ridges; the escalation only adds extra damping on top of them when the factor fails. PCG / AdaptiveCorrection failures are left untouched because they are not factorization-recoverable.

pub fn solve_with_options( &self, ridge_t: f64, ridge_beta: f64, options: &ArrowSolveOptions, ) -> Result<(Array1<f64>, Array1<f64>, PcgDiagnostics), ArrowSchurError>

Solve with an explicit BA Schur mode, returning (Δt, Δβ, PcgDiagnostics).

ArrowSolverMode::Direct is the classic dense reduced-camera-system Cholesky path; ArrowSolverMode::SqrtBA forms the same dense system through Square-Root BA factors; ArrowSolverMode::InexactPCG runs inexact-step LM on the reduced system with Jacobi-preconditioned Steihaug-CG. PcgDiagnostics is zero-valued for Direct/SqrtBA and carries live counters for InexactPCG (iterations, matvec calls, preconditioner escalations, final relative residual, stopping reason).

Trait Implementations

impl Clone for ArrowSchurSystem

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.