aboutsummaryrefslogtreecommitdiff
path: root/src/heatmap.rs
blob: dd3ac3be776113f5ac0a035c1715db872462371c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
use std::io::Write;

use chrono::{Datelike, NaiveDate};

use crate::{Data, Result};

/// Renders an SVG heatmap.
pub struct Heatmap {
    pub projection: Box<dyn Projection>,
    pub colors: Vec<String>,
    pub unit: Option<String>,
    pub log_scale: bool,
}

impl Heatmap {
    pub fn render<T: Write>(&self, data: &Data, mut io: T) -> Result<()> {
        // No data, no image.
        if data.is_empty() {
            writeln!(
                io,
                r#"<svg xmlns="http://www.w3.org/2000/svg" width="0" height="0"></svg>"#
            )?;
            return Ok(());
        }

        let (y0, y1) = data.values().copied().fold((f64::MAX, f64::MIN), |acc, v| {
            let min = if acc.0 < v { acc.0 } else { v };
            let max = if acc.1 > v { acc.1 } else { v };
            (min, max)
        });

        // The color scale range.
        let y_range = if self.log_scale {
            y1.log10() - y0.log10()
        } else {
            y1 - y0
        };

        let (width, height) = self.projection.dimensions();
        write!(
            io,
            r#"<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}">"#
        )?;

        let (start, end) = self.projection.range();
        for date in start.iter_days() {
            if date > end {
                break;
            }

            let (x, y) = self.projection.map(date);
            let value = data.get(&date).copied().unwrap_or(0.0);

            // Scale
            let scaled_value = if self.log_scale {
                (value.signum() * (1.0 + value.abs()).log10())
                    / (y1.signum() * (1.0 + y1.abs()).log10())
            } else {
                (value + y0) / y_range
            };

            let color = &self.colors[(scaled_value * (self.colors.len() - 1) as f64) as usize];

            writeln!(
                io,
                r#"<rect x="{x}" y="{y}" width="{0}" height="{0}" fill="{color}" rx="2" ry="2">"#,
                self.projection.size()
            )?;
            if let Some(ref unit) = self.unit {
                writeln!(io, "<title>{value} {unit} on {date}</title>")?;
            } else {
                writeln!(io, "<title>{value} {date}</title>")?;
            }
            writeln!(io, "</rect>")?;
        }
        writeln!(io, "</svg>")?;

        Ok(())
    }
}

/// Projection instances determine where each date goes in the image.
pub trait Projection {
    /// Transform a date to x,y coordinates.
    fn map(&self, date: NaiveDate) -> (u32, u32);

    /// Return the width and height of the image.
    fn dimensions(&self) -> (u32, u32);

    /// Return the size of each cell in the image.
    fn size(&self) -> u32;

    /// Return the range of dates in the image.
    fn range(&self) -> (NaiveDate, NaiveDate);
}

/// CalendarProjection creates a heatmap that looks like a calendar, with all
/// of the months in a row.
pub struct CalendarProjection {
    start: NaiveDate,
    end: NaiveDate,
    size: u32,
    padding: u32,
}

impl CalendarProjection {
    pub fn new(start: NaiveDate, end: NaiveDate, size: u32, padding: u32) -> Self {
        Self {
            start,
            end,
            size,
            padding,
        }
    }
}

impl Projection for CalendarProjection {
    fn map(&self, date: NaiveDate) -> (u32, u32) {
        let fom = date.with_day(1).unwrap();
        let year_number = (date.year() - self.start.year()) as u32;
        let week_number = (fom.weekday().num_days_from_monday() + date.day() - 1) / 7 + 1;

        let x = self.padding
            + ((date.month() - self.start.month()) % 12)
                * (self.padding + 7 * (self.size + self.padding))
            + date.weekday().num_days_from_monday() * (self.size + self.padding);

        let y = self.padding
            + week_number * (self.size + self.padding)
            + year_number * 6 * (self.size + self.padding);

        (x, y)
    }

    fn dimensions(&self) -> (u32, u32) {
        let years = (self.end.year() - self.start.year() + 1) as u32;
        let width = self.padding + ((self.size + self.padding) * 8) * 12;
        let height = self.padding + (self.size + self.padding) * 6 * years;
        (width, height)
    }

    fn size(&self) -> u32 {
        self.size
    }

    fn range(&self) -> (NaiveDate, NaiveDate) {
        (self.start, self.end)
    }
}

/// IsoProjection creates a denser heatmap where each column is a week of the
/// ISO calendar year and each row is a day of the week (starting with Monday).
pub struct IsoProjection {
    start: NaiveDate,
    end: NaiveDate,
    size: u32,
    padding: u32,
}

impl IsoProjection {
    pub fn new(start: NaiveDate, end: NaiveDate, size: u32, padding: u32) -> Self {
        Self {
            start,
            end,
            size,
            padding,
        }
    }
}

impl Projection for IsoProjection {
    fn map(&self, date: NaiveDate) -> (u32, u32) {
        let start = self.start.iso_week();
        let iso_date = date.iso_week();

        let year_number = (iso_date.year() - start.year()) as u32;

        let x = self.padding + iso_date.week0() * (self.size + self.padding);

        let y = self.padding
            + date.weekday().num_days_from_monday() * (self.size + self.padding)
            + year_number * 8 * (self.size + self.padding);

        (x, y)
    }

    fn dimensions(&self) -> (u32, u32) {
        let start = self.start.iso_week();
        let end = self.end.iso_week();
        let width = self.padding + (self.size + self.padding) * 53;
        let height =
            self.padding + (self.size + self.padding) * 8 * (end.year() - start.year() + 1) as u32;
        (width, height)
    }

    fn size(&self) -> u32 {
        self.size
    }

    fn range(&self) -> (NaiveDate, NaiveDate) {
        (self.start, self.end)
    }
}