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                visibility: super::point::PointVisibility::Visible,
59                hidden_targets: true,
60            },
61        )
62    }
63
64    fn required_aes(&self) -> Vec<Aesthetic> {
65        vec![Aesthetic::X, Aesthetic::Y]
66    }
67    fn default_stat(&self) -> Box<dyn Stat> {
68        Box::new(StatIdentity)
69    }
70    fn default_position(&self) -> Box<dyn Position> {
71        Box::new(PositionIdentity)
72    }
73    fn default_params(&self) -> GeomParams {
74        GeomParams::default()
75    }
76    fn name(&self) -> &str {
77        "censor_marks"
78    }
79    fn apply_element_geom(&mut self, el: &crate::theme::elements::ElementGeom) {
80        if let (true, Some(ink)) = (self.color == (0, 0, 0), el.ink) {
81            self.color = ink;
82        }
83    }
84    fn set_series_color(&mut self, color: (u8, u8, u8)) {
85        self.color = color;
86    }
87}
88
89/// Keep only the rows whose `column` marks a censoring: a number `> 0`
90/// (e.g. `n_censor`), `true`, or the strings `"true"`/`"1"`/`"censored"`.
91#[derive(Clone, Debug)]
92pub struct StatCensored {
93    pub column: String,
94}
95
96impl StatCensored {
97    pub fn new(column: &str) -> Self {
98        StatCensored {
99            column: column.to_string(),
100        }
101    }
102}
103
104fn is_censored(v: &Value) -> bool {
105    match v {
106        Value::Bool(b) => *b,
107        Value::Str(s) => matches!(
108            s.trim().to_ascii_lowercase().as_str(),
109            "true" | "1" | "censored" | "yes"
110        ),
111        Value::Na => false,
112        other => other.as_f64().is_some_and(|f| f > 0.0 && f.is_finite()),
113    }
114}
115
116impl Stat for StatCensored {
117    fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
118        let Some(flag) = data.column(&self.column) else {
119            return DataFrame::new();
120        };
121        let keep: Vec<usize> = (0..data.nrows())
122            .filter(|&i| is_censored(&flag[i]))
123            .collect();
124        let mut out = DataFrame::new();
125        for name in data.column_names() {
126            if let Some(col) = data.column(name) {
127                out.add_column(
128                    name.to_string(),
129                    keep.iter().map(|&i| col[i].clone()).collect(),
130                );
131            }
132        }
133        out
134    }
135
136    fn required_aes(&self) -> Vec<Aesthetic> {
137        vec![Aesthetic::X, Aesthetic::Y]
138    }
139
140    fn name(&self) -> &str {
141        "censored"
142    }
143}