Coverage for jetgp/full_gddegp/gddegp.py: 89%
179 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-05-07 14:04 -0500
1import numpy as np
2import jetgp.utils as utils
3from jetgp.kernel_funcs.kernel_funcs import KernelFactory, get_oti_module
4from jetgp.full_gddegp.optimizer import Optimizer
5from jetgp.full_gddegp import gddegp_utils
6from scipy.linalg import cho_solve, cho_factor, solve_triangular
7import warnings
9class gddegp:
10 """
11 Global Directional Derivative-Enhanced Gaussian Process (GDDEGP) model.
13 Supports point-wise directional derivatives with unique rays per point,
14 hypercomplex representation, and automatic normalization. Includes methods
15 for training, prediction, and uncertainty quantification using kernel methods.
17 Parameters
18 ----------
19 x_train : ndarray
20 Training input data of shape (n_samples, n_features).
21 y_train : list or ndarray
22 Training targets or list of directional derivatives.
23 n_order : int
24 Maximum derivative order.
25 rays_list : list of ndarray
26 List of ray arrays. rays_list[i] has shape (d, len(derivative_locations[i])).
27 der_indices : list of lists
28 Derivative multi-indices corresponding to each derivative term.
29 derivative_locations : list of lists
30 Which training points have which derivatives.
31 n_bases : int, optional
32 Override the OTI space size. By default ``2 * n_direction_types`` (inferred
33 from ``der_indices``). Pass explicitly when training on function values only
34 (``der_indices=[]``) and you still want to predict directional derivatives:
35 set ``n_bases = 2 * n_prediction_direction_types``.
36 normalize : bool, default=True
37 Whether to normalize inputs and outputs.
38 sigma_data : float or array-like, optional
39 Observation noise standard deviation or diagonal noise values.
40 kernel : str, default='SE'
41 Kernel type ('SE', 'RQ', 'Matern', etc.).
42 kernel_type : str, default='anisotropic'
43 Kernel anisotropy ('anisotropic' or 'isotropic').
44 smoothness_parameter : float, optional
45 Smoothness parameter for Matern kernel.
46 """
48 def __init__(self, x_train, y_train, n_order, rays_list, der_indices,
49 derivative_locations=None, n_bases=None, normalize=True,
50 sigma_data=None, kernel="SE", kernel_type="anisotropic",
51 smoothness_parameter=None):
53 if n_order > 0 and derivative_locations is None:
54 import warnings
55 # Count total number of derivative components across all orders
56 n_derivs = sum(len(order_derivs) for order_derivs in der_indices)
57 n_train = len(x_train)
58 derivative_locations = [[i for i in range(n_train)] for _ in range(n_derivs)]
59 warnings.warn(
60 f"derivative_locations not provided. Assuming all {n_derivs} derivative(s) "
61 f"are available at all {n_train} training point(s).",
62 UserWarning
63 )
65 elif der_indices is None and n_order == 0:
66 der_indices = []
67 derivative_locations = []
69 self.x_train = x_train
70 self.y_train = y_train
71 self.sigma_data = sigma_data
72 self.n_order = n_order
73 self.max_order = n_order
74 self.rays_list = rays_list
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.normalize = normalize
80 self.derivative_locations = derivative_locations
81 self.der_indices = der_indices
83 # Flatten derivative indices first so we can size the OTI module correctly.
84 # GDDEGP needs 2 OTI bases per direction type (one odd tag for X1, one even
85 # tag for X2), so n_bases = 2 * n_direction_types by default.
86 # An explicit n_bases can be passed to support function-only training
87 # (der_indices=[]) while still reserving OTI space for derivative predictions.
88 self.flattened_der_indices = utils.flatten_der_indices(der_indices)
89 if n_bases is not None:
90 self.n_bases = n_bases
91 else:
92 self.n_bases = 2 * len(self.flattened_der_indices)
93 self.oti = get_oti_module(self.n_bases, n_order)
95 if normalize:
96 self.y_train, self.mu_y, self.sigma_y, self.sigmas_x, self.mus_x, sigma_data = \
97 utils.normalize_y_data_directional(
98 x_train, y_train, sigma_data, self.flattened_der_indices)
99 self.rays_list = utils.normalize_directions_2(self.sigmas_x, self.rays_list)
100 self.x_train = utils.normalize_x_data_train(x_train)
101 else:
102 self.x_train = x_train
103 self.y_train = utils.reshape_y_train(y_train)
105 self.differences_by_dim = gddegp_utils.differences_by_dim_func(
106 self.x_train, self.x_train,
107 self.rays_list, self.rays_list,
108 self.derivative_locations, self.derivative_locations,
109 n_order, self.oti, return_deriv=True
110 )
112 self.sigma_data = (
113 np.zeros((self.y_train.shape[0], self.y_train.shape[0]))
114 if sigma_data is None else np.diag(sigma_data)
115 )
116 self.sigma_data_sq_diag = (
117 np.zeros(self.y_train.shape[0])
118 if sigma_data is None
119 else np.asarray(sigma_data) ** 2
120 )
122 self.kernel_factory = KernelFactory(
123 dim=self.dim,
124 normalize=self.normalize,
125 n_order=self.max_order,
126 differences_by_dim=self.differences_by_dim,
127 smoothness_parameter=smoothness_parameter,
128 oti_module=self.oti,
129 sparse_diffs=False
130 )
131 self.kernel_func = self.kernel_factory.create_kernel(
132 kernel_name=self.kernel,
133 kernel_type=self.kernel_type
134 )
135 self.bounds = self.kernel_factory.bounds
136 self.optimizer = Optimizer(self)
138 def optimize_hyperparameters(self, *args, **kwargs):
139 """
140 Run the optimizer to find the best kernel hyperparameters.
141 Returns optimized hyperparameter vector.
142 """
143 self.params = self.optimizer.optimize_hyperparameters(*args, **kwargs)
144 return self.params
146 def predict(self, X_test, params, rays_predict=None, calc_cov=False,
147 return_deriv=False, derivs_to_predict=None,
148 return_full_cov=False):
149 """
150 Predict posterior mean and optional variance at test points.
152 Parameters
153 ----------
154 X_test : ndarray
155 Test input points of shape (n_test, n_features).
156 params : ndarray
157 Log-scaled kernel hyperparameters.
158 rays_predict : list of ndarray, optional
159 Rays at test points for derivative predictions.
160 calc_cov : bool, default=False
161 Whether to compute predictive variance.
162 return_deriv : bool, default=False
163 Whether to return derivative predictions.
164 derivs_to_predict : list, optional
165 Specific derivatives to predict. Can include derivatives not present in the
166 training set — the cross-covariance K_* is constructed from kernel derivatives
167 and does not require the requested derivative to have been observed during
168 training. Each entry must be a valid derivative spec within n_bases and n_order.
169 If None, defaults to all derivatives used in training.
170 return_full_cov : bool, default=False
171 If True, return the full posterior covariance matrix in addition to the
172 reshaped marginal variances. Requires calc_cov=True. This is opt-in to
173 preserve the existing predict return contract.
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 """
183 n_predict = X_test.shape[0]
185 # Handle missing rays_predict when derivatives are requested
186 if return_deriv and rays_predict is None:
187 n_rays = len(self.flattened_der_indices)
189 warnings.warn(
190 f"No rays_predict provided for derivative predictions. "
191 f"Predictions will be made along coordinate axes: "
192 f"[1,0,0,...], [0,1,0,...], etc. for {n_rays} directional derivative(s).",
193 UserWarning
194 )
196 # Construct coordinate axis rays for each entry in flattened_der_indices
197 # Each ray array has shape (n_bases, n_predict)
198 rays_predict = []
199 for i in range(n_rays):
200 # Cycle through coordinate axes if more rays than dimensions
201 axis_idx = i % self.dim
202 ray_array = np.zeros((self.dim, n_predict))
203 ray_array[axis_idx, :] = 1.0
204 rays_predict.append(ray_array)
206 # Warn if rays provided but not needed
207 if not return_deriv and rays_predict is not None:
208 warnings.warn(
209 "rays_predict was provided but return_deriv=False. "
210 "The provided rays will be ignored.",
211 UserWarning
212 )
214 # Validate rays_predict structure when predicting derivatives
215 if return_deriv and rays_predict is not None:
216 # Check number of requested prediction directions does not exceed
217 # the reserved OTI direction capacity. Prediction derivatives do not
218 # need to have been observed in training; function-only and sparse
219 # directional models can reserve extra bases for this purpose.
220 max_prediction_rays = self.n_bases // 2
221 if len(rays_predict) > max_prediction_rays:
222 raise ValueError(
223 f"Number of prediction rays ({len(rays_predict)}) exceeds the reserved "
224 f"direction capacity ({max_prediction_rays}). Increase n_bases when "
225 f"constructing the model to predict more derivative directions."
226 )
228 # Check shape of each ray array
229 for i, ray_array in enumerate(rays_predict):
230 if not isinstance(ray_array, np.ndarray):
231 raise TypeError(
232 f"Ray array {i} must be a numpy ndarray, got {type(ray_array).__name__}."
233 )
235 if ray_array.ndim != 2:
236 raise ValueError(
237 f"Ray array {i} must be 2-dimensional, got {ray_array.ndim} dimensions."
238 )
240 if ray_array.shape[0] != self.dim:
241 raise ValueError(
242 f"Ray array {i} has {ray_array.shape[0]} rows, expected {self.n_bases} "
243 f"(one per spatial dimension)."
244 )
246 if ray_array.shape[1] != n_predict:
247 raise ValueError(
248 f"Ray array {i} has {ray_array.shape[1]} columns, expected {n_predict} "
249 f"(one per test point)."
250 )
252 length_scales = params[:-1]
253 sigma_n = params[-1]
255 # Set up derivative prediction configuration
256 if return_deriv:
257 if derivs_to_predict is not None:
258 common_derivs = derivs_to_predict
259 else:
260 common_derivs = self.flattened_der_indices
262 # Determine prediction order from requested derivatives
263 required_order = max(
264 sum(pair[1] for pair in deriv_spec)
265 for deriv_spec in common_derivs
266 )
267 predict_order = max(required_order, self.n_order)
269 if predict_order > self.n_order:
270 predict_oti = get_oti_module(self.n_bases, predict_order)
271 smoothness_param = getattr(self.kernel_factory, 'alpha', None)
272 predict_kernel_factory = KernelFactory(
273 dim=self.dim,
274 normalize=self.normalize,
275 differences_by_dim=self.differences_by_dim,
276 n_order=predict_order,
277 smoothness_parameter=smoothness_param,
278 oti_module=predict_oti,
279 sparse_diffs=False
280 )
281 predict_kernel_func = predict_kernel_factory.create_kernel(
282 kernel_name=self.kernel, kernel_type=self.kernel_type
283 )
284 else:
285 predict_oti = self.oti
286 predict_kernel_func = self.kernel_func
287 else:
288 common_derivs = []
289 predict_order = self.n_order
290 predict_oti = self.oti
291 predict_kernel_func = self.kernel_func
293 # Check for cached Cholesky from optimizer
294 _cache_hit = (
295 hasattr(self, '_cached_params')
296 and self._cached_params is not None
297 and np.array_equal(self._cached_params, params)
298 )
300 if _cache_hit:
301 L = self._cached_L
302 low = self._cached_low
303 alpha = self._cached_alpha
304 cho_solve_failed = False
305 else:
306 # Build training kernel matrix
307 phi_train = self.kernel_func(self.differences_by_dim, length_scales)
308 if self.n_order == 0:
309 # Use a local variable so self.n_bases (which may have been set
310 # explicitly for derivative *prediction*) is never overwritten.
311 train_n_bases = 0
312 phi_exp_train = phi_train.real
313 phi_exp_train = phi_exp_train[np.newaxis,:,:]
314 else:
315 # Respect the configured GDDEGP directional slot count. The OTI
316 # kernel algebra may report many active scalar basis tags when a
317 # direction is dense in the ambient coordinates, but that should
318 # not expand the model beyond the number of trained directional
319 # slots.
320 train_n_bases = self.n_bases
321 phi_exp_train = phi_train.get_all_derivs(train_n_bases, 2 * self.n_order)
323 # Placeholder for powers (GDDEGP doesn't use sign powers like DEGP/DDEGP)
324 powers = [0] * (len(self.flattened_der_indices) + 1)
326 K = gddegp_utils.rbf_kernel(
327 phi_train, phi_exp_train, self.n_order, train_n_bases,
328 self.flattened_der_indices,
329 index=self.derivative_locations
330 )
331 K.flat[::K.shape[0] + 1] += (10 ** sigma_n) ** 2
332 K += self.sigma_data ** 2
333 self.K_train = K
334 # Solve linear system
335 try:
336 cho_solve_failed = False
337 L, low = cho_factor(K, lower=True)
338 alpha = cho_solve((L, low), self.y_train)
339 except:
340 cho_solve_failed = True
341 alpha = np.linalg.solve(K, self.y_train)
342 print('Warning: Cholesky decomposition failed via scipy, using standard np solve instead.')
344 # Normalize test inputs and rays
345 rays_test = rays_predict
347 if self.normalize:
348 X_test = utils.normalize_x_data_test(X_test, self.sigmas_x, self.mus_x)
350 if not return_deriv:
351 rays_test = None
352 derivative_locations_test = None
353 else:
354 derivative_locations_test = [
355 list(range(X_test.shape[0])) for _ in range(len(common_derivs))]
356 if self.normalize:
357 rays_test = utils.normalize_directions_2(self.sigmas_x, rays_test)
359 # Compute train-test differences
360 diff_x_train_x_test = gddegp_utils.differences_by_dim_func(
361 self.x_train, X_test,
362 self.rays_list, rays_test,
363 self.derivative_locations, derivative_locations_test,
364 predict_order, predict_oti, return_deriv=return_deriv
365 )
367 # Compute train-test kernel
368 phi_train_test = predict_kernel_func(diff_x_train_x_test, length_scales)
369 if predict_order > 0:
370 if return_deriv:
371 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases, 2 * predict_order)
372 else:
373 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases, predict_order)
374 else:
375 phi_exp_train_test = phi_train_test.real
376 phi_exp_train_test = phi_exp_train_test[np.newaxis, :, :]
377 K_s = gddegp_utils.rbf_kernel_predictions(
378 phi_train_test, phi_exp_train_test, predict_order, self.n_bases,
379 self.flattened_der_indices,
380 return_deriv=return_deriv,
381 index=self.derivative_locations,
382 common_derivs=common_derivs
383 )
385 # Compute posterior mean
386 f_mean = K_s @ alpha
388 # Denormalize predictions
389 if self.normalize:
390 if return_deriv:
391 f_mean = utils.transform_predictions_directional(
392 f_mean, self.mu_y, self.sigma_y, self.sigmas_x,
393 common_derivs, X_test)
394 else:
395 f_mean = self.mu_y + f_mean * self.sigma_y
397 # Reshape predictions
398 f_mean = f_mean.reshape(-1, 1)
399 n = X_test.shape[0]
400 m = f_mean.shape[0]
401 num_derivs = m // n
402 reshaped_mean = f_mean.reshape(num_derivs, n)
404 if return_full_cov and not calc_cov:
405 raise ValueError("return_full_cov=True requires calc_cov=True.")
407 if not calc_cov:
408 return reshaped_mean
410 # Compute test-test differences
411 diff_x_test_x_test = gddegp_utils.differences_by_dim_func(
412 X_test, X_test,
413 rays_test, rays_test,
414 derivative_locations_test, derivative_locations_test,
415 predict_order, predict_oti, return_deriv=return_deriv
416 )
418 # Compute test-test kernel
419 phi_test_test = predict_kernel_func(diff_x_test_x_test, length_scales)
420 if predict_order > 0:
421 phi_exp_test_test = phi_test_test.get_all_derivs(self.n_bases, 2 * predict_order)
422 else:
423 phi_exp_test_test = phi_test_test.real
424 phi_exp_test_test = phi_exp_test_test[np.newaxis, :, :]
426 K_ss = gddegp_utils.rbf_kernel_predictions(
427 phi_test_test, phi_exp_test_test, predict_order, self.n_bases,
428 self.flattened_der_indices,
429 return_deriv=return_deriv,
430 index=derivative_locations_test,
431 common_derivs=common_derivs,
432 calc_cov=True,
433 )
435 # Compute predictive covariance
436 if cho_solve_failed:
437 v_fallback = np.linalg.solve(K, K_s.T)
438 f_cov = K_ss - K_s @ v_fallback
439 else:
440 v = solve_triangular(L, K_s.T, lower=low)
441 f_cov = K_ss - v.T @ v
443 # Transform covariance
444 if self.normalize:
445 if return_deriv:
446 f_var = utils.transform_cov_directional(
447 f_cov, self.sigma_y, self.sigmas_x,
448 common_derivs, X_test)
449 f_cov_full = f_cov * self.sigma_y**2
450 else:
451 f_var = self.sigma_y ** 2 * np.diag(np.abs(f_cov))
452 f_cov_full = f_cov * self.sigma_y**2
453 else:
454 f_var = np.diag(np.abs(f_cov))
455 f_cov_full = f_cov
457 reshaped_var = f_var.reshape(num_derivs, n)
458 if return_full_cov:
459 return reshaped_mean, reshaped_var, f_cov_full
460 return reshaped_mean, reshaped_var