1use egui::{Align2, FontId, Pos2, Rect, Sense, Stroke, Ui, vec2};
7use facett_core::clip::{ClipKind, ClipPayload, CopySource};
8use facett_core::{Elm, FacetCaps, Semantics, a11y_node, theme};
9use serde::{Deserialize, Serialize};
10
11use crate::boxplot::paint_tooltip;
12use crate::ramp::{Ramp, normalize};
13
14#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
17pub struct HeatmapState {
18 pub hover: Option<(usize, usize)>,
20}
21
22#[derive(Clone, Debug, PartialEq)]
25pub enum HeatmapMsg {
26 Hover(Option<(usize, usize)>),
28}
29
30#[derive(Clone, Debug)]
32pub enum HeatmapEffect {}
33
34pub struct HeatmapView {
36 pub title: String,
37 pub rows: usize,
38 pub cols: usize,
39 pub data: Vec<f64>,
41 pub row_labels: Vec<String>,
42 pub col_labels: Vec<String>,
43 pub diverging: bool,
45 pub state: HeatmapState,
47}
48
49pub fn cell_rect(grid: Rect, rows: usize, cols: usize, row: usize, col: usize) -> Rect {
52 let cw = grid.width() / cols.max(1) as f32;
53 let ch = grid.height() / rows.max(1) as f32;
54 let x = grid.left() + cw * col as f32;
55 let y = grid.top() + ch * row as f32;
56 Rect::from_min_size(Pos2::new(x, y), vec2(cw, ch))
57}
58
59impl HeatmapView {
60 pub fn new(title: impl Into<String>, rows: usize, cols: usize, data: Vec<f64>) -> Self {
61 Self {
62 title: title.into(),
63 rows,
64 cols,
65 data,
66 row_labels: Vec::new(),
67 col_labels: Vec::new(),
68 diverging: false,
69 state: HeatmapState::default(),
70 }
71 }
72
73 pub fn state(&self) -> &HeatmapState {
75 &self.state
76 }
77
78 pub fn update(&mut self, msg: HeatmapMsg) -> Vec<HeatmapEffect> {
80 match msg {
81 HeatmapMsg::Hover(h) => self.state.hover = h,
82 }
83 Vec::new()
84 }
85
86 pub fn with_labels(mut self, rows: Vec<String>, cols: Vec<String>) -> Self {
87 self.row_labels = rows;
88 self.col_labels = cols;
89 self
90 }
91
92 pub fn diverging(mut self) -> Self {
93 self.diverging = true;
94 self
95 }
96
97 pub fn get(&self, row: usize, col: usize) -> Option<f64> {
99 if row >= self.rows || col >= self.cols {
100 return None;
101 }
102 self.data.get(row * self.cols + col).copied()
103 }
104
105 pub fn value_range(&self) -> Option<(f64, f64)> {
107 let mut r: Option<(f64, f64)> = None;
108 for &x in &self.data {
109 if x.is_finite() {
110 r = Some(match r {
111 None => (x, x),
112 Some((lo, hi)) => (lo.min(x), hi.max(x)),
113 });
114 }
115 }
116 r
117 }
118
119 fn to_csv(&self) -> String {
120 let mut s = String::new();
121 if !self.col_labels.is_empty() {
122 s.push_str(&format!(",{}\n", self.col_labels.join(",")));
123 }
124 for r in 0..self.rows {
125 if let Some(lbl) = self.row_labels.get(r) {
126 s.push_str(lbl);
127 } else {
128 s.push_str(&format!("r{r}"));
129 }
130 for c in 0..self.cols {
131 s.push(',');
132 s.push_str(&format!("{}", self.get(r, c).unwrap_or(f64::NAN)));
133 }
134 s.push('\n');
135 }
136 s
137 }
138}
139
140impl CopySource for HeatmapView {
142 fn copy_kinds(&self) -> &[ClipKind] {
143 &[ClipKind::Text]
144 }
145 fn copy_payload(&self) -> Option<ClipPayload> {
146 let csv = self.to_csv();
147 (!csv.is_empty()).then_some(ClipPayload::Text(csv))
148 }
149}
150
151impl HeatmapView {
152 fn render(&self, ui: &mut Ui) -> Vec<HeatmapMsg> {
155 let th = theme(ui);
156 let ramp = if self.diverging { Ramp::diverging(&th) } else { Ramp::sequential(&th) };
157 let (rect, resp) = ui.allocate_exact_size(ui.available_size(), Sense::hover());
158 let painter = ui.painter_at(rect);
159 painter.rect_filled(rect, egui::CornerRadius::ZERO, th.bg);
160 let Some((vmin, vmax)) = self.value_range() else {
161 #[cfg(feature = "testmatrix")]
163 facett_core::testmatrix::emit(
164 "facett-plotters::HeatmapView::ui",
165 "ui_render",
166 self.data.is_empty(),
167 &format!("drew=empty_hint rows={} cols={} cells={}", self.rows, self.cols, self.data.len()),
168 );
169 painter.text(rect.center(), Align2::CENTER_CENTER, "no data", FontId::proportional(13.0), th.text_dim);
170 return hover_msg(self.state.hover, None);
171 };
172 let legend_w = 56.0;
174 let grid = rect.shrink2(vec2(8.0, 8.0));
175 let grid = Rect::from_min_max(grid.min, Pos2::new(grid.right() - legend_w, grid.bottom()));
176 let hover = resp.hover_pos();
177 let mut new_hover: Option<(usize, usize)> = None;
178 let base = ui.id().with("heatmap");
179
180 for r in 0..self.rows {
181 for c in 0..self.cols {
182 let Some(v) = self.get(r, c) else { continue };
183 let cr = cell_rect(grid, self.rows, self.cols, r, c);
184 let color = ramp.at(normalize(v, vmin, vmax));
185 painter.rect_filled(cr.shrink(0.5), egui::CornerRadius::ZERO, color);
186 let rl = self.row_labels.get(r).cloned().unwrap_or_else(|| format!("r{r}"));
187 let cl = self.col_labels.get(c).cloned().unwrap_or_else(|| format!("c{c}"));
188 a11y_node(
192 ui,
193 base,
194 (r, c),
195 Sense::hover(),
196 cr,
197 Semantics::image(format!("{rl} × {cl}")).value(v),
198 );
199 if let Some(hp) = hover {
200 if cr.contains(hp) {
201 new_hover = Some((r, c));
202 }
203 }
204 }
205 }
206
207 let lx = rect.right() - legend_w + 14.0;
209 let lrect = Rect::from_min_max(Pos2::new(lx, grid.top()), Pos2::new(lx + 16.0, grid.bottom()));
210 let steps = 48;
211 for i in 0..steps {
212 let t = i as f32 / (steps - 1) as f32;
213 let y0 = lrect.bottom() - lrect.height() * (i as f32 / steps as f32);
214 let y1 = lrect.bottom() - lrect.height() * ((i + 1) as f32 / steps as f32);
215 painter.rect_filled(
216 Rect::from_min_max(Pos2::new(lrect.left(), y1), Pos2::new(lrect.right(), y0)),
217 egui::CornerRadius::ZERO,
218 ramp.at(t),
219 );
220 }
221 painter.rect_stroke(lrect, egui::CornerRadius::ZERO, Stroke::new(1.0_f32, th.panel_stroke), egui::StrokeKind::Outside);
222 painter.text(Pos2::new(lrect.right() + 4.0, lrect.top()), Align2::LEFT_CENTER, format!("{vmax:.2}"), FontId::proportional(10.0), th.text_dim);
223 painter.text(Pos2::new(lrect.right() + 4.0, lrect.bottom()), Align2::LEFT_CENTER, format!("{vmin:.2}"), FontId::proportional(10.0), th.text_dim);
224
225 if let Some(hp) = hover {
228 if let Some((r, c)) = self.state.hover {
229 if let Some(v) = self.get(r, c) {
230 let rl = self.row_labels.get(r).cloned().unwrap_or_else(|| format!("r{r}"));
231 let cl = self.col_labels.get(c).cloned().unwrap_or_else(|| format!("c{c}"));
232 paint_tooltip(&painter, &th, hp, &format!("{rl} × {cl}\n{v:.4}"));
233 }
234 }
235 }
236 #[cfg(feature = "testmatrix")]
238 facett_core::testmatrix::emit(
239 "facett-plotters::HeatmapView::ui",
240 "ui_render",
241 !self.data.is_empty() && self.rows * self.cols <= self.data.len(),
242 &format!("rows={} cols={} cells={}", self.rows, self.cols, self.data.len()),
243 );
244 hover_msg(self.state.hover, new_hover)
245 }
246}
247
248fn hover_msg(current: Option<(usize, usize)>, next: Option<(usize, usize)>) -> Vec<HeatmapMsg> {
250 if current != next { vec![HeatmapMsg::Hover(next)] } else { Vec::new() }
251}
252
253impl Elm for HeatmapView {
254 type Model = HeatmapState;
255 type Msg = HeatmapMsg;
256 type Effect = HeatmapEffect;
257
258 fn title(&self) -> &str {
259 &self.title
260 }
261 fn state(&self) -> &HeatmapState {
262 &self.state
263 }
264 fn update(&mut self, msg: HeatmapMsg) -> Vec<HeatmapEffect> {
265 HeatmapView::update(self, msg)
266 }
267 fn view(&self, ui: &mut Ui) -> Vec<HeatmapMsg> {
268 self.render(ui)
269 }
270}
271
272facett_core::impl_facet_via_elm!(HeatmapView, custom_state_json, {
275 fn caps(&self) -> FacetCaps {
276 FacetCaps::NONE.themeable().copyable().resizable()
277 }
278 fn copy(&mut self) -> Option<String> {
279 Some(self.to_csv())
280 }
281 fn state_json(&self) -> serde_json::Value {
282 let (mn, mx) = self.value_range().unwrap_or((0.0, 0.0));
283 serde_json::json!({
284 "kind": "heatmap",
285 "rows": self.rows, "cols": self.cols,
286 "cells": self.data.len(),
287 "min": mn, "max": mx,
288 "diverging": self.diverging,
289 })
290 }
291});
292
293#[cfg(test)]
294mod tests {
295 use super::*;
296
297 #[test]
298 fn typed_copy_is_the_dataset_csv() {
299 use facett_core::clip::{ClipKind, CopySource};
300 let v = HeatmapView::new("h", 2, 2, vec![1.0, 2.0, 3.0, 4.0]);
301 let p = v.copy_payload().expect("data copies");
302 assert_eq!(p.kind(), ClipKind::Text);
303 assert_eq!(p.as_text(), v.to_csv());
304 }
305
306 fn grid() -> Rect {
307 Rect::from_min_size(Pos2::new(0.0, 0.0), vec2(100.0, 80.0))
308 }
309
310 #[test]
311 fn cell_rects_tile_without_gaps() {
312 let g = grid();
313 let (rows, cols) = (4, 5);
314 let c00 = cell_rect(g, rows, cols, 0, 0);
316 assert_eq!(c00.min, Pos2::new(0.0, 0.0));
317 assert!((c00.width() - 20.0).abs() < 1e-4);
318 assert!((c00.height() - 20.0).abs() < 1e-4);
319 let c01 = cell_rect(g, rows, cols, 0, 1);
321 assert!((c01.left() - c00.right()).abs() < 1e-4);
322 let last = cell_rect(g, rows, cols, rows - 1, cols - 1);
324 assert!((last.right() - g.right()).abs() < 1e-4);
325 assert!((last.bottom() - g.bottom()).abs() < 1e-4);
326 }
327
328 #[test]
329 fn get_bounds_and_row_major() {
330 let h = HeatmapView::new("t", 2, 3, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
331 assert_eq!(h.get(0, 0), Some(0.0));
332 assert_eq!(h.get(0, 2), Some(2.0));
333 assert_eq!(h.get(1, 0), Some(3.0));
334 assert_eq!(h.get(1, 2), Some(5.0));
335 assert_eq!(h.get(2, 0), None); assert_eq!(h.get(0, 3), None); }
338
339 #[test]
340 fn value_range_ignores_nan() {
341 let h = HeatmapView::new("t", 1, 3, vec![1.0, f64::NAN, 9.0]);
342 assert_eq!(h.value_range(), Some((1.0, 9.0)));
343 }
344
345 #[test]
346 fn csv_round_trips_shape() {
347 let h = HeatmapView::new("t", 2, 2, vec![1.0, 2.0, 3.0, 4.0])
348 .with_labels(vec!["a".into(), "b".into()], vec!["x".into(), "y".into()]);
349 let csv = h.to_csv();
350 assert!(csv.contains(",x,y\n"));
351 assert!(csv.contains("a,1,2\n"));
352 assert!(csv.contains("b,3,4\n"));
353 }
354
355 #[test]
357 fn hover_transitions_are_observable() {
358 let mut v = HeatmapView::new("t", 2, 2, vec![1.0, 2.0, 3.0, 4.0]);
359 assert_eq!(v.state().hover, None);
360 let snap = facett_core::harness::snapshot(&mut v, [HeatmapMsg::Hover(Some((1, 0)))]);
361 assert_eq!(snap.hover, Some((1, 0)));
362 let snap = facett_core::harness::snapshot(&mut v, [HeatmapMsg::Hover(None)]);
363 assert_eq!(snap.hover, None);
364 }
365
366 #[test]
367 fn drive_produces_no_effects() {
368 let mut v = HeatmapView::new("t", 2, 2, vec![1.0, 2.0, 3.0, 4.0]);
369 let effects = facett_core::harness::drive(&mut v, [HeatmapMsg::Hover(Some((0, 1))), HeatmapMsg::Hover(None)]);
370 assert!(effects.is_empty());
371 }
372
373 #[test]
374 fn state_serde_round_trips() {
375 let st = HeatmapState { hover: Some((3, 2)) };
376 let back: HeatmapState = serde_json::from_value(serde_json::to_value(&st).unwrap()).unwrap();
377 assert_eq!(st, back);
378 }
379}