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