Skip to main content

xpict_core/
plotdot.rs

1//! Concentric plot-dot shading (xenopict `PlotDot` semantics).
2//!
3//! Keep in sync with `python/xpict/draw/plotdot.py`.
4
5/// One disk before paint: radius fraction in `[0, 1]`, color stop in `[-1, 1]`.
6#[derive(Debug, Clone, Copy, PartialEq)]
7pub struct ShadeDisk {
8    pub radius_frac: f64,
9    pub color_z: f64,
10    pub x: f64,
11    pub y: f64,
12}
13
14/// Map `z ∈ [-1, 1]` to concentric `(radius_frac, color_z)` rings.
15///
16/// Rings sort weak → strong so intense colors paint on top — same order as
17/// xenopict / Python `PlotDot`.
18#[derive(Debug, Clone)]
19pub struct PlotDot {
20    levels: usize,
21    stops: Vec<f64>,
22}
23
24impl Default for PlotDot {
25    fn default() -> Self {
26        Self::new(4)
27    }
28}
29
30impl PlotDot {
31    pub fn new(levels: usize) -> Self {
32        let levels = levels.max(1);
33        let stops = (0..levels).map(|i| (i + 1) as f64 / levels as f64).collect();
34        Self { levels, stops }
35    }
36
37    fn radius(&self, z: f64, level: usize) -> f64 {
38        let az = z.abs();
39        if level == 0 {
40            return self.stops[0].sqrt();
41        }
42        let offset = 1.0 - self.stops[level];
43        let r = az - offset;
44        if r < self.stops[0] {
45            0.0
46        } else {
47            r.sqrt()
48        }
49    }
50
51    fn color(&self, z: f64, level: usize) -> f64 {
52        let sign = if z < 0.0 { -1.0 } else { 1.0 };
53        if level == 0 {
54            z
55        } else {
56            sign * self.stops[self.stops.len() - level - 1]
57        }
58    }
59
60    /// Rings for one score; empty when `|z| < 0.05`.
61    pub fn rings(&self, z: f64) -> Vec<(f64, f64)> {
62        if z.abs() < 0.05 {
63            return Vec::new();
64        }
65        (0..self.levels)
66            .map(|lvl| (self.radius(z, lvl), self.color(z, lvl)))
67            .filter(|(r, _)| *r > 0.0)
68            .collect()
69    }
70
71    /// All disks for a molecule, sorted weak/small first.
72    pub fn disks(&self, zs: &[f64], coords: &[(f64, f64)]) -> Vec<ShadeDisk> {
73        let mut dots = Vec::new();
74        for (z, &(x, y)) in zs.iter().zip(coords.iter()) {
75            for (radius_frac, color_z) in self.rings(*z) {
76                dots.push(ShadeDisk {
77                    radius_frac,
78                    color_z,
79                    x,
80                    y,
81                });
82            }
83        }
84        dots.sort_by(|a, b| {
85            a.color_z
86                .abs()
87                .partial_cmp(&b.color_z.abs())
88                .unwrap_or(std::cmp::Ordering::Equal)
89                .then_with(|| {
90                    a.radius_frac
91                        .partial_cmp(&b.radius_frac)
92                        .unwrap_or(std::cmp::Ordering::Equal)
93                })
94        });
95        dots
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn skips_near_zero_scores() {
105        let dots = PlotDot::default().disks(
106            &[0.0, 0.01, 0.9],
107            &[(0.0, 0.0), (1.0, 0.0), (2.0, 0.0)],
108        );
109        assert!(dots.iter().all(|d| d.color_z.abs() >= 0.05));
110        assert!(dots.iter().any(|d| d.color_z.abs() > 0.5));
111    }
112
113    #[test]
114    fn full_score_has_four_rings() {
115        let rings = PlotDot::default().rings(1.0);
116        assert_eq!(rings.len(), 4);
117        assert!((rings[0].0 - 0.5).abs() < 1e-9);
118        assert!((rings[0].1 - 1.0).abs() < 1e-9);
119    }
120
121    /// Parity with Python ``test_plotdot_rings_match_python`` sample scores.
122    #[test]
123    fn plotdot_rings_sample_scores() {
124        let pd = PlotDot::default();
125        for z in [0.25, 0.5, 0.9, 1.0] {
126            let rings = pd.rings(z);
127            assert!(!rings.is_empty() || z < 0.05);
128            for &(r, c) in &rings {
129                assert!(r > 0.0 && r <= 1.0);
130                assert!(c.abs() <= 1.0 + 1e-9);
131            }
132        }
133    }
134}