swcgeom 0.19.4__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

swcgeom/analysis/sholl.py

@@ -0,0 +1,201 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Sholl analysis."""
+
+import warnings
+from typing import Literal
+
+import numpy as np
+import numpy.typing as npt
+import seaborn as sns
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from typing_extensions import deprecated
+
+from swcgeom.analysis.visualization import draw
+from swcgeom.core import Tree
+from swcgeom.transforms import TranslateOrigin
+from swcgeom.utils import draw_circles, get_fig_ax
+
+__all__ = ["Sholl"]
+
+XLABEL = "Radial Distance"
+YLABLE = "Count of Intersections"
+
+
+class Sholl:
+    """Sholl analysis.
+
+    Implementation of original Sholl analysis as described in [1]_. The Sholl analysis
+    is a method to quantify the spatial distribution of neuronal processes. It is
+    based on the number of intersections of concentric circles with the neuronal
+    processes.
+
+    References:
+    .. [1] Dendritic organization in the neurons of the visual and
+       motor cortices of the cat J. Anat., 87 (1953), pp. 387-406
+    """
+
+    tree: Tree
+    rs: npt.NDArray[np.float32]
+    rmax: float
+
+    # compat
+    step: float | None = None
+
+    def __init__(self, tree: Tree | str, step: float | None = None) -> None:
+        tree = Tree.from_swc(tree) if isinstance(tree, str) else tree
+        try:
+            self.tree = TranslateOrigin.transform(tree)  # shift
+            self.rs = np.linalg.norm(self.tree.get_segments().xyz(), axis=2)
+            self.rmax = self.rs.max()
+        except Exception as e:
+            raise ValueError(f"invalid tree: {tree.source or ''}") from e
+
+        if step is not None:
+            warnings.warn(
+                "`Sholl(x, step=...)` has been replaced by "
+                "`Sholl(x).get(steps=...)` since v0.6.0 because it has "
+                "been change to dynamic calculate, and will be removed "
+                "in next version",
+                DeprecationWarning,
+            )
+            self.step = step
+
+    def get(self, steps: int | npt.ArrayLike = 20) -> npt.NDArray[np.int64]:
+        intersections = [
+            np.logical_or(
+                np.logical_and(self.rs[:, 0] <= r, self.rs[:, 1] > r),
+                np.logical_and(self.rs[:, 1] <= r, self.rs[:, 0] > r),
+            )
+            for r in self._get_rs(steps=steps)
+        ]
+        return np.count_nonzero(intersections, axis=1)
+
+    def intersect(self, r: float) -> int:
+        return np.count_nonzero(
+            np.logical_or(
+                np.logical_and(self.rs[:, 0] <= r, self.rs[:, 1] > r),
+                np.logical_and(self.rs[:, 1] <= r, self.rs[:, 0] > r),
+            )
+        )
+
+    def plot(  # pylint: disable=too-many-arguments
+        self,
+        steps: list[float] | int = 20,
+        plot_type: str | None = None,
+        kind: Literal["bar", "linechart", "circles"] = "circles",
+        fig: Figure | None = None,
+        ax: Axes | None = None,
+        **kwargs,
+    ) -> tuple[Figure, Axes]:
+        """Plot Sholl analysis.
+
+        Args:
+            steps: Steps of raius of circle.
+                If steps is int, then it will be evenly divided into n radii.
+            kind: Kind of plot.
+            fig: The figure to plot on.
+            ax: The axes to plot on.
+            **kwargs: Forwarding to plot method.
+        """
+        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
+
+        xs = self._get_rs(steps=steps)
+        ys = self.get(steps=xs)
+        fig, ax = get_fig_ax(fig, ax)
+        match kind:
+            case "bar":
+                sns.barplot(x=xs, y=ys, ax=ax, **kwargs)
+                ax.set_ylabel(YLABLE)
+
+            case "linechart":
+                sns.lineplot(x=xs, y=ys, ax=ax, **kwargs)
+                ax.set_ylabel(YLABLE)
+
+            case "circles":
+                kwargs.setdefault("y_min", 0)
+                drawtree = kwargs.pop("drawtree", True)
+                colorbar = kwargs.pop("colorbar", True)
+                cmap = kwargs.pop("cmap", "Blues")
+                patches = draw_circles(ax, xs, ys, cmap=cmap, **kwargs)
+
+                if drawtree is True:
+                    draw(self.tree, ax=ax, direction_indicator=False)
+                elif isinstance(drawtree, str):
+                    draw(self.tree, ax=ax, color=drawtree, direction_indicator=False)
+
+                if colorbar is True:
+                    fig.colorbar(patches, ax=ax, label=YLABLE)
+                elif isinstance(colorbar, (Axes, np.ndarray, list)):
+                    fig.colorbar(patches, ax=colorbar, label=YLABLE)
+
+            case _:
+                raise ValueError(f"unsupported kind: {kind}")
+
+        ax.set_xlabel(XLABEL)
+        return fig, ax
+
+    @staticmethod
+    def get_rs(rmax: float, steps: int | npt.ArrayLike) -> npt.NDArray[np.float32]:
+        """Function to calculate the list of radius used by the sholl."""
+        if isinstance(steps, int):
+            s = rmax / (steps + 1)
+            return np.arange(s, rmax, s)
+
+        return np.array(steps)
+
+    def _get_rs(self, steps: int | npt.ArrayLike) -> npt.NDArray[np.float32]:
+        if self.step is not None:  # compat
+            return np.arange(self.step, int(np.ceil(self.rmax)), self.step)
+
+        return self.get_rs(self.rmax, steps)
+
+    @deprecated("Use `Sholl.get(x)` instead")
+    def get_count(self) -> npt.NDArray[np.int32]:
+        """Get the count of intersection.
+
+        .. deprecated:: 0.5.0
+            Use :meth:`Sholl(x).get()` instead.
+        """
+
+        return self.get().astype(np.int32)
+
+    @deprecated("Use `Shool(x).get().mean()` instead")
+    def avg(self) -> float:
+        """Get the average of the count of intersection.
+
+        .. deprecated:: 0.6.0
+            Use :meth:`Shool(x).get().mean()` instead.
+        """
+
+        return self.get().mean()
+
+    @deprecated("Use `Shool(x).get().std()` instead")
+    def std(self) -> float:
+        """Get the std of the count of intersection.
+
+        .. deprecated:: 0.6.0
+            Use :meth:`Shool(x).get().std()` instead.
+        """
+
+        return self.get().std()
+
+    @deprecated("Use `Shool(x).get().sum()` instead")
+    def sum(self) -> int:
+        """Get the sum of the count of intersection.
+
+        .. deprecated:: 0.6.0
+            Use :meth:`Shool(x).get().sum()` instead.
+        """
+
+        return self.get().sum()

swcgeom/analysis/trunk.py

@@ -0,0 +1,183 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Plot trunk and florets."""
+
+# pylint: disable=invalid-name
+
+from collections.abc import Iterable
+from itertools import chain
+from typing import Any, Literal, cast
+
+import numpy as np
+import numpy.typing as npt
+from matplotlib import cm
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from matplotlib.patches import Circle, Ellipse, Patch, Rectangle
+
+from swcgeom.analysis.visualization import draw
+from swcgeom.core import Tree, get_subtree, to_subtree
+from swcgeom.utils import get_fig_ax, mvee
+
+__all__ = ["draw_trunk"]
+
+Bounds = Literal["aabb", "ellipse"]
+Projection = Literal["2d"]
+
+
+def draw_trunk(
+    t: Tree,
+    florets: Iterable[int | Iterable[int]],
+    *,
+    fig: Figure | None = None,
+    ax: Axes | None = None,
+    bound: Bounds | tuple[Bounds, dict[str, Any]] | None = "ellipse",
+    point: bool | dict[str, Any] = True,
+    projection: Projection = "2d",
+    cmap: Any = "viridis",
+    **kwargs,
+) -> tuple[Figure, Axes]:
+    """Draw trunk tree.
+
+    Args:
+        t: Tree
+        florets: The florets that needs to be removed.
+            Each floret can be a subtree or multiple subtrees (e.g., dendrites are a
+            bunch of subtrees), each number is the id of a tree node.
+        fig: Figure to plot on.
+        ax: Axes to plot on.
+        bound: Kind of bound, support 'aabb', 'ellipse'.
+            If bound is None, no bound will be drawn. If bound is a tuple, the second
+            item will used as kwargs and forward to draw function.
+        point: Draw point at the start of a subtree.
+            If point is False, no point will be drawn. If point is a dict, this will
+            used a kwargs and forward to draw function.
+        cmap: Colormap.
+            Any value supported by ~matplotlib.cm.Colormap. We will use the ratio of
+            the length of the subtree to the total length of the tree to determine the
+            color.
+        **kwargs: Forward to ~swcgeom.analysis.draw.
+    """
+    # pylint: disable=too-many-locals
+    trunk, tss = split_florets(t, florets)
+    lens = get_length_ratio(t, tss)
+
+    cmap = cm.get_cmap(cmap)
+    c = cmap(lens)
+
+    fig, ax = get_fig_ax(fig, ax)
+    if bound is not None:
+        for ts, cc in zip(tss, c):
+            draw_bound(ts, ax, bound, projection, color=cc)
+
+    if point is not False:
+        point_kwargs = point if isinstance(point, dict) else {}
+        for ts, cc in zip(tss, c):
+            draw_point(ts, ax, projection, color=cc, **point_kwargs)
+
+    draw(trunk, ax=ax, color=cmap(1), **kwargs)
+    return fig, ax
+
+
+def split_florets(
+    t: Tree, florets: Iterable[int | Iterable[int]]
+) -> tuple[Tree, list[list[Tree]]]:
+    florets = [[i] if isinstance(i, (int, np.integer)) else i for i in florets]
+    subtrees = [[get_subtree(t, ff) for ff in f] for f in florets]
+    trunk = to_subtree(t, chain(*florets))
+    return trunk, subtrees
+
+
+def get_length_ratio(t: Tree, tss: list[list[Tree]]) -> Any:
+    lens = np.array([sum(t.length() for t in ts) for ts in tss])
+    return lens / t.length()
+
+
+# Bounds
+
+
+def draw_bound(
+    ts: Iterable[Tree],
+    ax: Axes,
+    bound: Bounds | tuple[Bounds, dict[str, Any]],
+    projection: Projection,
+    **kwargs,
+) -> None:
+    kind, bound_kwargs = (bound, {}) if isinstance(bound, str) else bound
+    if projection == "2d":
+        patch = create_bound_2d(ts, kind, **kwargs, **bound_kwargs)
+    else:
+        raise ValueError(f"unsupported projection {projection}")
+
+    ax.add_patch(patch)
+
+
+def create_bound_2d(ts: Iterable[Tree], bound: Bounds, **kwargs) -> Patch:
+    xyz = np.concatenate([t.xyz() for t in ts])
+    xy = xyz[:, :2]  # TODO: camera
+
+    if bound == "aabb":
+        return create_aabb_2d(xy, **kwargs)
+    if bound == "ellipse":
+        return create_ellipse_2d(xy, **kwargs)
+    raise ValueError(f"unsupported bound `{bound}` in 2d projection")
+
+
+def create_aabb_2d(xy: npt.NDArray, fill: bool = False, **kwargs) -> Rectangle:
+    xmin, ymin = xy[:, 0].min(), xy[:, 1].min()
+    xmax, ymax = xy[:, 0].max(), xy[:, 1].max()
+    width, height = xmax - xmin, ymax - ymin
+    rect = Rectangle(
+        xy=(xmin, ymin), width=width, height=height, angle=0, fill=fill, **kwargs
+    )
+    return rect
+
+
+def create_ellipse_2d(xy: npt.NDArray, fill: bool = False, **kwargs) -> Ellipse:
+    ellipse = mvee(xy)
+    patch = Ellipse(
+        xy=ellipse.centroid,  # type:ignore
+        width=ellipse.a,
+        height=ellipse.b,
+        angle=ellipse.alpha,
+        fill=fill,
+        **kwargs,
+    )
+    return patch
+
+
+# point
+
+
+def draw_point(ts: Iterable[Tree], ax: Axes, projection: Projection, **kwargs) -> None:
+    if projection == "2d":
+        patch = create_point_2d(ts, **kwargs)
+    else:
+        raise ValueError(f"unsupported projection {projection}")
+
+    ax.add_patch(patch)
+
+
+def create_point_2d(
+    ts: Iterable[Tree], radius: float | None = None, **kwargs
+) -> Circle:
+    if radius is None:
+        xyz = np.concatenate([t.xyz() for t in ts])  # TODO: cache
+        radius = 0.05 * min(
+            xyz[:, 0].max() - xyz[:, 0].min(),
+            xyz[:, 1].max() - xyz[:, 1].min(),
+        )
+        radius = cast(float, radius)
+
+    center = np.mean([t.xyz()[0, :2] for t in ts], axis=0)
+    return Circle(center, radius, **kwargs)
+
+
+# Helpers
+
+
+def get_dendrites(tree: Tree) -> Iterable[int]:
+    return (t.node(0).id for t in tree.get_dendrites())

swcgeom/analysis/visualization.py

@@ -0,0 +1,191 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Painter utils."""
+
+import os
+import weakref
+from typing import Any, Literal
+
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from matplotlib.legend import Legend
+
+from swcgeom.core import SWCLike, Tree
+from swcgeom.utils import (
+    CameraOptions,
+    SimpleCamera,
+    draw_direction_indicator,
+    draw_lines,
+    get_fig_ax,
+    palette,
+)
+
+__all__ = ["draw"]
+
+Positions = Literal["lt", "lb", "rt", "rb"] | tuple[float, float]
+locations: dict[Literal["lt", "lb", "rt", "rb"], tuple[float, float]] = {
+    "lt": (0.10, 0.90),
+    "lb": (0.10, 0.10),
+    "rt": (0.90, 0.90),
+    "rb": (0.90, 0.10),
+}
+
+ax_weak_memo = weakref.WeakKeyDictionary[Axes, dict[str, Any]]({})
+
+
+def draw(
+    swc: SWCLike | str,
+    *,
+    fig: Figure | None = None,
+    ax: Axes | None = None,
+    show: bool | None = None,
+    camera: CameraOptions = "xy",
+    color: dict[int, str] | str | None = None,
+    label: str | bool = True,
+    direction_indicator: Positions | Literal[False] = "rb",
+    unit: str | None = None,
+    **kwargs,
+) -> tuple[Figure, Axes]:
+    r"""Draw neuron tree.
+
+    Args:
+        swc: The swc tree to draw.
+            If it is str, then it is treated as the path of swc file.
+        fig: The figure to plot on.
+        ax: The axes to plot on.
+        show: Weather to call `plt.show()`.
+            If not specified, it will depend on if ax is passed in, it will not be
+            called, otherwise it will be called by default.
+        camera: Camera options (position, look-at, up).
+            One, two, or three vectors are supported, if only one vector, then threat
+            it as look-at, so camera is ((0, 0, 0), look-at, (0, 1, 0)); if two vector,
+            then then threat it as (look-at, up), so camera is ((0, 0, 0), look-at, up).
+            An easy way is to use the presets "xy", "yz" and "zx".
+        color: Color map.
+            If is dict, segments will be colored by the type of parent node.If is
+            string, the value will be use for any type.
+        label: Label of legend, disable if False.
+        direction_indicator: Draw a xyz direction indicator.
+            Can be place on 'lt', 'lb', 'rt', 'rb', or custom position.
+        unit: str, optional
+            Add unit text, e.g.: r"$\mu m$".
+        **kwargs: dict[str, Unknown]
+            Forwarded to `~matplotlib.collections.LineCollection`.
+
+    Returns:
+        fig: The figure to plot on.
+        ax: The axes to plot on.
+    """
+    # pylint: disable=too-many-locals
+    swc = Tree.from_swc(swc) if isinstance(swc, str) else swc
+
+    show = (show is True) or (show is None and ax is None)
+    fig, ax = get_fig_ax(fig, ax)
+
+    my_camera = SimpleCamera.from_options(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(ax, lines, camera=my_camera, color=my_color, **kwargs)
+
+    ax.autoscale()
+    _set_ax_memo(ax, swc, label=label, handle=collection)
+
+    if len(get_ax_swc(ax)) == 1:
+        ax.set_aspect(1)
+        ax.spines[["top", "right"]].set_visible(False)
+        if direction_indicator is not False:
+            loc = (
+                locations[direction_indicator]
+                if isinstance(direction_indicator, str)
+                else direction_indicator
+            )
+            draw_direction_indicator(ax, camera=my_camera, loc=loc)
+        if unit is not None:
+            ax.text(0.05, 0.95, unit, transform=ax.transAxes)
+    else:
+        set_ax_legend(ax, loc="upper right")  # enable legend
+
+    if show:
+        fig.show(warn=False)
+
+    return fig, ax
+
+
+def get_ax_swc(ax: Axes) -> list[SWCLike]:
+    ax_weak_memo.setdefault(ax, {})
+    return ax_weak_memo[ax]["swc"]
+
+
+def get_ax_color(
+    ax: Axes,
+    swc: SWCLike,
+    color: dict[int, str] | str | None = None,  # TODO: improve typing
+) -> str | list[str]:
+    if color == "vaa3d":
+        color = palette.vaa3d
+    elif isinstance(color, str):
+        return color  # user specified
+
+    # choose 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
+        return list(map(lambda type: color.get(type, c), types))
+
+    return c
+
+
+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] is not False]
+    labels = [a for a in labels if a is not False]
+
+    if len(labels) == 0:
+        return None
+
+    return ax.legend(handles, labels, *args, **kwargs)
+
+
+def _set_ax_memo(
+    ax: Axes, swc: SWCLike, label: str | bool | None = None, handle: Any | None = 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)

swcgeom/analysis/visualization3d.py

@@ -0,0 +1,81 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Painter utils.
+
+NOTE: This is a experimental function, it may be changed in the future.
+"""
+
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from mpl_toolkits.mplot3d import Axes3D
+
+from swcgeom.analysis.visualization import (
+    _set_ax_memo,
+    get_ax_color,
+    get_ax_swc,
+    set_ax_legend,
+)
+from swcgeom.core import SWCLike, Tree
+from swcgeom.utils.plotter_3d import draw_lines_3d
+
+__all__ = ["draw3d"]
+
+
+# TODO: support Camera
+def draw3d(
+    swc: SWCLike | str,
+    *,
+    ax: Axes,
+    show: bool | None = None,
+    color: dict[int, str] | str | None = None,  # TODO: improve typing
+    label: str | bool = True,
+    **kwargs,
+) -> tuple[Figure, Axes]:
+    r"""Draw neuron tree.
+
+    Args:
+        swc: The swc tree to draw.
+            If it is str, then it is treated as the path of swc file.
+        fig: The figure to plot on.
+        ax: The axes to plot on.
+        show: Weather to call `plt.show()`.
+            If not specified, it will depend on if ax is passed in, it will not be
+            called, otherwise it will be called by default.
+        color: Color map.
+            If is dict, segments will be colored by the type of parent node.If is
+            string, the value will be use for any type.
+        label: Label of legend, disable if False.
+        **kwargs: Forwarded to `~mpl_toolkits.mplot3d.art3d.Line3DCollection`.
+    """
+    assert isinstance(ax, Axes3D), "only support 3D axes."
+
+    swc = Tree.from_swc(swc) if isinstance(swc, str) else swc
+
+    show = (show is True) or (show is None and ax is None)
+    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_3d(ax, lines, color=my_color, **kwargs)
+
+    min_vals = lines.reshape(-1, 3).min(axis=0)
+    max_vals = lines.reshape(-1, 3).max(axis=0)
+    ax.set_xlim(min_vals[0], max_vals[0])
+    ax.set_ylim(min_vals[1], max_vals[1])
+    ax.set_zlim(min_vals[2], max_vals[2])
+
+    _set_ax_memo(ax, swc, label=label, handle=collection)
+
+    if len(get_ax_swc(ax)) == 1:
+        # ax.set_aspect(1)
+        ax.spines[["top", "right"]].set_visible(False)
+    else:
+        set_ax_legend(ax, loc="upper right")  # enable legend
+
+    fig = ax.figure
+    return fig, ax  # type: ignore