Coverage for jetgp/full_degp_sparse/optimizer.py: 51%
1029 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_degp_sparse import degp_utils as utils
5from jetgp.full_degp_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, cache=True)
15def _symmetrise_upper(A):
16 """Copy upper triangle to lower triangle in-place."""
17 n = A.shape[0]
18 for i in range(n):
19 for j in range(i + 1, n):
20 A[j, i] = A[i, j]
23@numba.jit(nopython=True, parallel=True, cache=True)
24def _permute_and_subtract_outer(K_inv_ord, alpha_v, P_full, W):
25 """Fused: W[P[i], P[j]] = K_inv_ord[i, j] - alpha_v[P[i]] * alpha_v[P[j]]
26 Reads only the lower triangle of K_inv_ord (as produced by dsyrk with lower=1
27 on a Fortran-order buffer). In column-major layout, lower-triangle entries
28 within each column are contiguous → cache-friendly sequential reads.
29 W is symmetric, so we write both (pi,pj) and (pj,pi).
30 """
31 N = len(P_full)
32 for i in numba.prange(N):
33 pi = P_full[i]
34 ai = alpha_v[pi]
35 # Diagonal
36 W[pi, pi] = K_inv_ord[i, i] - ai * ai
37 # j > i → read lower triangle: K_inv_ord[j, i] (row > col)
38 for j in range(i + 1, N):
39 pj = P_full[j]
40 val = K_inv_ord[j, i] - ai * alpha_v[pj]
41 W[pi, pj] = val
42 W[pj, pi] = val
45def _build_k_index_map(plan, n_rows_func):
46 """
47 Build arrays that map each K-matrix index to (deriv_type, physical_point).
49 Returns
50 -------
51 k_type : int64 array of shape (N_total,)
52 Derivative type for each K index (0 = function value, 1.. = derivatives).
53 k_phys : int64 array of shape (N_total,)
54 Physical training point index for each K index.
55 deriv_lookup : int64 array of shape (n_types, n_types)
56 phi_exp_3d derivative-dimension index for block (type_i, type_j).
57 sign_lookup : float64 array of shape (n_types,)
58 Sign multiplier for each derivative type.
59 """
60 n_dt = plan['n_deriv_types']
61 n_types = n_dt + 1
62 N_total = n_rows_func + plan['n_pts_with_derivs']
64 k_type = np.empty(N_total, dtype=np.int64)
65 k_phys = np.empty(N_total, dtype=np.int64)
67 # Function-value rows: type 0, phys = row index
68 k_type[:n_rows_func] = 0
69 k_phys[:n_rows_func] = np.arange(n_rows_func)
71 # Derivative rows
72 idx_flat = plan['idx_flat']
73 idx_offsets = plan['idx_offsets']
74 idx_sizes = plan['index_sizes']
75 row_offsets = plan.get('row_offsets_abs', plan['row_offsets'] + n_rows_func)
76 for j in range(n_dt):
77 ro = row_offsets[j]
78 sz = idx_sizes[j]
79 off = idx_offsets[j]
80 k_type[ro:ro + sz] = j + 1
81 k_phys[ro:ro + sz] = idx_flat[off:off + sz]
83 # Derivative lookup: which phi_exp_3d dimension for (type_i, type_j)
84 deriv_lookup = np.empty((n_types, n_types), dtype=np.int64)
85 deriv_lookup[0, 0] = 0
86 fd = plan['fd_flat_indices']
87 df = plan['df_flat_indices']
88 dd = plan['dd_flat_indices']
89 for j in range(n_dt):
90 deriv_lookup[0, j + 1] = fd[j]
91 deriv_lookup[j + 1, 0] = df[j]
92 for i in range(n_dt):
93 for j in range(n_dt):
94 deriv_lookup[i + 1, j + 1] = dd[i, j]
96 # Sign lookup
97 signs = plan['signs']
98 sign_lookup = np.empty(n_types, dtype=np.float64)
99 sign_lookup[0] = signs[0]
100 for j in range(n_dt):
101 sign_lookup[j + 1] = signs[j + 1]
103 return k_type, k_phys, deriv_lookup, sign_lookup
106@numba.jit(nopython=True, cache=True)
107def _extract_K_sub(phi_exp_3d, nb_type, nb_phys, deriv_lookup, sign_lookup,
108 sigma_n_sq, sigma_data_diag, m, K_sub):
109 """
110 Assemble K_sub directly from phi_exp_3d for a neighbourhood `nb`.
112 nb_type[a], nb_phys[a] give the derivative type and physical point
113 for the a-th row/column of K_sub.
114 """
115 for a in range(m):
116 ta = nb_type[a]
117 pa = nb_phys[a]
118 for b in range(m):
119 tb = nb_type[b]
120 pb = nb_phys[b]
121 d = deriv_lookup[ta, tb]
122 K_sub[a, b] = phi_exp_3d[d, pa, pb] * sign_lookup[tb]
123 # Add noise to diagonal
124 K_sub[a, a] += sigma_n_sq + sigma_data_diag[a]
127@numba.jit(nopython=True, cache=True)
128def _extract_dK_sub(dphi_exp_3d, nb_type, nb_phys, deriv_lookup, sign_lookup, m, dK_sub):
129 """
130 Assemble dK_sub/dtheta from dphi_exp_3d for a neighbourhood.
132 Same as _extract_K_sub but without noise diagonal (noise doesn't
133 depend on kernel hyperparameters).
134 """
135 for a in range(m):
136 ta = nb_type[a]
137 pa = nb_phys[a]
138 for b in range(m):
139 tb = nb_type[b]
140 pb = nb_phys[b]
141 d = deriv_lookup[ta, tb]
142 dK_sub[a, b] = dphi_exp_3d[d, pa, pb] * sign_lookup[tb]
145@numba.jit(nopython=True, cache=True)
146def _trace_term_all_blocks(dphi_exp_3d, deriv_lookup, sign_lookup,
147 block_nb_type, block_nb_phys, block_m,
148 block_V_flat, block_V_offsets, block_n_cols,
149 n_blocks):
150 """
151 Compute sum_blocks trace(V^T dK_sub V) in a single numba pass.
153 Avoids Python-level block loop and per-block np.empty/matmul overhead.
154 """
155 result = 0.0
156 for b in range(n_blocks):
157 m = block_m[b]
158 n_cols = block_n_cols[b]
159 off = block_V_offsets[b]
161 # For each column pair (i, j) of V, compute V[:,i]^T dK_sub V[:,j]
162 # trace = sum_i V[:,i]^T dK_sub V[:,i]
163 nb_type = block_nb_type[b]
164 nb_phys = block_nb_phys[b]
166 for col in range(n_cols):
167 # Compute V[:,col]^T @ dK_sub @ V[:,col]
168 # = sum_a sum_b V[a,col] * dK_sub[a,b] * V[b,col]
169 for a in range(m):
170 ta = nb_type[a]
171 pa = nb_phys[a]
172 va = block_V_flat[off + col * m + a]
173 for bb in range(m):
174 tb = nb_type[bb]
175 pb = nb_phys[bb]
176 d = deriv_lookup[ta, tb]
177 dk_ab = dphi_exp_3d[d, pa, pb] * sign_lookup[tb]
178 vb = block_V_flat[off + col * m + bb]
179 result += va * dk_ab * vb
180 return result
183@numba.jit(nopython=True, cache=True)
184def _build_K_inv_proj_from_blocks(K_inv_proj,
185 block_nb_type, block_nb_phys, block_m,
186 block_V_flat, block_V_offsets, block_n_cols,
187 deriv_lookup, sign_lookup, n_blocks):
188 """
189 Project K_inv = U U^T into phi-space directly from U's block structure.
191 For each block, computes V[a,:] @ V[bb,:] (dot over columns) and
192 accumulates into K_inv_proj[d, pa, pb] with sign correction.
194 K_inv_proj must be pre-zeroed. Cost: O(Σ_blocks m² · n_cols).
195 """
196 max_m = block_nb_type.shape[1]
197 for b in range(n_blocks):
198 m = block_m[b]
199 nc = block_n_cols[b]
200 off = block_V_offsets[b]
202 for a in range(m):
203 ta = block_nb_type[b, a]
204 pa = block_nb_phys[b, a]
205 for bb in range(m):
206 tb = block_nb_type[b, bb]
207 pb = block_nb_phys[b, bb]
209 # V[a,:] @ V[bb,:] — dot product over block columns
210 vv = 0.0
211 for col in range(nc):
212 vv += (block_V_flat[off + col * m + a]
213 * block_V_flat[off + col * m + bb])
215 d = deriv_lookup[ta, tb]
216 K_inv_proj[d, pa, pb] += sign_lookup[tb] * vv
219@numba.jit(nopython=True, cache=True)
220def _build_alpha_proj(alpha_proj, alpha_vecs, signed_vecs, deriv_lookup):
221 """
222 Project alpha*alpha^T into phi-space using per-type alpha vectors.
224 alpha_proj must be pre-zeroed. signed_vecs = sign_lookup[:, None] * alpha_vecs.
225 alpha_proj[d, i, j] += Σ_{ta,tb: deriv_lookup[ta,tb]==d} alpha_vecs[ta,i] * signed_vecs[tb,j]
226 """
227 n_types = alpha_vecs.shape[0]
228 n_func = alpha_vecs.shape[1]
229 for ta in range(n_types):
230 for tb in range(n_types):
231 d = deriv_lookup[ta, tb]
232 for i in range(n_func):
233 ai = alpha_vecs[ta, i]
234 if ai == 0.0:
235 continue
236 for j in range(n_func):
237 alpha_proj[d, i, j] += ai * signed_vecs[tb, j]
240def _alpha_quadratic_form(dphi_3d, alpha_v, k_type, k_phys,
241 deriv_lookup, sign_lookup, n_types, n_func):
242 """
243 Compute α^T dK α via BLAS, grouping by derivative type.
245 Builds per-type alpha vectors (n_types, n_func), then for each
246 type pair (ta, tb) computes sign[tb] * a_ta @ dphi_3d[d] @ a_tb
247 using BLAS gemv + ddot. Avoids forming the full alpha projection.
248 """
249 # Build per-type alpha vectors
250 alpha_vecs = np.zeros((n_types, n_func))
251 N = len(alpha_v)
252 for i in range(N):
253 alpha_vecs[k_type[i], k_phys[i]] = alpha_v[i]
255 result = 0.0
256 for ta in range(n_types):
257 a_vec = alpha_vecs[ta]
258 for tb in range(n_types):
259 b_vec = alpha_vecs[tb]
260 d = deriv_lookup[ta, tb]
261 s = sign_lookup[tb]
262 # a_vec @ dphi_3d[d] @ b_vec — uses BLAS dgemv + ddot
263 result += s * a_vec @ dphi_3d[d] @ b_vec
264 return result
267@numba.jit(nopython=True, cache=True)
268def _project_G_to_W_proj(W_proj, G_b, nb_type, nb_phys,
269 deriv_lookup, sign_lookup, m):
270 """
271 Project per-block sensitivity G_b (m × m) into W_proj (ndir × n_func × n_func).
273 W_proj[d, pa, pb] += sign_lookup[tb] * G_b[a, b]
275 where d = deriv_lookup[ta, tb], pa = nb_phys[a], pb = nb_phys[b].
276 """
277 n_func = W_proj.shape[1]
278 plane = n_func * W_proj.shape[2]
279 wptr = W_proj.ravel()
280 for a_i in range(m):
281 ta = nb_type[a_i]
282 pa = nb_phys[a_i]
283 for bb_i in range(m):
284 tb = nb_type[bb_i]
285 pb = nb_phys[bb_i]
286 d = deriv_lookup[ta, tb]
287 wptr[d * plane + pa * W_proj.shape[2] + pb] += sign_lookup[tb] * G_b[a_i, bb_i]
290class Optimizer:
291 """
292 Optimizer class to perform hyperparameter tuning for derivative-enhanced Gaussian Process models
293 by minimizing the negative log marginal likelihood (NLL).
295 Parameters
296 ----------
297 model : object
298 An instance of a model (e.g., ddegp) containing the necessary training data
299 and kernel configuration.
300 """
302 def __init__(self, model):
303 self.model = model
304 self._kernel_plan = None
305 self._deriv_buf = None
306 self._deriv_buf_shape = None
307 self._deriv_factors = None
308 self._deriv_factors_key = None
309 self._ndir = None
310 self._K_buf = None
311 self._dK_buf = None
312 self._kernel_buf_size = None
313 self._W_proj_buf = None
314 self._W_proj_shape = None
315 self._U_buf = None
316 self._K_inv_buf = None
317 self._P_ix = None
318 # Direct phi extraction maps (built lazily)
319 self._k_index_map = None
320 self._inv_P = None
321 self._sigma_data_diag_mmd = None
322 self._block_metadata = None
324 def _get_deriv_buf(self, phi, n_bases, order):
325 """Return a pre-allocated buffer for get_all_derivs, reusing if shape matches."""
326 if self._ndir is None:
327 from math import comb
328 self._ndir = comb(n_bases + order, order)
329 shape = (self._ndir, phi.shape[0], phi.shape[1])
330 if self._deriv_buf is None or self._deriv_buf_shape != shape:
331 self._deriv_buf = np.zeros(shape, dtype=np.float64)
332 self._deriv_buf_shape = shape
333 return self._deriv_buf
335 def _expand_derivs(self, phi, n_bases, deriv_order):
336 """Expand OTI derivatives, using fast struct path if available."""
337 if hasattr(phi, 'get_all_derivs_fast'):
338 buf = self._get_deriv_buf(phi, n_bases, deriv_order)
339 factors = self._get_deriv_factors(n_bases, deriv_order)
340 return phi.get_all_derivs_fast(factors, buf)
341 return phi.get_all_derivs(n_bases, deriv_order)
343 @staticmethod
344 def _enum_factors(max_basis, ordi):
345 """Enumerate derivative factors in struct memory order for a given order.
347 Yields the factorial factor prod(count_b!) for each multi-index of
348 the given order, enumerated in the same order as the OTI struct layout
349 (last-index-major: for last=1..max_basis, recurse prefix with max=last).
350 """
351 from math import factorial
352 from collections import Counter
353 if ordi == 1:
354 for _ in range(max_basis):
355 yield 1.0
356 return
357 for last in range(1, max_basis + 1):
358 if ordi == 2:
359 for i in range(1, last + 1):
360 counts = Counter((i, last))
361 f = 1
362 for c in counts.values():
363 f *= factorial(c)
364 yield float(f)
365 else:
366 for prefix_factor, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1):
367 counts = dict(prefix_counts)
368 counts[last] = counts.get(last, 0) + 1
369 f = 1
370 for c in counts.values():
371 f *= factorial(c)
372 yield float(f)
374 @staticmethod
375 def _enum_factors_with_counts(max_basis, ordi):
376 """Enumerate (factor, counts_dict) pairs in struct order."""
377 from math import factorial
378 from collections import Counter
379 if ordi == 1:
380 for i in range(1, max_basis + 1):
381 yield 1.0, {i: 1}
382 return
383 for last in range(1, max_basis + 1):
384 for _, prefix_counts in Optimizer._enum_factors_with_counts(last, ordi - 1):
385 counts = dict(prefix_counts)
386 counts[last] = counts.get(last, 0) + 1
387 f = 1
388 for c in counts.values():
389 f *= factorial(c)
390 yield float(f), counts
392 def _get_deriv_factors(self, n_bases, order):
393 """Return cached precomputed derivative factorial factors."""
394 from math import comb
395 key = (n_bases, order)
396 if self._deriv_factors is not None and self._deriv_factors_key == key:
397 return self._deriv_factors
398 factors = [1.0] # order 0: real part
399 for ordi in range(1, order + 1):
400 factors.extend(self._enum_factors(n_bases, ordi))
401 self._deriv_factors = np.array(factors, dtype=np.float64)
402 self._deriv_factors_key = key
403 self._ndir = comb(n_bases + order, order)
404 return self._deriv_factors
406 def _ensure_kernel_plan(self, n_bases):
407 """Lazily precompute kernel plan (once per n_bases)."""
408 if self._kernel_plan is not None and self._kernel_plan_n_bases == n_bases:
409 return
410 if not hasattr(utils, 'precompute_kernel_plan'):
411 self._kernel_plan = None
412 return
413 self._kernel_plan = utils.precompute_kernel_plan(
414 self.model.n_order, n_bases,
415 self.model.flattened_der_indices,
416 self.model.powers,
417 self.model.derivative_locations,
418 )
419 self._kernel_plan_n_bases = n_bases
420 # Reset kernel buffers when plan changes
421 self._K_buf = None
422 self._dK_buf = None
423 self._kernel_buf_size = None
425 def _ensure_kernel_bufs(self, n_rows_func):
426 """Pre-allocate reusable K and dK buffers (avoids repeated malloc)."""
427 if self._kernel_plan is None:
428 return
429 total = n_rows_func + self._kernel_plan['n_pts_with_derivs']
430 if self._kernel_buf_size != total:
431 self._K_buf = np.empty((total, total))
432 self._dK_buf = np.empty((total, total))
433 self._kernel_buf_size = total
434 # Cache absolute offsets in plan so rbf_kernel_fast doesn't recompute
435 if 'row_offsets_abs' not in self._kernel_plan:
436 self._kernel_plan['row_offsets_abs'] = self._kernel_plan['row_offsets'] + n_rows_func
437 self._kernel_plan['col_offsets_abs'] = self._kernel_plan['col_offsets'] + n_rows_func
439 def _ensure_phi_index_maps(self, n_rows_func):
440 """Lazily build the K-index-to-phi maps and inverse permutation."""
441 if self._k_index_map is not None:
442 return
443 plan = self._kernel_plan
444 k_type, k_phys, deriv_lookup, sign_lookup = _build_k_index_map(
445 plan, n_rows_func)
446 self._k_index_map = (k_type, k_phys, deriv_lookup, sign_lookup)
448 P_full = self.model.mmd_P_full
449 inv_P = np.empty_like(P_full)
450 inv_P[P_full] = np.arange(len(P_full))
451 self._inv_P = inv_P
453 # sigma_data diagonal in MMD order
454 sd = self.model.sigma_data
455 if sd.ndim == 2:
456 sd_diag_orig = np.diag(sd) ** 2 if np.any(sd) else np.zeros(len(P_full))
457 else:
458 sd_diag_orig = np.zeros(len(P_full))
459 self._sigma_data_diag_mmd = sd_diag_orig[P_full]
461 # Precompute flat index arrays for phi_exp_3d gather.
462 # At runtime, K_sub = phi_exp_3d.ravel()[flat_idx] * sign_mat + noise.
463 stride_d = n_rows_func * n_rows_func
464 stride_row = n_rows_func
466 if (self.model.use_supernodes
467 and self.model.sparse_supernodes_full is not None):
468 for sn in self.model.sparse_supernodes_full:
469 ch = sn.get('children_arr')
470 if ch is None:
471 ch = np.asarray(sn['children'])
472 orig_ch = P_full[ch]
473 ch_type = k_type[orig_ch]
474 ch_phys = k_phys[orig_ch]
475 m = len(ch)
477 d_mat = deriv_lookup[ch_type[:, None], ch_type[None, :]]
478 pa_mat = np.broadcast_to(ch_phys[:, None], (m, m))
479 pb_mat = np.broadcast_to(ch_phys[None, :], (m, m))
480 sn['phi_flat_idx'] = np.ascontiguousarray(
481 d_mat * stride_d + pa_mat * stride_row + pb_mat
482 )
483 sn['phi_sign_mat'] = np.ascontiguousarray(
484 sign_lookup[ch_type[None, :]] * np.ones((m, 1))
485 )
486 sn['phi_sd_diag'] = self._sigma_data_diag_mmd[ch]
488 # Same precomputation for non-supernode block path
489 if (not self.model.use_supernodes
490 and self.model.n_order > 0):
491 N_total = len(P_full)
492 block_size = self.model.n_bases + 1
493 S = self.model.sparse_S_full_arr
494 self._block_phi_maps = []
495 for start in range(0, N_total, block_size):
496 end = min(start + block_size, N_total)
497 nb = S[end - 1] if isinstance(S[end - 1], np.ndarray) else np.asarray(S[end - 1])
498 m = len(nb)
500 orig_nb = P_full[nb]
501 nb_type = k_type[orig_nb]
502 nb_phys = k_phys[orig_nb]
504 d_mat = deriv_lookup[nb_type[:, None], nb_type[None, :]]
505 pa_mat = np.broadcast_to(nb_phys[:, None], (m, m))
506 pb_mat = np.broadcast_to(nb_phys[None, :], (m, m))
508 self._block_phi_maps.append({
509 'nb': nb,
510 'start': start,
511 'flat_idx': np.ascontiguousarray(
512 d_mat * stride_d + pa_mat * stride_row + pb_mat
513 ),
514 'sign_mat': np.ascontiguousarray(
515 sign_lookup[nb_type[None, :]] * np.ones((m, 1))
516 ),
517 'sd_diag': self._sigma_data_diag_mmd[nb],
518 'positions': np.searchsorted(nb, np.arange(start, end)),
519 })
522 def negative_log_marginal_likelihood(self, x0):
523 """
524 Compute the negative log marginal likelihood (NLL) via sparse U.
526 NLL = 0.5 * ||U^T y||^2 - sum(log|diag(U)|) + 0.5 * N * log(2π)
528 Parameters
529 ----------
530 x0 : ndarray
531 Vector of log-scaled hyperparameters (length scales and noise).
533 Returns
534 -------
535 float
536 Value of the negative log marginal likelihood.
537 """
538 try:
539 if self.model._use_dense_factor:
540 # Dense path: single Cholesky, no sparse U
541 W, alpha, nll, *_ = self._dense_nll_and_W(x0)
542 if nll > 1e6:
543 return 1e6
544 return nll
546 # Use direct phi path (skip full K construction) when possible
547 if self.model.n_order > 0:
548 alpha, U, nlml, *_ = self._sparse_nlml_direct(x0)
549 else:
550 K, _, _, _, _ = self._build_K_and_phi(x0)
551 alpha, U, nlml = self._sparse_U_alpha_nll(K)
553 # Sparse U can silently produce bad factors when K is
554 # ill-conditioned (e.g. very small noise). Clamp to 1e6
555 # to match the dense fallback behaviour.
556 if nlml > 1e6:
557 return 1e6
559 self.model._cached_U = U
560 self.model._cached_P = self.model.mmd_P_full
561 self.model._cached_alpha = alpha
562 self.model._cached_L = None
563 self.model._cached_low = None
564 self.model._cached_params = x0.copy()
566 return nlml
567 except Exception:
568 return 1e6
570 def nll_wrapper(self, x0):
571 """
572 Wrapper function to compute NLL for optimizer.
574 Parameters
575 ----------
576 x0 : ndarray
577 Hyperparameter vector.
579 Returns
580 -------
581 float
582 NLL evaluated at x0.
583 """
584 return self.negative_log_marginal_likelihood(x0)
586 @profile
587 def _compute_grad(self, x0, W, phi, n_bases, oti, diffs):
588 """
589 Compute the NLL gradient given pre-factorised W = K^{-1} - α α^T.
591 Factoring this out allows nll_grad and nll_and_grad to share the
592 expensive Cholesky decomposition instead of each rebuilding it.
593 """
594 ln10 = np.log(10.0)
595 kernel = self.model.kernel
596 kernel_type = self.model.kernel_type
597 D = len(diffs)
598 sigma_n_sq = (10.0 ** x0[-1]) ** 2
600 grad = np.zeros(len(x0))
601 use_fast = self._kernel_plan is not None
602 base_shape = (W.shape[0] - self._kernel_plan['n_pts_with_derivs'],) * 2 if use_fast else None
604 deriv_order = 2 * self.model.n_order
606 # Precompute W projected into phi_exp space to avoid assembling
607 # the full dK matrix for each hyperparameter dimension.
608 W_proj = None
609 if use_fast and self.model.n_order > 0:
610 ndir = self._ndir
611 proj_shape = (ndir, base_shape[0], base_shape[1])
612 if self._W_proj_buf is None or self._W_proj_shape != proj_shape:
613 self._W_proj_buf = np.empty(proj_shape)
614 self._W_proj_shape = proj_shape
615 W_proj = self._W_proj_buf
617 plan = self._kernel_plan
618 row_off = plan.get('row_offsets_abs', plan['row_offsets'] + base_shape[0])
619 col_off = plan.get('col_offsets_abs', plan['col_offsets'] + base_shape[1])
621 utils._project_W_to_phi_space(
622 W, W_proj, base_shape[0], base_shape[1],
623 plan['fd_flat_indices'], plan['df_flat_indices'],
624 plan['dd_flat_indices'],
625 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'],
626 plan['signs'], plan['n_deriv_types'], row_off, col_off,
627 )
629 _use_vdot_fused = W_proj is not None and hasattr(phi, 'vdot_expand_fast')
630 FW_T = None
631 if _use_vdot_fused:
632 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order)
633 _vdot_arr = np.asarray(_vdot_factors)
634 ndir_d = len(_vdot_arr)
635 FW_T = np.empty((base_shape[0] * base_shape[1], ndir_d))
636 np.multiply(W_proj.reshape(ndir_d, -1).T, _vdot_arr, out=FW_T)
638 @profile
639 def _gc(dphi):
640 if _use_vdot_fused:
641 return 0.5 * dphi.vdot_expand_fast(_vdot_factors, W_proj)
642 if self.model.n_order == 0:
643 dphi_exp = dphi.real[np.newaxis, :, :]
644 else:
645 dphi_exp = self._expand_derivs(dphi, n_bases, deriv_order)
646 if W_proj is not None:
647 dphi_3d = dphi_exp.reshape(W_proj.shape)
648 return 0.5 * np.vdot(W_proj, dphi_3d)
649 elif use_fast:
650 dphi_3d = dphi_exp.reshape(dphi_exp.shape[0], base_shape[0], base_shape[1])
651 dK = utils.rbf_kernel_fast(dphi_3d, self._kernel_plan, out=self._dK_buf)
652 return 0.5 * np.vdot(W, dK)
653 else:
654 dK = utils.rbf_kernel(
655 dphi, dphi_exp,
656 self.model.n_order, n_bases,
657 self.model.flattened_der_indices, self.model.powers,
658 index=self.model.derivative_locations,
659 )
660 return 0.5 * np.vdot(W, dK)
662 # ── signal variance (common: d phi/d log_sf = 2*ln10 * phi) ──────
663 if _use_vdot_fused:
664 grad[-2] = ln10 * phi.vdot_expand_fast(_vdot_factors, W_proj)
665 else:
666 grad[-2] = _gc(oti.mul(2.0 * ln10, phi))
668 # ── noise variance (common: dK/d log_sn = diag(2*ln10*σ_n²)) ────
669 grad[-1] = ln10 * sigma_n_sq * np.trace(W)
671 # ── kernel-specific hyperparameter gradients ──────────────────────
673 if kernel == 'SE':
674 # phi = sf² * exp(-0.5 * Σ_d ell_d² * diff_d²)
675 # d phi/d log_ell_d = -ln10 * ell_d² * diff_d² * phi
676 if kernel_type == 'anisotropic':
677 ell = 10.0 ** x0[:D]
678 if _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'):
679 scales = np.array([-ln10 * ell[d] ** 2 for d in range(D)])
680 grad_buf = np.zeros(D)
681 phi.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T)
682 grad[:D] = grad_buf
683 elif hasattr(phi, 'fused_scale_sq_mul_sparse'):
684 dphi_buf = oti.zeros(phi.shape)
685 for d in range(D):
686 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], phi, -ln10 * ell[d] ** 2, d)
687 grad[d] = _gc(dphi_buf)
688 elif hasattr(phi, 'fused_scale_sq_mul'):
689 dphi_buf = oti.zeros(phi.shape)
690 for d in range(D):
691 dphi_buf.fused_scale_sq_mul(diffs[d], phi, -ln10 * ell[d] ** 2)
692 grad[d] = _gc(dphi_buf)
693 else:
694 for d in range(D):
695 d_sq = oti.mul(diffs[d], diffs[d])
696 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi))
697 grad[d] = _gc(dphi_d)
698 else: # isotropic: single ell
699 ell = 10.0 ** float(x0[0])
700 if hasattr(phi, 'fused_sum_sq_sparse'):
701 sum_sq = oti.zeros(phi.shape)
702 sum_sq.fused_sum_sq_sparse(diffs)
703 elif hasattr(phi, 'fused_sum_sq'):
704 sum_sq = oti.zeros(phi.shape)
705 sum_sq.fused_sum_sq(diffs)
706 else:
707 sum_sq = oti.mul(diffs[0], diffs[0])
708 for d in range(1, D):
709 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
710 grad[0] = _gc(oti.mul(-ln10 * ell ** 2, oti.mul(sum_sq, phi)))
712 elif kernel == 'RQ':
713 # phi = sf² * (1 + r²/(2α))^(-α), r² = Σ_d (ell_d * diff_d)²
714 # d phi/d log_ell_d = -ln10 * ell_d² * diff_d² * phi / base
715 # d phi/d log_α = ln10 * α * phi * [-log(base) + (1 - 1/base)]
716 if kernel_type == 'anisotropic':
717 ell = 10.0 ** x0[:D]
718 alpha_rq = 10.0 ** float(x0[D])
719 alpha_idx = D
720 else:
721 ell_val = 10.0 ** float(x0[0])
722 ell = np.full(D, ell_val)
723 alpha_rq = np.exp(float(x0[1])) # iso uses exp(x), not 10^x
724 alpha_idx = 1
726 # Recompute r² and base in OTI
727 if hasattr(phi, 'fused_sqdist_sparse'):
728 r2 = oti.zeros(phi.shape)
729 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
730 r2.fused_sqdist_sparse(diffs, ell_sq)
731 elif hasattr(phi, 'fused_sqdist'):
732 r2 = oti.zeros(phi.shape)
733 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
734 r2.fused_sqdist(diffs, ell_sq)
735 else:
736 r2 = oti.mul(ell[0], diffs[0])
737 r2 = oti.mul(r2, r2)
738 for d in range(1, D):
739 td = oti.mul(ell[d], diffs[d])
740 r2 = oti.sum(r2, oti.mul(td, td))
741 base = oti.sum(1.0, oti.mul(r2, 1.0 / (2.0 * alpha_rq)))
742 inv_base = oti.pow(base, -1)
743 phi_over_base = oti.mul(phi, inv_base)
745 if kernel_type == 'anisotropic':
746 if _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'):
747 scales = np.array([-ln10 * ell[d] ** 2 for d in range(D)])
748 grad_buf = np.zeros(D)
749 phi_over_base.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T)
750 grad[:D] = grad_buf
751 elif hasattr(phi, 'fused_scale_sq_mul_sparse'):
752 dphi_buf = oti.zeros(phi.shape)
753 for d in range(D):
754 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], phi_over_base, -ln10 * ell[d] ** 2, d)
755 grad[d] = _gc(dphi_buf)
756 elif hasattr(phi, 'fused_scale_sq_mul'):
757 dphi_buf = oti.zeros(phi.shape)
758 for d in range(D):
759 dphi_buf.fused_scale_sq_mul(diffs[d], phi_over_base, -ln10 * ell[d] ** 2)
760 grad[d] = _gc(dphi_buf)
761 else:
762 for d in range(D):
763 d_sq = oti.mul(diffs[d], diffs[d])
764 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi_over_base))
765 grad[d] = _gc(dphi_d)
766 else:
767 if hasattr(phi, 'fused_sum_sq_sparse'):
768 sum_sq = oti.zeros(phi.shape)
769 sum_sq.fused_sum_sq_sparse(diffs)
770 elif hasattr(phi, 'fused_sum_sq'):
771 sum_sq = oti.zeros(phi.shape)
772 sum_sq.fused_sum_sq(diffs)
773 else:
774 sum_sq = oti.mul(diffs[0], diffs[0])
775 for d in range(1, D):
776 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
777 grad[0] = _gc(oti.mul(-ln10 * ell[0] ** 2, oti.mul(sum_sq, phi_over_base)))
779 # alpha gradient: phi * α_factor * [-log(base) + (1 - 1/base)]
780 # aniso: alpha = 10^x → d alpha/dx = ln10 * alpha
781 # iso: alpha = exp(x) → d alpha/dx = alpha
782 log_base = oti.log(base)
783 term = oti.sub(oti.sub(1.0, inv_base), log_base)
784 alpha_factor = ln10 * alpha_rq if kernel_type == 'anisotropic' else alpha_rq
785 grad[alpha_idx] = _gc(oti.mul(alpha_factor, oti.mul(phi, term)))
787 elif kernel == 'SineExp':
788 # phi = sf² * exp(-2 * Σ_d (ell_d * sin(π/p_d * diff_d))²)
789 # d phi/d log_ell_d = -4*ln10 * ell_d² * sin_d² * phi
790 # d phi/d log_p_d = 4*ln10 * ell_d² * (π/p_d) * sin_d * cos_d * diff_d * phi
791 if kernel_type == 'anisotropic':
792 ell = 10.0 ** x0[:D]
793 p = 10.0 ** x0[D:2 * D]
794 pip = np.pi / p # π/p_d per dimension
795 p_start = D # index of first log_p in x0
796 else:
797 ell_val = 10.0 ** float(x0[0])
798 p_val = 10.0 ** float(x0[1])
799 pip_val = np.pi / p_val
800 ell = np.full(D, ell_val)
801 pip = np.full(D, pip_val)
802 p_start = 1
804 # Precompute sin and cos for each dimension
805 sin_d = []
806 cos_d = []
807 for d in range(D):
808 arg = oti.mul(pip[d], diffs[d])
809 sin_d.append(oti.sin(arg))
810 cos_d.append(oti.cos(arg))
812 # Length-scale gradients: d phi/d log_ell_d = -4*ln10 * ell_d² * sin_d² * phi
813 if kernel_type == 'anisotropic':
814 if hasattr(phi, 'fused_scale_sq_mul'):
815 dphi_buf = oti.zeros(phi.shape)
816 for d in range(D):
817 dphi_buf.fused_scale_sq_mul(sin_d[d], phi, -4.0 * ln10 * ell[d] ** 2)
818 grad[d] = _gc(dphi_buf)
819 else:
820 for d in range(D):
821 sin_sq = oti.mul(sin_d[d], sin_d[d])
822 grad[d] = _gc(oti.mul(-4.0 * ln10 * ell[d] ** 2,
823 oti.mul(sin_sq, phi)))
824 else:
825 if hasattr(phi, 'fused_sum_sq'):
826 sum_sin_sq = oti.zeros(phi.shape)
827 sum_sin_sq.fused_sum_sq(sin_d)
828 else:
829 sum_sin_sq = oti.mul(sin_d[0], sin_d[0])
830 for d in range(1, D):
831 sum_sin_sq = oti.sum(sum_sin_sq, oti.mul(sin_d[d], sin_d[d]))
832 grad[0] = _gc(oti.mul(-4.0 * ln10 * ell[0] ** 2,
833 oti.mul(sum_sin_sq, phi)))
835 # Period gradients
836 if kernel_type == 'anisotropic':
837 for d in range(D):
838 # d phi/d log_p_d = 4*ln10*ell_d²*(π/p_d)*sin_d*cos_d*diff_d * phi
839 sc_diff = oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d]))
840 scale = 4.0 * ln10 * ell[d] ** 2 * pip[d]
841 grad[p_start + d] = _gc(oti.mul(scale, oti.mul(sc_diff, phi)))
842 else:
843 # d phi/d log_p = 4*ln10*ell²*(π/p) * Σ_d(sin_d*cos_d*diff_d) * phi
844 sum_scd = oti.mul(sin_d[0], oti.mul(cos_d[0], diffs[0]))
845 for d in range(1, D):
846 sum_scd = oti.sum(sum_scd,
847 oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d])))
848 scale = 4.0 * ln10 * ell[0] ** 2 * pip[0]
849 grad[p_start] = _gc(oti.mul(scale, oti.mul(sum_scd, phi)))
851 elif kernel == 'Matern':
852 # phi = sf² * f(r), r = sqrt(Σ_d (ell_d*(diff_d+ε))²)
853 # d phi/d log_ell_d = sf² * f'(r) * ln10 * ell_d² * (diff_d+ε)² / r
854 kf = self.model.kernel_factory
856 # Build/cache the Matern derivative function
857 if not hasattr(kf, '_matern_grad_prebuild'):
858 kf._matern_grad_prebuild = matern_kernel_grad_builder(
859 kf.nu, oti_module=oti)
861 if kernel_type == 'anisotropic':
862 ell = 10.0 ** x0[:D]
863 else:
864 ell = np.full(D, 10.0 ** float(x0[0]))
866 sigma_f_sq = (10.0 ** float(x0[-2])) ** 2
867 _eps = 1e-10 # regularise r, not each diff (matches kernel_funcs.py)
869 # Recompute r in OTI (matches matern_kernel_anisotropic/isotropic)
870 if hasattr(phi, 'fused_sqdist_sparse'):
871 r2 = oti.zeros(phi.shape)
872 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
873 r2.fused_sqdist_sparse(diffs, ell_sq)
874 elif hasattr(phi, 'fused_sqdist'):
875 r2 = oti.zeros(phi.shape)
876 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
877 r2.fused_sqdist(diffs, ell_sq)
878 else:
879 r2 = oti.mul(ell[0], diffs[0])
880 r2 = oti.mul(r2, r2)
881 for d in range(1, D):
882 td = oti.mul(ell[d], diffs[d])
883 r2 = oti.sum(r2, oti.mul(td, td))
884 r_oti = oti.sqrt(oti.sum(r2, _eps ** 2))
885 f_prime_r = kf._matern_grad_prebuild(r_oti) # df/dr (OTI)
886 inv_r = oti.pow(r_oti, -1)
888 # Precompute base = sigma_f² * f'(r) * 1/r for length-scale gradients
889 # grad[d] = _gc(base * ln10 * ell_d² * diff_d²)
890 base_matern = oti.mul(sigma_f_sq, oti.mul(f_prime_r, inv_r))
891 if kernel_type == 'anisotropic':
892 if _use_vdot_fused and hasattr(phi, 'fused_grad_all_dims'):
893 scales = np.array([ln10 * ell[d] ** 2 for d in range(D)])
894 grad_buf = np.zeros(D)
895 base_matern.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T)
896 grad[:D] = grad_buf
897 elif hasattr(phi, 'fused_scale_sq_mul_sparse'):
898 dphi_buf = oti.zeros(phi.shape)
899 for d in range(D):
900 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], base_matern, ln10 * ell[d] ** 2, d)
901 grad[d] = _gc(dphi_buf)
902 elif hasattr(phi, 'fused_scale_sq_mul'):
903 dphi_buf = oti.zeros(phi.shape)
904 for d in range(D):
905 dphi_buf.fused_scale_sq_mul(diffs[d], base_matern, ln10 * ell[d] ** 2)
906 grad[d] = _gc(dphi_buf)
907 else:
908 for d in range(D):
909 d_sq = oti.mul(diffs[d], diffs[d])
910 dphi_d = oti.mul(ln10 * ell[d] ** 2, oti.mul(d_sq, base_matern))
911 grad[d] = _gc(dphi_d)
912 else:
913 ell_val = ell[0]
914 if hasattr(phi, 'fused_sum_sq_sparse'):
915 sum_dsq = oti.zeros(phi.shape)
916 sum_dsq.fused_sum_sq_sparse(diffs)
917 elif hasattr(phi, 'fused_sum_sq'):
918 sum_dsq = oti.zeros(phi.shape)
919 sum_dsq.fused_sum_sq(diffs)
920 else:
921 sum_dsq = oti.mul(diffs[0], diffs[0])
922 for d in range(1, D):
923 sum_dsq = oti.sum(sum_dsq, oti.mul(diffs[d], diffs[d]))
924 dphi_e = oti.mul(ln10 * ell_val ** 2, oti.mul(sum_dsq, base_matern))
925 grad[0] = _gc(dphi_e)
927 elif kernel == 'SI':
928 # phi = sf² * Π_d (1 + ell_d * B(diff_d))
929 # d phi/d log_ell_d = ln10 * ell_d * B(diff_d) / (1 + ell_d*B(diff_d)) * phi
930 kf = self.model.kernel_factory
931 si_prebuild = kf.SI_kernel_prebuild
933 if kernel_type == 'anisotropic':
934 ell = 10.0 ** x0[:D]
935 else:
936 ell = np.full(D, 10.0 ** float(x0[0]))
938 # Precompute SI values and factor terms for each dimension
939 si_vals = [si_prebuild(diffs[d]) for d in range(D)]
940 term_vals = [oti.sum(1.0, oti.mul(ell[d], si_vals[d])) for d in range(D)]
942 if kernel_type == 'anisotropic':
943 for d in range(D):
944 phi_over_term = oti.div(phi, term_vals[d])
945 dphi_d = oti.mul(ln10 * ell[d],
946 oti.mul(si_vals[d], phi_over_term))
947 grad[d] = _gc(dphi_d)
948 else:
949 ell_val = ell[0]
950 # d phi/d log_ell = ln10 * ell * Σ_d [B(diff_d)/(1+ell*B(diff_d))] * phi
951 acc = oti.mul(si_vals[0], oti.div(phi, term_vals[0]))
952 for d in range(1, D):
953 acc = oti.sum(acc, oti.mul(si_vals[d],
954 oti.div(phi, term_vals[d])))
955 grad[0] = _gc(oti.mul(ln10 * ell_val, acc))
957 return grad
959 @profile
960 def _compute_grad_blockwise(self, x0, U, alpha_v, phi, n_bases, oti, diffs):
961 """
962 Compute the NLL gradient by differentiating through each block's
963 Cholesky in the Vecchia decomposition.
965 For each block b the Vecchia NLL contribution is:
966 NLL_b = 0.5 * Σ_j [ α_b[p_j]² / s_j - log(s_j) ]
968 where α_b = K_sub⁻¹ y_nb, s_j = (K_sub⁻¹)[p_j, p_j].
970 Differentiating through K_sub⁻¹ gives the per-block sensitivity:
971 G_b = 0.5 * M V^T
973 where V = K_sub⁻¹ E (un-normalised U block), and
974 M[:, j] = γ_j V[:, j] - 2 β_j α_b
975 β_j = α_b[p_j] / s_j, γ_j = β_j² + 1/s_j
977 The gradient is dNLL/dθ = Σ_b tr(G_b dK_sub_b/dθ),
978 which is projected into phi-space as W_proj for vdot.
979 """
980 ln10 = np.log(10.0)
981 kernel = self.model.kernel
982 kernel_type = self.model.kernel_type
983 D = len(diffs)
984 sigma_n_sq = (10.0 ** x0[-1]) ** 2
986 grad = np.zeros(len(x0))
987 deriv_order = 2 * self.model.n_order
988 plan = self._kernel_plan
989 P_full = self.model.mmd_P_full
990 N_total = len(P_full)
991 n_func = phi.shape[0]
993 ndir = self._ndir
994 k_type, k_phys, deriv_lookup, sign_lookup = self._k_index_map
996 # ── phi_exp for K_sub reconstruction ──────────────────────
997 phi_exp = self._expand_derivs(phi, n_bases, deriv_order)
998 phi_3d = phi_exp.reshape(phi_exp.shape[0], n_func, n_func)
999 phi_flat = phi_3d.ravel()
1001 block_maps = self._block_phi_maps
1002 y_ord = self.model.y_train[P_full]
1004 # ── accumulate W_proj and noise trace from per-block G_b ──
1005 proj_shape = (ndir, n_func, n_func)
1006 W_proj = np.zeros(proj_shape)
1007 w_flat = W_proj.ravel()
1008 noise_trace = 0.0
1010 for bm in block_maps:
1011 nb = bm['nb']
1012 m = len(nb)
1013 positions = bm['positions']
1014 n_cols = len(positions)
1016 # Reconstruct K_sub for this block
1017 K_sub = phi_flat[bm['flat_idx']].reshape(m, m) * bm['sign_mat']
1018 diag_idx = np.arange(m)
1019 K_sub[diag_idx, diag_idx] += sigma_n_sq + bm['sd_diag']
1021 # Cholesky factor
1022 L_u, low_u = cho_factor(K_sub, lower=True)
1024 # V = K_sub⁻¹ E (un-normalised U block, m × n_cols)
1025 E = np.zeros((m, n_cols))
1026 E[positions, np.arange(n_cols)] = 1.0
1027 V = cho_solve((L_u, low_u), E)
1029 # α_b = K_sub⁻¹ y_nb (m,)
1030 y_nb = y_ord[nb]
1031 alpha_b = cho_solve((L_u, low_u), y_nb)
1033 # Per-parent scalars
1034 s = V[positions, np.arange(n_cols)] # s_j = S[p_j, p_j]
1035 a = alpha_b[positions] # α_b[p_j]
1036 beta = a / s # β_j
1037 gamma = beta ** 2 + 1.0 / s # γ_j
1039 # M[:, j] = γ_j V[:, j] - 2 β_j α_b (m × n_cols)
1040 M = V * gamma[np.newaxis, :] - 2.0 * alpha_b[:, np.newaxis] * beta[np.newaxis, :]
1042 # Noise trace: tr(G_b) = 0.5 Σ_{a,k} M[a,k] V[a,k]
1043 noise_trace += np.sum(M * V)
1045 # Project G_b into W_proj using precomputed flat indices
1046 G_b = M @ V.T # m × m
1047 np.add.at(w_flat, bm['flat_idx'].ravel(), (G_b * bm['sign_mat']).ravel())
1049 # ── noise gradient ───────────────────────────────────────
1050 # dK_sub/d(sigma_n_sq) = I ⇒ dNLL/d(sigma_n_sq) = Σ_b 0.5 tr(M V^T)
1051 # x[-1] = log10(sigma_n), sigma_n_sq = 10^(2x[-1])
1052 # chain rule: d(sigma_n_sq)/d(x[-1]) = 2 ln10 sigma_n_sq
1053 grad[-1] = ln10 * sigma_n_sq * noise_trace
1055 # ── deriv factors + fast vdot path ───────────────────────
1056 _vdot_factors = self._get_deriv_factors(n_bases, deriv_order)
1058 # Precompute transposed factor-weighted W_proj for cache-friendly access
1059 # FW_T[kk, c] = factors[c] * W_proj[c, kk] — shape (size, ndir)
1060 _vdot_arr = np.asarray(_vdot_factors)
1061 FW_T = np.empty((n_func * n_func, ndir))
1062 np.multiply(W_proj.reshape(ndir, -1).T, _vdot_arr, out=FW_T)
1064 @profile
1065 def _gc_block(dphi):
1066 return 0.5 * dphi.vdot_expand_fast(_vdot_factors, W_proj)
1068 # ── signal variance ──────────────────────────────────────────
1069 grad[-2] = ln10 * phi.vdot_expand_fast(_vdot_factors, W_proj)
1071 # ── kernel-specific hyperparameter gradients ─────────────────
1073 if kernel == 'SE':
1074 if kernel_type == 'anisotropic':
1075 ell = 10.0 ** x0[:D]
1076 if hasattr(phi, 'fused_grad_all_dims'):
1077 scales = np.array([-ln10 * ell[d] ** 2 for d in range(D)])
1078 grad_buf = np.zeros(D)
1079 phi.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T)
1080 grad[:D] = grad_buf
1081 elif hasattr(phi, 'fused_scale_sq_mul_sparse'):
1082 dphi_buf = oti.zeros(phi.shape)
1083 for d in range(D):
1084 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], phi, -ln10 * ell[d] ** 2, d)
1085 grad[d] = _gc_block(dphi_buf)
1086 elif hasattr(phi, 'fused_scale_sq_mul'):
1087 dphi_buf = oti.zeros(phi.shape)
1088 for d in range(D):
1089 dphi_buf.fused_scale_sq_mul(diffs[d], phi, -ln10 * ell[d] ** 2)
1090 grad[d] = _gc_block(dphi_buf)
1091 else:
1092 for d in range(D):
1093 d_sq = oti.mul(diffs[d], diffs[d])
1094 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi))
1095 grad[d] = _gc_block(dphi_d)
1096 else:
1097 ell = 10.0 ** float(x0[0])
1098 if hasattr(phi, 'fused_sum_sq_sparse'):
1099 sum_sq = oti.zeros(phi.shape)
1100 sum_sq.fused_sum_sq_sparse(diffs)
1101 elif hasattr(phi, 'fused_sum_sq'):
1102 sum_sq = oti.zeros(phi.shape)
1103 sum_sq.fused_sum_sq(diffs)
1104 else:
1105 sum_sq = oti.mul(diffs[0], diffs[0])
1106 for d in range(1, D):
1107 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
1108 grad[0] = _gc_block(oti.mul(-ln10 * ell ** 2, oti.mul(sum_sq, phi)))
1110 elif kernel == 'RQ':
1111 if kernel_type == 'anisotropic':
1112 ell = 10.0 ** x0[:D]
1113 alpha_rq = 10.0 ** float(x0[D])
1114 alpha_idx = D
1115 else:
1116 ell_val = 10.0 ** float(x0[0])
1117 ell = np.full(D, ell_val)
1118 alpha_rq = np.exp(float(x0[1]))
1119 alpha_idx = 1
1121 if hasattr(phi, 'fused_sqdist_sparse'):
1122 r2 = oti.zeros(phi.shape)
1123 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
1124 r2.fused_sqdist_sparse(diffs, ell_sq)
1125 elif hasattr(phi, 'fused_sqdist'):
1126 r2 = oti.zeros(phi.shape)
1127 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
1128 r2.fused_sqdist(diffs, ell_sq)
1129 else:
1130 r2 = oti.mul(ell[0], diffs[0])
1131 r2 = oti.mul(r2, r2)
1132 for d in range(1, D):
1133 td = oti.mul(ell[d], diffs[d])
1134 r2 = oti.sum(r2, oti.mul(td, td))
1135 base = oti.sum(1.0, oti.mul(r2, 1.0 / (2.0 * alpha_rq)))
1136 inv_base = oti.pow(base, -1)
1137 phi_over_base = oti.mul(phi, inv_base)
1139 if kernel_type == 'anisotropic':
1140 if hasattr(phi, 'fused_grad_all_dims'):
1141 scales = np.array([-ln10 * ell[d] ** 2 for d in range(D)])
1142 grad_buf = np.zeros(D)
1143 phi_over_base.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T)
1144 grad[:D] = grad_buf
1145 elif hasattr(phi, 'fused_scale_sq_mul_sparse'):
1146 dphi_buf = oti.zeros(phi.shape)
1147 for d in range(D):
1148 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], phi_over_base, -ln10 * ell[d] ** 2, d)
1149 grad[d] = _gc_block(dphi_buf)
1150 elif hasattr(phi, 'fused_scale_sq_mul'):
1151 dphi_buf = oti.zeros(phi.shape)
1152 for d in range(D):
1153 dphi_buf.fused_scale_sq_mul(diffs[d], phi_over_base, -ln10 * ell[d] ** 2)
1154 grad[d] = _gc_block(dphi_buf)
1155 else:
1156 for d in range(D):
1157 d_sq = oti.mul(diffs[d], diffs[d])
1158 dphi_d = oti.mul(-ln10 * ell[d] ** 2, oti.mul(d_sq, phi_over_base))
1159 grad[d] = _gc_block(dphi_d)
1160 else:
1161 if hasattr(phi, 'fused_sum_sq_sparse'):
1162 sum_sq = oti.zeros(phi.shape)
1163 sum_sq.fused_sum_sq_sparse(diffs)
1164 elif hasattr(phi, 'fused_sum_sq'):
1165 sum_sq = oti.zeros(phi.shape)
1166 sum_sq.fused_sum_sq(diffs)
1167 else:
1168 sum_sq = oti.mul(diffs[0], diffs[0])
1169 for d in range(1, D):
1170 sum_sq = oti.sum(sum_sq, oti.mul(diffs[d], diffs[d]))
1171 grad[0] = _gc_block(oti.mul(-ln10 * ell[0] ** 2, oti.mul(sum_sq, phi_over_base)))
1173 log_base = oti.log(base)
1174 term = oti.sub(oti.sub(1.0, inv_base), log_base)
1175 alpha_factor = ln10 * alpha_rq if kernel_type == 'anisotropic' else alpha_rq
1176 grad[alpha_idx] = _gc_block(oti.mul(alpha_factor, oti.mul(phi, term)))
1178 elif kernel == 'SineExp':
1179 if kernel_type == 'anisotropic':
1180 ell = 10.0 ** x0[:D]
1181 p = 10.0 ** x0[D:2 * D]
1182 pip = np.pi / p
1183 p_start = D
1184 else:
1185 ell_val = 10.0 ** float(x0[0])
1186 p_val = 10.0 ** float(x0[1])
1187 pip_val = np.pi / p_val
1188 ell = np.full(D, ell_val)
1189 pip = np.full(D, pip_val)
1190 p_start = 1
1192 sin_d = []
1193 cos_d = []
1194 for d in range(D):
1195 arg = oti.mul(pip[d], diffs[d])
1196 sin_d.append(oti.sin(arg))
1197 cos_d.append(oti.cos(arg))
1199 if kernel_type == 'anisotropic':
1200 if hasattr(phi, 'fused_scale_sq_mul'):
1201 dphi_buf = oti.zeros(phi.shape)
1202 for d in range(D):
1203 dphi_buf.fused_scale_sq_mul(sin_d[d], phi, -4.0 * ln10 * ell[d] ** 2)
1204 grad[d] = _gc_block(dphi_buf)
1205 else:
1206 for d in range(D):
1207 sin_sq = oti.mul(sin_d[d], sin_d[d])
1208 grad[d] = _gc_block(oti.mul(-4.0 * ln10 * ell[d] ** 2,
1209 oti.mul(sin_sq, phi)))
1210 else:
1211 if hasattr(phi, 'fused_sum_sq'):
1212 sum_sin_sq = oti.zeros(phi.shape)
1213 sum_sin_sq.fused_sum_sq(sin_d)
1214 else:
1215 sum_sin_sq = oti.mul(sin_d[0], sin_d[0])
1216 for d in range(1, D):
1217 sum_sin_sq = oti.sum(sum_sin_sq, oti.mul(sin_d[d], sin_d[d]))
1218 grad[0] = _gc_block(oti.mul(-4.0 * ln10 * ell[0] ** 2,
1219 oti.mul(sum_sin_sq, phi)))
1221 if kernel_type == 'anisotropic':
1222 for d in range(D):
1223 sc_diff = oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d]))
1224 scale = 4.0 * ln10 * ell[d] ** 2 * pip[d]
1225 grad[p_start + d] = _gc_block(oti.mul(scale, oti.mul(sc_diff, phi)))
1226 else:
1227 sum_scd = oti.mul(sin_d[0], oti.mul(cos_d[0], diffs[0]))
1228 for d in range(1, D):
1229 sum_scd = oti.sum(sum_scd,
1230 oti.mul(sin_d[d], oti.mul(cos_d[d], diffs[d])))
1231 scale = 4.0 * ln10 * ell[0] ** 2 * pip[0]
1232 grad[p_start] = _gc_block(oti.mul(scale, oti.mul(sum_scd, phi)))
1234 elif kernel == 'Matern':
1235 kf = self.model.kernel_factory
1236 if not hasattr(kf, '_matern_grad_prebuild'):
1237 from jetgp.kernel_funcs.kernel_funcs import matern_kernel_grad_builder
1238 kf._matern_grad_prebuild = matern_kernel_grad_builder(
1239 kf.nu, oti_module=oti)
1241 if kernel_type == 'anisotropic':
1242 ell = 10.0 ** x0[:D]
1243 else:
1244 ell = np.full(D, 10.0 ** float(x0[0]))
1246 sigma_f_sq = (10.0 ** float(x0[-2])) ** 2
1247 _eps = 1e-10
1249 if hasattr(phi, 'fused_sqdist_sparse'):
1250 r2 = oti.zeros(phi.shape)
1251 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
1252 r2.fused_sqdist_sparse(diffs, ell_sq)
1253 elif hasattr(phi, 'fused_sqdist'):
1254 r2 = oti.zeros(phi.shape)
1255 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64)
1256 r2.fused_sqdist(diffs, ell_sq)
1257 else:
1258 r2 = oti.mul(ell[0], diffs[0])
1259 r2 = oti.mul(r2, r2)
1260 for d in range(1, D):
1261 td = oti.mul(ell[d], diffs[d])
1262 r2 = oti.sum(r2, oti.mul(td, td))
1263 r_oti = oti.sqrt(oti.sum(r2, _eps ** 2))
1264 f_prime_r = kf._matern_grad_prebuild(r_oti)
1265 inv_r = oti.pow(r_oti, -1)
1267 base_matern = oti.mul(sigma_f_sq, oti.mul(f_prime_r, inv_r))
1268 if kernel_type == 'anisotropic':
1269 if hasattr(phi, 'fused_grad_all_dims'):
1270 scales = np.array([ln10 * ell[d] ** 2 for d in range(D)])
1271 grad_buf = np.zeros(D)
1272 base_matern.fused_grad_all_dims(diffs, scales, _vdot_factors, W_proj, grad_buf, FW_T)
1273 grad[:D] = grad_buf
1274 elif hasattr(phi, 'fused_scale_sq_mul_sparse'):
1275 dphi_buf = oti.zeros(phi.shape)
1276 for d in range(D):
1277 dphi_buf.fused_scale_sq_mul_sparse(diffs[d], base_matern, ln10 * ell[d] ** 2, d)
1278 grad[d] = _gc_block(dphi_buf)
1279 elif hasattr(phi, 'fused_scale_sq_mul'):
1280 dphi_buf = oti.zeros(phi.shape)
1281 for d in range(D):
1282 dphi_buf.fused_scale_sq_mul(diffs[d], base_matern, ln10 * ell[d] ** 2)
1283 grad[d] = _gc_block(dphi_buf)
1284 else:
1285 for d in range(D):
1286 d_sq = oti.mul(diffs[d], diffs[d])
1287 dphi_d = oti.mul(ln10 * ell[d] ** 2, oti.mul(d_sq, base_matern))
1288 grad[d] = _gc_block(dphi_d)
1289 else:
1290 ell_val = ell[0]
1291 if hasattr(phi, 'fused_sum_sq_sparse'):
1292 sum_dsq = oti.zeros(phi.shape)
1293 sum_dsq.fused_sum_sq_sparse(diffs)
1294 elif hasattr(phi, 'fused_sum_sq'):
1295 sum_dsq = oti.zeros(phi.shape)
1296 sum_dsq.fused_sum_sq(diffs)
1297 else:
1298 sum_dsq = oti.mul(diffs[0], diffs[0])
1299 for d in range(1, D):
1300 sum_dsq = oti.sum(sum_dsq, oti.mul(diffs[d], diffs[d]))
1301 dphi_e = oti.mul(ln10 * ell_val ** 2, oti.mul(sum_dsq, base_matern))
1302 grad[0] = _gc_block(dphi_e)
1304 elif kernel == 'SI':
1305 kf = self.model.kernel_factory
1306 si_prebuild = kf.SI_kernel_prebuild
1308 if kernel_type == 'anisotropic':
1309 ell = 10.0 ** x0[:D]
1310 else:
1311 ell = np.full(D, 10.0 ** float(x0[0]))
1313 si_vals = [si_prebuild(diffs[d]) for d in range(D)]
1314 term_vals = [oti.sum(1.0, oti.mul(ell[d], si_vals[d])) for d in range(D)]
1316 if kernel_type == 'anisotropic':
1317 for d in range(D):
1318 phi_over_term = oti.div(phi, term_vals[d])
1319 dphi_d = oti.mul(ln10 * ell[d],
1320 oti.mul(si_vals[d], phi_over_term))
1321 grad[d] = _gc_block(dphi_d)
1322 else:
1323 ell_val = ell[0]
1324 acc = oti.mul(si_vals[0], oti.div(phi, term_vals[0]))
1325 for d in range(1, D):
1326 acc = oti.sum(acc, oti.mul(si_vals[d],
1327 oti.div(phi, term_vals[d])))
1328 grad[0] = _gc_block(oti.mul(ln10 * ell_val, acc))
1330 return grad
1332 def _build_K_and_phi(self, x0):
1333 """
1334 Shared helper: build K (with noise), phi, n_bases, oti, diffs.
1336 Returns (K, phi, n_bases, oti, diffs).
1337 """
1338 diffs = self.model.differences_by_dim
1339 oti = self.model.kernel_factory.oti
1340 sigma_n_sq = (10.0 ** x0[-1]) ** 2
1342 phi = self.model.kernel_func(diffs, x0[:-1])
1343 if self.model.n_order == 0:
1344 n_bases = 0
1345 phi_exp = phi.real[np.newaxis, :, :]
1346 else:
1347 n_bases = phi.get_active_bases()[-1]
1348 deriv_order = 2 * self.model.n_order
1349 phi_exp = self._expand_derivs(phi, n_bases, deriv_order)
1351 self._ensure_kernel_plan(n_bases)
1352 if self._kernel_plan is not None:
1353 base_shape = phi.shape
1354 self._ensure_kernel_bufs(base_shape[0])
1355 phi_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1])
1356 K = utils.rbf_kernel_fast(phi_3d, self._kernel_plan, out=self._K_buf)
1357 else:
1358 K = utils.rbf_kernel(
1359 phi, phi_exp, self.model.n_order, n_bases,
1360 self.model.flattened_der_indices, self.model.powers,
1361 index=self.model.derivative_locations,
1362 )
1363 K.flat[::K.shape[0] + 1] += sigma_n_sq
1364 K.flat[::K.shape[0] + 1] += self.model.sigma_data_sq_diag
1365 return K, phi, n_bases, oti, diffs
1368 def _sparse_nlml_direct(self, x0):
1369 """
1370 Compute sparse NLML directly from phi_exp_3d, skipping full K
1371 construction and permutation.
1373 Returns (alpha_v, U, nll, phi, n_bases, oti, diffs) where the
1374 last four are needed by _compute_grad.
1375 """
1376 from jetgp.full_degp_sparse.sparse_cholesky import (
1377 build_U_from_phi, build_U_from_phi_flat,
1378 build_U_supernodes_from_phi,
1379 )
1381 diffs = self.model.differences_by_dim
1382 oti = self.model.kernel_factory.oti
1383 sigma_n_sq = (10.0 ** x0[-1]) ** 2
1385 phi = self.model.kernel_func(diffs, x0[:-1])
1386 n_bases = phi.get_active_bases()[-1]
1387 deriv_order = 2 * self.model.n_order
1388 phi_exp = self._expand_derivs(phi, n_bases, deriv_order)
1390 self._ensure_kernel_plan(n_bases)
1391 base_shape = phi.shape
1392 phi_3d = phi_exp.reshape(phi_exp.shape[0], base_shape[0], base_shape[1])
1394 # Build index maps (once)
1395 self._ensure_phi_index_maps(base_shape[0])
1396 k_type, k_phys, deriv_lookup, sign_lookup = self._k_index_map
1398 P_full = self.model.mmd_P_full
1399 N_total = len(P_full)
1401 if self.model.use_supernodes and self.model.sparse_supernodes_full is not None:
1402 U, _ = build_U_supernodes_from_phi(
1403 phi_3d, self.model.sparse_supernodes_full, N_total,
1404 sigma_n_sq,
1405 )
1406 else:
1407 if self._U_buf is None or self._U_buf.shape[0] != N_total:
1408 self._U_buf = np.zeros((N_total, N_total), order='F')
1410 U = build_U_from_phi_flat(
1411 phi_3d, self._block_phi_maps, N_total,
1412 sigma_n_sq, out=self._U_buf,
1413 )
1415 y_ord = self.model.y_train[P_full]
1416 nll = nlml_from_U(U, y_ord)
1418 alpha_ord = alpha_from_U(U, y_ord)
1419 alpha_v = np.empty_like(alpha_ord)
1420 alpha_v[P_full] = alpha_ord
1422 return alpha_v, U, nll, phi, n_bases, oti, diffs
1424 def _dense_nll_and_W(self, x0):
1425 """
1426 Dense Cholesky path: build full K, factor once, compute NLL and W.
1428 Used as a fallback when the sparsity pattern is too full for the
1429 block-wise sparse path to be efficient.
1431 Returns (W, alpha_v, nll, phi, n_bases, oti, diffs).
1432 """
1433 K, phi, n_bases, oti, diffs = self._build_K_and_phi(x0)
1434 N = K.shape[0]
1436 L, low = cho_factor(K, lower=True)
1437 alpha_v = cho_solve((L, low), self.model.y_train)
1439 nll = (0.5 * np.dot(self.model.y_train, alpha_v)
1440 + np.sum(np.log(np.diag(L)))
1441 + 0.5 * N * np.log(2 * np.pi))
1443 K_inv = cho_solve((L, low), np.eye(N))
1444 W = K_inv - np.outer(alpha_v, alpha_v)
1446 # Cache dense factors for prediction
1447 self.model._cached_L = L
1448 self.model._cached_low = low
1449 self.model._cached_alpha = alpha_v
1450 self.model._cached_U = None
1451 self.model._cached_P = None
1452 self.model._cached_params = x0.copy()
1454 return W, alpha_v, nll, phi, n_bases, oti, diffs
1455 def _sparse_U_alpha_nll(self, K):
1456 """
1457 Build sparse U, compute alpha and NLML. Does NOT form K^{-1}.
1459 Returns (alpha_v, U, nll) all in original index space.
1460 """
1461 P_full = self.model.mmd_P_full
1462 N_total = len(P_full)
1464 if self._P_ix is None:
1465 self._P_ix = np.ix_(P_full, P_full)
1466 K_ord = K[self._P_ix]
1467 y_ord = self.model.y_train[P_full]
1469 if self.model.use_supernodes and self.model.sparse_supernodes_full is not None:
1470 U, _ = build_U_supernodes(K_ord, self.model.sparse_supernodes_full, N_total)
1471 else:
1472 if self._U_buf is None or self._U_buf.shape[0] != N_total:
1473 self._U_buf = np.zeros((N_total, N_total), order='F')
1474 U = build_U(K_ord, self.model.sparse_S_full_arr, N_total,
1475 block_size=self.model.n_bases + 1, out=self._U_buf)
1477 nll = nlml_from_U(U, y_ord)
1479 # alpha in original space
1480 alpha_ord = alpha_from_U(U, y_ord)
1481 alpha_v = np.empty_like(alpha_ord)
1482 alpha_v[P_full] = alpha_ord
1484 return alpha_v, U, nll
1487 def _W_from_U(self, U, alpha_v):
1488 """
1489 Compute W = K^{-1} - αα^T from a pre-built sparse U.
1491 U is in MMD order, alpha_v is in original index space.
1492 Returns W in original index space.
1493 """
1494 P_full = self.model.mmd_P_full
1495 N_total = len(P_full)
1497 # K^{-1} = U @ U.T — exploit symmetry of result via dsyrk
1498 # (computes upper triangle only, ~2x fewer FLOPs than dgemm)
1499 if self._K_inv_buf is None or self._K_inv_buf.shape[0] != N_total:
1500 self._K_inv_buf = np.empty((N_total, N_total), order='F')
1501 K_inv_ord = blas.dsyrk(1.0, U, lower=1,
1502 c=self._K_inv_buf, overwrite_c=1)
1503 # No symmetrisation needed — _permute_and_subtract_outer reads
1504 # the lower triangle directly (contiguous in Fortran-order).
1506 W = np.empty((N_total, N_total))
1507 _permute_and_subtract_outer(K_inv_ord, alpha_v, P_full, W)
1508 return W
1510 def _sparse_W_and_alpha(self, K):
1511 """
1512 Compute W = K^{-1} - αα^T and alpha using the sparse U factor.
1514 K is in the ORIGINAL index space (size N_total × N_total).
1515 Returns (W, alpha_v, U, nll) all in original index space.
1516 """
1517 alpha_v, U, nll = self._sparse_U_alpha_nll(K)
1519 P_full = self.model.mmd_P_full
1520 N_total = len(P_full)
1522 # K^{-1} in original space: K_inv_ord = U @ U^T, then permute back.
1523 K_inv_ord = U @ U.T
1524 K_inv = np.empty_like(K_inv_ord)
1525 K_inv[np.ix_(P_full, P_full)] = K_inv_ord
1527 W = K_inv - np.outer(alpha_v, alpha_v)
1528 return W, alpha_v, U, nll
1530 def nll_grad(self, x0):
1531 """Analytic gradient of the NLL."""
1532 try:
1533 if self.model._use_dense_factor:
1534 W, alpha_v, nll, phi, n_bases, oti, diffs = self._dense_nll_and_W(x0)
1535 return self._compute_grad(x0, W, phi, n_bases, oti, diffs)
1536 else:
1537 alpha_v, U, nll, phi, n_bases, oti, diffs = self._sparse_nlml_direct(x0)
1538 return self._compute_grad_blockwise(x0, U, alpha_v, phi, n_bases, oti, diffs)
1539 except Exception:
1540 return np.zeros(len(x0))
1542 @profile
1543 def nll_and_grad(self, x0):
1544 """
1545 Compute NLL and its gradient in a single pass.
1547 Routes to either the dense Cholesky path or the sparse U path
1548 based on the sparsity pattern fill fraction.
1550 Returns
1551 -------
1552 nll : float
1553 grad : ndarray
1554 """
1555 try:
1556 if self.model._use_dense_factor:
1557 W, alpha_v, nll, phi, n_bases, oti, diffs = self._dense_nll_and_W(x0)
1558 grad = self._compute_grad(x0, W, phi, n_bases, oti, diffs)
1559 elif self.model.n_order > 0:
1560 alpha_v, U, nll, phi, n_bases, oti, diffs = self._sparse_nlml_direct(x0)
1562 # Cache for fast prediction (reused by degp.predict)
1563 self.model._cached_U = U
1564 self.model._cached_P = self.model.mmd_P_full
1565 self.model._cached_alpha = alpha_v
1566 self.model._cached_L = None
1567 self.model._cached_low = None
1568 self.model._cached_params = x0.copy()
1570 grad = self._compute_grad_blockwise(x0, U, alpha_v, phi, n_bases, oti, diffs)
1571 else:
1572 K, phi, n_bases, oti, diffs = self._build_K_and_phi(x0)
1573 alpha_v, U, nll = self._sparse_U_alpha_nll(K)
1575 self.model._cached_U = U
1576 self.model._cached_P = self.model.mmd_P_full
1577 self.model._cached_alpha = alpha_v
1578 self.model._cached_L = None
1579 self.model._cached_low = None
1580 self.model._cached_params = x0.copy()
1582 W = self._W_from_U(U, alpha_v)
1583 grad = self._compute_grad(x0, W, phi, n_bases, oti, diffs)
1584 except Exception:
1585 return 1e6, np.zeros(len(x0))
1587 if nll > 1e6:
1588 return 1e6, np.zeros(len(x0))
1589 return float(nll), grad
1591 def optimize_hyperparameters(self,
1592 optimizer="pso",
1593 **kwargs):
1594 """
1595 Optimize the DEGP model hyperparameters using Particle Swarm Optimization (PSO).
1597 Parameters:
1598 ----------
1599 n_restart_optimizer : int, default=20
1600 Maximum number of iterations for PSO.
1601 swarm_size : int, default=20
1602 Number of particles in the swarm.
1603 verbose : bool, default=True
1604 Controls verbosity of PSO output.
1606 Returns:
1607 -------
1608 best_x : ndarray
1609 The optimal set of hyperparameters found.
1610 """
1612 if isinstance(optimizer, str):
1613 if optimizer not in OPTIMIZERS:
1614 raise ValueError(
1615 f"Unknown optimizer '{optimizer}'. Available: {list(OPTIMIZERS.keys())}"
1616 )
1617 optimizer_fn = OPTIMIZERS[optimizer]
1618 else:
1619 optimizer_fn = optimizer # allow passing a callable directly
1621 bounds = self.model.bounds
1622 lb = [b[0] for b in bounds]
1623 ub = [b[1] for b in bounds]
1625 # Inject nll_and_grad (single Cholesky per step) for all gradient-aware optimizers.
1626 if optimizer in ('lbfgs', 'jade', 'pso') and 'func_and_grad' not in kwargs and 'grad_func' not in kwargs:
1627 kwargs['func_and_grad'] = self.nll_and_grad
1629 best_x, best_val = optimizer_fn(self.nll_wrapper, lb, ub, **kwargs)
1631 self.model.opt_x0 = best_x
1632 self.model.opt_nll = best_val
1635 return best_x