1use serde::{Deserialize, Serialize};
2use std::ops::{Add, AddAssign, Div, Mul, Neg, Sub, SubAssign};
3
4use crate::numerical::{linear_algebra, scalar};
5
6#[derive(Clone, Copy, Debug, Default, PartialEq, Deserialize, Serialize)]
7pub struct Vec3 {
9 pub x: f64,
11 pub y: f64,
13 pub z: f64,
15}
16
17impl Vec3 {
18 pub const ZERO: Self = Self::new(0.0, 0.0, 0.0);
20 pub const X: Self = Self::new(1.0, 0.0, 0.0);
22 pub const Y: Self = Self::new(0.0, 1.0, 0.0);
24 pub const Z: Self = Self::new(0.0, 0.0, 1.0);
26
27 pub const fn new(x: f64, y: f64, z: f64) -> Self {
29 Self { x, y, z }
30 }
31 pub fn dot(self, rhs: Self) -> f64 {
33 self.x * rhs.x + self.y * rhs.y + self.z * rhs.z
34 }
35 pub fn cross(self, rhs: Self) -> Self {
37 Self::new(
38 self.y * rhs.z - self.z * rhs.y,
39 self.z * rhs.x - self.x * rhs.z,
40 self.x * rhs.y - self.y * rhs.x,
41 )
42 }
43 pub fn length_squared(self) -> f64 {
45 self.dot(self)
46 }
47 pub fn length(self) -> f64 {
49 self.length_squared().sqrt()
50 }
51 pub fn normalized(self) -> Option<Self> {
54 let length = self.length();
55 (length.is_finite() && length > scalar::MIN_NORMALIZABLE_NORM).then(|| self / length)
56 }
57 pub fn distance(self, rhs: Self) -> f64 {
59 (self - rhs).length()
60 }
61 pub fn is_finite(self) -> bool {
63 self.x.is_finite() && self.y.is_finite() && self.z.is_finite()
64 }
65 pub fn component(self, index: usize) -> f64 {
71 [self.x, self.y, self.z][index]
72 }
73 pub fn canonicalized(self) -> Self {
76 let values = [self.x, self.y, self.z];
77 let index = (0..3)
78 .max_by(|&a, &b| values[a].abs().total_cmp(&values[b].abs()))
79 .unwrap();
80 if values[index] < 0.0 {
81 -self
82 } else {
83 self
84 }
85 }
86 pub fn orthonormal_basis(self) -> Option<(Self, Self)> {
90 let n = self.normalized()?;
91 let seed = if n.x.abs() <= n.y.abs() && n.x.abs() <= n.z.abs() {
92 Self::X
93 } else if n.y.abs() <= n.z.abs() {
94 Self::Y
95 } else {
96 Self::Z
97 };
98 let u = n.cross(seed).normalized()?;
99 Some((u, n.cross(u)))
100 }
101}
102
103impl Add for Vec3 {
104 type Output = Self;
105 fn add(self, r: Self) -> Self {
106 Self::new(self.x + r.x, self.y + r.y, self.z + r.z)
107 }
108}
109impl Sub for Vec3 {
110 type Output = Self;
111 fn sub(self, r: Self) -> Self {
112 Self::new(self.x - r.x, self.y - r.y, self.z - r.z)
113 }
114}
115impl Mul<f64> for Vec3 {
116 type Output = Self;
117 fn mul(self, s: f64) -> Self {
118 Self::new(self.x * s, self.y * s, self.z * s)
119 }
120}
121impl Div<f64> for Vec3 {
122 type Output = Self;
123 fn div(self, s: f64) -> Self {
124 self * (1.0 / s)
125 }
126}
127impl Neg for Vec3 {
128 type Output = Self;
129 fn neg(self) -> Self {
130 self * -1.0
131 }
132}
133impl AddAssign for Vec3 {
134 fn add_assign(&mut self, r: Self) {
135 *self = *self + r;
136 }
137}
138impl SubAssign for Vec3 {
139 fn sub_assign(&mut self, r: Self) {
140 *self = *self - r;
141 }
142}
143
144pub(crate) fn outer_accumulate(matrix: &mut [[f64; 3]; 3], value: Vec3, weight: f64) {
145 let v = [value.x, value.y, value.z];
146 for i in 0..3 {
147 for j in 0..3 {
148 matrix[i][j] += weight * v[i] * v[j];
149 }
150 }
151}
152
153pub(crate) fn eigen_symmetric3(mut a: [[f64; 3]; 3]) -> ([f64; 3], [Vec3; 3]) {
155 let mut v = [[0.0; 3]; 3];
156 for (i, row) in v.iter_mut().enumerate() {
157 row[i] = 1.0;
158 }
159 let scale = a
160 .iter()
161 .flatten()
162 .fold(0.0_f64, |m, x| m.max(x.abs()))
163 .max(linear_algebra::MATRIX_SCALE_FLOOR);
164 for _ in 0..64 {
165 let mut pair = (0, 1);
166 for candidate in [(0, 2), (1, 2)] {
167 if a[candidate.0][candidate.1].abs() > a[pair.0][pair.1].abs() {
168 pair = candidate;
169 }
170 }
171 let (p, q) = pair;
172 if a[p][q].abs() <= linear_algebra::JACOBI_RELATIVE_CONVERGENCE * scale {
173 break;
174 }
175 let tau = (a[q][q] - a[p][p]) / (2.0 * a[p][q]);
176 let t = tau.signum() / (tau.abs() + (1.0 + tau * tau).sqrt());
177 let c = 1.0 / (1.0 + t * t).sqrt();
178 let s = t * c;
179 let apq = a[p][q];
180 a[p][p] -= t * apq;
181 a[q][q] += t * apq;
182 a[p][q] = 0.0;
183 a[q][p] = 0.0;
184 for k in 0..3 {
185 if k != p && k != q {
186 let (akp, akq) = (a[k][p], a[k][q]);
187 a[k][p] = c * akp - s * akq;
188 a[p][k] = a[k][p];
189 a[k][q] = s * akp + c * akq;
190 a[q][k] = a[k][q];
191 }
192 let (vkp, vkq) = (v[k][p], v[k][q]);
193 v[k][p] = c * vkp - s * vkq;
194 v[k][q] = s * vkp + c * vkq;
195 }
196 }
197 let mut order = [0, 1, 2];
198 order.sort_by(|&i, &j| a[i][i].total_cmp(&a[j][j]));
199 let values = order.map(|i| a[i][i]);
200 let vectors = order.map(|i| {
201 Vec3::new(v[0][i], v[1][i], v[2][i])
202 .normalized()
203 .unwrap_or(Vec3::ZERO)
204 });
205 (values, vectors)
206}
207
208pub(crate) fn solve_linear(mut a: Vec<Vec<f64>>, mut b: Vec<f64>) -> Option<Vec<f64>> {
209 let n = b.len();
210 if a.len() != n || a.iter().any(|r| r.len() != n) {
211 return None;
212 }
213 let scale = a
214 .iter()
215 .flatten()
216 .fold(0.0_f64, |m, x| m.max(x.abs()))
217 .max(1.0);
218 for col in 0..n {
219 let pivot = (col..n).max_by(|&i, &j| a[i][col].abs().total_cmp(&a[j][col].abs()))?;
220 if a[pivot][col].abs() <= linear_algebra::LINEAR_PIVOT_RELATIVE_MIN * scale {
221 return None;
222 }
223 a.swap(col, pivot);
224 b.swap(col, pivot);
225 let d = a[col][col];
226 for j in col..n {
227 a[col][j] /= d;
228 }
229 b[col] /= d;
230 for i in 0..n {
231 if i != col {
232 let f = a[i][col];
233 for j in col..n {
234 a[i][j] -= f * a[col][j];
235 }
236 b[i] -= f * b[col];
237 }
238 }
239 }
240 b.iter().all(|x| x.is_finite()).then_some(b)
241}
242
243pub(crate) fn least_squares(
244 rows: &[Vec<f64>],
245 rhs: &[f64],
246 regularization: f64,
247) -> Option<Vec<f64>> {
248 let n = rows.first()?.len();
249 if rows.len() != rhs.len() || rows.iter().any(|r| r.len() != n) {
250 return None;
251 }
252 let mut ata = vec![vec![0.0; n]; n];
253 let mut atb = vec![0.0; n];
254 for (row, &value) in rows.iter().zip(rhs) {
255 for i in 0..n {
256 atb[i] += row[i] * value;
257 for j in 0..n {
258 ata[i][j] += row[i] * row[j];
259 }
260 }
261 }
262 for (i, row) in ata.iter_mut().enumerate() {
263 row[i] += regularization;
264 }
265 solve_linear(ata, atb)
266}
267
268pub(crate) fn least_squares_qr(
275 rows: &[Vec<f64>],
276 rhs: &[f64],
277 relative_rank_min: f64,
278) -> Option<Vec<f64>> {
279 let parameter_count = rows.first()?.len();
280 if parameter_count == 0
281 || rows.len() < parameter_count
282 || rows.len() != rhs.len()
283 || !relative_rank_min.is_finite()
284 || relative_rank_min < 0.0
285 || rows
286 .iter()
287 .any(|row| row.len() != parameter_count || row.iter().any(|value| !value.is_finite()))
288 || rhs.iter().any(|value| !value.is_finite())
289 {
290 return None;
291 }
292
293 let mut columns: Vec<Vec<f64>> = (0..parameter_count)
294 .map(|column| rows.iter().map(|row| row[column]).collect())
295 .collect();
296 let mut residual_rhs = rhs.to_vec();
297 let mut permutation: Vec<usize> = (0..parameter_count).collect();
298 let mut upper = vec![vec![0.0; parameter_count]; parameter_count];
299 let mut projected_rhs = vec![0.0; parameter_count];
300 let squared_norm = |values: &[f64]| values.iter().map(|value| value * value).sum::<f64>();
301 let reference_norm = columns
302 .iter()
303 .map(|column| squared_norm(column).sqrt())
304 .fold(0.0_f64, f64::max);
305 if !reference_norm.is_finite() || reference_norm <= 0.0 {
306 return None;
307 }
308
309 for rank in 0..parameter_count {
310 let pivot = (rank..parameter_count).max_by(|&left, &right| {
311 squared_norm(&columns[left]).total_cmp(&squared_norm(&columns[right]))
312 })?;
313 if pivot != rank {
314 columns.swap(rank, pivot);
315 permutation.swap(rank, pivot);
316 for row in upper.iter_mut().take(rank) {
317 row.swap(rank, pivot);
318 }
319 }
320 let norm = squared_norm(&columns[rank]).sqrt();
321 if !norm.is_finite() || norm <= relative_rank_min * reference_norm {
322 return None;
323 }
324 upper[rank][rank] = norm;
325 let direction: Vec<_> = columns[rank].iter().map(|value| value / norm).collect();
326 projected_rhs[rank] = direction
327 .iter()
328 .zip(&residual_rhs)
329 .map(|(left, right)| left * right)
330 .sum();
331 for (value, direction_value) in residual_rhs.iter_mut().zip(&direction) {
332 *value -= direction_value * projected_rhs[rank];
333 }
334 for (column, values) in columns.iter_mut().enumerate().skip(rank + 1) {
335 let mut projection = direction
338 .iter()
339 .zip(&*values)
340 .map(|(left, right)| left * right)
341 .sum::<f64>();
342 for (value, direction_value) in values.iter_mut().zip(&direction) {
343 *value -= direction_value * projection;
344 }
345 let correction = direction
346 .iter()
347 .zip(&*values)
348 .map(|(left, right)| left * right)
349 .sum::<f64>();
350 for (value, direction_value) in values.iter_mut().zip(&direction) {
351 *value -= direction_value * correction;
352 }
353 projection += correction;
354 upper[rank][column] = projection;
355 }
356 }
357
358 let mut pivoted_solution = vec![0.0; parameter_count];
359 for row in (0..parameter_count).rev() {
360 let remainder = (row + 1..parameter_count)
361 .map(|column| upper[row][column] * pivoted_solution[column])
362 .sum::<f64>();
363 pivoted_solution[row] = (projected_rhs[row] - remainder) / upper[row][row];
364 }
365 if pivoted_solution.iter().any(|value| !value.is_finite()) {
366 return None;
367 }
368 let mut solution = vec![0.0; parameter_count];
369 for (pivoted, original) in permutation.into_iter().enumerate() {
370 solution[original] = pivoted_solution[pivoted];
371 }
372 Some(solution)
373}
374
375