jaxlie

Package Contents

Classes

MatrixLieGroup

Interface definition for matrix Lie groups.

SE2

Special Euclidean group for proper rigid transforms in 2D.

SE3

Special Euclidean group for proper rigid transforms in 3D.

SO2

Special orthogonal group for 2D rotations.

SO3

Special orthogonal group for 3D rotations.

Functions

register_lie_group(*, matrix_dim: int, parameters_dim: int, tangent_dim: int, space_dim: int) → Callable[[Type[T]], Type[T]]

Decorator for registering Lie group dataclasses.

class jaxlie.MatrixLieGroup(parameters: jnp.ndarray)[source]

Bases: abc.ABC

Inheritance diagram of jaxlie.MatrixLieGroup

Interface definition for matrix Lie groups.

matrix_dim :int = 0

Dimension of square matrix output from .as_matrix().

parameters_dim :int = 0

Dimension of underlying parameters, .parameters.

tangent_dim :int = 0

Dimension of tangent space.

space_dim :int = 0

Dimension of coordinates that can be transformed.

__matmul__(self: T, other: T) → T[source]
__matmul__(self: T, other: types.Vector)types.Vector

Overload for the @ operator.

Switches between the group action (.apply()) and multiplication (.multiply()) based on the type of other.

abstract static identity()MatrixLieGroup[source]

Returns identity element.

Returns

types.Matrix – Identity.

abstract static from_matrix(matrix: types.Matrix)MatrixLieGroup[source]

Get group member from matrix representation.

Parameters

matrix (jnp.ndarray) – types.Matrix representaiton.

Returns

T – Group member.

abstract as_matrix(self)types.Matrix[source]

Get transformation as a matrix. Homogeneous for SE groups.

property parameters(self)types.Vector

Get underlying representation.

abstract apply(self: T, target: types.Vector)types.Vector[source]

Applies the group action.

Parameters

target (types.Vector) – types.Vector to transform.

Returns

types.Vector – Transformed vector.

abstract multiply(self: T, other: T) → T[source]

Left-multiplies this transformations with another.

Parameters

other (T) – other

Returns

T – self @ other

abstract static exp(tangent: types.TangentVector)MatrixLieGroup[source]

Computes expm(wedge(tangent)).

Parameters

tangent (types.TangentVector) – Input.

Returns

MatrixLieGroup – Output.

abstract log(self: T)types.TangentVector[source]

Computes vee(logm(transformation matrix)).

Returns

types.TangentVector – Output. Shape should be (tangent_dim,).

abstract adjoint(self: T)types.Matrix[source]

Computes the adjoint, which transforms tangent vectors between tangent spaces.

More precisely, for a transform T:

T @ exp(omega) = exp(Adj_T @ omega) @ T

For robotics, typically used for converting twists, wrenches, and Jacobians between our spatial and body representations.

Returns

types.Matrix – Output. Shape should be (tangent_dim, tangent_dim).

abstract inverse(self: T) → T[source]

Computes the inverse of our transform.

Returns

types.Matrix – Output.

abstract normalize(self: T) → T[source]

Normalize/projects values and returns.

Returns

T – Normalized group member.

class jaxlie.SE2(parameters: jnp.ndarray)[source]

Bases: jaxlie.MatrixLieGroup

Inheritance diagram of jaxlie.SE2

Special Euclidean group for proper rigid transforms in 2D.

xy_unit_complex :types.Vector

Internal parameters. (x, y, cos, sin).

__repr__(self) → str[source]

Return repr(self).

static from_xy_theta(x: float, y: float, theta: float)SE2[source]
static from_rotation_and_translation(rotation: SO2, translation: types.Vector)SE2[source]
property rotation(self)SO2
property translation(self)types.Vector
static identity()SE2[source]

Returns identity element.

Returns

types.Matrix – Identity.

static from_matrix(matrix: types.Matrix)SE2[source]

Get group member from matrix representation.

Parameters

matrix (jnp.ndarray) – types.Matrix representaiton.

Returns

T – Group member.

property parameters(self)types.Vector

Get underlying representation.

as_matrix(self)types.Matrix[source]

Get transformation as a matrix. Homogeneous for SE groups.

apply(self: SE2, target: types.Vector)types.Vector[source]

Applies the group action.

Parameters

target (types.Vector) – types.Vector to transform.

Returns

types.Vector – Transformed vector.

multiply(self: SE2, other: SE2)SE2[source]

Left-multiplies this transformations with another.

Parameters

other (T) – other

Returns

T – self @ other

static exp(tangent: types.TangentVector)SE2[source]

Computes expm(wedge(tangent)).

Parameters

tangent (types.TangentVector) – Input.

Returns

MatrixLieGroup – Output.

log(self: SE2)types.TangentVector[source]

Computes vee(logm(transformation matrix)).

Returns

types.TangentVector – Output. Shape should be (tangent_dim,).

adjoint(self: SE2)types.Matrix[source]

Computes the adjoint, which transforms tangent vectors between tangent spaces.

More precisely, for a transform T:

T @ exp(omega) = exp(Adj_T @ omega) @ T

For robotics, typically used for converting twists, wrenches, and Jacobians between our spatial and body representations.

Returns

types.Matrix – Output. Shape should be (tangent_dim, tangent_dim).

inverse(self: SE2)SE2[source]

Computes the inverse of our transform.

Returns

types.Matrix – Output.

normalize(self: SE2)SE2[source]

Normalize/projects values and returns.

Returns

T – Normalized group member.

class jaxlie.SE3(parameters: jnp.ndarray)[source]

Bases: jaxlie.MatrixLieGroup

Inheritance diagram of jaxlie.SE3

Special Euclidean group for proper rigid transforms in 3D.

xyz_wxyz :types.Vector

Internal parameters. Length-3 translation followed by wxyz quaternion.

__repr__(self) → str[source]

Return repr(self).

static from_rotation_and_translation(rotation: SO3, translation: types.Vector)SE3[source]
property rotation(self)SO3
property translation(self)types.Vector
static identity()SE3[source]

Returns identity element.

Returns

types.Matrix – Identity.

static from_matrix(matrix: types.Matrix)SE3[source]

Get group member from matrix representation.

Parameters

matrix (jnp.ndarray) – types.Matrix representaiton.

Returns

T – Group member.

as_matrix(self)types.Matrix[source]

Get transformation as a matrix. Homogeneous for SE groups.

property parameters(self)types.Vector

Get underlying representation.

apply(self: SE3, target: types.Vector)types.Vector[source]

Applies the group action.

Parameters

target (types.Vector) – types.Vector to transform.

Returns

types.Vector – Transformed vector.

multiply(self: SE3, other: SE3)SE3[source]

Left-multiplies this transformations with another.

Parameters

other (T) – other

Returns

T – self @ other

static exp(tangent: types.TangentVector)SE3[source]

Computes expm(wedge(tangent)).

Parameters

tangent (types.TangentVector) – Input.

Returns

MatrixLieGroup – Output.

log(self: SE3)types.TangentVector[source]

Computes vee(logm(transformation matrix)).

Returns

types.TangentVector – Output. Shape should be (tangent_dim,).

adjoint(self: SE3)types.Matrix[source]

Computes the adjoint, which transforms tangent vectors between tangent spaces.

More precisely, for a transform T:

T @ exp(omega) = exp(Adj_T @ omega) @ T

For robotics, typically used for converting twists, wrenches, and Jacobians between our spatial and body representations.

Returns

types.Matrix – Output. Shape should be (tangent_dim, tangent_dim).

inverse(self: SE3)SE3[source]

Computes the inverse of our transform.

Returns

types.Matrix – Output.

normalize(self: SE3)SE3[source]

Normalize/projects values and returns.

Returns

T – Normalized group member.

class jaxlie.SO2(parameters: jnp.ndarray)[source]

Bases: jaxlie.MatrixLieGroup

Inheritance diagram of jaxlie.SO2

Special orthogonal group for 2D rotations.

unit_complex :types.Vector

Internal parameters. (cos, sin).

__repr__(self) → str[source]

Return repr(self).

static from_radians(theta: float)SO2[source]
to_radians(self) → float[source]
static identity()SO2[source]

Returns identity element.

Returns

types.Matrix – Identity.

static from_matrix(matrix: types.Matrix)SO2[source]

Get group member from matrix representation.

Parameters

matrix (jnp.ndarray) – types.Matrix representaiton.

Returns

T – Group member.

as_matrix(self)types.Matrix[source]

Get transformation as a matrix. Homogeneous for SE groups.

property parameters(self)types.Vector

Get underlying representation.

apply(self: SO2, target: types.Vector)types.Vector[source]

Applies the group action.

Parameters

target (types.Vector) – types.Vector to transform.

Returns

types.Vector – Transformed vector.

multiply(self: SO2, other: SO2)SO2[source]

Left-multiplies this transformations with another.

Parameters

other (T) – other

Returns

T – self @ other

static exp(tangent: types.TangentVector)SO2[source]

Computes expm(wedge(tangent)).

Parameters

tangent (types.TangentVector) – Input.

Returns

MatrixLieGroup – Output.

log(self: SO2)types.TangentVector[source]

Computes vee(logm(transformation matrix)).

Returns

types.TangentVector – Output. Shape should be (tangent_dim,).

adjoint(self: SO2)types.Matrix[source]

Computes the adjoint, which transforms tangent vectors between tangent spaces.

More precisely, for a transform T:

T @ exp(omega) = exp(Adj_T @ omega) @ T

For robotics, typically used for converting twists, wrenches, and Jacobians between our spatial and body representations.

Returns

types.Matrix – Output. Shape should be (tangent_dim, tangent_dim).

inverse(self: SO2)SO2[source]

Computes the inverse of our transform.

Returns

types.Matrix – Output.

normalize(self: SO2)SO2[source]

Normalize/projects values and returns.

Returns

T – Normalized group member.

class jaxlie.SO3(parameters: jnp.ndarray)[source]

Bases: jaxlie.MatrixLieGroup

Inheritance diagram of jaxlie.SO3

Special orthogonal group for 3D rotations.

wxyz :types.Vector

Internal parameters. (w, x, y, z) quaternion.

__repr__(self) → str[source]

Return repr(self).

static from_x_radians(theta: float)SO3[source]

Generates a x-axis rotation.

Parameters

angle – X rotation, in radians.

Returns

SO3 – Output.

static from_y_radians(theta: float)SO3[source]

Generates a y-axis rotation.

Parameters

angle – Y rotation, in radians.

Returns

SO3 – Output.

static from_z_radians(theta: float)SO3[source]

Generates a z-axis rotation.

Parameters

angle – Z rotation, in radians.

Returns

SO3 – Output.

static from_rpy_radians(roll: float, pitch: float, yaw: float)SO3[source]

Generates a transform from a set of Euler angles. Uses the ZYX mobile robot convention.

Parameters
  • roll – X rotation, in radians. Applied first.

  • pitch – Y rotation, in radians. Applied second.

  • yaw – Z rotation, in radians. Applied last.

Returns

SO3 – Output.

static from_quaternion_xyzw(xyzw: jnp.ndarray)SO3[source]

Construct a rotation from an xyzw quaternion.

Note that wxyz quaternions can be constructed using the default dataclass constructor.

Parameters

xyzw (jnp.ndarray) – xyzw quaternion. Shape should be (4,).

Returns

SO3 – Output.

as_quaternion_xyzw(self) → jnp.ndarray[source]

Grab parameters as xyzw quaternion.

static identity()SO3[source]

Returns identity element.

Returns

types.Matrix – Identity.

static from_matrix(matrix: types.Matrix)SO3[source]

Get group member from matrix representation.

Parameters

matrix (jnp.ndarray) – types.Matrix representaiton.

Returns

T – Group member.

as_matrix(self)types.Matrix[source]

Get transformation as a matrix. Homogeneous for SE groups.

property parameters(self)types.Vector

Get underlying representation.

apply(self: SO3, target: types.Vector)types.Vector[source]

Applies the group action.

Parameters

target (types.Vector) – types.Vector to transform.

Returns

types.Vector – Transformed vector.

multiply(self: SO3, other: SO3)SO3[source]

Left-multiplies this transformations with another.

Parameters

other (T) – other

Returns

T – self @ other

static exp(tangent: types.TangentVector)SO3[source]

Computes expm(wedge(tangent)).

Parameters

tangent (types.TangentVector) – Input.

Returns

MatrixLieGroup – Output.

log(self: SO3)types.TangentVector[source]

Computes vee(logm(transformation matrix)).

Returns

types.TangentVector – Output. Shape should be (tangent_dim,).

adjoint(self: SO3)types.Matrix[source]

Computes the adjoint, which transforms tangent vectors between tangent spaces.

More precisely, for a transform T:

T @ exp(omega) = exp(Adj_T @ omega) @ T

For robotics, typically used for converting twists, wrenches, and Jacobians between our spatial and body representations.

Returns

types.Matrix – Output. Shape should be (tangent_dim, tangent_dim).

inverse(self: SO3)SO3[source]

Computes the inverse of our transform.

Returns

types.Matrix – Output.

normalize(self: SO3)SO3[source]

Normalize/projects values and returns.

Returns

T – Normalized group member.

jaxlie.register_lie_group(*, matrix_dim: int, parameters_dim: int, tangent_dim: int, space_dim: int) → Callable[[Type[T]], Type[T]][source]

Decorator for registering Lie group dataclasses.

  • Sets static dimensionality attributes

  • Makes the group hashable

  • Marks all functions for JIT compilation

  • Adds flattening/unflattening ops for use as a PyTree node

  • Adds serialization ops for flax.serialization

Note that a significant amount of functionality here could be replaced by flax.struct, but flax.struct doesn’t work very well with jedi or mypy.

Example:

@register_lie_group(
    matrix_dim=2,
    parameters_dim=2,
    tangent_dim=1,
    space_dim=2,
)
@dataclasses.dataclass(frozen=True)
class SO2(_base.MatrixLieGroup):
    ...