solow_duration/
survdiff.rs

1//! Log-rank and weighted log-rank tests for equality of survival distributions.
2//!
3//! [`survdiff`] compares the survival distributions of two or more groups using
4//! the (weighted) log-rank test. It mirrors the reference
5//! `survfunc.survdiff` for the single-stratum, no-entry case and supports the
6//! standard weight families:
7//!
8//! * [`WeightType::LogRank`] — the unweighted log-rank test (Mantel–Cox).
9//! * [`WeightType::GehanBreslow`] — weights by the number at risk.
10//! * [`WeightType::TaroneWare`] — weights by the square root of the number at
11//!   risk.
12//! * [`WeightType::FlemingHarrington`] — weights by `S(t-)^p`, where `S` is the
13//!   pooled Kaplan–Meier estimate at the previous event time.
14//!
15//! The statistic is `(O − E)' V⁻¹ (O − E)`, distributed as chi-square with
16//! `g − 1` degrees of freedom under the null, and the returned p-value is the
17//! upper tail of that chi-square distribution.
18
19use ndarray::{Array1, Array2};
20use solow_core::error::{Error, Result};
21use solow_distributions::chi2_sf;
22use solow_linalg::solve;
23
24/// Weight family for the (weighted) log-rank test.
25#[derive(Clone, Copy, Debug, PartialEq)]
26pub enum WeightType {
27    /// Unweighted log-rank test (all weights equal to 1).
28    LogRank,
29    /// Gehan–Breslow: weight by the total number at risk at each event time.
30    GehanBreslow,
31    /// Tarone–Ware: weight by the square root of the number at risk.
32    TaroneWare,
33    /// Fleming–Harrington: weight by `S(t-)^p` (pooled KM at previous time).
34    FlemingHarrington(f64),
35}
36
37/// Outcome of a (weighted) log-rank test.
38#[derive(Clone, Debug)]
39pub struct SurvDiffResult {
40    /// The chi-square test statistic.
41    pub chisq: f64,
42    /// The upper-tail p-value under the chi-square(`g - 1`) null.
43    pub pvalue: f64,
44    /// Degrees of freedom (`number of groups - 1`).
45    pub df: usize,
46}
47
48/// Test the equality of two or more survival distributions.
49///
50/// `time[i]` is the event or censoring time, `status[i]` is `1.0` for an
51/// observed event and `0.0` for right-censoring, and `group[i]` is a group
52/// label (any `f64`; distinct values define the groups, ordered ascending).
53///
54/// Returns the chi-square statistic, its p-value, and the degrees of freedom.
55pub fn survdiff(
56    time: &[f64],
57    status: &[f64],
58    group: &[f64],
59    weight_type: WeightType,
60) -> Result<SurvDiffResult> {
61    let n = time.len();
62    if status.len() != n || group.len() != n {
63        return Err(Error::Shape("time/status/group length mismatch".into()));
64    }
65    if n == 0 {
66        return Err(Error::Shape("empty sample".into()));
67    }
68
69    // Distinct group labels (ascending) -> `gr`.
70    let mut gr: Vec<f64> = group.to_vec();
71    gr.sort_by(|a, b| a.total_cmp(b));
72    gr.dedup();
73    let ng = gr.len();
74    if ng < 2 {
75        return Err(Error::Shape("survdiff requires at least two groups".into()));
76    }
77
78    // Unique event/censoring times (ascending) and the inverse map.
79    let mut utimes: Vec<f64> = time.to_vec();
80    utimes.sort_by(|a, b| a.total_cmp(b));
81    utimes.dedup();
82    let ml = utimes.len();
83    let time_rank = |t: f64| -> usize { utimes.partition_point(|&u| u < t) };
84
85    // Per-group event counts (obsv) and risk-set sizes (nrisk) at each time.
86    let mut obsv: Vec<Array1<f64>> = vec![Array1::zeros(ml); ng];
87    let mut nrisk: Vec<Array1<f64>> = vec![Array1::zeros(ml); ng];
88    let group_idx = |g: f64| -> usize { gr.iter().position(|&x| x == g).unwrap() };
89
90    // n[g][k] = number of subjects in group g whose time equals utimes[k].
91    let mut nbin: Vec<Array1<f64>> = vec![Array1::zeros(ml); ng];
92    for i in 0..n {
93        let gi = group_idx(group[i]);
94        let k = time_rank(time[i]);
95        nbin[gi][k] += 1.0;
96        if status[i].round() as i64 == 1 {
97            obsv[gi][k] += 1.0;
98        }
99    }
100    // Risk set: reverse cumulative sum of nbin (no entry/left truncation).
101    for g in 0..ng {
102        let mut acc = 0.0;
103        for k in (0..ml).rev() {
104            acc += nbin[g][k];
105            nrisk[g][k] = acc;
106        }
107    }
108
109    // Pooled observed events and total at risk at each time.
110    let mut obs = Array1::<f64>::zeros(ml);
111    let mut nrisk_tot = Array1::<f64>::zeros(ml);
112    for g in 0..ng {
113        for k in 0..ml {
114            obs[k] += obsv[g][k];
115            nrisk_tot[k] += nrisk[g][k];
116        }
117    }
118
119    // Indices where the total risk set exceeds 1 (others contribute nothing).
120    let ix: Vec<usize> = (0..ml).filter(|&k| nrisk_tot[k] > 1.0).collect();
121
122    // Weight series w[k].
123    let weights: Option<Array1<f64>> = match weight_type {
124        WeightType::LogRank => None,
125        WeightType::GehanBreslow => Some(nrisk_tot.clone()),
126        WeightType::TaroneWare => Some(nrisk_tot.mapv(f64::sqrt)),
127        WeightType::FlemingHarrington(p) => {
128            // sp = cumprod(1 - obs/nrisk_tot); weights = roll(sp^p, 1); w[0]=1.
129            let mut sp = Array1::<f64>::zeros(ml);
130            let mut logcum = 0.0;
131            for k in 0..ml {
132                let frac = 1.0 - obs[k] / nrisk_tot[k];
133                logcum += frac.ln();
134                sp[k] = logcum.exp();
135            }
136            let mut w = sp.mapv(|v| v.powf(p));
137            // np.roll(w, 1): shift right by one, wrap last -> first, then w[0]=1.
138            let mut rolled = Array1::<f64>::zeros(ml);
139            for k in 0..ml {
140                rolled[k] = w[(k + ml - 1) % ml];
141            }
142            rolled[0] = 1.0;
143            w = rolled;
144            Some(w)
145        }
146    };
147
148    let dfs = ng - 1;
149
150    // r[g][k] = nrisk[g][k] / clip(nrisk_tot[k], 1e-10).
151    let mut r: Vec<Array1<f64>> = vec![Array1::zeros(ml); ng];
152    for g in 0..ng {
153        for k in 0..ml {
154            let denom = nrisk_tot[k].max(1e-10);
155            r[g][k] = nrisk[g][k] / denom;
156        }
157    }
158
159    // var_denom = clip(nrisk_tot - 1, 1e-10).
160    let var_denom: Array1<f64> = nrisk_tot.mapv(|v| (v - 1.0).max(1e-10));
161    // var_scalar_part[k] = obs * (nrisk_tot - obs) / var_denom.
162    let var_scalar: Array1<f64> =
163        Array1::from_iter((0..ml).map(|k| obs[k] * (nrisk_tot[k] - obs[k]) / var_denom[k]));
164
165    // Build O-E vector and variance matrix using groups 1..dfs (reference uses
166    // the first group as reference).
167    let mut obs_vec = Array1::<f64>::zeros(dfs);
168    let mut var_mat = Array2::<f64>::zeros((dfs, dfs));
169
170    for g in 1..=dfs {
171        // oe[k] = obsv[g][k] - r[g][k] * obs[k]
172        let mut oe = Array1::<f64>::zeros(ml);
173        for k in 0..ml {
174            oe[k] = obsv[g][k] - r[g][k] * obs[k];
175        }
176
177        // var row over the dfs other groups: for column c (1..=dfs),
178        //   r[c][k] * (indicator(c == g) - r[g][k]) * var_scalar[k]
179        // accumulated over kept time indices.
180        let mut var_row = Array1::<f64>::zeros(dfs);
181
182        // Apply weights if present.
183        let (oe_w, w2): (Array1<f64>, Option<Array1<f64>>) = match &weights {
184            None => (oe, None),
185            Some(w) => {
186                let mut oew = Array1::<f64>::zeros(ml);
187                for k in 0..ml {
188                    oew[k] = w[k] * oe[k];
189                }
190                (oew, Some(w.mapv(|v| v * v)))
191            }
192        };
193
194        for &k in &ix {
195            obs_vec[g - 1] += oe_w[k];
196            for (ci, c) in (1..=dfs).enumerate() {
197                let ind = if c == g { 1.0 } else { 0.0 };
198                let mut v = r[c][k] * (ind - r[g][k]) * var_scalar[k];
199                if let Some(w2v) = &w2 {
200                    v *= w2v[k];
201                }
202                var_row[ci] += v;
203            }
204        }
205        for ci in 0..dfs {
206            var_mat[[g - 1, ci]] = var_row[ci];
207        }
208    }
209
210    // chisq = (O-E)' V^{-1} (O-E); pvalue = chi2 upper tail with dfs df.
211    let sol = solve(&var_mat, &obs_vec)?;
212    let chisq = obs_vec.dot(&sol);
213    let pvalue = chi2_sf(chisq, dfs as f64);
214
215    Ok(SurvDiffResult {
216        chisq,
217        pvalue,
218        df: dfs,
219    })
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[test]
227    fn identical_groups_zero_statistic() {
228        // Two groups with identical data -> O = E exactly -> chisq = 0.
229        let time = [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0];
230        let status = [1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0];
231        let group = [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0];
232        let res = survdiff(&time, &status, &group, WeightType::LogRank).unwrap();
233        assert!(res.chisq.abs() < 1e-12, "chisq={}", res.chisq);
234        assert!((res.pvalue - 1.0).abs() < 1e-12);
235        assert_eq!(res.df, 1);
236    }
237
238    #[test]
239    fn requires_two_groups() {
240        let time = [1.0, 2.0, 3.0];
241        let status = [1.0, 1.0, 1.0];
242        let group = [0.0, 0.0, 0.0];
243        assert!(survdiff(&time, &status, &group, WeightType::LogRank).is_err());
244    }
245
246    #[test]
247    fn statistic_nonnegative() {
248        let time = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
249        let status = [1.0, 1.0, 0.0, 1.0, 1.0, 1.0];
250        let group = [0.0, 1.0, 0.0, 1.0, 0.0, 1.0];
251        for wt in [
252            WeightType::LogRank,
253            WeightType::GehanBreslow,
254            WeightType::TaroneWare,
255            WeightType::FlemingHarrington(1.0),
256        ] {
257            let res = survdiff(&time, &status, &group, wt).unwrap();
258            assert!(res.chisq >= 0.0);
259            assert!((0.0..=1.0).contains(&res.pvalue));
260        }
261    }
262}
solow_duration/survdiff.rs

solow_duration/
survdiff.rs
