geometry.py

1import numpy as np2from copy import deepcopy3from typing import List, TYPE_CHECKING4from mrlypy.core.errors import MrlyError56if TYPE_CHECKING:7    from .models import Cell2d89# PUBLIC - IMMUTABLE1011def merge_2d(cells: List["Cell2d"], width: int, height: int) -> "Cell2d":12    from .models import Cell2d13    if not cells:14        raise MrlyError("Cannot merge an empty list of cells.")15    if len(cells) != width * height:16        raise MrlyError(f"Expected {width * height} cells, got {len(cells)}")17    first_cell = cells[0]18    cell_width, cell_height = first_cell.width, first_cell.height19    total_width = width * cell_width20    total_height = height * cell_height21    new_cell = Cell2d(width=total_width, height=total_height)22    for i, cell in enumerate(cells):23        if cell.width != cell_width or cell.height != cell_height:24            raise MrlyError("All cells in a merge operation must have the same dimensions.")25        x = i % width26        y = i // width27        start_x = x * cell_width28        end_x = start_x + cell_width29        start_y = y * cell_height30        end_y = start_y + cell_height31        if cell._types is not None:32            new_cell.types[start_y:end_y, start_x:end_x] = cell.types33        if cell._colors is not None:34            new_cell.colors[start_y:end_y, start_x:end_x] = cell.colors35        if cell._tags is not None:36            new_cell.tags[start_y:end_y, start_x:end_x] = cell.tags37    return new_cell3839def combine_2d(cell_1: "Cell2d", cell_2: "Cell2d") -> "Cell2d":40    from .models import Cell2d41    new_types = np.kron(cell_1.types, cell_2.types).astype(np.uint8)42    return Cell2d(types=new_types)4344def magic_2d(cells: List["Cell2d"]) -> "Cell2d":45    if len(cells) < 2:46        raise MrlyError("Magic composition requires at least two cells.")47    new_cell = combine_2d(cells[0], cells[1])48    for i in range(2, len(cells)):49        new_cell = combine_2d(new_cell, cells[i])50    return new_cell5152def special_2d(mask: np.ndarray, cell: "Cell2d") -> "Cell2d":53    height, width = mask.shape54    new_cells = []55    for y in range(height):56        for x in range(width):57            rotation_index = mask[y, x]58            if not (0 <= rotation_index <= 3):59                raise MrlyError(f"Invalid rotation value '{rotation_index}'. Must be 0, 1, 2, or 3.")60            new_cell = cell.copy().rotate(rotation_index)61            new_cells.append(new_cell)62    return merge_2d(new_cells, width, height)6364def mosaic_2d(mask: np.ndarray, cells: List["Cell2d"]) -> "Cell2d":65    height, width = mask.shape66    new_cells = []67    for y in range(height):68        for x in range(width):69            cell_index = mask[y, x]70            new_cell = cells[cell_index].copy()71            new_cells.append(new_cell)72    return merge_2d(new_cells, width, height)7374# PRIVATE - MUTABLE7576def invert_2d(cell: "Cell2d") -> "Cell2d":77    if cell._types is not None:78        cell.types = 1 - cell.types79    return cell8081def pad_2d(cell: "Cell2d", count: int = 1, value: int = 0) -> "Cell2d":82    if cell._types is not None:83        cell.types = np.pad(cell.types, count, mode="constant", constant_values=value)84    if cell._colors is not None:85        cell.colors = np.pad(cell.colors, count, mode="constant", constant_values=value)86    if cell._tags is not None:87        cell.tags = np.pad(cell.tags, count, mode="constant", constant_values=value)88    return cell8990def rotate_2d(cell: "Cell2d", k: int = 1) -> "Cell2d":91    if k % 4 == 0:92        return cell93    if cell._types is not None:94        cell.types = np.rot90(cell.types, k)95    if cell._colors is not None:96        cell.colors = np.rot90(cell.colors, k, axes=(1, 0))97    if cell._tags is not None:98        cell.tags = np.rot90(cell.tags, k)99    return cell100101def fractal_2d(cell: "Cell2d", level: int) -> "Cell2d":102    if level < 1:103        raise MrlyError("Fractal level must be at least 1.")104    if level == 1:105        return cell106    new_types = cell.types107    for _ in range(1, level):108        new_types = np.kron(new_types, cell.types)109    cell.types = new_types.astype(np.uint8)110    cell.colors = None111    cell.tags = None112    return cell113114def tile_2d(cell: "Cell2d", width: int, height: int) -> "Cell2d":115    if cell._types is not None:116        cell.types = np.tile(cell.types, (height, width))117    if cell._colors is not None:118        cell.colors = np.tile(cell.colors, (height, width, 1))119    if cell._tags is not None:120        cell.tags = np.tile(cell.tags, (height, width))121    return cell122123def layers_2d(cell: "Cell2d", dtype: np.dtype = np.dtype(np.uint8)) -> "Cell2d":124    height, width = cell.height, cell.width125    y_indices, x_indices = np.indices((height, width))126    center_y = (height - 1) / 2127    center_x = (width - 1) / 2128    distance_y = np.floor(np.abs(y_indices - center_y))129    distance_x = np.floor(np.abs(x_indices - center_x))130    tags = np.maximum(distance_x, distance_y)131    cell.tags = tags.astype(dtype)132    return cell133134def neighbors_2d(cell: "Cell2d", mask: np.ndarray, target: int = 1, mode: str = "constant", dtype: np.dtype = np.dtype(np.uint8)) -> "Cell2d":135    mask_height, mask_width = mask.shape136    if mask_height % 2 == 0 or mask_width % 2 == 0:137        raise MrlyError("Neighborhood (mask) dimensions must be odd.")138    if target not in [0, 1]:139        raise MrlyError("Bit to count (target) must be 0 or 1.")140    bit = (cell.types == target).astype(dtype)141    if mode not in ["constant", "wrap"]:142        raise MrlyError("Boundary (mode) must be 'constant' or 'wrap'.")143    py = mask_height // 2144    px = mask_width // 2145    pad_bits = np.pad(bit, pad_width=((py, py), (px, px)), mode=mode)146    neighbor_counts = np.zeros_like(cell.types, dtype=dtype)147    for r in range(mask_height):148        for c in range(mask_width):149            if mask[r, c] == 1:150                start_row, end_row = r, r + cell.types.shape[0]151                start_col, end_col = c, c + cell.types.shape[1]152                neighbor_counts += pad_bits[start_row:end_row, start_col:end_col]153    cell.tags = neighbor_counts.astype(dtype)154    return cell