swcgeom 0.19.4__cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

swcgeom/utils/plotter_2d.py

@@ -0,0 +1,134 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""2D Plotting utils."""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import numpy.typing as npt
+from matplotlib import cm
+from matplotlib.axes import Axes
+from matplotlib.collections import LineCollection, PatchCollection
+from matplotlib.colors import Colormap, Normalize
+from matplotlib.figure import Figure
+from matplotlib.patches import Circle
+
+from swcgeom.utils.renderer import Camera
+from swcgeom.utils.transforms import to_homogeneous, translate3d
+
+__all__ = ["draw_lines", "draw_direction_indicator", "draw_circles", "get_fig_ax"]
+
+
+def draw_lines(
+    ax: Axes,
+    lines: npt.NDArray[np.floating],
+    camera: Camera,
+    joinstyle="round",
+    capstyle="round",
+    **kwargs,
+) -> LineCollection:
+    """Draw lines.
+
+    Args:
+        ax: The plot axes.
+        lines: A collection of coords of lines
+            Excepting a ndarray of shape (N, 2, 3), the axis-2 holds two points,
+            and the axis-3 holds the coordinates (x, y, z).
+        camera: Camera position.
+        **kwargs: Forwarded to `~matplotlib.collections.LineCollection`.
+    """
+    T = camera.MVP
+    T = translate3d(*camera.position).dot(T)  # keep origin
+
+    starts, ends = lines[:, 0], lines[:, 1]
+    starts, ends = to_homogeneous(starts, 1), to_homogeneous(ends, 1)
+    starts, ends = np.dot(T, starts.T).T[:, 0:2], np.dot(T, ends.T).T[:, 0:2]
+
+    edges = np.stack([starts, ends], axis=1)
+    collection = LineCollection(edges, joinstyle=joinstyle, capstyle=capstyle, **kwargs)  # type: ignore
+    return ax.add_collection(collection)  # type: ignore
+
+
+def draw_direction_indicator(
+    ax: Axes, camera: Camera, loc: tuple[float, float]
+) -> None:
+    x, y = loc
+    direction = camera.MV.dot(
+        [
+            [1, 0, 0, 1],
+            [0, 1, 0, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+        ]
+    )
+
+    arrow_length, text_offset = 0.05, 0.05  # TODO: may still overlap
+    text_colors = [["x", "red"], ["y", "green"], ["z", "blue"]]
+    for (dx, dy, dz, _), (text, color) in zip(direction, text_colors):
+        if 1 - abs(dz) < 1e-5:
+            continue
+
+        ax.arrow(
+            x,
+            y,
+            arrow_length * dx,
+            arrow_length * dy,
+            head_length=0.02,
+            head_width=0.01,
+            color=color,
+            transform=ax.transAxes,
+        )
+
+        ax.text(
+            x + (arrow_length + text_offset) * dx,
+            y + (arrow_length + text_offset) * dy,
+            text,
+            color=color,
+            transform=ax.transAxes,
+            horizontalalignment="center",
+            verticalalignment="center",
+        )
+
+
+def draw_circles(
+    ax: Axes,
+    x: npt.NDArray,
+    y: npt.NDArray,
+    *,
+    y_min: float | None = None,
+    y_max: float | None = None,
+    cmap: str | Colormap = "viridis",
+) -> PatchCollection:
+    """Draw a sequential of circles."""
+
+    y_min = y.min() if y_min is None else y_min
+    y_max = y.max() if y_max is None else y_max
+    norm = Normalize(y_min, y_max)
+
+    color_map = cmap if isinstance(cmap, Colormap) else cm.get_cmap(name=cmap)
+    colors = color_map(norm(y))
+
+    circles = [
+        Circle((0, 0), xi, color=color) for xi, color in reversed(list(zip(x, colors)))
+    ]
+    patches = PatchCollection(circles, match_original=True)
+    patches.set_cmap(color_map)
+    patches.set_norm(norm)
+    patches: PatchCollection = ax.add_collection(patches)  # type: ignore
+
+    ax.set_aspect(1)
+    ax.autoscale()
+    return patches
+
+
+def get_fig_ax(
+    fig: Figure | None = None, ax: Axes | None = None
+) -> tuple[Figure, Axes]:
+    if fig is None and ax is not None:
+        fig = ax.get_figure()
+        assert fig is not None, "expecting a figure from the axes"
+
+    fig = fig or plt.gcf()
+    ax = ax or plt.gca()
+    return fig, ax

swcgeom/utils/plotter_3d.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""3D Plotting utils."""
+
+import numpy as np
+import numpy.typing as npt
+from mpl_toolkits.mplot3d import Axes3D
+from mpl_toolkits.mplot3d.art3d import Line3DCollection
+
+__all__ = ["draw_lines_3d"]
+
+
+def draw_lines_3d(
+    ax: Axes3D,
+    lines: npt.NDArray[np.floating],
+    joinstyle="round",
+    capstyle="round",
+    **kwargs,
+):
+    """Draw lines.
+
+    Args:
+        ax: The plot axes.
+        lines: A collection of coords of lines
+            Excepting a ndarray of shape (N, 2, 3), the axis-2 holds two points,
+            and the axis-3 holds the coordinates (x, y, z).
+        **kwargs: Forwarded to `~mpl_toolkits.mplot3d.art3d.Line3DCollection`.
+    """
+
+    line_collection = Line3DCollection(
+        lines, joinstyle=joinstyle, capstyle=capstyle, **kwargs
+    )
+    return ax.add_collection3d(line_collection)

swcgeom/utils/renderer.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Rendering related utils."""
+
+from functools import cached_property
+from typing import Literal, cast
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import Self
+
+from swcgeom.utils.transforms import (
+    Vec3f,
+    model_view_transformation,
+    orthographic_projection_simple,
+)
+
+__all__ = ["CameraOptions", "Camera", "SimpleCamera", "palette"]
+
+CameraOption = Vec3f | tuple[Vec3f, Vec3f] | tuple[Vec3f, Vec3f, Vec3f]
+CameraPreset = Literal["xy", "yz", "zx", "yx", "zy", "xz"]
+CameraPresets: dict[CameraPreset, tuple[Vec3f, Vec3f, Vec3f]] = {
+    "xy": ((0.0, 0.0, 0.0), (+0.0, +0.0, -1.0), (+0.0, +1.0, +0.0)),
+    "yz": ((0.0, 0.0, 0.0), (-1.0, +0.0, +0.0), (+0.0, +0.0, +1.0)),
+    "zx": ((0.0, 0.0, 0.0), (+0.0, -1.0, +0.0), (+1.0, +0.0, +0.0)),
+    "yx": ((0.0, 0.0, 0.0), (+0.0, +0.0, -1.0), (+0.0, -1.0, +0.0)),
+    "zy": ((0.0, 0.0, 0.0), (-1.0, +0.0, +0.0), (+0.0, +0.0, -1.0)),
+    "xz": ((0.0, 0.0, 0.0), (+0.0, -1.0, +0.0), (-1.0, +0.0, +0.0)),
+}
+CameraOptions = CameraOption | CameraPreset
+
+
+class Camera:
+    _position: Vec3f
+    _look_at: Vec3f
+    _up: Vec3f
+
+    @property
+    def position(self) -> Vec3f:
+        return self._position
+
+    @property
+    def look_at(self) -> Vec3f:
+        return self._look_at
+
+    @property
+    def up(self) -> Vec3f:
+        return self._up
+
+    @property
+    def MV(self) -> npt.NDArray[np.float32]:
+        raise NotImplementedError()
+
+    @property
+    def P(self) -> npt.NDArray[np.float32]:
+        raise NotImplementedError()
+
+    @property
+    def MVP(self) -> npt.NDArray[np.float32]:
+        return self.P.dot(self.MV)
+
+
+class SimpleCamera(Camera):
+    """Simplest camera."""
+
+    def __init__(self, position: Vec3f, look_at: Vec3f, up: Vec3f):
+        self._position = position
+        self._look_at = look_at
+        self._up = up
+
+    @cached_property
+    def MV(self) -> npt.NDArray[np.float32]:  # pylint: disable=invalid-name
+        return model_view_transformation(self.position, self.look_at, self.up)
+
+    @cached_property
+    def P(self) -> npt.NDArray[np.float32]:  # pylint: disable=invalid-name
+        return orthographic_projection_simple()
+
+    @classmethod
+    def from_options(cls, camera: CameraOptions) -> Self:
+        if isinstance(camera, str):
+            return cls(*CameraPresets[camera])
+
+        if len(camera) == 2:
+            return cls((0, 0, 0), camera[0], camera[1])
+
+        if isinstance(camera[0], tuple):
+            return cls((0, 0, 0), cast(Vec3f, camera), (0, 1, 0))
+
+        return cls(*cast(tuple[Vec3f, Vec3f, Vec3f], camera))
+
+
+class Palette:
+    """The palette provides default and vaa3d color matching."""
+
+    # pylint: disable=too-few-public-methods
+
+    default: dict[int, str]
+    vaa3d: dict[int, str]
+
+    def __init__(self):
+        default = [
+            "#F596AA",  # momo,
+            "#867835",  # kimirucha,
+            "#E2943B",  # kuchiba,
+            "#00896C",  # aotake,
+            "#B9887D",  # mizugaki,
+            "#2EA9DF",  # tsuyukusa,
+            "#66327C",  # sumire,
+            "#52433D",  # benikeshinezumi,
+        ]
+        self.default = dict(enumerate(default))
+
+        vaa3d = [
+            "#ffffff",  # white, 0-undefined
+            "#141414",  # black, 1-soma
+            "#c81400",  # red, 2-axon
+            "#0014c8",  # blue, 3-dendrite
+            "#c800c8",  # purple, 4-apical dendrite
+            # the following is Hanchuan’s extended color. 090331
+            "#00c8c8",  # cyan, 5
+            "#dcc800",  # yellow, 6
+            "#00c814",  # green, 7
+            "#bc5e25",  # coffee, 8
+            "#b4c878",  # asparagus, 9
+            "#fa6478",  # salmon, 10
+            "#78c8c8",  # ice, 11
+            "#6478c8",  # orchid, 12
+            # the following is Hanchuan’s further extended color. 111003
+            "#ff80a8",  # 13
+            "#80ffa8",  # 14
+            "#80a8ff",  # 15
+            "#a8ff80",  # 16
+            "#ffa880",  # 17
+            "#a880ff",  # 18
+            "#000000",  # 19 # totally black. PHC, 2012-02-15
+            # the following (20-275) is used for matlab heat map. 120209 by WYN
+            "#000083",
+        ]
+        self.vaa3d = dict(enumerate(vaa3d))
+
+
+palette = Palette()

swcgeom/utils/sdf.py

@@ -0,0 +1,324 @@
+
+# SPDX-FileCopyrightText: 2022 - 2025 Zexin Yuan <pypi@yzx9.xyz>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Signed distance functions.
+
+Refs: https://iquilezles.org/articles/distfunctions/
+
+NOTE: This module has been deprecated since v0.14.0, and will be removed in the future,
+use `sdflit` instead.
+"""
+
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import deprecated
+
+from swcgeom.utils.solid_geometry import project_vector_on_plane
+
+__all__ = [
+    "SDF",
+    "SDFUnion",
+    "SDFIntersection",
+    "SDFDifference",
+    "SDFCompose",
+    "SDFSphere",
+    "SDFFrustumCone",
+    "SDFRoundCone",
+]
+
+# Axis-aligned bounding box, tuple of array of shape (3,)
+AABB = tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]
+
+
+class SDF(ABC):
+    """Signed distance functions."""
+
+    bounding_box: AABB | None = None
+
+    def __call__(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        return self.distance(p)
+
+    @abstractmethod
+    def distance(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        """Calculate signed distance.
+
+        Args:
+            p: Hit point p of shape (N, 3).
+
+        Returns:
+            distance: Distance array of shape (3,).
+        """
+        raise NotImplementedError()
+
+    def is_in(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.bool_]:
+        p = np.array(p, dtype=np.float32)
+        assert p.ndim == 2 and p.shape[1] == 3, "p should be array of shape (N, 3)"
+
+        in_box = self.is_in_bounding_box(p)
+        flags = np.full((p.shape[0]), False, dtype=np.bool_)
+        flags[in_box] = self.distance(p[in_box]) <= 0
+        return flags
+
+    def is_in_bounding_box(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.bool_]:
+        """Is p in bounding box.
+
+        Returns:
+            is_in: Array of shape (N,).
+                If bounding box is `None`, `True` will be returned.
+        """
+
+        if self.bounding_box is None:
+            return np.full((p.shape[0]), True, dtype=np.bool_)
+
+        is_in = np.logical_and(
+            np.all(p >= self.bounding_box[0], axis=1),
+            np.all(p <= self.bounding_box[1], axis=1),
+        )
+        return is_in
+
+
+class SDFUnion(SDF):
+    """Union multiple SDFs."""
+
+    def __init__(self, *sdfs: SDF) -> None:
+        assert len(sdfs) != 0, "must combine at least one SDF"
+        super().__init__()
+
+        self.sdfs = sdfs
+
+        bounding_boxes = [sdf.bounding_box for sdf in sdfs if sdf.bounding_box]
+        if len(bounding_boxes) == len(self.sdfs):
+            self.bounding_box = (
+                np.min(np.stack([box[0] for box in bounding_boxes]).T, axis=1),
+                np.max(np.stack([box[1] for box in bounding_boxes]).T, axis=1),
+            )
+
+    def distance(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        return np.min([sdf(p) for sdf in self.sdfs], axis=0)
+
+    def is_in(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.bool_]:
+        p = np.array(p, dtype=np.float32)
+        assert p.ndim == 2 and p.shape[1] == 3, "p should be array of shape (N, 3)"
+
+        in_box = self.is_in_bounding_box(p)
+        p_in_box = p[in_box]
+        is_in = np.stack([sdf.is_in(p_in_box) for sdf in self.sdfs])
+        flags = np.full_like(in_box, False, dtype=np.bool_)
+        flags[in_box] = np.any(is_in, axis=0)
+        return flags
+
+
+class SDFIntersection(SDF):
+    def __init__(self, *sdfs: SDF) -> None:
+        assert len(sdfs) != 0, "must intersect at least one SDF"
+        super().__init__()
+        self.sdfs = sdfs
+
+        bounding_boxes = [sdf.bounding_box for sdf in self.sdfs if sdf.bounding_box]
+        if len(bounding_boxes) == len(self.sdfs):
+            self.bounding_box = (
+                np.max(np.stack([box[0] for box in bounding_boxes]).T, axis=1),
+                np.min(np.stack([box[1] for box in bounding_boxes]).T, axis=1),
+            )
+
+    def distance(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        distances = np.stack([sdf.distance(p) for sdf in self.sdfs])
+        return np.max(distances, axis=1)
+
+    def is_in(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.bool_]:
+        p = np.array(p, dtype=np.float32)
+        assert p.ndim == 2 and p.shape[1] == 3, "p should be array of shape (N, 3)"
+
+        in_box = self.is_in_bounding_box(p)
+        p_in_box = p[in_box]
+        is_in = np.stack([sdf.is_in(p_in_box) for sdf in self.sdfs])
+        flags = np.full_like(in_box, False, dtype=np.bool_)
+        flags[in_box] = np.all(is_in, axis=0)
+        return flags
+
+
+class SDFDifference(SDF):
+    """Difference of two SDFs A-B."""
+
+    def __init__(self, sdf_a: SDF, sdf_b: SDF) -> None:
+        super().__init__()
+        self.sdf_a = sdf_a
+        self.sdf_b = sdf_b
+
+        self.bounding_box = sdf_a.bounding_box
+
+    def distance(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        da = self.sdf_a.distance(p)
+        db = self.sdf_b.distance(p)
+        return np.maximum(da, -db)
+
+    def is_in(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.bool_]:
+        p = np.array(p, dtype=np.float32)
+        assert p.ndim == 2 and p.shape[1] == 3, "p should be array of shape (N, 3)"
+
+        in_box = self.is_in_bounding_box(p)
+        p_in_box = p[in_box]
+        is_in_a = self.sdf_a.is_in(p_in_box)
+        is_in_b = self.sdf_b.is_in(p_in_box)
+        flags = np.full_like(in_box, False, dtype=np.bool_)
+        flags[in_box] = np.logical_and(is_in_a, np.logical_not(is_in_b))
+        return flags
+
+
+@deprecated("Use `SDFUnion` instead")
+class SDFCompose(SDFUnion):
+    """Compose multiple SDFs.
+
+    .. deprecated:: 0.14.0
+        Use :cls:`SDFUnion` instead.
+    """
+
+    def __init__(self, sdfs: Iterable[SDF]) -> None:
+        sdfs = list(sdfs)
+        if len(sdfs) == 1:
+            warnings.warn("compose only one SDF, use SDFCompose.compose instead")
+
+        super().__init__(*sdfs)
+
+    @staticmethod
+    def compose(sdfs: Iterable[SDF]) -> SDF:
+        sdfs = list(sdfs)
+        return SDFCompose(sdfs) if len(sdfs) != 1 else sdfs[0]
+
+
+class SDFSphere(SDF):
+    """SDF of sphere."""
+
+    def __init__(self, center: npt.ArrayLike, radius: float) -> None:
+        super().__init__()
+        self.center = np.array(center)
+        self.radius = radius
+        assert tuple(self.center.shape) == (3,), "center should be vector of 3d"
+
+        self.bounding_box = (self.center - self.radius, self.center + self.radius)
+
+    def distance(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        return np.linalg.norm(p - self.center, axis=1) - self.radius
+
+
+class SDFFrustumCone(SDF):
+    """SDF of frustum cone."""
+
+    def __init__(
+        self, a: npt.ArrayLike, b: npt.ArrayLike, ra: float, rb: float
+    ) -> None:
+        super().__init__()
+        self.a = np.array(a)
+        self.b = np.array(b)
+        self.ra = ra
+        self.rb = rb
+        assert tuple(self.a.shape) == (3,), "a should be vector of 3d"
+        assert tuple(self.b.shape) == (3,), "b should be vector of 3d"
+
+        self.bounding_box = self.get_bounding_box()
+
+    def distance(self, p: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+        a, b, ra, rb = self.a, self.b, self.ra, self.rb
+
+        rba = rb - ra
+        baba = np.dot(b - a, b - a)
+        papa = np.einsum("ij,ij->i", p - a, p - a)
+        paba = np.dot(p - a, b - a) / baba
+        # maybe negative due to numerical error
+        x = np.sqrt(np.maximum(papa - paba * paba * baba, 0))
+        cax = np.maximum(0.0, x - np.where(paba < 0.5, ra, rb))
+        cay = np.abs(paba - 0.5) - 0.5
+        k = rba * rba + baba
+        f = np.clip((rba * (x - ra) + paba * baba) / k, 0.0, 1.0)
+        cbx = x - ra - f * rba
+        cby = paba - f
+        s = np.where(np.logical_and(cbx < 0.0, cay < 0.0), -1.0, 1.0)
+        return s * np.sqrt(
+            np.minimum(cax * cax + cay * cay * baba, cbx * cbx + cby * cby * baba)
+        )
+
+    def get_bounding_box(self) -> AABB | None:
+        a, b, ra, rb = self.a, self.b, self.ra, self.rb
+        up = a - b
+        vx = project_vector_on_plane((1, 0, 0), up)
+        vy = project_vector_on_plane((0, 1, 0), up)
+        vz = project_vector_on_plane((0, 0, 1), up)
+        a1 = a - ra * vx - ra * vy - ra * vz
+        a2 = a + ra * vx + ra * vy + ra * vz
+        b1 = b - rb * vx - rb * vy - rb * vz
+        b2 = b + rb * vx + rb * vy + rb * vz
+        return (
+            np.minimum(a1, b1).astype(np.float32),
+            np.maximum(a2, b2).astype(np.float32),
+        )
+
+
+class SDFRoundCone(SDF):
+    """Round cone is made up of two balls and a cylinder."""
+
+    def __init__(
+        self, a: npt.ArrayLike, b: npt.ArrayLike, ra: float, rb: float
+    ) -> None:
+        """SDF of round cone.
+
+        Args:
+            a, b: Coordinates of point A/B of shape (3,).
+            ra, rb: Radius of point A/B.
+        """
+
+        self.a = np.array(a, dtype=np.float32)
+        self.b = np.array(b, dtype=np.float32)
+        self.ra = ra
+        self.rb = rb
+
+        assert tuple(self.a.shape) == (3,), "a should be vector of 3d"
+        assert tuple(self.b.shape) == (3,), "b should be vector of 3d"
+
+        self.bounding_box = (
+            np.min([self.a - self.ra, self.b - self.rb], axis=0).astype(np.float32),
+            np.max([self.a + self.ra, self.b + self.rb], axis=0).astype(np.float32),
+        )
+
+    def distance(self, p: npt.ArrayLike) -> npt.NDArray[np.float32]:
+        # pylint: disable=too-many-locals
+        p = np.array(p, dtype=np.float32)
+        assert p.ndim == 2 and p.shape[1] == 3, "p should be array of shape (N, 3)"
+
+        a = self.a
+        b = self.b
+        ra = self.ra
+        rb = self.rb
+
+        # sampling independent computations (only depend on shape)
+        ba = b - a
+        l2 = np.dot(ba, ba)
+        rr = ra - rb
+        a2 = l2 - rr * rr
+        il2 = 1.0 / l2
+
+        # sampling dependant computations
+        pa = p - a
+        y = np.dot(pa, ba)
+        z = y - l2
+        x = pa * l2 - np.outer(y, ba)
+        x2 = np.sum(x * x, axis=1)
+        y2 = y * y * l2
+        z2 = z * z * l2
+
+        # single square root!
+        k = np.sign(rr) * rr * rr * x2
+        dis = (np.sqrt(x2 * a2 * il2) + y * rr) * il2 - ra
+
+        lt = np.sign(z) * a2 * z2 > k
+        dis[lt] = np.sqrt(x2[lt] + z2[lt]) * il2 - rb
+
+        rt = np.sign(y) * a2 * y2 < k
+        dis[rt] = np.sqrt(x2[rt] + y2[rt]) * il2 - ra
+
+        return dis