polyhorn_ui/linalg/
geometry.rs

1use num_traits::Float;
2use std::ops::{Index, IndexMut};
3
4/// A point in a 3D space.
5#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
6pub struct Point3D<T> {
7    /// The x-coordinate of this point.
8    pub x: T,
9
10    /// The y-coordinate of this point.
11    pub y: T,
12
13    /// The z-coordinate of this point.
14    pub z: T,
15}
16
17impl<T> Point3D<T> {
18    /// Returns a new point in a 3D space with the given coordinates.
19    pub fn new(x: T, y: T, z: T) -> Point3D<T> {
20        Point3D { x, y, z }
21    }
22}
23
24impl<T> Point3D<T>
25where
26    T: Float,
27{
28    /// Returns the L2 norm of this vector (note to MLs: incl. the square root).
29    pub fn l2_norm(self) -> T {
30        (self.x * self.x + self.y * self.y + self.z * self.z).sqrt()
31    }
32
33    /// Returns a normalized vector by dividing each component by the L2 norm
34    /// of this vector.
35    pub fn normalize(self) -> Point3D<T> {
36        let norm = self.l2_norm();
37
38        Point3D::new(self.x / norm, self.y / norm, self.z / norm)
39    }
40}
41
42/// A 3D rotation quaternion.
43#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
44pub struct Quaternion3D<T> {
45    /// The x-element of this quaternion.
46    pub x: T,
47
48    /// The y-element of this quaternion.
49    pub y: T,
50
51    /// The z-element of this quaternion.
52    pub z: T,
53
54    /// The w-element of this quaternion.
55    pub w: T,
56}
57
58impl<T> Quaternion3D<T> {
59    /// Returns a new quaternion with the given elements.
60    pub fn new(x: T, y: T, z: T, w: T) -> Quaternion3D<T> {
61        Quaternion3D { x, y, z, w }
62    }
63
64    /// Applies the given transformation to each element of the given quaternion
65    /// and returns the result.
66    pub fn map<F, O>(self, mut op: F) -> Quaternion3D<O>
67    where
68        F: FnMut(T) -> O,
69    {
70        Quaternion3D {
71            x: op(self.x),
72            y: op(self.y),
73            z: op(self.z),
74            w: op(self.w),
75        }
76    }
77
78    /// Returns a new quaternion with references to each of the original
79    /// elements. This is particularly useful when `T` does not implement
80    /// `Copy`.
81    pub fn as_ref(&self) -> Quaternion3D<&T> {
82        Quaternion3D {
83            x: &self.x,
84            y: &self.y,
85            z: &self.z,
86            w: &self.w,
87        }
88    }
89}
90
91impl<T> Quaternion3D<T>
92where
93    T: Float,
94{
95    /// Converts the rotation with the given angle radians around a vector with
96    /// the given coordinates, into a quaternion with the corresponding
97    /// elements.
98    pub fn with_angle(angle: T, rx: T, ry: T, rz: T) -> Quaternion3D<T> {
99        let r = Point3D::new(rx, ry, rz);
100        let u = r.normalize();
101        let angle = angle * r.l2_norm();
102
103        let one = T::one();
104        let two = one + one;
105
106        let x = (angle / two).sin() * u.x;
107        let y = (angle / two).sin() * u.y;
108        let z = (angle / two).sin() * u.z;
109        let w = (angle / two).cos();
110
111        Quaternion3D { x, y, z, w }
112    }
113
114    /// Computes the inner product of this quaternion and the given quaternion.
115    pub fn dot(&self, other: &Self) -> T {
116        self.x * other.x + self.y * other.y + self.z * other.z + self.w * other.w
117    }
118
119    /// Interpolates between the given quaternions with the given weight. If the
120    /// weight is one, `self` is returned. If the weight is zero, `other` is
121    /// returned.
122    pub fn mix(self, weight: T, other: Quaternion3D<T>) -> Quaternion3D<T> {
123        let mut product = self.dot(&other);
124        product = product.min(T::one());
125        product = product.max(T::one().neg());
126
127        if product.abs().is_one() {
128            return self;
129        }
130
131        let theta = product.acos();
132        let w = ((T::one() - weight) * theta).sin() * (T::one() - product * product).sqrt().recip();
133
134        let mut result = self;
135
136        for i in 0..4 {
137            let a = self[i] * (((T::one() - weight) * theta).cos() - product * w);
138            let b = other[i] * w;
139            result[i] = a + b;
140        }
141
142        result
143    }
144
145    /// Returns the element-wise addition of this quaternion and the given
146    /// quaternion.
147    pub fn addition(&self, other: &Quaternion3D<T>) -> Quaternion3D<T> {
148        Quaternion3D {
149            x: self.x + other.x,
150            y: self.y + other.y,
151            z: self.z + other.z,
152            w: self.w + other.w,
153        }
154    }
155
156    /// Returns the element-wise subtraction of this quaternion and the given
157    /// quaternion.
158    pub fn subtract(&self, other: &Quaternion3D<T>) -> Quaternion3D<T> {
159        Quaternion3D {
160            x: self.x - other.x,
161            y: self.y - other.y,
162            z: self.z - other.z,
163            w: self.w - other.w,
164        }
165    }
166}
167
168impl<T> Index<usize> for Quaternion3D<T> {
169    type Output = T;
170
171    fn index(&self, index: usize) -> &Self::Output {
172        assert!(index < 4);
173
174        match index {
175            0 => &self.x,
176            1 => &self.y,
177            2 => &self.z,
178            3 => &self.w,
179            _ => unreachable!(),
180        }
181    }
182}
183
184impl<T> IndexMut<usize> for Quaternion3D<T> {
185    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
186        assert!(index < 4);
187
188        match index {
189            0 => &mut self.x,
190            1 => &mut self.y,
191            2 => &mut self.z,
192            3 => &mut self.w,
193            _ => unreachable!(),
194        }
195    }
196}