Coverage for jetgp/full_degp_sparse/degp.py: 65%
181 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-10 23:11 -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_sparse import degp_utils # noqa: E402 (sparse copy)
5import jetgp.utils as utils
6from jetgp.kernel_funcs.kernel_funcs import KernelFactory, get_oti_module
7from jetgp.full_degp_sparse.optimizer import Optimizer
8from jetgp.full_degp_sparse.sparse_cholesky import mmd_ordering, build_sparsity_pattern
11class degp:
12 """
13 Derivative-Enhanced Gaussian Process (DEGP) model.
15 Supports coordinate-aligned partial derivatives, hypercomplex representation,
16 and automatic normalization. Includes methods for training, prediction,
17 and uncertainty quantification using kernel methods.
19 Parameters
20 ----------
21 x_train : ndarray
22 Training input data of shape (n_samples, n_features).
23 y_train : list or ndarray
24 Training targets or list of partial derivatives.
25 n_order : int
26 Maximum derivative order.
27 n_bases : int
28 Number of input dimensions.
29 der_indices : list of lists
30 Derivative multi-indices corresponding to each derivative term.
31 derivative_locations : list of lists
32 Which training points have which derivatives.
33 normalize : bool, default=True
34 Whether to normalize inputs and outputs.
35 sigma_data : float or array-like, optional
36 Observation noise standard deviation or diagonal noise values.
37 kernel : str, default='SE'
38 Kernel type ('SE', 'RQ', 'Matern', 'SI', etc.).
39 kernel_type : str, default='anisotropic'
40 Kernel anisotropy ('anisotropic' or 'isotropic').
41 smoothness_parameter : float, optional
42 Smoothness parameter for Matern kernel.
43 rho : float, default=3.0
44 Sparsity radius multiplier for the geometric criterion
45 dist(x_P(i), x_P(j)) <= rho * l(j). Larger values give denser
46 sparsity patterns and more accurate (but slower) approximations.
47 use_supernodes : bool, default=True
48 If True, aggregate columns into supernodes to reduce the number of
49 local factorisations during sparse U construction.
50 supernode_lam : float, default=1.5
51 Merging threshold for supernode construction.
52 """
54 def __init__(
55 self,
56 x_train,
57 y_train,
58 n_order,
59 n_bases,
60 der_indices,
61 derivative_locations=None,
62 normalize=True,
63 sigma_data=None,
64 kernel="SE",
65 kernel_type="anisotropic",
66 smoothness_parameter=None,
67 rho=1.0,
68 use_supernodes=True,
69 supernode_lam=1.5,
70 ):
71 if n_order > 0 and derivative_locations is None:
72 import warnings
73 # Count total number of derivative components across all orders
74 n_derivs = sum(len(order_derivs) for order_derivs in der_indices)
75 n_train = len(x_train)
76 derivative_locations = [[i for i in range(n_train)] for _ in range(n_derivs)]
77 warnings.warn(
78 f"derivative_locations not provided. Assuming all {n_derivs} derivative(s) "
79 f"are available at all {n_train} training point(s).",
80 UserWarning
81 )
83 elif der_indices is None and n_order == 0:
84 der_indices = []
85 derivative_locations = []
86 self.n_order = n_order
87 self.n_bases = n_bases
88 self.dim = x_train.shape[1]
89 self.num_points = x_train.shape[0]
90 self.kernel = kernel
91 self.kernel_type = kernel_type
92 self.der_indices = der_indices
93 self.normalize = normalize
94 self.derivative_locations = derivative_locations
95 self.oti = get_oti_module(n_bases, n_order)
96 self.y_train_input = y_train
97 self.x_train_input = x_train
99 # Prepare indices and powers
100 self.flattened_der_indices = utils.flatten_der_indices(der_indices)
101 self.powers = utils.build_companion_array(n_bases, n_order, der_indices)
103 # Normalize data if required
104 if normalize:
105 (
106 self.y_train,
107 self.mu_y,
108 self.sigma_y,
109 self.sigmas_x,
110 self.mus_x,
111 sigma_data,
112 ) = utils.normalize_y_data(
113 x_train, y_train, sigma_data, self.flattened_der_indices
114 )
115 self.x_train = utils.normalize_x_data_train(x_train)
116 else:
117 self.x_train = x_train
118 self.y_train = utils.reshape_y_train(y_train)
120 # Compute differences for the kernel
121 # if kernel == 'SI':
122 # self.differences_by_dim = degp_utils.differences_by_dim_func_SI(
123 # self.x_train, self.x_train, n_order
124 # )
125 # else:
126 self.differences_by_dim = degp_utils.differences_by_dim_func(
127 self.x_train, self.x_train, n_order, self.oti
128 )
130 # Initialize noise matrix
131 self.sigma_data = (
132 np.zeros((self.y_train.shape[0], self.y_train.shape[0]))
133 if sigma_data is None
134 else np.diag(sigma_data)
135 )
136 self.sigma_data_sq_diag = (
137 np.zeros(self.y_train.shape[0])
138 if sigma_data is None
139 else np.asarray(sigma_data) ** 2
140 )
142 # Initialize kernel factory and optimizer
143 self.kernel_factory = KernelFactory(
144 dim=n_bases,
145 normalize=normalize,
146 differences_by_dim=self.differences_by_dim,
147 n_order=n_order,
148 smoothness_parameter=smoothness_parameter,
149 oti_module=self.oti
150 )
151 self.kernel_func = self.kernel_factory.create_kernel(
152 kernel_name=self.kernel, kernel_type=self.kernel_type
153 )
154 self.bounds = self.kernel_factory.bounds
155 self.optimizer = Optimizer(self)
157 # Sparse Cholesky: precompute MMD ordering and sparsity pattern once.
158 # The pattern depends only on (x_train, rho) and NOT on hyperparameters,
159 # so it is safe to compute here and reuse across all NLML evaluations.
160 self.rho = rho
161 self.use_supernodes = use_supernodes
162 self.supernode_lam = supernode_lam
163 self._setup_sparse_cholesky()
165 def _setup_sparse_cholesky(self):
166 """
167 Precompute the MMD ordering, fill-distances, and sparsity pattern.
169 Called once during __init__. Stores:
170 self.mmd_P : physical permutation (size N)
171 self.mmd_l : fill-distances (size N)
172 self.sparse_S : physical sparsity pattern
173 self.mmd_P_full : full K-matrix permutation (size N_total)
174 self.sparse_S_full : sparsity pattern in P_full-indexed space
175 self.sparse_supernodes_full : supernodes in P_full-indexed space
176 """
177 from jetgp.full_degp_sparse.sparse_cholesky import (
178 mmd_ordering, build_sparsity_pattern, build_supernodes,
179 expand_mmd_permutation, expand_sparsity_to_blocks,
180 expand_supernodes_to_blocks,
181 )
182 X = self.x_train # already normalised if normalize=True
183 self.mmd_P, self.mmd_l = mmd_ordering(X)
184 X_ord = X[self.mmd_P]
185 self.sparse_S = build_sparsity_pattern(X_ord, self.mmd_l, self.rho)
187 # Expand physical ordering to cover all K-matrix rows (function + derivatives)
188 self.mmd_P_full, self._phys_to_rows = expand_mmd_permutation(
189 self.mmd_P, self.num_points, self.derivative_locations
190 )
191 self.sparse_S_full = expand_sparsity_to_blocks(self.sparse_S, self._phys_to_rows)
192 # Pre-convert sparsity sets to numpy arrays for build_U
193 self.sparse_S_full_arr = {
194 j: np.asarray(s, dtype=np.intp) for j, s in self.sparse_S_full.items()
195 }
197 # Compute fill fraction to decide sparse vs dense factorisation path.
198 # When neighbourhoods are nearly full, dense Cholesky is faster than
199 # many overlapping block factorisations.
200 N_total = len(self.mmd_P_full)
201 total_nb = sum(len(s) for s in self.sparse_S_full.values())
202 self.sparse_fill_fraction = total_nb / (N_total * N_total)
203 self._use_dense_factor = self.sparse_fill_fraction > 0.25
205 if self.use_supernodes:
206 phys_sns = build_supernodes(
207 X_ord, self.mmd_l, self.sparse_S, lam=self.supernode_lam
208 )
209 self.sparse_supernodes = phys_sns
210 self.sparse_supernodes_full = expand_supernodes_to_blocks(
211 phys_sns, self._phys_to_rows
212 )
213 # Pre-convert supernode index lists to numpy arrays and build
214 # position lookups so build_U_supernodes avoids per-call overhead.
215 for sn in self.sparse_supernodes_full:
216 sn['children_arr'] = np.asarray(sn['children'])
217 ch_pos = {c: i for i, c in enumerate(sn['children'])}
218 sn['ch_pos'] = ch_pos
219 sn['parent_positions'] = np.array(
220 [ch_pos[p] for p in sn['parents']]
221 )
222 else:
223 self.sparse_supernodes = None
224 self.sparse_supernodes_full = None
226 def optimize_hyperparameters(self, *args, **kwargs):
227 """
228 Optimize model hyperparameters using the optimizer.
229 Returns optimized hyperparameter vector.
230 """
231 self.params = self.optimizer.optimize_hyperparameters(*args, **kwargs)
232 return self.params
234 def predict(self, X_test, params, calc_cov=False, return_deriv=False, derivs_to_predict=None):
235 """
236 Compute posterior predictive mean and (optionally) covariance at X_test.
238 Parameters
239 ----------
240 X_test : ndarray
241 Test input points of shape (n_test, n_features).
242 params : ndarray
243 Log-scaled kernel hyperparameters.
244 calc_cov : bool, default=False
245 Whether to compute predictive variance.
246 return_deriv : bool, default=False
247 Whether to return derivative predictions.
248 derivs_to_predict : list, optional
249 Specific derivatives to predict. Can include derivatives not present in the
250 training set — the cross-covariance K_* is constructed from kernel derivatives
251 and does not require the requested derivative to have been observed during
252 training. Each entry must be a valid derivative spec within n_bases and n_order
253 (e.g. ``[[3, 1]]`` for df/dx3 in a first-order model).
254 If None, defaults to all derivatives used in training.
256 Returns
257 -------
258 f_mean : ndarray
259 Predictive mean vector.
260 f_var : ndarray, optional
261 Predictive variance vector (only if calc_cov=True).
262 """
263 length_scales = params[:-1]
264 sigma_n = params[-1]
266 # Set up derivative prediction configuration
267 if return_deriv:
268 if derivs_to_predict is not None:
269 common_derivs = derivs_to_predict
270 else:
271 common_derivs = self.flattened_der_indices
273 # Determine prediction order from requested derivatives
274 required_order = max(
275 sum(pair[1] for pair in deriv_spec)
276 for deriv_spec in common_derivs
277 )
278 predict_order = max(required_order, self.n_order)
280 if predict_order > self.n_order:
281 predict_oti = get_oti_module(self.n_bases, predict_order)
282 smoothness_param = getattr(self.kernel_factory, 'alpha', None)
283 predict_kernel_factory = KernelFactory(
284 dim=self.n_bases,
285 normalize=self.normalize,
286 differences_by_dim=self.differences_by_dim,
287 n_order=predict_order,
288 smoothness_parameter=smoothness_param,
289 oti_module=predict_oti
290 )
291 predict_kernel_func = predict_kernel_factory.create_kernel(
292 kernel_name=self.kernel, kernel_type=self.kernel_type
293 )
294 else:
295 predict_oti = self.oti
296 predict_kernel_func = self.kernel_func
298 self.powers_predict = utils.build_companion_array_predict(
299 self.n_bases, predict_order, common_derivs)
300 else:
301 common_derivs = []
302 self.powers_predict = None
303 predict_order = self.n_order
304 predict_oti = self.oti
305 predict_kernel_func = self.kernel_func
307 # Reuse cached exact Cholesky + alpha from a previous predict call
308 # if available. Skip the cache when _cached_L is None — that means
309 # the cache was set by the sparse optimiser path (approximate alpha).
310 _cache_hit = (
311 hasattr(self, '_cached_params')
312 and self._cached_params is not None
313 and np.array_equal(self._cached_params, params)
314 and getattr(self, '_cached_L', None) is not None
315 )
317 if _cache_hit:
318 L = self._cached_L
319 low = self._cached_low
320 alpha = self._cached_alpha
321 cho_solve_failed = False
322 else:
323 # Build training kernel matrix (no cache available)
324 phi_train = self.kernel_func(self.differences_by_dim, length_scales)
326 if self.n_order > 0:
327 phi_exp_train = phi_train.get_all_derivs(self.n_bases, 2 * self.n_order)
328 else:
329 phi_exp_train = phi_train.real
330 phi_exp_train = phi_exp_train[np.newaxis, :, :]
332 K = degp_utils.rbf_kernel(
333 phi_train, phi_exp_train, self.n_order, self.n_bases,
334 self.flattened_der_indices, self.powers,
335 index=self.derivative_locations
336 )
337 K.flat[::K.shape[0] + 1] += (10 ** sigma_n) ** 2
338 K += self.sigma_data ** 2
340 # Final prediction always uses exact K solve (the sparse
341 # approximation is only for NLML during hyperparameter optimisation).
342 try:
343 L, low = cho_factor(K, lower=True)
344 alpha = cho_solve((L, low), self.y_train)
345 cho_solve_failed = False
346 except Exception:
347 alpha = np.linalg.solve(K, self.y_train)
348 L, low = None, None
349 cho_solve_failed = True
351 # Cache the exact solve for subsequent predict calls
352 self._cached_L = L
353 self._cached_low = low
354 self._cached_alpha = alpha
355 self._cached_params = params.copy()
357 # Normalize test inputs
358 if self.normalize:
359 X_test = utils.normalize_x_data_test(X_test, self.sigmas_x, self.mus_x)
361 # Set up test derivative locations
362 if return_deriv:
363 derivative_locations_test = [
364 list(range(X_test.shape[0])) for _ in range(len(common_derivs))]
365 else:
366 derivative_locations_test = None
368 # Compute train-test differences
369 # if self.kernel == 'SI':
370 # diff_x_test_x_train = degp_utils.differences_by_dim_func_SI(
371 # self.x_train, X_test, self.n_order, return_deriv=return_deriv
372 # )
373 # else:
374 diff_x_test_x_train = degp_utils.differences_by_dim_func(
375 self.x_train, X_test, predict_order, predict_oti, return_deriv=return_deriv
376 )
378 # Compute train-test kernel
379 phi_train_test = predict_kernel_func(diff_x_test_x_train, length_scales)
380 if predict_order > 0:
381 if return_deriv:
382 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases, 2 * predict_order)
383 else:
384 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases, predict_order)
385 else:
386 phi_exp_train_test = phi_train_test.real
387 phi_exp_train_test = phi_exp_train_test[np.newaxis, :, :]
389 K_s = degp_utils.rbf_kernel_predictions(
390 phi_train_test, phi_exp_train_test, predict_order, self.n_bases,
391 self.flattened_der_indices, self.powers,
392 return_deriv=return_deriv,
393 index=self.derivative_locations,
394 common_derivs=common_derivs,
395 powers_predict=self.powers_predict
396 )
398 # Compute posterior mean
399 f_mean = K_s.T @ alpha
401 # Denormalize predictions
402 if self.normalize:
403 if return_deriv:
404 f_mean = utils.transform_predictions(
405 f_mean, self.mu_y, self.sigma_y, self.sigmas_x,
406 common_derivs, X_test)
407 else:
408 f_mean = self.mu_y + f_mean * self.sigma_y
410 # Reshape predictions
411 f_mean = f_mean.reshape(-1, 1)
412 n = X_test.shape[0]
413 m = f_mean.shape[0]
414 num_derivs = m // n
415 reshaped_mean = f_mean.reshape(num_derivs, n)
417 if not calc_cov:
418 return reshaped_mean
420 # Compute test-test differences
421 diff_x_test_x_test = degp_utils.differences_by_dim_func(
422 X_test, X_test, predict_order, predict_oti, return_deriv=return_deriv
423 )
425 # Compute test-test kernel
426 phi_test_test = predict_kernel_func(diff_x_test_x_test, length_scales)
427 if predict_order > 0:
428 phi_exp_test_test = phi_test_test.get_all_derivs(self.n_bases, 2 * predict_order)
429 else:
430 phi_exp_test_test = phi_test_test.real
431 phi_exp_test_test = phi_exp_test_test[np.newaxis,:,:]
433 K_ss = degp_utils.rbf_kernel_predictions(
434 phi_test_test, phi_exp_test_test, predict_order, self.n_bases,
435 self.flattened_der_indices, self.powers,
436 return_deriv=return_deriv,
437 index=derivative_locations_test,
438 common_derivs=common_derivs,
439 calc_cov=True,
440 powers_predict=self.powers_predict
441 )
443 # Compute predictive covariance using sparse U: K^{-1} ≈ U U^T,
444 # so K_s^T K^{-1} K_s ≈ (U^T K_s)^T (U^T K_s).
445 # U is stored in original (non-permuted) index space as self._cached_U
446 # with column order following P; apply the same permutation to K_s rows.
447 if cho_solve_failed:
448 if hasattr(self, '_cached_U') and self._cached_U is not None:
449 P_full = self.mmd_P_full
450 U = self._cached_U
451 K_s_ord = K_s[P_full, :] # reorder rows to MMD order
452 v = U.T @ K_s_ord # shape (N, n_test*)
453 f_cov = K_ss - v.T @ v
454 else:
455 # Last resort: rebuild K and invert (only hit if U was never cached)
456 phi_train_fb = self.kernel_func(self.differences_by_dim, length_scales)
457 if self.n_order > 0:
458 phi_exp_fb = phi_train_fb.get_all_derivs(self.n_bases, 2 * self.n_order)
459 else:
460 phi_exp_fb = phi_train_fb.real[np.newaxis, :, :]
461 K_fb = degp_utils.rbf_kernel(
462 phi_train_fb, phi_exp_fb, self.n_order, self.n_bases,
463 self.flattened_der_indices, self.powers,
464 index=self.derivative_locations
465 )
466 K_fb += (10 ** sigma_n) ** 2 * np.eye(K_fb.shape[0])
467 K_fb += self.sigma_data ** 2
468 f_cov = K_ss - K_s.T @ np.linalg.inv(K_fb) @ K_s
469 else:
470 v = solve_triangular(L, K_s, lower=low)
471 f_cov = K_ss - v.T @ v
473 # Transform covariance
474 if self.normalize:
475 if return_deriv:
476 f_var = utils.transform_cov(
477 f_cov, self.sigma_y, self.sigmas_x,
478 common_derivs, X_test)
479 else:
480 f_var = self.sigma_y ** 2 * np.diag(np.abs(f_cov))
481 else:
482 f_var = np.diag(np.abs(f_cov))
484 reshaped_var = f_var.reshape(num_derivs, n)
485 return reshaped_mean, reshaped_var