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}