cell.rs

18.2 kB · rust · 518 lines

1use super::colors::{Color, ALPHA, BLACK, BLUE, GREEN, RED, WHITE};2use super::enums::Mode;3use super::errors::{value_error, Result};4use super::state;5use super::tensor::{Dtype, Tensor};6use std::collections::HashMap;78/// A grid of type bytes with optional per-cell colors and tags.9#[derive(Clone, Debug, PartialEq, Eq)]10pub struct Cell {11    /// The type of every cell.12    pub types: Tensor,13    /// The painted rgba of every cell, once painted.14    pub colors: Option<Vec<[u8; 4]>>,15    /// The tag layer over the cells, once built.16    pub tags: Option<Tensor>,17}1819/// Returns the default mapping of the first six types to white, black, alpha, red, green, and blue.20pub fn mapping() -> HashMap<u8, Vec<Color>> {21    HashMap::from([22        (0, vec![WHITE]),23        (1, vec![BLACK]),24        (2, vec![ALPHA]),25        (3, vec![RED]),26        (4, vec![GREEN]),27        (5, vec![BLUE]),28    ])29}3031impl Cell {32    /// Wraps a tensor of types in a bare cell, colorless and tagless.33    pub fn new(types: Tensor) -> Cell {34        Cell {35            types,36            colors: None,37            tags: None,38        }39    }40    /// Returns the shape of the type tensor.41    pub fn shape(&self) -> &[usize] {42        &self.types.shape43    }44    /// Returns the number of cells.45    pub fn size(&self) -> usize {46        self.types.size()47    }48    /// Returns the painted color at a flat index, or transparent while unpainted.49    pub fn color_at(&self, flat: usize) -> [u8; 4] {50        match &self.colors {51            Some(colors) => colors[flat],52            None => [0, 0, 0, 0],53        }54    }55    /// Flips every type to one minus itself.56    pub fn invert(mut self) -> Cell {57        self.types = self.types.invert();58        self59    }60    /// Flips every type to one minus itself, same as invert.61    pub fn anti(self) -> Cell {62        self.invert()63    }64    /// Wraps the cell in count layers of value on every side, dropping colors.65    pub fn pad(mut self, count: usize, value: u8) -> Cell {66        self.types = self.types.pad(count, value);67        self.colors = None;68        self.tags = self.tags.map(|t| t.pad(count, value));69        self70    }71    /// Rotates the cell k quarter turns in the plane of the given axes, carrying colors and tags along.72    pub fn rotate(mut self, k: usize, axes: (usize, usize)) -> Cell {73        if let Some(colors) = &self.colors {74            let map = rot90_map(&self.types.shape, k, axes);75            self.colors = Some(map.iter().map(|&src| colors[src]).collect());76        }77        self.types = self.types.rot90(k, axes);78        self.tags = self.tags.map(|t| t.rot90(k, axes));79        self80    }81    /// Grows the types to the level-fold Kronecker power of themselves, dropping colors and tags.82    pub fn fractal(mut self, level: usize) -> Result<Cell> {83        if level < 1 {84            return value_error("Fractal level must be at least 1.");85        }86        self.types = self.types.fractal(level);87        self.colors = None;88        self.tags = None;89        Ok(self)90    }91    /// Repeats the cell reps times along each axis, carrying colors and tags along.92    pub fn tile(self, reps: &[usize]) -> Cell {93        let shape: Vec<usize> = self94            .types95            .shape96            .iter()97            .zip(reps)98            .map(|(n, r)| n * r)99            .collect();100        remap(&self, &tile_map(&self.types.shape, reps), &shape)101    }102    /// Tags every cell with its concentric shell distance from the center.103    pub fn layers(mut self, dtype: Dtype) -> Cell {104        self.tags = Some(self.types.layers(dtype));105        self106    }107    /// Tags every cell with its count of target-valued neighbors under the mask.108    pub fn neighbors(109        mut self,110        mask: &Tensor,111        target: u8,112        wrap: bool,113        dtype: Dtype,114    ) -> Result<Cell> {115        self.tags = Some(self.types.neighbors(mask, target, wrap, dtype)?);116        Ok(self)117    }118    /// Maps every type to one at or above the threshold and zero below, dropping colors.119    pub fn binarize(mut self, threshold: u8) -> Cell {120        self.types = self.types.binarize(threshold);121        self.colors = None;122        self123    }124    /// Binarizes the types at Otsu's threshold, dropping colors.125    pub fn binarize_otsu(mut self) -> Cell {126        self.types = self.types.binarize_otsu();127        self.colors = None;128        self129    }130    /// Replaces every type with the rounded mean of its masked neighborhood, dropping colors.131    pub fn blur(mut self, mask: &Tensor, wrap: bool) -> Result<Cell> {132        self.types = self.types.blur(mask, wrap)?;133        self.colors = None;134        Ok(self)135    }136    /// Stamps value wherever the tiled mask is on, dropping colors.137    pub fn perforate(mut self, mask: &Tensor, value: u8) -> Result<Cell> {138        self.types = self.types.perforate(mask, value)?;139        self.colors = None;140        Ok(self)141    }142    /// Builds the Kronecker product of the two cells' types.143    pub fn combine(&self, other: &Cell) -> Cell {144        Cell::new(self.types.kron(&other.types))145    }146    /// Colors every mapped cell, picking within each type's palette by the mode.147    pub fn paint(mut self, mapping: &HashMap<u8, Vec<Color>>, mode: Mode) -> Cell {148        let size = self.size();149        let mut colors = self150            .colors151            .take()152            .unwrap_or_else(|| vec![[0, 0, 0, 0]; size]);153        let mut keys: Vec<u8> = mapping.keys().copied().collect();154        keys.sort_unstable();155        for key in keys {156            let rgba: Vec<[u8; 4]> = mapping[&key].iter().map(|c| [c.r, c.g, c.b, c.a]).collect();157            if rgba.is_empty() {158                continue;159            }160            let mut enumerated = 0;161            for (flat, &t) in self.types.bytes().iter().enumerate() {162                if t != key {163                    continue;164                }165                let pick = match mode {166                    Mode::Type => 0,167                    Mode::Random => state::randint(0, rgba.len() as i64 - 1) as usize,168                    Mode::Enumerate => {169                        let i = enumerated;170                        enumerated += 1;171                        i % rgba.len()172                    }173                    Mode::Index => flat % rgba.len(),174                    Mode::Tag => match &self.tags {175                        Some(tags) => tags.at(flat) as usize % rgba.len(),176                        None => 0,177                    },178                    Mode::Row | Mode::Column | Mode::Depth => {179                        let axis = match mode {180                            Mode::Row => 0,181                            Mode::Column => 1,182                            _ => 2,183                        };184                        if axis < self.types.shape.len() {185                            axis_index(&self.types, flat, axis) % rgba.len()186                        } else {187                            0188                        }189                    }190                };191                colors[flat] = rgba[pick];192            }193        }194        self.colors = Some(colors);195        self196    }197}198199fn axis_index(t: &Tensor, flat: usize, axis: usize) -> usize {200    let mut stride = 1;201    for a in (axis + 1)..t.shape.len() {202        stride *= t.shape[a];203    }204    (flat / stride) % t.shape[axis]205}206207/// Builds the flat source index of every destination cell after tiling reps copies per axis.208pub fn tile_map(shape: &[usize], reps: &[usize]) -> Vec<usize> {209    let tiled: Vec<usize> = shape.iter().zip(reps).map(|(n, r)| n * r).collect();210    let size = tiled.iter().product();211    let mut map = Vec::with_capacity(size);212    for flat in 0..size {213        let mut rem = flat;214        let mut source = 0;215        for (axis, &n) in shape.iter().enumerate() {216            let stride: usize = tiled[(axis + 1)..].iter().product();217            let i = rem / stride;218            rem %= stride;219            source = source * n + i % n;220        }221        map.push(source);222    }223    map224}225226/// Rebuilds a cell's types, colors and tags at the new shape from one destination-to-source index map.227///228/// The map holds one source index per destination cell, so it must be as long as the shape's size.229pub fn remap(cell: &Cell, map: &[usize], shape: &[usize]) -> Cell {230    Cell {231        types: gather(&cell.types, map, shape),232        colors: cell233            .colors234            .as_ref()235            .map(|colors| map.iter().map(|&src| colors[src]).collect()),236        tags: cell.tags.as_ref().map(|tags| gather(tags, map, shape)),237    }238}239240fn gather(source: &Tensor, map: &[usize], shape: &[usize]) -> Tensor {241    let mut out = Tensor::typed(shape.to_vec(), source.dtype());242    for (flat, &src) in map.iter().enumerate() {243        out.put(flat, source.at(src));244    }245    out246}247248/// Builds the 3-wide Moore mask of the dimension, every site on but the center.249///250/// ```251/// let mask = mrlycore::cell::moore(2);252/// assert_eq!(mask.shape, vec![3, 3]);253/// assert_eq!(mask.sum(), 8);254/// ```255pub fn moore(dimension: usize) -> Tensor {256    let mut mask = Tensor::full(vec![3; dimension], 1);257    let mut center = 0;258    for _ in 0..dimension {259        center = center * 3 + 1;260    }261    mask.bytes_mut()[center] = 0;262    mask263}264265/// Builds the flat source index of every destination cell after k quarter turns in the plane of the axes.266pub fn rot90_map(shape: &[usize], k: usize, axes: (usize, usize)) -> Vec<usize> {267    let mut data: Vec<usize> = (0..shape.iter().product()).collect();268    let mut shape = shape.to_vec();269    for _ in 0..k % 4 {270        let (a, b) = axes;271        let mut next_shape = shape.clone();272        next_shape.swap(a, b);273        let mut next = vec![0; data.len()];274        for (flat, item) in next.iter_mut().enumerate() {275            let mut rem = flat;276            let mut multi = Vec::with_capacity(next_shape.len());277            for axis in 0..next_shape.len() {278                let stride: usize = next_shape[(axis + 1)..].iter().product();279                multi.push(rem / stride);280                rem %= stride;281            }282            multi[a] = next_shape[a] - 1 - multi[a];283            multi.swap(a, b);284            let mut source = 0;285            for axis in 0..shape.len() {286                source = source * shape[axis] + multi[axis];287            }288            *item = data[source];289        }290        data = next;291        shape = next_shape;292    }293    data294}295296/// Stitches same-shaped cells into one grid of reps blocks per axis, or an error when counts or shapes disagree.297pub fn merge(cells: &[Cell], reps: &[usize]) -> Result<Cell> {298    if cells.is_empty() {299        return value_error("Cannot merge an empty list of cells.");300    }301    let count: usize = reps.iter().product();302    if cells.len() != count {303        return value_error(format!("Expected {count} cells, got {}", cells.len()));304    }305    let inner = cells[0].types.shape.clone();306    for cell in cells {307        if cell.types.shape != inner {308            return value_error("All cells in a merge operation must have the same dimensions.");309        }310    }311    let shape: Vec<usize> = inner.iter().zip(reps).map(|(n, r)| n * r).collect();312    let mut out = Tensor::new(shape.clone());313    let dims = shape.len();314    for flat in 0..out.size() {315        let mut rem = flat;316        let mut block = 0;317        let mut local = Vec::with_capacity(dims);318        let mut block_multi = Vec::with_capacity(dims);319        for (axis, &inner_n) in inner.iter().enumerate() {320            let stride: usize = shape[(axis + 1)..].iter().product();321            let i = rem / stride;322            rem %= stride;323            block_multi.push(i / inner_n);324            local.push(i % inner_n);325        }326        for (axis, &b) in block_multi.iter().enumerate() {327            block = block * reps[axis] + b;328        }329        out.bytes_mut()[flat] = cells[block].types.get(&local);330    }331    Ok(Cell::new(out))332}333334/// Folds at least two cells into one by chained Kronecker products.335pub fn magic(cells: &[Cell]) -> Result<Cell> {336    if cells.len() < 2 {337        return value_error("Magic composition requires at least two cells.");338    }339    let mut out = cells[0].combine(&cells[1]);340    for cell in &cells[2..] {341        out = out.combine(cell);342    }343    Ok(out)344}345346/// Lays the cell each mask entry indexes into that entry's place and merges the lot.347pub fn mosaic(mask: &Tensor, cells: &[Cell]) -> Result<Cell> {348    let picked: Result<Vec<Cell>> = mask349        .bytes()350        .iter()351        .map(|&i| match cells.get(i as usize) {352            Some(cell) => Ok(cell.clone()),353            None => value_error(format!("mosaic index {i} out of range.")),354        })355        .collect();356    merge(&picked?, &mask.shape)357}358359#[cfg(test)]360mod tests {361    use super::*;362    use crate::atoms;363    use crate::state::{guard, seed};364    #[test]365    fn rot90_map_matches_tensor() {366        let t = Tensor::of((0..24).map(|v| v as u8).collect(), vec![2, 3, 4]);367        for axes in [(0, 1), (0, 2), (1, 2)] {368            for k in 0..5 {369                let rotated = t.rot90(k, axes);370                let map = rot90_map(&t.shape, k, axes);371                let mapped: Vec<u8> = map.iter().map(|&s| t.bytes()[s]).collect();372                assert_eq!(mapped, rotated.bytes());373            }374        }375    }376    #[test]377    fn remap_carries_types_colors_and_tags() {378        let painted = Cell::new(atoms::carpet_2d(3))379            .layers(Dtype::U8)380            .paint(&mapping(), Mode::Type);381        let map = rot90_map(&painted.types.shape, 1, (0, 1));382        let turned = remap(&painted, &map, &[3, 3]);383        assert_eq!(turned.types, painted.types.rot90(1, (0, 1)));384        assert_eq!(385            turned.tags.as_ref().unwrap(),386            &painted.tags.as_ref().unwrap().rot90(1, (0, 1))387        );388        let colors = turned.colors.as_ref().unwrap();389        let source = painted.colors.as_ref().unwrap();390        for (flat, &src) in map.iter().enumerate() {391            assert_eq!(colors[flat], source[src], "at {flat}");392        }393    }394    #[test]395    fn remap_keeps_a_wide_tag_layer_wide() {396        let mut grid = Cell::new(atoms::ones_2d(2));397        grid.tags = Some(Tensor::filled(vec![2, 2], 300, Dtype::U16));398        let tiled = grid.tile(&[2, 2]);399        let tags = tiled.tags.unwrap();400        assert_eq!(tags.dtype(), Dtype::U16);401        assert_eq!(tags.at(15), 300);402    }403    #[test]404    fn moore_masks_every_site_but_the_center() {405        let flat = moore(2);406        assert_eq!(flat.shape, vec![3, 3]);407        assert_eq!(flat.get(&[1, 1]), 0);408        assert_eq!(flat.sum(), 8);409        let cube = moore(3);410        assert_eq!(cube.shape, vec![3, 3, 3]);411        assert_eq!(cube.get(&[1, 1, 1]), 0);412        assert_eq!(cube.sum(), 26);413    }414    #[test]415    fn merge_two_by_two() {416        let a = Cell::new(atoms::ones_2d(2));417        let b = Cell::new(atoms::zeros_2d(2));418        let m = merge(&[a.clone(), b.clone(), b, a], &[2, 2]).unwrap();419        assert_eq!(m.types.shape, vec![4, 4]);420        assert_eq!(m.types.sum(), 8);421        assert_eq!(m.types.get(&[0, 0]), 1);422        assert_eq!(m.types.get(&[0, 2]), 0);423        assert_eq!(m.types.get(&[2, 0]), 0);424        assert_eq!(m.types.get(&[3, 3]), 1);425    }426    #[test]427    fn mosaic_picks_cells() {428        let a = Cell::new(atoms::ones_2d(2));429        let b = Cell::new(atoms::zeros_2d(2));430        let mask = Tensor::of(vec![0, 1, 1, 0], vec![2, 2]);431        let m = mosaic(&mask, &[a, b]).unwrap();432        assert_eq!(m.types.sum(), 8);433        assert_eq!(m.types.get(&[0, 0]), 1);434        assert_eq!(m.types.get(&[0, 2]), 0);435    }436    #[test]437    fn paint_type_mode() {438        let cell = Cell::new(atoms::carpet_2d(3)).paint(&mapping(), Mode::Type);439        let colors = cell.colors.as_ref().unwrap();440        assert_eq!(colors[0], [0, 0, 0, 255]);441        assert_eq!(colors[4], [255, 255, 255, 255]);442        let dark = cell443            .types444            .bytes()445            .iter()446            .zip(colors)447            .filter(|(&t, _)| t == 1)448            .count();449        assert_eq!(dark, 8);450    }451    #[test]452    fn tile_carries_colors_and_tags() {453        let painted = Cell::new(atoms::carpet_2d(3))454            .layers(Dtype::U8)455            .paint(&mapping(), Mode::Type);456        let tiled = painted.clone().tile(&[2, 3]);457        assert_eq!(tiled.types.shape, vec![6, 9]);458        let colors = tiled.colors.as_ref().unwrap();459        let source = painted.colors.as_ref().unwrap();460        assert_eq!(colors.len(), 54);461        for y in 0..6 {462            for x in 0..9 {463                assert_eq!(colors[y * 9 + x], source[(y % 3) * 3 + x % 3]);464            }465        }466        assert_eq!(tiled.tags.as_ref().unwrap().shape, vec![6, 9]);467    }468    #[test]469    fn paint_random_mode_is_seed_stable() {470        let _g = guard();471        seed(5);472        let types = Tensor::of(vec![0, 1, 1, 0], vec![2, 2]);473        let forward = HashMap::from([(0, vec![RED, GREEN]), (1, vec![BLUE, WHITE])]);474        let colors = Cell::new(types.clone())475            .paint(&forward, Mode::Random)476            .colors477            .unwrap();478        seed(5);479        let reversed = HashMap::from([(1, vec![BLUE, WHITE]), (0, vec![RED, GREEN])]);480        let again = Cell::new(types)481            .paint(&reversed, Mode::Random)482            .colors483            .unwrap();484        assert_eq!(colors, again);485        let pinned = [486            [255, 61, 64, 255],487            [255, 255, 255, 255],488            [255, 255, 255, 255],489            [50, 204, 88, 255],490        ];491        assert_eq!(colors, pinned);492    }493    #[test]494    fn magic_is_kron_chain() {495        let a = Cell::new(atoms::carpet_2d(2));496        let b = Cell::new(atoms::ones_2d(3));497        let m = magic(&[a.clone(), b]).unwrap();498        assert_eq!(m.types.shape, vec![6, 6]);499        assert_eq!(m.types.sum(), a.types.sum() * 9);500    }501    #[test]502    fn binarize_clears_colors_and_thresholds() {503        let cell = Cell::new(atoms::carpet_2d(3))504            .paint(&mapping(), Mode::Type)505            .binarize(1);506        assert!(cell.colors.is_none());507        assert_eq!(cell.types.bytes(), atoms::carpet_2d(3).bytes());508    }509    #[test]510    fn blur_and_perforate_wrappers_delegate_to_tensor() {511        let cell = Cell::new(atoms::carpet_2d(3));512        let mask = Tensor::full(vec![3, 3], 1);513        let blurred = cell.clone().blur(&mask, true).unwrap();514        assert_eq!(blurred.types.shape, cell.types.shape);515        let perforated = cell.clone().perforate(&Tensor::new(vec![3, 3]), 9).unwrap();516        assert_eq!(perforated.types, cell.types);517    }518}