jaxlie

Subpackages

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.

SE3

Special Euclidean group for proper rigid transforms in 3D.

SO2

Special orthogonal group for 2D rotations.

SO3

Special orthogonal group for 3D rotations.

class jaxlie.MatrixLieGroup(parameters)[source]

Bases: abc.ABC

Inheritance diagram of jaxlie.MatrixLieGroup

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: GroupType) GroupType[source]
__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()[source]

Returns identity element.

Returns:

Identity element.

Return type:

GroupType

abstract classmethod from_matrix(matrix)[source]

Get group member from matrix representation.

Parameters:

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

Returns:

Group member.

Return type:

GroupType

abstract as_matrix()[source]

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

abstract parameters()[source]

Get underlying representation.

Return type:

jax.Array

abstract apply(target)[source]

Applies group action to a point.

Parameters:

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

Returns:

Transformed point.

Return type:

jax.Array

abstract multiply(other)[source]

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (GroupType) –

Return type:

GroupType

abstract classmethod exp(tangent)[source]

Computes expm(wedge(tangent)).

Parameters:

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

Returns:

Output.

Return type:

GroupType

abstract log()[source]

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

abstract adjoint()[source]

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()[source]

Computes the inverse of our transform.

Returns:

Output.

Return type:

GroupType

abstract normalize()[source]

Normalize/projects values and returns.

Returns:

GroupType – Normalized group member.

Return type:

GroupType

abstract classmethod sample_uniform(key)[source]

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

Returns:

Sampled group member.

Return type:

GroupType

abstract get_batch_axes()[source]

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

This should generally be implemented by jdc.EnforcedAnnotationsMixin.

Return type:

Tuple[int, Ellipsis]

class jaxlie.SEBase(parameters)[source]

Bases: Generic[ContainedSOType], MatrixLieGroup

Inheritance diagram of jaxlie.SEBase

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)[source]
Parameters:

rotation (ContainedSOType) –

Return type:

SEGroupType

abstract rotation()[source]

Returns a transform’s rotation term.

Return type:

ContainedSOType

abstract translation()[source]

Returns a transform’s translation term.

Return type:

jax.Array

apply(target)[source]

Applies group action to a point.

Parameters:

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

Returns:

Transformed point.

Return type:

jax.Array

multiply(other)[source]

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (SEGroupType) –

Return type:

SEGroupType

inverse()[source]

Computes the inverse of our transform.

Returns:

Output.

Return type:

SEGroupType

normalize()[source]

Normalize/projects values and returns.

Returns:

GroupType – Normalized group member.

Return type:

SEGroupType

class jaxlie.SOBase(parameters)[source]

Bases: MatrixLieGroup

Inheritance diagram of jaxlie.SOBase

Base class for special orthogonal groups.

Parameters:

parameters (jax.Array) –

class jaxlie.SE2(parameters)[source]

Bases: jax_dataclasses.EnforcedAnnotationsMixin, jaxlie.SEBase[jaxlie.SO2]

Inheritance diagram of jaxlie.SE2

Special Euclidean group for proper rigid transforms in 2D.

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

type:

typing_extensions.Annotated[jax.Array, (Ellipsis, 4), jax.numpy.floating]

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

Parameters:

parameters (jax.Array) –

__repr__()[source]

Return repr(self).

Return type:

str

static from_xy_theta(x, y, theta)[source]

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

static from_rotation_and_translation(rotation, translation)[source]

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()[source]

Returns a transform’s rotation term.

Return type:

jaxlie.SO2

translation()[source]

Returns a transform’s translation term.

Return type:

jax.Array

static identity()[source]

Returns identity element.

Returns:

Identity element.

Return type:

SE2

static from_matrix(matrix)[source]

Get group member from matrix representation.

Parameters:

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

Returns:

Group member.

Return type:

SE2

parameters()[source]

Get underlying representation.

Return type:

jax.Array

as_matrix()[source]

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

static exp(tangent)[source]

Computes expm(wedge(tangent)).

Parameters:

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

Returns:

Output.

Return type:

SE2

log()[source]

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()[source]

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

static sample_uniform(key)[source]

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

Returns:

Sampled group member.

Return type:

SE2

class jaxlie.SE3(parameters)[source]

Bases: jax_dataclasses.EnforcedAnnotationsMixin, jaxlie.SEBase[jaxlie.SO3]

Inheritance diagram of jaxlie.SE3

Special Euclidean group for proper rigid transforms in 3D.

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:

typing_extensions.Annotated[jax.Array, (Ellipsis, 7), jax.numpy.floating]

Internal parameters. wxyz quaternion followed by xyz translation.

Parameters:

parameters (jax.Array) –

__repr__()[source]

Return repr(self).

Return type:

str

static from_rotation_and_translation(rotation, translation)[source]

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()[source]

Returns a transform’s rotation term.

Return type:

jaxlie.SO3

translation()[source]

Returns a transform’s translation term.

Return type:

jax.Array

static identity()[source]

Returns identity element.

Returns:

Identity element.

Return type:

SE3

static from_matrix(matrix)[source]

Get group member from matrix representation.

Parameters:

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

Returns:

Group member.

Return type:

SE3

as_matrix()[source]

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

parameters()[source]

Get underlying representation.

Return type:

jax.Array

static exp(tangent)[source]

Computes expm(wedge(tangent)).

Parameters:

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

Returns:

Output.

Return type:

SE3

log()[source]

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()[source]

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

static sample_uniform(key)[source]

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

Returns:

Sampled group member.

Return type:

SE3

class jaxlie.SO2(parameters)[source]

Bases: jax_dataclasses.EnforcedAnnotationsMixin, jaxlie.SOBase

Inheritance diagram of jaxlie.SO2

Special orthogonal group for 2D rotations.

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

type:

typing_extensions.Annotated[jax.Array, (Ellipsis, 2), jax.numpy.floating]

Internal parameters. (cos, sin).

Parameters:

parameters (jax.Array) –

__repr__()[source]

Return repr(self).

Return type:

str

static from_radians(theta)[source]

Construct a rotation object from a scalar angle.

Parameters:

theta (jaxlie.hints.Scalar) –

Return type:

SO2

as_radians()[source]

Compute a scalar angle from a rotation object.

Return type:

jax.Array

static identity()[source]

Returns identity element.

Returns:

Identity element.

Return type:

SO2

static from_matrix(matrix)[source]

Get group member from matrix representation.

Parameters:

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

Returns:

Group member.

Return type:

SO2

as_matrix()[source]

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

parameters()[source]

Get underlying representation.

Return type:

jax.Array

apply(target)[source]

Applies group action to a point.

Parameters:

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

Returns:

Transformed point.

Return type:

jax.Array

multiply(other)[source]

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (SO2) –

Return type:

SO2

static exp(tangent)[source]

Computes expm(wedge(tangent)).

Parameters:

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

Returns:

Output.

Return type:

SO2

log()[source]

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()[source]

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()[source]

Computes the inverse of our transform.

Returns:

Output.

Return type:

SO2

normalize()[source]

Normalize/projects values and returns.

Returns:

GroupType – Normalized group member.

Return type:

SO2

static sample_uniform(key)[source]

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

Returns:

Sampled group member.

Return type:

SO2

class jaxlie.SO3(parameters)[source]

Bases: jax_dataclasses.EnforcedAnnotationsMixin, jaxlie.SOBase

Inheritance diagram of jaxlie.SO3

Special orthogonal group for 3D rotations.

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

type:

typing_extensions.Annotated[jax.Array, (Ellipsis, 4), jax.numpy.floating]

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

Parameters:

parameters (jax.Array) –

__repr__()[source]

Return repr(self).

Return type:

str

static from_x_radians(theta)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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()[source]

Grab parameters as xyzw quaternion.

Return type:

jax.Array

as_rpy_radians()[source]

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()[source]

Compute roll angle. Uses the ZYX mobile robot convention.

Returns:

Euler angle in radians.

Return type:

jax.Array

compute_pitch_radians()[source]

Compute pitch angle. Uses the ZYX mobile robot convention.

Returns:

Euler angle in radians.

Return type:

jax.Array

compute_yaw_radians()[source]

Compute yaw angle. Uses the ZYX mobile robot convention.

Returns:

Euler angle in radians.

Return type:

jax.Array

static identity()[source]

Returns identity element.

Returns:

Identity element.

Return type:

SO3

static from_matrix(matrix)[source]

Get group member from matrix representation.

Parameters:

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

Returns:

Group member.

Return type:

SO3

as_matrix()[source]

Get transformation as a matrix. Homogeneous for SE groups.

Return type:

jax.Array

parameters()[source]

Get underlying representation.

Return type:

jax.Array

apply(target)[source]

Applies group action to a point.

Parameters:

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

Returns:

Transformed point.

Return type:

jax.Array

multiply(other)[source]

Composes this transformation with another.

Returns:

self @ other

Parameters:

other (SO3) –

Return type:

SO3

static exp(tangent)[source]

Computes expm(wedge(tangent)).

Parameters:

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

Returns:

Output.

Return type:

SO3

log()[source]

Computes vee(logm(transformation matrix)).

Returns:

Output. Shape should be (tangent_dim,).

Return type:

jax.Array

adjoint()[source]

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()[source]

Computes the inverse of our transform.

Returns:

Output.

Return type:

SO3

normalize()[source]

Normalize/projects values and returns.

Returns:

GroupType – Normalized group member.

Return type:

SO3

static sample_uniform(key)[source]

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

Returns:

Sampled group member.

Return type:

SO3