Coverage for jetgp/full_ddegp/optimizer.py: 69%

1import numpy as np 

2 

3from scipy.linalg import cho_solve, cho_factor 

4from jetgp.full_ddegp import ddegp_utils as utils 

5from line_profiler import profile 

6import jetgp.utils as gen_utils 

7from jetgp.hyperparameter_optimizers import OPTIMIZERS 

8from jetgp.utils import matern_kernel_grad_builder 

9 

10class Optimizer: 

11 """ 

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

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

14 

15 Parameters 

16 ---------- 

17 model : object 

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

19 and kernel configuration. 

20 """ 

21 

22 def __init__(self, model): 

23 self.model = model 

24 self._kernel_plan = None 

25 self._deriv_buf = None 

26 self._deriv_buf_shape = None 

27 self._deriv_buf_ndir = None 

28 self._deriv_factors = None 

29 self._deriv_factors_key = None 

30 self._K_buf = None 

31 self._dK_buf = None 

32 self._kernel_buf_size = None 

33 self._W_proj_buf = None 

34 self._W_proj_shape = None 

35 

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

37 if self._deriv_buf_ndir is None: 

38 from math import comb 

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

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

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

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

43 self._deriv_buf_shape = shape 

44 return self._deriv_buf 

45 

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

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

48 if hasattr(phi, 'get_all_derivs_fast'): 

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

50 factors = self._get_deriv_factors(n_bases, deriv_order) 

51 return phi.get_all_derivs_fast(factors, buf) 

52 return phi.get_all_derivs(n_bases, deriv_order) 

53 

54 @staticmethod 

55 def _enum_factors(max_basis, ordi): 

56 from math import factorial 

57 from collections import Counter 

58 if ordi == 1: 

59 for _ in range(max_basis): 

60 yield 1.0 

61 return 

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

63 if ordi == 2: 

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

65 counts = Counter((i, last)) 

66 f = 1 

67 for c in counts.values(): 

68 f *= factorial(c) 

69 yield float(f) 

70 else: 

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

72 counts = dict(prefix_counts) 

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

74 f = 1 

75 for c in counts.values(): 

76 f *= factorial(c) 

77 yield float(f) 

78 

79 @staticmethod 

80 def _enum_factors_with_counts(max_basis, ordi): 

81 from math import factorial 

82 from collections import Counter 

83 if ordi == 1: 

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

85 yield 1.0, {i: 1} 

86 return 

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

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

89 counts = dict(prefix_counts) 

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

91 f = 1 

92 for c in counts.values(): 

93 f *= factorial(c) 

94 yield float(f), counts 

95 

96 def _get_deriv_factors(self, n_bases, order): 

97 key = (n_bases, order) 

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

99 return self._deriv_factors 

100 factors = [1.0] 

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

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

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

104 self._deriv_factors_key = key 

105 return self._deriv_factors 

106 

107 def _ensure_kernel_plan(self, n_bases): 

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

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

110 return 

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

112 self._kernel_plan = None 

113 return 

114 self._kernel_plan = utils.precompute_kernel_plan( 

115 self.model.n_order, n_bases, 

116 self.model.flattened_der_indices, 

117 self.model.powers, 

118 self.model.derivative_locations, 

119 ) 

120 self._kernel_plan_n_bases = n_bases 

121 self._K_buf = None 

122 self._dK_buf = None 

123 self._kernel_buf_size = None 

124 

125 def _ensure_kernel_bufs(self, n_rows_func): 

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

127 if self._kernel_plan is None: 

128 return 

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

130 if self._kernel_buf_size != total: 

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

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

133 self._kernel_buf_size = total 

134 if 'row_offsets_abs' not in self._kernel_plan: 

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

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

137 

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

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

140 self._ensure_kernel_plan(n_bases) 

141 if self._kernel_plan is not None: 

142 base_shape = phi.shape 

143 self._ensure_kernel_bufs(base_shape[0]) 

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

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

146 return utils.rbf_kernel( 

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

148 self.model.flattened_der_indices, self.model.powers, 

149 self.model.derivative_locations, 

150 ) 

151 

152 @profile 

153 def negative_log_marginal_likelihood(self, x0): 

154 """ 

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

156 

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

158 

159 Parameters 

160 ---------- 

161 x0 : ndarray 

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

163 

164 Returns 

165 ------- 

166 float 

167 Value of the negative log marginal likelihood. 

168 """ 

169 ell = x0[:-1] 

170 sigma_n = x0[-1] 

171 llhood = 0 

172 diffs = self.model.differences_by_dim 

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

174 n_bases = phi.get_active_bases()[-1] 

175 

176 # Extract ALL derivative components 

177 deriv_order = 2 * self.model.n_order 

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

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

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

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

182 

183 try: 

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

185 alpha = cho_solve( 

186 (L, low), 

187 self.model.y_train 

188 ) 

189 

190 # Cache for fast prediction 

191 self.model._cached_L = L 

192 self.model._cached_low = low 

193 self.model._cached_alpha = alpha 

194 self.model._cached_n_bases_rays = n_bases 

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

196 

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

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

199 complexity = log_det_K 

200 N = len(self.model.y_train) 

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

202 return data_fit + complexity + const 

203 except Exception: 

204 return 1e6 

205 

206 def nll_wrapper(self, x0): 

207 """ 

208 Wrapper function to compute NLL for optimizer. 

209 

210 Parameters 

211 ---------- 

212 x0 : ndarray 

213 Hyperparameter vector. 

214 

215 Returns 

216 ------- 

217 float 

218 NLL evaluated at x0. 

219 """ 

220 return self.negative_log_marginal_likelihood(x0) 

221 

222 def nll_grad(self, x0): 

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

224 ln10 = np.log(10.0) 

225 

226 kernel = self.model.kernel 

227 kernel_type = self.model.kernel_type 

228 D = len(self.model.differences_by_dim) 

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

230 diffs = self.model.differences_by_dim 

231 oti = self.model.kernel_factory.oti 

232 

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

234 n_bases = phi.get_active_bases()[-1] 

235 deriv_order = 2 * self.model.n_order 

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

237 

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

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

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

241 

242 try: 

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

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

245 N = len(self.model.y_train) 

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

247 W = K_inv - np.outer(alpha_v, alpha_v) 

248 except Exception: 

249 return np.zeros(len(x0)) 

250 

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

252 use_fast = self._kernel_plan is not None 

253 base_shape = phi.shape 

254 

255 W_proj = None 

256 if use_fast: 

257 from math import comb 

258 ndir = comb(n_bases + deriv_order, deriv_order) 

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

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

261 self._W_proj_buf = np.empty(proj_shape) 

262 self._W_proj_shape = proj_shape 

263 W_proj = self._W_proj_buf 

264 plan = self._kernel_plan 

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

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

267 utils._project_W_to_phi_space( 

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

269 plan['fd_flat_indices'], plan['df_flat_indices'], 

270 plan['dd_flat_indices'], 

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

272 plan['signs'], plan['n_deriv_types'], row_off, col_off, 

273 ) 

274 

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

276 if _use_vdot_fused: 

277 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order) 

278 

279 def _gc(dphi): 

280 if _use_vdot_fused: 

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

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

283 if W_proj is not None: 

284 dphi_3d = dphi_exp.reshape(W_proj.shape) 

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

286 elif use_fast: 

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

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

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

290 else: 

291 dK = utils.rbf_kernel( 

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

293 self.model.flattened_der_indices, self.model.powers, 

294 self.model.derivative_locations, 

295 ) 

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

297 

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

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

300 

301 if kernel == 'SE': 

302 if kernel_type == 'anisotropic': 

303 ell = 10.0 ** x0[:D] 

304 if hasattr(phi, 'fused_scale_sq_mul'): 

305 dphi_buf = oti.zeros(phi.shape) 

306 for d in range(D): 

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

308 grad[d] = _gc(dphi_buf) 

309 else: 

310 for d in range(D): 

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

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

313 else: 

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

315 if hasattr(phi, 'fused_sum_sq'): 

316 sum_sq = oti.zeros(phi.shape) 

317 sum_sq.fused_sum_sq(diffs) 

318 else: 

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

320 for d in range(1, D): 

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

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

323 

324 elif kernel == 'RQ': 

325 if kernel_type == 'anisotropic': 

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

327 else: 

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

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

330 if hasattr(phi, 'fused_sqdist'): 

331 r2 = oti.zeros(phi.shape) 

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

333 r2.fused_sqdist(diffs, ell_sq) 

334 else: 

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

336 for d in range(1, D): 

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

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

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

340 phi_over_base = oti.mul(phi, inv_base) 

341 if kernel_type == 'anisotropic': 

342 if hasattr(phi, 'fused_scale_sq_mul'): 

343 dphi_buf = oti.zeros(phi.shape) 

344 for d in range(D): 

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

346 grad[d] = _gc(dphi_buf) 

347 else: 

348 for d in range(D): 

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

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

351 else: 

352 if hasattr(phi, 'fused_sum_sq'): 

353 sum_sq = oti.zeros(phi.shape) 

354 sum_sq.fused_sum_sq(diffs) 

355 else: 

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

357 for d in range(1, D): 

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

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

360 log_base = oti.log(base) 

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

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

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

364 

365 elif kernel == 'SineExp': 

366 if kernel_type == 'anisotropic': 

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

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

369 else: 

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

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

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

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

374 if kernel_type == 'anisotropic': 

375 if hasattr(phi, 'fused_scale_sq_mul'): 

376 dphi_buf = oti.zeros(phi.shape) 

377 for d in range(D): 

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

379 grad[d] = _gc(dphi_buf) 

380 else: 

381 for d in range(D): 

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

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

384 for d in range(D): 

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

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

387 oti.mul(sc, phi))) 

388 else: 

389 if hasattr(phi, 'fused_sum_sq'): 

390 ss = oti.zeros(phi.shape) 

391 ss.fused_sum_sq(sin_d) 

392 else: 

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

394 for d in range(1, D): 

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

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

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

398 for d in range(1, D): 

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

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

401 oti.mul(scd, phi))) 

402 

403 elif kernel == 'Matern': 

404 kf = self.model.kernel_factory 

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

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

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

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

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

410 _eps = 1e-10 

411 if hasattr(phi, 'fused_sqdist'): 

412 r2 = oti.zeros(phi.shape) 

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

414 r2.fused_sqdist(diffs, ell_sq) 

415 else: 

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

417 for d in range(1, D): 

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

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

420 f_prime_r = kf._matern_grad_prebuild(r_oti) 

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

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

423 if kernel_type == 'anisotropic': 

424 if hasattr(phi, 'fused_scale_sq_mul'): 

425 dphi_buf = oti.zeros(phi.shape) 

426 for d in range(D): 

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

428 grad[d] = _gc(dphi_buf) 

429 else: 

430 for d in range(D): 

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

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

433 grad[d] = _gc(dphi_d) 

434 else: 

435 if hasattr(phi, 'fused_sum_sq'): 

436 sum_dsq = oti.zeros(phi.shape) 

437 sum_dsq.fused_sum_sq(diffs) 

438 else: 

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

440 for d in range(1, D): 

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

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

443 grad[0] = _gc(dphi_e) 

444 

445 return grad 

446 

447 def nll_and_grad(self, x0): 

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

449 ln10 = np.log(10.0) 

450 

451 kernel = self.model.kernel 

452 kernel_type = self.model.kernel_type 

453 D = len(self.model.differences_by_dim) 

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

455 diffs = self.model.differences_by_dim 

456 oti = self.model.kernel_factory.oti 

457 

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

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

460 if self.model.n_order == 0: 

461 n_bases = 0 

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

463 else: 

464 n_bases = phi.get_active_bases()[-1] 

465 deriv_order = 2 * self.model.n_order 

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

467 

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

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

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

471 

472 try: 

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

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

475 N = len(self.model.y_train) 

476 

477 # NLL 

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

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

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

481 

482 # W matrix for gradient (reuse same Cholesky) 

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

484 W = K_inv - np.outer(alpha_v, alpha_v) 

485 except Exception: 

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

487 

488 # Cache for fast prediction (reused by ddegp.predict) 

489 self.model._cached_L = L 

490 self.model._cached_low = low 

491 self.model._cached_alpha = alpha_v 

492 self.model._cached_n_bases_rays = n_bases 

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

494 

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

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

497 use_fast = self._kernel_plan is not None 

498 base_shape = phi.shape 

499 deriv_order_gc = 2 * self.model.n_order 

500 

501 W_proj = None 

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

503 from math import comb 

504 ndir = comb(n_bases + deriv_order_gc, deriv_order_gc) 

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

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

507 self._W_proj_buf = np.empty(proj_shape) 

508 self._W_proj_shape = proj_shape 

509 W_proj = self._W_proj_buf 

510 plan = self._kernel_plan 

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

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

513 utils._project_W_to_phi_space( 

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

515 plan['fd_flat_indices'], plan['df_flat_indices'], 

516 plan['dd_flat_indices'], 

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

518 plan['signs'], plan['n_deriv_types'], row_off, col_off, 

519 ) 

520 

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

522 if _use_vdot_fused: 

523 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order_gc) 

524 

525 def _gc(dphi): 

526 if _use_vdot_fused: 

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

528 if self.model.n_order == 0: 

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

530 else: 

531 dphi_exp = self._expand_derivs(dphi, n_bases, deriv_order_gc) 

532 if W_proj is not None: 

533 dphi_3d = dphi_exp.reshape(W_proj.shape) 

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

535 elif use_fast: 

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

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

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

539 else: 

540 dK = utils.rbf_kernel( 

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

542 self.model.flattened_der_indices, self.model.powers, 

543 self.model.derivative_locations, 

544 ) 

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

546 

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

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

549 

550 if kernel == 'SE': 

551 if kernel_type == 'anisotropic': 

552 ell = 10.0 ** x0[:D] 

553 if hasattr(phi, 'fused_scale_sq_mul'): 

554 dphi_buf = oti.zeros(phi.shape) 

555 for d in range(D): 

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

557 grad[d] = _gc(dphi_buf) 

558 else: 

559 for d in range(D): 

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

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

562 else: 

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

564 if hasattr(phi, 'fused_sum_sq'): 

565 sum_sq = oti.zeros(phi.shape) 

566 sum_sq.fused_sum_sq(diffs) 

567 else: 

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

569 for d in range(1, D): 

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

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

572 

573 elif kernel == 'RQ': 

574 if kernel_type == 'anisotropic': 

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

576 else: 

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

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

579 if hasattr(phi, 'fused_sqdist'): 

580 r2 = oti.zeros(phi.shape) 

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

582 r2.fused_sqdist(diffs, ell_sq) 

583 else: 

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

585 for d in range(1, D): 

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

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

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

589 phi_over_base = oti.mul(phi, inv_base) 

590 if kernel_type == 'anisotropic':