"""
Trajectory Module
=================
This module provides data structures for representing time-series trajectories
of objects in world coordinates, including position, orientation, and optional
velocity information.
Classes
-------
WorldTrajectory
A time-indexed sequence of 4x4 transformation matrices representing
an object's pose over time in world coordinates.
Example
-------
Creating a simple trajectory::
import numpy as np
from mimic_data_api.trajectory_numpy import WorldTrajectory
# Create a 10-frame trajectory
timestamps = np.arange(0, 1.0, 0.1) # 10 timestamps at 10 Hz
transforms = np.stack([np.eye(4)] * 10) # Identity transforms
trajectory = WorldTrajectory(
timestamps=timestamps,
object2world=transforms
)
"""
from __future__ import annotations
from pydantic import BaseModel, model_validator, Field
from numpydantic import NDArray, Shape
import numpy as np
[docs]
class WorldTrajectory(BaseModel, validate_assignment=True):
"""Time-series trajectory of an object in world coordinates.
This class represents a sequence of poses (position and orientation)
for an object over time, stored as 4x4 homogeneous transformation matrices.
Optionally includes angular and linear velocity information.
The transformation matrices follow the convention where the matrix transforms
points from the object's local coordinate frame to the world coordinate frame:
.. math::
p_{world} = T_{object \\rightarrow world} \\cdot p_{object}
Attributes
----------
timestamps : NDArray[Shape["* x"], float]
1D array of N timestamps in seconds. Must be strictly increasing
(monotonic, no duplicates) and match the number of transformation matrices.
object2world : NDArray[Shape["* x, 4 y, 4 z"], float]
Array of shape (N, 4, 4) containing homogeneous transformation matrices.
Each matrix transforms points from object coordinates to world coordinates.
angular_velocity : NDArray[Shape["* x, 3 y"], float] or None
Optional array of shape (N, 3) containing angular velocity vectors
in world coordinates, measured in rad/s. Order is (roll, pitch, yaw) rate.
linear_velocity : NDArray[Shape["* x, 3 y"], float] or None
Optional array of shape (N, 3) containing linear velocity vectors
in world coordinates, measured in m/s. Order is (vx, vy, vz).
Raises
------
ValueError
If timestamps and transforms have different lengths, or if velocity
arrays (when provided) don't match the timestamp count.
Note
----
This class uses ``validate_assignment=True``, meaning validation runs
on attribute assignment, not just initialization.
Example
-------
Creating a trajectory with velocity information::
import numpy as np
from mimic_data_api.trajectory_numpy import WorldTrajectory
n_frames = 100
trajectory = WorldTrajectory(
timestamps=np.linspace(0, 10, n_frames),
object2world=np.stack([np.eye(4)] * n_frames),
linear_velocity=np.zeros((n_frames, 3)),
angular_velocity=np.zeros((n_frames, 3))
)
Accessing pose at a specific time index::
pose_at_t5 = trajectory.object2world[5] # 4x4 matrix
position = pose_at_t5[:3, 3] # Translation vector
rotation = pose_at_t5[:3, :3] # Rotation matrix
"""
timestamps: NDArray[Shape["* x"], float] = Field(
description="Vector of N timestamps in seconds"
)
object2world: NDArray[Shape["* x, 4 y, 4 z"], float] = Field(
description="Nx4x4 transform matrices."
)
angular_velocity: NDArray[Shape["* x, 3 y"], float] | None = Field(
default=None,
description="Optional Nx3 angular velocity in world coordinate vectors in rad/s.",
)
linear_velocity: NDArray[Shape["* x, 3 y"], float] | None = Field(
default=None,
description=(
"Optional Nx3 linear velocity in world coordinate vectors in m/s. Must be "
"d(translation)/dt of object2world: rolling-shutter/motion-blur consumers (e.g. the "
"GS renderer) advect points by it, so an independent sensor-measured velocity here "
"smears moving objects."
),
)
[docs]
@model_validator(mode="after")
def len_validation(self) -> WorldTrajectory:
"""Validate that all arrays have consistent lengths.
Returns
-------
WorldTrajectory
The validated trajectory instance.
Raises
------
ValueError
If array lengths are inconsistent, or if ``timestamps`` is not
strictly increasing (monotonic, no duplicates).
"""
if len(self.timestamps) != len(self.object2world):
raise ValueError("Timestamps and transforms should have the same size")
if self.angular_velocity is not None:
if len(self.angular_velocity) != len(self.timestamps):
raise ValueError(
"Angular velocity should have the same size as timestamps"
)
if self.linear_velocity is not None:
if len(self.linear_velocity) != len(self.timestamps):
raise ValueError(
"Linear velocity should have the same size as timestamps"
)
# Strictly increasing timestamps: scipy Slerp/RotationSpline raise on non-increasing
# times, and finite-difference velocity divides by dt=0 on duplicates. Fail early here.
timestamps = np.asarray(self.timestamps)
if timestamps.shape[0] >= 2:
diffs = np.diff(timestamps)
if not np.all(diffs > 0):
i = int(np.argmax(diffs <= 0)) # first offending step
raise ValueError(
"WorldTrajectory timestamps must be strictly increasing "
f"(monotonic, no duplicates); first violation at index {i + 1}: "
f"timestamps[{i + 1}]={timestamps[i + 1]!r} <= timestamps[{i}]={timestamps[i]!r}. "
"Sort and de-duplicate (timestamps, object2world, and any velocities) "
"together before constructing the trajectory."
)
return self
[docs]
def __eq__(self, other: object) -> bool:
"""Check equality with another WorldTrajectory.
Two trajectories are considered equal if they have the same timestamps,
transforms, and velocity arrays (using approximate comparison for floats).
Parameters
----------
other : object
The object to compare against.
Returns
-------
bool
``True`` if trajectories are equal, ``False`` otherwise.
"""
if not isinstance(other, WorldTrajectory):
return False
if not np.array_equal(self.timestamps, other.timestamps):
return False
if not np.allclose(self.timestamps, other.timestamps):
return False
# angular velocities optional checks
if self.angular_velocity is not None and other.angular_velocity is not None:
if not np.allclose(self.timestamps, other.timestamps):
return False
if self.angular_velocity is None and other.angular_velocity is not None:
return False
if self.angular_velocity is not None and other.angular_velocity is None:
return False
# linear velocities optional checks
if self.linear_velocity is not None and other.linear_velocity is not None:
if not np.allclose(self.timestamps, other.timestamps):
return False
if self.linear_velocity is None and other.linear_velocity is not None:
return False
if self.linear_velocity is not None and other.linear_velocity is None:
return False
return True