geometry.py

1import numpy as np2from copy import deepcopy3from typing import List, Tuple, TYPE_CHECKING4from mrlypy.core.errors import MrlyError56if TYPE_CHECKING:7    from .models import Cell3d89# PUBLIC - IMMUTABLE1011def merge_3d(cells: List["Cell3d"], width: int, height: int, depth: int) -> "Cell3d":12    from .models import Cell3d13    if not cells:14        raise MrlyError("Cannot merge an empty list of cells.")15    if len(cells) != width * height * depth:16        raise MrlyError(f"len(cells) != width * height * depth: {len(cells)} != {width} * {height} * {depth}")17    first_cell = cells[0]18    cell_width, cell_height, cell_depth = first_cell.width, first_cell.height, first_cell.depth19    total_width = width * cell_width20    total_height = height * cell_height21    total_depth = depth * cell_depth22    new_cell = Cell3d(width=total_width, height=total_height, depth=total_depth)23    for i, cell in enumerate(cells):24        if cell.width != cell_width or cell.height != cell_height or cell.depth != cell_depth:25            raise MrlyError("All cells in a merge operation must have the same dimensions.")26        x = i % width27        y = (i // width) % height28        z = i // (width * height)29        start_x = x * cell_width30        end_x = start_x + cell_width31        start_y = y * cell_height32        end_y = start_y + cell_height33        start_z = z * cell_depth34        end_z = start_z + cell_depth35        if cell._types is not None:36            new_cell.types[start_z:end_z, start_y:end_y, start_x:end_x] = cell.types37        if cell._colors is not None:38            new_cell.colors[start_z:end_z, start_y:end_y, start_x:end_x] = cell.colors39        if cell._tags is not None:40            new_cell.tags[start_z:end_z, start_y:end_y, start_x:end_x] = cell.tags41    return new_cell4243def combine_3d(cell_1: "Cell3d", cell_2: "Cell3d") -> "Cell3d":44    from .models import Cell3d45    new_types = np.kron(cell_1.types, cell_2.types).astype(np.uint8)46    return Cell3d(types=new_types)4748def magic_3d(cells: List["Cell3d"]) -> "Cell3d":49    if len(cells) < 2:50        raise MrlyError("Magic composition requires at least two cells.")51    new_cell = combine_3d(cells[0], cells[1])52    for i in range(2, len(cells)):53        new_cell = combine_3d(new_cell, cells[i])54    return new_cell5556def special_3d(mask: np.ndarray, cell: "Cell3d") -> "Cell3d":57    depth, height, width = mask.shape58    new_cells = []59    for z in range(depth):60        for y in range(height):61            for x in range(width):62                rotation_index = mask[z, y, x]63                new_cell = deepcopy(cell).rotate(k=int(rotation_index))64                new_cells.append(new_cell)65    return merge_3d(new_cells, width, height, depth)6667def mosaic_3d(mask: np.ndarray, cells: List["Cell3d"]) -> "Cell3d":68    depth, height, width = mask.shape69    new_cells = []70    for z in range(depth):71        for y in range(height):72            for x in range(width):73                cell_index = mask[z, y, x]74                new_cell = deepcopy(cells[cell_index])75                new_cells.append(new_cell)76    return merge_3d(new_cells, width, height, depth)7778# PRIVATE - MUTABLE7980def invert_3d(cell: "Cell3d") -> "Cell3d":81    if cell._types is not None:82        cell.types = 1 - cell.types83    return cell8485def pad_3d(cell: "Cell3d", count: int = 1, value: int = 0) -> "Cell3d":86    if cell._types is not None:87        cell.types = np.pad(cell.types, count, mode="constant", constant_values=value)88    if cell._colors is not None:89        cell.colors = np.pad(cell.colors, count, mode="constant", constant_values=value)90    if cell._tags is not None:91        cell.tags = np.pad(cell.tags, count, mode="constant", constant_values=value)92    return cell9394def rotate_3d(cell: "Cell3d", k: int = 1, axes: Tuple[int, int] = (1, 2)) -> "Cell3d":95    if cell._types is not None:96        cell.types = np.rot90(cell.types, k=k, axes=axes)97    if cell._colors is not None:98        cell.colors = np.rot90(cell.colors, k=k, axes=axes)99    if cell._tags is not None:100        cell.tags = np.rot90(cell.tags, k=k, axes=axes)101    return cell102103def fractal_3d(cell: "Cell3d", level: int = 1) -> "Cell3d":104    if level < 1:105        raise MrlyError("Fractal level must be at least 1.")106    if level == 1:107        return cell108    new_types = cell.types109    for _ in range(1, level):110        new_types = np.kron(new_types, cell.types)111    cell.types = new_types.astype(np.uint8)112    cell.colors = None113    cell.tags = None114    return cell115116def tile_3d(cell: "Cell3d", width: int, height: int, depth: int) -> "Cell3d":117    if cell._types is not None:118        cell.types = np.tile(cell.types, (depth, height, width))119    if cell._colors is not None:120        cell.colors = np.tile(cell.colors, (depth, height, width, 1))121    if cell._tags is not None:122        cell.tags = np.tile(cell.tags, (depth, height, width))123    return cell124125def layers_3d(cell: "Cell3d", dtype: np.dtype = np.dtype(np.uint8)) -> "Cell3d":126    depth, height, width = cell.depth, cell.height, cell.width127    z_indices, y_indices, x_indices = np.indices((depth, height, width))128    center_z = (depth - 1) / 2129    center_y = (height - 1) / 2130    center_x = (width - 1) / 2131    distance_z = np.abs(z_indices - center_z)132    distance_y = np.abs(y_indices - center_y)133    distance_x = np.abs(x_indices - center_x)134    tags = np.floor(distance_x + distance_y + distance_z)135    cell.tags = tags.astype(dtype)136    return cell137138def neighbors_3d(cell: "Cell3d", mode: str = "constant") -> "Cell3d":139    kernel = np.ones((3, 3, 3), dtype=np.uint8)140    kernel[1, 1, 1] = 0141    padded = np.pad(cell.types, pad_width=1, mode=mode)142    result = np.zeros_like(cell.types, dtype=np.uint8)143    d, h, w = cell.types.shape144    for z in range(d):145        for y in range(h):146            for x in range(w):147                window = padded[z:z+3, y:y+3, x:x+3]148                count = np.sum(window * kernel)149                result[z, y, x] = count150    cell.tags = result151    return cell