models.rs

1use crate::core::cell::Cell;2use crate::core::cell::Mode;3use crate::core::colors::Color;4use crate::core::error::{shape_error, value_error, Result};5use crate::core::tensor::{Dtype, Tensor};6use serde::{Deserialize, Serialize};7use std::collections::HashMap;89/// The two-dimensional cell.10pub type Cell2d = CellNd<2>;11/// The three-dimensional cell.12pub type Cell3d = CellNd<3>;1314/// Returns the narrowest unsigned dtype that holds the peak value.15///16/// ```17/// use mrlyrs::core::tensor::Dtype;18/// assert_eq!(mrlyrs::math::cell::models::dtype_for(255), Dtype::U8);19/// assert_eq!(mrlyrs::math::cell::models::dtype_for(256), Dtype::U16);20/// assert_eq!(mrlyrs::math::cell::models::dtype_for(70000), Dtype::U32);21/// ```22pub fn dtype_for(peak: i64) -> Dtype {23    if peak <= Dtype::U8.max() {24        Dtype::U825    } else if peak <= Dtype::U16.max() {26        Dtype::U1627    } else {28        Dtype::U3229    }30}3132/// Returns the narrowest count dtype that fits the mask's popcount.33pub fn counting_dtype(mask: &Tensor) -> Dtype {34    dtype_for(mask.sum() as i64)35}3637/// A cell whose tensor is pinned to N dimensions.38#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]39pub struct CellNd<const N: usize> {40    /// The wrapped cell.41    pub cell: Cell,42}4344impl<const N: usize> CellNd<N> {45    /// Builds a cell from an N-dimensional tensor of types.46    ///47    /// # Errors48    ///49    /// Errors when the tensor's rank is not N.50    pub fn new(types: Tensor) -> Result<CellNd<N>> {51        if types.shape.len() != N {52            return shape_error(format!(53                "a {N}d cell needs a {N}d tensor, got {}d.",54                types.shape.len()55            ));56        }57        Ok(CellNd {58            cell: Cell::new(types),59        })60    }61    /// Returns the size of axis 1.62    pub fn width(&self) -> usize {63        self.cell.types.shape[1]64    }65    /// Returns the size of axis 0.66    pub fn height(&self) -> usize {67        self.cell.types.shape[0]68    }69    /// Returns the tensor of types.70    pub fn types(&self) -> &Tensor {71        &self.cell.types72    }73    /// Swaps filled and empty sites.74    pub fn invert(self) -> CellNd<N> {75        CellNd {76            cell: self.cell.invert(),77        }78    }79    /// Inverts the cell.80    pub fn anti(self) -> CellNd<N> {81        self.invert()82    }83    /// Wraps the cell in count layers of the given value on every side.84    pub fn pad(self, count: usize, value: u8) -> CellNd<N> {85        CellNd {86            cell: self.cell.pad(count, value),87        }88    }89    /// Deepens the cell into its level-fold fractal.90    ///91    /// # Errors92    ///93    /// Errors below level one.94    pub fn fractal(self, level: usize) -> Result<CellNd<N>> {95        Ok(CellNd {96            cell: self.cell.fractal(level)?,97        })98    }99    /// Tags each site with its ring distance from the center.100    pub fn layers(self) -> CellNd<N> {101        CellNd {102            cell: self.cell.layers(Dtype::U8),103        }104    }105    /// Tags each site with its count of masked neighbors matching the target, wrapping on request.106    ///107    /// # Errors108    ///109    /// Errors when the mask does not match the cell's rank.110    pub fn neighbors(self, mask: &Tensor, target: u8, wrap: bool) -> Result<CellNd<N>> {111        let dtype = counting_dtype(mask);112        Ok(CellNd {113            cell: self.cell.neighbors(mask, target, wrap, dtype)?,114        })115    }116    /// Maps each site to one at or above the threshold, zero below.117    pub fn binarize(self, threshold: u8) -> CellNd<N> {118        CellNd {119            cell: self.cell.binarize(threshold),120        }121    }122    /// Binarizes the cell at the threshold Otsu's method picks.123    pub fn binarize_otsu(self) -> CellNd<N> {124        CellNd {125            cell: self.cell.binarize_otsu(),126        }127    }128    /// Rounds each site to the mean of its masked neighborhood, wrapping on request.129    ///130    /// # Errors131    ///132    /// Errors when the mask does not match the cell's rank.133    pub fn blur(self, mask: &Tensor, wrap: bool) -> Result<CellNd<N>> {134        Ok(CellNd {135            cell: self.cell.blur(mask, wrap)?,136        })137    }138    /// Writes the value wherever the tiled mask is nonzero.139    ///140    /// # Errors141    ///142    /// Errors when the mask does not tile the cell.143    pub fn perforate(self, mask: &Tensor, value: u8) -> Result<CellNd<N>> {144        Ok(CellNd {145            cell: self.cell.perforate(mask, value)?,146        })147    }148    /// Returns the Kronecker product of the two cells.149    pub fn combine(&self, other: &CellNd<N>) -> CellNd<N> {150        CellNd {151            cell: self.cell.combine(&other.cell),152        }153    }154    /// Colors each site by its type through the mapping in the given mode.155    pub fn paint(self, mapping: &HashMap<u8, Vec<Color>>, mode: Mode) -> CellNd<N> {156        CellNd {157            cell: self.cell.paint(mapping, mode),158        }159    }160}161162impl CellNd<2> {163    /// Rotates the cell k quarter turns in the plane.164    ///165    /// # Errors166    ///167    /// Errors for a cell without two axes.168    pub fn rotate(self, k: usize) -> Result<Cell2d> {169        Ok(CellNd {170            cell: self.cell.rotate(k, (0, 1))?,171        })172    }173    /// Repeats the cell into a width-by-height array of copies.174    ///175    /// # Errors176    ///177    /// Errors for a cell without two axes.178    pub fn tile(self, width: usize, height: usize) -> Result<Cell2d> {179        Ok(CellNd {180            cell: self.cell.tile(&[height, width])?,181        })182    }183}184185impl CellNd<3> {186    /// Returns the size of axis 2.187    pub fn depth(&self) -> usize {188        self.cell.types.shape[2]189    }190    /// Rotates the cell k quarter turns about the given pair of axes.191    ///192    /// # Errors193    ///194    /// Errors for axes the cell does not hold.195    pub fn rotate(self, k: usize, axes: (usize, usize)) -> Result<Cell3d> {196        Ok(CellNd {197            cell: self.cell.rotate(k, axes)?,198        })199    }200    /// Turns the cell into one of the 24 cube orientations.201    ///202    /// # Errors203    ///204    /// Errors past orientation twenty-three.205    pub fn orient(self, index: usize) -> Result<Cell3d> {206        let table = crate::math::three::orientations();207        match table.get(index) {208            Some(&(a, b, c)) => self.rotate(a, (1, 2))?.rotate(b, (0, 2))?.rotate(c, (0, 1)),209            None => value_error(format!("orientation index {index} out of range (0..23).")),210        }211    }212    /// Repeats the cell into a width-by-height-by-depth array of copies.213    ///214    /// # Errors215    ///216    /// Errors for a cell without three axes.217    pub fn tile(self, width: usize, height: usize, depth: usize) -> Result<Cell3d> {218        Ok(CellNd {219            cell: self.cell.tile(&[height, width, depth])?,220        })221    }222}223224#[cfg(test)]225mod tests {226    use super::*;227    use crate::math::atoms;228    #[test]229    fn binarize_wrapper_thresholds_pointwise() {230        let cell = Cell2d::new(atoms::carpet_2d(3)).unwrap();231        let binarized = cell.clone().binarize(1);232        assert_eq!(binarized.types(), cell.types());233    }234    #[test]235    fn blur_keeps_the_shape_at_every_rank() {236        let flat = Cell2d::new(atoms::carpet_2d(3)).unwrap();237        let blurred = flat238            .clone()239            .blur(&Tensor::full(vec![3, 3], 1), true)240            .unwrap();241        assert_eq!(blurred.types().shape, flat.types().shape);242        let cube = Cell3d::new(atoms::carpet_3d(3)).unwrap();243        let blurred = cube244            .clone()245            .blur(&Tensor::full(vec![3, 3, 3], 1), true)246            .unwrap();247        assert_eq!(blurred.types().shape, cube.types().shape);248    }249    #[test]250    fn a_zero_mask_perforates_nothing_at_every_rank() {251        let flat = Cell2d::new(atoms::carpet_2d(3)).unwrap();252        let mask = Tensor::new(flat.types().shape.clone());253        assert_eq!(254            flat.clone().perforate(&mask, 5).unwrap().types(),255            flat.types()256        );257        let cube = Cell3d::new(atoms::carpet_3d(3)).unwrap();258        let mask = Tensor::new(cube.types().shape.clone());259        assert_eq!(260            cube.clone().perforate(&mask, 5).unwrap().types(),261            cube.types()262        );263    }264    #[test]265    fn counting_dtype_widens_with_the_mask() {266        assert_eq!(counting_dtype(&Tensor::full(vec![3, 3], 1)), Dtype::U8);267        assert_eq!(counting_dtype(&Tensor::full(vec![15, 15], 1)), Dtype::U8);268        assert_eq!(counting_dtype(&Tensor::full(vec![17, 17], 1)), Dtype::U16);269        assert_eq!(counting_dtype(&Tensor::full(vec![63, 63], 1)), Dtype::U16);270    }271    #[test]272    fn neighbors_wrapper_survives_a_wide_mask() {273        let mut mask = Tensor::full(vec![17, 17], 1);274        mask.set(&[8, 8], 0).unwrap();275        let grid = Cell2d::new(Tensor::full(vec![21, 21], 1)).unwrap();276        let counted = grid.neighbors(&mask, 1, true).unwrap();277        let tags = counted.cell.tags.as_ref().unwrap();278        assert_eq!(tags.at(0), 17 * 17 - 1);279    }280281    #[test]282    fn refuses_a_tensor_of_the_wrong_rank() {283        assert!(Cell2d::new(Tensor::new(vec![2, 2, 2])).is_err());284        assert!(Cell3d::new(Tensor::new(vec![2, 2])).is_err());285        assert!(CellNd::<1>::new(Tensor::new(vec![2, 2])).is_err());286    }287}