Weighted DEGP (WDEGP)

Unified Weighted Derivative-Enhanced Gaussian Process (WDEGP)

Supports DEGP, DDEGP, or GDDEGP mode for all submodels.

Submodel Types: - ‘degp’: Coordinate-aligned derivatives (standard DEGP) - ‘ddegp’: Global directional derivatives (same rays at all points) - ‘gddegp’: Point-wise directional derivatives (unique rays per point)

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

Bases: object

Unified Weighted Derivative-Enhanced Gaussian Process (WDEGP) regression model.

Supports multiple submodels with DEGP, DDEGP, or GDDEGP derivative structure. All submodels use the same derivative type.

Parameters:

• x_train (ndarray of shape (n_samples, n_features)) – Input training points.

• y_train (list of lists of arrays) – Each element is a submodel’s data: [y_func, y_der1, y_der2, …]

• n_order (int) – Maximum derivative order to be supported.

• n_bases (int) – Number of OTI basis terms used.

• der_indices (list of lists) – Multi-indices of derivatives for each submodel.

• derivative_locations (list of lists of lists, optional) – For each submodel, which points have which derivatives. derivative_locations[submodel][deriv_type] = [point_indices] If None, all points have all derivatives.

• submodel_type (str, default='degp') – Type of derivative structure: ‘degp’, ‘ddegp’, or ‘gddegp’.

• rays (ndarray, optional) – For ‘ddegp’ mode: global ray directions, shape (d, n_directions). All submodels share these rays.

• rays_list (list of list of ndarray, optional) – For ‘gddegp’ mode: point-wise rays organized by submodel. rays_list[submodel_idx][dir_idx] has shape (d, n_points_with_dir). Example: rays_list[0] = [rays_dir1_sm1, rays_dir2_sm1] for submodel 1.

• normalize (bool, default=True) – If True, normalizes the input and output data.

• sigma_data (float or ndarray, optional) – Known observation noise or covariance matrix.

• kernel (str, default='SE') – Type of kernel to use: ‘SE’, ‘RQ’, ‘Matern’, or ‘SineExp’.

• kernel_type (str, default='anisotropic') – Whether kernel is ‘anisotropic’ or ‘isotropic’.

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

optimize_hyperparameters(*args, **kwargs)

Optimize hyperparameters via the configured optimizer.

Returns:

Optimized hyperparameter vector.

Return type:

ndarray

predict(X_test, length_scales, calc_cov=False, return_deriv=False, return_submodels=False, rays_predict=None, derivs_to_predict=None)

Compute posterior predictive mean and (optionally) covariance at test points.

Parameters:

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

• length_scales (ndarray) – Log-scaled kernel hyperparameters including noise level.

• calc_cov (bool, default=False) – If True, also compute and return predictive covariance.

• return_deriv (bool, default=False) – If True, also predict derivatives (requires rays_predict for GDDEGP).

• return_submodels (bool, default=False) – If True, return submodel-specific contributions.

• rays_predict (list of ndarray, optional) – For ‘gddegp’ mode with return_deriv=True: rays at test points. rays_predict[dir_idx] has shape (d, n_test).

• derivs_to_predict (list, optional) – Specific derivatives to predict. Can include derivatives not present in the training set of any submodel — each submodel constructs K_* from kernel derivatives directly. If None, defaults to all derivatives common to all submodels.

Returns:

• y_val (ndarray) – Predicted mean values. Shape depends on return_deriv.

• y_var (ndarray, optional) – Predictive variances (only if calc_cov=True).

• submodel_vals (list of ndarrays, optional) – Submodel predictions (only if return_submodels=True).

• submodel_cov (list of ndarrays, optional) – Submodel variances (only if calc_cov and return_submodels are True).

jetgp.wdegp.wdegp_utils.deriv_map(nbases, order)

Creates a mapping from (order, index_within_order) to a single flattened index for all derivative components.

jetgp.wdegp.wdegp_utils.determine_weights(diffs_by_dim, diffs_test, length_scales, kernel_func, sigma_n)

Vectorized version: compute interpolation weights for multiple test points at once.

Parameters:

• diffs_by_dim (list of ndarray) – Pairwise differences between training points (by dimension).

• diffs_test (list of ndarray) – Pairwise differences between test points and training points (by dimension). Shape: each array is (n_test, n_train) or similar batch dimension.

• length_scales (array-like) – Kernel hyperparameters.

• kernel_func (callable) – Kernel function.

• sigma_n (float) – Noise parameter (if needed).

Returns:

weights_matrix – Interpolation weights for each test point.

Return type:

ndarray of shape (n_test, n_train)

jetgp.wdegp.wdegp_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.wdegp.wdegp_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.wdegp.wdegp_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.wdegp.wdegp_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.wdegp.wdegp_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.wdegp.wdegp_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.wdegp.wdegp_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.wdegp.wdegp_utils.find_common_derivatives(all_indices)

Find derivative indices common to all submodels.

jetgp.wdegp.wdegp_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.wdegp.wdegp_utils.rbf_kernel(phi, phi_exp, n_order, n_bases, der_indices, powers, index=-1)

Constructs the RBF kernel matrix with derivative entries using an efficient pre-allocation strategy combined with a single call to extract all derivative components.

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.

• n_bases (int) – Number of OTI bases.

• der_indices (list) – Derivative specifications.

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

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

Returns:

K – Full RBF kernel matrix with mixed function and derivative entries.

Return type:

ndarray

jetgp.wdegp.wdegp_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. If None, a new array is allocated.

Returns:

K – Full kernel matrix.

Return type:

ndarray

jetgp.wdegp.wdegp_utils.rbf_kernel_predictions(phi, phi_exp, n_order, n_bases, der_indices, powers, return_deriv, index=-1, 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 OTI bases.

• 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 test points.

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

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

• calc_cov (bool) – If True, computing covariance.

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

Returns:

K – Prediction kernel matrix.

Return type:

ndarray

jetgp.wdegp.wdegp_utils.to_list(x)

Convert tuple to list recursively.

jetgp.wdegp.wdegp_utils.to_tuple(item)

Convert list to tuple recursively.

jetgp.wdegp.wdegp_utils.transform_der_indices(der_indices, der_map)

Transforms a list of user-facing derivative specifications into the internal (order, index) format and the final flattened index.

class jetgp.wdegp.optimizer.Optimizer(model)

Bases: object

Optimizer class for fitting the hyperparameters of a weighted derivative-enhanced GP model (wDEGP) by minimizing the negative log marginal likelihood (NLL).

Supports DEGP, DDEGP, and GDDEGP modes.

model

Instance of a weighted derivative-enhanced GP model (wDEGP) with attributes: x_train, y_train, n_order, n_bases, der_indices, index, bounds, submodel_type, etc.

Type:

object

negative_log_marginal_likelihood(x0, x_train, y_train, n_order, n_bases, der_indices, index)

Computes the negative log marginal likelihood (NLL) for a given hyperparameter vector.

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

Parameters:

• x0 (ndarray) – Log-scaled hyperparameter vector, where the last entry is log10(sigma_n).

• x_train (list of ndarrays) – Input training points (unused inside loop, included for general interface).

• y_train (list of ndarrays) – List of function and derivative training values for each submodel.

• n_order (int) – Maximum order of derivatives used.

• n_bases (int) – Number of Taylor bases used in the expansion.

• der_indices (list) – Multi-index derivative information.

• index (list of lists) – Indices partitioning the training data into submodels (derivative_locations).

Returns:

The computed negative log marginal likelihood.

Return type:

float

nll_and_grad(x0)

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

nll_grad(x0)

Analytic gradient of the NLL w.r.t. log10-scaled hyperparameters.

nll_wrapper(x0)

Wrapper for NLL function to fit PSO optimizer interface.

Parameters:

x0 (ndarray) – Hyperparameter vector.

Returns:

Computed NLL value.

Return type:

float

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

Optimize the DEGP model hyperparameters using the specified optimizer.

Parameters:

optimizerstr or callable, default=”pso”

Name of optimizer or callable. Available: ‘pso’, ‘lbfgs’, ‘jade’, etc.

**kwargsdict

Additional arguments passed to the optimizer.

Returns:

best_xndarray

The optimal set of hyperparameters found.