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