| |
| |
|
|
| from __future__ import annotations |
|
|
| import jax.numpy as jnp |
| from jaxtyping import Array, Float |
|
|
|
|
| def quaternion_multiply( |
| a: Float[Array, "... 4"], |
| b: Float[Array, "... 4"], |
| ) -> Float[Array, "... 4"]: |
| w = ( |
| a[..., 0] * b[..., 0] |
| - a[..., 1] * b[..., 1] |
| - a[..., 2] * b[..., 2] |
| - a[..., 3] * b[..., 3] |
| ) |
| x = ( |
| a[..., 0] * b[..., 1] |
| + a[..., 1] * b[..., 0] |
| + a[..., 2] * b[..., 3] |
| - a[..., 3] * b[..., 2] |
| ) |
| y = ( |
| a[..., 0] * b[..., 2] |
| - a[..., 1] * b[..., 3] |
| + a[..., 2] * b[..., 0] |
| + a[..., 3] * b[..., 1] |
| ) |
| z = ( |
| a[..., 0] * b[..., 3] |
| + a[..., 1] * b[..., 2] |
| - a[..., 2] * b[..., 1] |
| + a[..., 3] * b[..., 0] |
| ) |
| return jnp.stack([w, x, y, z], axis=-1) |
|
|
|
|
| def quaternion_exp_tangent( |
| v: Float[Array, "... 3"], |
| ) -> Float[Array, "... 4"]: |
| theta_sq = jnp.sum(v * v, axis=-1, keepdims=True) |
| theta = jnp.sqrt(theta_sq) |
| small = theta_sq < 1e-12 |
| sin_over_theta = jnp.where( |
| small, |
| 1.0 - theta_sq / 6.0 + theta_sq * theta_sq / 120.0, |
| jnp.sin(theta) / jnp.where(small, 1.0, theta), |
| ) |
| real = jnp.where( |
| small, |
| 1.0 - theta_sq / 2.0 + theta_sq * theta_sq / 24.0, |
| jnp.cos(theta), |
| ) |
| return jnp.concatenate([real, sin_over_theta * v], axis=-1) |
|
|
|
|
| def normalize_quaternion( |
| q: Float[Array, "... 4"], |
| eps: float = 1e-12, |
| ) -> Float[Array, "... 4"]: |
| return q / jnp.sqrt(jnp.sum(q * q, axis=-1, keepdims=True) + eps) |
|
|
|
|
| def quaternion_conjugate( |
| q: Float[Array, "... 4"], |
| ) -> Float[Array, "... 4"]: |
| return jnp.concatenate([q[..., 0:1], -q[..., 1:]], axis=-1) |
|
|
|
|
| def quaternion_log( |
| q: Float[Array, "... 4"], |
| eps: float = 1e-8, |
| ) -> Float[Array, "... 3"]: |
| w = q[..., 0:1] |
| v = q[..., 1:] |
| sin_norm_sq = jnp.sum(v * v, axis=-1, keepdims=True) |
| sin_norm = jnp.sqrt(sin_norm_sq) |
| theta = jnp.arctan2(sin_norm, w) |
| small = sin_norm < eps |
| factor_small = 1.0 + sin_norm_sq / 6.0 |
| factor_normal = theta / jnp.where( |
| small, |
| jnp.ones_like(sin_norm), |
| sin_norm, |
| ) |
| factor = jnp.where(small, factor_small, factor_normal) |
| return v * factor |
|
|
