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

swcgeom/transforms/base.py

@@ -18,12 +18,19 @@
 from abc import ABC, abstractmethod
 from typing import Any, Generic, TypeVar, overload
 
+from typing_extensions import override
+
 __all__ = ["Transform", "Transforms", "Identity"]
 
-T, K = TypeVar("T"), TypeVar("K")
+T = TypeVar("T")
+K = TypeVar("K")
 
-T1, T2, T3 = TypeVar("T1"), TypeVar("T2"), TypeVar("T3")  # pylint: disable=invalid-name
-T4, T5, T6 = TypeVar("T4"), TypeVar("T5"), TypeVar("T6")  # pylint: disable=invalid-name
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+T4 = TypeVar("T4")
+T5 = TypeVar("T5")
+T6 = TypeVar("T6")
 
 
 class Transform(ABC, Generic[T, K]):
@@ -36,9 +43,7 @@ class Transform(ABC, Generic[T, K]):
     def __call__(self, x: T) -> K:
         """Apply transform.
 
-        Notes
-        -----
-        All subclasses should overwrite :meth:`__call__`, supporting
+        NOTE: All subclasses should overwrite :meth:`__call__`, supporting
         applying transform in `x`.
         """
         raise NotImplementedError()
@@ -57,18 +62,14 @@ class Transform(ABC, Generic[T, K]):
         which can be particularly useful for debugging and model
         architecture introspection.
 
-        Examples
-        --------
-        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}"
+        >>> 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}"
 
-        Notes
-        -----
-        This method should be overridden in custom modules to provide
+        NOTE: This method should be overridden in custom modules to provide
         specific details relevant to the module's functionality and
         configuration.
         """
@@ -80,7 +81,7 @@ class Transforms(Transform[T, K]):
 
     transforms: list[Transform[Any, Any]]
 
-    # fmt:off
+    # fmt: off
     @overload
     def __init__(self, t1: Transform[T, K], /) -> None: ...
     @overload
@@ -100,11 +101,16 @@ class Transforms(Transform[T, K]):
                        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
+    # fmt: on
     def __init__(self, *transforms: Transform[Any, Any]) -> None:
         trans = []
         for t in transforms:
@@ -114,6 +120,7 @@ class Transforms(Transform[T, K]):
                 trans.append(t)
         self.transforms = trans
 
+    @override
     def __call__(self, x: T) -> K:
         """Apply transforms."""
         for transform in self.transforms:
@@ -127,6 +134,7 @@ class Transforms(Transform[T, K]):
     def __len__(self) -> int:
         return len(self.transforms)
 
+    @override
     def extra_repr(self) -> str:
         return ", ".join([str(transform) for transform in self])
 
@@ -134,5 +142,6 @@ class Transforms(Transform[T, K]):
 class Identity(Transform[T, T]):
     """Resurn input as-is."""
 
+    @override
     def __call__(self, x: T) -> T:
         return x

swcgeom/transforms/branch.py

@@ -21,6 +21,7 @@ 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
@@ -40,13 +41,13 @@ __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]:
-        raise NotImplementedError()
+    def resample(self, xyzr: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: ...
 
 
 class BranchLinearResampler(_BranchResampler):
@@ -55,26 +56,21 @@ class BranchLinearResampler(_BranchResampler):
     def __init__(self, n_nodes: int) -> None:
         """Resample branch to special num of nodes.
 
-        Parameters
-        ----------
-        n_nodes : int
-            Number of nodes after resample.
+        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.
 
-        Parameters
-        ----------
-        xyzr : np.ndarray[np.float32]
-            The array of shape (N, 4).
+        Args:
+            xyzr: The array of shape (N, 4).
 
-        Returns
-        -------
-        coordinates : ~numpy.NDArray[float64]
-            An array of shape (n_nodes, 4).
+        Returns:
+            coordinates: An array of shape (n_nodes, 4).
         """
 
         xp = np.cumsum(np.linalg.norm(xyzr[1:, :3] - xyzr[:-1, :3], axis=1))
@@ -87,6 +83,7 @@ class BranchLinearResampler(_BranchResampler):
         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}"
 
@@ -97,18 +94,15 @@ class BranchIsometricResampler(_BranchResampler):
         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.
 
-        Parameters
-        ----------
-        xyzr : np.ndarray[np.float32]
-            The array of shape (N, 4).
+        Args:
+            xyzr: The array of shape (N, 4).
 
-        Returns
-        -------
-        new_xyzr : ~numpy.NDArray[float32]
-            An array of shape (n_nodes, 4).
+        Returns:
+            new_xyzr: An array of shape (n_nodes, 4).
         """
 
         # Compute the cumulative distances between consecutive points
@@ -138,6 +132,7 @@ class BranchIsometricResampler(_BranchResampler):
         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}"
 
@@ -147,15 +142,14 @@ class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
 
     def __init__(self, n_nodes: int = 5) -> None:
         """
-        Parameters
-        ----------
-        n_nodes : int, default `5`
-            Window size.
+        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")
@@ -173,10 +167,11 @@ class BranchConvSmoother(Transform[Branch, Branch[DictSWC]]):
 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.
+    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]
@@ -191,23 +186,18 @@ class BranchStandardizer(Transform[Branch, Branch[DictSWC]]):
     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.
+        Standardized branch starts at (0, 0, 0), ends at (1, 0, 0), up at y.
 
-        Parameters
-        ----------
-        xyz : np.ndarray[np.float32]
-            The `x`, `y`, `z` matrix of shape (N, 3) of branch.
+        Args:
+            xyz: The `x`, `y`, `z` matrix of shape (N, 3) of branch.
 
-        Returns
-        -------
-        T : np.ndarray[np.float32]
-            An homogeneous transformation matrix of shape (4, 4).
+        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}"
+        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)

swcgeom/transforms/branch_tree.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 
-from typing import Iterable
+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
@@ -26,6 +27,7 @@ __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

swcgeom/transforms/geometry.py

@@ -16,10 +16,11 @@
 """SWC geometry operations."""
 
 import warnings
-from typing import Generic, Literal, Optional, TypeVar
+from typing import Generic, Literal, TypeVar
 
 import numpy as np
 import numpy.typing as npt
+from typing_extensions import override
 
 from swcgeom.core import DictSWC
 from swcgeom.core.swc_utils import SWCNames
@@ -54,7 +55,7 @@ Center = Literal["root", "soma", "origin"]
 class Normalizer(Generic[T], Transform[T, T]):
     """Noramlize coordinates and radius to 0-1."""
 
-    def __init__(self, *, names: Optional[SWCNames] = None) -> None:
+    def __init__(self, *, names: SWCNames | None = None) -> None:
         super().__init__()
         if names is not None:
             warnings.warn(
@@ -63,6 +64,7 @@ class Normalizer(Generic[T], Transform[T, T]):
                 DeprecationWarning,
             )
 
+    @override
     def __call__(self, x: T) -> T:
         """Scale the `x`, `y`, `z`, `r` of nodes to 0-1."""
         new_tree = x.copy()
@@ -87,6 +89,7 @@ class RadiusReseter(Generic[T], Transform[T, T]):
         new_tree.ndata[new_tree.names.r] = r
         return new_tree
 
+    @override
     def extra_repr(self) -> str:
         return f"r={self.r:.4f}"
 
@@ -103,8 +106,8 @@ class AffineTransform(Generic[T], Transform[T, T]):
         tm: npt.NDArray[np.float32],
         center: Center = "origin",
         *,
-        fmt: Optional[str] = None,
-        names: Optional[SWCNames] = None,
+        fmt: str | None = None,
+        names: SWCNames | None = None,
     ) -> None:
         self.tm, self.center = tm, center
 
@@ -122,6 +125,7 @@ class AffineTransform(Generic[T], Transform[T, T]):
                 DeprecationWarning,
             )
 
+    @override
     def __call__(self, x: T) -> T:
         match self.center:
             case "root" | "soma":
@@ -156,6 +160,7 @@ class Translate(Generic[T], AffineTransform[T]):
         super().__init__(translate3d(tx, ty, tz), **kwargs)
         self.tx, self.ty, self.tz = tx, ty, tz
 
+    @override
     def extra_repr(self) -> str:
         return f"tx={self.tx:.4f}, ty={self.ty:.4f}, tz={self.tz:.4f}"
 
@@ -209,6 +214,7 @@ class Rotate(Generic[T], AffineTransform[T]):
         self.theta = theta
         self.center = center
 
+    @override
     def extra_repr(self) -> str:
         return f"n={self.n}, theta={self.theta:.4f}, center={self.center}"  # TODO: improve format of n
 
@@ -231,6 +237,7 @@ class RotateX(Generic[T], AffineTransform[T]):
         super().__init__(rotate3d_x(theta), center=center, **kwargs)
         self.theta = theta
 
+    @override
     def extra_repr(self) -> str:
         return f"center={self.center}, theta={self.theta:.4f}"
 
@@ -247,6 +254,7 @@ class RotateY(Generic[T], AffineTransform[T]):
         self.theta = theta
         self.center = center
 
+    @override
     def extra_repr(self) -> str:
         return f"theta={self.theta:.4f}, center={self.center}"
 
@@ -263,6 +271,7 @@ class RotateZ(Generic[T], AffineTransform[T]):
         self.theta = theta
         self.center = center
 
+    @override
     def extra_repr(self) -> str:
         return f"theta={self.theta:.4f}, center={self.center}"
 

swcgeom/transforms/image_preprocess.py

@@ -19,6 +19,7 @@ import numpy as np
 import numpy.typing as npt
 from scipy.fftpack import fftn, fftshift, ifftn
 from scipy.ndimage import gaussian_filter, minimum_filter
+from typing_extensions import override
 
 from swcgeom.transforms.base import Transform
 
@@ -36,6 +37,7 @@ class SGuoImPreProcess(Transform[npt.NDArray[np.uint8], npt.NDArray[np.uint8]]):
     January 2022, Pages 503–512, https://doi.org/10.1093/bioinformatics/btab638
     """
 
+    @override
     def __call__(self, x: npt.NDArray[np.uint8]) -> npt.NDArray[np.uint8]:
         # TODO: support np.float32
         assert x.dtype == np.uint8, "Image must be in uint8 format"

swcgeom/transforms/image_stack.py

@@ -15,9 +15,7 @@
 
 """Create image stack from morphology.
 
-Notes
------
-All denpendencies need to be installed, try:
+NOTE: All denpendencies need to be installed, try:
 
 ```sh
 pip install swcgeom[all]
@@ -28,7 +26,7 @@ import os
 import re
 import time
 from collections.abc import Iterable
-from typing import Optional
+from typing import Sequence
 
 import numpy as np
 import numpy.typing as npt
@@ -42,7 +40,7 @@ from sdflit import (
     SDFObject,
 )
 from tqdm import tqdm
-from typing_extensions import deprecated
+from typing_extensions import deprecated, override
 
 from swcgeom.core import Population, Tree
 from swcgeom.transforms.base import Transform
@@ -58,46 +56,53 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
     def __init__(self, resolution: int | float | npt.ArrayLike = 1) -> None:
         """Transform tree to image stack.
 
-        Parameters
-        ----------
-        resolution : int | (x, y, z), default `(1, 1, 1)`
-            Resolution of image stack.
+        Args:
+            resolution: Resolution of image stack.
+                If a scalar, it will be broadcasted to a vector of 3d.
         """
-
         if isinstance(resolution, (int, float, np.integer, np.floating)):
             resolution = [resolution, resolution, resolution]  # type: ignore
 
         self.resolution = np.array(resolution, dtype=np.float32)
-        assert len(self.resolution) == 3, "resolution shoule be vector of 3d."
+        assert len(self.resolution) == 3, "resolution should be vector of 3d."
 
+    @override
     def __call__(self, x: Tree) -> npt.NDArray[np.uint8]:
         """Transform tree to image stack.
 
-        Notes
-        -----
-        This method loads the entire image stack into memory, so it
-        ONLY works for small image stacks, use :meth`transform_and_save`
-        for big image stack.
+        NOTE: This method loads the entire image stack into memory, so it ONLY works
+        for small image stacks, use :meth`transform_and_save` for big image stack.
         """
         return np.stack(list(self.transform(x, verbose=False)), axis=0)
 
     def transform(
         self,
-        x: Tree,
+        x: Tree | Sequence[Tree],
         verbose: bool = True,
         *,
-        ranges: Optional[tuple[npt.ArrayLike, npt.ArrayLike]] = None,
+        ranges: tuple[npt.ArrayLike, npt.ArrayLike] | None = None,
     ) -> Iterable[npt.NDArray[np.uint8]]:
+        trees = [x] if isinstance(x, Tree) else x
+        if not trees:
+            return iter([])  # Return empty iterator if sequence is empty
+
+        time_start = None
         if verbose:
-            print("To image stack: " + x.source)
+            sources = ", ".join(t.source for t in trees if t.source)
+            print(f"To image stack: {sources if sources else 'unnamed trees'}")
             time_start = time.time()
 
-        scene = self._get_scene(x)
+        scene = self._get_scene(trees)
 
         if ranges is None:
-            xyz, r = x.xyz(), x.r().reshape(-1, 1)
-            coord_min = np.floor(np.min(xyz - r, axis=0))
-            coord_max = np.ceil(np.max(xyz + r, axis=0))
+            all_xyz = np.concatenate([t.xyz() for t in trees], axis=0)
+            all_r = np.concatenate([t.r() for t in trees], axis=0).reshape(-1, 1)
+            if all_xyz.size == 0:  # Handle empty trees
+                coord_min = np.zeros(3, dtype=np.float32)
+                coord_max = np.zeros(3, dtype=np.float32)
+            else:
+                coord_min = np.floor(np.min(all_xyz - all_r, axis=0))
+                coord_max = np.ceil(np.max(all_xyz + all_r, axis=0))
         else:
             assert len(ranges) == 2
             coord_min = np.array(ranges[0])
@@ -106,12 +111,12 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
 
         samplers = self._get_samplers(coord_min, coord_max)
 
-        if verbose:
+        if verbose and time_start is not None:
             total = (coord_max[2] - coord_min[2]) / self.resolution[2]
             samplers = tqdm(samplers, total=total.astype(np.int64).item())
 
             time_end = time.time()
-            print("Prepare in: ", time_end - time_start, "s")  # type: ignore
+            print("Prepare in: ", time_end - time_start, "s")
 
         for sampler in samplers:
             voxel = sampler.sample(scene)  # should be shape of (x, y, z, 3) and z = 1
@@ -124,12 +129,12 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
         x: Tree,
         verbose: bool = True,
         *,
-        ranges: Optional[tuple[npt.ArrayLike, npt.ArrayLike]] = None,
+        ranges: tuple[npt.ArrayLike, npt.ArrayLike] | None = None,
     ) -> Iterable[npt.NDArray[np.uint8]]:
         return self.transform(x, verbose, ranges=ranges)
 
     def transform_and_save(
-        self, fname: str, x: Tree, verbose: bool = True, **kwargs
+        self, fname: str, x: Tree | Sequence[Tree], verbose: bool = True, **kwargs
     ) -> None:
         self.save_tif(fname, self.transform(x, verbose=verbose, **kwargs))
 
@@ -151,11 +156,12 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
             if not os.path.isfile(tif):
                 self.transform_and_save(tif, tree, verbose=False)
 
+    @override
     def extra_repr(self) -> str:
         res = ",".join(f"{a:.4f}" for a in self.resolution)
         return f"resolution=({res})"
 
-    def _get_scene(self, x: Tree) -> Scene:
+    def _get_scene(self, trees: Sequence[Tree]) -> Scene:
         material = ColoredMaterial((1, 0, 0)).into()
         scene = ObjectsScene()
         scene.set_background((0, 0, 0))
@@ -164,10 +170,11 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
             for c in children:
                 sdf = RoundCone(_tp3f(n.xyz()), _tp3f(c.xyz()), n.r, c.r).into()
                 scene.add_object(SDFObject(sdf, material).into())
-
             return n
 
-        x.traverse(leave=leave)
+        for tree in trees:
+            tree.traverse(leave=leave)
+
         scene.build_bvh()
         return scene.into()
 
@@ -175,14 +182,12 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
         self,
         coord_min: npt.NDArray,
         coord_max: npt.NDArray,
-        offset: Optional[npt.NDArray] = None,
+        offset: npt.NDArray | None = None,
     ) -> Iterable[RangeSampler]:
         """Get Samplers.
 
-        Parameters
-        ----------
-        coord_min, coord_max: npt.ArrayLike
-            Coordinates array of shape (3,).
+        Args:
+            coord_min, coord_max: Coordinates array of shape (3,).
         """
 
         eps = 1e-6