Skip to main content

gam_models/survival/latent/
survival.rs

1//! Jointly learned latent-frailty survival and binary deployment families with
2//! a live time/baseline block.
3//!
4//! Model:
5//!   H_0(a) = exp(q(a)),
6//!   h_0(a) = dq(a)/da,
7//!   H(a | U) = H_0(a) * exp(U),
8//!   U ~ N(mu, sigma^2),
9//!   mu = X beta + offset.
10//!
11//! Unlike the old compiled-row path, the cumulative masses and baseline hazard
12//! are rebuilt inside the optimizer from the current time-basis coefficients.
13//! The family-level fit surface supports exact events, right censoring, and
14//! interval censoring `T ∈ (L, R]` (contribution `log[S(L) − S(R)]`). Interval
15//! rows carry the reserved [`LATENT_SURVIVAL_EVENT_INTERVAL`] event code and a
16//! dedicated upper-bound time channel (`time_design_right` / `q_right`); the
17//! 3-way event dispatch is [`latent_survival_event_type_for`]. Reached from the
18//! formula DSL via `SurvInterval(L, R, event) ~ ...`.
19
20use crate::custom_family::{
21    BlockWorkingSet, BlockwiseFitOptions, CustomFamily, ExactNewtonJointGradientEvaluation,
22    ExactNewtonJointHessianWorkspace, FamilyEvaluation, ParameterBlockSpec, ParameterBlockState,
23    PenaltyMatrix, fit_custom_family, fit_custom_family_fixed_log_lambdas,
24};
25use crate::gamlss::{FamilyMetadata, ParameterLink};
26use crate::sigma_link::{exp_sigma_eta_for_sigma_scalar, exp_sigma_from_eta_scalar};
27use crate::survival::latent::interval::{
28    LatentFrailtyResolution, LatentIntervalModel, LatentIntervalRowView,
29    validate_latent_interval_inputs,
30};
31use crate::survival::location_scale::{
32    TimeBlockInput, project_onto_linear_constraints, structural_time_coefficient_constraints,
33};
34use crate::survival::lognormal_kernel::{
35    FrailtySpec, HazardLoading, LatentSurvivalEventType, LatentSurvivalRow, LatentSurvivalRowJet,
36    log_kernel_bundle,
37};
38use gam_linalg::matrix::{DenseDesignMatrix, DesignMatrix, SymmetricMatrix};
39use crate::model_types::UnifiedFitResult;
40use gam_solve::pirls::LinearInequalityConstraints;
41use crate::probability::signed_log_sum_exp;
42use crate::quadrature::{IntegratedExpectationMode, QuadratureContext};
43use gam_terms::smooth::{
44    TermCollectionDesign, TermCollectionSpec, build_term_collection_design,
45};
46use crate::fit_orchestration::drivers::freeze_term_collection_from_design;
47use gam_problem::MIN_WEIGHT;
48use ndarray::{Array1, Array2, ArrayView1, ArrayView2, s};
49use std::collections::BTreeMap;
50use std::sync::Arc;
51
52/// Typed error for the latent-survival / latent-binary family kernels and
53/// their fit-time and per-row validation helpers. Variants pick the semantic
54/// bucket while the inner `reason` carries the original byte-equivalent
55/// message so external callers that previously consumed `String` errors keep
56/// the same diagnostic text via `Display`.
57#[derive(Debug, Clone)]
58pub enum LatentSurvivalError {
59    /// The frailty spec supplied to a latent-survival or latent-binary
60    /// helper is incompatible (wrong variant, missing fixed sigma, non-finite
61    /// or negative fixed sigma).
62    InvalidFrailty { reason: String },
63    /// Per-row dataset validation failed: empty input, size mismatch across
64    /// the spec vectors, or invalid age / event / weight / unloaded-mass
65    /// values for an individual row.
66    InvalidDataset { reason: String },
67    /// A parameter-block state, eta vector, or directional-derivative
68    /// argument supplied to a family entry point has the wrong length.
69    BlockMismatch { reason: String },
70    /// A runtime numerical value (sigma, baseline hazard derivative, kernel
71    /// sum, event probability) became non-finite or out-of-domain.
72    NumericalFailure { reason: String },
73    /// The requested combination of time-block structure or event type is
74    /// not implemented (non-structural monotonicity, interval-censored rows
75    /// on the dynamic-derivative path).
76    UnsupportedConfiguration { reason: String },
77}
78
79impl_reason_error_boilerplate! {
80    LatentSurvivalError {
81        InvalidFrailty,
82        InvalidDataset,
83        BlockMismatch,
84        NumericalFailure,
85        UnsupportedConfiguration,
86    }
87}
88
89impl From<crate::block_layout::block_count::BlockCountMismatch> for LatentSurvivalError {
90    fn from(
91        err: crate::block_layout::block_count::BlockCountMismatch,
92    ) -> LatentSurvivalError {
93        LatentSurvivalError::BlockMismatch {
94            reason: err.message(),
95        }
96    }
97}
98
99impl From<String> for LatentSurvivalError {
100    /// Inbound conversion for the many `Result<_, String>` helpers this
101    /// module still calls into (term-collection design assembly, dense
102    /// chunk conversion, sparse linear constraints). The text is preserved
103    /// verbatim; we only pick a category so external messages flow through
104    /// `?` without per-callsite `.map_err`.
105    fn from(reason: String) -> LatentSurvivalError {
106        LatentSurvivalError::InvalidDataset { reason }
107    }
108}
109
110/// Reserved [`LatentSurvivalTermSpec::event_target`] code marking an
111/// interval-censored row `(L, R]`. Exact-event codes are `>= 1` and right
112/// censoring is `0`; the interval code is the sentinel `u8::MAX` so it never
113/// collides with an exact-event count and the dispatch is an explicit 3-way map
114/// `{0 → RightCensored, INTERVAL → IntervalCensored, k ≥ 1 → ExactEvent}`.
115pub const LATENT_SURVIVAL_EVENT_INTERVAL: u8 = u8::MAX;
116
117#[inline]
118fn latent_survival_event_type_for(code: u8) -> LatentSurvivalEventType {
119    match code {
120        0 => LatentSurvivalEventType::RightCensored,
121        LATENT_SURVIVAL_EVENT_INTERVAL => LatentSurvivalEventType::IntervalCensored,
122        _ => LatentSurvivalEventType::ExactEvent,
123    }
124}
125
126#[derive(Clone)]
127pub struct LatentSurvivalTermSpec {
128    pub age_entry: Array1<f64>,
129    pub age_exit: Array1<f64>,
130    pub event_target: Array1<u8>,
131    pub weights: Array1<f64>,
132    pub derivative_guard: f64,
133    pub time_block: TimeBlockInput,
134    /// Time-basis design evaluated at the interval upper bound `R` (so
135    /// `q_right = design_right · β_time + offset_right`). `None` when the data
136    /// carries no interval-censored rows; the family then reuses the exit design
137    /// for the unused `q_right` channel. When `Some`, rows whose
138    /// `event_target == LATENT_SURVIVAL_EVENT_INTERVAL` contribute the interval
139    /// likelihood `log[S(L) − S(R)]`.
140    pub time_design_right: Option<DesignMatrix>,
141    pub time_offset_right: Option<Array1<f64>>,
142    pub unloaded_mass_entry: Array1<f64>,
143    pub unloaded_mass_exit: Array1<f64>,
144    /// Unloaded (background) cumulative mass at the interval upper bound `R`.
145    /// Length-`n`; entries for non-interval rows are ignored. Empty/`None`
146    /// folds to zero (full-loading interval rows).
147    pub unloaded_mass_right: Array1<f64>,
148    pub unloaded_hazard_exit: Array1<f64>,
149    pub meanspec: TermCollectionSpec,
150    pub mean_offset: Array1<f64>,
151}
152
153pub struct LatentSurvivalTermFitResult {
154    pub fit: UnifiedFitResult,
155    pub design: TermCollectionDesign,
156    pub resolvedspec: TermCollectionSpec,
157    pub latent_sd: f64,
158    /// Per-row residuals of the unpenalized NLL w.r.t. the additive baseline
159    /// time-block offsets `(entry, exit, derivative)` at the converged β̂.
160    /// Contracted against `baseline_offset_theta_partials` by
161    /// `baseline_chain_rule_gradient` to give the exact θ-gradient of the
162    /// profile penalized NLL for the outer baseline-config optimizer.
163    pub baseline_offset_residuals: crate::survival::OffsetChannelResiduals,
164}
165
166#[derive(Clone)]
167pub struct LatentBinaryTermSpec {
168    pub age_entry: Array1<f64>,
169    pub age_exit: Array1<f64>,
170    pub event_target: Array1<u8>,
171    pub weights: Array1<f64>,
172    pub derivative_guard: f64,
173    pub time_block: TimeBlockInput,
174    pub unloaded_mass_entry: Array1<f64>,
175    pub unloaded_mass_exit: Array1<f64>,
176    pub meanspec: TermCollectionSpec,
177    pub mean_offset: Array1<f64>,
178}
179
180pub struct LatentBinaryTermFitResult {
181    pub fit: UnifiedFitResult,
182    pub design: TermCollectionDesign,
183    pub resolvedspec: TermCollectionSpec,
184    /// Per-row residuals of the unpenalized NLL w.r.t. the additive baseline
185    /// time-block offsets `(entry, exit)` at the converged β̂ (the derivative
186    /// channel is identically zero for the binary deployment likelihood).
187    pub baseline_offset_residuals: crate::survival::OffsetChannelResiduals,
188}
189
190#[derive(Clone)]
191struct PreparedLatentTimeBlock {
192    design_entry: Array2<f64>,
193    design_exit: Array2<f64>,
194    design_derivative_exit: Array2<f64>,
195    /// Dense time-basis design at the interval upper bound `R`. Falls back to a
196    /// clone of `design_exit` when the spec supplies no interval design, so the
197    /// `q_right` channel is always well-defined (and unused for non-interval
198    /// rows).
199    design_right: Array2<f64>,
200    linear_constraints: Option<LinearInequalityConstraints>,
201    penalties: Vec<Array2<f64>>,
202    initial_beta: Option<Array1<f64>>,
203}
204
205#[derive(Clone)]
206pub struct LatentSurvivalFamily {
207    pub event_target: Array1<u8>,
208    pub weights: Array1<f64>,
209    pub latent_sd_fixed: Option<f64>,
210    pub hazard_loading: HazardLoading,
211    pub unloaded_mass_entry: Array1<f64>,
212    pub unloaded_mass_exit: Array1<f64>,
213    pub unloaded_hazard_exit: Array1<f64>,
214    pub x_time_entry: Array2<f64>,
215    pub x_time_exit: Array2<f64>,
216    pub x_time_derivative_exit: Array2<f64>,
217    /// Time-basis design evaluated at the interval upper bound `R` (so
218    /// `q_right = x_time_right · β_time + time_offset_right`). For non-interval
219    /// rows this row equals `x_time_exit`'s row (`q_right` is then unused by the
220    /// likelihood), so the matrix always has `n` rows and the same column count
221    /// as the other time designs.
222    pub x_time_right: Array2<f64>,
223    /// Time-block offset at the interval upper bound `R` (length `n`).
224    pub time_offset_right: Array1<f64>,
225    /// Unloaded (background) cumulative mass at the interval upper bound `R`
226    /// (length `n`). Ignored for non-interval rows.
227    pub unloaded_mass_right: Array1<f64>,
228    pub x_mean: DesignMatrix,
229    pub time_linear_constraints: Option<LinearInequalityConstraints>,
230    pub quadctx: Arc<QuadratureContext>,
231}
232
233#[derive(Clone)]
234pub struct LatentBinaryFamily {
235    pub event_target: Array1<u8>,
236    pub weights: Array1<f64>,
237    pub latent_sd: f64,
238    pub hazard_loading: HazardLoading,
239    pub unloaded_mass_entry: Array1<f64>,
240    pub unloaded_mass_exit: Array1<f64>,
241    pub x_time_entry: Array2<f64>,
242    pub x_time_exit: Array2<f64>,
243    pub x_mean: DesignMatrix,
244    pub time_linear_constraints: Option<LinearInequalityConstraints>,
245    pub quadctx: Arc<QuadratureContext>,
246}
247
248impl LatentSurvivalFamily {
249    pub const BLOCK_TIME: usize = 0;
250    pub const BLOCK_MEAN: usize = 1;
251    pub const BLOCK_LOG_SIGMA: usize = 2;
252
253    pub fn parameter_names() -> &'static [&'static str] {
254        &["time_transform", "mean"]
255    }
256
257    pub fn parameter_links() -> &'static [ParameterLink] {
258        &[ParameterLink::Identity, ParameterLink::Identity]
259    }
260
261    pub fn metadata() -> FamilyMetadata {
262        FamilyMetadata {
263            name: "latent_survival",
264            parameternames: Self::parameter_names(),
265            parameter_links: Self::parameter_links(),
266        }
267    }
268
269    fn split_time_eta<'a>(
270        &self,
271        block_states: &'a [ParameterBlockState],
272    ) -> Result<
273        (
274            ArrayView1<'a, f64>,
275            ArrayView1<'a, f64>,
276            ArrayView1<'a, f64>,
277            &'a Array1<f64>,
278        ),
279        LatentSurvivalError,
280    > {
281        let expected_blocks = if self.latent_sd_fixed.is_some() { 2 } else { 3 };
282        crate::block_layout::block_count::validate_block_count::<LatentSurvivalError>(
283            "LatentSurvivalFamily",
284            expected_blocks,
285            block_states.len(),
286        )?;
287        let n = self.event_target.len();
288        let eta_time = &block_states[Self::BLOCK_TIME].eta;
289        let eta_mean = &block_states[Self::BLOCK_MEAN].eta;
290        if eta_time.len() != 3 * n {
291            return Err(LatentSurvivalError::BlockMismatch {
292                reason: format!(
293                    "latent survival time eta length mismatch: got {}, expected {}",
294                    eta_time.len(),
295                    3 * n
296                ),
297            });
298        }
299        if eta_mean.len() != n || self.weights.len() != n {
300            return Err(LatentSurvivalError::BlockMismatch {
301                reason: "latent survival mean eta dimension mismatch".to_string(),
302            });
303        }
304        Ok((
305            eta_time.slice(s![0..n]),
306            eta_time.slice(s![n..2 * n]),
307            eta_time.slice(s![2 * n..3 * n]),
308            eta_mean,
309        ))
310    }
311
312    /// Per-row interval upper-bound time transform `q_right = x_time_right · β_time
313    /// + time_offset_right`. Shares the time-block coefficients with `q_exit`
314    /// (same monotone basis, evaluated at `R`), so it is read off the time
315    /// block's `beta` rather than carried as an extra eta channel. For
316    /// non-interval rows `x_time_right` equals `x_time_exit`, so the (unused)
317    /// value is simply `q_exit`.
318    fn time_q_right(
319        &self,
320        block_states: &[ParameterBlockState],
321    ) -> Result<Array1<f64>, LatentSurvivalError> {
322        let n = self.event_target.len();
323        let beta_time = &block_states[Self::BLOCK_TIME].beta;
324        if self.x_time_right.ncols() != beta_time.len() {
325            return Err(LatentSurvivalError::BlockMismatch {
326                reason: format!(
327                    "latent survival interval right design has {} columns but time beta has {}",
328                    self.x_time_right.ncols(),
329                    beta_time.len()
330                ),
331            });
332        }
333        if self.x_time_right.nrows() != n || self.time_offset_right.len() != n {
334            return Err(LatentSurvivalError::BlockMismatch {
335                reason: "latent survival interval right design/offset row count mismatch"
336                    .to_string(),
337            });
338        }
339        let mut q_right = self.x_time_right.dot(beta_time);
340        q_right += &self.time_offset_right;
341        Ok(q_right)
342    }
343
344    fn latent_sd(&self, block_states: &[ParameterBlockState]) -> Result<f64, LatentSurvivalError> {
345        if let Some(sigma) = self.latent_sd_fixed {
346            return Ok(sigma);
347        }
348        let eta = *block_states
349            .get(Self::BLOCK_LOG_SIGMA)
350            .and_then(|state| state.eta.get(0))
351            .ok_or_else(|| LatentSurvivalError::BlockMismatch {
352                reason: "latent survival learnable log_sigma block is missing".to_string(),
353            })?;
354        let sigma = exp_sigma_from_eta_scalar(eta);
355        if !(sigma.is_finite() && sigma > 0.0) {
356            return Err(LatentSurvivalError::NumericalFailure {
357                reason: format!(
358                    "latent survival learnable sigma became invalid: log_sigma={eta}, sigma={sigma}"
359                ),
360            });
361        }
362        Ok(sigma)
363    }
364}
365
366impl LatentBinaryFamily {
367    pub const BLOCK_TIME: usize = 0;
368    pub const BLOCK_MEAN: usize = 1;
369
370    fn split_time_eta<'a>(
371        &self,
372        block_states: &'a [ParameterBlockState],
373    ) -> Result<(ArrayView1<'a, f64>, ArrayView1<'a, f64>, &'a Array1<f64>), LatentSurvivalError>
374    {
375        crate::block_layout::block_count::validate_block_count::<LatentSurvivalError>(
376            "LatentBinaryFamily",
377            2,
378            block_states.len(),
379        )?;
380        let n = self.event_target.len();
381        let eta_time = &block_states[Self::BLOCK_TIME].eta;
382        let eta_mean = &block_states[Self::BLOCK_MEAN].eta;
383        if eta_time.len() != 3 * n {
384            return Err(LatentSurvivalError::BlockMismatch {
385                reason: format!(
386                    "latent binary time eta length mismatch: got {}, expected {}",
387                    eta_time.len(),
388                    3 * n
389                ),
390            });
391        }
392        if eta_mean.len() != n || self.weights.len() != n {
393            return Err(LatentSurvivalError::BlockMismatch {
394                reason: "latent binary mean eta dimension mismatch".to_string(),
395            });
396        }
397        Ok((
398            eta_time.slice(s![0..n]),
399            eta_time.slice(s![n..2 * n]),
400            eta_mean,
401        ))
402    }
403}
404
405pub fn fixed_latent_hazard_frailty(
406    frailty: &FrailtySpec,
407    context: &str,
408) -> Result<(f64, HazardLoading), String> {
409    fixed_latent_hazard_frailty_typed(frailty, context).map_err(Into::into)
410}
411
412fn fixed_latent_hazard_frailty_typed(
413    frailty: &FrailtySpec,
414    context: &str,
415) -> Result<(f64, HazardLoading), LatentSurvivalError> {
416    match frailty {
417        FrailtySpec::HazardMultiplier {
418            sigma_fixed: Some(sigma),
419            loading,
420        } if sigma.is_finite() && *sigma >= 0.0 => Ok((*sigma, *loading)),
421        FrailtySpec::HazardMultiplier {
422            sigma_fixed: Some(sigma),
423            ..
424        } => Err(LatentSurvivalError::InvalidFrailty {
425            reason: format!(
426                "{context} requires a finite fixed hazard-multiplier sigma >= 0, got {sigma}"
427            ),
428        }),
429        FrailtySpec::HazardMultiplier {
430            sigma_fixed: None, ..
431        } => Err(LatentSurvivalError::InvalidFrailty {
432            reason: format!("{context} currently requires a fixed hazard-multiplier sigma"),
433        }),
434        FrailtySpec::GaussianShift { .. } => Err(LatentSurvivalError::InvalidFrailty {
435            reason: format!("{context} requires HazardMultiplier frailty, not GaussianShift"),
436        }),
437        FrailtySpec::None => Err(LatentSurvivalError::InvalidFrailty {
438            reason: format!("{context} requires a fixed HazardMultiplier frailty specification"),
439        }),
440    }
441}
442
443pub fn latent_hazard_loading(
444    frailty: &FrailtySpec,
445    context: &str,
446) -> Result<HazardLoading, String> {
447    latent_hazard_loading_typed(frailty, context).map_err(Into::into)
448}
449
450fn latent_hazard_loading_typed(
451    frailty: &FrailtySpec,
452    context: &str,
453) -> Result<HazardLoading, LatentSurvivalError> {
454    match frailty {
455        FrailtySpec::HazardMultiplier { loading, .. } => Ok(*loading),
456        FrailtySpec::GaussianShift { .. } => Err(LatentSurvivalError::InvalidFrailty {
457            reason: format!("{context} requires HazardMultiplier frailty, not GaussianShift"),
458        }),
459        FrailtySpec::None => Err(LatentSurvivalError::InvalidFrailty {
460            reason: format!("{context} requires a HazardMultiplier frailty specification"),
461        }),
462    }
463}
464
465#[derive(Clone, Copy)]
466struct LatentSurvivalTimeJet {
467    grad_entry: f64,
468    grad_exit: f64,
469    neg_hess_entry: f64,
470    neg_hess_exit: f64,
471}
472
473pub fn fit_latent_survival_terms(
474    data: ArrayView2<'_, f64>,
475    spec: LatentSurvivalTermSpec,
476    frailty: FrailtySpec,
477    options: &BlockwiseFitOptions,
478) -> Result<LatentSurvivalTermFitResult, String> {
479    let latent_sd = validate_latent_survival_inputs(data, &spec, &frailty)?;
480    let hazard_loading = latent_hazard_loading(&frailty, "latent-survival")?;
481    let mean_design =
482        build_term_collection_design(data, &spec.meanspec).map_err(|e| e.to_string())?;
483    let resolvedspec = freeze_term_collection_from_design(&spec.meanspec, &mean_design)
484        .map_err(|e| e.to_string())?;
485    let time_prepared = prepare_latent_time_block(
486        &spec.time_block,
487        spec.time_design_right.as_ref(),
488        spec.derivative_guard,
489    )?;
490
491    let n = spec.event_target.len();
492    let time_offset_right = match spec.time_offset_right.as_ref() {
493        Some(offset) => {
494            if offset.len() != n {
495                return Err(format!(
496                    "latent survival interval right time offset must have length {n}, got {}",
497                    offset.len()
498                ));
499            }
500            offset.clone()
501        }
502        None => Array1::zeros(n),
503    };
504    let unloaded_mass_right = if spec.unloaded_mass_right.is_empty() {
505        Array1::zeros(n)
506    } else {
507        if spec.unloaded_mass_right.len() != n {
508            return Err(format!(
509                "latent survival interval right unloaded mass must have length {n}, got {}",
510                spec.unloaded_mass_right.len()
511            ));
512        }
513        spec.unloaded_mass_right.clone()
514    };
515
516    let family = LatentSurvivalFamily {
517        event_target: spec.event_target.clone(),
518        weights: spec.weights.clone(),
519        latent_sd_fixed: latent_sd,
520        hazard_loading,
521        unloaded_mass_entry: spec.unloaded_mass_entry.clone(),
522        unloaded_mass_exit: spec.unloaded_mass_exit.clone(),
523        unloaded_hazard_exit: spec.unloaded_hazard_exit.clone(),
524        x_time_entry: time_prepared.design_entry.clone(),
525        x_time_exit: time_prepared.design_exit.clone(),
526        x_time_derivative_exit: time_prepared.design_derivative_exit.clone(),
527        x_time_right: time_prepared.design_right.clone(),
528        time_offset_right,
529        unloaded_mass_right,
530        x_mean: mean_design.design.clone(),
531        time_linear_constraints: time_prepared.linear_constraints.clone(),
532        quadctx: Arc::new(QuadratureContext::new()),
533    };
534
535    let mut blocks = vec![
536        build_time_blockspec(&time_prepared, &spec.time_block),
537        build_mean_blockspec(&mean_design, spec.mean_offset.clone()),
538    ];
539    if latent_sd.is_none() {
540        blocks.push(build_log_sigma_blockspec(
541            LEARNABLE_LATENT_SD_SEED,
542            mean_design.design.nrows(),
543        ));
544    }
545    // Interval warm start (issue #1108). Interval-censored rows contribute the
546    // NON-concave `ℓ = log[S(L) − S(R)]`; the coupled exact-joint inner Newton
547    // diverges from the cold seed (β_time = 1e-4, σ = 0.5) — the failure surfaces
548    // first as `fit_custom_family`'s outer ρ-seed startup validation rejecting
549    // every seed (`solver_started = 0`). We warm-start from a LOG-CONCAVE
550    // surrogate whose β/σ land in the interval basin, threaded via `initial_beta`
551    // (consumed by every inner solve, including each ρ-seed validation fit).
552    //
553    // Surrogate = right-censored at the bracket LOWER bound `L`. Its survival
554    // mass `S(L) = K_{0,B(L)}` is log-concave (PD Hessian) and — crucially —
555    // its time-block design is the SAME fixed-knot I-spline basis the interval
556    // fit uses, which is FULL RANK regardless of how heavily the inspection-grid
557    // `L` values are TIED (the basis columns are functions of the frozen knots,
558    // not of the observed time multiplicities). Unlike an exact-event surrogate
559    // it imposes NO per-row `q̇(L) > 0` hazard-derivative feasibility condition
560    // (which the tied/degenerate cold-start derivative design can violate), so it
561    // is robust where exact-event-at-L is not. The warm σ then refines from the
562    // bracket-width spread inside the (now in-basin) interval fit.
563    //
564    // Failure is NON-SILENT (#1108): a surrogate that errors or returns a
565    // non-finite / all-zero degenerate β is surfaced as a hard error rather than
566    // silently reverting to the diverging cold start (which masked the real
567    // failure across several attempts). Only `initial_beta` is seeded; the EXACT
568    // interval objective/gradient/Hessian are unchanged, so σ̂ is the true MLE.
569    let has_interval_rows = spec
570        .event_target
571        .iter()
572        .any(|&code| code == LATENT_SURVIVAL_EVENT_INTERVAL);
573    if has_interval_rows {
574        let censored_warm_event_target = spec.event_target.mapv(|code| {
575            if code == LATENT_SURVIVAL_EVENT_INTERVAL {
576                0u8
577            } else {
578                code
579            }
580        });
581        let mut warm_family = family.clone();
582        warm_family.event_target = censored_warm_event_target;
583        // Right-censored-at-L ignores the interval upper bound `R`, so the
584        // (unused) `q_right` channel cannot drift the fit; leaving the right
585        // design/mass in place is harmless (no interval row remains to read it).
586        let warm_fit_result = fit_custom_family_fixed_log_lambdas(
587            &warm_family,
588            &blocks,
589            options,
590            None,
591            0,
592            None,
593            false,
594        );
595        let warm_fit = match warm_fit_result {
596            Ok(fit) => fit,
597            Err(censored_error) => {
598                let has_finite_event_in_censored_surrogate =
599                    warm_family.event_target.iter().any(|&code| code != 0);
600                if has_finite_event_in_censored_surrogate {
601                    return Err(format!(
602                        "latent interval warm start: right-censored-at-L surrogate fit failed \
603                         (so the interval fit cannot be safely warm-started; this surrogate is \
604                         log-concave and should converge — investigate the surrogate, not the \
605                         interval kernel): {censored_error}"
606                    ));
607                }
608
609                // When every observed row is interval-censored, the
610                // right-censored-at-L surrogate contains no failures at all.
611                // Its likelihood is maximized only on the zero-hazard boundary
612                // (β_time -> -∞), so the fixed-λ Newton solve is correctly
613                // allowed to refuse it even though the objective is concave.
614                // Use the finite lower-endpoint event surrogate solely to obtain
615                // an interior β/σ seed for the exact interval likelihood below;
616                // no fitted surrogate likelihood or derivative is reused.
617                let lower_event_warm_target = spec.event_target.mapv(|code| {
618                    if code == LATENT_SURVIVAL_EVENT_INTERVAL {
619                        1u8
620                    } else {
621                        code
622                    }
623                });
624                let mut event_warm_family = family.clone();
625                event_warm_family.event_target = lower_event_warm_target;
626                fit_custom_family_fixed_log_lambdas(
627                    &event_warm_family,
628                    &blocks,
629                    options,
630                    None,
631                    0,
632                    None,
633                    false,
634                )
635                .map_err(|event_error| {
636                    format!(
637                        "latent interval warm start failed: the right-censored-at-L surrogate \
638                         has no finite failures and refused its boundary optimum ({censored_error}); \
639                         the finite lower-endpoint event surrogate also failed ({event_error})"
640                    )
641                })?
642            }
643        };
644        let warm_beta_usable = warm_fit
645            .block_states
646            .iter()
647            .any(|s| s.beta.iter().all(|v| v.is_finite()) && s.beta.iter().any(|&v| v != 0.0));
648        if !warm_beta_usable {
649            return Err(
650                "latent interval warm start: right-censored-at-L surrogate returned a \
651                 degenerate (non-finite or all-zero) β across every block; the warm start \
652                 cannot seed the interval fit. This indicates the surrogate's time-block \
653                 design is rank-deficient or the inner solve stalled at the seed — \
654                 investigate the surrogate before retrying the interval fit."
655                    .to_string(),
656            );
657        }
658        for (block, state) in blocks.iter_mut().zip(warm_fit.block_states.iter()) {
659            if state.beta.iter().all(|v| v.is_finite()) {
660                block.initial_beta = Some(state.beta.clone());
661            }
662        }
663    }
664    let fit = fit_custom_family(&family, &blocks, options).map_err(|e| e.to_string())?;
665    let latent_sd = family.latent_sd(&fit.block_states)?;
666    let baseline_offset_residuals = family.offset_channel_residuals(&fit.block_states)?;
667    Ok(LatentSurvivalTermFitResult {
668        fit,
669        design: mean_design,
670        resolvedspec,
671        latent_sd,
672        baseline_offset_residuals,
673    })
674}
675
676pub fn fit_latent_binary_terms(
677    data: ArrayView2<'_, f64>,
678    spec: LatentBinaryTermSpec,
679    frailty: FrailtySpec,
680    options: &BlockwiseFitOptions,
681) -> Result<LatentBinaryTermFitResult, String> {
682    let latent_sd = validate_latent_binary_inputs(data, &spec, &frailty)?;
683    let (_, hazard_loading) = fixed_latent_hazard_frailty(&frailty, "latent-binary")?;
684    let mean_design =
685        build_term_collection_design(data, &spec.meanspec).map_err(|e| e.to_string())?;
686    let resolvedspec = freeze_term_collection_from_design(&spec.meanspec, &mean_design)
687        .map_err(|e| e.to_string())?;
688    let time_prepared = prepare_latent_time_block(&spec.time_block, None, spec.derivative_guard)?;
689
690    let family = LatentBinaryFamily {
691        event_target: spec.event_target.clone(),
692        weights: spec.weights.clone(),
693        latent_sd,
694        hazard_loading,
695        unloaded_mass_entry: spec.unloaded_mass_entry.clone(),
696        unloaded_mass_exit: spec.unloaded_mass_exit.clone(),
697        x_time_entry: time_prepared.design_entry.clone(),
698        x_time_exit: time_prepared.design_exit.clone(),
699        x_mean: mean_design.design.clone(),
700        time_linear_constraints: time_prepared.linear_constraints.clone(),
701        quadctx: Arc::new(QuadratureContext::new()),
702    };
703
704    let blocks = vec![
705        build_time_blockspec(&time_prepared, &spec.time_block),
706        build_mean_blockspec(&mean_design, spec.mean_offset.clone()),
707    ];
708    let fit = fit_custom_family(&family, &blocks, options).map_err(|e| e.to_string())?;
709    let baseline_offset_residuals = family.offset_channel_residuals(&fit.block_states)?;
710    Ok(LatentBinaryTermFitResult {
711        fit,
712        design: mean_design,
713        resolvedspec,
714        baseline_offset_residuals,
715    })
716}
717
718/// Latent-survival adapter for the shared [`LatentIntervalModel`] driver.
719///
720/// Survival permits a learnable sigma (`sigma_fixed == None`) and carries the
721/// per-row unloaded baseline hazard at exit (which feeds the exact-event
722/// loaded/unloaded split); everything else is validated by the shared engine.
723struct LatentSurvivalModel;
724
725impl LatentIntervalModel for LatentSurvivalModel {
726    fn context() -> &'static str {
727        "latent-survival"
728    }
729
730    fn allows_interval() -> bool {
731        true
732    }
733
734    fn frailty_policy(
735        frailty: &FrailtySpec,
736    ) -> Result<LatentFrailtyResolution, LatentSurvivalError> {
737        match frailty {
738            FrailtySpec::HazardMultiplier {
739                sigma_fixed,
740                loading,
741            } => {
742                if let Some(sigma) = sigma_fixed
743                    && (!sigma.is_finite() || *sigma < 0.0)
744                {
745                    return Err(LatentSurvivalError::InvalidFrailty {
746                        reason: format!(
747                            "latent-survival requires a finite hazard-multiplier sigma >= 0, got {sigma}"
748                        ),
749                    });
750                }
751                Ok(LatentFrailtyResolution {
752                    sigma: *sigma_fixed,
753                    loading: *loading,
754                })
755            }
756            FrailtySpec::GaussianShift { .. } => Err(LatentSurvivalError::InvalidFrailty {
757                reason: "latent-survival requires HazardMultiplier frailty, not GaussianShift"
758                    .to_string(),
759            }),
760            FrailtySpec::None => Err(LatentSurvivalError::InvalidFrailty {
761                reason: "latent-survival requires a HazardMultiplier frailty specification"
762                    .to_string(),
763            }),
764        }
765    }
766}
767
768fn validate_latent_survival_inputs(
769    data: ArrayView2<'_, f64>,
770    spec: &LatentSurvivalTermSpec,
771    frailty: &FrailtySpec,
772) -> Result<Option<f64>, LatentSurvivalError> {
773    let row = LatentIntervalRowView {
774        frailty,
775        age_entry: &spec.age_entry,
776        age_exit: &spec.age_exit,
777        event_target: &spec.event_target,
778        weights: &spec.weights,
779        unloaded_mass_entry: &spec.unloaded_mass_entry,
780        unloaded_mass_exit: &spec.unloaded_mass_exit,
781        unloaded_hazard_exit: Some(&spec.unloaded_hazard_exit),
782        mean_offset: &spec.mean_offset,
783        derivative_guard: spec.derivative_guard,
784        time_block: &spec.time_block,
785    };
786    validate_latent_interval_inputs::<LatentSurvivalModel>(data, &row)
787}
788
789pub(crate) fn validate_unloaded_components_for_loading(
790    context: &str,
791    row_index: usize,
792    loading: HazardLoading,
793    unloaded_entry: f64,
794    unloaded_exit: f64,
795    unloaded_hazard: Option<f64>,
796) -> Result<(), LatentSurvivalError> {
797    match loading {
798        HazardLoading::Full => {
799            if unloaded_entry != 0.0
800                || unloaded_exit != 0.0
801                || unloaded_hazard.is_some_and(|hazard| hazard != 0.0)
802            {
803                return Err(LatentSurvivalError::InvalidDataset {
804                    reason: format!(
805                        "{context} row {} uses full hazard loading, so unloaded components must be exactly zero; got entry_mass={}, exit_mass={}, exit_hazard={}",
806                        row_index + 1,
807                        unloaded_entry,
808                        unloaded_exit,
809                        unloaded_hazard.unwrap_or(0.0)
810                    ),
811                });
812            }
813        }
814        HazardLoading::LoadedVsUnloaded => {}
815    }
816    Ok(())
817}
818
819/// Latent-binary adapter for the shared [`LatentIntervalModel`] driver.
820///
821/// Binary never evaluates an exact event, so it requires a finite *fixed*
822/// latent sigma (via [`fixed_latent_hazard_frailty_typed`]) and carries no
823/// per-row unloaded hazard; every other invariant is validated by the shared
824/// engine.
825struct LatentBinaryModel;
826
827impl LatentIntervalModel for LatentBinaryModel {
828    fn context() -> &'static str {
829        "latent-binary"
830    }
831
832    fn frailty_policy(
833        frailty: &FrailtySpec,
834    ) -> Result<LatentFrailtyResolution, LatentSurvivalError> {
835        let (sigma, loading) = fixed_latent_hazard_frailty_typed(frailty, "latent-binary")?;
836        Ok(LatentFrailtyResolution {
837            sigma: Some(sigma),
838            loading,
839        })
840    }
841}
842
843fn validate_latent_binary_inputs(
844    data: ArrayView2<'_, f64>,
845    spec: &LatentBinaryTermSpec,
846    frailty: &FrailtySpec,
847) -> Result<f64, LatentSurvivalError> {
848    let row = LatentIntervalRowView {
849        frailty,
850        age_entry: &spec.age_entry,
851        age_exit: &spec.age_exit,
852        event_target: &spec.event_target,
853        weights: &spec.weights,
854        unloaded_mass_entry: &spec.unloaded_mass_entry,
855        unloaded_mass_exit: &spec.unloaded_mass_exit,
856        unloaded_hazard_exit: None,
857        mean_offset: &spec.mean_offset,
858        derivative_guard: spec.derivative_guard,
859        time_block: &spec.time_block,
860    };
861    // The binary `frailty_policy` always yields `Some(sigma)` (it rejects the
862    // learnable-scale case), so the shared driver's `Option<f64>` is `Some`
863    // here; surface a structured error rather than unwrapping if that ever
864    // changes.
865    validate_latent_interval_inputs::<LatentBinaryModel>(data, &row)?.ok_or_else(|| {
866        LatentSurvivalError::InvalidFrailty {
867            reason: "latent-binary requires a fixed latent sigma".to_string(),
868        }
869    })
870}
871
872fn prepare_latent_time_block(
873    input: &TimeBlockInput,
874    design_right: Option<&DesignMatrix>,
875    derivative_guard: f64,
876) -> Result<PreparedLatentTimeBlock, LatentSurvivalError> {
877    if !input.time_monotonicity.is_coordinate_cone() {
878        return Err(LatentSurvivalError::UnsupportedConfiguration {
879            reason: format!(
880                "latent survival requires a coordinate-cone monotonicity strategy; got {:?}",
881                input.time_monotonicity
882            ),
883        });
884    }
885    let design_entry = input
886        .design_entry
887        .try_to_dense_by_chunks("latent survival entry time design")?;
888    let design_exit = input
889        .design_exit
890        .try_to_dense_by_chunks("latent survival exit time design")?;
891    let design_derivative_exit = input
892        .design_derivative_exit
893        .try_to_dense_by_chunks("latent survival derivative time design")?;
894    // The interval upper-bound design shares the time-block coefficients with
895    // the exit design; when the data has no interval rows we reuse the exit
896    // design so `q_right` stays well-defined (its likelihood contribution is
897    // gated off for non-interval rows). When present it must match the exit
898    // design's shape (same basis, evaluated at R).
899    let design_right = match design_right {
900        Some(matrix) => {
901            let dense =
902                matrix.try_to_dense_by_chunks("latent survival interval right time design")?;
903            if dense.nrows() != design_exit.nrows() || dense.ncols() != design_exit.ncols() {
904                return Err(LatentSurvivalError::InvalidDataset {
905                    reason: format!(
906                        "latent survival interval right time design must match exit design shape \
907                         {:?}, got {:?}",
908                        design_exit.dim(),
909                        dense.dim()
910                    ),
911                });
912            }
913            dense
914        }
915        None => design_exit.clone(),
916    };
917    let linear_constraints = structural_time_coefficient_constraints(
918        &input.design_derivative_exit,
919        &input.derivative_offset_exit,
920        derivative_guard,
921    )?;
922    let initial_beta = match linear_constraints.as_ref() {
923        // `project_onto_linear_constraints` validates that any supplied
924        // `initial_beta` matches `design_exit.ncols()`; surface a mismatch as a
925        // structured error rather than letting an ndarray broadcast panic
926        // (issue #374).
927        Some(constraints) => Some(project_onto_linear_constraints(
928            design_exit.ncols(),
929            constraints,
930            input.initial_beta.as_ref(),
931        )?),
932        None => None,
933    };
934    Ok(PreparedLatentTimeBlock {
935        design_entry,
936        design_exit,
937        design_derivative_exit,
938        design_right,
939        linear_constraints,
940        penalties: input.penalties.clone(),
941        initial_beta,
942    })
943}
944
945fn stack_rows(blocks: &[&Array2<f64>]) -> Array2<f64> {
946    let ncols = blocks.first().map_or(0, |m| m.ncols());
947    let nrows = blocks.iter().map(|m| m.nrows()).sum();
948    let mut out = Array2::<f64>::zeros((nrows, ncols));
949    let mut row = 0usize;
950    for block in blocks {
951        let end = row + block.nrows();
952        out.slice_mut(s![row..end, ..]).assign(block);
953        row = end;
954    }
955    out
956}
957
958fn build_time_blockspec(
959    prepared: &PreparedLatentTimeBlock,
960    input: &TimeBlockInput,
961) -> ParameterBlockSpec {
962    // The solver produces a `3·n`-long time `eta` (the `[entry; exit; deriv]`
963    // channel stack that `split_time_eta` slices). That stacked operator is
964    // the eta-producing matrix and so belongs in `stacked_design`, paired with
965    // the matching `3·n`-row stacked offset. The audit / shape-policy invariant
966    // `design.nrows() == n_obs` is satisfied by exposing the single-channel
967    // n-row exit design as `design`; the audit never inspects `stacked_design`.
968    //
969    // This mirrors the survival location-scale fix for the same #326 class
970    // (`survival_location_scale.rs`): the previous code put the `3·n`-row
971    // stack in `design`, which tripped the flat identifiability audit's
972    // row-equality invariant (`block 1 (mean) has n rows, expected 3n`).
973    let stacked_design = stack_rows(&[
974        &prepared.design_entry,
975        &prepared.design_exit,
976        &prepared.design_derivative_exit,
977    ]);
978    let stacked_offset = gam_linalg::utils::stack_offsets(&[
979        &input.offset_entry,
980        &input.offset_exit,
981        &input.derivative_offset_exit,
982    ]);
983    ParameterBlockSpec {
984        name: "time_transform".to_string(),
985        design: DesignMatrix::Dense(DenseDesignMatrix::from(Arc::new(
986            prepared.design_exit.clone(),
987        ))),
988        offset: input.offset_exit.clone(),
989        penalties: prepared
990            .penalties
991            .iter()
992            .cloned()
993            .map(PenaltyMatrix::Dense)
994            .collect(),
995        nullspace_dims: input.nullspace_dims.clone(),
996        initial_log_lambdas: input
997            .initial_log_lambdas
998            .clone()
999            .unwrap_or_else(|| Array1::zeros(prepared.penalties.len())),
1000        initial_beta: prepared.initial_beta.clone(),
1001        // Canonical-gauge ownership for the latent-survival joint design: the
1002        // time-transform block carries the structural monotone baseline that
1003        // anchors the parameterisation, so it owns any shared constant
1004        // direction (strictly higher than `mean`/`log_sigma` at 100). This
1005        // matches the survival location-scale gauge contract (time highest).
1006        gauge_priority: 200,
1007        jacobian_callback: None,
1008        stacked_design: Some(DesignMatrix::Dense(DenseDesignMatrix::from(Arc::new(
1009            stacked_design,
1010        )))),
1011        stacked_offset: Some(stacked_offset),
1012    }
1013}
1014
1015fn build_mean_blockspec(design: &TermCollectionDesign, offset: Array1<f64>) -> ParameterBlockSpec {
1016    ParameterBlockSpec {
1017        name: "mean".to_string(),
1018        design: design.design.clone(),
1019        offset,
1020        penalties: design.penalties_as_penalty_matrix(),
1021        nullspace_dims: design.nullspace_dims.clone(),
1022        initial_log_lambdas: Array1::zeros(design.penalties.len()),
1023        initial_beta: None,
1024        // Strictly below `time_transform` (200) so any constant direction
1025        // shared between the monotone time baseline and the mean intercept is
1026        // deterministically attributable to the lower-priority `mean` block by
1027        // the canonical-gauge RRQR (the descending-priority contract used by
1028        // survival location-scale; #366/#556 gauge story).
1029        gauge_priority: 150,
1030        jacobian_callback: None,
1031        stacked_design: None,
1032        stacked_offset: None,
1033    }
1034}
1035
1036/// Starting latent-frailty standard deviation when `sigma` is learnable
1037/// (`sigma_fixed == None`). The log-sigma block is seeded at `log(0.5)` so the
1038/// optimizer begins from a moderate, well-conditioned dispersion (σ = 0.5,
1039/// neither a near-degenerate σ → 0 that flattens the frailty integral nor a
1040/// large σ that makes the Gauss-Hermite quadrature heavy-tailed) and then
1041/// learns the data's actual scale. Only an initial value, not a constraint.
1042const LEARNABLE_LATENT_SD_SEED: f64 = 0.5;
1043
1044fn build_log_sigma_blockspec(initial_sigma: f64, n_obs: usize) -> ParameterBlockSpec {
1045    ParameterBlockSpec {
1046        name: "log_sigma".to_string(),
1047        // The frailty/dispersion scale is a single GLOBAL hyperparameter (one free
1048        // coefficient), but the identifiability audit — and the canonical-row
1049        // architecture generally — require every block's effective Jacobian to carry
1050        // `n_obs` rows. A global scalar is realised the same way the survival
1051        // location-scale `log_sigma` block is (see `BinomialLocationScaleFamily`): an
1052        // `n_obs × 1` constant column of ones, so `eta = design · β` is the same scalar
1053        // broadcast to every observation. This keeps it a single free parameter while
1054        // exposing the `n_obs`-row shape the audit checks, and `latent_sd` reads
1055        // `eta[0]` — identical across rows by construction.
1056        design: DesignMatrix::Dense(DenseDesignMatrix::from(Arc::new(Array2::from_elem(
1057            (n_obs, 1),
1058            1.0,
1059        )))),
1060        offset: Array1::zeros(n_obs),
1061        penalties: vec![],
1062        nullspace_dims: vec![],
1063        initial_log_lambdas: Array1::zeros(0),
1064        initial_beta: Some(Array1::from_elem(
1065            1,
1066            exp_sigma_eta_for_sigma_scalar(initial_sigma),
1067        )),
1068        // Lowest of the three (time=200, mean=150): the learnable-scale channel
1069        // yields any shared constant to the location blocks.
1070        gauge_priority: 120,
1071        jacobian_callback: None,
1072        stacked_design: None,
1073        stacked_offset: None,
1074    }
1075}
1076
1077const LATENT_SURVIVAL_PRIMARY_Q_ENTRY: usize = 0;
1078const LATENT_SURVIVAL_PRIMARY_Q_EXIT: usize = 1;
1079const LATENT_SURVIVAL_PRIMARY_QDOT_EXIT: usize = 2;
1080// Interval-censored right boundary R: q_right = log B(R) shares the time-block
1081// coefficients with q_exit (same monotone transform, different time point), so
1082// it is a fourth linear functional of the time block, NOT an independent eta
1083// channel. It sits before `mu`/`log_sigma` so the "trailing optional log_sigma"
1084// invariant used by `active_primary` (= `LATENT_SURVIVAL_PRIMARY_LOG_SIGMA`)
1085// keeps q_right always active.
1086const LATENT_SURVIVAL_PRIMARY_Q_RIGHT: usize = 3;
1087const LATENT_SURVIVAL_PRIMARY_MU: usize = 4;
1088const LATENT_SURVIVAL_PRIMARY_LOG_SIGMA: usize = 5;
1089const LATENT_SURVIVAL_PRIMARY_DIM: usize = 6;
1090
1091use gam_math::jet_partitions::MultiDirJet as LatentMultiDirJet;
1092
1093/// Derivatives of `log(x)` through 4th order.
1094///
1095/// # Contract
1096///
1097/// `x` must be strictly positive. This function does NOT clamp: a previous
1098/// version replaced `x` by `x.max(1e-300)`, which fabricated enormous finite
1099/// derivatives (`1/1e-300` etc.) that are the derivatives of neither `log(x)`
1100/// nor `log(max(x, floor))` and would silently mask an upstream domain
1101/// failure. Both callers guarantee `x > 0`: one composes at the literal `1.0`
1102/// (the normalised log-sum base); the other passes `base`, which is gated by
1103/// an explicit `base.is_finite() && base > 0.0` check immediately upstream. A
1104/// non-positive `x` therefore never reaches here on any supported path; were
1105/// one to, the function returns the honest IEEE result (`-inf`/`NaN`) —
1106/// identical in debug and release — rather than a finite fabrication. For all
1107/// valid `x > 0` the output is bit-identical to the previous clamped version.
1108#[inline]
1109fn latent_unary_derivatives_log(x: f64) -> [f64; 5] {
1110    let x2 = x * x;
1111    let x3 = x2 * x;
1112    let x4 = x3 * x;
1113    [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
1114}
1115
1116#[derive(Clone, Copy, Debug)]
1117struct LatentKernelPrimaryTerm {
1118    coeff: f64,
1119    q_exp: usize,
1120    qdot_power: usize,
1121    tau_exp: usize,
1122    k: usize,
1123}
1124
1125#[derive(Clone, Copy, Debug)]
1126struct LatentKernelPrimaryDirection {
1127    dq: f64,
1128    dqd: f64,
1129    dmu: f64,
1130    dtau: f64,
1131}
1132
1133#[derive(Clone, Copy, Debug)]
1134struct LatentSurvivalPrimaryDirection {
1135    dq_entry: f64,
1136    dq_exit: f64,
1137    dqdot_exit: f64,
1138    dq_right: f64,
1139    dmu: f64,
1140    dlog_sigma: f64,
1141}
1142
1143#[derive(Clone, Copy, Debug)]
1144struct LatentKernelPrimaryState {
1145    q: f64,
1146    qdot: f64,
1147    mu: f64,
1148    sigma: f64,
1149    log_sigma_factor: f64,
1150}
1151
1152fn latent_kernel_accumulate_term(
1153    terms: &mut BTreeMap<(usize, usize, usize, usize), f64>,
1154    term: LatentKernelPrimaryTerm,
1155    scale: f64,
1156) {
1157    if scale == 0.0 || term.coeff == 0.0 {
1158        return;
1159    }
1160    *terms
1161        .entry((term.q_exp, term.qdot_power, term.tau_exp, term.k))
1162        .or_insert(0.0) += scale * term.coeff;
1163}
1164
1165fn latent_kernel_differentiate_terms(
1166    terms: &[LatentKernelPrimaryTerm],
1167    dir: LatentKernelPrimaryDirection,
1168) -> Vec<LatentKernelPrimaryTerm> {
1169    let mut out = BTreeMap::<(usize, usize, usize, usize), f64>::new();
1170    for term in terms {
1171        if dir.dq != 0.0 {
1172            if term.q_exp > 0 {
1173                latent_kernel_accumulate_term(&mut out, *term, dir.dq * term.q_exp as f64);
1174            }
1175            latent_kernel_accumulate_term(
1176                &mut out,
1177                LatentKernelPrimaryTerm {
1178                    q_exp: term.q_exp + 1,
1179                    k: term.k + 1,
1180                    ..*term
1181                },
1182                -dir.dq,
1183            );
1184        }
1185        if dir.dmu != 0.0 {
1186            if term.k > 0 {
1187                latent_kernel_accumulate_term(&mut out, *term, dir.dmu * term.k as f64);
1188            }
1189            latent_kernel_accumulate_term(
1190                &mut out,
1191                LatentKernelPrimaryTerm {
1192                    q_exp: term.q_exp + 1,
1193                    k: term.k + 1,
1194                    ..*term
1195                },
1196                -dir.dmu,
1197            );
1198        }
1199        if dir.dtau != 0.0 {
1200            if term.tau_exp > 0 {
1201                latent_kernel_accumulate_term(&mut out, *term, dir.dtau * term.tau_exp as f64);
1202            }
1203            let kf = term.k as f64;
1204            latent_kernel_accumulate_term(
1205                &mut out,
1206                LatentKernelPrimaryTerm {
1207                    tau_exp: term.tau_exp + 2,
1208                    ..*term
1209                },
1210                dir.dtau * kf * kf,
1211            );
1212            latent_kernel_accumulate_term(
1213                &mut out,
1214                LatentKernelPrimaryTerm {
1215                    q_exp: term.q_exp + 1,
1216                    tau_exp: term.tau_exp + 2,
1217                    k: term.k + 1,
1218                    ..*term
1219                },
1220                -dir.dtau * (2.0 * kf + 1.0),
1221            );
1222            latent_kernel_accumulate_term(
1223                &mut out,
1224                LatentKernelPrimaryTerm {
1225                    q_exp: term.q_exp + 2,
1226                    tau_exp: term.tau_exp + 2,
1227                    k: term.k + 2,
1228                    ..*term
1229                },
1230                dir.dtau,
1231            );
1232        }
1233        if dir.dqd != 0.0 && term.qdot_power > 0 {
1234            latent_kernel_accumulate_term(
1235                &mut out,
1236                LatentKernelPrimaryTerm {
1237                    qdot_power: term.qdot_power - 1,
1238                    ..*term
1239                },
1240                dir.dqd * term.qdot_power as f64,
1241            );
1242        }
1243    }
1244    out.into_iter()
1245        .filter_map(|((q_exp, qdot_power, tau_exp, k), coeff)| {
1246            (coeff != 0.0).then_some(LatentKernelPrimaryTerm {
1247                coeff,
1248                q_exp,
1249                qdot_power,
1250                tau_exp,
1251                k,
1252            })
1253        })
1254        .collect()
1255}
1256
1257fn latent_kernel_term_lists_for_directions(
1258    base_terms: &[LatentKernelPrimaryTerm],
1259    directions: &[LatentKernelPrimaryDirection],
1260) -> Vec<Vec<LatentKernelPrimaryTerm>> {
1261    fn build_mask(
1262        mask: usize,
1263        base_terms: &[LatentKernelPrimaryTerm],
1264        directions: &[LatentKernelPrimaryDirection],
1265        cache: &mut [Option<Vec<LatentKernelPrimaryTerm>>],
1266    ) -> Vec<LatentKernelPrimaryTerm> {
1267        if let Some(existing) = &cache[mask] {
1268            return existing.clone();
1269        }
1270        let built = if mask == 0 {
1271            base_terms.to_vec()
1272        } else {
1273            let bit = 1usize << mask.trailing_zeros();
1274            let prev = build_mask(mask ^ bit, base_terms, directions, cache);
1275            latent_kernel_differentiate_terms(&prev, directions[bit.trailing_zeros() as usize])
1276        };
1277        cache[mask] = Some(built.clone());
1278        built
1279    }
1280
1281    let mut cache = vec![None; 1usize << directions.len()];
1282    (0..cache.len())
1283        .map(|mask| build_mask(mask, base_terms, directions, &mut cache))
1284        .collect()
1285}
1286
1287fn latent_kernel_sum_log_jet(
1288    quadctx: &QuadratureContext,
1289    base_terms: &[LatentKernelPrimaryTerm],
1290    state: LatentKernelPrimaryState,
1291    directions: &[LatentKernelPrimaryDirection],
1292    context: &str,
1293) -> Result<LatentMultiDirJet, LatentSurvivalError> {
1294    let term_lists = latent_kernel_term_lists_for_directions(base_terms, directions);
1295    let max_k = term_lists
1296        .iter()
1297        .flat_map(|terms| terms.iter().map(|term| term.k))
1298        .max()
1299        .unwrap_or(0);
1300    let bundle =
1301        log_kernel_bundle(quadctx, state.q.exp(), state.mu, state.sigma, max_k).map_err(|e| {
1302            LatentSurvivalError::NumericalFailure {
1303                reason: format!("{context} kernel evaluation failed: {e}"),
1304            }
1305        })?;
1306
1307    let evaluate_terms =
1308        |terms: &[LatentKernelPrimaryTerm]| -> Result<(f64, f64), LatentSurvivalError> {
1309            let mut log_mags = Vec::new();
1310            let mut signs = Vec::new();
1311            for term in terms {
1312                if term.coeff == 0.0 {
1313                    continue;
1314                }
1315                if term.qdot_power > 0 && !(state.qdot.is_finite() && state.qdot > 0.0) {
1316                    return Err(LatentSurvivalError::NumericalFailure {
1317                        reason: format!(
1318                            "{context} requires positive finite qdot for exact-event directional terms, got {}",
1319                            state.qdot
1320                        ),
1321                    });
1322                }
1323                let log_qdot = if term.qdot_power > 0 {
1324                    state.qdot.ln()
1325                } else {
1326                    0.0
1327                };
1328                let log_mag = term.coeff.abs().ln()
1329                    + term.q_exp as f64 * state.q
1330                    + term.tau_exp as f64 * state.log_sigma_factor
1331                    + term.qdot_power as f64 * log_qdot
1332                    + bundle.get(term.k);
1333                log_mags.push(log_mag);
1334                signs.push(term.coeff.signum());
1335            }
1336            if log_mags.is_empty() {
1337                return Ok((f64::NEG_INFINITY, 0.0));
1338            }
1339            Ok(signed_log_sum_exp(&log_mags, &signs))
1340        };
1341
1342    let (base_log_sum, base_sign) = evaluate_terms(&term_lists[0])?;
1343    if !(base_log_sum.is_finite() && base_sign > 0.0) {
1344        return Err(LatentSurvivalError::NumericalFailure {
1345            reason: format!("{context} produced a non-positive signed kernel sum"),
1346        });
1347    }
1348
1349    let mut normalized = LatentMultiDirJet::constant(directions.len(), 1.0);
1350    for mask in 1..term_lists.len() {
1351        let (log_abs, sign) = evaluate_terms(&term_lists[mask])?;
1352        normalized.coeffs[mask] = if !log_abs.is_finite() || sign == 0.0 {
1353            0.0
1354        } else {
1355            sign * (log_abs - base_log_sum).exp()
1356        };
1357    }
1358
1359    let mut out = normalized.compose_unary(latent_unary_derivatives_log(1.0));
1360    out.coeffs[0] += base_log_sum;
1361    Ok(out)
1362}
1363
1364fn latent_survival_basis_direction(primary_idx: usize) -> LatentSurvivalPrimaryDirection {
1365    match primary_idx {
1366        LATENT_SURVIVAL_PRIMARY_Q_ENTRY => LatentSurvivalPrimaryDirection {
1367            dq_entry: 1.0,
1368            dq_exit: 0.0,
1369            dqdot_exit: 0.0,
1370            dq_right: 0.0,
1371            dmu: 0.0,
1372            dlog_sigma: 0.0,
1373        },
1374        LATENT_SURVIVAL_PRIMARY_Q_EXIT => LatentSurvivalPrimaryDirection {
1375            dq_entry: 0.0,
1376            dq_exit: 1.0,
1377            dqdot_exit: 0.0,
1378            dq_right: 0.0,
1379            dmu: 0.0,
1380            dlog_sigma: 0.0,
1381        },
1382        LATENT_SURVIVAL_PRIMARY_QDOT_EXIT => LatentSurvivalPrimaryDirection {
1383            dq_entry: 0.0,
1384            dq_exit: 0.0,
1385            dqdot_exit: 1.0,
1386            dq_right: 0.0,
1387            dmu: 0.0,
1388            dlog_sigma: 0.0,
1389        },
1390        LATENT_SURVIVAL_PRIMARY_Q_RIGHT => LatentSurvivalPrimaryDirection {
1391            dq_entry: 0.0,
1392            dq_exit: 0.0,
1393            dqdot_exit: 0.0,
1394            dq_right: 1.0,
1395            dmu: 0.0,
1396            dlog_sigma: 0.0,
1397        },
1398        LATENT_SURVIVAL_PRIMARY_MU => LatentSurvivalPrimaryDirection {
1399            dq_entry: 0.0,
1400            dq_exit: 0.0,
1401            dqdot_exit: 0.0,
1402            dq_right: 0.0,
1403            dmu: 1.0,
1404            dlog_sigma: 0.0,
1405        },
1406        LATENT_SURVIVAL_PRIMARY_LOG_SIGMA => LatentSurvivalPrimaryDirection {
1407            dq_entry: 0.0,
1408            dq_exit: 0.0,
1409            dqdot_exit: 0.0,
1410            dq_right: 0.0,
1411            dmu: 0.0,
1412            dlog_sigma: 1.0,
1413        },
1414        // SAFETY: latent survival has exactly `LATENT_SURVIVAL_PRIMARY_DIM`
1415        // (= 5) primary directions, indexed 0..=4 via the module-private
1416        // `LATENT_SURVIVAL_PRIMARY_*` constants. All five are matched
1417        // above, so this wildcard fires only on an out-of-range index,
1418        // which the internal iteration bounds (`0..LATENT_SURVIVAL_PRIMARY_DIM`)
1419        // make unreachable.
1420        // SAFETY: primary_idx is bounded by LATENT_SURVIVAL_PRIMARY_DIM at every internal call site.
1421        _ => std::panic::panic_any(format!(
1422            "latent survival primary index out of bounds: primary_idx={primary_idx}, primary_dim={LATENT_SURVIVAL_PRIMARY_DIM}"
1423        )),
1424    }
1425}
1426
1427fn latent_survival_map_entry_direction(
1428    direction: LatentSurvivalPrimaryDirection,
1429) -> LatentKernelPrimaryDirection {
1430    LatentKernelPrimaryDirection {
1431        dq: direction.dq_entry,
1432        dqd: 0.0,
1433        dmu: direction.dmu,
1434        dtau: direction.dlog_sigma,
1435    }
1436}
1437
1438fn latent_survival_map_exit_direction(
1439    direction: LatentSurvivalPrimaryDirection,
1440    event_type: LatentSurvivalEventType,
1441) -> LatentKernelPrimaryDirection {
1442    LatentKernelPrimaryDirection {
1443        dq: direction.dq_exit,
1444        dqd: if matches!(event_type, LatentSurvivalEventType::ExactEvent) {
1445            direction.dqdot_exit
1446        } else {
1447            0.0
1448        },
1449        dmu: direction.dmu,
1450        dtau: direction.dlog_sigma,
1451    }
1452}
1453
1454/// Direction map for the interval-censored LEFT boundary state (mass `M_L =
1455/// exp(q_exit)`). The left boundary tracks the same `q_exit` time functional as
1456/// right-censoring (no hazard-derivative channel), plus the shared `mu`/`sigma`.
1457fn latent_survival_map_left_direction(
1458    direction: LatentSurvivalPrimaryDirection,
1459) -> LatentKernelPrimaryDirection {
1460    LatentKernelPrimaryDirection {
1461        dq: direction.dq_exit,
1462        dqd: 0.0,
1463        dmu: direction.dmu,
1464        dtau: direction.dlog_sigma,
1465    }
1466}
1467
1468/// Direction map for the interval-censored RIGHT boundary state (mass `M_R =
1469/// exp(q_right)`). The right boundary tracks the dedicated `q_right` functional
1470/// (which shares the time-block coefficients with `q_exit` but is evaluated at
1471/// the interval upper bound `R`), plus the shared `mu`/`sigma`.
1472fn latent_survival_map_right_direction(
1473    direction: LatentSurvivalPrimaryDirection,
1474) -> LatentKernelPrimaryDirection {
1475    LatentKernelPrimaryDirection {
1476        dq: direction.dq_right,
1477        dqd: 0.0,
1478        dmu: direction.dmu,
1479        dtau: direction.dlog_sigma,
1480    }
1481}
1482
1483fn latent_survival_row_primary_log_jet(
1484    quadctx: &QuadratureContext,
1485    row: &LatentSurvivalRow,
1486    q_entry: f64,
1487    q_exit: f64,
1488    qdot_exit: f64,
1489    q_right: f64,
1490    mu: f64,
1491    sigma: f64,
1492    log_sigma_factor: f64,
1493    directions: &[LatentSurvivalPrimaryDirection],
1494) -> Result<LatentMultiDirJet, String> {
1495    let entry_state = LatentKernelPrimaryState {
1496        q: q_entry,
1497        qdot: 1.0,
1498        mu,
1499        sigma,
1500        log_sigma_factor,
1501    };
1502    let entry_directions = directions
1503        .iter()
1504        .copied()
1505        .map(latent_survival_map_entry_direction)
1506        .collect::<Vec<_>>();
1507
1508    let denominator = latent_kernel_sum_log_jet(
1509        quadctx,
1510        &[LatentKernelPrimaryTerm {
1511            coeff: 1.0,
1512            q_exp: 0,
1513            qdot_power: 0,
1514            tau_exp: 0,
1515            k: 0,
1516        }],
1517        entry_state,
1518        &entry_directions,
1519        "latent survival denominator",
1520    )?;
1521
1522    // The numerator for right-censoring / exact events is a single-state log-sum
1523    // kernel at the exit mass. Interval censoring is the difference of two
1524    // single-state kernels at DIFFERENT masses (L at `q_exit`, R at `q_right`),
1525    // so it is assembled by `latent_survival_interval_numerator_log_jet` below.
1526    let numerator = match row.event_type {
1527        LatentSurvivalEventType::RightCensored | LatentSurvivalEventType::ExactEvent => {
1528            let exit_state = LatentKernelPrimaryState {
1529                q: q_exit,
1530                qdot: qdot_exit,
1531                mu,
1532                sigma,
1533                log_sigma_factor,
1534            };
1535            let exit_directions = directions
1536                .iter()
1537                .copied()
1538                .map(|dir| latent_survival_map_exit_direction(dir, row.event_type))
1539                .collect::<Vec<_>>();
1540            let numerator_terms = match row.event_type {
1541                LatentSurvivalEventType::RightCensored => vec![LatentKernelPrimaryTerm {
1542                    coeff: 1.0,
1543                    q_exp: 0,
1544                    qdot_power: 0,
1545                    tau_exp: 0,
1546                    k: 0,
1547                }],
1548                LatentSurvivalEventType::ExactEvent => {
1549                    let mut terms = Vec::new();
1550                    if row.hazard_unloaded > 0.0 {
1551                        terms.push(LatentKernelPrimaryTerm {
1552                            coeff: row.hazard_unloaded,
1553                            q_exp: 0,
1554                            qdot_power: 0,
1555                            tau_exp: 0,
1556                            k: 0,
1557                        });
1558                    }
1559                    terms.push(LatentKernelPrimaryTerm {
1560                        coeff: 1.0,
1561                        q_exp: 1,
1562                        qdot_power: 1,
1563                        tau_exp: 0,
1564                        k: 1,
1565                    });
1566                    terms
1567                }
1568                LatentSurvivalEventType::IntervalCensored => {
1569                    // Interval-censored rows are routed to the dedicated two-state
1570                    // numerator branch (the outer match arm below), so this inner
1571                    // arm is not reached; a clean error rather than a panic guards
1572                    // against a future routing change.
1573                    return Err(
1574                        "interval-censored row reached the single-state numerator branch; \
1575                         it must take the dedicated two-state branch"
1576                            .to_string(),
1577                    );
1578                }
1579            };
1580            latent_kernel_sum_log_jet(
1581                quadctx,
1582                &numerator_terms,
1583                exit_state,
1584                &exit_directions,
1585                "latent survival numerator",
1586            )?
1587        }
1588        LatentSurvivalEventType::IntervalCensored => latent_survival_interval_numerator_log_jet(
1589            quadctx,
1590            row,
1591            q_exit,
1592            q_right,
1593            mu,
1594            sigma,
1595            log_sigma_factor,
1596            directions,
1597        )?,
1598    };
1599
1600    let mut total = numerator.add(&denominator.scale(-1.0));
1601    // For interval rows the unloaded exit mass is folded into the per-boundary
1602    // coefficients `exp(-mass_unloaded_{left,right})` inside the two-state
1603    // numerator, so only the (constant) unloaded-entry term remains here; for
1604    // right-censoring / exact events the exit/entry unloaded masses are an
1605    // additive constant on the log-likelihood.
1606    match row.event_type {
1607        LatentSurvivalEventType::IntervalCensored => {
1608            total.coeffs[0] += row.mass_unloaded_entry;
1609        }
1610        _ => {
1611            total.coeffs[0] += -row.mass_unloaded_exit + row.mass_unloaded_entry;
1612        }
1613    }
1614    Ok(total)
1615}
1616
1617/// Interval-censored numerator jet `log[ c_L·K_{0,M_L} − c_R·K_{0,M_R} ]` where
1618/// `M_L = exp(q_exit)`, `M_R = exp(q_right)`, `c_L = exp(-mass_unloaded_left)`
1619/// and `c_R = exp(-mass_unloaded_right)`.
1620///
1621/// This is the dynamic-time analogue of the static
1622/// [`LatentSurvivalRowJet::interval_censored`] kernel: the interval likelihood
1623/// is the difference of two BOUNDARY survival masses, each a single-state
1624/// order-0 kernel, but at two DISTINCT cumulative masses. Because the two
1625/// boundaries respond to different time functionals (`q_exit` vs `q_right`) we
1626/// cannot fold them into one `latent_kernel_sum_log_jet` state. Instead we:
1627///   1. build each boundary's `log K_{0,M}` jet at its own state, with its own
1628///      direction map (left → `dq_exit`, right → `dq_right`; both share
1629///      `mu`/`sigma`),
1630///   2. lift each to the LINEAR domain via `exp` (a unary composition whose five
1631///      derivatives at value `v` are all `exp(v)`), scaled by its coefficient
1632///      `c_L` (resp. `−c_R`),
1633///   3. add the two linear-domain jets, and
1634///   4. drop back to the log domain via the same `log` unary composition the
1635///      single-state path uses.
1636/// Every multi-direction coefficient (value, score, neg-Hessian, 3rd, 4th)
1637/// follows by the Faà-di-Bruno composition already implemented in
1638/// `MultiDirJet::compose_unary`, so the derivative reductions are consistent
1639/// with the exact-event/right-censored branches by construction.
1640fn latent_survival_interval_numerator_log_jet(
1641    quadctx: &QuadratureContext,
1642    row: &LatentSurvivalRow,
1643    q_exit: f64,
1644    q_right: f64,
1645    mu: f64,
1646    sigma: f64,
1647    log_sigma_factor: f64,
1648    directions: &[LatentSurvivalPrimaryDirection],
1649) -> Result<LatentMultiDirJet, String> {
1650    let single_k0 = [LatentKernelPrimaryTerm {
1651        coeff: 1.0,
1652        q_exp: 0,
1653        qdot_power: 0,
1654        tau_exp: 0,
1655        k: 0,
1656    }];
1657
1658    let left_state = LatentKernelPrimaryState {
1659        q: q_exit,
1660        qdot: 1.0,
1661        mu,
1662        sigma,
1663        log_sigma_factor,
1664    };
1665    let right_state = LatentKernelPrimaryState {
1666        q: q_right,
1667        qdot: 1.0,
1668        mu,
1669        sigma,
1670        log_sigma_factor,
1671    };
1672    let left_directions = directions
1673        .iter()
1674        .copied()
1675        .map(latent_survival_map_left_direction)
1676        .collect::<Vec<_>>();
1677    let right_directions = directions
1678        .iter()
1679        .copied()
1680        .map(latent_survival_map_right_direction)
1681        .collect::<Vec<_>>();
1682
1683    let log_left = latent_kernel_sum_log_jet(
1684        quadctx,
1685        &single_k0,
1686        left_state,
1687        &left_directions,
1688        "latent survival interval left boundary",
1689    )?;
1690    let log_right = latent_kernel_sum_log_jet(
1691        quadctx,
1692        &single_k0,
1693        right_state,
1694        &right_directions,
1695        "latent survival interval right boundary",
1696    )?;
1697
1698    // Lift each boundary's log-kernel jet to the linear domain and scale by the
1699    // unloaded-mass prefactor. exp''''(v) = exp(v) for all orders, so the unary
1700    // derivative tower is `[exp(v); exp(v); exp(v); exp(v); exp(v)]`.
1701    let c_left = (-row.mass_unloaded_left).exp();
1702    let c_right = (-row.mass_unloaded_right).exp();
1703    let exp_left_value = log_left.coeff(0).exp();
1704    let exp_right_value = log_right.coeff(0).exp();
1705    let linear_left = log_left.compose_unary([exp_left_value; 5]).scale(c_left);
1706    let linear_right = log_right.compose_unary([exp_right_value; 5]).scale(c_right);
1707
1708    let linear_numerator = linear_left.add(&linear_right.scale(-1.0));
1709    let base = linear_numerator.coeff(0);
1710    if !(base.is_finite() && base > 0.0) {
1711        return Err(LatentSurvivalError::NumericalFailure {
1712            reason: format!(
1713                "latent survival interval numerator must be a positive survival-mass difference, \
1714                 got c_L*K0(M_L) - c_R*K0(M_R) = {base}; require M_L < M_R (i.e. L < R)"
1715            ),
1716        }
1717        .into());
1718    }
1719    // Drop back to the log domain. `latent_unary_derivatives_log(base)` is the
1720    // unary derivative tower of `ln` at the positive base value, so the composed
1721    // value channel is `ln(base)` and the higher coefficients are the
1722    // log-of-a-difference score / curvature, consistent with the single-state
1723    // log-sum path (which composes `ln` at its normalised base of 1).
1724    Ok(linear_numerator.compose_unary(latent_unary_derivatives_log(base)))
1725}
1726
1727fn latent_survival_row_primary_gradient_hessian(
1728    quadctx: &QuadratureContext,
1729    row: &LatentSurvivalRow,
1730    q_entry: f64,
1731    q_exit: f64,
1732    qdot_exit: f64,
1733    q_right: f64,
1734    mu: f64,
1735    sigma: f64,
1736    include_log_sigma: bool,
1737) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
1738    let log_sigma_factor = if sigma > 0.0 { sigma.ln() } else { 0.0 };
1739    let mut gradient = Array1::<f64>::zeros(LATENT_SURVIVAL_PRIMARY_DIM);
1740    let mut neg_hessian =
1741        Array2::<f64>::zeros((LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM));
1742    let active_primary = if include_log_sigma {
1743        LATENT_SURVIVAL_PRIMARY_DIM
1744    } else {
1745        LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
1746    };
1747    let log_lik = latent_survival_row_primary_log_jet(
1748        quadctx,
1749        row,
1750        q_entry,
1751        q_exit,
1752        qdot_exit,
1753        q_right,
1754        mu,
1755        sigma,
1756        log_sigma_factor,
1757        &[],
1758    )?
1759    .coeff(0);
1760    for a in 0..active_primary {
1761        let dir_a = latent_survival_basis_direction(a);
1762        gradient[a] = latent_survival_row_primary_log_jet(
1763            quadctx,
1764            row,
1765            q_entry,
1766            q_exit,
1767            qdot_exit,
1768            q_right,
1769            mu,
1770            sigma,
1771            log_sigma_factor,
1772            &[dir_a],
1773        )?
1774        .coeff(1);
1775        for b in a..active_primary {
1776            let coeff = latent_survival_row_primary_log_jet(
1777                quadctx,
1778                row,
1779                q_entry,
1780                q_exit,
1781                qdot_exit,
1782                q_right,
1783                mu,
1784                sigma,
1785                log_sigma_factor,
1786                &[dir_a, latent_survival_basis_direction(b)],
1787            )?
1788            .coeff(3);
1789            neg_hessian[[a, b]] = -coeff;
1790            neg_hessian[[b, a]] = -coeff;
1791        }
1792    }
1793    Ok((log_lik, gradient, neg_hessian))
1794}
1795
1796fn latent_survival_row_primary_third_contracted(
1797    quadctx: &QuadratureContext,
1798    row: &LatentSurvivalRow,
1799    q_entry: f64,
1800    q_exit: f64,
1801    qdot_exit: f64,
1802    q_right: f64,
1803    mu: f64,
1804    sigma: f64,
1805    direction: &Array1<f64>,
1806    include_log_sigma: bool,
1807) -> Result<Array2<f64>, String> {
1808    let log_sigma_factor = if sigma > 0.0 { sigma.ln() } else { 0.0 };
1809    let active_primary = if include_log_sigma {
1810        LATENT_SURVIVAL_PRIMARY_DIM
1811    } else {
1812        LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
1813    };
1814    let dir = LatentSurvivalPrimaryDirection {
1815        dq_entry: direction[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
1816        dq_exit: direction[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
1817        dqdot_exit: direction[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
1818        dq_right: direction[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
1819        dmu: direction[LATENT_SURVIVAL_PRIMARY_MU],
1820        dlog_sigma: direction[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
1821    };
1822    let mut out = Array2::<f64>::zeros((LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM));
1823    for a in 0..active_primary {
1824        let dir_a = latent_survival_basis_direction(a);
1825        for b in a..active_primary {
1826            let coeff = latent_survival_row_primary_log_jet(
1827                quadctx,
1828                row,
1829                q_entry,
1830                q_exit,
1831                qdot_exit,
1832                q_right,
1833                mu,
1834                sigma,
1835                log_sigma_factor,
1836                &[dir_a, latent_survival_basis_direction(b), dir],
1837            )?
1838            .coeff(7);
1839            out[[a, b]] = -coeff;
1840            out[[b, a]] = -coeff;
1841        }
1842    }
1843    Ok(out)
1844}
1845
1846fn latent_survival_row_primary_fourth_contracted(
1847    quadctx: &QuadratureContext,
1848    row: &LatentSurvivalRow,
1849    q_entry: f64,
1850    q_exit: f64,
1851    qdot_exit: f64,
1852    q_right: f64,
1853    mu: f64,
1854    sigma: f64,
1855    direction_u: &Array1<f64>,
1856    direction_v: &Array1<f64>,
1857    include_log_sigma: bool,
1858) -> Result<Array2<f64>, String> {
1859    let log_sigma_factor = if sigma > 0.0 { sigma.ln() } else { 0.0 };
1860    let active_primary = if include_log_sigma {
1861        LATENT_SURVIVAL_PRIMARY_DIM
1862    } else {
1863        LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
1864    };
1865    let dir_u = LatentSurvivalPrimaryDirection {
1866        dq_entry: direction_u[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
1867        dq_exit: direction_u[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
1868        dqdot_exit: direction_u[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
1869        dq_right: direction_u[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
1870        dmu: direction_u[LATENT_SURVIVAL_PRIMARY_MU],
1871        dlog_sigma: direction_u[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
1872    };
1873    let dir_v = LatentSurvivalPrimaryDirection {
1874        dq_entry: direction_v[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
1875        dq_exit: direction_v[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
1876        dqdot_exit: direction_v[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
1877        dq_right: direction_v[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
1878        dmu: direction_v[LATENT_SURVIVAL_PRIMARY_MU],
1879        dlog_sigma: direction_v[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
1880    };
1881    let mut out = Array2::<f64>::zeros((LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM));
1882    for a in 0..active_primary {
1883        let dir_a = latent_survival_basis_direction(a);
1884        for b in a..active_primary {
1885            let coeff = latent_survival_row_primary_log_jet(
1886                quadctx,
1887                row,
1888                q_entry,
1889                q_exit,
1890                qdot_exit,
1891                q_right,
1892                mu,
1893                sigma,
1894                log_sigma_factor,
1895                &[dir_a, latent_survival_basis_direction(b), dir_u, dir_v],
1896            )?
1897            .coeff(15);
1898            out[[a, b]] = -coeff;
1899            out[[b, a]] = -coeff;
1900        }
1901    }
1902    Ok(out)
1903}
1904
1905#[derive(Clone)]
1906struct LatentSurvivalJointSlices {
1907    time: std::ops::Range<usize>,
1908    mean: std::ops::Range<usize>,
1909    log_sigma: Option<std::ops::Range<usize>>,
1910    total: usize,
1911}
1912
1913#[derive(Clone)]
1914struct LatentSurvivalJointGradientAccum {
1915    ll: f64,
1916    gradient: Array1<f64>,
1917}
1918
1919#[derive(Clone)]
1920struct LatentSurvivalJointDenseAccum {
1921    ll: f64,
1922    gradient: Array1<f64>,
1923    hessian: Array2<f64>,
1924}
1925
1926#[derive(Clone)]
1927struct LatentSurvivalDenseHessianAccum {
1928    hessian: Array2<f64>,
1929}
1930
1931/// Process latent-survival rows in fixed contiguous chunks, using one
1932/// accumulator per rayon task and reducing those accumulators in chunk-index
1933/// order so gradient/Hessian assembly stays deterministic across runs.
1934fn deterministic_latent_survival_row_reduction<Acc, Init, Process, Combine>(
1935    n_rows: usize,
1936    init: Init,
1937    process_row: Process,
1938    mut combine: Combine,
1939) -> Result<Acc, String>
1940where
1941    Acc: Send,
1942    Init: Fn() -> Acc + Sync,
1943    Process: Fn(usize, &mut Acc) -> Result<(), String> + Sync,
1944    Combine: FnMut(&mut Acc, Acc),
1945{
1946    use rayon::iter::{IntoParallelIterator, ParallelIterator};
1947
1948    const TARGET_CHUNK_COUNT: usize = 32;
1949    if n_rows == 0 {
1950        return Ok(init());
1951    }
1952    let chunk_size = n_rows.div_ceil(TARGET_CHUNK_COUNT).max(1);
1953    let n_chunks = n_rows.div_ceil(chunk_size);
1954    let chunk_accumulators: Vec<Acc> = (0..n_chunks)
1955        .into_par_iter()
1956        .map(|chunk_idx| -> Result<Acc, String> {
1957            let start = chunk_idx * chunk_size;
1958            let end = (start + chunk_size).min(n_rows);
1959            let mut acc = init();
1960            for row_idx in start..end {
1961                process_row(row_idx, &mut acc)?;
1962            }
1963            Ok(acc)
1964        })
1965        .collect::<Result<Vec<_>, String>>()?;
1966
1967    let mut total = init();
1968    for acc in chunk_accumulators {
1969        combine(&mut total, acc);
1970    }
1971    Ok(total)
1972}
1973
1974impl LatentSurvivalFamily {
1975    /// Assemble the per-row [`LatentSurvivalRow`] for `row_idx` from the family's
1976    /// unloaded-mass/hazard fields and the supplied per-row time quantiles.
1977    ///
1978    /// Shared by every per-row reduction (log-likelihood, gradient, Hessian,
1979    /// directional third derivatives): each previously inlined an identical
1980    /// `event_type` lookup followed by the same 12-argument
1981    /// `build_latent_survival_row` call. Behavior is unchanged.
1982    fn build_row_at(
1983        &self,
1984        row_idx: usize,
1985        q_entry: f64,
1986        q_exit: f64,
1987        qdot_exit: f64,
1988        q_right: f64,
1989    ) -> Result<LatentSurvivalRow, LatentSurvivalError> {
1990        let event_type = latent_survival_event_type_for(self.event_target[row_idx]);
1991        build_latent_survival_row(
1992            row_idx,
1993            self.hazard_loading,
1994            event_type,
1995            q_entry,
1996            q_exit,
1997            qdot_exit,
1998            q_right,
1999            self.unloaded_mass_entry[row_idx],
2000            self.unloaded_mass_exit[row_idx],
2001            self.unloaded_mass_right[row_idx],
2002            self.unloaded_hazard_exit[row_idx],
2003        )
2004    }
2005
2006    fn joint_slices(&self) -> LatentSurvivalJointSlices {
2007        let p_time = self.x_time_exit.ncols();
2008        let p_mean = self.x_mean.ncols();
2009        let time = 0..p_time;
2010        let mean = p_time..p_time + p_mean;
2011        let log_sigma = self
2012            .latent_sd_fixed
2013            .is_none()
2014            .then_some((p_time + p_mean)..(p_time + p_mean + 1));
2015        LatentSurvivalJointSlices {
2016            total: log_sigma
2017                .as_ref()
2018                .map_or(p_time + p_mean, |range| range.end),
2019            time,
2020            mean,
2021            log_sigma,
2022        }
2023    }
2024
2025    fn row_primary_direction_from_flat(
2026        &self,
2027        row: usize,
2028        slices: &LatentSurvivalJointSlices,
2029        d_beta_flat: &Array1<f64>,
2030    ) -> Array1<f64> {
2031        let mut out = Array1::<f64>::zeros(LATENT_SURVIVAL_PRIMARY_DIM);
2032        let d_time = d_beta_flat.slice(s![slices.time.clone()]);
2033        out[LATENT_SURVIVAL_PRIMARY_Q_ENTRY] = self.x_time_entry.row(row).dot(&d_time);
2034        out[LATENT_SURVIVAL_PRIMARY_Q_EXIT] = self.x_time_exit.row(row).dot(&d_time);
2035        out[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT] = self.x_time_derivative_exit.row(row).dot(&d_time);
2036        out[LATENT_SURVIVAL_PRIMARY_Q_RIGHT] = self.x_time_right.row(row).dot(&d_time);
2037        out[LATENT_SURVIVAL_PRIMARY_MU] = self
2038            .x_mean
2039            .dot_row_view(row, d_beta_flat.slice(s![slices.mean.clone()]));
2040        if let Some(range) = &slices.log_sigma {
2041            out[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA] = d_beta_flat[range.start];
2042        }
2043        out
2044    }
2045
2046    fn joint_block_ranges(&self) -> Vec<std::ops::Range<usize>> {
2047        let slices = self.joint_slices();
2048        let mut ranges = vec![slices.time.clone(), slices.mean.clone()];
2049        if let Some(log_sigma) = slices.log_sigma {
2050            ranges.push(log_sigma);
2051        }
2052        ranges
2053    }
2054
2055    fn add_pullback_primary_gradient(
2056        &self,
2057        target: &mut Array1<f64>,
2058        row: usize,
2059        slices: &LatentSurvivalJointSlices,
2060        primary_gradient: &Array1<f64>,
2061        weight: f64,
2062    ) -> Result<(), String> {
2063        for (primary_idx, time_vec) in [
2064            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
2065            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
2066            (
2067                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2068                self.x_time_derivative_exit.row(row),
2069            ),
2070            (LATENT_SURVIVAL_PRIMARY_Q_RIGHT, self.x_time_right.row(row)),
2071        ] {
2072            let scale = weight * primary_gradient[primary_idx];
2073            if scale == 0.0 {
2074                continue;
2075            }
2076            for i in 0..time_vec.len() {
2077                let xi = time_vec[i];
2078                if xi != 0.0 {
2079                    target[slices.time.start + i] += scale * xi;
2080                }
2081            }
2082        }
2083
2084        let mean_scale = weight * primary_gradient[LATENT_SURVIVAL_PRIMARY_MU];
2085        if mean_scale != 0.0 {
2086            self.x_mean
2087                .axpy_row_into(
2088                    row,
2089                    mean_scale,
2090                    &mut target.slice_mut(s![slices.mean.clone()]),
2091                )
2092                .map_err(|error| {
2093                    format!(
2094                        "latent survival mean gradient pullback dimension mismatch: row={row}, mean_slice={:?}, target_len={}, x_mean_cols={}, error={error}",
2095                        slices.mean,
2096                        target.len(),
2097                        self.x_mean.ncols()
2098                    )
2099                })?;
2100        }
2101
2102        if let Some(log_sigma) = &slices.log_sigma {
2103            target[log_sigma.start] += weight * primary_gradient[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA];
2104        }
2105        Ok(())
2106    }
2107
2108    fn add_pullback_primary_hessian(
2109        &self,
2110        target: &mut Array2<f64>,
2111        row: usize,
2112        slices: &LatentSurvivalJointSlices,
2113        primary_hessian: &Array2<f64>,
2114    ) -> Result<(), String> {
2115        let time_weights = [
2116            primary_hessian[[
2117                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2118                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2119            ]],
2120            primary_hessian[[
2121                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2122                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2123            ]],
2124            primary_hessian[[
2125                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2126                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2127            ]],
2128            primary_hessian[[
2129                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2130                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2131            ]],
2132        ];
2133        let time_cross_weights = [
2134            (
2135                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2136                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2137                &self.x_time_entry,
2138                &self.x_time_exit,
2139            ),
2140            (
2141                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2142                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2143                &self.x_time_entry,
2144                &self.x_time_derivative_exit,
2145            ),
2146            (
2147                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2148                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2149                &self.x_time_exit,
2150                &self.x_time_derivative_exit,
2151            ),
2152            (
2153                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2154                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2155                &self.x_time_entry,
2156                &self.x_time_right,
2157            ),
2158            (
2159                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2160                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2161                &self.x_time_exit,
2162                &self.x_time_right,
2163            ),
2164            (
2165                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2166                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2167                &self.x_time_derivative_exit,
2168                &self.x_time_right,
2169            ),
2170        ];
2171        {
2172            let time_target = &mut target.slice_mut(s![slices.time.clone(), slices.time.clone()]);
2173            dense_outer_accumulate(time_target, time_weights[0], self.x_time_entry.row(row));
2174            dense_outer_accumulate(time_target, time_weights[1], self.x_time_exit.row(row));
2175            dense_outer_accumulate(
2176                time_target,
2177                time_weights[2],
2178                self.x_time_derivative_exit.row(row),
2179            );
2180            dense_outer_accumulate(time_target, time_weights[3], self.x_time_right.row(row));
2181            for (a, b, lhs, rhs) in time_cross_weights {
2182                let weight = primary_hessian[[a, b]];
2183                if weight == 0.0 {
2184                    continue;
2185                }
2186                dense_symmetric_cross_accumulate(time_target, weight, lhs.row(row), rhs.row(row));
2187            }
2188        }
2189
2190        let mean_weight = primary_hessian[[LATENT_SURVIVAL_PRIMARY_MU, LATENT_SURVIVAL_PRIMARY_MU]];
2191        self.x_mean
2192            .syr_row_into_view(
2193                row,
2194                mean_weight,
2195                target.slice_mut(s![slices.mean.clone(), slices.mean.clone()]),
2196            )
2197            .map_err(|error| {
2198                format!(
2199                    "latent survival mean Hessian pullback dimension mismatch: row={row}, mean_slice={:?}, target_dim={:?}, x_mean_cols={}, error={error}",
2200                    slices.mean,
2201                    target.dim(),
2202                    self.x_mean.ncols()
2203                )
2204            })?;
2205
2206        let mean_row = self
2207            .x_mean
2208            .try_row_chunk(row..row + 1)
2209            .map_err(|error| {
2210                format!(
2211                    "latent survival mean pullback row chunk failed: row={row}, x_mean_rows={}, x_mean_cols={}, error={error}",
2212                    self.x_mean.nrows(),
2213                    self.x_mean.ncols()
2214                )
2215            })?;
2216        let mean_vec = mean_row.row(0);
2217        let time_mean_weights = [
2218            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
2219            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
2220            (
2221                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2222                self.x_time_derivative_exit.row(row),
2223            ),
2224            (LATENT_SURVIVAL_PRIMARY_Q_RIGHT, self.x_time_right.row(row)),
2225        ];
2226        for (primary_idx, time_vec) in time_mean_weights {
2227            let weight = primary_hessian[[primary_idx, LATENT_SURVIVAL_PRIMARY_MU]];
2228            if weight == 0.0 {
2229                continue;
2230            }
2231            for i in 0..time_vec.len() {
2232                let xi = time_vec[i];
2233                if xi == 0.0 {
2234                    continue;
2235                }
2236                for j in 0..mean_vec.len() {
2237                    let xj = mean_vec[j];
2238                    if xj == 0.0 {
2239                        continue;
2240                    }
2241                    target[[slices.time.start + i, slices.mean.start + j]] += weight * xi * xj;
2242                    target[[slices.mean.start + j, slices.time.start + i]] += weight * xj * xi;
2243                }
2244            }
2245        }
2246
2247        if let Some(log_sigma) = &slices.log_sigma {
2248            let sigma_idx = log_sigma.start;
2249            target[[sigma_idx, sigma_idx]] += primary_hessian[[
2250                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
2251                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
2252            ]];
2253
2254            for (primary_idx, time_vec) in [
2255                (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
2256                (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
2257                (
2258                    LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2259                    self.x_time_derivative_exit.row(row),
2260                ),
2261                (LATENT_SURVIVAL_PRIMARY_Q_RIGHT, self.x_time_right.row(row)),
2262            ] {
2263                let weight = primary_hessian[[primary_idx, LATENT_SURVIVAL_PRIMARY_LOG_SIGMA]];
2264                if weight == 0.0 {
2265                    continue;
2266                }
2267                for i in 0..time_vec.len() {
2268                    let xi = time_vec[i];
2269                    if xi == 0.0 {
2270                        continue;
2271                    }
2272                    target[[slices.time.start + i, sigma_idx]] += weight * xi;
2273                    target[[sigma_idx, slices.time.start + i]] += weight * xi;
2274                }
2275            }
2276
2277            let mean_sigma_weight = primary_hessian[[
2278                LATENT_SURVIVAL_PRIMARY_MU,
2279                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
2280            ]];
2281            if mean_sigma_weight != 0.0 {
2282                for j in 0..mean_vec.len() {
2283                    let xj = mean_vec[j];
2284                    if xj == 0.0 {
2285                        continue;
2286                    }
2287                    target[[slices.mean.start + j, sigma_idx]] += mean_sigma_weight * xj;
2288                    target[[sigma_idx, slices.mean.start + j]] += mean_sigma_weight * xj;
2289                }
2290            }
2291        }
2292        Ok(())
2293    }
2294
2295    fn evaluate_exact_newton_joint_gradient_dense(
2296        &self,
2297        block_states: &[ParameterBlockState],
2298    ) -> Result<(f64, Array1<f64>), String> {
2299        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
2300        let q_right = self.time_q_right(block_states)?;
2301        let sigma = self.latent_sd(block_states)?;
2302        let slices = self.joint_slices();
2303        let include_log_sigma = slices.log_sigma.is_some();
2304        let total = slices.total;
2305        let acc = deterministic_latent_survival_row_reduction(
2306            self.event_target.len(),
2307            || LatentSurvivalJointGradientAccum {
2308                ll: 0.0,
2309                gradient: Array1::<f64>::zeros(total),
2310            },
2311            |row_idx, acc| {
2312                let wi = self.weights[row_idx];
2313                if wi <= MIN_WEIGHT {
2314                    return Ok(());
2315                }
2316                let row = self.build_row_at(
2317                    row_idx,
2318                    q_entry[row_idx],
2319                    q_exit[row_idx],
2320                    qdot_exit[row_idx],
2321                    q_right[row_idx],
2322                )?;
2323                let (row_ll, primary_gradient, _) = latent_survival_row_primary_gradient_hessian(
2324                    &self.quadctx,
2325                    &row,
2326                    q_entry[row_idx],
2327                    q_exit[row_idx],
2328                    qdot_exit[row_idx],
2329                    q_right[row_idx],
2330                    mu[row_idx],
2331                    sigma,
2332                    include_log_sigma,
2333                )?;
2334                acc.ll += wi * row_ll;
2335                self.add_pullback_primary_gradient(
2336                    &mut acc.gradient,
2337                    row_idx,
2338                    &slices,
2339                    &primary_gradient,
2340                    wi,
2341                )?;
2342                Ok(())
2343            },
2344            |total_acc, chunk_acc| {
2345                total_acc.ll += chunk_acc.ll;
2346                total_acc.gradient += &chunk_acc.gradient;
2347            },
2348        )?;
2349        Ok((acc.ll, acc.gradient))
2350    }
2351
2352    /// Per-row residuals of the unpenalized NLL with respect to the three
2353    /// additive baseline time-block offsets `(entry, exit, derivative)`.
2354    ///
2355    /// The baseline configuration θ enters the latent-survival working model
2356    /// only through the additive offsets on the three time channels
2357    ///   q_entry = x_time_entry·β_time + o_E(θ),
2358    ///   q_exit  = x_time_exit·β_time  + o_X(θ),
2359    ///   q̇_exit = x_time_deriv·β_time + o_D(θ),
2360    /// exactly the offset channel the transformation path carries through
2361    /// [`WorkingModelSurvival::offset_channel_residuals`]. Because
2362    /// `∂q_ch/∂o_ch = 1`, the residual `∂NLL/∂o_ch_i` equals
2363    /// `−∂(log-likelihood)/∂q_ch_i`, and the per-row primary log-likelihood
2364    /// gradient over `(q_entry, q_exit, q̇_exit)` is precisely the
2365    /// `Q_ENTRY`/`Q_EXIT`/`QDOT_EXIT` components returned by
2366    /// [`latent_survival_row_primary_gradient_hessian`]. Sampleweight-scaled to
2367    /// match the [`OffsetChannelResiduals`] contract consumed by
2368    /// `baseline_chain_rule_gradient`.
2369    ///
2370    /// At the converged (constrained) β̂ the envelope theorem makes this the
2371    /// exact θ-gradient of the profile penalized NLL `0.5·deviance + 0.5·βᵀSβ`.
2372    /// The interval upper-bound `q_right = x_time_right·β_time + o_R(θ)` channel
2373    /// DOES carry its own baseline-θ offset `o_R(θ)` (the time basis evaluated at
2374    /// the bracket upper bound `R`), distinct from the exit offset at `L`, so its
2375    /// residual `−∂(log-likelihood)/∂q_right` is returned in the dedicated
2376    /// [`OffsetChannelResiduals::right`] channel; it is exactly 0 on every
2377    /// non-interval row (the `Q_RIGHT` primary channel is inert there) and the
2378    /// baseline-θ chain rule contracts it against the `age_right`-evaluated
2379    /// η-partial.
2380    pub fn offset_channel_residuals(
2381        &self,
2382        block_states: &[ParameterBlockState],
2383    ) -> Result<crate::survival::OffsetChannelResiduals, String> {
2384        let n = self.event_target.len();
2385        if block_states.is_empty() {
2386            // Degraded-fit fallback mirroring the location-scale family: an
2387            // empty block-state slate (ARC deterministic-replay stall) yields
2388            // zero residuals so the outer baseline-θ BFGS sees ‖g‖ = 0 and
2389            // terminates cleanly at the current θ̂ rather than panicking.
2390            log::warn!(
2391                "LatentSurvivalFamily::offset_channel_residuals: block_states is empty \
2392                 (degraded fit); returning zero offset residuals (n={n})"
2393            );
2394            return Ok(crate::survival::OffsetChannelResiduals {
2395                exit: Array1::<f64>::zeros(n),
2396                entry: Array1::<f64>::zeros(n),
2397                derivative: Array1::<f64>::zeros(n),
2398                right: Array1::<f64>::zeros(n),
2399            });
2400        }
2401        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
2402        let q_right = self.time_q_right(block_states)?;
2403        let sigma = self.latent_sd(block_states)?;
2404        let include_log_sigma = self.joint_slices().log_sigma.is_some();
2405        let mut entry = Array1::<f64>::zeros(n);
2406        let mut exit = Array1::<f64>::zeros(n);
2407        let mut derivative = Array1::<f64>::zeros(n);
2408        let mut right = Array1::<f64>::zeros(n);
2409        for row_idx in 0..n {
2410            let wi = self.weights[row_idx];
2411            if wi <= MIN_WEIGHT {
2412                continue;
2413            }
2414            let row = self.build_row_at(
2415                row_idx,
2416                q_entry[row_idx],
2417                q_exit[row_idx],
2418                qdot_exit[row_idx],
2419                q_right[row_idx],
2420            )?;
2421            let (_, primary_gradient, _) = latent_survival_row_primary_gradient_hessian(
2422                &self.quadctx,
2423                &row,
2424                q_entry[row_idx],
2425                q_exit[row_idx],
2426                qdot_exit[row_idx],
2427                q_right[row_idx],
2428                mu[row_idx],
2429                sigma,
2430                include_log_sigma,
2431            )?;
2432            // ∂NLL/∂o_ch = −w · ∂(log-likelihood)/∂q_ch.
2433            entry[row_idx] = -wi * primary_gradient[LATENT_SURVIVAL_PRIMARY_Q_ENTRY];
2434            exit[row_idx] = -wi * primary_gradient[LATENT_SURVIVAL_PRIMARY_Q_EXIT];
2435            derivative[row_idx] = -wi * primary_gradient[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT];
2436            // Interval upper-bound (`R`) channel. `q_right` shares the time-block
2437            // coefficients but carries its OWN baseline-θ η-offset evaluated at
2438            // `R` (`o_R(θ)`), so the profile-NLL θ-gradient must include it.
2439            // `∂(log-likelihood)/∂q_right` is exactly 0 for non-interval rows
2440            // (the `Q_RIGHT` channel is inert there), so this is 0 except on
2441            // interval-censored rows.
2442            right[row_idx] = -wi * primary_gradient[LATENT_SURVIVAL_PRIMARY_Q_RIGHT];
2443        }
2444        Ok(crate::survival::OffsetChannelResiduals {
2445            exit,
2446            entry,
2447            derivative,
2448            right,
2449        })
2450    }
2451
2452    /// Block-diagonal-only pullback: writes only time-time, mean-mean, and
2453    /// log_sigma-log_sigma rowwise contributions into per-block targets.
2454    /// Used by `evaluate()` to populate per-block working sets without ever
2455    /// materializing the cross blocks the inner solver does not consume.
2456    fn add_pullback_primary_block_diagonals(
2457        &self,
2458        row: usize,
2459        primary_hessian: &Array2<f64>,
2460        time_target: &mut Array2<f64>,
2461        mean_target: &mut Array2<f64>,
2462        log_sigma_target: Option<&mut Array2<f64>>,
2463    ) -> Result<(), String> {
2464        let h = primary_hessian;
2465        // Time block: 4 squared rows (entry/exit/qdot/right) + 6 symmetric
2466        // crosses. The interval right-boundary functional `q_right` shares the
2467        // time-block coefficients, so it accumulates into the same time target.
2468        dense_outer_accumulate(
2469            time_target,
2470            h[[
2471                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2472                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2473            ]],
2474            self.x_time_entry.row(row),
2475        );
2476        dense_outer_accumulate(
2477            time_target,
2478            h[[
2479                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2480                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2481            ]],
2482            self.x_time_exit.row(row),
2483        );
2484        dense_outer_accumulate(
2485            time_target,
2486            h[[
2487                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2488                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2489            ]],
2490            self.x_time_derivative_exit.row(row),
2491        );
2492        dense_outer_accumulate(
2493            time_target,
2494            h[[
2495                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2496                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2497            ]],
2498            self.x_time_right.row(row),
2499        );
2500        for (a, b, lhs, rhs) in [
2501            (
2502                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2503                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2504                &self.x_time_entry,
2505                &self.x_time_exit,
2506            ),
2507            (
2508                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2509                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2510                &self.x_time_entry,
2511                &self.x_time_derivative_exit,
2512            ),
2513            (
2514                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2515                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2516                &self.x_time_exit,
2517                &self.x_time_derivative_exit,
2518            ),
2519            (
2520                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
2521                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2522                &self.x_time_entry,
2523                &self.x_time_right,
2524            ),
2525            (
2526                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
2527                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2528                &self.x_time_exit,
2529                &self.x_time_right,
2530            ),
2531            (
2532                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
2533                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
2534                &self.x_time_derivative_exit,
2535                &self.x_time_right,
2536            ),
2537        ] {
2538            let weight = h[[a, b]];
2539            if weight == 0.0 {
2540                continue;
2541            }
2542            dense_symmetric_cross_accumulate(time_target, weight, lhs.row(row), rhs.row(row));
2543        }
2544        // Mean block.
2545        let mean_weight = h[[LATENT_SURVIVAL_PRIMARY_MU, LATENT_SURVIVAL_PRIMARY_MU]];
2546        self.x_mean
2547            .syr_row_into_view(row, mean_weight, mean_target.view_mut())
2548            .map_err(|error| {
2549                format!(
2550                    "latent survival mean block-diagonal pullback dimension mismatch: row={row}, mean_target_dim={:?}, x_mean_cols={}, error={error}",
2551                    mean_target.dim(),
2552                    self.x_mean.ncols()
2553                )
2554            })?;
2555        // Log-σ block (scalar).
2556        if let Some(target) = log_sigma_target {
2557            target[[0, 0]] += h[[
2558                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
2559                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
2560            ]];
2561        }
2562        Ok(())
2563    }
2564
2565    /// Block-diagonal evaluator used by `evaluate()`. Returns the per-row
2566    /// log-likelihood, the joint gradient (sliced into block gradients by
2567    /// the caller), and the three per-block diagonal Hessians without ever
2568    /// materializing the full joint matrix.
2569    fn evaluate_exact_newton_block_diagonals(
2570        &self,
2571        block_states: &[ParameterBlockState],
2572    ) -> Result<
2573        (
2574            f64,
2575            Array1<f64>,
2576            Array2<f64>,
2577            Array2<f64>,
2578            Option<Array2<f64>>,
2579        ),
2580        String,
2581    > {
2582        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
2583        let q_right = self.time_q_right(block_states)?;
2584        let sigma = self.latent_sd(block_states)?;
2585        let slices = self.joint_slices();
2586        let include_log_sigma = slices.log_sigma.is_some();
2587        let mut ll = 0.0;
2588        let mut gradient = Array1::<f64>::zeros(slices.total);
2589        let p_time = slices.time.len();
2590        let p_mean = slices.mean.len();
2591        let mut hess_time = Array2::<f64>::zeros((p_time, p_time));
2592        let mut hess_mean = Array2::<f64>::zeros((p_mean, p_mean));
2593        let mut hess_log_sigma = if include_log_sigma {
2594            Some(Array2::<f64>::zeros((1, 1)))
2595        } else {
2596            None
2597        };
2598        for row_idx in 0..self.event_target.len() {
2599            let wi = self.weights[row_idx];
2600            if wi <= MIN_WEIGHT {
2601                continue;
2602            }
2603            let row = self.build_row_at(
2604                row_idx,
2605                q_entry[row_idx],
2606                q_exit[row_idx],
2607                qdot_exit[row_idx],
2608                q_right[row_idx],
2609            )?;
2610            let (row_ll, primary_gradient, primary_hessian) =
2611                latent_survival_row_primary_gradient_hessian(
2612                    &self.quadctx,
2613                    &row,
2614                    q_entry[row_idx],
2615                    q_exit[row_idx],
2616                    qdot_exit[row_idx],
2617                    q_right[row_idx],
2618                    mu[row_idx],
2619                    sigma,
2620                    include_log_sigma,
2621                )?;
2622            ll += wi * row_ll;
2623            self.add_pullback_primary_gradient(
2624                &mut gradient,
2625                row_idx,
2626                &slices,
2627                &primary_gradient,
2628                wi,
2629            )?;
2630            self.add_pullback_primary_block_diagonals(
2631                row_idx,
2632                &(wi * primary_hessian),
2633                &mut hess_time,
2634                &mut hess_mean,
2635                hess_log_sigma.as_mut(),
2636            )?;
2637        }
2638        Ok((ll, gradient, hess_time, hess_mean, hess_log_sigma))
2639    }
2640
2641    fn evaluate_exact_newton_joint_dense(
2642        &self,
2643        block_states: &[ParameterBlockState],
2644    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
2645        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
2646        let q_right = self.time_q_right(block_states)?;
2647        let sigma = self.latent_sd(block_states)?;
2648        let slices = self.joint_slices();
2649        let include_log_sigma = slices.log_sigma.is_some();
2650        let total = slices.total;
2651        let acc = deterministic_latent_survival_row_reduction(
2652            self.event_target.len(),
2653            || LatentSurvivalJointDenseAccum {
2654                ll: 0.0,
2655                gradient: Array1::<f64>::zeros(total),
2656                hessian: Array2::<f64>::zeros((total, total)),
2657            },
2658            |row_idx, acc| {
2659                let wi = self.weights[row_idx];
2660                if wi <= MIN_WEIGHT {
2661                    return Ok(());
2662                }
2663                let row = self.build_row_at(
2664                    row_idx,
2665                    q_entry[row_idx],
2666                    q_exit[row_idx],
2667                    qdot_exit[row_idx],
2668                    q_right[row_idx],
2669                )?;
2670                let (row_ll, primary_gradient, primary_hessian) =
2671                    latent_survival_row_primary_gradient_hessian(
2672                        &self.quadctx,
2673                        &row,
2674                        q_entry[row_idx],
2675                        q_exit[row_idx],
2676                        qdot_exit[row_idx],
2677                        q_right[row_idx],
2678                        mu[row_idx],
2679                        sigma,
2680                        include_log_sigma,
2681                    )?;
2682                acc.ll += wi * row_ll;
2683                self.add_pullback_primary_gradient(
2684                    &mut acc.gradient,
2685                    row_idx,
2686                    &slices,
2687                    &primary_gradient,
2688                    wi,
2689                )?;
2690                self.add_pullback_primary_hessian(
2691                    &mut acc.hessian,
2692                    row_idx,
2693                    &slices,
2694                    &(wi * primary_hessian),
2695                )?;
2696                Ok(())
2697            },
2698            |total_acc, chunk_acc| {
2699                total_acc.ll += chunk_acc.ll;
2700                total_acc.gradient += &chunk_acc.gradient;
2701                total_acc.hessian += &chunk_acc.hessian;
2702            },
2703        )?;
2704        Ok((acc.ll, acc.gradient, acc.hessian))
2705    }
2706
2707    fn exact_newton_joint_hessian_directional_derivative_dense(
2708        &self,
2709        block_states: &[ParameterBlockState],
2710        d_beta_flat: &Array1<f64>,
2711    ) -> Result<Array2<f64>, String> {
2712        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
2713        let q_right = self.time_q_right(block_states)?;
2714        let sigma = self.latent_sd(block_states)?;
2715        let slices = self.joint_slices();
2716        if d_beta_flat.len() != slices.total {
2717            return Err(format!(
2718                "latent survival joint dH direction length mismatch: got {}, expected {}",
2719                d_beta_flat.len(),
2720                slices.total
2721            ));
2722        }
2723        let include_log_sigma = slices.log_sigma.is_some();
2724        let total = slices.total;
2725        let acc = deterministic_latent_survival_row_reduction(
2726            self.event_target.len(),
2727            || LatentSurvivalDenseHessianAccum {
2728                hessian: Array2::<f64>::zeros((total, total)),
2729            },
2730            |row_idx, acc| {
2731                let wi = self.weights[row_idx];
2732                if wi <= MIN_WEIGHT {
2733                    return Ok(());
2734                }
2735                let row = self.build_row_at(
2736                    row_idx,
2737                    q_entry[row_idx],
2738                    q_exit[row_idx],
2739                    qdot_exit[row_idx],
2740                    q_right[row_idx],
2741                )?;
2742                let direction = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_flat);
2743                let third = latent_survival_row_primary_third_contracted(
2744                    &self.quadctx,
2745                    &row,
2746                    q_entry[row_idx],
2747                    q_exit[row_idx],
2748                    qdot_exit[row_idx],
2749                    q_right[row_idx],
2750                    mu[row_idx],
2751                    sigma,
2752                    &direction,
2753                    include_log_sigma,
2754                )?;
2755                self.add_pullback_primary_hessian(
2756                    &mut acc.hessian,
2757                    row_idx,
2758                    &slices,
2759                    &(wi * third),
2760                )?;
2761                Ok(())
2762            },
2763            |total_acc, chunk_acc| {
2764                total_acc.hessian += &chunk_acc.hessian;
2765            },
2766        )?;
2767        Ok(acc.hessian)
2768    }
2769
2770    fn exact_newton_joint_hessian_second_directional_derivative_dense(
2771        &self,
2772        block_states: &[ParameterBlockState],
2773        d_beta_u_flat: &Array1<f64>,
2774        d_beta_v_flat: &Array1<f64>,
2775    ) -> Result<Array2<f64>, String> {
2776        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
2777        let q_right = self.time_q_right(block_states)?;
2778        let sigma = self.latent_sd(block_states)?;
2779        let slices = self.joint_slices();
2780        if d_beta_u_flat.len() != slices.total || d_beta_v_flat.len() != slices.total {
2781            return Err(format!(
2782                "latent survival joint d2H direction length mismatch: got {} and {}, expected {}",
2783                d_beta_u_flat.len(),
2784                d_beta_v_flat.len(),
2785                slices.total
2786            ));
2787        }
2788        let include_log_sigma = slices.log_sigma.is_some();
2789        let total = slices.total;
2790        let acc = deterministic_latent_survival_row_reduction(
2791            self.event_target.len(),
2792            || LatentSurvivalDenseHessianAccum {
2793                hessian: Array2::<f64>::zeros((total, total)),
2794            },
2795            |row_idx, acc| {
2796                let wi = self.weights[row_idx];
2797                if wi <= MIN_WEIGHT {
2798                    return Ok(());
2799                }
2800                let row = self.build_row_at(
2801                    row_idx,
2802                    q_entry[row_idx],
2803                    q_exit[row_idx],
2804                    qdot_exit[row_idx],
2805                    q_right[row_idx],
2806                )?;
2807                let direction_u =
2808                    self.row_primary_direction_from_flat(row_idx, &slices, d_beta_u_flat);
2809                let direction_v =
2810                    self.row_primary_direction_from_flat(row_idx, &slices, d_beta_v_flat);
2811                let fourth = latent_survival_row_primary_fourth_contracted(
2812                    &self.quadctx,
2813                    &row,
2814                    q_entry[row_idx],
2815                    q_exit[row_idx],
2816                    qdot_exit[row_idx],
2817                    q_right[row_idx],
2818                    mu[row_idx],
2819                    sigma,
2820                    &direction_u,
2821                    &direction_v,
2822                    include_log_sigma,
2823                )?;
2824                self.add_pullback_primary_hessian(
2825                    &mut acc.hessian,
2826                    row_idx,
2827                    &slices,
2828                    &(wi * fourth),
2829                )?;
2830                Ok(())
2831            },
2832            |total_acc, chunk_acc| {
2833                total_acc.hessian += &chunk_acc.hessian;
2834            },
2835        )?;
2836        Ok(acc.hessian)
2837    }
2838}
2839
2840fn log_kernel_ratio(
2841    bundle: &crate::survival::lognormal_kernel::LogLognormalKernelBundle,
2842    num: usize,
2843    den: usize,
2844) -> f64 {
2845    let delta = bundle.get(num) - bundle.get(den);
2846    if delta.is_finite() {
2847        delta.exp()
2848    } else if delta > 0.0 {
2849        f64::INFINITY
2850    } else {
2851        0.0
2852    }
2853}
2854
2855fn logk_q_derivatives(
2856    quadctx: &QuadratureContext,
2857    k: usize,
2858    mass: f64,
2859    mu: f64,
2860    sigma: f64,
2861) -> Result<(f64, f64, IntegratedExpectationMode), LatentSurvivalError> {
2862    if mass <= 0.0 {
2863        return Ok((0.0, 0.0, IntegratedExpectationMode::ExactClosedForm));
2864    }
2865    let bundle = log_kernel_bundle(quadctx, mass, mu, sigma, k + 2).map_err(|e| {
2866        LatentSurvivalError::NumericalFailure {
2867            reason: format!("latent survival kernel evaluation failed: {e}"),
2868        }
2869    })?;
2870    let r1 = log_kernel_ratio(&bundle, k + 1, k);
2871    let r2 = log_kernel_ratio(&bundle, k + 2, k);
2872    let d1 = -mass * r1;
2873    let d2 = d1 + mass * mass * (r2 - r1 * r1);
2874    Ok((d1, d2, bundle.mode))
2875}
2876
2877fn latent_survival_time_jet(
2878    quadctx: &QuadratureContext,
2879    row: &LatentSurvivalRow,
2880    qdot_exit: f64,
2881    mu: f64,
2882    sigma: f64,
2883) -> Result<LatentSurvivalTimeJet, LatentSurvivalError> {
2884    let (entry_d1, entry_d2, _) = logk_q_derivatives(quadctx, 0, row.mass_entry, mu, sigma)?;
2885    match row.event_type {
2886        LatentSurvivalEventType::RightCensored => {
2887            let (exit_d1, exit_d2, _) = logk_q_derivatives(quadctx, 0, row.mass_exit, mu, sigma)?;
2888            Ok(LatentSurvivalTimeJet {
2889                grad_entry: -entry_d1,
2890                grad_exit: exit_d1,
2891                neg_hess_entry: entry_d2,
2892                neg_hess_exit: -exit_d2,
2893            })
2894        }
2895        LatentSurvivalEventType::ExactEvent => {
2896            if !(qdot_exit.is_finite() && qdot_exit > 0.0) {
2897                return Err(LatentSurvivalError::NumericalFailure {
2898                    reason: format!(
2899                        "latent survival requires positive finite baseline hazard derivative, got {qdot_exit}"
2900                    ),
2901                });
2902            }
2903            if row.hazard_unloaded > 0.0 {
2904                let bundle =
2905                    log_kernel_bundle(quadctx, row.mass_exit, mu, sigma, 3).map_err(|e| {
2906                        LatentSurvivalError::NumericalFailure {
2907                            reason: format!("latent survival kernel evaluation failed: {e}"),
2908                        }
2909                    })?;
2910                let (unloaded_d1, unloaded_d2, _) =
2911                    logk_q_derivatives(quadctx, 0, row.mass_exit, mu, sigma)?;
2912                let (loaded_log_d1, loaded_d2, _) =
2913                    logk_q_derivatives(quadctx, 1, row.mass_exit, mu, sigma)?;
2914                let loaded_d1 = 1.0 + loaded_log_d1;
2915                let log_loaded = row.hazard_loaded.ln() + bundle.get(1);
2916                let log_unloaded = row.hazard_unloaded.ln() + bundle.get(0);
2917                let shift = log_loaded.max(log_unloaded);
2918                let loaded_weight = (log_loaded - shift).exp();
2919                let unloaded_weight = (log_unloaded - shift).exp();
2920                let normalizer = loaded_weight + unloaded_weight;
2921                if !(normalizer.is_finite() && normalizer > 0.0) {
2922                    return Err(LatentSurvivalError::NumericalFailure {
2923                        reason: "latent survival exact-event numerator became non-finite under loaded/unloaded hazard decomposition"
2924                            .to_string(),
2925                    });
2926                }
2927                let w_loaded = loaded_weight / normalizer;
2928                let w_unloaded = unloaded_weight / normalizer;
2929                let grad_exit = w_loaded * loaded_d1 + w_unloaded * unloaded_d1;
2930                let d2_exit = w_loaded * (loaded_d2 + loaded_d1 * loaded_d1)
2931                    + w_unloaded * (unloaded_d2 + unloaded_d1 * unloaded_d1)
2932                    - grad_exit * grad_exit;
2933                Ok(LatentSurvivalTimeJet {
2934                    grad_entry: -entry_d1,
2935                    grad_exit,
2936                    neg_hess_entry: entry_d2,
2937                    neg_hess_exit: -d2_exit,
2938                })
2939            } else {
2940                let (exit_d1, exit_d2, _) =
2941                    logk_q_derivatives(quadctx, 1, row.mass_exit, mu, sigma)?;
2942                Ok(LatentSurvivalTimeJet {
2943                    grad_entry: -entry_d1,
2944                    grad_exit: 1.0 + exit_d1,
2945                    neg_hess_entry: entry_d2,
2946                    neg_hess_exit: -exit_d2,
2947                })
2948            }
2949        }
2950        LatentSurvivalEventType::IntervalCensored => {
2951            Err(LatentSurvivalError::UnsupportedConfiguration {
2952                reason:
2953                    "latent survival dynamic time derivatives do not implement interval censoring"
2954                        .to_string(),
2955            })
2956        }
2957    }
2958}
2959
2960fn dense_outer_accumulate<S>(
2961    target: &mut ndarray::ArrayBase<S, ndarray::Ix2>,
2962    weight: f64,
2963    x: ArrayView1<'_, f64>,
2964) where
2965    S: ndarray::DataMut<Elem = f64>,
2966{
2967    for a in 0..x.len() {
2968        let xa = x[a];
2969        if xa == 0.0 {
2970            continue;
2971        }
2972        for b in 0..x.len() {
2973            let xb = x[b];
2974            if xb == 0.0 {
2975                continue;
2976            }
2977            target[[a, b]] += weight * xa * xb;
2978        }
2979    }
2980}
2981
2982fn dense_symmetric_cross_accumulate<S>(
2983    target: &mut ndarray::ArrayBase<S, ndarray::Ix2>,
2984    weight: f64,
2985    x: ArrayView1<'_, f64>,
2986    y: ArrayView1<'_, f64>,
2987) where
2988    S: ndarray::DataMut<Elem = f64>,
2989{
2990    for a in 0..x.len() {
2991        let xa = x[a];
2992        let ya = y[a];
2993        if xa == 0.0 && ya == 0.0 {
2994            continue;
2995        }
2996        for b in 0..x.len() {
2997            let xb = x[b];
2998            let yb = y[b];
2999            let contribution = xa * yb + ya * xb;
3000            if contribution == 0.0 {
3001                continue;
3002            }
3003            target[[a, b]] += weight * contribution;
3004        }
3005    }
3006}
3007
3008fn build_latent_survival_row(
3009    row_index: usize,
3010    hazard_loading: HazardLoading,
3011    event_type: LatentSurvivalEventType,
3012    q_entry: f64,
3013    q_exit: f64,
3014    qdot_exit: f64,
3015    q_right: f64,
3016    unloaded_mass_entry: f64,
3017    unloaded_mass_exit: f64,
3018    unloaded_mass_right: f64,
3019    unloaded_hazard_exit: f64,
3020) -> Result<LatentSurvivalRow, LatentSurvivalError> {
3021    if !(q_entry.is_finite() && q_exit.is_finite()) {
3022        return Err(LatentSurvivalError::NumericalFailure {
3023            reason: format!(
3024                "latent survival requires finite q_entry and q_exit, got q_entry={q_entry}, q_exit={q_exit}"
3025            ),
3026        });
3027    }
3028    if q_exit < q_entry {
3029        return Err(LatentSurvivalError::NumericalFailure {
3030            reason: format!(
3031                "latent survival requires q_exit >= q_entry so cumulative mass is monotone, got q_entry={q_entry}, q_exit={q_exit}"
3032            ),
3033        });
3034    }
3035    if !(unloaded_mass_entry.is_finite()
3036        && unloaded_mass_exit.is_finite()
3037        && unloaded_hazard_exit.is_finite())
3038    {
3039        return Err(LatentSurvivalError::InvalidDataset {
3040            reason: format!(
3041                "latent survival requires finite unloaded components, got entry_mass={unloaded_mass_entry}, exit_mass={unloaded_mass_exit}, exit_hazard={unloaded_hazard_exit}"
3042            ),
3043        });
3044    }
3045    if unloaded_mass_entry < 0.0
3046        || unloaded_mass_exit < unloaded_mass_entry
3047        || unloaded_hazard_exit < 0.0
3048    {
3049        return Err(LatentSurvivalError::InvalidDataset {
3050            reason: format!(
3051                "latent survival requires unloaded masses/hazard to be non-negative and monotone, got entry_mass={unloaded_mass_entry}, exit_mass={unloaded_mass_exit}, exit_hazard={unloaded_hazard_exit}"
3052            ),
3053        });
3054    }
3055    let mass_entry = q_entry.exp();
3056    let mass_exit = q_exit.exp();
3057    let row = match event_type {
3058        LatentSurvivalEventType::RightCensored => {
3059            validate_unloaded_components_for_loading(
3060                "latent-survival",
3061                row_index,
3062                hazard_loading,
3063                unloaded_mass_entry,
3064                unloaded_mass_exit,
3065                Some(unloaded_hazard_exit),
3066            )?;
3067            LatentSurvivalRow::right_censored(
3068                mass_entry,
3069                mass_exit,
3070                unloaded_mass_entry,
3071                unloaded_mass_exit,
3072            )
3073        }
3074        LatentSurvivalEventType::ExactEvent => {
3075            validate_unloaded_components_for_loading(
3076                "latent-survival",
3077                row_index,
3078                hazard_loading,
3079                unloaded_mass_entry,
3080                unloaded_mass_exit,
3081                Some(unloaded_hazard_exit),
3082            )?;
3083            LatentSurvivalRow::exact_event(
3084                mass_entry,
3085                mass_exit,
3086                unloaded_mass_entry,
3087                unloaded_mass_exit,
3088                mass_exit
3089                    * if qdot_exit.is_finite() && qdot_exit > 0.0 {
3090                        qdot_exit
3091                    } else {
3092                        return Err(LatentSurvivalError::NumericalFailure {
3093                            reason: format!(
3094                                "latent survival exact event requires positive finite baseline hazard derivative, got {qdot_exit}"
3095                            ),
3096                        });
3097                    },
3098                unloaded_hazard_exit,
3099            )
3100        }
3101        LatentSurvivalEventType::IntervalCensored => {
3102            // Interval `(L, R]`: `q_exit` carries the LEFT boundary transform
3103            // `log B(L)` (so `mass_left = exp(q_exit)`) and `q_right` the RIGHT
3104            // boundary `log B(R)`. The likelihood is the survival-mass
3105            // difference `log[S(L) − S(R)]`, requiring `B(L) ≤ B(R)` i.e.
3106            // `q_exit ≤ q_right`. No event hazard participates, so the unloaded
3107            // exit hazard must be the full-loading zero (validated below via the
3108            // interval-specific unloaded check at the left/right boundaries).
3109            if !q_right.is_finite() {
3110                return Err(LatentSurvivalError::NumericalFailure {
3111                    reason: format!(
3112                        "latent survival interval row {} requires a finite q_right, got {q_right}",
3113                        row_index + 1
3114                    ),
3115                });
3116            }
3117            if q_right < q_exit {
3118                return Err(LatentSurvivalError::NumericalFailure {
3119                    reason: format!(
3120                        "latent survival interval row {} requires q_right >= q_exit (R >= L) so the \
3121                         survival-mass difference is non-negative, got q_left={q_exit}, q_right={q_right}",
3122                        row_index + 1
3123                    ),
3124                });
3125            }
3126            if !(unloaded_mass_right.is_finite()) || unloaded_mass_right < unloaded_mass_exit {
3127                return Err(LatentSurvivalError::InvalidDataset {
3128                    reason: format!(
3129                        "latent survival interval row {} requires a finite unloaded right mass >= unloaded left mass, got left={unloaded_mass_exit}, right={unloaded_mass_right}",
3130                        row_index + 1
3131                    ),
3132                });
3133            }
3134            // Interval rows carry no exit-event hazard; the loaded/unloaded
3135            // contract is validated by `LatentSurvivalRow::validate` (entry <=
3136            // left <= right monotonicity on both loaded and unloaded masses).
3137            let mass_right = q_right.exp();
3138            LatentSurvivalRow::interval_censored(
3139                mass_entry,
3140                mass_exit,
3141                mass_right,
3142                unloaded_mass_entry,
3143                unloaded_mass_exit,
3144                unloaded_mass_right,
3145            )
3146        }
3147    };
3148    row.validate()
3149        .map_err(|e| LatentSurvivalError::InvalidDataset {
3150            reason: e.to_string(),
3151        })?;
3152    Ok(row)
3153}
3154
3155#[derive(Clone, Copy)]
3156struct BinaryFromLogSurvival {
3157    log_lik: f64,
3158    /// dℓ/ds where s = log_survival and ℓ = log_lik. For event=1 this is
3159    /// ℓ' = -S/(1-S); for event=0 this is ℓ' = 1 (because ℓ ≡ s).
3160    grad_scale: f64,
3161    /// Coefficient applied to `survival_jet.neg_hessian` (which equals
3162    /// -d²s/dβ²) when assembling the negative Hessian of `wi * log_lik`
3163    /// against β. The Newton accumulator computes
3164    ///     neg_Hess(log_lik) = grad_scale * neg_hessian + outer_scale * score²
3165    /// so by the chain rule this MUST equal `grad_scale` (= ℓ'). Keeping
3166    /// the two fields separate is purely for readability at call sites;
3167    /// the `assert!` in [`binary_log_survival_scales`] enforces the
3168    /// equality.
3169    neg_hess_scale: f64,
3170    /// -ℓ''(s). For event=1 this is +S/(1-S)²; for event=0 it is 0.
3171    outer_scale: f64,
3172    /// ℓ''(s) — derivative of `grad_scale` w.r.t. s.
3173    grad_scale_prime: f64,
3174    /// ℓ'''(s) — second derivative of `grad_scale` w.r.t. s.
3175    grad_scale_second: f64,
3176    /// -ℓ'''(s) — derivative of `outer_scale` w.r.t. s.
3177    outer_scale_prime: f64,
3178    /// -ℓ''''(s) — second derivative of `outer_scale` w.r.t. s.
3179    outer_scale_second: f64,
3180}
3181
3182/// Analytic source of truth for the directional derivatives of
3183/// ℓ(s) = log(1 - exp(s)) at s = `log_survival`. Returns
3184/// `(ℓ, ℓ', ℓ'', ℓ''', ℓ'''')`. All consumer scales (`grad_scale`,
3185/// `neg_hess_scale`, `outer_scale`, and their two derivatives each)
3186/// are derived from this single function so the sign/algebra cannot
3187/// drift between sites.
3188#[inline]
3189fn binary_log_survival_scales(survival: f64, event_prob: f64) -> (f64, f64, f64, f64, f64) {
3190    // ℓ(s)   = log(1 - exp(s)) = log(event_prob)
3191    // dS/ds  = S,    dP/ds = -S        (S=survival, P=event_prob)
3192    // ℓ'(s)  = -S/P
3193    // ℓ''(s) = d/ds[-S/P] = -S/P²        (since P + S = 1)
3194    // ℓ'''(s) = d/ds[-S/P²] = -S(1 + S)/P³
3195    // ℓ''''(s) = d/ds[-S(1+S)/P³]
3196    //          = -S/P³ - 3S²/P³ - 6S²(1+S)/P⁴ - ... ; expanded form below.
3197    let log_lik = event_prob.ln();
3198    let p = event_prob;
3199    let p2 = p * p;
3200    let p3 = p2 * p;
3201    let p4 = p3 * p;
3202    let s = survival;
3203    let s2 = s * s;
3204    let s3 = s2 * s;
3205    let ell_prime = -s / p;
3206    let ell_pp = -s / p2;
3207    let ell_ppp = -s * (1.0 + s) / p3;
3208    // ℓ''''(s) = -S·(1 + 4S + S²) / P⁴ - 3·S²·(1+S)/P⁴? Use the equivalent
3209    // expansion that matches the prior closed form:
3210    //   d/ds[-S(1+S)/P³] = -(S + 2S²)/P³ - 3·S·(1+S)·S/P⁴
3211    //                    = -(S + 2S²)/P³ - 3S²(1+S)/P⁴
3212    // Combining over P⁴: -(S + 2S²)·P/P⁴ - 3S²(1+S)/P⁴
3213    //                  = -[S·P + 2S²·P + 3S² + 3S³] / P⁴
3214    // With P = 1 - S: S·P = S - S²; 2S²·P = 2S² - 2S³.
3215    //   numerator = -[S - S² + 2S² - 2S³ + 3S² + 3S³] = -[S + 4S² + S³].
3216    // So ℓ''''(s) = -(S + 4S² + S³) / P⁴.
3217    let ell_pppp = -(s + 4.0 * s2 + s3) / p4;
3218    (log_lik, ell_prime, ell_pp, ell_ppp, ell_pppp)
3219}
3220
3221fn binary_from_log_survival(
3222    log_survival: f64,
3223    event: u8,
3224) -> Result<BinaryFromLogSurvival, LatentSurvivalError> {
3225    if event == 0 {
3226        // ℓ(s) = s ⇒ ℓ' = 1, ℓ'' = ℓ''' = ℓ'''' = 0.
3227        return Ok(BinaryFromLogSurvival {
3228            log_lik: log_survival,
3229            grad_scale: 1.0,
3230            neg_hess_scale: 1.0,
3231            outer_scale: 0.0,
3232            grad_scale_prime: 0.0,
3233            grad_scale_second: 0.0,
3234            outer_scale_prime: 0.0,
3235            outer_scale_second: 0.0,
3236        });
3237    }
3238    if event != 1 {
3239        return Err(LatentSurvivalError::InvalidDataset {
3240            reason: format!("latent-binary requires event targets in {{0,1}}, got {event}"),
3241        });
3242    }
3243    // Cap log S(t) strictly below zero so the event probability
3244    // `1 - exp(log S)` stays strictly positive even when the survival
3245    // probability rounds to exactly 1 (log S == 0): a zero event probability
3246    // would make the binary log-likelihood `log(event_prob)` diverge. The cap
3247    // is at the f64 resolution near 1.0, so it never perturbs a genuinely
3248    // informative survival value.
3249    const MAX_LOG_SURVIVAL: f64 = -1e-15;
3250    let log_survival = log_survival.min(MAX_LOG_SURVIVAL);
3251    let survival = log_survival.exp();
3252    let event_prob = 1.0 - survival;
3253    if !(event_prob.is_finite() && event_prob > 0.0) {
3254        return Err(LatentSurvivalError::NumericalFailure {
3255            reason: format!(
3256                "latent-binary encountered non-positive event probability from log survival {log_survival}"
3257            ),
3258        });
3259    }
3260    let (log_lik, ell_prime, ell_pp, ell_ppp, ell_pppp) =
3261        binary_log_survival_scales(survival, event_prob);
3262    let grad_scale = ell_prime;
3263    let neg_hess_scale = ell_prime; // coefficient on (-d²s/dβ²); equals ℓ'.
3264    let outer_scale = -ell_pp;
3265    let grad_scale_prime = ell_pp;
3266    let grad_scale_second = ell_ppp;
3267    let outer_scale_prime = -ell_ppp;
3268    let outer_scale_second = -ell_pppp;
3269    // The Newton accumulator at the call sites computes
3270    //     neg_Hess(log_lik) = neg_hess_scale * (-d²s/dβ²) + outer_scale * (ds/dβ)²
3271    // For this identity to hold by the chain rule, the coefficient on the
3272    // neg_hessian term must equal ℓ' (== grad_scale). Document the invariant.
3273    assert!(
3274        (grad_scale - neg_hess_scale).abs() <= 1e-15 * grad_scale.abs().max(1.0),
3275        "binary_from_log_survival invariant: neg_hess_scale ({neg_hess_scale}) must equal grad_scale ({grad_scale}) so that grad_scale and the coefficient on neg_hessian share sign"
3276    );
3277    assert!(
3278        outer_scale >= 0.0 || !outer_scale.is_finite(),
3279        "binary_from_log_survival invariant: outer_scale (= -ℓ'') must be non-negative for event=1; got {outer_scale}"
3280    );
3281    Ok(BinaryFromLogSurvival {
3282        log_lik,
3283        grad_scale,
3284        neg_hess_scale,
3285        outer_scale,
3286        grad_scale_prime,
3287        grad_scale_second,
3288        outer_scale_prime,
3289        outer_scale_second,
3290    })
3291}
3292
3293impl LatentBinaryFamily {
3294    /// Assemble the per-row [`LatentSurvivalRow`] for a row treated as a pure
3295    /// right-censored survival contribution (exit time is the censoring
3296    /// boundary, unit exit-hazard derivative, no right / post-exit unloaded
3297    /// mass). Shared by every per-row binary-from-survival pullback reduction;
3298    /// behavior is identical to the previously inlined `RightCensored` call.
3299    fn build_right_censored_row_at(
3300        &self,
3301        row_idx: usize,
3302        q_entry: f64,
3303        q_exit: f64,
3304    ) -> Result<LatentSurvivalRow, LatentSurvivalError> {
3305        build_latent_survival_row(
3306            row_idx,
3307            self.hazard_loading,
3308            LatentSurvivalEventType::RightCensored,
3309            q_entry,
3310            q_exit,
3311            1.0,
3312            q_exit,
3313            self.unloaded_mass_entry[row_idx],
3314            self.unloaded_mass_exit[row_idx],
3315            0.0,
3316            0.0,
3317        )
3318    }
3319
3320    fn joint_slices(&self) -> LatentSurvivalJointSlices {
3321        let p_time = self.x_time_exit.ncols();
3322        let p_mean = self.x_mean.ncols();
3323        LatentSurvivalJointSlices {
3324            time: 0..p_time,
3325            mean: p_time..p_time + p_mean,
3326            log_sigma: None,
3327            total: p_time + p_mean,
3328        }
3329    }
3330
3331    fn row_primary_direction_from_flat(
3332        &self,
3333        row: usize,
3334        slices: &LatentSurvivalJointSlices,
3335        d_beta_flat: &Array1<f64>,
3336    ) -> Array1<f64> {
3337        let mut out = Array1::<f64>::zeros(LATENT_SURVIVAL_PRIMARY_DIM);
3338        let d_time = d_beta_flat.slice(s![slices.time.clone()]);
3339        out[LATENT_SURVIVAL_PRIMARY_Q_ENTRY] = self.x_time_entry.row(row).dot(&d_time);
3340        out[LATENT_SURVIVAL_PRIMARY_Q_EXIT] = self.x_time_exit.row(row).dot(&d_time);
3341        out[LATENT_SURVIVAL_PRIMARY_MU] = self
3342            .x_mean
3343            .dot_row_view(row, d_beta_flat.slice(s![slices.mean.clone()]));
3344        out
3345    }
3346
3347    fn add_pullback_primary_gradient(
3348        &self,
3349        target: &mut Array1<f64>,
3350        row: usize,
3351        slices: &LatentSurvivalJointSlices,
3352        primary_gradient: &Array1<f64>,
3353        weight: f64,
3354    ) {
3355        for (primary_idx, time_vec) in [
3356            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
3357            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
3358        ] {
3359            let scale = weight * primary_gradient[primary_idx];
3360            if scale == 0.0 {
3361                continue;
3362            }
3363            for i in 0..time_vec.len() {
3364                let xi = time_vec[i];
3365                if xi != 0.0 {
3366                    target[slices.time.start + i] += scale * xi;
3367                }
3368            }
3369        }
3370
3371        let mean_scale = weight * primary_gradient[LATENT_SURVIVAL_PRIMARY_MU];
3372        if mean_scale != 0.0 {
3373            self.x_mean
3374                .axpy_row_into(
3375                    row,
3376                    mean_scale,
3377                    &mut target.slice_mut(s![slices.mean.clone()]),
3378                )
3379                // SAFETY: `slices.mean` sized at construction to match
3380                // `x_mean.ncols()`; an error means caller-side shape drift,
3381                // an invariant violation. A swallowed sentinel would silently
3382                // corrupt the joint gradient, so fail loudly instead.
3383                .unwrap_or_else(|error| {
3384                    panic!(
3385                        "latent binary mean gradient pullback dimension mismatch: row={row}, mean_slice={:?}, target_len={}, x_mean_cols={}, error={error}",
3386                        slices.mean,
3387                        target.len(),
3388                        self.x_mean.ncols()
3389                    )
3390                });
3391        }
3392    }
3393
3394    fn add_pullback_primary_hessian(
3395        &self,
3396        target: &mut Array2<f64>,
3397        row: usize,
3398        slices: &LatentSurvivalJointSlices,
3399        primary_hessian: &Array2<f64>,
3400    ) {
3401        {
3402            let time_target = &mut target.slice_mut(s![slices.time.clone(), slices.time.clone()]);
3403            dense_outer_accumulate(
3404                time_target,
3405                primary_hessian[[
3406                    LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3407                    LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3408                ]],
3409                self.x_time_entry.row(row),
3410            );
3411            dense_outer_accumulate(
3412                time_target,
3413                primary_hessian[[
3414                    LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3415                    LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3416                ]],
3417                self.x_time_exit.row(row),
3418            );
3419            dense_symmetric_cross_accumulate(
3420                time_target,
3421                primary_hessian[[
3422                    LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3423                    LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3424                ]],
3425                self.x_time_entry.row(row),
3426                self.x_time_exit.row(row),
3427            );
3428        }
3429
3430        let mean_weight = primary_hessian[[LATENT_SURVIVAL_PRIMARY_MU, LATENT_SURVIVAL_PRIMARY_MU]];
3431        self.x_mean
3432            .syr_row_into_view(
3433                row,
3434                mean_weight,
3435                target.slice_mut(s![slices.mean.clone(), slices.mean.clone()]),
3436            )
3437            .unwrap_or_else(|error| {
3438                // SAFETY: `slices.mean` × `slices.mean` slab sized at
3439                // construction to `x_mean.ncols()` × `x_mean.ncols()`;
3440                // an error here is caller-side shape drift, an invariant
3441                // violation. A swallowed sentinel would silently corrupt the
3442                // joint Hessian, so fail loudly instead.
3443                panic!(
3444                    "latent binary mean Hessian pullback dimension mismatch: row={row}, mean_slice={:?}, target_dim={:?}, x_mean_cols={}, error={error}",
3445                    slices.mean,
3446                    target.dim(),
3447                    self.x_mean.ncols()
3448                )
3449            });
3450
3451        let mean_row = self
3452            .x_mean
3453            .try_row_chunk(row..row + 1)
3454            .unwrap_or_else(|error| {
3455                // SAFETY: row index comes from the enclosing `0..n` loop
3456                // bound by `self.x_mean.nrows()`, so `row..row+1` is
3457                // always a valid single-row chunk.
3458                panic!(
3459                    "latent binary mean pullback row chunk failed: row={row}, x_mean_rows={}, x_mean_cols={}, error={error}",
3460                    self.x_mean.nrows(),
3461                    self.x_mean.ncols()
3462                )
3463            });
3464        let mean_vec = mean_row.row(0);
3465        for (primary_idx, time_vec) in [
3466            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
3467            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
3468        ] {
3469            let weight = primary_hessian[[primary_idx, LATENT_SURVIVAL_PRIMARY_MU]];
3470            if weight == 0.0 {
3471                continue;
3472            }
3473            for i in 0..time_vec.len() {
3474                let xi = time_vec[i];
3475                if xi == 0.0 {
3476                    continue;
3477                }
3478                for j in 0..mean_vec.len() {
3479                    let xj = mean_vec[j];
3480                    if xj == 0.0 {
3481                        continue;
3482                    }
3483                    target[[slices.time.start + i, slices.mean.start + j]] += weight * xi * xj;
3484                    target[[slices.mean.start + j, slices.time.start + i]] += weight * xj * xi;
3485                }
3486            }
3487        }
3488    }
3489
3490    fn evaluate_exact_newton_joint_dense(
3491        &self,
3492        block_states: &[ParameterBlockState],
3493    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
3494        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
3495        let slices = self.joint_slices();
3496        let mut ll = 0.0;
3497        let mut gradient = Array1::<f64>::zeros(slices.total);
3498        let mut hessian = Array2::<f64>::zeros((slices.total, slices.total));
3499        for row_idx in 0..self.event_target.len() {
3500            let wi = self.weights[row_idx];
3501            if wi <= MIN_WEIGHT {
3502                continue;
3503            }
3504            let row =
3505                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
3506            let (row_log_survival, survival_gradient, survival_hessian) =
3507                latent_survival_row_primary_gradient_hessian(
3508                    &self.quadctx,
3509                    &row,
3510                    q_entry[row_idx],
3511                    q_exit[row_idx],
3512                    1.0,
3513                    q_exit[row_idx],
3514                    mu[row_idx],
3515                    self.latent_sd,
3516                    false,
3517                )?;
3518            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
3519            ll += wi * binary.log_lik;
3520            let primary_gradient = binary.grad_scale * &survival_gradient;
3521            let mut primary_hessian = binary.grad_scale * survival_hessian;
3522            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3523                for b in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3524                    primary_hessian[[a, b]] +=
3525                        binary.outer_scale * survival_gradient[a] * survival_gradient[b];
3526                }
3527            }
3528            self.add_pullback_primary_gradient(
3529                &mut gradient,
3530                row_idx,
3531                &slices,
3532                &primary_gradient,
3533                wi,
3534            );
3535            self.add_pullback_primary_hessian(
3536                &mut hessian,
3537                row_idx,
3538                &slices,
3539                &(wi * primary_hessian),
3540            );
3541        }
3542        Ok((ll, gradient, hessian))
3543    }
3544
3545    /// Per-row residuals of the unpenalized NLL with respect to the baseline
3546    /// time-block offsets `(entry, exit)`.
3547    ///
3548    /// The latent-binary deployment likelihood is a monotone scalar transform
3549    /// `ℓ_bin = b(log S_row)` of the latent-survival row log-survival, so by the
3550    /// chain rule `∂ℓ_bin/∂q_ch = b'(log S)·∂(log S)/∂q_ch = grad_scale·g_ch`,
3551    /// where `g_ch` are the `Q_ENTRY`/`Q_EXIT` components of the survival row
3552    /// primary gradient. The baseline θ enters only the additive entry/exit time
3553    /// offsets (`q̇_exit` is held at the constant deployment derivative `1`, so
3554    /// the derivative channel carries no baseline offset and its residual is 0).
3555    /// Sampleweight-scaled to match the [`OffsetChannelResiduals`] contract.
3556    pub fn offset_channel_residuals(
3557        &self,
3558        block_states: &[ParameterBlockState],
3559    ) -> Result<crate::survival::OffsetChannelResiduals, String> {
3560        let n = self.event_target.len();
3561        if block_states.is_empty() {
3562            log::warn!(
3563                "LatentBinaryFamily::offset_channel_residuals: block_states is empty \
3564                 (degraded fit); returning zero offset residuals (n={n})"
3565            );
3566            return Ok(crate::survival::OffsetChannelResiduals {
3567                exit: Array1::<f64>::zeros(n),
3568                entry: Array1::<f64>::zeros(n),
3569                derivative: Array1::<f64>::zeros(n),
3570                right: Array1::<f64>::zeros(n),
3571            });
3572        }
3573        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
3574        let mut entry = Array1::<f64>::zeros(n);
3575        let mut exit = Array1::<f64>::zeros(n);
3576        for row_idx in 0..n {
3577            let wi = self.weights[row_idx];
3578            if wi <= MIN_WEIGHT {
3579                continue;
3580            }
3581            let row =
3582                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
3583            let (row_log_survival, survival_gradient, _) =
3584                latent_survival_row_primary_gradient_hessian(
3585                    &self.quadctx,
3586                    &row,
3587                    q_entry[row_idx],
3588                    q_exit[row_idx],
3589                    1.0,
3590                    q_exit[row_idx],
3591                    mu[row_idx],
3592                    self.latent_sd,
3593                    false,
3594                )?;
3595            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
3596            // ∂NLL/∂o_ch = −w · grad_scale · ∂(log S)/∂q_ch.
3597            entry[row_idx] =
3598                -wi * binary.grad_scale * survival_gradient[LATENT_SURVIVAL_PRIMARY_Q_ENTRY];
3599            exit[row_idx] =
3600                -wi * binary.grad_scale * survival_gradient[LATENT_SURVIVAL_PRIMARY_Q_EXIT];
3601        }
3602        Ok(crate::survival::OffsetChannelResiduals {
3603            exit,
3604            entry,
3605            derivative: Array1::<f64>::zeros(n),
3606            // Latent-binary deployment has no interval upper bound; the `R`
3607            // channel is structurally absent (every row is right-censored).
3608            right: Array1::<f64>::zeros(n),
3609        })
3610    }
3611
3612    fn exact_newton_joint_hessian_directional_derivative_dense(
3613        &self,
3614        block_states: &[ParameterBlockState],
3615        d_beta_flat: &Array1<f64>,
3616    ) -> Result<Array2<f64>, String> {
3617        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
3618        let slices = self.joint_slices();
3619        if d_beta_flat.len() != slices.total {
3620            return Err(format!(
3621                "latent binary joint dH direction length mismatch: got {}, expected {}",
3622                d_beta_flat.len(),
3623                slices.total
3624            ));
3625        }
3626        let mut out = Array2::<f64>::zeros((slices.total, slices.total));
3627        for row_idx in 0..self.event_target.len() {
3628            let wi = self.weights[row_idx];
3629            if wi <= MIN_WEIGHT {
3630                continue;
3631            }
3632            let row =
3633                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
3634            let (row_log_survival, survival_gradient, survival_hessian) =
3635                latent_survival_row_primary_gradient_hessian(
3636                    &self.quadctx,
3637                    &row,
3638                    q_entry[row_idx],
3639                    q_exit[row_idx],
3640                    1.0,
3641                    q_exit[row_idx],
3642                    mu[row_idx],
3643                    self.latent_sd,
3644                    false,
3645                )?;
3646            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
3647            let direction = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_flat);
3648            let third = latent_survival_row_primary_third_contracted(
3649                &self.quadctx,
3650                &row,
3651                q_entry[row_idx],
3652                q_exit[row_idx],
3653                1.0,
3654                q_exit[row_idx],
3655                mu[row_idx],
3656                self.latent_sd,
3657                &direction,
3658                false,
3659            )?;
3660            let g_u = -survival_hessian.dot(&direction);
3661            let t_u = survival_gradient.dot(&direction);
3662            let mut primary = binary.grad_scale * third;
3663            primary.scaled_add(binary.grad_scale_prime * t_u, &survival_hessian);
3664            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3665                for b in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3666                    primary[[a, b]] += binary.outer_scale_prime
3667                        * t_u
3668                        * survival_gradient[a]
3669                        * survival_gradient[b]
3670                        + binary.outer_scale
3671                            * (g_u[a] * survival_gradient[b] + survival_gradient[a] * g_u[b]);
3672                }
3673            }
3674            self.add_pullback_primary_hessian(&mut out, row_idx, &slices, &(wi * primary));
3675        }
3676        Ok(out)
3677    }
3678
3679    fn exact_newton_joint_hessian_second_directional_derivative_dense(
3680        &self,
3681        block_states: &[ParameterBlockState],
3682        d_beta_u_flat: &Array1<f64>,
3683        d_beta_v_flat: &Array1<f64>,
3684    ) -> Result<Array2<f64>, String> {
3685        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
3686        let slices = self.joint_slices();
3687        if d_beta_u_flat.len() != slices.total || d_beta_v_flat.len() != slices.total {
3688            return Err(format!(
3689                "latent binary joint d2H direction length mismatch: got {} and {}, expected {}",
3690                d_beta_u_flat.len(),
3691                d_beta_v_flat.len(),
3692                slices.total
3693            ));
3694        }
3695        let mut out = Array2::<f64>::zeros((slices.total, slices.total));
3696        for row_idx in 0..self.event_target.len() {
3697            let wi = self.weights[row_idx];
3698            if wi <= MIN_WEIGHT {
3699                continue;
3700            }
3701            let row =
3702                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
3703            let (row_log_survival, survival_gradient, survival_hessian) =
3704                latent_survival_row_primary_gradient_hessian(
3705                    &self.quadctx,
3706                    &row,
3707                    q_entry[row_idx],
3708                    q_exit[row_idx],
3709                    1.0,
3710                    q_exit[row_idx],
3711                    mu[row_idx],
3712                    self.latent_sd,
3713                    false,
3714                )?;
3715            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
3716            let direction_u = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_u_flat);
3717            let direction_v = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_v_flat);
3718            let third_u = latent_survival_row_primary_third_contracted(
3719                &self.quadctx,
3720                &row,
3721                q_entry[row_idx],
3722                q_exit[row_idx],
3723                1.0,
3724                q_exit[row_idx],
3725                mu[row_idx],
3726                self.latent_sd,
3727                &direction_u,
3728                false,
3729            )?;
3730            let third_v = latent_survival_row_primary_third_contracted(
3731                &self.quadctx,
3732                &row,
3733                q_entry[row_idx],
3734                q_exit[row_idx],
3735                1.0,
3736                q_exit[row_idx],
3737                mu[row_idx],
3738                self.latent_sd,
3739                &direction_v,
3740                false,
3741            )?;
3742            let fourth = latent_survival_row_primary_fourth_contracted(
3743                &self.quadctx,
3744                &row,
3745                q_entry[row_idx],
3746                q_exit[row_idx],
3747                1.0,
3748                q_exit[row_idx],
3749                mu[row_idx],
3750                self.latent_sd,
3751                &direction_u,
3752                &direction_v,
3753                false,
3754            )?;
3755            let g_u = -survival_hessian.dot(&direction_u);
3756            let g_v = -survival_hessian.dot(&direction_v);
3757            let g_uv = -third_v.dot(&direction_u);
3758            let t_u = survival_gradient.dot(&direction_u);
3759            let t_v = survival_gradient.dot(&direction_v);
3760            let l_uv = -direction_u.dot(&survival_hessian.dot(&direction_v));
3761            let c_u = binary.grad_scale_prime * t_u;
3762            let c_v = binary.grad_scale_prime * t_v;
3763            let c_uv = binary.grad_scale_second * t_u * t_v + binary.grad_scale_prime * l_uv;
3764            let o_u = binary.outer_scale_prime * t_u;
3765            let o_v = binary.outer_scale_prime * t_v;
3766            let o_uv = binary.outer_scale_second * t_u * t_v + binary.outer_scale_prime * l_uv;
3767            let mut primary = binary.grad_scale * fourth;
3768            primary.scaled_add(c_u, &third_v);
3769            primary.scaled_add(c_v, &third_u);
3770            primary.scaled_add(c_uv, &survival_hessian);
3771            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3772                for b in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3773                    primary[[a, b]] += o_uv * survival_gradient[a] * survival_gradient[b]
3774                        + o_v * (g_u[a] * survival_gradient[b] + survival_gradient[a] * g_u[b])
3775                        + o_u * (g_v[a] * survival_gradient[b] + survival_gradient[a] * g_v[b])
3776                        + binary.outer_scale
3777                            * (g_uv[a] * survival_gradient[b]
3778                                + g_u[a] * g_v[b]
3779                                + g_v[a] * g_u[b]
3780                                + survival_gradient[a] * g_uv[b]);
3781                }
3782            }
3783            self.add_pullback_primary_hessian(&mut out, row_idx, &slices, &(wi * primary));
3784        }
3785        Ok(out)
3786    }
3787}
3788
3789/// Shared interface that both `LatentSurvivalFamily` and `LatentBinaryFamily`
3790/// expose to the joint Hessian workspace.
3791///
3792/// The two families produce the same `ExactNewtonJointHessianWorkspace`
3793/// shape — five of the six workspace methods are pure delegations to a
3794/// matching family method (dense evaluation, directional derivatives, and the
3795/// `slices` cache). The only family-specific piece is the per-row matvec body:
3796/// the survival family iterates over real (entry, exit, ḋ) triples and may
3797/// carry a log-σ block, while the binary family rewrites the same row kernel
3798/// through `binary_from_log_survival(·)` to recover the per-row binary
3799/// gradient/Hessian. That single difference is captured by `ws_matvec_into`;
3800/// every other method is shared by the generic `LatentHessianWorkspace<F>`
3801/// below.
3802trait LatentJointHessianFamily {
3803    fn ws_joint_slices(&self) -> LatentSurvivalJointSlices;
3804
3805    fn ws_evaluate_dense(
3806        &self,
3807        block_states: &[ParameterBlockState],
3808    ) -> Result<(f64, Array1<f64>, Array2<f64>), String>;
3809
3810    fn ws_dh_directional(
3811        &self,
3812        block_states: &[ParameterBlockState],
3813        d_beta_flat: &Array1<f64>,
3814    ) -> Result<Array2<f64>, String>;
3815
3816    fn ws_dh_second_directional(
3817        &self,
3818        block_states: &[ParameterBlockState],
3819        d_beta_u: &Array1<f64>,
3820        d_beta_v: &Array1<f64>,
3821    ) -> Result<Array2<f64>, String>;
3822
3823    /// Family-specific per-row Hessian matvec body, hoisted out of the
3824    /// workspace impl. Writes `out := H · v` (with `out.fill(0.0)` already
3825    /// performed by the caller) using the family's row kernel.
3826    fn ws_matvec_into(
3827        &self,
3828        slices: &LatentSurvivalJointSlices,
3829        block_states: &[ParameterBlockState],
3830        v: &Array1<f64>,
3831        out: &mut Array1<f64>,
3832    ) -> Result<bool, String>;
3833
3834    /// Family-name fragment used in the workspace's dimension-mismatch error
3835    /// message, so callers still see "latent survival …" / "latent binary …"
3836    /// after the workspace impl was unified.
3837    fn ws_label() -> &'static str;
3838}
3839
3840impl LatentJointHessianFamily for LatentSurvivalFamily {
3841    fn ws_joint_slices(&self) -> LatentSurvivalJointSlices {
3842        self.joint_slices()
3843    }
3844
3845    fn ws_evaluate_dense(
3846        &self,
3847        block_states: &[ParameterBlockState],
3848    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
3849        self.evaluate_exact_newton_joint_dense(block_states)
3850    }
3851
3852    fn ws_dh_directional(
3853        &self,
3854        block_states: &[ParameterBlockState],
3855        d_beta_flat: &Array1<f64>,
3856    ) -> Result<Array2<f64>, String> {
3857        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
3858    }
3859
3860    fn ws_dh_second_directional(
3861        &self,
3862        block_states: &[ParameterBlockState],
3863        d_beta_u: &Array1<f64>,
3864        d_beta_v: &Array1<f64>,
3865    ) -> Result<Array2<f64>, String> {
3866        self.exact_newton_joint_hessian_second_directional_derivative_dense(
3867            block_states,
3868            d_beta_u,
3869            d_beta_v,
3870        )
3871    }
3872
3873    fn ws_matvec_into(
3874        &self,
3875        slices: &LatentSurvivalJointSlices,
3876        block_states: &[ParameterBlockState],
3877        v: &Array1<f64>,
3878        out: &mut Array1<f64>,
3879    ) -> Result<bool, String> {
3880        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
3881        let q_right = self.time_q_right(block_states)?;
3882        let sigma = self.latent_sd(block_states)?;
3883        let include_log_sigma = slices.log_sigma.is_some();
3884        for row_idx in 0..self.event_target.len() {
3885            let wi = self.weights[row_idx];
3886            if wi <= MIN_WEIGHT {
3887                continue;
3888            }
3889            let row = self.build_row_at(
3890                row_idx,
3891                q_entry[row_idx],
3892                q_exit[row_idx],
3893                qdot_exit[row_idx],
3894                q_right[row_idx],
3895            )?;
3896            let (_, _, primary_hessian) = latent_survival_row_primary_gradient_hessian(
3897                &self.quadctx,
3898                &row,
3899                q_entry[row_idx],
3900                q_exit[row_idx],
3901                qdot_exit[row_idx],
3902                q_right[row_idx],
3903                mu[row_idx],
3904                sigma,
3905                include_log_sigma,
3906            )?;
3907            let primary_dir = self.row_primary_direction_from_flat(row_idx, slices, v);
3908            let primary_hv = primary_hessian.dot(&primary_dir);
3909            self.add_pullback_primary_gradient(out, row_idx, slices, &primary_hv, wi)?;
3910        }
3911        Ok(true)
3912    }
3913
3914    fn ws_label() -> &'static str {
3915        "survival"
3916    }
3917}
3918
3919impl LatentJointHessianFamily for LatentBinaryFamily {
3920    fn ws_joint_slices(&self) -> LatentSurvivalJointSlices {
3921        self.joint_slices()
3922    }
3923
3924    fn ws_evaluate_dense(
3925        &self,
3926        block_states: &[ParameterBlockState],
3927    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
3928        self.evaluate_exact_newton_joint_dense(block_states)
3929    }
3930
3931    fn ws_dh_directional(
3932        &self,
3933        block_states: &[ParameterBlockState],
3934        d_beta_flat: &Array1<f64>,
3935    ) -> Result<Array2<f64>, String> {
3936        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
3937    }
3938
3939    fn ws_dh_second_directional(
3940        &self,
3941        block_states: &[ParameterBlockState],
3942        d_beta_u: &Array1<f64>,
3943        d_beta_v: &Array1<f64>,
3944    ) -> Result<Array2<f64>, String> {
3945        self.exact_newton_joint_hessian_second_directional_derivative_dense(
3946            block_states,
3947            d_beta_u,
3948            d_beta_v,
3949        )
3950    }
3951
3952    fn ws_matvec_into(
3953        &self,
3954        slices: &LatentSurvivalJointSlices,
3955        block_states: &[ParameterBlockState],
3956        v: &Array1<f64>,
3957        out: &mut Array1<f64>,
3958    ) -> Result<bool, String> {
3959        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
3960        for row_idx in 0..self.event_target.len() {
3961            let wi = self.weights[row_idx];
3962            if wi <= MIN_WEIGHT {
3963                continue;
3964            }
3965            let row =
3966                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
3967            let (row_log_survival, survival_gradient, survival_hessian) =
3968                latent_survival_row_primary_gradient_hessian(
3969                    &self.quadctx,
3970                    &row,
3971                    q_entry[row_idx],
3972                    q_exit[row_idx],
3973                    1.0,
3974                    q_exit[row_idx],
3975                    mu[row_idx],
3976                    self.latent_sd,
3977                    false,
3978                )?;
3979            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
3980            let primary_dir = self.row_primary_direction_from_flat(row_idx, slices, v);
3981            let mut primary_hv = binary.grad_scale * survival_hessian.dot(&primary_dir);
3982            let outer_dot = survival_gradient.dot(&primary_dir);
3983            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
3984                primary_hv[a] += binary.outer_scale * survival_gradient[a] * outer_dot;
3985            }
3986            self.add_pullback_primary_gradient(out, row_idx, slices, &primary_hv, wi);
3987        }
3988        Ok(true)
3989    }
3990
3991    fn ws_label() -> &'static str {
3992        "binary"
3993    }
3994}
3995
3996/// Joint exact-Newton Hessian workspace shared by `LatentSurvivalFamily` and
3997/// `LatentBinaryFamily`. The two families plug into the workspace via
3998/// `LatentJointHessianFamily`; this struct holds the shared bookkeeping
3999/// (block states + cached slices) and routes every trait method either through
4000/// a thin family delegation or through the family's `ws_matvec_into` row
4001/// kernel.
4002struct LatentHessianWorkspace<F: LatentJointHessianFamily> {
4003    family: F,
4004    block_states: Vec<ParameterBlockState>,
4005    slices: LatentSurvivalJointSlices,
4006}
4007
4008impl<F: LatentJointHessianFamily> LatentHessianWorkspace<F> {
4009    fn new(family: F, block_states: Vec<ParameterBlockState>) -> Self {
4010        let slices = family.ws_joint_slices();
4011        Self {
4012            family,
4013            block_states,
4014            slices,
4015        }
4016    }
4017}
4018
4019impl<F> ExactNewtonJointHessianWorkspace for LatentHessianWorkspace<F>
4020where
4021    F: LatentJointHessianFamily + Send + Sync + 'static,
4022{
4023    fn hessian_dense(&self) -> Result<Option<Array2<f64>>, String> {
4024        self.family
4025            .ws_evaluate_dense(&self.block_states)
4026            .map(|(_, _, hessian)| Some(hessian))
4027    }
4028
4029    fn hessian_matvec(&self, v: &Array1<f64>) -> Result<Option<Array1<f64>>, String> {
4030        let mut out = Array1::<f64>::zeros(self.slices.total);
4031        self.hessian_matvec_into(v, &mut out)?;
4032        Ok(Some(out))
4033    }
4034
4035    fn hessian_matvec_into(&self, v: &Array1<f64>, out: &mut Array1<f64>) -> Result<bool, String> {
4036        if v.len() != self.slices.total || out.len() != self.slices.total {
4037            return Err(format!(
4038                "latent {} Hessian matvec dimension mismatch: v={} out={} expected={}",
4039                F::ws_label(),
4040                v.len(),
4041                out.len(),
4042                self.slices.total
4043            ));
4044        }
4045        out.fill(0.0);
4046        self.family
4047            .ws_matvec_into(&self.slices, &self.block_states, v, out)
4048    }
4049
4050    fn hessian_diagonal(&self) -> Result<Option<Array1<f64>>, String> {
4051        let dense = self.family.ws_evaluate_dense(&self.block_states)?.2;
4052        Ok(Some(dense.diag().to_owned()))
4053    }
4054
4055    fn directional_derivative(
4056        &self,
4057        d_beta_flat: &Array1<f64>,
4058    ) -> Result<Option<Array2<f64>>, String> {
4059        self.family
4060            .ws_dh_directional(&self.block_states, d_beta_flat)
4061            .map(Some)
4062    }
4063
4064    fn second_directional_derivative(
4065        &self,
4066        d_beta_u: &Array1<f64>,
4067        d_beta_v: &Array1<f64>,
4068    ) -> Result<Option<Array2<f64>>, String> {
4069        self.family
4070            .ws_dh_second_directional(&self.block_states, d_beta_u, d_beta_v)
4071            .map(Some)
4072    }
4073}
4074
4075type LatentSurvivalHessianWorkspace = LatentHessianWorkspace<LatentSurvivalFamily>;
4076type LatentBinaryHessianWorkspace = LatentHessianWorkspace<LatentBinaryFamily>;
4077
4078impl CustomFamily for LatentSurvivalFamily {
4079    // Latent survival fits keep the self-limiting Jeffreys/Firth curvature
4080    // active for their under-identification regime. The trait default flipped to
4081    // OFF in gam#1395 (flat-prior exact-Newton objective); opt back in here.
4082    fn joint_jeffreys_term_required(&self) -> bool {
4083        true
4084    }
4085
4086    fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
4087        true
4088    }
4089
4090    fn has_explicit_joint_hessian(&self) -> bool {
4091        true
4092    }
4093
4094    /// Route the pre-fit identifiability audit CHANNEL-AWARE across the latent
4095    /// survival blocks. The time-transform baseline `q(a)`, the frailty mean
4096    /// `μ = Xβ`, and the learnable log-σ scale are three STRUCTURALLY DISTINCT
4097    /// outputs — block-diagonal entries of the true joint Jacobian
4098    /// `blkdiag(X_time, X_mean, X_logσ)`, full rank `Σ p_b`. The blocks are
4099    /// hand-built (no `jacobian_callback`), so without this the flat single-
4100    /// channel audit runs.
4101    ///
4102    /// That flat audit is fatal for a learnable scale: the `log_sigma` block's
4103    /// design is a single structural constant-of-ones column (`eta = design·β`
4104    /// broadcasts the scalar log-σ to every row — see `build_log_sigma_blockspec`).
4105    /// The flat RRQR sees that constant as a repeated shared basis and mistakes
4106    /// it for a cross-block ALIAS with the `mean` intercept, then — because the
4107    /// log-σ block carries the lowest gauge priority — attributes the drop to
4108    /// `log_sigma[0]` and demotes it below the rank tolerance. That both FREEZES
4109    /// the frailty scale at its seed (the only handle on σ is deleted) and leaves
4110    /// the reduced spec width one short of the family's raw joint Hessian, so the
4111    /// outer LAML logdet aborts with `joint exact-newton Hessian validation …:
4112    /// got (p+1)×(p+1), expected p×p`. The mean intercept and log-σ are NOT
4113    /// aliased: they parameterise the frailty MEAN and VARIANCE respectively and
4114    /// are jointly identified; the collinearity is an artefact of auditing a
4115    /// nonlinear scale channel as if it were a linear predictor. Placing each
4116    /// block on its own output channel makes the audit see the genuine
4117    /// block-diagonal structure. Mirrors
4118    /// `SurvivalLocationScaleFamily::output_channel_assignment`.
4119    fn output_channel_assignment(&self, specs: &[ParameterBlockSpec]) -> Option<Vec<usize>> {
4120        Some(
4121            specs
4122                .iter()
4123                .map(|spec| match spec.name.as_str() {
4124                    "time_transform" => 0,
4125                    "mean" => 1,
4126                    "log_sigma" => 2,
4127                    _ => 0,
4128                })
4129                .collect(),
4130        )
4131    }
4132
4133    /// Engage the inner self-vanishing Levenberg–Marquardt μ on a full-rank but
4134    /// indefinite / ill-conditioned penalized joint Hessian, mirroring the
4135    /// sibling [`SurvivalMarginalSlopeFamily`]. Interval-censored rows contribute
4136    /// `ℓ = log[S(L) − S(R)]`, the log of a DIFFERENCE of two survival kernels:
4137    /// unlike the log-concave exact-event / right-censored contributions, its
4138    /// per-row Hessian is legitimately INDEFINITE away from the optimum, so the
4139    /// coupled exact-joint penalized Hessian on the constrained (monotone-cone)
4140    /// time block can be full-rank (`nullity == 0`) yet indefinite or severely
4141    /// ill-conditioned at the cold-start seed. The constrained-QP path already
4142    /// REFLECTS negative-curvature modes to `|λ|` (a convex modified-Newton
4143    /// model), but with this gate OFF it adds NO diagonal floor on a full-rank
4144    /// ill-conditioned reflected model, so the trust-region Newton oscillates on
4145    /// the near-singular mode and stalls out the inner budget before any KKT
4146    /// snapshot is taken ("exited the joint Newton path before convergence — no
4147    /// math snapshot"). Arming the gate adds the SAME self-vanishing μ
4148    /// (∝ the projected KKT residual `‖∇ℓ − Sβ + ∇Φ‖` → 0 at the fixed point) the
4149    /// marginal-slope survival inner relies on, so the step is a well-damped
4150    /// modified-Newton descent that converges, while the converged β̂ is the
4151    /// EXACT unconditioned optimum (μ → 0 there) — zero REML/LAML bias, exact
4152    /// gradient unchanged.
4153    fn levenberg_on_ill_conditioning(&self) -> bool {
4154        true
4155    }
4156
4157    fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
4158        // `evaluate_exact_newton_joint_dense` builds a fully dense joint
4159        // Hessian over (Σ p_b)² across time, mean, and optional log-σ blocks
4160        // via per-row pullback of the latent-survival primary kernel.
4161        crate::custom_family::joint_coupled_coefficient_hessian_cost(
4162            self.event_target.len() as u64,
4163            specs,
4164        )
4165    }
4166
4167    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
4168        let (ll, joint_gradient, hess_time, hess_mean, hess_log_sigma) =
4169            self.evaluate_exact_newton_block_diagonals(block_states)?;
4170        let block_ranges = self.joint_block_ranges();
4171        let mut blockworking_sets = vec![
4172            BlockWorkingSet::ExactNewton {
4173                gradient: joint_gradient.slice(s![block_ranges[0].clone()]).to_owned(),
4174                hessian: SymmetricMatrix::Dense(hess_time),
4175            },
4176            BlockWorkingSet::ExactNewton {
4177                gradient: joint_gradient.slice(s![block_ranges[1].clone()]).to_owned(),
4178                hessian: SymmetricMatrix::Dense(hess_mean),
4179            },
4180        ];
4181        if let (Some(range), Some(hessian)) = (block_ranges.get(2).cloned(), hess_log_sigma) {
4182            blockworking_sets.push(BlockWorkingSet::ExactNewton {
4183                gradient: joint_gradient.slice(s![range]).to_owned(),
4184                hessian: SymmetricMatrix::Dense(hessian),
4185            });
4186        }
4187        Ok(FamilyEvaluation {
4188            log_likelihood: ll,
4189            blockworking_sets,
4190        })
4191    }
4192
4193    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
4194        use rayon::iter::{IntoParallelIterator, ParallelIterator};
4195        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
4196        let q_right = self.time_q_right(block_states)?;
4197        let latent_sd = self.latent_sd(block_states)?;
4198        let n = self.event_target.len();
4199        // Per-row latent-survival jet + log-lik contribution. Independent
4200        // across rows; sum via parallel reduce. `?` propagation happens
4201        // through a Result-collecting fold.
4202        let contributions: Result<Vec<f64>, String> = (0..n)
4203            .into_par_iter()
4204            .map(|i| -> Result<f64, String> {
4205                let wi = self.weights[i];
4206                if wi <= MIN_WEIGHT {
4207                    return Ok(0.0);
4208                }
4209                let row = self.build_row_at(i, q_entry[i], q_exit[i], qdot_exit[i], q_right[i])?;
4210                let jet = LatentSurvivalRowJet::evaluate(&self.quadctx, &row, mu[i], latent_sd)
4211                    .map_err(|e| format!("LatentSurvivalFamily row {i}: {e}"))?;
4212                Ok(wi * jet.log_lik)
4213            })
4214            .collect();
4215        Ok(contributions?.into_iter().sum())
4216    }
4217
4218    fn block_linear_constraints(
4219        &self,
4220        _: &[ParameterBlockState],
4221        block_idx: usize,
4222        block_spec: &ParameterBlockSpec,
4223    ) -> Result<Option<LinearInequalityConstraints>, String> {
4224        assert!(!block_spec.name.is_empty());
4225        if block_idx == Self::BLOCK_TIME {
4226            Ok(self.time_linear_constraints.clone())
4227        } else {
4228            Ok(None)
4229        }
4230    }
4231
4232    fn exact_newton_joint_hessian(
4233        &self,
4234        block_states: &[ParameterBlockState],
4235    ) -> Result<Option<Array2<f64>>, String> {
4236        self.evaluate_exact_newton_joint_dense(block_states)
4237            .map(|(_, _, hessian)| Some(hessian))
4238    }
4239
4240    fn exact_newton_joint_hessian_workspace(
4241        &self,
4242        block_states: &[ParameterBlockState],
4243        _: &[ParameterBlockSpec],
4244    ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
4245        Ok(Some(Arc::new(LatentSurvivalHessianWorkspace::new(
4246            self.clone(),
4247            block_states.to_vec(),
4248        ))))
4249    }
4250
4251    fn exact_newton_joint_gradient_evaluation(
4252        &self,
4253        block_states: &[ParameterBlockState],
4254        _: &[ParameterBlockSpec],
4255    ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
4256        self.evaluate_exact_newton_joint_gradient_dense(block_states)
4257            .map(|(log_likelihood, gradient)| {
4258                Some(ExactNewtonJointGradientEvaluation {
4259                    log_likelihood,
4260                    gradient,
4261                })
4262            })
4263    }
4264
4265    fn exact_newton_joint_hessian_directional_derivative(
4266        &self,
4267        block_states: &[ParameterBlockState],
4268        d_beta_flat: &Array1<f64>,
4269    ) -> Result<Option<Array2<f64>>, String> {
4270        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
4271            .map(Some)
4272    }
4273
4274    fn exact_newton_joint_hessiansecond_directional_derivative(
4275        &self,
4276        block_states: &[ParameterBlockState],
4277        d_beta_u_flat: &Array1<f64>,
4278        d_beta_v_flat: &Array1<f64>,
4279    ) -> Result<Option<Array2<f64>>, String> {
4280        self.exact_newton_joint_hessian_second_directional_derivative_dense(
4281            block_states,
4282            d_beta_u_flat,
4283            d_beta_v_flat,
4284        )
4285        .map(Some)
4286    }
4287
4288    fn requires_joint_outer_hyper_path(&self) -> bool {
4289        true
4290    }
4291}
4292
4293impl CustomFamily for LatentBinaryFamily {
4294    // Latent binary fits have a separation regime; keep the self-limiting
4295    // Jeffreys/Firth curvature active. The trait default flipped to OFF in
4296    // gam#1395 (flat-prior exact-Newton objective); opt back in here.
4297    fn joint_jeffreys_term_required(&self) -> bool {
4298        true
4299    }
4300
4301    fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
4302        true
4303    }
4304
4305    fn has_explicit_joint_hessian(&self) -> bool {
4306        true
4307    }
4308
4309    /// Same self-vanishing Levenberg–Marquardt gate as
4310    /// [`LatentSurvivalFamily`]: the latent-binary deployment shares the
4311    /// constrained (monotone-cone) coupled time block, so a full-rank but
4312    /// ill-conditioned penalized joint Hessian at the cold-start seed must get
4313    /// the self-vanishing μ floor rather than oscillating the constrained-QP
4314    /// trust region into a snapshot-less stall. μ → 0 at the fixed point, so the
4315    /// converged β̂ is exact (no REML/LAML bias).
4316    fn levenberg_on_ill_conditioning(&self) -> bool {
4317        true
4318    }
4319
4320    fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
4321        crate::custom_family::joint_coupled_coefficient_hessian_cost(
4322            self.event_target.len() as u64,
4323            specs,
4324        )
4325    }
4326
4327    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
4328        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
4329        let n = self.event_target.len();
4330        let p_time = self.x_time_exit.ncols();
4331        let p_mean = self.x_mean.ncols();
4332
4333        let mut ll = 0.0;
4334        let mut grad_time = Array1::<f64>::zeros(p_time);
4335        let mut hess_time = Array2::<f64>::zeros((p_time, p_time));
4336        let mut grad_mean = Array1::<f64>::zeros(p_mean);
4337        let mut hess_mean = Array2::<f64>::zeros((p_mean, p_mean));
4338        // Reusable 1-row buffer for x_mean so we avoid allocating a fresh
4339        // Array2<f64> on every iteration via try_row_chunk(i..i+1).
4340        let mut mean_row_buf = Array2::<f64>::zeros((1, p_mean));
4341
4342        for i in 0..n {
4343            let wi = self.weights[i];
4344            if wi <= MIN_WEIGHT {
4345                continue;
4346            }
4347            if !(q_entry[i].is_finite() && q_exit[i].is_finite() && mu[i].is_finite()) {
4348                return Err(format!(
4349                    "latent-binary row {i} contains non-finite predictors: q_entry={}, q_exit={}, mu={}",
4350                    q_entry[i], q_exit[i], mu[i]
4351                ));
4352            }
4353            let row = self.build_right_censored_row_at(i, q_entry[i], q_exit[i])?;
4354            let survival_jet =
4355                LatentSurvivalRowJet::evaluate(&self.quadctx, &row, mu[i], self.latent_sd)
4356                    .map_err(|e| format!("LatentBinaryFamily row {i}: {e}"))?;
4357            let binary = binary_from_log_survival(survival_jet.log_lik, self.event_target[i])?;
4358            ll += wi * binary.log_lik;
4359
4360            self.x_mean
4361                .row_chunk_into(i..i + 1, mean_row_buf.view_mut())
4362                .map_err(|e| format!("LatentBinaryFamily row {i} mean row_chunk: {e}"))?;
4363            let mean_vec = mean_row_buf.row(0);
4364            let mean_grad_scale = wi * binary.grad_scale * survival_jet.score;
4365            for j in 0..p_mean {
4366                grad_mean[j] += mean_grad_scale * mean_vec[j];
4367            }
4368            let mean_neg_hess = wi
4369                * (binary.neg_hess_scale * survival_jet.neg_hessian
4370                    + binary.outer_scale * survival_jet.score * survival_jet.score);
4371            dense_outer_accumulate(&mut hess_mean, mean_neg_hess, mean_vec);
4372
4373            let time_jet =
4374                latent_survival_time_jet(&self.quadctx, &row, 0.0, mu[i], self.latent_sd)?;
4375            let t_entry = self.x_time_entry.row(i);
4376            let t_exit = self.x_time_exit.row(i);
4377            for j in 0..p_time {
4378                grad_time[j] += wi
4379                    * binary.grad_scale
4380                    * (time_jet.grad_entry * t_entry[j] + time_jet.grad_exit * t_exit[j]);
4381            }
4382            dense_outer_accumulate(
4383                &mut hess_time,
4384                wi * binary.neg_hess_scale * time_jet.neg_hess_entry,
4385                t_entry,
4386            );
4387            dense_outer_accumulate(
4388                &mut hess_time,
4389                wi * binary.neg_hess_scale * time_jet.neg_hess_exit,
4390                t_exit,
4391            );
4392            if binary.outer_scale != 0.0 {
4393                dense_outer_accumulate(
4394                    &mut hess_time,
4395                    wi * binary.outer_scale * time_jet.grad_entry * time_jet.grad_entry,
4396                    t_entry,
4397                );
4398                dense_outer_accumulate(
4399                    &mut hess_time,
4400                    wi * binary.outer_scale * time_jet.grad_exit * time_jet.grad_exit,
4401                    t_exit,
4402                );
4403                dense_symmetric_cross_accumulate(
4404                    &mut hess_time,
4405                    wi * binary.outer_scale * time_jet.grad_entry * time_jet.grad_exit,
4406                    t_entry,
4407                    t_exit,
4408                );
4409            }
4410        }
4411
4412        Ok(FamilyEvaluation {
4413            log_likelihood: ll,
4414            blockworking_sets: vec![
4415                BlockWorkingSet::ExactNewton {
4416                    gradient: grad_time,
4417                    hessian: SymmetricMatrix::Dense(hess_time),
4418                },
4419                BlockWorkingSet::ExactNewton {
4420                    gradient: grad_mean,
4421                    hessian: SymmetricMatrix::Dense(hess_mean),
4422                },
4423            ],
4424        })
4425    }
4426
4427    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
4428        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
4429        let mut ll = 0.0;
4430        for i in 0..self.event_target.len() {
4431            let wi = self.weights[i];
4432            if wi <= MIN_WEIGHT {
4433                continue;
4434            }
4435            let row = self.build_right_censored_row_at(i, q_entry[i], q_exit[i])?;
4436            let survival_jet =
4437                LatentSurvivalRowJet::evaluate(&self.quadctx, &row, mu[i], self.latent_sd)
4438                    .map_err(|e| format!("LatentBinaryFamily row {i}: {e}"))?;
4439            ll +=
4440                wi * binary_from_log_survival(survival_jet.log_lik, self.event_target[i])?.log_lik;
4441        }
4442        Ok(ll)
4443    }
4444
4445    fn block_linear_constraints(
4446        &self,
4447        _: &[ParameterBlockState],
4448        block_idx: usize,
4449        block_spec: &ParameterBlockSpec,
4450    ) -> Result<Option<LinearInequalityConstraints>, String> {
4451        assert!(!block_spec.name.is_empty());
4452        if block_idx == Self::BLOCK_TIME {
4453            Ok(self.time_linear_constraints.clone())
4454        } else {
4455            Ok(None)
4456        }
4457    }
4458
4459    fn exact_newton_joint_hessian(
4460        &self,
4461        block_states: &[ParameterBlockState],
4462    ) -> Result<Option<Array2<f64>>, String> {
4463        self.evaluate_exact_newton_joint_dense(block_states)
4464            .map(|(_, _, hessian)| Some(hessian))
4465    }
4466
4467    fn exact_newton_joint_hessian_workspace(
4468        &self,
4469        block_states: &[ParameterBlockState],
4470        _: &[ParameterBlockSpec],
4471    ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
4472        Ok(Some(Arc::new(LatentBinaryHessianWorkspace::new(
4473            self.clone(),
4474            block_states.to_vec(),
4475        ))))
4476    }
4477
4478    fn exact_newton_joint_gradient_evaluation(
4479        &self,
4480        block_states: &[ParameterBlockState],
4481        _: &[ParameterBlockSpec],
4482    ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
4483        self.evaluate_exact_newton_joint_dense(block_states)
4484            .map(|(log_likelihood, gradient, _)| {
4485                Some(ExactNewtonJointGradientEvaluation {
4486                    log_likelihood,
4487                    gradient,
4488                })
4489            })
4490    }
4491
4492    fn exact_newton_joint_hessian_directional_derivative(
4493        &self,
4494        block_states: &[ParameterBlockState],
4495        d_beta_flat: &Array1<f64>,
4496    ) -> Result<Option<Array2<f64>>, String> {
4497        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
4498            .map(Some)
4499    }
4500
4501    fn exact_newton_joint_hessiansecond_directional_derivative(
4502        &self,
4503        block_states: &[ParameterBlockState],
4504        d_beta_u_flat: &Array1<f64>,
4505        d_beta_v_flat: &Array1<f64>,
4506    ) -> Result<Option<Array2<f64>>, String> {
4507        self.exact_newton_joint_hessian_second_directional_derivative_dense(
4508            block_states,
4509            d_beta_u_flat,
4510            d_beta_v_flat,
4511        )
4512        .map(Some)
4513    }
4514
4515    fn requires_joint_outer_hyper_path(&self) -> bool {
4516        true
4517    }
4518}
4519
4520#[cfg(test)]
4521mod tests {
4522    use super::*;
4523    use crate::custom_family::BlockWorkingSet;
4524    use gam_linalg::matrix::DenseDesignMatrix;
4525    use ndarray::array;
4526
4527    fn learnable_sigma_test_family() -> LatentSurvivalFamily {
4528        LatentSurvivalFamily {
4529            event_target: array![1u8, 0u8],
4530            weights: array![1.0, 0.7],
4531            latent_sd_fixed: None,
4532            hazard_loading: HazardLoading::LoadedVsUnloaded,
4533            unloaded_mass_entry: array![0.02, 0.03],
4534            unloaded_mass_exit: array![0.05, 0.08],
4535            unloaded_hazard_exit: array![0.04, 0.0],
4536            x_time_entry: array![[1.0, -0.2], [0.4, 0.7]],
4537            x_time_exit: array![[1.3, 0.1], [0.9, 1.0]],
4538            x_time_derivative_exit: array![[0.8, 0.4], [0.6, 0.5]],
4539            x_time_right: array![[1.3, 0.1], [0.9, 1.0]],
4540            time_offset_right: Array1::zeros(2),
4541            unloaded_mass_right: Array1::zeros(2),
4542            x_mean: DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0, -0.3], [0.2, 0.9]])),
4543            time_linear_constraints: None,
4544            quadctx: Arc::new(QuadratureContext::new()),
4545        }
4546    }
4547
4548    fn learnable_sigma_test_joint_beta() -> Array1<f64> {
4549        array![0.15, 0.25, 0.1, -0.15, 0.35_f64.ln()]
4550    }
4551
4552    /// Regression (frailty scale block deletion): a learnable-σ latent-survival
4553    /// fit routes the pre-fit identifiability audit CHANNEL-AWARE, so the
4554    /// `log_sigma` scale block — realised as a single constant-of-ones column —
4555    /// is never aliased against the `mean` intercept and dropped. Before the
4556    /// `output_channel_assignment` override the family used the trait default
4557    /// (every block → channel 0); the flat single-channel RRQR then saw the two
4558    /// constant columns (mean intercept, log-σ constant) as a cross-block alias,
4559    /// attributed the drop to the lowest-priority `log_sigma` block, and deleted
4560    /// the ONLY handle on the frailty scale — which both froze σ and left the
4561    /// reduced spec width one short of the family's raw joint Hessian, aborting
4562    /// every outer LAML eval with `joint exact-newton Hessian validation … got
4563    /// 15x15, expected 14x14`. The contract this guards: the scale channel must
4564    /// be distinct from the location (mean) channel.
4565    #[test]
4566    fn latent_survival_learnable_sigma_block_lives_on_a_distinct_output_channel() {
4567        let family = learnable_sigma_test_family();
4568        assert!(
4569            family.latent_sd_fixed.is_none(),
4570            "test fixture must be the learnable-σ family"
4571        );
4572
4573        // The three blocks the learnable-σ builder emits, in order. Only the
4574        // block NAME drives `output_channel_assignment`, so borrow the real
4575        // `build_log_sigma_blockspec` shape and relabel for time/mean.
4576        let mut time_spec = build_log_sigma_blockspec(0.5, family.event_target.len());
4577        time_spec.name = "time_transform".to_string();
4578        let mut mean_spec = build_log_sigma_blockspec(0.5, family.event_target.len());
4579        mean_spec.name = "mean".to_string();
4580        let log_sigma_spec = build_log_sigma_blockspec(0.5, family.event_target.len());
4581        assert_eq!(log_sigma_spec.name, "log_sigma");
4582        let specs = vec![time_spec, mean_spec, log_sigma_spec];
4583
4584        let channels = family
4585            .output_channel_assignment(&specs)
4586            .expect("latent survival must declare an explicit channel assignment");
4587        assert_eq!(channels.len(), specs.len());
4588
4589        let (time_ch, mean_ch, log_sigma_ch) = (channels[0], channels[1], channels[2]);
4590        // The load-bearing assertion: the scale channel is NOT the location
4591        // channel, so the channel-aware audit never aliases σ's constant column
4592        // against the mean intercept.
4593        assert_ne!(
4594            log_sigma_ch, mean_ch,
4595            "log_sigma (frailty scale) must not share the mean's output channel, or the \
4596             identifiability audit will alias its constant column against the mean intercept \
4597             and delete the scale parameter"
4598        );
4599        // Each latent-survival block drives a structurally distinct output, so
4600        // all three channels are distinct (block-diagonal true Jacobian).
4601        assert_ne!(time_ch, mean_ch);
4602        assert_ne!(time_ch, log_sigma_ch);
4603        let n_outputs = channels.iter().copied().max().unwrap() + 1;
4604        assert!(
4605            n_outputs >= 3,
4606            "learnable-σ latent survival must expose ≥3 output channels (time, mean, scale), \
4607             got {n_outputs}"
4608        );
4609    }
4610
4611    fn survival_stress_test_family(n: usize) -> LatentSurvivalFamily {
4612        LatentSurvivalFamily {
4613            event_target: Array1::from_iter((0..n).map(|i| if i % 3 == 0 { 1u8 } else { 0u8 })),
4614            weights: Array1::from_iter((0..n).map(|i| 0.55 + 0.03 * ((i % 7) as f64))),
4615            latent_sd_fixed: None,
4616            hazard_loading: HazardLoading::LoadedVsUnloaded,
4617            unloaded_mass_entry: Array1::from_iter(
4618                (0..n).map(|i| 0.015 + 0.0015 * ((i % 11) as f64)),
4619            ),
4620            unloaded_mass_exit: Array1::from_iter((0..n).map(|i| 0.06 + 0.002 * ((i % 13) as f64))),
4621            unloaded_hazard_exit: Array1::from_iter((0..n).map(|i| {
4622                if i % 4 == 0 {
4623                    0.018 + 0.001 * ((i % 5) as f64)
4624                } else {
4625                    0.0
4626                }
4627            })),
4628            x_time_entry: Array2::from_shape_fn((n, 4), |(i, j)| {
4629                0.2 + 0.03 * ((i + 2 * j) % 9) as f64 - if j == 1 { 0.12 } else { 0.0 }
4630            }),
4631            x_time_exit: Array2::from_shape_fn((n, 4), |(i, j)| {
4632                0.35 + 0.025 * ((2 * i + j) % 10) as f64 - if j == 2 { 0.08 } else { 0.0 }
4633            }),
4634            x_time_derivative_exit: Array2::from_shape_fn((n, 4), |(i, j)| {
4635                0.45 + 0.015 * ((i + 3 * j) % 8) as f64
4636            }),
4637            x_time_right: Array2::from_shape_fn((n, 4), |(i, j)| {
4638                0.35 + 0.025 * ((2 * i + j) % 10) as f64 - if j == 2 { 0.08 } else { 0.0 }
4639            }),
4640            time_offset_right: Array1::zeros(n),
4641            unloaded_mass_right: Array1::zeros(n),
4642            x_mean: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::from_shape_fn(
4643                (n, 3),
4644                |(i, j)| 0.1 + 0.04 * ((3 * i + j) % 7) as f64 - if j == 0 { 0.18 } else { 0.0 },
4645            ))),
4646            time_linear_constraints: None,
4647            quadctx: Arc::new(QuadratureContext::new()),
4648        }
4649    }
4650
4651    fn survival_stress_test_joint_beta() -> Array1<f64> {
4652        array![0.18, 0.11, 0.07, 0.13, -0.09, 0.05, 0.12, 0.42_f64.ln()]
4653    }
4654
4655    fn latent_survival_states_from_joint_beta(
4656        family: &LatentSurvivalFamily,
4657        joint_beta: &Array1<f64>,
4658    ) -> Vec<ParameterBlockState> {
4659        let slices = family.joint_slices();
4660        let n = family.event_target.len();
4661        let beta_time = joint_beta.slice(s![slices.time.clone()]).to_owned();
4662        let beta_mean = joint_beta.slice(s![slices.mean.clone()]).to_owned();
4663
4664        let mut eta_time = Array1::<f64>::zeros(3 * n);
4665        eta_time
4666            .slice_mut(s![0..n])
4667            .assign(&gam_linalg::faer_ndarray::fast_av(
4668                &family.x_time_entry,
4669                &beta_time,
4670            ));
4671        eta_time
4672            .slice_mut(s![n..2 * n])
4673            .assign(&gam_linalg::faer_ndarray::fast_av(
4674                &family.x_time_exit,
4675                &beta_time,
4676            ));
4677        eta_time
4678            .slice_mut(s![2 * n..3 * n])
4679            .assign(&gam_linalg::faer_ndarray::fast_av(
4680                &family.x_time_derivative_exit,
4681                &beta_time,
4682            ));
4683
4684        let mut states = vec![
4685            ParameterBlockState {
4686                beta: beta_time,
4687                eta: eta_time,
4688            },
4689            ParameterBlockState {
4690                beta: beta_mean.clone(),
4691                eta: family.x_mean.dot(&beta_mean),
4692            },
4693        ];
4694        if let Some(log_sigma) = slices.log_sigma {
4695            let beta_log_sigma = array![joint_beta[log_sigma.start]];
4696            states.push(ParameterBlockState {
4697                beta: beta_log_sigma.clone(),
4698                eta: beta_log_sigma,
4699            });
4700        }
4701        states
4702    }
4703
4704    fn max_relative_array1(left: &Array1<f64>, right: &Array1<f64>) -> f64 {
4705        left.iter()
4706            .zip(right.iter())
4707            .map(|(l, r)| (l - r).abs() / l.abs().max(r.abs()).max(1e-12))
4708            .fold(0.0_f64, f64::max)
4709    }
4710
4711    fn max_relative_array2(left: &Array2<f64>, right: &Array2<f64>) -> f64 {
4712        left.iter()
4713            .zip(right.iter())
4714            .map(|(l, r)| (l - r).abs() / l.abs().max(r.abs()).max(1e-12))
4715            .fold(0.0_f64, f64::max)
4716    }
4717
4718    fn frobenius_relative_array2(left: &Array2<f64>, right: &Array2<f64>) -> f64 {
4719        let mut diff2 = 0.0_f64;
4720        let mut scale2 = 0.0_f64;
4721        for (l, r) in left.iter().zip(right.iter()) {
4722            let d = l - r;
4723            diff2 += d * d;
4724            scale2 += l * l + r * r;
4725        }
4726        diff2.sqrt() / scale2.sqrt().max(1e-12)
4727    }
4728
4729    fn latent_survival_row_loglik_from_primary(
4730        quadctx: &QuadratureContext,
4731        row: &LatentSurvivalRow,
4732        primary: &Array1<f64>,
4733    ) -> f64 {
4734        let q_entry = primary[LATENT_SURVIVAL_PRIMARY_Q_ENTRY];
4735        let q_exit = primary[LATENT_SURVIVAL_PRIMARY_Q_EXIT];
4736        let qdot_exit = primary[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT];
4737        let q_right = primary[LATENT_SURVIVAL_PRIMARY_Q_RIGHT];
4738        let mu = primary[LATENT_SURVIVAL_PRIMARY_MU];
4739        let sigma = primary[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA].exp();
4740        latent_survival_row_primary_gradient_hessian(
4741            quadctx, row, q_entry, q_exit, qdot_exit, q_right, mu, sigma, true,
4742        )
4743        .expect("row primary evaluation")
4744        .0
4745    }
4746
4747    fn latent_test_specs(n: usize, block_dims: &[(&str, usize)]) -> Vec<ParameterBlockSpec> {
4748        block_dims
4749            .iter()
4750            .map(|(name, p)| ParameterBlockSpec {
4751                name: (*name).to_string(),
4752                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((n, *p)))),
4753                offset: Array1::zeros(n),
4754                penalties: Vec::new(),
4755                nullspace_dims: Vec::new(),
4756                initial_log_lambdas: Array1::zeros(0),
4757                initial_beta: None,
4758                gauge_priority: 100,
4759                jacobian_callback: None,
4760                stacked_design: None,
4761                stacked_offset: None,
4762            })
4763            .collect()
4764    }
4765
4766    fn fixed_sigma_binary_test_family() -> LatentBinaryFamily {
4767        LatentBinaryFamily {
4768            event_target: array![1u8, 0u8],
4769            weights: array![1.0, 0.7],
4770            latent_sd: 0.35,
4771            hazard_loading: HazardLoading::LoadedVsUnloaded,
4772            unloaded_mass_entry: array![0.02, 0.03],
4773            unloaded_mass_exit: array![0.05, 0.08],
4774            x_time_entry: array![[1.0, -0.2], [0.4, 0.7]],
4775            x_time_exit: array![[1.3, 0.1], [0.9, 1.0]],
4776            x_mean: DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0, -0.3], [0.2, 0.9]])),
4777            time_linear_constraints: None,
4778            quadctx: Arc::new(QuadratureContext::new()),
4779        }
4780    }
4781
4782    fn latent_binary_states_from_joint_beta(
4783        family: &LatentBinaryFamily,
4784        joint_beta: &Array1<f64>,
4785    ) -> Vec<ParameterBlockState> {
4786        let slices = family.joint_slices();
4787        let n = family.event_target.len();
4788        let beta_time = joint_beta.slice(s![slices.time.clone()]).to_owned();
4789        let beta_mean = joint_beta.slice(s![slices.mean.clone()]).to_owned();
4790
4791        let mut eta_time = Array1::<f64>::zeros(3 * n);
4792        eta_time
4793            .slice_mut(s![0..n])
4794            .assign(&gam_linalg::faer_ndarray::fast_av(
4795                &family.x_time_entry,
4796                &beta_time,
4797            ));
4798        eta_time
4799            .slice_mut(s![n..2 * n])
4800            .assign(&gam_linalg::faer_ndarray::fast_av(
4801                &family.x_time_exit,
4802                &beta_time,
4803            ));
4804
4805        vec![
4806            ParameterBlockState {
4807                beta: beta_time,
4808                eta: eta_time,
4809            },
4810            ParameterBlockState {
4811                beta: beta_mean.clone(),
4812                eta: family.x_mean.dot(&beta_mean),
4813            },
4814        ]
4815    }
4816
4817    // --- shared latent-interval validation engine: parity / contract tests ---
4818
4819    use crate::survival::location_scale::{TimeBlockInput, TimeBlockMonotonicity};
4820
4821    /// Minimal, structurally valid `TimeBlockInput` for `n` rows and `p_time`
4822    /// columns, used to exercise the shared validation driver without standing
4823    /// up a full term-collection design.
4824    fn validation_time_block(n: usize, p_time: usize) -> TimeBlockInput {
4825        let design = |fill: f64| {
4826            DesignMatrix::Dense(DenseDesignMatrix::from(Array2::from_elem(
4827                (n, p_time),
4828                fill,
4829            )))
4830        };
4831        TimeBlockInput {
4832            design_entry: design(0.1),
4833            design_exit: design(0.2),
4834            design_derivative_exit: design(0.3),
4835            offset_entry: Array1::zeros(n),
4836            offset_exit: Array1::zeros(n),
4837            derivative_offset_exit: Array1::zeros(n),
4838            time_monotonicity: TimeBlockMonotonicity::EnforcedByCoordinateCone,
4839            penalties: Vec::new(),
4840            nullspace_dims: Vec::new(),
4841            initial_log_lambdas: None,
4842            initial_beta: None,
4843        }
4844    }
4845
4846    fn empty_meanspec() -> TermCollectionSpec {
4847        TermCollectionSpec {
4848            linear_terms: Vec::new(),
4849            random_effect_terms: Vec::new(),
4850            smooth_terms: Vec::new(),
4851        }
4852    }
4853
4854    /// A valid two-row latent-survival term spec (one exact event under loaded
4855    /// hazard, one right-censored row).
4856    fn valid_survival_spec(n: usize, p_time: usize) -> LatentSurvivalTermSpec {
4857        LatentSurvivalTermSpec {
4858            age_entry: Array1::zeros(n),
4859            age_exit: Array1::from_elem(n, 1.0),
4860            event_target: Array1::from_shape_fn(n, |i| (i % 2) as u8),
4861            weights: Array1::from_elem(n, 1.0),
4862            derivative_guard: 0.0,
4863            time_block: validation_time_block(n, p_time),
4864            time_design_right: None,
4865            time_offset_right: None,
4866            unloaded_mass_entry: Array1::from_elem(n, 0.01),
4867            unloaded_mass_exit: Array1::from_elem(n, 0.05),
4868            unloaded_mass_right: Array1::zeros(0),
4869            unloaded_hazard_exit: Array1::from_elem(n, 0.02),
4870            meanspec: empty_meanspec(),
4871            mean_offset: Array1::zeros(n),
4872        }
4873    }
4874
4875    /// A valid latent-binary term spec mirroring `valid_survival_spec` but
4876    /// without the per-row unloaded hazard.
4877    fn valid_binary_spec(n: usize, p_time: usize) -> LatentBinaryTermSpec {
4878        LatentBinaryTermSpec {
4879            age_entry: Array1::zeros(n),
4880            age_exit: Array1::from_elem(n, 1.0),
4881            event_target: Array1::from_shape_fn(n, |i| (i % 2) as u8),
4882            weights: Array1::from_elem(n, 1.0),
4883            derivative_guard: 0.0,
4884            time_block: validation_time_block(n, p_time),
4885            unloaded_mass_entry: Array1::from_elem(n, 0.01),
4886            unloaded_mass_exit: Array1::from_elem(n, 0.05),
4887            meanspec: empty_meanspec(),
4888            mean_offset: Array1::zeros(n),
4889        }
4890    }
4891
4892    fn loaded_frailty() -> FrailtySpec {
4893        FrailtySpec::HazardMultiplier {
4894            sigma_fixed: Some(0.3),
4895            loading: HazardLoading::LoadedVsUnloaded,
4896        }
4897    }
4898
4899    /// Both adapters route through the shared `validate_latent_interval_inputs`
4900    /// engine, but each must still emit its own context prefix and (for the
4901    /// size-mismatch / unloaded-decomposition diagnostics) the hazard-aware vs
4902    /// mass-only message variant. This pins the byte-for-byte contract the
4903    /// unification had to preserve, the property the issue's "old vs new
4904    /// validation errors" parity test guards.
4905    #[test]
4906    fn latent_interval_validation_parity_across_models() {
4907        let n = 2;
4908        let p_time = 2;
4909        let data = Array2::<f64>::zeros((n, 3));
4910
4911        // 1. A clean spec validates and round-trips the resolved sigma.
4912        //    Survival keeps the (possibly learnable) Option; binary unwraps to
4913        //    the fixed scalar.
4914        let surv_sigma = validate_latent_survival_inputs(
4915            data.view(),
4916            &valid_survival_spec(n, p_time),
4917            &loaded_frailty(),
4918        )
4919        .expect("valid survival spec must validate");
4920        assert_eq!(surv_sigma, Some(0.3));
4921        let bin_sigma = validate_latent_binary_inputs(
4922            data.view(),
4923            &valid_binary_spec(n, p_time),
4924            &loaded_frailty(),
4925        )
4926        .expect("valid binary spec must validate");
4927        assert_eq!(bin_sigma, 0.3);
4928
4929        // 2. Empty data: shared driver, per-model context prefix.
4930        let empty = Array2::<f64>::zeros((0, 3));
4931        let surv_empty = validate_latent_survival_inputs(
4932            empty.view(),
4933            &valid_survival_spec(n, p_time),
4934            &loaded_frailty(),
4935        )
4936        .expect_err("empty data must be rejected");
4937        assert_eq!(
4938            surv_empty.to_string(),
4939            "latent-survival requires a non-empty dataset"
4940        );
4941        let bin_empty = validate_latent_binary_inputs(
4942            empty.view(),
4943            &valid_binary_spec(n, p_time),
4944            &loaded_frailty(),
4945        )
4946        .expect_err("empty data must be rejected");
4947        assert_eq!(
4948            bin_empty.to_string(),
4949            "latent-binary requires a non-empty dataset"
4950        );
4951
4952        // 3. Size mismatch: survival's message carries `unloaded_hazard=`,
4953        //    binary's does not. This is the one shape that distinguishes the
4954        //    two row views feeding the shared driver.
4955        let mut surv_bad = valid_survival_spec(n, p_time);
4956        surv_bad.weights = Array1::from_elem(n + 1, 1.0);
4957        let surv_size = validate_latent_survival_inputs(data.view(), &surv_bad, &loaded_frailty())
4958            .expect_err("size mismatch must be rejected");
4959        let surv_msg = surv_size.to_string();
4960        assert!(
4961            surv_msg.starts_with("latent-survival size mismatch")
4962                && surv_msg.contains("unloaded_hazard="),
4963            "survival size-mismatch message must include unloaded_hazard: {surv_msg}"
4964        );
4965        let mut bin_bad = valid_binary_spec(n, p_time);
4966        bin_bad.weights = Array1::from_elem(n + 1, 1.0);
4967        let bin_size = validate_latent_binary_inputs(data.view(), &bin_bad, &loaded_frailty())
4968            .expect_err("size mismatch must be rejected");
4969        let bin_msg = bin_size.to_string();
4970        assert!(
4971            bin_msg.starts_with("latent-binary size mismatch")
4972                && !bin_msg.contains("unloaded_hazard"),
4973            "binary size-mismatch message must omit unloaded_hazard: {bin_msg}"
4974        );
4975
4976        // 4. Invalid unloaded decomposition: survival reports `exit_hazard=`,
4977        //    binary reports only the two masses.
4978        let mut surv_neg_hazard = valid_survival_spec(n, p_time);
4979        surv_neg_hazard.unloaded_hazard_exit[0] = -1.0;
4980        let surv_decomp =
4981            validate_latent_survival_inputs(data.view(), &surv_neg_hazard, &loaded_frailty())
4982                .expect_err("negative unloaded hazard must be rejected");
4983        assert_eq!(
4984            surv_decomp.to_string(),
4985            "latent-survival row 1 has invalid unloaded hazard decomposition: entry_mass=0.01, exit_mass=0.05, exit_hazard=-1"
4986        );
4987        let mut bin_bad_mass = valid_binary_spec(n, p_time);
4988        bin_bad_mass.unloaded_mass_exit[0] = 0.0; // exit < entry
4989        let bin_decomp =
4990            validate_latent_binary_inputs(data.view(), &bin_bad_mass, &loaded_frailty())
4991                .expect_err("non-monotone unloaded mass must be rejected");
4992        assert_eq!(
4993            bin_decomp.to_string(),
4994            "latent-binary row 1 has invalid unloaded mass decomposition: entry_mass=0.01, exit_mass=0"
4995        );
4996
4997        // 5. Per-row interval/event/weight diagnostics share one engine, so an
4998        //    identical invalid input yields identical (modulo prefix) text.
4999        let mut surv_event = valid_survival_spec(n, p_time);
5000        surv_event.event_target[1] = 7;
5001        let surv_event_err =
5002            validate_latent_survival_inputs(data.view(), &surv_event, &loaded_frailty())
5003                .expect_err("invalid event target must be rejected");
5004        assert_eq!(
5005            surv_event_err.to_string(),
5006            "latent-survival row 2 has invalid event target 7; expected 0 or 1"
5007        );
5008        let mut bin_event = valid_binary_spec(n, p_time);
5009        bin_event.event_target[1] = 7;
5010        let bin_event_err =
5011            validate_latent_binary_inputs(data.view(), &bin_event, &loaded_frailty())
5012                .expect_err("invalid event target must be rejected");
5013        assert_eq!(
5014            bin_event_err.to_string(),
5015            "latent-binary row 2 has invalid event target 7; expected 0 or 1"
5016        );
5017
5018        // 6. Frailty policy divergence: survival accepts a learnable scale
5019        //    (`sigma_fixed = None` ⇒ `Ok(None)`), binary rejects it.
5020        let learnable = FrailtySpec::HazardMultiplier {
5021            sigma_fixed: None,
5022            loading: HazardLoading::LoadedVsUnloaded,
5023        };
5024        let surv_learnable = validate_latent_survival_inputs(
5025            data.view(),
5026            &valid_survival_spec(n, p_time),
5027            &learnable,
5028        )
5029        .expect("survival accepts a learnable latent scale");
5030        assert_eq!(surv_learnable, None);
5031        let bin_learnable =
5032            validate_latent_binary_inputs(data.view(), &valid_binary_spec(n, p_time), &learnable)
5033                .expect_err("binary requires a fixed latent scale");
5034        assert_eq!(
5035            bin_learnable.to_string(),
5036            "latent-binary currently requires a fixed hazard-multiplier sigma"
5037        );
5038
5039        // 7. The time-block shape check is owned by the shared driver: a
5040        //    column-count mismatch is reported with the per-model prefix.
5041        let mut surv_time_bad = valid_survival_spec(n, p_time);
5042        surv_time_bad.time_block.design_entry = DesignMatrix::Dense(DenseDesignMatrix::from(
5043            Array2::from_elem((n, p_time + 1), 0.1),
5044        ));
5045        let surv_time_err =
5046            validate_latent_survival_inputs(data.view(), &surv_time_bad, &loaded_frailty())
5047                .expect_err("time block column mismatch must be rejected");
5048        assert!(
5049            surv_time_err
5050                .to_string()
5051                .starts_with("latent-survival time block column mismatch"),
5052            "unexpected survival time-block message: {surv_time_err}"
5053        );
5054    }
5055
5056    #[test]
5057    fn latent_survival_coefficient_cost_uses_joint_coupled_formula() {
5058        // `evaluate_exact_newton_joint_dense` builds a fully dense joint
5059        // Hessian over (Σ p_b)² across the time, mean, and log-σ blocks via
5060        // per-row pullback of the latent-survival primary kernel. The override
5061        // must reflect that joint coupling rather than the block-diagonal
5062        // default.
5063        let family = learnable_sigma_test_family();
5064        let n = family.event_target.len() as u64;
5065        let p_time = 2u64;
5066        let p_mean = 2u64;
5067        let p_log_sigma = 1u64;
5068        let specs = vec![
5069            ParameterBlockSpec {
5070                name: "time".to_string(),
5071                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((
5072                    n as usize,
5073                    p_time as usize,
5074                )))),
5075                offset: Array1::zeros(n as usize),
5076                penalties: Vec::new(),
5077                nullspace_dims: Vec::new(),
5078                initial_log_lambdas: Array1::zeros(0),
5079                initial_beta: None,
5080                gauge_priority: 100,
5081                jacobian_callback: None,
5082                stacked_design: None,
5083                stacked_offset: None,
5084            },
5085            ParameterBlockSpec {
5086                name: "mean".to_string(),
5087                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((
5088                    n as usize,
5089                    p_mean as usize,
5090                )))),
5091                offset: Array1::zeros(n as usize),
5092                penalties: Vec::new(),
5093                nullspace_dims: Vec::new(),
5094                initial_log_lambdas: Array1::zeros(0),
5095                initial_beta: None,
5096                gauge_priority: 100,
5097                jacobian_callback: None,
5098                stacked_design: None,
5099                stacked_offset: None,
5100            },
5101            ParameterBlockSpec {
5102                name: "log_sigma".to_string(),
5103                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((
5104                    n as usize,
5105                    p_log_sigma as usize,
5106                )))),
5107                offset: Array1::zeros(n as usize),
5108                penalties: Vec::new(),
5109                nullspace_dims: Vec::new(),
5110                initial_log_lambdas: Array1::zeros(0),
5111                initial_beta: None,
5112                gauge_priority: 100,
5113                jacobian_callback: None,
5114                stacked_design: None,
5115                stacked_offset: None,
5116            },
5117        ];
5118        let p_total = p_time + p_mean + p_log_sigma;
5119        let expected_joint = n * p_total * p_total;
5120        let expected_block_diag =
5121            n * (p_time * p_time + p_mean * p_mean + p_log_sigma * p_log_sigma);
5122        assert_eq!(family.coefficient_hessian_cost(&specs), expected_joint);
5123        // Cross-block fill (time–mean, time–log_sigma, mean–log_sigma) makes
5124        // the joint cost strictly larger than the block-diagonal default.
5125        assert!(expected_joint > expected_block_diag);
5126    }
5127
5128    #[test]
5129    fn latent_family_planner_keeps_outer_hessian_at_large_n() {
5130        use crate::custom_family::custom_family_outer_derivatives;
5131        use gam_problem::{DeclaredHessianForm, Derivative};
5132
5133        let options = BlockwiseFitOptions::default();
5134        let large_n = 50_001;
5135
5136        let survival = learnable_sigma_test_family();
5137        let survival_specs =
5138            latent_test_specs(large_n, &[("time", 2), ("mean", 2), ("log_sigma", 1)]);
5139        let (surv_grad, surv_hess) =
5140            custom_family_outer_derivatives(&survival, &survival_specs, &options);
5141        assert_eq!(surv_grad, Derivative::Analytic);
5142        assert_eq!(surv_hess, DeclaredHessianForm::Either);
5143
5144        let binary = fixed_sigma_binary_test_family();
5145        let binary_specs = latent_test_specs(large_n, &[("time", 2), ("mean", 2)]);
5146        let (bin_grad, bin_hess) =
5147            custom_family_outer_derivatives(&binary, &binary_specs, &options);
5148        assert_eq!(bin_grad, Derivative::Analytic);
5149        assert_eq!(bin_hess, DeclaredHessianForm::Either);
5150    }
5151
5152    #[test]
5153    fn latent_families_arm_self_vanishing_levenberg_on_ill_conditioning() {
5154        // Regression guard for #1108. The interval-censored row contribution
5155        // `ℓ = log[S(L) − S(R)]` is the log of a DIFFERENCE of survival kernels and
5156        // is legitimately NON-concave (indefinite per-row Hessian) away from the
5157        // optimum; on the constrained (monotone-cone) coupled time block this can
5158        // make the penalized joint Hessian full-rank yet indefinite / severely
5159        // ill-conditioned at the cold-start seed. The coupled exact-joint inner
5160        // solver only adds the self-vanishing Levenberg–Marquardt diagonal floor
5161        // (the cure for a full-rank ill-conditioned reflected QP that otherwise
5162        // oscillates the trust region into a snapshot-less stall) when the family
5163        // opts in via `levenberg_on_ill_conditioning()`. Both latent families MUST
5164        // keep this armed (the default is `false`, which leaves the interval inner
5165        // solve diverging with "exited the joint Newton path before convergence").
5166        assert!(
5167            learnable_sigma_test_family().levenberg_on_ill_conditioning(),
5168            "LatentSurvivalFamily must arm the self-vanishing Levenberg floor so the \
5169             indefinite interval-censored joint Hessian converges (see #1108)"
5170        );
5171        assert!(
5172            fixed_sigma_binary_test_family().levenberg_on_ill_conditioning(),
5173            "LatentBinaryFamily must arm the self-vanishing Levenberg floor on its \
5174             constrained coupled time block (see #1108)"
5175        );
5176    }
5177
5178    #[test]
5179    fn latent_binary_exact_joint_hessian_and_workspace_matvec_match_fd() {
5180        let family = fixed_sigma_binary_test_family();
5181        let beta = array![0.15, 0.25, 0.1, -0.15];
5182        let states = latent_binary_states_from_joint_beta(&family, &beta);
5183        let h = 1e-6;
5184
5185        let analytic_hessian = family
5186            .exact_newton_joint_hessian(&states)
5187            .expect("analytic latent binary joint hessian evaluation")
5188            .expect("latent binary should expose exact joint hessian");
5189
5190        for j in 0..beta.len() {
5191            let mut beta_plus = beta.clone();
5192            beta_plus[j] += h;
5193            let gradient_plus = family
5194                .exact_newton_joint_gradient_evaluation(
5195                    &latent_binary_states_from_joint_beta(&family, &beta_plus),
5196                    &[],
5197                )
5198                .expect("joint gradient plus")
5199                .expect("joint gradient should exist")
5200                .gradient;
5201
5202            let mut beta_minus = beta.clone();
5203            beta_minus[j] -= h;
5204            let gradient_minus = family
5205                .exact_newton_joint_gradient_evaluation(
5206                    &latent_binary_states_from_joint_beta(&family, &beta_minus),
5207                    &[],
5208                )
5209                .expect("joint gradient minus")
5210                .expect("joint gradient should exist")
5211                .gradient;
5212
5213            let fd_column = -((&gradient_plus - &gradient_minus) / (2.0 * h));
5214            let analytic_column = analytic_hessian.column(j).to_owned();
5215            let rel = max_relative_array1(&analytic_column, &fd_column);
5216            assert!(
5217                rel < 5e-4,
5218                "latent binary joint Hessian column {j} mismatch: rel={rel}, analytic={analytic_column:?}, fd={fd_column:?}"
5219            );
5220        }
5221
5222        let workspace = family
5223            .exact_newton_joint_hessian_workspace(&states, &[])
5224            .expect("latent binary hessian workspace")
5225            .expect("workspace should exist");
5226        let direction = array![0.4, -0.2, 0.3, 0.1];
5227        let hv = workspace
5228            .hessian_matvec(&direction)
5229            .expect("workspace matvec")
5230            .expect("workspace should support matvec");
5231        let dense_hv = analytic_hessian.dot(&direction);
5232        assert!(
5233            max_relative_array1(&hv, &dense_hv) < 1e-12,
5234            "latent binary workspace HVP mismatch: hv={hv:?}, dense={dense_hv:?}"
5235        );
5236
5237        let dh = workspace
5238            .directional_derivative(&direction)
5239            .expect("workspace dH")
5240            .expect("workspace should support dH");
5241        let fd_step = 1e-5;
5242        let h_plus = family
5243            .exact_newton_joint_hessian(&latent_binary_states_from_joint_beta(
5244                &family,
5245                &(beta.clone() + &(fd_step * &direction)),
5246            ))
5247            .expect("hessian plus")
5248            .expect("hessian plus should exist");
5249        let h_minus = family
5250            .exact_newton_joint_hessian(&latent_binary_states_from_joint_beta(
5251                &family,
5252                &(beta - &(fd_step * &direction)),
5253            ))
5254            .expect("hessian minus")
5255            .expect("hessian minus should exist");
5256        let fd_dh = (&h_plus - &h_minus) / (2.0 * fd_step);
5257        assert!(
5258            max_relative_array2(&dh, &fd_dh) < 2e-4,
5259            "latent binary workspace dH mismatch: dh={dh:?}, fd={fd_dh:?}"
5260        );
5261    }
5262
5263    #[test]
5264    fn latent_survival_learnable_sigma_block_matches_family_fd() {
5265        let family = learnable_sigma_test_family();
5266        let beta = learnable_sigma_test_joint_beta();
5267        let states = latent_survival_states_from_joint_beta(&family, &beta);
5268        let slices = family.joint_slices();
5269        let sigma_idx = slices
5270            .log_sigma
5271            .as_ref()
5272            .expect("learnable sigma test family should expose log_sigma")
5273            .start;
5274        let h = 2e-4;
5275
5276        let eval = family
5277            .evaluate(&states)
5278            .expect("learnable latent survival evaluation");
5279        let joint_gradient = family
5280            .exact_newton_joint_gradient_evaluation(&states, &[])
5281            .expect("joint gradient evaluation")
5282            .expect("joint gradient should exist")
5283            .gradient;
5284        let joint_hessian = family
5285            .exact_newton_joint_hessian(&states)
5286            .expect("joint hessian evaluation")
5287            .expect("joint hessian should exist");
5288        assert_eq!(eval.blockworking_sets.len(), 3);
5289
5290        let (block_grad, block_neg_hess) =
5291            match &eval.blockworking_sets[LatentSurvivalFamily::BLOCK_LOG_SIGMA] {
5292                BlockWorkingSet::ExactNewton { gradient, hessian } => {
5293                    let neg_hess = match hessian {
5294                        SymmetricMatrix::Dense(mat) => mat[[0, 0]],
5295                        _ => panic!("log_sigma block should use a dense exact-Newton Hessian"),
5296                    };
5297                    (gradient[0], neg_hess)
5298                }
5299                _ => panic!("log_sigma block should use ExactNewton"),
5300            };
5301
5302        assert!((block_grad - joint_gradient[sigma_idx]).abs() < 1e-12);
5303        assert!((block_neg_hess - joint_hessian[[sigma_idx, sigma_idx]]).abs() < 1e-12);
5304
5305        let mut beta_plus = beta.clone();
5306        beta_plus[sigma_idx] += h;
5307        let ll_plus = family
5308            .log_likelihood_only(&latent_survival_states_from_joint_beta(&family, &beta_plus))
5309            .expect("ll plus");
5310        let ll_0 = family.log_likelihood_only(&states).expect("ll base");
5311        let mut beta_minus = beta.clone();
5312        beta_minus[sigma_idx] -= h;
5313        let ll_minus = family
5314            .log_likelihood_only(&latent_survival_states_from_joint_beta(
5315                &family,
5316                &beta_minus,
5317            ))
5318            .expect("ll minus");
5319
5320        let fd_grad = (ll_plus - ll_minus) / (2.0 * h);
5321        let fd_neg_hess = -(ll_plus - 2.0 * ll_0 + ll_minus) / (h * h);
5322        assert!(
5323            (joint_gradient[sigma_idx] - fd_grad).abs()
5324                / joint_gradient[sigma_idx]
5325                    .abs()
5326                    .max(fd_grad.abs())
5327                    .max(1e-12)
5328                < 2e-3,
5329            "family log_sigma grad={}, fd={fd_grad}",
5330            joint_gradient[sigma_idx]
5331        );
5332        assert!(
5333            (joint_hessian[[sigma_idx, sigma_idx]] - fd_neg_hess).abs()
5334                / joint_hessian[[sigma_idx, sigma_idx]]
5335                    .abs()
5336                    .max(fd_neg_hess.abs())
5337                    .max(1e-10)
5338                < 2e-2,
5339            "family log_sigma neg_hess={}, fd={fd_neg_hess}",
5340            joint_hessian[[sigma_idx, sigma_idx]]
5341        );
5342    }
5343
5344    #[test]
5345    fn latent_survival_exact_joint_hessian_matches_gradient_fd() {
5346        let family = learnable_sigma_test_family();
5347        let beta = learnable_sigma_test_joint_beta();
5348        let states = latent_survival_states_from_joint_beta(&family, &beta);
5349        let h = 1e-6;
5350
5351        let analytic_hessian = family
5352            .exact_newton_joint_hessian(&states)
5353            .expect("analytic joint hessian evaluation")
5354            .expect("latent survival should expose exact joint hessian");
5355
5356        for j in 0..beta.len() {
5357            let mut beta_plus = beta.clone();
5358            beta_plus[j] += h;
5359            let gradient_plus = family
5360                .exact_newton_joint_gradient_evaluation(
5361                    &latent_survival_states_from_joint_beta(&family, &beta_plus),
5362                    &[],
5363                )
5364                .expect("joint gradient plus")
5365                .expect("joint gradient should exist")
5366                .gradient;
5367
5368            let mut beta_minus = beta.clone();
5369            beta_minus[j] -= h;
5370            let gradient_minus = family
5371                .exact_newton_joint_gradient_evaluation(
5372                    &latent_survival_states_from_joint_beta(&family, &beta_minus),
5373                    &[],
5374                )
5375                .expect("joint gradient minus")
5376                .expect("joint gradient should exist")
5377                .gradient;
5378
5379            let fd_column = (&gradient_plus - &gradient_minus) / (2.0 * h);
5380            let analytic_column = analytic_hessian.column(j).to_owned();
5381            let rel = max_relative_array1(&analytic_column, &(-fd_column));
5382            assert!(
5383                rel < 5e-4,
5384                "joint Hessian column {j} mismatch: rel={rel}, analytic={analytic_column:?}, fd={:?}",
5385                -((&gradient_plus - &gradient_minus) / (2.0 * h))
5386            );
5387        }
5388    }
5389
5390    /// FD check for `LatentSurvivalFamily::offset_channel_residuals`: each
5391    /// channel residual sums to `∂(−ℓ)/∂o_ch` for a uniform additive offset on
5392    /// that time channel (the baseline-θ enters only through these offsets).
5393    /// `o_ch` shifts `eta_time[ch-slice]` uniformly, so `Σ_i r^ch_i` is exactly
5394    /// the directional derivative of `−ℓ` along a constant offset on channel ch.
5395    /// This validates the envelope-theorem latent baseline-θ gradient primitive.
5396    #[test]
5397    fn latent_survival_offset_channel_residuals_match_finite_difference() {
5398        let family = survival_stress_test_family(24);
5399        let beta = survival_stress_test_joint_beta();
5400        let states = latent_survival_states_from_joint_beta(&family, &beta);
5401        let n = family.event_target.len();
5402
5403        let residuals = family
5404            .offset_channel_residuals(&states)
5405            .expect("offset channel residuals");
5406        let sum_entry: f64 = residuals.entry.sum();
5407        let sum_exit: f64 = residuals.exit.sum();
5408        let sum_deriv: f64 = residuals.derivative.sum();
5409
5410        // `−ℓ` after shifting one time channel's eta by a constant δ.
5411        let neg_ll_with_offset = |channel: usize, delta: f64| -> f64 {
5412            let mut shifted = states.clone();
5413            let slice = match channel {
5414                0 => s![0..n],
5415                1 => s![n..2 * n],
5416                2 => s![2 * n..3 * n],
5417                _ => unreachable!(),
5418            };
5419            shifted[LatentSurvivalFamily::BLOCK_TIME]
5420                .eta
5421                .slice_mut(slice)
5422                .mapv_inplace(|v| v + delta);
5423            let (ll, _) = family
5424                .evaluate_exact_newton_joint_gradient_dense(&shifted)
5425                .expect("shifted joint gradient evaluation");
5426            -ll
5427        };
5428
5429        let h = 1e-6;
5430        let fd_entry = (neg_ll_with_offset(0, h) - neg_ll_with_offset(0, -h)) / (2.0 * h);
5431        let fd_exit = (neg_ll_with_offset(1, h) - neg_ll_with_offset(1, -h)) / (2.0 * h);
5432        let fd_deriv = (neg_ll_with_offset(2, h) - neg_ll_with_offset(2, -h)) / (2.0 * h);
5433
5434        assert!(
5435            (sum_entry - fd_entry).abs() <= 1e-5 * fd_entry.abs().max(1.0),
5436            "entry-channel residual sum mismatch: analytic={sum_entry}, fd={fd_entry}"
5437        );
5438        assert!(
5439            (sum_exit - fd_exit).abs() <= 1e-5 * fd_exit.abs().max(1.0),
5440            "exit-channel residual sum mismatch: analytic={sum_exit}, fd={fd_exit}"
5441        );
5442        assert!(
5443            (sum_deriv - fd_deriv).abs() <= 1e-5 * fd_deriv.abs().max(1.0),
5444            "derivative-channel residual sum mismatch: analytic={sum_deriv}, fd={fd_deriv}"
5445        );
5446    }
5447
5448    #[test]
5449    fn latent_survival_exact_joint_parallel_stress_is_repeatable() {
5450        let family = survival_stress_test_family(96);
5451        let beta = survival_stress_test_joint_beta();
5452        let states = latent_survival_states_from_joint_beta(&family, &beta);
5453        let direction_u = array![0.03, -0.02, 0.01, 0.04, -0.015, 0.025, -0.005, 0.02];
5454        let direction_v = array![-0.01, 0.035, -0.025, 0.015, 0.02, -0.01, 0.03, -0.015];
5455
5456        let (ll_a, grad_a) = family
5457            .evaluate_exact_newton_joint_gradient_dense(&states)
5458            .expect("stress joint gradient evaluation");
5459        let (ll_b, grad_b) = family
5460            .evaluate_exact_newton_joint_gradient_dense(&states)
5461            .expect("repeat stress joint gradient evaluation");
5462        assert_eq!(ll_a.to_bits(), ll_b.to_bits());
5463        assert_eq!(grad_a, grad_b);
5464
5465        let (joint_ll_a, joint_grad_a, hess_a) = family
5466            .evaluate_exact_newton_joint_dense(&states)
5467            .expect("stress joint dense evaluation");
5468        let (joint_ll_b, joint_grad_b, hess_b) = family
5469            .evaluate_exact_newton_joint_dense(&states)
5470            .expect("repeat stress joint dense evaluation");
5471        assert_eq!(joint_ll_a.to_bits(), joint_ll_b.to_bits());
5472        assert_eq!(joint_grad_a, joint_grad_b);
5473        assert_eq!(hess_a, hess_b);
5474        assert!(hess_a.iter().all(|value| value.is_finite()));
5475        assert!(max_relative_array2(&hess_a, &hess_a.t().to_owned()) < 1e-12);
5476
5477        let dh_a = family
5478            .exact_newton_joint_hessian_directional_derivative_dense(&states, &direction_u)
5479            .expect("stress joint dH evaluation");
5480        let dh_b = family
5481            .exact_newton_joint_hessian_directional_derivative_dense(&states, &direction_u)
5482            .expect("repeat stress joint dH evaluation");
5483        assert_eq!(dh_a, dh_b);
5484        assert!(dh_a.iter().all(|value| value.is_finite()));
5485        assert!(max_relative_array2(&dh_a, &dh_a.t().to_owned()) < 1e-12);
5486
5487        let d2h_a = family
5488            .exact_newton_joint_hessian_second_directional_derivative_dense(
5489                &states,
5490                &direction_u,
5491                &direction_v,
5492            )
5493            .expect("stress joint d2H evaluation");
5494        let d2h_b = family
5495            .exact_newton_joint_hessian_second_directional_derivative_dense(
5496                &states,
5497                &direction_u,
5498                &direction_v,
5499            )
5500            .expect("repeat stress joint d2H evaluation");
5501        assert_eq!(d2h_a, d2h_b);
5502        assert!(d2h_a.iter().all(|value| value.is_finite()));
5503        assert!(max_relative_array2(&d2h_a, &d2h_a.t().to_owned()) < 1e-12);
5504    }
5505
5506    #[test]
5507    fn latent_survival_exact_joint_dh_matches_hessian_fd() {
5508        let family = learnable_sigma_test_family();
5509        let beta = learnable_sigma_test_joint_beta();
5510        let states = latent_survival_states_from_joint_beta(&family, &beta);
5511        let h = 2e-4;
5512        let direction = array![0.07, -0.03, 0.05, 0.02, -0.04];
5513
5514        let analytic = family
5515            .exact_newton_joint_hessian_directional_derivative(&states, &direction)
5516            .expect("analytic joint dH evaluation")
5517            .expect("latent survival should expose exact joint dH");
5518
5519        let hessian_plus = family
5520            .exact_newton_joint_hessian(&latent_survival_states_from_joint_beta(
5521                &family,
5522                &(beta.clone() + h * &direction),
5523            ))
5524            .expect("joint hessian plus")
5525            .expect("joint hessian should exist");
5526        let hessian_minus = family
5527            .exact_newton_joint_hessian(&latent_survival_states_from_joint_beta(
5528                &family,
5529                &(beta.clone() - h * &direction),
5530            ))
5531            .expect("joint hessian minus")
5532            .expect("joint hessian should exist");
5533
5534        let fd = (&hessian_plus - &hessian_minus) / (2.0 * h);
5535        let rel = frobenius_relative_array2(&analytic, &fd);
5536        assert!(rel < 2e-3, "joint dH mismatch: rel={rel}");
5537    }
5538
5539    #[test]
5540    fn latent_survival_exact_joint_d2h_matches_directional_fd() {
5541        let family = learnable_sigma_test_family();
5542        let beta = learnable_sigma_test_joint_beta();
5543        let states = latent_survival_states_from_joint_beta(&family, &beta);
5544        let h = 5e-4;
5545        let direction_u = array![0.07, -0.03, 0.05, 0.02, -0.04];
5546        let direction_v = array![-0.02, 0.06, -0.01, 0.03, 0.05];
5547
5548        let analytic = family
5549            .exact_newton_joint_hessiansecond_directional_derivative(
5550                &states,
5551                &direction_u,
5552                &direction_v,
5553            )
5554            .expect("analytic joint d2H evaluation")
5555            .expect("latent survival should expose exact joint d2H");
5556        let swapped = family
5557            .exact_newton_joint_hessiansecond_directional_derivative(
5558                &states,
5559                &direction_v,
5560                &direction_u,
5561            )
5562            .expect("swapped analytic joint d2H evaluation")
5563            .expect("latent survival should expose exact joint d2H");
5564        let symmetry_rel = max_relative_array2(&analytic, &swapped);
5565        assert!(
5566            symmetry_rel < 1e-10,
5567            "joint d2H should be symmetric in directions, got rel={symmetry_rel}"
5568        );
5569
5570        let dh_plus = family
5571            .exact_newton_joint_hessian_directional_derivative(
5572                &latent_survival_states_from_joint_beta(
5573                    &family,
5574                    &(beta.clone() + h * &direction_v),
5575                ),
5576                &direction_u,
5577            )
5578            .expect("joint dH plus")
5579            .expect("joint dH should exist");
5580        let dh_minus = family
5581            .exact_newton_joint_hessian_directional_derivative(
5582                &latent_survival_states_from_joint_beta(
5583                    &family,
5584                    &(beta.clone() - h * &direction_v),
5585                ),
5586                &direction_u,
5587            )
5588            .expect("joint dH minus")
5589            .expect("joint dH should exist");
5590
5591        let fd = (&dh_plus - &dh_minus) / (2.0 * h);
5592        let rel = frobenius_relative_array2(&analytic, &fd);
5593        assert!(rel < 2.5e-2, "joint d2H mismatch: rel={rel}");
5594    }
5595
5596    #[test]
5597    fn latent_survival_row_primary_derivatives_match_fd() {
5598        let quadctx = QuadratureContext::new();
5599        let row = LatentSurvivalRow::exact_event(0.35, 1.4, 0.1, 0.45, 0.8, 0.12);
5600        // [q_entry, q_exit, qdot_exit, q_right, mu, log_sigma]. This is an
5601        // exact-event row, so the `q_right` channel is inert (the likelihood
5602        // does not depend on it); the FD loop below confirms its gradient/Hessian
5603        // entries are zero.
5604        let primary = array![
5605            0.35f64.ln(),
5606            1.4f64.ln(),
5607            0.8,
5608            1.6f64.ln(),
5609            -0.2,
5610            0.4f64.ln()
5611        ];
5612        let sigma = primary[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA].exp();
5613        let h_grad = 1e-6;
5614        let h_hess = 2e-4;
5615
5616        let (_, gradient, neg_hessian) = latent_survival_row_primary_gradient_hessian(
5617            &quadctx,
5618            &row,
5619            primary[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
5620            primary[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
5621            primary[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
5622            primary[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
5623            primary[LATENT_SURVIVAL_PRIMARY_MU],
5624            sigma,
5625            true,
5626        )
5627        .expect("analytic row primary gradient/hessian");
5628
5629        for j in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5630            let mut plus = primary.clone();
5631            plus[j] += h_grad;
5632            let mut minus = primary.clone();
5633            minus[j] -= h_grad;
5634            let fd_grad = (latent_survival_row_loglik_from_primary(&quadctx, &row, &plus)
5635                - latent_survival_row_loglik_from_primary(&quadctx, &row, &minus))
5636                / (2.0 * h_grad);
5637            let rel_grad =
5638                (gradient[j] - fd_grad).abs() / gradient[j].abs().max(fd_grad.abs()).max(1e-12);
5639            assert!(
5640                rel_grad < 2e-4,
5641                "row primary grad[{j}] mismatch: analytic={}, fd={fd_grad}, rel={rel_grad}",
5642                gradient[j]
5643            );
5644
5645            for k in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5646                let mut pp = primary.clone();
5647                pp[j] += h_hess;
5648                pp[k] += h_hess;
5649                let mut pm = primary.clone();
5650                pm[j] += h_hess;
5651                pm[k] -= h_hess;
5652                let mut mp = primary.clone();
5653                mp[j] -= h_hess;
5654                mp[k] += h_hess;
5655                let mut mm = primary.clone();
5656                mm[j] -= h_hess;
5657                mm[k] -= h_hess;
5658                let fd_neg_hess = -(latent_survival_row_loglik_from_primary(&quadctx, &row, &pp)
5659                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &pm)
5660                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &mp)
5661                    + latent_survival_row_loglik_from_primary(&quadctx, &row, &mm))
5662                    / (4.0 * h_hess * h_hess);
5663                let analytic = neg_hessian[[j, k]];
5664                let abs_err = (analytic - fd_neg_hess).abs();
5665                let rel = abs_err / analytic.abs().max(fd_neg_hess.abs()).max(1e-10);
5666                assert!(
5667                    abs_err < 2e-5 || rel < 2e-3,
5668                    "row primary neg_hess[{j},{k}] mismatch: analytic={analytic}, fd={fd_neg_hess}, abs_err={abs_err}, rel={rel}"
5669                );
5670            }
5671        }
5672    }
5673
5674    #[test]
5675    fn latent_survival_interval_row_primary_derivatives_match_fd() {
5676        // Interval-censored row jet `ℓ = log[S(L) − S(R)] − log S(entry)`. The
5677        // dynamic two-state numerator differentiates BOTH boundary masses
5678        // `M_L = exp(q_exit)` (left, `q_exit`) and `M_R = exp(q_right)` (right,
5679        // `q_right`) independently — channels that the static
5680        // `LatentSurvivalRowJet::interval_censored` (μ-only) never exercises. This
5681        // FD-verifies the gradient AND neg-Hessian of the interval contribution
5682        // w.r.t. ALL six primary coordinates (q_entry, q_exit/L, qdot_exit,
5683        // q_right/R, mu, log_sigma) on a WELL-POSED bracket where `S(L) − S(R)` is
5684        // comfortably positive (M_L = e^{−0.4} ≈ 0.67 well below M_R = e^{0.5} ≈
5685        // 1.65, so the survival-mass difference is large and the log-of-a-
5686        // difference curvature is well-conditioned).
5687        let quadctx = QuadratureContext::new();
5688        // Bracket masses: entry < L < R with comfortable gaps.
5689        let q_entry = -1.2_f64; // M_entry = e^{−1.2} ≈ 0.30
5690        let q_exit = -0.4_f64; // L: M_L = e^{−0.4} ≈ 0.67
5691        let q_right = 0.5_f64; // R: M_R = e^{0.5} ≈ 1.65 (> M_L)
5692        let mu = -0.15_f64;
5693        let log_sigma = 0.3_f64; // σ ≈ 1.35
5694        // Small, monotone unloaded masses (entry ≤ left ≤ right); qdot is inert
5695        // for the interval contribution.
5696        let row = LatentSurvivalRow::interval_censored(
5697            q_entry.exp(), // mass_entry (consistency only; jet reads q's)
5698            q_exit.exp(),  // mass_left
5699            q_right.exp(), // mass_right
5700            0.01,          // mass_unloaded_entry
5701            0.02,          // mass_unloaded_left
5702            0.05,          // mass_unloaded_right
5703        );
5704        assert!(matches!(
5705            row.event_type,
5706            LatentSurvivalEventType::IntervalCensored
5707        ));
5708
5709        // [q_entry, q_exit/L, qdot_exit, q_right/R, mu, log_sigma]. qdot_exit is
5710        // inert for interval rows (no hazard-derivative channel); the FD loop
5711        // confirms its gradient/Hessian entries are 0.
5712        let primary = array![q_entry, q_exit, 0.7, q_right, mu, log_sigma];
5713        let sigma = primary[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA].exp();
5714        let h_grad = 1e-6;
5715        let h_hess = 2e-4;
5716
5717        let (_, gradient, neg_hessian) = latent_survival_row_primary_gradient_hessian(
5718            &quadctx,
5719            &row,
5720            primary[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
5721            primary[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
5722            primary[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
5723            primary[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
5724            primary[LATENT_SURVIVAL_PRIMARY_MU],
5725            sigma,
5726            true,
5727        )
5728        .expect("analytic interval row primary gradient/hessian");
5729
5730        // The interval contribution must be a positive survival-mass difference
5731        // at this bracket, so the value channel is finite.
5732        let value = latent_survival_row_loglik_from_primary(&quadctx, &row, &primary);
5733        assert!(
5734            value.is_finite(),
5735            "interval row log-likelihood must be finite on a well-posed bracket, got {value}"
5736        );
5737
5738        for j in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5739            let mut plus = primary.clone();
5740            plus[j] += h_grad;
5741            let mut minus = primary.clone();
5742            minus[j] -= h_grad;
5743            let fd_grad = (latent_survival_row_loglik_from_primary(&quadctx, &row, &plus)
5744                - latent_survival_row_loglik_from_primary(&quadctx, &row, &minus))
5745                / (2.0 * h_grad);
5746            let rel_grad =
5747                (gradient[j] - fd_grad).abs() / gradient[j].abs().max(fd_grad.abs()).max(1e-12);
5748            assert!(
5749                rel_grad < 2e-4,
5750                "interval row primary grad[{j}] mismatch: analytic={}, fd={fd_grad}, rel={rel_grad}",
5751                gradient[j]
5752            );
5753
5754            for k in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5755                let mut pp = primary.clone();
5756                pp[j] += h_hess;
5757                pp[k] += h_hess;
5758                let mut pm = primary.clone();
5759                pm[j] += h_hess;
5760                pm[k] -= h_hess;
5761                let mut mp = primary.clone();
5762                mp[j] -= h_hess;
5763                mp[k] += h_hess;
5764                let mut mm = primary.clone();
5765                mm[j] -= h_hess;
5766                mm[k] -= h_hess;
5767                let fd_neg_hess = -(latent_survival_row_loglik_from_primary(&quadctx, &row, &pp)
5768                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &pm)
5769                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &mp)
5770                    + latent_survival_row_loglik_from_primary(&quadctx, &row, &mm))
5771                    / (4.0 * h_hess * h_hess);
5772                let analytic = neg_hessian[[j, k]];
5773                let abs_err = (analytic - fd_neg_hess).abs();
5774                let rel = abs_err / analytic.abs().max(fd_neg_hess.abs()).max(1e-10);
5775                assert!(
5776                    abs_err < 5e-5 || rel < 3e-3,
5777                    "interval row primary neg_hess[{j},{k}] mismatch: analytic={analytic}, fd={fd_neg_hess}, abs_err={abs_err}, rel={rel}"
5778                );
5779            }
5780        }
5781    }
5782}