heatmap.rs

2.6 kB · rust · 83 lines

1use crate::two::{self, Cell2d};2use mrlycore::errors::{value_error, Result};3use mrlycore::ramp::Colorizer;4use mrlycore::tensor::Tensor;56/// Renders cumulative-visit heatmap frames over a slice of a run, or an error at a bad range.7pub fn heatmap_range(8    grids: &[Cell2d],9    start: usize,10    end: usize,11    colorizer: &Colorizer,12    scale: usize,13) -> Result<Vec<Vec<u8>>> {14    if grids.is_empty() {15        return Ok(Vec::new());16    }17    if start >= end || end > grids.len() {18        return value_error("heatmap range must satisfy 0 <= start < end <= frame count.");19    }20    let span = &grids[start..end];21    let shape = span[0].types().shape.clone();22    let size = span[0].types().size();23    let mut total = vec![0usize; size];24    for grid in span {25        for (i, &v) in grid.types().bytes().iter().enumerate() {26            total[i] += v as usize;27        }28    }29    let max = (*total.iter().max().unwrap_or(&0)).max(1);30    let mut cumulative = vec![0usize; size];31    let mut out = Vec::with_capacity(span.len());32    for grid in span {33        for (i, &v) in grid.types().bytes().iter().enumerate() {34            cumulative[i] += v as usize;35        }36        let colors = colorizer.colors(&cumulative, max);37        let mut cell = Cell2d::new(Tensor::new(shape.clone()));38        cell.cell.colors = Some(colors);39        out.push(two::png(&cell, scale)?);40    }41    Ok(out)42}4344/// Renders a whole run's cumulative-visit heatmap frames with the heat ramp.45pub fn heatmap(grids: &[Cell2d], scale: usize) -> Result<Vec<Vec<u8>>> {46    heatmap_range(grids, 0, grids.len(), &Colorizer::heat(), scale)47}4849#[cfg(test)]50mod tests {51    use super::*;52    use crate::life::{animate, moore, Boundary, Config};53    fn run() -> crate::life::Life {54        let config = Config {55            boundary: Boundary::Constant,56            max_generations: 8,57            ..Config::new(moore(), vec![3], vec![2, 3])58        };59        let mut t = Tensor::new(vec![5, 5]);60        t.set(&[1, 2], 1);61        t.set(&[2, 2], 1);62        t.set(&[3, 2], 1);63        animate(&Cell2d::new(t), &config).unwrap()64    }65    #[test]66    fn heatmap_frames_are_pngs() {67        let life = run();68        let pngs = heatmap(&life.grids, 4).unwrap();69        assert_eq!(pngs.len(), life.count);70        assert_eq!(&pngs[0][1..4], b"PNG");71    }72    #[test]73    fn range_renders_a_slice() {74        let life = run();75        let pngs = heatmap_range(&life.grids, 0, 1, &Colorizer::heat(), 4).unwrap();76        assert_eq!(pngs.len(), 1);77    }78    #[test]79    fn bad_range_errors() {80        let life = run();81        assert!(heatmap_range(&life.grids, 2, 1, &Colorizer::heat(), 4).is_err());82    }83}