swcgeom 0.19.4__cp311-cp311-win_amd64.whl

swcgeom/core/swc_utils/io.py

@@ -0,0 +1,204 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Read and write swc format."""
+
+import re
+import warnings
+from collections.abc import Callable, Iterable
+from typing import Literal
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+from swcgeom.core.swc_utils.base import SWCNames, get_names
+from swcgeom.core.swc_utils.checker import is_single_root
+from swcgeom.core.swc_utils.normalizer import (
+    link_roots_to_nearest_,
+    mark_roots_as_somas_,
+    reset_index_,
+    sort_nodes_,
+)
+from swcgeom.utils import FileReader, PathOrIO
+
+__all__ = ["read_swc", "to_swc"]
+
+
+def read_swc(
+    swc_file: PathOrIO,
+    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: SWCNames | None = None,
+) -> tuple[pd.DataFrame, list[str]]:
+    """Read swc file.
+
+    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
+    )
+
+    # fix swc
+    if fix_roots is not False and np.count_nonzero(df[names.pid] == -1) > 1:
+        match fix_roots:
+            case "somas":
+                mark_roots_as_somas_(df)
+            case "nearest":
+                link_roots_to_nearest_(df)
+            case _:
+                raise ValueError(f"unknown fix type `{fix_roots}`")
+
+    if sort_nodes:
+        sort_nodes_(df)
+    elif reset_index:
+        reset_index_(df)
+
+    # check swc
+    if not is_single_root(df, names=names):
+        warnings.warn(f"not a simple tree in `{swc_file}`")
+
+    if (df[names.pid] == -1).argmax() != 0:
+        warnings.warn(f"root is not the first node in `{swc_file}`")
+
+    if (df[names.r] <= 0).any():
+        warnings.warn(f"non-positive radius in `{swc_file}`")
+
+    return df, comments
+
+
+def to_swc(
+    get_ndata: Callable[[str], npt.NDArray],
+    *,
+    extra_cols: Iterable[str] | None = None,
+    id_offset: int = 1,
+    comments: Iterable[str] | None = None,
+    names: SWCNames | None = None,
+) -> Iterable[str]:
+    """Convert to swc format."""
+
+    if comments is not None:
+        for c in comments:
+            if not c.isspace():
+                yield f"# {c.lstrip()}\n"
+            else:
+                yield "#"
+
+    names = get_names(names)
+    cols = names.cols() + (list(extra_cols) if extra_cols is not None else [])
+    yield f"# {' '.join(cols)}\n"
+
+    def get_v(k: str, idx: int) -> str:
+        vs = get_ndata(k)
+        v = vs[idx]
+        if np.issubdtype(vs.dtype, np.floating):
+            return f"{v:.4f}"
+
+        if k == names.id or (k == names.pid and v != -1):
+            v += id_offset
+
+        return str(v)
+
+    for idx in get_ndata(names.id):
+        yield " ".join(get_v(k, idx) for k in cols) + "\n"
+
+
+RE_COMMENT = re.compile(r"^\s*#")
+RE_FLOAT = r"([+-]?(?:\d+(?:[.]\d*)?(?:[eE][+-]?\d+)?|[.]\d+(?:[eE][+-]?\d+)?))"
+
+
+def parse_swc(
+    fname: PathOrIO,
+    *,
+    names: SWCNames,
+    extra_cols: Iterable[str] | None = None,
+    encoding: Literal["detect"] | str = "utf-8",
+) -> tuple[pd.DataFrame, list[str]]:
+    """Parse swc file.
+
+    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 []
+
+    keys = names.cols() + extras
+    vals = [[] for _ in keys]
+    transforms = [int, int, float, float, float, float, int] + [float for _ in extras]
+
+    re_swc_cols = [
+        r"([0-9]+)",  # id
+        r"([0-9]+)",  # type
+        RE_FLOAT,  # x
+        RE_FLOAT,  # y
+        RE_FLOAT,  # z
+        RE_FLOAT,  # r
+        r"(-?[0-9]+)",  # pid
+    ] + [
+        RE_FLOAT
+        for _ in extras  # assert float
+    ]
+
+    re_swc_cols_str = r"\s+".join(re_swc_cols)
+    # Leading spaces are allowed, as this is part of the data in
+    # neuromorpho.org. More fields at the end is allowed, such as
+    # reading eswc as swc, but with a warning.
+    re_swc = re.compile(rf"^\s*{re_swc_cols_str}\s*([\s+-.0-9]*)$")
+
+    last_group = 7 + len(extras) + 1
+    ignored_comment = f"# {' '.join(names.cols())}"
+    flag = True
+
+    comments = []
+    with FileReader(fname, encoding=encoding) as f:
+        try:
+            for i, line in enumerate(f):
+                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}`"
+                        )
+                        flag = False
+
+                    for i, trans in enumerate(transforms):
+                        vals[i].append(trans(match.group(i + 1)))
+                elif match := RE_COMMENT.match(line):
+                    comment = line[len(match.group(0)) :].removesuffix("\n")
+                    if not comment.startswith(ignored_comment):
+                        comments.append(comment)
+                elif not line.isspace():
+                    raise ValueError(f"invalid row {i + 1} in `{fname}`")
+        except UnicodeDecodeError as e:
+            raise ValueError(
+                "decode failed, try to enable auto detect `encoding='detect'`"
+            ) from e
+
+    df = pd.DataFrame.from_dict(dict(zip(keys, vals)))
+    return df, comments

swcgeom/core/swc_utils/normalizer.py

@@ -0,0 +1,163 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""SWC format utils.
+
+Methods ending with a underline imply an in-place transformation.
+"""
+
+from collections.abc import Callable
+from typing import Literal
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+from swcgeom.core.swc_utils.base import SWCNames, Topology, get_dsu, get_names
+
+__all__ = [
+    "mark_roots_as_somas",
+    "mark_roots_as_somas_",
+    "link_roots_to_nearest",
+    "link_roots_to_nearest_",
+    "sort_nodes",
+    "sort_nodes_",
+    "sort_nodes_impl",
+    "reset_index",
+    "reset_index_",
+]
+
+
+def mark_roots_as_somas(
+    df: pd.DataFrame,
+    update_type: int | Literal[False] = 1,
+    *,
+    names: SWCNames | None = None,
+) -> pd.DataFrame:
+    return _copy_and_apply(
+        mark_roots_as_somas_, df, update_type=update_type, names=names
+    )
+
+
+def mark_roots_as_somas_(
+    df: pd.DataFrame,
+    update_type: int | Literal[False] = 1,
+    *,
+    names: SWCNames | None = None,
+) -> None:
+    """Merge multiple roots in swc.
+
+    The first root are reserved and others are linked to it.
+    """
+    names = get_names(names)
+    roots = df[names.pid] == -1
+    root_loc = roots.argmax()
+    root_id = df.loc[root_loc, names.id]  # type:ignore
+    df[names.pid] = np.where(df[names.pid] != -1, df[names.pid], root_id)
+    if update_type is not False:
+        df[names.type] = np.where(df[names.pid] != -1, df[names.type], update_type)
+    df.loc[root_loc, names.pid] = -1  # type:ignore
+
+
+def link_roots_to_nearest(
+    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: SWCNames | None = None) -> None:
+    """Merge multiple roots in swc.
+
+    The first root are reserved, and the others was.
+    """
+    names = get_names(names)
+    dsu = get_dsu(df)
+    roots = df[df[names.pid] == -1].iterrows()
+    next(roots)  # type: ignore # skip the first one
+    for i, row in roots:
+        vs = df[[names.x, names.y, names.z]] - row[[names.x, names.y, names.z]]
+        dis = np.linalg.norm(vs.to_numpy(), axis=1)
+        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()]
+
+
+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
+    """
+    return _copy_and_apply(sort_nodes_, df, names=names)
+
+
+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
+    """
+    names = get_names(names)
+    ids, pids = df[names.id].to_numpy(), df[names.pid].to_numpy()
+    (new_ids, new_pids), indices = sort_nodes_impl((ids, pids))
+    for col in df.columns:
+        df[col] = df[col][indices].to_numpy()
+
+    df[names.id], df[names.pid] = new_ids, new_pids
+
+
+def sort_nodes_impl(topology: Topology) -> tuple[Topology, npt.NDArray[np.int32]]:
+    """Sort the indices of neuron tree.
+
+    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"
+
+    id_map = np.full_like(old_ids, fill_value=-3)  # new_id to old_id
+    new_pids = np.full_like(old_ids, fill_value=-3)
+    new_id = 0
+    first_root = old_ids[(old_pids == -1).argmax()]
+    s: list[tuple[npt.NDArray[np.int32], int]] = [(first_root, -1)]
+    while len(s) != 0:
+        old_id, new_pid = s.pop()
+        id_map[new_id] = old_id
+        new_pids[new_id] = new_pid
+        s.extend((j, new_id) for j in old_ids[old_pids == old_id])  # (old_id, new_pid)
+        new_id = new_id + 1
+
+    id2idx = dict(zip(old_ids, range(len(old_ids))))  # old_id to old_idx
+    indices = np.array([id2idx[i] for i in id_map], dtype=np.int32)  # new_id to old_idx
+    new_ids = np.arange(len(new_pids))
+    return (new_ids, new_pids), indices
+
+
+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: SWCNames | None = None) -> None:
+    """Reset node index to start with zero."""
+    names = get_names(names)
+    roots = df[names.pid] == -1
+    root_loc = roots.argmax()
+    root_id = df.loc[root_loc, names.id]  # type:ignore
+    df[names.id] = df[names.id] - root_id
+    df[names.pid] = df[names.pid] - root_id
+    df.loc[root_loc, names.pid] = -1  # type:ignore
+
+
+def _copy_and_apply(fn: Callable, df: pd.DataFrame, *args, **kwargs):
+    df = df.copy()
+    fn(df, *args, **kwargs)
+    return df

swcgeom/core/swc_utils/subtree.py

@@ -0,0 +1,70 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""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.
+"""
+
+from typing import cast
+
+import numpy as np
+import numpy.typing as npt
+
+from swcgeom.core.swc_utils.base import Topology, traverse
+
+__all__ = ["REMOVAL", "to_sub_topology", "propagate_removal"]
+
+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.
+
+    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.
+    """
+    sub_id = np.array(sub[0], dtype=np.int32)
+    sub_pid = np.array(sub[1], dtype=np.int32)
+
+    # remove nodes
+    keeped_id = cast(npt.NDArray[np.bool_], sub_id != REMOVAL)
+    sub_id, sub_pid = sub_id[keeped_id], sub_pid[keeped_id]
+
+    old2new = {idx: i for i, idx in enumerate(sub_id)}  # old idx to new id
+    new_id = np.arange(0, sub_id.shape[0], dtype=np.int32)
+    new_pid = np.array([old2new[i] if i != -1 else -1 for i in sub_pid], dtype=np.int32)
+
+    return (new_id, new_pid), sub_id
+
+
+def propagate_removal(topology: Topology) -> Topology:
+    """Mark all children when parent is marked as removed.
+
+    Returns:
+        new_topology: Topology
+    """
+
+    new_ids, pids = topology
+    ids = np.arange(0, pids.shape[0])
+
+    def propagate(n: int, parent: bool | None) -> bool:
+        if remove := bool(parent) or (new_ids[n] == REMOVAL):
+            new_ids[n] = REMOVAL
+
+        return remove
+
+    traverse((ids, pids), enter=propagate)
+    return (new_ids, pids.copy())