use super::*;
pub struct PreparedSurvivalTimeStack {
pub eta_offset_entry: Array1<f64>,
pub eta_offset_exit: Array1<f64>,
pub derivative_offset_exit: Array1<f64>,
pub unloaded_mass_entry: Array1<f64>,
pub unloaded_mass_exit: Array1<f64>,
pub unloaded_hazard_exit: Array1<f64>,
pub time_design_entry: crate::matrix::DesignMatrix,
pub time_design_exit: crate::matrix::DesignMatrix,
pub time_design_derivative_exit: crate::matrix::DesignMatrix,
pub time_penalties: Vec<Array2<f64>>,
pub time_nullspace_dims: Vec<usize>,
pub timewiggle_build: Option<crate::families::survival::construction::SurvivalTimeWiggleBuild>,
pub timewiggle_block: Option<TimeWiggleBlockInput>,
}
pub fn prepare_survival_time_stack(
age_entry: &Array1<f64>,
age_exit: &Array1<f64>,
baseline_cfg: &crate::families::survival::construction::SurvivalBaselineConfig,
likelihood_mode: SurvivalLikelihoodMode,
inverse_link: Option<&InverseLink>,
time_anchor: f64,
derivative_guard: f64,
time_build: &crate::families::survival::construction::SurvivalTimeBuildOutput,
effective_timewiggle: Option<&LinkWiggleFormulaSpec>,
latent_loading: Option<crate::families::survival::lognormal_kernel::HazardLoading>,
) -> Result<PreparedSurvivalTimeStack, String> {
let (
mut eta_offset_entry,
mut eta_offset_exit,
mut derivative_offset_exit,
unloaded_mass_entry,
unloaded_mass_exit,
unloaded_hazard_exit,
) = if let Some(loading) = latent_loading {
let offsets =
build_latent_survival_baseline_offsets(age_entry, age_exit, baseline_cfg, loading)?;
(
offsets.loaded_eta_entry,
offsets.loaded_eta_exit,
offsets.loaded_derivative_exit,
offsets.unloaded_mass_entry,
offsets.unloaded_mass_exit,
offsets.unloaded_hazard_exit,
)
} else {
let conditioning_cfg;
let offset_cfg = if likelihood_mode == SurvivalLikelihoodMode::MarginalSlope
&& baseline_cfg.target == SurvivalBaselineTarget::Linear
{
let scale =
crate::families::survival::construction::positive_survival_time_seed(age_exit);
conditioning_cfg = crate::families::survival::construction::SurvivalBaselineConfig {
target: SurvivalBaselineTarget::Weibull,
scale: Some(scale),
shape: Some(1.0),
rate: None,
makeham: None,
};
&conditioning_cfg
} else {
baseline_cfg
};
let (eta_offset_entry, eta_offset_exit, derivative_offset_exit) =
build_survival_time_offsets_for_likelihood(
age_entry,
age_exit,
offset_cfg,
likelihood_mode,
inverse_link,
)?;
let n = age_entry.len();
(
eta_offset_entry,
eta_offset_exit,
derivative_offset_exit,
Array1::zeros(n),
Array1::zeros(n),
Array1::zeros(n),
)
};
add_survival_time_derivative_guard_offset(
age_entry,
age_exit,
time_anchor,
derivative_guard,
&mut eta_offset_entry,
&mut eta_offset_exit,
&mut derivative_offset_exit,
)?;
let timewiggle_build = if let Some(cfg) = effective_timewiggle {
Some(build_survival_timewiggle_from_baseline(
&eta_offset_entry,
&eta_offset_exit,
&derivative_offset_exit,
cfg,
)?)
} else {
None
};
let mut time_design_entry = time_build.x_entry_time.clone();
let mut time_design_exit = time_build.x_exit_time.clone();
let mut time_design_derivative_exit = time_build.x_derivative_time.clone();
let mut time_penalties = time_build.penalties.clone();
let mut time_nullspace_dims = time_build.nullspace_dims.clone();
let mut timewiggle_block = None;
if let Some(wiggle) = timewiggle_build.as_ref() {
let p_base = time_design_exit.ncols();
append_zero_tail_columns(
&mut time_design_entry,
&mut time_design_exit,
&mut time_design_derivative_exit,
wiggle.ncols,
);
for (idx, penalty) in wiggle.penalties.iter().enumerate() {
let mut embedded = Array2::<f64>::zeros((p_base + wiggle.ncols, p_base + wiggle.ncols));
embedded
.slice_mut(s![
p_base..p_base + wiggle.ncols,
p_base..p_base + wiggle.ncols
])
.assign(penalty);
time_penalties.push(embedded);
time_nullspace_dims.push(wiggle.nullspace_dims.get(idx).copied().unwrap_or(0));
}
timewiggle_block = Some(TimeWiggleBlockInput {
knots: wiggle.knots.clone(),
degree: wiggle.degree,
ncols: wiggle.ncols,
});
}
Ok(PreparedSurvivalTimeStack {
eta_offset_entry,
eta_offset_exit,
derivative_offset_exit,
unloaded_mass_entry,
unloaded_mass_exit,
unloaded_hazard_exit,
time_design_entry,
time_design_exit,
time_design_derivative_exit,
time_penalties,
time_nullspace_dims,
timewiggle_build,
timewiggle_block,
})
}