swcgeom 0.19.4__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

swcgeom/analysis/volume.py

@@ -0,0 +1,143 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Analysis of volume of a SWC tree."""
+
+from typing import Literal
+
+import numpy as np
+from sdflit import ColoredMaterial, ObjectsScene, SDFObject, UniformSampler
+
+from swcgeom.core import Tree
+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}
+
+
+def get_volume(
+    tree: Tree,
+    *,
+    method: Literal["frustum_cone"] = "frustum_cone",
+    accuracy: int | ACCURACY_LEVEL = "middle",
+) -> float:
+    """Get the volume of the tree.
+
+    Args:
+        tree: Neuronal tree.
+        method: Method for volume calculation.
+        accuracy: Accuracy level for volume calculation. The higher the accuracy,
+            the more accurate the volume calculation, but the slower the
+            calculation. The accuracy level can be specified either as an
+            integer or as a string.
+
+            The string values correspond to the following accuracy levels:
+
+            - "low": 3
+            - "middle": 5
+            - "high": 8
+
+    Returns:
+        volume: Volume of the tree.
+
+    NOTE: 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.
+
+    We welcome additional representation methods through pull requests.
+    """
+    if isinstance(accuracy, str):
+        accuracy = ACCURACY_LEVELS[accuracy]
+
+    assert 0 < accuracy <= 10
+
+    match method:
+        case "frustum_cone":
+            return _get_volume_frustum_cone(tree, accuracy=accuracy)
+        case _:
+            raise ValueError(f"Unsupported method: {method}")
+
+
+def _get_volume_frustum_cone(tree: Tree, *, accuracy: int) -> float:
+    """Get the volume of the tree using the frustum cone method.
+
+    Args:
+        tree: Neuronal tree.
+        accuracy: Accuracy level.
+            1 : Sphere only
+            2 : Sphere and Frustum Cone
+            3 : Sphere, Frustum Cone, and intersection in single-branch
+            5 : Above and Sphere-Frustum Cone intersection in multi-branch
+            10 : Fully calculated by Monte Carlo method
+    """
+    if accuracy == 10:
+        return _get_volume_frustum_cone_mc_only(tree)
+
+    volume = 0.0
+
+    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]
+
+        v = sphere.get_volume()
+        if accuracy >= 2:
+            v += sum(fc.get_volume() for fc in cones)
+
+        if accuracy >= 3:
+            v -= sum(sphere.intersect(fc).get_volume() for fc in cones)
+            v -= sum(s.intersect(fc).get_volume() for s, fc in zip(children, cones))
+            v += sum(s.intersect(sphere).get_volume() for s in children)
+
+        if accuracy >= 5:
+            v -= sum(
+                cones[i].intersect(cones[j]).subtract(sphere).get_volume()
+                for i in range(len(cones))
+                for j in range(i + 1, len(cones))
+            )
+
+        nonlocal volume
+        volume += v
+        return sphere
+
+    tree.traverse(leave=leave)
+    return volume
+
+
+def _get_volume_frustum_cone_mc_only(tree: Tree) -> float:
+    if tree.number_of_nodes() == 0:
+        return 0
+
+    material = ColoredMaterial((1, 0, 0)).into()
+    scene = ObjectsScene()
+    scene.set_background((0, 0, 0))
+
+    def leave(n: Tree.Node, children: list[VolSphere]) -> VolSphere:
+        sphere = VolSphere(n.xyz(), n.r)
+        scene.add_object(SDFObject(sphere.sdf, material).into())
+
+        for c in children:
+            fc = VolFrustumCone(n.xyz(), n.r, c.center, c.radius)
+            scene.add_object(SDFObject(fc.sdf, material).into())
+
+        return sphere
+
+    tree.traverse(leave=leave)
+    scene.build_bvh()
+
+    # TODO: estimate the number of samples needed
+    n_samples = 100_000_000
+
+    vmin, vmax = scene.bounding_box()
+    sampler = UniformSampler(vmin, vmax)
+    data = sampler.sample(scene.into(), n_samples)
+    volume = data.sum() / n_samples * np.subtract(vmax, vmin).prod()
+    return volume

swcgeom/core/__init__.py

@@ -0,0 +1,19 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Neuron trees."""
+
+from swcgeom.core import swc_utils as swc_utils
+from swcgeom.core.branch import *  # noqa: F403
+from swcgeom.core.branch_tree import *  # noqa: F403
+
+# Segment and Segments don't expose
+from swcgeom.core.compartment import Compartment, Compartments  # noqa: F401
+from swcgeom.core.node import *  # noqa: F403
+from swcgeom.core.path import *  # noqa: F403
+from swcgeom.core.population import *  # noqa: F403
+from swcgeom.core.swc import *  # noqa: F403
+from swcgeom.core.tree import *  # noqa: F403
+from swcgeom.core.tree_utils import *  # noqa: F403

swcgeom/core/branch.py

@@ -0,0 +1,129 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Branch is a set of node points."""
+
+from collections.abc import Iterable
+from typing import Generic
+
+import numpy as np
+import numpy.typing as npt
+
+from swcgeom.core.compartment import Compartment, Compartments
+from swcgeom.core.path import Path
+from swcgeom.core.swc import DictSWC, SWCTypeVar
+
+__all__ = ["Branch"]
+
+
+class Branch(Path, Generic[SWCTypeVar]):
+    r"""Neural branch.
+
+    NOTE: Only a part of data of branch nodes is valid, such as `x`, `y`, `z` and `r`,
+    but the `id` and `pid` is usually invalid.
+    """
+
+    attach: SWCTypeVar
+    idx: npt.NDArray[np.int32]
+
+    class Compartment(Compartment["Branch"]):
+        """Segment of branch."""
+
+    Segment = Compartment  # Alias
+
+    def __repr__(self) -> str:
+        return f"Neuron branch with {len(self)} nodes."
+
+    def keys(self) -> Iterable[str]:
+        return self.attach.keys()
+
+    def get_ndata(self, key: str) -> npt.NDArray:
+        return self.attach.get_ndata(key)[self.idx]
+
+    def get_compartments(self) -> Compartments[Compartment]:
+        return Compartments(self.Compartment(self, n.pid, n.id) for n in self[1:])
+
+    def get_segments(self) -> Compartments[Compartment]:
+        return self.get_compartments()  # Alias
+
+    def detach(self) -> "Branch[DictSWC]":
+        """Detach from current attached object."""
+        # pylint: disable=consider-using-dict-items
+        attact = DictSWC(
+            **{k: self[k] for k in self.keys()},
+            source=self.attach.source,
+            names=self.names,
+        )
+        attact.ndata[self.names.id] = self.id()
+        attact.ndata[self.names.pid] = self.pid()
+        return Branch(attact, self.id())
+
+    @classmethod
+    def from_xyzr(cls, xyzr: npt.NDArray[np.float32]) -> "Branch[DictSWC]":
+        r"""Create a branch from ~numpy.ndarray.
+
+        Args:
+            xyzr: Collection of nodes.
+                If shape (n, 4), both `x`, `y`, `z`, `r` of nodes is enabled. If shape
+                (n, 3), only `x`, `y`, `z` is enabled and `r` will fill by 1.
+        """
+        assert xyzr.ndim == 2 and xyzr.shape[1] in (
+            3,
+            4,
+        ), f"xyzr should be of shape (N, 3) or (N, 4), got {xyzr.shape}"
+
+        n_nodes = xyzr.shape[0]
+        if xyzr.shape[1] == 3:
+            ones = np.ones([n_nodes, 1], dtype=np.float32)
+            xyzr = np.concatenate([xyzr, ones], axis=1)
+
+        idx = np.arange(0, n_nodes, step=1, dtype=np.int32)
+        attact = DictSWC(
+            id=idx,
+            type=np.full((n_nodes), fill_value=3, dtype=np.int32),
+            x=xyzr[:, 0],
+            y=xyzr[:, 1],
+            z=xyzr[:, 2],
+            r=xyzr[:, 3],
+            pid=np.arange(-1, n_nodes - 1, step=1, dtype=np.int32),
+        )
+        return Branch(attact, idx)
+
+    @classmethod
+    def from_xyzr_batch(
+        cls, xyzr_batch: npt.NDArray[np.float32]
+    ) -> list["Branch[DictSWC]"]:
+        r"""Create list of branch form ~numpy.ndarray.
+
+        Args:
+            xyzr: Batch of collection of nodes.
+                If shape (bs, n, 4), both `x`, `y`, `z`, `r` of nodes is enabled. If
+                shape (bs, n, 3), only `x`, `y`, `z` is enabled and `r` will fill by 1.
+        """
+        assert xyzr_batch.ndim == 3
+        assert xyzr_batch.shape[1] >= 3
+
+        if xyzr_batch.shape[2] == 3:
+            ones = np.ones(
+                [xyzr_batch.shape[0], xyzr_batch.shape[1], 1], dtype=np.float32
+            )
+            xyzr_batch = np.concatenate([xyzr_batch, ones], axis=2)
+
+        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)
+            attact = DictSWC(
+                id=idx,
+                type=np.full((n_nodes), fill_value=3, dtype=np.int32),
+                x=xyzr[:, 0],
+                y=xyzr[:, 1],
+                z=xyzr[:, 2],
+                r=xyzr[:, 3],
+                pid=np.arange(-1, n_nodes - 1, step=1, dtype=np.int32),
+            )
+            branches.append(Branch(attact, idx))
+
+        return branches

swcgeom/core/branch_tree.py

@@ -0,0 +1,65 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Branch tree is a simplified neuron tree."""
+
+import itertools
+
+import numpy as np
+import pandas as pd
+from typing_extensions import Self
+
+from swcgeom.core.branch import Branch
+from swcgeom.core.swc_utils import to_sub_topology
+from swcgeom.core.tree import Tree
+
+__all__ = ["BranchTree"]
+
+
+class BranchTree(Tree):
+    """Branch tree keeps topology of tree.
+
+    A branch tree that contains only soma, branch, and tip nodes.
+    """
+
+    branches: dict[int, 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]:
+        """Get branches of node of original tree."""
+        return self.branches[idx]
+
+    @classmethod
+    def from_tree(cls, tree: Tree) -> Self:
+        """Generating a branch tree from tree."""
+
+        branches = tree.get_branches()
+
+        sub_id = np.array([0] + [br[-1].id for br in branches], dtype=np.int32)
+        sub_pid = np.array([-1] + [br[0].id for br in branches], dtype=np.int32)
+
+        (new_id, new_pid), id_map = to_sub_topology((sub_id, sub_pid))
+
+        n_nodes = new_id.shape[0]
+        ndata = {k: tree.get_ndata(k)[id_map].copy() for k in tree.keys()}
+        ndata.update(id=new_id, pid=new_pid)
+
+        branch_tree = cls(n_nodes, **ndata, source=tree.source, names=tree.names)
+
+        branch_tree.branches = {}
+        for br in branches:
+            idx = np.nonzero(id_map == br[0].id)[0][0].item()
+            branch_tree.branches.setdefault(idx, [])
+            branch_tree.branches[idx].append(br.detach())
+
+        return branch_tree
+
+    @classmethod
+    def from_data_frame(cls, df: pd.DataFrame, *args, **kwargs) -> Self:
+        tree = super().from_data_frame(df, *args, **kwargs)
+        return cls.from_tree(tree)

swcgeom/core/compartment.py

@@ -0,0 +1,107 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""The segment is a branch with two nodes."""
+
+from collections.abc import Iterable
+from typing import Generic, TypeVar
+
+import numpy as np
+import numpy.typing as npt
+
+from swcgeom.core.path import Path
+from swcgeom.core.swc import DictSWC, SWCTypeVar
+from swcgeom.core.swc_utils import SWCNames, get_names
+
+__all__ = ["Compartment", "Compartments", "Segment", "Segments"]
+
+
+class Compartment(Path, Generic[SWCTypeVar]):
+    r"""Compartment attached to external object."""
+
+    attach: SWCTypeVar
+    idx: npt.NDArray[np.int32]
+
+    def __init__(self, attach: SWCTypeVar, pid: int, idx: int) -> None:
+        super().__init__(attach, np.array([pid, idx]))
+
+    def keys(self) -> Iterable[str]:
+        return self.attach.keys()
+
+    def get_ndata(self, key: str) -> npt.NDArray:
+        return self.attach.get_ndata(key)[self.idx]
+
+    def detach(self) -> "Compartment[DictSWC]":
+        """Detach from current attached object."""
+        # pylint: disable=consider-using-dict-items
+        attact = DictSWC(
+            **{k: self[k] for k in self.keys()},
+            source=self.attach.source,
+            names=self.names,
+        )
+        attact.ndata[self.names.id] = self.id()
+        attact.ndata[self.names.pid] = self.pid()
+        return Compartment(attact, 0, 1)
+
+
+T = TypeVar("T", bound=Compartment)
+
+
+class Compartments(list[T]):
+    r"""Comparments contains a set of compartment."""
+
+    names: SWCNames
+
+    def __init__(self, segments: Iterable[T]) -> None:
+        super().__init__(segments)
+        self.names = self[0].names if len(self) > 0 else get_names()
+
+    def id(self) -> npt.NDArray[np.int32]:  # pylint: disable=invalid-name
+        """Get the ids of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.id)
+
+    def type(self) -> npt.NDArray[np.int32]:
+        """Get the types of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.type)
+
+    def x(self) -> npt.NDArray[np.float32]:
+        """Get the x coordinates of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.x)
+
+    def y(self) -> npt.NDArray[np.float32]:
+        """Get the y coordinates of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.y)
+
+    def z(self) -> npt.NDArray[np.float32]:
+        """Get the z coordinates of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.z)
+
+    def r(self) -> npt.NDArray[np.float32]:
+        """Get the radius of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.r)
+
+    def pid(self) -> npt.NDArray[np.int32]:
+        """Get the ids of parent of shape (n_sample, 2)."""
+        return self.get_ndata(self.names.pid)
+
+    def xyz(self) -> npt.NDArray[np.float32]:
+        """Get the coordinates of shape (n_sample, 2, 3)."""
+        return np.stack([self.x(), self.y(), self.z()], axis=2)
+
+    def xyzr(self) -> npt.NDArray[np.float32]:
+        """Get the xyzr array of shape (n_sample, 2, 4)."""
+        return np.stack([self.x(), self.y(), self.z(), self.r()], axis=2)
+
+    def get_ndata(self, key: str) -> npt.NDArray:
+        """Get ndata of shape (n_sample, 2).
+
+        The order of axis 1 is (parent, current node).
+        """
+        return np.array([s.get_ndata(key) for s in self])
+
+
+# Aliases
+Segment = Compartment
+Segments = Compartments

swcgeom/core/node.py

@@ -0,0 +1,130 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Nueron node."""
+
+from collections.abc import Iterable
+from typing import Any, Generic
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import deprecated
+
+from swcgeom.core.swc import DictSWC, SWCTypeVar
+from swcgeom.core.swc_utils import SWCNames
+
+__all__ = ["Node"]
+
+
+class Node(Generic[SWCTypeVar]):
+    """Neural node."""
+
+    attach: SWCTypeVar
+    idx: int | np.integer
+    names: SWCNames
+
+    # fmt: off
+    @property
+    def id(self) -> int: return self[self.names.id]
+    @id.setter
+    def id(self, v: int): self[self.names.id] = v
+
+    @property
+    def type(self) -> int: return self[self.names.type]
+    @type.setter
+    def type(self, v: int): self[self.names.type] = v
+
+    @property
+    def x(self) -> float: return self[self.names.x]
+    @x.setter
+    def x(self, v: float): self[self.names.x] = v
+
+    @property
+    def y(self) -> float: return self[self.names.y]
+    @y.setter
+    def y(self, v: float): self[self.names.y] = v
+
+    @property
+    def z(self) -> float: return self[self.names.z]
+    @z.setter
+    def z(self, v: float): self[self.names.z] = v
+
+    @property
+    def r(self) -> float: return self[self.names.r]
+    @r.setter
+    def r(self, v: float): self[self.names.r] = v
+
+    @property
+    def pid(self) -> int: return self[self.names.pid]
+    @pid.setter
+    def pid(self, v: int): self[self.names.pid] = v
+    # fmt: on
+
+    def __init__(self, attach: SWCTypeVar, idx: int | np.integer) -> None:
+        super().__init__()
+        self.attach = attach
+        self.idx = idx
+        self.names = attach.names
+
+    def __getitem__(self, key: str) -> Any:
+        return self.attach.get_ndata(key)[self.idx]
+
+    def __setitem__(self, k: str, v: Any) -> None:
+        self.attach.get_ndata(k)[self.idx] = v
+
+    def __str__(self) -> str:
+        return self.format_swc()
+
+    def __repr__(self) -> str:
+        return self.format_swc()
+
+    def xyz(self) -> npt.NDArray[np.float32]:
+        """Get the `x`, `y`, `z` of node, an array of shape (3,)"""
+        return np.array([self.x, self.y, self.z], dtype=np.float32)
+
+    def xyzr(self) -> npt.NDArray[np.float32]:
+        """Get the `x`, `y`, `z`, `r` of node, an array of shape (4,)"""
+        return np.array([self.x, self.y, self.z, self.r], dtype=np.float32)
+
+    def keys(self) -> Iterable[str]:
+        return self.attach.keys()
+
+    def distance(self, b: "Node") -> float:
+        """Get the distance of two nodes."""
+        return np.linalg.norm(self.xyz() - b.xyz()).item()
+
+    def format_swc(self) -> str:
+        """Get the SWC format string."""
+        x, y, z, r = [f"{f:.4f}" for f in [self.x, self.y, self.z, self.r]]
+        items = [self.id, self.type, x, y, z, r, self.pid]
+        return " ".join(map(str, items))
+
+    def is_furcation(self) -> bool:
+        """Is furcation node."""
+        return np.count_nonzero(self.attach.pid() == self.id) > 1
+
+    @deprecated("Use is_furcation instead")
+    def is_bifurcation(self) -> bool:
+        """Is furcation node.
+
+        NOTE: Deprecated due to the wrong spelling of furcation. For now, it is just an
+        alias of `is_furcation` and raise a warning. It will be change to raise an
+        error in the future.
+        """
+        return self.is_furcation()
+
+    def is_tip(self) -> bool:
+        return self.id not in self.attach.pid()
+
+    def detach(self) -> "Node[DictSWC]":
+        """Detach from current attached object."""
+        attact = DictSWC(
+            **{k: np.array([self[k]]) for k in self.keys()},
+            source=self.attach.source,
+            names=self.names,
+        )
+        attact.ndata[self.names.id] = np.array([0])
+        attact.ndata[self.names.pid] = np.array([-1])
+        return Node(attact, 0)