use crate::jit::eval_interp_checked;
use crate::kernel::{ExprId, ExprPool};
use crate::matrix::zero_test::{zero_status, ZeroStatus};
use crate::matrix::{Matrix, MatrixError};
use crate::simplify::engine::simplify;
use std::fmt;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum QuaternionError {
ZeroNorm {
norm_squared: String,
proven_zero: bool,
},
UndefinedAxis {
quaternion: String,
},
NotARotation {
reason: String,
},
}
impl fmt::Display for QuaternionError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
QuaternionError::ZeroNorm {
norm_squared,
proven_zero,
} => {
if *proven_zero {
write!(
f,
"quaternion has zero norm (|q|² = `{norm_squared}`); it has no inverse"
)
} else {
write!(
f,
"|q|² = `{norm_squared}` could not be shown to be non-zero, and this operation divides by it"
)
}
}
QuaternionError::UndefinedAxis { quaternion } => write!(
f,
"`{quaternion}` has zero vector part: it is the identity rotation, and every unit vector is an axis for it"
),
QuaternionError::NotARotation { reason } => {
write!(f, "not a proper rotation matrix: {reason}")
}
}
}
}
impl std::error::Error for QuaternionError {}
impl crate::errors::AlkahestError for QuaternionError {
fn code(&self) -> &'static str {
match self {
QuaternionError::ZeroNorm { .. } => "E-QUAT-001",
QuaternionError::UndefinedAxis { .. } => "E-QUAT-002",
QuaternionError::NotARotation { .. } => "E-QUAT-003",
}
}
fn remediation(&self) -> Option<&'static str> {
match self {
QuaternionError::ZeroNorm { .. } => Some(
"use a non-zero quaternion; if the components are symbolic, substitute concrete values or declare the parameters so the norm can be decided",
),
QuaternionError::UndefinedAxis { .. } => Some(
"the rotation is the identity — report the angle as 0 and pick whatever axis your convention prefers, explicitly, at the call site",
),
QuaternionError::NotARotation { .. } => Some(
"pass a 3×3 matrix of numbers with RᵀR = I and det R = +1; a reflection (det = −1) is not a rotation and has no quaternion",
),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Quaternion {
w: ExprId,
x: ExprId,
y: ExprId,
z: ExprId,
}
const ROTATION_TOL: f64 = 1e-9;
impl Quaternion {
pub fn new(w: ExprId, x: ExprId, y: ExprId, z: ExprId) -> Self {
Quaternion { w, x, y, z }
}
pub fn scalar(w: ExprId, pool: &ExprPool) -> Self {
let zero = pool.integer(0_i32);
Quaternion {
w,
x: zero,
y: zero,
z: zero,
}
}
pub fn pure(v: &[ExprId; 3], pool: &ExprPool) -> Self {
Quaternion {
w: pool.integer(0_i32),
x: v[0],
y: v[1],
z: v[2],
}
}
pub fn identity(pool: &ExprPool) -> Self {
Quaternion::scalar(pool.integer(1_i32), pool)
}
pub fn w(&self) -> ExprId {
self.w
}
pub fn x(&self) -> ExprId {
self.x
}
pub fn y(&self) -> ExprId {
self.y
}
pub fn z(&self) -> ExprId {
self.z
}
pub fn components(&self) -> [ExprId; 4] {
[self.w, self.x, self.y, self.z]
}
pub fn vector_part(&self) -> [ExprId; 3] {
[self.x, self.y, self.z]
}
pub fn simplified(&self, pool: &ExprPool) -> Self {
Quaternion {
w: simplify(self.w, pool).value,
x: simplify(self.x, pool).value,
y: simplify(self.y, pool).value,
z: simplify(self.z, pool).value,
}
}
pub fn add(&self, other: &Quaternion, pool: &ExprPool) -> Self {
Quaternion {
w: simplify(pool.add(vec![self.w, other.w]), pool).value,
x: simplify(pool.add(vec![self.x, other.x]), pool).value,
y: simplify(pool.add(vec![self.y, other.y]), pool).value,
z: simplify(pool.add(vec![self.z, other.z]), pool).value,
}
}
pub fn sub(&self, other: &Quaternion, pool: &ExprPool) -> Self {
let neg = pool.integer(-1_i32);
let negated = other.scale(neg, pool);
self.add(&negated, pool)
}
pub fn scale(&self, s: ExprId, pool: &ExprPool) -> Self {
Quaternion {
w: simplify(pool.mul(vec![s, self.w]), pool).value,
x: simplify(pool.mul(vec![s, self.x]), pool).value,
y: simplify(pool.mul(vec![s, self.y]), pool).value,
z: simplify(pool.mul(vec![s, self.z]), pool).value,
}
}
pub fn mul(&self, other: &Quaternion, pool: &ExprPool) -> Self {
let (a, b) = (self, other);
let neg = pool.integer(-1_i32);
let m = |p: ExprId, q: ExprId| pool.mul(vec![p, q]);
let nm = |p: ExprId, q: ExprId| pool.mul(vec![neg, p, q]);
let w = pool.add(vec![m(a.w, b.w), nm(a.x, b.x), nm(a.y, b.y), nm(a.z, b.z)]);
let x = pool.add(vec![m(a.w, b.x), m(a.x, b.w), m(a.y, b.z), nm(a.z, b.y)]);
let y = pool.add(vec![m(a.w, b.y), nm(a.x, b.z), m(a.y, b.w), m(a.z, b.x)]);
let z = pool.add(vec![m(a.w, b.z), m(a.x, b.y), nm(a.y, b.x), m(a.z, b.w)]);
Quaternion {
w: simplify(w, pool).value,
x: simplify(x, pool).value,
y: simplify(y, pool).value,
z: simplify(z, pool).value,
}
}
pub fn conjugate(&self, pool: &ExprPool) -> Self {
let neg = pool.integer(-1_i32);
Quaternion {
w: self.w,
x: simplify(pool.mul(vec![neg, self.x]), pool).value,
y: simplify(pool.mul(vec![neg, self.y]), pool).value,
z: simplify(pool.mul(vec![neg, self.z]), pool).value,
}
}
pub fn norm_squared(&self, pool: &ExprPool) -> ExprId {
let two = pool.integer(2_i32);
let terms: Vec<ExprId> = self
.components()
.iter()
.map(|&c| pool.pow(c, two))
.collect();
simplify(pool.add(terms), pool).value
}
pub fn norm(&self, pool: &ExprPool) -> ExprId {
let n2 = self.norm_squared(pool);
simplify(pool.func("sqrt", vec![n2]), pool).value
}
fn nonzero_norm_squared(&self, pool: &ExprPool) -> Result<ExprId, QuaternionError> {
let n2 = self.norm_squared(pool);
match zero_status(pool, n2) {
ZeroStatus::NonZero => Ok(n2),
status => Err(QuaternionError::ZeroNorm {
norm_squared: pool.display(n2).to_string(),
proven_zero: status == ZeroStatus::Zero,
}),
}
}
pub fn inverse(&self, pool: &ExprPool) -> Result<Self, QuaternionError> {
let n2 = self.nonzero_norm_squared(pool)?;
let inv = pool.pow(n2, pool.integer(-1_i32));
Ok(self.conjugate(pool).scale(inv, pool))
}
pub fn normalize(&self, pool: &ExprPool) -> Result<Self, QuaternionError> {
let n2 = self.nonzero_norm_squared(pool)?;
let n = simplify(pool.func("sqrt", vec![n2]), pool).value;
let inv = pool.pow(n, pool.integer(-1_i32));
Ok(self.scale(inv, pool))
}
pub fn rotate(&self, v: &[ExprId; 3], pool: &ExprPool) -> Result<[ExprId; 3], QuaternionError> {
let inv = self.inverse(pool)?;
let p = Quaternion::pure(v, pool);
let out = self.mul(&p, pool).mul(&inv, pool);
Ok(out.vector_part())
}
pub fn to_rotation_matrix(&self, pool: &ExprPool) -> Result<Matrix, QuaternionError> {
let n2 = self.nonzero_norm_squared(pool)?;
let inv = pool.pow(n2, pool.integer(-1_i32));
let two = pool.integer(2_i32);
let neg = pool.integer(-1_i32);
let sq = |c: ExprId| pool.pow(c, two);
let nsq = |c: ExprId| pool.mul(vec![neg, pool.pow(c, two)]);
let p2 = |a: ExprId, b: ExprId| pool.mul(vec![two, a, b]);
let m2 = |a: ExprId, b: ExprId| pool.mul(vec![neg, two, a, b]);
let (w, x, y, z) = (self.w, self.x, self.y, self.z);
let raw = [
[
pool.add(vec![sq(w), sq(x), nsq(y), nsq(z)]),
pool.add(vec![p2(x, y), m2(w, z)]),
pool.add(vec![p2(x, z), p2(w, y)]),
],
[
pool.add(vec![p2(x, y), p2(w, z)]),
pool.add(vec![sq(w), nsq(x), sq(y), nsq(z)]),
pool.add(vec![p2(y, z), m2(w, x)]),
],
[
pool.add(vec![p2(x, z), m2(w, y)]),
pool.add(vec![p2(y, z), p2(w, x)]),
pool.add(vec![sq(w), nsq(x), nsq(y), sq(z)]),
],
];
let rows: Vec<Vec<ExprId>> = raw
.iter()
.map(|row| {
row.iter()
.map(|&e| simplify(pool.mul(vec![inv, e]), pool).value)
.collect()
})
.collect();
Matrix::new(rows).map_err(|e: MatrixError| QuaternionError::NotARotation {
reason: e.to_string(),
})
}
pub fn from_axis_angle(
axis: &[ExprId; 3],
angle: ExprId,
pool: &ExprPool,
) -> Result<Self, QuaternionError> {
let n2 = crate::vector::norm_squared(axis, pool);
let unit = match zero_status(pool, n2) {
ZeroStatus::NonZero => {
let one = pool.integer(1_i32);
if n2 == one {
*axis
} else {
let n = simplify(pool.func("sqrt", vec![n2]), pool).value;
let inv = pool.pow(n, pool.integer(-1_i32));
crate::vector::scale(inv, axis, pool)
}
}
status => {
return Err(QuaternionError::ZeroNorm {
norm_squared: pool.display(n2).to_string(),
proven_zero: status == ZeroStatus::Zero,
})
}
};
let half = pool.mul(vec![pool.rational(1, 2), angle]);
let c = simplify(pool.func("cos", vec![half]), pool).value;
let s = simplify(pool.func("sin", vec![half]), pool).value;
Ok(Quaternion {
w: c,
x: simplify(pool.mul(vec![s, unit[0]]), pool).value,
y: simplify(pool.mul(vec![s, unit[1]]), pool).value,
z: simplify(pool.mul(vec![s, unit[2]]), pool).value,
})
}
pub fn to_axis_angle(&self, pool: &ExprPool) -> Result<([ExprId; 3], ExprId), QuaternionError> {
let v = self.vector_part();
let n2 = crate::vector::norm_squared(&v, pool);
if zero_status(pool, n2) != ZeroStatus::NonZero {
return Err(QuaternionError::UndefinedAxis {
quaternion: format!(
"{} + {}i + {}j + {}k",
pool.display(self.w),
pool.display(self.x),
pool.display(self.y),
pool.display(self.z)
),
});
}
let n = simplify(pool.func("sqrt", vec![n2]), pool).value;
let inv = pool.pow(n, pool.integer(-1_i32));
let axis = crate::vector::scale(inv, &v, pool);
let angle = simplify(
pool.mul(vec![
pool.integer(2_i32),
pool.func("atan2", vec![n, self.w]),
]),
pool,
)
.value;
Ok((axis, angle))
}
pub fn from_rotation_matrix(m: &Matrix, pool: &ExprPool) -> Result<Self, QuaternionError> {
if m.rows != 3 || m.cols != 3 {
return Err(QuaternionError::NotARotation {
reason: format!("expected a 3×3 matrix, got {}×{}", m.rows, m.cols),
});
}
let mut r = [[0.0_f64; 3]; 3];
for (i, row) in r.iter_mut().enumerate() {
for (j, cell) in row.iter_mut().enumerate() {
let entry = m.get(i, j);
match eval_interp_checked(entry, &std::collections::HashMap::new(), pool) {
Ok(v) if v.is_finite() => *cell = v,
_ => {
return Err(QuaternionError::NotARotation {
reason: format!(
"entry ({i},{j}) = `{}` is not a finite number; the branch \
selection is a comparison between entries and there is none \
to make on a symbol",
pool.display(entry)
),
})
}
}
}
}
check_proper_rotation(&r)?;
let q = shepperd(&r);
let rebuilt = rotation_matrix_f64(&q);
let worst = (0..3)
.flat_map(|i| (0..3).map(move |j| (i, j)))
.map(|(i, j)| (rebuilt[i][j] - r[i][j]).abs())
.fold(0.0_f64, f64::max);
if worst > ROTATION_TOL {
return Err(QuaternionError::NotARotation {
reason: format!(
"the recovered quaternion does not reproduce the matrix \
(worst entry discrepancy {worst:.3e} > {ROTATION_TOL:.0e})"
),
});
}
Ok(Quaternion {
w: pool.float(q[0], 53),
x: pool.float(q[1], 53),
y: pool.float(q[2], 53),
z: pool.float(q[3], 53),
})
}
}
fn check_proper_rotation(r: &[[f64; 3]; 3]) -> Result<(), QuaternionError> {
let mut worst = 0.0_f64;
for i in 0..3 {
for j in 0..3 {
let dot: f64 = (0..3).map(|k| r[k][i] * r[k][j]).sum();
let target = if i == j { 1.0 } else { 0.0 };
worst = worst.max((dot - target).abs());
}
}
if worst > ROTATION_TOL {
return Err(QuaternionError::NotARotation {
reason: format!(
"columns are not orthonormal: worst |RᵀR − I| entry is {worst:.3e} \
(tolerance {ROTATION_TOL:.0e})"
),
});
}
let det = r[0][0] * (r[1][1] * r[2][2] - r[1][2] * r[2][1])
- r[0][1] * (r[1][0] * r[2][2] - r[1][2] * r[2][0])
+ r[0][2] * (r[1][0] * r[2][1] - r[1][1] * r[2][0]);
if (det - 1.0).abs() > ROTATION_TOL {
return Err(QuaternionError::NotARotation {
reason: format!(
"det R = {det:.6}, not +1; an orthogonal matrix with det = −1 is a \
reflection and is not in the image of the quaternion rotation map"
),
});
}
Ok(())
}
fn shepperd(r: &[[f64; 3]; 3]) -> [f64; 4] {
let trace = r[0][0] + r[1][1] + r[2][2];
if trace > 0.0 {
let s = (trace + 1.0).sqrt() * 2.0;
[
0.25 * s,
(r[2][1] - r[1][2]) / s,
(r[0][2] - r[2][0]) / s,
(r[1][0] - r[0][1]) / s,
]
} else if r[0][0] > r[1][1] && r[0][0] > r[2][2] {
let s = (1.0 + r[0][0] - r[1][1] - r[2][2]).sqrt() * 2.0;
[
(r[2][1] - r[1][2]) / s,
0.25 * s,
(r[0][1] + r[1][0]) / s,
(r[0][2] + r[2][0]) / s,
]
} else if r[1][1] > r[2][2] {
let s = (1.0 + r[1][1] - r[0][0] - r[2][2]).sqrt() * 2.0;
[
(r[0][2] - r[2][0]) / s,
(r[0][1] + r[1][0]) / s,
0.25 * s,
(r[1][2] + r[2][1]) / s,
]
} else {
let s = (1.0 + r[2][2] - r[0][0] - r[1][1]).sqrt() * 2.0;
[
(r[1][0] - r[0][1]) / s,
(r[0][2] + r[2][0]) / s,
(r[1][2] + r[2][1]) / s,
0.25 * s,
]
}
}
fn rotation_matrix_f64(q: &[f64; 4]) -> [[f64; 3]; 3] {
let [w, x, y, z] = *q;
let n2 = w * w + x * x + y * y + z * z;
let inv = 1.0 / n2;
[
[
(w * w + x * x - y * y - z * z) * inv,
2.0 * (x * y - w * z) * inv,
2.0 * (x * z + w * y) * inv,
],
[
2.0 * (x * y + w * z) * inv,
(w * w - x * x + y * y - z * z) * inv,
2.0 * (y * z - w * x) * inv,
],
[
2.0 * (x * z - w * y) * inv,
2.0 * (y * z + w * x) * inv,
(w * w - x * x - y * y + z * z) * inv,
],
]
}