extract.py

1import numpy as np2from typing import Tuple3from mrlypy.core.errors import MrlyError4from .models import Network56# CORE GRAPH - ONE NODE PER OCCUPIED CELL, EDGES JOIN FACE-ADJACENT CELLS78def core_graph(cell) -> Network:9    grid = _grid(cell)10    occ = (grid != 0)11    dim = occ.ndim12    if dim not in (2, 3):13        raise MrlyError("core_graph expects a 2D or 3D cell.")14    coords = np.argwhere(occ)15    index_of = -np.ones(occ.shape, dtype=np.int64)16    for i, coord in enumerate(coords):17        index_of[tuple(coord)] = i18    network = Network(dim=dim)19    for coord in coords:20        network.add_node(_center(coord))21    for axis in range(dim):22        shifted = _shift(occ, axis)23        both = occ & shifted24        for coord in np.argwhere(both):25            here = index_of[tuple(coord)]26            neighbor_coord = coord.copy()27            neighbor_coord[axis] += 128            there = index_of[tuple(neighbor_coord)]29            network.add_branch(int(here), int(there))30    return network3132# EDGE GRAPH - NODES AT CELL CORNERS, EDGES ALONG CELL BOUNDARIES3334def edge_graph(cell) -> Network:35    grid = _grid(cell)36    occ = (grid != 0)37    dim = occ.ndim38    if dim not in (2, 3):39        raise MrlyError("edge_graph expects a 2D or 3D cell.")40    corner_shape = tuple(s + 1 for s in occ.shape)41    corner_used = np.zeros(corner_shape, dtype=bool)42    for offset in _corner_offsets(dim):43        slices = tuple(slice(o, o + s) for o, s in zip(offset, occ.shape))44        corner_used[slices] |= occ45    corner_index = -np.ones(corner_shape, dtype=np.int64)46    network = Network(dim=dim)47    for coord in np.argwhere(corner_used):48        corner_index[tuple(coord)] = network.add_node(tuple(float(c) for c in coord[::-1]))49    seen = set()50    for cell_coord in np.argwhere(occ):51        for a, b in _cell_edges(cell_coord, dim):52            ia = corner_index[a]53            ib = corner_index[b]54            key = (min(ia, ib), max(ia, ib))55            if key in seen:56                continue57            seen.add(key)58            network.add_branch(int(ia), int(ib))59    return network6061# TUNNEL GRAPH - CORE GRAPH OF THE INVERTED (VOID) CELL6263def tunnel_graph(cell) -> Network:64    grid = _grid(cell)65    inverted = 1 - (grid != 0).astype(np.uint8)66    return core_graph(inverted)6768# HELPERS6970def _grid(cell) -> np.ndarray:71    if hasattr(cell, "types"):72        return np.asarray(cell.types)73    return np.asarray(cell)7475def _center(coord: np.ndarray) -> Tuple[float, ...]:76    return tuple(float(c) + 0.5 for c in coord[::-1])7778def _shift(occ: np.ndarray, axis: int) -> np.ndarray:79    shifted = np.zeros_like(occ)80    src = [slice(None)] * occ.ndim81    dst = [slice(None)] * occ.ndim82    src[axis] = slice(1, None)83    dst[axis] = slice(0, -1)84    shifted[tuple(dst)] = occ[tuple(src)]85    return shifted8687def _corner_offsets(dim: int):88    if dim == 2:89        return [(dy, dx) for dy in (0, 1) for dx in (0, 1)]90    return [(dz, dy, dx) for dz in (0, 1) for dy in (0, 1) for dx in (0, 1)]9192def _cell_edges(coord: np.ndarray, dim: int):93    if dim == 2:94        y, x = int(coord[0]), int(coord[1])95        corners = {96            (0, 0): (y, x), (0, 1): (y, x + 1),97            (1, 0): (y + 1, x), (1, 1): (y + 1, x + 1),98        }99        pairs = [((0, 0), (0, 1)), ((0, 1), (1, 1)), ((1, 1), (1, 0)), ((1, 0), (0, 0))]100        return [(corners[a], corners[b]) for a, b in pairs]101    z, y, x = int(coord[0]), int(coord[1]), int(coord[2])102    def corner(dz, dy, dx):103        return (z + dz, y + dy, x + dx)104    verts = {(dz, dy, dx): corner(dz, dy, dx) for dz in (0, 1) for dy in (0, 1) for dx in (0, 1)}105    pairs = []106    for fixed_axis in range(3):107        for a in ((0, 0), (0, 1), (1, 0), (1, 1)):108            key_lo = list(a)109            key_lo.insert(fixed_axis, 0)110            key_hi = list(a)111            key_hi.insert(fixed_axis, 1)112            pairs.append((tuple(key_lo), tuple(key_hi)))113    return [(verts[a], verts[b]) for a, b in pairs]