pymlg
pymlg copied to clipboard
Pure static Lie groups in Numpy, Jax, and C++
PyMLG - Matrix Lie groups with Numpy, Jax, and C++ implementations!
An instantiation-free python package for common matrix Lie group operations implemented as pure static classes. Using pure static classes keeps the usage extremely simple while still allowing for abstraction and inheritance. We do not introduce new objects with stateful behavior that must be learnt. Everything operates directly on arrays/tensors. This allows users to implement their own more sophisticated objects using these classes as back-end mathematical implementations.
Installation
Begin by cloning this repo somewhere. To install, go to the clone directory and run
pip install -e .
Documentation
Documentation can be found here: https://decargroup.github.io/pymlg
Example
from pymlg import SE3
import numpy as np
# Random pose
T = SE3.random()
# R^n to group directly (using "Capital" notation)
x = np.array([0.1, 0.2, 0.3, 4, 5, 6])
T = SE3.Exp(x)
# Group to R^n directly
x = SE3.Log(T)
# Wedge, vee
Xi = SE3.wedge(x)
x = SE3.vee(Xi)
# Actual exp/log maps
T = SE3.exp(Xi)
Xi = SE3.log(T)
# Adjoint matrix representation of group element
A = SE3.adjoint(T)
# Adjoint representation of algebra element
ad = SE3.adjoint_algebra(Xi)
# Inverse of group element
T_inv = SE3.inverse(T)
# Group left/right jacobians, and their inverses
J_L = SE3.left_jacobian(x)
J_R = SE3.right_jacobian(x)
J_L_inv = SE3.left_jacobian_inv(x)
J_R_inv = SE3.right_jacobian_inv(x)
# ... and more.
Using Numpy/C++/Jax
To explicitly access pure numpy implementations use
from pymlg.numpy import SO2, SO3, SE2, SE3, SE23
To explicitly access classes which internally use C++ use
from pymlg.cpp import SO3, SE3, SE23
To explicitly access Jax implementations use
from pymlg.jax import SE2
Currently, only SO3, SE3, and SE23 are implemented in C++, with the functions accepting and returning numpy arrays. They are also the default internal implementations when simply using from pymlg import SO3, SE3, SE23. For the JAX implementation, the return types will be jax.numpy arrays. All operations in the jax implementation can be JIT-compiled.
For all implementations (jax, numpy, C++), the user API is exactly the same! This means that by changing the import statement the example still works.
Note: functions which output "vectors", such as SE2.Log(T) all return a 2D numpy array with dimensions (n, 1).
Running Tests
If you use VS Code, you should be able to enable the VS Code testing feature using pytest. Otherwise, you can run tests from the command line when inside this folder using
pytest tests
Credit to UTIAS's STARS group
Some specific implementations came from the UTIAS STARS Lie group package.. We wanted a different API and variable ordering, which led to us making our own package. Eventually, this repo evolved to contain more groups, as well as Jax and C++ implementations.