Coverage for jetgp/wdegp/wdegp.py: 86%

1""" 

2Unified Weighted Derivative-Enhanced Gaussian Process (WDEGP) 

3============================================================== 

4 

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

6 

7Submodel Types: 

8- 'degp': Coordinate-aligned derivatives (standard DEGP) 

9- 'ddegp': Global directional derivatives (same rays at all points) 

10- 'gddegp': Point-wise directional derivatives (unique rays per point) 

11""" 

12 

13import numpy as np 

14from numpy.linalg import cholesky 

15from scipy.linalg import cho_solve, cho_factor, solve_triangular 

16import jetgp.utils as utils 

17from jetgp.kernel_funcs.kernel_funcs import KernelFactory, get_oti_module 

18 

19 

20class wdegp: 

21 """ 

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

23 

24 Supports multiple submodels with DEGP, DDEGP, or GDDEGP derivative structure. 

25 All submodels use the same derivative type. 

26 

27 Parameters 

28 ---------- 

29 x_train : ndarray of shape (n_samples, n_features) 

30 Input training points. 

31 y_train : list of lists of arrays 

32 Each element is a submodel's data: [y_func, y_der1, y_der2, ...] 

33 n_order : int 

34 Maximum derivative order to be supported. 

35 n_bases : int 

36 Number of OTI basis terms used. 

37 der_indices : list of lists 

38 Multi-indices of derivatives for each submodel. 

39 derivative_locations : list of lists of lists, optional 

40 For each submodel, which points have which derivatives. 

41 derivative_locations[submodel][deriv_type] = [point_indices] 

42 If None, all points have all derivatives. 

43 submodel_type : str, default='degp' 

44 Type of derivative structure: 'degp', 'ddegp', or 'gddegp'. 

45 rays : ndarray, optional 

46 For 'ddegp' mode: global ray directions, shape (d, n_directions). 

47 All submodels share these rays. 

48 rays_list : list of list of ndarray, optional 

49 For 'gddegp' mode: point-wise rays organized by submodel. 

50 rays_list[submodel_idx][dir_idx] has shape (d, n_points_with_dir). 

51 Example: rays_list[0] = [rays_dir1_sm1, rays_dir2_sm1] for submodel 1. 

52 normalize : bool, default=True 

53 If True, normalizes the input and output data. 

54 sigma_data : float or ndarray, optional 

55 Known observation noise or covariance matrix. 

56 kernel : str, default='SE' 

57 Type of kernel to use: 'SE', 'RQ', 'Matern', or 'SineExp'. 

58 kernel_type : str, default='anisotropic' 

59 Whether kernel is 'anisotropic' or 'isotropic'. 

60 smoothness_parameter : float, optional 

61 Smoothness parameter for Matern kernel. 

62 """ 

63 

64 def __init__( 

65 self, 

66 x_train, 

67 y_train, 

68 n_order, 

69 n_bases, 

70 der_indices, 

71 derivative_locations=None, 

72 submodel_type='degp', 

73 rays=None, 

74 rays_list=None, 

75 normalize=True, 

76 sigma_data=None, 

77 kernel="SE", 

78 kernel_type="anisotropic", 

79 smoothness_parameter=None 

80 ): 

81 # Store basic parameters 

82 self.x_train = x_train 

83 self.y_train = y_train 

84 self.n_order = n_order 

85 self.n_bases = n_bases 

86 self.der_indices = der_indices 

87 self.derivative_locations = derivative_locations 

88 self.submodel_type = submodel_type 

89 self.rays = rays 

90 self.rays_list = rays_list 

91 self.normalize = normalize 

92 self.kernel = kernel 

93 self.kernel_type = kernel_type 

94 

95 if submodel_type == 'degp' or submodel_type == 'ddegp': 

96 self.oti = get_oti_module(self.n_bases, n_order) 

97 elif submodel_type == 'gddegp': 

98 self.oti = get_oti_module(2*self.n_bases, n_order) 

99 

100 self.num_points = len(x_train) 

101 self.dim = x_train.shape[1] 

102 self.num_submodels = len(y_train) 

103 

104 # Store original input for reference 

105 self.y_train_input = [yt.copy() if hasattr(yt, 'copy') else yt for yt in y_train] 

106 self.x_train_input = x_train.copy() 

107 

108 # Validate configuration 

109 self._validate_config() 

110 

111 # Set up derivative_locations defaults 

112 self._setup_derivative_locations() 

113 

114 # Process derivative indices 

115 self._setup_derivative_indices() 

116 

117 # Handle sigma_data 

118 if sigma_data is None: 

119 sigma_data = np.zeros(self._compute_total_constraints()) 

120 self.sigma_data = np.diag(sigma_data) 

121 self.sigma_data_sq_diag = np.asarray(sigma_data) ** 2 

122 

123 # Normalize if requested 

124 if normalize: 

125 self._normalize_data() 

126 else: 

127 self.y_train_normalized = [utils.reshape_y_train(submodel) for submodel in y_train] 

128 self.x_train_normalized = x_train 

129 

130 # Precompute differences 

131 self._precompute_differences() 

132 

133 # Set up kernel 

134 self.kernel_factory = KernelFactory( 

135 dim=self.dim, 

136 normalize=self.normalize, 

137 n_order=self.n_order, 

138 differences_by_dim=self.differences_by_dim, 

139 smoothness_parameter=smoothness_parameter, 

140 oti_module=self.oti, 

141 sparse_diffs=(self.submodel_type == 'degp') 

142 ) 

143 self.kernel_func = self.kernel_factory.create_kernel( 

144 kernel_name=self.kernel, 

145 kernel_type=self.kernel_type, 

146 ) 

147 self.bounds = self.kernel_factory.bounds 

148 

149 # Set up optimizer 

150 from jetgp.wdegp.optimizer import Optimizer 

151 self.optimizer = Optimizer(self) 

152 

153 def _validate_config(self): 

154 """Validate the configuration parameters.""" 

155 valid_types = ['degp', 'ddegp', 'gddegp'] 

156 if self.submodel_type not in valid_types: 

157 raise ValueError(f"submodel_type must be one of {valid_types}, got '{self.submodel_type}'") 

158 

159 if self.submodel_type == 'ddegp': 

160 if self.rays is None: 

161 raise ValueError("rays parameter is required for submodel_type='ddegp'") 

162 if self.rays.shape[0] != self.dim: 

163 raise ValueError(f"rays must have shape (d, n_directions), got {self.rays.shape}") 

164 

165 if self.submodel_type == 'gddegp': 

166 if self.rays_list is None: 

167 raise ValueError("rays_list parameter is required for submodel_type='gddegp'") 

168 if len(self.rays_list) != self.num_submodels: 

169 raise ValueError( 

170 f"rays_list must have {self.num_submodels} entries (one per submodel), " 

171 f"got {len(self.rays_list)}" 

172 ) 

173 for sm_idx, sm_rays in enumerate(self.rays_list): 

174 for dir_idx, r in enumerate(sm_rays): 

175 if r.shape[0] != self.dim: 

176 raise ValueError( 

177 f"rays_list[{sm_idx}][{dir_idx}] must have shape (d, n_points), " 

178 f"got {r.shape}" 

179 ) 

180 

181 def _setup_derivative_locations(self): 

182 """Set up derivative_locations with defaults.""" 

183 if self.derivative_locations is None: 

184 # Default: all points have all derivatives for all submodels 

185 self.derivative_locations = [] 

186 for submodel_idx in range(self.num_submodels): 

187 n_derivs = self._count_derivatives_for_submodel(submodel_idx) 

188 self.derivative_locations.append( 

189 [list(range(self.num_points))] * n_derivs 

190 ) 

191 

192 # For GDDEGP, validate rays_list matches derivative_locations per submodel 

193 if self.submodel_type == 'gddegp': 

194 for sm_idx, sm_rays in enumerate(self.rays_list): 

195 submodel_locs = self.derivative_locations[sm_idx] 

196 for dir_idx, r in enumerate(sm_rays): 

197 if dir_idx < len(submodel_locs): 

198 expected_size = len(submodel_locs[dir_idx]) 

199 if r.shape[1] != expected_size: 

200 raise ValueError( 

201 f"rays_list[{sm_idx}][{dir_idx}] has {r.shape[1]} columns but " 

202 f"derivative_locations[{sm_idx}][{dir_idx}] expects {expected_size} points" 

203 ) 

204 

205 def _count_derivatives_for_submodel(self, submodel_idx): 

206 """Count the number of derivative types for a submodel.""" 

207 # Flatten der_indices for this submodel 

208 der_idx = self.der_indices[submodel_idx] 

209 count = 0 

210 for group in der_idx: 

211 count += len(group) 

212 return count 

213 

214 def _setup_derivative_indices(self): 

215 """Process and flatten derivative indices for each submodel.""" 

216 self.flattened_der_indices = [] 

217 self.powers = [] 

218 

219 base_der_indices = utils.gen_OTI_indices(self.n_bases, self.n_order) 

220 

221 for submodel_idx, ders in enumerate(self.der_indices): 

222 self.powers.append( 

223 utils.build_companion_array(self.n_bases, self.n_order, ders) 

224 ) 

225 flat_indices = [i for sublist in ders for i in sublist] 

226 self.flattened_der_indices.append(flat_indices) 

227 

228 def _compute_total_constraints(self): 

229 """Compute total number of constraints across all submodels.""" 

230 total = 0 

231 for submodel_idx in range(self.num_submodels): 

232 # Function values 

233 total += self.num_points 

234 # Derivative values 

235 for locs in self.derivative_locations[submodel_idx]: 

236 total += len(locs) 

237 return total 

238 

239 def _normalize_data(self): 

240 """Normalize input and output data.""" 

241 self.y_train_normalized = [] 

242 

243 if self.submodel_type == 'degp': 

244 # Standard DEGP normalization 

245 for k, submodel_data in enumerate(self.y_train): 

246 y_norm, self.mu_y, self.sigma_y, self.sigmas_x, self.mus_x, _ = \ 

247 utils.normalize_y_data( 

248 self.x_train, 

249 submodel_data, 

250 np.zeros(len(submodel_data)), # placeholder 

251 self.flattened_der_indices[k] 

252 ) 

253 self.y_train_normalized.append(y_norm) 

254 

255 elif self.submodel_type == 'ddegp': 

256 # DDEGP uses directional normalization 

257 for k, submodel_data in enumerate(self.y_train): 

258 y_norm, self.mu_y, self.sigma_y, self.sigmas_x, self.mus_x, _ = \ 

259 utils.normalize_y_data_directional( 

260 self.x_train, 

261 submodel_data, 

262 np.zeros(len(submodel_data)), # placeholder 

263 self.flattened_der_indices[k] 

264 ) 

265 

266 self.y_train_normalized.append(y_norm) 

267 self.rays = self.rays / self.sigmas_x.flatten()[:, None] 

268 elif self.submodel_type == 'gddegp': 

269 # GDDEGP uses directional normalization 

270 for k, submodel_data in enumerate(self.y_train): 

271 y_norm, self.mu_y, self.sigma_y, self.sigmas_x, self.mus_x, _ = \ 

272 utils.normalize_y_data_directional( 

273 self.x_train, 

274 submodel_data, 

275 np.zeros(len(submodel_data)), # placeholder 

276 self.flattened_der_indices[k] 

277 ) 

278 self.y_train_normalized.append(y_norm) 

279 # Normalize rays_list using normalize_directions_2 

280 self.rays_list = [ 

281 utils.normalize_directions_2(self.sigmas_x, sm_rays) 

282 for sm_rays in self.rays_list 

283 ] 

284 

285 self.x_train_normalized = utils.normalize_x_data_train(self.x_train) 

286 

287 def _precompute_differences(self): 

288 """Precompute differences_by_dim based on submodel_type.""" 

289 x = self.x_train_normalized if self.normalize else self.x_train 

290 

291 if self.submodel_type == 'degp': 

292 from jetgp.wdegp import wdegp_utils 

293 self.differences_by_dim = wdegp_utils.differences_by_dim_func( 

294 x, x, self.n_order, self.oti 

295 ) 

296 elif self.submodel_type == 'ddegp': 

297 from jetgp.full_ddegp import wddegp_utils 

298 # Use first submodel's derivative_locations as reference 

299 # (all submodels share same rays structure) 

300 self.differences_by_dim = wddegp_utils.differences_by_dim_func( 

301 x, x, 

302 self.rays, 

303 self.n_order, 

304 self.oti, 

305 return_deriv=True 

306 ) 

307 elif self.submodel_type == 'gddegp': 

308 from jetgp.full_gddegp import wgddegp_utils 

309 

310 # For GDDEGP, combine rays from all submodels into global structure 

311 # Each submodel may have a different number of direction types 

312 max_n_dirs = max(len(sm_rays) for sm_rays in self.rays_list) 

313 

314 global_rays = [] 

315 global_derivative_locations = [] 

316 

317 for dir_idx in range(max_n_dirs): 

318 # Concatenate rays and locations for this direction from all submodels 

319 rays_for_dir = [] 

320 locs_for_dir = [] 

321 for sm_idx in range(self.num_submodels): 

322 sm_rays = self.rays_list[sm_idx] 

323 sm_locs = self.derivative_locations[sm_idx] 

324 # Only include if this submodel has this direction type 

325 if dir_idx < len(sm_rays): 

326 rays_for_dir.append(sm_rays[dir_idx]) 

327 locs_for_dir.extend(sm_locs[dir_idx]) 

328 

329 if rays_for_dir: # Only add if at least one submodel has this direction 

330 global_rays.append(np.hstack(rays_for_dir)) 

331 global_derivative_locations.append(locs_for_dir) 

332 

333 # Store global structures for use in kernel computations 

334 self.global_rays = global_rays 

335 self.global_derivative_locations = global_derivative_locations 

336 

337 self.differences_by_dim = wgddegp_utils.differences_by_dim_func( 

338 x, x, 

339 global_rays, global_rays, 

340 global_derivative_locations, global_derivative_locations, 

341 self.n_order, 

342 self.oti, 

343 return_deriv=True 

344 ) 

345 

346 def _get_utils_module(self): 

347 """Get the appropriate utils module based on submodel_type.""" 

348 if self.submodel_type == 'degp': 

349 from jetgp.wdegp import wdegp_utils 

350 return wdegp_utils 

351 elif self.submodel_type == 'ddegp': 

352 from jetgp.full_ddegp import wddegp_utils 

353 return wddegp_utils 

354 elif self.submodel_type == 'gddegp': 

355 from jetgp.full_gddegp import wgddegp_utils 

356 return wgddegp_utils 

357 

358 def optimize_hyperparameters(self, *args, **kwargs): 

359 """ 

360 Optimize hyperparameters via the configured optimizer. 

361 

362 Returns 

363 ------- 

364 ndarray 

365 Optimized hyperparameter vector. 

366 """ 

367 return self.optimizer.optimize_hyperparameters(*args, **kwargs) 

368 

369 def predict( 

370 self, 

371 X_test, 

372 length_scales, 

373 calc_cov=False, 

374 return_deriv=False, 

375 return_submodels=False, 

376 rays_predict=None, 

377 derivs_to_predict=None 

378 ): 

379 """ 

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

381  

382 Parameters 

383 ---------- 

384 X_test : ndarray of shape (n_test, n_features) 

385 Test input points. 

386 length_scales : ndarray 

387 Log-scaled kernel hyperparameters including noise level. 

388 calc_cov : bool, default=False 

389 If True, also compute and return predictive covariance. 

390 return_deriv : bool, default=False 

391 If True, also predict derivatives (requires rays_predict for GDDEGP). 

392 return_submodels : bool, default=False 

393 If True, return submodel-specific contributions. 

394 rays_predict : list of ndarray, optional 

395 For 'gddegp' mode with return_deriv=True: rays at test points. 

396 rays_predict[dir_idx] has shape (d, n_test). 

397 derivs_to_predict : list, optional 

398 Specific derivatives to predict. Can include derivatives not present in the 

399 training set of any submodel — each submodel constructs K_* from kernel 

400 derivatives directly. If None, defaults to all derivatives common to all 

401 submodels. 

402  

403 Returns 

404 ------- 

405 y_val : ndarray 

406 Predicted mean values. Shape depends on return_deriv. 

407 y_var : ndarray, optional 

408 Predictive variances (only if calc_cov=True). 

409 submodel_vals : list of ndarrays, optional 

410 Submodel predictions (only if return_submodels=True). 

411 submodel_cov : list of ndarrays, optional 

412 Submodel variances (only if calc_cov and return_submodels are True). 

413 """ 

414 import warnings 

415 

416 gp_utils = self._get_utils_module() 

417 

418 ell = length_scales[:-1] 

419 sigma_n = length_scales[-1] 

420 n_test = X_test.shape[0] 

421 n_train = self.x_train.shape[0] 

422 

423 # ========================================================================= 

424 # Validation checks for rays_predict parameter 

425 # ========================================================================= 

426 

427 # Check 1: Warning when rays_predict provided but return_deriv=False 

428 if not return_deriv and rays_predict is not None: 

429 warnings.warn( 

430 "rays_predict provided but return_deriv=False. " 

431 "The rays will be ignored.", 

432 UserWarning 

433 ) 

434 

435 # Check 2: Handle missing rays_predict for GDDEGP/DDEGP with return_deriv=True 

436 if return_deriv and rays_predict is None: 

437 if self.submodel_type == 'gddegp': 

438 # Determine number of ray directions from training 

439 n_rays = len(self.global_rays) 

440 warnings.warn( 

441 f"No rays_predict provided for GDDEGP with return_deriv=True. " 

442 f"Using coordinate axes as default prediction rays ({n_rays} directions).", 

443 UserWarning 

444 ) 

445 # Build default coordinate axis rays 

446 rays_predict = [] 

447 for i in range(n_rays): 

448 axis_idx = i % self.n_bases 

449 ray_array = np.zeros((self.n_bases, n_test)) 

450 ray_array[axis_idx, :] = 1.0 

451 rays_predict.append(ray_array) 

452 

453 elif self.submodel_type == 'ddegp': 

454 # DDEGP uses global rays, no rays_predict needed 

455 pass 

456 

457 # Check 3: Validate rays_predict structure when provided and return_deriv=True 

458 if return_deriv and rays_predict is not None: 

459 if self.submodel_type == 'gddegp': 

460 # Validate number of ray directions 

461 n_expected_rays = len(self.global_rays) 

462 if len(rays_predict) != n_expected_rays: 

463 raise ValueError( 

464 f"Number of prediction ray directions ({len(rays_predict)}) " 

465 f"does not match training ray directions ({n_expected_rays})" 

466 ) 

467 

468 # Validate each ray array 

469 for i, ray_array in enumerate(rays_predict): 

470 if not isinstance(ray_array, np.ndarray): 

471 raise TypeError( 

472 f"rays_predict[{i}] must be a numpy ndarray, " 

473 f"got {type(ray_array).__name__}" 

474 ) 

475 if ray_array.ndim != 2: 

476 raise ValueError( 

477 f"rays_predict[{i}] must be 2-dimensional, " 

478 f"got {ray_array.ndim} dimensions" 

479 ) 

480 if ray_array.shape[0] != self.dim: 

481 raise ValueError( 

482 f"rays_predict[{i}] has {ray_array.shape[0]} rows, " 

483 f"expected {self.dim} (number of input dimensions)" 

484 ) 

485 if ray_array.shape[1] != n_test: 

486 raise ValueError( 

487 f"rays_predict[{i}] has {ray_array.shape[1]} columns, " 

488 f"expected {n_test} (number of test points)" 

489 ) 

490 

491 elif self.submodel_type == 'ddegp': 

492 # For DDEGP, rays_predict should match global rays structure 

493 if not isinstance(rays_predict, np.ndarray): 

494 raise TypeError( 

495 f"rays_predict for DDEGP must be a numpy ndarray, " 

496 f"got {type(rays_predict).__name__}" 

497 ) 

498 if rays_predict.shape[0] != self.dim: 

499 raise ValueError( 

500 f"rays_predict has {rays_predict.shape[0]} rows, " 

501 f"expected {self.dim} (number of input dimensions)" 

502 ) 

503 

504 elif self.submodel_type == 'degp': 

505 # DEGP doesn't use rays, warn if provided 

506 warnings.warn( 

507 "rays_predict provided but submodel_type='degp' does not use rays. " 

508 "The rays will be ignored.", 

509 UserWarning 

510 ) 

511 

512 # ========================================================================= 

513 # End of validation checks 

514 # ========================================================================= 

515 

516 # Normalize test inputs 

517 if self.normalize: 

518 X_test_norm = utils.normalize_x_data_test(X_test, self.sigmas_x, self.mus_x) 

519 if rays_predict is not None and self.submodel_type == 'gddegp': 

520 rays_predict = utils.normalize_directions_2( 

521 self.sigmas_x, rays_predict) 

522 else: 

523 X_test_norm = X_test 

524 

525 x_train = self.x_train_normalized if self.normalize else self.x_train 

526 

527 # Determine which derivatives to predict across submodels 

528 if return_deriv: 

529 if derivs_to_predict is not None: 

530 common_derivs = derivs_to_predict 

531 else: 

532 if self.submodel_type in ('gddegp', 'wgddegp'): 

533 raise ValueError( 

534 "derivs_to_predict and rays_predict must be provided explicitly " 

535 "for WDEGP with GDDEGP submodels when return_deriv=True, " 

536 "because the length of rays_predict must match derivs_to_predict." 

537 ) 

538 # Default: predict all derivatives within n_bases and n_order, 

539 # deterministically sorted by total derivative order 

540 common_derivs = utils.flatten_der_indices( 

541 utils.gen_OTI_indices(self.n_bases, self.n_order) 

542 ) 

543 

544 # Determine prediction order from requested derivatives 

545 required_order = max( 

546 sum(pair[1] for pair in deriv_spec) 

547 for deriv_spec in common_derivs 

548 ) 

549 predict_order = max(required_order, self.n_order) 

550 

551 if predict_order > self.n_order: 

552 if self.submodel_type == 'degp' or self.submodel_type == 'ddegp': 

553 predict_oti = get_oti_module(self.n_bases, predict_order) 

554 elif self.submodel_type == 'gddegp': 

555 predict_oti = get_oti_module(2 * self.n_bases, predict_order) 

556 smoothness_param = getattr(self.kernel_factory, 'alpha', None) 

557 predict_kernel_factory = KernelFactory( 

558 dim=self.dim, 

559 normalize=self.normalize, 

560 differences_by_dim=self.differences_by_dim, 

561 n_order=predict_order, 

562 smoothness_parameter=smoothness_param, 

563 oti_module=predict_oti, 

564 sparse_diffs=(self.submodel_type == 'degp') 

565 ) 

566 predict_kernel_func = predict_kernel_factory.create_kernel( 

567 kernel_name=self.kernel, kernel_type=self.kernel_type 

568 ) 

569 else: 

570 predict_oti = self.oti 

571 predict_kernel_func = self.kernel_func 

572 

573 self.powers_predict = utils.build_companion_array_predict( 

574 self.n_bases, predict_order, common_derivs) 

575 else: 

576 common_derivs = [] 

577 self.powers_predict = None 

578 predict_order = self.n_order 

579 predict_oti = self.oti 

580 predict_kernel_func = self.kernel_func 

581 

582 y_val = 0 

583 y_var = 0 

584 submodel_vals = [] 

585 submodel_cov = [] 

586 

587 # For multiple submodels, compute weights (using function-only differences) 

588 if self.num_submodels > 1: 

589 if self.submodel_type == 'degp': 

590 from jetgp.wdegp import wdegp_utils 

591 diffs_for_weights = wdegp_utils.differences_by_dim_func( 

592 X_test_norm, x_train, 0,self.oti, return_deriv=False 

593 ) 

594 else: 

595 diffs_for_weights = self._compute_weight_differences(X_test_norm, x_train) 

596 

597 # Use train-train differences for weight computation 

598 diffs_train_for_weights = self.differences_by_dim 

599 

600 weights_matrix = gp_utils.determine_weights( 

601 diffs_train_for_weights, diffs_for_weights, ell, self.kernel_func, sigma_n 

602 ) 

603 

604 # Hoist shared computations out of the submodel loop. 

605 # OTI kernel_func reuses internal buffers, so a second call corrupts 

606 # the imaginary parts of previously returned OTI objects. We use 

607 # .copy() to snapshot each phi into independent memory before the 

608 # next kernel_func call overwrites the shared buffer. 

609 diffs_train_train = self.differences_by_dim 

610 diffs_train_test = self._compute_train_test_differences( 

611 x_train, X_test_norm, return_deriv, rays_predict, 

612 predict_order=predict_order, predict_oti=predict_oti 

613 ) 

614 

615 phi_train_train = self.kernel_func(diffs_train_train, ell).copy() 

616 if self.n_order == 0: 

617 self.n_bases_rays = 0 

618 phi_exp_train_train = phi_train_train.real 

619 phi_exp_train_train = phi_exp_train_train[np.newaxis, :, :] 

620 else: 

621 self.n_bases_rays = phi_train_train.get_active_bases()[-1] 

622 phi_exp_train_train = phi_train_train.get_all_derivs(self.n_bases_rays, 2 * self.n_order) 

623 

624 phi_train_test = predict_kernel_func(diffs_train_test, ell).copy() 

625 if predict_order > 0: 

626 if return_deriv: 

627 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases_rays, 2 * predict_order) 

628 else: 

629 phi_exp_train_test = phi_train_test.get_all_derivs(self.n_bases_rays, predict_order) 

630 else: 

631 phi_exp_train_test = phi_train_test.real 

632 phi_exp_train_test = phi_exp_train_test[np.newaxis, :, :] 

633 

634 # Pre-compute test-test kernel if covariance is requested 

635 phi_test_test = None 

636 phi_exp_test_test = None 

637 if calc_cov: 

638 diffs_test_test = self._compute_test_test_differences( 

639 X_test_norm, return_deriv, rays_predict, 

640 predict_order=predict_order, predict_oti=predict_oti 

641 ) 

642 phi_test_test = predict_kernel_func(diffs_test_test, ell).copy() 

643 if predict_order == 0: 

644 phi_exp_test_test = phi_test_test.real[np.newaxis, :, :] 

645 else: 

646 phi_exp_test_test = phi_test_test.get_all_derivs(self.n_bases_rays, 2 * predict_order) 

647 

648 # Loop over submodels 

649 for i in range(self.num_submodels): 

650 deriv_locs_i = self.derivative_locations[i] 

651 

652 # Build training kernel matrix (per-submodel: different indices/powers) 

653 K = gp_utils.rbf_kernel( 

654 phi_train_train, phi_exp_train_train, self.n_order, self.n_bases, 

655 self.flattened_der_indices[i], self.powers[i], 

656 index=deriv_locs_i 

657 ) 

658 K.flat[::K.shape[0] + 1] += (10 ** sigma_n) ** 2 

659 

660 # Solve linear system 

661 try: 

662 cho_solve_failed = False 

663 L, low = cho_factor(K, lower=True) 

664 alpha = cho_solve((L, low), self.y_train_normalized[i]) 

665 except Exception as e: 

666 cho_solve_failed = True 

667 alpha = np.linalg.solve(K, self.y_train_normalized[i]) 

668 print('Warning: Cholesky decomposition failed, using standard solve.') 

669 

670 K_s = gp_utils.rbf_kernel_predictions( 

671 phi_train_test, phi_exp_train_test, predict_order, self.n_bases_rays, 

672 self.flattened_der_indices[i], self.powers[i], 

673 return_deriv=return_deriv, 

674 index=deriv_locs_i, 

675 common_derivs=common_derivs, 

676 powers_predict=self.powers_predict 

677 ) 

678 

679 # Compute predictive mean 

680 if self.submodel_type == 'gddegp': 

681 K_s = K_s.T 

682 f_mean = K_s.T @ alpha 

683 else: 

684 f_mean = K_s.T @ alpha 

685 f_mean = f_mean.reshape(-1, 1) 

686 

687 # Denormalize predictions 

688 if self.normalize: 

689 if return_deriv: 

690 if self.submodel_type == 'gddegp' or self.submodel_type == 'ddegp': 

691 f_mean = utils.transform_predictions_directional( 

692 f_mean, self.mu_y, self.sigma_y, self.sigmas_x, 

693 common_derivs, X_test) 

694 else: 

695 f_mean = utils.transform_predictions( 

696 f_mean, self.mu_y, self.sigma_y, self.sigmas_x, 

697 common_derivs, X_test 

698 ) 

699 else: 

700 f_mean = self.mu_y + f_mean * self.sigma_y 

701 

702 # Reshape predictions 

703 n = X_test.shape[0] 

704 m = f_mean.shape[0] 

705 num_derivs = m // n 

706 reshaped = f_mean.reshape(num_derivs, n) 

707 

708 if self.n_order == 0 and not calc_cov: 

709 return reshaped 

710 

711 # Apply weights for multiple submodels 

712 if self.num_submodels > 1: 

713 # Compute weight for this submodel 

714 unique_indices = set() 

715 for subindex in deriv_locs_i: 

716 unique_indices.update(subindex) 

717 unique_indices = sorted(unique_indices) 

718 

719 weight = np.zeros(weights_matrix.shape[0]) 

720 for idx in unique_indices: 

721 weight += weights_matrix[:, idx] 

722 

723 if return_submodels: 

724 submodel_vals.append(reshaped.copy()) 

725 

726 reshaped_weighted = reshaped * weight 

727 else: 

728 reshaped_weighted = reshaped 

729 if return_submodels: 

730 raise ValueError('Cannot return submodels for a single model') 

731 

732 y_val += reshaped_weighted 

733 

734 # Compute covariance if requested 

735 if calc_cov: 

736 f_var = self._compute_predictive_variance( 

737 X_test_norm, deriv_locs_i, common_derivs, self.powers_predict, 

738 ell, i, K, K_s, L if not cho_solve_failed else None, 

739 low if not cho_solve_failed else None, cho_solve_failed, 

740 return_deriv, rays_predict, 

741 predict_order=predict_order, predict_oti=predict_oti, 

742 predict_kernel_func=predict_kernel_func, 

743 phi_test_test_cached=phi_test_test, 

744 phi_exp_test_test_cached=phi_exp_test_test 

745 ) 

746 

747 if self.num_submodels > 1: 

748 f_var_reshaped = f_var.reshape(num_derivs, n) 

749 if self.n_order == 0: 

750 return reshaped, f_var_reshaped 

751 if return_submodels: 

752 submodel_cov.append(f_var_reshaped.copy()) 

753 f_var_reshaped = f_var_reshaped * weight 

754 y_var += f_var_reshaped 

755 else: 

756 y_var = f_var.reshape(num_derivs, n) 

757 

758 # Return results 

759 if self.num_submodels == 1: 

760 if calc_cov: 

761 return (y_val, y_var ** 2) 

762 return y_val 

763 else: 

764 if return_submodels: 

765 if calc_cov: 

766 return (y_val, y_var ** 2, submodel_vals, submodel_cov) 

767 return (y_val, submodel_vals) 

768 else: 

769 if calc_cov: 

770 return (y_val, y_var ** 2) 

771 return y_val 

772 

773 def _compute_train_test_differences(self, x_train, X_test, return_deriv, rays_predict, 

774 submodel_idx=0, predict_order=None, predict_oti=None): 

775 """Compute train-test differences based on submodel_type.""" 

776 p_order = predict_order if predict_order is not None else self.n_order 

777 p_oti = predict_oti if predict_oti is not None else self.oti 

778 if self.submodel_type == 'degp': 

779 from jetgp.wdegp import wdegp_utils 

780 return wdegp_utils.differences_by_dim_func( 

781 x_train, X_test, p_order, p_oti, return_deriv=return_deriv 

782 ) 

783 elif self.submodel_type == 'ddegp': 

784 from jetgp.full_ddegp import wddegp_utils 

785 deriv_locs_test = [list(range(len(X_test)))] * self.rays.shape[1] if return_deriv else None 

786 return wddegp_utils.differences_by_dim_func( 

787 x_train, X_test, 

788 self.rays, 

789 p_order, 

790 p_oti, 

791 return_deriv=return_deriv 

792 ) 

793 elif self.submodel_type == 'gddegp': 

794 from jetgp.full_gddegp import gddegp_utils 

795 n_dirs = len(self.global_rays) 

796 rays_test = rays_predict if return_deriv and rays_predict is not None else None 

797 deriv_locs_test = [list(range(len(X_test)))] * n_dirs if return_deriv else None 

798 return gddegp_utils.differences_by_dim_func( 

799 x_train, X_test, 

800 self.global_rays, rays_test, 

801 self.global_derivative_locations, deriv_locs_test, 

802 p_order, 

803 p_oti, 

804 return_deriv=return_deriv 

805 ) 

806 

807 def _compute_weight_differences(self, X_test, x_train): 

808 """Compute differences for weight calculation (always without derivatives).""" 

809 if self.submodel_type == 'degp': 

810 from jetgp.wdegp import wdegp_utils 

811 return wdegp_utils.differences_by_dim_func(X_test, x_train, 0,self.oti, return_deriv=False) 

812 elif self.submodel_type == 'ddegp': 

813 from jetgp.full_ddegp import wddegp_utils 

814 return wddegp_utils.differences_by_dim_func( 

815 X_test, x_train, self.rays, 0,self.oti, return_deriv=False 

816 ) 

817 elif self.submodel_type == 'gddegp': 

818 from jetgp.full_gddegp import wgddegp_utils 

819 return wgddegp_utils.differences_by_dim_func( 

820 X_test, x_train, None, None, None, None, 0,self.oti, return_deriv=False 

821 ) 

822 

823 def _compute_predictive_variance( 

824 self, X_test, deriv_locs_i, common_derivs, powers_predict, ell, submodel_idx, K, K_s, L, low, cho_solve_failed, 

825 return_deriv, rays_predict, predict_order=None, predict_oti=None, predict_kernel_func=None, 

826 phi_test_test_cached=None, phi_exp_test_test_cached=None 

827 ): 

828 """Compute predictive variance for a submodel.""" 

829 gp_utils = self._get_utils_module() 

830 p_order = predict_order if predict_order is not None else self.n_order 

831 

832 # Use pre-computed test-test kernel if available, otherwise compute 

833 if phi_test_test_cached is not None: 

834 phi_test_test = phi_test_test_cached 

835 phi_exp_test_test = phi_exp_test_test_cached 

836 else: 

837 p_kernel_func = predict_kernel_func if predict_kernel_func is not None else self.kernel_func 

838 diffs_test_test = self._compute_test_test_differences( 

839 X_test, return_deriv, rays_predict, submodel_idx=submodel_idx, 

840 predict_order=p_order, predict_oti=predict_oti 

841 ) 

842 phi_test_test = p_kernel_func(diffs_test_test, ell).copy() 

843 if p_order == 0: 

844 phi_exp_test_test = phi_test_test.real[np.newaxis, :, :]