diff --git a/src/turnbull.rs b/src/turnbull.rs index 9d0b49d..fd67669 100644 --- a/src/turnbull.rs +++ b/src/turnbull.rs @@ -27,6 +27,7 @@ use prettytable::{Table, format, row}; use rayon::prelude::*; use serde::{Serialize, Deserialize}; +use crate::pava::monotonic_regression_pava; use crate::term::UnconditionalTermLike; #[derive(Args)] @@ -205,7 +206,7 @@ pub fn fit_turnbull(data_times: MatrixXx2, progress_bar: ProgressBar, max_i progress_bar.set_style(ProgressStyle::with_template("[{elapsed_precise}] {bar:40} {msg}").unwrap()); progress_bar.set_length(u64::MAX); progress_bar.reset(); - progress_bar.println("Running iterative algorithm to fit Turnbull estimator"); + progress_bar.println("Running EM-ICM algorithm to fit Turnbull estimator"); let (s, ll) = fit_turnbull_estimator(&mut data, progress_bar.clone(), max_iterations, ll_tolerance, s); @@ -285,24 +286,38 @@ fn get_turnbull_intervals(data_times: &MatrixXx2) -> Vec<(f64, f64)> { return intervals; } -fn fit_turnbull_estimator(data: &mut TurnbullData, progress_bar: ProgressBar, max_iterations: u32, ll_tolerance: f64, mut s: Vec) -> (Vec, f64) { - // Get likelihood for each observation (denominator of μ_ij) - let mut likelihood_obs = get_likelihood_obs(data, &s); +fn fit_turnbull_estimator(data: &mut TurnbullData, progress_bar: ProgressBar, max_iterations: u32, ll_tolerance: f64, mut p: Vec) -> (Vec, f64) { + // Pre-compute S, the survival probability at the start of each interval + let mut s = p_to_s(&p); + + // Get likelihood for each observation + let likelihood_obs = get_likelihood_obs(data, &s); let mut ll_model: f64 = likelihood_obs.iter().map(|l| l.ln()).sum(); let mut iteration = 1; loop { - // Compute π_j to update s - let pi = compute_pi(data, &s, likelihood_obs); + // ------- + // EM step - let likelihood_obs_new = get_likelihood_obs(data, &pi); - let ll_model_new = likelihood_obs_new.iter().map(|l| l.ln()).sum(); + let p_after_em = do_em_step(data, &p, &s); + let s_after_em = p_to_s(&p_after_em); + + let likelihood_obs_after_em = get_likelihood_obs(data, &s_after_em); + let ll_model_after_em: f64 = likelihood_obs_after_em.iter().map(|l| l.ln()).sum(); + + p = p_after_em; + s = s_after_em; + + // -------- + // ICM step + + let (p_new, s_new, ll_model_new) = do_icm_step(data, &p, &s, ll_model_after_em); let ll_change = ll_model_new - ll_model; let converged = ll_change <= ll_tolerance; - s = pi; - likelihood_obs = likelihood_obs_new; + p = p_new; + s = s_new; ll_model = ll_model_new; // Estimate progress bar according to either the order of magnitude of the ll_change relative to tolerance, or iteration/max_iterations @@ -324,48 +339,150 @@ fn fit_turnbull_estimator(data: &mut TurnbullData, progress_bar: ProgressBar, ma } } - return (s, ll_model); + return (p, ll_model); +} + +fn p_to_s(p: &Vec) -> Vec { + let mut s = Vec::with_capacity(p.len() + 1); // Survival probabilities + let mut survival = 1.0; + s.push(1.0); + for p_j in p.iter() { + survival -= p_j; + s.push(survival); + } + return s; +} + +fn s_to_lambda(s: &Vec) -> Vec { + // S = 1 means Λ = -inf and S = 0 means Λ = inf so skip these + let mut lambda = Vec::with_capacity(s.len() - 2); // Cumulative hazard + for s_j in &s[1..(s.len() - 1)] { + lambda.push((-s_j.ln()).ln()); + } + return lambda; } fn get_likelihood_obs(data: &TurnbullData, s: &Vec) -> Vec { return data.data_time_interval_indexes .par_iter() - .map(|(idx_left, idx_right)| s[*idx_left..(*idx_right + 1)].iter().sum()) - .collect(); + .map(|(idx_left, idx_right)| s[*idx_left] - s[*idx_right + 1]) + .collect(); // TODO: Return iterator directly } -fn compute_pi(data: &TurnbullData, s: &Vec, likelihood_obs: Vec) -> Vec { - /* - let mut pi: Vec = vec![0.0; data.num_intervals()]; - for ((idx_left, idx_right), likelihood_obs_i) in data.data_time_interval_indexes.iter().zip(likelihood_obs.iter()) { - for j in *idx_left..(*idx_right + 1) { - pi[j] += s[j] / likelihood_obs_i / data.num_obs() as f64; +fn do_em_step(data: &TurnbullData, p: &Vec, s: &Vec) -> Vec { + // Compute contributions to m + let mut m_contrib = vec![0.0; data.num_intervals()]; + for (idx_left, idx_right) in data.data_time_interval_indexes.iter() { + let contrib = 1.0 / (s[*idx_left] - s[*idx_right + 1]); + + // Adds to m for the first interval in the observation + m_contrib[*idx_left] += contrib; + + // Subtracts from m for the first interval beyond the observation + if *idx_right + 1 < data.num_intervals() { + m_contrib[*idx_right + 1] -= contrib; } } - */ - let pi = data.data_time_interval_indexes.par_iter().zip(likelihood_obs.par_iter()) - .fold_with( - // Compute the contributions to pi[j] for each observation and sum them in parallel using fold_with - vec![0.0; data.num_intervals()], - |mut acc, ((idx_left, idx_right), likelihood_obs_i)| { - // Contributions to pi[j] for the i-th observation - for j in *idx_left..(*idx_right + 1) { - acc[j] += s[j] / likelihood_obs_i / data.num_obs() as f64; - } - acc - } - ) - .reduce( - // Reduce all the sub-sums from fold_with into the total sum - || vec![0.0; data.num_intervals()], - |mut acc, subsum| { - acc.iter_mut().zip(subsum.iter()).for_each(|(x, y)| *x += y); - acc - } - ); + // Compute m + let mut m = Vec::with_capacity(data.num_intervals()); + let mut m_last = 0.0; + for m_contrib_j in m_contrib { + let m_next = m_last + m_contrib_j / (data.num_obs() as f64); + m.push(m_next); + m_last = m_next; + } - return pi; + // Update p + // p := p * m + let p_new = p.par_iter().zip(m.into_par_iter()).map(|(p_j, m_j)| p_j * m_j).collect(); + + return p_new; +} + +fn do_icm_step(data: &TurnbullData, _p: &Vec, s: &Vec, ll_model: f64) -> (Vec, Vec, f64) { + // Compute Λ, the cumulative hazard + // Since Λ = -inf when survival is 1, and Λ = inf when survival is 0, these are omitted + // The entry at lambda[j] corresponds to the survival immediately before time point j + 1 + let lambda = s_to_lambda(&s); + + // Compute gradient and diagonal of Hessian + let mut gradient = vec![0.0; data.num_intervals() - 1]; + let mut hessdiag = vec![0.0; data.num_intervals() - 1]; + for (idx_left, idx_right) in data.data_time_interval_indexes.iter() { + let denom = s[*idx_left] - s[*idx_right + 1]; + + // Add to gradient[j] when j + 1 == idx_right + 1 + // Add to hessdiag[j] when j + 1 == idx_right + 1 + if *idx_right < gradient.len() { + let j = *idx_right; + gradient[j] += (-lambda[j].exp() + lambda[j]).exp() / denom; + + let a = ((lambda[j] - lambda[j].exp()).exp() * (1.0 - lambda[j].exp())) / denom; + let b = (2.0 * lambda[j] - 2.0 * lambda[j].exp()).exp() / denom.powi(2); + hessdiag[j] += a - b; + } + + // Subtract from gradient[j] when j + 1 == idx_left + // Add to hessdiag[j] when j + 1 == idx_left + if *idx_left > 0 { + let j = *idx_left - 1; + gradient[j] -= (-lambda[j].exp() + lambda[j]).exp() / denom; + + let a = ((lambda[j] - lambda[j].exp()).exp() * (1.0 - lambda[j].exp())) / denom; + let b = (2.0 * lambda[j] - 2.0 * lambda[j].exp()).exp() / denom.powi(2); + hessdiag[j] += -a - b; + } + } + + // Description in Anderson-Bergman (2017) is slightly misleading + // Since we are maximising the likelihood, the second derivatives will be negative + // And we will move in the direction of the gradient + // So there are a few more negative signs here than suggested + + let weights = -DVector::from_vec(hessdiag.clone()) / 2.0; + let gradient_over_hessdiag = DVector::from_vec(gradient.par_iter().zip(hessdiag.par_iter()).map(|(g, h)| g / h).collect()); + + let mut s_new; + let mut p_new; + let mut ll_model_new: f64; + + // Take as large a step as possible while the log-likelihood increases + let mut step_size_exponent: i32 = 0; + loop { + let step_size = 0.5_f64.powi(step_size_exponent); + let lambda_target = -gradient_over_hessdiag.clone() * step_size + DVector::from_vec(lambda.clone()); + + let lambda_new = monotonic_regression_pava(lambda_target, weights.clone()); + + // Convert Λ to S to p + s_new = Vec::with_capacity(data.num_intervals() + 1); + p_new = Vec::with_capacity(data.num_intervals()); + + let mut survival = 1.0; + s_new.push(1.0); + for lambda_j in lambda_new.iter() { + let next_survival = (-lambda_j.exp()).exp(); + s_new.push(next_survival); + p_new.push(survival - next_survival); + survival = next_survival; + } + s_new.push(0.0); + p_new.push(survival); + + let likelihood_obs_new = get_likelihood_obs(data, &s_new); + ll_model_new = likelihood_obs_new.iter().map(|l| l.ln()).sum(); + + if ll_model_new > ll_model { + return (p_new, s_new, ll_model_new); + } + + step_size_exponent += 1; + + if step_size_exponent > 10 { + panic!("ICM fails to increase log-likelihood"); + } + } } fn compute_hessian(data: &TurnbullData, s: &Vec) -> DMatrix {