swcgeom 0.13.2__py3-none-any.whl → 0.15.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.13.2'
-__version_tuple__ = version_tuple = (0, 13, 2)
+__version__ = version = '0.15.0'
+__version_tuple__ = version_tuple = (0, 15, 0)

swcgeom/analysis/volume.py

@@ -1,22 +1,44 @@
 """Analysis of volume of a SWC tree."""
 
-from typing import List
+from typing import Dict, List, Literal
 
 import numpy as np
+from sdflit import ColoredMaterial, ObjectsScene, SDFObject, UniformSampler
 
 from swcgeom.core import Tree
-from swcgeom.utils import GeomFrustumCone, GeomSphere
+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):
+
+def get_volume(
+    tree: Tree,
+    *,
+    method: Literal["frustum_cone"] = "frustum_cone",
+    accuracy: int | ACCURACY_LEVEL = "middle",
+) -> float:
     """Get the volume of the tree.
 
     Parameters
     ----------
     tree : Tree
-        SWC tree.
+        Neuronal tree.
+    method : {"frustum_cone"}, optional
+        Method for volume calculation.
+    accuracy : int or {"low", "middle", "high"}, optional
+        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
     -------
@@ -38,26 +60,58 @@ def get_volume(tree: Tree):
 
     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.
+
+    Parameters
+    ----------
+    tree : Tree
+        Neuronal tree.
+    accuracy : int
+        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[GeomSphere]) -> GeomSphere:
-        sphere = GeomSphere(n.xyz(), n.r)
-        cones = [GeomFrustumCone(n.xyz(), n.r, c.center, c.radius) for c in children]
+    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()
-        v += sum(fc.get_volume() for fc in cones)
-        v -= sum(sphere.get_intersect_volume(fc) for fc in cones)
-        v -= sum(s.get_intersect_volume(fc) for s, fc in zip(children, cones))
-        v += sum(s.get_intersect_volume(sphere) for s in children)
-
-        # 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
-        # )
+        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
@@ -65,3 +119,34 @@ def get_volume(tree: Tree):
 
     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/branch_tree.py

@@ -5,7 +5,6 @@ from typing import Dict, List
 
 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
@@ -31,7 +30,7 @@ class BranchTree(Tree):
         return self.branches[idx]
 
     @classmethod
-    def from_tree(cls, tree: Tree) -> Self:
+    def from_tree(cls, tree: Tree) -> "BranchTree":
         """Generating a branch tree from tree."""
 
         branches = tree.get_branches()
@@ -56,6 +55,6 @@ class BranchTree(Tree):
         return branch_tree
 
     @classmethod
-    def from_data_frame(cls, df: pd.DataFrame, *args, **kwargs) -> Self:
+    def from_data_frame(cls, df: pd.DataFrame, *args, **kwargs) -> "BranchTree":
         tree = super().from_data_frame(df, *args, **kwargs)
         return cls.from_tree(tree)

swcgeom/core/population.py

@@ -17,7 +17,6 @@ from typing import (
 
 import numpy as np
 import numpy.typing as npt
-from typing_extensions import Self
 
 from swcgeom.core.swc import eswc_cols
 from swcgeom.core.tree import Tree
@@ -155,7 +154,7 @@ class Population:
         return f"Neuron population in '{self.root}'"
 
     @classmethod
-    def from_swc(cls, root: str, ext: str = ".swc", **kwargs) -> Self:
+    def from_swc(cls, root: str, ext: str = ".swc", **kwargs) -> "Population":
         if not os.path.exists(root):
             raise FileNotFoundError(
                 f"the root does not refers to an existing directory: {root}"
@@ -171,7 +170,7 @@ class Population:
         ext: str = ".eswc",
         extra_cols: Optional[Iterable[str]] = None,
         **kwargs,
-    ) -> Self:
+    ) -> "Population":
         extra_cols = list(extra_cols) if extra_cols is not None else []
         extra_cols.extend(k for k, t in eswc_cols)
         return cls.from_swc(root, ext, extra_cols=extra_cols, **kwargs)
@@ -244,7 +243,7 @@ class Populations:
         check_same: bool = False,
         labels: Optional[Iterable[str]] = None,
         **kwargs,
-    ) -> Self:
+    ) -> "Populations":
         """Get population from dirs.
 
         Parameters
@@ -268,7 +267,7 @@ class Populations:
 
             fs = [inter for _ in roots]
         elif check_same:
-            assert reduce(lambda a, b: a == b, fs), "not the same among populations"
+            assert [fs[0] == a for a in fs[1:]], "not the same among populations"
 
         populations = [
             Population(
@@ -276,7 +275,7 @@ class Populations:
             )
             for i, d in enumerate(roots)
         ]
-        return cls(populations, labels=labels)
+        return Populations(populations, labels=labels)
 
     @classmethod
     def from_eswc(
@@ -286,7 +285,7 @@ class Populations:
         *,
         ext: str = ".eswc",
         **kwargs,
-    ) -> Self:
+    ) -> "Populations":
         extra_cols = list(extra_cols) if extra_cols is not None else []
         extra_cols.extend(k for k, t in eswc_cols)
         return cls.from_swc(roots, extra_cols=extra_cols, ext=ext, **kwargs)

swcgeom/core/swc_utils/assembler.py

@@ -1,155 +1,26 @@
-"""Assemble lines to swc."""
+"""Assemble lines to swc.
 
-import warnings
-from copy import copy
-from typing import Iterable, List, Optional, Tuple
+Notes
+-----
+This module is deprecated, please use `~.transforms.LinesToTree`
+instead.
+"""
 
-import numpy as np
-import pandas as pd
-
-from swcgeom.core.swc_utils.base import SWCNames, get_names
-from swcgeom.core.swc_utils.normalizer import link_roots_to_nearest_, sort_nodes_
 
 __all__ = ["assemble_lines", "try_assemble_lines"]
 
 
-def assemble_lines(*args, **kwargs) -> pd.DataFrame:
-    """Assemble lines to a tree.
-
-    Assemble all the lines into a set of subtrees, and then connect
-    them.
-
-    Parameters
-    ----------
-    lines : List of ~pd.DataFrame
-        An array of tables containing a line, columns should follwing
-        the swc.
-    **kwargs
-        Forwarding to `try_assemble_lines`
-
-    Returns
-    -------
-    tree : ~pd.DataFrame
-
-    See Also
-    --------
-    ~swcgeom.core.swc_utils.try_assemble_lines
-    """
-    warnings.warn(
+def assemble_lines(*args, **kwargs):
+    raise DeprecationWarning(
         "`assemble_lines` has been replaced by "
         "`~.transforms.LinesToTree` because it can be easy assemble "
-        "with other tansformations, and this will be removed in next "
-        "version.",
-        DeprecationWarning,
+        "with other tansformations.",
     )
-    return assemble_lines_impl(*args, **kwargs)
-
 
-def try_assemble_lines(*args, **kwargs) -> Tuple[pd.DataFrame, List[pd.DataFrame]]:
-    """Trying assemble lines to a tree.
 
-    Treat the first line as a tree, find a line whose shortest distance
-    between the tree and the line is less than threshold, merge it into
-    the tree, repeat until there are no line to merge, return tree and
-    the remaining lines.
-
-    Parameters
-    ----------
-    lines : List of ~pd.DataFrame
-        An array of tables containing a line, columns should follwing
-        the swc.
-    undirected : bool, default `True`
-        Both ends of a line can be considered connection point. If
-        `False`, only the starting point.
-    thre : float, default `0.2`
-        Connection threshold.
-    id_offset : int, default `0`
-        The offset of the line node id.
-    sort_nodes : bool, default `True`
-        sort nodes of subtree
-    names : SWCNames, optional
-
-    Returns
-    -------
-    tree : ~pd.DataFrame
-    remaining_lines : List of ~pd.DataFrame
-    """
-    warnings.warn(
+def try_assemble_lines(*args, **kwargs):
+    raise DeprecationWarning(
         "`try_assemble_lines` has been replaced by "
         "`~.transforms.LinesToTree` because it can be easy assemble "
-        "with other tansformations, and this will be removed in next "
-        "version.",
-        DeprecationWarning,
+        "with other tansformations.",
     )
-    return try_assemble_lines_impl(*args, **kwargs)
-
-
-# TODO: move the following codes to `transforms` module
-
-EPS = 1e-5
-
-
-def assemble_lines_impl(lines: Iterable[pd.DataFrame], **kwargs) -> pd.DataFrame:
-    tree, lines = try_assemble_lines_impl(lines, sort_nodes=False, **kwargs)
-    while len(lines) > 0:
-        t, lines = try_assemble_lines_impl(
-            lines, id_offset=len(tree), sort_nodes=False, **kwargs
-        )
-        tree = pd.concat([tree, t])
-
-    tree = tree.reset_index()
-    link_roots_to_nearest_(tree)
-    sort_nodes_(tree)
-    return tree
-
-
-def try_assemble_lines_impl(  # pylint: disable=too-many-arguments
-    lines: Iterable[pd.DataFrame],
-    undirected: bool = True,
-    thre: float = 0.2,
-    id_offset: int = 0,
-    sort_nodes: bool = True,
-    *,
-    names: Optional[SWCNames] = None,
-) -> Tuple[pd.DataFrame, List[pd.DataFrame]]:
-    names = get_names(names)
-    lines = copy(list(lines))
-
-    tree = lines[0]
-    tree[names.id] = id_offset + np.arange(len(tree))
-    tree[names.pid] = tree[names.id] - 1
-    tree.at[0, names.pid] = -1
-    del lines[0]
-
-    while True:
-        for i, line in enumerate(lines):
-            for p in [0, -1] if undirected else [0]:
-                xyz = [names.x, names.y, names.z]
-                vs = tree[xyz] - line.iloc[p][xyz]
-                dis = np.linalg.norm(vs, axis=1)
-                ind = np.argmin(dis)
-                if dis[ind] > thre:
-                    continue
-
-                if dis[ind] < EPS:
-                    line = line.drop((p + len(line)) % len(line)).reset_index(drop=True)
-
-                line[names.id] = id_offset + len(tree) + np.arange(len(line))
-                line[names.pid] = line[names.id] + (-1 if p == 0 else 1)
-                line.at[(p + len(line)) % len(line), names.pid] = tree.iloc[ind][
-                    names.id
-                ]
-                tree = pd.concat([tree, line])
-                del lines[i]
-                break
-            else:
-                continue
-
-            break
-        else:
-            break
-
-    if sort_nodes:
-        sort_nodes_(tree)
-
-    return tree, lines

swcgeom/core/swc_utils/subtree.py

@@ -27,8 +27,8 @@ def to_sub_topology(sub: Topology) -> Tuple[Topology, npt.NDArray[np.int32]]:
     Returns
     -------
     sub_topology : Topology
-    id_map : List of int
-        Map from new id to original id.
+    mapping : List of int
+        Map from new id to old id.
 
     See Also
     --------

swcgeom/core/tree.py

@@ -20,7 +20,6 @@ from typing import (
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
-from typing_extensions import Self
 
 from swcgeom.core.branch import Branch
 from swcgeom.core.node import Node
@@ -28,7 +27,7 @@ from swcgeom.core.path import Path
 from swcgeom.core.segment import Segment, Segments
 from swcgeom.core.swc import DictSWC, eswc_cols
 from swcgeom.core.swc_utils import SWCNames, get_names, read_swc, traverse
-from swcgeom.core.tree_utils_impl import get_subtree_impl
+from swcgeom.core.tree_utils_impl import Mapping, get_subtree_impl
 from swcgeom.utils import PathOrIO, padding1d
 
 __all__ = ["Tree"]
@@ -73,9 +72,18 @@ class Tree(DictSWC):
             """The end-to-end straight-line distance to soma."""
             return self.distance(self.attach.soma())
 
-        def subtree(self) -> "Tree":
-            """Get subtree from node."""
-            n_nodes, ndata, source, names = get_subtree_impl(self.attach, self.id)
+        def subtree(self, *, out_mapping: Optional[Mapping] = None) -> "Tree":
+            """Get subtree from node.
+
+            Parameters
+            ----------
+            out_mapping : List of int or Dict[int, int], optional
+                Map from new id to old id.
+            """
+
+            n_nodes, ndata, source, names = get_subtree_impl(
+                self.attach, self.id, out_mapping=out_mapping
+            )
             return Tree(n_nodes, **ndata, source=source, names=names)
 
         def is_root(self) -> bool:
@@ -260,11 +268,11 @@ class Tree(DictSWC):
         paths = self.traverse(enter=assign_path, leave=collect_path)
         return [self.Path(self, idx) for idx in paths]
 
-    def get_neurites(self, type_check: bool = True) -> Iterable[Self]:
+    def get_neurites(self, type_check: bool = True) -> Iterable["Tree"]:
         """Get neurites from soma."""
         return (n.subtree() for n in self.soma(type_check).children())
 
-    def get_dendrites(self, type_check: bool = True) -> Iterable[Self]:
+    def get_dendrites(self, type_check: bool = True) -> Iterable["Tree"]:
         """Get dendrites."""
         types = [self.types.apical_dendrite, self.types.basal_dendrite]
         children = self.soma(type_check).children()
@@ -312,15 +320,14 @@ class Tree(DictSWC):
         """Get length of tree."""
         return sum(s.length() for s in self.get_segments())
 
-    @classmethod
+    @staticmethod
     def from_data_frame(
-        cls,
         df: pd.DataFrame,
         source: str = "",
         *,
         comments: Optional[Iterable[str]] = None,
         names: Optional[SWCNames] = None,
-    ) -> Self:
+    ) -> "Tree":
         """Read neuron tree from data frame."""
         names = get_names(names)
         tree = Tree(
@@ -333,7 +340,7 @@ class Tree(DictSWC):
         return tree
 
     @classmethod
-    def from_swc(cls, swc_file: PathOrIO, **kwargs) -> Self:
+    def from_swc(cls, swc_file: PathOrIO, **kwargs) -> "Tree":
         """Read neuron tree from swc file.
 
         See Also
@@ -352,7 +359,7 @@ class Tree(DictSWC):
     @classmethod
     def from_eswc(
         cls, swc_file: str, extra_cols: Optional[List[str]] = None, **kwargs
-    ) -> Self:
+    ) -> "Tree":
         """Read neuron tree from eswc file.
 
         See Also

swcgeom/core/tree_utils.py

@@ -17,7 +17,7 @@ from swcgeom.core.swc_utils import (
     to_sub_topology,
 )
 from swcgeom.core.tree import Tree
-from swcgeom.core.tree_utils_impl import get_subtree_impl, to_subtree_impl
+from swcgeom.core.tree_utils_impl import Mapping, get_subtree_impl, to_subtree_impl
 
 __all__ = [
     "sort_tree",
@@ -128,7 +128,12 @@ def to_sub_tree(swc_like: SWCLike, sub: Topology) -> Tuple[Tree, Dict[int, int]]
     return subtree, id_map
 
 
-def to_subtree(swc_like: SWCLike, removals: Iterable[int]) -> Tree:
+def to_subtree(
+    swc_like: SWCLike,
+    removals: Iterable[int],
+    *,
+    out_mapping: Optional[Mapping] = None,
+) -> Tree:
     """Create subtree from origin tree.
 
     Parameters
@@ -136,17 +141,24 @@ def to_subtree(swc_like: SWCLike, removals: Iterable[int]) -> Tree:
     swc_like : SWCLike
     removals : List of int
         A list of id of nodes to be removed.
+    out_mapping: List of int or Dict[int, int], optional
+        Map new id to old id.
     """
+
     new_ids = swc_like.id().copy()
     for i in removals:
         new_ids[i] = REMOVAL
 
     sub = propagate_removal((new_ids, swc_like.pid()))
-    n_nodes, ndata, source, names = to_subtree_impl(swc_like, sub)
+    n_nodes, ndata, source, names = to_subtree_impl(
+        swc_like, sub, out_mapping=out_mapping
+    )
     return Tree(n_nodes, **ndata, source=source, names=names)
 
 
-def get_subtree(swc_like: SWCLike, n: int) -> Tree:
+def get_subtree(
+    swc_like: SWCLike, n: int, *, out_mapping: Optional[Mapping] = None
+) -> Tree:
     """Get subtree rooted at n.
 
     Parameters
@@ -154,8 +166,13 @@ def get_subtree(swc_like: SWCLike, n: int) -> Tree:
     swc_like : SWCLike
     n : int
         Id of the root of the subtree.
+    out_mapping: List of int or Dict[int, int], optional
+        Map new id to old id.
     """
-    n_nodes, ndata, source, names = get_subtree_impl(swc_like, n)
+
+    n_nodes, ndata, source, names = get_subtree_impl(
+        swc_like, n, out_mapping=out_mapping
+    )
     return Tree(n_nodes, **ndata, source=source, names=names)
 
 
@@ -171,6 +188,7 @@ def redirect_tree(tree: Tree, new_root: int, sort: bool = True) -> Tree:
     sort : bool, default `True`
         If true, sort indices of nodes after redirect.
     """
+
     tree = tree.copy()
     path = [tree.node(new_root)]
     while (p := path[-1].parent()) is not None:

swcgeom/core/tree_utils_impl.py

@@ -5,7 +5,7 @@ Notes
 Do not import `Tree` and keep this file minimized.
 """
 
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -15,10 +15,13 @@ from swcgeom.core.swc_utils import Topology, to_sub_topology, traverse
 
 __all__ = ["get_subtree_impl", "to_subtree_impl"]
 
+Mapping = Dict[int, int] | List[int]
 TreeArgs = Tuple[int, Dict[str, npt.NDArray[Any]], str, SWCNames]
 
 
-def get_subtree_impl(swc_like: SWCLike, n: int) -> TreeArgs:
+def get_subtree_impl(
+    swc_like: SWCLike, n: int, *, out_mapping: Optional[Mapping] = None
+) -> TreeArgs:
     ids = []
     topo = (swc_like.id(), swc_like.pid())
     traverse(topo, enter=lambda n, _: ids.append(n), root=n)
@@ -26,14 +29,27 @@ def get_subtree_impl(swc_like: SWCLike, n: int) -> TreeArgs:
     sub_ids = np.array(ids, dtype=np.int32)
     sub_pid = swc_like.pid()[sub_ids]
     sub_pid[0] = -1
-    return to_subtree_impl(swc_like, (sub_ids, sub_pid))
+    return to_subtree_impl(swc_like, (sub_ids, sub_pid), out_mapping=out_mapping)
 
 
-def to_subtree_impl(swc_like: SWCLike, sub: Topology) -> TreeArgs:
-    (new_id, new_pid), id_map = to_sub_topology(sub)
+def to_subtree_impl(
+    swc_like: SWCLike,
+    sub: Topology,
+    *,
+    out_mapping: Optional[Mapping] = None,
+) -> TreeArgs:
+    (new_id, new_pid), mapping = to_sub_topology(sub)
 
     n_nodes = new_id.shape[0]
-    ndata = {k: swc_like.get_ndata(k)[id_map].copy() for k in swc_like.keys()}
+    ndata = {k: swc_like.get_ndata(k)[mapping].copy() for k in swc_like.keys()}
     ndata.update(id=new_id, pid=new_pid)
 
+    if isinstance(out_mapping, list):
+        out_mapping.clear()
+        out_mapping.extend(mapping)
+    elif isinstance(out_mapping, dict):
+        out_mapping.clear()
+        for new_id, old_id in enumerate(mapping):
+            out_mapping[new_id] = old_id  # returning a dict may leads to bad perf
+
     return n_nodes, ndata, swc_like.source, swc_like.names