"""Spatial types: points, orientations, poses, and transforms."""
from collections.abc import Sequence
from dataclasses import dataclass, field, replace
from datetime import datetime
from math import isfinite
from typing import overload
import numpy as np
from typing_extensions import Self
import olo.spatial._utils as _utils
_REPR_FLOAT_DP = 5
def _repr_float(value: float) -> str:
"""Format a coordinate for repr output (fixed decimal places, trim trailing zeros)."""
if value == 0.0:
return "0.0"
text = f"{value:.{_REPR_FLOAT_DP}f}".rstrip("0").rstrip(".")
return text or "0.0"
[docs]
@dataclass(frozen=True, slots=True, repr=False)
class Point:
"""Cartesian position in metres."""
x: float = 0.0
y: float = 0.0
z: float = 0.0
[docs]
@classmethod
def from_array(cls, values: Sequence[float]) -> "Point":
"""Create a point from an ``[x, y, z]`` array-like."""
arr = np.asarray(values, dtype=np.float64)
if arr.shape != (3,):
raise ValueError(f"Expected 3 values [x, y, z], got shape {arr.shape}")
return cls(x=float(arr[0]), y=float(arr[1]), z=float(arr[2]))
[docs]
def array(self) -> np.ndarray:
"""Return ``[x, y, z]`` as a float64 numpy array."""
return np.array([self.x, self.y, self.z], dtype=np.float64)
[docs]
def isclose(self, other: Self, *, tol: float = 1e-9) -> bool:
"""Return whether two points are equal within an absolute tolerance."""
if not isinstance(other, Point):
raise TypeError(f"Cannot compare Point with {type(other).__name__}; expected Point")
return bool(np.allclose(self.array(), other.array(), rtol=0.0, atol=tol))
def __add__(self, other: "Vector") -> "Point":
if not isinstance(other, Vector):
return NotImplemented
return Point(x=self.x + other.x, y=self.y + other.y, z=self.z + other.z)
def __sub__(self, other: Self) -> "Vector":
if not isinstance(other, Point):
return NotImplemented
return Vector(x=self.x - other.x, y=self.y - other.y, z=self.z - other.z)
def __repr__(self) -> str:
return f"Point(x={_repr_float(self.x)}, y={_repr_float(self.y)}, z={_repr_float(self.z)})"
[docs]
@dataclass(frozen=True, slots=True, repr=False)
class Vector:
"""Cartesian vector in metres."""
x: float = 0.0
y: float = 0.0
z: float = 0.0
[docs]
@classmethod
def from_array(cls, values: Sequence[float]) -> "Vector":
"""Create a vector from an ``[x, y, z]`` array-like."""
arr = np.asarray(values, dtype=np.float64)
if arr.shape != (3,):
raise ValueError(f"Expected 3 values [x, y, z], got shape {arr.shape}")
return cls(x=float(arr[0]), y=float(arr[1]), z=float(arr[2]))
[docs]
@classmethod
def from_point(cls, point: Point) -> "Vector":
"""Create a vector with the same components as a point."""
return cls(x=point.x, y=point.y, z=point.z)
[docs]
def as_point(self) -> Point:
"""Return a point with the same components."""
return Point(x=self.x, y=self.y, z=self.z)
[docs]
def array(self) -> np.ndarray:
"""Return ``[x, y, z]`` as a float64 numpy array."""
return np.array([self.x, self.y, self.z], dtype=np.float64)
[docs]
def isclose(self, other: Self, *, tol: float = 1e-9) -> bool:
"""Return whether two vectors are equal within an absolute tolerance."""
if not isinstance(other, Vector):
raise TypeError(f"Cannot compare Vector with {type(other).__name__}; expected Vector")
return bool(np.allclose(self.array(), other.array(), rtol=0.0, atol=tol))
[docs]
def magnitude(self) -> float:
"""Return the Euclidean norm."""
return float(np.linalg.norm(self.array()))
[docs]
def normalized(self) -> "Vector":
"""Return a unit vector."""
return self / self.magnitude()
[docs]
def dot(self, other: "Vector") -> float:
"""Return the dot product."""
if not isinstance(other, Vector):
raise TypeError(f"Cannot dot Vector with {type(other).__name__}; expected Vector")
return float(np.dot(self.array(), other.array()))
[docs]
def cross(self, other: "Vector") -> "Vector":
"""Return the cross product."""
if not isinstance(other, Vector):
raise TypeError(f"Cannot cross Vector with {type(other).__name__}; expected Vector")
x, y, z = np.cross(self.array(), other.array())
return Vector(x=float(x), y=float(y), z=float(z))
[docs]
def scale(self, scalar: float) -> "Vector":
"""Return this vector scaled by a scalar."""
return Vector(x=self.x * scalar, y=self.y * scalar, z=self.z * scalar)
def __add__(self, other: "Vector") -> "Vector":
if not isinstance(other, Vector):
return NotImplemented
return Vector(x=self.x + other.x, y=self.y + other.y, z=self.z + other.z)
def __sub__(self, other: "Vector") -> "Vector":
if not isinstance(other, Vector):
return NotImplemented
return Vector(x=self.x - other.x, y=self.y - other.y, z=self.z - other.z)
def __mul__(self, scalar: float) -> "Vector":
if not isinstance(scalar, (int, float)):
return NotImplemented
return self.scale(float(scalar))
def __rmul__(self, scalar: float) -> "Vector":
return self * scalar
def __truediv__(self, scalar: float) -> "Vector":
return Vector(x=self.x / scalar, y=self.y / scalar, z=self.z / scalar)
def __neg__(self) -> "Vector":
return Vector(x=-self.x, y=-self.y, z=-self.z)
def __repr__(self) -> str:
return f"Vector(x={_repr_float(self.x)}, y={_repr_float(self.y)}, z={_repr_float(self.z)})"
def _coerce_vector(value: Vector | Point) -> Vector:
if isinstance(value, Vector):
return value
if isinstance(value, Point):
return Vector.from_point(value)
raise TypeError(f"Expected Vector or Point, got {type(value).__name__}")
[docs]
@dataclass(frozen=True, slots=True, repr=False)
class Twist:
"""Linear and angular velocity vectors."""
linear: Vector = field(default_factory=Vector)
angular: Vector = field(default_factory=Vector)
def __post_init__(self) -> None:
object.__setattr__(self, "linear", _coerce_vector(self.linear))
object.__setattr__(self, "angular", _coerce_vector(self.angular))
[docs]
def isclose(self, other: Self, *, tol: float = 1e-9) -> bool:
"""Return whether two twists are equal within an absolute tolerance."""
if not isinstance(other, Twist):
raise TypeError(f"Cannot compare Twist with {type(other).__name__}; expected Twist")
return self.linear.isclose(other.linear, tol=tol) and self.angular.isclose(other.angular, tol=tol)
def __repr__(self) -> str:
return f"Twist(linear={self.linear!r}, angular={self.angular!r})"
def wrap_error(value: float, period: float) -> float:
"""Wrap a periodic scalar error into ``[-period / 2, period / 2)``."""
if not isfinite(value):
raise ValueError("value must be finite")
if period <= 0 or not isfinite(period):
raise ValueError("period must be a positive finite number")
return ((value + period / 2.0) % period) - period / 2.0
[docs]
@dataclass(frozen=True, slots=True, repr=False)
class Quaternion:
"""Orientation quaternion (xyzw)."""
x: float = 0.0
y: float = 0.0
z: float = 0.0
w: float = 1.0
[docs]
@classmethod
def identity(cls) -> "Quaternion":
"""Return the identity quaternion."""
return cls()
[docs]
@classmethod
def from_array(cls, values: Sequence[float]) -> "Quaternion":
"""Create a quaternion from an ``[x, y, z, w]`` array-like."""
arr = np.asarray(values, dtype=np.float64)
if arr.shape != (4,):
raise ValueError(f"Expected 4 values [x, y, z, w], got shape {arr.shape}")
return cls(x=float(arr[0]), y=float(arr[1]), z=float(arr[2]), w=float(arr[3]))
[docs]
@classmethod
def from_rpy(cls, roll: float, pitch: float, yaw: float) -> "Quaternion":
"""Create a quaternion from roll-pitch-yaw angles (radians, xyz extrinsic)."""
x, y, z, w = _utils.rpy_to_quat(roll, pitch, yaw)
return cls(x=x, y=y, z=z, w=w)
[docs]
@classmethod
def from_axis_angle(cls, axis: Sequence[float], angle: float) -> "Quaternion":
"""Create a quaternion from an axis-angle rotation."""
x, y, z, w = _utils.axis_angle_to_quat(axis, angle)
return cls(x=x, y=y, z=z, w=w)
[docs]
@classmethod
def from_matrix(cls, matrix: np.ndarray) -> "Quaternion":
"""Create a quaternion from a 3x3 rotation matrix."""
x, y, z, w = _utils.matrix_to_quat(matrix)
return cls(x=x, y=y, z=z, w=w)
[docs]
def array(self) -> np.ndarray:
"""Return ``[x, y, z, w]`` as a float64 numpy array."""
return np.array([self.x, self.y, self.z, self.w], dtype=np.float64)
[docs]
def matrix(self) -> np.ndarray:
"""Return the equivalent 3x3 rotation matrix."""
return _utils.quat_to_matrix(self.x, self.y, self.z, self.w)
[docs]
def normalized(self) -> "Quaternion":
"""Return a unit quaternion."""
x, y, z, w = _utils.quat_normalize(self.x, self.y, self.z, self.w)
return Quaternion(x=x, y=y, z=z, w=w)
[docs]
def rotate(self, vector: Vector) -> Vector:
"""Rotate a vector by this quaternion."""
if not isinstance(vector, Vector):
raise TypeError(f"Cannot rotate {type(vector).__name__}; expected Vector")
x, y, z = _utils.rotate_vector(self.matrix(), vector.x, vector.y, vector.z)
return Vector(x=x, y=y, z=z)
[docs]
def to_rpy(self) -> tuple[float, float, float]:
"""Return roll, pitch, yaw (radians, xyz extrinsic)."""
return _utils.quat_to_rpy(self.x, self.y, self.z, self.w)
[docs]
def isclose(self, other: Self, *, tol: float = 1e-9) -> bool:
"""Return whether two quaternions represent the same rotation within a tolerance."""
if not isinstance(other, Quaternion):
raise TypeError(f"Cannot compare Quaternion with {type(other).__name__}; expected Quaternion")
left = self.normalized().array()
right = other.normalized().array()
return bool(np.allclose(left, right, rtol=0.0, atol=tol) or np.allclose(left, -right, rtol=0.0, atol=tol))
def __mul__(self, other: "Quaternion") -> "Quaternion":
"""Compose with another quaternion."""
if not isinstance(other, Quaternion):
raise TypeError(f"Cannot compose Quaternion with {type(other).__name__}; expected Quaternion")
x, y, z, w = _utils.quat_multiply(
self.x,
self.y,
self.z,
self.w,
other.x,
other.y,
other.z,
other.w,
)
return Quaternion(x=x, y=y, z=z, w=w)
def __repr__(self) -> str:
return (
f"Quaternion(x={_repr_float(self.x)}, y={_repr_float(self.y)}, "
f"z={_repr_float(self.z)}, w={_repr_float(self.w)})"
)
[docs]
@dataclass(frozen=True, slots=True, repr=False)
class Pose:
"""A combined 3D position and orientation."""
position: Point = field(default_factory=Point)
orientation: Quaternion = field(default_factory=Quaternion)
[docs]
@classmethod
def identity(cls) -> "Pose":
"""Return the identity pose."""
return cls()
[docs]
@classmethod
def from_matrix(cls, matrix: np.ndarray) -> "Pose":
"""Create a pose from a 4x4 homogeneous transform matrix."""
tx, ty, tz, qx, qy, qz, qw = _utils.matrix_to_pose(matrix)
return cls(
position=Point(x=tx, y=ty, z=tz),
orientation=Quaternion(x=qx, y=qy, z=qz, w=qw),
)
[docs]
@classmethod
def from_xy_yaw(cls, x: float, y: float, yaw: float, *, z: float = 0.0) -> "Pose":
"""Create a planar pose from x, y, and yaw (radians)."""
return cls(
position=Point(x=float(x), y=float(y), z=float(z)),
orientation=Quaternion.from_rpy(0.0, 0.0, float(yaw)),
)
[docs]
def project_to_xy(self, *, z: float | None = None) -> "Pose":
"""Return a pose flattened to the XY plane, preserving yaw only."""
_, _, yaw = self.orientation.to_rpy()
return Pose.from_xy_yaw(
self.position.x,
self.position.y,
yaw,
z=self.position.z if z is None else float(z),
)
[docs]
def matrix(self) -> np.ndarray:
"""Return a 4x4 homogeneous transform matrix."""
return _utils.pose_matrix(
self.position.x,
self.position.y,
self.position.z,
self.orientation.x,
self.orientation.y,
self.orientation.z,
self.orientation.w,
)
[docs]
def inverse(self) -> "Pose":
"""Return the inverse pose."""
inv = _utils.invert_matrix(self.matrix())
return Pose.from_matrix(inv)
[docs]
def apply(self, point: Point) -> Point:
"""Transform a point by this pose."""
x, y, z = _utils.transform_point(self.matrix(), point.x, point.y, point.z)
return Point(x=x, y=y, z=z)
[docs]
def rotate(self, vector: Vector) -> Vector:
"""Rotate a vector by this pose without applying translation."""
if not isinstance(vector, Vector):
raise TypeError(f"Cannot rotate {type(vector).__name__}; expected Vector")
x, y, z = _utils.rotate_vector(self.matrix(), vector.x, vector.y, vector.z)
return Vector(x=x, y=y, z=z)
[docs]
def isclose(self, other: Self, *, tol: float = 1e-9) -> bool:
"""Return whether two poses are geometrically equal within a tolerance."""
if not isinstance(other, Pose):
raise TypeError(f"Cannot compare Pose with {type(other).__name__}; expected Pose")
return self.position.isclose(other.position, tol=tol) and self.orientation.isclose(
other.orientation,
tol=tol,
)
@overload
def __mul__(self, other: Point) -> Point: ...
@overload
def __mul__(self, other: "Pose") -> "Pose": ...
def __mul__(self, other: "Pose | Point") -> "Pose | Point":
"""Compose with another pose or apply to a point.
To offset a pose by a :class:`Transform`, be explicit about intent:
``transform * pose`` re-expresses the pose in the transform's parent
frame, or use ``pose * transform.as_pose()`` for a frameless offset.
"""
if isinstance(other, Point):
return self.apply(other)
if isinstance(other, Pose):
return Pose.from_matrix(_utils.compose_matrices(self.matrix(), other.matrix()))
raise TypeError(f"Cannot compose Pose with {type(other).__name__}; expected Pose or Point")
def __repr__(self) -> str:
return f"Pose(position={self.position!r}, orientation={self.orientation!r})"
[docs]
@dataclass(frozen=True, slots=True)
class FrameInfo:
"""A frame in the TF tree."""
name: str
parent: str = ""
transform: Transform | None = None
[docs]
@dataclass(frozen=True, slots=True, repr=False)
class FrameTreeNode:
"""A node in a robot-scoped TF subtree."""
frame_id: str
relative_name: str
transform: Transform | None = None
children: tuple["FrameTreeNode", ...] = ()
def __repr__(self) -> str:
return self._format_tree()
def __str__(self) -> str:
return self._format_tree()
def _format_tree(self) -> str:
lines = [self.relative_name]
lines.extend(self._child_lines(" "))
return "\n".join(lines)
def _child_lines(self, prefix: str) -> list[str]:
lines: list[str] = []
for i, child in enumerate(self.children):
is_last = i == len(self.children) - 1
branch = "└─ " if is_last else "├─ "
cont = " " if is_last else "│ "
lines.append(f"{prefix}{branch}{child.relative_name}")
lines.extend(child._child_lines(prefix + cont))
return lines
[docs]
class FrameTree(tuple[FrameTreeNode, ...]):
"""A TF tree; may have multiple roots when subtrees are disconnected."""
[docs]
def flatten(self) -> list[FrameInfo]:
"""Flatten the tree into FrameInfo entries (child -> parent)."""
frames: list[FrameInfo] = []
def _walk(node: FrameTreeNode, parent: str) -> None:
frames.append(FrameInfo(name=node.frame_id, parent=parent, transform=node.transform))
for child in node.children:
_walk(child, node.frame_id)
for root in self:
_walk(root, "")
return frames
def __repr__(self) -> str:
return self._format_tree()
def __str__(self) -> str:
return self._format_tree()
def _format_tree(self) -> str:
if not self:
return ""
if len(self) == 1:
return self[0]._format_tree()
lines: list[str] = []
for i, root in enumerate(self):
is_last = i == len(self) - 1
branch = "└─ " if is_last else "├─ "
cont = " " if is_last else "│ "
lines.append(f"{branch}{root.relative_name}")
lines.extend(root._child_lines(cont))
return "\n".join(lines)
__all__ = [
"Point",
"Vector",
"Twist",
"wrap_error",
"Quaternion",
"Pose",
"Transform",
"FrameInfo",
"FrameTreeNode",
"FrameTree",
]