Skip to main content

gam_models/fit_orchestration/drivers/
mod.rs

1// GAM fit-orchestration drivers, relocated from `gam-terms/src/smooth/`
2// (`design_construction.rs` + `spatial_optimization.rs`) up into `gam-models`
3// per #1521. They were `include!`d into `gam_terms::smooth` (one flat module
4// alongside `prelude.rs` + `term_specs.rs`); to preserve that single-module
5// flat namespace (and the heavy cross-references between the two files) byte
6// for byte, they are `include!`d here as well. The shared import surface that
7// `prelude.rs`/`term_specs.rs` used to provide is reconstructed below with the
8// relocated paths (families now resolve as `crate::*`, the solver as
9// `gam_solve::*`, basis/term machinery as `gam_terms::*`).
10use gam_terms::basis::{
11    BasisError, BasisMetadata, BasisPsiDerivativeResult, BasisPsiSecondDerivativeResult,
12    BasisWorkspace, CenterStrategy, MaternIdentifiability, PenaltyInfo,
13    PenaltySource, build_constant_curvature_basis_kappa_derivatives,
14    build_matern_basis_log_kappa_aniso_derivatives, build_matern_basis_log_kappa_derivatives,
15    build_matern_collocation_operator_matrices, build_measure_jet_basis_psi_derivatives,
16    build_thin_plate_basis_log_kappa_derivatives, estimate_penalty_nullity, initial_aniso_contrasts,
17};
18
19use gam_custom_family::{
20    BlockEffectiveJacobian, BlockGeometryDirectionalDerivative, BlockWorkingSet, BlockwiseFitOptions,
21    CustomFamily, CustomFamilyBlockPsiDerivative, CustomFamilyWarmStart, ExactNewtonJointPsiTerms,
22    ExactNewtonOuterObjective, FamilyEvaluation, FamilyLinearizationState, ParameterBlockSpec,
23    ParameterBlockState, PenaltyMatrix, evaluate_custom_family_joint_hyper,
24    evaluate_custom_family_joint_hyper_efs, fit_custom_family,
25};
26
27use gam_solve::estimate::{
28    EstimationError, ExternalOptimOptions, FitInference, FitOptions, FittedLinkState, PenaltySpec,
29    UnifiedFitResult, UnifiedFitResultParts, fit_gamwith_heuristic_lambdas,
30};
31
32use gam_solve::estimate::reml::DirectionalHyperParam;
33
34// #1521: `freeze_term_collection_from_design` relocated DOWN into gam_terms::smooth
35// (was an `include!`d `pub fn` in spatial_optimization.rs). Re-export here so the
36// `crate::fit_orchestration::drivers::freeze_term_collection_from_design` path used
37// by families + pyffi resolves unchanged.
38pub use gam_terms::smooth::freeze_term_collection_from_design;
39
40use crate::family_runtime::{FamilyStrategy, strategy_for_spec};
41
42use gam_solve::mixture_link::{
43    logit_inverse_link_jet5, state_from_beta_logisticspec, state_from_sasspec, state_fromspec,
44};
45
46use gam_math::quantile::quantile_from_sorted;
47
48use gam_linalg::faer_ndarray::{fast_ab, fast_atb, fast_atv};
49
50use gam_linalg::matrix::{
51    DesignBlock, DesignMatrix, RandomEffectOperator, SymmetricMatrix,
52};
53
54use gam_problem::LinearInequalityConstraints;
55
56use gam_spec::{
57    InverseLink, LatentCLogLogState, LikelihoodSpec, MixtureLinkState, ResponseFamily, SasLinkState,
58    StandardLink,
59};
60
61use gam_terms::smooth::input_standardization::{
62    apply_input_standardization, compensate_length_scale_for_standardization,
63    compensate_optional_length_scale_for_standardization,
64};
65
66use gam_terms::smooth::penalty_priors::{
67    realize_keyed_penalty_block_gamma_priors,
68    realize_penalty_block_gamma_priors,
69};
70
71use gam_terms::smooth::shape_constraints::{
72    linear_constraints_from_lower_bounds_global, merge_linear_constraints_global,
73    shape_lower_bounds_local,
74};
75
76// Every `pub` item that `gam_terms::smooth` exposes (the `term_specs.rs`
77// spec/design machinery, `SmoothError`, the `penalty_priors`/`structure_analysis`
78// re-exports, …). This reconstructs the sibling-module visibility the drivers
79// had while textually pasted inside `gam_terms::smooth`.
80use gam_terms::smooth::*;
81
82use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis, s};
83
84use std::collections::BTreeSet;
85use std::ops::Range;
86use std::sync::atomic::AtomicUsize;
87use std::sync::{Arc, Mutex};
88
89// Fit-result carriers relocated out of `gam_terms::smooth::term_specs` with the
90// drivers (they hold a `gam_solve` `UnifiedFitResult` and are consumed only by
91// the drivers / the surrounding fit-orchestration layer).
92#[derive(Clone)]
93pub struct FittedTermCollection {
94    pub fit: UnifiedFitResult,
95    pub design: TermCollectionDesign,
96    pub adaptive_diagnostics: Option<AdaptiveRegularizationDiagnostics>,
97}
98
99#[derive(Clone, Copy, Debug, Default)]
100pub struct SpatialLengthScaleOptimizationTiming {
101    pub log_kappa_dim: usize,
102    pub cost_calls: usize,
103    pub cost_total_s: f64,
104    pub eval_calls: usize,
105    pub eval_total_s: f64,
106    pub efs_calls: usize,
107    pub efs_total_s: f64,
108    pub slow_path_resets: u64,
109    pub design_revision_delta: u64,
110    pub nfree_miss_shape: u64,
111    pub nfree_miss_value: u64,
112    pub nfree_miss_gradient: u64,
113    pub nfree_miss_penalty: u64,
114    pub nfree_miss_revision: u64,
115    pub nfree_miss_second_order: u64,
116    pub nfree_miss_other: u64,
117    pub optim_total_s: f64,
118}
119
120impl SpatialLengthScaleOptimizationTiming {
121    pub fn trial_total_s(self) -> f64 {
122        self.cost_total_s + self.eval_total_s + self.efs_total_s
123    }
124}
125
126#[derive(Clone)]
127pub struct FittedTermCollectionWithSpec {
128    pub fit: UnifiedFitResult,
129    pub design: TermCollectionDesign,
130    pub resolvedspec: TermCollectionSpec,
131    pub adaptive_diagnostics: Option<AdaptiveRegularizationDiagnostics>,
132    pub kappa_timing: Option<SpatialLengthScaleOptimizationTiming>,
133}
134
135include!("design_construction.rs");
136include!("spatial_optimization.rs");