RunasSudo/hpstat
1
Fork 0
Code Issues Pull Requests Releases Activity

turnbull: Refactor for profiling

Browse Source
This commit is contained in:
RunasSudo 2023-10-28 23:16:14 +11:00
parent 81b0b3f9b5
commit 37c904bf34
Signed by: RunasSudo
GPG Key ID: 7234E476BF21C61A
1 changed files with 117 additions and 106 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

223
src/turnbull.rs
View File

@ -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<f64>) -> Vec<f64> {
.collect(); // TODO: Return iterator directly
}
fn do_em_step(data: &TurnbullData, p: &Vec<f64>, s: &Vec<f64>) -> Vec<f64> {
// 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<f64>, s: &Vec<f64>, ll_model: f64) -> (Vec<f64>, Vec<f64>, 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<f64>) -> DMatrix<f64> {
let mut hessian: DMatrix<f64> = DMatrix::zeros(data.num_intervals() - 1, data.num_intervals() - 1);