swcgeom 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

swcgeom/_version.py

@@ -1,4 +1,4 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = version = '0.4.0'
-__version_tuple__ = version_tuple = (0, 4, 0)
+__version__ = version = '0.5.0'
+__version_tuple__ = version_tuple = (0, 5, 0)

swcgeom/analysis/feature_extractor.py

@@ -2,32 +2,28 @@
 
 Notes
 -----
-For development, see method `Features.get_evaluator`
-to confirm the naming specification.
+For development, see method `Features.get_evaluator` to confirm the
+naming specification.
 """
 
 from functools import cached_property
-from typing import Any, Callable, Dict, List, Literal, Tuple, cast, overload
+from itertools import chain
+from os.path import basename
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, overload
 
 import numpy as np
 import numpy.typing as npt
 import seaborn as sns
 from matplotlib.axes import Axes
-from matplotlib.figure import Figure
 
-from ..core import Population, Tree
-from ..utils import XYPair, get_fig_ax, padding1d, to_distribution
+from ..core import Population, Populations, Tree
+from ..utils import padding1d
 from .branch_features import BranchFeatures
 from .node_features import NodeFeatures
 from .path_features import PathFeatures
 from .sholl import Sholl
 
-__all__ = [
-    "Feature",
-    "FeatureExtractor",
-    "PopulationFeatureExtractor",
-    "extract_feature",
-]
+__all__ = ["Feature", "extract_feature"]
 
 Feature = Literal[
     "length",
@@ -43,6 +39,11 @@ Feature = Literal[
     "path_tortuosity",
 ]
 
+Bins = int | npt.ArrayLike | str
+Range = Optional[Tuple[float, float]]
+HistAndBinEdges = Tuple[npt.NDArray, npt.NDArray]
+FeatureWithKwargs = Feature | Tuple[Feature, Dict[str, Any]]
+
 
 class Features:
     """Tree features"""
@@ -62,10 +63,10 @@ class Features:
     def __init__(self, tree: Tree) -> None:
         self.tree = tree
 
-    def get(self, feature: Feature, **kwargs) -> npt.NDArray[np.float32]:
-        evaluator = self.get_evaluator(feature)
-        feat = evaluator(**kwargs).astype(np.float32)
-        return feat
+    def get(self, feature: FeatureWithKwargs, **kwargs) -> npt.NDArray[np.float32]:
+        feat, kwargs = _get_feature_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)):
@@ -79,196 +80,204 @@ class Features:
 
         raise ValueError(f"Invalid feature: {feature}")
 
-    def get_distribution(
-        self, feature: Feature, step: float | None, **kwargs
-    ) -> XYPair:
-        if callable(method := getattr(self, f"get_{feature}_distribution", None)):
-            if step is not None:
-                kwargs.setdefault("step", step)
-            return method(**kwargs)  # custom feature distribution
-
-        feat = self.get(feature, **kwargs)
-        step = cast(float, feat.max() / 100) if step is None else step
-        x, y = to_distribution(feat, step)
-        return x, y
-
     # Features
 
     def get_length(self, **kwargs) -> npt.NDArray[np.float32]:
         return np.array([self.tree.length(**kwargs)], dtype=np.float32)
 
-    def get_sholl(self, **kwargs) -> npt.NDArray[np.int32]:
-        return Sholl(self.tree, **kwargs).get_count()
-
-    def get_sholl_distribution(self, **kwargs) -> XYPair:
-        x, y = Sholl(self.tree, **kwargs).get_distribution()
-        return x, y.astype(np.float32)
+    def get_sholl(self, **kwargs) -> npt.NDArray[np.float32]:
+        return Sholl(self.tree, **kwargs).get().astype(np.float32)
 
 
 class FeatureExtractor:
-    """Extract feature from tree."""
-
-    features: Features
-
-    def __init__(self, tree: Tree) -> None:
-        self.features = Features(tree)
-
     # fmt:off
     @overload
     def get(self, feature: Feature, **kwargs) -> npt.NDArray[np.float32]: ...
     @overload
-    def get(self, feature: List[Feature], **kwargs) -> List[npt.NDArray[np.float32]]: ...
+    def get(self, feature: List[FeatureWithKwargs]) -> List[npt.NDArray[np.float32]]: ...
     @overload
-    def get(self, feature: Dict[Feature, Dict[str, Any]], **kwargs) -> Dict[str, npt.NDArray[np.float32]]: ...
+    def get(self, feature: Dict[Feature, Dict[str, Any]]) -> Dict[str, npt.NDArray[np.float32]]: ...
     # fmt:on
     def get(self, feature, **kwargs):
-        """Get feature of shape (L,)."""
+        """Get feature.
+
+        Notes
+        -----
+        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.features.get(k, **v) for k, v in feature.items()}
+            return {k: self._get(k, **v) for k, v in feature.items()}
 
         if isinstance(feature, list):
-            return [self.features.get(k) for k in feature]
+            return [self._get(k) for k in feature]
 
-        return self.features.get(feature, **kwargs)
+        return self._get(feature, **kwargs)
 
-    # fmt:off
-    @overload
-    def get_distribution(self, feature: Feature, step: float = ..., **kwargs) -> XYPair: ...
-    @overload
-    def get_distribution(self, feature: List[Feature], step: float = ..., **kwargs) -> List[XYPair]: ...
-    @overload
-    def get_distribution(self, feature: Dict[Feature, Dict[str, Any]], step: float = ..., **kwargs) -> Dict[str, XYPair]: ...
-    # fmt:on
-    def get_distribution(self, feature, step: float | None = None, **kwargs):
-        """Get feature distribution of shape (S,)."""
-        if isinstance(feature, dict):
-            return {
-                k: self.features.get_distribution(k, step, **v)
-                for k, v in feature.items()
-            }
+    def plot(
+        self, feature: FeatureWithKwargs, title: str | bool = True, **kwargs
+    ) -> Axes:  # TODO: sholl
+        feat, _ = _get_feature_and_kwargs(feature)
+        if not callable(plot := getattr(self, f"_plot_{feat}", None)):
+            plot = self._plot  # default plot
 
-        if isinstance(feature, list):
-            return [self.features.get_distribution(k, step) for k in feature]
+        vals = self._get(feature)
+        ax = plot(vals, **kwargs)
 
-        return self.features.get_distribution(feature, step, **kwargs)
+        if isinstance(title, str):
+            ax.set_title(title)
+        elif title is True:
+            ax.set_title(feat.replace("_", " ").title())
 
-    def plot_distribution(
-        self,
-        feature: Feature,
-        feature_args: Dict[Any, Any] = {},
-        fig: Figure | None = None,
-        ax: Axes | None = None,
-        **kwargs,
-    ) -> Tuple[Figure, Axes]:
-        x, y = self.get_distribution(feature, **feature_args)
-        fig, ax = get_fig_ax(fig, ax)
-        sns.lineplot(x=x, y=y, ax=ax, **kwargs)
-        return fig, ax
+        return ax
 
+    def _get(self, feature: FeatureWithKwargs, **kwargs) -> npt.NDArray[np.float32]:
+        raise NotImplementedError()
 
-class PopulationFeatureExtractor:
-    """Extract feature from population."""
+    def _plot(self, vals: npt.NDArray[np.float32], **kwargs) -> Axes:
+        raise NotImplementedError()
 
-    population: Population
 
-    @cached_property
-    def _trees(self) -> List[Features]:
-        return [Features(tree) for tree in self.population]
+class TreeFeatureExtractor(FeatureExtractor):
+    """Extract feature from tree."""
 
-    def __init__(self, population: Population) -> None:
-        self.population = population
+    _tree: Tree
+    _features: Features
 
-    # fmt:off
-    @overload
-    def get(self, feature: Feature, **kwargs) -> List[npt.NDArray[np.float32]]: ...
-    @overload
-    def get(self, feature: Dict[Feature, Dict[str, Any]], **kwargs) -> Dict[str, List[npt.NDArray[np.float32]]]: ...
-    # fmt:on
-    def get(self, feature, **kwargs):
-        """Get feature list of array of shape (N, L_i).
+    def __init__(self, tree: Tree) -> None:
+        super().__init__()
+        self._tree = tree
+        self._features = Features(tree)
 
-        Which N is the number of tree of population, L is length of
-        nodes.
-        """
-        if isinstance(feature, dict):
-            return {k: self._get(k, **v) for k, v in feature.items()}
+    def _get(self, feature: FeatureWithKwargs, **kwargs) -> npt.NDArray[np.float32]:
+        return self._features.get(feature, **kwargs)
 
-        return self._get(feature, **kwargs)
+    def _plot(self, vals: npt.NDArray[np.float32], **kwargs) -> Axes:
+        ax: Axes = sns.histplot(x=vals, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
 
-    # fmt:off
-    @overload
-    def get_distribution(self, feature: Feature, step: float = ..., **kwargs) -> XYPair: ...
-    @overload
-    def get_distribution(self, feature: List[Feature], step: float = ..., **kwargs) -> List[XYPair]: ...
-    @overload
-    def get_distribution(self, feature: Dict[Feature, Dict[str, Any]], step: float = ..., **kwargs) -> Dict[str, XYPair]: ...
-    # fmt:on
-    def get_distribution(self, feature, step: float | None = None, **kwargs):
-        """Get feature distribution of shape (N, S).
-
-        Which N is the number of tree of population, S is size of
-        distrtibution.
-
-        Returns
-        -------
-        x : npt.NDArray[np.float32]
-            Array of shape (S,).
-        y : npt.NDArray[np.float32]
-            Array of shape (N, S).
-        """
-        if isinstance(feature, dict):
-            return {
-                k: self._get_distribution(k, step=step, **v) for k, v in feature.items()
-            }
+    def _plot_length(self, vals: npt.NDArray[np.float32], **kwargs) -> Axes:
+        name = basename(self._tree.source)
+        ax: Axes = sns.barplot(x=[name], y=vals.squeeze(), **kwargs)
+        ax.set_ylabel("Length")
+        return ax
 
-        if isinstance(feature, list):
-            return [self._get_distribution(k, step=step) for k in feature]
-
-        return self._get_distribution(feature, step=step, **kwargs)
-
-    def plot_distribution(
-        self,
-        feature: Feature,
-        feature_args: Dict[Any, Any] = {},
-        fig: Figure | None = None,
-        ax: Axes | None = None,
-        **kwargs,
-    ) -> Tuple[Figure, Axes]:
-        x, y = self.get_distribution(feature, **feature_args)
-        x, y = np.tile(x, y.shape[0]), y.flatten()
-
-        fig, ax = get_fig_ax(fig, ax)
-        sns.lineplot(x=x, y=y, ax=ax, **kwargs)
-        return fig, ax
-
-    def _get(self, feature: Feature, **kwargs) -> List[npt.NDArray[np.float32]]:
-        return [ex.get(feature, **kwargs) for ex in self._trees]
-
-    def _get_distribution(self, feature: Feature, **kwargs) -> XYPair:
-        assert len(self._trees) != 0
-
-        x, ys = np.array([], dtype=np.float32), list[npt.NDArray[np.float32]]()
-        for features in self._trees:
-            xx, y = features.get_distribution(feature, **kwargs)
-            x = xx if xx.shape[0] > x.shape[0] else x
-            ys.append(y)
-
-        max_len_y = max(y.shape[0] for y in ys)
-        y = np.stack([padding1d(max_len_y, y, 0) for y in ys])
-        return x, y
-
-
-# fmt: off
-@overload
-def extract_feature(obj: Tree) -> FeatureExtractor: ...
-@overload
-def extract_feature(obj: Population) -> PopulationFeatureExtractor: ...
-# fmt: on
-def extract_feature(obj):
+
+class PopulationFeatureExtractor(FeatureExtractor):
+    """Extract feature 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]
+
+    def _get(self, feature: FeatureWithKwargs, **kwargs) -> npt.NDArray[np.float32]:
+        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 _plot(
+        self, vals: npt.NDArray[np.float32], bins="auto", range=None, **kwargs
+    ) -> Axes:
+        bin_edges = np.histogram_bin_edges(vals, bins, range)
+        hists = [
+            np.histogram(v, bins=bin_edges, weights=(v != 0).astype(np.int32))[0]
+            for v in vals
+        ]
+        hist = np.concatenate(hists)
+        x = np.tile((bin_edges[:-1] + bin_edges[1:]) / 2, len(self._population))
+
+        ax: Axes = sns.lineplot(x=x, y=hist, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
+
+    def _plot_length(self, vals: npt.NDArray[np.float32], **kwargs) -> Axes:
+        vals = vals.squeeze(axis=1)
+        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_ylabel("Length")
+        ax.set_xticks([])
+        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 self._populations.populations
+        ]
+
+    def _get(self, feature: FeatureWithKwargs, **kwargs) -> npt.NDArray[np.float32]:
+        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 _plot(
+        self, vals: npt.NDArray[np.float32], bins="auto", range=None, **kwargs
+    ) -> Axes:
+        bin_edges = np.histogram_bin_edges(vals, bins, range)
+        histogram = lambda v: np.histogram(
+            v, bins=bin_edges, weights=(v != 0).astype(np.int32)
+        )
+        hists = [[histogram(t)[0] for t in p] for p in vals]
+        hist = np.concatenate(hists).flatten()
+
+        repeats = np.prod(vals.shape[:2]).item()
+        x = np.tile((bin_edges[:-1] + bin_edges[1:]) / 2, repeats)
+
+        labels = self._populations.labels
+        length = (len(bin_edges) - 1) * vals.shape[1]
+        hue = np.concatenate([np.full(length, fill_value=i) for i in labels])
+
+        ax: Axes = sns.lineplot(x=x, y=hist, hue=hue, **kwargs)
+        ax.set_ylabel("Count")
+        return ax
+
+    def _plot_length(self, vals: npt.NDArray[np.float32], **kwargs) -> Axes:
+        vals = vals.squeeze(axis=2)
+        labels = self._populations.labels
+        x = np.concatenate([np.full(vals.shape[1], fill_value=i) for i in labels])
+        y = vals.flatten()
+        ax: Axes = sns.boxplot(x=x, y=y, **kwargs)
+        ax.set_ylabel("Length")
+        return ax
+
+
+def extract_feature(obj: Tree | Population) -> FeatureExtractor:
     if isinstance(obj, Tree):
-        return FeatureExtractor(obj)
+        return TreeFeatureExtractor(obj)
 
     if isinstance(obj, Population):
         return PopulationFeatureExtractor(obj)
 
+    if isinstance(obj, Populations):
+        return PopulationsFeatureExtractor(obj)
+
     raise TypeError("Invalid argument type.")
+
+
+def _get_feature_and_kwargs(feature: FeatureWithKwargs, **kwargs):
+    if isinstance(feature, tuple):
+        return feature[0], {**feature[1], **kwargs}
+    else:
+        return feature, kwargs

swcgeom/analysis/node_features.py

@@ -1,14 +1,11 @@
 """Depth distribution of tree."""
 
-
 from functools import cached_property
-from typing import List
 
 import numpy as np
 import numpy.typing as npt
 
 from ..core import BranchTree, Tree
-from ..utils import XYPair, to_distribution
 
 __all__ = ["NodeFeatures"]
 
@@ -27,54 +24,22 @@ class NodeFeatures:
         Returns
         -------
         radial_distance : npt.NDArray[np.float32]
-            Array of shape (N,), while N is the number of nodes.
+            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_radial_distance_distribution(
-        self,
-        step: float = 1,
-        /,
-        filter_bifurcation: bool = False,
-        filter_tip: bool = False,
-        filter_other: bool = True,
-    ) -> XYPair:
-        """Get radial distance distribution of tree.
-
-        Parameters
-        ----------
-        filter_bifurcation : bool, default `False`
-            Filter bifurcation nodes.
-        filter_tip : bool, default `False`
-            Filter tip nodes.
-        filter_other : bool, default `False`
-            Filter nodes that are not bifurcations or tips.
-
-        Returns
-        -------
-        radial_distance : npt.NDArray[np.float32]
-            Array of shape (N,), while N is the number of nodes.
-        """
-        return self._to_distribution(
-            self.get_radial_distance(),
-            step,
-            filter_bifurcation=filter_bifurcation,
-            filter_tip=filter_tip,
-            filter_other=filter_other,
-        )
-
     def get_branch_order(self) -> npt.NDArray[np.int32]:
         """Get branch order of tree.
 
-        Bifurcation order is the number of bifurcations between current position
-        and the root.
+        Bifurcation order is the number of bifurcations between current
+        position and the root.
 
         Returns
         -------
         order : npt.NDArray[np.int32]
-            Array of shape (N,), while N is the number of nodes.
+            Array of shape (k,), which k is the number of branchs.
         """
         order = np.zeros_like(self._branch_tree.id(), dtype=np.int32)
 
@@ -86,69 +51,6 @@ class NodeFeatures:
         self._branch_tree.traverse(enter=assign_depth)
         return order
 
-    def get_branch_order_distribution(
-        self,
-        step: int = 1,
-        /,
-        filter_bifurcation: bool = False,
-        filter_tip: bool = False,
-        filter_other: bool = True,
-    ) -> XYPair:
-        """Get branch order distribution of tree.
-
-        Parameters
-        ----------
-        filter_bifurcation : bool, default `False`
-            Filter bifurcation nodes.
-        filter_tip : bool, default `False`
-            Filter tip nodes.
-        filter_other : bool, default `False`
-            Filter nodes that are not bifurcations or tips.
-        """
-        return self._to_distribution(
-            self.get_branch_order(),
-            step,
-            filter_bifurcation=filter_bifurcation,
-            filter_tip=filter_tip,
-            filter_other=filter_other,
-        )
-
-    def _to_distribution(
-        self,
-        x: npt.NDArray,
-        step: float,
-        /,
-        filter_bifurcation: bool,
-        filter_tip: bool,
-        filter_other: bool,
-    ) -> XYPair:
-        if filter_bifurcation:
-            x[self._bifurcations] = -1
-
-        if filter_tip:
-            x[self._tips] = -1
-
-        if filter_other:
-            x[self._other] = -1
-
-        x = x[x != -1]
-        return to_distribution(x, step)
-
     @cached_property
     def _branch_tree(self) -> BranchTree:
         return BranchTree.from_tree(self.tree)
-
-    @cached_property
-    def _bifurcations(self) -> List[int]:
-        return [n.id for n in self.tree.get_bifurcations()]
-
-    @cached_property
-    def _tips(self) -> List[int]:
-        return [n.id for n in self.tree.get_tips()]
-
-    @cached_property
-    def _other(self) -> npt.NDArray[np.int32]:
-        other = self.tree.id()
-        other = np.setdiff1d(other, self._bifurcations)
-        other = np.setdiff1d(other, self._tips)
-        return other