ggplot_rs/scale/

size.rs

use crate::aes::Aesthetic;
use crate::data::Value;

use super::Scale;

/// Continuous size scale — maps data values to a point radius range.
#[derive(Clone, Debug)]
pub struct ScaleSizeContinuous {
    name: String,
    min: f64,
    max: f64,
    trained: bool,
    /// Set by [`with_limits`](Self::with_limits): the domain is fixed and
    /// no longer widened by [`train`](Scale::train) — mirrors
    /// `ScaleLinewidth::with_limits` (sipemu/ggplot-rs#270).
    limits_fixed: bool,
    /// Output range in pixels: (min_radius, max_radius).
    range: (f64, f64),
    /// Locale for the default break labels (`None` = default marks).
    locale: Option<crate::locale::Locale>,
    /// `scale_size_area(max_size)` (B3-6): private, set only via
    /// [`area`](Self::area). `None` (the default) keeps the linear radius
    /// mapping (`range`); `Some(max_size)` maps so *area* — not radius —
    /// is proportional to the value, with 0 mapping to 0.
    area_max: Option<f64>,
}

impl ScaleSizeContinuous {
    pub fn new() -> Self {
        ScaleSizeContinuous {
            name: String::new(),
            min: f64::INFINITY,
            max: f64::NEG_INFINITY,
            trained: false,
            limits_fixed: false,
            range: (1.0, 6.0),
            locale: None,
            area_max: None,
        }
    }

    pub fn with_range(mut self, min: f64, max: f64) -> Self {
        self.range = (min, max);
        self
    }

    /// Fix the input domain; unlike the default (which widens to whatever
    /// `train` sees), values outside `[min, max]` are clamped to the
    /// nearest bound and later `train` calls don't widen it further
    /// (#270) — a size encoding shared across facets/frames/renders needs
    /// one domain fixed in advance, the same way `ScaleLinewidth::
    /// with_limits` already lets linewidth do.
    pub fn with_limits(mut self, min: f64, max: f64) -> Self {
        self.min = min;
        self.max = max;
        self.trained = true;
        self.limits_fixed = true;
        self
    }

    pub fn with_name(mut self, name: &str) -> Self {
        self.name = name.to_string();
        self
    }

    /// `scale_size_area(max_size)` (ggplot2): maps so that *area* — not
    /// radius — is proportional to `|value|`, with 0 mapping to 0:
    /// `size = max_size · √(|value| / max(|lo|, |hi|))`, where `lo`/`hi`
    /// are the trained domain. A non-finite or non-positive `max_size` is
    /// corrected to the ggplot2 default (6) by [`validate`](Self::validate)
    /// and by the mapping itself.
    pub fn area(mut self, max_size: f64) -> Self {
        self.area_max = Some(max_size);
        self
    }

    fn sanitized_area_max(max_size: f64) -> f64 {
        if max_size.is_finite() && max_size > 0.0 {
            max_size
        } else {
            6.0
        }
    }
}

impl Default for ScaleSizeContinuous {
    fn default() -> Self {
        Self::new()
    }
}

impl Scale for ScaleSizeContinuous {
    fn set_default_locale(&mut self, locale: &crate::locale::Locale) {
        if self.locale.is_none() {
            self.locale = Some(locale.clone());
        }
    }

    fn locale(&self) -> Option<&crate::locale::Locale> {
        self.locale.as_ref()
    }

    fn format_value(&self, value: &Value) -> Option<String> {
        let f = value.as_f64()?;
        if !f.is_finite() {
            return None;
        }
        Some(super::util::format_number_opt(f, self.locale.as_ref()))
    }

    fn aesthetic(&self) -> Aesthetic {
        Aesthetic::Size
    }

    fn train(&mut self, values: &[Value]) {
        if self.limits_fixed {
            self.trained = true;
            return;
        }
        for v in values {
            if let Some(f) = v.as_f64() {
                if f.is_finite() {
                    if f < self.min {
                        self.min = f;
                    }
                    if f > self.max {
                        self.max = f;
                    }
                }
            }
        }
        self.trained = true;
    }

    fn map(&self, value: &Value) -> f64 {
        let f = match value.as_f64() {
            Some(f) => f,
            None => return 0.0,
        };
        if !f.is_finite() {
            return 0.0;
        }
        let range = self.max - self.min;
        if range.abs() < f64::EPSILON {
            0.5
        } else {
            // #270: a `with_limits` domain can be narrower than the data
            // (unlike the default, data-trained domain, which never is)
            // — clamping to [0, 1] here squishes an out-of-domain value
            // to the nearest end, the same as `ScaleLinewidth`.
            ((f - self.min) / range).clamp(0.0, 1.0)
        }
    }

    fn breaks(&self) -> Vec<(f64, String)> {
        if !self.trained || self.min > self.max {
            return vec![];
        }
        let range = self.max - self.min;
        if super::util::is_degenerate_range(self.min, self.max) {
            return vec![(
                0.5,
                super::util::format_number_opt(self.min, self.locale.as_ref()),
            )];
        }
        let n_breaks = 4;
        let step = super::util::nice_step(range / n_breaks as f64);
        let start = (self.min / step).ceil() * step;
        let values = super::util::stepped_breaks(start, self.max, step);
        let labels = super::util::format_numbers_opt(&values, self.locale.as_ref());
        values
            .iter()
            .zip(labels)
            .map(|(&v, l)| (self.map(&Value::Float(v)), l))
            .collect()
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn set_name(&mut self, name: &str) {
        self.name = name.to_string();
    }

    fn map_to_size(&self, value: &Value) -> Option<f64> {
        if let Some(max_size) = self.area_max {
            let v = value.as_f64()?;
            if !v.is_finite() {
                return None;
            }
            let max_size = Self::sanitized_area_max(max_size);
            let denom = self.min.abs().max(self.max.abs());
            return Some(if denom < f64::EPSILON {
                0.0
            } else {
                max_size * (v.abs() / denom).sqrt()
            });
        }
        let t = self.map(value);
        let (lo, hi) = self.range;
        Some(lo + t * (hi - lo))
    }

    fn validate(&self, warnings: &mut Vec<String>) {
        if let Some(max_size) = self.area_max {
            if !(max_size.is_finite() && max_size > 0.0) {
                warnings.push(format!(
                    "scale_size_area: max_size({max_size}) must be finite and positive; falling back to 6"
                ));
            }
        }
    }

    /// #180: under `scale_size_area` (`area_max` set) this must use the
    /// same area-proportional formula as [`map_to_size`](Self::map_to_size)
    /// — not the plain linear `range` — or a graduated-symbol legend's
    /// circles don't match the radii the data itself draws at (a legend
    /// key for "break at 50% of the domain" rendered at 50% of the
    /// *radius* range, when the data point at that value is actually drawn
    /// at `max_size · √0.5 ≈ 71%` of `max_size`, say).
    fn map_fraction_to_size(&self, fraction: f64) -> Option<f64> {
        let t = fraction.clamp(0.0, 1.0);
        if let Some(max_size) = self.area_max {
            if !self.trained || self.max <= self.min {
                return Some(0.0);
            }
            let max_size = Self::sanitized_area_max(max_size);
            let v = self.min + t * (self.max - self.min);
            let denom = self.min.abs().max(self.max.abs());
            return Some(if denom < f64::EPSILON {
                0.0
            } else {
                max_size * (v.abs() / denom).sqrt()
            });
        }
        let (lo, hi) = self.range;
        Some(lo + t * (hi - lo))
    }

    fn clone_box(&self) -> Box<dyn Scale> {
        Box::new(self.clone())
    }

    fn reset_training(&mut self) {
        if !self.limits_fixed {
            self.min = f64::INFINITY;
            self.max = f64::NEG_INFINITY;
            self.trained = false;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn format_value_formats_a_finite_number_and_rejects_non_finite() {
        let s = ScaleSizeContinuous::new();
        assert_eq!(
            s.format_value(&Value::Float(1234.0)),
            Some("1234".to_string())
        );
        assert_eq!(s.format_value(&Value::Float(f64::NAN)), None);
        assert_eq!(s.format_value(&Value::Na), None);
    }

    #[test]
    fn scale_size_area_matches_r_ggplot2() {
        let mut s = ScaleSizeContinuous::new().area(6.0);
        s.train(&[
            Value::Float(0.0),
            Value::Float(1.0),
            Value::Float(4.0),
            Value::Float(16.0),
        ]);
        let got: Vec<f64> = [0.0, 1.0, 4.0, 16.0]
            .iter()
            .map(|&v| s.map_to_size(&Value::Float(v)).unwrap())
            .collect();
        for (g, e) in got.iter().zip([0.0, 1.5, 3.0, 6.0]) {
            assert!((g - e).abs() < 1e-9, "got {got:?}");
        }
    }

    /// #180: a graduated-symbol legend (`draw_legend_swatch` calls
    /// `map_fraction_to_size(pos)` where `pos = map(value)`) must draw the
    /// same radius `map_to_size(value)` does for that value — not the
    /// unrelated default linear `range`.
    #[test]
    fn scale_size_area_legend_fraction_matches_map_to_size() {
        let mut s = ScaleSizeContinuous::new().area(6.0);
        s.train(&[Value::Float(0.0), Value::Float(16.0)]);
        for v in [0.0, 1.0, 4.0, 9.0, 16.0] {
            let direct = s.map_to_size(&Value::Float(v)).unwrap();
            let pos = s.map(&Value::Float(v));
            let via_fraction = s.map_fraction_to_size(pos).unwrap();
            assert!(
                (direct - via_fraction).abs() < 1e-9,
                "value {v}: map_to_size={direct}, map_fraction_to_size={via_fraction}"
            );
        }
    }

    /// Without `.area(..)` (the plain linear size scale), `range` still
    /// drives both the data mapping and the legend fraction mapping.
    #[test]
    fn linear_size_legend_fraction_still_uses_range() {
        let mut s = ScaleSizeContinuous::new().with_range(2.0, 10.0);
        s.train(&[Value::Float(0.0), Value::Float(100.0)]);
        assert_eq!(s.map_fraction_to_size(0.0), Some(2.0));
        assert_eq!(s.map_fraction_to_size(1.0), Some(10.0));
        assert_eq!(s.map_fraction_to_size(0.5), Some(6.0));
    }

    #[test]
    fn scale_size_area_rejects_na_and_non_finite() {
        let mut s = ScaleSizeContinuous::new().area(6.0);
        s.train(&[Value::Float(0.0), Value::Float(10.0)]);
        assert_eq!(s.map_to_size(&Value::Na), None);
        assert_eq!(s.map_to_size(&Value::Float(f64::NAN)), None);
    }

    #[test]
    fn scale_size_area_bad_max_warns_and_falls_back_to_six() {
        let mut s = ScaleSizeContinuous::new().area(-1.0);
        s.train(&[Value::Float(0.0), Value::Float(16.0)]);
        let mut warnings = Vec::new();
        s.validate(&mut warnings);
        assert!(warnings.iter().any(|w| w.contains("max_size")));
        // Falls back to the ggplot2 default (6) regardless of `validate`
        // being called — `map_to_size` guards it inline too.
        assert_eq!(
            s.map_to_size(&Value::Float(16.0)),
            Some(6.0),
            "falls back to max_size 6"
        );
    }

    // ─── #270: `with_limits` ───────────────────────────────────────────

    /// Two renders of different data, sharing the same `with_limits`
    /// domain, give identical radii for equal values and identical
    /// legend breaks — the motivating case (a size shared across
    /// frames/facets/renders).
    #[test]
    fn with_limits_gives_identical_radii_and_breaks_across_different_data() {
        let mut a = ScaleSizeContinuous::new()
            .with_range(2.0, 10.0)
            .with_limits(0.0, 100.0);
        let mut b = ScaleSizeContinuous::new()
            .with_range(2.0, 10.0)
            .with_limits(0.0, 100.0);
        a.train(&[Value::Float(1.0), Value::Float(2.0)]);
        b.train(&[Value::Float(50.0), Value::Float(99.0), Value::Float(0.0)]);

        for v in [0.0, 25.0, 50.0, 75.0, 100.0] {
            assert_eq!(
                a.map_to_size(&Value::Float(v)),
                b.map_to_size(&Value::Float(v)),
                "value {v} should map to the same radius regardless of trained data"
            );
        }
        assert_eq!(a.breaks(), b.breaks());
    }

    #[test]
    fn with_limits_clamps_out_of_range_values() {
        let s = ScaleSizeContinuous::new()
            .with_range(2.0, 10.0)
            .with_limits(0.0, 100.0);
        assert_eq!(s.map_to_size(&Value::Float(-50.0)), Some(2.0));
        assert_eq!(s.map_to_size(&Value::Float(500.0)), Some(10.0));
    }

    /// `area(max_size)` under `with_limits`: the denominator is
    /// `max(|min|, |max|)` of the *fixed* limits, not whatever happens to
    /// be trained — so two different data batches sharing limits give
    /// the same area-proportional radius for the same value.
    #[test]
    fn with_limits_and_area_share_radii_across_different_data() {
        let mut a = ScaleSizeContinuous::new().area(6.0).with_limits(0.0, 16.0);
        let mut b = ScaleSizeContinuous::new().area(6.0).with_limits(0.0, 16.0);
        a.train(&[Value::Float(1.0)]);
        b.train(&[Value::Float(1.0), Value::Float(16.0), Value::Float(8.0)]);
        for v in [0.0, 1.0, 4.0, 16.0] {
            assert_eq!(
                a.map_to_size(&Value::Float(v)),
                b.map_to_size(&Value::Float(v))
            );
        }
    }

    #[test]
    fn with_limits_training_does_not_change_map_to_size() {
        let mut fixed = ScaleSizeContinuous::new()
            .with_range(1.0, 6.0)
            .with_limits(0.0, 10.0);
        let before = fixed.map_to_size(&Value::Float(5.0));
        fixed.train(&[Value::Float(-1000.0), Value::Float(1000.0)]);
        let after = fixed.map_to_size(&Value::Float(5.0));
        assert_eq!(before, after, "training must not widen a fixed domain");
        assert_eq!(before, Some(3.5));
    }
}