File size: 2,331 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
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
|