Coverage for jetgp/full_degp/degp.py: 94%
151 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-05-07 14:04 -0500
1import numpy as np
2from numpy.linalg import cholesky, solve
3from scipy.linalg import cho_solve, cho_factor, solve_triangular
4from jetgp.full_degp import degp_utils
5import jetgp.utils as utils
6from jetgp.kernel_funcs.kernel_funcs import KernelFactory, get_oti_module
7from jetgp.full_degp.optimizer import Optimizer
10class degp:
11 """
12 Derivative-Enhanced Gaussian Process (DEGP) model.
14 Supports coordinate-aligned partial derivatives, hypercomplex representation,
15 and automatic normalization. Includes methods for training, prediction,
16 and uncertainty quantification using kernel methods.
18 Parameters
19 ----------
20 x_train : ndarray
21 Training input data of shape (n_samples, n_features).
22 y_train : list or ndarray
23 Training targets or list of partial derivatives.
24 n_order : int
25 Maximum derivative order.
26 n_bases : int
27 Number of input dimensions.
28 der_indices : list of lists
29 Derivative multi-indices corresponding to each derivative term.
30 derivative_locations : list of lists
31 Which training points have which derivatives.
32 normalize : bool, default=True
33 Whether to normalize inputs and outputs.
34 sigma_data : float or array-like, optional
35 Observation noise standard deviation or diagonal noise values.
36 kernel : str, default='SE'
37 Kernel type ('SE', 'RQ', 'Matern', 'SI', etc.).
38 kernel_type : str, default='anisotropic'
39 Kernel anisotropy ('anisotropic' or 'isotropic').
40 smoothness_parameter : float, optional
41 Smoothness parameter for Matern kernel.
42 """
44 def __init__(
45 self,
46 x_train,
47 y_train,
48 n_order,
49 n_bases,
50 der_indices,
51 derivative_locations=None,
52 normalize=True,
53 sigma_data=None,
54 kernel="SE",
55 kernel_type="anisotropic",
56 smoothness_parameter=None
57 ):
58 if n_order > 0 and derivative_locations is None:
59 import warnings
60 # Count total number of derivative components across all orders
61 n_derivs = sum(len(order_derivs) for order_derivs in der_indices)
62 n_train = len(x_train)
63 derivative_locations = [[i for i in range(n_train)] for _ in range(n_derivs)]
64 warnings.warn(
65 f"derivative_locations not provided. Assuming all {n_derivs} derivative(s) "
66 f"are available at all {n_train} training point(s).",
67 UserWarning
68 )
70 elif der_indices is None and n_order == 0:
71 der_indices = []
72 derivative_locations = []
73 self.n_order = n_order
74 self.n_bases = n_bases
75 self.dim = x_train.shape[1]
76 self.num_points = x_train.shape[0]
77 self.kernel = kernel
78 self.kernel_type = kernel_type
79 self.der_indices = der_indices
80 self.normalize = normalize
81 self.derivative_locations = derivative_locations
82 self.oti = get_oti_module(n_bases, n_order)
83 self.y_train_input = y_train
84 self.x_train_input = x_train
86 # Prepare indices and powers
87 self.flattened_der_indices = utils.flatten_der_indices(der_indices)
88 self.powers = utils.build_companion_array(n_bases, n_order, der_indices)
90 # Normalize data if required
91 if normalize:
92 (
93 self.y_train,
94 self.mu_y,
95 self.sigma_y,
96 self.sigmas_x,
97 self.mus_x,
98 sigma_data,
99 ) = utils.normalize_y_data(
100 x_train, y_train, sigma_data, self.flattened_der_indices
101 )
102 self.x_train = utils.normalize_x_data_train(x_train)
103 else:
104 self.x_train = x_train
105 self.y_train = utils.reshape_y_train(y_train)
107 # Compute differences for the kernel
108 # if kernel == 'SI':
109 # self.differences_by_dim = degp_utils.differences_by_dim_func_SI(
110 # self.x_train, self.x_train, n_order
111 # )
112 # else:
113 self.differences_by_dim = degp_utils.differences_by_dim_func(
114 self.x_train, self.x_train, n_order, self.oti
115 )
117 # Initialize noise matrix
118 self.sigma_data = (
119 np.zeros((self.y_train.shape[0], self.y_train.shape[0]))
120 if sigma_data is None
121 else np.diag(sigma_data)
122 )
123 self.sigma_data_sq_diag = (
124 np.zeros(self.y_train.shape[0])
125 if sigma_data is None
126 else np.asarray(sigma_data) ** 2
127 )
129 # Initialize kernel factory and optimizer
130 self.kernel_factory = KernelFactory(
131 dim=n_bases,
132 normalize=normalize,
133 differences_by_dim=self.differences_by_dim,
134 n_order=n_order,
135 smoothness_parameter=smoothness_parameter,
136 oti_module=self.oti
137 )
138 self.kernel_func = self.kernel_factory.create_kernel(
139 kernel_name=self.kernel, kernel_type=self.kernel_type
140 )
141 self.bounds = self.kernel_factory.bounds
142 self.optimizer = Optimizer(self)
144 def optimize_hyperparameters(self, *args, **kwargs):
145 """
146 Optimize model hyperparameters using the optimizer.
147 Returns optimized hyperparameter vector.
148 """
149 self.params = self.optimizer.optimize_hyperparameters(*args, **kwargs)
150 return self.params
152 def predict(self, X_test, params, calc_cov=False, return_deriv=False,
153 derivs_to_predict=None, return_full_cov=False):
154 """
155 Compute posterior predictive mean and (optionally) covariance at X_test.
157 Parameters
158 ----------
159 X_test : ndarray
160 Test input points of shape (n_test, n_features).
161 params : ndarray
162 Log-scaled kernel hyperparameters.
163 calc_cov : bool, default=False
164 Whether to compute predictive variance.
165 return_deriv : bool, default=False
166 Whether to return derivative predictions.
167 derivs_to_predict : list, optional
168 Specific derivatives to predict. Can include derivatives not present in the
169 training set — the cross-covariance K_* is constructed from kernel derivatives
170 and does not require the requested derivative to have been observed during
171 training. Each entry must be a valid derivative spec within n_bases and n_order
172 (e.g. ``[[3, 1]]`` for df/dx3 in a first-order model).
173 If None, defaults to all derivatives used in training.
175 Returns
176 -------
177 f_mean : ndarray
178 Predictive mean vector.
179 f_var : ndarray, optional
180 Predictive variance vector (only if calc_cov=True).
181 f_cov_full : ndarray, optional
182 Full predictive covariance matrix (only if calc_cov=True and
183 return_full_cov=True).
184 """
185 length_scales = params[:-1]
186 sigma_n = params[-1]
188 # Set up derivative prediction configuration
189 if return_deriv:
190 if derivs_to_predict is not None:
191 common_derivs = derivs_to_predict
192 else:
193 common_derivs = self.flattened_der_indices
195 # Determine prediction order from requested derivatives
196 required_order = max(
197 sum(pair[1] for pair in deriv_spec)
198 for deriv_spec in common_derivs
199 )
200 predict_order = max(required_order, self.n_order)
202 if predict_order > self.n_order:
203 predict_oti = get_oti_module(self.n_bases, predict_order)
204 smoothness_param = getattr(self.kernel_factory, 'alpha', None)
205 predict_kernel_factory = KernelFactory(
206 dim=self.n_bases,
207 normalize=self.normalize,
208 differences_by_dim=self.differences_by_dim,
209 n_order=predict_order,
210 smoothness_parameter=smoothness_param,
211 oti_module=predict_oti
212 )
213 predict_kernel_func = predict_kernel_factory.create_kernel(
214 kernel_name=self.kernel, kernel_type=self.kernel_type
215 )
216 else:
217 predict_oti = self.oti
218 predict_kernel_func = self.kernel_func
220 self.powers_predict = utils.build_companion_array_predict(
221 self.n_bases, predict_order, common_derivs)
222 else:
223 common_derivs = []
224 self.powers_predict = None
225 predict_order = self.n_order
226 predict_oti = self.oti
227 predict_kernel_func = self.kernel_func
229 # Reuse cached Cholesky + alpha from optimizer if available
230 _cache_hit = (
231 hasattr(self, '_cached_params')
232 and self._cached_params is not None
233 and np.array_equal(self._cached_params, params)
234 )
236 if _cache_hit:
237 L = self._cached_L
238 low = self._cached_low
239 alpha = self._cached_alpha
240 cho_solve_failed = False
241 else:
242 # Build training kernel matrix (no cache available)
243 phi_train = self.kernel_func(self.differences_by_dim, length_scales)
245 if self.n_order > 0:
246 phi_exp_train = phi_train.get_all_derivs(self.n_bases, 2 * self.n_order)
247 else:
248 phi_exp_train = phi_train.real
249 phi_exp_train = phi_exp_train[np.newaxis, :, :]
251 K = degp_utils.rbf_kernel(
252 phi_train, phi_exp_train, self.n_order, self.n_bases,
253 self.flattened_der_indices, self.powers,
254 index=self.derivative_locations
255 )
256 K.flat[::K.shape[0] + 1] += (10 ** sigma_n) ** 2
257 K += self.sigma_data ** 2
259 # Solve linear system
260 try:
261 cho_solve_failed = False
262 L, low = cho_factor(K, lower=True)
263 alpha = cho_solve((L, low), self.y_train)
264 except:
265 cho_solve_failed = True
266 alpha = np.linalg.solve(K, self.y_train)
267 print('Warning: Cholesky decomposition failed via scipy, using standard np solve instead.')
269 # Normalize test inputs
270 if self.normalize:
271 X_test = utils.normalize_x_data_test(X_test, self.sigmas_x, self.mus_x)
273 # Set up test derivative locations
274 if return_deriv:
275 derivative_locations_test = [
276 list(range(X_test.shape[0])) for _ in range(len(common_derivs))]
277 else:
278 derivative_locations_test = None
280 # Compute train-test differences
281 # if self.kernel == 'SI':
282 # diff_x_test_x_train = degp_utils.differences_by_dim_func_SI(
283 # self.x_train, X_test, self.n_order, return_deriv=return_deriv
284 # )
285 # else:
286 diff_x_test_x_train = degp_utils.differences_by_dim_func(
287 self.x_train, X_test, predict_order, predict_oti, return_deriv=return_deriv
288 )
290 # Compute train-test kernel
291 phi_train_test = predict_kernel_func(diff_x_test_x_train, length_scales)
292 if predict_order > 0:
293 if return_deriv:
294 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases, 2 * predict_order)
295 else:
296 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases, predict_order)
297 else:
298 phi_exp_train_test = phi_train_test.real
299 phi_exp_train_test = phi_exp_train_test[np.newaxis, :, :]
301 K_s = degp_utils.rbf_kernel_predictions(
302 phi_train_test, phi_exp_train_test, predict_order, self.n_bases,
303 self.flattened_der_indices, self.powers,
304 return_deriv=return_deriv,
305 index=self.derivative_locations,
306 common_derivs=common_derivs,
307 powers_predict=self.powers_predict
308 )
310 # Compute posterior mean
311 f_mean = K_s.T @ alpha
313 # Denormalize predictions
314 if self.normalize:
315 if return_deriv:
316 f_mean = utils.transform_predictions(
317 f_mean, self.mu_y, self.sigma_y, self.sigmas_x,
318 common_derivs, X_test)
319 else:
320 f_mean = self.mu_y + f_mean * self.sigma_y
322 # Reshape predictions
323 f_mean = f_mean.reshape(-1, 1)
324 n = X_test.shape[0]
325 m = f_mean.shape[0]
326 num_derivs = m // n
327 reshaped_mean = f_mean.reshape(num_derivs, n)
329 if return_full_cov and not calc_cov:
330 raise ValueError("return_full_cov=True requires calc_cov=True.")
332 if not calc_cov:
333 return reshaped_mean
335 # Compute test-test differences
336 diff_x_test_x_test = degp_utils.differences_by_dim_func(
337 X_test, X_test, predict_order, predict_oti, return_deriv=return_deriv
338 )
340 # Compute test-test kernel
341 phi_test_test = predict_kernel_func(diff_x_test_x_test, length_scales)
342 if predict_order > 0:
343 phi_exp_test_test = phi_test_test.get_all_derivs(self.n_bases, 2 * predict_order)
344 else:
345 phi_exp_test_test = phi_test_test.real
346 phi_exp_test_test = phi_exp_test_test[np.newaxis,:,:]
348 K_ss = degp_utils.rbf_kernel_predictions(
349 phi_test_test, phi_exp_test_test, predict_order, self.n_bases,
350 self.flattened_der_indices, self.powers,
351 return_deriv=return_deriv,
352 index=derivative_locations_test,
353 common_derivs=common_derivs,
354 calc_cov=True,
355 powers_predict=self.powers_predict
356 )
358 # Compute predictive covariance
359 if cho_solve_failed:
360 f_cov = K_ss - K_s.T @ np.linalg.inv(K) @ K_s
361 else:
362 v = solve_triangular(L, K_s, lower=low)
363 f_cov = K_ss - v.T @ v
365 # Transform covariance
366 if self.normalize:
367 if return_deriv:
368 f_var = utils.transform_cov(
369 f_cov, self.sigma_y, self.sigmas_x,
370 common_derivs, X_test)
371 scale_vec = np.full(f_cov.shape[0], self.sigma_y, dtype=float)
372 n_test = X_test.shape[0]
373 for i, deriv_spec in enumerate(common_derivs):
374 factor = self.sigma_y
375 for pair in deriv_spec:
376 factor = factor / (
377 self.sigmas_x[0][pair[0] - 1] ** pair[1]
378 )
379 start = (i + 1) * n_test
380 stop = (i + 2) * n_test
381 scale_vec[start:stop] = factor
382 f_cov_full = f_cov * np.outer(scale_vec, scale_vec)
383 else:
384 f_var = self.sigma_y ** 2 * np.diag(np.abs(f_cov))
385 f_cov_full = f_cov * self.sigma_y ** 2
386 else:
387 f_var = np.diag(np.abs(f_cov))
388 f_cov_full = f_cov
390 reshaped_var = f_var.reshape(num_derivs, n)
391 if return_full_cov:
392 return reshaped_mean, reshaped_var, f_cov_full
393 return reshaped_mean, reshaped_var