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())