roche/
vec3.rs

1use pyo3::prelude::*;
2use std::{ops, fmt};
3
4// useful for defining python operators
5#[derive(FromPyObject)]
6enum Vec3OrF64 {
7    Vec3(Vec3),
8    F64(f64),
9}
10
11#[pyclass(from_py_object)]
12#[derive(Debug, Copy, Clone, PartialEq)]
13pub struct Vec3 {
14    #[pyo3(get, set)]
15    pub x: f64,
16    #[pyo3(get, set)]
17    pub y: f64,
18    #[pyo3(get, set)]
19    pub z: f64,
20}
21
22#[pymethods]
23impl Vec3 {
24    #[new]
25    pub fn new(x: f64, y: f64, z: f64) -> Self {
26        Self { x, y, z }
27    }
28
29    #[staticmethod]
30    pub fn cofm1() -> Self {
31        Self {
32            x: 0.0,
33            y: 0.0,
34            z: 0.0,
35        }
36    }
37
38    #[staticmethod]
39    pub fn cofm2() -> Self {
40        Self {
41            x: 1.0,
42            y: 0.0,
43            z: 0.0,
44        }
45    }
46
47    pub fn set(&mut self, x: f64, y: f64, z: f64) {
48        self.x = x;
49        self.y = y;
50        self.z = z;
51    }
52
53    // Normalises the vector in place
54    pub fn unit(&mut self) {
55        let norm = self.length();
56        self.x /= norm;
57        self.y /= norm;
58        self.z /= norm;
59    }
60
61    // Returns a normalised version of the vector
62    pub fn norm(&self) -> Self {
63        let norm = self.length();
64        Self {
65            x: self.x / norm,
66            y: self.y / norm,
67            z: self.z / norm,
68        }
69    }
70
71    // For `__repr__` we want to return a string that Python code could use to recreate
72    // the `Vec3`, like `Vec3(5, 6, 7)` for example.
73    fn __repr__(&self) -> String {
74        // We use the `format!` macro to create a string. Its first argument is a
75        // format string, followed by any number of parameters which replace the
76        // `{}`'s in the format string.
77        format!("Vec3({}, {}, {})", self.x, self.y, self.z)
78    }
79
80    // Returns the length of the vector
81    pub fn length(&self) -> f64 {
82        (self.x.powi(2) + self.y.powi(2) + self.z.powi(2)).sqrt()
83    }
84
85    // Returns the squared length of the vector
86    pub fn sqr(&self) -> f64 {
87        self.x.powi(2) + self.y.powi(2) + self.z.powi(2)
88    }
89
90    // Returns the dot product of two vectors
91    pub fn dot(&self, other: &Vec3) -> f64 {
92        self.x * other.x + self.y * other.y + self.z * other.z
93    }
94
95    // Returns the cross product of two vectors
96    pub fn cross(&self, other: &Vec3) -> Vec3 {
97        let temp_x = self.y * other.z - self.z * other.y;
98        let temp_y = self.z * other.x - self.x * other.z;
99        let temp_z = self.x * other.y - self.y * other.x;
100        Vec3::new(temp_x, temp_y, temp_z)
101    }
102
103    // Python dunder methods for arithmetic
104    fn __add__(&self, other: Vec3OrF64) -> Self {
105        match other {
106            Vec3OrF64::Vec3(v) => *self + v,
107            Vec3OrF64::F64(f) => *self + f,
108        }
109    }
110    fn __radd__(&self, other: Vec3OrF64) -> Self {
111        match other {
112            Vec3OrF64::Vec3(v) => v + *self,
113            Vec3OrF64::F64(f) => f + *self,
114        }
115    }
116    fn __sub__(&self, other: Vec3OrF64) -> Self {
117        match other {
118            Vec3OrF64::Vec3(v) => *self - v,
119            Vec3OrF64::F64(f) => *self - f,
120        }
121    }
122    fn __rsub__(&self, other: Vec3OrF64) -> Self {
123        match other {
124            Vec3OrF64::Vec3(v) => v - *self,
125            Vec3OrF64::F64(f) => f - *self,
126        }
127    }
128    fn __mul__(&self, other: f64) -> Self {
129        *self * other
130    }
131    fn __rmul__(&self, other: f64) -> Self {
132        other * *self
133    }
134    fn __truediv__(&self, other: f64) -> Self {
135        *self / other
136    }
137    fn __rtruediv__(&self, other: f64) -> Self {
138        other / *self
139    }
140    fn __neg__(&self) -> Self {
141        -*self
142    }
143}
144
145// in-place multiplication by f64
146impl ops::MulAssign<f64> for Vec3 {
147    fn mul_assign(&mut self, rhs: f64) {
148        self.x *= rhs;
149        self.y *= rhs;
150        self.z *= rhs;
151    }
152}
153
154// in-place division by f64
155impl ops::DivAssign<f64> for Vec3 {
156    fn div_assign(&mut self, rhs: f64) {
157        self.x /= rhs;
158        self.y /= rhs;
159        self.z /= rhs;
160    }
161}
162
163// in-place addition by f64
164impl ops::AddAssign<f64> for Vec3 {
165    fn add_assign(&mut self, rhs: f64) {
166        self.x += rhs;
167        self.y += rhs;
168        self.z += rhs;
169    }
170}
171
172// in-place subtraction by f64
173impl ops::SubAssign<f64> for Vec3 {
174    fn sub_assign(&mut self, rhs: f64) {
175        self.x -= rhs;
176        self.y -= rhs;
177        self.z -= rhs;
178    }
179}
180
181// in-place addition of vector
182impl ops::AddAssign<Vec3> for Vec3 {
183    fn add_assign(&mut self, rhs: Vec3) {
184        self.x += rhs.x;
185        self.y += rhs.y;
186        self.z += rhs.z;
187    }
188}
189
190// in-place subtraction of vector
191impl ops::SubAssign<Vec3> for Vec3 {
192    fn sub_assign(&mut self, rhs: Vec3) {
193        self.x -= rhs.x;
194        self.y -= rhs.y;
195        self.z -= rhs.z;
196    }
197}
198
199// Sum of two vectors
200impl ops::Add for Vec3 {
201    type Output = Self;
202    fn add(self, other: Self) -> Self {
203        Self {
204            x: self.x + other.x,
205            y: self.y + other.y,
206            z: self.z + other.z,
207        }
208    }
209}
210
211// Sum of vector and f64
212impl ops::Add<f64> for Vec3 {
213    type Output = Self;
214    fn add(self, other: f64) -> Self {
215        Self {
216            x: self.x + other,
217            y: self.y + other,
218            z: self.z + other,
219        }
220    }
221}
222
223// Sum of f64 and vector
224impl ops::Add<Vec3> for f64 {
225    type Output = Vec3;
226    fn add(self, other: Vec3) -> Vec3 {
227        Vec3 {
228            x: self + other.x,
229            y: self + other.y,
230            z: self + other.z,
231        }
232    }
233}
234
235// Difference between two vectors
236impl ops::Sub for Vec3 {
237    type Output = Self;
238    fn sub(self, other: Self) -> Self {
239        Self {
240            x: self.x - other.x,
241            y: self.y - other.y,
242            z: self.z - other.z,
243        }
244    }
245}
246
247// Difference between vector and f64
248impl ops::Sub<f64> for Vec3 {
249    type Output = Self;
250    fn sub(self, other: f64) -> Self {
251        Self {
252            x: self.x - other,
253            y: self.y - other,
254            z: self.z - other,
255        }
256    }
257}
258
259// Difference between f64 and vector
260impl ops::Sub<Vec3> for f64 {
261    type Output = Vec3;
262    fn sub(self, other: Vec3) -> Vec3 {
263        Vec3 {
264            x: self - other.x,
265            y: self - other.y,
266            z: self - other.z,
267        }
268    }
269}
270
271// Multiplication of Vec3 by f64
272impl ops::Mul<f64> for Vec3 {
273    type Output = Self;
274
275    fn mul(self, rhs: f64) -> Self {
276        Self {
277            x: self.x * rhs,
278            y: self.y * rhs,
279            z: self.z * rhs,
280        }
281    }
282}
283
284// Multiplication of f64 by Vec3
285impl ops::Mul<Vec3> for f64 {
286    type Output = Vec3;
287
288    fn mul(self, rhs: Vec3) -> Vec3 {
289        Vec3 {
290            x: self * rhs.x,
291            y: self * rhs.y,
292            z: self * rhs.z,
293        }
294    }
295}
296
297// Division of Vec3 by f64
298impl ops::Div<f64> for Vec3 {
299    type Output = Self;
300
301    fn div(self, rhs: f64) -> Self {
302        Self {
303            x: self.x / rhs,
304            y: self.y / rhs,
305            z: self.z / rhs,
306        }
307    }
308}
309
310// Division of f64 by Vec3
311impl ops::Div<Vec3> for f64 {
312    type Output = Vec3;
313
314    fn div(self, rhs: Vec3) -> Vec3 {
315        Vec3 {
316            x: self / rhs.x,
317            y: self / rhs.y,
318            z: self / rhs.z,
319        }
320    }
321}
322
323// Negative of a vector
324impl ops::Neg for Vec3 {
325    type Output = Self;
326
327    fn neg(self) -> Self {
328        Vec3 {
329            x: -self.x,
330            y: -self.y,
331            z: -self.z,
332        }
333    }
334}
335
336impl fmt::Display for Vec3 {
337    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
338        write!(
339            f,
340            "Vec3({}, {}, {})",
341            self.x,
342            self.y,
343            self.z
344        )
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    #[test]
353    fn set() {
354        let mut v1 = Vec3::new(0.0, 0.0, 0.0);
355        v1.set(1.0, 1.0, 1.0);
356        assert_eq!(v1, Vec3::new(1.0, 1.0, 1.0));
357    }
358
359    #[test]
360    fn unit() {
361        let mut v1 = Vec3::new(2.0, 3.0, 6.0);
362        v1.unit();
363        assert_eq!(v1, Vec3::new(2.0 / 7.0, 3.0 / 7.0, 6.0 / 7.0));
364    }
365
366    #[test]
367    fn norm() {
368        let v1 = Vec3::new(2.0, 3.0, 6.0);
369        assert_eq!(v1.norm(), Vec3::new(2.0 / 7.0, 3.0 / 7.0, 6.0 / 7.0));
370    }
371
372    #[test]
373    fn length() {
374        let v1 = Vec3::new(2.0, 3.0, 6.0);
375        assert_eq!(v1.length(), 7.0);
376    }
377
378    #[test]
379    fn sqr() {
380        let v1 = Vec3::new(2.0, 3.0, 6.0);
381        assert_eq!(v1.sqr(), 49.0);
382    }
383
384    #[test]
385    fn dot() {
386        let v1 = Vec3::new(1.0, 3.0, -5.0);
387        let v2 = Vec3::new(4.0, -2.0, -1.0);
388        assert_eq!(v1.dot(&v2), 3.0);
389        assert_eq!(v2.dot(&v1), 3.0);
390    }
391
392    #[test]
393    fn cross() {
394        let v1 = Vec3::new(3.0, -3.0, 1.0);
395        let v2 = Vec3::new(4.0, 9.0, 2.0);
396        assert_eq!(v1.cross(&v2), Vec3::new(-15.0, -2.0, 39.0));
397    }
398
399    #[test]
400    fn mulassign_f64() {
401        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
402        v1 *= 2.0;
403        assert_eq!(v1, Vec3::new(6.0, -6.0, 2.0));
404    }
405
406    #[test]
407    fn addassign_f64() {
408        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
409        v1 += 2.0;
410        assert_eq!(v1, Vec3::new(5.0, -1.0, 3.0));
411    }
412
413    #[test]
414    fn subassign_f64() {
415        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
416        v1 -= 2.0;
417        assert_eq!(v1, Vec3::new(1.0, -5.0, -1.0));
418    }
419
420    #[test]
421    fn addassign_vec3() {
422        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
423        let v2 = Vec3::new(1.0, 2.0, -5.0);
424        v1 += v2;
425        assert_eq!(v1, Vec3::new(4.0, -1.0, -4.0));
426    }
427
428    #[test]
429    fn subassign_vec3() {
430        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
431        let v2 = Vec3::new(1.0, 2.0, -5.0);
432        v1 -= v2;
433        assert_eq!(v1, Vec3::new(2.0, -5.0, 6.0));
434    }
435
436    #[test]
437    fn add_vec3vec3() {
438        let v1 = Vec3::new(3.0, -3.0, 1.0);
439        let v2 = Vec3::new(1.0, 2.0, -5.0);
440        assert_eq!(v1 + v2, Vec3::new(4.0, -1.0, -4.0));
441    }
442
443    #[test]
444    fn sub_vec3vec3() {
445        let v1 = Vec3::new(3.0, -3.0, 1.0);
446        let v2 = Vec3::new(1.0, 2.0, -5.0);
447        assert_eq!(v1 - v2, Vec3::new(2.0, -5.0, 6.0));
448    }
449
450    #[test]
451    fn divassign_f64() {
452        let mut v1 = Vec3::new(4.0, -2.0, 1.0);
453        v1 /= 2.0;
454        assert_eq!(v1, Vec3::new(2.0, -1.0, 0.5));
455    }
456
457    #[test]
458    fn div_vec3f64() {
459        let v1 = Vec3::new(3.0, -3.0, 1.0);
460        assert_eq!(v1 / 2.0, Vec3::new(1.5, -1.5, 0.5));
461    }
462
463    #[test]
464    fn div_f64vec3() {
465        let v1 = Vec3::new(3.0, -3.0, 1.0);
466        assert_eq!(9.0 / v1, Vec3::new(3.0, -3.0, 9.0));
467    }
468
469    #[test]
470    fn mul_vec3f64() {
471        let v1 = Vec3::new(3.0, -3.0, 1.0);
472        assert_eq!(v1 * 2.0, Vec3::new(6.0, -6.0, 2.0));
473    }
474
475    #[test]
476    fn mul_f64vec3() {
477        let v1 = Vec3::new(3.0, -3.0, 1.0);
478        assert_eq!(2.0 * v1, Vec3::new(6.0, -6.0, 2.0));
479    }
480
481    #[test]
482    fn add_vec3f64() {
483        let v1 = Vec3::new(3.0, -3.0, 1.0);
484        assert_eq!(v1 + 2.0, Vec3::new(5.0, -1.0, 3.0));
485    }
486
487    #[test]
488    fn add_f64vec3() {
489        let v1 = Vec3::new(3.0, -3.0, 1.0);
490        assert_eq!(2.0 + v1, Vec3::new(5.0, -1.0, 3.0));
491    }
492
493    #[test]
494    fn sub_vec3f64() {
495        let v1 = Vec3::new(3.0, -3.0, 1.0);
496        assert_eq!(v1 - 2.0, Vec3::new(1.0, -5.0, -1.0));
497    }
498
499    #[test]
500    fn sub_f64vec3() {
501        let v1 = Vec3::new(3.0, -3.0, 1.0);
502        assert_eq!(2.0 - v1, Vec3::new(-1.0, 5.0, 1.0));
503    }
504
505    #[test]
506    fn neg() {
507        let v1 = Vec3::new(3.0, -3.0, 1.0);
508        assert_eq!(-v1, Vec3::new(-3.0, 3.0, -1.0));
509    }
510}
roche/vec3.rs

roche/
vec3.rs
