def get_transformation(self, timestamps=None, arrays_first=True):
""" Get the transformation across all reference frames.
Parameters
----------
timestamps: array_like, shape (n_timestamps,), optional
Timestamps to which the transformation should be matched. If not
provided the matcher will call `get_timestamps` for the target
timestamps.
arrays_first: bool, default True
If True and timestamps aren't provided, the first array in the
list defines the sampling of the timestamps. Otherwise, the first
reference frame in the list defines the sampling.
Returns
-------
translation: array_like, shape (3,) or (n_timestamps, 3)
The translation across all reference frames.
rotation: array_like, shape (4,) or (n_timestamps, 4)
The rotation across all reference frames.
timestamps: array_like, shape (n_timestamps,) or None
The timestamps for which the transformation is defined.
"""
from rigid_body_motion.utils import rotate_vectors
if timestamps is None:
timestamps = self.get_timestamps(arrays_first)
translation = np.zeros(3) if timestamps is None else np.zeros((1, 3))
rotation = quaternion(1.0, 0.0, 0.0, 0.0)
for frame in self.frames:
t, r = self._transform_from_frame(frame, timestamps)
if frame.inverse:
translation = rotate_vectors(
1 / as_quat_array(r), translation - np.array(t)
)
rotation = 1 / as_quat_array(r) * rotation
else:
translation = rotate_vectors(
as_quat_array(r), translation
) + np.array(t)
rotation = as_quat_array(r) * rotation
return translation, as_float_array(rotation), timestamps
def transform_vectors(
self,
arr,
to_frame,
axis=-1,
time_axis=0,
timestamps=None,
return_timestamps=False,
):
""" Transform array of vectors from this frame to another.
Parameters
----------
arr: array_like
The array to transform.
to_frame: str or ReferenceFrame
The target reference frame. If str, the frame will be looked up
in the registry under that name.
axis: int, default -1
The axis of the array representing the spatial coordinates of the
vectors.
time_axis: int, default 0
The axis of the array representing the timestamps of the vectors.
timestamps: array_like, optional
The timestamps of the vectors, corresponding to the `time_axis`
of the array. If not None, the axis defined by `time_axis` will be
re-sampled to the timestamps for which the transformation is
defined.
return_timestamps: bool, default False
If True, also return the timestamps after the transformation.
Returns
-------
arr_transformed: array_like
The transformed array.
ts: array_like, shape (n_timestamps,) or None
The timestamps after the transformation.
"""
arr, arr_ts = self._validate_input(arr, axis, 3, timestamps, time_axis)
matcher = self._get_matcher(to_frame, arrays=[(arr, arr_ts)])
t, r, ts = matcher.get_transformation()
arr, _ = matcher.get_arrays(ts)
r = self._expand_singleton_axes(r, arr.ndim)
arr = rotate_vectors(r, arr, axis=axis)
# undo time axis swap
if time_axis is not None:
arr = np.swapaxes(arr, 0, time_axis)
if not return_timestamps:
return arr
else:
return arr, ts
def plot_quaternions(arr, base=None, ax=None, figsize=(6, 6), **kwargs):
""" Plot an array of quaternions.
Parameters
----------
arr: array_like, shape (4,) or (N, 4)
Array of quaternions to plot.
base: array_like, shape (4,) or (N, 4), optional
If provided, base points of the quaternions.
ax: matplotlib.axes.Axes instance, optional
If provided, plot the points onto these axes.
figsize:
If `ax` is not provided, create a figure of this size.
kwargs:
Additional keyword arguments passed to ax.quiver().
Returns
-------
lines: list of Line3DCollection
A list of lines representing the plotted data.
"""
if ax is None:
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(111, projection="3d")
vx = rotate_vectors(arr, np.array((1, 0, 0)), one_to_one=False)
vy = rotate_vectors(arr, np.array((0, 1, 0)), one_to_one=False)
vz = rotate_vectors(arr, np.array((0, 0, 1)), one_to_one=False)
lines = [
plot_vectors(vx, base, ax, color="r", length=0.5, **kwargs),
plot_vectors(vy, base, ax, color="g", length=0.5, **kwargs),
plot_vectors(vz, base, ax, color="b", length=0.5, **kwargs),
]
return lines
def make_test_motion(
n_samples,
freq=1,
max_angle=np.pi / 2,
fs=1000,
stack=True,
inverse=False,
offset=None,
):
""" Create sinusoidal linear and angular motion around all three axes. """
import pandas as pd
if inverse:
max_angle = -max_angle
trajectory = (
max_angle *
np.sin(2 * np.pi * freq * np.arange(n_samples) / fs)[:, np.newaxis])
if stack:
ax = np.array((1.0, 0.0, 0.0))[np.newaxis, :]
ay = np.array((0.0, 1.0, 0.0))[np.newaxis, :]
az = np.array((0.0, 0.0, 1.0))[np.newaxis, :]
tx, ty, tz = np.array_split(trajectory, 3)
translation = np.vstack((tx * ax, ty * ay, tz * az))
else:
translation = np.tile(trajectory, (1, 3))
rotation = as_float_array(from_rotation_vector(translation))
if offset is not None:
translation += rotate_vectors(as_quat_array(rotation),
np.array(offset)[np.newaxis, :])
timestamps = pd.date_range(start=0, periods=n_samples, freq=f"{1/fs}S")
return translation, rotation, timestamps
def _init_arrays(translation, rotation, timestamps, inverse):
""" Initialize translation, rotation and timestamp arrays. """
if timestamps is not None:
timestamps = np.asarray(timestamps)
if timestamps.ndim != 1:
raise ValueError("timestamps must be one-dimensional.")
t_shape = (len(timestamps), 3)
r_shape = (len(timestamps), 4)
else:
t_shape = (3, )
r_shape = (4, )
if translation is not None:
translation = np.asarray(translation)
if translation.shape != t_shape:
raise ValueError(
f"Expected translation to be of shape {t_shape}, got "
f"{translation.shape}")
else:
translation = np.zeros(t_shape)
if rotation is not None:
rotation = np.asarray(rotation)
if rotation.shape != r_shape:
raise ValueError(
f"Expected rotation to be of shape {r_shape}, got "
f"{rotation.shape}")
else:
rotation = np.zeros(r_shape)
rotation[..., 0] = 1.0
if inverse:
rotation = qinv(rotation)
translation = -rotate_vectors(rotation, translation)
return translation, rotation, timestamps
def iterative_closest_point(
v1,
v2,
dim=None,
axis=None,
init_transform=None,
max_iterations=20,
tolerance=1e-3,
):
""" Iterative closest point algorithm matching two arrays of vectors.
Finds the rotation `r` and the translation `t` such that:
.. math:: v_2 \simeq rot(r, v_1) + t
Parameters
----------
v1: array_like, shape (..., 3, ...)
The first array of vectors.
v2: array_like, shape (..., 3, ...)
The second array of vectors.
dim: str, optional
If the first array is a DataArray, the name of the dimension
representing the spatial coordinates of the vectors.
axis: int, optional
The axis of the arrays representing the spatial coordinates of the
vectors. Defaults to the last axis of the arrays.
init_transform: tuple, optional
Initial guess as (translation, rotation) tuple.
max_iterations: int, default 20
Maximum number of iterations.
tolerance: float, default 1e-3
Abort if the mean distance error between the transformed arrays does
not improve by more than this threshold between iterations.
Returns
-------
translation: array_like, shape (3,)
Translation of transform.
rotation: array_like, shape (4,)
Rotation of transform.
References
----------
Adapted from https://github.com/ClayFlannigan/icp
Notes
-----
For points with known correspondences (e.g. timeseries of positions), it is
recommended to interpolate the points to a common sampling base and use the
`best_fit_transform` method.
See Also
--------
best_fit_transform, best_fit_rotation
""" # noqa
v1, v2, was_dataarray = _reshape_vectors(v1,
v2,
axis,
dim,
same_shape=False)
v1_new = np.copy(v1)
# apply the initial pose estimation
if init_transform is not None:
t, r = init_transform
v1_new = rotate_vectors(np.asarray(r), v1_new) + np.asarray(t)
prev_error = 0
for i in range(max_iterations):
# find the nearest neighbors between the current source and destination
# points
idx, distances = _nearest_neighbor(v1_new, v2)
# compute the transformation between the current source and nearest
# destination points
t, r = best_fit_transform(v1_new, v2[idx])
# update the current source
v1_new = rotate_vectors(r, v1_new) + t
# check error
mean_error = np.mean(distances)
if np.abs(prev_error - mean_error) < tolerance:
break
prev_error = mean_error
# calculate final transformation
translation, rotation = best_fit_transform(v1, v1_new)
if was_dataarray:
translation, rotation = _make_transform_dataarrays(
translation, rotation)
return translation, rotation
def test_rotate_vectors(self):
""""""
v = np.ones((10, 3))
q = np.tile(from_euler_angles(yaw=np.pi / 4, return_quaternion=True),
10)
vr = np.vstack((np.zeros(10), np.sqrt(2) * np.ones(10), np.ones(10))).T
# single quaternion, single vector
vr_act = rotate_vectors(q[0], v[0])
np.testing.assert_allclose(vr[0], vr_act, rtol=1.0)
# single quaternion, multiple vectors
vr_act = rotate_vectors(q[0], v)
np.testing.assert_allclose(vr, vr_act, rtol=1.0)
# single quaternion, explicit axis
vr_act = rotate_vectors(q[0], v, axis=1)
np.testing.assert_allclose(vr, vr_act, rtol=1.0)
# multiple quaternions, multiple vectors
vr_act = rotate_vectors(q, v)
np.testing.assert_allclose(vr, vr_act)
# different axis
vr_act = rotate_vectors(q, v.T, axis=0)
np.testing.assert_allclose(vr.T, vr_act)
# singleton expansion
vr_act = rotate_vectors(q[:, None], v[None, ...])
np.testing.assert_allclose(np.tile(vr, (10, 1, 1)), vr_act)
# float dtype
vr_act = rotate_vectors(as_float_array(q[0]), v[0])
np.testing.assert_allclose(vr[0], vr_act, rtol=1.0)
with pytest.raises(ValueError):
rotate_vectors(q, v.T)
with pytest.raises(ValueError):
rotate_vectors(q, np.ones((10, 4)))