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_solve::pirls::{
10 LinearInequalityConstraints, WorkingModel as PirlsWorkingModel, WorkingState, array1_l2_norm,
11};
12use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayView3, Axis};
13use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, constants, escalate_ridge};
14use serde::{Deserialize, Serialize};
15use std::collections::BTreeMap;
16use std::convert::Infallible;
17use std::ops::Range;
18use std::sync::LazyLock;
19use thiserror::Error;
20
21#[derive(Debug, Error)]
22pub enum SurvivalError {
23 #[error("input dimensions are inconsistent")]
24 DimensionMismatch,
25 #[error("inputs contain non-finite values")]
26 NonFiniteInput,
27 #[error("survival spec '{0}' is not supported by the one-hazard survival engine")]
28 UnsupportedSpec(&'static str),
29 #[error("crude risk integration setup is invalid")]
30 InvalidIntegrationSetup,
31 #[error("survival time grid must be finite, non-negative, and strictly increasing")]
32 InvalidTimeGrid,
33 #[error("cumulative hazard must be nondecreasing")]
34 NonMonotoneCumulativeHazard,
35 #[error("instantaneous hazard must stay strictly positive during integration")]
36 NonPositiveHazard,
37 #[error("{reason}")]
38 InvalidInput { reason: String },
39 #[error("{reason}")]
40 CauseSpecificDimensionMismatch { reason: String },
41 #[error("{reason}")]
42 NumericalFailure { reason: String },
43 #[error("{reason}")]
44 EventCodeInvalid { reason: String },
45 #[error("{reason}")]
46 EventDegenerate { reason: String },
47 #[error("cause-specific survival block {block}: {source}")]
48 CauseSpecificBlock {
49 block: usize,
50 #[source]
51 source: Box<SurvivalError>,
52 },
53}
54
55impl From<SurvivalError> for String {
56 fn from(err: SurvivalError) -> Self {
57 err.to_string()
58 }
59}
60
61impl From<crate::block_layout::block_count::BlockCountMismatch> for SurvivalError {
62 fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> SurvivalError {
63 SurvivalError::CauseSpecificDimensionMismatch {
64 reason: err.message(),
65 }
66 }
67}
68
69#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
70pub enum SurvivalSpec {
71 #[default]
72 Net,
73 Crude,
74}
75
76#[derive(Debug, Clone)]
77pub struct SurvivalEngineInputs<'a> {
78 pub age_entry: ArrayView1<'a, f64>,
79 pub age_exit: ArrayView1<'a, f64>,
80 pub event_target: ArrayView1<'a, u8>,
81 pub event_competing: ArrayView1<'a, u8>,
82 pub sampleweight: ArrayView1<'a, f64>,
83 pub x_entry: ArrayView2<'a, f64>,
84 pub x_exit: ArrayView2<'a, f64>,
85 pub x_derivative: ArrayView2<'a, f64>,
86 pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
90 pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
92}
93
94#[derive(Debug, Clone)]
95pub struct SurvivalTimeCovarInputs<'a> {
96 pub age_entry: ArrayView1<'a, f64>,
97 pub age_exit: ArrayView1<'a, f64>,
98 pub event_target: ArrayView1<'a, u8>,
99 pub event_competing: ArrayView1<'a, u8>,
100 pub sampleweight: ArrayView1<'a, f64>,
101 pub time_entry: ArrayView2<'a, f64>,
102 pub time_exit: ArrayView2<'a, f64>,
103 pub time_derivative: ArrayView2<'a, f64>,
104 pub covariates: ArrayView2<'a, f64>,
105 pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
109 pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
111}
112
113#[derive(Debug, Clone)]
114pub struct SurvivalBaselineOffsets<'a> {
115 pub eta_entry: ArrayView1<'a, f64>,
117 pub eta_exit: ArrayView1<'a, f64>,
119 pub derivative_exit: ArrayView1<'a, f64>,
127}
128
129#[derive(Debug, Clone)]
130pub struct PenaltyBlock {
131 pub matrix: Array2<f64>,
132 pub lambda: f64,
133 pub range: Range<usize>,
134 pub nullspace_dim: usize,
137}
138
139#[derive(Debug, Clone)]
140pub struct PenaltyBlocks {
141 pub blocks: Vec<PenaltyBlock>,
142}
143
144impl PenaltyBlocks {
145 pub fn new(blocks: Vec<PenaltyBlock>) -> Self {
146 Self { blocks }
147 }
148
149 pub fn gradient(&self, beta: &Array1<f64>) -> Array1<f64> {
150 let mut grad = Array1::zeros(beta.len());
151 for block in &self.blocks {
152 if block.lambda == 0.0 {
153 continue;
154 }
155 let b = beta.slice(ndarray::s![block.range.clone()]);
156 let g = block.matrix.dot(&b);
157 let mut dst = grad.slice_mut(ndarray::s![block.range.clone()]);
158 dst += &(block.lambda * g);
159 }
160 grad
161 }
162
163 pub fn hessian(&self, dim: usize) -> Array2<f64> {
164 let mut h = Array2::zeros((dim, dim));
165 self.addhessian_inplace(&mut h);
166 h
167 }
168
169 pub fn deviance(&self, beta: &Array1<f64>) -> f64 {
170 let mut value = 0.0;
171 for block in &self.blocks {
172 if block.lambda == 0.0 {
173 continue;
174 }
175 let b = beta.slice(ndarray::s![block.range.clone()]);
176 value += 0.5 * block.lambda * b.dot(&block.matrix.dot(&b));
177 }
178 value
179 }
180
181 pub fn addhessian_inplace(&self, h: &mut Array2<f64>) {
182 for block in &self.blocks {
183 if block.lambda == 0.0 {
184 continue;
185 }
186 let start = block.range.start;
187 let end = block.range.end;
188 h.slice_mut(ndarray::s![start..end, start..end])
189 .scaled_add(block.lambda, &block.matrix);
190 }
191 }
192}
193
194pub const ENTRY_AT_ORIGIN_THRESHOLD: f64 = 1e-8;
207
208const DERIVATIVE_FRACTION_TO_BOUNDARY: f64 = 0.995;
216
217#[derive(Debug, Clone)]
218pub struct CauseSpecificRoystonParmarBlock {
219 pub age_entry: Array1<f64>,
220 pub age_exit: Array1<f64>,
221 pub event_target: Array1<u8>,
222 pub sampleweight: Array1<f64>,
223 pub x_entry: Array2<f64>,
224 pub x_exit: Array2<f64>,
225 pub x_derivative: Array2<f64>,
226 pub offset_eta_entry: Array1<f64>,
227 pub offset_eta_exit: Array1<f64>,
228 pub offset_derivative_exit: Array1<f64>,
229 pub derivative_floor: f64,
230}
231
232#[derive(Debug, Clone)]
238pub struct CauseSpecificRoystonParmarFamily {
239 blocks: Vec<CauseSpecificRoystonParmarBlock>,
240}
241
242impl CauseSpecificRoystonParmarFamily {
243 pub fn new(blocks: Vec<CauseSpecificRoystonParmarBlock>) -> Result<Self, String> {
244 if blocks.is_empty() {
245 return Err(SurvivalError::InvalidInput {
246 reason: "cause-specific survival family requires at least one endpoint".to_string(),
247 }
248 .into());
249 }
250 for (idx, block) in blocks.iter().enumerate() {
251 validate_cause_specific_block(block).map_err(|err| {
252 SurvivalError::CauseSpecificBlock {
253 block: idx + 1,
254 source: Box::new(err),
255 }
256 .to_string()
257 })?;
258 }
259 Ok(Self { blocks })
260 }
261
262 pub fn cause_count(&self) -> usize {
263 self.blocks.len()
264 }
265}
266
267fn validate_cause_specific_block(
268 block: &CauseSpecificRoystonParmarBlock,
269) -> Result<(), SurvivalError> {
270 let n = block.event_target.len();
271 let p = block.x_exit.ncols();
272 if n == 0 || p == 0 {
273 bail_invalid_surv!("empty event vector or coefficient block");
274 }
275 if block.age_entry.len() != n
276 || block.age_exit.len() != n
277 || block.sampleweight.len() != n
278 || block.x_entry.nrows() != n
279 || block.x_exit.nrows() != n
280 || block.x_derivative.nrows() != n
281 || block.x_entry.ncols() != p
282 || block.x_derivative.ncols() != p
283 || block.offset_eta_entry.len() != n
284 || block.offset_eta_exit.len() != n
285 || block.offset_derivative_exit.len() != n
286 {
287 return Err(SurvivalError::CauseSpecificDimensionMismatch {
288 reason: "dimension mismatch".to_string(),
289 });
290 }
291 if let Some(&label) = block.event_target.iter().find(|&&v| v > 1) {
297 return Err(SurvivalError::EventCodeInvalid {
298 reason: format!(
299 "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"
300 ),
301 });
302 }
303 if block.age_entry.iter().any(|v| !v.is_finite())
304 || block.age_exit.iter().any(|v| !v.is_finite())
305 || block
306 .sampleweight
307 .iter()
308 .any(|v| !v.is_finite() || *v < 0.0)
309 || block.x_entry.iter().any(|v| !v.is_finite())
310 || block.x_exit.iter().any(|v| !v.is_finite())
311 || block.x_derivative.iter().any(|v| !v.is_finite())
312 || block.offset_eta_entry.iter().any(|v| !v.is_finite())
313 || block.offset_eta_exit.iter().any(|v| !v.is_finite())
314 || block.offset_derivative_exit.iter().any(|v| !v.is_finite())
315 || !block.derivative_floor.is_finite()
316 || block.derivative_floor < 0.0
317 {
318 bail_invalid_surv!("non-finite input");
319 }
320 Ok(())
321}
322
323fn evaluate_cause_specific_block(
324 block: &CauseSpecificRoystonParmarBlock,
325 beta: &Array1<f64>,
326) -> Result<(f64, Array1<f64>, Array2<f64>), SurvivalError> {
327 let n = block.event_target.len();
328 let p = block.x_exit.ncols();
329 if beta.len() != p {
330 return Err(SurvivalError::CauseSpecificDimensionMismatch {
331 reason: format!("beta length mismatch: got {}, expected {p}", beta.len()),
332 });
333 }
334 let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
335 let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
336 let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
337 let mut log_likelihood = 0.0;
338 let mut w_exit = Array1::<f64>::zeros(n);
339 let mut w_entry = Array1::<f64>::zeros(n);
340 let mut w_event = Array1::<f64>::zeros(n);
341 let mut w_event_inv_deriv = Array1::<f64>::zeros(n);
342 let mut w_event_outer = Array1::<f64>::zeros(n);
343
344 for i in 0..n {
345 let weight = block.sampleweight[i];
346 if weight <= 0.0 {
347 continue;
348 }
349 if block.age_exit[i] < block.age_entry[i] {
350 bail_invalid_surv!("age_exit < age_entry at row {i}");
351 }
352 let has_entry = block.age_entry[i] > ENTRY_AT_ORIGIN_THRESHOLD;
353 let h_exit = eta_exit[i].exp();
354 let h_entry = if has_entry { eta_entry[i].exp() } else { 0.0 };
355 if !(h_exit.is_finite() && h_entry.is_finite()) {
356 return Err(SurvivalError::NumericalFailure {
357 reason: format!("non-finite cumulative hazard at row {i}"),
358 });
359 }
360 log_likelihood -= weight * (h_exit - h_entry);
361 w_exit[i] = weight * h_exit;
362 w_entry[i] = weight * h_entry;
363 if block.event_target[i] > 0 {
364 let deriv = derivative[i];
365 if !(deriv.is_finite() && deriv > 0.0) {
366 return Err(SurvivalError::NumericalFailure {
367 reason: format!(
368 "cause-specific survival derivative must be positive at row {i}, got {deriv}"
369 ),
370 });
371 }
372 log_likelihood += weight * (eta_exit[i] + deriv.ln());
373 w_event[i] = weight;
374 w_event_inv_deriv[i] = weight / deriv;
375 w_event_outer[i] = weight / (deriv * deriv);
376 }
377 }
378
379 let mut nll_gradient = fast_atv(&block.x_exit, &w_exit);
380 nll_gradient -= &fast_atv(&block.x_entry, &w_entry);
381 nll_gradient -= &fast_atv(&block.x_exit, &w_event);
382 nll_gradient -= &fast_atv(&block.x_derivative, &w_event_inv_deriv);
383 let gradient = -nll_gradient;
384
385 let mut hessian = fast_xt_diag_x(&block.x_exit, &w_exit);
386 hessian -= &fast_xt_diag_x(&block.x_entry, &w_entry);
387 hessian += &fast_xt_diag_x(&block.x_derivative, &w_event_outer);
388 Ok((log_likelihood, gradient, hessian))
389}
390
391impl CustomFamily for CauseSpecificRoystonParmarFamily {
392 fn joint_jeffreys_term_required(&self) -> bool {
396 true
397 }
398
399 fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
400 crate::block_layout::block_count::validate_block_count::<SurvivalError>(
401 "cause-specific survival",
402 self.blocks.len(),
403 block_states.len(),
404 )?;
405 let mut log_likelihood = 0.0;
406 let mut blockworking_sets = Vec::with_capacity(self.blocks.len());
407 for (block, state) in self.blocks.iter().zip(block_states.iter()) {
408 let (ll, gradient, hessian) = evaluate_cause_specific_block(block, &state.beta)?;
409 log_likelihood += ll;
410 blockworking_sets.push(BlockWorkingSet::ExactNewton {
411 gradient,
412 hessian: SymmetricMatrix::Dense(hessian),
413 });
414 }
415 Ok(FamilyEvaluation {
416 log_likelihood,
417 blockworking_sets,
418 })
419 }
420
421 fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
422 crate::block_layout::block_count::validate_block_count::<SurvivalError>(
423 "cause-specific survival",
424 self.blocks.len(),
425 block_states.len(),
426 )?;
427 let mut log_likelihood = 0.0;
428 for (block, state) in self.blocks.iter().zip(block_states.iter()) {
429 let (ll, _, _) = evaluate_cause_specific_block(block, &state.beta)?;
430 log_likelihood += ll;
431 }
432 Ok(log_likelihood)
433 }
434
435 fn likelihood_blocks_uncoupled(&self) -> bool {
436 true
437 }
438
439 fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
440 true
441 }
442
443 fn output_channel_assignment(
444 &self,
445 specs: &[crate::custom_family::ParameterBlockSpec],
446 ) -> Option<Vec<usize>> {
447 if specs.len() != self.blocks.len() {
448 return Some((0..self.blocks.len()).collect());
449 }
450 Some((0..specs.len()).collect())
451 }
452
453 fn coefficient_hessian_cost(&self, specs: &[crate::custom_family::ParameterBlockSpec]) -> u64 {
454 crate::custom_family::default_coefficient_hessian_cost(specs)
455 }
456
457 fn block_linear_constraints(
458 &self,
459 _: &[ParameterBlockState],
460 block_idx: usize,
461 spec: &crate::custom_family::ParameterBlockSpec,
462 ) -> Result<Option<LinearInequalityConstraints>, String> {
463 let block = self.blocks.get(block_idx).ok_or_else(|| {
464 SurvivalError::CauseSpecificDimensionMismatch {
465 reason: format!(
466 "cause-specific survival expected block index < {}, got {block_idx}",
467 self.blocks.len()
468 ),
469 }
470 .to_string()
471 })?;
472 if block.x_derivative.ncols() != spec.design.ncols() {
473 return Err(SurvivalError::CauseSpecificDimensionMismatch {
474 reason: format!(
475 "cause-specific survival derivative design has {} columns but block '{}' has {}",
476 block.x_derivative.ncols(),
477 spec.name,
478 spec.design.ncols()
479 ),
480 }
481 .into());
482 }
483 let rhs = block
484 .offset_derivative_exit
485 .mapv(|offset| block.derivative_floor - offset);
486 Ok(Some(LinearInequalityConstraints {
487 a: block.x_derivative.clone(),
488 b: rhs,
489 }))
490 }
491
492 fn max_feasible_step_size(
493 &self,
494 block_states: &[ParameterBlockState],
495 block_idx: usize,
496 delta: &Array1<f64>,
497 ) -> Result<Option<f64>, String> {
498 let block = self.blocks.get(block_idx).ok_or_else(|| {
499 SurvivalError::CauseSpecificDimensionMismatch {
500 reason: format!(
501 "cause-specific survival expected block index < {}, got {block_idx}",
502 self.blocks.len()
503 ),
504 }
505 .to_string()
506 })?;
507 let state = block_states.get(block_idx).ok_or_else(|| {
508 SurvivalError::CauseSpecificDimensionMismatch {
509 reason: format!(
510 "cause-specific survival expected {} block states, got {}",
511 self.blocks.len(),
512 block_states.len()
513 ),
514 }
515 .to_string()
516 })?;
517 if delta.len() != state.beta.len() || block.x_derivative.ncols() != delta.len() {
518 return Err(SurvivalError::CauseSpecificDimensionMismatch {
519 reason: "cause-specific survival feasible-step dimension mismatch".to_string(),
520 }
521 .into());
522 }
523 let derivative = fast_av(&block.x_derivative, &state.beta) + &block.offset_derivative_exit;
524 let derivative_delta = fast_av(&block.x_derivative, delta);
525 let mut alpha_max = 1.0_f64;
526 for i in 0..derivative.len() {
527 if block.sampleweight[i] <= 0.0 {
528 continue;
529 }
530 let current = derivative[i] - block.derivative_floor;
531 let slope = derivative_delta[i];
532 if slope < 0.0 {
533 if current <= 0.0 {
534 return Ok(Some(0.0));
535 }
536 alpha_max = alpha_max.min(DERIVATIVE_FRACTION_TO_BOUNDARY * current / -slope);
537 }
538 }
539 Ok(Some(alpha_max.clamp(0.0, 1.0)))
540 }
541
542 fn exact_newton_hessian_directional_derivative(
543 &self,
544 block_states: &[ParameterBlockState],
545 block_idx: usize,
546 d_beta: &Array1<f64>,
547 ) -> Result<Option<Array2<f64>>, String> {
548 let block = self.blocks.get(block_idx).ok_or_else(|| {
549 SurvivalError::CauseSpecificDimensionMismatch {
550 reason: format!(
551 "cause-specific survival expected block index < {}, got {block_idx}",
552 self.blocks.len()
553 ),
554 }
555 .to_string()
556 })?;
557 let state = block_states.get(block_idx).ok_or_else(|| {
558 SurvivalError::CauseSpecificDimensionMismatch {
559 reason: format!(
560 "cause-specific survival expected {} block states, got {}",
561 self.blocks.len(),
562 block_states.len()
563 ),
564 }
565 .to_string()
566 })?;
567 Ok(Some(cause_specific_hessian_directional_derivative(
568 block,
569 &state.beta,
570 d_beta,
571 )?))
572 }
573
574 fn exact_newton_hessian_second_directional_derivative(
575 &self,
576 block_states: &[ParameterBlockState],
577 block_idx: usize,
578 d_beta_u: &Array1<f64>,
579 d_beta_v: &Array1<f64>,
580 ) -> Result<Option<Array2<f64>>, String> {
581 let block = self.blocks.get(block_idx).ok_or_else(|| {
582 SurvivalError::CauseSpecificDimensionMismatch {
583 reason: format!(
584 "cause-specific survival expected block index < {}, got {block_idx}",
585 self.blocks.len()
586 ),
587 }
588 .to_string()
589 })?;
590 let state = block_states.get(block_idx).ok_or_else(|| {
591 SurvivalError::CauseSpecificDimensionMismatch {
592 reason: format!(
593 "cause-specific survival expected {} block states, got {}",
594 self.blocks.len(),
595 block_states.len()
596 ),
597 }
598 .to_string()
599 })?;
600 Ok(Some(cause_specific_hessian_second_directional_derivative(
601 block,
602 &state.beta,
603 d_beta_u,
604 d_beta_v,
605 )?))
606 }
607}
608
609fn cause_specific_hessian_directional_derivative(
621 block: &CauseSpecificRoystonParmarBlock,
622 beta: &Array1<f64>,
623 d_beta: &Array1<f64>,
624) -> Result<Array2<f64>, SurvivalError> {
625 let p = block.x_exit.ncols();
626 if beta.len() != p || d_beta.len() != p {
627 return Err(SurvivalError::CauseSpecificDimensionMismatch {
628 reason: "cause-specific survival Hessian derivative dimension mismatch".to_string(),
629 });
630 }
631 let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
632 let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
633 let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
634 let d_eta_entry = fast_av(&block.x_entry, d_beta);
635 let d_eta_exit = fast_av(&block.x_exit, d_beta);
636 let d_derivative = fast_av(&block.x_derivative, d_beta);
637 let mut w_exit = Array1::<f64>::zeros(block.event_target.len());
638 let mut w_entry = Array1::<f64>::zeros(block.event_target.len());
639 let mut w_derivative = Array1::<f64>::zeros(block.event_target.len());
640
641 for i in 0..block.event_target.len() {
642 let weight = block.sampleweight[i];
643 if weight <= 0.0 {
644 continue;
645 }
646 let has_entry = block.age_entry[i] > ENTRY_AT_ORIGIN_THRESHOLD;
647 w_exit[i] = weight * eta_exit[i].exp() * d_eta_exit[i];
648 if has_entry {
649 w_entry[i] = weight * eta_entry[i].exp() * d_eta_entry[i];
650 }
651 if block.event_target[i] > 0 {
652 let deriv = derivative[i];
653 if !(deriv.is_finite() && deriv > 0.0) {
654 return Err(SurvivalError::NumericalFailure {
655 reason: format!(
656 "cause-specific survival derivative must be positive at row {i}, got {deriv}"
657 ),
658 });
659 }
660 w_derivative[i] = -2.0 * weight * d_derivative[i] / (deriv * deriv * deriv);
661 }
662 }
663
664 let mut d_hessian = fast_xt_diag_x(&block.x_exit, &w_exit);
665 d_hessian -= &fast_xt_diag_x(&block.x_entry, &w_entry);
666 d_hessian += &fast_xt_diag_x(&block.x_derivative, &w_derivative);
667 Ok(d_hessian)
668}
669
670fn cause_specific_hessian_second_directional_derivative(
671 block: &CauseSpecificRoystonParmarBlock,
672 beta: &Array1<f64>,
673 d_beta_u: &Array1<f64>,
674 d_beta_v: &Array1<f64>,
675) -> Result<Array2<f64>, SurvivalError> {
676 let p = block.x_exit.ncols();
677 if beta.len() != p || d_beta_u.len() != p || d_beta_v.len() != p {
678 return Err(SurvivalError::CauseSpecificDimensionMismatch {
679 reason: "cause-specific survival second Hessian derivative dimension mismatch"
680 .to_string(),
681 });
682 }
683 let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
684 let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
685 let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
686 let u_eta_entry = fast_av(&block.x_entry, d_beta_u);
687 let u_eta_exit = fast_av(&block.x_exit, d_beta_u);
688 let u_derivative = fast_av(&block.x_derivative, d_beta_u);
689 let v_eta_entry = fast_av(&block.x_entry, d_beta_v);
690 let v_eta_exit = fast_av(&block.x_exit, d_beta_v);
691 let v_derivative = fast_av(&block.x_derivative, d_beta_v);
692 let mut w_exit = Array1::<f64>::zeros(block.event_target.len());
693 let mut w_entry = Array1::<f64>::zeros(block.event_target.len());
694 let mut w_derivative = Array1::<f64>::zeros(block.event_target.len());
695
696 for i in 0..block.event_target.len() {
697 let weight = block.sampleweight[i];
698 if weight <= 0.0 {
699 continue;
700 }
701 let has_entry = block.age_entry[i] > ENTRY_AT_ORIGIN_THRESHOLD;
702 w_exit[i] = weight * eta_exit[i].exp() * u_eta_exit[i] * v_eta_exit[i];
703 if has_entry {
704 w_entry[i] = weight * eta_entry[i].exp() * u_eta_entry[i] * v_eta_entry[i];
705 }
706 if block.event_target[i] > 0 {
707 let deriv = derivative[i];
708 if !(deriv.is_finite() && deriv > 0.0) {
709 return Err(SurvivalError::NumericalFailure {
710 reason: format!(
711 "cause-specific survival derivative must be positive at row {i}, got {deriv}"
712 ),
713 });
714 }
715 w_derivative[i] = 6.0 * weight * u_derivative[i] * v_derivative[i] / deriv.powi(4);
716 }
717 }
718
719 let mut d2_hessian = fast_xt_diag_x(&block.x_exit, &w_exit);
720 d2_hessian -= &fast_xt_diag_x(&block.x_entry, &w_entry);
721 d2_hessian += &fast_xt_diag_x(&block.x_derivative, &w_derivative);
722 Ok(d2_hessian)
723}
724
725pub fn survival_event_code_from_value(value: f64, row_index: usize) -> Result<u8, String> {
726 const INTEGER_TOL: f64 = 1e-8;
727 const MAX_AUTO_CAUSES: u8 = 32;
728 if !value.is_finite() {
729 return Err(SurvivalError::EventCodeInvalid {
730 reason: format!(
731 "survival event value at row {} is non-finite",
732 row_index + 1
733 ),
734 }
735 .into());
736 }
737 if value < 0.0 {
738 return Err(SurvivalError::EventCodeInvalid {
739 reason: format!(
740 "survival event value at row {} is negative: {value}",
741 row_index + 1
742 ),
743 }
744 .into());
745 }
746 let rounded = value.round();
747 if (value - rounded).abs() > INTEGER_TOL {
748 return Err(SurvivalError::EventCodeInvalid {
749 reason: format!(
750 "survival event value at row {} must be an integer code with 0=censored, got {value}",
751 row_index + 1
752 ),
753 }
754 .into());
755 }
756 if rounded > f64::from(MAX_AUTO_CAUSES) {
757 return Err(SurvivalError::EventCodeInvalid {
758 reason: format!(
759 "survival event value at row {} has code {rounded}; automatic competing-risks detection supports codes 0..={MAX_AUTO_CAUSES}",
760 row_index + 1
761 ),
762 }
763 .into());
764 }
765 Ok(rounded as u8)
766}
767
768pub fn cause_count_from_event_codes(
769 event_codes: ArrayView1<'_, u8>,
770) -> Result<usize, SurvivalError> {
771 let max_code = event_codes.iter().copied().max().map_or(0, usize::from);
772 if max_code == 0 {
773 return Ok(1);
774 }
775
776 let mut present = vec![false; max_code + 1];
777 for code in event_codes.iter().copied() {
778 present[usize::from(code)] = true;
779 }
780 if (1..=max_code).any(|code| !present[code]) {
781 let actual = present
782 .iter()
783 .enumerate()
784 .skip(1)
785 .filter_map(|(code, &seen)| seen.then_some(code.to_string()))
786 .collect::<Vec<_>>()
787 .join(", ");
788 return Err(SurvivalError::EventCodeInvalid {
789 reason: format!(
790 "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."
791 ),
792 });
793 }
794
795 Ok(max_code)
796}
797
798pub fn pooled_any_event_indicator(event_codes: ArrayView1<'_, u8>) -> Array1<u8> {
811 event_codes.mapv(|label| u8::from(label > 0))
812}
813
814pub fn cause_specific_event_indicator(event_codes: ArrayView1<'_, u8>, cause: usize) -> Array1<u8> {
824 let cause_code = cause as u8;
825 event_codes.mapv(|observed| u8::from(observed == cause_code))
826}
827
828fn compress_positive_collinear_constraints(
829 a: &Array2<f64>,
830 b: &Array1<f64>,
831) -> LinearInequalityConstraints {
832 const SCALE_TOL: f64 = 1e-14;
833 const KEY_TOL: f64 = 1e-8;
834
835 let mut grouped: BTreeMap<Vec<i64>, (Vec<f64>, f64)> = BTreeMap::new();
836 let mut fallbackrows: Vec<(Vec<f64>, f64)> = Vec::new();
837
838 for i in 0..a.nrows() {
839 let row = a.row(i);
840 let scale = row.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
841 if !scale.is_finite() || scale <= SCALE_TOL {
842 if b[i] > 0.0 {
843 fallbackrows.push((row.to_vec(), b[i]));
844 }
845 continue;
846 }
847
848 let normalizedrow: Vec<f64> = row
849 .iter()
850 .map(|&v| {
851 let scaled = v / scale;
852 if scaled.abs() <= KEY_TOL { 0.0 } else { scaled }
853 })
854 .collect();
855 let normalized_rhs = b[i] / scale;
856 let key: Vec<i64> = normalizedrow
857 .iter()
858 .map(|&v| (v / KEY_TOL).round() as i64)
859 .collect();
860
861 match grouped.get_mut(&key) {
862 Some((_, rhs_max)) => {
863 if normalized_rhs > *rhs_max {
864 *rhs_max = normalized_rhs;
865 }
866 }
867 None => {
868 grouped.insert(key, (normalizedrow, normalized_rhs));
869 }
870 }
871 }
872
873 let nrows = grouped.len() + fallbackrows.len();
874 let n_cols = a.ncols();
875 let mut a_out = Array2::<f64>::zeros((nrows, n_cols));
876 let mut b_out = Array1::<f64>::zeros(nrows);
877
878 let mut outrow = 0usize;
879 for (_, (row, rhs)) in grouped {
880 for (j, value) in row.into_iter().enumerate() {
881 a_out[[outrow, j]] = value;
882 }
883 b_out[outrow] = rhs;
884 outrow += 1;
885 }
886 for (row, rhs) in fallbackrows {
887 for (j, value) in row.into_iter().enumerate() {
888 a_out[[outrow, j]] = value;
889 }
890 b_out[outrow] = rhs;
891 outrow += 1;
892 }
893
894 LinearInequalityConstraints { a: a_out, b: b_out }
895}
896
897#[derive(Debug, Clone, Copy, Default)]
898pub struct SurvivalMonotonicityPenalty {
899 pub tolerance: f64,
900}
901
902#[derive(Debug, Clone)]
903enum SurvivalDesign {
904 Flat {
905 x_entry: Array2<f64>,
906 x_exit: Array2<f64>,
907 x_derivative: Array2<f64>,
908 },
909 TimeCovariateShared {
910 time_entry: Array2<f64>,
911 time_exit: Array2<f64>,
912 time_derivative: Array2<f64>,
913 covariates: Array2<f64>,
914 },
915}
916
917impl SurvivalDesign {
918 fn p_total(&self) -> usize {
919 match self {
920 Self::Flat { x_exit, .. } => x_exit.ncols(),
921 Self::TimeCovariateShared {
922 time_exit,
923 covariates,
924 ..
925 } => time_exit.ncols() + covariates.ncols(),
926 }
927 }
928
929 fn design_dot(&self, time_mat: &Array2<f64>, beta: &Array1<f64>) -> Array1<f64> {
930 match self {
931 Self::Flat { .. } => time_mat.dot(beta),
932 Self::TimeCovariateShared { covariates, .. } => {
933 let p_time = time_mat.ncols();
934 let mut out = time_mat.dot(&beta.slice(ndarray::s![..p_time]));
935 if covariates.ncols() > 0 {
936 out += &covariates.dot(&beta.slice(ndarray::s![p_time..]));
937 }
938 out
939 }
940 }
941 }
942
943 fn fill_row(&self, time_mat: &Array2<f64>, i: usize, out: &mut [f64]) {
944 match self {
945 Self::Flat { .. } => {
946 for (dst, &src) in out.iter_mut().zip(time_mat.row(i).iter()) {
947 *dst = src;
948 }
949 }
950 Self::TimeCovariateShared { covariates, .. } => {
951 let p_time = time_mat.ncols();
952 for j in 0..p_time {
953 out[j] = time_mat[[i, j]];
954 }
955 for j in 0..covariates.ncols() {
956 out[p_time + j] = covariates[[i, j]];
957 }
958 }
959 }
960 }
961}
962
963#[derive(Debug, Clone)]
965struct SurvivalWorkspace {
966 w_event: Array1<f64>,
967 w_event_inv_deriv: Array1<f64>,
968 w_event_outer: Array1<f64>,
969 w_hess_exit: Array1<f64>,
970 w_hess_entry: Array1<f64>,
971}
972
973impl SurvivalWorkspace {
974 fn new(n: usize) -> Self {
975 Self {
976 w_event: Array1::zeros(n),
977 w_event_inv_deriv: Array1::zeros(n),
978 w_event_outer: Array1::zeros(n),
979 w_hess_exit: Array1::zeros(n),
980 w_hess_entry: Array1::zeros(n),
981 }
982 }
983
984 fn reset(&mut self, n: usize) {
985 if self.w_event.len() != n {
986 *self = Self::new(n);
987 } else {
988 self.w_event.fill(0.0);
989 self.w_event_inv_deriv.fill(0.0);
990 self.w_event_outer.fill(0.0);
991 self.w_hess_exit.fill(0.0);
992 self.w_hess_entry.fill(0.0);
993 }
994 }
995}
996
997#[derive(Clone, Debug)]
1010pub struct OffsetChannelResiduals {
1011 pub exit: Array1<f64>,
1013 pub entry: Array1<f64>,
1015 pub derivative: Array1<f64>,
1017 pub right: Array1<f64>,
1021}
1022
1023#[derive(Clone, Debug)]
1026pub struct OffsetChannelCurvatures {
1027 pub rows: Vec<[[f64; 3]; 3]>,
1028}
1029
1030#[derive(Debug)]
1031pub struct WorkingModelSurvival {
1032 age_entry: Array1<f64>,
1033 age_exit: Array1<f64>,
1034 entry_at_origin: Array1<bool>,
1035 event_target: Array1<u8>,
1036 sampleweight: Array1<f64>,
1037 design: SurvivalDesign,
1038 offset_eta_entry: Array1<f64>,
1039 offset_eta_exit: Array1<f64>,
1040 offset_derivative_exit: Array1<f64>,
1041 penalties: PenaltyBlocks,
1042 monotonicity: SurvivalMonotonicityPenalty,
1043 structurally_monotonic: bool,
1044 structural_time_columns: usize,
1045 monotonicity_constraint_rows: Option<Array2<f64>>,
1046 monotonicity_constraint_offsets: Option<Array1<f64>>,
1047 workspace: std::sync::Mutex<SurvivalWorkspace>,
1048}
1049
1050impl Clone for WorkingModelSurvival {
1051 fn clone(&self) -> Self {
1052 let workspace = self.workspace.lock().unwrap().clone();
1053 Self {
1054 age_entry: self.age_entry.clone(),
1055 age_exit: self.age_exit.clone(),
1056 entry_at_origin: self.entry_at_origin.clone(),
1057 event_target: self.event_target.clone(),
1058 sampleweight: self.sampleweight.clone(),
1059 design: self.design.clone(),
1060 offset_eta_entry: self.offset_eta_entry.clone(),
1061 offset_eta_exit: self.offset_eta_exit.clone(),
1062 offset_derivative_exit: self.offset_derivative_exit.clone(),
1063 penalties: self.penalties.clone(),
1064 monotonicity: self.monotonicity,
1065 structurally_monotonic: self.structurally_monotonic,
1066 structural_time_columns: self.structural_time_columns,
1067 monotonicity_constraint_rows: self.monotonicity_constraint_rows.clone(),
1068 monotonicity_constraint_offsets: self.monotonicity_constraint_offsets.clone(),
1069 workspace: std::sync::Mutex::new(workspace),
1070 }
1071 }
1072}
1073
1074impl WorkingModelSurvival {
1075 const LOG_F64_MAX: f64 = 709.782712893384;
1076
1077 #[inline]
1078 fn scaled_exp_component(log_scale: f64, base: f64) -> Result<f64, EstimationError> {
1079 if base == 0.0 {
1080 return Ok(0.0);
1081 }
1082 let log_abs = log_scale + base.abs().ln();
1083 if !log_abs.is_finite() {
1084 crate::bail_invalid_estim!("survival interval term produced non-finite log-magnitude");
1085 }
1086 if log_abs > Self::LOG_F64_MAX {
1087 crate::bail_invalid_estim!(
1088 "survival interval term exceeds f64 range (log-magnitude={log_abs:.3e})"
1089 );
1090 }
1091 Ok(base.signum() * log_abs.exp())
1092 }
1093
1094 fn coefficient_dim(&self) -> usize {
1095 self.design.p_total()
1096 }
1097
1098 fn nrows(&self) -> usize {
1099 self.sampleweight.len()
1100 }
1101
1102 fn entry_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1103 let time_mat = match &self.design {
1104 SurvivalDesign::Flat { x_entry, .. } => x_entry,
1105 SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
1106 };
1107 self.design.design_dot(time_mat, beta)
1108 }
1109
1110 fn exit_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1111 let time_mat = match &self.design {
1112 SurvivalDesign::Flat { x_exit, .. } => x_exit,
1113 SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
1114 };
1115 self.design.design_dot(time_mat, beta)
1116 }
1117
1118 fn derivative_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1119 match &self.design {
1120 SurvivalDesign::Flat { x_derivative, .. } => x_derivative.dot(beta),
1121 SurvivalDesign::TimeCovariateShared {
1122 time_derivative, ..
1123 } => time_derivative.dot(&beta.slice(ndarray::s![..time_derivative.ncols()])),
1124 }
1125 }
1126
1127 fn fill_entry_row(&self, i: usize, out: &mut [f64]) {
1128 let time_mat = match &self.design {
1129 SurvivalDesign::Flat { x_entry, .. } => x_entry,
1130 SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
1131 };
1132 self.design.fill_row(time_mat, i, out);
1133 }
1134
1135 fn fill_exit_row(&self, i: usize, out: &mut [f64]) {
1136 let time_mat = match &self.design {
1137 SurvivalDesign::Flat { x_exit, .. } => x_exit,
1138 SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
1139 };
1140 self.design.fill_row(time_mat, i, out);
1141 }
1142
1143 fn fill_derivative_row(&self, i: usize, out: &mut [f64]) {
1144 match &self.design {
1145 SurvivalDesign::Flat { x_derivative, .. } => {
1146 for (dst, &src) in out.iter_mut().zip(x_derivative.row(i).iter()) {
1147 *dst = src;
1148 }
1149 }
1150 SurvivalDesign::TimeCovariateShared {
1151 time_derivative, ..
1152 } => {
1153 let p_time = time_derivative.ncols();
1154 for j in 0..p_time {
1155 out[j] = time_derivative[[i, j]];
1156 }
1157 for dst in out.iter_mut().skip(p_time) {
1158 *dst = 0.0;
1159 }
1160 }
1161 }
1162 }
1163
1164 fn derivative_xt_diag_x(&self, weights: &Array1<f64>) -> Array2<f64> {
1165 match &self.design {
1166 SurvivalDesign::Flat { x_derivative, .. } => fast_xt_diag_x(x_derivative, weights),
1167 SurvivalDesign::TimeCovariateShared {
1168 time_derivative,
1169 covariates,
1170 ..
1171 } => {
1172 let p_time = time_derivative.ncols();
1173 let p_cov = covariates.ncols();
1174 let mut out = Array2::<f64>::zeros((p_time + p_cov, p_time + p_cov));
1175 let time_block = fast_xt_diag_x(time_derivative, weights);
1176 out.slice_mut(ndarray::s![..p_time, ..p_time])
1177 .assign(&time_block);
1178 out
1179 }
1180 }
1181 }
1182
1183 fn interval_hessian_blas(&self, w_exit: &Array1<f64>, w_entry: &Array1<f64>) -> Array2<f64> {
1187 match &self.design {
1188 SurvivalDesign::Flat {
1189 x_entry, x_exit, ..
1190 } => {
1191 let mut h = fast_xt_diag_x(x_exit, w_exit);
1192 h -= &fast_xt_diag_x(x_entry, w_entry);
1193 h
1194 }
1195 SurvivalDesign::TimeCovariateShared {
1196 time_entry,
1197 time_exit,
1198 covariates,
1199 ..
1200 } => {
1201 let p_time = time_exit.ncols();
1202 let p_cov = covariates.ncols();
1203 let p = p_time + p_cov;
1204 let mut h = Array2::<f64>::zeros((p, p));
1205 let tt = {
1207 let mut block = fast_xt_diag_x(time_exit, w_exit);
1208 block -= &fast_xt_diag_x(time_entry, w_entry);
1209 block
1210 };
1211 h.slice_mut(ndarray::s![..p_time, ..p_time]).assign(&tt);
1212 if p_cov > 0 {
1213 let tc = {
1215 let mut block = fast_xt_diag_y(time_exit, w_exit, covariates);
1216 block -= &fast_xt_diag_y(time_entry, w_entry, covariates);
1217 block
1218 };
1219 h.slice_mut(ndarray::s![..p_time, p_time..]).assign(&tc);
1220 h.slice_mut(ndarray::s![p_time.., ..p_time]).assign(&tc.t());
1221 let w_diff = w_exit - w_entry;
1223 let cc = fast_xt_diag_x(covariates, &w_diff);
1224 h.slice_mut(ndarray::s![p_time.., p_time..]).assign(&cc);
1225 }
1226 h
1227 }
1228 }
1229 }
1230
1231 fn stabilized_structural_derivative(&self, deriv: f64) -> Option<(f64, f64)> {
1245 const STRUCTURAL_MONO_ROUNDOFF_TOL: f64 = 1e-7;
1246 const STRUCTURAL_DERIV_FLOOR: f64 = 1e-12;
1247 if !self.structurally_monotonic {
1248 return None;
1249 }
1250 if deriv >= STRUCTURAL_DERIV_FLOOR {
1251 return Some((deriv, 1.0));
1252 }
1253 if deriv >= -STRUCTURAL_MONO_ROUNDOFF_TOL {
1254 return Some((STRUCTURAL_DERIV_FLOOR, 0.0));
1255 }
1256 None
1257 }
1258
1259 fn validate_penalties(
1260 penalties: &PenaltyBlocks,
1261 coefficient_dim: usize,
1262 ) -> Result<(), SurvivalError> {
1263 for block in &penalties.blocks {
1264 if !block.lambda.is_finite() || block.lambda < 0.0 {
1265 return Err(SurvivalError::NonFiniteInput);
1266 }
1267 if block.range.start > block.range.end || block.range.end > coefficient_dim {
1268 return Err(SurvivalError::DimensionMismatch);
1269 }
1270 let block_dim = block.range.end - block.range.start;
1271 if block.matrix.nrows() != block_dim || block.matrix.ncols() != block_dim {
1272 return Err(SurvivalError::DimensionMismatch);
1273 }
1274 if block.matrix.iter().any(|v| !v.is_finite()) {
1275 return Err(SurvivalError::NonFiniteInput);
1276 }
1277 }
1278 Ok(())
1279 }
1280
1281 fn derivative_guard(&self) -> f64 {
1282 if self.structurally_monotonic {
1283 return 0.0;
1287 }
1288 self.monotonicity.tolerance.max(0.0)
1289 }
1290
1291 fn derivative_guard_numerical(&self) -> f64 {
1292 let derivative_guard = self.derivative_guard();
1293 if derivative_guard <= 0.0 {
1294 if self.structurally_monotonic {
1303 -1e-10
1304 } else {
1305 1e-12
1306 }
1307 } else {
1308 (derivative_guard - (1e-10_f64).min(0.01 * derivative_guard)).max(1e-12)
1309 }
1310 }
1311
1312 fn interval_increment_guard(&self, h_entry: f64, h_exit: f64) -> f64 {
1313 let scale = h_entry.abs().max(h_exit.abs()).max(1.0);
1314 1e-10 * scale
1315 }
1316
1317 fn structural_time_coefficient_constraints(&self) -> Option<LinearInequalityConstraints> {
1318 if !self.structurally_monotonic {
1319 return None;
1320 }
1321 let p = self.coefficient_dim();
1322 let time_columns = self.structural_time_columns.min(p);
1323 if time_columns == 0 {
1324 return None;
1325 }
1326 const STRUCTURAL_DERIV_TOL: f64 = 1e-12;
1327 let mut active_columns = vec![false; time_columns];
1328 let mut derivative_row = vec![0.0_f64; p];
1329 for i in 0..self.nrows() {
1330 if self.sampleweight[i] <= 0.0 {
1331 continue;
1332 }
1333 self.fill_derivative_row(i, &mut derivative_row);
1334 for j in 0..time_columns {
1335 if derivative_row[j] > STRUCTURAL_DERIV_TOL {
1336 active_columns[j] = true;
1337 }
1338 }
1339 }
1340 if let Some(rows) = self.monotonicity_constraint_rows.as_ref() {
1341 for i in 0..rows.nrows() {
1342 for j in 0..time_columns {
1343 if rows[[i, j]] > STRUCTURAL_DERIV_TOL {
1344 active_columns[j] = true;
1345 }
1346 }
1347 }
1348 }
1349 let active_columns: Vec<usize> = active_columns
1350 .into_iter()
1351 .enumerate()
1352 .filter_map(|(j, active)| active.then_some(j))
1353 .collect();
1354 if active_columns.is_empty() {
1355 return None;
1356 }
1357 let mut a = Array2::<f64>::zeros((active_columns.len(), p));
1358 let b = Array1::<f64>::zeros(active_columns.len());
1359 for (row, &col) in active_columns.iter().enumerate() {
1360 a[[row, col]] = 1.0;
1361 }
1362 Some(LinearInequalityConstraints { a, b })
1363 }
1364
1365 pub fn monotonicity_linear_constraints(&self) -> Option<LinearInequalityConstraints> {
1366 let p = self.coefficient_dim();
1367 const DERIVATIVE_ROW_NORM_TOL: f64 = 1e-12;
1368 if p == 0 {
1369 return None;
1370 }
1371 if self.structurally_monotonic {
1372 return self.structural_time_coefficient_constraints();
1373 }
1374 if let (Some(rows), Some(offsets)) = (
1375 self.monotonicity_constraint_rows.as_ref(),
1376 self.monotonicity_constraint_offsets.as_ref(),
1377 ) {
1378 let activerows: Vec<usize> = (0..rows.nrows())
1379 .filter(|&i| {
1380 rows.row(i).iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()))
1381 > DERIVATIVE_ROW_NORM_TOL
1382 })
1383 .collect();
1384 if activerows.is_empty() {
1385 return None;
1386 }
1387 let mut a = Array2::<f64>::zeros((activerows.len(), p));
1388 let mut b = Array1::<f64>::zeros(activerows.len());
1389 for (r, &i) in activerows.iter().enumerate() {
1390 a.row_mut(r).assign(&rows.row(i));
1391 b[r] = self.derivative_guard() - offsets[i];
1392 }
1393 return Some(compress_positive_collinear_constraints(&a, &b));
1394 }
1395 None
1396 }
1397
1398 pub fn from_engine_inputs(
1399 inputs: SurvivalEngineInputs<'_>,
1400 penalties: PenaltyBlocks,
1401 monotonicity: SurvivalMonotonicityPenalty,
1402 spec: SurvivalSpec,
1403 ) -> Result<Self, SurvivalError> {
1404 Self::from_engine_inputswith_offsets(inputs, None, penalties, monotonicity, spec)
1405 }
1406
1407 fn validate_offsets(
1408 offsets: Option<SurvivalBaselineOffsets<'_>>,
1409 n: usize,
1410 ) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), SurvivalError> {
1411 if let Some(off) = offsets {
1412 if off.eta_entry.len() != n || off.eta_exit.len() != n || off.derivative_exit.len() != n
1413 {
1414 return Err(SurvivalError::DimensionMismatch);
1415 }
1416 if off.eta_entry.iter().any(|v| !v.is_finite())
1417 || off.eta_exit.iter().any(|v| !v.is_finite())
1418 || off.derivative_exit.iter().any(|v| !v.is_finite())
1419 {
1420 return Err(SurvivalError::NonFiniteInput);
1421 }
1422 Ok((
1423 off.eta_entry.to_owned(),
1424 off.eta_exit.to_owned(),
1425 off.derivative_exit.to_owned(),
1426 ))
1427 } else {
1428 Ok((Array1::zeros(n), Array1::zeros(n), Array1::zeros(n)))
1429 }
1430 }
1431
1432 fn validate_common_inputs(
1433 age_entry: &ArrayView1<f64>,
1434 age_exit: &ArrayView1<f64>,
1435 event_target: &ArrayView1<u8>,
1436 event_competing: &ArrayView1<u8>,
1437 sampleweight: &ArrayView1<f64>,
1438 ) -> Result<(), SurvivalError> {
1439 if age_entry.iter().any(|v| !v.is_finite())
1440 || age_exit.iter().any(|v| !v.is_finite())
1441 || sampleweight.iter().any(|v| !v.is_finite() || *v < 0.0)
1442 {
1443 return Err(SurvivalError::NonFiniteInput);
1444 }
1445 if let Some(&label) = event_target.iter().find(|&&v| v > 1) {
1452 return Err(SurvivalError::EventCodeInvalid {
1453 reason: format!(
1454 "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"
1455 ),
1456 });
1457 }
1458 if let Some(&label) = event_competing.iter().find(|&&v| v > 1) {
1459 return Err(SurvivalError::EventCodeInvalid {
1460 reason: format!(
1461 "single-hazard survival engine requires a binary {{0, 1}} event_competing, got multi-cause label {label}"
1462 ),
1463 });
1464 }
1465 if event_target
1466 .iter()
1467 .zip(event_competing.iter())
1468 .any(|(&target, &competing)| target > 0 && competing > 0)
1469 {
1470 return Err(SurvivalError::EventCodeInvalid {
1471 reason: "a row cannot be simultaneously a target event and a competing event"
1472 .to_string(),
1473 });
1474 }
1475 if age_entry
1489 .iter()
1490 .zip(age_exit.iter())
1491 .any(|(&entry, &exit)| entry < 0.0 || exit <= 0.0)
1492 {
1493 return Err(SurvivalError::NonFiniteInput);
1494 }
1495 Ok::<(), _>(())
1496 }
1497
1498 fn validate_monotonicity_constraints(
1499 rows: Option<ArrayView2<'_, f64>>,
1500 offsets: Option<ArrayView1<'_, f64>>,
1501 coefficient_dim: usize,
1502 ) -> Result<(Option<Array2<f64>>, Option<Array1<f64>>), SurvivalError> {
1503 match (rows, offsets) {
1504 (None, None) => Ok((None, None)),
1505 (Some(rows), Some(offsets)) => {
1506 if rows.ncols() != coefficient_dim
1507 || rows.nrows() != offsets.len()
1508 || rows.iter().any(|v| !v.is_finite())
1509 || offsets.iter().any(|v| !v.is_finite())
1510 {
1511 return Err(SurvivalError::DimensionMismatch);
1512 }
1513 Ok((Some(rows.to_owned()), Some(offsets.to_owned())))
1514 }
1515 _ => Err(SurvivalError::DimensionMismatch),
1516 }
1517 }
1518
1519 fn finish_construction(
1520 age_entry: ArrayView1<f64>,
1521 age_exit: ArrayView1<f64>,
1522 event_target: ArrayView1<u8>,
1523 sampleweight: ArrayView1<f64>,
1524 design: SurvivalDesign,
1525 offset_eta_entry: Array1<f64>,
1526 offset_eta_exit: Array1<f64>,
1527 offset_derivative_exit: Array1<f64>,
1528 penalties: PenaltyBlocks,
1529 monotonicity: SurvivalMonotonicityPenalty,
1530 monotonicity_constraint_rows: Option<Array2<f64>>,
1531 monotonicity_constraint_offsets: Option<Array1<f64>>,
1532 ) -> Self {
1533 let n = age_entry.len();
1534 Self {
1535 age_entry: age_entry.to_owned(),
1536 age_exit: age_exit.to_owned(),
1537 entry_at_origin: age_entry.mapv(|t| t <= ENTRY_AT_ORIGIN_THRESHOLD),
1538 event_target: event_target.to_owned(),
1539 sampleweight: sampleweight.to_owned(),
1540 design,
1541 offset_eta_entry,
1542 offset_eta_exit,
1543 offset_derivative_exit,
1544 penalties,
1545 monotonicity,
1546 structurally_monotonic: false,
1547 structural_time_columns: 0,
1548 monotonicity_constraint_rows,
1549 monotonicity_constraint_offsets,
1550 workspace: std::sync::Mutex::new(SurvivalWorkspace::new(n)),
1551 }
1552 }
1553
1554 pub fn from_engine_inputswith_offsets(
1555 inputs: SurvivalEngineInputs<'_>,
1556 offsets: Option<SurvivalBaselineOffsets<'_>>,
1557 penalties: PenaltyBlocks,
1558 monotonicity: SurvivalMonotonicityPenalty,
1559 spec: SurvivalSpec,
1560 ) -> Result<Self, SurvivalError> {
1561 if spec == SurvivalSpec::Crude {
1562 return Err(SurvivalError::UnsupportedSpec("crude"));
1563 }
1564 let n = inputs.age_entry.len();
1565 let p = inputs.x_entry.ncols();
1566 if inputs.age_exit.len() != n
1567 || inputs.event_target.len() != n
1568 || inputs.event_competing.len() != n
1569 || inputs.sampleweight.len() != n
1570 || inputs.x_entry.nrows() != n
1571 || inputs.x_exit.nrows() != n
1572 || inputs.x_derivative.nrows() != n
1573 || inputs.x_entry.ncols() != inputs.x_exit.ncols()
1574 || inputs.x_entry.ncols() != inputs.x_derivative.ncols()
1575 {
1576 return Err(SurvivalError::DimensionMismatch);
1577 }
1578 Self::validate_penalties(&penalties, p)?;
1579 Self::validate_common_inputs(
1580 &inputs.age_entry,
1581 &inputs.age_exit,
1582 &inputs.event_target,
1583 &inputs.event_competing,
1584 &inputs.sampleweight,
1585 )?;
1586 if inputs.x_entry.iter().any(|v| !v.is_finite())
1587 || inputs.x_exit.iter().any(|v| !v.is_finite())
1588 || inputs.x_derivative.iter().any(|v| !v.is_finite())
1589 {
1590 return Err(SurvivalError::NonFiniteInput);
1591 }
1592 let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
1593 Self::validate_offsets(offsets, n)?;
1594 let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
1595 Self::validate_monotonicity_constraints(
1596 inputs.monotonicity_constraint_rows,
1597 inputs.monotonicity_constraint_offsets,
1598 p,
1599 )?;
1600
1601 Ok(Self::finish_construction(
1602 inputs.age_entry,
1603 inputs.age_exit,
1604 inputs.event_target,
1605 inputs.sampleweight,
1606 SurvivalDesign::Flat {
1607 x_entry: inputs.x_entry.to_owned(),
1608 x_exit: inputs.x_exit.to_owned(),
1609 x_derivative: inputs.x_derivative.to_owned(),
1610 },
1611 offset_eta_entry,
1612 offset_eta_exit,
1613 offset_derivative_exit,
1614 penalties,
1615 monotonicity,
1616 monotonicity_constraint_rows,
1617 monotonicity_constraint_offsets,
1618 ))
1619 }
1620
1621 pub fn from_time_covariate_inputswith_offsets(
1622 inputs: SurvivalTimeCovarInputs<'_>,
1623 offsets: Option<SurvivalBaselineOffsets<'_>>,
1624 penalties: PenaltyBlocks,
1625 monotonicity: SurvivalMonotonicityPenalty,
1626 spec: SurvivalSpec,
1627 ) -> Result<Self, SurvivalError> {
1628 if spec == SurvivalSpec::Crude {
1629 return Err(SurvivalError::UnsupportedSpec("crude"));
1630 }
1631 let n = inputs.age_entry.len();
1632 let p_time = inputs.time_entry.ncols();
1633 let p_cov = inputs.covariates.ncols();
1634 let p = p_time + p_cov;
1635 if inputs.age_exit.len() != n
1636 || inputs.event_target.len() != n
1637 || inputs.event_competing.len() != n
1638 || inputs.sampleweight.len() != n
1639 || inputs.time_entry.nrows() != n
1640 || inputs.time_exit.nrows() != n
1641 || inputs.time_derivative.nrows() != n
1642 || inputs.covariates.nrows() != n
1643 || inputs.time_entry.ncols() != inputs.time_exit.ncols()
1644 || inputs.time_entry.ncols() != inputs.time_derivative.ncols()
1645 {
1646 return Err(SurvivalError::DimensionMismatch);
1647 }
1648 Self::validate_penalties(&penalties, p)?;
1649 Self::validate_common_inputs(
1650 &inputs.age_entry,
1651 &inputs.age_exit,
1652 &inputs.event_target,
1653 &inputs.event_competing,
1654 &inputs.sampleweight,
1655 )?;
1656 if inputs.time_entry.iter().any(|v| !v.is_finite())
1657 || inputs.time_exit.iter().any(|v| !v.is_finite())
1658 || inputs.time_derivative.iter().any(|v| !v.is_finite())
1659 || inputs.covariates.iter().any(|v| !v.is_finite())
1660 {
1661 return Err(SurvivalError::NonFiniteInput);
1662 }
1663 let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
1664 Self::validate_offsets(offsets, n)?;
1665 let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
1666 Self::validate_monotonicity_constraints(
1667 inputs.monotonicity_constraint_rows,
1668 inputs.monotonicity_constraint_offsets,
1669 p,
1670 )?;
1671
1672 Ok(Self::finish_construction(
1673 inputs.age_entry,
1674 inputs.age_exit,
1675 inputs.event_target,
1676 inputs.sampleweight,
1677 SurvivalDesign::TimeCovariateShared {
1678 time_entry: inputs.time_entry.to_owned(),
1679 time_exit: inputs.time_exit.to_owned(),
1680 time_derivative: inputs.time_derivative.to_owned(),
1681 covariates: inputs.covariates.to_owned(),
1682 },
1683 offset_eta_entry,
1684 offset_eta_exit,
1685 offset_derivative_exit,
1686 penalties,
1687 monotonicity,
1688 monotonicity_constraint_rows,
1689 monotonicity_constraint_offsets,
1690 ))
1691 }
1692
1693 pub fn set_penalty_lambdas(&mut self, lambdas: &[f64]) -> Result<(), EstimationError> {
1707 if lambdas.len() != self.penalties.blocks.len() {
1708 crate::bail_invalid_estim!(
1709 "set_penalty_lambdas expects {} lambdas, got {}",
1710 self.penalties.blocks.len(),
1711 lambdas.len()
1712 );
1713 }
1714 for (block, &lambda) in self.penalties.blocks.iter_mut().zip(lambdas.iter()) {
1715 if !lambda.is_finite() || lambda < 0.0 {
1716 crate::bail_invalid_estim!("penalty lambda must be finite and >= 0, got {lambda}");
1717 }
1718 block.lambda = lambda;
1719 }
1720 Ok(())
1721 }
1722
1723 pub fn set_structural_monotonicity(
1724 &mut self,
1725 enabled: bool,
1726 time_columns: usize,
1727 ) -> Result<(), EstimationError> {
1728 let p = self.coefficient_dim();
1729 if time_columns > p {
1730 crate::bail_invalid_estim!(
1731 "structural time columns {} exceed coefficient dimension {}",
1732 time_columns,
1733 p
1734 );
1735 }
1736 if enabled && time_columns == 0 {
1737 crate::bail_invalid_estim!("structural monotonicity requires at least one time column");
1738 }
1739 if enabled {
1740 const STRUCTURAL_DERIV_TOL: f64 = 1e-12;
1741 for (i, &offset) in self.offset_derivative_exit.iter().enumerate() {
1742 if offset < -STRUCTURAL_DERIV_TOL {
1743 crate::bail_invalid_estim!(
1744 "structural monotonicity requires nonnegative derivative offsets; found offset_derivative_exit[{i}]={offset:.3e}"
1745 );
1746 }
1747 }
1748 let mut derivative_row = vec![0.0_f64; p];
1749 for i in 0..self.nrows() {
1750 self.fill_derivative_row(i, &mut derivative_row);
1751 for j in 0..time_columns {
1752 let v = derivative_row[j];
1753 if v < -STRUCTURAL_DERIV_TOL {
1754 crate::bail_invalid_estim!(
1755 "structural monotonicity requires nonnegative time-derivative basis entries; found x_derivative[{i},{j}]={v:.3e}"
1756 );
1757 }
1758 }
1759 for j in time_columns..p {
1760 let v = derivative_row[j];
1761 if v.abs() > STRUCTURAL_DERIV_TOL {
1762 crate::bail_invalid_estim!(
1763 "structural monotonicity requires zero derivative contribution outside the time block; found x_derivative[{i},{j}]={v:.3e}"
1764 );
1765 }
1766 }
1767 }
1768 if let (Some(rows), Some(offsets)) = (
1769 self.monotonicity_constraint_rows.as_ref(),
1770 self.monotonicity_constraint_offsets.as_ref(),
1771 ) {
1772 for (i, &offset) in offsets.iter().enumerate() {
1773 if offset < -STRUCTURAL_DERIV_TOL {
1774 crate::bail_invalid_estim!(
1775 "structural monotonicity requires nonnegative collocation derivative offsets; found monotonicity_constraint_offsets[{i}]={offset:.3e}"
1776 );
1777 }
1778 }
1779 for i in 0..rows.nrows() {
1780 for j in 0..time_columns {
1781 let v = rows[[i, j]];
1782 if v < -STRUCTURAL_DERIV_TOL {
1783 crate::bail_invalid_estim!(
1784 "structural monotonicity requires nonnegative collocation derivative basis entries; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
1785 );
1786 }
1787 }
1788 for j in time_columns..p {
1789 let v = rows[[i, j]];
1790 if v.abs() > STRUCTURAL_DERIV_TOL {
1791 crate::bail_invalid_estim!(
1792 "structural monotonicity requires zero collocation derivative contribution outside the time block; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
1793 );
1794 }
1795 }
1796 }
1797 }
1798 }
1799 self.structurally_monotonic = enabled;
1800 self.structural_time_columns = if enabled { time_columns } else { 0 };
1801 Ok(())
1802 }
1803
1804 pub fn update_state(&self, beta: &Array1<f64>) -> Result<WorkingState, EstimationError> {
1805 if beta.len() != self.coefficient_dim() {
1806 crate::bail_invalid_estim!("survival beta dimension mismatch");
1807 }
1808
1809 let n = self.nrows();
1810 let p = self.coefficient_dim();
1811
1812 let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
1838 let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
1839 let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
1840
1841 let mut nll = 0.0;
1842 let derivative_guard = self.derivative_guard();
1843 let derivative_guard_numerical = self.derivative_guard_numerical();
1844 let mut workspace = self.workspace.lock().unwrap();
1845 workspace.reset(n);
1846 let SurvivalWorkspace {
1847 w_event,
1848 w_event_inv_deriv,
1849 w_event_outer,
1850 w_hess_exit,
1851 w_hess_entry,
1852 } = &mut *workspace;
1853
1854 for i in 0..n {
1856 let w = self.sampleweight[i];
1857 if w <= 0.0 {
1858 continue;
1859 }
1860 let entry_age = self.age_entry[i];
1861 let exit_age = self.age_exit[i];
1862 if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
1863 crate::bail_invalid_estim!(
1864 "survival ages must be finite with age_exit >= age_entry"
1865 );
1866 }
1867 let d = f64::from(self.event_target[i]);
1868
1869 let has_entry_interval = !self.entry_at_origin[i];
1870 let interval_scale = if has_entry_interval {
1871 eta_exit[i].max(eta_entry[i])
1872 } else {
1873 eta_exit[i]
1874 };
1875 let h_e_scaled = (eta_exit[i] - interval_scale).exp();
1876 let h_s_scaled = if has_entry_interval {
1877 (eta_entry[i] - interval_scale).exp()
1878 } else {
1879 0.0
1880 };
1881 let interval_scaled = h_e_scaled - h_s_scaled;
1882 let interval = Self::scaled_exp_component(interval_scale, interval_scaled)?;
1883 let (deriv, deriv_slope) = self
1884 .stabilized_structural_derivative(derivative_raw[i])
1885 .unwrap_or((derivative_raw[i], 1.0));
1886 let mono_floor = if d > 0.0 {
1895 derivative_guard_numerical
1896 } else {
1897 0.0
1898 };
1899 if !deriv.is_finite() || deriv < mono_floor {
1900 return Err(EstimationError::ParameterConstraintViolation(format!(
1901 "survival monotonicity violated at row {}: d_eta/dt={:.3e} <= tolerance={:.3e}",
1902 i, deriv, derivative_guard
1903 )));
1904 }
1905 if has_entry_interval {
1906 let increment_guard = self.interval_increment_guard(h_s_scaled, h_e_scaled);
1907 if interval_scaled + increment_guard < 0.0 {
1908 return Err(EstimationError::ParameterConstraintViolation(format!(
1909 "survival cumulative hazard decreased over row {}: H(exit)-H(entry)={:.6e}",
1910 i, interval
1911 )));
1912 }
1913 }
1914 nll += w * interval;
1915
1916 let w_exit_i = w * eta_exit[i].exp();
1920 let w_entry_i = if has_entry_interval {
1921 w * eta_entry[i].exp()
1922 } else {
1923 0.0
1924 };
1925 if !w_exit_i.is_finite() {
1926 crate::bail_invalid_estim!(
1927 "survival interval term exceeds f64 range at row {i} (w*exp(eta_exit)={w_exit_i:.3e})"
1928 );
1929 }
1930 w_hess_exit[i] = w_exit_i;
1931 w_hess_entry[i] = w_entry_i;
1932
1933 if d > 0.0 {
1934 let inv_deriv = deriv_slope / deriv;
1938 nll += -w * (eta_exit[i] + deriv.ln());
1939 w_event[i] = w;
1940 w_event_inv_deriv[i] = w * inv_deriv;
1941 w_event_outer[i] = w * inv_deriv * inv_deriv;
1942 }
1943 }
1944
1945 let mut h = self.interval_hessian_blas(w_hess_exit, w_hess_entry);
1949 let mut grad = Array1::<f64>::zeros(p);
1953 let mut grad_comp = Array1::<f64>::zeros(p);
1954 let mut row_exit = vec![0.0_f64; p];
1955 let mut row_entry = vec![0.0_f64; p];
1956 let mut row_derivative = vec![0.0_f64; p];
1957 for i in 0..n {
1958 let w_interval_exit = w_hess_exit[i];
1959 let w_interval_entry = w_hess_entry[i];
1960 let w_event_exit = w_event[i];
1961 let w_event_derivative = w_event_inv_deriv[i];
1962 if w_interval_exit == 0.0
1963 && w_interval_entry == 0.0
1964 && w_event_exit == 0.0
1965 && w_event_derivative == 0.0
1966 {
1967 continue;
1968 }
1969 self.fill_exit_row(i, &mut row_exit);
1970 self.fill_entry_row(i, &mut row_entry);
1971 self.fill_derivative_row(i, &mut row_derivative);
1972 for j in 0..p {
1973 let contribution = w_interval_exit * row_exit[j]
1974 - w_interval_entry * row_entry[j]
1975 - w_event_exit * row_exit[j]
1976 - w_event_derivative * row_derivative[j];
1977 let t = grad[j] + contribution;
1978 if grad[j].abs() >= contribution.abs() {
1979 grad_comp[j] += (grad[j] - t) + contribution;
1980 } else {
1981 grad_comp[j] += (contribution - t) + grad[j];
1982 }
1983 grad[j] = t;
1984 }
1985 }
1986 grad += &grad_comp;
1987
1988 h += &self.derivative_xt_diag_x(w_event_outer);
1989
1990 let score_norm = array1_l2_norm(&grad);
1994
1995 let penaltygrad = self.penalties.gradient(beta);
1996 let penalty_dev = self.penalties.deviance(beta);
1997 let penaltygrad_norm = array1_l2_norm(&penaltygrad);
1998
1999 let mut totalgrad = grad;
2000 totalgrad += &penaltygrad;
2001
2002 self.penalties.addhessian_inplace(&mut h);
2003 let log_likelihood = -nll;
2010 let deviance = 2.0 * nll;
2011
2012 Ok(WorkingState {
2013 eta: LinearPredictor::new(eta_exit),
2014 gradient: totalgrad,
2015 hessian: gam_linalg::matrix::SymmetricMatrix::Dense(h),
2016 log_likelihood,
2017 deviance,
2018 penalty_term: penalty_dev,
2019 firth: gam_solve::pirls::FirthDiagnostics::Inactive,
2020 ridge_used: 0.0,
2021 hessian_curvature: gam_solve::pirls::HessianCurvatureKind::Observed,
2022 gradient_natural_scale: score_norm + penaltygrad_norm,
2023 })
2024 }
2025
2026 pub(crate) fn survival_hessian_derivative_correction(
2036 &self,
2037 beta: &Array1<f64>,
2038 u_k: &Array1<f64>,
2039 ) -> Result<Array2<f64>, EstimationError> {
2040 let p = beta.len();
2041 let n = self.nrows();
2042
2043 let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2044 let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2045 let deriv_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2046 let exp_entry = eta_entry.mapv(f64::exp);
2047 let exp_exit = eta_exit.mapv(f64::exp);
2048 let guard = self.derivative_guard();
2049 let guard_numerical = self.derivative_guard_numerical();
2050
2051 let jac = Array1::<f64>::ones(p);
2052 let curvature = Array1::<f64>::zeros(p);
2053 let third = Array1::<f64>::zeros(p);
2054
2055 let mut row_exit = vec![0.0_f64; p];
2056 let mut row_entry = vec![0.0_f64; p];
2057 let mut row_derivative = vec![0.0_f64; p];
2058 let mut ge = vec![0.0_f64; p];
2059 let mut gs = vec![0.0_f64; p];
2060 let mut gsd = vec![0.0_f64; p];
2061 let mut he = vec![0.0_f64; p];
2062 let mut hs = vec![0.0_f64; p];
2063 let mut hsd = vec![0.0_f64; p];
2064 let mut te = vec![0.0_f64; p];
2065 let mut ts = vec![0.0_f64; p];
2066 let mut tsd = vec![0.0_f64; p];
2067
2068 let mut b_dir = Array2::<f64>::zeros((p, p));
2069
2070 for i in 0..n {
2071 let w_i = self.sampleweight[i];
2072 if w_i <= 0.0 {
2073 continue;
2074 }
2075 let has_entry = !self.entry_at_origin[i];
2076 let mut deta_e = 0.0_f64;
2077 let mut deta_s = 0.0_f64;
2078 let mut ds = 0.0_f64;
2079 self.fill_exit_row(i, &mut row_exit);
2080 self.fill_entry_row(i, &mut row_entry);
2081 self.fill_derivative_row(i, &mut row_derivative);
2082 for j in 0..p {
2083 ge[j] = row_exit[j] * jac[j];
2084 gs[j] = row_entry[j] * jac[j];
2085 gsd[j] = row_derivative[j] * jac[j];
2086 he[j] = row_exit[j] * curvature[j];
2087 hs[j] = row_entry[j] * curvature[j];
2088 hsd[j] = row_derivative[j] * curvature[j];
2089 te[j] = row_exit[j] * third[j];
2090 ts[j] = row_entry[j] * third[j];
2091 tsd[j] = row_derivative[j] * third[j];
2092 deta_e += ge[j] * u_k[j];
2093 if has_entry {
2094 deta_s += gs[j] * u_k[j];
2095 }
2096 ds += gsd[j] * u_k[j];
2097 }
2098
2099 for r in 0..p {
2101 let dge_r = he[r] * u_k[r];
2102 let dgs_r = hs[r] * u_k[r];
2103 let dhe_r = te[r] * u_k[r];
2104 let dhs_r = ts[r] * u_k[r];
2105 for c in 0..p {
2106 let dge_c = he[c] * u_k[c];
2107 let dgs_c = hs[c] * u_k[c];
2108 let mut d_h_rc =
2109 exp_exit[i] * (deta_e * ge[r] * ge[c] + dge_r * ge[c] + ge[r] * dge_c);
2110 if r == c {
2111 d_h_rc += exp_exit[i] * (deta_e * he[r] + dhe_r);
2112 }
2113 if has_entry {
2114 d_h_rc -=
2115 exp_entry[i] * (deta_s * gs[r] * gs[c] + dgs_r * gs[c] + gs[r] * dgs_c);
2116 if r == c {
2117 d_h_rc -= exp_entry[i] * (deta_s * hs[r] + dhs_r);
2118 }
2119 }
2120 b_dir[[r, c]] += w_i * d_h_rc;
2121 }
2122 }
2123
2124 let (s_i, s_slope) = self
2126 .stabilized_structural_derivative(deriv_raw[i])
2127 .unwrap_or((deriv_raw[i], 1.0));
2128 if !s_i.is_finite() {
2129 return Err(EstimationError::ParameterConstraintViolation(format!(
2130 "survival monotonicity violated in unified trace contraction at row {i}: \
2131 d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
2132 )));
2133 }
2134 if self.event_target[i] > 0 && s_slope != 0.0 {
2135 if s_i < guard_numerical {
2140 return Err(EstimationError::ParameterConstraintViolation(format!(
2141 "survival monotonicity violated in unified trace contraction at row {i}: \
2142 d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
2143 )));
2144 }
2145 let inv_s = 1.0 / s_i;
2146 let inv_s2 = inv_s * inv_s;
2147 let inv_s3 = inv_s2 * inv_s;
2148 for r in 0..p {
2149 let dgd_r = hsd[r] * u_k[r];
2150 let dtsd_r = tsd[r] * u_k[r];
2151 let dte_r = te[r] * u_k[r];
2152 for c in 0..p {
2153 let dgd_c = hsd[c] * u_k[c];
2154 let mut d_h_rc = (dgd_r * gsd[c] + gsd[r] * dgd_c) * inv_s2
2155 - 2.0 * gsd[r] * gsd[c] * ds * inv_s3;
2156 if r == c {
2157 d_h_rc += -dte_r;
2158 d_h_rc += -(dtsd_r * inv_s - hsd[r] * ds * inv_s2);
2159 }
2160 b_dir[[r, c]] += w_i * d_h_rc;
2161 }
2162 }
2163 }
2164 }
2165
2166 Ok(b_dir)
2167 }
2168
2169 pub fn offset_channel_residuals(
2207 &self,
2208 beta: &Array1<f64>,
2209 ) -> Result<OffsetChannelResiduals, EstimationError> {
2210 if beta.len() != self.coefficient_dim() {
2211 crate::bail_invalid_estim!(
2212 "survival beta dimension mismatch in offset_channel_residuals"
2213 );
2214 }
2215 let n = self.nrows();
2216 let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2217 let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2218 let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2219
2220 let derivative_guard_numerical = self.derivative_guard_numerical();
2221 let mut r_exit = Array1::<f64>::zeros(n);
2222 let mut r_entry = Array1::<f64>::zeros(n);
2223 let mut r_deriv = Array1::<f64>::zeros(n);
2224
2225 for i in 0..n {
2226 let w = self.sampleweight[i];
2227 if w <= 0.0 {
2228 continue;
2229 }
2230 let entry_age = self.age_entry[i];
2231 let exit_age = self.age_exit[i];
2232 if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
2233 crate::bail_invalid_estim!(
2234 "survival ages must be finite with age_exit >= age_entry"
2235 );
2236 }
2237 let has_entry_interval = !self.entry_at_origin[i];
2238 let d = f64::from(self.event_target[i]);
2239 let w_exit_i = w * eta_exit[i].exp();
2243 let w_entry_i = if has_entry_interval {
2244 w * eta_entry[i].exp()
2245 } else {
2246 0.0
2247 };
2248 if !w_exit_i.is_finite() {
2249 crate::bail_invalid_estim!(
2250 "offset_channel_residuals: w*exp(eta_exit)={w_exit_i:.3e} non-finite at row {i}"
2251 );
2252 }
2253 r_exit[i] = w_exit_i - d * w;
2254 r_entry[i] = -w_entry_i;
2255 let deriv_raw = derivative_raw[i];
2260 let (deriv, deriv_slope) = self
2261 .stabilized_structural_derivative(deriv_raw)
2262 .unwrap_or((deriv_raw, 1.0));
2263 let mono_floor = if d > 0.0 {
2264 derivative_guard_numerical
2265 } else {
2266 0.0
2267 };
2268 if !deriv.is_finite() || deriv < mono_floor {
2269 return Err(EstimationError::ParameterConstraintViolation(format!(
2270 "offset_channel_residuals: derivative ≤ numerical guard at row {i}: {deriv:.3e}"
2271 )));
2272 }
2273 if d > 0.0 {
2274 r_deriv[i] = -w * d * deriv_slope / deriv;
2277 }
2278 }
2279
2280 let right = Array1::<f64>::zeros(r_exit.len());
2281 Ok(OffsetChannelResiduals {
2282 exit: r_exit,
2283 entry: r_entry,
2284 derivative: r_deriv,
2285 right,
2286 })
2287 }
2288
2289 pub fn unified_lamlobjective_and_rhogradient(
2295 &self,
2296 beta: &Array1<f64>,
2297 state: &WorkingState,
2298 rho: &Array1<f64>,
2299 ) -> Result<(f64, Array1<f64>), EstimationError> {
2300 use gam_problem::{EvalMode, PseudoLogdetMode};
2301 use gam_solve::estimate::reml::assembly::{
2302 InnerAssembly, PenaltyBlockDesc, penalty_coords_from_blocks,
2303 };
2304 use gam_solve::estimate::reml::reml_outer_engine::{
2305 DenseSpectralOperator, DispersionHandling, PenaltyLogdetDerivs,
2306 compute_block_penalty_logdet_derivs,
2307 };
2308
2309 let p = beta.len();
2310 let active_penalty_blocks: Vec<&PenaltyBlock> = self
2311 .penalties
2312 .blocks
2313 .iter()
2314 .filter(|b| b.lambda > 0.0)
2315 .collect();
2316 if rho.len() != active_penalty_blocks.len() {
2317 crate::bail_invalid_estim!(
2318 "survival LAML rho dimension {} does not match active penalty block count {}",
2319 rho.len(),
2320 active_penalty_blocks.len()
2321 );
2322 }
2323 let k_count = active_penalty_blocks.len();
2324
2325 let h_dense = state.hessian.to_dense();
2327 let has_left_truncation = self
2328 .age_entry
2329 .iter()
2330 .any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD);
2331 let hessian_logdet_mode = if has_left_truncation {
2343 PseudoLogdetMode::HardPseudo
2344 } else {
2345 PseudoLogdetMode::Smooth
2346 };
2347 let hop = DenseSpectralOperator::from_symmetric_with_mode(&h_dense, hessian_logdet_mode)
2348 .map_err(EstimationError::InvalidInput)?;
2349
2350 let block_descs: Vec<PenaltyBlockDesc> = self
2352 .penalties
2353 .blocks
2354 .iter()
2355 .filter(|b| b.lambda > 0.0)
2356 .map(|b| PenaltyBlockDesc {
2357 matrix: &b.matrix,
2358 range_start: b.range.start,
2359 range_end: b.range.end,
2360 })
2361 .collect();
2362 let penalty_coords =
2363 penalty_coords_from_blocks(&block_descs, p).map_err(EstimationError::InvalidInput)?;
2364
2365 let per_block_rho: Vec<Array1<f64>> =
2367 rho.iter().map(|&r| Array1::from_vec(vec![r])).collect();
2368 let per_block_penalty_matrices: Vec<Vec<Array2<f64>>> = active_penalty_blocks
2369 .iter()
2370 .map(|b| vec![b.matrix.clone()])
2371 .collect();
2372 let per_block_penalty_refs: Vec<&[Array2<f64>]> = per_block_penalty_matrices
2373 .iter()
2374 .map(|v| v.as_slice())
2375 .collect();
2376 let penalty_logdet = if k_count > 0 {
2377 compute_block_penalty_logdet_derivs(&per_block_rho, &per_block_penalty_refs, 0.0)
2378 .map_err(EstimationError::InvalidInput)?
2379 } else {
2380 PenaltyLogdetDerivs {
2381 value: 0.0,
2382 first: Array1::zeros(0),
2383 second: Some(Array2::zeros((0, 0))),
2384 }
2385 };
2386
2387 let penalty_quadratic = 2.0 * state.penalty_term;
2389 let provider = SurvivalDerivProvider::new(self.clone(), beta.clone());
2390
2391 const SURVIVAL_LAML_IFT_RELATIVE_KKT_GATE: f64 = 1.0e-8;
2405 let kkt_residual = {
2406 let raw = state.gradient.clone();
2407 let projected = match self.monotonicity_linear_constraints() {
2408 Some(constraints) => {
2409 projected_linear_constraint_stationarity_vector(&raw, beta, &constraints, None)
2410 .ok_or_else(|| {
2411 EstimationError::InvalidInput(
2412 "survival LAML could not project the monotonicity KKT residual"
2413 .to_string(),
2414 )
2415 })?
2416 }
2417 None => raw,
2418 };
2419 let projected_norm = array1_l2_norm(&projected);
2420 let relative_projected_norm = state.relative_gradient_norm(projected_norm);
2421 if relative_projected_norm <= SURVIVAL_LAML_IFT_RELATIVE_KKT_GATE {
2422 Some(crate::model_types::ProjectedKktResidual::from_active_projected(projected))
2423 } else {
2424 None
2425 }
2426 };
2427
2428 let result = InnerAssembly {
2429 log_likelihood: state.log_likelihood,
2430 penalty_quadratic,
2431 beta: beta.clone(),
2432 n_observations: self.nrows(),
2433 hessian_op: std::sync::Arc::new(hop),
2434 penalty_coords,
2435 penalty_logdet,
2436 dispersion: DispersionHandling::Fixed {
2437 phi: 1.0,
2438 include_logdet_h: true,
2439 include_logdet_s: true,
2440 },
2441 rho_curvature_scale: 1.0,
2442 rho_prior: gam_problem::RhoPrior::Flat,
2443 hessian_logdet_correction: 0.0,
2444 penalty_subspace_trace: None,
2445 deriv_provider: Some(Box::new(provider)),
2446 firth: None,
2447 nullspace_dim: None,
2448 barrier_config: None,
2449 ext_coords: Vec::new(),
2450 ext_coord_pair_fn: None,
2451 rho_ext_pair_fn: None,
2452 fixed_drift_deriv: None,
2453 contracted_psi_second_order: None,
2454 kkt_residual,
2455 active_constraints: None,
2456 }
2457 .evaluate(
2458 rho.as_slice().expect("rho must be contiguous"),
2459 EvalMode::ValueAndGradient,
2460 None,
2461 )
2462 .map_err(EstimationError::InvalidInput)?;
2463
2464 let gradient = result.gradient.unwrap_or_else(|| Array1::zeros(rho.len()));
2465 Ok((result.cost, gradient))
2466 }
2467
2468 pub fn evaluate_survival_lamlcost_and_gradient(
2489 &self,
2490 rho: &[f64],
2491 beta0: &Array1<f64>,
2492 ) -> Result<(f64, Array1<f64>), EstimationError> {
2493 let (candidate, beta) = self.reconverge_survival_inner_mode(rho, beta0)?;
2494 let rho_arr = Array1::from_vec(rho.to_vec());
2499 let state = candidate.update_state(&beta)?;
2500 candidate.unified_lamlobjective_and_rhogradient(&beta, &state, &rho_arr)
2501 }
2502
2503 fn reconverge_survival_inner_mode(
2513 &self,
2514 rho: &[f64],
2515 beta0: &Array1<f64>,
2516 ) -> Result<(WorkingModelSurvival, Array1<f64>), EstimationError> {
2517 const SHIM_PIRLS_MAX_ITERATIONS: usize = 600;
2522 const SHIM_PIRLS_CONVERGENCE_TOL: f64 = 1e-12;
2523 const SHIM_PIRLS_MAX_STEP_HALVING: usize = 40;
2524 const SHIM_PIRLS_MIN_STEP_SIZE: f64 = 1e-12;
2525
2526 let active_block_count = self
2527 .penalties
2528 .blocks
2529 .iter()
2530 .filter(|b| b.lambda > 0.0)
2531 .count();
2532 if rho.len() != active_block_count {
2533 crate::bail_invalid_estim!(
2534 "reconverge_survival_inner_mode: rho dimension {} does not match active penalty block count {}",
2535 rho.len(),
2536 active_block_count
2537 );
2538 }
2539 if beta0.len() != self.coefficient_dim() {
2540 crate::bail_invalid_estim!(
2541 "reconverge_survival_inner_mode: beta0 dimension {} does not match coefficient dimension {}",
2542 beta0.len(),
2543 self.coefficient_dim()
2544 );
2545 }
2546
2547 let mut candidate = self.clone();
2550 let mut lambdas: Vec<f64> = candidate
2551 .penalties
2552 .blocks
2553 .iter()
2554 .map(|b| b.lambda)
2555 .collect();
2556 let mut active_idx = 0usize;
2557 for (block, lambda) in candidate.penalties.blocks.iter().zip(lambdas.iter_mut()) {
2558 if block.lambda > 0.0 {
2559 *lambda = rho[active_idx].exp();
2560 active_idx += 1;
2561 }
2562 }
2563 candidate.set_penalty_lambdas(&lambdas)?;
2564
2565 let opts = gam_solve::pirls::WorkingModelPirlsOptions {
2566 max_iterations: SHIM_PIRLS_MAX_ITERATIONS,
2567 convergence_tolerance: SHIM_PIRLS_CONVERGENCE_TOL,
2568 adaptive_kkt_tolerance: None,
2569 max_step_halving: SHIM_PIRLS_MAX_STEP_HALVING,
2570 min_step_size: SHIM_PIRLS_MIN_STEP_SIZE,
2571 firth_bias_reduction: false,
2572 coefficient_lower_bounds: None,
2573 linear_constraints: None,
2574 initial_lm_lambda: None,
2575 arrow_schur: None,
2576 };
2577 let summary = gam_solve::pirls::runworking_model_pirls(
2578 &mut candidate,
2579 Coefficients::new(beta0.clone()),
2580 &opts,
2581 |_| {},
2582 )?;
2583 let mut beta = summary.beta.as_ref().to_owned();
2584
2585 {
2607 const POLISH_MAX_ITERS: usize = 400;
2608 const POLISH_TOL: f64 = 1e-13;
2609 const ARMIJO_C: f64 = constants::ARMIJO_C1;
2613 const BACKTRACK: f64 = constants::BACKTRACK_CONTRACTION;
2614 const MAX_BACKTRACK: usize = 80;
2615 let p = beta.len();
2616 let penalized_objective =
2621 |st: &WorkingState| -> f64 { -st.log_likelihood + st.penalty_term };
2622 for _ in 0..POLISH_MAX_ITERS {
2623 let st = match candidate.update_state(&beta) {
2624 Ok(st) => st,
2625 Err(_) => break,
2626 };
2627 let r = st.gradient.clone();
2628 let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
2629 if !r_norm.is_finite() || r_norm < POLISH_TOL {
2630 break;
2631 }
2632 let h = st.hessian.to_dense();
2633 let f0 = penalized_objective(&st);
2634 let h_scale = (0..p)
2649 .map(|d| h[[d, d]].abs())
2650 .fold(0.0_f64, f64::max)
2651 .max(1.0);
2652 let try_lm = |lambda_lm: f64| -> Option<(Array1<f64>, f64)> {
2668 let mut h_reg = h.clone();
2669 for d in 0..p {
2670 h_reg[[d, d]] += lambda_lm;
2671 }
2672 let factor =
2673 gam_linalg::faer_ndarray::FaerCholesky::cholesky(&h_reg, faer::Side::Lower)
2674 .ok()?;
2675 let candidate_step = factor.solvevec(&r);
2676 if candidate_step.iter().any(|v| !v.is_finite()) {
2677 return None;
2678 }
2679 let dd = -r.dot(&candidate_step);
2680 (dd.is_finite() && dd < -1e-14 * r_norm * r_norm)
2681 .then_some((candidate_step, dd))
2682 };
2683 let (step, dir_deriv) = try_lm(0.0)
2686 .or_else(|| {
2687 escalate_ridge(RidgeSchedule::geometric(1e-11 * h_scale, 17), try_lm)
2688 .ok()
2689 .map(|success| success.value)
2690 })
2691 .unwrap_or_else(|| {
2692 (r.clone(), -r_norm * r_norm)
2695 });
2696 let accepted = match backtracking_line_search::<_, Infallible>(
2713 BacktrackConfig {
2714 contraction: BACKTRACK,
2715 max_steps: MAX_BACKTRACK,
2716 ..BacktrackConfig::default()
2717 },
2718 |alpha| {
2719 let trial = &beta - &(alpha * &step);
2720 let Ok(ts) = candidate.update_state(&trial) else {
2721 return Ok(None);
2722 };
2723 let ft = penalized_objective(&ts);
2724 let tn = ts.gradient.iter().map(|v| v * v).sum::<f64>().sqrt();
2725 let armijo_ok = ft.is_finite() && ft <= f0 + ARMIJO_C * alpha * dir_deriv;
2726 let residual_ok = tn.is_finite() && tn < r_norm;
2727 Ok((armijo_ok || residual_ok).then_some((ft, trial)))
2728 },
2729 |_alpha, _ft| true,
2730 ) {
2731 Ok(result) => result,
2732 Err(never) => match never {},
2733 };
2734 let Some(ls) = accepted else {
2735 break;
2736 };
2737 beta = ls.payload;
2738 }
2739 }
2740
2741 Ok((candidate, beta))
2742 }
2743}
2744
2745pub(crate) struct SurvivalDerivProvider {
2754 model: WorkingModelSurvival,
2755 beta: Array1<f64>,
2756}
2757
2758impl SurvivalDerivProvider {
2759 pub(crate) fn new(model: WorkingModelSurvival, beta: Array1<f64>) -> Self {
2760 Self { model, beta }
2761 }
2762}
2763
2764impl gam_solve::estimate::reml::reml_outer_engine::HessianDerivativeProvider
2765 for SurvivalDerivProvider
2766{
2767 fn hessian_derivative_correction(
2768 &self,
2769 v_k: &Array1<f64>,
2770 ) -> Result<Option<Array2<f64>>, String> {
2771 let u_k = -v_k;
2774 match self
2775 .model
2776 .survival_hessian_derivative_correction(&self.beta, &u_k)
2777 {
2778 Ok(correction) => Ok(Some(correction)),
2779 Err(e) => Err(e.to_string()),
2780 }
2781 }
2782
2783 fn has_corrections(&self) -> bool {
2784 true
2785 }
2786}
2787
2788#[derive(Debug, Clone)]
2789pub struct CrudeRiskResult {
2790 pub risk: f64,
2791 pub diseasegradient: Array1<f64>,
2792 pub mortalitygradient: Array1<f64>,
2793}
2794
2795#[derive(Debug, Clone)]
2796pub struct CompetingRisksCifResult {
2797 pub cif: Vec<Array2<f64>>,
2802 pub overall_survival: Array2<f64>,
2803}
2804
2805const COMPETING_RISKS_CIF_PARALLEL_ROW_MIN: usize = 256;
2810
2811pub fn assemble_competing_risks_cif(
2812 times: ArrayView1<'_, f64>,
2813 cumulative_hazard: ArrayView3<'_, f64>,
2814) -> Result<CompetingRisksCifResult, SurvivalError> {
2815 let (n_endpoints, n_rows, n_times) = cumulative_hazard.dim();
2816 if n_endpoints == 0 {
2817 return Err(SurvivalError::DimensionMismatch);
2818 }
2819 let endpoint_hazards = cumulative_hazard
2820 .axis_iter(Axis(0))
2821 .map(|view| view.to_owned())
2822 .collect::<Vec<_>>();
2823 assemble_competing_risks_cif_from_endpoints(times, &endpoint_hazards).and_then(|result| {
2824 if result.overall_survival.dim() != (n_rows, n_times) {
2825 Err(SurvivalError::DimensionMismatch)
2826 } else {
2827 Ok(result)
2828 }
2829 })
2830}
2831
2832pub fn assemble_competing_risks_cif_from_endpoints(
2833 times: ArrayView1<'_, f64>,
2834 cumulative_hazards: &[Array2<f64>],
2835) -> Result<CompetingRisksCifResult, SurvivalError> {
2836 let n_endpoints = cumulative_hazards.len();
2837 if n_endpoints == 0 || times.is_empty() {
2838 return Err(SurvivalError::DimensionMismatch);
2839 }
2840 let (n_rows, n_times) = cumulative_hazards[0].dim();
2841 if n_rows == 0 || n_times == 0 || times.len() != n_times {
2842 return Err(SurvivalError::DimensionMismatch);
2843 }
2844 if times.iter().any(|time| !time.is_finite() || *time < 0.0) {
2845 return Err(SurvivalError::InvalidTimeGrid);
2846 }
2847 if times
2848 .iter()
2849 .zip(times.iter().skip(1))
2850 .any(|(previous, current)| current <= previous)
2851 {
2852 return Err(SurvivalError::InvalidTimeGrid);
2853 }
2854 for endpoint_hazard in cumulative_hazards {
2855 if endpoint_hazard.dim() != (n_rows, n_times) {
2856 return Err(SurvivalError::DimensionMismatch);
2857 }
2858 if endpoint_hazard.iter().any(|value| !value.is_finite()) {
2859 return Err(SurvivalError::NonFiniteInput);
2860 }
2861 }
2862
2863 let max_abs_hazard = cumulative_hazards
2864 .iter()
2865 .flat_map(|endpoint_hazard| endpoint_hazard.iter())
2866 .fold(0.0_f64, |acc, value| acc.max(value.abs()));
2867 let monotone_tolerance = 1.0e-10_f64 * max_abs_hazard.max(1.0);
2868 let mut cif: Vec<Array2<f64>> = (0..n_endpoints)
2869 .map(|_| Array2::<f64>::zeros((n_rows, n_times)))
2870 .collect();
2871 let mut overall_survival = Array2::<f64>::zeros((n_rows, n_times));
2872
2873 let assemble_row = |row: usize| -> Result<(Vec<f64>, Vec<f64>), SurvivalError> {
2885 let mut cif_flat = vec![0.0_f64; n_endpoints * n_times];
2886 let mut surv_row = vec![0.0_f64; n_times];
2887 let mut previous_cif = vec![0.0_f64; n_endpoints];
2888 let mut previous_cumulative = vec![0.0_f64; n_endpoints];
2889 let mut increments = vec![0.0_f64; n_endpoints];
2890 let mut previous_total_cumulative = 0.0_f64;
2891 for time_idx in 0..n_times {
2892 let mut total_increment = 0.0_f64;
2893 for endpoint in 0..n_endpoints {
2894 let current = cumulative_hazards[endpoint][[row, time_idx]];
2895 if current < -monotone_tolerance {
2896 return Err(SurvivalError::NonMonotoneCumulativeHazard);
2897 }
2898 let raw_increment = current - previous_cumulative[endpoint];
2899 if raw_increment < -monotone_tolerance {
2900 return Err(SurvivalError::NonMonotoneCumulativeHazard);
2901 }
2902 let increment = raw_increment.max(0.0);
2903 increments[endpoint] = increment;
2904 total_increment += increment;
2905 previous_cumulative[endpoint] += increment;
2906 }
2907
2908 let survival_left = (-previous_total_cumulative).exp();
2909 let interval_failure = -(-total_increment).exp_m1();
2910 for endpoint in 0..n_endpoints {
2911 if total_increment > 0.0 {
2912 previous_cif[endpoint] +=
2913 survival_left * interval_failure * increments[endpoint] / total_increment;
2914 }
2915 cif_flat[endpoint * n_times + time_idx] = previous_cif[endpoint].clamp(0.0, 1.0);
2916 }
2917 previous_total_cumulative += total_increment;
2918 let mut fsum_at_t = 0.0_f64;
2935 for endpoint in 0..n_endpoints {
2936 fsum_at_t += cif_flat[endpoint * n_times + time_idx];
2937 }
2938 surv_row[time_idx] = (1.0_f64 - fsum_at_t).clamp(0.0, 1.0);
2939 }
2940 Ok((cif_flat, surv_row))
2941 };
2942
2943 let rows: Vec<(Vec<f64>, Vec<f64>)> = if n_rows >= COMPETING_RISKS_CIF_PARALLEL_ROW_MIN
2947 && rayon::current_thread_index().is_none()
2948 {
2949 use rayon::prelude::*;
2950 (0..n_rows)
2951 .into_par_iter()
2952 .map(assemble_row)
2953 .collect::<Result<_, _>>()?
2954 } else {
2955 (0..n_rows).map(assemble_row).collect::<Result<_, _>>()?
2956 };
2957
2958 for (row, (cif_flat, surv_row)) in rows.into_iter().enumerate() {
2959 for endpoint in 0..n_endpoints {
2960 for time_idx in 0..n_times {
2961 cif[endpoint][[row, time_idx]] = cif_flat[endpoint * n_times + time_idx];
2962 }
2963 }
2964 for time_idx in 0..n_times {
2965 overall_survival[[row, time_idx]] = surv_row[time_idx];
2966 }
2967 }
2968
2969 Ok(CompetingRisksCifResult {
2970 cif,
2971 overall_survival,
2972 })
2973}
2974
2975fn compute_gauss_legendre_nodes(n: usize) -> Vec<(f64, f64)> {
2979 let (nodes, weights) = gam_math::special::gauss_legendre(n);
2980 nodes.into_iter().zip(weights).collect()
2981}
2982
2983fn gauss_legendre_quadrature() -> &'static [(f64, f64)] {
2984 static CACHE: LazyLock<Vec<(f64, f64)>> = LazyLock::new(|| compute_gauss_legendre_nodes(40));
2990 &CACHE
2991}
2992
2993pub fn calculate_crude_risk_quadrature<F>(
3017 t0: f64,
3018 t1: f64,
3019 breakpoints: &[f64],
3020 h_dis_t0: f64,
3021 h_mor_t0: f64,
3022 design_d_t0: ArrayView1<'_, f64>,
3023 design_m_t0: ArrayView1<'_, f64>,
3024 mut eval_at: F,
3025) -> Result<CrudeRiskResult, SurvivalError>
3026where
3027 F: FnMut(
3028 f64,
3029 &mut Array1<f64>,
3030 &mut Array1<f64>,
3031 &mut Array1<f64>,
3032 ) -> Result<(f64, f64, f64), SurvivalError>,
3033{
3034 let coeff_len_d = design_d_t0.len();
3035 let coeff_len_m = design_m_t0.len();
3036 if coeff_len_d == 0 || coeff_len_m == 0 {
3037 return Err(SurvivalError::InvalidIntegrationSetup);
3038 }
3039 if !t0.is_finite()
3040 || !t1.is_finite()
3041 || !h_dis_t0.is_finite()
3042 || !h_mor_t0.is_finite()
3043 || design_d_t0.iter().any(|v| !v.is_finite())
3044 || design_m_t0.iter().any(|v| !v.is_finite())
3045 {
3046 return Err(SurvivalError::NonFiniteInput);
3047 }
3048 if t1 <= t0 {
3049 return Ok(CrudeRiskResult {
3050 risk: 0.0,
3051 diseasegradient: Array1::zeros(coeff_len_d),
3052 mortalitygradient: Array1::zeros(coeff_len_m),
3053 });
3054 }
3055
3056 let mut sorted_breaks: Vec<f64> = breakpoints
3057 .iter()
3058 .copied()
3059 .filter(|x| x.is_finite() && *x >= t0 && *x <= t1)
3060 .collect();
3061 sorted_breaks.push(t0);
3062 sorted_breaks.push(t1);
3063 sorted_breaks.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
3064 sorted_breaks.dedup_by(|a, b| (*a - *b).abs() < 1e-6);
3065 if sorted_breaks.len() < 2 {
3066 return Err(SurvivalError::InvalidIntegrationSetup);
3067 }
3068
3069 let mut total_risk = 0.0;
3070 let mut diseasegradient = Array1::zeros(coeff_len_d);
3071 let mut mortalitygradient = Array1::zeros(coeff_len_m);
3072 let nodesweights = gauss_legendre_quadrature();
3073
3074 let mut design_d = Array1::<f64>::zeros(coeff_len_d);
3075 let mut deriv_d = Array1::<f64>::zeros(coeff_len_d);
3076 let mut design_m = Array1::<f64>::zeros(coeff_len_m);
3077
3078 for segment in sorted_breaks.windows(2) {
3079 let a = segment[0];
3080 let b = segment[1];
3081 let center = 0.5 * (b + a);
3082 let halfwidth = 0.5 * (b - a);
3083 if halfwidth <= 0.0 {
3084 continue;
3085 }
3086
3087 for &(x, w) in nodesweights {
3088 let u = center + halfwidth * x;
3089 let (inst_hazard_d, hazard_d, hazard_m) =
3090 eval_at(u, &mut design_d, &mut deriv_d, &mut design_m)?;
3091 if !inst_hazard_d.is_finite() || !hazard_d.is_finite() || !hazard_m.is_finite() {
3092 return Err(SurvivalError::NonFiniteInput);
3093 }
3094 if inst_hazard_d <= 0.0 {
3095 return Err(SurvivalError::NonPositiveHazard);
3096 }
3097
3098 if hazard_d < h_dis_t0 || hazard_m < h_mor_t0 {
3099 return Err(SurvivalError::NonMonotoneCumulativeHazard);
3100 }
3101
3102 let h_dis_cond = hazard_d - h_dis_t0;
3103 let h_mor_cond = hazard_m - h_mor_t0;
3104 let s_total = (-(h_dis_cond + h_mor_cond)).exp();
3105
3106 total_risk += w * inst_hazard_d * s_total * halfwidth;
3107
3108 let weight = w * s_total * halfwidth;
3114 for j in 0..coeff_len_d {
3115 let d_inst_hazard = inst_hazard_d * design_d[j] + hazard_d * deriv_d[j];
3116 let d_hazard_cond = hazard_d * design_d[j] - h_dis_t0 * design_d_t0[j];
3117 let g = d_inst_hazard - inst_hazard_d * d_hazard_cond;
3118 diseasegradient[j] += weight * g;
3119 }
3120
3121 let weight = w * inst_hazard_d * s_total * halfwidth;
3124 for j in 0..coeff_len_m {
3125 let g = -hazard_m * design_m[j] + h_mor_t0 * design_m_t0[j];
3126 mortalitygradient[j] += weight * g;
3127 }
3128 }
3129 }
3130
3131 Ok(CrudeRiskResult {
3132 risk: total_risk,
3133 diseasegradient,
3134 mortalitygradient,
3135 })
3136}
3137
3138impl PirlsWorkingModel for WorkingModelSurvival {
3139 fn update(&mut self, beta: &Coefficients) -> Result<WorkingState, EstimationError> {
3140 self.update_state(beta)
3141 }
3142}
3143
3144#[cfg(test)]
3145mod tests {
3146 use super::*;
3147 use ndarray::{Array1, Array2, Array3, array, s};
3148
3149 mod jet_cause_specific_production_parity {
3160 use super::*;
3161 use gam_math::jet_scalar::JetScalar;
3162 use gam_math::jet_tower::{
3163 RowNllProgramGeneric, generic_fourth_contracted, generic_row_kernel,
3164 generic_third_contracted,
3165 };
3166
3167 struct CauseSpecificJetRow {
3173 has_entry: bool,
3174 event: bool,
3175 w: f64,
3176 base: [f64; 3],
3177 }
3178
3179 impl RowNllProgramGeneric<3> for CauseSpecificJetRow {
3180 fn n_rows(&self) -> usize {
3181 1
3182 }
3183 fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
3184 if row != 0 {
3185 return Err(format!(
3186 "CauseSpecificJetRow holds exactly one row; got row {row}"
3187 ));
3188 }
3189 Ok(self.base)
3190 }
3191 fn row_nll_generic<S: JetScalar<3>>(
3192 &self,
3193 row: usize,
3194 p: &[S; 3],
3195 ) -> Result<S, String> {
3196 if row != 0 {
3197 return Err(format!(
3198 "CauseSpecificJetRow holds exactly one row; got row {row}"
3199 ));
3200 }
3201 let mut ell = p[0].exp();
3202 if self.has_entry {
3203 ell = ell.sub(&p[1].exp());
3204 }
3205 if self.event {
3206 ell = ell.sub(&p[0].add(&p[2].ln()));
3207 }
3208 Ok(ell.scale(self.w))
3209 }
3210 }
3211
3212 fn identity_block(w: f64, has_entry: bool, event: bool) -> CauseSpecificRoystonParmarBlock {
3219 let age_entry = if has_entry { 1.0 } else { 0.0 };
3220 CauseSpecificRoystonParmarBlock {
3221 age_entry: array![age_entry],
3222 age_exit: array![2.0],
3223 event_target: array![if event { 1u8 } else { 0u8 }],
3224 sampleweight: array![w],
3225 x_entry: array![[0.0, 1.0, 0.0]],
3226 x_exit: array![[1.0, 0.0, 0.0]],
3227 x_derivative: array![[0.0, 0.0, 1.0]],
3228 offset_eta_entry: array![0.0],
3229 offset_eta_exit: array![0.0],
3230 offset_derivative_exit: array![0.0],
3231 derivative_floor: 0.0,
3232 }
3233 }
3234
3235 fn close(hand: f64, jet: f64, tol: f64, label: &str) {
3236 let band = tol + tol * hand.abs().max(jet.abs());
3237 assert!(
3238 (hand - jet).abs() <= band,
3239 "{label}: hand {hand:+.15e} vs jet {jet:+.15e} (|Δ|={:.3e} band {band:.3e})",
3240 (hand - jet).abs()
3241 );
3242 }
3243
3244 const JET_TOL: f64 = 1e-9;
3245
3246 fn run_corner(has_entry: bool, event: bool) {
3247 let beta = array![0.4_f64, -0.3_f64, 1.3_f64];
3249 let d_beta = array![0.7_f64, -0.5_f64, 0.6_f64];
3250 let v_beta = array![-0.2_f64, 0.8_f64, -0.4_f64];
3251 let w = 1.4_f64;
3252 let block = identity_block(w, has_entry, event);
3253 let prog = CauseSpecificJetRow {
3254 has_entry,
3255 event,
3256 w,
3257 base: [beta[0], beta[1], beta[2]],
3258 };
3259 let label = format!("entry={has_entry} event={event}");
3260
3261 let (ll, grad, hess) =
3263 evaluate_cause_specific_block(&block, &beta).expect("evaluate block");
3264 let (jet_v, jet_g, jet_h) = generic_row_kernel(&prog, 0).expect("jet kernel");
3265 close(jet_v, -ll, JET_TOL, &format!("{label} value"));
3266 for a in 0..3 {
3267 close(jet_g[a], -grad[a], JET_TOL, &format!("{label} grad[{a}]"));
3268 for b in 0..3 {
3269 close(jet_h[a][b], hess[[a, b]], JET_TOL, &format!("{label} H[{a}][{b}]"));
3270 }
3271 }
3272
3273 let dh = cause_specific_hessian_directional_derivative(&block, &beta, &d_beta)
3275 .expect("live third");
3276 let dir = [d_beta[0], d_beta[1], d_beta[2]];
3277 let jet_t3 = generic_third_contracted(&prog, 0, &dir).expect("jet third");
3278 for a in 0..3 {
3279 for b in 0..3 {
3280 close(
3281 jet_t3[a][b],
3282 dh[[a, b]],
3283 JET_TOL,
3284 &format!("{label} third[{a}][{b}]"),
3285 );
3286 }
3287 }
3288
3289 let d2h =
3291 cause_specific_hessian_second_directional_derivative(&block, &beta, &d_beta, &v_beta)
3292 .expect("live fourth");
3293 let uu = [d_beta[0], d_beta[1], d_beta[2]];
3294 let vv = [v_beta[0], v_beta[1], v_beta[2]];
3295 let jet_t4 = generic_fourth_contracted(&prog, 0, &uu, &vv).expect("jet fourth");
3296 for a in 0..3 {
3297 for b in 0..3 {
3298 close(
3299 jet_t4[a][b],
3300 d2h[[a, b]],
3301 JET_TOL,
3302 &format!("{label} fourth[{a}][{b}]"),
3303 );
3304 }
3305 }
3306
3307 let h_fd = 1e-5;
3310 let bp = &beta + &(&d_beta * h_fd);
3311 let bm = &beta - &(&d_beta * h_fd);
3312 let (_, _, hp) = evaluate_cause_specific_block(&block, &bp).expect("evaluate +");
3313 let (_, _, hm) = evaluate_cause_specific_block(&block, &bm).expect("evaluate -");
3314 for a in 0..3 {
3315 for b in 0..3 {
3316 let fd = (hp[[a, b]] - hm[[a, b]]) / (2.0 * h_fd);
3317 close(dh[[a, b]], fd, 1e-5, &format!("{label} FD third[{a}][{b}]"));
3318 }
3319 }
3320 let dhp = cause_specific_hessian_directional_derivative(&block, &bp_along(&beta, &v_beta, h_fd), &d_beta)
3322 .expect("live third +");
3323 let dhm = cause_specific_hessian_directional_derivative(&block, &bm_along(&beta, &v_beta, h_fd), &d_beta)
3324 .expect("live third -");
3325 for a in 0..3 {
3326 for b in 0..3 {
3327 let fd = (dhp[[a, b]] - dhm[[a, b]]) / (2.0 * h_fd);
3328 close(d2h[[a, b]], fd, 1e-5, &format!("{label} FD fourth[{a}][{b}]"));
3329 }
3330 }
3331 }
3332
3333 fn bp_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
3334 beta + &(v * h)
3335 }
3336 fn bm_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
3337 beta - &(v * h)
3338 }
3339
3340 #[test]
3346 fn cause_specific_live_tower_matches_jet_and_fd() {
3347 for &has_entry in &[false, true] {
3348 for &event in &[false, true] {
3349 run_corner(has_entry, event);
3350 }
3351 }
3352 }
3353 }
3354
3355 #[test]
3356 fn competing_risks_cif_constant_hazard_matches_closed_form() {
3357 let times = array![0.0, 2.0, 5.0, 10.0];
3358 let disease_rates = [0.12, 0.06];
3359 let death_rates = [0.05, 0.02];
3360 let cumulative = Array3::from_shape_fn((2, 2, times.len()), |(endpoint, row, time_idx)| {
3361 let rate = if endpoint == 0 {
3362 disease_rates[row]
3363 } else {
3364 death_rates[row]
3365 };
3366 rate * times[time_idx]
3367 });
3368
3369 let result =
3370 assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
3371
3372 for row in 0..2 {
3373 let total_rate = disease_rates[row] + death_rates[row];
3374 for time_idx in 0..times.len() {
3375 let failure = 1.0 - (-total_rate * times[time_idx]).exp();
3376 let expected_disease = disease_rates[row] / total_rate * failure;
3377 let expected_death = death_rates[row] / total_rate * failure;
3378 assert!((result.cif[0][[row, time_idx]] - expected_disease).abs() < 1e-12);
3379 assert!((result.cif[1][[row, time_idx]] - expected_death).abs() < 1e-12);
3380 assert!(
3381 (result.cif[0][[row, time_idx]]
3382 + result.cif[1][[row, time_idx]]
3383 + result.overall_survival[[row, time_idx]]
3384 - 1.0)
3385 .abs()
3386 < 1e-12
3387 );
3388 }
3389 }
3390 }
3391
3392 #[test]
3393 fn competing_risks_cif_rejects_nonmonotone_hazards() {
3394 let times = array![0.0, 1.0, 2.0];
3395 let cumulative = Array3::from_shape_vec((1, 1, 3), vec![0.0, 0.2, 0.1]).expect("shape");
3396 let err = assemble_competing_risks_cif(times.view(), cumulative.view())
3397 .expect_err("nonmonotone cumulative hazard should be rejected");
3398 assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
3399 }
3400
3401 #[test]
3402 fn competing_risks_cif_plateaus_and_three_causes_conserve_probability() {
3403 let times = array![0.0, 1.0, 3.0, 7.0, 12.0];
3404 let cumulative = Array3::from_shape_vec(
3405 (3, 2, 5),
3406 vec![
3407 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,
3411 ],
3412 )
3413 .expect("shape");
3414
3415 let result =
3416 assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
3417
3418 for row in 0..2 {
3419 for time_idx in 0..times.len() {
3420 let total_cif = result.cif[0][[row, time_idx]]
3421 + result.cif[1][[row, time_idx]]
3422 + result.cif[2][[row, time_idx]];
3423 assert!(
3424 (total_cif + result.overall_survival[[row, time_idx]] - 1.0).abs() < 1e-12,
3425 "probability mass mismatch at row={row}, time_idx={time_idx}"
3426 );
3427 assert!((0.0..=1.0).contains(&result.overall_survival[[row, time_idx]]));
3428 for cause in 0..3 {
3429 assert!((0.0..=1.0).contains(&result.cif[cause][[row, time_idx]]));
3430 if time_idx > 0 {
3431 assert!(
3432 result.cif[cause][[row, time_idx]] + 1e-12
3433 >= result.cif[cause][[row, time_idx - 1]],
3434 "CIF decreased for cause={cause}, row={row}, time_idx={time_idx}"
3435 );
3436 }
3437 }
3438 }
3439 }
3440
3441 assert_eq!(result.cif[0][[0, 1]], result.cif[0][[0, 2]]);
3444 assert_eq!(result.cif[0][[1, 2]], result.cif[0][[1, 3]]);
3447 assert_eq!(result.cif[2][[1, 2]], result.cif[2][[1, 3]]);
3448 }
3449
3450 #[test]
3451 fn competing_risks_cif_rejects_bad_time_grids_and_nonfinite_hazards() {
3452 let cumulative = Array3::zeros((2, 1, 2));
3453
3454 for times in [array![0.0, 0.0], array![1.0, 0.5], array![-1.0, 1.0]] {
3455 let err = assemble_competing_risks_cif(times.view(), cumulative.view())
3456 .expect_err("bad time grid should be rejected");
3457 assert!(matches!(err, SurvivalError::InvalidTimeGrid));
3458 }
3459
3460 let times = array![0.0, 1.0];
3461 let nonfinite = Array3::from_shape_vec((1, 1, 2), vec![0.0, f64::NAN]).expect("shape");
3462 let err = assemble_competing_risks_cif(times.view(), nonfinite.view())
3463 .expect_err("nonfinite hazard should be rejected");
3464 assert!(matches!(err, SurvivalError::NonFiniteInput));
3465 }
3466
3467 #[test]
3468 fn competing_risks_cif_extreme_hazards_remain_bounded() {
3469 let times = array![0.0, 1.0, 2.0];
3470 let cumulative =
3471 Array3::from_shape_vec((2, 1, 3), vec![0.0, 500.0, 1000.0, 0.0, 250.0, 1000.0])
3472 .expect("shape");
3473
3474 let result =
3475 assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
3476
3477 for value in result
3478 .cif
3479 .iter()
3480 .flat_map(|m| m.iter())
3481 .chain(result.overall_survival.iter())
3482 {
3483 assert!(value.is_finite());
3484 assert!((0.0..=1.0).contains(value));
3485 }
3486 assert!((result.cif[0][[0, 2]] + result.cif[1][[0, 2]] - 1.0).abs() < 1e-12);
3487 assert_eq!(result.overall_survival[[0, 2]], 0.0);
3488 }
3489
3490 fn toy_penalties() -> PenaltyBlocks {
3491 let s = array![[2.0, 0.5], [0.5, 3.0]];
3492 PenaltyBlocks::new(vec![PenaltyBlock {
3493 matrix: s,
3494 lambda: 1.7,
3495 range: 1..3,
3496 nullspace_dim: 0,
3497 }])
3498 }
3499
3500 fn survival_inputs<'a>(
3501 age_entry: &'a Array1<f64>,
3502 age_exit: &'a Array1<f64>,
3503 event_target: &'a Array1<u8>,
3504 event_competing: &'a Array1<u8>,
3505 sampleweight: &'a Array1<f64>,
3506 x_entry: &'a Array2<f64>,
3507 x_exit: &'a Array2<f64>,
3508 x_derivative: &'a Array2<f64>,
3509 ) -> SurvivalEngineInputs<'a> {
3510 SurvivalEngineInputs {
3511 age_entry: age_entry.view(),
3512 age_exit: age_exit.view(),
3513 event_target: event_target.view(),
3514 event_competing: event_competing.view(),
3515 sampleweight: sampleweight.view(),
3516 x_entry: x_entry.view(),
3517 x_exit: x_exit.view(),
3518 x_derivative: x_derivative.view(),
3519 monotonicity_constraint_rows: None,
3520 monotonicity_constraint_offsets: None,
3521 }
3522 }
3523
3524 fn survival_model(
3525 inputs: SurvivalEngineInputs<'_>,
3526 penalties: PenaltyBlocks,
3527 monotonicity: SurvivalMonotonicityPenalty,
3528 spec: SurvivalSpec,
3529 ) -> Result<WorkingModelSurvival, SurvivalError> {
3530 WorkingModelSurvival::from_engine_inputs(inputs, penalties, monotonicity, spec)
3531 }
3532
3533 fn survival_model_with_offsets(
3534 inputs: SurvivalEngineInputs<'_>,
3535 offsets: Option<SurvivalBaselineOffsets<'_>>,
3536 penalties: PenaltyBlocks,
3537 monotonicity: SurvivalMonotonicityPenalty,
3538 spec: SurvivalSpec,
3539 ) -> Result<WorkingModelSurvival, SurvivalError> {
3540 WorkingModelSurvival::from_engine_inputswith_offsets(
3541 inputs,
3542 offsets,
3543 penalties,
3544 monotonicity,
3545 spec,
3546 )
3547 }
3548
3549 #[test]
3550 fn penaltyhessian_matchesgradient_jacobian() {
3551 let penalties = toy_penalties();
3552 let beta = array![10.0, -0.3, 1.2, 7.0];
3553
3554 let grad = penalties.gradient(&beta);
3555 let h = penalties.hessian(beta.len());
3556 let b_block = beta.slice(s![1..3]).to_owned();
3557 let expected = 1.7 * array![[2.0, 0.5], [0.5, 3.0]].dot(&b_block);
3558
3559 assert!((grad[1] - expected[0]).abs() < 1e-12);
3560 assert!((grad[2] - expected[1]).abs() < 1e-12);
3561 assert!((h[[1, 1]] - 1.7 * 2.0).abs() < 1e-12);
3562 assert!((h[[1, 2]] - 1.7 * 0.5).abs() < 1e-12);
3563 assert!((h[[2, 1]] - 1.7 * 0.5).abs() < 1e-12);
3564 assert!((h[[2, 2]] - 1.7 * 3.0).abs() < 1e-12);
3565 }
3566
3567 #[test]
3568 fn penaltygradient_matches_deviance_finite_difference() {
3569 let penalties = toy_penalties();
3570 let beta = array![10.0, -0.3, 1.2, 7.0];
3571 let grad = penalties.gradient(&beta);
3572 let eps = 1e-7;
3573
3574 for idx in 0..beta.len() {
3575 let mut plus = beta.clone();
3576 let mut minus = beta.clone();
3577 plus[idx] += eps;
3578 minus[idx] -= eps;
3579 let fd = (penalties.deviance(&plus) - penalties.deviance(&minus)) / (2.0 * eps);
3580 assert_eq!(
3581 grad[idx].signum(),
3582 fd.signum(),
3583 "gradient/deviance sign mismatch at idx={idx}: grad={} fd={fd}",
3584 grad[idx]
3585 );
3586 assert!(
3587 (grad[idx] - fd).abs() < 1e-6,
3588 "gradient/deviance mismatch at idx={idx}: grad={} fd={fd}",
3589 grad[idx]
3590 );
3591 }
3592 }
3593
3594 #[test]
3595 fn zero_offsets_match_default_survival_state() {
3596 let age_entry = array![1.0_f64, 2.0_f64];
3597 let age_exit = array![2.0_f64, 3.5_f64];
3598 let event_target = array![1u8, 0u8];
3599 let event_competing = array![0u8, 0u8];
3600 let sampleweight = array![1.0, 1.0];
3601 let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
3602 let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
3603 let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
3604 let penalties = PenaltyBlocks::new(Vec::new());
3605 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3606 let beta = array![-1.0, 0.8];
3607
3608 let base = survival_model(
3609 survival_inputs(
3610 &age_entry,
3611 &age_exit,
3612 &event_target,
3613 &event_competing,
3614 &sampleweight,
3615 &x_entry,
3616 &x_exit,
3617 &x_derivative,
3618 ),
3619 penalties.clone(),
3620 mono,
3621 SurvivalSpec::Net,
3622 )
3623 .expect("construct base survival model");
3624
3625 let zero_offsets = survival_model_with_offsets(
3626 survival_inputs(
3627 &age_entry,
3628 &age_exit,
3629 &event_target,
3630 &event_competing,
3631 &sampleweight,
3632 &x_entry,
3633 &x_exit,
3634 &x_derivative,
3635 ),
3636 Some(SurvivalBaselineOffsets {
3637 eta_entry: array![0.0, 0.0].view(),
3638 eta_exit: array![0.0, 0.0].view(),
3639 derivative_exit: array![0.0, 0.0].view(),
3640 }),
3641 penalties,
3642 mono,
3643 SurvivalSpec::Net,
3644 )
3645 .expect("construct offset survival model");
3646
3647 let state_base = base.update_state(&beta).expect("base state");
3648 let statezero = zero_offsets.update_state(&beta).expect("zero-offset state");
3649 assert!((state_base.deviance - statezero.deviance).abs() < 1e-12);
3650 assert!(
3651 state_base
3652 .gradient
3653 .iter()
3654 .zip(statezero.gradient.iter())
3655 .all(|(a, b)| (a - b).abs() < 1e-12)
3656 );
3657 }
3658
3659 #[test]
3660 fn competing_risk_cause_labels_collapse_to_pooled_baseline_indicator() {
3661 let age_entry = array![0.0_f64, 0.0, 0.0, 0.0];
3675 let age_exit = array![1.2_f64, 0.8, 2.1, 1.5];
3676 let cause_labels = array![0u8, 1u8, 2u8, 0u8];
3678 let event_competing = Array1::<u8>::zeros(cause_labels.len());
3679 let sampleweight = array![1.0_f64, 1.0, 1.0, 1.0];
3680 let x_entry = array![
3681 [1.0, age_entry[0].max(1e-8).ln()],
3682 [1.0, age_entry[1].max(1e-8).ln()],
3683 [1.0, age_entry[2].max(1e-8).ln()],
3684 [1.0, age_entry[3].max(1e-8).ln()],
3685 ];
3686 let x_exit = array![
3687 [1.0, age_exit[0].ln()],
3688 [1.0, age_exit[1].ln()],
3689 [1.0, age_exit[2].ln()],
3690 [1.0, age_exit[3].ln()],
3691 ];
3692 let x_derivative = array![
3693 [0.0, 1.0 / age_exit[0]],
3694 [0.0, 1.0 / age_exit[1]],
3695 [0.0, 1.0 / age_exit[2]],
3696 [0.0, 1.0 / age_exit[3]],
3697 ];
3698 let penalties = PenaltyBlocks::new(Vec::new());
3699 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3700
3701 let raw = survival_model(
3706 survival_inputs(
3707 &age_entry,
3708 &age_exit,
3709 &cause_labels,
3710 &event_competing,
3711 &sampleweight,
3712 &x_entry,
3713 &x_exit,
3714 &x_derivative,
3715 ),
3716 penalties.clone(),
3717 mono,
3718 SurvivalSpec::Net,
3719 );
3720 assert!(
3721 matches!(raw, Err(SurvivalError::EventCodeInvalid { .. })),
3722 "raw competing-risks cause labels must be rejected as EventCodeInvalid (not NonFiniteInput), got {raw:?}"
3723 );
3724
3725 let any_event = pooled_any_event_indicator(cause_labels.view());
3728 assert_eq!(any_event, array![0u8, 1u8, 1u8, 0u8]);
3729 assert_eq!(
3731 cause_specific_event_indicator(cause_labels.view(), 1),
3732 array![0u8, 1u8, 0u8, 0u8]
3733 );
3734 assert_eq!(
3735 cause_specific_event_indicator(cause_labels.view(), 2),
3736 array![0u8, 0u8, 1u8, 0u8]
3737 );
3738 let model = survival_model(
3739 survival_inputs(
3740 &age_entry,
3741 &age_exit,
3742 &any_event,
3743 &event_competing,
3744 &sampleweight,
3745 &x_entry,
3746 &x_exit,
3747 &x_derivative,
3748 ),
3749 penalties,
3750 mono,
3751 SurvivalSpec::Net,
3752 )
3753 .expect("pooled any-event baseline model must construct from competing-risks data");
3754
3755 let beta = array![-1.0_f64, 0.8];
3758 let state = model.update_state(&beta).expect("pooled baseline state");
3759 assert!(
3760 state.deviance.is_finite(),
3761 "pooled baseline deviance must be finite, got {}",
3762 state.deviance
3763 );
3764 assert!(
3765 state.gradient.iter().all(|g| g.is_finite()),
3766 "pooled baseline gradient must be finite"
3767 );
3768 }
3769
3770 #[test]
3771 fn offset_channel_residuals_match_central_fd_of_nll() {
3772 let age_entry = array![0.5_f64, 0.0, 0.3];
3777 let age_exit = array![1.4_f64, 1.0, 2.0];
3778 let event_target = array![1u8, 1u8, 0u8];
3779 let event_competing = array![0u8, 0u8, 0u8];
3780 let sampleweight = array![1.0_f64, 2.5, 0.7];
3781 let x_entry = array![
3782 [1.0, age_entry[0].ln()],
3783 [1.0, age_entry[1].max(1e-8).ln()],
3784 [1.0, age_entry[2].ln()]
3785 ];
3786 let x_exit = array![
3787 [1.0, age_exit[0].ln()],
3788 [1.0, age_exit[1].ln()],
3789 [1.0, age_exit[2].ln()]
3790 ];
3791 let x_derivative = array![
3792 [0.0, 1.0 / age_exit[0]],
3793 [0.0, 1.0 / age_exit[1]],
3794 [0.0, 1.0 / age_exit[2]]
3795 ];
3796 let o_entry = array![0.2_f64, 0.0, 0.1];
3799 let o_exit = array![0.4_f64, 0.5, 0.7];
3800 let o_deriv = array![0.3_f64, 0.8, 0.5];
3801 let penalties = PenaltyBlocks::new(Vec::new());
3802 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3803 let beta = array![-0.7_f64, 0.6];
3804
3805 let build = |o_e: &Array1<f64>, o_x: &Array1<f64>, o_d: &Array1<f64>| {
3806 survival_model_with_offsets(
3807 survival_inputs(
3808 &age_entry,
3809 &age_exit,
3810 &event_target,
3811 &event_competing,
3812 &sampleweight,
3813 &x_entry,
3814 &x_exit,
3815 &x_derivative,
3816 ),
3817 Some(SurvivalBaselineOffsets {
3818 eta_entry: o_e.view(),
3819 eta_exit: o_x.view(),
3820 derivative_exit: o_d.view(),
3821 }),
3822 penalties.clone(),
3823 mono,
3824 SurvivalSpec::Net,
3825 )
3826 .expect("model build")
3827 };
3828
3829 let base = build(&o_entry, &o_exit, &o_deriv);
3830 let resid = base
3831 .offset_channel_residuals(&beta)
3832 .expect("offset residuals");
3833 assert_eq!(resid.exit.len(), 3);
3834 assert_eq!(resid.entry.len(), 3);
3835 assert_eq!(resid.derivative.len(), 3);
3836
3837 let nll = |m: &WorkingModelSurvival| 0.5 * m.update_state(&beta).expect("state").deviance;
3840 let h = 1e-6;
3841
3842 assert_eq!(resid.entry[1], 0.0);
3846 assert_eq!(resid.derivative[2], 0.0);
3847
3848 for i in 0..3 {
3849 {
3851 let mut op = o_exit.clone();
3852 let mut om = o_exit.clone();
3853 op[i] += h;
3854 om[i] -= h;
3855 let fd = (nll(&build(&o_entry, &op, &o_deriv))
3856 - nll(&build(&o_entry, &om, &o_deriv)))
3857 / (2.0 * h);
3858 assert!(
3859 (resid.exit[i] - fd).abs() < 1e-6,
3860 "∂NLL/∂o_X[{i}]: analytic={:.6e} fd={:.6e}",
3861 resid.exit[i],
3862 fd
3863 );
3864 }
3865 {
3869 let mut op = o_entry.clone();
3870 let mut om = o_entry.clone();
3871 op[i] += h;
3872 om[i] -= h;
3873 let fd = (nll(&build(&op, &o_exit, &o_deriv))
3874 - nll(&build(&om, &o_exit, &o_deriv)))
3875 / (2.0 * h);
3876 assert!(
3877 (resid.entry[i] - fd).abs() < 1e-6,
3878 "∂NLL/∂o_E[{i}]: analytic={:.6e} fd={:.6e}",
3879 resid.entry[i],
3880 fd
3881 );
3882 }
3883 {
3885 let mut op = o_deriv.clone();
3886 let mut om = o_deriv.clone();
3887 op[i] += h;
3888 om[i] -= h;
3889 let fd = (nll(&build(&o_entry, &o_exit, &op))
3890 - nll(&build(&o_entry, &o_exit, &om)))
3891 / (2.0 * h);
3892 assert!(
3893 (resid.derivative[i] - fd).abs() < 1e-6,
3894 "∂NLL/∂o_D[{i}]: analytic={:.6e} fd={:.6e}",
3895 resid.derivative[i],
3896 fd
3897 );
3898 }
3899 }
3900 }
3901
3902 #[test]
3903 fn offset_channel_residuals_respect_zero_sampleweight() {
3904 let age_entry = array![1.0_f64, 2.0];
3905 let age_exit = array![2.0_f64, 3.5];
3906 let event_target = array![1u8, 1u8];
3907 let event_competing = array![0u8, 0u8];
3908 let sampleweight = array![0.0_f64, 1.2]; let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
3910 let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
3911 let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
3912 let penalties = PenaltyBlocks::new(Vec::new());
3913 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3914 let beta = array![-1.0_f64, 0.8];
3915
3916 let model = survival_model_with_offsets(
3917 survival_inputs(
3918 &age_entry,
3919 &age_exit,
3920 &event_target,
3921 &event_competing,
3922 &sampleweight,
3923 &x_entry,
3924 &x_exit,
3925 &x_derivative,
3926 ),
3927 Some(SurvivalBaselineOffsets {
3928 eta_entry: array![0.0_f64, 0.1].view(),
3929 eta_exit: array![0.0_f64, 0.2].view(),
3930 derivative_exit: array![0.0_f64, 0.1].view(),
3931 }),
3932 penalties,
3933 mono,
3934 SurvivalSpec::Net,
3935 )
3936 .expect("model");
3937 let r = model.offset_channel_residuals(&beta).expect("resid");
3938 assert_eq!(r.exit[0], 0.0);
3940 assert_eq!(r.entry[0], 0.0);
3941 assert_eq!(r.derivative[0], 0.0);
3942 assert!(r.exit[1] != 0.0);
3944 }
3945
3946 #[test]
3947 fn offset_channel_residuals_reject_beta_dim_mismatch() {
3948 let age_entry = array![1.0_f64];
3949 let age_exit = array![2.0_f64];
3950 let event_target = array![1u8];
3951 let event_competing = array![0u8];
3952 let sampleweight = array![1.0_f64];
3953 let x_entry = array![[1.0, 0.0]];
3954 let x_exit = array![[1.0, 0.7]];
3955 let x_derivative = array![[0.0, 0.5]];
3956 let penalties = PenaltyBlocks::new(Vec::new());
3957 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3958 let model = survival_model(
3959 survival_inputs(
3960 &age_entry,
3961 &age_exit,
3962 &event_target,
3963 &event_competing,
3964 &sampleweight,
3965 &x_entry,
3966 &x_exit,
3967 &x_derivative,
3968 ),
3969 penalties,
3970 mono,
3971 SurvivalSpec::Net,
3972 )
3973 .expect("model");
3974 let bad_beta = array![0.0_f64]; let err = model
3976 .offset_channel_residuals(&bad_beta)
3977 .expect_err("mismatch must error");
3978 match err {
3979 EstimationError::InvalidInput(msg) => {
3980 assert!(msg.contains("beta dimension mismatch"), "msg={msg}")
3981 }
3982 other => panic!("expected InvalidInput, got {other:?}"),
3983 }
3984 }
3985
3986 #[test]
3987 fn crudespec_is_rejected_by_one_hazard_engine() {
3988 let age_entry = array![1.0_f64];
3989 let age_exit = array![2.0_f64];
3990 let event_target = array![0u8];
3991 let event_competing = array![1u8];
3992 let sampleweight = array![1.0];
3993 let x_entry = array![[0.1]];
3994 let x_exit = array![[0.4]];
3995 let x_derivative = array![[1.0]];
3996 let penalties = PenaltyBlocks::new(Vec::new());
3997 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3998
3999 let err = survival_model(
4000 survival_inputs(
4001 &age_entry,
4002 &age_exit,
4003 &event_target,
4004 &event_competing,
4005 &sampleweight,
4006 &x_entry,
4007 &x_exit,
4008 &x_derivative,
4009 ),
4010 penalties,
4011 mono,
4012 SurvivalSpec::Crude,
4013 )
4014 .expect_err("crude fitting should be rejected by the one-hazard engine");
4015 assert!(matches!(err, SurvivalError::UnsupportedSpec("crude")));
4016 }
4017
4018 #[test]
4019 fn nonstructural_models_require_explicit_monotonicity_collocation() {
4020 let age_entry = array![1.0_f64, 1.5_f64];
4021 let age_exit = array![2.0_f64, 2.5_f64];
4022 let event_target = array![0u8, 0u8];
4023 let event_competing = array![0u8, 1u8];
4024 let sampleweight = array![1.0, 1.0];
4025 let x_entry = array![[0.2], [0.1]];
4026 let x_exit = array![[0.3], [0.2]];
4027 let x_derivative = array![[1.0], [1.0]];
4028
4029 let model = survival_model(
4030 survival_inputs(
4031 &age_entry,
4032 &age_exit,
4033 &event_target,
4034 &event_competing,
4035 &sampleweight,
4036 &x_entry,
4037 &x_exit,
4038 &x_derivative,
4039 ),
4040 PenaltyBlocks::new(Vec::new()),
4041 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4042 SurvivalSpec::Net,
4043 )
4044 .expect("construct censored survival model");
4045
4046 assert!(
4047 model.monotonicity_linear_constraints().is_none(),
4048 "non-structural survival models must not fabricate rowwise monotonicity constraints"
4049 );
4050 }
4051
4052 #[test]
4053 fn decreasing_interval_is_rejectedwithout_target_events() {
4054 let age_entry = array![1.0_f64];
4055 let age_exit = array![2.0_f64];
4056 let event_target = array![0u8];
4057 let event_competing = array![0u8];
4058 let sampleweight = array![1.0];
4059 let x_entry = array![[0.5]];
4060 let x_exit = array![[0.0]];
4061 let x_derivative = array![[1.0]];
4062
4063 let model = survival_model(
4064 survival_inputs(
4065 &age_entry,
4066 &age_exit,
4067 &event_target,
4068 &event_competing,
4069 &sampleweight,
4070 &x_entry,
4071 &x_exit,
4072 &x_derivative,
4073 ),
4074 PenaltyBlocks::new(Vec::new()),
4075 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4076 SurvivalSpec::Net,
4077 )
4078 .expect("construct censored survival model");
4079
4080 let err = model
4081 .update_state(&array![1.0])
4082 .expect_err("decreasing cumulative hazard increment should be rejected");
4083 assert!(
4084 err.to_string().contains("cumulative hazard decreased"),
4085 "unexpected error: {err}"
4086 );
4087 }
4088
4089 fn smooth_crude_risk(beta_d: f64, beta_m: f64) -> CrudeRiskResult {
4090 calculate_crude_risk_quadrature(
4091 0.0,
4092 1.0,
4093 &[0.0, 1.0],
4094 beta_d.exp(),
4095 beta_m.exp(),
4096 array![1.0].view(),
4097 array![1.0].view(),
4098 |u, design_d, deriv_d, design_m| {
4099 let cumulative_d = beta_d.exp() * (1.0 + 0.2 * u);
4100 let cumulative_m = beta_m.exp() * (1.0 + 0.1 * u);
4101 let inst_hazard_d = 0.2 * beta_d.exp();
4102 design_d[0] = 1.0;
4103 deriv_d[0] = 0.0;
4106 design_m[0] = 1.0;
4107 Ok((inst_hazard_d, cumulative_d, cumulative_m))
4108 },
4109 )
4110 .expect("smooth crude-risk quadrature should succeed")
4111 }
4112
4113 #[test]
4114 fn crude_riskgradient_matches_monotoneobjective() {
4115 let beta_d = -0.2_f64;
4116 let beta_m = -0.5_f64;
4117 let result = smooth_crude_risk(beta_d, beta_m);
4118 let eps = 1e-6;
4119
4120 let fd_d = (smooth_crude_risk(beta_d + eps, beta_m).risk
4121 - smooth_crude_risk(beta_d - eps, beta_m).risk)
4122 / (2.0 * eps);
4123 let fd_m = (smooth_crude_risk(beta_d, beta_m + eps).risk
4124 - smooth_crude_risk(beta_d, beta_m - eps).risk)
4125 / (2.0 * eps);
4126
4127 assert!(
4128 (result.diseasegradient[0] - fd_d).abs() < 1e-5,
4129 "disease gradient mismatch for monotone crude risk: analytic={} fd={fd_d}",
4130 result.diseasegradient[0]
4131 );
4132 assert!(
4133 (result.mortalitygradient[0] - fd_m).abs() < 1e-5,
4134 "mortality gradient mismatch for monotone crude risk: analytic={} fd={fd_m}",
4135 result.mortalitygradient[0]
4136 );
4137 }
4138
4139 #[test]
4140 fn survival_working_state_is_ridge_free() {
4141 let age_entry = array![1.0_f64, 2.0_f64];
4142 let age_exit = array![2.0_f64, 3.5_f64];
4143 let event_target = array![1u8, 0u8];
4144 let event_competing = array![0u8, 0u8];
4145 let sampleweight = array![1.0, 1.0];
4146 let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
4147 let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
4148 let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
4149 let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4150 matrix: array![[2.0]],
4151 lambda: 1.7,
4152 range: 1..2,
4153 nullspace_dim: 0,
4154 }]);
4155 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4156 let beta = array![-1.2, 0.4];
4157
4158 let model = survival_model(
4159 survival_inputs(
4160 &age_entry,
4161 &age_exit,
4162 &event_target,
4163 &event_competing,
4164 &sampleweight,
4165 &x_entry,
4166 &x_exit,
4167 &x_derivative,
4168 ),
4169 penalties.clone(),
4170 mono,
4171 SurvivalSpec::Net,
4172 )
4173 .expect("construct survival model");
4174
4175 let state = model.update_state(&beta).expect("survival state");
4176 assert_eq!(
4177 state.ridge_used, 0.0,
4178 "survival objective must not fuse a coefficient ridge"
4179 );
4180 let expected_penalty = penalties.deviance(&beta);
4181 assert!(
4182 (state.penalty_term - expected_penalty).abs() < 1e-12,
4183 "penalty_term mismatch: state={} expected={}",
4184 state.penalty_term,
4185 expected_penalty
4186 );
4187 }
4188
4189 #[test]
4190 fn negative_penalty_lambda_is_rejected() {
4191 let age_entry = array![1.0_f64];
4192 let age_exit = array![2.0_f64];
4193 let event_target = array![1u8];
4194 let event_competing = array![0u8];
4195 let sampleweight = array![1.0];
4196 let x_entry = array![[1.0, 0.0]];
4197 let x_exit = array![[1.0, 0.5]];
4198 let x_derivative = array![[0.0, 1.0]];
4199 let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4200 matrix: array![[1.0]],
4201 lambda: -0.1,
4202 range: 1..2,
4203 nullspace_dim: 0,
4204 }]);
4205
4206 let err = survival_model(
4207 survival_inputs(
4208 &age_entry,
4209 &age_exit,
4210 &event_target,
4211 &event_competing,
4212 &sampleweight,
4213 &x_entry,
4214 &x_exit,
4215 &x_derivative,
4216 ),
4217 penalties,
4218 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4219 SurvivalSpec::Net,
4220 )
4221 .expect_err("negative lambda must be rejected");
4222
4223 assert!(matches!(err, SurvivalError::NonFiniteInput));
4224 }
4225
4226 #[test]
4227 fn penalty_block_range_and_shapemust_match_coefficients() {
4228 let age_entry = array![1.0_f64];
4229 let age_exit = array![2.0_f64];
4230 let event_target = array![1u8];
4231 let event_competing = array![0u8];
4232 let sampleweight = array![1.0];
4233 let x_entry = array![[1.0, 0.0]];
4234 let x_exit = array![[1.0, 0.5]];
4235 let x_derivative = array![[0.0, 1.0]];
4236 let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4237 matrix: array![[1.0]],
4238 lambda: 0.5,
4239 range: 0..2,
4240 nullspace_dim: 0,
4241 }]);
4242
4243 let err = survival_model(
4244 survival_inputs(
4245 &age_entry,
4246 &age_exit,
4247 &event_target,
4248 &event_competing,
4249 &sampleweight,
4250 &x_entry,
4251 &x_exit,
4252 &x_derivative,
4253 ),
4254 penalties,
4255 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
4256 SurvivalSpec::Net,
4257 )
4258 .expect_err("penalty block geometry must match coefficient support");
4259
4260 assert!(matches!(err, SurvivalError::DimensionMismatch));
4261 }
4262
4263 #[test]
4264 fn survivalgradient_matches_ridge_free_objective_fd() {
4265 let age_entry = array![1.0_f64, 2.0_f64, 3.0_f64];
4266 let age_exit = array![2.0_f64, 3.5_f64, 4.0_f64];
4267 let event_target = array![1u8, 0u8, 1u8];
4268 let event_competing = array![0u8, 0u8, 0u8];
4269 let sampleweight = array![1.0, 1.0, 1.0];
4270 let x_entry = array![
4271 [1.0, age_entry[0].ln()],
4272 [1.0, age_entry[1].ln()],
4273 [1.0, age_entry[2].ln()]
4274 ];
4275 let x_exit = array![
4276 [1.0, age_exit[0].ln()],
4277 [1.0, age_exit[1].ln()],
4278 [1.0, age_exit[2].ln()]
4279 ];
4280 let x_derivative = array![
4281 [0.0, 1.0 / age_exit[0]],
4282 [0.0, 1.0 / age_exit[1]],
4283 [0.0, 1.0 / age_exit[2]]
4284 ];
4285 let penalties = PenaltyBlocks::new(Vec::new());
4286 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4287 let beta = array![-1.0, 3.0];
4288
4289 let model = survival_model(
4290 survival_inputs(
4291 &age_entry,
4292 &age_exit,
4293 &event_target,
4294 &event_competing,
4295 &sampleweight,
4296 &x_entry,
4297 &x_exit,
4298 &x_derivative,
4299 ),
4300 penalties,
4301 mono,
4302 SurvivalSpec::Net,
4303 )
4304 .expect("construct survival model");
4305
4306 let state = model.update_state(&beta).expect("state at beta");
4307 let eps = 1e-7;
4308 for j in 0..beta.len() {
4309 let mut plus = beta.clone();
4310 let mut minus = beta.clone();
4311 plus[j] += eps;
4312 minus[j] -= eps;
4313 let state_plus = model.update_state(&plus).expect("state at beta + eps");
4314 let state_minus = model.update_state(&minus).expect("state at beta - eps");
4315 let obj_plus = 0.5 * state_plus.deviance + state_plus.penalty_term;
4316 let obj_minus = 0.5 * state_minus.deviance + state_minus.penalty_term;
4317 let fd = (obj_plus - obj_minus) / (2.0 * eps);
4318 assert_eq!(
4319 state.gradient[j].signum(),
4320 fd.signum(),
4321 "objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
4322 state.gradient[j]
4323 );
4324 assert!(
4325 (state.gradient[j] - fd).abs() < 1e-5,
4326 "objective/gradient mismatch at j={j}: grad={} fd={fd}",
4327 state.gradient[j]
4328 );
4329 }
4330 }
4331
4332 fn laml_fd_test_model(lambda: f64) -> WorkingModelSurvival {
4333 let age_entry: Array1<f64> = Array1::from(vec![
4340 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,
4341 34.0, 39.0, 44.0, 49.0, 54.0, 59.0,
4342 ]);
4343 let age_exit: Array1<f64> = Array1::from(vec![
4344 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,
4345 48.0, 51.0, 58.0, 62.0, 66.0, 69.0,
4346 ]);
4347 let event_target = Array1::from(vec![
4348 1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
4349 ]);
4350 let event_competing = Array1::<u8>::zeros(age_entry.len());
4351 let sampleweight = Array1::from_elem(age_entry.len(), 1.0_f64);
4352 let n = age_entry.len();
4353 let ln_age_mean: f64 = {
4354 let mut sum = 0.0;
4355 for i in 0..n {
4356 sum += age_entry[i].ln() + age_exit[i].ln();
4357 }
4358 sum / (2.0 * n as f64)
4359 };
4360 let mut x_entry = Array2::<f64>::zeros((n, 2));
4361 let mut x_exit = Array2::<f64>::zeros((n, 2));
4362 let mut x_derivative = Array2::<f64>::zeros((n, 2));
4363 for i in 0..n {
4364 x_entry[[i, 0]] = 1.0;
4365 x_exit[[i, 0]] = 1.0;
4366 x_entry[[i, 1]] = age_entry[i].ln() - ln_age_mean;
4367 x_exit[[i, 1]] = age_exit[i].ln() - ln_age_mean;
4368 x_derivative[[i, 0]] = 0.0;
4369 x_derivative[[i, 1]] = 1.0 / age_exit[i];
4370 }
4371 let penalties = PenaltyBlocks::new(vec![
4372 PenaltyBlock {
4373 matrix: array![[3.0]],
4374 lambda: 0.0,
4375 range: 0..1,
4376 nullspace_dim: 0,
4377 },
4378 PenaltyBlock {
4379 matrix: array![[2.5]],
4380 lambda,
4381 range: 1..2,
4382 nullspace_dim: 0,
4383 },
4384 ]);
4385 survival_model(
4386 survival_inputs(
4387 &age_entry,
4388 &age_exit,
4389 &event_target,
4390 &event_competing,
4391 &sampleweight,
4392 &x_entry,
4393 &x_exit,
4394 &x_derivative,
4395 ),
4396 penalties,
4397 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
4398 SurvivalSpec::Net,
4399 )
4400 .expect("construct LAML FD survival model")
4401 }
4402
4403 fn laml_test_logdet_h(state: &WorkingState) -> f64 {
4404 use gam_linalg::faer_ndarray::FaerEigh;
4405 use gam_solve::estimate::reml::reml_outer_engine::{spectral_epsilon, spectral_regularize};
4406
4407 let h_dense = state.hessian.to_dense();
4408 let (evals, _) = h_dense.eigh(faer::Side::Lower).expect("eigh");
4409 let eps = spectral_epsilon(evals.as_slice().unwrap());
4410 evals
4411 .iter()
4412 .map(|&sigma| spectral_regularize(sigma, eps).ln())
4413 .sum()
4414 }
4415
4416 #[test]
4417 fn survival_solver_damping_converges_undamped_objective() {
4418 let rho = -0.35_f64;
4419 let model = laml_fd_test_model(rho.exp());
4420 let beta0 = array![-2.5_f64, 1.0];
4421 let (converged_model, beta) = model
4422 .reconverge_survival_inner_mode(&[rho], &beta0)
4423 .expect("converge survival mode with solver-only damping");
4424 let state = converged_model
4425 .update_state(&beta)
4426 .expect("evaluate undamped objective at converged mode");
4427
4428 assert_eq!(
4429 state.ridge_used, 0.0,
4430 "solver damping must not enter the converged statistical objective"
4431 );
4432 let undamped_stationarity = array1_l2_norm(&state.gradient);
4433 assert!(
4434 undamped_stationarity <= 1.0e-9,
4435 "solver must converge the undamped objective; ||gradient||={undamped_stationarity:.3e}"
4436 );
4437 }
4438
4439 #[test]
4440 fn laml_gradient_and_objective_ignore_inactive_penalty_prefix_blocks() {
4441 let rho0 = -0.35_f64;
4455 let beta = array![-2.5_f64, 1.0];
4456 let model = laml_fd_test_model(rho0.exp());
4457 let state = model
4458 .update_state(&beta)
4459 .expect("state for LAML prefix-skip test");
4460
4461 assert_eq!(model.penalties.blocks.len(), 2);
4466 assert_eq!(model.penalties.blocks[0].lambda, 0.0);
4467 assert!(model.penalties.blocks[1].lambda > 0.0);
4468
4469 let rho = Array1::from_iter(
4470 model
4471 .penalties
4472 .blocks
4473 .iter()
4474 .filter(|b| b.lambda > 0.0)
4475 .map(|b| b.lambda.ln()),
4476 );
4477 assert_eq!(
4478 rho.len(),
4479 1,
4480 "fixture should expose exactly one active penalty block for the rho vector"
4481 );
4482
4483 let (obj, grad) = model
4484 .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
4485 .expect("survival LAML objective and gradient");
4486
4487 let expected = 0.5 * state.deviance + state.penalty_term + 0.5 * laml_test_logdet_h(&state)
4488 - 0.5 * (rho0 + 2.5_f64.ln());
4489 assert_eq!(
4490 grad.len(),
4491 1,
4492 "rho-gradient must match the active-penalty count, not the full block list"
4493 );
4494 assert!(
4495 (obj - expected).abs() < 1e-10,
4496 "survival LAML objective mismatch with inactive prefix block: obj={obj} expected={expected}",
4497 );
4498 assert!(
4499 grad[0].is_finite(),
4500 "rho-gradient must be finite: {}",
4501 grad[0]
4502 );
4503 }
4504
4505 #[test]
4506 fn structural_monotonicgradient_matchesobjectivefd() {
4507 let age_entry = array![1.0_f64, 1.3_f64, 1.8_f64];
4508 let age_exit = array![1.6_f64, 2.1_f64, 2.7_f64];
4509 let event_target = array![1u8, 0u8, 1u8];
4510 let event_competing = array![0u8, 0u8, 0u8];
4511 let sampleweight = array![1.0, 1.0, 1.0];
4512
4513 let x_entry = array![
4516 [1.0, 0.2, 0.05, -0.7],
4517 [1.0, 0.5, 0.20, 0.1],
4518 [1.0, 0.9, 0.60, 1.2]
4519 ];
4520 let x_exit = array![
4521 [1.0, 0.4, 0.16, -0.7],
4522 [1.0, 0.8, 0.64, 0.1],
4523 [1.0, 1.1, 1.21, 1.2]
4524 ];
4525 let x_derivative = array![
4526 [0.0, 0.8, 0.64, 0.0],
4527 [0.0, 0.7, 1.12, 0.0],
4528 [0.0, 0.6, 1.32, 0.0]
4529 ];
4530 let penalties = PenaltyBlocks::new(Vec::new());
4531 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4532 let mut model = survival_model(
4533 survival_inputs(
4534 &age_entry,
4535 &age_exit,
4536 &event_target,
4537 &event_competing,
4538 &sampleweight,
4539 &x_entry,
4540 &x_exit,
4541 &x_derivative,
4542 ),
4543 penalties,
4544 mono,
4545 SurvivalSpec::Net,
4546 )
4547 .expect("construct structural survival model");
4548 model
4549 .set_structural_monotonicity(true, 3)
4550 .expect("enable structural monotonicity");
4551 let constraints = model
4552 .monotonicity_linear_constraints()
4553 .expect("structural derivative constraints");
4554 assert_eq!(constraints.a.nrows(), 2);
4555 assert_eq!(constraints.a.ncols(), 4);
4556 assert_eq!(constraints.a.row(0).to_vec(), vec![0.0, 1.0, 0.0, 0.0]);
4557 assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 0.0, 1.0, 0.0]);
4558 assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
4559
4560 let beta = array![0.2, 0.2, 0.1, 0.2];
4561 let state = model.update_state(&beta).expect("state at structural beta");
4562 let eps = 1e-7;
4563 for j in 0..beta.len() {
4564 let mut plus = beta.clone();
4565 let mut minus = beta.clone();
4566 plus[j] += eps;
4567 minus[j] -= eps;
4568 let state_plus = model.update_state(&plus).expect("state at beta + eps");
4569 let state_minus = model.update_state(&minus).expect("state at beta - eps");
4570 let obj_plus = 0.5 * state_plus.deviance + state_plus.penalty_term;
4571 let obj_minus = 0.5 * state_minus.deviance + state_minus.penalty_term;
4572 let fd = (obj_plus - obj_minus) / (2.0 * eps);
4573 assert_eq!(
4574 state.gradient[j].signum(),
4575 fd.signum(),
4576 "structural objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
4577 state.gradient[j]
4578 );
4579 assert!(
4580 (state.gradient[j] - fd).abs() < 2e-5,
4581 "structural objective/gradient mismatch at j={j}: grad={} fd={fd}",
4582 state.gradient[j]
4583 );
4584 }
4585 }
4586
4587 #[test]
4588 fn structural_monotonic_lamlgradient_returns_finitevalues() {
4589 let age_entry = array![1.0_f64, 1.2_f64];
4590 let age_exit = array![1.5_f64, 2.0_f64];
4591 let event_target = array![1u8, 0u8];
4592 let event_competing = array![0u8, 0u8];
4593 let sampleweight = array![1.0, 1.0];
4594 let x_entry = array![[1.0, 0.2, -0.5], [1.0, 0.4, 0.2]];
4595 let x_exit = array![[1.0, 0.5, -0.5], [1.0, 0.8, 0.2]];
4596 let x_derivative = array![[0.0, 0.9, 0.0], [0.0, 0.7, 0.0]];
4597 let penalties = PenaltyBlocks::new(Vec::new());
4598 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4599 let mut model = survival_model(
4600 survival_inputs(
4601 &age_entry,
4602 &age_exit,
4603 &event_target,
4604 &event_competing,
4605 &sampleweight,
4606 &x_entry,
4607 &x_exit,
4608 &x_derivative,
4609 ),
4610 penalties,
4611 mono,
4612 SurvivalSpec::Net,
4613 )
4614 .expect("construct structural survival model");
4615 model
4616 .set_structural_monotonicity(true, 2)
4617 .expect("enable structural monotonicity");
4618 model.penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4620 matrix: array![[1.0]],
4621 lambda: 0.7,
4622 range: 1..2,
4623 nullspace_dim: 0,
4624 }]);
4625 let beta = array![0.2, 0.2, 0.1];
4626 let state = model.update_state(&beta).expect("state at structural beta");
4627 let rho = Array1::from_iter(
4628 model
4629 .penalties
4630 .blocks
4631 .iter()
4632 .filter(|b| b.lambda > 0.0)
4633 .map(|b| b.lambda.ln()),
4634 );
4635 let (obj, grad) = model
4636 .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
4637 .expect("laml gradient should work in structural mode");
4638 assert!(obj.is_finite());
4639 assert_eq!(grad.len(), 1);
4640 assert!(grad[0].is_finite());
4641 }
4642
4643 #[test]
4644 fn structural_monotonicity_switches_to_tiny_derivative_guard_constraints() {
4645 let age_entry = array![1.0_f64];
4646 let age_exit = array![2.0_f64];
4647 let event_target = array![1u8];
4648 let event_competing = array![0u8];
4649 let sampleweight = array![1.0];
4650 let x_entry = array![[0.0]];
4651 let x_exit = array![[0.2]];
4652 let x_derivative = array![[1.0]];
4653
4654 let penalties = PenaltyBlocks::new(Vec::new());
4655 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4656 let mut model = survival_model(
4657 survival_inputs(
4658 &age_entry,
4659 &age_exit,
4660 &event_target,
4661 &event_competing,
4662 &sampleweight,
4663 &x_entry,
4664 &x_exit,
4665 &x_derivative,
4666 ),
4667 penalties,
4668 mono,
4669 SurvivalSpec::Net,
4670 )
4671 .expect("construct structural survival model");
4672
4673 let beta = array![-3.0];
4674 assert!(
4675 model.update_state(&beta).is_err(),
4676 "negative derivative coefficient should violate derivative guard"
4677 );
4678
4679 model
4680 .set_structural_monotonicity(true, 1)
4681 .expect("enable structural monotonicity");
4682 let constraints = model
4683 .monotonicity_linear_constraints()
4684 .expect("structural derivative constraints");
4685 assert_eq!(constraints.a.nrows(), 1);
4686 assert_eq!(constraints.a.ncols(), 1);
4687 assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
4688 assert!(constraints.b[0].abs() <= 1e-12);
4690 let state = model
4691 .update_state(&array![1e-6])
4692 .expect("small positive derivative coefficient should remain feasible");
4693 assert!(state.deviance.is_finite());
4694 }
4695
4696 #[test]
4697 fn derivative_offset_must_clear_nonstructural_monotonicity_threshold() {
4698 let age_entry = array![1.0_f64];
4699 let age_exit = array![2.0_f64];
4700 let event_target = array![1u8];
4701 let event_competing = array![0u8];
4702 let sampleweight = array![1.0];
4703 let x_entry = array![[1.0, 0.0]];
4704 let x_exit = array![[1.0, 0.0]];
4705 let x_derivative = array![[0.0, 0.0]];
4706 let penalties = PenaltyBlocks::new(Vec::new());
4707 let monotonicity = SurvivalMonotonicityPenalty { tolerance: 3.0 };
4708 let eta_entry_offset = array![0.0];
4709 let eta_exit_offset = array![0.0];
4710 let derivative_offset_below_guard = array![2.0];
4711 let derivative_offset_above_guard = array![3.1];
4712 let offsets_below_guard = SurvivalBaselineOffsets {
4713 eta_entry: eta_entry_offset.view(),
4714 eta_exit: eta_exit_offset.view(),
4715 derivative_exit: derivative_offset_below_guard.view(),
4716 };
4717 let offsets_above_guard = SurvivalBaselineOffsets {
4718 eta_entry: eta_entry_offset.view(),
4719 eta_exit: eta_exit_offset.view(),
4720 derivative_exit: derivative_offset_above_guard.view(),
4721 };
4722
4723 let model_below_guard = survival_model_with_offsets(
4724 survival_inputs(
4725 &age_entry,
4726 &age_exit,
4727 &event_target,
4728 &event_competing,
4729 &sampleweight,
4730 &x_entry,
4731 &x_exit,
4732 &x_derivative,
4733 ),
4734 Some(offsets_below_guard),
4735 penalties.clone(),
4736 monotonicity,
4737 SurvivalSpec::Net,
4738 )
4739 .expect("construct model with derivative offset below guard");
4740 let err = model_below_guard
4741 .update_state(&array![0.0, 0.0])
4742 .expect_err("derivative offset below guard should be rejected");
4743 let err_text = err.to_string();
4744 assert!(
4745 err_text.contains("d_eta/dt=2.000e0") && err_text.contains("tolerance=3.000e0"),
4746 "expected derivative guard rejection to report the offset-driven derivative: {err_text}"
4747 );
4748
4749 let model_above_guard = survival_model_with_offsets(
4750 survival_inputs(
4751 &age_entry,
4752 &age_exit,
4753 &event_target,
4754 &event_competing,
4755 &sampleweight,
4756 &x_entry,
4757 &x_exit,
4758 &x_derivative,
4759 ),
4760 Some(offsets_above_guard),
4761 penalties,
4762 SurvivalMonotonicityPenalty { tolerance: 3.0 },
4763 SurvivalSpec::Net,
4764 )
4765 .expect("construct model with derivative offset above guard");
4766 let state = model_above_guard
4767 .update_state(&array![0.0, 0.0])
4768 .expect("derivative offset above guard should remain feasible");
4769 assert!(state.deviance.is_finite());
4770 }
4771
4772 #[test]
4773 fn structural_monotonicity_rejects_negative_derivative_offsets() {
4774 let age_entry = array![1.0_f64];
4775 let age_exit = array![2.0_f64];
4776 let event_target = array![1u8];
4777 let event_competing = array![0u8];
4778 let sampleweight = array![1.0];
4779 let x_entry = array![[0.0]];
4780 let x_exit = array![[0.2]];
4781 let x_derivative = array![[1.0]];
4782 let eta_entry = array![0.0];
4783 let eta_exit = array![0.0];
4784 let derivative_exit = array![-1e-3];
4785 let offsets = SurvivalBaselineOffsets {
4786 eta_entry: eta_entry.view(),
4787 eta_exit: eta_exit.view(),
4788 derivative_exit: derivative_exit.view(),
4789 };
4790
4791 let mut model = survival_model_with_offsets(
4792 survival_inputs(
4793 &age_entry,
4794 &age_exit,
4795 &event_target,
4796 &event_competing,
4797 &sampleweight,
4798 &x_entry,
4799 &x_exit,
4800 &x_derivative,
4801 ),
4802 Some(offsets),
4803 PenaltyBlocks::new(Vec::new()),
4804 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4805 SurvivalSpec::Net,
4806 )
4807 .expect("construct structural survival model");
4808 let err = model
4809 .set_structural_monotonicity(true, 1)
4810 .expect_err("negative derivative offsets must be rejected");
4811 assert!(
4812 err.to_string()
4813 .contains("structural monotonicity requires nonnegative derivative offsets"),
4814 "unexpected error: {err}"
4815 );
4816 }
4817
4818 #[test]
4819 fn structural_monotonicity_emits_coefficient_constraints() {
4820 let age_entry = array![1.0_f64, 1.5_f64];
4821 let age_exit = array![2.0_f64, 3.0_f64];
4822 let event_target = array![1u8, 0u8];
4823 let event_competing = array![0u8, 0u8];
4824 let sampleweight = array![1.0, 1.0];
4825 let x_entry = array![[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]];
4826 let x_exit = array![[0.2, 0.4, 1.0], [0.3, 0.5, 1.0]];
4827 let x_derivative = array![[0.3, 0.2, 0.0], [0.4, 0.1, 0.0]];
4828
4829 let mut model = survival_model(
4830 survival_inputs(
4831 &age_entry,
4832 &age_exit,
4833 &event_target,
4834 &event_competing,
4835 &sampleweight,
4836 &x_entry,
4837 &x_exit,
4838 &x_derivative,
4839 ),
4840 PenaltyBlocks::new(Vec::new()),
4841 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4842 SurvivalSpec::Net,
4843 )
4844 .expect("construct structural survival model");
4845 model
4846 .set_structural_monotonicity(true, 2)
4847 .expect("enable structural monotonicity");
4848
4849 let constraints = model
4850 .monotonicity_linear_constraints()
4851 .expect("structural derivative constraints");
4852
4853 assert_eq!(constraints.a.nrows(), 2);
4854 assert_eq!(constraints.a.ncols(), 3);
4855 assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0, 0.0]);
4856 assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0, 0.0]);
4857 assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
4858 }
4859
4860 #[test]
4861 fn structural_monotonicity_preserves_inactive_time_columns_in_constraints() {
4862 let age_entry = array![1.0_f64];
4863 let age_exit = array![2.0_f64];
4864 let event_target = array![1u8];
4865 let event_competing = array![0u8];
4866 let sampleweight = array![1.0];
4867 let x_entry = array![[1.0, 0.2]];
4868 let x_exit = array![[1.0, 0.6]];
4869 let x_derivative = array![[0.0, 1.0]];
4870
4871 let mut model = survival_model(
4872 survival_inputs(
4873 &age_entry,
4874 &age_exit,
4875 &event_target,
4876 &event_competing,
4877 &sampleweight,
4878 &x_entry,
4879 &x_exit,
4880 &x_derivative,
4881 ),
4882 PenaltyBlocks::new(Vec::new()),
4883 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4884 SurvivalSpec::Net,
4885 )
4886 .expect("construct structural survival model");
4887 model
4888 .set_structural_monotonicity(true, 2)
4889 .expect("enable structural monotonicity");
4890
4891 let constraints = model
4892 .monotonicity_linear_constraints()
4893 .expect("structural derivative constraints");
4894
4895 assert_eq!(constraints.a.nrows(), 1);
4896 assert!(
4897 constraints.a[[0, 0]].abs() <= 1e-12,
4898 "inactive time column should remain unconstrained"
4899 );
4900 assert!(
4901 (constraints.a[[0, 1]] - 1.0).abs() <= 1e-12,
4902 "active time column should remain constrained"
4903 );
4904 }
4905
4906 #[test]
4907 fn structural_monotonicity_preserves_sparse_row_patterns() {
4908 let age_entry = array![1.0_f64, 1.5_f64];
4909 let age_exit = array![2.0_f64, 2.5_f64];
4910 let event_target = array![1u8, 1u8];
4911 let event_competing = array![0u8, 0u8];
4912 let sampleweight = array![1.0, 1.0];
4913 let x_entry = array![[0.0, 0.0], [0.0, 0.0]];
4914 let x_exit = array![[0.4, 0.2], [0.6, 0.3]];
4915 let x_derivative = array![[1.0, 0.0], [1.0, 0.5]];
4916
4917 let mut model = survival_model(
4918 survival_inputs(
4919 &age_entry,
4920 &age_exit,
4921 &event_target,
4922 &event_competing,
4923 &sampleweight,
4924 &x_entry,
4925 &x_exit,
4926 &x_derivative,
4927 ),
4928 PenaltyBlocks::new(Vec::new()),
4929 SurvivalMonotonicityPenalty { tolerance: 0.0 },
4930 SurvivalSpec::Net,
4931 )
4932 .expect("construct structural survival model");
4933 model
4934 .set_structural_monotonicity(true, 2)
4935 .expect("enable structural monotonicity");
4936
4937 let constraints = model
4938 .monotonicity_linear_constraints()
4939 .expect("structural derivative constraints");
4940
4941 assert_eq!(constraints.a.nrows(), 2);
4942 assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0]);
4943 assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0]);
4944 }
4945
4946 #[test]
4947 fn update_state_rejects_negative_exit_derivative_for_censoredrows() {
4948 let age_entry = array![1.0_f64];
4949 let age_exit = array![1.1_f64];
4950 let event_target = array![0u8];
4951 let event_competing = array![0u8];
4952 let sampleweight = array![1.0];
4953 let x_entry = array![[0.0]];
4954 let x_exit = array![[0.0]];
4955 let x_derivative = array![[-1.0]];
4956 let penalties = PenaltyBlocks::new(Vec::new());
4957 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4958 let model = survival_model(
4959 survival_inputs(
4960 &age_entry,
4961 &age_exit,
4962 &event_target,
4963 &event_competing,
4964 &sampleweight,
4965 &x_entry,
4966 &x_exit,
4967 &x_derivative,
4968 ),
4969 penalties,
4970 mono,
4971 SurvivalSpec::Net,
4972 )
4973 .expect("construct censored survival model");
4974
4975 let err = model
4976 .update_state(&array![1.0])
4977 .expect_err("censored row should still enforce monotonic derivative");
4978 assert!(
4979 matches!(err, EstimationError::ParameterConstraintViolation(_)),
4980 "unexpected error: {err:?}"
4981 );
4982 }
4983
4984 fn crude_risk_quadrature_error(
4985 cumulative_entry: f64,
4986 cumulative_exit: f64,
4987 hazard_exit: f64,
4988 ) -> SurvivalError {
4989 calculate_crude_risk_quadrature(
4990 1.0,
4991 2.0,
4992 &[],
4993 0.4,
4994 0.2,
4995 array![1.0].view(),
4996 array![1.0].view(),
4997 |_, design_d, deriv_d, design_m| {
4998 design_d[0] = 1.0;
4999 deriv_d[0] = 0.0;
5000 design_m[0] = 1.0;
5001 Ok((cumulative_entry, cumulative_exit, hazard_exit))
5002 },
5003 )
5004 .expect_err("invalid hazards should fail")
5005 }
5006
5007 #[test]
5008 fn crude_risk_quadrature_rejects_decreasing_cumulative_hazard() {
5009 let err = crude_risk_quadrature_error(0.1, 0.3, 0.25);
5010 assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
5011 }
5012
5013 #[test]
5014 fn crude_risk_quadrature_rejects_nonpositive_instantaneous_hazard() {
5015 let err = crude_risk_quadrature_error(0.0, 0.4, 0.25);
5016 assert!(matches!(err, SurvivalError::NonPositiveHazard));
5017 }
5018
5019 #[test]
5020 fn laml_no_penalties_matches_documentedobjective() {
5021 let age_entry = array![40.0, 45.0, 50.0, 55.0];
5022 let age_exit = array![44.0, 49.0, 54.0, 59.0];
5023 let event_target = array![1u8, 0u8, 1u8, 0u8];
5024 let event_competing = Array1::<u8>::zeros(4);
5025 let sampleweight = Array1::ones(4);
5026 let x_entry = array![
5027 [1.0, -0.2, 0.04],
5028 [1.0, -0.1, 0.01],
5029 [1.0, 0.0, 0.0],
5030 [1.0, 0.1, 0.01]
5031 ];
5032 let x_exit = array![
5033 [1.0, -0.12, 0.0144],
5034 [1.0, -0.02, 0.0004],
5035 [1.0, 0.08, 0.0064],
5036 [1.0, 0.18, 0.0324]
5037 ];
5038 let x_derivative = array![
5039 [0.0, 0.02, 0.001],
5040 [0.0, 0.02, 0.001],
5041 [0.0, 0.02, 0.001],
5042 [0.0, 0.02, 0.001]
5043 ];
5044 let penalties = PenaltyBlocks::new(Vec::new());
5045 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5046 let beta = array![-2.0, 0.7, 0.2];
5047
5048 let model = survival_model(
5049 survival_inputs(
5050 &age_entry,
5051 &age_exit,
5052 &event_target,
5053 &event_competing,
5054 &sampleweight,
5055 &x_entry,
5056 &x_exit,
5057 &x_derivative,
5058 ),
5059 penalties,
5060 mono,
5061 SurvivalSpec::Net,
5062 )
5063 .expect("construct survival model");
5064
5065 let state = model.update_state(&beta).expect("state at beta");
5066 let rho = Array1::from_iter(
5067 model
5068 .penalties
5069 .blocks
5070 .iter()
5071 .filter(|b| b.lambda > 0.0)
5072 .map(|b| b.lambda.ln()),
5073 );
5074 let (obj, grad) = model
5075 .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
5076 .expect("laml objective for no-penalty model");
5077
5078 let h_dense = state.hessian.to_dense();
5079 let logdet_h: f64 = {
5095 use gam_problem::PseudoLogdetMode;
5096 use gam_solve::estimate::reml::reml_outer_engine::{
5097 DenseSpectralOperator, HessianFactorization,
5098 };
5099 let has_left_truncation = age_entry.iter().any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD);
5100 let mode = if has_left_truncation {
5101 PseudoLogdetMode::HardPseudo
5102 } else {
5103 PseudoLogdetMode::Smooth
5104 };
5105 DenseSpectralOperator::from_symmetric_with_mode(&h_dense, mode)
5106 .expect("survival LAML Hessian operator")
5107 .logdet()
5108 };
5109 let expected = 0.5 * state.deviance + state.penalty_term + 0.5 * logdet_h;
5110
5111 assert_eq!(grad.len(), 0);
5112 assert!(
5113 (obj - expected).abs() < 1e-10,
5114 "no-penalty LAML objective mismatch: obj={} expected={}",
5115 obj,
5116 expected
5117 );
5118 }
5119
5120 #[test]
5121 fn monotonicity_constraints_collapse_positive_collinearrows() {
5122 let a = array![[0.0, 0.5, 0.0], [0.0, 0.25, 0.0], [0.0, 0.125, 0.0]];
5123 let b = array![1e-8, 1e-8, 1e-8];
5124
5125 let compressed = compress_positive_collinear_constraints(&a, &b);
5126
5127 assert_eq!(compressed.a.nrows(), 1);
5128 assert_eq!(compressed.a.ncols(), 3);
5129 assert!(compressed.a[[0, 0]].abs() <= 1e-12);
5130 assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
5131 assert!(compressed.a[[0, 2]].abs() <= 1e-12);
5132 assert!((compressed.b[0] - 8e-8).abs() <= 1e-18);
5133 }
5134
5135 #[test]
5136 fn monotonicity_constraints_preserve_distinct_directions() {
5137 let a = array![[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]];
5138 let b = array![0.2, 0.3, 0.1];
5139
5140 let compressed = compress_positive_collinear_constraints(&a, &b);
5141
5142 assert_eq!(compressed.a.nrows(), 2);
5143 let mut saw_x = false;
5144 let mut saw_y = false;
5145 for i in 0..compressed.a.nrows() {
5146 if (compressed.a[[i, 0]] - 1.0).abs() <= 1e-12 && compressed.a[[i, 1]].abs() <= 1e-12 {
5147 saw_x = true;
5148 assert!((compressed.b[i] - 0.2).abs() <= 1e-12);
5149 }
5150 if compressed.a[[i, 0]].abs() <= 1e-12 && (compressed.a[[i, 1]] - 1.0).abs() <= 1e-12 {
5151 saw_y = true;
5152 assert!((compressed.b[i] - 0.3).abs() <= 1e-12);
5153 }
5154 }
5155 assert!(saw_x);
5156 assert!(saw_y);
5157 }
5158
5159 #[test]
5160 fn monotonicity_constraints_cluster_near_collinearrows() {
5161 let a = array![
5162 [0.0, 0.5, 0.0],
5163 [0.0, 0.50000000003, 0.0],
5164 [0.0, 0.49999999997, 0.0]
5165 ];
5166 let b = array![1e-8, 1.00000000005e-8, 0.99999999995e-8];
5167
5168 let compressed = compress_positive_collinear_constraints(&a, &b);
5169
5170 assert_eq!(compressed.a.nrows(), 1);
5171 assert_eq!(compressed.a.ncols(), 3);
5172 assert!(compressed.a[[0, 0]].abs() <= 1e-12);
5173 assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
5174 assert!(compressed.a[[0, 2]].abs() <= 1e-12);
5175 assert!((compressed.b[0] - 2.0e-8).abs() <= 1e-18);
5176 }
5177
5178 #[test]
5179 fn monotonicity_constraints_cluster_spline_like_near_duplicates() {
5180 let a = array![
5181 [0.0, 0.401, 0.302, 0.197],
5182 [0.0, 0.40100000003, 0.30199999998, 0.19700000001],
5183 [0.0, 0.40099999997, 0.30200000002, 0.19699999999],
5184 [0.0, 0.125, 0.500, 0.375]
5185 ];
5186 let b = array![2.0e-8, 2.00000000004e-8, 1.99999999996e-8, 3.0e-8];
5187
5188 let compressed = compress_positive_collinear_constraints(&a, &b);
5189
5190 assert_eq!(compressed.a.nrows(), 2);
5191 let mut clustered_face = false;
5192 let mut distinct_face = false;
5193 for i in 0..compressed.a.nrows() {
5194 let row = compressed.a.row(i);
5195 if row[1] > 0.99 && row[2] > 0.7 && row[3] > 0.49 {
5196 clustered_face = true;
5197 assert!((compressed.b[i] - (2.0e-8 / 0.401)).abs() <= 1e-12);
5198 } else {
5199 distinct_face = true;
5200 assert!((row[1] - 0.25).abs() <= 1e-12);
5201 assert!((row[2] - 1.0).abs() <= 1e-12);
5202 assert!((row[3] - 0.75).abs() <= 1e-12);
5203 assert!((compressed.b[i] - 6.0e-8).abs() <= 1e-18);
5204 }
5205 }
5206 assert!(clustered_face);
5207 assert!(distinct_face);
5208 }
5209
5210 #[test]
5211 fn linear_time_monotonicity_constraints_reduce_to_single_halfspace() {
5212 let age_entry = array![1.0_f64, 1.0, 1.0];
5213 let age_exit = array![2.0_f64, 4.0, 8.0];
5214 let event_target = array![0u8, 1u8, 0u8];
5215 let event_competing = array![0u8, 0u8, 0u8];
5216 let sampleweight = array![1.0, 1.0, 1.0];
5217 let x_entry = array![
5218 [1.0, age_entry[0].ln()],
5219 [1.0, age_entry[1].ln()],
5220 [1.0, age_entry[2].ln()]
5221 ];
5222 let x_exit = array![
5223 [1.0, age_exit[0].ln()],
5224 [1.0, age_exit[1].ln()],
5225 [1.0, age_exit[2].ln()]
5226 ];
5227 let x_derivative = array![[0.0, 0.5], [0.0, 0.25], [0.0, 0.125]];
5228 let penalties = PenaltyBlocks::new(Vec::new());
5229 let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5230
5231 let collocation_offsets = Array1::zeros(x_derivative.nrows());
5232 let mut inputs = survival_inputs(
5233 &age_entry,
5234 &age_exit,
5235 &event_target,
5236 &event_competing,
5237 &sampleweight,
5238 &x_entry,
5239 &x_exit,
5240 &x_derivative,
5241 );
5242 inputs.monotonicity_constraint_rows = Some(x_derivative.view());
5243 inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
5244
5245 let model = survival_model(inputs, penalties, mono, SurvivalSpec::Net)
5246 .expect("construct linear survival model");
5247
5248 let constraints = model
5249 .monotonicity_linear_constraints()
5250 .expect("monotonicity constraints");
5251 assert_eq!(constraints.a.nrows(), 1);
5252 assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
5253 assert!((constraints.b[0] - 8e-8).abs() <= 1e-12);
5254 }
5255
5256 #[test]
5257 fn monotonicity_constraints_skip_numericallyzerorows() {
5258 let age_entry = array![1.0_f64, 1.0, 1.0];
5259 let age_exit = array![2.0_f64, 3.0, 4.0];
5260 let event_target = array![0u8, 0u8, 0u8];
5261 let event_competing = array![0u8, 0u8, 0u8];
5262 let sampleweight = array![1.0, 1.0, 1.0];
5263 let x_entry = array![[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]];
5264 let x_exit = x_entry.clone();
5265 let x_derivative = array![[0.0, 0.0], [0.0, 1e-16], [0.0, 0.25]];
5266
5267 let collocation_offsets = Array1::zeros(x_derivative.nrows());
5268 let mut inputs = survival_inputs(
5269 &age_entry,
5270 &age_exit,
5271 &event_target,
5272 &event_competing,
5273 &sampleweight,
5274 &x_entry,
5275 &x_exit,
5276 &x_derivative,
5277 );
5278 inputs.monotonicity_constraint_rows = Some(x_derivative.view());
5279 inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
5280
5281 let model = survival_model(
5282 inputs,
5283 PenaltyBlocks::new(Vec::new()),
5284 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5285 SurvivalSpec::Net,
5286 )
5287 .expect("construct survival model");
5288
5289 let constraints = model
5290 .monotonicity_linear_constraints()
5291 .expect("nonzero derivative row should remain");
5292 assert_eq!(constraints.a.nrows(), 1);
5293 assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
5294 assert!(constraints.b[0].abs() <= 1e-18);
5295 }
5296
5297 #[test]
5298 fn censoredrows_allowzero_boundary_derivative() {
5299 let age_entry = array![1.0_f64];
5300 let age_exit = array![2.0_f64];
5301 let event_target = array![0u8];
5302 let event_competing = array![0u8];
5303 let sampleweight = array![1.0];
5304 let x_entry = array![[0.0]];
5305 let x_exit = array![[0.0]];
5306 let x_derivative = array![[1.0]];
5307
5308 let model = survival_model(
5309 survival_inputs(
5310 &age_entry,
5311 &age_exit,
5312 &event_target,
5313 &event_competing,
5314 &sampleweight,
5315 &x_entry,
5316 &x_exit,
5317 &x_derivative,
5318 ),
5319 PenaltyBlocks::new(Vec::new()),
5320 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5321 SurvivalSpec::Net,
5322 )
5323 .expect("construct censored survival model");
5324
5325 let state = model
5326 .update_state(&array![0.0])
5327 .expect("censored boundary derivative should remain feasible with zero tolerance");
5328 assert!(state.deviance.is_finite());
5329 }
5330
5331 #[test]
5332 fn eventrows_keep_positive_derivative_constraint() {
5333 let age_entry = array![1.0_f64, 1.0];
5334 let age_exit = array![2.0_f64, 4.0];
5335 let event_target = array![0u8, 1u8];
5336 let event_competing = array![0u8, 0u8];
5337 let sampleweight = array![1.0, 1.0];
5338 let x_entry = array![[0.0], [0.0]];
5339 let x_exit = array![[0.0], [0.0]];
5340 let x_derivative = array![[0.5], [0.25]];
5341
5342 let collocation_offsets = Array1::zeros(x_derivative.nrows());
5343 let mut inputs = survival_inputs(
5344 &age_entry,
5345 &age_exit,
5346 &event_target,
5347 &event_competing,
5348 &sampleweight,
5349 &x_entry,
5350 &x_exit,
5351 &x_derivative,
5352 );
5353 inputs.monotonicity_constraint_rows = Some(x_derivative.view());
5354 inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
5355
5356 let model = survival_model(
5357 inputs,
5358 PenaltyBlocks::new(Vec::new()),
5359 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
5360 SurvivalSpec::Net,
5361 )
5362 .expect("construct mixed survival model");
5363
5364 let constraints = model
5365 .monotonicity_linear_constraints()
5366 .expect("event row should induce positive lower bound");
5367 assert_eq!(constraints.a.nrows(), 1);
5368 assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
5369 assert!((constraints.b[0] - 4e-8).abs() <= 1e-18);
5370 }
5371
5372 #[test]
5373 fn structural_monotonicity_clamps_tiny_negative_roundoff() {
5374 let age_entry = array![1.0_f64];
5375 let age_exit = array![2.0_f64];
5376 let event_target = array![1u8];
5377 let event_competing = array![0u8];
5378 let sampleweight = array![1.0];
5379 let x_entry = array![[0.0]];
5380 let x_exit = array![[0.0]];
5381 let x_derivative = array![[1.0]];
5382 let mut model = survival_model(
5383 survival_inputs(
5384 &age_entry,
5385 &age_exit,
5386 &event_target,
5387 &event_competing,
5388 &sampleweight,
5389 &x_entry,
5390 &x_exit,
5391 &x_derivative,
5392 ),
5393 PenaltyBlocks::new(Vec::new()),
5394 SurvivalMonotonicityPenalty { tolerance: 1e-8 },
5395 SurvivalSpec::Net,
5396 )
5397 .expect("construct survival model");
5398 model
5399 .set_structural_monotonicity(true, 1)
5400 .expect("enable structural monotonicity");
5401
5402 let state = model
5403 .update_state(&array![-1e-8])
5404 .expect("tiny structural roundoff should be clamped");
5405 assert!(state.deviance.is_finite());
5406 }
5407
5408 #[test]
5409 fn compressed_monotonicity_constraints_preserve_uncompressed_feasible_region() {
5410 let uncompressed_constraints = LinearInequalityConstraints {
5411 a: array![
5412 [0.0, 0.5, 0.0],
5413 [0.0, 1.0 / 3.0, 0.0],
5414 [0.0, 0.2, 0.0],
5415 [0.0, 0.125, 0.0]
5416 ],
5417 b: Array1::from_elem(4, 1e-8),
5418 };
5419 let compressed_constraints = compress_positive_collinear_constraints(
5420 &uncompressed_constraints.a,
5421 &uncompressed_constraints.b,
5422 );
5423
5424 let candidates = [
5425 array![0.0, 1e-9, 0.0],
5426 array![0.0, 4e-8, 0.0],
5427 array![0.0, 8e-8, 0.0],
5428 array![0.0, 2e-7, 1.5],
5429 ];
5430 for beta in candidates {
5431 let uncompressed_ok = (0..uncompressed_constraints.a.nrows()).all(|i| {
5432 uncompressed_constraints.a.row(i).dot(&beta) >= uncompressed_constraints.b[i]
5433 });
5434 let compressed_ok = (0..compressed_constraints.a.nrows())
5435 .all(|i| compressed_constraints.a.row(i).dot(&beta) >= compressed_constraints.b[i]);
5436 assert_eq!(compressed_ok, uncompressed_ok);
5437 }
5438 }
5439
5440 #[test]
5441 fn exact_survival_derivatives_are_time_unit_invariant_up_to_constant_shift() {
5442 let age_entry = array![10.0_f64, 20.0, 25.0];
5443 let age_exit = array![15.0_f64, 30.0, 40.0];
5444 let event_target = array![1u8, 0u8, 1u8];
5445 let event_competing = array![0u8, 0u8, 0u8];
5446 let sampleweight = array![1.0, 2.0, 0.5];
5447 let x_entry = array![[0.1, 0.2, 1.0], [0.3, 0.4, 1.0], [0.2, 0.6, 1.0]];
5448 let x_exit = array![[0.2, 0.3, 1.0], [0.5, 0.7, 1.0], [0.4, 0.8, 1.0]];
5449 let x_derivative = array![[0.04, 0.02, 0.0], [0.03, 0.01, 0.0], [0.02, 0.03, 0.0]];
5450 let beta = array![0.8, 1.1, -0.2];
5451
5452 let base_model = survival_model(
5453 survival_inputs(
5454 &age_entry,
5455 &age_exit,
5456 &event_target,
5457 &event_competing,
5458 &sampleweight,
5459 &x_entry,
5460 &x_exit,
5461 &x_derivative,
5462 ),
5463 PenaltyBlocks::new(Vec::new()),
5464 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5465 SurvivalSpec::Net,
5466 )
5467 .expect("construct base survival model");
5468 let base_state = base_model
5469 .update_state(&beta)
5470 .expect("evaluate base survival state");
5471
5472 let time_scale = 365.25;
5473 let scaled_age_entry = age_entry.mapv(|v| v * time_scale);
5474 let scaled_age_exit = age_exit.mapv(|v| v * time_scale);
5475 let scaled_x_derivative = x_derivative.mapv(|v| v / time_scale);
5476 let scaled_model = survival_model(
5477 survival_inputs(
5478 &scaled_age_entry,
5479 &scaled_age_exit,
5480 &event_target,
5481 &event_competing,
5482 &sampleweight,
5483 &x_entry,
5484 &x_exit,
5485 &scaled_x_derivative,
5486 ),
5487 PenaltyBlocks::new(Vec::new()),
5488 SurvivalMonotonicityPenalty { tolerance: 0.0 },
5489 SurvivalSpec::Net,
5490 )
5491 .expect("construct scaled survival model");
5492 let scaled_state = scaled_model
5493 .update_state(&beta)
5494 .expect("evaluate scaled survival state");
5495
5496 let weighted_events = sampleweight
5497 .iter()
5498 .zip(event_target.iter())
5499 .map(|(w, d)| *w * f64::from(*d))
5500 .sum::<f64>();
5501 let expected_deviance_shift = 2.0 * weighted_events * time_scale.ln();
5502 assert!(
5503 (scaled_state.deviance - base_state.deviance - expected_deviance_shift).abs() <= 1e-10,
5504 "deviance shift mismatch: scaled={} base={} expected_shift={expected_deviance_shift}",
5505 scaled_state.deviance,
5506 base_state.deviance
5507 );
5508
5509 for j in 0..beta.len() {
5510 assert!(
5511 (scaled_state.gradient[j] - base_state.gradient[j]).abs() <= 1e-12,
5512 "gradient mismatch at j={j}: scaled={} base={}",
5513 scaled_state.gradient[j],
5514 base_state.gradient[j]
5515 );
5516 }
5517
5518 let base_hessian = base_state.hessian.to_dense();
5519 let scaled_hessian = scaled_state.hessian.to_dense();
5520 for r in 0..beta.len() {
5521 for c in 0..beta.len() {
5522 assert!(
5523 (scaled_hessian[[r, c]] - base_hessian[[r, c]]).abs() <= 1e-12,
5524 "hessian mismatch at ({r},{c}): scaled={} base={}",
5525 scaled_hessian[[r, c]],
5526 base_hessian[[r, c]]
5527 );
5528 }
5529 }
5530 }
5531}