def test_se3_compose(_random_module):
"""Compare SE3 composition in matrix form vs compact form."""
T1 = sample_transform(jaxlie.SE3)
T2 = sample_transform(jaxlie.SE3)
assert_arrays_close(T1.as_matrix() @ T2.as_matrix(), (T1 @ T2).as_matrix())
assert_transforms_close(
jaxlie.SE3.from_matrix(T1.as_matrix() @ T2.as_matrix()), T1 @ T2)
def test_rplus_rminus(Group: Type[jaxlie.MatrixLieGroup]):
"""Check closure property."""
T_wa = sample_transform(Group)
T_wb = sample_transform(Group)
T_ab = T_wa.inverse() @ T_wb
assert_transforms_close(jaxlie.manifold.rplus(T_wa, T_ab.log()), T_wb)
assert_arrays_close(jaxlie.manifold.rminus(T_wa, T_wb), T_ab.log())
def test_associative(Group: Type[jaxlie.MatrixLieGroup]):
"""Check associative property."""
transform_a = sample_transform(Group)
transform_b = sample_transform(Group)
transform_c = sample_transform(Group)
assert_transforms_close(
(transform_a @ transform_b) @ transform_c,
transform_a @ (transform_b @ transform_c),
)
def test_multiply(Group: Type[jaxlie.MatrixLieGroup]):
"""Check multiply interfaces."""
T_w_b = sample_transform(Group)
T_b_a = sample_transform(Group)
assert_arrays_close(T_w_b.as_matrix() @ T_w_b.inverse().as_matrix(),
onp.eye(Group.matrix_dim))
assert_arrays_close(T_w_b.as_matrix() @ jnp.linalg.inv(T_w_b.as_matrix()),
onp.eye(Group.matrix_dim))
assert_transforms_close(T_w_b @ T_b_a, Group.multiply(T_w_b, T_b_a))
def test_closure(Group: Type[jaxlie.MatrixLieGroup]):
"""Check closure property."""
transform_a = sample_transform(Group)
transform_b = sample_transform(Group)
composed = transform_a @ transform_b
assert_transforms_close(composed, composed.normalize())
composed = transform_b @ transform_a
assert_transforms_close(composed, composed.normalize())
composed = Group.multiply(transform_a, transform_b)
assert_transforms_close(composed, composed.normalize())
composed = Group.multiply(transform_b, transform_a)
assert_transforms_close(composed, composed.normalize())
def test_serialization_state_dict_bijective(
Group: Type[jaxlie.MatrixLieGroup]):
"""Check bijectivity of state dict representation conversations."""
T = sample_transform(Group)
T_recovered = flax.serialization.from_state_dict(
T, flax.serialization.to_state_dict(T))
assert_transforms_close(T, T_recovered)
def test_adjoint(Group: Type[jaxlie.MatrixLieGroup]):
"""Check adjoint definition."""
transform = sample_transform(Group)
omega = onp.random.randn(Group.tangent_dim)
assert_transforms_close(
transform @ Group.exp(omega),
Group.exp(transform.adjoint() @ omega) @ transform,
)
def test_rplus_jacobian(Group: Type[jaxlie.MatrixLieGroup]):
"""Check analytical rplus Jacobian.."""
T_wa = sample_transform(Group)
J_ours = jaxlie.manifold.rplus_jacobian_parameters_wrt_delta(T_wa)
J_jacfwd = _rplus_jacobian_parameters_wrt_delta(T_wa)
assert_arrays_close(J_ours, J_jacfwd)
def test_log_exp_bijective(Group: Type[jaxlie.MatrixLieGroup]):
"""Check 1-to-1 mapping for log <=> exp operations."""
transform = sample_transform(Group)
tangent = transform.log()
assert tangent.shape == (Group.tangent_dim,)
exp_transform = Group.exp(tangent)
assert_transforms_close(transform, exp_transform)
assert_arrays_close(tangent, exp_transform.log())
def test_identity(Group: Type[jaxlie.MatrixLieGroup]):
"""Check identity property."""
transform = sample_transform(Group)
identity = Group.identity()
assert_transforms_close(transform, identity @ transform)
assert_transforms_close(transform, transform @ identity)
assert_arrays_close(transform.as_matrix(),
identity.as_matrix() @ transform.as_matrix())
assert_arrays_close(transform.as_matrix(),
transform.as_matrix() @ identity.as_matrix())
def test_inverse(Group: Type[jaxlie.MatrixLieGroup]):
"""Check inverse property."""
transform = sample_transform(Group)
identity = Group.identity()
assert_transforms_close(identity, transform @ transform.inverse())
assert_transforms_close(identity, transform.inverse() @ transform)
assert_transforms_close(identity,
Group.multiply(transform, transform.inverse()))
assert_transforms_close(identity,
Group.multiply(transform.inverse(), transform))
assert_arrays_close(
onp.eye(Group.matrix_dim),
transform.as_matrix() @ transform.inverse().as_matrix(),
)
assert_arrays_close(
onp.eye(Group.matrix_dim),
transform.inverse().as_matrix() @ transform.as_matrix(),
)
def test_apply(Group: Type[jaxlie.MatrixLieGroup]):
"""Check group action interfaces."""
T_w_b = sample_transform(Group)
p_b = onp.random.randn(Group.space_dim)
if Group.matrix_dim == Group.space_dim:
assert_arrays_close(
T_w_b @ p_b,
T_w_b.apply(p_b),
T_w_b.as_matrix() @ p_b,
)
else:
# Homogeneous coordinates
assert Group.matrix_dim == Group.space_dim + 1
assert_arrays_close(
T_w_b @ p_b,
T_w_b.apply(p_b),
(T_w_b.as_matrix() @ onp.append(p_b, 1.0))[:-1],
)
def test_so3_xyzw_bijective(_random_module):
"""Check that we can convert between xyzw and wxyz quaternions."""
T = sample_transform(jaxlie.SO3)
assert_transforms_close(T, jaxlie.SO3.from_quaternion_xyzw(T.as_quaternion_xyzw()))
def test_repr(Group: Type[jaxlie.MatrixLieGroup]):
"""Smoke test for __repr__ implementations."""
transform = sample_transform(Group)
print(transform)
def test_matrix_bijective(Group: Type[jaxlie.MatrixLieGroup]):
"""Check that we can convert to and from matrices."""
transform = sample_transform(Group)
assert_transforms_close(transform, Group.from_matrix(transform.as_matrix()))
def test_inverse_bijective(Group: Type[jaxlie.MatrixLieGroup]):
"""Check inverse of inverse."""
transform = sample_transform(Group)
assert_transforms_close(transform, transform.inverse().inverse())
def test_so3_rpy_bijective(_random_module):
"""Check that we can convert between quaternions and Euler angles."""
T = sample_transform(jaxlie.SO3)
assert_transforms_close(T, jaxlie.SO3.from_rpy_radians(*T.as_rpy_radians()))
def test_sample_uniform_valid(Group: Type[jaxlie.MatrixLieGroup]):
"""Check that sample_uniform() returns valid group members."""
T = sample_transform(Group) # Calls sample_uniform under the hood
assert_transforms_close(T, T.normalize())
