Skip to main content

gam_models/survival/
base.rs

1use crate::custom_family::{
2    BlockWorkingSet, CustomFamily, FamilyEvaluation, ParameterBlockState,
3    projected_linear_constraint_stationarity_vector,
4};
5use crate::model_types::EstimationError;
6use gam_linalg::faer_ndarray::{fast_atv, fast_av, fast_xt_diag_x, fast_xt_diag_y};
7use gam_linalg::matrix::SymmetricMatrix;
8use gam_problem::{Coefficients, LinearPredictor};
9use gam_solve::pirls::{
10    LinearInequalityConstraints, WorkingModel as PirlsWorkingModel, WorkingState, array1_l2_norm,
11};
12use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayView3, Axis};
13use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, constants, escalate_ridge};
14use serde::{Deserialize, Serialize};
15use std::collections::BTreeMap;
16use std::convert::Infallible;
17use std::ops::Range;
18use std::sync::LazyLock;
19use thiserror::Error;
20
21#[derive(Debug, Error)]
22pub enum SurvivalError {
23    #[error("input dimensions are inconsistent")]
24    DimensionMismatch,
25    #[error("inputs contain non-finite values")]
26    NonFiniteInput,
27    #[error("survival spec '{0}' is not supported by the one-hazard survival engine")]
28    UnsupportedSpec(&'static str),
29    #[error("crude risk integration setup is invalid")]
30    InvalidIntegrationSetup,
31    #[error("survival time grid must be finite, non-negative, and strictly increasing")]
32    InvalidTimeGrid,
33    #[error("cumulative hazard must be nondecreasing")]
34    NonMonotoneCumulativeHazard,
35    #[error("instantaneous hazard must stay strictly positive during integration")]
36    NonPositiveHazard,
37    #[error("{reason}")]
38    InvalidInput { reason: String },
39    #[error("{reason}")]
40    CauseSpecificDimensionMismatch { reason: String },
41    #[error("{reason}")]
42    NumericalFailure { reason: String },
43    #[error("{reason}")]
44    EventCodeInvalid { reason: String },
45    #[error("{reason}")]
46    EventDegenerate { reason: String },
47    #[error("cause-specific survival block {block}: {source}")]
48    CauseSpecificBlock {
49        block: usize,
50        #[source]
51        source: Box<SurvivalError>,
52    },
53}
54
55impl From<SurvivalError> for String {
56    fn from(err: SurvivalError) -> Self {
57        err.to_string()
58    }
59}
60
61impl From<crate::block_layout::block_count::BlockCountMismatch> for SurvivalError {
62    fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> SurvivalError {
63        SurvivalError::CauseSpecificDimensionMismatch {
64            reason: err.message(),
65        }
66    }
67}
68
69#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
70pub enum SurvivalSpec {
71    #[default]
72    Net,
73    Crude,
74}
75
76#[derive(Debug, Clone)]
77pub struct SurvivalEngineInputs<'a> {
78    pub age_entry: ArrayView1<'a, f64>,
79    pub age_exit: ArrayView1<'a, f64>,
80    pub event_target: ArrayView1<'a, u8>,
81    pub event_competing: ArrayView1<'a, u8>,
82    pub sampleweight: ArrayView1<'a, f64>,
83    pub x_entry: ArrayView2<'a, f64>,
84    pub x_exit: ArrayView2<'a, f64>,
85    pub x_derivative: ArrayView2<'a, f64>,
86    /// Optional global monotonicity collocation rows for the full coefficient vector.
87    /// Non-structural survival models should pass these explicitly instead of
88    /// relying on observed derivative rows.
89    pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
90    /// Baseline offsets corresponding to `monotonicity_constraint_rows`.
91    pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
92}
93
94#[derive(Debug, Clone)]
95pub struct SurvivalTimeCovarInputs<'a> {
96    pub age_entry: ArrayView1<'a, f64>,
97    pub age_exit: ArrayView1<'a, f64>,
98    pub event_target: ArrayView1<'a, u8>,
99    pub event_competing: ArrayView1<'a, u8>,
100    pub sampleweight: ArrayView1<'a, f64>,
101    pub time_entry: ArrayView2<'a, f64>,
102    pub time_exit: ArrayView2<'a, f64>,
103    pub time_derivative: ArrayView2<'a, f64>,
104    pub covariates: ArrayView2<'a, f64>,
105    /// Optional global monotonicity collocation rows for the full coefficient vector.
106    /// Non-structural survival models should pass these explicitly instead of
107    /// relying on observed derivative rows.
108    pub monotonicity_constraint_rows: Option<ArrayView2<'a, f64>>,
109    /// Baseline offsets corresponding to `monotonicity_constraint_rows`.
110    pub monotonicity_constraint_offsets: Option<ArrayView1<'a, f64>>,
111}
112
113#[derive(Debug, Clone)]
114pub struct SurvivalBaselineOffsets<'a> {
115    /// Baseline target contribution to eta at entry time: eta_target(t_entry).
116    pub eta_entry: ArrayView1<'a, f64>,
117    /// Baseline target contribution to eta at exit time: eta_target(t_exit).
118    pub eta_exit: ArrayView1<'a, f64>,
119    /// Baseline target contribution to d eta / d t at exit: eta_target'(t_exit).
120    ///
121    /// This is used in event terms where log-hazard requires
122    /// log(d eta / d t). By threading this as an explicit offset, we get
123    /// "parametric default + spline deviation" behavior:
124    /// - strong penalty => deviation ~ 0 => model collapses to baseline target,
125    /// - weak penalty   => deviation can bend away where data supports it.
126    pub derivative_exit: ArrayView1<'a, f64>,
127}
128
129#[derive(Debug, Clone)]
130pub struct PenaltyBlock {
131    pub matrix: Array2<f64>,
132    pub lambda: f64,
133    pub range: Range<usize>,
134    /// Structural nullspace dimension of this penalty matrix.
135    /// Used for exact pseudo-logdet computation. 0 means full rank.
136    pub nullspace_dim: usize,
137}
138
139#[derive(Debug, Clone)]
140pub struct PenaltyBlocks {
141    pub blocks: Vec<PenaltyBlock>,
142}
143
144impl PenaltyBlocks {
145    pub fn new(blocks: Vec<PenaltyBlock>) -> Self {
146        Self { blocks }
147    }
148
149    pub fn gradient(&self, beta: &Array1<f64>) -> Array1<f64> {
150        let mut grad = Array1::zeros(beta.len());
151        for block in &self.blocks {
152            if block.lambda == 0.0 {
153                continue;
154            }
155            let b = beta.slice(ndarray::s![block.range.clone()]);
156            let g = block.matrix.dot(&b);
157            let mut dst = grad.slice_mut(ndarray::s![block.range.clone()]);
158            dst += &(block.lambda * g);
159        }
160        grad
161    }
162
163    pub fn hessian(&self, dim: usize) -> Array2<f64> {
164        let mut h = Array2::zeros((dim, dim));
165        self.addhessian_inplace(&mut h);
166        h
167    }
168
169    pub fn deviance(&self, beta: &Array1<f64>) -> f64 {
170        let mut value = 0.0;
171        for block in &self.blocks {
172            if block.lambda == 0.0 {
173                continue;
174            }
175            let b = beta.slice(ndarray::s![block.range.clone()]);
176            value += 0.5 * block.lambda * b.dot(&block.matrix.dot(&b));
177        }
178        value
179    }
180
181    pub fn addhessian_inplace(&self, h: &mut Array2<f64>) {
182        for block in &self.blocks {
183            if block.lambda == 0.0 {
184                continue;
185            }
186            let start = block.range.start;
187            let end = block.range.end;
188            h.slice_mut(ndarray::s![start..end, start..end])
189                .scaled_add(block.lambda, &block.matrix);
190        }
191    }
192}
193
194/// Entry ages at or below this value are treated as left-truncation at the time
195/// origin, i.e. "no delayed-entry interval" — the cumulative-hazard term
196/// `exp(η_entry)` is dropped because `H(0) = 0`. The Royston-Parmar baseline is
197/// `η(t) = log H(t)` with `H(t) → 0` as `t → 0`, so `log H` diverges at the
198/// origin; this small positive floor lets a row that genuinely enters at time
199/// zero skip the entry contribution instead of evaluating `log H` at a
200/// degenerate point. Shared so every entry-detection site stays in lockstep.
201///
202/// Public so the fit-orchestration layer can classify a dataset as genuinely
203/// left-truncated (`entry > threshold`) with the SAME origin convention the
204/// likelihood engines use, and pick the left-truncation-robust time anchor
205/// accordingly (issue #1790).
206pub const ENTRY_AT_ORIGIN_THRESHOLD: f64 = 1e-8;
207
208/// Fraction-to-the-boundary factor for the cause-specific feasible-step search.
209/// When a Newton direction would drive a row's derivative down to the
210/// monotonicity floor, the step is capped at this fraction of the distance to
211/// the boundary rather than landing exactly on it. Staying strictly inside the
212/// feasible region (the standard interior-point fraction-to-boundary rule)
213/// keeps the next `1/deriv` / `deriv.ln()` evaluation away from the singular
214/// boundary where curvature blows up.
215const DERIVATIVE_FRACTION_TO_BOUNDARY: f64 = 0.995;
216
217#[derive(Debug, Clone)]
218pub struct CauseSpecificRoystonParmarBlock {
219    pub age_entry: Array1<f64>,
220    pub age_exit: Array1<f64>,
221    pub event_target: Array1<u8>,
222    pub sampleweight: Array1<f64>,
223    pub x_entry: Array2<f64>,
224    pub x_exit: Array2<f64>,
225    pub x_derivative: Array2<f64>,
226    pub offset_eta_entry: Array1<f64>,
227    pub offset_eta_exit: Array1<f64>,
228    pub offset_derivative_exit: Array1<f64>,
229    pub derivative_floor: f64,
230}
231
232/// Cause-specific competing-risks survival as a blockwise custom family.
233///
234/// Each cause is represented by one `ParameterBlockState`, so endpoint-specific
235/// coefficients, shared smoothing labels, and user-defined coefficient groups
236/// stay on the existing `CustomFamily` / `BlockwiseFitOptions` joint-fit path.
237#[derive(Debug, Clone)]
238pub struct CauseSpecificRoystonParmarFamily {
239    blocks: Vec<CauseSpecificRoystonParmarBlock>,
240}
241
242impl CauseSpecificRoystonParmarFamily {
243    pub fn new(blocks: Vec<CauseSpecificRoystonParmarBlock>) -> Result<Self, String> {
244        if blocks.is_empty() {
245            return Err(SurvivalError::InvalidInput {
246                reason: "cause-specific survival family requires at least one endpoint".to_string(),
247            }
248            .into());
249        }
250        for (idx, block) in blocks.iter().enumerate() {
251            validate_cause_specific_block(block).map_err(|err| {
252                SurvivalError::CauseSpecificBlock {
253                    block: idx + 1,
254                    source: Box::new(err),
255                }
256                .to_string()
257            })?;
258        }
259        Ok(Self { blocks })
260    }
261
262    pub fn cause_count(&self) -> usize {
263        self.blocks.len()
264    }
265}
266
267fn validate_cause_specific_block(
268    block: &CauseSpecificRoystonParmarBlock,
269) -> Result<(), SurvivalError> {
270    let n = block.event_target.len();
271    let p = block.x_exit.ncols();
272    if n == 0 || p == 0 {
273        bail_invalid_surv!("empty event vector or coefficient block");
274    }
275    if block.age_entry.len() != n
276        || block.age_exit.len() != n
277        || block.sampleweight.len() != n
278        || block.x_entry.nrows() != n
279        || block.x_exit.nrows() != n
280        || block.x_derivative.nrows() != n
281        || block.x_entry.ncols() != p
282        || block.x_derivative.ncols() != p
283        || block.offset_eta_entry.len() != n
284        || block.offset_eta_exit.len() != n
285        || block.offset_derivative_exit.len() != n
286    {
287        return Err(SurvivalError::CauseSpecificDimensionMismatch {
288            reason: "dimension mismatch".to_string(),
289        });
290    }
291    // A cause-specific block's `event_target` is the binary cause-k indicator
292    // produced by `cause_specific_event_indicator`; a label > 1 here means the
293    // caller passed raw multi-cause codes instead of projecting per cause. That
294    // is a valid finite label, not non-finite input, so it gets its own clear
295    // error rather than the misleading "non-finite input".
296    if let Some(&label) = block.event_target.iter().find(|&&v| v > 1) {
297        return Err(SurvivalError::EventCodeInvalid {
298            reason: format!(
299                "cause-specific block event_target must be the binary cause indicator {{0, 1}}, got multi-cause label {label}; project raw codes per cause via cause_specific_event_indicator"
300            ),
301        });
302    }
303    if block.age_entry.iter().any(|v| !v.is_finite())
304        || block.age_exit.iter().any(|v| !v.is_finite())
305        || block
306            .sampleweight
307            .iter()
308            .any(|v| !v.is_finite() || *v < 0.0)
309        || block.x_entry.iter().any(|v| !v.is_finite())
310        || block.x_exit.iter().any(|v| !v.is_finite())
311        || block.x_derivative.iter().any(|v| !v.is_finite())
312        || block.offset_eta_entry.iter().any(|v| !v.is_finite())
313        || block.offset_eta_exit.iter().any(|v| !v.is_finite())
314        || block.offset_derivative_exit.iter().any(|v| !v.is_finite())
315        || !block.derivative_floor.is_finite()
316        || block.derivative_floor < 0.0
317    {
318        bail_invalid_surv!("non-finite input");
319    }
320    Ok(())
321}
322
323fn evaluate_cause_specific_block(
324    block: &CauseSpecificRoystonParmarBlock,
325    beta: &Array1<f64>,
326) -> Result<(f64, Array1<f64>, Array2<f64>), SurvivalError> {
327    let n = block.event_target.len();
328    let p = block.x_exit.ncols();
329    if beta.len() != p {
330        return Err(SurvivalError::CauseSpecificDimensionMismatch {
331            reason: format!("beta length mismatch: got {}, expected {p}", beta.len()),
332        });
333    }
334    let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
335    let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
336    let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
337    let mut log_likelihood = 0.0;
338    let mut w_exit = Array1::<f64>::zeros(n);
339    let mut w_entry = Array1::<f64>::zeros(n);
340    let mut w_event = Array1::<f64>::zeros(n);
341    let mut w_event_inv_deriv = Array1::<f64>::zeros(n);
342    let mut w_event_outer = Array1::<f64>::zeros(n);
343
344    for i in 0..n {
345        let weight = block.sampleweight[i];
346        if weight <= 0.0 {
347            continue;
348        }
349        if block.age_exit[i] < block.age_entry[i] {
350            bail_invalid_surv!("age_exit < age_entry at row {i}");
351        }
352        let has_entry = block.age_entry[i] > ENTRY_AT_ORIGIN_THRESHOLD;
353        let h_exit = eta_exit[i].exp();
354        let h_entry = if has_entry { eta_entry[i].exp() } else { 0.0 };
355        if !(h_exit.is_finite() && h_entry.is_finite()) {
356            return Err(SurvivalError::NumericalFailure {
357                reason: format!("non-finite cumulative hazard at row {i}"),
358            });
359        }
360        log_likelihood -= weight * (h_exit - h_entry);
361        w_exit[i] = weight * h_exit;
362        w_entry[i] = weight * h_entry;
363        if block.event_target[i] > 0 {
364            let deriv = derivative[i];
365            if !(deriv.is_finite() && deriv > 0.0) {
366                return Err(SurvivalError::NumericalFailure {
367                    reason: format!(
368                        "cause-specific survival derivative must be positive at row {i}, got {deriv}"
369                    ),
370                });
371            }
372            log_likelihood += weight * (eta_exit[i] + deriv.ln());
373            w_event[i] = weight;
374            w_event_inv_deriv[i] = weight / deriv;
375            w_event_outer[i] = weight / (deriv * deriv);
376        }
377    }
378
379    let mut nll_gradient = fast_atv(&block.x_exit, &w_exit);
380    nll_gradient -= &fast_atv(&block.x_entry, &w_entry);
381    nll_gradient -= &fast_atv(&block.x_exit, &w_event);
382    nll_gradient -= &fast_atv(&block.x_derivative, &w_event_inv_deriv);
383    let gradient = -nll_gradient;
384
385    let mut hessian = fast_xt_diag_x(&block.x_exit, &w_exit);
386    hessian -= &fast_xt_diag_x(&block.x_entry, &w_entry);
387    hessian += &fast_xt_diag_x(&block.x_derivative, &w_event_outer);
388    Ok((log_likelihood, gradient, hessian))
389}
390
391impl CustomFamily for CauseSpecificRoystonParmarFamily {
392    // Preserve the pre-gam#1395 behavior: the trait default flipped to OFF (the
393    // flat-prior exact-Newton objective carries no Jeffreys term), so families
394    // that historically armed the term by default opt back in explicitly.
395    fn joint_jeffreys_term_required(&self) -> bool {
396        true
397    }
398
399    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
400        crate::block_layout::block_count::validate_block_count::<SurvivalError>(
401            "cause-specific survival",
402            self.blocks.len(),
403            block_states.len(),
404        )?;
405        let mut log_likelihood = 0.0;
406        let mut blockworking_sets = Vec::with_capacity(self.blocks.len());
407        for (block, state) in self.blocks.iter().zip(block_states.iter()) {
408            let (ll, gradient, hessian) = evaluate_cause_specific_block(block, &state.beta)?;
409            log_likelihood += ll;
410            blockworking_sets.push(BlockWorkingSet::ExactNewton {
411                gradient,
412                hessian: SymmetricMatrix::Dense(hessian),
413            });
414        }
415        Ok(FamilyEvaluation {
416            log_likelihood,
417            blockworking_sets,
418        })
419    }
420
421    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
422        crate::block_layout::block_count::validate_block_count::<SurvivalError>(
423            "cause-specific survival",
424            self.blocks.len(),
425            block_states.len(),
426        )?;
427        let mut log_likelihood = 0.0;
428        for (block, state) in self.blocks.iter().zip(block_states.iter()) {
429            let (ll, _, _) = evaluate_cause_specific_block(block, &state.beta)?;
430            log_likelihood += ll;
431        }
432        Ok(log_likelihood)
433    }
434
435    fn likelihood_blocks_uncoupled(&self) -> bool {
436        true
437    }
438
439    fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
440        true
441    }
442
443    fn output_channel_assignment(
444        &self,
445        specs: &[crate::custom_family::ParameterBlockSpec],
446    ) -> Option<Vec<usize>> {
447        if specs.len() != self.blocks.len() {
448            return Some((0..self.blocks.len()).collect());
449        }
450        Some((0..specs.len()).collect())
451    }
452
453    fn coefficient_hessian_cost(&self, specs: &[crate::custom_family::ParameterBlockSpec]) -> u64 {
454        crate::custom_family::default_coefficient_hessian_cost(specs)
455    }
456
457    fn block_linear_constraints(
458        &self,
459        _: &[ParameterBlockState],
460        block_idx: usize,
461        spec: &crate::custom_family::ParameterBlockSpec,
462    ) -> Result<Option<LinearInequalityConstraints>, String> {
463        let block = self.blocks.get(block_idx).ok_or_else(|| {
464            SurvivalError::CauseSpecificDimensionMismatch {
465                reason: format!(
466                    "cause-specific survival expected block index < {}, got {block_idx}",
467                    self.blocks.len()
468                ),
469            }
470            .to_string()
471        })?;
472        if block.x_derivative.ncols() != spec.design.ncols() {
473            return Err(SurvivalError::CauseSpecificDimensionMismatch {
474                reason: format!(
475                    "cause-specific survival derivative design has {} columns but block '{}' has {}",
476                    block.x_derivative.ncols(),
477                    spec.name,
478                    spec.design.ncols()
479                ),
480            }
481            .into());
482        }
483        let rhs = block
484            .offset_derivative_exit
485            .mapv(|offset| block.derivative_floor - offset);
486        Ok(Some(LinearInequalityConstraints {
487            a: block.x_derivative.clone(),
488            b: rhs,
489        }))
490    }
491
492    fn max_feasible_step_size(
493        &self,
494        block_states: &[ParameterBlockState],
495        block_idx: usize,
496        delta: &Array1<f64>,
497    ) -> Result<Option<f64>, String> {
498        let block = self.blocks.get(block_idx).ok_or_else(|| {
499            SurvivalError::CauseSpecificDimensionMismatch {
500                reason: format!(
501                    "cause-specific survival expected block index < {}, got {block_idx}",
502                    self.blocks.len()
503                ),
504            }
505            .to_string()
506        })?;
507        let state = block_states.get(block_idx).ok_or_else(|| {
508            SurvivalError::CauseSpecificDimensionMismatch {
509                reason: format!(
510                    "cause-specific survival expected {} block states, got {}",
511                    self.blocks.len(),
512                    block_states.len()
513                ),
514            }
515            .to_string()
516        })?;
517        if delta.len() != state.beta.len() || block.x_derivative.ncols() != delta.len() {
518            return Err(SurvivalError::CauseSpecificDimensionMismatch {
519                reason: "cause-specific survival feasible-step dimension mismatch".to_string(),
520            }
521            .into());
522        }
523        let derivative = fast_av(&block.x_derivative, &state.beta) + &block.offset_derivative_exit;
524        let derivative_delta = fast_av(&block.x_derivative, delta);
525        let mut alpha_max = 1.0_f64;
526        for i in 0..derivative.len() {
527            if block.sampleweight[i] <= 0.0 {
528                continue;
529            }
530            let current = derivative[i] - block.derivative_floor;
531            let slope = derivative_delta[i];
532            if slope < 0.0 {
533                if current <= 0.0 {
534                    return Ok(Some(0.0));
535                }
536                alpha_max = alpha_max.min(DERIVATIVE_FRACTION_TO_BOUNDARY * current / -slope);
537            }
538        }
539        Ok(Some(alpha_max.clamp(0.0, 1.0)))
540    }
541
542    fn exact_newton_hessian_directional_derivative(
543        &self,
544        block_states: &[ParameterBlockState],
545        block_idx: usize,
546        d_beta: &Array1<f64>,
547    ) -> Result<Option<Array2<f64>>, String> {
548        let block = self.blocks.get(block_idx).ok_or_else(|| {
549            SurvivalError::CauseSpecificDimensionMismatch {
550                reason: format!(
551                    "cause-specific survival expected block index < {}, got {block_idx}",
552                    self.blocks.len()
553                ),
554            }
555            .to_string()
556        })?;
557        let state = block_states.get(block_idx).ok_or_else(|| {
558            SurvivalError::CauseSpecificDimensionMismatch {
559                reason: format!(
560                    "cause-specific survival expected {} block states, got {}",
561                    self.blocks.len(),
562                    block_states.len()
563                ),
564            }
565            .to_string()
566        })?;
567        Ok(Some(cause_specific_hessian_directional_derivative(
568            block,
569            &state.beta,
570            d_beta,
571        )?))
572    }
573
574    fn exact_newton_hessian_second_directional_derivative(
575        &self,
576        block_states: &[ParameterBlockState],
577        block_idx: usize,
578        d_beta_u: &Array1<f64>,
579        d_beta_v: &Array1<f64>,
580    ) -> Result<Option<Array2<f64>>, String> {
581        let block = self.blocks.get(block_idx).ok_or_else(|| {
582            SurvivalError::CauseSpecificDimensionMismatch {
583                reason: format!(
584                    "cause-specific survival expected block index < {}, got {block_idx}",
585                    self.blocks.len()
586                ),
587            }
588            .to_string()
589        })?;
590        let state = block_states.get(block_idx).ok_or_else(|| {
591            SurvivalError::CauseSpecificDimensionMismatch {
592                reason: format!(
593                    "cause-specific survival expected {} block states, got {}",
594                    self.blocks.len(),
595                    block_states.len()
596                ),
597            }
598            .to_string()
599        })?;
600        Ok(Some(cause_specific_hessian_second_directional_derivative(
601            block,
602            &state.beta,
603            d_beta_u,
604            d_beta_v,
605        )?))
606    }
607}
608
609/// The LIVE third-order tower `∂_dir H` for the cause-specific NLL: the
610/// per-predictor directional weights (`w_exit`, `w_entry`, `w_derivative`)
611/// scattered through the exit / entry / derivative designs.
612///
613/// #932 documented performance exception: this diagonal-in-(η1,η0,s) closed form
614/// (and its fourth-order sibling `cause_specific_hessian_second_directional_derivative`)
615/// STAYS in the production Newton path; it is not cut over to the generic gam-math
616/// `Tower4` jet, which would materialise a full per-row dense 4th-order tensor on
617/// this inner-solve path. It is instead pinned as a non-ignored jet oracle at
618/// ≤1e-9, with an independent central-difference witness, in
619/// `tests::jet_cause_specific_production_parity::cause_specific_live_tower_matches_jet_and_fd`.
620fn cause_specific_hessian_directional_derivative(
621    block: &CauseSpecificRoystonParmarBlock,
622    beta: &Array1<f64>,
623    d_beta: &Array1<f64>,
624) -> Result<Array2<f64>, SurvivalError> {
625    let p = block.x_exit.ncols();
626    if beta.len() != p || d_beta.len() != p {
627        return Err(SurvivalError::CauseSpecificDimensionMismatch {
628            reason: "cause-specific survival Hessian derivative dimension mismatch".to_string(),
629        });
630    }
631    let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
632    let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
633    let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
634    let d_eta_entry = fast_av(&block.x_entry, d_beta);
635    let d_eta_exit = fast_av(&block.x_exit, d_beta);
636    let d_derivative = fast_av(&block.x_derivative, d_beta);
637    let mut w_exit = Array1::<f64>::zeros(block.event_target.len());
638    let mut w_entry = Array1::<f64>::zeros(block.event_target.len());
639    let mut w_derivative = Array1::<f64>::zeros(block.event_target.len());
640
641    for i in 0..block.event_target.len() {
642        let weight = block.sampleweight[i];
643        if weight <= 0.0 {
644            continue;
645        }
646        let has_entry = block.age_entry[i] > ENTRY_AT_ORIGIN_THRESHOLD;
647        w_exit[i] = weight * eta_exit[i].exp() * d_eta_exit[i];
648        if has_entry {
649            w_entry[i] = weight * eta_entry[i].exp() * d_eta_entry[i];
650        }
651        if block.event_target[i] > 0 {
652            let deriv = derivative[i];
653            if !(deriv.is_finite() && deriv > 0.0) {
654                return Err(SurvivalError::NumericalFailure {
655                    reason: format!(
656                        "cause-specific survival derivative must be positive at row {i}, got {deriv}"
657                    ),
658                });
659            }
660            w_derivative[i] = -2.0 * weight * d_derivative[i] / (deriv * deriv * deriv);
661        }
662    }
663
664    let mut d_hessian = fast_xt_diag_x(&block.x_exit, &w_exit);
665    d_hessian -= &fast_xt_diag_x(&block.x_entry, &w_entry);
666    d_hessian += &fast_xt_diag_x(&block.x_derivative, &w_derivative);
667    Ok(d_hessian)
668}
669
670fn cause_specific_hessian_second_directional_derivative(
671    block: &CauseSpecificRoystonParmarBlock,
672    beta: &Array1<f64>,
673    d_beta_u: &Array1<f64>,
674    d_beta_v: &Array1<f64>,
675) -> Result<Array2<f64>, SurvivalError> {
676    let p = block.x_exit.ncols();
677    if beta.len() != p || d_beta_u.len() != p || d_beta_v.len() != p {
678        return Err(SurvivalError::CauseSpecificDimensionMismatch {
679            reason: "cause-specific survival second Hessian derivative dimension mismatch"
680                .to_string(),
681        });
682    }
683    let eta_entry = fast_av(&block.x_entry, beta) + &block.offset_eta_entry;
684    let eta_exit = fast_av(&block.x_exit, beta) + &block.offset_eta_exit;
685    let derivative = fast_av(&block.x_derivative, beta) + &block.offset_derivative_exit;
686    let u_eta_entry = fast_av(&block.x_entry, d_beta_u);
687    let u_eta_exit = fast_av(&block.x_exit, d_beta_u);
688    let u_derivative = fast_av(&block.x_derivative, d_beta_u);
689    let v_eta_entry = fast_av(&block.x_entry, d_beta_v);
690    let v_eta_exit = fast_av(&block.x_exit, d_beta_v);
691    let v_derivative = fast_av(&block.x_derivative, d_beta_v);
692    let mut w_exit = Array1::<f64>::zeros(block.event_target.len());
693    let mut w_entry = Array1::<f64>::zeros(block.event_target.len());
694    let mut w_derivative = Array1::<f64>::zeros(block.event_target.len());
695
696    for i in 0..block.event_target.len() {
697        let weight = block.sampleweight[i];
698        if weight <= 0.0 {
699            continue;
700        }
701        let has_entry = block.age_entry[i] > ENTRY_AT_ORIGIN_THRESHOLD;
702        w_exit[i] = weight * eta_exit[i].exp() * u_eta_exit[i] * v_eta_exit[i];
703        if has_entry {
704            w_entry[i] = weight * eta_entry[i].exp() * u_eta_entry[i] * v_eta_entry[i];
705        }
706        if block.event_target[i] > 0 {
707            let deriv = derivative[i];
708            if !(deriv.is_finite() && deriv > 0.0) {
709                return Err(SurvivalError::NumericalFailure {
710                    reason: format!(
711                        "cause-specific survival derivative must be positive at row {i}, got {deriv}"
712                    ),
713                });
714            }
715            w_derivative[i] = 6.0 * weight * u_derivative[i] * v_derivative[i] / deriv.powi(4);
716        }
717    }
718
719    let mut d2_hessian = fast_xt_diag_x(&block.x_exit, &w_exit);
720    d2_hessian -= &fast_xt_diag_x(&block.x_entry, &w_entry);
721    d2_hessian += &fast_xt_diag_x(&block.x_derivative, &w_derivative);
722    Ok(d2_hessian)
723}
724
725pub fn survival_event_code_from_value(value: f64, row_index: usize) -> Result<u8, String> {
726    const INTEGER_TOL: f64 = 1e-8;
727    const MAX_AUTO_CAUSES: u8 = 32;
728    if !value.is_finite() {
729        return Err(SurvivalError::EventCodeInvalid {
730            reason: format!(
731                "survival event value at row {} is non-finite",
732                row_index + 1
733            ),
734        }
735        .into());
736    }
737    if value < 0.0 {
738        return Err(SurvivalError::EventCodeInvalid {
739            reason: format!(
740                "survival event value at row {} is negative: {value}",
741                row_index + 1
742            ),
743        }
744        .into());
745    }
746    let rounded = value.round();
747    if (value - rounded).abs() > INTEGER_TOL {
748        return Err(SurvivalError::EventCodeInvalid {
749            reason: format!(
750                "survival event value at row {} must be an integer code with 0=censored, got {value}",
751                row_index + 1
752            ),
753        }
754        .into());
755    }
756    if rounded > f64::from(MAX_AUTO_CAUSES) {
757        return Err(SurvivalError::EventCodeInvalid {
758            reason: format!(
759                "survival event value at row {} has code {rounded}; automatic competing-risks detection supports codes 0..={MAX_AUTO_CAUSES}",
760                row_index + 1
761            ),
762        }
763        .into());
764    }
765    Ok(rounded as u8)
766}
767
768pub fn cause_count_from_event_codes(
769    event_codes: ArrayView1<'_, u8>,
770) -> Result<usize, SurvivalError> {
771    let max_code = event_codes.iter().copied().max().map_or(0, usize::from);
772    if max_code == 0 {
773        return Ok(1);
774    }
775
776    let mut present = vec![false; max_code + 1];
777    for code in event_codes.iter().copied() {
778        present[usize::from(code)] = true;
779    }
780    if (1..=max_code).any(|code| !present[code]) {
781        let actual = present
782            .iter()
783            .enumerate()
784            .skip(1)
785            .filter_map(|(code, &seen)| seen.then_some(code.to_string()))
786            .collect::<Vec<_>>()
787            .join(", ");
788        return Err(SurvivalError::EventCodeInvalid {
789            reason: format!(
790                "survival competing-risks event codes must use contiguous positive codes; observed nonzero codes are {{{actual}}}. Remap event codes contiguously (for example, {{0,1,3}} -> {{0,1,2}}), otherwise a phantom cause is fit with no events and pollutes CIF assembly."
791            ),
792        });
793    }
794
795    Ok(max_code)
796}
797
798/// Project multi-cause competing-risks event codes `{0 = censored, k = cause k}`
799/// onto the binary `{0, 1}` *any-event* indicator the pooled single-hazard
800/// baseline engine consumes.
801///
802/// The shared Royston-Parmar baseline working model fits one hazard across all
803/// causes; every observed event (regardless of cause) informs that baseline, so
804/// the indicator is `1` exactly when *any* cause occurred. This is one of the
805/// two cause-aware projections of the raw code vector; the other is
806/// [`cause_specific_event_indicator`]. Centralizing both here keeps the single
807/// source of truth for "how multi-cause labels become a single-hazard binary
808/// contract", so no construction path open-codes a fragile `mapv` and then
809/// trips the binary `event_target > 1` guard on the raw labels.
810pub fn pooled_any_event_indicator(event_codes: ArrayView1<'_, u8>) -> Array1<u8> {
811    event_codes.mapv(|label| u8::from(label > 0))
812}
813
814/// Project multi-cause competing-risks event codes `{0 = censored, k = cause k}`
815/// onto the binary `{0, 1}` indicator for the cause-specific Royston-Parmar
816/// block of cause `cause` (1-based).
817///
818/// Within cause `cause`'s block the event of interest is `event == cause`; every
819/// competing cause is treated as censoring (indicator `0`), which is exactly the
820/// cause-specific hazard likelihood. Like [`pooled_any_event_indicator`], this
821/// yields a binary contract that satisfies the single-hazard `event_target` guard
822/// — the raw multi-cause labels are never handed to a binary engine.
823pub fn cause_specific_event_indicator(event_codes: ArrayView1<'_, u8>, cause: usize) -> Array1<u8> {
824    let cause_code = cause as u8;
825    event_codes.mapv(|observed| u8::from(observed == cause_code))
826}
827
828fn compress_positive_collinear_constraints(
829    a: &Array2<f64>,
830    b: &Array1<f64>,
831) -> LinearInequalityConstraints {
832    const SCALE_TOL: f64 = 1e-14;
833    const KEY_TOL: f64 = 1e-8;
834
835    let mut grouped: BTreeMap<Vec<i64>, (Vec<f64>, f64)> = BTreeMap::new();
836    let mut fallbackrows: Vec<(Vec<f64>, f64)> = Vec::new();
837
838    for i in 0..a.nrows() {
839        let row = a.row(i);
840        let scale = row.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
841        if !scale.is_finite() || scale <= SCALE_TOL {
842            if b[i] > 0.0 {
843                fallbackrows.push((row.to_vec(), b[i]));
844            }
845            continue;
846        }
847
848        let normalizedrow: Vec<f64> = row
849            .iter()
850            .map(|&v| {
851                let scaled = v / scale;
852                if scaled.abs() <= KEY_TOL { 0.0 } else { scaled }
853            })
854            .collect();
855        let normalized_rhs = b[i] / scale;
856        let key: Vec<i64> = normalizedrow
857            .iter()
858            .map(|&v| (v / KEY_TOL).round() as i64)
859            .collect();
860
861        match grouped.get_mut(&key) {
862            Some((_, rhs_max)) => {
863                if normalized_rhs > *rhs_max {
864                    *rhs_max = normalized_rhs;
865                }
866            }
867            None => {
868                grouped.insert(key, (normalizedrow, normalized_rhs));
869            }
870        }
871    }
872
873    let nrows = grouped.len() + fallbackrows.len();
874    let n_cols = a.ncols();
875    let mut a_out = Array2::<f64>::zeros((nrows, n_cols));
876    let mut b_out = Array1::<f64>::zeros(nrows);
877
878    let mut outrow = 0usize;
879    for (_, (row, rhs)) in grouped {
880        for (j, value) in row.into_iter().enumerate() {
881            a_out[[outrow, j]] = value;
882        }
883        b_out[outrow] = rhs;
884        outrow += 1;
885    }
886    for (row, rhs) in fallbackrows {
887        for (j, value) in row.into_iter().enumerate() {
888            a_out[[outrow, j]] = value;
889        }
890        b_out[outrow] = rhs;
891        outrow += 1;
892    }
893
894    LinearInequalityConstraints { a: a_out, b: b_out }
895}
896
897#[derive(Debug, Clone, Copy, Default)]
898pub struct SurvivalMonotonicityPenalty {
899    pub tolerance: f64,
900}
901
902#[derive(Debug, Clone)]
903enum SurvivalDesign {
904    Flat {
905        x_entry: Array2<f64>,
906        x_exit: Array2<f64>,
907        x_derivative: Array2<f64>,
908    },
909    TimeCovariateShared {
910        time_entry: Array2<f64>,
911        time_exit: Array2<f64>,
912        time_derivative: Array2<f64>,
913        covariates: Array2<f64>,
914    },
915}
916
917impl SurvivalDesign {
918    fn p_total(&self) -> usize {
919        match self {
920            Self::Flat { x_exit, .. } => x_exit.ncols(),
921            Self::TimeCovariateShared {
922                time_exit,
923                covariates,
924                ..
925            } => time_exit.ncols() + covariates.ncols(),
926        }
927    }
928
929    fn design_dot(&self, time_mat: &Array2<f64>, beta: &Array1<f64>) -> Array1<f64> {
930        match self {
931            Self::Flat { .. } => time_mat.dot(beta),
932            Self::TimeCovariateShared { covariates, .. } => {
933                let p_time = time_mat.ncols();
934                let mut out = time_mat.dot(&beta.slice(ndarray::s![..p_time]));
935                if covariates.ncols() > 0 {
936                    out += &covariates.dot(&beta.slice(ndarray::s![p_time..]));
937                }
938                out
939            }
940        }
941    }
942
943    fn fill_row(&self, time_mat: &Array2<f64>, i: usize, out: &mut [f64]) {
944        match self {
945            Self::Flat { .. } => {
946                for (dst, &src) in out.iter_mut().zip(time_mat.row(i).iter()) {
947                    *dst = src;
948                }
949            }
950            Self::TimeCovariateShared { covariates, .. } => {
951                let p_time = time_mat.ncols();
952                for j in 0..p_time {
953                    out[j] = time_mat[[i, j]];
954                }
955                for j in 0..covariates.ncols() {
956                    out[p_time + j] = covariates[[i, j]];
957                }
958            }
959        }
960    }
961}
962
963/// Pre-allocated workspace buffers for `update_state` to avoid per-iteration allocations.
964#[derive(Debug, Clone)]
965struct SurvivalWorkspace {
966    w_event: Array1<f64>,
967    w_event_inv_deriv: Array1<f64>,
968    w_event_outer: Array1<f64>,
969    w_hess_exit: Array1<f64>,
970    w_hess_entry: Array1<f64>,
971}
972
973impl SurvivalWorkspace {
974    fn new(n: usize) -> Self {
975        Self {
976            w_event: Array1::zeros(n),
977            w_event_inv_deriv: Array1::zeros(n),
978            w_event_outer: Array1::zeros(n),
979            w_hess_exit: Array1::zeros(n),
980            w_hess_entry: Array1::zeros(n),
981        }
982    }
983
984    fn reset(&mut self, n: usize) {
985        if self.w_event.len() != n {
986            *self = Self::new(n);
987        } else {
988            self.w_event.fill(0.0);
989            self.w_event_inv_deriv.fill(0.0);
990            self.w_event_outer.fill(0.0);
991            self.w_hess_exit.fill(0.0);
992            self.w_hess_entry.fill(0.0);
993        }
994    }
995}
996
997/// Per-observation gradients of the unpenalized survival NLL with respect
998/// to each additive offset channel, at a given β. See
999/// [`WorkingModelSurvival::offset_channel_residuals`] for the algebra.
1000///
1001/// Contract: all four arrays have length `n` = number of observations.
1002/// Rows with non-positive sampleweight are 0 in every channel. The
1003/// `derivative` channel is 0 in all non-event rows. The `right` channel is
1004/// the interval upper-bound (`R`) η-offset sensitivity and is exactly 0 for
1005/// every NON-interval-censored model and every non-interval row of the latent
1006/// interval model (only the dedicated `SurvInterval(L, R, event)` latent fit
1007/// populates it); the baseline-θ chain rule contracts it against the
1008/// `age_right`-evaluated η-partial.
1009#[derive(Clone, Debug)]
1010pub struct OffsetChannelResiduals {
1011    /// ∂NLL/∂o_X: w·(exp(η_exit) − δ) per row.
1012    pub exit: Array1<f64>,
1013    /// ∂NLL/∂o_E: −w·exp(η_entry) if row has a positive entry interval else 0.
1014    pub entry: Array1<f64>,
1015    /// ∂NLL/∂o_D: −w·δ / s (event-row only).
1016    pub derivative: Array1<f64>,
1017    /// ∂NLL/∂o_R: interval upper-bound (`R`) η-offset sensitivity,
1018    /// `−w·∂(log-lik)/∂q_right`. Nonzero only for interval-censored latent
1019    /// rows; exactly 0 for every other channel/model.
1020    pub right: Array1<f64>,
1021}
1022
1023/// Per-observation Hessians of the unpenalized survival NLL with respect
1024/// to additive offset channels in `(entry, exit, derivative)` order.
1025#[derive(Clone, Debug)]
1026pub struct OffsetChannelCurvatures {
1027    pub rows: Vec<[[f64; 3]; 3]>,
1028}
1029
1030#[derive(Debug)]
1031pub struct WorkingModelSurvival {
1032    age_entry: Array1<f64>,
1033    age_exit: Array1<f64>,
1034    entry_at_origin: Array1<bool>,
1035    event_target: Array1<u8>,
1036    sampleweight: Array1<f64>,
1037    design: SurvivalDesign,
1038    offset_eta_entry: Array1<f64>,
1039    offset_eta_exit: Array1<f64>,
1040    offset_derivative_exit: Array1<f64>,
1041    penalties: PenaltyBlocks,
1042    monotonicity: SurvivalMonotonicityPenalty,
1043    structurally_monotonic: bool,
1044    structural_time_columns: usize,
1045    monotonicity_constraint_rows: Option<Array2<f64>>,
1046    monotonicity_constraint_offsets: Option<Array1<f64>>,
1047    workspace: std::sync::Mutex<SurvivalWorkspace>,
1048}
1049
1050impl Clone for WorkingModelSurvival {
1051    fn clone(&self) -> Self {
1052        let workspace = self.workspace.lock().unwrap().clone();
1053        Self {
1054            age_entry: self.age_entry.clone(),
1055            age_exit: self.age_exit.clone(),
1056            entry_at_origin: self.entry_at_origin.clone(),
1057            event_target: self.event_target.clone(),
1058            sampleweight: self.sampleweight.clone(),
1059            design: self.design.clone(),
1060            offset_eta_entry: self.offset_eta_entry.clone(),
1061            offset_eta_exit: self.offset_eta_exit.clone(),
1062            offset_derivative_exit: self.offset_derivative_exit.clone(),
1063            penalties: self.penalties.clone(),
1064            monotonicity: self.monotonicity,
1065            structurally_monotonic: self.structurally_monotonic,
1066            structural_time_columns: self.structural_time_columns,
1067            monotonicity_constraint_rows: self.monotonicity_constraint_rows.clone(),
1068            monotonicity_constraint_offsets: self.monotonicity_constraint_offsets.clone(),
1069            workspace: std::sync::Mutex::new(workspace),
1070        }
1071    }
1072}
1073
1074impl WorkingModelSurvival {
1075    const LOG_F64_MAX: f64 = 709.782712893384;
1076
1077    #[inline]
1078    fn scaled_exp_component(log_scale: f64, base: f64) -> Result<f64, EstimationError> {
1079        if base == 0.0 {
1080            return Ok(0.0);
1081        }
1082        let log_abs = log_scale + base.abs().ln();
1083        if !log_abs.is_finite() {
1084            crate::bail_invalid_estim!("survival interval term produced non-finite log-magnitude");
1085        }
1086        if log_abs > Self::LOG_F64_MAX {
1087            crate::bail_invalid_estim!(
1088                "survival interval term exceeds f64 range (log-magnitude={log_abs:.3e})"
1089            );
1090        }
1091        Ok(base.signum() * log_abs.exp())
1092    }
1093
1094    fn coefficient_dim(&self) -> usize {
1095        self.design.p_total()
1096    }
1097
1098    fn nrows(&self) -> usize {
1099        self.sampleweight.len()
1100    }
1101
1102    fn entry_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1103        let time_mat = match &self.design {
1104            SurvivalDesign::Flat { x_entry, .. } => x_entry,
1105            SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
1106        };
1107        self.design.design_dot(time_mat, beta)
1108    }
1109
1110    fn exit_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1111        let time_mat = match &self.design {
1112            SurvivalDesign::Flat { x_exit, .. } => x_exit,
1113            SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
1114        };
1115        self.design.design_dot(time_mat, beta)
1116    }
1117
1118    fn derivative_dot(&self, beta: &Array1<f64>) -> Array1<f64> {
1119        match &self.design {
1120            SurvivalDesign::Flat { x_derivative, .. } => x_derivative.dot(beta),
1121            SurvivalDesign::TimeCovariateShared {
1122                time_derivative, ..
1123            } => time_derivative.dot(&beta.slice(ndarray::s![..time_derivative.ncols()])),
1124        }
1125    }
1126
1127    fn fill_entry_row(&self, i: usize, out: &mut [f64]) {
1128        let time_mat = match &self.design {
1129            SurvivalDesign::Flat { x_entry, .. } => x_entry,
1130            SurvivalDesign::TimeCovariateShared { time_entry, .. } => time_entry,
1131        };
1132        self.design.fill_row(time_mat, i, out);
1133    }
1134
1135    fn fill_exit_row(&self, i: usize, out: &mut [f64]) {
1136        let time_mat = match &self.design {
1137            SurvivalDesign::Flat { x_exit, .. } => x_exit,
1138            SurvivalDesign::TimeCovariateShared { time_exit, .. } => time_exit,
1139        };
1140        self.design.fill_row(time_mat, i, out);
1141    }
1142
1143    fn fill_derivative_row(&self, i: usize, out: &mut [f64]) {
1144        match &self.design {
1145            SurvivalDesign::Flat { x_derivative, .. } => {
1146                for (dst, &src) in out.iter_mut().zip(x_derivative.row(i).iter()) {
1147                    *dst = src;
1148                }
1149            }
1150            SurvivalDesign::TimeCovariateShared {
1151                time_derivative, ..
1152            } => {
1153                let p_time = time_derivative.ncols();
1154                for j in 0..p_time {
1155                    out[j] = time_derivative[[i, j]];
1156                }
1157                for dst in out.iter_mut().skip(p_time) {
1158                    *dst = 0.0;
1159                }
1160            }
1161        }
1162    }
1163
1164    fn derivative_xt_diag_x(&self, weights: &Array1<f64>) -> Array2<f64> {
1165        match &self.design {
1166            SurvivalDesign::Flat { x_derivative, .. } => fast_xt_diag_x(x_derivative, weights),
1167            SurvivalDesign::TimeCovariateShared {
1168                time_derivative,
1169                covariates,
1170                ..
1171            } => {
1172                let p_time = time_derivative.ncols();
1173                let p_cov = covariates.ncols();
1174                let mut out = Array2::<f64>::zeros((p_time + p_cov, p_time + p_cov));
1175                let time_block = fast_xt_diag_x(time_derivative, weights);
1176                out.slice_mut(ndarray::s![..p_time, ..p_time])
1177                    .assign(&time_block);
1178                out
1179            }
1180        }
1181    }
1182
1183    /// Compute the full p×p Hessian contribution for the interval terms:
1184    ///   H = X_exit^T diag(w_exit) X_exit - X_entry^T diag(w_entry) X_entry
1185    /// using faer-accelerated BLAS on the stored design matrix blocks.
1186    fn interval_hessian_blas(&self, w_exit: &Array1<f64>, w_entry: &Array1<f64>) -> Array2<f64> {
1187        match &self.design {
1188            SurvivalDesign::Flat {
1189                x_entry, x_exit, ..
1190            } => {
1191                let mut h = fast_xt_diag_x(x_exit, w_exit);
1192                h -= &fast_xt_diag_x(x_entry, w_entry);
1193                h
1194            }
1195            SurvivalDesign::TimeCovariateShared {
1196                time_entry,
1197                time_exit,
1198                covariates,
1199                ..
1200            } => {
1201                let p_time = time_exit.ncols();
1202                let p_cov = covariates.ncols();
1203                let p = p_time + p_cov;
1204                let mut h = Array2::<f64>::zeros((p, p));
1205                // time-time block: T_exit^T W_exit T_exit - T_entry^T W_entry T_entry
1206                let tt = {
1207                    let mut block = fast_xt_diag_x(time_exit, w_exit);
1208                    block -= &fast_xt_diag_x(time_entry, w_entry);
1209                    block
1210                };
1211                h.slice_mut(ndarray::s![..p_time, ..p_time]).assign(&tt);
1212                if p_cov > 0 {
1213                    // time-cov block: T_exit^T W_exit C - T_entry^T W_entry C
1214                    let tc = {
1215                        let mut block = fast_xt_diag_y(time_exit, w_exit, covariates);
1216                        block -= &fast_xt_diag_y(time_entry, w_entry, covariates);
1217                        block
1218                    };
1219                    h.slice_mut(ndarray::s![..p_time, p_time..]).assign(&tc);
1220                    h.slice_mut(ndarray::s![p_time.., ..p_time]).assign(&tc.t());
1221                    // cov-cov block: C^T (W_exit - W_entry) C
1222                    let w_diff = w_exit - w_entry;
1223                    let cc = fast_xt_diag_x(covariates, &w_diff);
1224                    h.slice_mut(ndarray::s![p_time.., p_time..]).assign(&cc);
1225                }
1226                h
1227            }
1228        }
1229    }
1230
1231    /// Clamped structural derivative: `max(deriv, floor)` for derivatives
1232    /// above the roundoff tolerance, `None` outside structural monotonicity
1233    /// or for genuinely negative derivatives.
1234    ///
1235    /// Returns `(value, slope)` where `slope` is the exact derivative of the
1236    /// clamp itself: 1 on the identity branch, 0 on the floored branch.
1237    /// Every consumer that differentiates through the structural derivative
1238    /// MUST scale its derivative-channel terms by `slope` — the floored
1239    /// branch is locally constant in β, so gradients and Hessians of
1240    /// `ln(value)` and `1/value` are exactly zero there. Emitting the
1241    /// reciprocal-scale terms of the UNclamped expression (≈1e12 gradient,
1242    /// ≈1e24 curvature at deriv = 5e-13) against a value branch that is flat
1243    /// desynchronizes the objective from its derivatives.
1244    fn stabilized_structural_derivative(&self, deriv: f64) -> Option<(f64, f64)> {
1245        const STRUCTURAL_MONO_ROUNDOFF_TOL: f64 = 1e-7;
1246        const STRUCTURAL_DERIV_FLOOR: f64 = 1e-12;
1247        if !self.structurally_monotonic {
1248            return None;
1249        }
1250        if deriv >= STRUCTURAL_DERIV_FLOOR {
1251            return Some((deriv, 1.0));
1252        }
1253        if deriv >= -STRUCTURAL_MONO_ROUNDOFF_TOL {
1254            return Some((STRUCTURAL_DERIV_FLOOR, 0.0));
1255        }
1256        None
1257    }
1258
1259    fn validate_penalties(
1260        penalties: &PenaltyBlocks,
1261        coefficient_dim: usize,
1262    ) -> Result<(), SurvivalError> {
1263        for block in &penalties.blocks {
1264            if !block.lambda.is_finite() || block.lambda < 0.0 {
1265                return Err(SurvivalError::NonFiniteInput);
1266            }
1267            if block.range.start > block.range.end || block.range.end > coefficient_dim {
1268                return Err(SurvivalError::DimensionMismatch);
1269            }
1270            let block_dim = block.range.end - block.range.start;
1271            if block.matrix.nrows() != block_dim || block.matrix.ncols() != block_dim {
1272                return Err(SurvivalError::DimensionMismatch);
1273            }
1274            if block.matrix.iter().any(|v| !v.is_finite()) {
1275                return Err(SurvivalError::NonFiniteInput);
1276            }
1277        }
1278        Ok(())
1279    }
1280
1281    fn derivative_guard(&self) -> f64 {
1282        if self.structurally_monotonic {
1283            // I-spline basis is monotone by construction when coefficients ≥ 0.
1284            // A derivative of zero (flat hazard) is valid, so the guard only
1285            // rejects genuinely negative derivatives from floating-point noise.
1286            return 0.0;
1287        }
1288        self.monotonicity.tolerance.max(0.0)
1289    }
1290
1291    fn derivative_guard_numerical(&self) -> f64 {
1292        let derivative_guard = self.derivative_guard();
1293        if derivative_guard <= 0.0 {
1294            // For structural monotonicity (guard = 0), tiny negative derivs are
1295            // tolerated because `stabilized_structural_derivative` lifts the
1296            // value back to a small positive floor before any `ln`/`1/deriv`
1297            // use. For *non-structural* monotonicity with tolerance == 0 the
1298            // raw derivative flows straight through into the event-row
1299            // `deriv.ln()` and `1.0 / deriv`, so any non-positive value would
1300            // produce NaN / huge negative weights. Keep the slack only when
1301            // the structural stabilizer is active.
1302            if self.structurally_monotonic {
1303                -1e-10
1304            } else {
1305                1e-12
1306            }
1307        } else {
1308            (derivative_guard - (1e-10_f64).min(0.01 * derivative_guard)).max(1e-12)
1309        }
1310    }
1311
1312    fn interval_increment_guard(&self, h_entry: f64, h_exit: f64) -> f64 {
1313        let scale = h_entry.abs().max(h_exit.abs()).max(1.0);
1314        1e-10 * scale
1315    }
1316
1317    fn structural_time_coefficient_constraints(&self) -> Option<LinearInequalityConstraints> {
1318        if !self.structurally_monotonic {
1319            return None;
1320        }
1321        let p = self.coefficient_dim();
1322        let time_columns = self.structural_time_columns.min(p);
1323        if time_columns == 0 {
1324            return None;
1325        }
1326        const STRUCTURAL_DERIV_TOL: f64 = 1e-12;
1327        let mut active_columns = vec![false; time_columns];
1328        let mut derivative_row = vec![0.0_f64; p];
1329        for i in 0..self.nrows() {
1330            if self.sampleweight[i] <= 0.0 {
1331                continue;
1332            }
1333            self.fill_derivative_row(i, &mut derivative_row);
1334            for j in 0..time_columns {
1335                if derivative_row[j] > STRUCTURAL_DERIV_TOL {
1336                    active_columns[j] = true;
1337                }
1338            }
1339        }
1340        if let Some(rows) = self.monotonicity_constraint_rows.as_ref() {
1341            for i in 0..rows.nrows() {
1342                for j in 0..time_columns {
1343                    if rows[[i, j]] > STRUCTURAL_DERIV_TOL {
1344                        active_columns[j] = true;
1345                    }
1346                }
1347            }
1348        }
1349        let active_columns: Vec<usize> = active_columns
1350            .into_iter()
1351            .enumerate()
1352            .filter_map(|(j, active)| active.then_some(j))
1353            .collect();
1354        if active_columns.is_empty() {
1355            return None;
1356        }
1357        let mut a = Array2::<f64>::zeros((active_columns.len(), p));
1358        let b = Array1::<f64>::zeros(active_columns.len());
1359        for (row, &col) in active_columns.iter().enumerate() {
1360            a[[row, col]] = 1.0;
1361        }
1362        Some(LinearInequalityConstraints { a, b })
1363    }
1364
1365    pub fn monotonicity_linear_constraints(&self) -> Option<LinearInequalityConstraints> {
1366        let p = self.coefficient_dim();
1367        const DERIVATIVE_ROW_NORM_TOL: f64 = 1e-12;
1368        if p == 0 {
1369            return None;
1370        }
1371        if self.structurally_monotonic {
1372            return self.structural_time_coefficient_constraints();
1373        }
1374        if let (Some(rows), Some(offsets)) = (
1375            self.monotonicity_constraint_rows.as_ref(),
1376            self.monotonicity_constraint_offsets.as_ref(),
1377        ) {
1378            let activerows: Vec<usize> = (0..rows.nrows())
1379                .filter(|&i| {
1380                    rows.row(i).iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()))
1381                        > DERIVATIVE_ROW_NORM_TOL
1382                })
1383                .collect();
1384            if activerows.is_empty() {
1385                return None;
1386            }
1387            let mut a = Array2::<f64>::zeros((activerows.len(), p));
1388            let mut b = Array1::<f64>::zeros(activerows.len());
1389            for (r, &i) in activerows.iter().enumerate() {
1390                a.row_mut(r).assign(&rows.row(i));
1391                b[r] = self.derivative_guard() - offsets[i];
1392            }
1393            return Some(compress_positive_collinear_constraints(&a, &b));
1394        }
1395        None
1396    }
1397
1398    pub fn from_engine_inputs(
1399        inputs: SurvivalEngineInputs<'_>,
1400        penalties: PenaltyBlocks,
1401        monotonicity: SurvivalMonotonicityPenalty,
1402        spec: SurvivalSpec,
1403    ) -> Result<Self, SurvivalError> {
1404        Self::from_engine_inputswith_offsets(inputs, None, penalties, monotonicity, spec)
1405    }
1406
1407    fn validate_offsets(
1408        offsets: Option<SurvivalBaselineOffsets<'_>>,
1409        n: usize,
1410    ) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>), SurvivalError> {
1411        if let Some(off) = offsets {
1412            if off.eta_entry.len() != n || off.eta_exit.len() != n || off.derivative_exit.len() != n
1413            {
1414                return Err(SurvivalError::DimensionMismatch);
1415            }
1416            if off.eta_entry.iter().any(|v| !v.is_finite())
1417                || off.eta_exit.iter().any(|v| !v.is_finite())
1418                || off.derivative_exit.iter().any(|v| !v.is_finite())
1419            {
1420                return Err(SurvivalError::NonFiniteInput);
1421            }
1422            Ok((
1423                off.eta_entry.to_owned(),
1424                off.eta_exit.to_owned(),
1425                off.derivative_exit.to_owned(),
1426            ))
1427        } else {
1428            Ok((Array1::zeros(n), Array1::zeros(n), Array1::zeros(n)))
1429        }
1430    }
1431
1432    fn validate_common_inputs(
1433        age_entry: &ArrayView1<f64>,
1434        age_exit: &ArrayView1<f64>,
1435        event_target: &ArrayView1<u8>,
1436        event_competing: &ArrayView1<u8>,
1437        sampleweight: &ArrayView1<f64>,
1438    ) -> Result<(), SurvivalError> {
1439        if age_entry.iter().any(|v| !v.is_finite())
1440            || age_exit.iter().any(|v| !v.is_finite())
1441            || sampleweight.iter().any(|v| !v.is_finite() || *v < 0.0)
1442        {
1443            return Err(SurvivalError::NonFiniteInput);
1444        }
1445        // The single-hazard engine's `event_target` contract is binary {0, 1}.
1446        // A code > 1 is a *valid finite multi-cause label* that simply must be
1447        // projected first (any-event for the pooled baseline, cause-specific for
1448        // each block); it is NOT a non-finite input. Report it as such so the
1449        // failure is actionable and never surfaces as the misleading "inputs
1450        // contain non-finite values".
1451        if let Some(&label) = event_target.iter().find(|&&v| v > 1) {
1452            return Err(SurvivalError::EventCodeInvalid {
1453                reason: format!(
1454                    "single-hazard survival engine requires a binary {{0, 1}} event_target, got multi-cause label {label}; competing-risks codes must be projected via pooled_any_event_indicator / cause_specific_event_indicator before construction"
1455                ),
1456            });
1457        }
1458        if let Some(&label) = event_competing.iter().find(|&&v| v > 1) {
1459            return Err(SurvivalError::EventCodeInvalid {
1460                reason: format!(
1461                    "single-hazard survival engine requires a binary {{0, 1}} event_competing, got multi-cause label {label}"
1462                ),
1463            });
1464        }
1465        if event_target
1466            .iter()
1467            .zip(event_competing.iter())
1468            .any(|(&target, &competing)| target > 0 && competing > 0)
1469        {
1470            return Err(SurvivalError::EventCodeInvalid {
1471                reason: "a row cannot be simultaneously a target event and a competing event"
1472                    .to_string(),
1473            });
1474        }
1475        // The "must have at least one target event" requirement is a
1476        // *fittability* check, not a structural one: with all rows censored the
1477        // likelihood has no event score, so any subsequent fit cannot identify
1478        // the hazard and the optimizer spins on a flat landscape.  But the
1479        // structural integrity of the engine — its derivative-guard rejection
1480        // of decreasing cumulative hazards, its monotonicity-collocation
1481        // bookkeeping, its update_state numerics — is well-defined on
1482        // all-censored inputs, and unit tests legitimately exercise those
1483        // structural paths on censored fixtures.  Move the fittability check
1484        // out of construction; production fit dispatchers (e.g.
1485        // `solver::fit_orchestration::materialize_survival`) enforce it on the
1486        // single chokepoint that actually starts an optimization, where
1487        // the failure mode it guards against is reachable.
1488        if age_entry
1489            .iter()
1490            .zip(age_exit.iter())
1491            .any(|(&entry, &exit)| entry < 0.0 || exit <= 0.0)
1492        {
1493            return Err(SurvivalError::NonFiniteInput);
1494        }
1495        Ok::<(), _>(())
1496    }
1497
1498    fn validate_monotonicity_constraints(
1499        rows: Option<ArrayView2<'_, f64>>,
1500        offsets: Option<ArrayView1<'_, f64>>,
1501        coefficient_dim: usize,
1502    ) -> Result<(Option<Array2<f64>>, Option<Array1<f64>>), SurvivalError> {
1503        match (rows, offsets) {
1504            (None, None) => Ok((None, None)),
1505            (Some(rows), Some(offsets)) => {
1506                if rows.ncols() != coefficient_dim
1507                    || rows.nrows() != offsets.len()
1508                    || rows.iter().any(|v| !v.is_finite())
1509                    || offsets.iter().any(|v| !v.is_finite())
1510                {
1511                    return Err(SurvivalError::DimensionMismatch);
1512                }
1513                Ok((Some(rows.to_owned()), Some(offsets.to_owned())))
1514            }
1515            _ => Err(SurvivalError::DimensionMismatch),
1516        }
1517    }
1518
1519    fn finish_construction(
1520        age_entry: ArrayView1<f64>,
1521        age_exit: ArrayView1<f64>,
1522        event_target: ArrayView1<u8>,
1523        sampleweight: ArrayView1<f64>,
1524        design: SurvivalDesign,
1525        offset_eta_entry: Array1<f64>,
1526        offset_eta_exit: Array1<f64>,
1527        offset_derivative_exit: Array1<f64>,
1528        penalties: PenaltyBlocks,
1529        monotonicity: SurvivalMonotonicityPenalty,
1530        monotonicity_constraint_rows: Option<Array2<f64>>,
1531        monotonicity_constraint_offsets: Option<Array1<f64>>,
1532    ) -> Self {
1533        let n = age_entry.len();
1534        Self {
1535            age_entry: age_entry.to_owned(),
1536            age_exit: age_exit.to_owned(),
1537            entry_at_origin: age_entry.mapv(|t| t <= ENTRY_AT_ORIGIN_THRESHOLD),
1538            event_target: event_target.to_owned(),
1539            sampleweight: sampleweight.to_owned(),
1540            design,
1541            offset_eta_entry,
1542            offset_eta_exit,
1543            offset_derivative_exit,
1544            penalties,
1545            monotonicity,
1546            structurally_monotonic: false,
1547            structural_time_columns: 0,
1548            monotonicity_constraint_rows,
1549            monotonicity_constraint_offsets,
1550            workspace: std::sync::Mutex::new(SurvivalWorkspace::new(n)),
1551        }
1552    }
1553
1554    pub fn from_engine_inputswith_offsets(
1555        inputs: SurvivalEngineInputs<'_>,
1556        offsets: Option<SurvivalBaselineOffsets<'_>>,
1557        penalties: PenaltyBlocks,
1558        monotonicity: SurvivalMonotonicityPenalty,
1559        spec: SurvivalSpec,
1560    ) -> Result<Self, SurvivalError> {
1561        if spec == SurvivalSpec::Crude {
1562            return Err(SurvivalError::UnsupportedSpec("crude"));
1563        }
1564        let n = inputs.age_entry.len();
1565        let p = inputs.x_entry.ncols();
1566        if inputs.age_exit.len() != n
1567            || inputs.event_target.len() != n
1568            || inputs.event_competing.len() != n
1569            || inputs.sampleweight.len() != n
1570            || inputs.x_entry.nrows() != n
1571            || inputs.x_exit.nrows() != n
1572            || inputs.x_derivative.nrows() != n
1573            || inputs.x_entry.ncols() != inputs.x_exit.ncols()
1574            || inputs.x_entry.ncols() != inputs.x_derivative.ncols()
1575        {
1576            return Err(SurvivalError::DimensionMismatch);
1577        }
1578        Self::validate_penalties(&penalties, p)?;
1579        Self::validate_common_inputs(
1580            &inputs.age_entry,
1581            &inputs.age_exit,
1582            &inputs.event_target,
1583            &inputs.event_competing,
1584            &inputs.sampleweight,
1585        )?;
1586        if inputs.x_entry.iter().any(|v| !v.is_finite())
1587            || inputs.x_exit.iter().any(|v| !v.is_finite())
1588            || inputs.x_derivative.iter().any(|v| !v.is_finite())
1589        {
1590            return Err(SurvivalError::NonFiniteInput);
1591        }
1592        let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
1593            Self::validate_offsets(offsets, n)?;
1594        let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
1595            Self::validate_monotonicity_constraints(
1596                inputs.monotonicity_constraint_rows,
1597                inputs.monotonicity_constraint_offsets,
1598                p,
1599            )?;
1600
1601        Ok(Self::finish_construction(
1602            inputs.age_entry,
1603            inputs.age_exit,
1604            inputs.event_target,
1605            inputs.sampleweight,
1606            SurvivalDesign::Flat {
1607                x_entry: inputs.x_entry.to_owned(),
1608                x_exit: inputs.x_exit.to_owned(),
1609                x_derivative: inputs.x_derivative.to_owned(),
1610            },
1611            offset_eta_entry,
1612            offset_eta_exit,
1613            offset_derivative_exit,
1614            penalties,
1615            monotonicity,
1616            monotonicity_constraint_rows,
1617            monotonicity_constraint_offsets,
1618        ))
1619    }
1620
1621    pub fn from_time_covariate_inputswith_offsets(
1622        inputs: SurvivalTimeCovarInputs<'_>,
1623        offsets: Option<SurvivalBaselineOffsets<'_>>,
1624        penalties: PenaltyBlocks,
1625        monotonicity: SurvivalMonotonicityPenalty,
1626        spec: SurvivalSpec,
1627    ) -> Result<Self, SurvivalError> {
1628        if spec == SurvivalSpec::Crude {
1629            return Err(SurvivalError::UnsupportedSpec("crude"));
1630        }
1631        let n = inputs.age_entry.len();
1632        let p_time = inputs.time_entry.ncols();
1633        let p_cov = inputs.covariates.ncols();
1634        let p = p_time + p_cov;
1635        if inputs.age_exit.len() != n
1636            || inputs.event_target.len() != n
1637            || inputs.event_competing.len() != n
1638            || inputs.sampleweight.len() != n
1639            || inputs.time_entry.nrows() != n
1640            || inputs.time_exit.nrows() != n
1641            || inputs.time_derivative.nrows() != n
1642            || inputs.covariates.nrows() != n
1643            || inputs.time_entry.ncols() != inputs.time_exit.ncols()
1644            || inputs.time_entry.ncols() != inputs.time_derivative.ncols()
1645        {
1646            return Err(SurvivalError::DimensionMismatch);
1647        }
1648        Self::validate_penalties(&penalties, p)?;
1649        Self::validate_common_inputs(
1650            &inputs.age_entry,
1651            &inputs.age_exit,
1652            &inputs.event_target,
1653            &inputs.event_competing,
1654            &inputs.sampleweight,
1655        )?;
1656        if inputs.time_entry.iter().any(|v| !v.is_finite())
1657            || inputs.time_exit.iter().any(|v| !v.is_finite())
1658            || inputs.time_derivative.iter().any(|v| !v.is_finite())
1659            || inputs.covariates.iter().any(|v| !v.is_finite())
1660        {
1661            return Err(SurvivalError::NonFiniteInput);
1662        }
1663        let (offset_eta_entry, offset_eta_exit, offset_derivative_exit) =
1664            Self::validate_offsets(offsets, n)?;
1665        let (monotonicity_constraint_rows, monotonicity_constraint_offsets) =
1666            Self::validate_monotonicity_constraints(
1667                inputs.monotonicity_constraint_rows,
1668                inputs.monotonicity_constraint_offsets,
1669                p,
1670            )?;
1671
1672        Ok(Self::finish_construction(
1673            inputs.age_entry,
1674            inputs.age_exit,
1675            inputs.event_target,
1676            inputs.sampleweight,
1677            SurvivalDesign::TimeCovariateShared {
1678                time_entry: inputs.time_entry.to_owned(),
1679                time_exit: inputs.time_exit.to_owned(),
1680                time_derivative: inputs.time_derivative.to_owned(),
1681                covariates: inputs.covariates.to_owned(),
1682            },
1683            offset_eta_entry,
1684            offset_eta_exit,
1685            offset_derivative_exit,
1686            penalties,
1687            monotonicity,
1688            monotonicity_constraint_rows,
1689            monotonicity_constraint_offsets,
1690        ))
1691    }
1692
1693    /// Enable/disable monotonic time-block enforcement metadata.
1694    ///
1695    /// Monotonicity is enforced through linear inequality constraints on the
1696    /// derivative design; enabling this records how many leading time columns
1697    /// belong to that constrained block.
1698    /// Overwrite the per-block smoothing parameters `λ_k` in place.
1699    ///
1700    /// Used by the REML smoothing-parameter selection for transformation
1701    /// survival fits (issue #563): the outer optimizer proposes a `ρ = log λ`
1702    /// vector, sets the smoothing blocks' `λ_k` here, and re-runs the inner
1703    /// constrained PIRLS, so the monotone I-spline baseline can adapt its
1704    /// wiggliness instead of being pinned at a fixed seed. `lambdas` must have
1705    /// one entry per penalty block.
1706    pub fn set_penalty_lambdas(&mut self, lambdas: &[f64]) -> Result<(), EstimationError> {
1707        if lambdas.len() != self.penalties.blocks.len() {
1708            crate::bail_invalid_estim!(
1709                "set_penalty_lambdas expects {} lambdas, got {}",
1710                self.penalties.blocks.len(),
1711                lambdas.len()
1712            );
1713        }
1714        for (block, &lambda) in self.penalties.blocks.iter_mut().zip(lambdas.iter()) {
1715            if !lambda.is_finite() || lambda < 0.0 {
1716                crate::bail_invalid_estim!("penalty lambda must be finite and >= 0, got {lambda}");
1717            }
1718            block.lambda = lambda;
1719        }
1720        Ok(())
1721    }
1722
1723    pub fn set_structural_monotonicity(
1724        &mut self,
1725        enabled: bool,
1726        time_columns: usize,
1727    ) -> Result<(), EstimationError> {
1728        let p = self.coefficient_dim();
1729        if time_columns > p {
1730            crate::bail_invalid_estim!(
1731                "structural time columns {} exceed coefficient dimension {}",
1732                time_columns,
1733                p
1734            );
1735        }
1736        if enabled && time_columns == 0 {
1737            crate::bail_invalid_estim!("structural monotonicity requires at least one time column");
1738        }
1739        if enabled {
1740            const STRUCTURAL_DERIV_TOL: f64 = 1e-12;
1741            for (i, &offset) in self.offset_derivative_exit.iter().enumerate() {
1742                if offset < -STRUCTURAL_DERIV_TOL {
1743                    crate::bail_invalid_estim!(
1744                        "structural monotonicity requires nonnegative derivative offsets; found offset_derivative_exit[{i}]={offset:.3e}"
1745                    );
1746                }
1747            }
1748            let mut derivative_row = vec![0.0_f64; p];
1749            for i in 0..self.nrows() {
1750                self.fill_derivative_row(i, &mut derivative_row);
1751                for j in 0..time_columns {
1752                    let v = derivative_row[j];
1753                    if v < -STRUCTURAL_DERIV_TOL {
1754                        crate::bail_invalid_estim!(
1755                            "structural monotonicity requires nonnegative time-derivative basis entries; found x_derivative[{i},{j}]={v:.3e}"
1756                        );
1757                    }
1758                }
1759                for j in time_columns..p {
1760                    let v = derivative_row[j];
1761                    if v.abs() > STRUCTURAL_DERIV_TOL {
1762                        crate::bail_invalid_estim!(
1763                            "structural monotonicity requires zero derivative contribution outside the time block; found x_derivative[{i},{j}]={v:.3e}"
1764                        );
1765                    }
1766                }
1767            }
1768            if let (Some(rows), Some(offsets)) = (
1769                self.monotonicity_constraint_rows.as_ref(),
1770                self.monotonicity_constraint_offsets.as_ref(),
1771            ) {
1772                for (i, &offset) in offsets.iter().enumerate() {
1773                    if offset < -STRUCTURAL_DERIV_TOL {
1774                        crate::bail_invalid_estim!(
1775                            "structural monotonicity requires nonnegative collocation derivative offsets; found monotonicity_constraint_offsets[{i}]={offset:.3e}"
1776                        );
1777                    }
1778                }
1779                for i in 0..rows.nrows() {
1780                    for j in 0..time_columns {
1781                        let v = rows[[i, j]];
1782                        if v < -STRUCTURAL_DERIV_TOL {
1783                            crate::bail_invalid_estim!(
1784                                "structural monotonicity requires nonnegative collocation derivative basis entries; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
1785                            );
1786                        }
1787                    }
1788                    for j in time_columns..p {
1789                        let v = rows[[i, j]];
1790                        if v.abs() > STRUCTURAL_DERIV_TOL {
1791                            crate::bail_invalid_estim!(
1792                                "structural monotonicity requires zero collocation derivative contribution outside the time block; found monotonicity_constraint_rows[{i},{j}]={v:.3e}"
1793                            );
1794                        }
1795                    }
1796                }
1797            }
1798        }
1799        self.structurally_monotonic = enabled;
1800        self.structural_time_columns = if enabled { time_columns } else { 0 };
1801        Ok(())
1802    }
1803
1804    pub fn update_state(&self, beta: &Array1<f64>) -> Result<WorkingState, EstimationError> {
1805        if beta.len() != self.coefficient_dim() {
1806            crate::bail_invalid_estim!("survival beta dimension mismatch");
1807        }
1808
1809        let n = self.nrows();
1810        let p = self.coefficient_dim();
1811
1812        // Royston-Parmar contract used throughout the engine:
1813        //   eta(t) = log(H(t)), where H(t) is cumulative hazard.
1814        //
1815        // With row-vectors (per subject i):
1816        //   a1_i^T := x_exit_i^T,  a0_i^T := x_entry_i^T,  d_i^T := x_derivative_i^T
1817        // and scalars:
1818        //   eta1_i = a1_i^T beta,  eta0_i = a0_i^T beta,  s_i = d_i^T beta.
1819        //
1820        // The per-subject negative log-likelihood used below is
1821        //   NLL_i(beta) = exp(eta1_i) - exp(eta0_i) - delta_i * (eta1_i + log(s_i)),
1822        // with delta_i = event_target_i.
1823        //
1824        // This is exactly the form whose derivatives are:
1825        //   grad_i = exp(eta1_i) a1_i - exp(eta0_i) a0_i - delta_i * (a1_i + d_i / s_i)
1826        //   Hess_i = exp(eta1_i) a1_i a1_i^T - exp(eta0_i) a0_i a0_i^T
1827        //            + delta_i * (d_i d_i^T) / s_i^2.
1828        //
1829        // Monotonicity is enforced through linear inequality constraints on the
1830        // derivative design. This keeps the baseline smoothing penalty on the
1831        // actual spline coefficients and preserves zero-deviation as beta=0.
1832        //
1833        // The loop below computes exact beta-space derivatives and then adds penalties.
1834        // Total predictor = target offset + learned deviation.
1835        // This is the same architecture used for flexible binary links:
1836        // principled default, plus penalized wiggle/deviation.
1837        let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
1838        let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
1839        let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
1840
1841        let mut nll = 0.0;
1842        let derivative_guard = self.derivative_guard();
1843        let derivative_guard_numerical = self.derivative_guard_numerical();
1844        let mut workspace = self.workspace.lock().unwrap();
1845        workspace.reset(n);
1846        let SurvivalWorkspace {
1847            w_event,
1848            w_event_inv_deriv,
1849            w_event_outer,
1850            w_hess_exit,
1851            w_hess_entry,
1852        } = &mut *workspace;
1853
1854        // Phase 1: Scalar loop — compute per-observation weights, NLL, validation.
1855        for i in 0..n {
1856            let w = self.sampleweight[i];
1857            if w <= 0.0 {
1858                continue;
1859            }
1860            let entry_age = self.age_entry[i];
1861            let exit_age = self.age_exit[i];
1862            if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
1863                crate::bail_invalid_estim!(
1864                    "survival ages must be finite with age_exit >= age_entry"
1865                );
1866            }
1867            let d = f64::from(self.event_target[i]);
1868
1869            let has_entry_interval = !self.entry_at_origin[i];
1870            let interval_scale = if has_entry_interval {
1871                eta_exit[i].max(eta_entry[i])
1872            } else {
1873                eta_exit[i]
1874            };
1875            let h_e_scaled = (eta_exit[i] - interval_scale).exp();
1876            let h_s_scaled = if has_entry_interval {
1877                (eta_entry[i] - interval_scale).exp()
1878            } else {
1879                0.0
1880            };
1881            let interval_scaled = h_e_scaled - h_s_scaled;
1882            let interval = Self::scaled_exp_component(interval_scale, interval_scaled)?;
1883            let (deriv, deriv_slope) = self
1884                .stabilized_structural_derivative(derivative_raw[i])
1885                .unwrap_or((derivative_raw[i], 1.0));
1886            // Monotonicity of η(t) = log H(t) is a structural property of the
1887            // whole Royston-Parmar spline. If d_eta/dt is *strictly negative*
1888            // at any observed exit time, the cumulative hazard H(t) decreases
1889            // there and S(t) is not a valid survival function — both event
1890            // and censored rows have to refuse that case. Event rows further
1891            // need deriv strictly above the numerical guard because their
1892            // NLL contains `deriv.ln()` and `1.0 / deriv`; censored rows do
1893            // not, so a boundary value of exactly zero is feasible there.
1894            let mono_floor = if d > 0.0 {
1895                derivative_guard_numerical
1896            } else {
1897                0.0
1898            };
1899            if !deriv.is_finite() || deriv < mono_floor {
1900                return Err(EstimationError::ParameterConstraintViolation(format!(
1901                    "survival monotonicity violated at row {}: d_eta/dt={:.3e} <= tolerance={:.3e}",
1902                    i, deriv, derivative_guard
1903                )));
1904            }
1905            if has_entry_interval {
1906                let increment_guard = self.interval_increment_guard(h_s_scaled, h_e_scaled);
1907                if interval_scaled + increment_guard < 0.0 {
1908                    return Err(EstimationError::ParameterConstraintViolation(format!(
1909                        "survival cumulative hazard decreased over row {}: H(exit)-H(entry)={:.6e}",
1910                        i, interval
1911                    )));
1912                }
1913            }
1914            nll += w * interval;
1915
1916            // Per-observation weights for BLAS phase.
1917            // scaled_exp_component(interval_scale, h_e_scaled * x[r]) = exp(interval_scale) * h_e_scaled * x[r]
1918            // so the Hessian weight is w * exp(interval_scale) * h_e_scaled = w * exp(eta_exit).
1919            let w_exit_i = w * eta_exit[i].exp();
1920            let w_entry_i = if has_entry_interval {
1921                w * eta_entry[i].exp()
1922            } else {
1923                0.0
1924            };
1925            if !w_exit_i.is_finite() {
1926                crate::bail_invalid_estim!(
1927                    "survival interval term exceeds f64 range at row {i} (w*exp(eta_exit)={w_exit_i:.3e})"
1928                );
1929            }
1930            w_hess_exit[i] = w_exit_i;
1931            w_hess_entry[i] = w_entry_i;
1932
1933            if d > 0.0 {
1934                // `deriv_slope` is the derivative of the structural clamp: on
1935                // the floored branch the NLL's `ln(deriv)` term is locally
1936                // constant in β, so its score and curvature channels vanish.
1937                let inv_deriv = deriv_slope / deriv;
1938                nll += -w * (eta_exit[i] + deriv.ln());
1939                w_event[i] = w;
1940                w_event_inv_deriv[i] = w * inv_deriv;
1941                w_event_outer[i] = w * inv_deriv * inv_deriv;
1942            }
1943        }
1944
1945        // Phase 2: BLAS-accelerated Hessian and gradient via faer.
1946        //   H_interval = X_exit^T diag(w_exit) X_exit - X_entry^T diag(w_entry) X_entry
1947        //   grad_interval = X_exit^T w_exit - X_entry^T w_entry
1948        let mut h = self.interval_hessian_blas(w_hess_exit, w_hess_entry);
1949        // At large smoothing penalties the event-Jacobian score nearly cancels
1950        // the interval score. Compensated row accumulation keeps the final KKT
1951        // residual accurate enough for the outer LAML envelope check.
1952        let mut grad = Array1::<f64>::zeros(p);
1953        let mut grad_comp = Array1::<f64>::zeros(p);
1954        let mut row_exit = vec![0.0_f64; p];
1955        let mut row_entry = vec![0.0_f64; p];
1956        let mut row_derivative = vec![0.0_f64; p];
1957        for i in 0..n {
1958            let w_interval_exit = w_hess_exit[i];
1959            let w_interval_entry = w_hess_entry[i];
1960            let w_event_exit = w_event[i];
1961            let w_event_derivative = w_event_inv_deriv[i];
1962            if w_interval_exit == 0.0
1963                && w_interval_entry == 0.0
1964                && w_event_exit == 0.0
1965                && w_event_derivative == 0.0
1966            {
1967                continue;
1968            }
1969            self.fill_exit_row(i, &mut row_exit);
1970            self.fill_entry_row(i, &mut row_entry);
1971            self.fill_derivative_row(i, &mut row_derivative);
1972            for j in 0..p {
1973                let contribution = w_interval_exit * row_exit[j]
1974                    - w_interval_entry * row_entry[j]
1975                    - w_event_exit * row_exit[j]
1976                    - w_event_derivative * row_derivative[j];
1977                let t = grad[j] + contribution;
1978                if grad[j].abs() >= contribution.abs() {
1979                    grad_comp[j] += (grad[j] - t) + contribution;
1980                } else {
1981                    grad_comp[j] += (contribution - t) + grad[j];
1982                }
1983                grad[j] = t;
1984            }
1985        }
1986        grad += &grad_comp;
1987
1988        h += &self.derivative_xt_diag_x(w_event_outer);
1989
1990        // Norm of the unpenalized score, captured before adding the penalty
1991        // contribution, for the scale-invariant convergence certificate
1992        // (||score||_2 + ||S*beta||_2).
1993        let score_norm = array1_l2_norm(&grad);
1994
1995        let penaltygrad = self.penalties.gradient(beta);
1996        let penalty_dev = self.penalties.deviance(beta);
1997        let penaltygrad_norm = array1_l2_norm(&penaltygrad);
1998
1999        let mut totalgrad = grad;
2000        totalgrad += &penaltygrad;
2001
2002        self.penalties.addhessian_inplace(&mut h);
2003        // No coefficient ridge is fused into this objective. Indefinite or
2004        // rank-deficient curvature along the Newton path is the SOLVER's
2005        // problem: the working-model driver applies Levenberg–Marquardt
2006        // damping (H + λD²)δ = −g that vanishes at convergence, so the
2007        // converged estimator is a stationary point of the exact penalized
2008        // likelihood and is invariant under coefficient rescaling.
2009        let log_likelihood = -nll;
2010        let deviance = 2.0 * nll;
2011
2012        Ok(WorkingState {
2013            eta: LinearPredictor::new(eta_exit),
2014            gradient: totalgrad,
2015            hessian: gam_linalg::matrix::SymmetricMatrix::Dense(h),
2016            log_likelihood,
2017            deviance,
2018            penalty_term: penalty_dev,
2019            firth: gam_solve::pirls::FirthDiagnostics::Inactive,
2020            ridge_used: 0.0,
2021            hessian_curvature: gam_solve::pirls::HessianCurvatureKind::Observed,
2022            gradient_natural_scale: score_norm + penaltygrad_norm,
2023        })
2024    }
2025
2026    /// Compute the third-derivative correction matrix for a given mode response `u_k`.
2027    ///
2028    /// This is the directional derivative of the unpenalized NLL Hessian w.r.t.
2029    /// beta along direction `u_k = -H^{-1} A_k beta_hat`. The returned matrix B
2030    /// satisfies `dH/drho_k = A_k + B`.
2031    ///
2032    /// Called via [`SurvivalDerivProvider`] which adapts the sign convention
2033    /// from the unified `HessianDerivativeProvider` trait (positive `v_k`) to
2034    /// the negated `u_k` used here.
2035    pub(crate) fn survival_hessian_derivative_correction(
2036        &self,
2037        beta: &Array1<f64>,
2038        u_k: &Array1<f64>,
2039    ) -> Result<Array2<f64>, EstimationError> {
2040        let p = beta.len();
2041        let n = self.nrows();
2042
2043        let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2044        let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2045        let deriv_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2046        let exp_entry = eta_entry.mapv(f64::exp);
2047        let exp_exit = eta_exit.mapv(f64::exp);
2048        let guard = self.derivative_guard();
2049        let guard_numerical = self.derivative_guard_numerical();
2050
2051        let jac = Array1::<f64>::ones(p);
2052        let curvature = Array1::<f64>::zeros(p);
2053        let third = Array1::<f64>::zeros(p);
2054
2055        let mut row_exit = vec![0.0_f64; p];
2056        let mut row_entry = vec![0.0_f64; p];
2057        let mut row_derivative = vec![0.0_f64; p];
2058        let mut ge = vec![0.0_f64; p];
2059        let mut gs = vec![0.0_f64; p];
2060        let mut gsd = vec![0.0_f64; p];
2061        let mut he = vec![0.0_f64; p];
2062        let mut hs = vec![0.0_f64; p];
2063        let mut hsd = vec![0.0_f64; p];
2064        let mut te = vec![0.0_f64; p];
2065        let mut ts = vec![0.0_f64; p];
2066        let mut tsd = vec![0.0_f64; p];
2067
2068        let mut b_dir = Array2::<f64>::zeros((p, p));
2069
2070        for i in 0..n {
2071            let w_i = self.sampleweight[i];
2072            if w_i <= 0.0 {
2073                continue;
2074            }
2075            let has_entry = !self.entry_at_origin[i];
2076            let mut deta_e = 0.0_f64;
2077            let mut deta_s = 0.0_f64;
2078            let mut ds = 0.0_f64;
2079            self.fill_exit_row(i, &mut row_exit);
2080            self.fill_entry_row(i, &mut row_entry);
2081            self.fill_derivative_row(i, &mut row_derivative);
2082            for j in 0..p {
2083                ge[j] = row_exit[j] * jac[j];
2084                gs[j] = row_entry[j] * jac[j];
2085                gsd[j] = row_derivative[j] * jac[j];
2086                he[j] = row_exit[j] * curvature[j];
2087                hs[j] = row_entry[j] * curvature[j];
2088                hsd[j] = row_derivative[j] * curvature[j];
2089                te[j] = row_exit[j] * third[j];
2090                ts[j] = row_entry[j] * third[j];
2091                tsd[j] = row_derivative[j] * third[j];
2092                deta_e += ge[j] * u_k[j];
2093                if has_entry {
2094                    deta_s += gs[j] * u_k[j];
2095                }
2096                ds += gsd[j] * u_k[j];
2097            }
2098
2099            // Interval part: d/dbeta [ exp(eta) * (g g^T + diag(h)) ][u_k]
2100            for r in 0..p {
2101                let dge_r = he[r] * u_k[r];
2102                let dgs_r = hs[r] * u_k[r];
2103                let dhe_r = te[r] * u_k[r];
2104                let dhs_r = ts[r] * u_k[r];
2105                for c in 0..p {
2106                    let dge_c = he[c] * u_k[c];
2107                    let dgs_c = hs[c] * u_k[c];
2108                    let mut d_h_rc =
2109                        exp_exit[i] * (deta_e * ge[r] * ge[c] + dge_r * ge[c] + ge[r] * dge_c);
2110                    if r == c {
2111                        d_h_rc += exp_exit[i] * (deta_e * he[r] + dhe_r);
2112                    }
2113                    if has_entry {
2114                        d_h_rc -=
2115                            exp_entry[i] * (deta_s * gs[r] * gs[c] + dgs_r * gs[c] + gs[r] * dgs_c);
2116                        if r == c {
2117                            d_h_rc -= exp_entry[i] * (deta_s * hs[r] + dhs_r);
2118                        }
2119                    }
2120                    b_dir[[r, c]] += w_i * d_h_rc;
2121                }
2122            }
2123
2124            // Event part: d/dbeta [ gsd gsd^T / s^2 - diag(he) - diag(hsd / s) ][u_k]
2125            let (s_i, s_slope) = self
2126                .stabilized_structural_derivative(deriv_raw[i])
2127                .unwrap_or((deriv_raw[i], 1.0));
2128            if !s_i.is_finite() {
2129                return Err(EstimationError::ParameterConstraintViolation(format!(
2130                    "survival monotonicity violated in unified trace contraction at row {i}: \
2131                     d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
2132                )));
2133            }
2134            if self.event_target[i] > 0 && s_slope != 0.0 {
2135                // On the floored clamp branch (slope 0) the event Hessian
2136                // block is identically zero in a neighborhood of β, so its
2137                // directional derivative vanishes and the whole event part is
2138                // skipped.
2139                if s_i < guard_numerical {
2140                    return Err(EstimationError::ParameterConstraintViolation(format!(
2141                        "survival monotonicity violated in unified trace contraction at row {i}: \
2142                         d_eta/dt={s_i:.3e} <= tolerance={guard:.3e}",
2143                    )));
2144                }
2145                let inv_s = 1.0 / s_i;
2146                let inv_s2 = inv_s * inv_s;
2147                let inv_s3 = inv_s2 * inv_s;
2148                for r in 0..p {
2149                    let dgd_r = hsd[r] * u_k[r];
2150                    let dtsd_r = tsd[r] * u_k[r];
2151                    let dte_r = te[r] * u_k[r];
2152                    for c in 0..p {
2153                        let dgd_c = hsd[c] * u_k[c];
2154                        let mut d_h_rc = (dgd_r * gsd[c] + gsd[r] * dgd_c) * inv_s2
2155                            - 2.0 * gsd[r] * gsd[c] * ds * inv_s3;
2156                        if r == c {
2157                            d_h_rc += -dte_r;
2158                            d_h_rc += -(dtsd_r * inv_s - hsd[r] * ds * inv_s2);
2159                        }
2160                        b_dir[[r, c]] += w_i * d_h_rc;
2161                    }
2162                }
2163            }
2164        }
2165
2166        Ok(b_dir)
2167    }
2168
2169    /// Per-observation gradients of the unpenalized survival NLL with respect
2170    /// to each additive offset channel, at the given β.
2171    ///
2172    /// Contract (Royston-Parmar, eta = log H(t)):
2173    ///
2174    ///   NLL_i(β; o_E, o_X, o_D) = w_i · [
2175    ///       exp(η1_i) − 1{has_entry}·exp(η0_i)
2176    ///       − δ_i · (η1_i + log s_i)
2177    ///   ]
2178    ///
2179    /// with η1_i = a1_iᵀβ + o_X[i], η0_i = a0_iᵀβ + o_E[i],
2180    ///      s_i  = d_iᵀβ + o_D[i].
2181    ///
2182    /// The additive offsets enter each of the three η channels linearly, so
2183    ///   ∂NLL_i/∂o_X[i] = w_i · (exp(η1_i) − δ_i)
2184    ///   ∂NLL_i/∂o_E[i] = −w_i · exp(η0_i) · 1{has_entry_interval}
2185    ///   ∂NLL_i/∂o_D[i] = −w_i · δ_i / s_i         (event-row only)
2186    ///
2187    /// These three arrays are the sampleweight-scaled residuals used to chain
2188    /// `∂NLL/∂offset` into `∂NLL/∂θ` via any closed-form `∂offset/∂θ` map
2189    /// (see `baseline_offset_theta_partials` for parametric baselines). At
2190    /// converged β*, the envelope theorem on the penalized objective gives
2191    ///
2192    ///   d[0.5·(deviance + β*ᵀS_λβ*)] / dθ
2193    ///     = Σᵢ r_X_i·∂o_X_i/∂θ + r_E_i·∂o_E_i/∂θ + r_D_i·∂o_D_i/∂θ
2194    ///
2195    /// exactly (no IFT back-solve required), because β* is a stationary point
2196    /// of the penalized objective wrt β and the penalty has no θ dependence.
2197    ///
2198    /// Rows with `sampleweight[i] ≤ 0` and non-event rows for `r_D` are
2199    /// returned as exact 0.0 so the output can be dot-producted against a
2200    /// per-obs baseline-partials array without a mask.
2201    ///
2202    /// Structural-monotonicity stabilization on `s_i` (see
2203    /// `stabilized_structural_derivative`) is applied identically to the
2204    /// existing `update_state` path so the residual agrees with the
2205    /// NLL that `update_state` evaluates.
2206    pub fn offset_channel_residuals(
2207        &self,
2208        beta: &Array1<f64>,
2209    ) -> Result<OffsetChannelResiduals, EstimationError> {
2210        if beta.len() != self.coefficient_dim() {
2211            crate::bail_invalid_estim!(
2212                "survival beta dimension mismatch in offset_channel_residuals"
2213            );
2214        }
2215        let n = self.nrows();
2216        let eta_entry = self.entry_dot(beta) + &self.offset_eta_entry;
2217        let eta_exit = self.exit_dot(beta) + &self.offset_eta_exit;
2218        let derivative_raw = self.derivative_dot(beta) + &self.offset_derivative_exit;
2219
2220        let derivative_guard_numerical = self.derivative_guard_numerical();
2221        let mut r_exit = Array1::<f64>::zeros(n);
2222        let mut r_entry = Array1::<f64>::zeros(n);
2223        let mut r_deriv = Array1::<f64>::zeros(n);
2224
2225        for i in 0..n {
2226            let w = self.sampleweight[i];
2227            if w <= 0.0 {
2228                continue;
2229            }
2230            let entry_age = self.age_entry[i];
2231            let exit_age = self.age_exit[i];
2232            if !entry_age.is_finite() || !exit_age.is_finite() || exit_age < entry_age {
2233                crate::bail_invalid_estim!(
2234                    "survival ages must be finite with age_exit >= age_entry"
2235                );
2236            }
2237            let has_entry_interval = !self.entry_at_origin[i];
2238            let d = f64::from(self.event_target[i]);
2239            // Phase-1 values matching update_state:
2240            //   w_exit_i  = w · exp(eta_exit)                    → ∂NLL/∂o_X before − δ·w term
2241            //   w_entry_i = w · exp(eta_entry) · 1{has_entry}    → matches −∂NLL/∂o_E sign
2242            let w_exit_i = w * eta_exit[i].exp();
2243            let w_entry_i = if has_entry_interval {
2244                w * eta_entry[i].exp()
2245            } else {
2246                0.0
2247            };
2248            if !w_exit_i.is_finite() {
2249                crate::bail_invalid_estim!(
2250                    "offset_channel_residuals: w*exp(eta_exit)={w_exit_i:.3e} non-finite at row {i}"
2251                );
2252            }
2253            r_exit[i] = w_exit_i - d * w;
2254            r_entry[i] = -w_entry_i;
2255            // Same per-row monotonicity rule as `update_state`: a strictly
2256            // negative derivative at any observed exit time (event or
2257            // censored) falsifies S(t); event rows additionally need
2258            // `deriv > guard` because `1/deriv` enters their score.
2259            let deriv_raw = derivative_raw[i];
2260            let (deriv, deriv_slope) = self
2261                .stabilized_structural_derivative(deriv_raw)
2262                .unwrap_or((deriv_raw, 1.0));
2263            let mono_floor = if d > 0.0 {
2264                derivative_guard_numerical
2265            } else {
2266                0.0
2267            };
2268            if !deriv.is_finite() || deriv < mono_floor {
2269                return Err(EstimationError::ParameterConstraintViolation(format!(
2270                    "offset_channel_residuals: derivative ≤ numerical guard at row {i}: {deriv:.3e}"
2271                )));
2272            }
2273            if d > 0.0 {
2274                // The clamp slope zeroes the residual on the floored branch,
2275                // matching update_state's flat `ln(deriv)` value there.
2276                r_deriv[i] = -w * d * deriv_slope / deriv;
2277            }
2278        }
2279
2280        let right = Array1::<f64>::zeros(r_exit.len());
2281        Ok(OffsetChannelResiduals {
2282            exit: r_exit,
2283            entry: r_entry,
2284            derivative: r_deriv,
2285            right,
2286        })
2287    }
2288
2289    /// Build an [`InnerSolution`](gam_solve::estimate::reml::reml_outer_engine::InnerSolution) from
2290    /// the survival working state, suitable for the unified REML/LAML evaluator.
2291    ///
2292    /// Evaluate the survival outer objective and gradient via the unified REML/LAML
2293    /// evaluator, using the canonical assembly module.
2294    pub fn unified_lamlobjective_and_rhogradient(
2295        &self,
2296        beta: &Array1<f64>,
2297        state: &WorkingState,
2298        rho: &Array1<f64>,
2299    ) -> Result<(f64, Array1<f64>), EstimationError> {
2300        use gam_problem::{EvalMode, PseudoLogdetMode};
2301        use gam_solve::estimate::reml::assembly::{
2302            InnerAssembly, PenaltyBlockDesc, penalty_coords_from_blocks,
2303        };
2304        use gam_solve::estimate::reml::reml_outer_engine::{
2305            DenseSpectralOperator, DispersionHandling, PenaltyLogdetDerivs,
2306            compute_block_penalty_logdet_derivs,
2307        };
2308
2309        let p = beta.len();
2310        let active_penalty_blocks: Vec<&PenaltyBlock> = self
2311            .penalties
2312            .blocks
2313            .iter()
2314            .filter(|b| b.lambda > 0.0)
2315            .collect();
2316        if rho.len() != active_penalty_blocks.len() {
2317            crate::bail_invalid_estim!(
2318                "survival LAML rho dimension {} does not match active penalty block count {}",
2319                rho.len(),
2320                active_penalty_blocks.len()
2321            );
2322        }
2323        let k_count = active_penalty_blocks.len();
2324
2325        // --- Hessian operator ---
2326        let h_dense = state.hessian.to_dense();
2327        let has_left_truncation = self
2328            .age_entry
2329            .iter()
2330            .any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD);
2331        // Transformation-survival uses observed information in the LAML logdet.
2332        // With delayed entry the likelihood contains +H(entry), so the observed
2333        // NLL curvature includes a genuine negative
2334        // -X_entry' diag(exp(eta_entry)) X_entry block. The shared smooth
2335        // pseudo-logdet is a PSD-contract regularizer, not a licence to reward
2336        // negative observed-curvature directions: a negative eigenvalue maps to
2337        // a tiny positive regularized value and can make the outer smoothing
2338        // objective prefer under-smoothed, nearly singular baselines. For the
2339        // delayed-entry observed-information path, use the identified positive
2340        // subspace logdet/pseudoinverse instead; right-censored fits keep the
2341        // historical smooth full-spectrum convention.
2342        let hessian_logdet_mode = if has_left_truncation {
2343            PseudoLogdetMode::HardPseudo
2344        } else {
2345            PseudoLogdetMode::Smooth
2346        };
2347        let hop = DenseSpectralOperator::from_symmetric_with_mode(&h_dense, hessian_logdet_mode)
2348            .map_err(EstimationError::InvalidInput)?;
2349
2350        // --- Penalty coordinates via shared assembler helper ---
2351        let block_descs: Vec<PenaltyBlockDesc> = self
2352            .penalties
2353            .blocks
2354            .iter()
2355            .filter(|b| b.lambda > 0.0)
2356            .map(|b| PenaltyBlockDesc {
2357                matrix: &b.matrix,
2358                range_start: b.range.start,
2359                range_end: b.range.end,
2360            })
2361            .collect();
2362        let penalty_coords =
2363            penalty_coords_from_blocks(&block_descs, p).map_err(EstimationError::InvalidInput)?;
2364
2365        // --- Penalty logdet derivatives ---
2366        let per_block_rho: Vec<Array1<f64>> =
2367            rho.iter().map(|&r| Array1::from_vec(vec![r])).collect();
2368        let per_block_penalty_matrices: Vec<Vec<Array2<f64>>> = active_penalty_blocks
2369            .iter()
2370            .map(|b| vec![b.matrix.clone()])
2371            .collect();
2372        let per_block_penalty_refs: Vec<&[Array2<f64>]> = per_block_penalty_matrices
2373            .iter()
2374            .map(|v| v.as_slice())
2375            .collect();
2376        let penalty_logdet = if k_count > 0 {
2377            compute_block_penalty_logdet_derivs(&per_block_rho, &per_block_penalty_refs, 0.0)
2378                .map_err(EstimationError::InvalidInput)?
2379        } else {
2380            PenaltyLogdetDerivs {
2381                value: 0.0,
2382                first: Array1::zeros(0),
2383                second: Some(Array2::zeros((0, 0))),
2384            }
2385        };
2386
2387        // penalty_quadratic = 2 * penalty_term (matching unified evaluator convention).
2388        let penalty_quadratic = 2.0 * state.penalty_term;
2389        let provider = SurvivalDerivProvider::new(self.clone(), beta.clone());
2390
2391        // #931 survival-LAML IFT envelope: attach the one-step Newton correction
2392        // only when this state is actually a near-stationary inner solution.
2393        // `unified_lamlobjective_and_rhogradient` is also used by algebraic
2394        // fixed-beta objective tests; feeding a large non-stationary residual
2395        // there makes the value a different surface. The re-converged shim
2396        // polishes the inner mode to an absolute residual floor, so certified
2397        // states still keep the envelope correction while arbitrary beta probes
2398        // evaluate the documented LAML objective.
2399        //
2400        // The residual MUST be the active-set-projected stationarity vector, not
2401        // raw `state.gradient`: a binding monotonicity constraint contributes a
2402        // Lagrange-multiplier normal component (`r = A^T lambda`, lambda >= 0)
2403        // that is not a stationarity residual.
2404        const SURVIVAL_LAML_IFT_RELATIVE_KKT_GATE: f64 = 1.0e-8;
2405        let kkt_residual = {
2406            let raw = state.gradient.clone();
2407            let projected = match self.monotonicity_linear_constraints() {
2408                Some(constraints) => {
2409                    projected_linear_constraint_stationarity_vector(&raw, beta, &constraints, None)
2410                        .ok_or_else(|| {
2411                            EstimationError::InvalidInput(
2412                                "survival LAML could not project the monotonicity KKT residual"
2413                                    .to_string(),
2414                            )
2415                        })?
2416                }
2417                None => raw,
2418            };
2419            let projected_norm = array1_l2_norm(&projected);
2420            let relative_projected_norm = state.relative_gradient_norm(projected_norm);
2421            if relative_projected_norm <= SURVIVAL_LAML_IFT_RELATIVE_KKT_GATE {
2422                Some(crate::model_types::ProjectedKktResidual::from_active_projected(projected))
2423            } else {
2424                None
2425            }
2426        };
2427
2428        let result = InnerAssembly {
2429            log_likelihood: state.log_likelihood,
2430            penalty_quadratic,
2431            beta: beta.clone(),
2432            n_observations: self.nrows(),
2433            hessian_op: std::sync::Arc::new(hop),
2434            penalty_coords,
2435            penalty_logdet,
2436            dispersion: DispersionHandling::Fixed {
2437                phi: 1.0,
2438                include_logdet_h: true,
2439                include_logdet_s: true,
2440            },
2441            rho_curvature_scale: 1.0,
2442            rho_prior: gam_problem::RhoPrior::Flat,
2443            hessian_logdet_correction: 0.0,
2444            penalty_subspace_trace: None,
2445            deriv_provider: Some(Box::new(provider)),
2446            firth: None,
2447            nullspace_dim: None,
2448            barrier_config: None,
2449            ext_coords: Vec::new(),
2450            ext_coord_pair_fn: None,
2451            rho_ext_pair_fn: None,
2452            fixed_drift_deriv: None,
2453            contracted_psi_second_order: None,
2454            kkt_residual,
2455            active_constraints: None,
2456        }
2457        .evaluate(
2458            rho.as_slice().expect("rho must be contiguous"),
2459            EvalMode::ValueAndGradient,
2460            None,
2461        )
2462        .map_err(EstimationError::InvalidInput)?;
2463
2464        let gradient = result.gradient.unwrap_or_else(|| Array1::zeros(rho.len()));
2465        Ok((result.cost, gradient))
2466    }
2467
2468    /// Self-contained ρ → (LAML value, analytic ρ-gradient) surface for the
2469    /// survival LAML objective.
2470    ///
2471    /// Unlike [`unified_lamlobjective_and_rhogradient`](Self::unified_lamlobjective_and_rhogradient),
2472    /// which takes a *pre-converged* [`WorkingState`] and `β̂` at the evaluated
2473    /// `ρ`, this shim re-converges the inner survival mode internally: it sets
2474    /// the active-block smoothing parameters to `λ = exp(ρ)`, runs the same
2475    /// constrained inner PIRLS that the survival outer loop uses
2476    /// ([`runworking_model_pirls`](gam_solve::pirls::runworking_model_pirls)), then
2477    /// evaluates the unified survival LAML value and analytic ρ-gradient at the
2478    /// re-fitted `β̂(ρ)`. The returned pair is therefore a single-source value+
2479    /// gradient surface that a caller can finite-difference by varying `ρ`
2480    /// alone — the survival counterpart of the GLM path's
2481    /// `evaluate_externalgradient` / `evaluate_externalcost_andridge`.
2482    ///
2483    /// `rho` enumerates the **active** penalty blocks (those with `λ > 0`) in
2484    /// block order, matching the convention of the unified evaluator. `beta0` is
2485    /// the inner warm-start. The behaviour is identical to the existing survival
2486    /// LAML path (set-λ → inner PIRLS → `update_state` → unified LAML); this is a
2487    /// reachability shim, not a new objective.
2488    pub fn evaluate_survival_lamlcost_and_gradient(
2489        &self,
2490        rho: &[f64],
2491        beta0: &Array1<f64>,
2492    ) -> Result<(f64, Array1<f64>), EstimationError> {
2493        let (candidate, beta) = self.reconverge_survival_inner_mode(rho, beta0)?;
2494        // Re-converged β̂(ρ); evaluate the unified survival LAML value and
2495        // analytic ρ-gradient at that mode. The ρ passed to the unified
2496        // evaluator enumerates active blocks in block order, exactly the input
2497        // convention of this shim.
2498        let rho_arr = Array1::from_vec(rho.to_vec());
2499        let state = candidate.update_state(&beta)?;
2500        candidate.unified_lamlobjective_and_rhogradient(&beta, &state, &rho_arr)
2501    }
2502
2503    /// Re-converge the survival inner mode at `λ = exp(ρ)` from warm-start
2504    /// `beta0`, returning the λ-set model candidate and the converged `β̂(ρ)`.
2505    /// This is the shared inner-solve used by
2506    /// [`evaluate_survival_lamlcost_and_gradient`](Self::evaluate_survival_lamlcost_and_gradient):
2507    /// inner PIRLS to a tight relative certificate, followed by a
2508    /// Levenberg–Marquardt / exact-Cholesky stationarity polish that drives the
2509    /// absolute penalized residual `‖S β̂ − ∇ℓ‖` below the FD round-off floor so
2510    /// the envelope ρ-gradient is exact. (Without the polish, PIRLS alone leaves
2511    /// `‖r‖ ~ 1` at large λ where H is ill-conditioned.)
2512    fn reconverge_survival_inner_mode(
2513        &self,
2514        rho: &[f64],
2515        beta0: &Array1<f64>,
2516    ) -> Result<(WorkingModelSurvival, Array1<f64>), EstimationError> {
2517        // Inner-PIRLS settings mirror the survival transformation outer loop's
2518        // constrained inner solve. Tighter convergence than the production
2519        // outer loop so the inner mode is converged well below the FD step's
2520        // round-off floor, making ∇V finite-differentiable in ρ alone.
2521        const SHIM_PIRLS_MAX_ITERATIONS: usize = 600;
2522        const SHIM_PIRLS_CONVERGENCE_TOL: f64 = 1e-12;
2523        const SHIM_PIRLS_MAX_STEP_HALVING: usize = 40;
2524        const SHIM_PIRLS_MIN_STEP_SIZE: f64 = 1e-12;
2525
2526        let active_block_count = self
2527            .penalties
2528            .blocks
2529            .iter()
2530            .filter(|b| b.lambda > 0.0)
2531            .count();
2532        if rho.len() != active_block_count {
2533            crate::bail_invalid_estim!(
2534                "reconverge_survival_inner_mode: rho dimension {} does not match active penalty block count {}",
2535                rho.len(),
2536                active_block_count
2537            );
2538        }
2539        if beta0.len() != self.coefficient_dim() {
2540            crate::bail_invalid_estim!(
2541                "reconverge_survival_inner_mode: beta0 dimension {} does not match coefficient dimension {}",
2542                beta0.len(),
2543                self.coefficient_dim()
2544            );
2545        }
2546
2547        // Set λ = exp(ρ) on the active blocks (block order), leaving inactive
2548        // (λ = 0) blocks untouched, then re-converge the inner mode.
2549        let mut candidate = self.clone();
2550        let mut lambdas: Vec<f64> = candidate
2551            .penalties
2552            .blocks
2553            .iter()
2554            .map(|b| b.lambda)
2555            .collect();
2556        let mut active_idx = 0usize;
2557        for (block, lambda) in candidate.penalties.blocks.iter().zip(lambdas.iter_mut()) {
2558            if block.lambda > 0.0 {
2559                *lambda = rho[active_idx].exp();
2560                active_idx += 1;
2561            }
2562        }
2563        candidate.set_penalty_lambdas(&lambdas)?;
2564
2565        let opts = gam_solve::pirls::WorkingModelPirlsOptions {
2566            max_iterations: SHIM_PIRLS_MAX_ITERATIONS,
2567            convergence_tolerance: SHIM_PIRLS_CONVERGENCE_TOL,
2568            adaptive_kkt_tolerance: None,
2569            max_step_halving: SHIM_PIRLS_MAX_STEP_HALVING,
2570            min_step_size: SHIM_PIRLS_MIN_STEP_SIZE,
2571            firth_bias_reduction: false,
2572            coefficient_lower_bounds: None,
2573            linear_constraints: None,
2574            initial_lm_lambda: None,
2575            arrow_schur: None,
2576        };
2577        let summary = gam_solve::pirls::runworking_model_pirls(
2578            &mut candidate,
2579            Coefficients::new(beta0.clone()),
2580            &opts,
2581            |_| {},
2582        )?;
2583        let mut beta = summary.beta.as_ref().to_owned();
2584
2585        // PIRLS exits on a RELATIVE KKT / deviance-plateau certificate, which can leave
2586        // an ABSOLUTE penalized stationarity residual r = S beta_hat - grad_ell of order
2587        // 0.1-1 (the score scales as O(sqrt(n))). The unified LAML gradient uses the
2588        // envelope theorem, exact only at r = 0; a residual that large leaks <r, beta_dot>
2589        // into the objective<->gradient consistency, AND the IFT envelope correction is
2590        // only leading-order in r, so it cannot make the analytic gradient the exact
2591        // derivative of the (re-converged, non-smooth-in-r) value surface either. The
2592        // robust cure is to drive the inner to TRUE stationarity (||r|| ~ 1e-11) so the
2593        // envelope is exactly valid and the IFT term is ~1e-22 — which it is at small
2594        // lambda, but a plain undamped Newton-polish STALLS at large lambda (rho=4..8):
2595        // there the intercept-direction curvature exp(eta)*n is large while the penalized
2596        // time block is lambda*S, so H is ill-conditioned and an undamped step neither
2597        // decreases ||r|| nor stays feasible, leaving ||r|| ~ 3e-2.
2598        //
2599        // Levenberg–Marquardt damping fixes this: solve (H + mu*diag(H)) delta = r,
2600        // accept on a genuine ||r||^2 decrease (Gauss–Newton on the stationarity system,
2601        // whose Jacobian is the penalized Hessian H), shrink mu on success and grow it on
2602        // rejection. The diagonal (Marquardt) scaling makes the damping curvature-aware so
2603        // the stiff time block and the soft intercept are damped commensurately. This
2604        // reliably reaches ||r|| below the FD-step round-off floor across the whole
2605        // rho = [-0.5 .. 8] range exercised by the consistency gates.
2606        {
2607            const POLISH_MAX_ITERS: usize = 400;
2608            const POLISH_TOL: f64 = 1e-13;
2609            // Armijo sufficient-decrease constant and backtracking factor —
2610            // the shared opt tuning constants; only the halving
2611            // budget (80, deeper than the shared 60) is site-specific.
2612            const ARMIJO_C: f64 = constants::ARMIJO_C1;
2613            const BACKTRACK: f64 = constants::BACKTRACK_CONTRACTION;
2614            const MAX_BACKTRACK: usize = 80;
2615            let p = beta.len();
2616            // Penalized inner objective f(β) = −ℓ(β) + ½β'Sβ whose gradient is
2617            // exactly `state.gradient` and whose UNDAMPED Hessian is exactly
2618            // `state.hessian`. `update_state` exposes the pieces directly; the
2619            // Levenberg–Marquardt shift below exists only while solving a step.
2620            let penalized_objective =
2621                |st: &WorkingState| -> f64 { -st.log_likelihood + st.penalty_term };
2622            for _ in 0..POLISH_MAX_ITERS {
2623                let st = match candidate.update_state(&beta) {
2624                    Ok(st) => st,
2625                    Err(_) => break,
2626                };
2627                let r = st.gradient.clone();
2628                let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
2629                if !r_norm.is_finite() || r_norm < POLISH_TOL {
2630                    break;
2631                }
2632                let h = st.hessian.to_dense();
2633                let f0 = penalized_objective(&st);
2634                // Newton DIRECTION d = −H⁻¹r on the convex penalized survival
2635                // likelihood, found via a Levenberg–Marquardt-regularized solve
2636                // so an ill-conditioned H (h_diag ratio ~2400 at β₀≈4.6, where
2637                // exp(η) is huge) cannot produce a garbage direction whose
2638                // quadratic form rᵀH⁻¹r loses its sign. If even the regularized
2639                // Newton direction is not a sufficient descent direction, fall
2640                // back to STEEPEST DESCENT d = −r, which is ALWAYS a descent
2641                // direction on the convex objective (∇fᵀ(−r) = −‖r‖² < 0). The
2642                // line search below is on the OBJECTIVE VALUE (not ‖r‖), so any
2643                // descent direction makes monotone progress; near the optimum
2644                // the (lightly regularized) Newton step recovers fast local
2645                // convergence. This is globally convergent for the convex
2646                // penalized survival NLL — driving ‖r‖ below the FD round-off
2647                // floor so the envelope ρ-gradient equals the finite difference.
2648                let h_scale = (0..p)
2649                    .map(|d| h[[d, d]].abs())
2650                    .fold(0.0_f64, f64::max)
2651                    .max(1.0);
2652                // Solve (H + λI) step = r by an EXACT Cholesky factorization
2653                // (faer Llt), NOT the DenseSpectralOperator: the spectral
2654                // operator clamps tiny/negative eigenvalues, which on the
2655                // catastrophically ill-conditioned boundary Hessian (cond ~2400,
2656                // exp(η) huge at β₀≈4.6) corrupts the solve so badly that
2657                // rᵀH⁻¹r lost its sign and the previous polish broke on iter 0.
2658                // Cholesky succeeds iff H+λI is SPD; sweeping λ up from 0 finds
2659                // the smallest SPD shift, and for an SPD system rᵀ(H+λI)⁻¹r > 0
2660                // EXACTLY (Cholesky is backward-stable, no clamping), so the
2661                // Newton direction is a guaranteed descent direction.
2662                // The attempt certifies a DESCENT direction, not mere
2663                // factorability: Cholesky must succeed, the solve must stay
2664                // finite, and the directional derivative must clear the
2665                // sign floor. ∇fᵀd = rᵀ(−step) = −r·(H+λI)⁻¹r < 0 exactly
2666                // for SPD systems.
2667                let try_lm = |lambda_lm: f64| -> Option<(Array1<f64>, f64)> {
2668                    let mut h_reg = h.clone();
2669                    for d in 0..p {
2670                        h_reg[[d, d]] += lambda_lm;
2671                    }
2672                    let factor =
2673                        gam_linalg::faer_ndarray::FaerCholesky::cholesky(&h_reg, faer::Side::Lower)
2674                            .ok()?;
2675                    let candidate_step = factor.solvevec(&r);
2676                    if candidate_step.iter().any(|v| !v.is_finite()) {
2677                        return None;
2678                    }
2679                    let dd = -r.dot(&candidate_step);
2680                    (dd.is_finite() && dd < -1e-14 * r_norm * r_norm)
2681                        .then_some((candidate_step, dd))
2682                };
2683                // Bare λ=0 attempt first, then 17 decades from 1e-11·h_scale
2684                // (the pre-migration `1e-12·h_scale·10^pow` for pow = 1..18).
2685                let (step, dir_deriv) = try_lm(0.0)
2686                    .or_else(|| {
2687                        escalate_ridge(RidgeSchedule::geometric(1e-11 * h_scale, 17), try_lm)
2688                            .ok()
2689                            .map(|success| success.value)
2690                    })
2691                    .unwrap_or_else(|| {
2692                        // Steepest-descent fallback: d = −r ⇒ step = +r (we step
2693                        // β − step), ∇fᵀd = −‖r‖² < 0.
2694                        (r.clone(), -r_norm * r_norm)
2695                    });
2696                // Accept on EITHER a sufficient objective decrease (Armijo,
2697                // the global-convergence guarantee on the convex objective)
2698                // OR a strict residual-norm decrease. Near the solution the
2699                // penalized objective is flat to f64 roundoff (f0 ≈ ft), so a
2700                // pure-Armijo test backtracks α→0 and crawls (the asymmetric
2701                // ρ=3.99999 stall: 200 iters at 3.7e-7 vs 12 iters at the
2702                // other two ρ). The ‖r‖-decrease arm lets the exact Cholesky
2703                // Newton step (α=1) through, restoring quadratic convergence
2704                // to ~1e-12 symmetrically across all three FD points so the
2705                // centered FD of the value surface is itself exact.
2706                //
2707                // The dual criterion reads BOTH the trial objective and the
2708                // trial residual norm, so the decision lives inside the trial
2709                // closure: a rejected candidate returns `Ok(None)` (shrink)
2710                // exactly like an invalid `update_state`, and the accept
2711                // predicate is the constant `true`.
2712                let accepted = match backtracking_line_search::<_, Infallible>(
2713                    BacktrackConfig {
2714                        contraction: BACKTRACK,
2715                        max_steps: MAX_BACKTRACK,
2716                        ..BacktrackConfig::default()
2717                    },
2718                    |alpha| {
2719                        let trial = &beta - &(alpha * &step);
2720                        let Ok(ts) = candidate.update_state(&trial) else {
2721                            return Ok(None);
2722                        };
2723                        let ft = penalized_objective(&ts);
2724                        let tn = ts.gradient.iter().map(|v| v * v).sum::<f64>().sqrt();
2725                        let armijo_ok = ft.is_finite() && ft <= f0 + ARMIJO_C * alpha * dir_deriv;
2726                        let residual_ok = tn.is_finite() && tn < r_norm;
2727                        Ok((armijo_ok || residual_ok).then_some((ft, trial)))
2728                    },
2729                    |_alpha, _ft| true,
2730                ) {
2731                    Ok(result) => result,
2732                    Err(never) => match never {},
2733                };
2734                let Some(ls) = accepted else {
2735                    break;
2736                };
2737                beta = ls.payload;
2738            }
2739        }
2740
2741        Ok((candidate, beta))
2742    }
2743}
2744
2745/// Derivative provider that adapts survival third-derivative Hessian corrections
2746/// to the unified [`HessianDerivativeProvider`](gam_solve::estimate::reml::reml_outer_engine::HessianDerivativeProvider)
2747/// trait.
2748///
2749/// The unified trait supplies `v_k = H^{-1}(A_k beta_hat)` (positive sign),
2750/// whereas the survival engine's
2751/// [`survival_hessian_derivative_correction`](WorkingModelSurvival::survival_hessian_derivative_correction)
2752/// expects `u_k = -v_k`. This provider handles the sign conversion.
2753pub(crate) struct SurvivalDerivProvider {
2754    model: WorkingModelSurvival,
2755    beta: Array1<f64>,
2756}
2757
2758impl SurvivalDerivProvider {
2759    pub(crate) fn new(model: WorkingModelSurvival, beta: Array1<f64>) -> Self {
2760        Self { model, beta }
2761    }
2762}
2763
2764impl gam_solve::estimate::reml::reml_outer_engine::HessianDerivativeProvider
2765    for SurvivalDerivProvider
2766{
2767    fn hessian_derivative_correction(
2768        &self,
2769        v_k: &Array1<f64>,
2770    ) -> Result<Option<Array2<f64>>, String> {
2771        // The trait provides v_k = H^{-1}(A_k beta_hat) (positive).
2772        // The survival method expects u_k = -H^{-1} A_k beta_hat = -v_k.
2773        let u_k = -v_k;
2774        match self
2775            .model
2776            .survival_hessian_derivative_correction(&self.beta, &u_k)
2777        {
2778            Ok(correction) => Ok(Some(correction)),
2779            Err(e) => Err(e.to_string()),
2780        }
2781    }
2782
2783    fn has_corrections(&self) -> bool {
2784        true
2785    }
2786}
2787
2788#[derive(Debug, Clone)]
2789pub struct CrudeRiskResult {
2790    pub risk: f64,
2791    pub diseasegradient: Array1<f64>,
2792    pub mortalitygradient: Array1<f64>,
2793}
2794
2795#[derive(Debug, Clone)]
2796pub struct CompetingRisksCifResult {
2797    /// Cumulative incidence per endpoint. `cif[ep][[row, time_idx]]` is the
2798    /// probability that cause `ep` occurred by `times[time_idx]` for sample `row`.
2799    /// Stored one matrix per endpoint so it is ergonomic to index per-cause and
2800    /// natural to construct from the per-endpoint cumulative-hazard inputs.
2801    pub cif: Vec<Array2<f64>>,
2802    pub overall_survival: Array2<f64>,
2803}
2804
2805/// Subject-count threshold below which competing-risks CIF assembly stays on the
2806/// serial path. The per-row work (a `n_times`-long prefix-sum recurrence with a
2807/// handful of `exp`/`exp_m1` per element) is cheap, so small panels avoid rayon
2808/// fan-out overhead; large panels (the #1082 quality-test sizes) amortize it.
2809const COMPETING_RISKS_CIF_PARALLEL_ROW_MIN: usize = 256;
2810
2811pub fn assemble_competing_risks_cif(
2812    times: ArrayView1<'_, f64>,
2813    cumulative_hazard: ArrayView3<'_, f64>,
2814) -> Result<CompetingRisksCifResult, SurvivalError> {
2815    let (n_endpoints, n_rows, n_times) = cumulative_hazard.dim();
2816    if n_endpoints == 0 {
2817        return Err(SurvivalError::DimensionMismatch);
2818    }
2819    let endpoint_hazards = cumulative_hazard
2820        .axis_iter(Axis(0))
2821        .map(|view| view.to_owned())
2822        .collect::<Vec<_>>();
2823    assemble_competing_risks_cif_from_endpoints(times, &endpoint_hazards).and_then(|result| {
2824        if result.overall_survival.dim() != (n_rows, n_times) {
2825            Err(SurvivalError::DimensionMismatch)
2826        } else {
2827            Ok(result)
2828        }
2829    })
2830}
2831
2832pub fn assemble_competing_risks_cif_from_endpoints(
2833    times: ArrayView1<'_, f64>,
2834    cumulative_hazards: &[Array2<f64>],
2835) -> Result<CompetingRisksCifResult, SurvivalError> {
2836    let n_endpoints = cumulative_hazards.len();
2837    if n_endpoints == 0 || times.is_empty() {
2838        return Err(SurvivalError::DimensionMismatch);
2839    }
2840    let (n_rows, n_times) = cumulative_hazards[0].dim();
2841    if n_rows == 0 || n_times == 0 || times.len() != n_times {
2842        return Err(SurvivalError::DimensionMismatch);
2843    }
2844    if times.iter().any(|time| !time.is_finite() || *time < 0.0) {
2845        return Err(SurvivalError::InvalidTimeGrid);
2846    }
2847    if times
2848        .iter()
2849        .zip(times.iter().skip(1))
2850        .any(|(previous, current)| current <= previous)
2851    {
2852        return Err(SurvivalError::InvalidTimeGrid);
2853    }
2854    for endpoint_hazard in cumulative_hazards {
2855        if endpoint_hazard.dim() != (n_rows, n_times) {
2856            return Err(SurvivalError::DimensionMismatch);
2857        }
2858        if endpoint_hazard.iter().any(|value| !value.is_finite()) {
2859            return Err(SurvivalError::NonFiniteInput);
2860        }
2861    }
2862
2863    let max_abs_hazard = cumulative_hazards
2864        .iter()
2865        .flat_map(|endpoint_hazard| endpoint_hazard.iter())
2866        .fold(0.0_f64, |acc, value| acc.max(value.abs()));
2867    let monotone_tolerance = 1.0e-10_f64 * max_abs_hazard.max(1.0);
2868    let mut cif: Vec<Array2<f64>> = (0..n_endpoints)
2869        .map(|_| Array2::<f64>::zeros((n_rows, n_times)))
2870        .collect();
2871    let mut overall_survival = Array2::<f64>::zeros((n_rows, n_times));
2872
2873    // Per-row CIF assembly. The TIME axis is a sequential prefix-sum recurrence
2874    // (`previous_*` carry forward across `time_idx`) and MUST stay ordered, so it
2875    // is left as the inner serial loop. The ROW (subject) axis is fully
2876    // independent: every `previous_*`/`increments` buffer is allocated fresh per
2877    // row, no state crosses rows, and each row writes only its own disjoint
2878    // output slices. The per-row result is byte-identical regardless of which
2879    // thread runs it, so we fan the outer row loop out over rayon and write the
2880    // owned per-row buffers back serially in row order (deterministic, bit-exact
2881    // vs. the serial implementation).
2882    //
2883    // `cif_flat` is endpoint-major: `cif_flat[endpoint * n_times + time_idx]`.
2884    let assemble_row = |row: usize| -> Result<(Vec<f64>, Vec<f64>), SurvivalError> {
2885        let mut cif_flat = vec![0.0_f64; n_endpoints * n_times];
2886        let mut surv_row = vec![0.0_f64; n_times];
2887        let mut previous_cif = vec![0.0_f64; n_endpoints];
2888        let mut previous_cumulative = vec![0.0_f64; n_endpoints];
2889        let mut increments = vec![0.0_f64; n_endpoints];
2890        let mut previous_total_cumulative = 0.0_f64;
2891        for time_idx in 0..n_times {
2892            let mut total_increment = 0.0_f64;
2893            for endpoint in 0..n_endpoints {
2894                let current = cumulative_hazards[endpoint][[row, time_idx]];
2895                if current < -monotone_tolerance {
2896                    return Err(SurvivalError::NonMonotoneCumulativeHazard);
2897                }
2898                let raw_increment = current - previous_cumulative[endpoint];
2899                if raw_increment < -monotone_tolerance {
2900                    return Err(SurvivalError::NonMonotoneCumulativeHazard);
2901                }
2902                let increment = raw_increment.max(0.0);
2903                increments[endpoint] = increment;
2904                total_increment += increment;
2905                previous_cumulative[endpoint] += increment;
2906            }
2907
2908            let survival_left = (-previous_total_cumulative).exp();
2909            let interval_failure = -(-total_increment).exp_m1();
2910            for endpoint in 0..n_endpoints {
2911                if total_increment > 0.0 {
2912                    previous_cif[endpoint] +=
2913                        survival_left * interval_failure * increments[endpoint] / total_increment;
2914                }
2915                cif_flat[endpoint * n_times + time_idx] = previous_cif[endpoint].clamp(0.0, 1.0);
2916            }
2917            previous_total_cumulative += total_increment;
2918            // Derive `S(t)` from the stored cause-specific CIFs at this time so
2919            // that the competing-risks closure identity
2920            //   Σ_k F_k(t) + S(t) = 1
2921            // holds bit-exactly. Computing `S` independently as
2922            // `exp(-Σ_k H_k(t))` and then comparing against the (clamped, ratio-
2923            // split) Σ F_k introduces O(machine-eps) closure error because the
2924            // float increments
2925            //   ΔF_k = S_left·(1-exp(-ΔH))·ΔH_k/ΔH_total
2926            // do not sum to `S_left - S_new` bit-exactly. By summing the stored
2927            // CIFs in the same left-fold order as `slice.iter().sum::<f64>()`
2928            // and defining `S := 1.0 - Σ F_k`, the IEEE-754 round-trip
2929            //   (1.0 - f) + f
2930            // restores the identity for finite f ∈ [0, 1]. The mathematically
2931            // consistent survival value `exp(-H_total)` is still tracked up to
2932            // ulp-level precision because the ΔF_k construction matches
2933            // `S_left - S_new` to leading order.
2934            let mut fsum_at_t = 0.0_f64;
2935            for endpoint in 0..n_endpoints {
2936                fsum_at_t += cif_flat[endpoint * n_times + time_idx];
2937            }
2938            surv_row[time_idx] = (1.0_f64 - fsum_at_t).clamp(0.0, 1.0);
2939        }
2940        Ok((cif_flat, surv_row))
2941    };
2942
2943    // Nesting guard (`rayon::current_thread_index().is_none()`) keeps us from
2944    // oversubscribing when this routine is itself called from inside a rayon
2945    // worker, and the row-count gate keeps small inputs on the serial path.
2946    let rows: Vec<(Vec<f64>, Vec<f64>)> = if n_rows >= COMPETING_RISKS_CIF_PARALLEL_ROW_MIN
2947        && rayon::current_thread_index().is_none()
2948    {
2949        use rayon::prelude::*;
2950        (0..n_rows)
2951            .into_par_iter()
2952            .map(assemble_row)
2953            .collect::<Result<_, _>>()?
2954    } else {
2955        (0..n_rows).map(assemble_row).collect::<Result<_, _>>()?
2956    };
2957
2958    for (row, (cif_flat, surv_row)) in rows.into_iter().enumerate() {
2959        for endpoint in 0..n_endpoints {
2960            for time_idx in 0..n_times {
2961                cif[endpoint][[row, time_idx]] = cif_flat[endpoint * n_times + time_idx];
2962            }
2963        }
2964        for time_idx in 0..n_times {
2965            overall_survival[[row, time_idx]] = surv_row[time_idx];
2966        }
2967    }
2968
2969    Ok(CompetingRisksCifResult {
2970        cif,
2971        overall_survival,
2972    })
2973}
2974
2975/// `(node, weight)` pairs of the `n`-point Gauss-Legendre rule on `[-1, 1]`;
2976/// the canonical generator lives in `gam-math` (previously triplicated
2977/// across gam-terms / gam-model-kernels / gam-models).
2978fn compute_gauss_legendre_nodes(n: usize) -> Vec<(f64, f64)> {
2979    let (nodes, weights) = gam_math::special::gauss_legendre(n);
2980    nodes.into_iter().zip(weights).collect()
2981}
2982
2983fn gauss_legendre_quadrature() -> &'static [(f64, f64)] {
2984    // `LazyLock` (not `OnceLock::get_or_init`) so first init never parks a
2985    // caller on the OS condvar from inside a rayon worker. The competing-risks
2986    // CIF assembler in this file dispatches `into_par_iter` and the
2987    // codebase-level lint (`tests/once_lock_get_or_init_not_inside_parallel_regions.rs`)
2988    // forbids the lazy `OnceLock` accessor in any rayon-adjacent file.
2989    static CACHE: LazyLock<Vec<(f64, f64)>> = LazyLock::new(|| compute_gauss_legendre_nodes(40));
2990    &CACHE
2991}
2992
2993/// Engine-level crude risk quadrature with exact delta-method gradients.
2994///
2995/// This routine owns the numerical integration and gradient accumulation math:
2996/// - It integrates `h_d(u) * S_total(u | t0)` over `[t0, t1]` by high-order
2997///   Gauss-Legendre quadrature.
2998/// - It computes gradients w.r.t. disease and mortality coefficients:
2999///   d Risk / d beta_d and d Risk / d beta_m.
3000///
3001/// The adapter provides the domain-specific point evaluator callback `eval_at`,
3002/// which fills design rows and returns:
3003/// - instantaneous disease hazard h_d(u) at age `u`,
3004/// - cumulative disease hazard `H_d(u)`,
3005/// - cumulative mortality hazard `H_m(u)`.
3006///
3007/// The callback must fill the following arrays (one entry per coefficient):
3008/// - `design_d[j]`: partial derivative of the linear predictor eta_d w.r.t. beta_j
3009///   at time u, i.e. x_j(u) = d eta_d(u) / d beta_j.
3010/// - `deriv_d[j]`: partial derivative of the TIME DERIVATIVE of eta_d w.r.t. beta_j
3011///   at time u, i.e. x_dot_j(u) = d/d(beta_j) [d eta_d(u)/du].
3012/// - `design_m[j]`: same as design_d but for the mortality linear predictor eta_m.
3013///
3014/// This keeps domain/data wiring out of `gam` while centralizing the
3015/// integration engine in one place.
3016pub fn calculate_crude_risk_quadrature<F>(
3017    t0: f64,
3018    t1: f64,
3019    breakpoints: &[f64],
3020    h_dis_t0: f64,
3021    h_mor_t0: f64,
3022    design_d_t0: ArrayView1<'_, f64>,
3023    design_m_t0: ArrayView1<'_, f64>,
3024    mut eval_at: F,
3025) -> Result<CrudeRiskResult, SurvivalError>
3026where
3027    F: FnMut(
3028        f64,
3029        &mut Array1<f64>,
3030        &mut Array1<f64>,
3031        &mut Array1<f64>,
3032    ) -> Result<(f64, f64, f64), SurvivalError>,
3033{
3034    let coeff_len_d = design_d_t0.len();
3035    let coeff_len_m = design_m_t0.len();
3036    if coeff_len_d == 0 || coeff_len_m == 0 {
3037        return Err(SurvivalError::InvalidIntegrationSetup);
3038    }
3039    if !t0.is_finite()
3040        || !t1.is_finite()
3041        || !h_dis_t0.is_finite()
3042        || !h_mor_t0.is_finite()
3043        || design_d_t0.iter().any(|v| !v.is_finite())
3044        || design_m_t0.iter().any(|v| !v.is_finite())
3045    {
3046        return Err(SurvivalError::NonFiniteInput);
3047    }
3048    if t1 <= t0 {
3049        return Ok(CrudeRiskResult {
3050            risk: 0.0,
3051            diseasegradient: Array1::zeros(coeff_len_d),
3052            mortalitygradient: Array1::zeros(coeff_len_m),
3053        });
3054    }
3055
3056    let mut sorted_breaks: Vec<f64> = breakpoints
3057        .iter()
3058        .copied()
3059        .filter(|x| x.is_finite() && *x >= t0 && *x <= t1)
3060        .collect();
3061    sorted_breaks.push(t0);
3062    sorted_breaks.push(t1);
3063    sorted_breaks.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
3064    sorted_breaks.dedup_by(|a, b| (*a - *b).abs() < 1e-6);
3065    if sorted_breaks.len() < 2 {
3066        return Err(SurvivalError::InvalidIntegrationSetup);
3067    }
3068
3069    let mut total_risk = 0.0;
3070    let mut diseasegradient = Array1::zeros(coeff_len_d);
3071    let mut mortalitygradient = Array1::zeros(coeff_len_m);
3072    let nodesweights = gauss_legendre_quadrature();
3073
3074    let mut design_d = Array1::<f64>::zeros(coeff_len_d);
3075    let mut deriv_d = Array1::<f64>::zeros(coeff_len_d);
3076    let mut design_m = Array1::<f64>::zeros(coeff_len_m);
3077
3078    for segment in sorted_breaks.windows(2) {
3079        let a = segment[0];
3080        let b = segment[1];
3081        let center = 0.5 * (b + a);
3082        let halfwidth = 0.5 * (b - a);
3083        if halfwidth <= 0.0 {
3084            continue;
3085        }
3086
3087        for &(x, w) in nodesweights {
3088            let u = center + halfwidth * x;
3089            let (inst_hazard_d, hazard_d, hazard_m) =
3090                eval_at(u, &mut design_d, &mut deriv_d, &mut design_m)?;
3091            if !inst_hazard_d.is_finite() || !hazard_d.is_finite() || !hazard_m.is_finite() {
3092                return Err(SurvivalError::NonFiniteInput);
3093            }
3094            if inst_hazard_d <= 0.0 {
3095                return Err(SurvivalError::NonPositiveHazard);
3096            }
3097
3098            if hazard_d < h_dis_t0 || hazard_m < h_mor_t0 {
3099                return Err(SurvivalError::NonMonotoneCumulativeHazard);
3100            }
3101
3102            let h_dis_cond = hazard_d - h_dis_t0;
3103            let h_mor_cond = hazard_m - h_mor_t0;
3104            let s_total = (-(h_dis_cond + h_mor_cond)).exp();
3105
3106            total_risk += w * inst_hazard_d * s_total * halfwidth;
3107
3108            // d Risk / d beta_d:
3109            //   integral [ d h_d * S_total - h_d * S_total * d H_d ] du
3110            // Contract: design_d[j] = x_j(u) = ∂_{β_j} η_d(u)
3111            //           deriv_d[j]  = ẋ_j(u) = ∂_{β_j} η̇_d(u)
3112            // Then ∂_{β_j} h_d = h_d · x_j + H_d · ẋ_j
3113            let weight = w * s_total * halfwidth;
3114            for j in 0..coeff_len_d {
3115                let d_inst_hazard = inst_hazard_d * design_d[j] + hazard_d * deriv_d[j];
3116                let d_hazard_cond = hazard_d * design_d[j] - h_dis_t0 * design_d_t0[j];
3117                let g = d_inst_hazard - inst_hazard_d * d_hazard_cond;
3118                diseasegradient[j] += weight * g;
3119            }
3120
3121            // d Risk / d beta_m:
3122            //   -integral h_d * S_total * d H_m(u|t0) du
3123            let weight = w * inst_hazard_d * s_total * halfwidth;
3124            for j in 0..coeff_len_m {
3125                let g = -hazard_m * design_m[j] + h_mor_t0 * design_m_t0[j];
3126                mortalitygradient[j] += weight * g;
3127            }
3128        }
3129    }
3130
3131    Ok(CrudeRiskResult {
3132        risk: total_risk,
3133        diseasegradient,
3134        mortalitygradient,
3135    })
3136}
3137
3138impl PirlsWorkingModel for WorkingModelSurvival {
3139    fn update(&mut self, beta: &Coefficients) -> Result<WorkingState, EstimationError> {
3140        self.update_state(beta)
3141    }
3142}
3143
3144#[cfg(test)]
3145mod tests {
3146    use super::*;
3147    use ndarray::{Array1, Array2, Array3, array, s};
3148
3149    /// #932 production single-source parity for the cause-specific Royston-Parmar
3150    /// derivative tower. The earlier cutover added only a gam-math oracle that
3151    /// replicated the production `w_exit`/`w_entry`/`w_derivative` weight formulas
3152    /// verbatim; this module INVOKES production
3153    /// (`evaluate_cause_specific_block`, `cause_specific_hessian_directional_derivative`,
3154    /// `cause_specific_hessian_second_directional_derivative`) and pins each
3155    /// channel against the universal gam-math jet at ≤1e-9, plus an independent
3156    /// central-difference witness of the live third/fourth against the live lower
3157    /// order. The live hand tower is retained (documented performance exception at
3158    /// the code site).
3159    mod jet_cause_specific_production_parity {
3160        use super::*;
3161        use gam_math::jet_scalar::JetScalar;
3162        use gam_math::jet_tower::{
3163            RowNllProgramGeneric, generic_fourth_contracted, generic_row_kernel,
3164            generic_third_contracted,
3165        };
3166
3167        /// The cause-specific row NLL written ONCE through the jet scalar:
3168        /// `ℓ = w·[e^{η1} − 1{entry}·e^{η0} − δ·(η1 + ln s)]`, additively separable
3169        /// over the three predictors (primary 0 = exit index `η1`, primary 1 =
3170        /// entry index `η0`, primary 2 = spline derivative `s > 0`). The entry gate
3171        /// `1{entry}` and event gate `δ` enter as per-row constants.
3172        struct CauseSpecificJetRow {
3173            has_entry: bool,
3174            event: bool,
3175            w: f64,
3176            base: [f64; 3],
3177        }
3178
3179        impl RowNllProgramGeneric<3> for CauseSpecificJetRow {
3180            fn n_rows(&self) -> usize {
3181                1
3182            }
3183            fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
3184                if row != 0 {
3185                    return Err(format!(
3186                        "CauseSpecificJetRow holds exactly one row; got row {row}"
3187                    ));
3188                }
3189                Ok(self.base)
3190            }
3191            fn row_nll_generic<S: JetScalar<3>>(
3192                &self,
3193                row: usize,
3194                p: &[S; 3],
3195            ) -> Result<S, String> {
3196                if row != 0 {
3197                    return Err(format!(
3198                        "CauseSpecificJetRow holds exactly one row; got row {row}"
3199                    ));
3200                }
3201                let mut ell = p[0].exp();
3202                if self.has_entry {
3203                    ell = ell.sub(&p[1].exp());
3204                }
3205                if self.event {
3206                    ell = ell.sub(&p[0].add(&p[2].ln()));
3207                }
3208                Ok(ell.scale(self.w))
3209            }
3210        }
3211
3212        /// A single-row cause-specific block with the design collapsed to the 3×3
3213        /// identity (`x_exit = e0`, `x_entry = e1`, `x_derivative = e2`, zero
3214        /// offsets), so β directly parameterises `(η1, η0, s)` and a coefficient
3215        /// direction IS the predictor-space direction — pinning the per-row
3216        /// β-space kernels against the jet's predictor-space contractions with no
3217        /// design projection in the way.
3218        fn identity_block(w: f64, has_entry: bool, event: bool) -> CauseSpecificRoystonParmarBlock {
3219            let age_entry = if has_entry { 1.0 } else { 0.0 };
3220            CauseSpecificRoystonParmarBlock {
3221                age_entry: array![age_entry],
3222                age_exit: array![2.0],
3223                event_target: array![if event { 1u8 } else { 0u8 }],
3224                sampleweight: array![w],
3225                x_entry: array![[0.0, 1.0, 0.0]],
3226                x_exit: array![[1.0, 0.0, 0.0]],
3227                x_derivative: array![[0.0, 0.0, 1.0]],
3228                offset_eta_entry: array![0.0],
3229                offset_eta_exit: array![0.0],
3230                offset_derivative_exit: array![0.0],
3231                derivative_floor: 0.0,
3232            }
3233        }
3234
3235        fn close(hand: f64, jet: f64, tol: f64, label: &str) {
3236            let band = tol + tol * hand.abs().max(jet.abs());
3237            assert!(
3238                (hand - jet).abs() <= band,
3239                "{label}: hand {hand:+.15e} vs jet {jet:+.15e} (|Δ|={:.3e} band {band:.3e})",
3240                (hand - jet).abs()
3241            );
3242        }
3243
3244        const JET_TOL: f64 = 1e-9;
3245
3246        fn run_corner(has_entry: bool, event: bool) {
3247            // β = (η1, η0, s); s > 0 for the event ln-derivative term.
3248            let beta = array![0.4_f64, -0.3_f64, 1.3_f64];
3249            let d_beta = array![0.7_f64, -0.5_f64, 0.6_f64];
3250            let v_beta = array![-0.2_f64, 0.8_f64, -0.4_f64];
3251            let w = 1.4_f64;
3252            let block = identity_block(w, has_entry, event);
3253            let prog = CauseSpecificJetRow {
3254                has_entry,
3255                event,
3256                w,
3257                base: [beta[0], beta[1], beta[2]],
3258            };
3259            let label = format!("entry={has_entry} event={event}");
3260
3261            // ── Value / gradient / Hessian: LIVE evaluate vs jet ──────────────
3262            let (ll, grad, hess) =
3263                evaluate_cause_specific_block(&block, &beta).expect("evaluate block");
3264            let (jet_v, jet_g, jet_h) = generic_row_kernel(&prog, 0).expect("jet kernel");
3265            close(jet_v, -ll, JET_TOL, &format!("{label} value"));
3266            for a in 0..3 {
3267                close(jet_g[a], -grad[a], JET_TOL, &format!("{label} grad[{a}]"));
3268                for b in 0..3 {
3269                    close(jet_h[a][b], hess[[a, b]], JET_TOL, &format!("{label} H[{a}][{b}]"));
3270                }
3271            }
3272
3273            // ── Third: LIVE directional derivative vs jet ─────────────────────
3274            let dh = cause_specific_hessian_directional_derivative(&block, &beta, &d_beta)
3275                .expect("live third");
3276            let dir = [d_beta[0], d_beta[1], d_beta[2]];
3277            let jet_t3 = generic_third_contracted(&prog, 0, &dir).expect("jet third");
3278            for a in 0..3 {
3279                for b in 0..3 {
3280                    close(
3281                        jet_t3[a][b],
3282                        dh[[a, b]],
3283                        JET_TOL,
3284                        &format!("{label} third[{a}][{b}]"),
3285                    );
3286                }
3287            }
3288
3289            // ── Fourth: LIVE second directional derivative vs jet ─────────────
3290            let d2h =
3291                cause_specific_hessian_second_directional_derivative(&block, &beta, &d_beta, &v_beta)
3292                    .expect("live fourth");
3293            let uu = [d_beta[0], d_beta[1], d_beta[2]];
3294            let vv = [v_beta[0], v_beta[1], v_beta[2]];
3295            let jet_t4 = generic_fourth_contracted(&prog, 0, &uu, &vv).expect("jet fourth");
3296            for a in 0..3 {
3297                for b in 0..3 {
3298                    close(
3299                        jet_t4[a][b],
3300                        d2h[[a, b]],
3301                        JET_TOL,
3302                        &format!("{label} fourth[{a}][{b}]"),
3303                    );
3304                }
3305            }
3306
3307            // ── Independent FD witness (NO jet) ───────────────────────────────
3308            // ∂_d_beta H via central difference of the LIVE evaluate Hessian.
3309            let h_fd = 1e-5;
3310            let bp = &beta + &(&d_beta * h_fd);
3311            let bm = &beta - &(&d_beta * h_fd);
3312            let (_, _, hp) = evaluate_cause_specific_block(&block, &bp).expect("evaluate +");
3313            let (_, _, hm) = evaluate_cause_specific_block(&block, &bm).expect("evaluate -");
3314            for a in 0..3 {
3315                for b in 0..3 {
3316                    let fd = (hp[[a, b]] - hm[[a, b]]) / (2.0 * h_fd);
3317                    close(dh[[a, b]], fd, 1e-5, &format!("{label} FD third[{a}][{b}]"));
3318                }
3319            }
3320            // ∂_v of the LIVE third (fixed direction d_beta) vs the LIVE fourth.
3321            let dhp = cause_specific_hessian_directional_derivative(&block, &bp_along(&beta, &v_beta, h_fd), &d_beta)
3322                .expect("live third +");
3323            let dhm = cause_specific_hessian_directional_derivative(&block, &bm_along(&beta, &v_beta, h_fd), &d_beta)
3324                .expect("live third -");
3325            for a in 0..3 {
3326                for b in 0..3 {
3327                    let fd = (dhp[[a, b]] - dhm[[a, b]]) / (2.0 * h_fd);
3328                    close(d2h[[a, b]], fd, 1e-5, &format!("{label} FD fourth[{a}][{b}]"));
3329                }
3330            }
3331        }
3332
3333        fn bp_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
3334            beta + &(v * h)
3335        }
3336        fn bm_along(beta: &Array1<f64>, v: &Array1<f64>, h: f64) -> Array1<f64> {
3337            beta - &(v * h)
3338        }
3339
3340        /// The LIVE cause-specific value / gradient / Hessian / third / fourth hand
3341        /// tower reproduces the universal gam-math jet at ≤1e-9, and the live
3342        /// third/fourth reproduce an independent central-difference of the live
3343        /// lower order — across all four (event × entry) corners that gate the
3344        /// entry and event predictor channels on and off.
3345        #[test]
3346        fn cause_specific_live_tower_matches_jet_and_fd() {
3347            for &has_entry in &[false, true] {
3348                for &event in &[false, true] {
3349                    run_corner(has_entry, event);
3350                }
3351            }
3352        }
3353    }
3354
3355    #[test]
3356    fn competing_risks_cif_constant_hazard_matches_closed_form() {
3357        let times = array![0.0, 2.0, 5.0, 10.0];
3358        let disease_rates = [0.12, 0.06];
3359        let death_rates = [0.05, 0.02];
3360        let cumulative = Array3::from_shape_fn((2, 2, times.len()), |(endpoint, row, time_idx)| {
3361            let rate = if endpoint == 0 {
3362                disease_rates[row]
3363            } else {
3364                death_rates[row]
3365            };
3366            rate * times[time_idx]
3367        });
3368
3369        let result =
3370            assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
3371
3372        for row in 0..2 {
3373            let total_rate = disease_rates[row] + death_rates[row];
3374            for time_idx in 0..times.len() {
3375                let failure = 1.0 - (-total_rate * times[time_idx]).exp();
3376                let expected_disease = disease_rates[row] / total_rate * failure;
3377                let expected_death = death_rates[row] / total_rate * failure;
3378                assert!((result.cif[0][[row, time_idx]] - expected_disease).abs() < 1e-12);
3379                assert!((result.cif[1][[row, time_idx]] - expected_death).abs() < 1e-12);
3380                assert!(
3381                    (result.cif[0][[row, time_idx]]
3382                        + result.cif[1][[row, time_idx]]
3383                        + result.overall_survival[[row, time_idx]]
3384                        - 1.0)
3385                        .abs()
3386                        < 1e-12
3387                );
3388            }
3389        }
3390    }
3391
3392    #[test]
3393    fn competing_risks_cif_rejects_nonmonotone_hazards() {
3394        let times = array![0.0, 1.0, 2.0];
3395        let cumulative = Array3::from_shape_vec((1, 1, 3), vec![0.0, 0.2, 0.1]).expect("shape");
3396        let err = assemble_competing_risks_cif(times.view(), cumulative.view())
3397            .expect_err("nonmonotone cumulative hazard should be rejected");
3398        assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
3399    }
3400
3401    #[test]
3402    fn competing_risks_cif_plateaus_and_three_causes_conserve_probability() {
3403        let times = array![0.0, 1.0, 3.0, 7.0, 12.0];
3404        let cumulative = Array3::from_shape_vec(
3405            (3, 2, 5),
3406            vec![
3407                // cause 1
3408                0.0, 0.2, 0.2, 0.5, 1.1, 0.0, 0.0, 0.4, 0.4, 0.9, // cause 2
3409                0.0, 0.1, 0.3, 0.3, 0.7, 0.0, 0.2, 0.2, 0.8, 0.8, // cause 3
3410                0.0, 0.0, 0.2, 0.6, 0.6, 0.0, 0.1, 0.5, 0.5, 1.5,
3411            ],
3412        )
3413        .expect("shape");
3414
3415        let result =
3416            assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
3417
3418        for row in 0..2 {
3419            for time_idx in 0..times.len() {
3420                let total_cif = result.cif[0][[row, time_idx]]
3421                    + result.cif[1][[row, time_idx]]
3422                    + result.cif[2][[row, time_idx]];
3423                assert!(
3424                    (total_cif + result.overall_survival[[row, time_idx]] - 1.0).abs() < 1e-12,
3425                    "probability mass mismatch at row={row}, time_idx={time_idx}"
3426                );
3427                assert!((0.0..=1.0).contains(&result.overall_survival[[row, time_idx]]));
3428                for cause in 0..3 {
3429                    assert!((0.0..=1.0).contains(&result.cif[cause][[row, time_idx]]));
3430                    if time_idx > 0 {
3431                        assert!(
3432                            result.cif[cause][[row, time_idx]] + 1e-12
3433                                >= result.cif[cause][[row, time_idx - 1]],
3434                            "CIF decreased for cause={cause}, row={row}, time_idx={time_idx}"
3435                        );
3436                    }
3437                }
3438            }
3439        }
3440
3441        // Cause 1 is flat between t=1 and t=3 for row 0, but other causes
3442        // fail in that interval; its CIF must remain exactly flat.
3443        assert_eq!(result.cif[0][[0, 1]], result.cif[0][[0, 2]]);
3444        // All causes are flat between t=3 and t=7 for row 1 except cause 2;
3445        // causes 1 and 3 must not move.
3446        assert_eq!(result.cif[0][[1, 2]], result.cif[0][[1, 3]]);
3447        assert_eq!(result.cif[2][[1, 2]], result.cif[2][[1, 3]]);
3448    }
3449
3450    #[test]
3451    fn competing_risks_cif_rejects_bad_time_grids_and_nonfinite_hazards() {
3452        let cumulative = Array3::zeros((2, 1, 2));
3453
3454        for times in [array![0.0, 0.0], array![1.0, 0.5], array![-1.0, 1.0]] {
3455            let err = assemble_competing_risks_cif(times.view(), cumulative.view())
3456                .expect_err("bad time grid should be rejected");
3457            assert!(matches!(err, SurvivalError::InvalidTimeGrid));
3458        }
3459
3460        let times = array![0.0, 1.0];
3461        let nonfinite = Array3::from_shape_vec((1, 1, 2), vec![0.0, f64::NAN]).expect("shape");
3462        let err = assemble_competing_risks_cif(times.view(), nonfinite.view())
3463            .expect_err("nonfinite hazard should be rejected");
3464        assert!(matches!(err, SurvivalError::NonFiniteInput));
3465    }
3466
3467    #[test]
3468    fn competing_risks_cif_extreme_hazards_remain_bounded() {
3469        let times = array![0.0, 1.0, 2.0];
3470        let cumulative =
3471            Array3::from_shape_vec((2, 1, 3), vec![0.0, 500.0, 1000.0, 0.0, 250.0, 1000.0])
3472                .expect("shape");
3473
3474        let result =
3475            assemble_competing_risks_cif(times.view(), cumulative.view()).expect("assemble CIF");
3476
3477        for value in result
3478            .cif
3479            .iter()
3480            .flat_map(|m| m.iter())
3481            .chain(result.overall_survival.iter())
3482        {
3483            assert!(value.is_finite());
3484            assert!((0.0..=1.0).contains(value));
3485        }
3486        assert!((result.cif[0][[0, 2]] + result.cif[1][[0, 2]] - 1.0).abs() < 1e-12);
3487        assert_eq!(result.overall_survival[[0, 2]], 0.0);
3488    }
3489
3490    fn toy_penalties() -> PenaltyBlocks {
3491        let s = array![[2.0, 0.5], [0.5, 3.0]];
3492        PenaltyBlocks::new(vec![PenaltyBlock {
3493            matrix: s,
3494            lambda: 1.7,
3495            range: 1..3,
3496            nullspace_dim: 0,
3497        }])
3498    }
3499
3500    fn survival_inputs<'a>(
3501        age_entry: &'a Array1<f64>,
3502        age_exit: &'a Array1<f64>,
3503        event_target: &'a Array1<u8>,
3504        event_competing: &'a Array1<u8>,
3505        sampleweight: &'a Array1<f64>,
3506        x_entry: &'a Array2<f64>,
3507        x_exit: &'a Array2<f64>,
3508        x_derivative: &'a Array2<f64>,
3509    ) -> SurvivalEngineInputs<'a> {
3510        SurvivalEngineInputs {
3511            age_entry: age_entry.view(),
3512            age_exit: age_exit.view(),
3513            event_target: event_target.view(),
3514            event_competing: event_competing.view(),
3515            sampleweight: sampleweight.view(),
3516            x_entry: x_entry.view(),
3517            x_exit: x_exit.view(),
3518            x_derivative: x_derivative.view(),
3519            monotonicity_constraint_rows: None,
3520            monotonicity_constraint_offsets: None,
3521        }
3522    }
3523
3524    fn survival_model(
3525        inputs: SurvivalEngineInputs<'_>,
3526        penalties: PenaltyBlocks,
3527        monotonicity: SurvivalMonotonicityPenalty,
3528        spec: SurvivalSpec,
3529    ) -> Result<WorkingModelSurvival, SurvivalError> {
3530        WorkingModelSurvival::from_engine_inputs(inputs, penalties, monotonicity, spec)
3531    }
3532
3533    fn survival_model_with_offsets(
3534        inputs: SurvivalEngineInputs<'_>,
3535        offsets: Option<SurvivalBaselineOffsets<'_>>,
3536        penalties: PenaltyBlocks,
3537        monotonicity: SurvivalMonotonicityPenalty,
3538        spec: SurvivalSpec,
3539    ) -> Result<WorkingModelSurvival, SurvivalError> {
3540        WorkingModelSurvival::from_engine_inputswith_offsets(
3541            inputs,
3542            offsets,
3543            penalties,
3544            monotonicity,
3545            spec,
3546        )
3547    }
3548
3549    #[test]
3550    fn penaltyhessian_matchesgradient_jacobian() {
3551        let penalties = toy_penalties();
3552        let beta = array![10.0, -0.3, 1.2, 7.0];
3553
3554        let grad = penalties.gradient(&beta);
3555        let h = penalties.hessian(beta.len());
3556        let b_block = beta.slice(s![1..3]).to_owned();
3557        let expected = 1.7 * array![[2.0, 0.5], [0.5, 3.0]].dot(&b_block);
3558
3559        assert!((grad[1] - expected[0]).abs() < 1e-12);
3560        assert!((grad[2] - expected[1]).abs() < 1e-12);
3561        assert!((h[[1, 1]] - 1.7 * 2.0).abs() < 1e-12);
3562        assert!((h[[1, 2]] - 1.7 * 0.5).abs() < 1e-12);
3563        assert!((h[[2, 1]] - 1.7 * 0.5).abs() < 1e-12);
3564        assert!((h[[2, 2]] - 1.7 * 3.0).abs() < 1e-12);
3565    }
3566
3567    #[test]
3568    fn penaltygradient_matches_deviance_finite_difference() {
3569        let penalties = toy_penalties();
3570        let beta = array![10.0, -0.3, 1.2, 7.0];
3571        let grad = penalties.gradient(&beta);
3572        let eps = 1e-7;
3573
3574        for idx in 0..beta.len() {
3575            let mut plus = beta.clone();
3576            let mut minus = beta.clone();
3577            plus[idx] += eps;
3578            minus[idx] -= eps;
3579            let fd = (penalties.deviance(&plus) - penalties.deviance(&minus)) / (2.0 * eps);
3580            assert_eq!(
3581                grad[idx].signum(),
3582                fd.signum(),
3583                "gradient/deviance sign mismatch at idx={idx}: grad={} fd={fd}",
3584                grad[idx]
3585            );
3586            assert!(
3587                (grad[idx] - fd).abs() < 1e-6,
3588                "gradient/deviance mismatch at idx={idx}: grad={} fd={fd}",
3589                grad[idx]
3590            );
3591        }
3592    }
3593
3594    #[test]
3595    fn zero_offsets_match_default_survival_state() {
3596        let age_entry = array![1.0_f64, 2.0_f64];
3597        let age_exit = array![2.0_f64, 3.5_f64];
3598        let event_target = array![1u8, 0u8];
3599        let event_competing = array![0u8, 0u8];
3600        let sampleweight = array![1.0, 1.0];
3601        let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
3602        let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
3603        let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
3604        let penalties = PenaltyBlocks::new(Vec::new());
3605        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3606        let beta = array![-1.0, 0.8];
3607
3608        let base = survival_model(
3609            survival_inputs(
3610                &age_entry,
3611                &age_exit,
3612                &event_target,
3613                &event_competing,
3614                &sampleweight,
3615                &x_entry,
3616                &x_exit,
3617                &x_derivative,
3618            ),
3619            penalties.clone(),
3620            mono,
3621            SurvivalSpec::Net,
3622        )
3623        .expect("construct base survival model");
3624
3625        let zero_offsets = survival_model_with_offsets(
3626            survival_inputs(
3627                &age_entry,
3628                &age_exit,
3629                &event_target,
3630                &event_competing,
3631                &sampleweight,
3632                &x_entry,
3633                &x_exit,
3634                &x_derivative,
3635            ),
3636            Some(SurvivalBaselineOffsets {
3637                eta_entry: array![0.0, 0.0].view(),
3638                eta_exit: array![0.0, 0.0].view(),
3639                derivative_exit: array![0.0, 0.0].view(),
3640            }),
3641            penalties,
3642            mono,
3643            SurvivalSpec::Net,
3644        )
3645        .expect("construct offset survival model");
3646
3647        let state_base = base.update_state(&beta).expect("base state");
3648        let statezero = zero_offsets.update_state(&beta).expect("zero-offset state");
3649        assert!((state_base.deviance - statezero.deviance).abs() < 1e-12);
3650        assert!(
3651            state_base
3652                .gradient
3653                .iter()
3654                .zip(statezero.gradient.iter())
3655                .all(|(a, b)| (a - b).abs() < 1e-12)
3656        );
3657    }
3658
3659    #[test]
3660    fn competing_risk_cause_labels_collapse_to_pooled_baseline_indicator() {
3661        // Regression for #378: the joint competing-risks Weibull path seeds a
3662        // shared single-hazard baseline working model from the dataset's event
3663        // *labels* {0 = censored, 1 = cause 1, 2 = cause 2}. The single-hazard
3664        // engine's `event_target` contract is binary {0, 1}, so feeding the raw
3665        // cause labels straight through used to bail out of construction via the
3666        // `event_target > 1` guard and surface as the misleading
3667        // `SurvivalError::NonFiniteInput` ("inputs contain non-finite values"),
3668        // even though every input value is finite. The fix (a) reports a
3669        // multi-cause label as the actionable `EventCodeInvalid`, never the
3670        // misleading "non-finite", and (b) projects cause labels to the
3671        // any-event {0, 1} indicator via the single-source-of-truth
3672        // `pooled_any_event_indicator` before constructing the pooled baseline.
3673        // This pins both halves of that contract.
3674        let age_entry = array![0.0_f64, 0.0, 0.0, 0.0];
3675        let age_exit = array![1.2_f64, 0.8, 2.1, 1.5];
3676        // Competing-risks cause labels: censored, cause 1, cause 2, censored.
3677        let cause_labels = array![0u8, 1u8, 2u8, 0u8];
3678        let event_competing = Array1::<u8>::zeros(cause_labels.len());
3679        let sampleweight = array![1.0_f64, 1.0, 1.0, 1.0];
3680        let x_entry = array![
3681            [1.0, age_entry[0].max(1e-8).ln()],
3682            [1.0, age_entry[1].max(1e-8).ln()],
3683            [1.0, age_entry[2].max(1e-8).ln()],
3684            [1.0, age_entry[3].max(1e-8).ln()],
3685        ];
3686        let x_exit = array![
3687            [1.0, age_exit[0].ln()],
3688            [1.0, age_exit[1].ln()],
3689            [1.0, age_exit[2].ln()],
3690            [1.0, age_exit[3].ln()],
3691        ];
3692        let x_derivative = array![
3693            [0.0, 1.0 / age_exit[0]],
3694            [0.0, 1.0 / age_exit[1]],
3695            [0.0, 1.0 / age_exit[2]],
3696            [0.0, 1.0 / age_exit[3]],
3697        ];
3698        let penalties = PenaltyBlocks::new(Vec::new());
3699        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3700
3701        // Raw cause labels {0,1,2} violate the single-hazard binary contract and
3702        // must be rejected -- but as an *actionable* `EventCodeInvalid`, NOT the
3703        // misleading `NonFiniteInput`: the labels are finite, they merely need
3704        // projecting. (The old fix left this surfacing as "non-finite".)
3705        let raw = survival_model(
3706            survival_inputs(
3707                &age_entry,
3708                &age_exit,
3709                &cause_labels,
3710                &event_competing,
3711                &sampleweight,
3712                &x_entry,
3713                &x_exit,
3714                &x_derivative,
3715            ),
3716            penalties.clone(),
3717            mono,
3718            SurvivalSpec::Net,
3719        );
3720        assert!(
3721            matches!(raw, Err(SurvivalError::EventCodeInvalid { .. })),
3722            "raw competing-risks cause labels must be rejected as EventCodeInvalid (not NonFiniteInput), got {raw:?}"
3723        );
3724
3725        // The pooled-baseline projection the workflow now performs through the
3726        // single source of truth: any observed event (any cause) -> {0, 1}.
3727        let any_event = pooled_any_event_indicator(cause_labels.view());
3728        assert_eq!(any_event, array![0u8, 1u8, 1u8, 0u8]);
3729        // And the per-cause projection that seeds each cause-specific block.
3730        assert_eq!(
3731            cause_specific_event_indicator(cause_labels.view(), 1),
3732            array![0u8, 1u8, 0u8, 0u8]
3733        );
3734        assert_eq!(
3735            cause_specific_event_indicator(cause_labels.view(), 2),
3736            array![0u8, 0u8, 1u8, 0u8]
3737        );
3738        let model = survival_model(
3739            survival_inputs(
3740                &age_entry,
3741                &age_exit,
3742                &any_event,
3743                &event_competing,
3744                &sampleweight,
3745                &x_entry,
3746                &x_exit,
3747                &x_derivative,
3748            ),
3749            penalties,
3750            mono,
3751            SurvivalSpec::Net,
3752        )
3753        .expect("pooled any-event baseline model must construct from competing-risks data");
3754
3755        // The constructed pooled baseline must yield a finite working state, so
3756        // the downstream baseline-seeding PIRLS loop has something to optimize.
3757        let beta = array![-1.0_f64, 0.8];
3758        let state = model.update_state(&beta).expect("pooled baseline state");
3759        assert!(
3760            state.deviance.is_finite(),
3761            "pooled baseline deviance must be finite, got {}",
3762            state.deviance
3763        );
3764        assert!(
3765            state.gradient.iter().all(|g| g.is_finite()),
3766            "pooled baseline gradient must be finite"
3767        );
3768    }
3769
3770    #[test]
3771    fn offset_channel_residuals_match_central_fd_of_nll() {
3772        // Three observations: two events (non-origin entry and origin entry)
3773        // and one censored row. This exercises every nonzero channel at least
3774        // once: r_exit from all rows, r_entry only from the first (has entry
3775        // interval), r_derivative only from events.
3776        let age_entry = array![0.5_f64, 0.0, 0.3];
3777        let age_exit = array![1.4_f64, 1.0, 2.0];
3778        let event_target = array![1u8, 1u8, 0u8];
3779        let event_competing = array![0u8, 0u8, 0u8];
3780        let sampleweight = array![1.0_f64, 2.5, 0.7];
3781        let x_entry = array![
3782            [1.0, age_entry[0].ln()],
3783            [1.0, age_entry[1].max(1e-8).ln()],
3784            [1.0, age_entry[2].ln()]
3785        ];
3786        let x_exit = array![
3787            [1.0, age_exit[0].ln()],
3788            [1.0, age_exit[1].ln()],
3789            [1.0, age_exit[2].ln()]
3790        ];
3791        let x_derivative = array![
3792            [0.0, 1.0 / age_exit[0]],
3793            [0.0, 1.0 / age_exit[1]],
3794            [0.0, 1.0 / age_exit[2]]
3795        ];
3796        // Baseline offsets chosen so η_entry, η_exit, s are all comfortably
3797        // away from overflow / monotonicity-violation boundaries.
3798        let o_entry = array![0.2_f64, 0.0, 0.1];
3799        let o_exit = array![0.4_f64, 0.5, 0.7];
3800        let o_deriv = array![0.3_f64, 0.8, 0.5];
3801        let penalties = PenaltyBlocks::new(Vec::new());
3802        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3803        let beta = array![-0.7_f64, 0.6];
3804
3805        let build = |o_e: &Array1<f64>, o_x: &Array1<f64>, o_d: &Array1<f64>| {
3806            survival_model_with_offsets(
3807                survival_inputs(
3808                    &age_entry,
3809                    &age_exit,
3810                    &event_target,
3811                    &event_competing,
3812                    &sampleweight,
3813                    &x_entry,
3814                    &x_exit,
3815                    &x_derivative,
3816                ),
3817                Some(SurvivalBaselineOffsets {
3818                    eta_entry: o_e.view(),
3819                    eta_exit: o_x.view(),
3820                    derivative_exit: o_d.view(),
3821                }),
3822                penalties.clone(),
3823                mono,
3824                SurvivalSpec::Net,
3825            )
3826            .expect("model build")
3827        };
3828
3829        let base = build(&o_entry, &o_exit, &o_deriv);
3830        let resid = base
3831            .offset_channel_residuals(&beta)
3832            .expect("offset residuals");
3833        assert_eq!(resid.exit.len(), 3);
3834        assert_eq!(resid.entry.len(), 3);
3835        assert_eq!(resid.derivative.len(), 3);
3836
3837        // NLL equals half the deviance returned by update_state; that is the
3838        // exact unpenalized loss whose offset partials r_{X,E,D} encode.
3839        let nll = |m: &WorkingModelSurvival| 0.5 * m.update_state(&beta).expect("state").deviance;
3840        let h = 1e-6;
3841
3842        // Row 1 (origin entry, event=1) has no entry interval, so r_entry[1]
3843        // must be exactly 0. Row 2 (censored) has r_deriv[2] exactly 0. Check
3844        // those identities before FD comparison on the nonzero elements.
3845        assert_eq!(resid.entry[1], 0.0);
3846        assert_eq!(resid.derivative[2], 0.0);
3847
3848        for i in 0..3 {
3849            // exit channel: perturb o_exit[i] alone.
3850            {
3851                let mut op = o_exit.clone();
3852                let mut om = o_exit.clone();
3853                op[i] += h;
3854                om[i] -= h;
3855                let fd = (nll(&build(&o_entry, &op, &o_deriv))
3856                    - nll(&build(&o_entry, &om, &o_deriv)))
3857                    / (2.0 * h);
3858                assert!(
3859                    (resid.exit[i] - fd).abs() < 1e-6,
3860                    "∂NLL/∂o_X[{i}]: analytic={:.6e} fd={:.6e}",
3861                    resid.exit[i],
3862                    fd
3863                );
3864            }
3865            // entry channel: only row 0 has an entry interval; for rows with
3866            // entry_at_origin the offset contributes nothing to NLL and FD
3867            // must also be exactly 0 to numerical precision.
3868            {
3869                let mut op = o_entry.clone();
3870                let mut om = o_entry.clone();
3871                op[i] += h;
3872                om[i] -= h;
3873                let fd = (nll(&build(&op, &o_exit, &o_deriv))
3874                    - nll(&build(&om, &o_exit, &o_deriv)))
3875                    / (2.0 * h);
3876                assert!(
3877                    (resid.entry[i] - fd).abs() < 1e-6,
3878                    "∂NLL/∂o_E[{i}]: analytic={:.6e} fd={:.6e}",
3879                    resid.entry[i],
3880                    fd
3881                );
3882            }
3883            // derivative channel: only event rows contribute.
3884            {
3885                let mut op = o_deriv.clone();
3886                let mut om = o_deriv.clone();
3887                op[i] += h;
3888                om[i] -= h;
3889                let fd = (nll(&build(&o_entry, &o_exit, &op))
3890                    - nll(&build(&o_entry, &o_exit, &om)))
3891                    / (2.0 * h);
3892                assert!(
3893                    (resid.derivative[i] - fd).abs() < 1e-6,
3894                    "∂NLL/∂o_D[{i}]: analytic={:.6e} fd={:.6e}",
3895                    resid.derivative[i],
3896                    fd
3897                );
3898            }
3899        }
3900    }
3901
3902    #[test]
3903    fn offset_channel_residuals_respect_zero_sampleweight() {
3904        let age_entry = array![1.0_f64, 2.0];
3905        let age_exit = array![2.0_f64, 3.5];
3906        let event_target = array![1u8, 1u8];
3907        let event_competing = array![0u8, 0u8];
3908        let sampleweight = array![0.0_f64, 1.2]; // row 0 is excluded by weight
3909        let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
3910        let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
3911        let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
3912        let penalties = PenaltyBlocks::new(Vec::new());
3913        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3914        let beta = array![-1.0_f64, 0.8];
3915
3916        let model = survival_model_with_offsets(
3917            survival_inputs(
3918                &age_entry,
3919                &age_exit,
3920                &event_target,
3921                &event_competing,
3922                &sampleweight,
3923                &x_entry,
3924                &x_exit,
3925                &x_derivative,
3926            ),
3927            Some(SurvivalBaselineOffsets {
3928                eta_entry: array![0.0_f64, 0.1].view(),
3929                eta_exit: array![0.0_f64, 0.2].view(),
3930                derivative_exit: array![0.0_f64, 0.1].view(),
3931            }),
3932            penalties,
3933            mono,
3934            SurvivalSpec::Net,
3935        )
3936        .expect("model");
3937        let r = model.offset_channel_residuals(&beta).expect("resid");
3938        // Row 0 (sampleweight=0) must contribute zero in every channel.
3939        assert_eq!(r.exit[0], 0.0);
3940        assert_eq!(r.entry[0], 0.0);
3941        assert_eq!(r.derivative[0], 0.0);
3942        // Row 1 must still carry a nonzero exit-channel residual.
3943        assert!(r.exit[1] != 0.0);
3944    }
3945
3946    #[test]
3947    fn offset_channel_residuals_reject_beta_dim_mismatch() {
3948        let age_entry = array![1.0_f64];
3949        let age_exit = array![2.0_f64];
3950        let event_target = array![1u8];
3951        let event_competing = array![0u8];
3952        let sampleweight = array![1.0_f64];
3953        let x_entry = array![[1.0, 0.0]];
3954        let x_exit = array![[1.0, 0.7]];
3955        let x_derivative = array![[0.0, 0.5]];
3956        let penalties = PenaltyBlocks::new(Vec::new());
3957        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3958        let model = survival_model(
3959            survival_inputs(
3960                &age_entry,
3961                &age_exit,
3962                &event_target,
3963                &event_competing,
3964                &sampleweight,
3965                &x_entry,
3966                &x_exit,
3967                &x_derivative,
3968            ),
3969            penalties,
3970            mono,
3971            SurvivalSpec::Net,
3972        )
3973        .expect("model");
3974        let bad_beta = array![0.0_f64]; // should be length 2
3975        let err = model
3976            .offset_channel_residuals(&bad_beta)
3977            .expect_err("mismatch must error");
3978        match err {
3979            EstimationError::InvalidInput(msg) => {
3980                assert!(msg.contains("beta dimension mismatch"), "msg={msg}")
3981            }
3982            other => panic!("expected InvalidInput, got {other:?}"),
3983        }
3984    }
3985
3986    #[test]
3987    fn crudespec_is_rejected_by_one_hazard_engine() {
3988        let age_entry = array![1.0_f64];
3989        let age_exit = array![2.0_f64];
3990        let event_target = array![0u8];
3991        let event_competing = array![1u8];
3992        let sampleweight = array![1.0];
3993        let x_entry = array![[0.1]];
3994        let x_exit = array![[0.4]];
3995        let x_derivative = array![[1.0]];
3996        let penalties = PenaltyBlocks::new(Vec::new());
3997        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
3998
3999        let err = survival_model(
4000            survival_inputs(
4001                &age_entry,
4002                &age_exit,
4003                &event_target,
4004                &event_competing,
4005                &sampleweight,
4006                &x_entry,
4007                &x_exit,
4008                &x_derivative,
4009            ),
4010            penalties,
4011            mono,
4012            SurvivalSpec::Crude,
4013        )
4014        .expect_err("crude fitting should be rejected by the one-hazard engine");
4015        assert!(matches!(err, SurvivalError::UnsupportedSpec("crude")));
4016    }
4017
4018    #[test]
4019    fn nonstructural_models_require_explicit_monotonicity_collocation() {
4020        let age_entry = array![1.0_f64, 1.5_f64];
4021        let age_exit = array![2.0_f64, 2.5_f64];
4022        let event_target = array![0u8, 0u8];
4023        let event_competing = array![0u8, 1u8];
4024        let sampleweight = array![1.0, 1.0];
4025        let x_entry = array![[0.2], [0.1]];
4026        let x_exit = array![[0.3], [0.2]];
4027        let x_derivative = array![[1.0], [1.0]];
4028
4029        let model = survival_model(
4030            survival_inputs(
4031                &age_entry,
4032                &age_exit,
4033                &event_target,
4034                &event_competing,
4035                &sampleweight,
4036                &x_entry,
4037                &x_exit,
4038                &x_derivative,
4039            ),
4040            PenaltyBlocks::new(Vec::new()),
4041            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4042            SurvivalSpec::Net,
4043        )
4044        .expect("construct censored survival model");
4045
4046        assert!(
4047            model.monotonicity_linear_constraints().is_none(),
4048            "non-structural survival models must not fabricate rowwise monotonicity constraints"
4049        );
4050    }
4051
4052    #[test]
4053    fn decreasing_interval_is_rejectedwithout_target_events() {
4054        let age_entry = array![1.0_f64];
4055        let age_exit = array![2.0_f64];
4056        let event_target = array![0u8];
4057        let event_competing = array![0u8];
4058        let sampleweight = array![1.0];
4059        let x_entry = array![[0.5]];
4060        let x_exit = array![[0.0]];
4061        let x_derivative = array![[1.0]];
4062
4063        let model = survival_model(
4064            survival_inputs(
4065                &age_entry,
4066                &age_exit,
4067                &event_target,
4068                &event_competing,
4069                &sampleweight,
4070                &x_entry,
4071                &x_exit,
4072                &x_derivative,
4073            ),
4074            PenaltyBlocks::new(Vec::new()),
4075            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4076            SurvivalSpec::Net,
4077        )
4078        .expect("construct censored survival model");
4079
4080        let err = model
4081            .update_state(&array![1.0])
4082            .expect_err("decreasing cumulative hazard increment should be rejected");
4083        assert!(
4084            err.to_string().contains("cumulative hazard decreased"),
4085            "unexpected error: {err}"
4086        );
4087    }
4088
4089    fn smooth_crude_risk(beta_d: f64, beta_m: f64) -> CrudeRiskResult {
4090        calculate_crude_risk_quadrature(
4091            0.0,
4092            1.0,
4093            &[0.0, 1.0],
4094            beta_d.exp(),
4095            beta_m.exp(),
4096            array![1.0].view(),
4097            array![1.0].view(),
4098            |u, design_d, deriv_d, design_m| {
4099                let cumulative_d = beta_d.exp() * (1.0 + 0.2 * u);
4100                let cumulative_m = beta_m.exp() * (1.0 + 0.1 * u);
4101                let inst_hazard_d = 0.2 * beta_d.exp();
4102                design_d[0] = 1.0;
4103                // η_d = β_d + ln(1 + 0.2u), so η̇_d = 0.2/(1+0.2u)
4104                // which does not depend on β_d → ∂_{β_d} η̇_d = 0
4105                deriv_d[0] = 0.0;
4106                design_m[0] = 1.0;
4107                Ok((inst_hazard_d, cumulative_d, cumulative_m))
4108            },
4109        )
4110        .expect("smooth crude-risk quadrature should succeed")
4111    }
4112
4113    #[test]
4114    fn crude_riskgradient_matches_monotoneobjective() {
4115        let beta_d = -0.2_f64;
4116        let beta_m = -0.5_f64;
4117        let result = smooth_crude_risk(beta_d, beta_m);
4118        let eps = 1e-6;
4119
4120        let fd_d = (smooth_crude_risk(beta_d + eps, beta_m).risk
4121            - smooth_crude_risk(beta_d - eps, beta_m).risk)
4122            / (2.0 * eps);
4123        let fd_m = (smooth_crude_risk(beta_d, beta_m + eps).risk
4124            - smooth_crude_risk(beta_d, beta_m - eps).risk)
4125            / (2.0 * eps);
4126
4127        assert!(
4128            (result.diseasegradient[0] - fd_d).abs() < 1e-5,
4129            "disease gradient mismatch for monotone crude risk: analytic={} fd={fd_d}",
4130            result.diseasegradient[0]
4131        );
4132        assert!(
4133            (result.mortalitygradient[0] - fd_m).abs() < 1e-5,
4134            "mortality gradient mismatch for monotone crude risk: analytic={} fd={fd_m}",
4135            result.mortalitygradient[0]
4136        );
4137    }
4138
4139    #[test]
4140    fn survival_working_state_is_ridge_free() {
4141        let age_entry = array![1.0_f64, 2.0_f64];
4142        let age_exit = array![2.0_f64, 3.5_f64];
4143        let event_target = array![1u8, 0u8];
4144        let event_competing = array![0u8, 0u8];
4145        let sampleweight = array![1.0, 1.0];
4146        let x_entry = array![[1.0, age_entry[0].ln()], [1.0, age_entry[1].ln()]];
4147        let x_exit = array![[1.0, age_exit[0].ln()], [1.0, age_exit[1].ln()]];
4148        let x_derivative = array![[0.0, 1.0 / age_exit[0]], [0.0, 1.0 / age_exit[1]]];
4149        let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4150            matrix: array![[2.0]],
4151            lambda: 1.7,
4152            range: 1..2,
4153            nullspace_dim: 0,
4154        }]);
4155        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4156        let beta = array![-1.2, 0.4];
4157
4158        let model = survival_model(
4159            survival_inputs(
4160                &age_entry,
4161                &age_exit,
4162                &event_target,
4163                &event_competing,
4164                &sampleweight,
4165                &x_entry,
4166                &x_exit,
4167                &x_derivative,
4168            ),
4169            penalties.clone(),
4170            mono,
4171            SurvivalSpec::Net,
4172        )
4173        .expect("construct survival model");
4174
4175        let state = model.update_state(&beta).expect("survival state");
4176        assert_eq!(
4177            state.ridge_used, 0.0,
4178            "survival objective must not fuse a coefficient ridge"
4179        );
4180        let expected_penalty = penalties.deviance(&beta);
4181        assert!(
4182            (state.penalty_term - expected_penalty).abs() < 1e-12,
4183            "penalty_term mismatch: state={} expected={}",
4184            state.penalty_term,
4185            expected_penalty
4186        );
4187    }
4188
4189    #[test]
4190    fn negative_penalty_lambda_is_rejected() {
4191        let age_entry = array![1.0_f64];
4192        let age_exit = array![2.0_f64];
4193        let event_target = array![1u8];
4194        let event_competing = array![0u8];
4195        let sampleweight = array![1.0];
4196        let x_entry = array![[1.0, 0.0]];
4197        let x_exit = array![[1.0, 0.5]];
4198        let x_derivative = array![[0.0, 1.0]];
4199        let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4200            matrix: array![[1.0]],
4201            lambda: -0.1,
4202            range: 1..2,
4203            nullspace_dim: 0,
4204        }]);
4205
4206        let err = survival_model(
4207            survival_inputs(
4208                &age_entry,
4209                &age_exit,
4210                &event_target,
4211                &event_competing,
4212                &sampleweight,
4213                &x_entry,
4214                &x_exit,
4215                &x_derivative,
4216            ),
4217            penalties,
4218            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4219            SurvivalSpec::Net,
4220        )
4221        .expect_err("negative lambda must be rejected");
4222
4223        assert!(matches!(err, SurvivalError::NonFiniteInput));
4224    }
4225
4226    #[test]
4227    fn penalty_block_range_and_shapemust_match_coefficients() {
4228        let age_entry = array![1.0_f64];
4229        let age_exit = array![2.0_f64];
4230        let event_target = array![1u8];
4231        let event_competing = array![0u8];
4232        let sampleweight = array![1.0];
4233        let x_entry = array![[1.0, 0.0]];
4234        let x_exit = array![[1.0, 0.5]];
4235        let x_derivative = array![[0.0, 1.0]];
4236        let penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4237            matrix: array![[1.0]],
4238            lambda: 0.5,
4239            range: 0..2,
4240            nullspace_dim: 0,
4241        }]);
4242
4243        let err = survival_model(
4244            survival_inputs(
4245                &age_entry,
4246                &age_exit,
4247                &event_target,
4248                &event_competing,
4249                &sampleweight,
4250                &x_entry,
4251                &x_exit,
4252                &x_derivative,
4253            ),
4254            penalties,
4255            SurvivalMonotonicityPenalty { tolerance: 1e-8 },
4256            SurvivalSpec::Net,
4257        )
4258        .expect_err("penalty block geometry must match coefficient support");
4259
4260        assert!(matches!(err, SurvivalError::DimensionMismatch));
4261    }
4262
4263    #[test]
4264    fn survivalgradient_matches_ridge_free_objective_fd() {
4265        let age_entry = array![1.0_f64, 2.0_f64, 3.0_f64];
4266        let age_exit = array![2.0_f64, 3.5_f64, 4.0_f64];
4267        let event_target = array![1u8, 0u8, 1u8];
4268        let event_competing = array![0u8, 0u8, 0u8];
4269        let sampleweight = array![1.0, 1.0, 1.0];
4270        let x_entry = array![
4271            [1.0, age_entry[0].ln()],
4272            [1.0, age_entry[1].ln()],
4273            [1.0, age_entry[2].ln()]
4274        ];
4275        let x_exit = array![
4276            [1.0, age_exit[0].ln()],
4277            [1.0, age_exit[1].ln()],
4278            [1.0, age_exit[2].ln()]
4279        ];
4280        let x_derivative = array![
4281            [0.0, 1.0 / age_exit[0]],
4282            [0.0, 1.0 / age_exit[1]],
4283            [0.0, 1.0 / age_exit[2]]
4284        ];
4285        let penalties = PenaltyBlocks::new(Vec::new());
4286        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4287        let beta = array![-1.0, 3.0];
4288
4289        let model = survival_model(
4290            survival_inputs(
4291                &age_entry,
4292                &age_exit,
4293                &event_target,
4294                &event_competing,
4295                &sampleweight,
4296                &x_entry,
4297                &x_exit,
4298                &x_derivative,
4299            ),
4300            penalties,
4301            mono,
4302            SurvivalSpec::Net,
4303        )
4304        .expect("construct survival model");
4305
4306        let state = model.update_state(&beta).expect("state at beta");
4307        let eps = 1e-7;
4308        for j in 0..beta.len() {
4309            let mut plus = beta.clone();
4310            let mut minus = beta.clone();
4311            plus[j] += eps;
4312            minus[j] -= eps;
4313            let state_plus = model.update_state(&plus).expect("state at beta + eps");
4314            let state_minus = model.update_state(&minus).expect("state at beta - eps");
4315            let obj_plus = 0.5 * state_plus.deviance + state_plus.penalty_term;
4316            let obj_minus = 0.5 * state_minus.deviance + state_minus.penalty_term;
4317            let fd = (obj_plus - obj_minus) / (2.0 * eps);
4318            assert_eq!(
4319                state.gradient[j].signum(),
4320                fd.signum(),
4321                "objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
4322                state.gradient[j]
4323            );
4324            assert!(
4325                (state.gradient[j] - fd).abs() < 1e-5,
4326                "objective/gradient mismatch at j={j}: grad={} fd={fd}",
4327                state.gradient[j]
4328            );
4329        }
4330    }
4331
4332    fn laml_fd_test_model(lambda: f64) -> WorkingModelSurvival {
4333        // 20-subject survival fixture with mean-centered log-age time
4334        // covariate, balanced events/censorings, and moderate hazard levels.
4335        // The fixture is large enough that the observed-information Hessian
4336        // is well-conditioned at the MLE so PIRLS reaches the 1e-10 KKT
4337        // tolerance in well under 80 iterations from the starting beta used
4338        // below.
4339        let age_entry: Array1<f64> = Array1::from(vec![
4340            30.0, 35.0, 40.0, 45.0, 50.0, 55.0, 60.0, 32.0, 37.0, 42.0, 47.0, 52.0, 57.0, 62.0,
4341            34.0, 39.0, 44.0, 49.0, 54.0, 59.0,
4342        ]);
4343        let age_exit: Array1<f64> = Array1::from(vec![
4344            45.0, 48.0, 55.0, 58.0, 62.0, 66.0, 68.0, 47.0, 52.0, 53.0, 55.0, 60.0, 63.0, 70.0,
4345            48.0, 51.0, 58.0, 62.0, 66.0, 69.0,
4346        ]);
4347        let event_target = Array1::from(vec![
4348            1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
4349        ]);
4350        let event_competing = Array1::<u8>::zeros(age_entry.len());
4351        let sampleweight = Array1::from_elem(age_entry.len(), 1.0_f64);
4352        let n = age_entry.len();
4353        let ln_age_mean: f64 = {
4354            let mut sum = 0.0;
4355            for i in 0..n {
4356                sum += age_entry[i].ln() + age_exit[i].ln();
4357            }
4358            sum / (2.0 * n as f64)
4359        };
4360        let mut x_entry = Array2::<f64>::zeros((n, 2));
4361        let mut x_exit = Array2::<f64>::zeros((n, 2));
4362        let mut x_derivative = Array2::<f64>::zeros((n, 2));
4363        for i in 0..n {
4364            x_entry[[i, 0]] = 1.0;
4365            x_exit[[i, 0]] = 1.0;
4366            x_entry[[i, 1]] = age_entry[i].ln() - ln_age_mean;
4367            x_exit[[i, 1]] = age_exit[i].ln() - ln_age_mean;
4368            x_derivative[[i, 0]] = 0.0;
4369            x_derivative[[i, 1]] = 1.0 / age_exit[i];
4370        }
4371        let penalties = PenaltyBlocks::new(vec![
4372            PenaltyBlock {
4373                matrix: array![[3.0]],
4374                lambda: 0.0,
4375                range: 0..1,
4376                nullspace_dim: 0,
4377            },
4378            PenaltyBlock {
4379                matrix: array![[2.5]],
4380                lambda,
4381                range: 1..2,
4382                nullspace_dim: 0,
4383            },
4384        ]);
4385        survival_model(
4386            survival_inputs(
4387                &age_entry,
4388                &age_exit,
4389                &event_target,
4390                &event_competing,
4391                &sampleweight,
4392                &x_entry,
4393                &x_exit,
4394                &x_derivative,
4395            ),
4396            penalties,
4397            SurvivalMonotonicityPenalty { tolerance: 1e-8 },
4398            SurvivalSpec::Net,
4399        )
4400        .expect("construct LAML FD survival model")
4401    }
4402
4403    fn laml_test_logdet_h(state: &WorkingState) -> f64 {
4404        use gam_linalg::faer_ndarray::FaerEigh;
4405        use gam_solve::estimate::reml::reml_outer_engine::{spectral_epsilon, spectral_regularize};
4406
4407        let h_dense = state.hessian.to_dense();
4408        let (evals, _) = h_dense.eigh(faer::Side::Lower).expect("eigh");
4409        let eps = spectral_epsilon(evals.as_slice().unwrap());
4410        evals
4411            .iter()
4412            .map(|&sigma| spectral_regularize(sigma, eps).ln())
4413            .sum()
4414    }
4415
4416    #[test]
4417    fn survival_solver_damping_converges_undamped_objective() {
4418        let rho = -0.35_f64;
4419        let model = laml_fd_test_model(rho.exp());
4420        let beta0 = array![-2.5_f64, 1.0];
4421        let (converged_model, beta) = model
4422            .reconverge_survival_inner_mode(&[rho], &beta0)
4423            .expect("converge survival mode with solver-only damping");
4424        let state = converged_model
4425            .update_state(&beta)
4426            .expect("evaluate undamped objective at converged mode");
4427
4428        assert_eq!(
4429            state.ridge_used, 0.0,
4430            "solver damping must not enter the converged statistical objective"
4431        );
4432        let undamped_stationarity = array1_l2_norm(&state.gradient);
4433        assert!(
4434            undamped_stationarity <= 1.0e-9,
4435            "solver must converge the undamped objective; ||gradient||={undamped_stationarity:.3e}"
4436        );
4437    }
4438
4439    #[test]
4440    fn laml_gradient_and_objective_ignore_inactive_penalty_prefix_blocks() {
4441        // The core claim under test: the survival LAML rho-gradient and the
4442        // documented LAML objective enumerate only penalty blocks whose
4443        // lambda is actually active (> 0). An inactive prefix block must
4444        // therefore contribute neither a log|lambda * S| term to the
4445        // objective nor an entry to the rho-gradient vector.
4446        //
4447        // We verify the objective formula and the gradient dimensionality at
4448        // a fixed beta rather than a fitted one: the bug this test guards
4449        // against was purely algebraic enumeration over penalty blocks and
4450        // has no dependence on PIRLS convergence quality. A gradient-vs-FD
4451        // comparison would require beta to sit at the joint MLE of a tiny
4452        // synthetic survival fixture, which the analytic Newton/PIRLS path
4453        // cannot reach to 1e-10 KKT tolerance without a much richer design.
4454        let rho0 = -0.35_f64;
4455        let beta = array![-2.5_f64, 1.0];
4456        let model = laml_fd_test_model(rho0.exp());
4457        let state = model
4458            .update_state(&beta)
4459            .expect("state for LAML prefix-skip test");
4460
4461        // Sanity: the fixture has two penalty blocks; the first has
4462        // lambda = 0 (inactive prefix) and the second has lambda > 0
4463        // (active). If a future refactor flips this ordering, the prefix
4464        // skip being exercised here would silently become an identity test.
4465        assert_eq!(model.penalties.blocks.len(), 2);
4466        assert_eq!(model.penalties.blocks[0].lambda, 0.0);
4467        assert!(model.penalties.blocks[1].lambda > 0.0);
4468
4469        let rho = Array1::from_iter(
4470            model
4471                .penalties
4472                .blocks
4473                .iter()
4474                .filter(|b| b.lambda > 0.0)
4475                .map(|b| b.lambda.ln()),
4476        );
4477        assert_eq!(
4478            rho.len(),
4479            1,
4480            "fixture should expose exactly one active penalty block for the rho vector"
4481        );
4482
4483        let (obj, grad) = model
4484            .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
4485            .expect("survival LAML objective and gradient");
4486
4487        let expected = 0.5 * state.deviance + state.penalty_term + 0.5 * laml_test_logdet_h(&state)
4488            - 0.5 * (rho0 + 2.5_f64.ln());
4489        assert_eq!(
4490            grad.len(),
4491            1,
4492            "rho-gradient must match the active-penalty count, not the full block list"
4493        );
4494        assert!(
4495            (obj - expected).abs() < 1e-10,
4496            "survival LAML objective mismatch with inactive prefix block: obj={obj} expected={expected}",
4497        );
4498        assert!(
4499            grad[0].is_finite(),
4500            "rho-gradient must be finite: {}",
4501            grad[0]
4502        );
4503    }
4504
4505    #[test]
4506    fn structural_monotonicgradient_matchesobjectivefd() {
4507        let age_entry = array![1.0_f64, 1.3_f64, 1.8_f64];
4508        let age_exit = array![1.6_f64, 2.1_f64, 2.7_f64];
4509        let event_target = array![1u8, 0u8, 1u8];
4510        let event_competing = array![0u8, 0u8, 0u8];
4511        let sampleweight = array![1.0, 1.0, 1.0];
4512
4513        // Time block has 3 structural-monotone columns.
4514        // Final column is a covariate, left unconstrained.
4515        let x_entry = array![
4516            [1.0, 0.2, 0.05, -0.7],
4517            [1.0, 0.5, 0.20, 0.1],
4518            [1.0, 0.9, 0.60, 1.2]
4519        ];
4520        let x_exit = array![
4521            [1.0, 0.4, 0.16, -0.7],
4522            [1.0, 0.8, 0.64, 0.1],
4523            [1.0, 1.1, 1.21, 1.2]
4524        ];
4525        let x_derivative = array![
4526            [0.0, 0.8, 0.64, 0.0],
4527            [0.0, 0.7, 1.12, 0.0],
4528            [0.0, 0.6, 1.32, 0.0]
4529        ];
4530        let penalties = PenaltyBlocks::new(Vec::new());
4531        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4532        let mut model = survival_model(
4533            survival_inputs(
4534                &age_entry,
4535                &age_exit,
4536                &event_target,
4537                &event_competing,
4538                &sampleweight,
4539                &x_entry,
4540                &x_exit,
4541                &x_derivative,
4542            ),
4543            penalties,
4544            mono,
4545            SurvivalSpec::Net,
4546        )
4547        .expect("construct structural survival model");
4548        model
4549            .set_structural_monotonicity(true, 3)
4550            .expect("enable structural monotonicity");
4551        let constraints = model
4552            .monotonicity_linear_constraints()
4553            .expect("structural derivative constraints");
4554        assert_eq!(constraints.a.nrows(), 2);
4555        assert_eq!(constraints.a.ncols(), 4);
4556        assert_eq!(constraints.a.row(0).to_vec(), vec![0.0, 1.0, 0.0, 0.0]);
4557        assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 0.0, 1.0, 0.0]);
4558        assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
4559
4560        let beta = array![0.2, 0.2, 0.1, 0.2];
4561        let state = model.update_state(&beta).expect("state at structural beta");
4562        let eps = 1e-7;
4563        for j in 0..beta.len() {
4564            let mut plus = beta.clone();
4565            let mut minus = beta.clone();
4566            plus[j] += eps;
4567            minus[j] -= eps;
4568            let state_plus = model.update_state(&plus).expect("state at beta + eps");
4569            let state_minus = model.update_state(&minus).expect("state at beta - eps");
4570            let obj_plus = 0.5 * state_plus.deviance + state_plus.penalty_term;
4571            let obj_minus = 0.5 * state_minus.deviance + state_minus.penalty_term;
4572            let fd = (obj_plus - obj_minus) / (2.0 * eps);
4573            assert_eq!(
4574                state.gradient[j].signum(),
4575                fd.signum(),
4576                "structural objective/gradient sign mismatch at j={j}: grad={} fd={fd}",
4577                state.gradient[j]
4578            );
4579            assert!(
4580                (state.gradient[j] - fd).abs() < 2e-5,
4581                "structural objective/gradient mismatch at j={j}: grad={} fd={fd}",
4582                state.gradient[j]
4583            );
4584        }
4585    }
4586
4587    #[test]
4588    fn structural_monotonic_lamlgradient_returns_finitevalues() {
4589        let age_entry = array![1.0_f64, 1.2_f64];
4590        let age_exit = array![1.5_f64, 2.0_f64];
4591        let event_target = array![1u8, 0u8];
4592        let event_competing = array![0u8, 0u8];
4593        let sampleweight = array![1.0, 1.0];
4594        let x_entry = array![[1.0, 0.2, -0.5], [1.0, 0.4, 0.2]];
4595        let x_exit = array![[1.0, 0.5, -0.5], [1.0, 0.8, 0.2]];
4596        let x_derivative = array![[0.0, 0.9, 0.0], [0.0, 0.7, 0.0]];
4597        let penalties = PenaltyBlocks::new(Vec::new());
4598        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4599        let mut model = survival_model(
4600            survival_inputs(
4601                &age_entry,
4602                &age_exit,
4603                &event_target,
4604                &event_competing,
4605                &sampleweight,
4606                &x_entry,
4607                &x_exit,
4608                &x_derivative,
4609            ),
4610            penalties,
4611            mono,
4612            SurvivalSpec::Net,
4613        )
4614        .expect("construct structural survival model");
4615        model
4616            .set_structural_monotonicity(true, 2)
4617            .expect("enable structural monotonicity");
4618        // One simple penalty block to exercise rho-gradient path.
4619        model.penalties = PenaltyBlocks::new(vec![PenaltyBlock {
4620            matrix: array![[1.0]],
4621            lambda: 0.7,
4622            range: 1..2,
4623            nullspace_dim: 0,
4624        }]);
4625        let beta = array![0.2, 0.2, 0.1];
4626        let state = model.update_state(&beta).expect("state at structural beta");
4627        let rho = Array1::from_iter(
4628            model
4629                .penalties
4630                .blocks
4631                .iter()
4632                .filter(|b| b.lambda > 0.0)
4633                .map(|b| b.lambda.ln()),
4634        );
4635        let (obj, grad) = model
4636            .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
4637            .expect("laml gradient should work in structural mode");
4638        assert!(obj.is_finite());
4639        assert_eq!(grad.len(), 1);
4640        assert!(grad[0].is_finite());
4641    }
4642
4643    #[test]
4644    fn structural_monotonicity_switches_to_tiny_derivative_guard_constraints() {
4645        let age_entry = array![1.0_f64];
4646        let age_exit = array![2.0_f64];
4647        let event_target = array![1u8];
4648        let event_competing = array![0u8];
4649        let sampleweight = array![1.0];
4650        let x_entry = array![[0.0]];
4651        let x_exit = array![[0.2]];
4652        let x_derivative = array![[1.0]];
4653
4654        let penalties = PenaltyBlocks::new(Vec::new());
4655        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4656        let mut model = survival_model(
4657            survival_inputs(
4658                &age_entry,
4659                &age_exit,
4660                &event_target,
4661                &event_competing,
4662                &sampleweight,
4663                &x_entry,
4664                &x_exit,
4665                &x_derivative,
4666            ),
4667            penalties,
4668            mono,
4669            SurvivalSpec::Net,
4670        )
4671        .expect("construct structural survival model");
4672
4673        let beta = array![-3.0];
4674        assert!(
4675            model.update_state(&beta).is_err(),
4676            "negative derivative coefficient should violate derivative guard"
4677        );
4678
4679        model
4680            .set_structural_monotonicity(true, 1)
4681            .expect("enable structural monotonicity");
4682        let constraints = model
4683            .monotonicity_linear_constraints()
4684            .expect("structural derivative constraints");
4685        assert_eq!(constraints.a.nrows(), 1);
4686        assert_eq!(constraints.a.ncols(), 1);
4687        assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
4688        // Structural monotonicity uses derivative_guard() == 0.0
4689        assert!(constraints.b[0].abs() <= 1e-12);
4690        let state = model
4691            .update_state(&array![1e-6])
4692            .expect("small positive derivative coefficient should remain feasible");
4693        assert!(state.deviance.is_finite());
4694    }
4695
4696    #[test]
4697    fn derivative_offset_must_clear_nonstructural_monotonicity_threshold() {
4698        let age_entry = array![1.0_f64];
4699        let age_exit = array![2.0_f64];
4700        let event_target = array![1u8];
4701        let event_competing = array![0u8];
4702        let sampleweight = array![1.0];
4703        let x_entry = array![[1.0, 0.0]];
4704        let x_exit = array![[1.0, 0.0]];
4705        let x_derivative = array![[0.0, 0.0]];
4706        let penalties = PenaltyBlocks::new(Vec::new());
4707        let monotonicity = SurvivalMonotonicityPenalty { tolerance: 3.0 };
4708        let eta_entry_offset = array![0.0];
4709        let eta_exit_offset = array![0.0];
4710        let derivative_offset_below_guard = array![2.0];
4711        let derivative_offset_above_guard = array![3.1];
4712        let offsets_below_guard = SurvivalBaselineOffsets {
4713            eta_entry: eta_entry_offset.view(),
4714            eta_exit: eta_exit_offset.view(),
4715            derivative_exit: derivative_offset_below_guard.view(),
4716        };
4717        let offsets_above_guard = SurvivalBaselineOffsets {
4718            eta_entry: eta_entry_offset.view(),
4719            eta_exit: eta_exit_offset.view(),
4720            derivative_exit: derivative_offset_above_guard.view(),
4721        };
4722
4723        let model_below_guard = survival_model_with_offsets(
4724            survival_inputs(
4725                &age_entry,
4726                &age_exit,
4727                &event_target,
4728                &event_competing,
4729                &sampleweight,
4730                &x_entry,
4731                &x_exit,
4732                &x_derivative,
4733            ),
4734            Some(offsets_below_guard),
4735            penalties.clone(),
4736            monotonicity,
4737            SurvivalSpec::Net,
4738        )
4739        .expect("construct model with derivative offset below guard");
4740        let err = model_below_guard
4741            .update_state(&array![0.0, 0.0])
4742            .expect_err("derivative offset below guard should be rejected");
4743        let err_text = err.to_string();
4744        assert!(
4745            err_text.contains("d_eta/dt=2.000e0") && err_text.contains("tolerance=3.000e0"),
4746            "expected derivative guard rejection to report the offset-driven derivative: {err_text}"
4747        );
4748
4749        let model_above_guard = survival_model_with_offsets(
4750            survival_inputs(
4751                &age_entry,
4752                &age_exit,
4753                &event_target,
4754                &event_competing,
4755                &sampleweight,
4756                &x_entry,
4757                &x_exit,
4758                &x_derivative,
4759            ),
4760            Some(offsets_above_guard),
4761            penalties,
4762            SurvivalMonotonicityPenalty { tolerance: 3.0 },
4763            SurvivalSpec::Net,
4764        )
4765        .expect("construct model with derivative offset above guard");
4766        let state = model_above_guard
4767            .update_state(&array![0.0, 0.0])
4768            .expect("derivative offset above guard should remain feasible");
4769        assert!(state.deviance.is_finite());
4770    }
4771
4772    #[test]
4773    fn structural_monotonicity_rejects_negative_derivative_offsets() {
4774        let age_entry = array![1.0_f64];
4775        let age_exit = array![2.0_f64];
4776        let event_target = array![1u8];
4777        let event_competing = array![0u8];
4778        let sampleweight = array![1.0];
4779        let x_entry = array![[0.0]];
4780        let x_exit = array![[0.2]];
4781        let x_derivative = array![[1.0]];
4782        let eta_entry = array![0.0];
4783        let eta_exit = array![0.0];
4784        let derivative_exit = array![-1e-3];
4785        let offsets = SurvivalBaselineOffsets {
4786            eta_entry: eta_entry.view(),
4787            eta_exit: eta_exit.view(),
4788            derivative_exit: derivative_exit.view(),
4789        };
4790
4791        let mut model = survival_model_with_offsets(
4792            survival_inputs(
4793                &age_entry,
4794                &age_exit,
4795                &event_target,
4796                &event_competing,
4797                &sampleweight,
4798                &x_entry,
4799                &x_exit,
4800                &x_derivative,
4801            ),
4802            Some(offsets),
4803            PenaltyBlocks::new(Vec::new()),
4804            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4805            SurvivalSpec::Net,
4806        )
4807        .expect("construct structural survival model");
4808        let err = model
4809            .set_structural_monotonicity(true, 1)
4810            .expect_err("negative derivative offsets must be rejected");
4811        assert!(
4812            err.to_string()
4813                .contains("structural monotonicity requires nonnegative derivative offsets"),
4814            "unexpected error: {err}"
4815        );
4816    }
4817
4818    #[test]
4819    fn structural_monotonicity_emits_coefficient_constraints() {
4820        let age_entry = array![1.0_f64, 1.5_f64];
4821        let age_exit = array![2.0_f64, 3.0_f64];
4822        let event_target = array![1u8, 0u8];
4823        let event_competing = array![0u8, 0u8];
4824        let sampleweight = array![1.0, 1.0];
4825        let x_entry = array![[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]];
4826        let x_exit = array![[0.2, 0.4, 1.0], [0.3, 0.5, 1.0]];
4827        let x_derivative = array![[0.3, 0.2, 0.0], [0.4, 0.1, 0.0]];
4828
4829        let mut model = survival_model(
4830            survival_inputs(
4831                &age_entry,
4832                &age_exit,
4833                &event_target,
4834                &event_competing,
4835                &sampleweight,
4836                &x_entry,
4837                &x_exit,
4838                &x_derivative,
4839            ),
4840            PenaltyBlocks::new(Vec::new()),
4841            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4842            SurvivalSpec::Net,
4843        )
4844        .expect("construct structural survival model");
4845        model
4846            .set_structural_monotonicity(true, 2)
4847            .expect("enable structural monotonicity");
4848
4849        let constraints = model
4850            .monotonicity_linear_constraints()
4851            .expect("structural derivative constraints");
4852
4853        assert_eq!(constraints.a.nrows(), 2);
4854        assert_eq!(constraints.a.ncols(), 3);
4855        assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0, 0.0]);
4856        assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0, 0.0]);
4857        assert!(constraints.b.iter().all(|&v| v.abs() <= 1e-12));
4858    }
4859
4860    #[test]
4861    fn structural_monotonicity_preserves_inactive_time_columns_in_constraints() {
4862        let age_entry = array![1.0_f64];
4863        let age_exit = array![2.0_f64];
4864        let event_target = array![1u8];
4865        let event_competing = array![0u8];
4866        let sampleweight = array![1.0];
4867        let x_entry = array![[1.0, 0.2]];
4868        let x_exit = array![[1.0, 0.6]];
4869        let x_derivative = array![[0.0, 1.0]];
4870
4871        let mut model = survival_model(
4872            survival_inputs(
4873                &age_entry,
4874                &age_exit,
4875                &event_target,
4876                &event_competing,
4877                &sampleweight,
4878                &x_entry,
4879                &x_exit,
4880                &x_derivative,
4881            ),
4882            PenaltyBlocks::new(Vec::new()),
4883            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4884            SurvivalSpec::Net,
4885        )
4886        .expect("construct structural survival model");
4887        model
4888            .set_structural_monotonicity(true, 2)
4889            .expect("enable structural monotonicity");
4890
4891        let constraints = model
4892            .monotonicity_linear_constraints()
4893            .expect("structural derivative constraints");
4894
4895        assert_eq!(constraints.a.nrows(), 1);
4896        assert!(
4897            constraints.a[[0, 0]].abs() <= 1e-12,
4898            "inactive time column should remain unconstrained"
4899        );
4900        assert!(
4901            (constraints.a[[0, 1]] - 1.0).abs() <= 1e-12,
4902            "active time column should remain constrained"
4903        );
4904    }
4905
4906    #[test]
4907    fn structural_monotonicity_preserves_sparse_row_patterns() {
4908        let age_entry = array![1.0_f64, 1.5_f64];
4909        let age_exit = array![2.0_f64, 2.5_f64];
4910        let event_target = array![1u8, 1u8];
4911        let event_competing = array![0u8, 0u8];
4912        let sampleweight = array![1.0, 1.0];
4913        let x_entry = array![[0.0, 0.0], [0.0, 0.0]];
4914        let x_exit = array![[0.4, 0.2], [0.6, 0.3]];
4915        let x_derivative = array![[1.0, 0.0], [1.0, 0.5]];
4916
4917        let mut model = survival_model(
4918            survival_inputs(
4919                &age_entry,
4920                &age_exit,
4921                &event_target,
4922                &event_competing,
4923                &sampleweight,
4924                &x_entry,
4925                &x_exit,
4926                &x_derivative,
4927            ),
4928            PenaltyBlocks::new(Vec::new()),
4929            SurvivalMonotonicityPenalty { tolerance: 0.0 },
4930            SurvivalSpec::Net,
4931        )
4932        .expect("construct structural survival model");
4933        model
4934            .set_structural_monotonicity(true, 2)
4935            .expect("enable structural monotonicity");
4936
4937        let constraints = model
4938            .monotonicity_linear_constraints()
4939            .expect("structural derivative constraints");
4940
4941        assert_eq!(constraints.a.nrows(), 2);
4942        assert_eq!(constraints.a.row(0).to_vec(), vec![1.0, 0.0]);
4943        assert_eq!(constraints.a.row(1).to_vec(), vec![0.0, 1.0]);
4944    }
4945
4946    #[test]
4947    fn update_state_rejects_negative_exit_derivative_for_censoredrows() {
4948        let age_entry = array![1.0_f64];
4949        let age_exit = array![1.1_f64];
4950        let event_target = array![0u8];
4951        let event_competing = array![0u8];
4952        let sampleweight = array![1.0];
4953        let x_entry = array![[0.0]];
4954        let x_exit = array![[0.0]];
4955        let x_derivative = array![[-1.0]];
4956        let penalties = PenaltyBlocks::new(Vec::new());
4957        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
4958        let model = survival_model(
4959            survival_inputs(
4960                &age_entry,
4961                &age_exit,
4962                &event_target,
4963                &event_competing,
4964                &sampleweight,
4965                &x_entry,
4966                &x_exit,
4967                &x_derivative,
4968            ),
4969            penalties,
4970            mono,
4971            SurvivalSpec::Net,
4972        )
4973        .expect("construct censored survival model");
4974
4975        let err = model
4976            .update_state(&array![1.0])
4977            .expect_err("censored row should still enforce monotonic derivative");
4978        assert!(
4979            matches!(err, EstimationError::ParameterConstraintViolation(_)),
4980            "unexpected error: {err:?}"
4981        );
4982    }
4983
4984    fn crude_risk_quadrature_error(
4985        cumulative_entry: f64,
4986        cumulative_exit: f64,
4987        hazard_exit: f64,
4988    ) -> SurvivalError {
4989        calculate_crude_risk_quadrature(
4990            1.0,
4991            2.0,
4992            &[],
4993            0.4,
4994            0.2,
4995            array![1.0].view(),
4996            array![1.0].view(),
4997            |_, design_d, deriv_d, design_m| {
4998                design_d[0] = 1.0;
4999                deriv_d[0] = 0.0;
5000                design_m[0] = 1.0;
5001                Ok((cumulative_entry, cumulative_exit, hazard_exit))
5002            },
5003        )
5004        .expect_err("invalid hazards should fail")
5005    }
5006
5007    #[test]
5008    fn crude_risk_quadrature_rejects_decreasing_cumulative_hazard() {
5009        let err = crude_risk_quadrature_error(0.1, 0.3, 0.25);
5010        assert!(matches!(err, SurvivalError::NonMonotoneCumulativeHazard));
5011    }
5012
5013    #[test]
5014    fn crude_risk_quadrature_rejects_nonpositive_instantaneous_hazard() {
5015        let err = crude_risk_quadrature_error(0.0, 0.4, 0.25);
5016        assert!(matches!(err, SurvivalError::NonPositiveHazard));
5017    }
5018
5019    #[test]
5020    fn laml_no_penalties_matches_documentedobjective() {
5021        let age_entry = array![40.0, 45.0, 50.0, 55.0];
5022        let age_exit = array![44.0, 49.0, 54.0, 59.0];
5023        let event_target = array![1u8, 0u8, 1u8, 0u8];
5024        let event_competing = Array1::<u8>::zeros(4);
5025        let sampleweight = Array1::ones(4);
5026        let x_entry = array![
5027            [1.0, -0.2, 0.04],
5028            [1.0, -0.1, 0.01],
5029            [1.0, 0.0, 0.0],
5030            [1.0, 0.1, 0.01]
5031        ];
5032        let x_exit = array![
5033            [1.0, -0.12, 0.0144],
5034            [1.0, -0.02, 0.0004],
5035            [1.0, 0.08, 0.0064],
5036            [1.0, 0.18, 0.0324]
5037        ];
5038        let x_derivative = array![
5039            [0.0, 0.02, 0.001],
5040            [0.0, 0.02, 0.001],
5041            [0.0, 0.02, 0.001],
5042            [0.0, 0.02, 0.001]
5043        ];
5044        let penalties = PenaltyBlocks::new(Vec::new());
5045        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5046        let beta = array![-2.0, 0.7, 0.2];
5047
5048        let model = survival_model(
5049            survival_inputs(
5050                &age_entry,
5051                &age_exit,
5052                &event_target,
5053                &event_competing,
5054                &sampleweight,
5055                &x_entry,
5056                &x_exit,
5057                &x_derivative,
5058            ),
5059            penalties,
5060            mono,
5061            SurvivalSpec::Net,
5062        )
5063        .expect("construct survival model");
5064
5065        let state = model.update_state(&beta).expect("state at beta");
5066        let rho = Array1::from_iter(
5067            model
5068                .penalties
5069                .blocks
5070                .iter()
5071                .filter(|b| b.lambda > 0.0)
5072                .map(|b| b.lambda.ln()),
5073        );
5074        let (obj, grad) = model
5075            .unified_lamlobjective_and_rhogradient(&beta, &state, &rho)
5076            .expect("laml objective for no-penalty model");
5077
5078        let h_dense = state.hessian.to_dense();
5079        // Mirror the production LAML Hessian logdet EXACTLY (call production, not
5080        // replay): `unified_lamlobjective_and_rhogradient` assembles the logdet
5081        // through a `DenseSpectralOperator` whose pseudo-logdet mode is selected by
5082        // delayed entry. Left-truncated (delayed-entry) transformation survival
5083        // uses the identified positive-subspace HardPseudo logdet (#1915: with the
5084        // +H(entry) term the observed information carries genuine negative
5085        // curvature, which the smooth full-spectrum regularizer must NOT reward by
5086        // mapping to a tiny positive value); right-censored keeps Smooth. This
5087        // test's data is left-truncated (age_entry ≫ origin), so the documented
5088        // objective's logdet is the HardPseudo one — deriving `expected` from the
5089        // same operator+mode keeps this a genuine
5090        // obj = ½·deviance + penalty + ½·logdet(H) decomposition check instead of a
5091        // stale hand-rolled Smooth-mode formula (which pre-dated #1915 and summed
5092        // ln(spectral_regularize(σ)) over the whole spectrum, including the
5093        // non-positive delayed-entry modes the objective now excludes).
5094        let logdet_h: f64 = {
5095            use gam_problem::PseudoLogdetMode;
5096            use gam_solve::estimate::reml::reml_outer_engine::{
5097                DenseSpectralOperator, HessianFactorization,
5098            };
5099            let has_left_truncation = age_entry.iter().any(|&t| t > ENTRY_AT_ORIGIN_THRESHOLD);
5100            let mode = if has_left_truncation {
5101                PseudoLogdetMode::HardPseudo
5102            } else {
5103                PseudoLogdetMode::Smooth
5104            };
5105            DenseSpectralOperator::from_symmetric_with_mode(&h_dense, mode)
5106                .expect("survival LAML Hessian operator")
5107                .logdet()
5108        };
5109        let expected = 0.5 * state.deviance + state.penalty_term + 0.5 * logdet_h;
5110
5111        assert_eq!(grad.len(), 0);
5112        assert!(
5113            (obj - expected).abs() < 1e-10,
5114            "no-penalty LAML objective mismatch: obj={} expected={}",
5115            obj,
5116            expected
5117        );
5118    }
5119
5120    #[test]
5121    fn monotonicity_constraints_collapse_positive_collinearrows() {
5122        let a = array![[0.0, 0.5, 0.0], [0.0, 0.25, 0.0], [0.0, 0.125, 0.0]];
5123        let b = array![1e-8, 1e-8, 1e-8];
5124
5125        let compressed = compress_positive_collinear_constraints(&a, &b);
5126
5127        assert_eq!(compressed.a.nrows(), 1);
5128        assert_eq!(compressed.a.ncols(), 3);
5129        assert!(compressed.a[[0, 0]].abs() <= 1e-12);
5130        assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
5131        assert!(compressed.a[[0, 2]].abs() <= 1e-12);
5132        assert!((compressed.b[0] - 8e-8).abs() <= 1e-18);
5133    }
5134
5135    #[test]
5136    fn monotonicity_constraints_preserve_distinct_directions() {
5137        let a = array![[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]];
5138        let b = array![0.2, 0.3, 0.1];
5139
5140        let compressed = compress_positive_collinear_constraints(&a, &b);
5141
5142        assert_eq!(compressed.a.nrows(), 2);
5143        let mut saw_x = false;
5144        let mut saw_y = false;
5145        for i in 0..compressed.a.nrows() {
5146            if (compressed.a[[i, 0]] - 1.0).abs() <= 1e-12 && compressed.a[[i, 1]].abs() <= 1e-12 {
5147                saw_x = true;
5148                assert!((compressed.b[i] - 0.2).abs() <= 1e-12);
5149            }
5150            if compressed.a[[i, 0]].abs() <= 1e-12 && (compressed.a[[i, 1]] - 1.0).abs() <= 1e-12 {
5151                saw_y = true;
5152                assert!((compressed.b[i] - 0.3).abs() <= 1e-12);
5153            }
5154        }
5155        assert!(saw_x);
5156        assert!(saw_y);
5157    }
5158
5159    #[test]
5160    fn monotonicity_constraints_cluster_near_collinearrows() {
5161        let a = array![
5162            [0.0, 0.5, 0.0],
5163            [0.0, 0.50000000003, 0.0],
5164            [0.0, 0.49999999997, 0.0]
5165        ];
5166        let b = array![1e-8, 1.00000000005e-8, 0.99999999995e-8];
5167
5168        let compressed = compress_positive_collinear_constraints(&a, &b);
5169
5170        assert_eq!(compressed.a.nrows(), 1);
5171        assert_eq!(compressed.a.ncols(), 3);
5172        assert!(compressed.a[[0, 0]].abs() <= 1e-12);
5173        assert!((compressed.a[[0, 1]] - 1.0).abs() <= 1e-12);
5174        assert!(compressed.a[[0, 2]].abs() <= 1e-12);
5175        assert!((compressed.b[0] - 2.0e-8).abs() <= 1e-18);
5176    }
5177
5178    #[test]
5179    fn monotonicity_constraints_cluster_spline_like_near_duplicates() {
5180        let a = array![
5181            [0.0, 0.401, 0.302, 0.197],
5182            [0.0, 0.40100000003, 0.30199999998, 0.19700000001],
5183            [0.0, 0.40099999997, 0.30200000002, 0.19699999999],
5184            [0.0, 0.125, 0.500, 0.375]
5185        ];
5186        let b = array![2.0e-8, 2.00000000004e-8, 1.99999999996e-8, 3.0e-8];
5187
5188        let compressed = compress_positive_collinear_constraints(&a, &b);
5189
5190        assert_eq!(compressed.a.nrows(), 2);
5191        let mut clustered_face = false;
5192        let mut distinct_face = false;
5193        for i in 0..compressed.a.nrows() {
5194            let row = compressed.a.row(i);
5195            if row[1] > 0.99 && row[2] > 0.7 && row[3] > 0.49 {
5196                clustered_face = true;
5197                assert!((compressed.b[i] - (2.0e-8 / 0.401)).abs() <= 1e-12);
5198            } else {
5199                distinct_face = true;
5200                assert!((row[1] - 0.25).abs() <= 1e-12);
5201                assert!((row[2] - 1.0).abs() <= 1e-12);
5202                assert!((row[3] - 0.75).abs() <= 1e-12);
5203                assert!((compressed.b[i] - 6.0e-8).abs() <= 1e-18);
5204            }
5205        }
5206        assert!(clustered_face);
5207        assert!(distinct_face);
5208    }
5209
5210    #[test]
5211    fn linear_time_monotonicity_constraints_reduce_to_single_halfspace() {
5212        let age_entry = array![1.0_f64, 1.0, 1.0];
5213        let age_exit = array![2.0_f64, 4.0, 8.0];
5214        let event_target = array![0u8, 1u8, 0u8];
5215        let event_competing = array![0u8, 0u8, 0u8];
5216        let sampleweight = array![1.0, 1.0, 1.0];
5217        let x_entry = array![
5218            [1.0, age_entry[0].ln()],
5219            [1.0, age_entry[1].ln()],
5220            [1.0, age_entry[2].ln()]
5221        ];
5222        let x_exit = array![
5223            [1.0, age_exit[0].ln()],
5224            [1.0, age_exit[1].ln()],
5225            [1.0, age_exit[2].ln()]
5226        ];
5227        let x_derivative = array![[0.0, 0.5], [0.0, 0.25], [0.0, 0.125]];
5228        let penalties = PenaltyBlocks::new(Vec::new());
5229        let mono = SurvivalMonotonicityPenalty { tolerance: 1e-8 };
5230
5231        let collocation_offsets = Array1::zeros(x_derivative.nrows());
5232        let mut inputs = survival_inputs(
5233            &age_entry,
5234            &age_exit,
5235            &event_target,
5236            &event_competing,
5237            &sampleweight,
5238            &x_entry,
5239            &x_exit,
5240            &x_derivative,
5241        );
5242        inputs.monotonicity_constraint_rows = Some(x_derivative.view());
5243        inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
5244
5245        let model = survival_model(inputs, penalties, mono, SurvivalSpec::Net)
5246            .expect("construct linear survival model");
5247
5248        let constraints = model
5249            .monotonicity_linear_constraints()
5250            .expect("monotonicity constraints");
5251        assert_eq!(constraints.a.nrows(), 1);
5252        assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
5253        assert!((constraints.b[0] - 8e-8).abs() <= 1e-12);
5254    }
5255
5256    #[test]
5257    fn monotonicity_constraints_skip_numericallyzerorows() {
5258        let age_entry = array![1.0_f64, 1.0, 1.0];
5259        let age_exit = array![2.0_f64, 3.0, 4.0];
5260        let event_target = array![0u8, 0u8, 0u8];
5261        let event_competing = array![0u8, 0u8, 0u8];
5262        let sampleweight = array![1.0, 1.0, 1.0];
5263        let x_entry = array![[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]];
5264        let x_exit = x_entry.clone();
5265        let x_derivative = array![[0.0, 0.0], [0.0, 1e-16], [0.0, 0.25]];
5266
5267        let collocation_offsets = Array1::zeros(x_derivative.nrows());
5268        let mut inputs = survival_inputs(
5269            &age_entry,
5270            &age_exit,
5271            &event_target,
5272            &event_competing,
5273            &sampleweight,
5274            &x_entry,
5275            &x_exit,
5276            &x_derivative,
5277        );
5278        inputs.monotonicity_constraint_rows = Some(x_derivative.view());
5279        inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
5280
5281        let model = survival_model(
5282            inputs,
5283            PenaltyBlocks::new(Vec::new()),
5284            SurvivalMonotonicityPenalty { tolerance: 0.0 },
5285            SurvivalSpec::Net,
5286        )
5287        .expect("construct survival model");
5288
5289        let constraints = model
5290            .monotonicity_linear_constraints()
5291            .expect("nonzero derivative row should remain");
5292        assert_eq!(constraints.a.nrows(), 1);
5293        assert!((constraints.a[[0, 1]] - 1.0).abs() <= 1e-12);
5294        assert!(constraints.b[0].abs() <= 1e-18);
5295    }
5296
5297    #[test]
5298    fn censoredrows_allowzero_boundary_derivative() {
5299        let age_entry = array![1.0_f64];
5300        let age_exit = array![2.0_f64];
5301        let event_target = array![0u8];
5302        let event_competing = array![0u8];
5303        let sampleweight = array![1.0];
5304        let x_entry = array![[0.0]];
5305        let x_exit = array![[0.0]];
5306        let x_derivative = array![[1.0]];
5307
5308        let model = survival_model(
5309            survival_inputs(
5310                &age_entry,
5311                &age_exit,
5312                &event_target,
5313                &event_competing,
5314                &sampleweight,
5315                &x_entry,
5316                &x_exit,
5317                &x_derivative,
5318            ),
5319            PenaltyBlocks::new(Vec::new()),
5320            SurvivalMonotonicityPenalty { tolerance: 0.0 },
5321            SurvivalSpec::Net,
5322        )
5323        .expect("construct censored survival model");
5324
5325        let state = model
5326            .update_state(&array![0.0])
5327            .expect("censored boundary derivative should remain feasible with zero tolerance");
5328        assert!(state.deviance.is_finite());
5329    }
5330
5331    #[test]
5332    fn eventrows_keep_positive_derivative_constraint() {
5333        let age_entry = array![1.0_f64, 1.0];
5334        let age_exit = array![2.0_f64, 4.0];
5335        let event_target = array![0u8, 1u8];
5336        let event_competing = array![0u8, 0u8];
5337        let sampleweight = array![1.0, 1.0];
5338        let x_entry = array![[0.0], [0.0]];
5339        let x_exit = array![[0.0], [0.0]];
5340        let x_derivative = array![[0.5], [0.25]];
5341
5342        let collocation_offsets = Array1::zeros(x_derivative.nrows());
5343        let mut inputs = survival_inputs(
5344            &age_entry,
5345            &age_exit,
5346            &event_target,
5347            &event_competing,
5348            &sampleweight,
5349            &x_entry,
5350            &x_exit,
5351            &x_derivative,
5352        );
5353        inputs.monotonicity_constraint_rows = Some(x_derivative.view());
5354        inputs.monotonicity_constraint_offsets = Some(collocation_offsets.view());
5355
5356        let model = survival_model(
5357            inputs,
5358            PenaltyBlocks::new(Vec::new()),
5359            SurvivalMonotonicityPenalty { tolerance: 1e-8 },
5360            SurvivalSpec::Net,
5361        )
5362        .expect("construct mixed survival model");
5363
5364        let constraints = model
5365            .monotonicity_linear_constraints()
5366            .expect("event row should induce positive lower bound");
5367        assert_eq!(constraints.a.nrows(), 1);
5368        assert!((constraints.a[[0, 0]] - 1.0).abs() <= 1e-12);
5369        assert!((constraints.b[0] - 4e-8).abs() <= 1e-18);
5370    }
5371
5372    #[test]
5373    fn structural_monotonicity_clamps_tiny_negative_roundoff() {
5374        let age_entry = array![1.0_f64];
5375        let age_exit = array![2.0_f64];
5376        let event_target = array![1u8];
5377        let event_competing = array![0u8];
5378        let sampleweight = array![1.0];
5379        let x_entry = array![[0.0]];
5380        let x_exit = array![[0.0]];
5381        let x_derivative = array![[1.0]];
5382        let mut model = survival_model(
5383            survival_inputs(
5384                &age_entry,
5385                &age_exit,
5386                &event_target,
5387                &event_competing,
5388                &sampleweight,
5389                &x_entry,
5390                &x_exit,
5391                &x_derivative,
5392            ),
5393            PenaltyBlocks::new(Vec::new()),
5394            SurvivalMonotonicityPenalty { tolerance: 1e-8 },
5395            SurvivalSpec::Net,
5396        )
5397        .expect("construct survival model");
5398        model
5399            .set_structural_monotonicity(true, 1)
5400            .expect("enable structural monotonicity");
5401
5402        let state = model
5403            .update_state(&array![-1e-8])
5404            .expect("tiny structural roundoff should be clamped");
5405        assert!(state.deviance.is_finite());
5406    }
5407
5408    #[test]
5409    fn compressed_monotonicity_constraints_preserve_uncompressed_feasible_region() {
5410        let uncompressed_constraints = LinearInequalityConstraints {
5411            a: array![
5412                [0.0, 0.5, 0.0],
5413                [0.0, 1.0 / 3.0, 0.0],
5414                [0.0, 0.2, 0.0],
5415                [0.0, 0.125, 0.0]
5416            ],
5417            b: Array1::from_elem(4, 1e-8),
5418        };
5419        let compressed_constraints = compress_positive_collinear_constraints(
5420            &uncompressed_constraints.a,
5421            &uncompressed_constraints.b,
5422        );
5423
5424        let candidates = [
5425            array![0.0, 1e-9, 0.0],
5426            array![0.0, 4e-8, 0.0],
5427            array![0.0, 8e-8, 0.0],
5428            array![0.0, 2e-7, 1.5],
5429        ];
5430        for beta in candidates {
5431            let uncompressed_ok = (0..uncompressed_constraints.a.nrows()).all(|i| {
5432                uncompressed_constraints.a.row(i).dot(&beta) >= uncompressed_constraints.b[i]
5433            });
5434            let compressed_ok = (0..compressed_constraints.a.nrows())
5435                .all(|i| compressed_constraints.a.row(i).dot(&beta) >= compressed_constraints.b[i]);
5436            assert_eq!(compressed_ok, uncompressed_ok);
5437        }
5438    }
5439
5440    #[test]
5441    fn exact_survival_derivatives_are_time_unit_invariant_up_to_constant_shift() {
5442        let age_entry = array![10.0_f64, 20.0, 25.0];
5443        let age_exit = array![15.0_f64, 30.0, 40.0];
5444        let event_target = array![1u8, 0u8, 1u8];
5445        let event_competing = array![0u8, 0u8, 0u8];
5446        let sampleweight = array![1.0, 2.0, 0.5];
5447        let x_entry = array![[0.1, 0.2, 1.0], [0.3, 0.4, 1.0], [0.2, 0.6, 1.0]];
5448        let x_exit = array![[0.2, 0.3, 1.0], [0.5, 0.7, 1.0], [0.4, 0.8, 1.0]];
5449        let x_derivative = array![[0.04, 0.02, 0.0], [0.03, 0.01, 0.0], [0.02, 0.03, 0.0]];
5450        let beta = array![0.8, 1.1, -0.2];
5451
5452        let base_model = survival_model(
5453            survival_inputs(
5454                &age_entry,
5455                &age_exit,
5456                &event_target,
5457                &event_competing,
5458                &sampleweight,
5459                &x_entry,
5460                &x_exit,
5461                &x_derivative,
5462            ),
5463            PenaltyBlocks::new(Vec::new()),
5464            SurvivalMonotonicityPenalty { tolerance: 0.0 },
5465            SurvivalSpec::Net,
5466        )
5467        .expect("construct base survival model");
5468        let base_state = base_model
5469            .update_state(&beta)
5470            .expect("evaluate base survival state");
5471
5472        let time_scale = 365.25;
5473        let scaled_age_entry = age_entry.mapv(|v| v * time_scale);
5474        let scaled_age_exit = age_exit.mapv(|v| v * time_scale);
5475        let scaled_x_derivative = x_derivative.mapv(|v| v / time_scale);
5476        let scaled_model = survival_model(
5477            survival_inputs(
5478                &scaled_age_entry,
5479                &scaled_age_exit,
5480                &event_target,
5481                &event_competing,
5482                &sampleweight,
5483                &x_entry,
5484                &x_exit,
5485                &scaled_x_derivative,
5486            ),
5487            PenaltyBlocks::new(Vec::new()),
5488            SurvivalMonotonicityPenalty { tolerance: 0.0 },
5489            SurvivalSpec::Net,
5490        )
5491        .expect("construct scaled survival model");
5492        let scaled_state = scaled_model
5493            .update_state(&beta)
5494            .expect("evaluate scaled survival state");
5495
5496        let weighted_events = sampleweight
5497            .iter()
5498            .zip(event_target.iter())
5499            .map(|(w, d)| *w * f64::from(*d))
5500            .sum::<f64>();
5501        let expected_deviance_shift = 2.0 * weighted_events * time_scale.ln();
5502        assert!(
5503            (scaled_state.deviance - base_state.deviance - expected_deviance_shift).abs() <= 1e-10,
5504            "deviance shift mismatch: scaled={} base={} expected_shift={expected_deviance_shift}",
5505            scaled_state.deviance,
5506            base_state.deviance
5507        );
5508
5509        for j in 0..beta.len() {
5510            assert!(
5511                (scaled_state.gradient[j] - base_state.gradient[j]).abs() <= 1e-12,
5512                "gradient mismatch at j={j}: scaled={} base={}",
5513                scaled_state.gradient[j],
5514                base_state.gradient[j]
5515            );
5516        }
5517
5518        let base_hessian = base_state.hessian.to_dense();
5519        let scaled_hessian = scaled_state.hessian.to_dense();
5520        for r in 0..beta.len() {
5521            for c in 0..beta.len() {
5522                assert!(
5523                    (scaled_hessian[[r, c]] - base_hessian[[r, c]]).abs() <= 1e-12,
5524                    "hessian mismatch at ({r},{c}): scaled={} base={}",
5525                    scaled_hessian[[r, c]],
5526                    base_hessian[[r, c]]
5527                );
5528            }
5529        }
5530    }
5531}