"""Routing geometry primitives: the voxel frame (world/grid chart) and shapes that voxelize themselves into it."""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import cached_property
from itertools import pairwise
import numpy as np
from scipy import ndimage
from volmdlr import Frame3D, Point3D
from volmdlr.discrete_representation import MatrixBasedVoxelization
from volmdlr.primitives3d import Cylinder as CadCylinder
from volmdlr.shapes import Solid
from ._types import Coord3D, Index3D
from .functions import clipped_grid_window, cylinder_bounding_box, grid_window
from .functions_compiled import bresenham_3d_c
_MIN_POLYLINE_POINTS = 2 # below this, voxelize/smoothing return early
_MIN_SMOOTHABLE_POINTS = 2 # below this, smoothing has no interior point to average
_DEGENERATE_SEGMENT = 1e-6 # segment length below which two polyline points are the same place
_HALF_VOXEL = 0.5 # in fractional indices, the distance from a voxel centre to its faces
[docs]
@dataclass(frozen=True)
class VoxelFrame:
"""Where a voxel grid sits in space: its origin, its step and its extent.
The one chart between world coordinates and grid indices, shared by everything shaped
like a grid — a voxelization's matrix, a pathfinding grid's nodes, a region's fields.
Its distinct roundings are distinct operations: the voxel nearest to a point
(:meth:`nearest_indices`), the exact fractional position (:meth:`fractional_indices`),
and the covering bounds of a world box (:meth:`floor_indices` / :meth:`ceil_indices`).
:param origin: The centre of voxel (0, 0, 0) in world coordinates.
:param voxel_size: The edge length of a voxel, in world units.
:param shape: The grid extent in voxels per axis.
"""
origin: Coord3D
voxel_size: float
shape: Index3D
[docs]
@classmethod
def of(cls, voxelization: MatrixBasedVoxelization) -> "VoxelFrame":
"""Read the frame a voxelization is sampled on."""
return cls(
tuple(float(coordinate) for coordinate in voxelization.min_grid_center),
float(voxelization.voxel_size),
tuple(int(side) for side in voxelization.matrix.shape),
)
[docs]
def fractional_indices(self, points: np.ndarray) -> np.ndarray:
"""Exact grid position of world ``points``, ``(3,)`` or ``(N, 3)``, without rounding to a voxel."""
return (np.asarray(points, dtype=float) - self.origin) / self.voxel_size
[docs]
def nearest_indices(self, points: np.ndarray) -> np.ndarray:
"""Convert world ``points``, ``(3,)`` or ``(N, 3)``, to the indices of the voxels containing them."""
return np.rint(self.fractional_indices(points)).astype(int)
[docs]
def nearest_index(self, point: Coord3D | np.ndarray) -> Index3D:
"""Index of the single voxel containing ``point``, as a plain int tuple."""
index = self.nearest_indices(point)
return (int(index[0]), int(index[1]), int(index[2]))
[docs]
def floor_indices(self, points: np.ndarray) -> np.ndarray:
"""Highest voxel indices at or below world ``points`` — the lower bound of a covering window."""
return np.floor(self.fractional_indices(points)).astype(int)
[docs]
def ceil_indices(self, points: np.ndarray) -> np.ndarray:
"""Lowest voxel indices at or above world ``points`` — the upper bound of a covering window."""
return np.ceil(self.fractional_indices(points)).astype(int)
[docs]
def contains(self, points: np.ndarray) -> np.ndarray:
"""Whether world ``points``, ``(3,)`` or ``(N, 3)``, fall inside the grid's extent — the outer faces of its boundary voxels."""
fractional = self.fractional_indices(np.atleast_2d(points))
return np.all((fractional >= -_HALF_VOXEL) & (fractional <= np.asarray(self.shape) - _HALF_VOXEL), axis=1)
[docs]
def cylindrical_window(
self, start: np.ndarray, end: np.ndarray, half_width: float
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Dense grid window around the ``[start, end]`` segment, in cylindrical coordinates.
Covers the axis-aligned box enclosing the segment inflated by ``half_width``, clipped
to the grid. Radial distances and abscissas are separable along the grid axes, so no
per-voxel point array is ever built.
:return: Tuple ``(lower (3,) int, radial_distances (ni, nj, nk), abscissas (ni, nj, nk))``
— the window's lower corner index and, per window voxel, the world distance to the
segment's supporting line and the world abscissa along it (0 at ``start``). Empty
arrays when the window misses the grid.
"""
start = np.asarray(start, dtype=float)
corners = np.stack([start, np.asarray(end, dtype=float)])
lower = np.maximum(self.floor_indices(corners.min(axis=0) - half_width), 0)
upper = np.minimum(self.ceil_indices(corners.max(axis=0) + half_width) + 1, np.asarray(self.shape))
if np.any(lower >= upper):
empty = np.empty((0, 0, 0))
return lower, empty, empty
direction = corners[1] - corners[0]
direction /= np.linalg.norm(direction)
relative_axes = [
self.origin[axis] + np.arange(lower[axis], upper[axis]) * self.voxel_size - start[axis] for axis in range(3)
]
abscissas = (
relative_axes[0][:, None, None] * direction[0]
+ relative_axes[1][None, :, None] * direction[1]
+ relative_axes[2][None, None, :] * direction[2]
)
squared_norms = (
relative_axes[0][:, None, None] ** 2
+ relative_axes[1][None, :, None] ** 2
+ relative_axes[2][None, None, :] ** 2
)
radial_distances = np.sqrt(np.maximum(squared_norms - abscissas**2, 0.0))
return lower, radial_distances, abscissas
[docs]
def indices_around(self, center: np.ndarray, radius: float) -> np.ndarray:
"""List the grid's voxels within ``radius`` (in voxels) of the fractional index ``center``, in C order."""
center = np.asarray(center, dtype=float)
lower = np.maximum(np.floor(center - radius), 0).astype(int)
upper = np.minimum(np.ceil(center + radius), np.asarray(self.shape) - 1).astype(int)
if np.any(lower > upper):
return np.empty((0, 3), dtype=int)
axes = [np.arange(low, high + 1) for low, high in zip(lower, upper, strict=True)]
indices = np.stack(np.meshgrid(*axes, indexing="ij"), axis=-1).reshape(-1, 3)
return indices[np.linalg.norm(indices - center, axis=1) <= radius]
[docs]
def centers(self, indices: np.ndarray) -> np.ndarray:
"""World centres of the voxels at integer or fractional ``indices``, ``(3,)`` or ``(N, 3)``."""
return np.asarray(self.origin) + np.asarray(indices, dtype=float) * self.voxel_size
[docs]
def cells_covered_by(self, coarser: "VoxelFrame", indices: np.ndarray) -> np.ndarray:
"""Give the cells of this frame whose centre lies in the voxels ``indices`` ``(N, 3)`` of the ``coarser`` frame.
The multi-scale map anchors its grids on one centre and refines by an odd factor, so a coarse voxel
covers a block of cells centred on its own centre; a frame covers itself cell for cell. Not clipped.
"""
if not len(indices):
return np.empty((0, 3), dtype=int)
centres = self.nearest_indices(coarser.centers(indices))
half_span = round((coarser.voxel_size / self.voxel_size - 1) / 2)
if half_span == 0:
return centres
offsets = np.mgrid[-half_span : half_span + 1, -half_span : half_span + 1, -half_span : half_span + 1]
return (centres[:, None, :] + offsets.reshape(3, -1).T[None, :, :]).reshape(-1, 3)
[docs]
def center(self, index: Index3D | np.ndarray) -> Coord3D:
"""World centre of the single voxel at ``index``, as a plain float tuple."""
center = self.centers(index)
return (float(center[0]), float(center[1]), float(center[2]))
[docs]
class RoutingShape(ABC):
"""A class to handle shapes in routing context."""
[docs]
@abstractmethod
def bounding_box(self) -> np.ndarray:
"""Axis-aligned bounding box as ``[minimums, maximums]``."""
[docs]
@abstractmethod
def voxelize_into(
self,
volume: np.ndarray,
frame: VoxelFrame,
dilation_radius: float = 0,
) -> None:
"""Add this shape's voxels to ``volume``, sampled on ``frame``, touching only the cells the shape occupies."""
[docs]
def voxelize(self, frame: VoxelFrame, dilation_radius: float = 0) -> np.ndarray:
"""Boolean mask of the ``frame`` voxels intersecting the shape, on a grid of its own."""
volume = np.zeros(frame.shape, dtype=bool)
self.voxelize_into(volume, frame, dilation_radius)
return volume
@staticmethod
def _stamp_thickened(volume: np.ndarray, indices: np.ndarray, ball_radius: int) -> None:
"""
Add to ``volume`` every voxel within ``ball_radius`` of ``indices``, working in their own window.
:param volume: The grid to write into.
:param indices: The ``(N, 3)`` voxel indices to thicken, all inside ``volume``.
:param ball_radius: Thickening radius in voxels; ``0`` stamps ``indices`` as they are.
"""
window_low = np.maximum(indices.min(axis=0) - ball_radius, 0)
window_high = np.minimum(indices.max(axis=0) + ball_radius + 1, volume.shape)
thickened = np.zeros(window_high - window_low, dtype=bool)
local_indices = indices - window_low
thickened[local_indices[:, 0], local_indices[:, 1], local_indices[:, 2]] = True
if ball_radius > 0:
thickened = ndimage.distance_transform_edt(~thickened) <= ball_radius
volume[grid_window(window_low, window_high)] |= thickened
[docs]
@dataclass
class Box(RoutingShape):
"""A class to handle boxes in routing context."""
minimums: np.ndarray # min_x, min_y, min_z
maximums: np.ndarray # max_x, max_y, max_z
@cached_property
def size(self) -> np.ndarray:
"""Box extent along each axis."""
return self.maximums - self.minimums
@cached_property
def size_x(self) -> float:
"""Box extent along x."""
return self.size[0]
@cached_property
def size_y(self) -> float:
"""Box extent along y."""
return self.size[1]
@cached_property
def size_z(self) -> float:
"""Box extent along z."""
return self.size[2]
@cached_property
def center(self) -> np.ndarray:
"""Box geometric center."""
return (self.minimums + self.maximums) / 2
@cached_property
def origin(self) -> np.ndarray:
"""Box origin (min corner)."""
return self.minimums.copy()
@cached_property
def bounding_box(self) -> np.ndarray:
"""Box vertices as ``[minimums, maximums]``."""
return np.array([self.minimums, self.maximums])
@cached_property
def points(self) -> np.ndarray:
"""All box's points."""
minimums, maximums = self.minimums, self.maximums
return np.array(
[
[minimums[0], minimums[1], minimums[2]],
[maximums[0], minimums[1], minimums[2]],
[minimums[0], maximums[1], minimums[2]],
[maximums[0], maximums[1], minimums[2]],
[minimums[0], minimums[1], maximums[2]],
[maximums[0], minimums[1], maximums[2]],
[minimums[0], maximums[1], maximums[2]],
[maximums[0], maximums[1], maximums[2]],
]
)
@cached_property
def volume(self) -> float:
"""The volume of the box."""
return float(np.prod(self.size))
[docs]
def contains(self, points: np.ndarray) -> np.ndarray:
"""Whether each of ``points`` ``(n, 3)`` lies in the box, its faces included."""
return np.all((points >= self.minimums) & (points <= self.maximums), axis=1)
[docs]
def meets(self, boxes: np.ndarray) -> np.ndarray:
"""Whether each of ``boxes``, ``(n, 2, 3)`` lowest and highest corners, shares a point with this box."""
return np.all((boxes[:, 0] <= self.maximums) & (boxes[:, 1] >= self.minimums), axis=1)
[docs]
def grown(self, margin: float | np.ndarray) -> "Box":
"""Give the box widened by ``margin`` on each of its six faces, one value per axis or one for all."""
return Box(self.minimums - margin, self.maximums + margin)
[docs]
@classmethod
def from_bounding_box(cls, bounding_box: np.ndarray) -> "Box":
"""Create a box from its extreme vertices."""
return cls(bounding_box[0], bounding_box[1])
[docs]
@classmethod
def from_points(cls, points: np.ndarray) -> "Box":
"""Create a box from points."""
return cls(np.min(points, axis=0), np.max(points, axis=0))
[docs]
@classmethod
def cubes_at(
cls,
centers: np.ndarray,
edge: float,
color: tuple = (0.5, 0.5, 0.5),
alpha: float = 1.0,
) -> list[Solid]:
"""
Render one cube solid of side ``edge`` centred on each of ``centers``.
:param centers: World centres, as any array of ``(x, y, z)`` rows.
:param edge: The side of every cube.
"""
half_edge = edge / 2
return [
cls(centre - half_edge, centre + half_edge).volmdlr_primitive(color, alpha)
for centre in np.asarray(centers, dtype=float).reshape(-1, 3)
]
[docs]
def volmdlr_primitive(self, color: tuple = (0.5, 0.5, 0.5), alpha: float = 1.0) -> Solid:
"""Render this box as a volmdlr CAD solid."""
primitive = Solid.make_box(
*(float(side) for side in self.size),
frame=Frame3D.from_origin_xyz(Point3D(*self.center)),
frame_centered=True,
)
primitive.color = color
primitive.alpha = alpha
return primitive
[docs]
def voxelize_into(
self,
volume: np.ndarray,
frame: VoxelFrame,
dilation_radius: float = 0,
) -> None:
"""
Fill the index block the box spans, widened by ``dilation_radius`` and clipped to the grid.
:param volume: The grid to write into.
:param frame: The chart ``volume`` is sampled on.
:param dilation_radius: Extra half-width in voxel units (expands the filled index range).
"""
grid_shape = np.array(volume.shape)
dilatation_offset = round(float(dilation_radius))
lower_bounds = frame.floor_indices(self.minimums) - dilatation_offset
upper_bounds = frame.ceil_indices(self.maximums) + dilatation_offset
volume[clipped_grid_window(lower_bounds, np.minimum(upper_bounds, grid_shape - 1) + 1, grid_shape)] = True
[docs]
@dataclass
class OrientedBox(RoutingShape):
"""A box lying in a frame of its own: the tightest box around points, seen from there.
An oriented box is an aligned one seen from its frame, so it carries that frame and the :class:`Box`
it spans in it. A shape lying along a diagonal keeps a tight box, where a world-aligned one would
swallow the space around it.
:param origin: Where the frame sits in the world.
:param axes: The frame's three unit axes, as columns.
:param extent: The box the shape spans in that frame.
"""
origin: np.ndarray
axes: np.ndarray
extent: Box
[docs]
@classmethod
def from_points(cls, points: np.ndarray, origin: np.ndarray, axes: np.ndarray) -> "OrientedBox":
"""Build the tightest box around ``points`` in the frame ``(origin, axes)``."""
return cls(origin, axes, Box.from_points((points - origin) @ axes))
[docs]
def local_coordinates(self, points: np.ndarray) -> np.ndarray:
"""Give ``points`` in the box's own frame."""
return (points - self.origin) @ self.axes
[docs]
def contains(self, points: np.ndarray) -> np.ndarray:
"""Whether each of ``points`` ``(n, 3)`` lies in the box, its faces included."""
return self.extent.contains(self.local_coordinates(points))
[docs]
def grown(self, margin: float | np.ndarray) -> "OrientedBox":
"""Give the box widened by ``margin`` on each of its six faces, one value per axis or one for all."""
return OrientedBox(self.origin, self.axes, self.extent.grown(margin))
@cached_property
def points(self) -> np.ndarray:
"""The box's eight vertices, in the world."""
return self.extent.points @ self.axes.T + self.origin
[docs]
def bounding_box(self) -> np.ndarray:
"""Axis-aligned bounding box of the eight vertices, as ``[minimums, maximums]``."""
return np.array([self.points.min(axis=0), self.points.max(axis=0)])
[docs]
def cells(self, frame: VoxelFrame, dilation_radius: float = 0) -> np.ndarray:
"""Give the cells of ``frame`` whose centre lies in the box, widened by ``dilation_radius`` voxels.
The candidates are read off the aligned block the box spans, so an oblique box costs its own
neighbourhood and not the grid.
"""
widened = self.grown(dilation_radius * frame.voxel_size) if dilation_radius else self
corners = widened.bounding_box()
window = clipped_grid_window(
frame.floor_indices(corners[0]), frame.ceil_indices(corners[1]) + 1, np.array(frame.shape)
)
candidates = np.mgrid[window].reshape(3, -1).T
return candidates[widened.contains(frame.centers(candidates))] if len(candidates) else candidates
[docs]
def voxelize_into(self, volume: np.ndarray, frame: VoxelFrame, dilation_radius: float = 0) -> None:
"""Fill the cells whose centre lies in the box, widened by ``dilation_radius`` voxels."""
cells = self.cells(frame, dilation_radius)
if len(cells):
volume[tuple(cells.T)] = True
[docs]
@dataclass
class Cylinder(RoutingShape):
"""A class to handle cylinders in routing context."""
start: np.ndarray
end: np.ndarray
radius: float
@cached_property
def bounding_box(self) -> np.ndarray:
"""Get the bounding box of the cylinder."""
return cylinder_bounding_box(self.start, self.end, self.radius)
@cached_property
def length(self) -> float:
"""Cylinder axis length."""
return float(np.linalg.norm(self.end - self.start))
[docs]
def volmdlr_primitive(self, color: tuple = (0.5, 0.5, 0.5), alpha: float = 1.0) -> CadCylinder:
"""Render this cylinder as a volmdlr CAD primitive."""
primitive = CadCylinder.from_end_points(Point3D(*self.start), Point3D(*self.end), self.radius)
primitive.color = color
primitive.alpha = alpha
return primitive
[docs]
def fiber_indices(self, frame: VoxelFrame) -> np.ndarray:
"""Voxel-snapped Bresenham rasterization of the axis (neutral fiber), clipped to the grid."""
upper = np.array(frame.shape) - 1
start = np.clip(frame.nearest_indices(self.start), 0, upper)
end = np.clip(frame.nearest_indices(self.end), 0, upper)
return bresenham_3d_c(
tuple(int(coordinate) for coordinate in start), tuple(int(coordinate) for coordinate in end)
)
[docs]
def ball_radius(self, voxel_size: float, dilation_radius: float = 0) -> int:
"""Integer thickening radius in voxels: ``ceil(radius + dilation)``.
The thickening ball needs an INTEGER radius: a non-integer yields an even-sized footprint
with no centre voxel, so thickening shifts the tube half a voxel off the fiber and brings
the skeletonization vanish straight back.
"""
return int(np.ceil(self.radius / voxel_size + dilation_radius))
[docs]
def voxelize_into(
self,
volume: np.ndarray,
frame: VoxelFrame,
dilation_radius: float = 0,
) -> None:
"""Rasterize the axis (neutral fiber) via Bresenham, then thicken it into an exact-Euclidean-ball tube."""
fiber = self.fiber_indices(frame)
self._stamp_thickened(volume, fiber, self.ball_radius(frame.voxel_size, dilation_radius))
[docs]
@dataclass
class Polyline(RoutingShape):
"""An ordered sequence of points forming a broken line in 3D."""
points: np.ndarray # (N, 3)
@cached_property
def length(self) -> float:
"""Sum of segment lengths."""
return float(np.linalg.norm(np.diff(self.points, axis=0), axis=1).sum())
@cached_property
def bounding_box(self) -> np.ndarray:
"""Polyline AABB as ``[minimums, maximums]``."""
return np.array([self.points.min(axis=0), self.points.max(axis=0)])
[docs]
def voxelize_into(
self,
volume: np.ndarray,
frame: VoxelFrame,
dilation_radius: float = 0,
) -> None:
"""Rasterize each segment via Bresenham, then dilate over the polyline's own bounding window."""
if len(self.points) < _MIN_POLYLINE_POINTS:
return
grid_shape = np.array(volume.shape)
voxel_coords = np.clip(frame.nearest_indices(self.points), 0, grid_shape - 1)
crossed = np.concatenate(
[
bresenham_3d_c(tuple(int(c) for c in start), tuple(int(c) for c in end))
for start, end in pairwise(voxel_coords)
]
)
self._stamp_thickened(volume, crossed, round(float(dilation_radius)))
[docs]
def to_cylinders(self, radius: float = 1e-6) -> list["Cylinder"]:
"""
One cylinder per segment, all of ``radius``; degenerate segments are dropped.
:param radius: A polyline carries no thickness of its own, so this states the scale it
should be seen at; the default draws it as a hairline.
"""
return [
Cylinder(start, end, radius)
for start, end in pairwise(self.points)
if np.linalg.norm(end - start) > _DEGENERATE_SEGMENT
]
[docs]
def distances_to(self, points: np.ndarray) -> np.ndarray:
"""Min distance from each of ``points`` to this polyline."""
starts = self.points[:-1]
ends = self.points[1:]
edges = ends - starts
squared_lengths = np.einsum("ij,ij->i", edges, edges)
squared_lengths = np.where(squared_lengths > 0, squared_lengths, 1.0)
offsets = points[:, None, :] - starts[None, :, :]
parameters = np.clip(np.einsum("nsj,sj->ns", offsets, edges) / squared_lengths[None, :], 0.0, 1.0)
closest = starts[None, :, :] + parameters[:, :, None] * edges[None, :, :]
differences = points[:, None, :] - closest
return np.sqrt(np.einsum("nsj,nsj->ns", differences, differences).min(axis=1))
[docs]
def smoothed(self, window: int) -> "Polyline":
"""Moving-average smoothing with endpoints preserved."""
count = len(self.points)
if count <= _MIN_SMOOTHABLE_POINTS or window <= 1:
return Polyline(self.points.copy())
half_window = window // 2
smoothed = self.points.copy()
for index in range(1, count - 1):
lower = max(1, index - half_window)
upper = min(count - 1, index + half_window + 1)
smoothed[index] = self.points[lower:upper].mean(axis=0)
return Polyline(smoothed)
[docs]
def local_angles(self, tangent_window: int) -> np.ndarray:
"""Per-vertex bending angle measured over a ``tangent_window``-step neighborhood.
Returns an array of length ``len(points)``. NaN at vertices closer than
``tangent_window`` from the endpoints (no back/forward tangent available)
or where back/forward is degenerate (zero-length).
"""
count = len(self.points)
window = tangent_window
angles = np.full(count, np.nan, dtype=float)
if count < 2 * window + 1:
return angles
for index in range(window, count - window):
back = self.points[index] - self.points[index - window]
forward = self.points[index + window] - self.points[index]
back_length = float(np.linalg.norm(back))
forward_length = float(np.linalg.norm(forward))
if back_length == 0.0 or forward_length == 0.0:
continue
cosine = float(np.clip(np.dot(back, forward) / (back_length * forward_length), -1.0, 1.0))
angles[index] = np.arccos(cosine)
return angles
[docs]
def turn_indices(self, tangent_window: int, min_turn_angle: float) -> list[int]:
"""Interior indices where :meth:`local_angles` peaks above ``min_turn_angle``.
Detected turns are at least ``tangent_window`` apart so a single messy
corner yields a single cut.
"""
count = len(self.points)
window = tangent_window
if count < 2 * window + 1:
return []
angles = self.local_angles(window)
turns: list[int] = []
index = window
while index < count - window:
if not np.isnan(angles[index]) and angles[index] > min_turn_angle:
local = angles[max(window, index - window) : min(count - window, index + window + 1)]
if angles[index] == np.nanmax(local):
turns.append(index)
index += window
continue
index += 1
return turns
