swcgeom 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

swcgeom/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.12.0'
-__version_tuple__ = version_tuple = (0, 12, 0)
+__version__ = version = '0.13.0'
+__version_tuple__ = version_tuple = (0, 13, 0)

swcgeom/analysis/feature_extractor.py

@@ -25,6 +25,7 @@ from swcgeom.analysis.node_features import (
 )
 from swcgeom.analysis.path_features import PathFeatures
 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
 
@@ -32,6 +33,7 @@ __all__ = ["Feature", "extract_feature"]
 
 Feature = Literal[
     "length",
+    "volume",
     "sholl",
     # node
     "node_count",
@@ -54,7 +56,7 @@ Feature = Literal[
 NDArrayf32 = npt.NDArray[np.float32]
 FeatAndKwargs = Feature | Tuple[Feature, Dict[str, Any]]
 
-Feature1D = set(["length", "node_count", "bifurcation_count", "tip_count"])
+Feature1D = set(["length", "volume", "node_count", "bifurcation_count", "tip_count"])
 
 
 class Features:
@@ -105,6 +107,9 @@ class 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)
 
@@ -288,9 +293,11 @@ class PopulationFeatureExtractor(FeatureExtractor):
         v = np.stack([padding1d(len_max, v, dtype=np.float32) for v in vals])
         return v
 
-    def _get_sholl_impl(self, **kwargs) -> Tuple[NDArrayf32, NDArrayf32]:
+    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=20)
+        rs = Sholl.get_rs(rmax=rmax, steps=steps)
         vals = self._get_impl("sholl", steps=rs, **kwargs)
         return vals, rs
 
@@ -360,10 +367,12 @@ class PopulationsFeatureExtractor(FeatureExtractor):
 
         return out
 
-    def _get_sholl_impl(self, **kwargs) -> Tuple[NDArrayf32, NDArrayf32]:
+    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=20)
+        rs = Sholl.get_rs(rmax=rmax, steps=steps)
         vals = self._get_impl("sholl", steps=rs, **kwargs)
         return vals, rs
 
@@ -408,7 +417,7 @@ class PopulationsFeatureExtractor(FeatureExtractor):
         return ax
 
 
-def extract_feature(obj: Tree | Population) -> FeatureExtractor:
+def extract_feature(obj: Tree | Population | Populations) -> FeatureExtractor:
     if isinstance(obj, Tree):
         return TreeFeatureExtractor(obj)
 

swcgeom/analysis/sholl.py

@@ -42,9 +42,12 @@ class Sholl:
         step: Optional[float] = None,
     ) -> None:
         tree = Tree.from_swc(tree) if isinstance(tree, str) else tree
-        self.tree = TranslateOrigin.transform(tree)  # shift
-        self.rs = np.linalg.norm(self.tree.get_segments().xyz(), axis=2)
-        self.rmax = self.rs.max()
+        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(
@@ -96,6 +99,7 @@ class Sholl:
         **kwargs :
             Forwarding to plot method.
         """
+
         if plot_type is not None:
             warnings.warn(
                 "`plot_type` has been renamed to `kind` since v0.5.0, "

swcgeom/analysis/volume.py

@@ -0,0 +1,80 @@
+"""Analysis of volume of a SWC tree."""
+
+from typing import List
+
+from swcgeom.core import Tree
+from swcgeom.utils import GeomFrustumCone, GeomSphere
+
+__all__ = ["get_volume"]
+
+
+def get_volume(tree: Tree):
+    """Get the volume of the tree.
+
+    Parameters
+    ----------
+    tree : Tree
+        SWC tree.
+
+    Returns
+    -------
+    volume : float
+        Volume of the tree.
+
+    Notes
+    -----
+    The SWC format is a method for representing neurons, which includes
+    both the radius of individual points and their interconnectivity.
+    Consequently, there are multiple distinct approaches to
+    representation within this framework.
+
+    Currently, we support a standard approach to volume calculation.
+    This method involves treating each node as a sphere and
+    representing the connections between them as truncated cone-like
+    structures, or frustums, with varying radii at their top and bottom
+    surfaces.
+
+    More representation methods will be supported in the future.
+    """
+    volume = 0.0
+
+    def leave(node: Tree.Node, children: List[GeomSphere]) -> GeomSphere:
+        sphere = GeomSphere(node.xyz(), node.r)
+        frustum_cones = [
+            GeomFrustumCone(node.xyz(), node.r, c.center, c.radius) for c in children
+        ]
+
+        v = sphere.get_volume()
+        v += sum(fc.get_volume() for fc in frustum_cones)
+        v -= sum(sphere.get_intersect_volume(fc) for fc in frustum_cones)
+        v -= sum(s.get_intersect_volume(fc) for s, fc in zip(children, frustum_cones))
+
+        # TODO
+        # remove volume of intersection between frustum cones
+        # v -= sum(
+        #     fc1.get_intersect_volume(fc2)
+        #     for fc1 in frustum_cones
+        #     for fc2 in frustum_cones
+        #     if fc1 != fc2
+        # )
+
+        nonlocal volume
+        volume += v
+        return sphere
+
+    tree.traverse(leave=leave)
+    return volume
+
+
+if __name__ == "__main__":
+    from io import StringIO
+
+    swc = """
+1 1 0 0 0 1 -1
+2 1 2 0 0 1  1
+3 1 4 0 0 1  2
+"""
+
+    tree = Tree.from_swc(StringIO(swc))
+    volume = get_volume(tree)
+    print(volume)

swcgeom/core/swc_utils/checker.py

@@ -1,17 +1,22 @@
 """Check common """
 
 import warnings
+from collections import defaultdict
 from typing import Optional
 
 import numpy as np
 import pandas as pd
 
 from swcgeom.core.swc_utils.base import SWCNames, Topology, get_dsu, get_names, traverse
+from swcgeom.utils import DisjointSetUnion
+
 
 __all__ = [
     "is_single_root",
     "is_bifurcate",
     "is_sorted",
+    "has_cyclic",
+    # legacy
     "is_binary_tree",
     "check_single_root",
 ]
@@ -25,13 +30,11 @@ def is_single_root(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> boo
 def is_bifurcate(topology: Topology, *, exclude_root: bool = True) -> bool:
     """Check is it a bifurcate topology."""
 
-    children = {}
+    children = defaultdict(list)
     for idx, pid in zip(*topology):
-        s = children.get(pid, [])
-        s.append(idx)
-        children[pid] = s
+        children[pid].append(idx)
 
-    root = children.get(-1, [])
+    root = children[-1]
     for k, v in children.items():
         if len(v) > 1 and (not exclude_root or k in root):
             return False
@@ -58,6 +61,27 @@ def is_sorted(topology: Topology) -> bool:
     return flag
 
 
+def has_cyclic(topology: Topology) -> bool:
+    """Has cyclic in topology."""
+    node_num = len(topology[0])
+    dsu = DisjointSetUnion(node_number=node_num)
+
+    for i in range(node_num):
+        node_a = topology[0][i]
+        node_b = topology[1][i]
+        # skip the root node
+        if node_b == -1:
+            continue
+
+        # check whether it is circle
+        if dsu.is_same_set(node_a, node_b):
+            return True
+
+        dsu.union_sets(node_a, node_b)
+
+    return False
+
+
 def check_single_root(*args, **kwargs) -> bool:
     warnings.warn(
         "`check_single_root` has been renamed to `is_single_root` since"

swcgeom/core/swc_utils/io.py

@@ -35,7 +35,7 @@ def read_swc(
 
     Parameters
     ----------
-    swc_file : str
+    swc_file : PathOrIO
         Path of swc file, the id should be consecutively incremented.
     extra_cols : Iterable[str], optional
         Read more cols in swc file.
@@ -142,7 +142,7 @@ def parse_swc(
 
     Parameters
     ----------
-    fname : str
+    fname : PathOrIO
     names : SWCNames
     extra_cols : list of str, optional
     encoding : str | 'detect', default `utf-8`

swcgeom/utils/__init__.py

@@ -1,8 +1,10 @@
 """Utils."""
 
 from swcgeom.utils.debug import *
+from swcgeom.utils.dsu import *
 from swcgeom.utils.ellipse import *
 from swcgeom.utils.file import *
+from swcgeom.utils.geometry_object import *
 from swcgeom.utils.neuromorpho import *
 from swcgeom.utils.numpy_helper import *
 from swcgeom.utils.renderer import *

swcgeom/utils/dsu.py

@@ -0,0 +1,42 @@
+"""Disjoint Set Union Impl."""
+
+__all__ = ["DisjointSetUnion"]
+
+
+class DisjointSetUnion:
+    """Disjoint Set Union.
+
+    DSU with path compression and union by rank.
+    """
+
+    def __init__(self, node_number: int):
+        self.element_parent = [i for i in range(node_number)]
+        self.rank = [0 for _ in range(node_number)]
+
+    def find_parent(self, node_id: int) -> int:
+        if node_id != self.element_parent[node_id]:
+            self.element_parent[node_id] = self.find_parent(
+                self.element_parent[node_id]
+            )
+        return self.element_parent[node_id]
+
+    def union_sets(self, node_a: int, node_b: int) -> None:
+        assert self.validate_node(node_a) and self.validate_node(node_b)
+
+        root_a = self.find_parent(node_a)
+        root_b = self.find_parent(node_b)
+        if root_a != root_b:
+            # union by rank
+            if self.rank[root_a] < self.rank[root_b]:
+                self.element_parent[root_a] = root_b
+            elif self.rank[root_a] > self.rank[root_b]:
+                self.element_parent[root_b] = root_a
+            else:
+                self.element_parent[root_b] = root_a
+                self.rank[root_a] += 1
+
+    def is_same_set(self, node_a: int, node_b: int) -> bool:
+        return self.find_parent(node_a) == self.find_parent(node_b)
+
+    def validate_node(self, node_id: int) -> bool:
+        return 0 <= node_id < len(self.element_parent)

swcgeom/utils/file.py

@@ -11,12 +11,12 @@ pip install swcgeom[all]
 """
 
 import warnings
-from io import BytesIO, TextIOWrapper
+from io import BytesIO, TextIOBase, TextIOWrapper
 from typing import Literal
 
 __all__ = ["FileReader", "PathOrIO"]
 
-PathOrIO = int | str | bytes | BytesIO | TextIOWrapper
+PathOrIO = int | str | bytes | BytesIO | TextIOBase
 
 
 class FileReader:
@@ -32,7 +32,7 @@ class FileReader:
 
         Parameters
         ----------
-        fname : str
+        fname : PathOrIO
         encoding : str | 'detect', default `utf-8`
             The name of the encoding used to decode the file. If is
             `detect`, we will try to detect the character encoding.
@@ -40,8 +40,9 @@ class FileReader:
             Used for detect character endocing, raising warning when
             parsing with low confidence.
         """
+        # TODO: support StringIO
         self.fname, self.fb, self.f = "", None, None
-        if isinstance(fname, TextIOWrapper):
+        if isinstance(fname, TextIOBase):
             self.f = fname
             encoding = fname.encoding  # skip detect
         elif isinstance(fname, BytesIO):
@@ -54,7 +55,7 @@ class FileReader:
         self.encoding = encoding
         self.kwargs = kwargs
 
-    def __enter__(self) -> TextIOWrapper:
+    def __enter__(self) -> TextIOBase:
         if isinstance(self.fb, BytesIO):
             self.f = TextIOWrapper(self.fb, encoding=self.encoding)
         elif self.f is None:
@@ -71,7 +72,7 @@ class FileReader:
 def detect_encoding(fname: PathOrIO, *, low_confidence: float = 0.9) -> str:
     import chardet
 
-    if isinstance(fname, TextIOWrapper):
+    if isinstance(fname, TextIOBase):
         return fname.encoding
     elif isinstance(fname, BytesIO):
         data = fname.read()

swcgeom/utils/geometry_object.py

@@ -0,0 +1,299 @@
+"""Geometry object."""
+
+from abc import ABC, abstractmethod
+from functools import lru_cache
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+from sympy import Eq, solve, symbols
+
+__all__ = ["GeomObject", "GeomSphere", "GeomFrustumCone"]
+
+eps = 1e-6
+
+
+class GeomObject(ABC):
+    """Geometry object."""
+
+    @abstractmethod
+    def get_volume(self) -> float:
+        """Get volume."""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_intersect_volume(self, obj: "GeomObject") -> float:
+        """Get intersect volume.
+
+        Parameters
+        ----------
+        obj : GeometryObject
+            Another geometry object.
+
+        Returns
+        -------
+        volume : float
+            Intersect volume.
+        """
+        raise NotImplementedError()
+
+
+class GeomSphere(GeomObject):
+    """Geometry Sphere."""
+
+    def __init__(self, center: npt.ArrayLike, radius: float):
+        super().__init__()
+
+        self.center = np.array(center)
+        assert len(self.center) == 3
+
+        self.radius = radius
+
+    def get_volume(self) -> float:
+        return self.calc_volume(self.radius)
+
+    def get_volume_spherical_cap(self, h: float) -> float:
+        return self.calc_volume_spherical_cap(self.radius, h)
+
+    def get_intersect_volume_sphere(self, obj: "GeomSphere") -> float:
+        return self.calc_intersect_volume_sphere(self, obj)
+
+    def get_intersect_volume_sphere_frustum_cone(
+        self, frustum_cone: "GeomFrustumCone"
+    ) -> float:
+        return calc_intersect_volume_sphere_frustum_cone(self, frustum_cone)
+
+    def get_intersect_volume(self, obj: GeomObject) -> float:
+        if isinstance(obj, GeomSphere):
+            return self.get_intersect_volume_sphere(obj)
+
+        if isinstance(obj, GeomFrustumCone):
+            return self.get_intersect_volume_sphere_frustum_cone(obj)
+
+        classname = obj.__class__.__name__
+        raise NotImplementedError(f"unsupported geometry object: {classname}")
+
+    @staticmethod
+    def calc_volume(radius: float) -> float:
+        r"""Calculate volume of sphere.
+
+        \being{equation}
+        V = \frac{4}{3} * π * r^3
+        \end{equation}
+
+        Returns
+        -------
+        volume : float
+            Volume.
+        """
+        return 4 / 3 * np.pi * radius**3
+
+    @staticmethod
+    def calc_volume_spherical_cap(r: float, h: float) -> float:
+        r"""Calculate the volume of a spherical cap.
+
+        \being{equation}
+        V = π * h^2 * (3r - h) / 3
+        \end{equation}
+
+        Parameters
+        ----------
+        r : float
+            radius of the sphere
+        h : float
+            height of the spherical cap
+
+        Returns
+        -------
+        volume : float
+            volume of the spherical cap
+        """
+        return np.pi * h**2 * (3 * r - h) / 3
+
+    @classmethod
+    def calc_intersect_volume_sphere(
+        cls, obj1: "GeomSphere", obj2: "GeomSphere"
+    ) -> float:
+        r"""Calculate intersect volume of two spheres.
+
+        \being{equation}
+        V = \frac{\pi}{12d} * (r_1 + r_2 - d)^2 (d^2 + 2d r_1 - 3r_1^2 + 2d r_2 - 3r_2^2 + 6 r_1r_2)
+        \end{equation}
+
+        Returns
+        -------
+        volume : float
+            Intersect volume.
+        """
+
+        r1, r2 = obj1.radius, obj2.radius
+        d = np.linalg.norm(obj1.center - obj2.center).item()
+        if d > r1 + r2:
+            return 0
+
+        if d <= abs(r1 - r2):
+            return cls.calc_volume(min(r1, r2))
+
+        part1 = (np.pi / (12 * d)) * (r1 + r2 - d) ** 2
+        part2 = (
+            d**2 + 2 * d * r1 - 3 * r1**2 + 2 * d * r2 - 3 * r2**2 + 6 * r1 * r2
+        )
+        return part1 * part2
+
+
+class GeomFrustumCone(GeomObject):
+    """Geometry Frustum."""
+
+    def __init__(self, c1: npt.ArrayLike, r1: float, c2: npt.ArrayLike, r2: float):
+        super().__init__()
+
+        self.c1 = np.array(c1)
+        assert len(self.c1) == 3
+
+        self.c2 = np.array(c2)
+        assert len(self.c2) == 3
+
+        self.r1 = r1
+        self.r2 = r2
+
+    def height(self) -> float:
+        """Get height of frustum."""
+        return np.linalg.norm(self.c1 - self.c2).item()
+
+    def get_volume(self) -> float:
+        return self.calc_volume(self.r1, self.r2, self.height())
+
+    def get_intersect_volume_sphere(self, sphere: GeomSphere) -> float:
+        return calc_intersect_volume_sphere_frustum_cone(sphere, self)
+
+    def get_intersect_volume(self, obj: GeomObject) -> float:
+        if isinstance(obj, GeomSphere):
+            return self.get_intersect_volume_sphere(obj)
+
+        classname = obj.__class__.__name__
+        raise NotImplementedError(f"unsupported geometry object: {classname}")
+
+    @staticmethod
+    def calc_volume(r1: float, r2: float, height: float) -> float:
+        r"""Calculate volume of frustum.
+
+        \being{equation}
+        V = \frac{1}{3} * π * h * (r^2 + r * R + R^2)
+        \end{equation}
+
+        Returns
+        -------
+        volume : float
+            Volume.
+        """
+        return (1 / 3) * np.pi * height * (r1**2 + r1 * r2 + r2**2)
+
+
+@lru_cache
+def calc_intersect_volume_sphere_frustum_cone(
+    sphere: GeomSphere, frustum_cone: GeomFrustumCone
+) -> float:
+    r"""Calculate intersect volume of sphere and frustum cone.
+
+    Returns
+    -------
+    volume : float
+        Intersect volume.
+    """
+    h = frustum_cone.height()
+    c1, r1 = sphere.center, sphere.radius
+    if np.allclose(c1, frustum_cone.c1) and np.allclose(r1, frustum_cone.r1):
+        c2, r2 = frustum_cone.c2, frustum_cone.r2
+    elif np.allclose(c1, frustum_cone.c2) and np.allclose(r1, frustum_cone.r2):
+        c2, r2 = frustum_cone.c1, frustum_cone.r1
+    else:
+        raise NotImplementedError("unsupported to calculate intersect volume")
+
+    # Fast-Path: The surface of the frustum concentric with the sphere
+    # is the surface with smaller radius
+    if r2 >= r1:
+        if h >= r1:
+            # The hemisphere is completely inside the frustum cone
+            return GeomSphere.calc_volume_spherical_cap(r1, r1)
+
+        # The frustum cone is lower than the hemisphere
+        v_himisphere = GeomSphere.calc_volume_spherical_cap(r1, r1)
+        v_cap = GeomSphere.calc_volume_spherical_cap(r1, r1 - h)
+        return v_himisphere - v_cap
+
+    up = (c2 - c1) / np.linalg.norm(c2 - c1)
+    v = _find_unit_vector_on_plane(up)
+
+    intersections = _find_sphere_line_intersection(c1, r1, c1 + r1 * v, c2 + r2 * v)
+    assert len(intersections) == 2
+    t, p = max(intersections, key=lambda x: x[0])
+
+    # Fast-Path: The frustum cone is completely inside the sphere
+    if t > 1 + eps:
+        return frustum_cone.get_volume()
+
+    M = _project_point_on_line(c1, up, p)
+    h1 = np.linalg.norm(M - c1).item()
+    r3 = np.linalg.norm(M - p).item()
+    v_cap1 = GeomSphere.calc_volume_spherical_cap(r1, r1 - h1)
+    v_frustum = GeomFrustumCone.calc_volume(r1, r3, h1)
+
+    # Fast-Path: The frustum cone is higher than the sphere
+    if h >= r1:
+        return v_cap1 + v_frustum
+
+    v_cap2 = GeomSphere.calc_volume_spherical_cap(r1, r1 - h)
+    return v_cap1 + v_frustum - v_cap2
+
+
+def _find_unit_vector_on_plane(normal_vec3: npt.NDArray) -> npt.NDArray:
+    r = np.random.rand(3)
+    while np.allclose(r, normal_vec3) or np.allclose(r, -normal_vec3):
+        r = np.random.rand(3)
+
+    u = np.cross(r, normal_vec3)
+    unit_vector = u / np.linalg.norm(u)
+    return unit_vector
+
+
+def _find_sphere_line_intersection(
+    sphere_center: npt.NDArray,
+    sphere_radius: float,
+    line_point_a: npt.NDArray,
+    line_point_b: npt.NDArray,
+) -> List[Tuple[float, npt.NDArray[np.float64]]]:
+    x1, y1, z1 = sphere_center
+    x2, y2, z2 = line_point_a
+    x3, y3, z3 = line_point_b
+    t = symbols("t")
+
+    # line
+    x = x2 + t * (x3 - x2)
+    y = y2 + t * (y3 - y2)
+    z = z2 + t * (z3 - z2)
+
+    # sphere
+    sphere_eq = Eq((x - x1) ** 2 + (y - y1) ** 2 + (z - z1) ** 2, sphere_radius**2)
+
+    # solve
+    t_values = solve(sphere_eq, t)
+    intersections = [
+        np.array(
+            [float(x.subs(t, t_val)), float(y.subs(t, t_val)), float(z.subs(t, t_val))]
+        )
+        for t_val in t_values
+    ]
+    return list(zip(t_values, intersections))
+
+
+def _project_point_on_line(
+    point_a: npt.ArrayLike, direction_vector: npt.ArrayLike, point_p: npt.ArrayLike
+) -> npt.NDArray:
+    A = np.array(point_a)
+    n = np.array(direction_vector)
+    P = np.array(point_p)
+    print(A.dtype, n.dtype, P.dtype)
+
+    AP = P - A
+    projection = A + np.dot(AP, n) / np.dot(n, n) * n
+    return projection

swcgeom/utils/neuromorpho.py

@@ -97,14 +97,19 @@ URL_CNG_VERSION = (
 )
 API_NEURON_MAX_SIZE = 500
 
-# about 1.1 GB and 18 GB in version 8.5.25 released in 2023-08-01
 KB = 1024
 MB = 1024 * KB
 GB = 1024 * MB
+
+# Test version: 8.5.25 (2023-08-01)
+# About 1.1 GB and 18 GB
+# No ETAs for future version
 SIZE_METADATA = 2 * GB
 SIZE_DATA = 20 * GB
 
 # fmt:off
+# Test version: 8.5.25 (2023-08-01)
+# No ETAs for future version
 invalid_ids = [
     # bad file
     81062, 86970, 79791,
@@ -171,7 +176,7 @@ def neuromorpho_convert_lmdb_to_swc(
     See Also
     --------
     neuromorpho_is_valid :
-        Recommend filter function, use `where=neuromorpho_is_valid`
+        Recommended filter function, try `where=neuromorpho_is_valid`
     """
     import lmdb
     from tqdm import tqdm

swcgeom/utils/renderer.py

@@ -9,7 +9,7 @@ import numpy.typing as npt
 from matplotlib import cm
 from matplotlib.axes import Axes
 from matplotlib.collections import LineCollection, PatchCollection
-from matplotlib.colors import Normalize
+from matplotlib.colors import Colormap, Normalize
 from matplotlib.figure import Figure
 from matplotlib.patches import Circle
 from typing_extensions import Self
@@ -208,14 +208,15 @@ def draw_circles(
     *,
     y_min: Optional[float] = None,
     y_max: Optional[float] = None,
-    cmap: str = "viridis",
+    cmap: str | Colormap = "viridis",
 ) -> PatchCollection:
     """Draw a sequential of circles."""
+
     y_min = y.min() if y_min is None else y_min
     y_max = y.max() if y_max is None else y_max
     norm = Normalize(y_min, y_max)
 
-    color_map = cm.get_cmap(name=cmap)
+    color_map = cmap if isinstance(cmap, Colormap) else cm.get_cmap(name=cmap)
     colors = color_map(norm(y))
 
     circles = [

swcgeom/utils/transforms.py

@@ -212,6 +212,31 @@ def to_homogeneous(xyz: npt.ArrayLike, w: float) -> npt.NDArray[np.float32]:
     Parameters
     ----------
     xyz : ArrayLike
+        Coordinate of shape (..., 3)
+    w : float
+        w of homogeneous coordinate, 1 for dot, 0 for vector.
+
+    Returns
+    -------
+    xyz4 : npt.NDArray[np.float32]
+        Array of shape (..., 4)
+    """
+    xyz = np.array(xyz)
+    if xyz.ndim == 1:
+        return _to_homogeneous(xyz[None, ...], w)[0]
+
+    shape = xyz.shape[:-1]
+    xyz = xyz.reshape(-1, xyz.shape[-1])
+    xyz4 = _to_homogeneous(xyz, w).reshape(*shape, 4)
+    return xyz4
+
+
+def _to_homogeneous(xyz: npt.NDArray, w: float) -> npt.NDArray[np.float32]:
+    """Fill xyz to homogeneous coordinates.
+
+    Parameters
+    ----------
+    xyz : npt.NDArray
         Coordinate of shape (N, 3)
     w : float
         w of homogeneous coordinate, 1 for dot, 0 for vector.
@@ -221,10 +246,10 @@ def to_homogeneous(xyz: npt.ArrayLike, w: float) -> npt.NDArray[np.float32]:
     xyz4 : npt.NDArray[np.float32]
         Array of shape (N, 4)
     """
-    xyz = np.array(xyz)
     if xyz.shape[1] == 4:
         return xyz
 
+    assert xyz.shape[1] == 3
     filled = np.full((xyz.shape[0], 1), fill_value=w)
     xyz4 = np.concatenate([xyz, filled], axis=1)
     return xyz4

{swcgeom-0.12.0.dist-info → swcgeom-0.13.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: swcgeom
-Version: 0.12.0
+Version: 0.13.0
 Summary: Neuron geometry library for swc format
 Author-email: yzx9 <yuan.zx@outlook.com>
 License: CC4.0 BY-NC-SA
@@ -16,6 +16,7 @@ Requires-Dist: pandas >=1.4.2
 Requires-Dist: pynrrd >=1.0.0
 Requires-Dist: scipy >=1.9.1
 Requires-Dist: seaborn >=0.12.0
+Requires-Dist: sympy >=1.12
 Requires-Dist: tifffile >=2022.8.12
 Requires-Dist: typing-extensions >=4.4.0
 Requires-Dist: v3d-py-helper-0.1.0

{swcgeom-0.12.0.dist-info → swcgeom-0.13.0.dist-info}/RECORD

@@ -1,13 +1,14 @@
 swcgeom/__init__.py,sha256=z88Zwcjv-ii7c7dYd9QPg9XrUVorQjtrgGbQCsEnQhc,265
-swcgeom/_version.py,sha256=HJyq47ePtlF3WPYAhFFp5LPsrHyF6oju_vSEfklwWro,413
+swcgeom/_version.py,sha256=Y4u7iBqF7QJAbpBNSFA2tk5t2mrMGrI-nonxUAhVkPU,413
 swcgeom/analysis/__init__.py,sha256=c8X_oMxaOijEVjHexFseYs49WMalG_mSxmlQZr5CZeI,339
 swcgeom/analysis/branch_features.py,sha256=s6PTMwwvxrVtZXRZlQbUSIw4M9-1IG63kf-Nxc0tMB0,1958
-swcgeom/analysis/feature_extractor.py,sha256=HKVHCnNQsAgnyX1BimEDUX0q-caZ0_VZ8wqPV7lC_R4,13750
+swcgeom/analysis/feature_extractor.py,sha256=coe07_bJau96BkimcnXzuf4KqjY5_QRLwqaumFsu_tQ,14031
 swcgeom/analysis/node_features.py,sha256=fevnyrF-t4PX39ifLypiDW6EUWB8i-a3PpBnQZU3VOc,3407
 swcgeom/analysis/path_features.py,sha256=iE21HBxAoGLxk_qK7MBwQhyUOBqNPcnk4urVHr9SVqk,889
-swcgeom/analysis/sholl.py,sha256=QlPb2aCSzNCXmQkVFH_W1GRATPIngTZdNtZgBXvFEHk,6259
+swcgeom/analysis/sholl.py,sha256=9hSW8rZm1gvSIgcEZg8IVPT8kzBgBfwqbwP4E8R7L44,6390
 swcgeom/analysis/trunk.py,sha256=L2tjUIUmrRQpah_W3ZETGWd16bDXJ5F8Sk2XBNGms0Q,5558
 swcgeom/analysis/visualization.py,sha256=mKOpzTPkLpr1ggGL1MZPZRTG92GEg4idLT4eN5z5KOs,5654
+swcgeom/analysis/volume.py,sha256=bMQ9N9Ow1-h27SLQ9AzUm9dyF_jhoTNBMNVqQWDcHMw,2103
 swcgeom/core/__init__.py,sha256=ZUudZavxAIUU6Q0lBHrQ4ybmL5lBfvzyYsTtpuih9wg,332
 swcgeom/core/branch.py,sha256=uuJCxaByRu-OdDZVWEffSFcmZWY-6ZWUhHN1M2Awj1s,3980
 swcgeom/core/branch_tree.py,sha256=sN0viBVg5A4r9dMCkGNAaVttrdR4bEoZZBbHZFKdXj4,1892
@@ -22,8 +23,8 @@ swcgeom/core/tree_utils_impl.py,sha256=5Cb63ziVVLADnkjfuq1T-ePw2TQQ5TKk4gcPZR6f_
 swcgeom/core/swc_utils/__init__.py,sha256=qghRxjtzvq5KKfN4HhvLpZNsGPfZQu-Jj2x62_5-TbQ,575
 swcgeom/core/swc_utils/assembler.py,sha256=RoMJ3RjLC4O7mk62QxXVTQ5SUHagrFmpEw3nOnnqeJo,4563
 swcgeom/core/swc_utils/base.py,sha256=huVxjuMLlTHbEb-KSEFDLgU0Ss3723t2Gr4Z_gQtl00,4737
-swcgeom/core/swc_utils/checker.py,sha256=piyvIvskJi7SoyURLZV68q1mDX7FYHHf5bW2BPbc-QM,2092
-swcgeom/core/swc_utils/io.py,sha256=-O2lIATRfcK_eQ_jkJN790oieoqY4pwaNEqhlCnevFM,6434
+swcgeom/core/swc_utils/checker.py,sha256=E72GtLZ_1IqQQ7aWQGs0dZ3Z609__bw3EYQqeWrk-EI,2657
+swcgeom/core/swc_utils/io.py,sha256=6_--Qoe8kDja4PWsjwqRAvPJZNMFILFgauHaeWeGikU,6444
 swcgeom/core/swc_utils/normalizer.py,sha256=_Ysi8bSJ2JBnIGB8o6BvAg2mcz6xuJp9rgNLZqpLuR8,5083
 swcgeom/core/swc_utils/subtree.py,sha256=bd4XOLmRDfQSn_ktfQM3Hn8ONpCuZ_TdTWhE9-7QXW4,1999
 swcgeom/images/__init__.py,sha256=QBP1ZGGo2nWAcV7Krz-vbvW_jN4ChqXrrpoScXcUURs,96
@@ -41,18 +42,20 @@ swcgeom/transforms/path.py,sha256=Gk2iunGQMX7vE83bdo8xoDO-KAT1Vvep0iZs7oFLzFQ,10
 swcgeom/transforms/population.py,sha256=ZrKfMAMx4l729f-JLgw0dnGIPtPUoV0ZZoNNyA5cBw8,826
 swcgeom/transforms/tree.py,sha256=Q5OnSti0ZeTNb-WpA_UZsDN7dhLN8YweEF_Siyrj66c,6420
 swcgeom/transforms/tree_assembler.py,sha256=UZ9OUg1bQNecYIY_7ippg3S8gpuoi617ZUE0jg6BrQE,3177
-swcgeom/utils/__init__.py,sha256=g0bfRbRa4iC_ZovrJgQfEmP6mclMt2vYlZiaAZA5Rqg,306
+swcgeom/utils/__init__.py,sha256=Ys04ObDfoNOX0mft2s4PGTpgelkdgwuM0qfeZRSx6VM,382
 swcgeom/utils/debug.py,sha256=qay2qJpViLX82mzxdndxQFn-pi1vaEj9CbLGuGt8Y9k,465
 swcgeom/utils/download.py,sha256=By2qZezo6h1Ke_4YpSIhDgcisOrpjVqRmNzbhynC2xs,2834
+swcgeom/utils/dsu.py,sha256=3aCbtpnl_D0OXnowTS8-kuwnCS4BKBYL5ECiFQ1fUW8,1435
 swcgeom/utils/ellipse.py,sha256=LB3q5CIy75GEUdTauIpKySwIHaDrwXzzkBhOCnjJ8Vw,3259
-swcgeom/utils/file.py,sha256=qrJYhTscJqQAbJ-R0vpHPDP9Hjf9u-woY4l4lQLLG8E,2477
-swcgeom/utils/neuromorpho.py,sha256=PDSNVgXqKWO4qQI-VjtwL10--cVzxp7k7OXGHihRtdw,14328
+swcgeom/utils/file.py,sha256=1hchQDsPgn-i-Vz5OQtcogxav_ajCQ_OaEZCLmqczRg,2515
+swcgeom/utils/geometry_object.py,sha256=y3Ikg0ywYQ6c_KKZOb0oq5ovHhuFiwI_VH-f4tAWZAI,8520
+swcgeom/utils/neuromorpho.py,sha256=CDK2tUM2pNwHv_lEserHhQs_VlY3Rn557-jtV63EFlk,14420
 swcgeom/utils/numpy_helper.py,sha256=A-F-eFdGktCHVAQ_HcXiFB3Y1YhhSNfAmtOl8483Dvo,1292
-swcgeom/utils/renderer.py,sha256=AEbPSmIxxrHf4y3Y6FtLwVWgq8akh-AQls1LbpqcXmA,7228
+swcgeom/utils/renderer.py,sha256=xHVZ06Z1MeKBPC3nKzuwA1HryzR0ga79y6johZA4-q0,7290
 swcgeom/utils/sdf.py,sha256=cmkHwdnh33u1gAXlzol5u9ZI3_SdMoMiY4pIkOgEcK0,5206
-swcgeom/utils/transforms.py,sha256=4U6flX8lfnMoAAR-uc4uSYjSd2vg-AOWGLgZ6_RVxJU,6352
-swcgeom-0.12.0.dist-info/LICENSE,sha256=aiTJ_1to1Xx6PaByy-pSXg42VYzE4FvF6GIt69WSDDI,202
-swcgeom-0.12.0.dist-info/METADATA,sha256=lyjjYsqQr4AjKf0O1nhY9hOUMz9QxX3vpB9S0gLOLb8,2577
-swcgeom-0.12.0.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-swcgeom-0.12.0.dist-info/top_level.txt,sha256=hmLyUXWS61Gxl07haswFEKKefYPBVJYlUlol8ghNkjY,8
-swcgeom-0.12.0.dist-info/RECORD,,
+swcgeom/utils/transforms.py,sha256=PmP5fL_iVguq4GR2aqXhM0TeCsiFVnrPZMZG6zLohrE,6983
+swcgeom-0.13.0.dist-info/LICENSE,sha256=aiTJ_1to1Xx6PaByy-pSXg42VYzE4FvF6GIt69WSDDI,202
+swcgeom-0.13.0.dist-info/METADATA,sha256=vsgLGIGz97A-V8n1DZsAhGAc-uVJuX_CdFbB2BCQZ5g,2605
+swcgeom-0.13.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+swcgeom-0.13.0.dist-info/top_level.txt,sha256=hmLyUXWS61Gxl07haswFEKKefYPBVJYlUlol8ghNkjY,8
+swcgeom-0.13.0.dist-info/RECORD,,

{swcgeom-0.12.0.dist-info → swcgeom-0.13.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.2)
+Generator: bdist_wheel (0.42.0)
 Root-Is-Purelib: true
 Tag: py3-none-any