survival 1.1.37

A high-performance survival analysis library written in Rust with Python bindings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use pyo3::prelude::*;

/// Result of adjusting near ties in survival times
#[derive(Debug, Clone)]
#[pyclass]
pub struct AeqSurvResult {
    /// Adjusted survival times with near-ties resolved
    #[pyo3(get)]
    pub time: Vec<f64>,
    /// Number of values that were adjusted
    #[pyo3(get)]
    pub adjusted_count: usize,
    /// Indices of values that were adjusted
    #[pyo3(get)]
    pub adjusted_indices: Vec<usize>,
}

/// Adjudicate near ties in survival times.
///
/// This function handles floating-point precision issues that can cause
/// survival times that should be equal to be treated as different.
/// It compares values and replaces near-ties with the smaller value.
///
/// # Arguments
/// * `time` - Vector of survival times
/// * `tolerance` - Tolerance for considering values as tied (default: 1e-8 * range)
///
/// # Returns
/// * `AeqSurvResult` containing adjusted times and adjustment info
#[pyfunction]
#[pyo3(signature = (time, tolerance=None))]
pub fn aeq_surv(time: Vec<f64>, tolerance: Option<f64>) -> PyResult<AeqSurvResult> {
    let n = time.len();
    if n == 0 {
        return Ok(AeqSurvResult {
            time: vec![],
            adjusted_count: 0,
            adjusted_indices: vec![],
        });
    }

    let tol = tolerance.unwrap_or_else(|| {
        let min_val = time.iter().fold(f64::INFINITY, |a, &b| a.min(b));
        let max_val = time.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
        let range = max_val - min_val;
        if range > 0.0 { range * 1e-8 } else { 1e-8 }
    });

    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 adjusted_time = time.clone();
    let mut adjusted_indices = Vec::new();

    let mut i = 0;
    while i < n {
        let base_val = adjusted_time[indices[i]];
        let mut j = i + 1;

        while j < n {
            let current_val = adjusted_time[indices[j]];
            if (current_val - base_val).abs() <= tol {
                // Only count as adjusted if the value actually changes
                if current_val != base_val {
                    adjusted_time[indices[j]] = base_val;
                    adjusted_indices.push(indices[j]);
                }
                j += 1;
            } else {
                break;
            }
        }
        i = j;
    }

    let adjusted_count = adjusted_indices.len();

    Ok(AeqSurvResult {
        time: adjusted_time,
        adjusted_count,
        adjusted_indices,
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_aeq_surv_no_ties() {
        let time = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let result = aeq_surv(time.clone(), None).unwrap();
        assert_eq!(result.time, time);
        assert_eq!(result.adjusted_count, 0);
    }

    #[test]
    fn test_aeq_surv_near_ties() {
        let time = vec![1.0, 1.0 + 1e-10, 2.0, 3.0];
        let result = aeq_surv(time, Some(1e-8)).unwrap();
        assert_eq!(result.adjusted_count, 1);
        assert!((result.time[0] - result.time[1]).abs() < 1e-15);
    }

    #[test]
    fn test_aeq_surv_empty() {
        let time: Vec<f64> = vec![];
        let result = aeq_surv(time, None).unwrap();
        assert_eq!(result.time.len(), 0);
        assert_eq!(result.adjusted_count, 0);
    }

    #[test]
    fn test_aeq_surv_all_same() {
        let time = vec![1.0, 1.0, 1.0, 1.0];
        let result = aeq_surv(time, None).unwrap();
        assert_eq!(result.adjusted_count, 0); // Already exactly tied
    }
}