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 set_series_color(&mut self, color: (u8, u8, u8)) {
78        self.color = color;
79    }
80}
81
82/// Keep only the rows whose `column` marks a censoring: a number `> 0`
83/// (e.g. `n_censor`), `true`, or the strings `"true"`/`"1"`/`"censored"`.
84#[derive(Clone, Debug)]
85pub struct StatCensored {
86    pub column: String,
87}
88
89impl StatCensored {
90    pub fn new(column: &str) -> Self {
91        StatCensored {
92            column: column.to_string(),
93        }
94    }
95}
96
97fn is_censored(v: &Value) -> bool {
98    match v {
99        Value::Bool(b) => *b,
100        Value::Str(s) => matches!(
101            s.trim().to_ascii_lowercase().as_str(),
102            "true" | "1" | "censored" | "yes"
103        ),
104        Value::Na => false,
105        other => other.as_f64().is_some_and(|f| f > 0.0 && f.is_finite()),
106    }
107}
108
109impl Stat for StatCensored {
110    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
111        let Some(flag) = data.column(&self.column) else {
112            return DataFrame::new();
113        };
114        let keep: Vec<usize> = (0..data.nrows())
115            .filter(|&i| is_censored(&flag[i]))
116            .collect();
117        let mut out = DataFrame::new();
118        for name in data.column_names() {
119            if let Some(col) = data.column(name) {
120                out.add_column(
121                    name.to_string(),
122                    keep.iter().map(|&i| col[i].clone()).collect(),
123                );
124            }
125        }
126        out
127    }
128
129    fn required_aes(&self) -> Vec<Aesthetic> {
130        vec![Aesthetic::X, Aesthetic::Y]
131    }
132
133    fn name(&self) -> &str {
134        "censored"
135    }
136}