1use crate::small_linalg::{Matrix2, Matrix2x3, Matrix3, Vector3};
2use crate::{small_linalg::Vector2, PhysicsQuantity};
3use ndarray::{Array2, ArrayView2};
4use std::ops::{Add, Div, Index, IndexMut, Mul, Neg, Sub};
5
6#[derive(Clone, Copy, PartialEq, Debug)]
7pub struct Matrix3x2<T: PhysicsQuantity> {
8 pub data: [[T; 2]; 3],
9}
10
11impl<T: PhysicsQuantity> Matrix3x2<T> {
12 pub fn new(data: [[T; 2]; 3]) -> Self {
13 Self { data }
14 }
15
16 pub fn zero() -> Self {
17 Self {
18 data: [[T::zero(); 2]; 3],
19 }
20 }
21
22 pub fn from_f64(data: [[f64; 2]; 3]) -> Self {
23 let mut converted = [[T::zero(); 2]; 3];
24 for i in 0..3 {
25 for j in 0..2 {
26 converted[i][j] = T::from_raw(data[i][j]);
27 }
28 }
29 Self::new(converted)
30 }
31
32 pub fn as_f64(&self) -> Matrix3x2<f64> {
33 Matrix3x2::new([
34 [self.data[0][0].as_f64(), self.data[0][1].as_f64()],
35 [self.data[1][0].as_f64(), self.data[1][1].as_f64()],
36 [self.data[2][0].as_f64(), self.data[2][1].as_f64()],
37 ])
38 }
39
40 pub fn transpose(&self) -> Matrix2x3<T> {
41 Matrix2x3::new([
42 [self.data[0][0], self.data[1][0], self.data[2][0]],
43 [self.data[0][1], self.data[1][1], self.data[2][1]],
44 ])
45 }
46
47 pub fn column(&self, i: usize) -> Vector3<T> {
48 assert!(i < 2);
49 Vector3::new([self.data[0][i], self.data[1][i], self.data[2][i]])
50 }
51
52 pub fn row(&self, i: usize) -> Vector2<T> {
53 assert!(i < 2);
54 Vector2::new(self.data[i])
55 }
56
57 pub fn optimize(&mut self) {
58 for row in &mut self.data {
59 for value in row {
60 value.optimize();
61 }
62 }
63 }
64
65 pub fn abs(&self) -> Self {
66 Self {
67 data: [
68 [self.data[0][0].abs(), self.data[0][1].abs()],
69 [self.data[1][0].abs(), self.data[1][1].abs()],
70 [self.data[2][0].abs(), self.data[2][1].abs()],
71 ],
72 }
73 }
74
75 pub fn to_ndarray(&self) -> Result<Array2<T>, String> {
76 Array2::from_shape_vec(
77 (3, 2),
78 vec![
79 self.data[0][0],
80 self.data[0][1],
81 self.data[1][0],
82 self.data[1][1],
83 self.data[2][0],
84 self.data[2][1],
85 ],
86 )
87 .map_err(|_| "Matrix3x2::to_ndarray: shape mismatch".to_string())
88 }
89
90 pub fn from_ndarray(array: ArrayView2<T>) -> Result<Self, String> {
91 if array.shape() != [3, 2] {
92 return Err("Expected 3x2 ndarray for Matrix3x2".to_string());
93 }
94
95 Ok(Self::new([
96 [array[[0, 0]], array[[0, 1]]],
97 [array[[1, 0]], array[[1, 1]]],
98 [array[[2, 0]], array[[2, 1]]],
99 ]))
100 }
101}
102
103impl<T: PhysicsQuantity + Neg<Output = T>> Neg for Matrix3x2<T> {
104 type Output = Self;
105 fn neg(self) -> Self {
106 Self::new([
107 [-self.data[0][0], -self.data[0][1]],
108 [-self.data[1][0], -self.data[1][1]],
109 [-self.data[2][0], -self.data[2][1]],
110 ])
111 }
112}
113
114impl<T: PhysicsQuantity + Add<Output = T>> Add for Matrix3x2<T> {
115 type Output = Self;
116
117 fn add(self, rhs: Self) -> Self {
118 let mut out = [[T::zero(); 2]; 3];
119 for (i, row) in out.iter_mut().enumerate() {
120 for (j, val) in row.iter_mut().enumerate() {
121 *val = self.data[i][j] + rhs.data[i][j];
122 }
123 }
124 Self::new(out)
125 }
126}
127
128impl<T: PhysicsQuantity + Sub<Output = T>> Sub for Matrix3x2<T> {
129 type Output = Self;
130
131 fn sub(self, rhs: Self) -> Self {
132 let mut out = [[T::zero(); 2]; 3];
133 for (i, row) in out.iter_mut().enumerate() {
134 for (j, val) in row.iter_mut().enumerate() {
135 *val = self.data[i][j] - rhs.data[i][j];
136 }
137 }
138 Self::new(out)
139 }
140}
141
142impl<T: PhysicsQuantity> Index<(usize, usize)> for Matrix3x2<T> {
143 type Output = T;
144 fn index(&self, idx: (usize, usize)) -> &Self::Output {
145 &self.data[idx.0][idx.1]
146 }
147}
148
149impl<T: PhysicsQuantity> IndexMut<(usize, usize)> for Matrix3x2<T> {
150 fn index_mut(&mut self, idx: (usize, usize)) -> &mut Self::Output {
151 &mut self.data[idx.0][idx.1]
152 }
153}
154
155impl<T: PhysicsQuantity> Matrix3x2<T> {
156 pub fn dot_vector2<U, V>(&self, vec: &Vector2<U>) -> Vector3<V>
157 where
158 U: PhysicsQuantity,
159 V: PhysicsQuantity + Add<Output = V>,
160 T: Mul<U, Output = V>,
161 {
162 let r0 = self.data[0][0] * vec[0] + self.data[0][1] * vec[1];
163 let r1 = self.data[1][0] * vec[0] + self.data[1][1] * vec[1];
164 let r2 = self.data[2][0] * vec[0] + self.data[2][1] * vec[1];
165 Vector3::new([r0, r1, r2])
166 }
167}
168
169impl<T, U, V> Mul<Vector2<U>> for Matrix3x2<T>
170where
171 T: PhysicsQuantity + Mul<U, Output = V>,
172 U: PhysicsQuantity,
173 V: PhysicsQuantity + Add<Output = V>,
174{
175 type Output = Vector3<V>;
176
177 fn mul(self, rhs: Vector2<U>) -> Self::Output {
178 self.dot_vector2(&rhs)
179 }
180}
181
182impl<T, U, V> Mul<Matrix3x2<U>> for Matrix3<T>
183where
184 T: PhysicsQuantity + Mul<U, Output = V>,
185 U: PhysicsQuantity,
186 V: PhysicsQuantity + Add<Output = V>,
187{
188 type Output = Matrix3x2<V>;
189
190 fn mul(self, rhs: Matrix3x2<U>) -> Self::Output {
191 let mut result = [[<V as num_traits::Zero>::zero(); 2]; 3];
192 for i in 0..3 {
193 for j in 0..2 {
194 result[i][j] = self[(i, 0)] * rhs[(0, j)]
195 + self[(i, 1)] * rhs[(1, j)]
196 + self[(i, 2)] * rhs[(2, j)];
197 }
198 }
199 Matrix3x2::new(result)
200 }
201}
202
203impl<T, U, V> Mul<Matrix2x3<U>> for Matrix3x2<T>
204where
205 T: PhysicsQuantity + Mul<U, Output = V>,
206 U: PhysicsQuantity,
207 V: PhysicsQuantity + Add<Output = V>,
208{
209 type Output = Matrix3<V>;
210
211 fn mul(self, rhs: Matrix2x3<U>) -> Self::Output {
212 let mut result = [[<V as num_traits::Zero>::zero(); 3]; 3];
213 for i in 0..3 {
214 for j in 0..3 {
215 result[i][j] = self[(i, 0)] * rhs[(0, j)] + self[(i, 1)] * rhs[(1, j)];
216 }
217 }
218 Matrix3::new(result)
219 }
220}
221
222impl<T> Matrix3x2<T>
223where
224 T: PhysicsQuantity + Copy,
225{
226 pub fn pseudoinverse<U, D, E, V>(&self) -> Option<Matrix2x3<V>>
227 where
228 T: Mul<T, Output = U>,
229 U: PhysicsQuantity + Copy + Mul<U, Output = D> + Div<D, Output = E>,
230
231 D: PhysicsQuantity + Copy + Sub<Output = D>,
232 E: PhysicsQuantity + Copy,
233
234 E: Mul<T, Output = V>,
235 V: PhysicsQuantity + Copy + Add<Output = V>,
236 {
237 let a_t: Matrix2x3<T> = self.transpose();
238 let a_ta: Matrix2<U> = a_t * *self;
239
240 let det: D = a_ta[(0, 0)] * a_ta[(1, 1)] - a_ta[(0, 1)] * a_ta[(1, 0)];
241 if det.as_f64() == 0.0 {
242 return None;
243 }
244
245 let inv_00: E = a_ta[(1, 1)] / det;
246 let inv_01: E = -a_ta[(0, 1)] / det;
247 let inv_10: E = -a_ta[(1, 0)] / det;
248 let inv_11: E = a_ta[(0, 0)] / det;
249 let inv = Matrix2::new([[inv_00, inv_01], [inv_10, inv_11]]);
250
251 let pinv: Matrix2x3<V> = inv * a_t;
252 Some(pinv)
253 }
254}
255
256impl<T, U, V> Mul<Matrix2<U>> for Matrix3x2<T>
257where
258 T: PhysicsQuantity + Mul<U, Output = V>,
259 U: PhysicsQuantity,
260 V: PhysicsQuantity + Add<Output = V>,
261{
262 type Output = Matrix3x2<V>;
263
264 fn mul(self, rhs: Matrix2<U>) -> Self::Output {
265 let mut result = [[V::zero(); 2]; 3];
266 for i in 0..3 {
267 for j in 0..2 {
268 result[i][j] = self[(i, 0)] * rhs[(0, j)] + self[(i, 1)] * rhs[(1, j)];
269 }
270 }
271 Matrix3x2::new(result)
272 }
273}