use inset_foundation::PRECISION_ERROR_TOLERANCE;
struct Vector {
offset: usize,
length: usize,
elements: Vec<f64>,
}
impl Vector {
fn new(size: usize) -> Vector {
Vector {
offset: 0,
length: size,
elements: vec![0.0; size],
}
}
fn from_vol(values: Vec<f64>, offset: usize, length: usize) -> Vector {
Vector {
offset,
length,
elements: values,
}
}
fn get(&self, i: usize) -> f64 {
self.elements[i + self.offset]
}
fn set(&mut self, i: usize, value: f64) {
self.elements[i + self.offset] = value;
}
fn dot(&self, a: &Vector) -> f64 {
let mut result = 0.0;
for i in 0..self.length {
result += self.get(i) * a.get(i);
}
result
}
fn norm(&self) -> f64 {
self.dot(self).sqrt()
}
}
struct Matrix {
columns: usize,
elements: Vec<f64>,
}
impl Matrix {
fn new(rows: usize, cols: usize) -> Matrix {
Matrix {
columns: cols,
elements: vec![0.0; rows * cols],
}
}
fn get(&self, row: usize, col: usize) -> f64 {
self.elements[row * self.columns + col]
}
fn set(&mut self, row: usize, col: usize, value: f64) {
self.elements[row * self.columns + col] = value;
}
fn get_row(&self, row: usize) -> Vector {
let start = row * self.columns;
let end = start + self.columns;
Vector::from_vol(self.elements[start..end].to_vec(), 0, self.columns)
}
}
pub struct PolynomialFit {
pub coefficients: Vec<f64>,
pub confidence: f64,
}
impl PolynomialFit {
pub fn new(degree: usize) -> PolynomialFit {
PolynomialFit {
coefficients: vec![0.0; degree + 1],
confidence: 0.0,
}
}
}
impl std::fmt::Display for PolynomialFit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let coefficient_string: Vec<String> = self
.coefficients
.iter()
.map(|c| format!("{c:.3}"))
.collect();
write!(
f,
"PolynomialFit({coefficient_string:?}, confidence: {:.3})",
self.confidence
)
}
}
pub struct LeastSquaresSolver {
pub x: Vec<f64>,
pub y: Vec<f64>,
pub w: Vec<f64>,
}
impl LeastSquaresSolver {
pub fn new(x: Vec<f64>, y: Vec<f64>, w: Vec<f64>) -> LeastSquaresSolver {
debug_assert_eq!(x.len(), y.len());
debug_assert_eq!(y.len(), w.len());
LeastSquaresSolver { x, y, w }
}
pub fn solve(&self, degree: usize) -> Option<PolynomialFit> {
if degree > self.x.len() {
return None;
}
let mut result = PolynomialFit::new(degree);
let m = self.x.len();
let n = degree + 1;
let mut a = Matrix::new(n, m);
for h in 0..m {
a.set(0, h, self.w[h]);
for i in 1..n {
a.set(i, h, a.get(i - 1, h) * self.x[h]);
}
}
let mut q = Matrix::new(n, m);
let mut r = Matrix::new(n, n);
for j in 0..n {
for h in 0..m {
q.set(j, h, a.get(j, h));
}
for i in 0..j {
let dot = q.get_row(j).dot(&q.get_row(i));
for h in 0..m {
q.set(j, h, q.get(j, h) - dot * q.get(i, h));
}
}
let norm = q.get_row(j).norm();
if norm < PRECISION_ERROR_TOLERANCE {
return None;
}
let inverse_norm = 1.0 / norm;
for h in 0..m {
q.set(j, h, q.get(j, h) * inverse_norm);
}
for i in 0..n {
r.set(
j,
i,
if i < j {
0.0
} else {
q.get_row(j).dot(&a.get_row(i))
},
);
}
}
let mut wy = Vector::new(m);
for h in 0..m {
wy.set(h, self.y[h] * self.w[h]);
}
for i in (0..n).rev() {
result.coefficients[i] = q.get_row(i).dot(&wy);
for j in ((i + 1)..n).rev() {
result.coefficients[i] -= r.get(i, j) * result.coefficients[j];
}
result.coefficients[i] /= r.get(i, i);
}
let y_mean = self.y[..m].iter().sum::<f64>() / m as f64;
let mut sum_squared_error = 0.0;
let mut sum_squared_total = 0.0;
for h in 0..m {
let mut term = 1.0;
let mut err = self.y[h] - result.coefficients[0];
for coefficient in &result.coefficients[1..n] {
term *= self.x[h];
err -= term * coefficient;
}
sum_squared_error += self.w[h] * self.w[h] * err * err;
let v = self.y[h] - y_mean;
sum_squared_total += self.w[h] * self.w[h] * v * v;
}
result.confidence = if sum_squared_total <= PRECISION_ERROR_TOLERANCE {
1.0
} else {
1.0 - (sum_squared_error / sum_squared_total)
};
Some(result)
}
}