use crate::custom_family::{
BlockWorkingSet, CustomFamily, FamilyEvaluation, ParameterBlockState,
projected_linear_constraint_stationarity_vector,
};
use crate::model_types::EstimationError;
use gam_linalg::faer_ndarray::{fast_atv, fast_av, fast_xt_diag_x, fast_xt_diag_y};
use gam_linalg::matrix::SymmetricMatrix;
use gam_problem::{Coefficients, LinearPredictor};
use gam_row_macros::row_atom;
use gam_solve::pirls::{
ConstraintSet, LinearInequalityConstraints, WorkingModel as PirlsWorkingModel, WorkingState,
array1_l2_norm,
};
use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayView3, Axis};
use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, constants, escalate_ridge};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::ops::Range;
use std::sync::LazyLock;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum SurvivalError {
#[error("input dimensions are inconsistent")]
DimensionMismatch,
#[error("inputs contain non-finite values")]
NonFiniteInput,
#[error("survival spec '{0}' is not supported by the one-hazard survival engine")]
UnsupportedSpec(&'static str),
#[error("crude risk integration setup is invalid")]
InvalidIntegrationSetup,
#[error("survival time grid must be finite, non-negative, and strictly increasing")]
InvalidTimeGrid,
#[error("cumulative hazard must be nondecreasing")]
NonMonotoneCumulativeHazard,
#[error("instantaneous hazard must stay strictly positive during integration")]
NonPositiveHazard,
#[error("{reason}")]
InvalidInput { reason: String },
#[error("{reason}")]
CauseSpecificDimensionMismatch { reason: String },
#[error("{reason}")]
NumericalFailure { reason: String },
#[error("{reason}")]
EventCodeInvalid { reason: String },
#[error("{reason}")]
EventDegenerate { reason: String },
#[error("cause-specific survival block {block}: {source}")]
CauseSpecificBlock {
block: usize,
#[source]
source: Box<SurvivalError>,
},
}
impl From<SurvivalError> for String {
fn from(err: SurvivalError) -> Self {
err.to_string()
}
}
impl From<crate::block_layout::block_count::BlockCountMismatch> for SurvivalError {
fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> SurvivalError {
SurvivalError::CauseSpecificDimensionMismatch {
reason: err.message(),
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
pub enum SurvivalSpec {
#[default]
Net,
Crude,
}
#[derive(Debug, Clone)]
pub struct SurvivalEngineInputs<'a> {
pub age_entry: ArrayView1<'a, f64>,
pub age_exit: ArrayView1<'a, f64>,
pub event_target: ArrayView1<'a, u8>,
pub event_competing: ArrayView1<'a, u8>,
pub sampleweight: ArrayView1<'a, f64>,
pub x_entry: ArrayView2<'a, f64>,
pub x_exit: ArrayView2<'a, f64>,
pub x_derivative: ArrayView2<'a, f64>,
pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
}
#[derive(Debug, Clone)]
pub struct SurvivalTimeCovarInputs<'a> {
pub age_entry: ArrayView1<'a, f64>,
pub age_exit: ArrayView1<'a, f64>,
pub event_target: ArrayView1<'a, u8>,
pub event_competing: ArrayView1<'a, u8>,
pub sampleweight: ArrayView1<'a, f64>,
pub time_entry: ArrayView2<'a, f64>,
pub time_exit: ArrayView2<'a, f64>,
pub time_derivative: ArrayView2<'a, f64>,
pub covariates: ArrayView2<'a, f64>,
pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
}
#[derive(Debug, Clone)]
pub struct SurvivalBaselineOffsets<'a> {
pub eta_entry: ArrayView1<'a, f64>,
pub eta_exit: ArrayView1<'a, f64>,
pub derivative_exit: ArrayView1<'a, f64>,
}
#[derive(Debug, Clone)]
pub struct PenaltyBlock {
pub matrix: Array2<f64>,
pub lambda: f64,
pub range: Range<usize>,
pub nullspace_dim: usize,
}
#[derive(Debug, Clone)]
pub struct PenaltyBlocks {
pub blocks: Vec<PenaltyBlock>,
}
impl PenaltyBlocks {
pub fn new(blocks: Vec<PenaltyBlock>) -> Self {
Self { blocks }
}
pub fn gradient(&self, beta: &Array1<f64>) -> Array1<f64> {
let mut grad = Array1::zeros(beta.len());
for block in &self.blocks {
if block.lambda == 0.0 {
continue;
}
let b = beta.slice(ndarray::s![block.range.clone()]);
let g = block.matrix.dot(&b);
let mut dst = grad.slice_mut(ndarray::s![block.range.clone()]);
dst += &(block.lambda * g);
}
grad
}
pub fn hessian(&self, dim: usize) -> Array2<f64> {
let mut h = Array2::zeros((dim, dim));
self.addhessian_inplace(&mut h);
h
}
pub fn deviance(&self, beta: &Array1<f64>) -> f64 {
let mut value = 0.0;
for block in &self.blocks {
if block.lambda == 0.0 {
continue;
}
let b = beta.slice(ndarray::s![block.range.clone()]);
value += 0.5 * block.lambda * b.dot(&block.matrix.dot(&b));
}
value
}
pub fn addhessian_inplace(&self, h: &mut Array2<f64>) {
for block in &self.blocks {
if block.lambda == 0.0 {
continue;
}
let start = block.range.start;
let end = block.range.end;
h.slice_mut(ndarray::s![start..end, start..end])
.scaled_add(block.lambda, &block.matrix);
}
}
}
pub const ENTRY_AT_ORIGIN_THRESHOLD: f64 = 1e-8;
const DERIVATIVE_FRACTION_TO_BOUNDARY: f64 = 0.995;
pub(crate) const SURVIVAL_LAML_STATIONARITY_RELATIVE_TOL: f64 = 1.0e-8;
#[derive(Debug, Clone)]
pub struct CauseSpecificRoystonParmarBlock {
pub age_entry: Array1<f64>,
pub age_exit: Array1<f64>,
pub event_target: Array1<u8>,
pub sampleweight: Array1<f64>,
pub x_entry: Array2<f64>,
pub x_exit: Array2<f64>,
pub x_derivative: Array2<f64>,
pub offset_eta_entry: Array1<f64>,
pub offset_eta_exit: Array1<f64>,
pub offset_derivative_exit: Array1<f64>,
pub derivative_floor: f64,
pub structural_time_columns: usize,
}
#[derive(Debug, Clone)]
pub struct CauseSpecificRoystonParmarFamily {
blocks: Vec<CauseSpecificRoystonParmarBlock>,
}
impl CauseSpecificRoystonParmarFamily {
pub fn new(blocks: Vec<CauseSpecificRoystonParmarBlock>) -> Result<Self, String> {
if blocks.is_empty() {
return Err(SurvivalError::InvalidInput {
reason: "cause-specific survival family requires at least one endpoint".to_string(),
}
.into());
}
for (idx, block) in blocks.iter().enumerate() {
validate_cause_specific_block(block).map_err(|err| {
SurvivalError::CauseSpecificBlock {
block: idx + 1,
source: Box::new(err),
}
.to_string()
})?;
}
Ok(Self { blocks })
}
pub fn cause_count(&self) -> usize {
self.blocks.len()
}
}
fn validate_cause_specific_block(
block: &CauseSpecificRoystonParmarBlock,
) -> Result<(), SurvivalError> {
let n = block.event_target.len();
let p = block.x_exit.ncols();
if n == 0 || p == 0 {
bail_invalid_surv!("empty event vector or coefficient block");
}
if block.age_entry.len() != n
|| block.age_exit.len() != n
|| block.sampleweight.len() != n
|| block.x_entry.nrows() != n
|| block.x_exit.nrows() != n
|| block.x_derivative.nrows() != n
|| block.x_entry.ncols() != p
|| block.x_derivative.ncols() != p
|| block.offset_eta_entry.len() != n
|| block.offset_eta_exit.len() != n
|| block.offset_derivative_exit.len() != n
{
return Err(SurvivalError::CauseSpecificDimensionMismatch {
reason: "dimension mismatch".to_string(),
});
}
if let Some(&label) = block.event_target.iter().find(|&&v| v > 1) {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"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"
),
});
}
if block.age_entry.iter().any(|v| !v.is_finite())
|| block.age_exit.iter().any(|v| !v.is_finite())
|| block
.sampleweight
.iter()
.any(|v| !v.is_finite() || *v < 0.0)
|| block.x_entry.iter().any(|v| !v.is_finite())
|| block.x_exit.iter().any(|v| !v.is_finite())
|| block.x_derivative.iter().any(|v| !v.is_finite())
|| block.offset_eta_entry.iter().any(|v| !v.is_finite())
|| block.offset_eta_exit.iter().any(|v| !v.is_finite())
|| block.offset_derivative_exit.iter().any(|v| !v.is_finite())
|| !block.derivative_floor.is_finite()
|| block.derivative_floor < 0.0
{
bail_invalid_surv!("non-finite input");
}
Ok(())
}
row_atom! {
fn cause_specific_row [generic, order2, third, fourth](
eta_exit,
eta_entry,
derivative;
weight,
entry_active,
event
) {
weight
* (exp(eta_exit)
- entry_active * exp(eta_entry)
- event * (eta_exit + ln(derivative)))
}
}
pub struct CauseSpecificSurvivalAloRowInput {
pub eta_exit: f64,
pub eta_entry: f64,
pub derivative_exit: f64,
pub prior_weight: f64,
pub entry_active: bool,
pub event: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CauseSpecificSurvivalAloRowGeometry {
pub negative_log_likelihood: f64,
pub nll_score: [f64; 3],
pub observed_hessian: [[f64; 3]; 3],
}
pub fn cause_specific_survival_alo_row_geometry(
input: CauseSpecificSurvivalAloRowInput,
) -> Result<CauseSpecificSurvivalAloRowGeometry, String> {
if !input.prior_weight.is_finite() || input.prior_weight < 0.0 {
return Err(format!(
"cause-specific saved ALO prior weight must be finite and non-negative, got {}",
input.prior_weight
));
}
if input.prior_weight == 0.0 {
return Ok(CauseSpecificSurvivalAloRowGeometry {
negative_log_likelihood: 0.0,
nll_score: [0.0; 3],
observed_hessian: [[0.0; 3]; 3],
});
}
if !input.eta_exit.is_finite() {
return Err(format!(
"cause-specific saved ALO exit index must be finite, got {}",
input.eta_exit
));
}
let eta_entry = if input.entry_active {
if !input.eta_entry.is_finite() {
return Err(format!(
"cause-specific saved ALO active entry index must be finite, got {}",
input.eta_entry
));
}
input.eta_entry
} else {
0.0
};
let derivative_exit = if input.event {
if !input.derivative_exit.is_finite() || input.derivative_exit <= 0.0 {
return Err(format!(
"cause-specific saved ALO event derivative must be positive and finite, got {}",
input.derivative_exit
));
}
input.derivative_exit
} else {
1.0
};
let atom = cause_specific_row_order2(
input.eta_exit,
eta_entry,
derivative_exit,
input.prior_weight,
f64::from(input.entry_active),
f64::from(input.event),
);
let gradient = atom.gradient();
let observed_hessian =
std::array::from_fn(|row| std::array::from_fn(|column| atom.hessian_at(row, column)));
if !atom.value().is_finite()
|| gradient.iter().any(|value| !value.is_finite())
|| observed_hessian
.iter()
.flatten()
.any(|value| !value.is_finite())
{
return Err(format!(
"cause-specific saved ALO row geometry is non-finite: nll={}, score={gradient:?}, hessian={observed_hessian:?}",
atom.value(),
));
}
Ok(CauseSpecificSurvivalAloRowGeometry {
negative_log_likelihood: atom.value(),
nll_score: gradient,
observed_hessian,
})
}
#[derive(Clone, Copy)]
struct CauseSpecificAtomInput {
primary: [f64; 3],
weight: f64,
entry_active: f64,
event: f64,
}
pub struct CauseSpecificRowProgram {
primary: [f64; 3],
weight: f64,
entry_active: f64,
event: f64,
}
impl CauseSpecificRowProgram {
pub fn new(primary: [f64; 3], weight: f64, entry_active: bool, event: bool) -> Self {
Self {
primary,
weight,
entry_active: f64::from(entry_active),
event: f64::from(event),
}
}
fn require_row(row: usize) -> Result<(), String> {
if row != 0 {
return Err(format!(
"CauseSpecificRowProgram holds exactly one row; got row {row}"
));
}
Ok(())
}
}
impl gam_math::jet_tower::RowProgram<3> for CauseSpecificRowProgram {
fn n_rows(&self) -> usize {
1
}
fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
Self::require_row(row)?;
Ok(self.primary)
}
fn eval<S: gam_math::jet_scalar::JetScalar<3>>(
&self,
row: usize,
p: &[S; 3],
) -> Result<S, String> {
Self::require_row(row)?;
Ok(cause_specific_row(
&p[0],
&p[1],
&p[2],
self.weight,
self.entry_active,
self.event,
))
}
}
fn cause_specific_atom_input(
block: &CauseSpecificRoystonParmarBlock,
row: usize,
eta_entry: f64,
eta_exit: f64,
derivative: f64,
) -> Result<Option<CauseSpecificAtomInput>, SurvivalError> {
let weight = block.sampleweight[row];
if weight <= 0.0 {
return Ok(None);
}
if block.age_exit[row] < block.age_entry[row] {
bail_invalid_surv!("age_exit < age_entry at row {row}");
}
let entry_active = block.age_entry[row] > ENTRY_AT_ORIGIN_THRESHOLD;
let event = block.event_target[row] > 0;
let eta_entry = if entry_active { eta_entry } else { 0.0 };
let derivative = if event {
if !(derivative.is_finite() && derivative > 0.0) {
return Err(SurvivalError::NumericalFailure {
reason: format!(
"cause-specific survival derivative must be positive at row {row}, got {derivative}"
),
});
}
derivative
} else {
1.0
};
let h_exit = eta_exit.exp();
let h_entry = eta_entry.exp();
if !(h_exit.is_finite() && h_entry.is_finite()) {
return Err(SurvivalError::NumericalFailure {
reason: format!("non-finite cumulative hazard at row {row}"),
});
}
Ok(Some(CauseSpecificAtomInput {
primary: [eta_exit, eta_entry, derivative],
weight,
entry_active: f64::from(entry_active),
event: f64::from(event),
}))
}
const CAUSE_SPECIFIC_PRIMARY_PAIRS: [(usize, usize); 6] =
[(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)];
fn cause_specific_pullback_hessian(
block: &CauseSpecificRoystonParmarBlock,
weights: &[Array1<f64>; 6],
) -> Array2<f64> {
let designs = [&block.x_exit, &block.x_entry, &block.x_derivative];
let p = block.x_exit.ncols();
let mut hessian = Array2::<f64>::zeros((p, p));
for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
let channel = &weights[slot];
if channel.iter().all(|&value| value == 0.0) {
continue;
}
if left == right {
hessian += &fast_xt_diag_x(designs[left], channel);
} else {
let cross = fast_xt_diag_y(designs[left], channel, designs[right]);
hessian += ✗
hessian += &cross.t();
}
}
hessian
}
fn evaluate_cause_specific_block(
block: &CauseSpecificRoystonParmarBlock,
beta: &Array1<f64>,
) -> Result<(f64, Array1<f64>, Array2<f64>), SurvivalError> {
let n = block.event_target.len();
let p = block.x_exit.ncols();
if beta.len() != p {
return Err(SurvivalError::CauseSpecificDimensionMismatch {
reason: format!("beta length mismatch: got {}, expected {p}", beta.len()),
});
}
let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
let mut log_likelihood = 0.0;
let mut gradient_weights: [Array1<f64>; 3] = std::array::from_fn(|_| Array1::<f64>::zeros(n));
let mut hessian_weights: [Array1<f64>; 6] = std::array::from_fn(|_| Array1::<f64>::zeros(n));
for i in 0..n {
let Some(input) =
cause_specific_atom_input(block, i, eta_entry[i], eta_exit[i], derivative[i])?
else {
continue;
};
let atom = cause_specific_row_order2(
input.primary[0],
input.primary[1],
input.primary[2],
input.weight,
input.entry_active,
input.event,
);
log_likelihood -= atom.value();
let gradient = atom.gradient();
for axis in 0..3 {
gradient_weights[axis][i] = -gradient[axis];
}
for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
hessian_weights[slot][i] = atom.hessian_at(left, right);
}
}
let designs = [&block.x_exit, &block.x_entry, &block.x_derivative];
let mut gradient = Array1::<f64>::zeros(p);
for axis in 0..3 {
gradient += &fast_atv(designs[axis], &gradient_weights[axis]);
}
let hessian = cause_specific_pullback_hessian(block, &hessian_weights);
Ok((log_likelihood, gradient, hessian))
}
impl CustomFamily for CauseSpecificRoystonParmarFamily {
fn joint_jeffreys_term_required(&self) -> bool {
true
}
fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
crate::block_layout::block_count::validate_block_count::<SurvivalError>(
"cause-specific survival",
self.blocks.len(),
block_states.len(),
)?;
let mut log_likelihood = 0.0;
let mut blockworking_sets = Vec::with_capacity(self.blocks.len());
for (block, state) in self.blocks.iter().zip(block_states.iter()) {
let (ll, gradient, hessian) = evaluate_cause_specific_block(block, &state.beta)?;
log_likelihood += ll;
blockworking_sets.push(BlockWorkingSet::ExactNewton {
gradient,
hessian: SymmetricMatrix::Dense(hessian),
});
}
Ok(FamilyEvaluation {
log_likelihood,
blockworking_sets,
})
}
fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
crate::block_layout::block_count::validate_block_count::<SurvivalError>(
"cause-specific survival",
self.blocks.len(),
block_states.len(),
)?;
let mut log_likelihood = 0.0;
for (block, state) in self.blocks.iter().zip(block_states.iter()) {
let (ll, _, _) = evaluate_cause_specific_block(block, &state.beta)?;
log_likelihood += ll;
}
Ok(log_likelihood)
}
fn likelihood_blocks_uncoupled(&self) -> bool {
true
}
fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
true
}
fn output_channel_assignment(
&self,
specs: &[crate::custom_family::ParameterBlockSpec],
) -> Option<Vec<usize>> {
if specs.len() != self.blocks.len() {
return Some((0..self.blocks.len()).collect());
}
Some((0..specs.len()).collect())
}
fn coefficient_hessian_cost(&self, specs: &[crate::custom_family::ParameterBlockSpec]) -> u64 {
crate::custom_family::default_coefficient_hessian_cost(specs)
}
fn block_linear_constraints(
&self,
_: &[ParameterBlockState],
block_idx: usize,
spec: &crate::custom_family::ParameterBlockSpec,
) -> Result<Option<ConstraintSet>, String> {
let block = self.blocks.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected block index < {}, got {block_idx}",
self.blocks.len()
),
}
.to_string()
})?;
if block.x_derivative.ncols() != spec.design.ncols() {
return Err(SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival derivative design has {} columns but block '{}' has {}",
block.x_derivative.ncols(),
spec.name,
spec.design.ncols()
),
}
.into());
}
let rhs = block
.offset_derivative_exit
.mapv(|offset| block.derivative_floor - offset);
let p = block.x_derivative.ncols();
let n_rows = block.x_derivative.nrows();
let structural_cols = block.structural_time_columns.min(p);
if structural_cols == 0 {
return Ok(Some(ConstraintSet::Dense(LinearInequalityConstraints {
a: block.x_derivative.clone(),
b: rhs,
})));
}
let mut a = Array2::<f64>::zeros((n_rows + structural_cols, p));
a.slice_mut(ndarray::s![..n_rows, ..])
.assign(&block.x_derivative);
for j in 0..structural_cols {
a[[n_rows + j, j]] = 1.0;
}
let mut b = Array1::<f64>::zeros(n_rows + structural_cols);
b.slice_mut(ndarray::s![..n_rows]).assign(&rhs);
Ok(Some(ConstraintSet::Dense(LinearInequalityConstraints {
a,
b,
})))
}
fn max_feasible_step_size(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
delta: &Array1<f64>,
) -> Result<Option<f64>, String> {
let block = self.blocks.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected block index < {}, got {block_idx}",
self.blocks.len()
),
}
.to_string()
})?;
let state = block_states.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected {} block states, got {}",
self.blocks.len(),
block_states.len()
),
}
.to_string()
})?;
if delta.len() != state.beta.len() || block.x_derivative.ncols() != delta.len() {
return Err(SurvivalError::CauseSpecificDimensionMismatch {
reason: "cause-specific survival feasible-step dimension mismatch".to_string(),
}
.into());
}
let derivative = fast_av(&block.x_derivative, &state.beta) + &block.offset_derivative_exit;
let derivative_delta = fast_av(&block.x_derivative, delta);
let mut alpha_max = 1.0_f64;
for i in 0..derivative.len() {
if block.sampleweight[i] <= 0.0 {
continue;
}
let current = derivative[i] - block.derivative_floor;
let slope = derivative_delta[i];
if slope < 0.0 {
if current <= 0.0 {
return Ok(Some(0.0));
}
alpha_max = alpha_max.min(DERIVATIVE_FRACTION_TO_BOUNDARY * current / -slope);
}
}
Ok(Some(alpha_max.clamp(0.0, 1.0)))
}
fn exact_newton_hessian_directional_derivative(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
d_beta: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let block = self.blocks.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected block index < {}, got {block_idx}",
self.blocks.len()
),
}
.to_string()
})?;
let state = block_states.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected {} block states, got {}",
self.blocks.len(),
block_states.len()
),
}
.to_string()
})?;
Ok(Some(cause_specific_hessian_directional_derivative(
block,
&state.beta,
d_beta,
)?))
}
fn exact_newton_hessian_second_directional_derivative(
&self,
block_states: &[ParameterBlockState],
block_idx: usize,
d_beta_u: &Array1<f64>,
d_beta_v: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let block = self.blocks.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected block index < {}, got {block_idx}",
self.blocks.len()
),
}
.to_string()
})?;
let state = block_states.get(block_idx).ok_or_else(|| {
SurvivalError::CauseSpecificDimensionMismatch {
reason: format!(
"cause-specific survival expected {} block states, got {}",
self.blocks.len(),
block_states.len()
),
}
.to_string()
})?;
Ok(Some(cause_specific_hessian_second_directional_derivative(
block,
&state.beta,
d_beta_u,
d_beta_v,
)?))
}
}
fn cause_specific_hessian_directional_derivative(
block: &CauseSpecificRoystonParmarBlock,
beta: &Array1<f64>,
d_beta: &Array1<f64>,
) -> Result<Array2<f64>, SurvivalError> {
let p = block.x_exit.ncols();
if beta.len() != p || d_beta.len() != p {
return Err(SurvivalError::CauseSpecificDimensionMismatch {
reason: "cause-specific survival Hessian derivative dimension mismatch".to_string(),
});
}
let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
let d_eta_entry = fast_av(&block.x_entry, d_beta);
let d_eta_exit = fast_av(&block.x_exit, d_beta);
let d_derivative = fast_av(&block.x_derivative, d_beta);
let n = block.event_target.len();
let mut weights: [Array1<f64>; 6] = std::array::from_fn(|_| Array1::zeros(n));
for i in 0..n {
let Some(input) =
cause_specific_atom_input(block, i, eta_entry[i], eta_exit[i], derivative[i])?
else {
continue;
};
let direction = [
d_eta_exit[i],
d_eta_entry[i] * input.entry_active,
d_derivative[i] * input.event,
];
let matrix = cause_specific_row_third_contracted(
input.primary[0],
input.primary[1],
input.primary[2],
input.weight,
input.entry_active,
input.event,
&direction,
);
for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
weights[slot][i] = matrix[left][right];
}
}
Ok(cause_specific_pullback_hessian(block, &weights))
}
fn cause_specific_hessian_second_directional_derivative(
block: &CauseSpecificRoystonParmarBlock,
beta: &Array1<f64>,
d_beta_u: &Array1<f64>,
d_beta_v: &Array1<f64>,
) -> Result<Array2<f64>, SurvivalError> {
let p = block.x_exit.ncols();
if beta.len() != p || d_beta_u.len() != p || d_beta_v.len() != p {
return Err(SurvivalError::CauseSpecificDimensionMismatch {
reason: "cause-specific survival second Hessian derivative dimension mismatch"
.to_string(),
});
}
let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
let u_eta_entry = fast_av(&block.x_entry, d_beta_u);
let u_eta_exit = fast_av(&block.x_exit, d_beta_u);
let u_derivative = fast_av(&block.x_derivative, d_beta_u);
let v_eta_entry = fast_av(&block.x_entry, d_beta_v);
let v_eta_exit = fast_av(&block.x_exit, d_beta_v);
let v_derivative = fast_av(&block.x_derivative, d_beta_v);
let n = block.event_target.len();
let mut weights: [Array1<f64>; 6] = std::array::from_fn(|_| Array1::zeros(n));
for i in 0..n {
let Some(input) =
cause_specific_atom_input(block, i, eta_entry[i], eta_exit[i], derivative[i])?
else {
continue;
};
let direction_u = [
u_eta_exit[i],
u_eta_entry[i] * input.entry_active,
u_derivative[i] * input.event,
];
let direction_v = [
v_eta_exit[i],
v_eta_entry[i] * input.entry_active,
v_derivative[i] * input.event,
];
let matrix = cause_specific_row_fourth_contracted(
input.primary[0],
input.primary[1],
input.primary[2],
input.weight,
input.entry_active,
input.event,
&direction_u,
&direction_v,
);
for (slot, &(left, right)) in CAUSE_SPECIFIC_PRIMARY_PAIRS.iter().enumerate() {
weights[slot][i] = matrix[left][right];
}
}
Ok(cause_specific_pullback_hessian(block, &weights))
}
pub fn survival_event_code_from_value(value: f64, row_index: usize) -> Result<u8, String> {
const INTEGER_TOL: f64 = 1e-8;
const MAX_AUTO_CAUSES: u8 = 32;
if !value.is_finite() {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"survival event value at row {} is non-finite",
row_index + 1
),
}
.into());
}
if value < 0.0 {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"survival event value at row {} is negative: {value}",
row_index + 1
),
}
.into());
}
let rounded = value.round();
if (value - rounded).abs() > INTEGER_TOL {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"survival event value at row {} must be an integer code with 0=censored, got {value}",
row_index + 1
),
}
.into());
}
if rounded > f64::from(MAX_AUTO_CAUSES) {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"survival event value at row {} has code {rounded}; automatic competing-risks detection supports codes 0..={MAX_AUTO_CAUSES}",
row_index + 1
),
}
.into());
}
Ok(rounded as u8)
}
pub fn cause_count_from_event_codes(
event_codes: ArrayView1<'_, u8>,
) -> Result<usize, SurvivalError> {
let max_code = event_codes.iter().copied().max().map_or(0, usize::from);
if max_code == 0 {
return Ok(1);
}
let mut present = vec![false; max_code + 1];
for code in event_codes.iter().copied() {
present[usize::from(code)] = true;
}
if (1..=max_code).any(|code| !present[code]) {
let actual = present
.iter()
.enumerate()
.skip(1)
.filter_map(|(code, &seen)| seen.then_some(code.to_string()))
.collect::<Vec<_>>()
.join(", ");
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"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."
),
});
}
Ok(max_code)
}
pub fn pooled_any_event_indicator(event_codes: ArrayView1<'_, u8>) -> Array1<u8> {
event_codes.mapv(|label| u8::from(label > 0))
}
pub fn cause_specific_event_indicator(event_codes: ArrayView1<'_, u8>, cause: usize) -> Array1<u8> {
let cause_code = cause as u8;
event_codes.mapv(|observed| u8::from(observed == cause_code))
}
fn compress_positive_collinear_constraints(
a: &Array2<f64>,
b: &Array1<f64>,
) -> LinearInequalityConstraints {
const SCALE_TOL: f64 = 1e-14;
const KEY_TOL: f64 = 1e-8;
let mut grouped: BTreeMap<Vec<i64>, (Vec<f64>, f64)> = BTreeMap::new();
let mut fallbackrows: Vec<(Vec<f64>, f64)> = Vec::new();
for i in 0..a.nrows() {
let row = a.row(i);
let scale = row.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
if !scale.is_finite() || scale <= SCALE_TOL {
if b[i] > 0.0 {
fallbackrows.push((row.to_vec(), b[i]));
}
continue;
}
let normalizedrow: Vec<f64> = row
.iter()
.map(|&v| {
let scaled = v / scale;
if scaled.abs() <= KEY_TOL { 0.0 } else { scaled }
})
.collect();
let normalized_rhs = b[i] / scale;
let key: Vec<i64> = normalizedrow
.iter()
.map(|&v| (v / KEY_TOL).round() as i64)
.collect();
match grouped.get_mut(&key) {
Some((_, rhs_max)) => {
if normalized_rhs > *rhs_max {
*rhs_max = normalized_rhs;
}
}
None => {
grouped.insert(key, (normalizedrow, normalized_rhs));
}
}
}
let nrows = grouped.len() + fallbackrows.len();
let n_cols = a.ncols();
let mut a_out = Array2::<f64>::zeros((nrows, n_cols));
let mut b_out = Array1::<f64>::zeros(nrows);
let mut outrow = 0usize;
for (_, (row, rhs)) in grouped {
for (j, value) in row.into_iter().enumerate() {
a_out[[outrow, j]] = value;
}
b_out[outrow] = rhs;
outrow += 1;
}
for (row, rhs) in fallbackrows {
for (j, value) in row.into_iter().enumerate() {
a_out[[outrow, j]] = value;
}
b_out[outrow] = rhs;
outrow += 1;
}
LinearInequalityConstraints { a: a_out, b: b_out }
}
#[derive(Debug, Clone, Copy, Default)]
pub struct SurvivalMonotonicityPenalty {
pub tolerance: f64,
}
#[derive(Debug, Clone)]
enum SurvivalDesign {
Flat {
x_entry: Array2<f64>,
x_exit: Array2<f64>,
x_derivative: Array2<f64>,
},
TimeCovariateShared {
time_entry: Array2<f64>,
time_exit: Array2<f64>,
time_derivative: Array2<f64>,
covariates: Array2<f64>,
},
}
impl SurvivalDesign {
fn p_total(&self) -> usize {
match self {
Self::Flat { x_exit, .. } => x_exit.ncols(),
Self::TimeCovariateShared {
time_exit,
covariates,
..
} => time_exit.ncols() + covariates.ncols(),
}
}
fn design_dot(&self, time_mat: &Array2<f64>, beta: &Array1<f64>) -> Array1<f64> {
match self {
Self::Flat { .. } => time_mat.dot(beta),
Self::TimeCovariateShared { covariates, .. } => {
let p_time = time_mat.ncols();
let mut out = time_mat.dot(&beta.slice(ndarray::s![..p_time]));
if covariates.ncols() > 0 {
out += &covariates.dot(&beta.slice(ndarray::s![p_time..]));
}
out
}
}
}
fn fill_row(&self, time_mat: &Array2<f64>, i: usize, out: &mut [f64]) {
match self {
Self::Flat { .. } => {
for (dst, &src) in out.iter_mut().zip(time_mat.row(i).iter()) {
*dst = src;
}
}
Self::TimeCovariateShared { covariates, .. } => {
let p_time = time_mat.ncols();
for j in 0..p_time {
out[j] = time_mat[[i, j]];
}
for j in 0..covariates.ncols() {
out[p_time + j] = covariates[[i, j]];
}
}
}
}
}
#[derive(Debug, Clone)]
struct SurvivalWorkspace {
w_event: Array1<f64>,
w_event_inv_deriv: Array1<f64>,
w_event_outer: Array1<f64>,
w_hess_exit: Array1<f64>,
w_hess_entry: Array1<f64>,
}
impl SurvivalWorkspace {
fn new(n: usize) -> Self {
Self {
w_event: Array1::zeros(n),
w_event_inv_deriv: Array1::zeros(n),
w_event_outer: Array1::zeros(n),
w_hess_exit: Array1::zeros(n),
w_hess_entry: Array1::zeros(n),
}
}
fn reset(&mut self, n: usize) {
if self.w_event.len() != n {
*self = Self::new(n);
} else {
self.w_event.fill(0.0);
self.w_event_inv_deriv.fill(0.0);
self.w_event_outer.fill(0.0);
self.w_hess_exit.fill(0.0);
self.w_hess_entry.fill(0.0);
}
}
}
#[derive(Clone, Debug)]
pub struct OffsetChannelResiduals {
pub exit: Array1<f64>,
pub entry: Array1<f64>,
pub derivative: Array1<f64>,
pub right: Array1<f64>,
}
#[derive(Clone, Debug)]
pub struct OffsetChannelCurvatures {
pub rows: Vec<[[f64; 3]; 3]>,
}
#[derive(Debug)]
pub struct WorkingModelSurvival {
age_entry: Array1<f64>,
age_exit: Array1<f64>,
entry_at_origin: Array1<bool>,
event_target: Array1<u8>,
sampleweight: Array1<f64>,
design: SurvivalDesign,
offset_eta_entry: Array1<f64>,
offset_eta_exit: Array1<f64>,
offset_derivative_exit: Array1<f64>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
structurally_monotonic: bool,
structural_time_columns: usize,
monotonicity_constraint_rows: Option<Array2<f64>>,
monotonicity_constraint_offsets: Option<Array1<f64>>,
workspace: std::sync::Mutex<SurvivalWorkspace>,
}
impl Clone for WorkingModelSurvival {
fn clone(&self) -> Self {
let workspace = self.workspace.lock().unwrap().clone();
Self {
age_entry: self.age_entry.clone(),
age_exit: self.age_exit.clone(),
entry_at_origin: self.entry_at_origin.clone(),
event_target: self.event_target.clone(),
sampleweight: self.sampleweight.clone(),
design: self.design.clone(),
offset_eta_entry: self.offset_eta_entry.clone(),
offset_eta_exit: self.offset_eta_exit.clone(),
offset_derivative_exit: self.offset_derivative_exit.clone(),
penalties: self.penalties.clone(),
monotonicity: self.monotonicity,
structurally_monotonic: self.structurally_monotonic,
structural_time_columns: self.structural_time_columns,
monotonicity_constraint_rows: self.monotonicity_constraint_rows.clone(),
monotonicity_constraint_offsets: self.monotonicity_constraint_offsets.clone(),
workspace: std::sync::Mutex::new(workspace),
}
}
}
impl WorkingModelSurvival {
const LOG_F64_MAX: f64 = 709.782712893384;
#[inline]
fn scaled_exp_component(log_scale: f64, base: f64) -> Result<f64, EstimationError> {
if base == 0.0 {
return Ok(0.0);
}
let log_abs = log_scale + base.abs().ln();
if !log_abs.is_finite() {
crate::bail_invalid_estim!("survival interval term produced non-finite log-magnitude");
}
if log_abs > Self::LOG_F64_MAX {
crate::bail_invalid_estim!(
"survival interval term exceeds f64 range (log-magnitude={log_abs:.3e})"
);
}
Ok(base.signum() * log_abs.exp())
}
fn coefficient_dim(&self) -> usize {
self.design.p_total()
}
fn nrows(&self) -> usize {
self.sampleweight.len()
}
fn entry_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
let time_mat = match &self.design {
SurvivalDesign::Flat { x_entry, .. } => x_entry,
SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
};
self.design.design_dot(time_mat, beta)
}
fn exit_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
let time_mat = match &self.design {
SurvivalDesign::Flat { x_exit, .. } => x_exit,
SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
};
self.design.design_dot(time_mat, beta)
}
fn derivative_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
match &self.design {
SurvivalDesign::Flat { x_derivative, .. } => x_derivative.dot(beta),
SurvivalDesign::TimeCovariateShared {
time_derivative, ..
} => time_derivative.dot(&beta.slice(ndarray::s![..time_derivative.ncols()])),
}
}
fn fill_entry_row(&self, i: usize, out: &mut [f64]) {
let time_mat = match &self.design {
SurvivalDesign::Flat { x_entry, .. } => x_entry,
SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
};
self.design.fill_row(time_mat, i, out);
}
fn fill_exit_row(&self, i: usize, out: &mut [f64]) {
let time_mat = match &self.design {
SurvivalDesign::Flat { x_exit, .. } => x_exit,
SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
};
self.design.fill_row(time_mat, i, out);
}
fn fill_derivative_row(&self, i: usize, out: &mut [f64]) {
match &self.design {
SurvivalDesign::Flat { x_derivative, .. } => {
for (dst, &src) in out.iter_mut().zip(x_derivative.row(i).iter()) {
*dst = src;
}
}
SurvivalDesign::TimeCovariateShared {
time_derivative, ..
} => {
let p_time = time_derivative.ncols();
for j in 0..p_time {
out[j] = time_derivative[[i, j]];
}
for dst in out.iter_mut().skip(p_time) {
*dst = 0.0;
}
}
}
}
fn derivative_xt_diag_x(&self, weights: &Array1<f64>) -> Array2<f64> {
match &self.design {
SurvivalDesign::Flat { x_derivative, .. } => fast_xt_diag_x(x_derivative, weights),
SurvivalDesign::TimeCovariateShared {
time_derivative,
covariates,
..
} => {
let p_time = time_derivative.ncols();
let p_cov = covariates.ncols();
let mut out = Array2::<f64>::zeros((p_time + p_cov, p_time + p_cov));
let time_block = fast_xt_diag_x(time_derivative, weights);
out.slice_mut(ndarray::s![..p_time, ..p_time])
.assign(&time_block);
out
}
}
}
fn interval_hessian_blas(&self, w_exit: &Array1<f64>, w_entry: &Array1<f64>) -> Array2<f64> {
match &self.design {
SurvivalDesign::Flat {
x_entry, x_exit, ..
} => {
let mut h = fast_xt_diag_x(x_exit, w_exit);
h -= &fast_xt_diag_x(x_entry, w_entry);
h
}
SurvivalDesign::TimeCovariateShared {
time_entry,
time_exit,
covariates,
..
} => {
let p_time = time_exit.ncols();
let p_cov = covariates.ncols();
let p = p_time + p_cov;
let mut h = Array2::<f64>::zeros((p, p));
let tt = {
let mut block = fast_xt_diag_x(time_exit, w_exit);
block -= &fast_xt_diag_x(time_entry, w_entry);
block
};
h.slice_mut(ndarray::s![..p_time, ..p_time]).assign(&tt);
if p_cov > 0 {
let tc = {
let mut block = fast_xt_diag_y(time_exit, w_exit, covariates);
block -= &fast_xt_diag_y(time_entry, w_entry, covariates);
block
};
h.slice_mut(ndarray::s![..p_time, p_time..]).assign(&tc);
h.slice_mut(ndarray::s![p_time.., ..p_time]).assign(&tc.t());
let w_diff = w_exit - w_entry;
let cc = fast_xt_diag_x(covariates, &w_diff);
h.slice_mut(ndarray::s![p_time.., p_time..]).assign(&cc);
}
h
}
}
}
fn stabilized_structural_derivative(&self, deriv: f64) -> Option<(f64, f64)> {
const STRUCTURAL_MONO_ROUNDOFF_TOL: f64 = 1e-7;
const STRUCTURAL_DERIV_FLOOR: f64 = 1e-12;
if !self.structurally_monotonic {
return None;
}
if deriv >= STRUCTURAL_DERIV_FLOOR {
return Some((deriv, 1.0));
}
if deriv >= -STRUCTURAL_MONO_ROUNDOFF_TOL {
return Some((STRUCTURAL_DERIV_FLOOR, 0.0));
}
None
}
fn validate_penalties(
penalties: &PenaltyBlocks,
coefficient_dim: usize,
) -> Result<(), SurvivalError> {
for block in &penalties.blocks {
if !block.lambda.is_finite() || block.lambda < 0.0 {
return Err(SurvivalError::NonFiniteInput);
}
if block.range.start > block.range.end || block.range.end > coefficient_dim {
return Err(SurvivalError::DimensionMismatch);
}
let block_dim = block.range.end - block.range.start;
if block.matrix.nrows() != block_dim || block.matrix.ncols() != block_dim {
return Err(SurvivalError::DimensionMismatch);
}
if block.matrix.iter().any(|v| !v.is_finite()) {
return Err(SurvivalError::NonFiniteInput);
}
}
Ok(())
}
fn derivative_guard(&self) -> f64 {
if self.structurally_monotonic {
return 0.0;
}
self.monotonicity.tolerance.max(0.0)
}
fn derivative_guard_numerical(&self) -> f64 {
let derivative_guard = self.derivative_guard();
if derivative_guard <= 0.0 {
if self.structurally_monotonic {
-1e-10
} else {
1e-12
}
} else {
(derivative_guard - (1e-10_f64).min(0.01 * derivative_guard)).max(1e-12)
}
}
fn interval_increment_guard(&self, h_entry: f64, h_exit: f64) -> f64 {
let scale = h_entry.abs().max(h_exit.abs()).max(1.0);
1e-10 * scale
}
fn structural_time_coefficient_constraints(&self) -> Option<LinearInequalityConstraints> {
if !self.structurally_monotonic {
return None;
}
let p = self.coefficient_dim();
let time_columns = self.structural_time_columns.min(p);
if time_columns == 0 {
return None;
}
let mut a = Array2::<f64>::zeros((time_columns, p));
let b = Array1::<f64>::zeros(time_columns);
for j in 0..time_columns {
a[[j, j]] = 1.0;
}
Some(LinearInequalityConstraints { a, b })
}
pub fn monotonicity_linear_constraints(&self) -> Option<LinearInequalityConstraints> {
let p = self.coefficient_dim();
const DERIVATIVE_ROW_NORM_TOL: f64 = 1e-12;
if p == 0 {
return None;
}
if self.structurally_monotonic {
return self.structural_time_coefficient_constraints();
}
if let (Some(rows), Some(offsets)) = (
self.monotonicity_constraint_rows.as_ref(),
self.monotonicity_constraint_offsets.as_ref(),
) {
let activerows: Vec<usize> = (0..rows.nrows())
.filter(|&i| {
rows.row(i).iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()))
> DERIVATIVE_ROW_NORM_TOL
})
.collect();
if activerows.is_empty() {
return None;
}
let mut a = Array2::<f64>::zeros((activerows.len(), p));
let mut b = Array1::<f64>::zeros(activerows.len());
for (r, &i) in activerows.iter().enumerate() {
a.row_mut(r).assign(&rows.row(i));
b[r] = self.derivative_guard() - offsets[i];
}
return Some(compress_positive_collinear_constraints(&a, &b));
}
None
}
pub fn from_engine_inputs(
inputs: SurvivalEngineInputs<'_>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
spec: SurvivalSpec,
) -> Result<Self, SurvivalError> {
Self::from_engine_inputswith_offsets(inputs, None, penalties, monotonicity, spec)
}
fn validate_offsets(
offsets: Option<SurvivalBaselineOffsets<'_>>,
n: usize,
) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), SurvivalError> {
if let Some(off) = offsets {
if off.eta_entry.len() != n || off.eta_exit.len() != n || off.derivative_exit.len() != n
{
return Err(SurvivalError::DimensionMismatch);
}
if off.eta_entry.iter().any(|v| !v.is_finite())
|| off.eta_exit.iter().any(|v| !v.is_finite())
|| off.derivative_exit.iter().any(|v| !v.is_finite())
{
return Err(SurvivalError::NonFiniteInput);
}
Ok((
off.eta_entry.to_owned(),
off.eta_exit.to_owned(),
off.derivative_exit.to_owned(),
))
} else {
Ok((Array1::zeros(n), Array1::zeros(n), Array1::zeros(n)))
}
}
fn validate_common_inputs(
age_entry: &ArrayView1<f64>,
age_exit: &ArrayView1<f64>,
event_target: &ArrayView1<u8>,
event_competing: &ArrayView1<u8>,
sampleweight: &ArrayView1<f64>,
) -> Result<(), SurvivalError> {
if age_entry.iter().any(|v| !v.is_finite())
|| age_exit.iter().any(|v| !v.is_finite())
|| sampleweight.iter().any(|v| !v.is_finite() || *v < 0.0)
{
return Err(SurvivalError::NonFiniteInput);
}
if let Some(&label) = event_target.iter().find(|&&v| v > 1) {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"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"
),
});
}
if let Some(&label) = event_competing.iter().find(|&&v| v > 1) {
return Err(SurvivalError::EventCodeInvalid {
reason: format!(
"single-hazard survival engine requires a binary {{0, 1}} event_competing, got multi-cause label {label}"
),
});
}
if event_target
.iter()
.zip(event_competing.iter())
.any(|(&target, &competing)| target > 0 && competing > 0)
{
return Err(SurvivalError::EventCodeInvalid {
reason: "a row cannot be simultaneously a target event and a competing event"
.to_string(),
});
}
if age_entry
.iter()
.zip(age_exit.iter())
.any(|(&entry, &exit)| entry < 0.0 || exit <= 0.0)
{
return Err(SurvivalError::NonFiniteInput);
}
Ok::<(), _>(())
}
fn validate_monotonicity_constraints(
rows: Option<ArrayView2<'_, f64>>,
offsets: Option<ArrayView1<'_, f64>>,
coefficient_dim: usize,
) -> Result<(Option<Array2<f64>>, Option<Array1<f64>>), SurvivalError> {
match (rows, offsets) {
(None, None) => Ok((None, None)),
(Some(rows), Some(offsets)) => {
if rows.ncols() != coefficient_dim
|| rows.nrows() != offsets.len()
|| rows.iter().any(|v| !v.is_finite())
|| offsets.iter().any(|v| !v.is_finite())
{
return Err(SurvivalError::DimensionMismatch);
}
Ok((Some(rows.to_owned()), Some(offsets.to_owned())))
}
_ => Err(SurvivalError::DimensionMismatch),
}
}
fn finish_construction(
age_entry: ArrayView1<f64>,
age_exit: ArrayView1<f64>,
event_target: ArrayView1<u8>,
sampleweight: ArrayView1<f64>,
design: SurvivalDesign,
offset_eta_entry: Array1<f64>,
offset_eta_exit: Array1<f64>,
offset_derivative_exit: Array1<f64>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
monotonicity_constraint_rows: Option<Array2<f64>>,
monotonicity_constraint_offsets: Option<Array1<f64>>,
) -> Self {
let n = age_entry.len();
Self {
age_entry: age_entry.to_owned(),
age_exit: age_exit.to_owned(),
entry_at_origin: age_entry.mapv(|t| t <= ENTRY_AT_ORIGIN_THRESHOLD),
event_target: event_target.to_owned(),
sampleweight: sampleweight.to_owned(),
design,
offset_eta_entry,
offset_eta_exit,
offset_derivative_exit,
penalties,
monotonicity,
structurally_monotonic: false,
structural_time_columns: 0,
monotonicity_constraint_rows,
monotonicity_constraint_offsets,
workspace: std::sync::Mutex::new(SurvivalWorkspace::new(n)),
}
}
pub fn from_engine_inputswith_offsets(
inputs: SurvivalEngineInputs<'_>,
offsets: Option<SurvivalBaselineOffsets<'_>>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
spec: SurvivalSpec,
) -> Result<Self, SurvivalError> {
if spec == SurvivalSpec::Crude {
return Err(SurvivalError::UnsupportedSpec("crude"));
}
let n = inputs.age_entry.len();
let p = inputs.x_entry.ncols();
if inputs.age_exit.len() != n
|| inputs.event_target.len() != n
|| inputs.event_competing.len() != n
|| inputs.sampleweight.len() != n
|| inputs.x_entry.nrows() != n
|| inputs.x_exit.nrows() != n
|| inputs.x_derivative.nrows() != n
|| inputs.x_entry.ncols() != inputs.x_exit.ncols()
|| inputs.x_entry.ncols() != inputs.x_derivative.ncols()
{
return Err(SurvivalError::DimensionMismatch);
}
Self::validate_penalties(&penalties, p)?;
Self::validate_common_inputs(
&inputs.age_entry,
&inputs.age_exit,
&inputs.event_target,
&inputs.event_competing,
&inputs.sampleweight,
)?;
if inputs.x_entry.iter().any(|v| !v.is_finite())
|| inputs.x_exit.iter().any(|v| !v.is_finite())
|| inputs.x_derivative.iter().any(|v| !v.is_finite())
{
return Err(SurvivalError::NonFiniteInput);
}
let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
Self::validate_offsets(offsets, n)?;
let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
Self::validate_monotonicity_constraints(
inputs.monotonicity_constraint_rows,
inputs.monotonicity_constraint_offsets,
p,
)?;
Ok(Self::finish_construction(
inputs.age_entry,
inputs.age_exit,
inputs.event_target,
inputs.sampleweight,
SurvivalDesign::Flat {
x_entry: inputs.x_entry.to_owned(),
x_exit: inputs.x_exit.to_owned(),
x_derivative: inputs.x_derivative.to_owned(),
},
offset_eta_entry,
offset_eta_exit,
offset_derivative_exit,
penalties,
monotonicity,
monotonicity_constraint_rows,
monotonicity_constraint_offsets,
))
}
pub fn from_time_covariate_inputswith_offsets(
inputs: SurvivalTimeCovarInputs<'_>,
offsets: Option<SurvivalBaselineOffsets<'_>>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
spec: SurvivalSpec,
) -> Result<Self, SurvivalError> {
if spec == SurvivalSpec::Crude {
return Err(SurvivalError::UnsupportedSpec("crude"));
}
let n = inputs.age_entry.len();
let p_time = inputs.time_entry.ncols();
let p_cov = inputs.covariates.ncols();
let p = p_time + p_cov;
if inputs.age_exit.len() != n
|| inputs.event_target.len() != n
|| inputs.event_competing.len() != n
|| inputs.sampleweight.len() != n
|| inputs.time_entry.nrows() != n
|| inputs.time_exit.nrows() != n
|| inputs.time_derivative.nrows() != n
|| inputs.covariates.nrows() != n
|| inputs.time_entry.ncols() != inputs.time_exit.ncols()
|| inputs.time_entry.ncols() != inputs.time_derivative.ncols()
{
return Err(SurvivalError::DimensionMismatch);
}
Self::validate_penalties(&penalties, p)?;
Self::validate_common_inputs(
&inputs.age_entry,
&inputs.age_exit,
&inputs.event_target,
&inputs.event_competing,
&inputs.sampleweight,
)?;
if inputs.time_entry.iter().any(|v| !v.is_finite())
|| inputs.time_exit.iter().any(|v| !v.is_finite())
|| inputs.time_derivative.iter().any(|v| !v.is_finite())
|| inputs.covariates.iter().any(|v| !v.is_finite())
{
return Err(SurvivalError::NonFiniteInput);
}
let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
Self::validate_offsets(offsets, n)?;
let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
Self::validate_monotonicity_constraints(
inputs.monotonicity_constraint_rows,
inputs.monotonicity_constraint_offsets,
p,
)?;
Ok(Self::finish_construction(
inputs.age_entry,
inputs.age_exit,
inputs.event_target,
inputs.sampleweight,
SurvivalDesign::TimeCovariateShared {
time_entry: inputs.time_entry.to_owned(),
time_exit: inputs.time_exit.to_owned(),
time_derivative: inputs.time_derivative.to_owned(),
covariates: inputs.covariates.to_owned(),
},
offset_eta_entry,
offset_eta_exit,
offset_derivative_exit,
penalties,
monotonicity,
monotonicity_constraint_rows,
monotonicity_constraint_offsets,
))
}
pub fn set_penalty_lambdas(&mut self, lambdas: &[f64]) -> Result<(), EstimationError> {
if lambdas.len() != self.penalties.blocks.len() {
crate::bail_invalid_estim!(
"set_penalty_lambdas expects {} lambdas, got {}",
self.penalties.blocks.len(),
lambdas.len()
);
}
for (block, &lambda) in self.penalties.blocks.iter_mut().zip(lambdas.iter()) {
if !lambda.is_finite() || lambda < 0.0 {
crate::bail_invalid_estim!("penalty lambda must be finite and >= 0, got {lambda}");
}
block.lambda = lambda;
}
Ok(())
}
pub fn set_structural_monotonicity(
&mut self,
enabled: bool,
time_columns: usize,
) -> Result<(), EstimationError> {
let p = self.coefficient_dim();
if time_columns > p {
crate::bail_invalid_estim!(
"structural time columns {} exceed coefficient dimension {}",
time_columns,
p
);
}
if enabled && time_columns == 0 {
crate::bail_invalid_estim!("structural monotonicity requires at least one time column");
}
if enabled {
const STRUCTURAL_DERIV_TOL: f64 = 1e-12;
for (i, &offset) in self.offset_derivative_exit.iter().enumerate() {
if offset < -STRUCTURAL_DERIV_TOL {
crate::bail_invalid_estim!(
"structural monotonicity requires nonnegative derivative offsets; found offset_derivative_exit[{i}]={offset:.3e}"
);
}
}
let mut derivative_row = vec![0.0_f64; p];
for i in 0..self.nrows() {
self.fill_derivative_row(i, &mut derivative_row);
for j in 0..time_columns {
let v = derivative_row[j];
if v < -STRUCTURAL_DERIV_TOL {
crate::bail_invalid_estim!(
"structural monotonicity requires nonnegative time-derivative basis entries; found x_derivative[{i},{j}]={v:.3e}"
);
}
}
for j in time_columns..p {
let v = derivative_row[j];
if v.abs() > STRUCTURAL_DERIV_TOL {
crate::bail_invalid_estim!(
"structural monotonicity requires zero derivative contribution outside the time block; found x_derivative[{i},{j}]={v:.3e}"
);
}
}
}
if let (Some(rows), Some(offsets)) = (
self.monotonicity_constraint_rows.as_ref(),
self.monotonicity_constraint_offsets.as_ref(),
) {
for (i, &offset) in offsets.iter().enumerate() {
if offset < -STRUCTURAL_DERIV_TOL {
crate::bail_invalid_estim!(
"structural monotonicity requires nonnegative collocation derivative offsets; found monotonicity_constraint_offsets[{i}]={offset:.3e}"
);
}
}
for i in 0..rows.nrows() {
for j in 0..time_columns {
let v = rows[[i, j]];
if v < -STRUCTURAL_DERIV_TOL {
crate::bail_invalid_estim!(
"structural monotonicity requires nonnegative collocation derivative basis entries; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
);
}
}
for j in time_columns..p {
let v = rows[[i, j]];
if v.abs() > STRUCTURAL_DERIV_TOL {
crate::bail_invalid_estim!(
"structural monotonicity requires zero collocation derivative contribution outside the time block; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
);
}
}
}
}
}
self.structurally_monotonic = enabled;
self.structural_time_columns = if enabled { time_columns } else { 0 };
Ok(())
}
pub fn update_state(&self, beta: &Array1<f64>) -> Result<WorkingState, EstimationError> {
if beta.len() != self.coefficient_dim() {
crate::bail_invalid_estim!("survival beta dimension mismatch");
}
let n = self.nrows();
let p = self.coefficient_dim();
let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
let mut nll = 0.0;
let derivative_guard = self.derivative_guard();
let derivative_guard_numerical = self.derivative_guard_numerical();
let mut workspace = self.workspace.lock().unwrap();
workspace.reset(n);
let SurvivalWorkspace {
w_event,
w_event_inv_deriv,
w_event_outer,
w_hess_exit,
w_hess_entry,
} = &mut *workspace;
for i in 0..n {
let w = self.sampleweight[i];
if w <= 0.0 {
continue;
}
let entry_age = self.age_entry[i];
let exit_age = self.age_exit[i];
if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
crate::bail_invalid_estim!(
"survival ages must be finite with age_exit >= age_entry"
);
}
let d = f64::from(self.event_target[i]);
let has_entry_interval = !self.entry_at_origin[i];
let interval_scale = if has_entry_interval {
eta_exit[i].max(eta_entry[i])
} else {
eta_exit[i]
};
let h_e_scaled = (eta_exit[i] - interval_scale).exp();
let h_s_scaled = if has_entry_interval {
(eta_entry[i] - interval_scale).exp()
} else {
0.0
};
let interval_scaled = h_e_scaled - h_s_scaled;
let interval = Self::scaled_exp_component(interval_scale, interval_scaled)?;
let (deriv, deriv_slope) = self
.stabilized_structural_derivative(derivative_raw[i])
.unwrap_or((derivative_raw[i], 1.0));
let mono_floor = if d > 0.0 {
derivative_guard_numerical
} else {
0.0
};
if !deriv.is_finite() || deriv < mono_floor {
return Err(EstimationError::ParameterConstraintViolation(format!(
"survival monotonicity violated at row {}: d_eta/dt={:.3e} <= tolerance={:.3e}",
i, deriv, derivative_guard
)));
}
if has_entry_interval {
let increment_guard = self.interval_increment_guard(h_s_scaled, h_e_scaled);
if interval_scaled + increment_guard < 0.0 {
return Err(EstimationError::ParameterConstraintViolation(format!(
"survival cumulative hazard decreased over row {}: H(exit)-H(entry)={:.6e}",
i, interval
)));
}
}
nll += w * interval;
let w_exit_i = w * eta_exit[i].exp();
let w_entry_i = if has_entry_interval {
w * eta_entry[i].exp()
} else {
0.0
};
if !w_exit_i.is_finite() {
crate::bail_invalid_estim!(
"survival interval term exceeds f64 range at row {i} (w*exp(eta_exit)={w_exit_i:.3e})"
);
}
w_hess_exit[i] = w_exit_i;
w_hess_entry[i] = w_entry_i;
if d > 0.0 {
let inv_deriv = deriv_slope / deriv;
nll += -w * (eta_exit[i] + deriv.ln());
w_event[i] = w;
w_event_inv_deriv[i] = w * inv_deriv;
w_event_outer[i] = w * inv_deriv * inv_deriv;
}
}
let mut h = self.interval_hessian_blas(w_hess_exit, w_hess_entry);
let mut grad = Array1::<f64>::zeros(p);
let mut grad_comp = Array1::<f64>::zeros(p);
let mut row_exit = vec![0.0_f64; p];
let mut row_entry = vec![0.0_f64; p];
let mut row_derivative = vec![0.0_f64; p];
for i in 0..n {
let w_interval_exit = w_hess_exit[i];
let w_interval_entry = w_hess_entry[i];
let w_event_exit = w_event[i];
let w_event_derivative = w_event_inv_deriv[i];
if w_interval_exit == 0.0
&& w_interval_entry == 0.0
&& w_event_exit == 0.0
&& w_event_derivative == 0.0
{
continue;
}
self.fill_exit_row(i, &mut row_exit);
self.fill_entry_row(i, &mut row_entry);
self.fill_derivative_row(i, &mut row_derivative);
for j in 0..p {
let contribution = w_interval_exit * row_exit[j]
- w_interval_entry * row_entry[j]
- w_event_exit * row_exit[j]
- w_event_derivative * row_derivative[j];
let t = grad[j] + contribution;
if grad[j].abs() >= contribution.abs() {
grad_comp[j] += (grad[j] - t) + contribution;
} else {
grad_comp[j] += (contribution - t) + grad[j];
}
grad[j] = t;
}
}
grad += &grad_comp;
h += &self.derivative_xt_diag_x(w_event_outer);
let score_norm = array1_l2_norm(&grad);
let penaltygrad = self.penalties.gradient(beta);
let penalty_quadratic_form = 2.0 * self.penalties.deviance(beta);
let penaltygrad_norm = array1_l2_norm(&penaltygrad);
let mut totalgrad = grad;
totalgrad += &penaltygrad;
self.penalties.addhessian_inplace(&mut h);
let log_likelihood = -nll;
let deviance = 2.0 * nll;
Ok(WorkingState {
eta: LinearPredictor::new(eta_exit),
gradient: totalgrad,
hessian: gam_linalg::matrix::SymmetricMatrix::Dense(h),
log_likelihood,
deviance,
penalty_term: penalty_quadratic_form,
firth: gam_solve::pirls::FirthDiagnostics::Inactive,
ridge_used: 0.0,
hessian_curvature: gam_solve::pirls::HessianCurvatureKind::Observed,
gradient_natural_scale: score_norm + penaltygrad_norm,
})
}
pub(crate) fn survival_hessian_derivative_correction(
&self,
beta: &Array1<f64>,
u_k: &Array1<f64>,
) -> Result<Array2<f64>, EstimationError> {
let p = beta.len();
let n = self.nrows();
let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
let deriv_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
let exp_entry = eta_entry.mapv(f64::exp);
let exp_exit = eta_exit.mapv(f64::exp);
let guard = self.derivative_guard();
let guard_numerical = self.derivative_guard_numerical();
let jac = Array1::<f64>::ones(p);
let curvature = Array1::<f64>::zeros(p);
let third = Array1::<f64>::zeros(p);
let mut row_exit = vec![0.0_f64; p];
let mut row_entry = vec![0.0_f64; p];
let mut row_derivative = vec![0.0_f64; p];
let mut ge = vec![0.0_f64; p];
let mut gs = vec![0.0_f64; p];
let mut gsd = vec![0.0_f64; p];
let mut he = vec![0.0_f64; p];
let mut hs = vec![0.0_f64; p];
let mut hsd = vec![0.0_f64; p];
let mut te = vec![0.0_f64; p];
let mut ts = vec![0.0_f64; p];
let mut tsd = vec![0.0_f64; p];
let mut b_dir = Array2::<f64>::zeros((p, p));
for i in 0..n {
let w_i = self.sampleweight[i];
if w_i <= 0.0 {
continue;
}
let has_entry = !self.entry_at_origin[i];
let mut deta_e = 0.0_f64;
let mut deta_s = 0.0_f64;
let mut ds = 0.0_f64;
self.fill_exit_row(i, &mut row_exit);
self.fill_entry_row(i, &mut row_entry);
self.fill_derivative_row(i, &mut row_derivative);
for j in 0..p {
ge[j] = row_exit[j] * jac[j];
gs[j] = row_entry[j] * jac[j];
gsd[j] = row_derivative[j] * jac[j];
he[j] = row_exit[j] * curvature[j];
hs[j] = row_entry[j] * curvature[j];
hsd[j] = row_derivative[j] * curvature[j];
te[j] = row_exit[j] * third[j];
ts[j] = row_entry[j] * third[j];
tsd[j] = row_derivative[j] * third[j];
deta_e += ge[j] * u_k[j];
if has_entry {
deta_s += gs[j] * u_k[j];
}
ds += gsd[j] * u_k[j];
}
for r in 0..p {
let dge_r = he[r] * u_k[r];
let dgs_r = hs[r] * u_k[r];
let dhe_r = te[r] * u_k[r];
let dhs_r = ts[r] * u_k[r];
for c in 0..p {
let dge_c = he[c] * u_k[c];
let dgs_c = hs[c] * u_k[c];
let mut d_h_rc =
exp_exit[i] * (deta_e * ge[r] * ge[c] + dge_r * ge[c] + ge[r] * dge_c);
if r == c {
d_h_rc += exp_exit[i] * (deta_e * he[r] + dhe_r);
}
if has_entry {
d_h_rc -=
exp_entry[i] * (deta_s * gs[r] * gs[c] + dgs_r * gs[c] + gs[r] * dgs_c);
if r == c {
d_h_rc -= exp_entry[i] * (deta_s * hs[r] + dhs_r);
}
}
b_dir[[r, c]] += w_i * d_h_rc;
}
}
let (s_i, s_slope) = self
.stabilized_structural_derivative(deriv_raw[i])
.unwrap_or((deriv_raw[i], 1.0));
if !s_i.is_finite() {
return Err(EstimationError::ParameterConstraintViolation(format!(
"survival monotonicity violated in unified trace contraction at row {i}: \
d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
)));
}
if self.event_target[i] > 0 && s_slope != 0.0 {
if s_i < guard_numerical {
return Err(EstimationError::ParameterConstraintViolation(format!(
"survival monotonicity violated in unified trace contraction at row {i}: \
d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
)));
}
let inv_s = 1.0 / s_i;
let inv_s2 = inv_s * inv_s;
let inv_s3 = inv_s2 * inv_s;
for r in 0..p {
let dgd_r = hsd[r] * u_k[r];
let dtsd_r = tsd[r] * u_k[r];
let dte_r = te[r] * u_k[r];
for c in 0..p {
let dgd_c = hsd[c] * u_k[c];
let mut d_h_rc = (dgd_r * gsd[c] + gsd[r] * dgd_c) * inv_s2
- 2.0 * gsd[r] * gsd[c] * ds * inv_s3;
if r == c {
d_h_rc += -dte_r;
d_h_rc += -(dtsd_r * inv_s - hsd[r] * ds * inv_s2);
}
b_dir[[r, c]] += w_i * d_h_rc;
}
}
}
}
Ok(b_dir)
}
pub fn offset_channel_residuals(
&self,
beta: &Array1<f64>,
) -> Result<OffsetChannelResiduals, EstimationError> {
if beta.len() != self.coefficient_dim() {
crate::bail_invalid_estim!(
"survival beta dimension mismatch in offset_channel_residuals"
);
}
let n = self.nrows();
let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
let derivative_guard_numerical = self.derivative_guard_numerical();
let mut r_exit = Array1::<f64>::zeros(n);
let mut r_entry = Array1::<f64>::zeros(n);
let mut r_deriv = Array1::<f64>::zeros(n);
for i in 0..n {
let w = self.sampleweight[i];
if w <= 0.0 {
continue;
}
let entry_age = self.age_entry[i];
let exit_age = self.age_exit[i];
if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
crate::bail_invalid_estim!(
"survival ages must be finite with age_exit >= age_entry"
);
}
let has_entry_interval = !self.entry_at_origin[i];
let d = f64::from(self.event_target[i]);
let w_exit_i = w * eta_exit[i].exp();
let w_entry_i = if has_entry_interval {
w * eta_entry[i].exp()
} else {
0.0
};
if !w_exit_i.is_finite() {
crate::bail_invalid_estim!(
"offset_channel_residuals: w*exp(eta_exit)={w_exit_i:.3e} non-finite at row {i}"
);
}
r_exit[i] = w_exit_i - d * w;
r_entry[i] = -w_entry_i;
let deriv_raw = derivative_raw[i];
let (deriv, deriv_slope) = self
.stabilized_structural_derivative(deriv_raw)
.unwrap_or((deriv_raw, 1.0));
let mono_floor = if d > 0.0 {
derivative_guard_numerical
} else {
0.0
};
if !deriv.is_finite() || deriv < mono_floor {
return Err(EstimationError::ParameterConstraintViolation(format!(
"offset_channel_residuals: derivative ≤ numerical guard at row {i}: {deriv:.3e}"
)));
}
if d > 0.0 {
r_deriv[i] = -w * d * deriv_slope / deriv;
}
}
let right = Array1::<f64>::zeros(r_exit.len());
Ok(OffsetChannelResiduals {
exit: r_exit,
entry: r_entry,
derivative: r_deriv,
right,
})
}
pub fn unified_lamlobjective_and_rhogradient(
&self,
beta: &Array1<f64>,
state: &WorkingState,
rho: &Array1<f64>,
) -> Result<(f64, Array1<f64>), EstimationError> {
use gam_problem::{EvalMode, PseudoLogdetMode};
use gam_solve::estimate::reml::assembly::InnerAssembly;
use gam_solve::estimate::reml::reml_outer_engine::{
DenseSpectralOperator, DispersionHandling,
};
use gam_solve::estimate::reml::reparameterized_inner::{
RawInnerReparamContext, assemble_reparameterized_inner,
};
use gam_terms::construction::{
canonicalize_penalty_specs, precompute_reparam_invariant_from_canonical,
stable_reparameterizationwith_invariant,
};
use gam_terms::penalty_spec::PenaltySpec;
let p = beta.len();
let active_penalty_blocks: Vec<&PenaltyBlock> = self
.penalties
.blocks
.iter()
.filter(|b| b.lambda > 0.0)
.collect();
if rho.len() != active_penalty_blocks.len() {
crate::bail_invalid_estim!(
"survival LAML rho dimension {} does not match active penalty block count {}",
rho.len(),
active_penalty_blocks.len()
);
}
let k_count = active_penalty_blocks.len();
let relative_projected_norm = {
let raw = state.gradient.clone();
let projected = match self.monotonicity_linear_constraints() {
Some(constraints) => {
let constraints = ConstraintSet::Dense(constraints);
projected_linear_constraint_stationarity_vector(&raw, beta, &constraints, None)
.ok_or_else(|| {
EstimationError::InvalidInput(
"survival LAML could not project the monotonicity KKT residual"
.to_string(),
)
})?
}
None => raw,
};
state.relative_gradient_norm(array1_l2_norm(&projected))
};
if !relative_projected_norm.is_finite()
|| relative_projected_norm > SURVIVAL_LAML_STATIONARITY_RELATIVE_TOL
{
return Err(EstimationError::InvalidInput(format!(
"survival LAML requires a stationary inner mode: projected relative KKT \
residual {relative_projected_norm:.3e} exceeds \
{SURVIVAL_LAML_STATIONARITY_RELATIVE_TOL:.3e}; a one-step residual \
surrogate is not a differentiable substitute for the Laplace mode"
)));
}
let lambdas: Vec<f64> = rho.iter().map(|&r| r.exp()).collect();
let h_dense = state.hessian.to_dense();
let hessian_logdet_mode = PseudoLogdetMode::PositiveDefinite;
let s_k_embedded: Vec<Array2<f64>> = active_penalty_blocks
.iter()
.map(|b| {
let mut s = Array2::<f64>::zeros((p, p));
let (rs, re) = (b.range.start, b.range.end);
s.slice_mut(ndarray::s![rs..re, rs..re]).assign(&b.matrix);
s
})
.collect();
let penalty_specs: Vec<PenaltySpec> = active_penalty_blocks
.iter()
.map(|b| PenaltySpec::Block {
local: b.matrix.clone(),
col_range: b.range.clone(),
prior_mean: gam_problem::CoefficientPriorMean::Zero,
structure_hint: None,
op: None,
})
.collect();
let nullspace_dims: Vec<usize> = active_penalty_blocks
.iter()
.map(|b| b.nullspace_dim)
.collect();
let (canonical_penalties, _canonical_nullspace) = canonicalize_penalty_specs(
&penalty_specs,
&nullspace_dims,
p,
"survival LAML seam-A reparameterization",
)
.map_err(|e| EstimationError::InvalidInput(e.to_string()))?;
if canonical_penalties.len() != k_count {
return Err(EstimationError::InvalidInput(format!(
"survival LAML reparameterization dropped {} of {} active (λ>0) penalty \
block(s) as numerically rank-0; cannot align transformed penalty \
coordinates with ρ",
k_count - canonical_penalties.len(),
k_count
)));
}
let reparam_invariant =
precompute_reparam_invariant_from_canonical(&canonical_penalties, p)
.map_err(|e| EstimationError::InvalidInput(e.to_string()))?;
let reparam = stable_reparameterizationwith_invariant(
&canonical_penalties,
&lambdas,
p,
&reparam_invariant,
None,
)
.map_err(|e| EstimationError::InvalidInput(e.to_string()))?;
let provider = SurvivalDerivProvider::new(self.clone(), beta.clone());
let ctx = RawInnerReparamContext {
hessian: &h_dense,
beta,
penalties_embedded: &s_k_embedded,
lambdas: &lambdas,
};
let reparam_inner = assemble_reparameterized_inner(
&ctx,
Some(Box::new(provider)),
&reparam,
)
.map_err(EstimationError::InvalidInput)?;
let hop = DenseSpectralOperator::from_symmetric_with_mode(
&reparam_inner.hessian_transformed,
hessian_logdet_mode,
)
.map_err(EstimationError::InvalidInput)?;
let penalty_coords = reparam
.canonical_transformed
.iter()
.map(|cp| cp.to_penalty_coordinate())
.collect::<Vec<_>>();
let penalty_quadratic = state.penalty_term;
let result = InnerAssembly {
log_likelihood: state.log_likelihood,
penalty_quadratic,
beta: reparam_inner.beta_transformed,
n_observations: self.nrows(),
hessian_op: std::sync::Arc::new(hop),
penalty_coords,
penalty_logdet: reparam_inner.penalty_logdet,
dispersion: DispersionHandling::Fixed {
phi: 1.0,
include_logdet_h: true,
include_logdet_s: true,
},
rho_curvature_scale: 1.0,
rho_prior: gam_problem::RhoPrior::Flat,
hessian_logdet_correction: 0.0,
penalty_subspace_trace: None,
deriv_provider: reparam_inner.deriv_provider,
firth: None,
nullspace_dim: None,
barrier_config: None,
ext_coords: Vec::new(),
ext_coord_pair_fn: None,
rho_ext_pair_fn: None,
fixed_drift_deriv: None,
contracted_psi_second_order: None,
kkt_residual: None,
active_constraints: None,
}
.evaluate(
rho.as_slice().expect("rho must be contiguous"),
EvalMode::ValueAndGradient,
None,
)
.map_err(EstimationError::InvalidInput)?;
let gradient = result.gradient.unwrap_or_else(|| Array1::zeros(rho.len()));
Ok((result.cost, gradient))
}
pub fn evaluate_survival_lamlcost_and_gradient(
&self,
rho: &[f64],
beta0: &Array1<f64>,
) -> Result<(f64, Array1<f64>), EstimationError> {
let (candidate, beta) = self.reconverge_survival_inner_mode(rho, beta0)?;
let rho_arr = Array1::from_vec(rho.to_vec());
let state = candidate.update_state(&beta)?;
candidate.unified_lamlobjective_and_rhogradient(&beta, &state, &rho_arr)
}
fn reconverge_survival_inner_mode(
&self,
rho: &[f64],
beta0: &Array1<f64>,
) -> Result<(WorkingModelSurvival, Array1<f64>), EstimationError> {
const SHIM_PIRLS_MAX_ITERATIONS: usize = 600;
const SHIM_PIRLS_CONVERGENCE_TOL: f64 = 1e-12;
const SHIM_PIRLS_MAX_STEP_HALVING: usize = 40;
const SHIM_PIRLS_MIN_STEP_SIZE: f64 = 1e-12;
let active_block_count = self
.penalties
.blocks
.iter()
.filter(|b| b.lambda > 0.0)
.count();
if rho.len() != active_block_count {
crate::bail_invalid_estim!(
"reconverge_survival_inner_mode: rho dimension {} does not match active penalty block count {}",
rho.len(),
active_block_count
);
}
if beta0.len() != self.coefficient_dim() {
crate::bail_invalid_estim!(
"reconverge_survival_inner_mode: beta0 dimension {} does not match coefficient dimension {}",
beta0.len(),
self.coefficient_dim()
);
}
let active_lambdas = gam_problem::checked_exp_log_strengths(rho.iter().copied())?;
let mut candidate = self.clone();
let mut lambdas: Vec<f64> = candidate
.penalties
.blocks
.iter()
.map(|b| b.lambda)
.collect();
let mut active_idx = 0usize;
for (block, lambda) in candidate.penalties.blocks.iter().zip(lambdas.iter_mut()) {
if block.lambda > 0.0 {
*lambda = active_lambdas[active_idx];
active_idx += 1;
}
}
candidate.set_penalty_lambdas(&lambdas)?;
let opts = gam_solve::pirls::WorkingModelPirlsOptions {
max_iterations: SHIM_PIRLS_MAX_ITERATIONS,
convergence_tolerance: SHIM_PIRLS_CONVERGENCE_TOL,
adaptive_kkt_tolerance: None,
max_step_halving: SHIM_PIRLS_MAX_STEP_HALVING,
min_step_size: SHIM_PIRLS_MIN_STEP_SIZE,
firth_bias_reduction: false,
coefficient_lower_bounds: None,
linear_constraints: None,
initial_lm_lambda: None,
arrow_schur: None,
};
let summary = gam_solve::pirls::runworking_model_pirls(
&mut candidate,
Coefficients::new(beta0.clone()),
&opts,
|_| {},
)?;
let mut beta = summary.beta.as_ref().to_owned();
{
const POLISH_MAX_ITERS: usize = 400;
const POLISH_TOL: f64 = 1e-13;
const ARMIJO_C: f64 = constants::ARMIJO_C1;
const BACKTRACK: f64 = constants::BACKTRACK_CONTRACTION;
const MAX_BACKTRACK: usize = 80;
let p = beta.len();
let penalized_objective =
|st: &WorkingState| -> f64 { -st.log_likelihood + 0.5 * st.penalty_term };
for _ in 0..POLISH_MAX_ITERS {
let st = match candidate.update_state(&beta) {
Ok(st) => st,
Err(_) => break,
};
let r = st.gradient.clone();
let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
if !r_norm.is_finite() || r_norm < POLISH_TOL {
break;
}
let h = st.hessian.to_dense();
let f0 = penalized_objective(&st);
let h_scale = (0..p)
.map(|d| h[[d, d]].abs())
.fold(0.0_f64, f64::max)
.max(1.0);
let try_lm = |lambda_lm: f64| -> Option<(Array1<f64>, f64)> {
let mut h_reg = h.clone();
for d in 0..p {
h_reg[[d, d]] += lambda_lm;
}
let factor =
gam_linalg::faer_ndarray::FaerCholesky::cholesky(&h_reg, faer::Side::Lower)
.ok()?;
let candidate_step = factor.solvevec(&r);
if candidate_step.iter().any(|v| !v.is_finite()) {
return None;
}
let dd = -r.dot(&candidate_step);
(dd.is_finite() && dd < -1e-14 * r_norm * r_norm)
.then_some((candidate_step, dd))
};
let (step, dir_deriv) = try_lm(0.0)
.or_else(|| {
escalate_ridge(RidgeSchedule::geometric(1e-11 * h_scale, 17), try_lm)
.ok()
.map(|success| success.value)
})
.unwrap_or_else(|| {
(r.clone(), -r_norm * r_norm)
});
let accepted = match backtracking_line_search::<_, Infallible>(
BacktrackConfig {
contraction: BACKTRACK,
max_steps: MAX_BACKTRACK,
..BacktrackConfig::default()
},
|alpha| {
let trial = &beta - &(alpha * &step);
let Ok(ts) = candidate.update_state(&trial) else {
return Ok(None);
};
let ft = penalized_objective(&ts);
let tn = ts.gradient.iter().map(|v| v * v).sum::<f64>().sqrt();
let armijo_ok = ft.is_finite() && ft <= f0 + ARMIJO_C * alpha * dir_deriv;
let residual_ok = tn.is_finite() && tn < r_norm;
Ok((armijo_ok || residual_ok).then_some((ft, trial)))
},
|_alpha, _ft| true,
) {
Ok(result) => result,
Err(never) => match never {},
};
let Some(ls) = accepted else {
break;
};
beta = ls.payload;
}
}
Ok((candidate, beta))
}
}
pub(crate) struct SurvivalDerivProvider {
model: WorkingModelSurvival,
beta: Array1<f64>,
}
impl SurvivalDerivProvider {
pub(crate) fn new(model: WorkingModelSurvival, beta: Array1<f64>) -> Self {
Self { model, beta }
}
}
impl gam_solve::estimate::reml::reml_outer_engine::HessianDerivativeProvider
for SurvivalDerivProvider
{
fn hessian_derivative_correction(
&self,
v_k: &Array1<f64>,
) -> Result<Option<Array2<f64>>, String> {
let u_k = -v_k;
match self
.model
.survival_hessian_derivative_correction(&self.beta, &u_k)
{
Ok(correction) => Ok(Some(correction)),
Err(e) => Err(e.to_string()),
}
}
fn has_corrections(&self) -> bool {
true
}
}
#[derive(Debug, Clone)]
pub struct CrudeRiskResult {
pub risk: f64,
pub diseasegradient: Array1<f64>,
pub mortalitygradient: Array1<f64>,
}
#[derive(Debug, Clone)]
pub struct CompetingRisksCifResult {
pub cif: Vec<Array2<f64>>,
pub overall_survival: Array2<f64>,
}
const COMPETING_RISKS_CIF_PARALLEL_ROW_MIN: usize = 256;
pub fn assemble_competing_risks_cif(
times: ArrayView1<'_, f64>,
cumulative_hazard: ArrayView3<'_, f64>,
) -> Result<CompetingRisksCifResult, SurvivalError> {
let (n_endpoints, n_rows, n_times) = cumulative_hazard.dim();
if n_endpoints == 0 {
return Err(SurvivalError::DimensionMismatch);
}
let endpoint_hazards = cumulative_hazard
.axis_iter(Axis(0))
.map(|view| view.to_owned())
.collect::<Vec<_>>();
assemble_competing_risks_cif_from_endpoints(times, &endpoint_hazards).and_then(|result| {
if result.overall_survival.dim() != (n_rows, n_times) {
Err(SurvivalError::DimensionMismatch)
} else {
Ok(result)
}
})
}
pub fn assemble_competing_risks_cif_from_endpoints(
times: ArrayView1<'_, f64>,
cumulative_hazards: &[Array2<f64>],
) -> Result<CompetingRisksCifResult, SurvivalError> {
let n_endpoints = cumulative_hazards.len();
if n_endpoints == 0 || times.is_empty() {
return Err(SurvivalError::DimensionMismatch);
}
let (n_rows, n_times) = cumulative_hazards[0].dim();
if n_rows == 0 || n_times == 0 || times.len() != n_times {
return Err(SurvivalError::DimensionMismatch);
}
if times.iter().any(|time| !time.is_finite() || *time < 0.0) {
return Err(SurvivalError::InvalidTimeGrid);
}
if times
.iter()
.zip(times.iter().skip(1))
.any(|(previous, current)| current <= previous)
{
return Err(SurvivalError::InvalidTimeGrid);
}
for endpoint_hazard in cumulative_hazards {
if endpoint_hazard.dim() != (n_rows, n_times) {
return Err(SurvivalError::DimensionMismatch);
}
if endpoint_hazard.iter().any(|value| !value.is_finite()) {
return Err(SurvivalError::NonFiniteInput);
}
}
let max_abs_hazard = cumulative_hazards
.iter()
.flat_map(|endpoint_hazard| endpoint_hazard.iter())
.fold(0.0_f64, |acc, value| acc.max(value.abs()));
let monotone_tolerance = 1.0e-10_f64 * max_abs_hazard.max(1.0);
let mut cif: Vec<Array2<f64>> = (0..n_endpoints)
.map(|_| Array2::<f64>::zeros((n_rows, n_times)))
.collect();
let mut overall_survival = Array2::<f64>::zeros((n_rows, n_times));
let assemble_row = |row: usize| -> Result<(Vec<f64>, Vec<f64>), SurvivalError> {
let mut cif_flat = vec![0.0_f64; n_endpoints * n_times];
let mut surv_row = vec![0.0_f64; n_times];
let mut previous_cif = vec![0.0_f64; n_endpoints];
let mut previous_cumulative = vec![0.0_f64; n_endpoints];
let mut increments = vec![0.0_f64; n_endpoints];
let mut previous_total_cumulative = 0.0_f64;
for time_idx in 0..n_times {
let mut total_increment = 0.0_f64;
for endpoint in 0..n_endpoints {
let current = cumulative_hazards[endpoint][[row, time_idx]];
if current < -monotone_tolerance {
return Err(SurvivalError::NonMonotoneCumulativeHazard);
}
let raw_increment = current - previous_cumulative[endpoint];
if raw_increment < -monotone_tolerance {
return Err(SurvivalError::NonMonotoneCumulativeHazard);
}
let increment = raw_increment.max(0.0);
increments[endpoint] = increment;
total_increment += increment;
previous_cumulative[endpoint] += increment;
}
let survival_left = (-previous_total_cumulative).exp();
let interval_failure = -(-total_increment).exp_m1();
for endpoint in 0..n_endpoints {
if total_increment > 0.0 {
previous_cif[endpoint] +=
survival_left * interval_failure * increments[endpoint] / total_increment;
}
cif_flat[endpoint * n_times + time_idx] = previous_cif[endpoint].clamp(0.0, 1.0);
}
previous_total_cumulative += total_increment;
let mut fsum_at_t = 0.0_f64;
for endpoint in 0..n_endpoints {
fsum_at_t += cif_flat[endpoint * n_times + time_idx];
}
surv_row[time_idx] = (1.0_f64 - fsum_at_t).clamp(0.0, 1.0);
}
Ok((cif_flat, surv_row))
};
let rows: Vec<(Vec<f64>, Vec<f64>)> = if n_rows >= COMPETING_RISKS_CIF_PARALLEL_ROW_MIN
&& rayon::current_thread_index().is_none()
{
use rayon::prelude::*;
(0..n_rows)
.into_par_iter()
.map(assemble_row)
.collect::<Result<_, _>>()?
} else {
(0..n_rows).map(assemble_row).collect::<Result<_, _>>()?
};
for (row, (cif_flat, surv_row)) in rows.into_iter().enumerate() {
for endpoint in 0..n_endpoints {
for time_idx in 0..n_times {
cif[endpoint][[row, time_idx]] = cif_flat[endpoint * n_times + time_idx];
}
}
for time_idx in 0..n_times {
overall_survival[[row, time_idx]] = surv_row[time_idx];
}
}
Ok(CompetingRisksCifResult {
cif,
overall_survival,
})
}
fn compute_gauss_legendre_nodes(n: usize) -> Vec<(f64, f64)> {
let (nodes, weights) = gam_math::special::gauss_legendre(n);
nodes.into_iter().zip(weights).collect()
}
fn gauss_legendre_quadrature() -> &'static [(f64, f64)] {
static CACHE: LazyLock<Vec<(f64, f64)>> = LazyLock::new(|| compute_gauss_legendre_nodes(40));
&CACHE
}
pub fn calculate_crude_risk_quadrature<F>(
t0: f64,
t1: f64,
breakpoints: &[f64],
h_dis_t0: f64,
h_mor_t0: f64,
design_d_t0: ArrayView1<'_, f64>,
design_m_t0: ArrayView1<'_, f64>,
mut eval_at: F,
) -> Result<CrudeRiskResult, SurvivalError>
where
F: FnMut(
f64,
&mut Array1<f64>,
&mut Array1<f64>,
&mut Array1<f64>,
) -> Result<(f64, f64, f64), SurvivalError>,
{
let coeff_len_d = design_d_t0.len();
let coeff_len_m = design_m_t0.len();
if coeff_len_d == 0 || coeff_len_m == 0 {
return Err(SurvivalError::InvalidIntegrationSetup);
}
if !t0.is_finite()
|| !t1.is_finite()
|| !h_dis_t0.is_finite()
|| !h_mor_t0.is_finite()
|| design_d_t0.iter().any(|v| !v.is_finite())
|| design_m_t0.iter().any(|v| !v.is_finite())
{
return Err(SurvivalError::NonFiniteInput);
}
if t1 <= t0 {
return Ok(CrudeRiskResult {
risk: 0.0,
diseasegradient: Array1::zeros(coeff_len_d),
mortalitygradient: Array1::zeros(coeff_len_m),
});
}
let mut sorted_breaks: Vec<f64> = breakpoints
.iter()
.copied()
.filter(|x| x.is_finite() && *x >= t0 && *x <= t1)
.collect();
sorted_breaks.push(t0);
sorted_breaks.push(t1);
sorted_breaks.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
sorted_breaks.dedup_by(|a, b| (*a - *b).abs() < 1e-6);
if sorted_breaks.len() < 2 {
return Err(SurvivalError::InvalidIntegrationSetup);
}
let mut total_risk = 0.0;
let mut diseasegradient = Array1::zeros(coeff_len_d);
let mut mortalitygradient = Array1::zeros(coeff_len_m);
let nodesweights = gauss_legendre_quadrature();
let mut design_d = Array1::<f64>::zeros(coeff_len_d);
let mut deriv_d = Array1::<f64>::zeros(coeff_len_d);
let mut design_m = Array1::<f64>::zeros(coeff_len_m);
for segment in sorted_breaks.windows(2) {
let a = segment[0];
let b = segment[1];
let center = 0.5 * (b + a);
let halfwidth = 0.5 * (b - a);
if halfwidth <= 0.0 {
continue;
}
for &(x, w) in nodesweights {
let u = center + halfwidth * x;
let (inst_hazard_d, hazard_d, hazard_m) =
eval_at(u, &mut design_d, &mut deriv_d, &mut design_m)?;
if !inst_hazard_d.is_finite() || !hazard_d.is_finite() || !hazard_m.is_finite() {
return Err(SurvivalError::NonFiniteInput);
}
if inst_hazard_d <= 0.0 {
return Err(SurvivalError::NonPositiveHazard);
}
if hazard_d < h_dis_t0 || hazard_m < h_mor_t0 {
return Err(SurvivalError::NonMonotoneCumulativeHazard);
}
let h_dis_cond = hazard_d - h_dis_t0;
let h_mor_cond = hazard_m - h_mor_t0;
let s_total = (-(h_dis_cond + h_mor_cond)).exp();
total_risk += w * inst_hazard_d * s_total * halfwidth;
let weight = w * s_total * halfwidth;
for j in 0..coeff_len_d {
let d_inst_hazard = inst_hazard_d * design_d[j] + hazard_d * deriv_d[j];
let d_hazard_cond = hazard_d * design_d[j] - h_dis_t0 * design_d_t0[j];
let g = d_inst_hazard - inst_hazard_d * d_hazard_cond;
diseasegradient[j] += weight * g;
}
let weight = w * inst_hazard_d * s_total * halfwidth;
for j in 0..coeff_len_m {
let g = -hazard_m * design_m[j] + h_mor_t0 * design_m_t0[j];
mortalitygradient[j] += weight * g;
}
}
}
Ok(CrudeRiskResult {
risk: total_risk,
diseasegradient,
mortalitygradient,
})
}
impl PirlsWorkingModel for WorkingModelSurvival {
fn update(&mut self, beta: &Coefficients) -> Result<WorkingState, EstimationError> {
self.update_state(beta)
}
}
#[cfg(test)]
mod tests {
use super::*;
use ndarray::{Array1, Array2, Array3, array, s};
#[test]
fn saved_cause_specific_alo_matches_independent_closed_form() {
let eta_exit = 0.4_f64;
let eta_entry = -0.3_f64;
let derivative_exit = 1.7_f64;
let weight = 2.2_f64;
let geometry = cause_specific_survival_alo_row_geometry(CauseSpecificSurvivalAloRowInput {
eta_exit,
eta_entry,
derivative_exit,
prior_weight: weight,
entry_active: true,
event: true,
})
.expect("valid cause-specific row");
let expected_nll =
weight * (eta_exit.exp() - eta_entry.exp() - eta_exit - derivative_exit.ln());
let expected_score = [
weight * (eta_exit.exp() - 1.0),
-weight * eta_entry.exp(),
-weight / derivative_exit,
];
let expected_hessian = [
[weight * eta_exit.exp(), 0.0, 0.0],
[0.0, -weight * eta_entry.exp(), 0.0],
[0.0, 0.0, weight / derivative_exit.powi(2)],
];
assert!((geometry.negative_log_likelihood - expected_nll).abs() <= 2.0e-14);
for row in 0..3 {
assert!((geometry.nll_score[row] - expected_score[row]).abs() <= 2.0e-14);
for column in 0..3 {
assert!(
(geometry.observed_hessian[row][column] - expected_hessian[row][column]).abs()
<= 2.0e-14
);
}
}
let score_meat = geometry.nll_score[0] * geometry.nll_score[0];
assert!(
(geometry.observed_hessian[0][0] - score_meat).abs() > 1.0e-2,
"survival observed W and empirical score meat C must remain separate"
);
}
mod jet_cause_specific_production_parity {
use super::*;
use gam_math::jet_tower::{
program_fourth_contracted, program_row_kernel, program_third_contracted,
};
fn identity_block(w: f64, has_entry: bool, event: bool) -> CauseSpecificRoystonParmarBlock {
let age_entry = if has_entry { 1.0 } else { 0.0 };
CauseSpecificRoystonParmarBlock {
age_entry: array![age_entry],
age_exit: array![2.0],
event_target: array![if event { 1u8 } else { 0u8 }],
sampleweight: array![w],
x_entry: array![[0.0, 1.0, 0.0]],
x_exit: array![[1.0, 0.0, 0.0]],
x_derivative: array![[0.0, 0.0, 1.0]],
offset_eta_entry: array![0.0],
offset_eta_exit: array![0.0],
offset_derivative_exit: array![0.0],
derivative_floor: 0.0,
structural_time_columns: 0,
}
}
fn close(hand: f64, jet: f64, tol: f64, label: &str) {
let band = tol + tol * hand.abs().max(jet.abs());
assert!(
(hand - jet).abs() <= band,
"{label}: hand {hand:+.15e} vs jet {jet:+.15e} (|Δ|={:.3e} band {band:.3e})",
(hand - jet).abs()
);
}
const JET_TOL: f64 = 1e-9;
fn run_corner(has_entry: bool, event: bool) {
let beta = array![0.4_f64, -0.3_f64, 1.3_f64];
let d_beta = array![0.7_f64, -0.5_f64, 0.6_f64];
let v_beta = array![-0.2_f64, 0.8_f64, -0.4_f64];
let w = 1.4_f64;
let block = identity_block(w, has_entry, event);
let prog = crate::survival::CauseSpecificRowProgram::new(
[beta[0], beta[1], beta[2]],
w,
has_entry,
event,
);
let label = format!("entry={has_entry} event={event}");
let (ll, grad, hess) =
evaluate_cause_specific_block(&block, &beta).expect("evaluate block");
let (jet_v, jet_g, jet_h) = program_row_kernel(&prog, 0).expect("jet kernel");
close(jet_v, -ll, JET_TOL, &format!("{label} value"));
for a in 0..3 {
close(jet_g[a], -grad[a], JET_TOL, &format!("{label} grad[{a}]"));
for b in 0..3 {
close(
jet_h[a][b],
hess[[a, b]],
JET_TOL,
&format!("{label} H[{a}][{b}]"),
);
}
}
let dh = cause_specific_hessian_directional_derivative(&block, &beta, &d_beta)
.expect("live third");
let dir = [d_beta[0], d_beta[1], d_beta[2]];
let jet_t3 = program_third_contracted(&prog, 0, &dir).expect("jet third");
for a in 0..3 {
for b in 0..3 {
close(
jet_t3[a][b],
dh[[a, b]],
JET_TOL,
&format!("{label} third[{a}][{b}]"),
);
}
}
let d2h = cause_specific_hessian_second_directional_derivative(
&block, &beta, &d_beta, &v_beta,
)
.expect("live fourth");
let uu = [d_beta[0], d_beta[1], d_beta[2]];
let vv = [v_beta[0], v_beta[1], v_beta[2]];
let jet_t4 = program_fourth_contracted(&prog, 0, &uu, &vv).expect("jet fourth");
for a in 0..3 {
for b in 0..3 {
close(
jet_t4[a][b],
d2h[[a, b]],
JET_TOL,
&format!("{label} fourth[{a}][{b}]"),
);
}
}
let h_fd = 1e-5;
let bp = &beta + &(&d_beta * h_fd);
let bm = &beta - &(&d_beta * h_fd);
let (_, _, hp) = evaluate_cause_specific_block(&block, &bp).expect("evaluate +");
let (_, _, hm) = evaluate_cause_specific_block(&block, &bm).expect("evaluate -");
for a in 0..3 {
for b in 0..3 {
let fd = (hp[[a, b]] - hm[[a, b]]) / (2.0 * h_fd);
close(dh[[a, b]], fd, 1e-5, &format!("{label} FD third[{a}][{b}]"));
}
}
let dhp = cause_specific_hessian_directional_derivative(
&block,
&bp_along(&beta, &v_beta, h_fd),
&d_beta,
)
.expect("live third +");
let dhm = cause_specific_hessian_directional_derivative(
&block,
&bm_along(&beta, &v_beta, h_fd),
&d_beta,
)
.expect("live third -");
for a in 0..3 {
for b in 0..3 {
let fd = (dhp[[a, b]] - dhm[[a, b]]) / (2.0 * h_fd);
close(
d2h[[a, b]],
fd,
1e-5,
&format!("{label} FD fourth[{a}][{b}]"),
);
}
}
}
fn bp_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
beta + &(v * h)
}
fn bm_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
beta - &(v * h)
}
#[test]
fn cause_specific_live_tower_matches_jet_and_fd() {
for &has_entry in &[false, true] {
for &event in &[false, true] {
run_corner(has_entry, event);
}
}
}
#[test]
fn release_measure_cause_specific_specialized_vs_generic_tower_932() {
use std::time::Instant;
const ROWS: usize = 512;
let mut rows: Vec<([f64; 3], f64, bool, bool)> = Vec::with_capacity(ROWS);
for idx in 0..ROWS {
let f = idx as f64;
let eta_exit = 1.6 * (f * 0.17 + 0.3).sin() - 0.4 * (f * 0.09).cos();
let eta_entry = 1.1 * (f * 0.13 + 0.7).cos() + 0.35 * (f * 0.05).sin();
let derivative = 0.5 + 0.45 * (f * 0.31 + 0.2).sin().abs();
let weight = 0.6 + 0.4 * (f * 0.07 + 1.0).sin().abs();
let entry_active = idx % 2 == 0;
let event = (idx / 2) % 2 == 0;
rows.push((
[eta_exit, eta_entry, derivative],
weight,
entry_active,
event,
));
}
let programs: Vec<crate::survival::CauseSpecificRowProgram> = rows
.iter()
.map(|&(primary, weight, entry_active, event)| {
crate::survival::CauseSpecificRowProgram::new(
primary,
weight,
entry_active,
event,
)
})
.collect();
let dir_u: Vec<[f64; 3]> = (0..ROWS)
.map(|idx| {
let f = idx as f64;
[
0.7 * (f * 0.23 + 0.4).cos() - 0.2 * (f * 0.03).sin(),
-0.6 * (f * 0.29 + 0.1).sin() + 0.25 * (f * 0.15).cos(),
0.5 * (f * 0.19 + 0.6).cos() - 0.3 * (f * 0.08).sin(),
]
})
.collect();
let dir_v: Vec<[f64; 3]> = (0..ROWS)
.map(|idx| {
let f = idx as f64;
[
-0.5 * (f * 0.21 + 0.9).sin() + 0.3 * (f * 0.06).cos(),
0.8 * (f * 0.27 + 0.5).cos() - 0.15 * (f * 0.04).sin(),
0.4 * (f * 0.13 + 0.3).sin() - 0.2 * (f * 0.11).cos(),
]
})
.collect();
for (idx, (row, program)) in rows.iter().zip(programs.iter()).enumerate() {
let (primary, weight, entry_active, event) = *row;
let atom = cause_specific_row_order2(
primary[0],
primary[1],
primary[2],
weight,
f64::from(entry_active),
f64::from(event),
);
let (tower_value, tower_gradient, tower_hessian) =
program_row_kernel(program, 0).expect("tower warm kernel");
close(
atom.value(),
tower_value,
JET_TOL,
"release-measure value parity",
);
let production_gradient = atom.gradient();
for a in 0..3 {
close(
production_gradient[a],
tower_gradient[a],
JET_TOL,
"release-measure gradient parity",
);
for b in 0..3 {
close(
atom.hessian_at(a, b),
tower_hessian[a][b],
JET_TOL,
"release-measure hessian parity",
);
}
}
let production_third = cause_specific_row_third_contracted(
primary[0],
primary[1],
primary[2],
weight,
f64::from(entry_active),
f64::from(event),
&dir_u[idx],
);
let tower_third =
program_third_contracted(program, 0, &dir_u[idx]).expect("tower warm third");
let production_fourth = cause_specific_row_fourth_contracted(
primary[0],
primary[1],
primary[2],
weight,
f64::from(entry_active),
f64::from(event),
&dir_u[idx],
&dir_v[idx],
);
let tower_fourth = program_fourth_contracted(program, 0, &dir_u[idx], &dir_v[idx])
.expect("tower warm fourth");
for a in 0..3 {
for b in 0..3 {
close(
production_third[a][b],
tower_third[a][b],
JET_TOL,
"release-measure third parity",
);
close(
production_fourth[a][b],
tower_fourth[a][b],
JET_TOL,
"release-measure fourth parity",
);
}
}
}
let best_secs = |sweep: &mut dyn FnMut() -> f64| -> f64 {
let mut best = f64::INFINITY;
for _ in 0..5 {
let started = Instant::now();
let checksum = sweep();
assert!(
checksum.is_finite(),
"cause-specific release-measure checksum must stay finite"
);
best = best.min(started.elapsed().as_secs_f64());
}
best
};
let mut production_sweep = || {
let mut checksum = 0.0_f64;
for &(primary, weight, entry_active, event) in &rows {
let atom = cause_specific_row_order2(
primary[0],
primary[1],
primary[2],
weight,
f64::from(entry_active),
f64::from(event),
);
checksum += atom.value() + atom.gradient()[0] + atom.hessian_at(0, 0);
}
checksum
};
let production_secs = best_secs(&mut production_sweep);
let mut tower_sweep = || {
let mut checksum = 0.0_f64;
for program in &programs {
let (value, gradient, hessian) =
program_row_kernel(program, 0).expect("tower kernel");
checksum += value + gradient[0] + hessian[0][0];
}
checksum
};
let tower_secs = best_secs(&mut tower_sweep);
let mut production_third_sweep = || {
let mut checksum = 0.0_f64;
for (idx, &(primary, weight, entry_active, event)) in rows.iter().enumerate() {
let third = cause_specific_row_third_contracted(
primary[0],
primary[1],
primary[2],
weight,
f64::from(entry_active),
f64::from(event),
&dir_u[idx],
);
checksum += third[0][0] + third[0][1] + third[1][1];
}
checksum
};
let production_third_secs = best_secs(&mut production_third_sweep);
let mut tower_third_sweep = || {
let mut checksum = 0.0_f64;
for (idx, program) in programs.iter().enumerate() {
let third = program_third_contracted(program, 0, &dir_u[idx])
.expect("tower third kernel");
checksum += third[0][0] + third[0][1] + third[1][1];
}
checksum
};
let tower_third_secs = best_secs(&mut tower_third_sweep);
let mut production_fourth_sweep = || {
let mut checksum = 0.0_f64;
for (idx, &(primary, weight, entry_active, event)) in rows.iter().enumerate() {
let fourth = cause_specific_row_fourth_contracted(
primary[0],
primary[1],
primary[2],
weight,
f64::from(entry_active),
f64::from(event),
&dir_u[idx],
&dir_v[idx],
);
checksum += fourth[0][0] + fourth[0][1] + fourth[1][1];
}
checksum
};
let production_fourth_secs = best_secs(&mut production_fourth_sweep);
let mut tower_fourth_sweep = || {
let mut checksum = 0.0_f64;
for (idx, program) in programs.iter().enumerate() {
let fourth = program_fourth_contracted(program, 0, &dir_u[idx], &dir_v[idx])
.expect("tower fourth kernel");
checksum += fourth[0][0] + fourth[0][1] + fourth[1][1];
}
checksum
};
let tower_fourth_secs = best_secs(&mut tower_fourth_sweep);
for (channel, production_secs, tower_secs) in [
("order2", production_secs, tower_secs),
("third", production_third_secs, tower_third_secs),
("fourth", production_fourth_secs, tower_fourth_secs),
] {
let production_ns = production_secs * 1e9 / ROWS as f64;
let tower_ns = tower_secs * 1e9 / ROWS as f64;
eprintln!(
"CAUSE-SPECIFIC-RELEASE-932 channel={channel} rows={ROWS} \
production_ns={production_ns:.3} generic_tower_ns={tower_ns:.3} \
hand_over_production={:.6}",
tower_ns / production_ns,
);
}
}
}
#[test]
fn competing_risks_cif_constant_hazard_matches_closed_form() {
let times = array![0.0, 2.0, 5.0, 10.0];
let disease_rates = [0.12, 0.06];
let death_rates = [0.05, 0.02];
let cumulative = Array3::from_shape_fn((2, 2, times.len()), |(endpoint, row, time_idx)| {
let rate = if endpoint == 0 {
disease_rates[row]
} else {
death_rates[row]
};
rate * times[time_idx]
});
let result =
assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
for row in 0..2 {
let total_rate = disease_rates[row] + death_rates[row];
for time_idx in 0..times.len() {
let failure = 1.0 - (-total_rate * times[time_idx]).exp();
let expected_disease = disease_rates[row] / total_rate * failure;
let expected_death = death_rates[row] / total_rate * failure;
assert!((result.cif[0][[row, time_idx]] - expected_disease).abs() < 1e-12);
assert!((result.cif[1][[row, time_idx]] - expected_death).abs() < 1e-12);
assert!(
(result.cif[0][[row, time_idx]]
+ result.cif[1][[row, time_idx]]
+ result.overall_survival[[row, time_idx]]
- 1.0)
.abs()
< 1e-12
);
}
}
}
#[test]
fn competing_risks_cif_rejects_nonmonotone_hazards() {
let times = array![0.0, 1.0, 2.0];
let cumulative = Array3::from_shape_vec((1, 1, 3), vec![0.0, 0.2, 0.1]).expect("shape");
let err = assemble_competing_risks_cif(times.view(), cumulative.view())
.expect_err("nonmonotone cumulative hazard should be rejected");
assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
}
#[test]
fn competing_risks_cif_plateaus_and_three_causes_conserve_probability() {
let times = array![0.0, 1.0, 3.0, 7.0, 12.0];
let cumulative = Array3::from_shape_vec(
(3, 2, 5),
vec![
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,
],
)
.expect("shape");
let result =
assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
for row in 0..2 {
for time_idx in 0..times.len() {
let total_cif = result.cif[0][[row, time_idx]]
+ result.cif[1][[row, time_idx]]
+ result.cif[2][[row, time_idx]];
assert!(
(total_cif + result.overall_survival[[row, time_idx]] - 1.0).abs() < 1e-12,
"probability mass mismatch at row={row}, time_idx={time_idx}"
);
assert!((0.0..=1.0).contains(&result.overall_survival[[row, time_idx]]));
for cause in 0..3 {
assert!((0.0..=1.0).contains(&result.cif[cause][[row, time_idx]]));
if time_idx > 0 {
assert!(
result.cif[cause][[row, time_idx]] + 1e-12
>= result.cif[cause][[row, time_idx - 1]],
"CIF decreased for cause={cause}, row={row}, time_idx={time_idx}"
);
}
}
}
}
assert_eq!(result.cif[0][[0, 1]], result.cif[0][[0, 2]]);
assert_eq!(result.cif[0][[1, 2]], result.cif[0][[1, 3]]);
assert_eq!(result.cif[2][[1, 2]], result.cif[2][[1, 3]]);
}
#[test]
fn competing_risks_cif_rejects_bad_time_grids_and_nonfinite_hazards() {
let cumulative = Array3::zeros((2, 1, 2));
for times in [array![0.0, 0.0], array![1.0, 0.5], array![-1.0, 1.0]] {
let err = assemble_competing_risks_cif(times.view(), cumulative.view())
.expect_err("bad time grid should be rejected");
assert!(matches!(err, SurvivalError::InvalidTimeGrid));
}
let times = array![0.0, 1.0];
let nonfinite = Array3::from_shape_vec((1, 1, 2), vec![0.0, f64::NAN]).expect("shape");
let err = assemble_competing_risks_cif(times.view(), nonfinite.view())
.expect_err("nonfinite hazard should be rejected");
assert!(matches!(err, SurvivalError::NonFiniteInput));
}
#[test]
fn competing_risks_cif_extreme_hazards_remain_bounded() {
let times = array![0.0, 1.0, 2.0];
let cumulative =
Array3::from_shape_vec((2, 1, 3), vec![0.0, 500.0, 1000.0, 0.0, 250.0, 1000.0])
.expect("shape");
let result =
assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
for value in result
.cif
.iter()
.flat_map(|m| m.iter())
.chain(result.overall_survival.iter())
{
assert!(value.is_finite());
assert!((0.0..=1.0).contains(value));
}
assert!((result.cif[0][[0, 2]] + result.cif[1][[0, 2]] - 1.0).abs() < 1e-12);
assert_eq!(result.overall_survival[[0, 2]], 0.0);
}
fn toy_penalties() -> PenaltyBlocks {
let s = array![[2.0, 0.5], [0.5, 3.0]];
PenaltyBlocks::new(vec![PenaltyBlock {
matrix: s,
lambda: 1.7,
range: 1..3,
nullspace_dim: 0,
}])
}
fn survival_inputs<'a>(
age_entry: &'a Array1<f64>,
age_exit: &'a Array1<f64>,
event_target: &'a Array1<u8>,
event_competing: &'a Array1<u8>,
sampleweight: &'a Array1<f64>,
x_entry: &'a Array2<f64>,
x_exit: &'a Array2<f64>,
x_derivative: &'a Array2<f64>,
) -> SurvivalEngineInputs<'a> {
SurvivalEngineInputs {
age_entry: age_entry.view(),
age_exit: age_exit.view(),
event_target: event_target.view(),
event_competing: event_competing.view(),
sampleweight: sampleweight.view(),
x_entry: x_entry.view(),
x_exit: x_exit.view(),
x_derivative: x_derivative.view(),
monotonicity_constraint_rows: None,
monotonicity_constraint_offsets: None,
}
}
fn survival_model(
inputs: SurvivalEngineInputs<'_>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
spec: SurvivalSpec,
) -> Result<WorkingModelSurvival, SurvivalError> {
WorkingModelSurvival::from_engine_inputs(inputs, penalties, monotonicity, spec)
}
fn survival_model_with_offsets(
inputs: SurvivalEngineInputs<'_>,
offsets: Option<SurvivalBaselineOffsets<'_>>,
penalties: PenaltyBlocks,
monotonicity: SurvivalMonotonicityPenalty,
spec: SurvivalSpec,
) -> Result<WorkingModelSurvival, SurvivalError> {
WorkingModelSurvival::from_engine_inputswith_offsets(
inputs,
offsets,
penalties,
monotonicity,
spec,
)
}
#[test]
fn penaltyhessian_matchesgradient_jacobian() {
let penalties = toy_penalties();
let beta = array![10.0, -0.3, 1.2, 7.0];
let grad = penalties.gradient(&beta);
let h = penalties.hessian(beta.len());
let b_block = beta.slice(s![1..3]).to_owned();
let expected = 1.7 * array![[2.0, 0.5], [0.5, 3.0]].dot(&b_block);
assert!((grad[1] - expected[0]).abs() < 1e-12);
assert!((grad[2] - expected[1]).abs() < 1e-12);
assert!((h[[1, 1]] - 1.7 * 2.0).abs() < 1e-12);
assert!((h[[1, 2]] - 1.7 * 0.5).abs() < 1e-12);
assert!((h[[2, 1]] - 1.7 * 0.5).abs() < 1e-12);
assert!((h[[2, 2]] - 1.7 * 3.0).abs() < 1e-12);
}
#[test]
fn penaltygradient_matches_deviance_finite_difference() {
let penalties = toy_penalties();
let beta = array![10.0, -0.3, 1.2, 7.0];
let grad = penalties.gradient(&beta);
let eps = 1e-7;
for idx in 0..beta.len() {
let mut plus = beta.clone();
let mut minus = beta.clone();
plus[idx] += eps;
minus[idx] -= eps;
let fd = (penalties.deviance(&plus) - penalties.deviance(&minus)) / (2.0 * eps);
assert_eq!(
grad[idx].signum(),
fd.signum(),
"gradient/deviance sign mismatch at idx={idx}: grad={} fd={fd}",
grad[idx]
);
assert!(
(grad[idx] - fd).abs() < 1e-6,
"gradient/deviance mismatch at idx={idx}: grad={} fd={fd}",
grad[idx]
);
}
}
#[test]
fn zero_offsets_match_default_survival_state() {
let age_entry = array![1.0_f64, 2.0_f64];
let age_exit = array![2.0_f64, 3.5_f64];
let event_target = array![1u8, 0u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let beta = array![-1.0, 0.8];
let base = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties.clone(),
mono,
SurvivalSpec::Net,
)
.expect("construct base survival model");
let zero_offsets = survival_model_with_offsets(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
Some(SurvivalBaselineOffsets {
eta_entry: array![0.0, 0.0].view(),
eta_exit: array![0.0, 0.0].view(),
derivative_exit: array![0.0, 0.0].view(),
}),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("construct offset survival model");
let state_base = base.update_state(&beta).expect("base state");
let statezero = zero_offsets.update_state(&beta).expect("zero-offset state");
assert!((state_base.deviance - statezero.deviance).abs() < 1e-12);
assert!(
state_base
.gradient
.iter()
.zip(statezero.gradient.iter())
.all(|(a, b)| (a - b).abs() < 1e-12)
);
}
#[test]
fn competing_risk_cause_labels_collapse_to_pooled_baseline_indicator() {
let age_entry = array![0.0_f64, 0.0, 0.0, 0.0];
let age_exit = array![1.2_f64, 0.8, 2.1, 1.5];
let cause_labels = array![0u8, 1u8, 2u8, 0u8];
let event_competing = Array1::<u8>::zeros(cause_labels.len());
let sampleweight = array![1.0_f64, 1.0, 1.0, 1.0];
let x_entry = array![
[1.0, age_entry[0].max(1e-8).ln()],
[1.0, age_entry[1].max(1e-8).ln()],
[1.0, age_entry[2].max(1e-8).ln()],
[1.0, age_entry[3].max(1e-8).ln()],
];
let x_exit = array![
[1.0, age_exit[0].ln()],
[1.0, age_exit[1].ln()],
[1.0, age_exit[2].ln()],
[1.0, age_exit[3].ln()],
];
let x_derivative = array![
[0.0, 1.0 / age_exit[0]],
[0.0, 1.0 / age_exit[1]],
[0.0, 1.0 / age_exit[2]],
[0.0, 1.0 / age_exit[3]],
];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let raw = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&cause_labels,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties.clone(),
mono,
SurvivalSpec::Net,
);
assert!(
matches!(raw, Err(SurvivalError::EventCodeInvalid { .. })),
"raw competing-risks cause labels must be rejected as EventCodeInvalid (not NonFiniteInput), got {raw:?}"
);
let any_event = pooled_any_event_indicator(cause_labels.view());
assert_eq!(any_event, array![0u8, 1u8, 1u8, 0u8]);
assert_eq!(
cause_specific_event_indicator(cause_labels.view(), 1),
array![0u8, 1u8, 0u8, 0u8]
);
assert_eq!(
cause_specific_event_indicator(cause_labels.view(), 2),
array![0u8, 0u8, 1u8, 0u8]
);
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&any_event,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("pooled any-event baseline model must construct from competing-risks data");
let beta = array![-1.0_f64, 0.8];
let state = model.update_state(&beta).expect("pooled baseline state");
assert!(
state.deviance.is_finite(),
"pooled baseline deviance must be finite, got {}",
state.deviance
);
assert!(
state.gradient.iter().all(|g| g.is_finite()),
"pooled baseline gradient must be finite"
);
}
#[test]
fn offset_channel_residuals_match_central_fd_of_nll() {
let age_entry = array![0.5_f64, 0.0, 0.3];
let age_exit = array![1.4_f64, 1.0, 2.0];
let event_target = array![1u8, 1u8, 0u8];
let event_competing = array![0u8, 0u8, 0u8];
let sampleweight = array![1.0_f64, 2.5, 0.7];
let x_entry = array![
[1.0, age_entry[0].ln()],
[1.0, age_entry[1].max(1e-8).ln()],
[1.0, age_entry[2].ln()]
];
let x_exit = array![
[1.0, age_exit[0].ln()],
[1.0, age_exit[1].ln()],
[1.0, age_exit[2].ln()]
];
let x_derivative = array![
[0.0, 1.0 / age_exit[0]],
[0.0, 1.0 / age_exit[1]],
[0.0, 1.0 / age_exit[2]]
];
let o_entry = array![0.2_f64, 0.0, 0.1];
let o_exit = array![0.4_f64, 0.5, 0.7];
let o_deriv = array![0.3_f64, 0.8, 0.5];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let beta = array![-0.7_f64, 0.6];
let build = |o_e: &Array1<f64>, o_x: &Array1<f64>, o_d: &Array1<f64>| {
survival_model_with_offsets(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
Some(SurvivalBaselineOffsets {
eta_entry: o_e.view(),
eta_exit: o_x.view(),
derivative_exit: o_d.view(),
}),
penalties.clone(),
mono,
SurvivalSpec::Net,
)
.expect("model build")
};
let base = build(&o_entry, &o_exit, &o_deriv);
let resid = base
.offset_channel_residuals(&beta)
.expect("offset residuals");
assert_eq!(resid.exit.len(), 3);
assert_eq!(resid.entry.len(), 3);
assert_eq!(resid.derivative.len(), 3);
let nll = |m: &WorkingModelSurvival| 0.5 * m.update_state(&beta).expect("state").deviance;
let h = 1e-6;
assert_eq!(resid.entry[1], 0.0);
assert_eq!(resid.derivative[2], 0.0);
for i in 0..3 {
{
let mut op = o_exit.clone();
let mut om = o_exit.clone();
op[i] += h;
om[i] -= h;
let fd = (nll(&build(&o_entry, &op, &o_deriv))
- nll(&build(&o_entry, &om, &o_deriv)))
/ (2.0 * h);
assert!(
(resid.exit[i] - fd).abs() < 1e-6,
"∂NLL/∂o_X[{i}]: analytic={:.6e} fd={:.6e}",
resid.exit[i],
fd
);
}
{
let mut op = o_entry.clone();
let mut om = o_entry.clone();
op[i] += h;
om[i] -= h;
let fd = (nll(&build(&op, &o_exit, &o_deriv))
- nll(&build(&om, &o_exit, &o_deriv)))
/ (2.0 * h);
assert!(
(resid.entry[i] - fd).abs() < 1e-6,
"∂NLL/∂o_E[{i}]: analytic={:.6e} fd={:.6e}",
resid.entry[i],
fd
);
}
{
let mut op = o_deriv.clone();
let mut om = o_deriv.clone();
op[i] += h;
om[i] -= h;
let fd = (nll(&build(&o_entry, &o_exit, &op))
- nll(&build(&o_entry, &o_exit, &om)))
/ (2.0 * h);
assert!(
(resid.derivative[i] - fd).abs() < 1e-6,
"∂NLL/∂o_D[{i}]: analytic={:.6e} fd={:.6e}",
resid.derivative[i],
fd
);
}
}
}
#[test]
fn offset_channel_residuals_respect_zero_sampleweight() {
let age_entry = array![1.0_f64, 2.0];
let age_exit = array![2.0_f64, 3.5];
let event_target = array![1u8, 1u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![0.0_f64, 1.2]; let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let beta = array![-1.0_f64, 0.8];
let model = survival_model_with_offsets(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
Some(SurvivalBaselineOffsets {
eta_entry: array![0.0_f64, 0.1].view(),
eta_exit: array![0.0_f64, 0.2].view(),
derivative_exit: array![0.0_f64, 0.1].view(),
}),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("model");
let r = model.offset_channel_residuals(&beta).expect("resid");
assert_eq!(r.exit[0], 0.0);
assert_eq!(r.entry[0], 0.0);
assert_eq!(r.derivative[0], 0.0);
assert!(r.exit[1] != 0.0);
}
#[test]
fn offset_channel_residuals_reject_beta_dim_mismatch() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0_f64];
let x_entry = array![[1.0, 0.0]];
let x_exit = array![[1.0, 0.7]];
let x_derivative = array![[0.0, 0.5]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("model");
let bad_beta = array![0.0_f64]; let err = model
.offset_channel_residuals(&bad_beta)
.expect_err("mismatch must error");
match err {
EstimationError::InvalidInput(msg) => {
assert!(msg.contains("beta dimension mismatch"), "msg={msg}")
}
other => panic!("expected InvalidInput, got {other:?}"),
}
}
#[test]
fn crudespec_is_rejected_by_one_hazard_engine() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![0u8];
let event_competing = array![1u8];
let sampleweight = array![1.0];
let x_entry = array![[0.1]];
let x_exit = array![[0.4]];
let x_derivative = array![[1.0]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let err = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Crude,
)
.expect_err("crude fitting should be rejected by the one-hazard engine");
assert!(matches!(err, SurvivalError::UnsupportedSpec("crude")));
}
#[test]
fn nonstructural_models_require_explicit_monotonicity_collocation() {
let age_entry = array![1.0_f64, 1.5_f64];
let age_exit = array![2.0_f64, 2.5_f64];
let event_target = array![0u8, 0u8];
let event_competing = array![0u8, 1u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[0.2], [0.1]];
let x_exit = array![[0.3], [0.2]];
let x_derivative = array![[1.0], [1.0]];
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct censored survival model");
assert!(
model.monotonicity_linear_constraints().is_none(),
"non-structural survival models must not fabricate rowwise monotonicity constraints"
);
}
#[test]
fn decreasing_interval_is_rejectedwithout_target_events() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![0u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[0.5]];
let x_exit = array![[0.0]];
let x_derivative = array![[1.0]];
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct censored survival model");
let err = model
.update_state(&array![1.0])
.expect_err("decreasing cumulative hazard increment should be rejected");
assert!(
err.to_string().contains("cumulative hazard decreased"),
"unexpected error: {err}"
);
}
fn smooth_crude_risk(beta_d: f64, beta_m: f64) -> CrudeRiskResult {
calculate_crude_risk_quadrature(
0.0,
1.0,
&[0.0, 1.0],
beta_d.exp(),
beta_m.exp(),
array![1.0].view(),
array![1.0].view(),
|u, design_d, deriv_d, design_m| {
let cumulative_d = beta_d.exp() * (1.0 + 0.2 * u);
let cumulative_m = beta_m.exp() * (1.0 + 0.1 * u);
let inst_hazard_d = 0.2 * beta_d.exp();
design_d[0] = 1.0;
deriv_d[0] = 0.0;
design_m[0] = 1.0;
Ok((inst_hazard_d, cumulative_d, cumulative_m))
},
)
.expect("smooth crude-risk quadrature should succeed")
}
#[test]
fn crude_riskgradient_matches_monotoneobjective() {
let beta_d = -0.2_f64;
let beta_m = -0.5_f64;
let result = smooth_crude_risk(beta_d, beta_m);
let eps = 1e-6;
let fd_d = (smooth_crude_risk(beta_d + eps, beta_m).risk
- smooth_crude_risk(beta_d - eps, beta_m).risk)
/ (2.0 * eps);
let fd_m = (smooth_crude_risk(beta_d, beta_m + eps).risk
- smooth_crude_risk(beta_d, beta_m - eps).risk)
/ (2.0 * eps);
assert!(
(result.diseasegradient[0] - fd_d).abs() < 1e-5,
"disease gradient mismatch for monotone crude risk: analytic={} fd={fd_d}",
result.diseasegradient[0]
);
assert!(
(result.mortalitygradient[0] - fd_m).abs() < 1e-5,
"mortality gradient mismatch for monotone crude risk: analytic={} fd={fd_m}",
result.mortalitygradient[0]
);
}
#[test]
fn survival_working_state_is_ridge_free() {
let age_entry = array![1.0_f64, 2.0_f64];
let age_exit = array![2.0_f64, 3.5_f64];
let event_target = array![1u8, 0u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
matrix: array![[2.0]],
lambda: 1.7,
range: 1..2,
nullspace_dim: 0,
}]);
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let beta = array![-1.2, 0.4];
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties.clone(),
mono,
SurvivalSpec::Net,
)
.expect("construct survival model");
let state = model.update_state(&beta).expect("survival state");
assert_eq!(
state.ridge_used, 0.0,
"survival objective must not fuse a coefficient ridge"
);
let expected_penalty = 2.0 * penalties.deviance(&beta);
assert!(
(state.penalty_term - expected_penalty).abs() < 1e-12,
"penalty_term mismatch: state={} expected={}",
state.penalty_term,
expected_penalty
);
}
#[test]
fn negative_penalty_lambda_is_rejected() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[1.0, 0.0]];
let x_exit = array![[1.0, 0.5]];
let x_derivative = array![[0.0, 1.0]];
let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
matrix: array![[1.0]],
lambda: -0.1,
range: 1..2,
nullspace_dim: 0,
}]);
let err = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect_err("negative lambda must be rejected");
assert!(matches!(err, SurvivalError::NonFiniteInput));
}
#[test]
fn penalty_block_range_and_shapemust_match_coefficients() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[1.0, 0.0]];
let x_exit = array![[1.0, 0.5]];
let x_derivative = array![[0.0, 1.0]];
let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
matrix: array![[1.0]],
lambda: 0.5,
range: 0..2,
nullspace_dim: 0,
}]);
let err = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
SurvivalMonotonicityPenalty { tolerance: 1e-8 },
SurvivalSpec::Net,
)
.expect_err("penalty block geometry must match coefficient support");
assert!(matches!(err, SurvivalError::DimensionMismatch));
}
#[test]
fn survivalgradient_matches_ridge_free_objective_fd() {
let age_entry = array![1.0_f64, 2.0_f64, 3.0_f64];
let age_exit = array![2.0_f64, 3.5_f64, 4.0_f64];
let event_target = array![1u8, 0u8, 1u8];
let event_competing = array![0u8, 0u8, 0u8];
let sampleweight = array![1.0, 1.0, 1.0];
let x_entry = array![
[1.0, age_entry[0].ln()],
[1.0, age_entry[1].ln()],
[1.0, age_entry[2].ln()]
];
let x_exit = array![
[1.0, age_exit[0].ln()],
[1.0, age_exit[1].ln()],
[1.0, age_exit[2].ln()]
];
let x_derivative = array![
[0.0, 1.0 / age_exit[0]],
[0.0, 1.0 / age_exit[1]],
[0.0, 1.0 / age_exit[2]]
];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let beta = array![-1.0, 3.0];
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("construct survival model");
let state = model.update_state(&beta).expect("state at beta");
let eps = 1e-7;
for j in 0..beta.len() {
let mut plus = beta.clone();
let mut minus = beta.clone();
plus[j] += eps;
minus[j] -= eps;
let state_plus = model.update_state(&plus).expect("state at beta + eps");
let state_minus = model.update_state(&minus).expect("state at beta - eps");
let obj_plus = 0.5 * (state_plus.deviance + state_plus.penalty_term);
let obj_minus = 0.5 * (state_minus.deviance + state_minus.penalty_term);
let fd = (obj_plus - obj_minus) / (2.0 * eps);
assert_eq!(
state.gradient[j].signum(),
fd.signum(),
"objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
state.gradient[j]
);
assert!(
(state.gradient[j] - fd).abs() < 1e-5,
"objective/gradient mismatch at j={j}: grad={} fd={fd}",
state.gradient[j]
);
}
}
fn laml_fd_test_model(lambda: f64) -> WorkingModelSurvival {
let age_entry: Array1<f64> = Array1::from(vec![
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,
34.0, 39.0, 44.0, 49.0, 54.0, 59.0,
]);
let age_exit: Array1<f64> = Array1::from(vec![
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,
48.0, 51.0, 58.0, 62.0, 66.0, 69.0,
]);
let event_target = Array1::from(vec![
1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
]);
let event_competing = Array1::<u8>::zeros(age_entry.len());
let sampleweight = Array1::from_elem(age_entry.len(), 1.0_f64);
let n = age_entry.len();
let ln_age_mean: f64 = {
let mut sum = 0.0;
for i in 0..n {
sum += age_entry[i].ln() + age_exit[i].ln();
}
sum / (2.0 * n as f64)
};
let mut x_entry = Array2::<f64>::zeros((n, 2));
let mut x_exit = Array2::<f64>::zeros((n, 2));
let mut x_derivative = Array2::<f64>::zeros((n, 2));
for i in 0..n {
x_entry[[i, 0]] = 1.0;
x_exit[[i, 0]] = 1.0;
x_entry[[i, 1]] = age_entry[i].ln() - ln_age_mean;
x_exit[[i, 1]] = age_exit[i].ln() - ln_age_mean;
x_derivative[[i, 0]] = 0.0;
x_derivative[[i, 1]] = 1.0 / age_exit[i];
}
let penalties = PenaltyBlocks::new(vec![
PenaltyBlock {
matrix: array![[3.0]],
lambda: 0.0,
range: 0..1,
nullspace_dim: 0,
},
PenaltyBlock {
matrix: array![[2.5]],
lambda,
range: 1..2,
nullspace_dim: 0,
},
]);
survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
SurvivalMonotonicityPenalty { tolerance: 1e-8 },
SurvivalSpec::Net,
)
.expect("construct LAML FD survival model")
}
fn laml_test_logdet_h(state: &WorkingState) -> f64 {
use gam_problem::PseudoLogdetMode;
use gam_solve::estimate::reml::reml_outer_engine::{
DenseSpectralOperator, HessianFactorization,
};
DenseSpectralOperator::from_symmetric_with_mode(
&state.hessian.to_dense(),
PseudoLogdetMode::PositiveDefinite,
)
.expect("positive-definite fitted survival Hessian")
.logdet()
}
#[test]
fn survival_laml_mode_response_and_hessian_drift_match_finite_differences() {
use gam_linalg::faer_ndarray::FaerCholesky;
use gam_problem::PseudoLogdetMode;
use gam_solve::estimate::reml::reml_outer_engine::{
DenseSpectralOperator, HessianFactorization,
};
const RHO: f64 = 4.0;
const RHO_STEP: f64 = 1.0e-5;
const BETA_STEP: f64 = 1.0e-5;
const REL_TOL: f64 = 2.0e-4;
let beta0 = array![-2.5_f64, 1.0];
let model = laml_fd_test_model(1.0);
let (center_model, beta_hat) = model
.reconverge_survival_inner_mode(&[RHO], &beta0)
.expect("reconverge survival mode at rho=4");
let center_state = center_model
.update_state(&beta_hat)
.expect("center survival state");
let h_center = center_state.hessian.to_dense();
let p = beta_hat.len();
let active: Vec<&PenaltyBlock> = center_model
.penalties
.blocks
.iter()
.filter(|block| block.lambda > 0.0)
.collect();
assert_eq!(active.len(), 1, "fixture must have one active penalty block");
let block = active[0];
let mut a = Array2::<f64>::zeros((p, p));
for i in 0..block.matrix.nrows() {
for j in 0..block.matrix.ncols() {
a[[block.range.start + i, block.range.start + j]] =
block.lambda * block.matrix[[i, j]];
}
}
let factor = h_center
.cholesky(faer::Side::Lower)
.expect("center Hessian Cholesky");
let v = factor.solvevec(&a.dot(&beta_hat));
let u = -&v;
let (plus_model, beta_plus) = model
.reconverge_survival_inner_mode(&[RHO + RHO_STEP], &beta_hat)
.expect("rho-plus survival mode");
let (minus_model, beta_minus) = model
.reconverge_survival_inner_mode(&[RHO - RHO_STEP], &beta_hat)
.expect("rho-minus survival mode");
let beta_fd = (&beta_plus - &beta_minus) / (2.0 * RHO_STEP);
let correction = center_model
.survival_hessian_derivative_correction(&beta_hat, &u)
.expect("analytic survival Hessian correction");
let beta_dir_plus = &beta_hat + &u.mapv(|value| BETA_STEP * value);
let beta_dir_minus = &beta_hat - &u.mapv(|value| BETA_STEP * value);
let h_beta_plus = center_model
.update_state(&beta_dir_plus)
.expect("beta-direction plus state")
.hessian
.to_dense();
let h_beta_minus = center_model
.update_state(&beta_dir_minus)
.expect("beta-direction minus state")
.hessian
.to_dense();
let correction_fd = (&h_beta_plus - &h_beta_minus) / (2.0 * BETA_STEP);
let state_plus = plus_model
.update_state(&beta_plus)
.expect("rho-plus state");
let state_minus = minus_model
.update_state(&beta_minus)
.expect("rho-minus state");
let total_fd =
(state_plus.hessian.to_dense() - state_minus.hessian.to_dense()) / (2.0 * RHO_STEP);
let total_analytic = &a + &correction;
let total_sign_reversed = &a - &correction;
let relative_vector_error = |actual: &Array1<f64>, expected: &Array1<f64>| {
let difference = actual
.iter()
.zip(expected.iter())
.map(|(&lhs, &rhs)| (lhs - rhs) * (lhs - rhs))
.sum::<f64>()
.sqrt();
let scale = expected.iter().map(|value| value * value).sum::<f64>().sqrt();
difference / scale.max(1.0e-12)
};
let relative_matrix_error = |actual: &Array2<f64>, expected: &Array2<f64>| {
let difference = actual
.iter()
.zip(expected.iter())
.map(|(&lhs, &rhs)| (lhs - rhs) * (lhs - rhs))
.sum::<f64>()
.sqrt();
let scale = expected.iter().map(|value| value * value).sum::<f64>().sqrt();
difference / scale.max(1.0e-12)
};
let mode_error = relative_vector_error(&u, &beta_fd);
let correction_error = relative_matrix_error(&correction, &correction_fd);
let correction_reversed_error = relative_matrix_error(&(-&correction), &correction_fd);
let total_error = relative_matrix_error(&total_analytic, &total_fd);
let total_reversed_error = relative_matrix_error(&total_sign_reversed, &total_fd);
let t1_plus = 0.5 * (state_plus.deviance + state_plus.penalty_term);
let t1_minus = 0.5 * (state_minus.deviance + state_minus.penalty_term);
let t1_fd = (t1_plus - t1_minus) / (2.0 * RHO_STEP);
let t2_fd = 0.5
* (laml_test_logdet_h(&state_plus) - laml_test_logdet_h(&state_minus))
/ (2.0 * RHO_STEP);
let t3_fd = -0.5_f64;
let positive_hop = DenseSpectralOperator::from_symmetric_with_mode(
&h_center,
PseudoLogdetMode::PositiveDefinite,
)
.expect("positive-definite operator at the fitted survival mode");
let half_trace_a = 0.5 * positive_hop.trace_hinv_product(&a);
let half_trace_c = 0.5 * positive_hop.trace_hinv_product(&correction);
let t1_analytic = 0.5 * beta_hat.dot(&a.dot(&beta_hat));
let expected_gradient = t1_analytic + half_trace_a + half_trace_c - 0.5;
let rho = array![RHO];
let (_, public_gradient) = center_model
.unified_lamlobjective_and_rhogradient(&beta_hat, ¢er_state, &rho)
.expect("public survival LAML gradient at fitted mode");
eprintln!(
"survival rho-chain decomposition: mode_error={mode_error:.6e} \
correction_error={correction_error:.6e} correction_reversed_error={correction_reversed_error:.6e} \
total_error={total_error:.6e} total_reversed_error={total_reversed_error:.6e} \
t1_fd={t1_fd:+.12e} t2_fd={t2_fd:+.12e} t3_fd={t3_fd:+.12e} \
t1_analytic={t1_analytic:+.12e} half_trace_a={half_trace_a:+.12e} \
half_trace_c={half_trace_c:+.12e} expected_gradient={expected_gradient:+.12e} \
public_gradient={:?} beta_analytic={:?} beta_fd={:?} correction={:?} correction_fd={:?} \
total_analytic={:?} total_fd={:?}",
public_gradient.to_vec(),
u.to_vec(),
beta_fd.to_vec(),
correction,
correction_fd,
total_analytic,
total_fd,
);
let public_gradient_error = (public_gradient[0] - expected_gradient).abs()
/ expected_gradient.abs().max(1.0);
assert!(
mode_error <= REL_TOL
&& correction_error <= REL_TOL
&& total_error <= REL_TOL
&& public_gradient_error <= REL_TOL,
"survival rho chain-rule identity failed: mode_error={mode_error:.6e}, \
correction_error={correction_error:.6e} (sign-reversed={correction_reversed_error:.6e}), \
total_error={total_error:.6e} (sign-reversed={total_reversed_error:.6e}), \
public_gradient_error={public_gradient_error:.6e}"
);
}
fn laml_rail_fd_test_model(lambda0: f64, lambda1: f64) -> WorkingModelSurvival {
let age_entry: Array1<f64> = Array1::from(vec![
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,
34.0, 39.0, 44.0, 49.0, 54.0, 59.0,
]);
let age_exit: Array1<f64> = Array1::from(vec![
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,
48.0, 51.0, 58.0, 62.0, 66.0, 69.0,
]);
let event_target = Array1::from(vec![
1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
]);
let event_competing = Array1::<u8>::zeros(age_entry.len());
let sampleweight = Array1::from_elem(age_entry.len(), 1.0_f64);
let n = age_entry.len();
let ln_age_mean: f64 = {
let mut sum = 0.0;
for i in 0..n {
sum += age_entry[i].ln() + age_exit[i].ln();
}
sum / (2.0 * n as f64)
};
let mut x_entry = Array2::<f64>::zeros((n, 2));
let mut x_exit = Array2::<f64>::zeros((n, 2));
let mut x_derivative = Array2::<f64>::zeros((n, 2));
for i in 0..n {
x_entry[[i, 0]] = 1.0;
x_exit[[i, 0]] = 1.0;
x_entry[[i, 1]] = age_entry[i].ln() - ln_age_mean;
x_exit[[i, 1]] = age_exit[i].ln() - ln_age_mean;
x_derivative[[i, 0]] = 0.0;
x_derivative[[i, 1]] = 1.0 / age_exit[i];
}
let penalties = PenaltyBlocks::new(vec![
PenaltyBlock {
matrix: array![[3.0]],
lambda: lambda0,
range: 0..1,
nullspace_dim: 0,
},
PenaltyBlock {
matrix: array![[2.5]],
lambda: lambda1,
range: 1..2,
nullspace_dim: 0,
},
]);
survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
SurvivalMonotonicityPenalty { tolerance: 1e-8 },
SurvivalSpec::Net,
)
.expect("construct two-active-block rail LAML FD model")
}
#[test]
fn survival_laml_rho_gradient_matches_fd_at_the_over_smoothing_rail() {
use gam_linalg::faer_ndarray::FaerEigh;
const RAIL_RHO0: f64 = 7.394829814011909;
const FREE_RHO1: f64 = -2.45;
const FD_STEP: f64 = 1.0e-4;
let beta0 = array![-2.5_f64, 1.0];
let rho = array![RAIL_RHO0, FREE_RHO1];
let model = laml_rail_fd_test_model(RAIL_RHO0.exp(), FREE_RHO1.exp());
let (value, analytic) = model
.evaluate_survival_lamlcost_and_gradient(
rho.as_slice().expect("contiguous rho"),
&beta0,
)
.expect("rail LAML analytic value+gradient (inner solve must converge at the rail)");
let (rail_model, beta_hat) = model
.reconverge_survival_inner_mode(rho.as_slice().expect("contiguous rho"), &beta0)
.expect("reconverge inner mode at the rail");
let state = rail_model
.update_state(&beta_hat)
.expect("inner state at the rail");
let h_dense = state.hessian.to_dense();
let (evals, _) = h_dense.eigh(faer::Side::Lower).expect("eigh at rail");
let min_ev = evals.iter().copied().fold(f64::INFINITY, f64::min);
let max_ev = evals.iter().copied().fold(f64::NEG_INFINITY, f64::max);
let cond = max_ev / min_ev.abs().max(f64::MIN_POSITIVE);
let term_values = |r: &Array1<f64>| -> (f64, f64, f64) {
let (cand, b) = model
.reconverge_survival_inner_mode(r.as_slice().expect("contiguous rho"), &beta0)
.expect("reconverge for per-term FD");
let st = cand.update_state(&b).expect("state for per-term FD");
let t1 = 0.5 * (st.deviance + st.penalty_term);
let t2 = 0.5 * laml_test_logdet_h(&st);
let t3 = -0.5 * (r[0] + 3.0_f64.ln() + r[1] + 2.5_f64.ln());
(t1, t2, t3)
};
let mut fd = vec![0.0_f64; rho.len()];
let mut fd_terms = vec![(0.0_f64, 0.0_f64, 0.0_f64); rho.len()];
for j in 0..rho.len() {
let mut plus = rho.clone();
plus[j] += FD_STEP;
let mut minus = rho.clone();
minus[j] -= FD_STEP;
let fp = model
.evaluate_survival_lamlcost_and_gradient(
plus.as_slice().expect("contiguous rho"),
&beta0,
)
.expect("rail LAML f+ (probe ρ inner solve must converge)")
.0;
let fm = model
.evaluate_survival_lamlcost_and_gradient(
minus.as_slice().expect("contiguous rho"),
&beta0,
)
.expect("rail LAML f- (probe ρ inner solve must converge)")
.0;
fd[j] = (fp - fm) / (2.0 * FD_STEP);
let (p1, p2, p3) = term_values(&plus);
let (m1, m2, m3) = term_values(&minus);
fd_terms[j] = (
(p1 - m1) / (2.0 * FD_STEP),
(p2 - m2) / (2.0 * FD_STEP),
(p3 - m3) / (2.0 * FD_STEP),
);
}
eprintln!(
"[RAIL-FD] rho=[{RAIL_RHO0:.9}, {FREE_RHO1}] value={value:.9} inner_H min_ev={min_ev:.3e} max_ev={max_ev:.3e} cond={cond:.3e}"
);
let mut amplification = vec![0.0_f64; rho.len()];
for j in 0..rho.len() {
let (dt1, dt2, dt3) = fd_terms[j];
amplification[j] = dt2.abs().max(dt3.abs()) / analytic[j].abs().max(f64::MIN_POSITIVE);
let amp = amplification[j];
eprintln!(
"[RAIL-FD] rho[{j}]: analytic_total={:.6e} fd_total={:.6e} abs_err={:.3e} amplification={amp:.3e} | fd_terms: half_d(dev+pen)={dt1:.6e} half_dlogdetH={dt2:.6e} neg_half_dlogdetS={dt3:.6e}",
analytic[j],
fd[j],
(analytic[j] - fd[j]).abs()
);
}
for j in 0..rho.len() {
let tol = 1.0e-4 * (1.0 + analytic[j].abs().max(fd[j].abs()));
assert!(
(analytic[j] - fd[j]).abs() <= tol,
"survival LAML ρ-gradient desync at coordinate {j} in the over-smoothing rail regime: \
analytic={:.6e} fd={:.6e} (inner H cond={cond:.3e}); see the per-term [RAIL-FD] grid above",
analytic[j],
fd[j],
);
}
let max_amplification = amplification
.iter()
.copied()
.fold(0.0_f64, f64::max);
assert!(
max_amplification.is_finite() && max_amplification <= 1.0e8,
"rail logdet-gradient amplification ratio not finite/bounded: {max_amplification:.3e}"
);
}
#[test]
fn survival_laml_rho_gradient_matches_fd_at_interior_rho() {
const INTERIOR_RHO0: f64 = 0.3;
const INTERIOR_RHO1: f64 = -0.5;
const FD_STEP: f64 = 1.0e-4;
let beta0 = array![-2.5_f64, 1.0];
let rho = array![INTERIOR_RHO0, INTERIOR_RHO1];
let model = laml_rail_fd_test_model(INTERIOR_RHO0.exp(), INTERIOR_RHO1.exp());
let (_value, analytic) = model
.evaluate_survival_lamlcost_and_gradient(
rho.as_slice().expect("contiguous rho"),
&beta0,
)
.expect("interior LAML analytic value+gradient");
for j in 0..rho.len() {
let mut plus = rho.clone();
plus[j] += FD_STEP;
let mut minus = rho.clone();
minus[j] -= FD_STEP;
let fp = model
.evaluate_survival_lamlcost_and_gradient(
plus.as_slice().expect("contiguous rho"),
&beta0,
)
.expect("interior LAML f+")
.0;
let fm = model
.evaluate_survival_lamlcost_and_gradient(
minus.as_slice().expect("contiguous rho"),
&beta0,
)
.expect("interior LAML f-")
.0;
let fd = (fp - fm) / (2.0 * FD_STEP);
let tol = 1.0e-4 * (1.0 + analytic[j].abs().max(fd.abs()));
assert!(
(analytic[j] - fd).abs() <= tol,
"interior survival LAML ρ-gradient mismatch at coordinate {j}: \
analytic={:.6e} fd={:.6e}",
analytic[j],
fd,
);
}
}
#[test]
fn survival_solver_damping_converges_undamped_objective() {
let rho = -0.35_f64;
let model = laml_fd_test_model(rho.exp());
let beta0 = array![-2.5_f64, 1.0];
let (converged_model, beta) = model
.reconverge_survival_inner_mode(&[rho], &beta0)
.expect("converge survival mode with solver-only damping");
let state = converged_model
.update_state(&beta)
.expect("evaluate undamped objective at converged mode");
assert_eq!(
state.ridge_used, 0.0,
"solver damping must not enter the converged statistical objective"
);
let undamped_stationarity = array1_l2_norm(&state.gradient);
assert!(
undamped_stationarity <= 1.0e-9,
"solver must converge the undamped objective; ||gradient||={undamped_stationarity:.3e}"
);
}
#[test]
fn laml_gradient_and_objective_ignore_inactive_penalty_prefix_blocks() {
let rho0 = -0.35_f64;
let beta0 = array![-2.5_f64, 1.0];
let model = laml_fd_test_model(rho0.exp());
let (model, beta) = model
.reconverge_survival_inner_mode(&[rho0], &beta0)
.expect("converge inner mode for LAML prefix-skip test");
let state = model
.update_state(&beta)
.expect("state for LAML prefix-skip test");
assert_eq!(model.penalties.blocks.len(), 2);
assert_eq!(model.penalties.blocks[0].lambda, 0.0);
assert!(model.penalties.blocks[1].lambda > 0.0);
let rho = Array1::from_iter(
model
.penalties
.blocks
.iter()
.filter(|b| b.lambda > 0.0)
.map(|b| b.lambda.ln()),
);
assert_eq!(
rho.len(),
1,
"fixture should expose exactly one active penalty block for the rho vector"
);
let (obj, grad) = model
.unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
.expect("survival LAML objective and gradient");
let expected = 0.5 * (state.deviance + state.penalty_term)
+ 0.5 * laml_test_logdet_h(&state)
- 0.5 * (rho0 + 2.5_f64.ln());
assert_eq!(
grad.len(),
1,
"rho-gradient must match the active-penalty count, not the full block list"
);
assert!(
(obj - expected).abs() < 1e-10,
"survival LAML objective mismatch with inactive prefix block: obj={obj} expected={expected}",
);
assert!(
grad[0].is_finite(),
"rho-gradient must be finite: {}",
grad[0]
);
}
#[test]
fn survival_laml_refuses_nonstationary_inner_state() {
let rho0 = -0.35_f64;
let beta0 = array![-2.5_f64, 1.0];
let model = laml_fd_test_model(rho0.exp());
let (model, beta_hat) = model
.reconverge_survival_inner_mode(&[rho0], &beta0)
.expect("converge reference survival mode");
let mut beta_off_mode = beta_hat;
beta_off_mode[0] += 0.25;
let state = model
.update_state(&beta_off_mode)
.expect("off-mode state remains in the survival domain");
let rho = array![rho0];
let error = model
.unified_lamlobjective_and_rhogradient(&beta_off_mode, &state, &rho)
.expect_err("LAML must refuse a nonstationary inner state");
assert!(
error
.to_string()
.contains("survival LAML requires a stationary inner mode"),
"unexpected off-mode refusal: {error}"
);
}
#[test]
fn structural_monotonicgradient_matchesobjectivefd() {
let age_entry = array![1.0_f64, 1.3_f64, 1.8_f64];
let age_exit = array![1.6_f64, 2.1_f64, 2.7_f64];
let event_target = array![1u8, 0u8, 1u8];
let event_competing = array![0u8, 0u8, 0u8];
let sampleweight = array![1.0, 1.0, 1.0];
let x_entry = array![
[1.0, 0.2, 0.05, -0.7],
[1.0, 0.5, 0.20, 0.1],
[1.0, 0.9, 0.60, 1.2]
];
let x_exit = array![
[1.0, 0.4, 0.16, -0.7],
[1.0, 0.8, 0.64, 0.1],
[1.0, 1.1, 1.21, 1.2]
];
let x_derivative = array![
[0.0, 0.8, 0.64, 0.0],
[0.0, 0.7, 1.12, 0.0],
[0.0, 0.6, 1.32, 0.0]
];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("construct structural survival model");
model
.set_structural_monotonicity(true, 3)
.expect("enable structural monotonicity");
let constraints = model
.monotonicity_linear_constraints()
.expect("structural derivative constraints");
assert_eq!(constraints.a.nrows(), 2);
assert_eq!(constraints.a.ncols(), 4);
assert_eq!(constraints.a.row(0).to_vec(), vec![0.0, 1.0, 0.0, 0.0]);
assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 0.0, 1.0, 0.0]);
assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
let beta = array![0.2, 0.2, 0.1, 0.2];
let state = model.update_state(&beta).expect("state at structural beta");
let eps = 1e-7;
for j in 0..beta.len() {
let mut plus = beta.clone();
let mut minus = beta.clone();
plus[j] += eps;
minus[j] -= eps;
let state_plus = model.update_state(&plus).expect("state at beta + eps");
let state_minus = model.update_state(&minus).expect("state at beta - eps");
let obj_plus = 0.5 * (state_plus.deviance + state_plus.penalty_term);
let obj_minus = 0.5 * (state_minus.deviance + state_minus.penalty_term);
let fd = (obj_plus - obj_minus) / (2.0 * eps);
assert_eq!(
state.gradient[j].signum(),
fd.signum(),
"structural objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
state.gradient[j]
);
assert!(
(state.gradient[j] - fd).abs() < 2e-5,
"structural objective/gradient mismatch at j={j}: grad={} fd={fd}",
state.gradient[j]
);
}
}
#[test]
fn structural_monotonic_lamlgradient_returns_finitevalues() {
let age_entry = array![1.0_f64, 1.2_f64];
let age_exit = array![1.5_f64, 2.0_f64];
let event_target = array![1u8, 0u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[1.0, 0.2, -0.5], [1.0, 0.4, 0.2]];
let x_exit = array![[1.0, 0.5, -0.5], [1.0, 0.8, 0.2]];
let x_derivative = array![[0.0, 0.9, 0.0], [0.0, 0.7, 0.0]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("construct structural survival model");
model
.set_structural_monotonicity(true, 2)
.expect("enable structural monotonicity");
model.penalties = PenaltyBlocks::new(vec![PenaltyBlock {
matrix: array![[1.0]],
lambda: 0.7,
range: 1..2,
nullspace_dim: 0,
}]);
let beta0 = array![0.2, 0.2, 0.1];
let rho = Array1::from_iter(
model
.penalties
.blocks
.iter()
.filter(|b| b.lambda > 0.0)
.map(|b| b.lambda.ln()),
);
let (model, beta) = model
.reconverge_survival_inner_mode(
rho.as_slice().expect("contiguous structural rho"),
&beta0,
)
.expect("converge structural survival mode");
let state = model.update_state(&beta).expect("state at structural mode");
let (obj, grad) = model
.unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
.expect("laml gradient should work in structural mode");
assert!(obj.is_finite());
assert_eq!(grad.len(), 1);
assert!(grad[0].is_finite());
}
#[test]
fn structural_monotonicity_switches_to_tiny_derivative_guard_constraints() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[0.0]];
let x_exit = array![[0.2]];
let x_derivative = array![[1.0]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("construct structural survival model");
let beta = array![-3.0];
assert!(
model.update_state(&beta).is_err(),
"negative derivative coefficient should violate derivative guard"
);
model
.set_structural_monotonicity(true, 1)
.expect("enable structural monotonicity");
let constraints = model
.monotonicity_linear_constraints()
.expect("structural derivative constraints");
assert_eq!(constraints.a.nrows(), 1);
assert_eq!(constraints.a.ncols(), 1);
assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
assert!(constraints.b[0].abs() <= 1e-12);
let state = model
.update_state(&array![1e-6])
.expect("small positive derivative coefficient should remain feasible");
assert!(state.deviance.is_finite());
}
#[test]
fn derivative_offset_must_clear_nonstructural_monotonicity_threshold() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[1.0, 0.0]];
let x_exit = array![[1.0, 0.0]];
let x_derivative = array![[0.0, 0.0]];
let penalties = PenaltyBlocks::new(Vec::new());
let monotonicity = SurvivalMonotonicityPenalty { tolerance: 3.0 };
let eta_entry_offset = array![0.0];
let eta_exit_offset = array![0.0];
let derivative_offset_below_guard = array![2.0];
let derivative_offset_above_guard = array![3.1];
let offsets_below_guard = SurvivalBaselineOffsets {
eta_entry: eta_entry_offset.view(),
eta_exit: eta_exit_offset.view(),
derivative_exit: derivative_offset_below_guard.view(),
};
let offsets_above_guard = SurvivalBaselineOffsets {
eta_entry: eta_entry_offset.view(),
eta_exit: eta_exit_offset.view(),
derivative_exit: derivative_offset_above_guard.view(),
};
let model_below_guard = survival_model_with_offsets(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
Some(offsets_below_guard),
penalties.clone(),
monotonicity,
SurvivalSpec::Net,
)
.expect("construct model with derivative offset below guard");
let err = model_below_guard
.update_state(&array![0.0, 0.0])
.expect_err("derivative offset below guard should be rejected");
let err_text = err.to_string();
assert!(
err_text.contains("d_eta/dt=2.000e0") && err_text.contains("tolerance=3.000e0"),
"expected derivative guard rejection to report the offset-driven derivative: {err_text}"
);
let model_above_guard = survival_model_with_offsets(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
Some(offsets_above_guard),
penalties,
SurvivalMonotonicityPenalty { tolerance: 3.0 },
SurvivalSpec::Net,
)
.expect("construct model with derivative offset above guard");
let state = model_above_guard
.update_state(&array![0.0, 0.0])
.expect("derivative offset above guard should remain feasible");
assert!(state.deviance.is_finite());
}
#[test]
fn structural_monotonicity_rejects_negative_derivative_offsets() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[0.0]];
let x_exit = array![[0.2]];
let x_derivative = array![[1.0]];
let eta_entry = array![0.0];
let eta_exit = array![0.0];
let derivative_exit = array![-1e-3];
let offsets = SurvivalBaselineOffsets {
eta_entry: eta_entry.view(),
eta_exit: eta_exit.view(),
derivative_exit: derivative_exit.view(),
};
let mut model = survival_model_with_offsets(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
Some(offsets),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct structural survival model");
let err = model
.set_structural_monotonicity(true, 1)
.expect_err("negative derivative offsets must be rejected");
assert!(
err.to_string()
.contains("structural monotonicity requires nonnegative derivative offsets"),
"unexpected error: {err}"
);
}
#[test]
fn structural_monotonicity_emits_coefficient_constraints() {
let age_entry = array![1.0_f64, 1.5_f64];
let age_exit = array![2.0_f64, 3.0_f64];
let event_target = array![1u8, 0u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]];
let x_exit = array![[0.2, 0.4, 1.0], [0.3, 0.5, 1.0]];
let x_derivative = array![[0.3, 0.2, 0.0], [0.4, 0.1, 0.0]];
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct structural survival model");
model
.set_structural_monotonicity(true, 2)
.expect("enable structural monotonicity");
let constraints = model
.monotonicity_linear_constraints()
.expect("structural derivative constraints");
assert_eq!(constraints.a.nrows(), 2);
assert_eq!(constraints.a.ncols(), 3);
assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0, 0.0]);
assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0, 0.0]);
assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
}
#[test]
fn structural_monotonicity_preserves_inactive_time_columns_in_constraints() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[1.0, 0.2]];
let x_exit = array![[1.0, 0.6]];
let x_derivative = array![[0.0, 1.0]];
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct structural survival model");
model
.set_structural_monotonicity(true, 2)
.expect("enable structural monotonicity");
let constraints = model
.monotonicity_linear_constraints()
.expect("structural derivative constraints");
assert_eq!(constraints.a.nrows(), 1);
assert!(
constraints.a[[0, 0]].abs() <= 1e-12,
"inactive time column should remain unconstrained"
);
assert!(
(constraints.a[[0, 1]] - 1.0).abs() <= 1e-12,
"active time column should remain constrained"
);
}
#[test]
fn structural_monotonicity_preserves_sparse_row_patterns() {
let age_entry = array![1.0_f64, 1.5_f64];
let age_exit = array![2.0_f64, 2.5_f64];
let event_target = array![1u8, 1u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[0.0, 0.0], [0.0, 0.0]];
let x_exit = array![[0.4, 0.2], [0.6, 0.3]];
let x_derivative = array![[1.0, 0.0], [1.0, 0.5]];
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct structural survival model");
model
.set_structural_monotonicity(true, 2)
.expect("enable structural monotonicity");
let constraints = model
.monotonicity_linear_constraints()
.expect("structural derivative constraints");
assert_eq!(constraints.a.nrows(), 2);
assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0]);
assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0]);
}
#[test]
fn update_state_rejects_negative_exit_derivative_for_censoredrows() {
let age_entry = array![1.0_f64];
let age_exit = array![1.1_f64];
let event_target = array![0u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[0.0]];
let x_exit = array![[0.0]];
let x_derivative = array![[-1.0]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
penalties,
mono,
SurvivalSpec::Net,
)
.expect("construct censored survival model");
let err = model
.update_state(&array![1.0])
.expect_err("censored row should still enforce monotonic derivative");
assert!(
matches!(err, EstimationError::ParameterConstraintViolation(_)),
"unexpected error: {err:?}"
);
}
fn crude_risk_quadrature_error(
cumulative_entry: f64,
cumulative_exit: f64,
hazard_exit: f64,
) -> SurvivalError {
calculate_crude_risk_quadrature(
1.0,
2.0,
&[],
0.4,
0.2,
array![1.0].view(),
array![1.0].view(),
|_, design_d, deriv_d, design_m| {
design_d[0] = 1.0;
deriv_d[0] = 0.0;
design_m[0] = 1.0;
Ok((cumulative_entry, cumulative_exit, hazard_exit))
},
)
.expect_err("invalid hazards should fail")
}
#[test]
fn crude_risk_quadrature_rejects_decreasing_cumulative_hazard() {
let err = crude_risk_quadrature_error(0.1, 0.3, 0.25);
assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
}
#[test]
fn crude_risk_quadrature_rejects_nonpositive_instantaneous_hazard() {
let err = crude_risk_quadrature_error(0.0, 0.4, 0.25);
assert!(matches!(err, SurvivalError::NonPositiveHazard));
}
#[test]
fn monotonicity_constraints_collapse_positive_collinearrows() {
let a = array![[0.0, 0.5, 0.0], [0.0, 0.25, 0.0], [0.0, 0.125, 0.0]];
let b = array![1e-8, 1e-8, 1e-8];
let compressed = compress_positive_collinear_constraints(&a, &b);
assert_eq!(compressed.a.nrows(), 1);
assert_eq!(compressed.a.ncols(), 3);
assert!(compressed.a[[0, 0]].abs() <= 1e-12);
assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
assert!(compressed.a[[0, 2]].abs() <= 1e-12);
assert!((compressed.b[0] - 8e-8).abs() <= 1e-18);
}
#[test]
fn monotonicity_constraints_preserve_distinct_directions() {
let a = array![[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]];
let b = array![0.2, 0.3, 0.1];
let compressed = compress_positive_collinear_constraints(&a, &b);
assert_eq!(compressed.a.nrows(), 2);
let mut saw_x = false;
let mut saw_y = false;
for i in 0..compressed.a.nrows() {
if (compressed.a[[i, 0]] - 1.0).abs() <= 1e-12 && compressed.a[[i, 1]].abs() <= 1e-12 {
saw_x = true;
assert!((compressed.b[i] - 0.2).abs() <= 1e-12);
}
if compressed.a[[i, 0]].abs() <= 1e-12 && (compressed.a[[i, 1]] - 1.0).abs() <= 1e-12 {
saw_y = true;
assert!((compressed.b[i] - 0.3).abs() <= 1e-12);
}
}
assert!(saw_x);
assert!(saw_y);
}
#[test]
fn monotonicity_constraints_cluster_near_collinearrows() {
let a = array![
[0.0, 0.5, 0.0],
[0.0, 0.50000000003, 0.0],
[0.0, 0.49999999997, 0.0]
];
let b = array![1e-8, 1.00000000005e-8, 0.99999999995e-8];
let compressed = compress_positive_collinear_constraints(&a, &b);
assert_eq!(compressed.a.nrows(), 1);
assert_eq!(compressed.a.ncols(), 3);
assert!(compressed.a[[0, 0]].abs() <= 1e-12);
assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
assert!(compressed.a[[0, 2]].abs() <= 1e-12);
assert!((compressed.b[0] - 2.0e-8).abs() <= 1e-18);
}
#[test]
fn monotonicity_constraints_cluster_spline_like_near_duplicates() {
let a = array![
[0.0, 0.401, 0.302, 0.197],
[0.0, 0.40100000003, 0.30199999998, 0.19700000001],
[0.0, 0.40099999997, 0.30200000002, 0.19699999999],
[0.0, 0.125, 0.500, 0.375]
];
let b = array![2.0e-8, 2.00000000004e-8, 1.99999999996e-8, 3.0e-8];
let compressed = compress_positive_collinear_constraints(&a, &b);
assert_eq!(compressed.a.nrows(), 2);
let mut clustered_face = false;
let mut distinct_face = false;
for i in 0..compressed.a.nrows() {
let row = compressed.a.row(i);
if row[1] > 0.99 && row[2] > 0.7 && row[3] > 0.49 {
clustered_face = true;
assert!((compressed.b[i] - (2.0e-8 / 0.401)).abs() <= 1e-12);
} else {
distinct_face = true;
assert!((row[1] - 0.25).abs() <= 1e-12);
assert!((row[2] - 1.0).abs() <= 1e-12);
assert!((row[3] - 0.75).abs() <= 1e-12);
assert!((compressed.b[i] - 6.0e-8).abs() <= 1e-18);
}
}
assert!(clustered_face);
assert!(distinct_face);
}
#[test]
fn linear_time_monotonicity_constraints_reduce_to_single_halfspace() {
let age_entry = array![1.0_f64, 1.0, 1.0];
let age_exit = array![2.0_f64, 4.0, 8.0];
let event_target = array![0u8, 1u8, 0u8];
let event_competing = array![0u8, 0u8, 0u8];
let sampleweight = array![1.0, 1.0, 1.0];
let x_entry = array![
[1.0, age_entry[0].ln()],
[1.0, age_entry[1].ln()],
[1.0, age_entry[2].ln()]
];
let x_exit = array![
[1.0, age_exit[0].ln()],
[1.0, age_exit[1].ln()],
[1.0, age_exit[2].ln()]
];
let x_derivative = array![[0.0, 0.5], [0.0, 0.25], [0.0, 0.125]];
let penalties = PenaltyBlocks::new(Vec::new());
let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
let collocation_offsets = Array1::zeros(x_derivative.nrows());
let mut inputs = survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
);
inputs.monotonicity_constraint_rows = Some(x_derivative.view());
inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
let model = survival_model(inputs, penalties, mono, SurvivalSpec::Net)
.expect("construct linear survival model");
let constraints = model
.monotonicity_linear_constraints()
.expect("monotonicity constraints");
assert_eq!(constraints.a.nrows(), 1);
assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
assert!((constraints.b[0] - 8e-8).abs() <= 1e-12);
}
#[test]
fn monotonicity_constraints_skip_numericallyzerorows() {
let age_entry = array![1.0_f64, 1.0, 1.0];
let age_exit = array![2.0_f64, 3.0, 4.0];
let event_target = array![0u8, 0u8, 0u8];
let event_competing = array![0u8, 0u8, 0u8];
let sampleweight = array![1.0, 1.0, 1.0];
let x_entry = array![[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]];
let x_exit = x_entry.clone();
let x_derivative = array![[0.0, 0.0], [0.0, 1e-16], [0.0, 0.25]];
let collocation_offsets = Array1::zeros(x_derivative.nrows());
let mut inputs = survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
);
inputs.monotonicity_constraint_rows = Some(x_derivative.view());
inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
let model = survival_model(
inputs,
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct survival model");
let constraints = model
.monotonicity_linear_constraints()
.expect("nonzero derivative row should remain");
assert_eq!(constraints.a.nrows(), 1);
assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
assert!(constraints.b[0].abs() <= 1e-18);
}
#[test]
fn censoredrows_allowzero_boundary_derivative() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![0u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[0.0]];
let x_exit = array![[0.0]];
let x_derivative = array![[1.0]];
let model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct censored survival model");
let state = model
.update_state(&array![0.0])
.expect("censored boundary derivative should remain feasible with zero tolerance");
assert_eq!(state.deviance, 0.0);
assert_eq!(state.log_likelihood, 0.0);
assert_eq!(state.gradient, array![0.0]);
}
#[test]
fn eventrows_keep_positive_derivative_constraint() {
let age_entry = array![1.0_f64, 1.0];
let age_exit = array![2.0_f64, 4.0];
let event_target = array![0u8, 1u8];
let event_competing = array![0u8, 0u8];
let sampleweight = array![1.0, 1.0];
let x_entry = array![[0.0], [0.0]];
let x_exit = array![[0.0], [0.0]];
let x_derivative = array![[0.5], [0.25]];
let collocation_offsets = Array1::zeros(x_derivative.nrows());
let mut inputs = survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
);
inputs.monotonicity_constraint_rows = Some(x_derivative.view());
inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
let model = survival_model(
inputs,
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 1e-8 },
SurvivalSpec::Net,
)
.expect("construct mixed survival model");
let constraints = model
.monotonicity_linear_constraints()
.expect("event row should induce positive lower bound");
assert_eq!(constraints.a.nrows(), 1);
assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
assert!((constraints.b[0] - 4e-8).abs() <= 1e-18);
}
#[test]
fn structural_monotonicity_clamps_tiny_negative_roundoff() {
let age_entry = array![1.0_f64];
let age_exit = array![2.0_f64];
let event_target = array![1u8];
let event_competing = array![0u8];
let sampleweight = array![1.0];
let x_entry = array![[0.0]];
let x_exit = array![[0.0]];
let x_derivative = array![[1.0]];
let mut model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 1e-8 },
SurvivalSpec::Net,
)
.expect("construct survival model");
model
.set_structural_monotonicity(true, 1)
.expect("enable structural monotonicity");
let state = model
.update_state(&array![-1e-8])
.expect("tiny structural roundoff should be clamped");
let expected_deviance = -2.0 * (1.0e-12_f64).ln();
assert!(
(state.deviance - expected_deviance).abs() <= 1e-12,
"floored structural event deviance: expected {expected_deviance}, got {}",
state.deviance
);
assert_eq!(state.gradient, array![0.0]);
}
#[test]
fn compressed_monotonicity_constraints_preserve_uncompressed_feasible_region() {
let uncompressed_constraints = LinearInequalityConstraints {
a: array![
[0.0, 0.5, 0.0],
[0.0, 1.0 / 3.0, 0.0],
[0.0, 0.2, 0.0],
[0.0, 0.125, 0.0]
],
b: Array1::from_elem(4, 1e-8),
};
let compressed_constraints = compress_positive_collinear_constraints(
&uncompressed_constraints.a,
&uncompressed_constraints.b,
);
let candidates = [
array![0.0, 1e-9, 0.0],
array![0.0, 4e-8, 0.0],
array![0.0, 8e-8, 0.0],
array![0.0, 2e-7, 1.5],
];
for beta in candidates {
let uncompressed_ok = (0..uncompressed_constraints.a.nrows()).all(|i| {
uncompressed_constraints.a.row(i).dot(&beta) >= uncompressed_constraints.b[i]
});
let compressed_ok = (0..compressed_constraints.a.nrows())
.all(|i| compressed_constraints.a.row(i).dot(&beta) >= compressed_constraints.b[i]);
assert_eq!(compressed_ok, uncompressed_ok);
}
}
#[test]
fn exact_survival_derivatives_are_time_unit_invariant_up_to_constant_shift() {
let age_entry = array![10.0_f64, 20.0, 25.0];
let age_exit = array![15.0_f64, 30.0, 40.0];
let event_target = array![1u8, 0u8, 1u8];
let event_competing = array![0u8, 0u8, 0u8];
let sampleweight = array![1.0, 2.0, 0.5];
let x_entry = array![[0.1, 0.2, 1.0], [0.3, 0.4, 1.0], [0.2, 0.6, 1.0]];
let x_exit = array![[0.2, 0.3, 1.0], [0.5, 0.7, 1.0], [0.4, 0.8, 1.0]];
let x_derivative = array![[0.04, 0.02, 0.0], [0.03, 0.01, 0.0], [0.02, 0.03, 0.0]];
let beta = array![0.8, 1.1, -0.2];
let base_model = survival_model(
survival_inputs(
&age_entry,
&age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct base survival model");
let base_state = base_model
.update_state(&beta)
.expect("evaluate base survival state");
let time_scale = 365.25;
let scaled_age_entry = age_entry.mapv(|v| v * time_scale);
let scaled_age_exit = age_exit.mapv(|v| v * time_scale);
let scaled_x_derivative = x_derivative.mapv(|v| v / time_scale);
let scaled_model = survival_model(
survival_inputs(
&scaled_age_entry,
&scaled_age_exit,
&event_target,
&event_competing,
&sampleweight,
&x_entry,
&x_exit,
&scaled_x_derivative,
),
PenaltyBlocks::new(Vec::new()),
SurvivalMonotonicityPenalty { tolerance: 0.0 },
SurvivalSpec::Net,
)
.expect("construct scaled survival model");
let scaled_state = scaled_model
.update_state(&beta)
.expect("evaluate scaled survival state");
let weighted_events = sampleweight
.iter()
.zip(event_target.iter())
.map(|(w, d)| *w * f64::from(*d))
.sum::<f64>();
let expected_deviance_shift = 2.0 * weighted_events * time_scale.ln();
assert!(
(scaled_state.deviance - base_state.deviance - expected_deviance_shift).abs() <= 1e-10,
"deviance shift mismatch: scaled={} base={} expected_shift={expected_deviance_shift}",
scaled_state.deviance,
base_state.deviance
);
for j in 0..beta.len() {
assert!(
(scaled_state.gradient[j] - base_state.gradient[j]).abs() <= 1e-12,
"gradient mismatch at j={j}: scaled={} base={}",
scaled_state.gradient[j],
base_state.gradient[j]
);
}
let base_hessian = base_state.hessian.to_dense();
let scaled_hessian = scaled_state.hessian.to_dense();
for r in 0..beta.len() {
for c in 0..beta.len() {
assert!(
(scaled_hessian[[r, c]] - base_hessian[[r, c]]).abs() <= 1e-12,
"hessian mismatch at ({r},{c}): scaled={} base={}",
scaled_hessian[[r, c]],
base_hessian[[r, c]]
);
}
}
}
#[test]
fn survival_laml_rho_gradient_invariant_under_injected_orthogonal_frame_at_the_rail() {
use gam_linalg::faer_ndarray::FaerEigh;
use gam_problem::{EvalMode, PseudoLogdetMode};
use gam_solve::estimate::reml::assembly::InnerAssembly;
use gam_solve::estimate::reml::reml_outer_engine::{
DenseSpectralOperator, DispersionHandling,
};
use gam_solve::estimate::reml::reparameterized_inner::{
RawInnerReparamContext, assemble_reparameterized_inner,
};
use gam_terms::construction::{
canonicalize_penalty_specs, precompute_reparam_invariant_from_canonical,
stable_reparameterizationwith_invariant,
};
use gam_terms::penalty_spec::PenaltySpec;
const RAIL_RHO0: f64 = 7.394829814011909;
const FREE_RHO1: f64 = -2.45;
const DECISION_MARGIN: f64 = 1.0e-9;
let beta0 = array![-2.5_f64, 1.0];
let rho = array![RAIL_RHO0, FREE_RHO1];
let model = laml_rail_fd_test_model(RAIL_RHO0.exp(), FREE_RHO1.exp());
let (rail_model, beta_hat) = model
.reconverge_survival_inner_mode(rho.as_slice().expect("contiguous rho"), &beta0)
.expect("reconverge inner mode at the rail");
let state = rail_model
.update_state(&beta_hat)
.expect("inner state at the rail");
let p = beta_hat.len();
let h_dense = state.hessian.to_dense();
let lambdas: Vec<f64> = rho.iter().map(|&r| r.exp()).collect();
let active_blocks: Vec<&PenaltyBlock> = rail_model
.penalties
.blocks
.iter()
.filter(|b| b.lambda > 0.0)
.collect();
let s_k_embedded: Vec<Array2<f64>> = active_blocks
.iter()
.map(|b| {
let mut s = Array2::<f64>::zeros((p, p));
let (rs, re) = (b.range.start, b.range.end);
s.slice_mut(ndarray::s![rs..re, rs..re]).assign(&b.matrix);
s
})
.collect();
let penalty_specs: Vec<PenaltySpec> = active_blocks
.iter()
.map(|b| PenaltySpec::Block {
local: b.matrix.clone(),
col_range: b.range.clone(),
prior_mean: gam_problem::CoefficientPriorMean::Zero,
structure_hint: None,
op: None,
})
.collect();
let nullspace_dims: Vec<usize> = active_blocks.iter().map(|b| b.nullspace_dim).collect();
let (canonical, _) = canonicalize_penalty_specs(
&penalty_specs,
&nullspace_dims,
p,
"rail-stability gate reparameterization",
)
.expect("canonicalize rail penalties");
let invariant =
precompute_reparam_invariant_from_canonical(&canonical, p).expect("reparam invariant");
let reparam_prod =
stable_reparameterizationwith_invariant(&canonical, &lambdas, p, &invariant, None)
.expect("production reparameterization");
let hessian_logdet_mode = if rail_model
.age_entry
.iter()
.any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD)
{
PseudoLogdetMode::HardPseudo
} else {
PseudoLogdetMode::Smooth
};
let g = {
let sym = Array2::<f64>::from_shape_fn((p, p), |(i, j)| {
(((i * j) as f64) * 0.37).sin() + (((i + j) as f64) * 0.11 + 1.0).cos()
});
let (_, evecs) = sym.eigh(faer::Side::Lower).expect("orthogonal factor from eigh");
evecs
};
let grad_for_reparam = |reparam: &gam_terms::construction::ReparamResult| -> Array1<f64> {
let provider = SurvivalDerivProvider::new(rail_model.clone(), beta_hat.clone());
let ctx = RawInnerReparamContext {
hessian: &h_dense,
beta: &beta_hat,
penalties_embedded: &s_k_embedded,
lambdas: &lambdas,
};
let reparam_inner =
assemble_reparameterized_inner(&ctx, Some(Box::new(provider)), reparam)
.expect("reparameterized inner assembly");
let hop = DenseSpectralOperator::from_symmetric_with_mode(
&reparam_inner.hessian_transformed,
hessian_logdet_mode,
)
.expect("transformed Hessian operator");
let penalty_coords = reparam
.canonical_transformed
.iter()
.map(|cp| cp.to_penalty_coordinate())
.collect::<Vec<_>>();
let result = InnerAssembly {
log_likelihood: state.log_likelihood,
penalty_quadratic: state.penalty_term,
beta: reparam_inner.beta_transformed,
n_observations: rail_model.nrows(),
hessian_op: std::sync::Arc::new(hop),
penalty_coords,
penalty_logdet: reparam_inner.penalty_logdet,
dispersion: DispersionHandling::Fixed {
phi: 1.0,
include_logdet_h: true,
include_logdet_s: true,
},
rho_curvature_scale: 1.0,
rho_prior: gam_problem::RhoPrior::Flat,
hessian_logdet_correction: 0.0,
penalty_subspace_trace: None,
deriv_provider: reparam_inner.deriv_provider,
firth: None,
nullspace_dim: None,
barrier_config: None,
ext_coords: Vec::new(),
ext_coord_pair_fn: None,
rho_ext_pair_fn: None,
fixed_drift_deriv: None,
contracted_psi_second_order: None,
kkt_residual: None,
active_constraints: None,
}
.evaluate(
rho.as_slice().expect("contiguous rho"),
EvalMode::ValueAndGradient,
None,
)
.expect("transformed-frame LAML evaluate");
result.gradient.expect("analytic ρ-gradient present")
};
let g_prod = grad_for_reparam(&reparam_prod);
let mut reparam_conj = reparam_prod.clone();
reparam_conj.qs = reparam_prod.qs.dot(&g);
reparam_conj.canonical_transformed = reparam_prod
.canonical_transformed
.iter()
.map(|cp| {
let mut rotated = cp.clone();
rotated.root = cp.root.dot(&g);
rotated.local = rotated.root.t().dot(&rotated.root);
rotated
})
.collect();
let g_conj = grad_for_reparam(&reparam_conj);
for k in 0..rho.len() {
let drift = (g_prod[k] - g_conj[k]).abs();
assert!(
drift <= DECISION_MARGIN * (1.0 + g_prod[k].abs()),
"rail ρ-gradient not frame-invariant at coordinate {k}: \
Q_s frame {} vs Q_s·G frame {} (drift {:.3e} > margin {:.3e})",
g_prod[k],
g_conj[k],
drift,
DECISION_MARGIN * (1.0 + g_prod[k].abs())
);
}
}
}