swcgeom 0.19.4__cp313-cp313-win_amd64.whl

swcgeom/core/swc.py

@@ -0,0 +1,247 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWC format."""
+
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from copy import deepcopy
+from typing import Any, TypeVar, overload
+
+import numpy as np
+import numpy.typing as npt
+import scipy.sparse as sp
+from typing_extensions import Self
+
+from swcgeom.core.swc_utils import (
+    SWCNames,
+    SWCTypes,
+    get_names,
+    get_types,
+    read_swc,
+    to_swc,
+)
+
+__all__ = [
+    "swc_cols",
+    "eswc_cols",
+    "SWCLike",
+    "DictSWC",
+    "SWCTypeVar",
+    # TODO: `read_swc` has been deprecated and will be removed in next
+    # version, import from `swcgeom.core.swc_utils` instead
+    "read_swc",
+]
+
+
+swc_names_default = get_names()
+swc_cols: list[tuple[str, npt.DTypeLike]] = [
+    (swc_names_default.id, np.int32),
+    (swc_names_default.type, np.int32),
+    (swc_names_default.x, np.float32),
+    (swc_names_default.y, np.float32),
+    (swc_names_default.z, np.float32),
+    (swc_names_default.r, np.float32),
+    (swc_names_default.pid, np.int32),
+]
+
+eswc_cols: list[tuple[str, npt.DTypeLike]] = [
+    ("level", np.int32),
+    ("mode", np.int32),
+    ("timestamp", np.int32),
+    ("teraflyindex", np.int32),
+    ("feature_value", np.int32),
+]
+
+
+class SWCLike(ABC):
+    """ABC of SWC."""
+
+    source: str = ""
+    comments: list[str] = []
+    names: SWCNames
+    types: SWCTypes
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.types = get_types()
+
+    def __len__(self) -> int:
+        return self.number_of_nodes()
+
+    def id(self) -> npt.NDArray[np.int32]:  # pylint: disable=invalid-name
+        """Get the ids of shape (n_sample,)."""
+        return self.get_ndata(self.names.id)
+
+    def type(self) -> npt.NDArray[np.int32]:
+        """Get the types of shape (n_sample,)."""
+        return self.get_ndata(self.names.type)
+
+    def x(self) -> npt.NDArray[np.float32]:
+        """Get the x coordinates of shape (n_sample,)."""
+        return self.get_ndata(self.names.x)
+
+    def y(self) -> npt.NDArray[np.float32]:
+        """Get the y coordinates of shape (n_sample,)."""
+        return self.get_ndata(self.names.y)
+
+    def z(self) -> npt.NDArray[np.float32]:
+        """Get the z coordinates of shape (n_sample,)."""
+        return self.get_ndata(self.names.z)
+
+    def r(self) -> npt.NDArray[np.float32]:
+        """Get the radius of shape (n_sample,)."""
+        return self.get_ndata(self.names.r)
+
+    def pid(self) -> npt.NDArray[np.int32]:
+        """Get the ids of parent of shape (n_sample,)."""
+        return self.get_ndata(self.names.pid)
+
+    def xyz(self) -> npt.NDArray[np.float32]:
+        """Get the coordinates of shape(n_sample, 3)."""
+        return np.stack([self.x(), self.y(), self.z()], axis=1)
+
+    def xyzw(self) -> npt.NDArray[np.float32]:
+        """Get the homogeneous coordinates of shape(n_sample, 4)."""
+        w = np.ones_like(self.x())
+        return np.stack([self.x(), self.y(), self.z(), w], axis=1)
+
+    def xyzr(self) -> npt.NDArray[np.float32]:
+        """Get the coordinates and radius array of shape(n_sample, 4)."""
+        return np.stack([self.x(), self.y(), self.z(), self.r()], axis=1)
+
+    @abstractmethod
+    def keys(self) -> Iterable[str]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_ndata(self, key: str) -> npt.NDArray[Any]:
+        raise NotImplementedError()
+
+    def get_adjacency_matrix(self) -> sp.coo_matrix:
+        n_nodes = len(self)
+        row, col = self.pid()[1:], self.id()[1:]  # ignore root
+        triad = (np.ones_like(row), (row, col))
+        return sp.coo_matrix(triad, shape=(n_nodes, n_nodes), dtype=np.int32)
+
+    def number_of_nodes(self) -> int:
+        """Get the number of nodes."""
+        return self.id().shape[0]
+
+    def number_of_edges(self) -> int:
+        """Get the number of edges."""
+        return self.number_of_nodes() - 1  # for tree structure: n = e + 1
+
+    @overload
+    def to_swc(
+        self,
+        fname: str,
+        *,
+        extra_cols: list[str] | None = ...,
+        source: bool | str = ...,
+        id_offset: int = ...,
+    ) -> None: ...
+    @overload
+    def to_swc(
+        self,
+        *,
+        extra_cols: list[str] | None = ...,
+        source: bool | str = ...,
+        id_offset: int = ...,
+    ) -> str: ...
+    def to_swc(
+        self,
+        fname: str | None = None,
+        *,
+        extra_cols: list[str] | None = None,
+        source: bool | str = True,
+        comments: bool = True,
+        id_offset: int = 1,
+    ) -> str | None:
+        """Write swc file."""
+        data = []
+        if source is not False:
+            if not isinstance(source, str):
+                source = self.source if self.source else "Unknown"
+            data.append(f"source: {source}")
+            data.append("")
+
+        if comments is True:
+            data.extend(self.comments)
+
+        it = to_swc(
+            self.get_ndata, comments=data, extra_cols=extra_cols, id_offset=id_offset
+        )
+
+        if fname is None:
+            return "".join(it)
+
+        with open(fname, "w", encoding="utf-8") as f:
+            f.writelines(it)
+
+        return None
+
+    @overload
+    def to_eswc(self, fname: str, **kwargs) -> None: ...
+    @overload
+    def to_eswc(self, fname: None = ..., **kwargs) -> str: ...
+    def to_eswc(
+        self,
+        fname: str | None = None,
+        swc_path: str | None = None,
+        extra_cols: list[str] | None = None,
+        **kwargs,
+    ) -> str | None:
+        if swc_path is None:
+            warnings.warn(
+                "`swc_path` has been renamed to `fname` since v0.5.1, "
+                "and will be removed in next version",
+                DeprecationWarning,
+            )
+            fname = swc_path
+
+        extra_cols = extra_cols or []
+        extra_cols.extend(k for k, _ in eswc_cols)
+        return self.to_swc(fname, extra_cols=extra_cols, **kwargs)  # type: ignore
+
+
+SWCTypeVar = TypeVar("SWCTypeVar", bound=SWCLike)
+
+
+class DictSWC(SWCLike):
+    """SWC implementation on dict."""
+
+    ndata: dict[str, npt.NDArray]
+
+    def __init__(
+        self,
+        *,
+        source: str = "",
+        comments: Iterable[str] | None = None,
+        names: SWCNames | None = None,
+        **kwargs: npt.NDArray,
+    ):
+        super().__init__()
+        self.source = source
+        self.comments = list(comments) if comments is not None else []
+        self.names = get_names(names)
+        self.ndata = kwargs
+
+    def keys(self) -> Iterable[str]:
+        return self.ndata.keys()
+
+    def values(self) -> Iterable[npt.NDArray[Any]]:
+        return self.ndata.values()
+
+    def items(self) -> Iterable[tuple[str, npt.NDArray[Any]]]:
+        return self.ndata.items()
+
+    def get_ndata(self, key: str) -> npt.NDArray[Any]:
+        return self.ndata[key]
+
+    def copy(self) -> Self:
+        """Make a copy."""
+        return deepcopy(self)

swcgeom/core/swc_utils/__init__.py

@@ -0,0 +1,19 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWC format utils.
+
+NOTE: This module provides a bunch of methods to manipulating swc files, they are
+always trivial and unstabled, so we are NOT export it by default. If you use the method
+here, please review the code more frequently, we will try to flag all breaking changes
+but NO promises.
+"""
+
+from swcgeom.core.swc_utils.assembler import *  # noqa: F403
+from swcgeom.core.swc_utils.base import *  # noqa: F403
+from swcgeom.core.swc_utils.checker import *  # noqa: F403
+from swcgeom.core.swc_utils.io import *  # noqa: F403
+from swcgeom.core.swc_utils.normalizer import *  # noqa: F403
+from swcgeom.core.swc_utils.subtree import *  # noqa: F403

swcgeom/core/swc_utils/assembler.py

@@ -0,0 +1,35 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Assemble lines to swc.
+
+NOTE: This module is deprecated, please use `~.transforms.LinesToTree` instead.
+"""
+
+__all__ = ["assemble_lines", "try_assemble_lines"]
+
+
+def assemble_lines(*args, **kwargs):
+    """Assemble lines to tree.
+
+    .. deprecated:: 0.15.0
+        Use :meth:`~.transforms.LinesToTree` instead.
+    """
+    raise DeprecationWarning(
+        "`assemble_lines` has been replaced by `~.transforms.LinesToTree` because it "
+        "can be easy assemble with other tansformations.",
+    )
+
+
+def try_assemble_lines(*args, **kwargs):
+    """Try assemble lines to tree.
+
+    .. deprecated:: 0.15.0
+        Use :meth:`~.transforms.LinesToTree` instead.
+    """
+    raise DeprecationWarning(
+        "`try_assemble_lines` has been replaced by `~.transforms.LinesToTree` because "
+        "it can be easy assemble with other tansformations.",
+    )

swcgeom/core/swc_utils/base.py

@@ -0,0 +1,180 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Base SWC format utils."""
+
+from collections.abc import Callable
+from typing import Literal, NamedTuple, TypeVar, overload
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+__all__ = [
+    "Topology",
+    "SWCNames",
+    "swc_names",  # may not need to export
+    "get_names",
+    "SWCTypes",
+    "get_types",
+    "get_topology",
+    "get_dsu",
+    "traverse",
+]
+
+T = TypeVar("T")
+K = TypeVar("K")
+Topology = tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]  # (id, pid)
+
+
+class SWCNames(NamedTuple):
+    """SWC format column names."""
+
+    id: str = "id"
+    type: str = "type"
+    x: str = "x"
+    y: str = "y"
+    z: str = "z"
+    r: str = "r"
+    pid: str = "pid"
+
+    def cols(self) -> list[str]:
+        return [self.id, self.type, self.x, self.y, self.z, self.r, self.pid]
+
+
+swc_names = SWCNames()
+
+
+def get_names(names: SWCNames | None = None) -> SWCNames:
+    return names or swc_names
+
+
+class SWCTypes(NamedTuple):
+    """SWC format types.
+
+    See Also:
+        NeuroMoprho.org - What is SWC format?: https://neuromorpho.org/myfaq.jsp
+    """
+
+    undefined: int = 0
+    soma: int = 1
+    axon: int = 2
+    basal_dendrite: int = 3
+    apical_dendrite: int = 4
+    custom: int = 5  # user-defined preferences
+    unspecified_neurites: int = 6
+    glia_processes: int = 7
+
+
+swc_types = SWCTypes()
+
+
+def get_types(types: SWCTypes | None = None) -> SWCTypes:
+    return types or swc_types
+
+
+def get_topology(df: pd.DataFrame, *, names: SWCNames | None = None) -> Topology:
+    names = get_names(names)
+    return (df[names.id].to_numpy(), df[names.pid].to_numpy())
+
+
+def get_dsu(
+    df: pd.DataFrame, *, names: SWCNames | None = None
+) -> npt.NDArray[np.int32]:
+    """Get disjoint set union."""
+    names = get_names(names)
+    dsu = np.where(
+        df[names.pid] == -1, df[names.id], df[names.pid]
+    )  # Disjoint Set Union
+
+    id2idx = dict(zip(df[names.id], range(len(df))))
+    dsu = np.array([id2idx[i] for i in dsu], dtype=np.int32)
+
+    while True:
+        flag = True
+        for i, p in enumerate(dsu):
+            if dsu[i] != dsu[p]:
+                dsu[i] = dsu[p]
+                flag = False
+
+        if flag:
+            break
+
+    return dsu
+
+
+@overload
+def traverse(
+    topology: Topology,
+    *,
+    enter: Callable[[int, T | None], T],
+    root: int | np.integer = ...,
+    mode: Literal["dfs"] = ...,
+) -> None: ...
+@overload
+def traverse(
+    topology: Topology,
+    *,
+    leave: Callable[[int, list[K]], K],
+    root: int | np.integer = ...,
+    mode: Literal["dfs"] = ...,
+) -> K: ...
+@overload
+def traverse(
+    topology: Topology,
+    *,
+    enter: Callable[[int, T | None], T],
+    leave: Callable[[int, list[K]], K],
+    root: int | np.integer = ...,
+    mode: Literal["dfs"] = ...,
+) -> K: ...
+def traverse(topology: Topology, *, mode="dfs", **kwargs):
+    """Traverse nodes.
+
+    Args:
+    enter: (id: int, parent: T | None) => T, optional
+        The callback when entering node, which accepts two parameters, the current node
+        id and the return value of it parent node. In particular, the root node
+        receives an `None`.
+    leave: (id: int, children: list[T]) => T, optional
+        The callback when leaving node. When leaving a node, subtree has already been
+        traversed. Callback accepts two parameters, the current node id and list of the
+        return value of children, In particular, the leaf node receives an empty list.
+    root: Start from the root node of the subtree
+    mode: The traverse mode, only support "dfs" now.
+    """
+    match mode:
+        case "dfs":
+            return _traverse_dfs(topology, **kwargs)
+        case _:
+            raise ValueError(f"unsupported mode: `{mode}`")
+
+
+def _traverse_dfs(topology: Topology, *, enter=None, leave=None, root=0):
+    """Traverse each nodes by dfs."""
+    children_map = dict[int, list[int]]()
+    for idx, pid in zip(*topology):
+        children_map.setdefault(pid, [])
+        children_map[pid].append(idx)
+
+    # manual dfs to avoid stack overflow in long branch
+    stack: list[tuple[int, bool]] = [(root, True)]  # (idx, is_enter)
+    params = {root: None}
+    vals = {}
+
+    while len(stack) != 0:
+        idx, is_enter = stack.pop()
+        if is_enter:  # enter
+            pre = params.pop(idx)
+            cur = enter(idx, pre) if enter is not None else None
+            stack.append((idx, False))
+            for child in children_map.get(idx, []):
+                stack.append((child, True))
+                params[child] = cur
+        else:  # leave
+            children = [vals.pop(i) for i in children_map.get(idx, [])]
+            vals[idx] = leave(idx, children) if leave is not None else None
+
+    return vals[root]

swcgeom/core/swc_utils/checker.py

@@ -0,0 +1,107 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Check common"""
+
+from collections import defaultdict
+
+import numpy as np
+import pandas as pd
+from typing_extensions import deprecated
+
+from swcgeom.core.swc_utils.base import SWCNames, Topology, get_dsu, get_names, traverse
+from swcgeom.utils import DisjointSetUnion
+
+__all__ = [
+    "is_single_root",
+    "is_bifurcate",
+    "is_sorted",
+    "has_cyclic",
+    # legacy
+    "is_binary_tree",
+    "check_single_root",
+]
+
+
+def is_single_root(df: pd.DataFrame, *, names: SWCNames | None = None) -> bool:
+    """Check is it only one root."""
+    return len(np.unique(get_dsu(df, names=names))) == 1
+
+
+def is_bifurcate(topology: Topology, *, exclude_root: bool = True) -> bool:
+    """Check is it a bifurcate topology."""
+    children = defaultdict(list)
+    for idx, pid in zip(*topology):
+        children[pid].append(idx)
+
+    root = children[-1]
+    for k, v in children.items():
+        if len(v) > 1 and (not exclude_root or k in root):
+            return False
+
+    return True
+
+
+def is_sorted(topology: Topology) -> bool:
+    """Check is it sorted.
+
+    In a sorted topology, parent samples should appear before any child samples.
+    """
+    flag = True
+
+    def enter(idx: int, parent: int | None) -> int:
+        nonlocal flag
+        if parent is not None and idx < parent:
+            flag = False
+
+        return idx
+
+    traverse(topology=topology, enter=enter)
+    return flag
+
+
+def has_cyclic(topology: Topology) -> bool:
+    """Has cyclic in topology."""
+    node_num = len(topology[0])
+    dsu = DisjointSetUnion(node_number=node_num)
+
+    for i in range(node_num):
+        node_a = topology[0][i]
+        node_b = topology[1][i]
+        # skip the root node
+        if node_b == -1:
+            continue
+
+        # check whether it is circle
+        if dsu.is_same_set(node_a, node_b):
+            return True
+
+        dsu.union_sets(node_a, node_b)
+
+    return False
+
+
+@deprecated("Use `is_single_root` instead")
+def check_single_root(*args, **kwargs) -> bool:
+    """Check if the tree is single root.
+
+    .. deprecated:: 0.5.0
+        Use :meth:`is_single_root` instead.
+    """
+    return is_single_root(*args, **kwargs)
+
+
+@deprecated("Use `is_bifurcate` instead")
+def is_binary_tree(
+    df: pd.DataFrame, exclude_root: bool = True, *, names: SWCNames | None = None
+) -> bool:
+    """Check is it a binary tree.
+
+    .. deprecated:: 0.8.0
+        Use :meth:`is_bifurcate` instead.
+    """
+    names = get_names(names)
+    topo = (df[names.id].to_numpy(), df[names.pid].to_numpy())
+    return is_bifurcate(topo, exclude_root=exclude_root)