[ API reference ] [ PyPI ]
jaxlie
is a library containing implementations of Lie groups commonly used for
rigid body transformations, targeted at computer vision & robotics
applications written in JAX. Heavily inspired by the C++ library
Sophus.
We implement Lie groups as high-level (data)classes:
Group | Description | Parameterization |
---|---|---|
jaxlie.SO2 |
Rotations in 2D. | (real, imaginary): unit complex (โ S1) |
jaxlie.SE2 |
Proper rigid transforms in 2D. | (real, imaginary, x, y): unit complex & translation |
jaxlie.SO3 |
Rotations in 3D. | (qw, qx, qy, qz): wxyz quaternion (โ S3) |
jaxlie.SE3 |
Proper rigid transforms in 3D. | (qw, qx, qy, qz, x, y, z): wxyz quaternion & translation |
Where each group supports:
- Forward- and reverse-mode AD-friendly
exp()
,log()
,adjoint()
,apply()
,multiply()
,inverse()
,identity()
,from_matrix()
, andas_matrix()
operations. - Helpers + analytical Jacobians for manifold optimization
(
jaxlie.manifold
). - (Un)flattening as pytree nodes.
- Serialization using flax.
- Compatibility with standard JAX function transformations. (we've included some
examples for use with
jax.vmap
)
We also implement various common utilities for things like uniform random
sampling (sample_uniform()
) and converting from/to Euler angles (in the
SO3
class).
# Python 3.6 releases also exist, but are no longer being updated.
pip install jaxlie
import numpy as onp
from jaxlie import SE3
#############################
# (1) Constructing transforms.
#############################
# We can compute a w<-b transform by integrating over an se(3) screw, equivalent
# to `SE3.from_matrix(expm(wedge(twist)))`.
twist = onp.array([1.0, 0.0, 0.2, 0.0, 0.5, 0.0])
T_w_b = SE3.exp(twist)
# We can print the (quaternion) rotation term; this is an `SO3` object:
print(T_w_b.rotation())
# Or print the translation; this is a simple array with shape (3,):
print(T_w_b.translation())
# Or the underlying parameters; this is a length-7 (quaternion, translation) array:
print(T_w_b.wxyz_xyz) # SE3-specific field.
print(T_w_b.parameters()) # Helper shared by all groups.
# There are also other helpers to generate transforms, eg from matrices:
T_w_b = SE3.from_matrix(T_w_b.as_matrix())
# Or from explicit rotation and translation terms:
T_w_b = SE3.from_rotation_and_translation(
rotation=T_w_b.rotation(),
translation=T_w_b.translation(),
)
# Or with the dataclass constructor + the underlying length-7 parameterization:
T_w_b = SE3(wxyz_xyz=T_w_b.wxyz_xyz)
#############################
# (2) Applying transforms.
#############################
# Transform points with the `@` operator:
p_b = onp.random.randn(3)
p_w = T_w_b @ p_b
print(p_w)
# or `.apply()`:
p_w = T_w_b.apply(p_b)
print(p_w)
# or the homogeneous matrix form:
p_w = (T_w_b.as_matrix() @ onp.append(p_b, 1.0))[:-1]
print(p_w)
#############################
# (3) Composing transforms.
#############################
# Compose transforms with the `@` operator:
T_b_a = SE3.identity()
T_w_a = T_w_b @ T_b_a
print(T_w_a)
# or `.multiply()`:
T_w_a = T_w_b.multiply(T_b_a)
print(T_w_a)
#############################
# (4) Misc.
#############################
# Compute inverses:
T_b_w = T_w_b.inverse()
identity = T_w_b @ T_b_w
print(identity)
# Compute adjoints:
adjoint_T_w_b = T_w_b.adjoint()
print(adjoint_T_w_b)
# Recover our twist, equivalent to `vee(logm(T_w_b.as_matrix()))`:
twist_recovered = T_w_b.log()
print(twist_recovered)