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, ConstraintSet, CustomFamily,
22    ExactNewtonJointGradientEvaluation, ExactNewtonJointHessianWorkspace, FamilyEvaluation,
23    ParameterBlockSpec, ParameterBlockState, PenaltyMatrix, fit_custom_family,
24    fit_custom_family_fixed_log_lambdas,
25};
26use crate::fit_orchestration::drivers::freeze_term_collection_from_design;
27use crate::gamlss::{FamilyMetadata, ParameterLink};
28use crate::model_types::UnifiedFitResult;
29use crate::probability::signed_log_sum_exp;
30use crate::quadrature::{IntegratedExpectationMode, QuadratureContext};
31use crate::sigma_link::{exp_sigma_eta_for_sigma_scalar, exp_sigma_from_eta_scalar};
32use crate::survival::latent::interval::{
33    LatentFrailtyResolution, LatentIntervalModel, LatentIntervalRowView,
34    validate_latent_interval_inputs,
35};
36use crate::survival::location_scale::{
37    TimeBlockInput, project_onto_linear_constraints, structural_time_coefficient_constraints,
38};
39use crate::survival::lognormal_kernel::{
40    FrailtyScale, FrailtySpec, HazardLoading, LatentSurvivalEventType, LatentSurvivalRow,
41    LatentSurvivalRowJet, log_kernel_bundle,
42};
43use gam_linalg::matrix::{DenseDesignMatrix, DesignMatrix, LinearOperator, SymmetricMatrix};
44use gam_math::jet_scalar::{JetScalar, OneSeed, Order2, TwoSeed};
45// `value`/`compose_unary`/… now live on the shared `JetField` base (JetScalar: JetField);
46// the concrete `row_jet.base.value()` reads below need it in scope.
47use gam_math::nested_dual::JetField;
48use gam_solve::pirls::LinearInequalityConstraints;
49use gam_terms::smooth::{TermCollectionDesign, TermCollectionSpec, build_term_collection_design};
50use ndarray::{Array1, Array2, ArrayView1, ArrayView2, s};
51use smallvec::SmallVec;
52use std::sync::Arc;
53
54/// Typed error for the latent-survival / latent-binary family kernels and
55/// their fit-time and per-row validation helpers. Variants pick the semantic
56/// bucket while the inner `reason` carries the original byte-equivalent
57/// message so external callers that previously consumed `String` errors keep
58/// the same diagnostic text via `Display`.
59#[derive(Debug, Clone)]
60pub enum LatentSurvivalError {
61    /// The frailty spec supplied to a latent-survival or latent-binary
62    /// helper is incompatible (wrong variant, missing fixed sigma, non-finite
63    /// or negative fixed sigma).
64    InvalidFrailty { reason: String },
65    /// Per-row dataset validation failed: empty input, size mismatch across
66    /// the spec vectors, or invalid age / event / weight / unloaded-mass
67    /// values for an individual row.
68    InvalidDataset { reason: String },
69    /// A parameter-block state, eta vector, or directional-derivative
70    /// argument supplied to a family entry point has the wrong length.
71    BlockMismatch { reason: String },
72    /// A runtime numerical value (sigma, baseline hazard derivative, kernel
73    /// sum, event probability) became non-finite or out-of-domain.
74    NumericalFailure { reason: String },
75    /// The requested combination of time-block structure or event type is
76    /// not implemented (non-structural monotonicity, interval-censored rows
77    /// on the dynamic-derivative path).
78    UnsupportedConfiguration { reason: String },
79}
80
81impl_reason_error_boilerplate! {
82    LatentSurvivalError {
83        InvalidFrailty,
84        InvalidDataset,
85        BlockMismatch,
86        NumericalFailure,
87        UnsupportedConfiguration,
88    }
89}
90
91impl From<crate::block_layout::block_count::BlockCountMismatch> for LatentSurvivalError {
92    fn from(err: crate::block_layout::block_count::BlockCountMismatch) -> LatentSurvivalError {
93        LatentSurvivalError::BlockMismatch {
94            reason: err.message(),
95        }
96    }
97}
98
99impl From<String> for LatentSurvivalError {
100    /// Inbound conversion for the many `Result<_, String>` helpers this
101    /// module still calls into (term-collection design assembly, dense
102    /// chunk conversion, sparse linear constraints). The text is preserved
103    /// verbatim; we only pick a category so external messages flow through
104    /// `?` without per-callsite `.map_err`.
105    fn from(reason: String) -> LatentSurvivalError {
106        LatentSurvivalError::InvalidDataset { reason }
107    }
108}
109
110/// Reserved [`LatentSurvivalTermSpec::event_target`] code marking an
111/// interval-censored row `(L, R]`. Exact-event codes are `>= 1` and right
112/// censoring is `0`; the interval code is the sentinel `u8::MAX` so it never
113/// collides with an exact-event count and the dispatch is an explicit 3-way map
114/// `{0 → RightCensored, INTERVAL → IntervalCensored, k ≥ 1 → ExactEvent}`.
115pub const LATENT_SURVIVAL_EVENT_INTERVAL: u8 = u8::MAX;
116
117#[inline]
118fn latent_survival_event_type_for(code: u8) -> LatentSurvivalEventType {
119    match code {
120        0 => LatentSurvivalEventType::RightCensored,
121        LATENT_SURVIVAL_EVENT_INTERVAL => LatentSurvivalEventType::IntervalCensored,
122        _ => LatentSurvivalEventType::ExactEvent,
123    }
124}
125
126/// Whole-input proof that likelihood row weights are finite and non-negative.
127///
128/// Construct this before evaluating any response or predictor row. Once
129/// constructed, `weight == 0` is the only dormant-row case; every positive
130/// value, including subnormals, remains part of the likelihood.
131#[derive(Clone, Copy)]
132struct ValidatedLikelihoodWeights<'a> {
133    values: &'a Array1<f64>,
134}
135
136impl<'a> ValidatedLikelihoodWeights<'a> {
137    fn new(values: &'a Array1<f64>, context: &str) -> Result<Self, LatentSurvivalError> {
138        if let Some((row, &weight)) = values
139            .iter()
140            .enumerate()
141            .find(|(_, weight)| !weight.is_finite() || **weight < 0.0)
142        {
143            return Err(LatentSurvivalError::InvalidDataset {
144                reason: format!(
145                    "{context} row {} has invalid likelihood weight {weight:?}; expected finite weight >= 0",
146                    row + 1
147                ),
148            });
149        }
150        Ok(Self { values })
151    }
152
153    #[inline]
154    fn at(self, row: usize) -> f64 {
155        self.values[row]
156    }
157}
158
159/// Multiply one already-finite row quantity by a validated positive weight.
160/// Overflow and non-zero underflow are explicit errors rather than silently
161/// changing the row's contribution. Exact zero row quantities remain zero.
162fn checked_weighted_row_value(
163    weight: f64,
164    value: f64,
165    row: usize,
166    quantity: &str,
167) -> Result<f64, String> {
168    assert!(weight.is_finite() && weight > 0.0);
169    if !value.is_finite() {
170        return Err(format!(
171            "latent likelihood row {} has non-finite unweighted {quantity}: {value:?}",
172            row + 1
173        ));
174    }
175    let weighted = weight * value;
176    if !weighted.is_finite() {
177        return Err(format!(
178            "latent likelihood row {} weighted {quantity} is not representable: {weight:?} * {value:?}",
179            row + 1
180        ));
181    }
182    if value != 0.0 && weighted == 0.0 {
183        return Err(format!(
184            "latent likelihood row {} weighted {quantity} underflowed and is not representable: {weight:?} * {value:?}",
185            row + 1
186        ));
187    }
188    Ok(weighted)
189}
190
191fn checked_weighted_row_matrix(
192    weight: f64,
193    values: &Array2<f64>,
194    row: usize,
195    quantity: &str,
196) -> Result<Array2<f64>, String> {
197    let mut weighted = Array2::<f64>::zeros(values.dim());
198    for ((left, right), &value) in values.indexed_iter() {
199        if !value.is_finite() {
200            return Err(format!(
201                "latent likelihood row {} has non-finite unweighted {quantity}[{left},{right}]: {value:?}",
202                row + 1
203            ));
204        }
205        let product = weight * value;
206        if !product.is_finite() || (value != 0.0 && product == 0.0) {
207            return Err(format!(
208                "latent likelihood row {} weighted {quantity}[{left},{right}] is not representable: {weight:?} * {value:?}",
209                row + 1
210            ));
211        }
212        weighted[[left, right]] = product;
213    }
214    Ok(weighted)
215}
216
217fn require_finite_likelihood_scalar(value: f64, quantity: &str) -> Result<f64, String> {
218    if value.is_finite() {
219        Ok(value)
220    } else {
221        Err(format!(
222            "latent likelihood accumulated {quantity} is not representable: {value:?}"
223        ))
224    }
225}
226
227fn require_finite_likelihood_vector(values: &Array1<f64>, quantity: &str) -> Result<(), String> {
228    if let Some((index, &value)) = values
229        .iter()
230        .enumerate()
231        .find(|(_, value)| !value.is_finite())
232    {
233        return Err(format!(
234            "latent likelihood accumulated {quantity}[{index}] is not representable: {value:?}"
235        ));
236    }
237    Ok(())
238}
239
240fn require_finite_likelihood_matrix(values: &Array2<f64>, quantity: &str) -> Result<(), String> {
241    if let Some(((row, col), &value)) = values.indexed_iter().find(|(_, value)| !value.is_finite())
242    {
243        return Err(format!(
244            "latent likelihood accumulated {quantity}[{row},{col}] is not representable: {value:?}"
245        ));
246    }
247    Ok(())
248}
249
250#[derive(Clone)]
251pub struct LatentSurvivalTermSpec {
252    pub age_entry: Array1<f64>,
253    pub age_exit: Array1<f64>,
254    pub event_target: Array1<u8>,
255    pub weights: Array1<f64>,
256    pub derivative_guard: f64,
257    pub time_block: TimeBlockInput,
258    /// Time-basis design evaluated at the interval upper bound `R` (so
259    /// `q_right = design_right · β_time + offset_right`). `None` when the data
260    /// carries no interval-censored rows; the family then reuses the exit design
261    /// for the unused `q_right` channel. When `Some`, rows whose
262    /// `event_target == LATENT_SURVIVAL_EVENT_INTERVAL` contribute the interval
263    /// likelihood `log[S(L) − S(R)]`.
264    pub time_design_right: Option<DesignMatrix>,
265    pub time_offset_right: Option<Array1<f64>>,
266    pub unloaded_mass_entry: Array1<f64>,
267    pub unloaded_mass_exit: Array1<f64>,
268    /// Unloaded (background) cumulative mass at the interval upper bound `R`.
269    /// Length-`n`; entries for non-interval rows are ignored. Empty/`None`
270    /// folds to zero (full-loading interval rows).
271    pub unloaded_mass_right: Array1<f64>,
272    pub unloaded_hazard_exit: Array1<f64>,
273    pub meanspec: TermCollectionSpec,
274    pub mean_offset: Array1<f64>,
275}
276
277pub struct LatentSurvivalTermFitResult {
278    pub fit: UnifiedFitResult,
279    pub design: TermCollectionDesign,
280    pub resolvedspec: TermCollectionSpec,
281    pub latent_sd: f64,
282    /// Per-row residuals of the unpenalized NLL w.r.t. the additive baseline
283    /// time-block offsets `(entry, exit, derivative)` at the converged β̂.
284    /// Contracted against `baseline_offset_theta_partials` by
285    /// `baseline_chain_rule_gradient` to give the exact θ-gradient of the
286    /// profile penalized NLL for the outer baseline-config optimizer.
287    pub baseline_offset_residuals: crate::survival::OffsetChannelResiduals,
288}
289
290#[derive(Clone)]
291pub struct LatentBinaryTermSpec {
292    pub age_entry: Array1<f64>,
293    pub age_exit: Array1<f64>,
294    pub event_target: Array1<u8>,
295    pub weights: Array1<f64>,
296    pub derivative_guard: f64,
297    pub time_block: TimeBlockInput,
298    pub unloaded_mass_entry: Array1<f64>,
299    pub unloaded_mass_exit: Array1<f64>,
300    pub meanspec: TermCollectionSpec,
301    pub mean_offset: Array1<f64>,
302}
303
304pub struct LatentBinaryTermFitResult {
305    pub fit: UnifiedFitResult,
306    pub design: TermCollectionDesign,
307    pub resolvedspec: TermCollectionSpec,
308    /// Per-row residuals of the unpenalized NLL w.r.t. the additive baseline
309    /// time-block offsets `(entry, exit)` at the converged β̂ (the derivative
310    /// channel is identically zero for the binary deployment likelihood).
311    pub baseline_offset_residuals: crate::survival::OffsetChannelResiduals,
312}
313
314#[derive(Clone)]
315struct PreparedLatentTimeBlock {
316    design_entry: Array2<f64>,
317    design_exit: Array2<f64>,
318    design_derivative_exit: Array2<f64>,
319    /// Dense time-basis design at the interval upper bound `R`. Falls back to a
320    /// clone of `design_exit` when the spec supplies no interval design, so the
321    /// `q_right` channel is always well-defined (and unused for non-interval
322    /// rows).
323    design_right: Array2<f64>,
324    linear_constraints: Option<LinearInequalityConstraints>,
325    penalties: Vec<Array2<f64>>,
326    initial_beta: Option<Array1<f64>>,
327}
328
329#[derive(Clone)]
330pub struct LatentSurvivalFamily {
331    pub event_target: Array1<u8>,
332    pub weights: Array1<f64>,
333    pub latent_sd_fixed: Option<f64>,
334    pub hazard_loading: HazardLoading,
335    pub unloaded_mass_entry: Array1<f64>,
336    pub unloaded_mass_exit: Array1<f64>,
337    pub unloaded_hazard_exit: Array1<f64>,
338    pub x_time_entry: Array2<f64>,
339    pub x_time_exit: Array2<f64>,
340    pub x_time_derivative_exit: Array2<f64>,
341    /// Time-basis design evaluated at the interval upper bound `R` (so
342    /// `q_right = x_time_right · β_time + time_offset_right`). For non-interval
343    /// rows this row equals `x_time_exit`'s row (`q_right` is then unused by the
344    /// likelihood), so the matrix always has `n` rows and the same column count
345    /// as the other time designs.
346    pub x_time_right: Array2<f64>,
347    /// Time-block offset at the interval upper bound `R` (length `n`).
348    pub time_offset_right: Array1<f64>,
349    /// Unloaded (background) cumulative mass at the interval upper bound `R`
350    /// (length `n`). Ignored for non-interval rows.
351    pub unloaded_mass_right: Array1<f64>,
352    pub x_mean: DesignMatrix,
353    pub time_linear_constraints: Option<LinearInequalityConstraints>,
354    pub quadctx: Arc<QuadratureContext>,
355}
356
357#[derive(Clone)]
358pub struct LatentBinaryFamily {
359    pub event_target: Array1<u8>,
360    pub weights: Array1<f64>,
361    pub latent_sd: f64,
362    pub hazard_loading: HazardLoading,
363    pub unloaded_mass_entry: Array1<f64>,
364    pub unloaded_mass_exit: Array1<f64>,
365    pub x_time_entry: Array2<f64>,
366    pub x_time_exit: Array2<f64>,
367    pub x_mean: DesignMatrix,
368    pub time_linear_constraints: Option<LinearInequalityConstraints>,
369    pub quadctx: Arc<QuadratureContext>,
370}
371
372impl LatentSurvivalFamily {
373    pub const BLOCK_TIME: usize = 0;
374    pub const BLOCK_MEAN: usize = 1;
375    pub const BLOCK_LOG_SIGMA: usize = 2;
376
377    pub fn parameter_names() -> &'static [&'static str] {
378        &["time_transform", "mean"]
379    }
380
381    pub fn parameter_links() -> &'static [ParameterLink] {
382        &[ParameterLink::Identity, ParameterLink::Identity]
383    }
384
385    pub fn metadata() -> FamilyMetadata {
386        FamilyMetadata {
387            name: "latent_survival",
388            parameternames: Self::parameter_names(),
389            parameter_links: Self::parameter_links(),
390        }
391    }
392
393    fn split_time_eta<'a>(
394        &self,
395        block_states: &'a [ParameterBlockState],
396    ) -> Result<
397        (
398            ArrayView1<'a, f64>,
399            ArrayView1<'a, f64>,
400            ArrayView1<'a, f64>,
401            &'a Array1<f64>,
402        ),
403        LatentSurvivalError,
404    > {
405        let expected_blocks = if self.latent_sd_fixed.is_some() { 2 } else { 3 };
406        crate::block_layout::block_count::validate_block_count::<LatentSurvivalError>(
407            "LatentSurvivalFamily",
408            expected_blocks,
409            block_states.len(),
410        )?;
411        let n = self.event_target.len();
412        let eta_time = &block_states[Self::BLOCK_TIME].eta;
413        let eta_mean = &block_states[Self::BLOCK_MEAN].eta;
414        if eta_time.len() != 3 * n {
415            return Err(LatentSurvivalError::BlockMismatch {
416                reason: format!(
417                    "latent survival time eta length mismatch: got {}, expected {}",
418                    eta_time.len(),
419                    3 * n
420                ),
421            });
422        }
423        if eta_mean.len() != n || self.weights.len() != n {
424            return Err(LatentSurvivalError::BlockMismatch {
425                reason: "latent survival mean eta dimension mismatch".to_string(),
426            });
427        }
428        Ok((
429            eta_time.slice(s![0..n]),
430            eta_time.slice(s![n..2 * n]),
431            eta_time.slice(s![2 * n..3 * n]),
432            eta_mean,
433        ))
434    }
435
436    /// Per-row interval upper-bound time transform `q_right = x_time_right · β_time
437    /// + time_offset_right`. Shares the time-block coefficients with `q_exit`
438    /// (same monotone basis, evaluated at `R`), so it is read off the time
439    /// block's `beta` rather than carried as an extra eta channel. For
440    /// non-interval rows `x_time_right` equals `x_time_exit`, so the (unused)
441    /// value is simply `q_exit`.
442    fn time_q_right(
443        &self,
444        block_states: &[ParameterBlockState],
445    ) -> Result<Array1<f64>, LatentSurvivalError> {
446        let n = self.event_target.len();
447        let beta_time = &block_states[Self::BLOCK_TIME].beta;
448        if self.x_time_right.ncols() != beta_time.len() {
449            return Err(LatentSurvivalError::BlockMismatch {
450                reason: format!(
451                    "latent survival interval right design has {} columns but time beta has {}",
452                    self.x_time_right.ncols(),
453                    beta_time.len()
454                ),
455            });
456        }
457        if self.x_time_right.nrows() != n || self.time_offset_right.len() != n {
458            return Err(LatentSurvivalError::BlockMismatch {
459                reason: "latent survival interval right design/offset row count mismatch"
460                    .to_string(),
461            });
462        }
463        let mut q_right = self.x_time_right.dot(beta_time);
464        q_right += &self.time_offset_right;
465        Ok(q_right)
466    }
467
468    fn latent_sd(&self, block_states: &[ParameterBlockState]) -> Result<f64, LatentSurvivalError> {
469        if let Some(sigma) = self.latent_sd_fixed {
470            return Ok(sigma);
471        }
472        let eta = *block_states
473            .get(Self::BLOCK_LOG_SIGMA)
474            .and_then(|state| state.eta.get(0))
475            .ok_or_else(|| LatentSurvivalError::BlockMismatch {
476                reason: "latent survival learnable log_sigma block is missing".to_string(),
477            })?;
478        let sigma = exp_sigma_from_eta_scalar(eta);
479        if !(sigma.is_finite() && sigma > 0.0) {
480            return Err(LatentSurvivalError::NumericalFailure {
481                reason: format!(
482                    "latent survival learnable sigma became invalid: log_sigma={eta}, sigma={sigma}"
483                ),
484            });
485        }
486        Ok(sigma)
487    }
488}
489
490impl LatentBinaryFamily {
491    pub const BLOCK_TIME: usize = 0;
492    pub const BLOCK_MEAN: usize = 1;
493
494    fn split_time_eta<'a>(
495        &self,
496        block_states: &'a [ParameterBlockState],
497    ) -> Result<(ArrayView1<'a, f64>, ArrayView1<'a, f64>, &'a Array1<f64>), LatentSurvivalError>
498    {
499        crate::block_layout::block_count::validate_block_count::<LatentSurvivalError>(
500            "LatentBinaryFamily",
501            2,
502            block_states.len(),
503        )?;
504        let n = self.event_target.len();
505        let eta_time = &block_states[Self::BLOCK_TIME].eta;
506        let eta_mean = &block_states[Self::BLOCK_MEAN].eta;
507        if eta_time.len() != 3 * n {
508            return Err(LatentSurvivalError::BlockMismatch {
509                reason: format!(
510                    "latent binary time eta length mismatch: got {}, expected {}",
511                    eta_time.len(),
512                    3 * n
513                ),
514            });
515        }
516        if eta_mean.len() != n || self.weights.len() != n {
517            return Err(LatentSurvivalError::BlockMismatch {
518                reason: "latent binary mean eta dimension mismatch".to_string(),
519            });
520        }
521        Ok((
522            eta_time.slice(s![0..n]),
523            eta_time.slice(s![n..2 * n]),
524            eta_mean,
525        ))
526    }
527}
528
529pub fn fixed_latent_hazard_frailty(
530    frailty: &FrailtySpec,
531    context: &str,
532) -> Result<(f64, HazardLoading), String> {
533    fixed_latent_hazard_frailty_typed(frailty, context).map_err(Into::into)
534}
535
536fn fixed_latent_hazard_frailty_typed(
537    frailty: &FrailtySpec,
538    context: &str,
539) -> Result<(f64, HazardLoading), LatentSurvivalError> {
540    frailty
541        .validate()
542        .map_err(|err| LatentSurvivalError::InvalidFrailty {
543            reason: err.to_string(),
544        })?;
545    match frailty {
546        FrailtySpec::HazardMultiplier {
547            scale: FrailtyScale::Fixed { sigma },
548            loading,
549        } => Ok((*sigma, *loading)),
550        FrailtySpec::HazardMultiplier {
551            scale: FrailtyScale::Learned { .. },
552            ..
553        } => Err(LatentSurvivalError::InvalidFrailty {
554            reason: format!("{context} requires a fixed hazard-multiplier sigma"),
555        }),
556        FrailtySpec::GaussianShift { .. } => Err(LatentSurvivalError::InvalidFrailty {
557            reason: format!("{context} requires HazardMultiplier frailty, not GaussianShift"),
558        }),
559        FrailtySpec::None => Err(LatentSurvivalError::InvalidFrailty {
560            reason: format!("{context} requires a fixed HazardMultiplier frailty specification"),
561        }),
562    }
563}
564
565pub fn latent_hazard_loading(
566    frailty: &FrailtySpec,
567    context: &str,
568) -> Result<HazardLoading, String> {
569    latent_hazard_loading_typed(frailty, context).map_err(Into::into)
570}
571
572fn latent_hazard_loading_typed(
573    frailty: &FrailtySpec,
574    context: &str,
575) -> Result<HazardLoading, LatentSurvivalError> {
576    match frailty {
577        FrailtySpec::HazardMultiplier { loading, .. } => Ok(*loading),
578        FrailtySpec::GaussianShift { .. } => Err(LatentSurvivalError::InvalidFrailty {
579            reason: format!("{context} requires HazardMultiplier frailty, not GaussianShift"),
580        }),
581        FrailtySpec::None => Err(LatentSurvivalError::InvalidFrailty {
582            reason: format!("{context} requires a HazardMultiplier frailty specification"),
583        }),
584    }
585}
586
587#[derive(Clone, Copy)]
588struct LatentSurvivalTimeJet {
589    grad_entry: f64,
590    grad_exit: f64,
591    neg_hess_entry: f64,
592    neg_hess_exit: f64,
593}
594
595pub fn fit_latent_survival_terms(
596    data: ArrayView2<'_, f64>,
597    mut spec: LatentSurvivalTermSpec,
598    frailty: FrailtySpec,
599    options: &BlockwiseFitOptions,
600) -> Result<LatentSurvivalTermFitResult, String> {
601    let frailty_scale = validate_latent_survival_inputs(data, &spec, &frailty)?;
602    // Cover the monotone I-spline baseline's unpenalized affine null direction
603    // with a REML-selected function-space shrinkage ridge, so the interval
604    // warm-start surrogates (whose likelihood has no curvature along that
605    // direction) have a unique MAP instead of refusing. No-op on blocks whose
606    // penalties already span their column space. See the installer's doc.
607    install_latent_time_nullspace_shrinkage_penalty(&mut spec.time_block)?;
608    let (latent_sd, learned_initial_sigma) = match frailty_scale {
609        FrailtyScale::Fixed { sigma } => (Some(sigma), None),
610        FrailtyScale::Learned { initial_sigma } => (None, Some(initial_sigma)),
611    };
612    let hazard_loading = latent_hazard_loading(&frailty, "latent-survival")?;
613    let mean_design =
614        build_term_collection_design(data, &spec.meanspec).map_err(|e| e.to_string())?;
615    let mean_offset = mean_design
616        .compose_offset(spec.mean_offset.view(), "latent-survival mean block")
617        .map_err(|e| e.to_string())?;
618    let resolvedspec = freeze_term_collection_from_design(&spec.meanspec, &mean_design)
619        .map_err(|e| e.to_string())?;
620    let time_prepared = prepare_latent_time_block(
621        &spec.time_block,
622        spec.time_design_right.as_ref(),
623        spec.derivative_guard,
624    )?;
625
626    let n = spec.event_target.len();
627    let time_offset_right = match spec.time_offset_right.as_ref() {
628        Some(offset) => {
629            if offset.len() != n {
630                return Err(format!(
631                    "latent survival interval right time offset must have length {n}, got {}",
632                    offset.len()
633                ));
634            }
635            offset.clone()
636        }
637        None => Array1::zeros(n),
638    };
639    let unloaded_mass_right = if spec.unloaded_mass_right.is_empty() {
640        Array1::zeros(n)
641    } else {
642        if spec.unloaded_mass_right.len() != n {
643            return Err(format!(
644                "latent survival interval right unloaded mass must have length {n}, got {}",
645                spec.unloaded_mass_right.len()
646            ));
647        }
648        spec.unloaded_mass_right.clone()
649    };
650
651    let family = LatentSurvivalFamily {
652        event_target: spec.event_target.clone(),
653        weights: spec.weights.clone(),
654        latent_sd_fixed: latent_sd,
655        hazard_loading,
656        unloaded_mass_entry: spec.unloaded_mass_entry.clone(),
657        unloaded_mass_exit: spec.unloaded_mass_exit.clone(),
658        unloaded_hazard_exit: spec.unloaded_hazard_exit.clone(),
659        x_time_entry: time_prepared.design_entry.clone(),
660        x_time_exit: time_prepared.design_exit.clone(),
661        x_time_derivative_exit: time_prepared.design_derivative_exit.clone(),
662        x_time_right: time_prepared.design_right.clone(),
663        time_offset_right,
664        unloaded_mass_right,
665        x_mean: mean_design.design.clone(),
666        time_linear_constraints: time_prepared.linear_constraints.clone(),
667        quadctx: Arc::new(QuadratureContext::new()),
668    };
669
670    let mut blocks = vec![
671        build_time_blockspec(&time_prepared, &spec.time_block),
672        build_mean_blockspec(&mean_design, mean_offset),
673    ];
674    if let Some(initial_sigma) = learned_initial_sigma {
675        blocks.push(build_log_sigma_blockspec(
676            initial_sigma,
677            mean_design.design.nrows(),
678        ));
679    }
680    // Interval warm start (issue #1108). Interval-censored rows contribute the
681    // NON-concave `ℓ = log[S(L) − S(R)]`; the coupled exact-joint inner Newton
682    // diverges from the cold seed (β_time = 1e-4, σ = 0.5) — the failure surfaces
683    // first as `fit_custom_family`'s outer ρ-seed startup validation rejecting
684    // every seed (`solver_started = 0`). We warm-start from a LOG-CONCAVE
685    // surrogate whose β/σ land in the interval basin, threaded via `initial_beta`
686    // (consumed by every inner solve, including each ρ-seed validation fit).
687    //
688    // Surrogate = right-censored at the bracket LOWER bound `L`. Its survival
689    // mass `S(L) = K_{0,B(L)}` is log-concave (PD Hessian) and — crucially —
690    // its time-block design is the SAME fixed-knot I-spline basis the interval
691    // fit uses, which is FULL RANK regardless of how heavily the inspection-grid
692    // `L` values are TIED (the basis columns are functions of the frozen knots,
693    // not of the observed time multiplicities). Unlike an exact-event surrogate
694    // it imposes NO per-row `q̇(L) > 0` hazard-derivative feasibility condition
695    // (which the tied/degenerate cold-start derivative design can violate), so it
696    // is robust where exact-event-at-L is not. The warm σ then refines from the
697    // bracket-width spread inside the (now in-basin) interval fit.
698    //
699    // Failure is NON-SILENT (#1108): a surrogate that errors or returns a
700    // non-finite / all-zero degenerate β is surfaced as a hard error rather than
701    // silently reverting to the diverging cold start (which masked the real
702    // failure across several attempts). Only `initial_beta` is seeded; the EXACT
703    // interval objective/gradient/Hessian are unchanged, so σ̂ is the true MLE.
704    let has_interval_rows = spec
705        .event_target
706        .iter()
707        .any(|&code| code == LATENT_SURVIVAL_EVENT_INTERVAL);
708    if has_interval_rows {
709        let censored_warm_event_target = spec.event_target.mapv(|code| {
710            if code == LATENT_SURVIVAL_EVENT_INTERVAL {
711                0u8
712            } else {
713                code
714            }
715        });
716        let mut warm_family = family.clone();
717        warm_family.event_target = censored_warm_event_target;
718        // Right-censored-at-L ignores the interval upper bound `R`, so the
719        // (unused) `q_right` channel cannot drift the fit; leaving the right
720        // design/mass in place is harmless (no interval row remains to read it).
721        // Fixed-λ surrogate: no outer smoothing loop runs here, and the inner
722        // solve is convergence-gated inside `fit_custom_family_fixed_log_lambdas`,
723        // so the assembled surrogate fit is (vacuously) outer-converged.
724        let warm_fit_result = fit_custom_family_fixed_log_lambdas(
725            &warm_family,
726            &blocks,
727            options,
728            None,
729        );
730        let warm_fit = match warm_fit_result {
731            Ok(fit) => fit,
732            Err(censored_error) => {
733                let has_finite_event_in_censored_surrogate =
734                    warm_family.event_target.iter().any(|&code| code != 0);
735                if has_finite_event_in_censored_surrogate {
736                    return Err(format!(
737                        "latent interval warm start: right-censored-at-L surrogate fit failed \
738                         (so the interval fit cannot be safely warm-started; this surrogate is \
739                         log-concave and should converge — investigate the surrogate, not the \
740                         interval kernel): {censored_error}"
741                    ));
742                }
743
744                // When every observed row is interval-censored, the
745                // right-censored-at-L surrogate contains no failures at all.
746                // Its likelihood is maximized only on the zero-hazard boundary
747                // (β_time -> -∞), so the fixed-λ Newton solve is correctly
748                // allowed to refuse it even though the objective is concave.
749                // Use the finite lower-endpoint event surrogate solely to obtain
750                // an interior β/σ seed for the exact interval likelihood below;
751                // no fitted surrogate likelihood or derivative is reused.
752                let lower_event_warm_target = spec.event_target.mapv(|code| {
753                    if code == LATENT_SURVIVAL_EVENT_INTERVAL {
754                        1u8
755                    } else {
756                        code
757                    }
758                });
759                let mut event_warm_family = family.clone();
760                event_warm_family.event_target = lower_event_warm_target;
761                fit_custom_family_fixed_log_lambdas(
762                    &event_warm_family,
763                    &blocks,
764                    options,
765                    None,
766                )
767                .map_err(|event_error| {
768                    format!(
769                        "latent interval warm start failed: the right-censored-at-L surrogate \
770                         has no finite failures and refused its boundary optimum ({censored_error}); \
771                         the finite lower-endpoint event surrogate also failed ({event_error})"
772                    )
773                })?
774            }
775        };
776        let warm_beta_usable = warm_fit
777            .block_states
778            .iter()
779            .any(|s| s.beta.iter().all(|v| v.is_finite()) && s.beta.iter().any(|&v| v != 0.0));
780        if !warm_beta_usable {
781            return Err(
782                "latent interval warm start: right-censored-at-L surrogate returned a \
783                 degenerate (non-finite or all-zero) β across every block; the warm start \
784                 cannot seed the interval fit. This indicates the surrogate's time-block \
785                 design is rank-deficient or the inner solve stalled at the seed — \
786                 investigate the surrogate before retrying the interval fit."
787                    .to_string(),
788            );
789        }
790        for (block, state) in blocks.iter_mut().zip(warm_fit.block_states.iter()) {
791            if state.beta.iter().all(|v| v.is_finite()) {
792                block.initial_beta = Some(state.beta.clone());
793            }
794        }
795    }
796    let fit = fit_custom_family(&family, &blocks, options).map_err(|e| e.to_string())?;
797    let latent_sd = family.latent_sd(&fit.block_states)?;
798    let baseline_offset_residuals = family.offset_channel_residuals(&fit.block_states)?;
799    Ok(LatentSurvivalTermFitResult {
800        fit,
801        design: mean_design,
802        resolvedspec,
803        latent_sd,
804        baseline_offset_residuals,
805    })
806}
807
808pub fn fit_latent_binary_terms(
809    data: ArrayView2<'_, f64>,
810    spec: LatentBinaryTermSpec,
811    frailty: FrailtySpec,
812    options: &BlockwiseFitOptions,
813) -> Result<LatentBinaryTermFitResult, String> {
814    let latent_sd = validate_latent_binary_inputs(data, &spec, &frailty)?;
815    let (_, hazard_loading) = fixed_latent_hazard_frailty(&frailty, "latent-binary")?;
816    let mean_design =
817        build_term_collection_design(data, &spec.meanspec).map_err(|e| e.to_string())?;
818    let mean_offset = mean_design
819        .compose_offset(spec.mean_offset.view(), "latent-binary mean block")
820        .map_err(|e| e.to_string())?;
821    let resolvedspec = freeze_term_collection_from_design(&spec.meanspec, &mean_design)
822        .map_err(|e| e.to_string())?;
823    let time_prepared = prepare_latent_time_block(&spec.time_block, None, spec.derivative_guard)?;
824
825    let family = LatentBinaryFamily {
826        event_target: spec.event_target.clone(),
827        weights: spec.weights.clone(),
828        latent_sd,
829        hazard_loading,
830        unloaded_mass_entry: spec.unloaded_mass_entry.clone(),
831        unloaded_mass_exit: spec.unloaded_mass_exit.clone(),
832        x_time_entry: time_prepared.design_entry.clone(),
833        x_time_exit: time_prepared.design_exit.clone(),
834        x_mean: mean_design.design.clone(),
835        time_linear_constraints: time_prepared.linear_constraints.clone(),
836        quadctx: Arc::new(QuadratureContext::new()),
837    };
838
839    let blocks = vec![
840        build_time_blockspec(&time_prepared, &spec.time_block),
841        build_mean_blockspec(&mean_design, mean_offset),
842    ];
843    let fit = fit_custom_family(&family, &blocks, options).map_err(|e| e.to_string())?;
844    let baseline_offset_residuals = family.offset_channel_residuals(&fit.block_states)?;
845    Ok(LatentBinaryTermFitResult {
846        fit,
847        design: mean_design,
848        resolvedspec,
849        baseline_offset_residuals,
850    })
851}
852
853/// Latent-survival adapter for the shared [`LatentIntervalModel`] driver.
854///
855/// Survival permits [`FrailtyScale::Learned`] and carries the
856/// per-row unloaded baseline hazard at exit (which feeds the exact-event
857/// loaded/unloaded split); everything else is validated by the shared engine.
858struct LatentSurvivalModel;
859
860impl LatentIntervalModel for LatentSurvivalModel {
861    fn context() -> &'static str {
862        "latent-survival"
863    }
864
865    fn allows_interval() -> bool {
866        true
867    }
868
869    fn frailty_policy(
870        frailty: &FrailtySpec,
871    ) -> Result<LatentFrailtyResolution, LatentSurvivalError> {
872        frailty
873            .validate()
874            .map_err(|err| LatentSurvivalError::InvalidFrailty {
875                reason: err.to_string(),
876            })?;
877        match frailty {
878            FrailtySpec::HazardMultiplier {
879                scale,
880                loading,
881            } => Ok(LatentFrailtyResolution {
882                scale: *scale,
883                loading: *loading,
884            }),
885            FrailtySpec::GaussianShift { .. } => Err(LatentSurvivalError::InvalidFrailty {
886                reason: "latent-survival requires HazardMultiplier frailty, not GaussianShift"
887                    .to_string(),
888            }),
889            FrailtySpec::None => Err(LatentSurvivalError::InvalidFrailty {
890                reason: "latent-survival requires a HazardMultiplier frailty specification"
891                    .to_string(),
892            }),
893        }
894    }
895}
896
897fn validate_latent_survival_inputs(
898    data: ArrayView2<'_, f64>,
899    spec: &LatentSurvivalTermSpec,
900    frailty: &FrailtySpec,
901) -> Result<FrailtyScale, LatentSurvivalError> {
902    let row = LatentIntervalRowView {
903        frailty,
904        age_entry: &spec.age_entry,
905        age_exit: &spec.age_exit,
906        event_target: &spec.event_target,
907        weights: &spec.weights,
908        unloaded_mass_entry: &spec.unloaded_mass_entry,
909        unloaded_mass_exit: &spec.unloaded_mass_exit,
910        unloaded_hazard_exit: Some(&spec.unloaded_hazard_exit),
911        mean_offset: &spec.mean_offset,
912        derivative_guard: spec.derivative_guard,
913        time_block: &spec.time_block,
914    };
915    validate_latent_interval_inputs::<LatentSurvivalModel>(data, &row)
916}
917
918pub(crate) fn validate_unloaded_components_for_loading(
919    context: &str,
920    row_index: usize,
921    loading: HazardLoading,
922    unloaded_entry: f64,
923    unloaded_exit: f64,
924    unloaded_hazard: Option<f64>,
925) -> Result<(), LatentSurvivalError> {
926    match loading {
927        HazardLoading::Full => {
928            if unloaded_entry != 0.0
929                || unloaded_exit != 0.0
930                || unloaded_hazard.is_some_and(|hazard| hazard != 0.0)
931            {
932                return Err(LatentSurvivalError::InvalidDataset {
933                    reason: format!(
934                        "{context} row {} uses full hazard loading, so unloaded components must be exactly zero; got entry_mass={}, exit_mass={}, exit_hazard={}",
935                        row_index + 1,
936                        unloaded_entry,
937                        unloaded_exit,
938                        unloaded_hazard.unwrap_or(0.0)
939                    ),
940                });
941            }
942        }
943        HazardLoading::LoadedVsUnloaded => {}
944    }
945    Ok(())
946}
947
948/// Latent-binary adapter for the shared [`LatentIntervalModel`] driver.
949///
950/// Binary never evaluates an exact event, so it requires a finite *fixed*
951/// latent sigma (via [`fixed_latent_hazard_frailty_typed`]) and carries no
952/// per-row unloaded hazard; every other invariant is validated by the shared
953/// engine.
954struct LatentBinaryModel;
955
956impl LatentIntervalModel for LatentBinaryModel {
957    fn context() -> &'static str {
958        "latent-binary"
959    }
960
961    fn frailty_policy(
962        frailty: &FrailtySpec,
963    ) -> Result<LatentFrailtyResolution, LatentSurvivalError> {
964        let (sigma, loading) = fixed_latent_hazard_frailty_typed(frailty, "latent-binary")?;
965        Ok(LatentFrailtyResolution {
966            scale: FrailtyScale::Fixed { sigma },
967            loading,
968        })
969    }
970}
971
972fn validate_latent_binary_inputs(
973    data: ArrayView2<'_, f64>,
974    spec: &LatentBinaryTermSpec,
975    frailty: &FrailtySpec,
976) -> Result<f64, LatentSurvivalError> {
977    let row = LatentIntervalRowView {
978        frailty,
979        age_entry: &spec.age_entry,
980        age_exit: &spec.age_exit,
981        event_target: &spec.event_target,
982        weights: &spec.weights,
983        unloaded_mass_entry: &spec.unloaded_mass_entry,
984        unloaded_mass_exit: &spec.unloaded_mass_exit,
985        unloaded_hazard_exit: None,
986        mean_offset: &spec.mean_offset,
987        derivative_guard: spec.derivative_guard,
988        time_block: &spec.time_block,
989    };
990    match validate_latent_interval_inputs::<LatentBinaryModel>(data, &row)? {
991        FrailtyScale::Fixed { sigma } => Ok(sigma),
992        FrailtyScale::Learned { .. } => Err(LatentSurvivalError::InvalidFrailty {
993            reason: "latent-binary requires a fixed latent sigma".to_string(),
994        }),
995    }
996}
997
998fn prepare_latent_time_block(
999    input: &TimeBlockInput,
1000    design_right: Option<&DesignMatrix>,
1001    derivative_guard: f64,
1002) -> Result<PreparedLatentTimeBlock, LatentSurvivalError> {
1003    if !input.time_monotonicity.is_coordinate_cone() {
1004        return Err(LatentSurvivalError::UnsupportedConfiguration {
1005            reason: format!(
1006                "latent survival requires a coordinate-cone monotonicity strategy; got {:?}",
1007                input.time_monotonicity
1008            ),
1009        });
1010    }
1011    let design_entry = input
1012        .design_entry
1013        .try_to_dense_by_chunks("latent survival entry time design")?;
1014    let design_exit = input
1015        .design_exit
1016        .try_to_dense_by_chunks("latent survival exit time design")?;
1017    let design_derivative_exit = input
1018        .design_derivative_exit
1019        .try_to_dense_by_chunks("latent survival derivative time design")?;
1020    // The interval upper-bound design shares the time-block coefficients with
1021    // the exit design; when the data has no interval rows we reuse the exit
1022    // design so `q_right` stays well-defined (its likelihood contribution is
1023    // gated off for non-interval rows). When present it must match the exit
1024    // design's shape (same basis, evaluated at R).
1025    let design_right = match design_right {
1026        Some(matrix) => {
1027            let dense =
1028                matrix.try_to_dense_by_chunks("latent survival interval right time design")?;
1029            if dense.nrows() != design_exit.nrows() || dense.ncols() != design_exit.ncols() {
1030                return Err(LatentSurvivalError::InvalidDataset {
1031                    reason: format!(
1032                        "latent survival interval right time design must match exit design shape \
1033                         {:?}, got {:?}",
1034                        design_exit.dim(),
1035                        dense.dim()
1036                    ),
1037                });
1038            }
1039            dense
1040        }
1041        None => design_exit.clone(),
1042    };
1043    let linear_constraints = structural_time_coefficient_constraints(
1044        &input.design_derivative_exit,
1045        &input.derivative_offset_exit,
1046        derivative_guard,
1047    )?;
1048    let initial_beta = match linear_constraints.as_ref() {
1049        // `project_onto_linear_constraints` validates that any supplied
1050        // `initial_beta` matches `design_exit.ncols()`; surface a mismatch as a
1051        // structured error rather than letting an ndarray broadcast panic
1052        // (issue #374).
1053        Some(constraints) => Some(project_onto_linear_constraints(
1054            design_exit.ncols(),
1055            constraints,
1056            input.initial_beta.as_ref(),
1057        )?),
1058        None => None,
1059    };
1060    Ok(PreparedLatentTimeBlock {
1061        design_entry,
1062        design_exit,
1063        design_derivative_exit,
1064        design_right,
1065        linear_constraints,
1066        penalties: input.penalties.clone(),
1067        initial_beta,
1068    })
1069}
1070
1071fn stack_rows(blocks: &[&Array2<f64>]) -> Array2<f64> {
1072    let ncols = blocks.first().map_or(0, |m| m.ncols());
1073    let nrows = blocks.iter().map(|m| m.nrows()).sum();
1074    let mut out = Array2::<f64>::zeros((nrows, ncols));
1075    let mut row = 0usize;
1076    for block in blocks {
1077        let end = row + block.nrows();
1078        out.slice_mut(s![row..end, ..]).assign(block);
1079        row = end;
1080    }
1081    out
1082}
1083
1084fn build_time_blockspec(
1085    prepared: &PreparedLatentTimeBlock,
1086    input: &TimeBlockInput,
1087) -> ParameterBlockSpec {
1088    // The solver produces a `3·n`-long time `eta` (the `[entry; exit; deriv]`
1089    // channel stack that `split_time_eta` slices). That stacked operator is
1090    // the eta-producing matrix and so belongs in `stacked_design`, paired with
1091    // the matching `3·n`-row stacked offset. The audit / shape-policy invariant
1092    // `design.nrows() == n_obs` is satisfied by exposing the single-channel
1093    // n-row exit design as `design`; the audit never inspects `stacked_design`.
1094    //
1095    // This mirrors the survival location-scale fix for the same #326 class
1096    // (`survival_location_scale.rs`): the previous code put the `3·n`-row
1097    // stack in `design`, which tripped the flat identifiability audit's
1098    // row-equality invariant (`block 1 (mean) has n rows, expected 3n`).
1099    let stacked_design = stack_rows(&[
1100        &prepared.design_entry,
1101        &prepared.design_exit,
1102        &prepared.design_derivative_exit,
1103    ]);
1104    let stacked_offset = gam_linalg::utils::stack_offsets(&[
1105        &input.offset_entry,
1106        &input.offset_exit,
1107        &input.derivative_offset_exit,
1108    ]);
1109    ParameterBlockSpec {
1110        name: "time_transform".to_string(),
1111        design: DesignMatrix::Dense(DenseDesignMatrix::from(Arc::new(
1112            prepared.design_exit.clone(),
1113        ))),
1114        offset: input.offset_exit.clone(),
1115        penalties: prepared
1116            .penalties
1117            .iter()
1118            .cloned()
1119            .map(PenaltyMatrix::Dense)
1120            .collect(),
1121        nullspace_dims: input.nullspace_dims.clone(),
1122        initial_log_lambdas: input
1123            .initial_log_lambdas
1124            .clone()
1125            .unwrap_or_else(|| Array1::zeros(prepared.penalties.len())),
1126        initial_beta: prepared.initial_beta.clone(),
1127        // Canonical-gauge ownership for the latent-survival joint design: the
1128        // time-transform block carries the structural monotone baseline that
1129        // anchors the parameterisation, so it owns any shared constant
1130        // direction (strictly higher than `mean`/`log_sigma` at 100). This
1131        // matches the survival location-scale gauge contract (time highest).
1132        gauge_priority: 200,
1133        jacobian_callback: None,
1134        stacked_design: Some(DesignMatrix::Dense(DenseDesignMatrix::from(Arc::new(
1135            stacked_design,
1136        )))),
1137        stacked_offset: Some(stacked_offset),
1138    }
1139}
1140
1141fn build_mean_blockspec(design: &TermCollectionDesign, offset: Array1<f64>) -> ParameterBlockSpec {
1142    ParameterBlockSpec {
1143        name: "mean".to_string(),
1144        design: design.design.clone(),
1145        offset,
1146        penalties: design.penalties_as_penalty_matrix(),
1147        nullspace_dims: design.nullspace_dims.clone(),
1148        initial_log_lambdas: Array1::zeros(design.penalties.len()),
1149        initial_beta: None,
1150        // Strictly below `time_transform` (200) so any constant direction
1151        // shared between the monotone time baseline and the mean intercept is
1152        // deterministically attributable to the lower-priority `mean` block by
1153        // the canonical-gauge RRQR (the descending-priority contract used by
1154        // survival location-scale; #366/#556 gauge story).
1155        gauge_priority: 150,
1156        jacobian_callback: None,
1157        stacked_design: None,
1158        stacked_offset: None,
1159    }
1160}
1161
1162fn build_log_sigma_blockspec(initial_sigma: f64, n_obs: usize) -> ParameterBlockSpec {
1163    ParameterBlockSpec {
1164        name: "log_sigma".to_string(),
1165        // The frailty/dispersion scale is a single GLOBAL hyperparameter (one free
1166        // coefficient), but the identifiability audit — and the canonical-row
1167        // architecture generally — require every block's effective Jacobian to carry
1168        // `n_obs` rows. A global scalar is realised the same way the survival
1169        // location-scale `log_sigma` block is (see `BinomialLocationScaleFamily`): an
1170        // `n_obs × 1` constant column of ones, so `eta = design · β` is the same scalar
1171        // broadcast to every observation. This keeps it a single free parameter while
1172        // exposing the `n_obs`-row shape the audit checks, and `latent_sd` reads
1173        // `eta[0]` — identical across rows by construction.
1174        design: DesignMatrix::Dense(DenseDesignMatrix::from(Arc::new(Array2::from_elem(
1175            (n_obs, 1),
1176            1.0,
1177        )))),
1178        offset: Array1::zeros(n_obs),
1179        penalties: vec![],
1180        nullspace_dims: vec![],
1181        initial_log_lambdas: Array1::zeros(0),
1182        initial_beta: Some(Array1::from_elem(
1183            1,
1184            exp_sigma_eta_for_sigma_scalar(initial_sigma),
1185        )),
1186        // Lowest of the three (time=200, mean=150): the learnable-scale channel
1187        // yields any shared constant to the location blocks.
1188        gauge_priority: 120,
1189        jacobian_callback: None,
1190        stacked_design: None,
1191        stacked_offset: None,
1192    }
1193}
1194
1195/// Install the Marra & Wood function-space null-space shrinkage penalty (the
1196/// "double penalty") on the latent-survival time block, so the monotone I-spline
1197/// baseline's unpenalized affine null direction carries its own REML-selected
1198/// curvature instead of being left flat.
1199///
1200/// The I-spline value-space penalty `S_I = Lᵀ S_B[1:,1:] L` (see the `ISpline`
1201/// arm of [`crate::survival::construction::build_survival_time_basis`])
1202/// deliberately leaves the affine trend `d(log Λ)/d(log t)` in its 1-D null space
1203/// (`constant γ ↦ affine log Λ ↦ D₂ = 0`, gam#1076): on an event-rich fit the
1204/// likelihood identifies that trend, so penalizing it would only bias the
1205/// baseline. But the interval-censored warm-start evaluates the same direction
1206/// where the surrogate likelihood has no curvature there — the
1207/// right-censored-at-`L` surrogate has no finite failures at all, and the
1208/// finite-lower-endpoint surrogate's grid-tied events leave `Jᵀ W J` near-flat
1209/// along the affine mode — so the MAP is non-unique and the fit refuses
1210/// (`gam_identifiability::check_map_uniqueness`: the affine null direction of
1211/// `Jᵀ W J` carries `nᵀ S n < tol`, dominant block `time_transform`).
1212///
1213/// This is the treatment the survival marginal-slope time block already installs
1214/// (`install_time_nullspace_shrinkage_penalty`): the shared
1215/// [`gam_terms::basis::function_space_nullspace_shrinkage`] builds the
1216/// function-metric ridge `G Z (ZᵀGZ)⁻¹ ZᵀG` (`Z` spanning the primary penalty's
1217/// null space, `G` the endpoint-averaged basis Gram), whose range is exactly that
1218/// null direction, so `nᵀ S n > 0` there. It is a *second* REML coordinate:
1219/// where the affine trend is identified by the data REML drives its λ toward zero
1220/// (no bias — the gam#1076 behaviour is preserved), and where it is not (the
1221/// interval warm-start) the ridge supplies the curvature that makes the MAP
1222/// unique. Returns `Ok(false)` when the block carries no penalty with a null
1223/// space (nothing to cover).
1224fn install_latent_time_nullspace_shrinkage_penalty(
1225    time_block: &mut TimeBlockInput,
1226) -> Result<bool, String> {
1227    let p = time_block.design_exit.ncols();
1228    if p == 0 || time_block.penalties.is_empty() {
1229        return Ok(false);
1230    }
1231    if time_block.nullspace_dims.len() != time_block.penalties.len() {
1232        return Err(format!(
1233            "latent-survival time_block nullspace_dims length {} does not match penalties {}",
1234            time_block.nullspace_dims.len(),
1235            time_block.penalties.len(),
1236        ));
1237    }
1238
1239    // Aggregate the existing wiggliness penalties, each normalized by its own
1240    // max-abs scale so the shared null space (not a scale-weighted average) is
1241    // what the ridge covers. Mirrors the marginal-slope installer.
1242    let mut aggregate = Array2::<f64>::zeros((p, p));
1243    for (idx, penalty) in time_block.penalties.iter().enumerate() {
1244        if penalty.nrows() != p || penalty.ncols() != p {
1245            return Err(format!(
1246                "latent-survival time_block penalty {idx} must be {p}x{p}, got {}x{}",
1247                penalty.nrows(),
1248                penalty.ncols(),
1249            ));
1250        }
1251        let scale = penalty
1252            .iter()
1253            .try_fold(0.0_f64, |acc, &value| {
1254                value.is_finite().then_some(acc.max(value.abs()))
1255            })
1256            .ok_or_else(|| {
1257                format!("latent-survival time_block penalty {idx} contains non-finite values")
1258            })?;
1259        if scale > 0.0 {
1260            ndarray::Zip::from(&mut aggregate)
1261                .and(penalty)
1262                .for_each(|agg, &value| *agg += value / scale);
1263        }
1264    }
1265
1266    // Endpoint-averaged function metric: entry and exit are the two value
1267    // channels through which the baseline enters the survival likelihood, so
1268    // averaging their Grams makes the ridge invariant to whole-sample
1269    // replication and covariant under any coefficient-chart change
1270    // (`G -> MᵀGM`), exactly like the marginal-slope time-block ridge.
1271    if time_block.design_entry.ncols() != p {
1272        return Err(format!(
1273            "latent-survival time_block entry design has {} columns, expected {p}",
1274            time_block.design_entry.ncols(),
1275        ));
1276    }
1277    let entry_mass = time_block.design_entry.nrows();
1278    let exit_mass = time_block.design_exit.nrows();
1279    let total_mass = entry_mass.saturating_add(exit_mass);
1280    if total_mass == 0 {
1281        return Err(
1282            "latent-survival time_block cannot define a function metric from zero endpoint rows"
1283                .to_string(),
1284        );
1285    }
1286    let entry_gram = time_block
1287        .design_entry
1288        .diag_xtw_x(&Array1::ones(entry_mass))
1289        .map_err(|err| format!("latent-survival time_block entry function Gram: {err}"))?;
1290    let exit_gram = time_block
1291        .design_exit
1292        .diag_xtw_x(&Array1::ones(exit_mass))
1293        .map_err(|err| format!("latent-survival time_block exit function Gram: {err}"))?;
1294    let function_gram = (entry_gram + exit_gram).mapv(|value| value / total_mass as f64);
1295
1296    let Some(shrinkage) =
1297        gam_terms::basis::function_space_nullspace_shrinkage(&aggregate, &function_gram)
1298            .map_err(|err| format!("latent-survival time_block nullspace shrinkage: {err}"))?
1299    else {
1300        return Ok(false);
1301    };
1302    if shrinkage.nrows() != p || shrinkage.ncols() != p {
1303        return Err(format!(
1304            "latent-survival time_block nullspace shrinkage penalty must be {p}x{p}, got {}x{}",
1305            shrinkage.nrows(),
1306            shrinkage.ncols(),
1307        ));
1308    }
1309    time_block.penalties.push(shrinkage);
1310    time_block.nullspace_dims.push(0);
1311    // Keep the seed ρ vector consistent with the widened penalty list. The
1312    // latent block builder reads `initial_log_lambdas` directly (unlike the
1313    // marginal-slope path, which rebuilds its seed from the penalty list), so a
1314    // `Some` seed must gain a coordinate for the new null-space penalty; seed it
1315    // at the same smoothing scale as the existing time penalties.
1316    if let Some(seeds) = time_block.initial_log_lambdas.as_mut() {
1317        let seed = seeds.iter().copied().last().unwrap_or(0.0);
1318        let mut widened = seeds.to_vec();
1319        widened.push(seed);
1320        *seeds = Array1::from_vec(widened);
1321    }
1322    Ok(true)
1323}
1324
1325const LATENT_SURVIVAL_PRIMARY_Q_ENTRY: usize = 0;
1326const LATENT_SURVIVAL_PRIMARY_Q_EXIT: usize = 1;
1327const LATENT_SURVIVAL_PRIMARY_QDOT_EXIT: usize = 2;
1328// Interval-censored right boundary R: q_right = log B(R) shares the time-block
1329// coefficients with q_exit (same monotone transform, different time point), so
1330// it is a fourth linear functional of the time block, NOT an independent eta
1331// channel. It sits before `mu`/`log_sigma` so the "trailing optional log_sigma"
1332// invariant used by `active_primary` (= `LATENT_SURVIVAL_PRIMARY_LOG_SIGMA`)
1333// keeps q_right always active.
1334const LATENT_SURVIVAL_PRIMARY_Q_RIGHT: usize = 3;
1335const LATENT_SURVIVAL_PRIMARY_MU: usize = 4;
1336const LATENT_SURVIVAL_PRIMARY_LOG_SIGMA: usize = 5;
1337const LATENT_SURVIVAL_PRIMARY_DIM: usize = 6;
1338
1339/// Derivatives of `log(x)` through 4th order.
1340///
1341/// # Contract
1342///
1343/// `x` must be strictly positive. This function does NOT clamp: a previous
1344/// version replaced `x` by `x.max(1e-300)`, which fabricated enormous finite
1345/// derivatives (`1/1e-300` etc.) that are the derivatives of neither `log(x)`
1346/// nor `log(max(x, floor))` and would silently mask an upstream domain
1347/// failure. Both callers guarantee `x > 0`: one composes at the literal `1.0`
1348/// (the normalised log-sum base); the other passes `base`, which is gated by
1349/// an explicit `base.is_finite() && base > 0.0` check immediately upstream. A
1350/// non-positive `x` therefore never reaches here on any supported path; were
1351/// one to, the function returns the honest IEEE result (`-inf`/`NaN`) —
1352/// identical in debug and release — rather than a finite fabrication. For all
1353/// valid `x > 0` the output is bit-identical to the previous clamped version.
1354#[inline]
1355fn latent_unary_derivatives_log(x: f64) -> [f64; 5] {
1356    let x2 = x * x;
1357    let x3 = x2 * x;
1358    let x4 = x3 * x;
1359    [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
1360}
1361
1362#[derive(Clone, Copy, Debug)]
1363struct LatentKernelPrimaryTerm {
1364    coeff: f64,
1365    q_exp: usize,
1366    qdot_power: usize,
1367    tau_exp: usize,
1368    k: usize,
1369}
1370
1371#[derive(Clone, Copy, Debug)]
1372struct LatentKernelPrimaryDirection {
1373    dq: f64,
1374    dqd: f64,
1375    dmu: f64,
1376    dtau: f64,
1377}
1378
1379#[derive(Clone, Copy, Debug)]
1380struct LatentSurvivalPrimaryDirection {
1381    dq_entry: f64,
1382    dq_exit: f64,
1383    dqdot_exit: f64,
1384    dq_right: f64,
1385    dmu: f64,
1386    dlog_sigma: f64,
1387}
1388
1389#[derive(Clone, Copy, Debug)]
1390struct LatentKernelPrimaryState {
1391    q: f64,
1392    qdot: f64,
1393    mu: f64,
1394    sigma: f64,
1395    log_sigma_factor: f64,
1396}
1397
1398/// Complete primary-coordinate state for one latent-survival row.
1399///
1400/// Keeping these coupled channels together prevents boundary reordering and
1401/// cross-row mean/scale mismatches when selecting a derivative backend.
1402#[derive(Clone, Copy, Debug)]
1403struct LatentSurvivalPrimaryPoint {
1404    q_entry: f64,
1405    q_exit: f64,
1406    qdot_exit: f64,
1407    q_right: f64,
1408    mu: f64,
1409    sigma: f64,
1410}
1411
1412impl LatentSurvivalPrimaryPoint {
1413    /// Log-scale factor used only by derivatives along the learnable-sigma
1414    /// coordinate. Fixed zero frailty has no such coordinate, so its factor is
1415    /// the neutral finite placeholder; every other invalid scale remains NaN
1416    /// or infinite and is rejected by the kernel's numerical checks.
1417    #[inline]
1418    fn log_sigma_factor(self) -> f64 {
1419        if self.sigma == 0.0 {
1420            0.0
1421        } else {
1422            self.sigma.ln()
1423        }
1424    }
1425}
1426
1427#[cfg(test)]
1428mod tests_kernel_recurrence {
1429    use super::*;
1430    use std::collections::BTreeMap;
1431
1432    fn latent_kernel_accumulate_term(
1433        terms: &mut BTreeMap<(usize, usize, usize, usize), f64>,
1434        term: LatentKernelPrimaryTerm,
1435        scale: f64,
1436    ) {
1437        if scale == 0.0 || term.coeff == 0.0 {
1438            return;
1439        }
1440        *terms
1441            .entry((term.q_exp, term.qdot_power, term.tau_exp, term.k))
1442            .or_insert(0.0) += scale * term.coeff;
1443    }
1444
1445    pub(super) fn latent_kernel_differentiate_terms(
1446        terms: &[LatentKernelPrimaryTerm],
1447        dir: LatentKernelPrimaryDirection,
1448    ) -> Vec<LatentKernelPrimaryTerm> {
1449        let mut out = BTreeMap::<(usize, usize, usize, usize), f64>::new();
1450        for term in terms {
1451            if dir.dq != 0.0 {
1452                if term.q_exp > 0 {
1453                    latent_kernel_accumulate_term(&mut out, *term, dir.dq * term.q_exp as f64);
1454                }
1455                latent_kernel_accumulate_term(
1456                    &mut out,
1457                    LatentKernelPrimaryTerm {
1458                        q_exp: term.q_exp + 1,
1459                        k: term.k + 1,
1460                        ..*term
1461                    },
1462                    -dir.dq,
1463                );
1464            }
1465            if dir.dmu != 0.0 {
1466                if term.k > 0 {
1467                    latent_kernel_accumulate_term(&mut out, *term, dir.dmu * term.k as f64);
1468                }
1469                latent_kernel_accumulate_term(
1470                    &mut out,
1471                    LatentKernelPrimaryTerm {
1472                        q_exp: term.q_exp + 1,
1473                        k: term.k + 1,
1474                        ..*term
1475                    },
1476                    -dir.dmu,
1477                );
1478            }
1479            if dir.dtau != 0.0 {
1480                if term.tau_exp > 0 {
1481                    latent_kernel_accumulate_term(&mut out, *term, dir.dtau * term.tau_exp as f64);
1482                }
1483                let kf = term.k as f64;
1484                latent_kernel_accumulate_term(
1485                    &mut out,
1486                    LatentKernelPrimaryTerm {
1487                        tau_exp: term.tau_exp + 2,
1488                        ..*term
1489                    },
1490                    dir.dtau * kf * kf,
1491                );
1492                latent_kernel_accumulate_term(
1493                    &mut out,
1494                    LatentKernelPrimaryTerm {
1495                        q_exp: term.q_exp + 1,
1496                        tau_exp: term.tau_exp + 2,
1497                        k: term.k + 1,
1498                        ..*term
1499                    },
1500                    -dir.dtau * (2.0 * kf + 1.0),
1501                );
1502                latent_kernel_accumulate_term(
1503                    &mut out,
1504                    LatentKernelPrimaryTerm {
1505                        q_exp: term.q_exp + 2,
1506                        tau_exp: term.tau_exp + 2,
1507                        k: term.k + 2,
1508                        ..*term
1509                    },
1510                    dir.dtau,
1511                );
1512            }
1513            if dir.dqd != 0.0 && term.qdot_power > 0 {
1514                latent_kernel_accumulate_term(
1515                    &mut out,
1516                    LatentKernelPrimaryTerm {
1517                        qdot_power: term.qdot_power - 1,
1518                        ..*term
1519                    },
1520                    dir.dqd * term.qdot_power as f64,
1521                );
1522            }
1523        }
1524        out.into_iter()
1525            .filter_map(|((q_exp, qdot_power, tau_exp, k), coeff)| {
1526                (coeff != 0.0).then_some(LatentKernelPrimaryTerm {
1527                    coeff,
1528                    q_exp,
1529                    qdot_power,
1530                    tau_exp,
1531                    k,
1532                })
1533            })
1534            .collect()
1535    }
1536}
1537
1538// Fourth-order latent-kernel recurrences have a small finite support. Keeping
1539// the sorted support inline removes the BTreeMap node allocation and output-Vec
1540// allocation that the pre-cutover directional oracle intentionally retains.
1541// The all-event/K=5/K=6 oracle below asserts this capacity never spills.
1542const LATENT_TERM_INLINE_CAPACITY: usize = 64;
1543type LatentTermBuffer = SmallVec<[LatentKernelPrimaryTerm; LATENT_TERM_INLINE_CAPACITY]>;
1544
1545#[inline]
1546fn latent_kernel_accumulate_term_inline(
1547    terms: &mut LatentTermBuffer,
1548    term: LatentKernelPrimaryTerm,
1549    scale: f64,
1550) {
1551    if scale == 0.0 || term.coeff == 0.0 {
1552        return;
1553    }
1554    let contribution = scale * term.coeff;
1555    if let Some(existing) = terms.iter_mut().find(|existing| {
1556        existing.q_exp == term.q_exp
1557            && existing.qdot_power == term.qdot_power
1558            && existing.tau_exp == term.tau_exp
1559            && existing.k == term.k
1560    }) {
1561        existing.coeff += contribution;
1562    } else {
1563        terms.push(LatentKernelPrimaryTerm {
1564            coeff: contribution,
1565            ..term
1566        });
1567    }
1568}
1569
1570fn latent_kernel_differentiate_terms_inline(
1571    terms: &[LatentKernelPrimaryTerm],
1572    dir: LatentKernelPrimaryDirection,
1573) -> LatentTermBuffer {
1574    let mut out = LatentTermBuffer::new();
1575    for term in terms {
1576        if dir.dq != 0.0 {
1577            if term.q_exp > 0 {
1578                latent_kernel_accumulate_term_inline(&mut out, *term, dir.dq * term.q_exp as f64);
1579            }
1580            latent_kernel_accumulate_term_inline(
1581                &mut out,
1582                LatentKernelPrimaryTerm {
1583                    q_exp: term.q_exp + 1,
1584                    k: term.k + 1,
1585                    ..*term
1586                },
1587                -dir.dq,
1588            );
1589        }
1590        if dir.dmu != 0.0 {
1591            if term.k > 0 {
1592                latent_kernel_accumulate_term_inline(&mut out, *term, dir.dmu * term.k as f64);
1593            }
1594            latent_kernel_accumulate_term_inline(
1595                &mut out,
1596                LatentKernelPrimaryTerm {
1597                    q_exp: term.q_exp + 1,
1598                    k: term.k + 1,
1599                    ..*term
1600                },
1601                -dir.dmu,
1602            );
1603        }
1604        if dir.dtau != 0.0 {
1605            if term.tau_exp > 0 {
1606                latent_kernel_accumulate_term_inline(
1607                    &mut out,
1608                    *term,
1609                    dir.dtau * term.tau_exp as f64,
1610                );
1611            }
1612            let kf = term.k as f64;
1613            latent_kernel_accumulate_term_inline(
1614                &mut out,
1615                LatentKernelPrimaryTerm {
1616                    tau_exp: term.tau_exp + 2,
1617                    ..*term
1618                },
1619                dir.dtau * kf * kf,
1620            );
1621            latent_kernel_accumulate_term_inline(
1622                &mut out,
1623                LatentKernelPrimaryTerm {
1624                    q_exp: term.q_exp + 1,
1625                    tau_exp: term.tau_exp + 2,
1626                    k: term.k + 1,
1627                    ..*term
1628                },
1629                -dir.dtau * (2.0 * kf + 1.0),
1630            );
1631            latent_kernel_accumulate_term_inline(
1632                &mut out,
1633                LatentKernelPrimaryTerm {
1634                    q_exp: term.q_exp + 2,
1635                    tau_exp: term.tau_exp + 2,
1636                    k: term.k + 2,
1637                    ..*term
1638                },
1639                dir.dtau,
1640            );
1641        }
1642        if dir.dqd != 0.0 && term.qdot_power > 0 {
1643            latent_kernel_accumulate_term_inline(
1644                &mut out,
1645                LatentKernelPrimaryTerm {
1646                    qdot_power: term.qdot_power - 1,
1647                    ..*term
1648                },
1649                dir.dqd * term.qdot_power as f64,
1650            );
1651        }
1652    }
1653    out.retain(|term| term.coeff != 0.0);
1654    out.sort_unstable_by_key(|term| (term.q_exp, term.qdot_power, term.tau_exp, term.k));
1655    out
1656}
1657
1658fn latent_kernel_term_sequence_inline(
1659    base_terms: &[LatentKernelPrimaryTerm],
1660    axes: &[LatentKernelPrimaryDirection],
1661    suffix: &[LatentKernelPrimaryDirection],
1662) -> LatentTermBuffer {
1663    let mut terms = LatentTermBuffer::from_slice(base_terms);
1664    terms.retain(|term| term.coeff != 0.0);
1665    // The canonical subset-cache recurrence strips the least-significant
1666    // selected slot and applies it after recursively building the remaining
1667    // mask. Its deterministic floating-point order is therefore highest slot
1668    // to lowest slot. Preserve that order here so the allocation-free packed
1669    // path and the independent MultiDir layout accumulate identical analytic
1670    // coefficients, including cancellation-heavy tail derivatives.
1671    for direction in axes.iter().chain(suffix.iter()).rev() {
1672        terms = latent_kernel_differentiate_terms_inline(&terms, *direction);
1673    }
1674    terms
1675}
1676
1677#[cfg(test)]
1678mod tests_multidir_kernel {
1679    use super::tests_kernel_recurrence::latent_kernel_differentiate_terms;
1680    use super::*;
1681    use gam_math::jet_partitions::MultiDirJet as LatentMultiDirJet;
1682
1683    fn latent_kernel_term_lists_for_directions(
1684        base_terms: &[LatentKernelPrimaryTerm],
1685        directions: &[LatentKernelPrimaryDirection],
1686    ) -> Vec<Vec<LatentKernelPrimaryTerm>> {
1687        fn build_mask(
1688            mask: usize,
1689            base_terms: &[LatentKernelPrimaryTerm],
1690            directions: &[LatentKernelPrimaryDirection],
1691            cache: &mut [Option<Vec<LatentKernelPrimaryTerm>>],
1692        ) -> Vec<LatentKernelPrimaryTerm> {
1693            if let Some(existing) = &cache[mask] {
1694                return existing.clone();
1695            }
1696            let built = if mask == 0 {
1697                base_terms.to_vec()
1698            } else {
1699                let bit = 1usize << mask.trailing_zeros();
1700                let prev = build_mask(mask ^ bit, base_terms, directions, cache);
1701                latent_kernel_differentiate_terms(&prev, directions[bit.trailing_zeros() as usize])
1702            };
1703            cache[mask] = Some(built.clone());
1704            built
1705        }
1706
1707        let mut cache = vec![None; 1usize << directions.len()];
1708        (0..cache.len())
1709            .map(|mask| build_mask(mask, base_terms, directions, &mut cache))
1710            .collect()
1711    }
1712
1713    pub(super) fn latent_kernel_sum_log_jet(
1714        quadctx: &QuadratureContext,
1715        base_terms: &[LatentKernelPrimaryTerm],
1716        state: LatentKernelPrimaryState,
1717        directions: &[LatentKernelPrimaryDirection],
1718        context: &str,
1719    ) -> Result<LatentMultiDirJet, LatentSurvivalError> {
1720        let term_lists = latent_kernel_term_lists_for_directions(base_terms, directions);
1721        let max_k = term_lists
1722            .iter()
1723            .flat_map(|terms| terms.iter().map(|term| term.k))
1724            .max()
1725            .unwrap_or(0);
1726        let bundle = log_kernel_bundle(quadctx, state.q.exp(), state.mu, state.sigma, max_k)
1727            .map_err(|e| LatentSurvivalError::NumericalFailure {
1728                reason: format!("{context} kernel evaluation failed: {e}"),
1729            })?;
1730
1731        let evaluate_terms =
1732            |terms: &[LatentKernelPrimaryTerm]| -> Result<(f64, f64), LatentSurvivalError> {
1733                let mut log_mags = Vec::new();
1734                let mut signs = Vec::new();
1735                for term in terms {
1736                    if term.coeff == 0.0 {
1737                        continue;
1738                    }
1739                    if term.qdot_power > 0 && !(state.qdot.is_finite() && state.qdot > 0.0) {
1740                        return Err(LatentSurvivalError::NumericalFailure {
1741                            reason: format!(
1742                                "{context} requires positive finite qdot for exact-event directional terms, got {}",
1743                                state.qdot
1744                            ),
1745                        });
1746                    }
1747                    let log_qdot = if term.qdot_power > 0 {
1748                        state.qdot.ln()
1749                    } else {
1750                        0.0
1751                    };
1752                    let log_mag = term.coeff.abs().ln()
1753                        + term.q_exp as f64 * state.q
1754                        + term.tau_exp as f64 * state.log_sigma_factor
1755                        + term.qdot_power as f64 * log_qdot
1756                        + bundle.get(term.k);
1757                    log_mags.push(log_mag);
1758                    signs.push(term.coeff.signum());
1759                }
1760                if log_mags.is_empty() {
1761                    return Ok((f64::NEG_INFINITY, 0.0));
1762                }
1763                Ok(signed_log_sum_exp(&log_mags, &signs))
1764            };
1765
1766        let (base_log_sum, base_sign) = evaluate_terms(&term_lists[0])?;
1767        if !(base_log_sum.is_finite() && base_sign > 0.0) {
1768            return Err(LatentSurvivalError::NumericalFailure {
1769                reason: format!("{context} produced a non-positive signed kernel sum"),
1770            });
1771        }
1772
1773        let mut normalized = LatentMultiDirJet::constant(directions.len(), 1.0);
1774        for mask in 1..term_lists.len() {
1775            let (log_abs, sign) = evaluate_terms(&term_lists[mask])?;
1776            normalized.coeffs[mask] = if !log_abs.is_finite() || sign == 0.0 {
1777                0.0
1778            } else {
1779                sign * (log_abs - base_log_sum).exp()
1780            };
1781        }
1782
1783        let mut out = normalized.compose_unary(latent_unary_derivatives_log(1.0));
1784        out.coeffs[0] += base_log_sum;
1785        Ok(out)
1786    }
1787}
1788
1789/// A one-pass analytic lift of a latent kernel sum into an order-specific jet.
1790///
1791/// `suffixes` describes the nilpotent parts carried by the requested scalar:
1792/// `[]` for the ordinary order-two base, `[u]` for the `OneSeed` epsilon part,
1793/// and `[u]`, `[v]`, `[u,v]` for the three non-base `TwoSeed` parts.  For each
1794/// part we differentiate the SAME kernel-term program in the canonical
1795/// highest-slot-to-lowest-slot order used by the pre-cutover `MultiDirJet`
1796/// subset cache. Every requested raw derivative is therefore assembled by the
1797/// same recurrence, accumulation order, and signed-log reduction as its oracle.
1798/// The expensive quadrature bundle is then evaluated ONCE at the maximum `k`
1799/// required by the complete output instead of once per Hessian cell.
1800fn latent_kernel_sum_order2_parts<const K: usize>(
1801    quadctx: &QuadratureContext,
1802    base_terms: &[LatentKernelPrimaryTerm],
1803    state: LatentKernelPrimaryState,
1804    primary_directions: &[LatentKernelPrimaryDirection; K],
1805    suffixes: &[&[LatentKernelPrimaryDirection]],
1806    context: &str,
1807) -> Result<[Order2<K>; 4], LatentSurvivalError> {
1808    assert!(
1809        !suffixes.is_empty() && suffixes.len() <= 4,
1810        "latent kernel lift supports one to four order-two parts"
1811    );
1812    let base_max_k = base_terms.iter().map(|term| term.k).max().unwrap_or(0);
1813    let k_increment = |direction: &LatentKernelPrimaryDirection| {
1814        if direction.dtau != 0.0 {
1815            2
1816        } else if direction.dq != 0.0 || direction.dmu != 0.0 {
1817            1
1818        } else {
1819            0
1820        }
1821    };
1822    // Every emitted part carries a full order-two base.  The largest requested
1823    // recurrence is therefore two primary differentiations plus the largest
1824    // nilpotent suffix.  This exact support bound lets us build the one shared
1825    // bundle before constructing any individual derivative term list.
1826    let max_primary_increment = primary_directions
1827        .iter()
1828        .map(&k_increment)
1829        .max()
1830        .unwrap_or(0);
1831    let max_suffix_increment = suffixes
1832        .iter()
1833        .map(|suffix| suffix.iter().map(&k_increment).sum::<usize>())
1834        .max()
1835        .unwrap_or(0);
1836    let max_k = base_max_k + 2 * max_primary_increment + max_suffix_increment;
1837    let bundle =
1838        log_kernel_bundle(quadctx, state.q.exp(), state.mu, state.sigma, max_k).map_err(|e| {
1839            LatentSurvivalError::NumericalFailure {
1840                reason: format!("{context} kernel evaluation failed: {e}"),
1841            }
1842        })?;
1843
1844    let evaluate_terms =
1845        |terms: &[LatentKernelPrimaryTerm]| -> Result<(f64, f64), LatentSurvivalError> {
1846            let mut log_mags = SmallVec::<[f64; LATENT_TERM_INLINE_CAPACITY]>::new();
1847            let mut signs = SmallVec::<[f64; LATENT_TERM_INLINE_CAPACITY]>::new();
1848            for term in terms {
1849                if term.coeff == 0.0 {
1850                    continue;
1851                }
1852                if term.qdot_power > 0 && !(state.qdot.is_finite() && state.qdot > 0.0) {
1853                    return Err(LatentSurvivalError::NumericalFailure {
1854                        reason: format!(
1855                            "{context} requires positive finite qdot for exact-event directional terms, got {}",
1856                            state.qdot
1857                        ),
1858                    });
1859                }
1860                let log_qdot = if term.qdot_power > 0 {
1861                    state.qdot.ln()
1862                } else {
1863                    0.0
1864                };
1865                log_mags.push(
1866                    term.coeff.abs().ln()
1867                        + term.q_exp as f64 * state.q
1868                        + term.tau_exp as f64 * state.log_sigma_factor
1869                        + term.qdot_power as f64 * log_qdot
1870                        + bundle.get(term.k),
1871                );
1872                signs.push(term.coeff.signum());
1873            }
1874            if log_mags.is_empty() {
1875                return Ok((f64::NEG_INFINITY, 0.0));
1876            }
1877            Ok(signed_log_sum_exp(&log_mags, &signs))
1878        };
1879
1880    let (base_log_sum, base_sign) = evaluate_terms(base_terms)?;
1881    if !(base_log_sum.is_finite() && base_sign > 0.0) {
1882        return Err(LatentSurvivalError::NumericalFailure {
1883            reason: format!("{context} produced a non-positive signed kernel sum"),
1884        });
1885    }
1886    let normalized = |axes: &[LatentKernelPrimaryDirection],
1887                      suffix: &[LatentKernelPrimaryDirection]|
1888     -> Result<f64, LatentSurvivalError> {
1889        let is_zero = |direction: &LatentKernelPrimaryDirection| {
1890            direction.dq == 0.0
1891                && direction.dqd == 0.0
1892                && direction.dmu == 0.0
1893                && direction.dtau == 0.0
1894        };
1895        if axes.iter().chain(suffix.iter()).any(is_zero) {
1896            return Ok(0.0);
1897        }
1898        let terms = latent_kernel_term_sequence_inline(base_terms, axes, suffix);
1899        assert!(
1900            !terms.spilled(),
1901            "latent derivative support exceeded the inline allocation-free capacity: {} > {}",
1902            terms.len(),
1903            LATENT_TERM_INLINE_CAPACITY
1904        );
1905        let (log_abs, sign) = evaluate_terms(&terms)?;
1906        Ok(if !log_abs.is_finite() || sign == 0.0 {
1907            0.0
1908        } else {
1909            sign * (log_abs - base_log_sum).exp()
1910        })
1911    };
1912
1913    let mut parts = [Order2::<K>::constant(0.0); 4];
1914    for (part, suffix) in suffixes.iter().enumerate() {
1915        let value = if part == 0 {
1916            // The base is the kernel divided by itself.  Keep this literal 1.0,
1917            // matching the old `MultiDirJet::constant(..., 1.0)` construction.
1918            1.0
1919        } else {
1920            normalized(&[], suffix)?
1921        };
1922        let mut tower = gam_math::jet_tower::Tower2::<K>::constant(value);
1923        for a in 0..K {
1924            tower.g[a] = normalized(&[primary_directions[a]], suffix)?;
1925        }
1926        for a in 0..K {
1927            for b in a..K {
1928                let derivative =
1929                    normalized(&[primary_directions[a], primary_directions[b]], suffix)?;
1930                tower.h[a][b] = derivative;
1931                tower.h[b][a] = derivative;
1932            }
1933        }
1934        parts[part] = Order2(tower);
1935    }
1936    Ok(latent_kernel_normalized_log_parts(
1937        base_log_sum,
1938        parts,
1939        suffixes.len(),
1940    ))
1941}
1942
1943/// Convert normalized kernel-sum moments into derivatives of the log sum.
1944///
1945/// `normalized_parts` is the single analytic recurrence's compact moment
1946/// layout: the base carries `(1, S_a/S, S_ab/S)`, the one-seed parts carry
1947/// `(S_u/S, S_au/S, S_abu/S)`, and the two-seed cross part carries
1948/// `(S_uv/S, S_auv/S, S_abuv/S)`. For each requested output channel we expose
1949/// those moments as a four-slot derivative table `(a,b,u,v)` and let the shared
1950/// compensated truncated-Taylor reducer perform `log` composition. This avoids
1951/// the severe tail cancellation caused by composing projected `Order2` products
1952/// in stages, while retaining one recurrence, one moment layout, and one general
1953/// composition rule for Order2, OneSeed, and TwoSeed.
1954fn latent_kernel_normalized_log_parts<const K: usize>(
1955    base_log_sum: f64,
1956    normalized_parts: [Order2<K>; 4],
1957    part_count: usize,
1958) -> [Order2<K>; 4] {
1959    assert!(matches!(part_count, 1 | 2 | 4));
1960    let log_stack = latent_unary_derivatives_log(1.0);
1961    let compose_log = |moments: [f64; 16]| {
1962        gam_math::jet_partitions::compose_unary_four_slot_coefficients(moments, log_stack)
1963    };
1964    let moments_for = |a: usize, b: usize| {
1965        let base = &normalized_parts[0].0;
1966        let u = &normalized_parts[1].0;
1967        let v = &normalized_parts[2].0;
1968        let uv = &normalized_parts[3].0;
1969        [
1970            1.0,
1971            base.g[a],
1972            base.g[b],
1973            base.h[a][b],
1974            u.v,
1975            u.g[a],
1976            u.g[b],
1977            u.h[a][b],
1978            v.v,
1979            v.g[a],
1980            v.g[b],
1981            v.h[a][b],
1982            uv.v,
1983            uv.g[a],
1984            uv.g[b],
1985            uv.h[a][b],
1986        ]
1987    };
1988
1989    let mut out = [Order2::<K>::constant(0.0); 4];
1990    out[0].0.v = base_log_sum;
1991    if part_count >= 2 {
1992        out[1].0.v = normalized_parts[1].0.v;
1993    }
1994    if part_count == 4 {
1995        out[2].0.v = normalized_parts[2].0.v;
1996        let composed = compose_log(moments_for(0, 0));
1997        out[3].0.v = composed[0b1100];
1998    }
1999
2000    for a in 0..K {
2001        let composed = compose_log(moments_for(a, a));
2002        out[0].0.g[a] = composed[0b0001];
2003        if part_count >= 2 {
2004            out[1].0.g[a] = composed[0b0101];
2005        }
2006        if part_count == 4 {
2007            out[2].0.g[a] = composed[0b1001];
2008            out[3].0.g[a] = composed[0b1101];
2009        }
2010        for b in a..K {
2011            let composed = compose_log(moments_for(a, b));
2012            out[0].0.h[a][b] = composed[0b0011];
2013            if part_count >= 2 {
2014                out[1].0.h[a][b] = composed[0b0111];
2015            }
2016            if part_count == 4 {
2017                out[2].0.h[a][b] = composed[0b1011];
2018                out[3].0.h[a][b] = composed[0b1111];
2019            }
2020            for part in 0..part_count {
2021                out[part].0.h[b][a] = out[part].0.h[a][b];
2022            }
2023        }
2024    }
2025    out
2026}
2027
2028#[inline]
2029fn latent_kernel_direction_linear_combination<const K: usize>(
2030    primary_directions: &[LatentKernelPrimaryDirection; K],
2031    coefficients: &[f64; K],
2032) -> LatentKernelPrimaryDirection {
2033    let mut out = LatentKernelPrimaryDirection {
2034        dq: 0.0,
2035        dqd: 0.0,
2036        dmu: 0.0,
2037        dtau: 0.0,
2038    };
2039    for a in 0..K {
2040        out.dq += coefficients[a] * primary_directions[a].dq;
2041        out.dqd += coefficients[a] * primary_directions[a].dqd;
2042        out.dmu += coefficients[a] * primary_directions[a].dmu;
2043        out.dtau += coefficients[a] * primary_directions[a].dtau;
2044    }
2045    out
2046}
2047
2048/// Backend seam for the single latent-survival row expression.  Only the
2049/// analytic multivariate kernel primitive differs by requested channel; all
2050/// numerator/denominator/event algebra below is instantiated unchanged.
2051trait LatentPrimaryJetBackend<const K: usize> {
2052    type Jet: JetScalar<K>;
2053
2054    fn kernel_sum_log(
2055        &self,
2056        quadctx: &QuadratureContext,
2057        base_terms: &[LatentKernelPrimaryTerm],
2058        state: LatentKernelPrimaryState,
2059        primary_directions: &[LatentKernelPrimaryDirection; K],
2060        context: &str,
2061    ) -> Result<Self::Jet, LatentSurvivalError>;
2062}
2063
2064#[derive(Clone, Copy)]
2065struct LatentOrder2Backend;
2066
2067impl<const K: usize> LatentPrimaryJetBackend<K> for LatentOrder2Backend {
2068    type Jet = Order2<K>;
2069
2070    fn kernel_sum_log(
2071        &self,
2072        quadctx: &QuadratureContext,
2073        base_terms: &[LatentKernelPrimaryTerm],
2074        state: LatentKernelPrimaryState,
2075        primary_directions: &[LatentKernelPrimaryDirection; K],
2076        context: &str,
2077    ) -> Result<Self::Jet, LatentSurvivalError> {
2078        let suffixes: [&[LatentKernelPrimaryDirection]; 1] = [&[]];
2079        let parts = latent_kernel_sum_order2_parts(
2080            quadctx,
2081            base_terms,
2082            state,
2083            primary_directions,
2084            &suffixes,
2085            context,
2086        )?;
2087        Ok(parts[0])
2088    }
2089}
2090
2091#[derive(Clone, Copy)]
2092struct LatentOneSeedBackend<const K: usize> {
2093    direction: [f64; K],
2094}
2095
2096impl<const K: usize> LatentPrimaryJetBackend<K> for LatentOneSeedBackend<K> {
2097    type Jet = OneSeed<K>;
2098
2099    fn kernel_sum_log(
2100        &self,
2101        quadctx: &QuadratureContext,
2102        base_terms: &[LatentKernelPrimaryTerm],
2103        state: LatentKernelPrimaryState,
2104        primary_directions: &[LatentKernelPrimaryDirection; K],
2105        context: &str,
2106    ) -> Result<Self::Jet, LatentSurvivalError> {
2107        let seed = latent_kernel_direction_linear_combination(primary_directions, &self.direction);
2108        let seed_suffix = [seed];
2109        let suffixes: [&[LatentKernelPrimaryDirection]; 2] = [&[], &seed_suffix];
2110        let parts = latent_kernel_sum_order2_parts(
2111            quadctx,
2112            base_terms,
2113            state,
2114            primary_directions,
2115            &suffixes,
2116            context,
2117        )?;
2118        Ok(OneSeed {
2119            base: parts[0],
2120            eps: parts[1],
2121        })
2122    }
2123}
2124
2125#[derive(Clone, Copy)]
2126struct LatentTwoSeedBackend<const K: usize> {
2127    direction_u: [f64; K],
2128    direction_v: [f64; K],
2129}
2130
2131impl<const K: usize> LatentPrimaryJetBackend<K> for LatentTwoSeedBackend<K> {
2132    type Jet = TwoSeed<K>;
2133
2134    fn kernel_sum_log(
2135        &self,
2136        quadctx: &QuadratureContext,
2137        base_terms: &[LatentKernelPrimaryTerm],
2138        state: LatentKernelPrimaryState,
2139        primary_directions: &[LatentKernelPrimaryDirection; K],
2140        context: &str,
2141    ) -> Result<Self::Jet, LatentSurvivalError> {
2142        let seed_u =
2143            latent_kernel_direction_linear_combination(primary_directions, &self.direction_u);
2144        let seed_v =
2145            latent_kernel_direction_linear_combination(primary_directions, &self.direction_v);
2146        let suffix_u = [seed_u];
2147        let suffix_v = [seed_v];
2148        let suffix_uv = [seed_u, seed_v];
2149        let suffixes: [&[LatentKernelPrimaryDirection]; 4] =
2150            [&[], &suffix_u, &suffix_v, &suffix_uv];
2151        let parts = latent_kernel_sum_order2_parts(
2152            quadctx,
2153            base_terms,
2154            state,
2155            primary_directions,
2156            &suffixes,
2157            context,
2158        )?;
2159        Ok(TwoSeed {
2160            base: parts[0],
2161            eps: parts[1],
2162            del: parts[2],
2163            eps_del: parts[3],
2164        })
2165    }
2166}
2167
2168fn latent_survival_basis_direction(primary_idx: usize) -> LatentSurvivalPrimaryDirection {
2169    match primary_idx {
2170        LATENT_SURVIVAL_PRIMARY_Q_ENTRY => LatentSurvivalPrimaryDirection {
2171            dq_entry: 1.0,
2172            dq_exit: 0.0,
2173            dqdot_exit: 0.0,
2174            dq_right: 0.0,
2175            dmu: 0.0,
2176            dlog_sigma: 0.0,
2177        },
2178        LATENT_SURVIVAL_PRIMARY_Q_EXIT => LatentSurvivalPrimaryDirection {
2179            dq_entry: 0.0,
2180            dq_exit: 1.0,
2181            dqdot_exit: 0.0,
2182            dq_right: 0.0,
2183            dmu: 0.0,
2184            dlog_sigma: 0.0,
2185        },
2186        LATENT_SURVIVAL_PRIMARY_QDOT_EXIT => LatentSurvivalPrimaryDirection {
2187            dq_entry: 0.0,
2188            dq_exit: 0.0,
2189            dqdot_exit: 1.0,
2190            dq_right: 0.0,
2191            dmu: 0.0,
2192            dlog_sigma: 0.0,
2193        },
2194        LATENT_SURVIVAL_PRIMARY_Q_RIGHT => LatentSurvivalPrimaryDirection {
2195            dq_entry: 0.0,
2196            dq_exit: 0.0,
2197            dqdot_exit: 0.0,
2198            dq_right: 1.0,
2199            dmu: 0.0,
2200            dlog_sigma: 0.0,
2201        },
2202        LATENT_SURVIVAL_PRIMARY_MU => LatentSurvivalPrimaryDirection {
2203            dq_entry: 0.0,
2204            dq_exit: 0.0,
2205            dqdot_exit: 0.0,
2206            dq_right: 0.0,
2207            dmu: 1.0,
2208            dlog_sigma: 0.0,
2209        },
2210        LATENT_SURVIVAL_PRIMARY_LOG_SIGMA => LatentSurvivalPrimaryDirection {
2211            dq_entry: 0.0,
2212            dq_exit: 0.0,
2213            dqdot_exit: 0.0,
2214            dq_right: 0.0,
2215            dmu: 0.0,
2216            dlog_sigma: 1.0,
2217        },
2218        // SAFETY: latent survival has exactly `LATENT_SURVIVAL_PRIMARY_DIM`
2219        // (= 5) primary directions, indexed 0..=4 via the module-private
2220        // `LATENT_SURVIVAL_PRIMARY_*` constants. All five are matched
2221        // above, so this wildcard fires only on an out-of-range index,
2222        // which the internal iteration bounds (`0..LATENT_SURVIVAL_PRIMARY_DIM`)
2223        // make unreachable.
2224        // SAFETY: primary_idx is bounded by LATENT_SURVIVAL_PRIMARY_DIM at every internal call site.
2225        _ => std::panic::panic_any(format!(
2226            "latent survival primary index out of bounds: primary_idx={primary_idx}, primary_dim={LATENT_SURVIVAL_PRIMARY_DIM}"
2227        )),
2228    }
2229}
2230
2231fn latent_survival_map_entry_direction(
2232    direction: LatentSurvivalPrimaryDirection,
2233) -> LatentKernelPrimaryDirection {
2234    LatentKernelPrimaryDirection {
2235        dq: direction.dq_entry,
2236        dqd: 0.0,
2237        dmu: direction.dmu,
2238        dtau: direction.dlog_sigma,
2239    }
2240}
2241
2242fn latent_survival_map_exit_direction(
2243    direction: LatentSurvivalPrimaryDirection,
2244    event_type: LatentSurvivalEventType,
2245) -> LatentKernelPrimaryDirection {
2246    LatentKernelPrimaryDirection {
2247        dq: direction.dq_exit,
2248        dqd: if matches!(event_type, LatentSurvivalEventType::ExactEvent) {
2249            direction.dqdot_exit
2250        } else {
2251            0.0
2252        },
2253        dmu: direction.dmu,
2254        dtau: direction.dlog_sigma,
2255    }
2256}
2257
2258/// Direction map for the interval-censored LEFT boundary state (mass `M_L =
2259/// exp(q_exit)`). The left boundary tracks the same `q_exit` time functional as
2260/// right-censoring (no hazard-derivative channel), plus the shared `mu`/`sigma`.
2261fn latent_survival_map_left_direction(
2262    direction: LatentSurvivalPrimaryDirection,
2263) -> LatentKernelPrimaryDirection {
2264    LatentKernelPrimaryDirection {
2265        dq: direction.dq_exit,
2266        dqd: 0.0,
2267        dmu: direction.dmu,
2268        dtau: direction.dlog_sigma,
2269    }
2270}
2271
2272/// Direction map for the interval-censored RIGHT boundary state (mass `M_R =
2273/// exp(q_right)`). The right boundary tracks the dedicated `q_right` functional
2274/// (which shares the time-block coefficients with `q_exit` but is evaluated at
2275/// the interval upper bound `R`), plus the shared `mu`/`sigma`.
2276fn latent_survival_map_right_direction(
2277    direction: LatentSurvivalPrimaryDirection,
2278) -> LatentKernelPrimaryDirection {
2279    LatentKernelPrimaryDirection {
2280        dq: direction.dq_right,
2281        dqd: 0.0,
2282        dmu: direction.dmu,
2283        dtau: direction.dlog_sigma,
2284    }
2285}
2286
2287#[cfg(test)]
2288mod tests_multidir_row {
2289    use super::tests_multidir_kernel::latent_kernel_sum_log_jet;
2290    use super::*;
2291    use gam_math::jet_partitions::MultiDirJet as LatentMultiDirJet;
2292
2293    pub(super) fn latent_survival_row_primary_log_jet_multidir_reference(
2294        quadctx: &QuadratureContext,
2295        row: &LatentSurvivalRow,
2296        point: LatentSurvivalPrimaryPoint,
2297        directions: &[LatentSurvivalPrimaryDirection],
2298    ) -> Result<LatentMultiDirJet, String> {
2299        let LatentSurvivalPrimaryPoint {
2300            q_entry,
2301            q_exit,
2302            qdot_exit,
2303            mu,
2304            sigma,
2305            ..
2306        } = point;
2307        let log_sigma_factor = point.log_sigma_factor();
2308        let entry_state = LatentKernelPrimaryState {
2309            q: q_entry,
2310            qdot: 1.0,
2311            mu,
2312            sigma,
2313            log_sigma_factor,
2314        };
2315        let entry_directions = directions
2316            .iter()
2317            .copied()
2318            .map(latent_survival_map_entry_direction)
2319            .collect::<Vec<_>>();
2320
2321        let denominator = latent_kernel_sum_log_jet(
2322            quadctx,
2323            &[LatentKernelPrimaryTerm {
2324                coeff: 1.0,
2325                q_exp: 0,
2326                qdot_power: 0,
2327                tau_exp: 0,
2328                k: 0,
2329            }],
2330            entry_state,
2331            &entry_directions,
2332            "latent survival denominator",
2333        )?;
2334
2335        // The numerator for right-censoring / exact events is a single-state log-sum
2336        // kernel at the exit mass. Interval censoring is the difference of two
2337        // single-state kernels at DIFFERENT masses (L at `q_exit`, R at `q_right`),
2338        // so it is assembled by `latent_survival_interval_numerator_log_jet` below.
2339        let numerator = match row.event_type {
2340            LatentSurvivalEventType::RightCensored | LatentSurvivalEventType::ExactEvent => {
2341                let exit_state = LatentKernelPrimaryState {
2342                    q: q_exit,
2343                    qdot: qdot_exit,
2344                    mu,
2345                    sigma,
2346                    log_sigma_factor,
2347                };
2348                let exit_directions = directions
2349                    .iter()
2350                    .copied()
2351                    .map(|dir| latent_survival_map_exit_direction(dir, row.event_type))
2352                    .collect::<Vec<_>>();
2353                let numerator_terms = match row.event_type {
2354                    LatentSurvivalEventType::RightCensored => vec![LatentKernelPrimaryTerm {
2355                        coeff: 1.0,
2356                        q_exp: 0,
2357                        qdot_power: 0,
2358                        tau_exp: 0,
2359                        k: 0,
2360                    }],
2361                    LatentSurvivalEventType::ExactEvent => {
2362                        let mut terms = Vec::new();
2363                        if row.hazard_unloaded > 0.0 {
2364                            terms.push(LatentKernelPrimaryTerm {
2365                                coeff: row.hazard_unloaded,
2366                                q_exp: 0,
2367                                qdot_power: 0,
2368                                tau_exp: 0,
2369                                k: 0,
2370                            });
2371                        }
2372                        terms.push(LatentKernelPrimaryTerm {
2373                            coeff: 1.0,
2374                            q_exp: 1,
2375                            qdot_power: 1,
2376                            tau_exp: 0,
2377                            k: 1,
2378                        });
2379                        terms
2380                    }
2381                    LatentSurvivalEventType::IntervalCensored => {
2382                        // Interval-censored rows are routed to the dedicated two-state
2383                        // numerator branch (the outer match arm below), so this inner
2384                        // arm is not reached; a clean error rather than a panic guards
2385                        // against a future routing change.
2386                        return Err(
2387                            "interval-censored row reached the single-state numerator branch; \
2388                         it must take the dedicated two-state branch"
2389                                .to_string(),
2390                        );
2391                    }
2392                };
2393                latent_kernel_sum_log_jet(
2394                    quadctx,
2395                    &numerator_terms,
2396                    exit_state,
2397                    &exit_directions,
2398                    "latent survival numerator",
2399                )?
2400            }
2401            LatentSurvivalEventType::IntervalCensored => {
2402                latent_survival_interval_numerator_log_jet_multidir_reference(
2403                    quadctx, row, point, directions,
2404                )?
2405            }
2406        };
2407
2408        let mut total = numerator.add(&denominator.scale(-1.0));
2409        // For interval rows the unloaded exit mass is folded into the per-boundary
2410        // coefficients `exp(-mass_unloaded_{left,right})` inside the two-state
2411        // numerator, so only the (constant) unloaded-entry term remains here; for
2412        // right-censoring / exact events the exit/entry unloaded masses are an
2413        // additive constant on the log-likelihood.
2414        match row.event_type {
2415            LatentSurvivalEventType::IntervalCensored => {
2416                total.coeffs[0] += row.mass_unloaded_entry;
2417            }
2418            _ => {
2419                total.coeffs[0] += -row.mass_unloaded_exit + row.mass_unloaded_entry;
2420            }
2421        }
2422        Ok(total)
2423    }
2424
2425    /// Interval-censored numerator jet `log[ c_L·K_{0,M_L} − c_R·K_{0,M_R} ]` where
2426    /// `M_L = exp(q_exit)`, `M_R = exp(q_right)`, `c_L = exp(-mass_unloaded_left)`
2427    /// and `c_R = exp(-mass_unloaded_right)`.
2428    ///
2429    /// This is the dynamic-time analogue of the static
2430    /// [`LatentSurvivalRowJet::interval_censored`] kernel: the interval likelihood
2431    /// is the difference of two BOUNDARY survival masses, each a single-state
2432    /// order-0 kernel, but at two DISTINCT cumulative masses. Because the two
2433    /// boundaries respond to different time functionals (`q_exit` vs `q_right`) we
2434    /// cannot fold them into one `latent_kernel_sum_log_jet` state. Instead we:
2435    ///   1. build each boundary's `log K_{0,M}` jet at its own state, with its own
2436    ///      direction map (left → `dq_exit`, right → `dq_right`; both share
2437    ///      `mu`/`sigma`),
2438    ///   2. lift each to the LINEAR domain via `exp` (a unary composition whose five
2439    ///      derivatives at value `v` are all `exp(v)`), scaled by its coefficient
2440    ///      `c_L` (resp. `−c_R`),
2441    ///   3. add the two linear-domain jets, and
2442    ///   4. drop back to the log domain via the same `log` unary composition the
2443    ///      single-state path uses.
2444    /// Every multi-direction coefficient (value, score, neg-Hessian, 3rd, 4th)
2445    /// follows by the Faà-di-Bruno composition already implemented in
2446    /// `MultiDirJet::compose_unary`, so the derivative reductions are consistent
2447    /// with the exact-event/right-censored branches by construction.
2448    fn latent_survival_interval_numerator_log_jet_multidir_reference(
2449        quadctx: &QuadratureContext,
2450        row: &LatentSurvivalRow,
2451        point: LatentSurvivalPrimaryPoint,
2452        directions: &[LatentSurvivalPrimaryDirection],
2453    ) -> Result<LatentMultiDirJet, String> {
2454        let LatentSurvivalPrimaryPoint {
2455            q_exit,
2456            q_right,
2457            mu,
2458            sigma,
2459            ..
2460        } = point;
2461        let log_sigma_factor = point.log_sigma_factor();
2462        let single_k0 = [LatentKernelPrimaryTerm {
2463            coeff: 1.0,
2464            q_exp: 0,
2465            qdot_power: 0,
2466            tau_exp: 0,
2467            k: 0,
2468        }];
2469
2470        let left_state = LatentKernelPrimaryState {
2471            q: q_exit,
2472            qdot: 1.0,
2473            mu,
2474            sigma,
2475            log_sigma_factor,
2476        };
2477        let right_state = LatentKernelPrimaryState {
2478            q: q_right,
2479            qdot: 1.0,
2480            mu,
2481            sigma,
2482            log_sigma_factor,
2483        };
2484        let left_directions = directions
2485            .iter()
2486            .copied()
2487            .map(latent_survival_map_left_direction)
2488            .collect::<Vec<_>>();
2489        let right_directions = directions
2490            .iter()
2491            .copied()
2492            .map(latent_survival_map_right_direction)
2493            .collect::<Vec<_>>();
2494
2495        let log_left = latent_kernel_sum_log_jet(
2496            quadctx,
2497            &single_k0,
2498            left_state,
2499            &left_directions,
2500            "latent survival interval left boundary",
2501        )?;
2502        let log_right = latent_kernel_sum_log_jet(
2503            quadctx,
2504            &single_k0,
2505            right_state,
2506            &right_directions,
2507            "latent survival interval right boundary",
2508        )?;
2509
2510        // Lift each boundary's log-kernel jet to the linear domain and scale by the
2511        // unloaded-mass prefactor. exp''''(v) = exp(v) for all orders, so the unary
2512        // derivative tower is `[exp(v); exp(v); exp(v); exp(v); exp(v)]`.
2513        let c_left = (-row.mass_unloaded_left).exp();
2514        let c_right = (-row.mass_unloaded_right).exp();
2515        let exp_left_value = log_left.coeff(0).exp();
2516        let exp_right_value = log_right.coeff(0).exp();
2517        let linear_left = log_left.compose_unary([exp_left_value; 5]).scale(c_left);
2518        let linear_right = log_right.compose_unary([exp_right_value; 5]).scale(c_right);
2519
2520        let linear_numerator = linear_left.add(&linear_right.scale(-1.0));
2521        let base = linear_numerator.coeff(0);
2522        if !(base.is_finite() && base > 0.0) {
2523            return Err(LatentSurvivalError::NumericalFailure {
2524            reason: format!(
2525                "latent survival interval numerator must be a positive survival-mass difference, \
2526                 got c_L*K0(M_L) - c_R*K0(M_R) = {base}; require M_L < M_R (i.e. L < R)"
2527            ),
2528        }
2529        .into());
2530        }
2531        // Drop back to the log domain. `latent_unary_derivatives_log(base)` is the
2532        // unary derivative tower of `ln` at the positive base value, so the composed
2533        // value channel is `ln(base)` and the higher coefficients are the
2534        // log-of-a-difference score / curvature, consistent with the single-state
2535        // log-sum path (which composes `ln` at its normalised base of 1).
2536        Ok(linear_numerator.compose_unary(latent_unary_derivatives_log(base)))
2537    }
2538}
2539
2540/// The single latent-survival row program, instantiated at an order-two,
2541/// one-seed, or two-seed backend.  Event dispatch and the interval
2542/// log-difference are intentionally expressed once here; a backend changes only
2543/// the derivative layout used to lift each analytic kernel-sum primitive.
2544fn latent_survival_row_primary_jet<const K: usize, B: LatentPrimaryJetBackend<K>>(
2545    backend: &B,
2546    quadctx: &QuadratureContext,
2547    row: &LatentSurvivalRow,
2548    point: LatentSurvivalPrimaryPoint,
2549) -> Result<B::Jet, String> {
2550    let LatentSurvivalPrimaryPoint {
2551        q_entry,
2552        q_exit,
2553        qdot_exit,
2554        mu,
2555        sigma,
2556        ..
2557    } = point;
2558    let log_sigma_factor = point.log_sigma_factor();
2559    let entry_state = LatentKernelPrimaryState {
2560        q: q_entry,
2561        qdot: 1.0,
2562        mu,
2563        sigma,
2564        log_sigma_factor,
2565    };
2566    let entry_directions: [LatentKernelPrimaryDirection; K] = std::array::from_fn(|a| {
2567        latent_survival_map_entry_direction(latent_survival_basis_direction(a))
2568    });
2569    let denominator = backend
2570        .kernel_sum_log(
2571            quadctx,
2572            &[LatentKernelPrimaryTerm {
2573                coeff: 1.0,
2574                q_exp: 0,
2575                qdot_power: 0,
2576                tau_exp: 0,
2577                k: 0,
2578            }],
2579            entry_state,
2580            &entry_directions,
2581            "latent survival denominator",
2582        )
2583        .map_err(|error| error.to_string())?;
2584
2585    let numerator = match row.event_type {
2586        LatentSurvivalEventType::RightCensored => {
2587            let exit_state = LatentKernelPrimaryState {
2588                q: q_exit,
2589                qdot: qdot_exit,
2590                mu,
2591                sigma,
2592                log_sigma_factor,
2593            };
2594            let exit_directions: [LatentKernelPrimaryDirection; K] = std::array::from_fn(|a| {
2595                latent_survival_map_exit_direction(
2596                    latent_survival_basis_direction(a),
2597                    row.event_type,
2598                )
2599            });
2600            backend
2601                .kernel_sum_log(
2602                    quadctx,
2603                    &[LatentKernelPrimaryTerm {
2604                        coeff: 1.0,
2605                        q_exp: 0,
2606                        qdot_power: 0,
2607                        tau_exp: 0,
2608                        k: 0,
2609                    }],
2610                    exit_state,
2611                    &exit_directions,
2612                    "latent survival numerator",
2613                )
2614                .map_err(|error| error.to_string())?
2615        }
2616        LatentSurvivalEventType::ExactEvent => {
2617            let exit_state = LatentKernelPrimaryState {
2618                q: q_exit,
2619                qdot: qdot_exit,
2620                mu,
2621                sigma,
2622                log_sigma_factor,
2623            };
2624            let exit_directions: [LatentKernelPrimaryDirection; K] = std::array::from_fn(|a| {
2625                latent_survival_map_exit_direction(
2626                    latent_survival_basis_direction(a),
2627                    LatentSurvivalEventType::ExactEvent,
2628                )
2629            });
2630            // A zero unloaded-hazard term remains in this stack array but the
2631            // signed-log evaluator and derivative recurrences discard it.
2632            let numerator_terms = [
2633                LatentKernelPrimaryTerm {
2634                    coeff: row.hazard_unloaded,
2635                    q_exp: 0,
2636                    qdot_power: 0,
2637                    tau_exp: 0,
2638                    k: 0,
2639                },
2640                LatentKernelPrimaryTerm {
2641                    coeff: 1.0,
2642                    q_exp: 1,
2643                    qdot_power: 1,
2644                    tau_exp: 0,
2645                    k: 1,
2646                },
2647            ];
2648            backend
2649                .kernel_sum_log(
2650                    quadctx,
2651                    &numerator_terms,
2652                    exit_state,
2653                    &exit_directions,
2654                    "latent survival numerator",
2655                )
2656                .map_err(|error| error.to_string())?
2657        }
2658        LatentSurvivalEventType::IntervalCensored => {
2659            latent_survival_interval_numerator_jet(backend, quadctx, row, point)?
2660        }
2661    };
2662
2663    let unloaded_offset = match row.event_type {
2664        LatentSurvivalEventType::IntervalCensored => row.mass_unloaded_entry,
2665        _ => -row.mass_unloaded_exit + row.mass_unloaded_entry,
2666    };
2667    Ok(numerator
2668        .sub(&denominator)
2669        .add(&B::Jet::constant(unloaded_offset)))
2670}
2671
2672fn latent_survival_interval_numerator_jet<const K: usize, B: LatentPrimaryJetBackend<K>>(
2673    backend: &B,
2674    quadctx: &QuadratureContext,
2675    row: &LatentSurvivalRow,
2676    point: LatentSurvivalPrimaryPoint,
2677) -> Result<B::Jet, String> {
2678    let LatentSurvivalPrimaryPoint {
2679        q_exit,
2680        q_right,
2681        mu,
2682        sigma,
2683        ..
2684    } = point;
2685    let log_sigma_factor = point.log_sigma_factor();
2686    let single_k0 = [LatentKernelPrimaryTerm {
2687        coeff: 1.0,
2688        q_exp: 0,
2689        qdot_power: 0,
2690        tau_exp: 0,
2691        k: 0,
2692    }];
2693    let left_state = LatentKernelPrimaryState {
2694        q: q_exit,
2695        qdot: 1.0,
2696        mu,
2697        sigma,
2698        log_sigma_factor,
2699    };
2700    let right_state = LatentKernelPrimaryState {
2701        q: q_right,
2702        qdot: 1.0,
2703        mu,
2704        sigma,
2705        log_sigma_factor,
2706    };
2707    let left_directions: [LatentKernelPrimaryDirection; K] = std::array::from_fn(|a| {
2708        latent_survival_map_left_direction(latent_survival_basis_direction(a))
2709    });
2710    let right_directions: [LatentKernelPrimaryDirection; K] = std::array::from_fn(|a| {
2711        latent_survival_map_right_direction(latent_survival_basis_direction(a))
2712    });
2713    let log_left = backend
2714        .kernel_sum_log(
2715            quadctx,
2716            &single_k0,
2717            left_state,
2718            &left_directions,
2719            "latent survival interval left boundary",
2720        )
2721        .map_err(|error| error.to_string())?;
2722    let log_right = backend
2723        .kernel_sum_log(
2724            quadctx,
2725            &single_k0,
2726            right_state,
2727            &right_directions,
2728            "latent survival interval right boundary",
2729        )
2730        .map_err(|error| error.to_string())?;
2731
2732    let linear_left = log_left.exp().scale((-row.mass_unloaded_left).exp());
2733    let linear_right = log_right.exp().scale((-row.mass_unloaded_right).exp());
2734    let linear_numerator = linear_left.sub(&linear_right);
2735    let base = linear_numerator.value();
2736    if !(base.is_finite() && base > 0.0) {
2737        return Err(LatentSurvivalError::NumericalFailure {
2738            reason: format!(
2739                "latent survival interval numerator must be a positive survival-mass difference, \
2740                 got c_L*K0(M_L) - c_R*K0(M_R) = {base}; require M_L < M_R (i.e. L < R)"
2741            ),
2742        }
2743        .into());
2744    }
2745    Ok(linear_numerator.compose_unary(latent_unary_derivatives_log(base)))
2746}
2747
2748fn latent_survival_row_primary_gradient_hessian(
2749    quadctx: &QuadratureContext,
2750    row: &LatentSurvivalRow,
2751    point: LatentSurvivalPrimaryPoint,
2752    include_log_sigma: bool,
2753) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
2754    if include_log_sigma {
2755        let out = latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_DIM, _>(
2756            &LatentOrder2Backend,
2757            quadctx,
2758            row,
2759            point,
2760        )?;
2761        let out_gradient = out.g();
2762        let hessian = out.h();
2763        Ok((
2764            out.value(),
2765            Array1::from_shape_fn(LATENT_SURVIVAL_PRIMARY_DIM, |a| out_gradient[a]),
2766            Array2::from_shape_fn(
2767                (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
2768                |(a, b)| -hessian[a][b],
2769            ),
2770        ))
2771    } else {
2772        let out = latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_LOG_SIGMA, _>(
2773            &LatentOrder2Backend,
2774            quadctx,
2775            row,
2776            point,
2777        )?;
2778        let out_gradient = out.g();
2779        let out_hessian = out.h();
2780        Ok((
2781            out.value(),
2782            Array1::from_shape_fn(LATENT_SURVIVAL_PRIMARY_DIM, |a| {
2783                if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA {
2784                    out_gradient[a]
2785                } else {
2786                    0.0
2787                }
2788            }),
2789            Array2::from_shape_fn(
2790                (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
2791                |(a, b)| {
2792                    if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
2793                        && b < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
2794                    {
2795                        -out_hessian[a][b]
2796                    } else {
2797                        0.0
2798                    }
2799                },
2800            ),
2801        ))
2802    }
2803}
2804
2805fn latent_survival_row_primary_one_seed_fixed_sigma(
2806    quadctx: &QuadratureContext,
2807    row: &LatentSurvivalRow,
2808    point: LatentSurvivalPrimaryPoint,
2809    direction: &Array1<f64>,
2810) -> Result<OneSeed<LATENT_SURVIVAL_PRIMARY_LOG_SIGMA>, String> {
2811    let backend = LatentOneSeedBackend {
2812        direction: std::array::from_fn(|a| direction[a]),
2813    };
2814    latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_LOG_SIGMA, _>(
2815        &backend, quadctx, row, point,
2816    )
2817}
2818
2819fn latent_survival_row_primary_two_seed_fixed_sigma(
2820    quadctx: &QuadratureContext,
2821    row: &LatentSurvivalRow,
2822    point: LatentSurvivalPrimaryPoint,
2823    direction_u: &Array1<f64>,
2824    direction_v: &Array1<f64>,
2825) -> Result<TwoSeed<LATENT_SURVIVAL_PRIMARY_LOG_SIGMA>, String> {
2826    let backend = LatentTwoSeedBackend {
2827        direction_u: std::array::from_fn(|a| direction_u[a]),
2828        direction_v: std::array::from_fn(|a| direction_v[a]),
2829    };
2830    latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_LOG_SIGMA, _>(
2831        &backend, quadctx, row, point,
2832    )
2833}
2834
2835fn latent_survival_row_primary_third_contracted(
2836    quadctx: &QuadratureContext,
2837    row: &LatentSurvivalRow,
2838    point: LatentSurvivalPrimaryPoint,
2839    direction: &Array1<f64>,
2840    include_log_sigma: bool,
2841) -> Result<Array2<f64>, String> {
2842    if include_log_sigma {
2843        let backend = LatentOneSeedBackend {
2844            direction: std::array::from_fn(|a| direction[a]),
2845        };
2846        let out = latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_DIM, _>(
2847            &backend, quadctx, row, point,
2848        )?;
2849        let third = out.contracted_third();
2850        Ok(Array2::from_shape_fn(
2851            (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
2852            |(a, b)| -third[a][b],
2853        ))
2854    } else {
2855        let out = latent_survival_row_primary_one_seed_fixed_sigma(quadctx, row, point, direction)?;
2856        let third = out.contracted_third();
2857        Ok(Array2::from_shape_fn(
2858            (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
2859            |(a, b)| {
2860                if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA && b < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA {
2861                    -third[a][b]
2862                } else {
2863                    0.0
2864                }
2865            },
2866        ))
2867    }
2868}
2869
2870fn latent_survival_row_primary_fourth_contracted(
2871    quadctx: &QuadratureContext,
2872    row: &LatentSurvivalRow,
2873    point: LatentSurvivalPrimaryPoint,
2874    direction_u: &Array1<f64>,
2875    direction_v: &Array1<f64>,
2876    include_log_sigma: bool,
2877) -> Result<Array2<f64>, String> {
2878    if include_log_sigma {
2879        let backend = LatentTwoSeedBackend {
2880            direction_u: std::array::from_fn(|a| direction_u[a]),
2881            direction_v: std::array::from_fn(|a| direction_v[a]),
2882        };
2883        let out = latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_DIM, _>(
2884            &backend, quadctx, row, point,
2885        )?;
2886        let fourth = out.contracted_fourth();
2887        Ok(Array2::from_shape_fn(
2888            (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
2889            |(a, b)| -fourth[a][b],
2890        ))
2891    } else {
2892        let out = latent_survival_row_primary_two_seed_fixed_sigma(
2893            quadctx,
2894            row,
2895            point,
2896            direction_u,
2897            direction_v,
2898        )?;
2899        let fourth = out.contracted_fourth();
2900        Ok(Array2::from_shape_fn(
2901            (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
2902            |(a, b)| {
2903                if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA && b < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA {
2904                    -fourth[a][b]
2905                } else {
2906                    0.0
2907                }
2908            },
2909        ))
2910    }
2911}
2912
2913#[cfg(test)]
2914mod tests_multidir_channels {
2915    use super::tests_multidir_row::latent_survival_row_primary_log_jet_multidir_reference;
2916    use super::*;
2917
2918    pub(super) fn latent_survival_row_primary_gradient_hessian_multidir_reference(
2919        quadctx: &QuadratureContext,
2920        row: &LatentSurvivalRow,
2921        point: LatentSurvivalPrimaryPoint,
2922        include_log_sigma: bool,
2923    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
2924        let mut gradient = Array1::<f64>::zeros(LATENT_SURVIVAL_PRIMARY_DIM);
2925        let mut neg_hessian =
2926            Array2::<f64>::zeros((LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM));
2927        let active_primary = if include_log_sigma {
2928            LATENT_SURVIVAL_PRIMARY_DIM
2929        } else {
2930            LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
2931        };
2932        let log_lik =
2933            latent_survival_row_primary_log_jet_multidir_reference(quadctx, row, point, &[])?
2934                .coeff(0);
2935        for a in 0..active_primary {
2936            let dir_a = latent_survival_basis_direction(a);
2937            gradient[a] = latent_survival_row_primary_log_jet_multidir_reference(
2938                quadctx,
2939                row,
2940                point,
2941                &[dir_a],
2942            )?
2943            .coeff(1);
2944            for b in a..active_primary {
2945                let coeff = latent_survival_row_primary_log_jet_multidir_reference(
2946                    quadctx,
2947                    row,
2948                    point,
2949                    &[dir_a, latent_survival_basis_direction(b)],
2950                )?
2951                .coeff(3);
2952                neg_hessian[[a, b]] = -coeff;
2953                neg_hessian[[b, a]] = -coeff;
2954            }
2955        }
2956        Ok((log_lik, gradient, neg_hessian))
2957    }
2958
2959    pub(super) fn latent_survival_row_primary_third_contracted_multidir_reference(
2960        quadctx: &QuadratureContext,
2961        row: &LatentSurvivalRow,
2962        point: LatentSurvivalPrimaryPoint,
2963        direction: &Array1<f64>,
2964        include_log_sigma: bool,
2965    ) -> Result<Array2<f64>, String> {
2966        let active_primary = if include_log_sigma {
2967            LATENT_SURVIVAL_PRIMARY_DIM
2968        } else {
2969            LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
2970        };
2971        let dir = LatentSurvivalPrimaryDirection {
2972            dq_entry: direction[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
2973            dq_exit: direction[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
2974            dqdot_exit: direction[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
2975            dq_right: direction[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
2976            dmu: direction[LATENT_SURVIVAL_PRIMARY_MU],
2977            dlog_sigma: direction[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
2978        };
2979        let mut out =
2980            Array2::<f64>::zeros((LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM));
2981        for a in 0..active_primary {
2982            let dir_a = latent_survival_basis_direction(a);
2983            for b in a..active_primary {
2984                let coeff = latent_survival_row_primary_log_jet_multidir_reference(
2985                    quadctx,
2986                    row,
2987                    point,
2988                    &[dir_a, latent_survival_basis_direction(b), dir],
2989                )?
2990                .coeff(7);
2991                out[[a, b]] = -coeff;
2992                out[[b, a]] = -coeff;
2993            }
2994        }
2995        Ok(out)
2996    }
2997
2998    pub(super) fn latent_survival_row_primary_fourth_contracted_multidir_reference(
2999        quadctx: &QuadratureContext,
3000        row: &LatentSurvivalRow,
3001        point: LatentSurvivalPrimaryPoint,
3002        direction_u: &Array1<f64>,
3003        direction_v: &Array1<f64>,
3004        include_log_sigma: bool,
3005    ) -> Result<Array2<f64>, String> {
3006        let active_primary = if include_log_sigma {
3007            LATENT_SURVIVAL_PRIMARY_DIM
3008        } else {
3009            LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
3010        };
3011        let dir_u = LatentSurvivalPrimaryDirection {
3012            dq_entry: direction_u[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
3013            dq_exit: direction_u[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
3014            dqdot_exit: direction_u[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
3015            dq_right: direction_u[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
3016            dmu: direction_u[LATENT_SURVIVAL_PRIMARY_MU],
3017            dlog_sigma: direction_u[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
3018        };
3019        let dir_v = LatentSurvivalPrimaryDirection {
3020            dq_entry: direction_v[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
3021            dq_exit: direction_v[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
3022            dqdot_exit: direction_v[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
3023            dq_right: direction_v[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
3024            dmu: direction_v[LATENT_SURVIVAL_PRIMARY_MU],
3025            dlog_sigma: direction_v[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
3026        };
3027        let mut out =
3028            Array2::<f64>::zeros((LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM));
3029        for a in 0..active_primary {
3030            let dir_a = latent_survival_basis_direction(a);
3031            for b in a..active_primary {
3032                let coeff = latent_survival_row_primary_log_jet_multidir_reference(
3033                    quadctx,
3034                    row,
3035                    point,
3036                    &[dir_a, latent_survival_basis_direction(b), dir_u, dir_v],
3037                )?
3038                .coeff(15);
3039                out[[a, b]] = -coeff;
3040                out[[b, a]] = -coeff;
3041            }
3042        }
3043        Ok(out)
3044    }
3045}
3046
3047#[derive(Clone)]
3048struct LatentSurvivalJointSlices {
3049    time: std::ops::Range<usize>,
3050    mean: std::ops::Range<usize>,
3051    log_sigma: Option<std::ops::Range<usize>>,
3052    total: usize,
3053}
3054
3055#[derive(Clone)]
3056struct LatentSurvivalJointGradientAccum {
3057    ll: CompensatedRowSum,
3058    gradient: Array1<f64>,
3059}
3060
3061#[derive(Clone)]
3062struct LatentSurvivalJointDenseAccum {
3063    ll: CompensatedRowSum,
3064    gradient: Array1<f64>,
3065    hessian: Array2<f64>,
3066}
3067
3068#[derive(Clone)]
3069struct LatentSurvivalDenseHessianAccum {
3070    hessian: Array2<f64>,
3071}
3072
3073#[derive(Clone, Copy, Default)]
3074struct CompensatedRowSum {
3075    sum: f64,
3076    correction: f64,
3077}
3078
3079impl CompensatedRowSum {
3080    #[inline]
3081    fn add(&mut self, value: f64) {
3082        let next = self.sum + value;
3083        if self.sum.abs() >= value.abs() {
3084            self.correction += (self.sum - next) + value;
3085        } else {
3086            self.correction += (value - next) + self.sum;
3087        }
3088        self.sum = next;
3089    }
3090
3091    #[inline]
3092    fn value(self) -> f64 {
3093        self.sum + self.correction
3094    }
3095}
3096
3097/// Process latent-survival rows in fixed contiguous chunks, using one
3098/// accumulator per rayon task and reducing those accumulators in chunk-index
3099/// order so gradient/Hessian assembly stays deterministic across runs.
3100fn deterministic_latent_survival_row_reduction<Acc, Init, Process, Combine>(
3101    n_rows: usize,
3102    init: Init,
3103    process_row: Process,
3104    mut combine: Combine,
3105) -> Result<Acc, String>
3106where
3107    Acc: Send,
3108    Init: Fn() -> Acc + Sync,
3109    Process: Fn(usize, &mut Acc) -> Result<(), String> + Sync,
3110    Combine: FnMut(&mut Acc, Acc),
3111{
3112    use rayon::iter::{IntoParallelIterator, ParallelIterator};
3113
3114    const TARGET_CHUNK_COUNT: usize = 32;
3115    if n_rows == 0 {
3116        return Ok(init());
3117    }
3118    let chunk_size = n_rows.div_ceil(TARGET_CHUNK_COUNT).max(1);
3119    let n_chunks = n_rows.div_ceil(chunk_size);
3120    let chunk_accumulators: Vec<Acc> = (0..n_chunks)
3121        .into_par_iter()
3122        .map(|chunk_idx| -> Result<Acc, String> {
3123            let start = chunk_idx * chunk_size;
3124            let end = (start + chunk_size).min(n_rows);
3125            let mut acc = init();
3126            for row_idx in start..end {
3127                process_row(row_idx, &mut acc)?;
3128            }
3129            Ok(acc)
3130        })
3131        .collect::<Result<Vec<_>, String>>()?;
3132
3133    let mut total = init();
3134    for acc in chunk_accumulators {
3135        combine(&mut total, acc);
3136    }
3137    Ok(total)
3138}
3139
3140impl LatentSurvivalFamily {
3141    /// Assemble the per-row [`LatentSurvivalRow`] for `row_idx` from the family's
3142    /// unloaded-mass/hazard fields and the supplied per-row time quantiles.
3143    ///
3144    /// Shared by every per-row reduction (log-likelihood, gradient, Hessian,
3145    /// directional third derivatives): each previously inlined an identical
3146    /// `event_type` lookup followed by the same 12-argument
3147    /// `build_latent_survival_row` call. Behavior is unchanged.
3148    fn build_row_at(
3149        &self,
3150        row_idx: usize,
3151        q_entry: f64,
3152        q_exit: f64,
3153        qdot_exit: f64,
3154        q_right: f64,
3155    ) -> Result<LatentSurvivalRow, LatentSurvivalError> {
3156        let event_type = latent_survival_event_type_for(self.event_target[row_idx]);
3157        build_latent_survival_row(
3158            row_idx,
3159            self.hazard_loading,
3160            event_type,
3161            q_entry,
3162            q_exit,
3163            qdot_exit,
3164            q_right,
3165            self.unloaded_mass_entry[row_idx],
3166            self.unloaded_mass_exit[row_idx],
3167            self.unloaded_mass_right[row_idx],
3168            self.unloaded_hazard_exit[row_idx],
3169        )
3170    }
3171
3172    fn joint_slices(&self) -> LatentSurvivalJointSlices {
3173        let p_time = self.x_time_exit.ncols();
3174        let p_mean = self.x_mean.ncols();
3175        let time = 0..p_time;
3176        let mean = p_time..p_time + p_mean;
3177        let log_sigma = self
3178            .latent_sd_fixed
3179            .is_none()
3180            .then_some((p_time + p_mean)..(p_time + p_mean + 1));
3181        LatentSurvivalJointSlices {
3182            total: log_sigma
3183                .as_ref()
3184                .map_or(p_time + p_mean, |range| range.end),
3185            time,
3186            mean,
3187            log_sigma,
3188        }
3189    }
3190
3191    fn row_primary_direction_from_flat(
3192        &self,
3193        row: usize,
3194        slices: &LatentSurvivalJointSlices,
3195        d_beta_flat: &Array1<f64>,
3196    ) -> Array1<f64> {
3197        let mut out = Array1::<f64>::zeros(LATENT_SURVIVAL_PRIMARY_DIM);
3198        let d_time = d_beta_flat.slice(s![slices.time.clone()]);
3199        out[LATENT_SURVIVAL_PRIMARY_Q_ENTRY] = self.x_time_entry.row(row).dot(&d_time);
3200        out[LATENT_SURVIVAL_PRIMARY_Q_EXIT] = self.x_time_exit.row(row).dot(&d_time);
3201        out[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT] = self.x_time_derivative_exit.row(row).dot(&d_time);
3202        out[LATENT_SURVIVAL_PRIMARY_Q_RIGHT] = self.x_time_right.row(row).dot(&d_time);
3203        out[LATENT_SURVIVAL_PRIMARY_MU] = self
3204            .x_mean
3205            .dot_row_view(row, d_beta_flat.slice(s![slices.mean.clone()]));
3206        if let Some(range) = &slices.log_sigma {
3207            out[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA] = d_beta_flat[range.start];
3208        }
3209        out
3210    }
3211
3212    fn joint_block_ranges(&self) -> Vec<std::ops::Range<usize>> {
3213        let slices = self.joint_slices();
3214        let mut ranges = vec![slices.time.clone(), slices.mean.clone()];
3215        if let Some(log_sigma) = slices.log_sigma {
3216            ranges.push(log_sigma);
3217        }
3218        ranges
3219    }
3220
3221    fn add_pullback_primary_gradient(
3222        &self,
3223        target: &mut Array1<f64>,
3224        row: usize,
3225        slices: &LatentSurvivalJointSlices,
3226        primary_gradient: &Array1<f64>,
3227        weight: f64,
3228    ) -> Result<(), String> {
3229        for (primary_idx, time_vec) in [
3230            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
3231            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
3232            (
3233                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3234                self.x_time_derivative_exit.row(row),
3235            ),
3236            (LATENT_SURVIVAL_PRIMARY_Q_RIGHT, self.x_time_right.row(row)),
3237        ] {
3238            let scale = checked_weighted_row_value(
3239                weight,
3240                primary_gradient[primary_idx],
3241                row,
3242                "primary gradient",
3243            )?;
3244            if scale == 0.0 {
3245                continue;
3246            }
3247            for i in 0..time_vec.len() {
3248                let xi = time_vec[i];
3249                if xi != 0.0 {
3250                    target[slices.time.start + i] += scale * xi;
3251                }
3252            }
3253        }
3254
3255        let mean_scale = checked_weighted_row_value(
3256            weight,
3257            primary_gradient[LATENT_SURVIVAL_PRIMARY_MU],
3258            row,
3259            "mean gradient",
3260        )?;
3261        if mean_scale != 0.0 {
3262            self.x_mean
3263                .axpy_row_into(
3264                    row,
3265                    mean_scale,
3266                    &mut target.slice_mut(s![slices.mean.clone()]),
3267                )
3268                .map_err(|error| {
3269                    format!(
3270                        "latent survival mean gradient pullback dimension mismatch: row={row}, mean_slice={:?}, target_len={}, x_mean_cols={}, error={error}",
3271                        slices.mean,
3272                        target.len(),
3273                        self.x_mean.ncols()
3274                    )
3275                })?;
3276        }
3277
3278        if let Some(log_sigma) = &slices.log_sigma {
3279            target[log_sigma.start] += checked_weighted_row_value(
3280                weight,
3281                primary_gradient[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA],
3282                row,
3283                "log-sigma gradient",
3284            )?;
3285        }
3286        Ok(())
3287    }
3288
3289    fn add_pullback_primary_hessian(
3290        &self,
3291        target: &mut Array2<f64>,
3292        row: usize,
3293        slices: &LatentSurvivalJointSlices,
3294        primary_hessian: &Array2<f64>,
3295    ) -> Result<(), String> {
3296        let time_weights = [
3297            primary_hessian[[
3298                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3299                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3300            ]],
3301            primary_hessian[[
3302                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3303                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3304            ]],
3305            primary_hessian[[
3306                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3307                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3308            ]],
3309            primary_hessian[[
3310                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3311                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3312            ]],
3313        ];
3314        let time_cross_weights = [
3315            (
3316                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3317                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3318                &self.x_time_entry,
3319                &self.x_time_exit,
3320            ),
3321            (
3322                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3323                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3324                &self.x_time_entry,
3325                &self.x_time_derivative_exit,
3326            ),
3327            (
3328                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3329                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3330                &self.x_time_exit,
3331                &self.x_time_derivative_exit,
3332            ),
3333            (
3334                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3335                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3336                &self.x_time_entry,
3337                &self.x_time_right,
3338            ),
3339            (
3340                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3341                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3342                &self.x_time_exit,
3343                &self.x_time_right,
3344            ),
3345            (
3346                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3347                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3348                &self.x_time_derivative_exit,
3349                &self.x_time_right,
3350            ),
3351        ];
3352        {
3353            let time_target = &mut target.slice_mut(s![slices.time.clone(), slices.time.clone()]);
3354            dense_outer_accumulate(time_target, time_weights[0], self.x_time_entry.row(row));
3355            dense_outer_accumulate(time_target, time_weights[1], self.x_time_exit.row(row));
3356            dense_outer_accumulate(
3357                time_target,
3358                time_weights[2],
3359                self.x_time_derivative_exit.row(row),
3360            );
3361            dense_outer_accumulate(time_target, time_weights[3], self.x_time_right.row(row));
3362            for (a, b, lhs, rhs) in time_cross_weights {
3363                let weight = primary_hessian[[a, b]];
3364                if weight == 0.0 {
3365                    continue;
3366                }
3367                dense_symmetric_cross_accumulate(time_target, weight, lhs.row(row), rhs.row(row));
3368            }
3369        }
3370
3371        let mean_weight = primary_hessian[[LATENT_SURVIVAL_PRIMARY_MU, LATENT_SURVIVAL_PRIMARY_MU]];
3372        self.x_mean
3373            .syr_row_into_view(
3374                row,
3375                mean_weight,
3376                target.slice_mut(s![slices.mean.clone(), slices.mean.clone()]),
3377            )
3378            .map_err(|error| {
3379                format!(
3380                    "latent survival mean Hessian pullback dimension mismatch: row={row}, mean_slice={:?}, target_dim={:?}, x_mean_cols={}, error={error}",
3381                    slices.mean,
3382                    target.dim(),
3383                    self.x_mean.ncols()
3384                )
3385            })?;
3386
3387        let mean_row = self
3388            .x_mean
3389            .try_row_chunk(row..row + 1)
3390            .map_err(|error| {
3391                format!(
3392                    "latent survival mean pullback row chunk failed: row={row}, x_mean_rows={}, x_mean_cols={}, error={error}",
3393                    self.x_mean.nrows(),
3394                    self.x_mean.ncols()
3395                )
3396            })?;
3397        let mean_vec = mean_row.row(0);
3398        let time_mean_weights = [
3399            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
3400            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
3401            (
3402                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3403                self.x_time_derivative_exit.row(row),
3404            ),
3405            (LATENT_SURVIVAL_PRIMARY_Q_RIGHT, self.x_time_right.row(row)),
3406        ];
3407        for (primary_idx, time_vec) in time_mean_weights {
3408            let weight = primary_hessian[[primary_idx, LATENT_SURVIVAL_PRIMARY_MU]];
3409            if weight == 0.0 {
3410                continue;
3411            }
3412            for i in 0..time_vec.len() {
3413                let xi = time_vec[i];
3414                if xi == 0.0 {
3415                    continue;
3416                }
3417                for j in 0..mean_vec.len() {
3418                    let xj = mean_vec[j];
3419                    if xj == 0.0 {
3420                        continue;
3421                    }
3422                    target[[slices.time.start + i, slices.mean.start + j]] += weight * xi * xj;
3423                    target[[slices.mean.start + j, slices.time.start + i]] += weight * xj * xi;
3424                }
3425            }
3426        }
3427
3428        if let Some(log_sigma) = &slices.log_sigma {
3429            let sigma_idx = log_sigma.start;
3430            target[[sigma_idx, sigma_idx]] += primary_hessian[[
3431                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
3432                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
3433            ]];
3434
3435            for (primary_idx, time_vec) in [
3436                (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
3437                (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
3438                (
3439                    LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3440                    self.x_time_derivative_exit.row(row),
3441                ),
3442                (LATENT_SURVIVAL_PRIMARY_Q_RIGHT, self.x_time_right.row(row)),
3443            ] {
3444                let weight = primary_hessian[[primary_idx, LATENT_SURVIVAL_PRIMARY_LOG_SIGMA]];
3445                if weight == 0.0 {
3446                    continue;
3447                }
3448                for i in 0..time_vec.len() {
3449                    let xi = time_vec[i];
3450                    if xi == 0.0 {
3451                        continue;
3452                    }
3453                    target[[slices.time.start + i, sigma_idx]] += weight * xi;
3454                    target[[sigma_idx, slices.time.start + i]] += weight * xi;
3455                }
3456            }
3457
3458            let mean_sigma_weight = primary_hessian[[
3459                LATENT_SURVIVAL_PRIMARY_MU,
3460                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
3461            ]];
3462            if mean_sigma_weight != 0.0 {
3463                for j in 0..mean_vec.len() {
3464                    let xj = mean_vec[j];
3465                    if xj == 0.0 {
3466                        continue;
3467                    }
3468                    target[[slices.mean.start + j, sigma_idx]] += mean_sigma_weight * xj;
3469                    target[[sigma_idx, slices.mean.start + j]] += mean_sigma_weight * xj;
3470                }
3471            }
3472        }
3473        Ok(())
3474    }
3475
3476    fn evaluate_exact_newton_joint_gradient_dense(
3477        &self,
3478        block_states: &[ParameterBlockState],
3479    ) -> Result<(f64, Array1<f64>), String> {
3480        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
3481            .map_err(String::from)?;
3482        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
3483        let q_right = self.time_q_right(block_states)?;
3484        let sigma = self.latent_sd(block_states)?;
3485        let slices = self.joint_slices();
3486        let include_log_sigma = slices.log_sigma.is_some();
3487        let total = slices.total;
3488        let acc = deterministic_latent_survival_row_reduction(
3489            self.event_target.len(),
3490            || LatentSurvivalJointGradientAccum {
3491                ll: CompensatedRowSum::default(),
3492                gradient: Array1::<f64>::zeros(total),
3493            },
3494            |row_idx, acc| {
3495                let wi = weights.at(row_idx);
3496                if wi == 0.0 {
3497                    return Ok(());
3498                }
3499                let row = self.build_row_at(
3500                    row_idx,
3501                    q_entry[row_idx],
3502                    q_exit[row_idx],
3503                    qdot_exit[row_idx],
3504                    q_right[row_idx],
3505                )?;
3506                let point = LatentSurvivalPrimaryPoint {
3507                    q_entry: q_entry[row_idx],
3508                    q_exit: q_exit[row_idx],
3509                    qdot_exit: qdot_exit[row_idx],
3510                    q_right: q_right[row_idx],
3511                    mu: mu[row_idx],
3512                    sigma,
3513                };
3514                let (row_ll, primary_gradient, _) = latent_survival_row_primary_gradient_hessian(
3515                    &self.quadctx,
3516                    &row,
3517                    point,
3518                    include_log_sigma,
3519                )?;
3520                acc.ll.add(checked_weighted_row_value(
3521                    wi,
3522                    row_ll,
3523                    row_idx,
3524                    "log likelihood",
3525                )?);
3526                self.add_pullback_primary_gradient(
3527                    &mut acc.gradient,
3528                    row_idx,
3529                    &slices,
3530                    &primary_gradient,
3531                    wi,
3532                )?;
3533                Ok(())
3534            },
3535            |total_acc, chunk_acc| {
3536                total_acc.ll.add(chunk_acc.ll.value());
3537                total_acc.gradient += &chunk_acc.gradient;
3538            },
3539        )?;
3540        let ll = require_finite_likelihood_scalar(acc.ll.value(), "log likelihood")?;
3541        require_finite_likelihood_vector(&acc.gradient, "gradient")?;
3542        Ok((ll, acc.gradient))
3543    }
3544
3545    /// Per-row residuals of the unpenalized NLL with respect to the three
3546    /// additive baseline time-block offsets `(entry, exit, derivative)`.
3547    ///
3548    /// The baseline configuration θ enters the latent-survival working model
3549    /// only through the additive offsets on the three time channels
3550    ///   q_entry = x_time_entry·β_time + o_E(θ),
3551    ///   q_exit  = x_time_exit·β_time  + o_X(θ),
3552    ///   q̇_exit = x_time_deriv·β_time + o_D(θ),
3553    /// exactly the offset channel the transformation path carries through
3554    /// [`WorkingModelSurvival::offset_channel_residuals`]. Because
3555    /// `∂q_ch/∂o_ch = 1`, the residual `∂NLL/∂o_ch_i` equals
3556    /// `−∂(log-likelihood)/∂q_ch_i`, and the per-row primary log-likelihood
3557    /// gradient over `(q_entry, q_exit, q̇_exit)` is precisely the
3558    /// `Q_ENTRY`/`Q_EXIT`/`QDOT_EXIT` components returned by
3559    /// [`latent_survival_row_primary_gradient_hessian`]. Sampleweight-scaled to
3560    /// match the [`OffsetChannelResiduals`] contract consumed by
3561    /// `baseline_chain_rule_gradient`.
3562    ///
3563    /// At the converged (constrained) β̂ the envelope theorem makes this the
3564    /// exact θ-gradient of the profile penalized NLL `0.5·deviance + 0.5·βᵀSβ`.
3565    /// The interval upper-bound `q_right = x_time_right·β_time + o_R(θ)` channel
3566    /// DOES carry its own baseline-θ offset `o_R(θ)` (the time basis evaluated at
3567    /// the bracket upper bound `R`), distinct from the exit offset at `L`, so its
3568    /// residual `−∂(log-likelihood)/∂q_right` is returned in the dedicated
3569    /// [`OffsetChannelResiduals::right`] channel; it is exactly 0 on every
3570    /// non-interval row (the `Q_RIGHT` primary channel is inert there) and the
3571    /// baseline-θ chain rule contracts it against the `age_right`-evaluated
3572    /// η-partial.
3573    pub fn offset_channel_residuals(
3574        &self,
3575        block_states: &[ParameterBlockState],
3576    ) -> Result<crate::survival::OffsetChannelResiduals, LatentSurvivalError> {
3577        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")?;
3578        let n = self.event_target.len();
3579        // `split_time_eta` validates the complete block slate before indexing
3580        // it and returns `LatentSurvivalError::BlockMismatch` when fitted state
3581        // is missing. There is deliberately no zero-residual fallback: zeros
3582        // would manufacture a stationary outer baseline gradient.
3583        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
3584        let q_right = self.time_q_right(block_states)?;
3585        let sigma = self.latent_sd(block_states)?;
3586        let include_log_sigma = self.joint_slices().log_sigma.is_some();
3587        let mut entry = Array1::<f64>::zeros(n);
3588        let mut exit = Array1::<f64>::zeros(n);
3589        let mut derivative = Array1::<f64>::zeros(n);
3590        let mut right = Array1::<f64>::zeros(n);
3591        for row_idx in 0..n {
3592            let wi = weights.at(row_idx);
3593            if wi == 0.0 {
3594                continue;
3595            }
3596            let row = self.build_row_at(
3597                row_idx,
3598                q_entry[row_idx],
3599                q_exit[row_idx],
3600                qdot_exit[row_idx],
3601                q_right[row_idx],
3602            )?;
3603            let point = LatentSurvivalPrimaryPoint {
3604                q_entry: q_entry[row_idx],
3605                q_exit: q_exit[row_idx],
3606                qdot_exit: qdot_exit[row_idx],
3607                q_right: q_right[row_idx],
3608                mu: mu[row_idx],
3609                sigma,
3610            };
3611            let (_, primary_gradient, _) = latent_survival_row_primary_gradient_hessian(
3612                &self.quadctx,
3613                &row,
3614                point,
3615                include_log_sigma,
3616            )
3617            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
3618            // ∂NLL/∂o_ch = −w · ∂(log-likelihood)/∂q_ch.
3619            entry[row_idx] = -checked_weighted_row_value(
3620                wi,
3621                primary_gradient[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
3622                row_idx,
3623                "entry-offset score",
3624            )
3625            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
3626            exit[row_idx] = -checked_weighted_row_value(
3627                wi,
3628                primary_gradient[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
3629                row_idx,
3630                "exit-offset score",
3631            )
3632            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
3633            derivative[row_idx] = -checked_weighted_row_value(
3634                wi,
3635                primary_gradient[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
3636                row_idx,
3637                "derivative-offset score",
3638            )
3639            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
3640            // Interval upper-bound (`R`) channel. `q_right` shares the time-block
3641            // coefficients but carries its OWN baseline-θ η-offset evaluated at
3642            // `R` (`o_R(θ)`), so the profile-NLL θ-gradient must include it.
3643            // `∂(log-likelihood)/∂q_right` is exactly 0 for non-interval rows
3644            // (the `Q_RIGHT` channel is inert there), so this is 0 except on
3645            // interval-censored rows.
3646            right[row_idx] = -checked_weighted_row_value(
3647                wi,
3648                primary_gradient[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
3649                row_idx,
3650                "right-offset score",
3651            )
3652            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
3653        }
3654        Ok(crate::survival::OffsetChannelResiduals {
3655            exit,
3656            entry,
3657            derivative,
3658            right,
3659        })
3660    }
3661
3662    /// Block-diagonal-only pullback: writes only time-time, mean-mean, and
3663    /// log_sigma-log_sigma rowwise contributions into per-block targets.
3664    /// Used by `evaluate()` to populate per-block working sets without ever
3665    /// materializing the cross blocks the inner solver does not consume.
3666    fn add_pullback_primary_block_diagonals(
3667        &self,
3668        row: usize,
3669        primary_hessian: &Array2<f64>,
3670        time_target: &mut Array2<f64>,
3671        mean_target: &mut Array2<f64>,
3672        log_sigma_target: Option<&mut Array2<f64>>,
3673    ) -> Result<(), String> {
3674        let h = primary_hessian;
3675        // Time block: 4 squared rows (entry/exit/qdot/right) + 6 symmetric
3676        // crosses. The interval right-boundary functional `q_right` shares the
3677        // time-block coefficients, so it accumulates into the same time target.
3678        dense_outer_accumulate(
3679            time_target,
3680            h[[
3681                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3682                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3683            ]],
3684            self.x_time_entry.row(row),
3685        );
3686        dense_outer_accumulate(
3687            time_target,
3688            h[[
3689                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3690                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3691            ]],
3692            self.x_time_exit.row(row),
3693        );
3694        dense_outer_accumulate(
3695            time_target,
3696            h[[
3697                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3698                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3699            ]],
3700            self.x_time_derivative_exit.row(row),
3701        );
3702        dense_outer_accumulate(
3703            time_target,
3704            h[[
3705                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3706                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3707            ]],
3708            self.x_time_right.row(row),
3709        );
3710        for (a, b, lhs, rhs) in [
3711            (
3712                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3713                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3714                &self.x_time_entry,
3715                &self.x_time_exit,
3716            ),
3717            (
3718                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3719                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3720                &self.x_time_entry,
3721                &self.x_time_derivative_exit,
3722            ),
3723            (
3724                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3725                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3726                &self.x_time_exit,
3727                &self.x_time_derivative_exit,
3728            ),
3729            (
3730                LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
3731                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3732                &self.x_time_entry,
3733                &self.x_time_right,
3734            ),
3735            (
3736                LATENT_SURVIVAL_PRIMARY_Q_EXIT,
3737                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3738                &self.x_time_exit,
3739                &self.x_time_right,
3740            ),
3741            (
3742                LATENT_SURVIVAL_PRIMARY_QDOT_EXIT,
3743                LATENT_SURVIVAL_PRIMARY_Q_RIGHT,
3744                &self.x_time_derivative_exit,
3745                &self.x_time_right,
3746            ),
3747        ] {
3748            let weight = h[[a, b]];
3749            if weight == 0.0 {
3750                continue;
3751            }
3752            dense_symmetric_cross_accumulate(time_target, weight, lhs.row(row), rhs.row(row));
3753        }
3754        // Mean block.
3755        let mean_weight = h[[LATENT_SURVIVAL_PRIMARY_MU, LATENT_SURVIVAL_PRIMARY_MU]];
3756        self.x_mean
3757            .syr_row_into_view(row, mean_weight, mean_target.view_mut())
3758            .map_err(|error| {
3759                format!(
3760                    "latent survival mean block-diagonal pullback dimension mismatch: row={row}, mean_target_dim={:?}, x_mean_cols={}, error={error}",
3761                    mean_target.dim(),
3762                    self.x_mean.ncols()
3763                )
3764            })?;
3765        // Log-σ block (scalar).
3766        if let Some(target) = log_sigma_target {
3767            target[[0, 0]] += h[[
3768                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
3769                LATENT_SURVIVAL_PRIMARY_LOG_SIGMA,
3770            ]];
3771        }
3772        Ok(())
3773    }
3774
3775    /// Block-diagonal evaluator used by `evaluate()`. Returns the per-row
3776    /// log-likelihood, the joint gradient (sliced into block gradients by
3777    /// the caller), and the three per-block diagonal Hessians without ever
3778    /// materializing the full joint matrix.
3779    fn evaluate_exact_newton_block_diagonals(
3780        &self,
3781        block_states: &[ParameterBlockState],
3782    ) -> Result<
3783        (
3784            f64,
3785            Array1<f64>,
3786            Array2<f64>,
3787            Array2<f64>,
3788            Option<Array2<f64>>,
3789        ),
3790        String,
3791    > {
3792        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
3793            .map_err(String::from)?;
3794        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
3795        let q_right = self.time_q_right(block_states)?;
3796        let sigma = self.latent_sd(block_states)?;
3797        let slices = self.joint_slices();
3798        let include_log_sigma = slices.log_sigma.is_some();
3799        let mut ll = CompensatedRowSum::default();
3800        let mut gradient = Array1::<f64>::zeros(slices.total);
3801        let p_time = slices.time.len();
3802        let p_mean = slices.mean.len();
3803        let mut hess_time = Array2::<f64>::zeros((p_time, p_time));
3804        let mut hess_mean = Array2::<f64>::zeros((p_mean, p_mean));
3805        let mut hess_log_sigma = if include_log_sigma {
3806            Some(Array2::<f64>::zeros((1, 1)))
3807        } else {
3808            None
3809        };
3810        for row_idx in 0..self.event_target.len() {
3811            let wi = weights.at(row_idx);
3812            if wi == 0.0 {
3813                continue;
3814            }
3815            let row = self.build_row_at(
3816                row_idx,
3817                q_entry[row_idx],
3818                q_exit[row_idx],
3819                qdot_exit[row_idx],
3820                q_right[row_idx],
3821            )?;
3822            let (row_ll, primary_gradient, primary_hessian) =
3823                latent_survival_row_primary_gradient_hessian(
3824                    &self.quadctx,
3825                    &row,
3826                    LatentSurvivalPrimaryPoint {
3827                        q_entry: q_entry[row_idx],
3828                        q_exit: q_exit[row_idx],
3829                        qdot_exit: qdot_exit[row_idx],
3830                        q_right: q_right[row_idx],
3831                        mu: mu[row_idx],
3832                        sigma,
3833                    },
3834                    include_log_sigma,
3835                )?;
3836            ll.add(checked_weighted_row_value(
3837                wi,
3838                row_ll,
3839                row_idx,
3840                "log likelihood",
3841            )?);
3842            self.add_pullback_primary_gradient(
3843                &mut gradient,
3844                row_idx,
3845                &slices,
3846                &primary_gradient,
3847                wi,
3848            )?;
3849            let weighted_primary_hessian =
3850                checked_weighted_row_matrix(wi, &primary_hessian, row_idx, "primary Hessian")?;
3851            self.add_pullback_primary_block_diagonals(
3852                row_idx,
3853                &weighted_primary_hessian,
3854                &mut hess_time,
3855                &mut hess_mean,
3856                hess_log_sigma.as_mut(),
3857            )?;
3858        }
3859        let ll = require_finite_likelihood_scalar(ll.value(), "log likelihood")?;
3860        require_finite_likelihood_vector(&gradient, "gradient")?;
3861        require_finite_likelihood_matrix(&hess_time, "time Hessian")?;
3862        require_finite_likelihood_matrix(&hess_mean, "mean Hessian")?;
3863        if let Some(hessian) = hess_log_sigma.as_ref() {
3864            require_finite_likelihood_matrix(hessian, "log-sigma Hessian")?;
3865        }
3866        Ok((ll, gradient, hess_time, hess_mean, hess_log_sigma))
3867    }
3868
3869    fn evaluate_exact_newton_joint_dense(
3870        &self,
3871        block_states: &[ParameterBlockState],
3872    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
3873        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
3874            .map_err(String::from)?;
3875        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
3876        let q_right = self.time_q_right(block_states)?;
3877        let sigma = self.latent_sd(block_states)?;
3878        let slices = self.joint_slices();
3879        let include_log_sigma = slices.log_sigma.is_some();
3880        let total = slices.total;
3881        let acc = deterministic_latent_survival_row_reduction(
3882            self.event_target.len(),
3883            || LatentSurvivalJointDenseAccum {
3884                ll: CompensatedRowSum::default(),
3885                gradient: Array1::<f64>::zeros(total),
3886                hessian: Array2::<f64>::zeros((total, total)),
3887            },
3888            |row_idx, acc| {
3889                let wi = weights.at(row_idx);
3890                if wi == 0.0 {
3891                    return Ok(());
3892                }
3893                let row = self.build_row_at(
3894                    row_idx,
3895                    q_entry[row_idx],
3896                    q_exit[row_idx],
3897                    qdot_exit[row_idx],
3898                    q_right[row_idx],
3899                )?;
3900                let (row_ll, primary_gradient, primary_hessian) =
3901                    latent_survival_row_primary_gradient_hessian(
3902                        &self.quadctx,
3903                        &row,
3904                        LatentSurvivalPrimaryPoint {
3905                            q_entry: q_entry[row_idx],
3906                            q_exit: q_exit[row_idx],
3907                            qdot_exit: qdot_exit[row_idx],
3908                            q_right: q_right[row_idx],
3909                            mu: mu[row_idx],
3910                            sigma,
3911                        },
3912                        include_log_sigma,
3913                    )?;
3914                acc.ll.add(checked_weighted_row_value(
3915                    wi,
3916                    row_ll,
3917                    row_idx,
3918                    "log likelihood",
3919                )?);
3920                self.add_pullback_primary_gradient(
3921                    &mut acc.gradient,
3922                    row_idx,
3923                    &slices,
3924                    &primary_gradient,
3925                    wi,
3926                )?;
3927                let weighted_primary_hessian =
3928                    checked_weighted_row_matrix(wi, &primary_hessian, row_idx, "primary Hessian")?;
3929                self.add_pullback_primary_hessian(
3930                    &mut acc.hessian,
3931                    row_idx,
3932                    &slices,
3933                    &weighted_primary_hessian,
3934                )?;
3935                Ok(())
3936            },
3937            |total_acc, chunk_acc| {
3938                total_acc.ll.add(chunk_acc.ll.value());
3939                total_acc.gradient += &chunk_acc.gradient;
3940                total_acc.hessian += &chunk_acc.hessian;
3941            },
3942        )?;
3943        let ll = require_finite_likelihood_scalar(acc.ll.value(), "log likelihood")?;
3944        require_finite_likelihood_vector(&acc.gradient, "gradient")?;
3945        require_finite_likelihood_matrix(&acc.hessian, "Hessian")?;
3946        Ok((ll, acc.gradient, acc.hessian))
3947    }
3948
3949    fn exact_newton_joint_hessian_directional_derivative_dense(
3950        &self,
3951        block_states: &[ParameterBlockState],
3952        d_beta_flat: &Array1<f64>,
3953    ) -> Result<Array2<f64>, String> {
3954        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
3955            .map_err(String::from)?;
3956        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
3957        let q_right = self.time_q_right(block_states)?;
3958        let sigma = self.latent_sd(block_states)?;
3959        let slices = self.joint_slices();
3960        if d_beta_flat.len() != slices.total {
3961            return Err(format!(
3962                "latent survival joint dH direction length mismatch: got {}, expected {}",
3963                d_beta_flat.len(),
3964                slices.total
3965            ));
3966        }
3967        let include_log_sigma = slices.log_sigma.is_some();
3968        let total = slices.total;
3969        let acc = deterministic_latent_survival_row_reduction(
3970            self.event_target.len(),
3971            || LatentSurvivalDenseHessianAccum {
3972                hessian: Array2::<f64>::zeros((total, total)),
3973            },
3974            |row_idx, acc| {
3975                let wi = weights.at(row_idx);
3976                if wi == 0.0 {
3977                    return Ok(());
3978                }
3979                let row = self.build_row_at(
3980                    row_idx,
3981                    q_entry[row_idx],
3982                    q_exit[row_idx],
3983                    qdot_exit[row_idx],
3984                    q_right[row_idx],
3985                )?;
3986                let direction = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_flat);
3987                let third = latent_survival_row_primary_third_contracted(
3988                    &self.quadctx,
3989                    &row,
3990                    LatentSurvivalPrimaryPoint {
3991                        q_entry: q_entry[row_idx],
3992                        q_exit: q_exit[row_idx],
3993                        qdot_exit: qdot_exit[row_idx],
3994                        q_right: q_right[row_idx],
3995                        mu: mu[row_idx],
3996                        sigma,
3997                    },
3998                    &direction,
3999                    include_log_sigma,
4000                )?;
4001                let weighted_third =
4002                    checked_weighted_row_matrix(wi, &third, row_idx, "contracted third")?;
4003                self.add_pullback_primary_hessian(
4004                    &mut acc.hessian,
4005                    row_idx,
4006                    &slices,
4007                    &weighted_third,
4008                )?;
4009                Ok(())
4010            },
4011            |total_acc, chunk_acc| {
4012                total_acc.hessian += &chunk_acc.hessian;
4013            },
4014        )?;
4015        require_finite_likelihood_matrix(&acc.hessian, "directional Hessian derivative")?;
4016        Ok(acc.hessian)
4017    }
4018
4019    fn exact_newton_joint_hessian_second_directional_derivative_dense(
4020        &self,
4021        block_states: &[ParameterBlockState],
4022        d_beta_u_flat: &Array1<f64>,
4023        d_beta_v_flat: &Array1<f64>,
4024    ) -> Result<Array2<f64>, String> {
4025        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
4026            .map_err(String::from)?;
4027        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
4028        let q_right = self.time_q_right(block_states)?;
4029        let sigma = self.latent_sd(block_states)?;
4030        let slices = self.joint_slices();
4031        if d_beta_u_flat.len() != slices.total || d_beta_v_flat.len() != slices.total {
4032            return Err(format!(
4033                "latent survival joint d2H direction length mismatch: got {} and {}, expected {}",
4034                d_beta_u_flat.len(),
4035                d_beta_v_flat.len(),
4036                slices.total
4037            ));
4038        }
4039        let include_log_sigma = slices.log_sigma.is_some();
4040        let total = slices.total;
4041        let acc = deterministic_latent_survival_row_reduction(
4042            self.event_target.len(),
4043            || LatentSurvivalDenseHessianAccum {
4044                hessian: Array2::<f64>::zeros((total, total)),
4045            },
4046            |row_idx, acc| {
4047                let wi = weights.at(row_idx);
4048                if wi == 0.0 {
4049                    return Ok(());
4050                }
4051                let row = self.build_row_at(
4052                    row_idx,
4053                    q_entry[row_idx],
4054                    q_exit[row_idx],
4055                    qdot_exit[row_idx],
4056                    q_right[row_idx],
4057                )?;
4058                let direction_u =
4059                    self.row_primary_direction_from_flat(row_idx, &slices, d_beta_u_flat);
4060                let direction_v =
4061                    self.row_primary_direction_from_flat(row_idx, &slices, d_beta_v_flat);
4062                let fourth = latent_survival_row_primary_fourth_contracted(
4063                    &self.quadctx,
4064                    &row,
4065                    LatentSurvivalPrimaryPoint {
4066                        q_entry: q_entry[row_idx],
4067                        q_exit: q_exit[row_idx],
4068                        qdot_exit: qdot_exit[row_idx],
4069                        q_right: q_right[row_idx],
4070                        mu: mu[row_idx],
4071                        sigma,
4072                    },
4073                    &direction_u,
4074                    &direction_v,
4075                    include_log_sigma,
4076                )?;
4077                let weighted_fourth =
4078                    checked_weighted_row_matrix(wi, &fourth, row_idx, "contracted fourth")?;
4079                self.add_pullback_primary_hessian(
4080                    &mut acc.hessian,
4081                    row_idx,
4082                    &slices,
4083                    &weighted_fourth,
4084                )?;
4085                Ok(())
4086            },
4087            |total_acc, chunk_acc| {
4088                total_acc.hessian += &chunk_acc.hessian;
4089            },
4090        )?;
4091        require_finite_likelihood_matrix(&acc.hessian, "second directional Hessian derivative")?;
4092        Ok(acc.hessian)
4093    }
4094}
4095
4096fn log_kernel_ratio(
4097    bundle: &crate::survival::lognormal_kernel::LogLognormalKernelBundle,
4098    num: usize,
4099    den: usize,
4100) -> f64 {
4101    let delta = bundle.get(num) - bundle.get(den);
4102    if delta.is_finite() {
4103        delta.exp()
4104    } else if delta > 0.0 {
4105        f64::INFINITY
4106    } else {
4107        0.0
4108    }
4109}
4110
4111fn logk_q_derivatives(
4112    quadctx: &QuadratureContext,
4113    k: usize,
4114    mass: f64,
4115    mu: f64,
4116    sigma: f64,
4117) -> Result<(f64, f64, IntegratedExpectationMode), LatentSurvivalError> {
4118    if mass <= 0.0 {
4119        return Ok((0.0, 0.0, IntegratedExpectationMode::ExactClosedForm));
4120    }
4121    let bundle = log_kernel_bundle(quadctx, mass, mu, sigma, k + 2).map_err(|e| {
4122        LatentSurvivalError::NumericalFailure {
4123            reason: format!("latent survival kernel evaluation failed: {e}"),
4124        }
4125    })?;
4126    let r1 = log_kernel_ratio(&bundle, k + 1, k);
4127    let r2 = log_kernel_ratio(&bundle, k + 2, k);
4128    let d1 = -mass * r1;
4129    let d2 = d1 + mass * mass * (r2 - r1 * r1);
4130    Ok((d1, d2, bundle.mode))
4131}
4132
4133fn latent_survival_time_jet(
4134    quadctx: &QuadratureContext,
4135    row: &LatentSurvivalRow,
4136    qdot_exit: f64,
4137    mu: f64,
4138    sigma: f64,
4139) -> Result<LatentSurvivalTimeJet, LatentSurvivalError> {
4140    let (entry_d1, entry_d2, _) = logk_q_derivatives(quadctx, 0, row.mass_entry, mu, sigma)?;
4141    match row.event_type {
4142        LatentSurvivalEventType::RightCensored => {
4143            let (exit_d1, exit_d2, _) = logk_q_derivatives(quadctx, 0, row.mass_exit, mu, sigma)?;
4144            Ok(LatentSurvivalTimeJet {
4145                grad_entry: -entry_d1,
4146                grad_exit: exit_d1,
4147                neg_hess_entry: entry_d2,
4148                neg_hess_exit: -exit_d2,
4149            })
4150        }
4151        LatentSurvivalEventType::ExactEvent => {
4152            if !(qdot_exit.is_finite() && qdot_exit > 0.0) {
4153                return Err(LatentSurvivalError::NumericalFailure {
4154                    reason: format!(
4155                        "latent survival requires positive finite baseline hazard derivative, got {qdot_exit}"
4156                    ),
4157                });
4158            }
4159            if row.hazard_unloaded > 0.0 {
4160                let bundle =
4161                    log_kernel_bundle(quadctx, row.mass_exit, mu, sigma, 3).map_err(|e| {
4162                        LatentSurvivalError::NumericalFailure {
4163                            reason: format!("latent survival kernel evaluation failed: {e}"),
4164                        }
4165                    })?;
4166                let (unloaded_d1, unloaded_d2, _) =
4167                    logk_q_derivatives(quadctx, 0, row.mass_exit, mu, sigma)?;
4168                let (loaded_log_d1, loaded_d2, _) =
4169                    logk_q_derivatives(quadctx, 1, row.mass_exit, mu, sigma)?;
4170                let loaded_d1 = 1.0 + loaded_log_d1;
4171                let log_loaded = row.hazard_loaded.ln() + bundle.get(1);
4172                let log_unloaded = row.hazard_unloaded.ln() + bundle.get(0);
4173                let shift = log_loaded.max(log_unloaded);
4174                let loaded_weight = (log_loaded - shift).exp();
4175                let unloaded_weight = (log_unloaded - shift).exp();
4176                let normalizer = loaded_weight + unloaded_weight;
4177                if !(normalizer.is_finite() && normalizer > 0.0) {
4178                    return Err(LatentSurvivalError::NumericalFailure {
4179                        reason: "latent survival exact-event numerator became non-finite under loaded/unloaded hazard decomposition"
4180                            .to_string(),
4181                    });
4182                }
4183                let w_loaded = loaded_weight / normalizer;
4184                let w_unloaded = unloaded_weight / normalizer;
4185                let grad_exit = w_loaded * loaded_d1 + w_unloaded * unloaded_d1;
4186                let d2_exit = w_loaded * (loaded_d2 + loaded_d1 * loaded_d1)
4187                    + w_unloaded * (unloaded_d2 + unloaded_d1 * unloaded_d1)
4188                    - grad_exit * grad_exit;
4189                Ok(LatentSurvivalTimeJet {
4190                    grad_entry: -entry_d1,
4191                    grad_exit,
4192                    neg_hess_entry: entry_d2,
4193                    neg_hess_exit: -d2_exit,
4194                })
4195            } else {
4196                let (exit_d1, exit_d2, _) =
4197                    logk_q_derivatives(quadctx, 1, row.mass_exit, mu, sigma)?;
4198                Ok(LatentSurvivalTimeJet {
4199                    grad_entry: -entry_d1,
4200                    grad_exit: 1.0 + exit_d1,
4201                    neg_hess_entry: entry_d2,
4202                    neg_hess_exit: -exit_d2,
4203                })
4204            }
4205        }
4206        LatentSurvivalEventType::IntervalCensored => {
4207            Err(LatentSurvivalError::UnsupportedConfiguration {
4208                reason:
4209                    "latent survival dynamic time derivatives do not implement interval censoring"
4210                        .to_string(),
4211            })
4212        }
4213    }
4214}
4215
4216fn dense_outer_accumulate<S>(
4217    target: &mut ndarray::ArrayBase<S, ndarray::Ix2>,
4218    weight: f64,
4219    x: ArrayView1<'_, f64>,
4220) where
4221    S: ndarray::DataMut<Elem = f64>,
4222{
4223    for a in 0..x.len() {
4224        let xa = x[a];
4225        if xa == 0.0 {
4226            continue;
4227        }
4228        for b in 0..x.len() {
4229            let xb = x[b];
4230            if xb == 0.0 {
4231                continue;
4232            }
4233            target[[a, b]] += weight * xa * xb;
4234        }
4235    }
4236}
4237
4238fn dense_symmetric_cross_accumulate<S>(
4239    target: &mut ndarray::ArrayBase<S, ndarray::Ix2>,
4240    weight: f64,
4241    x: ArrayView1<'_, f64>,
4242    y: ArrayView1<'_, f64>,
4243) where
4244    S: ndarray::DataMut<Elem = f64>,
4245{
4246    for a in 0..x.len() {
4247        let xa = x[a];
4248        let ya = y[a];
4249        if xa == 0.0 && ya == 0.0 {
4250            continue;
4251        }
4252        for b in 0..x.len() {
4253            let xb = x[b];
4254            let yb = y[b];
4255            let contribution = xa * yb + ya * xb;
4256            if contribution == 0.0 {
4257                continue;
4258            }
4259            target[[a, b]] += weight * contribution;
4260        }
4261    }
4262}
4263
4264fn build_latent_survival_row(
4265    row_index: usize,
4266    hazard_loading: HazardLoading,
4267    event_type: LatentSurvivalEventType,
4268    q_entry: f64,
4269    q_exit: f64,
4270    qdot_exit: f64,
4271    q_right: f64,
4272    unloaded_mass_entry: f64,
4273    unloaded_mass_exit: f64,
4274    unloaded_mass_right: f64,
4275    unloaded_hazard_exit: f64,
4276) -> Result<LatentSurvivalRow, LatentSurvivalError> {
4277    if !(q_entry.is_finite() && q_exit.is_finite()) {
4278        return Err(LatentSurvivalError::NumericalFailure {
4279            reason: format!(
4280                "latent survival requires finite q_entry and q_exit, got q_entry={q_entry}, q_exit={q_exit}"
4281            ),
4282        });
4283    }
4284    if q_exit < q_entry {
4285        return Err(LatentSurvivalError::NumericalFailure {
4286            reason: format!(
4287                "latent survival requires q_exit >= q_entry so cumulative mass is monotone, got q_entry={q_entry}, q_exit={q_exit}"
4288            ),
4289        });
4290    }
4291    if !(unloaded_mass_entry.is_finite()
4292        && unloaded_mass_exit.is_finite()
4293        && unloaded_hazard_exit.is_finite())
4294    {
4295        return Err(LatentSurvivalError::InvalidDataset {
4296            reason: format!(
4297                "latent survival requires finite unloaded components, got entry_mass={unloaded_mass_entry}, exit_mass={unloaded_mass_exit}, exit_hazard={unloaded_hazard_exit}"
4298            ),
4299        });
4300    }
4301    if unloaded_mass_entry < 0.0
4302        || unloaded_mass_exit < unloaded_mass_entry
4303        || unloaded_hazard_exit < 0.0
4304    {
4305        return Err(LatentSurvivalError::InvalidDataset {
4306            reason: format!(
4307                "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}"
4308            ),
4309        });
4310    }
4311    let mass_entry = q_entry.exp();
4312    let mass_exit = q_exit.exp();
4313    let row = match event_type {
4314        LatentSurvivalEventType::RightCensored => {
4315            validate_unloaded_components_for_loading(
4316                "latent-survival",
4317                row_index,
4318                hazard_loading,
4319                unloaded_mass_entry,
4320                unloaded_mass_exit,
4321                Some(unloaded_hazard_exit),
4322            )?;
4323            LatentSurvivalRow::right_censored(
4324                mass_entry,
4325                mass_exit,
4326                unloaded_mass_entry,
4327                unloaded_mass_exit,
4328            )
4329        }
4330        LatentSurvivalEventType::ExactEvent => {
4331            validate_unloaded_components_for_loading(
4332                "latent-survival",
4333                row_index,
4334                hazard_loading,
4335                unloaded_mass_entry,
4336                unloaded_mass_exit,
4337                Some(unloaded_hazard_exit),
4338            )?;
4339            LatentSurvivalRow::exact_event(
4340                mass_entry,
4341                mass_exit,
4342                unloaded_mass_entry,
4343                unloaded_mass_exit,
4344                mass_exit
4345                    * if qdot_exit.is_finite() && qdot_exit > 0.0 {
4346                        qdot_exit
4347                    } else {
4348                        return Err(LatentSurvivalError::NumericalFailure {
4349                            reason: format!(
4350                                "latent survival exact event requires positive finite baseline hazard derivative, got {qdot_exit}"
4351                            ),
4352                        });
4353                    },
4354                unloaded_hazard_exit,
4355            )
4356        }
4357        LatentSurvivalEventType::IntervalCensored => {
4358            // Interval `(L, R]`: `q_exit` carries the LEFT boundary transform
4359            // `log B(L)` (so `mass_left = exp(q_exit)`) and `q_right` the RIGHT
4360            // boundary `log B(R)`. The likelihood is the survival-mass
4361            // difference `log[S(L) − S(R)]`, requiring `B(L) ≤ B(R)` i.e.
4362            // `q_exit ≤ q_right`. No event hazard participates, so the unloaded
4363            // exit hazard must be the full-loading zero (validated below via the
4364            // interval-specific unloaded check at the left/right boundaries).
4365            if !q_right.is_finite() {
4366                return Err(LatentSurvivalError::NumericalFailure {
4367                    reason: format!(
4368                        "latent survival interval row {} requires a finite q_right, got {q_right}",
4369                        row_index + 1
4370                    ),
4371                });
4372            }
4373            if q_right < q_exit {
4374                return Err(LatentSurvivalError::NumericalFailure {
4375                    reason: format!(
4376                        "latent survival interval row {} requires q_right >= q_exit (R >= L) so the \
4377                         survival-mass difference is non-negative, got q_left={q_exit}, q_right={q_right}",
4378                        row_index + 1
4379                    ),
4380                });
4381            }
4382            if !(unloaded_mass_right.is_finite()) || unloaded_mass_right < unloaded_mass_exit {
4383                return Err(LatentSurvivalError::InvalidDataset {
4384                    reason: format!(
4385                        "latent survival interval row {} requires a finite unloaded right mass >= unloaded left mass, got left={unloaded_mass_exit}, right={unloaded_mass_right}",
4386                        row_index + 1
4387                    ),
4388                });
4389            }
4390            // Interval rows carry no exit-event hazard; the loaded/unloaded
4391            // contract is validated by `LatentSurvivalRow::validate` (entry <=
4392            // left <= right monotonicity on both loaded and unloaded masses).
4393            let mass_right = q_right.exp();
4394            LatentSurvivalRow::interval_censored(
4395                mass_entry,
4396                mass_exit,
4397                mass_right,
4398                unloaded_mass_entry,
4399                unloaded_mass_exit,
4400                unloaded_mass_right,
4401            )
4402        }
4403    };
4404    row.validate()
4405        .map_err(|e| LatentSurvivalError::InvalidDataset {
4406            reason: e.to_string(),
4407        })?;
4408    Ok(row)
4409}
4410
4411#[derive(Clone, Copy, Debug)]
4412struct BinaryFromLogSurvival {
4413    log_lik: f64,
4414    /// dℓ/ds where s = log_survival and ℓ = log_lik. For event=1 this is
4415    /// ℓ' = -S/(1-S); for event=0 this is ℓ' = 1 (because ℓ ≡ s).
4416    grad_scale: f64,
4417    /// Coefficient applied to `survival_jet.neg_hessian` (which equals
4418    /// -d²s/dβ²) when assembling the negative Hessian of `wi * log_lik`
4419    /// against β. The Newton accumulator computes
4420    ///     neg_Hess(log_lik) = grad_scale * neg_hessian + outer_scale * score²
4421    /// so by the chain rule this MUST equal `grad_scale` (= ℓ'). Keeping
4422    /// the two fields separate is purely for readability at call sites;
4423    /// the `assert!` in [`binary_from_log_survival`] enforces the
4424    /// equality.
4425    neg_hess_scale: f64,
4426    /// -ℓ''(s). For event=1 this is +S/(1-S)²; for event=0 it is 0.
4427    outer_scale: f64,
4428}
4429
4430/// Exact binary log likelihood from a row log-survival `s`.
4431///
4432/// This value-only path deliberately does not evaluate derivatives: near
4433/// `s = 0`, `log(1-exp(s))` can remain finite after one or more derivatives
4434/// cease to be representable. A likelihood-only caller must not fail because
4435/// of an output it did not request.
4436fn binary_log_likelihood_from_log_survival(
4437    log_survival: f64,
4438    event: u8,
4439) -> Result<f64, LatentSurvivalError> {
4440    match event {
4441        0 => {
4442            if !log_survival.is_finite() || log_survival > 0.0 {
4443                return Err(LatentSurvivalError::NumericalFailure {
4444                    reason: format!(
4445                        "latent-binary requires finite log survival <= 0 for a censored row, got {log_survival:?}"
4446                    ),
4447                });
4448            }
4449            Ok(log_survival)
4450        }
4451        1 => {
4452            if !log_survival.is_finite() || log_survival >= 0.0 {
4453                return Err(LatentSurvivalError::NumericalFailure {
4454                    reason: format!(
4455                        "latent-binary requires finite log survival < 0 for an observed event, got {log_survival:?}"
4456                    ),
4457                });
4458            }
4459            let event_prob = -log_survival.exp_m1();
4460            if !(event_prob.is_finite() && event_prob > 0.0) {
4461                return Err(LatentSurvivalError::NumericalFailure {
4462                    reason: format!(
4463                        "latent-binary event probability is not representable from log survival {log_survival:?}"
4464                    ),
4465                });
4466            }
4467            Ok(event_prob.ln())
4468        }
4469        _ => Err(LatentSurvivalError::InvalidDataset {
4470            reason: format!("latent-binary requires event targets in {{0,1}}, got {event}"),
4471        }),
4472    }
4473}
4474
4475/// Value and first log-survival derivative for the binary row transform.
4476fn binary_from_log_survival_through_first(
4477    log_survival: f64,
4478    event: u8,
4479) -> Result<(f64, f64), LatentSurvivalError> {
4480    let log_lik = binary_log_likelihood_from_log_survival(log_survival, event)?;
4481    if event == 0 {
4482        return Ok((log_lik, 1.0));
4483    }
4484    let odds = (log_survival - log_lik).exp();
4485    if !odds.is_finite() {
4486        return Err(LatentSurvivalError::NumericalFailure {
4487            reason: format!(
4488                "latent-binary log-survival derivative order 1 is not representable at {log_survival:?}: {odds:?}"
4489            ),
4490        });
4491    }
4492    Ok((log_lik, -odds))
4493}
4494
4495/// Analytic source of truth for derivatives of
4496/// `ell(s) = log(1 - exp(s))`, evaluated directly in the log-survival
4497/// coordinate `s < 0`.
4498///
4499/// `P = -expm1(s)` avoids cancellation when survival is near one. Writing the
4500/// derivative algebra in terms of the odds `r = exp(s) / P` avoids the `P²`
4501/// intermediate that previously underflowed before a finite ratio could be
4502/// formed. This base routine computes only through order two; third/fourth
4503/// derivatives are validated lazily by the directional-Hessian paths that
4504/// consume them, so an unrepresentable unused fourth derivative cannot reject
4505/// an otherwise representable likelihood, gradient, and Hessian.
4506fn binary_log_survival_scales(log_survival: f64) -> Result<(f64, f64, f64), LatentSurvivalError> {
4507    let (log_lik, ell_prime) = binary_from_log_survival_through_first(log_survival, 1)?;
4508    let odds = -ell_prime;
4509    let one_plus_odds = 1.0 + odds;
4510    let ell_pp = -odds * one_plus_odds;
4511    let scales = [log_lik, ell_prime, ell_pp];
4512    if let Some((order, value)) = scales
4513        .iter()
4514        .enumerate()
4515        .find(|(_, value)| !value.is_finite())
4516    {
4517        return Err(LatentSurvivalError::NumericalFailure {
4518            reason: format!(
4519                "latent-binary log-survival derivative order {order} is not representable at {log_survival:?}: {value:?}"
4520            ),
4521        });
4522    }
4523    Ok((log_lik, ell_prime, ell_pp))
4524}
4525
4526fn binary_from_log_survival(
4527    log_survival: f64,
4528    event: u8,
4529) -> Result<BinaryFromLogSurvival, LatentSurvivalError> {
4530    if event == 0 {
4531        // ℓ(s) = s ⇒ ℓ' = 1, ℓ'' = ℓ''' = ℓ'''' = 0.
4532        return Ok(BinaryFromLogSurvival {
4533            log_lik: binary_log_likelihood_from_log_survival(log_survival, event)?,
4534            grad_scale: 1.0,
4535            neg_hess_scale: 1.0,
4536            outer_scale: 0.0,
4537        });
4538    }
4539    if event != 1 {
4540        return Err(LatentSurvivalError::InvalidDataset {
4541            reason: format!("latent-binary requires event targets in {{0,1}}, got {event}"),
4542        });
4543    }
4544    let (log_lik, ell_prime, ell_pp) = binary_log_survival_scales(log_survival)?;
4545    let grad_scale = ell_prime;
4546    let neg_hess_scale = ell_prime; // coefficient on (-d²s/dβ²); equals ℓ'.
4547    let outer_scale = -ell_pp;
4548    // The Newton accumulator at the call sites computes
4549    //     neg_Hess(log_lik) = neg_hess_scale * (-d²s/dβ²) + outer_scale * (ds/dβ)²
4550    // For this identity to hold by the chain rule, the coefficient on the
4551    // neg_hessian term must equal ℓ' (== grad_scale). Document the invariant.
4552    assert!(
4553        (grad_scale - neg_hess_scale).abs() <= 1e-15 * grad_scale.abs().max(1.0),
4554        "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"
4555    );
4556    assert!(
4557        outer_scale >= 0.0 || !outer_scale.is_finite(),
4558        "binary_from_log_survival invariant: outer_scale (= -ℓ'') must be non-negative for event=1; got {outer_scale}"
4559    );
4560    Ok(BinaryFromLogSurvival {
4561        log_lik,
4562        grad_scale,
4563        neg_hess_scale,
4564        outer_scale,
4565    })
4566}
4567
4568/// Binary log-survival chain rule through third order. The extra scalar is
4569/// `d outer_scale / ds = -ℓ'''(s)`. The derivative of `grad_scale` is already
4570/// available exactly as `-base.outer_scale = ℓ''(s)`.
4571fn binary_from_log_survival_through_third(
4572    log_survival: f64,
4573    event: u8,
4574) -> Result<(BinaryFromLogSurvival, f64), LatentSurvivalError> {
4575    let base = binary_from_log_survival(log_survival, event)?;
4576    if event == 0 {
4577        return Ok((base, 0.0));
4578    }
4579    let odds = -base.grad_scale;
4580    let ell_pp = -base.outer_scale;
4581    let ell_ppp = ell_pp * (1.0 + 2.0 * odds);
4582    if !ell_ppp.is_finite() {
4583        return Err(LatentSurvivalError::NumericalFailure {
4584            reason: format!(
4585                "latent-binary log-survival derivative order 3 is not representable at {log_survival:?}: {ell_ppp:?}"
4586            ),
4587        });
4588    }
4589    Ok((base, -ell_ppp))
4590}
4591
4592/// Binary log-survival chain rule through fourth order. Returns
4593/// `(base, -ℓ''', -ℓ'''')`, the first and second derivatives of
4594/// `base.outer_scale` with respect to log survival.
4595fn binary_from_log_survival_through_fourth(
4596    log_survival: f64,
4597    event: u8,
4598) -> Result<(BinaryFromLogSurvival, f64, f64), LatentSurvivalError> {
4599    let (base, outer_scale_prime) = binary_from_log_survival_through_third(log_survival, event)?;
4600    if event == 0 {
4601        return Ok((base, 0.0, 0.0));
4602    }
4603    let odds = -base.grad_scale;
4604    let ell_pp = -base.outer_scale;
4605    let ell_pppp = ell_pp * (1.0 + 6.0 * odds + 6.0 * odds * odds);
4606    if !ell_pppp.is_finite() {
4607        return Err(LatentSurvivalError::NumericalFailure {
4608            reason: format!(
4609                "latent-binary log-survival derivative order 4 is not representable at {log_survival:?}: {ell_pppp:?}"
4610            ),
4611        });
4612    }
4613    Ok((base, outer_scale_prime, -ell_pppp))
4614}
4615
4616/// Fitted frailty-scale coordinate used by exact saved latent-survival ALO.
4617///
4618/// A fixed scale is likelihood metadata and therefore contributes no fitted
4619/// coordinate. A learned scale is represented by the exact raw `log_sigma`
4620/// coefficient consumed by the fitter; replay evaluates `sigma = exp(eta)`
4621/// inside the same primary row program.
4622#[derive(Clone, Copy, Debug)]
4623pub enum LatentSurvivalAloSigma {
4624    Fixed(f64),
4625    LearnedLogScale(f64),
4626}
4627
4628/// One saved latent-survival row in the fitter's affine primary coordinates.
4629pub struct LatentSurvivalAloRowInput<'a> {
4630    pub quadrature: &'a QuadratureContext,
4631    pub hazard_loading: HazardLoading,
4632    pub event_code: u8,
4633    pub prior_weight: f64,
4634    pub q_entry: f64,
4635    pub q_exit: f64,
4636    pub qdot_exit: f64,
4637    pub q_right: f64,
4638    pub mu: f64,
4639    pub sigma: LatentSurvivalAloSigma,
4640    pub unloaded_mass_entry: f64,
4641    pub unloaded_mass_exit: f64,
4642    pub unloaded_mass_right: f64,
4643    pub unloaded_hazard_exit: f64,
4644}
4645
4646/// One saved latent-binary row in its three live affine coordinates
4647/// `[q_entry, q_exit, mu]`.
4648pub struct LatentBinaryAloRowInput<'a> {
4649    pub quadrature: &'a QuadratureContext,
4650    pub hazard_loading: HazardLoading,
4651    pub event: u8,
4652    pub prior_weight: f64,
4653    pub q_entry: f64,
4654    pub q_exit: f64,
4655    pub mu: f64,
4656    pub sigma: f64,
4657    pub unloaded_mass_entry: f64,
4658    pub unloaded_mass_exit: f64,
4659}
4660
4661/// Exact NLL row geometry returned by the saved latent-window replay seam.
4662pub struct LatentWindowAloRowGeometry {
4663    pub nll_score: Array1<f64>,
4664    pub observed_hessian: Array2<f64>,
4665    pub coordinate_values: Array1<f64>,
4666}
4667
4668fn validate_saved_alo_weight(weight: f64, context: &str) -> Result<(), String> {
4669    if weight.is_finite() && weight >= 0.0 {
4670        Ok(())
4671    } else {
4672        Err(format!(
4673            "{context} prior weight must be finite and non-negative, got {weight}"
4674        ))
4675    }
4676}
4677
4678fn checked_saved_alo_scale_vector(
4679    values: Array1<f64>,
4680    scale: f64,
4681    context: &str,
4682) -> Result<Array1<f64>, String> {
4683    let mut out = Array1::<f64>::zeros(values.len());
4684    for (axis, value) in values.into_iter().enumerate() {
4685        let product = scale * value;
4686        if !product.is_finite() || (scale != 0.0 && value != 0.0 && product == 0.0) {
4687            return Err(format!(
4688                "{context}[{axis}] is not representable: {scale:?} * {value:?}"
4689            ));
4690        }
4691        out[axis] = product;
4692    }
4693    Ok(out)
4694}
4695
4696fn checked_saved_alo_scale_matrix(
4697    values: Array2<f64>,
4698    scale: f64,
4699    context: &str,
4700) -> Result<Array2<f64>, String> {
4701    let mut out = Array2::<f64>::zeros(values.dim());
4702    for ((row, column), value) in values.indexed_iter() {
4703        let product = scale * value;
4704        if !product.is_finite() || (scale != 0.0 && *value != 0.0 && product == 0.0) {
4705            return Err(format!(
4706                "{context}[{row},{column}] is not representable: {scale:?} * {value:?}"
4707            ));
4708        }
4709        out[[row, column]] = product;
4710    }
4711    Ok(out)
4712}
4713
4714/// Replay one saved latent-survival likelihood row through the exact fitting
4715/// program.
4716///
4717/// Coordinates are `[q_entry, q_exit, qdot_exit, q_right, mu]` followed by
4718/// `log_sigma` only when that scale was learned. The primary authority returns
4719/// log-likelihood score and negative log-likelihood Hessian; this boundary
4720/// applies the row weight and flips only the score to the NLL convention.
4721pub fn latent_survival_alo_row_geometry(
4722    input: LatentSurvivalAloRowInput<'_>,
4723) -> Result<LatentWindowAloRowGeometry, String> {
4724    validate_saved_alo_weight(input.prior_weight, "latent-survival ALO")?;
4725    let mut coordinate_values = vec![
4726        input.q_entry,
4727        input.q_exit,
4728        input.qdot_exit,
4729        input.q_right,
4730        input.mu,
4731    ];
4732    let (sigma, include_log_sigma) = match input.sigma {
4733        LatentSurvivalAloSigma::Fixed(sigma) => (sigma, false),
4734        LatentSurvivalAloSigma::LearnedLogScale(log_sigma) => {
4735            coordinate_values.push(log_sigma);
4736            (log_sigma.exp(), true)
4737        }
4738    };
4739    let coordinate_values = Array1::from_vec(coordinate_values);
4740    let dimension = coordinate_values.len();
4741    if input.prior_weight == 0.0 {
4742        return Ok(LatentWindowAloRowGeometry {
4743            nll_score: Array1::zeros(dimension),
4744            observed_hessian: Array2::zeros((dimension, dimension)),
4745            coordinate_values,
4746        });
4747    }
4748    if !matches!(input.event_code, 0 | 1 | LATENT_SURVIVAL_EVENT_INTERVAL) {
4749        return Err(format!(
4750            "latent-survival ALO event code must be 0, 1, or the interval sentinel {LATENT_SURVIVAL_EVENT_INTERVAL}, got {}",
4751            input.event_code
4752        ));
4753    }
4754    if !sigma.is_finite()
4755        || sigma < 0.0
4756        || (include_log_sigma && (sigma == 0.0 || !coordinate_values[5].is_finite()))
4757    {
4758        return Err(format!(
4759            "latent-survival ALO frailty scale is invalid: sigma={sigma:?}, learned={include_log_sigma}"
4760        ));
4761    }
4762    if coordinate_values.iter().any(|value| !value.is_finite()) {
4763        return Err("latent-survival ALO affine coordinates must be finite".to_string());
4764    }
4765    let event_type = latent_survival_event_type_for(input.event_code);
4766    let row = build_latent_survival_row(
4767        0,
4768        input.hazard_loading,
4769        event_type,
4770        input.q_entry,
4771        input.q_exit,
4772        input.qdot_exit,
4773        input.q_right,
4774        input.unloaded_mass_entry,
4775        input.unloaded_mass_exit,
4776        input.unloaded_mass_right,
4777        input.unloaded_hazard_exit,
4778    )
4779    .map_err(String::from)?;
4780    let (_, log_likelihood_score, negative_log_likelihood_hessian) =
4781        latent_survival_row_primary_gradient_hessian(
4782            input.quadrature,
4783            &row,
4784            LatentSurvivalPrimaryPoint {
4785                q_entry: input.q_entry,
4786                q_exit: input.q_exit,
4787                qdot_exit: input.qdot_exit,
4788                q_right: input.q_right,
4789                mu: input.mu,
4790                sigma,
4791            },
4792            include_log_sigma,
4793        )?;
4794    let nll_score = checked_saved_alo_scale_vector(
4795        log_likelihood_score.slice(s![0..dimension]).to_owned(),
4796        -input.prior_weight,
4797        "latent-survival ALO NLL score",
4798    )?;
4799    let observed_hessian = checked_saved_alo_scale_matrix(
4800        negative_log_likelihood_hessian
4801            .slice(s![0..dimension, 0..dimension])
4802            .to_owned(),
4803        input.prior_weight,
4804        "latent-survival ALO observed Hessian",
4805    )?;
4806    Ok(LatentWindowAloRowGeometry {
4807        nll_score,
4808        observed_hessian,
4809        coordinate_values,
4810    })
4811}
4812
4813/// Replay one saved latent-binary row through the exact right-censored latent
4814/// survival authority and the fitter's analytic binary-from-log-survival
4815/// chain. `W` is the observed NLL Hessian; the score outer product remains a
4816/// separate downstream ALO covariance channel.
4817pub fn latent_binary_alo_row_geometry(
4818    input: LatentBinaryAloRowInput<'_>,
4819) -> Result<LatentWindowAloRowGeometry, String> {
4820    validate_saved_alo_weight(input.prior_weight, "latent-binary ALO")?;
4821    let coordinate_values = Array1::from_vec(vec![input.q_entry, input.q_exit, input.mu]);
4822    const DIMENSION: usize = 3;
4823    if input.prior_weight == 0.0 {
4824        return Ok(LatentWindowAloRowGeometry {
4825            nll_score: Array1::zeros(DIMENSION),
4826            observed_hessian: Array2::zeros((DIMENSION, DIMENSION)),
4827            coordinate_values,
4828        });
4829    }
4830    if input.event > 1 {
4831        return Err(format!(
4832            "latent-binary ALO event must be 0 or 1, got {}",
4833            input.event
4834        ));
4835    }
4836    if !input.sigma.is_finite() || input.sigma < 0.0 {
4837        return Err(format!(
4838            "latent-binary ALO frailty sigma must be finite and non-negative, got {:?}",
4839            input.sigma
4840        ));
4841    }
4842    if coordinate_values.iter().any(|value| !value.is_finite()) {
4843        return Err("latent-binary ALO affine coordinates must be finite".to_string());
4844    }
4845    let row = build_latent_survival_row(
4846        0,
4847        input.hazard_loading,
4848        LatentSurvivalEventType::RightCensored,
4849        input.q_entry,
4850        input.q_exit,
4851        1.0,
4852        input.q_exit,
4853        input.unloaded_mass_entry,
4854        input.unloaded_mass_exit,
4855        0.0,
4856        0.0,
4857    )
4858    .map_err(String::from)?;
4859    let (log_survival, survival_score, survival_negative_hessian) =
4860        latent_survival_row_primary_gradient_hessian(
4861            input.quadrature,
4862            &row,
4863            LatentSurvivalPrimaryPoint {
4864                q_entry: input.q_entry,
4865                q_exit: input.q_exit,
4866                qdot_exit: 1.0,
4867                q_right: input.q_exit,
4868                mu: input.mu,
4869                sigma: input.sigma,
4870            },
4871            false,
4872        )?;
4873    let binary = binary_from_log_survival(log_survival, input.event).map_err(String::from)?;
4874    let primary_indices = [
4875        LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
4876        LATENT_SURVIVAL_PRIMARY_Q_EXIT,
4877        LATENT_SURVIVAL_PRIMARY_MU,
4878    ];
4879    let binary_log_likelihood_score = Array1::from_shape_fn(DIMENSION, |axis| {
4880        binary.grad_scale * survival_score[primary_indices[axis]]
4881    });
4882    let binary_negative_log_likelihood_hessian =
4883        Array2::from_shape_fn((DIMENSION, DIMENSION), |(left, right)| {
4884            let source_left = primary_indices[left];
4885            let source_right = primary_indices[right];
4886            binary.neg_hess_scale * survival_negative_hessian[[source_left, source_right]]
4887                + binary.outer_scale * survival_score[source_left] * survival_score[source_right]
4888        });
4889    let nll_score = checked_saved_alo_scale_vector(
4890        binary_log_likelihood_score,
4891        -input.prior_weight,
4892        "latent-binary ALO NLL score",
4893    )?;
4894    let observed_hessian = checked_saved_alo_scale_matrix(
4895        binary_negative_log_likelihood_hessian,
4896        input.prior_weight,
4897        "latent-binary ALO observed Hessian",
4898    )?;
4899    Ok(LatentWindowAloRowGeometry {
4900        nll_score,
4901        observed_hessian,
4902        coordinate_values,
4903    })
4904}
4905
4906impl LatentBinaryFamily {
4907    /// Assemble the per-row [`LatentSurvivalRow`] for a row treated as a pure
4908    /// right-censored survival contribution (exit time is the censoring
4909    /// boundary, unit exit-hazard derivative, no right / post-exit unloaded
4910    /// mass). Shared by every per-row binary-from-survival pullback reduction;
4911    /// behavior is identical to the previously inlined `RightCensored` call.
4912    fn build_right_censored_row_at(
4913        &self,
4914        row_idx: usize,
4915        q_entry: f64,
4916        q_exit: f64,
4917    ) -> Result<LatentSurvivalRow, LatentSurvivalError> {
4918        build_latent_survival_row(
4919            row_idx,
4920            self.hazard_loading,
4921            LatentSurvivalEventType::RightCensored,
4922            q_entry,
4923            q_exit,
4924            1.0,
4925            q_exit,
4926            self.unloaded_mass_entry[row_idx],
4927            self.unloaded_mass_exit[row_idx],
4928            0.0,
4929            0.0,
4930        )
4931    }
4932
4933    fn joint_slices(&self) -> LatentSurvivalJointSlices {
4934        let p_time = self.x_time_exit.ncols();
4935        let p_mean = self.x_mean.ncols();
4936        LatentSurvivalJointSlices {
4937            time: 0..p_time,
4938            mean: p_time..p_time + p_mean,
4939            log_sigma: None,
4940            total: p_time + p_mean,
4941        }
4942    }
4943
4944    fn row_primary_direction_from_flat(
4945        &self,
4946        row: usize,
4947        slices: &LatentSurvivalJointSlices,
4948        d_beta_flat: &Array1<f64>,
4949    ) -> Array1<f64> {
4950        let mut out = Array1::<f64>::zeros(LATENT_SURVIVAL_PRIMARY_DIM);
4951        let d_time = d_beta_flat.slice(s![slices.time.clone()]);
4952        out[LATENT_SURVIVAL_PRIMARY_Q_ENTRY] = self.x_time_entry.row(row).dot(&d_time);
4953        out[LATENT_SURVIVAL_PRIMARY_Q_EXIT] = self.x_time_exit.row(row).dot(&d_time);
4954        out[LATENT_SURVIVAL_PRIMARY_MU] = self
4955            .x_mean
4956            .dot_row_view(row, d_beta_flat.slice(s![slices.mean.clone()]));
4957        out
4958    }
4959
4960    fn add_pullback_primary_gradient(
4961        &self,
4962        target: &mut Array1<f64>,
4963        row: usize,
4964        slices: &LatentSurvivalJointSlices,
4965        primary_gradient: &Array1<f64>,
4966        weight: f64,
4967    ) -> Result<(), String> {
4968        for (primary_idx, time_vec) in [
4969            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
4970            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
4971        ] {
4972            let scale = checked_weighted_row_value(
4973                weight,
4974                primary_gradient[primary_idx],
4975                row,
4976                "binary primary gradient",
4977            )?;
4978            if scale == 0.0 {
4979                continue;
4980            }
4981            for i in 0..time_vec.len() {
4982                let xi = time_vec[i];
4983                if xi != 0.0 {
4984                    target[slices.time.start + i] += scale * xi;
4985                }
4986            }
4987        }
4988
4989        let mean_scale = checked_weighted_row_value(
4990            weight,
4991            primary_gradient[LATENT_SURVIVAL_PRIMARY_MU],
4992            row,
4993            "binary mean gradient",
4994        )?;
4995        if mean_scale != 0.0 {
4996            self.x_mean
4997                .axpy_row_into(
4998                    row,
4999                    mean_scale,
5000                    &mut target.slice_mut(s![slices.mean.clone()]),
5001                )
5002                .map_err(|error| {
5003                    format!(
5004                        "latent binary mean gradient pullback dimension mismatch: row={row}, mean_slice={:?}, target_len={}, x_mean_cols={}, error={error}",
5005                        slices.mean,
5006                        target.len(),
5007                        self.x_mean.ncols()
5008                    )
5009                })?;
5010        }
5011        Ok(())
5012    }
5013
5014    fn add_pullback_primary_hessian(
5015        &self,
5016        target: &mut Array2<f64>,
5017        row: usize,
5018        slices: &LatentSurvivalJointSlices,
5019        primary_hessian: &Array2<f64>,
5020    ) {
5021        {
5022            let time_target = &mut target.slice_mut(s![slices.time.clone(), slices.time.clone()]);
5023            dense_outer_accumulate(
5024                time_target,
5025                primary_hessian[[
5026                    LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
5027                    LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
5028                ]],
5029                self.x_time_entry.row(row),
5030            );
5031            dense_outer_accumulate(
5032                time_target,
5033                primary_hessian[[
5034                    LATENT_SURVIVAL_PRIMARY_Q_EXIT,
5035                    LATENT_SURVIVAL_PRIMARY_Q_EXIT,
5036                ]],
5037                self.x_time_exit.row(row),
5038            );
5039            dense_symmetric_cross_accumulate(
5040                time_target,
5041                primary_hessian[[
5042                    LATENT_SURVIVAL_PRIMARY_Q_ENTRY,
5043                    LATENT_SURVIVAL_PRIMARY_Q_EXIT,
5044                ]],
5045                self.x_time_entry.row(row),
5046                self.x_time_exit.row(row),
5047            );
5048        }
5049
5050        let mean_weight = primary_hessian[[LATENT_SURVIVAL_PRIMARY_MU, LATENT_SURVIVAL_PRIMARY_MU]];
5051        self.x_mean
5052            .syr_row_into_view(
5053                row,
5054                mean_weight,
5055                target.slice_mut(s![slices.mean.clone(), slices.mean.clone()]),
5056            )
5057            .unwrap_or_else(|error| {
5058                // SAFETY: `slices.mean` × `slices.mean` slab sized at
5059                // construction to `x_mean.ncols()` × `x_mean.ncols()`;
5060                // an error here is caller-side shape drift, an invariant
5061                // violation. A swallowed sentinel would silently corrupt the
5062                // joint Hessian, so fail loudly instead.
5063                panic!(
5064                    "latent binary mean Hessian pullback dimension mismatch: row={row}, mean_slice={:?}, target_dim={:?}, x_mean_cols={}, error={error}",
5065                    slices.mean,
5066                    target.dim(),
5067                    self.x_mean.ncols()
5068                )
5069            });
5070
5071        let mean_row = self
5072            .x_mean
5073            .try_row_chunk(row..row + 1)
5074            .unwrap_or_else(|error| {
5075                // SAFETY: row index comes from the enclosing `0..n` loop
5076                // bound by `self.x_mean.nrows()`, so `row..row+1` is
5077                // always a valid single-row chunk.
5078                panic!(
5079                    "latent binary mean pullback row chunk failed: row={row}, x_mean_rows={}, x_mean_cols={}, error={error}",
5080                    self.x_mean.nrows(),
5081                    self.x_mean.ncols()
5082                )
5083            });
5084        let mean_vec = mean_row.row(0);
5085        for (primary_idx, time_vec) in [
5086            (LATENT_SURVIVAL_PRIMARY_Q_ENTRY, self.x_time_entry.row(row)),
5087            (LATENT_SURVIVAL_PRIMARY_Q_EXIT, self.x_time_exit.row(row)),
5088        ] {
5089            let weight = primary_hessian[[primary_idx, LATENT_SURVIVAL_PRIMARY_MU]];
5090            if weight == 0.0 {
5091                continue;
5092            }
5093            for i in 0..time_vec.len() {
5094                let xi = time_vec[i];
5095                if xi == 0.0 {
5096                    continue;
5097                }
5098                for j in 0..mean_vec.len() {
5099                    let xj = mean_vec[j];
5100                    if xj == 0.0 {
5101                        continue;
5102                    }
5103                    target[[slices.time.start + i, slices.mean.start + j]] += weight * xi * xj;
5104                    target[[slices.mean.start + j, slices.time.start + i]] += weight * xj * xi;
5105                }
5106            }
5107        }
5108    }
5109
5110    fn evaluate_exact_newton_joint_dense(
5111        &self,
5112        block_states: &[ParameterBlockState],
5113    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5114        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")
5115            .map_err(String::from)?;
5116        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
5117        let slices = self.joint_slices();
5118        let mut ll = CompensatedRowSum::default();
5119        let mut gradient = Array1::<f64>::zeros(slices.total);
5120        let mut hessian = Array2::<f64>::zeros((slices.total, slices.total));
5121        for row_idx in 0..self.event_target.len() {
5122            let wi = weights.at(row_idx);
5123            if wi == 0.0 {
5124                continue;
5125            }
5126            let row =
5127                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
5128            let (row_log_survival, survival_gradient, survival_hessian) =
5129                latent_survival_row_primary_gradient_hessian(
5130                    &self.quadctx,
5131                    &row,
5132                    LatentSurvivalPrimaryPoint {
5133                        q_entry: q_entry[row_idx],
5134                        q_exit: q_exit[row_idx],
5135                        qdot_exit: 1.0,
5136                        q_right: q_exit[row_idx],
5137                        mu: mu[row_idx],
5138                        sigma: self.latent_sd,
5139                    },
5140                    false,
5141                )?;
5142            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
5143            ll.add(checked_weighted_row_value(
5144                wi,
5145                binary.log_lik,
5146                row_idx,
5147                "binary log likelihood",
5148            )?);
5149            let primary_gradient = binary.grad_scale * &survival_gradient;
5150            let mut primary_hessian = binary.grad_scale * survival_hessian;
5151            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5152                for b in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5153                    primary_hessian[[a, b]] +=
5154                        binary.outer_scale * survival_gradient[a] * survival_gradient[b];
5155                }
5156            }
5157            self.add_pullback_primary_gradient(
5158                &mut gradient,
5159                row_idx,
5160                &slices,
5161                &primary_gradient,
5162                wi,
5163            )?;
5164            let weighted_primary_hessian = checked_weighted_row_matrix(
5165                wi,
5166                &primary_hessian,
5167                row_idx,
5168                "binary primary Hessian",
5169            )?;
5170            self.add_pullback_primary_hessian(
5171                &mut hessian,
5172                row_idx,
5173                &slices,
5174                &weighted_primary_hessian,
5175            );
5176        }
5177        let ll = require_finite_likelihood_scalar(ll.value(), "binary log likelihood")?;
5178        require_finite_likelihood_vector(&gradient, "binary gradient")?;
5179        require_finite_likelihood_matrix(&hessian, "binary Hessian")?;
5180        Ok((ll, gradient, hessian))
5181    }
5182
5183    /// Per-row residuals of the unpenalized NLL with respect to the baseline
5184    /// time-block offsets `(entry, exit)`.
5185    ///
5186    /// The latent-binary deployment likelihood is a monotone scalar transform
5187    /// `ℓ_bin = b(log S_row)` of the latent-survival row log-survival, so by the
5188    /// chain rule `∂ℓ_bin/∂q_ch = b'(log S)·∂(log S)/∂q_ch = grad_scale·g_ch`,
5189    /// where `g_ch` are the `Q_ENTRY`/`Q_EXIT` components of the survival row
5190    /// primary gradient. The baseline θ enters only the additive entry/exit time
5191    /// offsets (`q̇_exit` is held at the constant deployment derivative `1`, so
5192    /// the derivative channel carries no baseline offset and its residual is 0).
5193    /// Sampleweight-scaled to match the [`OffsetChannelResiduals`] contract.
5194    pub fn offset_channel_residuals(
5195        &self,
5196        block_states: &[ParameterBlockState],
5197    ) -> Result<crate::survival::OffsetChannelResiduals, LatentSurvivalError> {
5198        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")?;
5199        let n = self.event_target.len();
5200        // `split_time_eta` returns a typed block-count error before indexing.
5201        // Missing state is never translated into zero residuals, because that
5202        // would falsely certify the enclosing baseline optimization.
5203        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
5204        let mut entry = Array1::<f64>::zeros(n);
5205        let mut exit = Array1::<f64>::zeros(n);
5206        for row_idx in 0..n {
5207            let wi = weights.at(row_idx);
5208            if wi == 0.0 {
5209                continue;
5210            }
5211            let row =
5212                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
5213            let (row_log_survival, survival_gradient, _) =
5214                latent_survival_row_primary_gradient_hessian(
5215                    &self.quadctx,
5216                    &row,
5217                    LatentSurvivalPrimaryPoint {
5218                        q_entry: q_entry[row_idx],
5219                        q_exit: q_exit[row_idx],
5220                        qdot_exit: 1.0,
5221                        q_right: q_exit[row_idx],
5222                        mu: mu[row_idx],
5223                        sigma: self.latent_sd,
5224                    },
5225                    false,
5226                )
5227                .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
5228            let (_, grad_scale) = binary_from_log_survival_through_first(
5229                row_log_survival,
5230                self.event_target[row_idx],
5231            )?;
5232            // ∂NLL/∂o_ch = −w · grad_scale · ∂(log S)/∂q_ch.
5233            entry[row_idx] = -checked_weighted_row_value(
5234                wi,
5235                grad_scale * survival_gradient[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
5236                row_idx,
5237                "binary entry-offset score",
5238            )
5239            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
5240            exit[row_idx] = -checked_weighted_row_value(
5241                wi,
5242                grad_scale * survival_gradient[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
5243                row_idx,
5244                "binary exit-offset score",
5245            )
5246            .map_err(|reason| LatentSurvivalError::NumericalFailure { reason })?;
5247        }
5248        Ok(crate::survival::OffsetChannelResiduals {
5249            exit,
5250            entry,
5251            derivative: Array1::<f64>::zeros(n),
5252            // Latent-binary deployment has no interval upper bound; the `R`
5253            // channel is structurally absent (every row is right-censored).
5254            right: Array1::<f64>::zeros(n),
5255        })
5256    }
5257
5258    fn exact_newton_joint_hessian_directional_derivative_dense(
5259        &self,
5260        block_states: &[ParameterBlockState],
5261        d_beta_flat: &Array1<f64>,
5262    ) -> Result<Array2<f64>, String> {
5263        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")
5264            .map_err(String::from)?;
5265        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
5266        let slices = self.joint_slices();
5267        if d_beta_flat.len() != slices.total {
5268            return Err(format!(
5269                "latent binary joint dH direction length mismatch: got {}, expected {}",
5270                d_beta_flat.len(),
5271                slices.total
5272            ));
5273        }
5274        let mut out = Array2::<f64>::zeros((slices.total, slices.total));
5275        for row_idx in 0..self.event_target.len() {
5276            let wi = weights.at(row_idx);
5277            if wi == 0.0 {
5278                continue;
5279            }
5280            let row =
5281                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
5282            let direction = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_flat);
5283            // OneSeed already carries the ordinary value/gradient/Hessian in
5284            // its base part.  Reuse those channels for the binary outer chain
5285            // instead of running a separate Order2 row first.
5286            let row_jet = latent_survival_row_primary_one_seed_fixed_sigma(
5287                &self.quadctx,
5288                &row,
5289                LatentSurvivalPrimaryPoint {
5290                    q_entry: q_entry[row_idx],
5291                    q_exit: q_exit[row_idx],
5292                    qdot_exit: 1.0,
5293                    q_right: q_exit[row_idx],
5294                    mu: mu[row_idx],
5295                    sigma: self.latent_sd,
5296                },
5297                &direction,
5298            )?;
5299            let (binary, outer_scale_prime) = binary_from_log_survival_through_third(
5300                row_jet.base.value(),
5301                self.event_target[row_idx],
5302            )?;
5303            let base_gradient = row_jet.base.g();
5304            let base_hessian = row_jet.base.h();
5305            let contracted_third = row_jet.contracted_third();
5306            let survival_gradient = Array1::from_shape_fn(LATENT_SURVIVAL_PRIMARY_DIM, |a| {
5307                if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA {
5308                    base_gradient[a]
5309                } else {
5310                    0.0
5311                }
5312            });
5313            let survival_hessian = Array2::from_shape_fn(
5314                (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
5315                |(a, b)| {
5316                    if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
5317                        && b < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
5318                    {
5319                        -base_hessian[a][b]
5320                    } else {
5321                        0.0
5322                    }
5323                },
5324            );
5325            let third = Array2::from_shape_fn(
5326                (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
5327                |(a, b)| {
5328                    if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
5329                        && b < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
5330                    {
5331                        -contracted_third[a][b]
5332                    } else {
5333                        0.0
5334                    }
5335                },
5336            );
5337            let g_u = -survival_hessian.dot(&direction);
5338            let t_u = survival_gradient.dot(&direction);
5339            let mut primary = binary.grad_scale * third;
5340            primary.scaled_add(-binary.outer_scale * t_u, &survival_hessian);
5341            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5342                for b in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5343                    primary[[a, b]] +=
5344                        outer_scale_prime * t_u * survival_gradient[a] * survival_gradient[b]
5345                            + binary.outer_scale
5346                                * (g_u[a] * survival_gradient[b] + survival_gradient[a] * g_u[b]);
5347                }
5348            }
5349            let weighted_primary =
5350                checked_weighted_row_matrix(wi, &primary, row_idx, "binary contracted third")?;
5351            self.add_pullback_primary_hessian(&mut out, row_idx, &slices, &weighted_primary);
5352        }
5353        require_finite_likelihood_matrix(&out, "binary directional Hessian derivative")?;
5354        Ok(out)
5355    }
5356
5357    fn exact_newton_joint_hessian_second_directional_derivative_dense(
5358        &self,
5359        block_states: &[ParameterBlockState],
5360        d_beta_u_flat: &Array1<f64>,
5361        d_beta_v_flat: &Array1<f64>,
5362    ) -> Result<Array2<f64>, String> {
5363        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")
5364            .map_err(String::from)?;
5365        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
5366        let slices = self.joint_slices();
5367        if d_beta_u_flat.len() != slices.total || d_beta_v_flat.len() != slices.total {
5368            return Err(format!(
5369                "latent binary joint d2H direction length mismatch: got {} and {}, expected {}",
5370                d_beta_u_flat.len(),
5371                d_beta_v_flat.len(),
5372                slices.total
5373            ));
5374        }
5375        let mut out = Array2::<f64>::zeros((slices.total, slices.total));
5376        for row_idx in 0..self.event_target.len() {
5377            let wi = weights.at(row_idx);
5378            if wi == 0.0 {
5379                continue;
5380            }
5381            let row =
5382                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
5383            let direction_u = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_u_flat);
5384            let direction_v = self.row_primary_direction_from_flat(row_idx, &slices, d_beta_v_flat);
5385            // One TwoSeed row contains the base VGH, both one-seed Hessians,
5386            // and the mixed two-seed Hessian.  The previous composition ran
5387            // four complete rows (Order2 + OneSeed(u) + OneSeed(v) +
5388            // TwoSeed(u,v)) to recover these same channels.
5389            let row_jet = latent_survival_row_primary_two_seed_fixed_sigma(
5390                &self.quadctx,
5391                &row,
5392                LatentSurvivalPrimaryPoint {
5393                    q_entry: q_entry[row_idx],
5394                    q_exit: q_exit[row_idx],
5395                    qdot_exit: 1.0,
5396                    q_right: q_exit[row_idx],
5397                    mu: mu[row_idx],
5398                    sigma: self.latent_sd,
5399                },
5400                &direction_u,
5401                &direction_v,
5402            )?;
5403            let (binary, outer_scale_prime, outer_scale_second) =
5404                binary_from_log_survival_through_fourth(
5405                    row_jet.base.value(),
5406                    self.event_target[row_idx],
5407                )?;
5408            let base_gradient = row_jet.base.g();
5409            let base_hessian = row_jet.base.h();
5410            let contracted_third_u = row_jet.eps.h();
5411            let contracted_third_v = row_jet.del.h();
5412            let contracted_fourth = row_jet.contracted_fourth();
5413            let survival_gradient = Array1::from_shape_fn(LATENT_SURVIVAL_PRIMARY_DIM, |a| {
5414                if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA {
5415                    base_gradient[a]
5416                } else {
5417                    0.0
5418                }
5419            });
5420            let pad_matrix =
5421                |matrix: &[[f64; LATENT_SURVIVAL_PRIMARY_LOG_SIGMA];
5422                      LATENT_SURVIVAL_PRIMARY_LOG_SIGMA]| {
5423                    Array2::from_shape_fn(
5424                        (LATENT_SURVIVAL_PRIMARY_DIM, LATENT_SURVIVAL_PRIMARY_DIM),
5425                        |(a, b)| {
5426                            if a < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
5427                                && b < LATENT_SURVIVAL_PRIMARY_LOG_SIGMA
5428                            {
5429                                -matrix[a][b]
5430                            } else {
5431                                0.0
5432                            }
5433                        },
5434                    )
5435                };
5436            let survival_hessian = pad_matrix(&base_hessian);
5437            let third_u = pad_matrix(&contracted_third_u);
5438            let third_v = pad_matrix(&contracted_third_v);
5439            let fourth = pad_matrix(&contracted_fourth);
5440            let g_u = -survival_hessian.dot(&direction_u);
5441            let g_v = -survival_hessian.dot(&direction_v);
5442            let g_uv = -third_v.dot(&direction_u);
5443            let t_u = survival_gradient.dot(&direction_u);
5444            let t_v = survival_gradient.dot(&direction_v);
5445            let l_uv = -direction_u.dot(&survival_hessian.dot(&direction_v));
5446            let grad_scale_prime = -binary.outer_scale;
5447            let grad_scale_second = -outer_scale_prime;
5448            let c_u = grad_scale_prime * t_u;
5449            let c_v = grad_scale_prime * t_v;
5450            let c_uv = grad_scale_second * t_u * t_v + grad_scale_prime * l_uv;
5451            let o_u = outer_scale_prime * t_u;
5452            let o_v = outer_scale_prime * t_v;
5453            let o_uv = outer_scale_second * t_u * t_v + outer_scale_prime * l_uv;
5454            let mut primary = binary.grad_scale * fourth;
5455            primary.scaled_add(c_u, &third_v);
5456            primary.scaled_add(c_v, &third_u);
5457            primary.scaled_add(c_uv, &survival_hessian);
5458            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5459                for b in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5460                    primary[[a, b]] += o_uv * survival_gradient[a] * survival_gradient[b]
5461                        + o_v * (g_u[a] * survival_gradient[b] + survival_gradient[a] * g_u[b])
5462                        + o_u * (g_v[a] * survival_gradient[b] + survival_gradient[a] * g_v[b])
5463                        + binary.outer_scale
5464                            * (g_uv[a] * survival_gradient[b]
5465                                + g_u[a] * g_v[b]
5466                                + g_v[a] * g_u[b]
5467                                + survival_gradient[a] * g_uv[b]);
5468                }
5469            }
5470            let weighted_primary =
5471                checked_weighted_row_matrix(wi, &primary, row_idx, "binary contracted fourth")?;
5472            self.add_pullback_primary_hessian(&mut out, row_idx, &slices, &weighted_primary);
5473        }
5474        require_finite_likelihood_matrix(&out, "binary second directional Hessian derivative")?;
5475        Ok(out)
5476    }
5477}
5478
5479/// Shared interface that both `LatentSurvivalFamily` and `LatentBinaryFamily`
5480/// expose to the joint Hessian workspace.
5481///
5482/// The two families produce the same `ExactNewtonJointHessianWorkspace`
5483/// shape — five of the six workspace methods are pure delegations to a
5484/// matching family method (dense evaluation, directional derivatives, and the
5485/// `slices` cache). The only family-specific piece is the per-row matvec body:
5486/// the survival family iterates over real (entry, exit, ḋ) triples and may
5487/// carry a log-σ block, while the binary family rewrites the same row kernel
5488/// through `binary_from_log_survival(·)` to recover the per-row binary
5489/// gradient/Hessian. That single difference is captured by `ws_matvec_into`;
5490/// every other method is shared by the generic `LatentHessianWorkspace<F>`
5491/// below.
5492trait LatentJointHessianFamily {
5493    fn ws_joint_slices(&self) -> LatentSurvivalJointSlices;
5494
5495    fn ws_evaluate_dense(
5496        &self,
5497        block_states: &[ParameterBlockState],
5498    ) -> Result<(f64, Array1<f64>, Array2<f64>), String>;
5499
5500    fn ws_dh_directional(
5501        &self,
5502        block_states: &[ParameterBlockState],
5503        d_beta_flat: &Array1<f64>,
5504    ) -> Result<Array2<f64>, String>;
5505
5506    fn ws_dh_second_directional(
5507        &self,
5508        block_states: &[ParameterBlockState],
5509        d_beta_u: &Array1<f64>,
5510        d_beta_v: &Array1<f64>,
5511    ) -> Result<Array2<f64>, String>;
5512
5513    /// Family-specific per-row Hessian matvec body, hoisted out of the
5514    /// workspace impl. Writes `out := H · v` (with `out.fill(0.0)` already
5515    /// performed by the caller) using the family's row kernel.
5516    fn ws_matvec_into(
5517        &self,
5518        slices: &LatentSurvivalJointSlices,
5519        block_states: &[ParameterBlockState],
5520        v: &Array1<f64>,
5521        out: &mut Array1<f64>,
5522    ) -> Result<bool, String>;
5523
5524    /// Family-name fragment used in the workspace's dimension-mismatch error
5525    /// message, so callers still see "latent survival …" / "latent binary …"
5526    /// after the workspace impl was unified.
5527    fn ws_label() -> &'static str;
5528}
5529
5530impl LatentJointHessianFamily for LatentSurvivalFamily {
5531    fn ws_joint_slices(&self) -> LatentSurvivalJointSlices {
5532        self.joint_slices()
5533    }
5534
5535    fn ws_evaluate_dense(
5536        &self,
5537        block_states: &[ParameterBlockState],
5538    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5539        self.evaluate_exact_newton_joint_dense(block_states)
5540    }
5541
5542    fn ws_dh_directional(
5543        &self,
5544        block_states: &[ParameterBlockState],
5545        d_beta_flat: &Array1<f64>,
5546    ) -> Result<Array2<f64>, String> {
5547        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
5548    }
5549
5550    fn ws_dh_second_directional(
5551        &self,
5552        block_states: &[ParameterBlockState],
5553        d_beta_u: &Array1<f64>,
5554        d_beta_v: &Array1<f64>,
5555    ) -> Result<Array2<f64>, String> {
5556        self.exact_newton_joint_hessian_second_directional_derivative_dense(
5557            block_states,
5558            d_beta_u,
5559            d_beta_v,
5560        )
5561    }
5562
5563    fn ws_matvec_into(
5564        &self,
5565        slices: &LatentSurvivalJointSlices,
5566        block_states: &[ParameterBlockState],
5567        v: &Array1<f64>,
5568        out: &mut Array1<f64>,
5569    ) -> Result<bool, String> {
5570        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
5571            .map_err(String::from)?;
5572        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
5573        let q_right = self.time_q_right(block_states)?;
5574        let sigma = self.latent_sd(block_states)?;
5575        let include_log_sigma = slices.log_sigma.is_some();
5576        for row_idx in 0..self.event_target.len() {
5577            let wi = weights.at(row_idx);
5578            if wi == 0.0 {
5579                continue;
5580            }
5581            let row = self.build_row_at(
5582                row_idx,
5583                q_entry[row_idx],
5584                q_exit[row_idx],
5585                qdot_exit[row_idx],
5586                q_right[row_idx],
5587            )?;
5588            let (_, _, primary_hessian) = latent_survival_row_primary_gradient_hessian(
5589                &self.quadctx,
5590                &row,
5591                LatentSurvivalPrimaryPoint {
5592                    q_entry: q_entry[row_idx],
5593                    q_exit: q_exit[row_idx],
5594                    qdot_exit: qdot_exit[row_idx],
5595                    q_right: q_right[row_idx],
5596                    mu: mu[row_idx],
5597                    sigma,
5598                },
5599                include_log_sigma,
5600            )?;
5601            let primary_dir = self.row_primary_direction_from_flat(row_idx, slices, v);
5602            let primary_hv = primary_hessian.dot(&primary_dir);
5603            self.add_pullback_primary_gradient(out, row_idx, slices, &primary_hv, wi)?;
5604        }
5605        require_finite_likelihood_vector(out, "Hessian matvec")?;
5606        Ok(true)
5607    }
5608
5609    fn ws_label() -> &'static str {
5610        "survival"
5611    }
5612}
5613
5614impl LatentJointHessianFamily for LatentBinaryFamily {
5615    fn ws_joint_slices(&self) -> LatentSurvivalJointSlices {
5616        self.joint_slices()
5617    }
5618
5619    fn ws_evaluate_dense(
5620        &self,
5621        block_states: &[ParameterBlockState],
5622    ) -> Result<(f64, Array1<f64>, Array2<f64>), String> {
5623        self.evaluate_exact_newton_joint_dense(block_states)
5624    }
5625
5626    fn ws_dh_directional(
5627        &self,
5628        block_states: &[ParameterBlockState],
5629        d_beta_flat: &Array1<f64>,
5630    ) -> Result<Array2<f64>, String> {
5631        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
5632    }
5633
5634    fn ws_dh_second_directional(
5635        &self,
5636        block_states: &[ParameterBlockState],
5637        d_beta_u: &Array1<f64>,
5638        d_beta_v: &Array1<f64>,
5639    ) -> Result<Array2<f64>, String> {
5640        self.exact_newton_joint_hessian_second_directional_derivative_dense(
5641            block_states,
5642            d_beta_u,
5643            d_beta_v,
5644        )
5645    }
5646
5647    fn ws_matvec_into(
5648        &self,
5649        slices: &LatentSurvivalJointSlices,
5650        block_states: &[ParameterBlockState],
5651        v: &Array1<f64>,
5652        out: &mut Array1<f64>,
5653    ) -> Result<bool, String> {
5654        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")
5655            .map_err(String::from)?;
5656        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
5657        for row_idx in 0..self.event_target.len() {
5658            let wi = weights.at(row_idx);
5659            if wi == 0.0 {
5660                continue;
5661            }
5662            let row =
5663                self.build_right_censored_row_at(row_idx, q_entry[row_idx], q_exit[row_idx])?;
5664            let (row_log_survival, survival_gradient, survival_hessian) =
5665                latent_survival_row_primary_gradient_hessian(
5666                    &self.quadctx,
5667                    &row,
5668                    LatentSurvivalPrimaryPoint {
5669                        q_entry: q_entry[row_idx],
5670                        q_exit: q_exit[row_idx],
5671                        qdot_exit: 1.0,
5672                        q_right: q_exit[row_idx],
5673                        mu: mu[row_idx],
5674                        sigma: self.latent_sd,
5675                    },
5676                    false,
5677                )?;
5678            let binary = binary_from_log_survival(row_log_survival, self.event_target[row_idx])?;
5679            let primary_dir = self.row_primary_direction_from_flat(row_idx, slices, v);
5680            let mut primary_hv = binary.grad_scale * survival_hessian.dot(&primary_dir);
5681            let outer_dot = survival_gradient.dot(&primary_dir);
5682            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
5683                primary_hv[a] += binary.outer_scale * survival_gradient[a] * outer_dot;
5684            }
5685            self.add_pullback_primary_gradient(out, row_idx, slices, &primary_hv, wi)?;
5686        }
5687        require_finite_likelihood_vector(out, "binary Hessian matvec")?;
5688        Ok(true)
5689    }
5690
5691    fn ws_label() -> &'static str {
5692        "binary"
5693    }
5694}
5695
5696/// Joint exact-Newton Hessian workspace shared by `LatentSurvivalFamily` and
5697/// `LatentBinaryFamily`. The two families plug into the workspace via
5698/// `LatentJointHessianFamily`; this struct holds the shared bookkeeping
5699/// (block states + cached slices) and routes every trait method either through
5700/// a thin family delegation or through the family's `ws_matvec_into` row
5701/// kernel.
5702struct LatentHessianWorkspace<F: LatentJointHessianFamily> {
5703    family: F,
5704    block_states: Vec<ParameterBlockState>,
5705    slices: LatentSurvivalJointSlices,
5706}
5707
5708impl<F: LatentJointHessianFamily> LatentHessianWorkspace<F> {
5709    fn new(family: F, block_states: Vec<ParameterBlockState>) -> Self {
5710        let slices = family.ws_joint_slices();
5711        Self {
5712            family,
5713            block_states,
5714            slices,
5715        }
5716    }
5717}
5718
5719impl<F> ExactNewtonJointHessianWorkspace for LatentHessianWorkspace<F>
5720where
5721    F: LatentJointHessianFamily + Send + Sync + 'static,
5722{
5723    fn warm_up_outer_caches_for_mode(
5724        &self,
5725        eval_mode: gam_problem::EvalMode,
5726    ) -> Result<(), String> {
5727        match eval_mode {
5728            gam_problem::EvalMode::ValueOnly
5729            | gam_problem::EvalMode::ValueAndGradient
5730            | gam_problem::EvalMode::ValueGradientHessian => Ok(()),
5731        }
5732    }
5733
5734    fn hessian_dense(&self) -> Result<Option<Array2<f64>>, String> {
5735        self.family
5736            .ws_evaluate_dense(&self.block_states)
5737            .map(|(_, _, hessian)| Some(hessian))
5738    }
5739
5740    fn hessian_matvec(&self, v: &Array1<f64>) -> Result<Option<Array1<f64>>, String> {
5741        let mut out = Array1::<f64>::zeros(self.slices.total);
5742        self.hessian_matvec_into(v, &mut out)?;
5743        Ok(Some(out))
5744    }
5745
5746    fn hessian_matvec_into(&self, v: &Array1<f64>, out: &mut Array1<f64>) -> Result<bool, String> {
5747        if v.len() != self.slices.total || out.len() != self.slices.total {
5748            return Err(format!(
5749                "latent {} Hessian matvec dimension mismatch: v={} out={} expected={}",
5750                F::ws_label(),
5751                v.len(),
5752                out.len(),
5753                self.slices.total
5754            ));
5755        }
5756        out.fill(0.0);
5757        self.family
5758            .ws_matvec_into(&self.slices, &self.block_states, v, out)
5759    }
5760
5761    fn hessian_diagonal(&self) -> Result<Option<Array1<f64>>, String> {
5762        let dense = self.family.ws_evaluate_dense(&self.block_states)?.2;
5763        Ok(Some(dense.diag().to_owned()))
5764    }
5765
5766    fn directional_derivative(
5767        &self,
5768        d_beta_flat: &Array1<f64>,
5769    ) -> Result<Option<Array2<f64>>, String> {
5770        self.family
5771            .ws_dh_directional(&self.block_states, d_beta_flat)
5772            .map(Some)
5773    }
5774
5775    fn second_directional_derivative(
5776        &self,
5777        d_beta_u: &Array1<f64>,
5778        d_beta_v: &Array1<f64>,
5779    ) -> Result<Option<Array2<f64>>, String> {
5780        self.family
5781            .ws_dh_second_directional(&self.block_states, d_beta_u, d_beta_v)
5782            .map(Some)
5783    }
5784}
5785
5786type LatentSurvivalHessianWorkspace = LatentHessianWorkspace<LatentSurvivalFamily>;
5787type LatentBinaryHessianWorkspace = LatentHessianWorkspace<LatentBinaryFamily>;
5788
5789impl CustomFamily for LatentSurvivalFamily {
5790    // Latent survival fits keep the self-limiting Jeffreys/Firth curvature
5791    // active for their under-identification regime. The trait default flipped to
5792    // OFF in gam#1395 (flat-prior exact-Newton objective); opt back in here.
5793    fn joint_jeffreys_term_required(&self) -> bool {
5794        true
5795    }
5796
5797    fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
5798        true
5799    }
5800
5801    fn has_explicit_joint_hessian(&self) -> bool {
5802        true
5803    }
5804
5805    /// Route the pre-fit identifiability audit CHANNEL-AWARE across the latent
5806    /// survival blocks. The time-transform baseline `q(a)`, the frailty mean
5807    /// `μ = Xβ`, and the learnable log-σ scale are three STRUCTURALLY DISTINCT
5808    /// outputs — block-diagonal entries of the true joint Jacobian
5809    /// `blkdiag(X_time, X_mean, X_logσ)`, full rank `Σ p_b`. The blocks are
5810    /// hand-built (no `jacobian_callback`), so without this the flat single-
5811    /// channel audit runs.
5812    ///
5813    /// That flat audit is fatal for a learnable scale: the `log_sigma` block's
5814    /// design is a single structural constant-of-ones column (`eta = design·β`
5815    /// broadcasts the scalar log-σ to every row — see `build_log_sigma_blockspec`).
5816    /// The flat RRQR sees that constant as a repeated shared basis and mistakes
5817    /// it for a cross-block ALIAS with the `mean` intercept, then — because the
5818    /// log-σ block carries the lowest gauge priority — attributes the drop to
5819    /// `log_sigma[0]` and demotes it below the rank tolerance. That both FREEZES
5820    /// the frailty scale at its seed (the only handle on σ is deleted) and leaves
5821    /// the reduced spec width one short of the family's raw joint Hessian, so the
5822    /// outer LAML logdet aborts with `joint exact-newton Hessian validation …:
5823    /// got (p+1)×(p+1), expected p×p`. The mean intercept and log-σ are NOT
5824    /// aliased: they parameterise the frailty MEAN and VARIANCE respectively and
5825    /// are jointly identified; the collinearity is an artefact of auditing a
5826    /// nonlinear scale channel as if it were a linear predictor. Placing each
5827    /// block on its own output channel makes the audit see the genuine
5828    /// block-diagonal structure. Mirrors
5829    /// `SurvivalLocationScaleFamily::output_channel_assignment`.
5830    fn output_channel_assignment(&self, specs: &[ParameterBlockSpec]) -> Option<Vec<usize>> {
5831        Some(
5832            specs
5833                .iter()
5834                .map(|spec| match spec.name.as_str() {
5835                    "time_transform" => 0,
5836                    "mean" => 1,
5837                    "log_sigma" => 2,
5838                    _ => 0,
5839                })
5840                .collect(),
5841        )
5842    }
5843
5844    /// Engage the inner self-vanishing Levenberg–Marquardt μ on a full-rank but
5845    /// indefinite / ill-conditioned penalized joint Hessian, mirroring the
5846    /// sibling [`SurvivalMarginalSlopeFamily`]. Interval-censored rows contribute
5847    /// `ℓ = log[S(L) − S(R)]`, the log of a DIFFERENCE of two survival kernels:
5848    /// unlike the log-concave exact-event / right-censored contributions, its
5849    /// per-row Hessian is legitimately INDEFINITE away from the optimum, so the
5850    /// coupled exact-joint penalized Hessian on the constrained (monotone-cone)
5851    /// time block can be full-rank (`nullity == 0`) yet indefinite or severely
5852    /// ill-conditioned at the cold-start seed. The constrained-QP path already
5853    /// REFLECTS negative-curvature modes to `|λ|` (a convex modified-Newton
5854    /// model), but with this gate OFF it adds NO diagonal floor on a full-rank
5855    /// ill-conditioned reflected model, so the trust-region Newton oscillates on
5856    /// the near-singular mode and stalls out the inner budget before any KKT
5857    /// snapshot is taken ("exited the joint Newton path before convergence — no
5858    /// math snapshot"). Arming the gate adds the SAME self-vanishing μ
5859    /// (∝ the projected KKT residual `‖∇ℓ − Sβ + ∇Φ‖` → 0 at the fixed point) the
5860    /// marginal-slope survival inner relies on, so the step is a well-damped
5861    /// modified-Newton descent that converges, while the converged β̂ is the
5862    /// EXACT unconditioned optimum (μ → 0 there) — zero REML/LAML bias, exact
5863    /// gradient unchanged.
5864    fn levenberg_on_ill_conditioning(&self) -> bool {
5865        true
5866    }
5867
5868    fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
5869        // `evaluate_exact_newton_joint_dense` builds a fully dense joint
5870        // Hessian over (Σ p_b)² across time, mean, and optional log-σ blocks
5871        // via per-row pullback of the latent-survival primary kernel.
5872        crate::custom_family::joint_coupled_coefficient_hessian_cost(
5873            self.event_target.len() as u64,
5874            specs,
5875        )
5876    }
5877
5878    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
5879        let (ll, joint_gradient, hess_time, hess_mean, hess_log_sigma) =
5880            self.evaluate_exact_newton_block_diagonals(block_states)?;
5881        let block_ranges = self.joint_block_ranges();
5882        let mut blockworking_sets = vec![
5883            BlockWorkingSet::ExactNewton {
5884                gradient: joint_gradient.slice(s![block_ranges[0].clone()]).to_owned(),
5885                hessian: SymmetricMatrix::Dense(hess_time),
5886            },
5887            BlockWorkingSet::ExactNewton {
5888                gradient: joint_gradient.slice(s![block_ranges[1].clone()]).to_owned(),
5889                hessian: SymmetricMatrix::Dense(hess_mean),
5890            },
5891        ];
5892        if let (Some(range), Some(hessian)) = (block_ranges.get(2).cloned(), hess_log_sigma) {
5893            blockworking_sets.push(BlockWorkingSet::ExactNewton {
5894                gradient: joint_gradient.slice(s![range]).to_owned(),
5895                hessian: SymmetricMatrix::Dense(hessian),
5896            });
5897        }
5898        Ok(FamilyEvaluation {
5899            log_likelihood: ll,
5900            blockworking_sets,
5901        })
5902    }
5903
5904    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
5905        use rayon::iter::{IntoParallelIterator, ParallelIterator};
5906        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-survival")
5907            .map_err(String::from)?;
5908        let (q_entry, q_exit, qdot_exit, mu) = self.split_time_eta(block_states)?;
5909        let q_right = self.time_q_right(block_states)?;
5910        let latent_sd = self.latent_sd(block_states)?;
5911        let n = self.event_target.len();
5912        // Per-row latent-survival jet + log-lik contribution. Independent
5913        // across rows; sum via parallel reduce. `?` propagation happens
5914        // through a Result-collecting fold.
5915        let contributions: Result<Vec<f64>, String> = (0..n)
5916            .into_par_iter()
5917            .map(|i| -> Result<f64, String> {
5918                let wi = weights.at(i);
5919                if wi == 0.0 {
5920                    return Ok(0.0);
5921                }
5922                let row = self.build_row_at(i, q_entry[i], q_exit[i], qdot_exit[i], q_right[i])?;
5923                let jet = LatentSurvivalRowJet::evaluate(&self.quadctx, &row, mu[i], latent_sd)
5924                    .map_err(|e| format!("LatentSurvivalFamily row {i}: {e}"))?;
5925                checked_weighted_row_value(wi, jet.log_lik, i, "log likelihood")
5926            })
5927            .collect();
5928        let mut total = CompensatedRowSum::default();
5929        for contribution in contributions? {
5930            total.add(contribution);
5931        }
5932        require_finite_likelihood_scalar(total.value(), "log likelihood")
5933    }
5934
5935    fn block_linear_constraints(
5936        &self,
5937        _: &[ParameterBlockState],
5938        block_idx: usize,
5939        block_spec: &ParameterBlockSpec,
5940    ) -> Result<Option<ConstraintSet>, String> {
5941        assert!(!block_spec.name.is_empty());
5942        if block_idx == Self::BLOCK_TIME {
5943            Ok(self
5944                .time_linear_constraints
5945                .clone()
5946                .map(ConstraintSet::Dense))
5947        } else {
5948            Ok(None)
5949        }
5950    }
5951
5952    fn exact_newton_joint_hessian(
5953        &self,
5954        block_states: &[ParameterBlockState],
5955    ) -> Result<Option<Array2<f64>>, String> {
5956        self.evaluate_exact_newton_joint_dense(block_states)
5957            .map(|(_, _, hessian)| Some(hessian))
5958    }
5959
5960    fn exact_newton_joint_hessian_workspace(
5961        &self,
5962        block_states: &[ParameterBlockState],
5963        _: &[ParameterBlockSpec],
5964    ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
5965        Ok(Some(Arc::new(LatentSurvivalHessianWorkspace::new(
5966            self.clone(),
5967            block_states.to_vec(),
5968        ))))
5969    }
5970
5971    fn exact_newton_joint_gradient_evaluation(
5972        &self,
5973        block_states: &[ParameterBlockState],
5974        _: &[ParameterBlockSpec],
5975    ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
5976        self.evaluate_exact_newton_joint_gradient_dense(block_states)
5977            .map(|(log_likelihood, gradient)| {
5978                Some(ExactNewtonJointGradientEvaluation {
5979                    log_likelihood,
5980                    gradient,
5981                })
5982            })
5983    }
5984
5985    fn exact_newton_joint_hessian_directional_derivative(
5986        &self,
5987        block_states: &[ParameterBlockState],
5988        d_beta_flat: &Array1<f64>,
5989    ) -> Result<Option<Array2<f64>>, String> {
5990        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
5991            .map(Some)
5992    }
5993
5994    fn exact_newton_joint_hessiansecond_directional_derivative(
5995        &self,
5996        block_states: &[ParameterBlockState],
5997        d_beta_u_flat: &Array1<f64>,
5998        d_beta_v_flat: &Array1<f64>,
5999    ) -> Result<Option<Array2<f64>>, String> {
6000        self.exact_newton_joint_hessian_second_directional_derivative_dense(
6001            block_states,
6002            d_beta_u_flat,
6003            d_beta_v_flat,
6004        )
6005        .map(Some)
6006    }
6007
6008    fn requires_joint_outer_hyper_path(&self) -> bool {
6009        true
6010    }
6011}
6012
6013impl CustomFamily for LatentBinaryFamily {
6014    // Latent binary fits have a separation regime; keep the self-limiting
6015    // Jeffreys/Firth curvature active. The trait default flipped to OFF in
6016    // gam#1395 (flat-prior exact-Newton objective); opt back in here.
6017    fn joint_jeffreys_term_required(&self) -> bool {
6018        true
6019    }
6020
6021    fn exact_newton_joint_hessian_beta_dependent(&self) -> bool {
6022        true
6023    }
6024
6025    fn has_explicit_joint_hessian(&self) -> bool {
6026        true
6027    }
6028
6029    /// Same self-vanishing Levenberg–Marquardt gate as
6030    /// [`LatentSurvivalFamily`]: the latent-binary deployment shares the
6031    /// constrained (monotone-cone) coupled time block, so a full-rank but
6032    /// ill-conditioned penalized joint Hessian at the cold-start seed must get
6033    /// the self-vanishing μ floor rather than oscillating the constrained-QP
6034    /// trust region into a snapshot-less stall. μ → 0 at the fixed point, so the
6035    /// converged β̂ is exact (no REML/LAML bias).
6036    fn levenberg_on_ill_conditioning(&self) -> bool {
6037        true
6038    }
6039
6040    fn coefficient_hessian_cost(&self, specs: &[ParameterBlockSpec]) -> u64 {
6041        crate::custom_family::joint_coupled_coefficient_hessian_cost(
6042            self.event_target.len() as u64,
6043            specs,
6044        )
6045    }
6046
6047    fn evaluate(&self, block_states: &[ParameterBlockState]) -> Result<FamilyEvaluation, String> {
6048        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")
6049            .map_err(String::from)?;
6050        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
6051        let n = self.event_target.len();
6052        let p_time = self.x_time_exit.ncols();
6053        let p_mean = self.x_mean.ncols();
6054
6055        let mut ll = CompensatedRowSum::default();
6056        let mut grad_time = Array1::<f64>::zeros(p_time);
6057        let mut hess_time = Array2::<f64>::zeros((p_time, p_time));
6058        let mut grad_mean = Array1::<f64>::zeros(p_mean);
6059        let mut hess_mean = Array2::<f64>::zeros((p_mean, p_mean));
6060        // Reusable 1-row buffer for x_mean so we avoid allocating a fresh
6061        // Array2<f64> on every iteration via try_row_chunk(i..i+1).
6062        let mut mean_row_buf = Array2::<f64>::zeros((1, p_mean));
6063
6064        for i in 0..n {
6065            let wi = weights.at(i);
6066            if wi == 0.0 {
6067                continue;
6068            }
6069            if !(q_entry[i].is_finite() && q_exit[i].is_finite() && mu[i].is_finite()) {
6070                return Err(format!(
6071                    "latent-binary row {i} contains non-finite predictors: q_entry={}, q_exit={}, mu={}",
6072                    q_entry[i], q_exit[i], mu[i]
6073                ));
6074            }
6075            let row = self.build_right_censored_row_at(i, q_entry[i], q_exit[i])?;
6076            let survival_jet =
6077                LatentSurvivalRowJet::evaluate(&self.quadctx, &row, mu[i], self.latent_sd)
6078                    .map_err(|e| format!("LatentBinaryFamily row {i}: {e}"))?;
6079            let binary = binary_from_log_survival(survival_jet.log_lik, self.event_target[i])?;
6080            ll.add(checked_weighted_row_value(
6081                wi,
6082                binary.log_lik,
6083                i,
6084                "binary log likelihood",
6085            )?);
6086
6087            self.x_mean
6088                .row_chunk_into(i..i + 1, mean_row_buf.view_mut())
6089                .map_err(|e| format!("LatentBinaryFamily row {i} mean row_chunk: {e}"))?;
6090            let mean_vec = mean_row_buf.row(0);
6091            let mean_grad_scale = checked_weighted_row_value(
6092                wi,
6093                binary.grad_scale * survival_jet.score,
6094                i,
6095                "binary mean gradient scale",
6096            )?;
6097            for j in 0..p_mean {
6098                grad_mean[j] += mean_grad_scale * mean_vec[j];
6099            }
6100            let mean_neg_hess = checked_weighted_row_value(
6101                wi,
6102                binary.neg_hess_scale * survival_jet.neg_hessian
6103                    + binary.outer_scale * survival_jet.score * survival_jet.score,
6104                i,
6105                "binary mean Hessian scale",
6106            )?;
6107            dense_outer_accumulate(&mut hess_mean, mean_neg_hess, mean_vec);
6108
6109            let time_jet =
6110                latent_survival_time_jet(&self.quadctx, &row, 0.0, mu[i], self.latent_sd)?;
6111            let t_entry = self.x_time_entry.row(i);
6112            let t_exit = self.x_time_exit.row(i);
6113            let time_gradient_scale =
6114                checked_weighted_row_value(wi, binary.grad_scale, i, "binary time gradient scale")?;
6115            for j in 0..p_time {
6116                grad_time[j] += time_gradient_scale
6117                    * (time_jet.grad_entry * t_entry[j] + time_jet.grad_exit * t_exit[j]);
6118            }
6119            let entry_hessian_scale = checked_weighted_row_value(
6120                wi,
6121                binary.neg_hess_scale * time_jet.neg_hess_entry,
6122                i,
6123                "binary entry Hessian scale",
6124            )?;
6125            dense_outer_accumulate(&mut hess_time, entry_hessian_scale, t_entry);
6126            let exit_hessian_scale = checked_weighted_row_value(
6127                wi,
6128                binary.neg_hess_scale * time_jet.neg_hess_exit,
6129                i,
6130                "binary exit Hessian scale",
6131            )?;
6132            dense_outer_accumulate(&mut hess_time, exit_hessian_scale, t_exit);
6133            if binary.outer_scale != 0.0 {
6134                let entry_outer_scale = checked_weighted_row_value(
6135                    wi,
6136                    binary.outer_scale * time_jet.grad_entry * time_jet.grad_entry,
6137                    i,
6138                    "binary entry outer Hessian scale",
6139                )?;
6140                dense_outer_accumulate(&mut hess_time, entry_outer_scale, t_entry);
6141                let exit_outer_scale = checked_weighted_row_value(
6142                    wi,
6143                    binary.outer_scale * time_jet.grad_exit * time_jet.grad_exit,
6144                    i,
6145                    "binary exit outer Hessian scale",
6146                )?;
6147                dense_outer_accumulate(&mut hess_time, exit_outer_scale, t_exit);
6148                let cross_outer_scale = checked_weighted_row_value(
6149                    wi,
6150                    binary.outer_scale * time_jet.grad_entry * time_jet.grad_exit,
6151                    i,
6152                    "binary cross outer Hessian scale",
6153                )?;
6154                dense_symmetric_cross_accumulate(
6155                    &mut hess_time,
6156                    cross_outer_scale,
6157                    t_entry,
6158                    t_exit,
6159                );
6160            }
6161        }
6162
6163        let ll = require_finite_likelihood_scalar(ll.value(), "binary log likelihood")?;
6164        require_finite_likelihood_vector(&grad_time, "binary time gradient")?;
6165        require_finite_likelihood_vector(&grad_mean, "binary mean gradient")?;
6166        require_finite_likelihood_matrix(&hess_time, "binary time Hessian")?;
6167        require_finite_likelihood_matrix(&hess_mean, "binary mean Hessian")?;
6168        Ok(FamilyEvaluation {
6169            log_likelihood: ll,
6170            blockworking_sets: vec![
6171                BlockWorkingSet::ExactNewton {
6172                    gradient: grad_time,
6173                    hessian: SymmetricMatrix::Dense(hess_time),
6174                },
6175                BlockWorkingSet::ExactNewton {
6176                    gradient: grad_mean,
6177                    hessian: SymmetricMatrix::Dense(hess_mean),
6178                },
6179            ],
6180        })
6181    }
6182
6183    fn log_likelihood_only(&self, block_states: &[ParameterBlockState]) -> Result<f64, String> {
6184        let weights = ValidatedLikelihoodWeights::new(&self.weights, "latent-binary")
6185            .map_err(String::from)?;
6186        let (q_entry, q_exit, mu) = self.split_time_eta(block_states)?;
6187        let mut ll = CompensatedRowSum::default();
6188        for i in 0..self.event_target.len() {
6189            let wi = weights.at(i);
6190            if wi == 0.0 {
6191                continue;
6192            }
6193            let row = self.build_right_censored_row_at(i, q_entry[i], q_exit[i])?;
6194            let survival_jet =
6195                LatentSurvivalRowJet::evaluate(&self.quadctx, &row, mu[i], self.latent_sd)
6196                    .map_err(|e| format!("LatentBinaryFamily row {i}: {e}"))?;
6197            let binary_log_lik = binary_log_likelihood_from_log_survival(
6198                survival_jet.log_lik,
6199                self.event_target[i],
6200            )?;
6201            ll.add(checked_weighted_row_value(
6202                wi,
6203                binary_log_lik,
6204                i,
6205                "binary log likelihood",
6206            )?);
6207        }
6208        require_finite_likelihood_scalar(ll.value(), "binary log likelihood")
6209    }
6210
6211    fn block_linear_constraints(
6212        &self,
6213        _: &[ParameterBlockState],
6214        block_idx: usize,
6215        block_spec: &ParameterBlockSpec,
6216    ) -> Result<Option<ConstraintSet>, String> {
6217        assert!(!block_spec.name.is_empty());
6218        if block_idx == Self::BLOCK_TIME {
6219            Ok(self
6220                .time_linear_constraints
6221                .clone()
6222                .map(ConstraintSet::Dense))
6223        } else {
6224            Ok(None)
6225        }
6226    }
6227
6228    fn exact_newton_joint_hessian(
6229        &self,
6230        block_states: &[ParameterBlockState],
6231    ) -> Result<Option<Array2<f64>>, String> {
6232        self.evaluate_exact_newton_joint_dense(block_states)
6233            .map(|(_, _, hessian)| Some(hessian))
6234    }
6235
6236    fn exact_newton_joint_hessian_workspace(
6237        &self,
6238        block_states: &[ParameterBlockState],
6239        _: &[ParameterBlockSpec],
6240    ) -> Result<Option<Arc<dyn ExactNewtonJointHessianWorkspace>>, String> {
6241        Ok(Some(Arc::new(LatentBinaryHessianWorkspace::new(
6242            self.clone(),
6243            block_states.to_vec(),
6244        ))))
6245    }
6246
6247    fn exact_newton_joint_gradient_evaluation(
6248        &self,
6249        block_states: &[ParameterBlockState],
6250        _: &[ParameterBlockSpec],
6251    ) -> Result<Option<ExactNewtonJointGradientEvaluation>, String> {
6252        self.evaluate_exact_newton_joint_dense(block_states)
6253            .map(|(log_likelihood, gradient, _)| {
6254                Some(ExactNewtonJointGradientEvaluation {
6255                    log_likelihood,
6256                    gradient,
6257                })
6258            })
6259    }
6260
6261    fn exact_newton_joint_hessian_directional_derivative(
6262        &self,
6263        block_states: &[ParameterBlockState],
6264        d_beta_flat: &Array1<f64>,
6265    ) -> Result<Option<Array2<f64>>, String> {
6266        self.exact_newton_joint_hessian_directional_derivative_dense(block_states, d_beta_flat)
6267            .map(Some)
6268    }
6269
6270    fn exact_newton_joint_hessiansecond_directional_derivative(
6271        &self,
6272        block_states: &[ParameterBlockState],
6273        d_beta_u_flat: &Array1<f64>,
6274        d_beta_v_flat: &Array1<f64>,
6275    ) -> Result<Option<Array2<f64>>, String> {
6276        self.exact_newton_joint_hessian_second_directional_derivative_dense(
6277            block_states,
6278            d_beta_u_flat,
6279            d_beta_v_flat,
6280        )
6281        .map(Some)
6282    }
6283
6284    fn requires_joint_outer_hyper_path(&self) -> bool {
6285        true
6286    }
6287}
6288
6289#[cfg(test)]
6290mod tests {
6291    use super::tests_multidir_channels::{
6292        latent_survival_row_primary_fourth_contracted_multidir_reference,
6293        latent_survival_row_primary_gradient_hessian_multidir_reference,
6294        latent_survival_row_primary_third_contracted_multidir_reference,
6295    };
6296    use super::*;
6297    use crate::custom_family::BlockWorkingSet;
6298    use gam_linalg::matrix::DenseDesignMatrix;
6299    use ndarray::array;
6300
6301    fn learnable_sigma_test_family() -> LatentSurvivalFamily {
6302        LatentSurvivalFamily {
6303            event_target: array![1u8, 0u8],
6304            weights: array![1.0, 0.7],
6305            latent_sd_fixed: None,
6306            hazard_loading: HazardLoading::LoadedVsUnloaded,
6307            unloaded_mass_entry: array![0.02, 0.03],
6308            unloaded_mass_exit: array![0.05, 0.08],
6309            unloaded_hazard_exit: array![0.04, 0.0],
6310            x_time_entry: array![[1.0, -0.2], [0.4, 0.7]],
6311            x_time_exit: array![[1.3, 0.1], [0.9, 1.0]],
6312            x_time_derivative_exit: array![[0.8, 0.4], [0.6, 0.5]],
6313            x_time_right: array![[1.3, 0.1], [0.9, 1.0]],
6314            time_offset_right: Array1::zeros(2),
6315            unloaded_mass_right: Array1::zeros(2),
6316            x_mean: DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0, -0.3], [0.2, 0.9]])),
6317            time_linear_constraints: None,
6318            quadctx: Arc::new(QuadratureContext::new()),
6319        }
6320    }
6321
6322    fn learnable_sigma_test_joint_beta() -> Array1<f64> {
6323        array![0.15, 0.25, 0.1, -0.15, 0.35_f64.ln()]
6324    }
6325
6326    #[test]
6327    fn saved_latent_survival_alo_matches_deterministic_frailty_oracle() {
6328        const K: usize = 5;
6329        let weight = 1.7;
6330        let quadrature = QuadratureContext::new();
6331        let geometry = latent_survival_alo_row_geometry(LatentSurvivalAloRowInput {
6332            quadrature: &quadrature,
6333            hazard_loading: HazardLoading::Full,
6334            event_code: 0,
6335            prior_weight: weight,
6336            q_entry: -0.4,
6337            q_exit: 0.2,
6338            qdot_exit: 1.1,
6339            q_right: 0.2,
6340            mu: -0.1,
6341            sigma: LatentSurvivalAloSigma::Fixed(0.0),
6342            unloaded_mass_entry: 0.0,
6343            unloaded_mass_exit: 0.0,
6344            unloaded_mass_right: 0.0,
6345            unloaded_hazard_exit: 0.0,
6346        })
6347        .expect("exact saved latent-survival row");
6348
6349        let values: [f64; K] = geometry
6350            .coordinate_values
6351            .as_slice()
6352            .expect("owned coordinates are contiguous")
6353            .try_into()
6354            .expect("fixed-scale latent survival has five primaries");
6355        let variables: [Order2<K>; K] =
6356            std::array::from_fn(|axis| Order2::variable(values[axis], axis));
6357        // At sigma=0 and full loading, right-censoring has
6358        // NLL = w[exp(q_exit + mu) - exp(q_entry + mu)].
6359        let oracle = variables[1]
6360            .add(&variables[4])
6361            .exp()
6362            .sub(&variables[0].add(&variables[4]).exp())
6363            .scale(weight);
6364        for left in 0..K {
6365            assert!((geometry.nll_score[left] - oracle.g()[left]).abs() <= 3e-12);
6366            for right in 0..K {
6367                assert!(
6368                    (geometry.observed_hessian[[left, right]] - oracle.h()[left][right]).abs()
6369                        <= 4e-12
6370                );
6371            }
6372        }
6373    }
6374
6375    #[test]
6376    fn saved_latent_binary_alo_matches_deterministic_frailty_oracle() {
6377        const K: usize = 3;
6378        let weight = 0.8;
6379        let quadrature = QuadratureContext::new();
6380        let geometry = latent_binary_alo_row_geometry(LatentBinaryAloRowInput {
6381            quadrature: &quadrature,
6382            hazard_loading: HazardLoading::Full,
6383            event: 1,
6384            prior_weight: weight,
6385            q_entry: -0.6,
6386            q_exit: 0.3,
6387            mu: -0.2,
6388            sigma: 0.0,
6389            unloaded_mass_entry: 0.0,
6390            unloaded_mass_exit: 0.0,
6391        })
6392        .expect("exact saved latent-binary row");
6393
6394        let values: [f64; K] = geometry
6395            .coordinate_values
6396            .as_slice()
6397            .expect("owned coordinates are contiguous")
6398            .try_into()
6399            .expect("latent binary has three live primaries");
6400        let variables: [Order2<K>; K] =
6401            std::array::from_fn(|axis| Order2::variable(values[axis], axis));
6402        let log_survival = variables[0]
6403            .add(&variables[2])
6404            .exp()
6405            .sub(&variables[1].add(&variables[2]).exp());
6406        let one = variables[0].compose_unary([1.0, 0.0, 0.0, 0.0, 0.0]);
6407        // Event NLL = -w log(1 - exp(log_survival)).
6408        let oracle = one.sub(&log_survival.exp()).ln().neg().scale(weight);
6409        for left in 0..K {
6410            assert!((geometry.nll_score[left] - oracle.g()[left]).abs() <= 4e-12);
6411            for right in 0..K {
6412                assert!(
6413                    (geometry.observed_hessian[[left, right]] - oracle.h()[left][right]).abs()
6414                        <= 6e-12
6415                );
6416            }
6417        }
6418    }
6419
6420    /// Regression (frailty scale block deletion): a learnable-σ latent-survival
6421    /// fit routes the pre-fit identifiability audit CHANNEL-AWARE, so the
6422    /// `log_sigma` scale block — realised as a single constant-of-ones column —
6423    /// is never aliased against the `mean` intercept and dropped. Before the
6424    /// `output_channel_assignment` override the family used the trait default
6425    /// (every block → channel 0); the flat single-channel RRQR then saw the two
6426    /// constant columns (mean intercept, log-σ constant) as a cross-block alias,
6427    /// attributed the drop to the lowest-priority `log_sigma` block, and deleted
6428    /// the ONLY handle on the frailty scale — which both froze σ and left the
6429    /// reduced spec width one short of the family's raw joint Hessian, aborting
6430    /// every outer LAML eval with `joint exact-newton Hessian validation … got
6431    /// 15x15, expected 14x14`. The contract this guards: the scale channel must
6432    /// be distinct from the location (mean) channel.
6433    #[test]
6434    fn latent_survival_learnable_sigma_block_lives_on_a_distinct_output_channel() {
6435        let family = learnable_sigma_test_family();
6436        assert!(
6437            family.latent_sd_fixed.is_none(),
6438            "test fixture must be the learnable-σ family"
6439        );
6440
6441        // The three blocks the learnable-σ builder emits, in order. Only the
6442        // block NAME drives `output_channel_assignment`, so borrow the real
6443        // `build_log_sigma_blockspec` shape and relabel for time/mean.
6444        let mut time_spec = build_log_sigma_blockspec(0.5, family.event_target.len());
6445        time_spec.name = "time_transform".to_string();
6446        let mut mean_spec = build_log_sigma_blockspec(0.5, family.event_target.len());
6447        mean_spec.name = "mean".to_string();
6448        let log_sigma_spec = build_log_sigma_blockspec(0.5, family.event_target.len());
6449        assert_eq!(log_sigma_spec.name, "log_sigma");
6450        let specs = vec![time_spec, mean_spec, log_sigma_spec];
6451
6452        let channels = family
6453            .output_channel_assignment(&specs)
6454            .expect("latent survival must declare an explicit channel assignment");
6455        assert_eq!(channels.len(), specs.len());
6456
6457        let (time_ch, mean_ch, log_sigma_ch) = (channels[0], channels[1], channels[2]);
6458        // The load-bearing assertion: the scale channel is NOT the location
6459        // channel, so the channel-aware audit never aliases σ's constant column
6460        // against the mean intercept.
6461        assert_ne!(
6462            log_sigma_ch, mean_ch,
6463            "log_sigma (frailty scale) must not share the mean's output channel, or the \
6464             identifiability audit will alias its constant column against the mean intercept \
6465             and delete the scale parameter"
6466        );
6467        // Each latent-survival block drives a structurally distinct output, so
6468        // all three channels are distinct (block-diagonal true Jacobian).
6469        assert_ne!(time_ch, mean_ch);
6470        assert_ne!(time_ch, log_sigma_ch);
6471        let n_outputs = channels.iter().copied().max().unwrap() + 1;
6472        assert!(
6473            n_outputs >= 3,
6474            "learnable-σ latent survival must expose ≥3 output channels (time, mean, scale), \
6475             got {n_outputs}"
6476        );
6477    }
6478
6479    fn survival_stress_test_family(n: usize) -> LatentSurvivalFamily {
6480        LatentSurvivalFamily {
6481            event_target: Array1::from_iter((0..n).map(|i| if i % 3 == 0 { 1u8 } else { 0u8 })),
6482            weights: Array1::from_iter((0..n).map(|i| 0.55 + 0.03 * ((i % 7) as f64))),
6483            latent_sd_fixed: None,
6484            hazard_loading: HazardLoading::LoadedVsUnloaded,
6485            unloaded_mass_entry: Array1::from_iter(
6486                (0..n).map(|i| 0.015 + 0.0015 * ((i % 11) as f64)),
6487            ),
6488            unloaded_mass_exit: Array1::from_iter((0..n).map(|i| 0.06 + 0.002 * ((i % 13) as f64))),
6489            unloaded_hazard_exit: Array1::from_iter((0..n).map(|i| {
6490                if i % 4 == 0 {
6491                    0.018 + 0.001 * ((i % 5) as f64)
6492                } else {
6493                    0.0
6494                }
6495            })),
6496            x_time_entry: Array2::from_shape_fn((n, 4), |(i, j)| {
6497                0.2 + 0.03 * ((i + 2 * j) % 9) as f64 - if j == 1 { 0.12 } else { 0.0 }
6498            }),
6499            x_time_exit: Array2::from_shape_fn((n, 4), |(i, j)| {
6500                0.35 + 0.025 * ((2 * i + j) % 10) as f64 - if j == 2 { 0.08 } else { 0.0 }
6501            }),
6502            x_time_derivative_exit: Array2::from_shape_fn((n, 4), |(i, j)| {
6503                0.45 + 0.015 * ((i + 3 * j) % 8) as f64
6504            }),
6505            x_time_right: Array2::from_shape_fn((n, 4), |(i, j)| {
6506                0.35 + 0.025 * ((2 * i + j) % 10) as f64 - if j == 2 { 0.08 } else { 0.0 }
6507            }),
6508            time_offset_right: Array1::zeros(n),
6509            unloaded_mass_right: Array1::zeros(n),
6510            x_mean: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::from_shape_fn(
6511                (n, 3),
6512                |(i, j)| 0.1 + 0.04 * ((3 * i + j) % 7) as f64 - if j == 0 { 0.18 } else { 0.0 },
6513            ))),
6514            time_linear_constraints: None,
6515            quadctx: Arc::new(QuadratureContext::new()),
6516        }
6517    }
6518
6519    fn survival_stress_test_joint_beta() -> Array1<f64> {
6520        array![0.18, 0.11, 0.07, 0.13, -0.09, 0.05, 0.12, 0.42_f64.ln()]
6521    }
6522
6523    fn latent_survival_states_from_joint_beta(
6524        family: &LatentSurvivalFamily,
6525        joint_beta: &Array1<f64>,
6526    ) -> Vec<ParameterBlockState> {
6527        let slices = family.joint_slices();
6528        let n = family.event_target.len();
6529        let beta_time = joint_beta.slice(s![slices.time.clone()]).to_owned();
6530        let beta_mean = joint_beta.slice(s![slices.mean.clone()]).to_owned();
6531
6532        let mut eta_time = Array1::<f64>::zeros(3 * n);
6533        eta_time
6534            .slice_mut(s![0..n])
6535            .assign(&gam_linalg::faer_ndarray::fast_av(
6536                &family.x_time_entry,
6537                &beta_time,
6538            ));
6539        eta_time
6540            .slice_mut(s![n..2 * n])
6541            .assign(&gam_linalg::faer_ndarray::fast_av(
6542                &family.x_time_exit,
6543                &beta_time,
6544            ));
6545        eta_time
6546            .slice_mut(s![2 * n..3 * n])
6547            .assign(&gam_linalg::faer_ndarray::fast_av(
6548                &family.x_time_derivative_exit,
6549                &beta_time,
6550            ));
6551
6552        let mut states = vec![
6553            ParameterBlockState {
6554                beta: beta_time,
6555                eta: eta_time,
6556            },
6557            ParameterBlockState {
6558                beta: beta_mean.clone(),
6559                eta: family.x_mean.dot(&beta_mean),
6560            },
6561        ];
6562        if let Some(log_sigma) = slices.log_sigma {
6563            let beta_log_sigma = array![joint_beta[log_sigma.start]];
6564            states.push(ParameterBlockState {
6565                beta: beta_log_sigma.clone(),
6566                eta: beta_log_sigma,
6567            });
6568        }
6569        states
6570    }
6571
6572    fn max_relative_array1(left: &Array1<f64>, right: &Array1<f64>) -> f64 {
6573        left.iter()
6574            .zip(right.iter())
6575            .map(|(l, r)| (l - r).abs() / l.abs().max(r.abs()).max(1e-12))
6576            .fold(0.0_f64, f64::max)
6577    }
6578
6579    fn max_relative_array2(left: &Array2<f64>, right: &Array2<f64>) -> f64 {
6580        left.iter()
6581            .zip(right.iter())
6582            .map(|(l, r)| (l - r).abs() / l.abs().max(r.abs()).max(1e-12))
6583            .fold(0.0_f64, f64::max)
6584    }
6585
6586    fn frobenius_relative_array2(left: &Array2<f64>, right: &Array2<f64>) -> f64 {
6587        let mut diff2 = 0.0_f64;
6588        let mut scale2 = 0.0_f64;
6589        for (l, r) in left.iter().zip(right.iter()) {
6590            let d = l - r;
6591            diff2 += d * d;
6592            scale2 += l * l + r * r;
6593        }
6594        diff2.sqrt() / scale2.sqrt().max(1e-12)
6595    }
6596
6597    fn latent_survival_row_loglik_from_primary(
6598        quadctx: &QuadratureContext,
6599        row: &LatentSurvivalRow,
6600        primary: &Array1<f64>,
6601    ) -> f64 {
6602        let q_entry = primary[LATENT_SURVIVAL_PRIMARY_Q_ENTRY];
6603        let q_exit = primary[LATENT_SURVIVAL_PRIMARY_Q_EXIT];
6604        let qdot_exit = primary[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT];
6605        let q_right = primary[LATENT_SURVIVAL_PRIMARY_Q_RIGHT];
6606        let mu = primary[LATENT_SURVIVAL_PRIMARY_MU];
6607        let sigma = primary[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA].exp();
6608        latent_survival_row_primary_gradient_hessian(
6609            quadctx,
6610            row,
6611            LatentSurvivalPrimaryPoint {
6612                q_entry,
6613                q_exit,
6614                qdot_exit,
6615                q_right,
6616                mu,
6617                sigma,
6618            },
6619            true,
6620        )
6621        .expect("row primary evaluation")
6622        .0
6623    }
6624
6625    fn latent_test_specs(n: usize, block_dims: &[(&str, usize)]) -> Vec<ParameterBlockSpec> {
6626        block_dims
6627            .iter()
6628            .map(|(name, p)| ParameterBlockSpec {
6629                name: (*name).to_string(),
6630                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((n, *p)))),
6631                offset: Array1::zeros(n),
6632                penalties: Vec::new(),
6633                nullspace_dims: Vec::new(),
6634                initial_log_lambdas: Array1::zeros(0),
6635                initial_beta: None,
6636                gauge_priority: 100,
6637                jacobian_callback: None,
6638                stacked_design: None,
6639                stacked_offset: None,
6640            })
6641            .collect()
6642    }
6643
6644    fn fixed_sigma_binary_test_family() -> LatentBinaryFamily {
6645        LatentBinaryFamily {
6646            event_target: array![1u8, 0u8],
6647            weights: array![1.0, 0.7],
6648            latent_sd: 0.35,
6649            hazard_loading: HazardLoading::LoadedVsUnloaded,
6650            unloaded_mass_entry: array![0.02, 0.03],
6651            unloaded_mass_exit: array![0.05, 0.08],
6652            x_time_entry: array![[1.0, -0.2], [0.4, 0.7]],
6653            x_time_exit: array![[1.3, 0.1], [0.9, 1.0]],
6654            x_mean: DesignMatrix::Dense(DenseDesignMatrix::from(array![[1.0, -0.3], [0.2, 0.9]])),
6655            time_linear_constraints: None,
6656            quadctx: Arc::new(QuadratureContext::new()),
6657        }
6658    }
6659
6660    #[test]
6661    fn latent_survival_offset_residuals_reject_missing_block_state() {
6662        let error = learnable_sigma_test_family()
6663            .offset_channel_residuals(&[])
6664            .expect_err("missing fitted blocks must not become zero residuals");
6665        match error {
6666            LatentSurvivalError::BlockMismatch { reason } => {
6667                assert!(reason.contains("got 0"), "unexpected mismatch: {reason}");
6668            }
6669            other => panic!("missing fitted blocks must be a block mismatch, got {other:?}"),
6670        }
6671    }
6672
6673    #[test]
6674    fn latent_binary_offset_residuals_reject_missing_block_state() {
6675        let error = fixed_sigma_binary_test_family()
6676            .offset_channel_residuals(&[])
6677            .expect_err("missing fitted blocks must not become zero residuals");
6678        match error {
6679            LatentSurvivalError::BlockMismatch { reason } => {
6680                assert!(reason.contains("got 0"), "unexpected mismatch: {reason}");
6681            }
6682            other => panic!("missing fitted blocks must be a block mismatch, got {other:?}"),
6683        }
6684    }
6685
6686    fn latent_binary_states_from_joint_beta(
6687        family: &LatentBinaryFamily,
6688        joint_beta: &Array1<f64>,
6689    ) -> Vec<ParameterBlockState> {
6690        let slices = family.joint_slices();
6691        let n = family.event_target.len();
6692        let beta_time = joint_beta.slice(s![slices.time.clone()]).to_owned();
6693        let beta_mean = joint_beta.slice(s![slices.mean.clone()]).to_owned();
6694
6695        let mut eta_time = Array1::<f64>::zeros(3 * n);
6696        eta_time
6697            .slice_mut(s![0..n])
6698            .assign(&gam_linalg::faer_ndarray::fast_av(
6699                &family.x_time_entry,
6700                &beta_time,
6701            ));
6702        eta_time
6703            .slice_mut(s![n..2 * n])
6704            .assign(&gam_linalg::faer_ndarray::fast_av(
6705                &family.x_time_exit,
6706                &beta_time,
6707            ));
6708
6709        vec![
6710            ParameterBlockState {
6711                beta: beta_time,
6712                eta: eta_time,
6713            },
6714            ParameterBlockState {
6715                beta: beta_mean.clone(),
6716                eta: family.x_mean.dot(&beta_mean),
6717            },
6718        ]
6719    }
6720
6721    fn assert_scalar_is_scaled(got: f64, unweighted: f64, scale: f64, quantity: &str) {
6722        let expected = unweighted * scale;
6723        let tolerance = 128.0 * f64::EPSILON * expected.abs().max(f64::MIN_POSITIVE);
6724        assert!(
6725            (got - expected).abs() <= tolerance,
6726            "{quantity} did not scale with its positive row weight: got={got:?}, expected={expected:?}, unweighted={unweighted:?}, scale={scale:?}, tolerance={tolerance:?}"
6727        );
6728    }
6729
6730    fn assert_vector_is_scaled(
6731        got: &Array1<f64>,
6732        unweighted: &Array1<f64>,
6733        scale: f64,
6734        quantity: &str,
6735    ) {
6736        assert_eq!(got.len(), unweighted.len());
6737        for (index, (&actual, &base)) in got.iter().zip(unweighted.iter()).enumerate() {
6738            assert_scalar_is_scaled(actual, base, scale, &format!("{quantity}[{index}]"));
6739        }
6740    }
6741
6742    fn assert_matrix_is_scaled(
6743        got: &Array2<f64>,
6744        unweighted: &Array2<f64>,
6745        scale: f64,
6746        quantity: &str,
6747    ) {
6748        assert_eq!(got.dim(), unweighted.dim());
6749        for ((row, col), &actual) in got.indexed_iter() {
6750            assert_scalar_is_scaled(
6751                actual,
6752                unweighted[[row, col]],
6753                scale,
6754                &format!("{quantity}[{row},{col}]"),
6755            );
6756        }
6757    }
6758
6759    #[test]
6760    fn binary_log_survival_math_is_cancellation_free_and_derivative_order_aware() {
6761        let near_one_survival: f64 = -1.0e-16;
6762        let expected = (-near_one_survival.exp_m1()).ln();
6763        let got = binary_log_likelihood_from_log_survival(near_one_survival, 1)
6764            .expect("near-boundary binary event likelihood");
6765        assert_eq!(got.to_bits(), expected.to_bits());
6766
6767        // At s=-1e-100 the value, score, Hessian, and third derivative are
6768        // representable, while the fourth derivative is mathematically beyond
6769        // f64 range. Value/order-2 callers must succeed; only the order-4 path
6770        // is allowed to refuse the unrepresentable result.
6771        let extreme = -1.0e-100;
6772        assert!(
6773            binary_log_likelihood_from_log_survival(extreme, 1)
6774                .expect("value remains representable")
6775                .is_finite()
6776        );
6777        let (_, first) = binary_from_log_survival_through_first(extreme, 1)
6778            .expect("first derivative remains representable");
6779        assert!(first.is_finite());
6780        let second =
6781            binary_from_log_survival(extreme, 1).expect("second derivative remains representable");
6782        assert!(second.outer_scale.is_finite());
6783        let (_, third) = binary_from_log_survival_through_third(extreme, 1)
6784            .expect("third derivative remains representable");
6785        assert!(third.is_finite());
6786        let fourth = binary_from_log_survival_through_fourth(extreme, 1)
6787            .expect_err("unrepresentable fourth derivative must be explicit");
6788        assert!(
6789            fourth.to_string().contains("derivative order 4"),
6790            "unexpected fourth-derivative error: {fourth}"
6791        );
6792    }
6793
6794    #[test]
6795    fn latent_likelihood_preserves_every_positive_weight_and_scales_all_derivatives() {
6796        // This power-of-two scale is far below the deleted 1e-12 omission
6797        // threshold. Multiplication by it changes only the exponent, making
6798        // value/gradient/Hessian scaling a sharp regression oracle rather than
6799        // a loose floating-point approximation.
6800        let tiny_normal = 2.0_f64.powi(-48);
6801
6802        let mut survival_unit = learnable_sigma_test_family();
6803        survival_unit.weights = array![1.0, 0.0];
6804        let survival_states = latent_survival_states_from_joint_beta(
6805            &survival_unit,
6806            &learnable_sigma_test_joint_beta(),
6807        );
6808        let (survival_ll, survival_gradient, survival_hessian) = survival_unit
6809            .evaluate_exact_newton_joint_dense(&survival_states)
6810            .expect("unit-weight latent-survival evaluation");
6811        let mut survival_tiny = survival_unit.clone();
6812        survival_tiny.weights[0] = tiny_normal;
6813        let (tiny_survival_ll, tiny_survival_gradient, tiny_survival_hessian) = survival_tiny
6814            .evaluate_exact_newton_joint_dense(&survival_states)
6815            .expect("tiny-positive latent-survival evaluation");
6816        assert_scalar_is_scaled(
6817            tiny_survival_ll,
6818            survival_ll,
6819            tiny_normal,
6820            "survival log likelihood",
6821        );
6822        assert_vector_is_scaled(
6823            &tiny_survival_gradient,
6824            &survival_gradient,
6825            tiny_normal,
6826            "survival gradient",
6827        );
6828        assert_matrix_is_scaled(
6829            &tiny_survival_hessian,
6830            &survival_hessian,
6831            tiny_normal,
6832            "survival Hessian",
6833        );
6834
6835        let mut binary_unit = fixed_sigma_binary_test_family();
6836        binary_unit.weights = array![1.0, 0.0];
6837        let binary_beta = array![0.15, 0.25, 0.1, -0.15];
6838        let binary_states = latent_binary_states_from_joint_beta(&binary_unit, &binary_beta);
6839        let (binary_ll, binary_gradient, binary_hessian) = binary_unit
6840            .evaluate_exact_newton_joint_dense(&binary_states)
6841            .expect("unit-weight latent-binary evaluation");
6842        let mut binary_tiny = binary_unit.clone();
6843        binary_tiny.weights[0] = tiny_normal;
6844        let (tiny_binary_ll, tiny_binary_gradient, tiny_binary_hessian) = binary_tiny
6845            .evaluate_exact_newton_joint_dense(&binary_states)
6846            .expect("tiny-positive latent-binary evaluation");
6847        assert_scalar_is_scaled(
6848            tiny_binary_ll,
6849            binary_ll,
6850            tiny_normal,
6851            "binary log likelihood",
6852        );
6853        assert_vector_is_scaled(
6854            &tiny_binary_gradient,
6855            &binary_gradient,
6856            tiny_normal,
6857            "binary gradient",
6858        );
6859        assert_matrix_is_scaled(
6860            &tiny_binary_hessian,
6861            &binary_hessian,
6862            tiny_normal,
6863            "binary Hessian",
6864        );
6865
6866        // A largest-subnormal sample weight remains a real likelihood row.
6867        // Its single-row log-likelihood product is representable and must not
6868        // collapse to the all-zero dormant-row result.
6869        let largest_subnormal = f64::from_bits((1_u64 << 52) - 1);
6870        survival_tiny.weights[0] = largest_subnormal;
6871        let subnormal_survival_ll = survival_tiny
6872            .log_likelihood_only(&survival_states)
6873            .expect("subnormal latent-survival likelihood");
6874        assert_ne!(subnormal_survival_ll, 0.0);
6875        assert_scalar_is_scaled(
6876            subnormal_survival_ll,
6877            survival_ll,
6878            largest_subnormal,
6879            "subnormal survival log likelihood",
6880        );
6881        binary_tiny.weights[0] = largest_subnormal;
6882        let subnormal_binary_ll = binary_tiny
6883            .log_likelihood_only(&binary_states)
6884            .expect("subnormal latent-binary likelihood");
6885        assert_ne!(subnormal_binary_ll, 0.0);
6886        assert_scalar_is_scaled(
6887            subnormal_binary_ll,
6888            binary_ll,
6889            largest_subnormal,
6890            "subnormal binary log likelihood",
6891        );
6892
6893        // Even the smallest positive subnormal survives when the mathematical
6894        // product is representable. If it is not representable, the checked
6895        // scaling primitive refuses it explicitly instead of converting the
6896        // row into a semantic zero.
6897        let smallest_subnormal = f64::from_bits(1);
6898        assert_eq!(
6899            checked_weighted_row_value(smallest_subnormal, 1.0, 0, "test")
6900                .expect("representable smallest-subnormal product")
6901                .to_bits(),
6902            smallest_subnormal.to_bits()
6903        );
6904        let underflow = checked_weighted_row_value(smallest_subnormal, 0.25, 0, "test")
6905            .expect_err("non-zero underflow must be explicit");
6906        assert!(
6907            underflow.contains("underflowed"),
6908            "unexpected error: {underflow}"
6909        );
6910    }
6911
6912    #[test]
6913    fn zero_weight_likelihood_rows_are_dormant_before_response_and_predictor_access() {
6914        let mut survival = learnable_sigma_test_family();
6915        survival.weights = array![1.0, 0.0];
6916        let beta = learnable_sigma_test_joint_beta();
6917        let states = latent_survival_states_from_joint_beta(&survival, &beta);
6918        let expected = survival
6919            .evaluate_exact_newton_joint_dense(&states)
6920            .expect("clean zero-weight survival reference");
6921
6922        survival.event_target[1] = 17;
6923        survival.unloaded_mass_entry[1] = f64::NAN;
6924        survival.unloaded_mass_exit[1] = f64::NEG_INFINITY;
6925        survival.unloaded_mass_right[1] = -1.0;
6926        survival.unloaded_hazard_exit[1] = f64::NAN;
6927        survival.time_offset_right[1] = f64::NAN;
6928        survival.x_time_right.row_mut(1).fill(f64::NAN);
6929        let mut dormant_states = states.clone();
6930        let n = survival.event_target.len();
6931        dormant_states[LatentSurvivalFamily::BLOCK_TIME].eta[1] = f64::NAN;
6932        dormant_states[LatentSurvivalFamily::BLOCK_TIME].eta[n + 1] = f64::INFINITY;
6933        dormant_states[LatentSurvivalFamily::BLOCK_TIME].eta[2 * n + 1] = f64::NEG_INFINITY;
6934        dormant_states[LatentSurvivalFamily::BLOCK_MEAN].eta[1] = f64::NAN;
6935        let got = survival
6936            .evaluate_exact_newton_joint_dense(&dormant_states)
6937            .expect("zero-weight survival row must not inspect dormant response/predictors");
6938        assert_eq!(got, expected);
6939
6940        let mut binary = fixed_sigma_binary_test_family();
6941        binary.weights = array![1.0, 0.0];
6942        let binary_beta = array![0.15, 0.25, 0.1, -0.15];
6943        let binary_states = latent_binary_states_from_joint_beta(&binary, &binary_beta);
6944        let expected = binary
6945            .evaluate_exact_newton_joint_dense(&binary_states)
6946            .expect("clean zero-weight binary reference");
6947        binary.event_target[1] = 17;
6948        binary.unloaded_mass_entry[1] = f64::NAN;
6949        binary.unloaded_mass_exit[1] = f64::NEG_INFINITY;
6950        let mut dormant_binary_states = binary_states.clone();
6951        let n = binary.event_target.len();
6952        dormant_binary_states[LatentBinaryFamily::BLOCK_TIME].eta[1] = f64::NAN;
6953        dormant_binary_states[LatentBinaryFamily::BLOCK_TIME].eta[n + 1] = f64::INFINITY;
6954        dormant_binary_states[LatentBinaryFamily::BLOCK_MEAN].eta[1] = f64::NAN;
6955        let got = binary
6956            .evaluate_exact_newton_joint_dense(&dormant_binary_states)
6957            .expect("zero-weight binary row must not inspect dormant response/predictors");
6958        assert_eq!(got, expected);
6959    }
6960
6961    #[test]
6962    fn invalid_likelihood_weight_preflight_is_atomic_and_precedes_row_evaluation() {
6963        let mut family = fixed_sigma_binary_test_family();
6964        let beta = array![0.15, 0.25, 0.1, -0.15];
6965        let mut states = latent_binary_states_from_joint_beta(&family, &beta);
6966        // If row evaluation ran first, this active row would fail on its NaN
6967        // predictor before the invalid second weight was discovered.
6968        states[LatentBinaryFamily::BLOCK_MEAN].eta[0] = f64::NAN;
6969        for invalid in [-1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
6970            family.weights = array![1.0, invalid];
6971            let error = family
6972                .evaluate_exact_newton_joint_dense(&states)
6973                .expect_err("invalid likelihood weight must refuse the whole call");
6974            assert!(
6975                error.contains("latent-binary row 2 has invalid likelihood weight"),
6976                "weight preflight did not win atomically for {invalid:?}: {error}"
6977            );
6978            assert!(
6979                !error.contains("predictor"),
6980                "row evaluation ran before weight preflight for {invalid:?}: {error}"
6981            );
6982        }
6983    }
6984
6985    // --- shared latent-interval validation engine: parity / contract tests ---
6986
6987    use crate::survival::location_scale::{TimeBlockInput, TimeBlockMonotonicity};
6988
6989    /// Minimal, structurally valid `TimeBlockInput` for `n` rows and `p_time`
6990    /// columns, used to exercise the shared validation driver without standing
6991    /// up a full term-collection design.
6992    fn validation_time_block(n: usize, p_time: usize) -> TimeBlockInput {
6993        let design = |fill: f64| {
6994            DesignMatrix::Dense(DenseDesignMatrix::from(Array2::from_elem(
6995                (n, p_time),
6996                fill,
6997            )))
6998        };
6999        TimeBlockInput {
7000            design_entry: design(0.1),
7001            design_exit: design(0.2),
7002            design_derivative_exit: design(0.3),
7003            offset_entry: Array1::zeros(n),
7004            offset_exit: Array1::zeros(n),
7005            derivative_offset_exit: Array1::zeros(n),
7006            time_monotonicity: TimeBlockMonotonicity::EnforcedByCoordinateCone,
7007            penalties: Vec::new(),
7008            nullspace_dims: Vec::new(),
7009            initial_log_lambdas: None,
7010            initial_beta: None,
7011        }
7012    }
7013
7014    fn empty_meanspec() -> TermCollectionSpec {
7015        TermCollectionSpec {
7016            linear_terms: Vec::new(),
7017            random_effect_terms: Vec::new(),
7018            smooth_terms: Vec::new(),
7019        }
7020    }
7021
7022    /// Quadratic form `nᵀ M n`.
7023    fn quad_form(m: &Array2<f64>, n_vec: &Array1<f64>) -> f64 {
7024        n_vec.dot(&m.dot(n_vec))
7025    }
7026
7027    /// The latent-survival time-block ridge must COVER the monotone baseline's
7028    /// unpenalized affine null direction — the exact direction whose flat MAP
7029    /// makes the interval-censored warm-start refuse
7030    /// (`check_map_uniqueness`: `nᵀ S n < tol`, dominant block `time_transform`).
7031    ///
7032    /// The I-spline value-space penalty leaves a 1-D null space spanned by the
7033    /// constant coefficient direction (`constant γ ↦ affine log Λ`). We stand it
7034    /// in here with the first-difference penalty `D₁ᵀD₁`, whose null space is
7035    /// exactly that constant/level direction, on a full-rank endpoint design so
7036    /// the function Gram is SPD. Before installation the direction carries no
7037    /// penalty curvature; after installation the appended shrinkage penalty
7038    /// gives it the block's function-norm curvature, so the assembled penalty is
7039    /// no longer flat there.
7040    #[test]
7041    fn latent_time_nullspace_shrinkage_covers_affine_null_direction() {
7042        let n = 12usize;
7043        let p = 4usize;
7044
7045        // Full column-rank endpoint designs (a Vandermonde over a spread grid),
7046        // so the endpoint-averaged function Gram is SPD.
7047        let design_at = |shift: f64| {
7048            Array2::<f64>::from_shape_fn((n, p), |(i, j)| {
7049                let t = shift + (i as f64) / (n as f64 - 1.0);
7050                t.powi(j as i32)
7051            })
7052        };
7053        let x_entry = design_at(0.0);
7054        let x_exit = design_at(0.3);
7055
7056        // First-difference penalty D₁ᵀD₁ (order-1 P-spline): PSD with a 1-D null
7057        // space = the constant (level) direction, mirroring the I-spline
7058        // value-space penalty's affine null.
7059        let mut penalty = Array2::<f64>::zeros((p, p));
7060        for k in 0..p - 1 {
7061            penalty[[k, k]] += 1.0;
7062            penalty[[k + 1, k + 1]] += 1.0;
7063            penalty[[k, k + 1]] -= 1.0;
7064            penalty[[k + 1, k]] -= 1.0;
7065        }
7066
7067        let mut time_block = TimeBlockInput {
7068            design_entry: DesignMatrix::Dense(DenseDesignMatrix::from(x_entry.clone())),
7069            design_exit: DesignMatrix::Dense(DenseDesignMatrix::from(x_exit.clone())),
7070            design_derivative_exit: DesignMatrix::Dense(DenseDesignMatrix::from(x_exit.clone())),
7071            offset_entry: Array1::zeros(n),
7072            offset_exit: Array1::zeros(n),
7073            derivative_offset_exit: Array1::zeros(n),
7074            time_monotonicity: TimeBlockMonotonicity::EnforcedByCoordinateCone,
7075            penalties: vec![penalty.clone()],
7076            nullspace_dims: vec![1],
7077            initial_log_lambdas: Some(Array1::from_elem(1, 0.5)),
7078            initial_beta: None,
7079        };
7080
7081        // The constant (level) direction is the penalty's null direction.
7082        let null_dir = Array1::from_elem(p, 1.0 / (p as f64).sqrt());
7083        let primary_before = quad_form(&penalty, &null_dir);
7084        let penalty_scale = penalty.iter().fold(0.0_f64, |a, &b| a.max(b.abs()));
7085        assert!(
7086            primary_before < 1e-9 * penalty_scale,
7087            "affine direction must start unpenalized: nᵀ S n = {primary_before:.3e}"
7088        );
7089
7090        let installed = install_latent_time_nullspace_shrinkage_penalty(&mut time_block)
7091            .expect("shrinkage installation must succeed on a full-rank endpoint design");
7092        assert!(installed, "a penalty with a null space must gain a ridge");
7093
7094        // Bookkeeping stays self-consistent: one extra penalty, its structural
7095        // nullspace dim recorded as 0 (a full-rank-on-support ridge), and the
7096        // seed ρ vector widened to match.
7097        assert_eq!(time_block.penalties.len(), 2);
7098        assert_eq!(time_block.nullspace_dims, vec![1, 0]);
7099        assert_eq!(
7100            time_block.initial_log_lambdas.as_ref().map(|s| s.len()),
7101            Some(2)
7102        );
7103
7104        // The appended ridge gives the previously-null direction the block's
7105        // function-norm curvature: `nᵀ R n ≈ nᵀ G n`, the endpoint-averaged
7106        // function Gram of the null direction (the Marra & Wood identity for
7107        // `n` spanning the 1-D null space).
7108        let shrinkage = &time_block.penalties[1];
7109        let function_gram =
7110            (x_entry.t().dot(&x_entry) + x_exit.t().dot(&x_exit)).mapv(|v| v / (2 * n) as f64);
7111        let ridge_curvature = quad_form(shrinkage, &null_dir);
7112        let function_norm = quad_form(&function_gram, &null_dir);
7113        assert!(
7114            function_norm > 0.0,
7115            "the null direction has positive function norm on a full-rank design"
7116        );
7117        assert!(
7118            ridge_curvature > 0.5 * function_norm,
7119            "ridge must cover the affine null direction: nᵀ R n = {ridge_curvature:.3e} vs \
7120             nᵀ G n = {function_norm:.3e}"
7121        );
7122        assert!(
7123            (ridge_curvature - function_norm).abs() < 1e-6 * function_norm.max(1.0),
7124            "ridge curvature on the 1-D null must equal the function norm: \
7125             nᵀ R n = {ridge_curvature:.3e}, nᵀ G n = {function_norm:.3e}"
7126        );
7127
7128        // The assembled time-block penalty is no longer flat along the direction
7129        // the MAP-uniqueness audit flagged.
7130        let total: Array2<f64> = &penalty + shrinkage;
7131        assert!(
7132            quad_form(&total, &null_dir) > 0.5 * function_norm,
7133            "assembled penalty must cover the previously-null direction"
7134        );
7135    }
7136
7137    /// A valid two-row latent-survival term spec (one exact event under loaded
7138    /// hazard, one right-censored row).
7139    fn valid_survival_spec(n: usize, p_time: usize) -> LatentSurvivalTermSpec {
7140        LatentSurvivalTermSpec {
7141            age_entry: Array1::zeros(n),
7142            age_exit: Array1::from_elem(n, 1.0),
7143            event_target: Array1::from_shape_fn(n, |i| (i % 2) as u8),
7144            weights: Array1::from_elem(n, 1.0),
7145            derivative_guard: 0.0,
7146            time_block: validation_time_block(n, p_time),
7147            time_design_right: None,
7148            time_offset_right: None,
7149            unloaded_mass_entry: Array1::from_elem(n, 0.01),
7150            unloaded_mass_exit: Array1::from_elem(n, 0.05),
7151            unloaded_mass_right: Array1::zeros(0),
7152            unloaded_hazard_exit: Array1::from_elem(n, 0.02),
7153            meanspec: empty_meanspec(),
7154            mean_offset: Array1::zeros(n),
7155        }
7156    }
7157
7158    /// A valid latent-binary term spec mirroring `valid_survival_spec` but
7159    /// without the per-row unloaded hazard.
7160    fn valid_binary_spec(n: usize, p_time: usize) -> LatentBinaryTermSpec {
7161        LatentBinaryTermSpec {
7162            age_entry: Array1::zeros(n),
7163            age_exit: Array1::from_elem(n, 1.0),
7164            event_target: Array1::from_shape_fn(n, |i| (i % 2) as u8),
7165            weights: Array1::from_elem(n, 1.0),
7166            derivative_guard: 0.0,
7167            time_block: validation_time_block(n, p_time),
7168            unloaded_mass_entry: Array1::from_elem(n, 0.01),
7169            unloaded_mass_exit: Array1::from_elem(n, 0.05),
7170            meanspec: empty_meanspec(),
7171            mean_offset: Array1::zeros(n),
7172        }
7173    }
7174
7175    fn loaded_frailty() -> FrailtySpec {
7176        FrailtySpec::HazardMultiplier {
7177            scale: FrailtyScale::Fixed { sigma: 0.3 },
7178            loading: HazardLoading::LoadedVsUnloaded,
7179        }
7180    }
7181
7182    /// Both adapters route through the shared `validate_latent_interval_inputs`
7183    /// engine, but each must still emit its own context prefix and (for the
7184    /// size-mismatch / unloaded-decomposition diagnostics) the hazard-aware vs
7185    /// mass-only message variant. This pins the byte-for-byte contract the
7186    /// unification had to preserve, the property the issue's "old vs new
7187    /// validation errors" parity test guards.
7188    #[test]
7189    fn latent_interval_validation_parity_across_models() {
7190        let n = 2;
7191        let p_time = 2;
7192        let data = Array2::<f64>::zeros((n, 3));
7193
7194        // 1. A clean spec validates and round-trips the typed scale.
7195        //    Binary then extracts the required fixed scalar.
7196        let surv_sigma = validate_latent_survival_inputs(
7197            data.view(),
7198            &valid_survival_spec(n, p_time),
7199            &loaded_frailty(),
7200        )
7201        .expect("valid survival spec must validate");
7202        assert_eq!(surv_sigma, FrailtyScale::Fixed { sigma: 0.3 });
7203        let bin_sigma = validate_latent_binary_inputs(
7204            data.view(),
7205            &valid_binary_spec(n, p_time),
7206            &loaded_frailty(),
7207        )
7208        .expect("valid binary spec must validate");
7209        assert_eq!(bin_sigma, 0.3);
7210
7211        // 2. Empty data: shared driver, per-model context prefix.
7212        let empty = Array2::<f64>::zeros((0, 3));
7213        let surv_empty = validate_latent_survival_inputs(
7214            empty.view(),
7215            &valid_survival_spec(n, p_time),
7216            &loaded_frailty(),
7217        )
7218        .expect_err("empty data must be rejected");
7219        assert_eq!(
7220            surv_empty.to_string(),
7221            "latent-survival requires a non-empty dataset"
7222        );
7223        let bin_empty = validate_latent_binary_inputs(
7224            empty.view(),
7225            &valid_binary_spec(n, p_time),
7226            &loaded_frailty(),
7227        )
7228        .expect_err("empty data must be rejected");
7229        assert_eq!(
7230            bin_empty.to_string(),
7231            "latent-binary requires a non-empty dataset"
7232        );
7233
7234        // 3. Size mismatch: survival's message carries `unloaded_hazard=`,
7235        //    binary's does not. This is the one shape that distinguishes the
7236        //    two row views feeding the shared driver.
7237        let mut surv_bad = valid_survival_spec(n, p_time);
7238        surv_bad.weights = Array1::from_elem(n + 1, 1.0);
7239        let surv_size = validate_latent_survival_inputs(data.view(), &surv_bad, &loaded_frailty())
7240            .expect_err("size mismatch must be rejected");
7241        let surv_msg = surv_size.to_string();
7242        assert!(
7243            surv_msg.starts_with("latent-survival size mismatch")
7244                && surv_msg.contains("unloaded_hazard="),
7245            "survival size-mismatch message must include unloaded_hazard: {surv_msg}"
7246        );
7247        let mut bin_bad = valid_binary_spec(n, p_time);
7248        bin_bad.weights = Array1::from_elem(n + 1, 1.0);
7249        let bin_size = validate_latent_binary_inputs(data.view(), &bin_bad, &loaded_frailty())
7250            .expect_err("size mismatch must be rejected");
7251        let bin_msg = bin_size.to_string();
7252        assert!(
7253            bin_msg.starts_with("latent-binary size mismatch")
7254                && !bin_msg.contains("unloaded_hazard"),
7255            "binary size-mismatch message must omit unloaded_hazard: {bin_msg}"
7256        );
7257
7258        // 4. Invalid unloaded decomposition: survival reports `exit_hazard=`,
7259        //    binary reports only the two masses.
7260        let mut surv_neg_hazard = valid_survival_spec(n, p_time);
7261        surv_neg_hazard.unloaded_hazard_exit[0] = -1.0;
7262        let surv_decomp =
7263            validate_latent_survival_inputs(data.view(), &surv_neg_hazard, &loaded_frailty())
7264                .expect_err("negative unloaded hazard must be rejected");
7265        assert_eq!(
7266            surv_decomp.to_string(),
7267            "latent-survival row 1 has invalid unloaded hazard decomposition: entry_mass=0.01, exit_mass=0.05, exit_hazard=-1"
7268        );
7269        let mut bin_bad_mass = valid_binary_spec(n, p_time);
7270        bin_bad_mass.unloaded_mass_exit[0] = 0.0; // exit < entry
7271        let bin_decomp =
7272            validate_latent_binary_inputs(data.view(), &bin_bad_mass, &loaded_frailty())
7273                .expect_err("non-monotone unloaded mass must be rejected");
7274        assert_eq!(
7275            bin_decomp.to_string(),
7276            "latent-binary row 1 has invalid unloaded mass decomposition: entry_mass=0.01, exit_mass=0"
7277        );
7278
7279        // 5. Per-row interval/event/weight diagnostics share one engine, so an
7280        //    identical invalid input yields identical (modulo prefix) text.
7281        let mut surv_event = valid_survival_spec(n, p_time);
7282        surv_event.event_target[1] = 7;
7283        let surv_event_err =
7284            validate_latent_survival_inputs(data.view(), &surv_event, &loaded_frailty())
7285                .expect_err("invalid event target must be rejected");
7286        assert_eq!(
7287            surv_event_err.to_string(),
7288            "latent-survival row 2 has invalid event target 7; expected 0 or 1"
7289        );
7290        let mut bin_event = valid_binary_spec(n, p_time);
7291        bin_event.event_target[1] = 7;
7292        let bin_event_err =
7293            validate_latent_binary_inputs(data.view(), &bin_event, &loaded_frailty())
7294                .expect_err("invalid event target must be rejected");
7295        assert_eq!(
7296            bin_event_err.to_string(),
7297            "latent-binary row 2 has invalid event target 7; expected 0 or 1"
7298        );
7299
7300        // 6. Frailty policy divergence: survival accepts a learned scale,
7301        //    binary rejects it.
7302        let learnable = FrailtySpec::HazardMultiplier {
7303            scale: FrailtyScale::Learned { initial_sigma: 0.5 },
7304            loading: HazardLoading::LoadedVsUnloaded,
7305        };
7306        let surv_learnable = validate_latent_survival_inputs(
7307            data.view(),
7308            &valid_survival_spec(n, p_time),
7309            &learnable,
7310        )
7311        .expect("survival accepts a learnable latent scale");
7312        assert_eq!(
7313            surv_learnable,
7314            FrailtyScale::Learned { initial_sigma: 0.5 }
7315        );
7316        let bin_learnable =
7317            validate_latent_binary_inputs(data.view(), &valid_binary_spec(n, p_time), &learnable)
7318                .expect_err("binary requires a fixed latent scale");
7319        assert_eq!(
7320            bin_learnable.to_string(),
7321            "latent-binary currently requires a fixed hazard-multiplier sigma"
7322        );
7323
7324        // 7. The time-block shape check is owned by the shared driver: a
7325        //    column-count mismatch is reported with the per-model prefix.
7326        let mut surv_time_bad = valid_survival_spec(n, p_time);
7327        surv_time_bad.time_block.design_entry = DesignMatrix::Dense(DenseDesignMatrix::from(
7328            Array2::from_elem((n, p_time + 1), 0.1),
7329        ));
7330        let surv_time_err =
7331            validate_latent_survival_inputs(data.view(), &surv_time_bad, &loaded_frailty())
7332                .expect_err("time block column mismatch must be rejected");
7333        assert!(
7334            surv_time_err
7335                .to_string()
7336                .starts_with("latent-survival time block column mismatch"),
7337            "unexpected survival time-block message: {surv_time_err}"
7338        );
7339    }
7340
7341    #[test]
7342    fn latent_interval_validation_treats_exact_zero_rows_as_response_dormant() {
7343        let n = 2;
7344        let data = Array2::<f64>::zeros((n, 1));
7345
7346        let mut survival = valid_survival_spec(n, 1);
7347        survival.weights[0] = 0.0;
7348        survival.age_entry[0] = f64::NAN;
7349        survival.age_exit[0] = f64::NEG_INFINITY;
7350        survival.event_target[0] = 19;
7351        survival.unloaded_mass_entry[0] = f64::NAN;
7352        survival.unloaded_mass_exit[0] = -1.0;
7353        survival.unloaded_hazard_exit[0] = f64::INFINITY;
7354        validate_latent_survival_inputs(data.view(), &survival, &loaded_frailty())
7355            .expect("zero-weight survival response row must be dormant");
7356
7357        let mut binary = valid_binary_spec(n, 1);
7358        binary.weights[0] = 0.0;
7359        binary.age_entry[0] = f64::NAN;
7360        binary.age_exit[0] = f64::NEG_INFINITY;
7361        binary.event_target[0] = 19;
7362        binary.unloaded_mass_entry[0] = f64::NAN;
7363        binary.unloaded_mass_exit[0] = -1.0;
7364        validate_latent_binary_inputs(data.view(), &binary, &loaded_frailty())
7365            .expect("zero-weight binary response row must be dormant");
7366
7367        // The weight scan is whole-vector and precedes response geometry. A
7368        // later invalid weight therefore wins deterministically over the bad
7369        // active response in row 1.
7370        survival.weights = array![1.0, f64::NAN];
7371        let error = validate_latent_survival_inputs(data.view(), &survival, &loaded_frailty())
7372            .expect_err("non-finite weight must atomically refuse validation");
7373        assert!(
7374            error
7375                .to_string()
7376                .contains("latent-survival row 2 has invalid weight"),
7377            "unexpected atomic preflight error: {error}"
7378        );
7379    }
7380
7381    #[test]
7382    fn latent_survival_coefficient_cost_uses_joint_coupled_formula() {
7383        // `evaluate_exact_newton_joint_dense` builds a fully dense joint
7384        // Hessian over (Σ p_b)² across the time, mean, and log-σ blocks via
7385        // per-row pullback of the latent-survival primary kernel. The override
7386        // must reflect that joint coupling rather than the block-diagonal
7387        // default.
7388        let family = learnable_sigma_test_family();
7389        let n = family.event_target.len() as u64;
7390        let p_time = 2u64;
7391        let p_mean = 2u64;
7392        let p_log_sigma = 1u64;
7393        let specs = vec![
7394            ParameterBlockSpec {
7395                name: "time".to_string(),
7396                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((
7397                    n as usize,
7398                    p_time as usize,
7399                )))),
7400                offset: Array1::zeros(n as usize),
7401                penalties: Vec::new(),
7402                nullspace_dims: Vec::new(),
7403                initial_log_lambdas: Array1::zeros(0),
7404                initial_beta: None,
7405                gauge_priority: 100,
7406                jacobian_callback: None,
7407                stacked_design: None,
7408                stacked_offset: None,
7409            },
7410            ParameterBlockSpec {
7411                name: "mean".to_string(),
7412                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((
7413                    n as usize,
7414                    p_mean as usize,
7415                )))),
7416                offset: Array1::zeros(n as usize),
7417                penalties: Vec::new(),
7418                nullspace_dims: Vec::new(),
7419                initial_log_lambdas: Array1::zeros(0),
7420                initial_beta: None,
7421                gauge_priority: 100,
7422                jacobian_callback: None,
7423                stacked_design: None,
7424                stacked_offset: None,
7425            },
7426            ParameterBlockSpec {
7427                name: "log_sigma".to_string(),
7428                design: DesignMatrix::Dense(DenseDesignMatrix::from(Array2::zeros((
7429                    n as usize,
7430                    p_log_sigma as usize,
7431                )))),
7432                offset: Array1::zeros(n as usize),
7433                penalties: Vec::new(),
7434                nullspace_dims: Vec::new(),
7435                initial_log_lambdas: Array1::zeros(0),
7436                initial_beta: None,
7437                gauge_priority: 100,
7438                jacobian_callback: None,
7439                stacked_design: None,
7440                stacked_offset: None,
7441            },
7442        ];
7443        let p_total = p_time + p_mean + p_log_sigma;
7444        let expected_joint = n * p_total * p_total;
7445        let expected_block_diag =
7446            n * (p_time * p_time + p_mean * p_mean + p_log_sigma * p_log_sigma);
7447        assert_eq!(family.coefficient_hessian_cost(&specs), expected_joint);
7448        // Cross-block fill (time–mean, time–log_sigma, mean–log_sigma) makes
7449        // the joint cost strictly larger than the block-diagonal default.
7450        assert!(expected_joint > expected_block_diag);
7451    }
7452
7453    #[test]
7454    fn latent_family_planner_keeps_outer_hessian_at_large_n() {
7455        use crate::custom_family::custom_family_outer_derivatives;
7456        use gam_problem::{DeclaredHessianForm, Derivative};
7457
7458        let options = BlockwiseFitOptions::default();
7459        let large_n = 50_001;
7460
7461        let survival = learnable_sigma_test_family();
7462        let survival_specs =
7463            latent_test_specs(large_n, &[("time", 2), ("mean", 2), ("log_sigma", 1)]);
7464        let (surv_grad, surv_hess) =
7465            custom_family_outer_derivatives(&survival, &survival_specs, &options);
7466        assert_eq!(surv_grad, Derivative::Analytic);
7467        assert_eq!(surv_hess, DeclaredHessianForm::Either);
7468
7469        let binary = fixed_sigma_binary_test_family();
7470        let binary_specs = latent_test_specs(large_n, &[("time", 2), ("mean", 2)]);
7471        let (bin_grad, bin_hess) =
7472            custom_family_outer_derivatives(&binary, &binary_specs, &options);
7473        assert_eq!(bin_grad, Derivative::Analytic);
7474        assert_eq!(bin_hess, DeclaredHessianForm::Either);
7475    }
7476
7477    #[test]
7478    fn latent_families_arm_self_vanishing_levenberg_on_ill_conditioning() {
7479        // Regression guard for #1108. The interval-censored row contribution
7480        // `ℓ = log[S(L) − S(R)]` is the log of a DIFFERENCE of survival kernels and
7481        // is legitimately NON-concave (indefinite per-row Hessian) away from the
7482        // optimum; on the constrained (monotone-cone) coupled time block this can
7483        // make the penalized joint Hessian full-rank yet indefinite / severely
7484        // ill-conditioned at the cold-start seed. The coupled exact-joint inner
7485        // solver only adds the self-vanishing Levenberg–Marquardt diagonal floor
7486        // (the cure for a full-rank ill-conditioned reflected QP that otherwise
7487        // oscillates the trust region into a snapshot-less stall) when the family
7488        // opts in via `levenberg_on_ill_conditioning()`. Both latent families MUST
7489        // keep this armed (the default is `false`, which leaves the interval inner
7490        // solve diverging with "exited the joint Newton path before convergence").
7491        assert!(
7492            learnable_sigma_test_family().levenberg_on_ill_conditioning(),
7493            "LatentSurvivalFamily must arm the self-vanishing Levenberg floor so the \
7494             indefinite interval-censored joint Hessian converges (see #1108)"
7495        );
7496        assert!(
7497            fixed_sigma_binary_test_family().levenberg_on_ill_conditioning(),
7498            "LatentBinaryFamily must arm the self-vanishing Levenberg floor on its \
7499             constrained coupled time block (see #1108)"
7500        );
7501    }
7502
7503    #[test]
7504    fn latent_binary_exact_joint_hessian_and_workspace_matvec_match_fd() {
7505        let family = fixed_sigma_binary_test_family();
7506        let beta = array![0.15, 0.25, 0.1, -0.15];
7507        let states = latent_binary_states_from_joint_beta(&family, &beta);
7508        let h = 1e-6;
7509
7510        let analytic_hessian = family
7511            .exact_newton_joint_hessian(&states)
7512            .expect("analytic latent binary joint hessian evaluation")
7513            .expect("latent binary should expose exact joint hessian");
7514
7515        for j in 0..beta.len() {
7516            let mut beta_plus = beta.clone();
7517            beta_plus[j] += h;
7518            let gradient_plus = family
7519                .exact_newton_joint_gradient_evaluation(
7520                    &latent_binary_states_from_joint_beta(&family, &beta_plus),
7521                    &[],
7522                )
7523                .expect("joint gradient plus")
7524                .expect("joint gradient should exist")
7525                .gradient;
7526
7527            let mut beta_minus = beta.clone();
7528            beta_minus[j] -= h;
7529            let gradient_minus = family
7530                .exact_newton_joint_gradient_evaluation(
7531                    &latent_binary_states_from_joint_beta(&family, &beta_minus),
7532                    &[],
7533                )
7534                .expect("joint gradient minus")
7535                .expect("joint gradient should exist")
7536                .gradient;
7537
7538            let fd_column = -((&gradient_plus - &gradient_minus) / (2.0 * h));
7539            let analytic_column = analytic_hessian.column(j).to_owned();
7540            let rel = max_relative_array1(&analytic_column, &fd_column);
7541            assert!(
7542                rel < 5e-4,
7543                "latent binary joint Hessian column {j} mismatch: rel={rel}, analytic={analytic_column:?}, fd={fd_column:?}"
7544            );
7545        }
7546
7547        let workspace = family
7548            .exact_newton_joint_hessian_workspace(&states, &[])
7549            .expect("latent binary hessian workspace")
7550            .expect("workspace should exist");
7551        let direction = array![0.4, -0.2, 0.3, 0.1];
7552        let hv = workspace
7553            .hessian_matvec(&direction)
7554            .expect("workspace matvec")
7555            .expect("workspace should support matvec");
7556        let dense_hv = analytic_hessian.dot(&direction);
7557        assert!(
7558            max_relative_array1(&hv, &dense_hv) < 1e-12,
7559            "latent binary workspace HVP mismatch: hv={hv:?}, dense={dense_hv:?}"
7560        );
7561
7562        let dh = workspace
7563            .directional_derivative(&direction)
7564            .expect("workspace dH")
7565            .expect("workspace should support dH");
7566        let fd_step = 1e-5;
7567        let h_plus = family
7568            .exact_newton_joint_hessian(&latent_binary_states_from_joint_beta(
7569                &family,
7570                &(beta.clone() + &(fd_step * &direction)),
7571            ))
7572            .expect("hessian plus")
7573            .expect("hessian plus should exist");
7574        let h_minus = family
7575            .exact_newton_joint_hessian(&latent_binary_states_from_joint_beta(
7576                &family,
7577                &(beta - &(fd_step * &direction)),
7578            ))
7579            .expect("hessian minus")
7580            .expect("hessian minus should exist");
7581        let fd_dh = (&h_plus - &h_minus) / (2.0 * fd_step);
7582        assert!(
7583            max_relative_array2(&dh, &fd_dh) < 2e-4,
7584            "latent binary workspace dH mismatch: dh={dh:?}, fd={fd_dh:?}"
7585        );
7586
7587        let direction_v = array![-0.15, 0.25, 0.08, -0.12];
7588        let d2h = family
7589            .exact_newton_joint_hessiansecond_directional_derivative(
7590                &states,
7591                &direction,
7592                &direction_v,
7593            )
7594            .expect("latent binary d2H")
7595            .expect("latent binary should expose d2H");
7596        let d2_step = 4e-4;
7597        let dh_plus = family
7598            .exact_newton_joint_hessian_directional_derivative(
7599                &latent_binary_states_from_joint_beta(
7600                    &family,
7601                    &(array![0.15, 0.25, 0.1, -0.15] + d2_step * &direction_v),
7602                ),
7603                &direction,
7604            )
7605            .expect("latent binary dH plus")
7606            .expect("latent binary should expose dH plus");
7607        let dh_minus = family
7608            .exact_newton_joint_hessian_directional_derivative(
7609                &latent_binary_states_from_joint_beta(
7610                    &family,
7611                    &(array![0.15, 0.25, 0.1, -0.15] - d2_step * &direction_v),
7612                ),
7613                &direction,
7614            )
7615            .expect("latent binary dH minus")
7616            .expect("latent binary should expose dH minus");
7617        let fd_d2h = (dh_plus - dh_minus) / (2.0 * d2_step);
7618        let d2h_rel = frobenius_relative_array2(&d2h, &fd_d2h);
7619        assert!(
7620            d2h_rel < 2e-2,
7621            "latent binary combined TwoSeed d2H mismatch: rel={d2h_rel}, analytic={d2h:?}, fd={fd_d2h:?}"
7622        );
7623    }
7624
7625    #[test]
7626    fn latent_survival_learnable_sigma_block_matches_family_fd() {
7627        let family = learnable_sigma_test_family();
7628        let beta = learnable_sigma_test_joint_beta();
7629        let states = latent_survival_states_from_joint_beta(&family, &beta);
7630        let slices = family.joint_slices();
7631        let sigma_idx = slices
7632            .log_sigma
7633            .as_ref()
7634            .expect("learnable sigma test family should expose log_sigma")
7635            .start;
7636        let h = 2e-4;
7637
7638        let eval = family
7639            .evaluate(&states)
7640            .expect("learnable latent survival evaluation");
7641        let joint_gradient = family
7642            .exact_newton_joint_gradient_evaluation(&states, &[])
7643            .expect("joint gradient evaluation")
7644            .expect("joint gradient should exist")
7645            .gradient;
7646        let joint_hessian = family
7647            .exact_newton_joint_hessian(&states)
7648            .expect("joint hessian evaluation")
7649            .expect("joint hessian should exist");
7650        assert_eq!(eval.blockworking_sets.len(), 3);
7651
7652        let (block_grad, block_neg_hess) =
7653            match &eval.blockworking_sets[LatentSurvivalFamily::BLOCK_LOG_SIGMA] {
7654                BlockWorkingSet::ExactNewton { gradient, hessian } => {
7655                    let neg_hess = match hessian {
7656                        SymmetricMatrix::Dense(mat) => mat[[0, 0]],
7657                        _ => panic!("log_sigma block should use a dense exact-Newton Hessian"),
7658                    };
7659                    (gradient[0], neg_hess)
7660                }
7661                _ => panic!("log_sigma block should use ExactNewton"),
7662            };
7663
7664        assert!((block_grad - joint_gradient[sigma_idx]).abs() < 1e-12);
7665        assert!((block_neg_hess - joint_hessian[[sigma_idx, sigma_idx]]).abs() < 1e-12);
7666
7667        let mut beta_plus = beta.clone();
7668        beta_plus[sigma_idx] += h;
7669        let ll_plus = family
7670            .log_likelihood_only(&latent_survival_states_from_joint_beta(&family, &beta_plus))
7671            .expect("ll plus");
7672        let ll_0 = family.log_likelihood_only(&states).expect("ll base");
7673        let mut beta_minus = beta.clone();
7674        beta_minus[sigma_idx] -= h;
7675        let ll_minus = family
7676            .log_likelihood_only(&latent_survival_states_from_joint_beta(
7677                &family,
7678                &beta_minus,
7679            ))
7680            .expect("ll minus");
7681
7682        let fd_grad = (ll_plus - ll_minus) / (2.0 * h);
7683        let fd_neg_hess = -(ll_plus - 2.0 * ll_0 + ll_minus) / (h * h);
7684        assert!(
7685            (joint_gradient[sigma_idx] - fd_grad).abs()
7686                / joint_gradient[sigma_idx]
7687                    .abs()
7688                    .max(fd_grad.abs())
7689                    .max(1e-12)
7690                < 2e-3,
7691            "family log_sigma grad={}, fd={fd_grad}",
7692            joint_gradient[sigma_idx]
7693        );
7694        assert!(
7695            (joint_hessian[[sigma_idx, sigma_idx]] - fd_neg_hess).abs()
7696                / joint_hessian[[sigma_idx, sigma_idx]]
7697                    .abs()
7698                    .max(fd_neg_hess.abs())
7699                    .max(1e-10)
7700                < 2e-2,
7701            "family log_sigma neg_hess={}, fd={fd_neg_hess}",
7702            joint_hessian[[sigma_idx, sigma_idx]]
7703        );
7704    }
7705
7706    #[test]
7707    fn latent_survival_exact_joint_hessian_matches_gradient_fd() {
7708        let family = learnable_sigma_test_family();
7709        let beta = learnable_sigma_test_joint_beta();
7710        let states = latent_survival_states_from_joint_beta(&family, &beta);
7711        let h = 1e-6;
7712
7713        let analytic_hessian = family
7714            .exact_newton_joint_hessian(&states)
7715            .expect("analytic joint hessian evaluation")
7716            .expect("latent survival should expose exact joint hessian");
7717
7718        for j in 0..beta.len() {
7719            let mut beta_plus = beta.clone();
7720            beta_plus[j] += h;
7721            let gradient_plus = family
7722                .exact_newton_joint_gradient_evaluation(
7723                    &latent_survival_states_from_joint_beta(&family, &beta_plus),
7724                    &[],
7725                )
7726                .expect("joint gradient plus")
7727                .expect("joint gradient should exist")
7728                .gradient;
7729
7730            let mut beta_minus = beta.clone();
7731            beta_minus[j] -= h;
7732            let gradient_minus = family
7733                .exact_newton_joint_gradient_evaluation(
7734                    &latent_survival_states_from_joint_beta(&family, &beta_minus),
7735                    &[],
7736                )
7737                .expect("joint gradient minus")
7738                .expect("joint gradient should exist")
7739                .gradient;
7740
7741            let fd_column = (&gradient_plus - &gradient_minus) / (2.0 * h);
7742            let analytic_column = analytic_hessian.column(j).to_owned();
7743            let rel = max_relative_array1(&analytic_column, &(-fd_column));
7744            assert!(
7745                rel < 5e-4,
7746                "joint Hessian column {j} mismatch: rel={rel}, analytic={analytic_column:?}, fd={:?}",
7747                -((&gradient_plus - &gradient_minus) / (2.0 * h))
7748            );
7749        }
7750    }
7751
7752    /// FD check for `LatentSurvivalFamily::offset_channel_residuals`: each
7753    /// channel residual sums to `∂(−ℓ)/∂o_ch` for a uniform additive offset on
7754    /// that time channel (the baseline-θ enters only through these offsets).
7755    /// `o_ch` shifts `eta_time[ch-slice]` uniformly, so `Σ_i r^ch_i` is exactly
7756    /// the directional derivative of `−ℓ` along a constant offset on channel ch.
7757    /// This validates the envelope-theorem latent baseline-θ gradient primitive.
7758    #[test]
7759    fn latent_survival_offset_channel_residuals_match_finite_difference() {
7760        let family = survival_stress_test_family(24);
7761        let beta = survival_stress_test_joint_beta();
7762        let states = latent_survival_states_from_joint_beta(&family, &beta);
7763        let n = family.event_target.len();
7764
7765        let residuals = family
7766            .offset_channel_residuals(&states)
7767            .expect("offset channel residuals");
7768        let sum_entry: f64 = residuals.entry.sum();
7769        let sum_exit: f64 = residuals.exit.sum();
7770        let sum_deriv: f64 = residuals.derivative.sum();
7771
7772        #[derive(Clone, Copy)]
7773        enum TimeOffsetChannel {
7774            Entry,
7775            Exit,
7776            Derivative,
7777        }
7778
7779        // `−ℓ` after shifting one time channel's eta by a constant δ.
7780        let neg_ll_with_offset = |channel: TimeOffsetChannel, delta: f64| -> f64 {
7781            let mut shifted = states.clone();
7782            let slice = match channel {
7783                TimeOffsetChannel::Entry => s![0..n],
7784                TimeOffsetChannel::Exit => s![n..2 * n],
7785                TimeOffsetChannel::Derivative => s![2 * n..3 * n],
7786            };
7787            shifted[LatentSurvivalFamily::BLOCK_TIME]
7788                .eta
7789                .slice_mut(slice)
7790                .mapv_inplace(|v| v + delta);
7791            let (ll, _) = family
7792                .evaluate_exact_newton_joint_gradient_dense(&shifted)
7793                .expect("shifted joint gradient evaluation");
7794            -ll
7795        };
7796
7797        let h = 1e-6;
7798        let fd_entry = (neg_ll_with_offset(TimeOffsetChannel::Entry, h)
7799            - neg_ll_with_offset(TimeOffsetChannel::Entry, -h))
7800            / (2.0 * h);
7801        let fd_exit = (neg_ll_with_offset(TimeOffsetChannel::Exit, h)
7802            - neg_ll_with_offset(TimeOffsetChannel::Exit, -h))
7803            / (2.0 * h);
7804        let fd_deriv = (neg_ll_with_offset(TimeOffsetChannel::Derivative, h)
7805            - neg_ll_with_offset(TimeOffsetChannel::Derivative, -h))
7806            / (2.0 * h);
7807
7808        assert!(
7809            (sum_entry - fd_entry).abs() <= 1e-5 * fd_entry.abs().max(1.0),
7810            "entry-channel residual sum mismatch: analytic={sum_entry}, fd={fd_entry}"
7811        );
7812        assert!(
7813            (sum_exit - fd_exit).abs() <= 1e-5 * fd_exit.abs().max(1.0),
7814            "exit-channel residual sum mismatch: analytic={sum_exit}, fd={fd_exit}"
7815        );
7816        assert!(
7817            (sum_deriv - fd_deriv).abs() <= 1e-5 * fd_deriv.abs().max(1.0),
7818            "derivative-channel residual sum mismatch: analytic={sum_deriv}, fd={fd_deriv}"
7819        );
7820    }
7821
7822    #[test]
7823    fn latent_survival_exact_joint_parallel_stress_is_repeatable() {
7824        let family = survival_stress_test_family(96);
7825        let beta = survival_stress_test_joint_beta();
7826        let states = latent_survival_states_from_joint_beta(&family, &beta);
7827        let direction_u = array![0.03, -0.02, 0.01, 0.04, -0.015, 0.025, -0.005, 0.02];
7828        let direction_v = array![-0.01, 0.035, -0.025, 0.015, 0.02, -0.01, 0.03, -0.015];
7829
7830        let (ll_a, grad_a) = family
7831            .evaluate_exact_newton_joint_gradient_dense(&states)
7832            .expect("stress joint gradient evaluation");
7833        let (ll_b, grad_b) = family
7834            .evaluate_exact_newton_joint_gradient_dense(&states)
7835            .expect("repeat stress joint gradient evaluation");
7836        assert_eq!(ll_a.to_bits(), ll_b.to_bits());
7837        assert_eq!(grad_a, grad_b);
7838
7839        let (joint_ll_a, joint_grad_a, hess_a) = family
7840            .evaluate_exact_newton_joint_dense(&states)
7841            .expect("stress joint dense evaluation");
7842        let (joint_ll_b, joint_grad_b, hess_b) = family
7843            .evaluate_exact_newton_joint_dense(&states)
7844            .expect("repeat stress joint dense evaluation");
7845        assert_eq!(joint_ll_a.to_bits(), joint_ll_b.to_bits());
7846        assert_eq!(joint_grad_a, joint_grad_b);
7847        assert_eq!(hess_a, hess_b);
7848        assert!(hess_a.iter().all(|value| value.is_finite()));
7849        assert!(max_relative_array2(&hess_a, &hess_a.t().to_owned()) < 1e-12);
7850
7851        let dh_a = family
7852            .exact_newton_joint_hessian_directional_derivative_dense(&states, &direction_u)
7853            .expect("stress joint dH evaluation");
7854        let dh_b = family
7855            .exact_newton_joint_hessian_directional_derivative_dense(&states, &direction_u)
7856            .expect("repeat stress joint dH evaluation");
7857        assert_eq!(dh_a, dh_b);
7858        assert!(dh_a.iter().all(|value| value.is_finite()));
7859        assert!(max_relative_array2(&dh_a, &dh_a.t().to_owned()) < 1e-12);
7860
7861        let d2h_a = family
7862            .exact_newton_joint_hessian_second_directional_derivative_dense(
7863                &states,
7864                &direction_u,
7865                &direction_v,
7866            )
7867            .expect("stress joint d2H evaluation");
7868        let d2h_b = family
7869            .exact_newton_joint_hessian_second_directional_derivative_dense(
7870                &states,
7871                &direction_u,
7872                &direction_v,
7873            )
7874            .expect("repeat stress joint d2H evaluation");
7875        assert_eq!(d2h_a, d2h_b);
7876        assert!(d2h_a.iter().all(|value| value.is_finite()));
7877        assert!(max_relative_array2(&d2h_a, &d2h_a.t().to_owned()) < 1e-12);
7878    }
7879
7880    #[test]
7881    fn latent_survival_exact_joint_dh_matches_hessian_fd() {
7882        let family = learnable_sigma_test_family();
7883        let beta = learnable_sigma_test_joint_beta();
7884        let states = latent_survival_states_from_joint_beta(&family, &beta);
7885        let h = 2e-4;
7886        let direction = array![0.07, -0.03, 0.05, 0.02, -0.04];
7887
7888        let analytic = family
7889            .exact_newton_joint_hessian_directional_derivative(&states, &direction)
7890            .expect("analytic joint dH evaluation")
7891            .expect("latent survival should expose exact joint dH");
7892
7893        let hessian_plus = family
7894            .exact_newton_joint_hessian(&latent_survival_states_from_joint_beta(
7895                &family,
7896                &(beta.clone() + h * &direction),
7897            ))
7898            .expect("joint hessian plus")
7899            .expect("joint hessian should exist");
7900        let hessian_minus = family
7901            .exact_newton_joint_hessian(&latent_survival_states_from_joint_beta(
7902                &family,
7903                &(beta.clone() - h * &direction),
7904            ))
7905            .expect("joint hessian minus")
7906            .expect("joint hessian should exist");
7907
7908        let fd = (&hessian_plus - &hessian_minus) / (2.0 * h);
7909        let rel = frobenius_relative_array2(&analytic, &fd);
7910        assert!(rel < 2e-3, "joint dH mismatch: rel={rel}");
7911    }
7912
7913    #[test]
7914    fn latent_survival_exact_joint_d2h_matches_directional_fd() {
7915        let family = learnable_sigma_test_family();
7916        let beta = learnable_sigma_test_joint_beta();
7917        let states = latent_survival_states_from_joint_beta(&family, &beta);
7918        let h = 5e-4;
7919        let direction_u = array![0.07, -0.03, 0.05, 0.02, -0.04];
7920        let direction_v = array![-0.02, 0.06, -0.01, 0.03, 0.05];
7921
7922        let analytic = family
7923            .exact_newton_joint_hessiansecond_directional_derivative(
7924                &states,
7925                &direction_u,
7926                &direction_v,
7927            )
7928            .expect("analytic joint d2H evaluation")
7929            .expect("latent survival should expose exact joint d2H");
7930        let swapped = family
7931            .exact_newton_joint_hessiansecond_directional_derivative(
7932                &states,
7933                &direction_v,
7934                &direction_u,
7935            )
7936            .expect("swapped analytic joint d2H evaluation")
7937            .expect("latent survival should expose exact joint d2H");
7938        let symmetry_rel = max_relative_array2(&analytic, &swapped);
7939        assert!(
7940            symmetry_rel < 1e-10,
7941            "joint d2H should be symmetric in directions, got rel={symmetry_rel}"
7942        );
7943
7944        let dh_plus = family
7945            .exact_newton_joint_hessian_directional_derivative(
7946                &latent_survival_states_from_joint_beta(
7947                    &family,
7948                    &(beta.clone() + h * &direction_v),
7949                ),
7950                &direction_u,
7951            )
7952            .expect("joint dH plus")
7953            .expect("joint dH should exist");
7954        let dh_minus = family
7955            .exact_newton_joint_hessian_directional_derivative(
7956                &latent_survival_states_from_joint_beta(
7957                    &family,
7958                    &(beta.clone() - h * &direction_v),
7959                ),
7960                &direction_u,
7961            )
7962            .expect("joint dH minus")
7963            .expect("joint dH should exist");
7964
7965        let fd = (&dh_plus - &dh_minus) / (2.0 * h);
7966        let rel = frobenius_relative_array2(&analytic, &fd);
7967        assert!(rel < 2.5e-2, "joint d2H mismatch: rel={rel}");
7968    }
7969
7970    #[test]
7971    fn latent_survival_row_primary_derivatives_match_fd() {
7972        let quadctx = QuadratureContext::new();
7973        let row = LatentSurvivalRow::exact_event(0.35, 1.4, 0.1, 0.45, 0.8, 0.12);
7974        // [q_entry, q_exit, qdot_exit, q_right, mu, log_sigma]. This is an
7975        // exact-event row, so the `q_right` channel is inert (the likelihood
7976        // does not depend on it); the FD loop below confirms its gradient/Hessian
7977        // entries are zero.
7978        let primary = array![
7979            0.35f64.ln(),
7980            1.4f64.ln(),
7981            0.8,
7982            1.6f64.ln(),
7983            -0.2,
7984            0.4f64.ln()
7985        ];
7986        let sigma = primary[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA].exp();
7987        let h_grad = 1e-6;
7988        let h_hess = 2e-4;
7989
7990        let (_, gradient, neg_hessian) = latent_survival_row_primary_gradient_hessian(
7991            &quadctx,
7992            &row,
7993            LatentSurvivalPrimaryPoint {
7994                q_entry: primary[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
7995                q_exit: primary[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
7996                qdot_exit: primary[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
7997                q_right: primary[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
7998                mu: primary[LATENT_SURVIVAL_PRIMARY_MU],
7999                sigma,
8000            },
8001            true,
8002        )
8003        .expect("analytic row primary gradient/hessian");
8004
8005        for j in 0..LATENT_SURVIVAL_PRIMARY_DIM {
8006            let mut plus = primary.clone();
8007            plus[j] += h_grad;
8008            let mut minus = primary.clone();
8009            minus[j] -= h_grad;
8010            let fd_grad = (latent_survival_row_loglik_from_primary(&quadctx, &row, &plus)
8011                - latent_survival_row_loglik_from_primary(&quadctx, &row, &minus))
8012                / (2.0 * h_grad);
8013            let rel_grad =
8014                (gradient[j] - fd_grad).abs() / gradient[j].abs().max(fd_grad.abs()).max(1e-12);
8015            assert!(
8016                rel_grad < 2e-4,
8017                "row primary grad[{j}] mismatch: analytic={}, fd={fd_grad}, rel={rel_grad}",
8018                gradient[j]
8019            );
8020
8021            for k in 0..LATENT_SURVIVAL_PRIMARY_DIM {
8022                let mut pp = primary.clone();
8023                pp[j] += h_hess;
8024                pp[k] += h_hess;
8025                let mut pm = primary.clone();
8026                pm[j] += h_hess;
8027                pm[k] -= h_hess;
8028                let mut mp = primary.clone();
8029                mp[j] -= h_hess;
8030                mp[k] += h_hess;
8031                let mut mm = primary.clone();
8032                mm[j] -= h_hess;
8033                mm[k] -= h_hess;
8034                let fd_neg_hess = -(latent_survival_row_loglik_from_primary(&quadctx, &row, &pp)
8035                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &pm)
8036                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &mp)
8037                    + latent_survival_row_loglik_from_primary(&quadctx, &row, &mm))
8038                    / (4.0 * h_hess * h_hess);
8039                let analytic = neg_hessian[[j, k]];
8040                let abs_err = (analytic - fd_neg_hess).abs();
8041                let rel = abs_err / analytic.abs().max(fd_neg_hess.abs()).max(1e-10);
8042                assert!(
8043                    abs_err < 2e-5 || rel < 2e-3,
8044                    "row primary neg_hess[{j},{k}] mismatch: analytic={analytic}, fd={fd_neg_hess}, abs_err={abs_err}, rel={rel}"
8045                );
8046            }
8047        }
8048    }
8049
8050    #[test]
8051    fn latent_survival_interval_row_primary_derivatives_match_fd() {
8052        // Interval-censored row jet `ℓ = log[S(L) − S(R)] − log S(entry)`. The
8053        // dynamic two-state numerator differentiates BOTH boundary masses
8054        // `M_L = exp(q_exit)` (left, `q_exit`) and `M_R = exp(q_right)` (right,
8055        // `q_right`) independently — channels that the static
8056        // `LatentSurvivalRowJet::interval_censored` (μ-only) never exercises. This
8057        // FD-verifies the gradient AND neg-Hessian of the interval contribution
8058        // w.r.t. ALL six primary coordinates (q_entry, q_exit/L, qdot_exit,
8059        // q_right/R, mu, log_sigma) on a WELL-POSED bracket where `S(L) − S(R)` is
8060        // comfortably positive (M_L = e^{−0.4} ≈ 0.67 well below M_R = e^{0.5} ≈
8061        // 1.65, so the survival-mass difference is large and the log-of-a-
8062        // difference curvature is well-conditioned).
8063        let quadctx = QuadratureContext::new();
8064        // Bracket masses: entry < L < R with comfortable gaps.
8065        let q_entry = -1.2_f64; // M_entry = e^{−1.2} ≈ 0.30
8066        let q_exit = -0.4_f64; // L: M_L = e^{−0.4} ≈ 0.67
8067        let q_right = 0.5_f64; // R: M_R = e^{0.5} ≈ 1.65 (> M_L)
8068        let mu = -0.15_f64;
8069        let log_sigma = 0.3_f64; // σ ≈ 1.35
8070        // Small, monotone unloaded masses (entry ≤ left ≤ right); qdot is inert
8071        // for the interval contribution.
8072        let row = LatentSurvivalRow::interval_censored(
8073            q_entry.exp(), // mass_entry (consistency only; jet reads q's)
8074            q_exit.exp(),  // mass_left
8075            q_right.exp(), // mass_right
8076            0.01,          // mass_unloaded_entry
8077            0.02,          // mass_unloaded_left
8078            0.05,          // mass_unloaded_right
8079        );
8080        assert!(matches!(
8081            row.event_type,
8082            LatentSurvivalEventType::IntervalCensored
8083        ));
8084
8085        // [q_entry, q_exit/L, qdot_exit, q_right/R, mu, log_sigma]. qdot_exit is
8086        // inert for interval rows (no hazard-derivative channel); the FD loop
8087        // confirms its gradient/Hessian entries are 0.
8088        let primary = array![q_entry, q_exit, 0.7, q_right, mu, log_sigma];
8089        let sigma = primary[LATENT_SURVIVAL_PRIMARY_LOG_SIGMA].exp();
8090        let h_grad = 1e-6;
8091        let h_hess = 2e-4;
8092
8093        let (_, gradient, neg_hessian) = latent_survival_row_primary_gradient_hessian(
8094            &quadctx,
8095            &row,
8096            LatentSurvivalPrimaryPoint {
8097                q_entry: primary[LATENT_SURVIVAL_PRIMARY_Q_ENTRY],
8098                q_exit: primary[LATENT_SURVIVAL_PRIMARY_Q_EXIT],
8099                qdot_exit: primary[LATENT_SURVIVAL_PRIMARY_QDOT_EXIT],
8100                q_right: primary[LATENT_SURVIVAL_PRIMARY_Q_RIGHT],
8101                mu: primary[LATENT_SURVIVAL_PRIMARY_MU],
8102                sigma,
8103            },
8104            true,
8105        )
8106        .expect("analytic interval row primary gradient/hessian");
8107
8108        // The interval contribution must be a positive survival-mass difference
8109        // at this bracket, so the value channel is finite.
8110        let value = latent_survival_row_loglik_from_primary(&quadctx, &row, &primary);
8111        assert!(
8112            value.is_finite(),
8113            "interval row log-likelihood must be finite on a well-posed bracket, got {value}"
8114        );
8115
8116        for j in 0..LATENT_SURVIVAL_PRIMARY_DIM {
8117            let mut plus = primary.clone();
8118            plus[j] += h_grad;
8119            let mut minus = primary.clone();
8120            minus[j] -= h_grad;
8121            let fd_grad = (latent_survival_row_loglik_from_primary(&quadctx, &row, &plus)
8122                - latent_survival_row_loglik_from_primary(&quadctx, &row, &minus))
8123                / (2.0 * h_grad);
8124            let rel_grad =
8125                (gradient[j] - fd_grad).abs() / gradient[j].abs().max(fd_grad.abs()).max(1e-12);
8126            assert!(
8127                rel_grad < 2e-4,
8128                "interval row primary grad[{j}] mismatch: analytic={}, fd={fd_grad}, rel={rel_grad}",
8129                gradient[j]
8130            );
8131
8132            for k in 0..LATENT_SURVIVAL_PRIMARY_DIM {
8133                let mut pp = primary.clone();
8134                pp[j] += h_hess;
8135                pp[k] += h_hess;
8136                let mut pm = primary.clone();
8137                pm[j] += h_hess;
8138                pm[k] -= h_hess;
8139                let mut mp = primary.clone();
8140                mp[j] -= h_hess;
8141                mp[k] += h_hess;
8142                let mut mm = primary.clone();
8143                mm[j] -= h_hess;
8144                mm[k] -= h_hess;
8145                let fd_neg_hess = -(latent_survival_row_loglik_from_primary(&quadctx, &row, &pp)
8146                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &pm)
8147                    - latent_survival_row_loglik_from_primary(&quadctx, &row, &mp)
8148                    + latent_survival_row_loglik_from_primary(&quadctx, &row, &mm))
8149                    / (4.0 * h_hess * h_hess);
8150                let analytic = neg_hessian[[j, k]];
8151                let abs_err = (analytic - fd_neg_hess).abs();
8152                let rel = abs_err / analytic.abs().max(fd_neg_hess.abs()).max(1e-10);
8153                assert!(
8154                    abs_err < 5e-5 || rel < 3e-3,
8155                    "interval row primary neg_hess[{j},{k}] mismatch: analytic={analytic}, fd={fd_neg_hess}, abs_err={abs_err}, rel={rel}"
8156                );
8157            }
8158        }
8159    }
8160
8161    type LatentFullPrimaryChannels = (f64, Array1<f64>, Array2<f64>, Array2<f64>, Array2<f64>);
8162
8163    fn latent_full_primary_channels_at(
8164        quadctx: &QuadratureContext,
8165        row: &LatentSurvivalRow,
8166        include_log_sigma: bool,
8167        use_multidir_reference: bool,
8168        point: LatentSurvivalPrimaryPoint,
8169    ) -> LatentFullPrimaryChannels {
8170        let direction_u = array![
8171            0.17,
8172            -0.11,
8173            0.09,
8174            0.13,
8175            -0.07,
8176            if include_log_sigma { 0.05 } else { 0.0 }
8177        ];
8178        let direction_v = array![
8179            -0.08,
8180            0.14,
8181            -0.06,
8182            0.04,
8183            0.12,
8184            if include_log_sigma { -0.09 } else { 0.0 }
8185        ];
8186        if use_multidir_reference {
8187            let (value, gradient, hessian) =
8188                latent_survival_row_primary_gradient_hessian_multidir_reference(
8189                    quadctx,
8190                    row,
8191                    point,
8192                    include_log_sigma,
8193                )
8194                .expect("pre-cutover MultiDirJet VGH reference");
8195            let third = latent_survival_row_primary_third_contracted_multidir_reference(
8196                quadctx,
8197                row,
8198                point,
8199                &direction_u,
8200                include_log_sigma,
8201            )
8202            .expect("pre-cutover MultiDirJet third reference");
8203            let fourth = latent_survival_row_primary_fourth_contracted_multidir_reference(
8204                quadctx,
8205                row,
8206                point,
8207                &direction_u,
8208                &direction_v,
8209                include_log_sigma,
8210            )
8211            .expect("pre-cutover MultiDirJet fourth reference");
8212            (value, gradient, hessian, third, fourth)
8213        } else {
8214            let (value, gradient, hessian) = latent_survival_row_primary_gradient_hessian(
8215                quadctx,
8216                row,
8217                point,
8218                include_log_sigma,
8219            )
8220            .expect("one-pass Order2 VGH");
8221            let third = latent_survival_row_primary_third_contracted(
8222                quadctx,
8223                row,
8224                point,
8225                &direction_u,
8226                include_log_sigma,
8227            )
8228            .expect("one-pass OneSeed third");
8229            let fourth = latent_survival_row_primary_fourth_contracted(
8230                quadctx,
8231                row,
8232                point,
8233                &direction_u,
8234                &direction_v,
8235                include_log_sigma,
8236            )
8237            .expect("one-pass TwoSeed fourth");
8238            (value, gradient, hessian, third, fourth)
8239        }
8240    }
8241
8242    fn latent_full_primary_channels(
8243        quadctx: &QuadratureContext,
8244        row: &LatentSurvivalRow,
8245        include_log_sigma: bool,
8246        use_multidir_reference: bool,
8247    ) -> LatentFullPrimaryChannels {
8248        latent_full_primary_channels_at(
8249            quadctx,
8250            row,
8251            include_log_sigma,
8252            use_multidir_reference,
8253            LatentSurvivalPrimaryPoint {
8254                q_entry: -1.2,
8255                q_exit: -0.4,
8256                qdot_exit: 0.73,
8257                q_right: 0.5,
8258                mu: -0.15,
8259                sigma: 0.3_f64.exp(),
8260            },
8261        )
8262    }
8263
8264    fn assert_latent_full_channels_close(
8265        label: &str,
8266        got: &LatentFullPrimaryChannels,
8267        reference: &LatentFullPrimaryChannels,
8268    ) {
8269        let mut max_abs = (got.0 - reference.0).abs();
8270        let mut max_rel = max_abs / got.0.abs().max(reference.0.abs()).max(1e-13);
8271        let mut max_abs_channel = "value".to_string();
8272        let mut max_abs_values = (got.0, reference.0);
8273        let mut max_rel_channel = "value".to_string();
8274        let mut max_rel_values = (got.0, reference.0);
8275        let mut record = |channel: String, left: f64, right: f64| {
8276            let absolute = (left - right).abs();
8277            let relative = absolute / left.abs().max(right.abs()).max(1e-13);
8278            if absolute > max_abs {
8279                max_abs = absolute;
8280                max_abs_channel = channel.clone();
8281                max_abs_values = (left, right);
8282            }
8283            if relative > max_rel {
8284                max_rel = relative;
8285                max_rel_channel = channel;
8286                max_rel_values = (left, right);
8287            }
8288        };
8289        for (a, (&left, &right)) in got.1.iter().zip(reference.1.iter()).enumerate() {
8290            record(format!("gradient[{a}]"), left, right);
8291        }
8292        for ((a, b), &left) in got.2.indexed_iter() {
8293            record(format!("hessian[{a},{b}]"), left, reference.2[[a, b]]);
8294        }
8295        for ((a, b), &left) in got.3.indexed_iter() {
8296            record(format!("third[{a},{b}]"), left, reference.3[[a, b]]);
8297        }
8298        for ((a, b), &left) in got.4.indexed_iter() {
8299            record(format!("fourth[{a},{b}]"), left, reference.4[[a, b]]);
8300        }
8301        assert!(
8302            max_abs <= 5e-11 || max_rel <= 5e-10,
8303            "{label}: one-pass channels differ from the pre-cutover MultiDirJet oracle: \
8304             max_abs={max_abs:e} at {max_abs_channel} (one-pass={}, oracle={}); \
8305             max_rel={max_rel:e} at {max_rel_channel} (one-pass={}, oracle={})",
8306            max_abs_values.0,
8307            max_abs_values.1,
8308            max_rel_values.0,
8309            max_rel_values.1,
8310        );
8311    }
8312
8313    /// Pins the one-pass scalar layouts to the complete pre-cutover output on
8314    /// every event branch and at both live primary dimensions.  This is stronger
8315    /// than an FD-only oracle: it covers value, gradient, negative Hessian,
8316    /// contracted third, and contracted fourth simultaneously.
8317    #[test]
8318    fn latent_survival_one_pass_matches_multidir_all_events_all_channels_932() {
8319        let quadctx = QuadratureContext::new();
8320        let rows = [
8321            (
8322                "right",
8323                LatentSurvivalRow::right_censored(0.3, 0.67, 0.01, 0.02),
8324            ),
8325            (
8326                "exact",
8327                LatentSurvivalRow::exact_event(0.3, 0.67, 0.01, 0.02, 0.73, 0.08),
8328            ),
8329            (
8330                "interval",
8331                LatentSurvivalRow::interval_censored(0.3, 0.67, 1.65, 0.01, 0.02, 0.05),
8332            ),
8333        ];
8334        for (event, row) in &rows {
8335            for include_log_sigma in [false, true] {
8336                let reference =
8337                    latent_full_primary_channels(&quadctx, row, include_log_sigma, true);
8338                let got = latent_full_primary_channels(&quadctx, row, include_log_sigma, false);
8339                let dimension = if include_log_sigma { 6 } else { 5 };
8340                assert_latent_full_channels_close(
8341                    &format!("event={event}, K={dimension}"),
8342                    &got,
8343                    &reference,
8344                );
8345            }
8346        }
8347    }
8348
8349    /// Exact-event tail audit: small/large loaded masses, displaced frailty
8350    /// locations, narrow/wide scales, and the learnable-scale axis all retain
8351    /// every channel of the signed-log MultiDirJet oracle.  These are analytic
8352    /// cross-oracles, not finite differences in a cancellation-prone tail.
8353    #[test]
8354    fn latent_survival_one_pass_exact_tails_match_multidir_all_channels_932() {
8355        let quadctx = QuadratureContext::new();
8356        let regimes: [(&str, f64, f64, f64, f64, f64, f64); 3] = [
8357            (
8358                "tiny-mass-left",
8359                -14.0_f64,
8360                -10.0_f64,
8361                0.31,
8362                0.0,
8363                -5.0,
8364                0.08,
8365            ),
8366            ("large-mass-right", 2.0, 6.0, 1.7, 0.0, 3.5, 0.45),
8367            ("wide-frailty", -3.0, 1.5, 0.62, 0.0, -1.8, 4.0),
8368        ];
8369        for (name, q_entry, q_exit, qdot, q_right, mu, sigma) in regimes {
8370            let row =
8371                LatentSurvivalRow::exact_event(q_entry.exp(), q_exit.exp(), 0.01, 0.04, qdot, 0.07);
8372            let point = LatentSurvivalPrimaryPoint {
8373                q_entry,
8374                q_exit,
8375                qdot_exit: qdot,
8376                q_right,
8377                mu,
8378                sigma,
8379            };
8380            for include_log_sigma in [false, true] {
8381                let reference =
8382                    latent_full_primary_channels_at(&quadctx, &row, include_log_sigma, true, point);
8383                let got = latent_full_primary_channels_at(
8384                    &quadctx,
8385                    &row,
8386                    include_log_sigma,
8387                    false,
8388                    point,
8389                );
8390                let dimension = if include_log_sigma { 6 } else { 5 };
8391                assert_latent_full_channels_close(
8392                    &format!("exact-tail={name}, K={dimension}"),
8393                    &got,
8394                    &reference,
8395                );
8396            }
8397        }
8398    }
8399
8400    #[test]
8401    fn latent_survival_derivative_support_stays_inline_932() {
8402        let base_terms = [
8403            LatentKernelPrimaryTerm {
8404                coeff: 0.08,
8405                q_exp: 0,
8406                qdot_power: 0,
8407                tau_exp: 0,
8408                k: 0,
8409            },
8410            LatentKernelPrimaryTerm {
8411                coeff: 1.0,
8412                q_exp: 1,
8413                qdot_power: 1,
8414                tau_exp: 0,
8415                k: 1,
8416            },
8417        ];
8418        let primary: [LatentKernelPrimaryDirection; LATENT_SURVIVAL_PRIMARY_DIM] =
8419            std::array::from_fn(|a| {
8420                latent_survival_map_exit_direction(
8421                    latent_survival_basis_direction(a),
8422                    LatentSurvivalEventType::ExactEvent,
8423                )
8424            });
8425        let u_coeff = [0.17, -0.11, 0.09, 0.13, -0.07, 0.05];
8426        let v_coeff = [-0.08, 0.14, -0.06, 0.04, 0.12, -0.09];
8427        let u = latent_kernel_direction_linear_combination(&primary, &u_coeff);
8428        let v = latent_kernel_direction_linear_combination(&primary, &v_coeff);
8429        let suffix_u = [u];
8430        let suffix_v = [v];
8431        let suffix_uv = [u, v];
8432        let suffixes: [&[LatentKernelPrimaryDirection]; 4] =
8433            [&[], &suffix_u, &suffix_v, &suffix_uv];
8434        let mut maximum_support = 0usize;
8435        for suffix in suffixes {
8436            for a in 0..LATENT_SURVIVAL_PRIMARY_DIM {
8437                for b in a..LATENT_SURVIVAL_PRIMARY_DIM {
8438                    let terms = latent_kernel_term_sequence_inline(
8439                        &base_terms,
8440                        &[primary[a], primary[b]],
8441                        suffix,
8442                    );
8443                    assert!(!terms.spilled(), "derivative term support spilled to heap");
8444                    maximum_support = maximum_support.max(terms.len());
8445                }
8446            }
8447        }
8448        eprintln!(
8449            "LATENT-ONE-PASS-932 derivative-support max_terms={maximum_support} inline_capacity={LATENT_TERM_INLINE_CAPACITY} heap_allocations=0"
8450        );
8451        assert!(maximum_support <= LATENT_TERM_INLINE_CAPACITY);
8452    }
8453
8454    fn best_elapsed_seconds(mut run: impl FnMut(), iterations: usize, samples: usize) -> f64 {
8455        let mut best = f64::INFINITY;
8456        for _ in 0..samples {
8457            let started = std::time::Instant::now();
8458            for _ in 0..iterations {
8459                run();
8460            }
8461            best = best.min(started.elapsed().as_secs_f64());
8462        }
8463        best
8464    }
8465
8466    fn measured_channel_ratio(
8467        mut reference: impl FnMut(),
8468        mut one_pass: impl FnMut(),
8469        iterations: usize,
8470        samples: usize,
8471    ) -> (f64, f64, f64) {
8472        reference();
8473        one_pass();
8474        let reference_seconds = best_elapsed_seconds(&mut reference, iterations, samples);
8475        let one_pass_seconds = best_elapsed_seconds(&mut one_pass, iterations, samples);
8476        (
8477            reference_seconds * 1e6 / iterations as f64,
8478            one_pass_seconds * 1e6 / iterations as f64,
8479            one_pass_seconds / reference_seconds,
8480        )
8481    }
8482
8483    /// Full-output pre-cutover benchmark: the baseline includes all 21/28 VGH
8484    /// row sweeps and all 15/21 third/fourth pair sweeps; the candidate returns
8485    /// the identical five-channel payload via one Order2, one OneSeed, and one
8486    /// TwoSeed row evaluation.  Run with `--release -- --nocapture` for the
8487    /// reported ratios; the debug configuration uses one iteration to keep the
8488    /// ordinary suite bounded while still enforcing the direction of the win.
8489    #[test]
8490    fn measure_latent_survival_one_pass_full_output_k5_k6_932() {
8491        let quadctx = QuadratureContext::new();
8492        let row = LatentSurvivalRow::exact_event(0.3, 0.67, 0.01, 0.02, 0.73, 0.08);
8493        let q_entry = -1.2;
8494        let q_exit = -0.4;
8495        let qdot_exit = 0.73;
8496        let q_right = 0.5;
8497        let mu = -0.15;
8498        let sigma = 0.3_f64.exp();
8499        let point = LatentSurvivalPrimaryPoint {
8500            q_entry,
8501            q_exit,
8502            qdot_exit,
8503            q_right,
8504            mu,
8505            sigma,
8506        };
8507        let iterations = if cfg!(debug_assertions) { 1 } else { 5 };
8508        let samples = if cfg!(debug_assertions) { 1 } else { 3 };
8509
8510        for include_log_sigma in [false, true] {
8511            let direction_u = array![
8512                0.17,
8513                -0.11,
8514                0.09,
8515                0.13,
8516                -0.07,
8517                if include_log_sigma { 0.05 } else { 0.0 }
8518            ];
8519            let direction_v = array![
8520                -0.08,
8521                0.14,
8522                -0.06,
8523                0.04,
8524                0.12,
8525                if include_log_sigma { -0.09 } else { 0.0 }
8526            ];
8527            let dimension = if include_log_sigma { 6 } else { 5 };
8528            let (vgh_reference_us, vgh_one_pass_us, vgh_ratio) = measured_channel_ratio(
8529                || {
8530                    std::hint::black_box(
8531                        latent_survival_row_primary_gradient_hessian_multidir_reference(
8532                            std::hint::black_box(&quadctx),
8533                            std::hint::black_box(&row),
8534                            point,
8535                            include_log_sigma,
8536                        )
8537                        .expect("prechange VGH benchmark"),
8538                    );
8539                },
8540                || {
8541                    std::hint::black_box(
8542                        latent_survival_row_primary_gradient_hessian(
8543                            std::hint::black_box(&quadctx),
8544                            std::hint::black_box(&row),
8545                            point,
8546                            include_log_sigma,
8547                        )
8548                        .expect("one-pass VGH benchmark"),
8549                    );
8550                },
8551                iterations,
8552                samples,
8553            );
8554            let (third_reference_us, third_one_pass_us, third_ratio) = measured_channel_ratio(
8555                || {
8556                    std::hint::black_box(
8557                        latent_survival_row_primary_third_contracted_multidir_reference(
8558                            std::hint::black_box(&quadctx),
8559                            std::hint::black_box(&row),
8560                            point,
8561                            std::hint::black_box(&direction_u),
8562                            include_log_sigma,
8563                        )
8564                        .expect("prechange third benchmark"),
8565                    );
8566                },
8567                || {
8568                    std::hint::black_box(
8569                        latent_survival_row_primary_third_contracted(
8570                            std::hint::black_box(&quadctx),
8571                            std::hint::black_box(&row),
8572                            point,
8573                            std::hint::black_box(&direction_u),
8574                            include_log_sigma,
8575                        )
8576                        .expect("one-pass third benchmark"),
8577                    );
8578                },
8579                iterations,
8580                samples,
8581            );
8582            let (fourth_reference_us, fourth_one_pass_us, fourth_ratio) = measured_channel_ratio(
8583                || {
8584                    std::hint::black_box(
8585                        latent_survival_row_primary_fourth_contracted_multidir_reference(
8586                            std::hint::black_box(&quadctx),
8587                            std::hint::black_box(&row),
8588                            point,
8589                            std::hint::black_box(&direction_u),
8590                            std::hint::black_box(&direction_v),
8591                            include_log_sigma,
8592                        )
8593                        .expect("prechange fourth benchmark"),
8594                    );
8595                },
8596                || {
8597                    std::hint::black_box(
8598                        latent_survival_row_primary_fourth_contracted(
8599                            std::hint::black_box(&quadctx),
8600                            std::hint::black_box(&row),
8601                            point,
8602                            std::hint::black_box(&direction_u),
8603                            std::hint::black_box(&direction_v),
8604                            include_log_sigma,
8605                        )
8606                        .expect("one-pass fourth benchmark"),
8607                    );
8608                },
8609                iterations,
8610                samples,
8611            );
8612            let (full_reference_us, full_one_pass_us, full_ratio) = measured_channel_ratio(
8613                || {
8614                    std::hint::black_box(latent_full_primary_channels(
8615                        std::hint::black_box(&quadctx),
8616                        std::hint::black_box(&row),
8617                        include_log_sigma,
8618                        true,
8619                    ));
8620                },
8621                || {
8622                    std::hint::black_box(latent_full_primary_channels(
8623                        std::hint::black_box(&quadctx),
8624                        std::hint::black_box(&row),
8625                        include_log_sigma,
8626                        false,
8627                    ));
8628                },
8629                iterations,
8630                samples,
8631            );
8632            let (combined_reference_us, combined_one_pass_us, combined_ratio) =
8633                measured_channel_ratio(
8634                    || {
8635                        std::hint::black_box(
8636                            latent_survival_row_primary_gradient_hessian_multidir_reference(
8637                                &quadctx,
8638                                &row,
8639                                point,
8640                                include_log_sigma,
8641                            )
8642                            .expect("prechange combined VGH"),
8643                        );
8644                        for direction in [&direction_u, &direction_v] {
8645                            std::hint::black_box(
8646                                latent_survival_row_primary_third_contracted_multidir_reference(
8647                                    &quadctx,
8648                                    &row,
8649                                    point,
8650                                    direction,
8651                                    include_log_sigma,
8652                                )
8653                                .expect("prechange combined third"),
8654                            );
8655                        }
8656                        std::hint::black_box(
8657                            latent_survival_row_primary_fourth_contracted_multidir_reference(
8658                                &quadctx,
8659                                &row,
8660                                point,
8661                                &direction_u,
8662                                &direction_v,
8663                                include_log_sigma,
8664                            )
8665                            .expect("prechange combined fourth"),
8666                        );
8667                    },
8668                    || {
8669                        if include_log_sigma {
8670                            let backend = LatentTwoSeedBackend {
8671                                direction_u: std::array::from_fn(|a| direction_u[a]),
8672                                direction_v: std::array::from_fn(|a| direction_v[a]),
8673                            };
8674                            std::hint::black_box(
8675                                latent_survival_row_primary_jet::<LATENT_SURVIVAL_PRIMARY_DIM, _>(
8676                                    &backend, &quadctx, &row, point,
8677                                )
8678                                .expect("combined K6 TwoSeed"),
8679                            );
8680                        } else {
8681                            std::hint::black_box(
8682                                latent_survival_row_primary_two_seed_fixed_sigma(
8683                                    &quadctx,
8684                                    &row,
8685                                    point,
8686                                    &direction_u,
8687                                    &direction_v,
8688                                )
8689                                .expect("combined K5 TwoSeed"),
8690                            );
8691                        }
8692                    },
8693                    iterations,
8694                    samples,
8695                );
8696            let pair_count = dimension * (dimension + 1) / 2;
8697            let order2_width = 1 + dimension + pair_count;
8698            let vgh_reference_bundle_allocs = 2 * (1 + dimension + pair_count);
8699            let contracted_reference_bundle_allocs = 2 * pair_count;
8700            eprintln!(
8701                "LATENT-ONE-PASS-OPS-932 K={dimension} signed-term-reductions/state VGH {}->{order2_width} T3 {}->{} T4 {}->{}",
8702                1 + 2 * dimension + 4 * pair_count,
8703                8 * pair_count,
8704                2 * order2_width,
8705                16 * pair_count,
8706                4 * order2_width,
8707            );
8708            eprintln!(
8709                "LATENT-ONE-PASS-932 K={dimension} VGH prechange={vgh_reference_us:.3}us one-pass={vgh_one_pass_us:.3}us ratio={vgh_ratio:.4} speedup={:.2}x bundle-Vec-allocs={vgh_reference_bundle_allocs}->2; T3 prechange={third_reference_us:.3}us one-pass={third_one_pass_us:.3}us ratio={third_ratio:.4} speedup={:.2}x bundle-Vec-allocs={contracted_reference_bundle_allocs}->2; T4 prechange={fourth_reference_us:.3}us one-pass={fourth_one_pass_us:.3}us ratio={fourth_ratio:.4} speedup={:.2}x bundle-Vec-allocs={contracted_reference_bundle_allocs}->2; FULL-3PASS prechange={full_reference_us:.3}us one-pass={full_one_pass_us:.3}us ratio={full_ratio:.4} speedup={:.2}x bundle-Vec-allocs={}->6; FULL-COMBINED prechange={combined_reference_us:.3}us one-pass={combined_one_pass_us:.3}us ratio={combined_ratio:.4} speedup={:.2}x bundle-Vec-allocs={}->2; derivative-plan-heap-allocs=0 three-pass-output-ndarray-allocs=4->4 combined-output-ndarray-allocs=5->0",
8710                1.0 / vgh_ratio,
8711                1.0 / third_ratio,
8712                1.0 / fourth_ratio,
8713                1.0 / full_ratio,
8714                vgh_reference_bundle_allocs + 2 * contracted_reference_bundle_allocs,
8715                1.0 / combined_ratio,
8716                vgh_reference_bundle_allocs + 3 * contracted_reference_bundle_allocs,
8717            );
8718            for (channel, ratio) in [
8719                ("VGH", vgh_ratio),
8720                ("T3", third_ratio),
8721                ("T4", fourth_ratio),
8722                ("full-3pass", full_ratio),
8723                ("full-combined", combined_ratio),
8724            ] {
8725                assert!(
8726                    ratio < 1.0,
8727                    "K={dimension} one-pass {channel} must beat the exact pre-cutover path: ratio={ratio}"
8728                );
8729            }
8730        }
8731    }
8732}