Skip to main content

ggplot_rs/scale/
discrete.rs

1use crate::aes::Aesthetic;
2use crate::data::Value;
3use indexmap::IndexSet;
4
5use super::Scale;
6
7/// Discrete scale: maps categorical values to evenly-spaced positions.
8#[derive(Clone, Debug)]
9pub struct ScaleDiscrete {
10    aesthetic: Aesthetic,
11    name: String,
12    /// Insertion-ordered set: O(1) training and lookup, so many-category
13    /// scales stay linear in the number of rows.
14    levels: IndexSet<String>,
15    custom_labels: Option<Vec<String>>,
16    /// Pre-set level order/filter. When set, only these levels are shown (in this order).
17    limits: Option<IndexSet<String>>,
18    /// Extra room below the first and above the last level, in levels
19    /// ([`ScaleDiscrete::with_expand`]).
20    expand: (f64, f64),
21}
22
23impl ScaleDiscrete {
24    pub fn new() -> Self {
25        ScaleDiscrete {
26            aesthetic: Aesthetic::X,
27            name: String::new(),
28            levels: IndexSet::new(),
29            custom_labels: None,
30            limits: None,
31            expand: (0.0, 0.0),
32        }
33    }
34
35    /// Extra room below the first level and above the last one, in levels
36    /// (0 by default: each level sits in the middle of its `1/n` slot). A
37    /// ridgeline's top ridge needs `scale - 0.5` above
38    /// ([`GeomDensityRidges`](crate::geom::density_ridges::GeomDensityRidges)).
39    pub fn with_expand(mut self, lower: f64, upper: f64) -> Self {
40        let ok = |v: f64| if v.is_finite() { v.max(0.0) } else { 0.0 };
41        self.expand = (ok(lower), ok(upper));
42        self
43    }
44
45    /// The mapped position of the `i`-th of `n` levels.
46    fn slot(&self, i: usize, n: usize) -> f64 {
47        let (lo, hi) = self.expand;
48        (i as f64 + 0.5 + lo) / (n as f64 + lo + hi)
49    }
50
51    pub fn for_aesthetic(mut self, aes: Aesthetic) -> Self {
52        self.aesthetic = aes;
53        self
54    }
55
56    pub fn with_name(mut self, name: &str) -> Self {
57        self.name = name.to_string();
58        self
59    }
60
61    /// Set custom display labels for each level. Must match the number of levels.
62    pub fn with_labels(mut self, labels: Vec<String>) -> Self {
63        self.custom_labels = Some(labels);
64        self
65    }
66
67    /// Set level order and filter. Only these levels are shown, in this order.
68    /// Data values not in limits are mapped to the middle (0.5).
69    pub fn with_limits(mut self, limits: Vec<&str>) -> Self {
70        self.limits = Some(limits.into_iter().map(|s| s.to_string()).collect());
71        self
72    }
73}
74
75impl ScaleDiscrete {
76    /// Get the effective levels (filtered by limits if set).
77    fn effective_levels(&self) -> &IndexSet<String> {
78        if let Some(ref limits) = self.limits {
79            limits
80        } else {
81            &self.levels
82        }
83    }
84}
85
86impl Default for ScaleDiscrete {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92impl Scale for ScaleDiscrete {
93    fn aesthetic(&self) -> Aesthetic {
94        self.aesthetic.clone()
95    }
96
97    fn train(&mut self, values: &[Value]) {
98        if let Some(ref limits) = self.limits {
99            // When limits are set, use them as the level order (ignore data order)
100            self.levels = limits.clone();
101        } else {
102            for v in values {
103                // ±Inf (a rect extending to the panel edge) is not a level.
104                if matches!(v, Value::Float(f) if !f.is_finite()) {
105                    continue;
106                }
107                let key = v.key_str();
108                if !self.levels.contains(key.as_ref()) {
109                    self.levels.insert(key.into_owned());
110                }
111            }
112        }
113    }
114
115    fn map(&self, value: &Value) -> f64 {
116        match value {
117            Value::Float(f) if *f == f64::INFINITY => return 1.0,
118            Value::Float(f) if *f == f64::NEG_INFINITY => return 0.0,
119            _ => {}
120        }
121        let key = value.key_str();
122        let effective = self.effective_levels();
123        let n = effective.len();
124        if n == 0 {
125            return 0.5;
126        }
127        match effective.get_index_of(key.as_ref()) {
128            Some(idx) => self.slot(idx, n),
129            None => 0.5, // Not in limits → maps to middle
130        }
131    }
132
133    fn breaks(&self) -> Vec<(f64, String)> {
134        let effective = self.effective_levels();
135        let n = effective.len();
136        if n == 0 {
137            return vec![];
138        }
139        effective
140            .iter()
141            .enumerate()
142            .map(|(i, level)| {
143                let pos = self.slot(i, n);
144                let label = if let Some(ref labels) = self.custom_labels {
145                    labels.get(i).cloned().unwrap_or_else(|| level.clone())
146                } else {
147                    level.clone()
148                };
149                (pos, label)
150            })
151            .collect()
152    }
153
154    fn name(&self) -> &str {
155        &self.name
156    }
157
158    fn set_name(&mut self, name: &str) {
159        self.name = name.to_string();
160    }
161
162    fn is_discrete(&self) -> bool {
163        true
164    }
165
166    fn clone_box(&self) -> Box<dyn Scale> {
167        Box::new(self.clone())
168    }
169
170    fn reset_training(&mut self) {
171        if self.limits.is_none() {
172            self.levels.clear();
173        }
174    }
175}