loam_math/
spherical.rs

1//! Upper-hemisphere chart `(p, sqrt(1 - |p|²))` with Vec3 points.
2//! Crossing the equator loses the hemisphere sign; use SphericalS3Embedded for full coverage.
3
4use std::borrow::Cow;
5
6use glam::{Mat3, Mat4, Quat, Vec3, Vec4};
7use serde::{Deserialize, Serialize};
8
9use crate::space::{IsometryGroup, Space, WgslAccuracy, WgslSpace};
10
11const SPHERE_R2_MAX: f32 = 1.0 - 1e-6;
12
13const EXP_TANGENT_MIN_SQ: f32 = 1e-14;
14
15const LOG_PERP_MIN: f32 = 1e-7;
16
17const ISO_TRANSLATION_MIN_ARC: f32 = 1e-7;
18
19fn clamp_to_hemisphere(p: Vec3) -> Vec3 {
20    let r2 = p.length_squared();
21    if r2 <= SPHERE_R2_MAX {
22        p
23    } else {
24        #[cfg(debug_assertions)]
25        tracing::warn!("SphericalS3: point outside upper hemisphere clamped (|p|²={r2:.4})");
26        p * (SPHERE_R2_MAX.sqrt() / r2.sqrt())
27    }
28}
29
30fn to_sphere(p: Vec3) -> Vec4 {
31    let r2 = p.length_squared().min(SPHERE_R2_MAX);
32    Vec4::new(p.x, p.y, p.z, (1.0 - r2).sqrt())
33}
34
35fn from_sphere(q: Vec4) -> Vec3 {
36    #[cfg(debug_assertions)]
37    if q.w < 0.0 {
38        tracing::warn!(
39            "SphericalS3: iso_apply moved point to lower hemisphere (w={:.4}); \
40             result will be out-of-domain",
41            q.w
42        );
43    }
44    q.truncate()
45}
46
47/// SO(4) acting on the ambient embedding of S³.
48#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
49pub struct Iso4 {
50    pub matrix: Mat4,
51}
52
53impl Iso4 {
54    pub const IDENTITY: Self = Self {
55        matrix: Mat4::IDENTITY,
56    };
57
58    /// Fixes the north pole.
59    pub fn from_rotation(rotation: Quat) -> Self {
60        let r = Mat3::from_quat(rotation);
61        Self {
62            matrix: Mat4::from_cols(
63                r.col(0).extend(0.0),
64                r.col(1).extend(0.0),
65                r.col(2).extend(0.0),
66                Vec4::W,
67            ),
68        }
69    }
70
71    /// Maps the north pole to `target`, clamped to the upper-hemisphere chart.
72    pub fn from_translation(target: Vec3) -> Self {
73        let qt = to_sphere(clamp_to_hemisphere(target));
74        let c = qt.w;
75        let s = qt.truncate().length();
76        if s < ISO_TRANSLATION_MIN_ARC {
77            return Self::IDENTITY;
78        }
79        let n = qt.truncate() / s;
80        let k = c - 1.0;
81
82        Self {
83            matrix: Mat4::from_cols(
84                Vec4::new(1.0 + k * n.x * n.x, k * n.x * n.y, k * n.x * n.z, -s * n.x),
85                Vec4::new(k * n.y * n.x, 1.0 + k * n.y * n.y, k * n.y * n.z, -s * n.y),
86                Vec4::new(k * n.z * n.x, k * n.z * n.y, 1.0 + k * n.z * n.z, -s * n.z),
87                Vec4::new(s * n.x, s * n.y, s * n.z, c),
88            ),
89        }
90    }
91
92    pub fn left_translation(point: Vec4) -> Self {
93        let (x, y, z, w) = (point.x, point.y, point.z, point.w);
94        Self {
95            matrix: Mat4::from_cols(
96                Vec4::new(w, z, -y, -x),
97                Vec4::new(-z, w, x, -y),
98                Vec4::new(y, -x, w, -z),
99                point,
100            ),
101        }
102    }
103}
104
105#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
106pub struct SphericalS3;
107
108impl Space for SphericalS3 {
109    type Point = Vec3;
110    type Vector = Vec3;
111    type Frame = Iso4;
112
113    fn frame_at(&self, at: Vec3) -> Iso4 {
114        Iso4::from_translation(at)
115    }
116
117    fn distance(&self, a: Vec3, b: Vec3) -> f32 {
118        let qa = to_sphere(clamp_to_hemisphere(a));
119        let qb = to_sphere(clamp_to_hemisphere(b));
120
121        let half_chord = (qa - qb).length() * 0.5;
122        2.0 * half_chord.clamp(0.0, 1.0).asin()
123    }
124
125    fn exp(&self, at: Vec3, v: Vec3) -> Vec3 {
126        let at = clamp_to_hemisphere(at);
127        if v.length_squared() < EXP_TANGENT_MIN_SQ {
128            return at;
129        }
130        let q = to_sphere(at);
131
132        let vw = -v.dot(at) / q.w;
133        let v4 = Vec4::new(v.x, v.y, v.z, vw);
134        let mag = v4.length();
135        let result4 = (q * mag.cos() + v4 * (mag.sin() / mag)).normalize();
136        clamp_to_hemisphere(result4.truncate())
137    }
138
139    fn log(&self, from: Vec3, to: Vec3) -> Vec3 {
140        let qf = to_sphere(clamp_to_hemisphere(from));
141        let qt = to_sphere(clamp_to_hemisphere(to));
142        let d_dot = qf.dot(qt).clamp(-1.0, 1.0);
143        let perp4 = qt - d_dot * qf;
144        let n = perp4.length();
145        if n < LOG_PERP_MIN {
146            return Vec3::ZERO;
147        }
148        let half_chord = (qt - qf).length() * 0.5;
149        let d = 2.0 * half_chord.clamp(0.0, 1.0).asin();
150
151        perp4.truncate() * (d / n)
152    }
153
154    fn chart_envelope(&self) -> f32 {
155        std::f32::consts::FRAC_PI_2
156    }
157
158    fn valid_point(&self, p: Vec3) -> bool {
159        p.is_finite() && p.length_squared() <= SPHERE_R2_MAX
160    }
161
162    fn parallel_transport(&self, from: Vec3, to: Vec3, v: Vec3) -> Vec3 {
163        let from = clamp_to_hemisphere(from);
164        let to = clamp_to_hemisphere(to);
165        let qf = to_sphere(from);
166        let qt = to_sphere(to);
167        let vw = -v.dot(from) / qf.w;
168        let v4 = Vec4::new(v.x, v.y, v.z, vw);
169        // do Carmo, Riemannian Geometry, ch. 2.
170        let sum = qf + qt;
171        let denom = sum.length_squared() * 0.5;
172        let v4_transported = v4 - v4.dot(qt) / denom * sum;
173        v4_transported.truncate()
174    }
175}
176
177impl IsometryGroup for SphericalS3 {
178    type Iso = Iso4;
179
180    fn iso_identity(&self) -> Iso4 {
181        Iso4::IDENTITY
182    }
183
184    fn iso_compose(&self, a: Iso4, b: Iso4) -> Iso4 {
185        Iso4 {
186            matrix: a.matrix * b.matrix,
187        }
188    }
189
190    fn iso_inverse(&self, a: Iso4) -> Iso4 {
191        Iso4 {
192            matrix: a.matrix.transpose(),
193        }
194    }
195
196    fn iso_apply(&self, iso: Iso4, p: Vec3) -> Vec3 {
197        from_sphere(iso.matrix * to_sphere(clamp_to_hemisphere(p)))
198    }
199
200    fn iso_transport(&self, iso: Iso4, at: Vec3, v: Vec3) -> Vec3 {
201        let target = self.exp(at, v);
202        let m_at = self.iso_apply(iso, at);
203        let m_target = self.iso_apply(iso, target);
204        self.log(m_at, m_target)
205    }
206}
207
208impl WgslSpace for SphericalS3 {
209    fn wgsl_impl(&self) -> Cow<'static, str> {
210        Cow::Owned(format!(
211            r#"
212
213const LOAM_MAX_ARC: f32 = 1.5;
214const LOAM_S3_R2_MAX: f32 = {SPHERE_R2_MAX};
215const LOAM_S3_EXP_TANGENT_MIN_SQ: f32 = {EXP_TANGENT_MIN_SQ:e};
216const LOAM_S3_LOG_PERP_MIN: f32 = {LOG_PERP_MIN:e};
217{WGSL_FUNCTIONS}"#
218        ))
219    }
220
221    fn wgsl_accuracy(&self) -> WgslAccuracy {
222        WgslAccuracy::Exact
223    }
224}
225
226const WGSL_FUNCTIONS: &str = r#"
227fn loam_s3_clamp(p: vec3<f32>) -> vec3<f32> {
228    let r2 = dot(p, p);
229    if (r2 <= LOAM_S3_R2_MAX) { return p; }
230    return p * (sqrt(LOAM_S3_R2_MAX) / sqrt(r2));
231}
232
233fn loam_s3_lift(p: vec3<f32>) -> vec4<f32> {
234    let r2 = min(dot(p, p), LOAM_S3_R2_MAX);
235    return vec4<f32>(p.x, p.y, p.z, sqrt(1.0 - r2));
236}
237
238fn loam_origin_distance(p: vec3<f32>) -> f32 {
239
240    let r2 = min(dot(p, p), LOAM_S3_R2_MAX);
241    return asin(sqrt(r2));
242}
243
244fn loam_distance(a: vec3<f32>, b: vec3<f32>) -> f32 {
245    let qa = loam_s3_lift(loam_s3_clamp(a));
246    let qb = loam_s3_lift(loam_s3_clamp(b));
247    let half_chord = length(qa - qb) * 0.5;
248    return 2.0 * asin(clamp(half_chord, 0.0, 1.0));
249}
250
251fn loam_exp(at: vec3<f32>, v: vec3<f32>) -> vec3<f32> {
252    let p = loam_s3_clamp(at);
253    let n2 = dot(v, v);
254    if (n2 < LOAM_S3_EXP_TANGENT_MIN_SQ) { return p; }
255    let q = loam_s3_lift(p);
256    let vw = -dot(v, p) / q.w;
257    let v4 = vec4<f32>(v.x, v.y, v.z, vw);
258    let mag = length(v4);
259    let result4 = normalize(q * cos(mag) + v4 * (sin(mag) / mag));
260    return loam_s3_clamp(result4.xyz);
261}
262
263fn loam_log(p_from: vec3<f32>, p_to: vec3<f32>) -> vec3<f32> {
264    let qf = loam_s3_lift(loam_s3_clamp(p_from));
265    let qt = loam_s3_lift(loam_s3_clamp(p_to));
266    let d_dot = clamp(dot(qf, qt), -1.0, 1.0);
267    let perp4 = qt - d_dot * qf;
268    let n = length(perp4);
269    if (n < LOAM_S3_LOG_PERP_MIN) { return vec3<f32>(0.0, 0.0, 0.0); }
270    let half_chord = length(qt - qf) * 0.5;
271    let d = 2.0 * asin(clamp(half_chord, 0.0, 1.0));
272    return perp4.xyz * (d / n);
273}
274
275fn loam_s3_transport(p_from: vec3<f32>, p_to: vec3<f32>, v: vec3<f32>) -> vec3<f32> {
276    let pf = loam_s3_clamp(p_from);
277    let pt = loam_s3_clamp(p_to);
278    let qf = loam_s3_lift(pf);
279    let qt = loam_s3_lift(pt);
280    let vw = -dot(v, pf) / qf.w;
281    let v4 = vec4<f32>(v.x, v.y, v.z, vw);
282
283    let sum = qf + qt;
284    let denom = dot(sum, sum) * 0.5;
285    let v4t = v4 - (dot(v4, qt) / denom) * sum;
286    return v4t.xyz;
287}
288
289struct LoamGeodesicStep { p: vec3<f32>, v: vec3<f32> }
290
291fn loam_geodesic_step(p: vec3<f32>, v: vec3<f32>, s: f32) -> LoamGeodesicStep {
292    let next = loam_exp(p, v * s);
293    return LoamGeodesicStep(next, loam_s3_transport(p, next, v));
294}
295"#;
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use approx::assert_relative_eq;
301
302    fn s3() -> SphericalS3 {
303        SphericalS3
304    }
305
306    #[test]
307    fn distance_at_origin_matches_arc_length() {
308        let s = s3();
309
310        let r = 0.4;
311        let p = Vec3::new(r, 0.0, 0.0);
312        assert_relative_eq!(s.distance(Vec3::ZERO, p), r.asin(), epsilon = 1e-5);
313    }
314
315    #[test]
316    fn exp_tiny_vector_clamps_out_of_domain_basepoint() {
317        let s = s3();
318        let at = Vec3::new(2.0, 0.0, 0.0);
319        let tiny = Vec3::new(EXP_TANGENT_MIN_SQ.sqrt() * 0.1, 0.0, 0.0);
320        let got = s.exp(at, tiny);
321        let want = clamp_to_hemisphere(at);
322        assert_relative_eq!(got.x, want.x, epsilon = 1e-6);
323        assert_relative_eq!(got.y, want.y, epsilon = 1e-6);
324        assert_relative_eq!(got.z, want.z, epsilon = 1e-6);
325    }
326
327    #[test]
328    fn iso_translation_moves_origin_to_target() {
329        let s = s3();
330        let target = Vec3::new(0.2, -0.1, 0.15);
331        let iso = Iso4::from_translation(target);
332        let moved = s.iso_apply(iso, Vec3::ZERO);
333        assert_relative_eq!(moved.x, target.x, epsilon = 1e-5);
334        assert_relative_eq!(moved.y, target.y, epsilon = 1e-5);
335        assert_relative_eq!(moved.z, target.z, epsilon = 1e-5);
336    }
337
338    #[test]
339    fn parallel_transport_preserves_norm_near_antipode() {
340        let s = s3();
341        let lifted_norm = |p: Vec3, v: Vec3| {
342            let vw = -v.dot(p) / to_sphere(p).w;
343            Vec4::new(v.x, v.y, v.z, vw).length()
344        };
345        for w in [5e-3_f32, 2e-3, 1.2e-3] {
346            let b = w;
347            let a = (1.0 - w * w - b * b).sqrt();
348            let from = Vec3::new(a, b, 0.0);
349            let to = Vec3::new(-a, b, 0.0);
350
351            let v = Vec3::X;
352            let vt = s.parallel_transport(from, to, v);
353            let norm_from = lifted_norm(from, v);
354            assert_relative_eq!(lifted_norm(to, vt), norm_from, max_relative = 1e-5);
355            assert!(
356                (vt - v).length() > 0.5 * norm_from,
357                "transport should rotate an in-plane vector, got {vt:?}"
358            );
359        }
360    }
361
362    #[test]
363    fn small_scale_distance_matches_euclidean() {
364        let s = s3();
365        let eps = 1e-3;
366        let p = Vec3::new(eps, 0.0, 0.0);
367        assert_relative_eq!(s.distance(Vec3::ZERO, p), eps, epsilon = 1e-6);
368    }
369
370    #[test]
371    fn angle_excess_in_small_triangle_scales_with_area() {
372        let s = s3();
373        let l = 0.05_f32;
374        let a = Vec3::ZERO;
375        let b = s.exp(a, Vec3::new(l, 0.0, 0.0));
376        let c = s.exp(a, Vec3::new(l * 0.5, l * 3.0_f32.sqrt() * 0.5, 0.0));
377
378        let angle_at = |p: Vec3, q: Vec3, r: Vec3| -> f32 {
379            let u3 = s.log(p, q);
380            let w3 = s.log(p, r);
381
382            let qp = to_sphere(p);
383            let u4 = Vec4::new(u3.x, u3.y, u3.z, -u3.dot(p) / qp.w);
384            let w4 = Vec4::new(w3.x, w3.y, w3.z, -w3.dot(p) / qp.w);
385            (u4.dot(w4) / (u4.length() * w4.length()))
386                .clamp(-1.0, 1.0)
387                .acos()
388        };
389
390        let alpha = angle_at(a, b, c);
391        let beta = angle_at(b, a, c);
392        let gamma = angle_at(c, a, b);
393        let excess = (alpha + beta + gamma) - std::f32::consts::PI;
394        let expected_area = 3.0_f32.sqrt() / 4.0 * l * l;
395
396        assert!(
397            excess > 0.0,
398            "spherical triangle should have positive angle excess, got {excess}"
399        );
400        assert_relative_eq!(excess, expected_area, epsilon = 5e-4);
401    }
402
403    const PARITY_DIR: Vec3 = Vec3::new(0.6, -0.48, 0.64);
404
405    fn parity_points() -> [Vec3; 11] {
406        let shell = SPHERE_R2_MAX.sqrt();
407        [
408            Vec3::ZERO,
409            Vec3::new(LOG_PERP_MIN * 0.5, 0.0, 0.0),
410            Vec3::new(LOG_PERP_MIN * 1.5, 0.0, 0.0),
411            Vec3::new(LOG_PERP_MIN * 100.0, 0.0, 0.0),
412            Vec3::new(0.2, -0.3, 0.1),
413            Vec3::new(-0.45, 0.5, 0.35),
414            PARITY_DIR * 0.95,
415            PARITY_DIR * shell,
416            -PARITY_DIR * shell,
417            Vec3::new(1.0, 0.0, 0.0),
418            Vec3::new(2.0, -1.0, 0.5),
419        ]
420    }
421
422    #[test]
423    fn lift_floors_w_at_the_shell_even_where_the_clamp_overshoots_it() {
424        let floor = (1.0 - SPHERE_R2_MAX).sqrt();
425        let mut worst_w = f32::INFINITY;
426        for i in 0..6 {
427            for j in 0..6 {
428                for k in 0..6 {
429                    let dir = Vec3::new(i as f32 - 2.5, j as f32 - 2.5, k as f32 - 2.5).normalize();
430                    for scale in [1.0 + 1e-6, 1.01, 2.0, 1e3, 1e18] {
431                        let raw = dir * scale;
432                        let clamped = clamp_to_hemisphere(raw);
433                        worst_w = worst_w.min(to_sphere(raw).w).min(to_sphere(clamped).w);
434                    }
435                }
436            }
437        }
438        assert!(
439            worst_w >= floor,
440            "lift w reached {worst_w:e}, under the shell floor {floor:e}"
441        );
442
443        assert!(to_sphere(Vec3::splat(f32::NAN)).w >= floor);
444        assert!(to_sphere(Vec3::splat(f32::INFINITY)).w >= floor);
445    }
446
447    #[test]
448    fn transport_denominator_is_bounded_below_by_the_saturation_shell() {
449        let chart_min = 2.0 * (1.0 - SPHERE_R2_MAX);
450        let mut worst = f32::INFINITY;
451        for from in parity_points() {
452            for to in parity_points() {
453                let sum = to_sphere(clamp_to_hemisphere(from)) + to_sphere(clamp_to_hemisphere(to));
454                worst = worst.min(sum.length_squared() * 0.5);
455            }
456        }
457        assert!(
458            worst >= chart_min,
459            "denominator reached {worst:e}, under the shell bound {chart_min:e}"
460        );
461
462        assert!(
463            worst <= chart_min * 1.5,
464            "closest approach {worst:e} is not the chart minimum {chart_min:e}"
465        );
466    }
467
468    #[test]
469    fn translation_guard_separates_degenerate_targets_from_representable_ones() {
470        let s = s3();
471        let below = Vec3::new(ISO_TRANSLATION_MIN_ARC * 0.5, 0.0, 0.0);
472        assert_eq!(Iso4::from_translation(below).matrix, Mat4::IDENTITY);
473
474        let above = Vec3::new(ISO_TRANSLATION_MIN_ARC * 1.5, 0.0, 0.0);
475        let moved = s.iso_apply(Iso4::from_translation(above), Vec3::ZERO);
476        assert_relative_eq!(moved.x, above.x, max_relative = 1e-5);
477        assert_eq!(moved.y, 0.0);
478        assert_eq!(moved.z, 0.0);
479    }
480}
loam_math/spherical.rs

loam_math/
spherical.rs
