1use crate::bms::{LatentMeasureKind, LatentZConditionalCalibration, LatentZRankIntCalibration};
2use crate::survival::construction::{
3 SurvivalBaselineConfig, SurvivalTimeBasisConfig, parse_survival_baseline_config,
4};
5use crate::survival::location_scale::{
6 ResidualDistribution, SurvivalCovariateTimeBasis, SurvivalLocationScaleTimeParameterization,
7};
8use crate::survival::lognormal_kernel::{FrailtyScale, FrailtySpec};
9use crate::wiggle::{
10 WigglePenaltyMetadata, canonical_wiggle_function_penalties,
11 monotone_wiggle_basis_with_derivative_order, validate_monotone_wiggle_beta_nonnegative,
12};
13use gam_linalg::faer_ndarray::{FaerCholesky, array2_to_nested_vec};
14use gam_linalg::matrix::DesignMatrix;
15use gam_problem::types::{
16 InverseLink, LatentCLogLogState, LikelihoodSpec, MixtureLinkState, ResponseFamily, SasLinkSpec,
17 SasLinkState, StandardLink,
18};
19use gam_runtime::span::span_index_for_breakpoints;
20use gam_solve::estimate::{BlockRole, FittedLinkState, UnifiedFitResult};
21use gam_solve::mixture_link::{state_from_beta_logisticspec, state_from_sasspec};
22use gam_terms::basis::BasisOptions;
23use gam_terms::inference::formula_dsl::{
24 inverse_link_supports_joint_wiggle, joint_wiggle_unsupported_link_message, parse_formula,
25 parse_surv_interval_response, parse_surv_response, parsed_term_column_names,
26};
27use gam_terms::smooth::{AdaptiveRegularizationDiagnostics, TermCollectionSpec};
28pub use gam_data::{ColumnKindTag, DataSchema, SchemaColumn};
34use ndarray::{Array1, Array2, ArrayView1};
35use serde::{Deserialize, Serialize};
36use serde_json::Value as JsonValue;
37use std::collections::{BTreeMap, HashMap, HashSet};
38use std::fs;
39use std::ops::{Deref, DerefMut};
40use std::path::Path;
41
42pub const MODEL_PAYLOAD_VERSION: u32 = 14;
62
63#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
72#[serde(rename_all = "kebab-case")]
73pub enum TransformationNormalParameterization {
74 DirectAlpha,
77}
78
79#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
93pub struct SavedTransformationNormalGeometry {
94 pub parameterization: TransformationNormalParameterization,
96 pub response_degree: usize,
98 pub response_knot_count: usize,
102 pub shape_coordinate_count: usize,
105 pub cone_carrier_covariate_width: usize,
108 pub cone_carrier_row_count: usize,
111 pub certified_response_support: (f64, f64),
115 pub response_median: f64,
117}
118
119impl SavedTransformationNormalGeometry {
120 pub fn validate(&self, context: &str) -> Result<(), FittedModelError> {
124 match self.parameterization {
125 TransformationNormalParameterization::DirectAlpha => {}
126 }
127 if self.response_degree < 1 {
128 return Err(FittedModelError::PayloadCorrupt {
129 reason: format!(
130 "{context} CTN geometry response_degree must be >= 1, got {}",
131 self.response_degree
132 ),
133 });
134 }
135 if self.shape_coordinate_count == 0 {
136 return Err(FittedModelError::PayloadCorrupt {
137 reason: format!("{context} CTN geometry needs at least one shape coordinate"),
138 });
139 }
140 if self.cone_carrier_covariate_width == 0 || self.cone_carrier_row_count == 0 {
141 return Err(FittedModelError::PayloadCorrupt {
142 reason: format!(
143 "{context} CTN geometry cone carrier must be non-empty: {} rows x {} covariate columns",
144 self.cone_carrier_row_count, self.cone_carrier_covariate_width
145 ),
146 });
147 }
148 let (lo, hi) = self.certified_response_support;
149 if !(lo.is_finite() && hi.is_finite() && lo < hi) {
150 return Err(FittedModelError::PayloadCorrupt {
151 reason: format!(
152 "{context} CTN geometry certified response support must be finite and ordered lo < hi, got [{lo}, {hi}]"
153 ),
154 });
155 }
156 if !self.response_median.is_finite() {
157 return Err(FittedModelError::PayloadCorrupt {
158 reason: format!(
159 "{context} CTN geometry response_median must be finite, got {}",
160 self.response_median
161 ),
162 });
163 }
164 Ok::<(), _>(())
165 }
166}
167
168#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
172pub struct SavedSurvivalLocationScaleStructure {
173 pub time_parameterization: SurvivalLocationScaleTimeParameterization,
174 pub threshold_time_basis: Option<SurvivalCovariateTimeBasis>,
175 pub log_sigma_time_basis: Option<SurvivalCovariateTimeBasis>,
176}
177
178pub type GroupMetadata = BTreeMap<String, JsonValue>;
185
186#[derive(Clone, Debug, Serialize, Deserialize)]
189pub struct SavedSplineScan {
190 pub feature_column: String,
192 pub state: gam_solve::spline_scan::SplineScanState,
193}
194
195#[derive(Clone, Debug, Serialize, Deserialize)]
200pub struct SavedResidualCascade {
201 pub feature_columns: Vec<String>,
203 pub state: gam_solve::residual_cascade::ResidualCascadeState,
204}
205
206#[derive(Clone, Debug, PartialEq, Eq)]
216pub enum FittedModelError {
217 SchemaMismatch { reason: String },
221 PayloadCorrupt { reason: String },
225 MissingField { reason: String },
228 IncompatibleConfig { reason: String },
232 InvalidInput { reason: String },
235}
236
237impl_reason_error_boilerplate! {
238 FittedModelError {
239 SchemaMismatch,
240 PayloadCorrupt,
241 MissingField,
242 IncompatibleConfig,
243 InvalidInput,
244 }
245}
246
247impl From<FittedModelError> for gam_solve::model_types::EstimationError {
252 fn from(err: FittedModelError) -> Self {
253 gam_solve::model_types::EstimationError::InvalidInput(err.to_string())
254 }
255}
256
257impl From<FittedModelError> for crate::survival::predict::SurvivalPredictError {
258 fn from(err: FittedModelError) -> Self {
259 crate::survival::predict::SurvivalPredictError::ModelPayload {
260 context: "saved-model survival prediction payload",
261 source: err,
262 }
263 }
264}
265
266#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
267pub struct SavedLatentZNormalization {
268 pub mean: f64,
269 pub sd: f64,
270}
271
272impl SavedLatentZNormalization {
273 pub fn validate(&self, context: &str) -> Result<(), FittedModelError> {
274 if !self.mean.is_finite() {
275 return Err(FittedModelError::PayloadCorrupt {
276 reason: format!("{context} latent z mean must be finite"),
277 });
278 }
279 if !(self.sd.is_finite() && self.sd > 1e-12) {
280 return Err(FittedModelError::PayloadCorrupt {
281 reason: format!(
282 "{context} latent z sd must be finite and > 1e-12; got {}",
283 self.sd
284 ),
285 });
286 }
287 Ok::<(), _>(())
288 }
289
290 pub fn apply(&self, z: &Array1<f64>, context: &str) -> Result<Array1<f64>, FittedModelError> {
291 self.validate(context)?;
292 if z.iter().any(|value| !value.is_finite()) {
293 return Err(FittedModelError::PayloadCorrupt {
294 reason: format!("{context} requires finite z values"),
295 });
296 }
297 Ok(z.mapv(|zi| (zi - self.mean) / self.sd))
298 }
299}
300
301pub const TRANSFORMATION_SCORE_PIT_CLIP_EPS: f64 = 1.0e-12;
302
303#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
304#[serde(rename_all = "kebab-case")]
305#[derive(Default)]
306pub enum TransformationScoreKind {
307 #[default]
308 FiniteSupportPit,
309}
310
311#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
312pub struct TransformationScoreCalibration {
313 #[serde(default)]
314 pub score_kind: TransformationScoreKind,
315 #[serde(default = "default_transformation_score_pit_clip_eps")]
316 pub clip_eps: f64,
317}
318
319const fn default_transformation_score_pit_clip_eps() -> f64 {
320 TRANSFORMATION_SCORE_PIT_CLIP_EPS
321}
322
323impl TransformationScoreCalibration {
324 pub fn finite_support_pit() -> Self {
325 Self {
326 score_kind: TransformationScoreKind::FiniteSupportPit,
327 clip_eps: TRANSFORMATION_SCORE_PIT_CLIP_EPS,
328 }
329 }
330
331 pub fn validate(&self, context: &str) -> Result<(), FittedModelError> {
332 if self.score_kind != TransformationScoreKind::FiniteSupportPit {
333 return Err(FittedModelError::IncompatibleConfig {
334 reason: format!("{context} supports only finite-support CTN PIT score semantics"),
335 });
336 }
337 if !(self.clip_eps.is_finite() && self.clip_eps > 0.0 && self.clip_eps < 0.5) {
338 return Err(FittedModelError::IncompatibleConfig {
339 reason: format!(
340 "{context} requires PIT clip_eps in (0, 0.5), got {}",
341 self.clip_eps
342 ),
343 });
344 }
345 Ok::<(), _>(())
346 }
347}
348
349#[derive(Clone, Serialize, Deserialize)]
350pub struct FittedModelPayload {
351 pub version: u32,
352 pub formula: String,
353 pub model_kind: ModelKind,
354 pub family_state: FittedFamily,
355 pub family: String,
356 pub estimator: FittedEstimator,
363 #[serde(default)]
372 pub inference_notes: Vec<String>,
373 #[serde(default)]
374 pub used_device: bool,
375 #[serde(default)]
376 pub fit_result: Option<UnifiedFitResult>,
377 #[serde(default)]
379 pub unified: Option<UnifiedFitResult>,
380 #[serde(default)]
388 pub spline_scan: Option<SavedSplineScan>,
389 #[serde(default)]
395 pub residual_cascade: Option<SavedResidualCascade>,
396 #[serde(default)]
397 pub data_schema: Option<DataSchema>,
398 pub link: Option<InverseLink>,
399 #[serde(default)]
400 pub mixture_link_param_covariance: Option<Vec<Vec<f64>>>,
401 #[serde(default)]
402 pub sas_param_covariance: Option<Vec<Vec<f64>>>,
403 #[serde(default)]
404 pub formula_noise: Option<String>,
405 #[serde(default)]
406 pub formula_logslope: Option<String>,
407 #[serde(default)]
408 pub formula_logslopes: Option<Vec<String>>,
409 #[serde(default)]
410 pub offset_column: Option<String>,
411 #[serde(default)]
412 pub noise_offset_column: Option<String>,
413 #[serde(default)]
418 pub weight_column: Option<String>,
419 #[serde(default)]
420 pub beta_noise: Option<Vec<f64>>,
421 #[serde(default)]
422 pub noise_projection: Option<Vec<Vec<f64>>>,
423 #[serde(default)]
424 pub noise_center: Option<Vec<f64>>,
425 #[serde(default)]
426 pub noise_scale: Option<Vec<f64>>,
427 #[serde(default)]
428 pub noise_non_intercept_start: Option<usize>,
429 #[serde(default)]
433 pub noise_projection_ridge_alpha: Option<f64>,
434 #[serde(default)]
435 pub gaussian_response_scale: Option<f64>,
436 #[serde(default)]
437 pub linkwiggle_knots: Option<Vec<f64>>,
438 #[serde(default)]
439 pub linkwiggle_degree: Option<usize>,
440 #[serde(default)]
445 pub linkwiggle_penalty_metadata: Option<WigglePenaltyMetadata>,
446 #[serde(default)]
447 pub beta_link_wiggle: Option<Vec<f64>>,
448 #[serde(default)]
456 pub link_wiggle_index_shift: Option<Vec<f64>>,
457 #[serde(default)]
458 pub baseline_timewiggle_knots: Option<Vec<f64>>,
459 #[serde(default)]
460 pub baseline_timewiggle_degree: Option<usize>,
461 #[serde(default)]
462 pub baseline_timewiggle_penalty_orders: Option<Vec<usize>>,
463 #[serde(default)]
464 pub baseline_timewiggle_double_penalty: Option<bool>,
465 #[serde(default)]
466 pub beta_baseline_timewiggle: Option<Vec<f64>>,
467 #[serde(default)]
468 pub beta_baseline_timewiggle_by_cause: Option<Vec<Vec<f64>>>,
469 #[serde(default)]
470 pub z_column: Option<String>,
471 #[serde(default)]
472 pub z_columns: Option<Vec<String>>,
473 #[serde(default)]
474 pub latent_z_normalization: Option<SavedLatentZNormalization>,
475 #[serde(default)]
476 pub latent_score_contract: Option<SavedLatentScoreContract>,
477 #[serde(default)]
478 pub latent_measure: Option<LatentMeasureKind>,
479 #[serde(default)]
486 pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
487 #[serde(default)]
496 pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
497 #[serde(default)]
498 pub marginal_baseline: Option<f64>,
499 #[serde(default)]
500 pub logslope_baseline: Option<f64>,
501 #[serde(default)]
502 pub logslope_baselines: Option<Vec<f64>>,
503 #[serde(default)]
504 pub score_warp_runtime: Option<SavedCompiledFlexBlock>,
505 #[serde(default)]
506 pub link_deviation_runtime: Option<SavedCompiledFlexBlock>,
507 #[serde(default)]
512 pub influence_absorber_width: Option<usize>,
513 pub influence_absorber_design: Option<Vec<Vec<f64>>>,
518 pub survival_marginal_slope_score_covariance: Option<Vec<Vec<f64>>>,
521 #[serde(default)]
522 pub survival_entry: Option<String>,
523 #[serde(default)]
524 pub survival_exit: Option<String>,
525 #[serde(default)]
526 pub survival_event: Option<String>,
527 #[serde(default)]
528 pub survivalspec: Option<String>,
529 #[serde(default)]
530 pub survival_cause_count: Option<usize>,
531 #[serde(default)]
532 pub survival_endpoint_names: Option<Vec<String>>,
533 #[serde(default)]
534 pub survival_baseline_target: Option<String>,
535 #[serde(default)]
536 pub survival_baseline_scale: Option<f64>,
537 #[serde(default)]
538 pub survival_baseline_shape: Option<f64>,
539 #[serde(default)]
540 pub survival_baseline_rate: Option<f64>,
541 #[serde(default)]
542 pub survival_baseline_makeham: Option<f64>,
543 #[serde(default)]
544 pub survival_time_basis: Option<String>,
545 #[serde(default)]
546 pub survival_time_degree: Option<usize>,
547 #[serde(default)]
548 pub survival_time_knots: Option<Vec<f64>>,
549 #[serde(default)]
550 pub survival_time_keep_cols: Option<Vec<usize>>,
551 #[serde(default)]
552 pub survival_time_smooth_lambda: Option<f64>,
553 #[serde(default)]
554 pub survival_time_anchor: Option<f64>,
555 #[serde(default)]
556 pub survivalridge_lambda: Option<f64>,
557 #[serde(default)]
558 pub survival_likelihood: Option<String>,
559 pub survival_location_scale_structure: Option<SavedSurvivalLocationScaleStructure>,
563 #[serde(default)]
564 pub survival_beta_time: Option<Vec<f64>>,
565 #[serde(default)]
566 pub survival_beta_threshold: Option<Vec<f64>>,
567 #[serde(default)]
568 pub survival_beta_log_sigma: Option<Vec<f64>>,
569 #[serde(default)]
570 pub survival_distribution: Option<ResidualDistribution>,
571 #[serde(default)]
572 pub training_headers: Option<Vec<String>>,
573 pub training_table_kind: String,
579 #[serde(default)]
587 pub training_feature_ranges: Option<Vec<(f64, f64)>>,
588 #[serde(default, skip_serializing_if = "Option::is_none")]
594 pub group_metadata: Option<GroupMetadata>,
595 #[serde(default, skip_serializing_if = "Vec::is_empty")]
602 pub deployment_extensions: Vec<SavedDeploymentExtension>,
603 #[serde(default)]
605 pub transformation_response_knots: Option<Vec<f64>>,
606 #[serde(default)]
608 pub transformation_response_transform: Option<Vec<Vec<f64>>>,
609 #[serde(default)]
611 pub transformation_response_degree: Option<usize>,
612 #[serde(default)]
614 pub transformation_response_median: Option<f64>,
615 #[serde(default)]
620 pub transformation_geometry: Option<SavedTransformationNormalGeometry>,
621 #[serde(default)]
631 pub transformation_cone_carrier: Option<Vec<f64>>,
632 #[serde(default)]
636 pub transformation_score_calibration: Option<TransformationScoreCalibration>,
637 #[serde(default)]
638 pub resolved_termspec: Option<TermCollectionSpec>,
639 #[serde(default)]
640 pub resolved_termspec_noise: Option<TermCollectionSpec>,
641 #[serde(default)]
642 pub resolved_termspec_logslope: Option<TermCollectionSpec>,
643 #[serde(default)]
644 pub resolved_termspec_logslopes: Option<Vec<TermCollectionSpec>>,
645 #[serde(default)]
646 pub adaptive_regularization_diagnostics: Option<AdaptiveRegularizationDiagnostics>,
647 #[serde(default)]
661 pub gaussian_jackknife_plus:
662 Option<crate::inference::full_conformal::GaussianJackknifePlusStats>,
663 #[serde(default)]
681 pub full_conformal: Option<crate::inference::full_conformal::ExactFullConformalSubstrate>,
682}
683
684#[derive(Clone, Debug, Serialize, Deserialize)]
685pub struct SavedDeploymentExtension {
686 pub name: String,
687 pub kind: String,
688 pub term: String,
689 pub level: JsonValue,
690 pub level_bits: u64,
691 pub coefficient_index: usize,
692 pub coefficient_mean: f64,
693 pub coefficient_variance: f64,
694 #[serde(default, skip_serializing_if = "Option::is_none")]
695 pub metadata: Option<JsonValue>,
696 #[serde(default, skip_serializing_if = "Option::is_none")]
697 pub prior: Option<JsonValue>,
698}
699
700pub fn append_deployment_extension_columns(
717 model: &FittedModelPayload,
718 data: ndarray::ArrayView2<'_, f64>,
719 col_map: &HashMap<String, usize>,
720 training_headers: Option<&Vec<String>>,
721 base_design: Array2<f64>,
722) -> Result<Array2<f64>, FittedModelError> {
723 if model.deployment_extensions.is_empty() {
724 return Ok(base_design);
725 }
726 if base_design.nrows() != data.nrows() {
727 return Err(FittedModelError::SchemaMismatch {
728 reason: format!(
729 "deployment extension design row mismatch: base design has {} rows but data has {}",
730 base_design.nrows(),
731 data.nrows()
732 ),
733 });
734 }
735 let spec = model
736 .resolved_termspec
737 .as_ref()
738 .ok_or_else(|| FittedModelError::MissingField {
739 reason: "deployment extension prediction requires saved resolved_termspec; refit"
740 .to_string(),
741 })?;
742 let n = base_design.nrows();
743 let p_old = base_design.ncols();
744 let mut extensions: Vec<&SavedDeploymentExtension> =
745 model.deployment_extensions.iter().collect();
746 extensions.sort_by_key(|extension| extension.coefficient_index);
747 for (tail_idx, extension) in extensions.iter().enumerate() {
748 let expected = p_old + tail_idx;
749 if extension.coefficient_index != expected {
750 return Err(FittedModelError::SchemaMismatch {
751 reason: format!(
752 "deployment extension '{}' has coefficient index {}, expected append-only index {}",
753 extension.name, extension.coefficient_index, expected
754 ),
755 });
756 }
757 }
758
759 let mut out = Array2::<f64>::zeros((n, p_old + extensions.len()));
760 out.slice_mut(ndarray::s![.., ..p_old]).assign(&base_design);
761 for (tail_idx, extension) in extensions.into_iter().enumerate() {
762 if extension.kind != "random-effect-level" {
763 return Err(FittedModelError::IncompatibleConfig {
764 reason: format!(
765 "unsupported deployment extension kind '{}' for '{}'",
766 extension.kind, extension.name
767 ),
768 });
769 }
770 let term = spec
771 .random_effect_terms
772 .iter()
773 .find(|term| term.name == extension.term)
774 .ok_or_else(|| FittedModelError::MissingField {
775 reason: format!(
776 "deployment extension '{}' references unknown random-effect term '{}'",
777 extension.name, extension.term
778 ),
779 })?;
780 let prediction_col = training_headers
781 .and_then(|headers| headers.get(term.feature_col))
782 .and_then(|name| col_map.get(name))
783 .copied()
784 .unwrap_or(term.feature_col);
785 if prediction_col >= data.ncols() {
786 return Err(FittedModelError::SchemaMismatch {
787 reason: format!(
788 "deployment extension '{}' feature column {} out of bounds for {} prediction columns",
789 extension.name,
790 prediction_col,
791 data.ncols()
792 ),
793 });
794 }
795 let col = p_old + tail_idx;
796 let level_bits = gam_data::canonical_level_bits(f64::from_bits(extension.level_bits));
797 for row in 0..n {
798 if gam_data::canonical_level_bits(data[[row, prediction_col]]) == level_bits {
799 out[[row, col]] = 1.0;
800 }
801 }
802 }
803 Ok(out)
804}
805
806#[derive(Clone, Debug, Serialize, Deserialize)]
807pub struct SavedLatentScoreContract {
808 pub semantics: String,
809 pub source_transform_id: Option<String>,
810 pub normalization_mean: f64,
811 pub normalization_sd: f64,
812 pub clip_eps: Option<f64>,
813 pub conditioning_columns: Vec<String>,
814}
815
816impl FittedModelPayload {
817 pub fn new(
818 version: u32,
819 formula: String,
820 model_kind: ModelKind,
821 family_state: FittedFamily,
822 family: String,
823 ) -> Self {
824 Self {
825 version,
826 formula,
827 model_kind,
828 family_state,
829 family,
830 estimator: FittedEstimator::Likelihood,
831 inference_notes: Vec::new(),
832 used_device: false,
833 fit_result: None,
834 unified: None,
835 spline_scan: None,
836 residual_cascade: None,
837 data_schema: None,
838 link: None,
839 mixture_link_param_covariance: None,
840 sas_param_covariance: None,
841 formula_noise: None,
842 formula_logslope: None,
843 formula_logslopes: None,
844 offset_column: None,
845 noise_offset_column: None,
846 weight_column: None,
847 beta_noise: None,
848 noise_projection: None,
849 noise_center: None,
850 noise_scale: None,
851 noise_non_intercept_start: None,
852 noise_projection_ridge_alpha: None,
853 gaussian_response_scale: None,
854 linkwiggle_knots: None,
855 linkwiggle_degree: None,
856 linkwiggle_penalty_metadata: None,
857 beta_link_wiggle: None,
858 link_wiggle_index_shift: None,
859 baseline_timewiggle_knots: None,
860 baseline_timewiggle_degree: None,
861 baseline_timewiggle_penalty_orders: None,
862 baseline_timewiggle_double_penalty: None,
863 beta_baseline_timewiggle: None,
864 beta_baseline_timewiggle_by_cause: None,
865 z_column: None,
866 z_columns: None,
867 latent_z_normalization: None,
868 latent_score_contract: None,
869 latent_measure: None,
870 latent_z_rank_int_calibration: None,
871 latent_z_conditional_calibration: None,
872 marginal_baseline: None,
873 logslope_baseline: None,
874 logslope_baselines: None,
875 score_warp_runtime: None,
876 link_deviation_runtime: None,
877 influence_absorber_width: None,
878 influence_absorber_design: None,
879 survival_marginal_slope_score_covariance: None,
880 survival_entry: None,
881 survival_exit: None,
882 survival_event: None,
883 survivalspec: None,
884 survival_cause_count: None,
885 survival_endpoint_names: None,
886 survival_baseline_target: None,
887 survival_baseline_scale: None,
888 survival_baseline_shape: None,
889 survival_baseline_rate: None,
890 survival_baseline_makeham: None,
891 survival_time_basis: None,
892 survival_time_degree: None,
893 survival_time_knots: None,
894 survival_time_keep_cols: None,
895 survival_time_smooth_lambda: None,
896 survival_time_anchor: None,
897 survivalridge_lambda: None,
898 survival_likelihood: None,
899 survival_location_scale_structure: None,
900 survival_beta_time: None,
901 survival_beta_threshold: None,
902 survival_beta_log_sigma: None,
903 survival_distribution: None,
904 training_headers: None,
905 training_table_kind: "unknown".to_string(),
906 training_feature_ranges: None,
907 group_metadata: None,
908 deployment_extensions: Vec::new(),
909 transformation_response_knots: None,
910 transformation_response_transform: None,
911 transformation_response_degree: None,
912 transformation_response_median: None,
913 transformation_geometry: None,
914 transformation_cone_carrier: None,
915 transformation_score_calibration: None,
916 resolved_termspec: None,
917 resolved_termspec_noise: None,
918 resolved_termspec_logslope: None,
919 resolved_termspec_logslopes: None,
920 adaptive_regularization_diagnostics: None,
921 gaussian_jackknife_plus: None,
922 full_conformal: None,
923 }
924 }
925
926 pub fn set_training_feature_metadata(
927 &mut self,
928 headers: Vec<String>,
929 feature_ranges: Vec<(f64, f64)>,
930 ) {
931 self.training_headers = Some(headers);
932 self.training_feature_ranges = Some(feature_ranges);
933 }
934
935 fn synchronize_empty_feature_contract(&mut self) {
936 if self.fit_result.is_none() {
937 return;
938 }
939 let Some(schema) = self.data_schema.as_ref() else {
940 return;
941 };
942 if !schema.columns.is_empty() {
943 return;
944 }
945 self.training_headers.get_or_insert_with(Vec::new);
946 self.resolved_termspec
947 .get_or_insert_with(|| TermCollectionSpec {
948 linear_terms: Vec::new(),
949 smooth_terms: Vec::new(),
950 random_effect_terms: Vec::new(),
951 });
952 }
953
954 pub fn apply_survival_time_basis(
962 &mut self,
963 snapshot: &crate::survival::construction::SavedSurvivalTimeBasis,
964 ) {
965 self.survival_time_basis = Some(snapshot.basisname.clone());
966 self.survival_time_degree = snapshot.degree;
967 self.survival_time_knots = snapshot.knots.clone();
968 self.survival_time_keep_cols = snapshot.keep_cols.clone();
969 self.survival_time_smooth_lambda = snapshot.smooth_lambda;
970 self.survival_time_anchor = Some(snapshot.anchor);
971 }
972
973 fn validate_payload_version(&self) -> Result<(), FittedModelError> {
974 if self.version != MODEL_PAYLOAD_VERSION {
975 return Err(FittedModelError::SchemaMismatch {
976 reason: format!(
977 "saved model payload schema mismatch: file has version={}, \
978 this binary expects MODEL_PAYLOAD_VERSION={}. \
979 Refit with the current CLI, or rebuild the reader at the same \
980 version the model was written with.",
981 self.version, MODEL_PAYLOAD_VERSION
982 ),
983 });
984 }
985 Ok(())
986 }
987}
988
989#[derive(Clone, Serialize, Deserialize)]
990#[serde(tag = "model_type", rename_all = "kebab-case")]
991pub enum FittedModel {
992 Standard { payload: FittedModelPayload },
993 LocationScale { payload: FittedModelPayload },
994 MarginalSlope { payload: FittedModelPayload },
995 Survival { payload: FittedModelPayload },
996 TransformationNormal { payload: FittedModelPayload },
997}
998
999#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
1000#[serde(rename_all = "kebab-case")]
1001pub enum ModelKind {
1002 Standard,
1003 LocationScale,
1004 MarginalSlope,
1005 Survival,
1006 TransformationNormal,
1007}
1008
1009#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq)]
1015#[serde(tag = "estimator_kind", rename_all = "kebab-case")]
1016pub enum FittedEstimator {
1017 Likelihood,
1018 Expectile { tau: f64 },
1019}
1020
1021#[derive(Clone, Debug, Serialize, Deserialize)]
1022#[serde(tag = "family_kind", rename_all = "kebab-case")]
1023pub enum FittedFamily {
1024 Standard {
1025 likelihood: LikelihoodSpec,
1026 #[serde(default)]
1027 link: Option<StandardLink>,
1028 #[serde(default)]
1029 latent_cloglog_state: Option<LatentCLogLogState>,
1030 #[serde(default)]
1031 mixture_state: Option<MixtureLinkState>,
1032 #[serde(default)]
1033 sas_state: Option<SasLinkState>,
1034 },
1035 LocationScale {
1036 likelihood: LikelihoodSpec,
1037 #[serde(default)]
1038 base_link: Option<InverseLink>,
1039 },
1040 MarginalSlope {
1041 likelihood: LikelihoodSpec,
1042 base_link: InverseLink,
1043 frailty: FrailtySpec,
1044 },
1045 Survival {
1046 likelihood: LikelihoodSpec,
1047 #[serde(default)]
1048 survival_likelihood: Option<String>,
1049 #[serde(default)]
1050 survival_distribution: Option<ResidualDistribution>,
1051 frailty: FrailtySpec,
1052 },
1053 LatentSurvival {
1054 frailty: FrailtySpec,
1055 },
1056 LatentBinary {
1057 frailty: FrailtySpec,
1058 },
1059 TransformationNormal {
1060 likelihood: LikelihoodSpec,
1061 },
1062}
1063
1064#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1065pub enum PredictModelClass {
1066 Standard,
1067 GaussianLocationScale,
1068 BinomialLocationScale,
1069 DispersionLocationScale,
1074 BernoulliMarginalSlope,
1075 Survival,
1076 TransformationNormal,
1077}
1078
1079impl PredictModelClass {
1080 #[inline]
1081 pub const fn name(self) -> &'static str {
1082 match self {
1083 Self::Standard => "standard",
1084 Self::GaussianLocationScale => "gaussian location-scale",
1085 Self::BinomialLocationScale => "binomial location-scale",
1086 Self::DispersionLocationScale => "dispersion location-scale",
1087 Self::BernoulliMarginalSlope => "bernoulli marginal-slope",
1088 Self::Survival => "survival",
1089 Self::TransformationNormal => "transformation-normal",
1090 }
1091 }
1092}
1093
1094#[derive(Clone, Debug)]
1095pub struct SavedLinkWiggleRuntime {
1096 pub knots: Vec<f64>,
1097 pub degree: usize,
1098 pub penalty_metadata: Option<WigglePenaltyMetadata>,
1102 pub beta: Vec<f64>,
1103 pub index_shift: Option<Vec<f64>>,
1111}
1112
1113#[derive(Clone, Debug)]
1114pub struct SavedBaselineTimeWiggleRuntime {
1115 pub knots: Vec<f64>,
1116 pub degree: usize,
1117 pub penalty_orders: Vec<usize>,
1118 pub double_penalty: bool,
1119 pub beta: Vec<f64>,
1120}
1121
1122pub use crate::bms::deviation_runtime::ParametricAnchorBlock;
1125
1126#[derive(Clone, Debug, Serialize, Deserialize)]
1127pub struct SavedCompiledFlexBlock {
1128 pub kernel: String,
1129 pub breakpoints: Vec<f64>,
1130 pub basis_dim: usize,
1131 pub span_c0: Vec<Vec<f64>>,
1132 pub span_c1: Vec<Vec<f64>>,
1133 pub span_c2: Vec<Vec<f64>>,
1134 pub span_c3: Vec<Vec<f64>>,
1135 #[serde(default)]
1141 pub anchor_correction: Option<Vec<Vec<f64>>>,
1142 #[serde(default)]
1146 pub anchor_components: Vec<SavedAnchorComponent>,
1147}
1148
1149#[derive(Clone, Debug, Serialize, Deserialize)]
1150pub struct SavedAnchorComponent {
1151 pub kind: SavedAnchorKind,
1152}
1153
1154#[derive(Clone, Debug, Serialize, Deserialize)]
1155pub enum SavedAnchorKind {
1156 Parametric {
1157 block: ParametricAnchorBlock,
1158 ncols: usize,
1159 },
1160 FlexEvaluation { ncols: usize },
1165}
1166
1167#[derive(Clone, Debug)]
1168pub struct SavedPredictionRuntime {
1169 pub model_class: PredictModelClass,
1170 pub likelihood: LikelihoodSpec,
1171 pub inverse_link: Option<InverseLink>,
1172 pub link_wiggle: Option<SavedLinkWiggleRuntime>,
1173 pub baseline_time_wiggle: Option<SavedBaselineTimeWiggleRuntime>,
1174 pub score_warp: Option<SavedCompiledFlexBlock>,
1175 pub link_deviation: Option<SavedCompiledFlexBlock>,
1176 pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
1180 pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
1186 pub influence_absorber_width: Option<usize>,
1196}
1197
1198pub fn gaussian_location_scale_mean_beta(fit: &UnifiedFitResult) -> Option<Array1<f64>> {
1199 fit.block_by_role(BlockRole::Location)
1200 .or_else(|| fit.block_by_role(BlockRole::Mean))
1201 .map(|block| block.beta.clone())
1202}
1203
1204pub fn binomial_location_scale_threshold_beta(fit: &UnifiedFitResult) -> Option<Array1<f64>> {
1205 fit.block_by_role(BlockRole::Threshold)
1206 .or_else(|| fit.block_by_role(BlockRole::Location))
1207 .or_else(|| fit.block_by_role(BlockRole::Mean))
1208 .map(|block| block.beta.clone())
1209}
1210
1211pub fn location_scale_noise_beta(fit: &UnifiedFitResult) -> Option<Array1<f64>> {
1212 fit.block_by_role(BlockRole::Scale)
1213 .map(|block| block.beta.clone())
1214}
1215
1216fn is_dispersion_location_scale_response(response: &gam_problem::types::ResponseFamily) -> bool {
1225 use gam_problem::types::ResponseFamily;
1226 matches!(
1227 response,
1228 ResponseFamily::NegativeBinomial { .. }
1229 | ResponseFamily::Gamma
1230 | ResponseFamily::Beta { .. }
1231 | ResponseFamily::Tweedie { .. }
1232 )
1233}
1234
1235fn validate_location_scale_saved_fit(
1236 fit: &UnifiedFitResult,
1237 model_class: PredictModelClass,
1238 link_wiggle: Option<&SavedLinkWiggleRuntime>,
1239) -> Result<(), FittedModelError> {
1240 let primary = match model_class {
1241 PredictModelClass::GaussianLocationScale | PredictModelClass::DispersionLocationScale => {
1245 gaussian_location_scale_mean_beta(fit)
1246 }
1247 PredictModelClass::BinomialLocationScale => binomial_location_scale_threshold_beta(fit),
1248 _ => None,
1249 }
1250 .ok_or_else(|| FittedModelError::MissingField {
1251 reason: match model_class {
1252 PredictModelClass::GaussianLocationScale => {
1253 "gaussian-location-scale saved fit is missing mean/location block".to_string()
1254 }
1255 PredictModelClass::DispersionLocationScale => {
1256 "dispersion-location-scale saved fit is missing mean/location block".to_string()
1257 }
1258 PredictModelClass::BinomialLocationScale => {
1259 "binomial-location-scale saved fit is missing threshold/location block".to_string()
1260 }
1261 _ => "location-scale saved fit is missing primary block".to_string(),
1262 },
1263 })?;
1264
1265 let scale = location_scale_noise_beta(fit).ok_or_else(|| FittedModelError::MissingField {
1266 reason: "location-scale saved fit is missing scale block".to_string(),
1267 })?;
1268 let expected =
1269 primary.len() + scale.len() + link_wiggle.map_or(0, |runtime| runtime.beta.len());
1270
1271 if let Some(cov) = fit.beta_covariance()
1272 && (cov.nrows() != expected || cov.ncols() != expected)
1273 {
1274 return Err(FittedModelError::SchemaMismatch {
1275 reason: format!(
1276 "location-scale saved conditional covariance shape mismatch: got {}x{}, expected {}x{}",
1277 cov.nrows(),
1278 cov.ncols(),
1279 expected,
1280 expected
1281 ),
1282 });
1283 }
1284 if let Some(cov) = fit.beta_covariance_corrected()
1285 && (cov.nrows() != expected || cov.ncols() != expected)
1286 {
1287 return Err(FittedModelError::SchemaMismatch {
1288 reason: format!(
1289 "location-scale saved corrected covariance shape mismatch: got {}x{}, expected {}x{}",
1290 cov.nrows(),
1291 cov.ncols(),
1292 expected,
1293 expected
1294 ),
1295 });
1296 }
1297 Ok(())
1298}
1299
1300fn validate_survival_saved_block_matches_payload(
1301 fit: &UnifiedFitResult,
1302 role: BlockRole,
1303 payload_beta: Option<&Vec<f64>>,
1304 label: &str,
1305) -> Result<usize, FittedModelError> {
1306 let block = fit
1307 .block_by_role(role)
1308 .ok_or_else(|| FittedModelError::MissingField {
1309 reason: format!("location-scale survival saved fit is missing {label} block"),
1310 })?;
1311 if let Some(saved) = payload_beta
1312 && block.beta.to_vec() != *saved
1313 {
1314 return Err(FittedModelError::SchemaMismatch {
1315 reason: format!(
1316 "location-scale survival saved {label} coefficients disagree with fit_result"
1317 ),
1318 });
1319 }
1320 Ok(block.beta.len())
1321}
1322
1323fn validate_survival_covariate_time_basis(
1324 basis: &SurvivalCovariateTimeBasis,
1325 label: &str,
1326) -> Result<usize, FittedModelError> {
1327 let minimum_knots =
1328 basis
1329 .degree
1330 .checked_add(2)
1331 .ok_or_else(|| FittedModelError::SchemaMismatch {
1332 reason: format!("location-scale survival saved {label} degree overflows"),
1333 })?;
1334 if basis.knots.len() < minimum_knots {
1335 return Err(FittedModelError::SchemaMismatch {
1336 reason: format!(
1337 "location-scale survival saved {label} knot vector has length {}, but degree {} requires at least {minimum_knots}",
1338 basis.knots.len(),
1339 basis.degree
1340 ),
1341 });
1342 }
1343 if basis.knots.iter().any(|value| !value.is_finite())
1344 || basis.knots.windows(2).any(|pair| pair[1] < pair[0])
1345 || basis.knots.first() == basis.knots.last()
1346 {
1347 return Err(FittedModelError::SchemaMismatch {
1348 reason: format!(
1349 "location-scale survival saved {label} knots must be finite, nondecreasing, and span a nonzero interval"
1350 ),
1351 });
1352 }
1353 Ok(basis.knots.len() - basis.degree - 1)
1354}
1355
1356fn validate_survival_location_scale_saved_fit(
1357 payload: &FittedModelPayload,
1358 link_wiggle: Option<&SavedLinkWiggleRuntime>,
1359) -> Result<(), FittedModelError> {
1360 let structure = payload
1361 .survival_location_scale_structure
1362 .as_ref()
1363 .ok_or_else(|| FittedModelError::MissingField {
1364 reason: "location-scale survival model is missing exact replay structure".to_string(),
1365 })?;
1366 let fit = payload
1367 .fit_result
1368 .as_ref()
1369 .ok_or_else(|| FittedModelError::MissingField {
1370 reason: "location-scale survival model is missing canonical fit_result payload"
1371 .to_string(),
1372 })?;
1373 let p_time = validate_survival_saved_block_matches_payload(
1374 fit,
1375 BlockRole::Time,
1376 payload.survival_beta_time.as_ref(),
1377 "time",
1378 )?;
1379 let p_threshold = validate_survival_saved_block_matches_payload(
1380 fit,
1381 BlockRole::Threshold,
1382 payload.survival_beta_threshold.as_ref(),
1383 "threshold",
1384 )?;
1385 let p_log_sigma = validate_survival_saved_block_matches_payload(
1386 fit,
1387 BlockRole::Scale,
1388 payload.survival_beta_log_sigma.as_ref(),
1389 "log-sigma",
1390 )?;
1391 if let Some(basis) = structure.threshold_time_basis.as_ref() {
1392 let width = validate_survival_covariate_time_basis(basis, "threshold time basis")?;
1393 if p_threshold % width != 0 {
1394 return Err(FittedModelError::SchemaMismatch {
1395 reason: format!(
1396 "location-scale survival threshold width {p_threshold} is not divisible by its saved time-basis width {width}"
1397 ),
1398 });
1399 }
1400 }
1401 if let Some(basis) = structure.log_sigma_time_basis.as_ref() {
1402 let width = validate_survival_covariate_time_basis(basis, "log-sigma time basis")?;
1403 if p_log_sigma % width != 0 {
1404 return Err(FittedModelError::SchemaMismatch {
1405 reason: format!(
1406 "location-scale survival log-sigma width {p_log_sigma} is not divisible by its saved time-basis width {width}"
1407 ),
1408 });
1409 }
1410 }
1411 let p_wiggle = match link_wiggle {
1412 Some(runtime) => {
1413 let block = fit.block_by_role(BlockRole::LinkWiggle).ok_or_else(|| {
1414 FittedModelError::MissingField {
1415 reason: "location-scale survival saved fit is missing link-wiggle block"
1416 .to_string(),
1417 }
1418 })?;
1419 if block.beta.to_vec() != runtime.beta {
1420 return Err(FittedModelError::SchemaMismatch {
1421 reason:
1422 "location-scale survival saved link-wiggle coefficients disagree with fit_result"
1423 .to_string(),
1424 });
1425 }
1426 runtime.beta.len()
1427 }
1428 None => {
1429 if fit.block_by_role(BlockRole::LinkWiggle).is_some() {
1430 return Err(FittedModelError::SchemaMismatch {
1431 reason:
1432 "location-scale survival saved fit has a LinkWiggle block without payload metadata"
1433 .to_string(),
1434 });
1435 }
1436 0
1437 }
1438 };
1439 let expected = p_time + p_threshold + p_log_sigma + p_wiggle;
1440
1441 match structure.time_parameterization {
1442 SurvivalLocationScaleTimeParameterization::MonotoneWarp => {}
1443 SurvivalLocationScaleTimeParameterization::ReducedParametricAft => {
1444 if payload.beta_baseline_timewiggle.is_some() || link_wiggle.is_some() {
1445 return Err(FittedModelError::SchemaMismatch {
1446 reason: "reduced parametric-AFT location-scale survival cannot carry a time or link wiggle"
1447 .to_string(),
1448 });
1449 }
1450 let time = fit
1451 .block_by_role(BlockRole::Time)
1452 .expect("time block was validated above");
1453 if time.beta.iter().any(|value| *value != 0.0) {
1454 return Err(FittedModelError::SchemaMismatch {
1455 reason: "reduced parametric-AFT location-scale survival time block must be the exact zero affine lift"
1456 .to_string(),
1457 });
1458 }
1459 }
1460 }
1461 if let Some(timewiggle_beta) = payload.beta_baseline_timewiggle.as_ref() {
1462 let time = fit
1463 .block_by_role(BlockRole::Time)
1464 .expect("time block was validated above");
1465 if timewiggle_beta.len() > time.beta.len()
1466 || time
1467 .beta
1468 .slice(ndarray::s![time.beta.len() - timewiggle_beta.len()..])
1469 .to_vec()
1470 != *timewiggle_beta
1471 {
1472 return Err(FittedModelError::SchemaMismatch {
1473 reason: "location-scale survival baseline-timewiggle coefficients must equal the protected tail of the time block"
1474 .to_string(),
1475 });
1476 }
1477 }
1478
1479 if let Some(cov) = fit.beta_covariance()
1480 && (cov.nrows() != expected || cov.ncols() != expected)
1481 {
1482 return Err(FittedModelError::SchemaMismatch {
1483 reason: format!(
1484 "location-scale survival saved conditional covariance shape mismatch: got {}x{}, expected {}x{}",
1485 cov.nrows(),
1486 cov.ncols(),
1487 expected,
1488 expected
1489 ),
1490 });
1491 }
1492 if let Some(cov) = fit.beta_covariance_corrected()
1493 && (cov.nrows() != expected || cov.ncols() != expected)
1494 {
1495 return Err(FittedModelError::SchemaMismatch {
1496 reason: format!(
1497 "location-scale survival saved corrected covariance shape mismatch: got {}x{}, expected {}x{}",
1498 cov.nrows(),
1499 cov.ncols(),
1500 expected,
1501 expected
1502 ),
1503 });
1504 }
1505 Ok(())
1506}
1507
1508fn validate_marginal_slope_saved_fit(
1509 fit: &UnifiedFitResult,
1510 score_warp: Option<&SavedCompiledFlexBlock>,
1511 link_deviation: Option<&SavedCompiledFlexBlock>,
1512 fit_label: &str,
1513) -> Result<(), FittedModelError> {
1514 validate_marginal_slope_saved_fit_impl(
1515 fit,
1516 score_warp,
1517 link_deviation,
1518 fit_label,
1519 "bernoulli",
1520 2,
1521 "marginal, logslope",
1522 None,
1523 )
1524}
1525
1526fn validate_survival_marginal_slope_saved_fit(
1527 payload: &FittedModelPayload,
1528 fit: &UnifiedFitResult,
1529 fit_label: &str,
1530) -> Result<(), FittedModelError> {
1531 validate_marginal_slope_saved_fit_impl(
1532 fit,
1533 payload.score_warp_runtime.as_ref(),
1534 payload.link_deviation_runtime.as_ref(),
1535 fit_label,
1536 "survival",
1537 3,
1538 "time, marginal, slope",
1539 payload.influence_absorber_width,
1540 )
1541 .and_then(|()| validate_survival_marginal_slope_replay_state(payload, fit, fit_label))
1542}
1543
1544fn validate_marginal_slope_saved_fit_impl(
1552 fit: &UnifiedFitResult,
1553 score_warp: Option<&SavedCompiledFlexBlock>,
1554 link_deviation: Option<&SavedCompiledFlexBlock>,
1555 fit_label: &str,
1556 family_kind: &str,
1557 base_block_count: usize,
1558 base_block_role_list: &str,
1559 influence_absorber_width: Option<usize>,
1560) -> Result<(), FittedModelError> {
1561 let expected_blocks = base_block_count
1562 + usize::from(score_warp.is_some())
1563 + usize::from(link_deviation.is_some())
1564 + usize::from(influence_absorber_width.is_some());
1565 if fit.blocks.len() != expected_blocks {
1566 let score_warp_suffix = if score_warp.is_some() {
1567 ", score-warp"
1568 } else {
1569 ""
1570 };
1571 let link_deviation_suffix = if link_deviation.is_some() {
1572 ", link-deviation"
1573 } else {
1574 ""
1575 };
1576 let influence_suffix = if influence_absorber_width.is_some() {
1577 ", influence-absorber"
1578 } else {
1579 ""
1580 };
1581 return Err(FittedModelError::SchemaMismatch {
1582 reason: format!(
1583 "{family_kind} marginal-slope saved {fit_label} requires {expected_blocks} blocks [{base_block_role_list}{score_warp_suffix}{link_deviation_suffix}{influence_suffix}], got {}",
1584 fit.blocks.len(),
1585 ),
1586 });
1587 }
1588 if let Some(runtime) = score_warp {
1589 let beta = &fit.blocks[base_block_count].beta;
1590 if beta.len() != runtime.basis_dim {
1591 return Err(FittedModelError::SchemaMismatch {
1592 reason: format!(
1593 "{family_kind} marginal-slope saved {fit_label} score-warp coefficient mismatch: beta has {} entries but runtime expects {}",
1594 beta.len(),
1595 runtime.basis_dim
1596 ),
1597 });
1598 }
1599 }
1600 if let Some(runtime) = link_deviation {
1601 let idx = base_block_count + usize::from(score_warp.is_some());
1602 let beta = &fit.blocks[idx].beta;
1603 if beta.len() != runtime.basis_dim {
1604 return Err(FittedModelError::SchemaMismatch {
1605 reason: format!(
1606 "{family_kind} marginal-slope saved {fit_label} link-deviation coefficient mismatch: beta has {} entries but runtime expects {}",
1607 beta.len(),
1608 runtime.basis_dim
1609 ),
1610 });
1611 }
1612 }
1613 if let Some(width) = influence_absorber_width {
1614 let idx = base_block_count
1615 + usize::from(score_warp.is_some())
1616 + usize::from(link_deviation.is_some());
1617 if width == 0 || fit.blocks[idx].beta.len() != width {
1618 return Err(FittedModelError::SchemaMismatch {
1619 reason: format!(
1620 "{family_kind} marginal-slope saved {fit_label} influence absorber width is {width}, but its fitted block has {} coefficients",
1621 fit.blocks[idx].beta.len(),
1622 ),
1623 });
1624 }
1625 }
1626 Ok(())
1627}
1628
1629fn validate_survival_marginal_slope_replay_state(
1630 payload: &FittedModelPayload,
1631 fit: &UnifiedFitResult,
1632 fit_label: &str,
1633) -> Result<(), FittedModelError> {
1634 let score_covariance = payload
1635 .survival_marginal_slope_score_covariance
1636 .as_ref()
1637 .ok_or_else(|| FittedModelError::MissingField {
1638 reason: format!(
1639 "survival marginal-slope saved {fit_label} is missing its exact latent-score covariance"
1640 ),
1641 })?;
1642 if score_covariance.len() != 1
1643 || score_covariance[0].len() != 1
1644 || !score_covariance[0][0].is_finite()
1645 || score_covariance[0][0] < 0.0
1646 {
1647 return Err(FittedModelError::SchemaMismatch {
1648 reason: format!(
1649 "survival marginal-slope saved {fit_label} scalar latent-score covariance must be a finite non-negative 1x1 matrix"
1650 ),
1651 });
1652 }
1653 match (
1654 payload.influence_absorber_width,
1655 payload.influence_absorber_design.as_ref(),
1656 ) {
1657 (None, None) => {}
1658 (Some(width), Some(rows)) => {
1659 if rows.is_empty()
1660 || rows
1661 .iter()
1662 .any(|row| row.len() != width || row.iter().any(|value| !value.is_finite()))
1663 {
1664 return Err(FittedModelError::SchemaMismatch {
1665 reason: format!(
1666 "survival marginal-slope saved {fit_label} influence absorber must be a non-empty finite rectangular matrix with width {width}"
1667 ),
1668 });
1669 }
1670 }
1671 _ => {
1672 return Err(FittedModelError::SchemaMismatch {
1673 reason: format!(
1674 "survival marginal-slope saved {fit_label} influence absorber width and exact training-row design must be present together"
1675 ),
1676 });
1677 }
1678 }
1679
1680 let timewiggle_metadata = (
1681 payload.baseline_timewiggle_knots.as_ref(),
1682 payload.baseline_timewiggle_degree,
1683 payload.beta_baseline_timewiggle.as_ref(),
1684 );
1685 match timewiggle_metadata {
1686 (None, None, None) => {}
1687 (Some(knots), Some(degree), Some(beta)) => {
1688 if beta.is_empty()
1689 || knots.len() < degree.saturating_add(2)
1690 || knots.iter().chain(beta).any(|value| !value.is_finite())
1691 {
1692 return Err(FittedModelError::SchemaMismatch {
1693 reason: format!(
1694 "survival marginal-slope saved {fit_label} has invalid exact baseline-timewiggle authority"
1695 ),
1696 });
1697 }
1698 let time_beta = &fit.blocks[0].beta;
1699 if time_beta.len() < beta.len()
1700 || time_beta
1701 .slice(ndarray::s![time_beta.len() - beta.len()..])
1702 .to_vec()
1703 != *beta
1704 {
1705 return Err(FittedModelError::SchemaMismatch {
1706 reason: format!(
1707 "survival marginal-slope saved {fit_label} baseline-timewiggle beta does not equal the protected tail of the fitted time block"
1708 ),
1709 });
1710 }
1711 }
1712 _ => {
1713 return Err(FittedModelError::SchemaMismatch {
1714 reason: format!(
1715 "survival marginal-slope saved {fit_label} baseline-timewiggle knots, degree, and beta must be present together"
1716 ),
1717 });
1718 }
1719 }
1720 if payload.beta_baseline_timewiggle_by_cause.is_some() {
1721 return Err(FittedModelError::SchemaMismatch {
1722 reason: "survival marginal-slope saved fit cannot carry cause-specific timewiggle coefficients"
1723 .to_string(),
1724 });
1725 }
1726 Ok(())
1727}
1728
1729impl SavedLinkWiggleRuntime {
1730 fn validate_monotone_derivative(
1731 &self,
1732 q0: &Array1<f64>,
1733 ) -> Result<Array1<f64>, FittedModelError> {
1734 let d_constrained = self.constrained_basis(q0, BasisOptions::first_derivative())?;
1741 let beta_link_wiggle = Array1::from_vec(self.beta.clone());
1742 let dq_dq0 = d_constrained.dot(&beta_link_wiggle) + 1.0;
1743 if let Some((idx, value)) = dq_dq0.iter().copied().enumerate().find(|(_, v)| *v <= 0.0) {
1744 return Err(FittedModelError::PayloadCorrupt {
1745 reason: format!(
1746 "saved link-wiggle is not monotone at row {idx}: dq/dq0={value:.3e} <= 0"
1747 ),
1748 });
1749 }
1750 Ok(dq_dq0)
1751 }
1752
1753 pub fn constrained_basis(
1754 &self,
1755 q0: &Array1<f64>,
1756 basis_options: BasisOptions,
1757 ) -> Result<Array2<f64>, FittedModelError> {
1758 let knot_arr = Array1::from_vec(self.knots.clone());
1759 let constrained = monotone_wiggle_basis_with_derivative_order(
1760 q0.view(),
1761 &knot_arr,
1762 self.degree,
1763 basis_options.derivative_order,
1764 )
1765 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
1766 if constrained.ncols() != self.beta.len() {
1767 return Err(FittedModelError::SchemaMismatch {
1768 reason: format!(
1769 "saved link-wiggle dimension mismatch: coefficients have {} entries but basis has {} columns",
1770 self.beta.len(),
1771 constrained.ncols()
1772 ),
1773 });
1774 }
1775 Ok(constrained)
1776 }
1777
1778 pub fn design(&self, q0: &Array1<f64>) -> Result<Array2<f64>, FittedModelError> {
1779 self.validate_monotone_derivative(q0)?;
1780 self.constrained_basis(q0, BasisOptions::value())
1781 }
1782
1783 pub fn warp_index(
1792 &self,
1793 base: &Array1<f64>,
1794 mean_design: &DesignMatrix,
1795 ) -> Result<Array1<f64>, FittedModelError> {
1796 if mean_design.nrows() != base.len() {
1797 return Err(FittedModelError::SchemaMismatch {
1798 reason: format!(
1799 "link-wiggle base predictor has {} rows but mean design has {}",
1800 base.len(),
1801 mean_design.nrows()
1802 ),
1803 });
1804 }
1805 if let Some((row, value)) = base
1806 .iter()
1807 .copied()
1808 .enumerate()
1809 .find(|(_, value)| !value.is_finite())
1810 {
1811 return Err(FittedModelError::InvalidInput {
1812 reason: format!("link-wiggle base predictor is non-finite at row {row}: {value}"),
1813 });
1814 }
1815 let Some(shift) = self.index_shift.as_ref() else {
1816 return Ok(base.clone());
1817 };
1818 if shift.len() != mean_design.ncols() {
1819 return Err(FittedModelError::SchemaMismatch {
1820 reason: format!(
1821 "link-wiggle frozen-index shift has {} entries but the mean design has {} columns",
1822 shift.len(),
1823 mean_design.ncols()
1824 ),
1825 });
1826 }
1827 if let Some((column, value)) = shift
1828 .iter()
1829 .copied()
1830 .enumerate()
1831 .find(|(_, value)| !value.is_finite())
1832 {
1833 return Err(FittedModelError::PayloadCorrupt {
1834 reason: format!(
1835 "link-wiggle frozen-index shift is non-finite at column {column}: {value}"
1836 ),
1837 });
1838 }
1839 let shift = Array1::from_vec(shift.clone());
1840 Ok(base + &mean_design.dot(&shift))
1841 }
1842
1843 pub fn basis_row_scalar(&self, q0: f64) -> Result<Array1<f64>, FittedModelError> {
1844 let q = Array1::from_vec(vec![q0]);
1845 let x = self.design(&q)?;
1846 if x.nrows() != 1 {
1847 return Err(FittedModelError::SchemaMismatch {
1848 reason: format!(
1849 "saved link-wiggle scalar evaluation expected 1 row, got {}",
1850 x.nrows()
1851 ),
1852 });
1853 }
1854 Ok(x.row(0).to_owned())
1855 }
1856
1857 pub fn apply(&self, q0: &Array1<f64>) -> Result<Array1<f64>, FittedModelError> {
1858 self.apply_with_index(q0, q0)
1859 }
1860
1861 pub fn apply_with_index(
1871 &self,
1872 base: &Array1<f64>,
1873 warp_index: &Array1<f64>,
1874 ) -> Result<Array1<f64>, FittedModelError> {
1875 if base.len() != warp_index.len() {
1876 return Err(FittedModelError::SchemaMismatch {
1877 reason: format!(
1878 "link-wiggle base predictor has {} rows but warp index has {}",
1879 base.len(),
1880 warp_index.len()
1881 ),
1882 });
1883 }
1884 self.validate_monotone_derivative(warp_index)?;
1885 let xwiggle = self.constrained_basis(warp_index, BasisOptions::value())?;
1886 let beta_link_wiggle = Array1::from_vec(self.beta.clone());
1887 Ok(base + &xwiggle.dot(&beta_link_wiggle))
1888 }
1889
1890 pub fn derivative_q0(&self, q0: &Array1<f64>) -> Result<Array1<f64>, FittedModelError> {
1891 self.validate_monotone_derivative(q0)
1892 }
1893}
1894
1895impl SavedBaselineTimeWiggleRuntime {
1896 pub fn validate_global_monotonicity(&self) -> Result<(), FittedModelError> {
1897 validate_monotone_wiggle_beta_nonnegative(&self.beta, "saved baseline-timewiggle")
1898 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })
1899 }
1900}
1901
1902impl SavedCompiledFlexBlock {
1903 pub(crate) fn validate_exact_replay_contract(&self) -> Result<(), FittedModelError> {
1904 if self.kernel.is_empty() {
1905 return Err(FittedModelError::SchemaMismatch {
1906 reason: "saved anchored deviation runtime is missing the exact kernel marker"
1907 .to_string(),
1908 });
1909 }
1910 if self.kernel != crate::cubic_cell_kernel::ANCHORED_DEVIATION_KERNEL {
1911 return Err(FittedModelError::IncompatibleConfig {
1912 reason: format!(
1913 "saved anchored deviation runtime uses unsupported kernel '{}'; expected {}",
1914 self.kernel,
1915 crate::cubic_cell_kernel::ANCHORED_DEVIATION_KERNEL
1916 ),
1917 });
1918 }
1919 if self.basis_dim == 0 {
1920 return Err(FittedModelError::SchemaMismatch {
1921 reason: format!(
1922 "saved anchored deviation runtime basis_dim must be positive, got {}",
1923 self.basis_dim
1924 ),
1925 });
1926 }
1927 if self.breakpoints.len() < 2 {
1928 return Err(FittedModelError::SchemaMismatch {
1929 reason: format!(
1930 "saved anchored deviation runtime requires at least two breakpoints, got {}",
1931 self.breakpoints.len()
1932 ),
1933 });
1934 }
1935 for window in self.breakpoints.windows(2) {
1936 let left = window[0];
1937 let right = window[1];
1938 if !left.is_finite() || !right.is_finite() || right <= left {
1939 return Err(FittedModelError::PayloadCorrupt {
1940 reason: format!(
1941 "saved anchored deviation runtime breakpoints must be finite and strictly increasing, got [{left}, {right}]"
1942 ),
1943 });
1944 }
1945 }
1946 let span_count = self.breakpoints.len() - 1;
1947 self.validate_coefficient_matrix(&self.span_c0, "c0", span_count)?;
1948 self.validate_coefficient_matrix(&self.span_c1, "c1", span_count)?;
1949 self.validate_coefficient_matrix(&self.span_c2, "c2", span_count)?;
1950 self.validate_coefficient_matrix(&self.span_c3, "c3", span_count)?;
1951 self.validate_c2_span_continuity()?;
1952 self.validate_anchor_residual_shape()?;
1953 Ok(())
1954 }
1955
1956 fn validate_anchor_residual_shape(&self) -> Result<(), FittedModelError> {
1957 let coeffs = match self.anchor_correction.as_ref() {
1958 Some(c) => c,
1959 None => {
1960 if !self.anchor_components.is_empty() {
1961 return Err(FittedModelError::SchemaMismatch {
1962 reason:
1963 "saved anchored deviation runtime has anchor_components but no anchor_correction"
1964 .to_string(),
1965 });
1966 }
1967 return Ok(());
1968 }
1969 };
1970 let d: usize = self
1971 .anchor_components
1972 .iter()
1973 .map(|c| match &c.kind {
1974 SavedAnchorKind::Parametric { ncols, .. } => *ncols,
1975 SavedAnchorKind::FlexEvaluation { ncols } => *ncols,
1976 })
1977 .sum();
1978 if coeffs.len() != d {
1979 return Err(FittedModelError::SchemaMismatch {
1980 reason: format!(
1981 "saved anchored deviation runtime anchor_correction has {} rows; expected {} (sum of component ncols)",
1982 coeffs.len(),
1983 d,
1984 ),
1985 });
1986 }
1987 for (i, row) in coeffs.iter().enumerate() {
1988 if row.len() != self.basis_dim {
1989 return Err(FittedModelError::SchemaMismatch {
1990 reason: format!(
1991 "saved anchored deviation runtime anchor_correction row {} has width {}, expected basis_dim {}",
1992 i,
1993 row.len(),
1994 self.basis_dim,
1995 ),
1996 });
1997 }
1998 for (j, &v) in row.iter().enumerate() {
1999 if !v.is_finite() {
2000 return Err(FittedModelError::PayloadCorrupt {
2001 reason: format!(
2002 "saved anchored deviation runtime anchor_correction ({i},{j}) is non-finite"
2003 ),
2004 });
2005 }
2006 }
2007 }
2008 Ok(())
2009 }
2010
2011 fn validate_c2_span_continuity(&self) -> Result<(), FittedModelError> {
2012 const TOL: f64 = 1e-8;
2013 for span_idx in 1..self.breakpoints.len() - 1 {
2014 let left_span = span_idx - 1;
2015 let right_span = span_idx;
2016 let width = self.breakpoints[span_idx] - self.breakpoints[left_span];
2017 for basis_idx in 0..self.basis_dim {
2018 let left_value = self.span_c0[left_span][basis_idx]
2019 + self.span_c1[left_span][basis_idx] * width
2020 + self.span_c2[left_span][basis_idx] * width * width
2021 + self.span_c3[left_span][basis_idx] * width * width * width;
2022 let left_d1 = self.span_c1[left_span][basis_idx]
2023 + 2.0 * self.span_c2[left_span][basis_idx] * width
2024 + 3.0 * self.span_c3[left_span][basis_idx] * width * width;
2025 let left_d2 = 2.0 * self.span_c2[left_span][basis_idx]
2026 + 6.0 * self.span_c3[left_span][basis_idx] * width;
2027 let right_value = self.span_c0[right_span][basis_idx];
2028 let right_d1 = self.span_c1[right_span][basis_idx];
2029 let right_d2 = 2.0 * self.span_c2[right_span][basis_idx];
2030 if (left_value - right_value).abs() > TOL
2031 || (left_d1 - right_d1).abs() > TOL
2032 || (left_d2 - right_d2).abs() > TOL
2033 {
2034 return Err(FittedModelError::SchemaMismatch {
2035 reason: format!(
2036 "saved anchored deviation runtime must be C2 cubic at breakpoint {span_idx}, basis {basis_idx}: value jump={:.3e}, d1 jump={:.3e}, d2 jump={:.3e}",
2037 left_value - right_value,
2038 left_d1 - right_d1,
2039 left_d2 - right_d2
2040 ),
2041 });
2042 }
2043 }
2044 }
2045 Ok(())
2046 }
2047
2048 fn validate_coefficient_matrix(
2049 &self,
2050 matrix: &[Vec<f64>],
2051 label: &str,
2052 expected_rows: usize,
2053 ) -> Result<(), FittedModelError> {
2054 if matrix.len() != expected_rows {
2055 return Err(FittedModelError::SchemaMismatch {
2056 reason: format!(
2057 "saved anchored deviation runtime {label} row count mismatch: got {}, expected {}",
2058 matrix.len(),
2059 expected_rows
2060 ),
2061 });
2062 }
2063 for (row_idx, row) in matrix.iter().enumerate() {
2064 if row.len() != self.basis_dim {
2065 return Err(FittedModelError::SchemaMismatch {
2066 reason: format!(
2067 "saved anchored deviation runtime {label} row {} has width {}, expected {}",
2068 row_idx,
2069 row.len(),
2070 self.basis_dim
2071 ),
2072 });
2073 }
2074 for (j, &value) in row.iter().enumerate() {
2075 if !value.is_finite() {
2076 return Err(FittedModelError::PayloadCorrupt {
2077 reason: format!(
2078 "saved anchored deviation runtime {label} entry ({row_idx},{j}) is non-finite"
2079 ),
2080 });
2081 }
2082 }
2083 }
2084 Ok(())
2085 }
2086
2087 fn right_boundary_basis_value(&self, basis_idx: usize) -> f64 {
2088 let last_span = self.breakpoints.len() - 2;
2089 let width = self.breakpoints[last_span + 1] - self.breakpoints[last_span];
2090 self.span_c0[last_span][basis_idx]
2091 + self.span_c1[last_span][basis_idx] * width
2092 + self.span_c2[last_span][basis_idx] * width * width
2093 + self.span_c3[last_span][basis_idx] * width * width * width
2094 }
2095
2096 fn evaluate_span_polynomial_design(
2097 &self,
2098 values: &Array1<f64>,
2099 derivative_order: usize,
2100 ) -> Result<Array2<f64>, FittedModelError> {
2101 self.validate_exact_replay_contract()?;
2102 let (left_ep, right_ep) = self.support_interval()?;
2103 let mut out = Array2::<f64>::zeros((values.len(), self.basis_dim));
2104 for (row_idx, &value) in values.iter().enumerate() {
2105 if !value.is_finite() {
2106 return Err(FittedModelError::PayloadCorrupt {
2107 reason: format!(
2108 "saved anchored deviation runtime design value at row {row_idx} is non-finite ({value})"
2109 ),
2110 });
2111 }
2112 if value < left_ep {
2113 if derivative_order == 0 {
2114 for basis_idx in 0..self.basis_dim {
2115 out[[row_idx, basis_idx]] = self.span_c0[0][basis_idx];
2116 }
2117 }
2118 continue;
2119 }
2120 if value > right_ep {
2121 if derivative_order == 0 {
2122 for basis_idx in 0..self.basis_dim {
2123 out[[row_idx, basis_idx]] = self.right_boundary_basis_value(basis_idx);
2124 }
2125 }
2126 continue;
2127 }
2128 let span_idx = self.left_biased_span_index_for(value)?;
2129 let t = value - self.breakpoints[span_idx];
2130 for basis_idx in 0..self.basis_dim {
2131 let c0 = self.span_c0[span_idx][basis_idx];
2132 let c1 = self.span_c1[span_idx][basis_idx];
2133 let c2 = self.span_c2[span_idx][basis_idx];
2134 let c3 = self.span_c3[span_idx][basis_idx];
2135 out[[row_idx, basis_idx]] = match derivative_order {
2136 0 => c0 + c1 * t + c2 * t * t + c3 * t * t * t,
2137 1 => c1 + 2.0 * c2 * t + 3.0 * c3 * t * t,
2138 2 => 2.0 * c2 + 6.0 * c3 * t,
2139 3 => 6.0 * c3,
2140 4 => 0.0,
2141 other => {
2142 return Err(FittedModelError::IncompatibleConfig {
2143 reason: format!(
2144 "saved anchored deviation runtime only supports derivative orders up to 4, got {other}"
2145 ),
2146 });
2147 }
2148 };
2149 }
2150 }
2151 Ok(out)
2152 }
2153
2154 pub fn breakpoints(&self) -> Result<Vec<f64>, FittedModelError> {
2155 self.validate_exact_replay_contract()?;
2156 Ok(self.breakpoints.clone())
2157 }
2158
2159 pub fn span_count(&self) -> Result<usize, FittedModelError> {
2160 Ok(self.breakpoints()?.windows(2).count())
2161 }
2162
2163 pub fn span_index_for(&self, value: f64) -> Result<usize, FittedModelError> {
2164 let points = self.breakpoints()?;
2165 span_index_for_breakpoints(&points, value, "saved anchored deviation span lookup")
2166 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })
2167 }
2168
2169 fn left_biased_span_index_for(&self, value: f64) -> Result<usize, FittedModelError> {
2170 let mut span_idx = span_index_for_breakpoints(
2171 &self.breakpoints,
2172 value,
2173 "saved anchored deviation span lookup",
2174 )
2175 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
2176 if span_idx > 0 && value == self.breakpoints[span_idx] {
2179 span_idx -= 1;
2180 }
2181 Ok(span_idx)
2182 }
2183
2184 pub fn local_cubic_on_span(
2185 &self,
2186 beta: ArrayView1<'_, f64>,
2187 span_idx: usize,
2188 ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
2189 self.validate_exact_replay_contract()?;
2190 if beta.len() != self.basis_dim {
2191 return Err(FittedModelError::SchemaMismatch {
2192 reason: format!(
2193 "saved anchored deviation coefficient length mismatch: got {}, expected {}",
2194 beta.len(),
2195 self.basis_dim
2196 ),
2197 });
2198 }
2199 self.local_cubic_on_span_validated(beta, span_idx)
2200 }
2201
2202 fn local_cubic_on_span_validated(
2203 &self,
2204 beta: ArrayView1<'_, f64>,
2205 span_idx: usize,
2206 ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
2207 let points = &self.breakpoints;
2208 if span_idx + 1 >= points.len() {
2209 return Err(FittedModelError::SchemaMismatch {
2210 reason: format!(
2211 "saved anchored deviation span index {} out of range for {} spans",
2212 span_idx,
2213 points.len() - 1
2214 ),
2215 });
2216 }
2217 let left = points[span_idx];
2218 let right = points[span_idx + 1];
2219 Ok(crate::cubic_cell_kernel::LocalSpanCubic {
2220 left,
2221 right,
2222 c0: self.span_c0[span_idx]
2223 .iter()
2224 .zip(beta.iter())
2225 .map(|(coeff, weight)| coeff * weight)
2226 .sum(),
2227 c1: self.span_c1[span_idx]
2228 .iter()
2229 .zip(beta.iter())
2230 .map(|(coeff, weight)| coeff * weight)
2231 .sum(),
2232 c2: self.span_c2[span_idx]
2233 .iter()
2234 .zip(beta.iter())
2235 .map(|(coeff, weight)| coeff * weight)
2236 .sum(),
2237 c3: self.span_c3[span_idx]
2238 .iter()
2239 .zip(beta.iter())
2240 .map(|(coeff, weight)| coeff * weight)
2241 .sum(),
2242 })
2243 }
2244
2245 pub fn basis_span_cubic(
2246 &self,
2247 span_idx: usize,
2248 basis_idx: usize,
2249 ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
2250 self.validate_exact_replay_contract()?;
2251 if basis_idx >= self.basis_dim {
2252 return Err(FittedModelError::SchemaMismatch {
2253 reason: format!(
2254 "saved anchored deviation basis index {} out of range for {} coefficients",
2255 basis_idx, self.basis_dim
2256 ),
2257 });
2258 }
2259 self.basis_span_cubic_validated(span_idx, basis_idx)
2260 }
2261
2262 fn basis_span_cubic_validated(
2263 &self,
2264 span_idx: usize,
2265 basis_idx: usize,
2266 ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
2267 let points = &self.breakpoints;
2268 if span_idx + 1 >= points.len() {
2269 return Err(FittedModelError::SchemaMismatch {
2270 reason: format!(
2271 "saved anchored deviation span index {} out of range for {} spans",
2272 span_idx,
2273 points.len() - 1
2274 ),
2275 });
2276 }
2277 Ok(crate::cubic_cell_kernel::LocalSpanCubic {
2278 left: points[span_idx],
2279 right: points[span_idx + 1],
2280 c0: self.span_c0[span_idx][basis_idx],
2281 c1: self.span_c1[span_idx][basis_idx],
2282 c2: self.span_c2[span_idx][basis_idx],
2283 c3: self.span_c3[span_idx][basis_idx],
2284 })
2285 }
2286
2287 pub fn basis_cubic_at(
2288 &self,
2289 basis_idx: usize,
2290 value: f64,
2291 ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
2292 self.validate_exact_replay_contract()?;
2293 if basis_idx >= self.basis_dim {
2294 return Err(FittedModelError::SchemaMismatch {
2295 reason: format!(
2296 "saved anchored deviation basis index {} out of range for {} coefficients",
2297 basis_idx, self.basis_dim
2298 ),
2299 });
2300 }
2301 let (left_ep, right_ep) = self.support_interval()?;
2302 if value < left_ep {
2303 return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
2304 left: left_ep,
2305 right: left_ep + 1.0,
2306 c0: self.span_c0[0][basis_idx],
2307 c1: 0.0,
2308 c2: 0.0,
2309 c3: 0.0,
2310 });
2311 }
2312 if value > right_ep {
2313 return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
2314 left: right_ep,
2315 right: right_ep + 1.0,
2316 c0: self.right_boundary_basis_value(basis_idx),
2317 c1: 0.0,
2318 c2: 0.0,
2319 c3: 0.0,
2320 });
2321 }
2322 let span_idx = self.left_biased_span_index_for(value)?;
2323 self.basis_span_cubic_validated(span_idx, basis_idx)
2324 }
2325
2326 pub fn local_cubic_at(
2327 &self,
2328 beta: ArrayView1<'_, f64>,
2329 value: f64,
2330 ) -> Result<crate::cubic_cell_kernel::LocalSpanCubic, FittedModelError> {
2331 self.validate_exact_replay_contract()?;
2332 if beta.len() != self.basis_dim {
2333 return Err(FittedModelError::SchemaMismatch {
2334 reason: format!(
2335 "saved anchored deviation coefficient length mismatch: got {}, expected {}",
2336 beta.len(),
2337 self.basis_dim
2338 ),
2339 });
2340 }
2341 let (left_ep, right_ep) = self.support_interval()?;
2342 if value < left_ep {
2343 return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
2344 left: left_ep,
2345 right: left_ep + 1.0,
2346 c0: self.span_c0[0]
2347 .iter()
2348 .zip(beta.iter())
2349 .map(|(coeff, weight)| coeff * weight)
2350 .sum(),
2351 c1: 0.0,
2352 c2: 0.0,
2353 c3: 0.0,
2354 });
2355 }
2356 if value > right_ep {
2357 return Ok(crate::cubic_cell_kernel::LocalSpanCubic {
2358 left: right_ep,
2359 right: right_ep + 1.0,
2360 c0: (0..self.basis_dim)
2361 .map(|basis_idx| self.right_boundary_basis_value(basis_idx) * beta[basis_idx])
2362 .sum(),
2363 c1: 0.0,
2364 c2: 0.0,
2365 c3: 0.0,
2366 });
2367 }
2368 let span_idx = self.left_biased_span_index_for(value)?;
2369 self.local_cubic_on_span_validated(beta, span_idx)
2370 }
2371
2372 fn support_interval(&self) -> Result<(f64, f64), FittedModelError> {
2373 let points = self.breakpoints()?;
2374 match (points.first(), points.last()) {
2375 (Some(&left), Some(&right)) => Ok((left, right)),
2376 _ => Err(FittedModelError::MissingField {
2377 reason: "saved anchored deviation runtime is missing support breakpoints"
2378 .to_string(),
2379 }),
2380 }
2381 }
2382
2383 pub fn design(&self, values: &Array1<f64>) -> Result<Array2<f64>, FittedModelError> {
2384 self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)
2393 }
2394
2395 pub fn design_uncorrected(
2403 &self,
2404 values: &Array1<f64>,
2405 ) -> Result<Array2<f64>, FittedModelError> {
2406 self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)
2407 }
2408
2409 pub fn design_with_anchor_rows(
2418 &self,
2419 values: &Array1<f64>,
2420 anchor_rows: ndarray::ArrayView2<f64>,
2421 ) -> Result<Array2<f64>, FittedModelError> {
2422 let mut out =
2423 self.evaluate_span_polynomial_design(values, BasisOptions::value().derivative_order)?;
2424 if let Some(m_rows) = self.anchor_correction.as_ref() {
2425 let d = m_rows.len();
2426 if anchor_rows.nrows() != values.len() {
2427 return Err(FittedModelError::SchemaMismatch {
2428 reason: format!(
2429 "design_with_anchor_rows: anchor_rows has {} rows, expected {} (matching values)",
2430 anchor_rows.nrows(),
2431 values.len(),
2432 ),
2433 });
2434 }
2435 if anchor_rows.ncols() != d {
2436 return Err(FittedModelError::SchemaMismatch {
2437 reason: format!(
2438 "design_with_anchor_rows: anchor_rows has {} cols, expected {} (sum of component ncols)",
2439 anchor_rows.ncols(),
2440 d,
2441 ),
2442 });
2443 }
2444 let mut m_dense = Array2::<f64>::zeros((d, self.basis_dim));
2446 for (i, row) in m_rows.iter().enumerate() {
2447 if row.len() != self.basis_dim {
2448 return Err(FittedModelError::SchemaMismatch {
2449 reason: format!(
2450 "design_with_anchor_rows: anchor_correction row {} has length {}, expected basis_dim {}",
2451 i,
2452 row.len(),
2453 self.basis_dim,
2454 ),
2455 });
2456 }
2457 for (j, &v) in row.iter().enumerate() {
2458 m_dense[[i, j]] = v;
2459 }
2460 }
2461 let subtract = anchor_rows.dot(&m_dense);
2464 out = out - subtract;
2465 } else if anchor_rows.ncols() != 0 {
2466 return Err(FittedModelError::SchemaMismatch {
2467 reason: format!(
2468 "design_with_anchor_rows: runtime has no anchor residual but anchor_rows has {} cols",
2469 anchor_rows.ncols(),
2470 ),
2471 });
2472 }
2473 Ok(out)
2474 }
2475
2476 pub fn anchor_correction_matrix(
2484 &self,
2485 n_anchor_rows: ndarray::ArrayView2<f64>,
2486 ) -> Result<Option<Array2<f64>>, FittedModelError> {
2487 let Some(m_rows) = self.anchor_correction.as_ref() else {
2488 return Ok(None);
2489 };
2490 let d = m_rows.len();
2491 if n_anchor_rows.ncols() != d {
2492 return Err(FittedModelError::SchemaMismatch {
2493 reason: format!(
2494 "anchor_correction_matrix: anchor_rows has {} cols, expected {} (sum of component ncols)",
2495 n_anchor_rows.ncols(),
2496 d,
2497 ),
2498 });
2499 }
2500 let mut m_dense = Array2::<f64>::zeros((d, self.basis_dim));
2501 for (i, row) in m_rows.iter().enumerate() {
2502 if row.len() != self.basis_dim {
2503 return Err(FittedModelError::SchemaMismatch {
2504 reason: format!(
2505 "anchor_correction_matrix: M row {} has length {}, expected basis_dim {}",
2506 i,
2507 row.len(),
2508 self.basis_dim,
2509 ),
2510 });
2511 }
2512 for (j, &v) in row.iter().enumerate() {
2513 m_dense[[i, j]] = v;
2514 }
2515 }
2516 Ok(Some(n_anchor_rows.dot(&m_dense)))
2519 }
2520
2521 pub fn first_derivative_design(
2522 &self,
2523 values: &Array1<f64>,
2524 ) -> Result<Array2<f64>, FittedModelError> {
2525 self.evaluate_span_polynomial_design(
2526 values,
2527 BasisOptions::first_derivative().derivative_order,
2528 )
2529 }
2530
2531 pub fn second_derivative_design(
2532 &self,
2533 values: &Array1<f64>,
2534 ) -> Result<Array2<f64>, FittedModelError> {
2535 self.evaluate_span_polynomial_design(
2536 values,
2537 BasisOptions::second_derivative().derivative_order,
2538 )
2539 }
2540}
2541
2542impl FittedFamily {
2543 #[inline]
2544 pub fn likelihood(&self) -> LikelihoodSpec {
2545 let spec = match self {
2546 Self::Standard { likelihood, .. }
2547 | Self::LocationScale { likelihood, .. }
2548 | Self::MarginalSlope { likelihood, .. }
2549 | Self::Survival { likelihood, .. }
2550 | Self::TransformationNormal { likelihood, .. } => likelihood,
2551 Self::LatentSurvival { .. } | Self::LatentBinary { .. } => {
2552 return LikelihoodSpec::royston_parmar();
2553 }
2554 };
2555 spec.clone()
2556 }
2557
2558 #[inline]
2559 pub fn frailty(&self) -> Option<&FrailtySpec> {
2560 match self {
2561 Self::MarginalSlope { frailty, .. }
2562 | Self::Survival { frailty, .. }
2563 | Self::LatentSurvival { frailty }
2564 | Self::LatentBinary { frailty } => Some(frailty),
2565 _ => None,
2566 }
2567 }
2568}
2569
2570fn re_factor_smooth_group_col(basis: &gam_terms::smooth::SmoothBasisSpec) -> Option<usize> {
2577 use gam_terms::smooth::{FactorSmoothFlavour, SmoothBasisSpec};
2578 match basis {
2579 SmoothBasisSpec::FactorSmooth { spec } => {
2580 matches!(spec.flavour, FactorSmoothFlavour::Re).then_some(spec.group_col)
2581 }
2582 SmoothBasisSpec::ByVariable { inner, .. }
2583 | SmoothBasisSpec::FactorSumToZero { inner, .. } => re_factor_smooth_group_col(inner),
2584 SmoothBasisSpec::BySmooth { smooth, .. } => re_factor_smooth_group_col(smooth),
2585 _ => None,
2586 }
2587}
2588
2589fn collect_smooth_extrapolation_axes(
2595 basis: &gam_terms::smooth::SmoothBasisSpec,
2596 n_training_headers: usize,
2597 out: &mut std::collections::HashSet<usize>,
2598) {
2599 use gam_terms::smooth::SmoothBasisSpec;
2600 let push = |col: usize, out: &mut std::collections::HashSet<usize>| {
2601 if col < n_training_headers {
2602 out.insert(col);
2603 }
2604 };
2605 match basis {
2606 SmoothBasisSpec::BSpline1D { feature_col, .. } => push(*feature_col, out),
2608 SmoothBasisSpec::TensorBSpline { feature_cols, .. } => {
2611 for &c in feature_cols {
2612 push(c, out);
2613 }
2614 }
2615 SmoothBasisSpec::ThinPlate { feature_cols, .. }
2622 | SmoothBasisSpec::Matern { feature_cols, .. }
2623 | SmoothBasisSpec::MeasureJet { feature_cols, .. }
2624 | SmoothBasisSpec::Duchon { feature_cols, .. } => {
2625 for &c in feature_cols {
2626 push(c, out);
2627 }
2628 }
2629 SmoothBasisSpec::FactorSmooth { spec } => {
2633 for &c in &spec.continuous_cols {
2634 push(c, out);
2635 }
2636 }
2637 SmoothBasisSpec::ByVariable { inner, .. }
2639 | SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2640 collect_smooth_extrapolation_axes(inner, n_training_headers, out)
2641 }
2642 SmoothBasisSpec::BySmooth { smooth, .. } => {
2643 collect_smooth_extrapolation_axes(smooth, n_training_headers, out)
2644 }
2645 SmoothBasisSpec::Sphere { .. }
2651 | SmoothBasisSpec::ConstantCurvature { .. }
2652 | SmoothBasisSpec::Pca { .. } => {}
2653 }
2654}
2655
2656fn collect_by_variable_numeric_axes(
2675 basis: &gam_terms::smooth::SmoothBasisSpec,
2676 n_training_headers: usize,
2677 out: &mut std::collections::HashSet<usize>,
2678) {
2679 use gam_terms::smooth::{BySmoothKind, ByVarKind, SmoothBasisSpec};
2680 match basis {
2681 SmoothBasisSpec::ByVariable {
2682 inner,
2683 by_col,
2684 kind,
2685 ..
2686 } => {
2687 if matches!(kind, BySmoothKind::Numeric) && *by_col < n_training_headers {
2688 out.insert(*by_col);
2689 }
2690 collect_by_variable_numeric_axes(inner, n_training_headers, out);
2691 }
2692 SmoothBasisSpec::BySmooth { smooth, by_kind } => {
2693 if let ByVarKind::Numeric { feature_col } = by_kind
2694 && *feature_col < n_training_headers
2695 {
2696 out.insert(*feature_col);
2697 }
2698 collect_by_variable_numeric_axes(smooth, n_training_headers, out);
2699 }
2700 SmoothBasisSpec::FactorSumToZero { inner, .. } => {
2701 collect_by_variable_numeric_axes(inner, n_training_headers, out);
2702 }
2703 _ => {}
2704 }
2705}
2706
2707impl FittedModel {
2708 pub fn axis_clip_to_training_ranges(
2717 &self,
2718 data: ndarray::ArrayView2<'_, f64>,
2719 col_map: &std::collections::HashMap<String, usize>,
2720 ) -> Option<ndarray::Array2<f64>> {
2721 let training_headers = self.training_headers.as_ref()?;
2722 let ranges = self.training_feature_ranges.as_ref()?;
2723 if training_headers.len() != ranges.len() {
2724 return None;
2725 }
2726 let mut kind_by_header: std::collections::HashMap<&str, ColumnKindTag> =
2727 std::collections::HashMap::new();
2728 if let Some(schema) = self.data_schema.as_ref() {
2729 for col in &schema.columns {
2730 kind_by_header.insert(col.name.as_str(), col.kind);
2731 }
2732 }
2733 let periodic_axes = self.training_periodic_axes(training_headers);
2739 let linear_axes = self.training_linear_axes(training_headers.len());
2747 let random_effect_axes = self.training_random_effect_axes(training_headers.len());
2752 let smooth_extrapolation_axes =
2760 self.training_smooth_extrapolation_axes(training_headers.len());
2761 let by_variable_axes = self.training_by_variable_numeric_axes(training_headers.len());
2767 let sphere_lat_bounds = self.training_sphere_latitude_bounds(training_headers);
2772 let mut clipped = data.to_owned();
2773 let mut any_clipped = false;
2774 for (col_in_training, (header, &(lo, hi))) in
2775 training_headers.iter().zip(ranges.iter()).enumerate()
2776 {
2777 let (lo, hi) = sphere_lat_bounds
2778 .get(&col_in_training)
2779 .copied()
2780 .unwrap_or((lo, hi));
2781 if !(lo.is_finite() && hi.is_finite()) || hi <= lo {
2782 continue;
2783 }
2784 if !matches!(
2785 kind_by_header.get(header.as_str()).copied(),
2786 Some(ColumnKindTag::Continuous)
2787 ) {
2788 continue;
2789 }
2790 if periodic_axes.contains(&col_in_training) {
2791 continue;
2792 }
2793 if linear_axes.contains(&col_in_training) {
2794 continue;
2795 }
2796 if random_effect_axes.contains(&col_in_training) {
2797 continue;
2798 }
2799 if smooth_extrapolation_axes.contains(&col_in_training) {
2800 continue;
2801 }
2802 if by_variable_axes.contains(&col_in_training) {
2803 continue;
2804 }
2805 let Some(&col_idx) = col_map.get(header) else {
2806 continue;
2807 };
2808 if col_idx >= clipped.ncols() {
2809 continue;
2810 }
2811 let mut col = clipped.column_mut(col_idx);
2812 for v in col.iter_mut() {
2813 if v.is_finite() {
2814 if *v < lo {
2815 *v = lo;
2816 any_clipped = true;
2817 } else if *v > hi {
2818 *v = hi;
2819 any_clipped = true;
2820 }
2821 }
2822 }
2823 }
2824 if any_clipped { Some(clipped) } else { None }
2825 }
2826
2827 fn saved_term_specs(&self) -> Vec<&TermCollectionSpec> {
2828 let mut specs: Vec<&TermCollectionSpec> = [
2829 self.resolved_termspec.as_ref(),
2830 self.resolved_termspec_noise.as_ref(),
2831 self.resolved_termspec_logslope.as_ref(),
2832 ]
2833 .into_iter()
2834 .flatten()
2835 .collect();
2836 if let Some(logslopes) = self.resolved_termspec_logslopes.as_ref() {
2837 specs.extend(logslopes.iter());
2838 }
2839 specs
2840 }
2841
2842 fn training_periodic_axes(
2849 &self,
2850 training_headers: &[String],
2851 ) -> std::collections::HashSet<usize> {
2852 use gam_terms::basis::BSplineKnotSpec;
2853 use gam_terms::smooth::SmoothBasisSpec;
2854 let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2855 let Some(spec) = self.resolved_termspec.as_ref() else {
2856 return out;
2857 };
2858 for term in &spec.smooth_terms {
2859 match &term.basis {
2860 SmoothBasisSpec::Sphere { feature_cols, .. } => {
2866 if let Some(&lon_col) = feature_cols.get(1)
2867 && lon_col < training_headers.len()
2868 {
2869 out.insert(lon_col);
2870 }
2871 }
2872 SmoothBasisSpec::BSpline1D { feature_col, spec } => {
2874 if matches!(spec.knotspec, BSplineKnotSpec::PeriodicUniform { .. })
2875 && *feature_col < training_headers.len()
2876 {
2877 out.insert(*feature_col);
2878 }
2879 }
2880 SmoothBasisSpec::TensorBSpline { feature_cols, spec } => {
2883 for (i, marginal) in spec.marginalspecs.iter().enumerate() {
2884 if matches!(marginal.knotspec, BSplineKnotSpec::PeriodicUniform { .. })
2885 && let Some(&col) = feature_cols.get(i)
2886 && col < training_headers.len()
2887 {
2888 out.insert(col);
2889 }
2890 }
2891 }
2892 _ => {}
2893 }
2894 }
2895 out
2896 }
2897
2898 fn training_linear_axes(&self, n_training_headers: usize) -> std::collections::HashSet<usize> {
2909 let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2910 for spec in self.saved_term_specs() {
2911 for term in &spec.linear_terms {
2912 for col in term.effective_feature_cols() {
2913 if col < n_training_headers {
2914 out.insert(col);
2915 }
2916 }
2917 }
2918 }
2919 out
2920 }
2921
2922 fn training_random_effect_axes(
2928 &self,
2929 n_training_headers: usize,
2930 ) -> std::collections::HashSet<usize> {
2931 let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2932 for spec in self.saved_term_specs() {
2933 for term in &spec.random_effect_terms {
2934 if term.feature_col < n_training_headers {
2935 out.insert(term.feature_col);
2936 }
2937 }
2938 }
2939 out
2940 }
2941
2942 fn training_smooth_extrapolation_axes(
2975 &self,
2976 n_training_headers: usize,
2977 ) -> std::collections::HashSet<usize> {
2978 let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2979 for spec in self.saved_term_specs() {
2980 for term in &spec.smooth_terms {
2981 collect_smooth_extrapolation_axes(&term.basis, n_training_headers, &mut out);
2982 }
2983 }
2984 out
2985 }
2986
2987 fn training_by_variable_numeric_axes(
2993 &self,
2994 n_training_headers: usize,
2995 ) -> std::collections::HashSet<usize> {
2996 let mut out: std::collections::HashSet<usize> = std::collections::HashSet::new();
2997 for spec in self.saved_term_specs() {
2998 for term in &spec.smooth_terms {
2999 collect_by_variable_numeric_axes(&term.basis, n_training_headers, &mut out);
3000 }
3001 }
3002 out
3003 }
3004
3005 fn training_sphere_latitude_bounds(
3027 &self,
3028 training_headers: &[String],
3029 ) -> std::collections::HashMap<usize, (f64, f64)> {
3030 use gam_terms::smooth::SmoothBasisSpec;
3031 let mut out: std::collections::HashMap<usize, (f64, f64)> =
3032 std::collections::HashMap::new();
3033 let Some(spec) = self.resolved_termspec.as_ref() else {
3034 return out;
3035 };
3036 for term in &spec.smooth_terms {
3037 if let SmoothBasisSpec::Sphere { feature_cols, spec } = &term.basis
3038 && let Some(&lat_col) = feature_cols.first()
3039 && lat_col < training_headers.len()
3040 {
3041 let bound = if spec.radians {
3042 std::f64::consts::FRAC_PI_2
3043 } else {
3044 90.0
3045 };
3046 out.insert(lat_col, (-bound, bound));
3047 }
3048 }
3049 out
3050 }
3051
3052 pub fn from_payload(mut payload: FittedModelPayload) -> Self {
3053 let likelihood = payload.family_state.likelihood();
3054 let class = match payload.model_kind {
3055 ModelKind::Survival => PredictModelClass::Survival,
3056 ModelKind::MarginalSlope => PredictModelClass::BernoulliMarginalSlope,
3057 ModelKind::TransformationNormal => PredictModelClass::TransformationNormal,
3058 ModelKind::LocationScale => {
3059 if likelihood == LikelihoodSpec::gaussian_identity() {
3060 PredictModelClass::GaussianLocationScale
3061 } else if is_dispersion_location_scale_response(&likelihood.response) {
3062 PredictModelClass::DispersionLocationScale
3063 } else {
3064 PredictModelClass::BinomialLocationScale
3065 }
3066 }
3067 ModelKind::Standard => PredictModelClass::Standard,
3068 };
3069 match class {
3070 PredictModelClass::Survival => {
3071 payload.model_kind = ModelKind::Survival;
3072 Self::Survival { payload }
3073 }
3074 PredictModelClass::BernoulliMarginalSlope => {
3075 payload.model_kind = ModelKind::MarginalSlope;
3076 Self::MarginalSlope { payload }
3077 }
3078 PredictModelClass::TransformationNormal => {
3079 payload.model_kind = ModelKind::TransformationNormal;
3080 Self::TransformationNormal { payload }
3081 }
3082 PredictModelClass::GaussianLocationScale
3083 | PredictModelClass::BinomialLocationScale
3084 | PredictModelClass::DispersionLocationScale => {
3085 payload.model_kind = ModelKind::LocationScale;
3086 Self::LocationScale { payload }
3087 }
3088 PredictModelClass::Standard => {
3089 payload.model_kind = ModelKind::Standard;
3090 Self::Standard { payload }
3091 }
3092 }
3093 .with_synchronized_stateful_link_metadata()
3094 }
3095
3096 #[inline]
3097 pub fn payload(&self) -> &FittedModelPayload {
3098 match self {
3099 Self::Standard { payload }
3100 | Self::LocationScale { payload }
3101 | Self::MarginalSlope { payload }
3102 | Self::Survival { payload }
3103 | Self::TransformationNormal { payload } => payload,
3104 }
3105 }
3106
3107 #[inline]
3108 fn payload_mut(&mut self) -> &mut FittedModelPayload {
3109 match self {
3110 Self::Standard { payload }
3111 | Self::LocationScale { payload }
3112 | Self::MarginalSlope { payload }
3113 | Self::Survival { payload }
3114 | Self::TransformationNormal { payload } => payload,
3115 }
3116 }
3117
3118 fn with_synchronized_stateful_link_metadata(mut self) -> Self {
3119 self.synchronize_stateful_link_metadata();
3120 self
3121 }
3122
3123 fn synchronize_stateful_link_metadata(&mut self) {
3124 let payload = self.payload_mut();
3125 match (payload.fit_result.is_none(), payload.unified.is_none()) {
3137 (true, false) => payload.fit_result = payload.unified.clone(),
3138 (false, true) => payload.unified = payload.fit_result.clone(),
3139 _ => {}
3140 }
3141 payload.used_device = payload
3142 .fit_result
3143 .as_ref()
3144 .or(payload.unified.as_ref())
3145 .is_some_and(|fit| fit.used_device);
3146 payload.synchronize_empty_feature_contract();
3147 let Some(fit) = payload.fit_result.as_ref().or(payload.unified.as_ref()) else {
3148 return;
3149 };
3150 match (&mut payload.family_state, &fit.fitted_link) {
3151 (
3152 FittedFamily::Standard {
3153 likelihood,
3154 latent_cloglog_state,
3155 ..
3156 },
3157 FittedLinkState::LatentCLogLog { state },
3158 ) if likelihood.is_latent_cloglog() => {
3159 *latent_cloglog_state = Some(*state);
3160 }
3161 (
3162 FittedFamily::Standard {
3163 likelihood,
3164 sas_state,
3165 ..
3166 },
3167 FittedLinkState::Sas { state, covariance },
3168 ) if likelihood.is_binomial_sas() => {
3169 *sas_state = Some(*state);
3170 payload.sas_param_covariance = covariance.as_ref().map(array2_to_nested_vec);
3171 }
3172 (
3173 FittedFamily::Standard {
3174 likelihood,
3175 sas_state,
3176 ..
3177 },
3178 FittedLinkState::BetaLogistic { state, covariance },
3179 ) if likelihood.is_binomial_beta_logistic() => {
3180 *sas_state = Some(*state);
3181 payload.sas_param_covariance = covariance.as_ref().map(array2_to_nested_vec);
3182 }
3183 (
3184 FittedFamily::Standard {
3185 likelihood,
3186 mixture_state,
3187 ..
3188 },
3189 FittedLinkState::Mixture { state, covariance },
3190 ) if likelihood.is_binomial_mixture() => {
3191 *mixture_state = Some(state.clone());
3192 payload.mixture_link_param_covariance =
3193 covariance.as_ref().map(array2_to_nested_vec);
3194 }
3195 _ => {}
3196 }
3197 }
3198
3199 #[inline]
3200 pub fn likelihood(&self) -> LikelihoodSpec {
3201 self.payload().family_state.likelihood()
3202 }
3203
3204 #[inline]
3205 pub fn estimator(&self) -> FittedEstimator {
3206 self.payload().estimator
3207 }
3208
3209 pub fn prediction_required_columns(
3227 &self,
3228 ) -> Result<std::collections::BTreeSet<String>, String> {
3229 let payload = self.payload();
3230 let parsed = parse_formula(payload.formula.as_str()).map_err(|e| e.to_string())?;
3231 let mut required = std::collections::BTreeSet::<String>::new();
3232 parsed_term_column_names(&parsed.terms, &mut required);
3233
3234 if let Some((entry, exit, _event)) =
3235 parse_surv_response(parsed.response.as_str()).map_err(|e| e.to_string())?
3236 {
3237 if let Some(entry) = entry {
3238 required.insert(entry);
3239 }
3240 required.insert(exit);
3241 } else if let Some((left, right, _event)) =
3242 parse_surv_interval_response(parsed.response.as_str()).map_err(|e| e.to_string())?
3243 {
3244 required.insert(left);
3245 required.insert(right);
3246 }
3247 if let Some(offset) = payload.offset_column.as_ref() {
3255 required.insert(offset.clone());
3256 }
3257 if let Some(noise_offset) = payload.noise_offset_column.as_ref() {
3258 required.insert(noise_offset.clone());
3259 }
3260 if matches!(
3261 self.predict_model_class(),
3262 PredictModelClass::BernoulliMarginalSlope | PredictModelClass::Survival
3263 ) {
3264 if let Some(z_column) = payload.z_column.as_ref() {
3265 required.remove("z");
3266 required.insert(z_column.clone());
3267 }
3268 }
3269 if let Some(noise_formula) = payload.formula_noise.as_ref() {
3270 self.add_auxiliary_formula_columns(
3271 &mut required,
3272 noise_formula,
3273 parsed.response.as_str(),
3274 )?;
3275 }
3276 if let Some(logslope_formula) = payload.formula_logslope.as_ref() {
3277 if logslope_formula != "same-as-main" {
3278 self.add_auxiliary_formula_columns(
3279 &mut required,
3280 logslope_formula,
3281 parsed.response.as_str(),
3282 )?;
3283 }
3284 }
3285 Ok(required)
3286 }
3287
3288 pub fn diagnostic_extra_columns(&self) -> Result<Vec<String>, String> {
3305 let payload = self.payload();
3306 let parsed = parse_formula(payload.formula.as_str()).map_err(|e| e.to_string())?;
3307 let mut extras: Vec<String> = Vec::new();
3317 if let Some(weight_column) = payload.weight_column.as_ref() {
3318 extras.push(weight_column.clone());
3319 }
3320 if parse_surv_response(parsed.response.as_str())
3323 .map_err(|e| e.to_string())?
3324 .is_some()
3325 || parse_surv_interval_response(parsed.response.as_str())
3326 .map_err(|e| e.to_string())?
3327 .is_some()
3328 {
3329 return Ok(extras);
3330 }
3331 let response = parsed.response.trim();
3332 if response.is_empty() || response.contains('(') {
3335 return Ok(extras);
3336 }
3337 if self.prediction_required_columns()?.contains(response) {
3340 return Ok(extras);
3341 }
3342 extras.push(response.to_string());
3343 Ok(extras)
3344 }
3345
3346 fn add_auxiliary_formula_columns(
3349 &self,
3350 required: &mut std::collections::BTreeSet<String>,
3351 formula_or_rhs: &str,
3352 response: &str,
3353 ) -> Result<(), String> {
3354 let trimmed = formula_or_rhs.trim();
3355 if trimmed.is_empty() || trimmed == "1" {
3356 return Ok(());
3357 }
3358 let formula = if trimmed.contains('~') {
3359 trimmed.to_string()
3360 } else {
3361 format!("{response} ~ {trimmed}")
3362 };
3363 let parsed = parse_formula(formula.as_str()).map_err(|e| e.to_string())?;
3364 parsed_term_column_names(&parsed.terms, required);
3365 Ok(())
3366 }
3367
3368 #[inline]
3369 pub fn predict_model_class(&self) -> PredictModelClass {
3370 match &self.payload().family_state {
3371 FittedFamily::Survival { .. }
3372 | FittedFamily::LatentSurvival { .. }
3373 | FittedFamily::LatentBinary { .. } => PredictModelClass::Survival,
3374 FittedFamily::MarginalSlope { .. } => PredictModelClass::BernoulliMarginalSlope,
3375 FittedFamily::TransformationNormal { .. } => PredictModelClass::TransformationNormal,
3376 FittedFamily::LocationScale { likelihood, .. } if likelihood.is_gaussian_identity() => {
3377 PredictModelClass::GaussianLocationScale
3378 }
3379 FittedFamily::LocationScale { likelihood, .. }
3380 if is_dispersion_location_scale_response(&likelihood.response) =>
3381 {
3382 PredictModelClass::DispersionLocationScale
3383 }
3384 FittedFamily::LocationScale { .. } => PredictModelClass::BinomialLocationScale,
3385 FittedFamily::Standard { .. } => PredictModelClass::Standard,
3386 }
3387 }
3388
3389 pub fn saved_link_wiggle(&self) -> Result<Option<SavedLinkWiggleRuntime>, FittedModelError> {
3390 let payload = self.payload();
3391 let (knots, degree) = match (
3392 payload.linkwiggle_knots.as_ref(),
3393 payload.linkwiggle_degree,
3394 ) {
3395 (None, None) => return Ok(None),
3396 (Some(knots), Some(degree)) => (knots.clone(), degree),
3397 _ => {
3398 return Err(FittedModelError::SchemaMismatch {
3399 reason:
3400 "saved model has partial link-wiggle metadata; expected linkwiggle_knots and linkwiggle_degree together"
3401 .to_string(),
3402 })
3403 }
3404 };
3405 let resolved_link = self.resolved_inverse_link()?;
3406 let saved_link_disallows_wiggle = resolved_link
3407 .as_ref()
3408 .is_some_and(|link| !inverse_link_supports_joint_wiggle(link))
3409 || payload
3410 .link
3411 .as_ref()
3412 .is_some_and(|link| !inverse_link_supports_joint_wiggle(link));
3413 if saved_link_disallows_wiggle {
3414 return Err(FittedModelError::IncompatibleConfig {
3415 reason: joint_wiggle_unsupported_link_message("link wiggle"),
3416 });
3417 }
3418 let model_class = self.predict_model_class();
3419 let beta = match model_class {
3420 PredictModelClass::Standard => {
3429 let fit = payload.fit_result.as_ref().ok_or_else(|| {
3430 FittedModelError::MissingField {
3431 reason:
3432 "standard link-wiggle model is missing canonical fit_result payload"
3433 .to_string(),
3434 }
3435 })?;
3436 if fit.blocks.len() != 2
3437 || fit.blocks[0].role != BlockRole::Mean
3438 || fit.blocks[1].role != BlockRole::LinkWiggle
3439 {
3440 return Err(FittedModelError::SchemaMismatch {
3441 reason:
3442 "standard link-wiggle models must store blocks in [Mean, LinkWiggle] order"
3443 .to_string(),
3444 });
3445 }
3446 let block = fit.block_by_role(BlockRole::LinkWiggle).ok_or_else(|| {
3447 FittedModelError::MissingField {
3448 reason:
3449 "standard link-wiggle model is missing LinkWiggle coefficient block"
3450 .to_string(),
3451 }
3452 })?;
3453 let payload_beta = payload.beta_link_wiggle.as_ref().ok_or_else(|| {
3454 FittedModelError::MissingField {
3455 reason: "standard link-wiggle model is missing its exact saved prediction coefficients; refit"
3456 .to_string(),
3457 }
3458 })?;
3459 if payload_beta.len() != block.beta.len()
3460 || payload_beta
3461 .iter()
3462 .zip(block.beta.iter())
3463 .any(|(saved, fitted)| saved.to_bits() != fitted.to_bits())
3464 {
3465 return Err(FittedModelError::SchemaMismatch {
3466 reason: "standard link-wiggle payload coefficients disagree with the fitted LinkWiggle block"
3467 .to_string(),
3468 });
3469 }
3470 let shift = payload.link_wiggle_index_shift.as_ref().ok_or_else(|| {
3471 FittedModelError::MissingField {
3472 reason: "standard link-wiggle model is missing its frozen-index shift; refit"
3473 .to_string(),
3474 }
3475 })?;
3476 if shift.len() != fit.blocks[0].beta.len() {
3477 return Err(FittedModelError::SchemaMismatch {
3478 reason: format!(
3479 "standard link-wiggle frozen-index shift has {} entries but the Mean block has {} coefficients",
3480 shift.len(),
3481 fit.blocks[0].beta.len(),
3482 ),
3483 });
3484 }
3485 block.beta.to_vec()
3486 }
3487 _ => payload
3488 .beta_link_wiggle
3489 .clone()
3490 .ok_or_else(|| FittedModelError::MissingField {
3491 reason:
3492 "saved model has link-wiggle metadata but is missing payload.beta_link_wiggle"
3493 .to_string(),
3494 })?,
3495 };
3496 let penalty_metadata = payload.linkwiggle_penalty_metadata.clone();
3497 if let Some(metadata) = penalty_metadata.as_ref() {
3498 let canonical = canonical_wiggle_function_penalties(
3499 &Array1::from_vec(knots.clone()),
3500 degree,
3501 &metadata.derivative_orders,
3502 metadata.double_penalty,
3503 )
3504 .map_err(|reason| FittedModelError::PayloadCorrupt {
3505 reason: format!("saved link-wiggle penalty metadata is invalid: {reason}"),
3506 })?;
3507 if canonical.metadata != *metadata {
3508 return Err(FittedModelError::SchemaMismatch {
3509 reason: format!(
3510 "saved link-wiggle penalty topology {:?} disagrees with canonical topology {:?}",
3511 metadata.blocks, canonical.metadata.blocks,
3512 ),
3513 });
3514 }
3515 }
3516 let index_shift = payload.link_wiggle_index_shift.clone();
3521 Ok(Some(SavedLinkWiggleRuntime {
3522 knots,
3523 degree,
3524 penalty_metadata,
3525 beta,
3526 index_shift,
3527 }))
3528 }
3529
3530 pub fn saved_baseline_time_wiggle(
3531 &self,
3532 ) -> Result<Option<SavedBaselineTimeWiggleRuntime>, FittedModelError> {
3533 let payload = self.payload();
3534 if payload
3535 .survival_cause_count
3536 .is_some_and(|cause_count| cause_count > 1)
3537 && payload.beta_baseline_timewiggle.is_none()
3538 && payload.beta_baseline_timewiggle_by_cause.is_some()
3539 {
3540 return Err(FittedModelError::SchemaMismatch {
3541 reason:
3542 "joint cause-specific survival stores baseline-timewiggle coefficients per cause"
3543 .to_string(),
3544 });
3545 }
3546 match (
3547 payload.baseline_timewiggle_knots.as_ref(),
3548 payload.baseline_timewiggle_degree,
3549 payload.baseline_timewiggle_penalty_orders.as_ref(),
3550 payload.baseline_timewiggle_double_penalty,
3551 payload.beta_baseline_timewiggle.as_ref(),
3552 ) {
3553 (None, None, None, None, None) => Ok(None),
3554 (Some(knots), Some(degree), Some(penalty_orders), Some(double_penalty), Some(beta)) => {
3555 Ok(Some(SavedBaselineTimeWiggleRuntime {
3556 knots: knots.clone(),
3557 degree,
3558 penalty_orders: penalty_orders.clone(),
3559 double_penalty,
3560 beta: beta.clone(),
3561 }))
3562 }
3563 _ => Err(FittedModelError::SchemaMismatch {
3564 reason:
3565 "saved model has partial baseline-timewiggle metadata; expected knots+degree+penalty_order+double_penalty+beta_baseline_timewiggle together"
3566 .to_string(),
3567 }),
3568 }
3569 }
3570
3571 #[inline]
3573 pub fn has_link_wiggle(&self) -> bool {
3574 self.saved_link_wiggle()
3575 .map(|runtime| runtime.is_some())
3576 .unwrap_or(false)
3577 }
3578
3579 #[inline]
3581 pub fn has_baseline_time_wiggle(&self) -> bool {
3582 let payload = self.payload();
3583 if payload
3584 .survival_cause_count
3585 .is_some_and(|cause_count| cause_count > 1)
3586 {
3587 return payload.baseline_timewiggle_knots.is_some()
3588 && payload.baseline_timewiggle_degree.is_some()
3589 && payload.baseline_timewiggle_penalty_orders.is_some()
3590 && payload.baseline_timewiggle_double_penalty.is_some()
3591 && payload.beta_baseline_timewiggle_by_cause.is_some();
3592 }
3593 self.saved_baseline_time_wiggle()
3594 .map(|runtime| runtime.is_some())
3595 .unwrap_or(false)
3596 }
3597
3598 #[inline]
3623 pub fn prediction_uses_posterior_mean(&self) -> bool {
3624 let family = self.likelihood();
3625 let curved_family = match &family.response {
3626 ResponseFamily::Gaussian => false,
3629 ResponseFamily::Poisson
3631 | ResponseFamily::Gamma
3632 | ResponseFamily::Tweedie { .. }
3633 | ResponseFamily::NegativeBinomial { .. } => true,
3634 ResponseFamily::Beta { .. } => true,
3636 ResponseFamily::RoystonParmar => true,
3638 ResponseFamily::Binomial => matches!(
3641 &family.link,
3642 InverseLink::Standard(_)
3643 | InverseLink::Sas(_)
3644 | InverseLink::BetaLogistic(_)
3645 | InverseLink::Mixture(_)
3646 | InverseLink::LatentCLogLog(_)
3647 ),
3648 };
3649 curved_family || self.has_link_wiggle() || self.has_baseline_time_wiggle()
3650 }
3651
3652 pub fn saved_prediction_runtime(&self) -> Result<SavedPredictionRuntime, FittedModelError> {
3653 self.payload().validate_payload_version()?;
3654 if matches!(
3655 self.predict_model_class(),
3656 PredictModelClass::BernoulliMarginalSlope | PredictModelClass::Survival
3657 ) {
3658 if let Some(runtime) = self.payload().score_warp_runtime.as_ref() {
3659 runtime.validate_exact_replay_contract().map_err(|err| {
3660 FittedModelError::PayloadCorrupt {
3661 reason: format!("saved anchored score-warp runtime is invalid: {err}"),
3662 }
3663 })?;
3664 }
3665 if let Some(runtime) = self.payload().link_deviation_runtime.as_ref() {
3666 runtime.validate_exact_replay_contract().map_err(|err| {
3667 FittedModelError::PayloadCorrupt {
3668 reason: format!("saved anchored link-deviation runtime is invalid: {err}"),
3669 }
3670 })?;
3671 }
3672 }
3673 let runtime = SavedPredictionRuntime {
3674 model_class: self.predict_model_class(),
3675 likelihood: self.likelihood(),
3676 inverse_link: self.resolved_inverse_link()?,
3677 link_wiggle: self.saved_link_wiggle()?,
3678 baseline_time_wiggle: self.saved_baseline_time_wiggle()?,
3679 score_warp: self.payload().score_warp_runtime.clone(),
3680 link_deviation: self.payload().link_deviation_runtime.clone(),
3681 latent_z_rank_int_calibration: self.payload().latent_z_rank_int_calibration.clone(),
3682 latent_z_conditional_calibration: self
3683 .payload()
3684 .latent_z_conditional_calibration
3685 .clone(),
3686 influence_absorber_width: self.payload().influence_absorber_width,
3687 };
3688 if matches!(
3689 runtime.model_class,
3690 PredictModelClass::GaussianLocationScale
3691 | PredictModelClass::BinomialLocationScale
3692 | PredictModelClass::DispersionLocationScale
3693 ) {
3694 let fit = self.payload().fit_result.as_ref().ok_or_else(|| {
3695 FittedModelError::MissingField {
3696 reason: "location-scale model is missing canonical fit_result payload"
3697 .to_string(),
3698 }
3699 })?;
3700 validate_location_scale_saved_fit(
3701 fit,
3702 runtime.model_class,
3703 runtime.link_wiggle.as_ref(),
3704 )?;
3705 } else if matches!(runtime.model_class, PredictModelClass::Survival)
3706 && self
3707 .payload()
3708 .survival_likelihood
3709 .as_deref()
3710 .is_some_and(|value| value.eq_ignore_ascii_case("location-scale"))
3711 {
3712 validate_survival_location_scale_saved_fit(
3713 self.payload(),
3714 runtime.link_wiggle.as_ref(),
3715 )?;
3716 } else if matches!(
3717 runtime.model_class,
3718 PredictModelClass::BernoulliMarginalSlope
3719 ) {
3720 let unified =
3721 self.payload()
3722 .unified
3723 .as_ref()
3724 .ok_or_else(|| FittedModelError::MissingField {
3725 reason: "marginal-slope model is missing unified fit payload; refit"
3726 .to_string(),
3727 })?;
3728 validate_marginal_slope_saved_fit(
3729 unified,
3730 runtime.score_warp.as_ref(),
3731 runtime.link_deviation.as_ref(),
3732 "unified",
3733 )?;
3734 } else if matches!(runtime.model_class, PredictModelClass::Survival)
3735 && self
3736 .payload()
3737 .survival_likelihood
3738 .as_deref()
3739 .is_some_and(|value| value.eq_ignore_ascii_case("marginal-slope"))
3740 {
3741 let fit = self.payload().fit_result.as_ref().ok_or_else(|| {
3742 FittedModelError::MissingField {
3743 reason: "survival marginal-slope model is missing canonical fit_result payload"
3744 .to_string(),
3745 }
3746 })?;
3747 validate_survival_marginal_slope_saved_fit(self.payload(), fit, "fit_result")?;
3748 }
3749 Ok(runtime)
3750 }
3751
3752 pub fn saved_sas_state(&self) -> Result<Option<SasLinkState>, FittedModelError> {
3753 let payload = self.payload();
3754 let raw = match &payload.family_state {
3755 FittedFamily::Standard {
3756 likelihood,
3757 sas_state,
3758 ..
3759 } if likelihood.is_binomial_sas() => {
3760 (*sas_state).ok_or_else(|| FittedModelError::MissingField {
3761 reason: "binomial-sas model is missing state in family_state.sas_state"
3762 .to_string(),
3763 })?
3764 }
3765 FittedFamily::LocationScale {
3766 likelihood,
3767 base_link,
3768 } if likelihood.is_binomial_sas() => match base_link {
3769 Some(InverseLink::Sas(state)) => *state,
3770 _ => {
3771 return Err(FittedModelError::MissingField {
3772 reason: "binomial-sas location-scale model is missing SAS base_link state"
3773 .to_string(),
3774 });
3775 }
3776 },
3777 _ => return Ok(None),
3778 };
3779 state_from_sasspec(SasLinkSpec {
3780 initial_epsilon: raw.epsilon,
3781 initial_log_delta: raw.log_delta,
3782 })
3783 .map(Some)
3784 .map_err(|e| FittedModelError::PayloadCorrupt {
3785 reason: format!("invalid saved SAS link state: {e}"),
3786 })
3787 }
3788
3789 pub fn saved_beta_logistic_state(&self) -> Result<Option<SasLinkState>, FittedModelError> {
3790 let payload = self.payload();
3791 let raw = match &payload.family_state {
3792 FittedFamily::Standard {
3793 likelihood,
3794 sas_state,
3795 ..
3796 } if likelihood.is_binomial_beta_logistic() => {
3797 (*sas_state).ok_or_else(|| FittedModelError::MissingField {
3798 reason:
3799 "binomial-beta-logistic model is missing state in family_state.sas_state"
3800 .to_string(),
3801 })?
3802 }
3803 FittedFamily::LocationScale {
3804 likelihood,
3805 base_link,
3806 } if likelihood.is_binomial_beta_logistic() => match base_link {
3807 Some(InverseLink::BetaLogistic(state)) => *state,
3808 _ => {
3809 return Err(FittedModelError::MissingField {
3810 reason:
3811 "binomial-beta-logistic location-scale model is missing beta-logistic base_link state"
3812 .to_string(),
3813 });
3814 }
3815 },
3816 _ => return Ok(None),
3817 };
3818 state_from_beta_logisticspec(SasLinkSpec {
3819 initial_epsilon: raw.epsilon,
3820 initial_log_delta: raw.log_delta,
3821 })
3822 .map(Some)
3823 .map_err(|e| FittedModelError::PayloadCorrupt {
3824 reason: format!("invalid saved Beta-Logistic link state: {e}"),
3825 })
3826 }
3827
3828 pub fn saved_mixture_state(&self) -> Result<Option<MixtureLinkState>, FittedModelError> {
3829 let payload = self.payload();
3830 match &payload.family_state {
3831 FittedFamily::Standard {
3832 likelihood,
3833 mixture_state,
3834 ..
3835 } if likelihood.is_binomial_mixture() => mixture_state
3836 .clone()
3837 .ok_or_else(|| FittedModelError::MissingField {
3838 reason: "binomial-mixture model is missing state in family_state.mixture_state"
3839 .to_string(),
3840 })
3841 .map(Some),
3842 FittedFamily::LocationScale {
3843 likelihood,
3844 base_link,
3845 } if likelihood.is_binomial_mixture() => match base_link {
3846 Some(InverseLink::Mixture(state)) => Ok(Some(state.clone())),
3847 _ => Err(FittedModelError::MissingField {
3848 reason:
3849 "binomial-mixture location-scale model is missing mixture base_link state"
3850 .to_string(),
3851 }),
3852 },
3853 _ => Ok(None),
3854 }
3855 }
3856
3857 pub fn saved_latent_cloglog_state(
3858 &self,
3859 ) -> Result<Option<LatentCLogLogState>, FittedModelError> {
3860 let payload = self.payload();
3861 match &payload.family_state {
3862 FittedFamily::Standard {
3863 likelihood,
3864 latent_cloglog_state,
3865 ..
3866 } if likelihood.is_latent_cloglog() => latent_cloglog_state
3867 .ok_or_else(|| FittedModelError::MissingField {
3868 reason:
3869 "latent-cloglog-binomial model is missing state in family_state.latent_cloglog_state"
3870 .to_string(),
3871 })
3872 .map(Some),
3873 _ => Ok(None),
3874 }
3875 }
3876
3877 pub fn resolved_inverse_link(&self) -> Result<Option<InverseLink>, FittedModelError> {
3878 let stateful = if let Some(state) = self.saved_mixture_state()? {
3879 Some(InverseLink::Mixture(state))
3880 } else if let Some(state) = self.saved_latent_cloglog_state()? {
3881 Some(InverseLink::LatentCLogLog(state))
3882 } else if let Some(state) = self.saved_beta_logistic_state()? {
3883 Some(InverseLink::BetaLogistic(state))
3884 } else {
3885 self.saved_sas_state()?.map(InverseLink::Sas)
3886 };
3887 match &self.payload().family_state {
3888 FittedFamily::LocationScale { base_link, .. } => Ok(base_link.clone().or(stateful)),
3889 FittedFamily::Standard { link, .. } => {
3890 Ok(stateful.or_else(|| link.map(InverseLink::Standard)))
3891 }
3892 FittedFamily::MarginalSlope { base_link, .. } => Ok(Some(base_link.clone())),
3893 FittedFamily::Survival { .. }
3894 | FittedFamily::LatentSurvival { .. }
3895 | FittedFamily::LatentBinary { .. } => Ok(None),
3896 FittedFamily::TransformationNormal { .. } => Ok(None),
3897 }
3898 }
3899
3900 const MEASURE_JET_COVERAGE_FLOOR: f64 = 0.05;
3909
3910 pub fn measure_jet_extrapolation_variance(
3937 &self,
3938 data: ndarray::ArrayView2<'_, f64>,
3939 col_map: &HashMap<String, usize>,
3940 ) -> Result<Option<Array1<f64>>, FittedModelError> {
3941 use gam_terms::basis::{
3942 CenterStrategy, MeasureJetExtrapolationSpectrum, MeasureJetIdentifiability,
3943 PenaltySource,
3944 };
3945 use gam_terms::smooth::SmoothBasisSpec;
3946 use gam_terms::smooth::build_term_collection_design;
3947 let Some(saved_spec) = self.resolved_termspec.as_ref() else {
3948 return Ok(None);
3949 };
3950 if data.nrows() == 0
3951 || !saved_spec
3952 .smooth_terms
3953 .iter()
3954 .any(|t| matches!(t.basis, SmoothBasisSpec::MeasureJet { .. }))
3955 {
3956 return Ok(None);
3957 }
3958 let fit = self
3959 .fit_result
3960 .as_ref()
3961 .ok_or_else(|| FittedModelError::MissingField {
3962 reason: "measure-jet extrapolation variance requires the canonical \
3963 fit_result payload; refit"
3964 .to_string(),
3965 })?;
3966 let spec = crate::survival::predict::resolve_termspec_for_prediction(
3967 &self.resolved_termspec,
3968 self.training_headers.as_ref(),
3969 col_map,
3970 "resolved_termspec",
3971 )
3972 .map_err(|e| FittedModelError::SchemaMismatch {
3973 reason: format!("measure-jet extrapolation variance: {e}"),
3974 })?;
3975 let probe = data.slice(ndarray::s![0..1, ..]);
3982 let design = build_term_collection_design(probe, &spec).map_err(|e| {
3983 FittedModelError::SchemaMismatch {
3984 reason: format!(
3985 "measure-jet extrapolation variance: penalty-layout replay failed: {e}"
3986 ),
3987 }
3988 })?;
3989 let lambdas = &fit.lambdas;
3990 let phi_scale = fit.coefficient_covariance_scale().map_err(|err| {
3995 FittedModelError::SchemaMismatch {
3996 reason: format!(
3997 "measure-jet extrapolation variance has no valid coefficient-covariance scale: {err}"
3998 ),
3999 }
4000 })?;
4001 let mut total = Array1::<f64>::zeros(data.nrows());
4002 let mut contributed = false;
4003 for (smooth_idx, term) in spec.smooth_terms.iter().enumerate() {
4004 let SmoothBasisSpec::MeasureJet {
4005 feature_cols,
4006 spec: mj,
4007 input_scale,
4008 } = &term.basis
4009 else {
4010 continue;
4011 };
4012 let (Some(frozen), CenterStrategy::UserProvided(centers)) =
4013 (mj.frozen_quadrature.as_ref(), &mj.center_strategy)
4014 else {
4015 log::warn!(
4016 "measure-jet term '{}' is not frozen (UserProvided centers + frozen \
4017 quadrature); skipping its extrapolation variance",
4018 term.name
4019 );
4020 continue;
4021 };
4022 let n_levels = frozen.eps_band.len();
4023 let read_lambda = |global_index: usize| -> Result<f64, FittedModelError> {
4030 lambdas
4031 .get(global_index)
4032 .copied()
4033 .ok_or_else(|| FittedModelError::SchemaMismatch {
4034 reason: format!(
4035 "measure-jet term '{}': penalty global index {global_index} out \
4036 of bounds for {} fitted lambdas",
4037 term.name,
4038 lambdas.len()
4039 ),
4040 })
4041 };
4042 let mut per_scale: Vec<(usize, f64)> = Vec::new();
4043 let mut fused: Option<f64> = None;
4044 for info in &design.penaltyinfo {
4045 if info.termname.as_deref() != Some(term.name.as_str()) {
4046 continue;
4047 }
4048 match &info.penalty.source {
4049 PenaltySource::Other(label) => {
4050 if let Some(level_txt) = label.strip_prefix("measure_jet_scale_") {
4051 let level: usize = level_txt.parse().map_err(|_| {
4052 FittedModelError::SchemaMismatch {
4053 reason: format!(
4054 "measure-jet term '{}': unparseable penalty label \
4055 '{label}'",
4056 term.name
4057 ),
4058 }
4059 })?;
4060 per_scale.push((level, read_lambda(info.global_index)?));
4061 }
4062 }
4063 PenaltySource::Primary => {
4064 fused = Some(read_lambda(info.global_index)?);
4065 }
4066 _ => {}
4067 }
4068 }
4069 let mut lambda_phys = Vec::with_capacity(n_levels);
4070 let spectrum = if per_scale.is_empty() {
4071 let Some(lam) = fused else {
4072 log::warn!(
4073 "measure-jet term '{}' has no fitted amplitude in the penalty \
4074 layout; skipping its extrapolation variance",
4075 term.name
4076 );
4077 continue;
4078 };
4079 let Some(c) = frozen.fused_penalty_normalization_scale else {
4080 log::warn!(
4081 "measure-jet term '{}' is missing the fused penalty normalization scale; \
4082 skipping its extrapolation variance",
4083 term.name
4084 );
4085 continue;
4086 };
4087 MeasureJetExtrapolationSpectrum::Fused(lam / c)
4088 } else {
4089 per_scale.sort_by_key(|&(level, _)| level);
4090 let levels_complete = per_scale.len() == n_levels
4091 && per_scale
4092 .iter()
4093 .enumerate()
4094 .all(|(i, &(level, _))| level == i);
4095 if !levels_complete {
4096 log::warn!(
4097 "measure-jet term '{}': {} fitted per-scale amplitudes for {} band \
4098 scales; skipping its extrapolation variance",
4099 term.name,
4100 per_scale.len(),
4101 n_levels
4102 );
4103 continue;
4104 }
4105 if frozen.penalty_normalization_scales.len() != n_levels {
4106 log::warn!(
4107 "measure-jet term '{}': {} frozen penalty normalization scales for {} \
4108 band scales; skipping its extrapolation variance",
4109 term.name,
4110 frozen.penalty_normalization_scales.len(),
4111 n_levels
4112 );
4113 continue;
4114 }
4115 lambda_phys.extend(
4116 per_scale
4117 .iter()
4118 .map(|&(level, lam)| lam / frozen.penalty_normalization_scales[level]),
4119 );
4120 MeasureJetExtrapolationSpectrum::PerLevel(&lambda_phys)
4121 };
4122 let mut queries = Array2::<f64>::zeros((data.nrows(), feature_cols.len()));
4127 for (j, &col) in feature_cols.iter().enumerate() {
4128 if col >= data.ncols() {
4129 return Err(FittedModelError::SchemaMismatch {
4130 reason: format!(
4131 "measure-jet term '{}': prediction column {col} out of bounds \
4132 for {} data columns",
4133 term.name,
4134 data.ncols()
4135 ),
4136 });
4137 }
4138 queries.column_mut(j).assign(&data.column(col));
4139 }
4140 let scale = (*input_scale).ok_or_else(|| FittedModelError::SchemaMismatch {
4141 reason: format!(
4142 "measure-jet term '{}' is missing its frozen isotropic input scale",
4143 term.name
4144 ),
4145 })?;
4146 scale.standardize(&mut queries);
4147 let support = gam_terms::basis::measure_jet_support_curve(
4148 queries.view(),
4149 centers.view(),
4150 frozen.masses.view(),
4151 &frozen.eps_band,
4152 )
4153 .map_err(|e| FittedModelError::SchemaMismatch {
4154 reason: format!(
4155 "measure-jet term '{}': support curve failed: {e}",
4156 term.name
4157 ),
4158 })?;
4159 for i in 0..data.nrows() {
4160 let v = gam_terms::basis::measure_jet_extrapolation_variance(
4161 support.row(i),
4162 &frozen.eps_band,
4163 &frozen.support_means,
4164 spectrum,
4165 Self::MEASURE_JET_COVERAGE_FLOOR,
4166 )
4167 .map_err(|e| FittedModelError::SchemaMismatch {
4168 reason: format!(
4169 "measure-jet term '{}': extrapolation variance failed: {e}",
4170 term.name
4171 ),
4172 })?;
4173 total[i] += phi_scale * v;
4174 }
4175 contributed = true;
4176 if let Some(sigma_coord) = frozen.sigma_coord {
4189 'input_var: {
4190 let MeasureJetIdentifiability::FrozenTransform { transform } =
4191 &mj.identifiability
4192 else {
4193 log::warn!(
4194 "measure-jet term '{}': identifiability is not a frozen transform; \
4195 skipping its input-measurement-error variance",
4196 term.name
4197 );
4198 break 'input_var;
4199 };
4200 let full_cols = design.design.ncols();
4201 if fit.beta.len() != full_cols {
4202 log::warn!(
4203 "measure-jet term '{}': joint coefficient vector length {} disagrees \
4204 with the replayed design's {} columns; skipping its \
4205 input-measurement-error variance",
4206 term.name,
4207 fit.beta.len(),
4208 full_cols
4209 );
4210 break 'input_var;
4211 }
4212 if design.smooth.term_designs.len() != spec.smooth_terms.len() {
4213 log::warn!(
4214 "measure-jet term '{}': smooth design/term count mismatch ({} vs {}); \
4215 skipping its input-measurement-error variance",
4216 term.name,
4217 design.smooth.term_designs.len(),
4218 spec.smooth_terms.len()
4219 );
4220 break 'input_var;
4221 }
4222 let m = centers.nrows();
4223 let m_aug = transform.nrows();
4224 let reduced = transform.ncols();
4225 let term_cols = design.smooth.term_designs[smooth_idx].ncols();
4226 if term_cols != reduced {
4227 log::warn!(
4228 "measure-jet term '{}': replayed reduced width {term_cols} disagrees \
4229 with the frozen transform ({m_aug}×{reduced}); skipping its \
4230 input-measurement-error variance",
4231 term.name
4232 );
4233 break 'input_var;
4234 }
4235 let smooth_start = full_cols - design.smooth.total_smooth_cols();
4237 let offset_in_smooth: usize = design.smooth.term_designs[..smooth_idx]
4238 .iter()
4239 .map(|d| d.ncols())
4240 .sum();
4241 let g0 = smooth_start + offset_in_smooth;
4242 let beta_term = fit.beta.slice(ndarray::s![g0..g0 + term_cols]).to_owned();
4243 let z_full = transform.dot(&beta_term);
4245 let head_rank = m_aug - m;
4246 let rep = z_full.slice(ndarray::s![..m]).to_owned();
4247 let head_coeffs = z_full.slice(ndarray::s![m..]).to_owned();
4248 let head_t = if head_rank > 0 {
4249 Some(gam_terms::basis::measure_jet_affine_head_transform(
4250 centers.view(),
4251 frozen.masses.view(),
4252 ))
4253 } else {
4254 None
4255 };
4256 if let Some(t) = head_t.as_ref() {
4257 if t.ncols() != head_rank {
4258 log::warn!(
4259 "measure-jet term '{}': reconstructed head lift rank {} disagrees \
4260 with the frozen head block {head_rank}; skipping its \
4261 input-measurement-error variance",
4262 term.name,
4263 t.ncols()
4264 );
4265 break 'input_var;
4266 }
4267 }
4268 let sigma2 = sigma_coord * sigma_coord;
4269 let mut input_var = Array1::<f64>::zeros(data.nrows());
4270 let mut ok = true;
4271 for i in 0..data.nrows() {
4272 match gam_terms::basis::measure_jet_ambient_gradient(
4273 queries.row(i),
4274 centers.view(),
4275 rep.view(),
4276 mj.length_scale,
4277 head_t.as_ref().map(|t| t.view()),
4278 head_coeffs.view(),
4279 ) {
4280 Ok(grad) => {
4281 let norm_sq: f64 = grad.iter().map(|g| g * g).sum();
4282 input_var[i] = sigma2 * norm_sq;
4283 }
4284 Err(e) => {
4285 log::warn!(
4286 "measure-jet term '{}': ambient gradient failed ({e}); \
4287 skipping its input-measurement-error variance",
4288 term.name
4289 );
4290 ok = false;
4291 break;
4292 }
4293 }
4294 }
4295 if ok {
4296 total += &input_var;
4297 }
4298 }
4299 }
4300 }
4301 Ok(contributed.then_some(total))
4302 }
4303
4304 pub fn unified(&self) -> Option<&UnifiedFitResult> {
4306 self.payload().unified.as_ref()
4307 }
4308
4309 pub fn load_from_path(path: &Path) -> Result<Self, FittedModelError> {
4310 let payload = fs::read_to_string(path).map_err(|e| FittedModelError::PayloadCorrupt {
4311 reason: format!("failed to read model '{}': {e}", path.display()),
4312 })?;
4313 let model: Self =
4314 serde_json::from_str(&payload).map_err(|e| FittedModelError::PayloadCorrupt {
4315 reason: format!("failed to parse model json: {e}"),
4316 })?;
4317 let model = model.with_synchronized_stateful_link_metadata();
4318 model.validate_for_persistence()?;
4319 model.validate_numeric_finiteness()?;
4320 Ok(model)
4321 }
4322
4323 pub fn save_to_path(&self, path: &Path) -> Result<(), FittedModelError> {
4324 let normalized = self.clone().with_synchronized_stateful_link_metadata();
4325 normalized.validate_for_persistence()?;
4326 normalized.validate_numeric_finiteness()?;
4327 let parent = path.parent().unwrap_or_else(|| Path::new("."));
4334 let file_name = path
4335 .file_name()
4336 .and_then(|s| s.to_str())
4337 .unwrap_or("model.json");
4338 let pid = std::process::id();
4339 let nanos = std::time::SystemTime::now()
4340 .duration_since(std::time::UNIX_EPOCH)
4341 .map(|d| d.as_nanos())
4342 .unwrap_or(0);
4343 let tmp = parent.join(format!(".{file_name}.tmp.{pid}.{nanos:x}"));
4344 let file = fs::File::create(&tmp).map_err(|e| FittedModelError::PayloadCorrupt {
4345 reason: format!("failed to write model '{}': {e}", tmp.display()),
4346 })?;
4347 let mut writer = std::io::BufWriter::new(file);
4348 let ser_result = serde_json::to_writer(&mut writer, &normalized);
4349 if let Err(e) = ser_result {
4350 std::io::Write::flush(&mut writer).ok();
4353 drop(writer);
4354 fs::remove_file(&tmp).ok();
4355 return Err(FittedModelError::PayloadCorrupt {
4356 reason: format!("failed to serialize model: {e}"),
4357 });
4358 }
4359 std::io::Write::flush(&mut writer).map_err(|e| FittedModelError::PayloadCorrupt {
4360 reason: format!("failed to write model '{}': {e}", tmp.display()),
4361 })?;
4362 let inner = writer
4364 .into_inner()
4365 .map_err(|e| FittedModelError::PayloadCorrupt {
4366 reason: format!("failed to flush model '{}': {}", tmp.display(), e.error()),
4367 })?;
4368 inner.sync_all().ok();
4369 drop(inner);
4370 if let Err(e) = fs::rename(&tmp, path) {
4371 fs::remove_file(&tmp).ok();
4372 return Err(FittedModelError::PayloadCorrupt {
4373 reason: format!("failed to publish model '{}': {e}", path.display()),
4374 });
4375 }
4376 if let Ok(d) = fs::File::open(parent) {
4381 d.sync_all().ok();
4382 }
4383 Ok(())
4384 }
4385
4386 pub fn require_data_schema(&self) -> Result<&DataSchema, FittedModelError> {
4387 self.data_schema
4388 .as_ref()
4389 .ok_or_else(|| FittedModelError::MissingField {
4390 reason: "model is missing data_schema; refit".to_string(),
4391 })
4392 }
4393
4394 pub fn saved_spline_scan(
4398 &self,
4399 ) -> Result<Option<(&str, gam_solve::spline_scan::SplineScanFit)>, FittedModelError> {
4400 let Some(saved) = self.spline_scan.as_ref() else {
4401 return Ok(None);
4402 };
4403 let fit = gam_solve::spline_scan::SplineScanFit::from_state(&saved.state)
4404 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4405 Ok(Some((saved.feature_column.as_str(), fit)))
4406 }
4407
4408 pub fn saved_residual_cascade(
4413 &self,
4414 ) -> Result<
4415 Option<(&[String], gam_solve::residual_cascade::ResidualCascadeFit)>,
4416 FittedModelError,
4417 > {
4418 let Some(saved) = self.residual_cascade.as_ref() else {
4419 return Ok(None);
4420 };
4421 let fit = gam_solve::residual_cascade::ResidualCascadeFit::from_state(&saved.state)
4422 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4423 Ok(Some((saved.feature_columns.as_slice(), fit)))
4424 }
4425
4426 pub fn random_effect_group_columns(&self) -> HashSet<String> {
4441 let Some(training_headers) = self.training_headers.as_ref() else {
4442 return HashSet::new();
4443 };
4444 let mut out = HashSet::<String>::new();
4445 for spec in self.saved_term_specs() {
4446 for term in &spec.random_effect_terms {
4447 if !term.lenient_unseen {
4448 continue;
4449 }
4450 if let Some(name) = training_headers.get(term.feature_col) {
4451 out.insert(name.clone());
4452 }
4453 }
4454 for term in &spec.smooth_terms {
4461 if let Some(group_col) = re_factor_smooth_group_col(&term.basis)
4462 && let Some(name) = training_headers.get(group_col)
4463 {
4464 out.insert(name.clone());
4465 }
4466 }
4467 }
4468 out
4469 }
4470
4471 pub fn numeric_fixed_factor_vocabularies(&self) -> Vec<(String, HashSet<u64>)> {
4489 let Some(training_headers) = self.training_headers.as_ref() else {
4490 return Vec::new();
4491 };
4492 let Some(schema) = self.data_schema.as_ref() else {
4493 return Vec::new();
4494 };
4495 let mut out = Vec::<(String, HashSet<u64>)>::new();
4496 for spec in self.saved_term_specs() {
4497 for term in &spec.random_effect_terms {
4498 if term.lenient_unseen || term.drop_first_level {
4499 continue;
4500 }
4501 let Some(levels) = term.frozen_levels.as_ref() else {
4502 continue;
4503 };
4504 let Some(name) = training_headers.get(term.feature_col) else {
4505 continue;
4506 };
4507 let is_numeric = schema
4510 .columns
4511 .iter()
4512 .find(|c| &c.name == name)
4513 .map(|c| matches!(c.kind, ColumnKindTag::Continuous | ColumnKindTag::Binary))
4514 .unwrap_or(false);
4515 if !is_numeric {
4516 continue;
4517 }
4518 let vocab: HashSet<u64> = levels
4519 .iter()
4520 .map(|&b| gam_data::canonical_level_bits(f64::from_bits(b)))
4521 .collect();
4522 out.push((name.clone(), vocab));
4523 }
4524 }
4525 out
4526 }
4527
4528 pub fn validate_for_persistence(&self) -> Result<(), FittedModelError> {
4529 self.validate_payload_version()?;
4543 let expectile_family_tag = {
4544 let family = self.family.trim().to_ascii_lowercase();
4545 family == "expectile" || family.starts_with("expectile(")
4546 };
4547 match self.estimator {
4548 FittedEstimator::Likelihood if expectile_family_tag => {
4549 return Err(FittedModelError::SchemaMismatch {
4550 reason:
4551 "saved family is tagged expectile but estimator metadata says likelihood"
4552 .to_string(),
4553 });
4554 }
4555 FittedEstimator::Likelihood => {}
4556 FittedEstimator::Expectile { tau } => {
4557 if !tau.is_finite() || tau <= 0.0 || tau >= 1.0 {
4558 return Err(FittedModelError::SchemaMismatch {
4559 reason: format!(
4560 "saved expectile estimator requires finite tau strictly in (0, 1), got {tau}"
4561 ),
4562 });
4563 }
4564 let gaussian_identity_standard = self.model_kind == ModelKind::Standard
4565 && matches!(
4566 &self.family_state,
4567 FittedFamily::Standard { likelihood, .. }
4568 if likelihood == &LikelihoodSpec::gaussian_identity()
4569 );
4570 if !gaussian_identity_standard || !expectile_family_tag {
4571 return Err(FittedModelError::SchemaMismatch {
4572 reason: format!(
4573 "saved expectile estimator requires an expectile-tagged standard \
4574 Gaussian-identity fit; got model_kind={:?}, family={:?}, likelihood={:?}",
4575 self.model_kind,
4576 self.family,
4577 self.family_state.likelihood(),
4578 ),
4579 });
4580 }
4581 }
4582 }
4583 if self.training_table_kind.trim().is_empty() {
4584 return Err(FittedModelError::MissingField {
4585 reason: "saved model training_table_kind must be non-empty".to_string(),
4586 });
4587 }
4588 if let Some(scan) = self.spline_scan.as_ref() {
4589 if self.fit_result.is_some() || self.unified.is_some() {
4594 return Err(FittedModelError::SchemaMismatch {
4595 reason: "spline-scan model must not also carry a dense fit_result/unified \
4596 payload; the representations are mutually exclusive"
4597 .to_string(),
4598 });
4599 }
4600 if self.model_kind != ModelKind::Standard
4601 || self.family_state.likelihood() != LikelihoodSpec::gaussian_identity()
4602 {
4603 return Err(FittedModelError::SchemaMismatch {
4604 reason: format!(
4605 "spline-scan representation requires a standard Gaussian-identity model; \
4606 got model_kind={:?}, likelihood={:?}",
4607 self.model_kind,
4608 self.family_state.likelihood()
4609 ),
4610 });
4611 }
4612 if scan.feature_column.is_empty() {
4613 return Err(FittedModelError::MissingField {
4614 reason: "spline-scan model is missing its feature column name; refit"
4615 .to_string(),
4616 });
4617 }
4618 gam_solve::spline_scan::SplineScanFit::from_state(&scan.state)
4619 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4620 if self.data_schema.is_none() {
4626 return Err(FittedModelError::MissingField {
4627 reason: "spline-scan model is missing data_schema; refit".to_string(),
4628 });
4629 }
4630 if self.training_headers.is_none() {
4631 return Err(FittedModelError::MissingField {
4632 reason: "spline-scan model is missing training_headers; refit".to_string(),
4633 });
4634 }
4635 return Ok(());
4636 } else if let Some(cascade) = self.residual_cascade.as_ref() {
4637 if self.spline_scan.is_some() || self.fit_result.is_some() || self.unified.is_some() {
4641 return Err(FittedModelError::SchemaMismatch {
4642 reason: "residual-cascade model must not also carry spline_scan / \
4643 fit_result / unified payloads; the representations are \
4644 mutually exclusive"
4645 .to_string(),
4646 });
4647 }
4648 if self.model_kind != ModelKind::Standard
4649 || self.family_state.likelihood() != LikelihoodSpec::gaussian_identity()
4650 {
4651 return Err(FittedModelError::SchemaMismatch {
4652 reason: format!(
4653 "residual-cascade representation requires a standard Gaussian-identity \
4654 model; got model_kind={:?}, likelihood={:?}",
4655 self.model_kind,
4656 self.family_state.likelihood()
4657 ),
4658 });
4659 }
4660 if cascade.feature_columns.is_empty()
4661 || !(2..=3).contains(&cascade.feature_columns.len())
4662 {
4663 return Err(FittedModelError::MissingField {
4664 reason: format!(
4665 "residual-cascade model needs 2 or 3 feature columns; got {}; refit",
4666 cascade.feature_columns.len()
4667 ),
4668 });
4669 }
4670 gam_solve::residual_cascade::ResidualCascadeFit::from_state(&cascade.state)
4671 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4672 if self.data_schema.is_none() {
4673 return Err(FittedModelError::MissingField {
4674 reason: "residual-cascade model is missing data_schema; refit".to_string(),
4675 });
4676 }
4677 if self.training_headers.is_none() {
4678 return Err(FittedModelError::MissingField {
4679 reason: "residual-cascade model is missing training_headers; refit".to_string(),
4680 });
4681 }
4682 return Ok(());
4683 } else if self.fit_result.is_none() {
4684 return Err(FittedModelError::MissingField {
4685 reason: "model is missing canonical fit_result payload; refit".to_string(),
4686 });
4687 }
4688 if self.data_schema.is_none() {
4689 return Err(FittedModelError::MissingField {
4690 reason: "model is missing data_schema; refit".to_string(),
4691 });
4692 }
4693 if self.training_headers.is_none() {
4694 return Err(FittedModelError::MissingField {
4695 reason: "model is missing training_headers; refit to guarantee stable feature mapping at prediction time"
4696 .to_string(),
4697 });
4698 }
4699 let spec = self.resolved_termspec.as_ref().ok_or_else(|| {
4700 FittedModelError::MissingField {
4701 reason: "model is missing resolved_termspec; refit to guarantee train/predict design consistency"
4702 .to_string(),
4703 }
4704 })?;
4705 validate_frozen_term_collectionspec(spec, "resolved_termspec")?;
4706
4707 if self.formula_noise.is_some() && self.resolved_termspec_noise.is_none() {
4708 return Err(FittedModelError::MissingField {
4709 reason: "model defines formula_noise but is missing resolved_termspec_noise; refit"
4710 .to_string(),
4711 });
4712 }
4713 if let Some(spec_noise) = self.resolved_termspec_noise.as_ref() {
4714 validate_frozen_term_collectionspec(spec_noise, "resolved_termspec_noise")?;
4715 }
4716 if matches!(self.family_state, FittedFamily::TransformationNormal { .. }) {
4717 let score = self.transformation_score_calibration.ok_or_else(|| {
4718 FittedModelError::MissingField {
4719 reason: "transformation-normal model is missing transformation_score_calibration; refit"
4720 .to_string(),
4721 }
4722 })?;
4723 score.validate("transformation-normal model")?;
4724 let geometry = self.transformation_geometry.as_ref().ok_or_else(|| {
4729 FittedModelError::MissingField {
4730 reason: "transformation-normal model is missing the direct-α geometry record \
4731 (transformation_geometry); this is a pre-cutover (v12-or-older) CTN \
4732 model whose squared-γ chart is not replayable under the direct-α \
4733 cutover (gam#2306) — refit"
4734 .to_string(),
4735 }
4736 })?;
4737 geometry.validate("transformation-normal model")?;
4738 if let Some(knots) = self.transformation_response_knots.as_ref() {
4742 if geometry.response_knot_count != knots.len() {
4743 return Err(FittedModelError::SchemaMismatch {
4744 reason: format!(
4745 "transformation-normal geometry response_knot_count {} disagrees with \
4746 transformation_response_knots length {}",
4747 geometry.response_knot_count,
4748 knots.len()
4749 ),
4750 });
4751 }
4752 }
4753 if let Some(degree) = self.transformation_response_degree {
4754 if geometry.response_degree != degree {
4755 return Err(FittedModelError::SchemaMismatch {
4756 reason: format!(
4757 "transformation-normal geometry response_degree {} disagrees with \
4758 transformation_response_degree {degree}",
4759 geometry.response_degree
4760 ),
4761 });
4762 }
4763 }
4764 let carrier = self.transformation_cone_carrier.as_ref().ok_or_else(|| {
4769 FittedModelError::MissingField {
4770 reason: "transformation-normal model is missing the monotonicity-cone carrier \
4771 (transformation_cone_carrier); constrained posterior sampling cannot \
4772 certify draws against the positivity cone — refit"
4773 .to_string(),
4774 }
4775 })?;
4776 let expected = geometry
4777 .cone_carrier_row_count
4778 .checked_mul(geometry.cone_carrier_covariate_width)
4779 .ok_or_else(|| FittedModelError::SchemaMismatch {
4780 reason: "transformation-normal cone carrier dimensions overflow usize"
4781 .to_string(),
4782 })?;
4783 if carrier.len() != expected {
4784 return Err(FittedModelError::SchemaMismatch {
4785 reason: format!(
4786 "transformation-normal cone carrier length {} disagrees with geometry \
4787 {} rows x {} covariate columns = {expected}",
4788 carrier.len(),
4789 geometry.cone_carrier_row_count,
4790 geometry.cone_carrier_covariate_width,
4791 ),
4792 });
4793 }
4794 if carrier.iter().any(|value| !value.is_finite()) {
4795 return Err(FittedModelError::SchemaMismatch {
4796 reason: "transformation-normal cone carrier contains a non-finite entry"
4797 .to_string(),
4798 });
4799 }
4800 }
4801 if matches!(self.family_state, FittedFamily::MarginalSlope { .. }) {
4802 if self.formula_logslope.is_none() {
4803 return Err(FittedModelError::MissingField {
4804 reason: "marginal-slope model is missing formula_logslope; refit".to_string(),
4805 });
4806 }
4807 if self.z_column.is_none() {
4808 return Err(FittedModelError::MissingField {
4809 reason: "marginal-slope model is missing z_column; refit".to_string(),
4810 });
4811 }
4812 let z_normalization =
4813 self.latent_z_normalization
4814 .ok_or_else(|| FittedModelError::MissingField {
4815 reason: "marginal-slope model is missing latent_z_normalization; refit"
4816 .to_string(),
4817 })?;
4818 z_normalization.validate("marginal-slope model")?;
4819 let latent_measure =
4820 self.latent_measure
4821 .as_ref()
4822 .ok_or_else(|| FittedModelError::MissingField {
4823 reason: "marginal-slope model is missing latent_measure; refit".to_string(),
4824 })?;
4825 latent_measure
4826 .validate("marginal-slope model latent_measure")
4827 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4828 if self.marginal_baseline.is_none() || self.logslope_baseline.is_none() {
4829 return Err(FittedModelError::MissingField {
4830 reason: "marginal-slope model is missing baseline offsets; refit".to_string(),
4831 });
4832 }
4833 if self.resolved_termspec_logslope.as_ref().is_none() {
4834 return Err(FittedModelError::MissingField {
4835 reason: "marginal-slope model is missing resolved_termspec_logslope for the logslope surface"
4836 .to_string(),
4837 });
4838 }
4839 match self.family_state.frailty() {
4840 Some(FrailtySpec::None)
4841 | Some(FrailtySpec::GaussianShift {
4842 scale: FrailtyScale::Fixed { .. },
4843 }) => {}
4844 Some(FrailtySpec::GaussianShift {
4845 scale: FrailtyScale::Learned { .. },
4846 }) => {
4847 return Err(FittedModelError::IncompatibleConfig {
4848 reason: "marginal-slope model requires a fixed GaussianShift sigma in family_state.frailty"
4849 .to_string(),
4850 });
4851 }
4852 Some(FrailtySpec::HazardMultiplier { .. }) => {
4853 return Err(FittedModelError::IncompatibleConfig {
4854 reason: "marginal-slope model does not support HazardMultiplier frailty"
4855 .to_string(),
4856 });
4857 }
4858 None => {
4859 return Err(FittedModelError::MissingField {
4860 reason: "marginal-slope model is missing family_state.frailty; refit"
4861 .to_string(),
4862 });
4863 }
4864 }
4865 }
4866
4867 if let FittedFamily::Survival {
4868 survival_likelihood,
4869 frailty,
4870 ..
4871 } = &self.family_state
4872 {
4873 if matches!(
4874 survival_likelihood.as_deref(),
4875 Some("latent") | Some("latent-binary")
4876 ) {
4877 return Err(FittedModelError::SchemaMismatch {
4878 reason: "latent hazard-window models must persist explicit family_state metadata, not generic survival metadata"
4879 .to_string(),
4880 });
4881 }
4882 if survival_likelihood.as_deref() == Some("marginal-slope") {
4883 if self.formula_logslope.is_none() {
4884 return Err(FittedModelError::MissingField {
4885 reason: "survival marginal-slope model is missing formula_logslope; refit"
4886 .to_string(),
4887 });
4888 }
4889 if self.z_column.is_none() {
4890 return Err(FittedModelError::MissingField {
4891 reason: "survival marginal-slope model is missing z_column; refit"
4892 .to_string(),
4893 });
4894 }
4895 let z_normalization =
4896 self.latent_z_normalization
4897 .ok_or_else(|| {
4898 FittedModelError::MissingField {
4899 reason:
4900 "survival marginal-slope model is missing latent_z_normalization; refit"
4901 .to_string(),
4902 }
4903 })?;
4904 z_normalization.validate("survival marginal-slope model")?;
4905 let latent_measure =
4906 self.latent_measure
4907 .as_ref()
4908 .ok_or_else(|| FittedModelError::MissingField {
4909 reason:
4910 "survival marginal-slope model is missing latent_measure; refit"
4911 .to_string(),
4912 })?;
4913 latent_measure
4914 .validate("survival marginal-slope model latent_measure")
4915 .map_err(|reason| FittedModelError::PayloadCorrupt { reason })?;
4916 if self.logslope_baseline.is_none() {
4917 return Err(FittedModelError::MissingField {
4918 reason: "survival marginal-slope model is missing logslope_baseline; refit"
4919 .to_string(),
4920 });
4921 }
4922 if self.resolved_termspec_logslope.as_ref().is_none() {
4923 return Err(FittedModelError::MissingField {
4924 reason: "survival marginal-slope model is missing resolved_termspec_logslope for the logslope surface"
4925 .to_string(),
4926 });
4927 }
4928 match frailty {
4929 FrailtySpec::None
4930 | FrailtySpec::GaussianShift {
4931 scale: FrailtyScale::Fixed { .. },
4932 } => {}
4933 FrailtySpec::GaussianShift {
4934 scale: FrailtyScale::Learned { .. },
4935 } => {
4936 return Err(FittedModelError::IncompatibleConfig {
4937 reason: "survival marginal-slope model requires a fixed GaussianShift sigma in family_state.frailty"
4938 .to_string(),
4939 });
4940 }
4941 FrailtySpec::HazardMultiplier { .. } => {
4942 return Err(FittedModelError::IncompatibleConfig {
4943 reason: "survival marginal-slope model does not support HazardMultiplier frailty"
4944 .to_string(),
4945 });
4946 }
4947 }
4948 } else if !matches!(frailty, FrailtySpec::None) {
4949 return Err(FittedModelError::IncompatibleConfig {
4950 reason:
4951 "non-marginal survival models do not currently persist a frailty modifier"
4952 .to_string(),
4953 });
4954 }
4955 if self.survival_time_basis.is_none() {
4963 return Err(FittedModelError::MissingField {
4964 reason: "survival model is missing survival_time_basis; refit to persist the baseline-time basis configuration".to_string(),
4965 });
4966 }
4967 if self.survival_time_anchor.is_none() {
4968 return Err(FittedModelError::MissingField {
4969 reason: "survival model is missing survival_time_anchor; refit to persist the baseline-time anchor".to_string(),
4970 });
4971 }
4972 }
4973 if let FittedFamily::LatentSurvival { frailty } = &self.family_state {
4974 match frailty {
4975 FrailtySpec::HazardMultiplier {
4976 scale: FrailtyScale::Fixed { .. },
4977 ..
4978 } => {}
4979 FrailtySpec::HazardMultiplier {
4980 scale: FrailtyScale::Learned { .. },
4981 ..
4982 } => {
4983 return Err(FittedModelError::IncompatibleConfig {
4984 reason: "latent survival model requires a fixed HazardMultiplier sigma in family_state.frailty"
4985 .to_string(),
4986 });
4987 }
4988 FrailtySpec::GaussianShift { .. } | FrailtySpec::None => {
4989 return Err(FittedModelError::IncompatibleConfig {
4990 reason: "latent survival model requires a fixed HazardMultiplier frailty specification"
4991 .to_string(),
4992 });
4993 }
4994 }
4995 if self.survival_likelihood.as_deref() != Some("latent") {
4996 return Err(FittedModelError::SchemaMismatch {
4997 reason: "latent survival model must persist survival_likelihood=latent"
4998 .to_string(),
4999 });
5000 }
5001 }
5002 if let FittedFamily::LatentBinary { frailty } = &self.family_state {
5003 match frailty {
5004 FrailtySpec::HazardMultiplier {
5005 scale: FrailtyScale::Fixed { .. },
5006 ..
5007 } => {}
5008 FrailtySpec::HazardMultiplier {
5009 scale: FrailtyScale::Learned { .. },
5010 ..
5011 } => {
5012 return Err(FittedModelError::IncompatibleConfig {
5013 reason: "latent binary model requires a fixed HazardMultiplier sigma in family_state.frailty"
5014 .to_string(),
5015 });
5016 }
5017 FrailtySpec::GaussianShift { .. } | FrailtySpec::None => {
5018 return Err(FittedModelError::IncompatibleConfig {
5019 reason: "latent binary model requires a fixed HazardMultiplier frailty specification"
5020 .to_string(),
5021 });
5022 }
5023 }
5024 if self.survival_likelihood.as_deref() != Some("latent-binary") {
5025 return Err(FittedModelError::SchemaMismatch {
5026 reason: "latent binary model must persist survival_likelihood=latent-binary"
5027 .to_string(),
5028 });
5029 }
5030 }
5031
5032 let family_likelihood = match &self.family_state {
5033 FittedFamily::Standard { likelihood, .. }
5034 | FittedFamily::LocationScale { likelihood, .. }
5035 | FittedFamily::MarginalSlope { likelihood, .. }
5036 | FittedFamily::Survival { likelihood, .. }
5037 | FittedFamily::TransformationNormal { likelihood, .. } => Some(likelihood),
5038 FittedFamily::LatentSurvival { .. } | FittedFamily::LatentBinary { .. } => None,
5039 };
5040 let is_standard_or_location_scale = matches!(
5041 self.family_state,
5042 FittedFamily::Standard { .. } | FittedFamily::LocationScale { .. }
5043 );
5044 if is_standard_or_location_scale
5045 && family_likelihood.is_some_and(LikelihoodSpec::is_binomial_sas)
5046 {
5047 self.saved_sas_state()?;
5048 }
5049 if is_standard_or_location_scale
5050 && family_likelihood.is_some_and(LikelihoodSpec::is_binomial_beta_logistic)
5051 {
5052 self.saved_beta_logistic_state()?;
5053 }
5054 if is_standard_or_location_scale
5055 && family_likelihood.is_some_and(LikelihoodSpec::is_binomial_mixture)
5056 {
5057 self.saved_mixture_state()?;
5058 }
5059 if matches!(self.family_state, FittedFamily::Standard { .. })
5060 && family_likelihood.is_some_and(LikelihoodSpec::is_latent_cloglog)
5061 {
5062 self.saved_latent_cloglog_state()?;
5063 }
5064 if matches!(self.family_state, FittedFamily::LocationScale { .. })
5065 && family_likelihood.is_some_and(LikelihoodSpec::is_latent_cloglog)
5066 {
5067 return Err(FittedModelError::IncompatibleConfig {
5068 reason: "latent-cloglog-binomial is not supported for location-scale saved models"
5069 .to_string(),
5070 });
5071 }
5072 if matches!(self.family_state, FittedFamily::Survival { .. })
5073 && self.survival_likelihood.is_none()
5074 {
5075 return Err(FittedModelError::MissingField {
5076 reason: "saved survival model is missing survival_likelihood metadata; refit"
5077 .to_string(),
5078 });
5079 }
5080 let is_survival_location_scale = matches!(self.family_state, FittedFamily::Survival { .. })
5081 && self
5082 .survival_likelihood
5083 .as_deref()
5084 .is_some_and(|value| value.eq_ignore_ascii_case("location-scale"));
5085 if !is_survival_location_scale && self.survival_location_scale_structure.is_some() {
5086 return Err(FittedModelError::SchemaMismatch {
5087 reason: "non-location-scale model carries location-scale replay structure"
5088 .to_string(),
5089 });
5090 }
5091 let has_any_saved_link_wiggle = self.linkwiggle_knots.is_some()
5092 || self.linkwiggle_degree.is_some()
5093 || self.linkwiggle_penalty_metadata.is_some()
5094 || self.beta_link_wiggle.is_some()
5095 || self
5096 .fit_result
5097 .as_ref()
5098 .and_then(|fit| fit.block_by_role(BlockRole::LinkWiggle))
5099 .is_some();
5100 let saved_link_wiggle = self.saved_link_wiggle()?;
5101 if has_any_saved_link_wiggle && saved_link_wiggle.is_none() {
5102 return Err(FittedModelError::SchemaMismatch {
5103 reason: "saved model has incomplete link-wiggle state; expected metadata and coefficients"
5104 .to_string(),
5105 });
5106 }
5107 if matches!(self.family_state, FittedFamily::Standard { .. })
5108 && saved_link_wiggle.is_some()
5109 && self.linkwiggle_penalty_metadata.is_none()
5110 {
5111 return Err(FittedModelError::MissingField {
5112 reason: "standard link-wiggle model is missing canonical penalty metadata; refit"
5113 .to_string(),
5114 });
5115 }
5116 let has_any_saved_baseline_time_wiggle = self.baseline_timewiggle_knots.is_some()
5117 || self.baseline_timewiggle_degree.is_some()
5118 || self.baseline_timewiggle_penalty_orders.is_some()
5119 || self.baseline_timewiggle_double_penalty.is_some()
5120 || self.beta_baseline_timewiggle.is_some()
5121 || self.beta_baseline_timewiggle_by_cause.is_some();
5122 let is_joint_cause_specific = self
5123 .survival_cause_count
5124 .is_some_and(|cause_count| cause_count > 1);
5125 if has_any_saved_baseline_time_wiggle {
5126 if is_joint_cause_specific {
5127 let complete = self.baseline_timewiggle_knots.is_some()
5128 && self.baseline_timewiggle_degree.is_some()
5129 && self.baseline_timewiggle_penalty_orders.is_some()
5130 && self.baseline_timewiggle_double_penalty.is_some()
5131 && self.beta_baseline_timewiggle_by_cause.is_some();
5132 if !complete {
5133 return Err(FittedModelError::SchemaMismatch {
5134 reason: "saved joint cause-specific survival model has incomplete baseline-timewiggle state; expected metadata and per-cause coefficients"
5135 .to_string(),
5136 });
5137 }
5138 } else if self.saved_baseline_time_wiggle()?.is_none() {
5139 return Err(FittedModelError::SchemaMismatch {
5140 reason: "saved model has incomplete baseline-timewiggle state; expected metadata and coefficients"
5141 .to_string(),
5142 });
5143 }
5144 }
5145 if is_survival_location_scale {
5146 validate_survival_location_scale_saved_fit(self.payload(), saved_link_wiggle.as_ref())?;
5147 }
5148 self.validate_required_posterior_mean_state()?;
5149
5150 if let Some(runtime) = self.score_warp_runtime.as_ref() {
5160 runtime.validate_exact_replay_contract().map_err(|err| {
5161 FittedModelError::PayloadCorrupt {
5162 reason: format!("saved anchored score-warp runtime is invalid: {err}"),
5163 }
5164 })?;
5165 }
5166 if let Some(runtime) = self.link_deviation_runtime.as_ref() {
5167 runtime.validate_exact_replay_contract().map_err(|err| {
5168 FittedModelError::PayloadCorrupt {
5169 reason: format!("saved anchored link-deviation runtime is invalid: {err}"),
5170 }
5171 })?;
5172 }
5173 if matches!(self.family_state, FittedFamily::MarginalSlope { .. }) {
5174 validate_marginal_slope_saved_fit(
5175 self.fit_result.as_ref().expect("checked above"),
5176 self.score_warp_runtime.as_ref(),
5177 self.link_deviation_runtime.as_ref(),
5178 "fit_result",
5179 )?;
5180 let unified = self
5181 .unified
5182 .as_ref()
5183 .ok_or_else(|| FittedModelError::MissingField {
5184 reason: "marginal-slope model is missing unified fit payload; refit"
5185 .to_string(),
5186 })?;
5187 validate_marginal_slope_saved_fit(
5188 unified,
5189 self.score_warp_runtime.as_ref(),
5190 self.link_deviation_runtime.as_ref(),
5191 "unified",
5192 )?;
5193 }
5194 if self
5195 .survival_likelihood
5196 .as_deref()
5197 .is_some_and(|value| value.eq_ignore_ascii_case("marginal-slope"))
5198 {
5199 validate_survival_marginal_slope_saved_fit(
5200 self,
5201 self.fit_result.as_ref().expect("checked above"),
5202 "fit_result",
5203 )?;
5204 if let Some(unified) = self.unified.as_ref() {
5205 validate_survival_marginal_slope_saved_fit(self, unified, "unified")?;
5206 }
5207 }
5208
5209 Ok(())
5210 }
5211
5212 fn validate_required_posterior_mean_state(&self) -> Result<(), FittedModelError> {
5222 if !self.prediction_uses_posterior_mean() {
5223 return Ok(());
5224 }
5225 let fit = self
5226 .payload()
5227 .fit_result
5228 .as_ref()
5229 .ok_or_else(|| FittedModelError::MissingField {
5230 reason:
5231 "curved-link model is missing the fit state required for posterior-mean prediction"
5232 .to_string(),
5233 })?;
5234 let p = fit.beta.len();
5235 if let Some(covariance) = fit.beta_covariance() {
5236 if covariance.dim() == (p, p) {
5237 return Ok(());
5238 }
5239 return Err(FittedModelError::SchemaMismatch {
5240 reason: format!(
5241 "curved-link model conditional covariance has shape {}x{}, expected {p}x{p} in the saved coefficient frame",
5242 covariance.nrows(),
5243 covariance.ncols(),
5244 ),
5245 });
5246 }
5247
5248 if fit
5249 .geometry
5250 .as_ref()
5251 .is_some_and(|geometry| !geometry.coefficient_gauge.is_identity())
5252 {
5253 return Err(FittedModelError::SchemaMismatch {
5254 reason: format!(
5255 "curved-link model has no saved/raw-frame covariance and its penalized precision lives in an active gauge; persist the lifted {p}x{p} covariance required for posterior-mean prediction"
5256 ),
5257 });
5258 }
5259 let precision = fit.penalized_hessian().ok_or_else(|| FittedModelError::MissingField {
5260 reason: format!(
5261 "curved-link model must persist a {p}x{p} joint conditional covariance or penalized precision to compute the required posterior mean"
5262 ),
5263 })?;
5264 if precision.dim() != (p, p) {
5265 return Err(FittedModelError::SchemaMismatch {
5266 reason: format!(
5267 "curved-link model penalized precision has shape {}x{}, expected {p}x{p} in the saved coefficient frame",
5268 precision.nrows(),
5269 precision.ncols(),
5270 ),
5271 });
5272 }
5273 precision
5274 .cholesky(faer::Side::Lower)
5275 .map_err(|error| FittedModelError::PayloadCorrupt {
5276 reason: format!(
5277 "curved-link model penalized precision cannot define the required posterior mean: strict Cholesky failed: {error}"
5278 ),
5279 })?;
5280 fit.coefficient_covariance_scale()
5281 .map_err(|error| FittedModelError::PayloadCorrupt {
5282 reason: format!(
5283 "curved-link model cannot scale its saved penalized precision into a posterior covariance: {error}"
5284 ),
5285 })?;
5286 Ok(())
5287 }
5288
5289 pub fn validate_numeric_finiteness(&self) -> Result<(), FittedModelError> {
5290 let corrupt = |reason: String| FittedModelError::PayloadCorrupt { reason };
5291 if let Some(fit) = self.fit_result.as_ref() {
5292 fit.validate_numeric_finiteness()
5293 .map_err(|e| corrupt(e.to_string()))?;
5294 }
5295
5296 for (name, opt) in [
5297 ("survival_baseline_scale", self.survival_baseline_scale),
5298 ("survival_baseline_shape", self.survival_baseline_shape),
5299 ("survival_baseline_rate", self.survival_baseline_rate),
5300 ("survival_baseline_makeham", self.survival_baseline_makeham),
5301 (
5302 "survival_time_smooth_lambda",
5303 self.survival_time_smooth_lambda,
5304 ),
5305 ("survival_time_anchor", self.survival_time_anchor),
5306 ("survivalridge_lambda", self.survivalridge_lambda),
5307 ] {
5308 if let Some(v) = opt {
5309 ensure_finite_scalar(name, v).map_err(corrupt)?;
5310 }
5311 }
5312
5313 if let Some(v) = self.beta_noise.as_ref() {
5314 validate_all_finite("beta_noise", v.iter().copied()).map_err(corrupt)?;
5315 }
5316 if let Some(v) = self.noise_projection.as_ref() {
5317 validate_all_finite("noise_projection", v.iter().flatten().copied())
5318 .map_err(corrupt)?;
5319 if self.noise_projection_ridge_alpha.is_none() {
5320 return Err(FittedModelError::MissingField {
5321 reason:
5322 "model has noise_projection but is missing noise_projection_ridge_alpha; refit"
5323 .to_string(),
5324 });
5325 }
5326 }
5327 if let Some(v) = self.noise_center.as_ref() {
5328 validate_all_finite("noise_center", v.iter().copied()).map_err(corrupt)?;
5329 }
5330 if let Some(v) = self.noise_scale.as_ref() {
5331 validate_all_finite("noise_scale", v.iter().copied()).map_err(corrupt)?;
5332 }
5333 if let Some(v) = self.noise_projection_ridge_alpha {
5334 ensure_finite_scalar("noise_projection_ridge_alpha", v).map_err(corrupt)?;
5335 if v < 0.0 {
5336 return Err(FittedModelError::InvalidInput {
5337 reason: format!("noise_projection_ridge_alpha must be non-negative, got {v}"),
5338 });
5339 }
5340 }
5341 if let Some(v) = self.gaussian_response_scale {
5342 ensure_finite_scalar("gaussian_response_scale", v).map_err(corrupt)?;
5343 }
5344 if let Some(v) = self.beta_link_wiggle.as_ref() {
5345 validate_all_finite("beta_link_wiggle", v.iter().copied()).map_err(corrupt)?;
5346 }
5347 if let Some(v) = self.link_wiggle_index_shift.as_ref() {
5348 validate_all_finite("link_wiggle_index_shift", v.iter().copied()).map_err(corrupt)?;
5349 }
5350 if let Some(v) = self.beta_baseline_timewiggle.as_ref() {
5351 validate_all_finite("beta_baseline_timewiggle", v.iter().copied()).map_err(corrupt)?;
5352 }
5353 if let Some(v) = self.beta_baseline_timewiggle_by_cause.as_ref() {
5354 validate_all_finite(
5355 "beta_baseline_timewiggle_by_cause",
5356 v.iter().flatten().copied(),
5357 )
5358 .map_err(corrupt)?;
5359 }
5360 if let Some(v) = self.latent_z_normalization {
5361 v.validate("latent_z_normalization")?;
5362 }
5363 if let Some(v) = self.latent_measure.as_ref() {
5364 v.validate("latent_measure").map_err(corrupt)?;
5365 }
5366 if let Some(v) = self.survival_beta_time.as_ref() {
5367 validate_all_finite("survival_beta_time", v.iter().copied()).map_err(corrupt)?;
5368 }
5369 if let Some(v) = self.survival_beta_threshold.as_ref() {
5370 validate_all_finite("survival_beta_threshold", v.iter().copied()).map_err(corrupt)?;
5371 }
5372 if let Some(v) = self.survival_beta_log_sigma.as_ref() {
5373 validate_all_finite("survival_beta_log_sigma", v.iter().copied()).map_err(corrupt)?;
5374 }
5375 if let Some(v) = self.mixture_link_param_covariance.as_ref() {
5376 validate_all_finite("mixture_link_param_covariance", v.iter().flatten().copied())
5377 .map_err(corrupt)?;
5378 }
5379 if let Some(v) = self.sas_param_covariance.as_ref() {
5380 validate_all_finite("sas_param_covariance", v.iter().flatten().copied())
5381 .map_err(corrupt)?;
5382 }
5383 Ok(())
5384 }
5385}
5386
5387use gam_solve::estimate::{ensure_finite_scalar, validate_all_finite};
5388
5389fn validate_frozen_term_collectionspec(
5390 spec: &TermCollectionSpec,
5391 label: &str,
5392) -> Result<(), FittedModelError> {
5393 spec.validate_frozen(label)
5394 .map_err(|reason| FittedModelError::SchemaMismatch { reason })
5395}
5396
5397impl Deref for FittedModel {
5398 type Target = FittedModelPayload;
5399
5400 fn deref(&self) -> &Self::Target {
5401 self.payload()
5402 }
5403}
5404
5405impl DerefMut for FittedModel {
5406 fn deref_mut(&mut self) -> &mut Self::Target {
5407 self.payload_mut()
5408 }
5409}
5410
5411pub fn survival_baseline_config_from_model(
5416 model: &FittedModel,
5417) -> Result<SurvivalBaselineConfig, FittedModelError> {
5418 let target = model.survival_baseline_target.as_deref().ok_or_else(|| {
5419 FittedModelError::MissingField {
5420 reason: "saved survival model missing survival_baseline_target; refit".to_string(),
5421 }
5422 })?;
5423 parse_survival_baseline_config(
5424 target,
5425 model.survival_baseline_scale,
5426 model.survival_baseline_shape,
5427 model.survival_baseline_rate,
5428 model.survival_baseline_makeham,
5429 )
5430 .map_err(|reason| FittedModelError::IncompatibleConfig { reason })
5431}
5432
5433pub fn load_survival_time_basis_config_from_model(
5434 model: &FittedModel,
5435) -> Result<SurvivalTimeBasisConfig, FittedModelError> {
5436 match model
5437 .survival_time_basis
5438 .as_deref()
5439 .ok_or_else(|| FittedModelError::MissingField {
5440 reason: "saved survival model missing survival_time_basis".to_string(),
5441 })?
5442 .to_ascii_lowercase()
5443 .as_str()
5444 {
5445 "none" => Ok(SurvivalTimeBasisConfig::None),
5446 "linear" => Ok(SurvivalTimeBasisConfig::Linear),
5447 "bspline" => {
5448 let degree =
5449 model
5450 .survival_time_degree
5451 .ok_or_else(|| FittedModelError::MissingField {
5452 reason: "saved survival bspline model missing survival_time_degree"
5453 .to_string(),
5454 })?;
5455 let knots = model.survival_time_knots.clone().ok_or_else(|| {
5456 FittedModelError::MissingField {
5457 reason: "saved survival bspline model missing survival_time_knots".to_string(),
5458 }
5459 })?;
5460 let smooth_lambda = model.survival_time_smooth_lambda.unwrap_or(1e-2);
5461 if degree < 1 || knots.is_empty() {
5462 return Err(FittedModelError::SchemaMismatch {
5463 reason: "saved survival bspline time basis metadata is invalid".to_string(),
5464 });
5465 }
5466 Ok(SurvivalTimeBasisConfig::BSpline {
5467 degree,
5468 knots: Array1::from_vec(knots),
5469 smooth_lambda,
5470 })
5471 }
5472 "ispline" => {
5473 let degree =
5474 model
5475 .survival_time_degree
5476 .ok_or_else(|| FittedModelError::MissingField {
5477 reason: "saved survival ispline model missing survival_time_degree"
5478 .to_string(),
5479 })?;
5480 let knots = model.survival_time_knots.clone().ok_or_else(|| {
5481 FittedModelError::MissingField {
5482 reason: "saved survival ispline model missing survival_time_knots".to_string(),
5483 }
5484 })?;
5485 let keep_cols = model.survival_time_keep_cols.clone().ok_or_else(|| {
5486 FittedModelError::MissingField {
5487 reason: "saved survival ispline model missing survival_time_keep_cols"
5488 .to_string(),
5489 }
5490 })?;
5491 let smooth_lambda = model.survival_time_smooth_lambda.unwrap_or(1e-2);
5492 if degree < 1 || knots.is_empty() || keep_cols.is_empty() {
5493 return Err(FittedModelError::SchemaMismatch {
5494 reason: "saved survival ispline time basis metadata is invalid".to_string(),
5495 });
5496 }
5497 Ok(SurvivalTimeBasisConfig::ISpline {
5498 degree,
5499 knots: Array1::from_vec(knots),
5500 keep_cols,
5501 smooth_lambda,
5502 })
5503 }
5504 other => Err(FittedModelError::IncompatibleConfig {
5505 reason: format!("unsupported saved survival_time_basis '{other}'"),
5506 }),
5507 }
5508}
5509
5510#[cfg(test)]
5511mod tests {
5512 use super::*;
5513 use crate::cubic_cell_kernel::ANCHORED_DEVIATION_KERNEL;
5514 use crate::survival::lognormal_kernel::FrailtySpec;
5515 use gam_data::SchemaColumn;
5516 use gam_problem::types::{LikelihoodScaleMetadata, LogLikelihoodNormalization};
5517 use gam_solve::estimate::{FitArtifacts, FittedBlock, FittedLinkState};
5518 use gam_solve::pirls::PirlsStatus;
5519 use ndarray::{Array1, Array2, array};
5520
5521 fn empty_termspec() -> TermCollectionSpec {
5522 TermCollectionSpec {
5523 linear_terms: vec![],
5524 random_effect_terms: vec![],
5525 smooth_terms: vec![],
5526 }
5527 }
5528
5529 fn transformation_normal_payload(version: u32, fit: UnifiedFitResult) -> FittedModelPayload {
5534 let mut payload = FittedModelPayload::new(
5535 version,
5536 "y ~ s(x)".to_string(),
5537 ModelKind::TransformationNormal,
5538 FittedFamily::TransformationNormal {
5539 likelihood: LikelihoodSpec::gaussian_identity(),
5540 },
5541 "transformation-normal".to_string(),
5542 );
5543 payload.fit_result = Some(fit.clone());
5544 payload.unified = Some(fit);
5545 payload.data_schema = Some(DataSchema {
5546 columns: vec![
5547 SchemaColumn {
5548 name: "y".to_string(),
5549 kind: ColumnKindTag::Continuous,
5550 levels: vec![],
5551 },
5552 SchemaColumn {
5553 name: "x".to_string(),
5554 kind: ColumnKindTag::Continuous,
5555 levels: vec![],
5556 },
5557 ],
5558 });
5559 payload.set_training_feature_metadata(vec!["x".to_string()], vec![(0.0, 1.0)]);
5560 payload.resolved_termspec = Some(empty_termspec());
5561 let knots = vec![0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0];
5562 payload.transformation_response_knots = Some(knots.clone());
5563 payload.transformation_response_transform = Some(vec![vec![1.0]]);
5564 payload.transformation_response_degree = Some(2);
5565 payload.transformation_response_median = Some(0.5);
5566 payload.transformation_score_calibration =
5567 Some(TransformationScoreCalibration::finite_support_pit());
5568 payload.transformation_geometry = Some(SavedTransformationNormalGeometry {
5569 parameterization: TransformationNormalParameterization::DirectAlpha,
5570 response_degree: 2,
5571 response_knot_count: knots.len(),
5572 shape_coordinate_count: 3,
5573 cone_carrier_covariate_width: 2,
5574 cone_carrier_row_count: 16,
5575 certified_response_support: (0.0, 1.0),
5576 response_median: 0.5,
5577 });
5578 payload.transformation_cone_carrier = Some(
5580 (0..16 * 2).map(|i| 1.0 + 0.01 * i as f64).collect(),
5581 );
5582 payload
5583 }
5584
5585 fn transformation_normal_fit() -> UnifiedFitResult {
5586 saved_fit(vec![FittedBlock {
5587 beta: Array1::from_vec(vec![0.1, 0.2, -0.3]),
5588 role: BlockRole::Mean,
5589 edf: 1.0,
5590 lambdas: Array1::zeros(0),
5591 }])
5592 }
5593
5594 #[test]
5597 fn transformation_normal_geometry_round_trips_and_validates() {
5598 let payload = transformation_normal_payload(MODEL_PAYLOAD_VERSION, transformation_normal_fit());
5599 let model = FittedModel::from_payload(payload);
5600 model
5601 .validate_for_persistence()
5602 .expect("CTN model carrying the direct-α geometry record validates");
5603
5604 let json = serde_json::to_string(&model).expect("serialize CTN model");
5605 let restored: FittedModel = serde_json::from_str(&json).expect("parse CTN model");
5606 restored
5607 .validate_for_persistence()
5608 .expect("restored CTN model validates");
5609 let geometry = restored
5610 .payload()
5611 .transformation_geometry
5612 .as_ref()
5613 .expect("restored payload carries the direct-α geometry record");
5614 assert_eq!(
5615 geometry.parameterization,
5616 TransformationNormalParameterization::DirectAlpha
5617 );
5618 assert_eq!(geometry.response_degree, 2);
5619 assert_eq!(geometry.response_knot_count, 7);
5620 assert_eq!(geometry.shape_coordinate_count, 3);
5621 assert_eq!(geometry.cone_carrier_covariate_width, 2);
5622 assert_eq!(geometry.cone_carrier_row_count, 16);
5623 assert_eq!(geometry.certified_response_support, (0.0, 1.0));
5624 assert_eq!(geometry.response_median, 0.5);
5625 }
5626
5627 #[test]
5631 fn validate_for_persistence_rejects_ctn_without_geometry_record() {
5632 let mut payload =
5633 transformation_normal_payload(MODEL_PAYLOAD_VERSION, transformation_normal_fit());
5634 payload.transformation_geometry = None;
5635 let err = FittedModel::from_payload(payload)
5636 .validate_for_persistence()
5637 .expect_err("CTN model without the direct-α geometry record must be rejected");
5638 assert!(
5639 err.to_string().contains("transformation_geometry"),
5640 "message names the field: {err}"
5641 );
5642 assert!(
5643 err.to_string().contains("pre-cutover"),
5644 "message explains the pre-cutover rejection: {err}"
5645 );
5646
5647 let mut mismatched =
5651 transformation_normal_payload(MODEL_PAYLOAD_VERSION, transformation_normal_fit());
5652 if let Some(geometry) = mismatched.transformation_geometry.as_mut() {
5653 geometry.response_knot_count += 1;
5654 }
5655 let err = FittedModel::from_payload(mismatched)
5656 .validate_for_persistence()
5657 .expect_err("geometry disagreeing with the persisted knots must be rejected");
5658 assert!(
5659 err.to_string().contains("response_knot_count"),
5660 "message names the mismatch: {err}"
5661 );
5662 }
5663
5664 #[test]
5668 fn validate_for_persistence_rejects_ctn_without_or_with_mismatched_cone_carrier() {
5669 let mut missing =
5670 transformation_normal_payload(MODEL_PAYLOAD_VERSION, transformation_normal_fit());
5671 missing.transformation_cone_carrier = None;
5672 let err = FittedModel::from_payload(missing)
5673 .validate_for_persistence()
5674 .expect_err("CTN model without the cone carrier must be rejected");
5675 assert!(
5676 err.to_string().contains("transformation_cone_carrier"),
5677 "message names the missing field: {err}"
5678 );
5679
5680 let mut mismatched =
5681 transformation_normal_payload(MODEL_PAYLOAD_VERSION, transformation_normal_fit());
5682 mismatched.transformation_cone_carrier = Some(vec![0.0; 31]);
5684 let err = FittedModel::from_payload(mismatched)
5685 .validate_for_persistence()
5686 .expect_err("a cone carrier disagreeing with the geometry dimensions must be rejected");
5687 assert!(
5688 err.to_string().contains("cone carrier length"),
5689 "message names the dimension mismatch: {err}"
5690 );
5691 }
5692
5693 #[test]
5697 fn spline_scan_payload_round_trips_and_validates() {
5698 let x: Vec<f64> = (0..40).map(|i| i as f64 / 39.0).collect();
5699 let y: Vec<f64> = x.iter().map(|&v| (4.0 * v).sin() + 0.1 * v).collect();
5700 let w = vec![1.0_f64; x.len()];
5701 let fit = gam_solve::spline_scan::fit_spline_scan(&x, &y, &w, 2).expect("scan fit");
5702 let make_payload = || {
5703 crate::inference::model_payload_builders::assemble_spline_scan_payload(
5704 "y ~ s(x)".to_string(),
5705 "x".to_string(),
5706 &fit,
5707 DataSchema {
5708 columns: vec![
5709 SchemaColumn {
5710 name: "y".to_string(),
5711 kind: ColumnKindTag::Continuous,
5712 levels: vec![],
5713 },
5714 SchemaColumn {
5715 name: "x".to_string(),
5716 kind: ColumnKindTag::Continuous,
5717 levels: vec![],
5718 },
5719 ],
5720 },
5721 vec!["x".to_string()],
5722 vec![(0.0, 1.0)],
5723 )
5724 };
5725 let model = FittedModel::from_payload(make_payload());
5728 model
5729 .validate_for_persistence()
5730 .expect("scan model validates");
5731 model
5732 .validate_numeric_finiteness()
5733 .expect("scan model is finite");
5734
5735 let json = serde_json::to_string(&model).expect("serialize model");
5736 let restored: FittedModel = serde_json::from_str(&json).expect("parse model");
5737 restored
5738 .validate_for_persistence()
5739 .expect("restored scan model validates");
5740 let (column, replay) = restored
5741 .saved_spline_scan()
5742 .expect("restore scan fit")
5743 .expect("payload carries the scan representation");
5744 assert_eq!(column, "x");
5745 for &xq in &[-0.1, 0.0, 0.31, 0.5, 0.77, 1.0, 1.4] {
5746 let (m0, v0) = fit.predict(xq).expect("predict original");
5747 let (m1, v1) = replay.predict(xq).expect("predict replayed");
5748 assert_eq!(m0.to_bits(), m1.to_bits(), "mean drift at x={xq}");
5749 assert_eq!(v0.to_bits(), v1.to_bits(), "variance drift at x={xq}");
5750 }
5751
5752 let mut dense = make_payload();
5754 dense.spline_scan = None;
5755 let err = FittedModel::from_payload(dense)
5756 .validate_for_persistence()
5757 .expect_err("dense payload without fit_result must be rejected");
5758 assert!(err.to_string().contains("fit_result"));
5759
5760 let mut corrupt = make_payload();
5762 corrupt
5763 .spline_scan
5764 .as_mut()
5765 .expect("scan channel present")
5766 .state
5767 .knots
5768 .truncate(2);
5769 FittedModel::from_payload(corrupt)
5770 .validate_for_persistence()
5771 .expect_err("corrupt scan state must be rejected");
5772 let mut unnamed = make_payload();
5773 unnamed
5774 .spline_scan
5775 .as_mut()
5776 .expect("scan channel present")
5777 .feature_column
5778 .clear();
5779 FittedModel::from_payload(unnamed)
5780 .validate_for_persistence()
5781 .expect_err("missing feature column must be rejected");
5782 }
5783
5784 fn standard_gaussian_payload() -> FittedModelPayload {
5785 FittedModelPayload::new(
5786 MODEL_PAYLOAD_VERSION,
5787 "y ~ 1".to_string(),
5788 ModelKind::Standard,
5789 FittedFamily::Standard {
5790 likelihood: LikelihoodSpec::gaussian_identity(),
5791 link: Some(StandardLink::Identity),
5792 latent_cloglog_state: None,
5793 mixture_state: None,
5794 sas_state: None,
5795 },
5796 "gaussian".to_string(),
5797 )
5798 }
5799
5800 fn anchored_runtime(basis_dim: usize) -> SavedCompiledFlexBlock {
5801 SavedCompiledFlexBlock {
5802 kernel: ANCHORED_DEVIATION_KERNEL.to_string(),
5803 breakpoints: vec![-1.0, 1.0],
5804 basis_dim,
5805 span_c0: vec![vec![0.0; basis_dim]],
5806 span_c1: vec![vec![0.0; basis_dim]],
5807 span_c2: vec![vec![0.0; basis_dim]],
5808 span_c3: vec![vec![0.0; basis_dim]],
5809 anchor_correction: None,
5810 anchor_components: Vec::new(),
5811 }
5812 }
5813
5814 fn saved_fit(blocks: Vec<FittedBlock>) -> UnifiedFitResult {
5815 let p: usize = blocks.iter().map(|block| block.beta.len()).sum();
5816 UnifiedFitResult::try_from_parts(gam_solve::estimate::UnifiedFitResultParts {
5817 blocks,
5818 log_lambdas: Array1::zeros(0),
5819 lambdas: Array1::zeros(0),
5820 likelihood_family: Some(LikelihoodSpec::binomial_probit()),
5821 likelihood_scale: LikelihoodScaleMetadata::FixedDispersion { phi: 1.0 },
5824 log_likelihood_normalization: LogLikelihoodNormalization::Full,
5825 log_likelihood: 0.0,
5826 deviance: 0.0,
5827 reml_score: 0.0,
5828 stable_penalty_term: 0.0,
5829 penalized_objective: 0.0,
5830 used_device: false,
5831 outer_iterations: 0,
5832 outer_converged: true,
5833 outer_gradient_norm: None,
5834 standard_deviation: 1.0,
5835 covariance_conditional: Some(Array2::zeros((p, p))),
5836 covariance_corrected: Some(Array2::zeros((p, p))),
5837 inference: None,
5838 fitted_link: FittedLinkState::Standard(None),
5839 geometry: None,
5840 block_states: vec![],
5841 pirls_status: PirlsStatus::Converged,
5842 max_abs_eta: 0.0,
5843 constraint_kkt: None,
5844 artifacts: FitArtifacts {
5845 pirls: None,
5846 null_space_logdet: None,
5847 null_space_dim: None,
5848 survival_link_wiggle_knots: None,
5849 survival_link_wiggle_degree: None,
5850 criterion_certificate: None,
5851 rho_posterior_certificate: None,
5852 rho_posterior_escalation: None,
5853 rho_covariance: None,
5854 joint_log_lambdas: None,
5855 firth_bias_reduction: false,
5856 },
5857 inner_cycles: 0,
5858 })
5859 .expect("test fixture fit must assemble")
5860 }
5861
5862 fn standard_binomial_model(fit: UnifiedFitResult) -> FittedModel {
5863 let mut payload = FittedModelPayload::new(
5864 MODEL_PAYLOAD_VERSION,
5865 "y ~ 1".to_string(),
5866 ModelKind::Standard,
5867 FittedFamily::Standard {
5868 likelihood: LikelihoodSpec::binomial_probit(),
5869 link: Some(StandardLink::Probit),
5870 latent_cloglog_state: None,
5871 mixture_state: None,
5872 sas_state: None,
5873 },
5874 "binomial".to_string(),
5875 );
5876 payload.fit_result = Some(fit.clone());
5877 payload.unified = Some(fit);
5878 FittedModel::from_payload(payload)
5879 }
5880
5881 #[test]
5882 fn curved_link_persistence_rejects_mode_without_posterior_state() {
5883 let mut fit = saved_fit(vec![FittedBlock {
5884 beta: array![0.25],
5885 role: BlockRole::Mean,
5886 edf: 1.0,
5887 lambdas: Array1::zeros(0),
5888 }]);
5889 fit.covariance_conditional = None;
5890 fit.inference = None;
5891 fit.geometry = None;
5892
5893 let error = standard_binomial_model(fit)
5894 .validate_required_posterior_mean_state()
5895 .expect_err("a curved-link mode alone is not a persistable fitted model");
5896 assert!(error.to_string().contains("posterior mean"));
5897 assert!(error.to_string().contains("covariance or penalized precision"));
5898 }
5899
5900 #[test]
5901 fn curved_link_persistence_accepts_and_round_trips_saved_posterior_state() {
5902 let fit = saved_fit(vec![FittedBlock {
5903 beta: array![0.25],
5904 role: BlockRole::Mean,
5905 edf: 1.0,
5906 lambdas: Array1::zeros(0),
5907 }]);
5908 let model = standard_binomial_model(fit);
5909 model
5910 .validate_required_posterior_mean_state()
5911 .expect("joint conditional covariance completes the curved-link fit");
5912
5913 let json = serde_json::to_string(&model).expect("serialize fitted model");
5914 let restored: FittedModel = serde_json::from_str(&json).expect("restore fitted model");
5915 restored
5916 .validate_required_posterior_mean_state()
5917 .expect("posterior state must survive the saved-model wire format");
5918 assert_eq!(
5919 restored
5920 .payload()
5921 .fit_result
5922 .as_ref()
5923 .and_then(UnifiedFitResult::beta_covariance),
5924 model
5925 .payload()
5926 .fit_result
5927 .as_ref()
5928 .and_then(UnifiedFitResult::beta_covariance),
5929 );
5930 }
5931
5932 #[test]
5933 fn curved_link_persistence_accepts_factorizable_saved_precision() {
5934 let mut fit = saved_fit(vec![FittedBlock {
5935 beta: array![0.25],
5936 role: BlockRole::Mean,
5937 edf: 1.0,
5938 lambdas: Array1::zeros(0),
5939 }]);
5940 fit.covariance_conditional = None;
5941 fit.geometry = Some(gam_solve::estimate::FitGeometry {
5942 coefficient_gauge: gam_problem::gauge::Gauge::identity(&[1]),
5943 penalized_hessian: gam_problem::dispersion_cov::UnscaledPrecision::wrap(array![[2.0]]),
5944 constrained_posterior: None,
5945 working: None,
5946 });
5947
5948 standard_binomial_model(fit)
5949 .validate_required_posterior_mean_state()
5950 .expect("a same-frame strictly-SPD precision can reconstruct posterior covariance");
5951 }
5952
5953 #[test]
5954 fn curved_link_persistence_rejects_active_frame_precision_without_lifted_covariance() {
5955 let mut fit = saved_fit(vec![FittedBlock {
5956 beta: array![0.25],
5957 role: BlockRole::Mean,
5958 edf: 1.0,
5959 lambdas: Array1::zeros(0),
5960 }]);
5961 fit.covariance_conditional = None;
5962 fit.geometry = Some(gam_solve::estimate::FitGeometry {
5963 coefficient_gauge: gam_problem::gauge::Gauge::from_block_transforms(&[array![[2.0]]]),
5964 penalized_hessian: gam_problem::dispersion_cov::UnscaledPrecision::wrap(array![[2.0]]),
5965 constrained_posterior: None,
5966 working: None,
5967 });
5968
5969 let error = standard_binomial_model(fit)
5970 .validate_required_posterior_mean_state()
5971 .expect_err("active-frame precision cannot be paired with raw prediction rows");
5972 assert!(error.to_string().contains("active gauge"));
5973 assert!(error.to_string().contains("lifted"));
5974 }
5975
5976 fn marginal_slope_payload(version: u32, fit: UnifiedFitResult) -> FittedModelPayload {
5977 let mut payload = FittedModelPayload::new(
5978 version,
5979 "y ~ 1".to_string(),
5980 ModelKind::MarginalSlope,
5981 FittedFamily::MarginalSlope {
5982 likelihood: LikelihoodSpec::binomial_probit(),
5983 base_link: InverseLink::Standard(StandardLink::Probit),
5984 frailty: FrailtySpec::None,
5985 },
5986 "bernoulli-marginal-slope".to_string(),
5987 );
5988 payload.fit_result = Some(fit.clone());
5989 payload.unified = Some(fit);
5990 payload.data_schema = Some(DataSchema {
5991 columns: vec![SchemaColumn {
5992 name: "z".to_string(),
5993 kind: ColumnKindTag::Continuous,
5994 levels: vec![],
5995 }],
5996 });
5997 payload.set_training_feature_metadata(vec!["z".to_string()], vec![(0.0, 0.0)]);
5998 payload.resolved_termspec = Some(empty_termspec());
5999 payload.resolved_termspec_logslope = Some(empty_termspec());
6000 payload.formula_logslope = Some("1".to_string());
6001 payload.z_column = Some("z".to_string());
6002 payload.latent_z_normalization = Some(SavedLatentZNormalization { mean: 0.0, sd: 1.0 });
6003 payload.latent_measure = Some(LatentMeasureKind::StandardNormal);
6004 payload.marginal_baseline = Some(0.0);
6005 payload.logslope_baseline = Some(0.0);
6006 payload.link = Some(InverseLink::Standard(StandardLink::Probit));
6007 payload
6008 }
6009
6010 #[test]
6011 fn from_payload_synchronizes_used_device_from_saved_fit() {
6012 let mut fit = saved_fit(vec![
6013 FittedBlock {
6014 beta: Array1::from_vec(vec![0.25]),
6015 role: BlockRole::Mean,
6016 edf: 1.0,
6017 lambdas: Array1::zeros(0),
6018 },
6019 FittedBlock {
6020 beta: Array1::from_vec(vec![0.5]),
6021 role: BlockRole::Scale,
6022 edf: 1.0,
6023 lambdas: Array1::zeros(0),
6024 },
6025 ]);
6026 fit.used_device = true;
6027 let mut payload = marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
6028 payload.used_device = false;
6029
6030 let model = FittedModel::from_payload(payload);
6031
6032 assert!(model.payload().used_device);
6033 }
6034
6035 fn survival_marginal_slope_payload(version: u32, fit: UnifiedFitResult) -> FittedModelPayload {
6036 let mut payload = FittedModelPayload::new(
6037 version,
6038 "Surv(entry, exit, event) ~ 1".to_string(),
6039 ModelKind::Survival,
6040 FittedFamily::Survival {
6041 likelihood: LikelihoodSpec::royston_parmar(),
6042 survival_likelihood: Some("marginal-slope".to_string()),
6043 survival_distribution: Some(ResidualDistribution::Gaussian),
6044 frailty: FrailtySpec::None,
6045 },
6046 "survival".to_string(),
6047 );
6048 payload.fit_result = Some(fit.clone());
6049 payload.unified = Some(fit);
6050 payload.survival_likelihood = Some("marginal-slope".to_string());
6051 payload.survival_distribution = Some(ResidualDistribution::Gaussian);
6052 payload.latent_measure = Some(LatentMeasureKind::StandardNormal);
6053 payload.data_schema = Some(DataSchema {
6054 columns: vec![SchemaColumn {
6055 name: "z".to_string(),
6056 kind: ColumnKindTag::Continuous,
6057 levels: vec![],
6058 }],
6059 });
6060 payload.set_training_feature_metadata(vec!["z".to_string()], vec![(0.0, 0.0)]);
6061 payload.resolved_termspec = Some(empty_termspec());
6062 payload.resolved_termspec_logslope = Some(empty_termspec());
6063 payload.formula_logslope = Some("1".to_string());
6064 payload.z_column = Some("z".to_string());
6065 payload.latent_z_normalization = Some(SavedLatentZNormalization { mean: 0.0, sd: 1.0 });
6066 payload.survival_marginal_slope_score_covariance = Some(vec![vec![1.0]]);
6067 payload.logslope_baseline = Some(0.0);
6068 payload.link = Some(InverseLink::Standard(StandardLink::Probit));
6069 payload
6070 }
6071
6072 fn gamma_dispersion_location_scale_payload() -> FittedModelPayload {
6073 let mut payload = FittedModelPayload::new(
6079 MODEL_PAYLOAD_VERSION,
6080 "y ~ x".to_string(),
6081 ModelKind::LocationScale,
6082 FittedFamily::LocationScale {
6083 likelihood: LikelihoodSpec::gamma_log(),
6084 base_link: Some(InverseLink::Standard(StandardLink::Log)),
6085 },
6086 "gamma-location-scale".to_string(),
6087 );
6088 payload.data_schema = Some(DataSchema {
6089 columns: vec![
6090 SchemaColumn {
6091 name: "y".to_string(),
6092 kind: ColumnKindTag::Continuous,
6093 levels: vec![],
6094 },
6095 SchemaColumn {
6096 name: "x".to_string(),
6097 kind: ColumnKindTag::Continuous,
6098 levels: vec![],
6099 },
6100 ],
6101 });
6102 payload.set_training_feature_metadata(vec!["x".to_string()], vec![(-1.0, 1.0)]);
6103 payload.resolved_termspec = Some(empty_termspec());
6104 payload.resolved_termspec_noise = Some(empty_termspec());
6105 payload.formula_noise = Some("x".to_string());
6106 payload.beta_noise = Some(vec![0.0]);
6107 payload.link = Some(InverseLink::Standard(StandardLink::Log));
6108 payload
6109 }
6110
6111 #[test]
6118 fn dispersion_location_scale_payload_is_not_classified_binomial() {
6119 let model = FittedModel::from_payload(gamma_dispersion_location_scale_payload());
6120 assert_eq!(
6121 model.predict_model_class(),
6122 PredictModelClass::DispersionLocationScale,
6123 "Gamma dispersion location-scale must route through the dispersion \
6124 predictor, not the binomial threshold-scale class",
6125 );
6126 assert!(
6127 !matches!(
6128 model.predict_model_class(),
6129 PredictModelClass::BinomialLocationScale
6130 ),
6131 "dispersion location-scale must never be classified as binomial",
6132 );
6133
6134 for likelihood in [
6136 LikelihoodSpec::gamma_log(),
6137 LikelihoodSpec::new(
6138 ResponseFamily::NegativeBinomial {
6139 theta: 1.0,
6140 theta_fixed: false,
6141 },
6142 InverseLink::Standard(StandardLink::Log),
6143 ),
6144 LikelihoodSpec::new(
6145 ResponseFamily::Beta { phi: 1.0 },
6146 InverseLink::Standard(StandardLink::Logit),
6147 ),
6148 LikelihoodSpec::new(
6149 ResponseFamily::Tweedie { p: 1.5 },
6150 InverseLink::Standard(StandardLink::Log),
6151 ),
6152 ] {
6153 let mut payload = gamma_dispersion_location_scale_payload();
6154 payload.family_state = FittedFamily::LocationScale {
6155 base_link: Some(likelihood.link.clone()),
6156 likelihood: likelihood.clone(),
6157 };
6158 let model = FittedModel::from_payload(payload);
6159 assert_eq!(
6160 model.predict_model_class(),
6161 PredictModelClass::DispersionLocationScale,
6162 "dispersion family {:?} mis-classified",
6163 likelihood.response,
6164 );
6165 }
6166 }
6167
6168 #[test]
6169 fn axis_clip_leaves_numeric_random_effect_group_axis_unclipped() {
6170 let data = array![[100.0], [-100.0]];
6171 let col_map = HashMap::from([("g".to_string(), 0usize)]);
6172
6173 let mut plain_payload = standard_gaussian_payload();
6174 plain_payload.data_schema = Some(DataSchema {
6175 columns: vec![SchemaColumn {
6176 name: "g".to_string(),
6177 kind: ColumnKindTag::Continuous,
6178 levels: vec![],
6179 }],
6180 });
6181 plain_payload.set_training_feature_metadata(vec!["g".to_string()], vec![(0.0, 7.0)]);
6182 plain_payload.resolved_termspec = Some(empty_termspec());
6183 let plain = FittedModel::from_payload(plain_payload.clone());
6184 let clipped = plain
6185 .axis_clip_to_training_ranges(data.view(), &col_map)
6186 .expect("ordinary continuous axis should clip outside the training range");
6187 assert_eq!(clipped.column(0).to_vec(), vec![7.0, 0.0]);
6188
6189 let mut group_payload = plain_payload;
6190 let mut group_spec = empty_termspec();
6191 group_spec
6192 .random_effect_terms
6193 .push(gam_terms::smooth::RandomEffectTermSpec {
6194 name: "g".to_string(),
6195 feature_col: 0,
6196 drop_first_level: false,
6197 penalized: true,
6198 frozen_levels: Some(vec![0.0_f64.to_bits(), 7.0_f64.to_bits()]),
6199 lenient_unseen: true,
6200 });
6201 group_payload.resolved_termspec = Some(group_spec);
6202 let group_model = FittedModel::from_payload(group_payload);
6203
6204 assert_eq!(
6205 group_model.random_effect_group_columns(),
6206 HashSet::from(["g".to_string()])
6207 );
6208
6209 assert_eq!(
6210 group_model.axis_clip_to_training_ranges(data.view(), &col_map),
6211 None,
6212 "numeric group labels must reach RandomEffectOperator as unseen levels, not be clipped to boundary seen levels"
6213 );
6214 }
6215
6216 #[test]
6230 fn bare_categorical_fixed_factor_unseen_level_rejected_by_predict_encode() {
6231 use csv::StringRecord;
6232 use gam_data::{EncodedDataset, UnseenCategoryPolicy, encode_recordswith_schema};
6233 use gam_runtime::resource::ResourcePolicy;
6234 use gam_terms::inference::formula_dsl::parse_formula;
6235 use gam_terms::term_builder::build_termspec;
6236
6237 let train_schema = DataSchema {
6239 columns: vec![
6240 SchemaColumn {
6241 name: "y".to_string(),
6242 kind: ColumnKindTag::Continuous,
6243 levels: vec![],
6244 },
6245 SchemaColumn {
6246 name: "g".to_string(),
6247 kind: ColumnKindTag::Categorical,
6248 levels: vec!["a".to_string(), "b".to_string(), "c".to_string()],
6249 },
6250 ],
6251 };
6252 let train = EncodedDataset {
6253 headers: vec!["y".to_string(), "g".to_string()],
6254 values: Array2::from_shape_vec(
6255 (6, 2),
6256 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],
6257 )
6258 .expect("rectangular training frame"),
6259 schema: train_schema.clone(),
6260 column_kinds: vec![ColumnKindTag::Continuous, ColumnKindTag::Categorical],
6261 };
6262 let build_col_map = train.column_map();
6263
6264 let model_for = |formula: &str| -> FittedModel {
6265 let parsed = parse_formula(formula).expect("formula parses");
6266 let mut notes = Vec::new();
6267 let spec = build_termspec(
6268 &parsed.terms,
6269 &train,
6270 &build_col_map,
6271 &mut notes,
6272 &ResourcePolicy::default_library(),
6273 )
6274 .unwrap_or_else(|err| panic!("`{formula}` must build a term spec, got: {err:?}"));
6275 let mut payload = standard_gaussian_payload();
6276 payload.data_schema = Some(train_schema.clone());
6277 payload.set_training_feature_metadata(
6278 vec!["y".to_string(), "g".to_string()],
6279 vec![(0.0, 2.0), (0.0, 2.0)],
6280 );
6281 payload.resolved_termspec = Some(spec);
6282 FittedModel::from_payload(payload)
6283 };
6284
6285 let g_schema = DataSchema {
6287 columns: vec![SchemaColumn {
6288 name: "g".to_string(),
6289 kind: ColumnKindTag::Categorical,
6290 levels: vec!["a".to_string(), "b".to_string(), "c".to_string()],
6291 }],
6292 };
6293 let encode_level = |model: &FittedModel, level: &str| -> Result<EncodedDataset, String> {
6294 let policy = UnseenCategoryPolicy::encode_unknown_for_columns(
6295 model.random_effect_group_columns(),
6296 );
6297 encode_recordswith_schema(
6298 vec!["g".to_string()],
6299 vec![StringRecord::from(vec![level])],
6300 &g_schema,
6301 policy,
6302 )
6303 };
6304
6305 let bare = model_for("y ~ g");
6308 assert!(
6309 !bare.random_effect_group_columns().contains("g"),
6310 "bare `+ g` is a fixed parametric factor; it must NOT be whitelisted for lenient \
6311 unseen-level encoding (#2102)"
6312 );
6313 encode_level(&bare, "a").expect("a seen level must still encode for the fixed factor");
6314 let err = encode_level(&bare, "TYPO")
6315 .expect_err("an unseen fixed-factor level must raise a schema mismatch (#2102)");
6316 assert!(
6317 err.contains("unseen level"),
6318 "expected an unseen-level schema mismatch naming the level, got: {err}"
6319 );
6320
6321 let factor = model_for("y ~ factor(g)");
6326 assert!(
6327 !factor.random_effect_group_columns().contains("g"),
6328 "factor(g) is a FIXED categorical factor; it must NOT be whitelisted for lenient \
6329 unseen-level encoding (#2137)"
6330 );
6331 encode_level(&factor, "a").expect("a seen level must still encode for factor(g)");
6332 let factor_err = encode_level(&factor, "TYPO")
6333 .expect_err("an unseen factor(g) level must raise a schema mismatch (#2137)");
6334 assert!(
6335 factor_err.contains("unseen level"),
6336 "expected an unseen-level schema mismatch naming the level, got: {factor_err}"
6337 );
6338
6339 for formula in ["y ~ group(g)", "y ~ re(g)", "y ~ s(g, bs=\"re\")"] {
6342 let grouped = model_for(formula);
6343 assert!(
6344 grouped.random_effect_group_columns().contains("g"),
6345 "`{formula}` is a random effect; it must remain lenient on unseen levels \
6346 (held-out-group policy)"
6347 );
6348 encode_level(&grouped, "TYPO").unwrap_or_else(|err| {
6349 panic!("`{formula}` must tolerate an unseen level, got: {err}")
6350 });
6351 }
6352 }
6353
6354 #[test]
6355 fn validate_for_persistence_rejects_marginal_slope_score_warp_basis_mismatch() {
6356 let fit = saved_fit(vec![
6357 FittedBlock {
6358 beta: array![0.1],
6359 role: BlockRole::Mean,
6360 edf: 1.0,
6361 lambdas: Array1::zeros(0),
6362 },
6363 FittedBlock {
6364 beta: array![0.2],
6365 role: BlockRole::Scale,
6366 edf: 1.0,
6367 lambdas: Array1::zeros(0),
6368 },
6369 FittedBlock {
6370 beta: array![0.3],
6371 role: BlockRole::Mean,
6372 edf: 1.0,
6373 lambdas: Array1::zeros(0),
6374 },
6375 ]);
6376 let mut payload = marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
6377 payload.score_warp_runtime = Some(anchored_runtime(2));
6378
6379 let err = FittedModel::from_payload(payload)
6380 .validate_for_persistence()
6381 .expect_err("marginal-slope score-warp basis mismatch should fail validation");
6382 assert!(err.to_string().contains("score-warp coefficient mismatch"));
6383 }
6384
6385 #[test]
6386 fn saved_prediction_runtime_rejects_survival_marginal_slope_link_basis_mismatch() {
6387 let fit = saved_fit(vec![
6388 FittedBlock {
6389 beta: array![0.1],
6390 role: BlockRole::Time,
6391 edf: 1.0,
6392 lambdas: Array1::zeros(0),
6393 },
6394 FittedBlock {
6395 beta: array![0.2],
6396 role: BlockRole::Mean,
6397 edf: 1.0,
6398 lambdas: Array1::zeros(0),
6399 },
6400 FittedBlock {
6401 beta: array![0.3],
6402 role: BlockRole::Scale,
6403 edf: 1.0,
6404 lambdas: Array1::zeros(0),
6405 },
6406 FittedBlock {
6407 beta: array![0.4],
6408 role: BlockRole::LinkWiggle,
6409 edf: 1.0,
6410 lambdas: Array1::zeros(0),
6411 },
6412 ]);
6413 let mut payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
6414 payload.link_deviation_runtime = Some(anchored_runtime(2));
6415
6416 let err = FittedModel::from_payload(payload)
6417 .saved_prediction_runtime()
6418 .expect_err(
6419 "survival marginal-slope link basis mismatch should fail runtime validation",
6420 );
6421 assert!(
6422 err.to_string()
6423 .contains("link-deviation coefficient mismatch")
6424 );
6425 }
6426
6427 #[test]
6428 fn apply_survival_time_basis_writes_all_required_fields() {
6429 use crate::survival::construction::SavedSurvivalTimeBasis;
6430
6431 let fit = saved_fit(vec![
6432 FittedBlock {
6433 beta: array![0.1],
6434 role: BlockRole::Time,
6435 edf: 1.0,
6436 lambdas: Array1::zeros(0),
6437 },
6438 FittedBlock {
6439 beta: array![0.2],
6440 role: BlockRole::Mean,
6441 edf: 1.0,
6442 lambdas: Array1::zeros(0),
6443 },
6444 FittedBlock {
6445 beta: array![0.3],
6446 role: BlockRole::Scale,
6447 edf: 1.0,
6448 lambdas: Array1::zeros(0),
6449 },
6450 ]);
6451 let mut payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
6452
6453 let snapshot = SavedSurvivalTimeBasis {
6458 basisname: "royston-parmar".to_string(),
6459 degree: Some(3),
6460 knots: Some(vec![0.0, 1.0, 2.0]),
6461 keep_cols: Some(vec![0, 2]),
6462 smooth_lambda: Some(0.5),
6463 anchor: 0.25,
6464 };
6465 payload.apply_survival_time_basis(&snapshot);
6466
6467 assert_eq!(
6468 payload.survival_time_basis.as_deref(),
6469 Some("royston-parmar")
6470 );
6471 assert_eq!(payload.survival_time_degree, Some(3));
6472 assert_eq!(payload.survival_time_knots, Some(vec![0.0, 1.0, 2.0]));
6473 assert_eq!(payload.survival_time_keep_cols, Some(vec![0, 2]));
6474 assert_eq!(payload.survival_time_smooth_lambda, Some(0.5));
6475 assert_eq!(payload.survival_time_anchor, Some(0.25));
6476 }
6477
6478 #[test]
6479 fn validate_for_persistence_rejects_survival_without_time_anchor_metadata() {
6480 let fit = saved_fit(vec![
6481 FittedBlock {
6482 beta: array![0.1],
6483 role: BlockRole::Time,
6484 edf: 1.0,
6485 lambdas: Array1::zeros(0),
6486 },
6487 FittedBlock {
6488 beta: array![0.2],
6489 role: BlockRole::Mean,
6490 edf: 1.0,
6491 lambdas: Array1::zeros(0),
6492 },
6493 FittedBlock {
6494 beta: array![0.3],
6495 role: BlockRole::Scale,
6496 edf: 1.0,
6497 lambdas: Array1::zeros(0),
6498 },
6499 ]);
6500 let mut payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
6501 payload.survival_time_basis = Some("ispline".to_string());
6507
6508 let err = FittedModel::from_payload(payload)
6509 .validate_for_persistence()
6510 .expect_err("survival model without time-anchor metadata should fail validation");
6511 assert!(err.to_string().contains("missing survival_time_anchor"));
6512 }
6513
6514 #[test]
6515 fn validate_for_persistence_rejects_survival_without_time_basis_metadata() {
6516 let fit = saved_fit(vec![
6517 FittedBlock {
6518 beta: array![0.1],
6519 role: BlockRole::Time,
6520 edf: 1.0,
6521 lambdas: Array1::zeros(0),
6522 },
6523 FittedBlock {
6524 beta: array![0.2],
6525 role: BlockRole::Mean,
6526 edf: 1.0,
6527 lambdas: Array1::zeros(0),
6528 },
6529 FittedBlock {
6530 beta: array![0.3],
6531 role: BlockRole::Scale,
6532 edf: 1.0,
6533 lambdas: Array1::zeros(0),
6534 },
6535 ]);
6536 let payload = survival_marginal_slope_payload(MODEL_PAYLOAD_VERSION, fit);
6537
6538 let err = FittedModel::from_payload(payload)
6539 .validate_for_persistence()
6540 .expect_err("survival model without time-basis metadata should fail validation");
6541 assert!(err.to_string().contains("missing survival_time_basis"));
6542 }
6543
6544 #[test]
6545 fn saved_prediction_runtime_rejects_stale_payload_version() {
6546 let fit = saved_fit(vec![
6547 FittedBlock {
6548 beta: array![0.1],
6549 role: BlockRole::Mean,
6550 edf: 1.0,
6551 lambdas: Array1::zeros(0),
6552 },
6553 FittedBlock {
6554 beta: array![0.2],
6555 role: BlockRole::Scale,
6556 edf: 1.0,
6557 lambdas: Array1::zeros(0),
6558 },
6559 ]);
6560 let payload = marginal_slope_payload(MODEL_PAYLOAD_VERSION - 1, fit);
6561
6562 let err = FittedModel::from_payload(payload)
6563 .saved_prediction_runtime()
6564 .expect_err("stale payload version should fail before runtime assembly");
6565 assert!(err.to_string().contains("payload schema mismatch"));
6566 }
6567
6568 #[test]
6569 fn saved_link_wiggle_warp_index_applies_exact_2141_mean_shift() {
6570 let runtime = SavedLinkWiggleRuntime {
6571 knots: vec![],
6572 degree: 0,
6573 penalty_metadata: None,
6574 beta: vec![],
6575 index_shift: Some(vec![0.25, -0.5]),
6576 };
6577 let design = DesignMatrix::from(array![[1.0, 2.0], [-3.0, 0.5]]);
6578 let base = array![0.75, -0.25];
6579 let index = runtime
6580 .warp_index(&base, &design)
6581 .expect("complete saved shift");
6582 assert_eq!(index, array![0.0, -1.25]);
6583 }
6584
6585 #[test]
6586 fn saved_link_wiggle_warp_index_rejects_partial_shift_coordinates() {
6587 let runtime = SavedLinkWiggleRuntime {
6588 knots: vec![],
6589 degree: 0,
6590 penalty_metadata: None,
6591 beta: vec![],
6592 index_shift: Some(vec![0.25]),
6593 };
6594 let design = DesignMatrix::from(array![[1.0, 2.0]]);
6595 let error = runtime
6596 .warp_index(&array![0.5], &design)
6597 .expect_err("partial #2141 shift metadata must fail loudly");
6598 assert!(error.to_string().contains("shift has 1 entries"));
6599 assert!(error.to_string().contains("mean design has 2 columns"));
6600 }
6601}