qualia_core_db/solvers/learning/survival/
kaplan_meier.rs

1//! Kaplan–Meier survival curve (ISL ch 11.3) — the nonparametric estimate of the
2//! survival function `S(t)` from right-censored event times.
3//!
4//! At each distinct event time `tᵢ`, `S` drops by the factor `(1 − dᵢ/nᵢ)`, where
5//! `dᵢ` is the number of events at `tᵢ` and `nᵢ` the number still at risk
6//! (time ≥ `tᵢ`). Censored observations leave the risk set without an event.
7
8use crate::solvers::learning::LearningError;
9
10/// A fitted Kaplan–Meier estimator: a right-continuous step function.
11#[derive(Debug, Clone)]
12pub struct KaplanMeier {
13    /// Distinct event times, ascending.
14    pub event_times: Vec<f64>,
15    /// Survival probability *after* each event time.
16    pub survival: Vec<f64>,
17    /// Number at risk at each event time.
18    pub at_risk: Vec<usize>,
19    /// Number of events at each event time.
20    pub events: Vec<usize>,
21}
22
23impl KaplanMeier {
24    /// Fit from `times` and `event` flags (`true` = event/death observed,
25    /// `false` = right-censored). Fails closed on length mismatch / empty input.
26    pub fn fit(times: &[f64], event: &[bool]) -> Result<Self, LearningError> {
27        let n = times.len();
28        if n == 0 || n != event.len() {
29            return Err(LearningError::InvalidDimension);
30        }
31        // Order by time; ties resolved with events before censorings is not needed
32        // for the standard estimator (we group by exact time).
33        let mut order: Vec<usize> = (0..n).collect();
34        order.sort_by(|&a, &b| {
35            times[a]
36                .partial_cmp(&times[b])
37                .unwrap_or(core::cmp::Ordering::Equal)
38        });
39
40        let mut event_times = Vec::new();
41        let mut survival = Vec::new();
42        let mut at_risk_v = Vec::new();
43        let mut events_v = Vec::new();
44
45        let mut s = 1.0;
46        let mut i = 0;
47        while i < n {
48            let t = times[order[i]];
49            // Count events and total observations at this exact time.
50            let mut d = 0usize; // events at t
51            let mut tied = 0usize; // total (events + censorings) at t
52            let mut j = i;
53            while j < n && times[order[j]] == t {
54                if event[order[j]] {
55                    d += 1;
56                }
57                tied += 1;
58                j += 1;
59            }
60            let n_at_risk = n - i; // everyone with time ≥ t is still at risk
61            if d > 0 {
62                s *= 1.0 - d as f64 / n_at_risk as f64;
63                event_times.push(t);
64                survival.push(s);
65                at_risk_v.push(n_at_risk);
66                events_v.push(d);
67            }
68            let _ = tied;
69            i = j;
70        }
71
72        Ok(Self {
73            event_times,
74            survival,
75            at_risk: at_risk_v,
76            events: events_v,
77        })
78    }
79
80    /// Estimated survival `S(t)` (right-continuous step). `1.0` before the first
81    /// event time.
82    pub fn survival_at(&self, t: f64) -> f64 {
83        let mut s = 1.0;
84        for (k, &et) in self.event_times.iter().enumerate() {
85            if et <= t {
86                s = self.survival[k];
87            } else {
88                break;
89            }
90        }
91        s
92    }
93
94    /// Median survival time — the first event time at which `S(t) ≤ 0.5`. `None` if
95    /// the curve never drops to 0.5 (heavy censoring).
96    pub fn median_survival(&self) -> Option<f64> {
97        self.event_times
98            .iter()
99            .zip(self.survival.iter())
100            .find(|(_, &s)| s <= 0.5)
101            .map(|(&t, _)| t)
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn no_censoring_matches_empirical() {
111        // 4 events at times 1,2,3,4 → S drops 1→.75→.5→.25→0.
112        let times = [1.0, 2.0, 3.0, 4.0];
113        let event = [true, true, true, true];
114        let km = KaplanMeier::fit(&times, &event).unwrap();
115        assert!((km.survival_at(1.0) - 0.75).abs() < 1e-12);
116        assert!((km.survival_at(2.0) - 0.5).abs() < 1e-12);
117        assert!((km.survival_at(3.0) - 0.25).abs() < 1e-12);
118        assert!((km.survival_at(0.5) - 1.0).abs() < 1e-12); // before first event
119        assert_eq!(km.median_survival(), Some(2.0));
120    }
121
122    #[test]
123    fn censoring_keeps_survival_higher() {
124        // Times 1(event),2(censored),3(event),4(censored).
125        let times = [1.0, 2.0, 3.0, 4.0];
126        let event = [true, false, true, false];
127        let km = KaplanMeier::fit(&times, &event).unwrap();
128        // At t=1: 1 event of 4 at risk → S=0.75.
129        assert!((km.survival_at(1.0) - 0.75).abs() < 1e-12);
130        // At t=3: 1 event of 2 at risk (1 and 2 already gone) → S = 0.75·(1−1/2)=0.375.
131        assert!((km.survival_at(3.0) - 0.375).abs() < 1e-12);
132    }
133
134    #[test]
135    fn tied_events_drop_together() {
136        // Two events at time 2.
137        let times = [1.0, 2.0, 2.0, 4.0];
138        let event = [true, true, true, true];
139        let km = KaplanMeier::fit(&times, &event).unwrap();
140        // t=1: S=0.75. t=2: 2 of 3 at risk → S=0.75·(1−2/3)=0.25.
141        assert!((km.survival_at(2.0) - 0.25).abs() < 1e-12);
142    }
143
144    #[test]
145    fn guards() {
146        assert_eq!(
147            KaplanMeier::fit(&[], &[]).unwrap_err(),
148            LearningError::InvalidDimension
149        );
150        assert_eq!(
151            KaplanMeier::fit(&[1.0], &[true, false]).unwrap_err(),
152            LearningError::InvalidDimension
153        );
154    }
155}
qualia_core_db/solvers/learning/survival/kaplan_meier.rs

qualia_core_db/solvers/learning/survival/
kaplan_meier.rs
