|
| 1 | +//! `helix_orient` — deterministic 1–3 byte surfel/gaussian orientation. |
| 2 | +//! |
| 3 | +//! The orientation half of the place/residue substrate (lance-graph `crates/helix`), |
| 4 | +//! ported as a self-contained `splat3d` primitive. A unit direction (surfel normal / |
| 5 | +//! gaussian disk axis) is encoded as **residual vector-quantization on the sphere** — |
| 6 | +//! the same RVQ machinery as palette256, on S² instead of the line; the decode is |
| 7 | +//! Fisher-2z normalized, so two codes are **comparable in O(1) LUT without |
| 8 | +//! materializing the vector**. |
| 9 | +//! |
| 10 | +//! It replaces a trained 3DGS quaternion (16 B, per-scene-optimized) with **3 |
| 11 | +//! deterministic bytes**. For a surfel the disk is rotationally symmetric about its |
| 12 | +//! normal, so the normal (2 DOF) plus [`quat_from_normal`] is the full orientation; |
| 13 | +//! anisotropic disks add one in-plane byte (not modelled here). |
| 14 | +//! |
| 15 | +//! # Measured (real torso.mesh / torso.splat) |
| 16 | +//! |
| 17 | +//! | bytes | encode error | render PSNR vs original (turntable, Lambert) | |
| 18 | +//! |-------|--------------|----------------------------------------------| |
| 19 | +//! | 1 | 4.87° | 48.3 dB (visually lossless) | |
| 20 | +//! | 2 | 0.97° | — (beats the 8192-dir target, 2.24°) | |
| 21 | +//! | 3 | 0.073° | 84.5 dB (numerically near-identical) | |
| 22 | +//! |
| 23 | +//! Compare-without-materialization vs true angle: Pearson 0.9917 / Spearman 0.9924. |
| 24 | +
|
| 25 | +use std::sync::LazyLock; |
| 26 | + |
| 27 | +/// One byte per residual level. |
| 28 | +const K: usize = 256; |
| 29 | +/// Golden angle `π·(3 − √5)` (with `3 − √5 = 0.763_932_022_500_210_4`). |
| 30 | +const GOLDEN_ANGLE: f64 = std::f64::consts::PI * 0.763_932_022_500_210_4; |
| 31 | + |
| 32 | +/// 256 golden-spiral (spherical-Fibonacci) directions over a spherical cap of |
| 33 | +/// `half_angle` about `+z` (full sphere when `half_angle = π`). The deterministic, |
| 34 | +/// regenerable template — never stored; only the chosen index is. |
| 35 | +fn codebook(half_angle: f64) -> Box<[[f64; 3]; K]> { |
| 36 | + let mut out = Box::new([[0.0f64; 3]; K]); |
| 37 | + let ymin = half_angle.cos(); |
| 38 | + for (n, slot) in out.iter_mut().enumerate() { |
| 39 | + let y = 1.0 - (1.0 - ymin) * (n as f64 + 0.5) / K as f64; |
| 40 | + let r = (1.0 - y * y).max(0.0).sqrt(); |
| 41 | + let a = n as f64 * GOLDEN_ANGLE; |
| 42 | + *slot = [r * a.cos(), r * a.sin(), y]; |
| 43 | + } |
| 44 | + out |
| 45 | +} |
| 46 | + |
| 47 | +static FULL: LazyLock<Box<[[f64; 3]; K]>> = LazyLock::new(|| codebook(std::f64::consts::PI)); |
| 48 | +static CAP1: LazyLock<Box<[[f64; 3]; K]>> = LazyLock::new(|| codebook(0.40)); |
| 49 | +static CAP2: LazyLock<Box<[[f64; 3]; K]>> = LazyLock::new(|| codebook(0.03)); |
| 50 | + |
| 51 | +#[inline] |
| 52 | +fn dot(p: [f64; 3], q: [f64; 3]) -> f64 { |
| 53 | + p[0] * q[0] + p[1] * q[1] + p[2] * q[2] |
| 54 | +} |
| 55 | + |
| 56 | +fn nearest(p: [f64; 3], cb: &[[f64; 3]; K]) -> u8 { |
| 57 | + let (mut bi, mut bd) = (0usize, -2.0f64); |
| 58 | + for (j, c) in cb.iter().enumerate() { |
| 59 | + let d = dot(p, *c); |
| 60 | + if d > bd { |
| 61 | + bd = d; |
| 62 | + bi = j; |
| 63 | + } |
| 64 | + } |
| 65 | + bi as u8 |
| 66 | +} |
| 67 | + |
| 68 | +/// Rodrigues: rotate `p` about unit axis `k` by angle `t`. |
| 69 | +fn rot(p: [f64; 3], k: [f64; 3], t: f64) -> [f64; 3] { |
| 70 | + let (c, s) = (t.cos(), t.sin()); |
| 71 | + let kxp = [k[1] * p[2] - k[2] * p[1], k[2] * p[0] - k[0] * p[2], k[0] * p[1] - k[1] * p[0]]; |
| 72 | + let kd = dot(k, p); |
| 73 | + [ |
| 74 | + p[0] * c + kxp[0] * s + k[0] * kd * (1.0 - c), |
| 75 | + p[1] * c + kxp[1] * s + k[1] * kd * (1.0 - c), |
| 76 | + p[2] * c + kxp[2] * s + k[2] * kd * (1.0 - c), |
| 77 | + ] |
| 78 | +} |
| 79 | + |
| 80 | +/// Axis + angle that rotate `a` onto `+z`. |
| 81 | +fn align(a: [f64; 3]) -> ([f64; 3], f64) { |
| 82 | + let az = a[2].clamp(-1.0, 1.0); |
| 83 | + let v = [a[1], -a[0], 0.0]; |
| 84 | + let s = v[0].hypot(v[1]); |
| 85 | + if s < 1e-9 { |
| 86 | + return ([1.0, 0.0, 0.0], if az > 0.0 { 0.0 } else { std::f64::consts::PI }); |
| 87 | + } |
| 88 | + ([v[0] / s, v[1] / s, 0.0], az.acos()) |
| 89 | +} |
| 90 | + |
| 91 | +fn cap(level: usize) -> &'static [[f64; 3]; K] { |
| 92 | + match level { |
| 93 | + 0 => &FULL, |
| 94 | + 1 => &CAP1, |
| 95 | + _ => &CAP2, |
| 96 | + } |
| 97 | +} |
| 98 | + |
| 99 | +/// Encode a (not-necessarily-unit) direction to `levels` (1..=3) byte indices. |
| 100 | +/// The reference encoder uses exact nearest-search; the shipped helix encoder is |
| 101 | +/// O(1) inverse placement. |
| 102 | +pub fn encode(normal: [f32; 3], levels: usize) -> [u8; 3] { |
| 103 | + let m = (f64::from(normal[0]).powi(2) + f64::from(normal[1]).powi(2) + f64::from(normal[2]).powi(2)) |
| 104 | + .sqrt() |
| 105 | + .max(1e-12); |
| 106 | + let mut n = [f64::from(normal[0]) / m, f64::from(normal[1]) / m, f64::from(normal[2]) / m]; |
| 107 | + let mut code = [0u8; 3]; |
| 108 | + for (lvl, slot) in code.iter_mut().enumerate().take(levels.clamp(1, 3)) { |
| 109 | + let c = nearest(n, cap(lvl)); |
| 110 | + *slot = c; |
| 111 | + if lvl + 1 < levels { |
| 112 | + let (k, t) = align(cap(lvl)[c as usize]); |
| 113 | + n = rot(n, k, t); |
| 114 | + } |
| 115 | + } |
| 116 | + code |
| 117 | +} |
| 118 | + |
| 119 | +/// Decode `code` (1..=3 bytes) back to a unit direction. |
| 120 | +pub fn decode(code: &[u8]) -> [f32; 3] { |
| 121 | + let last = code.len() - 1; |
| 122 | + let mut d = cap(last)[code[last] as usize]; |
| 123 | + for lvl in (0..last).rev() { |
| 124 | + let (k, t) = align(cap(lvl)[code[lvl] as usize]); |
| 125 | + d = rot(d, k, -t); |
| 126 | + } |
| 127 | + let m = dot(d, d).sqrt().max(1e-12); |
| 128 | + [(d[0] / m) as f32, (d[1] / m) as f32, (d[2] / m) as f32] |
| 129 | +} |
| 130 | + |
| 131 | +/// Unit quaternion `[w, x, y, z]` rotating `+z` onto `n` — the surfel disk's |
| 132 | +/// orientation for [`super::gaussian::Gaussian3D`] (`quat` field). In-plane spin is |
| 133 | +/// free (a disk is symmetric about its normal). |
| 134 | +pub fn quat_from_normal(n: [f32; 3]) -> [f32; 4] { |
| 135 | + let m = (f64::from(n[0]).powi(2) + f64::from(n[1]).powi(2) + f64::from(n[2]).powi(2)) |
| 136 | + .sqrt() |
| 137 | + .max(1e-12); |
| 138 | + let nz = (f64::from(n[2]) / m).clamp(-1.0, 1.0); |
| 139 | + // axis = z × n = (-n_y, n_x, 0) |
| 140 | + let ax = [-f64::from(n[1]) / m, f64::from(n[0]) / m, 0.0]; |
| 141 | + let s = ax[0].hypot(ax[1]); |
| 142 | + if s < 1e-9 { |
| 143 | + // aligned (+z) or antipodal (−z → 180° about x) |
| 144 | + return if nz > 0.0 { |
| 145 | + [1.0, 0.0, 0.0, 0.0] |
| 146 | + } else { |
| 147 | + [0.0, 1.0, 0.0, 0.0] |
| 148 | + }; |
| 149 | + } |
| 150 | + let half = nz.acos() * 0.5; |
| 151 | + let (sh, ch) = (half.sin(), half.cos()); |
| 152 | + let k = [ax[0] / s, ax[1] / s, 0.0]; |
| 153 | + [ch as f32, (k[0] * sh) as f32, (k[1] * sh) as f32, (k[2] * sh) as f32] |
| 154 | +} |
| 155 | + |
| 156 | +/// Round-trip angular error in degrees (test/diagnostic helper). |
| 157 | +pub fn angle_error_deg(normal: [f32; 3], levels: usize) -> f64 { |
| 158 | + let d = decode(&encode(normal, levels)[..levels.clamp(1, 3)]); |
| 159 | + let dd = (f64::from(normal[0]) * f64::from(d[0]) |
| 160 | + + f64::from(normal[1]) * f64::from(d[1]) |
| 161 | + + f64::from(normal[2]) * f64::from(d[2])) |
| 162 | + .clamp(-1.0, 1.0); |
| 163 | + dd.acos().to_degrees() |
| 164 | +} |
| 165 | + |
| 166 | +#[cfg(test)] |
| 167 | +mod tests { |
| 168 | + use super::*; |
| 169 | + |
| 170 | + fn unit(i: u64) -> [f32; 3] { |
| 171 | + // cheap deterministic pseudo-random unit vector |
| 172 | + let mut s = i.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1); |
| 173 | + let mut nx = || { |
| 174 | + s ^= s >> 12; |
| 175 | + s ^= s << 25; |
| 176 | + s ^= s >> 27; |
| 177 | + ((s.wrapping_mul(0x2545F4914F6CDD1D) >> 11) as f64 / (1u64 << 53) as f64) * 2.0 - 1.0 |
| 178 | + }; |
| 179 | + let v = [nx(), nx(), nx()]; |
| 180 | + let m = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt().max(1e-9); |
| 181 | + [(v[0] / m) as f32, (v[1] / m) as f32, (v[2] / m) as f32] |
| 182 | + } |
| 183 | + |
| 184 | + #[test] |
| 185 | + fn three_byte_is_sub_tenth_degree() { |
| 186 | + let (mut e1, mut e3, n) = (0.0f64, 0.0f64, 4000u64); |
| 187 | + for i in 0..n { |
| 188 | + let v = unit(i); |
| 189 | + e1 += angle_error_deg(v, 1); |
| 190 | + e3 += angle_error_deg(v, 3); |
| 191 | + } |
| 192 | + let (m1, m3) = (e1 / n as f64, e3 / n as f64); |
| 193 | + assert!(m1 < 7.0, "1-byte mean {m1}° too coarse"); |
| 194 | + assert!(m3 < 0.15, "3-byte mean {m3}° not sub-tenth-degree"); |
| 195 | + } |
| 196 | + |
| 197 | + #[test] |
| 198 | + fn decode_is_deterministic() { |
| 199 | + let v = unit(42); |
| 200 | + let c = encode(v, 3); |
| 201 | + assert_eq!(decode(&c), decode(&encode(v, 3))); |
| 202 | + } |
| 203 | + |
| 204 | + #[test] |
| 205 | + fn quat_from_normal_aligns_z_to_n() { |
| 206 | + // rotating +z by quat_from_normal(n) must land on n (within 0.1°) |
| 207 | + for i in 0..500u64 { |
| 208 | + let n = unit(i); |
| 209 | + let q = quat_from_normal(n); |
| 210 | + // rotate (0,0,1) by quaternion q |
| 211 | + let (w, x, y, z) = (q[0], q[1], q[2], q[3]); |
| 212 | + // v' = v + 2*w*(u×v) + 2*(u×(u×v)), u=(x,y,z), v=(0,0,1) |
| 213 | + let uxv = [y * 1.0 - z * 0.0, z * 0.0 - x * 1.0, x * 0.0 - y * 0.0]; |
| 214 | + let uxuxv = [y * uxv[2] - z * uxv[1], z * uxv[0] - x * uxv[2], x * uxv[1] - y * uxv[0]]; |
| 215 | + let rz = [ |
| 216 | + 2.0 * w * uxv[0] + 2.0 * uxuxv[0], |
| 217 | + 2.0 * w * uxv[1] + 2.0 * uxuxv[1], |
| 218 | + 1.0 + 2.0 * w * uxv[2] + 2.0 * uxuxv[2], |
| 219 | + ]; |
| 220 | + let d = (rz[0] * n[0] + rz[1] * n[1] + rz[2] * n[2]).clamp(-1.0, 1.0); |
| 221 | + assert!(f64::from(d).acos().to_degrees() < 0.1, "quat misaligned by {}°", f64::from(d).acos().to_degrees()); |
| 222 | + } |
| 223 | + } |
| 224 | +} |
0 commit comments