models.rs

1use mrlycore::cell::Cell;2use mrlycore::colors::Color;3use mrlycore::enums::Mode;4use mrlycore::errors::{value_error, Result};5use mrlycore::tensor::{Dtype, Tensor};6use std::collections::HashMap;78/// The two-dimensional cell.9pub type Cell2d = CellNd<2>;10/// The three-dimensional cell.11pub type Cell3d = CellNd<3>;1213/// Returns the narrowest unsigned dtype that holds the peak value.14///15/// ```16/// use mrlycore::tensor::Dtype;17/// assert_eq!(mrlymath::dim::models::dtype_for(255), Dtype::U8);18/// assert_eq!(mrlymath::dim::models::dtype_for(256), Dtype::U16);19/// assert_eq!(mrlymath::dim::models::dtype_for(70000), Dtype::U32);20/// ```21pub fn dtype_for(peak: i64) -> Dtype {22    if peak <= Dtype::U8.max() {23        Dtype::U824    } else if peak <= Dtype::U16.max() {25        Dtype::U1626    } else {27        Dtype::U3228    }29}3031/// Returns the narrowest count dtype that fits the mask's popcount.32pub fn counting_dtype(mask: &Tensor) -> Dtype {33    dtype_for(mask.sum() as i64)34}3536/// A cell whose tensor is pinned to N dimensions.37#[derive(Clone, Debug, PartialEq, Eq)]38pub struct CellNd<const N: usize> {39    /// The wrapped cell.40    pub cell: Cell,41}4243impl<const N: usize> CellNd<N> {44    /// Builds a cell from an N-dimensional tensor of types.45    pub fn new(types: Tensor) -> CellNd<N> {46        assert_eq!(types.shape.len(), N, "CellNd requires a {N}d tensor");47        CellNd {48            cell: Cell::new(types),49        }50    }51    /// Returns the size of axis 1.52    pub fn width(&self) -> usize {53        self.cell.types.shape[1]54    }55    /// Returns the size of axis 0.56    pub fn height(&self) -> usize {57        self.cell.types.shape[0]58    }59    /// Returns the tensor of types.60    pub fn types(&self) -> &Tensor {61        &self.cell.types62    }63    /// Swaps filled and empty sites.64    pub fn invert(self) -> CellNd<N> {65        CellNd {66            cell: self.cell.invert(),67        }68    }69    /// Inverts the cell.70    pub fn anti(self) -> CellNd<N> {71        self.invert()72    }73    /// Wraps the cell in count layers of the given value on every side.74    pub fn pad(self, count: usize, value: u8) -> CellNd<N> {75        CellNd {76            cell: self.cell.pad(count, value),77        }78    }79    /// Deepens the cell into its level-fold fractal, or an error below level one.80    pub fn fractal(self, level: usize) -> Result<CellNd<N>> {81        Ok(CellNd {82            cell: self.cell.fractal(level)?,83        })84    }85    /// Tags each site with its ring distance from the center.86    pub fn layers(self) -> CellNd<N> {87        CellNd {88            cell: self.cell.layers(Dtype::U8),89        }90    }91    /// Tags each site with its count of masked neighbors matching the target, wrapping on request.92    pub fn neighbors(self, mask: &Tensor, target: u8, wrap: bool) -> Result<CellNd<N>> {93        let dtype = counting_dtype(mask);94        Ok(CellNd {95            cell: self.cell.neighbors(mask, target, wrap, dtype)?,96        })97    }98    /// Maps each site to one at or above the threshold, zero below.99    pub fn binarize(self, threshold: u8) -> CellNd<N> {100        CellNd {101            cell: self.cell.binarize(threshold),102        }103    }104    /// Binarizes the cell at the threshold Otsu's method picks.105    pub fn binarize_otsu(self) -> CellNd<N> {106        CellNd {107            cell: self.cell.binarize_otsu(),108        }109    }110    /// Rounds each site to the mean of its masked neighborhood, wrapping on request.111    pub fn blur(self, mask: &Tensor, wrap: bool) -> Result<CellNd<N>> {112        Ok(CellNd {113            cell: self.cell.blur(mask, wrap)?,114        })115    }116    /// Writes the value wherever the tiled mask is nonzero.117    pub fn perforate(self, mask: &Tensor, value: u8) -> Result<CellNd<N>> {118        Ok(CellNd {119            cell: self.cell.perforate(mask, value)?,120        })121    }122    /// Returns the Kronecker product of the two cells.123    pub fn combine(&self, other: &CellNd<N>) -> CellNd<N> {124        CellNd {125            cell: self.cell.combine(&other.cell),126        }127    }128    /// Colors each site by its type through the mapping in the given mode.129    pub fn paint(self, mapping: &HashMap<u8, Vec<Color>>, mode: Mode) -> CellNd<N> {130        CellNd {131            cell: self.cell.paint(mapping, mode),132        }133    }134}135136impl CellNd<2> {137    /// Rotates the cell k quarter turns in the plane.138    pub fn rotate(self, k: usize) -> Cell2d {139        CellNd {140            cell: self.cell.rotate(k, (0, 1)),141        }142    }143    /// Repeats the cell into a width-by-height array of copies.144    pub fn tile(self, width: usize, height: usize) -> Cell2d {145        CellNd {146            cell: self.cell.tile(&[height, width]),147        }148    }149}150151impl CellNd<3> {152    /// Returns the size of axis 2.153    pub fn depth(&self) -> usize {154        self.cell.types.shape[2]155    }156    /// Rotates the cell k quarter turns about the given pair of axes.157    pub fn rotate(self, k: usize, axes: (usize, usize)) -> Cell3d {158        CellNd {159            cell: self.cell.rotate(k, axes),160        }161    }162    /// Turns the cell into one of the 24 cube orientations, or an error past the table.163    pub fn orient(self, index: usize) -> Result<Cell3d> {164        let table = crate::three::orientations();165        match table.get(index) {166            Some(&(a, b, c)) => Ok(self.rotate(a, (1, 2)).rotate(b, (0, 2)).rotate(c, (0, 1))),167            None => value_error(format!("orientation index {index} out of range (0..23).")),168        }169    }170    /// Repeats the cell into a width-by-height-by-depth array of copies.171    pub fn tile(self, width: usize, height: usize, depth: usize) -> Cell3d {172        CellNd {173            cell: self.cell.tile(&[height, width, depth]),174        }175    }176}177178#[cfg(test)]179mod tests {180    use super::*;181    use mrlycore::atoms;182    #[test]183    fn binarize_wrapper_thresholds_pointwise() {184        let cell = Cell2d::new(atoms::carpet_2d(3));185        let binarized = cell.clone().binarize(1);186        assert_eq!(binarized.types(), cell.types());187    }188    #[test]189    fn blur_wrapper_preserves_shape() {190        let cell = Cell2d::new(atoms::carpet_2d(3));191        let mask = Tensor::full(vec![3, 3], 1);192        let blurred = cell.clone().blur(&mask, true).unwrap();193        assert_eq!(blurred.types().shape, cell.types().shape);194    }195    #[test]196    fn perforate_wrapper_zero_mask_is_identity() {197        let cell = Cell2d::new(atoms::carpet_2d(3));198        let mask = Tensor::new(cell.types().shape.clone());199        let perforated = cell.clone().perforate(&mask, 5).unwrap();200        assert_eq!(perforated.types(), cell.types());201    }202    #[test]203    fn blur_wrapper_preserves_shape_3d() {204        let cell = Cell3d::new(atoms::carpet_3d(3));205        let mask = Tensor::full(vec![3, 3, 3], 1);206        let blurred = cell.clone().blur(&mask, true).unwrap();207        assert_eq!(blurred.types().shape, cell.types().shape);208    }209    #[test]210    fn perforate_wrapper_zero_mask_is_identity_3d() {211        let cell = Cell3d::new(atoms::carpet_3d(3));212        let mask = Tensor::new(cell.types().shape.clone());213        let perforated = cell.clone().perforate(&mask, 5).unwrap();214        assert_eq!(perforated.types(), cell.types());215    }216    #[test]217    fn counting_dtype_widens_with_the_mask() {218        assert_eq!(counting_dtype(&Tensor::full(vec![3, 3], 1)), Dtype::U8);219        assert_eq!(counting_dtype(&Tensor::full(vec![15, 15], 1)), Dtype::U8);220        assert_eq!(counting_dtype(&Tensor::full(vec![17, 17], 1)), Dtype::U16);221        assert_eq!(counting_dtype(&Tensor::full(vec![63, 63], 1)), Dtype::U16);222    }223    #[test]224    fn neighbors_wrapper_survives_a_wide_mask() {225        let mut mask = Tensor::full(vec![17, 17], 1);226        mask.set(&[8, 8], 0);227        let grid = Cell2d::new(Tensor::full(vec![21, 21], 1));228        let counted = grid.neighbors(&mask, 1, true).unwrap();229        let tags = counted.cell.tags.as_ref().unwrap();230        assert_eq!(tags.at(0), 17 * 17 - 1);231    }232}