Skip to main content

ggplot_rs/geom/
censor.rs

1//! Censor marks for survival curves: `+` glyphs at the (time, survival)
2//! points where observations were censored (ggsurvfit's `add_censor_mark`).
3
4use crate::aes::Aesthetic;
5use crate::coord::Coord;
6use crate::data::{DataFrame, Value};
7use crate::position::identity::PositionIdentity;
8use crate::position::Position;
9use crate::render::backend::{DrawBackend, PointShape};
10use crate::render::RenderError;
11use crate::scale::ScaleSet;
12use crate::stat::identity::StatIdentity;
13use crate::stat::Stat;
14use crate::theme::Theme;
15
16use super::point::{draw_points, PointDefaults};
17use super::{Geom, GeomParams};
18
19/// Points drawn with a fixed shape (default `+`) — censor marks on a
20/// Kaplan–Meier curve. Mapped `color`/`size`/`alpha`/`shape` still win.
21pub struct GeomCensorMarks {
22    pub size: f64,
23    pub color: (u8, u8, u8),
24    pub alpha: f64,
25    pub shape: PointShape,
26}
27
28impl Default for GeomCensorMarks {
29    fn default() -> Self {
30        GeomCensorMarks {
31            size: 4.0,
32            color: (0, 0, 0),
33            alpha: 1.0,
34            shape: PointShape::Plus,
35        }
36    }
37}
38
39impl Geom for GeomCensorMarks {
40    fn draw(
41        &self,
42        data: &DataFrame,
43        coord: &dyn Coord,
44        scales: &ScaleSet,
45        _theme: &Theme,
46        backend: &mut dyn DrawBackend,
47    ) -> Result<(), RenderError> {
48        draw_points(
49            data,
50            coord,
51            scales,
52            backend,
53            PointDefaults {
54                size: self.size,
55                color: self.color,
56                alpha: self.alpha,
57                shape: self.shape,
58            },
59        )
60    }
61
62    fn required_aes(&self) -> Vec<Aesthetic> {
63        vec![Aesthetic::X, Aesthetic::Y]
64    }
65    fn default_stat(&self) -> Box<dyn Stat> {
66        Box::new(StatIdentity)
67    }
68    fn default_position(&self) -> Box<dyn Position> {
69        Box::new(PositionIdentity)
70    }
71    fn default_params(&self) -> GeomParams {
72        GeomParams::default()
73    }
74    fn name(&self) -> &str {
75        "censor_marks"
76    }
77    fn apply_element_geom(&mut self, el: &crate::theme::elements::ElementGeom) {
78        if let (true, Some(ink)) = (self.color == (0, 0, 0), el.ink) {
79            self.color = ink;
80        }
81    }
82    fn set_series_color(&mut self, color: (u8, u8, u8)) {
83        self.color = color;
84    }
85}
86
87/// Keep only the rows whose `column` marks a censoring: a number `> 0`
88/// (e.g. `n_censor`), `true`, or the strings `"true"`/`"1"`/`"censored"`.
89#[derive(Clone, Debug)]
90pub struct StatCensored {
91    pub column: String,
92}
93
94impl StatCensored {
95    pub fn new(column: &str) -> Self {
96        StatCensored {
97            column: column.to_string(),
98        }
99    }
100}
101
102fn is_censored(v: &Value) -> bool {
103    match v {
104        Value::Bool(b) => *b,
105        Value::Str(s) => matches!(
106            s.trim().to_ascii_lowercase().as_str(),
107            "true" | "1" | "censored" | "yes"
108        ),
109        Value::Na => false,
110        other => other.as_f64().is_some_and(|f| f > 0.0 && f.is_finite()),
111    }
112}
113
114impl Stat for StatCensored {
115    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
116        let Some(flag) = data.column(&self.column) else {
117            return DataFrame::new();
118        };
119        let keep: Vec<usize> = (0..data.nrows())
120            .filter(|&i| is_censored(&flag[i]))
121            .collect();
122        let mut out = DataFrame::new();
123        for name in data.column_names() {
124            if let Some(col) = data.column(name) {
125                out.add_column(
126                    name.to_string(),
127                    keep.iter().map(|&i| col[i].clone()).collect(),
128                );
129            }
130        }
131        out
132    }
133
134    fn required_aes(&self) -> Vec<Aesthetic> {
135        vec![Aesthetic::X, Aesthetic::Y]
136    }
137
138    fn name(&self) -> &str {
139        "censored"
140    }
141}