Full DEGP

class jetgp.full_degp.degp.degp(x_train, y_train, n_order, n_bases, der_indices, derivative_locations=None, normalize=True, sigma_data=None, kernel='SE', kernel_type='anisotropic', smoothness_parameter=None)

Bases: object

Derivative-Enhanced Gaussian Process (DEGP) model.

Supports coordinate-aligned partial derivatives, hypercomplex representation, and automatic normalization. Includes methods for training, prediction, and uncertainty quantification using kernel methods.

Parameters:

• x_train (ndarray) – Training input data of shape (n_samples, n_features).

• y_train (list or ndarray) – Training targets or list of partial derivatives.

• n_order (int) – Maximum derivative order.

• n_bases (int) – Number of input dimensions.

• der_indices (list of lists) – Derivative multi-indices corresponding to each derivative term.

• derivative_locations (list of lists) – Which training points have which derivatives.

• normalize (bool, default=True) – Whether to normalize inputs and outputs.

• sigma_data (float or array-like, optional) – Observation noise standard deviation or diagonal noise values.

• kernel (str, default='SE') – Kernel type (‘SE’, ‘RQ’, ‘Matern’, ‘SI’, etc.).

• kernel_type (str, default='anisotropic') – Kernel anisotropy (‘anisotropic’ or ‘isotropic’).

• smoothness_parameter (float, optional) – Smoothness parameter for Matern kernel.

optimize_hyperparameters(*args, **kwargs)

Optimize model hyperparameters using the optimizer. Returns optimized hyperparameter vector.

predict(X_test, params, calc_cov=False, return_deriv=False, derivs_to_predict=None)

Compute posterior predictive mean and (optionally) covariance at X_test.

Parameters:

• X_test (ndarray) – Test input points of shape (n_test, n_features).

• params (ndarray) – Log-scaled kernel hyperparameters.

• calc_cov (bool, default=False) – Whether to compute predictive variance.

• return_deriv (bool, default=False) – Whether to return derivative predictions.

• derivs_to_predict (list, optional) – Specific derivatives to predict. Can include derivatives not present in the training set — the cross-covariance K_* is constructed from kernel derivatives and does not require the requested derivative to have been observed during training. Each entry must be a valid derivative spec within n_bases and n_order (e.g. [[3, 1]] for df/dx3 in a first-order model). If None, defaults to all derivatives used in training.

Returns:

• f_mean (ndarray) – Predictive mean vector.

• f_var (ndarray, optional) – Predictive variance vector (only if calc_cov=True).

jetgp.full_degp.degp_utils.deriv_map(nbases, order)

Create mapping from (order, index) to flattened index.

Parameters:

• nbases (int) – Number of base dimensions.

• order (int) – Maximum derivative order.

Returns:

map_deriv – Mapping where map_deriv[order][idx] gives the flattened index.

Return type:

list of lists

jetgp.full_degp.degp_utils.differences_by_dim_func(X1, X2, n_order, oti_module, return_deriv=True)

Compute pairwise differences between two input arrays X1 and X2 for each dimension, embedding hypercomplex units along each dimension for automatic differentiation.

For each dimension k, this function computes:

diff_k[i, j] = X1[i, k] + e_{k+1} - X2[j, k]

where e_{k+1} is a hypercomplex unit for the (k+1)-th dimension with order 2 * n_order.

Parameters:

• X1 (array_like of shape (n1, d)) – First set of input points with n1 samples in d dimensions.

• X2 (array_like of shape (n2, d)) – Second set of input points with n2 samples in d dimensions.

• n_order (int) – The base order used to construct hypercomplex units (e_{k+1}) with order 2 * n_order.

• oti_module (module) – The PyOTI static module (e.g., pyoti.static.onumm4n2).

• return_deriv (bool, optional) – If True, use 2*n_order for derivative predictions.

Returns:

differences_by_dim – A list where each element is an array of shape (n1, n2), containing the differences between corresponding dimensions of X1 and X2, augmented with hypercomplex units.

Return type:

list of length d

jetgp.full_degp.degp_utils.extract_and_assign(content_full, row_indices, col_indices, K, row_start, col_start, sign)

Extract submatrix and assign directly to K with sign multiplication. Combines extraction and assignment in one pass for better performance.

Parameters:

• content_full (ndarray of shape (n_rows_full, n_cols_full)) – Source matrix.

• row_indices (ndarray of int64) – Row indices to extract.

• col_indices (ndarray of int64) – Column indices to extract.

• K (ndarray) – Target matrix to fill.

• row_start (int) – Starting row index in K.

• col_start (int) – Starting column index in K.

• sign (float) – Sign multiplier (+1.0 or -1.0).

jetgp.full_degp.degp_utils.extract_cols(content_full, col_indices, n_rows)

Extract columns from content_full at specified indices.

Parameters:

• content_full (ndarray of shape (n_rows, n_cols_full)) – Source matrix.

• col_indices (ndarray of int64) – Column indices to extract.

• n_rows (int) – Number of rows.

Returns:

result – Extracted columns.

Return type:

ndarray of shape (n_rows, len(col_indices))

jetgp.full_degp.degp_utils.extract_cols_and_assign(content_full, col_indices, K, row_start, col_start, n_rows, sign)

Extract columns and assign directly to K with sign multiplication.

Parameters:

• content_full (ndarray of shape (n_rows, n_cols_full)) – Source matrix.

• col_indices (ndarray of int64) – Column indices to extract.

• K (ndarray) – Target matrix to fill.

• row_start (int) – Starting row index in K.

• col_start (int) – Starting column index in K.

• n_rows (int) – Number of rows to copy.

• sign (float) – Sign multiplier (+1.0 or -1.0).

jetgp.full_degp.degp_utils.extract_rows(content_full, row_indices, n_cols)

Extract rows from content_full at specified indices.

Parameters:

• content_full (ndarray of shape (n_rows_full, n_cols)) – Source matrix.

• row_indices (ndarray of int64) – Row indices to extract.

• n_cols (int) – Number of columns.

Returns:

result – Extracted rows.

Return type:

ndarray of shape (len(row_indices), n_cols)

jetgp.full_degp.degp_utils.extract_rows_and_assign(content_full, row_indices, K, row_start, col_start, n_cols, sign)

Extract rows and assign directly to K with sign multiplication.

Parameters:

• content_full (ndarray of shape (n_rows_full, n_cols)) – Source matrix.

• row_indices (ndarray of int64) – Row indices to extract.

• K (ndarray) – Target matrix to fill.

• row_start (int) – Starting row index in K.

• col_start (int) – Starting column index in K.

• n_cols (int) – Number of columns to copy.

• sign (float) – Sign multiplier (+1.0 or -1.0).

jetgp.full_degp.degp_utils.extract_submatrix(content_full, row_indices, col_indices)

Extract submatrix from content_full at specified row and column indices. Replaces the expensive np.ix_ operation.

Parameters:

• content_full (ndarray of shape (n_rows_full, n_cols_full)) – Source matrix.

• row_indices (ndarray of int64) – Row indices to extract.

• col_indices (ndarray of int64) – Column indices to extract.

Returns:

result – Extracted submatrix.

Return type:

ndarray of shape (len(row_indices), len(col_indices))

jetgp.full_degp.degp_utils.precompute_kernel_plan(n_order, n_bases, der_indices, powers, index)

Precompute all structural information needed by rbf_kernel so it can be reused across repeated calls with different phi_exp values.

Returns a dict containing flat indices, signs, index arrays, precomputed offsets/sizes, and mult_dir results for the dd block.

jetgp.full_degp.degp_utils.rbf_kernel(phi, phi_exp, n_order, n_bases, der_indices, powers, index=None)

Compute the derivative-enhanced RBF kernel matrix (optimized version).

This version uses Numba-accelerated functions for efficient matrix slicing, replacing expensive np.ix_ operations.

Parameters:

• phi (OTI array) – Base kernel matrix from kernel_func(differences, length_scales).

• phi_exp (ndarray) – Expanded derivative array from phi.get_all_derivs().

• n_order (int) – Maximum derivative order considered.

• n_bases (int) – Number of input dimensions.

• der_indices (list of lists) – Multi-index derivative structures for each derivative component.

• powers (list of int) – Powers of (-1) applied to each term.

• index (list of lists or None, optional) – If empty list, assumes uniform blocks. If provided, specifies which training point indices have each derivative type.

Returns:

K – Kernel matrix including function values and derivative terms.

Return type:

ndarray

jetgp.full_degp.degp_utils.rbf_kernel_fast(phi_exp_3d, plan, out=None)

Fast kernel assembly using a precomputed plan and fused numba kernel.

Parameters:

• phi_exp_3d (ndarray of shape (n_derivs, n_rows_func, n_cols_func)) – Pre-reshaped expanded derivative array.

• plan (dict) – Precomputed plan from precompute_kernel_plan().

• out (ndarray, optional) – Pre-allocated output array of shape (total, total). If None, a new array is allocated. Reusing a buffer avoids repeated allocation of large matrices during optimization loops.

Returns:

K – Full kernel matrix.

Return type:

ndarray

jetgp.full_degp.degp_utils.rbf_kernel_predictions(phi, phi_exp, n_order, n_bases, der_indices, powers, return_deriv, index=None, common_derivs=None, calc_cov=False, powers_predict=None)

Constructs the RBF kernel matrix for predictions with derivative entries.

This version uses Numba-accelerated functions for efficient matrix slicing.

Parameters:

• phi (OTI array) – Base kernel matrix between test and training points.

• phi_exp (ndarray) – Expanded derivative array from phi.get_all_derivs().

• n_order (int) – Maximum derivative order.

• n_bases (int) – Number of input dimensions.

• der_indices (list) – Derivative specifications for training data.

• powers (list of int) – Sign powers for each derivative type.

• return_deriv (bool) – If True, predict derivatives at ALL test points.

• index (list of lists or None) – Training point indices for each derivative type.

• common_derivs (list) – Common derivative indices to predict.

• calc_cov (bool) – If True, computing covariance (use all indices for rows).

• powers_predict (list of int, optional) – Sign powers for prediction derivatives.

Returns:

K – Prediction kernel matrix.

Return type:

ndarray

jetgp.full_degp.degp_utils.transform_der_indices(der_indices, der_map)

Transform derivative indices to flattened format.

Parameters:

• der_indices (list) – User-facing derivative specifications.

• der_map (list of lists) – Derivative mapping from deriv_map().

Returns:

• deriv_ind_transf (list) – Flattened indices for each derivative.

• deriv_ind_order (list) – (index, order) tuples for each derivative.

class jetgp.full_degp.optimizer.Optimizer(model)

Bases: object

Optimizer class to perform hyperparameter tuning for derivative-enhanced Gaussian Process models by minimizing the negative log marginal likelihood (NLL).

Parameters:

model (object) – An instance of a model (e.g., ddegp) containing the necessary training data and kernel configuration.

negative_log_marginal_likelihood(x0)

Compute the negative log marginal likelihood (NLL) of the model.

NLL = 0.5 * y^T K^-1 y + 0.5 * log|K| + 0.5 * N * log(2π)

Parameters:

x0 (ndarray) – Vector of log-scaled hyperparameters (length scales and noise).

Returns:

Value of the negative log marginal likelihood.

Return type:

float

nll_and_grad(x0)

Compute NLL and its gradient in a single pass, sharing one Cholesky.

Returns:

• nll (float)

• grad (ndarray)

nll_grad(x0)

Analytic gradient of the NLL (separate Cholesky from nll_wrapper).

nll_wrapper(x0)

Wrapper function to compute NLL for optimizer.

Parameters:

x0 (ndarray) – Hyperparameter vector.

Returns:

NLL evaluated at x0.

Return type:

float

optimize_hyperparameters(optimizer='pso', **kwargs)

Optimize the DEGP model hyperparameters using Particle Swarm Optimization (PSO).

Parameters:

n_restart_optimizerint, default=20

Maximum number of iterations for PSO.

swarm_sizeint, default=20

Number of particles in the swarm.

verbosebool, default=True

Controls verbosity of PSO output.

Returns:

best_xndarray

The optimal set of hyperparameters found.