generator.py

1import numpy as np2from typing import List, Tuple3from mrlypy.core.errors import MrlyError4from .models import Network56# ABSTRACT FRACTAL TREE78def tree_2d(9    depth: int = 6,10    angle: float = 30.0,11    ratio: float = 0.7,12    children: int = 2,13    length: float = 1.0,14) -> Network:15    if depth < 1:16        raise MrlyError("Tree depth must be at least 1.")17    if children < 1:18        raise MrlyError("Tree must have at least one child per node.")19    network = Network(dim=2)20    root = network.add_node((0.0, 0.0))21    spread = np.radians(angle)22    _grow_2d(network, root, np.array([0.0, 1.0]), length, depth, spread, ratio, children)23    return network2425def _grow_2d(26    network: Network,27    parent: int,28    direction: np.ndarray,29    length: float,30    depth: int,31    spread: float,32    ratio: float,33    children: int,34):35    if depth == 0:36        return37    base = network.nodes[parent].position38    if children == 1:39        offsets = [0.0]40    else:41        offsets = list(np.linspace(-spread, spread, children))42    for offset in offsets:43        rotated = _rotate(direction, offset)44        tip = (base[0] + rotated[0] * length, base[1] + rotated[1] * length)45        child = network.add_node(tip)46        network.add_branch(parent, child, radius=ratio ** (0) + depth)47        _grow_2d(network, child, rotated, length * ratio, depth - 1, spread, ratio, children)4849def _rotate(vector: np.ndarray, angle: float) -> np.ndarray:50    cos_a, sin_a = np.cos(angle), np.sin(angle)51    matrix = np.array([[cos_a, -sin_a], [sin_a, cos_a]])52    return matrix @ vector5354# L-SYSTEM PLANT5556def plant_2d(57    depth: int = 5,58    angle: float = 25.0,59    ratio: float = 0.65,60    length: float = 1.0,61) -> Network:62    return tree_2d(depth=depth, angle=angle, ratio=ratio, children=2, length=length)6364# 3D TREE6566def tree_3d(67    depth: int = 5,68    angle: float = 30.0,69    ratio: float = 0.7,70    children: int = 3,71    length: float = 1.0,72) -> Network:73    if depth < 1:74        raise MrlyError("Tree depth must be at least 1.")75    network = Network(dim=3)76    root = network.add_node((0.0, 0.0, 0.0))77    spread = np.radians(angle)78    _grow_3d(network, root, np.array([0.0, 0.0, 1.0]), length, depth, spread, ratio, children)79    return network8081def _grow_3d(82    network: Network,83    parent: int,84    direction: np.ndarray,85    length: float,86    depth: int,87    spread: float,88    ratio: float,89    children: int,90):91    if depth == 0:92        return93    base = network.nodes[parent].position94    axis_a, axis_b = _frame(direction)95    for k in range(children):96        phi = 2.0 * np.pi * k / children97        bend = axis_a * np.cos(phi) + axis_b * np.sin(phi)98        rotated = _normalize(direction * np.cos(spread) + bend * np.sin(spread))99        tip = tuple(base[i] + rotated[i] * length for i in range(3))100        child = network.add_node(tip)101        network.add_branch(parent, child, radius=float(depth))102        _grow_3d(network, child, rotated, length * ratio, depth - 1, spread, ratio, children)103104def _frame(direction: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:105    reference = np.array([1.0, 0.0, 0.0]) if abs(direction[0]) < 0.9 else np.array([0.0, 1.0, 0.0])106    axis_a = _normalize(np.cross(direction, reference))107    axis_b = _normalize(np.cross(direction, axis_a))108    return axis_a, axis_b109110def _normalize(vector: np.ndarray) -> np.ndarray:111    norm = np.linalg.norm(vector)112    if norm == 0:113        return vector114    return vector / norm