Skip to main content

gam_models/inference/
model.rs

1use crate::bms::{LatentMeasureKind, LatentZConditionalCalibration, LatentZRankIntCalibration};
2use crate::survival::construction::{
3    SurvivalBaselineConfig, SurvivalTimeBasisConfig, parse_survival_baseline_config,
4};
5use crate::survival::location_scale::ResidualDistribution;
6use crate::survival::lognormal_kernel::FrailtySpec;
7use crate::wiggle::{
8    monotone_wiggle_basis_with_derivative_order, validate_monotone_wiggle_beta_nonnegative,
9};
10use gam_linalg::faer_ndarray::array2_to_nested_vec;
11use gam_problem::types::{
12    InverseLink, LatentCLogLogState, LikelihoodSpec, MixtureLinkState, ResponseFamily, SasLinkSpec,
13    SasLinkState, StandardLink,
14};
15use gam_runtime::span::span_index_for_breakpoints;
16use gam_solve::estimate::{BlockRole, FittedLinkState, UnifiedFitResult};
17use gam_solve::mixture_link::{state_from_beta_logisticspec, state_from_sasspec};
18use gam_terms::basis::BasisOptions;
19use gam_terms::inference::formula_dsl::{
20    inverse_link_supports_joint_wiggle, joint_wiggle_unsupported_link_message, parse_formula,
21    parse_surv_interval_response, parse_surv_response, parsed_term_column_names,
22};
23use gam_terms::smooth::{AdaptiveRegularizationDiagnostics, TermCollectionSpec};
24// The data-schema value types live in the `gam-data` foundation crate; they
25// were previously authored here and are still named `gam::inference::model::{
26// ColumnKindTag, DataSchema, SchemaColumn}` by a broad set of integration tests
27// and by saved-payload consumers. Re-export them so that public path stays
28// valid rather than forcing every caller onto the relocated crate path.
29pub use gam_data::{ColumnKindTag, DataSchema, SchemaColumn};
30use ndarray::{Array1, Array2};
31use serde::{Deserialize, Serialize};
32use serde_json::Value as JsonValue;
33use std::collections::{BTreeMap, HashMap, HashSet};
34use std::fs;
35use std::ops::{Deref, DerefMut};
36use std::path::Path;
37
38/// Canonical saved-model payload schema version.
39///
40/// Every `FittedModelPayload` written by any binary (CLI `gam`, gam-pyffi,
41/// downstream library users) must set this as its `version` field, and every
42/// load path asserts equality via `validate_for_persistence`. Bump this when:
43///   - A required field is added to `FittedModelPayload` and the set of
44///     Option<T> fields that must be `Some(...)` for a given `family_state`
45///     changes (otherwise the `#[serde(default)]` decode would silently fill
46///     the new field with `None` when loading an older model and the CLI
47///     predict path would run with stale metadata).
48///   - The on-wire shape of any `serde`-tagged enum variant changes such that
49///     older payloads no longer round-trip losslessly.
50///   - The semantics of an existing field change (e.g. sign convention,
51///     coordinate frame) in a way that predict output would silently diverge
52///     between old and new readers.
53///
54/// Do NOT bump for purely additive `Option<T>` fields that the save-time
55/// invariant (`validate_for_persistence`) does not yet require. Those are
56/// forward-compatible.
57pub const MODEL_PAYLOAD_VERSION: u32 = 8;
58
59/// Schema-free saved-model metadata keyed by stable group id.
60///
61/// The values are JSON rather than a typed enum because group provenance is
62/// supplied by caller-owned catalogs. `FittedModelPayload::group_metadata`
63/// wraps this in `Option` with `#[serde(default)]`, so model files written
64/// before the field existed deserialize as `None`.
65pub type GroupMetadata = BTreeMap<String, JsonValue>;
66
67/// Saved exact spline-scan fit (#1030/#1034): the predict-time feature column
68/// plus the lossless smoother state the Gaussian-bridge `predict` replays.
69#[derive(Clone, Debug, Serialize, Deserialize)]
70pub struct SavedSplineScan {
71    /// Training column name feeding the single 1-D smooth at predict time.
72    pub feature_column: String,
73    pub state: gam_solve::spline_scan::SplineScanState,
74}
75
76/// Saved multiresolution residual-cascade fit (#1032): the predict-time feature
77/// columns (d ∈ {2, 3}) plus the serializable cascade state that `from_state`
78/// rebuilds a predict-capable `ResidualCascadeFit` from. The cascade is a
79/// DIFFERENT posterior from the dense Duchon/Matérn term — never a silent swap.
80#[derive(Clone, Debug, Serialize, Deserialize)]
81pub struct SavedResidualCascade {
82    /// Training column names for the d ∈ {2, 3} scattered-smooth coordinates.
83    pub feature_columns: Vec<String>,
84    pub state: gam_solve::residual_cascade::ResidualCascadeState,
85}
86
87/// Typed error surface for `src/inference/model.rs` saved-model code.
88///
89/// Every variant carries a free-form `reason: String` payload; `Display`
90/// emits exactly that payload, so converting a `FittedModelError` into
91/// `String` (via the `From` impl below) is byte-equivalent to the pre-
92/// refactor `Err(format!(...))` / `Err("...".to_string())` strings that
93/// the same call sites produced. This lets external callers keep using
94/// `?` against `Result<_, String>` without source changes — the typed
95/// enum is purely an in-module discipline gain.
96#[derive(Clone, Debug, PartialEq, Eq)]
97pub enum FittedModelError {
98    /// Saved payload structure / shape / version disagrees with what the
99    /// current binary expects (e.g. covariance shape, block ordering,
100    /// schema version, C2 continuity, out-of-range span/basis indices).
101    SchemaMismatch { reason: String },
102    /// Saved payload bytes / numeric content are corrupt or unreadable
103    /// (non-finite scalars, invalid JSON, IO failure, malformed stateful
104    /// link state).
105    PayloadCorrupt { reason: String },
106    /// A required field that the current code path needs is absent from
107    /// the payload (typically `..; refit` errors).
108    MissingField { reason: String },
109    /// A combination of saved-model options is not supported by the
110    /// current binary (unsupported deployment-extension kind, unsupported
111    /// kernel marker, unsupported survival_time_basis variant, etc.).
112    IncompatibleConfig { reason: String },
113    /// An input value rejected by a save-time sanity gate (e.g. negative
114    /// ridge alpha).
115    InvalidInput { reason: String },
116}
117
118impl_reason_error_boilerplate! {
119    FittedModelError {
120        SchemaMismatch,
121        PayloadCorrupt,
122        MissingField,
123        IncompatibleConfig,
124        InvalidInput,
125    }
126}
127
128// Boundary conversions so external `Result<_, EstimationError>` /
129// `Result<_, SurvivalPredictError>` call sites can propagate with `?`.
130// Survival prediction keeps the model-layer source so the chain identifies
131// the payload/schema failure that triggered the prediction error.
132impl From<FittedModelError> for gam_solve::model_types::EstimationError {
133    fn from(err: FittedModelError) -> Self {
134        gam_solve::model_types::EstimationError::InvalidInput(err.to_string())
135    }
136}
137
138impl From<FittedModelError> for crate::survival::predict::SurvivalPredictError {
139    fn from(err: FittedModelError) -> Self {
140        crate::survival::predict::SurvivalPredictError::ModelPayload {
141            context: "saved-model survival prediction payload",
142            source: err,
143        }
144    }
145}
146
147#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
148pub struct SavedLatentZNormalization {
149    pub mean: f64,
150    pub sd: f64,
151}
152
153impl SavedLatentZNormalization {
154    pub fn validate(&self, context: &str) -> Result<(), FittedModelError> {
155        if !self.mean.is_finite() {
156            return Err(FittedModelError::PayloadCorrupt {
157                reason: format!("{context} latent z mean must be finite"),
158            });
159        }
160        if !(self.sd.is_finite() && self.sd > 1e-12) {
161            return Err(FittedModelError::PayloadCorrupt {
162                reason: format!(
163                    "{context} latent z sd must be finite and > 1e-12; got {}",
164                    self.sd
165                ),
166            });
167        }
168        Ok::<(), _>(())
169    }
170
171    pub fn apply(&self, z: &Array1<f64>, context: &str) -> Result<Array1<f64>, FittedModelError> {
172        self.validate(context)?;
173        if z.iter().any(|value| !value.is_finite()) {
174            return Err(FittedModelError::PayloadCorrupt {
175                reason: format!("{context} requires finite z values"),
176            });
177        }
178        Ok(z.mapv(|zi| (zi - self.mean) / self.sd))
179    }
180}
181
182pub const TRANSFORMATION_SCORE_PIT_CLIP_EPS: f64 = 1.0e-12;
183
184#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
185#[serde(rename_all = "kebab-case")]
186#[derive(Default)]
187pub enum TransformationScoreKind {
188    #[default]
189    FiniteSupportPit,
190}
191
192#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
193pub struct TransformationScoreCalibration {
194    #[serde(default)]
195    pub score_kind: TransformationScoreKind,
196    #[serde(default = "default_transformation_score_pit_clip_eps")]
197    pub clip_eps: f64,
198}
199
200const fn default_transformation_score_pit_clip_eps() -> f64 {
201    TRANSFORMATION_SCORE_PIT_CLIP_EPS
202}
203
204impl TransformationScoreCalibration {
205    pub fn finite_support_pit() -> Self {
206        Self {
207            score_kind: TransformationScoreKind::FiniteSupportPit,
208            clip_eps: TRANSFORMATION_SCORE_PIT_CLIP_EPS,
209        }
210    }
211
212    pub fn validate(&self, context: &str) -> Result<(), FittedModelError> {
213        if self.score_kind != TransformationScoreKind::FiniteSupportPit {
214            return Err(FittedModelError::IncompatibleConfig {
215                reason: format!("{context} supports only finite-support CTN PIT score semantics"),
216            });
217        }
218        if !(self.clip_eps.is_finite() && self.clip_eps > 0.0 && self.clip_eps < 0.5) {
219            return Err(FittedModelError::IncompatibleConfig {
220                reason: format!(
221                    "{context} requires PIT clip_eps in (0, 0.5), got {}",
222                    self.clip_eps
223                ),
224            });
225        }
226        Ok::<(), _>(())
227    }
228}
229
230#[derive(Clone, Serialize, Deserialize)]
231pub struct FittedModelPayload {
232    pub version: u32,
233    pub formula: String,
234    pub model_kind: ModelKind,
235    pub family_state: FittedFamily,
236    pub family: String,
237    /// Human-readable advisories produced while materializing this model —
238    /// e.g. an mgcv-style "k was reduced to the data support" note when a
239    /// cubic-regression marginal is capped, or a basis-degradation note. These
240    /// are surfaced to CLI users via `print_inference_summary`; persisting them
241    /// here lets the Python (gamfit) interface surface the SAME advisories as
242    /// warnings / `model.notes` instead of silently dropping them at the FFI
243    /// boundary (#1543). `#[serde(default)]` keeps older payloads (which had no
244    /// such field) deserializing cleanly as "no notes".
245    #[serde(default)]
246    pub inference_notes: Vec<String>,
247    #[serde(default)]
248    pub used_device: bool,
249    #[serde(default)]
250    pub fit_result: Option<UnifiedFitResult>,
251    /// Unified (family-agnostic) representation of the fit result.
252    #[serde(default)]
253    pub unified: Option<UnifiedFitResult>,
254    /// Exact O(n) spline-scan fit representation (#1030/#1034): the
255    /// state-space smoothing-spline posterior of a single 1-D Gaussian
256    /// smooth. When `Some`, this standard Gaussian model's predictions
257    /// replay the Gaussian bridge from this state and the model carries no
258    /// dense `fit_result` (the representations are mutually exclusive —
259    /// enforced by `validate_for_persistence`). `#[serde(default)]` so older
260    /// payloads read as: not a scan model.
261    #[serde(default)]
262    pub spline_scan: Option<SavedSplineScan>,
263    /// O(n log n) multiresolution residual-cascade fit (#1032): the persisted
264    /// multilevel Wendland-frame state for a single scattered 2–3D Gaussian
265    /// smooth past the dense-kernel cliff. When `Some`, predictions replay the
266    /// cascade posterior; mutually exclusive with `spline_scan`/`fit_result`.
267    /// `#[serde(default)]` keeps forward-compatibility with older payloads.
268    #[serde(default)]
269    pub residual_cascade: Option<SavedResidualCascade>,
270    #[serde(default)]
271    pub data_schema: Option<DataSchema>,
272    pub link: Option<InverseLink>,
273    #[serde(default)]
274    pub mixture_link_param_covariance: Option<Vec<Vec<f64>>>,
275    #[serde(default)]
276    pub sas_param_covariance: Option<Vec<Vec<f64>>>,
277    #[serde(default)]
278    pub formula_noise: Option<String>,
279    #[serde(default)]
280    pub formula_logslope: Option<String>,
281    #[serde(default)]
282    pub formula_logslopes: Option<Vec<String>>,
283    #[serde(default)]
284    pub offset_column: Option<String>,
285    #[serde(default)]
286    pub noise_offset_column: Option<String>,
287    /// Name of the analytic prior-weights column used at fit time (`weights=`),
288    /// persisted so replicate/generative sampling can re-resolve the per-row
289    /// weights and draw heteroskedastic Gaussian observation noise
290    /// `sigma_i = sigma_hat / sqrt(w_i)` (#2025). `None` for an unweighted fit.
291    #[serde(default)]
292    pub weight_column: Option<String>,
293    #[serde(default)]
294    pub beta_noise: Option<Vec<f64>>,
295    #[serde(default)]
296    pub noise_projection: Option<Vec<Vec<f64>>>,
297    #[serde(default)]
298    pub noise_center: Option<Vec<f64>>,
299    #[serde(default)]
300    pub noise_scale: Option<Vec<f64>>,
301    #[serde(default)]
302    pub noise_non_intercept_start: Option<usize>,
303    /// Tikhonov ridge alpha used by `solve_scale_projection` when fitting
304    /// `noise_projection`.  Persisted so prediction-time replay is identical
305    /// to fit-time projection.
306    #[serde(default)]
307    pub noise_projection_ridge_alpha: Option<f64>,
308    #[serde(default)]
309    pub gaussian_response_scale: Option<f64>,
310    #[serde(default)]
311    pub linkwiggle_knots: Option<Vec<f64>>,
312    #[serde(default)]
313    pub linkwiggle_degree: Option<usize>,
314    #[serde(default)]
315    pub beta_link_wiggle: Option<Vec<f64>>,
316    /// Frozen-index mean-coordinate shift `s` for the standard binomial-mean
317    /// link-warp predict runtime (#2141). Predict evaluates the warp basis at
318    /// the frozen index `η̂ = X·(β_saved + s)` the fit pinned `B` at, rather than
319    /// at the de-aliased base predictor `X·β_saved` — reproducing the fitted `q`
320    /// (and thus the fitted deviance) at predict time. `None` for models whose
321    /// warp path never de-aliased (location-scale / dynamic-basis), where the
322    /// base predictor already is the warp index (back-compatible serde default).
323    #[serde(default)]
324    pub link_wiggle_index_shift: Option<Vec<f64>>,
325    #[serde(default)]
326    pub baseline_timewiggle_knots: Option<Vec<f64>>,
327    #[serde(default)]
328    pub baseline_timewiggle_degree: Option<usize>,
329    #[serde(default)]
330    pub baseline_timewiggle_penalty_orders: Option<Vec<usize>>,
331    #[serde(default)]
332    pub baseline_timewiggle_double_penalty: Option<bool>,
333    #[serde(default)]
334    pub beta_baseline_timewiggle: Option<Vec<f64>>,
335    #[serde(default)]
336    pub beta_baseline_timewiggle_by_cause: Option<Vec<Vec<f64>>>,
337    #[serde(default)]
338    pub z_column: Option<String>,
339    #[serde(default)]
340    pub z_columns: Option<Vec<String>>,
341    #[serde(default)]
342    pub latent_z_normalization: Option<SavedLatentZNormalization>,
343    #[serde(default)]
344    pub latent_score_contract: Option<SavedLatentScoreContract>,
345    #[serde(default)]
346    pub latent_measure: Option<LatentMeasureKind>,
347    /// Optional rank-INT calibration for the latent score (BMS family).
348    /// When `Some`, the marginal-slope predictor routes the input `z`
349    /// through [`LatentZRankIntCalibration::apply_at_predict`] before the
350    /// closed-form standard-normal kernel, matching fit-time semantics.
351    /// `#[serde(default)]` so models persisted before this field existed
352    /// continue to deserialize cleanly (interpreted as: no calibration).
353    #[serde(default)]
354    pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
355    /// Optional conditional location-scale calibration of the latent score
356    /// (#905, BMS family). When `Some`, the marginal-slope predictor replaces
357    /// the (normalized) input `z` by `ζ = (z − m(C))/√v(C)` — rebuilding the
358    /// conditioning span `a(C)` from the marginal prediction design — before
359    /// the closed-form standard-normal kernel, matching fit-time semantics.
360    /// Mutually exclusive with `latent_z_rank_int_calibration`. `#[serde(default)]`
361    /// so pre-existing models deserialize cleanly (interpreted as: no
362    /// conditional calibration).
363    #[serde(default)]
364    pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
365    #[serde(default)]
366    pub marginal_baseline: Option<f64>,
367    #[serde(default)]
368    pub logslope_baseline: Option<f64>,
369    #[serde(default)]
370    pub logslope_baselines: Option<Vec<f64>>,
371    #[serde(default)]
372    pub score_warp_runtime: Option<SavedCompiledFlexBlock>,
373    #[serde(default)]
374    pub link_deviation_runtime: Option<SavedCompiledFlexBlock>,
375    /// Width `p₁` of the survival marginal-slope absorbed Stage-1 influence block
376    /// (#461) when present (the dedicated trailing absorber block). Predict drops
377    /// its `γ`; this records the column count so the predictor can account for
378    /// the extra block and slice `γ` out of the joint covariance.
379    #[serde(default)]
380    pub influence_absorber_width: Option<usize>,
381    #[serde(default)]
382    pub survival_entry: Option<String>,
383    #[serde(default)]
384    pub survival_exit: Option<String>,
385    #[serde(default)]
386    pub survival_event: Option<String>,
387    #[serde(default)]
388    pub survivalspec: Option<String>,
389    #[serde(default)]
390    pub survival_cause_count: Option<usize>,
391    #[serde(default)]
392    pub survival_endpoint_names: Option<Vec<String>>,
393    #[serde(default)]
394    pub survival_baseline_target: Option<String>,
395    #[serde(default)]
396    pub survival_baseline_scale: Option<f64>,
397    #[serde(default)]
398    pub survival_baseline_shape: Option<f64>,
399    #[serde(default)]
400    pub survival_baseline_rate: Option<f64>,
401    #[serde(default)]
402    pub survival_baseline_makeham: Option<f64>,
403    #[serde(default)]
404    pub survival_time_basis: Option<String>,
405    #[serde(default)]
406    pub survival_time_degree: Option<usize>,
407    #[serde(default)]
408    pub survival_time_knots: Option<Vec<f64>>,
409    #[serde(default)]
410    pub survival_time_keep_cols: Option<Vec<usize>>,
411    #[serde(default)]
412    pub survival_time_smooth_lambda: Option<f64>,
413    #[serde(default)]
414    pub survival_time_anchor: Option<f64>,
415    #[serde(default)]
416    pub survivalridge_lambda: Option<f64>,
417    #[serde(default)]
418    pub survival_likelihood: Option<String>,
419    #[serde(default)]
420    pub survival_beta_time: Option<Vec<f64>>,
421    #[serde(default)]
422    pub survival_beta_threshold: Option<Vec<f64>>,
423    #[serde(default)]
424    pub survival_beta_log_sigma: Option<Vec<f64>>,
425    #[serde(default)]
426    pub survival_noise_projection: Option<Vec<Vec<f64>>>,
427    #[serde(default)]
428    pub survival_noise_center: Option<Vec<f64>>,
429    #[serde(default)]
430    pub survival_noise_scale: Option<Vec<f64>>,
431    #[serde(default)]
432    pub survival_noise_non_intercept_start: Option<usize>,
433    /// Survival analog of `noise_projection_ridge_alpha`: the Tikhonov ridge
434    /// used when fitting the survival log-sigma projection.
435    #[serde(default)]
436    pub survival_noise_projection_ridge_alpha: Option<f64>,
437    #[serde(default)]
438    pub survival_distribution: Option<ResidualDistribution>,
439    #[serde(default)]
440    pub training_headers: Option<Vec<String>>,
441    /// Container type of the table the model was fitted on, as detected by the
442    /// active frontend (`"pandas"`, `"polars"`, `"pyarrow"`, `"numpy"`, or
443    /// `"unknown"` outside a typed table frontend). This presentation provenance
444    /// is required in the current schema so save/load cannot silently change an
445    /// ambiguous predict input's output container.
446    pub training_table_kind: String,
447    /// Per-column (min, max) of the training input matrix, parallel to
448    /// `training_headers`. At predict time, inputs are axis-clipped to these
449    /// ranges so that out-of-distribution points evaluate at the nearest face
450    /// of the training bounding box rather than extrapolating polynomial
451    /// trends from polyharmonic / spline bases beyond the data envelope. Old
452    /// model JSONs that pre-date this field load with `None`, in which case
453    /// the predict path falls through unchanged (no clipping).
454    #[serde(default)]
455    pub training_feature_ranges: Option<Vec<(f64, f64)>>,
456    /// User-supplied per-group metadata, keyed by stable group identifier.
457    ///
458    /// This is intentionally schema-free JSON so provenance maps can carry
459    /// mixed scalar/list/object values. Missing in older payloads means no
460    /// group metadata was persisted.
461    #[serde(default, skip_serializing_if = "Option::is_none")]
462    pub group_metadata: Option<GroupMetadata>,
463    /// Deployment-time no-refit group extensions applied after fitting.
464    ///
465    /// Each entry records the requested group coordinate, caller metadata, and
466    /// prior used to initialize the inserted coefficient. The active
467    /// prediction contract lives in `data_schema` + `resolved_termspec`; this
468    /// ledger preserves provenance without requiring a refit.
469    #[serde(default, skip_serializing_if = "Vec::is_empty")]
470    pub deployment_extensions: Vec<SavedDeploymentExtension>,
471    /// Transformation-normal: B-spline knots for the response-direction basis.
472    #[serde(default)]
473    pub transformation_response_knots: Option<Vec<f64>>,
474    /// Transformation-normal: deviation nullspace transform matrix (row-major).
475    #[serde(default)]
476    pub transformation_response_transform: Option<Vec<Vec<f64>>>,
477    /// Transformation-normal: B-spline degree for the response basis.
478    #[serde(default)]
479    pub transformation_response_degree: Option<usize>,
480    /// Transformation-normal: median of the response used for anchoring.
481    #[serde(default)]
482    pub transformation_response_median: Option<f64>,
483    /// Transformation-normal saved score contract. The score is the exact
484    /// finite-support PIT:
485    /// z = Phi^{-1}((Phi(h) - Phi(h_L)) / (Phi(h_U) - Phi(h_L))).
486    #[serde(default)]
487    pub transformation_score_calibration: Option<TransformationScoreCalibration>,
488    #[serde(default)]
489    pub resolved_termspec: Option<TermCollectionSpec>,
490    #[serde(default)]
491    pub resolved_termspec_noise: Option<TermCollectionSpec>,
492    #[serde(default)]
493    pub resolved_termspec_logslope: Option<TermCollectionSpec>,
494    #[serde(default)]
495    pub resolved_termspec_logslopes: Option<Vec<TermCollectionSpec>>,
496    #[serde(default)]
497    pub adaptive_regularization_diagnostics: Option<AdaptiveRegularizationDiagnostics>,
498    /// Precomputed exact Gaussian-identity jackknife+ statistics (#942).
499    ///
500    /// Populated *only* for a standard Gaussian-identity model fit with unit
501    /// prior weights, where the closed-form Sherman–Morrison leave-one-out
502    /// substrate gives a distribution-free prediction interval with no held-out
503    /// fold, targeting ≈level coverage at α = 1 − level with the finite-sample
504    /// floor ≥ 2·level − 1 (Barber et al. 2021, ≥ 1 − 2α; see the pyffi
505    /// route for the calibration decision, #1546). When `Some`, `predict(interval=level)`
506    /// auto-routes through it (the MAGIC default); when `None` — any other
507    /// family/link, reweighted rows, or an older payload — predict falls back
508    /// to the model-based posterior band and labels the provenance honestly.
509    /// `#[serde(default)]` so pre-existing models deserialize as: no jackknife+
510    /// substrate available.
511    #[serde(default)]
512    pub gaussian_jackknife_plus:
513        Option<crate::inference::full_conformal::GaussianJackknifePlusStats>,
514    /// Precomputed substrate for the EXACT Gaussian-identity full-conformal set
515    /// (#942 Layer 1 + the frozen-ρ self-diagnostic).
516    ///
517    /// Populated under the SAME eligibility as `gaussian_jackknife_plus`
518    /// (Gaussian-identity, unit prior weights, offset-free, no link wiggle). It
519    /// persists the training design + response + frozen penalty `Sλ` so the
520    /// prediction set that is exact GIVEN `Sλ` (a union of intervals, valid for
521    /// any penalized smooth) can be replayed per test point — one Cholesky each,
522    /// zero refits. Because λ̂ was selected from all training responses, the
523    /// frozen-λ construction is not permutation symmetric in the augmented
524    /// points; the distribution-free finite-sample coverage theorem is asserted
525    /// only per row where the surfaced frozen-ρ certificate accepts (under the
526    /// global-ρ grid-Lipschitz assumption). `None` for any
527    /// ineligible model or an older payload, in which case the exact-set predict
528    /// path errors with a clear message and the caller uses jackknife+ or the
529    /// posterior band. `#[serde(default)]` so pre-existing models deserialize as
530    /// no exact substrate available.
531    #[serde(default)]
532    pub full_conformal: Option<crate::inference::full_conformal::ExactFullConformalSubstrate>,
533}
534
535#[derive(Clone, Debug, Serialize, Deserialize)]
536pub struct SavedDeploymentExtension {
537    pub name: String,
538    pub kind: String,
539    pub term: String,
540    pub level: JsonValue,
541    pub level_bits: u64,
542    pub coefficient_index: usize,
543    pub coefficient_mean: f64,
544    pub coefficient_variance: f64,
545    #[serde(default, skip_serializing_if = "Option::is_none")]
546    pub metadata: Option<JsonValue>,
547    #[serde(default, skip_serializing_if = "Option::is_none")]
548    pub prior: Option<JsonValue>,
549}
550
551/// Append deployment-only extension columns to the fitted design coordinate system.
552///
553/// No-refit group extension adds a new coefficient block after the fitted
554/// coefficient vector:
555///
556///   beta_ext = [beta_old, beta_new],    beta_new = mu_new.
557///
558/// For a new random-effect level g, the appended basis is the indicator
559/// e_g(x_i) = 1{x_i == g}.  The old fitted basis X_old is not rebuilt or
560/// reordered, so rows that do not exercise g have
561///
562///   eta_ext = X_old beta_old + 0 * beta_new = eta_old.
563///
564/// Rows at the new level get the exact prior-mean shift e_g beta_new.  This
565/// helper enforces the coordinate identity by requiring extension coefficient
566/// indices to be the consecutive tail columns of the base design.
567pub fn append_deployment_extension_columns(
568    model: &FittedModelPayload,
569    data: ndarray::ArrayView2<'_, f64>,
570    col_map: &HashMap<String, usize>,
571    training_headers: Option<&Vec<String>>,
572    base_design: Array2<f64>,
573) -> Result<Array2<f64>, FittedModelError> {
574    if model.deployment_extensions.is_empty() {
575        return Ok(base_design);
576    }
577    if base_design.nrows() != data.nrows() {
578        return Err(FittedModelError::SchemaMismatch {
579            reason: format!(
580                "deployment extension design row mismatch: base design has {} rows but data has {}",
581                base_design.nrows(),
582                data.nrows()
583            ),
584        });
585    }
586    let spec = model
587        .resolved_termspec
588        .as_ref()
589        .ok_or_else(|| FittedModelError::MissingField {
590            reason: "deployment extension prediction requires saved resolved_termspec; refit"
591                .to_string(),
592        })?;
593    let n = base_design.nrows();
594    let p_old = base_design.ncols();
595    let mut extensions: Vec<&SavedDeploymentExtension> =
596        model.deployment_extensions.iter().collect();
597    extensions.sort_by_key(|extension| extension.coefficient_index);
598    for (tail_idx, extension) in extensions.iter().enumerate() {
599        let expected = p_old + tail_idx;
600        if extension.coefficient_index != expected {
601            return Err(FittedModelError::SchemaMismatch {
602                reason: format!(
603                    "deployment extension '{}' has coefficient index {}, expected append-only index {}",
604                    extension.name, extension.coefficient_index, expected
605                ),
606            });
607        }
608    }
609
610    let mut out = Array2::<f64>::zeros((n, p_old + extensions.len()));
611    out.slice_mut(ndarray::s![.., ..p_old]).assign(&base_design);
612    for (tail_idx, extension) in extensions.into_iter().enumerate() {
613        if extension.kind != "random-effect-level" {
614            return Err(FittedModelError::IncompatibleConfig {
615                reason: format!(
616                    "unsupported deployment extension kind '{}' for '{}'",
617                    extension.kind, extension.name
618                ),
619            });
620        }
621        let term = spec
622            .random_effect_terms
623            .iter()
624            .find(|term| term.name == extension.term)
625            .ok_or_else(|| FittedModelError::MissingField {
626                reason: format!(
627                    "deployment extension '{}' references unknown random-effect term '{}'",
628                    extension.name, extension.term
629                ),
630            })?;
631        let prediction_col = training_headers
632            .and_then(|headers| headers.get(term.feature_col))
633            .and_then(|name| col_map.get(name))
634            .copied()
635            .unwrap_or(term.feature_col);
636        if prediction_col >= data.ncols() {
637            return Err(FittedModelError::SchemaMismatch {
638                reason: format!(
639                    "deployment extension '{}' feature column {} out of bounds for {} prediction columns",
640                    extension.name,
641                    prediction_col,
642                    data.ncols()
643                ),
644            });
645        }
646        let col = p_old + tail_idx;
647        let level_bits = gam_data::canonical_level_bits(f64::from_bits(extension.level_bits));
648        for row in 0..n {
649            if gam_data::canonical_level_bits(data[[row, prediction_col]]) == level_bits {
650                out[[row, col]] = 1.0;
651            }
652        }
653    }
654    Ok(out)
655}
656
657#[derive(Clone, Debug, Serialize, Deserialize)]
658pub struct SavedLatentScoreContract {
659    pub semantics: String,
660    pub source_transform_id: Option<String>,
661    pub normalization_mean: f64,
662    pub normalization_sd: f64,
663    pub clip_eps: Option<f64>,
664    pub conditioning_columns: Vec<String>,
665}
666
667impl FittedModelPayload {
668    pub fn new(
669        version: u32,
670        formula: String,
671        model_kind: ModelKind,
672        family_state: FittedFamily,
673        family: String,
674    ) -> Self {
675        Self {
676            version,
677            formula,
678            model_kind,
679            family_state,
680            family,
681            inference_notes: Vec::new(),
682            used_device: false,
683            fit_result: None,
684            unified: None,
685            spline_scan: None,
686            residual_cascade: None,
687            data_schema: None,
688            link: None,
689            mixture_link_param_covariance: None,
690            sas_param_covariance: None,
691            formula_noise: None,
692            formula_logslope: None,
693            formula_logslopes: None,
694            offset_column: None,
695            noise_offset_column: None,
696            weight_column: None,
697            beta_noise: None,
698            noise_projection: None,
699            noise_center: None,
700            noise_scale: None,
701            noise_non_intercept_start: None,
702            noise_projection_ridge_alpha: None,
703            gaussian_response_scale: None,
704            linkwiggle_knots: None,
705            linkwiggle_degree: None,
706            beta_link_wiggle: None,
707            link_wiggle_index_shift: None,
708            baseline_timewiggle_knots: None,
709            baseline_timewiggle_degree: None,
710            baseline_timewiggle_penalty_orders: None,
711            baseline_timewiggle_double_penalty: None,
712            beta_baseline_timewiggle: None,
713            beta_baseline_timewiggle_by_cause: None,
714            z_column: None,
715            z_columns: None,
716            latent_z_normalization: None,
717            latent_score_contract: None,
718            latent_measure: None,
719            latent_z_rank_int_calibration: None,
720            latent_z_conditional_calibration: None,
721            marginal_baseline: None,
722            logslope_baseline: None,
723            logslope_baselines: None,
724            score_warp_runtime: None,
725            link_deviation_runtime: None,
726            influence_absorber_width: None,
727            survival_entry: None,
728            survival_exit: None,
729            survival_event: None,
730            survivalspec: None,
731            survival_cause_count: None,
732            survival_endpoint_names: None,
733            survival_baseline_target: None,
734            survival_baseline_scale: None,
735            survival_baseline_shape: None,
736            survival_baseline_rate: None,
737            survival_baseline_makeham: None,
738            survival_time_basis: None,
739            survival_time_degree: None,
740            survival_time_knots: None,
741            survival_time_keep_cols: None,
742            survival_time_smooth_lambda: None,
743            survival_time_anchor: None,
744            survivalridge_lambda: None,
745            survival_likelihood: None,
746            survival_beta_time: None,
747            survival_beta_threshold: None,
748            survival_beta_log_sigma: None,
749            survival_noise_projection: None,
750            survival_noise_center: None,
751            survival_noise_scale: None,
752            survival_noise_non_intercept_start: None,
753            survival_noise_projection_ridge_alpha: None,
754            survival_distribution: None,
755            training_headers: None,
756            training_table_kind: "unknown".to_string(),
757            training_feature_ranges: None,
758            group_metadata: None,
759            deployment_extensions: Vec::new(),
760            transformation_response_knots: None,
761            transformation_response_transform: None,
762            transformation_response_degree: None,
763            transformation_response_median: None,
764            transformation_score_calibration: None,
765            resolved_termspec: None,
766            resolved_termspec_noise: None,
767            resolved_termspec_logslope: None,
768            resolved_termspec_logslopes: None,
769            adaptive_regularization_diagnostics: None,
770            gaussian_jackknife_plus: None,
771            full_conformal: None,
772        }
773    }
774
775    pub fn set_training_feature_metadata(
776        &mut self,
777        headers: Vec<String>,
778        feature_ranges: Vec<(f64, f64)>,
779    ) {
780        self.training_headers = Some(headers);
781        self.training_feature_ranges = Some(feature_ranges);
782    }
783
784    fn synchronize_empty_feature_contract(&mut self) {
785        if self.fit_result.is_none() {
786            return;
787        }
788        let Some(schema) = self.data_schema.as_ref() else {
789            return;
790        };
791        if !schema.columns.is_empty() {
792            return;
793        }
794        self.training_headers.get_or_insert_with(Vec::new);
795        self.resolved_termspec
796            .get_or_insert_with(|| TermCollectionSpec {
797                linear_terms: Vec::new(),
798                smooth_terms: Vec::new(),
799                random_effect_terms: Vec::new(),
800            });
801    }
802
803    /// Write the persistable time-basis snapshot for a survival model.
804    ///
805    /// This is the only path that should populate the `survival_time_*`
806    /// fields used by the loader. Routing every FFI builder through this
807    /// helper guarantees no builder can silently drop a field — the
808    /// marginal-slope save→load bug was a builder that
809    /// missed `survival_time_basis`.
810    pub fn apply_survival_time_basis(
811        &mut self,
812        snapshot: &crate::survival::construction::SavedSurvivalTimeBasis,
813    ) {
814        self.survival_time_basis = Some(snapshot.basisname.clone());
815        self.survival_time_degree = snapshot.degree;
816        self.survival_time_knots = snapshot.knots.clone();
817        self.survival_time_keep_cols = snapshot.keep_cols.clone();
818        self.survival_time_smooth_lambda = snapshot.smooth_lambda;
819        self.survival_time_anchor = Some(snapshot.anchor);
820    }
821
822    fn validate_payload_version(&self) -> Result<(), FittedModelError> {
823        if self.version != MODEL_PAYLOAD_VERSION {
824            return Err(FittedModelError::SchemaMismatch {
825                reason: format!(
826                    "saved model payload schema mismatch: file has version={}, \
827                 this binary expects MODEL_PAYLOAD_VERSION={}. \
828                 Refit with the current CLI, or rebuild the reader at the same \
829                 version the model was written with.",
830                    self.version, MODEL_PAYLOAD_VERSION
831                ),
832            });
833        }
834        Ok(())
835    }
836}
837
838#[derive(Clone, Serialize, Deserialize)]
839#[serde(tag = "model_type", rename_all = "kebab-case")]
840pub enum FittedModel {
841    Standard { payload: FittedModelPayload },
842    LocationScale { payload: FittedModelPayload },
843    MarginalSlope { payload: FittedModelPayload },
844    Survival { payload: FittedModelPayload },
845    TransformationNormal { payload: FittedModelPayload },
846}
847
848#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
849#[serde(rename_all = "kebab-case")]
850pub enum ModelKind {
851    Standard,
852    LocationScale,
853    MarginalSlope,
854    Survival,
855    TransformationNormal,
856}
857
858#[derive(Clone, Debug, Serialize, Deserialize)]
859#[serde(tag = "family_kind", rename_all = "kebab-case")]
860pub enum FittedFamily {
861    Standard {
862        likelihood: LikelihoodSpec,
863        #[serde(default)]
864        link: Option<StandardLink>,
865        #[serde(default)]
866        latent_cloglog_state: Option<LatentCLogLogState>,
867        #[serde(default)]
868        mixture_state: Option<MixtureLinkState>,
869        #[serde(default)]
870        sas_state: Option<SasLinkState>,
871    },
872    LocationScale {
873        likelihood: LikelihoodSpec,
874        #[serde(default)]
875        base_link: Option<InverseLink>,
876    },
877    MarginalSlope {
878        likelihood: LikelihoodSpec,
879        base_link: InverseLink,
880        frailty: FrailtySpec,
881    },
882    Survival {
883        likelihood: LikelihoodSpec,
884        #[serde(default)]
885        survival_likelihood: Option<String>,
886        #[serde(default)]
887        survival_distribution: Option<ResidualDistribution>,
888        frailty: FrailtySpec,
889    },
890    LatentSurvival {
891        frailty: FrailtySpec,
892    },
893    LatentBinary {
894        frailty: FrailtySpec,
895    },
896    TransformationNormal {
897        likelihood: LikelihoodSpec,
898    },
899}
900
901#[derive(Clone, Copy, Debug, Eq, PartialEq)]
902pub enum PredictModelClass {
903    Standard,
904    GaussianLocationScale,
905    BinomialLocationScale,
906    /// Genuine-dispersion location-scale (#913): NegativeBinomial / Gamma / Beta
907    /// / Tweedie mean families fitted with a `noise_formula` overdispersion
908    /// channel. Predicted through the GLM mean inverse link (not the binomial
909    /// threshold-scale predictor).
910    DispersionLocationScale,
911    BernoulliMarginalSlope,
912    Survival,
913    TransformationNormal,
914}
915
916impl PredictModelClass {
917    #[inline]
918    pub const fn name(self) -> &'static str {
919        match self {
920            Self::Standard => "standard",
921            Self::GaussianLocationScale => "gaussian location-scale",
922            Self::BinomialLocationScale => "binomial location-scale",
923            Self::DispersionLocationScale => "dispersion location-scale",
924            Self::BernoulliMarginalSlope => "bernoulli marginal-slope",
925            Self::Survival => "survival",
926            Self::TransformationNormal => "transformation-normal",
927        }
928    }
929}
930
931#[derive(Clone, Debug)]
932pub struct SavedLinkWiggleRuntime {
933    pub knots: Vec<f64>,
934    pub degree: usize,
935    pub beta: Vec<f64>,
936    /// Frozen-index mean-coordinate shift `s` (#2141). When present the predict
937    /// layer evaluates the warp basis at the frozen index
938    /// `η̂ = base + X·s` rather than at the de-aliased base predictor `base`, so
939    /// predict reproduces the exact `q` (and deviance) the fit computed. `None`
940    /// for warp paths that never de-aliased (location-scale / dynamic-basis),
941    /// where the base predictor already is the warp index.
942    pub index_shift: Option<Vec<f64>>,
943}
944
945#[derive(Clone, Debug)]
946pub struct SavedBaselineTimeWiggleRuntime {
947    pub knots: Vec<f64>,
948    pub degree: usize,
949    pub penalty_orders: Vec<usize>,
950    pub double_penalty: bool,
951    pub beta: Vec<f64>,
952}
953
954// Re-export so saved-model consumers can refer to the anchor-block tag
955// without reaching across module boundaries.
956pub use crate::bms::deviation_runtime::ParametricAnchorBlock;
957
958#[derive(Clone, Debug, Serialize, Deserialize)]
959pub struct SavedCompiledFlexBlock {
960    pub kernel: String,
961    pub breakpoints: Vec<f64>,
962    pub basis_dim: usize,
963    pub span_c0: Vec<Vec<f64>>,
964    pub span_c1: Vec<Vec<f64>>,
965    pub span_c2: Vec<Vec<f64>>,
966    pub span_c3: Vec<Vec<f64>>,
967    /// Cross-block anchor-residual coefficient matrix `M` of shape
968    /// `d × basis_dim`. When present, predict-time evaluation subtracts
969    /// `n_row · M` from each cubic-span row (where `n_row` stacks the
970    /// per-row parametric anchor values in the order given by
971    /// `anchor_components`).
972    #[serde(default)]
973    pub anchor_correction: Option<Vec<Vec<f64>>>,
974    /// Ordered list of parametric anchor components whose stacked row
975    /// values combine into `n_row`. Empty unless
976    /// `anchor_correction` is `Some`.
977    #[serde(default)]
978    pub anchor_components: Vec<SavedAnchorComponent>,
979}
980
981#[derive(Clone, Debug, Serialize, Deserialize)]
982pub struct SavedAnchorComponent {
983    pub kind: SavedAnchorKind,
984}
985
986#[derive(Clone, Debug, Serialize, Deserialize)]
987pub enum SavedAnchorKind {
988    Parametric {
989        block: ParametricAnchorBlock,
990        ncols: usize,
991    },
992    /// Flex-evaluation anchor (sibling flex block's reparameterised basis,
993    /// evaluated at training rows at fit time and at predict rows at
994    /// predict time). The predictor stacks `ncols` columns from the
995    /// sibling runtime's `design(arg)` into `n_row`.
996    FlexEvaluation { ncols: usize },
997}
998
999#[derive(Clone, Debug)]
1000pub struct SavedPredictionRuntime {
1001    pub model_class: PredictModelClass,
1002    pub likelihood: LikelihoodSpec,
1003    pub inverse_link: Option<InverseLink>,
1004    pub link_wiggle: Option<SavedLinkWiggleRuntime>,
1005    pub baseline_time_wiggle: Option<SavedBaselineTimeWiggleRuntime>,
1006    pub score_warp: Option<SavedCompiledFlexBlock>,
1007    pub link_deviation: Option<SavedCompiledFlexBlock>,
1008    /// Rank-INT latent-z calibration carried into the predictor build.
1009    /// `None` for non-BMS models and for BMS fits whose latent measure
1010    /// did not require rank-INT calibration.
1011    pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
1012    /// Conditional location-scale latent-z calibration (#905) carried into the
1013    /// predictor build. `None` for non-BMS models and for BMS fits whose Auto
1014    /// path did not detect a conditional `E[z|C]`/`Var(z|C)` shift. When
1015    /// `Some`, the predictor replaces the normalized `z` by `ζ = (z−m(C))/√v(C)`
1016    /// using the marginal prediction design as the conditioning span.
1017    pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
1018    /// Width `p₁` of the absorbed Stage-1 influence block (#461) when the
1019    /// survival marginal-slope fit hosted a dedicated additive absorber (the
1020    /// trailing block). `None` when no CTN Stage-1 chain produced an influence
1021    /// Jacobian. At predict the absorber's `γ` is DROPPED (the orthogonalized
1022    /// β̂ is a training-fit property), so the predictor uses this width only to
1023    /// (a) account for the extra trailing block in the saved block count and
1024    /// (b) slice `γ`'s rows/cols out of the joint covariance. Survival hosts the
1025    /// absorber as its own block (unlike the BMS A2 widened-marginal design),
1026    /// so it never widens any persisted prediction design.
1027    pub influence_absorber_width: Option<usize>,
1028}
1029
1030pub fn gaussian_location_scale_mean_beta(fit: &UnifiedFitResult) -> Option<Array1<f64>> {
1031    fit.block_by_role(BlockRole::Location)
1032        .or_else(|| fit.block_by_role(BlockRole::Mean))
1033        .map(|block| block.beta.clone())
1034}
1035
1036pub fn binomial_location_scale_threshold_beta(fit: &UnifiedFitResult) -> Option<Array1<f64>> {
1037    fit.block_by_role(BlockRole::Threshold)
1038        .or_else(|| fit.block_by_role(BlockRole::Location))
1039        .or_else(|| fit.block_by_role(BlockRole::Mean))
1040        .map(|block| block.beta.clone())
1041}
1042
1043pub fn location_scale_noise_beta(fit: &UnifiedFitResult) -> Option<Array1<f64>> {
1044    fit.block_by_role(BlockRole::Scale)
1045        .map(|block| block.beta.clone())
1046}
1047
1048/// Whether a `ModelKind::LocationScale` likelihood's response is one of the
1049/// genuine-dispersion mean families (#913) — NegativeBinomial, Gamma, Beta or
1050/// Tweedie. These carry a `noise_formula` overdispersion channel and must be
1051/// predicted through the GLM mean inverse link (the
1052/// [`PredictModelClass::DispersionLocationScale`] path), NOT the binomial
1053/// threshold-scale predictor. The binomial location-scale (BMS ordinal) path is
1054/// the only other non-Gaussian location-scale family, with a `Binomial`
1055/// response.
1056fn is_dispersion_location_scale_response(response: &gam_problem::types::ResponseFamily) -> bool {
1057    use gam_problem::types::ResponseFamily;
1058    matches!(
1059        response,
1060        ResponseFamily::NegativeBinomial { .. }
1061            | ResponseFamily::Gamma
1062            | ResponseFamily::Beta { .. }
1063            | ResponseFamily::Tweedie { .. }
1064    )
1065}
1066
1067fn validate_location_scale_saved_fit(
1068    fit: &UnifiedFitResult,
1069    model_class: PredictModelClass,
1070    link_wiggle: Option<&SavedLinkWiggleRuntime>,
1071) -> Result<(), FittedModelError> {
1072    let primary = match model_class {
1073        // Gaussian and dispersion (#913) location-scale both predict the mean
1074        // through the Location block; the binomial threshold-scale class reads
1075        // the Threshold block instead.
1076        PredictModelClass::GaussianLocationScale | PredictModelClass::DispersionLocationScale => {
1077            gaussian_location_scale_mean_beta(fit)
1078        }
1079        PredictModelClass::BinomialLocationScale => binomial_location_scale_threshold_beta(fit),
1080        _ => None,
1081    }
1082    .ok_or_else(|| FittedModelError::MissingField {
1083        reason: match model_class {
1084            PredictModelClass::GaussianLocationScale => {
1085                "gaussian-location-scale saved fit is missing mean/location block".to_string()
1086            }
1087            PredictModelClass::DispersionLocationScale => {
1088                "dispersion-location-scale saved fit is missing mean/location block".to_string()
1089            }
1090            PredictModelClass::BinomialLocationScale => {
1091                "binomial-location-scale saved fit is missing threshold/location block".to_string()
1092            }
1093            _ => "location-scale saved fit is missing primary block".to_string(),
1094        },
1095    })?;
1096
1097    let scale = location_scale_noise_beta(fit).ok_or_else(|| FittedModelError::MissingField {
1098        reason: "location-scale saved fit is missing scale block".to_string(),
1099    })?;
1100    let expected =
1101        primary.len() + scale.len() + link_wiggle.map_or(0, |runtime| runtime.beta.len());
1102
1103    if let Some(cov) = fit.beta_covariance()
1104        && (cov.nrows() != expected || cov.ncols() != expected)
1105    {
1106        return Err(FittedModelError::SchemaMismatch {
1107            reason: format!(
1108                "location-scale saved conditional covariance shape mismatch: got {}x{}, expected {}x{}",
1109                cov.nrows(),
1110                cov.ncols(),
1111                expected,
1112                expected
1113            ),
1114        });
1115    }
1116    if let Some(cov) = fit.beta_covariance_corrected()
1117        && (cov.nrows() != expected || cov.ncols() != expected)
1118    {
1119        return Err(FittedModelError::SchemaMismatch {
1120            reason: format!(
1121                "location-scale saved corrected covariance shape mismatch: got {}x{}, expected {}x{}",
1122                cov.nrows(),
1123                cov.ncols(),
1124                expected,
1125                expected
1126            ),
1127        });
1128    }
1129    Ok(())
1130}
1131
1132fn validate_survival_saved_block_matches_payload(
1133    fit: &UnifiedFitResult,
1134    role: BlockRole,
1135    payload_beta: Option<&Vec<f64>>,
1136    label: &str,
1137) -> Result<usize, FittedModelError> {
1138    let block = fit
1139        .block_by_role(role)
1140        .ok_or_else(|| FittedModelError::MissingField {
1141            reason: format!("location-scale survival saved fit is missing {label} block"),
1142        })?;
1143    if let Some(saved) = payload_beta
1144        && block.beta.to_vec() != *saved
1145    {
1146        return Err(FittedModelError::SchemaMismatch {
1147            reason: format!(
1148                "location-scale survival saved {label} coefficients disagree with fit_result"
1149            ),
1150        });
1151    }
1152    Ok(block.beta.len())
1153}
1154
1155fn validate_survival_location_scale_saved_fit(
1156    payload: &FittedModelPayload,
1157    link_wiggle: Option<&SavedLinkWiggleRuntime>,
1158) -> Result<(), FittedModelError> {
1159    let fit = payload
1160        .fit_result
1161        .as_ref()
1162        .ok_or_else(|| FittedModelError::MissingField {
1163            reason: "location-scale survival model is missing canonical fit_result payload"
1164                .to_string(),
1165        })?;
1166    let p_time = validate_survival_saved_block_matches_payload(
1167        fit,
1168        BlockRole::Time,
1169        payload.survival_beta_time.as_ref(),
1170        "time",
1171    )?;
1172    let p_threshold = validate_survival_saved_block_matches_payload(
1173        fit,
1174        BlockRole::Threshold,
1175        payload.survival_beta_threshold.as_ref(),
1176        "threshold",
1177    )?;
1178    let p_log_sigma = validate_survival_saved_block_matches_payload(
1179        fit,
1180        BlockRole::Scale,
1181        payload.survival_beta_log_sigma.as_ref(),
1182        "log-sigma",
1183    )?;
1184    let p_wiggle = match link_wiggle {
1185        Some(runtime) => {
1186            let block = fit.block_by_role(BlockRole::LinkWiggle).ok_or_else(|| {
1187                FittedModelError::MissingField {
1188                    reason: "location-scale survival saved fit is missing link-wiggle block"
1189                        .to_string(),
1190                }
1191            })?;
1192            if block.beta.to_vec() != runtime.beta {
1193                return Err(FittedModelError::SchemaMismatch {
1194                    reason:
1195                        "location-scale survival saved link-wiggle coefficients disagree with fit_result"
1196                            .to_string(),
1197                });
1198            }
1199            runtime.beta.len()
1200        }
1201        None => {
1202            if fit.block_by_role(BlockRole::LinkWiggle).is_some() {
1203                return Err(FittedModelError::SchemaMismatch {
1204                    reason:
1205                        "location-scale survival saved fit has a LinkWiggle block without payload metadata"
1206                            .to_string(),
1207                });
1208            }
1209            0
1210        }
1211    };
1212    let expected = p_time + p_threshold + p_log_sigma + p_wiggle;
1213
1214    if let Some(cov) = fit.beta_covariance()
1215        && (cov.nrows() != expected || cov.ncols() != expected)
1216    {
1217        return Err(FittedModelError::SchemaMismatch {
1218            reason: format!(
1219                "location-scale survival saved conditional covariance shape mismatch: got {}x{}, expected {}x{}",
1220                cov.nrows(),
1221                cov.ncols(),
1222                expected,
1223                expected
1224            ),
1225        });
1226    }
1227    if let Some(cov) = fit.beta_covariance_corrected()
1228        && (cov.nrows() != expected || cov.ncols() != expected)
1229    {
1230        return Err(FittedModelError::SchemaMismatch {
1231            reason: format!(
1232                "location-scale survival saved corrected covariance shape mismatch: got {}x{}, expected {}x{}",
1233                cov.nrows(),
1234                cov.ncols(),
1235                expected,
1236                expected
1237            ),
1238        });
1239    }
1240    Ok(())
1241}
1242
1243fn validate_marginal_slope_saved_fit(
1244    fit: &UnifiedFitResult,
1245    score_warp: Option<&SavedCompiledFlexBlock>,
1246    link_deviation: Option<&SavedCompiledFlexBlock>,
1247    fit_label: &str,
1248) -> Result<(), FittedModelError> {
1249    validate_marginal_slope_saved_fit_impl(
1250        fit,
1251        score_warp,
1252        link_deviation,
1253        fit_label,
1254        "bernoulli",
1255        2,
1256        "marginal, logslope",
1257    )
1258}
1259
1260fn validate_survival_marginal_slope_saved_fit(
1261    fit: &UnifiedFitResult,
1262    score_warp: Option<&SavedCompiledFlexBlock>,
1263    link_deviation: Option<&SavedCompiledFlexBlock>,
1264    fit_label: &str,
1265) -> Result<(), FittedModelError> {
1266    validate_marginal_slope_saved_fit_impl(
1267        fit,
1268        score_warp,
1269        link_deviation,
1270        fit_label,
1271        "survival",
1272        3,
1273        "time, marginal, slope",
1274    )
1275}
1276
1277/// Shared block-count + coefficient-dimension validation for the bernoulli
1278/// and survival marginal-slope saved-fit gates. The only family-specific
1279/// inputs are the family kind string ("bernoulli" / "survival"), the base
1280/// block count (2 for bernoulli, 3 for survival — the survival path has an
1281/// extra time block), and the base-block role list rendered in the error
1282/// message ("marginal, logslope" / "time, marginal, slope"). The score-warp
1283/// / link-deviation tail follows the same shape in both families.
1284fn validate_marginal_slope_saved_fit_impl(
1285    fit: &UnifiedFitResult,
1286    score_warp: Option<&SavedCompiledFlexBlock>,
1287    link_deviation: Option<&SavedCompiledFlexBlock>,
1288    fit_label: &str,
1289    family_kind: &str,
1290    base_block_count: usize,
1291    base_block_role_list: &str,
1292) -> Result<(), FittedModelError> {
1293    let expected_blocks = base_block_count
1294        + usize::from(score_warp.is_some())
1295        + usize::from(link_deviation.is_some());
1296    if fit.blocks.len() != expected_blocks {
1297        let score_warp_suffix = if score_warp.is_some() {
1298            ", score-warp"
1299        } else {
1300            ""
1301        };
1302        let link_deviation_suffix = if link_deviation.is_some() {
1303            ", link-deviation"
1304        } else {
1305            ""
1306        };
1307        return Err(FittedModelError::SchemaMismatch {
1308            reason: format!(
1309                "{family_kind} marginal-slope saved {fit_label} requires {expected_blocks} blocks [{base_block_role_list}{score_warp_suffix}{link_deviation_suffix}], got {}",
1310                fit.blocks.len(),
1311            ),
1312        });
1313    }
1314    if let Some(runtime) = score_warp {
1315        let beta = &fit.blocks[base_block_count].beta;
1316        if beta.len() != runtime.basis_dim {
1317            return Err(FittedModelError::SchemaMismatch {
1318                reason: format!(
1319                    "{family_kind} marginal-slope saved {fit_label} score-warp coefficient mismatch: beta has {} entries but runtime expects {}",
1320                    beta.len(),
1321                    runtime.basis_dim
1322                ),
1323            });
1324        }
1325    }
1326    if let Some(runtime) = link_deviation {
1327        let idx = base_block_count + usize::from(score_warp.is_some());
1328        let beta = &fit.blocks[idx].beta;
1329        if beta.len() != runtime.basis_dim {
1330            return Err(FittedModelError::SchemaMismatch {
1331                reason: format!(
1332                    "{family_kind} marginal-slope saved {fit_label} link-deviation coefficient mismatch: beta has {} entries but runtime expects {}",
1333                    beta.len(),
1334                    runtime.basis_dim
1335                ),
1336            });
1337        }
1338    }
1339    Ok(())
1340}
1341
1342impl SavedLinkWiggleRuntime {
1343    fn validate_monotone_derivative(
1344        &self,
1345        q0: &Array1<f64>,
1346    ) -> Result<Array1<f64>, FittedModelError> {
1347        // Monotonicity is verified pointwise at the actual evaluation grid `q0`
1348        // (the predict η). The fit guarantees a strictly-increasing warped link
1349        // across the training η range (#1596); checking at `q0` here flags an
1350        // extrapolation point where the learnable link genuinely turns
1351        // non-invertible, without rejecting the whole model for a sign dip in the
1352        // basis tail far outside any data or evaluation point.
1353        let d_constrained = self.constrained_basis(q0, BasisOptions::first_derivative())?;
1354        let beta_link_wiggle = Array1::from_vec(self.beta.clone());
1355        let dq_dq0 = d_constrained.dot(&beta_link_wiggle) + 1.0;
1356        if let Some((idx, value)) = dq_dq0.iter().copied().enumerate().find(|(_, v)| *v <= 0.0) {
1357            return Err(FittedModelError::PayloadCorrupt {
1358                reason: format!(
1359                    "saved link-wiggle is not monotone at row {idx}: dq/dq0={value:.3e} <= 0"
1360                ),
1361            });
1362        }
1363        Ok(dq_dq0)
1364    }
1365
1366    pub fn constrained_basis(
1367        &self,
1368        q0: &Array1<f64>,
1369        basis_options: BasisOptions,
1370    ) -> Result<Array2<f64>, FittedModelError> {
1371        let knot_arr = Array1::from_vec(self.knots.clone());
1372        let constrained = monotone_wiggle_basis_with_derivative_order(
1373            q0.view(),
1374            &knot_arr,
1375            self.degree,
1376            basis_options.derivative_order,
1377        )
1378        .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
1379        if constrained.ncols() != self.beta.len() {
1380            return Err(FittedModelError::SchemaMismatch {
1381                reason: format!(
1382                    "saved link-wiggle dimension mismatch: coefficients have {} entries but basis has {} columns",
1383                    self.beta.len(),
1384                    constrained.ncols()
1385                ),
1386            });
1387        }
1388        Ok(constrained)
1389    }
1390
1391    pub fn design(&self, q0: &Array1<f64>) -> Result<Array2<f64>, FittedModelError> {
1392        self.validate_monotone_derivative(q0)?;
1393        self.constrained_basis(q0, BasisOptions::value())
1394    }
1395
1396    pub fn basis_row_scalar(&self, q0: f64) -> Result<Array1<f64>, FittedModelError> {
1397        let q = Array1::from_vec(vec![q0]);
1398        let x = self.design(&q)?;
1399        if x.nrows() != 1 {
1400            return Err(FittedModelError::SchemaMismatch {
1401                reason: format!(
1402                    "saved link-wiggle scalar evaluation expected 1 row, got {}",
1403                    x.nrows()
1404                ),
1405            });
1406        }
1407        Ok(x.row(0).to_owned())
1408    }
1409
1410    pub fn apply(&self, q0: &Array1<f64>) -> Result<Array1<f64>, FittedModelError> {
1411        self.apply_with_index(q0, q0)
1412    }
1413
1414    /// Apply the warp with the additive base predictor `base` and the warp basis
1415    /// evaluated at a possibly distinct index `warp_index` (#2141):
1416    /// `q = base + B(warp_index)·β`. The frozen-basis de-aliased binomial-mean
1417    /// link fit pins `B` at the frozen index `η̂`, which differs from the
1418    /// de-aliased base predictor `base` by `X·s`; evaluating `B` at `warp_index`
1419    /// (= `base + X·s`) reproduces the exact `q` the fit scored. Monotonicity is
1420    /// certified at `warp_index`, where the link's `dq/dη = 1 + B'(warp_index)·β`
1421    /// is defined. When no shift is active `warp_index == base` and this reduces
1422    /// to the original `q = q0 + B(q0)·β`.
1423    pub fn apply_with_index(
1424        &self,
1425        base: &Array1<f64>,
1426        warp_index: &Array1<f64>,
1427    ) -> Result<Array1<f64>, FittedModelError> {
1428        if base.len() != warp_index.len() {
1429            return Err(FittedModelError::SchemaMismatch {
1430                reason: format!(
1431                    "link-wiggle base predictor has {} rows but warp index has {}",
1432                    base.len(),
1433                    warp_index.len()
1434                ),
1435            });
1436        }
1437        self.validate_monotone_derivative(warp_index)?;
1438        let xwiggle = self.constrained_basis(warp_index, BasisOptions::value())?;
1439        let beta_link_wiggle = Array1::from_vec(self.beta.clone());
1440        Ok(base + &xwiggle.dot(&beta_link_wiggle))
1441    }
1442
1443    pub fn derivative_q0(&self, q0: &Array1<f64>) -> Result<Array1<f64>, FittedModelError> {
1444        self.validate_monotone_derivative(q0)
1445    }
1446}
1447
1448impl SavedBaselineTimeWiggleRuntime {
1449    pub fn validate_global_monotonicity(&self) -> Result<(), FittedModelError> {
1450        validate_monotone_wiggle_beta_nonnegative(&self.beta, "saved baseline-timewiggle")
1451            .map_err(|reason| FittedModelError::PayloadCorrupt { reason })
1452    }
1453}
1454
1455impl SavedCompiledFlexBlock {
1456    pub(crate) fn validate_exact_replay_contract(&self) -> Result<(), FittedModelError> {
1457        if self.kernel.is_empty() {
1458            return Err(FittedModelError::SchemaMismatch {
1459                reason: "saved anchored deviation runtime is missing the exact kernel marker"
1460                    .to_string(),
1461            });
1462        }
1463        if self.kernel != crate::cubic_cell_kernel::ANCHORED_DEVIATION_KERNEL {
1464            return Err(FittedModelError::IncompatibleConfig {
1465                reason: format!(
1466                    "saved anchored deviation runtime uses unsupported kernel '{}'; expected {}",
1467                    self.kernel,
1468                    crate::cubic_cell_kernel::ANCHORED_DEVIATION_KERNEL
1469                ),
1470            });
1471        }
1472        if self.basis_dim == 0 {
1473            return Err(FittedModelError::SchemaMismatch {
1474                reason: format!(
1475                    "saved anchored deviation runtime basis_dim must be positive, got {}",
1476                    self.basis_dim
1477                ),
1478            });
1479        }
1480        if self.breakpoints.len() < 2 {
1481            return Err(FittedModelError::SchemaMismatch {
1482                reason: format!(
1483                    "saved anchored deviation runtime requires at least two breakpoints, got {}",
1484                    self.breakpoints.len()
1485                ),
1486            });
1487        }
1488        for window in self.breakpoints.windows(2) {
1489            let left = window[0];
1490            let right = window[1];
1491            if !left.is_finite() || !right.is_finite() || right <= left {
1492                return Err(FittedModelError::PayloadCorrupt {
1493                    reason: format!(
1494                        "saved anchored deviation runtime breakpoints must be finite and strictly increasing, got [{left}, {right}]"
1495                    ),
1496                });
1497            }
1498        }
1499        let span_count = self.breakpoints.len() - 1;
1500        self.validate_coefficient_matrix(&self.span_c0, "c0", span_count)?;
1501        self.validate_coefficient_matrix(&self.span_c1, "c1", span_count)?;
1502        self.validate_coefficient_matrix(&self.span_c2, "c2", span_count)?;
1503        self.validate_coefficient_matrix(&self.span_c3, "c3", span_count)?;
1504        self.validate_c2_span_continuity()?;
1505        self.validate_anchor_residual_shape()?;
1506        Ok(())
1507    }
1508
1509    fn validate_anchor_residual_shape(&self) -> Result<(), FittedModelError> {
1510        let coeffs = match self.anchor_correction.as_ref() {
1511            Some(c) => c,
1512            None => {
1513                if !self.anchor_components.is_empty() {
1514                    return Err(FittedModelError::SchemaMismatch {
1515                        reason:
1516                            "saved anchored deviation runtime has anchor_components but no anchor_correction"
1517                                .to_string(),
1518                    });
1519                }
1520                return Ok(());
1521            }
1522        };
1523        let d: usize = self
1524            .anchor_components
1525            .iter()
1526            .map(|c| match &c.kind {
1527                SavedAnchorKind::Parametric { ncols, .. } => *ncols,
1528                SavedAnchorKind::FlexEvaluation { ncols } => *ncols,
1529            })
1530            .sum();
1531        if coeffs.len() != d {
1532            return Err(FittedModelError::SchemaMismatch {
1533                reason: format!(
1534                    "saved anchored deviation runtime anchor_correction has {} rows; expected {} (sum of component ncols)",
1535                    coeffs.len(),
1536                    d,
1537                ),
1538            });
1539        }
1540        for (i, row) in coeffs.iter().enumerate() {
1541            if row.len() != self.basis_dim {
1542                return Err(FittedModelError::SchemaMismatch {
1543                    reason: format!(
1544                        "saved anchored deviation runtime anchor_correction row {} has width {}, expected basis_dim {}",
1545                        i,
1546                        row.len(),
1547                        self.basis_dim,
1548                    ),
1549                });
1550            }
1551            for (j, &v) in row.iter().enumerate() {
1552                if !v.is_finite() {
1553                    return Err(FittedModelError::PayloadCorrupt {
1554                        reason: format!(
1555                            "saved anchored deviation runtime anchor_correction ({i},{j}) is non-finite"
1556                        ),
1557                    });
1558                }
1559            }
1560        }
1561        Ok(())
1562    }
1563
1564    fn validate_c2_span_continuity(&self) -> Result<(), FittedModelError> {
1565        const TOL: f64 = 1e-8;
1566        for span_idx in 1..self.breakpoints.len() - 1 {
1567            let left_span = span_idx - 1;
1568            let right_span = span_idx;
1569            let width = self.breakpoints[span_idx] - self.breakpoints[left_span];
1570            for basis_idx in 0..self.basis_dim {
1571                let left_value = self.span_c0[left_span][basis_idx]
1572                    + self.span_c1[left_span][basis_idx] * width
1573                    + self.span_c2[left_span][basis_idx] * width * width
1574                    + self.span_c3[left_span][basis_idx] * width * width * width;
1575                let left_d1 = self.span_c1[left_span][basis_idx]
1576                    + 2.0 * self.span_c2[left_span][basis_idx] * width
1577                    + 3.0 * self.span_c3[left_span][basis_idx] * width * width;
1578                let left_d2 = 2.0 * self.span_c2[left_span][basis_idx]
1579                    + 6.0 * self.span_c3[left_span][basis_idx] * width;
1580                let right_value = self.span_c0[right_span][basis_idx];
1581                let right_d1 = self.span_c1[right_span][basis_idx];
1582                let right_d2 = 2.0 * self.span_c2[right_span][basis_idx];
1583                if (left_value - right_value).abs() > TOL
1584                    || (left_d1 - right_d1).abs() > TOL
1585                    || (left_d2 - right_d2).abs() > TOL
1586                {
1587                    return Err(FittedModelError::SchemaMismatch {
1588                        reason: format!(
1589                            "saved anchored deviation runtime must be C2 cubic at breakpoint {span_idx}, basis {basis_idx}: value jump={:.3e}, d1 jump={:.3e}, d2 jump={:.3e}",
1590                            left_value - right_value,
1591                            left_d1 - right_d1,
1592                            left_d2 - right_d2
1593                        ),
1594                    });
1595                }
1596            }
1597        }
1598        Ok(())
1599    }
1600
1601    fn validate_coefficient_matrix(
1602        &self,
1603        matrix: &[Vec<f64>],
1604        label: &str,
1605        expected_rows: usize,
1606    ) -> Result<(), FittedModelError> {
1607        if matrix.len() != expected_rows {
1608            return Err(FittedModelError::SchemaMismatch {
1609                reason: format!(
1610                    "saved anchored deviation runtime {label} row count mismatch: got {}, expected {}",
1611                    matrix.len(),
1612                    expected_rows
1613                ),
1614            });
1615        }
1616        for (row_idx, row) in matrix.iter().enumerate() {
1617            if row.len() != self.basis_dim {
1618                return Err(FittedModelError::SchemaMismatch {
1619                    reason: format!(
1620                        "saved anchored deviation runtime {label} row {} has width {}, expected {}",
1621                        row_idx,
1622                        row.len(),
1623                        self.basis_dim
1624                    ),
1625                });
1626            }
1627            for (j, &value) in row.iter().enumerate() {
1628                if !value.is_finite() {
1629                    return Err(FittedModelError::PayloadCorrupt {
1630                        reason: format!(
1631                            "saved anchored deviation runtime {label} entry ({row_idx},{j}) is non-finite"
1632                        ),
1633                    });
1634                }
1635            }
1636        }
1637        Ok(())
1638    }
1639
1640    fn right_boundary_basis_value(&self, basis_idx: usize) -> f64 {
1641        let last_span = self.breakpoints.len() - 2;
1642        let width = self.breakpoints[last_span + 1] - self.breakpoints[last_span];
1643        self.span_c0[last_span][basis_idx]
1644            + self.span_c1[last_span][basis_idx] * width
1645            + self.span_c2[last_span][basis_idx] * width * width
1646            + self.span_c3[last_span][basis_idx] * width * width * width
1647    }
1648
1649    fn evaluate_span_polynomial_design(
1650        &self,
1651        values: &Array1<f64>,
1652        derivative_order: usize,
1653    ) -> Result<Array2<f64>, FittedModelError> {
1654        self.validate_exact_replay_contract()?;
1655        let (left_ep, right_ep) = self.support_interval()?;
1656        let mut out = Array2::<f64>::zeros((values.len(), self.basis_dim));
1657        for (row_idx, &value) in values.iter().enumerate() {
1658            if !value.is_finite() {
1659                return Err(FittedModelError::PayloadCorrupt {
1660                    reason: format!(
1661                        "saved anchored deviation runtime design value at row {row_idx} is non-finite ({value})"
1662                    ),
1663                });
1664            }
1665            if value < left_ep {
1666                if derivative_order == 0 {
1667                    for basis_idx in 0..self.basis_dim {
1668                        out[[row_idx, basis_idx]] = self.span_c0[0][basis_idx];
1669                    }
1670                }
1671                continue;
1672            }
1673            if value > right_ep {
1674                if derivative_order == 0 {
1675                    for basis_idx in 0..self.basis_dim {
1676                        out[[row_idx, basis_idx]] = self.right_boundary_basis_value(basis_idx);
1677                    }
1678                }
1679                continue;
1680            }
1681            let span_idx = self.left_biased_span_index_for(value)?;
1682            let t = value - self.breakpoints[span_idx];
1683            for basis_idx in 0..self.basis_dim {
1684                let c0 = self.span_c0[span_idx][basis_idx];
1685                let c1 = self.span_c1[span_idx][basis_idx];
1686                let c2 = self.span_c2[span_idx][basis_idx];
1687                let c3 = self.span_c3[span_idx][basis_idx];
1688                out[[row_idx, basis_idx]] = match derivative_order {
1689                    0 => c0 + c1 * t + c2 * t * t + c3 * t * t * t,
1690                    1 => c1 + 2.0 * c2 * t + 3.0 * c3 * t * t,
1691                    2 => 2.0 * c2 + 6.0 * c3 * t,
1692                    3 => 6.0 * c3,
1693                    4 => 0.0,
1694                    other => {
1695                        return Err(FittedModelError::IncompatibleConfig {
1696                            reason: format!(
1697                                "saved anchored deviation runtime only supports derivative orders up to 4, got {other}"
1698                            ),
1699                        });
1700                    }
1701                };
1702            }
1703        }
1704        Ok(out)
1705    }
1706
1707    pub fn breakpoints(&self) -> Result<Vec<f64>, FittedModelError> {
1708        self.validate_exact_replay_contract()?;
1709        Ok(self.breakpoints.clone())
1710    }
1711
1712    pub fn span_count(&self) -> Result<usize, FittedModelError> {
1713        Ok(self.breakpoints()?.windows(2).count())
1714    }
1715
1716    pub fn span_index_for(&self, value: f64) -> Result<usize, FittedModelError> {
1717        let points = self.breakpoints()?;
1718        span_index_for_breakpoints(&points, value, "saved anchored deviation span lookup")
1719            .map_err(|reason| FittedModelError::PayloadCorrupt { reason })
1720    }
1721
1722    fn left_biased_span_index_for(&self, value: f64) -> Result<usize, FittedModelError> {
1723        let mut span_idx = span_index_for_breakpoints(
1724            &self.breakpoints,
1725            value,
1726            "saved anchored deviation span lookup",
1727        )
1728        .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
1729        // LEFT-bias at interior breakpoints mirrors DeviationRuntime. The
1730        // saved cubic basis is C2, but d3 remains span-local.
1731        if span_idx > 0 && value == self.breakpoints[span_idx] {
1732            span_idx -= 1;
1733        }
1734        Ok(span_idx)
1735    }
1736
1737    pub fn local_cubic_on_span(
1738        &self,
1739        beta: &Array1<f64>,
1740        span_idx: usize,
1741    ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
1742        self.validate_exact_replay_contract()?;
1743        if beta.len() != self.basis_dim {
1744            return Err(FittedModelError::SchemaMismatch {
1745                reason: format!(
1746                    "saved anchored deviation coefficient length mismatch: got {}, expected {}",
1747                    beta.len(),
1748                    self.basis_dim
1749                ),
1750            });
1751        }
1752        self.local_cubic_on_span_validated(beta, span_idx)
1753    }
1754
1755    fn local_cubic_on_span_validated(
1756        &self,
1757        beta: &Array1<f64>,
1758        span_idx: usize,
1759    ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
1760        let points = &self.breakpoints;
1761        if span_idx + 1 >= points.len() {
1762            return Err(FittedModelError::SchemaMismatch {
1763                reason: format!(
1764                    "saved anchored deviation span index {} out of range for {} spans",
1765                    span_idx,
1766                    points.len() - 1
1767                ),
1768            });
1769        }
1770        let left = points[span_idx];
1771        let right = points[span_idx + 1];
1772        Ok(crate::cubic_cell_kernel::LocalSpanCubic {
1773            left,
1774            right,
1775            c0: self.span_c0[span_idx]
1776                .iter()
1777                .zip(beta.iter())
1778                .map(|(coeff, weight)| coeff * weight)
1779                .sum(),
1780            c1: self.span_c1[span_idx]
1781                .iter()
1782                .zip(beta.iter())
1783                .map(|(coeff, weight)| coeff * weight)
1784                .sum(),
1785            c2: self.span_c2[span_idx]
1786                .iter()
1787                .zip(beta.iter())
1788                .map(|(coeff, weight)| coeff * weight)
1789                .sum(),
1790            c3: self.span_c3[span_idx]
1791                .iter()
1792                .zip(beta.iter())
1793                .map(|(coeff, weight)| coeff * weight)
1794                .sum(),
1795        })
1796    }
1797
1798    pub fn basis_span_cubic(
1799        &self,
1800        span_idx: usize,
1801        basis_idx: usize,
1802    ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
1803        self.validate_exact_replay_contract()?;
1804        if basis_idx >= self.basis_dim {
1805            return Err(FittedModelError::SchemaMismatch {
1806                reason: format!(
1807                    "saved anchored deviation basis index {} out of range for {} coefficients",
1808                    basis_idx, self.basis_dim
1809                ),
1810            });
1811        }
1812        self.basis_span_cubic_validated(span_idx, basis_idx)
1813    }
1814
1815    fn basis_span_cubic_validated(
1816        &self,
1817        span_idx: usize,
1818        basis_idx: usize,
1819    ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
1820        let points = &self.breakpoints;
1821        if span_idx + 1 >= points.len() {
1822            return Err(FittedModelError::SchemaMismatch {
1823                reason: format!(
1824                    "saved anchored deviation span index {} out of range for {} spans",
1825                    span_idx,
1826                    points.len() - 1
1827                ),
1828            });
1829        }
1830        Ok(crate::cubic_cell_kernel::LocalSpanCubic {
1831            left: points[span_idx],
1832            right: points[span_idx + 1],
1833            c0: self.span_c0[span_idx][basis_idx],
1834            c1: self.span_c1[span_idx][basis_idx],
1835            c2: self.span_c2[span_idx][basis_idx],
1836            c3: self.span_c3[span_idx][basis_idx],
1837        })
1838    }
1839
1840    pub fn basis_cubic_at(
1841        &self,
1842        basis_idx: usize,
1843        value: f64,
1844    ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
1845        self.validate_exact_replay_contract()?;
1846        if basis_idx >= self.basis_dim {
1847            return Err(FittedModelError::SchemaMismatch {
1848                reason: format!(
1849                    "saved anchored deviation basis index {} out of range for {} coefficients",
1850                    basis_idx, self.basis_dim
1851                ),
1852            });
1853        }
1854        let (left_ep, right_ep) = self.support_interval()?;
1855        if value < left_ep {
1856            return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
1857                left: left_ep,
1858                right: left_ep + 1.0,
1859                c0: self.span_c0[0][basis_idx],
1860                c1: 0.0,
1861                c2: 0.0,
1862                c3: 0.0,
1863            });
1864        }
1865        if value > right_ep {
1866            return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
1867                left: right_ep,
1868                right: right_ep + 1.0,
1869                c0: self.right_boundary_basis_value(basis_idx),
1870                c1: 0.0,
1871                c2: 0.0,
1872                c3: 0.0,
1873            });
1874        }
1875        let span_idx = self.left_biased_span_index_for(value)?;
1876        self.basis_span_cubic_validated(span_idx, basis_idx)
1877    }
1878
1879    pub fn local_cubic_at(
1880        &self,
1881        beta: &Array1<f64>,
1882        value: f64,
1883    ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
1884        self.validate_exact_replay_contract()?;
1885        if beta.len() != self.basis_dim {
1886            return Err(FittedModelError::SchemaMismatch {
1887                reason: format!(
1888                    "saved anchored deviation coefficient length mismatch: got {}, expected {}",
1889                    beta.len(),
1890                    self.basis_dim
1891                ),
1892            });
1893        }
1894        let (left_ep, right_ep) = self.support_interval()?;
1895        if value < left_ep {
1896            return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
1897                left: left_ep,
1898                right: left_ep + 1.0,
1899                c0: self.span_c0[0]
1900                    .iter()
1901                    .zip(beta.iter())
1902                    .map(|(coeff, weight)| coeff * weight)
1903                    .sum(),
1904                c1: 0.0,
1905                c2: 0.0,
1906                c3: 0.0,
1907            });
1908        }
1909        if value > right_ep {
1910            return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
1911                left: right_ep,
1912                right: right_ep + 1.0,
1913                c0: (0..self.basis_dim)
1914                    .map(|basis_idx| self.right_boundary_basis_value(basis_idx) * beta[basis_idx])
1915                    .sum(),
1916                c1: 0.0,
1917                c2: 0.0,
1918                c3: 0.0,
1919            });
1920        }
1921        let span_idx = self.left_biased_span_index_for(value)?;
1922        self.local_cubic_on_span_validated(beta, span_idx)
1923    }
1924
1925    fn support_interval(&self) -> Result<(f64, f64), FittedModelError> {
1926        let points = self.breakpoints()?;
1927        match (points.first(), points.last()) {
1928            (Some(&left), Some(&right)) => Ok((left, right)),
1929            _ => Err(FittedModelError::MissingField {
1930                reason: "saved anchored deviation runtime is missing support breakpoints"
1931                    .to_string(),
1932            }),
1933        }
1934    }
1935
1936    pub fn design(&self, values: &Array1<f64>) -> Result<Array2<f64>, FittedModelError> {
1937        // Note: when the saved runtime carries an anchor residual
1938        // (cross-block orthogonalisation), the value `design()` returns
1939        // is the raw cubic span output *without* the per-row `n_row · M`
1940        // subtraction. Callers used inside BMS prediction must either
1941        // switch to `design_with_anchor_rows` (when the per-row anchor
1942        // rows are available) or call `design_uncorrected` explicitly and
1943        // apply the subtraction at the call site. For runtimes without a
1944        // residual the two paths coincide.
1945        self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)
1946    }
1947
1948    /// Raw cubic-span design without any anchor-residual subtraction.
1949    ///
1950    /// Exposed for callers that intend to apply the `n_row · M` correction
1951    /// post-hoc (e.g., BMS `link_terms_value_d1` subtracts a precomputed
1952    /// `correction.dot(beta)` scalar from the linear-predictor contribution
1953    /// rather than building a full anchor-row matrix). Equivalent to
1954    /// `design()` when no residual is present.
1955    pub fn design_uncorrected(
1956        &self,
1957        values: &Array1<f64>,
1958    ) -> Result<Array2<f64>, FittedModelError> {
1959        self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)
1960    }
1961
1962    /// Evaluate the residual-corrected design at the supplied values.
1963    ///
1964    /// `anchor_rows` must be an `n × d` matrix where `n == values.len()`
1965    /// and `d == sum of anchor_components ncols`. Each row holds
1966    /// the concatenated parametric anchor design at the same prediction
1967    /// row as the corresponding `values[i]`. When the runtime has no
1968    /// anchor residual, `anchor_rows` must have zero columns (or be
1969    /// `Array2::zeros((n, 0))`).
1970    pub fn design_with_anchor_rows(
1971        &self,
1972        values: &Array1<f64>,
1973        anchor_rows: ndarray::ArrayView2<f64>,
1974    ) -> Result<Array2<f64>, FittedModelError> {
1975        let mut out =
1976            self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)?;
1977        if let Some(m_rows) = self.anchor_correction.as_ref() {
1978            let d = m_rows.len();
1979            if anchor_rows.nrows() != values.len() {
1980                return Err(FittedModelError::SchemaMismatch {
1981                    reason: format!(
1982                        "design_with_anchor_rows: anchor_rows has {} rows, expected {} (matching values)",
1983                        anchor_rows.nrows(),
1984                        values.len(),
1985                    ),
1986                });
1987            }
1988            if anchor_rows.ncols() != d {
1989                return Err(FittedModelError::SchemaMismatch {
1990                    reason: format!(
1991                        "design_with_anchor_rows: anchor_rows has {} cols, expected {} (sum of component ncols)",
1992                        anchor_rows.ncols(),
1993                        d,
1994                    ),
1995                });
1996            }
1997            // Materialise M (d × basis_dim) once.
1998            let mut m_dense = Array2::<f64>::zeros((d, self.basis_dim));
1999            for (i, row) in m_rows.iter().enumerate() {
2000                if row.len() != self.basis_dim {
2001                    return Err(FittedModelError::SchemaMismatch {
2002                        reason: format!(
2003                            "design_with_anchor_rows: anchor_correction row {} has length {}, expected basis_dim {}",
2004                            i,
2005                            row.len(),
2006                            self.basis_dim,
2007                        ),
2008                    });
2009                }
2010                for (j, &v) in row.iter().enumerate() {
2011                    m_dense[[i, j]] = v;
2012                }
2013            }
2014            // The compiler bakes the orthonormalising rotation into M, so
2015            // the predict-time subtraction is simply `n_anchor_rows · M`.
2016            let subtract = anchor_rows.dot(&m_dense);
2017            out = out - subtract;
2018        } else if anchor_rows.ncols() != 0 {
2019            return Err(FittedModelError::SchemaMismatch {
2020                reason: format!(
2021                    "design_with_anchor_rows: runtime has no anchor residual but anchor_rows has {} cols",
2022                    anchor_rows.ncols(),
2023                ),
2024            });
2025        }
2026        Ok(out)
2027    }
2028
2029    /// Build the n × basis_dim per-row, per-basis correction matrix
2030    /// `N · M` for a batch of predict rows.
2031    ///
2032    /// `n_anchor_rows` is the n × d matrix of stacked parametric anchor
2033    /// rows at the prediction rows (concatenation of the marginal and
2034    /// logslope design rows in component order). Returns `None` when the
2035    /// runtime has no anchor residual (zero-cost path).
2036    pub fn anchor_correction_matrix(
2037        &self,
2038        n_anchor_rows: ndarray::ArrayView2<f64>,
2039    ) -> Result<Option<Array2<f64>>, FittedModelError> {
2040        let Some(m_rows) = self.anchor_correction.as_ref() else {
2041            return Ok(None);
2042        };
2043        let d = m_rows.len();
2044        if n_anchor_rows.ncols() != d {
2045            return Err(FittedModelError::SchemaMismatch {
2046                reason: format!(
2047                    "anchor_correction_matrix: anchor_rows has {} cols, expected {} (sum of component ncols)",
2048                    n_anchor_rows.ncols(),
2049                    d,
2050                ),
2051            });
2052        }
2053        let mut m_dense = Array2::<f64>::zeros((d, self.basis_dim));
2054        for (i, row) in m_rows.iter().enumerate() {
2055            if row.len() != self.basis_dim {
2056                return Err(FittedModelError::SchemaMismatch {
2057                    reason: format!(
2058                        "anchor_correction_matrix: M row {} has length {}, expected basis_dim {}",
2059                        i,
2060                        row.len(),
2061                        self.basis_dim,
2062                    ),
2063                });
2064            }
2065            for (j, &v) in row.iter().enumerate() {
2066                m_dense[[i, j]] = v;
2067            }
2068        }
2069        // The compiler bakes the orthonormalising rotation into M, so
2070        // the predict-time correction is simply `n_anchor_rows · M`.
2071        Ok(Some(n_anchor_rows.dot(&m_dense)))
2072    }
2073
2074    pub fn first_derivative_design(
2075        &self,
2076        values: &Array1<f64>,
2077    ) -> Result<Array2<f64>, FittedModelError> {
2078        self.evaluate_span_polynomial_design(
2079            values,
2080            BasisOptions::first_derivative().derivative_order,
2081        )
2082    }
2083
2084    pub fn second_derivative_design(
2085        &self,
2086        values: &Array1<f64>,
2087    ) -> Result<Array2<f64>, FittedModelError> {
2088        self.evaluate_span_polynomial_design(
2089            values,
2090            BasisOptions::second_derivative().derivative_order,
2091        )
2092    }
2093}
2094
2095impl FittedFamily {
2096    #[inline]
2097    pub fn likelihood(&self) -> LikelihoodSpec {
2098        let spec = match self {
2099            Self::Standard { likelihood, .. }
2100            | Self::LocationScale { likelihood, .. }
2101            | Self::MarginalSlope { likelihood, .. }
2102            | Self::Survival { likelihood, .. }
2103            | Self::TransformationNormal { likelihood, .. } => likelihood,
2104            Self::LatentSurvival { .. } | Self::LatentBinary { .. } => {
2105                return LikelihoodSpec::royston_parmar();
2106            }
2107        };
2108        spec.clone()
2109    }
2110
2111    #[inline]
2112    pub fn frailty(&self) -> Option<&FrailtySpec> {
2113        match self {
2114            Self::MarginalSlope { frailty, .. }
2115            | Self::Survival { frailty, .. }
2116            | Self::LatentSurvival { frailty }
2117            | Self::LatentBinary { frailty } => Some(frailty),
2118            _ => None,
2119        }
2120    }
2121}
2122
2123/// Recursively collect the feature columns of a smooth basis whose out-of-hull
2124/// evaluation is bounded, so they can be exempted from the predict-time axis
2125/// clip (see [`FittedModel::training_smooth_extrapolation_axes`]). Wrapper bases
2126/// (`by=`, factor-smooth, sum-to-zero) delegate to their inner smooth; `Sphere`
2127/// and `Pca` are intentionally not collected.
2128fn collect_smooth_extrapolation_axes(
2129    basis: &gam_terms::smooth::SmoothBasisSpec,
2130    n_training_headers: usize,
2131    out: &mut std::collections::HashSet<usize>,
2132) {
2133    use gam_terms::smooth::SmoothBasisSpec;
2134    let push = |col: usize, out: &mut std::collections::HashSet<usize>| {
2135        if col < n_training_headers {
2136            out.insert(col);
2137        }
2138    };
2139    match basis {
2140        // 1D B-spline: first-order linear extension off the boundary slope.
2141        SmoothBasisSpec::BSpline1D { feature_col, .. } => push(*feature_col, out),
2142        // Tensor B-spline: each margin linear-extends independently. Periodic
2143        // margins are additionally (and harmlessly) exempted via the periodic set.
2144        SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
2145            for &c in feature_cols {
2146                push(c, out);
2147            }
2148        }
2149        // Radial bases with a bounded out-of-hull contract: Duchon / thin-plate
2150        // are linear outside the data span (natural-spline boundary conditions),
2151        // Matérn reverts to its mean as the kernel decays. Measure-jet shares
2152        // the Matérn contract (Gaussian representers decay to the parametric
2153        // layer off the data support) — and off-web queries are exactly the
2154        // ones its support diagnostic must see unclipped.
2155        SmoothBasisSpec::ThinPlate { feature_cols, .. }
2156        | SmoothBasisSpec::Matern { feature_cols, .. }
2157        | SmoothBasisSpec::MeasureJet { feature_cols, .. }
2158        | SmoothBasisSpec::Duchon { feature_cols, .. } => {
2159            for &c in feature_cols {
2160                push(c, out);
2161            }
2162        }
2163        // Factor-smooth: the continuous marginal axes are B-splines that
2164        // linear-extend; the group column is categorical and is left to the
2165        // random-effect / level-lookup machinery.
2166        SmoothBasisSpec::FactorSmooth { spec } => {
2167            for &c in &spec.continuous_cols {
2168                push(c, out);
2169            }
2170        }
2171        // Wrappers delegate to the inner smooth they modulate / replicate.
2172        SmoothBasisSpec::ByVariable { inner, .. }
2173        | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2174            collect_smooth_extrapolation_axes(inner, n_training_headers, out)
2175        }
2176        SmoothBasisSpec::BySmooth { smooth, .. } => {
2177            collect_smooth_extrapolation_axes(smooth, n_training_headers, out)
2178        }
2179        // Sphere: latitude is clipped to manifold bounds, longitude is periodic —
2180        // both handled elsewhere with non-plateau semantics. Pca: no extrapolation
2181        // contract, stays clipped. ConstantCurvature: chart coordinates must stay
2182        // inside the κ-stereographic chart (open ball for κ < 0), so clipping new
2183        // data to the training range is the safe out-of-hull behavior.
2184        SmoothBasisSpec::Sphere { .. }
2185        | SmoothBasisSpec::ConstantCurvature { .. }
2186        | SmoothBasisSpec::Pca { .. } => {}
2187    }
2188}
2189
2190/// Collect training-column indices that feed a *numeric* `by=` multiplier of a
2191/// varying-coefficient smooth (`s(x, by=z)`).
2192///
2193/// A numeric by-variable enters the model as a linear multiplier of the centred
2194/// smooth basis, so the prediction is exactly affine in `z` for any fixed `x`:
2195/// `pred(x, z) = intercept + z·f(x)`. The by-multiplier is therefore an
2196/// *unbounded linear* extrapolant — exactly like a parametric `linear` axis
2197/// (already exempt from the clip via `training_linear_axes`), not a
2198/// bounded-basis axis. Clamping it to the training range `[min(z), max(z)]`
2199/// before the design is built would make the varying-coefficient effect plateau
2200/// outside the sampled range (slope 0 above the max) and silently replace the
2201/// natural `z == 0` baseline with the `z == min` prediction — the prediction is
2202/// no longer affine in `z`. Exempt it from the predict-time axis clip.
2203///
2204/// Factor `by=` columns are categorical group labels handled by the
2205/// level-lookup / random-effect machinery and are deliberately *not* exempted
2206/// here. Returned indices reference `self.training_headers`, matching the
2207/// iteration in `axis_clip_to_training_ranges`.
2208fn collect_by_variable_numeric_axes(
2209    basis: &gam_terms::smooth::SmoothBasisSpec,
2210    n_training_headers: usize,
2211    out: &mut std::collections::HashSet<usize>,
2212) {
2213    use gam_terms::smooth::{BySmoothKind, ByVarKind, SmoothBasisSpec};
2214    match basis {
2215        SmoothBasisSpec::ByVariable {
2216            inner,
2217            by_col,
2218            kind,
2219            ..
2220        } => {
2221            if matches!(kind, BySmoothKind::Numeric) && *by_col < n_training_headers {
2222                out.insert(*by_col);
2223            }
2224            collect_by_variable_numeric_axes(inner, n_training_headers, out);
2225        }
2226        SmoothBasisSpec::BySmooth { smooth, by_kind } => {
2227            if let ByVarKind::Numeric { feature_col } = by_kind
2228                && *feature_col < n_training_headers
2229            {
2230                out.insert(*feature_col);
2231            }
2232            collect_by_variable_numeric_axes(smooth, n_training_headers, out);
2233        }
2234        SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2235            collect_by_variable_numeric_axes(inner, n_training_headers, out);
2236        }
2237        _ => {}
2238    }
2239}
2240
2241impl FittedModel {
2242    /// Axis-clip each continuous new-data column to the (min, max) range
2243    /// observed in training. Categorical and binary columns are left
2244    /// untouched so unseen levels surface rather than being silently remapped
2245    /// onto seen ones. Returns `Some(clipped_copy)` only if at least one
2246    /// value was actually clipped; otherwise `None` so callers can avoid
2247    /// owning a redundant copy. Pre-2026-04-29 model JSONs that lack the
2248    /// `training_feature_ranges` field deserialize to `None` and pass through
2249    /// unchanged.
2250    pub fn axis_clip_to_training_ranges(
2251        &self,
2252        data: ndarray::ArrayView2<'_, f64>,
2253        col_map: &std::collections::HashMap<String, usize>,
2254    ) -> Option<ndarray::Array2<f64>> {
2255        let training_headers = self.training_headers.as_ref()?;
2256        let ranges = self.training_feature_ranges.as_ref()?;
2257        if training_headers.len() != ranges.len() {
2258            return None;
2259        }
2260        let mut kind_by_header: std::collections::HashMap<&str, ColumnKindTag> =
2261            std::collections::HashMap::new();
2262        if let Some(schema) = self.data_schema.as_ref() {
2263            for col in &schema.columns {
2264                kind_by_header.insert(col.name.as_str(), col.kind);
2265            }
2266        }
2267        // Periodic axes (sphere longitude, periodic-B-spline 1D, periodic
2268        // tensor margins) must never be clipped to the training range:
2269        // clamping a value just past the seam to the training extreme breaks
2270        // the cyclic invariant f(x₀) = f(x₀ + period) at predict time and
2271        // shows up as a visible seam in surface plots.
2272        let periodic_axes = self.training_periodic_axes(training_headers);
2273        // Parametric/linear-term axes must never be clipped either: a linear
2274        // term's contract is η = β0 + β1·x, i.e. genuine linear extrapolation.
2275        // Clamping its input to the training extreme turns predict into a
2276        // piecewise-constant plateau outside the training hull and freezes the
2277        // prediction SE at the boundary (the clamped x feeds xᵀ Var(β) x), so
2278        // credible intervals stop widening with distance from the data. This
2279        // mirrors how periodic axes are exempted just above.
2280        let linear_axes = self.training_linear_axes(training_headers.len());
2281        // Random-effect grouping axes are categorical even when their source
2282        // column is numeric. Clipping them would remap an unseen group label to
2283        // a boundary training level instead of letting the random-effect block
2284        // encode it as the prior-mean zero effect.
2285        let random_effect_axes = self.training_random_effect_axes(training_headers.len());
2286        // Non-parametric smooth axes whose basis extrapolates boundedly on its
2287        // own (B-spline linear extension, Duchon/thin-plate natural-spline linear
2288        // tail, Matérn kernel decay). Clamping their input to the training extreme
2289        // hands the basis an already-clamped coordinate, so its extrapolation
2290        // never fires and predict freezes at a boundary plateau — diverging from
2291        // the raw design path, which does not clip. Exempt them so both paths go
2292        // through the single basis-layer extrapolation. See the method doc.
2293        let smooth_extrapolation_axes =
2294            self.training_smooth_extrapolation_axes(training_headers.len());
2295        // Numeric `by=` multipliers of a varying-coefficient smooth `s(x, by=z)`
2296        // are linear multipliers of the centred basis (prediction affine in z),
2297        // so — like parametric linear axes — clipping them to the training range
2298        // turns the varying-coefficient effect into a boundary plateau and
2299        // destroys the z==0 baseline. Exempt them from the clip.
2300        let by_variable_axes = self.training_by_variable_numeric_axes(training_headers.len());
2301        // Sphere latitude is a closed-manifold coordinate: its clip bounds are
2302        // the manifold's intrinsic domain ([-π/2, π/2] or [-90, 90]), not the
2303        // sampled range, so a pole prediction reaches the true pole instead of
2304        // being clamped to a near-pole latitude (see the method doc).
2305        let sphere_lat_bounds = self.training_sphere_latitude_bounds(training_headers);
2306        let mut clipped = data.to_owned();
2307        let mut any_clipped = false;
2308        for (col_in_training, (header, &(lo, hi))) in
2309            training_headers.iter().zip(ranges.iter()).enumerate()
2310        {
2311            let (lo, hi) = sphere_lat_bounds
2312                .get(&col_in_training)
2313                .copied()
2314                .unwrap_or((lo, hi));
2315            if !(lo.is_finite() && hi.is_finite()) || hi <= lo {
2316                continue;
2317            }
2318            if !matches!(
2319                kind_by_header.get(header.as_str()).copied(),
2320                Some(ColumnKindTag::Continuous)
2321            ) {
2322                continue;
2323            }
2324            if periodic_axes.contains(&col_in_training) {
2325                continue;
2326            }
2327            if linear_axes.contains(&col_in_training) {
2328                continue;
2329            }
2330            if random_effect_axes.contains(&col_in_training) {
2331                continue;
2332            }
2333            if smooth_extrapolation_axes.contains(&col_in_training) {
2334                continue;
2335            }
2336            if by_variable_axes.contains(&col_in_training) {
2337                continue;
2338            }
2339            let Some(&col_idx) = col_map.get(header) else {
2340                continue;
2341            };
2342            if col_idx >= clipped.ncols() {
2343                continue;
2344            }
2345            let mut col = clipped.column_mut(col_idx);
2346            for v in col.iter_mut() {
2347                if v.is_finite() {
2348                    if *v < lo {
2349                        *v = lo;
2350                        any_clipped = true;
2351                    } else if *v > hi {
2352                        *v = hi;
2353                        any_clipped = true;
2354                    }
2355                }
2356            }
2357        }
2358        if any_clipped { Some(clipped) } else { None }
2359    }
2360
2361    fn saved_term_specs(&self) -> Vec<&TermCollectionSpec> {
2362        let mut specs: Vec<&TermCollectionSpec> = [
2363            self.resolved_termspec.as_ref(),
2364            self.resolved_termspec_noise.as_ref(),
2365            self.resolved_termspec_logslope.as_ref(),
2366        ]
2367        .into_iter()
2368        .flatten()
2369        .collect();
2370        if let Some(logslopes) = self.resolved_termspec_logslopes.as_ref() {
2371            specs.extend(logslopes.iter());
2372        }
2373        specs
2374    }
2375
2376    /// Collect the set of training-column indices that are periodic axes —
2377    /// i.e. features for which a periodic basis (sphere longitude, periodic
2378    /// B-spline 1D, periodic tensor margin) must be allowed to take any
2379    /// real value at predict time and not be clamped to the training range.
2380    /// Returned indices reference `self.training_headers` (training-time
2381    /// layout), matching the iteration in `axis_clip_to_training_ranges`.
2382    fn training_periodic_axes(
2383        &self,
2384        training_headers: &[String],
2385    ) -> std::collections::HashSet<usize> {
2386        use gam_terms::basis::BSplineKnotSpec;
2387        use gam_terms::smooth::SmoothBasisSpec;
2388        let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2389        let Some(spec) = self.resolved_termspec.as_ref() else {
2390            return out;
2391        };
2392        for term in &spec.smooth_terms {
2393            match &term.basis {
2394                // Sphere terms: longitude (second feature col) is always
2395                // periodic and exempt from clipping. Latitude is not periodic
2396                // but is a closed-manifold coordinate, so it is clipped to the
2397                // manifold's intrinsic bounds rather than the sampled range —
2398                // see `training_sphere_latitude_bounds`.
2399                SmoothBasisSpec::Sphere { feature_cols, .. } => {
2400                    if let Some(&lon_col) = feature_cols.get(1)
2401                        && lon_col < training_headers.len()
2402                    {
2403                        out.insert(lon_col);
2404                    }
2405                }
2406                // 1D periodic B-spline: the single feature column is periodic.
2407                SmoothBasisSpec::BSpline1D { feature_col, spec } => {
2408                    if matches!(spec.knotspec, BSplineKnotSpec::PeriodicUniform { .. })
2409                        && *feature_col < training_headers.len()
2410                    {
2411                        out.insert(*feature_col);
2412                    }
2413                }
2414                // Tensor B-spline: each axis whose marginal knotspec is
2415                // PeriodicUniform is periodic; mark those columns.
2416                SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
2417                    for (i, marginal) in spec.marginalspecs.iter().enumerate() {
2418                        if matches!(marginal.knotspec, BSplineKnotSpec::PeriodicUniform { .. })
2419                            && let Some(&col) = feature_cols.get(i)
2420                            && col < training_headers.len()
2421                        {
2422                            out.insert(col);
2423                        }
2424                    }
2425                }
2426                _ => {}
2427            }
2428        }
2429        out
2430    }
2431
2432    /// Collect the set of training-column indices that feed a parametric/linear
2433    /// term — on *any* modelled surface (mean, noise/scale, log-slope). A
2434    /// linear term realises the design column `∏ feature_cols` and contributes
2435    /// `β·(that product)` to the linear predictor, so its inputs must be allowed
2436    /// to take any real value at predict time: clamping them to the training
2437    /// range would replace genuine linear extrapolation with a boundary plateau
2438    /// (and freeze the prediction SE at the hull edge). Returned indices
2439    /// reference `self.training_headers` (training-time layout), matching the
2440    /// iteration in `axis_clip_to_training_ranges`. Wilkinson-Rogers `:`
2441    /// interactions contribute every column in their `feature_cols` product.
2442    fn training_linear_axes(&self, n_training_headers: usize) -> std::collections::HashSet<usize> {
2443        let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2444        for spec in self.saved_term_specs() {
2445            for term in &spec.linear_terms {
2446                for col in term.effective_feature_cols() {
2447                    if col < n_training_headers {
2448                        out.insert(col);
2449                    }
2450                }
2451            }
2452        }
2453        out
2454    }
2455
2456    /// Collect the set of training-column indices that feed random-effect
2457    /// grouping terms. These columns are categorical model axes regardless of
2458    /// the ingest schema's scalar storage type, so prediction must leave them
2459    /// untouched and let the frozen random-effect levels decide whether a row is
2460    /// a seen level or the zero-effect unseen-level fallback.
2461    fn training_random_effect_axes(
2462        &self,
2463        n_training_headers: usize,
2464    ) -> std::collections::HashSet<usize> {
2465        let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2466        for spec in self.saved_term_specs() {
2467            for term in &spec.random_effect_terms {
2468                if term.feature_col < n_training_headers {
2469                    out.insert(term.feature_col);
2470                }
2471            }
2472        }
2473        out
2474    }
2475
2476    /// Collect the set of training-column indices that feed a non-parametric
2477    /// smooth whose basis performs its own *bounded* extrapolation outside the
2478    /// training hull — on any modelled surface (mean, noise/scale, log-slope).
2479    ///
2480    /// These columns must be exempt from the predict-time axis clip for the same
2481    /// reason periodic/linear/random-effect axes are: the clip clamps a new value
2482    /// to the training extreme *before* the design is built, so the basis is
2483    /// handed an already-clamped coordinate and its extrapolation machinery never
2484    /// fires. The result is a piecewise-constant plateau frozen at the boundary
2485    /// fitted value (with a prediction SE frozen at the hull edge), and — worse —
2486    /// a model that yields *different* predictions through the `FittedModel`
2487    /// predict pipeline than through the raw `build_term_collection_design` path,
2488    /// which does not clip. Exempting these axes routes both entry points through
2489    /// the single basis-layer extrapolation, restoring internal consistency.
2490    ///
2491    /// Only bases with a *bounded* out-of-hull contract are listed, so removing
2492    /// the clip cannot reintroduce the wild basis blow-up the clip guards against:
2493    ///   - B-spline 1D / tensor margins: first-order linear extension off the
2494    ///     boundary slope (`apply_linear_extension_from_first_derivative`) — grows
2495    ///     at most linearly.
2496    ///   - Duchon / thin-plate: the natural-spline boundary conditions make the
2497    ///     fit linear outside the data span — also at most linear growth.
2498    ///   - Matérn: the kernel decays with distance, so the fit reverts smoothly to
2499    ///     its (low-order polynomial / constant) mean far from the data — bounded.
2500    /// `sphere()` axes are deliberately *not* listed: latitude is a closed-manifold
2501    /// coordinate clipped to its intrinsic bounds (`training_sphere_latitude_bounds`)
2502    /// and longitude is periodic (`training_periodic_axes`); both already have the
2503    /// correct, non-plateau handling. `Pca` projections have no extrapolation
2504    /// contract and stay clipped.
2505    ///
2506    /// Returned indices reference `self.training_headers` (training-time layout),
2507    /// matching the iteration in `axis_clip_to_training_ranges`.
2508    fn training_smooth_extrapolation_axes(
2509        &self,
2510        n_training_headers: usize,
2511    ) -> std::collections::HashSet<usize> {
2512        let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2513        for spec in self.saved_term_specs() {
2514            for term in &spec.smooth_terms {
2515                collect_smooth_extrapolation_axes(&term.basis, n_training_headers, &mut out);
2516            }
2517        }
2518        out
2519    }
2520
2521    /// Collect the set of training-column indices that feed a *numeric* `by=`
2522    /// multiplier of a varying-coefficient smooth. These columns are linear
2523    /// multipliers (`pred = intercept + z·f(x)`), so they must be exempt from
2524    /// the predict-time axis clip for the same reason parametric linear axes
2525    /// are — see `collect_by_variable_numeric_axes` for the full rationale.
2526    fn training_by_variable_numeric_axes(
2527        &self,
2528        n_training_headers: usize,
2529    ) -> std::collections::HashSet<usize> {
2530        let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2531        for spec in self.saved_term_specs() {
2532            for term in &spec.smooth_terms {
2533                collect_by_variable_numeric_axes(&term.basis, n_training_headers, &mut out);
2534            }
2535        }
2536        out
2537    }
2538
2539    /// Manifold-intrinsic clip bounds for sphere *latitude* columns.
2540    ///
2541    /// A `sphere(lat, lon)` smooth charts the closed manifold S²: the poles
2542    /// (lat = ±π/2, or ±90°) are interior limit points of that manifold, not
2543    /// endpoints of an unbounded axis to extrapolate along. A finite training
2544    /// sample never reaches the pole exactly, so clamping a pole prediction to
2545    /// the *observed* latitude extreme lands at a near-pole latitude where the
2546    /// Wahba SOS kernel's `cos(lat)·cos(lat_c)·cos(Δlon)` term has not yet
2547    /// damped — and the predictor then sweeps a spurious `cos(lon)` profile at
2548    /// what is physically a single point, reintroducing the pole artefact the
2549    /// SOS basis exists to remove.
2550    ///
2551    /// The correct clip bound for this coordinate is therefore the manifold's
2552    /// intrinsic domain — `[-π/2, π/2]` radians or `[-90, 90]` degrees — not
2553    /// the sampled range. Clamping to those bounds keeps the pole reachable
2554    /// (single-valued in longitude) while still mapping any out-of-domain
2555    /// latitude onto the manifold boundary. Longitude needs no entry here: it
2556    /// is periodic and already exempted from clipping entirely
2557    /// (`training_periodic_axes`). Returned indices reference
2558    /// `self.training_headers`, matching the iteration in
2559    /// `axis_clip_to_training_ranges`.
2560    fn training_sphere_latitude_bounds(
2561        &self,
2562        training_headers: &[String],
2563    ) -> std::collections::HashMap<usize, (f64, f64)> {
2564        use gam_terms::smooth::SmoothBasisSpec;
2565        let mut out: std::collections::HashMap<usize, (f64, f64)> =
2566            std::collections::HashMap::new();
2567        let Some(spec) = self.resolved_termspec.as_ref() else {
2568            return out;
2569        };
2570        for term in &spec.smooth_terms {
2571            if let SmoothBasisSpec::Sphere { feature_cols, spec } = &term.basis
2572                && let Some(&lat_col) = feature_cols.first()
2573                && lat_col < training_headers.len()
2574            {
2575                let bound = if spec.radians {
2576                    std::f64::consts::FRAC_PI_2
2577                } else {
2578                    90.0
2579                };
2580                out.insert(lat_col, (-bound, bound));
2581            }
2582        }
2583        out
2584    }
2585
2586    pub fn from_payload(mut payload: FittedModelPayload) -> Self {
2587        let likelihood = payload.family_state.likelihood();
2588        let class = match payload.model_kind {
2589            ModelKind::Survival => PredictModelClass::Survival,
2590            ModelKind::MarginalSlope => PredictModelClass::BernoulliMarginalSlope,
2591            ModelKind::TransformationNormal => PredictModelClass::TransformationNormal,
2592            ModelKind::LocationScale => {
2593                if likelihood == LikelihoodSpec::gaussian_identity() {
2594                    PredictModelClass::GaussianLocationScale
2595                } else if is_dispersion_location_scale_response(&likelihood.response) {
2596                    PredictModelClass::DispersionLocationScale
2597                } else {
2598                    PredictModelClass::BinomialLocationScale
2599                }
2600            }
2601            ModelKind::Standard => PredictModelClass::Standard,
2602        };
2603        match class {
2604            PredictModelClass::Survival => {
2605                payload.model_kind = ModelKind::Survival;
2606                Self::Survival { payload }
2607            }
2608            PredictModelClass::BernoulliMarginalSlope => {
2609                payload.model_kind = ModelKind::MarginalSlope;
2610                Self::MarginalSlope { payload }
2611            }
2612            PredictModelClass::TransformationNormal => {
2613                payload.model_kind = ModelKind::TransformationNormal;
2614                Self::TransformationNormal { payload }
2615            }
2616            PredictModelClass::GaussianLocationScale
2617            | PredictModelClass::BinomialLocationScale
2618            | PredictModelClass::DispersionLocationScale => {
2619                payload.model_kind = ModelKind::LocationScale;
2620                Self::LocationScale { payload }
2621            }
2622            PredictModelClass::Standard => {
2623                payload.model_kind = ModelKind::Standard;
2624                Self::Standard { payload }
2625            }
2626        }
2627        .with_synchronized_stateful_link_metadata()
2628    }
2629
2630    #[inline]
2631    pub fn payload(&self) -> &FittedModelPayload {
2632        match self {
2633            Self::Standard { payload }
2634            | Self::LocationScale { payload }
2635            | Self::MarginalSlope { payload }
2636            | Self::Survival { payload }
2637            | Self::TransformationNormal { payload } => payload,
2638        }
2639    }
2640
2641    #[inline]
2642    fn payload_mut(&mut self) -> &mut FittedModelPayload {
2643        match self {
2644            Self::Standard { payload }
2645            | Self::LocationScale { payload }
2646            | Self::MarginalSlope { payload }
2647            | Self::Survival { payload }
2648            | Self::TransformationNormal { payload } => payload,
2649        }
2650    }
2651
2652    fn with_synchronized_stateful_link_metadata(mut self) -> Self {
2653        self.synchronize_stateful_link_metadata();
2654        self
2655    }
2656
2657    fn synchronize_stateful_link_metadata(&mut self) {
2658        let payload = self.payload_mut();
2659        // `fit_result` and `unified` are two names for the SAME canonical
2660        // UnifiedFitResult — every production builder (run_fit, the
2661        // model_payload_builders) sets both to the identical fit. Consumers and
2662        // the mutual-exclusivity persistence checks read `fit_result.or(unified)`,
2663        // but the dense-path persistence gate and the marginal-slope serialization
2664        // key on `fit_result` alone (`self.fit_result.as_ref().expect("checked
2665        // above")`). A payload constructed directly from just a `unified` fit
2666        // (nothing wrong with that — it is the same value) would then fail
2667        // `validate_for_persistence` with "missing canonical fit_result payload".
2668        // Mirror the two so whichever the caller populated, the canonical
2669        // `fit_result` slot (and its `unified` alias) is always present.
2670        match (payload.fit_result.is_none(), payload.unified.is_none()) {
2671            (true, false) => payload.fit_result = payload.unified.clone(),
2672            (false, true) => payload.unified = payload.fit_result.clone(),
2673            _ => {}
2674        }
2675        payload.used_device = payload
2676            .fit_result
2677            .as_ref()
2678            .or(payload.unified.as_ref())
2679            .is_some_and(|fit| fit.used_device);
2680        payload.synchronize_empty_feature_contract();
2681        let Some(fit) = payload.fit_result.as_ref().or(payload.unified.as_ref()) else {
2682            return;
2683        };
2684        match (&mut payload.family_state, &fit.fitted_link) {
2685            (
2686                FittedFamily::Standard {
2687                    likelihood,
2688                    latent_cloglog_state,
2689                    ..
2690                },
2691                FittedLinkState::LatentCLogLog { state },
2692            ) if likelihood.is_latent_cloglog() => {
2693                *latent_cloglog_state = Some(*state);
2694            }
2695            (
2696                FittedFamily::Standard {
2697                    likelihood,
2698                    sas_state,
2699                    ..
2700                },
2701                FittedLinkState::Sas { state, covariance },
2702            ) if likelihood.is_binomial_sas() => {
2703                *sas_state = Some(*state);
2704                payload.sas_param_covariance = covariance.as_ref().map(array2_to_nested_vec);
2705            }
2706            (
2707                FittedFamily::Standard {
2708                    likelihood,
2709                    sas_state,
2710                    ..
2711                },
2712                FittedLinkState::BetaLogistic { state, covariance },
2713            ) if likelihood.is_binomial_beta_logistic() => {
2714                *sas_state = Some(*state);
2715                payload.sas_param_covariance = covariance.as_ref().map(array2_to_nested_vec);
2716            }
2717            (
2718                FittedFamily::Standard {
2719                    likelihood,
2720                    mixture_state,
2721                    ..
2722                },
2723                FittedLinkState::Mixture { state, covariance },
2724            ) if likelihood.is_binomial_mixture() => {
2725                *mixture_state = Some(state.clone());
2726                payload.mixture_link_param_covariance =
2727                    covariance.as_ref().map(array2_to_nested_vec);
2728            }
2729            _ => {}
2730        }
2731    }
2732
2733    #[inline]
2734    pub fn likelihood(&self) -> LikelihoodSpec {
2735        self.payload().family_state.likelihood()
2736    }
2737
2738    /// Columns this model consumes from a prediction frame — its *input
2739    /// contract*.
2740    ///
2741    /// Every variable named by the main formula (features, interaction margins,
2742    /// random-effect groups, and a smooth's `by=` column), the survival
2743    /// entry/exit columns or the transformation-normal response, the auxiliary
2744    /// noise / logslope formula columns, and the offset / noise-offset /
2745    /// latent-`z` columns. The event-indicator and the plain response of a
2746    /// standard model are deliberately excluded: they are not needed to *form*
2747    /// a prediction (the conformal-calibration fold layers the response back on
2748    /// separately).
2749    ///
2750    /// This is the single authority shared by the CLI and PyFFI predict paths.
2751    /// A prediction frame column that is *not* in this set is irrelevant to the
2752    /// model and must be ignored rather than strict-encoded against the
2753    /// training schema — otherwise an unrelated ID/label column with a held-out
2754    /// categorical level aborts predict (#840).
2755    pub fn prediction_required_columns(
2756        &self,
2757    ) -> Result<std::collections::BTreeSet<String>, String> {
2758        let payload = self.payload();
2759        let parsed = parse_formula(payload.formula.as_str()).map_err(|e| e.to_string())?;
2760        let mut required = std::collections::BTreeSet::<String>::new();
2761        parsed_term_column_names(&parsed.terms, &mut required);
2762
2763        if let Some((entry, exit, _event)) =
2764            parse_surv_response(parsed.response.as_str()).map_err(|e| e.to_string())?
2765        {
2766            if let Some(entry) = entry {
2767                required.insert(entry);
2768            }
2769            required.insert(exit);
2770        } else if let Some((left, right, _event)) =
2771            parse_surv_interval_response(parsed.response.as_str()).map_err(|e| e.to_string())?
2772        {
2773            required.insert(left);
2774            required.insert(right);
2775        }
2776        // A transformation-normal (CTM) prediction returns the response-scale
2777        // conditional mean E[Y|x], a function of the covariates alone (issue
2778        // #1612). The earlier implementation precomputed the PIT h(y|x) of the
2779        // supplied response, which made the outcome column mandatory at predict
2780        // time; the response is no longer required, so a covariate-only frame
2781        // must predict without it.
2782
2783        if let Some(offset) = payload.offset_column.as_ref() {
2784            required.insert(offset.clone());
2785        }
2786        if let Some(noise_offset) = payload.noise_offset_column.as_ref() {
2787            required.insert(noise_offset.clone());
2788        }
2789        if matches!(
2790            self.predict_model_class(),
2791            PredictModelClass::BernoulliMarginalSlope | PredictModelClass::Survival
2792        ) {
2793            if let Some(z_column) = payload.z_column.as_ref() {
2794                required.remove("z");
2795                required.insert(z_column.clone());
2796            }
2797        }
2798        if let Some(noise_formula) = payload.formula_noise.as_ref() {
2799            self.add_auxiliary_formula_columns(
2800                &mut required,
2801                noise_formula,
2802                parsed.response.as_str(),
2803            )?;
2804        }
2805        if let Some(logslope_formula) = payload.formula_logslope.as_ref() {
2806            if logslope_formula != "same-as-main" {
2807                self.add_auxiliary_formula_columns(
2808                    &mut required,
2809                    logslope_formula,
2810                    parsed.response.as_str(),
2811                )?;
2812            }
2813        }
2814        Ok(required)
2815    }
2816
2817    /// Columns a *post-fit diagnostic* command (diagnose / sample / report)
2818    /// needs **beyond** [`Self::prediction_required_columns`].
2819    ///
2820    /// Prediction deliberately drops a standard GAM's bare response so a
2821    /// prediction frame may omit it (#840 / #864). Diagnostics are statements
2822    /// *about* that observed response — residuals, R², posterior likelihoods,
2823    /// leave-one-out — so the response must be present. This returns the bare
2824    /// response column when the prediction projection would otherwise drop it,
2825    /// and nothing when the response is already prediction-required (survival
2826    /// `Surv(...)` time/event columns, the transformation-normal response) or
2827    /// is not a plain data column.
2828    ///
2829    /// Centralising the intent here is what makes it *structurally impossible*
2830    /// for a diagnostic command to silently drop the response: callers use
2831    /// `load_dataset…_for_diagnostics`, which always folds these in, instead of
2832    /// each remembering to thread an `extra_required` response by hand.
2833    pub fn diagnostic_extra_columns(&self) -> Result<Vec<String>, String> {
2834        let payload = self.payload();
2835        let parsed = parse_formula(payload.formula.as_str()).map_err(|e| e.to_string())?;
2836        // Prior (case) weights never enter the linear predictor, so
2837        // `prediction_required_columns` deliberately omits the weight column and
2838        // a prediction frame may drop it. Diagnostics are weight-aware, though:
2839        // `diagnose` reconstructs the ALO working weights `w_i = prior_i ·
2840        // Fisher_i` (and the refit fallback re-weights the same way), so the
2841        // weight column must be loaded. Fold it in here — the single seam that
2842        // makes it structurally impossible for a diagnostic command to silently
2843        // drop a needed column — regardless of the response-shape early-outs
2844        // below, since it is orthogonal to the response.
2845        let mut extras: Vec<String> = Vec::new();
2846        if let Some(weight_column) = payload.weight_column.as_ref() {
2847            extras.push(weight_column.clone());
2848        }
2849        // Survival responses are `Surv(...)` expressions, not bare columns; the
2850        // underlying entry/exit columns are already prediction-required.
2851        if parse_surv_response(parsed.response.as_str())
2852            .map_err(|e| e.to_string())?
2853            .is_some()
2854            || parse_surv_interval_response(parsed.response.as_str())
2855                .map_err(|e| e.to_string())?
2856                .is_some()
2857        {
2858            return Ok(extras);
2859        }
2860        let response = parsed.response.trim();
2861        // A response that is empty, or a function-call expression rather than a
2862        // plain data column, has no bare column to re-add.
2863        if response.is_empty() || response.contains('(') {
2864            return Ok(extras);
2865        }
2866        // Already prediction-required (e.g. transformation-normal re-adds it):
2867        // nothing extra to fold in.
2868        if self.prediction_required_columns()?.contains(response) {
2869            return Ok(extras);
2870        }
2871        extras.push(response.to_string());
2872        Ok(extras)
2873    }
2874
2875    /// Add the columns referenced by an auxiliary (noise / logslope) formula,
2876    /// which may be supplied as a full `lhs ~ rhs` formula or as a bare RHS.
2877    fn add_auxiliary_formula_columns(
2878        &self,
2879        required: &mut std::collections::BTreeSet<String>,
2880        formula_or_rhs: &str,
2881        response: &str,
2882    ) -> Result<(), String> {
2883        let trimmed = formula_or_rhs.trim();
2884        if trimmed.is_empty() || trimmed == "1" {
2885            return Ok(());
2886        }
2887        let formula = if trimmed.contains('~') {
2888            trimmed.to_string()
2889        } else {
2890            format!("{response} ~ {trimmed}")
2891        };
2892        let parsed = parse_formula(formula.as_str()).map_err(|e| e.to_string())?;
2893        parsed_term_column_names(&parsed.terms, required);
2894        Ok(())
2895    }
2896
2897    #[inline]
2898    pub fn predict_model_class(&self) -> PredictModelClass {
2899        match &self.payload().family_state {
2900            FittedFamily::Survival { .. }
2901            | FittedFamily::LatentSurvival { .. }
2902            | FittedFamily::LatentBinary { .. } => PredictModelClass::Survival,
2903            FittedFamily::MarginalSlope { .. } => PredictModelClass::BernoulliMarginalSlope,
2904            FittedFamily::TransformationNormal { .. } => PredictModelClass::TransformationNormal,
2905            FittedFamily::LocationScale { likelihood, .. } if likelihood.is_gaussian_identity() => {
2906                PredictModelClass::GaussianLocationScale
2907            }
2908            FittedFamily::LocationScale { likelihood, .. }
2909                if is_dispersion_location_scale_response(&likelihood.response) =>
2910            {
2911                PredictModelClass::DispersionLocationScale
2912            }
2913            FittedFamily::LocationScale { .. } => PredictModelClass::BinomialLocationScale,
2914            FittedFamily::Standard { .. } => PredictModelClass::Standard,
2915        }
2916    }
2917
2918    pub fn saved_link_wiggle(&self) -> Result<Option<SavedLinkWiggleRuntime>, FittedModelError> {
2919        let payload = self.payload();
2920        let (knots, degree) = match (
2921            payload.linkwiggle_knots.as_ref(),
2922            payload.linkwiggle_degree,
2923        ) {
2924            (None, None) => return Ok(None),
2925            (Some(knots), Some(degree)) => (knots.clone(), degree),
2926            _ => {
2927                return Err(FittedModelError::SchemaMismatch {
2928                    reason:
2929                        "saved model has partial link-wiggle metadata; expected linkwiggle_knots and linkwiggle_degree together"
2930                            .to_string(),
2931                })
2932            }
2933        };
2934        let resolved_link = self.resolved_inverse_link()?;
2935        let saved_link_disallows_wiggle = resolved_link
2936            .as_ref()
2937            .is_some_and(|link| !inverse_link_supports_joint_wiggle(link))
2938            || payload
2939                .link
2940                .as_ref()
2941                .is_some_and(|link| !inverse_link_supports_joint_wiggle(link));
2942        if saved_link_disallows_wiggle {
2943            return Err(FittedModelError::IncompatibleConfig {
2944                reason: joint_wiggle_unsupported_link_message("link wiggle"),
2945            });
2946        }
2947        let beta = match self.predict_model_class() {
2948            // #1596: the frozen-basis de-aliased standard link-warp is fit in a
2949            // reduced, identifiable coordinate `γ` and the fit_result LinkWiggle
2950            // block stores `γ` (its true free parameters). The full-width
2951            // standard-basis lift `β_w = Z·γ`, the coefficients the predict-time
2952            // I-spline basis multiplies, is persisted in `payload.beta_link_wiggle`
2953            // — prefer it when present. Without it (the dynamic-basis path) the
2954            // block coefficients ARE the standard-basis warp, read directly.
2955            PredictModelClass::Standard if payload.beta_link_wiggle.is_some() => {
2956                payload.beta_link_wiggle.clone().expect("checked is_some")
2957            }
2958            PredictModelClass::Standard => {
2959                let fit = payload.fit_result.as_ref().ok_or_else(|| {
2960                    FittedModelError::MissingField {
2961                        reason:
2962                            "standard link-wiggle model is missing canonical fit_result payload"
2963                                .to_string(),
2964                    }
2965                })?;
2966                if fit.blocks.len() != 2
2967                    || fit.blocks[0].role != BlockRole::Mean
2968                    || fit.blocks[1].role != BlockRole::LinkWiggle
2969                {
2970                    return Err(FittedModelError::SchemaMismatch {
2971                        reason:
2972                            "standard link-wiggle models must store blocks in [Mean, LinkWiggle] order"
2973                                .to_string(),
2974                    });
2975                }
2976                fit.block_by_role(BlockRole::LinkWiggle)
2977                    .ok_or_else(|| FittedModelError::MissingField {
2978                        reason:
2979                            "standard link-wiggle model is missing LinkWiggle coefficient block"
2980                                .to_string(),
2981                    })?
2982                    .beta
2983                    .to_vec()
2984            }
2985            _ => payload
2986                .beta_link_wiggle
2987                .clone()
2988                .ok_or_else(|| FittedModelError::MissingField {
2989                    reason:
2990                        "saved model has link-wiggle metadata but is missing payload.beta_link_wiggle"
2991                            .to_string(),
2992                })?,
2993        };
2994        // #2141: the frozen-index shift lets predict evaluate the warp basis at
2995        // the index `η̂` the fit pinned `B` at (`base + X·s`). `None` for
2996        // pre-#2141 saves and for non-de-aliased warp paths, where the base
2997        // predictor already is the warp index.
2998        let index_shift = payload.link_wiggle_index_shift.clone();
2999        Ok(Some(SavedLinkWiggleRuntime {
3000            knots,
3001            degree,
3002            beta,
3003            index_shift,
3004        }))
3005    }
3006
3007    pub fn saved_baseline_time_wiggle(
3008        &self,
3009    ) -> Result<Option<SavedBaselineTimeWiggleRuntime>, FittedModelError> {
3010        let payload = self.payload();
3011        if payload
3012            .survival_cause_count
3013            .is_some_and(|cause_count| cause_count > 1)
3014            && payload.beta_baseline_timewiggle.is_none()
3015            && payload.beta_baseline_timewiggle_by_cause.is_some()
3016        {
3017            return Err(FittedModelError::SchemaMismatch {
3018                reason:
3019                    "joint cause-specific survival stores baseline-timewiggle coefficients per cause"
3020                        .to_string(),
3021            });
3022        }
3023        match (
3024            payload.baseline_timewiggle_knots.as_ref(),
3025            payload.baseline_timewiggle_degree,
3026            payload.baseline_timewiggle_penalty_orders.as_ref(),
3027            payload.baseline_timewiggle_double_penalty,
3028            payload.beta_baseline_timewiggle.as_ref(),
3029        ) {
3030            (None, None, None, None, None) => Ok(None),
3031            (Some(knots), Some(degree), Some(penalty_orders), Some(double_penalty), Some(beta)) => {
3032                Ok(Some(SavedBaselineTimeWiggleRuntime {
3033                    knots: knots.clone(),
3034                    degree,
3035                    penalty_orders: penalty_orders.clone(),
3036                    double_penalty,
3037                    beta: beta.clone(),
3038                }))
3039            }
3040            _ => Err(FittedModelError::SchemaMismatch {
3041                reason:
3042                    "saved model has partial baseline-timewiggle metadata; expected knots+degree+penalty_order+double_penalty+beta_baseline_timewiggle together"
3043                        .to_string(),
3044            }),
3045        }
3046    }
3047
3048    /// Whether this model has a link wiggle component with complete metadata.
3049    #[inline]
3050    pub fn has_link_wiggle(&self) -> bool {
3051        self.saved_link_wiggle()
3052            .map(|runtime| runtime.is_some())
3053            .unwrap_or(false)
3054    }
3055
3056    /// Whether this model has a baseline-time wiggle component with complete metadata.
3057    #[inline]
3058    pub fn has_baseline_time_wiggle(&self) -> bool {
3059        let payload = self.payload();
3060        if payload
3061            .survival_cause_count
3062            .is_some_and(|cause_count| cause_count > 1)
3063        {
3064            return payload.baseline_timewiggle_knots.is_some()
3065                && payload.baseline_timewiggle_degree.is_some()
3066                && payload.baseline_timewiggle_penalty_orders.is_some()
3067                && payload.baseline_timewiggle_double_penalty.is_some()
3068                && payload.beta_baseline_timewiggle_by_cause.is_some();
3069        }
3070        self.saved_baseline_time_wiggle()
3071            .map(|runtime| runtime.is_some())
3072            .unwrap_or(false)
3073    }
3074
3075    /// Whether the default point prediction must integrate the inverse link
3076    /// over the coefficient posterior — reporting the posterior mean
3077    /// `E[g⁻¹(Xβ)]` — rather than plugging in the posterior mode `g⁻¹(Xβ̂)`.
3078    ///
3079    /// SPEC (issue #960): the posterior mean is *always* the default point
3080    /// estimate (never MAP). It is observably distinct from the plug-in exactly
3081    /// when the inverse link is *curved* over the posterior's uncertainty, so
3082    /// `E[g⁻¹(η)] ≠ g⁻¹(E[η])` by Jensen. The curvature-based classification is:
3083    ///   * all log-link families (Poisson / Gamma / Tweedie / NegativeBinomial):
3084    ///     `E[exp η] = exp(η + se²/2) ≠ exp(η)` (log-normal MGF);
3085    ///   * all Binomial links (logit / probit / cloglog / SAS / BetaLogistic /
3086    ///     Mixture / LatentCLogLog): bounded sigmoidal inverse links;
3087    ///   * Beta (logit link): `E[σ(η)] ≠ σ(E[η])`;
3088    ///   * Royston–Parmar (curved survival-probability inverse link).
3089    /// The integral collapses to the plug-in (so the cheaper plug-in path is
3090    /// exact and taken instead) only for the effectively-linear identity-link
3091    /// Gaussian. Any model carrying a link wiggle or baseline-time wiggle is
3092    /// curved regardless of family. This curvature partition mirrors
3093    /// `families::family_runtime::posterior_mean`, the compute path that produces the
3094    /// corrected mean for each of these families.
3095    ///
3096    /// This is the single source of truth shared by the CLI (`gam predict`)
3097    /// and the Python FFI prediction path so the two can never drift on which
3098    /// models receive the posterior-mean correction.
3099    #[inline]
3100    pub fn prediction_uses_posterior_mean(&self) -> bool {
3101        let family = self.likelihood();
3102        let curved_family = match &family.response {
3103            // Identity-link Gaussian: inverse link is linear, so the posterior
3104            // mean equals the plug-in and the cheaper exact path is taken.
3105            ResponseFamily::Gaussian => false,
3106            // Log-link families: E[exp η] = exp(η + se²/2) ≠ exp(η).
3107            ResponseFamily::Poisson
3108            | ResponseFamily::Gamma
3109            | ResponseFamily::Tweedie { .. }
3110            | ResponseFamily::NegativeBinomial { .. } => true,
3111            // Beta (logit link): E[σ(η)] ≠ σ(E[η]).
3112            ResponseFamily::Beta { .. } => true,
3113            // Royston–Parmar: curved survival-probability inverse link.
3114            ResponseFamily::RoystonParmar => true,
3115            // Binomial: every link variant (logit / probit / cloglog / SAS /
3116            // BetaLogistic / Mixture / LatentCLogLog) is a curved sigmoid.
3117            ResponseFamily::Binomial => matches!(
3118                &family.link,
3119                InverseLink::Standard(_)
3120                    | InverseLink::Sas(_)
3121                    | InverseLink::BetaLogistic(_)
3122                    | InverseLink::Mixture(_)
3123                    | InverseLink::LatentCLogLog(_)
3124            ),
3125        };
3126        curved_family || self.has_link_wiggle() || self.has_baseline_time_wiggle()
3127    }
3128
3129    pub fn saved_prediction_runtime(&self) -> Result<SavedPredictionRuntime, FittedModelError> {
3130        self.payload().validate_payload_version()?;
3131        if matches!(
3132            self.predict_model_class(),
3133            PredictModelClass::BernoulliMarginalSlope | PredictModelClass::Survival
3134        ) {
3135            if let Some(runtime) = self.payload().score_warp_runtime.as_ref() {
3136                runtime.validate_exact_replay_contract().map_err(|err| {
3137                    FittedModelError::PayloadCorrupt {
3138                        reason: format!("saved anchored score-warp runtime is invalid: {err}"),
3139                    }
3140                })?;
3141            }
3142            if let Some(runtime) = self.payload().link_deviation_runtime.as_ref() {
3143                runtime.validate_exact_replay_contract().map_err(|err| {
3144                    FittedModelError::PayloadCorrupt {
3145                        reason: format!("saved anchored link-deviation runtime is invalid: {err}"),
3146                    }
3147                })?;
3148            }
3149        }
3150        let runtime = SavedPredictionRuntime {
3151            model_class: self.predict_model_class(),
3152            likelihood: self.likelihood(),
3153            inverse_link: self.resolved_inverse_link()?,
3154            link_wiggle: self.saved_link_wiggle()?,
3155            baseline_time_wiggle: self.saved_baseline_time_wiggle()?,
3156            score_warp: self.payload().score_warp_runtime.clone(),
3157            link_deviation: self.payload().link_deviation_runtime.clone(),
3158            latent_z_rank_int_calibration: self.payload().latent_z_rank_int_calibration.clone(),
3159            latent_z_conditional_calibration: self
3160                .payload()
3161                .latent_z_conditional_calibration
3162                .clone(),
3163            influence_absorber_width: self.payload().influence_absorber_width,
3164        };
3165        if matches!(
3166            runtime.model_class,
3167            PredictModelClass::GaussianLocationScale
3168                | PredictModelClass::BinomialLocationScale
3169                | PredictModelClass::DispersionLocationScale
3170        ) {
3171            let fit = self.payload().fit_result.as_ref().ok_or_else(|| {
3172                FittedModelError::MissingField {
3173                    reason: "location-scale model is missing canonical fit_result payload"
3174                        .to_string(),
3175                }
3176            })?;
3177            validate_location_scale_saved_fit(
3178                fit,
3179                runtime.model_class,
3180                runtime.link_wiggle.as_ref(),
3181            )?;
3182        } else if matches!(runtime.model_class, PredictModelClass::Survival)
3183            && self
3184                .payload()
3185                .survival_likelihood
3186                .as_deref()
3187                .is_some_and(|value| value.eq_ignore_ascii_case("location-scale"))
3188        {
3189            validate_survival_location_scale_saved_fit(
3190                self.payload(),
3191                runtime.link_wiggle.as_ref(),
3192            )?;
3193        } else if matches!(
3194            runtime.model_class,
3195            PredictModelClass::BernoulliMarginalSlope
3196        ) {
3197            let unified =
3198                self.payload()
3199                    .unified
3200                    .as_ref()
3201                    .ok_or_else(|| FittedModelError::MissingField {
3202                        reason: "marginal-slope model is missing unified fit payload; refit"
3203                            .to_string(),
3204                    })?;
3205            validate_marginal_slope_saved_fit(
3206                unified,
3207                runtime.score_warp.as_ref(),
3208                runtime.link_deviation.as_ref(),
3209                "unified",
3210            )?;
3211        } else if matches!(runtime.model_class, PredictModelClass::Survival)
3212            && self
3213                .payload()
3214                .survival_likelihood
3215                .as_deref()
3216                .is_some_and(|value| value.eq_ignore_ascii_case("marginal-slope"))
3217        {
3218            let fit = self.payload().fit_result.as_ref().ok_or_else(|| {
3219                FittedModelError::MissingField {
3220                    reason: "survival marginal-slope model is missing canonical fit_result payload"
3221                        .to_string(),
3222                }
3223            })?;
3224            validate_survival_marginal_slope_saved_fit(
3225                fit,
3226                runtime.score_warp.as_ref(),
3227                runtime.link_deviation.as_ref(),
3228                "fit_result",
3229            )?;
3230        }
3231        Ok(runtime)
3232    }
3233
3234    pub fn saved_sas_state(&self) -> Result<Option<SasLinkState>, FittedModelError> {
3235        let payload = self.payload();
3236        let raw = match &payload.family_state {
3237            FittedFamily::Standard {
3238                likelihood,
3239                sas_state,
3240                ..
3241            } if likelihood.is_binomial_sas() => {
3242                (*sas_state).ok_or_else(|| FittedModelError::MissingField {
3243                    reason: "binomial-sas model is missing state in family_state.sas_state"
3244                        .to_string(),
3245                })?
3246            }
3247            FittedFamily::LocationScale {
3248                likelihood,
3249                base_link,
3250            } if likelihood.is_binomial_sas() => match base_link {
3251                Some(InverseLink::Sas(state)) => *state,
3252                _ => {
3253                    return Err(FittedModelError::MissingField {
3254                        reason: "binomial-sas location-scale model is missing SAS base_link state"
3255                            .to_string(),
3256                    });
3257                }
3258            },
3259            _ => return Ok(None),
3260        };
3261        state_from_sasspec(SasLinkSpec {
3262            initial_epsilon: raw.epsilon,
3263            initial_log_delta: raw.log_delta,
3264        })
3265        .map(Some)
3266        .map_err(|e| FittedModelError::PayloadCorrupt {
3267            reason: format!("invalid saved SAS link state: {e}"),
3268        })
3269    }
3270
3271    pub fn saved_beta_logistic_state(&self) -> Result<Option<SasLinkState>, FittedModelError> {
3272        let payload = self.payload();
3273        let raw = match &payload.family_state {
3274            FittedFamily::Standard {
3275                likelihood,
3276                sas_state,
3277                ..
3278            } if likelihood.is_binomial_beta_logistic() => {
3279                (*sas_state).ok_or_else(|| FittedModelError::MissingField {
3280                    reason:
3281                        "binomial-beta-logistic model is missing state in family_state.sas_state"
3282                            .to_string(),
3283                })?
3284            }
3285            FittedFamily::LocationScale {
3286                likelihood,
3287                base_link,
3288            } if likelihood.is_binomial_beta_logistic() => match base_link {
3289                Some(InverseLink::BetaLogistic(state)) => *state,
3290                _ => {
3291                    return Err(FittedModelError::MissingField {
3292                        reason:
3293                            "binomial-beta-logistic location-scale model is missing beta-logistic base_link state"
3294                                .to_string(),
3295                    });
3296                }
3297            },
3298            _ => return Ok(None),
3299        };
3300        state_from_beta_logisticspec(SasLinkSpec {
3301            initial_epsilon: raw.epsilon,
3302            initial_log_delta: raw.log_delta,
3303        })
3304        .map(Some)
3305        .map_err(|e| FittedModelError::PayloadCorrupt {
3306            reason: format!("invalid saved Beta-Logistic link state: {e}"),
3307        })
3308    }
3309
3310    pub fn saved_mixture_state(&self) -> Result<Option<MixtureLinkState>, FittedModelError> {
3311        let payload = self.payload();
3312        match &payload.family_state {
3313            FittedFamily::Standard {
3314                likelihood,
3315                mixture_state,
3316                ..
3317            } if likelihood.is_binomial_mixture() => mixture_state
3318                .clone()
3319                .ok_or_else(|| FittedModelError::MissingField {
3320                    reason: "binomial-mixture model is missing state in family_state.mixture_state"
3321                        .to_string(),
3322                })
3323                .map(Some),
3324            FittedFamily::LocationScale {
3325                likelihood,
3326                base_link,
3327            } if likelihood.is_binomial_mixture() => match base_link {
3328                Some(InverseLink::Mixture(state)) => Ok(Some(state.clone())),
3329                _ => Err(FittedModelError::MissingField {
3330                    reason:
3331                        "binomial-mixture location-scale model is missing mixture base_link state"
3332                            .to_string(),
3333                }),
3334            },
3335            _ => Ok(None),
3336        }
3337    }
3338
3339    pub fn saved_latent_cloglog_state(
3340        &self,
3341    ) -> Result<Option<LatentCLogLogState>, FittedModelError> {
3342        let payload = self.payload();
3343        match &payload.family_state {
3344            FittedFamily::Standard {
3345                likelihood,
3346                latent_cloglog_state,
3347                ..
3348            } if likelihood.is_latent_cloglog() => latent_cloglog_state
3349                .ok_or_else(|| FittedModelError::MissingField {
3350                    reason:
3351                        "latent-cloglog-binomial model is missing state in family_state.latent_cloglog_state"
3352                            .to_string(),
3353                })
3354                .map(Some),
3355            _ => Ok(None),
3356        }
3357    }
3358
3359    pub fn resolved_inverse_link(&self) -> Result<Option<InverseLink>, FittedModelError> {
3360        let stateful = if let Some(state) = self.saved_mixture_state()? {
3361            Some(InverseLink::Mixture(state))
3362        } else if let Some(state) = self.saved_latent_cloglog_state()? {
3363            Some(InverseLink::LatentCLogLog(state))
3364        } else if let Some(state) = self.saved_beta_logistic_state()? {
3365            Some(InverseLink::BetaLogistic(state))
3366        } else {
3367            self.saved_sas_state()?.map(InverseLink::Sas)
3368        };
3369        match &self.payload().family_state {
3370            FittedFamily::LocationScale { base_link, .. } => Ok(base_link.clone().or(stateful)),
3371            FittedFamily::Standard { link, .. } => {
3372                Ok(stateful.or_else(|| link.map(InverseLink::Standard)))
3373            }
3374            FittedFamily::MarginalSlope { base_link, .. } => Ok(Some(base_link.clone())),
3375            FittedFamily::Survival { .. }
3376            | FittedFamily::LatentSurvival { .. }
3377            | FittedFamily::LatentBinary { .. } => Ok(None),
3378            FittedFamily::TransformationNormal { .. } => Ok(None),
3379        }
3380    }
3381
3382    /// V∞ §5 coverage floor for the measure-jet extrapolation variance: a
3383    /// band level "covers" a query once its kernel mass reaches this fraction
3384    /// of that level's web-averaged support. Magic-by-default (no dial):
3385    /// 0.05 keeps the ε★ gate's bounded discontinuity at ≤ 5 % of the
3386    /// spectrum's total prior ignorance (see the monotonicity theorem in
3387    /// `terms/basis/measure_jet_predict.rs`) while still refusing credit
3388    /// for stray sub-floor kernel mass at levels finer than the first
3389    /// covering scale.
3390    const MEASURE_JET_COVERAGE_FLOOR: f64 = 0.05;
3391
3392    /// V∞ §5 producer: per-row measure-jet extrapolation variance on the η
3393    /// scale for a prediction batch (`docs/measure_jet_v_infinity.md`).
3394    ///
3395    /// For every frozen measure-jet term in `resolved_termspec` this prices
3396    /// the off-support ignorance of the fitted multiscale spectrum at each
3397    /// query row: support curve from the frozen nodes/masses/band
3398    /// ([`gam_terms::basis::measure_jet_support_curve`]), fitted per-scale
3399    /// amplitudes λ̂_ℓ read from the fit's `lambdas` through the replayed
3400    /// design's penalty layout, folded through
3401    /// [`gam_terms::basis::measure_jet_extrapolation_variance`] and scaled by
3402    /// the fit's coefficient-covariance scale φ̂ so the result sits on Vp's
3403    /// η-variance scale. Terms not yet frozen (no `frozen_quadrature` or
3404    /// non-`UserProvided` centers) are skipped with a warning. Returns
3405    /// `Ok(None)` when no measure-jet term contributes, so callers leave
3406    /// `PredictUncertaintyOptions::extrapolation_variance` untouched.
3407    ///
3408    /// `data` must be the RAW (unclipped) prediction matrix in prediction
3409    /// column order — clipping to the training ranges would freeze the
3410    /// distance signal at the hull and defeat the honesty contract — and
3411    /// `col_map` the prediction header → column map (the same map handed to
3412    /// the design builder). This is the minimal-plumbing producer seam: the
3413    /// option-building callers (CLI predict, FFI) hold exactly
3414    /// `(model, data, col_map)` at the point where they assemble
3415    /// `PredictUncertaintyOptions`, and the fusion in
3416    /// `predict_gamwith_uncertainty` adds the array AFTER its multiplicative
3417    /// inflations: `Var_total = Var_Vp·inflation + Var_extrap`.
3418    pub fn measure_jet_extrapolation_variance(
3419        &self,
3420        data: ndarray::ArrayView2<'_, f64>,
3421        col_map: &HashMap<String, usize>,
3422    ) -> Result<Option<Array1<f64>>, FittedModelError> {
3423        use gam_terms::basis::{
3424            CenterStrategy, MeasureJetExtrapolationSpectrum, MeasureJetIdentifiability,
3425            PenaltySource,
3426        };
3427        use gam_terms::smooth::SmoothBasisSpec;
3428        use gam_terms::smooth::build_term_collection_design;
3429        let Some(saved_spec) = self.resolved_termspec.as_ref() else {
3430            return Ok(None);
3431        };
3432        if data.nrows() == 0
3433            || !saved_spec
3434                .smooth_terms
3435                .iter()
3436                .any(|t| matches!(t.basis, SmoothBasisSpec::MeasureJet { .. }))
3437        {
3438            return Ok(None);
3439        }
3440        let fit = self
3441            .fit_result
3442            .as_ref()
3443            .ok_or_else(|| FittedModelError::MissingField {
3444                reason: "measure-jet extrapolation variance requires the canonical \
3445                    fit_result payload; refit"
3446                    .to_string(),
3447            })?;
3448        let spec = crate::survival::predict::resolve_termspec_for_prediction(
3449            &self.resolved_termspec,
3450            self.training_headers.as_ref(),
3451            col_map,
3452            "resolved_termspec",
3453        )
3454        .map_err(|e| FittedModelError::SchemaMismatch {
3455            reason: format!("measure-jet extrapolation variance: {e}"),
3456        })?;
3457        // Penalty layout replay: the global penalty indices (→ `fit.lambdas`)
3458        // come from the SAME design builder the predict pipeline uses. One
3459        // probe row suffices — for a frozen spec the penalty layout is
3460        // row-count-invariant (centers, masses, band, and identifiability
3461        // transforms all replay verbatim) — keeping this O(centers²) instead
3462        // of duplicating the full O(rows·centers) prediction design build.
3463        let probe = data.slice(ndarray::s![0..1, ..]);
3464        let design = build_term_collection_design(probe, &spec).map_err(|e| {
3465            FittedModelError::SchemaMismatch {
3466                reason: format!(
3467                    "measure-jet extrapolation variance: penalty-layout replay failed: {e}"
3468                ),
3469            }
3470        })?;
3471        let lambdas = &fit.lambdas;
3472        // λ̂ are fitted on Frobenius-normalized penalties. The term loop
3473        // unnormalizes them to physical precisions before pricing; multiplying
3474        // by the coefficient-covariance scale puts Var_extrap on the same
3475        // η-variance scale as Vp.
3476        let phi_scale = fit.coefficient_covariance_scale();
3477        let mut total = Array1::<f64>::zeros(data.nrows());
3478        let mut contributed = false;
3479        for (smooth_idx, term) in spec.smooth_terms.iter().enumerate() {
3480            let SmoothBasisSpec::MeasureJet {
3481                feature_cols,
3482                spec: mj,
3483                input_scales,
3484            } = &term.basis
3485            else {
3486                continue;
3487            };
3488            let (Some(frozen), CenterStrategy::UserProvided(centers)) =
3489                (mj.frozen_quadrature.as_ref(), &mj.center_strategy)
3490            else {
3491                log::warn!(
3492                    "measure-jet term '{}' is not frozen (UserProvided centers + frozen \
3493                    quadrature); skipping its extrapolation variance",
3494                    term.name
3495                );
3496                continue;
3497            };
3498            let n_levels = frozen.eps_band.len();
3499            // λ̂ per level from the replayed layout: per-scale candidates carry
3500            // `PenaltySource::Other("measure_jet_scale_ℓ")`; fused
3501            // (pinned-order) mode carries one Primary charged once for the
3502            // whole band. The DoublePenaltyNullspace ridge is EXCLUDED — it shrinks
3503            // coefficients, it is not a scale amplitude, and counting it would
3504            // double-charge the spectrum.
3505            let read_lambda = |global_index: usize| -> Result<f64, FittedModelError> {
3506                lambdas
3507                    .get(global_index)
3508                    .copied()
3509                    .ok_or_else(|| FittedModelError::SchemaMismatch {
3510                        reason: format!(
3511                            "measure-jet term '{}': penalty global index {global_index} out \
3512                            of bounds for {} fitted lambdas",
3513                            term.name,
3514                            lambdas.len()
3515                        ),
3516                    })
3517            };
3518            let mut per_scale: Vec<(usize, f64)> = Vec::new();
3519            let mut fused: Option<f64> = None;
3520            for info in &design.penaltyinfo {
3521                if info.termname.as_deref() != Some(term.name.as_str()) {
3522                    continue;
3523                }
3524                match &info.penalty.source {
3525                    PenaltySource::Other(label) => {
3526                        if let Some(level_txt) = label.strip_prefix("measure_jet_scale_") {
3527                            let level: usize = level_txt.parse().map_err(|_| {
3528                                FittedModelError::SchemaMismatch {
3529                                    reason: format!(
3530                                        "measure-jet term '{}': unparseable penalty label \
3531                                        '{label}'",
3532                                        term.name
3533                                    ),
3534                                }
3535                            })?;
3536                            per_scale.push((level, read_lambda(info.global_index)?));
3537                        }
3538                    }
3539                    PenaltySource::Primary => {
3540                        fused = Some(read_lambda(info.global_index)?);
3541                    }
3542                    _ => {}
3543                }
3544            }
3545            let mut lambda_phys = Vec::with_capacity(n_levels);
3546            let spectrum = if per_scale.is_empty() {
3547                let Some(lam) = fused else {
3548                    log::warn!(
3549                        "measure-jet term '{}' has no fitted amplitude in the penalty \
3550                        layout; skipping its extrapolation variance",
3551                        term.name
3552                    );
3553                    continue;
3554                };
3555                let Some(c) = frozen.fused_penalty_normalization_scale else {
3556                    log::warn!(
3557                        "measure-jet term '{}' is missing the fused penalty normalization scale; \
3558                        skipping its extrapolation variance",
3559                        term.name
3560                    );
3561                    continue;
3562                };
3563                MeasureJetExtrapolationSpectrum::Fused(lam / c)
3564            } else {
3565                per_scale.sort_by_key(|&(level, _)| level);
3566                let levels_complete = per_scale.len() == n_levels
3567                    && per_scale
3568                        .iter()
3569                        .enumerate()
3570                        .all(|(i, &(level, _))| level == i);
3571                if !levels_complete {
3572                    log::warn!(
3573                        "measure-jet term '{}': {} fitted per-scale amplitudes for {} band \
3574                        scales; skipping its extrapolation variance",
3575                        term.name,
3576                        per_scale.len(),
3577                        n_levels
3578                    );
3579                    continue;
3580                }
3581                if frozen.penalty_normalization_scales.len() != n_levels {
3582                    log::warn!(
3583                        "measure-jet term '{}': {} frozen penalty normalization scales for {} \
3584                        band scales; skipping its extrapolation variance",
3585                        term.name,
3586                        frozen.penalty_normalization_scales.len(),
3587                        n_levels
3588                    );
3589                    continue;
3590                }
3591                lambda_phys.extend(
3592                    per_scale
3593                        .iter()
3594                        .map(|&(level, lam)| lam / frozen.penalty_normalization_scales[level]),
3595                );
3596                MeasureJetExtrapolationSpectrum::PerLevel(&lambda_phys)
3597            };
3598            // Query rows in the frozen geometry's coordinates: select the
3599            // term's axes and replay the per-axis standardization exactly as
3600            // the build dispatch does (divide by σ_a when input_scales is
3601            // Some; the persisted centers are already post-standardization).
3602            let mut queries = Array2::<f64>::zeros((data.nrows(), feature_cols.len()));
3603            for (j, &col) in feature_cols.iter().enumerate() {
3604                if col >= data.ncols() {
3605                    return Err(FittedModelError::SchemaMismatch {
3606                        reason: format!(
3607                            "measure-jet term '{}': prediction column {col} out of bounds \
3608                            for {} data columns",
3609                            term.name,
3610                            data.ncols()
3611                        ),
3612                    });
3613                }
3614                queries.column_mut(j).assign(&data.column(col));
3615            }
3616            if let Some(scales) = input_scales {
3617                if scales.len() != feature_cols.len() {
3618                    return Err(FittedModelError::SchemaMismatch {
3619                        reason: format!(
3620                            "measure-jet term '{}': {} input scales for {} axes",
3621                            term.name,
3622                            scales.len(),
3623                            feature_cols.len()
3624                        ),
3625                    });
3626                }
3627                for (j, &scale) in scales.iter().enumerate() {
3628                    queries.column_mut(j).mapv_inplace(|v| v / scale);
3629                }
3630            }
3631            let support = gam_terms::basis::measure_jet_support_curve(
3632                queries.view(),
3633                centers.view(),
3634                frozen.masses.view(),
3635                &frozen.eps_band,
3636            )
3637            .map_err(|e| FittedModelError::SchemaMismatch {
3638                reason: format!(
3639                    "measure-jet term '{}': support curve failed: {e}",
3640                    term.name
3641                ),
3642            })?;
3643            for i in 0..data.nrows() {
3644                let v = gam_terms::basis::measure_jet_extrapolation_variance(
3645                    support.row(i),
3646                    &frozen.eps_band,
3647                    &frozen.support_means,
3648                    spectrum,
3649                    Self::MEASURE_JET_COVERAGE_FLOOR,
3650                )
3651                .map_err(|e| FittedModelError::SchemaMismatch {
3652                    reason: format!(
3653                        "measure-jet term '{}': extrapolation variance failed: {e}",
3654                        term.name
3655                    ),
3656                })?;
3657                total[i] += phi_scale * v;
3658            }
3659            contributed = true;
3660            // #2225 errors-in-variables input-measurement-error term. When the
3661            // fit froze an ambient input-noise scale σ_coord, price the
3662            // delta-method propagation of that input noise through the fitted
3663            // surface: `Var_input(x★) = σ_coord²·‖∇f̂(x★)‖²`, evaluated with the
3664            // analytic ambient gradient ∇f̂ (representers + head). This prices the
3665            // TANGENTIAL movement of a noisy query along the manifold — where the
3666            // intrinsic response genuinely changes — complementary to the
3667            // PERPENDICULAR off-web ignorance the extrapolation term prices. It is
3668            // a pure propagation of the point estimate, so — unlike the
3669            // λ̂⁻¹-scaled extrapolation term — it carries NO φ̂ factor. Everything
3670            // is in the frozen centers' (standardized) frame: `queries`, centers,
3671            // `mj.length_scale`, and σ_coord are all standardized consistently.
3672            if let Some(sigma_coord) = frozen.sigma_coord {
3673                'input_var: {
3674                    let MeasureJetIdentifiability::FrozenTransform { transform } =
3675                        &mj.identifiability
3676                    else {
3677                        log::warn!(
3678                            "measure-jet term '{}': identifiability is not a frozen transform; \
3679                             skipping its input-measurement-error variance",
3680                            term.name
3681                        );
3682                        break 'input_var;
3683                    };
3684                    let full_cols = design.design.ncols();
3685                    if fit.beta.len() != full_cols {
3686                        log::warn!(
3687                            "measure-jet term '{}': joint coefficient vector length {} disagrees \
3688                             with the replayed design's {} columns; skipping its \
3689                             input-measurement-error variance",
3690                            term.name,
3691                            fit.beta.len(),
3692                            full_cols
3693                        );
3694                        break 'input_var;
3695                    }
3696                    if design.smooth.term_designs.len() != spec.smooth_terms.len() {
3697                        log::warn!(
3698                            "measure-jet term '{}': smooth design/term count mismatch ({} vs {}); \
3699                             skipping its input-measurement-error variance",
3700                            term.name,
3701                            design.smooth.term_designs.len(),
3702                            spec.smooth_terms.len()
3703                        );
3704                        break 'input_var;
3705                    }
3706                    let m = centers.nrows();
3707                    let m_aug = transform.nrows();
3708                    let reduced = transform.ncols();
3709                    let term_cols = design.smooth.term_designs[smooth_idx].ncols();
3710                    if term_cols != reduced {
3711                        log::warn!(
3712                            "measure-jet term '{}': replayed reduced width {term_cols} disagrees \
3713                             with the frozen transform ({m_aug}×{reduced}); skipping its \
3714                             input-measurement-error variance",
3715                            term.name
3716                        );
3717                        break 'input_var;
3718                    }
3719                    // Reduced-coefficient block of this term inside the joint β̂.
3720                    let smooth_start = full_cols - design.smooth.total_smooth_cols();
3721                    let offset_in_smooth: usize = design.smooth.term_designs[..smooth_idx]
3722                        .iter()
3723                        .map(|d| d.ncols())
3724                        .sum();
3725                    let g0 = smooth_start + offset_in_smooth;
3726                    let beta_term = fit.beta.slice(ndarray::s![g0..g0 + term_cols]).to_owned();
3727                    // Lift to raw representer+head coefficients z_full = Z·β̂_term.
3728                    let z_full = transform.dot(&beta_term);
3729                    let head_rank = m_aug - m;
3730                    let rep = z_full.slice(ndarray::s![..m]).to_owned();
3731                    let head_coeffs = z_full.slice(ndarray::s![m..]).to_owned();
3732                    let head_t = if head_rank > 0 {
3733                        Some(gam_terms::basis::measure_jet_affine_head_transform(
3734                            centers.view(),
3735                            frozen.masses.view(),
3736                        ))
3737                    } else {
3738                        None
3739                    };
3740                    if let Some(t) = head_t.as_ref() {
3741                        if t.ncols() != head_rank {
3742                            log::warn!(
3743                                "measure-jet term '{}': reconstructed head lift rank {} disagrees \
3744                                 with the frozen head block {head_rank}; skipping its \
3745                                 input-measurement-error variance",
3746                                term.name,
3747                                t.ncols()
3748                            );
3749                            break 'input_var;
3750                        }
3751                    }
3752                    let sigma2 = sigma_coord * sigma_coord;
3753                    let mut input_var = Array1::<f64>::zeros(data.nrows());
3754                    let mut ok = true;
3755                    for i in 0..data.nrows() {
3756                        match gam_terms::basis::measure_jet_ambient_gradient(
3757                            queries.row(i),
3758                            centers.view(),
3759                            rep.view(),
3760                            mj.length_scale,
3761                            head_t.as_ref().map(|t| t.view()),
3762                            head_coeffs.view(),
3763                        ) {
3764                            Ok(grad) => {
3765                                let norm_sq: f64 = grad.iter().map(|g| g * g).sum();
3766                                input_var[i] = sigma2 * norm_sq;
3767                            }
3768                            Err(e) => {
3769                                log::warn!(
3770                                    "measure-jet term '{}': ambient gradient failed ({e}); \
3771                                     skipping its input-measurement-error variance",
3772                                    term.name
3773                                );
3774                                ok = false;
3775                                break;
3776                            }
3777                        }
3778                    }
3779                    if ok {
3780                        total += &input_var;
3781                    }
3782                }
3783            }
3784        }
3785        Ok(contributed.then_some(total))
3786    }
3787
3788    /// Access the unified fit result, if stored.
3789    pub fn unified(&self) -> Option<&UnifiedFitResult> {
3790        self.payload().unified.as_ref()
3791    }
3792
3793    pub fn load_from_path(path: &Path) -> Result<Self, FittedModelError> {
3794        let payload = fs::read_to_string(path).map_err(|e| FittedModelError::PayloadCorrupt {
3795            reason: format!("failed to read model '{}': {e}", path.display()),
3796        })?;
3797        let model: Self =
3798            serde_json::from_str(&payload).map_err(|e| FittedModelError::PayloadCorrupt {
3799                reason: format!("failed to parse model json: {e}"),
3800            })?;
3801        let model = model.with_synchronized_stateful_link_metadata();
3802        model.validate_for_persistence()?;
3803        model.validate_numeric_finiteness()?;
3804        Ok(model)
3805    }
3806
3807    pub fn save_to_path(&self, path: &Path) -> Result<(), FittedModelError> {
3808        let normalized = self.clone().with_synchronized_stateful_link_metadata();
3809        normalized.validate_for_persistence()?;
3810        normalized.validate_numeric_finiteness()?;
3811        // Write to a sibling temp file, fsync, then rename into place so a
3812        // crash mid-write never corrupts the user's existing saved fit.
3813        // Concurrent writers to the same path each have a distinct temp
3814        // suffix (pid + nanos), so neither stomps the other's in-flight
3815        // bytes; the rename winner is last-rename-wins, which is the
3816        // expected last-write-wins semantics for a single canonical path.
3817        let parent = path.parent().unwrap_or_else(|| Path::new("."));
3818        let file_name = path
3819            .file_name()
3820            .and_then(|s| s.to_str())
3821            .unwrap_or("model.json");
3822        let pid = std::process::id();
3823        let nanos = std::time::SystemTime::now()
3824            .duration_since(std::time::UNIX_EPOCH)
3825            .map(|d| d.as_nanos())
3826            .unwrap_or(0);
3827        let tmp = parent.join(format!(".{file_name}.tmp.{pid}.{nanos:x}"));
3828        let file = fs::File::create(&tmp).map_err(|e| FittedModelError::PayloadCorrupt {
3829            reason: format!("failed to write model '{}': {e}", tmp.display()),
3830        })?;
3831        let mut writer = std::io::BufWriter::new(file);
3832        let ser_result = serde_json::to_writer(&mut writer, &normalized);
3833        if let Err(e) = ser_result {
3834            // Best-effort temp cleanup on serialization failure. flush
3835            // returns io::Result<()>; discarding via `.ok()` is enough.
3836            std::io::Write::flush(&mut writer).ok();
3837            drop(writer);
3838            fs::remove_file(&tmp).ok();
3839            return Err(FittedModelError::PayloadCorrupt {
3840                reason: format!("failed to serialize model: {e}"),
3841            });
3842        }
3843        std::io::Write::flush(&mut writer).map_err(|e| FittedModelError::PayloadCorrupt {
3844            reason: format!("failed to write model '{}': {e}", tmp.display()),
3845        })?;
3846        // Recover the underlying File to fsync its contents before rename.
3847        let inner = writer
3848            .into_inner()
3849            .map_err(|e| FittedModelError::PayloadCorrupt {
3850                reason: format!("failed to flush model '{}': {}", tmp.display(), e.error()),
3851            })?;
3852        inner.sync_all().ok();
3853        drop(inner);
3854        if let Err(e) = fs::rename(&tmp, path) {
3855            fs::remove_file(&tmp).ok();
3856            return Err(FittedModelError::PayloadCorrupt {
3857                reason: format!("failed to publish model '{}': {e}", path.display()),
3858            });
3859        }
3860        // fsync the parent directory so the rename itself is durable
3861        // across a crash; without this, the rename can be lost even though
3862        // file contents reached disk. Best-effort on platforms that don't
3863        // support opening a directory for fsync.
3864        if let Ok(d) = fs::File::open(parent) {
3865            d.sync_all().ok();
3866        }
3867        Ok(())
3868    }
3869
3870    pub fn require_data_schema(&self) -> Result<&DataSchema, FittedModelError> {
3871        self.data_schema
3872            .as_ref()
3873            .ok_or_else(|| FittedModelError::MissingField {
3874                reason: "model is missing data_schema; refit".to_string(),
3875            })
3876    }
3877
3878    /// Restore the exact in-memory spline-scan fit from a scan-bearing
3879    /// payload (#1030/#1034). `Ok(None)` for dense models; the returned
3880    /// `predict` replays the training Gaussian bridge bit-for-bit.
3881    pub fn saved_spline_scan(
3882        &self,
3883    ) -> Result<Option<(&str, gam_solve::spline_scan::SplineScanFit)>, FittedModelError> {
3884        let Some(saved) = self.spline_scan.as_ref() else {
3885            return Ok(None);
3886        };
3887        let fit = gam_solve::spline_scan::SplineScanFit::from_state(&saved.state)
3888            .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
3889        Ok(Some((saved.feature_column.as_str(), fit)))
3890    }
3891
3892    /// Restore the in-memory residual-cascade fit from a cascade-bearing
3893    /// payload (#1032). `Ok(None)` for non-cascade models; the returned fit
3894    /// replays the multilevel Wendland-frame posterior for the d ∈ {2, 3}
3895    /// feature columns at each predict point.
3896    pub fn saved_residual_cascade(
3897        &self,
3898    ) -> Result<
3899        Option<(&[String], gam_solve::residual_cascade::ResidualCascadeFit)>,
3900        FittedModelError,
3901    > {
3902        let Some(saved) = self.residual_cascade.as_ref() else {
3903            return Ok(None);
3904        };
3905        let fit = gam_solve::residual_cascade::ResidualCascadeFit::from_state(&saved.state)
3906            .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
3907        Ok(Some((saved.feature_columns.as_slice(), fit)))
3908    }
3909
3910    /// Grouping columns eligible for LENIENT unseen-level encoding at predict
3911    /// time (the held-out-group policy: an unseen group is encoded as an
3912    /// out-of-vocabulary code and shrunk toward the population mean).
3913    ///
3914    /// This is the whitelist the predict/`check` encode paths pass to
3915    /// [`UnseenCategoryPolicy::encode_unknown_for_columns`]. It intentionally
3916    /// covers ONLY genuine random effects (`group(g)`/`re(g)`/`s(g, bs="re")`).
3917    /// A FIXED categorical factor — a bare `+ g` OR an explicit `factor(g)` —
3918    /// is auto-promoted to a penalized random block internally but is still a
3919    /// fixed parametric factor: an unseen level of it must reach the strict
3920    /// schema encode and raise a `SchemaMismatchError` rather than be silently
3921    /// averaged to the factor's centering point (#2102/#2137). Such terms carry
3922    /// `lenient_unseen == false` and are excluded here so they hit the strict
3923    /// `UnseenCategoryPolicy::Error` arm.
3924    pub fn random_effect_group_columns(&self) -> HashSet<String> {
3925        let Some(training_headers) = self.training_headers.as_ref() else {
3926            return HashSet::new();
3927        };
3928        let mut out = HashSet::<String>::new();
3929        for spec in self.saved_term_specs() {
3930            for term in &spec.random_effect_terms {
3931                if !term.lenient_unseen {
3932                    continue;
3933                }
3934                if let Some(name) = training_headers.get(term.feature_col) {
3935                    out.insert(name.clone());
3936                }
3937            }
3938        }
3939        out
3940    }
3941
3942    /// Frozen level vocabularies for FIXED-factor terms (`factor(g)` or a bare
3943    /// `+ g`, i.e. `lenient_unseen == false`) whose feature column is *numeric*
3944    /// in the data schema.
3945    ///
3946    /// A string factor is a `Categorical` schema column, so the strict schema
3947    /// re-encode already rejects (and `check` reports) an out-of-vocabulary
3948    /// label. A numeric-coded `factor(year)`, however, reaches the model as a
3949    /// `Continuous`/`Binary` column with no categorical schema, so the encode
3950    /// path has no level set to validate against — the unseen-level guard is
3951    /// silently skipped (#2137). This exposes each such column's frozen numeric
3952    /// vocabulary (canonical `f64` bit patterns, signed-zero/NaN normalized) so
3953    /// the `check`/`predict` schema layer can enforce the same fixed-factor
3954    /// contract the design operator (`build_random_effect_block`) enforces.
3955    ///
3956    /// Only terms with concrete `frozen_levels` (captured at fit) and the full
3957    /// one-hot block (`!drop_first_level`, so the frozen set is the complete
3958    /// training vocabulary) are returned, matching the operator's strict gate.
3959    pub fn numeric_fixed_factor_vocabularies(&self) -> Vec<(String, HashSet<u64>)> {
3960        let Some(training_headers) = self.training_headers.as_ref() else {
3961            return Vec::new();
3962        };
3963        let Some(schema) = self.data_schema.as_ref() else {
3964            return Vec::new();
3965        };
3966        let mut out = Vec::<(String, HashSet<u64>)>::new();
3967        for spec in self.saved_term_specs() {
3968            for term in &spec.random_effect_terms {
3969                if term.lenient_unseen || term.drop_first_level {
3970                    continue;
3971                }
3972                let Some(levels) = term.frozen_levels.as_ref() else {
3973                    continue;
3974                };
3975                let Some(name) = training_headers.get(term.feature_col) else {
3976                    continue;
3977                };
3978                // Skip string factors: they are Categorical in the schema and
3979                // are already validated by the typed encode.
3980                let is_numeric = schema
3981                    .columns
3982                    .iter()
3983                    .find(|c| &c.name == name)
3984                    .map(|c| matches!(c.kind, ColumnKindTag::Continuous | ColumnKindTag::Binary))
3985                    .unwrap_or(false);
3986                if !is_numeric {
3987                    continue;
3988                }
3989                let vocab: HashSet<u64> = levels
3990                    .iter()
3991                    .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
3992                    .collect();
3993                out.push((name.clone(), vocab));
3994            }
3995        }
3996        out
3997    }
3998
3999    pub fn validate_for_persistence(&self) -> Result<(), FittedModelError> {
4000        // Hard version gate. The struct's ~40 Option<T> fields carry
4001        // `#[serde(default)]`, which is by design forward-compatible: old
4002        // payloads missing a new optional field decode with `None`. BUT:
4003        // when a new CLI release adds a required field for some family_state
4004        // (enforced below), an older model loaded by the newer CLI would have
4005        // `None` in that slot and the family-specific branch below would
4006        // correctly reject it — unless the new field also happens to slot
4007        // under a branch that hasn't been touched. Conversely, a newer model
4008        // loaded by an older CLI silently drops fields the older struct
4009        // doesn't know about. Both directions are silent-drift hazards. We
4010        // close them with an exact-version check anchored to the canonical
4011        // MODEL_PAYLOAD_VERSION constant — every payload must round-trip
4012        // identically between writers and readers running the same schema.
4013        self.validate_payload_version()?;
4014        if self.training_table_kind.trim().is_empty() {
4015            return Err(FittedModelError::MissingField {
4016                reason: "saved model training_table_kind must be non-empty".to_string(),
4017            });
4018        }
4019        if let Some(scan) = self.spline_scan.as_ref() {
4020            // Spline-scan representation (#1030/#1034): the smoother state IS
4021            // the fit. It is exclusive with the dense representation, only
4022            // standard Gaussian-identity models can carry it, and the state
4023            // must restore cleanly so predict never sees a corrupt snapshot.
4024            if self.fit_result.is_some() || self.unified.is_some() {
4025                return Err(FittedModelError::SchemaMismatch {
4026                    reason: "spline-scan model must not also carry a dense fit_result/unified \
4027                             payload; the representations are mutually exclusive"
4028                        .to_string(),
4029                });
4030            }
4031            if self.model_kind != ModelKind::Standard
4032                || self.family_state.likelihood() != LikelihoodSpec::gaussian_identity()
4033            {
4034                return Err(FittedModelError::SchemaMismatch {
4035                    reason: format!(
4036                        "spline-scan representation requires a standard Gaussian-identity model; \
4037                         got model_kind={:?}, likelihood={:?}",
4038                        self.model_kind,
4039                        self.family_state.likelihood()
4040                    ),
4041                });
4042            }
4043            if scan.feature_column.is_empty() {
4044                return Err(FittedModelError::MissingField {
4045                    reason: "spline-scan model is missing its feature column name; refit"
4046                        .to_string(),
4047                });
4048            }
4049            gam_solve::spline_scan::SplineScanFit::from_state(&scan.state)
4050                .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4051            // A scan model carries NO dense design, so the dense-path
4052            // requirements below (resolved_termspec, fit_result finiteness,
4053            // family-specific blocks) do not apply. Enforce only the metadata
4054            // predict actually consumes — the feature column resolves against
4055            // training_headers / data_schema — then accept.
4056            if self.data_schema.is_none() {
4057                return Err(FittedModelError::MissingField {
4058                    reason: "spline-scan model is missing data_schema; refit".to_string(),
4059                });
4060            }
4061            if self.training_headers.is_none() {
4062                return Err(FittedModelError::MissingField {
4063                    reason: "spline-scan model is missing training_headers; refit".to_string(),
4064                });
4065            }
4066            return Ok(());
4067        } else if let Some(cascade) = self.residual_cascade.as_ref() {
4068            // Residual-cascade representation (#1032): a multilevel
4069            // Wendland-frame model for a scattered d ∈ {2,3} Gaussian smooth.
4070            // Exclusive with the dense representation and with the scan.
4071            if self.spline_scan.is_some() || self.fit_result.is_some() || self.unified.is_some() {
4072                return Err(FittedModelError::SchemaMismatch {
4073                    reason: "residual-cascade model must not also carry spline_scan / \
4074                             fit_result / unified payloads; the representations are \
4075                             mutually exclusive"
4076                        .to_string(),
4077                });
4078            }
4079            if self.model_kind != ModelKind::Standard
4080                || self.family_state.likelihood() != LikelihoodSpec::gaussian_identity()
4081            {
4082                return Err(FittedModelError::SchemaMismatch {
4083                    reason: format!(
4084                        "residual-cascade representation requires a standard Gaussian-identity \
4085                         model; got model_kind={:?}, likelihood={:?}",
4086                        self.model_kind,
4087                        self.family_state.likelihood()
4088                    ),
4089                });
4090            }
4091            if cascade.feature_columns.is_empty()
4092                || !(2..=3).contains(&cascade.feature_columns.len())
4093            {
4094                return Err(FittedModelError::MissingField {
4095                    reason: format!(
4096                        "residual-cascade model needs 2 or 3 feature columns; got {}; refit",
4097                        cascade.feature_columns.len()
4098                    ),
4099                });
4100            }
4101            gam_solve::residual_cascade::ResidualCascadeFit::from_state(&cascade.state)
4102                .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4103            if self.data_schema.is_none() {
4104                return Err(FittedModelError::MissingField {
4105                    reason: "residual-cascade model is missing data_schema; refit".to_string(),
4106                });
4107            }
4108            if self.training_headers.is_none() {
4109                return Err(FittedModelError::MissingField {
4110                    reason: "residual-cascade model is missing training_headers; refit".to_string(),
4111                });
4112            }
4113            return Ok(());
4114        } else if self.fit_result.is_none() {
4115            return Err(FittedModelError::MissingField {
4116                reason: "model is missing canonical fit_result payload; refit".to_string(),
4117            });
4118        }
4119        if self.data_schema.is_none() {
4120            return Err(FittedModelError::MissingField {
4121                reason: "model is missing data_schema; refit".to_string(),
4122            });
4123        }
4124        if self.training_headers.is_none() {
4125            return Err(FittedModelError::MissingField {
4126                reason: "model is missing training_headers; refit to guarantee stable feature mapping at prediction time"
4127                    .to_string(),
4128            });
4129        }
4130        let spec = self.resolved_termspec.as_ref().ok_or_else(|| {
4131            FittedModelError::MissingField {
4132                reason: "model is missing resolved_termspec; refit to guarantee train/predict design consistency"
4133                    .to_string(),
4134            }
4135        })?;
4136        validate_frozen_term_collectionspec(spec, "resolved_termspec")?;
4137
4138        if self.formula_noise.is_some() && self.resolved_termspec_noise.is_none() {
4139            return Err(FittedModelError::MissingField {
4140                reason: "model defines formula_noise but is missing resolved_termspec_noise; refit"
4141                    .to_string(),
4142            });
4143        }
4144        if let Some(spec_noise) = self.resolved_termspec_noise.as_ref() {
4145            validate_frozen_term_collectionspec(spec_noise, "resolved_termspec_noise")?;
4146        }
4147        if matches!(self.family_state, FittedFamily::TransformationNormal { .. }) {
4148            let score = self.transformation_score_calibration.ok_or_else(|| {
4149                FittedModelError::MissingField {
4150                    reason: "transformation-normal model is missing transformation_score_calibration; refit"
4151                        .to_string(),
4152                }
4153            })?;
4154            score.validate("transformation-normal model")?;
4155        }
4156        if matches!(self.family_state, FittedFamily::MarginalSlope { .. }) {
4157            if self.formula_logslope.is_none() {
4158                return Err(FittedModelError::MissingField {
4159                    reason: "marginal-slope model is missing formula_logslope; refit".to_string(),
4160                });
4161            }
4162            if self.z_column.is_none() {
4163                return Err(FittedModelError::MissingField {
4164                    reason: "marginal-slope model is missing z_column; refit".to_string(),
4165                });
4166            }
4167            let z_normalization =
4168                self.latent_z_normalization
4169                    .ok_or_else(|| FittedModelError::MissingField {
4170                        reason: "marginal-slope model is missing latent_z_normalization; refit"
4171                            .to_string(),
4172                    })?;
4173            z_normalization.validate("marginal-slope model")?;
4174            let latent_measure =
4175                self.latent_measure
4176                    .as_ref()
4177                    .ok_or_else(|| FittedModelError::MissingField {
4178                        reason: "marginal-slope model is missing latent_measure; refit".to_string(),
4179                    })?;
4180            latent_measure
4181                .validate("marginal-slope model latent_measure")
4182                .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4183            if self.marginal_baseline.is_none() || self.logslope_baseline.is_none() {
4184                return Err(FittedModelError::MissingField {
4185                    reason: "marginal-slope model is missing baseline offsets; refit".to_string(),
4186                });
4187            }
4188            if self.resolved_termspec_logslope.as_ref().is_none() {
4189                return Err(FittedModelError::MissingField {
4190                    reason: "marginal-slope model is missing resolved_termspec_logslope for the logslope surface"
4191                        .to_string(),
4192                });
4193            }
4194            match self.family_state.frailty() {
4195                Some(FrailtySpec::None)
4196                | Some(FrailtySpec::GaussianShift {
4197                    sigma_fixed: Some(_),
4198                }) => {}
4199                Some(FrailtySpec::GaussianShift { sigma_fixed: None }) => {
4200                    return Err(FittedModelError::IncompatibleConfig {
4201                        reason: "marginal-slope model requires a fixed GaussianShift sigma in family_state.frailty"
4202                            .to_string(),
4203                    });
4204                }
4205                Some(FrailtySpec::HazardMultiplier { .. }) => {
4206                    return Err(FittedModelError::IncompatibleConfig {
4207                        reason: "marginal-slope model does not support HazardMultiplier frailty"
4208                            .to_string(),
4209                    });
4210                }
4211                None => {
4212                    return Err(FittedModelError::MissingField {
4213                        reason: "marginal-slope model is missing family_state.frailty; refit"
4214                            .to_string(),
4215                    });
4216                }
4217            }
4218        }
4219
4220        if let FittedFamily::Survival {
4221            survival_likelihood,
4222            frailty,
4223            ..
4224        } = &self.family_state
4225        {
4226            if matches!(
4227                survival_likelihood.as_deref(),
4228                Some("latent") | Some("latent-binary")
4229            ) {
4230                return Err(FittedModelError::SchemaMismatch {
4231                    reason: "latent hazard-window models must persist explicit family_state metadata, not generic survival metadata"
4232                        .to_string(),
4233                });
4234            }
4235            if survival_likelihood.as_deref() == Some("marginal-slope") {
4236                if self.formula_logslope.is_none() {
4237                    return Err(FittedModelError::MissingField {
4238                        reason: "survival marginal-slope model is missing formula_logslope; refit"
4239                            .to_string(),
4240                    });
4241                }
4242                if self.z_column.is_none() {
4243                    return Err(FittedModelError::MissingField {
4244                        reason: "survival marginal-slope model is missing z_column; refit"
4245                            .to_string(),
4246                    });
4247                }
4248                let z_normalization =
4249                    self.latent_z_normalization
4250                        .ok_or_else(|| {
4251                            FittedModelError::MissingField {
4252                        reason:
4253                            "survival marginal-slope model is missing latent_z_normalization; refit"
4254                                .to_string(),
4255                    }
4256                        })?;
4257                z_normalization.validate("survival marginal-slope model")?;
4258                let latent_measure =
4259                    self.latent_measure
4260                        .as_ref()
4261                        .ok_or_else(|| FittedModelError::MissingField {
4262                            reason:
4263                                "survival marginal-slope model is missing latent_measure; refit"
4264                                    .to_string(),
4265                        })?;
4266                latent_measure
4267                    .validate("survival marginal-slope model latent_measure")
4268                    .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4269                if self.logslope_baseline.is_none() {
4270                    return Err(FittedModelError::MissingField {
4271                        reason: "survival marginal-slope model is missing logslope_baseline; refit"
4272                            .to_string(),
4273                    });
4274                }
4275                if self.resolved_termspec_logslope.as_ref().is_none() {
4276                    return Err(FittedModelError::MissingField {
4277                        reason: "survival marginal-slope model is missing resolved_termspec_logslope for the logslope surface"
4278                            .to_string(),
4279                    });
4280                }
4281                match frailty {
4282                    FrailtySpec::None
4283                    | FrailtySpec::GaussianShift {
4284                        sigma_fixed: Some(_),
4285                    } => {}
4286                    FrailtySpec::GaussianShift { sigma_fixed: None } => {
4287                        return Err(FittedModelError::IncompatibleConfig {
4288                            reason: "survival marginal-slope model requires a fixed GaussianShift sigma in family_state.frailty"
4289                                .to_string(),
4290                        });
4291                    }
4292                    FrailtySpec::HazardMultiplier { .. } => {
4293                        return Err(FittedModelError::IncompatibleConfig {
4294                            reason: "survival marginal-slope model does not support HazardMultiplier frailty"
4295                                .to_string(),
4296                        });
4297                    }
4298                }
4299            } else if !matches!(frailty, FrailtySpec::None) {
4300                return Err(FittedModelError::IncompatibleConfig {
4301                    reason:
4302                        "non-marginal survival models do not currently persist a frailty modifier"
4303                            .to_string(),
4304                });
4305            }
4306            // Non-latent survival predict reconstructs the baseline-time
4307            // basis via `load_survival_time_basis_config_from_model` and
4308            // anchors that basis at `survival_time_anchor`; both are
4309            // required for the saved model to be loadable. The CLI's
4310            // marginal-slope+time-wiggle save path previously dropped one or
4311            // the other on partial-write, producing models that loaded but
4312            // would panic at the first predict. Enforce both before persisting.
4313            if self.survival_time_basis.is_none() {
4314                return Err(FittedModelError::MissingField {
4315                    reason: "survival model is missing survival_time_basis; refit to persist the baseline-time basis configuration".to_string(),
4316                });
4317            }
4318            if self.survival_time_anchor.is_none() {
4319                return Err(FittedModelError::MissingField {
4320                    reason: "survival model is missing survival_time_anchor; refit to persist the baseline-time anchor".to_string(),
4321                });
4322            }
4323        }
4324        if let FittedFamily::LatentSurvival { frailty } = &self.family_state {
4325            match frailty {
4326                FrailtySpec::HazardMultiplier {
4327                    sigma_fixed: Some(_),
4328                    ..
4329                } => {}
4330                FrailtySpec::HazardMultiplier {
4331                    sigma_fixed: None, ..
4332                } => {
4333                    return Err(FittedModelError::IncompatibleConfig {
4334                        reason: "latent survival model requires a fixed HazardMultiplier sigma in family_state.frailty"
4335                            .to_string(),
4336                    });
4337                }
4338                FrailtySpec::GaussianShift { .. } | FrailtySpec::None => {
4339                    return Err(FittedModelError::IncompatibleConfig {
4340                        reason: "latent survival model requires a fixed HazardMultiplier frailty specification"
4341                            .to_string(),
4342                    });
4343                }
4344            }
4345            if self.survival_likelihood.as_deref() != Some("latent") {
4346                return Err(FittedModelError::SchemaMismatch {
4347                    reason: "latent survival model must persist survival_likelihood=latent"
4348                        .to_string(),
4349                });
4350            }
4351        }
4352        if let FittedFamily::LatentBinary { frailty } = &self.family_state {
4353            match frailty {
4354                FrailtySpec::HazardMultiplier {
4355                    sigma_fixed: Some(_),
4356                    ..
4357                } => {}
4358                FrailtySpec::HazardMultiplier {
4359                    sigma_fixed: None, ..
4360                } => {
4361                    return Err(FittedModelError::IncompatibleConfig {
4362                        reason: "latent binary model requires a fixed HazardMultiplier sigma in family_state.frailty"
4363                            .to_string(),
4364                    });
4365                }
4366                FrailtySpec::GaussianShift { .. } | FrailtySpec::None => {
4367                    return Err(FittedModelError::IncompatibleConfig {
4368                        reason: "latent binary model requires a fixed HazardMultiplier frailty specification"
4369                            .to_string(),
4370                    });
4371                }
4372            }
4373            if self.survival_likelihood.as_deref() != Some("latent-binary") {
4374                return Err(FittedModelError::SchemaMismatch {
4375                    reason: "latent binary model must persist survival_likelihood=latent-binary"
4376                        .to_string(),
4377                });
4378            }
4379        }
4380
4381        let family_likelihood = match &self.family_state {
4382            FittedFamily::Standard { likelihood, .. }
4383            | FittedFamily::LocationScale { likelihood, .. }
4384            | FittedFamily::MarginalSlope { likelihood, .. }
4385            | FittedFamily::Survival { likelihood, .. }
4386            | FittedFamily::TransformationNormal { likelihood, .. } => Some(likelihood),
4387            FittedFamily::LatentSurvival { .. } | FittedFamily::LatentBinary { .. } => None,
4388        };
4389        let is_standard_or_location_scale = matches!(
4390            self.family_state,
4391            FittedFamily::Standard { .. } | FittedFamily::LocationScale { .. }
4392        );
4393        if is_standard_or_location_scale
4394            && family_likelihood.is_some_and(LikelihoodSpec::is_binomial_sas)
4395        {
4396            self.saved_sas_state()?;
4397        }
4398        if is_standard_or_location_scale
4399            && family_likelihood.is_some_and(LikelihoodSpec::is_binomial_beta_logistic)
4400        {
4401            self.saved_beta_logistic_state()?;
4402        }
4403        if is_standard_or_location_scale
4404            && family_likelihood.is_some_and(LikelihoodSpec::is_binomial_mixture)
4405        {
4406            self.saved_mixture_state()?;
4407        }
4408        if matches!(self.family_state, FittedFamily::Standard { .. })
4409            && family_likelihood.is_some_and(LikelihoodSpec::is_latent_cloglog)
4410        {
4411            self.saved_latent_cloglog_state()?;
4412        }
4413        if matches!(self.family_state, FittedFamily::LocationScale { .. })
4414            && family_likelihood.is_some_and(LikelihoodSpec::is_latent_cloglog)
4415        {
4416            return Err(FittedModelError::IncompatibleConfig {
4417                reason: "latent-cloglog-binomial is not supported for location-scale saved models"
4418                    .to_string(),
4419            });
4420        }
4421        if matches!(self.family_state, FittedFamily::Survival { .. })
4422            && self.survival_likelihood.is_none()
4423        {
4424            return Err(FittedModelError::MissingField {
4425                reason: "saved survival model is missing survival_likelihood metadata; refit"
4426                    .to_string(),
4427            });
4428        }
4429        let has_any_saved_link_wiggle = self.linkwiggle_knots.is_some()
4430            || self.linkwiggle_degree.is_some()
4431            || self.beta_link_wiggle.is_some()
4432            || self
4433                .fit_result
4434                .as_ref()
4435                .and_then(|fit| fit.block_by_role(BlockRole::LinkWiggle))
4436                .is_some();
4437        let saved_link_wiggle = self.saved_link_wiggle()?;
4438        if has_any_saved_link_wiggle && saved_link_wiggle.is_none() {
4439            return Err(FittedModelError::SchemaMismatch {
4440                reason: "saved model has incomplete link-wiggle state; expected metadata and coefficients"
4441                    .to_string(),
4442            });
4443        }
4444        let has_any_saved_baseline_time_wiggle = self.baseline_timewiggle_knots.is_some()
4445            || self.baseline_timewiggle_degree.is_some()
4446            || self.baseline_timewiggle_penalty_orders.is_some()
4447            || self.baseline_timewiggle_double_penalty.is_some()
4448            || self.beta_baseline_timewiggle.is_some()
4449            || self.beta_baseline_timewiggle_by_cause.is_some();
4450        let is_joint_cause_specific = self
4451            .survival_cause_count
4452            .is_some_and(|cause_count| cause_count > 1);
4453        if has_any_saved_baseline_time_wiggle {
4454            if is_joint_cause_specific {
4455                let complete = self.baseline_timewiggle_knots.is_some()
4456                    && self.baseline_timewiggle_degree.is_some()
4457                    && self.baseline_timewiggle_penalty_orders.is_some()
4458                    && self.baseline_timewiggle_double_penalty.is_some()
4459                    && self.beta_baseline_timewiggle_by_cause.is_some();
4460                if !complete {
4461                    return Err(FittedModelError::SchemaMismatch {
4462                        reason: "saved joint cause-specific survival model has incomplete baseline-timewiggle state; expected metadata and per-cause coefficients"
4463                            .to_string(),
4464                    });
4465                }
4466            } else if self.saved_baseline_time_wiggle()?.is_none() {
4467                return Err(FittedModelError::SchemaMismatch {
4468                    reason: "saved model has incomplete baseline-timewiggle state; expected metadata and coefficients"
4469                        .to_string(),
4470                });
4471            }
4472        }
4473        if self
4474            .survival_likelihood
4475            .as_deref()
4476            .is_some_and(|value| value.eq_ignore_ascii_case("location-scale"))
4477        {
4478            validate_survival_location_scale_saved_fit(self.payload(), saved_link_wiggle.as_ref())?;
4479        }
4480
4481        // Validate anchored-deviation replay contracts at LOAD/SAVE time rather
4482        // than waiting for first predict call. Previously these contracts
4483        // (span table dimensions, coefficient matrices, etc.) were only
4484        // asserted inside `saved_prediction_runtime`, which runs on the first
4485        // predict invocation. A corrupted runtime would therefore pass
4486        // `load_from_path` silently and fail later under a different error
4487        // surface. Enforcing the same check here makes the model self-
4488        // diagnostic: `gam fit` catches its own bad output at save, and
4489        // `gam predict` catches bad input at load rather than mid-pipeline.
4490        if let Some(runtime) = self.score_warp_runtime.as_ref() {
4491            runtime.validate_exact_replay_contract().map_err(|err| {
4492                FittedModelError::PayloadCorrupt {
4493                    reason: format!("saved anchored score-warp runtime is invalid: {err}"),
4494                }
4495            })?;
4496        }
4497        if let Some(runtime) = self.link_deviation_runtime.as_ref() {
4498            runtime.validate_exact_replay_contract().map_err(|err| {
4499                FittedModelError::PayloadCorrupt {
4500                    reason: format!("saved anchored link-deviation runtime is invalid: {err}"),
4501                }
4502            })?;
4503        }
4504        if matches!(self.family_state, FittedFamily::MarginalSlope { .. }) {
4505            validate_marginal_slope_saved_fit(
4506                self.fit_result.as_ref().expect("checked above"),
4507                self.score_warp_runtime.as_ref(),
4508                self.link_deviation_runtime.as_ref(),
4509                "fit_result",
4510            )?;
4511            let unified = self
4512                .unified
4513                .as_ref()
4514                .ok_or_else(|| FittedModelError::MissingField {
4515                    reason: "marginal-slope model is missing unified fit payload; refit"
4516                        .to_string(),
4517                })?;
4518            validate_marginal_slope_saved_fit(
4519                unified,
4520                self.score_warp_runtime.as_ref(),
4521                self.link_deviation_runtime.as_ref(),
4522                "unified",
4523            )?;
4524        }
4525        if self
4526            .survival_likelihood
4527            .as_deref()
4528            .is_some_and(|value| value.eq_ignore_ascii_case("marginal-slope"))
4529        {
4530            validate_survival_marginal_slope_saved_fit(
4531                self.fit_result.as_ref().expect("checked above"),
4532                self.score_warp_runtime.as_ref(),
4533                self.link_deviation_runtime.as_ref(),
4534                "fit_result",
4535            )?;
4536            if let Some(unified) = self.unified.as_ref() {
4537                validate_survival_marginal_slope_saved_fit(
4538                    unified,
4539                    self.score_warp_runtime.as_ref(),
4540                    self.link_deviation_runtime.as_ref(),
4541                    "unified",
4542                )?;
4543            }
4544        }
4545
4546        // Posterior-mean / uncertainty backends are validated at predict time
4547        // by `prediction_backend_from_model`, which has access to the actual
4548        // requested mode and emits the canonical "nonlinear posterior-mean
4549        // prediction requires either covariance or a saved penalized Hessian"
4550        // error.  Save-time we deliberately do NOT enforce that gate: a fit
4551        // produced for MAP / plug-in scoring can be persisted and replayed
4552        // without ever needing a covariance backend, and gating it here would
4553        // refuse legitimate MAP-only saves whose `UnifiedFitResult` carries
4554        // beta + lambdas without a stabilized Hessian.
4555
4556        Ok(())
4557    }
4558
4559    pub fn validate_numeric_finiteness(&self) -> Result<(), FittedModelError> {
4560        let corrupt = |reason: String| FittedModelError::PayloadCorrupt { reason };
4561        if let Some(fit) = self.fit_result.as_ref() {
4562            fit.validate_numeric_finiteness()
4563                .map_err(|e| corrupt(e.to_string()))?;
4564        }
4565
4566        for (name, opt) in [
4567            ("survival_baseline_scale", self.survival_baseline_scale),
4568            ("survival_baseline_shape", self.survival_baseline_shape),
4569            ("survival_baseline_rate", self.survival_baseline_rate),
4570            ("survival_baseline_makeham", self.survival_baseline_makeham),
4571            (
4572                "survival_time_smooth_lambda",
4573                self.survival_time_smooth_lambda,
4574            ),
4575            ("survival_time_anchor", self.survival_time_anchor),
4576            ("survivalridge_lambda", self.survivalridge_lambda),
4577        ] {
4578            if let Some(v) = opt {
4579                ensure_finite_scalar(name, v).map_err(corrupt)?;
4580            }
4581        }
4582
4583        if let Some(v) = self.beta_noise.as_ref() {
4584            validate_all_finite("beta_noise", v.iter().copied()).map_err(corrupt)?;
4585        }
4586        if let Some(v) = self.noise_projection.as_ref() {
4587            validate_all_finite("noise_projection", v.iter().flatten().copied())
4588                .map_err(corrupt)?;
4589            if self.noise_projection_ridge_alpha.is_none() {
4590                return Err(FittedModelError::MissingField {
4591                    reason:
4592                        "model has noise_projection but is missing noise_projection_ridge_alpha; refit"
4593                            .to_string(),
4594                });
4595            }
4596        }
4597        if let Some(v) = self.noise_center.as_ref() {
4598            validate_all_finite("noise_center", v.iter().copied()).map_err(corrupt)?;
4599        }
4600        if let Some(v) = self.noise_scale.as_ref() {
4601            validate_all_finite("noise_scale", v.iter().copied()).map_err(corrupt)?;
4602        }
4603        if let Some(v) = self.noise_projection_ridge_alpha {
4604            ensure_finite_scalar("noise_projection_ridge_alpha", v).map_err(corrupt)?;
4605            if v < 0.0 {
4606                return Err(FittedModelError::InvalidInput {
4607                    reason: format!("noise_projection_ridge_alpha must be non-negative, got {v}"),
4608                });
4609            }
4610        }
4611        if let Some(v) = self.gaussian_response_scale {
4612            ensure_finite_scalar("gaussian_response_scale", v).map_err(corrupt)?;
4613        }
4614        if let Some(v) = self.beta_link_wiggle.as_ref() {
4615            validate_all_finite("beta_link_wiggle", v.iter().copied()).map_err(corrupt)?;
4616        }
4617        if let Some(v) = self.link_wiggle_index_shift.as_ref() {
4618            validate_all_finite("link_wiggle_index_shift", v.iter().copied()).map_err(corrupt)?;
4619        }
4620        if let Some(v) = self.beta_baseline_timewiggle.as_ref() {
4621            validate_all_finite("beta_baseline_timewiggle", v.iter().copied()).map_err(corrupt)?;
4622        }
4623        if let Some(v) = self.beta_baseline_timewiggle_by_cause.as_ref() {
4624            validate_all_finite(
4625                "beta_baseline_timewiggle_by_cause",
4626                v.iter().flatten().copied(),
4627            )
4628            .map_err(corrupt)?;
4629        }
4630        if let Some(v) = self.latent_z_normalization {
4631            v.validate("latent_z_normalization")?;
4632        }
4633        if let Some(v) = self.latent_measure.as_ref() {
4634            v.validate("latent_measure").map_err(corrupt)?;
4635        }
4636        if let Some(v) = self.survival_beta_time.as_ref() {
4637            validate_all_finite("survival_beta_time", v.iter().copied()).map_err(corrupt)?;
4638        }
4639        if let Some(v) = self.survival_beta_threshold.as_ref() {
4640            validate_all_finite("survival_beta_threshold", v.iter().copied()).map_err(corrupt)?;
4641        }
4642        if let Some(v) = self.survival_beta_log_sigma.as_ref() {
4643            validate_all_finite("survival_beta_log_sigma", v.iter().copied()).map_err(corrupt)?;
4644        }
4645        if let Some(v) = self.survival_noise_projection.as_ref() {
4646            validate_all_finite("survival_noise_projection", v.iter().flatten().copied())
4647                .map_err(corrupt)?;
4648            if self.survival_noise_projection_ridge_alpha.is_none() {
4649                return Err(FittedModelError::MissingField {
4650                    reason:
4651                        "model has survival_noise_projection but is missing survival_noise_projection_ridge_alpha; refit"
4652                            .to_string(),
4653                });
4654            }
4655        }
4656        if let Some(v) = self.survival_noise_center.as_ref() {
4657            validate_all_finite("survival_noise_center", v.iter().copied()).map_err(corrupt)?;
4658        }
4659        if let Some(v) = self.survival_noise_projection_ridge_alpha {
4660            ensure_finite_scalar("survival_noise_projection_ridge_alpha", v).map_err(corrupt)?;
4661            if v < 0.0 {
4662                return Err(FittedModelError::InvalidInput {
4663                    reason: format!(
4664                        "survival_noise_projection_ridge_alpha must be non-negative, got {v}"
4665                    ),
4666                });
4667            }
4668        }
4669        if let Some(v) = self.survival_noise_scale.as_ref() {
4670            validate_all_finite("survival_noise_scale", v.iter().copied()).map_err(corrupt)?;
4671        }
4672        if let Some(v) = self.mixture_link_param_covariance.as_ref() {
4673            validate_all_finite("mixture_link_param_covariance", v.iter().flatten().copied())
4674                .map_err(corrupt)?;
4675        }
4676        if let Some(v) = self.sas_param_covariance.as_ref() {
4677            validate_all_finite("sas_param_covariance", v.iter().flatten().copied())
4678                .map_err(corrupt)?;
4679        }
4680        Ok(())
4681    }
4682}
4683
4684use gam_solve::estimate::{ensure_finite_scalar, validate_all_finite};
4685
4686fn validate_frozen_term_collectionspec(
4687    spec: &TermCollectionSpec,
4688    label: &str,
4689) -> Result<(), FittedModelError> {
4690    spec.validate_frozen(label)
4691        .map_err(|reason| FittedModelError::SchemaMismatch { reason })
4692}
4693
4694impl Deref for FittedModel {
4695    type Target = FittedModelPayload;
4696
4697    fn deref(&self) -> &Self::Target {
4698        self.payload()
4699    }
4700}
4701
4702impl DerefMut for FittedModel {
4703    fn deref_mut(&mut self) -> &mut Self::Target {
4704        self.payload_mut()
4705    }
4706}
4707
4708// ---------------------------------------------------------------------------
4709// Reconstruct library types from saved models
4710// ---------------------------------------------------------------------------
4711
4712pub fn survival_baseline_config_from_model(
4713    model: &FittedModel,
4714) -> Result<SurvivalBaselineConfig, FittedModelError> {
4715    let target = model.survival_baseline_target.as_deref().ok_or_else(|| {
4716        FittedModelError::MissingField {
4717            reason: "saved survival model missing survival_baseline_target; refit".to_string(),
4718        }
4719    })?;
4720    parse_survival_baseline_config(
4721        target,
4722        model.survival_baseline_scale,
4723        model.survival_baseline_shape,
4724        model.survival_baseline_rate,
4725        model.survival_baseline_makeham,
4726    )
4727    .map_err(|reason| FittedModelError::IncompatibleConfig { reason })
4728}
4729
4730pub fn load_survival_time_basis_config_from_model(
4731    model: &FittedModel,
4732) -> Result<SurvivalTimeBasisConfig, FittedModelError> {
4733    match model
4734        .survival_time_basis
4735        .as_deref()
4736        .ok_or_else(|| FittedModelError::MissingField {
4737            reason: "saved survival model missing survival_time_basis".to_string(),
4738        })?
4739        .to_ascii_lowercase()
4740        .as_str()
4741    {
4742        "none" => Ok(SurvivalTimeBasisConfig::None),
4743        "linear" => Ok(SurvivalTimeBasisConfig::Linear),
4744        "bspline" => {
4745            let degree =
4746                model
4747                    .survival_time_degree
4748                    .ok_or_else(|| FittedModelError::MissingField {
4749                        reason: "saved survival bspline model missing survival_time_degree"
4750                            .to_string(),
4751                    })?;
4752            let knots = model.survival_time_knots.clone().ok_or_else(|| {
4753                FittedModelError::MissingField {
4754                    reason: "saved survival bspline model missing survival_time_knots".to_string(),
4755                }
4756            })?;
4757            let smooth_lambda = model.survival_time_smooth_lambda.unwrap_or(1e-2);
4758            if degree < 1 || knots.is_empty() {
4759                return Err(FittedModelError::SchemaMismatch {
4760                    reason: "saved survival bspline time basis metadata is invalid".to_string(),
4761                });
4762            }
4763            Ok(SurvivalTimeBasisConfig::BSpline {
4764                degree,
4765                knots: Array1::from_vec(knots),
4766                smooth_lambda,
4767            })
4768        }
4769        "ispline" => {
4770            let degree =
4771                model
4772                    .survival_time_degree
4773                    .ok_or_else(|| FittedModelError::MissingField {
4774                        reason: "saved survival ispline model missing survival_time_degree"
4775                            .to_string(),
4776                    })?;
4777            let knots = model.survival_time_knots.clone().ok_or_else(|| {
4778                FittedModelError::MissingField {
4779                    reason: "saved survival ispline model missing survival_time_knots".to_string(),
4780                }
4781            })?;
4782            let keep_cols = model.survival_time_keep_cols.clone().ok_or_else(|| {
4783                FittedModelError::MissingField {
4784                    reason: "saved survival ispline model missing survival_time_keep_cols"
4785                        .to_string(),
4786                }
4787            })?;
4788            let smooth_lambda = model.survival_time_smooth_lambda.unwrap_or(1e-2);
4789            if degree < 1 || knots.is_empty() || keep_cols.is_empty() {
4790                return Err(FittedModelError::SchemaMismatch {
4791                    reason: "saved survival ispline time basis metadata is invalid".to_string(),
4792                });
4793            }
4794            Ok(SurvivalTimeBasisConfig::ISpline {
4795                degree,
4796                knots: Array1::from_vec(knots),
4797                keep_cols,
4798                smooth_lambda,
4799            })
4800        }
4801        other => Err(FittedModelError::IncompatibleConfig {
4802            reason: format!("unsupported saved survival_time_basis '{other}'"),
4803        }),
4804    }
4805}
4806
4807#[cfg(test)]
4808mod tests {
4809    use super::*;
4810    use crate::cubic_cell_kernel::ANCHORED_DEVIATION_KERNEL;
4811    use crate::survival::lognormal_kernel::FrailtySpec;
4812    use gam_data::SchemaColumn;
4813    use gam_problem::types::{LikelihoodScaleMetadata, LogLikelihoodNormalization};
4814    use gam_solve::estimate::{FitArtifacts, FittedBlock, FittedLinkState};
4815    use gam_solve::pirls::PirlsStatus;
4816    use ndarray::{Array1, Array2, array};
4817
4818    fn empty_termspec() -> TermCollectionSpec {
4819        TermCollectionSpec {
4820            linear_terms: vec![],
4821            random_effect_terms: vec![],
4822            smooth_terms: vec![],
4823        }
4824    }
4825
4826    /// #1030/#1034: a scan-bearing payload must round-trip through JSON +
4827    /// `validate_for_persistence` and replay the training Gaussian bridge
4828    /// bit-for-bit; structural corruption must fail loudly at validation.
4829    #[test]
4830    fn spline_scan_payload_round_trips_and_validates() {
4831        let x: Vec<f64> = (0..40).map(|i| i as f64 / 39.0).collect();
4832        let y: Vec<f64> = x.iter().map(|&v| (4.0 * v).sin() + 0.1 * v).collect();
4833        let w = vec![1.0_f64; x.len()];
4834        let fit = gam_solve::spline_scan::fit_spline_scan(&x, &y, &w, 2).expect("scan fit");
4835        let make_payload = || {
4836            crate::inference::model_payload_builders::assemble_spline_scan_payload(
4837                "y ~ s(x)".to_string(),
4838                "x".to_string(),
4839                &fit,
4840                DataSchema {
4841                    columns: vec![
4842                        SchemaColumn {
4843                            name: "y".to_string(),
4844                            kind: ColumnKindTag::Continuous,
4845                            levels: vec![],
4846                        },
4847                        SchemaColumn {
4848                            name: "x".to_string(),
4849                            kind: ColumnKindTag::Continuous,
4850                            levels: vec![],
4851                        },
4852                    ],
4853                },
4854                vec!["x".to_string()],
4855                vec![(0.0, 1.0)],
4856            )
4857        };
4858        // The on-disk form is the FittedModel tagged enum; validation and the
4859        // scan accessor live on FittedModel (Deref only goes Model -> Payload).
4860        let model = FittedModel::from_payload(make_payload());
4861        model
4862            .validate_for_persistence()
4863            .expect("scan model validates");
4864        model
4865            .validate_numeric_finiteness()
4866            .expect("scan model is finite");
4867
4868        let json = serde_json::to_string(&model).expect("serialize model");
4869        let restored: FittedModel = serde_json::from_str(&json).expect("parse model");
4870        restored
4871            .validate_for_persistence()
4872            .expect("restored scan model validates");
4873        let (column, replay) = restored
4874            .saved_spline_scan()
4875            .expect("restore scan fit")
4876            .expect("payload carries the scan representation");
4877        assert_eq!(column, "x");
4878        for &xq in &[-0.1, 0.0, 0.31, 0.5, 0.77, 1.0, 1.4] {
4879            let (m0, v0) = fit.predict(xq).expect("predict original");
4880            let (m1, v1) = replay.predict(xq).expect("predict replayed");
4881            assert_eq!(m0.to_bits(), m1.to_bits(), "mean drift at x={xq}");
4882            assert_eq!(v0.to_bits(), v1.to_bits(), "variance drift at x={xq}");
4883        }
4884
4885        // A dense model without the scan channel still requires fit_result.
4886        let mut dense = make_payload();
4887        dense.spline_scan = None;
4888        let err = FittedModel::from_payload(dense)
4889            .validate_for_persistence()
4890            .expect_err("dense payload without fit_result must be rejected");
4891        assert!(err.to_string().contains("fit_result"));
4892
4893        // Structural corruption fails at validation, not inside predict.
4894        let mut corrupt = make_payload();
4895        corrupt
4896            .spline_scan
4897            .as_mut()
4898            .expect("scan channel present")
4899            .state
4900            .knots
4901            .truncate(2);
4902        FittedModel::from_payload(corrupt)
4903            .validate_for_persistence()
4904            .expect_err("corrupt scan state must be rejected");
4905        let mut unnamed = make_payload();
4906        unnamed
4907            .spline_scan
4908            .as_mut()
4909            .expect("scan channel present")
4910            .feature_column
4911            .clear();
4912        FittedModel::from_payload(unnamed)
4913            .validate_for_persistence()
4914            .expect_err("missing feature column must be rejected");
4915    }
4916
4917    fn standard_gaussian_payload() -> FittedModelPayload {
4918        FittedModelPayload::new(
4919            MODEL_PAYLOAD_VERSION,
4920            "y ~ 1".to_string(),
4921            ModelKind::Standard,
4922            FittedFamily::Standard {
4923                likelihood: LikelihoodSpec::gaussian_identity(),
4924                link: Some(StandardLink::Identity),
4925                latent_cloglog_state: None,
4926                mixture_state: None,
4927                sas_state: None,
4928            },
4929            "gaussian".to_string(),
4930        )
4931    }
4932
4933    fn anchored_runtime(basis_dim: usize) -> SavedCompiledFlexBlock {
4934        SavedCompiledFlexBlock {
4935            kernel: ANCHORED_DEVIATION_KERNEL.to_string(),
4936            breakpoints: vec![-1.0, 1.0],
4937            basis_dim,
4938            span_c0: vec![vec![0.0; basis_dim]],
4939            span_c1: vec![vec![0.0; basis_dim]],
4940            span_c2: vec![vec![0.0; basis_dim]],
4941            span_c3: vec![vec![0.0; basis_dim]],
4942            anchor_correction: None,
4943            anchor_components: Vec::new(),
4944        }
4945    }
4946
4947    fn saved_fit(blocks: Vec<FittedBlock>) -> UnifiedFitResult {
4948        let p: usize = blocks.iter().map(|block| block.beta.len()).sum();
4949        UnifiedFitResult::try_from_parts(gam_solve::estimate::UnifiedFitResultParts {
4950            blocks,
4951            log_lambdas: Array1::zeros(0),
4952            lambdas: Array1::zeros(0),
4953            likelihood_family: Some(LikelihoodSpec::binomial_probit()),
4954            likelihood_scale: LikelihoodScaleMetadata::Unspecified,
4955            log_likelihood_normalization: LogLikelihoodNormalization::Full,
4956            log_likelihood: 0.0,
4957            deviance: 0.0,
4958            reml_score: 0.0,
4959            stable_penalty_term: 0.0,
4960            penalized_objective: 0.0,
4961            used_device: false,
4962            outer_iterations: 0,
4963            outer_converged: true,
4964            outer_gradient_norm: None,
4965            standard_deviation: 1.0,
4966            covariance_conditional: Some(Array2::zeros((p, p))),
4967            covariance_corrected: Some(Array2::zeros((p, p))),
4968            inference: None,
4969            fitted_link: FittedLinkState::Standard(None),
4970            geometry: None,
4971            block_states: vec![],
4972            pirls_status: PirlsStatus::Converged,
4973            max_abs_eta: 0.0,
4974            constraint_kkt: None,
4975            artifacts: FitArtifacts {
4976                pirls: None,
4977                null_space_logdet: None,
4978                null_space_dim: None,
4979                survival_link_wiggle_knots: None,
4980                survival_link_wiggle_degree: None,
4981                criterion_certificate: None,
4982                rho_posterior_certificate: None,
4983                rho_posterior_escalation: None,
4984                rho_covariance: None,
4985                joint_log_lambdas: None,
4986                firth_bias_reduction: false,
4987            },
4988            inner_cycles: 0,
4989        })
4990        .expect("test fixture fit must assemble")
4991    }
4992
4993    fn marginal_slope_payload(version: u32, fit: UnifiedFitResult) -> FittedModelPayload {
4994        let mut payload = FittedModelPayload::new(
4995            version,
4996            "y ~ 1".to_string(),
4997            ModelKind::MarginalSlope,
4998            FittedFamily::MarginalSlope {
4999                likelihood: LikelihoodSpec::binomial_probit(),
5000                base_link: InverseLink::Standard(StandardLink::Probit),
5001                frailty: FrailtySpec::None,
5002            },
5003            "bernoulli-marginal-slope".to_string(),
5004        );
5005        payload.fit_result = Some(fit.clone());
5006        payload.unified = Some(fit);
5007        payload.data_schema = Some(DataSchema {
5008            columns: vec![SchemaColumn {
5009                name: "z".to_string(),
5010                kind: ColumnKindTag::Continuous,
5011                levels: vec![],
5012            }],
5013        });
5014        payload.set_training_feature_metadata(vec!["z".to_string()], vec![(0.0, 0.0)]);
5015        payload.resolved_termspec = Some(empty_termspec());
5016        payload.resolved_termspec_logslope = Some(empty_termspec());
5017        payload.formula_logslope = Some("1".to_string());
5018        payload.z_column = Some("z".to_string());
5019        payload.latent_z_normalization = Some(SavedLatentZNormalization { mean: 0.0, sd: 1.0 });
5020        payload.latent_measure = Some(LatentMeasureKind::StandardNormal);
5021        payload.marginal_baseline = Some(0.0);
5022        payload.logslope_baseline = Some(0.0);
5023        payload.link = Some(InverseLink::Standard(StandardLink::Probit));
5024        payload
5025    }
5026
5027    #[test]
5028    fn from_payload_synchronizes_used_device_from_saved_fit() {
5029        let mut fit = saved_fit(vec![
5030            FittedBlock {
5031                beta: Array1::from_vec(vec![0.25]),
5032                role: BlockRole::Mean,
5033                edf: 1.0,
5034                lambdas: Array1::zeros(0),
5035            },
5036            FittedBlock {
5037                beta: Array1::from_vec(vec![0.5]),
5038                role: BlockRole::Scale,
5039                edf: 1.0,
5040                lambdas: Array1::zeros(0),
5041            },
5042        ]);
5043        fit.used_device = true;
5044        let mut payload = marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
5045        payload.used_device = false;
5046
5047        let model = FittedModel::from_payload(payload);
5048
5049        assert!(model.payload().used_device);
5050    }
5051
5052    fn survival_marginal_slope_payload(version: u32, fit: UnifiedFitResult) -> FittedModelPayload {
5053        let mut payload = FittedModelPayload::new(
5054            version,
5055            "Surv(entry, exit, event) ~ 1".to_string(),
5056            ModelKind::Survival,
5057            FittedFamily::Survival {
5058                likelihood: LikelihoodSpec::royston_parmar(),
5059                survival_likelihood: Some("marginal-slope".to_string()),
5060                survival_distribution: Some(ResidualDistribution::Gaussian),
5061                frailty: FrailtySpec::None,
5062            },
5063            "survival".to_string(),
5064        );
5065        payload.fit_result = Some(fit.clone());
5066        payload.unified = Some(fit);
5067        payload.survival_likelihood = Some("marginal-slope".to_string());
5068        payload.survival_distribution = Some(ResidualDistribution::Gaussian);
5069        payload.latent_measure = Some(LatentMeasureKind::StandardNormal);
5070        payload.data_schema = Some(DataSchema {
5071            columns: vec![SchemaColumn {
5072                name: "z".to_string(),
5073                kind: ColumnKindTag::Continuous,
5074                levels: vec![],
5075            }],
5076        });
5077        payload.set_training_feature_metadata(vec!["z".to_string()], vec![(0.0, 0.0)]);
5078        payload.resolved_termspec = Some(empty_termspec());
5079        payload.resolved_termspec_logslope = Some(empty_termspec());
5080        payload.formula_logslope = Some("1".to_string());
5081        payload.z_column = Some("z".to_string());
5082        payload.latent_z_normalization = Some(SavedLatentZNormalization { mean: 0.0, sd: 1.0 });
5083        payload.logslope_baseline = Some(0.0);
5084        payload.link = Some(InverseLink::Standard(StandardLink::Probit));
5085        payload
5086    }
5087
5088    fn gamma_dispersion_location_scale_payload() -> FittedModelPayload {
5089        // A #913 genuine-dispersion location-scale model: Gamma mean family with
5090        // a log-precision `noise_formula` channel. Its likelihood response is
5091        // non-Gaussian and non-Binomial, so the predict-path classifier must
5092        // route it to `DispersionLocationScale`, NOT the binomial threshold-scale
5093        // class (issue #1064).
5094        let mut payload = FittedModelPayload::new(
5095            MODEL_PAYLOAD_VERSION,
5096            "y ~ x".to_string(),
5097            ModelKind::LocationScale,
5098            FittedFamily::LocationScale {
5099                likelihood: LikelihoodSpec::gamma_log(),
5100                base_link: Some(InverseLink::Standard(StandardLink::Log)),
5101            },
5102            "gamma-location-scale".to_string(),
5103        );
5104        payload.data_schema = Some(DataSchema {
5105            columns: vec![
5106                SchemaColumn {
5107                    name: "y".to_string(),
5108                    kind: ColumnKindTag::Continuous,
5109                    levels: vec![],
5110                },
5111                SchemaColumn {
5112                    name: "x".to_string(),
5113                    kind: ColumnKindTag::Continuous,
5114                    levels: vec![],
5115                },
5116            ],
5117        });
5118        payload.set_training_feature_metadata(vec!["x".to_string()], vec![(-1.0, 1.0)]);
5119        payload.resolved_termspec = Some(empty_termspec());
5120        payload.resolved_termspec_noise = Some(empty_termspec());
5121        payload.formula_noise = Some("x".to_string());
5122        payload.beta_noise = Some(vec![0.0]);
5123        payload.link = Some(InverseLink::Standard(StandardLink::Log));
5124        payload
5125    }
5126
5127    /// #1064 regression: a dispersion location-scale (#913) payload must be
5128    /// classified as `DispersionLocationScale` at every predict-path entry —
5129    /// both `from_payload` (load) and `predict_model_class` (runtime) — and never
5130    /// fall through to the binomial threshold-scale class. Before the fix the
5131    /// non-Gaussian `else` arm mis-routed every dispersion model to
5132    /// `BinomialLocationScale`, predicting the wrong family/link.
5133    #[test]
5134    fn dispersion_location_scale_payload_is_not_classified_binomial() {
5135        let model = FittedModel::from_payload(gamma_dispersion_location_scale_payload());
5136        assert_eq!(
5137            model.predict_model_class(),
5138            PredictModelClass::DispersionLocationScale,
5139            "Gamma dispersion location-scale must route through the dispersion \
5140             predictor, not the binomial threshold-scale class",
5141        );
5142        assert!(
5143            !matches!(
5144                model.predict_model_class(),
5145                PredictModelClass::BinomialLocationScale
5146            ),
5147            "dispersion location-scale must never be classified as binomial",
5148        );
5149
5150        // Each of the four #913 dispersion mean families classifies the same way.
5151        for likelihood in [
5152            LikelihoodSpec::gamma_log(),
5153            LikelihoodSpec::new(
5154                ResponseFamily::NegativeBinomial {
5155                    theta: 1.0,
5156                    theta_fixed: false,
5157                },
5158                InverseLink::Standard(StandardLink::Log),
5159            ),
5160            LikelihoodSpec::new(
5161                ResponseFamily::Beta { phi: 1.0 },
5162                InverseLink::Standard(StandardLink::Logit),
5163            ),
5164            LikelihoodSpec::new(
5165                ResponseFamily::Tweedie { p: 1.5 },
5166                InverseLink::Standard(StandardLink::Log),
5167            ),
5168        ] {
5169            let mut payload = gamma_dispersion_location_scale_payload();
5170            payload.family_state = FittedFamily::LocationScale {
5171                base_link: Some(likelihood.link.clone()),
5172                likelihood: likelihood.clone(),
5173            };
5174            let model = FittedModel::from_payload(payload);
5175            assert_eq!(
5176                model.predict_model_class(),
5177                PredictModelClass::DispersionLocationScale,
5178                "dispersion family {:?} mis-classified",
5179                likelihood.response,
5180            );
5181        }
5182    }
5183
5184    #[test]
5185    fn axis_clip_leaves_numeric_random_effect_group_axis_unclipped() {
5186        let data = array![[100.0], [-100.0]];
5187        let col_map = HashMap::from([("g".to_string(), 0usize)]);
5188
5189        let mut plain_payload = standard_gaussian_payload();
5190        plain_payload.data_schema = Some(DataSchema {
5191            columns: vec![SchemaColumn {
5192                name: "g".to_string(),
5193                kind: ColumnKindTag::Continuous,
5194                levels: vec![],
5195            }],
5196        });
5197        plain_payload.set_training_feature_metadata(vec!["g".to_string()], vec![(0.0, 7.0)]);
5198        plain_payload.resolved_termspec = Some(empty_termspec());
5199        let plain = FittedModel::from_payload(plain_payload.clone());
5200        let clipped = plain
5201            .axis_clip_to_training_ranges(data.view(), &col_map)
5202            .expect("ordinary continuous axis should clip outside the training range");
5203        assert_eq!(clipped.column(0).to_vec(), vec![7.0, 0.0]);
5204
5205        let mut group_payload = plain_payload;
5206        let mut group_spec = empty_termspec();
5207        group_spec
5208            .random_effect_terms
5209            .push(gam_terms::smooth::RandomEffectTermSpec {
5210                name: "g".to_string(),
5211                feature_col: 0,
5212                drop_first_level: false,
5213                penalized: true,
5214                frozen_levels: Some(vec![0.0_f64.to_bits(), 7.0_f64.to_bits()]),
5215                lenient_unseen: true,
5216            });
5217        group_payload.resolved_termspec = Some(group_spec);
5218        let group_model = FittedModel::from_payload(group_payload);
5219
5220        assert_eq!(
5221            group_model.random_effect_group_columns(),
5222            HashSet::from(["g".to_string()])
5223        );
5224
5225        assert_eq!(
5226            group_model.axis_clip_to_training_ranges(data.view(), &col_map),
5227            None,
5228            "numeric group labels must reach RandomEffectOperator as unseen levels, not be clipped to boundary seen levels"
5229        );
5230    }
5231
5232    /// #2102/#2137: a FIXED categorical factor — a bare `y ~ g` or an explicit
5233    /// `y ~ factor(g)` — must reach the strict schema encode and raise a
5234    /// `SchemaMismatch` on an unseen level; it must NOT be silently mapped to
5235    /// the factor's centering point (the across-level average). Only a genuine
5236    /// random effect (`group(g)`/`re(g)`/`s(g, bs="re")`) is eligible for the
5237    /// lenient held-out-group policy, and it must stay lenient. `factor(g)`
5238    /// shared the `group()`/`re()` parse arm and so wrongly inherited the
5239    /// lenient policy (#2137); it is now lowered as the fixed factor it is.
5240    ///
5241    /// This drives the real predict/`check` encode contract: it derives the
5242    /// lenient whitelist from `random_effect_group_columns()` exactly as the
5243    /// predict and `schema_check` FFI paths do, then encodes a frame carrying an
5244    /// out-of-vocabulary level.
5245    #[test]
5246    fn bare_categorical_fixed_factor_unseen_level_rejected_by_predict_encode() {
5247        use csv::StringRecord;
5248        use gam_data::{EncodedDataset, UnseenCategoryPolicy, encode_recordswith_schema};
5249        use gam_runtime::resource::ResourcePolicy;
5250        use gam_terms::inference::formula_dsl::parse_formula;
5251        use gam_terms::term_builder::build_termspec;
5252
5253        // Training frame: response `y` + categorical `g` with levels {a,b,c}.
5254        let train_schema = DataSchema {
5255            columns: vec![
5256                SchemaColumn {
5257                    name: "y".to_string(),
5258                    kind: ColumnKindTag::Continuous,
5259                    levels: vec![],
5260                },
5261                SchemaColumn {
5262                    name: "g".to_string(),
5263                    kind: ColumnKindTag::Categorical,
5264                    levels: vec!["a".to_string(), "b".to_string(), "c".to_string()],
5265                },
5266            ],
5267        };
5268        let train = EncodedDataset {
5269            headers: vec!["y".to_string(), "g".to_string()],
5270            values: Array2::from_shape_vec(
5271                (6, 2),
5272                vec![0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 0.0, 0.0, 1.0, 1.0, 2.0, 2.0],
5273            )
5274            .expect("rectangular training frame"),
5275            schema: train_schema.clone(),
5276            column_kinds: vec![ColumnKindTag::Continuous, ColumnKindTag::Categorical],
5277        };
5278        let build_col_map = train.column_map();
5279
5280        let model_for = |formula: &str| -> FittedModel {
5281            let parsed = parse_formula(formula).expect("formula parses");
5282            let mut notes = Vec::new();
5283            let spec = build_termspec(
5284                &parsed.terms,
5285                &train,
5286                &build_col_map,
5287                &mut notes,
5288                &ResourcePolicy::default_library(),
5289            )
5290            .unwrap_or_else(|err| panic!("`{formula}` must build a term spec, got: {err:?}"));
5291            let mut payload = standard_gaussian_payload();
5292            payload.data_schema = Some(train_schema.clone());
5293            payload.set_training_feature_metadata(
5294                vec!["y".to_string(), "g".to_string()],
5295                vec![(0.0, 2.0), (0.0, 2.0)],
5296            );
5297            payload.resolved_termspec = Some(spec);
5298            FittedModel::from_payload(payload)
5299        };
5300
5301        // Predict frame with an out-of-vocabulary level for `g`.
5302        let g_schema = DataSchema {
5303            columns: vec![SchemaColumn {
5304                name: "g".to_string(),
5305                kind: ColumnKindTag::Categorical,
5306                levels: vec!["a".to_string(), "b".to_string(), "c".to_string()],
5307            }],
5308        };
5309        let encode_level = |model: &FittedModel, level: &str| -> Result<EncodedDataset, String> {
5310            let policy = UnseenCategoryPolicy::encode_unknown_for_columns(
5311                model.random_effect_group_columns(),
5312            );
5313            encode_recordswith_schema(
5314                vec!["g".to_string()],
5315                vec![StringRecord::from(vec![level])],
5316                &g_schema,
5317                policy,
5318            )
5319        };
5320
5321        // Fixed factor: `g` is NOT a lenient group column, so the strict encode
5322        // must reject the unseen level while still accepting a seen one.
5323        let bare = model_for("y ~ g");
5324        assert!(
5325            !bare.random_effect_group_columns().contains("g"),
5326            "bare `+ g` is a fixed parametric factor; it must NOT be whitelisted for lenient \
5327             unseen-level encoding (#2102)"
5328        );
5329        encode_level(&bare, "a").expect("a seen level must still encode for the fixed factor");
5330        let err = encode_level(&bare, "TYPO")
5331            .expect_err("an unseen fixed-factor level must raise a schema mismatch (#2102)");
5332        assert!(
5333            err.contains("unseen level"),
5334            "expected an unseen-level schema mismatch naming the level, got: {err}"
5335        );
5336
5337        // Explicit `factor(g)` is a FIXED categorical factor (#2137): it must
5338        // behave exactly like the bare `+ g` above — strict on unseen levels,
5339        // NOT whitelisted for lenient encoding — even though it shares the
5340        // penalized-categorical materialization with `group(g)`.
5341        let factor = model_for("y ~ factor(g)");
5342        assert!(
5343            !factor.random_effect_group_columns().contains("g"),
5344            "factor(g) is a FIXED categorical factor; it must NOT be whitelisted for lenient \
5345             unseen-level encoding (#2137)"
5346        );
5347        encode_level(&factor, "a").expect("a seen level must still encode for factor(g)");
5348        let factor_err = encode_level(&factor, "TYPO")
5349            .expect_err("an unseen factor(g) level must raise a schema mismatch (#2137)");
5350        assert!(
5351            factor_err.contains("unseen level"),
5352            "expected an unseen-level schema mismatch naming the level, got: {factor_err}"
5353        );
5354
5355        // Genuine random effects stay lenient (held-out group → population mean):
5356        // group(g), its re(g) alias, and the mgcv s(g, bs="re") spelling.
5357        for formula in ["y ~ group(g)", "y ~ re(g)", "y ~ s(g, bs=\"re\")"] {
5358            let grouped = model_for(formula);
5359            assert!(
5360                grouped.random_effect_group_columns().contains("g"),
5361                "`{formula}` is a random effect; it must remain lenient on unseen levels \
5362                 (held-out-group policy)"
5363            );
5364            encode_level(&grouped, "TYPO").unwrap_or_else(|err| {
5365                panic!("`{formula}` must tolerate an unseen level, got: {err}")
5366            });
5367        }
5368    }
5369
5370    #[test]
5371    fn validate_for_persistence_rejects_marginal_slope_score_warp_basis_mismatch() {
5372        let fit = saved_fit(vec![
5373            FittedBlock {
5374                beta: array![0.1],
5375                role: BlockRole::Mean,
5376                edf: 1.0,
5377                lambdas: Array1::zeros(0),
5378            },
5379            FittedBlock {
5380                beta: array![0.2],
5381                role: BlockRole::Scale,
5382                edf: 1.0,
5383                lambdas: Array1::zeros(0),
5384            },
5385            FittedBlock {
5386                beta: array![0.3],
5387                role: BlockRole::Mean,
5388                edf: 1.0,
5389                lambdas: Array1::zeros(0),
5390            },
5391        ]);
5392        let mut payload = marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
5393        payload.score_warp_runtime = Some(anchored_runtime(2));
5394
5395        let err = FittedModel::from_payload(payload)
5396            .validate_for_persistence()
5397            .expect_err("marginal-slope score-warp basis mismatch should fail validation");
5398        assert!(err.to_string().contains("score-warp coefficient mismatch"));
5399    }
5400
5401    #[test]
5402    fn saved_prediction_runtime_rejects_survival_marginal_slope_link_basis_mismatch() {
5403        let fit = saved_fit(vec![
5404            FittedBlock {
5405                beta: array![0.1],
5406                role: BlockRole::Time,
5407                edf: 1.0,
5408                lambdas: Array1::zeros(0),
5409            },
5410            FittedBlock {
5411                beta: array![0.2],
5412                role: BlockRole::Mean,
5413                edf: 1.0,
5414                lambdas: Array1::zeros(0),
5415            },
5416            FittedBlock {
5417                beta: array![0.3],
5418                role: BlockRole::Scale,
5419                edf: 1.0,
5420                lambdas: Array1::zeros(0),
5421            },
5422            FittedBlock {
5423                beta: array![0.4],
5424                role: BlockRole::LinkWiggle,
5425                edf: 1.0,
5426                lambdas: Array1::zeros(0),
5427            },
5428        ]);
5429        let mut payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
5430        payload.link_deviation_runtime = Some(anchored_runtime(2));
5431
5432        let err = FittedModel::from_payload(payload)
5433            .saved_prediction_runtime()
5434            .expect_err(
5435                "survival marginal-slope link basis mismatch should fail runtime validation",
5436            );
5437        assert!(
5438            err.to_string()
5439                .contains("link-deviation coefficient mismatch")
5440        );
5441    }
5442
5443    #[test]
5444    fn apply_survival_time_basis_writes_all_required_fields() {
5445        use crate::survival::construction::SavedSurvivalTimeBasis;
5446
5447        let fit = saved_fit(vec![
5448            FittedBlock {
5449                beta: array![0.1],
5450                role: BlockRole::Time,
5451                edf: 1.0,
5452                lambdas: Array1::zeros(0),
5453            },
5454            FittedBlock {
5455                beta: array![0.2],
5456                role: BlockRole::Mean,
5457                edf: 1.0,
5458                lambdas: Array1::zeros(0),
5459            },
5460            FittedBlock {
5461                beta: array![0.3],
5462                role: BlockRole::Scale,
5463                edf: 1.0,
5464                lambdas: Array1::zeros(0),
5465            },
5466        ]);
5467        let mut payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
5468
5469        // Snapshot writes must match every persisted survival_time_* field —
5470        // forgetting one is exactly the marginal-slope save
5471        // regression. Routing through `apply_survival_time_basis` is the
5472        // structural contract that prevents that recurrence.
5473        let snapshot = SavedSurvivalTimeBasis {
5474            basisname: "royston-parmar".to_string(),
5475            degree: Some(3),
5476            knots: Some(vec![0.0, 1.0, 2.0]),
5477            keep_cols: Some(vec![0, 2]),
5478            smooth_lambda: Some(0.5),
5479            anchor: 0.25,
5480        };
5481        payload.apply_survival_time_basis(&snapshot);
5482
5483        assert_eq!(
5484            payload.survival_time_basis.as_deref(),
5485            Some("royston-parmar")
5486        );
5487        assert_eq!(payload.survival_time_degree, Some(3));
5488        assert_eq!(payload.survival_time_knots, Some(vec![0.0, 1.0, 2.0]));
5489        assert_eq!(payload.survival_time_keep_cols, Some(vec![0, 2]));
5490        assert_eq!(payload.survival_time_smooth_lambda, Some(0.5));
5491        assert_eq!(payload.survival_time_anchor, Some(0.25));
5492    }
5493
5494    #[test]
5495    fn validate_for_persistence_rejects_survival_without_time_anchor_metadata() {
5496        let fit = saved_fit(vec![
5497            FittedBlock {
5498                beta: array![0.1],
5499                role: BlockRole::Time,
5500                edf: 1.0,
5501                lambdas: Array1::zeros(0),
5502            },
5503            FittedBlock {
5504                beta: array![0.2],
5505                role: BlockRole::Mean,
5506                edf: 1.0,
5507                lambdas: Array1::zeros(0),
5508            },
5509            FittedBlock {
5510                beta: array![0.3],
5511                role: BlockRole::Scale,
5512                edf: 1.0,
5513                lambdas: Array1::zeros(0),
5514            },
5515        ]);
5516        let mut payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
5517        // Pass the time_basis presence check but deliberately omit the
5518        // anchor — this is exactly the partial-write shape that the CLI's
5519        // marginal-slope+time-wiggle save path had before the structural
5520        // refactor (main.rs previously set basis/degree/knots/keep_cols/
5521        // smooth_lambda but forgot the anchor).
5522        payload.survival_time_basis = Some("ispline".to_string());
5523
5524        let err = FittedModel::from_payload(payload)
5525            .validate_for_persistence()
5526            .expect_err("survival model without time-anchor metadata should fail validation");
5527        assert!(err.to_string().contains("missing survival_time_anchor"));
5528    }
5529
5530    #[test]
5531    fn validate_for_persistence_rejects_survival_without_time_basis_metadata() {
5532        let fit = saved_fit(vec![
5533            FittedBlock {
5534                beta: array![0.1],
5535                role: BlockRole::Time,
5536                edf: 1.0,
5537                lambdas: Array1::zeros(0),
5538            },
5539            FittedBlock {
5540                beta: array![0.2],
5541                role: BlockRole::Mean,
5542                edf: 1.0,
5543                lambdas: Array1::zeros(0),
5544            },
5545            FittedBlock {
5546                beta: array![0.3],
5547                role: BlockRole::Scale,
5548                edf: 1.0,
5549                lambdas: Array1::zeros(0),
5550            },
5551        ]);
5552        let payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
5553
5554        let err = FittedModel::from_payload(payload)
5555            .validate_for_persistence()
5556            .expect_err("survival model without time-basis metadata should fail validation");
5557        assert!(err.to_string().contains("missing survival_time_basis"));
5558    }
5559
5560    #[test]
5561    fn saved_prediction_runtime_rejects_stale_payload_version() {
5562        let fit = saved_fit(vec![
5563            FittedBlock {
5564                beta: array![0.1],
5565                role: BlockRole::Mean,
5566                edf: 1.0,
5567                lambdas: Array1::zeros(0),
5568            },
5569            FittedBlock {
5570                beta: array![0.2],
5571                role: BlockRole::Scale,
5572                edf: 1.0,
5573                lambdas: Array1::zeros(0),
5574            },
5575        ]);
5576        let payload = marginal_slope_payload(MODEL_PAYLOAD_VERSION - 1, fit);
5577
5578        let err = FittedModel::from_payload(payload)
5579            .saved_prediction_runtime()
5580            .expect_err("stale payload version should fail before runtime assembly");
5581        assert!(err.to_string().contains("payload schema mismatch"));
5582    }
5583}