swcgeom 0.18.1__py3-none-any.whl → 0.19.0__py3-none-any.whl

swcgeom/core/swc_utils/io.py

@@ -18,7 +18,7 @@
 import re
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Literal, Optional
+from typing import Literal
 
 import numpy as np
 import numpy.typing as npt
@@ -39,41 +39,33 @@ __all__ = ["read_swc", "to_swc"]
 
 def read_swc(
     swc_file: PathOrIO,
-    extra_cols: Optional[Iterable[str]] = None,
+    extra_cols: Iterable[str] | None = None,
     fix_roots: Literal["somas", "nearest", False] = False,
     sort_nodes: bool = False,
     reset_index: bool = True,
     *,
     encoding: Literal["detect"] | str = "utf-8",
-    names: Optional[SWCNames] = None,
+    names: SWCNames | None = None,
 ) -> tuple[pd.DataFrame, list[str]]:
     """Read swc file.
 
-    Parameters
-    ----------
-    swc_file : PathOrIO
-        Path of swc file, the id should be consecutively incremented.
-    extra_cols : Iterable[str], optional
-        Read more cols in swc file.
-    fix_roots : `somas`|`nearest`|False, default `False`
-        Fix multiple roots.
-    sort_nodes : bool, default `False`
-        Sort the indices of neuron tree, the index for parent are
-        always less than children.
-    reset_index : bool, default `True`
-        Reset node index to start with zero, DO NOT set to false if
-        you are not sure what will happend.
-    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.
-    names : SWCNames, optional
-
-    Returns
-    -------
-    df : ~pandas.DataFrame
-    comments : List of string
+    NOTE: the id should be consecutively incremented.
+
+    Args:
+        extra_cols: Read more cols in swc file.
+        fix_roots: Fix multiple roots.
+        sort_nodes: Sort the indices of neuron tree.
+            After sorting the nodes, the index for each parent are always less than
+            that of its children.
+        reset_index: Reset node index to start with zero.
+            DO NOT set to false if you are not sure what will happened.
+        encoding: The name of the encoding used to decode the file.
+            If is `detect`, we will try to detect the character encoding.
+
+    Returns:
+        df: ~pandas.DataFrame
+        comments: List of string
     """
-
     names = get_names(names)
     df, comments = parse_swc(
         swc_file, names=names, extra_cols=extra_cols, encoding=encoding
@@ -110,10 +102,10 @@ def read_swc(
 def to_swc(
     get_ndata: Callable[[str], npt.NDArray],
     *,
-    extra_cols: Optional[Iterable[str]] = None,
+    extra_cols: Iterable[str] | None = None,
     id_offset: int = 1,
-    comments: Optional[Iterable[str]] = None,
-    names: Optional[SWCNames] = None,
+    comments: Iterable[str] | None = None,
+    names: SWCNames | None = None,
 ) -> Iterable[str]:
     """Convert to swc format."""
 
@@ -156,21 +148,14 @@ def parse_swc(
 ) -> tuple[pd.DataFrame, list[str]]:
     """Parse swc file.
 
-    Parameters
-    ----------
-    fname : PathOrIO
-    names : SWCNames
-    extra_cols : List of str, optional
-    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.
-
-    Returns
-    -------
-    df : ~pandas.DataFrame
-    comments : List of string
-    """
+    Args:
+        encoding: The name of the encoding used to decode the file.
+            If is `detect`, we will try to detect the character encoding.
 
+    Returns:
+        df: ~pandas.DataFrame
+        comments: List of string
+    """
     # pylint: disable=too-many-locals
     extras = list(extra_cols) if extra_cols else []
 
@@ -208,7 +193,7 @@ def parse_swc(
                 if (match := re_swc.search(line)) is not None:
                     if flag and match.group(last_group):
                         warnings.warn(
-                            f"some fields are ignored in row {i+1} of `{fname}`"
+                            f"some fields are ignored in row {i + 1} of `{fname}`"
                         )
                         flag = False
 
@@ -219,10 +204,10 @@ def parse_swc(
                     if not comment.startswith(ignored_comment):
                         comments.append(comment)
                 elif not line.isspace():
-                    raise ValueError(f"invalid row {i+1} in `{fname}`")
+                    raise ValueError(f"invalid row {i + 1} in `{fname}`")
         except UnicodeDecodeError as e:
             raise ValueError(
-                f"decode failed, try to enable auto detect `encoding='detect'`"
+                "decode failed, try to enable auto detect `encoding='detect'`"
             ) from e
 
     df = pd.DataFrame.from_dict(dict(zip(keys, vals)))

swcgeom/core/swc_utils/normalizer.py

@@ -19,7 +19,7 @@ Methods ending with a underline imply an in-place transformation.
 """
 
 from collections.abc import Callable
-from typing import Literal, Optional
+from typing import Literal
 
 import numpy as np
 import numpy.typing as npt
@@ -44,7 +44,7 @@ def mark_roots_as_somas(
     df: pd.DataFrame,
     update_type: int | Literal[False] = 1,
     *,
-    names: Optional[SWCNames] = None,
+    names: SWCNames | None = None,
 ) -> pd.DataFrame:
     return _copy_and_apply(
         mark_roots_as_somas_, df, update_type=update_type, names=names
@@ -55,7 +55,7 @@ def mark_roots_as_somas_(
     df: pd.DataFrame,
     update_type: int | Literal[False] = 1,
     *,
-    names: Optional[SWCNames] = None,
+    names: SWCNames | None = None,
 ) -> None:
     """Merge multiple roots in swc.
 
@@ -72,14 +72,12 @@ def mark_roots_as_somas_(
 
 
 def link_roots_to_nearest(
-    df: pd.DataFrame, *, names: Optional[SWCNames] = None
+    df: pd.DataFrame, *, names: SWCNames | None = None
 ) -> pd.DataFrame:
     return _copy_and_apply(link_roots_to_nearest_, df, names=names)
 
 
-def link_roots_to_nearest_(
-    df: pd.DataFrame, *, names: Optional[SWCNames] = None
-) -> None:
+def link_roots_to_nearest_(df: pd.DataFrame, *, names: SWCNames | None = None) -> None:
     """Merge multiple roots in swc.
 
     The first root are reserved, and the others was.
@@ -94,29 +92,27 @@ def link_roots_to_nearest_(
         subtree = dsu == dsu[i]  # type: ignore
         dis = np.where(subtree, np.inf, dis)  # avoid link to same tree
         dsu = np.where(subtree, dsu[dis.argmin()], dsu)  # merge set
-        df.loc[i, names.pid] = df[names.id].iloc[dis.argmin()]  # type: ignore
+        df.loc[i, names.pid] = df[names.id].iloc[dis.argmin()]
 
 
-def sort_nodes(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> pd.DataFrame:
+def sort_nodes(df: pd.DataFrame, *, names: SWCNames | None = None) -> pd.DataFrame:
     """Sort the indices of neuron tree.
 
     The index for parent are always less than children.
 
-    See Also
-    --------
-    ~.core.swc_utils.checker.is_sorted
+    See Also:
+        ~.core.swc_utils.checker.is_sorted
     """
     return _copy_and_apply(sort_nodes_, df, names=names)
 
 
-def sort_nodes_(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> None:
+def sort_nodes_(df: pd.DataFrame, *, names: SWCNames | None = None) -> None:
     """Sort the indices of neuron tree.
 
     The index for parent are always less than children.
 
-    See Also
-    --------
-    ~.core.swc_utils.checker.is_sorted
+    See Also:
+        ~.core.swc_utils.checker.is_sorted
     """
     names = get_names(names)
     ids, pids = df[names.id].to_numpy(), df[names.pid].to_numpy()
@@ -130,11 +126,9 @@ def sort_nodes_(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> None:
 def sort_nodes_impl(topology: Topology) -> tuple[Topology, npt.NDArray[np.int32]]:
     """Sort the indices of neuron tree.
 
-    Returns
-    -------
-    new_topology : Topology
-    id_map : List of int
-        Map from new id to original id.
+    Returns:
+        new_topology: Topology
+        id_map: Map from new id to original id.
     """
     old_ids, old_pids = topology
     assert np.count_nonzero(old_pids == -1) == 1, "should be single root"
@@ -157,12 +151,12 @@ def sort_nodes_impl(topology: Topology) -> tuple[Topology, npt.NDArray[np.int32]
     return (new_ids, new_pids), indices
 
 
-def reset_index(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> pd.DataFrame:
+def reset_index(df: pd.DataFrame, *, names: SWCNames | None = None) -> pd.DataFrame:
     """Reset node index to start with zero."""
     return _copy_and_apply(reset_index_, df, names=names)
 
 
-def reset_index_(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> None:
+def reset_index_(df: pd.DataFrame, *, names: SWCNames | None = None) -> None:
     """Reset node index to start with zero."""
     names = get_names(names)
     roots = df[names.pid] == -1

swcgeom/core/swc_utils/subtree.py

@@ -15,10 +15,9 @@
 
 """Cut subtree.
 
-This module provides a series of low-level topological subtree methods,
-but in more cases, you can use the high-level methods provided in
-`tree_utils`, which wrap the methods in this module and provide a
-high-level API.
+This module provides a series of low-level topological subtree methods, but in more
+cases, you can use the high-level methods provided in `tree_utils`, which wrap the
+methods in this module and provide a high-level API.
 """
 
 from typing import cast
@@ -36,22 +35,17 @@ REMOVAL = -2  # A marker in utils, place in the ids to mark it removal
 def to_sub_topology(sub: Topology) -> tuple[Topology, npt.NDArray[np.int32]]:
     """Create sub tree from origin tree.
 
-    Mark the node to be removed, then use this method to get a child
-    structure.
+    Mark the node to be removed, then use this method to get a child structure.
 
-    Returns
-    -------
-    sub_topology : Topology
-    mapping : List of int
-        Map from new id to old id.
+    Returns:
+        sub_topology: Topology
+        mapping: Map from new id to old id.
 
-    See Also
-    --------
-    propagate_removal :
-        If the node you remove is not a leaf node, you need to use it
-        to mark all child nodes.
+    See Also:
+        propagate_removal:
+            If the node you remove is not a leaf node, you need to use it
+            to mark all child nodes.
     """
-
     sub_id = np.array(sub[0], dtype=np.int32)
     sub_pid = np.array(sub[1], dtype=np.int32)
 
@@ -69,9 +63,8 @@ def to_sub_topology(sub: Topology) -> tuple[Topology, npt.NDArray[np.int32]]:
 def propagate_removal(topology: Topology) -> Topology:
     """Mark all children when parent is marked as removed.
 
-    Returns
-    -------
-    new_topology : Topology
+    Returns:
+        new_topology: Topology
     """
 
     new_ids, pids = topology

swcgeom/core/tree.py

@@ -18,7 +18,7 @@
 import itertools
 import os
 from collections.abc import Callable, Iterable, Iterator
-from typing import Literal, Optional, TypeVar, Union, overload
+from typing import Literal, TypeVar, Union, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -36,7 +36,8 @@ from swcgeom.utils import PathOrIO, padding1d
 
 __all__ = ["Tree"]
 
-T, K = TypeVar("T"), TypeVar("K")
+T = TypeVar("T")
+K = TypeVar("K")
 
 
 class Tree(DictSWC):
@@ -67,13 +68,11 @@ class Tree(DictSWC):
             """The end-to-end straight-line distance to soma."""
             return self.distance(self.attach.soma())
 
-        def subtree(self, *, out_mapping: Optional[Mapping] = None) -> "Tree":
+        def subtree(self, *, out_mapping: Mapping | None = None) -> "Tree":
             """Get subtree from node.
 
-            Parameters
-            ----------
-            out_mapping : List of int or dict[int, int], optional
-                Map from new id to old id.
+            Args:
+                out_mapping: Map from new id to old id.
             """
 
             n_nodes, ndata, source, names = get_subtree_impl(
@@ -87,22 +86,27 @@ class Tree(DictSWC):
         def is_soma(self) -> bool:  # TODO: support multi soma, e.g. 3 points
             return self.type == self.attach.types.soma and self.is_root()
 
-        # fmt: off
         @overload
-        def traverse(self, *, enter: Callable[[Node, T | None], T], mode: Literal["dfs"] = ...) -> None: ...
+        def traverse(
+            self, *, enter: Callable[[Node, T | None], T], mode: Literal["dfs"] = ...
+        ) -> None: ...
         @overload
-        def traverse(self, *, leave: Callable[[Node, list[K]], K], mode: Literal["dfs"] = ...) -> K: ...
+        def traverse(
+            self, *, leave: Callable[[Node, list[K]], K], mode: Literal["dfs"] = ...
+        ) -> K: ...
         @overload
-        def traverse(self, *,
-            enter: Callable[[Node, T | None], T], leave: Callable[[Node, list[K]], K], mode: Literal["dfs"] = ...,
+        def traverse(
+            self,
+            *,
+            enter: Callable[[Node, T | None], T],
+            leave: Callable[[Node, list[K]], K],
+            mode: Literal["dfs"] = ...,
         ) -> K: ...
-        # fmt: on
         def traverse(self, **kwargs):  # type: ignore
             """Traverse from node.
 
-            See Also
-            --------
-            ~Tree.traverse
+            See Also:
+                ~Tree.traverse
             """
             return self.attach.traverse(root=self.idx, **kwargs)
 
@@ -125,8 +129,8 @@ class Tree(DictSWC):
         n_nodes: int,
         *,
         source: str = "",
-        comments: Optional[Iterable[str]] = None,
-        names: Optional[SWCNames] = None,
+        comments: Iterable[str] | None = None,
+        names: SWCNames | None = None,
         **kwargs: npt.NDArray,
     ) -> None:
         names = get_names(names)
@@ -162,14 +166,12 @@ class Tree(DictSWC):
         n_nodes, n_edges = self.number_of_nodes(), self.number_of_edges()
         return f"Neuron Tree with {n_nodes} nodes and {n_edges} edges"
 
-    # fmt:off
     @overload
     def __getitem__(self, key: slice) -> list[Node]: ...
     @overload
     def __getitem__(self, key: int) -> Node: ...
     @overload
     def __getitem__(self, key: str) -> npt.NDArray: ...
-    # fmt:on
     def __getitem__(self, key):
         if isinstance(key, slice):
             return [self.node(i) for i in range(*key.indices(len(self)))]
@@ -218,11 +220,10 @@ class Tree(DictSWC):
     def get_bifurcations(self) -> list[Node]:
         """Get all node of furcations.
 
-        Notes
-        -----
-        Deprecated due to the wrong spelling of furcation. For now, it
-        is just an alias of `get_furcations` and raise a warning. It
-        will be change to raise an error in the future.
+        .. deprecated:: 0.17.2
+            Deprecated due to the wrong spelling of furcation. For now, it is just an
+            alias of `get_furcations` and raise a warning. It will be change to raise
+            an error in the future.
         """
         return self.get_furcations()
 
@@ -291,29 +292,40 @@ class Tree(DictSWC):
         children = self.soma(type_check).children()
         return (n.subtree() for n in children if n.type in types)
 
-    # fmt: off
     @overload
-    def traverse(self, *,
-                 enter: Callable[[Node, T | None], T],
-                 root: int | np.integer = ..., mode: Literal["dfs"] = ...) -> None: ...
+    def traverse(
+        self,
+        *,
+        enter: Callable[[Node, T | None], T],
+        root: int | np.integer = ...,
+        mode: Literal["dfs"] = ...,
+    ) -> None: ...
     @overload
-    def traverse(self, *,
-                 enter: Optional[Callable[[Node, T | None], T]] = ...,
-                 leave: Callable[[Node, list[K]], K],
-                 root: int | np.integer = ..., mode: Literal["dfs"] = ...) -> K: ...
-    # fmt: on
-
+    def traverse(
+        self,
+        *,
+        leave: Callable[[Node, list[K]], K],
+        root: int | np.integer = ...,
+        mode: Literal["dfs"] = ...,
+    ) -> K: ...
+    @overload
+    def traverse(
+        self,
+        *,
+        enter: Callable[[Node, T | None], T],
+        leave: Callable[[Node, list[K]], K],
+        root: int | np.integer = ...,
+        mode: Literal["dfs"] = ...,
+    ) -> K: ...
     def traverse(self, *, enter=None, leave=None, **kwargs):
         """Traverse nodes.
 
-        Parameters
-        ----------
-        enter : (n: Node, parent: T | None) => T, optional
-        leave : (n: Node, children: list[T]) => T, optional
+        Args:
+            enter: (n: Node, parent: T | None) => T
+            leave: (n: Node, children: list[T]) => T
 
-        See Also
-        --------
-        ~swc_utils.traverse
+        See Also:
+            ~swc_utils.traverse
         """
 
         def wrap(fn) -> Callable | None:
@@ -338,8 +350,8 @@ class Tree(DictSWC):
         df: pd.DataFrame,
         source: str = "",
         *,
-        comments: Optional[Iterable[str]] = None,
-        names: Optional[SWCNames] = None,
+        comments: Iterable[str] | None = None,
+        names: SWCNames | None = None,
     ) -> "Tree":
         """Read neuron tree from data frame."""
         names = get_names(names)
@@ -356,9 +368,8 @@ class Tree(DictSWC):
     def from_swc(cls, swc_file: PathOrIO, **kwargs) -> "Tree":
         """Read neuron tree from swc file.
 
-        See Also
-        --------
-        ~swcgeom.core.swc_utils.read_swc
+        See Also:
+            ~swcgeom.core.swc_utils.read_swc
         """
 
         try:
@@ -371,13 +382,12 @@ class Tree(DictSWC):
 
     @classmethod
     def from_eswc(
-        cls, swc_file: str, extra_cols: Optional[list[str]] = None, **kwargs
+        cls, swc_file: str, extra_cols: list[str] | None = None, **kwargs
     ) -> "Tree":
         """Read neuron tree from eswc file.
 
-        See Also
-        --------
-        ~swcgeom.Tree.from_swc
+        See Also:
+            ~swcgeom.Tree.from_swc
         """
         extra_cols = extra_cols or []
         extra_cols.extend(k for k, t in eswc_cols)

swcgeom/core/tree_utils.py

@@ -17,7 +17,7 @@
 
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Optional, TypeVar, overload
+from typing import TypeVar, overload
 
 import numpy as np
 from typing_extensions import deprecated
@@ -46,7 +46,8 @@ __all__ = [
     "cat_tree",
 ]
 
-T, K = TypeVar("T"), TypeVar("K")
+T = TypeVar("T")
+K = TypeVar("K")
 EPS = 1e-5
 
 
@@ -58,19 +59,20 @@ def is_binary_tree(tree: Tree, exclude_soma: bool = True) -> bool:
 def sort_tree(tree: Tree) -> Tree:
     """Sort the indices of neuron tree.
 
-    See Also
-    --------
-    ~.core.swc_utils.sort_nodes
+    See Also:
+        ~.core.swc_utils.sort_nodes
     """
     return _sort_tree(tree.copy())
 
 
-# fmt:off
 @overload
-def cut_tree(tree: Tree, *, enter: Callable[[Tree.Node, T | None], tuple[T, bool]]) -> Tree: ...
+def cut_tree(
+    tree: Tree, *, enter: Callable[[Tree.Node, T | None], tuple[T, bool]]
+) -> Tree: ...
 @overload
-def cut_tree(tree: Tree, *, leave: Callable[[Tree.Node, list[K]], tuple[K, bool]]) -> Tree: ...
-# fmt:on
+def cut_tree(
+    tree: Tree, *, leave: Callable[[Tree.Node, list[K]], tuple[K, bool]]
+) -> Tree: ...
 def cut_tree(tree: Tree, *, enter=None, leave=None):
     """Traverse and cut the tree.
 
@@ -123,9 +125,8 @@ def to_sub_tree(swc_like: SWCLike, sub: Topology) -> tuple[Tree, dict[int, int]]
         Use :meth:`to_subtree` instead.
 
     Returns
-    -------
-    tree : Tree
-    id_map : dict[int, int]
+        tree: Tree
+        id_map: dict[int, int]
     """
 
     sub = propagate_removal(sub)
@@ -147,17 +148,14 @@ def to_subtree(
     swc_like: SWCLike,
     removals: Iterable[int],
     *,
-    out_mapping: Optional[Mapping] = None,
+    out_mapping: Mapping | None = None,
 ) -> Tree:
     """Create subtree from origin tree.
 
-    Parameters
-    ----------
-    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.
+    Args:
+        swc_like: SWCLike
+        removals: A list of id of nodes to be removed.
+        out_mapping: Map new id to old id.
     """
 
     new_ids = swc_like.id().copy()
@@ -172,17 +170,14 @@ def to_subtree(
 
 
 def get_subtree(
-    swc_like: SWCLike, n: int, *, out_mapping: Optional[Mapping] = None
+    swc_like: SWCLike, n: int, *, out_mapping: Mapping | None = None
 ) -> Tree:
     """Get subtree rooted at n.
 
-    Parameters
-    ----------
-    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.
+    Args:
+        swc_like: SWCLike
+        n: Id of the root of the subtree.
+        out_mapping: Map new id to old id.
     """
 
     n_nodes, ndata, source, names = get_subtree_impl(
@@ -194,14 +189,10 @@ def get_subtree(
 def redirect_tree(tree: Tree, new_root: int, sort: bool = True) -> Tree:
     """Set root to new point and redirect tree graph.
 
-    Parameter
-    ---------
-    tree : Tree
-        The tree.
-    new_root : int
-        The id of new root.
-    sort : bool, default `True`
-        If true, sort indices of nodes after redirect.
+    Args:
+        tree: The tree.
+        new_root: The id of new root.
+        sort: If true, sort indices of nodes after redirect.
     """
 
     tree = tree.copy()
@@ -227,22 +218,18 @@ def cat_tree(  # pylint: disable=too-many-arguments
     node2: int = 0,
     *,
     translate: bool = True,
-    names: Optional[SWCNames] = None,
-    no_move: Optional[bool] = None,  # legacy
+    names: SWCNames | None = None,
+    no_move: bool | None = None,  # legacy
 ) -> Tree:
     """Concatenates the second tree onto the first one.
 
-    Paramters
-    ---------
-    tree1 : Tree
-    tree2 : Tree
-    node1 : int, default `0`
-        The node id of the tree to be connected.
-    node2 : int, default `0`
-        The node id of the connection point.
-    translate : bool, default `True`
-        Wheather to translate node_2 to node_1. If False, add link
-        between node_1 and node_2 without translate.
+    Args:
+        tree1: Tree
+        tree2: Tree
+        node1: The node id of the tree to be connected.
+        node2: The node id of the connection point.
+        translate: Weather to translate node_2 to node_1.
+            If False, add link between node_1 and node_2 without translate.
     """
     if no_move is not None:
         warnings.warn(