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