Coverage for jetgp/wdegp/optimizer.py: 80%

1import numpy as np 

2from scipy.linalg import cho_solve, cho_factor 

3from line_profiler import profile 

4import jetgp.utils as gen_utils 

5from jetgp.hyperparameter_optimizers import OPTIMIZERS 

6from jetgp.utils import matern_kernel_grad_builder 

7 

8 

9class Optimizer: 

10 """ 

11 Optimizer class for fitting the hyperparameters of a weighted derivative-enhanced GP model (wDEGP) 

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

13 

14 Supports DEGP, DDEGP, and GDDEGP modes. 

15 

16 Attributes 

17 ---------- 

18 model : object 

19 Instance of a weighted derivative-enhanced GP model (wDEGP) with attributes: 

20 x_train, y_train, n_order, n_bases, der_indices, index, bounds, submodel_type, etc. 

21 """ 

22 

23 def __init__(self, model): 

24 """ 

25 Parameters 

26 ---------- 

27 model : object 

28 An instance of a wDEGP model containing training data, hyperparameter bounds, 

29 and other model-specific structures required for kernel computation. 

30 """ 

31 self.model = model 

32 

33 # Import the appropriate utils module based on submodel_type 

34 self._setup_utils() 

35 

36 # Sparse fused OTI functions only valid for DEGP (axis-aligned diffs) 

37 self._sparse_safe = getattr(self.model, 'submodel_type', 'degp') == 'degp' 

38 

39 # Precompute kernel plans (structural info that never changes) 

40 self._kernel_plans = None # lazily initialized on first NLL call 

41 self._deriv_buf = None 

42 self._deriv_buf_shape = None 

43 self._deriv_buf_ndir = None 

44 self._deriv_factors = None 

45 self._deriv_factors_key = None 

46 self._K_bufs = None # per-submodel pre-allocated K buffers 

47 self._dK_bufs = None # per-submodel pre-allocated dK buffers 

48 self._W_proj_buf = None 

49 self._W_proj_shape = None 

50 

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

52 if self._deriv_buf_ndir is None: 

53 from math import comb 

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

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

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

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

58 self._deriv_buf_shape = shape 

59 return self._deriv_buf 

60 

61 @profile 

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

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

64 if hasattr(phi, 'get_all_derivs_fast'): 

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

66 factors = self._get_deriv_factors(n_bases, deriv_order) 

67 return phi.get_all_derivs_fast(factors, buf) 

68 return phi.get_all_derivs(n_bases, deriv_order) 

69 

70 @staticmethod 

71 def _enum_factors(max_basis, ordi): 

72 from math import factorial 

73 from collections import Counter 

74 if ordi == 1: 

75 for _ in range(max_basis): 

76 yield 1.0 

77 return 

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

79 if ordi == 2: 

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

81 counts = Counter((i, last)) 

82 f = 1 

83 for c in counts.values(): 

84 f *= factorial(c) 

85 yield float(f) 

86 else: 

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

88 counts = dict(prefix_counts) 

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

90 f = 1 

91 for c in counts.values(): 

92 f *= factorial(c) 

93 yield float(f) 

94 

95 @staticmethod 

96 def _enum_factors_with_counts(max_basis, ordi): 

97 from math import factorial 

98 from collections import Counter 

99 if ordi == 1: 

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

101 yield 1.0, {i: 1} 

102 return 

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

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

105 counts = dict(prefix_counts) 

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

107 f = 1 

108 for c in counts.values(): 

109 f *= factorial(c) 

110 yield float(f), counts 

111 

112 def _get_deriv_factors(self, n_bases, order): 

113 key = (n_bases, order) 

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

115 return self._deriv_factors 

116 factors = [1.0] 

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

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

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

120 self._deriv_factors_key = key 

121 return self._deriv_factors 

122 

123 def _setup_utils(self): 

124 """Set up the correct utils module based on submodel_type.""" 

125 submodel_type = getattr(self.model, 'submodel_type', 'degp') 

126 

127 if submodel_type == 'degp': 

128 from jetgp.wdegp import wdegp_utils 

129 self.utils = wdegp_utils 

130 self._uses_signs = True 

131 elif submodel_type == 'ddegp': 

132 from jetgp.full_ddegp import wddegp_utils 

133 self.utils = wddegp_utils 

134 self._uses_signs = True 

135 elif submodel_type == 'gddegp': 

136 from jetgp.full_gddegp import wgddegp_utils 

137 self.utils = wgddegp_utils 

138 self._uses_signs = False 

139 else: 

140 # Default to degp 

141 from jetgp.wdegp import wdegp_utils 

142 self.utils = wdegp_utils 

143 self._uses_signs = True 

144 

145 def _ensure_kernel_plans(self, n_bases): 

146 """Lazily precompute kernel plans for all submodels (once per n_bases).""" 

147 if self._kernel_plans is not None and self._kernel_plans_n_bases == n_bases: 

148 return 

149 if not hasattr(self.utils, 'precompute_kernel_plan'): 

150 self._kernel_plans = None 

151 return 

152 plans = [] 

153 index = self.model.derivative_locations 

154 for i in range(len(index)): 

155 plan = self.utils.precompute_kernel_plan( 

156 self.model.n_order, n_bases, 

157 self.model.flattened_der_indices[i], 

158 self.model.powers[i], 

159 index[i], 

160 ) 

161 plans.append(plan) 

162 self._kernel_plans = plans 

163 self._kernel_plans_n_bases = n_bases 

164 # Reset buffers when plans change 

165 self._K_bufs = None 

166 self._dK_bufs = None 

167 

168 def _ensure_kernel_bufs(self, n_rows_func): 

169 """Pre-allocate reusable K and dK buffers for each submodel.""" 

170 if self._kernel_plans is None: 

171 return 

172 if self._K_bufs is not None: 

173 return # already allocated 

174 self._K_bufs = [] 

175 self._dK_bufs = [] 

176 for plan in self._kernel_plans: 

177 total = n_rows_func + plan['n_pts_with_derivs'] 

178 self._K_bufs.append(np.empty((total, total))) 

179 self._dK_bufs.append(np.empty((total, total))) 

180 if 'row_offsets_abs' not in plan: 

181 plan['row_offsets_abs'] = plan['row_offsets'] + n_rows_func 

182 plan['col_offsets_abs'] = plan['col_offsets'] + n_rows_func 

183 

184 @profile 

185 def negative_log_marginal_likelihood( 

186 self, 

187 x0, 

188 x_train, 

189 y_train, 

190 n_order, 

191 n_bases, 

192 der_indices, 

193 index, 

194 ): 

195 """ 

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

197 

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

199 

200 Parameters 

201 ---------- 

202 x0 : ndarray 

203 Log-scaled hyperparameter vector, where the last entry is log10(sigma_n). 

204 x_train : list of ndarrays 

205 Input training points (unused inside loop, included for general interface). 

206 y_train : list of ndarrays 

207 List of function and derivative training values for each submodel. 

208 n_order : int 

209 Maximum order of derivatives used. 

210 n_bases : int 

211 Number of Taylor bases used in the expansion. 

212 der_indices : list 

213 Multi-index derivative information. 

214 index : list of lists 

215 Indices partitioning the training data into submodels (derivative_locations). 

216 

217 Returns 

218 ------- 

219 float 

220 The computed negative log marginal likelihood. 

221 """ 

222 ell = x0[:-1] 

223 sigma_n = x0[-1] 

224 llhood = 0 

225 # ell[0] = 0 

226 # ell[1] = 0 

227 # ell[2] = 0 

228 # sigma_n = -16 

229 diffs = self.model.differences_by_dim 

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

231 if self.model.n_order == 0: 

232 n_bases = 0 

233 phi_exp = phi.real 

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

235 else: 

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

237 

238 # Extract ALL derivative components 

239 deriv_order = 2 * n_order 

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

241 

242 # Ensure kernel plans are precomputed 

243 self._ensure_kernel_plans(n_bases) 

244 use_fast = self._kernel_plans is not None 

245 

246 # Pre-reshape phi_exp to 3D once 

247 if use_fast: 

248 base_shape = phi.shape 

249 self._ensure_kernel_bufs(base_shape[0]) 

250 phi_exp_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1]) 

251 

252 for i in range(len(index)): 

253 y_train_sub = y_train[i] 

254 

255 if use_fast: 

256 K = self.utils.rbf_kernel_fast(phi_exp_3d, self._kernel_plans[i], out=self._K_bufs[i]) 

257 else: 

258 K = self.utils.rbf_kernel( 

259 phi, phi_exp, n_order, n_bases, 

260 self.model.flattened_der_indices[i], 

261 self.model.powers[i], index=index[i] 

262 ) 

263 

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

265 

266 try: 

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

268 alpha = cho_solve( 

269 (L, low), 

270 y_train_sub 

271 ) 

272 

273 data_fit = 0.5 * np.dot(y_train_sub.flatten(), alpha.flatten()) 

274 log_det = np.sum(np.log(np.diag(L))) 

275 const = 0.5 * len(y_train_sub) * np.log(2 * np.pi) 

276 

277 llhood += data_fit + log_det + const 

278 except np.linalg.LinAlgError: 

279 llhood += 1e6 # Penalize badly conditioned matrices 

280 

281 return llhood 

282 

283 def nll_wrapper(self, x0): 

284 """ 

285 Wrapper for NLL function to fit PSO optimizer interface. 

286 

287 Parameters 

288 ---------- 

289 x0 : ndarray 

290 Hyperparameter vector. 

291 

292 Returns 

293 ------- 

294 float 

295 Computed NLL value. 

296 """ 

297 return self.negative_log_marginal_likelihood( 

298 x0, 

299 self.model.x_train, 

300 self.model.y_train_normalized, 

301 self.model.n_order, 

302 self.model.n_bases, 

303 self.model.der_indices, 

304 self.model.derivative_locations, 

305 ) 

306 

307 def nll_grad(self, x0): 

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

309 ln10 = np.log(10.0) 

310 

311 kernel = self.model.kernel 

312 kernel_type = self.model.kernel_type 

313 D = len(self.model.differences_by_dim) 

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

315 diffs = self.model.differences_by_dim 

316 oti = self.model.kernel_factory.oti 

317 index = self.model.derivative_locations 

318 

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

320 if self.model.n_order == 0: 

321 n_bases = 0 

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

323 else: 

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

325 deriv_order = 2 * self.model.n_order 

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

327 

328 # Ensure kernel plans are precomputed 

329 self._ensure_kernel_plans(n_bases) 

330 use_fast = self._kernel_plans is not None 

331 

332 # Pre-reshape phi_exp to 3D once 

333 if use_fast: 

334 base_shape = phi.shape 

335 self._ensure_kernel_bufs(base_shape[0]) 

336 phi_exp_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1]) 

337 

338 # Build per-submodel W matrices 

339 W_list = [] 

340 for i in range(len(index)): 

341 y_train_sub = self.model.y_train_normalized[i] 

342 if use_fast: 

343 K = self.utils.rbf_kernel_fast(phi_exp_3d, self._kernel_plans[i], out=self._K_bufs[i]) 

344 else: 

345 K = self.utils.rbf_kernel( 

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

347 self.model.flattened_der_indices[i], 

348 self.model.powers[i], index=index[i] 

349 ) 

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

351 try: 

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

353 alpha_v = cho_solve((L, low), y_train_sub) 

354 N = len(y_train_sub) 

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

356 W_list.append(K_inv - np.outer(alpha_v, alpha_v)) 

357 except Exception: 

358 return np.zeros(len(x0)) 

359 

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

361 

362 # Precompute W projected into phi_exp space (sum over submodels) 

363 W_proj = None 

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

365 from math import comb 

366 ndir = comb(n_bases + deriv_order, deriv_order) 

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

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

369 self._W_proj_buf = np.empty(proj_shape) 

370 self._W_proj_shape = proj_shape 

371 W_proj = self._W_proj_buf 

372 W_proj[:] = 0.0 

373 for i in range(len(index)): 

374 plan = self._kernel_plans[i] 

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

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

377 args = [ 

378 W_list[i], W_proj, base_shape[0], base_shape[1], 

379 plan['fd_flat_indices'], plan['df_flat_indices'], 

380 plan['dd_flat_indices'], 

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

382 ] 

383 if self._uses_signs: 

384 args.append(plan['signs']) 

385 args.extend([plan['n_deriv_types'], row_off, col_off]) 

386 self.utils._project_W_to_phi_space_accum(*args) 

387 

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

389 FW_T = None 

390 if _use_vdot_fused: 

391 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order) 

392 _vdot_arr = np.asarray(_vdot_factors) 

393 ndir_d = len(_vdot_arr) 

394 FW_T = np.empty((base_shape[0] * base_shape[1], ndir_d)) 

395 np.multiply(W_proj.reshape(ndir_d, -1).T, _vdot_arr, out=FW_T) 

396 

397 def _gc(dphi): 

398 if _use_vdot_fused: 

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

400 if self.model.n_order == 0: 

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

402 else: 

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

404 if W_proj is not None: 

405 dphi_3d = dphi_exp.reshape(W_proj.shape) 

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

407 elif use_fast: 

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

409 total = 0.0 

410 for i in range(len(index)): 

411 dK = self.utils.rbf_kernel_fast(dphi_3d, self._kernel_plans[i], out=self._dK_bufs[i]) 

412 total += np.vdot(W_list[i], dK) 

413 return 0.5 * total 

414 else: 

415 total = 0.0 

416 for i in range(len(index)): 

417 dK = self.utils.rbf_kernel( 

418 dphi, dphi_exp, 

419 self.model.n_order, n_bases, 

420 self.model.flattened_der_indices[i], 

421 self.model.powers[i], 

422 index=index[i], 

423 ) 

424 total += np.vdot(W_list[i], dK) 

425 return 0.5 * total 

426 

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

428 grad[-1] = ln10 * sigma_n_sq * sum(np.trace(W) for W in W_list) 

429 

430 if kernel == 'SE': 

431 if kernel_type == 'anisotropic': 

432 ell = 10.0 ** x0[:D] 

433 if self._sparse_safe and _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'): 

434 scales = np.array([-ln10 * ell[d] ** 2 for d in range(D)]) 

435 grad_buf = np.zeros(D) 

436 phi.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T) 

437 grad[:D] = grad_buf 

438 elif self._sparse_safe and hasattr(phi, 'fused_scale_sq_mul_sparse'): 

439 dphi_buf = oti.zeros(phi.shape) 

440 for d in range(D): 

441 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], phi, -ln10 * ell[d] ** 2, d) 

442 grad[d] = _gc(dphi_buf) 

443 elif hasattr(phi, 'fused_scale_sq_mul'): 

444 dphi_buf = oti.zeros(phi.shape) 

445 for d in range(D): 

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

447 grad[d] = _gc(dphi_buf) 

448 else: 

449 for d in range(D): 

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

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

452 else: 

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

454 if self._sparse_safe and hasattr(phi, 'fused_sum_sq_sparse'): 

455 sum_sq = oti.zeros(phi.shape) 

456 sum_sq.fused_sum_sq_sparse(diffs) 

457 elif hasattr(phi, 'fused_sum_sq'): 

458 sum_sq = oti.zeros(phi.shape) 

459 sum_sq.fused_sum_sq(diffs) 

460 else: 

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

462 for d in range(1, D): 

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

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

465 

466 elif kernel == 'RQ': 

467 if kernel_type == 'anisotropic': 

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

469 else: 

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

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

472 if self._sparse_safe and hasattr(phi, 'fused_sqdist_sparse'): 

473 r2 = oti.zeros(phi.shape) 

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

475 r2.fused_sqdist_sparse(diffs, ell_sq) 

476 elif hasattr(phi, 'fused_sqdist'): 

477 r2 = oti.zeros(phi.shape) 

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

479 r2.fused_sqdist(diffs, ell_sq) 

480 else: 

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

482 for d in range(1, D): 

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

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

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

486 phi_over_base = oti.mul(phi, inv_base) 

487 if kernel_type == 'anisotropic': 

488 if self._sparse_safe and _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'): 

489 scales = np.array([-ln10 * ell[d] ** 2 for d in range(D)]) 

490 grad_buf = np.zeros(D) 

491 phi_over_base.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T) 

492 grad[:D] = grad_buf 

493 elif self._sparse_safe and hasattr(phi, 'fused_scale_sq_mul_sparse'): 

494 dphi_buf = oti.zeros(phi.shape) 

495 for d in range(D): 

496 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], phi_over_base, -ln10 * ell[d] ** 2, d) 

497 grad[d] = _gc(dphi_buf) 

498 elif hasattr(phi, 'fused_scale_sq_mul'): 

499 dphi_buf = oti.zeros(phi.shape) 

500 for d in range(D): 

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

502 grad[d] = _gc(dphi_buf) 

503 else: 

504 for d in range(D): 

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

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

507 else: 

508 if self._sparse_safe and hasattr(phi, 'fused_sum_sq_sparse'): 

509 sum_sq = oti.zeros(phi.shape) 

510 sum_sq.fused_sum_sq_sparse(diffs) 

511 elif hasattr(phi, 'fused_sum_sq'): 

512 sum_sq = oti.zeros(phi.shape) 

513 sum_sq.fused_sum_sq(diffs) 

514 else: 

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

516 for d in range(1, D): 

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

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

519 log_base = oti.log(base) 

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

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

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

523 

524 elif kernel == 'SineExp': 

525 if kernel_type == 'anisotropic': 

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

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

528 else: 

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

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

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

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

533 if kernel_type == 'anisotropic': 

534 if hasattr(phi, 'fused_scale_sq_mul'): 

535 dphi_buf = oti.zeros(phi.shape) 

536 for d in range(D): 

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

538 grad[d] = _gc(dphi_buf) 

539 else: 

540 for d in range(D): 

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

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

543 for d in range(D): 

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

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

546 oti.mul(sc, phi))) 

547 else: 

548 if hasattr(phi, 'fused_sum_sq'): 

549 ss = oti.zeros(phi.shape) 

550 ss.fused_sum_sq(sin_d) 

551 else: 

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

553 for d in range(1, D): 

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

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

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

557 for d in range(1, D): 

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

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

560 oti.mul(scd, phi))) 

561 

562 elif kernel == 'Matern': 

563 kf = self.model.kernel_factory 

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

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

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

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

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

569 _eps = 1e-10 

570 if self._sparse_safe and hasattr(phi, 'fused_sqdist_sparse'): 

571 r2 = oti.zeros(phi.shape) 

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

573 r2.fused_sqdist_sparse(diffs, ell_sq) 

574 elif hasattr(phi, 'fused_sqdist'): 

575 r2 = oti.zeros(phi.shape) 

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

577 r2.fused_sqdist(diffs, ell_sq) 

578 else: 

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

580 for d in range(1, D): 

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

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

583 f_prime_r = kf._matern_grad_prebuild(r_oti) 

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

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

586 if kernel_type == 'anisotropic': 

587 if self._sparse_safe and _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'): 

588 scales = np.array([ln10 * ell[d] ** 2 for d in range(D)]) 

589 grad_buf = np.zeros(D) 

590 base_matern.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T) 

591 grad[:D] = grad_buf 

592 elif self._sparse_safe and hasattr(phi, 'fused_scale_sq_mul_sparse'): 

593 dphi_buf = oti.zeros(phi.shape) 

594 for d in range(D): 

595 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], base_matern, ln10 * ell[d] ** 2, d) 

596 grad[d] = _gc(dphi_buf) 

597 elif hasattr(phi, 'fused_scale_sq_mul'): 

598 dphi_buf = oti.zeros(phi.shape) 

599 for d in range(D): 

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

601 grad[d] = _gc(dphi_buf) 

602 else: 

603 for d in range(D): 

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

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

606 grad[d] = _gc(dphi_d) 

607 else: 

608 if self._sparse_safe and hasattr(phi, 'fused_sum_sq_sparse'): 

609 sum_dsq = oti.zeros(phi.shape) 

610 sum_dsq.fused_sum_sq_sparse(diffs) 

611 elif hasattr(phi, 'fused_sum_sq'): 

612 sum_dsq = oti.zeros(phi.shape) 

613 sum_dsq.fused_sum_sq(diffs) 

614 else: 

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

616 for d in range(1, D): 

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

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

619 grad[0] = _gc(dphi_e) 

620 

621 return grad 

622