use crate::math::array::Array;
use crate::math::matrix::Matrix;
use crate::types::Real;
pub trait CostFunction {
fn values(&self, x: &Array) -> Array;
fn value(&self, x: &Array) -> Real {
let v = self.values(x);
assert!(!v.is_empty(), "no residuals returned by the cost function");
(v.iter().map(|e| e * e).sum::<Real>() / v.size() as Real).sqrt()
}
fn gradient(&self, grad: &mut Array, x: &Array) {
let eps = self.finite_difference_epsilon();
assert!(
eps.is_finite() && eps > 0.0,
"finite-difference epsilon ({eps}) must be finite and positive"
);
assert_eq!(
grad.size(),
x.size(),
"gradient size ({}) must match x size ({})",
grad.size(),
x.size()
);
let mut xx = x.clone();
for i in 0..x.size() {
xx[i] += eps;
let fp = self.value(&xx);
xx[i] -= 2.0 * eps;
let fm = self.value(&xx);
grad[i] = 0.5 * (fp - fm) / eps;
xx[i] = x[i];
}
}
fn value_and_gradient(&self, grad: &mut Array, x: &Array) -> Real {
self.gradient(grad, x);
self.value(x)
}
fn jacobian(&self, jac: &mut Matrix, x: &Array) {
let eps = self.finite_difference_epsilon();
assert!(
eps.is_finite() && eps > 0.0,
"finite-difference epsilon ({eps}) must be finite and positive"
);
assert_eq!(
jac.columns(),
x.size(),
"jacobian columns ({}) must match x size ({})",
jac.columns(),
x.size()
);
if x.is_empty() {
let residuals = self.values(x);
assert_eq!(
jac.rows(),
residuals.size(),
"jacobian rows ({}) must match residual size ({})",
jac.rows(),
residuals.size()
);
return;
}
let mut xx = x.clone();
for i in 0..x.size() {
xx[i] += eps;
let fp = self.values(&xx);
xx[i] -= 2.0 * eps;
let fm = self.values(&xx);
assert_eq!(
fp.size(),
fm.size(),
"finite-difference residual size changed from {} to {}",
fp.size(),
fm.size()
);
assert_eq!(
jac.rows(),
fp.size(),
"jacobian rows ({}) must match residual size ({})",
jac.rows(),
fp.size()
);
for j in 0..fp.size() {
jac[(j, i)] = 0.5 * (fp[j] - fm[j]) / eps;
}
xx[i] = x[i];
}
}
fn values_and_jacobian(&self, jac: &mut Matrix, x: &Array) -> Array {
self.jacobian(jac, x);
self.values(x)
}
fn finite_difference_epsilon(&self) -> Real {
1e-8
}
}
#[cfg(test)]
mod tests {
use super::*;
struct Quadratic;
impl CostFunction for Quadratic {
fn values(&self, x: &Array) -> Array {
Array::from([x[0] * x[0] + x[1], 2.0 * x[1]])
}
}
#[test]
fn value_is_root_mean_square_of_values() {
let x = Array::from([2.0, 3.0]);
let expected = ((7.0 * 7.0 + 6.0 * 6.0) / 2.0_f64).sqrt();
assert!((Quadratic.value(&x) - expected).abs() < 1e-15);
}
#[test]
fn default_gradient_matches_analytic_derivative() {
let x = Array::from([2.0, 3.0]);
let mut grad = Array::with_size(2);
Quadratic.gradient(&mut grad, &x);
let f = Quadratic.value(&x);
let analytic_dx = 7.0 * 2.0 * x[0] / (2.0 * f);
let analytic_dy = (7.0 + 4.0 * x[1]) / (2.0 * f);
assert!((grad[0] - analytic_dx).abs() < 1e-6);
assert!((grad[1] - analytic_dy).abs() < 1e-6);
}
#[test]
fn default_jacobian_matches_analytic_derivatives() {
let x = Array::from([2.0, 3.0]);
let mut jac = Matrix::with_size(2, 2);
Quadratic.jacobian(&mut jac, &x);
assert!((jac[(0, 0)] - 4.0).abs() < 1e-6);
assert!((jac[(0, 1)] - 1.0).abs() < 1e-6);
assert!((jac[(1, 0)] - 0.0).abs() < 1e-6);
assert!((jac[(1, 1)] - 2.0).abs() < 1e-6);
}
#[test]
#[should_panic(expected = "no residuals returned by the cost function")]
fn default_value_rejects_an_empty_residual_vector() {
struct EmptyCost;
impl CostFunction for EmptyCost {
fn values(&self, _x: &Array) -> Array {
Array::new()
}
}
let _ = EmptyCost.value(&Array::from([1.0]));
}
#[test]
fn combined_evaluations_agree_with_separate_calls() {
let x = Array::from([2.0, 3.0]);
let mut grad = Array::with_size(2);
let v = Quadratic.value_and_gradient(&mut grad, &x);
assert_eq!(v, Quadratic.value(&x));
let mut jac = Matrix::with_size(2, 2);
let vals = Quadratic.values_and_jacobian(&mut jac, &x);
assert_eq!(vals, Quadratic.values(&x));
}
#[test]
#[should_panic(expected = "finite-difference epsilon")]
fn default_gradient_rejects_invalid_finite_difference_epsilon() {
struct BadStep;
impl CostFunction for BadStep {
fn values(&self, x: &Array) -> Array {
x.clone()
}
fn finite_difference_epsilon(&self) -> Real {
0.0
}
}
let mut grad = Array::with_size(1);
BadStep.gradient(&mut grad, &Array::from([1.0]));
}
#[test]
#[should_panic(expected = "gradient size")]
fn default_gradient_rejects_wrong_output_shape() {
let mut grad = Array::with_size(1);
Quadratic.gradient(&mut grad, &Array::from([2.0, 3.0]));
}
#[test]
#[should_panic(expected = "jacobian rows")]
fn default_jacobian_rejects_wrong_output_shape() {
let mut jac = Matrix::with_size(1, 2);
Quadratic.jacobian(&mut jac, &Array::from([2.0, 3.0]));
}
#[test]
#[should_panic(expected = "jacobian rows")]
fn default_jacobian_rejects_wrong_rows_with_zero_variables() {
struct ConstantResiduals;
impl CostFunction for ConstantResiduals {
fn values(&self, _x: &Array) -> Array {
Array::from([1.0, 2.0])
}
}
let mut jac = Matrix::with_size(1, 0);
ConstantResiduals.jacobian(&mut jac, &Array::new());
}
}