Skip to main content

tensorlogic_train/hyperparameter/
kernel.rs

1//! Gaussian Process kernels for Bayesian Optimization.
2
3use scirs2_core::ndarray::{Array1, Array2};
4
5/// Gaussian Process kernel for Bayesian Optimization.
6#[derive(Debug, Clone, Copy)]
7pub enum GpKernel {
8    /// Radial Basis Function (RBF) / Squared Exponential kernel.
9    /// K(x, x') = sigma^2 * exp(-||x - x'||^2 / (2 * l^2))
10    Rbf {
11        /// Signal variance (output scale).
12        sigma: f64,
13        /// Length scale (input scale).
14        length_scale: f64,
15    },
16    /// Matern kernel with nu = 3/2.
17    /// K(x, x') = sigma^2 * (1 + sqrt(3) * r / l) * exp(-sqrt(3) * r / l)
18    Matern32 {
19        /// Signal variance.
20        sigma: f64,
21        /// Length scale.
22        length_scale: f64,
23    },
24}
25
26impl Default for GpKernel {
27    fn default() -> Self {
28        Self::Rbf {
29            sigma: 1.0,
30            length_scale: 1.0,
31        }
32    }
33}
34
35impl GpKernel {
36    /// Compute kernel matrix K(X, X').
37    pub(super) fn compute_kernel(&self, x1: &Array2<f64>, x2: &Array2<f64>) -> Array2<f64> {
38        let n1 = x1.nrows();
39        let n2 = x2.nrows();
40        let mut k = Array2::zeros((n1, n2));
41        for i in 0..n1 {
42            for j in 0..n2 {
43                let x1_row = x1.row(i);
44                let x2_row = x2.row(j);
45                let dist_sq = x1_row
46                    .iter()
47                    .zip(x2_row.iter())
48                    .map(|(a, b)| (a - b).powi(2))
49                    .sum::<f64>();
50                k[[i, j]] = match self {
51                    Self::Rbf {
52                        sigma,
53                        length_scale,
54                    } => sigma.powi(2) * (-dist_sq / (2.0 * length_scale.powi(2))).exp(),
55                    Self::Matern32 {
56                        sigma,
57                        length_scale,
58                    } => {
59                        let r = dist_sq.sqrt();
60                        let sqrt3_r_l = (3.0_f64).sqrt() * r / length_scale;
61                        sigma.powi(2) * (1.0 + sqrt3_r_l) * (-sqrt3_r_l).exp()
62                    }
63                };
64            }
65        }
66        k
67    }
68
69    /// Compute kernel vector k(X, x).
70    pub(super) fn compute_kernel_vector(
71        &self,
72        x_train: &Array2<f64>,
73        x_test: &Array1<f64>,
74    ) -> Array1<f64> {
75        let n = x_train.nrows();
76        let mut k = Array1::zeros(n);
77        for i in 0..n {
78            let x_train_row = x_train.row(i);
79            let dist_sq = x_train_row
80                .iter()
81                .zip(x_test.iter())
82                .map(|(a, b)| (a - b).powi(2))
83                .sum::<f64>();
84            k[i] = match self {
85                Self::Rbf {
86                    sigma,
87                    length_scale,
88                } => sigma.powi(2) * (-dist_sq / (2.0 * length_scale.powi(2))).exp(),
89                Self::Matern32 {
90                    sigma,
91                    length_scale,
92                } => {
93                    let r = dist_sq.sqrt();
94                    let sqrt3_r_l = (3.0_f64).sqrt() * r / length_scale;
95                    sigma.powi(2) * (1.0 + sqrt3_r_l) * (-sqrt3_r_l).exp()
96                }
97            };
98        }
99        k
100    }
101}