1#[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#[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 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 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 #[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}