Coverage for jetgp/full_gddegp_sparse/optimizer.py: 63%

1import numpy as np

2import numba

3from scipy.linalg import cho_solve, cho_factor, blas

4from jetgp.full_gddegp_sparse import gddegp_utils as utils

5from jetgp.full_gddegp_sparse.sparse_cholesky import (

6 build_U, build_U_supernodes, nlml_from_U, alpha_from_U

8from line_profiler import profile

9import jetgp.utils as gen_utils

10from jetgp.hyperparameter_optimizers import OPTIMIZERS

11from jetgp.utils import matern_kernel_grad_builder

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

15def _permute_and_subtract_outer(K_inv_ord, alpha_v, P_full, W):

16 """Fused: W[P[i], P[j]] = K_inv_ord[i, j] - alpha_v[P[i]] * alpha_v[P[j]]

17 Reads only the lower triangle of K_inv_ord (as produced by dsyrk with lower=1

18 on a Fortran-order buffer).

19 """

20 N = len(P_full)

21 for i in numba.prange(N):

22 pi = P_full[i]

23 ai = alpha_v[pi]

24 W[pi, pi] = K_inv_ord[i, i] - ai * ai

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

26 pj = P_full[j]

27 val = K_inv_ord[j, i] - ai * alpha_v[pj]

28 W[pi, pj] = val

29 W[pj, pi] = val

32def _build_k_index_map(plan, n_rows_func):

33 """

34 Build arrays that map each K-matrix index to (deriv_type, physical_point).

36 Returns

37 -------

38 k_type : int64 array of shape (N_total,)

39 Derivative type for each K index (0 = function value, 1.. = derivatives).

40 k_phys : int64 array of shape (N_total,)

41 Physical training point index for each K index.

42 deriv_lookup : int64 array of shape (n_types, n_types)

43 phi_exp_3d derivative-dimension index for block (type_i, type_j).

44 sign_lookup : float64 array of shape (n_types,)

45 Sign multiplier for each derivative type (all 1.0 for GDDEGP).

46 """

47 n_dt = plan['n_deriv_types']

48 n_types = n_dt + 1

49 N_total = n_rows_func + plan['n_pts_with_derivs']

51 k_type = np.empty(N_total, dtype=np.int64)

52 k_phys = np.empty(N_total, dtype=np.int64)

54 # Function-value rows: type 0, phys = row index

55 k_type[:n_rows_func] = 0

56 k_phys[:n_rows_func] = np.arange(n_rows_func)

58 # Derivative rows

59 idx_flat = plan['idx_flat']

60 idx_offsets = plan['idx_offsets']

61 idx_sizes = plan['index_sizes']

62 row_offsets = plan.get('row_offsets_abs', plan['row_offsets'] + n_rows_func)

63 for j in range(n_dt):

64 ro = row_offsets[j]

65 sz = idx_sizes[j]

66 off = idx_offsets[j]

67 k_type[ro:ro + sz] = j + 1

68 k_phys[ro:ro + sz] = idx_flat[off:off + sz]

70 # Derivative lookup: which phi_exp_3d dimension for (type_i, type_j)

71 deriv_lookup = np.empty((n_types, n_types), dtype=np.int64)

72 deriv_lookup[0, 0] = 0

73 fd = plan['fd_flat_indices']

74 df = plan['df_flat_indices']

75 dd = plan['dd_flat_indices']

76 for j in range(n_dt):

77 deriv_lookup[0, j + 1] = fd[j]

78 deriv_lookup[j + 1, 0] = df[j]

79 for i in range(n_dt):

80 for j in range(n_dt):

81 deriv_lookup[i + 1, j + 1] = dd[i, j]

83 # Sign lookup (all 1.0 for GDDEGP — even/odd bases handle signs internally)

84 sign_lookup = np.ones(n_types, dtype=np.float64)

86 return k_type, k_phys, deriv_lookup, sign_lookup

89@numba.jit(nopython=True, cache=True)

90def _extract_K_sub(phi_exp_3d, nb_type, nb_phys, deriv_lookup, sign_lookup,

91 sigma_n_sq, sigma_data_diag, m, K_sub):

92 """

93 Assemble K_sub directly from phi_exp_3d for a neighbourhood `nb`.

95 nb_type[a], nb_phys[a] give the derivative type and physical point

96 for the a-th row/column of K_sub.

97 """

98 for a in range(m):

99 ta = nb_type[a]

100 pa = nb_phys[a]

101 for b in range(m):

102 tb = nb_type[b]

103 pb = nb_phys[b]

104 d = deriv_lookup[ta, tb]

105 K_sub[a, b] = phi_exp_3d[d, pa, pb] * sign_lookup[tb]

106 # Add noise to diagonal

107 K_sub[a, a] += sigma_n_sq + sigma_data_diag[a]

108

109

110@numba.jit(nopython=True, cache=True)

111def _project_G_to_W_proj(W_proj, G_b, nb_type, nb_phys,

112 deriv_lookup, sign_lookup, m):

113 """

114 Project per-block sensitivity G_b (m × m) into W_proj (ndir × n_func × n_func).

115 """

116 n_func = W_proj.shape[1]

117 plane = n_func * W_proj.shape[2]

118 wptr = W_proj.ravel()

119 for a_i in range(m):

120 ta = nb_type[a_i]

121 pa = nb_phys[a_i]

122 for bb_i in range(m):

123 tb = nb_type[bb_i]

124 pb = nb_phys[bb_i]

125 d = deriv_lookup[ta, tb]

126 wptr[d * plane + pa * W_proj.shape[2] + pb] += sign_lookup[tb] * G_b[a_i, bb_i]

127

128

129class Optimizer:

130 """

131 Optimizer class to perform hyperparameter tuning for sparse GDDEGP models

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

133

134 Parameters

135 ----------

136 model : object

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

138 and kernel configuration.

139 """

140

141 def __init__(self, model):

142 self.model = model

143 self._kernel_plan = None

144 self._deriv_buf = None

145 self._deriv_buf_shape = None

146 self._deriv_buf_ndir = None

147 self._deriv_factors = None

148 self._deriv_factors_key = None

149 self._K_buf = None

150 self._dK_buf = None

151 self._kernel_buf_size = None

152 self._W_proj_buf = None

153 self._W_proj_shape = None

154 self._U_buf = None

155 self._P_ix = None

156 self._K_inv_buf = None

157 # Direct phi extraction maps (built lazily)

158 self._k_index_map = None

159 self._inv_P = None

160 self._sigma_data_diag_mmd = None

161 self._block_phi_maps = None

162

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

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

165 if self._deriv_buf_ndir is None:

166 from math import comb

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

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

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

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

171 self._deriv_buf_shape = shape

172 return self._deriv_buf

173

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

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

176 if hasattr(phi, 'get_all_derivs_fast'):

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

178 factors = self._get_deriv_factors(n_bases, deriv_order)

179 return phi.get_all_derivs_fast(factors, buf)

180 return phi.get_all_derivs(n_bases, deriv_order)

181

182 @staticmethod

183 def _enum_factors(max_basis, ordi):

184 """Enumerate derivative factors in struct memory order for a given order."""

185 from math import factorial

186 from collections import Counter

187 if ordi == 1:

188 for _ in range(max_basis):

189 yield 1.0

190 return

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

192 if ordi == 2:

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

194 counts = Counter((i, last))

195 f = 1

196 for c in counts.values():

197 f *= factorial(c)

198 yield float(f)

199 else:

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

201 counts = dict(prefix_counts)

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

203 f = 1

204 for c in counts.values():

205 f *= factorial(c)

206 yield float(f)

207

208 @staticmethod

209 def _enum_factors_with_counts(max_basis, ordi):

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

211 from math import factorial

212 from collections import Counter

213 if ordi == 1:

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

215 yield 1.0, {i: 1}

216 return

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

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

219 counts = dict(prefix_counts)

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

221 f = 1

222 for c in counts.values():

223 f *= factorial(c)

224 yield float(f), counts

225

226 def _get_deriv_factors(self, n_bases, order):

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

228 key = (n_bases, order)

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

230 return self._deriv_factors

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

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

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

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

235 self._deriv_factors_key = key

236 return self._deriv_factors

237

238 def _ensure_kernel_plan(self, n_bases):

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

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

241 return

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

243 self._kernel_plan = None

244 return

245 self._kernel_plan = utils.precompute_kernel_plan(

246 self.model.n_order, n_bases,

247 self.model.flattened_der_indices,

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

249 self.model.derivative_locations,

250 )

251 self._kernel_plan_n_bases = n_bases

252 # Reset kernel buffers when plan changes

253 self._K_buf = None

254 self._dK_buf = None

255 self._kernel_buf_size = None

256

257 def _ensure_kernel_bufs(self, n_rows_func):

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

259 if self._kernel_plan is None:

260 return

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

262 if self._kernel_buf_size != total:

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

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

265 self._kernel_buf_size = total

266 if 'row_offsets_abs' not in self._kernel_plan:

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

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

269

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

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

272 self._ensure_kernel_plan(n_bases)

273 if self._kernel_plan is not None:

274 base_shape = phi.shape

275 self._ensure_kernel_bufs(base_shape[0])

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

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

278 return utils.rbf_kernel(

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

280 self.model.flattened_der_indices,

281 index=self.model.derivative_locations,

282 )

283

284 def _ensure_phi_index_maps(self, n_rows_func):

285 """Lazily build the K-index-to-phi maps and inverse permutation."""

286 if self._k_index_map is not None:

287 return

288 plan = self._kernel_plan

289 k_type, k_phys, deriv_lookup, sign_lookup = _build_k_index_map(

290 plan, n_rows_func)

291 self._k_index_map = (k_type, k_phys, deriv_lookup, sign_lookup)

292

293 P_full = self.model.mmd_P_full

294 inv_P = np.empty_like(P_full)

295 inv_P[P_full] = np.arange(len(P_full))

296 self._inv_P = inv_P

297

298 # sigma_data diagonal in MMD order

299 sd = self.model.sigma_data

300 if sd.ndim == 2:

301 sd_diag_orig = np.diag(sd) ** 2 if np.any(sd) else np.zeros(len(P_full))

302 else:

303 sd_diag_orig = np.zeros(len(P_full))

304 self._sigma_data_diag_mmd = sd_diag_orig[P_full]

305

306 # Precompute flat index arrays for phi_exp_3d gather.

307 stride_d = n_rows_func * n_rows_func

308 stride_row = n_rows_func

309

310 if (self.model.use_supernodes

311 and self.model.sparse_supernodes_full is not None):

312 for sn in self.model.sparse_supernodes_full:

313 ch = sn.get('children_arr')

314 if ch is None:

315 ch = np.asarray(sn['children'])

316 orig_ch = P_full[ch]

317 ch_type = k_type[orig_ch]

318 ch_phys = k_phys[orig_ch]

319 m = len(ch)

320

321 d_mat = deriv_lookup[ch_type[:, None], ch_type[None, :]]

322 pa_mat = np.broadcast_to(ch_phys[:, None], (m, m))

323 pb_mat = np.broadcast_to(ch_phys[None, :], (m, m))

324 sn['phi_flat_idx'] = np.ascontiguousarray(

325 d_mat * stride_d + pa_mat * stride_row + pb_mat

326 )

327 sn['phi_sign_mat'] = np.ascontiguousarray(

328 sign_lookup[ch_type[None, :]] * np.ones((m, 1))

329 )

330 sn['phi_sd_diag'] = self._sigma_data_diag_mmd[ch]

331

332 # Same precomputation for non-supernode block path

333 if (not self.model.use_supernodes

334 and self.model.n_order > 0):

335 N_total = len(P_full)

336 block_size = 1 + plan['n_deriv_types']

337 S = self.model.sparse_S_full_arr

338 self._block_phi_maps = []

339 for start in range(0, N_total, block_size):

340 end = min(start + block_size, N_total)

341 nb = S[end - 1] if isinstance(S[end - 1], np.ndarray) else np.asarray(S[end - 1])

342 m = len(nb)

343

344 orig_nb = P_full[nb]

345 nb_type = k_type[orig_nb]

346 nb_phys = k_phys[orig_nb]

347

348 d_mat = deriv_lookup[nb_type[:, None], nb_type[None, :]]

349 pa_mat = np.broadcast_to(nb_phys[:, None], (m, m))

350 pb_mat = np.broadcast_to(nb_phys[None, :], (m, m))

351

352 self._block_phi_maps.append({

353 'nb': nb,

354 'start': start,

355 'flat_idx': np.ascontiguousarray(

356 d_mat * stride_d + pa_mat * stride_row + pb_mat

357 ),

358 'sign_mat': np.ascontiguousarray(

359 sign_lookup[nb_type[None, :]] * np.ones((m, 1))

360 ),

361 'sd_diag': self._sigma_data_diag_mmd[nb],

362 'positions': np.searchsorted(nb, np.arange(start, end)),

363 })

364

365 @profile

366 def negative_log_marginal_likelihood(self, x0):

367 """

368 Compute the negative log marginal likelihood (NLL) via sparse U.

369

370 NLL = 0.5 * ||U^T y||^2 - sum(log|diag(U)|) + 0.5 * N * log(2pi)

371

372 Parameters

373 ----------

374 x0 : ndarray

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

376

377 Returns

378 -------

379 float

380 Value of the negative log marginal likelihood.

381 """

382 try:

383 if self.model._use_dense_factor:

384 W, alpha, nll, *_ = self._dense_nll_and_W(x0)

385 if nll > 1e6:

386 return 1e6

387 return nll

388

389 # Use direct phi path (skip full K construction) when possible

390 if self.model.n_order > 0:

391 alpha, U, nlml, *_ = self._sparse_nlml_direct(x0)

392 else:

393 K, _, _, _, _ = self._build_K_and_phi(x0)

394 alpha, U, nlml = self._sparse_U_alpha_nll(K)

395

396 # Sparse U can silently produce bad factors when K is

397 # ill-conditioned (e.g. very small noise). Clamp to 1e6

398 # to match the dense fallback behaviour.

399 if nlml > 1e6:

400 return 1e6

401

402 self.model._cached_U = U

403 self.model._cached_P = self.model.mmd_P_full

404 self.model._cached_alpha = alpha

405 self.model._cached_L = None

406 self.model._cached_low = None

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

408

409 return nlml

410 except Exception:

411 return 1e6

412

413 def nll_wrapper(self, x0):

414 """

415 Wrapper function to compute NLL for optimizer.

416

417 Parameters

418 ----------

419 x0 : ndarray

420 Hyperparameter vector.

421

422 Returns

423 -------

424 float

425 NLL evaluated at x0.

426 """

427 return self.negative_log_marginal_likelihood(x0)

428

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

430 """

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

432 """

433 ln10 = np.log(10.0)

434 kernel = self.model.kernel

435 kernel_type = self.model.kernel_type

436 D = len(diffs)

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

438

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

440 use_fast = self._kernel_plan is not None

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

442

443 deriv_order = 2 * self.model.n_order

444

445 # Precompute W projected into phi_exp space

446 W_proj = None

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

448 from math import comb

449 ndir = comb(n_bases + deriv_order, deriv_order)

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

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

452 self._W_proj_buf = np.empty(proj_shape)

453 self._W_proj_shape = proj_shape

454 W_proj = self._W_proj_buf

455

456 plan = self._kernel_plan

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

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

459

460 # GDDEGP _project_W_to_phi_space does NOT take signs

461 utils._project_W_to_phi_space(

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

463 plan['fd_flat_indices'], plan['df_flat_indices'],

464 plan['dd_flat_indices'],

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

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

467 )

468

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

470 if _use_vdot_fused:

471 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order)

472

473 def _gc(dphi):

474 if _use_vdot_fused:

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

476 if self.model.n_order == 0:

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

478 else:

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

480 if W_proj is not None:

481 dphi_3d = dphi_exp.reshape(W_proj.shape)

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

483 elif use_fast:

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

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

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

487 else:

488 dK = utils.rbf_kernel(

489 dphi, dphi_exp,

490 self.model.n_order, n_bases,

491 self.model.flattened_der_indices,

492 index=self.model.derivative_locations,

493 )

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

495

496 # signal variance

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

498 # noise variance

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

500

501 if kernel == 'SE':

502 if kernel_type == 'anisotropic':

503 ell = 10.0 ** x0[:D]

504 if hasattr(phi, 'fused_scale_sq_mul'):

505 dphi_buf = oti.zeros(phi.shape)

506 for d in range(D):

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

508 grad[d] = _gc(dphi_buf)

509 else:

510 for d in range(D):

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

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

513 grad[d] = _gc(dphi_d)

514 else:

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

516 if hasattr(phi, 'fused_sum_sq'):

517 sum_sq = oti.zeros(phi.shape)

518 sum_sq.fused_sum_sq(diffs)

519 else:

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

521 for d in range(1, D):

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

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

524

525 elif kernel == 'RQ':

526 if kernel_type == 'anisotropic':

527 ell = 10.0 ** x0[:D]

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

529 alpha_idx = D

530 else:

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

532 ell = np.full(D, ell_val)

533 alpha_rq = np.exp(float(x0[1]))

534 alpha_idx = 1

535

536 if hasattr(phi, 'fused_sqdist'):

537 r2 = oti.zeros(phi.shape)

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

539 r2.fused_sqdist(diffs, ell_sq)

540 else:

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

542 r2 = oti.mul(r2, r2)

543 for d in range(1, D):

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

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

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

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

548 phi_over_base = oti.mul(phi, inv_base)

549

550 if kernel_type == 'anisotropic':

551 if hasattr(phi, 'fused_scale_sq_mul'):

552 dphi_buf = oti.zeros(phi.shape)

553 for d in range(D):

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

555 grad[d] = _gc(dphi_buf)

556 else:

557 for d in range(D):

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

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

560 grad[d] = _gc(dphi_d)

561 else:

562 if hasattr(phi, 'fused_sum_sq'):

563 sum_sq = oti.zeros(phi.shape)

564 sum_sq.fused_sum_sq(diffs)

565 else:

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

567 for d in range(1, D):

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

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

570

571 log_base = oti.log(base)

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

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

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

575

576 elif kernel == 'SineExp':

577 if kernel_type == 'anisotropic':

578 ell = 10.0 ** x0[:D]

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

580 pip = np.pi / p

581 p_start = D

582 else:

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

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

585 pip_val = np.pi / p_val

586 ell = np.full(D, ell_val)

587 pip = np.full(D, pip_val)

588 p_start = 1

589

590 sin_d = []

591 cos_d = []

592 for d in range(D):

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

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

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

596

597 if kernel_type == 'anisotropic':

598 if hasattr(phi, 'fused_scale_sq_mul'):

599 dphi_buf = oti.zeros(phi.shape)

600 for d in range(D):

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

602 grad[d] = _gc(dphi_buf)

603 else:

604 for d in range(D):

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

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

607 oti.mul(sin_sq, phi)))

608 else:

609 if hasattr(phi, 'fused_sum_sq'):

610 sum_sin_sq = oti.zeros(phi.shape)

611 sum_sin_sq.fused_sum_sq(sin_d)

612 else:

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

614 for d in range(1, D):

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

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

617 oti.mul(sum_sin_sq, phi)))

618

619 if kernel_type == 'anisotropic':

620 for d in range(D):

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

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

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

624 else:

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

626 for d in range(1, D):

627 sum_scd = oti.sum(sum_scd,

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

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

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

631

632 elif kernel == 'Matern':

633 kf = self.model.kernel_factory

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

635 kf._matern_grad_prebuild = matern_kernel_grad_builder(

636 getattr(kf, "nu", 1.5), oti_module=oti)

637

638 if kernel_type == 'anisotropic':

639 ell = 10.0 ** x0[:D]

640 else:

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

642

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

644 _eps = 1e-10

645

646 if hasattr(phi, 'fused_sqdist'):

647 r2 = oti.zeros(phi.shape)

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

649 r2.fused_sqdist(diffs, ell_sq)

650 else:

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

652 r2 = oti.mul(r2, r2)

653 for d in range(1, D):

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

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

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

657 f_prime_r = kf._matern_grad_prebuild(r_oti)

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

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

660

661 if kernel_type == 'anisotropic':

662 if hasattr(phi, 'fused_scale_sq_mul'):

663 dphi_buf = oti.zeros(phi.shape)

664 for d in range(D):

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

666 grad[d] = _gc(dphi_buf)

667 else:

668 for d in range(D):

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

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

671 grad[d] = _gc(dphi_d)

672 else:

673 ell_val = ell[0]

674 if hasattr(phi, 'fused_sum_sq'):

675 sum_dsq = oti.zeros(phi.shape)

676 sum_dsq.fused_sum_sq(diffs)

677 else:

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

679 for d in range(1, D):

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

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

682 grad[0] = _gc(dphi_e)

683

684 return grad

685

686 @profile

687 def _compute_grad_blockwise(self, x0, U, alpha_v, phi, n_bases, oti, diffs):

688 """

689 Compute the NLL gradient by differentiating through each block's

690 Cholesky in the Vecchia decomposition.

691

692 For each block b the Vecchia NLL contribution is:

693 NLL_b = 0.5 * Σ_j [ α_b[p_j]² / s_j - log(s_j) ]

694

695 where α_b = K_sub⁻¹ y_nb, s_j = (K_sub⁻¹)[p_j, p_j].

696

697 Differentiating through K_sub⁻¹ gives the per-block sensitivity:

698 G_b = 0.5 * M V^T

699

700 where V = K_sub⁻¹ E (un-normalised U block), and

701 M[:, j] = γ_j V[:, j] - 2 β_j α_b

702 β_j = α_b[p_j] / s_j, γ_j = β_j² + 1/s_j

703

704 The gradient is dNLL/dθ = Σ_b tr(G_b dK_sub_b/dθ),

705 which is projected into phi-space as W_proj for vdot.

706 """

707 from math import comb

708

709 ln10 = np.log(10.0)

710 kernel = self.model.kernel

711 kernel_type = self.model.kernel_type

712 D = len(diffs)

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

714

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

716 deriv_order = 2 * self.model.n_order

717 plan = self._kernel_plan

718 P_full = self.model.mmd_P_full

719 N_total = len(P_full)

720 n_func = phi.shape[0]

721

722 ndir = comb(n_bases + deriv_order, deriv_order)

723 k_type, k_phys, deriv_lookup, sign_lookup = self._k_index_map

724

725 # ── phi_exp for K_sub reconstruction ──────────────────────

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

727 phi_3d = phi_exp.reshape(phi_exp.shape[0], n_func, n_func)

728 phi_flat = phi_3d.ravel()

729

730 block_maps = self._block_phi_maps

731 y_ord = self.model.y_train[P_full]

732

733 # ── accumulate W_proj and noise trace from per-block G_b ──

734 proj_shape = (ndir, n_func, n_func)

735 W_proj = np.zeros(proj_shape)

736 w_flat = W_proj.ravel()

737 noise_trace = 0.0

738

739 for bm in block_maps:

740 nb = bm['nb']

741 m = len(nb)

742 positions = bm['positions']

743 n_cols = len(positions)

744

745 # Reconstruct K_sub for this block

746 K_sub = phi_flat[bm['flat_idx']].reshape(m, m) * bm['sign_mat']

747 diag_idx = np.arange(m)

748 K_sub[diag_idx, diag_idx] += sigma_n_sq + bm['sd_diag']

749

750 # Cholesky factor

751 L_u, low_u = cho_factor(K_sub, lower=True)

752

753 # V = K_sub⁻¹ E (un-normalised U block, m × n_cols)

754 E = np.zeros((m, n_cols))

755 E[positions, np.arange(n_cols)] = 1.0

756 V = cho_solve((L_u, low_u), E)

757

758 # α_b = K_sub⁻¹ y_nb (m,)

759 y_nb = y_ord[nb]

760 alpha_b = cho_solve((L_u, low_u), y_nb)

761

762 # Per-parent scalars

763 s = V[positions, np.arange(n_cols)]

764 a = alpha_b[positions]

765 beta = a / s

766 gamma = beta ** 2 + 1.0 / s

767

768 # M[:, j] = γ_j V[:, j] - 2 β_j α_b

769 M = V * gamma[np.newaxis, :] - 2.0 * alpha_b[:, np.newaxis] * beta[np.newaxis, :]

770

771 # Noise trace

772 noise_trace += np.sum(M * V)

773

774 # Project G_b into W_proj using precomputed flat indices

775 G_b = M @ V.T

776 np.add.at(w_flat, bm['flat_idx'].ravel(), (G_b * bm['sign_mat']).ravel())

777

778 # ── noise gradient ───────────────────────────────────────

779 grad[-1] = ln10 * sigma_n_sq * noise_trace

780

781 # ── deriv factors + fast vdot path ───────────────────────

782 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order)

783

784 @profile

785 def _gc_block(dphi):

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

787

788 # ── signal variance ──────────────────────────────────────

789 grad[-2] = ln10 * phi.vdot_expand_fast(_vdot_factors, W_proj)

790

791 # ── kernel-specific hyperparameter gradients ─────────────

792 if kernel == 'SE':

793 if kernel_type == 'anisotropic':

794 ell = 10.0 ** x0[:D]

795 if hasattr(phi, 'fused_scale_sq_mul'):

796 dphi_buf = oti.zeros(phi.shape)

797 for d in range(D):

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

799 grad[d] = _gc_block(dphi_buf)

800 else:

801 for d in range(D):

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

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

804 grad[d] = _gc_block(dphi_d)

805 else:

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

807 if hasattr(phi, 'fused_sum_sq'):

808 sum_sq = oti.zeros(phi.shape)

809 sum_sq.fused_sum_sq(diffs)

810 else:

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

812 for d in range(1, D):

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

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

815

816 elif kernel == 'RQ':

817 if kernel_type == 'anisotropic':

818 ell = 10.0 ** x0[:D]

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

820 alpha_idx = D

821 else:

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

823 ell = np.full(D, ell_val)

824 alpha_rq = np.exp(float(x0[1]))

825 alpha_idx = 1

826

827 if hasattr(phi, 'fused_sqdist'):

828 r2 = oti.zeros(phi.shape)

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

830 r2.fused_sqdist(diffs, ell_sq)

831 else:

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

833 r2 = oti.mul(r2, r2)

834 for d in range(1, D):

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

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

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

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

839 phi_over_base = oti.mul(phi, inv_base)

840

841 if kernel_type == 'anisotropic':

842 if hasattr(phi, 'fused_scale_sq_mul'):

843 dphi_buf = oti.zeros(phi.shape)

844 for d in range(D):

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

846 grad[d] = _gc_block(dphi_buf)

847 else:

848 for d in range(D):

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

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

851 grad[d] = _gc_block(dphi_d)

852 else:

853 if hasattr(phi, 'fused_sum_sq'):

854 sum_sq = oti.zeros(phi.shape)

855 sum_sq.fused_sum_sq(diffs)

856 else:

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

858 for d in range(1, D):

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

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

861

862 log_base = oti.log(base)

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

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

865 grad[alpha_idx] = _gc_block(oti.mul(alpha_factor, oti.mul(phi, term)))

866

867 elif kernel == 'SineExp':

868 if kernel_type == 'anisotropic':

869 ell = 10.0 ** x0[:D]

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

871 pip = np.pi / p

872 p_start = D

873 else:

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

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

876 pip_val = np.pi / p_val

877 ell = np.full(D, ell_val)

878 pip = np.full(D, pip_val)

879 p_start = 1

880

881 sin_d = []

882 cos_d = []

883 for d in range(D):

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

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

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

887

888 if kernel_type == 'anisotropic':

889 if hasattr(phi, 'fused_scale_sq_mul'):

890 dphi_buf = oti.zeros(phi.shape)

891 for d in range(D):

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

893 grad[d] = _gc_block(dphi_buf)

894 else:

895 for d in range(D):

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

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

898 oti.mul(sin_sq, phi)))

899 else:

900 if hasattr(phi, 'fused_sum_sq'):

901 sum_sin_sq = oti.zeros(phi.shape)

902 sum_sin_sq.fused_sum_sq(sin_d)

903 else:

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

905 for d in range(1, D):

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

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

908 oti.mul(sum_sin_sq, phi)))

909

910 if kernel_type == 'anisotropic':

911 for d in range(D):

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

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

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

915 else:

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

917 for d in range(1, D):

918 sum_scd = oti.sum(sum_scd,

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

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

921 grad[p_start] = _gc_block(oti.mul(scale, oti.mul(sum_scd, phi)))

922

923 elif kernel == 'Matern':

924 kf = self.model.kernel_factory

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

926 kf._matern_grad_prebuild = matern_kernel_grad_builder(

927 getattr(kf, "nu", 1.5), oti_module=oti)

928

929 if kernel_type == 'anisotropic':

930 ell = 10.0 ** x0[:D]

931 else:

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

933

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

935 _eps = 1e-10

936

937 if hasattr(phi, 'fused_sqdist'):

938 r2 = oti.zeros(phi.shape)

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

940 r2.fused_sqdist(diffs, ell_sq)

941 else:

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

943 r2 = oti.mul(r2, r2)

944 for d in range(1, D):

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

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

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

948 f_prime_r = kf._matern_grad_prebuild(r_oti)

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

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

951

952 if kernel_type == 'anisotropic':

953 if hasattr(phi, 'fused_scale_sq_mul'):

954 dphi_buf = oti.zeros(phi.shape)

955 for d in range(D):

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

957 grad[d] = _gc_block(dphi_buf)

958 else:

959 for d in range(D):

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

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

962 grad[d] = _gc_block(dphi_d)

963 else:

964 ell_val = ell[0]

965 if hasattr(phi, 'fused_sum_sq'):

966 sum_dsq = oti.zeros(phi.shape)

967 sum_dsq.fused_sum_sq(diffs)

968 else:

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

970 for d in range(1, D):

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

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

973 grad[0] = _gc_block(dphi_e)

974

975 return grad

976

977 @profile

978 def _build_K_and_phi(self, x0):

979 """

980 Shared helper: build K (with noise), phi, n_bases, oti, diffs.

981

982 Returns (K, phi, n_bases, oti, diffs).

983 """

984 diffs = self.model.differences_by_dim

985 oti = self.model.kernel_factory.oti

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

987

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

989 if self.model.n_order == 0:

990 n_bases = 0

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

992 else:

993 active = phi.get_active_bases()

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

995 deriv_order = 2 * self.model.n_order

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

997

998 self._ensure_kernel_plan(n_bases)

999 if self._kernel_plan is not None:

1000 base_shape = phi.shape

1001 self._ensure_kernel_bufs(base_shape[0])

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

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

1004 else:

1005 K = utils.rbf_kernel(

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

1007 self.model.flattened_der_indices,

1008 index=self.model.derivative_locations,

1009 )

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

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

1012 return K, phi, n_bases, oti, diffs

1013

1014 def _sparse_nlml_direct(self, x0):

1015 """

1016 Compute sparse NLML directly from phi_exp_3d, skipping full K

1017 construction and permutation.

1018

1019 Returns (alpha_v, U, nll, phi, n_bases, oti, diffs).

1020 """

1021 from jetgp.full_gddegp_sparse.sparse_cholesky import (

1022 build_U_from_phi_flat, build_U_supernodes_from_phi,

1023 )

1024

1025 diffs = self.model.differences_by_dim

1026 oti = self.model.kernel_factory.oti

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

1028

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

1030 active = phi.get_active_bases()

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

1032 deriv_order = 2 * self.model.n_order

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

1034

1035 self._ensure_kernel_plan(n_bases)

1036 base_shape = phi.shape

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

1038

1039 # Build index maps (once)

1040 self._ensure_phi_index_maps(base_shape[0])

1041 k_type, k_phys, deriv_lookup, sign_lookup = self._k_index_map

1042

1043 P_full = self.model.mmd_P_full

1044 N_total = len(P_full)

1045

1046 if self.model.use_supernodes and self.model.sparse_supernodes_full is not None:

1047 U, _ = build_U_supernodes_from_phi(

1048 phi_3d, self.model.sparse_supernodes_full, N_total,

1049 sigma_n_sq,

1050 )

1051 else:

1052 if self._U_buf is None or self._U_buf.shape[0] != N_total:

1053 self._U_buf = np.zeros((N_total, N_total), order='F')

1054

1055 U = build_U_from_phi_flat(

1056 phi_3d, self._block_phi_maps, N_total,

1057 sigma_n_sq, out=self._U_buf,

1058 )

1059

1060 y_ord = self.model.y_train[P_full]

1061 nll = nlml_from_U(U, y_ord)

1062

1063 alpha_ord = alpha_from_U(U, y_ord)

1064 alpha_v = np.empty_like(alpha_ord)

1065 alpha_v[P_full] = alpha_ord

1066

1067 return alpha_v, U, nll, phi, n_bases, oti, diffs

1068

1069 def _dense_nll_and_W(self, x0):

1070 """

1071 Dense Cholesky path: build full K, factor once, compute NLL and W.

1072 Used as a fallback when the sparsity pattern is too full.

1073

1074 Returns (W, alpha_v, nll, phi, n_bases, oti, diffs).

1075 """

1076 K, phi, n_bases, oti, diffs = self._build_K_and_phi(x0)

1077 N = K.shape[0]

1078

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

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

1081

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

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

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

1085

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

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

1088

1089 # Don't cache L/low here — predict's cache expects _cached_n_bases

1090 # which is only set by predict itself.

1091 self.model._cached_alpha = alpha_v

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

1093

1094 return W, alpha_v, nll, phi, n_bases, oti, diffs

1095

1096 @profile

1097 def _sparse_U_alpha_nll(self, K):

1098 """

1099 Build sparse U, compute alpha and NLML. Does NOT form K^{-1}.

1100

1101 Returns (alpha_v, U, nll) all in original index space.

1102 """

1103 P_full = self.model.mmd_P_full

1104 N_total = len(P_full)

1105

1106 if self._P_ix is None:

1107 self._P_ix = np.ix_(P_full, P_full)

1108 K_ord = K[self._P_ix]

1109 y_ord = self.model.y_train[P_full]

1110

1111 if self.model.use_supernodes and self.model.sparse_supernodes_full is not None:

1112 U, _ = build_U_supernodes(K_ord, self.model.sparse_supernodes_full, N_total)

1113 else:

1114 if self._U_buf is None or self._U_buf.shape[0] != N_total:

1115 self._U_buf = np.zeros((N_total, N_total))

1116 U = build_U(K_ord, self.model.sparse_S_full_arr, N_total,

1117 block_size=self.model.n_bases + 1, out=self._U_buf)

1118

1119 nll = nlml_from_U(U, y_ord)

1120

1121 # alpha in original space

1122 alpha_ord = alpha_from_U(U, y_ord)

1123 alpha_v = np.empty_like(alpha_ord)

1124 alpha_v[P_full] = alpha_ord

1125

1126 return alpha_v, U, nll

1127

1128 def _W_from_U(self, U, alpha_v):

1129 """

1130 Compute W = K^{-1} - αα^T from a pre-built sparse U.

1131

1132 U is in MMD order, alpha_v is in original index space.

1133 Returns W in original index space.

1134 """

1135 P_full = self.model.mmd_P_full

1136 N_total = len(P_full)

1137

1138 if self._K_inv_buf is None or self._K_inv_buf.shape[0] != N_total:

1139 self._K_inv_buf = np.empty((N_total, N_total), order='F')

1140 K_inv_ord = blas.dsyrk(1.0, U, lower=1,

1141 c=self._K_inv_buf, overwrite_c=1)

1142

1143 W = np.empty((N_total, N_total))

1144 _permute_and_subtract_outer(K_inv_ord, alpha_v, P_full, W)

1145 return W

1146

1147 def _sparse_W_and_alpha(self, K):

1148 """

1149 Compute W = K^{-1} - alpha*alpha^T and alpha using the sparse U factor.

1150

1151 K is in the ORIGINAL index space (size N_total x N_total).

1152 Returns (W, alpha_v, U, nll) all in original index space.

1153 """

1154 alpha_v, U, nll = self._sparse_U_alpha_nll(K)

1155 W = self._W_from_U(U, alpha_v)

1156 return W, alpha_v, U, nll

1157

1158 def nll_grad(self, x0):

1159 """Analytic gradient of the NLL using the sparse U factor."""

1160 try:

1161 if self.model._use_dense_factor:

1162 W, alpha_v, nll, phi, n_bases, oti, diffs = self._dense_nll_and_W(x0)

1163 return self._compute_grad(x0, W, phi, n_bases, oti, diffs)

1164 else:

1165 alpha_v, U, nll, phi, n_bases, oti, diffs = self._sparse_nlml_direct(x0)

1166 return self._compute_grad_blockwise(x0, U, alpha_v, phi, n_bases, oti, diffs)

1167 except Exception:

1168 return np.zeros(len(x0))

1169

1170 def nll_and_grad(self, x0):

1171 """

1172 Compute NLL and its gradient in a single pass.

1173

1174 Routes to either the dense Cholesky path or the sparse U path.

1175

1176 Returns

1177 -------

1178 nll : float

1179 grad : ndarray

1180 """

1181 try:

1182 if self.model._use_dense_factor:

1183 W, alpha_v, nll, phi, n_bases, oti, diffs = self._dense_nll_and_W(x0)

1184 grad = self._compute_grad(x0, W, phi, n_bases, oti, diffs)

1185 elif self.model.n_order > 0:

1186 alpha_v, U, nll, phi, n_bases, oti, diffs = self._sparse_nlml_direct(x0)

1187

1188 self.model._cached_U = U

1189 self.model._cached_P = self.model.mmd_P_full

1190 self.model._cached_alpha = alpha_v

1191 self.model._cached_L = None

1192 self.model._cached_low = None

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

1194

1195 grad = self._compute_grad_blockwise(x0, U, alpha_v, phi, n_bases, oti, diffs)

1196 else:

1197 K, phi, n_bases, oti, diffs = self._build_K_and_phi(x0)

1198 alpha_v, U, nll = self._sparse_U_alpha_nll(K)

1199

1200 self.model._cached_U = U

1201 self.model._cached_P = self.model.mmd_P_full

1202 self.model._cached_alpha = alpha_v

1203 self.model._cached_L = None

1204 self.model._cached_low = None

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

1206

1207 W = self._W_from_U(U, alpha_v)

1208 grad = self._compute_grad(x0, W, phi, n_bases, oti, diffs)

1209 except Exception:

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

1211

1212 if nll > 1e6:

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

1214 return float(nll), grad

1215

1216 def optimize_hyperparameters(self,

1217 optimizer="pso",

1218 **kwargs):

1219 """

1220 Optimize the GDDEGP model hyperparameters.

1221

1222 Returns:

1223 -------

1224 best_x : ndarray

1225 The optimal set of hyperparameters found.

1226 """

1227

1228 if isinstance(optimizer, str):

1229 if optimizer not in OPTIMIZERS:

1230 raise ValueError(

1231 f"Unknown optimizer '{optimizer}'. Available: {list(OPTIMIZERS.keys())}"

1232 )

1233 optimizer_fn = OPTIMIZERS[optimizer]

1234 else:

1235 optimizer_fn = optimizer

1236

1237 bounds = self.model.bounds

1238 lb = [b[0] for b in bounds]

1239 ub = [b[1] for b in bounds]

1240

1241 # Inject nll_and_grad for gradient-aware optimizers.

1242 if optimizer in ('lbfgs', 'jade', 'pso') and 'func_and_grad' not in kwargs and 'grad_func' not in kwargs:

1243 kwargs['func_and_grad'] = self.nll_and_grad

1244

1245 best_x, best_val = optimizer_fn(self.nll_wrapper, lb, ub, **kwargs)

1246

1247 self.model.opt_x0 = best_x

1248 self.model.opt_nll = best_val

1249

1250 return best_x