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

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

swcgeom/transforms/images.py

@@ -17,7 +17,7 @@
 
 import numpy as np
 import numpy.typing as npt
-from typing_extensions import deprecated
+from typing_extensions import deprecated, override
 
 from swcgeom.transforms.base import Identity, Transform
 
@@ -49,12 +49,14 @@ class ImagesCenterCrop(Transform[NDArrayf32, NDArrayf32]):
             else (shape_out, shape_out, shape_out)
         )
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         diff = np.subtract(x.shape[:3], self.shape_out)
         s = diff // 2
         e = np.add(s, self.shape_out)
         return x[s[0] : e[0], s[1] : e[1], s[2] : e[2], :]
 
+    @override
     def extra_repr(self) -> str:
         return f"shape_out=({','.join(str(a) for a in self.shape_out)})"
 
@@ -73,9 +75,11 @@ class ImagesScale(Transform[NDArrayf32, NDArrayf32]):
         super().__init__()
         self.scaler = scaler
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         return self.scaler * x
 
+    @override
     def extra_repr(self) -> str:
         return f"scaler={self.scaler}"
 
@@ -85,9 +89,11 @@ class ImagesClip(Transform[NDArrayf32, NDArrayf32]):
         super().__init__()
         self.vmin, self.vmax = vmin, vmax
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         return np.clip(x, self.vmin, self.vmax)
 
+    @override
     def extra_repr(self) -> str:
         return f"vmin={self.vmin}, vmax={self.vmax}"
 
@@ -99,9 +105,11 @@ class ImagesFlip(Transform[NDArrayf32, NDArrayf32]):
         super().__init__()
         self.axis = axis
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         return np.flip(x, axis=self.axis)
 
+    @override
     def extra_repr(self) -> str:
         return f"axis={self.axis}"
 
@@ -109,15 +117,13 @@ class ImagesFlip(Transform[NDArrayf32, NDArrayf32]):
 class ImagesFlipY(ImagesFlip):
     """Flip image stack along Y-axis.
 
-    See Also
-    --------
-    ~.images.io.TeraflyImageStack:
-        Terafly and Vaa3d use a especial right-handed coordinate system
-        (with origin point in the left-top and z-axis points front),
-        but we flip y-axis to makes it a left-handed coordinate system
-        (with orgin point in the left-bottom and z-axis points front).
-        If you need to use its coordinate system, remember to FLIP
-        Y-AXIS BACK.
+    See Also:
+        ~.images.io.TeraflyImageStack:
+            Terafly and Vaa3d use a especial right-handed coordinate system (with
+            origin point in the left-top and z-axis points front), but we flip y-axis
+            to makes it a left-handed coordinate system (with origin point in the
+            left-bottom and z-axis points front). If you need to use its coordinate
+            system, remember to FLIP Y-AXIS BACK.
     """
 
     def __init__(self, axis: int = 1, /) -> None:
@@ -127,6 +133,7 @@ class ImagesFlipY(ImagesFlip):
 class ImagesNormalizer(Transform[NDArrayf32, NDArrayf32]):
     """Normalize image stack."""
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         mean = np.mean(x)
         variance = np.var(x)
@@ -136,9 +143,8 @@ class ImagesNormalizer(Transform[NDArrayf32, NDArrayf32]):
 class ImagesMeanVarianceAdjustment(Transform[NDArrayf32, NDArrayf32]):
     """Adjust image stack mean and variance.
 
-    See Also
-    --------
-    ~swcgeom.images.ImageStackFolder.stat
+    See Also:
+        ~swcgeom.images.ImageStackFolder.stat
     """
 
     def __init__(self, mean: float, variance: float) -> None:
@@ -146,9 +152,11 @@ class ImagesMeanVarianceAdjustment(Transform[NDArrayf32, NDArrayf32]):
         self.mean = mean
         self.variance = variance
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         return (x - self.mean) / self.variance
 
+    @override
     def extra_repr(self) -> str:
         return f"mean={self.mean}, variance={self.variance}"
 
@@ -159,14 +167,10 @@ class ImagesScaleToUnitRange(Transform[NDArrayf32, NDArrayf32]):
     def __init__(self, vmin: float, vmax: float, *, clip: bool = True) -> None:
         """Scale image stack to unit range.
 
-        Parameters
-        ----------
-        vmin : float
-            Minimum value.
-        vmax : float
-            Maximum value.
-        clip : bool, default True
-            Clip values to [0, 1] to avoid numerical issues.
+        Args:
+            vmin: Minimum value.
+            vmax: Maximum value.
+            clip: Clip values to [0, 1] to avoid numerical issues.
         """
 
         super().__init__()
@@ -176,9 +180,11 @@ class ImagesScaleToUnitRange(Transform[NDArrayf32, NDArrayf32]):
         self.clip = clip
         self.post = ImagesClip(0, 1) if self.clip else Identity()
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         return self.post((x - self.vmin) / self.diff)
 
+    @override
     def extra_repr(self) -> str:
         return f"vmin={self.vmin}, vmax={self.vmax}, clip={self.clip}"
 
@@ -186,15 +192,15 @@ class ImagesScaleToUnitRange(Transform[NDArrayf32, NDArrayf32]):
 class ImagesHistogramEqualization(Transform[NDArrayf32, NDArrayf32]):
     """Image histogram equalization.
 
-    References
-    ----------
-    http://www.janeriksolem.net/histogram-equalization-with-python-and.html
+    References:
+        http://www.janeriksolem.net/histogram-equalization-with-python-and.html
     """
 
     def __init__(self, bins: int = 256) -> None:
         super().__init__()
         self.bins = bins
 
+    @override
     def __call__(self, x: NDArrayf32) -> NDArrayf32:
         # get image histogram
         hist, bin_edges = np.histogram(x.flatten(), self.bins, density=True)
@@ -205,5 +211,6 @@ class ImagesHistogramEqualization(Transform[NDArrayf32, NDArrayf32]):
         equalized = np.interp(x.flatten(), bin_edges[:-1], cdf)
         return equalized.reshape(x.shape).astype(np.float32)
 
+    @override
     def extra_repr(self) -> str:
         return f"bins={self.bins}"

swcgeom/transforms/mst.py

@@ -16,12 +16,12 @@
 """Minimum spanning tree."""
 
 import warnings
-from typing import Optional
 
 import numpy as np
 import pandas as pd
 from numpy import ma
 from numpy import typing as npt
+from typing_extensions import override
 
 from swcgeom.core import Tree, sort_tree
 from swcgeom.core.swc_utils import SWCNames, SWCTypes, get_names, get_types
@@ -36,8 +36,7 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
     Creates trees corresponding to the minimum spanning tree keeping
     the path length to the root small (with balancing factor bf).
 
-    References
-    ----------
+    References:
     .. [1] Cuntz, H., Forstner, F., Borst, A. & Häusser, M. One Rule to
        Grow Them Al: A General Theory of Neuronal Branching and Its
        Practical Application. PLOS Comput Biol 6, e1000877 (2010).
@@ -52,21 +51,15 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
         furcations: int = 2,
         exclude_soma: bool = True,
         sort: bool = True,
-        names: Optional[SWCNames] = None,
-        types: Optional[SWCTypes] = None,
+        names: SWCNames | None = None,
+        types: SWCTypes | None = None,
     ) -> None:
         """
-        Parameters
-        ----------
-        bf : float, default `0.4`
-            Balancing factor between 0~1.
-        furcations : int, default `2`
-            Suppress multi-furcations which more than k. If set to -1,
-            no suppression.
-        exclude_soma : bool, default `True`
-            Suppress multi-furcations exclude soma.
-        names : SWCNames, optional
-        types : SWCTypes, optional
+        Args:
+            bf: Balancing factor between 0~1.
+            furcations: Suppress multi-furcations which more than k.
+                If set to -1, no suppression.
+            exclude_soma: Suppress multi-furcations exclude soma.
         """
         self.bf = np.clip(bf, 0, 1)
         self.furcations = furcations
@@ -75,29 +68,26 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
         self.names = get_names(names)
         self.types = get_types(types)
 
+    @override
     def __call__(  # pylint: disable=too-many-locals
         self,
         points: npt.NDArray[np.floating],
-        soma: Optional[npt.ArrayLike] = None,
+        soma: npt.ArrayLike | None = None,
         *,
-        names: Optional[SWCNames] = None,
+        names: SWCNames | None = None,
     ) -> Tree:
         """
-        Paramters
-        ---------
-        points : array of shape (N, 3)
-            Positions of points cloud.
-        soma : array of shape (3,), default `None`
-            Position of soma. If none, use the first point as soma.
-        names : SWCNames, optional
+        Args:
+            points: Positions of points cloud.
+            soma: Position of soma. If none, use the first point as soma.
         """
         if names is None:
             names = self.names
         else:
             warnings.warn(
-                "`PointsToCuntzMST(...)(names=...)` has been "
-                "replaced by `PointsToCuntzMST(...,names=...)` since "
-                "v0.12.0, and will be removed in next version",
+                "`PointsToCuntzMST(...)(names=...)` has been replaced by "
+                "`PointsToCuntzMST(...,names=...)` since v0.12.0, and will be removed "
+                "in next version",
                 DeprecationWarning,
             )
             names = get_names(names)  # TODO: remove it
@@ -155,6 +145,7 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
             t = sort_tree(t)
         return t
 
+    @override
     def extra_repr(self) -> str:  # TODO: names, types
         return f"bf={self.bf:.4f}, furcations={self.furcations}, exclude_soma={self.exclude_soma}, sort={self.sort}"
 
@@ -166,22 +157,17 @@ class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
         self,
         furcations: int = 2,
         *,
-        k_furcations: Optional[int] = None,
+        k_furcations: int | None = None,
         exclude_soma: bool = True,
-        names: Optional[SWCNames] = None,
-        types: Optional[SWCTypes] = None,
+        names: SWCNames | None = None,
+        types: SWCTypes | None = None,
         **kwargs,
     ) -> None:
         """
-        Parameters
-        ----------
-        furcations : int, default `2`
-            Suppress multifurcations which more than k. If set to -1,
-            no suppression.
-        exclude_soma : bool, default `True`
-            Suppress multi-furcations exclude soma.
-        names : SWCNames, optional
-        types : SWCTypes, optional
+        Args:
+            furcations: Suppress multifurcations which more than k.
+                If set to -1, no suppression.
+            exclude_soma: Suppress multi-furcations exclude soma.
         """
 
         if k_furcations is not None:
@@ -202,5 +188,6 @@ class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
             **kwargs,
         )
 
+    @override
     def extra_repr(self) -> str:
         return f"furcations-{self.furcations}, exclude-soma={self.exclude_soma}"