Skip to main content

ggplot_rs/scale/
scale_set.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::render::backend::{Linetype, PointShape};
4use crate::scale::alpha::ScaleAlphaContinuous;
5use crate::scale::color::{ScaleColorContinuous, ScaleColorDiscrete};
6use crate::scale::continuous::ScaleContinuous;
7use crate::scale::datetime::ScaleDateTime;
8use crate::scale::discrete::ScaleDiscrete;
9use crate::scale::linetype::ScaleLinetypeDiscrete;
10use crate::scale::linewidth::ScaleLinewidthContinuous;
11use crate::scale::shape::ScaleShapeDiscrete;
12use crate::scale::size::ScaleSizeContinuous;
13
14use super::Scale;
15
16/// Registry of all scales for a plot. Handles auto-detection and training.
17pub struct ScaleSet {
18    scales: Vec<Box<dyn Scale>>,
19}
20
21impl ScaleSet {
22    pub fn new() -> Self {
23        ScaleSet { scales: Vec::new() }
24    }
25
26    /// Add a user-specified scale.
27    pub fn add(&mut self, scale: Box<dyn Scale>) {
28        // Replace existing scale for same aesthetic
29        let aes = scale.aesthetic();
30        self.scales.retain(|s| s.aesthetic() != aes);
31        self.scales.push(scale);
32    }
33
34    /// Give every scale the plot-level default locale (scales with their own
35    /// locale keep it; scales without numeric/date labels ignore it).
36    pub fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
37        for s in &mut self.scales {
38            s.set_default_locale(locale);
39        }
40    }
41
42    /// Get a scale for a specific aesthetic.
43    pub fn get(&self, aes: &Aesthetic) -> Option<&dyn Scale> {
44        self.scales
45            .iter()
46            .find(|s| s.aesthetic() == *aes)
47            .map(|s| s.as_ref())
48    }
49
50    /// Get mutable scale for a specific aesthetic.
51    pub fn get_mut(&mut self, aes: &Aesthetic) -> Option<&mut Box<dyn Scale>> {
52        self.scales.iter_mut().find(|s| s.aesthetic() == *aes)
53    }
54
55    /// Ensure a scale exists for a given aesthetic. Auto-detect type from data.
56    pub fn ensure_scale(&mut self, aes: &Aesthetic, data: &DataFrame) {
57        if self.get(aes).is_some() {
58            return;
59        }
60        // OHLC values live on the y scale (via the geom's ymin/ymax extents);
61        // they get no scale of their own.
62        if matches!(
63            aes,
64            Aesthetic::Open
65                | Aesthetic::High
66                | Aesthetic::Low
67                | Aesthetic::Close
68                | Aesthetic::Slope
69                | Aesthetic::Intercept
70                // Reference-line intercepts live on the x / y scale; the
71                // build creates that scale from them only if no layer did
72                // (so a `geom_vline(0)` added first can't force a
73                // continuous scale onto a discrete axis).
74                | Aesthetic::Xintercept
75                | Aesthetic::Yintercept
76                // Interactive aesthetics carry their value straight through
77                // to the mark's attributes — no visual mapping, no legend.
78                | Aesthetic::Tooltip
79                | Aesthetic::DataAttr(_)
80                | Aesthetic::Class
81                | Aesthetic::Key
82                | Aesthetic::GroupWrap(_)
83                // Per-row fontface (#228) is parsed straight into
84                // `FontFace` — no visual mapping, no legend.
85                | Aesthetic::Fontface
86                // Per-row text justification: numbers used as they are.
87                | Aesthetic::Hjust
88                | Aesthetic::Vjust
89                // Boxplot stat="identity" hinges/outliers (B3-5) live on
90                // the y scale (via `column_aesthetic`/`Y_FAMILY`/the
91                // build's stat_aes training list), not a scale of their
92                // own.
93                | Aesthetic::Lower
94                | Aesthetic::Middle
95                | Aesthetic::Upper
96                | Aesthetic::Outliers
97                // Edge-list node names (flow charts) — no scale.
98                | Aesthetic::From
99                | Aesthetic::To
100                | Aesthetic::Subgroup
101                | Aesthetic::Sets
102        ) {
103            return;
104        }
105        // A position scale created from extent columns alone (`xmin`/`xmax`
106        // with no `x`, e.g. a Gantt bar) detects its type from them, so a
107        // date range still gets a date scale.
108        let fallback: &[&str] = match aes {
109            Aesthetic::X => &["xmin", "xmax", "xend"],
110            Aesthetic::Y => &["ymin", "ymax", "yend"],
111            _ => &[],
112        };
113        let values = data
114            .column(aes.col_name())
115            .or_else(|| fallback.iter().find_map(|c| data.column(c)));
116
117        let is_discrete = match values {
118            Some(vals) => vals
119                .iter()
120                .any(|v| matches!(v, Value::Str(_) | Value::Bool(_))),
121            None => false,
122        };
123
124        let is_datetime = match values {
125            Some(vals) => vals.iter().any(|v| v.is_datetime()),
126            None => false,
127        };
128
129        match aes {
130            Aesthetic::Color | Aesthetic::Fill => {
131                if is_discrete {
132                    let scale = ScaleColorDiscrete::new(aes.clone());
133                    self.scales.push(Box::new(scale));
134                } else {
135                    let scale = ScaleColorContinuous::new(aes.clone());
136                    self.scales.push(Box::new(scale));
137                }
138            }
139            Aesthetic::Shape => {
140                let scale = ScaleShapeDiscrete::new();
141                self.scales.push(Box::new(scale));
142            }
143            Aesthetic::Linetype => {
144                let scale = ScaleLinetypeDiscrete::new();
145                self.scales.push(Box::new(scale));
146            }
147            Aesthetic::Size => {
148                let scale = ScaleSizeContinuous::new();
149                self.scales.push(Box::new(scale));
150            }
151            Aesthetic::Linewidth => {
152                let scale = ScaleLinewidthContinuous::new();
153                self.scales.push(Box::new(scale));
154            }
155            Aesthetic::Alpha => {
156                let scale = ScaleAlphaContinuous::new();
157                self.scales.push(Box::new(scale));
158            }
159            _ => {
160                if is_discrete {
161                    let scale = ScaleDiscrete::new().for_aesthetic(aes.clone());
162                    self.scales.push(Box::new(scale));
163                } else if is_datetime {
164                    let scale = ScaleDateTime::new().for_aesthetic(aes.clone());
165                    self.scales.push(Box::new(scale));
166                } else {
167                    let scale = ScaleContinuous::new().for_aesthetic(aes.clone());
168                    self.scales.push(Box::new(scale));
169                }
170            }
171        }
172    }
173
174    /// Train all scales on data from one layer.
175    pub fn train_layer(&mut self, data: &DataFrame) {
176        for scale in &mut self.scales {
177            let col_name = scale.aesthetic().col_name().to_string();
178            if let Some(values) = data.column(&col_name) {
179                scale.train(values);
180            }
181        }
182    }
183
184    /// Map a single value through the appropriate scale.
185    pub fn map_value(&self, aes: &Aesthetic, value: &Value) -> f64 {
186        self.get(aes).map(|s| s.map(value)).unwrap_or(0.0)
187    }
188
189    /// Map a value to an RGB color through the appropriate color/fill scale.
190    pub fn map_color(&self, aes: &Aesthetic, value: &Value) -> Option<(u8, u8, u8)> {
191        self.get(aes).and_then(|s| s.map_to_color(value))
192    }
193
194    /// The `aes` scale's configured na-pattern, if any (`ScaleFillBinned::
195    /// na_pattern`, #206) — a geom checks this for a row/key it would
196    /// otherwise draw with the scale's flat na colour, before falling
197    /// back to [`map_color`](Self::map_color)'s plain result.
198    pub fn na_fill_pattern(&self, aes: &Aesthetic) -> Option<crate::render::backend::FillPattern> {
199        self.get(aes).and_then(|s| s.na_fill_pattern())
200    }
201
202    /// Map a value to a point shape through the shape scale.
203    pub fn map_shape(&self, value: &Value) -> Option<PointShape> {
204        self.get(&Aesthetic::Shape)
205            .and_then(|s| s.map_to_shape(value))
206    }
207
208    /// Map a value to a linetype through the linetype scale.
209    pub fn map_linetype(&self, value: &Value) -> Option<Linetype> {
210        self.get(&Aesthetic::Linetype)
211            .and_then(|s| s.map_to_linetype(value))
212    }
213
214    /// Map a value to a point size through the size scale.
215    pub fn map_size(&self, value: &Value) -> Option<f64> {
216        self.get(&Aesthetic::Size)
217            .and_then(|s| s.map_to_size(value))
218    }
219
220    /// Map a value to an alpha (opacity) through the alpha scale.
221    pub fn map_alpha(&self, value: &Value) -> Option<f64> {
222        self.get(&Aesthetic::Alpha)
223            .and_then(|s| s.map_to_alpha(value))
224    }
225
226    /// Map a value to a stroke width through the linewidth scale.
227    pub fn map_linewidth(&self, value: &Value) -> Option<f64> {
228        self.get(&Aesthetic::Linewidth)
229            .and_then(|s| s.map_to_size(value))
230    }
231
232    /// Override the domain limits for a scale (used by coord_cartesian zoom).
233    pub fn set_limits(&mut self, aes: &Aesthetic, min: f64, max: f64) {
234        if let Some(scale) = self.get_mut(aes) {
235            scale.set_limits(min, max);
236        }
237    }
238
239    /// Get the secondary axis for an aesthetic, if one exists.
240    pub fn sec_axis(&self, aes: &Aesthetic) -> Option<&crate::scale::sec_axis::SecAxis> {
241        self.get(aes).and_then(|s| s.sec_axis())
242    }
243
244    /// Get all scales.
245    pub fn iter(&self) -> impl Iterator<Item = &dyn Scale> {
246        self.scales.iter().map(|s| s.as_ref())
247    }
248}
249
250impl Clone for ScaleSet {
251    fn clone(&self) -> Self {
252        ScaleSet {
253            scales: self.scales.iter().map(|s| s.clone_box()).collect(),
254        }
255    }
256}
257
258impl Default for ScaleSet {
259    fn default() -> Self {
260        Self::new()
261    }
262}