jaxlie

Subpackages

• jaxlie.hints
• jaxlie.manifold
• jaxlie.utils

Package Contents

Classes

MatrixLieGroup	Interface definition for matrix Lie groups.
SEBase	Base class for special Euclidean groups.
SOBase	Base class for special orthogonal groups.
SE2	Special Euclidean group for proper rigid transforms in 2D. Broadcasting
SE3	Special Euclidean group for proper rigid transforms in 3D. Broadcasting
SO2	Special orthogonal group for 2D rotations. Broadcasting rules are the
SO3	Special orthogonal group for 3D rotations. Broadcasting rules are the same as

class jaxlie.MatrixLieGroup(parameters)

Bases: abc.ABC

Interface definition for matrix Lie groups. .. py:attribute:: matrix_dim

type:

ClassVar[int]

Dimension of square matrix output from .as_matrix().

Parameters:

parameters (jax.Array) –

parameters_dim: ClassVar[int]

Dimension of underlying parameters, .parameters().

tangent_dim: ClassVar[int]

Dimension of tangent space.

space_dim: ClassVar[int]

Dimension of coordinates that can be transformed.

__matmul__(other: typing_extensions.Self) → typing_extensions.Self

__matmul__(other: jaxlie.hints.Array) → jax.Array

Overload for the @ operator.

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

abstract classmethod identity(batch_axes=())

Returns identity element.

Parameters:

batch_axes (Tuple[int, Ellipsis]) – Any leading batch axes for the output transform.

Returns:

Identity element.

Return type:

typing_extensions.Self

abstract classmethod from_matrix(matrix)

Get group member from matrix representation.

Parameters:

matrix (jaxlie.hints.Array) – Matrix representaiton.

Returns:

Group member.

Return type:

typing_extensions.Self

abstract as_matrix()

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

abstract parameters()

Get underlying representation.

Return type:

jax.Array

abstract apply(target)

Applies group action to a point.

Parameters:

target (jaxlie.hints.Array) – Point to transform.

Returns:

Transformed point.

Return type:

jax.Array

abstract multiply(other)

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (typing_extensions.Self) –

Return type:

typing_extensions.Self

abstract classmethod exp(tangent)

Computes expm(wedge(tangent)).

Parameters:

tangent (jaxlie.hints.Array) – Tangent vector to take the exponential of.

Returns:

Output.

Return type:

typing_extensions.Self

abstract log()

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

abstract adjoint()

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

More precisely, for a transform GroupType:

GroupType @ exp(omega) = exp(Adj_T @ omega) @ GroupType

In robotics, typically used for transforming twists, wrenches, and Jacobians across different reference frames.

Returns:

Output. Shape should be (tangent_dim, tangent_dim).

Return type:

jax.Array

abstract inverse()

Computes the inverse of our transform.

Returns:

Output.

Return type:

typing_extensions.Self

abstract normalize()

Normalize/projects values and returns.

Returns:

Normalized group member.

Return type:

typing_extensions.Self

abstract jlog()

Computes the Jacobian of the logarithm of the group element when a local perturbation is applied.

This is equivalent to the inverse of the right Jacobian, or:

jax.jacrev(lambda x: (T @ exp(x)).log())(jnp.zeros(tangent_dim))

where T is the group element and exp(x) is the tangent vector.

Returns:

The Jacobian of the logarithm, having the dimensions (tangent_dim, tangent_dim,) or batch of these Jacobians.

Return type:

jax.Array

abstract classmethod sample_uniform(key, batch_axes=())

Draw a uniform sample from the group. Translations (if applicable) are in the range [-1, 1].

Parameters:

• key (jax.Array) – PRNG key, as returned by jax.random.PRNGKey().

• batch_axes (Tuple[int, Ellipsis]) – Any leading batch axes for the output transforms. Each sampled transform will be different.

Returns:

Sampled group member.

Return type:

typing_extensions.Self

get_batch_axes()

Return any leading batch axes in contained parameters. If an array of shape (100, 4) is placed in the wxyz field of an SO3 object, for example, this will return (100,).

Return type:

Tuple[int, Ellipsis]

class jaxlie.SEBase(parameters)

Bases: Generic[ContainedSOType], MatrixLieGroup

Base class for special Euclidean groups.

Each SE(N) group member contains an SO(N) rotation, as well as an N-dimensional translation vector. .. py:method:: from_rotation_and_translation(rotation, translation)

classmethod:

abstractmethod:

Construct a rigid transform from a rotation and a translation.

param rotation:

Rotation term.

param translation:

translation term.

returns:

Constructed transformation.

Parameters:

parameters (jax.Array) –

classmethod from_rotation(rotation)

Parameters:

rotation (ContainedSOType) –

Return type:

typing_extensions.Self

classmethod from_translation(translation)

Parameters:

translation (jaxlie.hints.Array) –

Return type:

typing_extensions.Self

abstract rotation()

Returns a transform’s rotation term.

Return type:

ContainedSOType

abstract translation()

Returns a transform’s translation term.

Return type:

jax.Array

apply(target)

Applies group action to a point.

Parameters:

target (jaxlie.hints.Array) – Point to transform.

Returns:

Transformed point.

Return type:

jax.Array

multiply(other)

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (typing_extensions.Self) –

Return type:

typing_extensions.Self

inverse()

Computes the inverse of our transform.

Returns:

Output.

Return type:

typing_extensions.Self

normalize()

Normalize/projects values and returns.

Returns:

Normalized group member.

Return type:

typing_extensions.Self

class jaxlie.SOBase(parameters)

Bases: MatrixLieGroup

Base class for special orthogonal groups.

Parameters:

parameters (jax.Array) –

class jaxlie.SE2(parameters)

Bases: jaxlie.SEBase[jaxlie.SO2]

Special Euclidean group for proper rigid transforms in 2D. Broadcasting rules are the same as for numpy.

Internal parameterization is (cos, sin, x, y). Tangent parameterization is (vx, vy, omega). .. py:attribute:: unit_complex_xy

type:

jax.Array

Internal parameters. (cos, sin, x, y). Shape should be (*, 4).

Parameters:

parameters (jax.Array) –

__repr__()

Return repr(self).

Return type:

str

static from_xy_theta(x, y, theta)

Construct a transformation from standard 2D pose parameters.

Note that this is not the same as integrating over a length-3 twist.

Parameters:

• x (jaxlie.hints.Scalar) –

• y (jaxlie.hints.Scalar) –

• theta (jaxlie.hints.Scalar) –

Return type:

SE2

classmethod from_rotation_and_translation(rotation, translation)

Construct a rigid transform from a rotation and a translation.

Parameters:

• rotation (jaxlie.SO2) – Rotation term.

• translation (jaxlie.hints.Array) – translation term.

Returns:

Constructed transformation.

Return type:

SE2

rotation()

Returns a transform’s rotation term.

Return type:

jaxlie.SO2

translation()

Returns a transform’s translation term.

Return type:

jax.Array

classmethod identity(batch_axes=())

Returns identity element.

Parameters:

batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transform.

Returns:

Identity element.

Return type:

SE2

classmethod from_matrix(matrix)

Get group member from matrix representation.

Parameters:

matrix (jaxlie.hints.Array) – Matrix representaiton.

Returns:

Group member.

Return type:

SE2

parameters()

Get underlying representation.

Return type:

jax.Array

as_matrix()

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

classmethod exp(tangent)

Computes expm(wedge(tangent)).

Parameters:

tangent (jaxlie.hints.Array) – Tangent vector to take the exponential of.

Returns:

Output.

Return type:

SE2

log()

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()

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

More precisely, for a transform GroupType:

GroupType @ exp(omega) = exp(Adj_T @ omega) @ GroupType

In robotics, typically used for transforming twists, wrenches, and Jacobians across different reference frames.

Returns:

Output. Shape should be (tangent_dim, tangent_dim).

Return type:

jax.Array

jlog()

Computes the Jacobian of the logarithm of the group element when a local perturbation is applied.

This is equivalent to the inverse of the right Jacobian, or:

jax.jacrev(lambda x: (T @ exp(x)).log())(jnp.zeros(tangent_dim))

where T is the group element and exp(x) is the tangent vector.

Returns:

The Jacobian of the logarithm, having the dimensions (tangent_dim, tangent_dim,) or batch of these Jacobians.

Return type:

jax.Array

classmethod sample_uniform(key, batch_axes=())

Draw a uniform sample from the group. Translations (if applicable) are in the range [-1, 1].

Parameters:

• key (jax.Array) – PRNG key, as returned by jax.random.PRNGKey().

• batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transforms. Each sampled transform will be different.

Returns:

Sampled group member.

Return type:

SE2

class jaxlie.SE3(parameters)

Bases: jaxlie.SEBase[jaxlie.SO3]

Special Euclidean group for proper rigid transforms in 3D. Broadcasting rules are the same as for numpy.

Internal parameterization is (qw, qx, qy, qz, x, y, z). Tangent parameterization is (vx, vy, vz, omega_x, omega_y, omega_z). .. py:attribute:: wxyz_xyz

type:

jax.Array

Internal parameters. wxyz quaternion followed by xyz translation. Shape should be (*, 7).

Parameters:

parameters (jax.Array) –

__repr__()

Return repr(self).

Return type:

str

classmethod from_rotation_and_translation(rotation, translation)

Construct a rigid transform from a rotation and a translation.

Parameters:

• rotation (jaxlie.SO3) – Rotation term.

• translation (jaxlie.hints.Array) – translation term.

Returns:

Constructed transformation.

Return type:

SE3

rotation()

Returns a transform’s rotation term.

Return type:

jaxlie.SO3

translation()

Returns a transform’s translation term.

Return type:

jax.Array

classmethod identity(batch_axes=())

Returns identity element.

Parameters:

batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transform.

Returns:

Identity element.

Return type:

SE3

classmethod from_matrix(matrix)

Get group member from matrix representation.

Parameters:

matrix (jaxlie.hints.Array) – Matrix representaiton.

Returns:

Group member.

Return type:

SE3

as_matrix()

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

parameters()

Get underlying representation.

Return type:

jax.Array

classmethod exp(tangent)

Computes expm(wedge(tangent)).

Parameters:

tangent (jaxlie.hints.Array) – Tangent vector to take the exponential of.

Returns:

Output.

Return type:

SE3

log()

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()

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

More precisely, for a transform GroupType:

GroupType @ exp(omega) = exp(Adj_T @ omega) @ GroupType

In robotics, typically used for transforming twists, wrenches, and Jacobians across different reference frames.

Returns:

Output. Shape should be (tangent_dim, tangent_dim).

Return type:

jax.Array

jlog()

Computes the Jacobian of the logarithm of the group element when a local perturbation is applied.

This is equivalent to the inverse of the right Jacobian, or:

jax.jacrev(lambda x: (T @ exp(x)).log())(jnp.zeros(tangent_dim))

where T is the group element and exp(x) is the tangent vector.

Returns:

The Jacobian of the logarithm, having the dimensions (tangent_dim, tangent_dim,) or batch of these Jacobians.

Return type:

jax.Array

classmethod sample_uniform(key, batch_axes=())

Draw a uniform sample from the group. Translations (if applicable) are in the range [-1, 1].

Parameters:

• key (jax.Array) – PRNG key, as returned by jax.random.PRNGKey().

• batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transforms. Each sampled transform will be different.

Returns:

Sampled group member.

Return type:

SE3

class jaxlie.SO2(parameters)

Bases: jaxlie.SOBase

Special orthogonal group for 2D rotations. Broadcasting rules are the same as for numpy.

Internal parameterization is (cos, sin). Tangent parameterization is (omega,). .. py:attribute:: unit_complex

type:

jax.Array

Internal parameters. (cos, sin). Shape should be (*, 2).

Parameters:

parameters (jax.Array) –

__repr__()

Return repr(self).

Return type:

str

static from_radians(theta)

Construct a rotation object from a scalar angle.

Parameters:

theta (jaxlie.hints.Scalar) –

Return type:

SO2

as_radians()

Compute a scalar angle from a rotation object.

Return type:

jax.Array

classmethod identity(batch_axes=())

Returns identity element.

Parameters:

batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transform.

Returns:

Identity element.

Return type:

SO2

classmethod from_matrix(matrix)

Get group member from matrix representation.

Parameters:

matrix (jaxlie.hints.Array) – Matrix representaiton.

Returns:

Group member.

Return type:

SO2

as_matrix()

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

parameters()

Get underlying representation.

Return type:

jax.Array

apply(target)

Applies group action to a point.

Parameters:

target (jaxlie.hints.Array) – Point to transform.

Returns:

Transformed point.

Return type:

jax.Array

multiply(other)

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (SO2) –

Return type:

SO2

classmethod exp(tangent)

Computes expm(wedge(tangent)).

Parameters:

tangent (jaxlie.hints.Array) – Tangent vector to take the exponential of.

Returns:

Output.

Return type:

SO2

log()

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()

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

More precisely, for a transform GroupType:

GroupType @ exp(omega) = exp(Adj_T @ omega) @ GroupType

In robotics, typically used for transforming twists, wrenches, and Jacobians across different reference frames.

Returns:

Output. Shape should be (tangent_dim, tangent_dim).

Return type:

jax.Array

inverse()

Computes the inverse of our transform.

Returns:

Output.

Return type:

SO2

normalize()

Normalize/projects values and returns.

Returns:

Normalized group member.

Return type:

SO2

jlog()

Computes the Jacobian of the logarithm of the group element when a local perturbation is applied.

This is equivalent to the inverse of the right Jacobian, or:

jax.jacrev(lambda x: (T @ exp(x)).log())(jnp.zeros(tangent_dim))

where T is the group element and exp(x) is the tangent vector.

Returns:

The Jacobian of the logarithm, having the dimensions (tangent_dim, tangent_dim,) or batch of these Jacobians.

Return type:

jax.Array

classmethod sample_uniform(key, batch_axes=())

Draw a uniform sample from the group. Translations (if applicable) are in the range [-1, 1].

Parameters:

• key (jax.Array) – PRNG key, as returned by jax.random.PRNGKey().

• batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transforms. Each sampled transform will be different.

Returns:

Sampled group member.

Return type:

SO2

class jaxlie.SO3(parameters)

Bases: jaxlie.SOBase

Special orthogonal group for 3D rotations. Broadcasting rules are the same as for numpy.

Internal parameterization is (qw, qx, qy, qz). Tangent parameterization is (omega_x, omega_y, omega_z). .. py:attribute:: wxyz

type:

jax.Array

Internal parameters. (w, x, y, z) quaternion. Shape should be (*, 4).

Parameters:

parameters (jax.Array) –

__repr__()

Return repr(self).

Return type:

str

static from_x_radians(theta)

Generates a x-axis rotation.

Parameters:

• angle – X rotation, in radians.

• theta (jaxlie.hints.Scalar) –

Returns:

Output.

Return type:

SO3

static from_y_radians(theta)

Generates a y-axis rotation.

Parameters:

• angle – Y rotation, in radians.

• theta (jaxlie.hints.Scalar) –

Returns:

Output.

Return type:

SO3

static from_z_radians(theta)

Generates a z-axis rotation.

Parameters:

• angle – Z rotation, in radians.

• theta (jaxlie.hints.Scalar) –

Returns:

Output.

Return type:

SO3

static from_rpy_radians(roll, pitch, yaw)

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

Parameters:

• roll (jaxlie.hints.Scalar) – X rotation, in radians. Applied first.

• pitch (jaxlie.hints.Scalar) – Y rotation, in radians. Applied second.

• yaw (jaxlie.hints.Scalar) – Z rotation, in radians. Applied last.

Returns:

Output.

Return type:

SO3

static from_quaternion_xyzw(xyzw)

Construct a rotation from an xyzw quaternion.

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

Parameters:

xyzw (jaxlie.hints.Array) – xyzw quaternion. Shape should be (*, 4).

Returns:

Output.

Return type:

SO3

as_quaternion_xyzw()

Grab parameters as xyzw quaternion.

Return type:

jax.Array

as_rpy_radians()

Computes roll, pitch, and yaw angles. Uses the ZYX mobile robot convention.

Returns:

Named tuple containing Euler angles in radians.

Return type:

jaxlie.hints.RollPitchYaw

compute_roll_radians()

Compute roll angle. Uses the ZYX mobile robot convention.

Returns:

Euler angle in radians.

Return type:

jax.Array

compute_pitch_radians()

Compute pitch angle. Uses the ZYX mobile robot convention.

Returns:

Euler angle in radians.

Return type:

jax.Array

compute_yaw_radians()

Compute yaw angle. Uses the ZYX mobile robot convention.

Returns:

Euler angle in radians.

Return type:

jax.Array

classmethod identity(batch_axes=())

Returns identity element.

Parameters:

batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transform.

Returns:

Identity element.

Return type:

SO3

classmethod from_matrix(matrix)

Get group member from matrix representation.

Parameters:

matrix (jaxlie.hints.Array) – Matrix representaiton.

Returns:

Group member.

Return type:

SO3

as_matrix()

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

parameters()

Get underlying representation.

Return type:

jax.Array

apply(target)

Applies group action to a point.

Parameters:

target (jaxlie.hints.Array) – Point to transform.

Returns:

Transformed point.

Return type:

jax.Array

multiply(other)

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (SO3) –

Return type:

SO3

classmethod exp(tangent)

Computes expm(wedge(tangent)).

Parameters:

tangent (jaxlie.hints.Array) – Tangent vector to take the exponential of.

Returns:

Output.

Return type:

SO3

log()

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()

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

More precisely, for a transform GroupType:

GroupType @ exp(omega) = exp(Adj_T @ omega) @ GroupType

In robotics, typically used for transforming twists, wrenches, and Jacobians across different reference frames.

Returns:

Output. Shape should be (tangent_dim, tangent_dim).

Return type:

jax.Array

inverse()

Computes the inverse of our transform.

Returns:

Output.

Return type:

SO3

normalize()

Normalize/projects values and returns.

Returns:

Normalized group member.

Return type:

SO3

jlog()

Computes the Jacobian of the logarithm of the group element when a local perturbation is applied.

This is equivalent to the inverse of the right Jacobian, or:

jax.jacrev(lambda x: (T @ exp(x)).log())(jnp.zeros(tangent_dim))

where T is the group element and exp(x) is the tangent vector.

Returns:

The Jacobian of the logarithm, having the dimensions (tangent_dim, tangent_dim,) or batch of these Jacobians.

Return type:

jax.Array

classmethod sample_uniform(key, batch_axes=())

Draw a uniform sample from the group. Translations (if applicable) are in the range [-1, 1].

Parameters:

• key (jax.Array) – PRNG key, as returned by jax.random.PRNGKey().

• batch_axes (jax_dataclasses.Static[Tuple[int, Ellipsis]]) – Any leading batch axes for the output transforms. Each sampled transform will be different.

Returns:

Sampled group member.

Return type:

SO3