Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,397 @@
|
||||
//! Gaussian Process regression implementation
|
||||
//!
|
||||
//! Gaussian Processes provide a probabilistic approach to regression
|
||||
//! that can quantify uncertainty in predictions.
|
||||
|
||||
use crate::error::Result;
|
||||
use rtx_tensor::Tensor;
|
||||
use std::f32;
|
||||
|
||||
/// Kernel functions for Gaussian Process
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Kernel {
|
||||
/// Radial Basis Function (RBF) kernel: k(x, x') = exp(-γ||x - x'||²)
|
||||
RBF { gamma: f32 },
|
||||
/// Linear kernel: k(x, x') = x^T x'
|
||||
Linear,
|
||||
/// Polynomial kernel: k(x, x') = (γ x^T x' + r)^d
|
||||
Polynomial { gamma: f32, degree: f32, coef0: f32 },
|
||||
/// Matern kernel with ν = 3/2
|
||||
Matern32 { length_scale: f32 },
|
||||
}
|
||||
|
||||
/// Gaussian Process regressor
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GaussianProcess {
|
||||
/// Kernel function
|
||||
kernel: Kernel,
|
||||
/// Noise level (alpha)
|
||||
alpha: f32,
|
||||
/// Normalization of target values
|
||||
normalize_y: bool,
|
||||
/// Training data
|
||||
x_train: Option<Tensor>,
|
||||
/// Training targets
|
||||
y_train: Option<Tensor>,
|
||||
/// Cholesky decomposition of K + αI
|
||||
l_matrix: Option<Tensor>,
|
||||
/// α = L^(-1) y
|
||||
alpha_vec: Option<Tensor>,
|
||||
/// Mean of training targets (if normalized)
|
||||
y_mean: f32,
|
||||
/// Standard deviation of training targets (if normalized)
|
||||
y_std: f32,
|
||||
/// Number of training samples
|
||||
n_train: usize,
|
||||
/// Number of features
|
||||
n_features: usize,
|
||||
}
|
||||
|
||||
impl GaussianProcess {
|
||||
/// Create new Gaussian Process regressor
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
kernel: Kernel::RBF { gamma: 1.0 },
|
||||
alpha: 1e-10,
|
||||
normalize_y: false,
|
||||
x_train: None,
|
||||
y_train: None,
|
||||
l_matrix: None,
|
||||
alpha_vec: None,
|
||||
y_mean: 0.0,
|
||||
y_std: 1.0,
|
||||
n_train: 0,
|
||||
n_features: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set kernel function
|
||||
pub fn kernel(mut self, kernel: Kernel) -> Self {
|
||||
self.kernel = kernel;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set noise level (regularization parameter)
|
||||
pub fn alpha(mut self, alpha: f32) -> Self {
|
||||
self.alpha = alpha;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set whether to normalize target values
|
||||
pub fn normalize_y(mut self, normalize: bool) -> Self {
|
||||
self.normalize_y = normalize;
|
||||
self
|
||||
}
|
||||
|
||||
/// Fit the Gaussian Process model
|
||||
pub fn fit(&mut self, x: &Tensor, y: &Tensor) -> Result<&mut Self> {
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if y.shape().ndim() != 1 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"y must be 1-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
self.n_train = x.shape().dims()[0];
|
||||
self.n_features = x.shape().dims()[1];
|
||||
|
||||
if y.shape().dims()[0] != self.n_train {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X and y must have same number of samples".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Store training data
|
||||
self.x_train = Some(x.clone());
|
||||
|
||||
// Normalize targets if requested
|
||||
let y_normalized = if self.normalize_y {
|
||||
let y_data = y.to_cpu()?;
|
||||
self.y_mean = y_data.iter().sum::<f32>() / self.n_train as f32;
|
||||
self.y_std = {
|
||||
let variance = y_data
|
||||
.iter()
|
||||
.map(|&val| (val - self.y_mean).powi(2))
|
||||
.sum::<f32>()
|
||||
/ self.n_train as f32;
|
||||
variance.sqrt().max(1e-8) // Avoid division by zero
|
||||
};
|
||||
|
||||
let normalized_data: Vec<f32> = y_data
|
||||
.iter()
|
||||
.map(|&val| (val - self.y_mean) / self.y_std)
|
||||
.collect();
|
||||
|
||||
Tensor::from_data(normalized_data, vec![self.n_train], y.device())?
|
||||
} else {
|
||||
self.y_mean = 0.0;
|
||||
self.y_std = 1.0;
|
||||
y.clone()
|
||||
};
|
||||
|
||||
self.y_train = Some(y_normalized.clone());
|
||||
|
||||
// Compute kernel matrix K
|
||||
let k_matrix = self.compute_kernel_matrix(x, x)?;
|
||||
|
||||
// Add noise term: K + αI
|
||||
let mut k_data = k_matrix.to_cpu()?;
|
||||
for i in 0..self.n_train {
|
||||
k_data[i * self.n_train + i] += self.alpha;
|
||||
}
|
||||
|
||||
let k_noisy = Tensor::from_data(k_data, vec![self.n_train, self.n_train], x.device())?;
|
||||
|
||||
// Cholesky decomposition: K + αI = L L^T
|
||||
self.l_matrix = Some(self.cholesky_decomposition(&k_noisy)?);
|
||||
|
||||
// Solve L α = y for α
|
||||
let l_matrix = self.l_matrix.as_ref().unwrap();
|
||||
self.alpha_vec = Some(self.solve_triangular(l_matrix, &y_normalized, true)?);
|
||||
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// Predict mean and standard deviation for test data
|
||||
pub fn predict(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let x_train = self.x_train.as_ref().ok_or_else(|| {
|
||||
crate::error::MLError::not_fitted("Model has not been fitted yet".to_string())
|
||||
})?;
|
||||
|
||||
let _l_matrix = self.l_matrix.as_ref().unwrap();
|
||||
let alpha_vec = self.alpha_vec.as_ref().unwrap();
|
||||
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_test = x.shape().dims()[0];
|
||||
let n_features = x.shape().dims()[1];
|
||||
|
||||
if n_features != self.n_features {
|
||||
return Err(crate::error::MLError::invalid_input(format!(
|
||||
"X has {} features but model was fitted with {} features",
|
||||
n_features, self.n_features
|
||||
)));
|
||||
}
|
||||
|
||||
// Compute kernel matrix between test and training data
|
||||
let k_star = self.compute_kernel_matrix(x, x_train)?;
|
||||
|
||||
// Compute mean predictions: μ* = K* α
|
||||
let mu_star = k_star.matmul(alpha_vec)?;
|
||||
|
||||
// Denormalize predictions if needed
|
||||
let predictions = if self.normalize_y {
|
||||
let mu_data = mu_star.to_cpu()?;
|
||||
let denormalized: Vec<f32> = mu_data
|
||||
.iter()
|
||||
.map(|&val| val * self.y_std + self.y_mean)
|
||||
.collect();
|
||||
Tensor::from_data(denormalized, vec![n_test], x.device())?
|
||||
} else {
|
||||
mu_star
|
||||
};
|
||||
|
||||
Ok(predictions)
|
||||
}
|
||||
|
||||
/// Predict with uncertainty (returns mean and standard deviation)
|
||||
pub fn predict_with_uncertainty(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let x_train = self.x_train.as_ref().ok_or_else(|| {
|
||||
crate::error::MLError::not_fitted("Model has not been fitted yet".to_string())
|
||||
})?;
|
||||
|
||||
let l_matrix = self.l_matrix.as_ref().unwrap();
|
||||
let alpha_vec = self.alpha_vec.as_ref().unwrap();
|
||||
|
||||
// Validate input
|
||||
if x.shape().ndim() != 2 {
|
||||
return Err(crate::error::MLError::invalid_input(
|
||||
"X must be 2-dimensional".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let n_test = x.shape().dims()[0];
|
||||
|
||||
// Compute kernel matrices
|
||||
let k_star = self.compute_kernel_matrix(x, x_train)?;
|
||||
let k_star_star = self.compute_kernel_matrix(x, x)?;
|
||||
|
||||
// Compute mean predictions
|
||||
let mu_star = k_star.matmul(alpha_vec)?;
|
||||
|
||||
// Compute variance: σ²* = K** - K* L^(-T) L^(-1) K*^T
|
||||
let v = self.solve_triangular(l_matrix, &k_star.transpose(0, 1)?, true)?;
|
||||
let var_star = k_star_star.subtract(&v.transpose(0, 1)?.matmul(&v)?)?;
|
||||
|
||||
// Extract diagonal for pointwise variances
|
||||
let var_data = var_star.to_cpu()?;
|
||||
let std_data: Vec<f32> = (0..n_test)
|
||||
.map(|i| var_data[i * n_test + i].max(0.0).sqrt())
|
||||
.collect();
|
||||
|
||||
let std_tensor = Tensor::from_data(std_data, vec![n_test], x.device())?;
|
||||
|
||||
// Denormalize if needed
|
||||
let (mean_final, std_final) = if self.normalize_y {
|
||||
let mu_data = mu_star.to_cpu()?;
|
||||
let denormalized_mean: Vec<f32> = mu_data
|
||||
.iter()
|
||||
.map(|&val| val * self.y_std + self.y_mean)
|
||||
.collect();
|
||||
|
||||
let std_data = std_tensor.to_cpu()?;
|
||||
let denormalized_std: Vec<f32> = std_data.iter().map(|&val| val * self.y_std).collect();
|
||||
|
||||
(
|
||||
Tensor::from_data(denormalized_mean, vec![n_test], x.device())?,
|
||||
Tensor::from_data(denormalized_std, vec![n_test], x.device())?,
|
||||
)
|
||||
} else {
|
||||
(mu_star, std_tensor)
|
||||
};
|
||||
|
||||
Ok((mean_final, std_final))
|
||||
}
|
||||
|
||||
/// Compute kernel matrix between two sets of points
|
||||
fn compute_kernel_matrix(&self, x1: &Tensor, x2: &Tensor) -> Result<Tensor> {
|
||||
let x1_data = x1.to_cpu()?;
|
||||
let x2_data = x2.to_cpu()?;
|
||||
|
||||
let n1 = x1.shape().dims()[0];
|
||||
let n2 = x2.shape().dims()[0];
|
||||
let n_features = x1.shape().dims()[1];
|
||||
|
||||
let mut k_data = vec![0.0f32; n1 * n2];
|
||||
|
||||
for i in 0..n1 {
|
||||
for j in 0..n2 {
|
||||
let x1_point = &x1_data[i * n_features..(i + 1) * n_features];
|
||||
let x2_point = &x2_data[j * n_features..(j + 1) * n_features];
|
||||
|
||||
k_data[i * n2 + j] = self.kernel_function(x1_point, x2_point);
|
||||
}
|
||||
}
|
||||
|
||||
Tensor::from_data(k_data, vec![n1, n2], x1.device()).map_err(std::convert::Into::into)
|
||||
}
|
||||
|
||||
/// Compute kernel function between two points
|
||||
fn kernel_function(&self, x1: &[f32], x2: &[f32]) -> f32 {
|
||||
match &self.kernel {
|
||||
Kernel::RBF { gamma } => {
|
||||
let sq_dist: f32 = x1
|
||||
.iter()
|
||||
.zip(x2.iter())
|
||||
.map(|(&a, &b)| (a - b).powi(2))
|
||||
.sum();
|
||||
(-gamma * sq_dist).exp()
|
||||
}
|
||||
Kernel::Linear => x1.iter().zip(x2.iter()).map(|(&a, &b)| a * b).sum(),
|
||||
Kernel::Polynomial {
|
||||
gamma,
|
||||
degree,
|
||||
coef0,
|
||||
} => {
|
||||
let dot_product: f32 = x1.iter().zip(x2.iter()).map(|(&a, &b)| a * b).sum();
|
||||
(gamma * dot_product + coef0).powf(*degree)
|
||||
}
|
||||
Kernel::Matern32 { length_scale } => {
|
||||
let dist: f32 = x1
|
||||
.iter()
|
||||
.zip(x2.iter())
|
||||
.map(|(&a, &b)| (a - b).powi(2))
|
||||
.sum::<f32>()
|
||||
.sqrt();
|
||||
let sqrt3_dist_l = (3.0_f32).sqrt() * dist / length_scale;
|
||||
(1.0 + sqrt3_dist_l) * (-sqrt3_dist_l).exp()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cholesky decomposition
|
||||
fn cholesky_decomposition(&self, matrix: &Tensor) -> Result<Tensor> {
|
||||
// Use the tensor's built-in Cholesky solve method as a reference
|
||||
// For now, implement a simple decomposition
|
||||
let n = matrix.shape().dims()[0];
|
||||
let data = matrix.to_cpu()?;
|
||||
let mut l_data = vec![0.0f32; n * n];
|
||||
|
||||
for i in 0..n {
|
||||
for j in 0..=i {
|
||||
if i == j {
|
||||
let mut sum = 0.0;
|
||||
for k in 0..j {
|
||||
sum += l_data[j * n + k] * l_data[j * n + k];
|
||||
}
|
||||
l_data[j * n + j] = (data[j * n + j] - sum).sqrt().max(1e-12);
|
||||
} else {
|
||||
let mut sum = 0.0;
|
||||
for k in 0..j {
|
||||
sum += l_data[i * n + k] * l_data[j * n + k];
|
||||
}
|
||||
l_data[i * n + j] = (data[i * n + j] - sum) / l_data[j * n + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Tensor::from_data(l_data, vec![n, n], matrix.device()).map_err(std::convert::Into::into)
|
||||
}
|
||||
|
||||
/// Solve triangular system
|
||||
fn solve_triangular(&self, l: &Tensor, b: &Tensor, lower: bool) -> Result<Tensor> {
|
||||
let n = l.shape().dims()[0];
|
||||
let l_data = l.to_cpu()?;
|
||||
let b_data = b.to_cpu()?;
|
||||
let mut x_data = vec![0.0f32; n];
|
||||
|
||||
if lower {
|
||||
// Forward substitution: L x = b
|
||||
for i in 0..n {
|
||||
let mut sum = 0.0;
|
||||
for j in 0..i {
|
||||
sum += l_data[i * n + j] * x_data[j];
|
||||
}
|
||||
x_data[i] = (b_data[i] - sum) / l_data[i * n + i];
|
||||
}
|
||||
} else {
|
||||
// Backward substitution: U x = b
|
||||
for i in (0..n).rev() {
|
||||
let mut sum = 0.0;
|
||||
for j in (i + 1)..n {
|
||||
sum += l_data[i * n + j] * x_data[j];
|
||||
}
|
||||
x_data[i] = (b_data[i] - sum) / l_data[i * n + i];
|
||||
}
|
||||
}
|
||||
|
||||
Tensor::from_data(x_data, vec![n], l.device()).map_err(std::convert::Into::into)
|
||||
}
|
||||
|
||||
/// Get fitted parameters
|
||||
pub fn params(&self) -> (Kernel, f32) {
|
||||
(self.kernel.clone(), self.alpha)
|
||||
}
|
||||
|
||||
/// Check if model is fitted
|
||||
pub fn is_fitted(&self) -> bool {
|
||||
self.x_train.is_some() && self.y_train.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for GaussianProcess {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user