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