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