slice.py

3.6 kB · python · 102 lines

1import numpy as np2from typing import Dict, List, Tuple3from mrlypy.core.errors import MrlyError4from .models import Network56VOID = 07FILL = 18GRID = 2910# TRIANGLE GEOMETRY1112def _north(x: int, y: int) -> List[Tuple[int, int]]:13    return [(x, 2 * y + 2), (x + 1, 2 * y), (x + 2, 2 * y + 2)]1415def _south(x: int, y: int) -> List[Tuple[int, int]]:16    return [(x, 2 * y), (x + 1, 2 * y + 2), (x + 2, 2 * y)]1718def _corners(x: int, y: int, start: int) -> List[Tuple[int, int]]:19    north = (x + y + start) % 2 == 020    return _north(x, y) if north else _south(x, y)2122def _edges_of(corners: List[Tuple[int, int]]) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]:23    a, b, c = corners24    return [tuple(sorted((a, b))), tuple(sorted((b, c))), tuple(sorted((a, c)))]2526def _centroid(corners: List[Tuple[int, int]]) -> Tuple[float, float]:27    xs = sum(p[0] for p in corners) / 3.028    ys = sum(p[1] for p in corners) / 3.029    return (xs, ys)3031# HELPERS3233def _unwrap(cell):34    inner = cell._cell if hasattr(cell, "_cell") else cell35    start = getattr(cell, "start", 0)36    if not hasattr(inner, "types"):37        raise MrlyError("slice graph expects a Cell6d or a 2D cell with a types array.")38    types = np.asarray(inner.types)39    if types.ndim != 2:40        raise MrlyError("slice graph expects a 2D (triangular) slice.")41    return types, int(start)4243def _adjacency_graph(types: np.ndarray, start: int, keep) -> Network:44    height, width = types.shape45    cells = [(x, y) for y in range(height) for x in range(width) if keep(int(types[y, x]))]46    index_of = {cell: i for i, cell in enumerate(cells)}47    network = Network(dim=2)48    for (x, y) in cells:49        network.add_node(_centroid(_corners(x, y, start)))50    edge_to_cells: Dict[Tuple, List[Tuple[int, int]]] = {}51    for (x, y) in cells:52        for edge in _edges_of(_corners(x, y, start)):53            edge_to_cells.setdefault(edge, []).append((x, y))54    for shared in edge_to_cells.values():55        if len(shared) == 2:56            a, b = shared57            network.add_branch(index_of[a], index_of[b])58    return network5960# CORE GRAPH - ONE NODE PER FILL TRIANGLE, EDGES JOIN EDGE-ADJACENT FILL TRIANGLES6162def slice_core_graph(cell) -> Network:63    types, start = _unwrap(cell)64    return _adjacency_graph(types, start, lambda v: v == FILL)6566# TUNNEL GRAPH - SAME ON THE VOID TRIANGLES (THE PORE NETWORK OF THE SLICE)6768def slice_tunnel_graph(cell) -> Network:69    types, start = _unwrap(cell)70    return _adjacency_graph(types, start, lambda v: v == VOID)7172# EDGE GRAPH - NODES AT TRIANGLE CORNERS, EDGES ALONG TRIANGLE SIDES (THE MESH)7374def slice_edge_graph(cell, value: int = FILL) -> Network:75    types, start = _unwrap(cell)76    height, width = types.shape77    if value is None:78        keep = lambda v: v in (FILL, VOID)79    else:80        keep = lambda v: v == value81    corner_index: Dict[Tuple[int, int], int] = {}82    network = Network(dim=2)83    def node(corner: Tuple[int, int]) -> int:84        if corner not in corner_index:85            corner_index[corner] = network.add_node((float(corner[0]), float(corner[1])))86        return corner_index[corner]87    seen = set()88    for y in range(height):89        for x in range(width):90            if not keep(int(types[y, x])):91                continue92            corners = _corners(x, y, start)93            for edge in _edges_of(corners):94                if edge in seen:95                    continue96                seen.add(edge)97                network.add_branch(node(edge[0]), node(edge[1]))98    return network99100def slice_dual_graph(cell) -> Network:101    types, start = _unwrap(cell)102    return _adjacency_graph(types, start, lambda v: v in (FILL, VOID))