swcgeom 0.21.0__cp313-cp313-win_amd64.whl

swcgeom/transforms/__init__.py

@@ -0,0 +1,20 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""A series of transformations to compose codes."""
+
+from swcgeom.transforms.base import *  # noqa: F403
+from swcgeom.transforms.branch import *  # noqa: F403
+from swcgeom.transforms.branch_tree import *  # noqa: F403
+from swcgeom.transforms.geometry import *  # noqa: F403
+from swcgeom.transforms.image_preprocess import *  # noqa: F403
+from swcgeom.transforms.image_stack import *  # noqa: F403
+from swcgeom.transforms.images import *  # noqa: F403
+from swcgeom.transforms.mst import *  # noqa: F403
+from swcgeom.transforms.neurolucida_asc import *  # noqa: F403
+from swcgeom.transforms.path import *  # noqa: F403
+from swcgeom.transforms.population import *  # noqa: F403
+from swcgeom.transforms.tree import *  # noqa: F403
+from swcgeom.transforms.tree_assembler import *  # noqa: F403

swcgeom/transforms/base.py

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Transformation in tree."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Generic, TypeVar, overload
+
+from typing_extensions import override
+
+__all__ = ["Transform", "Transforms", "Identity"]
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+T4 = TypeVar("T4")
+T5 = TypeVar("T5")
+T6 = TypeVar("T6")
+
+
+class Transform(ABC, Generic[T, K]):
+    r"""An abstract class representing a :class:`Transform`.
+
+    All transforms that represent a map from `T` to `K`.
+    """
+
+    @abstractmethod
+    def __call__(self, x: T) -> K:
+        """Apply transform.
+
+        NOTE: All subclasses should overwrite :meth:`__call__`, supporting
+        applying transform in `x`.
+        """
+        raise NotImplementedError()
+
+    def __repr__(self) -> str:
+        classname = self.__class__.__name__
+        repr_ = self.extra_repr()
+        return f"{classname}({repr_})"
+
+    def extra_repr(self) -> str:
+        """Provides a human-friendly representation of the module.
+
+        This method extends the basic string representation provided by
+        `__repr__` method. It is designed to display additional details
+        about the module's parameters or its specific configuration,
+        which can be particularly useful for debugging and model
+        architecture introspection.
+
+        >>> class Foo(Transform[T, K]):
+        ...     def __init__(self, my_parameter: int = 1):
+        ...         self.my_parameter = my_parameter
+        ...
+        ...     def extra_repr(self) -> str:
+        ...         return f"my_parameter={self.my_parameter}"
+
+        NOTE: This method should be overridden in custom modules to provide
+        specific details relevant to the module's functionality and
+        configuration.
+        """
+        return ""
+
+
+class Transforms(Transform[T, K]):
+    """A simple typed wrapper for transforms."""
+
+    transforms: list[Transform[Any, Any]]
+
+    # fmt: off
+    @overload
+    def __init__(self, t1: Transform[T, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, T5], t6: Transform[T5, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, T5], t6: Transform[T5, T6],
+                       t7: Transform[T6, K], /) -> None: ...
+    @overload
+    def __init__(self, t1: Transform[T, T1],  t2: Transform[T1, T2],
+                       t3: Transform[T2, T3], t4: Transform[T3, T4],
+                       t5: Transform[T4, T5], t6: Transform[T5, T6],
+                       t7: Transform[T6, Any], /, *transforms: Transform[Any, K]) -> None: ...
+    # fmt: on
+    def __init__(self, *transforms: Transform[Any, Any]) -> None:
+        trans = []
+        for t in transforms:
+            if isinstance(t, Transforms):
+                trans.extend(t.transforms)
+            else:
+                trans.append(t)
+        self.transforms = trans
+
+    @override
+    def __call__(self, x: T) -> K:
+        """Apply transforms."""
+        for transform in self.transforms:
+            x = transform(x)
+
+        return x  # type: ignore
+
+    def __getitem__(self, idx: int) -> Transform[Any, Any]:
+        return self.transforms[idx]
+
+    def __len__(self) -> int:
+        return len(self.transforms)
+
+    @override
+    def extra_repr(self) -> str:
+        return ", ".join([str(transform) for transform in self])
+
+
+class Identity(Transform[T, T]):
+    """Resurn input as-is."""
+
+    @override
+    def __call__(self, x: T) -> T:
+        return x

swcgeom/transforms/branch.py

@@ -0,0 +1,223 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Transformation in branch."""
+
+from abc import ABC, abstractmethod
+from typing import cast
+
+import numpy as np
+import numpy.typing as npt
+from scipy import signal
+from typing_extensions import override
+
+from swcgeom.core import Branch, DictSWC
+from swcgeom.transforms.base import Transform
+from swcgeom.utils import (
+    angle,
+    rotate3d_x,
+    rotate3d_y,
+    rotate3d_z,
+    scale3d,
+    to_homogeneous,
+    translate3d,
+)
+
+__all__ = ["BranchLinearResampler", "BranchConvSmoother", "BranchStandardizer"]
+
+
+class _BranchResampler(Transform[Branch, Branch], ABC):
+    r"""Resample branch."""
+
+    @override
+    def __call__(self, x: Branch) -> Branch:
+        xyzr = self.resample(x.xyzr())
+        return Branch.from_xyzr(xyzr)
+
+    @abstractmethod
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: ...
+
+
+class BranchLinearResampler(_BranchResampler):
+    r"""Resampling by linear interpolation, DO NOT keep original node."""
+
+    def __init__(self, n_nodes: int) -> None:
+        """Resample branch to special num of nodes.
+
+        Args:
+            n_nodes: Number of nodes after resample.
+        """
+        super().__init__()
+        self.n_nodes = n_nodes
+
+    @override
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        """Resampling by linear interpolation, DO NOT keep original node.
+
+        Args:
+            xyzr: The array of shape (N, 4).
+
+        Returns:
+            coordinates: An array of shape (n_nodes, 4).
+        """
+
+        xp = np.cumsum(np.linalg.norm(xyzr[1:, :3] - xyzr[:-1, :3], axis=1))
+        xp = np.insert(xp, 0, 0)
+        xvals = np.linspace(0, xp[-1], self.n_nodes)
+
+        x = np.interp(xvals, xp, xyzr[:, 0])
+        y = np.interp(xvals, xp, xyzr[:, 1])
+        z = np.interp(xvals, xp, xyzr[:, 2])
+        r = np.interp(xvals, xp, xyzr[:, 3])
+        return cast(npt.NDArray[np.float32], np.stack([x, y, z, r], axis=1))
+
+    @override
+    def extra_repr(self) -> str:
+        return f"n_nodes={self.n_nodes}"
+
+
+class BranchIsometricResampler(_BranchResampler):
+    def __init__(self, distance: float, *, adjust_last_gap: bool = True) -> None:
+        super().__init__()
+        self.distance = distance
+        self.adjust_last_gap = adjust_last_gap
+
+    @override
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        """Resampling by isometric interpolation, DO NOT keep original node.
+
+        Args:
+            xyzr: The array of shape (N, 4).
+
+        Returns:
+            new_xyzr: An array of shape (n_nodes, 4).
+        """
+
+        # Compute the cumulative distances between consecutive points
+        diffs = np.diff(xyzr[:, :3], axis=0)
+        distances = np.sqrt((diffs**2).sum(axis=1))
+        cumulative_distances = np.concatenate([[0], np.cumsum(distances)])
+
+        total_length = cumulative_distances[-1]
+        n_nodes = int(np.ceil(total_length / self.distance)) + 1
+
+        # Determine the new distances
+        if self.adjust_last_gap and n_nodes > 1:
+            new_distances = np.linspace(0, total_length, n_nodes)
+        else:
+            new_distances = np.arange(0, total_length, self.distance)
+            # keep endpoint
+            new_distances = np.concatenate([new_distances, total_length])
+
+        # Interpolate the new points
+        new_xyzr = np.zeros((n_nodes, 4), dtype=np.float32)
+        new_xyzr[:, :3] = np.array(
+            [
+                np.interp(new_distances, cumulative_distances, xyzr[:, i])
+                for i in range(3)
+            ]
+        ).T
+        new_xyzr[:, 3] = np.interp(new_distances, cumulative_distances, xyzr[:, 3])
+        return new_xyzr
+
+    @override
+    def extra_repr(self) -> str:
+        return f"distance={self.distance},adjust_last_gap={self.adjust_last_gap}"
+
+
+class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
+    r"""Smooth the branch by sliding window."""
+
+    def __init__(self, n_nodes: int = 5) -> None:
+        """
+        Args:
+            n_nodes: Window size.
+        """
+        super().__init__()
+        self.n_nodes = n_nodes
+        self.kernel = np.ones(n_nodes)
+
+    @override
+    def __call__(self, x: Branch) -> Branch[DictSWC]:
+        x = x.detach()
+        c = signal.convolve(np.ones(x.number_of_nodes()), self.kernel, mode="same")
+        for k in ["x", "y", "z"]:
+            v = x.get_ndata(k)
+            s = signal.convolve(v, self.kernel, mode="same")
+            x.attach.ndata[k][1:-1] = (s / c)[1:-1]
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"n_nodes={self.n_nodes}"
+
+
+class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):
+    r"""Standardize branch.
+
+    Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up at y, and scale max
+    radius to 1.
+    """
+
+    @override
+    def __call__(self, x: Branch) -> Branch:
+        xyzr = x.xyzr()
+        xyz, r = xyzr[:, 0:3], xyzr[:, 3:4]
+        T = self.get_matrix(xyz)
+
+        xyz4 = to_homogeneous(xyz, 1).transpose()  # (4, N)
+        new_xyz = np.dot(T, xyz4)[:3].transpose()
+        new_xyzr = np.concatenate([new_xyz, r / r.max()], axis=1)
+        return Branch.from_xyzr(new_xyzr)
+
+    @staticmethod
+    def get_matrix(xyz: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        r"""Get standardize transformation matrix.
+
+        Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up at y.
+
+        Args:
+            xyz: The `x`, `y`, `z` matrix of shape (N, 3) of branch.
+
+        Returns:
+            T: An homogeneous transformation matrix of shape (4, 4).
+        """
+
+        assert xyz.ndim == 2 and xyz.shape[1] == 3, (
+            f"xyz should be of shape (N, 3), got {xyz.shape}"
+        )
+
+        xyz = xyz[:, :3]
+        T = np.identity(4)
+        v = np.concatenate([xyz[-1] - xyz[0], np.zeros((1))])[:, None]
+
+        # translate to the origin
+        T = translate3d(-xyz[0, 0], -xyz[0, 1], -xyz[0, 2]).dot(T)
+
+        # scale to unit vector
+        s = (1 / np.linalg.norm(v[:3, 0])).item()
+        T = scale3d(s, s, s).dot(T)
+
+        # rotate v to the xz-plane, v should be (x, 0, z) now
+        vy = np.dot(T, v)[:, 0]
+        # when looking at the xz-plane along the positive y-axis, the
+        # coordinates should be (z, x)
+        T = rotate3d_y(angle([vy[2], vy[0]], [0, 1])).dot(T)
+
+        # rotate v to the x-axis, v should be (1, 0, 0) now
+        vx = np.dot(T, v)[:, 0]
+        T = rotate3d_z(angle([vx[0], vx[1]], [1, 0])).dot(T)
+
+        # rotate the farthest point to the xy-plane
+        if xyz.shape[0] > 2:
+            xyz4 = to_homogeneous(xyz, 1).transpose()  # (4, N)
+            new_xyz4 = np.dot(T, xyz4)  # (4, N)
+            max_index = np.argmax(np.linalg.norm(new_xyz4[1:3, :], axis=0)[1:-1]) + 1
+            max_xyz4 = xyz4[:, max_index].reshape(4, 1)
+            max_xyz4_t = np.dot(T, max_xyz4)  # (4, 1)
+            angle_x = angle(max_xyz4_t[1:3, 0], [1, 0])
+            T = rotate3d_x(angle_x).dot(T)
+
+        return T

swcgeom/transforms/branch_tree.py

@@ -0,0 +1,74 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+
+import numpy as np
+from typing_extensions import override
+
+from swcgeom.core import Branch, BranchTree, Node, Tree
+from swcgeom.transforms.base import Transform
+
+__all__ = ["BranchTreeAssembler"]
+
+
+class BranchTreeAssembler(Transform[BranchTree, Tree]):
+    EPS = 1e-6
+
+    @override
+    def __call__(self, x: BranchTree) -> Tree:
+        nodes = [x.soma().detach()]
+        stack = [(x.soma(), 0)]  # n_orig, id_new
+        while len(stack):
+            n_orig, pid_new = stack.pop()
+            children = n_orig.children()
+
+            for br, c in self.pair(x.branches.get(n_orig.id, []), children):
+                s = 1 if np.linalg.norm(br[0].xyz() - n_orig.xyz()) < self.EPS else 0
+                e = -2 if np.linalg.norm(br[-1].xyz() - c.xyz()) < self.EPS else -1
+
+                br_nodes = [n.detach() for n in br[s:e]] + [c.detach()]
+                for i, n in enumerate(br_nodes):
+                    # reindex
+                    n.id = len(nodes) + i
+                    n.pid = len(nodes) + i - 1
+
+                br_nodes[0].pid = pid_new
+                nodes.extend(br_nodes)
+                stack.append((c, br_nodes[-1].id))
+
+        return Tree(
+            len(nodes),
+            source=x.source,
+            comments=x.comments,
+            names=x.names,
+            **{
+                k: np.array([n.__getattribute__(k) for n in nodes])
+                for k in x.names.cols()
+            },
+        )
+
+    def pair(
+        self, branches: list[Branch], endpoints: list[Node]
+    ) -> Iterable[tuple[Branch, Node]]:
+        assert len(branches) == len(endpoints)
+        xyz1 = [br[-1].xyz() for br in branches]
+        xyz2 = [n.xyz() for n in endpoints]
+        v = np.reshape(xyz1, (-1, 1, 3)) - np.reshape(xyz2, (1, -1, 3))
+        dis = np.linalg.norm(v, axis=-1)
+
+        # greedy algorithm
+        pairs = []
+        for _ in range(len(branches)):
+            # find minimal
+            min_idx = np.argmin(dis)
+            min_branch_idx, min_endpoint_idx = np.unravel_index(min_idx, dis.shape)
+            pairs.append((branches[min_branch_idx], endpoints[min_endpoint_idx]))
+
+            # remove current node
+            dis[min_branch_idx, :] = np.inf
+            dis[:, min_endpoint_idx] = np.inf
+
+        return pairs