swcgeom 0.17.0__py3-none-any.whl → 0.17.1__py3-none-any.whl

swcgeom/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.17.0'
-__version_tuple__ = version_tuple = (0, 17, 0)
+__version__ = version = '0.17.1'
+__version_tuple__ = version_tuple = (0, 17, 1)

swcgeom/analysis/feature_extractor.py

@@ -7,10 +7,11 @@ 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, Callable, Dict, List, Literal, Tuple, overload
+from typing import Any, Literal, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -54,7 +55,7 @@ Feature = Literal[
 ]
 
 NDArrayf32 = npt.NDArray[np.float32]
-FeatAndKwargs = Feature | Tuple[Feature, Dict[str, Any]]
+FeatAndKwargs = Feature | tuple[Feature, dict[str, Any]]
 
 Feature1D = set(["length", "volume", "node_count", "bifurcation_count", "tip_count"])
 
@@ -121,9 +122,9 @@ class FeatureExtractor(ABC):
     @overload
     def get(self, feature: Feature, **kwargs) -> NDArrayf32: ...
     @overload
-    def get(self, feature: List[FeatAndKwargs]) -> List[NDArrayf32]: ...
+    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: dict[Feature, dict[str, Any]]) -> dict[str, NDArrayf32]: ...
     # fmt:on
     def get(self, feature, **kwargs):
         """Get feature.
@@ -168,7 +169,7 @@ class FeatureExtractor(ABC):
 
     # Custom Plots
 
-    def plot_node_branch_order(self, feature_kwargs: Dict[str, Any], **kwargs) -> Axes:
+    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)
@@ -234,7 +235,7 @@ class TreeFeatureExtractor(FeatureExtractor):
 
     def plot_sholl(
         self,
-        feature_kwargs: Dict[str, Any],  # pylint: disable=unused-argument
+        feature_kwargs: dict[str, Any],  # pylint: disable=unused-argument
         **kwargs,
     ) -> Axes:
         _, ax = self._features.sholl.plot(**kwargs)
@@ -264,7 +265,7 @@ class PopulationFeatureExtractor(FeatureExtractor):
     """Extract features from population."""
 
     _population: Population
-    _features: List[Features]
+    _features: list[Features]
 
     def __init__(self, population: Population) -> None:
         super().__init__()
@@ -279,7 +280,7 @@ class PopulationFeatureExtractor(FeatureExtractor):
 
     # Custom Plots
 
-    def plot_sholl(self, feature_kwargs: Dict[str, Any], **kwargs) -> Axes:
+    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")
@@ -295,7 +296,7 @@ class PopulationFeatureExtractor(FeatureExtractor):
 
     def _get_sholl_impl(
         self, steps: int = 20, **kwargs
-    ) -> Tuple[NDArrayf32, NDArrayf32]:
+    ) -> 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)
@@ -333,7 +334,7 @@ class PopulationsFeatureExtractor(FeatureExtractor):
     """Extract feature from population."""
 
     _populations: Populations
-    _features: List[List[Features]]
+    _features: list[list[Features]]
 
     def __init__(self, populations: Populations) -> None:
         super().__init__()
@@ -348,7 +349,7 @@ class PopulationsFeatureExtractor(FeatureExtractor):
 
     # Custom Plots
 
-    def plot_sholl(self, feature_kwargs: Dict[str, Any], **kwargs) -> Axes:
+    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")
@@ -369,7 +370,7 @@ class PopulationsFeatureExtractor(FeatureExtractor):
 
     def _get_sholl_impl(
         self, steps: int = 20, **kwargs
-    ) -> Tuple[NDArrayf32, NDArrayf32]:
+    ) -> 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)

swcgeom/analysis/features.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from functools import cached_property
-from typing import List, TypeVar
+from typing import TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -163,7 +163,7 @@ class PathFeatures:
         return np.array([path.tortuosity() for path in self._paths], dtype=np.float32)
 
     @cached_property
-    def _paths(self) -> List[Tree.Path]:
+    def _paths(self) -> list[Tree.Path]:
         return self.tree.get_paths()
 
 
@@ -201,7 +201,7 @@ class BranchFeatures:
         return self.calc_angle(self._branches, eps=eps)
 
     @staticmethod
-    def calc_angle(branches: List[T], eps: float = 1e-7) -> npt.NDArray[np.float32]:
+    def calc_angle(branches: list[T], eps: float = 1e-7) -> npt.NDArray[np.float32]:
         """Calc agnle between branches.
 
         Returns
@@ -219,5 +219,5 @@ class BranchFeatures:
         return angle
 
     @cached_property
-    def _branches(self) -> List[Tree.Branch]:
+    def _branches(self) -> list[Tree.Branch]:
         return self.tree.get_branches()

swcgeom/analysis/lmeasure.py

@@ -1,7 +1,7 @@
 """L-Measure analysis."""
 
 import math
-from typing import Literal, Tuple
+from typing import Literal
 
 import numpy as np
 import numpy.typing as npt
@@ -274,7 +274,7 @@ class LMeasure:
         rall_power, _, _, _ = self._rall_power(bif)
         return rall_power
 
-    def _rall_power_d(self, bif: Tree.Node) -> Tuple[float, float, float]:
+    def _rall_power_d(self, bif: Tree.Node) -> tuple[float, float, float]:
         children = bif.children()
         assert len(children) == 2, "Rall Power is only defined for bifurcations"
         parent = bif.parent()
@@ -284,7 +284,7 @@ class LMeasure:
         da, db = 2 * children[0].r, 2 * children[1].r
         return dp, da, db
 
-    def _rall_power(self, bif: Tree.Node) -> Tuple[float, float, float, float]:
+    def _rall_power(self, bif: Tree.Node) -> tuple[float, float, float, float]:
         dp, da, db = self._rall_power_d(bif)
         start, stop, step = 0, 5, 5 / 1000
         xs = np.arange(start, stop, step)
@@ -501,7 +501,7 @@ class LMeasure:
 
     def _bif_vector_local(
         self, bif: Tree.Node
-    ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
+    ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
         children = bif.children()
         assert len(children) == 2, "Only defined for bifurcations"
 
@@ -511,7 +511,7 @@ class LMeasure:
 
     def _bif_vector_remote(
         self, bif: Tree.Node
-    ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
+    ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]:
         children = bif.children()
         assert len(children) == 2, "Only defined for bifurcations"
 

swcgeom/analysis/sholl.py

@@ -1,7 +1,7 @@
 """Sholl analysis."""
 
 import warnings
-from typing import List, Literal, Optional, Tuple
+from typing import Literal, Optional
 
 import numpy as np
 import numpy.typing as npt
@@ -84,18 +84,18 @@ class Sholl:
 
     def plot(  # pylint: disable=too-many-arguments
         self,
-        steps: List[float] | int = 20,
+        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]:
+    ) -> tuple[Figure, Axes]:
         """Plot Sholl analysis.
 
         Parameters
         ----------
-        steps : int or List[float], default to 20
+        steps : int or list[float], default to 20
             Steps of raius of circle. If steps is int, then it will be
             evenly divided into n radii.
         kind : "bar" | "linechart" | "circles", default `circles`

swcgeom/analysis/trunk.py

@@ -2,8 +2,9 @@
 
 # pylint: disable=invalid-name
 
+from collections.abc import Iterable
 from itertools import chain
-from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, cast
+from typing import Any, Literal, Optional, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -28,28 +29,28 @@ def draw_trunk(
     *,
     fig: Optional[Figure] = None,
     ax: Optional[Axes] = None,
-    bound: Bounds | Tuple[Bounds, Dict[str, Any]] | None = "ellipse",
-    point: bool | Dict[str, Any] = True,
+    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]:
+) -> tuple[Figure, Axes]:
     """Draw trunk tree.
 
     Parameters
     ----------
     t : Tree
-    florets : List of (int | List of int)
+    florets : List of (int | list of int)
         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 : ~matplotlib.figure.Figure, optional
     ax : ~matplotlib.axes.Axes, optional
-    bound : Bounds | (Bounds, Dict[str, Any]) | None, default 'ellipse'
+    bound : Bounds | (Bounds, dict[str, Any]) | None, default 'ellipse'
         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 : bool | Dict[str, Any], default True
+    point : bool | dict[str, Any], default True
         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.
@@ -57,7 +58,7 @@ def draw_trunk(
         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 : Dict[str, Any]
+    **kwargs : dict[str, Any]
         Forward to ~swcgeom.analysis.draw.
     """
     # pylint: disable=too-many-locals
@@ -83,14 +84,14 @@ def draw_trunk(
 
 def split_florets(
     t: Tree, florets: Iterable[int | Iterable[int]]
-) -> Tuple[Tree, List[List[Tree]]]:
+) -> 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:
+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()
 
@@ -101,7 +102,7 @@ def get_length_ratio(t: Tree, tss: List[List[Tree]]) -> Any:
 def draw_bound(
     ts: Iterable[Tree],
     ax: Axes,
-    bound: Bounds | Tuple[Bounds, Dict[str, Any]],
+    bound: Bounds | tuple[Bounds, dict[str, Any]],
     projection: Projection,
     **kwargs,
 ) -> None:

swcgeom/analysis/visualization.py

@@ -2,7 +2,7 @@
 
 import os
 import weakref
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Literal, Optional
 
 import numpy as np
 from matplotlib.axes import Axes
@@ -21,15 +21,15 @@ from swcgeom.utils import (
 
 __all__ = ["draw"]
 
-Positions = Literal["lt", "lb", "rt", "rb"] | Tuple[float, float]
-locations: Dict[Literal["lt", "lb", "rt", "rb"], Tuple[float, float]] = {
+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]]({})
+ax_weak_memo = weakref.WeakKeyDictionary[Axes, dict[str, Any]]({})
 
 
 def draw(
@@ -39,7 +39,7 @@ def draw(
     ax: Optional[Axes] = None,
     show: bool | None = None,
     camera: CameraOptions = "xy",
-    color: Optional[Dict[int, str] | str] = None,
+    color: Optional[dict[int, str] | str] = None,
     label: str | bool = True,
     direction_indicator: Positions | Literal[False] = "rb",
     unit: Optional[str] = None,
@@ -64,7 +64,7 @@ def draw(
         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 : Dict[int, str] | "vaa3d" | str, optional
+    color : dict[int, str] | "vaa3d" | str, optional
         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 : str | bool, default True
@@ -120,14 +120,14 @@ def draw(
     return fig, ax
 
 
-def get_ax_swc(ax: Axes) -> List[SWCLike]:
+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: Optional[Dict[int, str] | str] = None
-) -> str | List[str]:
+    ax: Axes, swc: SWCLike, color: Optional[dict[int, str] | str] = None
+) -> str | list[str]:
     if color == "vaa3d":
         color = palette.vaa3d
     elif isinstance(color, str):

swcgeom/analysis/visualization3d.py

@@ -0,0 +1,85 @@
+"""Painter utils.
+
+Notes
+-----
+This is a experimental function, it may be changed in the future.
+"""
+
+from typing import Optional
+
+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: Optional[dict[int, str] | str] = None,
+    label: str | bool = True,
+    **kwargs,
+) -> tuple[Figure, Axes]:
+    r"""Draw neuron tree.
+
+    Parameters
+    ----------
+    swc : SWCLike | str
+        If it is str, then it is treated as the path of swc file.
+    fig : ~matplotlib.axes.Figure, optional
+    ax : ~matplotlib.axes.Axes, optional
+    show : bool | None, default `None`
+        Wheather 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 : dict[int, str] | "vaa3d" | str, optional
+        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 : str | bool, default True
+        Label of legend, disable if False.
+    **kwargs : dict[str, Unknown]
+        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)  # type: ignore
+
+    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

swcgeom/analysis/volume.py

@@ -1,6 +1,6 @@
 """Analysis of volume of a SWC tree."""
 
-from typing import Dict, List, Literal
+from typing import Literal
 
 import numpy as np
 from sdflit import ColoredMaterial, ObjectsScene, SDFObject, UniformSampler
@@ -11,7 +11,7 @@ from swcgeom.utils import VolFrustumCone, VolSphere
 __all__ = ["get_volume"]
 
 ACCURACY_LEVEL = Literal["low", "middle", "high"]
-ACCURACY_LEVELS: Dict[ACCURACY_LEVEL, int] = {"low": 3, "middle": 5, "high": 8}
+ACCURACY_LEVELS: dict[ACCURACY_LEVEL, int] = {"low": 3, "middle": 5, "high": 8}
 
 
 def get_volume(
@@ -93,7 +93,7 @@ def _get_volume_frustum_cone(tree: Tree, *, accuracy: int) -> float:
 
     volume = 0.0
 
-    def leave(n: Tree.Node, children: List[VolSphere]) -> VolSphere:
+    def leave(n: Tree.Node, children: list[VolSphere]) -> VolSphere:
         sphere = VolSphere(n.xyz(), n.r)
         cones = [VolFrustumCone(n.xyz(), n.r, c.center, c.radius) for c in children]
 
@@ -129,7 +129,7 @@ def _get_volume_frustum_cone_mc_only(tree: Tree) -> float:
     scene = ObjectsScene()
     scene.set_background((0, 0, 0))
 
-    def leave(n: Tree.Node, children: List[VolSphere]) -> VolSphere:
+    def leave(n: Tree.Node, children: list[VolSphere]) -> VolSphere:
         sphere = VolSphere(n.xyz(), n.r)
         scene.add_object(SDFObject(sphere.sdf, material).into())
 

swcgeom/core/branch.py

@@ -1,6 +1,7 @@
 """Branch is a set of node points."""
 
-from typing import Generic, Iterable, List
+from collections.abc import Iterable
+from typing import Generic
 
 import numpy as np
 import numpy.typing as npt
@@ -92,7 +93,7 @@ class Branch(Path, Generic[SWCTypeVar]):
     @classmethod
     def from_xyzr_batch(
         cls, xyzr_batch: npt.NDArray[np.float32]
-    ) -> List["Branch[DictSWC]"]:
+    ) -> list["Branch[DictSWC]"]:
         r"""Create list of branch form ~numpy.ndarray.
 
         Parameters
@@ -112,7 +113,7 @@ class Branch(Path, Generic[SWCTypeVar]):
             )
             xyzr_batch = np.concatenate([xyzr_batch, ones], axis=2)
 
-        branches: List[Branch[DictSWC]] = []
+        branches: list[Branch[DictSWC]] = []
         for xyzr in xyzr_batch:
             n_nodes = xyzr.shape[0]
             idx = np.arange(0, n_nodes, step=1, dtype=np.int32)

swcgeom/core/branch_tree.py

@@ -1,7 +1,6 @@
 """Branch tree is a simplified neuron tree."""
 
 import itertools
-from typing import Dict, List
 
 import numpy as np
 import pandas as pd
@@ -19,13 +18,13 @@ class BranchTree(Tree):
     A branch tree that contains only soma, branch, and tip nodes.
     """
 
-    branches: Dict[int, List[Branch]]
+    branches: dict[int, list[Branch]]
 
-    def get_origin_branches(self) -> List[Branch]:
+    def get_origin_branches(self) -> list[Branch]:
         """Get branches of original tree."""
         return list(itertools.chain(*self.branches.values()))
 
-    def get_origin_node_branches(self, idx: int) -> List[Branch]:
+    def get_origin_node_branches(self, idx: int) -> list[Branch]:
         """Get branches of node of original tree."""
         return self.branches[idx]
 

swcgeom/core/compartment.py

@@ -1,6 +1,7 @@
 """The segment is a branch with two nodes."""
 
-from typing import Generic, Iterable, List, TypeVar
+from collections.abc import Iterable
+from typing import Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -43,7 +44,7 @@ class Compartment(Path, Generic[SWCTypeVar]):
 T = TypeVar("T", bound=Compartment)
 
 
-class Compartments(List[T]):
+class Compartments(list[T]):
     r"""Comparments contains a set of comparment."""
 
     names: SWCNames

swcgeom/core/node.py

@@ -1,7 +1,7 @@
 """Nueron node."""
 
-import warnings
-from typing import Any, Generic, Iterable
+from collections.abc import Iterable
+from typing import Any, Generic
 
 import numpy as np
 import numpy.typing as npt

swcgeom/core/path.py

@@ -1,7 +1,8 @@
 """Nueron path."""
 
 import warnings
-from typing import Generic, Iterable, Iterator, List, overload
+from collections.abc import Iterable, Iterator
+from typing import Generic, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -44,7 +45,7 @@ class Path(SWCLike, Generic[SWCTypeVar]):
     @overload
     def __getitem__(self, key: int) -> Node: ...
     @overload
-    def __getitem__(self, key: slice) -> List[Node]: ...
+    def __getitem__(self, key: slice) -> list[Node]: ...
     @overload
     def __getitem__(self, key: str) -> npt.NDArray: ...
     # fmt:on