Coverage for jetgp/full_gddegp/optimizer.py: 77%

1import numpy as np 

2import numba 

3from numpy.linalg import cholesky, solve 

4from jetgp.full_gddegp import gddegp_utils as utils 

5import jetgp.utils as gen_utils 

6from scipy.linalg import cho_solve, cho_factor 

7from jetgp.hyperparameter_optimizers import OPTIMIZERS 

8from line_profiler import profile 

9from jetgp.utils import matern_kernel_grad_builder 

10 

11 

12@numba.jit(nopython=True, parallel=True, cache=True) 

13def _subtract_outer(K_inv, alpha_v, W): 

14 """Fused: W[i,j] = K_inv[i,j] - alpha_v[i]*alpha_v[j] 

15 Exploits symmetry — writes both (i,j) and (j,i) per inner iteration, 

16 halving inner loop work vs a naive double loop. 

17 Avoids the N×N temporary that np.outer allocates. 

18 """ 

19 N = K_inv.shape[0] 

20 for i in numba.prange(N): 

21 ai = alpha_v[i] 

22 W[i, i] = K_inv[i, i] - ai * ai 

23 for j in range(i + 1, N): 

24 val = K_inv[i, j] - ai * alpha_v[j] 

25 W[i, j] = val 

26 W[j, i] = val 

27 

28 

29class Optimizer: 

30 """ 

31 Optimizer class to perform hyperparameter tuning for derivative-enhanced Gaussian Process models 

32 by minimizing the negative log marginal likelihood (NLL). 

33 

34 Parameters 

35 ---------- 

36 model : object 

37 An instance of a model (e.g., ddegp) containing the necessary training data 

38 and kernel configuration. 

39 """ 

40 

41 def __init__(self, model): 

42 self.model = model 

43 self._kernel_plan = None 

44 self._deriv_buf = None 

45 self._deriv_buf_shape = None 

46 self._deriv_buf_ndir = None 

47 self._deriv_factors = None 

48 self._deriv_factors_key = None 

49 self._K_buf = None 

50 self._dK_buf = None 

51 self._kernel_buf_size = None 

52 self._W_proj_buf = None 

53 self._W_proj_shape = None 

54 

55 def _get_deriv_buf(self, phi, n_bases, order): 

56 if self._deriv_buf_ndir is None: 

57 from math import comb 

58 self._deriv_buf_ndir = comb(n_bases + order, order) 

59 shape = (self._deriv_buf_ndir, phi.shape[0], phi.shape[1]) 

60 if self._deriv_buf is None or self._deriv_buf_shape != shape: 

61 self._deriv_buf = np.zeros(shape, dtype=np.float64) 

62 self._deriv_buf_shape = shape 

63 return self._deriv_buf 

64 

65 def _expand_derivs(self, phi, n_bases, deriv_order): 

66 """Expand OTI derivatives, using fast struct path if available.""" 

67 if hasattr(phi, 'get_all_derivs_fast'): 

68 buf = self._get_deriv_buf(phi, n_bases, deriv_order) 

69 factors = self._get_deriv_factors(n_bases, deriv_order) 

70 return phi.get_all_derivs_fast(factors, buf) 

71 return phi.get_all_derivs(n_bases, deriv_order) 

72 

73 @staticmethod 

74 def _enum_factors(max_basis, ordi): 

75 from math import factorial 

76 from collections import Counter 

77 if ordi == 1: 

78 for _ in range(max_basis): 

79 yield 1.0 

80 return 

81 for last in range(1, max_basis + 1): 

82 if ordi == 2: 

83 for i in range(1, last + 1): 

84 counts = Counter((i, last)) 

85 f = 1 

86 for c in counts.values(): 

87 f *= factorial(c) 

88 yield float(f) 

89 else: 

90 for _, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1): 

91 counts = dict(prefix_counts) 

92 counts[last] = counts.get(last, 0) + 1 

93 f = 1 

94 for c in counts.values(): 

95 f *= factorial(c) 

96 yield float(f) 

97 

98 @staticmethod 

99 def _enum_factors_with_counts(max_basis, ordi): 

100 from math import factorial 

101 from collections import Counter 

102 if ordi == 1: 

103 for i in range(1, max_basis + 1): 

104 yield 1.0, {i: 1} 

105 return 

106 for last in range(1, max_basis + 1): 

107 for _, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1): 

108 counts = dict(prefix_counts) 

109 counts[last] = counts.get(last, 0) + 1 

110 f = 1 

111 for c in counts.values(): 

112 f *= factorial(c) 

113 yield float(f), counts 

114 

115 def _get_deriv_factors(self, n_bases, order): 

116 key = (n_bases, order) 

117 if self._deriv_factors is not None and self._deriv_factors_key == key: 

118 return self._deriv_factors 

119 factors = [1.0] 

120 for ordi in range(1, order + 1): 

121 factors.extend(self._enum_factors(n_bases, ordi)) 

122 self._deriv_factors = np.array(factors, dtype=np.float64) 

123 self._deriv_factors_key = key 

124 return self._deriv_factors 

125 

126 def _ensure_kernel_plan(self, n_bases): 

127 """Lazily precompute kernel plan (once per n_bases).""" 

128 if self._kernel_plan is not None and self._kernel_plan_n_bases == n_bases: 

129 return 

130 if not hasattr(utils, 'precompute_kernel_plan'): 

131 self._kernel_plan = None 

132 return 

133 self._kernel_plan = utils.precompute_kernel_plan( 

134 self.model.n_order, n_bases, 

135 self.model.flattened_der_indices, 

136 None, # GDDEGP uses even/odd bases, not powers 

137 self.model.derivative_locations, 

138 ) 

139 self._kernel_plan_n_bases = n_bases 

140 self._K_buf = None 

141 self._dK_buf = None 

142 self._kernel_buf_size = None 

143 

144 def _ensure_kernel_bufs(self, n_rows_func): 

145 """Pre-allocate reusable K and dK buffers (avoids repeated malloc).""" 

146 if self._kernel_plan is None: 

147 return 

148 total = n_rows_func + self._kernel_plan['n_pts_with_derivs'] 

149 if self._kernel_buf_size != total: 

150 self._K_buf = np.empty((total, total)) 

151 self._dK_buf = np.empty((total, total)) 

152 self._kernel_buf_size = total 

153 if 'row_offsets_abs' not in self._kernel_plan: 

154 self._kernel_plan['row_offsets_abs'] = self._kernel_plan['row_offsets'] + n_rows_func 

155 self._kernel_plan['col_offsets_abs'] = self._kernel_plan['col_offsets'] + n_rows_func 

156 @profile 

157 def _build_K(self, phi_exp, phi, n_bases): 

158 """Build kernel matrix using fast path if available.""" 

159 self._ensure_kernel_plan(n_bases) 

160 if self._kernel_plan is not None: 

161 base_shape = phi.shape 

162 self._ensure_kernel_bufs(base_shape[0]) 

163 phi_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1]) 

164 return utils.rbf_kernel_fast(phi_3d, self._kernel_plan, out=self._K_buf) 

165 return utils.rbf_kernel( 

166 phi, phi_exp, self.model.n_order, n_bases, 

167 self.model.flattened_der_indices, 

168 index=self.model.derivative_locations, 

169 ) 

170 

171 @profile 

172 def negative_log_marginal_likelihood(self, x0): 

173 """ 

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

175 

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

177 

178 Parameters 

179 ---------- 

180 x0 : ndarray 

181 Vector of log-scaled hyperparameters (length scales and noise). 

182 

183 Returns 

184 ------- 

185 float 

186 Value of the negative log marginal likelihood. 

187 """ 

188 ell = x0[:-1] 

189 sigma_n = x0[-1] 

190 llhood = 0 

191 diffs = self.model.differences_by_dim 

192 phi = self.model.kernel_func(diffs, ell) 

193 if self.model.n_order == 0: 

194 n_bases = 0 

195 phi_exp = phi.real 

196 phi_exp = phi_exp[np.newaxis,:,:] 

197 else: 

198 n_bases = self.model.n_bases 

199 

200 # Extract ALL derivative components using the configured number of 

201 # directional slots, not the ambient active basis count implied by 

202 # dense rays. 

203 deriv_order = 2 * self.model.n_order 

204 phi_exp = self._expand_derivs(phi, n_bases, deriv_order) 

205 K = self._build_K(phi_exp, phi, n_bases) 

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

207 K.flat[::K.shape[0] + 1] += self.model.sigma_data_sq_diag 

208 

209 try: 

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

211 alpha = cho_solve( 

212 (L, low), 

213 self.model.y_train 

214 ) 

215 

216 # Cache for fast prediction 

217 self.model._cached_L = L 

218 self.model._cached_low = low 

219 self.model._cached_alpha = alpha 

220 self.model._cached_n_bases = self.model.n_bases 

221 self.model._cached_params = x0.copy() 

222 

223 data_fit = 0.5 * np.dot(self.model.y_train, alpha) 

224 log_det_K = np.sum(np.log(np.diag(L))) 

225 complexity = log_det_K 

226 N = len(self.model.y_train) 

227 const = 0.5 * N * np.log(2 * np.pi) 

228 return data_fit + complexity + const 

229 except Exception: 

230 return 1e6 

231 

232 def nll_wrapper(self, x0): 

233 """ 

234 Wrapper function to compute NLL for optimizer. 

235 

236 Parameters 

237 ---------- 

238 x0 : ndarray 

239 Hyperparameter vector. 

240 

241 Returns 

242 ------- 

243 float 

244 NLL evaluated at x0. 

245 """ 

246 return self.negative_log_marginal_likelihood(x0) 

247 

248 def nll_grad(self, x0): 

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

250 ln10 = np.log(10.0) 

251 

252 kernel = self.model.kernel 

253 kernel_type = self.model.kernel_type 

254 D = len(self.model.differences_by_dim) 

255 sigma_n_sq = (10.0 ** x0[-1]) ** 2 

256 diffs = self.model.differences_by_dim 

257 oti = self.model.kernel_factory.oti 

258 

259 phi = self.model.kernel_func(diffs, x0[:-1]) 

260 if self.model.n_order == 0: 

261 n_bases = 0 

262 phi_exp = phi.real[np.newaxis, :, :] 

263 else: 

264 active = phi.get_active_bases() 

265 n_bases = active[-1] if active else self.model.n_bases 

266 deriv_order = 2 * self.model.n_order 

267 phi_exp = self._expand_derivs(phi, n_bases, deriv_order) 

268 

269 K = self._build_K(phi_exp, phi, n_bases) 

270 K.flat[::K.shape[0] + 1] += sigma_n_sq 

271 K += self.model.sigma_data ** 2 

272 

273 try: 

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

275 alpha_v = cho_solve((L, low), self.model.y_train) 

276 N = len(self.model.y_train) 

277 K_inv = cho_solve((L, low), np.eye(N)) 

278 W = np.empty((N, N)) 

279 _subtract_outer(K_inv, alpha_v, W) 

280 except Exception: 

281 return np.zeros(len(x0)) 

282 

283 grad = np.zeros(len(x0)) 

284 use_fast = self._kernel_plan is not None 

285 base_shape = phi.shape 

286 

287 W_proj = None 

288 if use_fast: 

289 from math import comb 

290 ndir = comb(n_bases + deriv_order, deriv_order) 

291 proj_shape = (ndir, base_shape[0], base_shape[1]) 

292 if self._W_proj_buf is None or self._W_proj_shape != proj_shape: 

293 self._W_proj_buf = np.empty(proj_shape) 

294 self._W_proj_shape = proj_shape 

295 W_proj = self._W_proj_buf 

296 plan = self._kernel_plan 

297 row_off = plan.get('row_offsets_abs', plan['row_offsets'] + base_shape[0]) 

298 col_off = plan.get('col_offsets_abs', plan['col_offsets'] + base_shape[1]) 

299 utils._project_W_to_phi_space( 

300 W, W_proj, base_shape[0], base_shape[1], 

301 plan['fd_flat_indices'], plan['df_flat_indices'], 

302 plan['dd_flat_indices'], 

303 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'], 

304 plan['n_deriv_types'], row_off, col_off, 

305 ) 

306 

307 _use_vdot_fused = W_proj is not None and hasattr(phi, 'vdot_expand_fast') 

308 if _use_vdot_fused: 

309 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order) 

310 

311 def _gc(dphi): 

312 if _use_vdot_fused: 

313 return 0.5 * dphi.vdot_expand_fast(_vdot_factors, W_proj) 

314 if self.model.n_order == 0: 

315 dphi_exp = dphi.real[np.newaxis, :, :] 

316 else: 

317 dphi_exp = self._expand_derivs(dphi, n_bases, deriv_order) 

318 if W_proj is not None: 

319 dphi_3d = dphi_exp.reshape(W_proj.shape) 

320 return 0.5 * np.vdot(W_proj, dphi_3d) 

321 elif use_fast: 

322 dphi_3d = dphi_exp.reshape(dphi_exp.shape[0], base_shape[0], base_shape[1]) 

323 dK = utils.rbf_kernel_fast(dphi_3d, self._kernel_plan, out=self._dK_buf) 

324 return 0.5 * np.vdot(W, dK) 

325 else: 

326 dK = utils.rbf_kernel( 

327 dphi, dphi_exp, self.model.n_order, n_bases, 

328 self.model.flattened_der_indices, 

329 index=self.model.derivative_locations, 

330 ) 

331 return 0.5 * np.vdot(W, dK) 

332 

333 grad[-2] = _gc(oti.mul(2.0 * ln10, phi)) 

334 grad[-1] = ln10 * sigma_n_sq * np.trace(W) 

335 

336 if kernel == 'SE': 

337 if kernel_type == 'anisotropic': 

338 ell = 10.0 ** x0[:D] 

339 if hasattr(phi, 'fused_scale_sq_mul'): 

340 dphi_buf = oti.zeros(phi.shape) 

341 for d in range(D): 

342 dphi_buf.fused_scale_sq_mul(diffs[d], phi, -ln10 * ell[d] ** 2) 

343 grad[d] = _gc(dphi_buf) 

344 else: 

345 for d in range(D): 

346 grad[d] = _gc(oti.mul(-ln10 * ell[d] ** 2, 

347 oti.mul(oti.mul(diffs[d], diffs[d]), phi))) 

348 else: 

349 ell = 10.0 ** float(x0[0]) 

350 if hasattr(phi, 'fused_sum_sq'): 

351 sum_sq = oti.zeros(phi.shape) 

352 sum_sq.fused_sum_sq(diffs) 

353 else: 

354 sum_sq = oti.mul(diffs[0], diffs[0]) 

355 for d in range(1, D): 

356 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d])) 

357 grad[0] = _gc(oti.mul(-ln10 * ell ** 2, oti.mul(sum_sq, phi))) 

358 

359 elif kernel == 'RQ': 

360 if kernel_type == 'anisotropic': 

361 ell = 10.0 ** x0[:D]; alpha_rq = 10.0 ** float(x0[D]); alpha_idx = D 

362 else: 

363 ell = np.full(D, 10.0 ** float(x0[0])) 

364 alpha_rq = np.exp(float(x0[1])); alpha_idx = 1 

365 if hasattr(phi, 'fused_sqdist'): 

366 r2 = oti.zeros(phi.shape) 

367 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64) 

368 r2.fused_sqdist(diffs, ell_sq) 

369 else: 

370 r2 = oti.mul(ell[0], diffs[0]); r2 = oti.mul(r2, r2) 

371 for d in range(1, D): 

372 td = oti.mul(ell[d], diffs[d]); r2 = oti.sum(r2, oti.mul(td, td)) 

373 base = oti.sum(1.0, oti.mul(r2, 1.0 / (2.0 * alpha_rq))) 

374 inv_base = oti.pow(base, -1) 

375 phi_over_base = oti.mul(phi, inv_base) 

376 if kernel_type == 'anisotropic': 

377 if hasattr(phi, 'fused_scale_sq_mul'): 

378 dphi_buf = oti.zeros(phi.shape) 

379 for d in range(D): 

380 dphi_buf.fused_scale_sq_mul(diffs[d], phi_over_base, -ln10 * ell[d] ** 2) 

381 grad[d] = _gc(dphi_buf) 

382 else: 

383 for d in range(D): 

384 grad[d] = _gc(oti.mul(-ln10 * ell[d] ** 2, 

385 oti.mul(oti.mul(diffs[d], diffs[d]), phi_over_base))) 

386 else: 

387 if hasattr(phi, 'fused_sum_sq'): 

388 sum_sq = oti.zeros(phi.shape) 

389 sum_sq.fused_sum_sq(diffs) 

390 else: 

391 sum_sq = oti.mul(diffs[0], diffs[0]) 

392 for d in range(1, D): 

393 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d])) 

394 grad[0] = _gc(oti.mul(-ln10 * ell[0] ** 2, oti.mul(sum_sq, phi_over_base))) 

395 log_base = oti.log(base) 

396 term = oti.sub(oti.sub(1.0, inv_base), log_base) 

397 alpha_factor = ln10 * alpha_rq if kernel_type == 'anisotropic' else alpha_rq 

398 grad[alpha_idx] = _gc(oti.mul(alpha_factor, oti.mul(phi, term))) 

399 

400 elif kernel == 'SineExp': 

401 if kernel_type == 'anisotropic': 

402 ell = 10.0 ** x0[:D]; p = 10.0 ** x0[D:2*D] 

403 pip = np.pi / p; p_start = D 

404 else: 

405 ell = np.full(D, 10.0 ** float(x0[0])) 

406 pip = np.full(D, np.pi / 10.0 ** float(x0[1])); p_start = 1 

407 sin_d = [oti.sin(oti.mul(pip[d], diffs[d])) for d in range(D)] 

408 cos_d = [oti.cos(oti.mul(pip[d], diffs[d])) for d in range(D)] 

409 if kernel_type == 'anisotropic': 

410 if hasattr(phi, 'fused_scale_sq_mul'): 

411 dphi_buf = oti.zeros(phi.shape) 

412 for d in range(D): 

413 dphi_buf.fused_scale_sq_mul(sin_d[d], phi, -4.0 * ln10 * ell[d] ** 2) 

414 grad[d] = _gc(dphi_buf) 

415 else: 

416 for d in range(D): 

417 grad[d] = _gc(oti.mul(-4.0 * ln10 * ell[d] ** 2, 

418 oti.mul(oti.mul(sin_d[d], sin_d[d]), phi))) 

419 for d in range(D): 

420 sc = oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d])) 

421 grad[p_start + d] = _gc(oti.mul(4.0 * ln10 * ell[d] ** 2 * pip[d], 

422 oti.mul(sc, phi))) 

423 else: 

424 if hasattr(phi, 'fused_sum_sq'): 

425 ss = oti.zeros(phi.shape) 

426 ss.fused_sum_sq(sin_d) 

427 else: 

428 ss = oti.mul(sin_d[0], sin_d[0]) 

429 for d in range(1, D): 

430 ss = oti.sum(ss, oti.mul(sin_d[d], sin_d[d])) 

431 grad[0] = _gc(oti.mul(-4.0 * ln10 * ell[0] ** 2, oti.mul(ss, phi))) 

432 scd = oti.mul(sin_d[0], oti.mul(cos_d[0], diffs[0])) 

433 for d in range(1, D): 

434 scd = oti.sum(scd, oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d]))) 

435 grad[p_start] = _gc(oti.mul(4.0 * ln10 * ell[0] ** 2 * pip[0], 

436 oti.mul(scd, phi))) 

437 

438 elif kernel == 'Matern': 

439 kf = self.model.kernel_factory 

440 if not hasattr(kf, '_matern_grad_prebuild'): 

441 kf._matern_grad_prebuild = matern_kernel_grad_builder(getattr(kf, "nu", 1.5), oti_module=oti) 

442 ell = (10.0 ** x0[:D] if kernel_type == 'anisotropic' 

443 else np.full(D, 10.0 ** float(x0[0]))) 

444 sigma_f_sq = (10.0 ** float(x0[-2])) ** 2 

445 _eps = 1e-10 

446 if hasattr(phi, 'fused_sqdist'): 

447 r2 = oti.zeros(phi.shape) 

448 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64) 

449 r2.fused_sqdist(diffs, ell_sq) 

450 else: 

451 r2 = oti.mul(ell[0], diffs[0]); r2 = oti.mul(r2, r2) 

452 for d in range(1, D): 

453 td = oti.mul(ell[d], diffs[d]); r2 = oti.sum(r2, oti.mul(td, td)) 

454 r_oti = oti.sqrt(oti.sum(r2, _eps ** 2)) 

455 f_prime_r = kf._matern_grad_prebuild(r_oti) 

456 inv_r = oti.pow(r_oti, -1) 

457 base_matern = oti.mul(sigma_f_sq, oti.mul(f_prime_r, inv_r)) 

458 if kernel_type == 'anisotropic': 

459 if hasattr(phi, 'fused_scale_sq_mul'): 

460 dphi_buf = oti.zeros(phi.shape) 

461 for d in range(D): 

462 dphi_buf.fused_scale_sq_mul(diffs[d], base_matern, ln10 * ell[d] ** 2) 

463 grad[d] = _gc(dphi_buf) 

464 else: 

465 for d in range(D): 

466 d_sq = oti.mul(diffs[d], diffs[d]) 

467 dphi_d = oti.mul(ln10 * ell[d] ** 2, oti.mul(d_sq, base_matern)) 

468 grad[d] = _gc(dphi_d) 

469 else: 

470 if hasattr(phi, 'fused_sum_sq'): 

471 sum_dsq = oti.zeros(phi.shape) 

472 sum_dsq.fused_sum_sq(diffs) 

473 else: 

474 sum_dsq = oti.mul(diffs[0], diffs[0]) 

475 for d in range(1, D): 

476 sum_dsq = oti.sum(sum_dsq, oti.mul(diffs[d], diffs[d])) 

477 dphi_e = oti.mul(ln10 * ell[0] ** 2, oti.mul(sum_dsq, base_matern)) 

478 grad[0] = _gc(dphi_e) 

479 

480 return grad 

481 

482 @profile 

483 def nll_and_grad(self, x0): 

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

485 ln10 = np.log(10.0) 

486 

487 kernel = self.model.kernel 

488 kernel_type = self.model.kernel_type 

489 D = len(self.model.differences_by_dim) 

490 sigma_n_sq = (10.0 ** x0[-1]) ** 2 

491 diffs = self.model.differences_by_dim 

492 oti = self.model.kernel_factory.oti 

493 

494 # --- shared kernel computation (done ONCE) --- 

495 phi = self.model.kernel_func(diffs, x0[:-1]) 

496 if self.model.n_order == 0: 

497 n_bases = 0 

498 phi_exp = phi.real[np.newaxis, :, :] 

499 else: 

500 active = phi.get_active_bases() 

501 n_bases = active[-1] if active else self.model.n_bases 

502 deriv_order = 2 * self.model.n_order 

503 phi_exp = self._expand_derivs(phi, n_bases, deriv_order) 

504 

505 K = self._build_K(phi_exp, phi, n_bases) 

506 K.flat[::K.shape[0] + 1] += sigma_n_sq 

507 K += self.model.sigma_data ** 2 

508 

509 try: 

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

511 alpha_v = cho_solve((L, low), self.model.y_train) 

512 N = len(self.model.y_train) 

513 

514 # NLL 

515 nll = (0.5 * np.dot(self.model.y_train, alpha_v) 

516 + np.sum(np.log(np.diag(L))) 

517 + 0.5 * N * np.log(2 * np.pi)) 

518 

519 # W matrix for gradient (reuse same Cholesky) 

520 K_inv = cho_solve((L, low), np.eye(N)) 

521 W = np.empty((N, N)) 

522 _subtract_outer(K_inv, alpha_v, W) 

523 except Exception: 

524 return 1e6, np.zeros(len(x0)) 

525 

526 # Cache for fast prediction (reused by gddegp.predict) 

527 self.model._cached_L = L 

528 self.model._cached_low = low 

529 self.model._cached_alpha = alpha_v 

530 # Preserve any explicitly-set n_bases (e.g. function-only training 

531 # with n_order=0 but n_bases>0 reserved for derivative prediction). 

532 self.model._cached_n_bases = max(n_bases, self.model.n_bases) 

533 self.model._cached_params = x0.copy() 

534 

535 # --- gradient from W (no second kernel build / Cholesky) --- 

536 grad = np.zeros(len(x0)) 

537 use_fast = self._kernel_plan is not None 

538 base_shape = phi.shape 

539 

540 W_proj = None 

541 if use_fast and self.model.n_order > 0: 

542 from math import comb 

543 ndir = comb(n_bases + deriv_order, deriv_order) 

544 proj_shape = (ndir, base_shape[0], base_shape[1]) 

545 if self._W_proj_buf is None or self._W_proj_shape != proj_shape: 

546 self._W_proj_buf = np.empty(proj_shape) 

547 self._W_proj_shape = proj_shape 

548 W_proj = self._W_proj_buf 

549 plan = self._kernel_plan 

550 row_off = plan.get('row_offsets_abs', plan['row_offsets'] + base_shape[0]) 

551 col_off = plan.get('col_offsets_abs', plan['col_offsets'] + base_shape[1]) 

552 utils._project_W_to_phi_space( 

553 W, W_proj, base_shape[0], base_shape[1], 

554 plan['fd_flat_indices'], plan['df_flat_indices'], 

555 plan['dd_flat_indices'], 

556 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'], 

557 plan['n_deriv_types'], row_off, col_off, 

558 ) 

559 

560 _use_vdot_fused = W_proj is not None and hasattr(phi, 'vdot_expand_fast') 

561 if _use_vdot_fused: 

562 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order) 

563 

564 def _gc(dphi): 

565 if _use_vdot_fused: 

566 return 0.5 * dphi.vdot_expand_fast(_vdot_factors, W_proj) 

567 if self.model.n_order == 0: 

568 dphi_exp = dphi.real[np.newaxis, :, :] 

569 else: 

570 dphi_exp = self._expand_derivs(dphi, n_bases, deriv_order) 

571 if W_proj is not None: 

572 dphi_3d = dphi_exp.reshape(W_proj.shape) 

573 return 0.5 * np.vdot(W_proj, dphi_3d) 

574 elif use_fast: 

575 dphi_3d = dphi_exp.reshape(dphi_exp.shape[0], base_shape[0], base_shape[1]) 

576 dK = utils.rbf_kernel_fast(dphi_3d, self._kernel_plan, out=self._dK_buf) 

577 return 0.5 * np.vdot(W, dK) 

578 else: 

579 dK = utils.rbf_kernel( 

580 dphi, dphi_exp, self.model.n_order, n_bases, 

581 self.model.flattened_der_indices, 

582 index=self.model.derivative_locations, 

583 ) 

584 return 0.5 * np.vdot(W, dK) 

585 

586 grad[-2] = _gc(oti.mul(2.0 * ln10, phi)) 

587 grad[-1] = ln10 * sigma_n_sq * np.trace(W) 

588 

589 if kernel == 'SE': 

590 if kernel_type == 'anisotropic': 

591 ell = 10.0 ** x0[:D] 

592 if hasattr(phi, 'fused_scale_sq_mul'): 

593 dphi_buf = oti.zeros(phi.shape) 

594 for d in range(D): 

595 dphi_buf.fused_scale_sq_mul(diffs[d], phi, -ln10 * ell[d] ** 2) 

596 grad[d] = _gc(dphi_buf) 

597 else: 

598 for d in range(D): 

599 grad[d] = _gc(oti.mul(-ln10 * ell[d] ** 2, 

600 oti.mul(oti.mul(diffs[d], diffs[d]), phi))) 

601 else: 

602 ell = 10.0 ** float(x0[0]) 

603 if hasattr(phi, 'fused_sum_sq'): 

604 sum_sq = oti.zeros(phi.shape) 

605 sum_sq.fused_sum_sq(diffs) 

606 else: 

607 sum_sq = oti.mul(diffs[0], diffs[0])