Quaternion Class

class fastquat.quaternion.Quaternion(w: float | Array, x: float | Array, y: float | Array, z: float | Array)

Bases: object

Class for manipulating quaternion tensors with JAX.

A quaternion is represented by [w, x, y, z] where w is the scalar part and (x, y, z) is the vector part.

__init__(w: float | Array, x: float | Array, y: float | Array, z: float | Array) → None

Initialize a tensor of quaternions.

Parameters:

• w – components of the quaternions.

• x – components of the quaternions.

• y – components of the quaternions.

• z – components of the quaternions.

tree_flatten() → tuple[tuple[Any, ...], Any]

Flatten the Quaternion PyTree.

classmethod tree_unflatten(aux_data, children) → Quaternion

Unflatten The Quaternion PyTree

classmethod from_array(array: Array) → Quaternion

Create a Quaternion array from a numeric array of shape (…, 4).

Parameters:

array – array of shape (…, 4) where the last dimension is [w, x, y, z]

classmethod from_scalar_vector(scalar: Array, vector: Array) → Quaternion

Create a quaternion from scalar and vector parts.

Parameters:

• scalar – Array of shape (…,) for the scalar part.

• vector – Array of shape (…, 3) for the vector part.

Returns:

Quaternion

classmethod from_rotation_matrix(rot: Array) → Quaternion

Create the quaternion associated to a rotation matrix.

Parameters:

rot – Array of shape (…, 3, 3) representing the rotation matrix

Returns:

The normalized Quaternion tensor representing the rotation matrix.

classmethod zeros(shape: tuple[int, ...] = (), dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Quaternion

Create quaternions with all components set to 0.

Parameters:

• shape – Shape of the tensor (without the last dimension).

• dtype – Data type of the quaternion components.

Returns:

Quaternion with all components equal to 0.

classmethod ones(shape: tuple[int, ...] = (), dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Quaternion

Create quaternions with scalar component set to 1 and vector components set to 0.

Parameters:

• shape – Shape of the tensor (without the last dimension).

• dtype – Data type of the quaternion components.

Returns:

Quaternions with w=1 and x=y=z=0.

classmethod full(shape: tuple[int, ...], fill_value: float, dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Quaternion

Create quaternions with scalar component set to a value and vector components set to 0.

Parameters:

• shape – Shape of the tensor (without the last dimension).

• fill_value – Value to fill the scalar component with.

• dtype – Data type of the quaternion components.

Returns:

Quaternions with w=fill_value and x=y=z=0.

classmethod random(key: PRNGKey, shape: tuple[int, ...] = ()) → Quaternion

Generate normalized random quaternions.

Parameters:

• key – Key PRNG.

• shape – Shape of the tensor (without the last dimension).

Returns:

Normalized Quaternion.

property w: Array

property x: Array

property y: Array

property z: Array

property vector: Array

Vector part (…, 3)

norm() → Array

Quaternion norm.

normalize() → Quaternion

Normalize the quaternion.

Returns the normalized quaternion. If the quaternion has zero norm, returns the zero quaternion [0, 0, 0, 0].

to_components() → tuple[Array, Array, Array, Array]

to_rotation_matrix() → Array

Convert quaternion to rotation matrix.

Returns:

Array of shape (…, 3, 3)

rotate_vector(v: Array) → Array

Apply quaternion rotation to a vector.

Parameters:

v – Array of shape (…, 3) representing vectors

Returns:

Array of shape (…, 3) representing rotated vectors

__len__()

Length of the first axis.

__iter__()

Iterate over the first axis.

__neg__() → Quaternion

Quaternion negation.

__add__(other: Any) → Quaternion

Quaternion addition.

__radd__(other: Any) → Quaternion

Quaternion addition.

__sub__(other: Any) → Quaternion

Quaternion subtraction.

__rsub__(other: Any) → Quaternion

Quaternion subtraction.

__mul__(other: Any) → Quaternion

Quaternion multiplication.

__rmul__(other: Any) → Quaternion

Quaternion multiplication.

__truediv__(other: Any) → Quaternion

Quaternion division.

__rtruediv__(other: Any) → Quaternion

Quaternion division.

__pow__(exponent: float | int | Array) → Quaternion

Quaternion exponentiation q^n.

For integer exponents, uses optimized special cases. For non-integer exponents, uses the general formula: q^n = exp(n * log(q))

Parameters:

exponent – The exponent (scalar or array)

Returns:

The quaternion raised to the given power

log() → Quaternion

Compute quaternion logarithm.

For a quaternion q = ‖q‖ * (cos(θ) + sin(θ)v), the logarithm is: log(q) = log(‖q‖) + θ * v

Returns:

The logarithm of the quaternion

exp() → Quaternion

Compute quaternion exponential.

For a quaternion q = s + v, the exponential is: exp(q) = exp(s) * (cos(‖v‖) + sin(‖v‖) * v/‖v‖)

Returns:

The exponential of the quaternion

property nbytes: int

Number of bytes in the tensor.

property itemsize: int

Size of one quaternion element in bytes.

property shape: tuple[int, ...]

Shape of the tensor.

property ndim

Number of dimensions of the quaternion tensor (without the quaternion dimension).

property size

Total number of quaternions.

property dtype: dtype

Data type.

reshape(*shape) → Quaternion

Redimensionne le tableau de quaternions

flatten() → Quaternion

Aplatis le tableau de quaternions

ravel() → Quaternion

Aplatis le tableau de quaternions

squeeze(axis=None) → Quaternion

Supprime les dimensions de taille 1

conjugate() → Quaternion

Quaternion conjugate.

conj() → Quaternion

Quaternion conjugate.

block_until_ready() → None

Block until all pending computations are done.

property device: jax.Device[Any]

devices() → set[jax.Device[Any]]

slerp(other: Quaternion, t: float | Array) → Quaternion

Spherical linear interpolation between two quaternions.

Parameters:

• other – Target quaternion to interpolate towards

• t – Interpolation parameter in [0, 1]. t=0 returns self, t=1 returns other

Returns:

Interpolated quaternion

The Quaternion class provides a comprehensive interface for quaternion operations optimized for JAX. All methods are compatible with JAX transformations including JIT compilation, automatic differentiation, and vectorization.

Constructor Methods

Quaternion.__init__(w: float | Array, x: float | Array, y: float | Array, z: float | Array) → None

Initialize a tensor of quaternions.

Parameters:

• w – components of the quaternions.

• x – components of the quaternions.

• y – components of the quaternions.

• z – components of the quaternions.

classmethod Quaternion.from_array(array: Array) → Quaternion

Create a Quaternion array from a numeric array of shape (…, 4).

Parameters:

array – array of shape (…, 4) where the last dimension is [w, x, y, z]

classmethod Quaternion.from_scalar_vector(scalar: Array, vector: Array) → Quaternion

Create a quaternion from scalar and vector parts.

Parameters:

• scalar – Array of shape (…,) for the scalar part.

• vector – Array of shape (…, 3) for the vector part.

Returns:

Quaternion

classmethod Quaternion.from_rotation_matrix(rot: Array) → Quaternion

Create the quaternion associated to a rotation matrix.

Parameters:

rot – Array of shape (…, 3, 3) representing the rotation matrix

Returns:

The normalized Quaternion tensor representing the rotation matrix.

classmethod Quaternion.zeros(shape: tuple[int, ...] = (), dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Quaternion

Create quaternions with all components set to 0.

Parameters:

• shape – Shape of the tensor (without the last dimension).

• dtype – Data type of the quaternion components.

Returns:

Quaternion with all components equal to 0.

classmethod Quaternion.ones(shape: tuple[int, ...] = (), dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Quaternion

Create quaternions with scalar component set to 1 and vector components set to 0.

Parameters:

• shape – Shape of the tensor (without the last dimension).

• dtype – Data type of the quaternion components.

Returns:

Quaternions with w=1 and x=y=z=0.

classmethod Quaternion.full(shape: tuple[int, ...], fill_value: float, dtype: ~numpy.dtype = <class 'jax.numpy.float32'>) → Quaternion

Create quaternions with scalar component set to a value and vector components set to 0.

Parameters:

• shape – Shape of the tensor (without the last dimension).

• fill_value – Value to fill the scalar component with.

• dtype – Data type of the quaternion components.

Returns:

Quaternions with w=fill_value and x=y=z=0.

classmethod Quaternion.random(key: PRNGKey, shape: tuple[int, ...] = ()) → Quaternion

Generate normalized random quaternions.

Parameters:

• key – Key PRNG.

• shape – Shape of the tensor (without the last dimension).

Returns:

Normalized Quaternion.

Properties

Quaternion.w

Quaternion.x

Quaternion.y

Quaternion.z

Quaternion.vector

Vector part (…, 3)

Quaternion.shape

Shape of the tensor.

Quaternion.dtype

Data type.

Core Operations

Quaternion.norm() → Array

Quaternion norm.

Quaternion.normalize() → Quaternion

Normalize the quaternion.

Returns the normalized quaternion. If the quaternion has zero norm, returns the zero quaternion [0, 0, 0, 0].

Quaternion.conjugate() → Quaternion

Quaternion conjugate.

Quaternion.conj() → Quaternion

Quaternion conjugate.

Quaternion.to_components() → tuple[Array, Array, Array, Array]

Rotation Operations

Quaternion.to_rotation_matrix() → Array

Convert quaternion to rotation matrix.

Returns:

Array of shape (…, 3, 3)

Quaternion.rotate_vector(v: Array) → Array

Apply quaternion rotation to a vector.

Parameters:

v – Array of shape (…, 3) representing vectors

Returns:

Array of shape (…, 3) representing rotated vectors

Interpolation

Quaternion.slerp(other: Quaternion, t: float | Array) → Quaternion

Spherical linear interpolation between two quaternions.

Parameters:

• other – Target quaternion to interpolate towards

• t – Interpolation parameter in [0, 1]. t=0 returns self, t=1 returns other

Returns:

Interpolated quaternion

Advanced Operations

Quaternion.log() → Quaternion

Compute quaternion logarithm.

For a quaternion q = ‖q‖ * (cos(θ) + sin(θ)v), the logarithm is: log(q) = log(‖q‖) + θ * v

Returns:

The logarithm of the quaternion

Quaternion.exp() → Quaternion

Compute quaternion exponential.

For a quaternion q = s + v, the exponential is: exp(q) = exp(s) * (cos(‖v‖) + sin(‖v‖) * v/‖v‖)

Returns:

The exponential of the quaternion

Quaternion.__pow__(exponent: float | int | Array) → Quaternion

Quaternion exponentiation q^n.

For integer exponents, uses optimized special cases. For non-integer exponents, uses the general formula: q^n = exp(n * log(q))

Parameters:

exponent – The exponent (scalar or array)

Returns:

The quaternion raised to the given power

Array Operations

Quaternion.reshape(*shape) → Quaternion

Redimensionne le tableau de quaternions

Quaternion.flatten() → Quaternion

Aplatis le tableau de quaternions

Quaternion.ravel() → Quaternion

Aplatis le tableau de quaternions

Quaternion.squeeze(axis=None) → Quaternion

Supprime les dimensions de taille 1

Quaternion.block_until_ready() → None

Block until all pending computations are done.

Device and Memory

Quaternion.device

Quaternion.devices() → set[jax.Device[Any]]

Quaternion.nbytes

Number of bytes in the tensor.

Quaternion.itemsize

Size of one quaternion element in bytes.

Quaternion.size

Total number of quaternions.

Quaternion.ndim

Number of dimensions of the quaternion tensor (without the quaternion dimension).