Coverage for jetgp/full_degp/optimizer.py: 76%

1import numpy as np 

2from scipy.linalg import cho_solve, cho_factor 

3from jetgp.full_degp import degp_utils as utils 

4from line_profiler import profile 

5import jetgp.utils as gen_utils 

6from jetgp.hyperparameter_optimizers import OPTIMIZERS 

7from jetgp.utils import matern_kernel_grad_builder 

8 

9class Optimizer: 

10 """ 

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

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

13 

14 Parameters 

15 ---------- 

16 model : object 

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

18 and kernel configuration. 

19 """ 

20 

21 def __init__(self, model): 

22 self.model = model 

23 self._kernel_plan = None 

24 self._deriv_buf = None 

25 self._deriv_buf_shape = None 

26 self._deriv_buf_ndir = None 

27 self._deriv_factors = None 

28 self._deriv_factors_key = None 

29 self._K_buf = None 

30 self._dK_buf = None 

31 self._kernel_buf_size = None 

32 self._W_proj_buf = None 

33 self._W_proj_shape = None 

34 

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

36 """Return a pre-allocated buffer for get_all_derivs, reusing if shape matches.""" 

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 @profile 

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 """Enumerate derivative factors in struct memory order for a given order. 

57 

58 Yields the factorial factor prod(count_b!) for each multi-index of 

59 the given order, enumerated in the same order as the OTI struct layout 

60 (last-index-major: for last=1..max_basis, recurse prefix with max=last). 

61 """ 

62 from math import factorial 

63 from collections import Counter 

64 if ordi == 1: 

65 for _ in range(max_basis): 

66 yield 1.0 

67 return 

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

69 if ordi == 2: 

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

71 counts = Counter((i, last)) 

72 f = 1 

73 for c in counts.values(): 

74 f *= factorial(c) 

75 yield float(f) 

76 else: 

77 for prefix_factor, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1): 

78 counts = dict(prefix_counts) 

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

80 f = 1 

81 for c in counts.values(): 

82 f *= factorial(c) 

83 yield float(f) 

84 

85 @staticmethod 

86 def _enum_factors_with_counts(max_basis, ordi): 

87 """Enumerate (factor, counts_dict) pairs in struct order.""" 

88 from math import factorial 

89 from collections import Counter 

90 if ordi == 1: 

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

92 yield 1.0, {i: 1} 

93 return 

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

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

96 counts = dict(prefix_counts) 

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

98 f = 1 

99 for c in counts.values(): 

100 f *= factorial(c) 

101 yield float(f), counts 

102 

103 def _get_deriv_factors(self, n_bases, order): 

104 """Return cached precomputed derivative factorial factors.""" 

105 key = (n_bases, order) 

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

107 return self._deriv_factors 

108 factors = [1.0] # order 0: real part 

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

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

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

112 self._deriv_factors_key = key 

113 return self._deriv_factors 

114 

115 def _ensure_kernel_plan(self, n_bases): 

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

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

118 return 

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

120 self._kernel_plan = None 

121 return 

122 self._kernel_plan = utils.precompute_kernel_plan( 

123 self.model.n_order, n_bases, 

124 self.model.flattened_der_indices, 

125 self.model.powers, 

126 self.model.derivative_locations, 

127 ) 

128 self._kernel_plan_n_bases = n_bases 

129 # Reset kernel buffers when plan changes 

130 self._K_buf = None 

131 self._dK_buf = None 

132 self._kernel_buf_size = None 

133 

134 def _ensure_kernel_bufs(self, n_rows_func): 

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

136 if self._kernel_plan is None: 

137 return 

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

139 if self._kernel_buf_size != total: 

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

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

142 self._kernel_buf_size = total 

143 # Cache absolute offsets in plan so rbf_kernel_fast doesn't recompute 

144 if 'row_offsets_abs' not in self._kernel_plan: 

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

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

147 

148 def negative_log_marginal_likelihood(self, x0): 

149 """ 

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

151 

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

153 

154 Parameters 

155 ---------- 

156 x0 : ndarray 

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

158 

159 Returns 

160 ------- 

161 float 

162 Value of the negative log marginal likelihood. 

163 """ 

164 ell = x0[:-1] 

165 sigma_n = x0[-1] 

166 llhood = 0 

167 diffs = self.model.differences_by_dim 

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

169 if self.model.n_order == 0: 

170 n_bases = 0 

171 phi_exp = phi.real 

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

173 else: 

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

175 deriv_order = 2 * self.model.n_order 

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

177 self._ensure_kernel_plan(n_bases) 

178 if self._kernel_plan is not None: 

179 base_shape = phi.shape 

180 self._ensure_kernel_bufs(base_shape[0]) 

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

182 K = utils.rbf_kernel_fast(phi_3d, self._kernel_plan, out=self._K_buf) 

183 else: 

184 K = utils.rbf_kernel( 

185 phi, phi_exp, 

186 self.model.n_order, n_bases, 

187 self.model.flattened_der_indices, self.model.powers, 

188 index=self.model.derivative_locations, 

189 ) 

190 noise_var = (10 ** sigma_n) ** 2 

191 K.flat[::K.shape[0] + 1] += noise_var 

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

193 

194 # Debug: check kernel matrix sparsity 

195 # near_zero = np.sum(np.abs(K) < 1e-10) 

196 # total = K.size 

197 # sparsity = near_zero / total 

198 # if sparsity > 0.5: 

199 # print(f" WARNING: K is {sparsity*100:.1f}% sparse | ell={10**np.array(ell)} | sigma_n={10**sigma_n:.2e} | cond={np.linalg.cond(K.real):.2e}") 

200 # input('i') 

201 try: 

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

203 alpha = cho_solve( 

204 (L, low), 

205 self.model.y_train 

206 ) 

207 

208 # Cache for fast prediction 

209 self.model._cached_L = L 

210 self.model._cached_low = low 

211 self.model._cached_alpha = alpha 

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

213 

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

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

216 complexity = log_det_K 

217 N = len(self.model.y_train) 

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

219 return data_fit + complexity + const 

220 except Exception: 

221 return 1e6 

222 

223 def nll_wrapper(self, x0): 

224 """ 

225 Wrapper function to compute NLL for optimizer. 

226 

227 Parameters 

228 ---------- 

229 x0 : ndarray 

230 Hyperparameter vector. 

231 

232 Returns 

233 ------- 

234 float 

235 NLL evaluated at x0. 

236 """ 

237 return self.negative_log_marginal_likelihood(x0) 

238 @profile 

239 def _compute_grad(self, x0, W, phi, n_bases, oti, diffs): 

240 """ 

241 Compute the NLL gradient given pre-factorised W = K^{-1} - α α^T. 

242 

243 Factoring this out allows nll_grad and nll_and_grad to share the 

244 expensive Cholesky decomposition instead of each rebuilding it. 

245 """ 

246 ln10 = np.log(10.0) 

247 kernel = self.model.kernel 

248 kernel_type = self.model.kernel_type 

249 D = len(diffs) 

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

251 

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

253 use_fast = self._kernel_plan is not None 

254 base_shape = (W.shape[0] - self._kernel_plan['n_pts_with_derivs'],) * 2 if use_fast else None 

255 

256 deriv_order = 2 * self.model.n_order 

257 

258 # Precompute W projected into phi_exp space to avoid assembling 

259 # the full dK matrix for each hyperparameter dimension. 

260 W_proj = None 

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

262 from math import comb 

263 ndir = comb(n_bases + deriv_order, deriv_order) 

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

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

266 self._W_proj_buf = np.empty(proj_shape) 

267 self._W_proj_shape = proj_shape 

268 W_proj = self._W_proj_buf 

269 

270 plan = self._kernel_plan 

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

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

273 

274 utils._project_W_to_phi_space( 

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

276 plan['fd_flat_indices'], plan['df_flat_indices'], 

277 plan['dd_flat_indices'], 

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

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

280 ) 

281 

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

283 FW_T = None 

284 if _use_vdot_fused: 

285 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order) 

286 from math import comb 

287 _vdot_arr = np.asarray(_vdot_factors) 

288 ndir_d = len(_vdot_arr) 

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

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

291 

292 def _gc(dphi): 

293 if _use_vdot_fused: 

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

295 if self.model.n_order == 0: 

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

297 else: 

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

299 if W_proj is not None: 

300 dphi_3d = dphi_exp.reshape(W_proj.shape) 

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

302 elif use_fast: 

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

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

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

306 else: 

307 dK = utils.rbf_kernel( 

308 dphi, dphi_exp, 

309 self.model.n_order, n_bases, 

310 self.model.flattened_der_indices, self.model.powers, 

311 index=self.model.derivative_locations, 

312 ) 

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

314 

315 # ── signal variance (common: d phi/d log_sf = 2*ln10 * phi) ────── 

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

317 

318 # ── noise variance (common: dK/d log_sn = diag(2*ln10*σ_n²)) ──── 

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

320 

321 # ── kernel-specific hyperparameter gradients ────────────────────── 

322 

323 if kernel == 'SE': 

324 # phi = sf² * exp(-0.5 * Σ_d ell_d² * diff_d²) 

325 # d phi/d log_ell_d = -ln10 * ell_d² * diff_d² * phi 

326 if kernel_type == 'anisotropic': 

327 ell = 10.0 ** x0[:D] 

328 if _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'): 

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

330 grad_buf = np.zeros(D) 

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

332 grad[:D] = grad_buf 

333 elif hasattr(phi, 'fused_scale_sq_mul_sparse'): 

334 dphi_buf = oti.zeros(phi.shape) 

335 for d in range(D): 

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

337 grad[d] = _gc(dphi_buf) 

338 elif hasattr(phi, 'fused_scale_sq_mul'): 

339 dphi_buf = oti.zeros(phi.shape) 

340 for d in range(D): 

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

342 grad[d] = _gc(dphi_buf) 

343 else: 

344 for d in range(D): 

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

346 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi)) 

347 grad[d] = _gc(dphi_d) 

348 else: # isotropic: single ell 

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

350 if hasattr(phi, 'fused_sum_sq_sparse'): 

351 sum_sq = oti.zeros(phi.shape) 

352 sum_sq.fused_sum_sq_sparse(diffs) 

353 elif hasattr(phi, 'fused_sum_sq'): 

354 sum_sq = oti.zeros(phi.shape) 

355 sum_sq.fused_sum_sq(diffs) 

356 else: 

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

358 for d in range(1, D): 

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

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

361 

362 elif kernel == 'RQ': 

363 # phi = sf² * (1 + r²/(2α))^(-α), r² = Σ_d (ell_d * diff_d)² 

364 # d phi/d log_ell_d = -ln10 * ell_d² * diff_d² * phi / base 

365 # d phi/d log_α = ln10 * α * phi * [-log(base) + (1 - 1/base)] 

366 if kernel_type == 'anisotropic': 

367 ell = 10.0 ** x0[:D] 

368 alpha_rq = 10.0 ** float(x0[D]) 

369 alpha_idx = D 

370 else: 

371 ell_val = 10.0 ** float(x0[0]) 

372 ell = np.full(D, ell_val) 

373 alpha_rq = np.exp(float(x0[1])) # iso uses exp(x), not 10^x 

374 alpha_idx = 1 

375 

376 # Recompute r² and base in OTI 

377 if hasattr(phi, 'fused_sqdist_sparse'): 

378 r2 = oti.zeros(phi.shape) 

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

380 r2.fused_sqdist_sparse(diffs, ell_sq) 

381 elif hasattr(phi, 'fused_sqdist'): 

382 r2 = oti.zeros(phi.shape) 

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

384 r2.fused_sqdist(diffs, ell_sq) 

385 else: 

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

387 r2 = oti.mul(r2, r2) 

388 for d in range(1, D): 

389 td = oti.mul(ell[d], diffs[d]) 

390 r2 = oti.sum(r2, oti.mul(td, td)) 

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

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

393 phi_over_base = oti.mul(phi, inv_base) 

394 

395 if kernel_type == 'anisotropic': 

396 if _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'): 

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

398 grad_buf = np.zeros(D) 

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

400 grad[:D] = grad_buf 

401 elif hasattr(phi, 'fused_scale_sq_mul_sparse'): 

402 dphi_buf = oti.zeros(phi.shape) 

403 for d in range(D): 

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

405 grad[d] = _gc(dphi_buf) 

406 elif hasattr(phi, 'fused_scale_sq_mul'): 

407 dphi_buf = oti.zeros(phi.shape) 

408 for d in range(D): 

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

410 grad[d] = _gc(dphi_buf) 

411 else: 

412 for d in range(D): 

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

414 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi_over_base)) 

415 grad[d] = _gc(dphi_d) 

416 else: 

417 if hasattr(phi, 'fused_sum_sq_sparse'): 

418 sum_sq = oti.zeros(phi.shape) 

419 sum_sq.fused_sum_sq_sparse(diffs) 

420 elif hasattr(phi, 'fused_sum_sq'): 

421 sum_sq = oti.zeros(phi.shape) 

422 sum_sq.fused_sum_sq(diffs) 

423 else: 

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

425 for d in range(1, D): 

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

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

428 

429 # alpha gradient: phi * α_factor * [-log(base) + (1 - 1/base)] 

430 # aniso: alpha = 10^x → d alpha/dx = ln10 * alpha 

431 # iso: alpha = exp(x) → d alpha/dx = alpha 

432 log_base = oti.log(base) 

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

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

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

436 

437 elif kernel == 'SineExp': 

438 # phi = sf² * exp(-2 * Σ_d (ell_d * sin(π/p_d * diff_d))²) 

439 # d phi/d log_ell_d = -4*ln10 * ell_d² * sin_d² * phi 

440 # d phi/d log_p_d = 4*ln10 * ell_d² * (π/p_d) * sin_d * cos_d * diff_d * phi 

441 if kernel_type == 'anisotropic': 

442 ell = 10.0 ** x0[:D] 

443 p = 10.0 ** x0[D:2 * D] 

444 pip = np.pi / p # π/p_d per dimension 

445 p_start = D # index of first log_p in x0 

446 else: 

447 ell_val = 10.0 ** float(x0[0]) 

448 p_val = 10.0 ** float(x0[1]) 

449 pip_val = np.pi / p_val 

450 ell = np.full(D, ell_val) 

451 pip = np.full(D, pip_val) 

452 p_start = 1 

453 

454 # Precompute sin and cos for each dimension 

455 sin_d = [] 

456 cos_d = [] 

457 for d in range(D): 

458 arg = oti.mul(pip[d], diffs[d]) 

459 sin_d.append(oti.sin(arg)) 

460 cos_d.append(oti.cos(arg)) 

461 

462 # Length-scale gradients: d phi/d log_ell_d = -4*ln10 * ell_d² * sin_d² * phi 

463 if kernel_type == 'anisotropic': 

464 if hasattr(phi, 'fused_scale_sq_mul'): 

465 dphi_buf = oti.zeros(phi.shape) 

466 for d in range(D): 

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

468 grad[d] = _gc(dphi_buf) 

469 else: 

470 for d in range(D): 

471 sin_sq = oti.mul(sin_d[d], sin_d[d]) 

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

473 oti.mul(sin_sq, phi))) 

474 else: 

475 if hasattr(phi, 'fused_sum_sq'): 

476 sum_sin_sq = oti.zeros(phi.shape) 

477 sum_sin_sq.fused_sum_sq(sin_d) 

478 else: 

479 sum_sin_sq = oti.mul(sin_d[0], sin_d[0]) 

480 for d in range(1, D): 

481 sum_sin_sq = oti.sum(sum_sin_sq, oti.mul(sin_d[d], sin_d[d])) 

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

483 oti.mul(sum_sin_sq, phi))) 

484 

485 # Period gradients 

486 if kernel_type == 'anisotropic': 

487 for d in range(D): 

488 # d phi/d log_p_d = 4*ln10*ell_d²*(π/p_d)*sin_d*cos_d*diff_d * phi 

489 sc_diff = oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d])) 

490 scale = 4.0 * ln10 * ell[d] ** 2 * pip[d] 

491 grad[p_start + d] = _gc(oti.mul(scale, oti.mul(sc_diff, phi))) 

492 else: 

493 # d phi/d log_p = 4*ln10*ell²*(π/p) * Σ_d(sin_d*cos_d*diff_d) * phi 

494 sum_scd = oti.mul(sin_d[0], oti.mul(cos_d[0], diffs[0])) 

495 for d in range(1, D): 

496 sum_scd = oti.sum(sum_scd, 

497 oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d]))) 

498 scale = 4.0 * ln10 * ell[0] ** 2 * pip[0] 

499 grad[p_start] = _gc(oti.mul(scale, oti.mul(sum_scd, phi))) 

500 

501 elif kernel == 'Matern': 

502 # phi = sf² * f(r), r = sqrt(Σ_d (ell_d*(diff_d+ε))²) 

503 # d phi/d log_ell_d = sf² * f'(r) * ln10 * ell_d² * (diff_d+ε)² / r 

504 kf = self.model.kernel_factory 

505 

506 # Build/cache the Matern derivative function 

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

508 kf._matern_grad_prebuild = matern_kernel_grad_builder( 

509 kf.nu, oti_module=oti) 

510 

511 if kernel_type == 'anisotropic': 

512 ell = 10.0 ** x0[:D] 

513 else: 

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

515 

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

517 _eps = 1e-10 # regularise r, not each diff (matches kernel_funcs.py) 

518 

519 # Recompute r in OTI (matches matern_kernel_anisotropic/isotropic) 

520 if hasattr(phi, 'fused_sqdist_sparse'): 

521 r2 = oti.zeros(phi.shape) 

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

523 r2.fused_sqdist_sparse(diffs, ell_sq) 

524 elif hasattr(phi, 'fused_sqdist'): 

525 r2 = oti.zeros(phi.shape) 

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

527 r2.fused_sqdist(diffs, ell_sq) 

528 else: 

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

530 r2 = oti.mul(r2, r2) 

531 for d in range(1, D): 

532 td = oti.mul(ell[d], diffs[d]) 

533 r2 = oti.sum(r2, oti.mul(td, td)) 

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

535 f_prime_r = kf._matern_grad_prebuild(r_oti) # df/dr (OTI) 

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

537 

538 # Precompute base = sigma_f² * f'(r) * 1/r for length-scale gradients 

539 # grad[d] = _gc(base * ln10 * ell_d² * diff_d²) 

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

541 if kernel_type == 'anisotropic': 

542 if _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'): 

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

544 grad_buf = np.zeros(D) 

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

546 grad[:D] = grad_buf 

547 elif hasattr(phi, 'fused_scale_sq_mul_sparse'): 

548 dphi_buf = oti.zeros(phi.shape) 

549 for d in range(D): 

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

551 grad[d] = _gc(dphi_buf) 

552 elif hasattr(phi, 'fused_scale_sq_mul'): 

553 dphi_buf = oti.zeros(phi.shape) 

554 for d in range(D): 

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

556 grad[d] = _gc(dphi_buf) 

557 else: 

558 for d in range(D): 

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

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

561 grad[d] = _gc(dphi_d) 

562 else: 

563 ell_val = ell[0] 

564 if hasattr(phi, 'fused_sum_sq_sparse'): 

565 sum_dsq = oti.zeros(phi.shape) 

566 sum_dsq.fused_sum_sq_sparse(diffs) 

567 elif hasattr(phi, 'fused_sum_sq'): 

568 sum_dsq = oti.zeros(phi.shape) 

569 sum_dsq.fused_sum_sq(diffs) 

570 else: 

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

572 for d in range(1, D): 

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

574 dphi_e = oti.mul(ln10 * ell_val ** 2, oti.mul(sum_dsq, base_matern)) 

575 grad[0] = _gc(dphi_e) 

576 

577 elif kernel == 'SI': 

578 # phi = sf² * Π_d (1 + ell_d * B(diff_d)) 

579 # d phi/d log_ell_d = ln10 * ell_d * B(diff_d) / (1 + ell_d*B(diff_d)) * phi 

580 kf = self.model.kernel_factory 

581 si_prebuild = kf.SI_kernel_prebuild 

582 

583 if kernel_type == 'anisotropic': 

584 ell = 10.0 ** x0[:D] 

585 else: 

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

587 

588 # Precompute SI values and factor terms for each dimension 

589 si_vals = [si_prebuild(diffs[d]) for d in range(D)] 

590 term_vals = [oti.sum(1.0, oti.mul(ell[d], si_vals[d])) for d in range(D)] 

591 

592 if kernel_type == 'anisotropic': 

593 for d in range(D): 

594 phi_over_term = oti.div(phi, term_vals[d]) 

595 dphi_d = oti.mul(ln10 * ell[d], 

596 oti.mul(si_vals[d], phi_over_term)) 

597 grad[d] = _gc(dphi_d) 

598 else: 

599 ell_val = ell[0] 

600 # d phi/d log_ell = ln10 * ell * Σ_d [B(diff_d)/(1+ell*B(diff_d))] * phi 

601 acc = oti.mul(si_vals[0], oti.div(phi, term_vals[0])) 

602 for d in range(1, D): 

603 acc = oti.sum(acc, oti.mul(si_vals[d], 

604 oti.div(phi, term_vals[d]))) 

605 grad[0] = _gc(oti.mul(ln10 * ell_val, acc)) 

606 

607 return grad 

608 

609 def nll_grad(self, x0): 

610 """Analytic gradient of the NLL (separate Cholesky from nll_wrapper).""" 

611 diffs = self.model.differences_by_dim 

612 oti = self.model.kernel_factory.oti 

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

614 

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

616 if self.model.n_order == 0: 

617 n_bases = 0 

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

619 else: 

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

621 deriv_order = 2 * self.model.n_order 

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

623 

624 self._ensure_kernel_plan(n_bases) 

625 if self._kernel_plan is not None: 

626 base_shape = phi.shape 

627 self._ensure_kernel_bufs(base_shape[0]) 

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

629 K = utils.rbf_kernel_fast(phi_3d, self._kernel_plan, out=self._K_buf) 

630 else: 

631 K = utils.rbf_kernel( 

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

633 self.model.flattened_der_indices, self.model.powers, 

634 index=self.model.derivative_locations, 

635 ) 

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

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