def compensated_tree():
""""""
n_samples = 10000
stack = False
rbm.register_frame("world", update=True)
t, r, ts = make_test_motion(n_samples, stack=stack)
rbm.ReferenceFrame(
translation=t,
rotation=r,
timestamps=ts,
parent="world",
name="head",
).register(update=True)
it, ir, _ = make_test_motion(n_samples,
inverse=True,
stack=stack,
offset=(1.0, 0.0, 0.0))
rbm.ReferenceFrame(
translation=it,
rotation=ir,
timestamps=ts,
parent="head",
name="eyes",
).register(update=True)
def test_lookup_transform_discrete(self):
""""""
t1 = np.ones((10, 3))
t2 = np.array([[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]])
rf_world = rbm.ReferenceFrame("world")
rf_child1 = rbm.ReferenceFrame(
"child1", rf_world, translation=t1, timestamps=np.arange(10),
)
rf_child2 = rbm.ReferenceFrame(
"child2",
rf_world,
translation=t2,
timestamps=[2, 5],
discrete=True,
)
# interpolated first
t_act, r_act, ts = rf_child1.lookup_transform(rf_child2)
npt.assert_equal(t_act, [[0.0, 1.0, 1.0]] * 5 + [[-1.0, 1.0, 1.0]] * 5)
npt.assert_equal(r_act, np.tile([[1.0, 0.0, 0.0, 0.0]], (10, 1)))
npt.assert_allclose(ts, np.arange(10))
# event-based first
t_act, r_act, ts = rf_child2.lookup_transform(rf_child1)
npt.assert_equal(
t_act, [[0.0, -1.0, -1.0]] * 5 + [[1.0, -1.0, -1.0]] * 5
)
npt.assert_equal(r_act, np.tile([[1.0, 0.0, 0.0, 0.0]], (10, 1)))
npt.assert_allclose(ts, np.arange(10))
def test_walk(self):
""""""
rf_world = rbm.ReferenceFrame("world")
rf_child = rbm.ReferenceFrame("child", parent=rf_world)
rf_child2 = rbm.ReferenceFrame("child2", parent=rf_world)
up, down = rf_child._walk(rf_child2)
assert up == (rf_child,)
assert down == (rf_child2,)
def test_register(self):
""""""
rf_world = rbm.ReferenceFrame("world")
rf_world.register()
assert "world" in rbm.registry
assert rbm.registry["world"] is rf_world
# update=True
rf_world2 = rbm.ReferenceFrame("world")
rf_world2.register(update=True)
assert rbm.registry["world"] is rf_world2
def test_constructor(self):
""""""
rf_world = rbm.ReferenceFrame("world")
rf_child = rbm.ReferenceFrame("child", parent=rf_world)
assert rf_child.parent is rf_world
_register(rf_world)
rf_child2 = rbm.ReferenceFrame("child2", parent="world")
assert rf_child2.parent is rf_world
# invalid parent
with pytest.raises(ValueError):
rbm.ReferenceFrame("child3", parent="not_an_rf")
def get_rf_tree(
tc1=(0.0, 0.0, 0.0),
rc1=(1.0, 0.0, 0.0, 0.0),
tc2=(0.0, 0.0, 0.0),
rc2=(1.0, 0.0, 0.0, 0.0),
):
""""""
rf_world = rbm.ReferenceFrame("world")
rf_child1 = rbm.ReferenceFrame(
"child1", parent=rf_world, translation=tc1, rotation=rc1
)
rf_child2 = rbm.ReferenceFrame(
"child2", parent=rf_world, translation=tc2, rotation=rc2
)
return rf_world, rf_child1, rf_child2
def test_transform_points(self, transform_grid, get_rf_tree):
""""""
o, ot, p, pt, rc1, rc2, tc1, tc2 = transform_grid
rf_world, rf_child1, rf_child2 = get_rf_tree(tc1, rc1, tc2, rc2)
rf_child3 = rbm.ReferenceFrame(
"child3", rf_child1, timestamps=np.arange(5) + 2.5
)
# static reference frame + single point
pt_act = rf_child1.transform_points(p, rf_child2)
np.testing.assert_allclose(pt_act, pt)
# moving reference frame + single point
pt_act = rf_child3.transform_points(p, rf_child2)
np.testing.assert_allclose(pt_act, np.tile(pt, (5, 1)), rtol=1.0)
# moving reference frame + multiple points
pt_act = rf_child3.transform_points(
np.tile(p, (10, 1)), rf_child2, timestamps=np.arange(10)
)
np.testing.assert_allclose(pt_act, np.tile(pt, (4, 1)), rtol=1.0)
# moving reference frame + multiple n-dimensional points
pt_act = rf_child3.transform_points(
np.tile(p, (10, 10, 1)),
rf_child2,
timestamps=np.arange(10),
time_axis=1,
)
np.testing.assert_allclose(pt_act, np.tile(pt, (10, 4, 1)), rtol=1.0)
def test_plot_reference_frame(self):
rf_world = rbm.ReferenceFrame("world")
rf_observer = rbm.ReferenceFrame(
"observer",
parent=rf_world,
translation=(5, 0, 0),
rotation=(np.sqrt(2) / 2, 0, 0, np.sqrt(2) / 2),
)
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
rbm.plot.reference_frame(rf_world, ax=ax)
rbm.plot.reference_frame(rf_observer, ax=ax)
fig.tight_layout()
plt.show()
def test_destructor(self):
""""""
rf_world = rbm.ReferenceFrame("world")
_register(rf_world)
del rf_world
assert "world" in rbm.registry
del rbm.registry["world"]
assert "world" not in rbm.registry
def test_from_rotation_matrix(self):
""""""
mat = np.array([[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
rf_world = rbm.ReferenceFrame("world")
rf_child = rbm.ReferenceFrame.from_rotation_matrix(mat, rf_world)
npt.assert_equal(rf_child.translation, (0.0, 0.0, 0.0))
npt.assert_allclose(rf_child.rotation, (-0.5, -0.5, 0.5, 0.5))
assert rf_child.timestamps is None
with pytest.raises(ValueError):
rbm.ReferenceFrame.from_rotation_matrix(np.zeros((4, 4)), rf_world)
def test_register(self):
""""""
rf_world = rbm.ReferenceFrame("world")
_register(rf_world)
assert rbm.registry["world"] is rf_world
# name is None
with pytest.raises(ValueError):
_register(rbm.ReferenceFrame())
# already registered
with pytest.raises(ValueError):
_register(rf_world)
# update=True
rf_child = rbm.ReferenceFrame("child", parent=rf_world)
_register(rf_child)
rf_child_new = rbm.ReferenceFrame("child", parent=rf_world)
_register(rf_child_new, update=True)
assert rbm.registry["child"] is rf_child_new
assert rbm.registry["child"].parent is rf_world
assert rbm.registry["world"].children == (rf_child_new,)
def test_lookup_transform(self, rf_grid, get_rf_tree):
""""""
r, rc1, rc2, t, tc1, tc2 = rf_grid
rf_world, rf_child1, rf_child2 = get_rf_tree(tc1, rc1, tc2, rc2)
# child1 to world
t_act, r_act, ts = rf_child1.lookup_transform(rf_world)
npt.assert_almost_equal(t_act, tc1)
npt.assert_almost_equal(r_act, rc1)
assert ts is None
# child1 to child2
t_act, r_act, ts = rf_child1.lookup_transform(rf_child2)
npt.assert_almost_equal(t_act, t)
npt.assert_almost_equal(r_act, r)
assert ts is None
# inverse
rf_child1_inv = rbm.ReferenceFrame(
parent=rf_world, translation=tc1, rotation=rc1, inverse=True
)
t_act, r_act, ts = rf_world.lookup_transform(rf_child1_inv)
npt.assert_almost_equal(t_act, tc1)
npt.assert_almost_equal(r_act, rc1)
assert ts is None
# with timestamps
rf_child3 = rbm.ReferenceFrame(
"child3", rf_child1, timestamps=np.arange(5) + 2.5
)
rf_child4 = rbm.ReferenceFrame(
"child4", rf_child2, timestamps=np.arange(10)
)
t_act, r_act, ts = rf_child3.lookup_transform(rf_child4)
npt.assert_almost_equal(t_act, np.tile(t, (5, 1)))
npt.assert_almost_equal(r_act, np.tile(r, (5, 1)))
npt.assert_equal(ts, np.arange(5) + 2.5)
def test_from_rotation_datarray(self):
""""""
xr = pytest.importorskip("xarray")
da = xr.DataArray(
np.ones((10, 4)), {"time": np.arange(10)}, ("time", "axis")
)
rf_world = rbm.ReferenceFrame("world")
rf_child = rbm.ReferenceFrame.from_rotation_dataarray(
da, "time", rf_world
)
npt.assert_equal(rf_child.rotation, np.ones((10, 4)))
npt.assert_equal(rf_child.timestamps, np.arange(10))
def _mock_frame(
t=None,
r=None,
ts=None,
name="mock",
parent="world",
inverse=False,
discrete=False,
):
""""""
if t is None and r is None:
t = (0.0, 0.0, 0.0)
r = (1.0, 0.0, 0.0, 0.0)
return rbm.ReferenceFrame(name, parent, t, r, ts, inverse, discrete)
def test_transform_vectors(self, transform_grid, get_rf_tree):
""""""
o, ot, p, pt, rc1, rc2, tc1, tc2 = transform_grid
_, rf_child1, rf_child2 = get_rf_tree(tc1, rc1, tc2, rc2)
rf_child3 = rbm.ReferenceFrame(
"child3", rf_child1, timestamps=np.arange(5) + 2.5
)
# static reference frame + single vector
vt_act = rf_child1.transform_vectors(p, rf_child2)
v0t = rf_child1.transform_points((0.0, 0.0, 0.0), rf_child2)
vt = np.array(pt) - np.array(v0t)
np.testing.assert_allclose(vt_act, vt, rtol=1.0, atol=1e-15)
# moving reference frame + single vector
vt_act = rf_child3.transform_vectors(p, rf_child2)
v0t = rf_child3.transform_points((0.0, 0.0, 0.0), rf_child2)
vt = np.tile(pt, (5, 1)) - np.array(v0t)
np.testing.assert_allclose(vt_act, vt, rtol=1.0, atol=1e-15)
# moving reference frame + multiple vectors
vt_act = rf_child3.transform_vectors(
np.tile(p, (10, 1)), rf_child2, timestamps=np.arange(10)
)
v0t = rf_child3.transform_points(
np.tile((0.0, 0.0, 0.0), (10, 1)),
rf_child2,
timestamps=np.arange(10),
)
vt = np.tile(pt, (4, 1)) - np.array(v0t)
np.testing.assert_allclose(vt_act, vt, rtol=1.0, atol=1e-15)
# moving reference frame + multiple n-dimensional vectors
vt_act = rf_child3.transform_vectors(
np.tile(p, (10, 10, 1)),
rf_child2,
timestamps=np.arange(10),
time_axis=1,
)
v0t = rf_child3.transform_points(
np.tile((0.0, 0.0, 0.0), (10, 1)),
rf_child2,
timestamps=np.arange(10),
)
vt = np.tile(pt, (10, 4, 1)) - np.array(v0t[np.newaxis, :, :])
np.testing.assert_allclose(vt_act, vt, rtol=1.0, atol=1e-15)
def test_from_dataset(self):
""""""
xr = pytest.importorskip("xarray")
ds = xr.Dataset(
{
"t": (["time", "cartesian_axis"], np.ones((10, 3))),
"r": (["time", "quaternion_axis"], np.ones((10, 4))),
},
{"time": np.arange(10)},
)
rf_world = rbm.ReferenceFrame("world")
rf_child = rbm.ReferenceFrame.from_dataset(
ds, "t", "r", "time", rf_world
)
npt.assert_equal(rf_child.translation, np.ones((10, 3)))
npt.assert_equal(rf_child.rotation, np.ones((10, 4)))
npt.assert_equal(rf_child.timestamps, np.arange(10))
def get_twist(
rf_world, transformer, t=(0.0, 0.0, 0.0), r=(1.0, 0.0, 0.0, 0.0), time=0.0
):
""""""
rf_child1 = rbm.ReferenceFrame(
"child1", parent=rf_world, translation=t, rotation=r
)
transformer.setTransform(static_rf_to_transform_msg(rf_child1, time=time))
try:
v, w = transformer.lookupTwist(
"child1", "world", rospy.Time(), rospy.Duration(0.01)
)
except tf.ExtrapolationException:
v, w = (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)
index = pd.MultiIndex.from_tuples(
list(product(("t", "v", "w"), ("x", "y", "z")))
+ ["rw", "rx", "ry", "rz"],
names=["var", "dim"],
)
return pd.Series(np.hstack((t, v, w, r)), index=index)
def test_transform_quaternions(self, transform_grid, get_rf_tree):
""""""
o, ot, p, pt, rc1, rc2, tc1, tc2 = transform_grid
_, rf_child1, rf_child2 = get_rf_tree(tc1, rc1, tc2, rc2)
rf_child3 = rbm.ReferenceFrame(
"child3", rf_child1, timestamps=np.arange(5) + 2.5
)
# static reference frame + single quaternion
ot_act = rf_child1.transform_quaternions(o, rf_child2)
npt.assert_allclose(np.abs(ot_act), np.abs(ot), rtol=1.0)
# moving reference frame + single quaternion
ot_act = rf_child3.transform_quaternions(o, rf_child2)
np.testing.assert_allclose(
np.abs(ot_act), np.tile(np.abs(ot), (5, 1)), rtol=1.0
)
# moving reference frame + multiple vectors
ot_act = rf_child3.transform_quaternions(
np.tile(o, (10, 1)), rf_child2, timestamps=np.arange(10)
)
np.testing.assert_allclose(
np.abs(ot_act), np.tile(np.abs(ot), (4, 1)), rtol=1.0
)
# moving reference frame + multiple n-dimensional points
ot_act = rf_child3.transform_quaternions(
np.tile(o, (10, 10, 1)),
rf_child2,
timestamps=np.arange(10),
time_axis=1,
)
np.testing.assert_allclose(
np.abs(ot_act), np.tile(np.abs(ot), (10, 4, 1)), rtol=1.0
)
rc2=mock_quaternion(ry, rz, rx),
o=mock_quaternion(rz, rx, ry),
)
for x, y, z, rx, ry, rz in it
)
df.to_csv(test_data_dir / "transform_test_grid.csv")
# twist sequences
twist_rg = np.linspace(0.0, 10.0, 100)
it = list(
product(twist_rg, (0.0,), (0.0,), rg * np.pi / 3, (0.0,), (0.0,))
)
times = np.arange(len(it))
rf_world = rbm.ReferenceFrame("world")
transformer = tf.TransformerROS(True)
df = pd.DataFrame(
get_twist(
rf_world,
transformer,
t=(x, y, z),
r=mock_quaternion(rx, ry, rz),
time=time,
)
for (x, y, z, rx, ry, rz), time in zip(it, times)
)
df.index = pd.to_datetime(times, unit="s")
df.to_csv(test_data_dir / "twist_test_grid.csv")
def test_str(self):
""""""
rf_world = rbm.ReferenceFrame("world")
assert str(rf_world) == "<ReferenceFrame 'world'>"
def test_deregister(self):
""""""
rf_world = rbm.ReferenceFrame("world")
rf_world.register()
rf_world.deregister()
assert "world" not in rbm.registry