Class based kernels for kriging (SMTorg#622)

* Begin writing hessian, unfinished, untested * Write derivatives, untested * Writes spatial derivatives, untested * Little modification * Solving bug in hessian in progress * Working on hessian * Fix hessian * Bug fixing for spatial derivative * Debugging of test prediction in progress * New file for extended kernel definition * Commit for debugging of spatial derivatives * Clean comments and print + ruff * Fix bug in order of test * Begin writing hessian, unfinished, untested * Write derivatives, untested * Writes spatial derivatives, untested * Little modification * Solving bug in hessian in progress * Working on hessian * Fix hessian * Bug fixing for spatial derivative * Debugging of test prediction in progress * New file for extended kernel definition * Commit for debugging of spatial derivatives * Clean comments and print + ruff * Fix bug in order of test * Remove kernel file * Typo * Readd kernel file * New test for squar_sin_exp new definition of a problem for squar_sin_exp, as the previous one was modeled by a constant line in one direction, creating a division by zero in the computation of the relative error. Also, a new computation for the error on the derivative of the variance is introduced, in order to account for when value can be so close to zero that the relative difference doesn't make much sense anymore. * First link between kernel and kriging * first complete lock * Interfacing of smt with kernels To do : finish cleaning, modify test_krg_training to test the actual kernels, remove kernels from utils.kriging * Bug fix in tests * Last commit before cleaning * Cleaning of comments * Apply ruff * Last removal of comments
MaelTremouille · Jul 10, 2024 · a7fb9fe · a7fb9fe
1 parent 1439cd7
commit a7fb9fe
Show file tree

Hide file tree

Showing 8 changed files with 598 additions and 842 deletions.
diff --git a/smt/applications/mfk.py b/smt/applications/mfk.py
@@ -514,9 +514,8 @@ def _predict_intermediate_values(self, X, lvl, descale=True):
         if self.is_continuous:
             dx = self._differences(X, Y=self.X_norma_all[0])
             d = self._componentwise_distance(dx)
-            r_ = self._correlation_types[self.options["corr"]](
-                self.optimal_theta[0], d
-            ).reshape(n_eval, self.nt_all[0])
+            self.corr.theta = self.optimal_theta[0]
+            r_ = self.corr(d).reshape(n_eval, self.nt_all[0])
         else:
             _, x_is_acting = self.design_space.correct_get_acting(X_usc)
             _, y_is_acting = self.design_space.correct_get_acting(self.X[0])
@@ -572,9 +571,8 @@ def _predict_intermediate_values(self, X, lvl, descale=True):
             if self.is_continuous:
                 dx = self._differences(X, Y=self.X_norma_all[i])
                 d = self._componentwise_distance(dx)
-                r_ = self._correlation_types[self.options["corr"]](
-                    self.optimal_theta[i], d
-                ).reshape(n_eval, self.nt_all[i])
+                self.corr.theta = self.optimal_theta[i]
+                r_ = self.corr(d).reshape(n_eval, self.nt_all[i])
             else:
                 _, x_is_acting = self.design_space.correct_get_acting(X_usc)
                 _, y_is_acting = self.design_space.correct_get_acting(self.X[i])
@@ -703,9 +701,8 @@ def predict_variances_all_levels(self, X, is_acting=None):
         if self.is_continuous:
             dx = self._differences(X, Y=self.X_norma_all[0])
             d = self._componentwise_distance(dx)
-            r_ = self._correlation_types[self.options["corr"]](
-                self.optimal_theta[0], d
-            ).reshape(n_eval, self.nt_all[0])
+            self.corr.theta = self.optimal_theta[0]
+            r_ = self.corr(d).reshape(n_eval, self.nt_all[0])
         else:
             _, y_is_acting = self.design_space.correct_get_acting(self.X[0])
             _, x_is_acting = self.design_space.correct_get_acting(X)
@@ -776,9 +773,8 @@ def predict_variances_all_levels(self, X, is_acting=None):
             if self.is_continuous:
                 dx = self._differences(X, Y=self.X_norma_all[i])
                 d = self._componentwise_distance(dx)
-                r_ = self._correlation_types[self.options["corr"]](
-                    self.optimal_theta[i], d
-                ).reshape(n_eval, self.nt_all[i])
+                self.corr.theta = self.optimal_theta[i]
+                r_ = self.corr(d).reshape(n_eval, self.nt_all[i])
             else:
                 _, y_is_acting = self.design_space.correct_get_acting(self.X[i])
                 _, x_is_acting = self.design_space.correct_get_acting(X)
@@ -916,9 +912,8 @@ def _predict_derivatives(self, x, kx):
         dx = self._differences(x, Y=self.X_norma_all[0])
         d = self._componentwise_distance(dx)
         # Compute the correlation function
-        r_ = self._correlation_types[self.options["corr"]](
-            self.optimal_theta[0], d
-        ).reshape(n_eval, self.nt_all[0])
+        self.corr.theta = self.optimal_theta[0]
+        r_ = self.corr(d).reshape(n_eval, self.nt_all[0])
 
         # Beta and gamma = R^-1(y-FBeta)
         beta = self.optimal_par[0]["beta"]
@@ -937,9 +932,8 @@ def _predict_derivatives(self, x, kx):
             g = self._regression_types[self.options["rho_regr"]](x)
             dx = self._differences(x, Y=self.X_norma_all[i])
             d = self._componentwise_distance(dx)
-            r_ = self._correlation_types[self.options["corr"]](
-                self.optimal_theta[i], d
-            ).reshape(n_eval, self.nt_all[i])
+            self.corr.theta = self.optimal_theta[i]
+            r_ = self.corr(d).reshape(n_eval, self.nt_all[i])
             df = np.vstack((g.T * dy_dx[:, i - 1], df0.T))
 
             beta = self.optimal_par[i]["beta"]

diff --git a/smt/surrogate_models/krg_based.py b/smt/surrogate_models/krg_based.py
@@ -14,6 +14,14 @@
 
 from smt.sampling_methods import LHS
 from smt.surrogate_models.surrogate_model import SurrogateModel
+
+from smt.utils.kernels import (
+    SquarSinExp,
+    PowExp,
+    Matern52,
+    Matern32,
+    ActExp,
+)
 from smt.utils.checks import check_support, ensure_2d_array
 from smt.utils.design_space import (
     BaseDesignSpace,
@@ -22,29 +30,23 @@
 )
 from smt.utils.kriging import (
     MixHrcKernelType,
-    abs_exp,
-    act_exp,
+    differences,
+    constant,
+    linear,
+    quadratic,
+    cross_distances,
+    gower_componentwise_distances,
     componentwise_distance,
     componentwise_distance_PLS,
     compute_X_cont,
     compute_X_cross,
-    constant,
-    cross_distances,
     cross_levels,
     cross_levels_homo_space,
-    differences,
-    gower_componentwise_distances,
-    linear,
-    matern32,
-    matern52,
     matrix_data_corr_levels_cat_matrix,
     matrix_data_corr_levels_cat_mod,
     matrix_data_corr_levels_cat_mod_comps,
-    pow_exp,
-    quadratic,
-    squar_exp,
-    squar_sin_exp,
 )
+
 from smt.utils.misc import standardization
 
 
@@ -59,16 +61,15 @@ class MixIntKernelType(Enum):
 class KrgBased(SurrogateModel):
     _regression_types = {"constant": constant, "linear": linear, "quadratic": quadratic}
 
-    _correlation_types = {
-        "pow_exp": pow_exp,
-        "abs_exp": abs_exp,
-        "squar_exp": squar_exp,
-        "squar_sin_exp": squar_sin_exp,
-        "act_exp": act_exp,
-        "matern52": matern52,
-        "matern32": matern32,
+    _correlation_class = {
+        "pow_exp": PowExp,
+        "abs_exp": PowExp,
+        "squar_exp": PowExp,
+        "squar_sin_exp": SquarSinExp,
+        "matern52": Matern52,
+        "matern32": Matern32,
+        "act_exp": ActExp,
     }
-
     name = "KrigingBased"
 
     def _initialize(self):
@@ -230,7 +231,9 @@ def _final_initialize(self):
             "matern52",
         ]:
             self.options["pow_exp_power"] = 1.0
-
+        self.corr = self._correlation_class[self.options["corr"]](
+            self.options["theta0"]
+        )
         # Check the pow_exp_power is >0 and <=2
         assert (
             self.options["pow_exp_power"] > 0 and self.options["pow_exp_power"] <= 2
@@ -667,16 +670,6 @@ def _matrix_data_corr(
             An array containing the values of the autocorrelation model.
         """
 
-        _correlation_types = {
-            "pow_exp": pow_exp,
-            "abs_exp": abs_exp,
-            "squar_exp": squar_exp,
-            "squar_sin_exp": squar_sin_exp,
-            "act_exp": act_exp,
-            "matern52": matern52,
-            "matern32": matern32,
-        }
-
         # Initialize static parameters
         (
             cat_kernel_comps,
@@ -721,7 +714,8 @@ def _matrix_data_corr(
                     theta=None,
                     return_derivative=False,
                 )
-                r = _correlation_types[corr](theta, d)
+                self.corr.theta = theta
+                r = self.corr(d)
                 return r
             else:
                 d_cont = componentwise_distance_PLS(
@@ -743,7 +737,8 @@ def _matrix_data_corr(
                 return_derivative=False,
             )
             if cat_kernel in [MixIntKernelType.GOWER, MixIntKernelType.CONT_RELAX]:
-                r = _correlation_types[corr](theta, d)
+                self.corr.theta = theta
+                r = self.corr(d)
                 return r
             else:
                 d_cont = d[:, np.logical_not(cat_features)]
@@ -761,7 +756,8 @@ def _matrix_data_corr(
                 )
 
         theta_cont = theta[theta_cont_features[:, 0]]
-        r_cont = _correlation_types[corr](theta_cont, d_cont)
+        self.corr.theta = theta_cont
+        r_cont = self.corr(d_cont)
         r_cat = np.copy(r_cont) * 0
         r = np.copy(r_cont)
         ##Theta_cat_i loop
@@ -980,9 +976,8 @@ def _reduced_likelihood_function(self, theta):
                 kplsk_second_loop=self.kplsk_second_loop,
             ).reshape(-1, 1)
         else:
-            r = self._correlation_types[self.options["corr"]](theta, self.D).reshape(
-                -1, 1
-            )
+            self.corr.theta = theta
+            r = self.corr(self.D).reshape(-1, 1)
         R = np.eye(self.nt) * (1.0 + nugget + noise)
         R[self.ij[:, 0], self.ij[:, 1]] = r[:, 0]
         R[self.ij[:, 1], self.ij[:, 0]] = r[:, 0]
@@ -1119,10 +1114,7 @@ def _reduced_likelihood_gradient(self, theta):
         dbeta_all = []
         for i_der in range(nb_theta):
             # Compute R derivatives
-            dr = self._correlation_types[self.options["corr"]](
-                theta, self.D, grad_ind=i_der
-            )
-
+            dr = self.corr(self.D, grad_ind=i_der)
             dr_all.append(dr)
 
             dR = np.zeros((self.nt, self.nt))
@@ -1264,10 +1256,7 @@ def _reduced_likelihood_hessian(self, theta):
                 dRdeta = np.zeros((self.nt, self.nt))
                 dRdeta[self.ij[:, 0], self.ij[:, 1]] = dr_all[eta][:, 0]
                 dRdeta[self.ij[:, 1], self.ij[:, 0]] = dr_all[eta][:, 0]
-
-                dr_eta_omega = self._correlation_types[self.options["corr"]](
-                    theta, self.D, grad_ind=omega, hess_ind=eta
-                )
+                dr_eta_omega = self.corr(self.D, grad_ind=omega, hess_ind=eta)
                 dRdetadomega = np.zeros((self.nt, self.nt))
                 dRdetadomega[self.ij[:, 0], self.ij[:, 1]] = dr_eta_omega[:, 0]
                 dRdetadomega[self.ij[:, 1], self.ij[:, 0]] = dr_eta_omega[:, 0]
@@ -1504,9 +1493,7 @@ def _predict_values(self, x: np.ndarray, is_acting=None) -> np.ndarray:
             X_cont = np.copy(x)
             d = self._componentwise_distance(dx)
             # Compute the correlation function
-            r = self._correlation_types[self.options["corr"]](
-                self.optimal_theta, d
-            ).reshape(n_eval, self.nt)
+            r = self.corr(d).reshape(n_eval, self.nt)
             y = np.zeros(n_eval)
         X_cont = (X_cont - self.X_offset) / self.X_scale
         # Compute the regression function
@@ -1550,10 +1537,7 @@ def _predict_derivatives(self, x, kx):
 
         # Compute the correlation function
         derivative_dic = {"dx": dx, "dd": dd}
-
-        r, dr = self._correlation_types[self.options["corr"]](
-            self.optimal_theta, d, derivative_params=derivative_dic
-        )
+        r, dr = self.corr(d, derivative_params=derivative_dic)
         r = r.reshape(n_eval, self.nt)
 
         drx = dr[:, kx].reshape(n_eval, self.nt)
@@ -1648,9 +1632,7 @@ def _predict_variances(self, x: np.ndarray, is_acting=None) -> np.ndarray:
             X_cont = np.copy(x)
             d = self._componentwise_distance(dx)
             # Compute the correlation function
-            r = self._correlation_types[self.options["corr"]](
-                self.optimal_theta, d
-            ).reshape(n_eval, self.nt)
+            r = self.corr(d).reshape(n_eval, self.nt)
         X_cont = (X_cont - self.X_offset) / self.X_scale
         C = self.optimal_par["C"]
         rt = linalg.solve_triangular(C, r.T, lower=True)
@@ -1714,7 +1696,6 @@ def _internal_predict_variance(self, x, kx=None):
         # Initialization
         n_eval, _ = x.shape
         x = (x - self.X_offset) / self.X_scale
-        theta = self.optimal_theta
         # Get pairwise componentwise L1-distances to the input training set
         dx = differences(x, Y=self.X_norma.copy())
         d = self._componentwise_distance(dx)
@@ -1730,9 +1711,7 @@ def _internal_predict_variance(self, x, kx=None):
         C = self.optimal_par["C"]
 
         # p1 : derivative of (rt**2.0).sum(axis=0)
-        r, dr = self._correlation_types[self.options["corr"]](
-            theta, d, derivative_params=derivative_dic
-        )
+        r, dr = self.corr(d, derivative_params=derivative_dic)
         if kx is None:
             rt = linalg.solve_triangular(C, r, lower=True)
             drx = dr.T
@@ -1885,6 +1864,7 @@ def hessian_minus_reduced_likelihood_function(log10t):
                 "KPLSK" in self.name and ii == 0
             ) or self.kplsk_second_loop
             self.theta0 = deepcopy(self.options["theta0"])
+            self.corr.theta = deepcopy(self.options["theta0"])
             for i in range(len(self.theta0)):
                 # In practice, in 1D and for X in [0,1], theta^{-2} in [1e-2,infty),
                 # i.e. theta in (0,1e1], is a good choice to avoid overfitting.

diff --git a/smt/surrogate_models/mgp.py b/smt/surrogate_models/mgp.py
@@ -117,9 +117,8 @@ def _predict_values(self, x, is_acting=None):
             d_x = None
 
         # Compute the correlation function
-        r = self._correlation_types[self.options["corr"]](theta, d, d_x=d_x).reshape(
-            n_eval, self.nt
-        )
+        self.corr.theta = theta
+        r = self.corr(d, d_x=d_x).reshape(n_eval, self.nt)
 
         f = self._regression_types[self.options["poly"]](x)
         # Scaled predictor
@@ -278,9 +277,8 @@ def _predict_value_derivatives_hyper(self, x, u=None):
             d_x = None
 
         # Compute the correlation function
-        r = self._correlation_types[self.options["corr"]](theta, d, d_x=d_x).reshape(
-            n_eval, self.nt
-        )
+        self.corr.theta = theta
+        r = self.corr(d, d_x=d_x).reshape(n_eval, self.nt)
         # Compute the regression function
         f = self._regression_types[self.options["poly"]](x)
 
@@ -291,9 +289,7 @@ def _predict_value_derivatives_hyper(self, x, u=None):
         Rinv_dmu = self.optimal_par["Rinv_dmu"]
 
         for omega in range(len(self.optimal_theta)):
-            drdomega = self._correlation_types[self.options["corr"]](
-                theta, d, grad_ind=omega, d_x=d_x
-            ).reshape(n_eval, self.nt)
+            drdomega = self.corr(d, grad_ind=omega, d_x=d_x).reshape(n_eval, self.nt)
 
             dbetadomega = self.optimal_par["dbeta_all"][omega]
 
@@ -346,11 +342,8 @@ def _predict_variance_derivatives_hyper(self, x, u=None):
             d_x = None
 
         # Compute the correlation function
-        r = (
-            self._correlation_types[self.options["corr"]](theta, d, d_x=d_x)
-            .reshape(n_eval, self.nt)
-            .T
-        )
+        self.corr.theta = theta
+        r = self.corr(d, d_x=d_x).reshape(n_eval, self.nt).T
         f = self._regression_types[self.options["poly"]](x).T
 
         C = self.optimal_par["C"]
@@ -382,13 +375,7 @@ def _predict_variance_derivatives_hyper(self, x, u=None):
         dsigma = self.optimal_par["dsigma"]
 
         for omega in range(len(self.optimal_theta)):
-            drdomega = (
-                self._correlation_types[self.options["corr"]](
-                    theta, d, grad_ind=omega, d_x=d_x
-                )
-                .reshape(n_eval, self.nt)
-                .T
-            )
+            drdomega = self.corr(d, grad_ind=omega, d_x=d_x).reshape(n_eval, self.nt).T
 
             dRdomega = np.zeros((self.nt, self.nt))
             dRdomega[self.ij[:, 0], self.ij[:, 1]] = dr_all[omega][:, 0]

diff --git a/smt/surrogate_models/sgp.py b/smt/surrogate_models/sgp.py
@@ -208,7 +208,8 @@ def _compute_K(self, A: np.ndarray, B: np.ndarray, theta, sigma2):
         dx = differences(A, B)
         d = self._componentwise_distance(dx)
         # Compute the correlation vector r and matrix R
-        r = self._correlation_types[self.options["corr"]](theta, d)
+        self.corr.theta = theta
+        r = self.corr(d)
         R = r.reshape(A.shape[0], B.shape[0])
         # Compute the covariance matrix K
         K = sigma2 * R