census.py

1import numpy as np2from typing import Dict, TYPE_CHECKING3from .errors import MrlyError45if TYPE_CHECKING:6    from mrlypy.two.models import Cell2d7    from mrlypy.three.models import Cell3d89# 2D PRIMITIVES1011def _occ_2d(grid: np.ndarray) -> np.ndarray:12    return (grid != 0)1314def cells_2d(grid: np.ndarray) -> int:15    return int(_occ_2d(grid).sum())1617def vertices_2d(grid: np.ndarray) -> int:18    occ = _occ_2d(grid)19    h, w = occ.shape20    corners = np.zeros((h + 1, w + 1), dtype=bool)21    corners[:-1, :-1] |= occ22    corners[:-1, 1:] |= occ23    corners[1:, :-1] |= occ24    corners[1:, 1:] |= occ25    return int(corners.sum())2627def edges_2d(grid: np.ndarray) -> int:28    occ = _occ_2d(grid)29    h, w = occ.shape30    horiz = np.zeros((h + 1, w), dtype=bool)31    horiz[:-1, :] |= occ32    horiz[1:, :] |= occ33    vert = np.zeros((h, w + 1), dtype=bool)34    vert[:, :-1] |= occ35    vert[:, 1:] |= occ36    return int(horiz.sum() + vert.sum())3738def perimeter_2d(grid: np.ndarray) -> int:39    occ = _occ_2d(grid).astype(np.int64)40    h, w = occ.shape41    top = np.pad(occ, ((1, 0), (0, 0)))[:-1]42    bottom = np.pad(occ, ((0, 1), (0, 0)))[1:]43    left = np.pad(occ, ((0, 0), (1, 0)))[:, :-1]44    right = np.pad(occ, ((0, 0), (0, 1)))[:, 1:]45    exposed = (occ - top) + (occ - bottom) + (occ - left) + (occ - right)46    return int(np.clip(exposed, 0, None)[occ.astype(bool)].sum())4748def faces_2d(grid: np.ndarray) -> int:49    return cells_2d(grid)5051def euler_2d(grid: np.ndarray) -> int:52    return vertices_2d(grid) - edges_2d(grid) + faces_2d(grid)5354def census_2d(grid: np.ndarray) -> Dict[str, int]:55    return {56        "cells": cells_2d(grid),57        "vertices": vertices_2d(grid),58        "edges": edges_2d(grid),59        "perimeter": perimeter_2d(grid),60        "faces": faces_2d(grid),61        "euler": euler_2d(grid),62    }6364# 3D PRIMITIVES6566def _occ_3d(grid: np.ndarray) -> np.ndarray:67    return (grid != 0)6869def cells_3d(grid: np.ndarray) -> int:70    return int(_occ_3d(grid).sum())7172def vertices_3d(grid: np.ndarray) -> int:73    occ = _occ_3d(grid)74    d, h, w = occ.shape75    corners = np.zeros((d + 1, h + 1, w + 1), dtype=bool)76    for dz in (0, 1):77        for dy in (0, 1):78            for dx in (0, 1):79                corners[dz:dz + d, dy:dy + h, dx:dx + w] |= occ80    return int(corners.sum())8182def edges_3d(grid: np.ndarray) -> int:83    occ = _occ_3d(grid)84    d, h, w = occ.shape85    ex = np.zeros((d + 1, h + 1, w), dtype=bool)86    for dz in (0, 1):87        for dy in (0, 1):88            ex[dz:dz + d, dy:dy + h, :] |= occ89    ey = np.zeros((d + 1, h, w + 1), dtype=bool)90    for dz in (0, 1):91        for dx in (0, 1):92            ey[dz:dz + d, :, dx:dx + w] |= occ93    ez = np.zeros((d, h + 1, w + 1), dtype=bool)94    for dy in (0, 1):95        for dx in (0, 1):96            ez[:, dy:dy + h, dx:dx + w] |= occ97    return int(ex.sum() + ey.sum() + ez.sum())9899def faces_3d(grid: np.ndarray) -> int:100    occ = _occ_3d(grid)101    d, h, w = occ.shape102    fz = np.zeros((d + 1, h, w), dtype=bool)103    fz[:-1] |= occ104    fz[1:] |= occ105    fy = np.zeros((d, h + 1, w), dtype=bool)106    fy[:, :-1] |= occ107    fy[:, 1:] |= occ108    fx = np.zeros((d, h, w + 1), dtype=bool)109    fx[:, :, :-1] |= occ110    fx[:, :, 1:] |= occ111    return int(fz.sum() + fy.sum() + fx.sum())112113def surface_3d(grid: np.ndarray) -> int:114    occ = _occ_3d(grid).astype(np.int64)115    total = 0116    for axis in (0, 1, 2):117        a = np.pad(occ, [(1, 0) if i == axis else (0, 0) for i in range(3)])118        a = np.take(a, range(occ.shape[axis]), axis=axis)119        b = np.pad(occ, [(0, 1) if i == axis else (0, 0) for i in range(3)])120        b = np.take(b, range(1, occ.shape[axis] + 1), axis=axis)121        total += int(np.clip(occ - a, 0, None).sum())122        total += int(np.clip(occ - b, 0, None).sum())123    return total124125def euler_3d(grid: np.ndarray) -> int:126    return vertices_3d(grid) - edges_3d(grid) + faces_3d(grid) - cells_3d(grid)127128def census_3d(grid: np.ndarray) -> Dict[str, int]:129    return {130        "cells": cells_3d(grid),131        "vertices": vertices_3d(grid),132        "edges": edges_3d(grid),133        "faces": faces_3d(grid),134        "surface": surface_3d(grid),135        "euler": euler_3d(grid),136    }137138# PUBLIC139140def census(cell) -> Dict[str, int]:141    grid = getattr(cell, "_cell", cell)142    if hasattr(grid, "depth"):143        return census_3d(grid.types)144    if hasattr(grid, "height"):145        return census_2d(grid.types)146    arr = np.asarray(cell)147    if arr.ndim == 3:148        return census_3d(arr)149    if arr.ndim == 2:150        return census_2d(arr)151    raise MrlyError("census expects a Cell2d, Cell3d, or 2D/3D array.")