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

swcgeom/analysis/sholl.py

@@ -1,6 +1,7 @@
 """Sholl analysis."""
 
 import math
+import warnings
 from typing import Literal, Tuple
 
 import numpy as np
@@ -21,52 +22,75 @@ class Sholl:
        cortices of the cat J. Anat., 87 (1953), pp. 387-406
     """
 
-    count: npt.NDArray[np.int32]
+    count: npt.NDArray[np.int64]
     step: float
+    steps: npt.NDArray[np.float32]
 
     def __init__(self, tree: Tree, step: float = 1) -> None:
         xyz = tree.get_segments().xyz() - tree.soma().xyz()  # shift
-        radius = np.linalg.norm(xyz, axis=2)
-        max_radius = math.ceil(np.max(radius[:, 1]))
-        count = [
-            np.count_nonzero(np.logical_and(radius[:, 0] <= i, radius[:, 1] > i))
-            for i in np.arange(0, max_radius, step)
-        ]
-
-        self.count = np.array(count, dtype=np.int32)
+        r = np.linalg.norm(xyz, axis=2)
+        steps = np.arange(step, int(np.ceil(r.max())), step)
+        intersections = [np.logical_and(r[:, 0] <= i, r[:, 1] > i) for i in steps]
+        count = np.count_nonzero(intersections, axis=1)
+
+        self.count = count
         self.step = step
+        self.steps = steps
 
     def __getitem__(self, idx: int) -> int:
         return self.count[idx] if 0 <= idx < len(self.count) else 0
 
-    def get_count(self) -> npt.NDArray[np.int32]:
+    def get(self) -> npt.NDArray[np.int64]:
         return self.count.copy()
 
-    def get_distribution(self) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.int32]]:
-        y = self.get_count()
-        x = self.step * np.arange(len(y))
-        return x, y
+    def get_count(self) -> npt.NDArray[np.int32]:
+        warnings.warn(
+            "`Sholl.get_count` has been renamed to `get` since v0.5.0, "
+            "and will be removed in next version",
+            DeprecationWarning,
+        )
+        return self.get().astype(np.int32)
 
     def avg(self) -> float:
-        return np.average(self.get_count()).item()
+        warnings.warn(
+            "`Sholl.avg` will be removed in next version",
+            DeprecationWarning,
+        )
+        return self.count.mean()
 
     def std(self) -> float:
-        return np.std(self.get_count()).item()
+        warnings.warn(
+            "`Sholl.std` will be removed in next version",
+            DeprecationWarning,
+        )
+        return self.count.std()
 
     def sum(self) -> int:
-        return np.sum(self.get_count()).item()
+        warnings.warn(
+            "`Sholl.sum` will be removed in next version",
+            DeprecationWarning,
+        )
+        return self.count.sum()
 
     def plot(
         self,
-        plot_type: Literal["bar", "linechart", "circles"] = "linechart",
+        plot_type: str | None = None,
+        kind: Literal["bar", "linechart", "circles"] = "linechart",
         fig: Figure | None = None,
         ax: Axes | None = None,
         **kwargs,
     ) -> Tuple[Figure, Axes]:
+        if plot_type is not None:
+            warnings.warn(
+                "`plot_type` has been renamed to `kind` since v0.5.0, "
+                "and will be removed in next version",
+                DeprecationWarning,
+            )
+            kind = plot_type  # type: ignore
+
+        x, y = self.steps, self.count
         fig, ax = get_fig_ax(fig, ax)
-        x, y = self.get_distribution()
-
-        match plot_type:
+        match kind:
             case "bar":
                 kwargs.setdefault("width", self.step)
                 ax.bar(x, y, **kwargs)
@@ -76,6 +100,6 @@ class Sholl:
                 kwargs.setdefault("y_min", 0)
                 draw_circles(fig, ax, x, y, **kwargs)
             case _:
-                raise ValueError(f"unsupported plot type: {plot_type}")
+                raise ValueError(f"unsupported plot kind: {kind}")
 
         return fig, ax

swcgeom/analysis/visualization.py

@@ -2,14 +2,22 @@
 
 import os
 import weakref
-from typing import Any, Dict, Literal, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 import numpy as np
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
+from matplotlib.legend import Legend
 
 from ..core import SWCLike, Tree
-from ..utils import Camera, Vector3D, draw_lines, draw_xyz_axes, get_fig_ax, palette
+from ..utils import (
+    Camera,
+    Vec3f,
+    draw_direction_indicator,
+    draw_lines,
+    get_fig_ax,
+    palette,
+)
 
 __all__ = ["draw"]
 
@@ -22,24 +30,31 @@ CameraPresets: Dict[CameraPreset, Camera] = {
     "zy": Camera((0, 0, 0), (-1, 0, 0), (0, 0, -1)),
     "xz": Camera((0, 0, 0), (0, -1, 0), (-1, 0, 0)),
 }
-CameraOptions = (
-    Vector3D | Tuple[Vector3D, Vector3D] | Tuple[Vector3D, Vector3D, Vector3D]
-)
+CameraOptions = Vec3f | Tuple[Vec3f, Vec3f] | Tuple[Vec3f, Vec3f, Vec3f]
+Positions = Literal["lt", "lb", "rt", "rb"] | Tuple[float, float]
+positions = {
+    "lt": (0.10, 0.90),
+    "lb": (0.10, 0.10),
+    "rt": (0.90, 0.90),
+    "rb": (0.90, 0.10),
+}
 
-ax_weak_dict = weakref.WeakKeyDictionary[Axes, Dict[str, Any]]({})
+ax_weak_memo = weakref.WeakKeyDictionary[Axes, Dict[str, Any]]({})
 
 
 def draw(
     swc: SWCLike | str,
     *,
-    fig: Figure | None = None,
-    ax: Axes | None = None,
+    fig: Optional[Figure] = None,
+    ax: Optional[Axes] = None,
     camera: CameraOptions | CameraPreset = "xy",
-    color: Dict[int, str] | str | None = None,
-    label: str | Literal[True] = True,  # TODO: support False
+    color: Optional[Dict[int, str] | str] = None,
+    label: str | bool = True,
+    direction_indicator: Positions | Literal[False] = "rb",
+    unit: Optional[str] = None,
     **kwargs,
 ) -> tuple[Figure, Axes]:
-    """Draw neuron tree.
+    r"""Draw neuron tree.
 
     Parameters
     ----------
@@ -59,6 +74,11 @@ def draw(
         parent node.If is string, the value will be use for any type.
     label : str | bool, default True
         Label of legend, disable if False.
+    direction_indicator : str or (float, float), default 'rb'
+        Draw a xyz direction indicator, can be place on 'lt', 'lb',
+        'rt', 'rb', or custom position.
+    unit : optional[str]
+        Add unit text, e.g.: r"$\mu m$".
     **kwargs : dict[str, Unknown]
         Forwarded to `~matplotlib.collections.LineCollection`.
 
@@ -68,70 +88,56 @@ def draw(
     ax : ~matplotlib.axes.Axes
     """
 
-    fig, ax = get_fig_ax(fig, ax)
-    ax_weak_dict.setdefault(ax, {})
-    ax_weak_dict[ax].setdefault("swc", [])
-
-    if isinstance(swc, str):
-        swc = Tree.from_swc(swc)
-    ax_weak_dict[ax]["swc"].append(swc)
+    swc = Tree.from_swc(swc) if isinstance(swc, str) else swc
 
-    my_camera = get_camera(camera)
-    my_color = get_color(ax, swc, color)
+    fig, ax = get_fig_ax(fig, ax)
+    my_camera = _get_camera(camera)
+    my_color = get_ax_color(ax, swc, color)
 
     xyz = swc.xyz()
     starts, ends = swc.id()[1:], swc.pid()[1:]
     lines = np.stack([xyz[starts], xyz[ends]], axis=1)
-
     collection = draw_lines(lines, ax=ax, camera=my_camera, color=my_color, **kwargs)
 
-    # legend
-    ax_weak_dict[ax].setdefault("handles", [])
-    ax_weak_dict[ax]["handles"].append(collection)
-    set_lable(ax, swc, label)
-
     ax.autoscale()
+    _set_ax_memo(ax, swc, label=label, handle=collection)  # legend
 
-    if len(ax_weak_dict[ax]["swc"]) == 1:
+    if len(get_ax_swc(ax)) == 1:
         ax.set_aspect(1)
         ax.spines[["top", "right"]].set_visible(False)
-        ax.text(0.05, 0.95, r"$\mu m$", transform=ax.transAxes)
-        draw_xyz_axes(ax=ax, camera=my_camera)
+        if direction_indicator is not False:
+            p = (
+                positions[direction_indicator]
+                if isinstance(direction_indicator, str)
+                else direction_indicator
+            )
+            draw_direction_indicator(ax=ax, camera=my_camera, position=p)
+        if unit is not None:
+            ax.text(0.05, 0.95, unit, transform=ax.transAxes)
     else:
-        # legend
-        handles = ax_weak_dict[ax].get("handles", [])
-        labels = ax_weak_dict[ax].get("labels", [])
-        ax.legend(handles, labels, loc="upper right")
+        set_ax_legend(ax, loc="upper right")  # enable legend
 
     return fig, ax
 
 
-def get_camera(camera: CameraOptions | CameraPreset) -> Camera:
-    if isinstance(camera, str):
-        return CameraPresets[camera]
-
-    if len(camera) == 1:
-        return Camera((0, 0, 0), camera, (0, 1, 0))
-
-    if len(camera) == 2:
-        return Camera((0, 0, 0), camera[0], camera[1])
-
-    return Camera(*camera)
+def get_ax_swc(ax: Axes) -> List[SWCLike]:
+    ax_weak_memo.setdefault(ax, {})
+    return ax_weak_memo[ax]["swc"]
 
 
-def get_color(
-    ax: Axes, swc: SWCLike, color: Dict[int, str] | str | None
-) -> str | list[str]:
+def get_ax_color(
+    ax: Axes, swc: SWCLike, color: Optional[Dict[int, str] | str] = None
+) -> str | List[str]:
     if color == "vaa3d":
         color = palette.vaa3d
-
-    if isinstance(color, str):
-        return color
+    elif isinstance(color, str):
+        return color  # user specified
 
     # choose default
-    ax_weak_dict[ax].setdefault("color", -1)
-    ax_weak_dict[ax]["color"] += 1
-    c = palette.default[ax_weak_dict[ax]["color"] % len(palette.default)]
+    ax_weak_memo.setdefault(ax, {})
+    ax_weak_memo[ax].setdefault("color", -1)
+    ax_weak_memo[ax]["color"] += 1
+    c = palette.default[ax_weak_memo[ax]["color"] % len(palette.default)]
 
     if isinstance(color, dict):
         types = swc.type()[:-1]  # colored by type of parent node
@@ -140,17 +146,48 @@ def get_color(
     return c
 
 
-def set_lable(ax: Axes, swc: SWCLike, label: str | bool):
-    ax_weak_dict[ax].setdefault("labels", [])
-    if label is False:
-        ax_weak_dict[ax]["labels"].append(False)
-        return
+def set_ax_legend(ax: Axes, *args, **kwargs) -> Legend | None:
+    labels = ax_weak_memo[ax].get("labels", [])
+    handles = ax_weak_memo[ax].get("handles", [])
+
+    # filter `label = False`
+    handles = [a for i, a in enumerate(handles) if labels[i] != False]
+    labels = [a for i, a in enumerate(labels) if labels[i] != False]
+
+    if len(labels) > 0:
+        return ax.legend(handles, labels, *args, **kwargs)
+    else:
+        return None
+
+
+def _set_ax_memo(
+    ax: Axes,
+    swc: SWCLike,
+    label: Optional[str | bool] = None,
+    handle: Optional[Any] = None,
+):
+    ax_weak_memo.setdefault(ax, {})
+    ax_weak_memo[ax].setdefault("swc", [])
+    ax_weak_memo[ax]["swc"].append(swc)
+
+    if label is not None:
+        label = os.path.basename(swc.source) if label is True else label
+        ax_weak_memo[ax].setdefault("labels", [])
+        ax_weak_memo[ax]["labels"].append(label)
+
+    if handle is not None:
+        ax_weak_memo[ax].setdefault("handles", [])
+        ax_weak_memo[ax]["handles"].append(handle)
+
 
-    if label is True:
-        try:
-            (_, tail) = os.path.split(swc.source)
-            label = tail
-        except:  # type: ignore
-            label = swc.source
+def _get_camera(camera: CameraOptions | CameraPreset) -> Camera:
+    if isinstance(camera, str):
+        return CameraPresets[camera]
+
+    if len(camera) == 1:
+        return Camera((0, 0, 0), camera, (0, 1, 0))
 
-    ax_weak_dict[ax]["labels"].append(label)
+    if len(camera) == 2:
+        return Camera((0, 0, 0), camera[0], camera[1])
+
+    return Camera(*camera)

swcgeom/core/branch_tree.py

@@ -3,6 +3,8 @@
 import itertools
 from typing import Dict, List
 
+import pandas as pd
+
 from .branch import Branch
 from .tree import Tree
 from .tree_utils import to_sub_tree
@@ -50,7 +52,7 @@ class BranchTree(Tree):
 
         return branch_tree
 
-    @staticmethod
-    def from_swc(swc_file: str, **kwargs) -> "BranchTree":
-        tree = super().from_swc(swc_file, **kwargs)
+    @classmethod
+    def from_data_frame(cls, df: pd.DataFrame, *args, **kwargs) -> "BranchTree":
+        tree = super().from_data_frame(df, *args, **kwargs)
         return BranchTree.from_tree(tree)

swcgeom/core/node.py

@@ -101,7 +101,7 @@ class Node(Generic[SWCTypeVar]):
         return self.attach.id()[self.attach.pid() == self.id]
 
     def is_bifurcation(self) -> bool:
-        return len(self.attach.pid() == self.id) > 1
+        return np.count_nonzero(self.attach.pid() == self.id) > 1
 
     def is_tip(self) -> bool:
         return self.id not in self.attach.pid()

swcgeom/core/population.py

@@ -1,26 +1,35 @@
 """Neuron population is a set of tree."""
 
 import os
-from typing import Any, Dict, Iterator, List, cast, overload
+import warnings
+from functools import reduce
+from itertools import chain
+from typing import Any, Dict, Iterator, List, Optional, cast, overload
+
+from typing_extensions import Self
 
 from .swc import eswc_cols
 from .tree import Tree
 
-__all__ = ["Population"]
+__all__ = ["Population", "Populations"]
 
 
 class Population:
     """Neuron population."""
 
-    source: str = ""
+    root: str
     swcs: List[str]
     trees: List[Tree | None]
-    read_kwargs: Dict[str, Any]
+    kwargs: Dict[str, Any]
 
-    def __init__(self, swcs: list[str], lazy_loading=True) -> None:
+    def __init__(
+        self, swcs: List[str], lazy_loading: bool = True, root: str = "", **kwargs
+    ) -> None:
         super().__init__()
+        self.root = root
         self.swcs = swcs
         self.trees = [None for _ in swcs]
+        self.kwargs = kwargs
         if not lazy_loading:
             self.load(slice(len(swcs)))
 
@@ -47,7 +56,7 @@ class Population:
         return (self[i] for i in range(self.__len__()))
 
     def __repr__(self) -> str:
-        return f"Neuron population in '{self.source}'"
+        return f"Neuron population in '{self.root}'"
 
     # fmt:off
     @overload
@@ -65,28 +74,121 @@ class Population:
 
         for i in idx:
             if self.trees[i] is None:
-                self.trees[i] = Tree.from_swc(self.swcs[i], **self.read_kwargs)
-
-    @staticmethod
-    def from_swc(swc_dir: str, suffix: str = ".swc", **kwargs) -> "Population":
-        swcs = Population.find_swcs(swc_dir, suffix)
-        population = Population(swcs)
-        population.source = swc_dir
-        population.read_kwargs = kwargs
-        return population
-
-    @staticmethod
-    def from_eswc(swc_dir: str, suffix: str = ".eswc", **kwargs) -> "Population":
-        kwargs.setdefault("extra_cols", [])
-        kwargs["extra_cols"].extend(k for k, t in eswc_cols)
-        return Population.from_swc(swc_dir, suffix, **kwargs)
+                self.trees[i] = Tree.from_swc(self.swcs[i], **self.kwargs)
+
+    @classmethod
+    def from_swc(cls, root: str, ext: str = ".swc", **kwargs) -> Self:
+        swcs = cls.find_swcs(root, ext)
+        if len(swcs) == 0:
+            warnings.warn(f"no trees in population from '{root}'")
+        return Population(swcs, root=root, **kwargs)
+
+    @classmethod
+    def from_eswc(
+        cls, root: str, ext: str = ".eswc", extra_cols: List[str] = [], **kwargs
+    ) -> Self:
+        extra_cols.extend(k for k, t in eswc_cols)
+        return cls.from_swc(root, ext, extra_cols=extra_cols, **kwargs)
 
     @staticmethod
-    def find_swcs(swc_dir: str, suffix: str = ".swc") -> list[str]:
+    def find_swcs(root: str, ext: str = ".swc", relpath: bool = False) -> List[str]:
         """Find all swc files."""
-        swcs = list[str]()
-        for root, _, files in os.walk(swc_dir):
-            files = [f for f in files if os.path.splitext(f)[-1] == suffix]
-            swcs.extend([os.path.join(root, f) for f in files])
+        swcs: List[str] = []
+        for r, _, files in os.walk(root):
+            rr = os.path.relpath(r, root) if relpath else r
+            fs = filter(lambda f: os.path.splitext(f)[-1] == ext, files)
+            swcs.extend(os.path.join(rr, f) for f in fs)
 
         return swcs
+
+
+class Populations:
+    len: int
+    populations: List[Population]
+    labels: List[str]
+
+    def __init__(
+        self, populations: List[Population], labels: Optional[List[str]] = None
+    ) -> None:
+        self.len = min(len(p) for p in populations)
+        self.populations = populations
+
+        labels = labels or ["" for i in populations]
+        assert len(labels) == len(
+            populations
+        ), f"got {len(populations)} populations, but {len(labels)} labels"
+        self.labels = labels
+
+    # fmt:off
+    @overload
+    def __getitem__(self, key: slice) -> List[List[Tree]]: ...
+    @overload
+    def __getitem__(self, key: int) -> List[Tree]: ...
+    # fmt:on
+    def __getitem__(self, key):
+        return [p[key] for p in self.populations]
+
+    def __len__(self) -> int:
+        return self.len
+
+    def __iter__(self) -> Iterator[List[Tree]]:
+        return (self[i] for i in range(self.len))
+
+    def __repr__(self) -> str:
+        return (
+            f"A cluster of {self.num_of_populations()} neuron populations, "
+            f"each containing at least {self.len} trees"
+        )
+
+    def num_of_populations(self) -> int:
+        return len(self.populations)
+
+    def to_population(self) -> Population:
+        swcs = list(chain.from_iterable(p.swcs for p in self.populations))
+        return Population(swcs)
+
+    @classmethod
+    def from_swc(
+        cls,
+        roots: List[str],
+        ext: str = ".swc",
+        intersect: bool = True,
+        check_same: bool = True,
+        labels: Optional[List[str]] = None,
+        **kwargs,
+    ) -> Self:
+        """Get population from dirs.
+
+        Parameters
+        ----------
+        roots : list of str
+        intersect : bool, default `False`
+            Take the intersection of these populations.
+        check_same : bool, default `True`
+            Check if the directories contains the same swc.
+        labels : List of str, optional
+            Label of populations.
+        **kwargs : Any
+            Forwarding to `Population`.
+        """
+
+        fs = [Population.find_swcs(d, ext=ext, relpath=True) for d in roots]
+        if intersect:
+            inter = list(reduce(lambda a, b: set(a).intersection(set(b)), fs))
+            if len(inter) == 0:
+                warnings.warn("no intersection among populations")
+
+            fs = [inter for _ in roots]
+        elif check_same:
+            assert reduce(lambda a, b: a == b, fs), "not the same among populations"
+
+        populations = [
+            Population([os.path.join(roots[i], p) for p in fs[i]], root=d, **kwargs)
+            for i, d in enumerate(roots)
+        ]
+        return cls(populations, labels=labels)
+
+    @classmethod
+    def from_eswc(cls, roots: List[str], extra_cols: List[str] = [], **kwargs) -> Self:
+        extra_cols.extend(k for k, t in eswc_cols)
+        return cls.from_swc(roots, extra_cols=extra_cols, **kwargs)

swcgeom/core/swc.py

@@ -8,9 +8,13 @@ import numpy.typing as npt
 import pandas as pd
 import scipy.sparse as sp
 
-from .swc_utils import check_single_root, link_roots_to_nearest, mark_roots_as_somas
-from .swc_utils import reset_index as _reset_index
-from .swc_utils import sort_nodes
+from .swc_utils import (
+    check_single_root,
+    link_roots_to_nearest_,
+    mark_roots_as_somas_,
+    reset_index_,
+    sort_nodes_,
+)
 
 __all__ = ["swc_cols", "eswc_cols", "read_swc", "SWCLike", "SWCTypeVar"]
 
@@ -176,7 +180,7 @@ def read_swc(
     swc_file: str,
     extra_cols: List[str] | None = None,
     fix_roots: Literal["somas", "nearest", False] = False,
-    sort: bool = False,
+    sort_nodes: bool = False,
     reset_index: bool = True,
 ) -> pd.DataFrame:
     """Read swc file.
@@ -189,7 +193,7 @@ def read_swc(
         Read more cols in swc file.
     fix_roots : `somas`|`nearest`|False, default `False`
         Fix multiple roots.
-    sort : bool, default `False`
+    sort_nodes : bool, default `False`
         Sort the indices of neuron tree, the index for parent are
         always less than children.
     reset_index : bool, default `True`
@@ -214,16 +218,16 @@ def read_swc(
     if fix_roots is not False and np.count_nonzero(df["pid"] == -1) > 1:
         match fix_roots:
             case "somas":
-                mark_roots_as_somas(df)
+                mark_roots_as_somas_(df)
             case "nearest":
-                link_roots_to_nearest(df)
+                link_roots_to_nearest_(df)
             case _:
                 raise ValueError(f"unknown fix type: {fix_roots}")
 
-    if sort:
-        sort_nodes(df)
+    if sort_nodes:
+        sort_nodes_(df)
     elif reset_index:
-        _reset_index(df)
+        reset_index_(df)
 
     # check swc
     if not check_single_root(df):