use learning::toolkit::kernel::{Kernel, SquaredExp};
use linalg::{Matrix, BaseMatrix};
use linalg::Vector;
use learning::{LearningResult, SupModel};
use learning::error::{Error, ErrorKind};
pub trait MeanFunc {
fn func(&self, x: Matrix<f64>) -> Vector<f64>;
}
#[derive(Clone, Copy, Debug)]
pub struct ConstMean {
a: f64,
}
impl Default for ConstMean {
fn default() -> ConstMean {
ConstMean { a: 0f64 }
}
}
impl MeanFunc for ConstMean {
fn func(&self, x: Matrix<f64>) -> Vector<f64> {
Vector::zeros(x.rows()) + self.a
}
}
#[derive(Debug)]
pub struct GaussianProcess<T: Kernel, U: MeanFunc> {
ker: T,
mean: U,
pub noise: f64,
alpha: Option<Vector<f64>>,
train_mat: Option<Matrix<f64>>,
train_data: Option<Matrix<f64>>,
}
impl Default for GaussianProcess<SquaredExp, ConstMean> {
fn default() -> GaussianProcess<SquaredExp, ConstMean> {
GaussianProcess {
ker: SquaredExp::default(),
mean: ConstMean::default(),
noise: 0f64,
train_mat: None,
train_data: None,
alpha: None,
}
}
}
impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
pub fn new(ker: T, mean: U, noise: f64) -> GaussianProcess<T, U> {
GaussianProcess {
ker: ker,
mean: mean,
noise: noise,
train_mat: None,
train_data: None,
alpha: None,
}
}
fn ker_mat(&self, m1: &Matrix<f64>, m2: &Matrix<f64>) -> LearningResult<Matrix<f64>> {
if m1.cols() != m2.cols() {
Err(Error::new(ErrorKind::InvalidState,
"Inputs to kernel matrices have different column counts."))
} else {
let dim1 = m1.rows();
let dim2 = m2.rows();
let mut ker_data = Vec::with_capacity(dim1 * dim2);
ker_data.extend(m1.iter_rows().flat_map(|row1| {
m2.iter_rows()
.map(move |row2| self.ker.kernel(row1, row2))
}));
Ok(Matrix::new(dim1, dim2, ker_data))
}
}
}
impl<T: Kernel, U: MeanFunc> SupModel<Matrix<f64>, Vector<f64>> for GaussianProcess<T, U> {
fn predict(&self, inputs: &Matrix<f64>) -> LearningResult<Vector<f64>> {
if let (&Some(ref alpha), &Some(ref t_data)) = (&self.alpha, &self.train_data) {
let mean = self.mean.func(inputs.clone());
let post_mean = try!(self.ker_mat(inputs, t_data)) * alpha;
Ok(mean + post_mean)
} else {
Err(Error::new(ErrorKind::UntrainedModel, "The model has not been trained."))
}
}
fn train(&mut self, inputs: &Matrix<f64>, targets: &Vector<f64>) -> LearningResult<()> {
let noise_mat = Matrix::identity(inputs.rows()) * self.noise;
let ker_mat = self.ker_mat(inputs, inputs).unwrap();
let train_mat = try!((ker_mat + noise_mat).cholesky().map_err(|_| {
Error::new(ErrorKind::InvalidState,
"Could not compute Cholesky decomposition.")
}));
let x = train_mat.solve_l_triangular(targets - self.mean.func(inputs.clone())).unwrap();
let alpha = train_mat.transpose().solve_u_triangular(x).unwrap();
self.train_mat = Some(train_mat);
self.train_data = Some(inputs.clone());
self.alpha = Some(alpha);
Ok(())
}
}
impl<T: Kernel, U: MeanFunc> GaussianProcess<T, U> {
pub fn get_posterior(&self,
inputs: &Matrix<f64>)
-> LearningResult<(Vector<f64>, Matrix<f64>)> {
if let (&Some(ref t_mat), &Some(ref alpha), &Some(ref t_data)) = (&self.train_mat,
&self.alpha,
&self.train_data) {
let mean = self.mean.func(inputs.clone());
let post_mean = mean + try!(self.ker_mat(inputs, t_data)) * alpha;
let test_mat = try!(self.ker_mat(inputs, t_data));
let mut var_data = Vec::with_capacity(inputs.rows() * inputs.cols());
for row in test_mat.iter_rows() {
let test_point = Vector::new(row.to_vec());
var_data.append(&mut t_mat.solve_l_triangular(test_point).unwrap().into_vec());
}
let v_mat = Matrix::new(test_mat.rows(), test_mat.cols(), var_data);
let post_var = try!(self.ker_mat(inputs, inputs)) - &v_mat * v_mat.transpose();
Ok((post_mean, post_var))
} else {
Err(Error::new_untrained())
}
}
}