swcgeom 0.14.0__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.14.0'
-__version_tuple__ = version_tuple = (0, 14, 0)
+__version__ = version = '0.15.0'
+__version_tuple__ = version_tuple = (0, 15, 0)

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

swcgeom/images/folder.py

@@ -2,14 +2,25 @@
 
 import os
 import re
+import warnings
 from abc import ABC, abstractmethod
-from typing import Callable, Generic, Iterable, List, Literal, Optional, Tuple, TypeVar
+from typing import (
+    Callable,
+    Generic,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    overload,
+)
 
 import numpy as np
 import numpy.typing as npt
 from typing_extensions import Self
 
-from swcgeom.images.io import read_imgs
+from swcgeom.images.io import ScalarType, read_imgs
 from swcgeom.transforms import Identity, Transform
 
 __all__ = [
@@ -21,20 +32,23 @@ __all__ = [
 T = TypeVar("T")
 
 
-class ImageStackFolderBase(Generic[T], ABC):
+class ImageStackFolderBase(Generic[ScalarType, T], ABC):
     """Image stack folder base."""
 
     files: List[str]
-    transform: Transform[npt.NDArray[np.float32], T]
+    transform: Transform[npt.NDArray[ScalarType], T]
 
-    def __init__(
-        self,
-        files: Iterable[str],
-        *,
-        transform: Optional[Transform[npt.NDArray[np.float32], T]] = None,
-    ) -> None:
+    # fmt: off
+    @overload
+    def __init__(self, files: Iterable[str], *, dtype: None = ..., transform: Optional[Transform[npt.NDArray[np.float32], T]] = None) -> None: ...
+    @overload
+    def __init__(self, files: Iterable[str], *, dtype: ScalarType, transform: Optional[Transform[npt.NDArray[ScalarType], T]] = None) -> None: ...
+    # fmt: on
+
+    def __init__(self, files: Iterable[str], *, dtype=None, transform=None) -> None:
         super().__init__()
         self.files = list(files)
+        self.dtype = dtype or np.float32
         self.transform = transform or Identity()  # type: ignore
 
     @abstractmethod
@@ -45,13 +59,12 @@ class ImageStackFolderBase(Generic[T], ABC):
         return len(self.files)
 
     def _get(self, fname: str) -> T:
-        imgs = self.read_imgs(fname)
+        imgs = self._read(fname)
         imgs = self.transform(imgs)
         return imgs
 
-    @staticmethod
-    def read_imgs(fname: str) -> npt.NDArray[np.float32]:
-        return read_imgs(fname).get_full()
+    def _read(self, fname: str) -> npt.NDArray[ScalarType]:
+        return read_imgs(fname, dtype=self.dtype).get_full()  # type: ignore
 
     @staticmethod
     def scan(root: str, *, pattern: Optional[str] = None) -> List[str]:
@@ -63,8 +76,20 @@ class ImageStackFolderBase(Generic[T], ABC):
 
         return fs
 
+    @staticmethod
+    def read_imgs(fname: str) -> npt.NDArray[np.float32]:
+        warnings.warn(
+            "`ImageStackFolderBase.read_imgs` serves as a "
+            "straightforward wrapper for `~swcgeom.images.io.read_imgs(fname).get_full()`. "
+            "However, as it is not utilized within our internal "
+            "processes, it is scheduled for removal in the "
+            "forthcoming version.",
+            DeprecationWarning,
+        )
+        return read_imgs(fname).get_full()
+
 
-class ImageStackFolder(Generic[T], ImageStackFolderBase[T]):
+class ImageStackFolder(ImageStackFolderBase[ScalarType, T]):
     """Image stack folder."""
 
     def __getitem__(self, idx: int, /) -> T:
@@ -84,7 +109,7 @@ class ImageStackFolder(Generic[T], ImageStackFolderBase[T]):
         return cls(cls.scan(root, pattern=pattern), **kwargs)
 
 
-class LabeledImageStackFolder(Generic[T], ImageStackFolderBase[T]):
+class LabeledImageStackFolder(ImageStackFolderBase[ScalarType, T]):
     """Image stack folder with label."""
 
     labels: List[int]
@@ -115,7 +140,7 @@ class LabeledImageStackFolder(Generic[T], ImageStackFolderBase[T]):
         return cls(files, labels, **kwargs)
 
 
-class PathImageStackFolder(Generic[T], ImageStackFolder[T]):
+class PathImageStackFolder(ImageStackFolder[ScalarType, T]):
     """Image stack folder with relpath."""
 
     root: str