Coverage for jetgp/full_ddegp_sparse/optimizer.py: 49%
840 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-11 00:01 -0500
1import numpy as np
2import numba
3from scipy.linalg import cho_solve, cho_factor, blas
4from jetgp.full_ddegp_sparse import ddegp_utils as utils
5from jetgp.full_ddegp_sparse.sparse_cholesky import (
6 build_U, build_U_supernodes, nlml_from_U, alpha_from_U
7)
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.
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
84 signs = plan['signs']
85 sign_lookup = np.empty(n_types, dtype=np.float64)
86 sign_lookup[0] = signs[0]
87 for j in range(n_dt):
88 sign_lookup[j + 1] = signs[j + 1]
90 return k_type, k_phys, deriv_lookup, sign_lookup
93@numba.jit(nopython=True, cache=True)
94def _extract_K_sub(phi_exp_3d, nb_type, nb_phys, deriv_lookup, sign_lookup,
95 sigma_n_sq, sigma_data_diag, m, K_sub):
96 """
97 Assemble K_sub directly from phi_exp_3d for a neighbourhood `nb`.
99 nb_type[a], nb_phys[a] give the derivative type and physical point
100 for the a-th row/column of K_sub.
101 """
102 for a in range(m):
103 ta = nb_type[a]
104 pa = nb_phys[a]
105 for b in range(m):
106 tb = nb_type[b]
107 pb = nb_phys[b]
108 d = deriv_lookup[ta, tb]
109 K_sub[a, b] = phi_exp_3d[d, pa, pb] * sign_lookup[tb]
110 # Add noise to diagonal
111 K_sub[a, a] += sigma_n_sq + sigma_data_diag[a]
114@numba.jit(nopython=True, cache=True)
115def _project_G_to_W_proj(W_proj, G_b, nb_type, nb_phys,
116 deriv_lookup, sign_lookup, m):
117 """
118 Project per-block sensitivity G_b (m × m) into W_proj (ndir × n_func × n_func).
119 """
120 n_func = W_proj.shape[1]
121 plane = n_func * W_proj.shape[2]
122 wptr = W_proj.ravel()
123 for a_i in range(m):
124 ta = nb_type[a_i]
125 pa = nb_phys[a_i]
126 for bb_i in range(m):
127 tb = nb_type[bb_i]
128 pb = nb_phys[bb_i]
129 d = deriv_lookup[ta, tb]
130 wptr[d * plane + pa * W_proj.shape[2] + pb] += sign_lookup[tb] * G_b[a_i, bb_i]
133class Optimizer:
134 """
135 Optimizer class to perform hyperparameter tuning for sparse DDEGP models
136 by minimizing the negative log marginal likelihood (NLL).
138 Parameters
139 ----------
140 model : object
141 An instance of a model (e.g., ddegp) containing the necessary training data
142 and kernel configuration.
143 """
145 def __init__(self, model):
146 self.model = model
147 self._kernel_plan = None
148 self._deriv_buf = None
149 self._deriv_buf_shape = None
150 self._deriv_buf_ndir = None
151 self._deriv_factors = None
152 self._deriv_factors_key = None
153 self._K_buf = None
154 self._dK_buf = None
155 self._kernel_buf_size = None
156 self._W_proj_buf = None
157 self._W_proj_shape = None
158 self._U_buf = None
159 self._P_ix = None
160 self._K_inv_buf = None
161 # Direct phi extraction maps (built lazily)
162 self._k_index_map = None
163 self._inv_P = None
164 self._sigma_data_diag_mmd = None
165 self._block_phi_maps = None
167 def _get_deriv_buf(self, phi, n_bases, order):
168 """Return a pre-allocated buffer for get_all_derivs, reusing if shape matches."""
169 if self._deriv_buf_ndir is None:
170 from math import comb
171 self._deriv_buf_ndir = comb(n_bases + order, order)
172 shape = (self._deriv_buf_ndir, phi.shape[0], phi.shape[1])
173 if self._deriv_buf is None or self._deriv_buf_shape != shape:
174 self._deriv_buf = np.zeros(shape, dtype=np.float64)
175 self._deriv_buf_shape = shape
176 return self._deriv_buf
178 def _expand_derivs(self, phi, n_bases, deriv_order):
179 """Expand OTI derivatives, using fast struct path if available."""
180 if hasattr(phi, 'get_all_derivs_fast'):
181 buf = self._get_deriv_buf(phi, n_bases, deriv_order)
182 factors = self._get_deriv_factors(n_bases, deriv_order)
183 return phi.get_all_derivs_fast(factors, buf)
184 return phi.get_all_derivs(n_bases, deriv_order)
186 @staticmethod
187 def _enum_factors(max_basis, ordi):
188 """Enumerate derivative factors in struct memory order for a given order."""
189 from math import factorial
190 from collections import Counter
191 if ordi == 1:
192 for _ in range(max_basis):
193 yield 1.0
194 return
195 for last in range(1, max_basis + 1):
196 if ordi == 2:
197 for i in range(1, last + 1):
198 counts = Counter((i, last))
199 f = 1
200 for c in counts.values():
201 f *= factorial(c)
202 yield float(f)
203 else:
204 for prefix_factor, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1):
205 counts = dict(prefix_counts)
206 counts[last] = counts.get(last, 0) + 1
207 f = 1
208 for c in counts.values():
209 f *= factorial(c)
210 yield float(f)
212 @staticmethod
213 def _enum_factors_with_counts(max_basis, ordi):
214 """Enumerate (factor, counts_dict) pairs in struct order."""
215 from math import factorial
216 from collections import Counter
217 if ordi == 1:
218 for i in range(1, max_basis + 1):
219 yield 1.0, {i: 1}
220 return
221 for last in range(1, max_basis + 1):
222 for _, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1):
223 counts = dict(prefix_counts)
224 counts[last] = counts.get(last, 0) + 1
225 f = 1
226 for c in counts.values():
227 f *= factorial(c)
228 yield float(f), counts
230 def _get_deriv_factors(self, n_bases, order):
231 """Return cached precomputed derivative factorial factors."""
232 key = (n_bases, order)
233 if self._deriv_factors is not None and self._deriv_factors_key == key:
234 return self._deriv_factors
235 factors = [1.0] # order 0: real part
236 for ordi in range(1, order + 1):
237 factors.extend(self._enum_factors(n_bases, ordi))
238 self._deriv_factors = np.array(factors, dtype=np.float64)
239 self._deriv_factors_key = key
240 return self._deriv_factors
242 def _ensure_kernel_plan(self, n_bases):
243 """Lazily precompute kernel plan (once per n_bases)."""
244 if self._kernel_plan is not None and self._kernel_plan_n_bases == n_bases:
245 return
246 if not hasattr(utils, 'precompute_kernel_plan'):
247 self._kernel_plan = None
248 return
249 self._kernel_plan = utils.precompute_kernel_plan(
250 self.model.n_order, n_bases,
251 self.model.flattened_der_indices,
252 self.model.powers,
253 self.model.derivative_locations,
254 )
255 self._kernel_plan_n_bases = n_bases
256 # Reset kernel buffers when plan changes
257 self._K_buf = None
258 self._dK_buf = None
259 self._kernel_buf_size = None
261 def _ensure_kernel_bufs(self, n_rows_func):
262 """Pre-allocate reusable K and dK buffers (avoids repeated malloc)."""
263 if self._kernel_plan is None:
264 return
265 total = n_rows_func + self._kernel_plan['n_pts_with_derivs']
266 if self._kernel_buf_size != total:
267 self._K_buf = np.empty((total, total))
268 self._dK_buf = np.empty((total, total))
269 self._kernel_buf_size = total
270 # Cache absolute offsets in plan so rbf_kernel_fast doesn't recompute
271 if 'row_offsets_abs' not in self._kernel_plan:
272 self._kernel_plan['row_offsets_abs'] = self._kernel_plan['row_offsets'] + n_rows_func
273 self._kernel_plan['col_offsets_abs'] = self._kernel_plan['col_offsets'] + n_rows_func
275 def _build_K(self, phi_exp, phi, n_bases):
276 """Build kernel matrix using fast path if available."""
277 self._ensure_kernel_plan(n_bases)
278 if self._kernel_plan is not None:
279 base_shape = phi.shape
280 self._ensure_kernel_bufs(base_shape[0])
281 phi_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1])
282 return utils.rbf_kernel_fast(phi_3d, self._kernel_plan, out=self._K_buf)
283 return utils.rbf_kernel(
284 phi, phi_exp, self.model.n_order, n_bases,
285 self.model.flattened_der_indices, self.model.powers,
286 self.model.derivative_locations,
287 )
289 def _ensure_phi_index_maps(self, n_rows_func):
290 """Lazily build the K-index-to-phi maps and inverse permutation."""
291 if self._k_index_map is not None:
292 return
293 plan = self._kernel_plan
294 k_type, k_phys, deriv_lookup, sign_lookup = _build_k_index_map(
295 plan, n_rows_func)
296 self._k_index_map = (k_type, k_phys, deriv_lookup, sign_lookup)
298 P_full = self.model.mmd_P_full
299 inv_P = np.empty_like(P_full)
300 inv_P[P_full] = np.arange(len(P_full))
301 self._inv_P = inv_P
303 # sigma_data diagonal in MMD order
304 sd = self.model.sigma_data
305 if sd.ndim == 2:
306 sd_diag_orig = np.diag(sd) ** 2 if np.any(sd) else np.zeros(len(P_full))
307 else:
308 sd_diag_orig = np.zeros(len(P_full))
309 self._sigma_data_diag_mmd = sd_diag_orig[P_full]
311 # Precompute flat index arrays for phi_exp_3d gather.
312 stride_d = n_rows_func * n_rows_func
313 stride_row = n_rows_func
315 if (self.model.use_supernodes
316 and self.model.sparse_supernodes_full is not None):
317 for sn in self.model.sparse_supernodes_full:
318 ch = sn.get('children_arr')
319 if ch is None:
320 ch = np.asarray(sn['children'])
321 orig_ch = P_full[ch]
322 ch_type = k_type[orig_ch]
323 ch_phys = k_phys[orig_ch]
324 m = len(ch)
326 d_mat = deriv_lookup[ch_type[:, None], ch_type[None, :]]
327 pa_mat = np.broadcast_to(ch_phys[:, None], (m, m))
328 pb_mat = np.broadcast_to(ch_phys[None, :], (m, m))
329 sn['phi_flat_idx'] = np.ascontiguousarray(
330 d_mat * stride_d + pa_mat * stride_row + pb_mat
331 )
332 sn['phi_sign_mat'] = np.ascontiguousarray(
333 sign_lookup[ch_type[None, :]] * np.ones((m, 1))
334 )
335 sn['phi_sd_diag'] = self._sigma_data_diag_mmd[ch]
337 # Same precomputation for non-supernode block path
338 if (not self.model.use_supernodes
339 and self.model.n_order > 0):
340 N_total = len(P_full)
341 block_size = self.model.n_bases + 1
342 S = self.model.sparse_S_full_arr
343 self._block_phi_maps = []
344 for start in range(0, N_total, block_size):
345 end = min(start + block_size, N_total)
346 nb = S[end - 1] if isinstance(S[end - 1], np.ndarray) else np.asarray(S[end - 1])
347 m = len(nb)
349 orig_nb = P_full[nb]
350 nb_type = k_type[orig_nb]
351 nb_phys = k_phys[orig_nb]
353 d_mat = deriv_lookup[nb_type[:, None], nb_type[None, :]]
354 pa_mat = np.broadcast_to(nb_phys[:, None], (m, m))
355 pb_mat = np.broadcast_to(nb_phys[None, :], (m, m))
357 self._block_phi_maps.append({
358 'nb': nb,
359 'start': start,
360 'flat_idx': np.ascontiguousarray(
361 d_mat * stride_d + pa_mat * stride_row + pb_mat
362 ),
363 'sign_mat': np.ascontiguousarray(
364 sign_lookup[nb_type[None, :]] * np.ones((m, 1))
365 ),
366 'sd_diag': self._sigma_data_diag_mmd[nb],
367 'positions': np.searchsorted(nb, np.arange(start, end)),
368 })
370 @profile
371 def negative_log_marginal_likelihood(self, x0):
372 """
373 Compute the negative log marginal likelihood (NLL) via sparse U.
375 NLL = 0.5 * ||U^T y||^2 - sum(log|diag(U)|) + 0.5 * N * log(2pi)
377 Parameters
378 ----------
379 x0 : ndarray
380 Vector of log-scaled hyperparameters (length scales and noise).
382 Returns
383 -------
384 float
385 Value of the negative log marginal likelihood.
386 """
387 try:
388 if self.model._use_dense_factor:
389 W, alpha, nll, *_ = self._dense_nll_and_W(x0)
390 if nll > 1e6:
391 return 1e6
392 return nll
394 # Use direct phi path (skip full K construction) when possible
395 if self.model.n_order > 0:
396 alpha, U, nlml, *_ = self._sparse_nlml_direct(x0)
397 else:
398 K, _, _, _, _ = self._build_K_and_phi(x0)
399 alpha, U, nlml = self._sparse_U_alpha_nll(K)
401 # Sparse U can silently produce bad factors when K is
402 # ill-conditioned (e.g. very small noise). Clamp to 1e6
403 # to match the dense fallback behaviour.
404 if nlml > 1e6:
405 return 1e6
407 self.model._cached_U = U
408 self.model._cached_P = self.model.mmd_P_full
409 self.model._cached_alpha = alpha
410 self.model._cached_L = None
411 self.model._cached_low = None
412 self.model._cached_params = x0.copy()
414 return nlml
415 except Exception:
416 return 1e6
418 def nll_wrapper(self, x0):
419 """
420 Wrapper function to compute NLL for optimizer.
422 Parameters
423 ----------
424 x0 : ndarray
425 Hyperparameter vector.
427 Returns
428 -------
429 float
430 NLL evaluated at x0.
431 """
432 return self.negative_log_marginal_likelihood(x0)
434 def _compute_grad(self, x0, W, phi, n_bases, oti, diffs):
435 """
436 Compute the NLL gradient given pre-factorised W = K^{-1} - alpha alpha^T.
437 """
438 ln10 = np.log(10.0)
439 kernel = self.model.kernel
440 kernel_type = self.model.kernel_type
441 D = len(diffs)
442 sigma_n_sq = (10.0 ** x0[-1]) ** 2
444 grad = np.zeros(len(x0))
445 use_fast = self._kernel_plan is not None
446 base_shape = (W.shape[0] - self._kernel_plan['n_pts_with_derivs'],) * 2 if use_fast else None
448 deriv_order = 2 * self.model.n_order
450 # Precompute W projected into phi_exp space
451 W_proj = None
452 if use_fast and self.model.n_order > 0:
453 from math import comb
454 ndir = comb(n_bases + deriv_order, deriv_order)
455 proj_shape = (ndir, base_shape[0], base_shape[1])
456 if self._W_proj_buf is None or self._W_proj_shape != proj_shape:
457 self._W_proj_buf = np.empty(proj_shape)
458 self._W_proj_shape = proj_shape
459 W_proj = self._W_proj_buf
461 plan = self._kernel_plan
462 row_off = plan.get('row_offsets_abs', plan['row_offsets'] + base_shape[0])
463 col_off = plan.get('col_offsets_abs', plan['col_offsets'] + base_shape[1])
465 utils._project_W_to_phi_space(
466 W, W_proj, base_shape[0], base_shape[1],
467 plan['fd_flat_indices'], plan['df_flat_indices'],
468 plan['dd_flat_indices'],
469 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'],
470 plan['signs'], plan['n_deriv_types'], row_off, col_off,
471 )
473 _use_vdot_fused = W_proj is not None and hasattr(phi, 'vdot_expand_fast')
474 if _use_vdot_fused:
475 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order)
477 def _gc(dphi):
478 if _use_vdot_fused:
479 return 0.5 * dphi.vdot_expand_fast(_vdot_factors, W_proj)
480 if self.model.n_order == 0:
481 dphi_exp = dphi.real[np.newaxis, :, :]
482 else:
483 dphi_exp = self._expand_derivs(dphi, n_bases, deriv_order)
484 if W_proj is not None:
485 dphi_3d = dphi_exp.reshape(W_proj.shape)
486 return 0.5 * np.vdot(W_proj, dphi_3d)
487 elif use_fast:
488 dphi_3d = dphi_exp.reshape(dphi_exp.shape[0], base_shape[0], base_shape[1])
489 dK = utils.rbf_kernel_fast(dphi_3d, self._kernel_plan, out=self._dK_buf)
490 return 0.5 * np.vdot(W, dK)
491 else:
492 dK = utils.rbf_kernel(
493 dphi, dphi_exp,
494 self.model.n_order, n_bases,
495 self.model.flattened_der_indices, self.model.powers,
496 index=self.model.derivative_locations,
497 )
498 return 0.5 * np.vdot(W, dK)
500 # signal variance
501 grad[-2] = _gc(oti.mul(2.0 * ln10, phi))
502 # noise variance
503 grad[-1] = ln10 * sigma_n_sq * np.trace(W)
505 if kernel == 'SE':
506 if kernel_type == 'anisotropic':
507 ell = 10.0 ** x0[:D]
508 if hasattr(phi, 'fused_scale_sq_mul'):
509 dphi_buf = oti.zeros(phi.shape)
510 for d in range(D):
511 dphi_buf.fused_scale_sq_mul(diffs[d], phi, -ln10 * ell[d] ** 2)
512 grad[d] = _gc(dphi_buf)
513 else:
514 for d in range(D):
515 d_sq = oti.mul(diffs[d], diffs[d])
516 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi))
517 grad[d] = _gc(dphi_d)
518 else:
519 ell = 10.0 ** float(x0[0])
520 if hasattr(phi, 'fused_sum_sq'):
521 sum_sq = oti.zeros(phi.shape)
522 sum_sq.fused_sum_sq(diffs)
523 else:
524 sum_sq = oti.mul(diffs[0], diffs[0])
525 for d in range(1, D):
526 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
527 grad[0] = _gc(oti.mul(-ln10 * ell ** 2, oti.mul(sum_sq, phi)))
529 elif kernel == 'RQ':
530 if kernel_type == 'anisotropic':
531 ell = 10.0 ** x0[:D]
532 alpha_rq = 10.0 ** float(x0[D])
533 alpha_idx = D
534 else:
535 ell_val = 10.0 ** float(x0[0])
536 ell = np.full(D, ell_val)
537 alpha_rq = np.exp(float(x0[1]))
538 alpha_idx = 1
540 if hasattr(phi, 'fused_sqdist'):
541 r2 = oti.zeros(phi.shape)
542 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
543 r2.fused_sqdist(diffs, ell_sq)
544 else:
545 r2 = oti.mul(ell[0], diffs[0])
546 r2 = oti.mul(r2, r2)
547 for d in range(1, D):
548 td = oti.mul(ell[d], diffs[d])
549 r2 = oti.sum(r2, oti.mul(td, td))
550 base = oti.sum(1.0, oti.mul(r2, 1.0 / (2.0 * alpha_rq)))
551 inv_base = oti.pow(base, -1)
552 phi_over_base = oti.mul(phi, inv_base)
554 if kernel_type == 'anisotropic':
555 if hasattr(phi, 'fused_scale_sq_mul'):
556 dphi_buf = oti.zeros(phi.shape)
557 for d in range(D):
558 dphi_buf.fused_scale_sq_mul(diffs[d], phi_over_base, -ln10 * ell[d] ** 2)
559 grad[d] = _gc(dphi_buf)
560 else:
561 for d in range(D):
562 d_sq = oti.mul(diffs[d], diffs[d])
563 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi_over_base))
564 grad[d] = _gc(dphi_d)
565 else:
566 if hasattr(phi, 'fused_sum_sq'):
567 sum_sq = oti.zeros(phi.shape)
568 sum_sq.fused_sum_sq(diffs)
569 else:
570 sum_sq = oti.mul(diffs[0], diffs[0])
571 for d in range(1, D):
572 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
573 grad[0] = _gc(oti.mul(-ln10 * ell[0] ** 2, oti.mul(sum_sq, phi_over_base)))
575 log_base = oti.log(base)
576 term = oti.sub(oti.sub(1.0, inv_base), log_base)
577 alpha_factor = ln10 * alpha_rq if kernel_type == 'anisotropic' else alpha_rq
578 grad[alpha_idx] = _gc(oti.mul(alpha_factor, oti.mul(phi, term)))
580 elif kernel == 'SineExp':
581 if kernel_type == 'anisotropic':
582 ell = 10.0 ** x0[:D]
583 p = 10.0 ** x0[D:2 * D]
584 pip = np.pi / p
585 p_start = D
586 else:
587 ell_val = 10.0 ** float(x0[0])
588 p_val = 10.0 ** float(x0[1])
589 pip_val = np.pi / p_val
590 ell = np.full(D, ell_val)
591 pip = np.full(D, pip_val)
592 p_start = 1
594 sin_d = []
595 cos_d = []
596 for d in range(D):
597 arg = oti.mul(pip[d], diffs[d])
598 sin_d.append(oti.sin(arg))
599 cos_d.append(oti.cos(arg))
601 if kernel_type == 'anisotropic':
602 if hasattr(phi, 'fused_scale_sq_mul'):
603 dphi_buf = oti.zeros(phi.shape)
604 for d in range(D):
605 dphi_buf.fused_scale_sq_mul(sin_d[d], phi, -4.0 * ln10 * ell[d] ** 2)
606 grad[d] = _gc(dphi_buf)
607 else:
608 for d in range(D):
609 sin_sq = oti.mul(sin_d[d], sin_d[d])
610 grad[d] = _gc(oti.mul(-4.0 * ln10 * ell[d] ** 2,
611 oti.mul(sin_sq, phi)))
612 else:
613 if hasattr(phi, 'fused_sum_sq'):
614 sum_sin_sq = oti.zeros(phi.shape)
615 sum_sin_sq.fused_sum_sq(sin_d)
616 else:
617 sum_sin_sq = oti.mul(sin_d[0], sin_d[0])
618 for d in range(1, D):
619 sum_sin_sq = oti.sum(sum_sin_sq, oti.mul(sin_d[d], sin_d[d]))
620 grad[0] = _gc(oti.mul(-4.0 * ln10 * ell[0] ** 2,
621 oti.mul(sum_sin_sq, phi)))
623 if kernel_type == 'anisotropic':
624 for d in range(D):
625 sc_diff = oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d]))
626 scale = 4.0 * ln10 * ell[d] ** 2 * pip[d]
627 grad[p_start + d] = _gc(oti.mul(scale, oti.mul(sc_diff, phi)))
628 else:
629 sum_scd = oti.mul(sin_d[0], oti.mul(cos_d[0], diffs[0]))
630 for d in range(1, D):
631 sum_scd = oti.sum(sum_scd,
632 oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d])))
633 scale = 4.0 * ln10 * ell[0] ** 2 * pip[0]
634 grad[p_start] = _gc(oti.mul(scale, oti.mul(sum_scd, phi)))
636 elif kernel == 'Matern':
637 kf = self.model.kernel_factory
638 if not hasattr(kf, '_matern_grad_prebuild'):
639 kf._matern_grad_prebuild = matern_kernel_grad_builder(
640 kf.nu, oti_module=oti)
642 if kernel_type == 'anisotropic':
643 ell = 10.0 ** x0[:D]
644 else:
645 ell = np.full(D, 10.0 ** float(x0[0]))
647 sigma_f_sq = (10.0 ** float(x0[-2])) ** 2
648 _eps = 1e-10
650 if hasattr(phi, 'fused_sqdist'):
651 r2 = oti.zeros(phi.shape)
652 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
653 r2.fused_sqdist(diffs, ell_sq)
654 else:
655 r2 = oti.mul(ell[0], diffs[0])
656 r2 = oti.mul(r2, r2)
657 for d in range(1, D):
658 td = oti.mul(ell[d], diffs[d])
659 r2 = oti.sum(r2, oti.mul(td, td))
660 r_oti = oti.sqrt(oti.sum(r2, _eps ** 2))
661 f_prime_r = kf._matern_grad_prebuild(r_oti)
662 inv_r = oti.pow(r_oti, -1)
663 base_matern = oti.mul(sigma_f_sq, oti.mul(f_prime_r, inv_r))
665 if kernel_type == 'anisotropic':
666 if hasattr(phi, 'fused_scale_sq_mul'):
667 dphi_buf = oti.zeros(phi.shape)
668 for d in range(D):
669 dphi_buf.fused_scale_sq_mul(diffs[d], base_matern, ln10 * ell[d] ** 2)
670 grad[d] = _gc(dphi_buf)
671 else:
672 for d in range(D):
673 d_sq = oti.mul(diffs[d], diffs[d])
674 dphi_d = oti.mul(ln10 * ell[d] ** 2, oti.mul(d_sq, base_matern))
675 grad[d] = _gc(dphi_d)
676 else:
677 ell_val = ell[0]
678 if hasattr(phi, 'fused_sum_sq'):
679 sum_dsq = oti.zeros(phi.shape)
680 sum_dsq.fused_sum_sq(diffs)
681 else:
682 sum_dsq = oti.mul(diffs[0], diffs[0])
683 for d in range(1, D):
684 sum_dsq = oti.sum(sum_dsq, oti.mul(diffs[d], diffs[d]))
685 dphi_e = oti.mul(ln10 * ell_val ** 2, oti.mul(sum_dsq, base_matern))
686 grad[0] = _gc(dphi_e)
688 elif kernel == 'SI':
689 kf = self.model.kernel_factory
690 si_prebuild = kf.SI_kernel_prebuild
691 if kernel_type == 'anisotropic':
692 ell = 10.0 ** x0[:D]
693 else:
694 ell = np.full(D, 10.0 ** float(x0[0]))
695 si_vals = [si_prebuild(diffs[d]) for d in range(D)]
696 term_vals = [oti.sum(1.0, oti.mul(ell[d], si_vals[d])) for d in range(D)]
697 if kernel_type == 'anisotropic':
698 for d in range(D):
699 phi_over_term = oti.div(phi, term_vals[d])
700 dphi_d = oti.mul(ln10 * ell[d],
701 oti.mul(si_vals[d], phi_over_term))
702 grad[d] = _gc(dphi_d)
703 else:
704 ell_val = ell[0]
705 acc = oti.mul(si_vals[0], oti.div(phi, term_vals[0]))
706 for d in range(1, D):
707 acc = oti.sum(acc, oti.mul(si_vals[d],
708 oti.div(phi, term_vals[d])))
709 grad[0] = _gc(oti.mul(ln10 * ell_val, acc))
711 return grad
713 @profile
714 def _compute_grad_blockwise(self, x0, U, alpha_v, phi, n_bases, oti, diffs):
715 """
716 Compute the NLL gradient by differentiating through each block's
717 Cholesky in the Vecchia decomposition.
719 For each block b the Vecchia NLL contribution is:
720 NLL_b = 0.5 * Σ_j [ α_b[p_j]² / s_j - log(s_j) ]
722 where α_b = K_sub⁻¹ y_nb, s_j = (K_sub⁻¹)[p_j, p_j].
724 Differentiating through K_sub⁻¹ gives the per-block sensitivity:
725 G_b = 0.5 * M V^T
727 where V = K_sub⁻¹ E (un-normalised U block), and
728 M[:, j] = γ_j V[:, j] - 2 β_j α_b
729 β_j = α_b[p_j] / s_j, γ_j = β_j² + 1/s_j
731 The gradient is dNLL/dθ = Σ_b tr(G_b dK_sub_b/dθ),
732 which is projected into phi-space as W_proj for vdot.
733 """
734 from math import comb
736 ln10 = np.log(10.0)
737 kernel = self.model.kernel
738 kernel_type = self.model.kernel_type
739 D = len(diffs)
740 sigma_n_sq = (10.0 ** x0[-1]) ** 2
742 grad = np.zeros(len(x0))
743 deriv_order = 2 * self.model.n_order
744 plan = self._kernel_plan
745 P_full = self.model.mmd_P_full
746 N_total = len(P_full)
747 n_func = phi.shape[0]
749 ndir = comb(n_bases + deriv_order, deriv_order)
750 k_type, k_phys, deriv_lookup, sign_lookup = self._k_index_map
752 # ── phi_exp for K_sub reconstruction ──────────────────────
753 phi_exp = self._expand_derivs(phi, n_bases, deriv_order)
754 phi_3d = phi_exp.reshape(phi_exp.shape[0], n_func, n_func)
755 phi_flat = phi_3d.ravel()
757 block_maps = self._block_phi_maps
758 y_ord = self.model.y_train[P_full]
760 # ── accumulate W_proj and noise trace from per-block G_b ──
761 proj_shape = (ndir, n_func, n_func)
762 W_proj = np.zeros(proj_shape)
763 w_flat = W_proj.ravel()
764 noise_trace = 0.0
766 for bm in block_maps:
767 nb = bm['nb']
768 m = len(nb)
769 positions = bm['positions']
770 n_cols = len(positions)
772 # Reconstruct K_sub for this block
773 K_sub = phi_flat[bm['flat_idx']].reshape(m, m) * bm['sign_mat']
774 diag_idx = np.arange(m)
775 K_sub[diag_idx, diag_idx] += sigma_n_sq + bm['sd_diag']
777 # Cholesky factor
778 L_u, low_u = cho_factor(K_sub, lower=True)
780 # V = K_sub⁻¹ E (un-normalised U block, m × n_cols)
781 E = np.zeros((m, n_cols))
782 E[positions, np.arange(n_cols)] = 1.0
783 V = cho_solve((L_u, low_u), E)
785 # α_b = K_sub⁻¹ y_nb (m,)
786 y_nb = y_ord[nb]
787 alpha_b = cho_solve((L_u, low_u), y_nb)
789 # Per-parent scalars
790 s = V[positions, np.arange(n_cols)]
791 a = alpha_b[positions]
792 beta = a / s
793 gamma = beta ** 2 + 1.0 / s
795 # M[:, j] = γ_j V[:, j] - 2 β_j α_b
796 M = V * gamma[np.newaxis, :] - 2.0 * alpha_b[:, np.newaxis] * beta[np.newaxis, :]
798 # Noise trace
799 noise_trace += np.sum(M * V)
801 # Project G_b into W_proj using precomputed flat indices
802 G_b = M @ V.T
803 np.add.at(w_flat, bm['flat_idx'].ravel(), (G_b * bm['sign_mat']).ravel())
805 # ── noise gradient ───────────────────────────────────────
806 grad[-1] = ln10 * sigma_n_sq * noise_trace
808 # ── deriv factors + fast vdot path ───────────────────────
809 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order)
811 @profile
812 def _gc_block(dphi):
813 return 0.5 * dphi.vdot_expand_fast(_vdot_factors, W_proj)
815 # ── signal variance ──────────────────────────────────────
816 grad[-2] = ln10 * phi.vdot_expand_fast(_vdot_factors, W_proj)
818 # ── kernel-specific hyperparameter gradients ─────────────
819 if kernel == 'SE':
820 if kernel_type == 'anisotropic':
821 ell = 10.0 ** x0[:D]
822 if hasattr(phi, 'fused_scale_sq_mul'):
823 dphi_buf = oti.zeros(phi.shape)
824 for d in range(D):
825 dphi_buf.fused_scale_sq_mul(diffs[d], phi, -ln10 * ell[d] ** 2)
826 grad[d] = _gc_block(dphi_buf)
827 else:
828 for d in range(D):
829 d_sq = oti.mul(diffs[d], diffs[d])
830 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi))
831 grad[d] = _gc_block(dphi_d)
832 else:
833 ell = 10.0 ** float(x0[0])
834 if hasattr(phi, 'fused_sum_sq'):
835 sum_sq = oti.zeros(phi.shape)
836 sum_sq.fused_sum_sq(diffs)
837 else:
838 sum_sq = oti.mul(diffs[0], diffs[0])
839 for d in range(1, D):
840 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
841 grad[0] = _gc_block(oti.mul(-ln10 * ell ** 2, oti.mul(sum_sq, phi)))
843 elif kernel == 'RQ':
844 if kernel_type == 'anisotropic':
845 ell = 10.0 ** x0[:D]
846 alpha_rq = 10.0 ** float(x0[D])
847 alpha_idx = D
848 else:
849 ell_val = 10.0 ** float(x0[0])
850 ell = np.full(D, ell_val)
851 alpha_rq = np.exp(float(x0[1]))
852 alpha_idx = 1
854 if hasattr(phi, 'fused_sqdist'):
855 r2 = oti.zeros(phi.shape)
856 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
857 r2.fused_sqdist(diffs, ell_sq)
858 else:
859 r2 = oti.mul(ell[0], diffs[0])
860 r2 = oti.mul(r2, r2)
861 for d in range(1, D):
862 td = oti.mul(ell[d], diffs[d])
863 r2 = oti.sum(r2, oti.mul(td, td))
864 base = oti.sum(1.0, oti.mul(r2, 1.0 / (2.0 * alpha_rq)))
865 inv_base = oti.pow(base, -1)
866 phi_over_base = oti.mul(phi, inv_base)
868 if kernel_type == 'anisotropic':
869 if hasattr(phi, 'fused_scale_sq_mul'):
870 dphi_buf = oti.zeros(phi.shape)
871 for d in range(D):
872 dphi_buf.fused_scale_sq_mul(diffs[d], phi_over_base, -ln10 * ell[d] ** 2)
873 grad[d] = _gc_block(dphi_buf)
874 else:
875 for d in range(D):
876 d_sq = oti.mul(diffs[d], diffs[d])
877 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi_over_base))
878 grad[d] = _gc_block(dphi_d)
879 else:
880 if hasattr(phi, 'fused_sum_sq'):
881 sum_sq = oti.zeros(phi.shape)
882 sum_sq.fused_sum_sq(diffs)
883 else:
884 sum_sq = oti.mul(diffs[0], diffs[0])
885 for d in range(1, D):
886 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
887 grad[0] = _gc_block(oti.mul(-ln10 * ell[0] ** 2, oti.mul(sum_sq, phi_over_base)))
889 log_base = oti.log(base)
890 term = oti.sub(oti.sub(1.0, inv_base), log_base)
891 alpha_factor = ln10 * alpha_rq if kernel_type == 'anisotropic' else alpha_rq
892 grad[alpha_idx] = _gc_block(oti.mul(alpha_factor, oti.mul(phi, term)))
894 elif kernel == 'SineExp':
895 if kernel_type == 'anisotropic':
896 ell = 10.0 ** x0[:D]
897 p = 10.0 ** x0[D:2 * D]
898 pip = np.pi / p
899 p_start = D
900 else:
901 ell_val = 10.0 ** float(x0[0])
902 p_val = 10.0 ** float(x0[1])
903 pip_val = np.pi / p_val
904 ell = np.full(D, ell_val)
905 pip = np.full(D, pip_val)
906 p_start = 1
908 sin_d = []
909 cos_d = []
910 for d in range(D):
911 arg = oti.mul(pip[d], diffs[d])
912 sin_d.append(oti.sin(arg))
913 cos_d.append(oti.cos(arg))
915 if kernel_type == 'anisotropic':
916 if hasattr(phi, 'fused_scale_sq_mul'):
917 dphi_buf = oti.zeros(phi.shape)
918 for d in range(D):
919 dphi_buf.fused_scale_sq_mul(sin_d[d], phi, -4.0 * ln10 * ell[d] ** 2)
920 grad[d] = _gc_block(dphi_buf)
921 else:
922 for d in range(D):
923 sin_sq = oti.mul(sin_d[d], sin_d[d])
924 grad[d] = _gc_block(oti.mul(-4.0 * ln10 * ell[d] ** 2,
925 oti.mul(sin_sq, phi)))
926 else:
927 if hasattr(phi, 'fused_sum_sq'):
928 sum_sin_sq = oti.zeros(phi.shape)
929 sum_sin_sq.fused_sum_sq(sin_d)
930 else:
931 sum_sin_sq = oti.mul(sin_d[0], sin_d[0])
932 for d in range(1, D):
933 sum_sin_sq = oti.sum(sum_sin_sq, oti.mul(sin_d[d], sin_d[d]))
934 grad[0] = _gc_block(oti.mul(-4.0 * ln10 * ell[0] ** 2,
935 oti.mul(sum_sin_sq, phi)))
937 if kernel_type == 'anisotropic':
938 for d in range(D):
939 sc_diff = oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d]))
940 scale = 4.0 * ln10 * ell[d] ** 2 * pip[d]
941 grad[p_start + d] = _gc_block(oti.mul(scale, oti.mul(sc_diff, phi)))
942 else:
943 sum_scd = oti.mul(sin_d[0], oti.mul(cos_d[0], diffs[0]))
944 for d in range(1, D):
945 sum_scd = oti.sum(sum_scd,
946 oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d])))
947 scale = 4.0 * ln10 * ell[0] ** 2 * pip[0]
948 grad[p_start] = _gc_block(oti.mul(scale, oti.mul(sum_scd, phi)))
950 elif kernel == 'Matern':
951 kf = self.model.kernel_factory
952 if not hasattr(kf, '_matern_grad_prebuild'):
953 kf._matern_grad_prebuild = matern_kernel_grad_builder(
954 kf.nu, oti_module=oti)
956 if kernel_type == 'anisotropic':
957 ell = 10.0 ** x0[:D]
958 else:
959 ell = np.full(D, 10.0 ** float(x0[0]))
961 sigma_f_sq = (10.0 ** float(x0[-2])) ** 2
962 _eps = 1e-10
964 if hasattr(phi, 'fused_sqdist'):
965 r2 = oti.zeros(phi.shape)
966 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
967 r2.fused_sqdist(diffs, ell_sq)
968 else:
969 r2 = oti.mul(ell[0], diffs[0])
970 r2 = oti.mul(r2, r2)
971 for d in range(1, D):
972 td = oti.mul(ell[d], diffs[d])
973 r2 = oti.sum(r2, oti.mul(td, td))
974 r_oti = oti.sqrt(oti.sum(r2, _eps ** 2))
975 f_prime_r = kf._matern_grad_prebuild(r_oti)
976 inv_r = oti.pow(r_oti, -1)
977 base_matern = oti.mul(sigma_f_sq, oti.mul(f_prime_r, inv_r))
979 if kernel_type == 'anisotropic':
980 if hasattr(phi, 'fused_scale_sq_mul'):
981 dphi_buf = oti.zeros(phi.shape)
982 for d in range(D):
983 dphi_buf.fused_scale_sq_mul(diffs[d], base_matern, ln10 * ell[d] ** 2)
984 grad[d] = _gc_block(dphi_buf)
985 else:
986 for d in range(D):
987 d_sq = oti.mul(diffs[d], diffs[d])
988 dphi_d = oti.mul(ln10 * ell[d] ** 2, oti.mul(d_sq, base_matern))
989 grad[d] = _gc_block(dphi_d)
990 else:
991 ell_val = ell[0]
992 if hasattr(phi, 'fused_sum_sq'):
993 sum_dsq = oti.zeros(phi.shape)
994 sum_dsq.fused_sum_sq(diffs)
995 else:
996 sum_dsq = oti.mul(diffs[0], diffs[0])
997 for d in range(1, D):
998 sum_dsq = oti.sum(sum_dsq, oti.mul(diffs[d], diffs[d]))
999 dphi_e = oti.mul(ln10 * ell_val ** 2, oti.mul(sum_dsq, base_matern))
1000 grad[0] = _gc_block(dphi_e)
1002 elif kernel == 'SI':
1003 kf = self.model.kernel_factory
1004 si_prebuild = kf.SI_kernel_prebuild
1005 if kernel_type == 'anisotropic':
1006 ell = 10.0 ** x0[:D]
1007 else:
1008 ell = np.full(D, 10.0 ** float(x0[0]))
1009 si_vals = [si_prebuild(diffs[d]) for d in range(D)]
1010 term_vals = [oti.sum(1.0, oti.mul(ell[d], si_vals[d])) for d in range(D)]
1011 if kernel_type == 'anisotropic':
1012 for d in range(D):
1013 phi_over_term = oti.div(phi, term_vals[d])
1014 dphi_d = oti.mul(ln10 * ell[d],
1015 oti.mul(si_vals[d], phi_over_term))
1016 grad[d] = _gc_block(dphi_d)
1017 else:
1018 ell_val = ell[0]
1019 acc = oti.mul(si_vals[0], oti.div(phi, term_vals[0]))
1020 for d in range(1, D):
1021 acc = oti.sum(acc, oti.mul(si_vals[d],
1022 oti.div(phi, term_vals[d])))
1023 grad[0] = _gc_block(oti.mul(ln10 * ell_val, acc))
1025 return grad
1027 @profile
1028 def _build_K_and_phi(self, x0):
1029 """
1030 Shared helper: build K (with noise), phi, n_bases, oti, diffs.
1032 Returns (K, phi, n_bases, oti, diffs).
1033 """
1034 diffs = self.model.differences_by_dim
1035 oti = self.model.kernel_factory.oti
1036 sigma_n_sq = (10.0 ** x0[-1]) ** 2
1038 phi = self.model.kernel_func(diffs, x0[:-1])
1039 if self.model.n_order == 0:
1040 n_bases = 0
1041 phi_exp = phi.real[np.newaxis, :, :]
1042 else:
1043 n_bases = phi.get_active_bases()[-1]
1044 deriv_order = 2 * self.model.n_order
1045 phi_exp = self._expand_derivs(phi, n_bases, deriv_order)
1047 self._ensure_kernel_plan(n_bases)
1048 if self._kernel_plan is not None:
1049 base_shape = phi.shape
1050 self._ensure_kernel_bufs(base_shape[0])
1051 phi_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1])
1052 K = utils.rbf_kernel_fast(phi_3d, self._kernel_plan, out=self._K_buf)
1053 else:
1054 K = utils.rbf_kernel(
1055 phi, phi_exp, self.model.n_order, n_bases,
1056 self.model.flattened_der_indices, self.model.powers,
1057 index=self.model.derivative_locations,
1058 )
1059 K.flat[::K.shape[0] + 1] += sigma_n_sq
1060 K.flat[::K.shape[0] + 1] += self.model.sigma_data_sq_diag
1061 return K, phi, n_bases, oti, diffs
1063 def _sparse_nlml_direct(self, x0):
1064 """
1065 Compute sparse NLML directly from phi_exp_3d, skipping full K
1066 construction and permutation.
1068 Returns (alpha_v, U, nll, phi, n_bases, oti, diffs).
1069 """
1070 from jetgp.full_ddegp_sparse.sparse_cholesky import (
1071 build_U_from_phi_flat, build_U_supernodes_from_phi,
1072 )
1074 diffs = self.model.differences_by_dim
1075 oti = self.model.kernel_factory.oti
1076 sigma_n_sq = (10.0 ** x0[-1]) ** 2
1078 phi = self.model.kernel_func(diffs, x0[:-1])
1079 n_bases = phi.get_active_bases()[-1]
1080 deriv_order = 2 * self.model.n_order
1081 phi_exp = self._expand_derivs(phi, n_bases, deriv_order)
1083 self._ensure_kernel_plan(n_bases)
1084 base_shape = phi.shape
1085 phi_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1])
1087 # Build index maps (once)
1088 self._ensure_phi_index_maps(base_shape[0])
1089 k_type, k_phys, deriv_lookup, sign_lookup = self._k_index_map
1091 P_full = self.model.mmd_P_full
1092 N_total = len(P_full)
1094 if self.model.use_supernodes and self.model.sparse_supernodes_full is not None:
1095 U, _ = build_U_supernodes_from_phi(
1096 phi_3d, self.model.sparse_supernodes_full, N_total,
1097 sigma_n_sq,
1098 )
1099 else:
1100 if self._U_buf is None or self._U_buf.shape[0] != N_total:
1101 self._U_buf = np.zeros((N_total, N_total), order='F')
1103 U = build_U_from_phi_flat(
1104 phi_3d, self._block_phi_maps, N_total,
1105 sigma_n_sq, out=self._U_buf,
1106 )
1108 y_ord = self.model.y_train[P_full]
1109 nll = nlml_from_U(U, y_ord)
1111 alpha_ord = alpha_from_U(U, y_ord)
1112 alpha_v = np.empty_like(alpha_ord)
1113 alpha_v[P_full] = alpha_ord
1115 return alpha_v, U, nll, phi, n_bases, oti, diffs
1117 def _dense_nll_and_W(self, x0):
1118 """
1119 Dense Cholesky path: build full K, factor once, compute NLL and W.
1120 Used as a fallback when the sparsity pattern is too full.
1122 Returns (W, alpha_v, nll, phi, n_bases, oti, diffs).
1123 """
1124 K, phi, n_bases, oti, diffs = self._build_K_and_phi(x0)
1125 N = K.shape[0]
1127 L, low = cho_factor(K, lower=True)
1128 alpha_v = cho_solve((L, low), self.model.y_train)
1130 nll = (0.5 * np.dot(self.model.y_train, alpha_v)
1131 + np.sum(np.log(np.diag(L)))
1132 + 0.5 * N * np.log(2 * np.pi))
1134 K_inv = cho_solve((L, low), np.eye(N))
1135 W = K_inv - np.outer(alpha_v, alpha_v)
1137 # Don't cache L/low here — predict's cache expects _cached_n_bases_rays
1138 # which is only set by predict itself.
1139 self.model._cached_alpha = alpha_v
1140 self.model._cached_params = x0.copy()
1142 return W, alpha_v, nll, phi, n_bases, oti, diffs
1144 @profile
1145 def _sparse_U_alpha_nll(self, K):
1146 """
1147 Build sparse U, compute alpha and NLML. Does NOT form K^{-1}.
1149 Returns (alpha_v, U, nll) all in original index space.
1150 """
1151 P_full = self.model.mmd_P_full
1152 N_total = len(P_full)
1154 if self._P_ix is None:
1155 self._P_ix = np.ix_(P_full, P_full)
1156 K_ord = K[self._P_ix]
1157 y_ord = self.model.y_train[P_full]
1159 if self.model.use_supernodes and self.model.sparse_supernodes_full is not None:
1160 U, _ = build_U_supernodes(K_ord, self.model.sparse_supernodes_full, N_total)
1161 else:
1162 if self._U_buf is None or self._U_buf.shape[0] != N_total:
1163 self._U_buf = np.zeros((N_total, N_total))
1164 U = build_U(K_ord, self.model.sparse_S_full_arr, N_total,
1165 block_size=self.model.n_bases + 1, out=self._U_buf)
1167 nll = nlml_from_U(U, y_ord)
1169 # alpha in original space
1170 alpha_ord = alpha_from_U(U, y_ord)
1171 alpha_v = np.empty_like(alpha_ord)
1172 alpha_v[P_full] = alpha_ord
1174 return alpha_v, U, nll
1176 def _W_from_U(self, U, alpha_v):
1177 """
1178 Compute W = K^{-1} - αα^T from a pre-built sparse U.
1180 U is in MMD order, alpha_v is in original index space.
1181 Returns W in original index space.
1182 """
1183 P_full = self.model.mmd_P_full
1184 N_total = len(P_full)
1186 if self._K_inv_buf is None or self._K_inv_buf.shape[0] != N_total:
1187 self._K_inv_buf = np.empty((N_total, N_total), order='F')
1188 K_inv_ord = blas.dsyrk(1.0, U, lower=1,
1189 c=self._K_inv_buf, overwrite_c=1)
1191 W = np.empty((N_total, N_total))
1192 _permute_and_subtract_outer(K_inv_ord, alpha_v, P_full, W)
1193 return W
1195 def _sparse_W_and_alpha(self, K):
1196 """
1197 Compute W = K^{-1} - alpha*alpha^T and alpha using the sparse U factor.
1199 K is in the ORIGINAL index space (size N_total x N_total).
1200 Returns (W, alpha_v, U, nll) all in original index space.
1201 """
1202 alpha_v, U, nll = self._sparse_U_alpha_nll(K)
1203 W = self._W_from_U(U, alpha_v)
1204 return W, alpha_v, U, nll
1206 def nll_grad(self, x0):
1207 """Analytic gradient of the NLL using the sparse U factor."""
1208 try:
1209 if self.model._use_dense_factor:
1210 W, alpha_v, nll, phi, n_bases, oti, diffs = self._dense_nll_and_W(x0)
1211 return self._compute_grad(x0, W, phi, n_bases, oti, diffs)
1212 else:
1213 alpha_v, U, nll, phi, n_bases, oti, diffs = self._sparse_nlml_direct(x0)
1214 return self._compute_grad_blockwise(x0, U, alpha_v, phi, n_bases, oti, diffs)
1215 except Exception:
1216 return np.zeros(len(x0))
1218 def nll_and_grad(self, x0):
1219 """
1220 Compute NLL and its gradient in a single pass.
1222 Routes to either the dense Cholesky path or the sparse U path.
1224 Returns
1225 -------
1226 nll : float
1227 grad : ndarray
1228 """
1229 try:
1230 if self.model._use_dense_factor:
1231 W, alpha_v, nll, phi, n_bases, oti, diffs = self._dense_nll_and_W(x0)
1232 grad = self._compute_grad(x0, W, phi, n_bases, oti, diffs)
1233 elif self.model.n_order > 0:
1234 alpha_v, U, nll, phi, n_bases, oti, diffs = self._sparse_nlml_direct(x0)
1236 self.model._cached_U = U
1237 self.model._cached_P = self.model.mmd_P_full
1238 self.model._cached_alpha = alpha_v
1239 self.model._cached_L = None
1240 self.model._cached_low = None
1241 self.model._cached_params = x0.copy()
1243 grad = self._compute_grad_blockwise(x0, U, alpha_v, phi, n_bases, oti, diffs)
1244 else:
1245 K, phi, n_bases, oti, diffs = self._build_K_and_phi(x0)
1246 alpha_v, U, nll = self._sparse_U_alpha_nll(K)
1248 self.model._cached_U = U
1249 self.model._cached_P = self.model.mmd_P_full
1250 self.model._cached_alpha = alpha_v
1251 self.model._cached_L = None
1252 self.model._cached_low = None
1253 self.model._cached_params = x0.copy()
1255 W = self._W_from_U(U, alpha_v)
1256 grad = self._compute_grad(x0, W, phi, n_bases, oti, diffs)
1257 except Exception:
1258 return 1e6, np.zeros(len(x0))
1260 if nll > 1e6:
1261 return 1e6, np.zeros(len(x0))
1262 return float(nll), grad
1264 def optimize_hyperparameters(self,
1265 optimizer="pso",
1266 **kwargs):
1267 """
1268 Optimize the DDEGP model hyperparameters.
1270 Returns:
1271 -------
1272 best_x : ndarray
1273 The optimal set of hyperparameters found.
1274 """
1276 if isinstance(optimizer, str):
1277 if optimizer not in OPTIMIZERS:
1278 raise ValueError(
1279 f"Unknown optimizer '{optimizer}'. Available: {list(OPTIMIZERS.keys())}"
1280 )
1281 optimizer_fn = OPTIMIZERS[optimizer]
1282 else:
1283 optimizer_fn = optimizer
1285 bounds = self.model.bounds
1286 lb = [b[0] for b in bounds]
1287 ub = [b[1] for b in bounds]
1289 # Inject nll_and_grad for gradient-aware optimizers.
1290 if optimizer in ('lbfgs', 'jade', 'pso') and 'func_and_grad' not in kwargs and 'grad_func' not in kwargs:
1291 kwargs['func_and_grad'] = self.nll_and_grad
1293 best_x, best_val = optimizer_fn(self.nll_wrapper, lb, ub, **kwargs)
1295 self.model.opt_x0 = best_x
1296 self.model.opt_nll = best_val
1298 return best_x