swcgeom 0.13.2__py3-none-any.whl → 0.15.0__py3-none-any.whl

swcgeom/transforms/image_stack.py

@@ -2,60 +2,55 @@
 
 Notes
 -----
-- This module has been deprecated, try `github.com/yzx9/swc2skeleton`
-- All denpendencies need to be installed, try:
+All denpendencies need to be installed, try:
 
 ```sh
 pip install swcgeom[all]
 ```
 """
 
-import math
 import os
 import re
 import time
-from typing import Any, Iterable, List, Tuple, cast
+from typing import Iterable, List, Optional, Tuple
 
 import numpy as np
 import numpy.typing as npt
 import tifffile
+from sdflit import (
+    ColoredMaterial,
+    ObjectsScene,
+    RangeSampler,
+    RoundCone,
+    Scene,
+    SDFObject,
+)
 
 from swcgeom.core import Population, Tree
 from swcgeom.transforms.base import Transform
-from swcgeom.utils import SDF, SDFCompose, SDFRoundCone
 
 __all__ = ["ToImageStack"]
 
 
-# TODO: migrate to github.com/yzx9/swc2skeleton
 class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
     r"""Transform tree to image stack."""
 
     resolution: npt.NDArray[np.float32]
-    msaa: int
-    z_per_iter: int = 1
 
-    def __init__(
-        self, resolution: float | npt.ArrayLike = 1, msaa: int = 8, z_per_iter: int = 1
-    ) -> None:
+    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.
-        mass : int, default `8`
-            Multi-sample anti-aliasing.
         """
 
-        if isinstance(resolution, float):
-            resolution = [resolution, resolution, resolution]
-
-        self.resolution = (resolution := np.array(resolution, dtype=np.float32))
-        assert tuple(resolution.shape) == (3,), "resolution shoule be vector of 3d."
+        if isinstance(resolution, (int, float, np.integer, np.floating)):
+            resolution = [resolution, resolution, resolution]  # type: ignore
 
-        self.msaa = msaa
-        self.z_per_iter = z_per_iter
+        self.resolution = np.array(resolution, dtype=np.float32)
+        assert len(self.resolution) == 3, "resolution shoule be vector of 3d."
 
     def __call__(self, x: Tree) -> npt.NDArray[np.uint8]:
         """Transform tree to image stack.
@@ -63,142 +58,121 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
         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.
+        ONLY works for small image stacks, use :meth`transform_and_save`
+        for big image stack.
         """
-        return np.concatenate(list(self.transfrom(x, verbose=False)), axis=0)
-
-    def __repr__(self) -> str:
-        return (
-            "ToImageStack"
-            + f"-resolution-{'-'.join(self.resolution)}"
-            + f"-mass-{self.msaa}"
-            + f"-z-{self.z_per_iter}"
-        )
+        return np.stack(list(self.transfrom(x, verbose=False)), axis=0)
 
     def transfrom(
-        self, x: Tree, verbose: bool = True
+        self,
+        x: Tree,
+        verbose: bool = True,
+        *,
+        ranges: Optional[Tuple[npt.ArrayLike, npt.ArrayLike]] = None,
     ) -> Iterable[npt.NDArray[np.uint8]]:
-        # pylint: disable=too-many-locals
-        from tqdm import tqdm
-
         if verbose:
             print("To image stack: " + x.source)
             time_start = time.time()
 
-        sdf = self.get_sdf(x)
+        scene = self._get_scene(x)
+
+        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))
+        else:
+            assert len(ranges) == 2
+            coord_min = np.array(ranges[0])
+            coord_max = np.array(ranges[1])
+            assert len(coord_min) == len(coord_max) == 3
 
-        xyz, r = x.xyz(), np.stack([x.r(), x.r(), x.r()], axis=1)  # TODO: perf
-        coord_min = np.floor(np.min(xyz - r, axis=0))  # TODO: snap to grid
-        coord_max = np.ceil(np.max(xyz + r, axis=0))
-        grids, total = self.get_grids(coord_min, coord_max)
+        samplers = self._get_samplers(coord_min, coord_max)
 
         if verbose:
+            from tqdm import tqdm
+
+            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
 
-        for grid in tqdm(grids, total=total) if verbose else grids:
-            is_in = sdf.is_in(grid.reshape(-1, 3)).reshape(*grid.shape[:4])
-            level = np.sum(is_in, axis=3, dtype=np.uint8)
-            voxel = (255 / self.msaa * level).astype(np.uint8)
-            for i in range(voxel.shape[2]):
-                yield voxel[:, :, i]
+        for sampler in samplers:
+            voxel = sampler.sample(scene)  # should be shape of (x, y, z, 3) and z = 1
+            frame = (255 * voxel[..., 0, 0]).astype(np.uint8)
+            yield frame
 
-    def transform_and_save(self, fname: str, x: Tree, verbose: bool = True) -> None:
-        self.save_tif(fname, self.transfrom(x, verbose=verbose))
+    def transform_and_save(
+        self, fname: str, x: Tree, verbose: bool = True, **kwargs
+    ) -> None:
+        self.save_tif(fname, self.transfrom(x, verbose=verbose, **kwargs))
 
     def transform_population(
         self, population: Population | str, verbose: bool = True
     ) -> None:
-        if isinstance(population, str):
-            population = Population.from_swc(population)
+        trees = (
+            Population.from_swc(population)
+            if isinstance(population, str)
+            else population
+        )
+
+        if verbose:
+            from tqdm import tqdm
+
+            trees = tqdm(trees)
 
         # TODO: multiprocess
-        for tree in population:
+        for tree in trees:
             tif = re.sub(r".swc$", ".tif", tree.source)
             if not os.path.isfile(tif):
-                self.transform_and_save(tif, tree, verbose=verbose)
+                self.transform_and_save(tif, tree, verbose=False)
+
+    def extra_repr(self):
+        res = ",".join(f"{a:.4f}" for a in self.resolution)
+        return f"resolution=({res})"
+
+    def _get_scene(self, x: Tree) -> Scene:
+        material = ColoredMaterial((1, 0, 0)).into()
+        scene = ObjectsScene()
+        scene.set_background((0, 0, 0))
+
+        def leave(n: Tree.Node, children: List[Tree.Node]) -> Tree.Node:
+            for c in children:
+                sdf = RoundCone(_tp3f(n.xyz()), _tp3f(c.xyz()), n.r, c.r).into()
+                scene.add_object(SDFObject(sdf, material).into())
 
-    def get_grids(
-        self, coord_min: npt.ArrayLike, coord_max: npt.ArrayLike
-    ) -> Tuple[Iterable[npt.NDArray[np.float32]], int]:
-        """Get point grid.
+            return n
+
+        x.traverse(leave=leave)
+        scene.build_bvh()
+        return scene.into()
+
+    def _get_samplers(
+        self,
+        coord_min: npt.NDArray,
+        coord_max: npt.NDArray,
+        offset: Optional[npt.NDArray] = None,
+    ) -> Iterable[RangeSampler]:
+        """Get Samplers.
 
         Parameters
         ----------
         coord_min, coord_max: npt.ArrayLike
             Coordinates array of shape (3,).
-        z_per_iter : int
-            Yeild z per iter, raising this option speeds up processing,
-            but consumes more memory.
-
-        Returns
-        -------
-        grid : npt.NDArray[np.float32]
-            Array of shape (nx, ny, z_per_iter, k, 3).
         """
 
-        k = np.cbrt(self.msaa)
-        coord_min, coord_max = np.array(coord_min), np.array(coord_max)
-        assert tuple(coord_min.shape) == (3,), "coord_min shoule be vector of 3d."
-        assert tuple(coord_max.shape) == (3,), "coord_max shoule be vector of 3d."
-
-        point_grid = np.mgrid[
-            coord_min[0] : coord_max[0] : self.resolution[0],
-            coord_min[1] : coord_max[1] : self.resolution[1],
-            coord_min[2] : coord_max[2] : self.resolution[2],
-        ]  # (3, nx, ny, nz)
-        point_grid = np.rollaxis(point_grid, 0, 4)  # (nx, ny, nz, 3)
-
-        step = self.resolution / (k + 1)
-        ends = self.resolution - step / 2
-        inter_grid = np.mgrid[
-            step[0] : ends[0] : step[0],
-            step[1] : ends[1] : step[1],
-            step[2] : ends[2] : step[2],
-        ]  # (3, kx, ky, kz)
-        inter_grid = np.rollaxis(inter_grid, 0, 4).reshape(-1, 3)  # (k, 3)
-
-        grids = np.expand_dims(point_grid, 3).repeat(k**3, axis=3)
-        grids = cast(Any, grids + inter_grid)
-
-        return (
-            grids[:, :, i : i + self.z_per_iter]
-            for i in range(0, grids.shape[2], self.z_per_iter)
-        ), math.ceil(grids.shape[2] / self.z_per_iter)
-
-    def get_sdf(self, x: Tree) -> SDF:
-        T = Tuple[Tree.Node, List[SDF], SDF | None]
-
-        def collect(n: Tree.Node, pre: List[T]) -> T:
-            if len(pre) == 0:
-                return (n, [], None)
-
-            if len(pre) == 1:
-                child, sub_sdfs, last = pre[0]
-                sub_sdfs.append(SDFRoundCone(n.xyz(), child.xyz(), n.r, child.r))
-                return (n, sub_sdfs, last)
-
-            sdfs: List[SDF] = []
-            for child, sub_sdfs, last in pre:
-                sub_sdfs.append(SDFRoundCone(n.xyz(), child.xyz(), n.r, child.r))
-                sdfs.append(SDFCompose.compose(sub_sdfs))
-                if last is not None:
-                    sdfs.append(last)
-
-            return (n, [], SDFCompose.compose(sdfs))
-
-        _, sdfs, last = x.traverse(leave=collect)
-        if len(sdfs) != 0:
-            sdf = SDFCompose.compose(sdfs)
-            if last is not None:
-                sdf = SDFCompose.compose([sdf, last])
-        elif last is not None:
-            sdf = last
-        else:
-            raise ValueError("empty tree")
+        eps = 1e-6
+        stride = self.resolution
+        offset = offset or (stride / 2)
 
-        return sdf
+        xmin, ymin, zmin = _tp3f(coord_min + offset)
+        xmax, ymax, zmax = _tp3f(coord_max)
+        z = zmin
+        while z < zmax:
+            yield RangeSampler(
+                (xmin, ymin, z), (xmax, ymax, z + stride[2] - eps), _tp3f(stride)
+            )
+            z += stride[2]
 
     @staticmethod
     def save_tif(
@@ -218,3 +192,9 @@ class ToImageStack(Transform[Tree, npt.NDArray[np.uint8]]):
                         "axes": "ZXY",
                     },
                 )
+
+
+def _tp3f(x: npt.NDArray) -> Tuple[float, float, float]:
+    """Convert to tuple of 3 floats."""
+    assert len(x) == 3
+    return (float(x[0]), float(x[1]), float(x[2]))

swcgeom/transforms/images.py

@@ -28,5 +28,5 @@ class Center(Transform[npt.NDArray[np.float32], npt.NDArray[np.float32]]):
         e = np.add(s, self.shape_out)
         return x[s[0] : e[0], s[1] : e[1], s[2] : e[2], :]
 
-    def __repr__(self) -> str:
-        return f"Center-{self.shape_out[0]}-{self.shape_out[1]}-{self.shape_out[2]}"
+    def extra_repr(self) -> str:
+        return f"shape_out=({','.join(str(a) for a in self.shape_out)})"

swcgeom/transforms/mst.py

@@ -140,13 +140,8 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
             t = sort_tree(t)
         return t
 
-    def __repr__(self) -> str:
-        return (
-            f"PointsToCuntzMST"
-            f"-bf-{self.bf}"
-            f"-furcations-{self.furcations}"
-            f"-{'exclude-soma' if self.exclude_soma else 'include-soma'}"
-        )  # TODO: names, types
+    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}"
 
 
 class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
@@ -173,6 +168,7 @@ class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
         names : SWCNames, optional
         types : SWCTypes, optional
         """
+
         if k_furcations is not None:
             warnings.warn(
                 "`PointsToMST(k_furcations=...)` has been renamed to "
@@ -191,9 +187,5 @@ class PointsToMST(PointsToCuntzMST):  # pylint: disable=too-few-public-methods
             **kwargs,
         )
 
-    def __repr__(self) -> str:
-        return (
-            f"PointsToMST"
-            f"-furcations-{self.furcations}"
-            f"-{'exclude-soma' if self.exclude_soma else 'include-soma'}"
-        )
+    def extra_repr(self) -> str:
+        return f"furcations-{self.furcations}, exclude-soma={self.exclude_soma}"

swcgeom/transforms/population.py

@@ -25,5 +25,5 @@ class PopulationTransform(Transform[Population, Population]):
 
         return Population(trees, root=population.root)
 
-    def __repr__(self) -> str:
-        return f"pop({self.transform})"
+    def extra_repr(self) -> str:
+        return f"transform={self.transform}"

swcgeom/transforms/tree.py

@@ -64,8 +64,8 @@ class TreeSmoother(Transform[Tree, Tree]):  # pylint: disable=missing-class-docs
 
         return x
 
-    def __repr__(self) -> str:
-        return f"TreeSmoother-{self.n_nodes}"
+    def extra_repr(self):
+        return f"n_nodes={self.n_nodes}"
 
 
 class TreeNormalizer(Normalizer[Tree]):
@@ -107,8 +107,8 @@ class CutByType(Transform[Tree, Tree]):
         y = to_subtree(x, removals)
         return y
 
-    def __repr__(self) -> str:
-        return f"CutByType-{self.type}"
+    def extra_repr(self):
+        return f"type={self.type}"
 
 
 class CutAxonTree(CutByType):
@@ -118,9 +118,6 @@ class CutAxonTree(CutByType):
         types = get_types(types)
         super().__init__(type=types.axon)
 
-    def __repr__(self) -> str:
-        return "CutAxonTree"
-
 
 class CutDendriteTree(CutByType):
     """Cut dendrite tree."""
@@ -129,9 +126,6 @@ class CutDendriteTree(CutByType):
         types = get_types(types)
         super().__init__(type=types.basal_dendrite)  # TODO: apical dendrite
 
-    def __repr__(self) -> str:
-        return "CutDenriteTree"
-
 
 class CutByBifurcationOrder(Transform[Tree, Tree]):
     """Cut tree by bifurcation order."""
@@ -177,9 +171,6 @@ class CutShortTipBranch(Transform[Tree, Tree]):
         if callback is not None:
             self.callbacks.append(callback)
 
-    def __repr__(self) -> str:
-        return f"CutShortTipBranch-{self.thre}"
-
     def __call__(self, x: Tree) -> Tree:
         removals: List[int] = []
         self.callbacks.append(lambda br: removals.append(br[1].id))
@@ -187,6 +178,9 @@ class CutShortTipBranch(Transform[Tree, Tree]):
         self.callbacks.pop()
         return to_subtree(x, removals)
 
+    def extra_repr(self):
+        return f"threshold={self.thre}"
+
     def _leave(
         self, n: Tree.Node, children: List[Tuple[float, Tree.Node] | None]
     ) -> Tuple[float, Tree.Node] | None:

swcgeom/transforms/tree_assembler.py

@@ -1,17 +1,22 @@
 """Assemble a tree."""
 
+from copy import copy
 from typing import Iterable, List, Optional, Tuple
 
+import numpy as np
 import pandas as pd
 
 from swcgeom.core import Tree
-from swcgeom.core.swc_utils import SWCNames
-from swcgeom.core.swc_utils.assembler import (
-    assemble_lines_impl,
-    try_assemble_lines_impl,
+from swcgeom.core.swc_utils import (
+    SWCNames,
+    get_names,
+    link_roots_to_nearest_,
+    sort_nodes_,
 )
 from swcgeom.transforms.base import Transform
 
+EPS = 1e-5
+
 
 class LinesToTree(Transform[List[pd.DataFrame], Tree]):
     """Assemble lines to swc."""
@@ -35,11 +40,12 @@ class LinesToTree(Transform[List[pd.DataFrame], Tree]):
     ):
         return self.assemble(lines, names=names)
 
-    def __repr__(self) -> str:
-        return f"LinesToTree-thre-{self.thre}-{'undirected' if self.undirected else 'directed'}"
-
     def assemble(
-        self, lines: Iterable[pd.DataFrame], *, names: Optional[SWCNames] = None
+        self,
+        lines: Iterable[pd.DataFrame],
+        *,
+        undirected: bool = True,
+        names: Optional[SWCNames] = None,
     ) -> pd.DataFrame:
         """Assemble lines to a tree.
 
@@ -51,6 +57,8 @@ class LinesToTree(Transform[List[pd.DataFrame], Tree]):
         lines : List of ~pd.DataFrame
             An array of tables containing a line, columns should follwing
             the swc.
+        undirected : bool, default `True`
+            Forwarding to `self.try_assemble`.
         names : SWCNames, optional
             Forwarding to `self.try_assemble`.
 
@@ -60,17 +68,33 @@ class LinesToTree(Transform[List[pd.DataFrame], Tree]):
 
         See Also
         --------
-        self.try_assemble_lines
+        self.try_assemble
         """
-        return assemble_lines_impl(
-            lines, thre=self.thre, undirected=self.undirected, names=names
+
+        tree, lines = self.try_assemble(
+            lines, sort_nodes=False, undirected=undirected, names=names
         )
+        while len(lines) > 0:
+            t, lines = self.try_assemble(
+                lines,
+                id_offset=len(tree),
+                sort_nodes=False,
+                undirected=undirected,
+                names=names,
+            )
+            tree = pd.concat([tree, t])
+
+        tree = tree.reset_index()
+        link_roots_to_nearest_(tree)
+        sort_nodes_(tree)
+        return tree
 
     def try_assemble(
         self,
         lines: Iterable[pd.DataFrame],
         *,
         id_offset: int = 0,
+        undirected: bool = True,
         sort_nodes: bool = True,
         names: Optional[SWCNames] = None,
     ) -> Tuple[pd.DataFrame, List[pd.DataFrame]]:
@@ -88,6 +112,9 @@ class LinesToTree(Transform[List[pd.DataFrame], Tree]):
             the swc.
         id_offset : int, default `0`
             The offset of the line node id.
+        undirected : bool, default `True`
+            Both ends of a line can be considered connection point. If
+            `False`, only the starting point.
         sort_nodes : bool, default `True`
             sort nodes of subtree.
         names : SWCNames, optional
@@ -97,11 +124,50 @@ class LinesToTree(Transform[List[pd.DataFrame], Tree]):
         tree : ~pandas.DataFrame
         remaining_lines : List of ~pandas.DataFrame
         """
-        return try_assemble_lines_impl(
-            lines,
-            undirected=self.undirected,
-            thre=self.thre,
-            id_offset=id_offset,
-            sort_nodes=sort_nodes,
-            names=names,
-        )
+
+        names = get_names(names)
+        lines = copy(list(lines))
+
+        tree = lines[0]
+        tree[names.id] = id_offset + np.arange(len(tree))
+        tree[names.pid] = tree[names.id] - 1
+        tree.at[0, names.pid] = -1
+        del lines[0]
+
+        while True:
+            for i, line in enumerate(lines):
+                for p in [0, -1] if undirected else [0]:
+                    xyz = [names.x, names.y, names.z]
+                    vs = tree[xyz] - line.iloc[p][xyz]
+                    dis = np.linalg.norm(vs, axis=1)
+                    ind = np.argmin(dis)
+                    if dis[ind] > self.thre:
+                        continue
+
+                    if dis[ind] < EPS:
+                        line = line.drop((p + len(line)) % len(line)).reset_index(
+                            drop=True
+                        )
+
+                    line[names.id] = id_offset + len(tree) + np.arange(len(line))
+                    line[names.pid] = line[names.id] + (-1 if p == 0 else 1)
+                    line.at[(p + len(line)) % len(line), names.pid] = tree.iloc[ind][
+                        names.id
+                    ]
+                    tree = pd.concat([tree, line])
+                    del lines[i]
+                    break
+                else:
+                    continue
+
+                break
+            else:
+                break
+
+        if sort_nodes:
+            sort_nodes_(tree)
+
+        return tree, lines
+
+    def extra_repr(self):
+        return f"thre={self.thre}, undirected={self.undirected}"

swcgeom/utils/__init__.py

@@ -4,10 +4,10 @@ from swcgeom.utils.debug import *
 from swcgeom.utils.dsu import *
 from swcgeom.utils.ellipse import *
 from swcgeom.utils.file import *
-from swcgeom.utils.geometry_object import *
 from swcgeom.utils.neuromorpho import *
 from swcgeom.utils.numpy_helper import *
 from swcgeom.utils.renderer import *
 from swcgeom.utils.sdf import *
 from swcgeom.utils.solid_geometry import *
 from swcgeom.utils.transforms import *
+from swcgeom.utils.volumetric_object import *