roche/
vec3.rs

1use pyo3::prelude::*;
2use std::ops;
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: x, y: y, z: 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.clone() + v,
107            Vec3OrF64::F64(f) => self.clone() + f,
108        }
109    }
110    fn __radd__(&self, other: Vec3OrF64) -> Self {
111        match other {
112            Vec3OrF64::Vec3(v) => v + self.clone(),
113            Vec3OrF64::F64(f) => f + self.clone(),
114        }
115    }
116    fn __sub__(&self, other: Vec3OrF64) -> Self {
117        match other {
118            Vec3OrF64::Vec3(v) => self.clone() - v,
119            Vec3OrF64::F64(f) => self.clone() - f,
120        }
121    }
122    fn __rsub__(&self, other: Vec3OrF64) -> Self {
123        match other {
124            Vec3OrF64::Vec3(v) => v - self.clone(),
125            Vec3OrF64::F64(f) => f - self.clone(),
126        }
127    }
128    fn __mul__(&self, other: f64) -> Self {
129        self.clone() * other
130    }
131    fn __rmul__(&self, other: f64) -> Self {
132        other * self.clone()
133    }
134    fn __truediv__(&self, other: f64) -> Self {
135        self.clone() / other
136    }
137    fn __rtruediv__(&self, other: f64) -> Self {
138        other / self.clone()
139    }
140    fn __neg__(&self) -> Self {
141        -self.clone()
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
336#[cfg(test)]
337mod tests {
338    use super::*;
339
340    #[test]
341    fn set() {
342        let mut v1 = Vec3::new(0.0, 0.0, 0.0);
343        v1.set(1.0, 1.0, 1.0);
344        assert_eq!(v1, Vec3::new(1.0, 1.0, 1.0));
345    }
346
347    #[test]
348    fn unit() {
349        let mut v1 = Vec3::new(2.0, 3.0, 6.0);
350        v1.unit();
351        assert_eq!(v1, Vec3::new(2.0 / 7.0, 3.0 / 7.0, 6.0 / 7.0));
352    }
353
354    #[test]
355    fn norm() {
356        let v1 = Vec3::new(2.0, 3.0, 6.0);
357        assert_eq!(v1.norm(), Vec3::new(2.0 / 7.0, 3.0 / 7.0, 6.0 / 7.0));
358    }
359
360    #[test]
361    fn length() {
362        let v1 = Vec3::new(2.0, 3.0, 6.0);
363        assert_eq!(v1.length(), 7.0);
364    }
365
366    #[test]
367    fn sqr() {
368        let v1 = Vec3::new(2.0, 3.0, 6.0);
369        assert_eq!(v1.sqr(), 49.0);
370    }
371
372    #[test]
373    fn dot() {
374        let v1 = Vec3::new(1.0, 3.0, -5.0);
375        let v2 = Vec3::new(4.0, -2.0, -1.0);
376        assert_eq!(v1.dot(&v2), 3.0);
377        assert_eq!(v2.dot(&v1), 3.0);
378    }
379
380    #[test]
381    fn cross() {
382        let v1 = Vec3::new(3.0, -3.0, 1.0);
383        let v2 = Vec3::new(4.0, 9.0, 2.0);
384        assert_eq!(v1.cross(&v2), Vec3::new(-15.0, -2.0, 39.0));
385    }
386
387    #[test]
388    fn mulassign_f64() {
389        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
390        v1 *= 2.0;
391        assert_eq!(v1, Vec3::new(6.0, -6.0, 2.0));
392    }
393
394    #[test]
395    fn addassign_f64() {
396        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
397        v1 += 2.0;
398        assert_eq!(v1, Vec3::new(5.0, -1.0, 3.0));
399    }
400
401    #[test]
402    fn subassign_f64() {
403        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
404        v1 -= 2.0;
405        assert_eq!(v1, Vec3::new(1.0, -5.0, -1.0));
406    }
407
408    #[test]
409    fn addassign_vec3() {
410        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
411        let v2 = Vec3::new(1.0, 2.0, -5.0);
412        v1 += v2;
413        assert_eq!(v1, Vec3::new(4.0, -1.0, -4.0));
414    }
415
416    #[test]
417    fn subassign_vec3() {
418        let mut v1 = Vec3::new(3.0, -3.0, 1.0);
419        let v2 = Vec3::new(1.0, 2.0, -5.0);
420        v1 -= v2;
421        assert_eq!(v1, Vec3::new(2.0, -5.0, 6.0));
422    }
423
424    #[test]
425    fn add_vec3vec3() {
426        let v1 = Vec3::new(3.0, -3.0, 1.0);
427        let v2 = Vec3::new(1.0, 2.0, -5.0);
428        assert_eq!(v1 + v2, Vec3::new(4.0, -1.0, -4.0));
429    }
430
431    #[test]
432    fn sub_vec3vec3() {
433        let v1 = Vec3::new(3.0, -3.0, 1.0);
434        let v2 = Vec3::new(1.0, 2.0, -5.0);
435        assert_eq!(v1 - v2, Vec3::new(2.0, -5.0, 6.0));
436    }
437
438    #[test]
439    fn divassign_f64() {
440        let mut v1 = Vec3::new(4.0, -2.0, 1.0);
441        v1 /= 2.0;
442        assert_eq!(v1, Vec3::new(2.0, -1.0, 0.5));
443    }
444
445    #[test]
446    fn div_vec3f64() {
447        let v1 = Vec3::new(3.0, -3.0, 1.0);
448        assert_eq!(v1 / 2.0, Vec3::new(1.5, -1.5, 0.5));
449    }
450
451    #[test]
452    fn div_f64vec3() {
453        let v1 = Vec3::new(3.0, -3.0, 1.0);
454        assert_eq!(9.0 / v1, Vec3::new(3.0, -3.0, 9.0));
455    }
456
457    #[test]
458    fn mul_vec3f64() {
459        let v1 = Vec3::new(3.0, -3.0, 1.0);
460        assert_eq!(v1 * 2.0, Vec3::new(6.0, -6.0, 2.0));
461    }
462
463    #[test]
464    fn mul_f64vec3() {
465        let v1 = Vec3::new(3.0, -3.0, 1.0);
466        assert_eq!(2.0 * v1, Vec3::new(6.0, -6.0, 2.0));
467    }
468
469    #[test]
470    fn add_vec3f64() {
471        let v1 = Vec3::new(3.0, -3.0, 1.0);
472        assert_eq!(v1 + 2.0, Vec3::new(5.0, -1.0, 3.0));
473    }
474
475    #[test]
476    fn add_f64vec3() {
477        let v1 = Vec3::new(3.0, -3.0, 1.0);
478        assert_eq!(2.0 + v1, Vec3::new(5.0, -1.0, 3.0));
479    }
480
481    #[test]
482    fn sub_vec3f64() {
483        let v1 = Vec3::new(3.0, -3.0, 1.0);
484        assert_eq!(v1 - 2.0, Vec3::new(1.0, -5.0, -1.0));
485    }
486
487    #[test]
488    fn sub_f64vec3() {
489        let v1 = Vec3::new(3.0, -3.0, 1.0);
490        assert_eq!(2.0 - v1, Vec3::new(-1.0, 5.0, 1.0));
491    }
492
493    #[test]
494    fn neg() {
495        let v1 = Vec3::new(3.0, -3.0, 1.0);
496        assert_eq!(-v1, Vec3::new(-3.0, 3.0, -1.0));
497    }
498}
roche/vec3.rs

roche/
vec3.rs
