From 37c904bf34af885e43e7170ec177367445ad4ec1 Mon Sep 17 00:00:00 2001 From: RunasSudo Date: Sat, 28 Oct 2023 23:16:14 +1100 Subject: [PATCH] turnbull: Refactor for profiling --- src/turnbull.rs | 223 +++++++++++++++++++++++++----------------------- 1 file changed, 117 insertions(+), 106 deletions(-) diff --git a/src/turnbull.rs b/src/turnbull.rs index d5b93f1..6760727 100644 --- a/src/turnbull.rs +++ b/src/turnbull.rs @@ -291,7 +291,7 @@ fn fit_turnbull_estimator(data: &mut TurnbullData, progress_bar: ProgressBar, ma let mut s = p_to_s(&p); // Get likelihood for each observation - let mut likelihood_obs = get_likelihood_obs(data, &s); + 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; @@ -299,118 +299,19 @@ fn fit_turnbull_estimator(data: &mut TurnbullData, progress_bar: ProgressBar, ma // ------- // EM step - // Update p - let mut p_new = Vec::with_capacity(data.num_intervals()); - for j in 0..data.num_intervals() { - let tmp: f64 = data.data_time_interval_indexes.iter() - .filter(|(idx_left, idx_right)| j >= *idx_left && j <= *idx_right) - .map(|(idx_left, idx_right)| 1.0 / (s[*idx_left] - s[*idx_right + 1])) - .sum(); - - p_new.push(p[j] * tmp / (data.num_obs() as f64)); - } + let p_after_em = do_em_step(data, &p, &s); + let s_after_em = p_to_s(&p_after_em); - let mut s_new = p_to_s(&p_new); - let likelihood_obs_after_em = get_likelihood_obs(data, &s_new); + 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_new; - s = s_new; + p = p_after_em; + s = s_after_em; // -------- // ICM step - // Compute Λ, the cumulative hazard - let lambda = s_to_lambda(&s); - - // Compute gradient - let mut gradient = DVector::zeros(data.num_intervals() - 1); - for j in 0..(data.num_intervals() - 1) { - let sum_right: f64 = data.data_time_interval_indexes.iter() - .filter(|(idx_left, idx_right)| j + 1 == *idx_right + 1) - .map(|(idx_left, idx_right)| (-lambda[j].exp() + lambda[j]).exp() / (s[*idx_left] - s[*idx_right + 1])) - .sum(); - - let sum_left: f64 = data.data_time_interval_indexes.iter() - .filter(|(idx_left, idx_right)| j + 1 == *idx_left) - .map(|(idx_left, idx_right)| (-lambda[j].exp() + lambda[j]).exp() / (s[*idx_left] - s[*idx_right + 1])) - .sum(); - - gradient[j] = sum_right - sum_left; - } - - // Compute diagonal of Hessian - let mut hessdiag = DVector::zeros(data.num_intervals() - 1); - for j in 0..(data.num_intervals() - 1) { - let sum_left: f64 = data.data_time_interval_indexes.iter() - .filter(|(idx_left, idx_right)| j + 1 == *idx_left) - .map(|(idx_left, idx_right)| { - let denom = s[*idx_left] - s[*idx_right + 1]; - 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); - -a - b - }) - .sum(); - - let sum_right: f64 = data.data_time_interval_indexes.iter() - .filter(|(idx_left, idx_right)| j + 1 == *idx_right + 1) - .map(|(idx_left, idx_right)| { - let denom = s[*idx_left] - s[*idx_right + 1]; - 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); - a - b - }) - .sum(); - - hessdiag[j] = sum_left + sum_right; - } - - // 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 = -hessdiag.clone() / 2.0; - - 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.component_div(&hessdiag) * 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_after_em { - break; - } - - step_size_exponent += 1; - - if step_size_exponent > 10 { - panic!("ICM fails to increase log-likelihood"); - } - } + 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; @@ -468,6 +369,116 @@ fn get_likelihood_obs(data: &TurnbullData, s: &Vec) -> Vec { .collect(); // TODO: Return iterator directly } +fn do_em_step(data: &TurnbullData, p: &Vec, s: &Vec) -> Vec { + // Update p + let mut p_new = Vec::with_capacity(data.num_intervals()); + for j in 0..data.num_intervals() { + let tmp: f64 = data.data_time_interval_indexes.iter() + .filter(|(idx_left, idx_right)| j >= *idx_left && j <= *idx_right) + .map(|(idx_left, idx_right)| 1.0 / (s[*idx_left] - s[*idx_right + 1])) + .sum(); + + p_new.push(p[j] * tmp / (data.num_obs() as f64)); + } + + return p_new; +} + +fn do_icm_step(data: &TurnbullData, _p: &Vec, s: &Vec, ll_model: f64) -> (Vec, Vec, f64) { + // Compute Λ, the cumulative hazard + let lambda = s_to_lambda(&s); + + // Compute gradient + let mut gradient = DVector::zeros(data.num_intervals() - 1); + for j in 0..(data.num_intervals() - 1) { + let sum_right: f64 = data.data_time_interval_indexes.iter() + .filter(|(idx_left, idx_right)| j + 1 == *idx_right + 1) + .map(|(idx_left, idx_right)| (-lambda[j].exp() + lambda[j]).exp() / (s[*idx_left] - s[*idx_right + 1])) + .sum(); + + let sum_left: f64 = data.data_time_interval_indexes.iter() + .filter(|(idx_left, idx_right)| j + 1 == *idx_left) + .map(|(idx_left, idx_right)| (-lambda[j].exp() + lambda[j]).exp() / (s[*idx_left] - s[*idx_right + 1])) + .sum(); + + gradient[j] = sum_right - sum_left; + } + + // Compute diagonal of Hessian + let mut hessdiag = DVector::zeros(data.num_intervals() - 1); + for j in 0..(data.num_intervals() - 1) { + let sum_left: f64 = data.data_time_interval_indexes.iter() + .filter(|(idx_left, idx_right)| j + 1 == *idx_left) + .map(|(idx_left, idx_right)| { + let denom = s[*idx_left] - s[*idx_right + 1]; + 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); + -a - b + }) + .sum(); + + let sum_right: f64 = data.data_time_interval_indexes.iter() + .filter(|(idx_left, idx_right)| j + 1 == *idx_right + 1) + .map(|(idx_left, idx_right)| { + let denom = s[*idx_left] - s[*idx_right + 1]; + 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); + a - b + }) + .sum(); + + hessdiag[j] = sum_left + sum_right; + } + + // 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 = -hessdiag.clone() / 2.0; + + 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.component_div(&hessdiag) * 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 { let mut hessian: DMatrix = DMatrix::zeros(data.num_intervals() - 1, data.num_intervals() - 1);