swcgeom 0.16.0__py3-none-any.whl → 0.18.3__py3-none-any.whl

swcgeom/images/io.py

@@ -1,33 +1,38 @@
+# Copyright 2022-2025 Zexin Yuan
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 """Read and write image stack."""
 
 import os
 import re
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
 from functools import cache, lru_cache
-from typing import (
-    Any,
-    Callable,
-    Generic,
-    Iterable,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-    TypeVar,
-    cast,
-    overload,
-)
+from typing import Any, Generic, Literal, Optional, TypeVar, cast, overload
 
 import nrrd
 import numpy as np
 import numpy.typing as npt
 import tifffile
+from typing_extensions import deprecated
 from v3dpy.loaders import PBD, Raw
 
 __all__ = ["read_imgs", "save_tiff", "read_images"]
 
-Vec3i = Tuple[int, int, int]
+Vec3i = tuple[int, int, int]
 ScalarType = TypeVar("ScalarType", bound=np.generic, covariant=True)
 
 RE_TERAFLY_ROOT = re.compile(r"^RES\((\d+)x(\d+)x(\d+)\)$")
@@ -58,17 +63,17 @@ class ImageStack(ABC, Generic[ScalarType]):
     def __getitem__(self, key: int) -> npt.NDArray[ScalarType]: ...                     # array of shape (Y, Z, C)
     @overload
     @abstractmethod
-    def __getitem__(self, key: Tuple[int, int]) -> npt.NDArray[ScalarType]: ...         # array of shape (Z, C)
+    def __getitem__(self, key: tuple[int, int]) -> npt.NDArray[ScalarType]: ...         # array of shape (Z, C)
     @overload
     @abstractmethod
-    def __getitem__(self, key: Tuple[int, int, int]) -> npt.NDArray[ScalarType]: ...    # array of shape (C,)
+    def __getitem__(self, key: tuple[int, int, int]) -> npt.NDArray[ScalarType]: ...    # array of shape (C,)
     @overload
     @abstractmethod
-    def __getitem__(self, key: Tuple[int, int, int, int]) -> ScalarType: ...            # value
+    def __getitem__(self, key: tuple[int, int, int, int]) -> ScalarType: ...            # value
     @overload
     @abstractmethod
     def __getitem__(
-        self, key: slice | Tuple[slice, slice] | Tuple[slice, slice, slice] |  Tuple[slice, slice, slice, slice],
+        self, key: slice | tuple[slice, slice] | tuple[slice, slice, slice] |  tuple[slice, slice, slice, slice],
     ) -> npt.NDArray[ScalarType]: ... # array of shape (X, Y, Z, C)
     @overload
     @abstractmethod
@@ -95,7 +100,7 @@ class ImageStack(ABC, Generic[ScalarType]):
         return self[:, :, :, :]
 
     @property
-    def shape(self) -> Tuple[int, int, int, int]:
+    def shape(self) -> tuple[int, int, int, int]:
         raise NotImplementedError()
 
 
@@ -107,26 +112,42 @@ def read_imgs(fname: str, *, dtype: None =..., **kwargs) -> ImageStack[np.float3
 # fmt:on
 
 
-def read_imgs(fname: str, *, dtype=None, **kwargs):  # type: ignore
-    """Read image stack."""
+def read_imgs(fname: str, **kwargs):  # type: ignore
+    """Read image stack.
 
-    kwargs["dtype"] = dtype or np.float32
+    Parameters
+    ----------
+    fname : str
+        The path of image stack.
+    dtype : np.dtype, default to `np.float32`
+        Casting data to specified dtype. If integer and float
+        conversions occur, they will be scaled (assuming floats are
+        between 0 and 1).
+    **kwargs : dict[str, Any]
+        Forwarding to the corresponding reader.
+    """
 
-    ext = os.path.splitext(fname)[-1]
-    if ext in [".tif", ".tiff"]:
-        return TiffImageStack(fname, **kwargs)
-    if ext in [".nrrd"]:
-        return NrrdImageStack(fname, **kwargs)
-    if ext in [".v3dpbd"]:
-        return V3dpbdImageStack(fname, **kwargs)
-    if ext in [".v3draw"]:
-        return V3drawImageStack(fname, **kwargs)
-    if ext in [".npy", ".npz"]:
-        return NDArrayImageStack(np.load(fname), **kwargs)
+    kwargs.setdefault("dtype", np.float32)
+    if not os.path.exists(fname):
+        raise ValueError(f"image stack not exists: {fname}")
+
+    # match file extension
+    match os.path.splitext(fname)[-1]:
+        case ".tif" | ".tiff":
+            return TiffImageStack(fname, **kwargs)
+        case ".nrrd":
+            return NrrdImageStack(fname, **kwargs)
+        case ".v3dpbd":
+            return V3dpbdImageStack(fname, **kwargs)
+        case ".v3draw":
+            return V3drawImageStack(fname, **kwargs)
+        case ".npy":
+            return NDArrayImageStack(np.load(fname), **kwargs)
+
+    # try to read as terafly
     if TeraflyImageStack.is_root(fname):
         return TeraflyImageStack(fname, **kwargs)
-    if not os.path.exists(fname):
-        raise ValueError("image stack not exists")
+
     raise ValueError("unsupported image stack")
 
 
@@ -135,7 +156,6 @@ def save_tiff(
     fname: str,
     *,
     dtype: Optional[np.unsignedinteger | np.floating] = None,
-    swap_xy: Optional[bool] = None,
     compression: str | Literal[False] = "zlib",
     **kwargs,
 ) -> None:
@@ -154,7 +174,7 @@ def save_tiff(
         Compression algorithm, forwarding to `tifffile.imwrite`. If no
         algorithnm is specify specified, we will use the zlib algorithm
         with compression level 6 by default.
-    **kwargs : Dict[str, Any]
+    **kwargs : dict[str, Any]
         Forwarding to `tifffile.imwrite`
     """
     if isinstance(data, ImageStack):
@@ -164,17 +184,6 @@ def save_tiff(
         data = np.expand_dims(data, -1)  # (_, _, _)  -> (_, _, _, C), C === 1
 
     axes = "ZXYC"
-    if swap_xy is not None:
-        warnings.warn(
-            "flag `swap_xy` is easy to implement in user space and "
-            "is more flexiable. Since this flag is rarely used, we "
-            "decided to remove it in the next version",
-            DeprecationWarning,
-        )
-        if swap_xy is True:
-            axes = "ZYXC"
-            data = data.swapaxes(0, 1)  # (X, Y, _, _) -> (Y, X, _, _)
-
     assert data.ndim == 4, "should be an array of shape (X, Y, Z, C)"
     assert data.shape[-1] in [1, 3], "support 'miniblack' or 'rgb'"
 
@@ -209,12 +218,7 @@ class NDArrayImageStack(ImageStack[ScalarType]):
     """NDArray image stack."""
 
     def __init__(
-        self,
-        imgs: npt.NDArray[Any],
-        swap_xy: Optional[bool] = None,
-        filp_xy: Optional[bool] = None,
-        *,
-        dtype: ScalarType,
+        self, imgs: npt.NDArray[Any], *, dtype: Optional[ScalarType] = None
     ) -> None:
         super().__init__()
 
@@ -222,34 +226,22 @@ class NDArrayImageStack(ImageStack[ScalarType]):
             imgs = np.expand_dims(imgs, -1)
         assert imgs.ndim == 4, "Should be shape of (X, Y, Z, C)"
 
-        if swap_xy is not None:
-            warnings.warn(
-                "flag `swap_xy` now is unnecessary, tifffile will "
-                "automatically adjust dimensions according to "
-                "`tags.axes`, so this flag will be removed in the next "
-                " version",
-                DeprecationWarning,
-            )
-            if swap_xy is True:
-                imgs = imgs.swapaxes(0, 1)  # (Y, X, _, _) -> (X, Y, _, _)
-
-        if filp_xy is not None:
-            warnings.warn(
-                "flag `filp_xy` is easy to implement in user space and "
-                "is more flexiable. Since this flag is rarely used, we "
-                "decided to remove it in the next version",
-                DeprecationWarning,
-            )
-            if filp_xy is True:
-                imgs = np.flip(imgs, (0, 1))  # (X, Y, Z, C)
+        if dtype is not None:
+            dtype_raw = imgs.dtype
+            if np.issubdtype(dtype, np.floating) and np.issubdtype(
+                dtype_raw, np.unsignedinteger
+            ):
+                sclar_factor = 1.0 / UINT_MAX[dtype_raw]
+                imgs = sclar_factor * imgs.astype(dtype)
+            elif np.issubdtype(dtype, np.unsignedinteger) and np.issubdtype(
+                dtype_raw, np.floating
+            ):
+                sclar_factor = UINT_MAX[dtype]  # type: ignore
+                imgs *= (sclar_factor * imgs).astype(dtype)
+            else:
+                imgs = imgs.astype(dtype)
 
-        dtype_raw = imgs.dtype
-        self.imgs = imgs.astype(dtype)
-        if np.issubdtype(dtype, np.floating) and np.issubdtype(
-            dtype_raw, np.unsignedinteger
-        ):  # TODO: add a option to disable this
-            sclar_factor = 1.0 / UINT_MAX[imgs.dtype]
-            self.imgs *= sclar_factor
+        self.imgs = imgs
 
     def __getitem__(self, key):
         return self.imgs.__getitem__(key)
@@ -258,22 +250,14 @@ class NDArrayImageStack(ImageStack[ScalarType]):
         return self.imgs
 
     @property
-    def shape(self) -> Tuple[int, int, int, int]:
-        return cast(Tuple[int, int, int, int], self.imgs.shape)
+    def shape(self) -> tuple[int, int, int, int]:
+        return cast(tuple[int, int, int, int], self.imgs.shape)
 
 
 class TiffImageStack(NDArrayImageStack[ScalarType]):
     """Tiff image stack."""
 
-    def __init__(
-        self,
-        fname: str,
-        swap_xy: Optional[bool] = None,
-        filp_xy: Optional[bool] = None,
-        *,
-        dtype: ScalarType,
-        **kwargs,
-    ) -> None:
+    def __init__(self, fname: str, *, dtype: ScalarType, **kwargs) -> None:
         with tifffile.TiffFile(fname, **kwargs) as f:
             s = f.series[0]
             imgs, axes = s.asarray(), s.axes
@@ -285,23 +269,15 @@ class TiffImageStack(NDArrayImageStack[ScalarType]):
 
         orders = [AXES_ORDER[c] for c in axes]
         imgs = imgs.transpose(np.argsort(orders))
-        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy, dtype=dtype)
+        super().__init__(imgs, dtype=dtype)
 
 
 class NrrdImageStack(NDArrayImageStack[ScalarType]):
     """Nrrd image stack."""
 
-    def __init__(
-        self,
-        fname: str,
-        swap_xy: Optional[bool] = None,
-        filp_xy: Optional[bool] = None,
-        *,
-        dtype: ScalarType,
-        **kwargs,
-    ) -> None:
+    def __init__(self, fname: str, *, dtype: ScalarType, **kwargs) -> None:
         imgs, header = nrrd.read(fname, **kwargs)
-        super().__init__(imgs, swap_xy=swap_xy, filp_xy=filp_xy, dtype=dtype)
+        super().__init__(imgs, dtype=dtype)
         self.header = header
 
 
@@ -353,7 +329,7 @@ class TeraflyImageStack(ImageStack[ScalarType]):
     use its coordinate system, remember to FLIP Y-AXIS BACK.
     """
 
-    _listdir: Callable[[str], List[str]]
+    _listdir: Callable[[str], list[str]]
     _read_patch: Callable[[str], npt.NDArray]
 
     def __init__(
@@ -379,12 +355,17 @@ class TeraflyImageStack(ImageStack[ScalarType]):
         self.res, self.res_dirs, self.res_patch_sizes = self.get_resolutions(root)
 
         @cache
-        def listdir(path: str) -> List[str]:
+        def listdir(path: str) -> list[str]:
             return os.listdir(path)
 
         @lru_cache(maxsize=lru_maxsize)
         def read_patch(path: str) -> npt.NDArray[ScalarType]:
-            return read_imgs(path, dtype=dtype).get_full()
+            match os.path.splitext(path)[-1]:
+                case "raw":
+                    # Treat it as a v3draw file
+                    return V3drawImageStack(path, dtype=dtype).get_full()
+                case _:
+                    return read_imgs(path, dtype=dtype).get_full()
 
         self._listdir, self._read_patch = listdir, read_patch
 
@@ -453,19 +434,19 @@ class TeraflyImageStack(ImageStack[ScalarType]):
         raise NotImplementedError()  # TODO
 
     @property
-    def shape(self) -> Tuple[int, int, int, int]:
+    def shape(self) -> tuple[int, int, int, int]:
         res_max = self.res[-1]
         return res_max[0], res_max[1], res_max[2], 1
 
     @classmethod
-    def get_resolutions(cls, root: str) -> Tuple[List[Vec3i], List[str], List[Vec3i]]:
+    def get_resolutions(cls, root: str) -> tuple[list[Vec3i], list[str], list[Vec3i]]:
         """Get all resolutions.
 
         Returns
         -------
         resolutions : List of (int, int, int)
             Sequence of sorted resolutions (from small to large).
-        roots : List[str]
+        roots : list[str]
             Sequence of root of resolutions respectively.
         patch_sizes : List of (int, int, int)
             Sequence of patch size of resolutions respectively.
@@ -572,7 +553,7 @@ class TeraflyImageStack(ImageStack[ScalarType]):
             if (invalid := diff > 10 * v).all():
                 return None, None
 
-            diff[invalid] = np.NINF  # remove values which greate than v
+            diff[invalid] = -np.inf  # remove values which greater than v
 
             # find the index of the value smaller than v and closest to v
             idx = np.argmax(diff)
@@ -605,7 +586,7 @@ class GrayImageStack:
     @overload
     def __getitem__(self, key: npt.NDArray[np.integer[Any]]) -> np.float32: ...
     @overload
-    def __getitem__(self, key: slice | Tuple[slice, slice] | Tuple[slice, slice, slice]) -> npt.NDArray[np.float32]: ...
+    def __getitem__(self, key: slice | tuple[slice, slice] | tuple[slice, slice, slice]) -> npt.NDArray[np.float32]: ...
     # fmt: on
     def __getitem__(self, key):
         """Get pixel/patch of image stack."""
@@ -632,10 +613,11 @@ class GrayImageStack:
         return self.imgs.get_full()[:, :, :, 0]
 
     @property
-    def shape(self) -> Tuple[int, int, int]:
+    def shape(self) -> tuple[int, int, int]:
         return self.imgs.shape[:-1]
 
 
+@deprecated("Use `read_imgs` instead")
 def read_images(*args, **kwargs) -> GrayImageStack:
     """Read images.
 
@@ -643,9 +625,4 @@ def read_images(*args, **kwargs) -> GrayImageStack:
         Use :meth:`read_imgs` instead.
     """
 
-    warnings.warn(
-        "`read_images` has been replaced by `read_imgs` because it"
-        "provide rgb support, and this will be removed in next version",
-        DeprecationWarning,
-    )
     return GrayImageStack(read_imgs(*args, **kwargs))

swcgeom/transforms/__init__.py

@@ -1,14 +1,30 @@
+# Copyright 2022-2025 Zexin Yuan
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 """A series of transformations to compose codes."""
 
-from swcgeom.transforms.base import *
-from swcgeom.transforms.branch import *
-from swcgeom.transforms.geometry import *
-from swcgeom.transforms.image_preprocess import *
-from swcgeom.transforms.image_stack import *
-from swcgeom.transforms.images import *
-from swcgeom.transforms.mst import *
-from swcgeom.transforms.neurolucida_asc import *
-from swcgeom.transforms.path import *
-from swcgeom.transforms.population import *
-from swcgeom.transforms.tree import *
-from swcgeom.transforms.tree_assembler import *
+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

@@ -1,3 +1,18 @@
+# Copyright 2022-2025 Zexin Yuan
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 """Transformation in tree."""
 
 from abc import ABC, abstractmethod
@@ -33,7 +48,7 @@ class Transform(ABC, Generic[T, K]):
         repr_ = self.extra_repr()
         return f"{classname}({repr_})"
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         """Provides a human-friendly representation of the module.
 
         This method extends the basic string representation provided by
@@ -48,7 +63,7 @@ class Transform(ABC, Generic[T, K]):
             def __init__(self, my_parameter: int = 1):
                 self.my_parameter = my_parameter
 
-            def extra_repr(self):
+            def extra_repr(self) -> str:
                 return f"my_parameter={self.my_parameter}"
 
         Notes

swcgeom/transforms/branch.py

@@ -1,3 +1,18 @@
+# Copyright 2022-2025 Zexin Yuan
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 """Transformation in branch."""
 
 from abc import ABC, abstractmethod
@@ -72,10 +87,61 @@ class BranchLinearResampler(_BranchResampler):
         r = np.interp(xvals, xp, xyzr[:, 3])
         return cast(npt.NDArray[np.float32], np.stack([x, y, z, r], axis=1))
 
-    def extra_repr(self):
+    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
+
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        """Resampling by isometric interpolation, DO NOT keep original node.
+
+        Parameters
+        ----------
+        xyzr : np.ndarray[np.float32]
+            The array of shape (N, 4).
+
+        Returns
+        -------
+        new_xyzr : ~numpy.NDArray[float32]
+            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
+
+    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."""
 
@@ -88,24 +154,24 @@ class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
         """
         super().__init__()
         self.n_nodes = n_nodes
-        self.kernal = np.ones(n_nodes)
+        self.kernel = np.ones(n_nodes)
 
     def __call__(self, x: Branch) -> Branch[DictSWC]:
         x = x.detach()
-        c = signal.convolve(np.ones(x.number_of_nodes()), self.kernal, mode="same")
+        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.kernal, mode="same")
+            s = signal.convolve(v, self.kernel, mode="same")
             x.attach.ndata[k][1:-1] = (s / c)[1:-1]
 
         return x
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         return f"n_nodes={self.n_nodes}"
 
 
 class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):
-    r"""Standarize branch.
+    r"""Standardize branch.
 
     Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up at
     y, and scale max radius to 1.
@@ -123,7 +189,7 @@ class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):
 
     @staticmethod
     def get_matrix(xyz: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
-        r"""Get standarize transformation matrix.
+        r"""Get standardize transformation matrix.
 
         Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up
         at y.
@@ -136,7 +202,7 @@ class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):
         Returns
         -------
         T : np.ndarray[np.float32]
-            An homogeneous transfomation matrix of shape (4, 4).
+            An homogeneous transformation matrix of shape (4, 4).
         """
 
         assert (

swcgeom/transforms/branch_tree.py

@@ -0,0 +1,82 @@
+# Copyright 2022-2025 Zexin Yuan
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Iterable
+
+import numpy as np
+
+from swcgeom.core import Branch, BranchTree, Node, Tree
+from swcgeom.transforms.base import Transform
+
+__all__ = ["BranchTreeAssembler"]
+
+
+class BranchTreeAssembler(Transform[BranchTree, Tree]):
+    EPS = 1e-6
+
+    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