Skip to main content

gam_models/survival/
construction.rs

1//! Survival model construction helpers.
2//!
3//! Types and functions for building survival model components:
4//! - Baseline hazard targets (Weibull, Gompertz, Gompertz-Makeham)
5//! - Time basis construction (I-spline on log-time)
6//! - Baseline offset computation
7//! - Time wiggle construction
8//!
9//! These are the building blocks a library consumer needs to construct
10//! a `FitRequest::SurvivalLocationScale` without going through the CLI.
11
12use crate::probability::{normal_pdf, standard_normal_quantile};
13use crate::survival::location_scale::{
14    DEFAULT_SURVIVAL_LOCATION_SCALE_DERIVATIVE_GUARD, ResidualDistribution,
15    SurvivalCovariateTermBlockTemplate, SurvivalCovariateTimeBasis,
16};
17use crate::survival::lognormal_kernel::HazardLoading;
18use crate::survival::marginal_slope::DEFAULT_SURVIVAL_MARGINAL_SLOPE_DERIVATIVE_GUARD;
19use crate::wiggle::{monotone_wiggle_basis_with_derivative_order, split_wiggle_penalty_orders};
20use gam_linalg::matrix::{
21    DenseDesignMatrix, DesignMatrix, SparseDesignMatrix, symmetrize_in_place,
22};
23use gam_problem::outer_subsample::RowSet;
24use gam_problem::{InverseLink, StandardLink};
25use gam_terms::basis::{
26    BSplineBasisSpec, BSplineBoundaryConditions, BSplineIdentifiability, BSplineKnotSpec,
27    BasisMetadata, BasisOptions, Dense, ISplineBoundary, KnotSource, OneDimensionalBoundary,
28    build_bspline_basis_1d, create_basis, evaluate_bspline_derivative_scalar,
29    ispline_modelling_interval, ispline_value, ispline_value_and_first_derivative,
30};
31use gam_terms::inference::formula_dsl::LinkWiggleFormulaSpec;
32use ndarray::{Array1, Array2, Array3, array, s};
33use rayon::prelude::*;
34
35// ---------------------------------------------------------------------------
36// Typed error
37// ---------------------------------------------------------------------------
38
39/// Structured failure surface for survival-model construction helpers
40/// (`parse_*`, baseline-config builders, time-basis construction). Every
41/// variant carries a free-form `reason: String` payload; `Display` emits
42/// that payload verbatim, so converting to `String` via the `From` impl
43/// produces text byte-equivalent to the pre-refactor `Err(format!(...))`
44/// call sites that were the only producers in this module.
45///
46/// The public CLI-input parsers (`parse_survival_distribution`,
47/// `parse_survival_likelihood_mode`, `parse_survival_baseline_config`)
48/// keep their `Result<_, String>` signatures — string is the natural
49/// failure type for free-form user input — and route through this enum
50/// internally via `From<SurvivalConstructionError> for String`.
51#[derive(Clone, Debug)]
52pub enum SurvivalConstructionError {
53    /// User-supplied configuration is malformed or out of range (knot
54    /// counts, anchor offsets, derivative guards, ranks).
55    InvalidConfig { reason: String },
56    /// A required column or block of metadata is absent (e.g. saved
57    /// survival ispline keep_cols, baseline target on a saved fit).
58    MissingColumn { reason: String },
59    /// Per-row / per-column shape disagreement (entry/exit lengths,
60    /// penalty rank vs basis width, basis vs coefficient counts).
61    IncompatibleDimensions { reason: String },
62    /// Numeric / domain rejection: non-finite ratios, non-positive
63    /// survival times, monotonicity violations, ispline-derivative
64    /// underflow.
65    DataValidationFailed { reason: String },
66    /// Underlying basis / penalty builder rejected the construction
67    /// request (invalid spline order, ispline keep_cols out of range,
68    /// internal empty ispline time basis).
69    BasisConstructionFailed { reason: String },
70    /// User-named distribution / likelihood-mode / baseline target /
71    /// time-basis kind is not one we recognise.
72    UnsupportedDistribution { reason: String },
73}
74
75impl_reason_error_boilerplate! {
76    SurvivalConstructionError {
77        InvalidConfig,
78        MissingColumn,
79        IncompatibleDimensions,
80        DataValidationFailed,
81        BasisConstructionFailed,
82        UnsupportedDistribution,
83    }
84}
85
86// ---------------------------------------------------------------------------
87// Types
88// ---------------------------------------------------------------------------
89
90#[derive(Clone, Copy, Debug, PartialEq, Eq)]
91pub enum SurvivalBaselineTarget {
92    /// No additional parametric target:
93    /// eta_target(t) = 0, so regularized model defaults to linear log-cumulative
94    /// hazard from the existing time basis.
95    Linear,
96    /// Parametric target: Weibull baseline.
97    ///
98    /// Transformation/cloglog survival uses `eta_target(t) = log(H0(t))`;
99    /// marginal-slope probit survival uses `q(t) = -Phi^-1(exp(-H0(t)))`.
100    Weibull,
101    /// Parametric target: Gompertz baseline.
102    ///
103    /// Transformation/cloglog survival uses `eta_target(t) = log(H0(t))`;
104    /// marginal-slope probit survival uses `q(t) = -Phi^-1(exp(-H0(t)))`.
105    Gompertz,
106    /// Parametric target: Gompertz-Makeham baseline.
107    ///
108    /// Transformation/cloglog survival uses `eta_target(t) = log(H0(t))`;
109    /// marginal-slope probit survival uses `q(t) = -Phi^-1(exp(-H0(t)))`.
110    GompertzMakeham,
111}
112
113#[derive(Clone, Debug)]
114pub struct SurvivalBaselineConfig {
115    pub target: SurvivalBaselineTarget,
116    pub scale: Option<f64>,
117    pub shape: Option<f64>,
118    pub rate: Option<f64>,
119    pub makeham: Option<f64>,
120}
121
122/// Recover the fitted Weibull baseline from the single-column `log(t)`
123/// time-basis coefficient.
124///
125/// The redundant constant column was dropped at design build (#2301 — it was the
126/// intercept-confounded `−shape·log_scale` location, now carried by the mean
127/// intercept), so the identified shape is the sole time coefficient `beta[0]` and
128/// the identified scale is the anchor itself. The fitted baseline is
129/// `shape * (log(t) - log(anchor))`.
130pub fn fitted_weibull_baseline_from_linear_time_beta(
131    beta: &Array1<f64>,
132    anchor: f64,
133) -> Option<SurvivalBaselineConfig> {
134    if beta.is_empty() {
135        return None;
136    }
137    let shape = beta[0];
138    if !shape.is_finite() || shape <= 0.0 || !anchor.is_finite() || anchor <= 0.0 {
139        return None;
140    }
141    Some(SurvivalBaselineConfig {
142        target: SurvivalBaselineTarget::Weibull,
143        scale: Some(anchor),
144        shape: Some(shape),
145        rate: None,
146        makeham: None,
147    })
148}
149
150#[derive(Clone, Debug)]
151pub enum SurvivalTimeBasisConfig {
152    None,
153    Linear,
154    BSpline {
155        degree: usize,
156        knots: Array1<f64>,
157        smooth_lambda: f64,
158    },
159    /// I-spline value rows on the `log(t)` axis with non-negative
160    /// coefficients (`γ ≥ 0`) enforcing structural monotonicity of
161    /// `q(t) = I_basis(log t) · γ`. This replaces the row-wise
162    /// `D β + o ≥ guard` derivative-guard constraints the marginal-slope
163    /// family previously relied on.
164    ///
165    /// The design builder lives below at `_build_time_block`'s
166    /// `SurvivalTimeBasisConfig::ISpline` arm and exposes:
167    ///
168    /// * `x_entry_time` / `x_exit_time` — I-spline value rows on the
169    ///   `log(t)` axis. Non-negative entries plus `γ ≥ 0` give a
170    ///   monotone-non-decreasing `q(t)`, the structural property the
171    ///   marginal-slope family needs.
172    /// * `x_derivative_time` — right-cumulative B-spline-derivative on
173    ///   `log(t)` scaled by `1/t`, again non-negative with `γ ≥ 0`, so
174    ///   `q'(t) ≥ 0` pointwise. The `derivative_guard` constant is added
175    ///   externally by [`add_survival_time_derivative_guard_offset`],
176    ///   leaving the derivative guarantee `q'(t) ≥ guard` exact.
177    /// * 2nd-difference penalty on the underlying degree-`(k+1)` B-spline
178    ///   coefficients, filtered through `keep_cols` for identifiability.
179    ///
180    /// `TimeBlockInput::time_monotonicity` declares to the consuming
181    /// family how monotonicity is enforced. The marginal-slope
182    /// construction site sets it to
183    /// [`crate::survival::location_scale::TimeBlockMonotonicity::StructuralISpline`]
184    /// so the family skips row-wise `D β + o ≥ guard` constraint
185    /// generation and treats `γ ≥ 0` as the sole derivative-guard
186    /// mechanism. The universal `validate_time_qd1_feasible` safety net
187    /// runs regardless.
188    ///
189    /// An earlier iteration proposed a separate C-spline antiderivative
190    /// parameterization that put `q'(t)` in the I-spline space and `q(t)`
191    /// in the integral-of-I-spline space. That was mathematically
192    /// equivalent but a strictly worse fit for the codebase (extra basis
193    /// degree, an extra antiderivative builder, an extra identifiability
194    /// path, an extra penalty); it was removed in favor of the canonical
195    /// I-spline-value path here.
196    ISpline {
197        degree: usize,
198        knots: Array1<f64>,
199        keep_cols: Vec<usize>,
200        smooth_lambda: f64,
201    },
202}
203
204/// Persistable snapshot of the time-basis state used by a survival fit.
205///
206/// Every survival family routes through [`SurvivalTimeBuildOutput`] during
207/// the fit, but the FFI save path needs only the metadata — not the full
208/// design matrices. This struct is the single source of truth that flows
209/// from the workflow-level basis construction, through the family-specific
210/// fit result, into the saved-model payload via
211/// [`crate::inference::model::FittedModelPayload::apply_survival_time_basis`].
212///
213/// Threading this snapshot end-to-end eliminates the prior bug pattern
214/// where each FFI builder had to reconstruct the metadata from
215/// `fit_config` + the formula (silent drift risk; one builder forgetting
216/// to do so caused the marginal-slope save→load break).
217#[derive(Clone, Debug, PartialEq)]
218pub struct SavedSurvivalTimeBasis {
219    pub basisname: String,
220    pub degree: Option<usize>,
221    pub knots: Option<Vec<f64>>,
222    pub keep_cols: Option<Vec<usize>>,
223    pub smooth_lambda: Option<f64>,
224    pub anchor: f64,
225}
226
227impl SavedSurvivalTimeBasis {
228    /// Build a snapshot from the realised time-basis state and the entry
229    /// anchor that was used during the fit.
230    pub fn from_build(build: &SurvivalTimeBuildOutput, anchor: f64) -> Self {
231        Self {
232            basisname: build.basisname.clone(),
233            degree: build.degree,
234            knots: build.knots.clone(),
235            keep_cols: build.keep_cols.clone(),
236            smooth_lambda: build.smooth_lambda,
237            anchor,
238        }
239    }
240}
241
242#[derive(Clone)]
243pub struct SurvivalTimeBuildOutput {
244    pub x_entry_time: DesignMatrix,
245    pub x_exit_time: DesignMatrix,
246    pub x_derivative_time: DesignMatrix,
247    pub penalties: Vec<Array2<f64>>,
248    /// Structural nullspace dimension of each penalty matrix.
249    pub nullspace_dims: Vec<usize>,
250    pub basisname: String,
251    pub degree: Option<usize>,
252    pub knots: Option<Vec<f64>>,
253    pub keep_cols: Option<Vec<usize>>,
254    pub smooth_lambda: Option<f64>,
255}
256
257pub const SURVIVAL_TIME_FLOOR: f64 = 1e-9;
258
259/// Seed smoothing penalty `λ` used when a survival time basis is reconstructed
260/// from a build (or saved model) that did not carry an explicit `smooth_lambda`.
261/// This is only an initial value for the REML smoothing search, not a fixed
262/// policy: a small positive seed keeps the baseline spline lightly regularized
263/// at the start so the outer optimizer begins from a well-conditioned point and
264/// then adapts `λ` to the data. Kept in one place so the b-spline and i-spline
265/// reconstruction paths cannot drift apart.
266const SURVIVAL_TIME_SMOOTH_LAMBDA_SEED: f64 = 1e-2;
267
268/// Default initial Gompertz / Gompertz-Makeham shape parameter when the user
269/// does not supply `--baseline-shape`. The Gompertz hazard is
270/// `h(t) = rate · exp(shape · t)`; a near-zero shape seeds the baseline at an
271/// almost-flat (exponential-like) hazard, letting the fit grow the
272/// age-acceleration term from the data rather than committing to a strong
273/// curvature up front. Shared by the parse and fit-seed paths so both start
274/// from the same neutral shape.
275const GOMPERTZ_DEFAULT_SHAPE_SEED: f64 = 0.01;
276
277#[derive(Clone, Copy, Debug, PartialEq, Eq)]
278pub enum SurvivalLikelihoodMode {
279    Transformation,
280    Weibull,
281    LocationScale,
282    MarginalSlope,
283    Latent,
284    LatentBinary,
285}
286
287/// Every survival likelihood mode, for the cross-mode contracts that must hold
288/// for all of them (e.g. the one time-basis anchor rule). Kept exhaustive by
289/// `survival_likelihood_modes_is_exhaustive`, which dispatches on the enum so a
290/// new variant fails to compile until it is listed here.
291pub const SURVIVAL_LIKELIHOOD_MODES: [SurvivalLikelihoodMode; 6] = [
292    SurvivalLikelihoodMode::Transformation,
293    SurvivalLikelihoodMode::Weibull,
294    SurvivalLikelihoodMode::LocationScale,
295    SurvivalLikelihoodMode::MarginalSlope,
296    SurvivalLikelihoodMode::Latent,
297    SurvivalLikelihoodMode::LatentBinary,
298];
299
300pub struct SurvivalTimeWiggleBuild {
301    pub penalties: Vec<Array2<f64>>,
302    pub nullspace_dims: Vec<usize>,
303    pub knots: Array1<f64>,
304    pub degree: usize,
305    pub ncols: usize,
306}
307
308// ---------------------------------------------------------------------------
309// Time normalization
310// ---------------------------------------------------------------------------
311
312pub fn normalize_survival_time_pair(
313    entry_raw: f64,
314    exit_raw: f64,
315    row_index: usize,
316) -> Result<(f64, f64), String> {
317    if !entry_raw.is_finite() || !exit_raw.is_finite() {
318        return Err(SurvivalConstructionError::DataValidationFailed {
319            reason: format!("non-finite survival times at row {}", row_index + 1),
320        }
321        .into());
322    }
323    if entry_raw < 0.0 || exit_raw < 0.0 {
324        return Err(SurvivalConstructionError::DataValidationFailed {
325            reason: format!("negative survival times at row {}", row_index + 1),
326        }
327        .into());
328    }
329
330    let entry = entry_raw.max(SURVIVAL_TIME_FLOOR);
331    let exit = exit_raw.max(entry + SURVIVAL_TIME_FLOOR);
332    Ok((entry, exit))
333}
334
335// ---------------------------------------------------------------------------
336// Basis monotonicity helpers
337// ---------------------------------------------------------------------------
338
339pub fn survival_basis_supports_structural_monotonicity(basisname: &str) -> bool {
340    basisname.eq_ignore_ascii_case("ispline")
341}
342
343pub fn require_structural_survival_time_basis(
344    basisname: &str,
345    context: &str,
346) -> Result<(), String> {
347    if survival_basis_supports_structural_monotonicity(basisname) {
348        return Ok(());
349    }
350    Err(SurvivalConstructionError::UnsupportedDistribution {
351        reason: format!(
352            "{context} requires a structural monotone survival time basis, but got '{basisname}'. \
353Only `ispline` is accepted here because its basis functions enforce a monotone cumulative time effect by construction. \
354`{basisname}` can fit non-monotone shapes, which can break survival semantics. \
355Re-run with `--time-basis ispline`."
356        ),
357    }
358    .into())
359}
360
361// ---------------------------------------------------------------------------
362// Baseline config parsing
363// ---------------------------------------------------------------------------
364
365pub fn parse_survival_baseline_config(
366    target_raw: &str,
367    scale: Option<f64>,
368    shape: Option<f64>,
369    rate: Option<f64>,
370    makeham: Option<f64>,
371) -> Result<SurvivalBaselineConfig, String> {
372    let target = match target_raw.to_ascii_lowercase().as_str() {
373        "linear" => SurvivalBaselineTarget::Linear,
374        "weibull" => SurvivalBaselineTarget::Weibull,
375        "gompertz" => SurvivalBaselineTarget::Gompertz,
376        "gompertz-makeham" => SurvivalBaselineTarget::GompertzMakeham,
377        other => {
378            return Err(SurvivalConstructionError::UnsupportedDistribution {
379                reason: format!(
380                    "unsupported --baseline-target '{other}'; use linear|weibull|gompertz|gompertz-makeham"
381                ),
382            }
383            .into());
384        }
385    };
386
387    match target {
388        SurvivalBaselineTarget::Linear => Ok(SurvivalBaselineConfig {
389            target,
390            scale: None,
391            shape: None,
392            rate: None,
393            makeham: None,
394        }),
395        SurvivalBaselineTarget::Weibull => {
396            let scale = scale.ok_or_else(|| {
397                "--baseline-target weibull requires --baseline-scale > 0".to_string()
398            })?;
399            let shape = shape.ok_or_else(|| {
400                "--baseline-target weibull requires --baseline-shape > 0".to_string()
401            })?;
402            if !scale.is_finite() || scale <= 0.0 || !shape.is_finite() || shape <= 0.0 {
403                return Err(
404                    "weibull baseline requires finite positive --baseline-scale and --baseline-shape"
405                        .to_string(),
406                );
407            }
408            Ok(SurvivalBaselineConfig {
409                target,
410                scale: Some(scale),
411                shape: Some(shape),
412                rate: None,
413                makeham: None,
414            })
415        }
416        SurvivalBaselineTarget::Gompertz => {
417            let rate = rate.unwrap_or(1.0);
418            let shape = shape.unwrap_or(GOMPERTZ_DEFAULT_SHAPE_SEED);
419            if !rate.is_finite() || rate <= 0.0 || !shape.is_finite() {
420                return Err(
421                    "gompertz baseline requires finite --baseline-shape and positive --baseline-rate"
422                        .to_string(),
423                );
424            }
425            Ok(SurvivalBaselineConfig {
426                target,
427                scale: None,
428                shape: Some(shape),
429                rate: Some(rate),
430                makeham: None,
431            })
432        }
433        SurvivalBaselineTarget::GompertzMakeham => {
434            let rate = rate.unwrap_or(0.5);
435            let shape = shape.unwrap_or(GOMPERTZ_DEFAULT_SHAPE_SEED);
436            let makeham = makeham.unwrap_or(0.5);
437            if !rate.is_finite()
438                || rate <= 0.0
439                || !shape.is_finite()
440                || !makeham.is_finite()
441                || makeham <= 0.0
442            {
443                return Err(
444                    "gompertz-makeham baseline requires finite --baseline-shape, positive --baseline-rate, and positive --baseline-makeham"
445                        .to_string(),
446                );
447            }
448            Ok(SurvivalBaselineConfig {
449                target,
450                scale: None,
451                shape: Some(shape),
452                rate: Some(rate),
453                makeham: Some(makeham),
454            })
455        }
456    }
457}
458
459// ---------------------------------------------------------------------------
460// Likelihood mode / distribution parsing
461// ---------------------------------------------------------------------------
462
463pub fn parse_survival_likelihood_mode(raw: &str) -> Result<SurvivalLikelihoodMode, String> {
464    match raw.to_ascii_lowercase().as_str() {
465        "transformation" => Ok(SurvivalLikelihoodMode::Transformation),
466        "weibull" => Ok(SurvivalLikelihoodMode::Weibull),
467        "location-scale" => Ok(SurvivalLikelihoodMode::LocationScale),
468        "marginal-slope" => Ok(SurvivalLikelihoodMode::MarginalSlope),
469        "latent" => Ok(SurvivalLikelihoodMode::Latent),
470        "latent-binary" => Ok(SurvivalLikelihoodMode::LatentBinary),
471        other => Err(SurvivalConstructionError::UnsupportedDistribution {
472            reason: format!(
473                "unsupported --survival-likelihood '{other}'; use transformation|weibull|location-scale|marginal-slope|latent|latent-binary"
474            ),
475        }
476        .into()),
477    }
478}
479
480pub const fn survival_likelihood_modename(mode: SurvivalLikelihoodMode) -> &'static str {
481    match mode {
482        SurvivalLikelihoodMode::Transformation => "transformation",
483        SurvivalLikelihoodMode::Weibull => "weibull",
484        SurvivalLikelihoodMode::LocationScale => "location-scale",
485        SurvivalLikelihoodMode::MarginalSlope => "marginal-slope",
486        SurvivalLikelihoodMode::Latent => "latent",
487        SurvivalLikelihoodMode::LatentBinary => "latent-binary",
488    }
489}
490
491pub fn parse_survival_distribution(raw: &str) -> Result<ResidualDistribution, String> {
492    match raw.to_ascii_lowercase().as_str() {
493        "gaussian" | "probit" => Ok(ResidualDistribution::Gaussian),
494        "gumbel" | "cloglog" => Ok(ResidualDistribution::Gumbel),
495        "logistic" | "logit" => Ok(ResidualDistribution::Logistic),
496        other => Err(SurvivalConstructionError::UnsupportedDistribution {
497            reason: format!(
498                "unsupported survmodel(distribution='{other}'); accepted: gaussian / probit, gumbel / cloglog, logistic / logit"
499            ),
500        }
501        .into()),
502    }
503}
504
505pub const fn survival_baseline_targetname(target: SurvivalBaselineTarget) -> &'static str {
506    match target {
507        SurvivalBaselineTarget::Linear => "linear",
508        SurvivalBaselineTarget::Weibull => "weibull",
509        SurvivalBaselineTarget::Gompertz => "gompertz",
510        SurvivalBaselineTarget::GompertzMakeham => "gompertz-makeham",
511    }
512}
513
514pub fn positive_survival_time_seed(age_exit: &Array1<f64>) -> f64 {
515    let sum = age_exit
516        .iter()
517        .copied()
518        .filter(|value| value.is_finite() && *value > 0.0)
519        .sum::<f64>();
520    let count = age_exit
521        .iter()
522        .filter(|value| value.is_finite() && **value > 0.0)
523        .count()
524        .max(1);
525    (sum / count as f64).max(SURVIVAL_TIME_FLOOR)
526}
527
528pub fn initial_survival_baseline_config_for_fit(
529    target_raw: &str,
530    scale: Option<f64>,
531    shape: Option<f64>,
532    rate: Option<f64>,
533    makeham: Option<f64>,
534    age_exit: &Array1<f64>,
535) -> Result<SurvivalBaselineConfig, String> {
536    let target = match target_raw.trim().to_ascii_lowercase().as_str() {
537        "linear" => SurvivalBaselineTarget::Linear,
538        "weibull" => SurvivalBaselineTarget::Weibull,
539        "gompertz" => SurvivalBaselineTarget::Gompertz,
540        "gompertz-makeham" => SurvivalBaselineTarget::GompertzMakeham,
541        other => {
542            return Err(SurvivalConstructionError::UnsupportedDistribution {
543                reason: format!(
544                    "unsupported --baseline-target '{other}'; use linear|weibull|gompertz|gompertz-makeham"
545                ),
546            }
547            .into());
548        }
549    };
550    let time_scale_seed = positive_survival_time_seed(age_exit);
551    let cfg = match target {
552        SurvivalBaselineTarget::Linear => SurvivalBaselineConfig {
553            target,
554            scale: None,
555            shape: None,
556            rate: None,
557            makeham: None,
558        },
559        SurvivalBaselineTarget::Weibull => SurvivalBaselineConfig {
560            target,
561            scale: Some(scale.unwrap_or(time_scale_seed)),
562            shape: Some(shape.unwrap_or(1.0)),
563            rate: None,
564            makeham: None,
565        },
566        SurvivalBaselineTarget::Gompertz => SurvivalBaselineConfig {
567            target,
568            scale: None,
569            shape: Some(shape.unwrap_or(GOMPERTZ_DEFAULT_SHAPE_SEED)),
570            rate: Some(rate.unwrap_or(1.0 / time_scale_seed)),
571            makeham: None,
572        },
573        SurvivalBaselineTarget::GompertzMakeham => SurvivalBaselineConfig {
574            target,
575            scale: None,
576            shape: Some(shape.unwrap_or(GOMPERTZ_DEFAULT_SHAPE_SEED)),
577            rate: Some(rate.unwrap_or(0.5 / time_scale_seed)),
578            makeham: Some(makeham.unwrap_or(0.5 / time_scale_seed)),
579        },
580    };
581    parse_survival_baseline_config(
582        survival_baseline_targetname(cfg.target),
583        cfg.scale,
584        cfg.shape,
585        cfg.rate,
586        cfg.makeham,
587    )
588}
589
590pub fn survival_baseline_theta_from_config(
591    cfg: &SurvivalBaselineConfig,
592) -> Result<Option<Array1<f64>>, String> {
593    let theta = match cfg.target {
594        SurvivalBaselineTarget::Linear => None,
595        SurvivalBaselineTarget::Weibull => Some(array![
596            cfg.scale
597                .ok_or_else(|| "missing weibull baseline scale".to_string())?
598                .ln(),
599            cfg.shape
600                .ok_or_else(|| "missing weibull baseline shape".to_string())?
601                .ln(),
602        ]),
603        SurvivalBaselineTarget::Gompertz => Some(array![
604            cfg.rate
605                .ok_or_else(|| "missing gompertz baseline rate".to_string())?
606                .ln(),
607            cfg.shape
608                .ok_or_else(|| "missing gompertz baseline shape".to_string())?,
609        ]),
610        SurvivalBaselineTarget::GompertzMakeham => Some(array![
611            cfg.rate
612                .ok_or_else(|| "missing gompertz-makeham baseline rate".to_string())?
613                .ln(),
614            cfg.shape
615                .ok_or_else(|| "missing gompertz-makeham baseline shape".to_string())?,
616            cfg.makeham
617                .ok_or_else(|| "missing gompertz-makeham baseline makeham".to_string())?
618                .ln(),
619        ]),
620    };
621    if let Some(theta) = theta.as_ref() {
622        if theta.iter().any(|value| !value.is_finite()) {
623            return Err(format!(
624                "{} baseline theta coordinates must be finite",
625                survival_baseline_targetname(cfg.target)
626            ));
627        }
628        // The inverse chart is also the target-specific domain validator. This
629        // keeps the public encoder from emitting coordinates for an invalid
630        // config (including when a caller has no data rows to evaluate).
631        survival_baseline_config_from_theta(cfg.target, theta)?;
632    }
633    Ok(theta)
634}
635
636pub fn survival_baseline_config_from_theta(
637    target: SurvivalBaselineTarget,
638    theta: &Array1<f64>,
639) -> Result<SurvivalBaselineConfig, String> {
640    let cfg = match target {
641        SurvivalBaselineTarget::Linear => SurvivalBaselineConfig {
642            target,
643            scale: None,
644            shape: None,
645            rate: None,
646            makeham: None,
647        },
648        SurvivalBaselineTarget::Weibull => {
649            if theta.len() != 2 {
650                return Err(SurvivalConstructionError::IncompatibleDimensions {
651                    reason: format!(
652                        "weibull baseline parameter dimension mismatch: expected 2, got {}",
653                        theta.len()
654                    ),
655                }
656                .into());
657            }
658            SurvivalBaselineConfig {
659                target,
660                scale: Some(theta[0].exp()),
661                shape: Some(theta[1].exp()),
662                rate: None,
663                makeham: None,
664            }
665        }
666        SurvivalBaselineTarget::Gompertz => {
667            if theta.len() != 2 {
668                return Err(SurvivalConstructionError::IncompatibleDimensions {
669                    reason: format!(
670                        "gompertz baseline parameter dimension mismatch: expected 2, got {}",
671                        theta.len()
672                    ),
673                }
674                .into());
675            }
676            SurvivalBaselineConfig {
677                target,
678                scale: None,
679                shape: Some(theta[1]),
680                rate: Some(theta[0].exp()),
681                makeham: None,
682            }
683        }
684        SurvivalBaselineTarget::GompertzMakeham => {
685            if theta.len() != 3 {
686                return Err(SurvivalConstructionError::IncompatibleDimensions {
687                    reason: format!(
688                        "gompertz-makeham baseline parameter dimension mismatch: expected 3, got {}",
689                        theta.len()
690                    ),
691                }
692                .into());
693            }
694            SurvivalBaselineConfig {
695                target,
696                scale: None,
697                shape: Some(theta[1]),
698                rate: Some(theta[0].exp()),
699                makeham: Some(theta[2].exp()),
700            }
701        }
702    };
703    parse_survival_baseline_config(
704        survival_baseline_targetname(cfg.target),
705        cfg.scale,
706        cfg.shape,
707        cfg.rate,
708        cfg.makeham,
709    )
710}
711
712/// Derivative contract for the shared baseline-θ outer optimizer.
713///
714/// The two public baseline optimizers (`…_with_gradient_only`,
715/// `…_with_gradient`) differ in exactly one axis: how much derivative
716/// information the objective closure supplies, and therefore which curvature
717/// declaration the `OuterProblem` must advertise. Every baseline-θ path now
718/// supplies an exact analytic gradient (profile-NLL envelope gradient), so both
719/// contracts route to a gradient-based solver. Everything else — θ↔config
720/// conversion, the ±6 log-space box,
721/// the single-seed config, the `run`/convergence/error-formatting boilerplate
722/// — is identical, so it lives once in [`run_baseline_theta_optimizer`] and
723/// this enum selects the per-contract `OuterProblem` configuration.
724#[derive(Clone, Copy, Debug, PartialEq, Eq)]
725enum BaselineDerivativeContract {
726    /// Cost + analytic gradient, no analytic Hessian. Routes to BFGS, which
727    /// builds its own quasi-Newton curvature from successive gradients.
728    GradientOnly,
729    /// Cost + analytic gradient + analytic Hessian. Search uses BFGS while the
730    /// terminal mint certificate evaluates the analytic Hessian once.
731    GradientHessian,
732}
733
734impl BaselineDerivativeContract {
735    /// Apply this contract's derivative declaration, solver class, tolerance,
736    /// and iteration budget to a freshly-constructed `OuterProblem`. The
737    /// bounds, initial ρ, and seed config are contract-independent and applied
738    /// by [`run_baseline_theta_optimizer`].
739    fn configure(
740        self,
741        problem: gam_solve::rho_optimizer::OuterProblem,
742    ) -> gam_solve::rho_optimizer::OuterProblem {
743        use gam_problem::{DeclaredHessianForm, Derivative};
744        match self {
745            // BFGS on a 2–3 dim problem with an exact gradient typically
746            // converges in 5–10 outer evaluations.
747            BaselineDerivativeContract::GradientOnly => problem
748                .with_gradient(Derivative::Analytic)
749                .with_hessian(DeclaredHessianForm::Unavailable)
750                .with_tolerance(1e-4)
751                .with_max_iter(240),
752            BaselineDerivativeContract::GradientHessian => problem
753                .with_gradient(Derivative::Analytic)
754                .with_hessian(DeclaredHessianForm::Either)
755                .with_tolerance(1e-4)
756                .with_max_iter(240),
757        }
758    }
759}
760
761/// Shared engine behind the three public baseline-config optimizers.
762///
763/// Owns every step that is identical across the cost-only, gradient-only, and
764/// gradient+Hessian contracts: config→θ seeding (with the linear/no-parameter
765/// early return), the ±6 log-space box, the single-seed `OuterProblem`
766/// skeleton, derivative-contract configuration, `build_objective` wiring,
767/// `run`, the convergence check + error formatting, and θ→config. The only
768/// contract-specific inputs are the already-wired `cost_fn`/`eval_fn` closures
769/// (which embed the derivative shape and dimension validation) and the
770/// `contract` selecting the `OuterProblem` derivative declaration.
771fn run_baseline_theta_optimizer<Fc, Fe>(
772    initial: &SurvivalBaselineConfig,
773    context: &str,
774    contract: BaselineDerivativeContract,
775    cost_fn: Fc,
776    eval_fn: Fe,
777) -> Result<SurvivalBaselineConfig, String>
778where
779    Fc: FnMut(&mut (), &Array1<f64>) -> Result<f64, crate::model_types::EstimationError>,
780    Fe: FnMut(
781        &mut (),
782        &Array1<f64>,
783    ) -> Result<gam_problem::OuterEval, crate::model_types::EstimationError>,
784{
785    use gam_solve::rho_optimizer::OuterProblem;
786    let Some(seed) = survival_baseline_theta_from_config(initial)? else {
787        return Ok(initial.clone());
788    };
789    let dim = seed.len();
790    let target = initial.target;
791    let lower = seed.mapv(|v| v - 6.0);
792    let upper = seed.mapv(|v| v + 6.0);
793    let problem = contract
794        .configure(OuterProblem::new(dim).with_prefer_gradient_only(true))
795        .with_bounds(lower, upper)
796        .with_initial_rho(seed.clone())
797        .with_seed_config(crate::seeding::SeedConfig {
798            max_seeds: 1,
799            seed_budget: 1,
800            num_auxiliary_trailing: dim,
801            ..Default::default()
802        });
803    let mut obj = problem.build_objective(
804        (),
805        cost_fn,
806        eval_fn,
807        None::<fn(&mut ())>,
808        None::<
809            fn(
810                &mut (),
811                &Array1<f64>,
812            ) -> Result<gam_problem::EfsEval, crate::model_types::EstimationError>,
813        >,
814    );
815    let result = problem
816        .run(&mut obj, context)
817        .map_err(|e| format!("{context} failed: {e}"))?;
818    if !result.converged() {
819        return Err(SurvivalConstructionError::InvalidConfig {
820            reason: format!(
821                "{context} did not converge after {} iterations (final_objective={:.6e}, final_grad_norm={})",
822                result.iterations,
823                result.final_value,
824                result.final_grad_norm_report(),
825            ),
826        }
827        .into());
828    }
829    survival_baseline_config_from_theta(target, &result.rho)
830}
831
832/// Shared engine for the two derivative-carrying baseline-config optimizers.
833///
834/// Both `…_with_gradient_only` and `…_with_gradient` route an objective that
835/// returns a fully-populated [`OuterEval`](gam_problem::OuterEval)
836/// (cost + analytic gradient, optionally + analytic Hessian) for a given
837/// config. Everything downstream of that — the `Rc<RefCell>` sharing that lets
838/// the same user closure back both the `cost_fn` and `eval_fn`, the θ→config
839/// conversion, and deriving the scalar `cost_fn` from the eval result — is
840/// identical, so it lives here once. The contract-specific axis is only which
841/// `HessianValue` the objective embeds, which the wrapper has already encoded
842/// in the returned `OuterEval`, so this helper is contract-agnostic beyond the
843/// `contract` it forwards to [`run_baseline_theta_optimizer`].
844fn run_baseline_theta_optimizer_with_eval<F>(
845    initial: &SurvivalBaselineConfig,
846    context: &str,
847    contract: BaselineDerivativeContract,
848    objective: F,
849) -> Result<SurvivalBaselineConfig, String>
850where
851    F: FnMut(&SurvivalBaselineConfig) -> Result<gam_problem::OuterEval, String>,
852{
853    let target = initial.target;
854    let engine_context = context.to_string();
855    let objective = std::rc::Rc::new(std::cell::RefCell::new(objective));
856    let eval_at = move |obj: &std::rc::Rc<std::cell::RefCell<F>>,
857                        theta: &Array1<f64>|
858          -> Result<gam_problem::OuterEval, crate::model_types::EstimationError> {
859        let cfg = survival_baseline_config_from_theta(target, theta)
860            .map_err(crate::model_types::EstimationError::InvalidInput)?;
861        let eval =
862            obj.borrow_mut()(&cfg).map_err(crate::model_types::EstimationError::InvalidInput)?;
863        if eval.gradient.len() != theta.len() {
864            return Err(crate::model_types::EstimationError::InvalidInput(format!(
865                "{engine_context}: baseline gradient dimension mismatch: got {}, expected {}",
866                eval.gradient.len(),
867                theta.len()
868            )));
869        }
870        if let gam_problem::HessianValue::Dense(ref h) = eval.hessian {
871            if h.nrows() != theta.len() || h.ncols() != theta.len() {
872                return Err(crate::model_types::EstimationError::InvalidInput(format!(
873                    "{engine_context}: baseline Hessian dimension mismatch: got {}x{}, expected {}x{}",
874                    h.nrows(),
875                    h.ncols(),
876                    theta.len(),
877                    theta.len()
878                )));
879            }
880        }
881        Ok(eval)
882    };
883    let cost_objective = std::rc::Rc::clone(&objective);
884    let cost_eval = eval_at.clone();
885    let cost_fn = move |_: &mut (), theta: &Array1<f64>| {
886        cost_eval(&cost_objective, theta).map(|eval| eval.cost)
887    };
888    let eval_fn = move |_: &mut (), theta: &Array1<f64>| eval_at(&objective, theta);
889    run_baseline_theta_optimizer(initial, context, contract, cost_fn, eval_fn)
890}
891
892/// Gradient-only outer baseline-config optimizer. Thin adapter over
893/// `run_baseline_theta_optimizer` under the
894/// `BaselineDerivativeContract::GradientOnly` contract, which advertises
895/// `DeclaredHessianForm::Unavailable`, so the planner routes to BFGS and
896/// builds its own quasi-Newton curvature from successive gradient
897/// evaluations. Used by the survival location-scale path which has a
898/// closed-form θ-gradient (`baseline_chain_rule_gradient` /
899/// `marginal_slope_baseline_chain_rule_gradient`) but no native analytic
900/// θ-Hessian; BFGS on a 2–3 dim problem with an exact gradient typically
901/// converges in 5–10 outer evaluations.
902pub fn optimize_survival_baseline_config_with_gradient_only<F>(
903    initial: &SurvivalBaselineConfig,
904    context: &str,
905    mut objective: F,
906) -> Result<SurvivalBaselineConfig, String>
907where
908    F: FnMut(&SurvivalBaselineConfig) -> Result<(f64, Array1<f64>), String>,
909{
910    use gam_problem::{HessianValue, OuterEval};
911    run_baseline_theta_optimizer_with_eval(
912        initial,
913        context,
914        BaselineDerivativeContract::GradientOnly,
915        move |cfg| {
916            let (cost, gradient) = objective(cfg)?;
917            Ok(OuterEval {
918                cost,
919                gradient,
920                hessian: HessianValue::Unavailable,
921                inner_beta_hint: None,
922            })
923        },
924    )
925}
926
927/// Gradient + Hessian outer baseline-config optimizer. Thin adapter over
928/// `run_baseline_theta_optimizer` under the
929/// `BaselineDerivativeContract::GradientHessian` contract, which advertises
930/// an analytic θ-Hessian so terminal mint certification can audit it.
931pub fn optimize_survival_baseline_config_with_gradient<F>(
932    initial: &SurvivalBaselineConfig,
933    context: &str,
934    mut objective: F,
935) -> Result<SurvivalBaselineConfig, String>
936where
937    F: FnMut(&SurvivalBaselineConfig) -> Result<(f64, Array1<f64>, Array2<f64>), String>,
938{
939    use gam_problem::{HessianValue, OuterEval};
940    run_baseline_theta_optimizer_with_eval(
941        initial,
942        context,
943        BaselineDerivativeContract::GradientHessian,
944        move |cfg| {
945            let (cost, gradient, hessian) = objective(cfg)?;
946            Ok(OuterEval {
947                cost,
948                gradient,
949                hessian: HessianValue::Dense(hessian),
950                inner_beta_hint: None,
951            })
952        },
953    )
954}
955
956// ---------------------------------------------------------------------------
957// Time basis config (library-friendly: takes primitives, not CLI args)
958// ---------------------------------------------------------------------------
959
960pub fn parse_survival_time_basis_config(
961    time_basis: &str,
962    time_degree: usize,
963    time_num_internal_knots: usize,
964    time_smooth_lambda: f64,
965) -> Result<SurvivalTimeBasisConfig, String> {
966    match time_basis.to_ascii_lowercase().as_str() {
967        "none" => Ok(SurvivalTimeBasisConfig::None),
968        "ispline" => {
969            if time_degree < 1 {
970                return Err(
971                    "time-basis degree must be >= 1 for ispline time basis (CLI: --time-degree; Python: time_degree=)"
972                        .to_string(),
973                );
974            }
975            if time_num_internal_knots == 0 {
976                return Err(
977                    "time-basis must have > 0 internal knots for ispline time basis (CLI: --time-num-internal-knots; Python: time_num_internal_knots=)"
978                        .to_string(),
979                );
980            }
981            if !time_smooth_lambda.is_finite() || time_smooth_lambda < 0.0 {
982                return Err(
983                    "time-basis smoothing lambda must be finite and >= 0 (CLI: --time-smooth-lambda; Python: time_smooth_lambda=)"
984                        .to_string(),
985                );
986            }
987            Ok(SurvivalTimeBasisConfig::ISpline {
988                degree: time_degree,
989                knots: Array1::zeros(0),
990                keep_cols: Vec::new(),
991                smooth_lambda: time_smooth_lambda,
992            })
993        }
994        "linear" | "bspline" => {
995            // Forward to the shared structural-basis check so error text
996            // stays consistent with every other call site. `linear` /
997            // `bspline` are not structural, so this always returns Err;
998            // we map a (currently impossible) `Ok` to an explicit error
999            // string instead of `unreachable!`, keeping the match total
1000            // without relying on a never-executes claim.
1001            match require_structural_survival_time_basis(time_basis, "survival model configuration")
1002            {
1003                Err(e) => Err(e),
1004                Ok(()) => Err(format!(
1005                    "internal: structural-basis check accepted non-structural \
1006                     survival time basis '{time_basis}'"
1007                )),
1008            }
1009        }
1010        other => Err(format!(
1011            "unsupported --time-basis '{other}'; accepted values: ispline, none"
1012        )),
1013    }
1014}
1015
1016// ---------------------------------------------------------------------------
1017// Time basis construction
1018// ---------------------------------------------------------------------------
1019
1020pub fn build_survival_time_basis(
1021    age_entry: &Array1<f64>,
1022    age_exit: &Array1<f64>,
1023    cfg: SurvivalTimeBasisConfig,
1024    infer_knots_if_needed: Option<(usize, f64)>,
1025) -> Result<SurvivalTimeBuildOutput, String> {
1026    fn checked_log_survival_times(times: &Array1<f64>, label: &str) -> Result<Array1<f64>, String> {
1027        if let Some(row) = times.iter().position(|t| !t.is_finite()) {
1028            return Err(SurvivalConstructionError::DataValidationFailed {
1029                reason: format!(
1030                    "survival time basis requires finite {label} times (row {})",
1031                    row + 1
1032                ),
1033            }
1034            .into());
1035        }
1036        if let Some(row) = times.iter().position(|t| *t < 0.0) {
1037            return Err(SurvivalConstructionError::DataValidationFailed {
1038                reason: format!(
1039                    "survival time basis requires non-negative {label} times (row {})",
1040                    row + 1
1041                ),
1042            }
1043            .into());
1044        }
1045        Ok(times.mapv(|t| t.max(SURVIVAL_TIME_FLOOR).ln()))
1046    }
1047
1048    let n = age_entry.len();
1049    if n != age_exit.len() {
1050        return Err(SurvivalConstructionError::IncompatibleDimensions {
1051            reason: "survival time basis requires matching entry/exit lengths".to_string(),
1052        }
1053        .into());
1054    }
1055    for i in 0..n {
1056        if age_exit[i] < age_entry[i] {
1057            return Err(format!(
1058                "survival time basis requires exit times >= entry times (row {})",
1059                i + 1
1060            ));
1061        }
1062    }
1063    let log_entry = checked_log_survival_times(age_entry, "entry")?;
1064    let log_exit = checked_log_survival_times(age_exit, "exit")?;
1065
1066    fn survival_time_knot_input(log_entry: &Array1<f64>, log_exit: &Array1<f64>) -> Array1<f64> {
1067        let n = log_entry.len();
1068        let entry_range = log_entry
1069            .iter()
1070            .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| {
1071                (lo.min(v), hi.max(v))
1072            });
1073        let entry_degenerate = (entry_range.1 - entry_range.0).abs() < 1e-8;
1074        if entry_degenerate {
1075            log_exit.clone()
1076        } else {
1077            let mut combined = Array1::<f64>::zeros(2 * n);
1078            for i in 0..n {
1079                combined[i] = log_entry[i];
1080                combined[n + i] = log_exit[i];
1081            }
1082            combined
1083        }
1084    }
1085
1086    /// Cap the requested monotone-baseline internal-knot count to what the
1087    /// observed time resolution can actually support.
1088    ///
1089    /// The survival location-scale baseline is a degree-`d` I-spline with
1090    /// `num_internal_knots + d` shape-varying columns. Its smoothing parameter
1091    /// is informed *only* by the distinct interior log-time points: with fewer
1092    /// distinct interior times than requested knots the baseline is
1093    /// rank-deficient, and the REML/LAML profile in the time smoothing
1094    /// parameter becomes a flat ridge — the exact-joint outer search then
1095    /// probes that ridge indefinitely (each inner constrained Newton burns its
1096    /// whole cycle budget without certifying convergence) and the fit never
1097    /// terminates. This is the survival analogue of the standard
1098    /// "df must not exceed the data resolution" guard (`mgcv` caps `k` at the
1099    /// number of unique covariate values; `flexsurv`/`rstpm2` use a handful of
1100    /// baseline knots): we never place more interior knots than there are
1101    /// distinct interior points, and we keep the total baseline dimension a
1102    /// bounded fraction of the sample so the smoothing profile stays curved.
1103    ///
1104    /// This clamp lives in the shared knot-inference routine so the fit and any
1105    /// independent rebuild of the time basis (e.g. a predictor reconstructing
1106    /// `design · β` at fresh covariates) resolve to the *same* knot vector from
1107    /// the same data — there is no raw/active dimension drift.
1108    fn data_capped_internal_knots(
1109        combined: &Array1<f64>,
1110        degree: usize,
1111        requested_internal_knots: usize,
1112    ) -> usize {
1113        if requested_internal_knots == 0 {
1114            return 0;
1115        }
1116        let mut sorted: Vec<f64> = combined.iter().copied().collect();
1117        sorted.sort_by(f64::total_cmp);
1118        let minval = sorted.first().copied().unwrap_or(0.0);
1119        let maxval = sorted.last().copied().unwrap_or(minval);
1120        if minval == maxval {
1121            // Degenerate (single distinct time): no interior structure to fit.
1122            return 1.min(requested_internal_knots);
1123        }
1124        let scale = (maxval - minval).abs().max(1.0);
1125        let tol = 1e-12 * scale;
1126        // Count distinct strictly-interior points (knots can only live strictly
1127        // between the data extremes).
1128        let mut distinct_interior = 0usize;
1129        let mut last: Option<f64> = None;
1130        for &x in &sorted {
1131            if x <= minval + tol || x >= maxval - tol {
1132                continue;
1133            }
1134            if last.is_some_and(|prev| (x - prev).abs() <= tol) {
1135                continue;
1136            }
1137            distinct_interior += 1;
1138            last = Some(x);
1139        }
1140        // Distinct-point ceiling: cannot place more interior knots than there
1141        // are distinct interior values.
1142        let mut cap = requested_internal_knots.min(distinct_interior.max(1));
1143        // Dimension-vs-resolution ceiling: keep the total baseline column count
1144        // `cap + degree` below ~1/4 of the distinct sample points so the
1145        // smoothing-parameter profile retains curvature (the data must be able
1146        // to identify the baseline shape, not just interpolate it). `n_distinct`
1147        // counts all distinct points (interior + the two extremes).
1148        let n_distinct = {
1149            let mut count = 0usize;
1150            let mut last: Option<f64> = None;
1151            for &x in &sorted {
1152                if last.is_some_and(|prev| (x - prev).abs() <= tol) {
1153                    continue;
1154                }
1155                count += 1;
1156                last = Some(x);
1157            }
1158            count
1159        };
1160        let dim_budget = n_distinct / 4;
1161        let dim_cap = dim_budget.saturating_sub(degree);
1162        cap = cap.min(dim_cap.max(1));
1163        cap.max(1)
1164    }
1165
1166    /// Infer a survival time knot vector, reporting the PUBLIC basis degree
1167    /// the returned vector can actually carry.
1168    ///
1169    /// `build_bspline_basis_1d` may auto-shrink the requested degree when the
1170    /// data cannot support it (issue #340) -- with 4 rows a degree-4 clamped
1171    /// vector is not constructible, so it silently returns a degree-3 one and
1172    /// records that in `BasisMetadata::BSpline1D::degree`. Callers must be told,
1173    /// or they will hand a shrunk vector to a consumer sized for the degree they
1174    /// asked for.
1175    fn infer_survival_time_knots_with_degree(
1176        combined: &Array1<f64>,
1177        knot_degree: usize,
1178        validation_degree: usize,
1179        num_internal_knots: usize,
1180        basis_options: BasisOptions,
1181    ) -> Result<(Array1<f64>, usize), String> {
1182        // Identifiability/termination guard: never request more baseline
1183        // internal knots than the observed time resolution supports. See
1184        // `data_capped_internal_knots` for the full rationale (a flat smoothing
1185        // ridge on an over-parameterized baseline is what makes the survival
1186        // location-scale exact-joint outer search fail to terminate).
1187        let num_internal_knots =
1188            data_capped_internal_knots(combined, validation_degree, num_internal_knots);
1189
1190        fn quantile_knot_inference_needs_uniform_fallback(
1191            combined: &Array1<f64>,
1192            num_internal_knots: usize,
1193        ) -> bool {
1194            if num_internal_knots == 0 || combined.is_empty() {
1195                return false;
1196            }
1197
1198            let mut sorted: Vec<f64> = combined.iter().copied().collect();
1199            sorted.sort_by(f64::total_cmp);
1200            let minval = sorted[0];
1201            let maxval = *sorted.last().unwrap_or(&minval);
1202            if minval == maxval {
1203                return false;
1204            }
1205
1206            let scale = (maxval - minval).abs().max(1.0);
1207            let tol = 1e-12 * scale;
1208            let mut support = Vec::with_capacity(sorted.len());
1209            let mut last: Option<f64> = None;
1210            for &x in &sorted {
1211                if x <= minval + tol || x >= maxval - tol {
1212                    continue;
1213                }
1214                if last.map(|prev| (x - prev).abs() <= tol).unwrap_or(false) {
1215                    continue;
1216                }
1217                support.push(x);
1218                last = Some(x);
1219            }
1220            if support.is_empty() {
1221                return true;
1222            }
1223
1224            let n = support.len();
1225            let mut prev_q = minval;
1226            for j in 1..=num_internal_knots {
1227                let p = j as f64 / (num_internal_knots + 1) as f64;
1228                let pos = p * (n.saturating_sub(1) as f64);
1229                let lo = pos.floor() as usize;
1230                let hi = pos.ceil() as usize;
1231                let frac = pos - lo as f64;
1232                let q = if lo == hi {
1233                    support[lo]
1234                } else {
1235                    support[lo] * (1.0 - frac) + support[hi] * frac
1236                }
1237                .clamp(minval, maxval);
1238                if q <= prev_q + tol || q >= maxval - tol {
1239                    return true;
1240                }
1241                prev_q = q;
1242            }
1243
1244            false
1245        }
1246
1247        let inferwith =
1248            |placement: gam_terms::basis::BSplineKnotPlacement|
1249             -> Result<(Array1<f64>, usize), String> {
1250                let built = build_bspline_basis_1d(
1251                    combined.view(),
1252                    &BSplineBasisSpec {
1253                        degree: knot_degree,
1254                        penalty_order: 2,
1255                        knotspec: BSplineKnotSpec::Automatic {
1256                            num_internal_knots: Some(num_internal_knots),
1257                            placement,
1258                        },
1259                        double_penalty: false,
1260                        identifiability: BSplineIdentifiability::None,
1261                        boundary: OneDimensionalBoundary::Open,
1262                        boundary_conditions: BSplineBoundaryConditions::default(),
1263                    },
1264                )
1265                .map_err(|e| format!("failed to infer survival time knots: {e}"))?;
1266                let (knots, built_degree) = match built.metadata {
1267                    BasisMetadata::BSpline1D { knots, degree, .. } => {
1268                        (knots, degree.unwrap_or(knot_degree))
1269                    }
1270                    _ => {
1271                        return Err(
1272                            "internal error: expected BSpline1D metadata for survival time basis"
1273                                .to_string(),
1274                        );
1275                    }
1276                };
1277                // `knot_degree` is the clamped B-spline degree used to size
1278                // the knot vector. `validation_degree` is the public basis
1279                // degree passed to the final evaluator. They differ for
1280                // I-splines because `create_basis(..., BasisOptions::i_spline())`
1281                // internally raises the public degree by one to its working
1282                // B-spline antiderivative degree. Validating with
1283                // `knot_degree` here would raise a second time and reject the
1284                // coherent knot vector we just inferred.
1285                // The caller's two degrees differ by a fixed raise: `i_spline()`
1286                // lifts the public degree to its working B-spline antiderivative
1287                // degree, so `knot_degree == validation_degree + raise`. When the
1288                // builder shrinks the vector, the public degree has to come down
1289                // by the same raise or the two stop describing one geometry.
1290                let raise = knot_degree.saturating_sub(validation_degree);
1291                let effective_validation_degree = built_degree.saturating_sub(raise);
1292                create_basis::<Dense>(
1293                    combined.view(),
1294                    KnotSource::Provided(knots.view()),
1295                    effective_validation_degree,
1296                    basis_options,
1297                )
1298                .map_err(|e| e.to_string())?;
1299                Ok((knots, effective_validation_degree))
1300            };
1301
1302        if quantile_knot_inference_needs_uniform_fallback(combined, num_internal_knots) {
1303            inferwith(gam_terms::basis::BSplineKnotPlacement::Uniform)
1304        } else {
1305            inferwith(gam_terms::basis::BSplineKnotPlacement::Quantile)
1306        }
1307    }
1308
1309    /// Knot vector only, for the callers whose consumer degree is the one they
1310    /// passed in (no i-spline raise, so a shrink cannot desynchronise anything).
1311    fn infer_survival_time_knots(
1312        combined: &Array1<f64>,
1313        knot_degree: usize,
1314        validation_degree: usize,
1315        num_internal_knots: usize,
1316        basis_options: BasisOptions,
1317    ) -> Result<Array1<f64>, String> {
1318        infer_survival_time_knots_with_degree(
1319            combined,
1320            knot_degree,
1321            validation_degree,
1322            num_internal_knots,
1323            basis_options,
1324        )
1325        .map(|(knots, _)| knots)
1326    }
1327
1328    match cfg {
1329        SurvivalTimeBasisConfig::None => Ok(SurvivalTimeBuildOutput {
1330            x_entry_time: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((n, 0)))),
1331            x_exit_time: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((n, 0)))),
1332            x_derivative_time: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((n, 0)))),
1333            penalties: Vec::new(),
1334            nullspace_dims: Vec::new(),
1335            basisname: "none".to_string(),
1336            degree: None,
1337            knots: None,
1338            keep_cols: None,
1339            smooth_lambda: None,
1340        }),
1341        SurvivalTimeBasisConfig::Linear => {
1342            // Single column `log t` — the Weibull baseline slope (shape). The
1343            // constant column `[1, ·]` this basis used to carry (#2301) is the
1344            // `−shape·log_scale` location, which is EXACTLY confounded with the
1345            // linear-predictor intercept, so it made the converged penalized
1346            // Hessian singular (the anchor gauge — the killed EDF trace solve and
1347            // the LM crawl were both downstream of that singularity). Dropping it
1348            // moves the whole location into the mean intercept and leaves `H`
1349            // nonsingular. This is valid ONLY because intercept removal (`~ x - 1`)
1350            // is a typed refusal at `formula_dsl.rs:2456` — the covariate block
1351            // ALWAYS carries an intercept to absorb the location. If intercept
1352            // suppression is ever implemented, the Weibull location becomes
1353            // unidentified without this column and the two features MUST be
1354            // reconciled here (re-add the constant and pin it, or keep the ban).
1355            let mut x_entry_time = Array2::<f64>::zeros((n, 1));
1356            let mut x_exit_time = Array2::<f64>::zeros((n, 1));
1357            let mut x_derivative_time = Array2::<f64>::zeros((n, 1));
1358            for i in 0..n {
1359                x_entry_time[[i, 0]] = log_entry[i];
1360                x_exit_time[[i, 0]] = log_exit[i];
1361                x_derivative_time[[i, 0]] = 1.0 / age_exit[i].max(SURVIVAL_TIME_FLOOR);
1362            }
1363            Ok(SurvivalTimeBuildOutput {
1364                x_entry_time: DesignMatrix::Dense(DenseDesignMatrix::from(x_entry_time)),
1365                x_exit_time: DesignMatrix::Dense(DenseDesignMatrix::from(x_exit_time)),
1366                x_derivative_time: DesignMatrix::Dense(DenseDesignMatrix::from(x_derivative_time)),
1367                penalties: Vec::new(),
1368                nullspace_dims: Vec::new(),
1369                basisname: "linear".to_string(),
1370                degree: None,
1371                knots: None,
1372                keep_cols: None,
1373                smooth_lambda: None,
1374            })
1375        }
1376        SurvivalTimeBasisConfig::BSpline {
1377            degree,
1378            knots,
1379            smooth_lambda,
1380        } => {
1381            let knotvec = if knots.is_empty() {
1382                let (num_internal_knots, _) = infer_knots_if_needed.ok_or_else(|| {
1383                    "internal error: bspline time basis requested without knot source".to_string()
1384                })?;
1385                let combined = survival_time_knot_input(&log_entry, &log_exit);
1386                infer_survival_time_knots(
1387                    &combined,
1388                    degree,
1389                    degree,
1390                    num_internal_knots,
1391                    BasisOptions::value(),
1392                )?
1393            } else {
1394                knots
1395            };
1396
1397            let entry_basis = build_bspline_basis_1d(
1398                log_entry.view(),
1399                &BSplineBasisSpec {
1400                    degree,
1401                    penalty_order: 2,
1402                    knotspec: BSplineKnotSpec::Provided(knotvec.clone()),
1403                    double_penalty: false,
1404                    identifiability: BSplineIdentifiability::None,
1405                    boundary: OneDimensionalBoundary::Open,
1406                    boundary_conditions: BSplineBoundaryConditions::default(),
1407                },
1408            )
1409            .map_err(|e| format!("failed to build bspline entry basis: {e}"))?;
1410            let exit_basis = build_bspline_basis_1d(
1411                log_exit.view(),
1412                &BSplineBasisSpec {
1413                    degree,
1414                    penalty_order: 2,
1415                    knotspec: BSplineKnotSpec::Provided(knotvec.clone()),
1416                    double_penalty: false,
1417                    identifiability: BSplineIdentifiability::None,
1418                    boundary: OneDimensionalBoundary::Open,
1419                    boundary_conditions: BSplineBoundaryConditions::default(),
1420                },
1421            )
1422            .map_err(|e| format!("failed to build bspline exit basis: {e}"))?;
1423
1424            let p_time = exit_basis.design.ncols();
1425            // Build derivative basis as sparse triplets — B-spline derivatives
1426            // have the same local support as the basis itself (at most degree+1
1427            // nonzeros per row), so building dense first wastes memory.
1428            let mut deriv_triplets = Vec::with_capacity(n * (degree + 1));
1429            let mut deriv_buf = vec![0.0_f64; p_time];
1430            for i in 0..n {
1431                deriv_buf.fill(0.0);
1432                evaluate_bspline_derivative_scalar(
1433                    log_exit[i],
1434                    knotvec.view(),
1435                    degree,
1436                    &mut deriv_buf,
1437                )
1438                .map_err(|e| format!("failed to evaluate bspline derivative: {e}"))?;
1439                let chain = 1.0 / age_exit[i].max(SURVIVAL_TIME_FLOOR);
1440                for j in 0..p_time {
1441                    let v = deriv_buf[j] * chain;
1442                    if v.abs() > 1e-15 {
1443                        deriv_triplets.push(faer::sparse::Triplet::new(i, j, v));
1444                    }
1445                }
1446            }
1447            let x_derivative_time =
1448                match faer::sparse::SparseColMat::try_new_from_triplets(n, p_time, &deriv_triplets)
1449                {
1450                    Ok(sparse) => DesignMatrix::Sparse(SparseDesignMatrix::new(sparse)),
1451                    Err(_) => {
1452                        // Fallback: build dense
1453                        let mut dense = Array2::<f64>::zeros((n, p_time));
1454                        for &faer::sparse::Triplet { row, col, val } in &deriv_triplets {
1455                            dense[[row, col]] = val;
1456                        }
1457                        DesignMatrix::Dense(DenseDesignMatrix::from(dense))
1458                    }
1459                };
1460
1461            let nullspace_dims = entry_basis
1462                .active_penalties
1463                .iter()
1464                .map(|penalty| penalty.nullity)
1465                .collect();
1466            let penalties = entry_basis
1467                .active_penalties
1468                .into_iter()
1469                .map(|penalty| penalty.matrix)
1470                .collect();
1471
1472            Ok(SurvivalTimeBuildOutput {
1473                x_entry_time: entry_basis.design,
1474                x_exit_time: exit_basis.design,
1475                x_derivative_time,
1476                nullspace_dims,
1477                penalties,
1478                basisname: "bspline".to_string(),
1479                degree: Some(degree),
1480                knots: Some(knotvec.to_vec()),
1481                keep_cols: None,
1482                smooth_lambda: Some(smooth_lambda),
1483            })
1484        }
1485        SurvivalTimeBasisConfig::ISpline {
1486            degree,
1487            knots,
1488            keep_cols,
1489            smooth_lambda,
1490        } => {
1491            let requested_bspline_degree = degree
1492                .checked_add(1)
1493                .ok_or_else(|| "ispline degree overflow while building knot basis".to_string())?;
1494            // Every consumer below -- the derivative basis at `bspline_degree`
1495            // and both i-spline bases at `degree` -- is sized from these two
1496            // numbers, so they must describe the vector we HAVE rather than the
1497            // one we asked for. Inference can shrink the degree on data too
1498            // sparse to carry it (4 rows cannot support degree 4), and the
1499            // shrunk vector was previously handed to a degree-4 consumer:
1500            //
1501            //   Insufficient knots for degree 4 spline: need at least 10 knots
1502            //   but only 9 were provided.
1503            //
1504            // An explicit knot vector is the user's own geometry and is never
1505            // re-derived, so it keeps the requested degrees.
1506            let (knotvec, degree, bspline_degree) = if knots.is_empty() {
1507                let (num_internal_knots, _) = infer_knots_if_needed.ok_or_else(|| {
1508                    "internal error: ispline time basis requested without knot source".to_string()
1509                })?;
1510                let combined = survival_time_knot_input(&log_entry, &log_exit);
1511                let (knotvec, effective_degree) = infer_survival_time_knots_with_degree(
1512                    &combined,
1513                    requested_bspline_degree,
1514                    degree,
1515                    num_internal_knots,
1516                    BasisOptions::i_spline(),
1517                )?;
1518                let effective_bspline_degree =
1519                    effective_degree.checked_add(1).ok_or_else(|| {
1520                        "ispline degree overflow while building knot basis".to_string()
1521                    })?;
1522                (knotvec, effective_degree, effective_bspline_degree)
1523            } else {
1524                (knots, degree, requested_bspline_degree)
1525            };
1526
1527            // ONE boundary convention for the baseline value AND its slope
1528            // (gam#2705).
1529            //
1530            // The Royston-Parmar baseline is `log Λ(t) = Σ_k γ_k·I_k(log t)`,
1531            // and the likelihood, the hazard and the predictive surface all read
1532            // BOTH `I_k` and `I'_k = M_k`. Those two used to be built by
1533            // different code paths here: the value by the shared I-spline
1534            // evaluator, which holds `I_k` CONSTANT past the boundary knots, and
1535            // the slope by a hand-rolled right-cumulative sum of a *clamped*
1536            // B-spline first-derivative basis, which returns the BOUNDARY SLOPE
1537            // there because a clamped B-spline's value extends linearly. So
1538            // outside the fitted knot span the two described different
1539            // functions, and a saved fit published a FLAT `Λ(t)` next to a
1540            // NONZERO `h(t) = Λ·d(log Λ)/dt` — measured on the #1564 heart-failure
1541            // fixture as `Λ ≡ 5.055558` with `t·h(t) ≡ 6.26088` from `t = 285`
1542            // out to `t = 2.85e6`, i.e. a surviving log-log slope of `1.23842`
1543            // in the derivative that the value does not have. `h = dΛ/dt`, so a
1544            // flat `Λ` forces `h = 0`; the two cannot both be the model.
1545            //
1546            // The convention is `LinearTails` rather than `Saturate` because
1547            // that IS the Royston-Parmar model: a *restricted* spline is linear
1548            // beyond its boundary knots by construction (Royston & Parmar 2002),
1549            // which gives the classical Weibull-shaped extrapolation
1550            // `Λ(t) ∝ t^c` used whenever a survival curve is projected past the
1551            // observed follow-up. Saturating instead asserts two things the data
1552            // never said: that the hazard drops to exactly zero at the last
1553            // observed exit time, and — on the lower tail, which
1554            // `default_survival_time_grid` reaches on its very first node —
1555            // that `Λ(t) → Λ(t_min) > 0` as `t → 0`, i.e. an atom of failures at
1556            // time zero and `S(0) < 1`.
1557            //
1558            // Nothing about a FIT moves: the knot vector is inferred from
1559            // `survival_time_knot_input(log_entry, log_exit)`, so every training
1560            // row is inside `[left, right]` where the two conventions are
1561            // bit-identical, `keep_cols` is inferred from those same interior
1562            // rows, and the penalty is built on `log_exit`. What moves is
1563            // evaluation OUTSIDE the fitted span — prediction grids, entry times
1564            // below the first knot, and any replay at a fresh time.
1565            let (x_exit_full, d_exit_log_full) = ispline_value_and_first_derivative(
1566                log_exit.view(),
1567                knotvec.view(),
1568                degree,
1569                ISplineBoundary::LinearTails,
1570            )
1571            .map_err(|e| format!("failed to build ispline exit basis and derivative: {e}"))?;
1572            // A row that ENTERS AT THE ORIGIN is not a row whose entry time
1573            // sits below the first knot — it is a row with no left truncation
1574            // at all, and the likelihood says so: `entry_active` is
1575            // `age_entry > ENTRY_AT_ORIGIN_THRESHOLD`, and the `S(entry)` factor
1576            // is dropped outright for the rest (`survival/base.rs`). Its entry
1577            // design row is therefore never read, and what it holds is a
1578            // conditioning choice rather than a model statement.
1579            //
1580            // `log_entry` for such a row is `ln(SURVIVAL_TIME_FLOOR) = −20.7`,
1581            // a NUMERICAL FLOOR and not a datum, so a linear tail evaluated
1582            // there would put a large arbitrary constant into a design column
1583            // (and into the column-scaling statistics computed from it) purely
1584            // as a readout of `1e-9`. The saturating basis got the right answer
1585            // here for the wrong reason: every time at or below the first knot
1586            // maps to the anchored ZERO row (`I_k(left) = 0` exactly). Keep that
1587            // answer, and keep it for the reason that holds — the likelihood's
1588            // own origin predicate — so a genuine delayed entry below the first
1589            // knot still receives the real extrapolation.
1590            let interval = ispline_modelling_interval(knotvec.view(), degree)
1591                .map_err(|e| format!("failed to resolve ispline modelling interval: {e}"))?;
1592            let mut log_entry_for_basis = log_entry.clone();
1593            if let Some((left, _right)) = interval {
1594                for i in 0..n {
1595                    if age_entry[i] <= crate::survival::base::ENTRY_AT_ORIGIN_THRESHOLD {
1596                        log_entry_for_basis[i] = left;
1597                    }
1598                }
1599            }
1600            let x_entry_full = ispline_value(
1601                log_entry_for_basis.view(),
1602                knotvec.view(),
1603                degree,
1604                ISplineBoundary::LinearTails,
1605            )
1606            .map_err(|e| format!("failed to build ispline entry basis: {e}"))?;
1607
1608            let (x_entry_time, x_exit_time, keep_cols, p_time, p_time_full) = {
1609                let p_time_full = x_exit_full.ncols();
1610                if p_time_full == 0 {
1611                    return Err(SurvivalConstructionError::BasisConstructionFailed {
1612                        reason: "internal error: empty ispline time basis".to_string(),
1613                    }
1614                    .into());
1615                }
1616                if d_exit_log_full.ncols() != p_time_full
1617                    || d_exit_log_full.nrows() != x_exit_full.nrows()
1618                {
1619                    return Err(format!(
1620                        "internal error: ispline time derivative basis is {:?} but its value basis \
1621                         is {:?}",
1622                        d_exit_log_full.dim(),
1623                        x_exit_full.dim()
1624                    ));
1625                }
1626
1627                let keep_cols = if keep_cols.is_empty() {
1628                    let constant_tol = 1e-12_f64;
1629                    let mut inferred_keep_cols: Vec<usize> = Vec::new();
1630                    for j in 0..p_time_full {
1631                        let mut minv = f64::INFINITY;
1632                        let mut maxv = f64::NEG_INFINITY;
1633                        for i in 0..n {
1634                            let ve = x_exit_full[[i, j]];
1635                            let vs = x_entry_full[[i, j]];
1636                            minv = minv.min(ve.min(vs));
1637                            maxv = maxv.max(ve.max(vs));
1638                        }
1639                        if (maxv - minv) > constant_tol {
1640                            inferred_keep_cols.push(j);
1641                        }
1642                    }
1643                    inferred_keep_cols
1644                } else {
1645                    keep_cols
1646                };
1647                if keep_cols.is_empty() {
1648                    return Err(
1649                        "internal error: ispline basis has no shape-varying time columns"
1650                            .to_string(),
1651                    );
1652                }
1653                if keep_cols.iter().any(|&j| j >= p_time_full) {
1654                    return Err(SurvivalConstructionError::MissingColumn {
1655                        reason: "saved survival ispline keep_cols exceed basis width".to_string(),
1656                    }
1657                    .into());
1658                }
1659
1660                let p_time = keep_cols.len();
1661                let x_entry_time = x_entry_full.select(ndarray::Axis(1), &keep_cols);
1662                let x_exit_time = x_exit_full.select(ndarray::Axis(1), &keep_cols);
1663                (x_entry_time, x_exit_time, keep_cols, p_time, p_time_full)
1664            };
1665            // The full-width VALUE bases are no longer needed; the retained
1666            // blocks above own their own storage. The full-width derivative is
1667            // still read below, one row at a time, so it stays.
1668            drop(x_entry_full);
1669            drop(x_exit_full);
1670
1671            // `d(log Λ)/dt = d(log Λ)/d(log t) · 1/t`. The `d/d(log t)` half is
1672            // the M-spline block the value basis was built with, so no second
1673            // opinion about the exterior can arise here.
1674            let mut deriv_triplets = Vec::with_capacity(n * p_time.min(16));
1675            let mut found_nonfinite: Option<(usize, usize)> = None;
1676            for i in 0..n {
1677                let chain = 1.0 / age_exit[i].max(SURVIVAL_TIME_FLOOR);
1678                for (j_new, &j_old) in keep_cols.iter().enumerate() {
1679                    let raw_v = d_exit_log_full[[i, j_old]] * chain;
1680                    let v = if (-1e-12..0.0).contains(&raw_v) {
1681                        0.0
1682                    } else {
1683                        raw_v
1684                    };
1685                    if !v.is_finite() {
1686                        found_nonfinite = Some((i, j_new));
1687                    }
1688                    if v < -1e-12 {
1689                        return Err(format!(
1690                            "survival ispline derivative basis must stay non-negative at row {}, column {}; found {:.3e}",
1691                            i + 1,
1692                            j_new + 1,
1693                            v
1694                        ));
1695                    }
1696                    if v.abs() > 1e-15 {
1697                        deriv_triplets.push(faer::sparse::Triplet::new(i, j_new, v));
1698                    }
1699                }
1700            }
1701            if let Some((row, col)) = found_nonfinite {
1702                return Err(format!(
1703                    "survival ispline derivative basis produced non-finite value at row {}, column {}",
1704                    row + 1,
1705                    col + 1
1706                ));
1707            }
1708            let x_derivative_time =
1709                match faer::sparse::SparseColMat::try_new_from_triplets(n, p_time, &deriv_triplets)
1710                {
1711                    Ok(sparse) => DesignMatrix::Sparse(SparseDesignMatrix::new(sparse)),
1712                    Err(_) => {
1713                        let mut dense = Array2::<f64>::zeros((n, p_time));
1714                        for &faer::sparse::Triplet { row, col, val } in &deriv_triplets {
1715                            dense[[row, col]] = val;
1716                        }
1717                        DesignMatrix::Dense(DenseDesignMatrix::from(dense))
1718                    }
1719                };
1720
1721            let penalty_basis = build_bspline_basis_1d(
1722                log_exit.view(),
1723                &BSplineBasisSpec {
1724                    degree: bspline_degree,
1725                    penalty_order: 2,
1726                    knotspec: BSplineKnotSpec::Provided(knotvec.clone()),
1727                    double_penalty: false,
1728                    identifiability: BSplineIdentifiability::None,
1729                    boundary: OneDimensionalBoundary::Open,
1730                    boundary_conditions: BSplineBoundaryConditions::default(),
1731                },
1732            )
1733            .map_err(|e| format!("failed to build ispline smoothing penalty: {e}"))?;
1734            if penalty_basis.design.ncols() != p_time_full + 1 {
1735                return Err("internal error: ispline penalty dimension mismatch".to_string());
1736            }
1737            // I-spline curvature penalty in the *value* space of the baseline
1738            // log-cumulative-hazard, restricted to the retained (non-dropped)
1739            // coefficient block.
1740            //
1741            // The I-spline coefficient γ is the consecutive increment of the B-spline
1742            // value coefficients `c`: `c_0 = 0`, `c_k = Σ_{j<k} γ_j = (L γ)_k`, where
1743            // `L` is the `p_time × p_time` lower-triangular cumsum matrix. The
1744            // second-difference penalty on the B-spline values is `S_B = D₂ᵀD₂`
1745            // (the active `penalty_basis` matrix block). The correct curvature penalty
1746            // on γ is the **value-space congruence transform**
1747            //
1748            //   `S_I = Lᵀ S_B[1:,1:] L`,
1749            //
1750            // which satisfies `γᵀ S_I γ = (Lγ)ᵀ S_B[1:,1:] (Lγ)`.
1751            //
1752            // A constant γ (γ_k = γ₀ ∀k) maps to the linear value sequence
1753            // `c_k = k·γ₀`, which is annihilated by D₂: `D₂c = 0`. Therefore
1754            // `γᵀ S_I γ = 0` for constant γ, i.e. the **affine trend lies in the
1755            // penalty null space**. REML does not penalize the baseline slope
1756            // `d(log Λ)/d(log t)` or the overall level, so it correctly lets the
1757            // data determine these quantities without bias. The previous increment-
1758            // space form `S_B[1:,1:]` (applied directly to γ instead of Lγ) did NOT
1759            // have constant γ in its null space and therefore over-penalized affine
1760            // baselines, causing the fitted log-cumulative-hazard to lose its tail
1761            // slope to the penalty and fail quality tests (#1076).
1762            //
1763            // The value-space form has a 1-dimensional null space (span{(1,…,1)}),
1764            // declared via `nullspace_dims` so the REML generalized-logdet picks it
1765            // up. The penalized inner PIRLS is well-conditioned because the
1766            // likelihood Hessian H_lik has O(n_events) curvature along the affine
1767            // direction (the overall baseline level is identified by the data), and
1768            // the global stabilization ridge (ridge_lambda) provides an absolute
1769            // positive-definite floor.
1770            let mut penalties = Vec::<Array2<f64>>::new();
1771            for active_penalty in &penalty_basis.active_penalties {
1772                let s_mat = &active_penalty.matrix;
1773                if s_mat.nrows() != p_time_full + 1 || s_mat.ncols() != p_time_full + 1 {
1774                    continue;
1775                }
1776                // I-spline value-space penalty, computed in the CORRECT order
1777                // (gam#979). The B-spline value coefficients are the cumulative
1778                // sum of the I-spline increment coefficients, `c = L γ_full`, where
1779                // `L` is the FULL `p_time_full × p_time_full` LOWER-triangular
1780                // all-ones cumsum matrix (`L[i,j] = 1 iff j ≤ i`, so
1781                // `c_i = Σ_{j≤i} γ_j`). The value-space curvature penalty on the
1782                // full increment vector is the symmetric congruence
1783                //
1784                //   `S_I_full = Lᵀ · S_B[1:,1:] · L`,
1785                //
1786                // which is PSD because `S_B[1:,1:]` is a principal submatrix of the
1787                // PSD `S_B = D₂ᵀD₂` and congruence by any matrix preserves PSD.
1788                //
1789                // CRITICAL ORDERING (the gam#979 indefiniteness bug): the retained
1790                // columns `keep_cols` must be selected as a PRINCIPAL SUBMATRIX of
1791                // the FULL congruence `S_I_full` — i.e. congruence FIRST, selection
1792                // SECOND. The previous code selected `keep_cols` from `S_B[1:,1:]`
1793                // first and then applied a `p_time × p_time` cumsum to that
1794                // already-reduced block. Because the cumsum `L` couples every
1795                // increment, restricting the increment index set BEFORE the cumsum
1796                // does NOT commute with it: the reduced operator is a different,
1797                // generally INDEFINITE matrix (measured `s0_min_eval = −9.8e7`),
1798                // which makes `½γᵀS_Iγ` unbounded below and the penalized survival
1799                // NLL diverge (β drifts up the negative-eigenvalue mode, the inner
1800                // joint-Newton follows the unbounded objective, the outer REML never
1801                // terminates — the #979 hang). Doing the congruence on the full γ
1802                // and then taking the `keep_cols` principal submatrix restores the
1803                // PSD guarantee (a principal submatrix of a PSD matrix is PSD).
1804                let s_increment = s_mat.slice(s![1.., 1..]);
1805                if s_increment.nrows() != p_time_full || s_increment.ncols() != p_time_full {
1806                    return Err(format!(
1807                        "internal error: ispline penalty increment block must be {p_time_full}x{p_time_full}, got {}x{}",
1808                        s_increment.nrows(),
1809                        s_increment.ncols(),
1810                    ));
1811                }
1812                // Symmetrize the (already-symmetric) source with the shared
1813                // matrix utility. The survival builder's value-space
1814                // congruence is domain-specific; only the low-level symmetric
1815                // cleanup is common with the generic and SAE construction code.
1816                let mut s_full = s_increment.to_owned();
1817                symmetrize_in_place(&mut s_full);
1818                // S_mid = S_B[1:,1:] · L  (right-multiply by lower-triangular
1819                // cumsum): (S·L)[i,j] = Σ_k S[i,k]·L[k,j] = Σ_{k≥j} S[i,k]
1820                // because L[k,j] = 1 iff j ≤ k.
1821                let mut s_mid_full = Array2::<f64>::zeros((p_time_full, p_time_full));
1822                for i in 0..p_time_full {
1823                    for j in 0..p_time_full {
1824                        let mut v = 0.0;
1825                        for k in j..p_time_full {
1826                            v += s_full[[i, k]];
1827                        }
1828                        s_mid_full[[i, j]] = v;
1829                    }
1830                }
1831                // S_I_full = Lᵀ · S_mid = Lᵀ · S · L:
1832                // (Lᵀ·S_mid)[i,j] = Σ_k Lᵀ[i,k]·S_mid[k,j] = Σ_{k≥i} S_mid[k,j]
1833                // because Lᵀ[i,k] = L[k,i] = 1 iff i ≤ k.
1834                let mut s_full_congruent = Array2::<f64>::zeros((p_time_full, p_time_full));
1835                for i in 0..p_time_full {
1836                    for j in 0..p_time_full {
1837                        let mut v = 0.0;
1838                        for k in i..p_time_full {
1839                            v += s_mid_full[[k, j]];
1840                        }
1841                        s_full_congruent[[i, j]] = v;
1842                    }
1843                }
1844                // Principal submatrix on the retained (shape-varying) columns.
1845                let mut local = Array2::<f64>::zeros((p_time, p_time));
1846                for (i_new, &i_old) in keep_cols.iter().enumerate() {
1847                    for (j_new, &j_old) in keep_cols.iter().enumerate() {
1848                        // Symmetrize on the way out to absorb residual
1849                        // floating-point asymmetry.
1850                        local[[i_new, j_new]] = 0.5
1851                            * (s_full_congruent[[i_old, j_old]] + s_full_congruent[[j_old, i_old]]);
1852                    }
1853                }
1854                penalties.push(local);
1855            }
1856
1857            // PSD contract (gam#979). The value-space congruence Lᵀ S_B[1:,1:] L,
1858            // restricted to a principal submatrix, is positive semidefinite by
1859            // construction. A negative eigenvalue here means the construction has
1860            // regressed to the increment-space / wrong-ordering form that made the
1861            // penalized survival NLL unbounded below (the #979 divergence). Verify
1862            // it here, at construction, so the defect can never silently reach the
1863            // inner solver again. The tolerance is the same relative scale the
1864            // nullspace detection below uses; a numerically tiny negative (round-off
1865            // on the genuine 1-D null direction) is allowed, a structural one is not.
1866            for (idx, s_mat) in penalties.iter().enumerate() {
1867                let p = s_mat.nrows();
1868                if p == 0 {
1869                    continue;
1870                }
1871                if let Ok((evals, _)) =
1872                    gam_linalg::faer_ndarray::FaerEigh::eigh(s_mat, faer::Side::Lower)
1873                {
1874                    let evals_slice: &[f64] = evals.as_slice().ok_or_else(|| {
1875                        "internal error: ispline penalty eigenvalues not contiguous".to_string()
1876                    })?;
1877                    let max_ev = evals_slice
1878                        .iter()
1879                        .copied()
1880                        .fold(0.0_f64, |a, b| a.max(b.abs()))
1881                        .max(1.0);
1882                    let min_ev = evals_slice.iter().copied().fold(f64::INFINITY, f64::min);
1883                    let neg_tol = -100.0 * (p as f64) * f64::EPSILON * max_ev;
1884                    if min_ev < neg_tol {
1885                        return Err(format!(
1886                            "internal error (gam#979): assembled ispline time-block penalty {idx} is \
1887                             indefinite (min eigenvalue {min_ev:.3e} < tol {neg_tol:.3e}, max |eig| \
1888                             {max_ev:.3e}); the value-space congruence Lᵀ S_B[1:,1:] L must be PSD"
1889                        ));
1890                    }
1891                }
1892            }
1893
1894            // The value-space penalty S_I = L^T S_B[1:,1:] L has a 1-dimensional
1895            // null space (constant γ ↦ affine c ↦ D₂c = 0). Detect it spectrally
1896            // so the REML uses the generalized logdet over the penalized subspace.
1897            let nullspace_dims: Vec<usize> = penalties
1898                .iter()
1899                .map(|s_mat| {
1900                    let p = s_mat.nrows();
1901                    if p == 0 {
1902                        return 0;
1903                    }
1904                    match gam_linalg::faer_ndarray::FaerEigh::eigh(s_mat, faer::Side::Lower) {
1905                        Ok((evals, _)) => {
1906                            let max_ev = evals
1907                                .iter()
1908                                .copied()
1909                                .fold(0.0_f64, |a, b| a.max(b.abs()))
1910                                .max(1.0);
1911                            let threshold = 100.0 * (p as f64) * f64::EPSILON * max_ev;
1912                            evals.iter().filter(|&&e| e <= threshold).count()
1913                        }
1914                        Err(_) => 0,
1915                    }
1916                })
1917                .collect();
1918            Ok(SurvivalTimeBuildOutput {
1919                x_entry_time: DesignMatrix::Dense(DenseDesignMatrix::from(x_entry_time)),
1920                x_exit_time: DesignMatrix::Dense(DenseDesignMatrix::from(x_exit_time)),
1921                x_derivative_time,
1922                penalties,
1923                nullspace_dims,
1924                basisname: "ispline".to_string(),
1925                degree: Some(degree),
1926                knots: Some(knotvec.to_vec()),
1927                keep_cols: Some(keep_cols),
1928                smooth_lambda: Some(smooth_lambda),
1929            })
1930        }
1931    }
1932}
1933
1934pub fn resolved_survival_time_basis_config_from_build(
1935    basisname: &str,
1936    degree: Option<usize>,
1937    knots: Option<&Vec<f64>>,
1938    keep_cols: Option<&Vec<usize>>,
1939    smooth_lambda: Option<f64>,
1940) -> Result<SurvivalTimeBasisConfig, String> {
1941    match basisname {
1942        "none" => Ok(SurvivalTimeBasisConfig::None),
1943        "linear" => Ok(SurvivalTimeBasisConfig::Linear),
1944        "bspline" => Ok(SurvivalTimeBasisConfig::BSpline {
1945            degree: degree.ok_or_else(|| "survival bspline basis is missing degree".to_string())?,
1946            knots: Array1::from_vec(
1947                knots
1948                    .cloned()
1949                    .ok_or_else(|| "survival bspline basis is missing knots".to_string())?,
1950            ),
1951            smooth_lambda: smooth_lambda.unwrap_or(SURVIVAL_TIME_SMOOTH_LAMBDA_SEED),
1952        }),
1953        "ispline" => Ok(SurvivalTimeBasisConfig::ISpline {
1954            degree: degree.ok_or_else(|| "survival ispline basis is missing degree".to_string())?,
1955            knots: Array1::from_vec(
1956                knots
1957                    .cloned()
1958                    .ok_or_else(|| "survival ispline basis is missing knots".to_string())?,
1959            ),
1960            keep_cols: keep_cols
1961                .cloned()
1962                .ok_or_else(|| "survival ispline basis is missing keep_cols".to_string())?,
1963            smooth_lambda: smooth_lambda.unwrap_or(SURVIVAL_TIME_SMOOTH_LAMBDA_SEED),
1964        }),
1965        other => Err(format!("unsupported survival time basis '{other}'")),
1966    }
1967}
1968
1969// ---------------------------------------------------------------------------
1970// Survival time-basis anchor: ONE rule, three primitives
1971// ---------------------------------------------------------------------------
1972//
1973// `center_survival_time_designs_at_anchor` subtracts the time-basis row at the
1974// anchor from every entry/exit design row, so the anchor sets the origin of the
1975// baseline-hazard reparameterization. WHICH time to anchor at is a function of
1976// exactly three things — the likelihood mode, the entry/exit data, and an
1977// optional caller override — so it is decided in exactly one place,
1978// [`resolve_survival_time_anchor_for_mode`], which every front end calls.
1979//
1980// It used to be decided in two. `materialize_survival` (the engine path behind
1981// `fit_from_formula` and the Python FFI) promoted the robust anchor for any
1982// left-truncated dataset and hardcoded the override to `None`; `gam-cli`'s
1983// survival path promoted it only for marginal-slope and owned the
1984// `--survival-time-anchor` override. Three consequences, all of them #2631:
1985//
1986//   1. The same formula, data and config produced a DIFFERENT fit depending on
1987//      which front end ran it — a left-truncated location-scale (or latent)
1988//      model centered at the robust median exit under `fit_from_formula` and at
1989//      the earliest entry under the CLI.
1990//   2. Because the override lived only in the CLI copy, and the CLI's own
1991//      default (transformation / Weibull) route delegates to the engine copy,
1992//      `--survival-time-anchor` was SILENTLY IGNORED on the default route.
1993//   3. `FitRequestConfigDocument` — the "complete scientific model
1994//      configuration" that `--survival-time-anchor` declares a conflict with —
1995//      had no field for the anchor at all, so a fit-request document could not
1996//      express what the flag it excludes expresses.
1997//
1998// A rule that lives in two places is a rule that disagrees with itself. The
1999// anchor is model configuration, not front-end transport, so the override now
2000// travels on `FitConfig` and the rule below is the only thing that reads it.
2001
2002/// Validate a caller-supplied survival time-anchor override.
2003///
2004/// Honored verbatim by every likelihood mode (subject to the `SURVIVAL_TIME_FLOOR`
2005/// clamp that keeps `log(anchor)` finite), because a caller who names the anchor
2006/// is overriding the conditioning heuristic on purpose.
2007pub fn validate_survival_time_anchor_override(time_anchor: f64) -> Result<f64, String> {
2008    if !time_anchor.is_finite() || time_anchor < 0.0 {
2009        return Err(format!(
2010            "survival time anchor must be finite and non-negative, got {time_anchor}"
2011        ));
2012    }
2013    Ok(time_anchor.max(SURVIVAL_TIME_FLOOR))
2014}
2015
2016/// Earliest-entry anchor — the default for data that is NOT left-truncated.
2017///
2018/// With every row entering at the time origin this is `≈ 0`, so for the
2019/// monotone I-spline time basis the anchor row is `≈ 0` and centering is a
2020/// near-no-op: the historical (pre-#751) behavior is preserved bit-for-bit on
2021/// ordinary right-censored data.
2022pub fn survival_earliest_entry_time_anchor(age_entry: &Array1<f64>) -> Result<f64, String> {
2023    let min_entry = age_entry
2024        .iter()
2025        .copied()
2026        .min_by(f64::total_cmp)
2027        .ok_or_else(|| "survival time anchor requires non-empty entry times".to_string())?;
2028    Ok(min_entry.max(SURVIVAL_TIME_FLOOR))
2029}
2030
2031/// Robust interior anchor — the median exit age, a time on the **exit** scale
2032/// where the at-risk mass concentrates.
2033///
2034/// Under left truncation the earliest entry age is a genuine positive
2035/// *left-tail* point, and centering the time basis there leaves the centered
2036/// linear-trend column `X(exit) − X(anchor)` large and one-signed across all
2037/// rows (every exit sits far to the right of the earliest entry). That column is
2038/// the unpenalized polynomial null space of the 2nd-difference time penalty, so
2039/// the inflated one-signed column multiplies the time-block score at the
2040/// smoothing seed up by hundreds: the marginal-slope constrained joint Newton
2041/// cannot certify KKT on it and REML rejects every seed (#751), and the
2042/// transformation (Royston-Parmar) smoothing selection rails a penalty
2043/// direction and collapses the baseline to a covariate-independent surface with
2044/// `H` inflated ~10³× and `S(t) ≡ 0` (#1790).
2045///
2046/// Centering at the median exit keeps that column small and two-signed (some
2047/// exits below the median, some above) so the exit-event likelihood pins the
2048/// linear trend and the seed score stays bounded. The median is chosen over the
2049/// mean for robustness to the heavy right tail of survival times.
2050pub fn survival_robust_interior_time_anchor(age_exit: &Array1<f64>) -> Result<f64, String> {
2051    if age_exit.is_empty() {
2052        return Err(
2053            "survival robust interior time anchor requires non-empty exit times".to_string(),
2054        );
2055    }
2056    let mut sorted: Vec<f64> = age_exit.iter().copied().collect();
2057    sorted.sort_by(f64::total_cmp);
2058    let m = sorted.len();
2059    let median = if m % 2 == 1 {
2060        sorted[m / 2]
2061    } else {
2062        0.5 * (sorted[m / 2 - 1] + sorted[m / 2])
2063    };
2064    Ok(median.max(SURVIVAL_TIME_FLOOR))
2065}
2066
2067/// The single definition of "this dataset is genuinely left-truncated".
2068///
2069/// **Any** row entering above `ENTRY_AT_ORIGIN_THRESHOLD` makes the data
2070/// left-truncated, not just the earliest one. Staggered entry — part of the
2071/// cohort observed from the time origin, the rest joining at positive delayed
2072/// entry times — is the ordinary shape of a real registry cohort, and it
2073/// exhibits the #751/#1790 inflation just as fully-delayed entry does: the
2074/// earliest-entry anchor is then `≈ 0`, the anchor row of a `log t` basis is
2075/// evaluated at `SURVIVAL_TIME_FLOOR`, and every centered exit column is
2076/// one-signed and large. Testing `min(entry) > threshold` instead would
2077/// under-trigger on exactly that shape, which is why the two former copies of
2078/// this predicate (`any` in the materializer, `min` inside the transformation
2079/// resolver) had to be collapsed to one.
2080///
2081/// The threshold is the likelihood engines' own origin convention, so "this row
2082/// has a delayed-entry interval" and "this dataset is left-truncated" cannot
2083/// drift apart.
2084pub fn survival_data_is_left_truncated(age_entry: &Array1<f64>) -> bool {
2085    age_entry
2086        .iter()
2087        .any(|&entry| entry > crate::survival::base::ENTRY_AT_ORIGIN_THRESHOLD)
2088}
2089
2090/// **The** survival time-basis anchor rule. Every front end calls this and
2091/// nothing else.
2092///
2093/// * An explicit `time_anchor` wins, in every mode.
2094/// * Marginal-slope always takes the robust interior anchor: its `γ = 0`
2095///   monotone-cone seed is where #751 was measured, and the fix was applied
2096///   there unconditionally.
2097/// * Every other time-basis-carrying likelihood takes the robust interior
2098///   anchor **iff the data is genuinely left-truncated**. Ordinary
2099///   right-censored data keeps the earliest-entry anchor, which is `≈` the time
2100///   origin, so centering stays a near-no-op and pre-#751 behavior is preserved
2101///   bit-for-bit.
2102///
2103/// Re-centering is an exact affine reparameterization of the baseline offset, so
2104/// this choice does not change the model being fitted — only the frame the
2105/// smoothing selection sees it in, and the frame the saved
2106/// `survival_time_anchor` must replay in.
2107pub fn resolve_survival_time_anchor_for_mode(
2108    survival_mode: SurvivalLikelihoodMode,
2109    age_entry: &Array1<f64>,
2110    age_exit: &Array1<f64>,
2111    time_anchor: Option<f64>,
2112) -> Result<f64, String> {
2113    if let Some(explicit) = time_anchor {
2114        return validate_survival_time_anchor_override(explicit);
2115    }
2116    if survival_mode == SurvivalLikelihoodMode::MarginalSlope
2117        || survival_data_is_left_truncated(age_entry)
2118    {
2119        survival_robust_interior_time_anchor(age_exit)
2120    } else {
2121        survival_earliest_entry_time_anchor(age_entry)
2122    }
2123}
2124
2125pub fn evaluate_survival_time_basis_row(
2126    age: f64,
2127    cfg: &SurvivalTimeBasisConfig,
2128) -> Result<Array1<f64>, String> {
2129    if !age.is_finite() || age < 0.0 {
2130        return Err(format!(
2131            "survival time basis row requires finite non-negative age, got {age}"
2132        ));
2133    }
2134    let age = age.max(SURVIVAL_TIME_FLOOR);
2135    let log_age = array![age.ln()];
2136    match cfg {
2137        SurvivalTimeBasisConfig::None => Ok(Array1::zeros(0)),
2138        // Single `log t` column (#2301): the confounded constant column is gone,
2139        // its location absorbed by the mean intercept. See the Linear arm of
2140        // `build_survival_time_basis`.
2141        SurvivalTimeBasisConfig::Linear => Ok(array![age.ln()]),
2142        SurvivalTimeBasisConfig::BSpline { degree, knots, .. } => {
2143            if knots.is_empty() {
2144                return Err(
2145                    "survival BSpline anchor evaluation requires resolved knot metadata"
2146                        .to_string(),
2147                );
2148            }
2149            let built = build_bspline_basis_1d(
2150                log_age.view(),
2151                &BSplineBasisSpec {
2152                    degree: *degree,
2153                    penalty_order: 2,
2154                    knotspec: BSplineKnotSpec::Provided(knots.clone()),
2155                    double_penalty: false,
2156                    identifiability: BSplineIdentifiability::None,
2157                    boundary: OneDimensionalBoundary::Open,
2158                    boundary_conditions: BSplineBoundaryConditions::default(),
2159                },
2160            )
2161            .map_err(|e| format!("failed to evaluate survival bspline anchor row: {e}"))?;
2162            Ok(built.design.to_dense().row(0).to_owned())
2163        }
2164        SurvivalTimeBasisConfig::ISpline {
2165            degree,
2166            knots,
2167            keep_cols,
2168            ..
2169        } => {
2170            if knots.is_empty() {
2171                return Err(
2172                    "survival ISpline anchor evaluation requires resolved knot metadata"
2173                        .to_string(),
2174                );
2175            }
2176            // The anchor is the ORIGIN of the baseline reparameterization, not
2177            // a prediction: `center_survival_time_designs_at_anchor` subtracts
2178            // this row from every entry and exit design row, and the fit is
2179            // invariant to it up to the baseline offset. So it has to be a time
2180            // the baseline is IDENTIFIED at.
2181            //
2182            // The default anchor for ordinary right-censored data is the
2183            // earliest entry, which is the time origin, which
2184            // `evaluate_survival_time_basis_row` floors to
2185            // `SURVIVAL_TIME_FLOOR = 1e-9` so `ln` stays finite — i.e. `−20.7`
2186            // in the basis's own coordinate, far below the first knot and a
2187            // readout of the floor rather than of the data. Under the saturating
2188            // convention that was invisible, because every time at or below the
2189            // first knot maps to the anchored ZERO row. Under the linear tails
2190            // the baseline now carries (gam#2705) it would instead re-center
2191            // every design column by a large constant — which is exactly the
2192            // #751 inflation the anchor rule exists to avoid.
2193            //
2194            // Clamping the anchor into the modelling interval says the thing
2195            // that is actually meant, and is numerically identical to what
2196            // shipped for every anchor at or below the first knot.
2197            let interval = ispline_modelling_interval(knots.view(), *degree)
2198                .map_err(|e| format!("failed to resolve ispline modelling interval: {e}"))?;
2199            let anchor_log_age = match interval {
2200                Some((left, right)) => array![log_age[0].clamp(left, right)],
2201                None => log_age.clone(),
2202            };
2203            let (basis_arc, _) = create_basis::<Dense>(
2204                anchor_log_age.view(),
2205                KnotSource::Provided(knots.view()),
2206                *degree,
2207                BasisOptions::i_spline(),
2208            )
2209            .map_err(|e| format!("failed to evaluate survival ispline anchor row: {e}"))?;
2210            let basis = basis_arc.as_ref();
2211            let row = basis.row(0);
2212            if keep_cols.is_empty() {
2213                return Ok(row.to_owned());
2214            }
2215            if keep_cols.iter().any(|&j| j >= row.len()) {
2216                return Err(SurvivalConstructionError::MissingColumn {
2217                    reason: "survival ISpline anchor keep_cols exceed basis width".to_string(),
2218                }
2219                .into());
2220            }
2221            Ok(Array1::from_iter(keep_cols.iter().map(|&j| row[j])))
2222        }
2223    }
2224}
2225
2226pub fn center_survival_time_designs_at_anchor(
2227    design_entry: &mut DesignMatrix,
2228    design_exit: &mut DesignMatrix,
2229    anchor_row: &Array1<f64>,
2230) -> Result<(), String> {
2231    if design_entry.ncols() != anchor_row.len() || design_exit.ncols() != anchor_row.len() {
2232        return Err(format!(
2233            "survival time anchoring column mismatch: entry={}, exit={}, anchor={}",
2234            design_entry.ncols(),
2235            design_exit.ncols(),
2236            anchor_row.len()
2237        ));
2238    }
2239    // Centering destroys sparsity (every row gets a dense offset), so
2240    // materialize to dense.  This only runs once at construction time.
2241    fn center_dense(dm: &mut DesignMatrix, anchor: &Array1<f64>) {
2242        let mut dense = dm.to_dense();
2243        for mut row in dense.rows_mut() {
2244            row -= &anchor.view();
2245        }
2246        *dm = DesignMatrix::Dense(DenseDesignMatrix::from(dense));
2247    }
2248    center_dense(design_entry, anchor_row);
2249    center_dense(design_exit, anchor_row);
2250    Ok(())
2251}
2252
2253// ---------------------------------------------------------------------------
2254// Baseline evaluation (Gompertz, Weibull, Gompertz-Makeham)
2255// ---------------------------------------------------------------------------
2256
2257/// Partial derivatives of the baseline offsets `(eta_target, d_eta_target/dt)`
2258/// with respect to the θ-parameters in the same parameterization that
2259/// [`survival_baseline_theta_from_config`] / [`survival_baseline_config_from_theta`]
2260/// use:
2261///
2262/// - **Weibull**: θ = (log_scale, log_shape).  `eta = shape·(log t − log scale)`,
2263///   `o_D = shape/t`.
2264/// - **Gompertz**: θ = (log_rate, shape).  `eta = log H_G(t)` with
2265///   `H_G(t) = (rate/shape)·(exp(shape·t) − 1)`, `o_D = h_G(t)/H_G(t) =
2266///   shape·E/(E−1)` where `E = exp(shape·t)`.
2267/// - **Gompertz–Makeham**: θ = (log_rate, shape, log_makeham).
2268///   `eta = log H(t)` with `H(t) = makeham·t + H_G(t)`,
2269///   `o_D = (makeham + h_G(t)) / H(t)`.
2270///
2271/// Returns a flat `(d_eta/dθ_k, d_oD/dθ_k)` pair for each component of θ,
2272/// in the same order as `survival_baseline_theta_from_config`.  Linear has
2273/// no θ-parameters so returns `Ok(None)`.
2274///
2275/// The `eta`-channel derivatives are closed-form for every branch.  The
2276/// `o_D`-channel derivatives use the log-derivative identity
2277/// `∂o_D/∂θ = o_D · ∂log(o_D)/∂θ` which is more numerically stable near
2278/// the small-shape limit (shape·t → 0).  Near shape = 0 we fall back to
2279/// a third-order Taylor expansion with the same 1e-10 pivot that
2280/// `evaluate_survival_baseline` uses, keeping the value/derivative pair
2281/// continuous and agreement with the linear-hazard limit exact at shape=0.
2282pub fn baseline_offset_theta_partials(
2283    age: f64,
2284    cfg: &SurvivalBaselineConfig,
2285) -> Result<Option<Vec<(f64, f64)>>, String> {
2286    let Some(params) = validated_baseline_params(age, cfg, "baseline derivative evaluation")?
2287    else {
2288        return Ok(None);
2289    };
2290
2291    match params {
2292        ValidatedBaselineTarget::Weibull { scale, shape } => {
2293            // eta = shape·(log t − log scale)
2294            //     = shape·log t − shape·log scale
2295            // o_D = shape / t
2296            //
2297            // θ = (log_scale, log_shape):
2298            //   ∂eta/∂log_scale  = −shape          ∂o_D/∂log_scale = 0
2299            //   ∂eta/∂log_shape  = shape·(log t − log scale) = eta
2300            //   ∂o_D/∂log_shape  = shape / t = o_D
2301            let eta = shape * (age.ln() - scale.ln());
2302            let o_d = shape / age;
2303            let d_eta_d_log_scale = -shape;
2304            let d_od_d_log_scale = 0.0;
2305            let d_eta_d_log_shape = eta;
2306            let d_od_d_log_shape = o_d;
2307            Ok(Some(vec![
2308                (d_eta_d_log_scale, d_od_d_log_scale),
2309                (d_eta_d_log_shape, d_od_d_log_shape),
2310            ]))
2311        }
2312        ValidatedBaselineTarget::Gompertz { shape, .. } => {
2313            // θ = (log_rate, shape):
2314            //   Rate cancels in o_D = h/H for Gompertz, so ∂o_D/∂log_rate = 0
2315            //   and ∂eta/∂log_rate = 1. The shape channel uses
2316            //     ∂eta/∂shape   = −1/shape + t·E/(E−1)
2317            //     ∂log(o_D)/∂shape = 1/shape − t/(E−1)
2318            //     ∂o_D/∂shape  = o_D · ∂log(o_D)/∂shape
2319            //   Near shape=0 both numerators are 1/shape cancellations. Use
2320            //   Taylor expansions with the same 1e-10 pivot that
2321            //   gompertz_components uses in evaluate_survival_baseline.
2322            let (d_eta_d_shape, d_od_d_shape) = gompertz_shape_derivatives(age, shape);
2323            Ok(Some(vec![(1.0, 0.0), (d_eta_d_shape, d_od_d_shape)]))
2324        }
2325        ValidatedBaselineTarget::GompertzMakeham {
2326            rate,
2327            shape,
2328            makeham,
2329        } => {
2330            // H(t) = M·t + H_G(t),   H_G(t) = (rate/shape)·(E−1),  E = exp(shape·t)
2331            // h(t) = M + h_G(t),     h_G(t) = rate·E
2332            // o_D  = h/H
2333            //
2334            // θ = (log_rate, shape, log_makeham):
2335            //   ∂H/∂log_rate    = rate · ∂H/∂rate = H_G               (scales with rate)
2336            //   ∂H/∂shape       = H_G_shape                            (closed form below)
2337            //   ∂H/∂log_makeham = makeham · t                          (linear in makeham)
2338            //   ∂h/∂log_rate    = rate · ∂h/∂rate = h_G
2339            //   ∂h/∂shape       = h_G_shape = rate·t·E + 0              (= rate·t·E)
2340            //   ∂h/∂log_makeham = makeham
2341            //   ∂eta/∂θ = (∂H/∂θ) / H
2342            //   ∂o_D/∂θ = (∂h/∂θ − o_D · ∂H/∂θ) / H
2343            //           = (∂h/∂θ)/H − o_D · (∂H/∂θ)/H
2344            let (cum_g, inst_g) = gompertz_hazard_components(age, rate, shape);
2345            let cum_total = makeham * age + cum_g;
2346            if cum_total <= 0.0 || !cum_total.is_finite() {
2347                return Err(SurvivalConstructionError::DataValidationFailed {
2348                    reason: "gm baseline produced non-positive cumulative hazard".to_string(),
2349                }
2350                .into());
2351            }
2352            let inst_total = makeham + inst_g;
2353            let o_d = inst_total / cum_total;
2354            let inv_cum = 1.0 / cum_total;
2355            // Each channel: ∂cum/∂θ and ∂inst/∂θ → ∂eta/∂θ = ∂cum/∂θ / cum
2356            //                                       ∂o_D/∂θ = (∂inst/∂θ − o_D·∂cum/∂θ) / cum
2357            // log_rate channel: cum is linear in rate through H_G; ∂cum/∂rate = H_G/rate,
2358            //   so ∂cum/∂log_rate = H_G (= cum_g here). Similarly ∂inst/∂log_rate = h_G (= inst_g).
2359            let d_cum_dlr = cum_g;
2360            let d_inst_dlr = inst_g;
2361            let d_eta_dlr = d_cum_dlr * inv_cum;
2362            let d_od_dlr = (d_inst_dlr - o_d * d_cum_dlr) * inv_cum;
2363            // shape channel: only H_G and h_G have shape dependence.
2364            let (d_cum_dshape, d_inst_dshape) =
2365                gompertz_cumulative_shape_derivative(age, rate, shape);
2366            let d_eta_dshape = d_cum_dshape * inv_cum;
2367            let d_od_dshape = (d_inst_dshape - o_d * d_cum_dshape) * inv_cum;
2368            // log_makeham channel: cum contributes M·t, inst contributes M.
2369            //   ∂cum/∂log_makeham = makeham·t,  ∂inst/∂log_makeham = makeham.
2370            let d_cum_dlm = makeham * age;
2371            let d_inst_dlm = makeham;
2372            let d_eta_dlm = d_cum_dlm * inv_cum;
2373            let d_od_dlm = (d_inst_dlm - o_d * d_cum_dlm) * inv_cum;
2374            Ok(Some(vec![
2375                (d_eta_dlr, d_od_dlr),
2376                (d_eta_dshape, d_od_dshape),
2377                (d_eta_dlm, d_od_dlm),
2378            ]))
2379        }
2380    }
2381}
2382
2383/// Shared chain-rule θ-gradient contraction for baseline offsets.
2384///
2385/// Both [`baseline_chain_rule_gradient`] (RP eta offsets) and
2386/// [`marginal_slope_baseline_chain_rule_gradient`] (probit q-offsets) reduce to
2387/// the same contraction of [`OffsetChannelResiduals`] against per-age baseline
2388/// θ-partials; only the `partials` provider differs. This engine owns the length
2389/// checks, the θ-dim probe, the parallel per-row reduction, the entry gating, and
2390/// the error handling. Each provider returns, per age, a length-`theta_dim` vector
2391/// of `(∂eta/∂θ_k, ∂(d eta/dt)/∂θ_k)` pairs (or `(∂q/∂θ_k, ∂(dq/dt)/∂θ_k)` for the
2392/// probit channel), and `None` when `cfg` has no θ-parameters (`Linear` target).
2393///
2394/// Contract (envelope theorem at converged β; the penalty has no θ dependence):
2395///
2396///   d[0.5·deviance + 0.5·βᵀS_λβ] / dθ_k
2397///     = Σᵢ r_X[i]·(∂o_X_i/∂θ_k) + r_D[i]·(∂o_D_i/∂θ_k) + r_E[i]·(∂o_E_i/∂θ_k)
2398///       + r_R[i]·(∂o_R_i/∂θ_k)
2399///
2400/// where `r_X = residuals.exit`, `r_D = residuals.derivative`, `r_E =
2401/// residuals.entry`, `r_R = residuals.right` (all sampleweight-scaled already).
2402/// Exit and derivative partials both come from the `age_exit[i]` evaluation;
2403/// the entry partial from `age_entry[i]`; the interval upper-bound (`R`)
2404/// η-partial from `age_right[i]`. Origin-entry rows have `r_E[i] == 0` exactly
2405/// and non-interval rows have `r_R[i] == 0` exactly, so those partials are
2406/// skipped for those rows (avoiding the `age > 0` precondition failure when an
2407/// inactive boundary age is 0 / a placeholder).
2408///
2409/// Returns `Ok(None)` when the provider reports no θ-parameters.
2410fn baseline_chain_rule_gradient_with_partials<F>(
2411    label: &'static str,
2412    age_entry: ndarray::ArrayView1<'_, f64>,
2413    age_exit: ndarray::ArrayView1<'_, f64>,
2414    age_right: ndarray::ArrayView1<'_, f64>,
2415    cfg: &SurvivalBaselineConfig,
2416    residuals: &crate::survival::OffsetChannelResiduals,
2417    partials: F,
2418) -> Result<Option<Array1<f64>>, String>
2419where
2420    F: Fn(f64, &SurvivalBaselineConfig) -> Result<Option<Vec<(f64, f64)>>, String> + Sync,
2421{
2422    let n = age_exit.len();
2423    if age_entry.len() != n
2424        || age_right.len() != n
2425        || residuals.exit.len() != n
2426        || residuals.entry.len() != n
2427        || residuals.derivative.len() != n
2428        || residuals.right.len() != n
2429    {
2430        return Err(format!(
2431            "{label}: length mismatch (age_entry={}, age_exit={}, age_right={}, r_exit={}, r_entry={}, r_deriv={}, r_right={})",
2432            age_entry.len(),
2433            n,
2434            age_right.len(),
2435            residuals.exit.len(),
2436            residuals.entry.len(),
2437            residuals.derivative.len(),
2438            residuals.right.len(),
2439        ));
2440    }
2441    // Probe θ-dim via any valid positive age. If the provider returns None the
2442    // config carries no θ-parameters (Linear target) and there is no θ-gradient.
2443    let probe_age = age_exit.iter().copied().find(|v| v.is_finite() && *v > 0.0);
2444    let theta_dim = match probe_age {
2445        Some(t) => match partials(t, cfg)? {
2446            None => return Ok(None),
2447            Some(v) => v.len(),
2448        },
2449        None => {
2450            return Err(format!("{label}: no valid positive age for dim probe"));
2451        }
2452    };
2453    // Per-row partial contractions are independent, but each row's
2454    // contribution is a `theta_dim`-vector of `O(theta_dim · partial_cost)`
2455    // flops — small enough that the rayon parallel reduction's split
2456    // overhead dominates for any plausible `theta_dim`, *and* the
2457    // non-associative IEEE-754 sum order across thread chunks made the
2458    // engine drift in the low-order bits from row to row. The serial
2459    // accumulator below mirrors the inline reference exactly (and remains
2460    // ~memory-bandwidth-bound at large-scale `n`), so the engine is now a
2461    // bit-for-bit replacement for the legacy path, not just a
2462    // floating-point-noise-equivalent one.
2463    let mut grad = Array1::<f64>::zeros(theta_dim);
2464    for i in 0..n {
2465        // Exit + derivative partials both come from the age_exit evaluation.
2466        let partials_exit = partials(age_exit[i], cfg)?
2467            .ok_or_else(|| format!("{label}: unexpected None from partials at exit"))?;
2468        if partials_exit.len() != theta_dim {
2469            return Err(format!(
2470                "{label}: theta_dim drifted ({} != {})",
2471                partials_exit.len(),
2472                theta_dim
2473            ));
2474        }
2475        let r_x = residuals.exit[i];
2476        let r_d = residuals.derivative[i];
2477        for k in 0..theta_dim {
2478            let (d_eta_dk, d_od_dk) = partials_exit[k];
2479            grad[k] += r_x * d_eta_dk + r_d * d_od_dk;
2480        }
2481        // Entry channel is nonzero only for rows with a positive entry
2482        // interval; for origin-entry rows age_entry may be 0 and calling
2483        // the provider would error. Gate on residual==0.
2484        let r_e = residuals.entry[i];
2485        if r_e != 0.0 {
2486            let partials_entry = partials(age_entry[i], cfg)?
2487                .ok_or_else(|| format!("{label}: unexpected None from partials at entry"))?;
2488            for k in 0..theta_dim {
2489                grad[k] += r_e * partials_entry[k].0;
2490            }
2491        }
2492        // Interval upper-bound (`R`) channel: `q_right = X_time(R)·β + o_R(θ)`
2493        // carries its own baseline-θ η-offset evaluated at `age_right[i]`. It is
2494        // an η-level offset with NO time-derivative channel (the interval
2495        // likelihood `log[S(L) − S(R)]` has no hazard-derivative term), so it
2496        // contracts against the η-partial `.0` only. Nonzero only for
2497        // interval-censored latent rows; for every other channel/model
2498        // `r_right[i] == 0` exactly, so the (possibly placeholder) `age_right[i]`
2499        // partial is never consulted.
2500        let r_r = residuals.right[i];
2501        if r_r != 0.0 {
2502            let partials_right = partials(age_right[i], cfg)?.ok_or_else(|| {
2503                format!("{label}: unexpected None from partials at right boundary")
2504            })?;
2505            if partials_right.len() != theta_dim {
2506                return Err(format!(
2507                    "{label}: theta_dim drifted at right boundary ({} != {})",
2508                    partials_right.len(),
2509                    theta_dim
2510                ));
2511            }
2512            for k in 0..theta_dim {
2513                grad[k] += r_r * partials_right[k].0;
2514            }
2515        }
2516    }
2517    Ok(Some(grad))
2518}
2519
2520/// Contract `OffsetChannelResiduals` against `baseline_offset_theta_partials`
2521/// to produce the closed-form θ-gradient of the unpenalized NLL at converged β.
2522///
2523/// Derivation (envelope theorem on the penalized objective, β* minimizes the
2524/// same cost wrt β and the penalty has no θ dependence):
2525///
2526///   d[0.5·deviance + 0.5·βᵀS_λβ] / dθ_k
2527///     = d[NLL(β*; o(θ))] / dθ_k
2528///     = Σᵢ (∂NLL_i/∂o_X\[i\])·(∂o_X_i/∂θ_k)
2529///       + (∂NLL_i/∂o_E\[i\])·(∂o_E_i/∂θ_k)
2530///       + (∂NLL_i/∂o_D\[i\])·(∂o_D_i/∂θ_k)
2531///       + (∂NLL_i/∂o_R\[i\])·(∂o_R_i/∂θ_k)
2532///
2533/// The four `∂NLL_i/∂o_channel` terms are the `exit`, `entry`, `derivative`,
2534/// `right` fields of `OffsetChannelResiduals` (sampleweight-scaled already).
2535/// The `∂o/∂θ_k` terms come from [`baseline_offset_theta_partials`] per obs at
2536/// the appropriate age.
2537///
2538/// Per the RP offset convention:
2539///   o_E\[i\] = eta_target(age_entry\[i\])
2540///   o_X\[i\] = eta_target(age_exit\[i\])
2541///   o_D\[i\] = d/dt eta_target(t) |_{t=age_exit\[i\]}
2542///   o_R\[i\] = eta_target(age_right\[i\])   (interval upper bound `R`; η-level only)
2543///
2544/// so the exit and derivative partials are both evaluated at `age_exit[i]`,
2545/// the entry partial at `age_entry[i]`, and the interval-right η-partial at
2546/// `age_right[i]`. The origin-entry case (`entry_at_origin[i]`) has
2547/// `r_entry[i] = 0` exactly and every non-interval row has `r_right[i] = 0`
2548/// exactly, so we skip the `baseline_offset_theta_partials(age, ..)` call for
2549/// those rows (avoiding the `age > 0` precondition failure when an inactive
2550/// boundary age is 0 / a placeholder).
2551///
2552/// Returns `Ok(None)` when `cfg.target == Linear` (no θ-parameters).
2553pub fn baseline_chain_rule_gradient(
2554    age_entry: ndarray::ArrayView1<'_, f64>,
2555    age_exit: ndarray::ArrayView1<'_, f64>,
2556    age_right: ndarray::ArrayView1<'_, f64>,
2557    cfg: &SurvivalBaselineConfig,
2558    residuals: &crate::survival::OffsetChannelResiduals,
2559) -> Result<Option<Array1<f64>>, String> {
2560    baseline_chain_rule_gradient_with_partials(
2561        "baseline_chain_rule_gradient",
2562        age_entry,
2563        age_exit,
2564        age_right,
2565        cfg,
2566        residuals,
2567        baseline_offset_theta_partials,
2568    )
2569}
2570
2571/// Chain-rule θ-gradient for marginal-slope probit baseline offsets.
2572///
2573/// This is the probit-survival counterpart of [`baseline_chain_rule_gradient`].
2574/// It contracts residuals against
2575/// [`marginal_slope_baseline_offset_theta_partials`], so the offset channels
2576/// are `(q_entry, q_exit, dq_exit/dt)` with `Phi(-q(t)) = exp(-H0(t))`.
2577pub fn marginal_slope_baseline_chain_rule_gradient(
2578    age_entry: ndarray::ArrayView1<'_, f64>,
2579    age_exit: ndarray::ArrayView1<'_, f64>,
2580    cfg: &SurvivalBaselineConfig,
2581    residuals: &crate::survival::OffsetChannelResiduals,
2582) -> Result<Option<Array1<f64>>, String> {
2583    // Marginal-slope has no interval upper-bound channel; `residuals.right` is
2584    // all-zero, so the right channel never contracts and `age_exit` serves as an
2585    // unconsulted placeholder for the (unused) `age_right` argument.
2586    baseline_chain_rule_gradient_with_partials(
2587        "marginal_slope_baseline_chain_rule_gradient",
2588        age_entry,
2589        age_exit,
2590        age_exit,
2591        cfg,
2592        residuals,
2593        marginal_slope_baseline_offset_theta_partials,
2594    )
2595}
2596
2597/// Shared Gompertz hazard components `(H_G(t), h_G(t))`.
2598/// Mirrors the private helper in `evaluate_survival_baseline` with the
2599/// same 1e-10 small-shape pivot.
2600#[inline]
2601fn gompertz_hazard_components(age: f64, rate: f64, shape: f64) -> (f64, f64) {
2602    if shape.abs() < 1e-10 {
2603        // Taylor at shape=0: H_G(t) = rate·t·(1 + shape·t/2 + (shape·t)²/6),
2604        // h_G(t) = rate·(1 + shape·t + (shape·t)²/2).
2605        let x = shape * age;
2606        (
2607            rate * age * (1.0 + 0.5 * x + x * x / 6.0),
2608            rate * (1.0 + x + 0.5 * x * x),
2609        )
2610    } else {
2611        let shape_age = shape * age;
2612        let cumulative_hazard = (rate / shape) * shape_age.exp_m1();
2613        let instant_hazard = rate * shape_age.exp();
2614        (cumulative_hazard, instant_hazard)
2615    }
2616}
2617
2618/// Partials of `(H_G(t), h_G(t))` with respect to the shape parameter.
2619///
2620/// H_G(t) = (rate/shape)·(E−1),  h_G(t) = rate·E,  E = exp(shape·t)
2621///
2622/// ∂H_G/∂shape  = −(rate/shape²)·(E−1) + (rate/shape)·t·E
2623///              = rate·[t·E/shape − (E−1)/shape²]
2624///              = rate·[t·E·shape − (E−1)] / shape²
2625/// ∂h_G/∂shape  = rate·t·E
2626///
2627/// Near shape=0 the first expression has a 1/shape² singularity that
2628/// cancels analytically. Using the series E−1 = Σₖ≥₁ (shape·t)ᵏ/k!:
2629///   t·E·shape − (E−1) = Σₖ≥₁ (shape·t)ᵏ·(k−1)/k!·shape⁰  [after simplification]
2630///                     = (shape·t)²/2 + 2(shape·t)³/6 + 3(shape·t)⁴/24 + ...
2631/// so ∂H_G/∂shape at shape→0 = rate·[t²/2 + shape·t³/3 + shape²·t⁴/8 + ...].
2632/// We use that Taylor expansion in the small-shape branch.
2633#[inline]
2634fn gompertz_cumulative_shape_derivative(age: f64, rate: f64, shape: f64) -> (f64, f64) {
2635    let x = shape * age;
2636    let dinstg_dshape = rate * age * x.exp();
2637    // The exact form rate·[t·E·shape − (E−1)]/shape² is a difference of two
2638    // O(1/shape) quantities whose leading terms cancel, so its accuracy is
2639    // governed by the dimensionless product x = shape·age, NOT by `shape`
2640    // alone. Pivoting on `shape < 1e-10` ignored `age`: for large ages a small
2641    // shape still yields a small x where the catastrophic cancellation has
2642    // already corrupted the difference. Pivot on x instead; the 3-term Taylor
2643    // (through O(x²)) is accurate to <1e-9 for |x| < 1e-4, and the exact branch
2644    // is clean above it.
2645    let dhg_dshape = if x.abs() < 1e-4 {
2646        let t = age;
2647        // Truncated to O(x³): t²/2 + x·t²/3 + x²·t²/8
2648        rate * t * t * (0.5 + x / 3.0 + x * x / 8.0)
2649    } else {
2650        // t·E·shape − (E−1) = t·e^x·shape − expm1(x)
2651        let e = x.exp();
2652        let em1 = x.exp_m1();
2653        let numerator = age * e * shape - em1;
2654        rate * numerator / (shape * shape)
2655    };
2656    (dhg_dshape, dinstg_dshape)
2657}
2658
2659/// Partials `(∂eta/∂shape, ∂o_D/∂shape)` for the pure Gompertz baseline.
2660/// Pure Gompertz has rate cancelling in o_D, so there is no log_rate
2661/// contribution in o_D. The rate channel for eta is trivially 1; this
2662/// helper only covers the shape channel.
2663#[inline]
2664fn gompertz_shape_derivatives(age: f64, shape: f64) -> (f64, f64) {
2665    if shape.abs() < 1e-10 {
2666        // Closed-form limits from the series t·E/(E−1) = 1/x + 1/2 + x/12 + ...
2667        // with E = e^x, x = shape·t:
2668        //   ∂eta/∂shape  = −1/shape + t·E/(E−1)
2669        //                = t/2 + shape·t²/12 + O(shape²)
2670        //   o_D         = shape·E/(E−1)
2671        //                = 1/t + shape/2 + shape²·t/12 + O(shape³)
2672        //   ∂log(o_D)/∂shape = 1/shape − t/(E−1)
2673        //                = t/2 − shape·t²/12 + O(shape²)
2674        //   ∂o_D/∂shape = o_D · ∂log(o_D)/∂shape
2675        let t = age;
2676        let d_eta = 0.5 * t + shape * t * t / 12.0;
2677        let dlog_od = 0.5 * t - shape * t * t / 12.0;
2678        let o_d = 1.0 / t + 0.5 * shape + shape * shape * t / 12.0;
2679        (d_eta, o_d * dlog_od)
2680    } else {
2681        let x = shape * age;
2682        let e = x.exp();
2683        let em1 = x.exp_m1(); // E − 1 via expm1 for accuracy at small x
2684        let d_eta = -1.0 / shape + age * e / em1;
2685        // o_D = shape · E/(E−1); ∂log(o_D)/∂shape = 1/shape − t/(E−1)
2686        let o_d = shape * e / em1;
2687        let dlog_od = 1.0 / shape - age / em1;
2688        (d_eta, o_d * dlog_od)
2689    }
2690}
2691
2692/// Per-target baseline parameters after the shared age guard and the per-target
2693/// required-field extraction + finiteness/positivity validation have passed.
2694///
2695/// This is the single source of truth for *which* config fields each baseline
2696/// target requires and *what* domain each must satisfy. Both the hazard-value
2697/// evaluator (`survival_cumulative_and_instant_hazard`) and the θ-partials
2698/// evaluator (`survival_hazard_theta_partials`) consume it and only differ in how
2699/// they assemble their (value vs derivative) outputs from these checked scalars.
2700#[derive(Clone, Copy, Debug)]
2701enum ValidatedBaselineTarget {
2702    Weibull { scale: f64, shape: f64 },
2703    Gompertz { rate: f64, shape: f64 },
2704    GompertzMakeham { rate: f64, shape: f64, makeham: f64 },
2705}
2706
2707/// Shared prologue for the survival baseline hazard evaluators: validate the age,
2708/// then extract and domain-check the per-target parameters from `cfg`.
2709///
2710/// `Ok(None)` is the `Linear` target (no parametric baseline). `context` is woven
2711/// into the age-guard error so each caller keeps its specific phrasing.
2712fn validated_baseline_params(
2713    age: f64,
2714    cfg: &SurvivalBaselineConfig,
2715    context: &str,
2716) -> Result<Option<ValidatedBaselineTarget>, String> {
2717    if !age.is_finite() || age <= 0.0 {
2718        return Err(format!(
2719            "survival ages must be finite and positive for {context}"
2720        ));
2721    }
2722
2723    match cfg.target {
2724        SurvivalBaselineTarget::Linear => Ok(None),
2725        SurvivalBaselineTarget::Weibull => {
2726            let scale = cfg
2727                .scale
2728                .ok_or_else(|| "weibull missing scale".to_string())?;
2729            let shape = cfg
2730                .shape
2731                .ok_or_else(|| "weibull missing shape".to_string())?;
2732            if !(scale.is_finite() && shape.is_finite() && scale > 0.0 && shape > 0.0) {
2733                return Err(SurvivalConstructionError::InvalidConfig {
2734                    reason: "weibull baseline requires finite positive scale and shape".to_string(),
2735                }
2736                .into());
2737            }
2738            Ok(Some(ValidatedBaselineTarget::Weibull { scale, shape }))
2739        }
2740        SurvivalBaselineTarget::Gompertz => {
2741            let rate = cfg
2742                .rate
2743                .ok_or_else(|| "gompertz missing rate".to_string())?;
2744            let shape = cfg
2745                .shape
2746                .ok_or_else(|| "gompertz missing shape".to_string())?;
2747            if !(rate.is_finite() && shape.is_finite() && rate > 0.0) {
2748                return Err(
2749                    "gompertz baseline requires finite positive rate and finite shape".to_string(),
2750                );
2751            }
2752            Ok(Some(ValidatedBaselineTarget::Gompertz { rate, shape }))
2753        }
2754        SurvivalBaselineTarget::GompertzMakeham => {
2755            let rate = cfg
2756                .rate
2757                .ok_or_else(|| "gompertz-makeham missing rate".to_string())?;
2758            let shape = cfg
2759                .shape
2760                .ok_or_else(|| "gompertz-makeham missing shape".to_string())?;
2761            let makeham = cfg
2762                .makeham
2763                .ok_or_else(|| "gompertz-makeham missing makeham".to_string())?;
2764            if !(rate.is_finite()
2765                && shape.is_finite()
2766                && makeham.is_finite()
2767                && rate > 0.0
2768                && makeham > 0.0)
2769            {
2770                return Err(
2771                    "gompertz-makeham baseline requires finite positive rate, makeham, and finite shape"
2772                        .to_string(),
2773                );
2774            }
2775            Ok(Some(ValidatedBaselineTarget::GompertzMakeham {
2776                rate,
2777                shape,
2778                makeham,
2779            }))
2780        }
2781    }
2782}
2783
2784fn survival_hazard_theta_partials(
2785    age: f64,
2786    cfg: &SurvivalBaselineConfig,
2787) -> Result<Option<Vec<(f64, f64)>>, String> {
2788    let Some(params) = validated_baseline_params(age, cfg, "baseline hazard partials")? else {
2789        return Ok(None);
2790    };
2791
2792    match params {
2793        ValidatedBaselineTarget::Weibull { scale, shape } => {
2794            let log_time_ratio = age.ln() - scale.ln();
2795            let cumulative_hazard = (age / scale).powf(shape);
2796            let instant_hazard = shape * cumulative_hazard / age;
2797            let eta = shape * log_time_ratio;
2798            Ok(Some(vec![
2799                (-shape * cumulative_hazard, -shape * instant_hazard),
2800                (eta * cumulative_hazard, (1.0 + eta) * instant_hazard),
2801            ]))
2802        }
2803        ValidatedBaselineTarget::Gompertz { rate, shape } => {
2804            let (cumulative_hazard, instant_hazard) = gompertz_hazard_components(age, rate, shape);
2805            let (d_cum_dshape, d_inst_dshape) =
2806                gompertz_cumulative_shape_derivative(age, rate, shape);
2807            Ok(Some(vec![
2808                (cumulative_hazard, instant_hazard),
2809                (d_cum_dshape, d_inst_dshape),
2810            ]))
2811        }
2812        ValidatedBaselineTarget::GompertzMakeham {
2813            rate,
2814            shape,
2815            makeham,
2816        } => {
2817            let (cum_gompertz, inst_gompertz) = gompertz_hazard_components(age, rate, shape);
2818            let (d_cum_dshape, d_inst_dshape) =
2819                gompertz_cumulative_shape_derivative(age, rate, shape);
2820            Ok(Some(vec![
2821                (cum_gompertz, inst_gompertz),
2822                (d_cum_dshape, d_inst_dshape),
2823                (makeham * age, makeham),
2824            ]))
2825        }
2826    }
2827}
2828
2829fn survival_cumulative_and_instant_hazard(
2830    age: f64,
2831    cfg: &SurvivalBaselineConfig,
2832) -> Result<Option<(f64, f64)>, String> {
2833    let Some(params) = validated_baseline_params(age, cfg, "baseline hazard evaluation")? else {
2834        return Ok(None);
2835    };
2836
2837    match params {
2838        ValidatedBaselineTarget::Weibull { scale, shape } => {
2839            let cumulative_hazard = (age / scale).powf(shape);
2840            let instant_hazard = shape * cumulative_hazard / age;
2841            Ok(Some((cumulative_hazard, instant_hazard)))
2842        }
2843        ValidatedBaselineTarget::Gompertz { rate, shape } => {
2844            let (cumulative_hazard, instant_hazard) = gompertz_hazard_components(age, rate, shape);
2845            Ok(Some((cumulative_hazard, instant_hazard)))
2846        }
2847        ValidatedBaselineTarget::GompertzMakeham {
2848            rate,
2849            shape,
2850            makeham,
2851        } => {
2852            let (h_gompertz, inst_gompertz) = gompertz_hazard_components(age, rate, shape);
2853            Ok(Some((makeham * age + h_gompertz, makeham + inst_gompertz)))
2854        }
2855    }
2856}
2857
2858#[derive(Clone, Copy, Debug)]
2859struct MarginalSlopeBaselinePoint {
2860    instant_hazard: f64,
2861    q: f64,
2862    q_t: f64,
2863}
2864
2865fn evaluate_marginal_slope_baseline_point(
2866    age: f64,
2867    cfg: &SurvivalBaselineConfig,
2868) -> Result<Option<MarginalSlopeBaselinePoint>, String> {
2869    let Some((cumulative_hazard, instant_hazard)) =
2870        survival_cumulative_and_instant_hazard(age, cfg)?
2871    else {
2872        return Ok(None);
2873    };
2874    if !(cumulative_hazard.is_finite() && cumulative_hazard > 0.0) {
2875        return Err(format!(
2876            "{} marginal-slope baseline produced non-positive cumulative hazard",
2877            survival_baseline_targetname(cfg.target)
2878        ));
2879    }
2880    if !(instant_hazard.is_finite() && instant_hazard > 0.0) {
2881        return Err(format!(
2882            "{} marginal-slope baseline produced non-positive instant hazard",
2883            survival_baseline_targetname(cfg.target)
2884        ));
2885    }
2886    let survival = (-cumulative_hazard).exp();
2887    if !(survival.is_finite() && survival > 0.0 && survival < 1.0) {
2888        return Err(format!(
2889            "{} marginal-slope baseline survival must be strictly inside (0,1), got {survival}",
2890            survival_baseline_targetname(cfg.target)
2891        ));
2892    }
2893    let q = -standard_normal_quantile(survival).map_err(|e| {
2894        format!(
2895            "{} marginal-slope baseline failed to invert survival probability {survival}: {e}",
2896            survival_baseline_targetname(cfg.target)
2897        )
2898    })?;
2899    let phi_q = normal_pdf(q);
2900    if !(phi_q.is_finite() && phi_q > 0.0) {
2901        return Err(format!(
2902            "{} marginal-slope baseline produced non-positive probit density phi(q)={phi_q} at q={q}",
2903            survival_baseline_targetname(cfg.target)
2904        ));
2905    }
2906    Ok(Some(MarginalSlopeBaselinePoint {
2907        instant_hazard,
2908        q,
2909        q_t: instant_hazard * survival / phi_q,
2910    }))
2911}
2912
2913/// Evaluate the parametric baseline target at a given age.
2914/// Returns `(eta_target(age), d eta_target / d age)` on the log-cumulative-hazard scale.
2915pub fn evaluate_survival_baseline(
2916    age: f64,
2917    cfg: &SurvivalBaselineConfig,
2918) -> Result<(f64, f64), String> {
2919    if !age.is_finite() || age < 0.0 {
2920        return Err(
2921            "survival ages must be finite and non-negative for baseline target evaluation"
2922                .to_string(),
2923        );
2924    }
2925
2926    // At t = 0 every parametric cumulative-hazard target satisfies H(0) = 0
2927    // exactly (this is the defining property of a cumulative hazard:
2928    // S(0) = 1 ⇒ H(0) = -log S(0) = 0). The log-cumulative-hazard offset is
2929    // therefore eta(0) = log H(0) = -inf, and we report a zero log-derivative
2930    // since `exp(eta(0)) = H(0) = 0` is the only physically valid value.
2931    // Returning `Ok((-inf, 0.0))` keeps the baseline cumulative hazard exactly
2932    // zero at the origin; downstream callers that need to multiply this offset
2933    // into a linear predictor are responsible for handling the origin row via
2934    // the `entry_at_origin` / `exit_at_origin` gating already wired through the
2935    // engine.
2936    if age == 0.0 {
2937        return match cfg.target {
2938            SurvivalBaselineTarget::Linear => Ok((0.0, 0.0)),
2939            SurvivalBaselineTarget::Weibull
2940            | SurvivalBaselineTarget::Gompertz
2941            | SurvivalBaselineTarget::GompertzMakeham => Ok((f64::NEG_INFINITY, 0.0)),
2942        };
2943    }
2944
2945    let Some(params) = validated_baseline_params(age, cfg, "baseline target evaluation")? else {
2946        return Ok((0.0, 0.0));
2947    };
2948
2949    match params {
2950        ValidatedBaselineTarget::Weibull { scale, shape } => {
2951            let eta = shape * (age.ln() - scale.ln());
2952            let derivative = shape / age;
2953            Ok((eta, derivative))
2954        }
2955        ValidatedBaselineTarget::Gompertz { rate, shape } => {
2956            let (h, inst) = gompertz_hazard_components(age, rate, shape);
2957            if h <= 0.0 || !h.is_finite() {
2958                return Err(if shape.abs() < 1e-10 {
2959                    "invalid gompertz baseline at near-zero shape".to_string()
2960                } else {
2961                    "gompertz baseline produced non-positive cumulative hazard".to_string()
2962                });
2963            }
2964            let derivative = inst / h;
2965            Ok((h.ln(), derivative))
2966        }
2967        ValidatedBaselineTarget::GompertzMakeham {
2968            rate,
2969            shape,
2970            makeham,
2971        } => {
2972            let (h_gompertz, inst_gompertz) = gompertz_hazard_components(age, rate, shape);
2973            let h = makeham * age + h_gompertz;
2974            if h <= 0.0 || !h.is_finite() {
2975                return Err(
2976                    "gompertz-makeham baseline produced non-positive cumulative hazard".to_string(),
2977                );
2978            }
2979            let inst = makeham + inst_gompertz;
2980            let derivative = inst / h;
2981            Ok((h.ln(), derivative))
2982        }
2983    }
2984}
2985
2986/// Evaluate the parametric baseline as the probit index whose marginal
2987/// survival is the true hazard survival `exp(-H0(t))`.
2988///
2989/// Returns `(q(age), dq / d age)` such that `Phi(-q(age)) = exp(-H0(age))`.
2990/// The derivative is `h0(t) * exp(-H0(t)) / phi(q(t))`.
2991pub fn evaluate_survival_marginal_slope_baseline(
2992    age: f64,
2993    cfg: &SurvivalBaselineConfig,
2994) -> Result<(f64, f64), String> {
2995    // Survival-curve origin. Every cumulative-hazard baseline satisfies
2996    // `H0(0) = 0` (`S0(0) = exp(-H0(0)) = 1`), so the probit index
2997    // `q(0) = -Phi^{-1}(S0(0)) = -Phi^{-1}(1) = -inf`: there is no *finite*
2998    // probit-survival offset at the origin. The survival surface anchors
2999    // `S(0) = 1` directly (see the `t <= 0` origin handling in the survival
3000    // predict paths), so the baseline contributes nothing here — report the
3001    // zero offset rather than aborting in the `age <= 0` hazard guard. This
3002    // mirrors `evaluate_survival_baseline`'s explicit `age == 0` branch on the
3003    // log-cumulative-hazard channel; without it the probit/marginal-slope
3004    // baseline path (location-scale + marginal-slope likelihoods) could not be
3005    // evaluated on a prediction grid whose first node is the origin (#1024).
3006    if age == 0.0 {
3007        return Ok((0.0, 0.0));
3008    }
3009    let Some(point) = evaluate_marginal_slope_baseline_point(age, cfg)? else {
3010        return Ok((0.0, 0.0));
3011    };
3012    Ok((point.q, point.q_t))
3013}
3014
3015/// Partial derivatives of the true survival marginal-slope probit offsets
3016/// `(q(t), dq(t)/dt)` with respect to the baseline θ-parameters.
3017///
3018/// The returned channels match `survival_baseline_theta_from_config`.  For
3019/// Gompertz-Makeham, θ is `(log_rate, shape, log_makeham)`.  If
3020/// `S(t)=exp(-H(t))`, `q(t)=-Phi^-1(S(t))`, `A(t)=S(t)/phi(q(t))`, and
3021/// `h(t)=dH/dt`, then
3022///
3023///   dq/dθ      = A * dH/dθ
3024///   d(q')/dθ   = A * (dh/dθ + h * (q*A - 1) * dH/dθ)
3025///
3026/// which keeps the probit transform and the hazard baseline analytically tied.
3027pub fn marginal_slope_baseline_offset_theta_partials(
3028    age: f64,
3029    cfg: &SurvivalBaselineConfig,
3030) -> Result<Option<Vec<(f64, f64)>>, String> {
3031    let Some(point) = evaluate_marginal_slope_baseline_point(age, cfg)? else {
3032        return Ok(None);
3033    };
3034    let hazard_partials = survival_hazard_theta_partials(age, cfg)?
3035        .ok_or_else(|| "unexpected missing hazard partials for nonlinear baseline".to_string())?;
3036    let a = point.q_t / point.instant_hazard;
3037    let a_log_derivative_factor = point.q * a - 1.0;
3038    Ok(Some(
3039        hazard_partials
3040            .into_iter()
3041            .map(|(d_h_cum, d_h_inst)| {
3042                (
3043                    a * d_h_cum,
3044                    a * (d_h_inst + point.instant_hazard * a_log_derivative_factor * d_h_cum),
3045                )
3046            })
3047            .collect(),
3048    ))
3049}
3050
3051/// Contract marginal-slope offset residuals and channel curvatures into the
3052/// exact Hessian with respect to baseline θ-parameters.
3053pub fn marginal_slope_baseline_chain_rule_hessian(
3054    age_entry: ndarray::ArrayView1<'_, f64>,
3055    age_exit: ndarray::ArrayView1<'_, f64>,
3056    cfg: &SurvivalBaselineConfig,
3057    residuals: &crate::survival::OffsetChannelResiduals,
3058    curvatures: &crate::survival::OffsetChannelCurvatures,
3059) -> Result<Option<Array2<f64>>, String> {
3060    let n = age_exit.len();
3061    if age_entry.len() != n
3062        || residuals.exit.len() != n
3063        || residuals.entry.len() != n
3064        || residuals.derivative.len() != n
3065        || curvatures.rows.len() != n
3066    {
3067        return Err(format!(
3068            "marginal_slope_baseline_chain_rule_hessian: length mismatch (age_entry={}, age_exit={}, r_exit={}, r_entry={}, r_deriv={}, h_rows={})",
3069            age_entry.len(),
3070            n,
3071            residuals.exit.len(),
3072            residuals.entry.len(),
3073            residuals.derivative.len(),
3074            curvatures.rows.len(),
3075        ));
3076    }
3077    let probe_age = age_exit.iter().copied().find(|v| v.is_finite() && *v > 0.0);
3078    let dim = match probe_age {
3079        Some(t) => match marginal_slope_baseline_offset_theta_geometry(t, cfg)? {
3080            None => return Ok(None),
3081            Some(parts) => parts.first.len(),
3082        },
3083        None => {
3084            return Err(
3085                "marginal_slope_baseline_chain_rule_hessian: no valid positive age for dim probe"
3086                    .to_string(),
3087            );
3088        }
3089    };
3090    // Per-row Hessian contractions are independent. Each row contributes a
3091    // dim×dim increment combining second partials (exit/entry channels) with
3092    // the curvature-weighted outer product of the (entry, exit, derivative)
3093    // first-partial Jacobians. Fixed row chunks are combined in chunk-index
3094    // order so floating-point addition stays deterministic across Rayon
3095    // scheduling decisions.
3096    let hessian = RowSet::All.par_try_reduce_fold(
3097        n,
3098        || Array2::<f64>::zeros((dim, dim)),
3099        |mut acc, i, _| -> Result<Array2<f64>, String> {
3100            let exit_parts = marginal_slope_baseline_offset_theta_geometry(age_exit[i], cfg)?
3101                .ok_or_else(|| {
3102                    "unexpected None from marginal-slope second partials at exit".to_string()
3103                })?;
3104            if exit_parts.first.len() != dim {
3105                return Err(
3106                    "marginal_slope_baseline_chain_rule_hessian: theta_dim drifted".to_string(),
3107                );
3108            }
3109            let mut entry_parts = None;
3110            if residuals.entry[i] != 0.0 {
3111                entry_parts = Some(
3112                    marginal_slope_baseline_offset_theta_geometry(age_entry[i], cfg)?.ok_or_else(
3113                        || {
3114                            "unexpected None from marginal-slope second partials at entry"
3115                                .to_string()
3116                        },
3117                    )?,
3118                );
3119            }
3120            for a in 0..dim {
3121                for b in 0..dim {
3122                    let j_exit_a = exit_parts.first[a].0;
3123                    let j_exit_b = exit_parts.first[b].0;
3124                    let j_deriv_a = exit_parts.first[a].1;
3125                    let j_deriv_b = exit_parts.first[b].1;
3126                    let mut value = residuals.exit[i] * exit_parts.second[a][b].0
3127                        + residuals.derivative[i] * exit_parts.second[a][b].1;
3128                    if let Some(parts) = entry_parts.as_ref() {
3129                        value += residuals.entry[i] * parts.second[a][b].0;
3130                    }
3131                    let curv = curvatures.rows[i];
3132                    let j_entry_a = entry_parts.as_ref().map_or(0.0, |parts| parts.first[a].0);
3133                    let j_entry_b = entry_parts.as_ref().map_or(0.0, |parts| parts.first[b].0);
3134                    let ja = [j_entry_a, j_exit_a, j_deriv_a];
3135                    let jb = [j_entry_b, j_exit_b, j_deriv_b];
3136                    for u in 0..3 {
3137                        for v in 0..3 {
3138                            value += ja[u] * curv[u][v] * jb[v];
3139                        }
3140                    }
3141                    acc[[a, b]] += value;
3142                }
3143            }
3144            Ok(acc)
3145        },
3146        |a, b| Ok(a + b),
3147    )?;
3148    Ok(Some(hessian))
3149}
3150
3151/// Complete analytic baseline chart at one age for survival marginal-slope.
3152///
3153/// `value` is `(q(t), dq(t)/dt)`. `first[k]` and `second[k][l]` are the
3154/// corresponding first and second partials with respect to the nonlinear
3155/// baseline coordinates returned by [`survival_baseline_theta_from_config`].
3156/// A nonlinear baseline evaluated at the survival origin has identically zero
3157/// value and derivatives because origin rows are anchored outside the finite
3158/// probit chart. Linear baselines have no coordinates and return `None`.
3159#[derive(Clone, Debug)]
3160pub struct MarginalSlopeBaselineOffsetThetaGeometry {
3161    pub value: (f64, f64),
3162    pub first: Vec<(f64, f64)>,
3163    pub second: Vec<Vec<(f64, f64)>>,
3164}
3165
3166pub fn marginal_slope_baseline_offset_theta_geometry(
3167    age: f64,
3168    cfg: &SurvivalBaselineConfig,
3169) -> Result<Option<MarginalSlopeBaselineOffsetThetaGeometry>, String> {
3170    if age == 0.0 {
3171        let Some(theta) = survival_baseline_theta_from_config(cfg)? else {
3172            return Ok(None);
3173        };
3174        let dim = theta.len();
3175        return Ok(Some(MarginalSlopeBaselineOffsetThetaGeometry {
3176            value: (0.0, 0.0),
3177            first: vec![(0.0, 0.0); dim],
3178            second: vec![vec![(0.0, 0.0); dim]; dim],
3179        }));
3180    }
3181    let Some(point) = evaluate_marginal_slope_baseline_point(age, cfg)? else {
3182        return Ok(None);
3183    };
3184    let Some((hazard, first, second)) = survival_hazard_theta_first_second(age, cfg)? else {
3185        return Ok(None);
3186    };
3187    let (cum_hazard, instant_hazard) = hazard;
3188    let survival = (-cum_hazard).exp();
3189    let a = survival / normal_pdf(point.q);
3190    let b = point.q * a - 1.0;
3191    let b_factor = a + point.q * b;
3192    let dim = first.len();
3193    let mut first_out = Vec::with_capacity(dim);
3194    let mut second_out = vec![vec![(0.0, 0.0); dim]; dim];
3195    for i in 0..dim {
3196        let (h_i, inst_i) = first[i];
3197        first_out.push((a * h_i, a * (inst_i + instant_hazard * b * h_i)));
3198    }
3199    for i in 0..dim {
3200        for j in i..dim {
3201            let (h_i, inst_i) = first[i];
3202            let (h_j, inst_j) = first[j];
3203            let (h_ij, inst_ij) = second[i][j];
3204            let a_j = a * b * h_j;
3205            let b_j = a * h_j * b_factor;
3206            let q_ij = a * h_ij + a * b * h_i * h_j;
3207            let qt_inner_i = inst_i + instant_hazard * b * h_i;
3208            let qt_ij = a_j * qt_inner_i
3209                + a * (inst_ij + inst_j * b * h_i + instant_hazard * (b_j * h_i + b * h_ij));
3210            let mixed = (q_ij, qt_ij);
3211            second_out[i][j] = mixed;
3212            second_out[j][i] = mixed;
3213        }
3214    }
3215    Ok(Some(MarginalSlopeBaselineOffsetThetaGeometry {
3216        value: (point.q, point.q_t),
3217        first: first_out,
3218        second: second_out,
3219    }))
3220}
3221
3222type HazardFirstSecond = ((f64, f64), Vec<(f64, f64)>, Vec<Vec<(f64, f64)>>);
3223
3224fn survival_hazard_theta_first_second(
3225    age: f64,
3226    cfg: &SurvivalBaselineConfig,
3227) -> Result<Option<HazardFirstSecond>, String> {
3228    let Some(hazard) = survival_cumulative_and_instant_hazard(age, cfg)? else {
3229        return Ok(None);
3230    };
3231    let first = survival_hazard_theta_partials(age, cfg)?
3232        .ok_or_else(|| "unexpected missing hazard partials".to_string())?;
3233    let dim = first.len();
3234    let mut second = vec![vec![(0.0, 0.0); dim]; dim];
3235    match cfg.target {
3236        SurvivalBaselineTarget::Linear => return Ok(None),
3237        SurvivalBaselineTarget::Weibull => {
3238            let scale = cfg
3239                .scale
3240                .ok_or_else(|| "weibull missing scale".to_string())?;
3241            let shape = cfg
3242                .shape
3243                .ok_or_else(|| "weibull missing shape".to_string())?;
3244            let log_time_ratio = age.ln() - scale.ln();
3245            let cumulative_hazard = hazard.0;
3246            let instant_hazard = hazard.1;
3247            let eta = shape * log_time_ratio;
3248            second[0][0] = (
3249                shape * shape * cumulative_hazard,
3250                shape * shape * instant_hazard,
3251            );
3252            second[0][1] = (
3253                -shape * cumulative_hazard * (1.0 + eta),
3254                -shape * instant_hazard * (2.0 + eta),
3255            );
3256            second[1][0] = second[0][1];
3257            second[1][1] = (
3258                eta * cumulative_hazard * (1.0 + eta),
3259                (eta + (1.0 + eta) * (1.0 + eta)) * instant_hazard,
3260            );
3261        }
3262        SurvivalBaselineTarget::Gompertz => {
3263            let rate = cfg
3264                .rate
3265                .ok_or_else(|| "gompertz missing rate".to_string())?;
3266            let shape = cfg
3267                .shape
3268                .ok_or_else(|| "gompertz missing shape".to_string())?;
3269            second[0][0] = first[0];
3270            second[0][1] = first[1];
3271            second[1][0] = first[1];
3272            second[1][1] = gompertz_cumulative_shape_second_derivative(age, rate, shape);
3273        }
3274        SurvivalBaselineTarget::GompertzMakeham => {
3275            let rate = cfg.rate.ok_or_else(|| "gm missing rate".to_string())?;
3276            let shape = cfg.shape.ok_or_else(|| "gm missing shape".to_string())?;
3277            second[0][0] = first[0];
3278            second[0][1] = first[1];
3279            second[1][0] = first[1];
3280            second[1][1] = gompertz_cumulative_shape_second_derivative(age, rate, shape);
3281            second[2][2] = first[2];
3282        }
3283    }
3284    Ok(Some((hazard, first, second)))
3285}
3286
3287#[inline]
3288fn gompertz_cumulative_shape_second_derivative(age: f64, rate: f64, shape: f64) -> (f64, f64) {
3289    let x = shape * age;
3290    // ∂²H_G/∂shape² = rate·[t²·E/shape − 2·(shape·t·E − (E−1))/shape³]. This is
3291    // a difference of O(1/shape³) terms whose leading parts cancel, so its
3292    // floating-point accuracy is governed by x = shape·age — and the
3293    // cancellation is FAR worse than the first derivative's 1/shape² form.
3294    // Empirically the exact branch is already garbage for |x| < ~1e-4 (e.g.
3295    // x=1e-9 gives a ~98% relative error; x=1e-10 a ~9700% error). The old
3296    // `shape < 1e-10` pivot ignored `age` and so routed those small-x cases
3297    // through the cancelling exact form, corrupting the marginal-slope baseline
3298    // Hessian near small shape. Pivot on x with a wider threshold than the
3299    // first derivative: the 3-term Taylor (through O(x²)) holds to <1e-8 for
3300    // |x| < 1e-3, and the exact branch is clean above it.
3301    if x.abs() < 1e-3 {
3302        let t = age;
3303        (
3304            rate * t * t * t * (1.0 / 3.0 + x / 4.0 + x * x / 10.0),
3305            rate * t * t * (1.0 + x + 0.5 * x * x),
3306        )
3307    } else {
3308        let e = x.exp();
3309        let em1 = x.exp_m1();
3310        let n = shape * age * e - em1;
3311        (
3312            rate * (age * age * e / shape - 2.0 * n / (shape * shape * shape)),
3313            rate * age * age * e,
3314        )
3315    }
3316}
3317
3318// ---------------------------------------------------------------------------
3319// Baseline offsets
3320// ---------------------------------------------------------------------------
3321
3322#[derive(Clone, Copy)]
3323enum BaselineOffsetEvaluator {
3324    LogCumulativeHazard,
3325    ProbitSurvival,
3326}
3327
3328impl BaselineOffsetEvaluator {
3329    fn length_error(self) -> String {
3330        match self {
3331            Self::LogCumulativeHazard => SurvivalConstructionError::IncompatibleDimensions {
3332                reason: "survival baseline offsets require matching entry/exit lengths".to_string(),
3333            }
3334            .into(),
3335            Self::ProbitSurvival => {
3336                "survival probit baseline offsets require matching entry/exit lengths".to_string()
3337            }
3338        }
3339    }
3340
3341    fn finite_error(self) -> &'static str {
3342        match self {
3343            Self::LogCumulativeHazard => "non-finite survival baseline offsets computed",
3344            Self::ProbitSurvival => "non-finite survival probit baseline offsets computed",
3345        }
3346    }
3347
3348    fn evaluate(self, age: f64, cfg: &SurvivalBaselineConfig) -> Result<(f64, f64), String> {
3349        match self {
3350            Self::LogCumulativeHazard => evaluate_survival_baseline(age, cfg),
3351            Self::ProbitSurvival => evaluate_survival_marginal_slope_baseline(age, cfg),
3352        }
3353    }
3354
3355    fn exit_is_finite(self, value: f64, age: f64) -> bool {
3356        match self {
3357            Self::LogCumulativeHazard => {
3358                value.is_finite() || (age == 0.0 && value == f64::NEG_INFINITY)
3359            }
3360            Self::ProbitSurvival => value.is_finite(),
3361        }
3362    }
3363}
3364
3365fn build_survival_offsets_with_evaluator(
3366    age_entry: &Array1<f64>,
3367    age_exit: &Array1<f64>,
3368    cfg: &SurvivalBaselineConfig,
3369    evaluator: BaselineOffsetEvaluator,
3370) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), String> {
3371    if age_entry.len() != age_exit.len() {
3372        return Err(evaluator.length_error());
3373    }
3374    let n = age_entry.len();
3375    // Each row's three offsets are independent across i. Compute the triplets
3376    // in parallel, then unpack into three Array1 outputs preserving order.
3377    let triples: Vec<(f64, f64, f64)> = (0..n)
3378        .into_par_iter()
3379        .map(|i| -> Result<(f64, f64, f64), String> {
3380            // Origin-entry rows are multiplied out by the survival engines, so
3381            // keep their entry channel finite even when the evaluator's natural
3382            // value at t=0 is undefined or -inf.
3383            let entry_age = age_entry[i];
3384            let e0 = if !entry_age.is_finite() {
3385                return Err(SurvivalConstructionError::DataValidationFailed {
3386                    reason: format!("non-finite entry age at row {i}"),
3387                }
3388                .into());
3389            } else if entry_age <= 0.0 {
3390                0.0
3391            } else {
3392                evaluator.evaluate(entry_age, cfg)?.0
3393            };
3394            let exit_age = age_exit[i];
3395            let (e1, d1) = evaluator.evaluate(exit_age, cfg)?;
3396            if !e0.is_finite() || !evaluator.exit_is_finite(e1, exit_age) || !d1.is_finite() {
3397                return Err(SurvivalConstructionError::DataValidationFailed {
3398                    reason: evaluator.finite_error().to_string(),
3399                }
3400                .into());
3401            }
3402            Ok((e0, e1, d1))
3403        })
3404        .collect::<Result<Vec<_>, String>>()?;
3405    let mut eta_entry = Array1::<f64>::zeros(n);
3406    let mut eta_exit = Array1::<f64>::zeros(n);
3407    let mut derivative_exit = Array1::<f64>::zeros(n);
3408    for (i, (e0, e1, d1)) in triples.into_iter().enumerate() {
3409        eta_entry[i] = e0;
3410        eta_exit[i] = e1;
3411        derivative_exit[i] = d1;
3412    }
3413    Ok((eta_entry, eta_exit, derivative_exit))
3414}
3415
3416/// Compute baseline target offsets for all observations.
3417/// Returns `(eta_entry, eta_exit, derivative_exit)`.
3418pub fn build_survival_baseline_offsets(
3419    age_entry: &Array1<f64>,
3420    age_exit: &Array1<f64>,
3421    cfg: &SurvivalBaselineConfig,
3422) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), String> {
3423    build_survival_offsets_with_evaluator(
3424        age_entry,
3425        age_exit,
3426        cfg,
3427        BaselineOffsetEvaluator::LogCumulativeHazard,
3428    )
3429}
3430
3431/// Compute probit-survival baseline target offsets for all observations.
3432/// Returns `(q_entry, q_exit, q_derivative_exit)` where `Phi(-q(t)) = exp(-H0(t))`.
3433pub fn build_survival_marginal_slope_baseline_offsets(
3434    age_entry: &Array1<f64>,
3435    age_exit: &Array1<f64>,
3436    cfg: &SurvivalBaselineConfig,
3437) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), String> {
3438    build_survival_offsets_with_evaluator(
3439        age_entry,
3440        age_exit,
3441        cfg,
3442        BaselineOffsetEvaluator::ProbitSurvival,
3443    )
3444}
3445
3446/// Rowwise value, gradient, and Hessian of the complete marginal-slope time
3447/// offset with respect to a nonlinear survival-baseline chart.
3448///
3449/// The first-derivative arrays have shape `n × d` and the second-derivative
3450/// arrays have shape `n × d × d`, where `d = theta.len()`. The value arrays may
3451/// include a frozen non-baseline residual; that residual has zero derivatives.
3452#[derive(Clone, Debug)]
3453pub struct SurvivalMarginalSlopeOffsetGeometry {
3454    pub baseline_config: SurvivalBaselineConfig,
3455    pub theta: Array1<f64>,
3456    pub offset_entry: Array1<f64>,
3457    pub offset_exit: Array1<f64>,
3458    pub derivative_offset_exit: Array1<f64>,
3459    pub offset_entry_theta_first: Array2<f64>,
3460    pub offset_exit_theta_first: Array2<f64>,
3461    pub derivative_offset_exit_theta_first: Array2<f64>,
3462    pub offset_entry_theta_second: Array3<f64>,
3463    pub offset_exit_theta_second: Array3<f64>,
3464    pub derivative_offset_exit_theta_second: Array3<f64>,
3465}
3466
3467fn validate_marginal_slope_baseline_row_geometry(
3468    row: &MarginalSlopeBaselineOffsetThetaGeometry,
3469    dim: usize,
3470    channel: &str,
3471) -> Result<(), String> {
3472    if row.first.len() != dim
3473        || row.second.len() != dim
3474        || row.second.iter().any(|axis| axis.len() != dim)
3475    {
3476        return Err(format!(
3477            "survival marginal-slope baseline {channel} theta dimension drifted"
3478        ));
3479    }
3480    if !row.value.0.is_finite()
3481        || !row.value.1.is_finite()
3482        || row
3483            .first
3484            .iter()
3485            .any(|&(value, derivative)| !value.is_finite() || !derivative.is_finite())
3486        || row
3487            .second
3488            .iter()
3489            .flatten()
3490            .any(|&(value, derivative)| !value.is_finite() || !derivative.is_finite())
3491    {
3492        return Err(format!(
3493            "survival marginal-slope baseline {channel} geometry must be finite"
3494        ));
3495    }
3496    Ok(())
3497}
3498
3499/// Evaluate the nonlinear parametric baseline on every marginal-slope row.
3500///
3501/// This function evaluates only baseline-dependent offset geometry. It never
3502/// constructs or mutates time designs, wiggle knots, penalties, or linear
3503/// constraints. Linear baselines have no hyperparameter chart and return
3504/// `None`.
3505pub fn build_survival_marginal_slope_baseline_geometry(
3506    age_entry: &Array1<f64>,
3507    age_exit: &Array1<f64>,
3508    cfg: &SurvivalBaselineConfig,
3509) -> Result<Option<SurvivalMarginalSlopeOffsetGeometry>, String> {
3510    let Some(theta) = survival_baseline_theta_from_config(cfg)? else {
3511        // A linear baseline has no hyperparameter chart. The length check below
3512        // is not reached in that case, and must not be: this arm is the
3513        // "no chart" answer, not an error.
3514        if age_entry.len() != age_exit.len() {
3515            return Err(
3516                "survival marginal-slope baseline geometry requires matching entry/exit lengths"
3517                    .to_string(),
3518            );
3519        }
3520        return Ok(None);
3521    };
3522    build_survival_marginal_slope_baseline_geometry_at_theta(age_entry, age_exit, cfg, theta)
3523}
3524
3525/// The θ-authored entry: realize the same geometry, but record the caller's own
3526/// `theta` VERBATIM instead of re-deriving it from `cfg`.
3527///
3528/// The two differ, and the difference is not academic. A chart evaluation is
3529/// `θ → cfg → rows`, and the config-authored entry above closes the loop with
3530/// `cfg → θ`. For a Weibull that loop is `ln(exp(θ))`, which is **not** the
3531/// identity in `f64`: measured over a grid on `[-3, 3]`, **17.3%** of
3532/// coordinates come back a ulp or more away, and `θ = 1e-5` comes back 57 269
3533/// ulps away.
3534///
3535/// `SurvivalMarginalSlopeFamilyHyperState` stores this `theta` as the family's
3536/// realized coordinates and `validate_layout` compares them to the outer
3537/// manifest with `to_bits()` equality — deliberately, so a workspace cannot
3538/// reuse row geometry from a neighbouring outer probe. With a re-derived `θ`
3539/// that exactness invariant fails for reasons that have nothing to do with the
3540/// geometry, and the inner solve refuses a point the outer optimizer is merely
3541/// trying to evaluate. See the #2765 measurement: at the acceptance fixture's
3542/// checkpoint the certificate probe refused coordinate 3 side `−` (round trip
3543/// `+1` ulp) and coordinate 4 on BOTH sides (`−57269` and `+3383` ulps), and
3544/// evaluated cleanly everywhere the round trip happened to be exact.
3545///
3546/// So: when a caller HAS a θ, that θ is the authority. `cfg` still drives every
3547/// row's arithmetic; only the recorded coordinates change.
3548pub fn build_survival_marginal_slope_baseline_geometry_at_theta(
3549    age_entry: &Array1<f64>,
3550    age_exit: &Array1<f64>,
3551    cfg: &SurvivalBaselineConfig,
3552    theta: Array1<f64>,
3553) -> Result<Option<SurvivalMarginalSlopeOffsetGeometry>, String> {
3554    if age_entry.len() != age_exit.len() {
3555        return Err(
3556            "survival marginal-slope baseline geometry requires matching entry/exit lengths"
3557                .to_string(),
3558        );
3559    }
3560    if theta.iter().any(|value| !value.is_finite()) {
3561        return Err(
3562            "survival marginal-slope baseline theta coordinates must be finite".to_string(),
3563        );
3564    }
3565    // Round-trip through the public chart decoder before touching row storage.
3566    // This validates every target-specific config value even when `n == 0`.
3567    survival_baseline_config_from_theta(cfg.target, &theta)?;
3568    let dim = theta.len();
3569    let zero = || MarginalSlopeBaselineOffsetThetaGeometry {
3570        value: (0.0, 0.0),
3571        first: vec![(0.0, 0.0); dim],
3572        second: vec![vec![(0.0, 0.0); dim]; dim],
3573    };
3574    let rows = (0..age_exit.len())
3575        .into_par_iter()
3576        .map(
3577            |row_index| -> Result<
3578                (
3579                    MarginalSlopeBaselineOffsetThetaGeometry,
3580                    MarginalSlopeBaselineOffsetThetaGeometry,
3581                ),
3582                String,
3583            > {
3584                let entry_age = age_entry[row_index];
3585                if !entry_age.is_finite() || entry_age < 0.0 {
3586                    return Err(format!(
3587                        "survival marginal-slope entry age must be finite and non-negative at row {row_index}"
3588                    ));
3589                }
3590                let exit_age = age_exit[row_index];
3591                if !exit_age.is_finite() || exit_age < 0.0 {
3592                    return Err(format!(
3593                        "survival marginal-slope exit age must be finite and non-negative at row {row_index}"
3594                    ));
3595                }
3596                let entry = if entry_age == 0.0 {
3597                    zero()
3598                } else {
3599                    marginal_slope_baseline_offset_theta_geometry(entry_age, cfg)?.ok_or_else(
3600                        || {
3601                            "nonlinear survival baseline unexpectedly has no entry geometry"
3602                                .to_string()
3603                        },
3604                    )?
3605                };
3606                let exit = marginal_slope_baseline_offset_theta_geometry(exit_age, cfg)?
3607                    .ok_or_else(|| {
3608                        "nonlinear survival baseline unexpectedly has no exit geometry".to_string()
3609                    })?;
3610                validate_marginal_slope_baseline_row_geometry(&entry, dim, "entry")?;
3611                validate_marginal_slope_baseline_row_geometry(&exit, dim, "exit")?;
3612                Ok((entry, exit))
3613            },
3614        )
3615        .collect::<Result<Vec<_>, String>>()?;
3616
3617    let n = rows.len();
3618    let mut offset_entry = Array1::<f64>::zeros(n);
3619    let mut offset_exit = Array1::<f64>::zeros(n);
3620    let mut derivative_offset_exit = Array1::<f64>::zeros(n);
3621    let mut offset_entry_theta_first = Array2::<f64>::zeros((n, dim));
3622    let mut offset_exit_theta_first = Array2::<f64>::zeros((n, dim));
3623    let mut derivative_offset_exit_theta_first = Array2::<f64>::zeros((n, dim));
3624    let mut offset_entry_theta_second = Array3::<f64>::zeros((n, dim, dim));
3625    let mut offset_exit_theta_second = Array3::<f64>::zeros((n, dim, dim));
3626    let mut derivative_offset_exit_theta_second = Array3::<f64>::zeros((n, dim, dim));
3627    for (row_index, (entry, exit)) in rows.into_iter().enumerate() {
3628        offset_entry[row_index] = entry.value.0;
3629        offset_exit[row_index] = exit.value.0;
3630        derivative_offset_exit[row_index] = exit.value.1;
3631        for axis in 0..dim {
3632            offset_entry_theta_first[[row_index, axis]] = entry.first[axis].0;
3633            offset_exit_theta_first[[row_index, axis]] = exit.first[axis].0;
3634            derivative_offset_exit_theta_first[[row_index, axis]] = exit.first[axis].1;
3635            for other_axis in 0..dim {
3636                offset_entry_theta_second[[row_index, axis, other_axis]] =
3637                    entry.second[axis][other_axis].0;
3638                offset_exit_theta_second[[row_index, axis, other_axis]] =
3639                    exit.second[axis][other_axis].0;
3640                derivative_offset_exit_theta_second[[row_index, axis, other_axis]] =
3641                    exit.second[axis][other_axis].1;
3642            }
3643        }
3644    }
3645    Ok(Some(SurvivalMarginalSlopeOffsetGeometry {
3646        baseline_config: cfg.clone(),
3647        theta,
3648        offset_entry,
3649        offset_exit,
3650        derivative_offset_exit,
3651        offset_entry_theta_first,
3652        offset_exit_theta_first,
3653        derivative_offset_exit_theta_first,
3654        offset_entry_theta_second,
3655        offset_exit_theta_second,
3656        derivative_offset_exit_theta_second,
3657    }))
3658}
3659
3660/// A nonlinear baseline chart over already-prepared marginal-slope offsets.
3661///
3662/// Construction subtracts the initial parametric baseline from the prepared
3663/// offset channels exactly once. Candidate evaluations add a new baseline to
3664/// that same frozen residual. Consequently candidate theta values cannot move
3665/// any prepared time design, wiggle knot, penalty, or feasibility cone; only
3666/// the three row-offset value channels move, with analytic first and second
3667/// derivatives supplied by the same evaluation.
3668#[derive(Clone, Debug)]
3669pub struct SurvivalMarginalSlopeFrozenOffsetChart {
3670    age_entry: Array1<f64>,
3671    age_exit: Array1<f64>,
3672    target: SurvivalBaselineTarget,
3673    initial_theta: Array1<f64>,
3674    lower_theta: Array1<f64>,
3675    upper_theta: Array1<f64>,
3676    fixed_offset_entry: Array1<f64>,
3677    fixed_offset_exit: Array1<f64>,
3678    fixed_derivative_offset_exit: Array1<f64>,
3679}
3680
3681impl SurvivalMarginalSlopeFrozenOffsetChart {
3682    pub fn new(
3683        age_entry: &Array1<f64>,
3684        age_exit: &Array1<f64>,
3685        initial_config: &SurvivalBaselineConfig,
3686        prepared_offset_entry: &Array1<f64>,
3687        prepared_offset_exit: &Array1<f64>,
3688        prepared_derivative_offset_exit: &Array1<f64>,
3689    ) -> Result<Self, String> {
3690        let n = age_exit.len();
3691        if age_entry.len() != n
3692            || prepared_offset_entry.len() != n
3693            || prepared_offset_exit.len() != n
3694            || prepared_derivative_offset_exit.len() != n
3695        {
3696            return Err(format!(
3697                "survival marginal-slope frozen offset chart length mismatch: entry={}, exit={n}, prepared_entry={}, prepared_exit={}, prepared_derivative={}",
3698                age_entry.len(),
3699                prepared_offset_entry.len(),
3700                prepared_offset_exit.len(),
3701                prepared_derivative_offset_exit.len(),
3702            ));
3703        }
3704        if prepared_offset_entry
3705            .iter()
3706            .chain(prepared_offset_exit.iter())
3707            .chain(prepared_derivative_offset_exit.iter())
3708            .any(|value| !value.is_finite())
3709        {
3710            return Err(
3711                "survival marginal-slope prepared offsets must be finite before freezing"
3712                    .to_string(),
3713            );
3714        }
3715        let initial_geometry =
3716            build_survival_marginal_slope_baseline_geometry(age_entry, age_exit, initial_config)?
3717                .ok_or_else(|| {
3718                String::from(
3719                    "survival marginal-slope frozen offset chart requires a nonlinear baseline",
3720                )
3721            })?;
3722        let lower_theta = initial_geometry.theta.mapv(|value| value - 6.0);
3723        let upper_theta = initial_geometry.theta.mapv(|value| value + 6.0);
3724        Ok(Self {
3725            age_entry: age_entry.clone(),
3726            age_exit: age_exit.clone(),
3727            target: initial_config.target,
3728            initial_theta: initial_geometry.theta,
3729            lower_theta,
3730            upper_theta,
3731            fixed_offset_entry: prepared_offset_entry - &initial_geometry.offset_entry,
3732            fixed_offset_exit: prepared_offset_exit - &initial_geometry.offset_exit,
3733            fixed_derivative_offset_exit: prepared_derivative_offset_exit
3734                - &initial_geometry.derivative_offset_exit,
3735        })
3736    }
3737
3738    pub fn target(&self) -> SurvivalBaselineTarget {
3739        self.target
3740    }
3741
3742    pub fn initial_theta(&self) -> &Array1<f64> {
3743        &self.initial_theta
3744    }
3745
3746    /// Declared finite domain of this frozen nonlinear chart. These are the
3747    /// same coordinate bounds used by the legacy standalone baseline solver,
3748    /// now owned by the chart so a joint solver and its terminal certificate
3749    /// cannot silently choose a different domain.
3750    pub fn theta_bounds(&self) -> (&Array1<f64>, &Array1<f64>) {
3751        (&self.lower_theta, &self.upper_theta)
3752    }
3753
3754    pub fn fixed_offsets(&self) -> (&Array1<f64>, &Array1<f64>, &Array1<f64>) {
3755        (
3756            &self.fixed_offset_entry,
3757            &self.fixed_offset_exit,
3758            &self.fixed_derivative_offset_exit,
3759        )
3760    }
3761
3762    pub fn evaluate_initial(&self) -> Result<SurvivalMarginalSlopeOffsetGeometry, String> {
3763        self.evaluate(&self.initial_theta)
3764    }
3765
3766    pub fn evaluate(
3767        &self,
3768        theta: &Array1<f64>,
3769    ) -> Result<SurvivalMarginalSlopeOffsetGeometry, String> {
3770        let config = survival_baseline_config_from_theta(self.target, theta)?;
3771        // θ-authored: this chart was ASKED to realize `theta`, so `theta` is
3772        // what the geometry records. Re-deriving it from `config` closes a
3773        // `ln(exp(·))` loop that is not the identity in `f64`, and the family's
3774        // `to_bits()` manifest check then refuses the point (#2765).
3775        let mut geometry = build_survival_marginal_slope_baseline_geometry_at_theta(
3776            &self.age_entry,
3777            &self.age_exit,
3778            &config,
3779            theta.clone(),
3780        )?
3781        .ok_or_else(|| {
3782            "survival marginal-slope nonlinear baseline chart lost its theta coordinates"
3783                .to_string()
3784        })?;
3785        geometry.offset_entry += &self.fixed_offset_entry;
3786        geometry.offset_exit += &self.fixed_offset_exit;
3787        geometry.derivative_offset_exit += &self.fixed_derivative_offset_exit;
3788        Ok(geometry)
3789    }
3790}
3791
3792pub fn location_scale_uses_probit_survival_baseline(inverse_link: Option<&InverseLink>) -> bool {
3793    matches!(
3794        inverse_link,
3795        Some(
3796            InverseLink::Standard(StandardLink::Probit)
3797                | InverseLink::LatentCLogLog(_)
3798                | InverseLink::Sas(_)
3799                | InverseLink::BetaLogistic(_)
3800                | InverseLink::Mixture(_)
3801        )
3802    )
3803}
3804
3805pub fn survival_derivative_guard_for_likelihood(likelihood_mode: SurvivalLikelihoodMode) -> f64 {
3806    match likelihood_mode {
3807        SurvivalLikelihoodMode::LocationScale
3808        | SurvivalLikelihoodMode::Latent
3809        | SurvivalLikelihoodMode::LatentBinary => DEFAULT_SURVIVAL_LOCATION_SCALE_DERIVATIVE_GUARD,
3810        SurvivalLikelihoodMode::MarginalSlope => DEFAULT_SURVIVAL_MARGINAL_SLOPE_DERIVATIVE_GUARD,
3811        SurvivalLikelihoodMode::Transformation | SurvivalLikelihoodMode::Weibull => 0.0,
3812    }
3813}
3814
3815/// Resolve the actual parametric offset chart used by a marginal-slope fit.
3816///
3817/// A nominal `Linear` target has zero derivative and therefore starts the
3818/// `-log(q')` barrier exactly on its guard.  Fitting consequently uses a
3819/// deterministic exponential-survival (Weibull shape one) offset at the
3820/// data-scale mean positive exit time.  This function is the shared authority
3821/// for fitting and persistence: saving the nominal `Linear` request would not
3822/// be enough to replay the fitted row likelihood.
3823pub fn survival_marginal_slope_offset_baseline_config(
3824    age_exit: &Array1<f64>,
3825    requested: &SurvivalBaselineConfig,
3826) -> SurvivalBaselineConfig {
3827    if requested.target == SurvivalBaselineTarget::Linear {
3828        SurvivalBaselineConfig {
3829            target: SurvivalBaselineTarget::Weibull,
3830            scale: Some(positive_survival_time_seed(age_exit)),
3831            shape: Some(1.0),
3832            rate: None,
3833            makeham: None,
3834        }
3835    } else {
3836        requested.clone()
3837    }
3838}
3839
3840pub fn build_survival_time_offsets_for_likelihood(
3841    age_entry: &Array1<f64>,
3842    age_exit: &Array1<f64>,
3843    baseline_cfg: &SurvivalBaselineConfig,
3844    likelihood_mode: SurvivalLikelihoodMode,
3845    inverse_link: Option<&InverseLink>,
3846) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), String> {
3847    if likelihood_mode == SurvivalLikelihoodMode::MarginalSlope
3848        || (likelihood_mode == SurvivalLikelihoodMode::LocationScale
3849            && location_scale_uses_probit_survival_baseline(inverse_link))
3850    {
3851        build_survival_marginal_slope_baseline_offsets(age_entry, age_exit, baseline_cfg)
3852    } else {
3853        build_survival_baseline_offsets(age_entry, age_exit, baseline_cfg)
3854    }
3855}
3856
3857pub fn add_survival_time_derivative_guard_offset(
3858    age_entry: &Array1<f64>,
3859    age_exit: &Array1<f64>,
3860    anchor_time: f64,
3861    derivative_guard: f64,
3862    eta_offset_entry: &mut Array1<f64>,
3863    eta_offset_exit: &mut Array1<f64>,
3864    derivative_offset_exit: &mut Array1<f64>,
3865) -> Result<(), String> {
3866    if derivative_guard <= 0.0 {
3867        return Ok(());
3868    }
3869    let n = age_entry.len();
3870    if age_exit.len() != n
3871        || eta_offset_entry.len() != n
3872        || eta_offset_exit.len() != n
3873        || derivative_offset_exit.len() != n
3874    {
3875        return Err(SurvivalConstructionError::IncompatibleDimensions {
3876            reason: "survival derivative-guard offset lengths must match".to_string(),
3877        }
3878        .into());
3879    }
3880    for i in 0..n {
3881        eta_offset_entry[i] += derivative_guard * (age_entry[i] - anchor_time);
3882        eta_offset_exit[i] += derivative_guard * (age_exit[i] - anchor_time);
3883        derivative_offset_exit[i] += derivative_guard;
3884    }
3885    Ok(())
3886}
3887
3888#[derive(Clone, Debug)]
3889pub struct LatentSurvivalBaselineOffsets {
3890    pub loaded_eta_entry: Array1<f64>,
3891    pub loaded_eta_exit: Array1<f64>,
3892    pub loaded_derivative_exit: Array1<f64>,
3893    pub unloaded_mass_entry: Array1<f64>,
3894    pub unloaded_mass_exit: Array1<f64>,
3895    pub unloaded_hazard_exit: Array1<f64>,
3896}
3897
3898pub fn build_latent_survival_baseline_offsets(
3899    age_entry: &Array1<f64>,
3900    age_exit: &Array1<f64>,
3901    cfg: &SurvivalBaselineConfig,
3902    loading: HazardLoading,
3903) -> Result<LatentSurvivalBaselineOffsets, String> {
3904    if age_entry.len() != age_exit.len() {
3905        return Err(
3906            "latent survival baseline offsets require matching entry/exit lengths".to_string(),
3907        );
3908    }
3909
3910    fn gompertz_components(age: f64, rate: f64, shape: f64) -> (f64, f64) {
3911        if shape.abs() < 1e-10 {
3912            // Taylor at shape=0 matching `gompertz_hazard_components`:
3913            //   H_G(t) = rate·t·(1 + (shape·t)/2 + (shape·t)²/6)
3914            //   h_G(t) = rate·(1 + shape·t + (shape·t)²/2)
3915            // Dropping the higher-order `shape*t` corrections silently
3916            // diverges this helper from its sibling for non-zero shape near
3917            // the cutoff and gives inconsistent loaded-vs-unloaded offsets.
3918            let x = shape * age;
3919            return (
3920                rate * age * (1.0 + 0.5 * x + x * x / 6.0),
3921                rate * (1.0 + x + 0.5 * x * x),
3922            );
3923        }
3924        let shape_age = shape * age;
3925        let cumulative_hazard = (rate / shape) * shape_age.exp_m1();
3926        let instant_hazard = rate * shape_age.exp();
3927        (cumulative_hazard, instant_hazard)
3928    }
3929
3930    let n = age_entry.len();
3931
3932    // Per-row 6-tuple is independent. Evaluate in parallel into a Vec and then
3933    // unpack into the six Array1 outputs in original order.
3934    let rows: Vec<[f64; 6]> = (0..n)
3935        .into_par_iter()
3936        .map(|i| -> Result<[f64; 6], String> {
3937            let entry = age_entry[i];
3938            let exit = age_exit[i];
3939            if !entry.is_finite()
3940                || !exit.is_finite()
3941                || entry <= 0.0
3942                || exit <= 0.0
3943                || exit < entry
3944            {
3945                return Err(format!(
3946                    "latent survival baseline offsets require finite positive entry/exit ages with exit >= entry (row {})",
3947                    i + 1
3948                ));
3949            }
3950            match loading {
3951                HazardLoading::Full => {
3952                    let (eta_entry, _) = evaluate_survival_baseline(entry, cfg)?;
3953                    let (eta_exit, derivative_exit) = evaluate_survival_baseline(exit, cfg)?;
3954                    Ok([eta_entry, eta_exit, derivative_exit, 0.0, 0.0, 0.0])
3955                }
3956                HazardLoading::LoadedVsUnloaded => {
3957                    if cfg.target != SurvivalBaselineTarget::GompertzMakeham {
3958                        return Err(format!(
3959                            "HazardLoading::LoadedVsUnloaded requires --baseline-target gompertz-makeham, got {}",
3960                            survival_baseline_targetname(cfg.target)
3961                        ));
3962                    }
3963                    let rate = cfg.rate.ok_or_else(|| {
3964                        "gompertz-makeham latent survival is missing baseline rate".to_string()
3965                    })?;
3966                    let shape = cfg.shape.ok_or_else(|| {
3967                        "gompertz-makeham latent survival is missing baseline shape".to_string()
3968                    })?;
3969                    let makeham = cfg.makeham.ok_or_else(|| {
3970                        "gompertz-makeham latent survival is missing baseline makeham".to_string()
3971                    })?;
3972                    let (loaded_entry, _) = gompertz_components(entry, rate, shape);
3973                    let (loaded_exit, loaded_hazard) = gompertz_components(exit, rate, shape);
3974                    if !(loaded_entry.is_finite()
3975                        && loaded_entry > 0.0
3976                        && loaded_exit.is_finite()
3977                        && loaded_exit > 0.0
3978                        && loaded_hazard.is_finite()
3979                        && loaded_hazard > 0.0)
3980                    {
3981                        return Err(format!(
3982                            "gompertz-makeham latent loaded component produced a non-positive or non-finite hazard decomposition at row {}",
3983                            i + 1
3984                        ));
3985                    }
3986                    Ok([
3987                        loaded_entry.ln(),
3988                        loaded_exit.ln(),
3989                        loaded_hazard / loaded_exit,
3990                        makeham * entry,
3991                        makeham * exit,
3992                        makeham,
3993                    ])
3994                }
3995            }
3996        })
3997        .collect::<Result<Vec<_>, String>>()?;
3998
3999    let mut loaded_eta_entry = Array1::<f64>::zeros(n);
4000    let mut loaded_eta_exit = Array1::<f64>::zeros(n);
4001    let mut loaded_derivative_exit = Array1::<f64>::zeros(n);
4002    let mut unloaded_mass_entry = Array1::<f64>::zeros(n);
4003    let mut unloaded_mass_exit = Array1::<f64>::zeros(n);
4004    let mut unloaded_hazard_exit = Array1::<f64>::zeros(n);
4005    for (i, row) in rows.into_iter().enumerate() {
4006        loaded_eta_entry[i] = row[0];
4007        loaded_eta_exit[i] = row[1];
4008        loaded_derivative_exit[i] = row[2];
4009        unloaded_mass_entry[i] = row[3];
4010        unloaded_mass_exit[i] = row[4];
4011        unloaded_hazard_exit[i] = row[5];
4012    }
4013
4014    Ok(LatentSurvivalBaselineOffsets {
4015        loaded_eta_entry,
4016        loaded_eta_exit,
4017        loaded_derivative_exit,
4018        unloaded_mass_entry,
4019        unloaded_mass_exit,
4020        unloaded_hazard_exit,
4021    })
4022}
4023
4024// ---------------------------------------------------------------------------
4025// Time wiggle construction
4026// ---------------------------------------------------------------------------
4027
4028pub fn build_survival_timewiggle_derivative_design(
4029    eta_exit: &Array1<f64>,
4030    derivative_exit: &Array1<f64>,
4031    knots: &Array1<f64>,
4032    degree: usize,
4033) -> Result<Array2<f64>, String> {
4034    let mut design_derivative_exit =
4035        monotone_wiggle_basis_with_derivative_order(eta_exit.view(), knots, degree, 1)?;
4036    for i in 0..design_derivative_exit.nrows() {
4037        let chain = derivative_exit[i];
4038        for j in 0..design_derivative_exit.ncols() {
4039            design_derivative_exit[[i, j]] *= chain;
4040        }
4041    }
4042    Ok(design_derivative_exit)
4043}
4044
4045/// Build the dynamic "baseline as prior" timewiggle runtime.
4046///
4047/// The baseline offsets are used only to initialize the wiggle knot placement
4048/// on a stable scalar scale.  The exact survival family evaluates the resulting
4049/// monotone wiggle dynamically on the current time predictor h0(t):
4050///
4051///   h(t) = g(h0(t)),   g(z) = z + w(z).
4052///
4053/// No fixed `B(eta_baseline)` design is constructed here.
4054pub fn build_survival_timewiggle_from_baseline(
4055    eta_entry: &Array1<f64>,
4056    eta_exit: &Array1<f64>,
4057    derivative_exit: &Array1<f64>,
4058    cfg: &LinkWiggleFormulaSpec,
4059) -> Result<SurvivalTimeWiggleBuild, String> {
4060    if eta_entry.len() != eta_exit.len() || eta_exit.len() != derivative_exit.len() {
4061        return Err(
4062            "baseline-timewiggle requires matching entry/exit/derivative lengths".to_string(),
4063        );
4064    }
4065    // Guard: if baseline offsets are all zero (linear baseline), the timewiggle
4066    // construction is degenerate — it adds only a constant, not time-varying structure.
4067    let all_zero = eta_entry.iter().all(|&v| v.abs() < 1e-15)
4068        && eta_exit.iter().all(|&v| v.abs() < 1e-15)
4069        && derivative_exit.iter().all(|&v| v.abs() < 1e-15);
4070    if all_zero {
4071        return Err(
4072            "timewiggle requires a non-linear scalar survival baseline target; \
4073             the provided baseline offsets are all zero (linear baseline)"
4074                .to_string(),
4075        );
4076    }
4077    let n = eta_exit.len();
4078    let mut seed = Array1::<f64>::zeros(2 * n);
4079    for i in 0..n {
4080        seed[i] = eta_entry[i];
4081        seed[n + i] = eta_exit[i];
4082    }
4083    // Use the smallest requested derivative order as the primary exact
4084    // function-space roughness so the fitted penalty system matches the public
4085    // formula exactly, including the slope (`order = 1`) case.
4086    let (primary_order, extra_orders) = split_wiggle_penalty_orders(2, &cfg.penalty_orders)?;
4087    let mut derivative_orders = Vec::with_capacity(1 + extra_orders.len());
4088    derivative_orders.push(primary_order);
4089    derivative_orders.extend(extra_orders);
4090    // A FIXED-index warp: `h = h_base + B(h_base)·β_w` composes onto the
4091    // BASELINE offsets, which are computed from time and do not move with β. So
4092    // the evaluation point never crosses a knot during a solve and the boundary
4093    // knot's multiplicity is invisible — this block keeps the clamped generator
4094    // rather than the simple-ended warp one (gam#2695).
4095    let knots = gam_terms::basis::initializewiggle_knots_from_seed(
4096        seed.view(),
4097        cfg.degree,
4098        cfg.num_internal_knots,
4099    )?;
4100    // One assembly for the WHOLE order list (gam#2647). The gauge-closure
4101    // coordinate is a property of the assembled set, so building the primary
4102    // order and appending the rest would decide it on a partial set.
4103    let combined_block = crate::wiggle::buildwiggle_block_input_from_orders(
4104        seed.view(),
4105        &knots,
4106        cfg.degree,
4107        &derivative_orders,
4108        cfg.double_penalty,
4109    )?;
4110    let ncols = combined_block.design.ncols();
4111    Ok(SurvivalTimeWiggleBuild {
4112        nullspace_dims: combined_block.nullspace_dims.clone(),
4113        penalties: {
4114            combined_block
4115                .penalties
4116                .into_iter()
4117                .map(|ps| ps.to_global(ncols))
4118                .collect()
4119        },
4120        knots,
4121        degree: cfg.degree,
4122        ncols,
4123    })
4124}
4125
4126pub fn append_zero_tail_columns(
4127    x_entry: &mut DesignMatrix,
4128    x_exit: &mut DesignMatrix,
4129    x_derivative: &mut DesignMatrix,
4130    tail_cols: usize,
4131) {
4132    if tail_cols == 0 {
4133        return;
4134    }
4135    // Wiggle tail columns are dense, so materialize everything to dense.
4136    // This only runs once at construction time when time-wiggles are active.
4137    fn append_dense(dm: &mut DesignMatrix, tail: usize) {
4138        let old = dm.to_dense();
4139        let n = old.nrows();
4140        let p_base = old.ncols();
4141        let mut out = Array2::<f64>::zeros((n, p_base + tail));
4142        out.slice_mut(s![.., 0..p_base]).assign(&old);
4143        *dm = DesignMatrix::Dense(DenseDesignMatrix::from(out));
4144    }
4145    append_dense(x_entry, tail_cols);
4146    append_dense(x_exit, tail_cols);
4147    append_dense(x_derivative, tail_cols);
4148}
4149
4150// ---------------------------------------------------------------------------
4151// Resolved config (from build output back to config for serialization)
4152// ---------------------------------------------------------------------------
4153
4154// ---------------------------------------------------------------------------
4155// Time-varying covariate template
4156// ---------------------------------------------------------------------------
4157
4158/// Build a time-varying covariate block by tensoring the covariate design
4159/// with a 1D B-spline basis on log(time).
4160pub fn build_time_varying_survival_covariate_template(
4161    age_entry: &Array1<f64>,
4162    age_exit: &Array1<f64>,
4163    time_k: usize,
4164    time_degree: usize,
4165    block_name: &str,
4166) -> Result<SurvivalCovariateTermBlockTemplate, String> {
4167    if time_k < time_degree + 1 {
4168        return Err(format!(
4169            "--{block_name}-time-k must be >= degree + 1 = {}, got {time_k}",
4170            time_degree + 1
4171        ));
4172    }
4173    let num_internal_knots = time_k - (time_degree + 1);
4174
4175    let log_exit = age_exit.mapv(|t| t.max(1e-12).ln());
4176
4177    let time_spec = BSplineBasisSpec {
4178        degree: time_degree,
4179        penalty_order: 2,
4180        knotspec: BSplineKnotSpec::Automatic {
4181            num_internal_knots: Some(num_internal_knots),
4182            placement: gam_terms::basis::BSplineKnotPlacement::Quantile,
4183        },
4184        double_penalty: false,
4185        identifiability: BSplineIdentifiability::None,
4186        boundary: OneDimensionalBoundary::Open,
4187        boundary_conditions: BSplineBoundaryConditions::default(),
4188    };
4189
4190    let time_build = build_bspline_basis_1d(log_exit.view(), &time_spec)
4191        .map_err(|e| format!("failed to build {block_name} time-margin B-spline basis: {e}"))?;
4192    let time_design_exit = time_build.design.to_dense();
4193
4194    let knots = match &time_build.metadata {
4195        BasisMetadata::BSpline1D { knots, .. } => knots.clone(),
4196        _ => {
4197            return Err(format!(
4198                "{block_name} time-margin basis returned unexpected metadata type"
4199            ));
4200        }
4201    };
4202
4203    let time_penalties = time_build
4204        .active_penalties
4205        .into_iter()
4206        .map(|penalty| penalty.matrix)
4207        .collect();
4208
4209    finish_time_varying_survival_covariate_template(
4210        age_entry,
4211        age_exit,
4212        time_degree,
4213        knots,
4214        time_design_exit,
4215        time_penalties,
4216        block_name,
4217    )
4218}
4219
4220/// Replay a fit-time threshold/log-scale time margin from its resolved knots.
4221/// Prediction and saved ALO use this path so the prediction sample can never
4222/// move the spline basis by re-estimating quantile knots.
4223pub fn replay_time_varying_survival_covariate_template(
4224    age_entry: &Array1<f64>,
4225    age_exit: &Array1<f64>,
4226    time_basis: &SurvivalCovariateTimeBasis,
4227    block_name: &str,
4228) -> Result<SurvivalCovariateTermBlockTemplate, String> {
4229    let log_exit = age_exit.mapv(|t| t.max(1e-12).ln());
4230    let knots = Array1::from_vec(time_basis.knots.clone());
4231    let time_build = build_bspline_basis_1d(
4232        log_exit.view(),
4233        &BSplineBasisSpec {
4234            degree: time_basis.degree,
4235            penalty_order: 2,
4236            knotspec: BSplineKnotSpec::Provided(knots.clone()),
4237            double_penalty: false,
4238            identifiability: BSplineIdentifiability::None,
4239            boundary: OneDimensionalBoundary::Open,
4240            boundary_conditions: BSplineBoundaryConditions::default(),
4241        },
4242    )
4243    .map_err(|e| format!("failed to replay {block_name} time-margin B-spline basis: {e}"))?;
4244    let time_design_exit = time_build.design.to_dense();
4245    let time_penalties = time_build
4246        .active_penalties
4247        .into_iter()
4248        .map(|penalty| penalty.matrix)
4249        .collect();
4250    finish_time_varying_survival_covariate_template(
4251        age_entry,
4252        age_exit,
4253        time_basis.degree,
4254        knots,
4255        time_design_exit,
4256        time_penalties,
4257        block_name,
4258    )
4259}
4260
4261/// The log-slope block's follow-up margin evaluated at arbitrary times
4262/// (gam#2765, gam#2767).
4263///
4264/// The margin is a B-spline in `log t` on the fit's own knots. At fit time the
4265/// knots are placed by quantile and the design is a by-product of that build; at
4266/// prediction time the knots are read back from the saved model and the SAME
4267/// spec is re-evaluated, so a prediction sample can never move the basis. This
4268/// is the only B-spline evaluation on the log-slope time axis, so the batch
4269/// replay below and the per-`(row, t)` survival-curve replay cannot disagree
4270/// about which basis they are asking for.
4271pub fn logslope_time_margin_rows(
4272    time_basis: &SurvivalCovariateTimeBasis,
4273    times: ndarray::ArrayView1<'_, f64>,
4274) -> Result<Array2<f64>, String> {
4275    let log_times = times.mapv(|t| t.max(1e-12).ln());
4276    let knots = Array1::from_vec(time_basis.knots.clone());
4277    let build = build_bspline_basis_1d(
4278        log_times.view(),
4279        &BSplineBasisSpec {
4280            degree: time_basis.degree,
4281            penalty_order: 2,
4282            knotspec: BSplineKnotSpec::Provided(knots),
4283            double_penalty: false,
4284            identifiability: BSplineIdentifiability::None,
4285            boundary: OneDimensionalBoundary::Open,
4286            boundary_conditions: BSplineBoundaryConditions::default(),
4287        },
4288    )
4289    .map_err(|e| format!("failed to replay the log-slope time margin: {e}"))?;
4290    Ok(build.design.to_dense())
4291}
4292
4293/// All three follow-up channels of a log-slope block, replayed from the saved
4294/// margin (gam#2765, gam#2767).
4295///
4296/// The row program reads the slope at three places — the row's entry time, its
4297/// exit time, and the exit-time rate — because the likelihood is
4298/// `log S(t₁) − log S(t₀)` and an event row also carries `log η′(t₁)`. Any
4299/// consumer that re-evaluates that program off a saved model (the leave-one-out
4300/// replay, above all) needs all three; handing it the exit design alone would
4301/// silently evaluate a time-CONSTANT slope, which is a different model.
4302pub struct LogslopeFollowUpReplayDesigns {
4303    pub entry: DesignMatrix,
4304    pub exit: DesignMatrix,
4305    pub derivative_exit: DesignMatrix,
4306}
4307
4308/// Replay every follow-up channel of a log-slope block from its saved margin.
4309pub fn replay_logslope_follow_up_designs(
4310    age_entry: &Array1<f64>,
4311    age_exit: &Array1<f64>,
4312    time_basis: &SurvivalCovariateTimeBasis,
4313    covariate_design: &DesignMatrix,
4314) -> Result<LogslopeFollowUpReplayDesigns, String> {
4315    let template = replay_time_varying_survival_covariate_template(
4316        age_entry, age_exit, time_basis, "logslope",
4317    )?;
4318    let SurvivalCovariateTermBlockTemplate::TimeVarying {
4319        time_basis_entry,
4320        time_basis_exit,
4321        time_basis_derivative_exit,
4322        ..
4323    } = &template
4324    else {
4325        return Err(
4326            "replaying a log-slope time margin produced a time-constant template".to_string(),
4327        );
4328    };
4329    if covariate_design.nrows() != time_basis_exit.nrows() {
4330        return Err(format!(
4331            "log-slope follow-up replay has {} covariate rows against {} time rows",
4332            covariate_design.nrows(),
4333            time_basis_exit.nrows(),
4334        ));
4335    }
4336    if covariate_design.ncols() == 0 || time_basis_exit.ncols() == 0 {
4337        return Err(format!(
4338            "a follow-up-varying log-slope needs a non-empty tensor product, got {}x{}",
4339            covariate_design.ncols(),
4340            time_basis_exit.ncols(),
4341        ));
4342    }
4343    let kron = |basis: &Array2<f64>| {
4344        crate::survival::location_scale::rowwise_kronecker(covariate_design, basis)
4345    };
4346    Ok(LogslopeFollowUpReplayDesigns {
4347        entry: kron(time_basis_entry),
4348        exit: kron(time_basis_exit),
4349        derivative_exit: kron(time_basis_derivative_exit),
4350    })
4351}
4352
4353/// The log-slope block's fitted design, rebuilt from the covariate factor and
4354/// the fit's own resolved time margin (gam#2765, gam#2767).
4355///
4356/// A follow-up-varying log-slope block does not own the covariate design its
4357/// term spec describes; it owns the row-wise Kronecker product
4358/// `X_cov ⊗ᵣ B(log t)`, evaluated at the time each row's slope is being read
4359/// at. At fit time that is the row's EXIT time, which is the convention the
4360/// block's `ParameterBlockSpec` eta already uses; at prediction time it is
4361/// whichever time the survival curve is being evaluated at, because `b(t)` moves
4362/// along the curve exactly as `q(t)` does.
4363///
4364/// Rebuilding it from the term spec alone — `p_cov` columns against a
4365/// `p_cov · p_time` coefficient vector — is the failure this function exists to
4366/// make impossible.
4367pub fn replay_logslope_time_margin_design(
4368    times: ndarray::ArrayView1<'_, f64>,
4369    time_basis: &SurvivalCovariateTimeBasis,
4370    covariate_design: &DesignMatrix,
4371) -> Result<DesignMatrix, String> {
4372    if covariate_design.nrows() != times.len() {
4373        return Err(format!(
4374            "log-slope time-margin replay has {} covariate rows against {} times",
4375            covariate_design.nrows(),
4376            times.len(),
4377        ));
4378    }
4379    let time_design = logslope_time_margin_rows(time_basis, times)?;
4380    if covariate_design.ncols() == 0 || time_design.ncols() == 0 {
4381        return Err(format!(
4382            "a follow-up-varying log-slope needs a non-empty tensor product, got {}x{}",
4383            covariate_design.ncols(),
4384            time_design.ncols(),
4385        ));
4386    }
4387    Ok(crate::survival::location_scale::rowwise_kronecker(
4388        covariate_design,
4389        &time_design,
4390    ))
4391}
4392
4393fn finish_time_varying_survival_covariate_template(
4394    age_entry: &Array1<f64>,
4395    age_exit: &Array1<f64>,
4396    time_degree: usize,
4397    knots: Array1<f64>,
4398    time_design_exit: Array2<f64>,
4399    time_penalties: Vec<Array2<f64>>,
4400    block_name: &str,
4401) -> Result<SurvivalCovariateTermBlockTemplate, String> {
4402    if age_entry.len() != age_exit.len() {
4403        return Err(format!(
4404            "{block_name} time-margin entry/exit row mismatch: {} versus {}",
4405            age_entry.len(),
4406            age_exit.len()
4407        ));
4408    }
4409    let log_entry = age_entry.mapv(|t| t.max(1e-12).ln());
4410    let log_exit = age_exit.mapv(|t| t.max(1e-12).ln());
4411    let time_build_entry = build_bspline_basis_1d(
4412        log_entry.view(),
4413        &BSplineBasisSpec {
4414            degree: time_degree,
4415            penalty_order: 2,
4416            knotspec: BSplineKnotSpec::Provided(knots.clone()),
4417            double_penalty: false,
4418            identifiability: BSplineIdentifiability::None,
4419            boundary: OneDimensionalBoundary::Open,
4420            boundary_conditions: BSplineBoundaryConditions::default(),
4421        },
4422    )
4423    .map_err(|e| format!("failed to evaluate {block_name} time-margin basis at entry: {e}"))?;
4424    let time_design_entry = time_build_entry.design.to_dense();
4425    let p_time = time_design_exit.ncols();
4426    if p_time == 0 {
4427        return Err(format!(
4428            "{block_name} time-margin basis resolved to zero columns"
4429        ));
4430    }
4431    let mut time_design_derivative_exit = Array2::<f64>::zeros((age_exit.len(), p_time));
4432    time_design_derivative_exit
4433        .as_slice_mut()
4434        .expect("zeros are contiguous")
4435        .par_chunks_mut(p_time)
4436        .enumerate()
4437        .try_for_each(|(i, row_out)| -> Result<(), String> {
4438            let mut deriv_buf = vec![0.0_f64; p_time];
4439            evaluate_bspline_derivative_scalar(
4440                log_exit[i],
4441                knots.view(),
4442                time_degree,
4443                &mut deriv_buf,
4444            )
4445            .map_err(|e| {
4446                format!("failed to evaluate {block_name} time-margin derivative basis: {e}")
4447            })?;
4448            let chain = 1.0 / age_exit[i].max(1e-12);
4449            for j in 0..p_time {
4450                row_out[j] = deriv_buf[j] * chain;
4451            }
4452            Ok(())
4453        })?;
4454
4455    Ok(SurvivalCovariateTermBlockTemplate::TimeVarying {
4456        time_basis: SurvivalCovariateTimeBasis {
4457            degree: time_degree,
4458            knots: knots.to_vec(),
4459        },
4460        time_basis_entry: time_design_entry,
4461        time_basis_exit: time_design_exit,
4462        time_basis_derivative_exit: time_design_derivative_exit,
4463        time_penalties,
4464    })
4465}
4466
4467#[cfg(test)]
4468mod tests {
4469    use super::{
4470        SURVIVAL_LIKELIHOOD_MODES, SURVIVAL_TIME_FLOOR, SurvivalBaselineConfig,
4471        SurvivalBaselineTarget, SurvivalLikelihoodMode, SurvivalMarginalSlopeFrozenOffsetChart,
4472        SurvivalTimeBasisConfig, baseline_chain_rule_gradient, baseline_offset_theta_partials,
4473        build_survival_marginal_slope_baseline_geometry,
4474        build_survival_marginal_slope_baseline_offsets, build_survival_time_basis,
4475        build_survival_timewiggle_from_baseline, evaluate_survival_baseline,
4476        evaluate_survival_marginal_slope_baseline, fitted_weibull_baseline_from_linear_time_beta,
4477        gompertz_cumulative_shape_derivative, gompertz_cumulative_shape_second_derivative,
4478        gompertz_hazard_components, marginal_slope_baseline_chain_rule_gradient,
4479        marginal_slope_baseline_chain_rule_hessian, marginal_slope_baseline_offset_theta_partials,
4480        optimize_survival_baseline_config_with_gradient,
4481        optimize_survival_baseline_config_with_gradient_only,
4482        resolve_survival_time_anchor_for_mode, survival_baseline_config_from_theta,
4483        survival_baseline_theta_from_config, survival_data_is_left_truncated,
4484        survival_earliest_entry_time_anchor, survival_robust_interior_time_anchor,
4485        validate_survival_time_anchor_override,
4486    };
4487    use super::{
4488        center_survival_time_designs_at_anchor, evaluate_survival_time_basis_row,
4489        resolved_survival_time_basis_config_from_build,
4490    };
4491    use super::{
4492        DesignMatrix, SurvivalCovariateTermBlockTemplate,
4493        build_time_varying_survival_covariate_template, logslope_time_margin_rows,
4494        replay_logslope_follow_up_designs, replay_logslope_time_margin_design,
4495    };
4496    use crate::probability::normal_cdf;
4497    use crate::survival::base::ENTRY_AT_ORIGIN_THRESHOLD;
4498    use crate::survival::{OffsetChannelCurvatures, OffsetChannelResiduals};
4499    use gam_terms::inference::formula_dsl::LinkWiggleFormulaSpec;
4500    use ndarray::{Array1, Array2, array};
4501
4502    #[test]
4503    fn fitted_weibull_baseline_uses_identified_anchor_and_slope() {
4504        // Single-column time basis (#2301): the shape is the sole coefficient.
4505        let fitted = fitted_weibull_baseline_from_linear_time_beta(&array![1.75], 4.5)
4506            .expect("valid Weibull baseline");
4507        assert_eq!(fitted.target, SurvivalBaselineTarget::Weibull);
4508        assert_eq!(fitted.scale, Some(4.5));
4509        assert_eq!(fitted.shape, Some(1.75));
4510        assert_eq!(fitted.rate, None);
4511        assert_eq!(fitted.makeham, None);
4512
4513        // Empty coefficient vector: no slope to recover.
4514        assert!(
4515            fitted_weibull_baseline_from_linear_time_beta(&Array1::<f64>::zeros(0), 4.5).is_none()
4516        );
4517        // Non-positive shape is not a valid Weibull baseline.
4518        assert!(fitted_weibull_baseline_from_linear_time_beta(&array![0.0], 4.5).is_none());
4519        // Non-positive anchor (scale) is invalid.
4520        assert!(fitted_weibull_baseline_from_linear_time_beta(&array![1.0], 0.0).is_none());
4521    }
4522
4523    #[test]
4524    fn survival_timewiggle_keeps_requested_order_one_penalty() {
4525        let eta_entry = array![0.1, 0.3, 0.5, 0.8];
4526        let eta_exit = array![0.4, 0.7, 1.0, 1.4];
4527        let derivative_exit = array![0.9, 1.1, 1.2, 1.3];
4528        let cfg = LinkWiggleFormulaSpec {
4529            degree: 3,
4530            num_internal_knots: 4,
4531            penalty_orders: vec![1, 2, 3],
4532            double_penalty: false,
4533        };
4534
4535        let build =
4536            build_survival_timewiggle_from_baseline(&eta_entry, &eta_exit, &derivative_exit, &cfg)
4537                .expect("build survival timewiggle");
4538
4539        assert_eq!(build.penalties.len(), 3);
4540        // Anchored I-spline value basis (#2306): the anchoring removes the
4541        // constant direction, so the order-m roughness nullity is m−1 — the
4542        // order-1 penalty is positive definite (nullity 0). The old [1, 2, 3]
4543        // encoded the unanchored convention.
4544        assert_eq!(build.nullspace_dims, vec![0, 1, 2]);
4545        assert!(build.ncols > 0);
4546    }
4547
4548    #[test]
4549    fn marginal_slope_frozen_offset_chart_moves_only_parametric_offsets() {
4550        let invalid_empty_config = SurvivalBaselineConfig {
4551            target: SurvivalBaselineTarget::Gompertz,
4552            scale: None,
4553            shape: Some(0.1),
4554            rate: Some(f64::NAN),
4555            makeham: None,
4556        };
4557        assert!(
4558            build_survival_marginal_slope_baseline_geometry(
4559                &Array1::zeros(0),
4560                &Array1::zeros(0),
4561                &invalid_empty_config,
4562            )
4563            .is_err(),
4564            "invalid baseline config must be rejected even with no rows"
4565        );
4566
4567        let age_entry = array![0.0, 0.75, 2.0];
4568        let age_exit = array![1.5, 3.0, 5.5];
4569        let initial_config = SurvivalBaselineConfig {
4570            target: SurvivalBaselineTarget::GompertzMakeham,
4571            scale: None,
4572            shape: Some(0.08),
4573            rate: Some(0.22),
4574            makeham: Some(0.04),
4575        };
4576        let initial_baseline =
4577            build_survival_marginal_slope_baseline_geometry(&age_entry, &age_exit, &initial_config)
4578                .expect("initial baseline geometry")
4579                .expect("nonlinear chart");
4580        let fixed_entry = array![0.125, -0.25, 0.375];
4581        let fixed_exit = array![-0.45, 0.55, 0.65];
4582        let fixed_derivative = array![0.015, 0.025, 0.035];
4583        let prepared_entry = &initial_baseline.offset_entry + &fixed_entry;
4584        let prepared_exit = &initial_baseline.offset_exit + &fixed_exit;
4585        let prepared_derivative = &initial_baseline.derivative_offset_exit + &fixed_derivative;
4586        let chart = SurvivalMarginalSlopeFrozenOffsetChart::new(
4587            &age_entry,
4588            &age_exit,
4589            &initial_config,
4590            &prepared_entry,
4591            &prepared_exit,
4592            &prepared_derivative,
4593        )
4594        .expect("freeze prepared offsets");
4595
4596        let initial = chart.evaluate_initial().expect("evaluate initial theta");
4597        for row in 0..age_exit.len() {
4598            assert!((initial.offset_entry[row] - prepared_entry[row]).abs() < 1e-14);
4599            assert!((initial.offset_exit[row] - prepared_exit[row]).abs() < 1e-14);
4600            assert!((initial.derivative_offset_exit[row] - prepared_derivative[row]).abs() < 1e-14);
4601        }
4602
4603        let mut candidate_theta = chart.initial_theta().clone();
4604        candidate_theta[0] += 0.3;
4605        candidate_theta[1] -= 0.025;
4606        candidate_theta[2] -= 0.2;
4607        let candidate = chart
4608            .evaluate(&candidate_theta)
4609            .expect("evaluate candidate theta");
4610        let candidate_baseline = build_survival_marginal_slope_baseline_geometry(
4611            &age_entry,
4612            &age_exit,
4613            &candidate.baseline_config,
4614        )
4615        .expect("candidate baseline geometry")
4616        .expect("nonlinear chart");
4617        let frozen = chart.fixed_offsets();
4618        for row in 0..age_exit.len() {
4619            assert!(
4620                (candidate.offset_entry[row]
4621                    - candidate_baseline.offset_entry[row]
4622                    - frozen.0[row])
4623                    .abs()
4624                    < 1e-14
4625            );
4626            assert!(
4627                (candidate.offset_exit[row] - candidate_baseline.offset_exit[row] - frozen.1[row])
4628                    .abs()
4629                    < 1e-14
4630            );
4631            assert!(
4632                (candidate.derivative_offset_exit[row]
4633                    - candidate_baseline.derivative_offset_exit[row]
4634                    - frozen.2[row])
4635                    .abs()
4636                    < 1e-14
4637            );
4638        }
4639        assert_eq!(
4640            candidate.offset_entry_theta_first,
4641            candidate_baseline.offset_entry_theta_first
4642        );
4643        assert_eq!(
4644            candidate.offset_exit_theta_first,
4645            candidate_baseline.offset_exit_theta_first
4646        );
4647        assert_eq!(
4648            candidate.derivative_offset_exit_theta_first,
4649            candidate_baseline.derivative_offset_exit_theta_first
4650        );
4651        assert_eq!(
4652            candidate.offset_entry_theta_second,
4653            candidate_baseline.offset_entry_theta_second
4654        );
4655        assert_eq!(
4656            candidate.offset_exit_theta_second,
4657            candidate_baseline.offset_exit_theta_second
4658        );
4659        assert_eq!(
4660            candidate.derivative_offset_exit_theta_second,
4661            candidate_baseline.derivative_offset_exit_theta_second
4662        );
4663        assert_eq!(candidate_baseline.offset_entry[0], 0.0);
4664        assert_eq!(
4665            candidate_baseline.offset_entry_theta_first.row(0).sum(),
4666            0.0
4667        );
4668        assert_eq!(
4669            candidate_baseline
4670                .offset_entry_theta_second
4671                .index_axis(ndarray::Axis(0), 0)
4672                .sum(),
4673            0.0
4674        );
4675        assert!(
4676            candidate
4677                .offset_entry_theta_first
4678                .row(0)
4679                .iter()
4680                .all(|value| *value == 0.0)
4681        );
4682        assert!(
4683            candidate
4684                .offset_entry_theta_second
4685                .index_axis(ndarray::Axis(0), 0)
4686                .iter()
4687                .all(|value| *value == 0.0)
4688        );
4689        for row in 0..age_exit.len() {
4690            for axis in 0..candidate_theta.len() {
4691                for other_axis in 0..candidate_theta.len() {
4692                    assert_eq!(
4693                        candidate.offset_entry_theta_second[[row, axis, other_axis]],
4694                        candidate.offset_entry_theta_second[[row, other_axis, axis]],
4695                    );
4696                    assert_eq!(
4697                        candidate.offset_exit_theta_second[[row, axis, other_axis]],
4698                        candidate.offset_exit_theta_second[[row, other_axis, axis]],
4699                    );
4700                    assert_eq!(
4701                        candidate.derivative_offset_exit_theta_second[[row, axis, other_axis]],
4702                        candidate.derivative_offset_exit_theta_second[[row, other_axis, axis]],
4703                    );
4704                }
4705            }
4706        }
4707    }
4708
4709    /// The one anchor rule (#2631), exercised across every mode and every
4710    /// truncation shape. Before the unification this behavior was spread over
4711    /// three resolvers and two front-end copies of the mode dispatch, and the
4712    /// copies disagreed.
4713    ///
4714    /// Marginal-slope takes the robust interior (median exit) anchor
4715    /// unconditionally: its `γ = 0` monotone-cone seed is where #751 was
4716    /// measured.
4717    #[test]
4718    fn marginal_slope_time_anchor_defaults_to_median_exit() {
4719        let age_entry = array![9.0, 1.0, 4.0, 6.0];
4720        let age_exit = array![20.0, 12.0, 18.0, 30.0];
4721        let anchor = resolve_survival_time_anchor_for_mode(
4722            SurvivalLikelihoodMode::MarginalSlope,
4723            &age_entry,
4724            &age_exit,
4725            None,
4726        )
4727        .expect("resolve marginal-slope default time anchor");
4728
4729        // Even count: mean of the two central exits, 18 and 20.
4730        assert!(
4731            (anchor - 19.0).abs() <= 1e-12,
4732            "marginal-slope default anchor should be median exit, got {anchor}"
4733        );
4734    }
4735
4736    /// An explicit anchor is the caller overriding the conditioning heuristic on
4737    /// purpose, so it wins in every mode — including the modes whose default
4738    /// would have been the median exit.
4739    #[test]
4740    fn explicit_time_anchor_wins_in_every_mode() {
4741        let age_entry = array![9.0, 1.0, 4.0, 6.0];
4742        let age_exit = array![20.0, 12.0, 18.0, 30.0];
4743        for mode in SURVIVAL_LIKELIHOOD_MODES {
4744            let anchor =
4745                resolve_survival_time_anchor_for_mode(mode, &age_entry, &age_exit, Some(7.5))
4746                    .expect("resolve explicit time anchor");
4747            assert!(
4748                (anchor - 7.5).abs() <= 1e-12,
4749                "explicit anchor must round-trip for {mode:?}, got {anchor}"
4750            );
4751        }
4752    }
4753
4754    /// Ordinary right-censored data (`entry == 0`) keeps the earliest-entry
4755    /// anchor in every non-marginal-slope mode, so centering stays the near-no-op
4756    /// it was before #751 and prior behavior is preserved bit-for-bit.
4757    #[test]
4758    fn right_censored_data_keeps_the_earliest_entry_anchor() {
4759        let age_entry = Array1::<f64>::zeros(4);
4760        let age_exit = array![20.0, 12.0, 18.0, 30.0];
4761        for mode in SURVIVAL_LIKELIHOOD_MODES {
4762            if mode == SurvivalLikelihoodMode::MarginalSlope {
4763                continue;
4764            }
4765            let anchor = resolve_survival_time_anchor_for_mode(mode, &age_entry, &age_exit, None)
4766                .expect("resolve right-censored time anchor");
4767            assert!(
4768                (anchor - SURVIVAL_TIME_FLOOR).abs() <= 1e-18,
4769                "un-truncated {mode:?} must anchor at the earliest entry (floored), got {anchor}"
4770            );
4771        }
4772    }
4773
4774    /// Genuine left truncation promotes the robust interior anchor for EVERY
4775    /// time-basis-carrying likelihood, not just marginal-slope. This is the rule
4776    /// the CLI's own copy did not implement, which is why the same data, formula
4777    /// and config produced a different location-scale fit on each front end
4778    /// (#2631).
4779    #[test]
4780    fn left_truncated_data_takes_the_robust_interior_anchor_in_every_mode() {
4781        let age_entry = array![9.0, 1.0, 4.0, 6.0];
4782        let age_exit = array![20.0, 12.0, 18.0, 30.0];
4783        assert!(survival_data_is_left_truncated(&age_entry));
4784        for mode in SURVIVAL_LIKELIHOOD_MODES {
4785            let anchor = resolve_survival_time_anchor_for_mode(mode, &age_entry, &age_exit, None)
4786                .expect("resolve left-truncated time anchor");
4787            assert!(
4788                (anchor - 19.0).abs() <= 1e-12,
4789                "left-truncated {mode:?} must anchor at the median exit (19.0), got {anchor}; \
4790                 the earliest entry (1.0) is the #751/#1790 defect"
4791            );
4792        }
4793    }
4794
4795    /// Staggered entry — part of the cohort followed from the time origin, the
4796    /// rest joining later — IS left truncation. The predicate is `any(entry >
4797    /// threshold)`, not `min(entry) > threshold`: with a `min` test this shape
4798    /// would fall back to an anchor at the time-origin floor, where the centered
4799    /// exit columns are maximally large and one-signed. The retired
4800    /// transformation-specific resolver used `min` and got exactly this case
4801    /// wrong.
4802    #[test]
4803    fn staggered_entry_counts_as_left_truncated() {
4804        let age_entry = array![0.0, 9.0, 0.0, 6.0];
4805        let age_exit = array![20.0, 12.0, 18.0, 30.0];
4806        assert!(
4807            survival_data_is_left_truncated(&age_entry),
4808            "a cohort with some rows entering at positive delayed-entry times is left-truncated"
4809        );
4810        for mode in SURVIVAL_LIKELIHOOD_MODES {
4811            let anchor = resolve_survival_time_anchor_for_mode(mode, &age_entry, &age_exit, None)
4812                .expect("resolve staggered-entry time anchor");
4813            assert!(
4814                (anchor - 19.0).abs() <= 1e-12,
4815                "staggered-entry {mode:?} must anchor at the median exit, got {anchor}"
4816            );
4817        }
4818    }
4819
4820    /// The origin convention is the likelihood engines' own: an entry exactly at
4821    /// the threshold is still "at the origin", one above it is delayed entry.
4822    #[test]
4823    fn left_truncation_predicate_uses_the_engine_origin_threshold() {
4824        assert!(!survival_data_is_left_truncated(&array![
4825            0.0,
4826            ENTRY_AT_ORIGIN_THRESHOLD
4827        ]));
4828        assert!(survival_data_is_left_truncated(&array![
4829            0.0,
4830            ENTRY_AT_ORIGIN_THRESHOLD * 1.000_001
4831        ]));
4832    }
4833
4834    /// Odd row counts take the true middle exit; the anchor is always floored so
4835    /// `log(anchor)` stays finite.
4836    #[test]
4837    fn robust_interior_anchor_is_the_median_and_is_floored() {
4838        assert!(
4839            (survival_robust_interior_time_anchor(&array![30.0, 12.0, 18.0])
4840                .expect("odd-count median")
4841                - 18.0)
4842                .abs()
4843                <= 1e-12
4844        );
4845        assert_eq!(
4846            survival_robust_interior_time_anchor(&array![0.0, 0.0]).expect("zero exits"),
4847            SURVIVAL_TIME_FLOOR
4848        );
4849        assert!(survival_robust_interior_time_anchor(&Array1::<f64>::zeros(0)).is_err());
4850        assert!(survival_earliest_entry_time_anchor(&Array1::<f64>::zeros(0)).is_err());
4851    }
4852
4853    /// The MECHANISM behind the rule, measured rather than asserted (#751/#1790,
4854    /// #2631).
4855    ///
4856    /// The robust interior anchor exists because centering a left-truncated
4857    /// design at the earliest entry leaves the exit columns large and ONE-SIGNED
4858    /// — that column is the unpenalized polynomial null space of the
4859    /// 2nd-difference time penalty, so its inflation multiplies the time-block
4860    /// score at the smoothing seed. This measures both properties directly on a
4861    /// staggered-entry cohort (half the rows entering at the time origin, half at
4862    /// positive delayed-entry times — the ordinary shape of a real registry
4863    /// cohort, and the case a `min(entry) > threshold` predicate would have
4864    /// missed):
4865    ///
4866    ///   * column magnitude `max |X_exit − X(anchor)|`
4867    ///   * sign balance: the fraction of rows sharing the dominant sign, per
4868    ///     column, worst column reported
4869    ///
4870    /// The earliest-entry anchor must be strictly worse on both, or the rule this
4871    /// module implements has no reason to exist.
4872    #[test]
4873    fn robust_interior_anchor_shrinks_and_balances_the_centered_time_design() {
4874        // Staggered entry: rows 0..5 observed from the origin, rows 6..11 with
4875        // positive delayed entry. Exits spread over a decade of time.
4876        let age_entry = array![
4877            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 5.0, 9.0, 14.0, 22.0, 30.0
4878        ];
4879        let age_exit = array![
4880            2.0, 4.0, 7.0, 11.0, 16.0, 23.0, 31.0, 40.0, 52.0, 68.0, 85.0, 110.0
4881        ];
4882        assert!(survival_data_is_left_truncated(&age_entry));
4883
4884        let earliest = survival_earliest_entry_time_anchor(&age_entry).expect("earliest anchor");
4885        let robust = survival_robust_interior_time_anchor(&age_exit).expect("robust anchor");
4886
4887        // Measure both anchorings on the SAME basis (knots inferred once from the
4888        // same times), so only the centering differs.
4889        let build = build_survival_time_basis(
4890            &age_entry,
4891            &age_exit,
4892            SurvivalTimeBasisConfig::ISpline {
4893                degree: 3,
4894                knots: Array1::zeros(0),
4895                keep_cols: Vec::new(),
4896                smooth_lambda: 1e-2,
4897            },
4898            Some((4, 1e-2)),
4899        )
4900        .expect("build survival time basis");
4901        let resolved = resolved_survival_time_basis_config_from_build(
4902            &build.basisname,
4903            build.degree,
4904            build.knots.as_ref(),
4905            build.keep_cols.as_ref(),
4906            build.smooth_lambda,
4907        )
4908        .expect("resolve time basis config");
4909
4910        // The quantity #1790 names is the centered design's component along the
4911        // TREND direction — the unpenalized null space of the 2nd-difference time
4912        // penalty — not any single raw column (an I-spline's last column
4913        // saturates at 1 over the observed range and is one-signed under every
4914        // anchor). For a monotone I-spline basis the row sum `Σ_j X_ij` is a
4915        // monotone increasing function of `t_i`, so it IS that trend coordinate,
4916        // and it is what "large and one-signed across all rows" is a statement
4917        // about.
4918        //
4919        // Returns `(max |row sum|, fraction of rows sharing the dominant sign)`.
4920        let measure = |anchor: f64| -> (f64, f64) {
4921            let mut centered = build.clone();
4922            let anchor_row =
4923                evaluate_survival_time_basis_row(anchor, &resolved).expect("anchor basis row");
4924            center_survival_time_designs_at_anchor(
4925                &mut centered.x_entry_time,
4926                &mut centered.x_exit_time,
4927                &anchor_row,
4928            )
4929            .expect("center at anchor");
4930            let dense = centered.x_exit_time.to_dense();
4931            let trend: Vec<f64> = dense.rows().into_iter().map(|row| row.sum()).collect();
4932            let magnitude = trend.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
4933            let positive = trend.iter().filter(|v| **v > 0.0).count() as f64;
4934            let rows = trend.len() as f64;
4935            let sign_fraction = (positive / rows).max(1.0 - positive / rows);
4936            eprintln!(
4937                "anchor {anchor:>10.4}: max|trend| = {magnitude:.6e}, \
4938                 dominant-sign fraction = {sign_fraction:.4}, trend = {trend:?}"
4939            );
4940            (magnitude, sign_fraction)
4941        };
4942
4943        let (earliest_magnitude, earliest_sign) = measure(earliest);
4944        let (robust_magnitude, robust_sign) = measure(robust);
4945
4946        // Measured on this fixture: earliest-entry anchor (the time-origin floor)
4947        // gives max|trend| = 5.000 with EVERY row positive; median-exit anchor
4948        // (27.0) gives 1.140 with a 6/6 sign split. The thresholds below are
4949        // loose around those numbers — they pin the phenomenon, not the digits.
4950        assert!(
4951            robust_magnitude < 0.5 * earliest_magnitude,
4952            "the robust interior anchor must materially shrink the centered trend \
4953             coordinate: earliest-entry anchor {earliest} gives max|trend| = \
4954             {earliest_magnitude}, median-exit anchor {robust} gives {robust_magnitude}"
4955        );
4956        assert!(
4957            (earliest_sign - 1.0).abs() <= 1e-12,
4958            "the earliest-entry anchor is expected to leave the trend coordinate \
4959             FULLY one-signed on left-truncated data — that is the #751/#1790 \
4960             mechanism, and a fixture where it does not hold is not exercising the \
4961             rule. Measured {earliest_sign}"
4962        );
4963        assert!(
4964            robust_sign <= 0.6,
4965            "the robust interior anchor must leave the trend coordinate two-signed \
4966             (the exit-event likelihood then pins the linear trend); measured \
4967             {robust_sign} of rows sharing one sign"
4968        );
4969    }
4970
4971    /// `SURVIVAL_LIKELIHOOD_MODES` must list every variant, or the cross-mode
4972    /// contracts above silently stop covering one. The `match` is what enforces
4973    /// it: adding a variant to the enum breaks this compilation until the new
4974    /// mode is added to the array too.
4975    #[test]
4976    fn survival_likelihood_modes_is_exhaustive() {
4977        fn slot(mode: SurvivalLikelihoodMode) -> usize {
4978            match mode {
4979                SurvivalLikelihoodMode::Transformation => 0,
4980                SurvivalLikelihoodMode::Weibull => 1,
4981                SurvivalLikelihoodMode::LocationScale => 2,
4982                SurvivalLikelihoodMode::MarginalSlope => 3,
4983                SurvivalLikelihoodMode::Latent => 4,
4984                SurvivalLikelihoodMode::LatentBinary => 5,
4985            }
4986        }
4987        let mut seen = [false; 6];
4988        for mode in SURVIVAL_LIKELIHOOD_MODES {
4989            let slot = slot(mode);
4990            assert!(!seen[slot], "{mode:?} listed twice");
4991            seen[slot] = true;
4992        }
4993        assert!(
4994            seen.iter().all(|&hit| hit),
4995            "SURVIVAL_LIKELIHOOD_MODES is missing a variant: {seen:?}"
4996        );
4997    }
4998
4999    /// A caller-supplied anchor is validated once, in one place, so every front
5000    /// end refuses the same values with the same message.
5001    #[test]
5002    fn time_anchor_override_rejects_non_finite_and_negative_values() {
5003        for bad in [-1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
5004            assert!(
5005                validate_survival_time_anchor_override(bad).is_err(),
5006                "override {bad} must be refused"
5007            );
5008            assert!(
5009                resolve_survival_time_anchor_for_mode(
5010                    SurvivalLikelihoodMode::Transformation,
5011                    &array![0.0, 1.0],
5012                    &array![5.0, 6.0],
5013                    Some(bad),
5014                )
5015                .is_err(),
5016                "the rule must refuse override {bad}"
5017            );
5018        }
5019        assert_eq!(
5020            validate_survival_time_anchor_override(0.0).expect("zero is a legal anchor"),
5021            SURVIVAL_TIME_FLOOR
5022        );
5023    }
5024
5025    /// Derivative-contract parity for the two public baseline optimizers.
5026    ///
5027    /// After the unification onto `run_baseline_theta_optimizer`, the
5028    /// gradient-only and gradient+Hessian entry points differ *only* in how
5029    /// much derivative information they hand the outer solver — not in the
5030    /// surface they minimize. We exercise that invariant on a known
5031    /// strictly-convex quadratic in θ-space (Weibull baseline: θ = (ln scale,
5032    /// ln shape)) whose unique minimizer is `theta_star`, supplying the same
5033    /// objective as `(f, ∇f)` and as `(f, ∇f, ∇²f)`. Both contracts must
5034    /// recover the same minimizer config, not weakened to pass.
5035    #[test]
5036    fn baseline_optimizer_contracts_agree_on_shared_surface() {
5037        // SPD curvature and interior minimizer in θ-space. A is well away from
5038        // singular so both the analytic-Hessian and BFGS paths see the same
5039        // unambiguous bowl; θ* sits comfortably inside the ±6 box around the
5040        // θ=(0,0) seed below.
5041        let curvature: Array2<f64> = array![[3.0, 0.5], [0.5, 2.0]];
5042        let theta_star: Array1<f64> = array![2.5_f64.ln(), 1.3_f64.ln()];
5043
5044        // Seed config at θ=(0,0) (scale=shape=1). The Linear early-return path
5045        // is not exercised here; Weibull has a genuine 2-dim θ to optimize.
5046        let initial = SurvivalBaselineConfig {
5047            target: SurvivalBaselineTarget::Weibull,
5048            scale: Some(1.0),
5049            shape: Some(1.0),
5050            rate: None,
5051            makeham: None,
5052        };
5053
5054        // θ recovered from a returned Weibull config, via the exact inverse of
5055        // the config→θ map the optimizers use internally.
5056        let recovered_theta = |cfg: &SurvivalBaselineConfig| -> Array1<f64> {
5057            survival_baseline_theta_from_config(cfg)
5058                .expect("config→θ")
5059                .expect("Weibull config has a θ")
5060        };
5061
5062        // Shared quadratic surface, evaluated by mapping config→θ so every
5063        // contract sees the identical objective.
5064        let curvature_cost = curvature.clone();
5065        let star_cost = theta_star.clone();
5066        let cost_at = move |cfg: &SurvivalBaselineConfig| -> Result<f64, String> {
5067            let theta = survival_baseline_theta_from_config(cfg)?
5068                .ok_or_else(|| "expected a θ for the cost surface".to_string())?;
5069            let d = &theta - &star_cost;
5070            let ad = curvature_cost.dot(&d);
5071            Ok(0.5 * d.dot(&ad))
5072        };
5073
5074        let curvature_grad = curvature.clone();
5075        let star_grad = theta_star.clone();
5076        let cost_for_grad = cost_at.clone();
5077        let result_grad_only = optimize_survival_baseline_config_with_gradient_only(
5078            &initial,
5079            "baseline parity (gradient-only)",
5080            move |cfg| {
5081                let cost = cost_for_grad(cfg)?;
5082                let theta = survival_baseline_theta_from_config(cfg)?
5083                    .ok_or_else(|| "expected a θ for the gradient".to_string())?;
5084                let gradient = curvature_grad.dot(&(&theta - &star_grad));
5085                Ok((cost, gradient))
5086            },
5087        )
5088        .expect("gradient-only baseline optimization converges");
5089
5090        let curvature_hess = curvature.clone();
5091        let star_hess = theta_star.clone();
5092        let cost_for_hess = cost_at.clone();
5093        let result_grad_hess = optimize_survival_baseline_config_with_gradient(
5094            &initial,
5095            "baseline parity (gradient+Hessian)",
5096            move |cfg| {
5097                let cost = cost_for_hess(cfg)?;
5098                let theta = survival_baseline_theta_from_config(cfg)?
5099                    .ok_or_else(|| "expected a θ for the gradient".to_string())?;
5100                let gradient = curvature_hess.dot(&(&theta - &star_hess));
5101                Ok((cost, gradient, curvature_hess.clone()))
5102            },
5103        )
5104        .expect("gradient+Hessian baseline optimization converges");
5105
5106        let theta_grad_only = recovered_theta(&result_grad_only);
5107        let theta_grad_hess = recovered_theta(&result_grad_hess);
5108
5109        // Each contract recovers the true minimizer. 2e-3 is a safe,
5110        // un-weakened bound; both gradient paths land far tighter.
5111        for (label, theta) in [
5112            ("gradient-only", &theta_grad_only),
5113            ("gradient+Hessian", &theta_grad_hess),
5114        ] {
5115            let err = (theta - &theta_star)
5116                .mapv(f64::abs)
5117                .fold(0.0_f64, |a, &v| a.max(v));
5118            assert!(
5119                err <= 2e-3,
5120                "{label} contract recovered θ {theta:?} off true minimizer {theta_star:?} by {err:e}"
5121            );
5122        }
5123
5124        // Cross-contract agreement: the three results must coincide, since the
5125        // only difference between the entry points is the derivative contract,
5126        // never the surface they minimize.
5127        let pairwise_max = |a: &Array1<f64>, b: &Array1<f64>| -> f64 {
5128            (a - b).mapv(f64::abs).fold(0.0_f64, |acc, &v| acc.max(v))
5129        };
5130        assert!(
5131            pairwise_max(&theta_grad_only, &theta_grad_hess) <= 2e-3,
5132            "gradient-only vs gradient+Hessian disagree: {theta_grad_only:?} vs {theta_grad_hess:?}"
5133        );
5134    }
5135
5136    #[test]
5137    fn automatic_ispline_time_knots_are_sized_for_antiderivative_degree() {
5138        let age_entry = array![1.0_f64, 1.0, 1.0, 1.0, 1.0, 1.0];
5139        let age_exit = array![2.0_f64, 3.0, 5.0, 8.0, 13.0, 21.0];
5140        let requested_degree = 3;
5141        let num_internal_knots = 1;
5142
5143        let built = build_survival_time_basis(
5144            &age_entry,
5145            &age_exit,
5146            SurvivalTimeBasisConfig::ISpline {
5147                degree: requested_degree,
5148                knots: Array1::zeros(0),
5149                keep_cols: Vec::new(),
5150                smooth_lambda: 1e-2,
5151            },
5152            Some((num_internal_knots, 1e-2)),
5153        )
5154        .expect("automatic cubic ispline with one interior knot builds");
5155
5156        let working_degree = requested_degree + 1;
5157        let knots = built.knots.expect("resolved ispline knots");
5158        assert_eq!(
5159            knots.len(),
5160            num_internal_knots + 2 * (working_degree + 1),
5161            "I-spline automatic knots must be clamped for the working B-spline degree"
5162        );
5163        assert_eq!(built.degree, Some(requested_degree));
5164        assert!(built.x_exit_time.ncols() > 0);
5165        assert_eq!(built.x_entry_time.ncols(), built.x_exit_time.ncols());
5166        assert_eq!(built.x_derivative_time.ncols(), built.x_exit_time.ncols());
5167    }
5168
5169    #[test]
5170    fn linear_weibull_time_basis_is_a_single_log_t_column() {
5171        // #2301 dropped the redundant `[1, ·]` constant column from the linear
5172        // Weibull time basis: it was EXACTLY confounded with the covariate
5173        // intercept (which absorbs the Weibull location) and made the converged
5174        // penalized Hessian singular. The surviving basis is the single `log t`
5175        // slope column. This pins the emitted width so a re-added constant column
5176        // (or any consumer drifting back to the old 2-column layout) is caught at
5177        // construction — it is the width every downstream consumer agrees on:
5178        // `fit_orchestration::fit` slices `beta[..x_exit_time.ncols()]` and hands
5179        // it to `fitted_weibull_baseline_from_linear_time_beta`, which reads the
5180        // shape from `beta[0]`, and `evaluate_survival_time_basis_row` emits a
5181        // matching one-element `[log t]` row.
5182        let age_entry = array![1.0_f64, 1.0, 1.0, 1.0];
5183        let age_exit = array![2.0_f64, 3.0, 5.0, 8.0];
5184
5185        let built =
5186            build_survival_time_basis(&age_entry, &age_exit, SurvivalTimeBasisConfig::Linear, None)
5187                .expect("build linear Weibull time basis");
5188
5189        assert_eq!(
5190            built.x_exit_time.ncols(),
5191            1,
5192            "the linear Weibull time basis must emit exactly one column (`log t`); \
5193             the confounded constant column was dropped in #2301"
5194        );
5195        assert_eq!(
5196            built.x_entry_time.ncols(),
5197            1,
5198            "entry basis width must match"
5199        );
5200        assert_eq!(
5201            built.x_derivative_time.ncols(),
5202            1,
5203            "derivative basis width must match"
5204        );
5205        assert_eq!(built.basisname, "linear");
5206        assert!(
5207            built.penalties.is_empty(),
5208            "the linear parametric time block is unpenalized"
5209        );
5210
5211        // The sole column is `log t` at the exit times.
5212        let exit = built.x_exit_time.as_dense_cow();
5213        for (i, &t) in age_exit.iter().enumerate() {
5214            assert!(
5215                (exit[[i, 0]] - t.ln()).abs() < 1e-12,
5216                "exit column must carry log t: row {i} got {} want {}",
5217                exit[[i, 0]],
5218                t.ln()
5219            );
5220        }
5221
5222        // The frozen anchor-row evaluator agrees on the one-column width so the
5223        // engine's centered `(b(t) − b(anchor))·β_time` reconstruction lines up.
5224        let anchor_row =
5225            super::evaluate_survival_time_basis_row(4.5, &SurvivalTimeBasisConfig::Linear)
5226                .expect("evaluate linear anchor row");
5227        assert_eq!(anchor_row.len(), 1, "linear anchor row must be one element");
5228        assert!((anchor_row[0] - 4.5_f64.ln()).abs() < 1e-12);
5229    }
5230
5231    #[test]
5232    fn ispline_time_derivative_is_nonzero_at_right_boundary() {
5233        let age_entry = array![1.0_f64, 1.0, 1.0];
5234        let age_exit = array![4.0_f64, 4.0, 4.0];
5235        let left = 1.0_f64.ln();
5236        let right = 4.0_f64.ln();
5237        let mid = left + 0.5 * (right - left);
5238        let knots = array![left, left, left, left, mid, right, right, right, right];
5239
5240        let built = build_survival_time_basis(
5241            &age_entry,
5242            &age_exit,
5243            SurvivalTimeBasisConfig::ISpline {
5244                degree: 2,
5245                knots,
5246                keep_cols: Vec::new(),
5247                smooth_lambda: 1e-2,
5248            },
5249            None,
5250        )
5251        .expect("build right-boundary ispline time basis");
5252
5253        let derivative = built.x_derivative_time.as_dense_cow();
5254        let max_abs = derivative.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
5255        assert!(
5256            max_abs > 1e-8,
5257            "right-boundary I-spline derivative must use the left-hand endpoint slope"
5258        );
5259        for row in derivative.rows() {
5260            assert!(
5261                row.iter().any(|v| *v > 1e-8),
5262                "each row at the right boundary needs a positive hazard derivative"
5263            );
5264        }
5265    }
5266
5267    #[test]
5268    fn ispline_time_penalty_is_psd_under_nontrivial_keep_cols() {
5269        // PSD-invariant forward guard for the gam#979 survival hang. The I-spline
5270        // value-space curvature penalty on the increment coefficients is the
5271        // congruence `S_I = Lᵀ S_B[1:,1:] L`. When identifiability drops columns,
5272        // the retained block MUST be taken as a PRINCIPAL SUBMATRIX of the FULL
5273        // congruence (congruence first, column selection second). The historical
5274        // regression assembled the reduced penalty in the wrong order, producing
5275        // a strongly INDEFINITE matrix (measured `s0_min_eval = −9.8e7`); an
5276        // indefinite time penalty makes `½γᵀ S_I γ` unbounded below, the inner
5277        // joint-Newton follows the divergence, and the outer REML never
5278        // terminates — the survival marginal-slope hang.
5279        //
5280        // This test exercises the reduction with a NON-TRIVIAL `keep_cols`
5281        // (a proper subset, an interior column dropped) and asserts the assembled
5282        // penalty satisfies the PSD contract the fix guarantees. It locks the
5283        // invariant on the shipped code path so a future reassembly that
5284        // reintroduces an indefinite reduction is caught at construction rather
5285        // than silently as an outer-loop hang. (It is a forward invariant lock,
5286        // not a bit-exact replay of the removed buggy assembly.)
5287        let age_entry = array![1.0_f64, 1.0, 1.0, 1.0, 1.0, 1.0];
5288        let age_exit = array![2.0_f64, 3.0, 5.0, 8.0, 13.0, 21.0];
5289        let left = 1.0_f64.ln();
5290        let right = 21.0_f64.ln();
5291        let q1 = left + 0.25 * (right - left);
5292        let mid = left + 0.5 * (right - left);
5293        let q3 = left + 0.75 * (right - left);
5294        // Degree-2 I-spline with three interior knots -> a value-space basis wide
5295        // enough to drop an interior column and still leave the reduction
5296        // non-trivial (p_time < p_time_full).
5297        let knots = array![
5298            left, left, left, left, q1, mid, q3, right, right, right, right
5299        ];
5300
5301        // Discover the full basis width by building with all columns retained.
5302        let full = build_survival_time_basis(
5303            &age_entry,
5304            &age_exit,
5305            SurvivalTimeBasisConfig::ISpline {
5306                degree: 2,
5307                knots: knots.clone(),
5308                keep_cols: Vec::new(),
5309                smooth_lambda: 1e-2,
5310            },
5311            None,
5312        )
5313        .expect("build full-width ispline time basis");
5314        let p_time_full = full
5315            .keep_cols
5316            .as_ref()
5317            .map(|k| k.len())
5318            .unwrap_or_else(|| full.x_exit_time.ncols());
5319        assert!(
5320            p_time_full >= 3,
5321            "test needs at least 3 shape-varying columns to drop an interior one; got {p_time_full}"
5322        );
5323
5324        // Retain everything except one interior column, forcing the
5325        // principal-submatrix-of-the-full-congruence path.
5326        let keep_cols: Vec<usize> = (0..p_time_full).filter(|&j| j != 1).collect();
5327
5328        let built = build_survival_time_basis(
5329            &age_entry,
5330            &age_exit,
5331            SurvivalTimeBasisConfig::ISpline {
5332                degree: 2,
5333                knots,
5334                keep_cols: keep_cols.clone(),
5335                smooth_lambda: 1e-2,
5336            },
5337            None,
5338        )
5339        .expect(
5340            "reduced ispline penalty must build (PSD contract must accept the \
5341             congruence-first / select-second ordering)",
5342        );
5343
5344        assert_eq!(
5345            built.penalties.len(),
5346            1,
5347            "the ispline time basis should carry exactly one curvature penalty"
5348        );
5349        let s = &built.penalties[0];
5350        assert_eq!(s.nrows(), keep_cols.len());
5351        assert_eq!(s.ncols(), keep_cols.len());
5352
5353        let (evals, _) = gam_linalg::faer_ndarray::FaerEigh::eigh(s, faer::Side::Lower)
5354            .expect("eigh of penalty");
5355        let evals_slice = evals.as_slice().expect("contiguous eigenvalues");
5356        let max_abs = evals_slice
5357            .iter()
5358            .copied()
5359            .fold(0.0_f64, |a, b| a.max(b.abs()))
5360            .max(1.0);
5361        let min_ev = evals_slice.iter().copied().fold(f64::INFINITY, f64::min);
5362        let tol = -100.0 * (s.nrows() as f64) * f64::EPSILON * max_abs;
5363        assert!(
5364            min_ev >= tol,
5365            "reduced I-spline time penalty must be PSD (gam#979): min eigenvalue \
5366             {min_ev:.3e} < tol {tol:.3e}, max|eig| {max_abs:.3e}"
5367        );
5368    }
5369
5370    #[test]
5371    fn marginal_slope_baseline_maps_gompertz_makeham_survival_to_probit_index() {
5372        let cfg = SurvivalBaselineConfig {
5373            target: SurvivalBaselineTarget::GompertzMakeham,
5374            scale: None,
5375            shape: Some(0.07),
5376            rate: Some(0.012),
5377            makeham: Some(0.003),
5378        };
5379        let age = 11.5;
5380        let (q, q_derivative) = evaluate_survival_marginal_slope_baseline(age, &cfg)
5381            .expect("evaluate marginal-slope gompertz-makeham baseline");
5382        let shape = cfg.shape.expect("shape");
5383        let rate = cfg.rate.expect("rate");
5384        let makeham = cfg.makeham.expect("makeham");
5385        let cumulative_hazard = makeham * age + (rate / shape) * ((shape * age).exp() - 1.0);
5386        let instant_hazard = makeham + rate * (shape * age).exp();
5387        let expected_survival = (-cumulative_hazard).exp();
5388        let actual_survival = normal_cdf(-q);
5389        assert!((actual_survival - expected_survival).abs() <= 1e-12);
5390
5391        let h = 1e-5;
5392        let q_plus = evaluate_survival_marginal_slope_baseline(age + h, &cfg)
5393            .expect("q plus")
5394            .0;
5395        let q_minus = evaluate_survival_marginal_slope_baseline(age - h, &cfg)
5396            .expect("q minus")
5397            .0;
5398        let fd = (q_plus - q_minus) / (2.0 * h);
5399        assert!((q_derivative - fd).abs() <= 1e-7);
5400        assert!(instant_hazard > 0.0);
5401    }
5402
5403    #[test]
5404    fn marginal_slope_baseline_is_evaluable_at_the_survival_curve_origin() {
5405        // Regression for #1024: the probit/marginal-slope baseline evaluator must
5406        // be defined at the survival-curve origin t = 0 (where S0(0) = 1, so the
5407        // probit index q(0) = -Phi^{-1}(1) = -inf and there is no finite offset),
5408        // exactly like its log-cumulative-hazard sibling `evaluate_survival_baseline`.
5409        // Before the fix the shared `age <= 0` hazard guard aborted, so a survival
5410        // prediction grid whose first node is the origin (the `Surv(time, event)`
5411        // right-censored shorthand) could not be evaluated for the location-scale /
5412        // marginal-slope likelihoods.
5413        let configs = [
5414            SurvivalBaselineConfig {
5415                target: SurvivalBaselineTarget::Linear,
5416                scale: None,
5417                shape: None,
5418                rate: None,
5419                makeham: None,
5420            },
5421            SurvivalBaselineConfig {
5422                target: SurvivalBaselineTarget::Weibull,
5423                scale: Some(2.5),
5424                shape: Some(1.3),
5425                rate: None,
5426                makeham: None,
5427            },
5428            SurvivalBaselineConfig {
5429                target: SurvivalBaselineTarget::Gompertz,
5430                scale: None,
5431                shape: Some(0.05),
5432                rate: Some(0.01),
5433                makeham: None,
5434            },
5435            SurvivalBaselineConfig {
5436                target: SurvivalBaselineTarget::GompertzMakeham,
5437                scale: None,
5438                shape: Some(0.07),
5439                rate: Some(0.012),
5440                makeham: Some(0.003),
5441            },
5442        ];
5443        for cfg in &configs {
5444            // The probit baseline returns a finite zero offset at the origin for
5445            // every target (the survival surface anchors S(0) = 1 directly).
5446            let (q0, q0_derivative) = evaluate_survival_marginal_slope_baseline(0.0, cfg)
5447                .expect("marginal-slope baseline must be evaluable at the origin");
5448            assert_eq!(q0, 0.0);
5449            assert_eq!(q0_derivative, 0.0);
5450
5451            // The log-cumulative-hazard sibling is likewise finite at the origin —
5452            // this parity is the whole point (the transformation likelihood already
5453            // worked because it rides this evaluator).
5454            let (eta0, eta0_derivative) =
5455                evaluate_survival_baseline(0.0, cfg).expect("log-cum-hazard baseline at origin");
5456            assert!(eta0_derivative.is_finite());
5457            assert!(eta0.is_finite() || eta0 == f64::NEG_INFINITY);
5458
5459            // The batched offset builder must not abort when a query exit age is the
5460            // origin (this is the exact call the location-scale predict path makes on
5461            // the default surface grid). Entry stays at the origin, exit spans 0 -> t.
5462            let age_entry = array![0.0, 0.0];
5463            let age_exit = array![0.0, 1.5];
5464            let (entry, exit, derivative) =
5465                build_survival_marginal_slope_baseline_offsets(&age_entry, &age_exit, cfg)
5466                    .expect("probit baseline offsets must build through the origin");
5467            assert!(entry.iter().all(|v| v.is_finite()));
5468            assert!(exit.iter().all(|v| v.is_finite()));
5469            assert!(derivative.iter().all(|v| v.is_finite()));
5470            // The origin exit column carries no probit offset.
5471            assert_eq!(exit[0], 0.0);
5472        }
5473    }
5474
5475    #[test]
5476    fn marginal_slope_baseline_offsets_use_true_gompertz_makeham_survival() {
5477        let cfg = SurvivalBaselineConfig {
5478            target: SurvivalBaselineTarget::GompertzMakeham,
5479            scale: None,
5480            shape: Some(0.03),
5481            rate: Some(0.01),
5482            makeham: Some(0.002),
5483        };
5484        let age_entry = array![2.0, 4.0];
5485        let age_exit = array![5.0, 9.0];
5486        let (entry, exit, derivative) =
5487            build_survival_marginal_slope_baseline_offsets(&age_entry, &age_exit, &cfg)
5488                .expect("marginal-slope baseline offsets");
5489        for i in 0..age_entry.len() {
5490            let entry_h = cfg.makeham.expect("makeham") * age_entry[i]
5491                + (cfg.rate.expect("rate") / cfg.shape.expect("shape"))
5492                    * ((cfg.shape.expect("shape") * age_entry[i]).exp() - 1.0);
5493            let exit_h = cfg.makeham.expect("makeham") * age_exit[i]
5494                + (cfg.rate.expect("rate") / cfg.shape.expect("shape"))
5495                    * ((cfg.shape.expect("shape") * age_exit[i]).exp() - 1.0);
5496            assert!((normal_cdf(-entry[i]) - (-entry_h).exp()).abs() <= 1e-12);
5497            assert!((normal_cdf(-exit[i]) - (-exit_h).exp()).abs() <= 1e-12);
5498            assert!(derivative[i].is_finite() && derivative[i] > 0.0);
5499        }
5500    }
5501
5502    fn fd_marginal_slope_baseline_offset(
5503        age: f64,
5504        cfg: &SurvivalBaselineConfig,
5505        steps: &[f64],
5506    ) -> Vec<(f64, f64)> {
5507        let theta = survival_baseline_theta_from_config(cfg)
5508            .expect("theta")
5509            .expect("non-linear baseline");
5510        assert_eq!(
5511            steps.len(),
5512            theta.len(),
5513            "fd_marginal_slope_baseline_offset: step vector length must match θ dimension"
5514        );
5515        (0..theta.len())
5516            .map(|k| {
5517                let h = steps[k];
5518                let mut theta_plus = theta.clone();
5519                theta_plus[k] += h;
5520                let mut theta_minus = theta.clone();
5521                theta_minus[k] -= h;
5522                let cfg_plus =
5523                    survival_baseline_config_from_theta(cfg.target, &theta_plus).expect("plus cfg");
5524                let cfg_minus = survival_baseline_config_from_theta(cfg.target, &theta_minus)
5525                    .expect("minus cfg");
5526                let (q_p, qt_p) =
5527                    evaluate_survival_marginal_slope_baseline(age, &cfg_plus).expect("q+");
5528                let (q_m, qt_m) =
5529                    evaluate_survival_marginal_slope_baseline(age, &cfg_minus).expect("q-");
5530                ((q_p - q_m) / (2.0 * h), (qt_p - qt_m) / (2.0 * h))
5531            })
5532            .collect()
5533    }
5534
5535    /// A frozen chart evaluated at `θ` records `θ` BITWISE (#2765).
5536    ///
5537    /// `SurvivalMarginalSlopeFamilyHyperState` stores the geometry's `theta` as
5538    /// the family's realized coordinates, and `validate_layout` compares them
5539    /// to the outer manifest with `to_bits()` equality — on purpose, so a
5540    /// workspace cannot reuse row geometry from a neighbouring outer probe.
5541    /// The chart used to close a `θ → cfg → θ` loop, which for a Weibull is
5542    /// `ln(exp(θ))` and is not the identity in `f64`. The exactness invariant
5543    /// then failed for a reason that has nothing to do with the geometry, the
5544    /// inner solve refused a point the outer optimizer was merely trying to
5545    /// evaluate, and the line search read that refusal as "no improvement" —
5546    /// 50 times, at every halving.
5547    ///
5548    /// The witnesses below are the coordinates the #2765 acceptance fixture
5549    /// actually refused at, with their measured round-trip error. The test
5550    /// asserts the round trip really is lossy at each (so it cannot quietly
5551    /// stop being a witness) and then that the chart is exact anyway.
5552    #[test]
5553    fn a_frozen_baseline_chart_records_the_theta_it_was_asked_for_2765() {
5554        let age_entry = array![0.0, 0.75, 2.0];
5555        let age_exit = array![1.5, 3.0, 5.5];
5556        let initial_config = SurvivalBaselineConfig {
5557            target: SurvivalBaselineTarget::Weibull,
5558            scale: Some(2.0),
5559            shape: Some(1.3),
5560            rate: None,
5561            makeham: None,
5562        };
5563        let baseline = build_survival_marginal_slope_baseline_geometry(
5564            &age_entry,
5565            &age_exit,
5566            &initial_config,
5567        )
5568        .expect("initial baseline geometry")
5569        .expect("Weibull is a nonlinear chart");
5570        let chart = SurvivalMarginalSlopeFrozenOffsetChart::new(
5571            &age_entry,
5572            &age_exit,
5573            &initial_config,
5574            &baseline.offset_entry,
5575            &baseline.offset_exit,
5576            &baseline.derivative_offset_exit,
5577        )
5578        .expect("freeze the Weibull chart");
5579
5580        // `θ₄ = ±1e-5` are the two coordinates the acceptance fixture's
5581        // certificate probe refused on BOTH sides; `1e-5` comes back 57_269
5582        // ulps away from itself through `ln(exp(·))`.
5583        for theta in [
5584            array![0.7574963781222602_f64, 1.0e-5],
5585            array![0.7574963781222602_f64, -1.0e-5],
5586            array![0.7574863781222603_f64, 0.0],
5587        ] {
5588            let lossy = theta
5589                .iter()
5590                .any(|value| value.exp().ln().to_bits() != value.to_bits());
5591            assert!(
5592                lossy,
5593                "this witness has stopped being one: every coordinate of {theta:?} now \
5594                 survives ln(exp(·)) bitwise, so it can no longer show the defect"
5595            );
5596            let realized = chart
5597                .evaluate(&theta)
5598                .expect("the chart evaluates inside its domain");
5599            for (axis, (want, got)) in theta.iter().zip(realized.theta.iter()).enumerate() {
5600                assert_eq!(
5601                    want.to_bits(),
5602                    got.to_bits(),
5603                    "chart axis {axis}: asked for {want:?}, recorded {got:?} — a chart must \
5604                     record the coordinate it was ASKED to realize, because the family's \
5605                     manifest check is bitwise"
5606                );
5607            }
5608        }
5609    }
5610
5611    #[test]
5612    fn marginal_slope_baseline_theta_partials_match_fd_for_gompertz_makeham() {
5613        let cfg = SurvivalBaselineConfig {
5614            target: SurvivalBaselineTarget::GompertzMakeham,
5615            scale: None,
5616            shape: Some(0.04),
5617            rate: Some(0.013),
5618            makeham: Some(0.002),
5619        };
5620        let age = 17.0;
5621        let analytic = marginal_slope_baseline_offset_theta_partials(age, &cfg)
5622            .expect("partials")
5623            .expect("nonlinear");
5624        let fd = fd_marginal_slope_baseline_offset(age, &cfg, &[1e-5, 1e-5, 1e-5]);
5625        assert_eq!(analytic.len(), fd.len());
5626        for (k, ((aq, aqt), (fq, fqt))) in analytic.iter().zip(fd.iter()).enumerate() {
5627            assert_close(*aq, *fq, 1e-6, &format!("gm-probit q theta[{k}]"));
5628            assert_close(*aqt, *fqt, 1e-6, &format!("gm-probit q' theta[{k}]"));
5629        }
5630    }
5631
5632    #[test]
5633    fn marginal_slope_baseline_theta_partials_match_fd_near_zero_gompertz_shape() {
5634        let cfg = SurvivalBaselineConfig {
5635            target: SurvivalBaselineTarget::GompertzMakeham,
5636            scale: None,
5637            shape: Some(1e-14),
5638            rate: Some(0.013),
5639            makeham: Some(0.002),
5640        };
5641        let age = 17.0;
5642        let analytic = marginal_slope_baseline_offset_theta_partials(age, &cfg)
5643            .expect("partials")
5644            .expect("nonlinear");
5645        let fd = fd_marginal_slope_baseline_offset(age, &cfg, &[1e-5, 1e-11, 1e-5]);
5646        assert_eq!(analytic.len(), fd.len());
5647        for (k, ((aq, aqt), (fq, fqt))) in analytic.iter().zip(fd.iter()).enumerate() {
5648            assert_close(*aq, *fq, 1e-5, &format!("near-zero gm-probit q theta[{k}]"));
5649            assert_close(
5650                *aqt,
5651                *fqt,
5652                1e-5,
5653                &format!("near-zero gm-probit q' theta[{k}]"),
5654            );
5655        }
5656    }
5657
5658    fn shifted_quadratic_offset_residuals(
5659        age_entry: ndarray::ArrayView1<'_, f64>,
5660        age_exit: ndarray::ArrayView1<'_, f64>,
5661        base_cfg: &SurvivalBaselineConfig,
5662        candidate_cfg: &SurvivalBaselineConfig,
5663        base: &OffsetChannelResiduals,
5664        curvatures: &OffsetChannelCurvatures,
5665    ) -> OffsetChannelResiduals {
5666        let n = age_exit.len();
5667        let mut entry = base.entry.clone();
5668        let mut exit = base.exit.clone();
5669        let mut derivative = base.derivative.clone();
5670        for row in 0..n {
5671            let (_, base_exit, base_deriv) =
5672                baseline_marginal_slope_channels(age_exit[row], base_cfg);
5673            let (_, cand_exit, cand_deriv) =
5674                baseline_marginal_slope_channels(age_exit[row], candidate_cfg);
5675            let base_entry = if base.entry[row] == 0.0 {
5676                0.0
5677            } else {
5678                baseline_marginal_slope_channels(age_entry[row], base_cfg).1
5679            };
5680            let cand_entry = if base.entry[row] == 0.0 {
5681                0.0
5682            } else {
5683                baseline_marginal_slope_channels(age_entry[row], candidate_cfg).1
5684            };
5685            let delta = [
5686                cand_entry - base_entry,
5687                cand_exit - base_exit,
5688                cand_deriv - base_deriv,
5689            ];
5690            let mut shift = [0.0; 3];
5691            for i in 0..3 {
5692                for j in 0..3 {
5693                    shift[i] += curvatures.rows[row][i][j] * delta[j];
5694                }
5695            }
5696            if base.entry[row] != 0.0 {
5697                entry[row] += shift[0];
5698            }
5699            exit[row] += shift[1];
5700            derivative[row] += shift[2];
5701        }
5702        OffsetChannelResiduals {
5703            entry,
5704            exit,
5705            derivative,
5706            right: base.right.clone(),
5707        }
5708    }
5709
5710    fn baseline_marginal_slope_channels(age: f64, cfg: &SurvivalBaselineConfig) -> (f64, f64, f64) {
5711        let (q, q_t) = evaluate_survival_marginal_slope_baseline(age, cfg).expect("baseline");
5712        (q, q, q_t)
5713    }
5714
5715    #[test]
5716    fn marginal_slope_baseline_chain_rule_hessian_matches_fd_gradient() {
5717        let cfg = SurvivalBaselineConfig {
5718            target: SurvivalBaselineTarget::GompertzMakeham,
5719            scale: None,
5720            shape: Some(0.025),
5721            rate: Some(0.012),
5722            makeham: Some(0.003),
5723        };
5724        let theta = survival_baseline_theta_from_config(&cfg)
5725            .expect("theta")
5726            .expect("nonlinear");
5727        let age_entry = array![2.5, 0.0, 5.0];
5728        let age_exit = array![7.5, 11.0, 15.0];
5729        let base_residuals = OffsetChannelResiduals {
5730            entry: array![0.2, 0.0, -0.1],
5731            exit: array![0.6, -0.3, 0.4],
5732            derivative: array![-0.5, 0.25, 0.15],
5733            right: Array1::<f64>::zeros(3),
5734        };
5735        let curvatures = OffsetChannelCurvatures {
5736            rows: vec![
5737                [[1.4, 0.2, -0.1], [0.2, 1.1, 0.05], [-0.1, 0.05, 0.7]],
5738                [[0.9, -0.15, 0.0], [-0.15, 1.3, 0.12], [0.0, 0.12, 0.8]],
5739                [[1.2, 0.05, 0.09], [0.05, 0.95, -0.04], [0.09, -0.04, 0.6]],
5740            ],
5741        };
5742        let analytic = marginal_slope_baseline_chain_rule_hessian(
5743            age_entry.view(),
5744            age_exit.view(),
5745            &cfg,
5746            &base_residuals,
5747            &curvatures,
5748        )
5749        .expect("hessian")
5750        .expect("nonlinear");
5751
5752        let gradient_at = |theta_candidate: &Array1<f64>| -> Array1<f64> {
5753            let candidate = survival_baseline_config_from_theta(cfg.target, theta_candidate)
5754                .expect("candidate cfg");
5755            let residuals = shifted_quadratic_offset_residuals(
5756                age_entry.view(),
5757                age_exit.view(),
5758                &cfg,
5759                &candidate,
5760                &base_residuals,
5761                &curvatures,
5762            );
5763            marginal_slope_baseline_chain_rule_gradient(
5764                age_entry.view(),
5765                age_exit.view(),
5766                &candidate,
5767                &residuals,
5768            )
5769            .expect("gradient")
5770            .expect("nonlinear")
5771        };
5772
5773        for j in 0..theta.len() {
5774            let step = if j == 1 { 2e-5 } else { 1e-5 };
5775            let mut plus = theta.clone();
5776            plus[j] += step;
5777            let mut minus = theta.clone();
5778            minus[j] -= step;
5779            let fd_col = (&gradient_at(&plus) - &gradient_at(&minus)) / (2.0 * step);
5780            for i in 0..theta.len() {
5781                assert_close(
5782                    analytic[[i, j]],
5783                    fd_col[i],
5784                    2e-5,
5785                    &format!("baseline Hessian ({i},{j})"),
5786                );
5787            }
5788        }
5789    }
5790
5791    #[test]
5792    fn marginal_slope_baseline_chain_rule_gradient_contracts_probit_partials() {
5793        let cfg = SurvivalBaselineConfig {
5794            target: SurvivalBaselineTarget::GompertzMakeham,
5795            scale: None,
5796            shape: Some(0.03),
5797            rate: Some(0.01),
5798            makeham: Some(0.002),
5799        };
5800        let age_entry = array![3.0, 6.0];
5801        let age_exit = array![8.0, 12.0];
5802        let residuals = OffsetChannelResiduals {
5803            exit: array![0.7, -0.2],
5804            entry: array![0.1, 0.4],
5805            derivative: array![1.3, -0.6],
5806            right: Array1::<f64>::zeros(2),
5807        };
5808        let grad = marginal_slope_baseline_chain_rule_gradient(
5809            age_entry.view(),
5810            age_exit.view(),
5811            &cfg,
5812            &residuals,
5813        )
5814        .expect("gradient")
5815        .expect("nonlinear");
5816
5817        let mut expected = Array1::<f64>::zeros(3);
5818        for i in 0..age_exit.len() {
5819            let exit_partials = marginal_slope_baseline_offset_theta_partials(age_exit[i], &cfg)
5820                .expect("exit partials")
5821                .expect("nonlinear");
5822            let entry_partials = marginal_slope_baseline_offset_theta_partials(age_entry[i], &cfg)
5823                .expect("entry partials")
5824                .expect("nonlinear");
5825            for k in 0..3 {
5826                expected[k] += residuals.exit[i] * exit_partials[k].0
5827                    + residuals.derivative[i] * exit_partials[k].1
5828                    + residuals.entry[i] * entry_partials[k].0;
5829            }
5830        }
5831        for k in 0..3 {
5832            assert_close(
5833                grad[k],
5834                expected[k],
5835                1e-12,
5836                &format!("gm-probit chain gradient theta[{k}]"),
5837            );
5838        }
5839    }
5840
5841    /// Parity guard for the shared `baseline_chain_rule_gradient_with_partials`
5842    /// engine (issue #429): both public gradient functions delegate to it with a
5843    /// different partials provider. This test reimplements the pre-unification
5844    /// inline contraction (the serial reference) and asserts bit-for-bit equality
5845    /// against the unified engine's output for BOTH providers on the same data —
5846    /// the RP-eta provider (`baseline_offset_theta_partials`) and the probit-q
5847    /// provider (`marginal_slope_baseline_offset_theta_partials`). Any drift in
5848    /// the extracted contraction (length checks, theta-dim probe, exit/derivative
5849    /// combination, or entry gating) breaks this with an exact (0.0) tolerance.
5850    #[test]
5851    fn baseline_chain_rule_gradient_engine_matches_inline_reference() {
5852        let cfg = SurvivalBaselineConfig {
5853            target: SurvivalBaselineTarget::GompertzMakeham,
5854            scale: None,
5855            shape: Some(0.028),
5856            rate: Some(0.011),
5857            makeham: Some(0.0025),
5858        };
5859        // Mixed entry interval: row 1 is origin-entry (age_entry==0, r_entry==0)
5860        // to exercise the entry-gating branch in the shared engine.
5861        let age_entry = array![3.0, 0.0, 5.5];
5862        let age_exit = array![8.0, 12.0, 16.0];
5863        let residuals = OffsetChannelResiduals {
5864            exit: array![0.7, -0.2, 0.45],
5865            entry: array![0.1, 0.0, -0.3],
5866            derivative: array![1.3, -0.6, 0.2],
5867            right: Array1::<f64>::zeros(3),
5868        };
5869
5870        // Serial reference contraction matching the original inline body. Mirrors
5871        // the engine's exit+derivative/entry split and origin-entry gating.
5872        let reference_gradient = |partials: &dyn Fn(
5873            f64,
5874            &SurvivalBaselineConfig,
5875        )
5876            -> Result<Option<Vec<(f64, f64)>>, String>|
5877         -> Array1<f64> {
5878            let theta_dim = partials(age_exit[0], &cfg)
5879                .expect("probe partials")
5880                .expect("nonlinear")
5881                .len();
5882            let mut acc = Array1::<f64>::zeros(theta_dim);
5883            for i in 0..age_exit.len() {
5884                let p_exit = partials(age_exit[i], &cfg)
5885                    .expect("exit partials")
5886                    .expect("nonlinear");
5887                let r_x = residuals.exit[i];
5888                let r_d = residuals.derivative[i];
5889                for k in 0..theta_dim {
5890                    acc[k] += r_x * p_exit[k].0 + r_d * p_exit[k].1;
5891                }
5892                let r_e = residuals.entry[i];
5893                if r_e != 0.0 {
5894                    let p_entry = partials(age_entry[i], &cfg)
5895                        .expect("entry partials")
5896                        .expect("nonlinear");
5897                    for k in 0..theta_dim {
5898                        acc[k] += r_e * p_entry[k].0;
5899                    }
5900                }
5901            }
5902            acc
5903        };
5904
5905        // RP-eta provider parity.
5906        let rp_engine = baseline_chain_rule_gradient(
5907            age_entry.view(),
5908            age_exit.view(),
5909            age_exit.view(),
5910            &cfg,
5911            &residuals,
5912        )
5913        .expect("rp gradient")
5914        .expect("rp nonlinear");
5915        let rp_reference = reference_gradient(&baseline_offset_theta_partials);
5916        assert_eq!(rp_engine.len(), rp_reference.len());
5917        for k in 0..rp_engine.len() {
5918            assert_close(
5919                rp_engine[k],
5920                rp_reference[k],
5921                0.0,
5922                &format!("rp engine vs inline reference theta[{k}]"),
5923            );
5924        }
5925
5926        // Probit-q provider parity.
5927        let probit_engine = marginal_slope_baseline_chain_rule_gradient(
5928            age_entry.view(),
5929            age_exit.view(),
5930            &cfg,
5931            &residuals,
5932        )
5933        .expect("probit gradient")
5934        .expect("probit nonlinear");
5935        let probit_reference = reference_gradient(&marginal_slope_baseline_offset_theta_partials);
5936        assert_eq!(probit_engine.len(), probit_reference.len());
5937        for k in 0..probit_engine.len() {
5938            assert_close(
5939                probit_engine[k],
5940                probit_reference[k],
5941                0.0,
5942                &format!("probit engine vs inline reference theta[{k}]"),
5943            );
5944        }
5945    }
5946
5947    /// Finite-difference verification of the analytic θ-gradient used by the
5948    /// survival location-scale workflow path.
5949    ///
5950    /// At a converged β, the envelope theorem reduces the profile-NLL gradient
5951    /// w.r.t. the baseline-config θ to a per-row residual contraction against
5952    /// the per-row offset-channel partials ∂o/∂θ:
5953    ///
5954    ///   d(NLL)/dθ_k = Σ_i [ r_X[i]·∂η_exit/∂θ_k + r_E[i]·∂η_entry/∂θ_k
5955    ///                       + r_D[i]·∂o_D_exit/∂θ_k ]
5956    ///
5957    /// (`baseline_chain_rule_gradient`). Because β is fixed, an explicit loss
5958    /// `L(θ) = Σ_i [ r_X[i]·η(t_exit_i; θ) + r_E[i]·η(t_entry_i; θ)
5959    ///              + r_D[i]·o_D(t_exit_i; θ) ]`
5960    /// has gradient identically equal to the chain-rule output. Comparing the
5961    /// analytic gradient to a central-difference of L over `evaluate_survival_baseline`
5962    /// therefore exercises every piece of the chain rule (incl. the Gompertz
5963    /// rate / shape / Makeham partials at both entry and exit ages) without
5964    /// needing the full location-scale fit pipeline inside this unit-test
5965    /// module. If the chain rule disagrees with FD here, the workflow's
5966    /// gradient is wrong by exactly the same amount.
5967    #[test]
5968    fn gompertz_makeham_baseline_chain_rule_gradient_matches_finite_difference() {
5969        let cfg = SurvivalBaselineConfig {
5970            target: SurvivalBaselineTarget::GompertzMakeham,
5971            scale: None,
5972            shape: Some(0.05),
5973            rate: Some(0.012),
5974            makeham: Some(0.003),
5975        };
5976        // n = 8 small synthetic dataset spanning a realistic age range.
5977        let age_entry = array![5.0, 8.0, 12.0, 0.5, 20.0, 30.0, 45.0, 60.0];
5978        let age_exit = array![10.0, 15.0, 25.0, 4.0, 35.0, 50.0, 65.0, 80.0];
5979        // Synthetic per-row NLL residuals on the three offset channels. Mix of
5980        // signs / magnitudes / one zero-entry row (origin entry → r_E=0).
5981        let residuals = OffsetChannelResiduals {
5982            exit: array![0.42, -0.18, 0.73, -0.91, 0.05, -0.27, 0.61, -0.34],
5983            entry: array![-0.12, 0.31, -0.44, 0.0, 0.16, -0.22, 0.07, -0.51],
5984            derivative: array![1.04, -0.65, 0.18, -1.21, 0.42, -0.13, 0.88, -0.27],
5985            right: Array1::<f64>::zeros(8),
5986        };
5987
5988        let analytic = baseline_chain_rule_gradient(
5989            age_entry.view(),
5990            age_exit.view(),
5991            age_exit.view(),
5992            &cfg,
5993            &residuals,
5994        )
5995        .expect("analytic gradient ok")
5996        .expect("GM baseline has a θ-gradient");
5997        assert_eq!(analytic.len(), 3, "GM θ has 3 components");
5998
5999        // Evaluate the offset-projected loss at a perturbed θ. Mirrors the
6000        // chain rule's algebra: the entry channel is only added for rows whose
6001        // r_E is nonzero (matching baseline_chain_rule_gradient's gating that
6002        // avoids calling evaluate_survival_baseline at age 0 for origin-entry
6003        // rows).
6004        let loss_at_cfg = |cfg_eval: &SurvivalBaselineConfig| -> f64 {
6005            let mut acc = 0.0;
6006            for i in 0..age_exit.len() {
6007                let (eta_exit_i, od_exit_i) =
6008                    evaluate_survival_baseline(age_exit[i], cfg_eval).expect("eval exit");
6009                acc += residuals.exit[i] * eta_exit_i + residuals.derivative[i] * od_exit_i;
6010                if residuals.entry[i] != 0.0 {
6011                    let (eta_entry_i, _) =
6012                        evaluate_survival_baseline(age_entry[i], cfg_eval).expect("eval entry");
6013                    acc += residuals.entry[i] * eta_entry_i;
6014                }
6015            }
6016            acc
6017        };
6018
6019        let theta0 = survival_baseline_theta_from_config(&cfg)
6020            .expect("theta seed")
6021            .expect("GM has θ");
6022        // Spec requested δ = 1e-4 per axis. Use central differences over θ.
6023        let delta = 1e-4;
6024        let mut fd = Array1::<f64>::zeros(analytic.len());
6025        for k in 0..analytic.len() {
6026            let mut theta_plus = theta0.clone();
6027            theta_plus[k] += delta;
6028            let mut theta_minus = theta0.clone();
6029            theta_minus[k] -= delta;
6030            let cfg_plus =
6031                survival_baseline_config_from_theta(cfg.target, &theta_plus).expect("cfg(θ+δ)");
6032            let cfg_minus =
6033                survival_baseline_config_from_theta(cfg.target, &theta_minus).expect("cfg(θ-δ)");
6034            let lp = loss_at_cfg(&cfg_plus);
6035            let lm = loss_at_cfg(&cfg_minus);
6036            fd[k] = (lp - lm) / (2.0 * delta);
6037        }
6038
6039        let analytic_norm = analytic.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
6040        let max_err = analytic
6041            .iter()
6042            .zip(fd.iter())
6043            .map(|(a, b)| (a - b).abs())
6044            .fold(0.0_f64, f64::max);
6045        let rel = max_err / (analytic_norm + 1e-12);
6046        // Print so the deliverable can quote the exact max-error number.
6047        eprintln!(
6048            "gompertz_makeham_baseline_chain_rule_gradient_matches_finite_difference: \
6049             analytic={analytic:?} fd={fd:?} max_err={max_err:.3e} \
6050             analytic_inf_norm={analytic_norm:.3e} rel={rel:.3e}"
6051        );
6052        assert!(
6053            rel < 1e-2,
6054            "analytic θ-gradient disagrees with central FD beyond 1%: \
6055             analytic={analytic:?}, fd={fd:?}, max_err={max_err:.3e}, \
6056             rel={rel:.3e} (analytic_inf_norm={analytic_norm:.3e})"
6057        );
6058    }
6059
6060    /// Weibull (dim=2) companion to
6061    /// `gompertz_makeham_baseline_chain_rule_gradient_matches_finite_difference`.
6062    ///
6063    /// This is the FD gate for the analytic outer θ-gradient that the
6064    /// transformation/Weibull survival baseline optimizers now feed to BFGS
6065    /// (`optimize_survival_baseline_config_with_gradient_only`). At a *fixed* β
6066    /// the profile-NLL surface is
6067    /// `L(θ) = Σ_i [ r_X[i]·η(t_exit_i;θ) + r_E[i]·η(t_entry_i;θ)
6068    ///              + r_D[i]·o_D(t_exit_i;θ) ]`,
6069    /// whose exact gradient is `baseline_chain_rule_gradient`. Comparing it to a
6070    /// central difference of `L` over `evaluate_survival_baseline` exercises the
6071    /// Weibull scale/shape partials at both entry and exit ages. If this
6072    /// disagrees with FD, the workflow's outer gradient is wrong by the same
6073    /// amount.
6074    #[test]
6075    fn weibull_baseline_chain_rule_gradient_matches_finite_difference() {
6076        let cfg = SurvivalBaselineConfig {
6077            target: SurvivalBaselineTarget::Weibull,
6078            scale: Some(11.0),
6079            shape: Some(1.4),
6080            rate: None,
6081            makeham: None,
6082        };
6083        let age_entry = array![5.0, 8.0, 12.0, 0.5, 20.0, 30.0, 45.0, 60.0];
6084        let age_exit = array![10.0, 15.0, 25.0, 4.0, 35.0, 50.0, 65.0, 80.0];
6085        let residuals = OffsetChannelResiduals {
6086            exit: array![0.42, -0.18, 0.73, -0.91, 0.05, -0.27, 0.61, -0.34],
6087            entry: array![-0.12, 0.31, -0.44, 0.0, 0.16, -0.22, 0.07, -0.51],
6088            derivative: array![1.04, -0.65, 0.18, -1.21, 0.42, -0.13, 0.88, -0.27],
6089            right: Array1::<f64>::zeros(8),
6090        };
6091
6092        let analytic = baseline_chain_rule_gradient(
6093            age_entry.view(),
6094            age_exit.view(),
6095            age_exit.view(),
6096            &cfg,
6097            &residuals,
6098        )
6099        .expect("analytic gradient ok")
6100        .expect("Weibull baseline has a θ-gradient");
6101        assert_eq!(analytic.len(), 2, "Weibull θ has 2 components");
6102
6103        let loss_at_cfg = |cfg_eval: &SurvivalBaselineConfig| -> f64 {
6104            let mut acc = 0.0;
6105            for i in 0..age_exit.len() {
6106                let (eta_exit_i, od_exit_i) =
6107                    evaluate_survival_baseline(age_exit[i], cfg_eval).expect("eval exit");
6108                acc += residuals.exit[i] * eta_exit_i + residuals.derivative[i] * od_exit_i;
6109                if residuals.entry[i] != 0.0 {
6110                    let (eta_entry_i, _) =
6111                        evaluate_survival_baseline(age_entry[i], cfg_eval).expect("eval entry");
6112                    acc += residuals.entry[i] * eta_entry_i;
6113                }
6114            }
6115            acc
6116        };
6117
6118        let theta0 = survival_baseline_theta_from_config(&cfg)
6119            .expect("theta seed")
6120            .expect("Weibull has θ");
6121        let delta = 1e-4;
6122        let mut fd = Array1::<f64>::zeros(analytic.len());
6123        for k in 0..analytic.len() {
6124            let mut theta_plus = theta0.clone();
6125            theta_plus[k] += delta;
6126            let mut theta_minus = theta0.clone();
6127            theta_minus[k] -= delta;
6128            let cfg_plus =
6129                survival_baseline_config_from_theta(cfg.target, &theta_plus).expect("cfg(θ+δ)");
6130            let cfg_minus =
6131                survival_baseline_config_from_theta(cfg.target, &theta_minus).expect("cfg(θ-δ)");
6132            let lp = loss_at_cfg(&cfg_plus);
6133            let lm = loss_at_cfg(&cfg_minus);
6134            fd[k] = (lp - lm) / (2.0 * delta);
6135        }
6136
6137        let analytic_norm = analytic.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
6138        let max_err = analytic
6139            .iter()
6140            .zip(fd.iter())
6141            .map(|(a, b)| (a - b).abs())
6142            .fold(0.0_f64, f64::max);
6143        let rel = max_err / (analytic_norm + 1e-12);
6144        eprintln!(
6145            "weibull_baseline_chain_rule_gradient_matches_finite_difference: \
6146             analytic={analytic:?} fd={fd:?} max_err={max_err:.3e} \
6147             analytic_inf_norm={analytic_norm:.3e} rel={rel:.3e}"
6148        );
6149        assert!(
6150            rel < 1e-2,
6151            "analytic θ-gradient disagrees with central FD beyond 1%: \
6152             analytic={analytic:?}, fd={fd:?}, max_err={max_err:.3e}, \
6153             rel={rel:.3e} (analytic_inf_norm={analytic_norm:.3e})"
6154        );
6155    }
6156
6157    // ─── baseline_offset_theta_partials — analytic vs central-difference ─
6158
6159    /// Central-difference of (eta, o_D) at fixed age wrt each θ component in
6160    /// the theta layout defined by `survival_baseline_theta_from_config`.
6161    ///
6162    /// `steps` is per-θ-component: the caller picks the step size appropriate
6163    /// for each channel. Gompertz / Gompertz–Makeham need a tiny step on the
6164    /// shape channel near the Taylor pivot |shape| < 1e-10 (so θ±h stays on
6165    /// the same branch), but a normal-scale step on log_rate / log_makeham;
6166    /// using the tiny shape-step on every channel corrupts the log_rate
6167    /// channel with `eps/(2h)` cancellation noise and has nothing to do with
6168    /// correctness of the analytic derivative.
6169    fn fd_baseline_offset(
6170        age: f64,
6171        cfg: &SurvivalBaselineConfig,
6172        steps: &[f64],
6173    ) -> Vec<(f64, f64)> {
6174        let theta = survival_baseline_theta_from_config(cfg)
6175            .expect("theta")
6176            .expect("non-linear baseline");
6177        assert_eq!(
6178            steps.len(),
6179            theta.len(),
6180            "fd_baseline_offset: step vector length must match θ dimension"
6181        );
6182        (0..theta.len())
6183            .map(|k| {
6184                let h = steps[k];
6185                let mut theta_plus = theta.clone();
6186                theta_plus[k] += h;
6187                let mut theta_minus = theta.clone();
6188                theta_minus[k] -= h;
6189                let cfg_plus =
6190                    survival_baseline_config_from_theta(cfg.target, &theta_plus).expect("plus cfg");
6191                let cfg_minus = survival_baseline_config_from_theta(cfg.target, &theta_minus)
6192                    .expect("minus cfg");
6193                let (eta_p, od_p) = evaluate_survival_baseline(age, &cfg_plus).expect("eta+");
6194                let (eta_m, od_m) = evaluate_survival_baseline(age, &cfg_minus).expect("eta-");
6195                ((eta_p - eta_m) / (2.0 * h), (od_p - od_m) / (2.0 * h))
6196            })
6197            .collect()
6198    }
6199
6200    fn assert_close(actual: f64, expected: f64, tol: f64, what: &str) {
6201        // `<=` so that bit-equal values satisfy tol = 0. With `<`, |a−e| < 0
6202        // is unsatisfiable and a zero-tolerance "must match exactly" call
6203        // would reject identical numbers.
6204        let ok = if expected.abs() < 1.0 {
6205            (actual - expected).abs() <= tol
6206        } else {
6207            (actual - expected).abs() <= tol * expected.abs().max(1.0)
6208        };
6209        assert!(
6210            ok,
6211            "{what}: analytic={actual:.6e} fd={expected:.6e} (tol={tol:.1e})"
6212        );
6213    }
6214
6215    #[test]
6216    fn gompertz_offset_partials_match_central_diff() {
6217        // Several (rate, shape, age) combinations spanning the small-shape
6218        // Taylor branch (|shape| < 1e-10) and the normal branch
6219        // (shape >> 1e-10), plus sign-reversed shape.
6220        let cases = [
6221            (0.5_f64, 0.01_f64, 30.0_f64),
6222            (0.2, 0.05, 60.0),
6223            (1.0, 0.001, 10.0),
6224            (0.4, 5e-11, 25.0),
6225            (0.4, -5e-11, 25.0),
6226            (0.3, -0.02, 40.0),
6227            (0.8, 0.2, 5.0),
6228        ];
6229        for &(rate, shape, age) in &cases {
6230            let cfg = SurvivalBaselineConfig {
6231                target: SurvivalBaselineTarget::Gompertz,
6232                scale: None,
6233                shape: Some(shape),
6234                rate: Some(rate),
6235                makeham: None,
6236            };
6237            let analytic = baseline_offset_theta_partials(age, &cfg)
6238                .expect("ok")
6239                .expect("non-linear");
6240            // Keep the FD probe inside the Taylor branch for tiny |shape| so
6241            // the numeric derivative matches the same small-shape map as the
6242            // analytic helper. log_rate always uses the normal step — rate
6243            // is a moderate-scale parameter and a 1e-11 step would swamp the
6244            // FD with cancellation noise.
6245            let h_shape = if shape.abs() < 1e-9 { 1e-11 } else { 1e-5 };
6246            let fd = fd_baseline_offset(age, &cfg, &[1e-5, h_shape]);
6247            assert_eq!(analytic.len(), 2);
6248            // Gompertz θ=(log_rate, shape). Rate channel: ∂eta/∂log_rate=1, ∂o_D/∂log_rate=0.
6249            assert_close(
6250                analytic[0].0,
6251                fd[0].0,
6252                1e-7,
6253                &format!("gompertz ∂eta/∂log_rate (rate={rate}, shape={shape}, age={age})"),
6254            );
6255            assert_close(
6256                analytic[0].1,
6257                fd[0].1,
6258                1e-7,
6259                &format!("gompertz ∂o_D/∂log_rate (rate={rate}, shape={shape}, age={age})"),
6260            );
6261            // shape channel — larger tol because finite-differencing near
6262            // shape=0 amplifies rounding; 1e-5 is fine.
6263            assert_close(
6264                analytic[1].0,
6265                fd[1].0,
6266                1e-5,
6267                &format!("gompertz ∂eta/∂shape (rate={rate}, shape={shape}, age={age})"),
6268            );
6269            assert_close(
6270                analytic[1].1,
6271                fd[1].1,
6272                1e-5,
6273                &format!("gompertz ∂o_D/∂shape (rate={rate}, shape={shape}, age={age})"),
6274            );
6275        }
6276    }
6277
6278    #[test]
6279    fn gompertz_offset_partials_log_rate_channel_is_trivial() {
6280        // Pure Gompertz: rate cancels in o_D, so ∂o_D/∂log_rate must be
6281        // exactly 0 and ∂eta/∂log_rate must be exactly 1. Verify the
6282        // analytic implementation returns the exact values, not FD-close.
6283        let cfg = SurvivalBaselineConfig {
6284            target: SurvivalBaselineTarget::Gompertz,
6285            scale: None,
6286            shape: Some(0.05),
6287            rate: Some(0.3),
6288            makeham: None,
6289        };
6290        let partials = baseline_offset_theta_partials(42.0, &cfg)
6291            .expect("ok")
6292            .expect("non-linear");
6293        assert_eq!(partials[0].0, 1.0);
6294        assert_eq!(partials[0].1, 0.0);
6295    }
6296
6297    #[test]
6298    fn gompertz_offset_partials_small_shape_taylor_agrees_with_direct_branch() {
6299        // Both branches of gompertz_shape_derivatives should agree to high
6300        // precision at shape = 1e-10 + epsilon on the direct side vs
6301        // shape = 1e-10 - epsilon on the Taylor side. Here we spot-check
6302        // the continuity at the branch cutoff: shape slightly above and
6303        // slightly below 1e-10 must give values within O(shape²·t²)
6304        // (the Taylor truncation error).
6305        let age = 25.0;
6306        let rate = 0.4;
6307        let cfg_taylor = SurvivalBaselineConfig {
6308            target: SurvivalBaselineTarget::Gompertz,
6309            scale: None,
6310            shape: Some(0.5e-10),
6311            rate: Some(rate),
6312            makeham: None,
6313        };
6314        let cfg_direct = SurvivalBaselineConfig {
6315            target: SurvivalBaselineTarget::Gompertz,
6316            scale: None,
6317            shape: Some(2.0e-10),
6318            rate: Some(rate),
6319            makeham: None,
6320        };
6321        let p_t = baseline_offset_theta_partials(age, &cfg_taylor)
6322            .expect("ok")
6323            .expect("nl");
6324        let p_d = baseline_offset_theta_partials(age, &cfg_direct)
6325            .expect("ok")
6326            .expect("nl");
6327        // ∂eta/∂shape at shape≈0 should be t/2 = 12.5 on both sides.
6328        assert_close(p_t[1].0, 12.5, 1e-8, "taylor ∂eta/∂shape near 0");
6329        assert_close(p_d[1].0, 12.5, 1e-8, "direct ∂eta/∂shape near 0");
6330        // ∂o_D/∂shape at shape≈0 should be 1/2.
6331        assert_close(p_t[1].1, 0.5, 1e-8, "taylor ∂o_D/∂shape near 0");
6332        assert_close(p_d[1].1, 0.5, 1e-8, "direct ∂o_D/∂shape near 0");
6333    }
6334
6335    // ----------------------------------------------------------------------
6336    // Gompertz hazard-channel shape derivatives: FD oracle + Taylor-branch
6337    // continuity. These feed `survival_hazard_theta_partials` /
6338    // `survival_hazard_theta_first_second` (the marginal-slope probit
6339    // baseline). Before this test, the only coverage of
6340    // `gompertz_cumulative_shape_{,second_}derivative` was the indirect
6341    // marginal-slope Hessian FD at shape=0.025, which never touches the
6342    // small-shape (`|shape| < 1e-10`) Taylor branch nor directly FD-checks
6343    // these analytic shape derivatives.
6344    // ----------------------------------------------------------------------
6345
6346    #[test]
6347    fn gompertz_hazard_shape_derivatives_match_central_diff() {
6348        // shape stays well above the 1e-10 Taylor cutoff so the exact
6349        // closed-form branch is exercised and the expm1/exp arithmetic is
6350        // numerically clean. FD on the analytic value/first-derivative
6351        // confirms the first and second shape derivatives.
6352        let cases = [
6353            (10.0_f64, 0.012_f64, 0.05_f64),
6354            (2.5, 0.5, 0.2),
6355            (15.0, 0.003, 0.01),
6356            (40.0, 0.3, 0.001),
6357        ];
6358        let h = 1e-6;
6359        for &(age, rate, shape) in &cases {
6360            // First shape derivative of (H_G, h_G) vs central diff of value.
6361            let (d_cum, d_inst) = gompertz_cumulative_shape_derivative(age, rate, shape);
6362            let (cum_p, inst_p) = gompertz_hazard_components(age, rate, shape + h);
6363            let (cum_m, inst_m) = gompertz_hazard_components(age, rate, shape - h);
6364            assert_close(
6365                d_cum,
6366                (cum_p - cum_m) / (2.0 * h),
6367                1e-6,
6368                &format!("∂H_G/∂shape (age={age}, rate={rate}, shape={shape})"),
6369            );
6370            assert_close(
6371                d_inst,
6372                (inst_p - inst_m) / (2.0 * h),
6373                1e-6,
6374                &format!("∂h_G/∂shape (age={age}, rate={rate}, shape={shape})"),
6375            );
6376
6377            // Second shape derivative vs central diff of the first derivative.
6378            let (d2_cum, d2_inst) = gompertz_cumulative_shape_second_derivative(age, rate, shape);
6379            let (dcum_p, dinst_p) = gompertz_cumulative_shape_derivative(age, rate, shape + h);
6380            let (dcum_m, dinst_m) = gompertz_cumulative_shape_derivative(age, rate, shape - h);
6381            assert_close(
6382                d2_cum,
6383                (dcum_p - dcum_m) / (2.0 * h),
6384                1e-5,
6385                &format!("∂²H_G/∂shape² (age={age}, rate={rate}, shape={shape})"),
6386            );
6387            assert_close(
6388                d2_inst,
6389                (dinst_p - dinst_m) / (2.0 * h),
6390                1e-5,
6391                &format!("∂²h_G/∂shape² (age={age}, rate={rate}, shape={shape})"),
6392            );
6393        }
6394    }
6395
6396    #[test]
6397    fn gompertz_hazard_shape_derivatives_small_shape_match_analytic_limit() {
6398        // At small x = shape·age the shape derivatives collapse to closed-form
6399        // limits. These MUST hold even for large ages with tiny shapes, which
6400        // is precisely the regime where the (cancelling) exact branch loses all
6401        // precision and the x-based pivot routes to the Taylor branch.
6402        //   ∂H_G/∂shape   -> rate·t²/2
6403        //   ∂h_G/∂shape   -> rate·t
6404        //   ∂²H_G/∂shape² -> rate·t³/3
6405        //   ∂²h_G/∂shape² -> rate·t²
6406        // The bug this guards: the second derivative's old `shape < 1e-10`
6407        // pivot ignored `age`, so e.g. (age=100, shape=1e-5 -> x=1e-3) took the
6408        // cancelling exact branch and returned a wildly wrong curvature.
6409        let cases = [
6410            (25.0_f64, 0.4_f64, 1e-9_f64),
6411            (100.0, 0.4, 1e-6),   // x = 1e-4
6412            (100.0, 0.012, 1e-6), // x = 1e-4, the old-pivot band (large age, tiny shape)
6413            (50.0, 1.2, 1e-8),
6414        ];
6415        // NOTE: every quantity below is compared against its shape->0 *limit*.
6416        // For the cancelling cumulative branches (∂H/∂shape, ∂²H/∂shape²,
6417        // ∂²h/∂shape²) the limit is the correct shape->0 target and the
6418        // implementation routes through Taylor in this band. But the
6419        // instantaneous first derivative ∂h_G/∂shape = rate·age·e^x carries NO
6420        // cancellation: it is exact, and its departure from the limit rate·t is
6421        // a genuine O(x) effect. At x=1e-3 that departure is ~1.2e-3 (> tol),
6422        // so the cases here keep x <= 1e-4 where the limit is a valid 1e-3
6423        // oracle for *all four* quantities. The cancelling-branch regression at
6424        // larger x is covered by gompertz_second_shape_derivative_is_accurate_in_old_pivot_gap.
6425        for &(age, rate, shape) in &cases {
6426            let t = age;
6427            let (d_cum, d_inst) = gompertz_cumulative_shape_derivative(age, rate, shape);
6428            assert_close(
6429                d_cum,
6430                rate * t * t / 2.0,
6431                1e-3,
6432                &format!("∂H_G/∂shape limit (age={age}, shape={shape})"),
6433            );
6434            assert_close(
6435                d_inst,
6436                rate * t,
6437                1e-3,
6438                &format!("∂h_G/∂shape limit (age={age}, shape={shape})"),
6439            );
6440
6441            let (d2_cum, d2_inst) = gompertz_cumulative_shape_second_derivative(age, rate, shape);
6442            assert_close(
6443                d2_cum,
6444                rate * t * t * t / 3.0,
6445                1e-3,
6446                &format!("∂²H_G/∂shape² limit (age={age}, shape={shape})"),
6447            );
6448            assert_close(
6449                d2_inst,
6450                rate * t * t,
6451                1e-3,
6452                &format!("∂²h_G/∂shape² limit (age={age}, shape={shape})"),
6453            );
6454        }
6455    }
6456
6457    #[test]
6458    fn gompertz_second_shape_derivative_is_accurate_in_old_pivot_gap() {
6459        // Regression: in the band shape ∈ [1e-10, ~1e-4] with a realistic age,
6460        // the OLD `shape < 1e-10` pivot sent ∂²H_G/∂shape² through the
6461        // catastrophically-cancelling exact branch. With age=100, shape=1e-9
6462        // (x=1e-7) the exact branch returned ~+5e1 vs the true ~rate·t³/3.
6463        // Assert the implementation now matches the closed-form limit to high
6464        // precision throughout that band, across several decades of shape.
6465        let age = 100.0;
6466        let rate = 0.4;
6467        let t = age;
6468        let truth = rate * t * t * t / 3.0; // 1.333e5
6469        // Start at shape=1e-5 (x=1e-3): below this the second derivative is,
6470        // to better than 1e-3 relative, equal to its shape->0 limit, so the
6471        // limit is a valid oracle. (At x=1e-2 the true value legitimately
6472        // departs from the limit by ~7e-3, which is a real O(x) correction,
6473        // not an error — so we do not extend the band up to shape=1e-4.)
6474        for k in 5..=12 {
6475            let shape = 10f64.powi(-(k as i32)); // 1e-5 .. 1e-12
6476            let (d2_cum, _) = gompertz_cumulative_shape_second_derivative(age, rate, shape);
6477            assert_close(
6478                d2_cum,
6479                truth,
6480                1e-3,
6481                &format!("∂²H_G/∂shape² in old-pivot gap (age={age}, shape=1e-{k})"),
6482            );
6483        }
6484    }
6485
6486    #[test]
6487    fn weibull_offset_partials_match_central_diff() {
6488        let cases = [
6489            (0.5_f64, 1.2_f64, 25.0_f64),
6490            (2.0, 0.8, 60.0),
6491            (0.1, 3.0, 10.0),
6492        ];
6493        for &(scale, shape, age) in &cases {
6494            let cfg = SurvivalBaselineConfig {
6495                target: SurvivalBaselineTarget::Weibull,
6496                scale: Some(scale),
6497                shape: Some(shape),
6498                rate: None,
6499                makeham: None,
6500            };
6501            let analytic = baseline_offset_theta_partials(age, &cfg)
6502                .expect("ok")
6503                .expect("nl");
6504            let fd = fd_baseline_offset(age, &cfg, &[1e-5, 1e-5]);
6505            assert_eq!(analytic.len(), 2);
6506            for k in 0..2 {
6507                assert_close(
6508                    analytic[k].0,
6509                    fd[k].0,
6510                    1e-7,
6511                    &format!("weibull ∂eta/∂θ[{k}] (scale={scale}, shape={shape}, age={age})"),
6512                );
6513                assert_close(
6514                    analytic[k].1,
6515                    fd[k].1,
6516                    1e-7,
6517                    &format!("weibull ∂o_D/∂θ[{k}] (scale={scale}, shape={shape}, age={age})"),
6518                );
6519            }
6520            // Weibull o_D = shape/t is independent of scale; verify exactly.
6521            assert_eq!(analytic[0].1, 0.0);
6522        }
6523    }
6524
6525    #[test]
6526    fn gompertz_makeham_offset_partials_match_central_diff() {
6527        let cases = [
6528            (0.3_f64, 0.05_f64, 0.002_f64, 40.0_f64),
6529            (0.5, 0.01, 0.01, 25.0),
6530            (0.2, 0.001, 0.005, 60.0),
6531            (0.4, 5e-11, 0.01, 25.0),
6532            (0.4, -5e-11, 0.01, 25.0),
6533            (0.8, 0.2, 0.05, 5.0),
6534        ];
6535        for &(rate, shape, makeham, age) in &cases {
6536            let cfg = SurvivalBaselineConfig {
6537                target: SurvivalBaselineTarget::GompertzMakeham,
6538                scale: None,
6539                shape: Some(shape),
6540                rate: Some(rate),
6541                makeham: Some(makeham),
6542            };
6543            let analytic = baseline_offset_theta_partials(age, &cfg)
6544                .expect("ok")
6545                .expect("nl");
6546            // See gompertz_offset_partials_match_central_diff: tiny shape-step
6547            // is only needed for the shape component; log_rate and
6548            // log_makeham take the normal-scale step.
6549            let h_shape = if shape.abs() < 1e-9 { 1e-11 } else { 1e-5 };
6550            let fd = fd_baseline_offset(age, &cfg, &[1e-5, h_shape, 1e-5]);
6551            assert_eq!(analytic.len(), 3);
6552            for k in 0..3 {
6553                assert_close(
6554                    analytic[k].0,
6555                    fd[k].0,
6556                    1e-5,
6557                    &format!(
6558                        "gm ∂eta/∂θ[{k}] (rate={rate}, shape={shape}, mk={makeham}, age={age})"
6559                    ),
6560                );
6561                assert_close(
6562                    analytic[k].1,
6563                    fd[k].1,
6564                    1e-5,
6565                    &format!(
6566                        "gm ∂o_D/∂θ[{k}] (rate={rate}, shape={shape}, mk={makeham}, age={age})"
6567                    ),
6568                );
6569            }
6570        }
6571    }
6572
6573    #[test]
6574    fn linear_baseline_has_no_theta_partials() {
6575        let cfg = SurvivalBaselineConfig {
6576            target: SurvivalBaselineTarget::Linear,
6577            scale: None,
6578            shape: None,
6579            rate: None,
6580            makeham: None,
6581        };
6582        assert!(baseline_offset_theta_partials(5.0, &cfg).unwrap().is_none());
6583    }
6584
6585    #[test]
6586    fn baseline_offset_partials_reject_non_positive_ages() {
6587        let cfg = SurvivalBaselineConfig {
6588            target: SurvivalBaselineTarget::Gompertz,
6589            scale: None,
6590            shape: Some(0.01),
6591            rate: Some(0.5),
6592            makeham: None,
6593        };
6594        assert!(baseline_offset_theta_partials(0.0, &cfg).is_err());
6595        assert!(baseline_offset_theta_partials(-1.0, &cfg).is_err());
6596        assert!(baseline_offset_theta_partials(f64::NAN, &cfg).is_err());
6597    }
6598
6599    // ─── baseline_chain_rule_gradient — mechanical and FD-vs-θ tests ─────
6600
6601    /// Mechanical sanity check: with only one event observation at known
6602    /// (r_X, r_E, r_D, age_exit, age_entry), the Gompertz chain-rule gradient
6603    /// reduces to the analytic linear combination of `baseline_offset_theta_partials`.
6604    #[test]
6605    fn chain_rule_gradient_single_obs_reduces_to_pointwise_contract() {
6606        let cfg = SurvivalBaselineConfig {
6607            target: SurvivalBaselineTarget::Gompertz,
6608            scale: None,
6609            shape: Some(0.05),
6610            rate: Some(0.3),
6611            makeham: None,
6612        };
6613        let age_entry = array![10.0_f64];
6614        let age_exit = array![25.0_f64];
6615        let residuals = OffsetChannelResiduals {
6616            exit: array![0.7_f64],
6617            entry: array![-0.2_f64],
6618            derivative: array![-0.4_f64],
6619            right: Array1::<f64>::zeros(1),
6620        };
6621        let grad = baseline_chain_rule_gradient(
6622            age_entry.view(),
6623            age_exit.view(),
6624            age_exit.view(),
6625            &cfg,
6626            &residuals,
6627        )
6628        .expect("ok")
6629        .expect("non-linear");
6630        // Hand-compute: grad[k] = r_X·∂eta_exit/∂θ_k + r_D·∂o_D_exit/∂θ_k + r_E·∂eta_entry/∂θ_k.
6631        let p_exit = baseline_offset_theta_partials(age_exit[0], &cfg)
6632            .unwrap()
6633            .unwrap();
6634        let p_entry = baseline_offset_theta_partials(age_entry[0], &cfg)
6635            .unwrap()
6636            .unwrap();
6637        for k in 0..p_exit.len() {
6638            let expected = 0.7 * p_exit[k].0 + (-0.4) * p_exit[k].1 + (-0.2) * p_entry[k].0;
6639            assert!(
6640                (grad[k] - expected).abs() < 1e-12,
6641                "chain-rule contract mismatch at k={k}: got={:.6e} expected={:.6e}",
6642                grad[k],
6643                expected
6644            );
6645        }
6646    }
6647
6648    /// Origin-entry rows (r_entry == 0) must skip the baseline partials call at
6649    /// `age_entry = 0`, which would otherwise fail the positive-age precondition.
6650    #[test]
6651    fn chain_rule_gradient_skips_entry_call_for_origin_entry_rows() {
6652        let cfg = SurvivalBaselineConfig {
6653            target: SurvivalBaselineTarget::Gompertz,
6654            scale: None,
6655            shape: Some(0.05),
6656            rate: Some(0.3),
6657            makeham: None,
6658        };
6659        let age_entry = array![0.0_f64, 5.0_f64];
6660        let age_exit = array![10.0_f64, 20.0_f64];
6661        let residuals = OffsetChannelResiduals {
6662            exit: array![0.5_f64, 0.3_f64],
6663            entry: array![0.0_f64, -0.1_f64], // row 0 is origin-entry (r_E = 0)
6664            derivative: array![-0.2_f64, 0.0_f64],
6665            right: Array1::<f64>::zeros(2),
6666        };
6667        // Must not error despite age_entry[0] == 0.
6668        let grad = baseline_chain_rule_gradient(
6669            age_entry.view(),
6670            age_exit.view(),
6671            age_exit.view(),
6672            &cfg,
6673            &residuals,
6674        )
6675        .expect("must not fail on origin-entry row with r_entry=0")
6676        .expect("non-linear");
6677        assert_eq!(grad.len(), 2);
6678        // Row 1's entry channel contributes, row 0's does not.
6679        let p_exit_0 = baseline_offset_theta_partials(10.0, &cfg).unwrap().unwrap();
6680        let p_exit_1 = baseline_offset_theta_partials(20.0, &cfg).unwrap().unwrap();
6681        let p_entry_1 = baseline_offset_theta_partials(5.0, &cfg).unwrap().unwrap();
6682        for k in 0..2 {
6683            let expected = 0.5 * p_exit_0[k].0
6684                + (-0.2) * p_exit_0[k].1
6685                + 0.3 * p_exit_1[k].0
6686                + (-0.1) * p_entry_1[k].0;
6687            assert!(
6688                (grad[k] - expected).abs() < 1e-12,
6689                "origin-entry contract at k={k}: got={:.6e} expected={:.6e}",
6690                grad[k],
6691                expected
6692            );
6693        }
6694    }
6695
6696    /// Linear target has no θ-parameters; contractor returns None.
6697    #[test]
6698    fn chain_rule_gradient_linear_target_returns_none() {
6699        let cfg = SurvivalBaselineConfig {
6700            target: SurvivalBaselineTarget::Linear,
6701            scale: None,
6702            shape: None,
6703            rate: None,
6704            makeham: None,
6705        };
6706        let age_entry = array![1.0_f64];
6707        let age_exit = array![2.0_f64];
6708        let residuals = OffsetChannelResiduals {
6709            exit: array![0.1_f64],
6710            entry: array![0.0_f64],
6711            derivative: array![0.0_f64],
6712            right: Array1::<f64>::zeros(1),
6713        };
6714        let grad = baseline_chain_rule_gradient(
6715            age_entry.view(),
6716            age_exit.view(),
6717            age_exit.view(),
6718            &cfg,
6719            &residuals,
6720        )
6721        .expect("ok");
6722        assert!(grad.is_none());
6723    }
6724
6725    /// End-to-end envelope-theorem check: the chain-rule gradient at
6726    /// residuals-evaluated-at-β-fixed matches the central FD of the
6727    /// unpenalized NLL with respect to θ when the OFFSETS are recomputed
6728    /// from the perturbed cfg and β is held at its base value.
6729    ///
6730    /// This is the mathematical content of the envelope theorem applied to
6731    /// the penalized-deviance cost at fixed β: if β solves ∂C/∂β = 0 at
6732    /// (θ, β*), then the total derivative of C at (θ±h) when β is held at
6733    /// β* equals the partial derivative of C wrt θ at the base — up to
6734    /// O(h²) in the truncation error of central differences. For THIS test
6735    /// we're directly differencing NLL (the unpenalized piece that carries
6736    /// all the θ dependence), so the envelope identity is exact up to FD
6737    /// truncation.
6738    ///
6739    /// The test synthesizes a plausible residual set by hand rather than
6740    /// running PIRLS — what we're validating is the chain-rule contractor,
6741    /// not the fit. A PIRLS-based end-to-end check belongs in an
6742    /// integration test, not this unit-test module.
6743    #[test]
6744    fn chain_rule_gradient_matches_fd_of_nll_through_offset_perturbation() {
6745        // Toy 3-observation case with two events (one origin-entry, one not)
6746        // and one censored row at large age.
6747        let cfg = SurvivalBaselineConfig {
6748            target: SurvivalBaselineTarget::Gompertz,
6749            scale: None,
6750            shape: Some(0.03),
6751            rate: Some(0.25),
6752            makeham: None,
6753        };
6754        let age_entry = array![0.0_f64, 5.0, 8.0];
6755        let age_exit = array![4.0_f64, 12.0, 20.0];
6756        // Weighted residuals at a notional β*. Values chosen in a plausible
6757        // range (~same order as w·exp(η)).
6758        let weights = array![1.0_f64, 2.0, 0.5];
6759        let events = [1.0_f64, 1.0, 0.0];
6760        // Fake a β* that yields finite eta_entry ± eta_exit ± s values by
6761        // directly specifying eta quantities. Contractor only consumes the
6762        // residuals, so the fake is sufficient.
6763        let eta_entry_vals = [-100.0_f64, 0.5, 0.8]; // row 0 doesn't matter (origin entry)
6764        let eta_exit_vals = [0.4_f64, 0.9, 1.3];
6765        let s_vals = [0.7_f64, 1.1, 1.5];
6766        let (r_x, r_e, r_d) = {
6767            let mut rx = Array1::<f64>::zeros(3);
6768            let mut re = Array1::<f64>::zeros(3);
6769            let mut rd = Array1::<f64>::zeros(3);
6770            for i in 0..3 {
6771                let w = weights[i];
6772                let d = events[i];
6773                rx[i] = w * (eta_exit_vals[i].exp() - d);
6774                re[i] = if i == 0 {
6775                    0.0 // origin entry
6776                } else {
6777                    -w * eta_entry_vals[i].exp()
6778                };
6779                rd[i] = if d > 0.0 { -w * d / s_vals[i] } else { 0.0 };
6780            }
6781            (rx, re, rd)
6782        };
6783        let residuals = OffsetChannelResiduals {
6784            exit: r_x.clone(),
6785            entry: r_e.clone(),
6786            derivative: r_d.clone(),
6787            right: Array1::<f64>::zeros(3),
6788        };
6789        let grad = baseline_chain_rule_gradient(
6790            age_entry.view(),
6791            age_exit.view(),
6792            age_exit.view(),
6793            &cfg,
6794            &residuals,
6795        )
6796        .expect("ok")
6797        .expect("non-linear");
6798
6799        // Construct NLL(θ) with β* held to the same eta/s values by treating
6800        // eta_i, s_i as fixed "linear predictor" samples and shifting by
6801        // (offset(θ) - offset(θ_base)). That's exactly the RP NLL with β*
6802        // held constant and offsets varied through θ.
6803        let nll = |theta_plus: &Array1<f64>| -> f64 {
6804            let cfg_p = survival_baseline_config_from_theta(cfg.target, theta_plus).expect("cfg_p");
6805            let mut sum = 0.0_f64;
6806            for i in 0..3 {
6807                let (eta_x_p, d_x_p) = evaluate_survival_baseline(age_exit[i], &cfg_p).unwrap();
6808                let base = evaluate_survival_baseline(age_exit[i], &cfg).unwrap();
6809                let d_eta_x = eta_x_p - base.0;
6810                let d_d_x = d_x_p - base.1;
6811                let eta_exit_new = eta_exit_vals[i] + d_eta_x;
6812                let s_new = s_vals[i] + d_d_x;
6813                let interval_entry = if i == 0 {
6814                    0.0_f64
6815                } else {
6816                    let (eta_e_p, _) = evaluate_survival_baseline(age_entry[i], &cfg_p).unwrap();
6817                    let base_e = evaluate_survival_baseline(age_entry[i], &cfg).unwrap();
6818                    let d_eta_e = eta_e_p - base_e.0;
6819                    let eta_entry_new = eta_entry_vals[i] + d_eta_e;
6820                    eta_entry_new.exp()
6821                };
6822                let w = weights[i];
6823                let d = events[i];
6824                let nll_i =
6825                    w * (eta_exit_new.exp() - interval_entry - d * (eta_exit_new + s_new.ln()));
6826                sum += nll_i;
6827            }
6828            sum
6829        };
6830
6831        let theta_base = survival_baseline_theta_from_config(&cfg).unwrap().unwrap();
6832        let h = 1e-6;
6833        for k in 0..theta_base.len() {
6834            let mut tp = theta_base.clone();
6835            let mut tm = theta_base.clone();
6836            tp[k] += h;
6837            tm[k] -= h;
6838            let fd = (nll(&tp) - nll(&tm)) / (2.0 * h);
6839            assert!(
6840                (grad[k] - fd).abs() < 1e-5 * grad[k].abs().max(1.0),
6841                "chain-rule θ[{k}]: analytic={:.6e} fd={:.6e}",
6842                grad[k],
6843                fd
6844            );
6845        }
6846    }
6847
6848    /// Length-mismatch surfaces as an error, not a silent contraction.
6849    #[test]
6850    fn chain_rule_gradient_rejects_length_mismatch() {
6851        let cfg = SurvivalBaselineConfig {
6852            target: SurvivalBaselineTarget::Gompertz,
6853            scale: None,
6854            shape: Some(0.05),
6855            rate: Some(0.3),
6856            makeham: None,
6857        };
6858        let age_entry = array![1.0_f64, 2.0]; // length 2
6859        let age_exit = array![5.0_f64, 6.0, 7.0]; // length 3
6860        let residuals = OffsetChannelResiduals {
6861            exit: array![0.1_f64, 0.2, 0.3],
6862            entry: array![0.0_f64, 0.0, 0.0],
6863            derivative: array![0.0_f64, 0.0, 0.0],
6864            right: Array1::<f64>::zeros(3),
6865        };
6866        let err = baseline_chain_rule_gradient(
6867            age_entry.view(),
6868            age_exit.view(),
6869            age_exit.view(),
6870            &cfg,
6871            &residuals,
6872        )
6873        .expect_err("length mismatch must error");
6874        assert!(err.contains("length mismatch"), "err={err}");
6875    }
6876
6877    // ── gam#2765 / gam#2767: the log-slope follow-up margin replays exactly ──
6878
6879    /// The predict-time replay must reproduce the fit-time margin bit for bit.
6880    ///
6881    /// This is the property the whole persistence contract rests on: at fit time
6882    /// the knots are placed by QUANTILE from the training exit times, and the
6883    /// design is a by-product of that build; at predict time only the knots
6884    /// survive, and the design is rebuilt from them. If those two designs are not
6885    /// the same matrix on the same rows, every saved follow-up-varying slope
6886    /// evaluates a different model than the one that was fitted — silently,
6887    /// because the widths still agree.
6888    #[test]
6889    fn logslope_time_margin_replay_reproduces_the_fit_time_design_2765() {
6890        let age_exit = Array1::from_iter((1..=40).map(|i| 0.25 + 0.35 * f64::from(i)));
6891        let age_entry = age_exit.mapv(|t| (t - 0.2).max(1e-3));
6892        let fitted = build_time_varying_survival_covariate_template(
6893            &age_entry,
6894            &age_exit,
6895            5,
6896            3,
6897            "logslope",
6898        )
6899        .expect("fit-time log-slope margin");
6900        let SurvivalCovariateTermBlockTemplate::TimeVarying {
6901            time_basis,
6902            time_basis_entry,
6903            time_basis_exit,
6904            time_basis_derivative_exit,
6905            ..
6906        } = &fitted
6907        else {
6908            panic!("a time-varying request must produce a time-varying template");
6909        };
6910
6911        let replayed_exit = logslope_time_margin_rows(time_basis, age_exit.view())
6912            .expect("replayed exit margin");
6913        assert_eq!(replayed_exit.dim(), time_basis_exit.dim());
6914        for (fit_value, replay_value) in time_basis_exit.iter().zip(replayed_exit.iter()) {
6915            assert_eq!(
6916                fit_value.to_bits(),
6917                replay_value.to_bits(),
6918                "the replayed exit margin must be the fitted one, not merely close"
6919            );
6920        }
6921
6922        // And the full three-channel replay the leave-one-out path consumes.
6923        let covariate = DesignMatrix::from(Array2::<f64>::from_shape_fn(
6924            (age_exit.len(), 2),
6925            |(row, col)| if col == 0 { 1.0 } else { (row as f64) * 0.05 - 1.0 },
6926        ));
6927        let replay =
6928            replay_logslope_follow_up_designs(&age_entry, &age_exit, time_basis, &covariate)
6929                .expect("three-channel replay");
6930        let p_time = time_basis_exit.ncols();
6931        assert_eq!(replay.exit.ncols(), 2 * p_time);
6932        for (channel, fitted_margin) in [
6933            (&replay.entry, time_basis_entry),
6934            (&replay.exit, time_basis_exit),
6935            (&replay.derivative_exit, time_basis_derivative_exit),
6936        ] {
6937            let dense = channel
6938                .try_to_dense_arc("replayed log-slope channel")
6939                .expect("dense channel");
6940            let covariate_dense = covariate
6941                .try_to_dense_arc("covariate factor")
6942                .expect("dense covariate");
6943            for row in 0..age_exit.len() {
6944                for cov_col in 0..2 {
6945                    for time_col in 0..p_time {
6946                        let expected =
6947                            covariate_dense[[row, cov_col]] * fitted_margin[[row, time_col]];
6948                        let got = dense[[row, cov_col * p_time + time_col]];
6949                        assert!(
6950                            (expected - got).abs() <= 1e-15 * (1.0 + expected.abs()),
6951                            "row-wise Kronecker mismatch at ({row}, {cov_col}, {time_col}): \
6952                             expected {expected} got {got}"
6953                        );
6954                    }
6955                }
6956            }
6957        }
6958    }
6959
6960    /// The exit design a batch replay produces is the same one a single-row
6961    /// replay produces at that row's time. The survival-curve path replays one
6962    /// `(row, t)` cell at a time, so if these disagreed a predicted curve would
6963    /// not pass through the batch-predicted point.
6964    #[test]
6965    fn logslope_time_margin_row_replay_matches_the_batch_replay_2765() {
6966        let age_exit = Array1::from_iter((1..=12).map(|i| 0.4 + 0.6 * f64::from(i)));
6967        let age_entry = age_exit.mapv(|t| (t - 0.15).max(1e-3));
6968        let fitted =
6969            build_time_varying_survival_covariate_template(&age_entry, &age_exit, 6, 2, "logslope")
6970                .expect("fit-time log-slope margin");
6971        let time_basis = fitted
6972            .resolved_time_basis()
6973            .expect("a time-varying template resolves a basis")
6974            .clone();
6975        let covariate = DesignMatrix::from(Array2::<f64>::from_shape_fn(
6976            (age_exit.len(), 2),
6977            |(row, col)| if col == 0 { 1.0 } else { 0.3 * (row as f64) },
6978        ));
6979        let batch = replay_logslope_time_margin_design(age_exit.view(), &time_basis, &covariate)
6980            .expect("batch replay")
6981            .try_to_dense_arc("batch replay")
6982            .expect("dense batch");
6983        let covariate_dense = covariate
6984            .try_to_dense_arc("covariate")
6985            .expect("dense covariate");
6986        for row in 0..age_exit.len() {
6987            let single_covariate = DesignMatrix::from(
6988                covariate_dense
6989                    .row(row)
6990                    .to_owned()
6991                    .into_shape_with_order((1, 2))
6992                    .expect("single covariate row"),
6993            );
6994            let single = replay_logslope_time_margin_design(
6995                Array1::from_elem(1, age_exit[row]).view(),
6996                &time_basis,
6997                &single_covariate,
6998            )
6999            .expect("single-row replay")
7000            .try_to_dense_arc("single-row replay")
7001            .expect("dense single row");
7002            for col in 0..batch.ncols() {
7003                assert_eq!(
7004                    batch[[row, col]].to_bits(),
7005                    single[[0, col]].to_bits(),
7006                    "single-row replay disagrees with the batch at ({row}, {col})"
7007                );
7008            }
7009        }
7010    }
7011
7012    /// The defect gam#2705 names, at the level it is created: the survival
7013    /// I-spline time block's `x_derivative_time` must be the derivative of its
7014    /// own `x_exit_time`, OUTSIDE the fitted knot span as well as inside it.
7015    ///
7016    /// Before the repair the value basis saturated past the boundary knots
7017    /// while the derivative — hand-rolled from a CLAMPED B-spline
7018    /// first-derivative basis — returned the boundary slope, so a saved
7019    /// Royston-Parmar fit published a flat `Λ(t)` beside a nonzero
7020    /// `h(t) = Λ·d(log Λ)/dt`. `h = dΛ/dt`, so those cannot both be one model.
7021    ///
7022    /// The check is a central difference in `t` (not in `log t`), because `t`
7023    /// is the variable `x_derivative_time` is a derivative with respect to —
7024    /// the `1/t` chain factor is part of what has to agree.
7025    #[test]
7026    fn ispline_time_derivative_is_a_finite_difference_of_its_value_2705() {
7027        let n = 24usize;
7028        let age_entry = Array1::<f64>::zeros(n);
7029        let age_exit =
7030            Array1::from_iter((0..n).map(|i| 4.0 + 40.0 * (i as f64) / ((n - 1) as f64)));
7031        let build = build_survival_time_basis(
7032            &age_entry,
7033            &age_exit,
7034            SurvivalTimeBasisConfig::ISpline {
7035                degree: 3,
7036                knots: Array1::zeros(0),
7037                keep_cols: Vec::new(),
7038                smooth_lambda: 1.0,
7039            },
7040            Some((3, 1.0)),
7041        )
7042        .expect("ispline time basis builds");
7043        let resolved = resolved_survival_time_basis_config_from_build(
7044            &build.basisname,
7045            build.degree,
7046            build.knots.as_ref(),
7047            build.keep_cols.as_ref(),
7048            build.smooth_lambda,
7049        )
7050        .expect("resolved ispline config");
7051
7052        // Inside the fitted span, far below it, and far above it — the last two
7053        // are where every `predict(...).survival_at(grid)` call lands, because
7054        // `default_survival_time_grid` starts at 0 and ends past `max(exit)`.
7055        //
7056        // The two BOUNDARY knots themselves (`t = 4` and `t = 44`, the extreme
7057        // training exits) are deliberately not central-differenced, and the
7058        // boundary is checked exactly below instead. An M-spline has a
7059        // one-sided kink at a clamped boundary knot — the trailing columns
7060        // vanish there like `(t_b − t)^k` — so the analytic derivative is the
7061        // one-sided LIMIT (zero for those columns) while a symmetric window of
7062        // half-width `h` averages the rising side and returns `O(h)`. Measured
7063        // at `t = 44`, `h = 1e-5`: analytic `0`, central difference `8.99e-8`,
7064        // which shrinks with `h` rather than marking a disagreement.
7065        let queries = [1.0_f64, 3.0, 12.0, 30.0, 43.0, 60.0, 400.0, 2_850.0];
7066        let step = 1.0e-5_f64;
7067        let mut exterior_rows_with_slope = 0usize;
7068        for &t in queries.iter() {
7069            let times = Array1::from_vec(vec![t - step, t, t + step]);
7070            let probe =
7071                build_survival_time_basis(&Array1::<f64>::zeros(3), &times, resolved.clone(), None)
7072                    .expect("ispline time basis replays at the query times");
7073            let value = probe.x_exit_time.to_dense();
7074            let derivative = probe.x_derivative_time.to_dense();
7075            let mut row_slope = 0.0_f64;
7076            for column in 0..value.ncols() {
7077                let difference = (value[[2, column]] - value[[0, column]]) / (2.0 * step);
7078                let analytic = derivative[[1, column]];
7079                row_slope += analytic.abs();
7080                let scale = analytic.abs().max(difference.abs()).max(1.0e-6);
7081                assert!(
7082                    (difference - analytic).abs() <= 1.0e-4 * scale,
7083                    "t={t}: column {column} analytic d/dt {analytic:.9e} disagrees with the \
7084                     central difference of its own value basis {difference:.9e}"
7085                );
7086            }
7087            if !(4.0..=44.0).contains(&t) {
7088                exterior_rows_with_slope += usize::from(row_slope > 0.0);
7089            }
7090        }
7091        // Non-vacuity: a basis whose exterior derivative is identically zero
7092        // passes every assertion above by agreeing with a flat value. The
7093        // Royston-Parmar tail is LINEAR, so the exterior slope is the boundary
7094        // slope and is nonzero on both sides.
7095        assert!(
7096            exterior_rows_with_slope >= 2,
7097            "the exterior must carry a nonzero boundary slope on both sides; \
7098             {exterior_rows_with_slope} of the exterior query times did"
7099        );
7100
7101        // The boundary itself, exactly rather than by difference: the tail is
7102        // ANCHORED at the spline's own one-sided value and slope there, so
7103        // crossing `t_b` must not move the derivative at all and must move the
7104        // value by exactly `Δ(log t) · slope`.
7105        let boundary = 44.0_f64;
7106        let outside = 44.05_f64;
7107        let pair = build_survival_time_basis(
7108            &Array1::<f64>::zeros(2),
7109            &Array1::from_vec(vec![boundary, outside]),
7110            resolved.clone(),
7111            None,
7112        )
7113        .expect("ispline time basis replays across the boundary knot");
7114        let value = pair.x_exit_time.to_dense();
7115        let derivative = pair.x_derivative_time.to_dense();
7116        // `x_derivative_time` carries the `1/t` chain factor, so undo it to
7117        // compare slopes in the basis's own `log t` coordinate.
7118        let log_gap = outside.ln() - boundary.ln();
7119        for column in 0..value.ncols() {
7120            let boundary_slope = derivative[[0, column]] * boundary;
7121            let outside_slope = derivative[[1, column]] * outside;
7122            assert!(
7123                (outside_slope - boundary_slope).abs() <= 1.0e-12 * boundary_slope.abs().max(1.0),
7124                "column {column}: the tail slope {outside_slope:.9e} is not the boundary slope \
7125                 {boundary_slope:.9e}"
7126            );
7127            let expected = value[[0, column]] + log_gap * boundary_slope;
7128            assert!(
7129                (value[[1, column]] - expected).abs() <= 1.0e-12 * expected.abs().max(1.0),
7130                "column {column}: the tail value {:.9e} is not the affine continuation \
7131                 {expected:.9e} of the boundary value {:.9e} at slope {boundary_slope:.9e}",
7132                value[[1, column]],
7133                value[[0, column]]
7134            );
7135        }
7136    }
7137
7138    /// The fit itself must not move. Every training row is inside the knot span
7139    /// the training rows themselves induced, so the linear-tail convention is
7140    /// inert there — and a row entering AT THE ORIGIN keeps the anchored zero
7141    /// entry row rather than a tail evaluated at `ln(SURVIVAL_TIME_FLOOR)`,
7142    /// which is a readout of `1e-9` and not of the data.
7143    #[test]
7144    fn the_linear_tail_convention_is_inert_on_the_training_rows_2705() {
7145        let n = 16usize;
7146        let age_entry = Array1::<f64>::zeros(n);
7147        let age_exit =
7148            Array1::from_iter((0..n).map(|i| 2.0 + 20.0 * (i as f64) / ((n - 1) as f64)));
7149        let build = build_survival_time_basis(
7150            &age_entry,
7151            &age_exit,
7152            SurvivalTimeBasisConfig::ISpline {
7153                degree: 3,
7154                knots: Array1::zeros(0),
7155                keep_cols: Vec::new(),
7156                smooth_lambda: 1.0,
7157            },
7158            Some((3, 1.0)),
7159        )
7160        .expect("ispline time basis builds");
7161        let entry = build.x_entry_time.to_dense();
7162        let exit = build.x_exit_time.to_dense();
7163        for row in 0..n {
7164            for column in 0..entry.ncols() {
7165                assert_eq!(
7166                    entry[[row, column]],
7167                    0.0,
7168                    "an entry-at-origin row must carry the anchored zero row at ({row}, {column})"
7169                );
7170            }
7171        }
7172        // The exit rows are I-spline values on their own knot span, so every
7173        // entry stays in [0, 1]: no tail is being evaluated at a training row.
7174        for row in 0..n {
7175            for column in 0..exit.ncols() {
7176                let value = exit[[row, column]];
7177                assert!(
7178                    (-1.0e-12..=1.0 + 1.0e-12).contains(&value),
7179                    "training exit row ({row}, {column}) = {value} is outside [0, 1], so the \
7180                     linear tail is being evaluated on the training data"
7181                );
7182            }
7183        }
7184    }
7185
7186    /// The anchor is the ORIGIN of the baseline reparameterization, and the
7187    /// default anchor for ordinary right-censored data is the time origin,
7188    /// which `evaluate_survival_time_basis_row` floors to `SURVIVAL_TIME_FLOOR`.
7189    /// Under a linear-tailed baseline an unclamped anchor there would re-center
7190    /// every design column by a large constant read off `1e-9` — the #751
7191    /// inflation the anchor rule exists to avoid. Clamping into the modelling
7192    /// interval keeps the shipped answer: `I_k(left) = 0` exactly.
7193    #[test]
7194    fn the_anchor_row_is_clamped_into_the_modelling_interval_2705() {
7195        let n = 16usize;
7196        let age_entry = Array1::<f64>::zeros(n);
7197        let age_exit =
7198            Array1::from_iter((0..n).map(|i| 5.0 + 50.0 * (i as f64) / ((n - 1) as f64)));
7199        let build = build_survival_time_basis(
7200            &age_entry,
7201            &age_exit,
7202            SurvivalTimeBasisConfig::ISpline {
7203                degree: 3,
7204                knots: Array1::zeros(0),
7205                keep_cols: Vec::new(),
7206                smooth_lambda: 1.0,
7207            },
7208            Some((3, 1.0)),
7209        )
7210        .expect("ispline time basis builds");
7211        let resolved = resolved_survival_time_basis_config_from_build(
7212            &build.basisname,
7213            build.degree,
7214            build.knots.as_ref(),
7215            build.keep_cols.as_ref(),
7216            build.smooth_lambda,
7217        )
7218        .expect("resolved ispline config");
7219        let anchor_at_origin = evaluate_survival_time_basis_row(0.0, &resolved)
7220            .expect("anchor row at the time origin");
7221        for (column, value) in anchor_at_origin.iter().enumerate() {
7222            assert_eq!(
7223                *value, 0.0,
7224                "the anchor row at the time origin must be exactly zero at column {column}, \
7225                 got {value}"
7226            );
7227        }
7228        // And an anchor INSIDE the span is still a real evaluation, so the
7229        // clamp has not turned the anchor into a constant.
7230        let interior =
7231            evaluate_survival_time_basis_row(30.0, &resolved).expect("anchor row inside the span");
7232        assert!(
7233            interior.iter().any(|value| *value > 1.0e-9),
7234            "an interior anchor must still evaluate the basis, got {interior:?}"
7235        );
7236    }
7237}