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}
19
20impl ScaleDiscrete {
21    pub fn new() -> Self {
22        ScaleDiscrete {
23            aesthetic: Aesthetic::X,
24            name: String::new(),
25            levels: IndexSet::new(),
26            custom_labels: None,
27            limits: None,
28        }
29    }
30
31    pub fn for_aesthetic(mut self, aes: Aesthetic) -> Self {
32        self.aesthetic = aes;
33        self
34    }
35
36    pub fn with_name(mut self, name: &str) -> Self {
37        self.name = name.to_string();
38        self
39    }
40
41    /// Set custom display labels for each level. Must match the number of levels.
42    pub fn with_labels(mut self, labels: Vec<String>) -> Self {
43        self.custom_labels = Some(labels);
44        self
45    }
46
47    /// Set level order and filter. Only these levels are shown, in this order.
48    /// Data values not in limits are mapped to the middle (0.5).
49    pub fn with_limits(mut self, limits: Vec<&str>) -> Self {
50        self.limits = Some(limits.into_iter().map(|s| s.to_string()).collect());
51        self
52    }
53}
54
55impl ScaleDiscrete {
56    /// Get the effective levels (filtered by limits if set).
57    fn effective_levels(&self) -> &IndexSet<String> {
58        if let Some(ref limits) = self.limits {
59            limits
60        } else {
61            &self.levels
62        }
63    }
64}
65
66impl Default for ScaleDiscrete {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl Scale for ScaleDiscrete {
73    fn aesthetic(&self) -> Aesthetic {
74        self.aesthetic.clone()
75    }
76
77    fn train(&mut self, values: &[Value]) {
78        if let Some(ref limits) = self.limits {
79            // When limits are set, use them as the level order (ignore data order)
80            self.levels = limits.clone();
81        } else {
82            for v in values {
83                // ±Inf (a rect extending to the panel edge) is not a level.
84                if matches!(v, Value::Float(f) if !f.is_finite()) {
85                    continue;
86                }
87                let key = v.key_str();
88                if !self.levels.contains(key.as_ref()) {
89                    self.levels.insert(key.into_owned());
90                }
91            }
92        }
93    }
94
95    fn map(&self, value: &Value) -> f64 {
96        match value {
97            Value::Float(f) if *f == f64::INFINITY => return 1.0,
98            Value::Float(f) if *f == f64::NEG_INFINITY => return 0.0,
99            _ => {}
100        }
101        let key = value.key_str();
102        let effective = self.effective_levels();
103        let n = effective.len();
104        if n == 0 {
105            return 0.5;
106        }
107        match effective.get_index_of(key.as_ref()) {
108            Some(idx) => (idx as f64 + 0.5) / n as f64,
109            None => 0.5, // Not in limits → maps to middle
110        }
111    }
112
113    fn breaks(&self) -> Vec<(f64, String)> {
114        let effective = self.effective_levels();
115        let n = effective.len();
116        if n == 0 {
117            return vec![];
118        }
119        effective
120            .iter()
121            .enumerate()
122            .map(|(i, level)| {
123                let pos = (i as f64 + 0.5) / n as f64;
124                let label = if let Some(ref labels) = self.custom_labels {
125                    labels.get(i).cloned().unwrap_or_else(|| level.clone())
126                } else {
127                    level.clone()
128                };
129                (pos, label)
130            })
131            .collect()
132    }
133
134    fn name(&self) -> &str {
135        &self.name
136    }
137
138    fn set_name(&mut self, name: &str) {
139        self.name = name.to_string();
140    }
141
142    fn is_discrete(&self) -> bool {
143        true
144    }
145
146    fn clone_box(&self) -> Box<dyn Scale> {
147        Box::new(self.clone())
148    }
149
150    fn reset_training(&mut self) {
151        if self.limits.is_none() {
152            self.levels.clear();
153        }
154    }
155}