swcgeom 0.20.0__cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl

swcgeom/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""A neuron geometry library for swc format."""
+
+from swcgeom import analysis, core, images, transforms
+from swcgeom.analysis import draw
+from swcgeom.core import BranchTree, Population, Populations, Tree
+
+__all__ = [
+    "analysis",
+    "core",
+    "images",
+    "transforms",
+    "draw",
+    "BranchTree",
+    "Population",
+    "Populations",
+    "Tree",
+]

swcgeom/analysis/__init__.py

@@ -0,0 +1,13 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Analysis for neuron trees."""
+
+from swcgeom.analysis.feature_extractor import *  # noqa: F403
+from swcgeom.analysis.features import *  # noqa: F403
+from swcgeom.analysis.sholl import *  # noqa: F403
+from swcgeom.analysis.trunk import *  # noqa: F403
+from swcgeom.analysis.visualization import *  # noqa: F403
+from swcgeom.analysis.volume import *  # noqa: F403

swcgeom/analysis/feature_extractor.py

@@ -0,0 +1,454 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Easy way to compute and visualize common features for feature.
+
+NOTE: For development, see method `Features.get_evaluator` to confirm the naming
+specification.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from functools import cached_property
+from itertools import chain
+from os.path import basename
+from typing import Any, Literal, overload
+
+import numpy as np
+import numpy.typing as npt
+import seaborn as sns
+from matplotlib.axes import Axes
+
+from swcgeom.analysis.features import (
+    BranchFeatures,
+    FurcationFeatures,
+    NodeFeatures,
+    PathFeatures,
+    TipFeatures,
+)
+from swcgeom.analysis.sholl import Sholl
+from swcgeom.analysis.volume import get_volume
+from swcgeom.core import Population, Populations, Tree
+from swcgeom.utils import padding1d
+
+__all__ = ["Feature", "extract_feature"]
+
+Feature = Literal[
+    "length",
+    "volume",
+    "sholl",
+    # node
+    "node_count",
+    "node_radial_distance",
+    "node_branch_order",
+    # furcation nodes
+    "furcation_count",
+    "furcation_radial_distance",
+    # bifurcation nodes
+    "bifurcation_count",
+    "bifurcation_radial_distance",
+    # tip nodes
+    "tip_count",
+    "tip_radial_distance",
+    # branch
+    "branch_length",
+    "branch_tortuosity",
+    # path
+    "path_length",
+    "path_tortuosity",
+]
+
+NDArrayf32 = npt.NDArray[np.float32]
+FeatAndKwargs = Feature | tuple[Feature, dict[str, Any]]
+
+Feature1D = set(["length", "volume", "node_count", "furcation_count", "tip_count"])
+
+
+class Features:
+    """Tree features"""
+
+    tree: Tree
+
+    # Modules
+    @cached_property
+    def node_features(self) -> NodeFeatures:
+        return NodeFeatures(self.tree)
+
+    @cached_property
+    def furcation_features(self) -> FurcationFeatures:
+        return FurcationFeatures(self.node_features)
+
+    @cached_property
+    def tip_features(self) -> TipFeatures:
+        return TipFeatures(self.node_features)
+
+    @cached_property
+    def branch_features(self) -> BranchFeatures:
+        return BranchFeatures(self.tree)
+
+    @cached_property
+    def path_features(self) -> PathFeatures:
+        return PathFeatures(self.tree)
+
+    # Caches
+    @cached_property
+    def sholl(self) -> Sholl:
+        return Sholl(self.tree)
+
+    def __init__(self, tree: Tree) -> None:
+        self.tree = tree
+
+    def get(self, feature: FeatAndKwargs, **kwargs) -> NDArrayf32:
+        feat, kwargs = _get_feat_and_kwargs(feature, **kwargs)
+        evaluator = self.get_evaluator(feat)
+        return evaluator(**kwargs)
+
+    def get_evaluator(self, feature: Feature) -> Callable[[], npt.NDArray]:
+        if callable(calc := getattr(self, f"get_{feature}", None)):
+            return calc  # custom features
+
+        components = feature.split("_")
+        if (module := getattr(self, f"{components[0]}_features", None)) and callable(
+            calc := getattr(module, f"get_{'_'.join(components[1:])}", None)
+        ):
+            return calc
+
+        raise ValueError(f"Invalid feature: {feature}")
+
+    # Custom Features
+
+    def get_length(self, **kwargs) -> NDArrayf32:
+        return np.array([self.tree.length(**kwargs)], dtype=np.float32)
+
+    def get_volume(self, **kwargs) -> NDArrayf32:
+        return np.array([get_volume(self.tree, **kwargs)], dtype=np.float32)
+
+    def get_sholl(self, **kwargs) -> NDArrayf32:
+        return self.sholl.get(**kwargs).astype(np.float32)
+
+
+class FeatureExtractor(ABC):
+    """Extract features from tree."""
+
+    @overload
+    def get(self, feature: Feature, **kwargs) -> NDArrayf32: ...
+    @overload
+    def get(self, feature: list[FeatAndKwargs]) -> list[NDArrayf32]: ...
+    @overload
+    def get(self, feature: dict[Feature, dict[str, Any]]) -> dict[str, NDArrayf32]: ...
+    def get(self, feature, **kwargs):
+        """Get feature.
+
+        NOTE: Shape of returned array is not uniform, `TreeFeatureExtractor` returns
+        array of shape (L, ), `PopulationFeatureExtracor` returns array of shape (N, L).
+        """
+        if isinstance(feature, dict):
+            return {k: self._get(k, **v) for k, v in feature.items()}
+
+        if isinstance(feature, list):
+            return [self._get(k) for k in feature]
+
+        return self._get(feature, **kwargs)
+
+    def plot(self, feature: FeatAndKwargs, title: str | bool = True, **kwargs) -> Axes:
+        """Plot feature with appropriate way.
+
+        NOTE: The drawing method is different in different classes, different in
+        different features, and may different between versions, there are NO guarantees.
+        """
+        feat, feat_kwargs = _get_feat_and_kwargs(feature)
+        if callable(custom_plot := getattr(self, f"plot_{feat}", None)):
+            ax = custom_plot(feat_kwargs, **kwargs)
+        elif feat in Feature1D:
+            ax = self._plot_1d(feature, **kwargs)
+        else:
+            ax = self._plot_histogram(feature, **kwargs)  # default plot
+
+        if isinstance(title, str):
+            ax.set_title(title)
+        elif title is True:
+            ax.set_title(_get_feature_name(feat))
+
+        return ax
+
+    # Custom Plots
+
+    def plot_node_branch_order(self, feature_kwargs: dict[str, Any], **kwargs) -> Axes:
+        vals = self._get("node_branch_order", **feature_kwargs)
+        bin_edges = np.arange(int(np.ceil(vals.max() + 1))) + 0.5
+        return self._plot_histogram_impl(vals, bin_edges, **kwargs)
+
+    # Implements
+
+    def _get(self, feature: FeatAndKwargs, **kwargs) -> NDArrayf32:
+        feat, kwargs = _get_feat_and_kwargs(feature, **kwargs)
+        if callable(custom_get := getattr(self, f"get_{feat}", None)):
+            return custom_get(**kwargs)
+
+        return self._get_impl(feat, **kwargs)  # default
+
+    def _plot_1d(self, feature: FeatAndKwargs, **kwargs) -> Axes:
+        vals = self._get(feature)
+        ax = self._plot_1d_impl(vals, **kwargs)
+        ax.set_ylabel(_get_feature_name(feature))
+        return ax
+
+    def _plot_histogram(
+        self,
+        feature: FeatAndKwargs,
+        bins=20,
+        range=None,  # pylint: disable=redefined-builtin
+        **kwargs,
+    ) -> Axes:
+        vals = self._get(feature)
+        bin_edges = np.histogram_bin_edges(vals, bins, range)
+        return self._plot_histogram_impl(vals, bin_edges, **kwargs)
+
+    @abstractmethod
+    def _get_impl(self, feature: Feature, **kwargs) -> NDArrayf32:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _plot_1d_impl(self, vals: NDArrayf32, **kwargs) -> Axes:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _plot_histogram_impl(
+        self, vals: NDArrayf32, bin_edges: npt.NDArray, **kwargs
+    ) -> Axes:
+        raise NotImplementedError()
+
+    def get_bifurcation_count(self, **kwargs):
+        raise DeprecationWarning("Use `furcation_count` instead.")
+
+    def get_bifurcation_radial_distance(self, **kwargs):
+        raise DeprecationWarning("Use `furcation_radial_distance` instead.")
+
+
+class TreeFeatureExtractor(FeatureExtractor):
+    """Extract feature from tree."""
+
+    _tree: Tree
+    _features: Features
+
+    def __init__(self, tree: Tree) -> None:
+        super().__init__()
+        self._tree = tree
+        self._features = Features(tree)
+
+    # Custom Features
+
+    def get_sholl(self, **kwargs) -> NDArrayf32:
+        return self._features.sholl.get(**kwargs).astype(np.float32)
+
+    # Custom Plots
+
+    def plot_sholl(
+        self,
+        feature_kwargs: dict[str, Any],  # pylint: disable=unused-argument
+        **kwargs,
+    ) -> Axes:
+        _, ax = self._features.sholl.plot(**kwargs)
+        return ax
+
+    # Implements
+
+    def _get_impl(self, feature: Feature, **kwargs) -> NDArrayf32:
+        return self._features.get(feature, **kwargs)
+
+    def _plot_histogram_impl(
+        self, vals: NDArrayf32, bin_edges: npt.NDArray, **kwargs
+    ) -> Axes:
+        hist, _ = np.histogram(vals[vals != 0], bins=bin_edges)
+        x = (bin_edges[:-1] + bin_edges[1:]) / 2
+
+        ax: Axes = sns.barplot(x=x, y=hist, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
+
+    def _plot_1d_impl(self, vals: NDArrayf32, **kwargs) -> Axes:
+        name = basename(self._tree.source)
+        return sns.barplot(x=[name], y=vals.squeeze(), **kwargs)
+
+
+class PopulationFeatureExtractor(FeatureExtractor):
+    """Extract features from population."""
+
+    _population: Population
+    _features: list[Features]
+
+    def __init__(self, population: Population) -> None:
+        super().__init__()
+        self._population = population
+        self._features = [Features(t) for t in self._population]
+
+    # Custom Features
+
+    def get_sholl(self, **kwargs) -> NDArrayf32:
+        vals, _ = self._get_sholl_impl(**kwargs)
+        return vals
+
+    # Custom Plots
+
+    def plot_sholl(self, feature_kwargs: dict[str, Any], **kwargs) -> Axes:
+        vals, rs = self._get_sholl_impl(**feature_kwargs)
+        ax = self._lineplot(xs=rs, ys=vals.flatten(), **kwargs)
+        ax.set_ylabel("Count of Intersections")
+        return ax
+
+    # Implements
+
+    def _get_impl(self, feature: Feature, **kwargs) -> NDArrayf32:
+        vals = [f.get(feature, **kwargs) for f in self._features]
+        len_max = max(len(v) for v in vals)
+        v = np.stack([padding1d(len_max, v, dtype=np.float32) for v in vals])
+        return v
+
+    def _get_sholl_impl(
+        self, steps: int = 20, **kwargs
+    ) -> tuple[NDArrayf32, NDArrayf32]:
+        rmax = max(t.sholl.rmax for t in self._features)
+        rs = Sholl.get_rs(rmax=rmax, steps=steps)
+        vals = self._get_impl("sholl", steps=rs, **kwargs)
+        return vals, rs
+
+    def _plot_histogram_impl(
+        self, vals: NDArrayf32, bin_edges: npt.NDArray, **kwargs
+    ) -> Axes:
+        def hist(v):
+            return np.histogram(v[v != 0], bins=bin_edges)[0]
+
+        xs = (bin_edges[:-1] + bin_edges[1:]) / 2
+        ys = np.stack([hist(v) for v in vals])
+
+        ax: Axes = self._lineplot(xs, ys, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
+
+    def _plot_1d_impl(self, vals: NDArrayf32, **kwargs) -> Axes:
+        x = [basename(t.source) for t in self._population]
+        y = vals.flatten()
+        ax: Axes = sns.barplot(x=x, y=y, **kwargs)
+        ax.axhline(y=y.mean(), ls="--", lw=1)
+        ax.set_xticks([])
+        return ax
+
+    def _lineplot(self, xs, ys, **kwargs) -> Axes:
+        xs = np.tile(xs, len(self._population))
+        ys = ys.flatten()
+        ax: Axes = sns.lineplot(x=xs, y=ys, **kwargs)
+        return ax
+
+
+class PopulationsFeatureExtractor(FeatureExtractor):
+    """Extract feature from population."""
+
+    _populations: Populations
+    _features: list[list[Features]]
+
+    def __init__(self, populations: Populations) -> None:
+        super().__init__()
+        self._populations = populations
+        self._features = [[Features(t) for t in p] for p in populations.populations]
+
+    # Custom Features
+
+    def get_sholl(self, **kwargs) -> NDArrayf32:
+        vals, _ = self._get_sholl_impl(**kwargs)
+        return vals
+
+    # Custom Plots
+
+    def plot_sholl(self, feature_kwargs: dict[str, Any], **kwargs) -> Axes:
+        vals, rs = self._get_sholl_impl(**feature_kwargs)
+        ax = self._lineplot(xs=rs, ys=vals, **kwargs)
+        ax.set_ylabel("Count of Intersections")
+        return ax
+
+    # Implements
+
+    def _get_impl(self, feature: Feature, **kwargs) -> NDArrayf32:
+        vals = [[f.get(feature, **kwargs) for f in fs] for fs in self._features]
+        len_max1 = max(len(v) for v in vals)
+        len_max2 = max(*chain.from_iterable(((len(vv) for vv in v) for v in vals)))
+        out = np.zeros((len(vals), len_max1, len_max2), dtype=np.float32)
+        for i, v in enumerate(vals):
+            for j, vv in enumerate(v):
+                out[i, j, : len(vv)] = vv
+
+        return out
+
+    def _get_sholl_impl(
+        self, steps: int = 20, **kwargs
+    ) -> tuple[NDArrayf32, NDArrayf32]:
+        rmaxs = chain.from_iterable((t.sholl.rmax for t in p) for p in self._features)
+        rmax = max(rmaxs)
+        rs = Sholl.get_rs(rmax=rmax, steps=steps)
+        vals = self._get_impl("sholl", steps=rs, **kwargs)
+        return vals, rs
+
+    def _plot_histogram_impl(
+        self, vals: NDArrayf32, bin_edges: npt.NDArray, **kwargs
+    ) -> Axes:
+        def hist(v):
+            return np.histogram(v[v != 0], bins=bin_edges)[0]
+
+        xs = (bin_edges[:-1] + bin_edges[1:]) / 2
+        ys = np.stack([np.stack([hist(t) for t in p]) for p in vals])
+
+        ax = self._lineplot(xs=xs, ys=ys, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
+
+    def _plot_1d_impl(self, vals: NDArrayf32, **kwargs) -> Axes:
+        labels = self._populations.labels
+        xs = np.concatenate([np.full(vals.shape[1], fill_value=i) for i in labels])
+        ys = vals.flatten()
+
+        # The numbers of tree in different populations may not be equal
+        valid = ys != 0
+        xs, ys = xs[valid], ys[valid]
+
+        ax: Axes = sns.boxplot(x=xs, y=ys, **kwargs)
+        return ax
+
+    def _lineplot(self, xs, ys, **kwargs) -> Axes:
+        p, t, f = ys.shape
+        labels = self._populations.labels  # (P,)
+        x = np.tile(xs, p * t)  # (F,) -> (P * T * F)
+        y = ys.flatten()  # (P, T, F) -> (P * T * F)
+        hue = np.concatenate([np.full(t * f, fill_value=i) for i in labels])
+
+        # The numbers of tree in different populations may not be equal
+        valid = np.repeat(np.any(ys != 0, axis=2), f)
+        x, y, hue = x[valid], y[valid], hue[valid]
+
+        ax: Axes = sns.lineplot(x=x, y=y, hue=hue, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
+
+
+def extract_feature(obj: Tree | Population | Populations) -> FeatureExtractor:
+    if isinstance(obj, Tree):
+        return TreeFeatureExtractor(obj)
+
+    if isinstance(obj, Population):
+        return PopulationFeatureExtractor(obj)
+
+    if isinstance(obj, Populations):
+        return PopulationsFeatureExtractor(obj)
+
+    raise TypeError("Invalid argument type.")
+
+
+def _get_feat_and_kwargs(feature: FeatAndKwargs, **kwargs):
+    if isinstance(feature, tuple):
+        return feature[0], {**feature[1], **kwargs}
+
+    return feature, kwargs
+
+
+def _get_feature_name(feature: FeatAndKwargs) -> str:
+    feat, _ = _get_feat_and_kwargs(feature)
+    return feat.replace("_", " ").title()

swcgeom/analysis/features.py

@@ -0,0 +1,218 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Feature analysis of tree."""
+
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import TypeVar
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import Self, deprecated
+
+from swcgeom.core import Branch, BranchTree, Tree
+
+__all__ = [
+    "NodeFeatures",
+    "BifurcationFeatures",
+    "TipFeatures",
+    "PathFeatures",
+    "BranchFeatures",
+]
+
+T = TypeVar("T", bound=Branch)
+
+# Node Level
+
+
+class NodeFeatures:
+    """Evaluate node feature of tree."""
+
+    tree: Tree
+
+    @cached_property
+    def _branch_tree(self) -> BranchTree:
+        return BranchTree.from_tree(self.tree)
+
+    def __init__(self, tree: Tree) -> None:
+        self.tree = tree
+
+    def get_count(self) -> npt.NDArray[np.float32]:
+        """Get number of nodes.
+
+        Args:
+            count: array of shape (1,)
+        """
+        return np.array([self.tree.number_of_nodes()], dtype=np.float32)
+
+    def get_radial_distance(self) -> npt.NDArray[np.float32]:
+        """Get the end-to-end straight-line distance to soma.
+
+        Returns:
+            radial_distance: Array of shape (N,).
+        """
+        xyz = self.tree.xyz() - self.tree.soma().xyz()
+        radial_distance = np.linalg.norm(xyz, axis=1)
+        return radial_distance
+
+    def get_branch_order(self) -> npt.NDArray[np.int32]:
+        """Get branch order of criticle nodes of tree.
+
+        Branch order is the number of bifurcations between current
+        position and the root.
+
+        Criticle node means that soma, bifucation nodes, tips.
+
+        Returns:
+            order: Array of shape (N,), which k is the number of branches.
+        """
+        order = np.zeros_like(self._branch_tree.id(), dtype=np.int32)
+
+        def assign_depth(n: Tree.Node, pre_depth: int | None) -> int:
+            cur_order = pre_depth + 1 if pre_depth is not None else 0
+            order[n.id] = cur_order
+            return cur_order
+
+        self._branch_tree.traverse(enter=assign_depth)
+        return order
+
+
+class _SubsetNodesFeatures(ABC):
+    _features: NodeFeatures
+
+    @property
+    @abstractmethod
+    def nodes(self) -> npt.NDArray[np.bool_]:
+        raise NotImplementedError()
+
+    def __init__(self, features: NodeFeatures) -> None:
+        self._features = features
+
+    def get_count(self) -> npt.NDArray[np.float32]:
+        """Get number of nodes.
+
+        Returns:
+            count: Array of shape (1,).
+        """
+        return np.array([np.count_nonzero(self.nodes)], dtype=np.float32)
+
+    def get_radial_distance(self) -> npt.NDArray[np.float32]:
+        """Get the end-to-end straight-line distance to soma.
+
+        Returns:
+            radial_distance: Array of shape (N,).
+        """
+        return self._features.get_radial_distance()[self.nodes]
+
+    @classmethod
+    def from_tree(cls, tree: Tree) -> Self:
+        return cls(NodeFeatures(tree))
+
+
+class FurcationFeatures(_SubsetNodesFeatures):
+    """Evaluate furcation node feature of tree."""
+
+    @cached_property
+    def nodes(self) -> npt.NDArray[np.bool_]:
+        return np.array([n.is_furcation() for n in self._features.tree])
+
+
+@deprecated("Use FurcationFeatures instead")
+class BifurcationFeatures(FurcationFeatures):
+    """Evaluate bifurcation node feature of tree.
+
+    NOTE: Deprecated due to the wrong spelling of furcation. For now, it is just an
+    alias of `FurcationFeatures` and raise a warning. It will be change to raise an
+    error in the future.
+    """
+
+
+class TipFeatures(_SubsetNodesFeatures):
+    """Evaluate tip node feature of tree."""
+
+    @cached_property
+    def nodes(self) -> npt.NDArray[np.bool_]:
+        return np.array([n.is_tip() for n in self._features.tree])
+
+
+# Path Level
+
+
+class PathFeatures:
+    """Path analysis of tree."""
+
+    tree: Tree
+
+    def __init__(self, tree: Tree) -> None:
+        self.tree = tree
+
+    def get_count(self) -> int:
+        return len(self._paths)
+
+    def get_length(self) -> npt.NDArray[np.float32]:
+        """Get length of paths."""
+
+        length = [path.length() for path in self._paths]
+        return np.array(length, dtype=np.float32)
+
+    def get_tortuosity(self) -> npt.NDArray[np.float32]:
+        """Get tortuosity of path."""
+
+        return np.array([path.tortuosity() for path in self._paths], dtype=np.float32)
+
+    @cached_property
+    def _paths(self) -> list[Tree.Path]:
+        return self.tree.get_paths()
+
+
+class BranchFeatures:
+    """Analysis bransh of tree."""
+
+    tree: Tree
+
+    def __init__(self, tree: Tree) -> None:
+        self.tree = tree
+
+    def get_count(self) -> int:
+        return len(self._branches)
+
+    def get_length(self) -> npt.NDArray[np.float32]:
+        """Get length of branches."""
+
+        length = [br.length() for br in self._branches]
+        return np.array(length, dtype=np.float32)
+
+    def get_tortuosity(self) -> npt.NDArray[np.float32]:
+        """Get tortuosity of path."""
+
+        return np.array([br.tortuosity() for br in self._branches], dtype=np.float32)
+
+    def get_angle(self, eps: float = 1e-7) -> npt.NDArray[np.float32]:
+        """Get agnle between branches.
+
+        Returns:
+            angle: An array of shape (N, N), which N is length of branches.
+        """
+        return self.calc_angle(self._branches, eps=eps)
+
+    @staticmethod
+    def calc_angle(branches: list[T], eps: float = 1e-7) -> npt.NDArray[np.float32]:
+        """Calc agnle between branches.
+
+        Returns:
+            angle: An array of shape (N, N), which N is length of branches.
+        """
+        vector = np.array([br[-1].xyz() - br[0].xyz() for br in branches])
+        vector_dot = np.matmul(vector, vector.T)
+        vector_norm = np.linalg.norm(vector, ord=2, axis=1, keepdims=True)
+        vector_norm_dot = np.matmul(vector_norm, vector_norm.T) + eps
+        arccos = np.clip(vector_dot / vector_norm_dot, -1, 1)
+        angle = np.arccos(arccos)
+        return angle
+
+    @cached_property
+    def _branches(self) -> list[Tree.Branch]:
+        return self.tree.get_branches()