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::shape::ScaleShapeDiscrete;
11use crate::scale::size::ScaleSizeContinuous;
12
13use super::Scale;
14
15/// Registry of all scales for a plot. Handles auto-detection and training.
16pub struct ScaleSet {
17    scales: Vec<Box<dyn Scale>>,
18}
19
20impl ScaleSet {
21    pub fn new() -> Self {
22        ScaleSet { scales: Vec::new() }
23    }
24
25    /// Add a user-specified scale.
26    pub fn add(&mut self, scale: Box<dyn Scale>) {
27        // Replace existing scale for same aesthetic
28        let aes = scale.aesthetic();
29        self.scales.retain(|s| s.aesthetic() != aes);
30        self.scales.push(scale);
31    }
32
33    /// Get a scale for a specific aesthetic.
34    pub fn get(&self, aes: &Aesthetic) -> Option<&dyn Scale> {
35        self.scales
36            .iter()
37            .find(|s| s.aesthetic() == *aes)
38            .map(|s| s.as_ref())
39    }
40
41    /// Get mutable scale for a specific aesthetic.
42    pub fn get_mut(&mut self, aes: &Aesthetic) -> Option<&mut Box<dyn Scale>> {
43        self.scales.iter_mut().find(|s| s.aesthetic() == *aes)
44    }
45
46    /// Ensure a scale exists for a given aesthetic. Auto-detect type from data.
47    pub fn ensure_scale(&mut self, aes: &Aesthetic, data: &DataFrame) {
48        if self.get(aes).is_some() {
49            return;
50        }
51
52        let col_name = aes.col_name();
53        let values = data.column(col_name);
54
55        let is_discrete = match values {
56            Some(vals) => vals
57                .iter()
58                .any(|v| matches!(v, Value::Str(_) | Value::Bool(_))),
59            None => false,
60        };
61
62        let is_datetime = match values {
63            Some(vals) => vals.iter().any(|v| v.is_datetime()),
64            None => false,
65        };
66
67        match aes {
68            Aesthetic::Color | Aesthetic::Fill => {
69                if is_discrete {
70                    let scale = ScaleColorDiscrete::new(aes.clone());
71                    self.scales.push(Box::new(scale));
72                } else {
73                    let scale = ScaleColorContinuous::new(aes.clone());
74                    self.scales.push(Box::new(scale));
75                }
76            }
77            Aesthetic::Shape => {
78                let scale = ScaleShapeDiscrete::new();
79                self.scales.push(Box::new(scale));
80            }
81            Aesthetic::Linetype => {
82                let scale = ScaleLinetypeDiscrete::new();
83                self.scales.push(Box::new(scale));
84            }
85            Aesthetic::Size => {
86                let scale = ScaleSizeContinuous::new();
87                self.scales.push(Box::new(scale));
88            }
89            Aesthetic::Alpha => {
90                let scale = ScaleAlphaContinuous::new();
91                self.scales.push(Box::new(scale));
92            }
93            _ => {
94                if is_discrete {
95                    let scale = ScaleDiscrete::new().for_aesthetic(aes.clone());
96                    self.scales.push(Box::new(scale));
97                } else if is_datetime {
98                    let scale = ScaleDateTime::new().for_aesthetic(aes.clone());
99                    self.scales.push(Box::new(scale));
100                } else {
101                    let scale = ScaleContinuous::new().for_aesthetic(aes.clone());
102                    self.scales.push(Box::new(scale));
103                }
104            }
105        }
106    }
107
108    /// Train all scales on data from one layer.
109    pub fn train_layer(&mut self, data: &DataFrame) {
110        for scale in &mut self.scales {
111            let col_name = scale.aesthetic().col_name().to_string();
112            if let Some(values) = data.column(&col_name) {
113                scale.train(values);
114            }
115        }
116    }
117
118    /// Map a single value through the appropriate scale.
119    pub fn map_value(&self, aes: &Aesthetic, value: &Value) -> f64 {
120        self.get(aes).map(|s| s.map(value)).unwrap_or(0.0)
121    }
122
123    /// Map a value to an RGB color through the appropriate color/fill scale.
124    pub fn map_color(&self, aes: &Aesthetic, value: &Value) -> Option<(u8, u8, u8)> {
125        self.get(aes).and_then(|s| s.map_to_color(value))
126    }
127
128    /// Map a value to a point shape through the shape scale.
129    pub fn map_shape(&self, value: &Value) -> Option<PointShape> {
130        self.get(&Aesthetic::Shape)
131            .and_then(|s| s.map_to_shape(value))
132    }
133
134    /// Map a value to a linetype through the linetype scale.
135    pub fn map_linetype(&self, value: &Value) -> Option<Linetype> {
136        self.get(&Aesthetic::Linetype)
137            .and_then(|s| s.map_to_linetype(value))
138    }
139
140    /// Map a value to a point size through the size scale.
141    pub fn map_size(&self, value: &Value) -> Option<f64> {
142        self.get(&Aesthetic::Size)
143            .and_then(|s| s.map_to_size(value))
144    }
145
146    /// Map a value to an alpha (opacity) through the alpha scale.
147    pub fn map_alpha(&self, value: &Value) -> Option<f64> {
148        self.get(&Aesthetic::Alpha)
149            .and_then(|s| s.map_to_alpha(value))
150    }
151
152    /// Override the domain limits for a scale (used by coord_cartesian zoom).
153    pub fn set_limits(&mut self, aes: &Aesthetic, min: f64, max: f64) {
154        if let Some(scale) = self.get_mut(aes) {
155            scale.set_limits(min, max);
156        }
157    }
158
159    /// Get the secondary axis for an aesthetic, if one exists.
160    pub fn sec_axis(&self, aes: &Aesthetic) -> Option<&crate::scale::sec_axis::SecAxis> {
161        self.get(aes).and_then(|s| s.sec_axis())
162    }
163
164    /// Get all scales.
165    pub fn iter(&self) -> impl Iterator<Item = &dyn Scale> {
166        self.scales.iter().map(|s| s.as_ref())
167    }
168}
169
170impl Clone for ScaleSet {
171    fn clone(&self) -> Self {
172        ScaleSet {
173            scales: self.scales.iter().map(|s| s.clone_box()).collect(),
174        }
175    }
176}
177
178impl Default for ScaleSet {
179    fn default() -> Self {
180        Self::new()
181    }
182}
ggplot_rs/scale/scale_set.rs

ggplot_rs/scale/
scale_set.rs
