FastQuat - High-Performance Quaternions with JAX

FastQuat provides optimized quaternion operations with full JAX compatibility, featuring:

• 🚀 Hardware-accelerated computations (CPU/GPU/TPU)
• 🔄 Automatic differentiation support
• 🧩 Seamless integration with JAX transformations (jit, grad, vmap)
• 📦 Efficient storage using interleaved memory layout

Installation

pip install fastquat

This will install FastQuat with CPU support. For GPU support, you may need to install JAX with CUDA support:

pip install "jax[cuda12]" fastquat

Quick Start

import jax.numpy as jnp
from fastquat import Quaternion

# Create quaternions
q1 = Quaternion(1)  # Identity quaternion
q2 = Quaternion(0.7071, 0.7071, 0.0, 0.0)  # 90° rotation around x-axis
q = Quaternion(1.0, 0.1, 0.2, 0.3)
p = 2

# Arithmetic
q_sum = q1 + q2
q_diff = q1 - q2
q_product = q1 * q2
q_power = q**p

# Comparison (element-wise)
same = q1 == q2  # Boolean array, True where all components are equal

# Normalization
norm = abs(q)  # Quaternion norm
q_unit = q.normalize()  # Unit quaternion

# Conjugation and inverse
q_conj = q.conj()  # Conjugate
q_inv = 1 / q  # Inverse, or q ** -1

# Other operations
q_log = q.log()
q_exp = q.exp()

# Rotate vectors
vector = jnp.array([1.0, 0.0, 0.0])
rotated = q2.rotate_vector(vector)

# Spherical interpolation (SLERP)
interpolated = q1.slerp(q2, t=0.5)  # Halfway between q1 and q2

Features

Core Operations

• Quaternion arithmetic: Addition, multiplication, conjugation, inverse, power, exponentiation, logarithm
• Normalization: Efficient unit quaternion computation
• Comparison: Element-wise == and !=, like NumPy arrays
• Conversion: To/from rotation matrices, axis-angle, rotation vectors, Euler angles
• Vector rotation: Direct vector transformation

Advanced Interpolation

• SLERP (Spherical Linear Interpolation): Smooth rotation interpolation
  • Automatically handles shortest path selection
  • Numerically stable for close quaternions
  • Supports batched operations and array-valued parameters

JAX Integration

• JIT compilation: Compile quaternion operations for maximum performance
• Automatic differentiation: Compute gradients through quaternion operations
• Vectorization: Process batches of quaternions efficiently
• Device support: Run on CPU, GPU, or TPU

Performance

FastQuat is optimized for high-performance computing:

• Memory-efficient interleaved storage
• SIMD-optimized operations on supported hardware
• Zero-copy integration with JAX arrays
• Minimal Python overhead through JIT compilation

Examples

Basic Usage

import jax
import jax.numpy as jnp
from fastquat import Quaternion

# Create random quaternions
key = jax.random.PRNGKey(42)
q_batch = Quaternion.random(key, shape=(1000,))


# JIT-compiled batch operations
@jax.jit
def batch_rotate(quaternions, vectors):
    return quaternions.rotate_vector(vectors)


vectors = jax.random.normal(key, (1000, 3))
rotated_batch = batch_rotate(q_batch, vectors)

SLERP

# Smooth rotation interpolation
q_start = Quaternion(1.0)
q_end = Quaternion.from_rotation_matrix(rotation_matrix)

# Generate smooth interpolation
t_values = jnp.linspace(0, 1, 100)
interpolated_rotations = q_start.slerp(q_end, t_values)

# Apply to object vertices for smooth animation
animated_vertices = interpolated_rotations.rotate_vector(object_vertices)

Documentation

Full documentation is available at fastquat.readthedocs.io

Contributing

Contributions are welcome! Please see our development guide for details.

License

MIT License - see LICENSE file for details.