swcgeom 0.19.4__cp311-cp311-win_amd64.whl

swcgeom/transforms/tree.py

@@ -0,0 +1,265 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Transformation in tree."""
+
+from collections.abc import Callable
+
+import numpy as np
+from typing_extensions import deprecated, override
+
+from swcgeom.core import Branch, BranchTree, DictSWC, Path, Tree, cut_tree, to_subtree
+from swcgeom.core.swc_utils import SWCTypes, get_types
+from swcgeom.transforms.base import Transform
+from swcgeom.transforms.branch import BranchConvSmoother, BranchIsometricResampler
+from swcgeom.transforms.branch_tree import BranchTreeAssembler
+from swcgeom.transforms.geometry import Normalizer
+
+__all__ = [
+    "ToBranchTree",
+    "ToLongestPath",
+    "TreeSmoother",
+    "TreeNormalizer",
+    "CutByType",
+    "CutAxonTree",
+    "CutDendriteTree",
+    "CutByFurcationOrder",
+    "CutShortTipBranch",
+    "IsometricResampler",
+]
+
+
+# pylint: disable=too-few-public-methods
+class ToBranchTree(Transform[Tree, BranchTree]):
+    """Transform tree to branch tree."""
+
+    @override
+    def __call__(self, x: Tree) -> BranchTree:
+        return BranchTree.from_tree(x)
+
+
+class ToLongestPath(Transform[Tree, Path[DictSWC]]):
+    """Transform tree to longest path."""
+
+    def __init__(self, *, detach: bool = True) -> None:
+        self.detach = detach
+
+    @override
+    def __call__(self, x: Tree) -> Path[DictSWC]:
+        paths = x.get_paths()
+        idx = np.argmax([p.length() for p in paths])
+        path = paths[idx]
+        if self.detach:
+            path = path.detach()
+        return path  # type: ignore
+
+
+class TreeSmoother(Transform[Tree, Tree]):  # pylint: disable=missing-class-docstring
+    def __init__(self, n_nodes: int = 5) -> None:
+        super().__init__()
+        self.n_nodes = n_nodes
+        self.trans = BranchConvSmoother(n_nodes=n_nodes)
+
+    @override
+    def __call__(self, x: Tree) -> Tree:
+        x = x.copy()
+        for br in x.get_branches():
+            # TODO: works but is weird
+            smoothed = self.trans(br)
+            x.ndata["x"][br.origin_id()] = smoothed.x()
+            x.ndata["y"][br.origin_id()] = smoothed.y()
+            x.ndata["z"][br.origin_id()] = smoothed.z()
+
+        return x
+
+    @override
+    def extra_repr(self) -> str:
+        return f"n_nodes={self.n_nodes}"
+
+
+@deprecated("Use `Normalizer` instead")
+class TreeNormalizer(Normalizer[Tree]):
+    """Noramlize coordinates and radius to 0-1.
+
+    .. deprecated:: 0.6.0
+        Use :cls:`Normalizer` instead.
+    """
+
+
+class CutByType(Transform[Tree, Tree]):
+    """Cut tree by type.
+
+    In order to preserve the tree structure, all ancestor nodes of the node to be preserved will be preserved.
+
+    NOTE: Not all reserved nodes are of the specified type.
+    """
+
+    def __init__(self, type: int) -> None:  # pylint: disable=redefined-builtin
+        super().__init__()
+        self.type = type
+
+    @override
+    def __call__(self, x: Tree) -> Tree:
+        removals = set(x.id()[x.type() != self.type])
+
+        def leave(n: Tree.Node, keep_children: list[bool]) -> bool:
+            if n.id in removals and any(keep_children):
+                removals.remove(n.id)
+            return n.id not in removals
+
+        x.traverse(leave=leave)
+        y = to_subtree(x, removals)
+        return y
+
+    @override
+    def extra_repr(self) -> str:
+        return f"type={self.type}"
+
+
+class CutAxonTree(CutByType):
+    """Cut axon tree."""
+
+    def __init__(self, types: SWCTypes | None = None) -> None:
+        types = get_types(types)
+        super().__init__(type=types.axon)
+
+
+class CutDendriteTree(CutByType):
+    """Cut dendrite tree."""
+
+    def __init__(self, types: SWCTypes | None = None) -> None:
+        types = get_types(types)
+        super().__init__(type=types.basal_dendrite)  # TODO: apical dendrite
+
+
+class CutByFurcationOrder(Transform[Tree, Tree]):
+    """Cut tree by furcation order."""
+
+    max_furcation_order: int
+
+    def __init__(self, max_bifurcation_order: int) -> None:
+        self.max_furcation_order = max_bifurcation_order
+
+    @override
+    def __call__(self, x: Tree) -> Tree:
+        return cut_tree(x, enter=self._enter)
+
+    def __repr__(self) -> str:
+        return f"CutByBifurcationOrder-{self.max_furcation_order}"
+
+    def _enter(self, n: Tree.Node, parent_level: int | None) -> tuple[int, bool]:
+        if parent_level is None:
+            level = 0
+        elif n.is_furcation():
+            level = parent_level + 1
+        else:
+            level = parent_level
+        return (level, level >= self.max_furcation_order)
+
+
+@deprecated("Use CutByFurcationOrder instead")
+class CutByBifurcationOrder(CutByFurcationOrder):
+    """Cut tree by bifurcation order.
+
+    NOTE: Deprecated due to the wrong spelling of furcation. For now, it
+    is just an alias of `CutByFurcationOrder` and raise a warning. It
+    will be change to raise an error in the future.
+    """
+
+    max_furcation_order: int
+
+    def __init__(self, max_bifurcation_order: int) -> None:
+        super().__init__(max_bifurcation_order)
+
+    def __repr__(self) -> str:
+        return f"CutByBifurcationOrder-{self.max_furcation_order}"
+
+
+class CutShortTipBranch(Transform[Tree, Tree]):
+    """Cut off too short terminal branches.
+
+    This method is usually applied in the post-processing of manual
+    reconstruction. When the user draw lines, a line head is often left
+    at the junction of two lines.
+    """
+
+    thre: float
+    callbacks: list[Callable[[Tree.Branch], None]]
+
+    def __init__(
+        self, thre: float = 5, callback: Callable[[Tree.Branch], None] | None = None
+    ) -> None:
+        self.thre = thre
+        self.callbacks = []
+
+        if callback is not None:
+            self.callbacks.append(callback)
+
+    @override
+    def __call__(self, x: Tree) -> Tree:
+        removals: list[int] = []
+        self.callbacks.append(lambda br: removals.append(br[1].id))
+        x.traverse(leave=self._leave)
+        self.callbacks.pop()
+        return to_subtree(x, removals)
+
+    @override
+    def extra_repr(self) -> str:
+        return f"threshold={self.thre}"
+
+    def _leave(
+        self, n: Tree.Node, children: list[tuple[float, Tree.Node] | None]
+    ) -> tuple[float, Tree.Node] | None:
+        if len(children) == 0:  # tip
+            return 0, n
+
+        if len(children) == 1 and children[0] is not None:  # elongation
+            dis, child = children[0]
+            dis += n.distance(child)
+            return dis, n
+
+        for c in children:
+            if c is None:
+                continue
+
+            dis, child = c
+            if dis + n.distance(child) > self.thre:
+                continue
+
+            path = [n.id]  # n does not delete, but will include in callback
+            while child is not None:  # TODO: perf
+                path.append(child.id)
+                child = cc[0] if len((cc := child.children())) > 0 else None
+
+            br = Tree.Branch(n.attach, path)
+            for cb in self.callbacks:
+                cb(br)
+
+        return None
+
+
+class Resampler(Transform[Tree, Tree]):
+    def __init__(self, branch_resampler: Transform[Branch, Branch]) -> None:
+        super().__init__()
+        self.resampler = branch_resampler
+        self.assembler = BranchTreeAssembler()
+
+    @override
+    def __call__(self, x: Tree) -> Tree:
+        t = BranchTree.from_tree(x)
+        t.branches = {
+            k: [self.resampler(br) for br in brs] for k, brs in t.branches.items()
+        }
+        return self.assembler(t)
+
+
+class IsometricResampler(Resampler):
+    def __init__(
+        self, distance: float, *, adjust_last_gap: bool = True, **kwargs
+    ) -> None:
+        branch_resampler = BranchIsometricResampler(
+            distance, adjust_last_gap=adjust_last_gap, **kwargs
+        )
+        super().__init__(branch_resampler)

swcgeom/transforms/tree_assembler.py

@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Assemble a tree."""
+
+from collections.abc import Iterable
+from copy import copy
+
+import numpy as np
+import pandas as pd
+from typing_extensions import override
+
+from swcgeom.core import Tree
+from swcgeom.core.swc_utils import (
+    SWCNames,
+    get_names,
+    link_roots_to_nearest_,
+    sort_nodes_,
+)
+from swcgeom.transforms.base import Transform
+
+EPS = 1e-5
+
+
+class LinesToTree(Transform[list[pd.DataFrame], Tree]):
+    """Assemble lines to swc."""
+
+    def __init__(self, *, thre: float = 0.2, undirected: bool = True):
+        """
+        Args:
+            thre: Connection threshold.
+            undirected: Both ends of a line can be considered connection point.
+                If `False`, only the starting point.
+        """
+        super().__init__()
+        self.thre = thre
+        self.undirected = undirected
+
+    @override
+    def __call__(
+        self, lines: Iterable[pd.DataFrame], *, names: SWCNames | None = None
+    ):  # TODO check this
+        return self.assemble(lines, names=names)
+
+    def assemble(
+        self,
+        lines: Iterable[pd.DataFrame],
+        *,
+        undirected: bool = True,
+        names: SWCNames | None = None,
+    ) -> pd.DataFrame:
+        """Assemble lines to a tree.
+
+        Assemble all the lines into a set of subtrees, and then connect them.
+
+        Args:
+            lines: An array of tables containing a line, columns should following the swc.
+            undirected: Forwarding to `self.try_assemble`.
+            names: Forwarding to `self.try_assemble`.
+
+        Returns:
+            tree: ~pd.DataFrame
+
+        See Also:
+            self.try_assemble
+        """
+
+        tree, lines = self.try_assemble(
+            lines, sort_nodes=False, undirected=undirected, names=names
+        )
+        while len(lines) > 0:
+            t, lines = self.try_assemble(
+                lines,
+                id_offset=len(tree),
+                sort_nodes=False,
+                undirected=undirected,
+                names=names,
+            )
+            tree = pd.concat([tree, t])
+
+        tree = tree.reset_index()
+        link_roots_to_nearest_(tree)
+        sort_nodes_(tree)
+        return tree
+
+    def try_assemble(
+        self,
+        lines: Iterable[pd.DataFrame],
+        *,
+        id_offset: int = 0,
+        undirected: bool = True,
+        sort_nodes: bool = True,
+        names: SWCNames | None = None,
+    ) -> tuple[pd.DataFrame, list[pd.DataFrame]]:
+        """Trying assemble lines to a tree.
+
+        Treat the first line as a tree, find a line whose shortest distance between
+        the tree and the line is less than threshold, merge it into the tree, repeat
+        until there are no line to merge, return tree and the remaining lines.
+
+        Args:
+            lines: An array of tables containing a line, columns should following the swc.
+            id_offset: The offset of the line node id.
+            undirected: Both ends of a line can be considered connection point.
+                If `False`, only the starting point.
+            sort_nodes: sort nodes of subtree.
+            names: SWCNames, optional
+
+        Returns:
+            tree: ~pandas.DataFrame
+            remaining_lines: List of ~pandas.DataFrame
+        """
+
+        names = get_names(names)
+        lines = copy(list(lines))
+
+        tree = lines[0]
+        tree[names.id] = id_offset + np.arange(len(tree))
+        tree[names.pid] = tree[names.id] - 1
+        tree.at[0, names.pid] = -1
+        del lines[0]
+
+        while True:
+            for i, line in enumerate(lines):
+                for p in [0, -1] if undirected else [0]:
+                    xyz = [names.x, names.y, names.z]
+                    vs = tree[xyz] - line.iloc[p][xyz]
+                    dis = np.linalg.norm(vs, axis=1)
+                    ind = np.argmin(dis)
+                    if dis[ind] > self.thre:
+                        continue
+
+                    if dis[ind] < EPS:
+                        line = line.drop((p + len(line)) % len(line)).reset_index(
+                            drop=True
+                        )
+
+                    line[names.id] = id_offset + len(tree) + np.arange(len(line))
+                    line[names.pid] = line[names.id] + (-1 if p == 0 else 1)
+                    line.at[(p + len(line)) % len(line), names.pid] = tree.iloc[ind][
+                        names.id
+                    ]
+                    tree = pd.concat([tree, line])
+                    del lines[i]
+                    break
+                else:
+                    continue
+
+                break
+            else:
+                break
+
+        if sort_nodes:
+            sort_nodes_(tree)
+
+        return tree, lines
+
+    @override
+    def extra_repr(self) -> str:
+        return f"thre={self.thre}, undirected={self.undirected}"

swcgeom/utils/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utils."""
+
+from swcgeom.utils.debug import *  # noqa: F403
+from swcgeom.utils.dsu import *  # noqa: F403
+from swcgeom.utils.ellipse import *  # noqa: F403
+from swcgeom.utils.file import *  # noqa: F403
+from swcgeom.utils.neuromorpho import *  # noqa: F403
+from swcgeom.utils.numpy_helper import *  # noqa: F403
+from swcgeom.utils.plotter_2d import *  # noqa: F403
+from swcgeom.utils.renderer import *  # noqa: F403
+from swcgeom.utils.sdf import *  # noqa: F403
+from swcgeom.utils.solid_geometry import *  # noqa: F403
+from swcgeom.utils.transforms import *  # noqa: F403
+from swcgeom.utils.volumetric_object import *  # noqa: F403

swcgeom/utils/debug.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Debug helpers"""
+
+import time
+from functools import wraps
+
+__all__ = ["func_timer"]
+
+
+def func_timer(function):
+    @wraps(function)
+    def function_timer(*args, **kwargs):
+        print(f"[Function: {function.__name__} start...]")
+        t0 = time.time()
+        result = function(*args, **kwargs)
+        t1 = time.time()
+        print(f"[Function: {function.__name__} finished, spent time: {t1 - t0:.2f}s]")
+        return result
+
+    return function_timer

swcgeom/utils/download.py

@@ -0,0 +1,119 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Download helpers.
+
+NOTE: All denpendencies need to be installed, try:
+
+```sh
+pip install swcgeom[all]
+```
+"""
+
+import itertools
+import logging
+import multiprocessing
+import os
+from functools import partial
+from urllib.parse import urljoin
+
+__all__ = ["download", "fetch_page", "clone_index_page"]
+
+
+def download(dst: str, url: str) -> None:
+    """Download a file."""
+    from urllib3 import connection_from_url
+
+    conn = connection_from_url(url)
+    r = conn.request("GET", url)
+
+    dirname = os.path.dirname(dst)
+    if dirname != "" and not os.path.exists(dirname):
+        os.makedirs(dirname)
+
+    with open(dst, "wb") as file:
+        file.write(r.data)
+
+
+def fetch_page(url: str):
+    """Fetch page content."""
+    from bs4 import BeautifulSoup
+    from urllib3 import connection_from_url
+
+    conn = connection_from_url(url)
+    r = conn.request("GET", url)
+    data = r.data.decode("utf-8")
+    return BeautifulSoup(data, features="html.parser")
+
+
+def clone_index_page(
+    index_url: str, dist_dir: str, override: bool = False, multiprocess: int = 4
+) -> None:
+    """Download directory from index page.
+
+    E.g: `https://download.brainimagelibrary.org/biccn/zeng/luo/fMOST/cells/`
+
+    Args:
+        index_url: URL of index page.
+        dist_dir: Directory of dist.
+        override: Override existing file, skip file if `False`.
+        multiprocess: How many process are available for download.
+    """
+    files = get_urls_in_index_page(index_url)
+    logging.info("downloader: search `%s`, found %s files.", index_url, len(files))
+
+    task = partial(
+        _clone_index_page, index_url=index_url, dist_dir=dist_dir, override=override
+    )
+    with multiprocessing.Pool(multiprocess) as p:
+        p.map(task, files)
+
+
+def _clone_index_page(url: str, index_url: str, dist_dir: str, override: bool) -> None:
+    from urllib3.exceptions import HTTPError
+
+    filepath = url.removeprefix(index_url)
+    dist = os.path.join(dist_dir, filepath)
+    if os.path.exists(dist):
+        if not override:
+            logging.info("downloader: file `%s` exits, skipped.", dist)
+            return
+
+        logging.info("downloader: file `%s` exits, deleted.", dist)
+        os.remove(dist)
+
+    try:
+        logging.info("downloader: downloading `%s` to `%s`", url, dist)
+        download(dist, url)
+        logging.info("downloader: download `%s` to `%s`", url, dist)
+    except HTTPError as ex:
+        logging.info("downloader: fails to download `%s`, except `%s`", url, ex)
+
+
+def get_urls_in_index_page(url: str) -> list[str]:
+    """Get all file links by dfs."""
+    soup = fetch_page(url)
+    links = [el.attrs["href"] for el in soup.find_all("a")]
+    files = [urljoin(url, a) for a in links if not a.endswith("/")]
+    dirs = [urljoin(url, a) for a in links if a != "../" and a.endswith("/")]
+    files.extend(itertools.chain(*[get_urls_in_index_page(dir) for dir in dirs]))
+    return files
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Download files from index page.")
+    parser.add_argument("url", type=str, help="URL of index page.")
+    parser.add_argument("dist", type=str, help="Directory of dist.")
+    parser.add_argument(
+        "--override", type=bool, default=False, help="Override existing file."
+    )
+    parser.add_argument(
+        "--multiprocess", type=int, default=4, help="How many process are available."
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+    clone_index_page(args.url, args.dist, args.override, args.multiprocess)

swcgeom/utils/dsu.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Disjoint Set Union Impl."""
+
+__all__ = ["DisjointSetUnion"]
+
+
+class DisjointSetUnion:
+    """Disjoint Set Union.
+
+    DSU with path compression and union by rank.
+
+    >>> dsu = DisjointSetUnion(3)
+    >>> dsu.is_same_set(0, 1)
+    False
+    >>> dsu.union_sets(0, 1)
+    >>> dsu.is_same_set(0, 1)
+    True
+    >>> dsu.is_same_set(0, 2)
+    False
+    >>> dsu.union_sets(1, 2)
+    >>> dsu.is_same_set(0, 2)
+    True
+    """
+
+    def __init__(self, node_number: int):
+        self.element_parent = [i for i in range(node_number)]
+        self.rank = [0 for _ in range(node_number)]
+
+    def find_parent(self, node_id: int) -> int:
+        if node_id != self.element_parent[node_id]:
+            self.element_parent[node_id] = self.find_parent(
+                self.element_parent[node_id]
+            )
+        return self.element_parent[node_id]
+
+    def union_sets(self, node_a: int, node_b: int) -> None:
+        assert self.validate_node(node_a) and self.validate_node(node_b)
+
+        root_a = self.find_parent(node_a)
+        root_b = self.find_parent(node_b)
+        if root_a != root_b:
+            # union by rank
+            if self.rank[root_a] < self.rank[root_b]:
+                self.element_parent[root_a] = root_b
+            elif self.rank[root_a] > self.rank[root_b]:
+                self.element_parent[root_b] = root_a
+            else:
+                self.element_parent[root_b] = root_a
+                self.rank[root_a] += 1
+
+    def is_same_set(self, node_a: int, node_b: int) -> bool:
+        return self.find_parent(node_a) == self.find_parent(node_b)
+
+    def validate_node(self, node_id: int) -> bool:
+        return 0 <= node_id < len(self.element_parent)