math_utils/
data.rs

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

math_utils/
data.rs
