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