use pyo3::prelude::*;
#[derive(Debug, Clone)]
#[pyclass]
pub struct NelsonAalenResult {
#[pyo3(get)]
pub time: Vec<f64>,
#[pyo3(get)]
pub cumulative_hazard: Vec<f64>,
#[pyo3(get)]
pub variance: Vec<f64>,
#[pyo3(get)]
pub ci_lower: Vec<f64>,
#[pyo3(get)]
pub ci_upper: Vec<f64>,
#[pyo3(get)]
pub n_risk: Vec<usize>,
#[pyo3(get)]
pub n_events: Vec<usize>,
}
#[pymethods]
impl NelsonAalenResult {
#[new]
fn new(
time: Vec<f64>,
cumulative_hazard: Vec<f64>,
variance: Vec<f64>,
ci_lower: Vec<f64>,
ci_upper: Vec<f64>,
n_risk: Vec<usize>,
n_events: Vec<usize>,
) -> Self {
Self {
time,
cumulative_hazard,
variance,
ci_lower,
ci_upper,
n_risk,
n_events,
}
}
fn survival(&self) -> Vec<f64> {
self.cumulative_hazard.iter().map(|&h| (-h).exp()).collect()
}
}
pub fn nelson_aalen(
time: &[f64],
status: &[i32],
weights: Option<&[f64]>,
confidence_level: f64,
) -> NelsonAalenResult {
let n = time.len();
if n == 0 {
return NelsonAalenResult {
time: vec![],
cumulative_hazard: vec![],
variance: vec![],
ci_lower: vec![],
ci_upper: vec![],
n_risk: vec![],
n_events: vec![],
};
}
let default_weights: Vec<f64> = vec![1.0; n];
let weights = weights.unwrap_or(&default_weights);
let mut indices: Vec<usize> = (0..n).collect();
indices.sort_by(|&a, &b| {
time[a]
.partial_cmp(&time[b])
.unwrap_or(std::cmp::Ordering::Equal)
});
let mut unique_times: Vec<f64> = Vec::new();
let mut events_at_time: Vec<f64> = Vec::new();
let mut at_risk: Vec<f64> = Vec::new();
let mut n_events_vec: Vec<usize> = Vec::new();
let mut n_risk_vec: Vec<usize> = Vec::new();
let mut total_weight: f64 = weights.iter().sum();
let mut total_count = n;
let mut i = 0;
while i < n {
let current_time = time[indices[i]];
let mut events = 0.0;
let mut event_count = 0usize;
let mut censored_weight = 0.0;
let mut censored_count = 0usize;
while i < n && time[indices[i]] == current_time {
let idx = indices[i];
if status[idx] == 1 {
events += weights[idx];
event_count += 1;
} else {
censored_weight += weights[idx];
censored_count += 1;
}
i += 1;
}
if event_count > 0 {
unique_times.push(current_time);
events_at_time.push(events);
at_risk.push(total_weight);
n_events_vec.push(event_count);
n_risk_vec.push(total_count);
}
total_weight -= events + censored_weight;
total_count -= event_count + censored_count;
}
let m = unique_times.len();
let mut cumulative_hazard = Vec::with_capacity(m);
let mut variance = Vec::with_capacity(m);
let mut cum_h = 0.0;
let mut cum_var = 0.0;
for j in 0..m {
let d = events_at_time[j];
let y = at_risk[j];
if y > 0.0 {
cum_h += d / y;
cum_var += d / (y * y);
}
cumulative_hazard.push(cum_h);
variance.push(cum_var);
}
let z = if confidence_level >= 0.99 {
2.576
} else if confidence_level >= 0.95 {
1.96
} else if confidence_level >= 0.90 {
1.645
} else {
1.28
};
let mut ci_lower = Vec::with_capacity(m);
let mut ci_upper = Vec::with_capacity(m);
for j in 0..m {
let se = variance[j].sqrt();
ci_lower.push((cumulative_hazard[j] - z * se).max(0.0));
ci_upper.push(cumulative_hazard[j] + z * se);
}
NelsonAalenResult {
time: unique_times,
cumulative_hazard,
variance,
ci_lower,
ci_upper,
n_risk: n_risk_vec,
n_events: n_events_vec,
}
}
#[pyfunction]
#[pyo3(signature = (time, status, weights=None, confidence_level=None))]
pub fn nelson_aalen_estimator(
time: Vec<f64>,
status: Vec<i32>,
weights: Option<Vec<f64>>,
confidence_level: Option<f64>,
) -> PyResult<NelsonAalenResult> {
let conf = confidence_level.unwrap_or(0.95);
let weights_ref = weights.as_deref();
Ok(nelson_aalen(&time, &status, weights_ref, conf))
}
#[derive(Debug, Clone)]
#[pyclass]
pub struct StratifiedKMResult {
#[pyo3(get)]
pub strata: Vec<i32>,
#[pyo3(get)]
pub times: Vec<Vec<f64>>,
#[pyo3(get)]
pub survival: Vec<Vec<f64>>,
#[pyo3(get)]
pub ci_lower: Vec<Vec<f64>>,
#[pyo3(get)]
pub ci_upper: Vec<Vec<f64>>,
#[pyo3(get)]
pub n_risk: Vec<Vec<usize>>,
#[pyo3(get)]
pub n_events: Vec<Vec<usize>>,
}
#[pymethods]
impl StratifiedKMResult {
#[new]
fn new(
strata: Vec<i32>,
times: Vec<Vec<f64>>,
survival: Vec<Vec<f64>>,
ci_lower: Vec<Vec<f64>>,
ci_upper: Vec<Vec<f64>>,
n_risk: Vec<Vec<usize>>,
n_events: Vec<Vec<usize>>,
) -> Self {
Self {
strata,
times,
survival,
ci_lower,
ci_upper,
n_risk,
n_events,
}
}
}
pub fn stratified_km(
time: &[f64],
status: &[i32],
strata: &[i32],
confidence_level: f64,
) -> StratifiedKMResult {
let n = time.len();
if n == 0 {
return StratifiedKMResult {
strata: vec![],
times: vec![],
survival: vec![],
ci_lower: vec![],
ci_upper: vec![],
n_risk: vec![],
n_events: vec![],
};
}
let mut unique_strata: Vec<i32> = strata.to_vec();
unique_strata.sort();
unique_strata.dedup();
let mut all_times = Vec::new();
let mut all_survival = Vec::new();
let mut all_ci_lower = Vec::new();
let mut all_ci_upper = Vec::new();
let mut all_n_risk = Vec::new();
let mut all_n_events = Vec::new();
for &stratum in &unique_strata {
let mask: Vec<bool> = strata.iter().map(|&s| s == stratum).collect();
let stratum_time: Vec<f64> = time
.iter()
.zip(&mask)
.filter(|&(_, m)| *m)
.map(|(&t, _)| t)
.collect();
let stratum_status: Vec<i32> = status
.iter()
.zip(&mask)
.filter(|&(_, m)| *m)
.map(|(&s, _)| s)
.collect();
let result = kaplan_meier(&stratum_time, &stratum_status, None, confidence_level);
all_times.push(result.time);
all_survival.push(result.survival);
all_ci_lower.push(result.ci_lower);
all_ci_upper.push(result.ci_upper);
all_n_risk.push(result.n_risk);
all_n_events.push(result.n_events);
}
StratifiedKMResult {
strata: unique_strata,
times: all_times,
survival: all_survival,
ci_lower: all_ci_lower,
ci_upper: all_ci_upper,
n_risk: all_n_risk,
n_events: all_n_events,
}
}
#[derive(Debug, Clone)]
struct KMResult {
time: Vec<f64>,
survival: Vec<f64>,
ci_lower: Vec<f64>,
ci_upper: Vec<f64>,
n_risk: Vec<usize>,
n_events: Vec<usize>,
}
fn kaplan_meier(
time: &[f64],
status: &[i32],
weights: Option<&[f64]>,
confidence_level: f64,
) -> KMResult {
let n = time.len();
if n == 0 {
return KMResult {
time: vec![],
survival: vec![],
ci_lower: vec![],
ci_upper: vec![],
n_risk: vec![],
n_events: vec![],
};
}
let default_weights: Vec<f64> = vec![1.0; n];
let weights = weights.unwrap_or(&default_weights);
let mut indices: Vec<usize> = (0..n).collect();
indices.sort_by(|&a, &b| {
time[a]
.partial_cmp(&time[b])
.unwrap_or(std::cmp::Ordering::Equal)
});
let mut unique_times: Vec<f64> = Vec::new();
let mut events_at_time: Vec<f64> = Vec::new();
let mut at_risk: Vec<f64> = Vec::new();
let mut n_events_vec: Vec<usize> = Vec::new();
let mut n_risk_vec: Vec<usize> = Vec::new();
let mut total_weight: f64 = weights.iter().sum();
let mut total_count = n;
let mut i = 0;
while i < n {
let current_time = time[indices[i]];
let mut events = 0.0;
let mut event_count = 0usize;
let mut removed_weight = 0.0;
let mut removed_count = 0usize;
while i < n && time[indices[i]] == current_time {
let idx = indices[i];
removed_weight += weights[idx];
removed_count += 1;
if status[idx] == 1 {
events += weights[idx];
event_count += 1;
}
i += 1;
}
if event_count > 0 {
unique_times.push(current_time);
events_at_time.push(events);
at_risk.push(total_weight);
n_events_vec.push(event_count);
n_risk_vec.push(total_count);
}
total_weight -= removed_weight;
total_count -= removed_count;
}
let m = unique_times.len();
let mut survival = Vec::with_capacity(m);
let mut greenwood_var = Vec::with_capacity(m);
let mut surv = 1.0;
let mut var_sum = 0.0;
for j in 0..m {
let d = events_at_time[j];
let y = at_risk[j];
if y > 0.0 {
surv *= 1.0 - d / y;
if y > d {
var_sum += d / (y * (y - d));
}
}
survival.push(surv);
greenwood_var.push(surv * surv * var_sum);
}
let z = if confidence_level >= 0.99 {
2.576
} else if confidence_level >= 0.95 {
1.96
} else if confidence_level >= 0.90 {
1.645
} else {
1.28
};
let mut ci_lower = Vec::with_capacity(m);
let mut ci_upper = Vec::with_capacity(m);
for j in 0..m {
let se = greenwood_var[j].sqrt();
ci_lower.push((survival[j] - z * se).clamp(0.0, 1.0));
ci_upper.push((survival[j] + z * se).clamp(0.0, 1.0));
}
KMResult {
time: unique_times,
survival,
ci_lower,
ci_upper,
n_risk: n_risk_vec,
n_events: n_events_vec,
}
}
#[pyfunction]
#[pyo3(signature = (time, status, strata, confidence_level=None))]
pub fn stratified_kaplan_meier(
time: Vec<f64>,
status: Vec<i32>,
strata: Vec<i32>,
confidence_level: Option<f64>,
) -> PyResult<StratifiedKMResult> {
let conf = confidence_level.unwrap_or(0.95);
Ok(stratified_km(&time, &status, &strata, conf))
}