gam_models/fit_orchestration/entry.rs
1use super::*;
2
3/// Request-specific inputs to the canonical standard-fit `FitOptions`.
4///
5/// Everything in here varies per call (the link state extracted from the
6/// formula/config, the linear constraints synthesized from `bounded()` /
7/// shape-constrained terms, the Firth / adaptive-regularization toggles read
8/// off the `FitConfig`). Every *policy* field of `FitOptions` — the ones that
9/// decide HOW the outer REML optimization behaves (`compute_inference`,
10/// `skip_rho_posterior_inference`, `tol`, the `max_iter` default, the penalty
11/// shrinkage floor) — is filled in by [`canonical_standard_fit_options`] and is
12/// NOT settable here, so the CLI binary and the Python/PyO3 path cannot resolve
13/// a different optimization policy for the same model (#1196). Before this seam
14/// existed the CLI hand-built `FitOptions` with `tol: 1e-6` /
15/// `skip_rho_posterior_inference: false` while the formula path used
16/// `tol: 1e-10` / `skip_rho_posterior_inference: true`, so the identical model
17/// fit *differently* depending on which entry point you called it from — the
18/// exact class of divergence #1191 surfaced.
19#[derive(Default)]
20pub struct StandardFitOptionsInputs {
21 pub latent_cloglog: Option<LatentCLogLogState>,
22 pub mixture_link: Option<MixtureLinkSpec>,
23 pub optimize_mixture: bool,
24 pub sas_link: Option<SasLinkSpec>,
25 pub optimize_sas: bool,
26 pub linear_constraints: Option<gam_solve::pirls::LinearInequalityConstraints>,
27 pub firth_bias_reduction: bool,
28 pub adaptive_regularization: Option<AdaptiveRegularizationOptions>,
29 /// `Some` only when a caller (the forced-Firth CLI branch) overrides the
30 /// canonical default. `None` keeps the single-source default `Some(1e-6)`.
31 pub penalty_shrinkage_floor_override: Option<Option<f64>>,
32 pub persist_warm_start_disk: bool,
33}
34
35/// The single source of truth for standard-fit `FitOptions` *policy*.
36///
37/// Both standard-fit entry points — `materialize_standard` (the formula /
38/// Python / PyO3 path) and the `gam` CLI's `run_fit` — construct their
39/// `StandardFitRequest` options through this function, so the outer REML
40/// optimization policy (`compute_inference`, `skip_rho_posterior_inference`,
41/// `tol`, `max_iter` default, `penalty_shrinkage_floor`) is identical by
42/// construction. New policy fields must be set HERE, never re-derived at a call
43/// site, which is what makes Python/CLI behavioral divergence structurally
44/// impossible rather than enforced by parallel-but-equal code (#1196).
45pub fn canonical_standard_fit_options(
46 config: &FitConfig,
47 inputs: StandardFitOptionsInputs,
48) -> FitOptions {
49 FitOptions {
50 latent_cloglog: inputs.latent_cloglog,
51 mixture_link: inputs.mixture_link,
52 optimize_mixture: inputs.optimize_mixture,
53 sas_link: inputs.sas_link,
54 optimize_sas: inputs.optimize_sas,
55 // Posterior covariance is always computed so `predict --uncertainty`
56 // works for every family (the `COV_MAX_P` diagonal fallback caps cost).
57 compute_inference: true,
58 // Formula/CLI fits are the interactive/default path: keep coefficient
59 // covariance and the smoothing correction, but do not run the optional
60 // live-rho posterior certificate/escalation, which can launch NUTS over
61 // rho and turn ordinary fits into sampler benchmarks. Lower-level
62 // callers that explicitly need the rho posterior opt in elsewhere.
63 skip_rho_posterior_inference: true,
64 max_iter: config.outer_max_iter.unwrap_or(200),
65 // Outer REML/LAML smoothing-selection tolerance. `1e-10` (effective
66 // projected-gradient threshold ≈ 1e-7) resolves λ̂ to optimiser
67 // precision and restores the `w=c ⇔ c-fold replication` invariance in
68 // smoothing selection (gam#893). The CLI previously used the stale
69 // `1e-6`, which over-smoothed relative to the formula path.
70 tol: 1e-10,
71 nullspace_dims: vec![],
72 linear_constraints: inputs.linear_constraints,
73 firth_bias_reduction: inputs.firth_bias_reduction,
74 adaptive_regularization: inputs.adaptive_regularization,
75 penalty_shrinkage_floor: inputs
76 .penalty_shrinkage_floor_override
77 .unwrap_or(Some(1e-6)),
78 rho_prior: Default::default(),
79 kronecker_penalty_system: None,
80 kronecker_factored: None,
81 persist_warm_start_disk: inputs.persist_warm_start_disk,
82 }
83}
84
85pub fn fit_model(request: FitRequest<'_>) -> Result<FitResult, WorkflowError> {
86 // Disk warm-start persistence is opt-in. The always-on in-memory warm start
87 // remains inside the fit engines, but the workflow dispatcher must not open
88 // the shared WarmStartStore for ordinary formula fits: refit-heavy quality
89 // tests get no cross-process reuse and previously paid cache lookup,
90 // checkpoint, and eviction scans on every replicate (#1082/#1114).
91 let request = request;
92 // Each `fit_*_model` helper still returns `Result<_, String>` internally;
93 // the boundary conversion happens here so the public API returns
94 // `WorkflowError::IntegrationFailed` carrying the underlying solver text.
95 let wrap_solver_err =
96 |reason: String| -> WorkflowError { WorkflowError::IntegrationFailed { reason } };
97 match request {
98 FitRequest::Standard(request) => fit_standard_model(request)
99 .map(FitResult::Standard)
100 .map_err(wrap_solver_err),
101 FitRequest::GaussianLocationScale(request) => fit_gaussian_location_scale_model(request)
102 .map(FitResult::GaussianLocationScale)
103 .map_err(wrap_solver_err),
104 FitRequest::BinomialLocationScale(request) => fit_binomial_location_scale_model(request)
105 .map(FitResult::BinomialLocationScale)
106 .map_err(wrap_solver_err),
107 FitRequest::DispersionLocationScale(request) => {
108 fit_dispersion_location_scale_model(request)
109 .map(FitResult::DispersionLocationScale)
110 .map_err(wrap_solver_err)
111 }
112 FitRequest::SurvivalLocationScale(request) => fit_survival_location_scale_model(request)
113 .map(FitResult::SurvivalLocationScale)
114 .map_err(wrap_solver_err),
115 FitRequest::SurvivalTransformation(request) => fit_survival_transformation_model(request)
116 .map(FitResult::SurvivalTransformation)
117 .map_err(wrap_solver_err),
118 FitRequest::BernoulliMarginalSlope(request) => fit_bernoulli_marginal_slope_model(request)
119 .map(FitResult::BernoulliMarginalSlope)
120 .map_err(wrap_solver_err),
121 FitRequest::SurvivalMarginalSlope(request) => fit_survival_marginal_slope_model(request)
122 .map(FitResult::SurvivalMarginalSlope)
123 .map_err(wrap_solver_err),
124 FitRequest::LatentSurvival(request) => fit_latent_survival_model(request)
125 .map(FitResult::LatentSurvival)
126 .map_err(wrap_solver_err),
127 FitRequest::LatentBinary(request) => fit_latent_binary_model(request)
128 .map(FitResult::LatentBinary)
129 .map_err(wrap_solver_err),
130 FitRequest::TransformationNormal(request) => fit_transformation_normal_model(request)
131 .map(FitResult::TransformationNormal)
132 .map_err(wrap_solver_err),
133 }
134}
135/// Resolve the [`gam_runtime::resource::ResourcePolicy`] backing term construction
136/// for a given [`FitConfig`] + dataset.
137///
138/// If the caller hasn't supplied an explicit policy override, derive one from
139/// the shape of the problem via
140/// [`gam_runtime::resource::ResourcePolicy::for_problem`]. At large scale (n_rows
141/// >= 100k or the marginal-slope large-scale path active) this returns
142/// > `analytic_operator_required` so that any silent dense materialization in
143/// > the term-construction layer fails fast rather than allocating tens of GiB;
144/// > at small scale it falls through to the permissive default-library policy
145/// > so that non-operator bases still build cleanly.
146///
147/// `p_estimate = 0` because the per-block coefficient count isn't known until
148/// the spec has been built; the n_rows and hints triggers are sufficient to
149/// flip strict mode for every shape that needs it.
150pub(crate) fn resolved_resource_policy(
151 config: &FitConfig,
152 data: &Dataset,
153 hints: gam_runtime::resource::ProblemHints,
154) -> gam_runtime::resource::ResourcePolicy {
155 if let Some(p) = config.resource_policy.clone() {
156 return p;
157 }
158 gam_runtime::resource::ResourcePolicy::for_problem(data.values.nrows(), 0, hints)
159}
160
161pub(crate) fn marginal_slope_hints(config: &FitConfig) -> gam_runtime::resource::ProblemHints {
162 gam_runtime::resource::ProblemHints {
163 marginal_slope_large_scale_active: config.logslope_formula.is_some()
164 || config.z_column.is_some(),
165 }
166}
167/// Parse, materialize, and fit a model in one call.
168/// Resolve the expectile asymmetry `τ` requested by `config`, if any.
169///
170/// Returns `Ok(Some(τ))` when `config.family` is `"expectile"` (optionally with
171/// an inline asymmetry, `"expectile(0.9)"`), `Ok(None)` for every other family,
172/// and `Err` when an expectile request carries an out-of-range `τ`. The inline
173/// form takes precedence over the explicit [`FitConfig::expectile_tau`] field
174/// only when both are present and disagree is rejected as a contradiction; when
175/// neither pins `τ`, the median expectile `τ = 0.5` (the ordinary mean fit) is
176/// the default.
177fn expectile_tau_for_config(config: &FitConfig) -> Result<Option<f64>, WorkflowError> {
178 let Some(raw) = config.family.as_deref() else {
179 return Ok(None);
180 };
181 let trimmed = raw.trim();
182 let lower = trimmed.to_ascii_lowercase();
183 if !(lower == "expectile" || lower.starts_with("expectile(")) {
184 return Ok(None);
185 }
186 let invalid = |reason: String| WorkflowError::InvalidConfig { reason };
187 // Optional inline asymmetry: `expectile(0.9)`.
188 let inline_tau = if let Some(rest) = lower.strip_prefix("expectile(") {
189 let inner = rest.strip_suffix(')').ok_or_else(|| {
190 invalid(format!(
191 "expectile family asymmetry must be written as `expectile(τ)`; got `{trimmed}`"
192 ))
193 })?;
194 let value: f64 = inner.trim().parse().map_err(|_| {
195 invalid(format!(
196 "expectile asymmetry `{}` is not a finite number",
197 inner.trim()
198 ))
199 })?;
200 Some(value)
201 } else {
202 None
203 };
204 let tau = match (inline_tau, config.expectile_tau) {
205 (Some(a), Some(b)) if (a - b).abs() > 0.0 => {
206 return Err(invalid(format!(
207 "expectile asymmetry given both inline (`expectile({a})`) and via expectile_tau \
208 ({b}); supply exactly one"
209 )));
210 }
211 (Some(a), _) => a,
212 (None, Some(b)) => b,
213 (None, None) => 0.5,
214 };
215 if !(tau.is_finite() && tau > 0.0 && tau < 1.0) {
216 return Err(invalid(format!(
217 "expectile asymmetry τ must be finite and strictly in (0, 1); got {tau}"
218 )));
219 }
220 Ok(Some(tau))
221}
222
223/// Per-row asymmetric LAWS weight `wᵢ(τ) = τ` if `yᵢ > μᵢ` else `1 − τ`, scaled
224/// by the base prior weight. At the boundary `yᵢ = μᵢ` the two half-weights
225/// agree in the limit only at `τ = 0.5`; the convention `yᵢ > μᵢ ⇒ τ` (strict)
226/// matches Newey–Powell's lower-closed asymmetric loss and is what `expectreg`
227/// uses. The fixed point is independent of the tie convention because ties form
228/// a measure-zero set under any continuous response.
229fn expectile_row_weights(
230 y: ArrayView1<f64>,
231 mu: ArrayView1<f64>,
232 base: ArrayView1<f64>,
233 tau: f64,
234) -> Array1<f64> {
235 Array1::from_shape_fn(y.len(), |i| {
236 let asym = if y[i] > mu[i] { tau } else { 1.0 - tau };
237 base[i] * asym
238 })
239}
240
241fn constant_gaussian_standard_fit(
242 request: &StandardFitRequest<'_>,
243) -> Result<StandardFitResult, WorkflowError> {
244 if !request.family.is_gaussian_identity() || request.y.is_empty() {
245 return Err(WorkflowError::InvalidConfig {
246 reason: "constant Gaussian shortcut requires a non-empty Gaussian identity request"
247 .to_string(),
248 });
249 }
250 if request.y.iter().any(|value| !value.is_finite())
251 || request.offset.iter().any(|value| !value.is_finite())
252 || request
253 .weights
254 .iter()
255 .any(|value| !value.is_finite() || *value < 0.0)
256 {
257 return Err(WorkflowError::InvalidConfig {
258 reason: "constant Gaussian shortcut requires finite response, offset, and non-negative weights"
259 .to_string(),
260 });
261 }
262 let weight_sum = request.weights.sum();
263 if !(weight_sum.is_finite() && weight_sum > 0.0) {
264 return Err(WorkflowError::InvalidConfig {
265 reason: "constant Gaussian shortcut requires positive total weight".to_string(),
266 });
267 }
268 let mut centered_sum = 0.0_f64;
269 for i in 0..request.y.len() {
270 centered_sum += request.weights[i] * (request.y[i] - request.offset[i]);
271 }
272 let intercept = centered_sum / weight_sum;
273 let design =
274 build_term_collection_design(request.data.view(), &request.spec).map_err(|err| {
275 WorkflowError::InvalidConfig {
276 reason: format!("constant Gaussian shortcut could not rebuild design: {err}"),
277 }
278 })?;
279 let p = design.design.ncols();
280 let mut beta = Array1::<f64>::zeros(p);
281 for col in design.intercept_range.clone() {
282 if col < p {
283 beta[col] = intercept;
284 }
285 }
286 let lambdas = Array1::<f64>::ones(design.penalties.len());
287 let log_lambdas = Array1::<f64>::zeros(design.penalties.len());
288 let fit =
289 gam_solve::estimate::UnifiedFitResult::try_from_parts(gam_solve::estimate::UnifiedFitResultParts {
290 blocks: vec![gam_solve::estimate::FittedBlock {
291 beta: beta.clone(),
292 role: gam_problem::BlockRole::Mean,
293 edf: design.intercept_range.len() as f64,
294 lambdas: lambdas.clone(),
295 }],
296 log_lambdas,
297 lambdas,
298 likelihood_family: Some(request.family.clone()),
299 likelihood_scale: gam_problem::LikelihoodScaleMetadata::ProfiledGaussian,
300 log_likelihood_normalization: gam_problem::LogLikelihoodNormalization::UserProvided,
301 log_likelihood: 0.0,
302 deviance: 0.0,
303 reml_score: 0.0,
304 stable_penalty_term: 0.0,
305 penalized_objective: 0.0,
306 used_device: false,
307 outer_iterations: 0,
308 outer_converged: true,
309 outer_gradient_norm: Some(0.0),
310 standard_deviation: 0.0,
311 covariance_conditional: None,
312 covariance_corrected: None,
313 inference: None,
314 fitted_link: gam_solve::estimate::FittedLinkState::Standard(None),
315 geometry: None,
316 block_states: Vec::new(),
317 pirls_status: gam_solve::pirls::PirlsStatus::Converged,
318 max_abs_eta: intercept.abs(),
319 constraint_kkt: None,
320 artifacts: gam_solve::estimate::FitArtifacts {
321 pirls: None,
322 ..Default::default()
323 },
324 inner_cycles: 0,
325 })
326 .map_err(|err| WorkflowError::IntegrationFailed {
327 reason: format!("constant Gaussian shortcut produced invalid fit: {err}"),
328 })?;
329 let resolvedspec =
330 freeze_term_collection_from_design(&request.spec, &design).map_err(|err| {
331 WorkflowError::InvalidConfig {
332 reason: format!("constant Gaussian shortcut could not freeze design: {err}"),
333 }
334 })?;
335 Ok(StandardFitResult {
336 fit,
337 design,
338 resolvedspec,
339 adaptive_diagnostics: None,
340 kappa_timing: None,
341 saved_link_state: gam_solve::estimate::FittedLinkState::Standard(None),
342 wiggle_knots: None,
343 wiggle_degree: None,
344 wiggle_saved_warp_beta: None,
345 })
346}
347
348fn gaussian_response_is_constant(request: &StandardFitRequest<'_>) -> bool {
349 if !request.family.is_gaussian_identity()
350 || request.y.is_empty()
351 || request.y.iter().any(|value| !value.is_finite())
352 {
353 return false;
354 }
355 let (lo, hi) = request
356 .y
357 .iter()
358 .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &value| {
359 (lo.min(value), hi.max(value))
360 });
361 (hi - lo).abs() <= 1.0e-12 * hi.abs().max(1.0)
362}
363
364pub fn fit_from_formula(
365 formula: &str,
366 data: &Dataset,
367 config: &FitConfig,
368) -> Result<FitResult, WorkflowError> {
369 // Expectile regression (Newey–Powell asymmetric least squares): when the
370 // family resolves to "expectile", the τ-expectile of `y | x` is the
371 // minimizer of `Σ wᵢ(τ)·(yᵢ − μᵢ)²`, `wᵢ(τ) = τ` if `yᵢ > μᵢ` else `1 − τ`
372 // — the smooth analogue of the τ-quantile. The minimizer is a Least
373 // Asymmetrically Weighted Squares (LAWS) fixed point: iterate the penalized
374 // Gaussian-identity GAM with `wᵢ(τ)` recomputed from the current `μᵢ` until
375 // the residual-sign pattern stabilizes. REML λ-selection runs inside each
376 // inner Gaussian solve, so every gam smooth/tensor/spatial basis becomes a
377 // penalized expectile smooth with data-driven smoothing for free. This is a
378 // genuine estimator route, not a silent swap: it fires only on the explicit
379 // `family = "expectile"`. Every other family falls through unchanged.
380 if let Some(result) = fit_expectile_if_requested(formula, data, config)? {
381 return Ok(FitResult::Standard(result));
382 }
383 let mat = materialize(formula, data, config)?;
384 // Exact O(n) spline-scan fast path (#1030): when the materialized request
385 // is the single 1-D Gaussian-identity penalized-smooth shape the
386 // state-space scan solves exactly, route through it and return the
387 // scan-bearing model directly — the same penalized posterior at O(n) per
388 // λ-trial instead of the dense design/Gram route. Detection is structural
389 // and conservative (see `spline_scan_fast_path`); every other shape falls
390 // through to the dense `fit_model` path unchanged. Mirrors the CLI
391 // (main.rs run_fit) and FFI consumers, which build the persistence payload
392 // from this same `SplineScanFit`.
393 if let FitRequest::Standard(request) = &mat.request {
394 if gaussian_response_is_constant(request) {
395 return constant_gaussian_standard_fit(request).map(FitResult::Standard);
396 }
397 if let Some(inputs) = spline_scan_fast_path(request) {
398 let scan = gam_solve::spline_scan::fit_spline_scan(
399 &inputs.x,
400 &inputs.y,
401 &inputs.w,
402 inputs.order,
403 )
404 .map_err(|reason| WorkflowError::IntegrationFailed { reason })?;
405 return Ok(FitResult::SplineScan(scan));
406 }
407 // O(n log n) multiresolution residual-cascade fast path (#1032): a
408 // scattered low-d Gaussian-identity Duchon/Matérn smooth past the
409 // dense-kernel cliff. UNLIKE the scan, the cascade is a DIFFERENT
410 // posterior from the dense radial term, so it only ever fires as an
411 // explicit alternative estimator on the exact structural signature
412 // (`residual_cascade_fast_path`) AND when the in-cascade quasi-uniformity
413 // guard certifies the metric — a rejected metric or any ineligible shape
414 // falls through to the dense `fit_model` path (a genuine estimator
415 // choice, never a silent swap). The save paths build the persistence
416 // payload from this `ResidualCascadeFit`'s `to_state` snapshot.
417 if let Some(inputs) = residual_cascade_fast_path(request) {
418 let coord_refs: Vec<&[f64]> = inputs.coords.iter().map(Vec::as_slice).collect();
419 if let Ok(fit) = gam_solve::residual_cascade::fit_residual_cascade(
420 &coord_refs,
421 &inputs.y,
422 &inputs.w,
423 &inputs.metric,
424 inputs.sobolev_s,
425 ) {
426 return Ok(FitResult::ResidualCascade(fit));
427 }
428 // The quasi-uniformity guard (caveat 2) or any degenerate-design
429 // signal surfaces as a build/solve error; fall through to the dense
430 // kernel path rather than failing the fit outright.
431 }
432 }
433 // `fit_model` already returns `WorkflowError` end-to-end; propagate it
434 // directly instead of stringifying then re-wrapping.
435 fit_model(mat.request)
436}
437
438/// THE single dispatch seam for the expectile (Newey–Powell LAWS) family.
439///
440/// Returns `Ok(Some(result))` with the converged τ-expectile as an ordinary
441/// [`StandardFitResult`] when `config.family` selects the expectile family
442/// (`"expectile"` or `"expectile(τ)"`, optionally pinned by
443/// [`FitConfig::expectile_tau`]), `Ok(None)` for every other family — in which
444/// case the caller runs its normal materialize/`fit_model` path — and `Err` on a
445/// malformed expectile request or an inner-fit failure.
446///
447/// Every public entry point that resolves a family routes through this seam
448/// *before* materializing: the in-process [`fit_from_formula`], the Python FFI
449/// (`gam-pyffi`), and the `gam` CLI. Centralizing the dispatch here is what makes
450/// the estimator reachable from every interface instead of only the library
451/// call — and what prevents the class of bug where a newly-added outer estimator
452/// is wired into one entry point and silently bypassed by the others (#1777).
453/// The returned [`StandardFitResult`] carries the full design / resolved spec /
454/// fit, so each caller builds its persistence payload from it exactly as it does
455/// for any other standard fit.
456pub fn fit_expectile_if_requested(
457 formula: &str,
458 data: &Dataset,
459 config: &FitConfig,
460) -> Result<Option<StandardFitResult>, WorkflowError> {
461 match expectile_tau_for_config(config)? {
462 Some(tau) => Ok(Some(fit_expectile_laws(formula, data, config, tau)?)),
463 None => Ok(None),
464 }
465}
466
467/// Least Asymmetrically Weighted Squares (LAWS) driver for expectile GAMs.
468///
469/// The τ-expectile surface minimizes `Σ wᵢ(τ)·(yᵢ − μᵢ)²` with the residual-
470/// sign asymmetric weight `wᵢ(τ)`. Because that weight is piecewise-constant in
471/// `sign(yᵢ − μᵢ)`, the objective is the supremum of a finite family of
472/// weighted least-squares problems and its minimizer is the unique fixed point
473/// of: *solve the penalized WLS with weights frozen at the current sign
474/// pattern, then recompute the sign pattern from the new fit*. The asymmetric
475/// loss is strictly convex (weights bounded in `[min(τ,1−τ), max(τ,1−τ)] > 0`),
476/// so this monotone-descent iteration converges, and since the sign pattern
477/// takes finitely many values it stabilizes in finitely many steps (Schnabel &
478/// Eilers 2009; the same Newton/IRLS-for-expectiles `expectreg` runs).
479///
480/// Each inner solve is the FULL standard Gaussian-identity GAM: any basis,
481/// tensor, spatial smooth, by-variable, random effect, plus REML λ-selection on
482/// the current asymmetric weights. The returned fit is an ordinary
483/// [`FitResult::Standard`] whose coefficients ARE the penalized τ-expectile —
484/// every downstream consumer (predict, posterior bands, persistence) works
485/// unchanged. The reported scale is the asymmetric working variance, so
486/// expectile standard errors are the sandwich-free Gaussian-form bands of the
487/// converged weighted problem (a deliberate first-rung choice; see #1100).
488fn fit_expectile_laws(
489 formula: &str,
490 data: &Dataset,
491 config: &FitConfig,
492 tau: f64,
493) -> Result<StandardFitResult, WorkflowError> {
494 use gam_linalg::matrix::LinearOperator;
495
496 if config.frailty.as_ref().is_some_and(FrailtySpec::is_active) {
497 return Err(WorkflowError::InvalidConfig {
498 reason: "expectile regression does not support frailty; use a survival/frailty-aware family instead"
499 .to_string(),
500 });
501 }
502
503 // Inner fits are ordinary Gaussian-identity GAMs; the τ asymmetry lives
504 // entirely in the per-iteration prior weights this driver injects.
505 let gaussian_config = FitConfig {
506 family: Some("gaussian".to_string()),
507 link: Some("identity".to_string()),
508 expectile_tau: None,
509 // The inner Gaussian-identity design carries no frailty. Normalize the
510 // CLI/config-layer null value (`Some(FrailtySpec::None)`) to `None` so
511 // the expectile driver does not leak survival-only plumbing into the
512 // standard-family materializer, while the active-frailty guard above
513 // still rejects unsupported frailty requests explicitly.
514 frailty: None,
515 ..config.clone()
516 };
517
518 // Materialize once to capture the fixed training design, response, offset,
519 // and base prior weights. The design (basis, penalties, identifiability
520 // transforms) does not depend on the prior weights, so it is reused across
521 // every LAWS iteration; only the weight vector and the resulting β change.
522 let base_mat = materialize(formula, data, &gaussian_config)?;
523 let FitRequest::Standard(base_request) = base_mat.request else {
524 return Err(WorkflowError::InvalidConfig {
525 reason: "expectile regression is only defined for standard (non-survival, \
526 non-location-scale) responses"
527 .to_string(),
528 });
529 };
530 let StandardFitRequest {
531 data: design_data,
532 y,
533 weights: base_weights,
534 offset,
535 spec,
536 family: materialized_family,
537 options,
538 kappa_options,
539 wiggle,
540 coefficient_groups,
541 penalty_block_gamma_priors,
542 latent_coord,
543 _marker,
544 } = base_request;
545 // The materializer already resolved the inner family to Gaussian-identity
546 // from `gaussian_config`; assert it so a future materializer change that
547 // silently picked a different family for `"gaussian"` is caught here rather
548 // than producing a non-expectile fit.
549 if !materialized_family.is_gaussian_identity() {
550 return Err(WorkflowError::InvalidConfig {
551 reason: format!(
552 "expectile LAWS requires a Gaussian-identity inner family; materializer produced {}",
553 materialized_family.name()
554 ),
555 });
556 }
557
558 if wiggle.is_some() || latent_coord.is_some() {
559 return Err(WorkflowError::InvalidConfig {
560 reason: "expectile regression does not support flexible-link wiggle or latent \
561 coordinates"
562 .to_string(),
563 });
564 }
565
566 let n = y.len();
567 let gaussian_family = LikelihoodSpec::gaussian_identity();
568 // Cold start: τ = 0.5 (symmetric) weights ⇒ the first inner fit is the OLS
569 // mean GAM, the natural warm start for any τ.
570 let mut weights = base_weights.clone();
571 let mut last_sign: Option<Vec<bool>> = None;
572 let mut last_result: Option<StandardFitResult> = None;
573
574 // The sign pattern has 2ⁿ values but LAWS visits a monotone-descent subset;
575 // empirically a handful of iterations suffice. The cap is a safety guard:
576 // on the rare oscillation between two equal-objective sign patterns (only
577 // possible when rows sit exactly on the fitted surface) the last fit is a
578 // valid τ-expectile of the perturbation-stable problem, so returning it is
579 // correct rather than an error.
580 const MAX_LAWS_ITERS: usize = 50;
581
582 for _iter in 0..MAX_LAWS_ITERS {
583 let request = StandardFitRequest {
584 data: design_data.clone(),
585 y: y.clone(),
586 weights: weights.clone(),
587 offset: offset.clone(),
588 spec: spec.clone(),
589 family: gaussian_family.clone(),
590 options: options.clone(),
591 kappa_options: kappa_options.clone(),
592 wiggle: None,
593 coefficient_groups: coefficient_groups.clone(),
594 penalty_block_gamma_priors: penalty_block_gamma_priors.clone(),
595 latent_coord: None,
596 _marker,
597 };
598 let result = fit_standard_model(request)
599 .map_err(|reason| WorkflowError::IntegrationFailed { reason })?;
600
601 // Training-scale fitted mean μ = X·β (identity link, zero-checked
602 // offset folded by the design path). The design columns match the
603 // combined coefficient vector exactly (the same contract `predict`
604 // and the safety tests rely on).
605 let mu = result.design.design.apply(&result.fit.beta);
606 if mu.len() != n {
607 return Err(WorkflowError::IntegrationFailed {
608 reason: format!(
609 "expectile LAWS: fitted mean length {} disagrees with response length {n}",
610 mu.len()
611 ),
612 });
613 }
614 let mut mu_off = mu;
615 mu_off += &offset;
616
617 let sign: Vec<bool> = (0..n).map(|i| y[i] > mu_off[i]).collect();
618 let converged = last_sign.as_ref().is_some_and(|prev| prev == &sign);
619 weights = expectile_row_weights(y.view(), mu_off.view(), base_weights.view(), tau);
620 last_sign = Some(sign);
621 last_result = Some(result);
622 if converged {
623 break;
624 }
625 }
626
627 let result = last_result.ok_or_else(|| WorkflowError::IntegrationFailed {
628 reason: "expectile LAWS produced no fit".to_string(),
629 })?;
630 Ok(result)
631}
632/// Detection seam for the exact O(n) cubic-smoothing-spline fast path.
633///
634/// This is the EARLIEST point in the standard workflow where a materialized
635/// fit request carries everything needed to prove the model is exactly the
636/// problem the scan solves: a Gaussian likelihood with identity link over
637/// `intercept + one 1-D cubic-class penalized smooth` — i.e. the penalized
638/// least-squares problem `min Σ w_i (y_i − f(x_i))² + λ∫f″²` with an
639/// unpenalized `{1, x}` null space. The Kalman/RTS scan computes that
640/// posterior (mean, pointwise variance, exact diffuse REML for λ) in O(n) per
641/// λ-trial instead of the dense design/Gram O(n·k²) + O(k³) route.
642///
643/// Returns `Some` only when ALL of the following hold; everything else falls
644/// through to the dense path:
645/// - family is Gaussian + identity link;
646/// - no link wiggle, no latent coordinates, no coefficient groups, no penalty
647/// hyperpriors, no linear/box constraints, no Firth, no adaptive
648/// regularization, no Kronecker systems, no externally injected null-space
649/// dims;
650/// - the term collection is exactly one smooth term — no linear terms, no
651/// random effects, no by-variables / factor interactions;
652/// - that smooth is a plain 1-D B-spline whose penalty order is compatible
653/// with the exact scan and whose null space is unshrunk
654/// (`double_penalty=false`). `double_penalty` (mgcv `select = TRUE`) on a free
655/// B-spline emits a second REML coordinate — the Marra & Wood (2011) null-space
656/// shrinkage block — that the scan cannot represent (its polynomial null space
657/// is an improper diffuse prior it can never shrink); routing such a fit
658/// through the scan would silently drop that penalty and select λ from the
659/// bending penalty alone, which is exactly the EDF inflation #1266 reports.
660/// Those fits fall through to the dense two-rho path, which owns both penalties
661/// jointly. Natural cubic regression (`bs="cr"`/`"cs"`) terms also fall
662/// through: their knot-value parameterization is a finite-rank regression
663/// spline, not the scan's full smoothing-spline state-space posterior;
664/// - the offset is identically zero and every weight is finite and positive;
665/// - at least 3 distinct finite abscissae (the scan's diffuse rank plus one).
666///
667/// λ-mapping note: the scan's penalty is exactly `λ∫f″²` (state-space
668/// `q = 1/λ` at unit σ²). The dense 1-D B-spline path penalizes the same
669/// cubic class through a reduced-rank discrete-difference Gram whose
670/// normalization differs by a basis-dependent constant, so a λ selected by
671/// one parameterization does not transfer numerically to the other. The scan
672/// therefore always re-selects λ by its own exact diffuse REML criterion
673/// (the optimizer of the same restricted likelihood, expressed in the scan's
674/// parameterization); user-pinned smoothing parameters are not representable
675/// at this seam (the formula DSL exposes none for this term class), so no
676/// pinned-λ mapping arises.
677///
678/// Identifiability transforms on the smooth (centering / linear-trend
679/// removal / orthogonality-to-intercept) are accepted as eligible: they only
680/// re-coordinate the unpenalized null space against the implicit intercept
681/// and do not change the fitted posterior of `E[y|x]`, which is what the
682/// scan returns directly.
683pub fn spline_scan_fast_path(request: &StandardFitRequest<'_>) -> Option<SplineScanInputs> {
684 if !request.family.is_gaussian_identity() {
685 return None;
686 }
687 if request.wiggle.is_some()
688 || request.latent_coord.is_some()
689 || !request.coefficient_groups.is_empty()
690 || !request.penalty_block_gamma_priors.is_empty()
691 {
692 return None;
693 }
694 let options = &request.options;
695 if options.latent_cloglog.is_some()
696 || options.mixture_link.is_some()
697 || options.sas_link.is_some()
698 || options.linear_constraints.is_some()
699 || options.adaptive_regularization.is_some()
700 || options.kronecker_penalty_system.is_some()
701 || options.kronecker_factored.is_some()
702 || options.firth_bias_reduction
703 || !options.nullspace_dims.is_empty()
704 {
705 return None;
706 }
707 let spec = &request.spec;
708 if !spec.linear_terms.is_empty()
709 || !spec.random_effect_terms.is_empty()
710 || spec.smooth_terms.len() != 1
711 {
712 return None;
713 }
714 let term = &spec.smooth_terms[0];
715 if !matches!(term.shape, gam_terms::smooth::ShapeConstraint::None)
716 || term.joint_null_rotation.is_some()
717 {
718 return None;
719 }
720 let gam_terms::smooth::SmoothBasisSpec::BSpline1D {
721 feature_col,
722 spec: bspec,
723 } = &term.basis
724 else {
725 return None;
726 };
727 // Smoothing-spline order m = penalty_order ∈ {1, 2, 3}. The exact scan
728 // integrates the order-m integrated-Wiener prior whose natural spline has
729 // degree 2m−1 (m=1 → linear, m=2 → cubic, m=3 → quintic), so require that
730 // degree to match user intent. The de Jong exact diffuse leading-block
731 // smoother (#1044) handles the m−1 partially-diffuse leading nodes for all
732 // m ≤ MAX_ORDER; m > MAX_ORDER falls through to the dense path.
733 let order = bspec.penalty_order;
734 // Double-penalty (mgcv `select = TRUE`) is NOT representable by the scan and
735 // must fall through to the dense two-rho path (#1266). On a free B-spline the
736 // double penalty emits a *second* REML coordinate — the Marra & Wood (2011)
737 // null-space shrinkage block `Z Zᵀ` (see `bspline_penalty_candidates`) —
738 // whose entire purpose is to let REML shrink the unpenalized `{1, x, …}`
739 // polynomial null space toward `EDF → 0` for an unsupported term. The scan,
740 // by construction, carries that null space as an *improper diffuse* prior it
741 // can never shrink (its EDF floor is the null-space dimension `order`), so
742 // routing a `double_penalty` fit through it silently DROPS the second penalty
743 // and selects λ from the single bending penalty alone. The scan's own exact
744 // diffuse REML then genuinely prefers a mildly wiggly fit at finite λ for
745 // some noise realizations (an interior REML optimum, EDF ≈ 3–4), which is the
746 // EDF inflation #1266 reports. The dense path owns both penalties jointly and
747 // its outer REML, seeded into the over-smoothing basin, drives the null space
748 // out (EDF → null-space dim) when the data are truly polynomial. Excluding
749 // `double_penalty` here keeps such a fit on the dense path; single-penalty
750 // and boundary-conditioned single-penalty B-splines keep the exact O(n) scan.
751 if !(1..=3).contains(&order)
752 || bspec.degree != 2 * order - 1
753 || bspec.double_penalty
754 || !bspec.boundary_conditions.is_free()
755 || !matches!(bspec.boundary, gam_terms::basis::OneDimensionalBoundary::Open)
756 || matches!(
757 bspec.knotspec,
758 gam_terms::basis::BSplineKnotSpec::PeriodicUniform { .. }
759 | gam_terms::basis::BSplineKnotSpec::NaturalCubicRegression { .. }
760 )
761 // mgcv `bs="cr"`/`"cs"` materialise a `NaturalCubicRegression` value-knot
762 // spec: a Lancaster–Salkauskas cubic-regression basis whose columns
763 // index `f(x*_i)` at `k` quantile knots — a genuinely DIFFERENT finite
764 // basis (and hence a different penalized posterior) from the free
765 // integrated-Wiener natural spline the exact scan solves on the raw data
766 // points. The scan builds its own knots from `x` and ignores this spec,
767 // so routing a cr fit through it would silently solve the wrong model and
768 // (per #1844) return a non-`Standard` `SplineScan` result the predict-time
769 // design replay cannot reconstruct. Keep cr/cs on the dense path.
770 || matches!(
771 bspec.knotspec,
772 gam_terms::basis::BSplineKnotSpec::NaturalCubicRegression { .. }
773 )
774 {
775 return None;
776 }
777 if request.offset.iter().any(|&v| v != 0.0) {
778 return None;
779 }
780 if request.weights.iter().any(|&v| !(v.is_finite() && v > 0.0)) {
781 return None;
782 }
783 if *feature_col >= request.data.ncols() || request.y.len() != request.data.nrows() {
784 return None;
785 }
786 let x: Vec<f64> = request.data.column(*feature_col).iter().copied().collect();
787 let y: Vec<f64> = request.y.iter().copied().collect();
788 let w: Vec<f64> = request.weights.iter().copied().collect();
789 if x.iter().any(|v| !v.is_finite()) || y.iter().any(|v| !v.is_finite()) {
790 return None;
791 }
792 // The diffuse polynomial null space consumes `order` innovations; the scan
793 // needs at least one proper innovation beyond them to profile σ².
794 let mut sorted = x.clone();
795 sorted.sort_by(f64::total_cmp);
796 sorted.dedup();
797 if sorted.len() < order + 1 {
798 return None;
799 }
800 Some(SplineScanInputs { x, y, w, order })
801}
802
803/// Formula-level entry for the exact O(n) cubic-smoothing-spline fast path.
804///
805/// Materializes the formula exactly like [`fit_from_formula`], then runs the
806/// [`spline_scan_fast_path`] detection on the resulting standard request.
807/// When detection fires the fit is routed through
808/// [`gam_solve::spline_scan::fit_spline_scan`] — the exact diffuse
809/// REML Kalman/RTS scan — and the full in-memory posterior
810/// ([`gam_solve::spline_scan::SplineScanFit`]: knots, smoothed
811/// states, pointwise variances, lag-one gains, σ², log λ, exact EDF, and an
812/// exact `predict`) is returned. `Ok(None)` means the model is not the
813/// scan-eligible shape and the caller should use the dense
814/// [`fit_from_formula`] path; this keeps every persistence-bearing consumer
815/// (model save, CLI, FFI) transparently on the dense fit, whose saved payload
816/// the scan does not yet have a schema for.
817pub fn fit_spline_scan_from_formula(
818 formula: &str,
819 data: &Dataset,
820 config: &FitConfig,
821) -> Result<Option<gam_solve::spline_scan::SplineScanFit>, WorkflowError> {
822 let mat = materialize(formula, data, config)?;
823 let FitRequest::Standard(request) = mat.request else {
824 return Ok(None);
825 };
826 let Some(inputs) = spline_scan_fast_path(&request) else {
827 return Ok(None);
828 };
829 gam_solve::spline_scan::fit_spline_scan(&inputs.x, &inputs.y, &inputs.w, inputs.order)
830 .map(Some)
831 .map_err(|reason| WorkflowError::IntegrationFailed { reason })
832}
833
834/// #1464 diagnostic entry point: evaluate the EXACT production fixed-κ
835/// profiled-REML criterion (`fixed_kappa_profiled_reml_score`, the same one the
836/// joint-fit κ-sign scan uses) at a list of pinned κ values for the first
837/// constant-curvature term of `formula`, materialised from `data`/`config`
838/// exactly like [`fit_from_formula`]. Returns `(κ, V_p(κ))` pairs.
839///
840/// This settles solver-vs-criterion for the railing bug: if `V_p(+κ) < V_p(−κ)`
841/// for a genuinely HYPERBOLIC dataset, the criterion itself prefers the collapsed
842/// +κ corner — the bug is in the constant-curvature REML/Occam term, not the
843/// optimiser. If `V_p(−κ) < V_p(+κ)` yet the full fit still returns +κ, the bug
844/// is in the solver/readback. The profiled fit pins κ and profiles only ρ
845/// (κ-optimisation disabled), so each returned score is the negative-log-evidence
846/// the outer loop minimises.
847pub fn constant_curvature_profiled_reml_scores(
848 formula: &str,
849 data: &Dataset,
850 config: &FitConfig,
851 kappas: &[f64],
852) -> Result<Vec<(f64, f64)>, WorkflowError> {
853 let mat = materialize(formula, data, config)?;
854 let FitRequest::Standard(request) = mat.request else {
855 return Err(WorkflowError::IntegrationFailed {
856 reason: "constant_curvature_profiled_reml_scores: formula did not materialise to a \
857 standard fit request"
858 .to_string(),
859 });
860 };
861 let term_idx = *crate::fit_orchestration::drivers::constant_curvature_term_indices(&request.spec)
862 .first()
863 .ok_or_else(|| WorkflowError::IntegrationFailed {
864 reason: "constant_curvature_profiled_reml_scores: formula has no constant-curvature \
865 curv() term"
866 .to_string(),
867 })?;
868 let mut out = Vec::with_capacity(kappas.len());
869 for &kappa in kappas {
870 let score = crate::fit_orchestration::drivers::fixed_kappa_profiled_reml_score(
871 request.data.view(),
872 request.y.view(),
873 request.weights.view(),
874 request.offset.view(),
875 &request.spec,
876 term_idx,
877 kappa,
878 request.family.clone(),
879 &request.options,
880 )
881 .map_err(|e| WorkflowError::IntegrationFailed {
882 reason: format!(
883 "constant_curvature_profiled_reml_scores: fixed-κ fit at κ={kappa} failed: {e}"
884 ),
885 })?;
886 out.push((kappa, score));
887 }
888 Ok(out)
889}
890
891/// Derived dense-kernel cliff: the cascade auto-route fires only once the dense
892/// radial basis the smooth would otherwise use has SATURATED at its center cap
893/// (`default_num_centers == K_MAX`), so the dense `O(n·K² + K³)` kernel solve
894/// can no longer grow resolution with `n` and the streaming cascade's
895/// `O(n·polylog)` is the only path that keeps improving. This is the structural
896/// "past the dense-kernel cliff" condition the issue names — derived from the
897/// dense sizing rule, NOT a magic n constant or a user flag.
898fn past_dense_kernel_cliff(n: usize, d: usize) -> bool {
899 // `default_num_centers` clamps to K_MAX = 2000; equality means the dense
900 // basis is pinned at the cap and cannot densify further with n.
901 const DENSE_CENTER_CAP: usize = 2000;
902 gam_terms::basis::default_num_centers(n, d) >= DENSE_CENTER_CAP
903}
904
905/// Map a Duchon/Matérn smoothness order onto the cascade's Sobolev order,
906/// clamped into the Wendland-(3,1) native window `(d/2, (d+3)/2]` (issue
907/// caveat 1: the multilevel frame can only represent up to `H^{(d+3)/2}`).
908fn cascade_sobolev_order(requested: f64, d: usize) -> f64 {
909 let lo = d as f64 / 2.0;
910 let hi = (d as f64 + 3.0) / 2.0;
911 // Nudge strictly inside the open lower bound when the request lands on it.
912 let eps = 1e-6 * (hi - lo);
913 requested.clamp(lo + eps, hi)
914}
915
916/// Detection seam for the O(n log n) multiresolution residual-cascade fast path
917/// (issue #1032).
918///
919/// This mirrors [`spline_scan_fast_path`] in shape but carries one CRITICAL
920/// difference dictated by the issue: the cascade is **not** the same posterior
921/// as the Duchon/Matérn term it stands in for (a different finite basis — the
922/// multilevel Wendland frame, not the reduced-rank radial kernel). So unlike
923/// the 1-D scan, which silently swaps an identical posterior, this path must
924/// only fire as an explicit alternative estimator on the structural signature
925/// the issue names, never as a transparent replacement. It returns `Some` only
926/// when ALL of the following hold:
927/// - family is Gaussian + identity link (the scattered low-d smooth the
928/// cascade solves);
929/// - none of the exotic-link / constraint / Firth / Kronecker / coefficient-
930/// group / hyperprior machinery is engaged;
931/// - the model is exactly one smooth term — no linear terms, no random
932/// effects, no by-variables;
933/// - that smooth is a scattered radial spatial smooth (`Duchon` or `Matern`)
934/// over `d ∈ {2, 3}` coordinates with no shape constraint;
935/// - the offset is identically zero and every weight is finite and positive;
936/// - `n` is past the derived dense-kernel cliff
937/// ([`past_dense_kernel_cliff`]) — below it the dense radial path is both
938/// exact-posterior and cheap, so there is no reason to change estimators.
939///
940/// The returned [`ResidualCascadeInputs`] carry a unit per-axis metric (the
941/// spec's isotropic radial distance); the quasi-uniformity guard inside
942/// [`gam_solve::residual_cascade::fit_residual_cascade`] (issue caveat 2)
943/// is the no-regression gate that refuses the iterative solve — and forces the
944/// caller back to the dense path — when a near-degenerate metric would break
945/// the BPX iteration bound.
946pub fn residual_cascade_fast_path(
947 request: &StandardFitRequest<'_>,
948) -> Option<ResidualCascadeInputs> {
949 if !request.family.is_gaussian_identity() {
950 return None;
951 }
952 if request.wiggle.is_some()
953 || request.latent_coord.is_some()
954 || !request.coefficient_groups.is_empty()
955 || !request.penalty_block_gamma_priors.is_empty()
956 {
957 return None;
958 }
959 let options = &request.options;
960 if options.latent_cloglog.is_some()
961 || options.mixture_link.is_some()
962 || options.sas_link.is_some()
963 || options.linear_constraints.is_some()
964 || options.adaptive_regularization.is_some()
965 || options.kronecker_penalty_system.is_some()
966 || options.kronecker_factored.is_some()
967 || options.firth_bias_reduction
968 || !options.nullspace_dims.is_empty()
969 {
970 return None;
971 }
972 let spec = &request.spec;
973 if !spec.linear_terms.is_empty()
974 || !spec.random_effect_terms.is_empty()
975 || spec.smooth_terms.len() != 1
976 {
977 return None;
978 }
979 let term = &spec.smooth_terms[0];
980 if !matches!(term.shape, gam_terms::smooth::ShapeConstraint::None)
981 || term.joint_null_rotation.is_some()
982 {
983 return None;
984 }
985 // Only scattered radial spatial smooths (Duchon / Matérn) over 2–3 axes.
986 // The Duchon spectral power `p + s` and the Matérn order set the requested
987 // Sobolev smoothness; both clamp into the Wendland native window.
988 let (feature_cols, requested_s) = match &term.basis {
989 gam_terms::smooth::SmoothBasisSpec::Duchon {
990 feature_cols, spec, ..
991 } => {
992 // Pure-Duchon native order is `p + s` (kernel exponent 2(p+s)−d);
993 // the multilevel frame targets the same continuum smoothness. `p`
994 // is the polynomial nullspace degree, `s` the spectral power.
995 let p = match spec.nullspace_order {
996 gam_terms::basis::DuchonNullspaceOrder::Zero => 0.0,
997 gam_terms::basis::DuchonNullspaceOrder::Linear => 1.0,
998 gam_terms::basis::DuchonNullspaceOrder::Degree(k) => k as f64,
999 };
1000 (feature_cols, spec.power + p)
1001 }
1002 gam_terms::smooth::SmoothBasisSpec::Matern {
1003 feature_cols, spec, ..
1004 } => {
1005 // Matérn smoothness ν sets native Sobolev order ν + d/2; the cascade
1006 // frame represents up to (d+3)/2, so the clamp below applies the
1007 // ceiling. (d is known just below from feature_cols.)
1008 let nu = spec.nu.half_integer_value();
1009 (feature_cols, nu + feature_cols.len() as f64 / 2.0)
1010 }
1011 _ => return None,
1012 };
1013 let d = feature_cols.len();
1014 if !(2..=3).contains(&d) {
1015 return None;
1016 }
1017 if request.offset.iter().any(|&v| v != 0.0) {
1018 return None;
1019 }
1020 if request.weights.iter().any(|&v| !(v.is_finite() && v > 0.0)) {
1021 return None;
1022 }
1023 let n = request.y.len();
1024 if n != request.data.nrows() || feature_cols.iter().any(|&c| c >= request.data.ncols()) {
1025 return None;
1026 }
1027 if !past_dense_kernel_cliff(n, d) {
1028 return None;
1029 }
1030 let coords: Vec<Vec<f64>> = feature_cols
1031 .iter()
1032 .map(|&c| request.data.column(c).iter().copied().collect())
1033 .collect();
1034 let y: Vec<f64> = request.y.iter().copied().collect();
1035 let w: Vec<f64> = request.weights.iter().copied().collect();
1036 if coords
1037 .iter()
1038 .any(|axis| axis.iter().any(|v| !v.is_finite()))
1039 || y.iter().any(|v| !v.is_finite())
1040 {
1041 return None;
1042 }
1043 let metric = vec![1.0_f64; d];
1044 let sobolev_s = cascade_sobolev_order(requested_s, d);
1045 Some(ResidualCascadeInputs {
1046 coords,
1047 y,
1048 w,
1049 metric,
1050 sobolev_s,
1051 })
1052}
1053
1054/// Formula-level library entry for the O(n log n) residual-cascade fast path
1055/// (issue #1032).
1056///
1057/// Materializes the formula exactly like [`fit_from_formula`], runs the
1058/// [`residual_cascade_fast_path`] detection, and — when it fires AND the
1059/// quasi-uniformity guard inside the cascade certifies the metric — returns the
1060/// certified [`ResidualCascadeFit`](gam_solve::residual_cascade::ResidualCascadeFit).
1061/// `Ok(None)` means EITHER the model is not the cascade-eligible shape OR the
1062/// quasi-uniformity guard rejected the metric; in both cases the caller falls
1063/// back to the dense [`fit_from_formula`] path (the cascade is a different
1064/// posterior, so the fallback is a genuine estimator choice, never a silent
1065/// swap). This keeps every persistence-bearing consumer on the dense fit until
1066/// the cascade payload schema lands.
1067pub fn fit_residual_cascade_from_formula(
1068 formula: &str,
1069 data: &Dataset,
1070 config: &FitConfig,
1071) -> Result<Option<gam_solve::residual_cascade::ResidualCascadeFit>, WorkflowError> {
1072 let mat = materialize(formula, data, config)?;
1073 let FitRequest::Standard(request) = mat.request else {
1074 return Ok(None);
1075 };
1076 let Some(inputs) = residual_cascade_fast_path(&request) else {
1077 return Ok(None);
1078 };
1079 let coord_refs: Vec<&[f64]> = inputs.coords.iter().map(Vec::as_slice).collect();
1080 match gam_solve::residual_cascade::fit_residual_cascade(
1081 &coord_refs,
1082 &inputs.y,
1083 &inputs.w,
1084 &inputs.metric,
1085 inputs.sobolev_s,
1086 ) {
1087 Ok(fit) => Ok(Some(fit)),
1088 // The quasi-uniformity guard (caveat 2) and any degenerate-design
1089 // signal both surface as a build/solve error; treat them as "not
1090 // cascade-eligible" so the caller falls back to the dense kernel path
1091 // rather than failing the fit outright.
1092 Err(_) => Ok(None),
1093 }
1094}
1095
1096/// Parse a formula, resolve it against a dataset, and produce a ready-to-fit `FitRequest`.
1097pub fn materialize<'a>(
1098 formula: &str,
1099 data: &'a Dataset,
1100 config: &FitConfig,
1101) -> Result<MaterializedModel<'a>, WorkflowError> {
1102 gam_gpu::configure_global_policy(config.gpu_policy);
1103 let parsed = parse_formula(formula)?;
1104 let col_map = data.column_map();
1105
1106 if let Some((left_col, right_col, event_col)) = parse_surv_interval_response(&parsed.response)?
1107 {
1108 if config.transformation_normal {
1109 return Err(WorkflowError::InvalidConfig {
1110 reason:
1111 "transformation_normal cannot be combined with a SurvInterval(...) response"
1112 .to_string(),
1113 });
1114 }
1115 // Interval censoring `T ∈ (L, R]` is only defined for the latent
1116 // hazard-window survival likelihood, whose kernel carries the
1117 // `log[S(L) − S(R)]` interval contribution. Route the left boundary `L`
1118 // through the standard exit channel and the right boundary `R` through
1119 // the dedicated interval-right channel; `event_col` distinguishes
1120 // bracketed (interval) rows from right-censored rows beyond the last
1121 // inspection (which carry an infinite/sentinel `R`).
1122 materialize_survival(
1123 &parsed,
1124 data,
1125 &col_map,
1126 config,
1127 None,
1128 &left_col,
1129 &event_col,
1130 Some(&right_col),
1131 )
1132 } else if let Some((entry_col, exit_col, event_col)) = parse_surv_response(&parsed.response)? {
1133 if config.transformation_normal {
1134 return Err(WorkflowError::InvalidConfig {
1135 reason: "transformation_normal cannot be combined with a Surv(...) response"
1136 .to_string(),
1137 });
1138 }
1139 // `materialize_*` now return `WorkflowError` directly so the typed
1140 // `ColumnNotFound` payload (and any future variant-typed leaf
1141 // errors) survive the dispatcher hop instead of being flattened
1142 // into `IntegrationFailed { reason: String }`.
1143 materialize_survival(
1144 &parsed,
1145 data,
1146 &col_map,
1147 config,
1148 entry_col.as_deref(),
1149 &exit_col,
1150 &event_col,
1151 None,
1152 )
1153 } else {
1154 // Non-survival response: `timewiggle(...)` and `survmodel(...)` are
1155 // structurally meaningless (there is no baseline hazard / time axis to
1156 // wiggle and no survival likelihood to configure). They are parsed into
1157 // `ParsedFormula` but consumed *only* by `materialize_survival`; without
1158 // this guard every non-survival materializer below would silently drop
1159 // them, fitting an ordinary GAM while the user believes they requested a
1160 // time-varying / survival model (#371). Reject here — the single
1161 // chokepoint for all non-survival paths — mirroring the symmetric
1162 // auxiliary-formula rejection in `validate_auxiliary_formula_controls`.
1163 reject_survival_only_terms_for_nonsurvival(&parsed)?;
1164 // Symmetrically, the `config.survival_likelihood` *knob* selects a
1165 // survival likelihood mode read only by `materialize_survival`. On this
1166 // non-survival branch a non-default value (e.g. "weibull") would be
1167 // discarded and the fit would silently degrade to an ordinary GAM
1168 // (#1767). Reject it at the same chokepoint.
1169 reject_survival_likelihood_for_nonsurvival(config)?;
1170 if config.transformation_normal {
1171 // Issue #789A: a Bernoulli marginal-slope request with
1172 // `transformation_normal=true` used to dispatch as a CTN fit while
1173 // retaining marginal-slope controls, leaving the transformation path
1174 // in a non-advancing loop. CTN score calibration now uses the
1175 // explicit `ctn_stage1` recipe instead, so the legacy boolean is a
1176 // hard configuration error for marginal-slope requests.
1177 reject_marginal_slope_controls_for_transformation_normal(config)?;
1178 if config.noise_formula.is_some() {
1179 return Err(WorkflowError::InvalidConfig {
1180 reason: "transformation_normal cannot be combined with noise_formula"
1181 .to_string(),
1182 });
1183 }
1184 materialize_transformation_normal(&parsed, data, &col_map, config)
1185 } else if config.logslope_formula.is_some() || config.z_column.is_some() {
1186 materialize_bernoulli_marginal_slope(&parsed, data, &col_map, config)
1187 } else if config.noise_formula.is_some() {
1188 materialize_location_scale(&parsed, data, &col_map, config)
1189 } else {
1190 materialize_standard(&parsed, data, &col_map, config)
1191 }
1192 }
1193}
1194
1195#[cfg(test)]
1196mod sz_factor_smooth_recovery_tests {
1197 // `super::*` brings in `Dataset` (= gam_data::EncodedDataset), `FitConfig`,
1198 // `FitResult`, `StandardFitResult`, and `fit_from_formula`.
1199 use super::*;
1200
1201 const NOISE_SD: f64 = 0.20;
1202 const N: usize = 4000;
1203 const N_GROUPS: usize = 4;
1204
1205 /// A simple deterministic LCG so the dataset is reproducible without pulling
1206 /// an RNG dependency into the test.
1207 struct Lcg(u64);
1208 impl Lcg {
1209 fn next_u64(&mut self) -> u64 {
1210 // Numerical Recipes LCG constants.
1211 self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
1212 self.0
1213 }
1214 /// Uniform in [0, 1).
1215 fn unif(&mut self) -> f64 {
1216 (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
1217 }
1218 /// Standard normal via Box–Muller (one of the pair).
1219 fn normal(&mut self) -> f64 {
1220 let u1 = (self.unif()).max(1e-12);
1221 let u2 = self.unif();
1222 (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
1223 }
1224 }
1225
1226 /// Data drawn from EXACTLY the `sz` model class: a shared smooth `f0(x)` plus
1227 /// zero-sum per-group deviations `d_g(x)` (phase-shifted sinusoids whose
1228 /// cross-group mean is removed at every `x`), plus observation noise. This
1229 /// mirrors the (blocked) Python bug-hunt test `tests/bug_hunt_sz_factor_
1230 /// smooth_underfits_own_model_class_test.py`.
1231 ///
1232 /// Written to a CSV and loaded through the real `load_dataset_projected`
1233 /// inferer so the grouping column `g` (string levels) is encoded as a genuine
1234 /// categorical exactly as production does — hand-built `EncodedDataset`s do
1235 /// not carry the categorical level map the factor-smooth level resolver needs.
1236 fn sz_class_dataset() -> (Dataset, tempfile::TempDir) {
1237 let mut rng = Lcg(0x5326_2026_0628_1605);
1238 let phases: Vec<f64> = (0..N_GROUPS)
1239 .map(|k| 1.2 * k as f64 / (N_GROUPS as f64 - 1.0))
1240 .collect();
1241 let deviations = |xi: f64| -> Vec<f64> {
1242 let vals: Vec<f64> = phases
1243 .iter()
1244 .map(|p| 0.6 * (std::f64::consts::TAU * xi + std::f64::consts::TAU * p).sin())
1245 .collect();
1246 let mean = vals.iter().sum::<f64>() / vals.len() as f64;
1247 vals.iter().map(|v| v - mean).collect()
1248 };
1249
1250 let mut csv = String::from("y,x,g\n");
1251 for _ in 0..N {
1252 let x = rng.unif();
1253 // Use the HIGH bits (via `unif`) for the group draw — an LCG's low
1254 // bits have a tiny period and would collapse `% N_GROUPS` to a near
1255 // constant.
1256 let g = ((rng.unif() * N_GROUPS as f64) as usize).min(N_GROUPS - 1);
1257 let f0 = (std::f64::consts::TAU * x).sin();
1258 let mu = f0 + deviations(x)[g];
1259 let y = mu + NOISE_SD * rng.normal();
1260 csv.push_str(&format!("{y},{x},g{g}\n"));
1261 }
1262 let td = tempfile::tempdir().expect("tempdir");
1263 let path = td.path().join("sz_class.csv");
1264 std::fs::write(&path, csv).expect("write sz-class csv");
1265 // Force `g` into a categorical role exactly as the formula intends so the
1266 // factor-smooth level resolver sees all `N_GROUPS` distinct levels.
1267 let mut roles = std::collections::HashSet::new();
1268 roles.insert("g");
1269 let data = gam_data::load_dataset_projected_with_categorical_roles(
1270 &path,
1271 &["y".to_string(), "x".to_string(), "g".to_string()],
1272 &roles,
1273 )
1274 .expect("load sz-class dataset");
1275 (data, td)
1276 }
1277
1278 fn gaussian_config() -> FitConfig {
1279 FitConfig { family: Some("gaussian".to_string()), ..FitConfig::default() }
1280 }
1281
1282 /// In-sample residual sd of a fitted standard GAM: `sd(y − Xβ̂)`.
1283 fn residual_sd(fit: &StandardFitResult, data: &Dataset) -> f64 {
1284 let beta = &fit.fit.beta;
1285 let design = &fit.design.design;
1286 let n = design.nrows();
1287 assert_eq!(design.ncols(), beta.len(), "design/beta width mismatch");
1288 let mut fitted = vec![0.0f64; n];
1289 // `try_row_chunk` materializes contiguous row blocks of whatever design
1290 // storage the fit used (dense or block-lazy) — robust to the storage kind.
1291 const CHUNK: usize = 512;
1292 let mut start = 0usize;
1293 while start < n {
1294 let end = (start + CHUNK).min(n);
1295 let block = design
1296 .try_row_chunk(start..end)
1297 .expect("materialize design row chunk");
1298 for (r, row) in block.rows().into_iter().enumerate() {
1299 let mut acc = 0.0;
1300 for (c, &xv) in row.iter().enumerate() {
1301 acc += xv * beta[c];
1302 }
1303 fitted[start + r] = acc;
1304 }
1305 start = end;
1306 }
1307 let y = data.values.column(0);
1308 let resid: Vec<f64> = y.iter().zip(fitted.iter()).map(|(&yi, &fi)| yi - fi).collect();
1309 let mean = resid.iter().sum::<f64>() / resid.len() as f64;
1310 let var = resid.iter().map(|r| (r - mean).powi(2)).sum::<f64>() / resid.len() as f64;
1311 var.sqrt()
1312 }
1313
1314 fn fit_standard(formula: &str, data: &Dataset) -> StandardFitResult {
1315 match fit_from_formula(formula, data, &gaussian_config())
1316 .unwrap_or_else(|e| panic!("fit `{formula}` failed: {e:?}"))
1317 {
1318 FitResult::Standard(r) => r,
1319 other => panic!("expected Standard fit for `{formula}`, got a different variant: {}",
1320 std::any::type_name_of_val(&other)),
1321 }
1322 }
1323
1324 /// #1605 (gold standard, end-to-end REML fit): the sum-to-zero factor smooth
1325 /// `s(x) + s(g, x, bs="sz")` must RECOVER data drawn from its own model class
1326 /// to the observation-noise floor, exactly as the strictly-more-general
1327 /// `s(x, g, bs="fs")` superset provably does.
1328 ///
1329 /// The recovery gap (`sz` resid ≈ 0.43 ≈ 2.1× the 0.20 floor while `fs`
1330 /// reaches the floor) was closed by THREE mgcv-faithful corrections, each
1331 /// necessary, that this end-to-end fit jointly exercises:
1332 /// 1. marginal basis (baef17e): cr → curvature-capable B-spline, so a
1333 /// deviation with non-zero boundary curvature is representable;
1334 /// 2. ownership/overlap residualization (b49bb5c): the `sz` deviation is
1335 /// sum-to-zero ACROSS the grouping factor, hence orthogonal to a
1336 /// factor-independent owner like the shared `s(x)`. Residualizing it
1337 /// against `s(x)`'s realized span (the #978 chart) collapsed every
1338 /// group's curve to a flat per-group contrast; skipping that ownership
1339 /// (same family as the #1276 factor-`by` level gate) restores the curve
1340 /// shape and stops REML railing the shared `s(x)` wiggliness λ;
1341 /// 3. null-space ridge (this change): the `sz` deviation blocks now carry
1342 /// the per-null-dimension ridge structure of `fs`, mapped into the
1343 /// zero-sum contrast space, so the {const, linear} null space is
1344 /// shrinkable per dimension (the #700/#712/#713 partial-pooling form)
1345 /// rather than left free — without breaking the zero-sum constraint.
1346 ///
1347 /// This is the gold-standard verification: it drives the real
1348 /// `fit_from_formula` REML λ-selection on data drawn from exactly the `sz`
1349 /// model class and asserts `sz` reaches the floor (and a `fs` control does
1350 /// too). It failed before the fixes and passes after.
1351 #[test]
1352 fn sz_factor_smooth_recovers_its_own_model_class_end_to_end() {
1353 let (data, _td) = sz_class_dataset();
1354
1355 // Control: bs="fs", a strict superset of the sz span, must reach the
1356 // noise floor — proves the data is well-posed and pins the floor.
1357 let fs_fit = fit_standard("y ~ s(x, g, bs='fs')", &data);
1358 let fs_resid = residual_sd(&fs_fit, &data);
1359 assert!(
1360 fs_resid < 1.2 * NOISE_SD,
1361 "control bs='fs' did not reach the noise floor: resid_sd={fs_resid:.4} \
1362 vs noise_sd={NOISE_SD} (data/floor sanity check)",
1363 );
1364
1365 // The documented sz idiom on data drawn from the sz model class.
1366 let sz_fit = fit_standard("y ~ s(x) + s(g, x, bs='sz')", &data);
1367 let sz_resid = residual_sd(&sz_fit, &data);
1368
1369 // A smoother whose span contains the truth, fit at large n, must explain
1370 // the systematic structure and leave ~only observation noise.
1371 assert!(
1372 sz_resid < 1.4 * NOISE_SD,
1373 "bs='sz' under-fits its own model class: resid_sd={sz_resid:.4} \
1374 ({:.2}x the noise floor {NOISE_SD}); the bs='fs' superset reached \
1375 {fs_resid:.4}. The sz fit leaves systematic signal in the residual.",
1376 sz_resid / NOISE_SD,
1377 );
1378
1379 // Comparative guard: sz must not be dramatically worse than the fs
1380 // superset that recovers the same data.
1381 assert!(
1382 sz_resid < 1.5 * fs_resid,
1383 "bs='sz' residual {sz_resid:.4} is {:.2}x the bs='fs' residual \
1384 {fs_resid:.4} on identical sz-class data",
1385 sz_resid / fs_resid,
1386 );
1387 }
1388}