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_se3_translation(_random_module):
"""Simple test for SE(3) translation terms."""
translation = onp.random.randn(3)
T = jaxlie.SE3.from_rotation_and_translation(
rotation=jaxlie.SO3.identity(),
translation=translation,
)
assert_arrays_close(T @ translation, translation * 2)
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_se3_rotation():
"""Simple test for SE(3) rotation terms."""
T_w_b = jaxlie.SE3.from_rotation_and_translation(
rotation=jaxlie.SO3.from_rpy_radians(onp.pi / 2.0, 0.0, 0.0),
translation=onp.zeros(3),
)
p_b = onp.array([0.0, 1.0, 0.0])
p_w = onp.array([0.0, 0.0, 1.0])
assert_arrays_close(T_w_b @ p_b, p_w)
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_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 _assert_jacobians_close(
Group: Type[jaxlie.MatrixLieGroup],
f: Callable[[Type[jaxlie.MatrixLieGroup], jnp.ndarray], jnp.ndarray],
primal: jaxlie.hints.Array,
) -> None:
jacobian_fwd = cached_jacfwd(f)(Group, primal)
jacobian_rev = cached_jacrev(f)(Group, primal)
jacobian_numerical = jacnumerical(
lambda primal: cached_jit(f, static_argnums=0)(Group, primal))(primal)
assert_arrays_close(jacobian_fwd, jacobian_rev)
assert_arrays_close(jacobian_fwd, jacobian_numerical, rtol=5e-4, atol=5e-4)
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_se2_translation(_random_module):
"""Simple test for SE(2) translation terms."""
translation = onp.random.randn(2)
T = jaxlie.SE2.from_xy_theta(*translation, theta=0.0)
assert_arrays_close(T @ translation, translation * 2)
def test_so2_from_to_radians_bijective(_random_module):
"""Check that we can convert from and to radians."""
radians = onp.random.uniform(low=-onp.pi, high=onp.pi)
assert_arrays_close(jaxlie.SO2.from_radians(radians).as_radians(), radians)
