use crate::{
traits::{CostFunction, GenericCostFunction, GenericGradient, Gradient, LogDensity},
DVector, Float,
};
use std::convert::Infallible;
pub struct Rosenbrock {
pub n: usize,
}
impl CostFunction for Rosenbrock {
fn evaluate(&self, x: &DVector<Float>, _args: &()) -> Result<Float, Infallible> {
#[allow(clippy::suboptimal_flops)]
Ok((0..(self.n - 1))
.map(|i| 100.0 * (x[i + 1] - x[i].powi(2)).powi(2) + (1.0 - x[i]).powi(2))
.sum())
}
}
impl Gradient for Rosenbrock {}
impl GenericCostFunction for Rosenbrock {
type Input = DVector<Float>;
fn evaluate_generic(&self, x: &Self::Input, args: &()) -> Result<Float, Infallible> {
<Self as CostFunction>::evaluate(self, x, args)
}
}
impl GenericGradient for Rosenbrock {
fn gradient_generic(&self, x: &Self::Input, args: &()) -> Result<DVector<Float>, Infallible> {
<Self as Gradient>::gradient(self, x, args)
}
fn hessian_generic(
&self,
x: &Self::Input,
args: &(),
) -> Result<nalgebra::DMatrix<Float>, Infallible> {
<Self as Gradient>::hessian(self, x, args)
}
}
impl LogDensity for Rosenbrock {
fn log_density(&self, x: &DVector<Float>, _args: &()) -> Result<Float, Infallible> {
#[allow(clippy::suboptimal_flops)]
Ok(-Float::ln(
(0..(self.n - 1))
.map(|i| 100.0 * (x[i + 1] - x[i].powi(2)).powi(2) + (1.0 - x[i]).powi(2))
.sum::<Float>(),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_rosenbrock_evaluate_at_minimum() {
let f = Rosenbrock { n: 2 };
let x = DVector::from_vec(vec![1.0, 1.0]);
let val = f.evaluate(&x, &()).unwrap();
assert_eq!(val, 0.0); }
#[test]
fn test_rosenbrock_evaluate_known_point() {
let f = Rosenbrock { n: 2 };
let x = DVector::from_vec(vec![0.0, 0.0]);
let val = f.evaluate(&x, &()).unwrap();
assert_eq!(val, 1.0);
}
#[test]
fn test_rosenbrock_evaluate_three_dimensions() {
let f = Rosenbrock { n: 3 };
let x = DVector::from_vec(vec![1.0, 1.0, 1.0]);
let val = f.evaluate(&x, &()).unwrap();
assert_eq!(val, 0.0);
}
#[test]
fn test_rosenbrock_log_density_at_minimum() {
let f = Rosenbrock { n: 2 };
let x = DVector::from_vec(vec![1.0, 1.0]);
let val = f.log_density(&x, &()).unwrap();
assert!(val.is_infinite() && val.is_sign_positive());
}
#[test]
fn test_rosenbrock_log_density_known_point() {
let f = Rosenbrock { n: 2 };
let x = DVector::from_vec(vec![0.0, 0.0]);
let val = f.log_density(&x, &()).unwrap();
assert_eq!(val, 0.0);
}
#[test]
fn test_rosenbrock_log_density_increasing_cost_decreases_density() {
let f = Rosenbrock { n: 2 };
let x1 = DVector::from_vec(vec![0.0, 0.0]); let x2 = DVector::from_vec(vec![2.0, 2.0]); let ld1 = f.log_density(&x1, &()).unwrap();
let ld2 = f.log_density(&x2, &()).unwrap();
assert!(ld2 < ld1);
}
}