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