1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
use super::*;
/// Cap on the number of coordinates at which a per-atom shape band is
/// materialized. The full per-atom decoder covariance is exact and exposed
/// regardless; this only bounds the cost of the convenience band, which is
/// evaluated at an evenly-strided subset of the atom's own on-atom coordinates.
pub const SHAPE_BAND_MAX_POINTS: usize = 512;
/// Entry budget for materializing one atom's dense `(M_k·p)²` decoder
/// covariance in the fit payload. Above it (LLM-scale ambient `p`) the band
/// quantities are computed exactly from the factored frame covariance and the
/// dense export is omitted (`decoder_covariance: None`) — the python reader
/// treats it as optional. 2^24 f64 entries = 128 MiB per atom.
pub const SAE_DECODER_COV_PAYLOAD_MAX_ENTRIES: usize = 1 << 24;
/// Posterior uncertainty of one fitted atom's manifold shape.
///
/// In the primary path — [`SaeManifoldTerm::assemble_shape_uncertainty`], and
/// after a structure/finalization change its final-state twin
/// [`SaeManifoldTerm::recompute_joint_shape_uncertainty`] — the covariance is
/// the φ-scaled β-block of the JOINT inverse Hessian (coordinates marginalized
/// out); the band is its closed-form push-forward through the linear
/// basis→ambient map `m_k(t) = Φ_k(t)·B_k`. This is what the production fit
/// returns.
///
/// When a streaming fit cannot provide the exact joint factor, every band field
/// is `None`. No per-atom marginal is substituted for the joint posterior.
#[derive(Debug, Clone)]
pub struct SaeAtomShapeUncertainty {
/// φ-scaled posterior covariance of this atom's decoder coefficients,
/// `Cov(β_k) = φ·S_β⁻¹[block_k]`, shape `(M_k·p, M_k·p)` in the decoder's
/// row-major `(basis, channel)` flat layout (flat index `b·p + c`).
///
/// `None` when materializing it would exceed
/// [`SAE_DECODER_COV_PAYLOAD_MAX_ENTRIES`] (LLM-scale ambient `p`: at
/// `(M=8, p=2048)` the dense block is 2 GiB *per atom*, at
/// `(M=16, p=5120)` ~50 GiB). The band quantities below are still exact
/// in that case — they are computed directly from the factored
/// `(M_k·r_k)²` frame covariance without ever lifting it.
pub decoder_covariance: Option<Array2<f64>>,
/// Coordinates at which the band is evaluated, shape `(G, d_k)`.
pub band_coords: Option<Array2<f64>>,
/// Fitted ambient point `m_k(t) = Φ_k(t)·B_k` at each band coordinate,
/// shape `(G, p)`.
pub band_mean: Option<Array2<f64>>,
/// Posterior standard deviation of each ambient channel at each band
/// coordinate, `sqrt(Var_c(t))` with
/// `Var_c(t) = Σ_{b1,b2} Φ[b1] Φ[b2] Cov(β_k)[(b1,c),(b2,c)]`, shape
/// `(G, p)`.
///
/// This is the MODEL-BASED band (`Cov = φ̂ H⁻¹`), correct only when the
/// working reconstruction likelihood is correctly specified. See
/// [`Self::band_sd_robust`] for the misspecification-robust companion.
pub band_sd: Option<Array2<f64>>,
/// Sandwich (Godambe / robust) posterior standard deviation of each ambient
/// channel, with the same shape as the model-based band, computed from the within-channel
/// sandwich covariance `A_c⁻¹ J_cc A_c⁻¹` (see `super::sandwich`). Reported
/// ALONGSIDE the model-based band, never in place of it; the two coincide
/// when the information-matrix equality holds and diverge exactly to the
/// degree the residuals violate the working likelihood (heteroskedastic
/// tokens, LayerNorm structure, cross-channel template correlation).
///
/// `None` when the robust band was not requested for this atom or could not
/// be formed (e.g. the huge-`p` factored path where the dense within-channel
/// meat is not materialized) — an honest gap, not a fabricated band.
pub band_sd_robust: Option<Array2<f64>>,
}
/// Posterior shape uncertainty for a whole SAE-manifold fit: one band per atom
/// plus the shared Gaussian reconstruction dispersion `φ̂` used to scale every
/// covariance. See [`SaeManifoldTerm::assemble_shape_uncertainty`].
#[derive(Debug, Clone)]
pub struct SaeShapeUncertainty {
/// Gaussian reconstruction scale `φ̂ = RSS / residual-dof`.
pub dispersion: f64,
/// One entry per atom, in atom order.
pub atoms: Vec<SaeAtomShapeUncertainty>,
}