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};
24pub 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
38pub const MODEL_PAYLOAD_VERSION: u32 = 8;
58
59pub type GroupMetadata = BTreeMap<String, JsonValue>;
66
67#[derive(Clone, Debug, Serialize, Deserialize)]
70pub struct SavedSplineScan {
71 pub feature_column: String,
73 pub state: gam_solve::spline_scan::SplineScanState,
74}
75
76#[derive(Clone, Debug, Serialize, Deserialize)]
81pub struct SavedResidualCascade {
82 pub feature_columns: Vec<String>,
84 pub state: gam_solve::residual_cascade::ResidualCascadeState,
85}
86
87#[derive(Clone, Debug, PartialEq, Eq)]
97pub enum FittedModelError {
98 SchemaMismatch { reason: String },
102 PayloadCorrupt { reason: String },
106 MissingField { reason: String },
109 IncompatibleConfig { reason: String },
113 InvalidInput { reason: String },
116}
117
118impl_reason_error_boilerplate! {
119 FittedModelError {
120 SchemaMismatch,
121 PayloadCorrupt,
122 MissingField,
123 IncompatibleConfig,
124 InvalidInput,
125 }
126}
127
128impl 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 #[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 #[serde(default)]
253 pub unified: Option<UnifiedFitResult>,
254 #[serde(default)]
262 pub spline_scan: Option<SavedSplineScan>,
263 #[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 #[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 #[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 #[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 #[serde(default)]
354 pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
355 #[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 #[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 #[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 pub training_table_kind: String,
447 #[serde(default)]
455 pub training_feature_ranges: Option<Vec<(f64, f64)>>,
456 #[serde(default, skip_serializing_if = "Option::is_none")]
462 pub group_metadata: Option<GroupMetadata>,
463 #[serde(default, skip_serializing_if = "Vec::is_empty")]
470 pub deployment_extensions: Vec<SavedDeploymentExtension>,
471 #[serde(default)]
473 pub transformation_response_knots: Option<Vec<f64>>,
474 #[serde(default)]
476 pub transformation_response_transform: Option<Vec<Vec<f64>>>,
477 #[serde(default)]
479 pub transformation_response_degree: Option<usize>,
480 #[serde(default)]
482 pub transformation_response_median: Option<f64>,
483 #[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 #[serde(default)]
512 pub gaussian_jackknife_plus:
513 Option<crate::inference::full_conformal::GaussianJackknifePlusStats>,
514 #[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
551pub 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 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 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 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
954pub 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 #[serde(default)]
973 pub anchor_correction: Option<Vec<Vec<f64>>>,
974 #[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 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 pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
1012 pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
1018 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
1048fn 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 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
1277fn 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 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 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 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 self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)
1946 }
1947
1948 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 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 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 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 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 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
2123fn 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 SmoothBasisSpec::BSpline1D { feature_col, .. } => push(*feature_col, out),
2142 SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
2145 for &c in feature_cols {
2146 push(c, out);
2147 }
2148 }
2149 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 SmoothBasisSpec::FactorSmooth { spec } => {
2167 for &c in &spec.continuous_cols {
2168 push(c, out);
2169 }
2170 }
2171 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 SmoothBasisSpec::Sphere { .. }
2185 | SmoothBasisSpec::ConstantCurvature { .. }
2186 | SmoothBasisSpec::Pca { .. } => {}
2187 }
2188}
2189
2190fn 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 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 let periodic_axes = self.training_periodic_axes(training_headers);
2273 let linear_axes = self.training_linear_axes(training_headers.len());
2281 let random_effect_axes = self.training_random_effect_axes(training_headers.len());
2286 let smooth_extrapolation_axes =
2294 self.training_smooth_extrapolation_axes(training_headers.len());
2295 let by_variable_axes = self.training_by_variable_numeric_axes(training_headers.len());
2301 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 if response.is_empty() || response.contains('(') {
2864 return Ok(extras);
2865 }
2866 if self.prediction_required_columns()?.contains(response) {
2869 return Ok(extras);
2870 }
2871 extras.push(response.to_string());
2872 Ok(extras)
2873 }
2874
2875 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 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 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 #[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 #[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 #[inline]
3100 pub fn prediction_uses_posterior_mean(&self) -> bool {
3101 let family = self.likelihood();
3102 let curved_family = match &family.response {
3103 ResponseFamily::Gaussian => false,
3106 ResponseFamily::Poisson
3108 | ResponseFamily::Gamma
3109 | ResponseFamily::Tweedie { .. }
3110 | ResponseFamily::NegativeBinomial { .. } => true,
3111 ResponseFamily::Beta { .. } => true,
3113 ResponseFamily::RoystonParmar => true,
3115 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 const MEASURE_JET_COVERAGE_FLOOR: f64 = 0.05;
3391
3392 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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
4708pub 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 #[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 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 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 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 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 #[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 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 #[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 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 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 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 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 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 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 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}