`jaxlie`¶

Subpackages¶

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

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

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

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

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):
    ...

jaxlie¶

Subpackages¶

Package Contents¶

Classes¶

Functions¶

`jaxlie`¶