math_utils/
data.rs

1//! Data analysis routines
2
3use std::array;
4use crate::{approx, algebra, num};
5use crate::traits::*;
6use crate::types::*;
7
8/// Compute the covariance matrix of a list of 2D points
9pub fn covariance_2 <S> (points : &[Point2 <S>]) -> Matrix2 <S> where
10  S : Field + num::NumCast + MaybeSerDes
11{
12  let npoints      = S::from (points.len()).unwrap();
13  let mean         = points.iter().map (|p| p.0).sum::<Vector2 <_>>() / npoints;
14  let centered     = points.iter().map (|p| p.0 - mean);
15  let [xx, xy, yy] = centered.map (|p| [p.x * p.x, p.x * p.y, p.y * p.y ])
16    .fold ([S::zero(); 3], |acc, item| array::from_fn (|i| acc[i] + item[i]));
17  let div = S::one() / (npoints - S::one());
18  let cov_xx = xx * div;
19  let cov_xy = xy * div;
20  let cov_yx = cov_xy;
21  let cov_yy = yy * div;
22  Matrix2::from_col_arrays ([
23    [cov_xx, cov_xy],
24    [cov_yx, cov_yy]
25  ])
26}
27
28/// Compute the covariance matrix of a list of 3D points
29pub fn covariance_3 <S> (points : &[Point3 <S>]) -> Matrix3 <S> where
30  S : Field + num::NumCast + MaybeSerDes
31{
32  let npoints      = S::from (points.len()).unwrap();
33  let mean         = points.iter().map (|p| p.0).sum::<Vector3 <_>>() / npoints;
34  let centered     = points.iter().map (|p| p.0 - mean);
35  let [xx, xy, xz, yy, yz, zz] = centered.map (|p| [
36    p.x * p.x,
37    p.x * p.y,
38    p.x * p.z,
39    p.y * p.y,
40    p.y * p.z,
41    p.z * p.z
42  ]).fold ([S::zero(); 6], |acc, item| array::from_fn (|i| acc[i] + item[i]));
43  let div = S::one() / (npoints - S::one());
44  Matrix3::from_col_arrays ([
45    [xx, xy, xz],
46    [xy, yy, yz],
47    [xz, yz, zz]
48  ]) * div
49}
50
51/// 2D principal component analysis
52pub fn pca_2 <S> (points : &[Point2 <S>]) -> Rotation2 <S> where
53  S : Real + num::NumCast + approx::RelativeEq <Epsilon=S> + std::fmt::Debug
54    + MaybeSerDes,
55  Vector2 <S> : NormedVectorSpace <S>
56{
57  let covariance = covariance_2 (points);
58  // covariance matrix is symmetric, therefore eigenvectors should be orthogonal and
59  // thus point in the direction of the principal components
60  let (eigenvalues, mut eigenvectors) : (Vec<_>, Vec<_>) =
61    algebra::linear::eigensystem_2 (covariance).iter().copied().map (Option::unwrap)
62      .unzip();
63  let eigenvectors = {
64    // normalize eigenvectors and create array
65    eigenvectors.iter_mut().for_each (|v| *v = NormedVectorSpace::normalize (*v).into());
66    <[Vector2 <S>; 2]>::try_from (eigenvectors.as_slice()).unwrap()
67      .map (Vector2::into_array)
68  };
69  // the principal component with the most variance corresponds to the largest
70  // eigenvalue, we take this to be the first principal component
71  debug_assert!(eigenvalues[0] >= eigenvalues[1]);
72  let mut mat = Matrix2::from_col_arrays (eigenvectors);
73  // ensure determinant == 1
74  if mat.determinant().is_negative() {
75    mat.cols[0] *= -S::one();
76  }
77  Rotation2::new_approx (mat).unwrap()
78}
79
80/// 3D principal component analysis
81pub fn pca_3 <S> (points : &[Point3 <S>]) -> Rotation3 <S> where
82  S : Real + num::NumCast + approx::RelativeEq <Epsilon=S> + std::fmt::Debug
83    + MaybeSerDes,
84  Vector2 <S> : NormedVectorSpace <S>
85{
86  let covariance = covariance_3 (points);
87  // covariance matrix is symmetric, therefore eigenvectors should be orthogonal and
88  // thus point in the direction of the principal components
89  let (eigenvalues, mut eigenvectors) : (Vec<_>, Vec<_>) =
90    algebra::linear::eigensystem_3 (covariance).iter().flatten().copied().unzip();
91  let eigenvectors = {
92    // normalize eigenvectors and create array
93    eigenvectors.iter_mut().for_each (|v| *v = NormedVectorSpace::normalize (*v).into());
94    <[Vector3 <S>; 3]>::try_from (eigenvectors.as_slice()).unwrap()
95      .map (Vector3::into_array)
96  };
97  // the principal component with the most variance corresponds to the largest
98  // eigenvalue, we take this to be the first principal component
99  debug_assert!(eigenvalues[0] >= eigenvalues[1]);
100  debug_assert!(eigenvalues[1] >= eigenvalues[2]);
101  let mut mat = Matrix3::from_col_arrays (eigenvectors);
102  // ensure determinant == 1
103  if mat.determinant().is_negative() {
104    mat.cols[0] *= -S::one();
105  }
106  Rotation3::new_approx (mat).unwrap()
107}
108
109#[cfg(test)]
110mod tests {
111  #![expect(clippy::unreadable_literal)]
112
113  use super::*;
114  #[test]
115  fn covariance_2d() {
116    let points = [
117      [ 1.0, 2.0],
118      [ 1.0, 3.0],
119      [10.0, 2.0]
120    ].map (Point2::<f32>::from);
121    let covariance = covariance_2 (points.as_slice());
122    assert_eq!(covariance.into_col_arrays(),
123      [ [27.0, -1.5],
124        [-1.5, 0.3333333] ]);
125  }
126
127  #[test]
128  fn pca_2d() {
129    let points = [
130      [-1.0, -1.0],
131      [-1.0,  1.0],
132      [ 1.0, -1.0],
133      [ 1.0,  1.0]
134    ].map (Point2::<f32>::from);
135    let components = pca_2 (points.as_slice()).into_col_arrays();
136    assert_eq!(components[0], [1.0, 0.0]);
137    assert_eq!(components[1], [0.0, 1.0]);
138  }
139
140  #[test]
141  fn pca_3d() {
142    use std::f32::consts::{FRAC_1_SQRT_2, FRAC_PI_4};
143    let points = [
144      [ 1.0, -1.0, -1.0],
145      [ 1.0, -1.0,  1.0],
146      [ 1.0,  1.0, -1.0],
147      [ 1.0,  1.0,  1.0],
148      [-1.0, -1.0, -1.0],
149      [-1.0, -1.0,  1.0],
150      [-1.0,  1.0, -1.0],
151      [-1.0,  1.0,  1.0]
152    ].map (Point3::<f32>::from);
153    let components = pca_3 (points.as_slice()).into_col_arrays();
154    assert_eq!(components[0], [1.0, 0.0, 0.0]);
155    assert_eq!(components[1], [0.0, 1.0, 0.0]);
156    assert_eq!(components[2], [0.0, 0.0, 1.0]);
157
158    // this case returns axis-aligned components, this agrees when numpy output (except
159    // that numpy returns the identity matrix)
160    let rotation = Rotation3::from_angle_y (FRAC_PI_4.into());
161    let points = [
162      [ 1.0, -1.0, -1.0],
163      [ 1.0, -1.0,  1.0],
164      [ 1.0,  1.0, -1.0],
165      [ 1.0,  1.0,  1.0],
166      [-1.0, -1.0, -1.0],
167      [-1.0, -1.0,  1.0],
168      [-1.0,  1.0, -1.0],
169      [-1.0,  1.0,  1.0]
170    ].map (Point3::<f32>::from).map (|p| rotation.rotate (p));
171    let components = pca_3 (points.as_slice()).into_col_arrays();
172    assert_eq!(components[0], [0.0, -1.0, 0.0]);
173    assert_eq!(components[1], [1.0,  0.0, 0.0]);
174    assert_eq!(components[2], [0.0,  0.0, 1.0]);
175
176    let rotation = Rotation3::from_angle_y (FRAC_PI_4.into());
177    let points = [
178      [ 1.0, -1.0, -4.0],
179      [ 1.0, -1.0,  4.0],
180      [ 1.0,  1.0, -4.0],
181      [ 1.0,  1.0,  4.0],
182      [-1.0, -1.0, -4.0],
183      [-1.0, -1.0,  4.0],
184      [-1.0,  1.0, -4.0],
185      [-1.0,  1.0,  4.0]
186    ].map (Point3::<f32>::from).map (|p| rotation.rotate (p));
187    let components = pca_3 (points.as_slice());
188    approx::assert_relative_eq!(components.mat(), Matrix3::from_col_arrays ([
189      [FRAC_1_SQRT_2, 0.0, FRAC_1_SQRT_2],
190      [0.0, 1.0, 0.0],
191      [-FRAC_1_SQRT_2, 0.0, FRAC_1_SQRT_2]
192    ]));
193  }
194}
math_utils/data.rs

math_utils/
data.rs
