models.py

1import numpy as np2from copy import deepcopy3from typing import Any, Dict, List, Optional, Tuple4from mrlypy.core.errors import MrlyError56class Node:78    def __init__(self, position: Tuple[float, ...], index: Optional[int] = None):9        self.position = tuple(float(c) for c in position)10        self.index = index1112    @property13    def dim(self) -> int:14        return len(self.position)1516    def __repr__(self) -> str:17        coords = ", ".join(f"{c:g}" for c in self.position)18        return f"Node({coords})"1920class Branch:2122    def __init__(self, parent: int, child: int, radius: float = 1.0):23        self.parent = parent24        self.child = child25        self.radius = float(radius)2627    def __repr__(self) -> str:28        return f"Branch({self.parent}->{self.child}, r={self.radius:g})"2930class Network:3132    def __init__(self, dim: int = 2):33        self._dim = dim34        self.nodes: List[Node] = []35        self.branches: List[Branch] = []3637    @property38    def dim(self) -> int:39        return self._dim4041    # MAIN4243    def add_node(self, position: Tuple[float, ...]) -> int:44        if len(position) != self._dim:45            raise MrlyError(f"Expected {self._dim}D position, got {len(position)}D")46        index = len(self.nodes)47        self.nodes.append(Node(position, index))48        return index4950    def add_branch(self, parent: int, child: int, radius: float = 1.0) -> Branch:51        n = len(self.nodes)52        if not (0 <= parent < n and 0 <= child < n):53            raise MrlyError(f"Branch endpoints out of range: {parent}, {child}")54        branch = Branch(parent, child, radius)55        self.branches.append(branch)56        return branch5758    def copy(self) -> "Network":59        return deepcopy(self)6061    def __repr__(self) -> str:62        return f"Network(dim={self._dim}, nodes={len(self.nodes)}, branches={len(self.branches)})"6364    # DERIVED6566    def positions(self) -> np.ndarray:67        if not self.nodes:68            return np.zeros((0, self._dim))69        return np.array([node.position for node in self.nodes], dtype=float)7071    def edge_array(self) -> np.ndarray:72        if not self.branches:73            return np.zeros((0, 2), dtype=np.int64)74        return np.array([(b.parent, b.child) for b in self.branches], dtype=np.int64)7576    def degree(self) -> np.ndarray:77        deg = np.zeros(len(self.nodes), dtype=np.int64)78        for branch in self.branches:79            deg[branch.parent] += 180            deg[branch.child] += 181        return deg8283    def adjacency(self) -> Dict[int, List[int]]:84        adj: Dict[int, List[int]] = {i: [] for i in range(len(self.nodes))}85        for branch in self.branches:86            adj[branch.parent].append(branch.child)87            adj[branch.child].append(branch.parent)88        return adj8990    # SERIALIZER9192    def to_dict(self) -> Dict[str, Any]:93        return {94            "dim": self._dim,95            "nodes": [list(node.position) for node in self.nodes],96            "branches": [[b.parent, b.child, b.radius] for b in self.branches],97        }9899    @classmethod100    def from_dict(cls, data: Dict[str, Any]) -> "Network":101        network = cls(dim=data["dim"])102        for position in data["nodes"]:103            network.add_node(tuple(position))104        for parent, child, radius in data["branches"]:105            network.add_branch(parent, child, radius)106        return network107108    # CENSUS109110    def census(self) -> Dict[str, float]:111        from . import census112        return census.census(self)