swcgeom 0.20.0__cp312-cp312-macosx_15_0_arm64.whl

swcgeom/core/tree.py

@@ -0,0 +1,387 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Neuron tree."""
+
+import itertools
+import os
+from collections.abc import Callable, Iterable, Iterator
+from typing import Literal, TypeVar, Union, overload
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+from typing_extensions import deprecated
+
+from swcgeom.core.branch import Branch
+from swcgeom.core.compartment import Compartment, Compartments
+from swcgeom.core.node import Node
+from swcgeom.core.path import Path
+from swcgeom.core.swc import DictSWC, eswc_cols
+from swcgeom.core.swc_utils import SWCNames, get_names, read_swc, traverse
+from swcgeom.core.tree_utils_impl import Mapping, get_subtree_impl
+from swcgeom.utils import PathOrIO, padding1d
+
+__all__ = ["Tree"]
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+
+class Tree(DictSWC):
+    """A neuron tree, which should be a binary tree in most cases."""
+
+    class Node(Node["Tree"]):
+        """Neural node."""
+
+        def parent(self) -> Union["Tree.Node", None]:
+            return Tree.Node(self.attach, self.pid) if self.pid != -1 else None
+
+        def children(self) -> list["Tree.Node"]:
+            children = self.attach.id()[self.attach.pid() == self.id]
+            return [Tree.Node(self.attach, idx) for idx in children]
+
+        def branch(self) -> "Tree.Branch":
+            ns: list["Tree.Node"] = [self]
+            while not ns[-1].is_furcation() and (p := ns[-1].parent()) is not None:
+                ns.append(p)
+
+            ns.reverse()
+            while not (ns[-1].is_furcation() or ns[-1].is_tip()):
+                ns.append(ns[-1].children()[0])
+
+            return Tree.Branch(self.attach, [n.id for n in ns])
+
+        def radial_distance(self) -> float:
+            """The end-to-end straight-line distance to soma."""
+            return self.distance(self.attach.soma())
+
+        def subtree(self, *, out_mapping: Mapping | None = None) -> "Tree":
+            """Get subtree from node.
+
+            Args:
+                out_mapping: Map from new id to old id.
+            """
+
+            n_nodes, ndata, source, names = get_subtree_impl(
+                self.attach, self.id, out_mapping=out_mapping
+            )
+            return Tree(n_nodes, **ndata, source=source, names=names)
+
+        def is_root(self) -> bool:
+            return self.parent() is None
+
+        def is_soma(self) -> bool:  # TODO: support multi soma, e.g. 3 points
+            return self.type == self.attach.types.soma and self.is_root()
+
+        @overload
+        def traverse(
+            self, *, enter: Callable[[Node, T | None], T], mode: Literal["dfs"] = ...
+        ) -> None: ...
+        @overload
+        def traverse(
+            self, *, leave: Callable[[Node, list[K]], K], mode: Literal["dfs"] = ...
+        ) -> K: ...
+        @overload
+        def traverse(
+            self,
+            *,
+            enter: Callable[[Node, T | None], T],
+            leave: Callable[[Node, list[K]], K],
+            mode: Literal["dfs"] = ...,
+        ) -> K: ...
+        def traverse(self, **kwargs):  # type: ignore
+            """Traverse from node.
+
+            See Also:
+                ~Tree.traverse
+            """
+            return self.attach.traverse(root=self.idx, **kwargs)
+
+    class Path(Path["Tree"]):
+        # TODO: should returns `Tree.Node`
+        """Neural path."""
+
+    class Compartment(Compartment["Tree"]):
+        # TODO: should returns `Tree.Node`
+        """Neural compartment."""
+
+    Segment = Compartment  # Alias
+
+    class Branch(Branch["Tree"]):
+        # TODO: should returns `Tree.Node`
+        """Neural branch."""
+
+    def __init__(
+        self,
+        n_nodes: int,
+        *,
+        source: str = "",
+        comments: Iterable[str] | None = None,
+        names: SWCNames | None = None,
+        **kwargs: npt.NDArray,
+    ) -> None:
+        names = get_names(names)
+
+        if names.id not in kwargs:
+            kwargs[names.id] = np.arange(0, n_nodes, step=1, dtype=np.int32)
+
+        if names.pid not in kwargs:
+            kwargs[names.pid] = np.arange(-1, n_nodes - 1, step=1, dtype=np.int32)
+
+        ndata = {
+            names.id: padding1d(n_nodes, kwargs.pop(names.id, None), dtype=np.int32),
+            names.type: padding1d(
+                n_nodes, kwargs.pop(names.type, None), dtype=np.int32
+            ),
+            names.x: padding1d(n_nodes, kwargs.pop(names.x, None), dtype=np.float32),
+            names.y: padding1d(n_nodes, kwargs.pop(names.y, None), dtype=np.float32),
+            names.z: padding1d(n_nodes, kwargs.pop(names.z, None), dtype=np.float32),
+            names.r: padding1d(
+                n_nodes, kwargs.pop(names.r, None), dtype=np.float32, padding_value=1
+            ),
+            names.pid: padding1d(n_nodes, kwargs.pop(names.pid, None), dtype=np.int32),
+        }
+        # ? padding other columns
+        super().__init__(
+            **ndata, **kwargs, source=source, comments=comments, names=names
+        )
+
+    def __iter__(self) -> Iterator[Node]:
+        return (self[i] for i in range(len(self)))
+
+    def __repr__(self) -> str:
+        n_nodes, n_edges = self.number_of_nodes(), self.number_of_edges()
+        return f"Neuron Tree with {n_nodes} nodes and {n_edges} edges"
+
+    @overload
+    def __getitem__(self, key: slice) -> list[Node]: ...
+    @overload
+    def __getitem__(self, key: int) -> Node: ...
+    @overload
+    def __getitem__(self, key: str) -> npt.NDArray: ...
+    def __getitem__(self, key):
+        if isinstance(key, slice):
+            return [self.node(i) for i in range(*key.indices(len(self)))]
+
+        if isinstance(key, (int, np.integer)):
+            length = len(self)
+            if key < -length or key >= length:
+                raise IndexError(f"The index ({key}) is out of range.")
+
+            if key < 0:  # Handle negative indices
+                key += length
+
+            return self.node(key)
+
+        if isinstance(key, str):
+            return self.get_ndata(key)
+
+        raise TypeError("Invalid argument type.")
+
+    def keys(self) -> Iterable[str]:
+        return self.ndata.keys()
+
+    def node(self, idx: int | np.integer) -> Node:
+        return self.Node(self, idx)
+
+    def soma(self, type_check: bool = True) -> Node:
+        """Get soma of neuron."""
+        # TODO: find soma, see also: https://neuromorpho.org/myfaq.jsp
+        n = self.node(0)
+        if type_check and n.type != self.types.soma:
+            raise ValueError(f"no soma found in: {self.source}")
+        return n
+
+    def get_furcations(self) -> list[Node]:
+        """Get all node of furcations."""
+        furcations: list[int] = []
+
+        def collect_furcations(n: Tree.Node, children: list[None]) -> None:
+            if len(children) > 1:
+                furcations.append(n.id)
+
+        self.traverse(leave=collect_furcations)
+        return [self.node(i) for i in furcations]
+
+    @deprecated("Use `get_furcations` instead")
+    def get_bifurcations(self) -> list[Node]:
+        """Get all node of furcations.
+
+        .. deprecated:: 0.17.2
+            Deprecated due to the wrong spelling of furcation. For now, it is just an
+            alias of `get_furcations` and raise a warning. It will be change to raise
+            an error in the future.
+        """
+        return self.get_furcations()
+
+    def get_tips(self) -> list[Node]:
+        """Get all node of tips."""
+        tip_ids = np.setdiff1d(self.id(), self.pid(), assume_unique=True)
+        return [self.node(i) for i in tip_ids]
+
+    def get_compartments(self) -> Compartments[Compartment]:
+        return Compartments(self.Compartment(self, n.pid, n.id) for n in self[1:])
+
+    def get_segments(self) -> Compartments[Compartment]:  # Alias
+        return self.get_compartments()
+
+    def get_branches(self) -> list[Branch]:
+        def collect_branches(
+            node: "Tree.Node", pre: list[tuple[list[Tree.Branch], list[int]]]
+        ) -> tuple[list[Tree.Branch], list[int]]:
+            if len(pre) == 1:
+                branches, children = pre[0]
+                children.append(node.id)
+                return branches, children
+
+            branches: list[Tree.Branch] = []
+            for sub_branches, children in pre:
+                children.append(node.id)
+                children.reverse()
+                sub_branches.append(
+                    Tree.Branch(self, np.array(children, dtype=np.int32))
+                )
+                sub_branches.reverse()
+                branches.extend(sub_branches)
+
+            return branches, [node.id]
+
+        branches, children = self.traverse(leave=collect_branches)
+        if len(children) > 0:
+            branches.append(Tree.Branch(self, np.array(children, dtype=np.int32)))
+
+        return branches
+
+    def get_paths(self) -> list[Path]:
+        """Get all path from soma to tips."""
+        path_dic: dict[int, list[int]] = {}
+
+        def assign_path(n: Tree.Node, pre_path: list[int] | None) -> list[int]:
+            path = [] if pre_path is None else pre_path.copy()
+            path.append(n.id)
+            path_dic[n.id] = path
+            return path
+
+        def collect_path(
+            n: Tree.Node, children: list[list[list[int]]]
+        ) -> list[list[int]]:
+            if len(children) == 0:
+                return [path_dic[n.id]]
+
+            return list(itertools.chain(*children))
+
+        paths = self.traverse(enter=assign_path, leave=collect_path)
+        return [self.Path(self, idx) for idx in paths]
+
+    def get_neurites(self, type_check: bool = True) -> Iterable["Tree"]:
+        """Get neurites from soma."""
+        return (n.subtree() for n in self.soma(type_check).children())
+
+    def get_dendrites(self, type_check: bool = True) -> Iterable["Tree"]:
+        """Get dendrites."""
+        types = [self.types.apical_dendrite, self.types.basal_dendrite]
+        children = self.soma(type_check).children()
+        return (n.subtree() for n in children if n.type in types)
+
+    @overload
+    def traverse(
+        self,
+        *,
+        enter: Callable[[Node, T | None], T],
+        root: int | np.integer = ...,
+        mode: Literal["dfs"] = ...,
+    ) -> None: ...
+    @overload
+    def traverse(
+        self,
+        *,
+        leave: Callable[[Node, list[K]], K],
+        root: int | np.integer = ...,
+        mode: Literal["dfs"] = ...,
+    ) -> K: ...
+    @overload
+    def traverse(
+        self,
+        *,
+        enter: Callable[[Node, T | None], T],
+        leave: Callable[[Node, list[K]], K],
+        root: int | np.integer = ...,
+        mode: Literal["dfs"] = ...,
+    ) -> K: ...
+    def traverse(self, *, enter=None, leave=None, **kwargs):
+        """Traverse nodes.
+
+        Args:
+            enter: (n: Node, parent: T | None) => T
+            leave: (n: Node, children: list[T]) => T
+
+        See Also:
+            ~swc_utils.traverse
+        """
+
+        def wrap(fn) -> Callable | None:
+            if fn is None:
+                return None
+
+            def fn_wrapped(idx, *args, **kwargs):
+                return fn(self[idx], *args, **kwargs)
+
+            return fn_wrapped
+
+        topology = (self.id(), self.pid())
+        enter, leave = wrap(enter), wrap(leave)
+        return traverse(topology, enter=enter, leave=leave, **kwargs)  # type: ignore
+
+    def length(self) -> float:
+        """Get length of tree."""
+        return sum(s.length() for s in self.get_segments())
+
+    @staticmethod
+    def from_data_frame(
+        df: pd.DataFrame,
+        source: str = "",
+        *,
+        comments: Iterable[str] | None = None,
+        names: SWCNames | None = None,
+    ) -> "Tree":
+        """Read neuron tree from data frame."""
+        names = get_names(names)
+        tree = Tree(
+            df.shape[0],
+            **{k: df[k].to_numpy() for k in names.cols()},
+            source=source,
+            comments=comments,
+            names=names,
+        )
+        return tree
+
+    @classmethod
+    def from_swc(cls, swc_file: PathOrIO, **kwargs) -> "Tree":
+        """Read neuron tree from swc file.
+
+        See Also:
+            ~swcgeom.core.swc_utils.read_swc
+        """
+
+        try:
+            df, comments = read_swc(swc_file, **kwargs)
+        except Exception as e:  # pylint: disable=broad-except
+            raise ValueError(f"fails to read swc: {swc_file}") from e
+
+        source = os.path.abspath(swc_file) if isinstance(swc_file, str) else ""
+        return cls.from_data_frame(df, source=source, comments=comments)
+
+    @classmethod
+    def from_eswc(
+        cls, swc_file: str, extra_cols: list[str] | None = None, **kwargs
+    ) -> "Tree":
+        """Read neuron tree from eswc file.
+
+        See Also:
+            ~swcgeom.Tree.from_swc
+        """
+        extra_cols = extra_cols or []
+        extra_cols.extend(k for k, t in eswc_cols)
+        return cls.from_swc(swc_file, extra_cols=extra_cols, **kwargs)

swcgeom/core/tree_utils.py

@@ -0,0 +1,277 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWC util wrapper for tree."""
+
+import warnings
+from collections.abc import Callable, Iterable
+from typing import TypeVar, overload
+
+import numpy as np
+from typing_extensions import deprecated
+
+from swcgeom.core.swc import SWCLike
+from swcgeom.core.swc_utils import (
+    REMOVAL,
+    SWCNames,
+    Topology,
+    get_names,
+    is_bifurcate,
+    propagate_removal,
+    sort_nodes_impl,
+    to_sub_topology,
+)
+from swcgeom.core.tree import Tree
+from swcgeom.core.tree_utils_impl import Mapping, get_subtree_impl, to_subtree_impl
+
+__all__ = [
+    "sort_tree",
+    "cut_tree",
+    "to_sub_tree",
+    "to_subtree",
+    "get_subtree",
+    "redirect_tree",
+    "cat_tree",
+]
+
+T = TypeVar("T")
+K = TypeVar("K")
+EPS = 1e-5
+
+
+def is_binary_tree(tree: Tree, exclude_soma: bool = True) -> bool:
+    """Check is it a bifurcate tree."""
+    return is_bifurcate((tree.id(), tree.pid()), exclude_root=exclude_soma)
+
+
+def sort_tree(tree: Tree) -> Tree:
+    """Sort the indices of neuron tree.
+
+    See Also:
+        ~.core.swc_utils.sort_nodes
+    """
+    return _sort_tree(tree.copy())
+
+
+@overload
+def cut_tree(
+    tree: Tree, *, enter: Callable[[Tree.Node, T | None], tuple[T, bool]]
+) -> Tree: ...
+@overload
+def cut_tree(
+    tree: Tree, *, leave: Callable[[Tree.Node, list[K]], tuple[K, bool]]
+) -> Tree: ...
+def cut_tree(tree: Tree, *, enter=None, leave=None):
+    """Traverse and cut the tree.
+
+    Returning a `True` can delete the current node and its children.
+    """
+
+    removals: list[int] = []
+
+    if enter:
+
+        def _enter(n: Tree.Node, parent: tuple[T, bool] | None) -> tuple[T, bool]:
+            if parent is not None and parent[1]:
+                removals.append(n.id)
+                return parent
+
+            res, removal = enter(n, parent[0] if parent else None)
+            if removal:
+                removals.append(n.id)
+
+            return res, removal
+
+        tree.traverse(enter=_enter)
+
+    elif leave:
+
+        def _leave(n: Tree.Node, children: list[K]) -> K:
+            res, removal = leave(n, children)
+            if removal:
+                removals.append(n.id)
+
+            return res
+
+        tree.traverse(leave=_leave)
+
+    else:
+        return tree.copy()
+
+    return to_subtree(tree, removals)
+
+
+@deprecated("Use `to_subtree` instead")
+def to_sub_tree(swc_like: SWCLike, sub: Topology) -> tuple[Tree, dict[int, int]]:
+    """Create subtree from origin tree.
+
+    You can directly mark the node for removal, and we will remove it,
+    but if the node you remove is not a leaf node, you need to use
+    `propagate_remove` to remove all children.
+
+    .. deprecated:: 0.6.0
+        Use :meth:`to_subtree` instead.
+
+    Returns
+        tree: Tree
+        id_map: dict[int, int]
+    """
+
+    sub = propagate_removal(sub)
+    (new_id, new_pid), id_map_arr = to_sub_topology(sub)
+
+    n_nodes = new_id.shape[0]
+    ndata = {k: swc_like.get_ndata(k)[id_map_arr].copy() for k in swc_like.keys()}
+    ndata.update(id=new_id, pid=new_pid)
+
+    subtree = Tree(n_nodes, **ndata, source=swc_like.source, names=swc_like.names)
+
+    id_map = {}
+    for i, idx in enumerate(id_map_arr):
+        id_map[idx] = i
+    return subtree, id_map
+
+
+def to_subtree(
+    swc_like: SWCLike,
+    removals: Iterable[int],
+    *,
+    out_mapping: Mapping | None = None,
+) -> Tree:
+    """Create subtree from origin tree.
+
+    Args:
+        swc_like: SWCLike
+        removals: A list of id of nodes to be removed.
+        out_mapping: Map new id to old id.
+    """
+
+    new_ids = swc_like.id().copy()
+    for i in removals:
+        new_ids[i] = REMOVAL
+
+    sub = propagate_removal((new_ids, swc_like.pid()))
+    n_nodes, ndata, source, names = to_subtree_impl(
+        swc_like, sub, out_mapping=out_mapping
+    )
+    return Tree(n_nodes, **ndata, source=source, names=names)
+
+
+def get_subtree(
+    swc_like: SWCLike, n: int, *, out_mapping: Mapping | None = None
+) -> Tree:
+    """Get subtree rooted at n.
+
+    Args:
+        swc_like: SWCLike
+        n: Id of the root of the subtree.
+        out_mapping: Map new id to old id.
+    """
+
+    n_nodes, ndata, source, names = get_subtree_impl(
+        swc_like, n, out_mapping=out_mapping
+    )
+    return Tree(n_nodes, **ndata, source=source, names=names)
+
+
+def redirect_tree(tree: Tree, new_root: int, sort: bool = True) -> Tree:
+    """Set root to new point and redirect tree graph.
+
+    Args:
+        tree: The tree.
+        new_root: The id of new root.
+        sort: If true, sort indices of nodes after redirect.
+    """
+
+    tree = tree.copy()
+    path = [tree.node(new_root)]
+    while (p := path[-1].parent()) is not None:
+        path.append(p)
+
+    path[0].pid = -1
+    path[0].type, path[-1].type = path[-1].type, path[0].type
+    for n, p in zip(path[1:], path[:-1]):
+        n.pid = p.id
+
+    if sort:
+        _sort_tree(tree)
+
+    return tree
+
+
+def cat_tree(  # pylint: disable=too-many-arguments
+    tree1: Tree,
+    tree2: Tree,
+    node1: int = 0,
+    node2: int = 0,
+    *,
+    translate: bool = True,
+    names: SWCNames | None = None,
+    no_move: bool | None = None,  # legacy
+) -> Tree:
+    """Concatenates the second tree onto the first one.
+
+    Args:
+        tree1: Tree
+        tree2: Tree
+        node1: The node id of the tree to be connected.
+        node2: The node id of the connection point.
+        translate: Weather to translate node_2 to node_1.
+            If False, add link between node_1 and node_2 without translate.
+    """
+    if no_move is not None:
+        warnings.warn(
+            "`no_move` has been, it is replaced by `translate` in "
+            "v0.12.0, and this will be removed in next version",
+            DeprecationWarning,
+        )
+        translate = not no_move
+
+    names = get_names(names)
+    tree, tree2 = tree1.copy(), tree2.copy()
+    if not tree2.node(node2).is_root():
+        tree2 = redirect_tree(tree2, node2, sort=False)
+
+    c = tree.node(node1)
+    if translate:
+        tree2.ndata[names.x] -= tree2.node(node2).x - c.x
+        tree2.ndata[names.y] -= tree2.node(node2).y - c.y
+        tree2.ndata[names.z] -= tree2.node(node2).z - c.z
+
+    ns = tree.number_of_nodes()
+    if np.linalg.norm(tree2.node(node2).xyz() - c.xyz()) < EPS:
+        remove = [node2 + ns]
+        link_to_root = [n.id + ns for n in tree2.node(node2).children()]
+    else:
+        remove = None
+        link_to_root = [node2 + ns]
+
+    # APIs of tree2 are no longer available since we modify the topology
+    tree2.ndata[names.id] += ns
+    tree2.ndata[names.pid] += ns
+
+    for k, v in tree.ndata.items():  # only keep keys in tree1
+        if k in tree2.ndata:
+            tree.ndata[k] = np.concatenate([v, tree2.ndata[k]])
+        else:
+            tree.ndata[k] = np.pad(v, (0, tree2.number_of_nodes()))
+
+    for n in link_to_root:
+        tree.node(n).pid = node1
+
+    if remove is not None:  # TODO: This should be easy to implement during sort
+        for k, v in tree.ndata.items():
+            tree.ndata[k] = np.delete(v, remove)
+
+    _sort_tree(tree)
+    return tree
+
+
+def _sort_tree(tree: Tree) -> Tree:
+    """Sort the indices of neuron tree inplace."""
+    (new_ids, new_pids), id_map = sort_nodes_impl((tree.id(), tree.pid()))
+    tree.ndata = {k: tree.ndata[k][id_map] for k in tree.ndata}
+    tree.ndata.update(id=new_ids, pid=new_pids)
+    return tree