Coverage for jetgp/full_ddegp/ddegp.py: 88%
140 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-10 23:19 -0500
1import numpy as np
2from numpy.linalg import cholesky, solve
3import jetgp.utils as utils
4from jetgp.kernel_funcs.kernel_funcs import KernelFactory, get_oti_module
5from jetgp.full_ddegp.optimizer import Optimizer
6from jetgp.full_ddegp import ddegp_utils
7from scipy.linalg import cho_solve, cho_factor, solve_triangular
10class ddegp:
11 """
12 Directional Derivative-Enhanced Gaussian Process (dDEGP) model.
14 Supports multiple directional 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 directional derivatives.
24 n_order : int
25 Maximum derivative order.
26 der_indices : list of lists
27 Derivative multi-indices corresponding to each derivative term.
28 rays : ndarray
29 Array of shape (d, n_rays), where each column is a direction vector.
30 **Important:** ``rays`` defines the OTI space dimension used internally
31 (``n_rays = rays.shape[1]``). Every direction you may ever want to
32 predict — including directions for which no training data exists — must
33 appear as a column here. A direction absent from ``rays`` cannot be
34 requested via ``derivs_to_predict`` at prediction time.
35 derivative_locations : list of lists
36 Which training points have which derivatives.
37 normalize : bool, default=True
38 Whether to normalize inputs and outputs.
39 sigma_data : float or array-like, optional
40 Observation noise standard deviation or diagonal noise values.
41 kernel : str, default='SE'
42 Kernel type ('SE', 'RQ', 'Matern', etc.).
43 kernel_type : str, default='anisotropic'
44 Kernel anisotropy ('anisotropic' or 'isotropic').
45 smoothness_parameter : float, optional
46 Smoothness parameter for Matern kernel.
47 """
49 def __init__(self, x_train, y_train, n_order, der_indices, rays,
50 derivative_locations=None, normalize=True, sigma_data=None,
51 kernel="SE", kernel_type="anisotropic", 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.rays = rays
74 self.n_rays = rays.shape[1]
75 self.dim = x_train.shape[1]
76 self.kernel = kernel
77 self.kernel_type = kernel_type
78 self.der_indices = der_indices
79 self.normalize = normalize
80 self.derivative_locations = derivative_locations
81 self.oti = get_oti_module(self.n_rays, n_order)
83 self.flattened_der_indices = utils.flatten_der_indices(der_indices)
85 if normalize:
86 self.y_train, self.mu_y, self.sigma_y, self.sigmas_x, self.mus_x, sigma_data = \
87 utils.normalize_y_data_directional(
88 x_train, y_train, sigma_data, self.flattened_der_indices)
89 self.rays = utils.normalize_directions(self.sigmas_x, self.rays)
90 self.x_train = utils.normalize_x_data_train(x_train)
91 else:
92 self.x_train = x_train
93 self.y_train = utils.reshape_y_train(y_train)
95 self.powers = utils.build_companion_array(self.n_rays, n_order, der_indices)
96 self.differences_by_dim = ddegp_utils.differences_by_dim_func(
97 self.x_train, self.x_train, self.rays, n_order, self.oti)
99 self.sigma_data = (
100 np.zeros((self.y_train.shape[0], self.y_train.shape[0]))
101 if sigma_data is None else np.diag(sigma_data)
102 )
103 self.sigma_data_sq_diag = (
104 np.zeros(self.y_train.shape[0])
105 if sigma_data is None
106 else np.asarray(sigma_data) ** 2
107 )
109 self.kernel_factory = KernelFactory(
110 dim=self.dim,
111 normalize=self.normalize,
112 n_order=self.n_order,
113 differences_by_dim=self.differences_by_dim,
114 smoothness_parameter=smoothness_parameter,
115 oti_module=self.oti,
116 sparse_diffs=False
117 )
118 self.kernel_func = self.kernel_factory.create_kernel(
119 kernel_name=self.kernel,
120 kernel_type=self.kernel_type
121 )
122 self.bounds = self.kernel_factory.bounds
123 self.optimizer = Optimizer(self)
125 def optimize_hyperparameters(self, *args, **kwargs):
126 """
127 Run the optimizer to find the best kernel hyperparameters.
128 Returns optimized hyperparameter vector.
129 """
130 return self.optimizer.optimize_hyperparameters(*args, **kwargs)
132 def predict(self, X_test, params, calc_cov=False, return_deriv=False, derivs_to_predict=None):
133 """
134 Predict posterior mean and optional variance at test points.
136 Parameters
137 ----------
138 X_test : ndarray
139 Test input points of shape (n_test, n_features).
140 params : ndarray
141 Log-scaled kernel hyperparameters.
142 calc_cov : bool, default=False
143 Whether to compute predictive variance.
144 return_deriv : bool, default=False
145 Whether to return derivative predictions.
146 derivs_to_predict : list, optional
147 Specific derivatives to predict. Can include derivatives not present in the
148 training set — the cross-covariance K_* is constructed from kernel derivatives
149 and does not require the requested derivative to have been observed during
150 training. Each entry must be a valid derivative spec within n_rays and n_order.
151 If None, defaults to all derivatives used in training.
153 **DDEGP-specific constraint:** each index must reference a ray that exists
154 in the ``rays`` array passed at construction. For example, ``[[4, 1]]``
155 requires ``rays`` to have at least 4 columns. Unlike DEGP — where the OTI
156 space always spans the fixed coordinate axes — the DDEGP OTI space is
157 spanned by the columns of ``rays``, so any direction not included there is
158 inaccessible at prediction time.
160 Returns
161 -------
162 f_mean : ndarray
163 Predictive mean vector.
164 f_var : ndarray, optional
165 Predictive variance vector (only if calc_cov=True).
166 """
167 length_scales = params[:-1]
168 sigma_n = params[-1]
170 # Set up derivative prediction configuration
171 if return_deriv:
172 if derivs_to_predict is not None:
173 common_derivs = derivs_to_predict
174 else:
175 common_derivs = self.flattened_der_indices
177 # Determine prediction order from requested derivatives
178 required_order = max(
179 sum(pair[1] for pair in deriv_spec)
180 for deriv_spec in common_derivs
181 )
182 predict_order = max(required_order, self.n_order)
184 if predict_order > self.n_order:
185 predict_oti = get_oti_module(self.n_rays, predict_order)
186 smoothness_param = getattr(self.kernel_factory, 'alpha', None)
187 predict_kernel_factory = KernelFactory(
188 dim=self.dim,
189 normalize=self.normalize,
190 differences_by_dim=self.differences_by_dim,
191 n_order=predict_order,
192 smoothness_parameter=smoothness_param,
193 oti_module=predict_oti,
194 sparse_diffs=False
195 )
196 predict_kernel_func = predict_kernel_factory.create_kernel(
197 kernel_name=self.kernel, kernel_type=self.kernel_type
198 )
199 else:
200 predict_oti = self.oti
201 predict_kernel_func = self.kernel_func
203 self.powers_predict = utils.build_companion_array_predict(
204 self.n_rays, predict_order, common_derivs)
205 else:
206 common_derivs = []
207 self.powers_predict = None
208 predict_order = self.n_order
209 predict_oti = self.oti
210 predict_kernel_func = self.kernel_func
212 # Check for cached Cholesky from optimizer
213 _cache_hit = (
214 hasattr(self, '_cached_params')
215 and self._cached_params is not None
216 and np.array_equal(self._cached_params, params)
217 )
219 if _cache_hit:
220 L = self._cached_L
221 low = self._cached_low
222 alpha = self._cached_alpha
223 self.n_bases_rays = self._cached_n_bases_rays
224 cho_solve_failed = False
225 else:
226 # Build training kernel matrix
227 phi_train = self.kernel_func(self.differences_by_dim, length_scales)
228 self.n_bases_rays = phi_train.get_active_bases()[-1]
229 if self.n_order > 0:
230 phi_exp_train = phi_train.get_all_derivs(self.n_bases_rays, 2 * self.n_order)
231 else:
232 phi_exp_train = phi_train.real
233 phi_exp_train = phi_exp_train[np.newaxis, :, :]
235 K = ddegp_utils.rbf_kernel(
236 phi_train, phi_exp_train, self.n_order, self.n_bases_rays,
237 self.flattened_der_indices, self.powers,
238 index=self.derivative_locations
239 )
240 K.flat[::K.shape[0] + 1] += (10 ** sigma_n) ** 2
241 K += self.sigma_data ** 2
243 # Solve linear system
244 try:
245 cho_solve_failed = False
246 L, low = cho_factor(K, lower=True)
247 alpha = cho_solve((L, low), self.y_train)
248 except:
249 cho_solve_failed = True
250 alpha = np.linalg.solve(K, self.y_train)
251 print('Warning: Cholesky decomposition failed via scipy, using standard np solve instead.')
253 # Normalize test inputs
254 if self.normalize:
255 X_test = utils.normalize_x_data_test(X_test, self.sigmas_x, self.mus_x)
257 # Set up test derivative locations
258 if return_deriv:
259 derivative_locations_test = [
260 list(range(X_test.shape[0])) for _ in range(len(common_derivs))]
261 else:
262 derivative_locations_test = None
264 # Compute train-test differences and kernel
265 diff_x_test_x_train = ddegp_utils.differences_by_dim_func(
266 self.x_train, X_test, self.rays, predict_order, predict_oti, return_deriv=return_deriv)
268 phi_train_test = predict_kernel_func(diff_x_test_x_train, length_scales)
269 if predict_order > 0:
270 if return_deriv:
271 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases_rays, 2 * predict_order)
272 else:
273 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases_rays, predict_order)
274 else:
275 phi_exp_train_test = phi_train_test.real
276 phi_exp_train_test = phi_exp_train_test[np.newaxis, :, :]
277 K_s = ddegp_utils.rbf_kernel_predictions(
278 phi_train_test, phi_exp_train_test, predict_order, self.n_bases_rays,
279 self.flattened_der_indices, self.powers,
280 return_deriv=return_deriv,
281 index=self.derivative_locations,
282 common_derivs=common_derivs,
283 powers_predict=self.powers_predict
284 )
286 f_mean = K_s.T @ alpha
288 # Denormalize predictions
289 if self.normalize:
290 if return_deriv:
291 f_mean = utils.transform_predictions_directional(
292 f_mean, self.mu_y, self.sigma_y, self.sigmas_x,
293 common_derivs, X_test)
294 else:
295 f_mean = self.mu_y + f_mean * self.sigma_y
297 # Reshape predictions
298 f_mean = f_mean.reshape(-1, 1)
299 n = X_test.shape[0]
300 m = f_mean.shape[0]
301 num_derivs = m // n
302 reshaped_mean = f_mean.reshape(num_derivs, n)
304 if not calc_cov:
305 return reshaped_mean
307 # Compute test-test differences and kernel for covariance
308 diff_x_test_x_test = ddegp_utils.differences_by_dim_func(
309 X_test, X_test, self.rays, predict_order, predict_oti, return_deriv=return_deriv)
311 phi_test_test = predict_kernel_func(diff_x_test_x_test, length_scales)
312 bases = phi_test_test.get_active_bases()
313 n_bases = bases[-1] if len(bases) > 0 else 0
315 if predict_order > 0:
316 phi_exp_test_test = phi_test_test.get_all_derivs(n_bases, 2 * predict_order)
317 else:
318 phi_exp_test_test = phi_test_test.real
319 phi_exp_test_test = phi_exp_test_test [np.newaxis, :, :]
320 K_ss = ddegp_utils.rbf_kernel_predictions(
321 phi_test_test, phi_exp_test_test, predict_order, n_bases,
322 self.flattened_der_indices, self.powers,
323 return_deriv=return_deriv,
324 index=derivative_locations_test,
325 common_derivs=common_derivs,
326 calc_cov=True,
327 powers_predict=self.powers_predict
328 )
330 # Compute predictive covariance
331 if cho_solve_failed:
332 f_cov = K_ss - K_s.T @ np.linalg.inv(K) @ K_s
333 else:
334 v = solve_triangular(L, K_s, lower=low)
335 f_cov = K_ss - v.T @ v
337 # Transform covariance
338 if self.normalize:
339 if return_deriv:
340 f_var = utils.transform_cov_directional(
341 f_cov, self.sigma_y, self.sigmas_x,
342 common_derivs, X_test)
343 else:
344 f_var = self.sigma_y ** 2 * np.diag(np.abs(f_cov))
345 else:
346 f_var = np.diag(np.abs(f_cov))
348 reshaped_var = f_var.reshape(num_derivs, n)
349 return reshaped_mean, reshaped_var