1use 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#[derive(Clone, Debug)]
52pub enum SurvivalConstructionError {
53 InvalidConfig { reason: String },
56 MissingColumn { reason: String },
59 IncompatibleDimensions { reason: String },
62 DataValidationFailed { reason: String },
66 BasisConstructionFailed { reason: String },
70 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
91pub enum SurvivalBaselineTarget {
92 Linear,
96 Weibull,
101 Gompertz,
106 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
122pub 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 ISpline {
197 degree: usize,
198 knots: Array1<f64>,
199 keep_cols: Vec<usize>,
200 smooth_lambda: f64,
201 },
202}
203
204#[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 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 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
259const SURVIVAL_TIME_SMOOTH_LAMBDA_SEED: f64 = 1e-2;
267
268const 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
287pub 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
308pub 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
335pub 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
361pub 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
459pub 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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
725enum BaselineDerivativeContract {
726 GradientOnly,
729 GradientHessian,
732}
733
734impl BaselineDerivativeContract {
735 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 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
761fn 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
832fn 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
892pub 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
927pub 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
956pub 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 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
1016pub 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 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 return 1.min(requested_internal_knots);
1123 }
1124 let scale = (maxval - minval).abs().max(1.0);
1125 let tol = 1e-12 * scale;
1126 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 let mut cap = requested_internal_knots.min(distinct_interior.max(1));
1143 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 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 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 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 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 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 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 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 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 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 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 drop(x_entry_full);
1669 drop(x_exit_full);
1670
1671 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 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 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 let mut s_full = s_increment.to_owned();
1817 symmetrize_in_place(&mut s_full);
1818 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 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 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 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 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 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
1969pub 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
2016pub 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
2031pub 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
2067pub 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
2090pub 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 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 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 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
2253pub 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 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 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 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 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 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 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
2383fn 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 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 let mut grad = Array1::<f64>::zeros(theta_dim);
2464 for i in 0..n {
2465 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 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 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
2520pub 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
2571pub 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 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#[inline]
2601fn gompertz_hazard_components(age: f64, rate: f64, shape: f64) -> (f64, f64) {
2602 if shape.abs() < 1e-10 {
2603 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#[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 let dhg_dshape = if x.abs() < 1e-4 {
2646 let t = age;
2647 rate * t * t * (0.5 + x / 3.0 + x * x / 8.0)
2649 } else {
2650 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#[inline]
2664fn gompertz_shape_derivatives(age: f64, shape: f64) -> (f64, f64) {
2665 if shape.abs() < 1e-10 {
2666 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(); let d_eta = -1.0 / shape + age * e / em1;
2685 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#[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
2707fn 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
2913pub 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 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
2986pub fn evaluate_survival_marginal_slope_baseline(
2992 age: f64,
2993 cfg: &SurvivalBaselineConfig,
2994) -> Result<(f64, f64), String> {
2995 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
3015pub 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
3051pub 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 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#[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 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#[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 let triples: Vec<(f64, f64, f64)> = (0..n)
3378 .into_par_iter()
3379 .map(|i| -> Result<(f64, f64, f64), String> {
3380 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
3416pub 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
3431pub 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#[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
3499pub 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 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
3525pub 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 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#[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 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 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
3815pub 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 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 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
4024pub 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
4045pub 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 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 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 let knots = gam_terms::basis::initializewiggle_knots_from_seed(
4096 seed.view(),
4097 cfg.degree,
4098 cfg.num_internal_knots,
4099 )?;
4100 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 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
4150pub 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
4220pub 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
4261pub 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
4293pub struct LogslopeFollowUpReplayDesigns {
4303 pub entry: DesignMatrix,
4304 pub exit: DesignMatrix,
4305 pub derivative_exit: DesignMatrix,
4306}
4307
4308pub 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
4353pub 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 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 assert!(
4515 fitted_weibull_baseline_from_linear_time_beta(&Array1::<f64>::zeros(0), 4.5).is_none()
4516 );
4517 assert!(fitted_weibull_baseline_from_linear_time_beta(&array![0.0], 4.5).is_none());
4519 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 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 #[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 assert!(
4731 (anchor - 19.0).abs() <= 1e-12,
4732 "marginal-slope default anchor should be median exit, got {anchor}"
4733 );
4734 }
4735
4736 #[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 #[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 #[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 #[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 #[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 #[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 #[test]
4873 fn robust_interior_anchor_shrinks_and_balances_the_centered_time_design() {
4874 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 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 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 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 #[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 #[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 #[test]
5036 fn baseline_optimizer_contracts_agree_on_shared_surface() {
5037 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 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 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 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 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 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 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 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 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 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 let knots = array![
5298 left, left, left, left, q1, mid, q3, right, right, right, right
5299 ];
5300
5301 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 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 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 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 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 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 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 #[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 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 #[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 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 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 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 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 #[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 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 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 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 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 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 #[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 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 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 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 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 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 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 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 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 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 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 #[test]
6347 fn gompertz_hazard_shape_derivatives_match_central_diff() {
6348 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 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 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 let cases = [
6410 (25.0_f64, 0.4_f64, 1e-9_f64),
6411 (100.0, 0.4, 1e-6), (100.0, 0.012, 1e-6), (50.0, 1.2, 1e-8),
6414 ];
6415 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 let age = 100.0;
6466 let rate = 0.4;
6467 let t = age;
6468 let truth = rate * t * t * t / 3.0; for k in 5..=12 {
6475 let shape = 10f64.powi(-(k as i32)); 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 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 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 #[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 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 #[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], derivative: array![-0.2_f64, 0.0_f64],
6665 right: Array1::<f64>::zeros(2),
6666 };
6667 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 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 #[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 #[test]
6744 fn chain_rule_gradient_matches_fd_of_nll_through_offset_perturbation() {
6745 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 let weights = array![1.0_f64, 2.0, 0.5];
6759 let events = [1.0_f64, 1.0, 0.0];
6760 let eta_entry_vals = [-100.0_f64, 0.5, 0.8]; 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 } 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 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 #[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]; let age_exit = array![5.0_f64, 6.0, 7.0]; 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 #[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 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 #[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 #[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 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), ×, 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 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 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 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 #[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 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 #[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 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}