1use crate::custom_family::{
2 BlockWorkingSet, CustomFamily, FamilyEvaluation, ParameterBlockState,
3 projected_linear_constraint_stationarity_vector,
4};
5use crate::model_types::EstimationError;
6use gam_linalg::faer_ndarray::{fast_atv, fast_av, fast_xt_diag_x, fast_xt_diag_y};
7use gam_linalg::matrix::SymmetricMatrix;
8use gam_problem::{Coefficients, LinearPredictor};
9use gam_row_macros::row_atom;
10use gam_solve::pirls::{
11 ConstraintSet, LinearInequalityConstraints, WorkingModel as PirlsWorkingModel, WorkingState,
12 array1_l2_norm,
13};
14use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayView3, Axis};
15use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, constants, escalate_ridge};
16use serde::{Deserialize, Serialize};
17use std::collections::BTreeMap;
18use std::convert::Infallible;
19use std::ops::Range;
20use std::sync::LazyLock;
21use thiserror::Error;
22
23#[derive(Debug, Error)]
24pub enum SurvivalError {
25 #[error("input dimensions are inconsistent")]
26 DimensionMismatch,
27 #[error("inputs contain non-finite values")]
28 NonFiniteInput,
29 #[error("survival spec '{0}' is not supported by the one-hazard survival engine")]
30 UnsupportedSpec(&'static str),
31 #[error("crude risk integration setup is invalid")]
32 InvalidIntegrationSetup,
33 #[error("survival time grid must be finite, non-negative, and strictly increasing")]
34 InvalidTimeGrid,
35 #[error("cumulative hazard must be nondecreasing")]
36 NonMonotoneCumulativeHazard,
37 #[error("instantaneous hazard must stay strictly positive during integration")]
38 NonPositiveHazard,
39 #[error("{reason}")]
40 InvalidInput { reason: String },
41 #[error("{reason}")]
42 CauseSpecificDimensionMismatch { reason: String },
43 #[error("{reason}")]
44 NumericalFailure { reason: String },
45 #[error("{reason}")]
46 EventCodeInvalid { reason: String },
47 #[error("{reason}")]
48 EventDegenerate { reason: String },
49 #[error("cause-specific survival block {block}: {source}")]
50 CauseSpecificBlock {
51 block: usize,
52 #[source]
53 source: Box<SurvivalError>,
54 },
55}
56
57impl From<SurvivalError> for String {
58 fn from(err: SurvivalError) -> Self {
59 err.to_string()
60 }
61}
62
63impl From<crate::block_layout::block_count::BlockCountMismatch> for SurvivalError {
64 fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> SurvivalError {
65 SurvivalError::CauseSpecificDimensionMismatch {
66 reason: err.message(),
67 }
68 }
69}
70
71#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
72pub enum SurvivalSpec {
73 #[default]
74 Net,
75 Crude,
76}
77
78#[derive(Debug, Clone)]
79pub struct SurvivalEngineInputs<'a> {
80 pub age_entry: ArrayView1<'a, f64>,
81 pub age_exit: ArrayView1<'a, f64>,
82 pub event_target: ArrayView1<'a, u8>,
83 pub event_competing: ArrayView1<'a, u8>,
84 pub sampleweight: ArrayView1<'a, f64>,
85 pub x_entry: ArrayView2<'a, f64>,
86 pub x_exit: ArrayView2<'a, f64>,
87 pub x_derivative: ArrayView2<'a, f64>,
88 pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
92 pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
94}
95
96#[derive(Debug, Clone)]
97pub struct SurvivalTimeCovarInputs<'a> {
98 pub age_entry: ArrayView1<'a, f64>,
99 pub age_exit: ArrayView1<'a, f64>,
100 pub event_target: ArrayView1<'a, u8>,
101 pub event_competing: ArrayView1<'a, u8>,
102 pub sampleweight: ArrayView1<'a, f64>,
103 pub time_entry: ArrayView2<'a, f64>,
104 pub time_exit: ArrayView2<'a, f64>,
105 pub time_derivative: ArrayView2<'a, f64>,
106 pub covariates: ArrayView2<'a, f64>,
107 pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
111 pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
113}
114
115#[derive(Debug, Clone)]
116pub struct SurvivalBaselineOffsets<'a> {
117 pub eta_entry: ArrayView1<'a, f64>,
119 pub eta_exit: ArrayView1<'a, f64>,
121 pub derivative_exit: ArrayView1<'a, f64>,
129}
130
131#[derive(Debug, Clone)]
132pub struct PenaltyBlock {
133 pub matrix: Array2<f64>,
134 pub lambda: f64,
135 pub range: Range<usize>,
136 pub nullspace_dim: usize,
139}
140
141#[derive(Debug, Clone)]
142pub struct PenaltyBlocks {
143 pub blocks: Vec<PenaltyBlock>,
144}
145
146impl PenaltyBlocks {
147 pub fn new(blocks: Vec<PenaltyBlock>) -> Self {
148 Self { blocks }
149 }
150
151 pub fn gradient(&self, beta: &Array1<f64>) -> Array1<f64> {
152 let mut grad = Array1::zeros(beta.len());
153 for block in &self.blocks {
154 if block.lambda == 0.0 {
155 continue;
156 }
157 let b = beta.slice(ndarray::s![block.range.clone()]);
158 let g = block.matrix.dot(&b);
159 let mut dst = grad.slice_mut(ndarray::s![block.range.clone()]);
160 dst += &(block.lambda * g);
161 }
162 grad
163 }
164
165 pub fn hessian(&self, dim: usize) -> Array2<f64> {
166 let mut h = Array2::zeros((dim, dim));
167 self.addhessian_inplace(&mut h);
168 h
169 }
170
171 pub fn deviance(&self, beta: &Array1<f64>) -> f64 {
172 let mut value = 0.0;
173 for block in &self.blocks {
174 if block.lambda == 0.0 {
175 continue;
176 }
177 let b = beta.slice(ndarray::s![block.range.clone()]);
178 value += 0.5 * block.lambda * b.dot(&block.matrix.dot(&b));
179 }
180 value
181 }
182
183 pub fn addhessian_inplace(&self, h: &mut Array2<f64>) {
184 for block in &self.blocks {
185 if block.lambda == 0.0 {
186 continue;
187 }
188 let start = block.range.start;
189 let end = block.range.end;
190 h.slice_mut(ndarray::s![start..end, start..end])
191 .scaled_add(block.lambda, &block.matrix);
192 }
193 }
194}
195
196pub const ENTRY_AT_ORIGIN_THRESHOLD: f64 = 1e-8;
209
210const DERIVATIVE_FRACTION_TO_BOUNDARY: f64 = 0.995;
218
219#[derive(Debug, Clone)]
220pub struct CauseSpecificRoystonParmarBlock {
221 pub age_entry: Array1<f64>,
222 pub age_exit: Array1<f64>,
223 pub event_target: Array1<u8>,
224 pub sampleweight: Array1<f64>,
225 pub x_entry: Array2<f64>,
226 pub x_exit: Array2<f64>,
227 pub x_derivative: Array2<f64>,
228 pub offset_eta_entry: Array1<f64>,
229 pub offset_eta_exit: Array1<f64>,
230 pub offset_derivative_exit: Array1<f64>,
231 pub derivative_floor: f64,
232 pub structural_time_columns: usize,
249}
250
251#[derive(Debug, Clone)]
257pub struct CauseSpecificRoystonParmarFamily {
258 blocks: Vec<CauseSpecificRoystonParmarBlock>,
259}
260
261impl CauseSpecificRoystonParmarFamily {
262 pub fn new(blocks: Vec<CauseSpecificRoystonParmarBlock>) -> Result<Self, String> {
263 if blocks.is_empty() {
264 return Err(SurvivalError::InvalidInput {
265 reason: "cause-specific survival family requires at least one endpoint".to_string(),
266 }
267 .into());
268 }
269 for (idx, block) in blocks.iter().enumerate() {
270 validate_cause_specific_block(block).map_err(|err| {
271 SurvivalError::CauseSpecificBlock {
272 block: idx + 1,
273 source: Box::new(err),
274 }
275 .to_string()
276 })?;
277 }
278 Ok(Self { blocks })
279 }
280
281 pub fn cause_count(&self) -> usize {
282 self.blocks.len()
283 }
284}
285
286fn validate_cause_specific_block(
287 block: &CauseSpecificRoystonParmarBlock,
288) -> Result<(), SurvivalError> {
289 let n = block.event_target.len();
290 let p = block.x_exit.ncols();
291 if n == 0 || p == 0 {
292 bail_invalid_surv!("empty event vector or coefficient block");
293 }
294 if block.age_entry.len() != n
295 || block.age_exit.len() != n
296 || block.sampleweight.len() != n
297 || block.x_entry.nrows() != n
298 || block.x_exit.nrows() != n
299 || block.x_derivative.nrows() != n
300 || block.x_entry.ncols() != p
301 || block.x_derivative.ncols() != p
302 || block.offset_eta_entry.len() != n
303 || block.offset_eta_exit.len() != n
304 || block.offset_derivative_exit.len() != n
305 {
306 return Err(SurvivalError::CauseSpecificDimensionMismatch {
307 reason: "dimension mismatch".to_string(),
308 });
309 }
310 if let Some(&label) = block.event_target.iter().find(|&&v| v > 1) {
316 return Err(SurvivalError::EventCodeInvalid {
317 reason: format!(
318 "cause-specific block event_target must be the binary cause indicator {{0, 1}}, got multi-cause label {label}; project raw codes per cause via cause_specific_event_indicator"
319 ),
320 });
321 }
322 if block.age_entry.iter().any(|v| !v.is_finite())
323 || block.age_exit.iter().any(|v| !v.is_finite())
324 || block
325 .sampleweight
326 .iter()
327 .any(|v| !v.is_finite() || *v < 0.0)
328 || block.x_entry.iter().any(|v| !v.is_finite())
329 || block.x_exit.iter().any(|v| !v.is_finite())
330 || block.x_derivative.iter().any(|v| !v.is_finite())
331 || block.offset_eta_entry.iter().any(|v| !v.is_finite())
332 || block.offset_eta_exit.iter().any(|v| !v.is_finite())
333 || block.offset_derivative_exit.iter().any(|v| !v.is_finite())
334 || !block.derivative_floor.is_finite()
335 || block.derivative_floor < 0.0
336 {
337 bail_invalid_surv!("non-finite input");
338 }
339 Ok(())
340}
341
342row_atom! {
343 fn cause_specific_row [generic, order2, third, fourth](
344 eta_exit,
345 eta_entry,
346 derivative;
347 weight,
348 entry_active,
349 event
350 ) {
351 weight
352 * (exp(eta_exit)
353 - entry_active * exp(eta_entry)
354 - event * (eta_exit + ln(derivative)))
355 }
356}
357
358pub struct CauseSpecificSurvivalAloRowInput {
361 pub eta_exit: f64,
362 pub eta_entry: f64,
363 pub derivative_exit: f64,
364 pub prior_weight: f64,
365 pub entry_active: bool,
366 pub event: bool,
367}
368
369#[derive(Clone, Debug, PartialEq)]
372pub struct CauseSpecificSurvivalAloRowGeometry {
373 pub negative_log_likelihood: f64,
374 pub nll_score: [f64; 3],
375 pub observed_hessian: [[f64; 3]; 3],
376}
377
378pub fn cause_specific_survival_alo_row_geometry(
384 input: CauseSpecificSurvivalAloRowInput,
385) -> Result<CauseSpecificSurvivalAloRowGeometry, String> {
386 if !input.prior_weight.is_finite() || input.prior_weight < 0.0 {
387 return Err(format!(
388 "cause-specific saved ALO prior weight must be finite and non-negative, got {}",
389 input.prior_weight
390 ));
391 }
392 if input.prior_weight == 0.0 {
398 return Ok(CauseSpecificSurvivalAloRowGeometry {
399 negative_log_likelihood: 0.0,
400 nll_score: [0.0; 3],
401 observed_hessian: [[0.0; 3]; 3],
402 });
403 }
404 if !input.eta_exit.is_finite() {
405 return Err(format!(
406 "cause-specific saved ALO exit index must be finite, got {}",
407 input.eta_exit
408 ));
409 }
410 let eta_entry = if input.entry_active {
411 if !input.eta_entry.is_finite() {
412 return Err(format!(
413 "cause-specific saved ALO active entry index must be finite, got {}",
414 input.eta_entry
415 ));
416 }
417 input.eta_entry
418 } else {
419 0.0
420 };
421 let derivative_exit = if input.event {
422 if !input.derivative_exit.is_finite() || input.derivative_exit <= 0.0 {
423 return Err(format!(
424 "cause-specific saved ALO event derivative must be positive and finite, got {}",
425 input.derivative_exit
426 ));
427 }
428 input.derivative_exit
429 } else {
430 1.0
431 };
432 let atom = cause_specific_row_order2(
433 input.eta_exit,
434 eta_entry,
435 derivative_exit,
436 input.prior_weight,
437 f64::from(input.entry_active),
438 f64::from(input.event),
439 );
440 let gradient = atom.gradient();
441 let observed_hessian =
442 std::array::from_fn(|row| std::array::from_fn(|column| atom.hessian_at(row, column)));
443 if !atom.value().is_finite()
444 || gradient.iter().any(|value| !value.is_finite())
445 || observed_hessian
446 .iter()
447 .flatten()
448 .any(|value| !value.is_finite())
449 {
450 return Err(format!(
451 "cause-specific saved ALO row geometry is non-finite: nll={}, score={gradient:?}, hessian={observed_hessian:?}",
452 atom.value(),
453 ));
454 }
455 Ok(CauseSpecificSurvivalAloRowGeometry {
456 negative_log_likelihood: atom.value(),
457 nll_score: gradient,
458 observed_hessian,
459 })
460}
461
462#[derive(Clone, Copy)]
463struct CauseSpecificAtomInput {
464 primary: [f64; 3],
465 weight: f64,
466 entry_active: f64,
467 event: f64,
468}
469
470pub struct CauseSpecificRowProgram {
478 primary: [f64; 3],
479 weight: f64,
480 entry_active: f64,
481 event: f64,
482}
483
484impl CauseSpecificRowProgram {
485 pub fn new(primary: [f64; 3], weight: f64, entry_active: bool, event: bool) -> Self {
487 Self {
488 primary,
489 weight,
490 entry_active: f64::from(entry_active),
491 event: f64::from(event),
492 }
493 }
494
495 fn require_row(row: usize) -> Result<(), String> {
496 if row != 0 {
497 return Err(format!(
498 "CauseSpecificRowProgram holds exactly one row; got row {row}"
499 ));
500 }
501 Ok(())
502 }
503}
504
505impl gam_math::jet_tower::RowProgram<3> for CauseSpecificRowProgram {
506 fn n_rows(&self) -> usize {
507 1
508 }
509
510 fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
511 Self::require_row(row)?;
512 Ok(self.primary)
513 }
514
515 fn eval<S: gam_math::jet_scalar::JetScalar<3>>(
516 &self,
517 row: usize,
518 p: &[S; 3],
519 ) -> Result<S, String> {
520 Self::require_row(row)?;
521 Ok(cause_specific_row(
522 &p[0],
523 &p[1],
524 &p[2],
525 self.weight,
526 self.entry_active,
527 self.event,
528 ))
529 }
530}
531
532fn cause_specific_atom_input(
537 block: &CauseSpecificRoystonParmarBlock,
538 row: usize,
539 eta_entry: f64,
540 eta_exit: f64,
541 derivative: f64,
542) -> Result<Option<CauseSpecificAtomInput>, SurvivalError> {
543 let weight = block.sampleweight[row];
544 if weight <= 0.0 {
545 return Ok(None);
546 }
547 if block.age_exit[row] < block.age_entry[row] {
548 bail_invalid_surv!("age_exit < age_entry at row {row}");
549 }
550 let entry_active = block.age_entry[row] > ENTRY_AT_ORIGIN_THRESHOLD;
551 let event = block.event_target[row] > 0;
552 let eta_entry = if entry_active { eta_entry } else { 0.0 };
553 let derivative = if event {
554 if !(derivative.is_finite() && derivative > 0.0) {
555 return Err(SurvivalError::NumericalFailure {
556 reason: format!(
557 "cause-specific survival derivative must be positive at row {row}, got {derivative}"
558 ),
559 });
560 }
561 derivative
562 } else {
563 1.0
564 };
565 let h_exit = eta_exit.exp();
566 let h_entry = eta_entry.exp();
567 if !(h_exit.is_finite() && h_entry.is_finite()) {
568 return Err(SurvivalError::NumericalFailure {
569 reason: format!("non-finite cumulative hazard at row {row}"),
570 });
571 }
572 Ok(Some(CauseSpecificAtomInput {
573 primary: [eta_exit, eta_entry, derivative],
574 weight,
575 entry_active: f64::from(entry_active),
576 event: f64::from(event),
577 }))
578}
579
580const CAUSE_SPECIFIC_PRIMARY_PAIRS: [(usize, usize); 6] =
581 [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)];
582
583fn cause_specific_pullback_hessian(
589 block: &CauseSpecificRoystonParmarBlock,
590 weights: &[Array1<f64>; 6],
591) -> Array2<f64> {
592 let designs = [&block.x_exit, &block.x_entry, &block.x_derivative];
593 let p = block.x_exit.ncols();
594 let mut hessian = Array2::<f64>::zeros((p, p));
595 for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
596 let channel = &weights[slot];
597 if channel.iter().all(|&value| value == 0.0) {
598 continue;
599 }
600 if left == right {
601 hessian += &fast_xt_diag_x(designs[left], channel);
602 } else {
603 let cross = fast_xt_diag_y(designs[left], channel, designs[right]);
604 hessian += ✗
605 hessian += &cross.t();
606 }
607 }
608 hessian
609}
610
611fn evaluate_cause_specific_block(
612 block: &CauseSpecificRoystonParmarBlock,
613 beta: &Array1<f64>,
614) -> Result<(f64, Array1<f64>, Array2<f64>), SurvivalError> {
615 let n = block.event_target.len();
616 let p = block.x_exit.ncols();
617 if beta.len() != p {
618 return Err(SurvivalError::CauseSpecificDimensionMismatch {
619 reason: format!("beta length mismatch: got {}, expected {p}", beta.len()),
620 });
621 }
622 let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
623 let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
624 let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
625 let mut log_likelihood = 0.0;
626 let mut gradient_weights: [Array1<f64>; 3] = std::array::from_fn(|_| Array1::<f64>::zeros(n));
627 let mut hessian_weights: [Array1<f64>; 6] = std::array::from_fn(|_| Array1::<f64>::zeros(n));
628
629 for i in 0..n {
630 let Some(input) =
631 cause_specific_atom_input(block, i, eta_entry[i], eta_exit[i], derivative[i])?
632 else {
633 continue;
634 };
635 let atom = cause_specific_row_order2(
636 input.primary[0],
637 input.primary[1],
638 input.primary[2],
639 input.weight,
640 input.entry_active,
641 input.event,
642 );
643 log_likelihood -= atom.value();
644 let gradient = atom.gradient();
645 for axis in 0..3 {
646 gradient_weights[axis][i] = -gradient[axis];
647 }
648 for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
649 hessian_weights[slot][i] = atom.hessian_at(left, right);
650 }
651 }
652
653 let designs = [&block.x_exit, &block.x_entry, &block.x_derivative];
654 let mut gradient = Array1::<f64>::zeros(p);
655 for axis in 0..3 {
656 gradient += &fast_atv(designs[axis], &gradient_weights[axis]);
657 }
658 let hessian = cause_specific_pullback_hessian(block, &hessian_weights);
659 Ok((log_likelihood, gradient, hessian))
660}
661
662impl CustomFamily for CauseSpecificRoystonParmarFamily {
663 fn joint_jeffreys_term_required(&self) -> bool {
667 true
668 }
669
670 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
671 crate::block_layout::block_count::validate_block_count::<SurvivalError>(
672 "cause-specific survival",
673 self.blocks.len(),
674 block_states.len(),
675 )?;
676 let mut log_likelihood = 0.0;
677 let mut blockworking_sets = Vec::with_capacity(self.blocks.len());
678 for (block, state) in self.blocks.iter().zip(block_states.iter()) {
679 let (ll, gradient, hessian) = evaluate_cause_specific_block(block, &state.beta)?;
680 log_likelihood += ll;
681 blockworking_sets.push(BlockWorkingSet::ExactNewton {
682 gradient,
683 hessian: SymmetricMatrix::Dense(hessian),
684 });
685 }
686 Ok(FamilyEvaluation {
687 log_likelihood,
688 blockworking_sets,
689 })
690 }
691
692 fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
693 crate::block_layout::block_count::validate_block_count::<SurvivalError>(
694 "cause-specific survival",
695 self.blocks.len(),
696 block_states.len(),
697 )?;
698 let mut log_likelihood = 0.0;
699 for (block, state) in self.blocks.iter().zip(block_states.iter()) {
700 let (ll, _, _) = evaluate_cause_specific_block(block, &state.beta)?;
701 log_likelihood += ll;
702 }
703 Ok(log_likelihood)
704 }
705
706 fn likelihood_blocks_uncoupled(&self) -> bool {
707 true
708 }
709
710 fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
711 true
712 }
713
714 fn output_channel_assignment(
715 &self,
716 specs: &[crate::custom_family::ParameterBlockSpec],
717 ) -> Option<Vec<usize>> {
718 if specs.len() != self.blocks.len() {
719 return Some((0..self.blocks.len()).collect());
720 }
721 Some((0..specs.len()).collect())
722 }
723
724 fn coefficient_hessian_cost(&self, specs: &[crate::custom_family::ParameterBlockSpec]) -> u64 {
725 crate::custom_family::default_coefficient_hessian_cost(specs)
726 }
727
728 fn block_linear_constraints(
729 &self,
730 _: &[ParameterBlockState],
731 block_idx: usize,
732 spec: &crate::custom_family::ParameterBlockSpec,
733 ) -> Result<Option<ConstraintSet>, String> {
734 let block = self.blocks.get(block_idx).ok_or_else(|| {
735 SurvivalError::CauseSpecificDimensionMismatch {
736 reason: format!(
737 "cause-specific survival expected block index < {}, got {block_idx}",
738 self.blocks.len()
739 ),
740 }
741 .to_string()
742 })?;
743 if block.x_derivative.ncols() != spec.design.ncols() {
744 return Err(SurvivalError::CauseSpecificDimensionMismatch {
745 reason: format!(
746 "cause-specific survival derivative design has {} columns but block '{}' has {}",
747 block.x_derivative.ncols(),
748 spec.name,
749 spec.design.ncols()
750 ),
751 }
752 .into());
753 }
754 let rhs = block
755 .offset_derivative_exit
756 .mapv(|offset| block.derivative_floor - offset);
757 let p = block.x_derivative.ncols();
758 let n_rows = block.x_derivative.nrows();
759 let structural_cols = block.structural_time_columns.min(p);
770 if structural_cols == 0 {
771 return Ok(Some(ConstraintSet::Dense(LinearInequalityConstraints {
772 a: block.x_derivative.clone(),
773 b: rhs,
774 })));
775 }
776 let mut a = Array2::<f64>::zeros((n_rows + structural_cols, p));
777 a.slice_mut(ndarray::s![..n_rows, ..])
778 .assign(&block.x_derivative);
779 for j in 0..structural_cols {
780 a[[n_rows + j, j]] = 1.0;
781 }
782 let mut b = Array1::<f64>::zeros(n_rows + structural_cols);
783 b.slice_mut(ndarray::s![..n_rows]).assign(&rhs);
784 Ok(Some(ConstraintSet::Dense(LinearInequalityConstraints {
785 a,
786 b,
787 })))
788 }
789
790 fn max_feasible_step_size(
791 &self,
792 block_states: &[ParameterBlockState],
793 block_idx: usize,
794 delta: &Array1<f64>,
795 ) -> Result<Option<f64>, String> {
796 let block = self.blocks.get(block_idx).ok_or_else(|| {
797 SurvivalError::CauseSpecificDimensionMismatch {
798 reason: format!(
799 "cause-specific survival expected block index < {}, got {block_idx}",
800 self.blocks.len()
801 ),
802 }
803 .to_string()
804 })?;
805 let state = block_states.get(block_idx).ok_or_else(|| {
806 SurvivalError::CauseSpecificDimensionMismatch {
807 reason: format!(
808 "cause-specific survival expected {} block states, got {}",
809 self.blocks.len(),
810 block_states.len()
811 ),
812 }
813 .to_string()
814 })?;
815 if delta.len() != state.beta.len() || block.x_derivative.ncols() != delta.len() {
816 return Err(SurvivalError::CauseSpecificDimensionMismatch {
817 reason: "cause-specific survival feasible-step dimension mismatch".to_string(),
818 }
819 .into());
820 }
821 let derivative = fast_av(&block.x_derivative, &state.beta) + &block.offset_derivative_exit;
822 let derivative_delta = fast_av(&block.x_derivative, delta);
823 let mut alpha_max = 1.0_f64;
824 for i in 0..derivative.len() {
825 if block.sampleweight[i] <= 0.0 {
826 continue;
827 }
828 let current = derivative[i] - block.derivative_floor;
829 let slope = derivative_delta[i];
830 if slope < 0.0 {
831 if current <= 0.0 {
832 return Ok(Some(0.0));
833 }
834 alpha_max = alpha_max.min(DERIVATIVE_FRACTION_TO_BOUNDARY * current / -slope);
835 }
836 }
837 Ok(Some(alpha_max.clamp(0.0, 1.0)))
838 }
839
840 fn exact_newton_hessian_directional_derivative(
841 &self,
842 block_states: &[ParameterBlockState],
843 block_idx: usize,
844 d_beta: &Array1<f64>,
845 ) -> Result<Option<Array2<f64>>, String> {
846 let block = self.blocks.get(block_idx).ok_or_else(|| {
847 SurvivalError::CauseSpecificDimensionMismatch {
848 reason: format!(
849 "cause-specific survival expected block index < {}, got {block_idx}",
850 self.blocks.len()
851 ),
852 }
853 .to_string()
854 })?;
855 let state = block_states.get(block_idx).ok_or_else(|| {
856 SurvivalError::CauseSpecificDimensionMismatch {
857 reason: format!(
858 "cause-specific survival expected {} block states, got {}",
859 self.blocks.len(),
860 block_states.len()
861 ),
862 }
863 .to_string()
864 })?;
865 Ok(Some(cause_specific_hessian_directional_derivative(
866 block,
867 &state.beta,
868 d_beta,
869 )?))
870 }
871
872 fn exact_newton_hessian_second_directional_derivative(
873 &self,
874 block_states: &[ParameterBlockState],
875 block_idx: usize,
876 d_beta_u: &Array1<f64>,
877 d_beta_v: &Array1<f64>,
878 ) -> Result<Option<Array2<f64>>, String> {
879 let block = self.blocks.get(block_idx).ok_or_else(|| {
880 SurvivalError::CauseSpecificDimensionMismatch {
881 reason: format!(
882 "cause-specific survival expected block index < {}, got {block_idx}",
883 self.blocks.len()
884 ),
885 }
886 .to_string()
887 })?;
888 let state = block_states.get(block_idx).ok_or_else(|| {
889 SurvivalError::CauseSpecificDimensionMismatch {
890 reason: format!(
891 "cause-specific survival expected {} block states, got {}",
892 self.blocks.len(),
893 block_states.len()
894 ),
895 }
896 .to_string()
897 })?;
898 Ok(Some(cause_specific_hessian_second_directional_derivative(
899 block,
900 &state.beta,
901 d_beta_u,
902 d_beta_v,
903 )?))
904 }
905}
906
907fn cause_specific_hessian_directional_derivative(
910 block: &CauseSpecificRoystonParmarBlock,
911 beta: &Array1<f64>,
912 d_beta: &Array1<f64>,
913) -> Result<Array2<f64>, SurvivalError> {
914 let p = block.x_exit.ncols();
915 if beta.len() != p || d_beta.len() != p {
916 return Err(SurvivalError::CauseSpecificDimensionMismatch {
917 reason: "cause-specific survival Hessian derivative dimension mismatch".to_string(),
918 });
919 }
920 let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
921 let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
922 let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
923 let d_eta_entry = fast_av(&block.x_entry, d_beta);
924 let d_eta_exit = fast_av(&block.x_exit, d_beta);
925 let d_derivative = fast_av(&block.x_derivative, d_beta);
926 let n = block.event_target.len();
927 let mut weights: [Array1<f64>; 6] = std::array::from_fn(|_| Array1::zeros(n));
928
929 for i in 0..n {
930 let Some(input) =
931 cause_specific_atom_input(block, i, eta_entry[i], eta_exit[i], derivative[i])?
932 else {
933 continue;
934 };
935 let direction = [
936 d_eta_exit[i],
937 d_eta_entry[i] * input.entry_active,
938 d_derivative[i] * input.event,
939 ];
940 let matrix = cause_specific_row_third_contracted(
941 input.primary[0],
942 input.primary[1],
943 input.primary[2],
944 input.weight,
945 input.entry_active,
946 input.event,
947 &direction,
948 );
949 for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
950 weights[slot][i] = matrix[left][right];
951 }
952 }
953 Ok(cause_specific_pullback_hessian(block, &weights))
954}
955
956fn cause_specific_hessian_second_directional_derivative(
959 block: &CauseSpecificRoystonParmarBlock,
960 beta: &Array1<f64>,
961 d_beta_u: &Array1<f64>,
962 d_beta_v: &Array1<f64>,
963) -> Result<Array2<f64>, SurvivalError> {
964 let p = block.x_exit.ncols();
965 if beta.len() != p || d_beta_u.len() != p || d_beta_v.len() != p {
966 return Err(SurvivalError::CauseSpecificDimensionMismatch {
967 reason: "cause-specific survival second Hessian derivative dimension mismatch"
968 .to_string(),
969 });
970 }
971 let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
972 let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
973 let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
974 let u_eta_entry = fast_av(&block.x_entry, d_beta_u);
975 let u_eta_exit = fast_av(&block.x_exit, d_beta_u);
976 let u_derivative = fast_av(&block.x_derivative, d_beta_u);
977 let v_eta_entry = fast_av(&block.x_entry, d_beta_v);
978 let v_eta_exit = fast_av(&block.x_exit, d_beta_v);
979 let v_derivative = fast_av(&block.x_derivative, d_beta_v);
980 let n = block.event_target.len();
981 let mut weights: [Array1<f64>; 6] = std::array::from_fn(|_| Array1::zeros(n));
982
983 for i in 0..n {
984 let Some(input) =
985 cause_specific_atom_input(block, i, eta_entry[i], eta_exit[i], derivative[i])?
986 else {
987 continue;
988 };
989 let direction_u = [
990 u_eta_exit[i],
991 u_eta_entry[i] * input.entry_active,
992 u_derivative[i] * input.event,
993 ];
994 let direction_v = [
995 v_eta_exit[i],
996 v_eta_entry[i] * input.entry_active,
997 v_derivative[i] * input.event,
998 ];
999 let matrix = cause_specific_row_fourth_contracted(
1000 input.primary[0],
1001 input.primary[1],
1002 input.primary[2],
1003 input.weight,
1004 input.entry_active,
1005 input.event,
1006 &direction_u,
1007 &direction_v,
1008 );
1009 for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
1010 weights[slot][i] = matrix[left][right];
1011 }
1012 }
1013 Ok(cause_specific_pullback_hessian(block, &weights))
1014}
1015
1016pub fn survival_event_code_from_value(value: f64, row_index: usize) -> Result<u8, String> {
1017 const INTEGER_TOL: f64 = 1e-8;
1018 const MAX_AUTO_CAUSES: u8 = 32;
1019 if !value.is_finite() {
1020 return Err(SurvivalError::EventCodeInvalid {
1021 reason: format!(
1022 "survival event value at row {} is non-finite",
1023 row_index + 1
1024 ),
1025 }
1026 .into());
1027 }
1028 if value < 0.0 {
1029 return Err(SurvivalError::EventCodeInvalid {
1030 reason: format!(
1031 "survival event value at row {} is negative: {value}",
1032 row_index + 1
1033 ),
1034 }
1035 .into());
1036 }
1037 let rounded = value.round();
1038 if (value - rounded).abs() > INTEGER_TOL {
1039 return Err(SurvivalError::EventCodeInvalid {
1040 reason: format!(
1041 "survival event value at row {} must be an integer code with 0=censored, got {value}",
1042 row_index + 1
1043 ),
1044 }
1045 .into());
1046 }
1047 if rounded > f64::from(MAX_AUTO_CAUSES) {
1048 return Err(SurvivalError::EventCodeInvalid {
1049 reason: format!(
1050 "survival event value at row {} has code {rounded}; automatic competing-risks detection supports codes 0..={MAX_AUTO_CAUSES}",
1051 row_index + 1
1052 ),
1053 }
1054 .into());
1055 }
1056 Ok(rounded as u8)
1057}
1058
1059pub fn cause_count_from_event_codes(
1060 event_codes: ArrayView1<'_, u8>,
1061) -> Result<usize, SurvivalError> {
1062 let max_code = event_codes.iter().copied().max().map_or(0, usize::from);
1063 if max_code == 0 {
1064 return Ok(1);
1065 }
1066
1067 let mut present = vec![false; max_code + 1];
1068 for code in event_codes.iter().copied() {
1069 present[usize::from(code)] = true;
1070 }
1071 if (1..=max_code).any(|code| !present[code]) {
1072 let actual = present
1073 .iter()
1074 .enumerate()
1075 .skip(1)
1076 .filter_map(|(code, &seen)| seen.then_some(code.to_string()))
1077 .collect::<Vec<_>>()
1078 .join(", ");
1079 return Err(SurvivalError::EventCodeInvalid {
1080 reason: format!(
1081 "survival competing-risks event codes must use contiguous positive codes; observed nonzero codes are {{{actual}}}. Remap event codes contiguously (for example, {{0,1,3}} -> {{0,1,2}}), otherwise a phantom cause is fit with no events and pollutes CIF assembly."
1082 ),
1083 });
1084 }
1085
1086 Ok(max_code)
1087}
1088
1089pub fn pooled_any_event_indicator(event_codes: ArrayView1<'_, u8>) -> Array1<u8> {
1102 event_codes.mapv(|label| u8::from(label > 0))
1103}
1104
1105pub fn cause_specific_event_indicator(event_codes: ArrayView1<'_, u8>, cause: usize) -> Array1<u8> {
1115 let cause_code = cause as u8;
1116 event_codes.mapv(|observed| u8::from(observed == cause_code))
1117}
1118
1119fn compress_positive_collinear_constraints(
1120 a: &Array2<f64>,
1121 b: &Array1<f64>,
1122) -> LinearInequalityConstraints {
1123 const SCALE_TOL: f64 = 1e-14;
1124 const KEY_TOL: f64 = 1e-8;
1125
1126 let mut grouped: BTreeMap<Vec<i64>, (Vec<f64>, f64)> = BTreeMap::new();
1127 let mut fallbackrows: Vec<(Vec<f64>, f64)> = Vec::new();
1128
1129 for i in 0..a.nrows() {
1130 let row = a.row(i);
1131 let scale = row.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
1132 if !scale.is_finite() || scale <= SCALE_TOL {
1133 if b[i] > 0.0 {
1134 fallbackrows.push((row.to_vec(), b[i]));
1135 }
1136 continue;
1137 }
1138
1139 let normalizedrow: Vec<f64> = row
1140 .iter()
1141 .map(|&v| {
1142 let scaled = v / scale;
1143 if scaled.abs() <= KEY_TOL { 0.0 } else { scaled }
1144 })
1145 .collect();
1146 let normalized_rhs = b[i] / scale;
1147 let key: Vec<i64> = normalizedrow
1148 .iter()
1149 .map(|&v| (v / KEY_TOL).round() as i64)
1150 .collect();
1151
1152 match grouped.get_mut(&key) {
1153 Some((_, rhs_max)) => {
1154 if normalized_rhs > *rhs_max {
1155 *rhs_max = normalized_rhs;
1156 }
1157 }
1158 None => {
1159 grouped.insert(key, (normalizedrow, normalized_rhs));
1160 }
1161 }
1162 }
1163
1164 let nrows = grouped.len() + fallbackrows.len();
1165 let n_cols = a.ncols();
1166 let mut a_out = Array2::<f64>::zeros((nrows, n_cols));
1167 let mut b_out = Array1::<f64>::zeros(nrows);
1168
1169 let mut outrow = 0usize;
1170 for (_, (row, rhs)) in grouped {
1171 for (j, value) in row.into_iter().enumerate() {
1172 a_out[[outrow, j]] = value;
1173 }
1174 b_out[outrow] = rhs;
1175 outrow += 1;
1176 }
1177 for (row, rhs) in fallbackrows {
1178 for (j, value) in row.into_iter().enumerate() {
1179 a_out[[outrow, j]] = value;
1180 }
1181 b_out[outrow] = rhs;
1182 outrow += 1;
1183 }
1184
1185 LinearInequalityConstraints { a: a_out, b: b_out }
1186}
1187
1188#[derive(Debug, Clone, Copy, Default)]
1189pub struct SurvivalMonotonicityPenalty {
1190 pub tolerance: f64,
1191}
1192
1193#[derive(Debug, Clone)]
1194enum SurvivalDesign {
1195 Flat {
1196 x_entry: Array2<f64>,
1197 x_exit: Array2<f64>,
1198 x_derivative: Array2<f64>,
1199 },
1200 TimeCovariateShared {
1201 time_entry: Array2<f64>,
1202 time_exit: Array2<f64>,
1203 time_derivative: Array2<f64>,
1204 covariates: Array2<f64>,
1205 },
1206}
1207
1208impl SurvivalDesign {
1209 fn p_total(&self) -> usize {
1210 match self {
1211 Self::Flat { x_exit, .. } => x_exit.ncols(),
1212 Self::TimeCovariateShared {
1213 time_exit,
1214 covariates,
1215 ..
1216 } => time_exit.ncols() + covariates.ncols(),
1217 }
1218 }
1219
1220 fn design_dot(&self, time_mat: &Array2<f64>, beta: &Array1<f64>) -> Array1<f64> {
1221 match self {
1222 Self::Flat { .. } => time_mat.dot(beta),
1223 Self::TimeCovariateShared { covariates, .. } => {
1224 let p_time = time_mat.ncols();
1225 let mut out = time_mat.dot(&beta.slice(ndarray::s![..p_time]));
1226 if covariates.ncols() > 0 {
1227 out += &covariates.dot(&beta.slice(ndarray::s![p_time..]));
1228 }
1229 out
1230 }
1231 }
1232 }
1233
1234 fn fill_row(&self, time_mat: &Array2<f64>, i: usize, out: &mut [f64]) {
1235 match self {
1236 Self::Flat { .. } => {
1237 for (dst, &src) in out.iter_mut().zip(time_mat.row(i).iter()) {
1238 *dst = src;
1239 }
1240 }
1241 Self::TimeCovariateShared { covariates, .. } => {
1242 let p_time = time_mat.ncols();
1243 for j in 0..p_time {
1244 out[j] = time_mat[[i, j]];
1245 }
1246 for j in 0..covariates.ncols() {
1247 out[p_time + j] = covariates[[i, j]];
1248 }
1249 }
1250 }
1251 }
1252}
1253
1254#[derive(Debug, Clone)]
1256struct SurvivalWorkspace {
1257 w_event: Array1<f64>,
1258 w_event_inv_deriv: Array1<f64>,
1259 w_event_outer: Array1<f64>,
1260 w_hess_exit: Array1<f64>,
1261 w_hess_entry: Array1<f64>,
1262}
1263
1264impl SurvivalWorkspace {
1265 fn new(n: usize) -> Self {
1266 Self {
1267 w_event: Array1::zeros(n),
1268 w_event_inv_deriv: Array1::zeros(n),
1269 w_event_outer: Array1::zeros(n),
1270 w_hess_exit: Array1::zeros(n),
1271 w_hess_entry: Array1::zeros(n),
1272 }
1273 }
1274
1275 fn reset(&mut self, n: usize) {
1276 if self.w_event.len() != n {
1277 *self = Self::new(n);
1278 } else {
1279 self.w_event.fill(0.0);
1280 self.w_event_inv_deriv.fill(0.0);
1281 self.w_event_outer.fill(0.0);
1282 self.w_hess_exit.fill(0.0);
1283 self.w_hess_entry.fill(0.0);
1284 }
1285 }
1286}
1287
1288#[derive(Clone, Debug)]
1301pub struct OffsetChannelResiduals {
1302 pub exit: Array1<f64>,
1304 pub entry: Array1<f64>,
1306 pub derivative: Array1<f64>,
1308 pub right: Array1<f64>,
1312}
1313
1314#[derive(Clone, Debug)]
1317pub struct OffsetChannelCurvatures {
1318 pub rows: Vec<[[f64; 3]; 3]>,
1319}
1320
1321#[derive(Debug)]
1322pub struct WorkingModelSurvival {
1323 age_entry: Array1<f64>,
1324 age_exit: Array1<f64>,
1325 entry_at_origin: Array1<bool>,
1326 event_target: Array1<u8>,
1327 sampleweight: Array1<f64>,
1328 design: SurvivalDesign,
1329 offset_eta_entry: Array1<f64>,
1330 offset_eta_exit: Array1<f64>,
1331 offset_derivative_exit: Array1<f64>,
1332 penalties: PenaltyBlocks,
1333 monotonicity: SurvivalMonotonicityPenalty,
1334 structurally_monotonic: bool,
1335 structural_time_columns: usize,
1336 monotonicity_constraint_rows: Option<Array2<f64>>,
1337 monotonicity_constraint_offsets: Option<Array1<f64>>,
1338 workspace: std::sync::Mutex<SurvivalWorkspace>,
1339}
1340
1341impl Clone for WorkingModelSurvival {
1342 fn clone(&self) -> Self {
1343 let workspace = self.workspace.lock().unwrap().clone();
1344 Self {
1345 age_entry: self.age_entry.clone(),
1346 age_exit: self.age_exit.clone(),
1347 entry_at_origin: self.entry_at_origin.clone(),
1348 event_target: self.event_target.clone(),
1349 sampleweight: self.sampleweight.clone(),
1350 design: self.design.clone(),
1351 offset_eta_entry: self.offset_eta_entry.clone(),
1352 offset_eta_exit: self.offset_eta_exit.clone(),
1353 offset_derivative_exit: self.offset_derivative_exit.clone(),
1354 penalties: self.penalties.clone(),
1355 monotonicity: self.monotonicity,
1356 structurally_monotonic: self.structurally_monotonic,
1357 structural_time_columns: self.structural_time_columns,
1358 monotonicity_constraint_rows: self.monotonicity_constraint_rows.clone(),
1359 monotonicity_constraint_offsets: self.monotonicity_constraint_offsets.clone(),
1360 workspace: std::sync::Mutex::new(workspace),
1361 }
1362 }
1363}
1364
1365impl WorkingModelSurvival {
1366 const LOG_F64_MAX: f64 = 709.782712893384;
1367
1368 #[inline]
1369 fn scaled_exp_component(log_scale: f64, base: f64) -> Result<f64, EstimationError> {
1370 if base == 0.0 {
1371 return Ok(0.0);
1372 }
1373 let log_abs = log_scale + base.abs().ln();
1374 if !log_abs.is_finite() {
1375 crate::bail_invalid_estim!("survival interval term produced non-finite log-magnitude");
1376 }
1377 if log_abs > Self::LOG_F64_MAX {
1378 crate::bail_invalid_estim!(
1379 "survival interval term exceeds f64 range (log-magnitude={log_abs:.3e})"
1380 );
1381 }
1382 Ok(base.signum() * log_abs.exp())
1383 }
1384
1385 fn coefficient_dim(&self) -> usize {
1386 self.design.p_total()
1387 }
1388
1389 fn nrows(&self) -> usize {
1390 self.sampleweight.len()
1391 }
1392
1393 fn entry_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1394 let time_mat = match &self.design {
1395 SurvivalDesign::Flat { x_entry, .. } => x_entry,
1396 SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
1397 };
1398 self.design.design_dot(time_mat, beta)
1399 }
1400
1401 fn exit_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1402 let time_mat = match &self.design {
1403 SurvivalDesign::Flat { x_exit, .. } => x_exit,
1404 SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
1405 };
1406 self.design.design_dot(time_mat, beta)
1407 }
1408
1409 fn derivative_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1410 match &self.design {
1411 SurvivalDesign::Flat { x_derivative, .. } => x_derivative.dot(beta),
1412 SurvivalDesign::TimeCovariateShared {
1413 time_derivative, ..
1414 } => time_derivative.dot(&beta.slice(ndarray::s![..time_derivative.ncols()])),
1415 }
1416 }
1417
1418 fn fill_entry_row(&self, i: usize, out: &mut [f64]) {
1419 let time_mat = match &self.design {
1420 SurvivalDesign::Flat { x_entry, .. } => x_entry,
1421 SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
1422 };
1423 self.design.fill_row(time_mat, i, out);
1424 }
1425
1426 fn fill_exit_row(&self, i: usize, out: &mut [f64]) {
1427 let time_mat = match &self.design {
1428 SurvivalDesign::Flat { x_exit, .. } => x_exit,
1429 SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
1430 };
1431 self.design.fill_row(time_mat, i, out);
1432 }
1433
1434 fn fill_derivative_row(&self, i: usize, out: &mut [f64]) {
1435 match &self.design {
1436 SurvivalDesign::Flat { x_derivative, .. } => {
1437 for (dst, &src) in out.iter_mut().zip(x_derivative.row(i).iter()) {
1438 *dst = src;
1439 }
1440 }
1441 SurvivalDesign::TimeCovariateShared {
1442 time_derivative, ..
1443 } => {
1444 let p_time = time_derivative.ncols();
1445 for j in 0..p_time {
1446 out[j] = time_derivative[[i, j]];
1447 }
1448 for dst in out.iter_mut().skip(p_time) {
1449 *dst = 0.0;
1450 }
1451 }
1452 }
1453 }
1454
1455 fn derivative_xt_diag_x(&self, weights: &Array1<f64>) -> Array2<f64> {
1456 match &self.design {
1457 SurvivalDesign::Flat { x_derivative, .. } => fast_xt_diag_x(x_derivative, weights),
1458 SurvivalDesign::TimeCovariateShared {
1459 time_derivative,
1460 covariates,
1461 ..
1462 } => {
1463 let p_time = time_derivative.ncols();
1464 let p_cov = covariates.ncols();
1465 let mut out = Array2::<f64>::zeros((p_time + p_cov, p_time + p_cov));
1466 let time_block = fast_xt_diag_x(time_derivative, weights);
1467 out.slice_mut(ndarray::s![..p_time, ..p_time])
1468 .assign(&time_block);
1469 out
1470 }
1471 }
1472 }
1473
1474 fn interval_hessian_blas(&self, w_exit: &Array1<f64>, w_entry: &Array1<f64>) -> Array2<f64> {
1478 match &self.design {
1479 SurvivalDesign::Flat {
1480 x_entry, x_exit, ..
1481 } => {
1482 let mut h = fast_xt_diag_x(x_exit, w_exit);
1483 h -= &fast_xt_diag_x(x_entry, w_entry);
1484 h
1485 }
1486 SurvivalDesign::TimeCovariateShared {
1487 time_entry,
1488 time_exit,
1489 covariates,
1490 ..
1491 } => {
1492 let p_time = time_exit.ncols();
1493 let p_cov = covariates.ncols();
1494 let p = p_time + p_cov;
1495 let mut h = Array2::<f64>::zeros((p, p));
1496 let tt = {
1498 let mut block = fast_xt_diag_x(time_exit, w_exit);
1499 block -= &fast_xt_diag_x(time_entry, w_entry);
1500 block
1501 };
1502 h.slice_mut(ndarray::s![..p_time, ..p_time]).assign(&tt);
1503 if p_cov > 0 {
1504 let tc = {
1506 let mut block = fast_xt_diag_y(time_exit, w_exit, covariates);
1507 block -= &fast_xt_diag_y(time_entry, w_entry, covariates);
1508 block
1509 };
1510 h.slice_mut(ndarray::s![..p_time, p_time..]).assign(&tc);
1511 h.slice_mut(ndarray::s![p_time.., ..p_time]).assign(&tc.t());
1512 let w_diff = w_exit - w_entry;
1514 let cc = fast_xt_diag_x(covariates, &w_diff);
1515 h.slice_mut(ndarray::s![p_time.., p_time..]).assign(&cc);
1516 }
1517 h
1518 }
1519 }
1520 }
1521
1522 fn stabilized_structural_derivative(&self, deriv: f64) -> Option<(f64, f64)> {
1536 const STRUCTURAL_MONO_ROUNDOFF_TOL: f64 = 1e-7;
1537 const STRUCTURAL_DERIV_FLOOR: f64 = 1e-12;
1538 if !self.structurally_monotonic {
1539 return None;
1540 }
1541 if deriv >= STRUCTURAL_DERIV_FLOOR {
1542 return Some((deriv, 1.0));
1543 }
1544 if deriv >= -STRUCTURAL_MONO_ROUNDOFF_TOL {
1545 return Some((STRUCTURAL_DERIV_FLOOR, 0.0));
1546 }
1547 None
1548 }
1549
1550 fn validate_penalties(
1551 penalties: &PenaltyBlocks,
1552 coefficient_dim: usize,
1553 ) -> Result<(), SurvivalError> {
1554 for block in &penalties.blocks {
1555 if !block.lambda.is_finite() || block.lambda < 0.0 {
1556 return Err(SurvivalError::NonFiniteInput);
1557 }
1558 if block.range.start > block.range.end || block.range.end > coefficient_dim {
1559 return Err(SurvivalError::DimensionMismatch);
1560 }
1561 let block_dim = block.range.end - block.range.start;
1562 if block.matrix.nrows() != block_dim || block.matrix.ncols() != block_dim {
1563 return Err(SurvivalError::DimensionMismatch);
1564 }
1565 if block.matrix.iter().any(|v| !v.is_finite()) {
1566 return Err(SurvivalError::NonFiniteInput);
1567 }
1568 }
1569 Ok(())
1570 }
1571
1572 fn derivative_guard(&self) -> f64 {
1573 if self.structurally_monotonic {
1574 return 0.0;
1578 }
1579 self.monotonicity.tolerance.max(0.0)
1580 }
1581
1582 fn derivative_guard_numerical(&self) -> f64 {
1583 let derivative_guard = self.derivative_guard();
1584 if derivative_guard <= 0.0 {
1585 if self.structurally_monotonic {
1594 -1e-10
1595 } else {
1596 1e-12
1597 }
1598 } else {
1599 (derivative_guard - (1e-10_f64).min(0.01 * derivative_guard)).max(1e-12)
1600 }
1601 }
1602
1603 fn interval_increment_guard(&self, h_entry: f64, h_exit: f64) -> f64 {
1604 let scale = h_entry.abs().max(h_exit.abs()).max(1.0);
1605 1e-10 * scale
1606 }
1607
1608 fn structural_time_coefficient_constraints(&self) -> Option<LinearInequalityConstraints> {
1609 if !self.structurally_monotonic {
1610 return None;
1611 }
1612 let p = self.coefficient_dim();
1613 let time_columns = self.structural_time_columns.min(p);
1614 if time_columns == 0 {
1615 return None;
1616 }
1617 let mut a = Array2::<f64>::zeros((time_columns, p));
1638 let b = Array1::<f64>::zeros(time_columns);
1639 for j in 0..time_columns {
1640 a[[j, j]] = 1.0;
1641 }
1642 Some(LinearInequalityConstraints { a, b })
1643 }
1644
1645 pub fn monotonicity_linear_constraints(&self) -> Option<LinearInequalityConstraints> {
1646 let p = self.coefficient_dim();
1647 const DERIVATIVE_ROW_NORM_TOL: f64 = 1e-12;
1648 if p == 0 {
1649 return None;
1650 }
1651 if self.structurally_monotonic {
1652 return self.structural_time_coefficient_constraints();
1653 }
1654 if let (Some(rows), Some(offsets)) = (
1655 self.monotonicity_constraint_rows.as_ref(),
1656 self.monotonicity_constraint_offsets.as_ref(),
1657 ) {
1658 let activerows: Vec<usize> = (0..rows.nrows())
1659 .filter(|&i| {
1660 rows.row(i).iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()))
1661 > DERIVATIVE_ROW_NORM_TOL
1662 })
1663 .collect();
1664 if activerows.is_empty() {
1665 return None;
1666 }
1667 let mut a = Array2::<f64>::zeros((activerows.len(), p));
1668 let mut b = Array1::<f64>::zeros(activerows.len());
1669 for (r, &i) in activerows.iter().enumerate() {
1670 a.row_mut(r).assign(&rows.row(i));
1671 b[r] = self.derivative_guard() - offsets[i];
1672 }
1673 return Some(compress_positive_collinear_constraints(&a, &b));
1674 }
1675 None
1676 }
1677
1678 pub fn from_engine_inputs(
1679 inputs: SurvivalEngineInputs<'_>,
1680 penalties: PenaltyBlocks,
1681 monotonicity: SurvivalMonotonicityPenalty,
1682 spec: SurvivalSpec,
1683 ) -> Result<Self, SurvivalError> {
1684 Self::from_engine_inputswith_offsets(inputs, None, penalties, monotonicity, spec)
1685 }
1686
1687 fn validate_offsets(
1688 offsets: Option<SurvivalBaselineOffsets<'_>>,
1689 n: usize,
1690 ) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), SurvivalError> {
1691 if let Some(off) = offsets {
1692 if off.eta_entry.len() != n || off.eta_exit.len() != n || off.derivative_exit.len() != n
1693 {
1694 return Err(SurvivalError::DimensionMismatch);
1695 }
1696 if off.eta_entry.iter().any(|v| !v.is_finite())
1697 || off.eta_exit.iter().any(|v| !v.is_finite())
1698 || off.derivative_exit.iter().any(|v| !v.is_finite())
1699 {
1700 return Err(SurvivalError::NonFiniteInput);
1701 }
1702 Ok((
1703 off.eta_entry.to_owned(),
1704 off.eta_exit.to_owned(),
1705 off.derivative_exit.to_owned(),
1706 ))
1707 } else {
1708 Ok((Array1::zeros(n), Array1::zeros(n), Array1::zeros(n)))
1709 }
1710 }
1711
1712 fn validate_common_inputs(
1713 age_entry: &ArrayView1<f64>,
1714 age_exit: &ArrayView1<f64>,
1715 event_target: &ArrayView1<u8>,
1716 event_competing: &ArrayView1<u8>,
1717 sampleweight: &ArrayView1<f64>,
1718 ) -> Result<(), SurvivalError> {
1719 if age_entry.iter().any(|v| !v.is_finite())
1720 || age_exit.iter().any(|v| !v.is_finite())
1721 || sampleweight.iter().any(|v| !v.is_finite() || *v < 0.0)
1722 {
1723 return Err(SurvivalError::NonFiniteInput);
1724 }
1725 if let Some(&label) = event_target.iter().find(|&&v| v > 1) {
1732 return Err(SurvivalError::EventCodeInvalid {
1733 reason: format!(
1734 "single-hazard survival engine requires a binary {{0, 1}} event_target, got multi-cause label {label}; competing-risks codes must be projected via pooled_any_event_indicator / cause_specific_event_indicator before construction"
1735 ),
1736 });
1737 }
1738 if let Some(&label) = event_competing.iter().find(|&&v| v > 1) {
1739 return Err(SurvivalError::EventCodeInvalid {
1740 reason: format!(
1741 "single-hazard survival engine requires a binary {{0, 1}} event_competing, got multi-cause label {label}"
1742 ),
1743 });
1744 }
1745 if event_target
1746 .iter()
1747 .zip(event_competing.iter())
1748 .any(|(&target, &competing)| target > 0 && competing > 0)
1749 {
1750 return Err(SurvivalError::EventCodeInvalid {
1751 reason: "a row cannot be simultaneously a target event and a competing event"
1752 .to_string(),
1753 });
1754 }
1755 if age_entry
1769 .iter()
1770 .zip(age_exit.iter())
1771 .any(|(&entry, &exit)| entry < 0.0 || exit <= 0.0)
1772 {
1773 return Err(SurvivalError::NonFiniteInput);
1774 }
1775 Ok::<(), _>(())
1776 }
1777
1778 fn validate_monotonicity_constraints(
1779 rows: Option<ArrayView2<'_, f64>>,
1780 offsets: Option<ArrayView1<'_, f64>>,
1781 coefficient_dim: usize,
1782 ) -> Result<(Option<Array2<f64>>, Option<Array1<f64>>), SurvivalError> {
1783 match (rows, offsets) {
1784 (None, None) => Ok((None, None)),
1785 (Some(rows), Some(offsets)) => {
1786 if rows.ncols() != coefficient_dim
1787 || rows.nrows() != offsets.len()
1788 || rows.iter().any(|v| !v.is_finite())
1789 || offsets.iter().any(|v| !v.is_finite())
1790 {
1791 return Err(SurvivalError::DimensionMismatch);
1792 }
1793 Ok((Some(rows.to_owned()), Some(offsets.to_owned())))
1794 }
1795 _ => Err(SurvivalError::DimensionMismatch),
1796 }
1797 }
1798
1799 fn finish_construction(
1800 age_entry: ArrayView1<f64>,
1801 age_exit: ArrayView1<f64>,
1802 event_target: ArrayView1<u8>,
1803 sampleweight: ArrayView1<f64>,
1804 design: SurvivalDesign,
1805 offset_eta_entry: Array1<f64>,
1806 offset_eta_exit: Array1<f64>,
1807 offset_derivative_exit: Array1<f64>,
1808 penalties: PenaltyBlocks,
1809 monotonicity: SurvivalMonotonicityPenalty,
1810 monotonicity_constraint_rows: Option<Array2<f64>>,
1811 monotonicity_constraint_offsets: Option<Array1<f64>>,
1812 ) -> Self {
1813 let n = age_entry.len();
1814 Self {
1815 age_entry: age_entry.to_owned(),
1816 age_exit: age_exit.to_owned(),
1817 entry_at_origin: age_entry.mapv(|t| t <= ENTRY_AT_ORIGIN_THRESHOLD),
1818 event_target: event_target.to_owned(),
1819 sampleweight: sampleweight.to_owned(),
1820 design,
1821 offset_eta_entry,
1822 offset_eta_exit,
1823 offset_derivative_exit,
1824 penalties,
1825 monotonicity,
1826 structurally_monotonic: false,
1827 structural_time_columns: 0,
1828 monotonicity_constraint_rows,
1829 monotonicity_constraint_offsets,
1830 workspace: std::sync::Mutex::new(SurvivalWorkspace::new(n)),
1831 }
1832 }
1833
1834 pub fn from_engine_inputswith_offsets(
1835 inputs: SurvivalEngineInputs<'_>,
1836 offsets: Option<SurvivalBaselineOffsets<'_>>,
1837 penalties: PenaltyBlocks,
1838 monotonicity: SurvivalMonotonicityPenalty,
1839 spec: SurvivalSpec,
1840 ) -> Result<Self, SurvivalError> {
1841 if spec == SurvivalSpec::Crude {
1842 return Err(SurvivalError::UnsupportedSpec("crude"));
1843 }
1844 let n = inputs.age_entry.len();
1845 let p = inputs.x_entry.ncols();
1846 if inputs.age_exit.len() != n
1847 || inputs.event_target.len() != n
1848 || inputs.event_competing.len() != n
1849 || inputs.sampleweight.len() != n
1850 || inputs.x_entry.nrows() != n
1851 || inputs.x_exit.nrows() != n
1852 || inputs.x_derivative.nrows() != n
1853 || inputs.x_entry.ncols() != inputs.x_exit.ncols()
1854 || inputs.x_entry.ncols() != inputs.x_derivative.ncols()
1855 {
1856 return Err(SurvivalError::DimensionMismatch);
1857 }
1858 Self::validate_penalties(&penalties, p)?;
1859 Self::validate_common_inputs(
1860 &inputs.age_entry,
1861 &inputs.age_exit,
1862 &inputs.event_target,
1863 &inputs.event_competing,
1864 &inputs.sampleweight,
1865 )?;
1866 if inputs.x_entry.iter().any(|v| !v.is_finite())
1867 || inputs.x_exit.iter().any(|v| !v.is_finite())
1868 || inputs.x_derivative.iter().any(|v| !v.is_finite())
1869 {
1870 return Err(SurvivalError::NonFiniteInput);
1871 }
1872 let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
1873 Self::validate_offsets(offsets, n)?;
1874 let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
1875 Self::validate_monotonicity_constraints(
1876 inputs.monotonicity_constraint_rows,
1877 inputs.monotonicity_constraint_offsets,
1878 p,
1879 )?;
1880
1881 Ok(Self::finish_construction(
1882 inputs.age_entry,
1883 inputs.age_exit,
1884 inputs.event_target,
1885 inputs.sampleweight,
1886 SurvivalDesign::Flat {
1887 x_entry: inputs.x_entry.to_owned(),
1888 x_exit: inputs.x_exit.to_owned(),
1889 x_derivative: inputs.x_derivative.to_owned(),
1890 },
1891 offset_eta_entry,
1892 offset_eta_exit,
1893 offset_derivative_exit,
1894 penalties,
1895 monotonicity,
1896 monotonicity_constraint_rows,
1897 monotonicity_constraint_offsets,
1898 ))
1899 }
1900
1901 pub fn from_time_covariate_inputswith_offsets(
1902 inputs: SurvivalTimeCovarInputs<'_>,
1903 offsets: Option<SurvivalBaselineOffsets<'_>>,
1904 penalties: PenaltyBlocks,
1905 monotonicity: SurvivalMonotonicityPenalty,
1906 spec: SurvivalSpec,
1907 ) -> Result<Self, SurvivalError> {
1908 if spec == SurvivalSpec::Crude {
1909 return Err(SurvivalError::UnsupportedSpec("crude"));
1910 }
1911 let n = inputs.age_entry.len();
1912 let p_time = inputs.time_entry.ncols();
1913 let p_cov = inputs.covariates.ncols();
1914 let p = p_time + p_cov;
1915 if inputs.age_exit.len() != n
1916 || inputs.event_target.len() != n
1917 || inputs.event_competing.len() != n
1918 || inputs.sampleweight.len() != n
1919 || inputs.time_entry.nrows() != n
1920 || inputs.time_exit.nrows() != n
1921 || inputs.time_derivative.nrows() != n
1922 || inputs.covariates.nrows() != n
1923 || inputs.time_entry.ncols() != inputs.time_exit.ncols()
1924 || inputs.time_entry.ncols() != inputs.time_derivative.ncols()
1925 {
1926 return Err(SurvivalError::DimensionMismatch);
1927 }
1928 Self::validate_penalties(&penalties, p)?;
1929 Self::validate_common_inputs(
1930 &inputs.age_entry,
1931 &inputs.age_exit,
1932 &inputs.event_target,
1933 &inputs.event_competing,
1934 &inputs.sampleweight,
1935 )?;
1936 if inputs.time_entry.iter().any(|v| !v.is_finite())
1937 || inputs.time_exit.iter().any(|v| !v.is_finite())
1938 || inputs.time_derivative.iter().any(|v| !v.is_finite())
1939 || inputs.covariates.iter().any(|v| !v.is_finite())
1940 {
1941 return Err(SurvivalError::NonFiniteInput);
1942 }
1943 let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
1944 Self::validate_offsets(offsets, n)?;
1945 let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
1946 Self::validate_monotonicity_constraints(
1947 inputs.monotonicity_constraint_rows,
1948 inputs.monotonicity_constraint_offsets,
1949 p,
1950 )?;
1951
1952 Ok(Self::finish_construction(
1953 inputs.age_entry,
1954 inputs.age_exit,
1955 inputs.event_target,
1956 inputs.sampleweight,
1957 SurvivalDesign::TimeCovariateShared {
1958 time_entry: inputs.time_entry.to_owned(),
1959 time_exit: inputs.time_exit.to_owned(),
1960 time_derivative: inputs.time_derivative.to_owned(),
1961 covariates: inputs.covariates.to_owned(),
1962 },
1963 offset_eta_entry,
1964 offset_eta_exit,
1965 offset_derivative_exit,
1966 penalties,
1967 monotonicity,
1968 monotonicity_constraint_rows,
1969 monotonicity_constraint_offsets,
1970 ))
1971 }
1972
1973 pub fn set_penalty_lambdas(&mut self, lambdas: &[f64]) -> Result<(), EstimationError> {
1987 if lambdas.len() != self.penalties.blocks.len() {
1988 crate::bail_invalid_estim!(
1989 "set_penalty_lambdas expects {} lambdas, got {}",
1990 self.penalties.blocks.len(),
1991 lambdas.len()
1992 );
1993 }
1994 for (block, &lambda) in self.penalties.blocks.iter_mut().zip(lambdas.iter()) {
1995 if !lambda.is_finite() || lambda < 0.0 {
1996 crate::bail_invalid_estim!("penalty lambda must be finite and >= 0, got {lambda}");
1997 }
1998 block.lambda = lambda;
1999 }
2000 Ok(())
2001 }
2002
2003 pub fn set_structural_monotonicity(
2004 &mut self,
2005 enabled: bool,
2006 time_columns: usize,
2007 ) -> Result<(), EstimationError> {
2008 let p = self.coefficient_dim();
2009 if time_columns > p {
2010 crate::bail_invalid_estim!(
2011 "structural time columns {} exceed coefficient dimension {}",
2012 time_columns,
2013 p
2014 );
2015 }
2016 if enabled && time_columns == 0 {
2017 crate::bail_invalid_estim!("structural monotonicity requires at least one time column");
2018 }
2019 if enabled {
2020 const STRUCTURAL_DERIV_TOL: f64 = 1e-12;
2021 for (i, &offset) in self.offset_derivative_exit.iter().enumerate() {
2022 if offset < -STRUCTURAL_DERIV_TOL {
2023 crate::bail_invalid_estim!(
2024 "structural monotonicity requires nonnegative derivative offsets; found offset_derivative_exit[{i}]={offset:.3e}"
2025 );
2026 }
2027 }
2028 let mut derivative_row = vec![0.0_f64; p];
2029 for i in 0..self.nrows() {
2030 self.fill_derivative_row(i, &mut derivative_row);
2031 for j in 0..time_columns {
2032 let v = derivative_row[j];
2033 if v < -STRUCTURAL_DERIV_TOL {
2034 crate::bail_invalid_estim!(
2035 "structural monotonicity requires nonnegative time-derivative basis entries; found x_derivative[{i},{j}]={v:.3e}"
2036 );
2037 }
2038 }
2039 for j in time_columns..p {
2040 let v = derivative_row[j];
2041 if v.abs() > STRUCTURAL_DERIV_TOL {
2042 crate::bail_invalid_estim!(
2043 "structural monotonicity requires zero derivative contribution outside the time block; found x_derivative[{i},{j}]={v:.3e}"
2044 );
2045 }
2046 }
2047 }
2048 if let (Some(rows), Some(offsets)) = (
2049 self.monotonicity_constraint_rows.as_ref(),
2050 self.monotonicity_constraint_offsets.as_ref(),
2051 ) {
2052 for (i, &offset) in offsets.iter().enumerate() {
2053 if offset < -STRUCTURAL_DERIV_TOL {
2054 crate::bail_invalid_estim!(
2055 "structural monotonicity requires nonnegative collocation derivative offsets; found monotonicity_constraint_offsets[{i}]={offset:.3e}"
2056 );
2057 }
2058 }
2059 for i in 0..rows.nrows() {
2060 for j in 0..time_columns {
2061 let v = rows[[i, j]];
2062 if v < -STRUCTURAL_DERIV_TOL {
2063 crate::bail_invalid_estim!(
2064 "structural monotonicity requires nonnegative collocation derivative basis entries; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
2065 );
2066 }
2067 }
2068 for j in time_columns..p {
2069 let v = rows[[i, j]];
2070 if v.abs() > STRUCTURAL_DERIV_TOL {
2071 crate::bail_invalid_estim!(
2072 "structural monotonicity requires zero collocation derivative contribution outside the time block; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
2073 );
2074 }
2075 }
2076 }
2077 }
2078 }
2079 self.structurally_monotonic = enabled;
2080 self.structural_time_columns = if enabled { time_columns } else { 0 };
2081 Ok(())
2082 }
2083
2084 pub fn update_state(&self, beta: &Array1<f64>) -> Result<WorkingState, EstimationError> {
2085 if beta.len() != self.coefficient_dim() {
2086 crate::bail_invalid_estim!("survival beta dimension mismatch");
2087 }
2088
2089 let n = self.nrows();
2090 let p = self.coefficient_dim();
2091
2092 let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2118 let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2119 let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2120
2121 let mut nll = 0.0;
2122 let derivative_guard = self.derivative_guard();
2123 let derivative_guard_numerical = self.derivative_guard_numerical();
2124 let mut workspace = self.workspace.lock().unwrap();
2125 workspace.reset(n);
2126 let SurvivalWorkspace {
2127 w_event,
2128 w_event_inv_deriv,
2129 w_event_outer,
2130 w_hess_exit,
2131 w_hess_entry,
2132 } = &mut *workspace;
2133
2134 for i in 0..n {
2136 let w = self.sampleweight[i];
2137 if w <= 0.0 {
2138 continue;
2139 }
2140 let entry_age = self.age_entry[i];
2141 let exit_age = self.age_exit[i];
2142 if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
2143 crate::bail_invalid_estim!(
2144 "survival ages must be finite with age_exit >= age_entry"
2145 );
2146 }
2147 let d = f64::from(self.event_target[i]);
2148
2149 let has_entry_interval = !self.entry_at_origin[i];
2150 let interval_scale = if has_entry_interval {
2151 eta_exit[i].max(eta_entry[i])
2152 } else {
2153 eta_exit[i]
2154 };
2155 let h_e_scaled = (eta_exit[i] - interval_scale).exp();
2156 let h_s_scaled = if has_entry_interval {
2157 (eta_entry[i] - interval_scale).exp()
2158 } else {
2159 0.0
2160 };
2161 let interval_scaled = h_e_scaled - h_s_scaled;
2162 let interval = Self::scaled_exp_component(interval_scale, interval_scaled)?;
2163 let (deriv, deriv_slope) = self
2164 .stabilized_structural_derivative(derivative_raw[i])
2165 .unwrap_or((derivative_raw[i], 1.0));
2166 let mono_floor = if d > 0.0 {
2175 derivative_guard_numerical
2176 } else {
2177 0.0
2178 };
2179 if !deriv.is_finite() || deriv < mono_floor {
2180 return Err(EstimationError::ParameterConstraintViolation(format!(
2181 "survival monotonicity violated at row {}: d_eta/dt={:.3e} <= tolerance={:.3e}",
2182 i, deriv, derivative_guard
2183 )));
2184 }
2185 if has_entry_interval {
2186 let increment_guard = self.interval_increment_guard(h_s_scaled, h_e_scaled);
2187 if interval_scaled + increment_guard < 0.0 {
2188 return Err(EstimationError::ParameterConstraintViolation(format!(
2189 "survival cumulative hazard decreased over row {}: H(exit)-H(entry)={:.6e}",
2190 i, interval
2191 )));
2192 }
2193 }
2194 nll += w * interval;
2195
2196 let w_exit_i = w * eta_exit[i].exp();
2200 let w_entry_i = if has_entry_interval {
2201 w * eta_entry[i].exp()
2202 } else {
2203 0.0
2204 };
2205 if !w_exit_i.is_finite() {
2206 crate::bail_invalid_estim!(
2207 "survival interval term exceeds f64 range at row {i} (w*exp(eta_exit)={w_exit_i:.3e})"
2208 );
2209 }
2210 w_hess_exit[i] = w_exit_i;
2211 w_hess_entry[i] = w_entry_i;
2212
2213 if d > 0.0 {
2214 let inv_deriv = deriv_slope / deriv;
2218 nll += -w * (eta_exit[i] + deriv.ln());
2219 w_event[i] = w;
2220 w_event_inv_deriv[i] = w * inv_deriv;
2221 w_event_outer[i] = w * inv_deriv * inv_deriv;
2222 }
2223 }
2224
2225 let mut h = self.interval_hessian_blas(w_hess_exit, w_hess_entry);
2229 let mut grad = Array1::<f64>::zeros(p);
2233 let mut grad_comp = Array1::<f64>::zeros(p);
2234 let mut row_exit = vec![0.0_f64; p];
2235 let mut row_entry = vec![0.0_f64; p];
2236 let mut row_derivative = vec![0.0_f64; p];
2237 for i in 0..n {
2238 let w_interval_exit = w_hess_exit[i];
2239 let w_interval_entry = w_hess_entry[i];
2240 let w_event_exit = w_event[i];
2241 let w_event_derivative = w_event_inv_deriv[i];
2242 if w_interval_exit == 0.0
2243 && w_interval_entry == 0.0
2244 && w_event_exit == 0.0
2245 && w_event_derivative == 0.0
2246 {
2247 continue;
2248 }
2249 self.fill_exit_row(i, &mut row_exit);
2250 self.fill_entry_row(i, &mut row_entry);
2251 self.fill_derivative_row(i, &mut row_derivative);
2252 for j in 0..p {
2253 let contribution = w_interval_exit * row_exit[j]
2254 - w_interval_entry * row_entry[j]
2255 - w_event_exit * row_exit[j]
2256 - w_event_derivative * row_derivative[j];
2257 let t = grad[j] + contribution;
2258 if grad[j].abs() >= contribution.abs() {
2259 grad_comp[j] += (grad[j] - t) + contribution;
2260 } else {
2261 grad_comp[j] += (contribution - t) + grad[j];
2262 }
2263 grad[j] = t;
2264 }
2265 }
2266 grad += &grad_comp;
2267
2268 h += &self.derivative_xt_diag_x(w_event_outer);
2269
2270 let score_norm = array1_l2_norm(&grad);
2274
2275 let penaltygrad = self.penalties.gradient(beta);
2276 let penalty_quadratic_form = 2.0 * self.penalties.deviance(beta);
2287 let penaltygrad_norm = array1_l2_norm(&penaltygrad);
2288
2289 let mut totalgrad = grad;
2290 totalgrad += &penaltygrad;
2291
2292 self.penalties.addhessian_inplace(&mut h);
2293 let log_likelihood = -nll;
2300 let deviance = 2.0 * nll;
2301
2302 Ok(WorkingState {
2303 eta: LinearPredictor::new(eta_exit),
2304 gradient: totalgrad,
2305 hessian: gam_linalg::matrix::SymmetricMatrix::Dense(h),
2306 log_likelihood,
2307 deviance,
2308 penalty_term: penalty_quadratic_form,
2309 firth: gam_solve::pirls::FirthDiagnostics::Inactive,
2310 ridge_used: 0.0,
2311 hessian_curvature: gam_solve::pirls::HessianCurvatureKind::Observed,
2312 gradient_natural_scale: score_norm + penaltygrad_norm,
2313 })
2314 }
2315
2316 pub(crate) fn survival_hessian_derivative_correction(
2326 &self,
2327 beta: &Array1<f64>,
2328 u_k: &Array1<f64>,
2329 ) -> Result<Array2<f64>, EstimationError> {
2330 let p = beta.len();
2331 let n = self.nrows();
2332
2333 let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2334 let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2335 let deriv_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2336 let exp_entry = eta_entry.mapv(f64::exp);
2337 let exp_exit = eta_exit.mapv(f64::exp);
2338 let guard = self.derivative_guard();
2339 let guard_numerical = self.derivative_guard_numerical();
2340
2341 let jac = Array1::<f64>::ones(p);
2342 let curvature = Array1::<f64>::zeros(p);
2343 let third = Array1::<f64>::zeros(p);
2344
2345 let mut row_exit = vec![0.0_f64; p];
2346 let mut row_entry = vec![0.0_f64; p];
2347 let mut row_derivative = vec![0.0_f64; p];
2348 let mut ge = vec![0.0_f64; p];
2349 let mut gs = vec![0.0_f64; p];
2350 let mut gsd = vec![0.0_f64; p];
2351 let mut he = vec![0.0_f64; p];
2352 let mut hs = vec![0.0_f64; p];
2353 let mut hsd = vec![0.0_f64; p];
2354 let mut te = vec![0.0_f64; p];
2355 let mut ts = vec![0.0_f64; p];
2356 let mut tsd = vec![0.0_f64; p];
2357
2358 let mut b_dir = Array2::<f64>::zeros((p, p));
2359
2360 for i in 0..n {
2361 let w_i = self.sampleweight[i];
2362 if w_i <= 0.0 {
2363 continue;
2364 }
2365 let has_entry = !self.entry_at_origin[i];
2366 let mut deta_e = 0.0_f64;
2367 let mut deta_s = 0.0_f64;
2368 let mut ds = 0.0_f64;
2369 self.fill_exit_row(i, &mut row_exit);
2370 self.fill_entry_row(i, &mut row_entry);
2371 self.fill_derivative_row(i, &mut row_derivative);
2372 for j in 0..p {
2373 ge[j] = row_exit[j] * jac[j];
2374 gs[j] = row_entry[j] * jac[j];
2375 gsd[j] = row_derivative[j] * jac[j];
2376 he[j] = row_exit[j] * curvature[j];
2377 hs[j] = row_entry[j] * curvature[j];
2378 hsd[j] = row_derivative[j] * curvature[j];
2379 te[j] = row_exit[j] * third[j];
2380 ts[j] = row_entry[j] * third[j];
2381 tsd[j] = row_derivative[j] * third[j];
2382 deta_e += ge[j] * u_k[j];
2383 if has_entry {
2384 deta_s += gs[j] * u_k[j];
2385 }
2386 ds += gsd[j] * u_k[j];
2387 }
2388
2389 for r in 0..p {
2391 let dge_r = he[r] * u_k[r];
2392 let dgs_r = hs[r] * u_k[r];
2393 let dhe_r = te[r] * u_k[r];
2394 let dhs_r = ts[r] * u_k[r];
2395 for c in 0..p {
2396 let dge_c = he[c] * u_k[c];
2397 let dgs_c = hs[c] * u_k[c];
2398 let mut d_h_rc =
2399 exp_exit[i] * (deta_e * ge[r] * ge[c] + dge_r * ge[c] + ge[r] * dge_c);
2400 if r == c {
2401 d_h_rc += exp_exit[i] * (deta_e * he[r] + dhe_r);
2402 }
2403 if has_entry {
2404 d_h_rc -=
2405 exp_entry[i] * (deta_s * gs[r] * gs[c] + dgs_r * gs[c] + gs[r] * dgs_c);
2406 if r == c {
2407 d_h_rc -= exp_entry[i] * (deta_s * hs[r] + dhs_r);
2408 }
2409 }
2410 b_dir[[r, c]] += w_i * d_h_rc;
2411 }
2412 }
2413
2414 let (s_i, s_slope) = self
2416 .stabilized_structural_derivative(deriv_raw[i])
2417 .unwrap_or((deriv_raw[i], 1.0));
2418 if !s_i.is_finite() {
2419 return Err(EstimationError::ParameterConstraintViolation(format!(
2420 "survival monotonicity violated in unified trace contraction at row {i}: \
2421 d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
2422 )));
2423 }
2424 if self.event_target[i] > 0 && s_slope != 0.0 {
2425 if s_i < guard_numerical {
2430 return Err(EstimationError::ParameterConstraintViolation(format!(
2431 "survival monotonicity violated in unified trace contraction at row {i}: \
2432 d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
2433 )));
2434 }
2435 let inv_s = 1.0 / s_i;
2436 let inv_s2 = inv_s * inv_s;
2437 let inv_s3 = inv_s2 * inv_s;
2438 for r in 0..p {
2439 let dgd_r = hsd[r] * u_k[r];
2440 let dtsd_r = tsd[r] * u_k[r];
2441 let dte_r = te[r] * u_k[r];
2442 for c in 0..p {
2443 let dgd_c = hsd[c] * u_k[c];
2444 let mut d_h_rc = (dgd_r * gsd[c] + gsd[r] * dgd_c) * inv_s2
2445 - 2.0 * gsd[r] * gsd[c] * ds * inv_s3;
2446 if r == c {
2447 d_h_rc += -dte_r;
2448 d_h_rc += -(dtsd_r * inv_s - hsd[r] * ds * inv_s2);
2449 }
2450 b_dir[[r, c]] += w_i * d_h_rc;
2451 }
2452 }
2453 }
2454 }
2455
2456 Ok(b_dir)
2457 }
2458
2459 pub fn offset_channel_residuals(
2497 &self,
2498 beta: &Array1<f64>,
2499 ) -> Result<OffsetChannelResiduals, EstimationError> {
2500 if beta.len() != self.coefficient_dim() {
2501 crate::bail_invalid_estim!(
2502 "survival beta dimension mismatch in offset_channel_residuals"
2503 );
2504 }
2505 let n = self.nrows();
2506 let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2507 let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2508 let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2509
2510 let derivative_guard_numerical = self.derivative_guard_numerical();
2511 let mut r_exit = Array1::<f64>::zeros(n);
2512 let mut r_entry = Array1::<f64>::zeros(n);
2513 let mut r_deriv = Array1::<f64>::zeros(n);
2514
2515 for i in 0..n {
2516 let w = self.sampleweight[i];
2517 if w <= 0.0 {
2518 continue;
2519 }
2520 let entry_age = self.age_entry[i];
2521 let exit_age = self.age_exit[i];
2522 if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
2523 crate::bail_invalid_estim!(
2524 "survival ages must be finite with age_exit >= age_entry"
2525 );
2526 }
2527 let has_entry_interval = !self.entry_at_origin[i];
2528 let d = f64::from(self.event_target[i]);
2529 let w_exit_i = w * eta_exit[i].exp();
2533 let w_entry_i = if has_entry_interval {
2534 w * eta_entry[i].exp()
2535 } else {
2536 0.0
2537 };
2538 if !w_exit_i.is_finite() {
2539 crate::bail_invalid_estim!(
2540 "offset_channel_residuals: w*exp(eta_exit)={w_exit_i:.3e} non-finite at row {i}"
2541 );
2542 }
2543 r_exit[i] = w_exit_i - d * w;
2544 r_entry[i] = -w_entry_i;
2545 let deriv_raw = derivative_raw[i];
2550 let (deriv, deriv_slope) = self
2551 .stabilized_structural_derivative(deriv_raw)
2552 .unwrap_or((deriv_raw, 1.0));
2553 let mono_floor = if d > 0.0 {
2554 derivative_guard_numerical
2555 } else {
2556 0.0
2557 };
2558 if !deriv.is_finite() || deriv < mono_floor {
2559 return Err(EstimationError::ParameterConstraintViolation(format!(
2560 "offset_channel_residuals: derivative ≤ numerical guard at row {i}: {deriv:.3e}"
2561 )));
2562 }
2563 if d > 0.0 {
2564 r_deriv[i] = -w * d * deriv_slope / deriv;
2567 }
2568 }
2569
2570 let right = Array1::<f64>::zeros(r_exit.len());
2571 Ok(OffsetChannelResiduals {
2572 exit: r_exit,
2573 entry: r_entry,
2574 derivative: r_deriv,
2575 right,
2576 })
2577 }
2578
2579 pub fn unified_lamlobjective_and_rhogradient(
2585 &self,
2586 beta: &Array1<f64>,
2587 state: &WorkingState,
2588 rho: &Array1<f64>,
2589 ) -> Result<(f64, Array1<f64>), EstimationError> {
2590 use gam_problem::{EvalMode, PseudoLogdetMode};
2591 use gam_solve::estimate::reml::assembly::InnerAssembly;
2592 use gam_solve::estimate::reml::reml_outer_engine::{
2593 DenseSpectralOperator, DispersionHandling,
2594 };
2595 use gam_solve::estimate::reml::reparameterized_inner::{
2596 RawInnerReparamContext, assemble_reparameterized_inner,
2597 };
2598 use gam_terms::construction::{
2599 canonicalize_penalty_specs, precompute_reparam_invariant_from_canonical,
2600 stable_reparameterizationwith_invariant,
2601 };
2602 use gam_terms::penalty_spec::PenaltySpec;
2603
2604 let p = beta.len();
2605 let active_penalty_blocks: Vec<&PenaltyBlock> = self
2606 .penalties
2607 .blocks
2608 .iter()
2609 .filter(|b| b.lambda > 0.0)
2610 .collect();
2611 if rho.len() != active_penalty_blocks.len() {
2612 crate::bail_invalid_estim!(
2613 "survival LAML rho dimension {} does not match active penalty block count {}",
2614 rho.len(),
2615 active_penalty_blocks.len()
2616 );
2617 }
2618 let k_count = active_penalty_blocks.len();
2619
2620 let lambdas: Vec<f64> = rho.iter().map(|&r| r.exp()).collect();
2623
2624 let h_dense = state.hessian.to_dense();
2631 let has_left_truncation = self
2632 .age_entry
2633 .iter()
2634 .any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD);
2635 let hessian_logdet_mode = if has_left_truncation {
2636 PseudoLogdetMode::HardPseudo
2637 } else {
2638 PseudoLogdetMode::Smooth
2639 };
2640
2641 let s_k_embedded: Vec<Array2<f64>> = active_penalty_blocks
2649 .iter()
2650 .map(|b| {
2651 let mut s = Array2::<f64>::zeros((p, p));
2652 let (rs, re) = (b.range.start, b.range.end);
2653 s.slice_mut(ndarray::s![rs..re, rs..re]).assign(&b.matrix);
2654 s
2655 })
2656 .collect();
2657
2658 let penalty_specs: Vec<PenaltySpec> = active_penalty_blocks
2678 .iter()
2679 .map(|b| PenaltySpec::Block {
2680 local: b.matrix.clone(),
2681 col_range: b.range.clone(),
2682 prior_mean: gam_problem::CoefficientPriorMean::Zero,
2683 structure_hint: None,
2684 op: None,
2685 })
2686 .collect();
2687 let nullspace_dims: Vec<usize> = active_penalty_blocks
2688 .iter()
2689 .map(|b| b.nullspace_dim)
2690 .collect();
2691 let (canonical_penalties, _canonical_nullspace) = canonicalize_penalty_specs(
2692 &penalty_specs,
2693 &nullspace_dims,
2694 p,
2695 "survival LAML seam-A reparameterization",
2696 )
2697 .map_err(|e| EstimationError::InvalidInput(e.to_string()))?;
2698 if canonical_penalties.len() != k_count {
2699 return Err(EstimationError::InvalidInput(format!(
2700 "survival LAML reparameterization dropped {} of {} active (λ>0) penalty \
2701 block(s) as numerically rank-0; cannot align transformed penalty \
2702 coordinates with ρ",
2703 k_count - canonical_penalties.len(),
2704 k_count
2705 )));
2706 }
2707 let reparam_invariant =
2708 precompute_reparam_invariant_from_canonical(&canonical_penalties, p)
2709 .map_err(|e| EstimationError::InvalidInput(e.to_string()))?;
2710 let reparam = stable_reparameterizationwith_invariant(
2711 &canonical_penalties,
2712 &lambdas,
2713 p,
2714 &reparam_invariant,
2715 None,
2719 )
2720 .map_err(|e| EstimationError::InvalidInput(e.to_string()))?;
2721
2722 let provider = SurvivalDerivProvider::new(self.clone(), beta.clone());
2728 let ctx = RawInnerReparamContext {
2729 hessian: &h_dense,
2730 beta,
2731 penalties_embedded: &s_k_embedded,
2732 lambdas: &lambdas,
2733 };
2734 let reparam_inner = assemble_reparameterized_inner(
2735 &ctx,
2736 Some(Box::new(provider)),
2737 &reparam,
2738 )
2739 .map_err(EstimationError::InvalidInput)?;
2740
2741 let hop = DenseSpectralOperator::from_symmetric_with_mode(
2744 &reparam_inner.hessian_transformed,
2745 hessian_logdet_mode,
2746 )
2747 .map_err(EstimationError::InvalidInput)?;
2748
2749 let penalty_coords = reparam
2756 .canonical_transformed
2757 .iter()
2758 .map(|cp| cp.to_penalty_coordinate())
2759 .collect::<Vec<_>>();
2760
2761 let penalty_quadratic = state.penalty_term;
2766
2767 const SURVIVAL_LAML_IFT_RELATIVE_KKT_GATE: f64 = 1.0e-8;
2778 let kkt_residual = {
2779 let raw = state.gradient.clone();
2780 let projected = match self.monotonicity_linear_constraints() {
2781 Some(constraints) => {
2782 let constraints = ConstraintSet::Dense(constraints);
2783 projected_linear_constraint_stationarity_vector(&raw, beta, &constraints, None)
2784 .ok_or_else(|| {
2785 EstimationError::InvalidInput(
2786 "survival LAML could not project the monotonicity KKT residual"
2787 .to_string(),
2788 )
2789 })?
2790 }
2791 None => raw,
2792 };
2793 let projected_norm = array1_l2_norm(&projected);
2794 let relative_projected_norm = state.relative_gradient_norm(projected_norm);
2795 if relative_projected_norm <= SURVIVAL_LAML_IFT_RELATIVE_KKT_GATE {
2796 let projected_transformed = reparam.qs.t().dot(&projected);
2797 Some(crate::model_types::ProjectedKktResidual::from_active_projected(
2798 projected_transformed,
2799 ))
2800 } else {
2801 None
2802 }
2803 };
2804
2805 let result = InnerAssembly {
2806 log_likelihood: state.log_likelihood,
2807 penalty_quadratic,
2808 beta: reparam_inner.beta_transformed,
2809 n_observations: self.nrows(),
2810 hessian_op: std::sync::Arc::new(hop),
2811 penalty_coords,
2812 penalty_logdet: reparam_inner.penalty_logdet,
2813 dispersion: DispersionHandling::Fixed {
2814 phi: 1.0,
2815 include_logdet_h: true,
2816 include_logdet_s: true,
2817 },
2818 rho_curvature_scale: 1.0,
2819 rho_prior: gam_problem::RhoPrior::Flat,
2820 hessian_logdet_correction: 0.0,
2821 penalty_subspace_trace: None,
2822 deriv_provider: reparam_inner.deriv_provider,
2823 firth: None,
2824 nullspace_dim: None,
2825 barrier_config: None,
2826 ext_coords: Vec::new(),
2827 ext_coord_pair_fn: None,
2828 rho_ext_pair_fn: None,
2829 fixed_drift_deriv: None,
2830 contracted_psi_second_order: None,
2831 kkt_residual,
2832 active_constraints: None,
2833 }
2834 .evaluate(
2835 rho.as_slice().expect("rho must be contiguous"),
2836 EvalMode::ValueAndGradient,
2837 None,
2838 )
2839 .map_err(EstimationError::InvalidInput)?;
2840
2841 let gradient = result.gradient.unwrap_or_else(|| Array1::zeros(rho.len()));
2842 Ok((result.cost, gradient))
2843 }
2844
2845
2846 pub fn evaluate_survival_lamlcost_and_gradient(
2867 &self,
2868 rho: &[f64],
2869 beta0: &Array1<f64>,
2870 ) -> Result<(f64, Array1<f64>), EstimationError> {
2871 let (candidate, beta) = self.reconverge_survival_inner_mode(rho, beta0)?;
2872 let rho_arr = Array1::from_vec(rho.to_vec());
2877 let state = candidate.update_state(&beta)?;
2878 candidate.unified_lamlobjective_and_rhogradient(&beta, &state, &rho_arr)
2879 }
2880
2881 fn reconverge_survival_inner_mode(
2891 &self,
2892 rho: &[f64],
2893 beta0: &Array1<f64>,
2894 ) -> Result<(WorkingModelSurvival, Array1<f64>), EstimationError> {
2895 const SHIM_PIRLS_MAX_ITERATIONS: usize = 600;
2900 const SHIM_PIRLS_CONVERGENCE_TOL: f64 = 1e-12;
2901 const SHIM_PIRLS_MAX_STEP_HALVING: usize = 40;
2902 const SHIM_PIRLS_MIN_STEP_SIZE: f64 = 1e-12;
2903
2904 let active_block_count = self
2905 .penalties
2906 .blocks
2907 .iter()
2908 .filter(|b| b.lambda > 0.0)
2909 .count();
2910 if rho.len() != active_block_count {
2911 crate::bail_invalid_estim!(
2912 "reconverge_survival_inner_mode: rho dimension {} does not match active penalty block count {}",
2913 rho.len(),
2914 active_block_count
2915 );
2916 }
2917 if beta0.len() != self.coefficient_dim() {
2918 crate::bail_invalid_estim!(
2919 "reconverge_survival_inner_mode: beta0 dimension {} does not match coefficient dimension {}",
2920 beta0.len(),
2921 self.coefficient_dim()
2922 );
2923 }
2924 let active_lambdas = gam_problem::checked_exp_log_strengths(rho.iter().copied())?;
2925
2926 let mut candidate = self.clone();
2929 let mut lambdas: Vec<f64> = candidate
2930 .penalties
2931 .blocks
2932 .iter()
2933 .map(|b| b.lambda)
2934 .collect();
2935 let mut active_idx = 0usize;
2936 for (block, lambda) in candidate.penalties.blocks.iter().zip(lambdas.iter_mut()) {
2937 if block.lambda > 0.0 {
2938 *lambda = active_lambdas[active_idx];
2939 active_idx += 1;
2940 }
2941 }
2942 candidate.set_penalty_lambdas(&lambdas)?;
2943
2944 let opts = gam_solve::pirls::WorkingModelPirlsOptions {
2945 max_iterations: SHIM_PIRLS_MAX_ITERATIONS,
2946 convergence_tolerance: SHIM_PIRLS_CONVERGENCE_TOL,
2947 adaptive_kkt_tolerance: None,
2948 max_step_halving: SHIM_PIRLS_MAX_STEP_HALVING,
2949 min_step_size: SHIM_PIRLS_MIN_STEP_SIZE,
2950 firth_bias_reduction: false,
2951 coefficient_lower_bounds: None,
2952 linear_constraints: None,
2953 initial_lm_lambda: None,
2954 arrow_schur: None,
2955 };
2956 let summary = gam_solve::pirls::runworking_model_pirls(
2957 &mut candidate,
2958 Coefficients::new(beta0.clone()),
2959 &opts,
2960 |_| {},
2961 )?;
2962 let mut beta = summary.beta.as_ref().to_owned();
2963
2964 {
2986 const POLISH_MAX_ITERS: usize = 400;
2987 const POLISH_TOL: f64 = 1e-13;
2988 const ARMIJO_C: f64 = constants::ARMIJO_C1;
2992 const BACKTRACK: f64 = constants::BACKTRACK_CONTRACTION;
2993 const MAX_BACKTRACK: usize = 80;
2994 let p = beta.len();
2995 let penalized_objective =
3002 |st: &WorkingState| -> f64 { -st.log_likelihood + 0.5 * st.penalty_term };
3003 for _ in 0..POLISH_MAX_ITERS {
3004 let st = match candidate.update_state(&beta) {
3005 Ok(st) => st,
3006 Err(_) => break,
3007 };
3008 let r = st.gradient.clone();
3009 let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
3010 if !r_norm.is_finite() || r_norm < POLISH_TOL {
3011 break;
3012 }
3013 let h = st.hessian.to_dense();
3014 let f0 = penalized_objective(&st);
3015 let h_scale = (0..p)
3030 .map(|d| h[[d, d]].abs())
3031 .fold(0.0_f64, f64::max)
3032 .max(1.0);
3033 let try_lm = |lambda_lm: f64| -> Option<(Array1<f64>, f64)> {
3049 let mut h_reg = h.clone();
3050 for d in 0..p {
3051 h_reg[[d, d]] += lambda_lm;
3052 }
3053 let factor =
3054 gam_linalg::faer_ndarray::FaerCholesky::cholesky(&h_reg, faer::Side::Lower)
3055 .ok()?;
3056 let candidate_step = factor.solvevec(&r);
3057 if candidate_step.iter().any(|v| !v.is_finite()) {
3058 return None;
3059 }
3060 let dd = -r.dot(&candidate_step);
3061 (dd.is_finite() && dd < -1e-14 * r_norm * r_norm)
3062 .then_some((candidate_step, dd))
3063 };
3064 let (step, dir_deriv) = try_lm(0.0)
3067 .or_else(|| {
3068 escalate_ridge(RidgeSchedule::geometric(1e-11 * h_scale, 17), try_lm)
3069 .ok()
3070 .map(|success| success.value)
3071 })
3072 .unwrap_or_else(|| {
3073 (r.clone(), -r_norm * r_norm)
3076 });
3077 let accepted = match backtracking_line_search::<_, Infallible>(
3094 BacktrackConfig {
3095 contraction: BACKTRACK,
3096 max_steps: MAX_BACKTRACK,
3097 ..BacktrackConfig::default()
3098 },
3099 |alpha| {
3100 let trial = &beta - &(alpha * &step);
3101 let Ok(ts) = candidate.update_state(&trial) else {
3102 return Ok(None);
3103 };
3104 let ft = penalized_objective(&ts);
3105 let tn = ts.gradient.iter().map(|v| v * v).sum::<f64>().sqrt();
3106 let armijo_ok = ft.is_finite() && ft <= f0 + ARMIJO_C * alpha * dir_deriv;
3107 let residual_ok = tn.is_finite() && tn < r_norm;
3108 Ok((armijo_ok || residual_ok).then_some((ft, trial)))
3109 },
3110 |_alpha, _ft| true,
3111 ) {
3112 Ok(result) => result,
3113 Err(never) => match never {},
3114 };
3115 let Some(ls) = accepted else {
3116 break;
3117 };
3118 beta = ls.payload;
3119 }
3120 }
3121
3122 Ok((candidate, beta))
3123 }
3124}
3125
3126pub(crate) struct SurvivalDerivProvider {
3135 model: WorkingModelSurvival,
3136 beta: Array1<f64>,
3137}
3138
3139impl SurvivalDerivProvider {
3140 pub(crate) fn new(model: WorkingModelSurvival, beta: Array1<f64>) -> Self {
3141 Self { model, beta }
3142 }
3143}
3144
3145impl gam_solve::estimate::reml::reml_outer_engine::HessianDerivativeProvider
3146 for SurvivalDerivProvider
3147{
3148 fn hessian_derivative_correction(
3149 &self,
3150 v_k: &Array1<f64>,
3151 ) -> Result<Option<Array2<f64>>, String> {
3152 let u_k = -v_k;
3155 match self
3156 .model
3157 .survival_hessian_derivative_correction(&self.beta, &u_k)
3158 {
3159 Ok(correction) => Ok(Some(correction)),
3160 Err(e) => Err(e.to_string()),
3161 }
3162 }
3163
3164 fn has_corrections(&self) -> bool {
3165 true
3166 }
3167}
3168
3169#[derive(Debug, Clone)]
3170pub struct CrudeRiskResult {
3171 pub risk: f64,
3172 pub diseasegradient: Array1<f64>,
3173 pub mortalitygradient: Array1<f64>,
3174}
3175
3176#[derive(Debug, Clone)]
3177pub struct CompetingRisksCifResult {
3178 pub cif: Vec<Array2<f64>>,
3183 pub overall_survival: Array2<f64>,
3184}
3185
3186const COMPETING_RISKS_CIF_PARALLEL_ROW_MIN: usize = 256;
3191
3192pub fn assemble_competing_risks_cif(
3193 times: ArrayView1<'_, f64>,
3194 cumulative_hazard: ArrayView3<'_, f64>,
3195) -> Result<CompetingRisksCifResult, SurvivalError> {
3196 let (n_endpoints, n_rows, n_times) = cumulative_hazard.dim();
3197 if n_endpoints == 0 {
3198 return Err(SurvivalError::DimensionMismatch);
3199 }
3200 let endpoint_hazards = cumulative_hazard
3201 .axis_iter(Axis(0))
3202 .map(|view| view.to_owned())
3203 .collect::<Vec<_>>();
3204 assemble_competing_risks_cif_from_endpoints(times, &endpoint_hazards).and_then(|result| {
3205 if result.overall_survival.dim() != (n_rows, n_times) {
3206 Err(SurvivalError::DimensionMismatch)
3207 } else {
3208 Ok(result)
3209 }
3210 })
3211}
3212
3213pub fn assemble_competing_risks_cif_from_endpoints(
3214 times: ArrayView1<'_, f64>,
3215 cumulative_hazards: &[Array2<f64>],
3216) -> Result<CompetingRisksCifResult, SurvivalError> {
3217 let n_endpoints = cumulative_hazards.len();
3218 if n_endpoints == 0 || times.is_empty() {
3219 return Err(SurvivalError::DimensionMismatch);
3220 }
3221 let (n_rows, n_times) = cumulative_hazards[0].dim();
3222 if n_rows == 0 || n_times == 0 || times.len() != n_times {
3223 return Err(SurvivalError::DimensionMismatch);
3224 }
3225 if times.iter().any(|time| !time.is_finite() || *time < 0.0) {
3226 return Err(SurvivalError::InvalidTimeGrid);
3227 }
3228 if times
3229 .iter()
3230 .zip(times.iter().skip(1))
3231 .any(|(previous, current)| current <= previous)
3232 {
3233 return Err(SurvivalError::InvalidTimeGrid);
3234 }
3235 for endpoint_hazard in cumulative_hazards {
3236 if endpoint_hazard.dim() != (n_rows, n_times) {
3237 return Err(SurvivalError::DimensionMismatch);
3238 }
3239 if endpoint_hazard.iter().any(|value| !value.is_finite()) {
3240 return Err(SurvivalError::NonFiniteInput);
3241 }
3242 }
3243
3244 let max_abs_hazard = cumulative_hazards
3245 .iter()
3246 .flat_map(|endpoint_hazard| endpoint_hazard.iter())
3247 .fold(0.0_f64, |acc, value| acc.max(value.abs()));
3248 let monotone_tolerance = 1.0e-10_f64 * max_abs_hazard.max(1.0);
3249 let mut cif: Vec<Array2<f64>> = (0..n_endpoints)
3250 .map(|_| Array2::<f64>::zeros((n_rows, n_times)))
3251 .collect();
3252 let mut overall_survival = Array2::<f64>::zeros((n_rows, n_times));
3253
3254 let assemble_row = |row: usize| -> Result<(Vec<f64>, Vec<f64>), SurvivalError> {
3266 let mut cif_flat = vec![0.0_f64; n_endpoints * n_times];
3267 let mut surv_row = vec![0.0_f64; n_times];
3268 let mut previous_cif = vec![0.0_f64; n_endpoints];
3269 let mut previous_cumulative = vec![0.0_f64; n_endpoints];
3270 let mut increments = vec![0.0_f64; n_endpoints];
3271 let mut previous_total_cumulative = 0.0_f64;
3272 for time_idx in 0..n_times {
3273 let mut total_increment = 0.0_f64;
3274 for endpoint in 0..n_endpoints {
3275 let current = cumulative_hazards[endpoint][[row, time_idx]];
3276 if current < -monotone_tolerance {
3277 return Err(SurvivalError::NonMonotoneCumulativeHazard);
3278 }
3279 let raw_increment = current - previous_cumulative[endpoint];
3280 if raw_increment < -monotone_tolerance {
3281 return Err(SurvivalError::NonMonotoneCumulativeHazard);
3282 }
3283 let increment = raw_increment.max(0.0);
3284 increments[endpoint] = increment;
3285 total_increment += increment;
3286 previous_cumulative[endpoint] += increment;
3287 }
3288
3289 let survival_left = (-previous_total_cumulative).exp();
3290 let interval_failure = -(-total_increment).exp_m1();
3291 for endpoint in 0..n_endpoints {
3292 if total_increment > 0.0 {
3293 previous_cif[endpoint] +=
3294 survival_left * interval_failure * increments[endpoint] / total_increment;
3295 }
3296 cif_flat[endpoint * n_times + time_idx] = previous_cif[endpoint].clamp(0.0, 1.0);
3297 }
3298 previous_total_cumulative += total_increment;
3299 let mut fsum_at_t = 0.0_f64;
3316 for endpoint in 0..n_endpoints {
3317 fsum_at_t += cif_flat[endpoint * n_times + time_idx];
3318 }
3319 surv_row[time_idx] = (1.0_f64 - fsum_at_t).clamp(0.0, 1.0);
3320 }
3321 Ok((cif_flat, surv_row))
3322 };
3323
3324 let rows: Vec<(Vec<f64>, Vec<f64>)> = if n_rows >= COMPETING_RISKS_CIF_PARALLEL_ROW_MIN
3328 && rayon::current_thread_index().is_none()
3329 {
3330 use rayon::prelude::*;
3331 (0..n_rows)
3332 .into_par_iter()
3333 .map(assemble_row)
3334 .collect::<Result<_, _>>()?
3335 } else {
3336 (0..n_rows).map(assemble_row).collect::<Result<_, _>>()?
3337 };
3338
3339 for (row, (cif_flat, surv_row)) in rows.into_iter().enumerate() {
3340 for endpoint in 0..n_endpoints {
3341 for time_idx in 0..n_times {
3342 cif[endpoint][[row, time_idx]] = cif_flat[endpoint * n_times + time_idx];
3343 }
3344 }
3345 for time_idx in 0..n_times {
3346 overall_survival[[row, time_idx]] = surv_row[time_idx];
3347 }
3348 }
3349
3350 Ok(CompetingRisksCifResult {
3351 cif,
3352 overall_survival,
3353 })
3354}
3355
3356fn compute_gauss_legendre_nodes(n: usize) -> Vec<(f64, f64)> {
3360 let (nodes, weights) = gam_math::special::gauss_legendre(n);
3361 nodes.into_iter().zip(weights).collect()
3362}
3363
3364fn gauss_legendre_quadrature() -> &'static [(f64, f64)] {
3365 static CACHE: LazyLock<Vec<(f64, f64)>> = LazyLock::new(|| compute_gauss_legendre_nodes(40));
3371 &CACHE
3372}
3373
3374pub fn calculate_crude_risk_quadrature<F>(
3398 t0: f64,
3399 t1: f64,
3400 breakpoints: &[f64],
3401 h_dis_t0: f64,
3402 h_mor_t0: f64,
3403 design_d_t0: ArrayView1<'_, f64>,
3404 design_m_t0: ArrayView1<'_, f64>,
3405 mut eval_at: F,
3406) -> Result<CrudeRiskResult, SurvivalError>
3407where
3408 F: FnMut(
3409 f64,
3410 &mut Array1<f64>,
3411 &mut Array1<f64>,
3412 &mut Array1<f64>,
3413 ) -> Result<(f64, f64, f64), SurvivalError>,
3414{
3415 let coeff_len_d = design_d_t0.len();
3416 let coeff_len_m = design_m_t0.len();
3417 if coeff_len_d == 0 || coeff_len_m == 0 {
3418 return Err(SurvivalError::InvalidIntegrationSetup);
3419 }
3420 if !t0.is_finite()
3421 || !t1.is_finite()
3422 || !h_dis_t0.is_finite()
3423 || !h_mor_t0.is_finite()
3424 || design_d_t0.iter().any(|v| !v.is_finite())
3425 || design_m_t0.iter().any(|v| !v.is_finite())
3426 {
3427 return Err(SurvivalError::NonFiniteInput);
3428 }
3429 if t1 <= t0 {
3430 return Ok(CrudeRiskResult {
3431 risk: 0.0,
3432 diseasegradient: Array1::zeros(coeff_len_d),
3433 mortalitygradient: Array1::zeros(coeff_len_m),
3434 });
3435 }
3436
3437 let mut sorted_breaks: Vec<f64> = breakpoints
3438 .iter()
3439 .copied()
3440 .filter(|x| x.is_finite() && *x >= t0 && *x <= t1)
3441 .collect();
3442 sorted_breaks.push(t0);
3443 sorted_breaks.push(t1);
3444 sorted_breaks.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
3445 sorted_breaks.dedup_by(|a, b| (*a - *b).abs() < 1e-6);
3446 if sorted_breaks.len() < 2 {
3447 return Err(SurvivalError::InvalidIntegrationSetup);
3448 }
3449
3450 let mut total_risk = 0.0;
3451 let mut diseasegradient = Array1::zeros(coeff_len_d);
3452 let mut mortalitygradient = Array1::zeros(coeff_len_m);
3453 let nodesweights = gauss_legendre_quadrature();
3454
3455 let mut design_d = Array1::<f64>::zeros(coeff_len_d);
3456 let mut deriv_d = Array1::<f64>::zeros(coeff_len_d);
3457 let mut design_m = Array1::<f64>::zeros(coeff_len_m);
3458
3459 for segment in sorted_breaks.windows(2) {
3460 let a = segment[0];
3461 let b = segment[1];
3462 let center = 0.5 * (b + a);
3463 let halfwidth = 0.5 * (b - a);
3464 if halfwidth <= 0.0 {
3465 continue;
3466 }
3467
3468 for &(x, w) in nodesweights {
3469 let u = center + halfwidth * x;
3470 let (inst_hazard_d, hazard_d, hazard_m) =
3471 eval_at(u, &mut design_d, &mut deriv_d, &mut design_m)?;
3472 if !inst_hazard_d.is_finite() || !hazard_d.is_finite() || !hazard_m.is_finite() {
3473 return Err(SurvivalError::NonFiniteInput);
3474 }
3475 if inst_hazard_d <= 0.0 {
3476 return Err(SurvivalError::NonPositiveHazard);
3477 }
3478
3479 if hazard_d < h_dis_t0 || hazard_m < h_mor_t0 {
3480 return Err(SurvivalError::NonMonotoneCumulativeHazard);
3481 }
3482
3483 let h_dis_cond = hazard_d - h_dis_t0;
3484 let h_mor_cond = hazard_m - h_mor_t0;
3485 let s_total = (-(h_dis_cond + h_mor_cond)).exp();
3486
3487 total_risk += w * inst_hazard_d * s_total * halfwidth;
3488
3489 let weight = w * s_total * halfwidth;
3495 for j in 0..coeff_len_d {
3496 let d_inst_hazard = inst_hazard_d * design_d[j] + hazard_d * deriv_d[j];
3497 let d_hazard_cond = hazard_d * design_d[j] - h_dis_t0 * design_d_t0[j];
3498 let g = d_inst_hazard - inst_hazard_d * d_hazard_cond;
3499 diseasegradient[j] += weight * g;
3500 }
3501
3502 let weight = w * inst_hazard_d * s_total * halfwidth;
3505 for j in 0..coeff_len_m {
3506 let g = -hazard_m * design_m[j] + h_mor_t0 * design_m_t0[j];
3507 mortalitygradient[j] += weight * g;
3508 }
3509 }
3510 }
3511
3512 Ok(CrudeRiskResult {
3513 risk: total_risk,
3514 diseasegradient,
3515 mortalitygradient,
3516 })
3517}
3518
3519impl PirlsWorkingModel for WorkingModelSurvival {
3520 fn update(&mut self, beta: &Coefficients) -> Result<WorkingState, EstimationError> {
3521 self.update_state(beta)
3522 }
3523}
3524
3525#[cfg(test)]
3526mod tests {
3527 use super::*;
3528 use ndarray::{Array1, Array2, Array3, array, s};
3529
3530 #[test]
3531 fn saved_cause_specific_alo_matches_independent_closed_form() {
3532 let eta_exit = 0.4_f64;
3533 let eta_entry = -0.3_f64;
3534 let derivative_exit = 1.7_f64;
3535 let weight = 2.2_f64;
3536 let geometry = cause_specific_survival_alo_row_geometry(CauseSpecificSurvivalAloRowInput {
3537 eta_exit,
3538 eta_entry,
3539 derivative_exit,
3540 prior_weight: weight,
3541 entry_active: true,
3542 event: true,
3543 })
3544 .expect("valid cause-specific row");
3545 let expected_nll =
3546 weight * (eta_exit.exp() - eta_entry.exp() - eta_exit - derivative_exit.ln());
3547 let expected_score = [
3548 weight * (eta_exit.exp() - 1.0),
3549 -weight * eta_entry.exp(),
3550 -weight / derivative_exit,
3551 ];
3552 let expected_hessian = [
3553 [weight * eta_exit.exp(), 0.0, 0.0],
3554 [0.0, -weight * eta_entry.exp(), 0.0],
3555 [0.0, 0.0, weight / derivative_exit.powi(2)],
3556 ];
3557 assert!((geometry.negative_log_likelihood - expected_nll).abs() <= 2.0e-14);
3558 for row in 0..3 {
3559 assert!((geometry.nll_score[row] - expected_score[row]).abs() <= 2.0e-14);
3560 for column in 0..3 {
3561 assert!(
3562 (geometry.observed_hessian[row][column] - expected_hessian[row][column]).abs()
3563 <= 2.0e-14
3564 );
3565 }
3566 }
3567 let score_meat = geometry.nll_score[0] * geometry.nll_score[0];
3568 assert!(
3569 (geometry.observed_hessian[0][0] - score_meat).abs() > 1.0e-2,
3570 "survival observed W and empirical score meat C must remain separate"
3571 );
3572 }
3573
3574 mod jet_cause_specific_production_parity {
3585 use super::*;
3586 use gam_math::jet_tower::{
3587 program_fourth_contracted, program_row_kernel, program_third_contracted,
3588 };
3589
3590 fn identity_block(w: f64, has_entry: bool, event: bool) -> CauseSpecificRoystonParmarBlock {
3597 let age_entry = if has_entry { 1.0 } else { 0.0 };
3598 CauseSpecificRoystonParmarBlock {
3599 age_entry: array![age_entry],
3600 age_exit: array![2.0],
3601 event_target: array![if event { 1u8 } else { 0u8 }],
3602 sampleweight: array![w],
3603 x_entry: array![[0.0, 1.0, 0.0]],
3604 x_exit: array![[1.0, 0.0, 0.0]],
3605 x_derivative: array![[0.0, 0.0, 1.0]],
3606 offset_eta_entry: array![0.0],
3607 offset_eta_exit: array![0.0],
3608 offset_derivative_exit: array![0.0],
3609 derivative_floor: 0.0,
3610 structural_time_columns: 0,
3611 }
3612 }
3613
3614 fn close(hand: f64, jet: f64, tol: f64, label: &str) {
3615 let band = tol + tol * hand.abs().max(jet.abs());
3616 assert!(
3617 (hand - jet).abs() <= band,
3618 "{label}: hand {hand:+.15e} vs jet {jet:+.15e} (|Δ|={:.3e} band {band:.3e})",
3619 (hand - jet).abs()
3620 );
3621 }
3622
3623 const JET_TOL: f64 = 1e-9;
3624
3625 fn run_corner(has_entry: bool, event: bool) {
3626 let beta = array![0.4_f64, -0.3_f64, 1.3_f64];
3628 let d_beta = array![0.7_f64, -0.5_f64, 0.6_f64];
3629 let v_beta = array![-0.2_f64, 0.8_f64, -0.4_f64];
3630 let w = 1.4_f64;
3631 let block = identity_block(w, has_entry, event);
3632 let prog = crate::survival::CauseSpecificRowProgram::new(
3633 [beta[0], beta[1], beta[2]],
3634 w,
3635 has_entry,
3636 event,
3637 );
3638 let label = format!("entry={has_entry} event={event}");
3639
3640 let (ll, grad, hess) =
3642 evaluate_cause_specific_block(&block, &beta).expect("evaluate block");
3643 let (jet_v, jet_g, jet_h) = program_row_kernel(&prog, 0).expect("jet kernel");
3644 close(jet_v, -ll, JET_TOL, &format!("{label} value"));
3645 for a in 0..3 {
3646 close(jet_g[a], -grad[a], JET_TOL, &format!("{label} grad[{a}]"));
3647 for b in 0..3 {
3648 close(
3649 jet_h[a][b],
3650 hess[[a, b]],
3651 JET_TOL,
3652 &format!("{label} H[{a}][{b}]"),
3653 );
3654 }
3655 }
3656
3657 let dh = cause_specific_hessian_directional_derivative(&block, &beta, &d_beta)
3659 .expect("live third");
3660 let dir = [d_beta[0], d_beta[1], d_beta[2]];
3661 let jet_t3 = program_third_contracted(&prog, 0, &dir).expect("jet third");
3662 for a in 0..3 {
3663 for b in 0..3 {
3664 close(
3665 jet_t3[a][b],
3666 dh[[a, b]],
3667 JET_TOL,
3668 &format!("{label} third[{a}][{b}]"),
3669 );
3670 }
3671 }
3672
3673 let d2h = cause_specific_hessian_second_directional_derivative(
3675 &block, &beta, &d_beta, &v_beta,
3676 )
3677 .expect("live fourth");
3678 let uu = [d_beta[0], d_beta[1], d_beta[2]];
3679 let vv = [v_beta[0], v_beta[1], v_beta[2]];
3680 let jet_t4 = program_fourth_contracted(&prog, 0, &uu, &vv).expect("jet fourth");
3681 for a in 0..3 {
3682 for b in 0..3 {
3683 close(
3684 jet_t4[a][b],
3685 d2h[[a, b]],
3686 JET_TOL,
3687 &format!("{label} fourth[{a}][{b}]"),
3688 );
3689 }
3690 }
3691
3692 let h_fd = 1e-5;
3695 let bp = &beta + &(&d_beta * h_fd);
3696 let bm = &beta - &(&d_beta * h_fd);
3697 let (_, _, hp) = evaluate_cause_specific_block(&block, &bp).expect("evaluate +");
3698 let (_, _, hm) = evaluate_cause_specific_block(&block, &bm).expect("evaluate -");
3699 for a in 0..3 {
3700 for b in 0..3 {
3701 let fd = (hp[[a, b]] - hm[[a, b]]) / (2.0 * h_fd);
3702 close(dh[[a, b]], fd, 1e-5, &format!("{label} FD third[{a}][{b}]"));
3703 }
3704 }
3705 let dhp = cause_specific_hessian_directional_derivative(
3707 &block,
3708 &bp_along(&beta, &v_beta, h_fd),
3709 &d_beta,
3710 )
3711 .expect("live third +");
3712 let dhm = cause_specific_hessian_directional_derivative(
3713 &block,
3714 &bm_along(&beta, &v_beta, h_fd),
3715 &d_beta,
3716 )
3717 .expect("live third -");
3718 for a in 0..3 {
3719 for b in 0..3 {
3720 let fd = (dhp[[a, b]] - dhm[[a, b]]) / (2.0 * h_fd);
3721 close(
3722 d2h[[a, b]],
3723 fd,
3724 1e-5,
3725 &format!("{label} FD fourth[{a}][{b}]"),
3726 );
3727 }
3728 }
3729 }
3730
3731 fn bp_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
3732 beta + &(v * h)
3733 }
3734 fn bm_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
3735 beta - &(v * h)
3736 }
3737
3738 #[test]
3744 fn cause_specific_live_tower_matches_jet_and_fd() {
3745 for &has_entry in &[false, true] {
3746 for &event in &[false, true] {
3747 run_corner(has_entry, event);
3748 }
3749 }
3750 }
3751
3752 #[test]
3773 fn release_measure_cause_specific_specialized_vs_generic_tower_932() {
3774 use std::time::Instant;
3775
3776 const ROWS: usize = 512;
3777 let mut rows: Vec<([f64; 3], f64, bool, bool)> = Vec::with_capacity(ROWS);
3778 for idx in 0..ROWS {
3779 let f = idx as f64;
3780 let eta_exit = 1.6 * (f * 0.17 + 0.3).sin() - 0.4 * (f * 0.09).cos();
3781 let eta_entry = 1.1 * (f * 0.13 + 0.7).cos() + 0.35 * (f * 0.05).sin();
3782 let derivative = 0.5 + 0.45 * (f * 0.31 + 0.2).sin().abs();
3784 let weight = 0.6 + 0.4 * (f * 0.07 + 1.0).sin().abs();
3785 let entry_active = idx % 2 == 0;
3786 let event = (idx / 2) % 2 == 0;
3787 rows.push((
3788 [eta_exit, eta_entry, derivative],
3789 weight,
3790 entry_active,
3791 event,
3792 ));
3793 }
3794 let programs: Vec<crate::survival::CauseSpecificRowProgram> = rows
3795 .iter()
3796 .map(|&(primary, weight, entry_active, event)| {
3797 crate::survival::CauseSpecificRowProgram::new(
3798 primary,
3799 weight,
3800 entry_active,
3801 event,
3802 )
3803 })
3804 .collect();
3805 let dir_u: Vec<[f64; 3]> = (0..ROWS)
3808 .map(|idx| {
3809 let f = idx as f64;
3810 [
3811 0.7 * (f * 0.23 + 0.4).cos() - 0.2 * (f * 0.03).sin(),
3812 -0.6 * (f * 0.29 + 0.1).sin() + 0.25 * (f * 0.15).cos(),
3813 0.5 * (f * 0.19 + 0.6).cos() - 0.3 * (f * 0.08).sin(),
3814 ]
3815 })
3816 .collect();
3817 let dir_v: Vec<[f64; 3]> = (0..ROWS)
3818 .map(|idx| {
3819 let f = idx as f64;
3820 [
3821 -0.5 * (f * 0.21 + 0.9).sin() + 0.3 * (f * 0.06).cos(),
3822 0.8 * (f * 0.27 + 0.5).cos() - 0.15 * (f * 0.04).sin(),
3823 0.4 * (f * 0.13 + 0.3).sin() - 0.2 * (f * 0.11).cos(),
3824 ]
3825 })
3826 .collect();
3827
3828 for (idx, (row, program)) in rows.iter().zip(programs.iter()).enumerate() {
3831 let (primary, weight, entry_active, event) = *row;
3832 let atom = cause_specific_row_order2(
3833 primary[0],
3834 primary[1],
3835 primary[2],
3836 weight,
3837 f64::from(entry_active),
3838 f64::from(event),
3839 );
3840 let (tower_value, tower_gradient, tower_hessian) =
3841 program_row_kernel(program, 0).expect("tower warm kernel");
3842 close(
3843 atom.value(),
3844 tower_value,
3845 JET_TOL,
3846 "release-measure value parity",
3847 );
3848 let production_gradient = atom.gradient();
3849 for a in 0..3 {
3850 close(
3851 production_gradient[a],
3852 tower_gradient[a],
3853 JET_TOL,
3854 "release-measure gradient parity",
3855 );
3856 for b in 0..3 {
3857 close(
3858 atom.hessian_at(a, b),
3859 tower_hessian[a][b],
3860 JET_TOL,
3861 "release-measure hessian parity",
3862 );
3863 }
3864 }
3865 let production_third = cause_specific_row_third_contracted(
3866 primary[0],
3867 primary[1],
3868 primary[2],
3869 weight,
3870 f64::from(entry_active),
3871 f64::from(event),
3872 &dir_u[idx],
3873 );
3874 let tower_third =
3875 program_third_contracted(program, 0, &dir_u[idx]).expect("tower warm third");
3876 let production_fourth = cause_specific_row_fourth_contracted(
3877 primary[0],
3878 primary[1],
3879 primary[2],
3880 weight,
3881 f64::from(entry_active),
3882 f64::from(event),
3883 &dir_u[idx],
3884 &dir_v[idx],
3885 );
3886 let tower_fourth = program_fourth_contracted(program, 0, &dir_u[idx], &dir_v[idx])
3887 .expect("tower warm fourth");
3888 for a in 0..3 {
3889 for b in 0..3 {
3890 close(
3891 production_third[a][b],
3892 tower_third[a][b],
3893 JET_TOL,
3894 "release-measure third parity",
3895 );
3896 close(
3897 production_fourth[a][b],
3898 tower_fourth[a][b],
3899 JET_TOL,
3900 "release-measure fourth parity",
3901 );
3902 }
3903 }
3904 }
3905
3906 let best_secs = |sweep: &mut dyn FnMut() -> f64| -> f64 {
3907 let mut best = f64::INFINITY;
3908 for _ in 0..5 {
3909 let started = Instant::now();
3910 let checksum = sweep();
3911 assert!(
3912 checksum.is_finite(),
3913 "cause-specific release-measure checksum must stay finite"
3914 );
3915 best = best.min(started.elapsed().as_secs_f64());
3916 }
3917 best
3918 };
3919
3920 let mut production_sweep = || {
3921 let mut checksum = 0.0_f64;
3922 for &(primary, weight, entry_active, event) in &rows {
3923 let atom = cause_specific_row_order2(
3924 primary[0],
3925 primary[1],
3926 primary[2],
3927 weight,
3928 f64::from(entry_active),
3929 f64::from(event),
3930 );
3931 checksum += atom.value() + atom.gradient()[0] + atom.hessian_at(0, 0);
3932 }
3933 checksum
3934 };
3935 let production_secs = best_secs(&mut production_sweep);
3936
3937 let mut tower_sweep = || {
3938 let mut checksum = 0.0_f64;
3939 for program in &programs {
3940 let (value, gradient, hessian) =
3941 program_row_kernel(program, 0).expect("tower kernel");
3942 checksum += value + gradient[0] + hessian[0][0];
3943 }
3944 checksum
3945 };
3946 let tower_secs = best_secs(&mut tower_sweep);
3947
3948 let mut production_third_sweep = || {
3949 let mut checksum = 0.0_f64;
3950 for (idx, &(primary, weight, entry_active, event)) in rows.iter().enumerate() {
3951 let third = cause_specific_row_third_contracted(
3952 primary[0],
3953 primary[1],
3954 primary[2],
3955 weight,
3956 f64::from(entry_active),
3957 f64::from(event),
3958 &dir_u[idx],
3959 );
3960 checksum += third[0][0] + third[0][1] + third[1][1];
3961 }
3962 checksum
3963 };
3964 let production_third_secs = best_secs(&mut production_third_sweep);
3965 let mut tower_third_sweep = || {
3966 let mut checksum = 0.0_f64;
3967 for (idx, program) in programs.iter().enumerate() {
3968 let third = program_third_contracted(program, 0, &dir_u[idx])
3969 .expect("tower third kernel");
3970 checksum += third[0][0] + third[0][1] + third[1][1];
3971 }
3972 checksum
3973 };
3974 let tower_third_secs = best_secs(&mut tower_third_sweep);
3975
3976 let mut production_fourth_sweep = || {
3977 let mut checksum = 0.0_f64;
3978 for (idx, &(primary, weight, entry_active, event)) in rows.iter().enumerate() {
3979 let fourth = cause_specific_row_fourth_contracted(
3980 primary[0],
3981 primary[1],
3982 primary[2],
3983 weight,
3984 f64::from(entry_active),
3985 f64::from(event),
3986 &dir_u[idx],
3987 &dir_v[idx],
3988 );
3989 checksum += fourth[0][0] + fourth[0][1] + fourth[1][1];
3990 }
3991 checksum
3992 };
3993 let production_fourth_secs = best_secs(&mut production_fourth_sweep);
3994 let mut tower_fourth_sweep = || {
3995 let mut checksum = 0.0_f64;
3996 for (idx, program) in programs.iter().enumerate() {
3997 let fourth = program_fourth_contracted(program, 0, &dir_u[idx], &dir_v[idx])
3998 .expect("tower fourth kernel");
3999 checksum += fourth[0][0] + fourth[0][1] + fourth[1][1];
4000 }
4001 checksum
4002 };
4003 let tower_fourth_secs = best_secs(&mut tower_fourth_sweep);
4004
4005 for (channel, production_secs, tower_secs) in [
4006 ("order2", production_secs, tower_secs),
4007 ("third", production_third_secs, tower_third_secs),
4008 ("fourth", production_fourth_secs, tower_fourth_secs),
4009 ] {
4010 let production_ns = production_secs * 1e9 / ROWS as f64;
4011 let tower_ns = tower_secs * 1e9 / ROWS as f64;
4012 eprintln!(
4013 "CAUSE-SPECIFIC-RELEASE-932 channel={channel} rows={ROWS} \
4014 production_ns={production_ns:.3} generic_tower_ns={tower_ns:.3} \
4015 hand_over_production={:.6}",
4016 tower_ns / production_ns,
4017 );
4018 }
4019 }
4020 }
4021
4022 #[test]
4023 fn competing_risks_cif_constant_hazard_matches_closed_form() {
4024 let times = array![0.0, 2.0, 5.0, 10.0];
4025 let disease_rates = [0.12, 0.06];
4026 let death_rates = [0.05, 0.02];
4027 let cumulative = Array3::from_shape_fn((2, 2, times.len()), |(endpoint, row, time_idx)| {
4028 let rate = if endpoint == 0 {
4029 disease_rates[row]
4030 } else {
4031 death_rates[row]
4032 };
4033 rate * times[time_idx]
4034 });
4035
4036 let result =
4037 assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
4038
4039 for row in 0..2 {
4040 let total_rate = disease_rates[row] + death_rates[row];
4041 for time_idx in 0..times.len() {
4042 let failure = 1.0 - (-total_rate * times[time_idx]).exp();
4043 let expected_disease = disease_rates[row] / total_rate * failure;
4044 let expected_death = death_rates[row] / total_rate * failure;
4045 assert!((result.cif[0][[row, time_idx]] - expected_disease).abs() < 1e-12);
4046 assert!((result.cif[1][[row, time_idx]] - expected_death).abs() < 1e-12);
4047 assert!(
4048 (result.cif[0][[row, time_idx]]
4049 + result.cif[1][[row, time_idx]]
4050 + result.overall_survival[[row, time_idx]]
4051 - 1.0)
4052 .abs()
4053 < 1e-12
4054 );
4055 }
4056 }
4057 }
4058
4059 #[test]
4060 fn competing_risks_cif_rejects_nonmonotone_hazards() {
4061 let times = array![0.0, 1.0, 2.0];
4062 let cumulative = Array3::from_shape_vec((1, 1, 3), vec![0.0, 0.2, 0.1]).expect("shape");
4063 let err = assemble_competing_risks_cif(times.view(), cumulative.view())
4064 .expect_err("nonmonotone cumulative hazard should be rejected");
4065 assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
4066 }
4067
4068 #[test]
4069 fn competing_risks_cif_plateaus_and_three_causes_conserve_probability() {
4070 let times = array![0.0, 1.0, 3.0, 7.0, 12.0];
4071 let cumulative = Array3::from_shape_vec(
4072 (3, 2, 5),
4073 vec![
4074 0.0, 0.2, 0.2, 0.5, 1.1, 0.0, 0.0, 0.4, 0.4, 0.9, 0.0, 0.1, 0.3, 0.3, 0.7, 0.0, 0.2, 0.2, 0.8, 0.8, 0.0, 0.0, 0.2, 0.6, 0.6, 0.0, 0.1, 0.5, 0.5, 1.5,
4078 ],
4079 )
4080 .expect("shape");
4081
4082 let result =
4083 assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
4084
4085 for row in 0..2 {
4086 for time_idx in 0..times.len() {
4087 let total_cif = result.cif[0][[row, time_idx]]
4088 + result.cif[1][[row, time_idx]]
4089 + result.cif[2][[row, time_idx]];
4090 assert!(
4091 (total_cif + result.overall_survival[[row, time_idx]] - 1.0).abs() < 1e-12,
4092 "probability mass mismatch at row={row}, time_idx={time_idx}"
4093 );
4094 assert!((0.0..=1.0).contains(&result.overall_survival[[row, time_idx]]));
4095 for cause in 0..3 {
4096 assert!((0.0..=1.0).contains(&result.cif[cause][[row, time_idx]]));
4097 if time_idx > 0 {
4098 assert!(
4099 result.cif[cause][[row, time_idx]] + 1e-12
4100 >= result.cif[cause][[row, time_idx - 1]],
4101 "CIF decreased for cause={cause}, row={row}, time_idx={time_idx}"
4102 );
4103 }
4104 }
4105 }
4106 }
4107
4108 assert_eq!(result.cif[0][[0, 1]], result.cif[0][[0, 2]]);
4111 assert_eq!(result.cif[0][[1, 2]], result.cif[0][[1, 3]]);
4114 assert_eq!(result.cif[2][[1, 2]], result.cif[2][[1, 3]]);
4115 }
4116
4117 #[test]
4118 fn competing_risks_cif_rejects_bad_time_grids_and_nonfinite_hazards() {
4119 let cumulative = Array3::zeros((2, 1, 2));
4120
4121 for times in [array![0.0, 0.0], array![1.0, 0.5], array![-1.0, 1.0]] {
4122 let err = assemble_competing_risks_cif(times.view(), cumulative.view())
4123 .expect_err("bad time grid should be rejected");
4124 assert!(matches!(err, SurvivalError::InvalidTimeGrid));
4125 }
4126
4127 let times = array![0.0, 1.0];
4128 let nonfinite = Array3::from_shape_vec((1, 1, 2), vec![0.0, f64::NAN]).expect("shape");
4129 let err = assemble_competing_risks_cif(times.view(), nonfinite.view())
4130 .expect_err("nonfinite hazard should be rejected");
4131 assert!(matches!(err, SurvivalError::NonFiniteInput));
4132 }
4133
4134 #[test]
4135 fn competing_risks_cif_extreme_hazards_remain_bounded() {
4136 let times = array![0.0, 1.0, 2.0];
4137 let cumulative =
4138 Array3::from_shape_vec((2, 1, 3), vec![0.0, 500.0, 1000.0, 0.0, 250.0, 1000.0])
4139 .expect("shape");
4140
4141 let result =
4142 assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
4143
4144 for value in result
4145 .cif
4146 .iter()
4147 .flat_map(|m| m.iter())
4148 .chain(result.overall_survival.iter())
4149 {
4150 assert!(value.is_finite());
4151 assert!((0.0..=1.0).contains(value));
4152 }
4153 assert!((result.cif[0][[0, 2]] + result.cif[1][[0, 2]] - 1.0).abs() < 1e-12);
4154 assert_eq!(result.overall_survival[[0, 2]], 0.0);
4155 }
4156
4157 fn toy_penalties() -> PenaltyBlocks {
4158 let s = array![[2.0, 0.5], [0.5, 3.0]];
4159 PenaltyBlocks::new(vec![PenaltyBlock {
4160 matrix: s,
4161 lambda: 1.7,
4162 range: 1..3,
4163 nullspace_dim: 0,
4164 }])
4165 }
4166
4167 fn survival_inputs<'a>(
4168 age_entry: &'a Array1<f64>,
4169 age_exit: &'a Array1<f64>,
4170 event_target: &'a Array1<u8>,
4171 event_competing: &'a Array1<u8>,
4172 sampleweight: &'a Array1<f64>,
4173 x_entry: &'a Array2<f64>,
4174 x_exit: &'a Array2<f64>,
4175 x_derivative: &'a Array2<f64>,
4176 ) -> SurvivalEngineInputs<'a> {
4177 SurvivalEngineInputs {
4178 age_entry: age_entry.view(),
4179 age_exit: age_exit.view(),
4180 event_target: event_target.view(),
4181 event_competing: event_competing.view(),
4182 sampleweight: sampleweight.view(),
4183 x_entry: x_entry.view(),
4184 x_exit: x_exit.view(),
4185 x_derivative: x_derivative.view(),
4186 monotonicity_constraint_rows: None,
4187 monotonicity_constraint_offsets: None,
4188 }
4189 }
4190
4191 fn survival_model(
4192 inputs: SurvivalEngineInputs<'_>,
4193 penalties: PenaltyBlocks,
4194 monotonicity: SurvivalMonotonicityPenalty,
4195 spec: SurvivalSpec,
4196 ) -> Result<WorkingModelSurvival, SurvivalError> {
4197 WorkingModelSurvival::from_engine_inputs(inputs, penalties, monotonicity, spec)
4198 }
4199
4200 fn survival_model_with_offsets(
4201 inputs: SurvivalEngineInputs<'_>,
4202 offsets: Option<SurvivalBaselineOffsets<'_>>,
4203 penalties: PenaltyBlocks,
4204 monotonicity: SurvivalMonotonicityPenalty,
4205 spec: SurvivalSpec,
4206 ) -> Result<WorkingModelSurvival, SurvivalError> {
4207 WorkingModelSurvival::from_engine_inputswith_offsets(
4208 inputs,
4209 offsets,
4210 penalties,
4211 monotonicity,
4212 spec,
4213 )
4214 }
4215
4216 #[test]
4217 fn penaltyhessian_matchesgradient_jacobian() {
4218 let penalties = toy_penalties();
4219 let beta = array![10.0, -0.3, 1.2, 7.0];
4220
4221 let grad = penalties.gradient(&beta);
4222 let h = penalties.hessian(beta.len());
4223 let b_block = beta.slice(s![1..3]).to_owned();
4224 let expected = 1.7 * array![[2.0, 0.5], [0.5, 3.0]].dot(&b_block);
4225
4226 assert!((grad[1] - expected[0]).abs() < 1e-12);
4227 assert!((grad[2] - expected[1]).abs() < 1e-12);
4228 assert!((h[[1, 1]] - 1.7 * 2.0).abs() < 1e-12);
4229 assert!((h[[1, 2]] - 1.7 * 0.5).abs() < 1e-12);
4230 assert!((h[[2, 1]] - 1.7 * 0.5).abs() < 1e-12);
4231 assert!((h[[2, 2]] - 1.7 * 3.0).abs() < 1e-12);
4232 }
4233
4234 #[test]
4235 fn penaltygradient_matches_deviance_finite_difference() {
4236 let penalties = toy_penalties();
4237 let beta = array![10.0, -0.3, 1.2, 7.0];
4238 let grad = penalties.gradient(&beta);
4239 let eps = 1e-7;
4240
4241 for idx in 0..beta.len() {
4242 let mut plus = beta.clone();
4243 let mut minus = beta.clone();
4244 plus[idx] += eps;
4245 minus[idx] -= eps;
4246 let fd = (penalties.deviance(&plus) - penalties.deviance(&minus)) / (2.0 * eps);
4247 assert_eq!(
4248 grad[idx].signum(),
4249 fd.signum(),
4250 "gradient/deviance sign mismatch at idx={idx}: grad={} fd={fd}",
4251 grad[idx]
4252 );
4253 assert!(
4254 (grad[idx] - fd).abs() < 1e-6,
4255 "gradient/deviance mismatch at idx={idx}: grad={} fd={fd}",
4256 grad[idx]
4257 );
4258 }
4259 }
4260
4261 #[test]
4262 fn zero_offsets_match_default_survival_state() {
4263 let age_entry = array![1.0_f64, 2.0_f64];
4264 let age_exit = array![2.0_f64, 3.5_f64];
4265 let event_target = array![1u8, 0u8];
4266 let event_competing = array![0u8, 0u8];
4267 let sampleweight = array![1.0, 1.0];
4268 let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
4269 let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
4270 let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
4271 let penalties = PenaltyBlocks::new(Vec::new());
4272 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4273 let beta = array![-1.0, 0.8];
4274
4275 let base = survival_model(
4276 survival_inputs(
4277 &age_entry,
4278 &age_exit,
4279 &event_target,
4280 &event_competing,
4281 &sampleweight,
4282 &x_entry,
4283 &x_exit,
4284 &x_derivative,
4285 ),
4286 penalties.clone(),
4287 mono,
4288 SurvivalSpec::Net,
4289 )
4290 .expect("construct base survival model");
4291
4292 let zero_offsets = survival_model_with_offsets(
4293 survival_inputs(
4294 &age_entry,
4295 &age_exit,
4296 &event_target,
4297 &event_competing,
4298 &sampleweight,
4299 &x_entry,
4300 &x_exit,
4301 &x_derivative,
4302 ),
4303 Some(SurvivalBaselineOffsets {
4304 eta_entry: array![0.0, 0.0].view(),
4305 eta_exit: array![0.0, 0.0].view(),
4306 derivative_exit: array![0.0, 0.0].view(),
4307 }),
4308 penalties,
4309 mono,
4310 SurvivalSpec::Net,
4311 )
4312 .expect("construct offset survival model");
4313
4314 let state_base = base.update_state(&beta).expect("base state");
4315 let statezero = zero_offsets.update_state(&beta).expect("zero-offset state");
4316 assert!((state_base.deviance - statezero.deviance).abs() < 1e-12);
4317 assert!(
4318 state_base
4319 .gradient
4320 .iter()
4321 .zip(statezero.gradient.iter())
4322 .all(|(a, b)| (a - b).abs() < 1e-12)
4323 );
4324 }
4325
4326 #[test]
4327 fn competing_risk_cause_labels_collapse_to_pooled_baseline_indicator() {
4328 let age_entry = array![0.0_f64, 0.0, 0.0, 0.0];
4342 let age_exit = array![1.2_f64, 0.8, 2.1, 1.5];
4343 let cause_labels = array![0u8, 1u8, 2u8, 0u8];
4345 let event_competing = Array1::<u8>::zeros(cause_labels.len());
4346 let sampleweight = array![1.0_f64, 1.0, 1.0, 1.0];
4347 let x_entry = array![
4348 [1.0, age_entry[0].max(1e-8).ln()],
4349 [1.0, age_entry[1].max(1e-8).ln()],
4350 [1.0, age_entry[2].max(1e-8).ln()],
4351 [1.0, age_entry[3].max(1e-8).ln()],
4352 ];
4353 let x_exit = array![
4354 [1.0, age_exit[0].ln()],
4355 [1.0, age_exit[1].ln()],
4356 [1.0, age_exit[2].ln()],
4357 [1.0, age_exit[3].ln()],
4358 ];
4359 let x_derivative = array![
4360 [0.0, 1.0 / age_exit[0]],
4361 [0.0, 1.0 / age_exit[1]],
4362 [0.0, 1.0 / age_exit[2]],
4363 [0.0, 1.0 / age_exit[3]],
4364 ];
4365 let penalties = PenaltyBlocks::new(Vec::new());
4366 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4367
4368 let raw = survival_model(
4373 survival_inputs(
4374 &age_entry,
4375 &age_exit,
4376 &cause_labels,
4377 &event_competing,
4378 &sampleweight,
4379 &x_entry,
4380 &x_exit,
4381 &x_derivative,
4382 ),
4383 penalties.clone(),
4384 mono,
4385 SurvivalSpec::Net,
4386 );
4387 assert!(
4388 matches!(raw, Err(SurvivalError::EventCodeInvalid { .. })),
4389 "raw competing-risks cause labels must be rejected as EventCodeInvalid (not NonFiniteInput), got {raw:?}"
4390 );
4391
4392 let any_event = pooled_any_event_indicator(cause_labels.view());
4395 assert_eq!(any_event, array![0u8, 1u8, 1u8, 0u8]);
4396 assert_eq!(
4398 cause_specific_event_indicator(cause_labels.view(), 1),
4399 array![0u8, 1u8, 0u8, 0u8]
4400 );
4401 assert_eq!(
4402 cause_specific_event_indicator(cause_labels.view(), 2),
4403 array![0u8, 0u8, 1u8, 0u8]
4404 );
4405 let model = survival_model(
4406 survival_inputs(
4407 &age_entry,
4408 &age_exit,
4409 &any_event,
4410 &event_competing,
4411 &sampleweight,
4412 &x_entry,
4413 &x_exit,
4414 &x_derivative,
4415 ),
4416 penalties,
4417 mono,
4418 SurvivalSpec::Net,
4419 )
4420 .expect("pooled any-event baseline model must construct from competing-risks data");
4421
4422 let beta = array![-1.0_f64, 0.8];
4425 let state = model.update_state(&beta).expect("pooled baseline state");
4426 assert!(
4427 state.deviance.is_finite(),
4428 "pooled baseline deviance must be finite, got {}",
4429 state.deviance
4430 );
4431 assert!(
4432 state.gradient.iter().all(|g| g.is_finite()),
4433 "pooled baseline gradient must be finite"
4434 );
4435 }
4436
4437 #[test]
4438 fn offset_channel_residuals_match_central_fd_of_nll() {
4439 let age_entry = array![0.5_f64, 0.0, 0.3];
4444 let age_exit = array![1.4_f64, 1.0, 2.0];
4445 let event_target = array![1u8, 1u8, 0u8];
4446 let event_competing = array![0u8, 0u8, 0u8];
4447 let sampleweight = array![1.0_f64, 2.5, 0.7];
4448 let x_entry = array![
4449 [1.0, age_entry[0].ln()],
4450 [1.0, age_entry[1].max(1e-8).ln()],
4451 [1.0, age_entry[2].ln()]
4452 ];
4453 let x_exit = array![
4454 [1.0, age_exit[0].ln()],
4455 [1.0, age_exit[1].ln()],
4456 [1.0, age_exit[2].ln()]
4457 ];
4458 let x_derivative = array![
4459 [0.0, 1.0 / age_exit[0]],
4460 [0.0, 1.0 / age_exit[1]],
4461 [0.0, 1.0 / age_exit[2]]
4462 ];
4463 let o_entry = array![0.2_f64, 0.0, 0.1];
4466 let o_exit = array![0.4_f64, 0.5, 0.7];
4467 let o_deriv = array![0.3_f64, 0.8, 0.5];
4468 let penalties = PenaltyBlocks::new(Vec::new());
4469 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4470 let beta = array![-0.7_f64, 0.6];
4471
4472 let build = |o_e: &Array1<f64>, o_x: &Array1<f64>, o_d: &Array1<f64>| {
4473 survival_model_with_offsets(
4474 survival_inputs(
4475 &age_entry,
4476 &age_exit,
4477 &event_target,
4478 &event_competing,
4479 &sampleweight,
4480 &x_entry,
4481 &x_exit,
4482 &x_derivative,
4483 ),
4484 Some(SurvivalBaselineOffsets {
4485 eta_entry: o_e.view(),
4486 eta_exit: o_x.view(),
4487 derivative_exit: o_d.view(),
4488 }),
4489 penalties.clone(),
4490 mono,
4491 SurvivalSpec::Net,
4492 )
4493 .expect("model build")
4494 };
4495
4496 let base = build(&o_entry, &o_exit, &o_deriv);
4497 let resid = base
4498 .offset_channel_residuals(&beta)
4499 .expect("offset residuals");
4500 assert_eq!(resid.exit.len(), 3);
4501 assert_eq!(resid.entry.len(), 3);
4502 assert_eq!(resid.derivative.len(), 3);
4503
4504 let nll = |m: &WorkingModelSurvival| 0.5 * m.update_state(&beta).expect("state").deviance;
4507 let h = 1e-6;
4508
4509 assert_eq!(resid.entry[1], 0.0);
4513 assert_eq!(resid.derivative[2], 0.0);
4514
4515 for i in 0..3 {
4516 {
4518 let mut op = o_exit.clone();
4519 let mut om = o_exit.clone();
4520 op[i] += h;
4521 om[i] -= h;
4522 let fd = (nll(&build(&o_entry, &op, &o_deriv))
4523 - nll(&build(&o_entry, &om, &o_deriv)))
4524 / (2.0 * h);
4525 assert!(
4526 (resid.exit[i] - fd).abs() < 1e-6,
4527 "∂NLL/∂o_X[{i}]: analytic={:.6e} fd={:.6e}",
4528 resid.exit[i],
4529 fd
4530 );
4531 }
4532 {
4536 let mut op = o_entry.clone();
4537 let mut om = o_entry.clone();
4538 op[i] += h;
4539 om[i] -= h;
4540 let fd = (nll(&build(&op, &o_exit, &o_deriv))
4541 - nll(&build(&om, &o_exit, &o_deriv)))
4542 / (2.0 * h);
4543 assert!(
4544 (resid.entry[i] - fd).abs() < 1e-6,
4545 "∂NLL/∂o_E[{i}]: analytic={:.6e} fd={:.6e}",
4546 resid.entry[i],
4547 fd
4548 );
4549 }
4550 {
4552 let mut op = o_deriv.clone();
4553 let mut om = o_deriv.clone();
4554 op[i] += h;
4555 om[i] -= h;
4556 let fd = (nll(&build(&o_entry, &o_exit, &op))
4557 - nll(&build(&o_entry, &o_exit, &om)))
4558 / (2.0 * h);
4559 assert!(
4560 (resid.derivative[i] - fd).abs() < 1e-6,
4561 "∂NLL/∂o_D[{i}]: analytic={:.6e} fd={:.6e}",
4562 resid.derivative[i],
4563 fd
4564 );
4565 }
4566 }
4567 }
4568
4569 #[test]
4570 fn offset_channel_residuals_respect_zero_sampleweight() {
4571 let age_entry = array![1.0_f64, 2.0];
4572 let age_exit = array![2.0_f64, 3.5];
4573 let event_target = array![1u8, 1u8];
4574 let event_competing = array![0u8, 0u8];
4575 let sampleweight = array![0.0_f64, 1.2]; let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
4577 let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
4578 let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
4579 let penalties = PenaltyBlocks::new(Vec::new());
4580 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4581 let beta = array![-1.0_f64, 0.8];
4582
4583 let model = survival_model_with_offsets(
4584 survival_inputs(
4585 &age_entry,
4586 &age_exit,
4587 &event_target,
4588 &event_competing,
4589 &sampleweight,
4590 &x_entry,
4591 &x_exit,
4592 &x_derivative,
4593 ),
4594 Some(SurvivalBaselineOffsets {
4595 eta_entry: array![0.0_f64, 0.1].view(),
4596 eta_exit: array![0.0_f64, 0.2].view(),
4597 derivative_exit: array![0.0_f64, 0.1].view(),
4598 }),
4599 penalties,
4600 mono,
4601 SurvivalSpec::Net,
4602 )
4603 .expect("model");
4604 let r = model.offset_channel_residuals(&beta).expect("resid");
4605 assert_eq!(r.exit[0], 0.0);
4607 assert_eq!(r.entry[0], 0.0);
4608 assert_eq!(r.derivative[0], 0.0);
4609 assert!(r.exit[1] != 0.0);
4611 }
4612
4613 #[test]
4614 fn offset_channel_residuals_reject_beta_dim_mismatch() {
4615 let age_entry = array![1.0_f64];
4616 let age_exit = array![2.0_f64];
4617 let event_target = array![1u8];
4618 let event_competing = array![0u8];
4619 let sampleweight = array![1.0_f64];
4620 let x_entry = array![[1.0, 0.0]];
4621 let x_exit = array![[1.0, 0.7]];
4622 let x_derivative = array![[0.0, 0.5]];
4623 let penalties = PenaltyBlocks::new(Vec::new());
4624 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4625 let model = survival_model(
4626 survival_inputs(
4627 &age_entry,
4628 &age_exit,
4629 &event_target,
4630 &event_competing,
4631 &sampleweight,
4632 &x_entry,
4633 &x_exit,
4634 &x_derivative,
4635 ),
4636 penalties,
4637 mono,
4638 SurvivalSpec::Net,
4639 )
4640 .expect("model");
4641 let bad_beta = array![0.0_f64]; let err = model
4643 .offset_channel_residuals(&bad_beta)
4644 .expect_err("mismatch must error");
4645 match err {
4646 EstimationError::InvalidInput(msg) => {
4647 assert!(msg.contains("beta dimension mismatch"), "msg={msg}")
4648 }
4649 other => panic!("expected InvalidInput, got {other:?}"),
4650 }
4651 }
4652
4653 #[test]
4654 fn crudespec_is_rejected_by_one_hazard_engine() {
4655 let age_entry = array![1.0_f64];
4656 let age_exit = array![2.0_f64];
4657 let event_target = array![0u8];
4658 let event_competing = array![1u8];
4659 let sampleweight = array![1.0];
4660 let x_entry = array![[0.1]];
4661 let x_exit = array![[0.4]];
4662 let x_derivative = array![[1.0]];
4663 let penalties = PenaltyBlocks::new(Vec::new());
4664 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4665
4666 let err = survival_model(
4667 survival_inputs(
4668 &age_entry,
4669 &age_exit,
4670 &event_target,
4671 &event_competing,
4672 &sampleweight,
4673 &x_entry,
4674 &x_exit,
4675 &x_derivative,
4676 ),
4677 penalties,
4678 mono,
4679 SurvivalSpec::Crude,
4680 )
4681 .expect_err("crude fitting should be rejected by the one-hazard engine");
4682 assert!(matches!(err, SurvivalError::UnsupportedSpec("crude")));
4683 }
4684
4685 #[test]
4686 fn nonstructural_models_require_explicit_monotonicity_collocation() {
4687 let age_entry = array![1.0_f64, 1.5_f64];
4688 let age_exit = array![2.0_f64, 2.5_f64];
4689 let event_target = array![0u8, 0u8];
4690 let event_competing = array![0u8, 1u8];
4691 let sampleweight = array![1.0, 1.0];
4692 let x_entry = array![[0.2], [0.1]];
4693 let x_exit = array![[0.3], [0.2]];
4694 let x_derivative = array![[1.0], [1.0]];
4695
4696 let model = survival_model(
4697 survival_inputs(
4698 &age_entry,
4699 &age_exit,
4700 &event_target,
4701 &event_competing,
4702 &sampleweight,
4703 &x_entry,
4704 &x_exit,
4705 &x_derivative,
4706 ),
4707 PenaltyBlocks::new(Vec::new()),
4708 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4709 SurvivalSpec::Net,
4710 )
4711 .expect("construct censored survival model");
4712
4713 assert!(
4714 model.monotonicity_linear_constraints().is_none(),
4715 "non-structural survival models must not fabricate rowwise monotonicity constraints"
4716 );
4717 }
4718
4719 #[test]
4720 fn decreasing_interval_is_rejectedwithout_target_events() {
4721 let age_entry = array![1.0_f64];
4722 let age_exit = array![2.0_f64];
4723 let event_target = array![0u8];
4724 let event_competing = array![0u8];
4725 let sampleweight = array![1.0];
4726 let x_entry = array![[0.5]];
4727 let x_exit = array![[0.0]];
4728 let x_derivative = array![[1.0]];
4729
4730 let model = survival_model(
4731 survival_inputs(
4732 &age_entry,
4733 &age_exit,
4734 &event_target,
4735 &event_competing,
4736 &sampleweight,
4737 &x_entry,
4738 &x_exit,
4739 &x_derivative,
4740 ),
4741 PenaltyBlocks::new(Vec::new()),
4742 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4743 SurvivalSpec::Net,
4744 )
4745 .expect("construct censored survival model");
4746
4747 let err = model
4748 .update_state(&array![1.0])
4749 .expect_err("decreasing cumulative hazard increment should be rejected");
4750 assert!(
4751 err.to_string().contains("cumulative hazard decreased"),
4752 "unexpected error: {err}"
4753 );
4754 }
4755
4756 fn smooth_crude_risk(beta_d: f64, beta_m: f64) -> CrudeRiskResult {
4757 calculate_crude_risk_quadrature(
4758 0.0,
4759 1.0,
4760 &[0.0, 1.0],
4761 beta_d.exp(),
4762 beta_m.exp(),
4763 array![1.0].view(),
4764 array![1.0].view(),
4765 |u, design_d, deriv_d, design_m| {
4766 let cumulative_d = beta_d.exp() * (1.0 + 0.2 * u);
4767 let cumulative_m = beta_m.exp() * (1.0 + 0.1 * u);
4768 let inst_hazard_d = 0.2 * beta_d.exp();
4769 design_d[0] = 1.0;
4770 deriv_d[0] = 0.0;
4773 design_m[0] = 1.0;
4774 Ok((inst_hazard_d, cumulative_d, cumulative_m))
4775 },
4776 )
4777 .expect("smooth crude-risk quadrature should succeed")
4778 }
4779
4780 #[test]
4781 fn crude_riskgradient_matches_monotoneobjective() {
4782 let beta_d = -0.2_f64;
4783 let beta_m = -0.5_f64;
4784 let result = smooth_crude_risk(beta_d, beta_m);
4785 let eps = 1e-6;
4786
4787 let fd_d = (smooth_crude_risk(beta_d + eps, beta_m).risk
4788 - smooth_crude_risk(beta_d - eps, beta_m).risk)
4789 / (2.0 * eps);
4790 let fd_m = (smooth_crude_risk(beta_d, beta_m + eps).risk
4791 - smooth_crude_risk(beta_d, beta_m - eps).risk)
4792 / (2.0 * eps);
4793
4794 assert!(
4795 (result.diseasegradient[0] - fd_d).abs() < 1e-5,
4796 "disease gradient mismatch for monotone crude risk: analytic={} fd={fd_d}",
4797 result.diseasegradient[0]
4798 );
4799 assert!(
4800 (result.mortalitygradient[0] - fd_m).abs() < 1e-5,
4801 "mortality gradient mismatch for monotone crude risk: analytic={} fd={fd_m}",
4802 result.mortalitygradient[0]
4803 );
4804 }
4805
4806 #[test]
4807 fn survival_working_state_is_ridge_free() {
4808 let age_entry = array![1.0_f64, 2.0_f64];
4809 let age_exit = array![2.0_f64, 3.5_f64];
4810 let event_target = array![1u8, 0u8];
4811 let event_competing = array![0u8, 0u8];
4812 let sampleweight = array![1.0, 1.0];
4813 let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
4814 let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
4815 let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
4816 let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4817 matrix: array![[2.0]],
4818 lambda: 1.7,
4819 range: 1..2,
4820 nullspace_dim: 0,
4821 }]);
4822 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4823 let beta = array![-1.2, 0.4];
4824
4825 let model = survival_model(
4826 survival_inputs(
4827 &age_entry,
4828 &age_exit,
4829 &event_target,
4830 &event_competing,
4831 &sampleweight,
4832 &x_entry,
4833 &x_exit,
4834 &x_derivative,
4835 ),
4836 penalties.clone(),
4837 mono,
4838 SurvivalSpec::Net,
4839 )
4840 .expect("construct survival model");
4841
4842 let state = model.update_state(&beta).expect("survival state");
4843 assert_eq!(
4844 state.ridge_used, 0.0,
4845 "survival objective must not fuse a coefficient ridge"
4846 );
4847 let expected_penalty = 2.0 * penalties.deviance(&beta);
4849 assert!(
4850 (state.penalty_term - expected_penalty).abs() < 1e-12,
4851 "penalty_term mismatch: state={} expected={}",
4852 state.penalty_term,
4853 expected_penalty
4854 );
4855 }
4856
4857 #[test]
4858 fn negative_penalty_lambda_is_rejected() {
4859 let age_entry = array![1.0_f64];
4860 let age_exit = array![2.0_f64];
4861 let event_target = array![1u8];
4862 let event_competing = array![0u8];
4863 let sampleweight = array![1.0];
4864 let x_entry = array![[1.0, 0.0]];
4865 let x_exit = array![[1.0, 0.5]];
4866 let x_derivative = array![[0.0, 1.0]];
4867 let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4868 matrix: array![[1.0]],
4869 lambda: -0.1,
4870 range: 1..2,
4871 nullspace_dim: 0,
4872 }]);
4873
4874 let err = survival_model(
4875 survival_inputs(
4876 &age_entry,
4877 &age_exit,
4878 &event_target,
4879 &event_competing,
4880 &sampleweight,
4881 &x_entry,
4882 &x_exit,
4883 &x_derivative,
4884 ),
4885 penalties,
4886 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4887 SurvivalSpec::Net,
4888 )
4889 .expect_err("negative lambda must be rejected");
4890
4891 assert!(matches!(err, SurvivalError::NonFiniteInput));
4892 }
4893
4894 #[test]
4895 fn penalty_block_range_and_shapemust_match_coefficients() {
4896 let age_entry = array![1.0_f64];
4897 let age_exit = array![2.0_f64];
4898 let event_target = array![1u8];
4899 let event_competing = array![0u8];
4900 let sampleweight = array![1.0];
4901 let x_entry = array![[1.0, 0.0]];
4902 let x_exit = array![[1.0, 0.5]];
4903 let x_derivative = array![[0.0, 1.0]];
4904 let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4905 matrix: array![[1.0]],
4906 lambda: 0.5,
4907 range: 0..2,
4908 nullspace_dim: 0,
4909 }]);
4910
4911 let err = survival_model(
4912 survival_inputs(
4913 &age_entry,
4914 &age_exit,
4915 &event_target,
4916 &event_competing,
4917 &sampleweight,
4918 &x_entry,
4919 &x_exit,
4920 &x_derivative,
4921 ),
4922 penalties,
4923 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
4924 SurvivalSpec::Net,
4925 )
4926 .expect_err("penalty block geometry must match coefficient support");
4927
4928 assert!(matches!(err, SurvivalError::DimensionMismatch));
4929 }
4930
4931 #[test]
4932 fn survivalgradient_matches_ridge_free_objective_fd() {
4933 let age_entry = array![1.0_f64, 2.0_f64, 3.0_f64];
4934 let age_exit = array![2.0_f64, 3.5_f64, 4.0_f64];
4935 let event_target = array![1u8, 0u8, 1u8];
4936 let event_competing = array![0u8, 0u8, 0u8];
4937 let sampleweight = array![1.0, 1.0, 1.0];
4938 let x_entry = array![
4939 [1.0, age_entry[0].ln()],
4940 [1.0, age_entry[1].ln()],
4941 [1.0, age_entry[2].ln()]
4942 ];
4943 let x_exit = array![
4944 [1.0, age_exit[0].ln()],
4945 [1.0, age_exit[1].ln()],
4946 [1.0, age_exit[2].ln()]
4947 ];
4948 let x_derivative = array![
4949 [0.0, 1.0 / age_exit[0]],
4950 [0.0, 1.0 / age_exit[1]],
4951 [0.0, 1.0 / age_exit[2]]
4952 ];
4953 let penalties = PenaltyBlocks::new(Vec::new());
4954 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4955 let beta = array![-1.0, 3.0];
4956
4957 let model = survival_model(
4958 survival_inputs(
4959 &age_entry,
4960 &age_exit,
4961 &event_target,
4962 &event_competing,
4963 &sampleweight,
4964 &x_entry,
4965 &x_exit,
4966 &x_derivative,
4967 ),
4968 penalties,
4969 mono,
4970 SurvivalSpec::Net,
4971 )
4972 .expect("construct survival model");
4973
4974 let state = model.update_state(&beta).expect("state at beta");
4975 let eps = 1e-7;
4976 for j in 0..beta.len() {
4977 let mut plus = beta.clone();
4978 let mut minus = beta.clone();
4979 plus[j] += eps;
4980 minus[j] -= eps;
4981 let state_plus = model.update_state(&plus).expect("state at beta + eps");
4982 let state_minus = model.update_state(&minus).expect("state at beta - eps");
4983 let obj_plus = 0.5 * (state_plus.deviance + state_plus.penalty_term);
4984 let obj_minus = 0.5 * (state_minus.deviance + state_minus.penalty_term);
4985 let fd = (obj_plus - obj_minus) / (2.0 * eps);
4986 assert_eq!(
4987 state.gradient[j].signum(),
4988 fd.signum(),
4989 "objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
4990 state.gradient[j]
4991 );
4992 assert!(
4993 (state.gradient[j] - fd).abs() < 1e-5,
4994 "objective/gradient mismatch at j={j}: grad={} fd={fd}",
4995 state.gradient[j]
4996 );
4997 }
4998 }
4999
5000 fn laml_fd_test_model(lambda: f64) -> WorkingModelSurvival {
5001 let age_entry: Array1<f64> = Array1::from(vec![
5008 30.0, 35.0, 40.0, 45.0, 50.0, 55.0, 60.0, 32.0, 37.0, 42.0, 47.0, 52.0, 57.0, 62.0,
5009 34.0, 39.0, 44.0, 49.0, 54.0, 59.0,
5010 ]);
5011 let age_exit: Array1<f64> = Array1::from(vec![
5012 45.0, 48.0, 55.0, 58.0, 62.0, 66.0, 68.0, 47.0, 52.0, 53.0, 55.0, 60.0, 63.0, 70.0,
5013 48.0, 51.0, 58.0, 62.0, 66.0, 69.0,
5014 ]);
5015 let event_target = Array1::from(vec![
5016 1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
5017 ]);
5018 let event_competing = Array1::<u8>::zeros(age_entry.len());
5019 let sampleweight = Array1::from_elem(age_entry.len(), 1.0_f64);
5020 let n = age_entry.len();
5021 let ln_age_mean: f64 = {
5022 let mut sum = 0.0;
5023 for i in 0..n {
5024 sum += age_entry[i].ln() + age_exit[i].ln();
5025 }
5026 sum / (2.0 * n as f64)
5027 };
5028 let mut x_entry = Array2::<f64>::zeros((n, 2));
5029 let mut x_exit = Array2::<f64>::zeros((n, 2));
5030 let mut x_derivative = Array2::<f64>::zeros((n, 2));
5031 for i in 0..n {
5032 x_entry[[i, 0]] = 1.0;
5033 x_exit[[i, 0]] = 1.0;
5034 x_entry[[i, 1]] = age_entry[i].ln() - ln_age_mean;
5035 x_exit[[i, 1]] = age_exit[i].ln() - ln_age_mean;
5036 x_derivative[[i, 0]] = 0.0;
5037 x_derivative[[i, 1]] = 1.0 / age_exit[i];
5038 }
5039 let penalties = PenaltyBlocks::new(vec![
5040 PenaltyBlock {
5041 matrix: array![[3.0]],
5042 lambda: 0.0,
5043 range: 0..1,
5044 nullspace_dim: 0,
5045 },
5046 PenaltyBlock {
5047 matrix: array![[2.5]],
5048 lambda,
5049 range: 1..2,
5050 nullspace_dim: 0,
5051 },
5052 ]);
5053 survival_model(
5054 survival_inputs(
5055 &age_entry,
5056 &age_exit,
5057 &event_target,
5058 &event_competing,
5059 &sampleweight,
5060 &x_entry,
5061 &x_exit,
5062 &x_derivative,
5063 ),
5064 penalties,
5065 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
5066 SurvivalSpec::Net,
5067 )
5068 .expect("construct LAML FD survival model")
5069 }
5070
5071 fn laml_test_logdet_h(state: &WorkingState) -> f64 {
5072 use gam_linalg::faer_ndarray::FaerEigh;
5073 use gam_solve::estimate::reml::reml_outer_engine::{spectral_epsilon, spectral_regularize};
5074
5075 let h_dense = state.hessian.to_dense();
5076 let (evals, _) = h_dense.eigh(faer::Side::Lower).expect("eigh");
5077 let eps = spectral_epsilon(evals.as_slice().unwrap());
5078 evals
5079 .iter()
5080 .map(|&sigma| spectral_regularize(sigma, eps).ln())
5081 .sum()
5082 }
5083
5084 fn laml_rail_fd_test_model(lambda0: f64, lambda1: f64) -> WorkingModelSurvival {
5091 let age_entry: Array1<f64> = Array1::from(vec![
5092 30.0, 35.0, 40.0, 45.0, 50.0, 55.0, 60.0, 32.0, 37.0, 42.0, 47.0, 52.0, 57.0, 62.0,
5093 34.0, 39.0, 44.0, 49.0, 54.0, 59.0,
5094 ]);
5095 let age_exit: Array1<f64> = Array1::from(vec![
5096 45.0, 48.0, 55.0, 58.0, 62.0, 66.0, 68.0, 47.0, 52.0, 53.0, 55.0, 60.0, 63.0, 70.0,
5097 48.0, 51.0, 58.0, 62.0, 66.0, 69.0,
5098 ]);
5099 let event_target = Array1::from(vec![
5100 1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
5101 ]);
5102 let event_competing = Array1::<u8>::zeros(age_entry.len());
5103 let sampleweight = Array1::from_elem(age_entry.len(), 1.0_f64);
5104 let n = age_entry.len();
5105 let ln_age_mean: f64 = {
5106 let mut sum = 0.0;
5107 for i in 0..n {
5108 sum += age_entry[i].ln() + age_exit[i].ln();
5109 }
5110 sum / (2.0 * n as f64)
5111 };
5112 let mut x_entry = Array2::<f64>::zeros((n, 2));
5113 let mut x_exit = Array2::<f64>::zeros((n, 2));
5114 let mut x_derivative = Array2::<f64>::zeros((n, 2));
5115 for i in 0..n {
5116 x_entry[[i, 0]] = 1.0;
5117 x_exit[[i, 0]] = 1.0;
5118 x_entry[[i, 1]] = age_entry[i].ln() - ln_age_mean;
5119 x_exit[[i, 1]] = age_exit[i].ln() - ln_age_mean;
5120 x_derivative[[i, 0]] = 0.0;
5121 x_derivative[[i, 1]] = 1.0 / age_exit[i];
5122 }
5123 let penalties = PenaltyBlocks::new(vec![
5124 PenaltyBlock {
5125 matrix: array![[3.0]],
5126 lambda: lambda0,
5127 range: 0..1,
5128 nullspace_dim: 0,
5129 },
5130 PenaltyBlock {
5131 matrix: array![[2.5]],
5132 lambda: lambda1,
5133 range: 1..2,
5134 nullspace_dim: 0,
5135 },
5136 ]);
5137 survival_model(
5138 survival_inputs(
5139 &age_entry,
5140 &age_exit,
5141 &event_target,
5142 &event_competing,
5143 &sampleweight,
5144 &x_entry,
5145 &x_exit,
5146 &x_derivative,
5147 ),
5148 penalties,
5149 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
5150 SurvivalSpec::Net,
5151 )
5152 .expect("construct two-active-block rail LAML FD model")
5153 }
5154
5155 #[test]
5180 fn survival_laml_rho_gradient_matches_fd_at_the_over_smoothing_rail() {
5181 use gam_linalg::faer_ndarray::FaerEigh;
5182
5183 const RAIL_RHO0: f64 = 7.394829814011909;
5184 const FREE_RHO1: f64 = -2.45;
5185 const FD_STEP: f64 = 1.0e-4;
5186
5187 let beta0 = array![-2.5_f64, 1.0];
5188 let rho = array![RAIL_RHO0, FREE_RHO1];
5189 let model = laml_rail_fd_test_model(RAIL_RHO0.exp(), FREE_RHO1.exp());
5190
5191 let (value, analytic) = model
5193 .evaluate_survival_lamlcost_and_gradient(
5194 rho.as_slice().expect("contiguous rho"),
5195 &beta0,
5196 )
5197 .expect("rail LAML analytic value+gradient (inner solve must converge at the rail)");
5198
5199 let (rail_model, beta_hat) = model
5201 .reconverge_survival_inner_mode(rho.as_slice().expect("contiguous rho"), &beta0)
5202 .expect("reconverge inner mode at the rail");
5203 let state = rail_model
5204 .update_state(&beta_hat)
5205 .expect("inner state at the rail");
5206 let h_dense = state.hessian.to_dense();
5207 let (evals, _) = h_dense.eigh(faer::Side::Lower).expect("eigh at rail");
5208 let min_ev = evals.iter().copied().fold(f64::INFINITY, f64::min);
5209 let max_ev = evals.iter().copied().fold(f64::NEG_INFINITY, f64::max);
5210 let cond = max_ev / min_ev.abs().max(f64::MIN_POSITIVE);
5211
5212 let term_values = |r: &Array1<f64>| -> (f64, f64, f64) {
5215 let (cand, b) = model
5216 .reconverge_survival_inner_mode(r.as_slice().expect("contiguous rho"), &beta0)
5217 .expect("reconverge for per-term FD");
5218 let st = cand.update_state(&b).expect("state for per-term FD");
5219 let t1 = 0.5 * (st.deviance + st.penalty_term);
5220 let t2 = 0.5 * laml_test_logdet_h(&st);
5221 let t3 = -0.5 * (r[0] + 3.0_f64.ln() + r[1] + 2.5_f64.ln());
5223 (t1, t2, t3)
5224 };
5225
5226 let mut fd = vec![0.0_f64; rho.len()];
5227 let mut fd_terms = vec![(0.0_f64, 0.0_f64, 0.0_f64); rho.len()];
5228 for j in 0..rho.len() {
5229 let mut plus = rho.clone();
5230 plus[j] += FD_STEP;
5231 let mut minus = rho.clone();
5232 minus[j] -= FD_STEP;
5233 let fp = model
5234 .evaluate_survival_lamlcost_and_gradient(
5235 plus.as_slice().expect("contiguous rho"),
5236 &beta0,
5237 )
5238 .expect("rail LAML f+ (probe ρ inner solve must converge)")
5239 .0;
5240 let fm = model
5241 .evaluate_survival_lamlcost_and_gradient(
5242 minus.as_slice().expect("contiguous rho"),
5243 &beta0,
5244 )
5245 .expect("rail LAML f- (probe ρ inner solve must converge)")
5246 .0;
5247 fd[j] = (fp - fm) / (2.0 * FD_STEP);
5248 let (p1, p2, p3) = term_values(&plus);
5249 let (m1, m2, m3) = term_values(&minus);
5250 fd_terms[j] = (
5251 (p1 - m1) / (2.0 * FD_STEP),
5252 (p2 - m2) / (2.0 * FD_STEP),
5253 (p3 - m3) / (2.0 * FD_STEP),
5254 );
5255 }
5256
5257 eprintln!(
5258 "[RAIL-FD] rho=[{RAIL_RHO0:.9}, {FREE_RHO1}] value={value:.9} inner_H min_ev={min_ev:.3e} max_ev={max_ev:.3e} cond={cond:.3e}"
5259 );
5260 let mut amplification = vec![0.0_f64; rho.len()];
5267 for j in 0..rho.len() {
5268 let (dt1, dt2, dt3) = fd_terms[j];
5269 amplification[j] = dt2.abs().max(dt3.abs()) / analytic[j].abs().max(f64::MIN_POSITIVE);
5270 let amp = amplification[j];
5271 eprintln!(
5272 "[RAIL-FD] rho[{j}]: analytic_total={:.6e} fd_total={:.6e} abs_err={:.3e} amplification={amp:.3e} | fd_terms: half_d(dev+pen)={dt1:.6e} half_dlogdetH={dt2:.6e} neg_half_dlogdetS={dt3:.6e}",
5273 analytic[j],
5274 fd[j],
5275 (analytic[j] - fd[j]).abs()
5276 );
5277 }
5278
5279 for j in 0..rho.len() {
5280 let tol = 1.0e-4 * (1.0 + analytic[j].abs().max(fd[j].abs()));
5281 assert!(
5282 (analytic[j] - fd[j]).abs() <= tol,
5283 "survival LAML ρ-gradient desync at coordinate {j} in the over-smoothing rail regime: \
5284 analytic={:.6e} fd={:.6e} (inner H cond={cond:.3e}); see the per-term [RAIL-FD] grid above",
5285 analytic[j],
5286 fd[j],
5287 );
5288 }
5289
5290 let max_amplification = amplification
5294 .iter()
5295 .copied()
5296 .fold(0.0_f64, f64::max);
5297 assert!(
5298 max_amplification.is_finite() && max_amplification <= 1.0e8,
5299 "rail logdet-gradient amplification ratio not finite/bounded: {max_amplification:.3e}"
5300 );
5301 }
5302
5303 #[test]
5310 fn survival_laml_rho_gradient_matches_fd_at_interior_rho() {
5311 const INTERIOR_RHO0: f64 = 0.3;
5312 const INTERIOR_RHO1: f64 = -0.5;
5313 const FD_STEP: f64 = 1.0e-4;
5314
5315 let beta0 = array![-2.5_f64, 1.0];
5316 let rho = array![INTERIOR_RHO0, INTERIOR_RHO1];
5317 let model = laml_rail_fd_test_model(INTERIOR_RHO0.exp(), INTERIOR_RHO1.exp());
5318 let (_value, analytic) = model
5319 .evaluate_survival_lamlcost_and_gradient(
5320 rho.as_slice().expect("contiguous rho"),
5321 &beta0,
5322 )
5323 .expect("interior LAML analytic value+gradient");
5324
5325 for j in 0..rho.len() {
5326 let mut plus = rho.clone();
5327 plus[j] += FD_STEP;
5328 let mut minus = rho.clone();
5329 minus[j] -= FD_STEP;
5330 let fp = model
5331 .evaluate_survival_lamlcost_and_gradient(
5332 plus.as_slice().expect("contiguous rho"),
5333 &beta0,
5334 )
5335 .expect("interior LAML f+")
5336 .0;
5337 let fm = model
5338 .evaluate_survival_lamlcost_and_gradient(
5339 minus.as_slice().expect("contiguous rho"),
5340 &beta0,
5341 )
5342 .expect("interior LAML f-")
5343 .0;
5344 let fd = (fp - fm) / (2.0 * FD_STEP);
5345 let tol = 1.0e-4 * (1.0 + analytic[j].abs().max(fd.abs()));
5346 assert!(
5347 (analytic[j] - fd).abs() <= tol,
5348 "interior survival LAML ρ-gradient mismatch at coordinate {j}: \
5349 analytic={:.6e} fd={:.6e}",
5350 analytic[j],
5351 fd,
5352 );
5353 }
5354 }
5355
5356 #[test]
5357 fn survival_solver_damping_converges_undamped_objective() {
5358 let rho = -0.35_f64;
5359 let model = laml_fd_test_model(rho.exp());
5360 let beta0 = array![-2.5_f64, 1.0];
5361 let (converged_model, beta) = model
5362 .reconverge_survival_inner_mode(&[rho], &beta0)
5363 .expect("converge survival mode with solver-only damping");
5364 let state = converged_model
5365 .update_state(&beta)
5366 .expect("evaluate undamped objective at converged mode");
5367
5368 assert_eq!(
5369 state.ridge_used, 0.0,
5370 "solver damping must not enter the converged statistical objective"
5371 );
5372 let undamped_stationarity = array1_l2_norm(&state.gradient);
5373 assert!(
5374 undamped_stationarity <= 1.0e-9,
5375 "solver must converge the undamped objective; ||gradient||={undamped_stationarity:.3e}"
5376 );
5377 }
5378
5379 #[test]
5380 fn laml_gradient_and_objective_ignore_inactive_penalty_prefix_blocks() {
5381 let rho0 = -0.35_f64;
5395 let beta = array![-2.5_f64, 1.0];
5396 let model = laml_fd_test_model(rho0.exp());
5397 let state = model
5398 .update_state(&beta)
5399 .expect("state for LAML prefix-skip test");
5400
5401 assert_eq!(model.penalties.blocks.len(), 2);
5406 assert_eq!(model.penalties.blocks[0].lambda, 0.0);
5407 assert!(model.penalties.blocks[1].lambda > 0.0);
5408
5409 let rho = Array1::from_iter(
5410 model
5411 .penalties
5412 .blocks
5413 .iter()
5414 .filter(|b| b.lambda > 0.0)
5415 .map(|b| b.lambda.ln()),
5416 );
5417 assert_eq!(
5418 rho.len(),
5419 1,
5420 "fixture should expose exactly one active penalty block for the rho vector"
5421 );
5422
5423 let (obj, grad) = model
5424 .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
5425 .expect("survival LAML objective and gradient");
5426
5427 let expected = 0.5 * (state.deviance + state.penalty_term)
5428 + 0.5 * laml_test_logdet_h(&state)
5429 - 0.5 * (rho0 + 2.5_f64.ln());
5430 assert_eq!(
5431 grad.len(),
5432 1,
5433 "rho-gradient must match the active-penalty count, not the full block list"
5434 );
5435 assert!(
5436 (obj - expected).abs() < 1e-10,
5437 "survival LAML objective mismatch with inactive prefix block: obj={obj} expected={expected}",
5438 );
5439 assert!(
5440 grad[0].is_finite(),
5441 "rho-gradient must be finite: {}",
5442 grad[0]
5443 );
5444 }
5445
5446 #[test]
5447 fn structural_monotonicgradient_matchesobjectivefd() {
5448 let age_entry = array![1.0_f64, 1.3_f64, 1.8_f64];
5449 let age_exit = array![1.6_f64, 2.1_f64, 2.7_f64];
5450 let event_target = array![1u8, 0u8, 1u8];
5451 let event_competing = array![0u8, 0u8, 0u8];
5452 let sampleweight = array![1.0, 1.0, 1.0];
5453
5454 let x_entry = array![
5457 [1.0, 0.2, 0.05, -0.7],
5458 [1.0, 0.5, 0.20, 0.1],
5459 [1.0, 0.9, 0.60, 1.2]
5460 ];
5461 let x_exit = array![
5462 [1.0, 0.4, 0.16, -0.7],
5463 [1.0, 0.8, 0.64, 0.1],
5464 [1.0, 1.1, 1.21, 1.2]
5465 ];
5466 let x_derivative = array![
5467 [0.0, 0.8, 0.64, 0.0],
5468 [0.0, 0.7, 1.12, 0.0],
5469 [0.0, 0.6, 1.32, 0.0]
5470 ];
5471 let penalties = PenaltyBlocks::new(Vec::new());
5472 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5473 let mut model = survival_model(
5474 survival_inputs(
5475 &age_entry,
5476 &age_exit,
5477 &event_target,
5478 &event_competing,
5479 &sampleweight,
5480 &x_entry,
5481 &x_exit,
5482 &x_derivative,
5483 ),
5484 penalties,
5485 mono,
5486 SurvivalSpec::Net,
5487 )
5488 .expect("construct structural survival model");
5489 model
5490 .set_structural_monotonicity(true, 3)
5491 .expect("enable structural monotonicity");
5492 let constraints = model
5493 .monotonicity_linear_constraints()
5494 .expect("structural derivative constraints");
5495 assert_eq!(constraints.a.nrows(), 2);
5496 assert_eq!(constraints.a.ncols(), 4);
5497 assert_eq!(constraints.a.row(0).to_vec(), vec![0.0, 1.0, 0.0, 0.0]);
5498 assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 0.0, 1.0, 0.0]);
5499 assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
5500
5501 let beta = array![0.2, 0.2, 0.1, 0.2];
5502 let state = model.update_state(&beta).expect("state at structural beta");
5503 let eps = 1e-7;
5504 for j in 0..beta.len() {
5505 let mut plus = beta.clone();
5506 let mut minus = beta.clone();
5507 plus[j] += eps;
5508 minus[j] -= eps;
5509 let state_plus = model.update_state(&plus).expect("state at beta + eps");
5510 let state_minus = model.update_state(&minus).expect("state at beta - eps");
5511 let obj_plus = 0.5 * (state_plus.deviance + state_plus.penalty_term);
5512 let obj_minus = 0.5 * (state_minus.deviance + state_minus.penalty_term);
5513 let fd = (obj_plus - obj_minus) / (2.0 * eps);
5514 assert_eq!(
5515 state.gradient[j].signum(),
5516 fd.signum(),
5517 "structural objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
5518 state.gradient[j]
5519 );
5520 assert!(
5521 (state.gradient[j] - fd).abs() < 2e-5,
5522 "structural objective/gradient mismatch at j={j}: grad={} fd={fd}",
5523 state.gradient[j]
5524 );
5525 }
5526 }
5527
5528 #[test]
5529 fn structural_monotonic_lamlgradient_returns_finitevalues() {
5530 let age_entry = array![1.0_f64, 1.2_f64];
5531 let age_exit = array![1.5_f64, 2.0_f64];
5532 let event_target = array![1u8, 0u8];
5533 let event_competing = array![0u8, 0u8];
5534 let sampleweight = array![1.0, 1.0];
5535 let x_entry = array![[1.0, 0.2, -0.5], [1.0, 0.4, 0.2]];
5536 let x_exit = array![[1.0, 0.5, -0.5], [1.0, 0.8, 0.2]];
5537 let x_derivative = array![[0.0, 0.9, 0.0], [0.0, 0.7, 0.0]];
5538 let penalties = PenaltyBlocks::new(Vec::new());
5539 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5540 let mut model = survival_model(
5541 survival_inputs(
5542 &age_entry,
5543 &age_exit,
5544 &event_target,
5545 &event_competing,
5546 &sampleweight,
5547 &x_entry,
5548 &x_exit,
5549 &x_derivative,
5550 ),
5551 penalties,
5552 mono,
5553 SurvivalSpec::Net,
5554 )
5555 .expect("construct structural survival model");
5556 model
5557 .set_structural_monotonicity(true, 2)
5558 .expect("enable structural monotonicity");
5559 model.penalties = PenaltyBlocks::new(vec![PenaltyBlock {
5561 matrix: array![[1.0]],
5562 lambda: 0.7,
5563 range: 1..2,
5564 nullspace_dim: 0,
5565 }]);
5566 let beta = array![0.2, 0.2, 0.1];
5567 let state = model.update_state(&beta).expect("state at structural beta");
5568 let rho = Array1::from_iter(
5569 model
5570 .penalties
5571 .blocks
5572 .iter()
5573 .filter(|b| b.lambda > 0.0)
5574 .map(|b| b.lambda.ln()),
5575 );
5576 let (obj, grad) = model
5577 .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
5578 .expect("laml gradient should work in structural mode");
5579 assert!(obj.is_finite());
5580 assert_eq!(grad.len(), 1);
5581 assert!(grad[0].is_finite());
5582 }
5583
5584 #[test]
5585 fn structural_monotonicity_switches_to_tiny_derivative_guard_constraints() {
5586 let age_entry = array![1.0_f64];
5587 let age_exit = array![2.0_f64];
5588 let event_target = array![1u8];
5589 let event_competing = array![0u8];
5590 let sampleweight = array![1.0];
5591 let x_entry = array![[0.0]];
5592 let x_exit = array![[0.2]];
5593 let x_derivative = array![[1.0]];
5594
5595 let penalties = PenaltyBlocks::new(Vec::new());
5596 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5597 let mut model = survival_model(
5598 survival_inputs(
5599 &age_entry,
5600 &age_exit,
5601 &event_target,
5602 &event_competing,
5603 &sampleweight,
5604 &x_entry,
5605 &x_exit,
5606 &x_derivative,
5607 ),
5608 penalties,
5609 mono,
5610 SurvivalSpec::Net,
5611 )
5612 .expect("construct structural survival model");
5613
5614 let beta = array![-3.0];
5615 assert!(
5616 model.update_state(&beta).is_err(),
5617 "negative derivative coefficient should violate derivative guard"
5618 );
5619
5620 model
5621 .set_structural_monotonicity(true, 1)
5622 .expect("enable structural monotonicity");
5623 let constraints = model
5624 .monotonicity_linear_constraints()
5625 .expect("structural derivative constraints");
5626 assert_eq!(constraints.a.nrows(), 1);
5627 assert_eq!(constraints.a.ncols(), 1);
5628 assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
5629 assert!(constraints.b[0].abs() <= 1e-12);
5631 let state = model
5632 .update_state(&array![1e-6])
5633 .expect("small positive derivative coefficient should remain feasible");
5634 assert!(state.deviance.is_finite());
5635 }
5636
5637 #[test]
5638 fn derivative_offset_must_clear_nonstructural_monotonicity_threshold() {
5639 let age_entry = array![1.0_f64];
5640 let age_exit = array![2.0_f64];
5641 let event_target = array![1u8];
5642 let event_competing = array![0u8];
5643 let sampleweight = array![1.0];
5644 let x_entry = array![[1.0, 0.0]];
5645 let x_exit = array![[1.0, 0.0]];
5646 let x_derivative = array![[0.0, 0.0]];
5647 let penalties = PenaltyBlocks::new(Vec::new());
5648 let monotonicity = SurvivalMonotonicityPenalty { tolerance: 3.0 };
5649 let eta_entry_offset = array![0.0];
5650 let eta_exit_offset = array![0.0];
5651 let derivative_offset_below_guard = array![2.0];
5652 let derivative_offset_above_guard = array![3.1];
5653 let offsets_below_guard = SurvivalBaselineOffsets {
5654 eta_entry: eta_entry_offset.view(),
5655 eta_exit: eta_exit_offset.view(),
5656 derivative_exit: derivative_offset_below_guard.view(),
5657 };
5658 let offsets_above_guard = SurvivalBaselineOffsets {
5659 eta_entry: eta_entry_offset.view(),
5660 eta_exit: eta_exit_offset.view(),
5661 derivative_exit: derivative_offset_above_guard.view(),
5662 };
5663
5664 let model_below_guard = survival_model_with_offsets(
5665 survival_inputs(
5666 &age_entry,
5667 &age_exit,
5668 &event_target,
5669 &event_competing,
5670 &sampleweight,
5671 &x_entry,
5672 &x_exit,
5673 &x_derivative,
5674 ),
5675 Some(offsets_below_guard),
5676 penalties.clone(),
5677 monotonicity,
5678 SurvivalSpec::Net,
5679 )
5680 .expect("construct model with derivative offset below guard");
5681 let err = model_below_guard
5682 .update_state(&array![0.0, 0.0])
5683 .expect_err("derivative offset below guard should be rejected");
5684 let err_text = err.to_string();
5685 assert!(
5686 err_text.contains("d_eta/dt=2.000e0") && err_text.contains("tolerance=3.000e0"),
5687 "expected derivative guard rejection to report the offset-driven derivative: {err_text}"
5688 );
5689
5690 let model_above_guard = survival_model_with_offsets(
5691 survival_inputs(
5692 &age_entry,
5693 &age_exit,
5694 &event_target,
5695 &event_competing,
5696 &sampleweight,
5697 &x_entry,
5698 &x_exit,
5699 &x_derivative,
5700 ),
5701 Some(offsets_above_guard),
5702 penalties,
5703 SurvivalMonotonicityPenalty { tolerance: 3.0 },
5704 SurvivalSpec::Net,
5705 )
5706 .expect("construct model with derivative offset above guard");
5707 let state = model_above_guard
5708 .update_state(&array![0.0, 0.0])
5709 .expect("derivative offset above guard should remain feasible");
5710 assert!(state.deviance.is_finite());
5711 }
5712
5713 #[test]
5714 fn structural_monotonicity_rejects_negative_derivative_offsets() {
5715 let age_entry = array![1.0_f64];
5716 let age_exit = array![2.0_f64];
5717 let event_target = array![1u8];
5718 let event_competing = array![0u8];
5719 let sampleweight = array![1.0];
5720 let x_entry = array![[0.0]];
5721 let x_exit = array![[0.2]];
5722 let x_derivative = array![[1.0]];
5723 let eta_entry = array![0.0];
5724 let eta_exit = array![0.0];
5725 let derivative_exit = array![-1e-3];
5726 let offsets = SurvivalBaselineOffsets {
5727 eta_entry: eta_entry.view(),
5728 eta_exit: eta_exit.view(),
5729 derivative_exit: derivative_exit.view(),
5730 };
5731
5732 let mut model = survival_model_with_offsets(
5733 survival_inputs(
5734 &age_entry,
5735 &age_exit,
5736 &event_target,
5737 &event_competing,
5738 &sampleweight,
5739 &x_entry,
5740 &x_exit,
5741 &x_derivative,
5742 ),
5743 Some(offsets),
5744 PenaltyBlocks::new(Vec::new()),
5745 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5746 SurvivalSpec::Net,
5747 )
5748 .expect("construct structural survival model");
5749 let err = model
5750 .set_structural_monotonicity(true, 1)
5751 .expect_err("negative derivative offsets must be rejected");
5752 assert!(
5753 err.to_string()
5754 .contains("structural monotonicity requires nonnegative derivative offsets"),
5755 "unexpected error: {err}"
5756 );
5757 }
5758
5759 #[test]
5760 fn structural_monotonicity_emits_coefficient_constraints() {
5761 let age_entry = array![1.0_f64, 1.5_f64];
5762 let age_exit = array![2.0_f64, 3.0_f64];
5763 let event_target = array![1u8, 0u8];
5764 let event_competing = array![0u8, 0u8];
5765 let sampleweight = array![1.0, 1.0];
5766 let x_entry = array![[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]];
5767 let x_exit = array![[0.2, 0.4, 1.0], [0.3, 0.5, 1.0]];
5768 let x_derivative = array![[0.3, 0.2, 0.0], [0.4, 0.1, 0.0]];
5769
5770 let mut model = survival_model(
5771 survival_inputs(
5772 &age_entry,
5773 &age_exit,
5774 &event_target,
5775 &event_competing,
5776 &sampleweight,
5777 &x_entry,
5778 &x_exit,
5779 &x_derivative,
5780 ),
5781 PenaltyBlocks::new(Vec::new()),
5782 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5783 SurvivalSpec::Net,
5784 )
5785 .expect("construct structural survival model");
5786 model
5787 .set_structural_monotonicity(true, 2)
5788 .expect("enable structural monotonicity");
5789
5790 let constraints = model
5791 .monotonicity_linear_constraints()
5792 .expect("structural derivative constraints");
5793
5794 assert_eq!(constraints.a.nrows(), 2);
5795 assert_eq!(constraints.a.ncols(), 3);
5796 assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0, 0.0]);
5797 assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0, 0.0]);
5798 assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
5799 }
5800
5801 #[test]
5802 fn structural_monotonicity_preserves_inactive_time_columns_in_constraints() {
5803 let age_entry = array![1.0_f64];
5804 let age_exit = array![2.0_f64];
5805 let event_target = array![1u8];
5806 let event_competing = array![0u8];
5807 let sampleweight = array![1.0];
5808 let x_entry = array![[1.0, 0.2]];
5809 let x_exit = array![[1.0, 0.6]];
5810 let x_derivative = array![[0.0, 1.0]];
5811
5812 let mut model = survival_model(
5813 survival_inputs(
5814 &age_entry,
5815 &age_exit,
5816 &event_target,
5817 &event_competing,
5818 &sampleweight,
5819 &x_entry,
5820 &x_exit,
5821 &x_derivative,
5822 ),
5823 PenaltyBlocks::new(Vec::new()),
5824 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5825 SurvivalSpec::Net,
5826 )
5827 .expect("construct structural survival model");
5828 model
5829 .set_structural_monotonicity(true, 2)
5830 .expect("enable structural monotonicity");
5831
5832 let constraints = model
5833 .monotonicity_linear_constraints()
5834 .expect("structural derivative constraints");
5835
5836 assert_eq!(constraints.a.nrows(), 1);
5837 assert!(
5838 constraints.a[[0, 0]].abs() <= 1e-12,
5839 "inactive time column should remain unconstrained"
5840 );
5841 assert!(
5842 (constraints.a[[0, 1]] - 1.0).abs() <= 1e-12,
5843 "active time column should remain constrained"
5844 );
5845 }
5846
5847 #[test]
5848 fn structural_monotonicity_preserves_sparse_row_patterns() {
5849 let age_entry = array![1.0_f64, 1.5_f64];
5850 let age_exit = array![2.0_f64, 2.5_f64];
5851 let event_target = array![1u8, 1u8];
5852 let event_competing = array![0u8, 0u8];
5853 let sampleweight = array![1.0, 1.0];
5854 let x_entry = array![[0.0, 0.0], [0.0, 0.0]];
5855 let x_exit = array![[0.4, 0.2], [0.6, 0.3]];
5856 let x_derivative = array![[1.0, 0.0], [1.0, 0.5]];
5857
5858 let mut model = survival_model(
5859 survival_inputs(
5860 &age_entry,
5861 &age_exit,
5862 &event_target,
5863 &event_competing,
5864 &sampleweight,
5865 &x_entry,
5866 &x_exit,
5867 &x_derivative,
5868 ),
5869 PenaltyBlocks::new(Vec::new()),
5870 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5871 SurvivalSpec::Net,
5872 )
5873 .expect("construct structural survival model");
5874 model
5875 .set_structural_monotonicity(true, 2)
5876 .expect("enable structural monotonicity");
5877
5878 let constraints = model
5879 .monotonicity_linear_constraints()
5880 .expect("structural derivative constraints");
5881
5882 assert_eq!(constraints.a.nrows(), 2);
5883 assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0]);
5884 assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0]);
5885 }
5886
5887 #[test]
5888 fn update_state_rejects_negative_exit_derivative_for_censoredrows() {
5889 let age_entry = array![1.0_f64];
5890 let age_exit = array![1.1_f64];
5891 let event_target = array![0u8];
5892 let event_competing = array![0u8];
5893 let sampleweight = array![1.0];
5894 let x_entry = array![[0.0]];
5895 let x_exit = array![[0.0]];
5896 let x_derivative = array![[-1.0]];
5897 let penalties = PenaltyBlocks::new(Vec::new());
5898 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5899 let model = survival_model(
5900 survival_inputs(
5901 &age_entry,
5902 &age_exit,
5903 &event_target,
5904 &event_competing,
5905 &sampleweight,
5906 &x_entry,
5907 &x_exit,
5908 &x_derivative,
5909 ),
5910 penalties,
5911 mono,
5912 SurvivalSpec::Net,
5913 )
5914 .expect("construct censored survival model");
5915
5916 let err = model
5917 .update_state(&array![1.0])
5918 .expect_err("censored row should still enforce monotonic derivative");
5919 assert!(
5920 matches!(err, EstimationError::ParameterConstraintViolation(_)),
5921 "unexpected error: {err:?}"
5922 );
5923 }
5924
5925 fn crude_risk_quadrature_error(
5926 cumulative_entry: f64,
5927 cumulative_exit: f64,
5928 hazard_exit: f64,
5929 ) -> SurvivalError {
5930 calculate_crude_risk_quadrature(
5931 1.0,
5932 2.0,
5933 &[],
5934 0.4,
5935 0.2,
5936 array![1.0].view(),
5937 array![1.0].view(),
5938 |_, design_d, deriv_d, design_m| {
5939 design_d[0] = 1.0;
5940 deriv_d[0] = 0.0;
5941 design_m[0] = 1.0;
5942 Ok((cumulative_entry, cumulative_exit, hazard_exit))
5943 },
5944 )
5945 .expect_err("invalid hazards should fail")
5946 }
5947
5948 #[test]
5949 fn crude_risk_quadrature_rejects_decreasing_cumulative_hazard() {
5950 let err = crude_risk_quadrature_error(0.1, 0.3, 0.25);
5951 assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
5952 }
5953
5954 #[test]
5955 fn crude_risk_quadrature_rejects_nonpositive_instantaneous_hazard() {
5956 let err = crude_risk_quadrature_error(0.0, 0.4, 0.25);
5957 assert!(matches!(err, SurvivalError::NonPositiveHazard));
5958 }
5959
5960 #[test]
5961 fn laml_no_penalties_matches_documentedobjective() {
5962 let age_entry = array![40.0, 45.0, 50.0, 55.0];
5963 let age_exit = array![44.0, 49.0, 54.0, 59.0];
5964 let event_target = array![1u8, 0u8, 1u8, 0u8];
5965 let event_competing = Array1::<u8>::zeros(4);
5966 let sampleweight = Array1::ones(4);
5967 let x_entry = array![
5968 [1.0, -0.2, 0.04],
5969 [1.0, -0.1, 0.01],
5970 [1.0, 0.0, 0.0],
5971 [1.0, 0.1, 0.01]
5972 ];
5973 let x_exit = array![
5974 [1.0, -0.12, 0.0144],
5975 [1.0, -0.02, 0.0004],
5976 [1.0, 0.08, 0.0064],
5977 [1.0, 0.18, 0.0324]
5978 ];
5979 let x_derivative = array![
5980 [0.0, 0.02, 0.001],
5981 [0.0, 0.02, 0.001],
5982 [0.0, 0.02, 0.001],
5983 [0.0, 0.02, 0.001]
5984 ];
5985 let penalties = PenaltyBlocks::new(Vec::new());
5986 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5987 let beta = array![-2.0, 0.7, 0.2];
5988
5989 let model = survival_model(
5990 survival_inputs(
5991 &age_entry,
5992 &age_exit,
5993 &event_target,
5994 &event_competing,
5995 &sampleweight,
5996 &x_entry,
5997 &x_exit,
5998 &x_derivative,
5999 ),
6000 penalties,
6001 mono,
6002 SurvivalSpec::Net,
6003 )
6004 .expect("construct survival model");
6005
6006 let state = model.update_state(&beta).expect("state at beta");
6007 let rho = Array1::from_iter(
6008 model
6009 .penalties
6010 .blocks
6011 .iter()
6012 .filter(|b| b.lambda > 0.0)
6013 .map(|b| b.lambda.ln()),
6014 );
6015 let (obj, grad) = model
6016 .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
6017 .expect("laml objective for no-penalty model");
6018
6019 let h_dense = state.hessian.to_dense();
6020 let logdet_h: f64 = {
6036 use gam_problem::PseudoLogdetMode;
6037 use gam_solve::estimate::reml::reml_outer_engine::{
6038 DenseSpectralOperator, HessianFactorization,
6039 };
6040 let has_left_truncation = age_entry.iter().any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD);
6041 let mode = if has_left_truncation {
6042 PseudoLogdetMode::HardPseudo
6043 } else {
6044 PseudoLogdetMode::Smooth
6045 };
6046 DenseSpectralOperator::from_symmetric_with_mode(&h_dense, mode)
6047 .expect("survival LAML Hessian operator")
6048 .logdet()
6049 };
6050 let expected = 0.5 * (state.deviance + state.penalty_term) + 0.5 * logdet_h;
6051
6052 assert_eq!(grad.len(), 0);
6053 assert!(
6054 (obj - expected).abs() < 1e-10,
6055 "no-penalty LAML objective mismatch: obj={} expected={}",
6056 obj,
6057 expected
6058 );
6059 }
6060
6061 #[test]
6062 fn monotonicity_constraints_collapse_positive_collinearrows() {
6063 let a = array![[0.0, 0.5, 0.0], [0.0, 0.25, 0.0], [0.0, 0.125, 0.0]];
6064 let b = array![1e-8, 1e-8, 1e-8];
6065
6066 let compressed = compress_positive_collinear_constraints(&a, &b);
6067
6068 assert_eq!(compressed.a.nrows(), 1);
6069 assert_eq!(compressed.a.ncols(), 3);
6070 assert!(compressed.a[[0, 0]].abs() <= 1e-12);
6071 assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
6072 assert!(compressed.a[[0, 2]].abs() <= 1e-12);
6073 assert!((compressed.b[0] - 8e-8).abs() <= 1e-18);
6074 }
6075
6076 #[test]
6077 fn monotonicity_constraints_preserve_distinct_directions() {
6078 let a = array![[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]];
6079 let b = array![0.2, 0.3, 0.1];
6080
6081 let compressed = compress_positive_collinear_constraints(&a, &b);
6082
6083 assert_eq!(compressed.a.nrows(), 2);
6084 let mut saw_x = false;
6085 let mut saw_y = false;
6086 for i in 0..compressed.a.nrows() {
6087 if (compressed.a[[i, 0]] - 1.0).abs() <= 1e-12 && compressed.a[[i, 1]].abs() <= 1e-12 {
6088 saw_x = true;
6089 assert!((compressed.b[i] - 0.2).abs() <= 1e-12);
6090 }
6091 if compressed.a[[i, 0]].abs() <= 1e-12 && (compressed.a[[i, 1]] - 1.0).abs() <= 1e-12 {
6092 saw_y = true;
6093 assert!((compressed.b[i] - 0.3).abs() <= 1e-12);
6094 }
6095 }
6096 assert!(saw_x);
6097 assert!(saw_y);
6098 }
6099
6100 #[test]
6101 fn monotonicity_constraints_cluster_near_collinearrows() {
6102 let a = array![
6103 [0.0, 0.5, 0.0],
6104 [0.0, 0.50000000003, 0.0],
6105 [0.0, 0.49999999997, 0.0]
6106 ];
6107 let b = array![1e-8, 1.00000000005e-8, 0.99999999995e-8];
6108
6109 let compressed = compress_positive_collinear_constraints(&a, &b);
6110
6111 assert_eq!(compressed.a.nrows(), 1);
6112 assert_eq!(compressed.a.ncols(), 3);
6113 assert!(compressed.a[[0, 0]].abs() <= 1e-12);
6114 assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
6115 assert!(compressed.a[[0, 2]].abs() <= 1e-12);
6116 assert!((compressed.b[0] - 2.0e-8).abs() <= 1e-18);
6117 }
6118
6119 #[test]
6120 fn monotonicity_constraints_cluster_spline_like_near_duplicates() {
6121 let a = array![
6122 [0.0, 0.401, 0.302, 0.197],
6123 [0.0, 0.40100000003, 0.30199999998, 0.19700000001],
6124 [0.0, 0.40099999997, 0.30200000002, 0.19699999999],
6125 [0.0, 0.125, 0.500, 0.375]
6126 ];
6127 let b = array![2.0e-8, 2.00000000004e-8, 1.99999999996e-8, 3.0e-8];
6128
6129 let compressed = compress_positive_collinear_constraints(&a, &b);
6130
6131 assert_eq!(compressed.a.nrows(), 2);
6132 let mut clustered_face = false;
6133 let mut distinct_face = false;
6134 for i in 0..compressed.a.nrows() {
6135 let row = compressed.a.row(i);
6136 if row[1] > 0.99 && row[2] > 0.7 && row[3] > 0.49 {
6137 clustered_face = true;
6138 assert!((compressed.b[i] - (2.0e-8 / 0.401)).abs() <= 1e-12);
6139 } else {
6140 distinct_face = true;
6141 assert!((row[1] - 0.25).abs() <= 1e-12);
6142 assert!((row[2] - 1.0).abs() <= 1e-12);
6143 assert!((row[3] - 0.75).abs() <= 1e-12);
6144 assert!((compressed.b[i] - 6.0e-8).abs() <= 1e-18);
6145 }
6146 }
6147 assert!(clustered_face);
6148 assert!(distinct_face);
6149 }
6150
6151 #[test]
6152 fn linear_time_monotonicity_constraints_reduce_to_single_halfspace() {
6153 let age_entry = array![1.0_f64, 1.0, 1.0];
6154 let age_exit = array![2.0_f64, 4.0, 8.0];
6155 let event_target = array![0u8, 1u8, 0u8];
6156 let event_competing = array![0u8, 0u8, 0u8];
6157 let sampleweight = array![1.0, 1.0, 1.0];
6158 let x_entry = array![
6159 [1.0, age_entry[0].ln()],
6160 [1.0, age_entry[1].ln()],
6161 [1.0, age_entry[2].ln()]
6162 ];
6163 let x_exit = array![
6164 [1.0, age_exit[0].ln()],
6165 [1.0, age_exit[1].ln()],
6166 [1.0, age_exit[2].ln()]
6167 ];
6168 let x_derivative = array![[0.0, 0.5], [0.0, 0.25], [0.0, 0.125]];
6169 let penalties = PenaltyBlocks::new(Vec::new());
6170 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
6171
6172 let collocation_offsets = Array1::zeros(x_derivative.nrows());
6173 let mut inputs = survival_inputs(
6174 &age_entry,
6175 &age_exit,
6176 &event_target,
6177 &event_competing,
6178 &sampleweight,
6179 &x_entry,
6180 &x_exit,
6181 &x_derivative,
6182 );
6183 inputs.monotonicity_constraint_rows = Some(x_derivative.view());
6184 inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
6185
6186 let model = survival_model(inputs, penalties, mono, SurvivalSpec::Net)
6187 .expect("construct linear survival model");
6188
6189 let constraints = model
6190 .monotonicity_linear_constraints()
6191 .expect("monotonicity constraints");
6192 assert_eq!(constraints.a.nrows(), 1);
6193 assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
6194 assert!((constraints.b[0] - 8e-8).abs() <= 1e-12);
6195 }
6196
6197 #[test]
6198 fn monotonicity_constraints_skip_numericallyzerorows() {
6199 let age_entry = array![1.0_f64, 1.0, 1.0];
6200 let age_exit = array![2.0_f64, 3.0, 4.0];
6201 let event_target = array![0u8, 0u8, 0u8];
6202 let event_competing = array![0u8, 0u8, 0u8];
6203 let sampleweight = array![1.0, 1.0, 1.0];
6204 let x_entry = array![[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]];
6205 let x_exit = x_entry.clone();
6206 let x_derivative = array![[0.0, 0.0], [0.0, 1e-16], [0.0, 0.25]];
6207
6208 let collocation_offsets = Array1::zeros(x_derivative.nrows());
6209 let mut inputs = survival_inputs(
6210 &age_entry,
6211 &age_exit,
6212 &event_target,
6213 &event_competing,
6214 &sampleweight,
6215 &x_entry,
6216 &x_exit,
6217 &x_derivative,
6218 );
6219 inputs.monotonicity_constraint_rows = Some(x_derivative.view());
6220 inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
6221
6222 let model = survival_model(
6223 inputs,
6224 PenaltyBlocks::new(Vec::new()),
6225 SurvivalMonotonicityPenalty { tolerance: 0.0 },
6226 SurvivalSpec::Net,
6227 )
6228 .expect("construct survival model");
6229
6230 let constraints = model
6231 .monotonicity_linear_constraints()
6232 .expect("nonzero derivative row should remain");
6233 assert_eq!(constraints.a.nrows(), 1);
6234 assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
6235 assert!(constraints.b[0].abs() <= 1e-18);
6236 }
6237
6238 #[test]
6239 fn censoredrows_allowzero_boundary_derivative() {
6240 let age_entry = array![1.0_f64];
6241 let age_exit = array![2.0_f64];
6242 let event_target = array![0u8];
6243 let event_competing = array![0u8];
6244 let sampleweight = array![1.0];
6245 let x_entry = array![[0.0]];
6246 let x_exit = array![[0.0]];
6247 let x_derivative = array![[1.0]];
6248
6249 let model = survival_model(
6250 survival_inputs(
6251 &age_entry,
6252 &age_exit,
6253 &event_target,
6254 &event_competing,
6255 &sampleweight,
6256 &x_entry,
6257 &x_exit,
6258 &x_derivative,
6259 ),
6260 PenaltyBlocks::new(Vec::new()),
6261 SurvivalMonotonicityPenalty { tolerance: 0.0 },
6262 SurvivalSpec::Net,
6263 )
6264 .expect("construct censored survival model");
6265
6266 let state = model
6267 .update_state(&array![0.0])
6268 .expect("censored boundary derivative should remain feasible with zero tolerance");
6269 assert_eq!(state.deviance, 0.0);
6270 assert_eq!(state.log_likelihood, 0.0);
6271 assert_eq!(state.gradient, array![0.0]);
6272 }
6273
6274 #[test]
6275 fn eventrows_keep_positive_derivative_constraint() {
6276 let age_entry = array![1.0_f64, 1.0];
6277 let age_exit = array![2.0_f64, 4.0];
6278 let event_target = array![0u8, 1u8];
6279 let event_competing = array![0u8, 0u8];
6280 let sampleweight = array![1.0, 1.0];
6281 let x_entry = array![[0.0], [0.0]];
6282 let x_exit = array![[0.0], [0.0]];
6283 let x_derivative = array![[0.5], [0.25]];
6284
6285 let collocation_offsets = Array1::zeros(x_derivative.nrows());
6286 let mut inputs = survival_inputs(
6287 &age_entry,
6288 &age_exit,
6289 &event_target,
6290 &event_competing,
6291 &sampleweight,
6292 &x_entry,
6293 &x_exit,
6294 &x_derivative,
6295 );
6296 inputs.monotonicity_constraint_rows = Some(x_derivative.view());
6297 inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
6298
6299 let model = survival_model(
6300 inputs,
6301 PenaltyBlocks::new(Vec::new()),
6302 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
6303 SurvivalSpec::Net,
6304 )
6305 .expect("construct mixed survival model");
6306
6307 let constraints = model
6308 .monotonicity_linear_constraints()
6309 .expect("event row should induce positive lower bound");
6310 assert_eq!(constraints.a.nrows(), 1);
6311 assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
6312 assert!((constraints.b[0] - 4e-8).abs() <= 1e-18);
6313 }
6314
6315 #[test]
6316 fn structural_monotonicity_clamps_tiny_negative_roundoff() {
6317 let age_entry = array![1.0_f64];
6318 let age_exit = array![2.0_f64];
6319 let event_target = array![1u8];
6320 let event_competing = array![0u8];
6321 let sampleweight = array![1.0];
6322 let x_entry = array![[0.0]];
6323 let x_exit = array![[0.0]];
6324 let x_derivative = array![[1.0]];
6325 let mut model = survival_model(
6326 survival_inputs(
6327 &age_entry,
6328 &age_exit,
6329 &event_target,
6330 &event_competing,
6331 &sampleweight,
6332 &x_entry,
6333 &x_exit,
6334 &x_derivative,
6335 ),
6336 PenaltyBlocks::new(Vec::new()),
6337 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
6338 SurvivalSpec::Net,
6339 )
6340 .expect("construct survival model");
6341 model
6342 .set_structural_monotonicity(true, 1)
6343 .expect("enable structural monotonicity");
6344
6345 let state = model
6346 .update_state(&array![-1e-8])
6347 .expect("tiny structural roundoff should be clamped");
6348 let expected_deviance = -2.0 * (1.0e-12_f64).ln();
6349 assert!(
6350 (state.deviance - expected_deviance).abs() <= 1e-12,
6351 "floored structural event deviance: expected {expected_deviance}, got {}",
6352 state.deviance
6353 );
6354 assert_eq!(state.gradient, array![0.0]);
6355 }
6356
6357 #[test]
6358 fn compressed_monotonicity_constraints_preserve_uncompressed_feasible_region() {
6359 let uncompressed_constraints = LinearInequalityConstraints {
6360 a: array![
6361 [0.0, 0.5, 0.0],
6362 [0.0, 1.0 / 3.0, 0.0],
6363 [0.0, 0.2, 0.0],
6364 [0.0, 0.125, 0.0]
6365 ],
6366 b: Array1::from_elem(4, 1e-8),
6367 };
6368 let compressed_constraints = compress_positive_collinear_constraints(
6369 &uncompressed_constraints.a,
6370 &uncompressed_constraints.b,
6371 );
6372
6373 let candidates = [
6374 array![0.0, 1e-9, 0.0],
6375 array![0.0, 4e-8, 0.0],
6376 array![0.0, 8e-8, 0.0],
6377 array![0.0, 2e-7, 1.5],
6378 ];
6379 for beta in candidates {
6380 let uncompressed_ok = (0..uncompressed_constraints.a.nrows()).all(|i| {
6381 uncompressed_constraints.a.row(i).dot(&beta) >= uncompressed_constraints.b[i]
6382 });
6383 let compressed_ok = (0..compressed_constraints.a.nrows())
6384 .all(|i| compressed_constraints.a.row(i).dot(&beta) >= compressed_constraints.b[i]);
6385 assert_eq!(compressed_ok, uncompressed_ok);
6386 }
6387 }
6388
6389 #[test]
6390 fn exact_survival_derivatives_are_time_unit_invariant_up_to_constant_shift() {
6391 let age_entry = array![10.0_f64, 20.0, 25.0];
6392 let age_exit = array![15.0_f64, 30.0, 40.0];
6393 let event_target = array![1u8, 0u8, 1u8];
6394 let event_competing = array![0u8, 0u8, 0u8];
6395 let sampleweight = array![1.0, 2.0, 0.5];
6396 let x_entry = array![[0.1, 0.2, 1.0], [0.3, 0.4, 1.0], [0.2, 0.6, 1.0]];
6397 let x_exit = array![[0.2, 0.3, 1.0], [0.5, 0.7, 1.0], [0.4, 0.8, 1.0]];
6398 let x_derivative = array![[0.04, 0.02, 0.0], [0.03, 0.01, 0.0], [0.02, 0.03, 0.0]];
6399 let beta = array![0.8, 1.1, -0.2];
6400
6401 let base_model = survival_model(
6402 survival_inputs(
6403 &age_entry,
6404 &age_exit,
6405 &event_target,
6406 &event_competing,
6407 &sampleweight,
6408 &x_entry,
6409 &x_exit,
6410 &x_derivative,
6411 ),
6412 PenaltyBlocks::new(Vec::new()),
6413 SurvivalMonotonicityPenalty { tolerance: 0.0 },
6414 SurvivalSpec::Net,
6415 )
6416 .expect("construct base survival model");
6417 let base_state = base_model
6418 .update_state(&beta)
6419 .expect("evaluate base survival state");
6420
6421 let time_scale = 365.25;
6422 let scaled_age_entry = age_entry.mapv(|v| v * time_scale);
6423 let scaled_age_exit = age_exit.mapv(|v| v * time_scale);
6424 let scaled_x_derivative = x_derivative.mapv(|v| v / time_scale);
6425 let scaled_model = survival_model(
6426 survival_inputs(
6427 &scaled_age_entry,
6428 &scaled_age_exit,
6429 &event_target,
6430 &event_competing,
6431 &sampleweight,
6432 &x_entry,
6433 &x_exit,
6434 &scaled_x_derivative,
6435 ),
6436 PenaltyBlocks::new(Vec::new()),
6437 SurvivalMonotonicityPenalty { tolerance: 0.0 },
6438 SurvivalSpec::Net,
6439 )
6440 .expect("construct scaled survival model");
6441 let scaled_state = scaled_model
6442 .update_state(&beta)
6443 .expect("evaluate scaled survival state");
6444
6445 let weighted_events = sampleweight
6446 .iter()
6447 .zip(event_target.iter())
6448 .map(|(w, d)| *w * f64::from(*d))
6449 .sum::<f64>();
6450 let expected_deviance_shift = 2.0 * weighted_events * time_scale.ln();
6451 assert!(
6452 (scaled_state.deviance - base_state.deviance - expected_deviance_shift).abs() <= 1e-10,
6453 "deviance shift mismatch: scaled={} base={} expected_shift={expected_deviance_shift}",
6454 scaled_state.deviance,
6455 base_state.deviance
6456 );
6457
6458 for j in 0..beta.len() {
6459 assert!(
6460 (scaled_state.gradient[j] - base_state.gradient[j]).abs() <= 1e-12,
6461 "gradient mismatch at j={j}: scaled={} base={}",
6462 scaled_state.gradient[j],
6463 base_state.gradient[j]
6464 );
6465 }
6466
6467 let base_hessian = base_state.hessian.to_dense();
6468 let scaled_hessian = scaled_state.hessian.to_dense();
6469 for r in 0..beta.len() {
6470 for c in 0..beta.len() {
6471 assert!(
6472 (scaled_hessian[[r, c]] - base_hessian[[r, c]]).abs() <= 1e-12,
6473 "hessian mismatch at ({r},{c}): scaled={} base={}",
6474 scaled_hessian[[r, c]],
6475 base_hessian[[r, c]]
6476 );
6477 }
6478 }
6479 }
6480
6481 #[test]
6482 fn survival_laml_rho_gradient_invariant_under_injected_orthogonal_frame_at_the_rail() {
6483 use gam_linalg::faer_ndarray::FaerEigh;
6484 use gam_problem::{EvalMode, PseudoLogdetMode};
6485 use gam_solve::estimate::reml::assembly::InnerAssembly;
6486 use gam_solve::estimate::reml::reml_outer_engine::{
6487 DenseSpectralOperator, DispersionHandling,
6488 };
6489 use gam_solve::estimate::reml::reparameterized_inner::{
6490 RawInnerReparamContext, assemble_reparameterized_inner,
6491 };
6492 use gam_terms::construction::{
6493 canonicalize_penalty_specs, precompute_reparam_invariant_from_canonical,
6494 stable_reparameterizationwith_invariant,
6495 };
6496 use gam_terms::penalty_spec::PenaltySpec;
6497
6498 const RAIL_RHO0: f64 = 7.394829814011909;
6499 const FREE_RHO1: f64 = -2.45;
6500 const DECISION_MARGIN: f64 = 1.0e-9;
6504
6505 let beta0 = array![-2.5_f64, 1.0];
6506 let rho = array![RAIL_RHO0, FREE_RHO1];
6507 let model = laml_rail_fd_test_model(RAIL_RHO0.exp(), FREE_RHO1.exp());
6508
6509 let (rail_model, beta_hat) = model
6511 .reconverge_survival_inner_mode(rho.as_slice().expect("contiguous rho"), &beta0)
6512 .expect("reconverge inner mode at the rail");
6513 let state = rail_model
6514 .update_state(&beta_hat)
6515 .expect("inner state at the rail");
6516 let p = beta_hat.len();
6517 let h_dense = state.hessian.to_dense();
6518 let lambdas: Vec<f64> = rho.iter().map(|&r| r.exp()).collect();
6519
6520 let active_blocks: Vec<&PenaltyBlock> = rail_model
6521 .penalties
6522 .blocks
6523 .iter()
6524 .filter(|b| b.lambda > 0.0)
6525 .collect();
6526 let s_k_embedded: Vec<Array2<f64>> = active_blocks
6527 .iter()
6528 .map(|b| {
6529 let mut s = Array2::<f64>::zeros((p, p));
6530 let (rs, re) = (b.range.start, b.range.end);
6531 s.slice_mut(ndarray::s![rs..re, rs..re]).assign(&b.matrix);
6532 s
6533 })
6534 .collect();
6535 let penalty_specs: Vec<PenaltySpec> = active_blocks
6536 .iter()
6537 .map(|b| PenaltySpec::Block {
6538 local: b.matrix.clone(),
6539 col_range: b.range.clone(),
6540 prior_mean: gam_problem::CoefficientPriorMean::Zero,
6541 structure_hint: None,
6542 op: None,
6543 })
6544 .collect();
6545 let nullspace_dims: Vec<usize> = active_blocks.iter().map(|b| b.nullspace_dim).collect();
6546 let (canonical, _) = canonicalize_penalty_specs(
6547 &penalty_specs,
6548 &nullspace_dims,
6549 p,
6550 "rail-stability gate reparameterization",
6551 )
6552 .expect("canonicalize rail penalties");
6553 let invariant =
6554 precompute_reparam_invariant_from_canonical(&canonical, p).expect("reparam invariant");
6555 let reparam_prod =
6556 stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &invariant, None)
6557 .expect("production reparameterization");
6558
6559 let hessian_logdet_mode = if rail_model
6560 .age_entry
6561 .iter()
6562 .any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD)
6563 {
6564 PseudoLogdetMode::HardPseudo
6565 } else {
6566 PseudoLogdetMode::Smooth
6567 };
6568
6569 let g = {
6572 let sym = Array2::<f64>::from_shape_fn((p, p), |(i, j)| {
6573 (((i * j) as f64) * 0.37).sin() + (((i + j) as f64) * 0.11 + 1.0).cos()
6574 });
6575 let (_, evecs) = sym.eigh(faer::Side::Lower).expect("orthogonal factor from eigh");
6576 evecs
6577 };
6578
6579 let grad_for_reparam = |reparam: &gam_terms::construction::ReparamResult| -> Array1<f64> {
6583 let provider = SurvivalDerivProvider::new(rail_model.clone(), beta_hat.clone());
6584 let ctx = RawInnerReparamContext {
6585 hessian: &h_dense,
6586 beta: &beta_hat,
6587 penalties_embedded: &s_k_embedded,
6588 lambdas: &lambdas,
6589 };
6590 let reparam_inner =
6591 assemble_reparameterized_inner(&ctx, Some(Box::new(provider)), reparam)
6592 .expect("reparameterized inner assembly");
6593 let hop = DenseSpectralOperator::from_symmetric_with_mode(
6594 &reparam_inner.hessian_transformed,
6595 hessian_logdet_mode,
6596 )
6597 .expect("transformed Hessian operator");
6598 let penalty_coords = reparam
6599 .canonical_transformed
6600 .iter()
6601 .map(|cp| cp.to_penalty_coordinate())
6602 .collect::<Vec<_>>();
6603 let result = InnerAssembly {
6604 log_likelihood: state.log_likelihood,
6605 penalty_quadratic: state.penalty_term,
6606 beta: reparam_inner.beta_transformed,
6607 n_observations: rail_model.nrows(),
6608 hessian_op: std::sync::Arc::new(hop),
6609 penalty_coords,
6610 penalty_logdet: reparam_inner.penalty_logdet,
6611 dispersion: DispersionHandling::Fixed {
6612 phi: 1.0,
6613 include_logdet_h: true,
6614 include_logdet_s: true,
6615 },
6616 rho_curvature_scale: 1.0,
6617 rho_prior: gam_problem::RhoPrior::Flat,
6618 hessian_logdet_correction: 0.0,
6619 penalty_subspace_trace: None,
6620 deriv_provider: reparam_inner.deriv_provider,
6621 firth: None,
6622 nullspace_dim: None,
6623 barrier_config: None,
6624 ext_coords: Vec::new(),
6625 ext_coord_pair_fn: None,
6626 rho_ext_pair_fn: None,
6627 fixed_drift_deriv: None,
6628 contracted_psi_second_order: None,
6629 kkt_residual: None,
6630 active_constraints: None,
6631 }
6632 .evaluate(
6633 rho.as_slice().expect("contiguous rho"),
6634 EvalMode::ValueAndGradient,
6635 None,
6636 )
6637 .expect("transformed-frame LAML evaluate");
6638 result.gradient.expect("analytic ρ-gradient present")
6639 };
6640
6641 let g_prod = grad_for_reparam(&reparam_prod);
6643
6644 let mut reparam_conj = reparam_prod.clone();
6648 reparam_conj.qs = reparam_prod.qs.dot(&g);
6649 reparam_conj.canonical_transformed = reparam_prod
6650 .canonical_transformed
6651 .iter()
6652 .map(|cp| {
6653 let mut rotated = cp.clone();
6654 rotated.root = cp.root.dot(&g);
6655 rotated.local = rotated.root.t().dot(&rotated.root);
6656 rotated
6657 })
6658 .collect();
6659 let g_conj = grad_for_reparam(&reparam_conj);
6660
6661 for k in 0..rho.len() {
6662 let drift = (g_prod[k] - g_conj[k]).abs();
6663 assert!(
6664 drift <= DECISION_MARGIN * (1.0 + g_prod[k].abs()),
6665 "rail ρ-gradient not frame-invariant at coordinate {k}: \
6666 Q_s frame {} vs Q_s·G frame {} (drift {:.3e} > margin {:.3e})",
6667 g_prod[k],
6668 g_conj[k],
6669 drift,
6670 DECISION_MARGIN * (1.0 + g_prod[k].abs())
6671 );
6672 }
6673 }
6674
6675}