Skip to main content

gam_solve/arrow_schur/
prelude.rs

1//! Shared prelude for the arrow-Schur solver: external imports re-exported
2//! crate-wide, the module-level tuning constants, and the matvec function-
3//! pointer type aliases. Every sibling concern module pulls these in through
4//! `use super::*;`, preserving the single-namespace resolution the previous
5//! `include!`-based layout relied on.
6
7pub(crate) use super::reduced_solve::ArrowSchurError;
8pub(crate) use super::system::ArrowRowBlock;
9pub(crate) use faer::Side;
10pub(crate) use gam_linalg::faer_ndarray::{FaerArrayView, FaerEigh, FaerLlt, FaerQr};
11pub(crate) use gam_linalg::triangular::{
12    cholesky_solve_matrix, cholesky_solve_vector, forward_substitution_lower_matrix,
13};
14pub(crate) use gam_runtime::warm_start::Fingerprinter;
15pub(crate) use gam_terms::analytic_penalties::AnalyticPenaltyKind;
16pub(crate) use gam_terms::latent::{LatentCoordValues, LatentManifold};
17pub(crate) use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
18pub(crate) use std::ops::Range;
19pub(crate) use std::sync::Arc;
20
21pub(crate) const DIRECT_SOLVE_MAX_K: usize = 2_000;
22
23pub(crate) const DEFAULT_PCG_MAX_ITERATIONS: usize = 200;
24
25pub(crate) const DEFAULT_PCG_RELATIVE_TOLERANCE: f64 = 1e-4;
26
27/// Absolute floor on the Steihaug-CG residual stopping threshold.
28///
29/// The native PCG criterion is purely relative: `tol = rel_tol · ‖rhs‖`. When
30/// `‖rhs‖` is tiny (degenerate / near-stationary reduced systems) this product
31/// can fall below the roundoff resolution of `metric_norm` (~1e-15 for f64),
32/// so the loop would "converge" on floating-point noise rather than a genuinely
33/// accurate solution. Floor the threshold at 1e-14: above machine epsilon
34/// (~2.2e-16) yet below any practical single-iteration residual reduction, so
35/// well-scaled problems are unaffected while degenerate ones stop cleanly.
36pub(crate) const PCG_ABSOLUTE_TOLERANCE_FLOOR: f64 = 1e-14;
37
38pub(crate) const DEFAULT_TRUST_REGION_RADIUS: f64 = f64::INFINITY;
39
40/// First rung of the geometric proximal-ridge ladder, and the `Default` for
41/// `ArrowSolveOptions::initial_ridge` (so a caller that needs a different
42/// starting damping sets one rather than editing this).
43///
44/// The ladder is only ever escalated on a REJECTED step and never lowered, so
45/// this rung's job is to be small enough that an accepted first attempt is
46/// effectively the undamped Newton step, while still being nonzero so
47/// [`DEFAULT_PROXIMAL_RIDGE_GROWTH`] has something to multiply. Its reach is
48/// what is actually derived: `1e-8 · 10^21 ≈ 1e14`, on
49/// [`DEFAULT_PROXIMAL_MAX_ATTEMPTS`].
50pub const DEFAULT_PROXIMAL_INITIAL_RIDGE: f64 = 1e-8;
51
52pub(crate) const F32_UNIT_ROUNDOFF: f64 = (f32::EPSILON as f64) * 0.5;
53
54pub(crate) const DEFAULT_MIXED_PRECISION_MAX_REFINEMENTS: usize = 6;
55
56pub(crate) const DEFAULT_MIXED_PRECISION_CERTIFICATE_TOLERANCE: f64 = 1e-11;
57
58pub(crate) const DEFAULT_MIXED_PRECISION_KAPPA_MARGIN: f64 = 0.5;
59
60/// Backward-error certificate floor, expressed as a small multiple of f64 epsilon.
61pub(crate) const MIXED_PRECISION_CERTIFICATE_EPSILON_MULTIPLIER: f64 = 64.0;
62
63/// User-supplied kappa margins above this are no stricter than the unit gate.
64pub(crate) const MIXED_PRECISION_KAPPA_MARGIN_CEILING: f64 = 1.0;
65/// Geometric ratio between consecutive proximal-ridge rungs, and the `Default`
66/// for `ArrowSolveOptions::ridge_growth`. One decade per rejection, which is
67/// what fixes the ladder's reach at
68/// `DEFAULT_PROXIMAL_INITIAL_RIDGE · 10^(DEFAULT_PROXIMAL_MAX_ATTEMPTS − 1)`;
69/// see [`DEFAULT_PROXIMAL_MAX_ATTEMPTS`] for why that reach is the requirement.
70pub const DEFAULT_PROXIMAL_RIDGE_GROWTH: f64 = 10.0;
71
72/// Number of geometric proximal-ridge escalations the adaptive correction
73/// attempts before giving up. Raised from 16 to 22 so the ridge can climb from
74/// `1e-8` to `~1e14` (`1e-8 · 10^21`): when the penalised Hessian curvature
75/// along the gradient exceeds `~1e9`, the damped Newton step at ridge `1e9`
76/// still overshoots, and the extra decades let the step length collapse far
77/// enough to either find descent or reach the near-stationary resolution floor
78/// that triggers the convergence exit. The cost of the extra attempts is paid
79/// only on configs that would otherwise have failed.
80pub const DEFAULT_PROXIMAL_MAX_ATTEMPTS: usize = 22;
81
82/// Armijo sufficient-decrease constant — sourced from the shared optimizer
83/// constants so the workspace has exactly one `c₁`.
84pub(crate) const DEFAULT_ARMIJO_C1: f64 = opt::constants::ARMIJO_C1;
85
86pub(crate) const DEFAULT_GRADIENT_TOLERANCE: f64 = 1e-10;
87
88/// Relative objective resolution for the proximal-correction convergence exit.
89///
90/// When the best achievable change in the penalised objective across all ridge
91/// attempts is within `rel_tol · (|f| + 1)` of the incumbent value, the damped
92/// Newton model has reached the floating-point resolution of the objective and
93/// no further productive decrease exists. `8e-12` sits a few decades above the
94/// `~2.2e-16` f64 epsilon (so genuine reductions of a well-scaled objective are
95/// never swallowed) yet comfortably above the accumulated rounding of the
96/// `O(N·M·p)` reductions that form the objective, so a truly stationary state
97/// is recognised rather than chased into a spurious failure.
98pub(crate) const DEFAULT_PROXIMAL_CONVERGENCE_REL_TOL: f64 = 8e-12;
99
100pub(crate) const EUCLIDEAN_MANIFOLD_MODE_FINGERPRINT: u64 = 0;
101
102pub(crate) const ARROW_FACTOR_CACHE_HTBETA_BUDGET_BYTES: usize = 256 * 1024 * 1024;
103
104/// Matrix-free shared-block multiply for large BA/SAE Schur PCG.
105///
106/// The closure writes `out = H_ββ x` without the LM ridge. This is the hook
107/// that lets SAE-manifold scale callers avoid materializing a dense `K × K`
108/// shared block before Agarwal-style inexact Schur PCG.
109pub type SharedBetaMatvec =
110    Arc<dyn for<'a> Fn(ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync>;
111
112pub type RowHtbetaMatvec =
113    Arc<dyn for<'a> Fn(usize, ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync>;
114
115/// Row-local matrix-free transpose multiply `out += H_βt^(i) · v` (length `K`).
116///
117/// This is the adjoint of [`RowHtbetaMatvec`]: it scatters a per-row latent
118/// vector `v` (length `d_i`) back into the shared β gradient, **adding** its
119/// contribution to `out`. For the SAE Kronecker form this is the sparse
120/// `scatter_jbeta_t` over the row's active atoms — `O(m_i · p)` per row, the
121/// per-row sparse apply that replaces the `O(K)` column-probe in the GPU and
122/// streaming Schur matvec.
123pub type RowHtbetaTransposeMatvec =
124    Arc<dyn for<'a> Fn(usize, ArrayView1<'a, f64>, &mut Array1<f64>) + Send + Sync>;
125
126pub type StreamingArrowRowBuilder =
127    Arc<dyn Fn(usize) -> Result<ArrowRowBlock, ArrowSchurError> + Send + Sync>;
128
129/// GPU-backed Schur matvec for CPU-driven PCG at K ≥ 5000.
130///
131/// The closure writes `out = S·x` where `S = H_ββ + ρ·I − Σ_i Y_i^T Y_i`
132/// is the reduced shared system, with `Y_i = L_i^{-1} H_tβ^(i)` pre-computed
133/// on device from the same forward kernel that Layer D uses for the dense Schur
134/// build. The CPU-driven Steihaug-CG outer loop uploads `x` (K doubles),
135/// receives `out` (K doubles), and handles the H_ββ contribution on the CPU side.
136///
137/// Constructed by `crate::gpu_kernels::arrow_schur::gpu_schur_matvec_backend` when
138/// `cuda_selected()` and K ≥ 5000. The closure is `Send + Sync` so PCG callers
139/// can hold it in an `Arc`.
140pub type GpuSchurMatvec = Arc<dyn Fn(&Array1<f64>, &mut Array1<f64>) + Send + Sync>;
141
142pub(crate) type MetricWeights = [f64];