Coverage for jetgp/full_gddegp/gddegp_utils.py: 65%
457 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-07 00:57 -0500
1import numpy as np
2import pyoti.core as coti
3from line_profiler import profile
4import numba
7# =============================================================================
8# Numba-accelerated helper functions for efficient matrix slicing
9# =============================================================================
11@numba.jit(nopython=True, cache=True)
12def extract_rows(content_full, row_indices, n_cols):
13 """
14 Extract rows from content_full at specified indices.
16 Parameters
17 ----------
18 content_full : ndarray of shape (n_rows_full, n_cols)
19 Source matrix.
20 row_indices : ndarray of int64
21 Row indices to extract.
22 n_cols : int
23 Number of columns.
25 Returns
26 -------
27 result : ndarray of shape (len(row_indices), n_cols)
28 Extracted rows.
29 """
30 n_rows = len(row_indices)
31 result = np.empty((n_rows, n_cols))
32 for i in range(n_rows):
33 ri = row_indices[i]
34 for j in range(n_cols):
35 result[i, j] = content_full[ri, j]
36 return result
39@numba.jit(nopython=True, cache=True)
40def extract_cols(content_full, col_indices, n_rows):
41 """
42 Extract columns from content_full at specified indices.
44 Parameters
45 ----------
46 content_full : ndarray of shape (n_rows, n_cols_full)
47 Source matrix.
48 col_indices : ndarray of int64
49 Column indices to extract.
50 n_rows : int
51 Number of rows.
53 Returns
54 -------
55 result : ndarray of shape (n_rows, len(col_indices))
56 Extracted columns.
57 """
58 n_cols = len(col_indices)
59 result = np.empty((n_rows, n_cols))
60 for i in range(n_rows):
61 for j in range(n_cols):
62 result[i, j] = content_full[i, col_indices[j]]
63 return result
66@numba.jit(nopython=True, cache=True)
67def extract_submatrix(content_full, row_indices, col_indices):
68 """
69 Extract submatrix from content_full at specified row and column indices.
70 Replaces the expensive np.ix_ operation.
72 Parameters
73 ----------
74 content_full : ndarray of shape (n_rows_full, n_cols_full)
75 Source matrix.
76 row_indices : ndarray of int64
77 Row indices to extract.
78 col_indices : ndarray of int64
79 Column indices to extract.
81 Returns
82 -------
83 result : ndarray of shape (len(row_indices), len(col_indices))
84 Extracted submatrix.
85 """
86 n_rows = len(row_indices)
87 n_cols = len(col_indices)
88 result = np.empty((n_rows, n_cols))
89 for i in range(n_rows):
90 ri = row_indices[i]
91 for j in range(n_cols):
92 result[i, j] = content_full[ri, col_indices[j]]
93 return result
96@numba.jit(nopython=True, cache=True)
97def extract_submatrix_transposed(content_full, row_indices, col_indices):
98 """
99 Extract submatrix and return its transpose.
100 Replaces content_full[np.ix_(row_indices, col_indices)].T
102 Parameters
103 ----------
104 content_full : ndarray of shape (n_rows_full, n_cols_full)
105 Source matrix.
106 row_indices : ndarray of int64
107 Row indices to extract.
108 col_indices : ndarray of int64
109 Column indices to extract.
111 Returns
112 -------
113 result : ndarray of shape (len(col_indices), len(row_indices))
114 Transposed extracted submatrix.
115 """
116 n_rows = len(row_indices)
117 n_cols = len(col_indices)
118 result = np.empty((n_cols, n_rows))
119 for i in range(n_rows):
120 ri = row_indices[i]
121 for j in range(n_cols):
122 result[j, i] = content_full[ri, col_indices[j]]
123 return result
126@numba.jit(nopython=True, cache=True)
127def extract_rows_transposed(content_full, row_indices, n_cols):
128 """
129 Extract rows and return transposed result.
130 Replaces content_full[row_indices, :].T
132 Parameters
133 ----------
134 content_full : ndarray of shape (n_rows_full, n_cols)
135 Source matrix.
136 row_indices : ndarray of int64
137 Row indices to extract.
138 n_cols : int
139 Number of columns.
141 Returns
142 -------
143 result : ndarray of shape (n_cols, len(row_indices))
144 Transposed extracted rows.
145 """
146 n_rows = len(row_indices)
147 result = np.empty((n_cols, n_rows))
148 for i in range(n_rows):
149 ri = row_indices[i]
150 for j in range(n_cols):
151 result[j, i] = content_full[ri, j]
152 return result
155@numba.jit(nopython=True, cache=True)
156def extract_cols_transposed(content_full, col_indices, n_rows):
157 """
158 Extract columns and return transposed result.
159 Replaces content_full[:, col_indices].T
161 Parameters
162 ----------
163 content_full : ndarray of shape (n_rows, n_cols_full)
164 Source matrix.
165 col_indices : ndarray of int64
166 Column indices to extract.
167 n_rows : int
168 Number of rows.
170 Returns
171 -------
172 result : ndarray of shape (len(col_indices), n_rows)
173 Transposed extracted columns.
174 """
175 n_cols = len(col_indices)
176 result = np.empty((n_cols, n_rows))
177 for i in range(n_rows):
178 for j in range(n_cols):
179 result[j, i] = content_full[i, col_indices[j]]
180 return result
183@numba.jit(nopython=True, cache=True, parallel=False)
184def extract_and_assign(content_full, row_indices, col_indices, K,
185 row_start, col_start):
186 """
187 Extract submatrix and assign directly to K.
189 Parameters
190 ----------
191 content_full : ndarray of shape (n_rows_full, n_cols_full)
192 Source matrix.
193 row_indices : ndarray of int64
194 Row indices to extract.
195 col_indices : ndarray of int64
196 Column indices to extract.
197 K : ndarray
198 Target matrix to fill.
199 row_start : int
200 Starting row index in K.
201 col_start : int
202 Starting column index in K.
203 """
204 n_rows = len(row_indices)
205 n_cols = len(col_indices)
206 for i in range(n_rows):
207 ri = row_indices[i]
208 for j in range(n_cols):
209 K[row_start + i, col_start + j] = content_full[ri, col_indices[j]]
212@numba.jit(nopython=True, cache=True, parallel=False)
213def extract_and_assign_transposed(content_full, row_indices, col_indices, K,
214 row_start, col_start):
215 """
216 Extract submatrix and assign its transpose directly to K.
217 Replaces K[...] = content_full[np.ix_(row_indices, col_indices)].T
219 Parameters
220 ----------
221 content_full : ndarray of shape (n_rows_full, n_cols_full)
222 Source matrix.
223 row_indices : ndarray of int64
224 Row indices to extract from content_full.
225 col_indices : ndarray of int64
226 Column indices to extract from content_full.
227 K : ndarray
228 Target matrix to fill.
229 row_start : int
230 Starting row index in K.
231 col_start : int
232 Starting column index in K.
233 """
234 n_rows = len(row_indices)
235 n_cols = len(col_indices)
236 for i in range(n_rows):
237 ri = row_indices[i]
238 for j in range(n_cols):
239 # Transposed assignment: K[col_idx, row_idx] = content[row_idx, col_idx]
240 K[row_start + j, col_start + i] = content_full[ri, col_indices[j]]
243@numba.jit(nopython=True, cache=True)
244def extract_rows_and_assign(content_full, row_indices, K,
245 row_start, col_start, n_cols):
246 """
247 Extract rows and assign directly to K.
249 Parameters
250 ----------
251 content_full : ndarray of shape (n_rows_full, n_cols)
252 Source matrix.
253 row_indices : ndarray of int64
254 Row indices to extract.
255 K : ndarray
256 Target matrix to fill.
257 row_start : int
258 Starting row index in K.
259 col_start : int
260 Starting column index in K.
261 n_cols : int
262 Number of columns to copy.
263 """
264 n_rows = len(row_indices)
265 for i in range(n_rows):
266 ri = row_indices[i]
267 for j in range(n_cols):
268 K[row_start + i, col_start + j] = content_full[ri, j]
271@numba.jit(nopython=True, cache=True)
272def extract_cols_and_assign(content_full, col_indices, K,
273 row_start, col_start, n_rows):
274 """
275 Extract columns and assign directly to K.
277 Parameters
278 ----------
279 content_full : ndarray of shape (n_rows, n_cols_full)
280 Source matrix.
281 col_indices : ndarray of int64
282 Column indices to extract.
283 K : ndarray
284 Target matrix to fill.
285 row_start : int
286 Starting row index in K.
287 col_start : int
288 Starting column index in K.
289 n_rows : int
290 Number of rows to copy.
291 """
292 n_cols = len(col_indices)
293 for i in range(n_rows):
294 for j in range(n_cols):
295 K[row_start + i, col_start + j] = content_full[i, col_indices[j]]
298@numba.jit(nopython=True, cache=True)
299def extract_rows_and_assign_transposed(content_full, row_indices, K,
300 row_start, col_start, n_cols):
301 """
302 Extract rows and assign transposed result directly to K.
303 Replaces K[...] = content_full[row_indices, :].T
305 Parameters
306 ----------
307 content_full : ndarray of shape (n_rows_full, n_cols)
308 Source matrix.
309 row_indices : ndarray of int64
310 Row indices to extract.
311 K : ndarray
312 Target matrix to fill.
313 row_start : int
314 Starting row index in K.
315 col_start : int
316 Starting column index in K.
317 n_cols : int
318 Number of columns in content_full.
319 """
320 n_rows = len(row_indices)
321 for i in range(n_rows):
322 ri = row_indices[i]
323 for j in range(n_cols):
324 K[row_start + j, col_start + i] = content_full[ri, j]
327@numba.jit(nopython=True, cache=True)
328def extract_cols_and_assign_transposed(content_full, col_indices, K,
329 row_start, col_start, n_rows):
330 """
331 Extract columns and assign transposed result directly to K.
332 Replaces K[...] = content_full[:, col_indices].T
334 Parameters
335 ----------
336 content_full : ndarray of shape (n_rows, n_cols_full)
337 Source matrix.
338 col_indices : ndarray of int64
339 Column indices to extract.
340 K : ndarray
341 Target matrix to fill.
342 row_start : int
343 Starting row index in K.
344 col_start : int
345 Starting column index in K.
346 n_rows : int
347 Number of rows in content_full.
348 """
349 n_cols = len(col_indices)
350 for i in range(n_rows):
351 for j in range(n_cols):
352 K[row_start + j, col_start + i] = content_full[i, col_indices[j]]
355# =============================================================================
356# Derivative index transformation utilities
357# =============================================================================
359def make_first_odd(der_indices):
360 """Transform derivative indices to use odd bases (1, 3, 5, ...)."""
361 result = []
362 for group in der_indices:
363 new_group = []
364 for pair in group:
365 first = pair[0]
366 new_group.append([2 * first - 1, pair[1]])
367 result.append(new_group)
368 return result
371def make_first_even(der_indices):
372 """Transform derivative indices to use even bases (2, 4, 6, ...)."""
373 result = []
374 for group in der_indices:
375 new_group = []
376 for pair in group:
377 first = pair[0]
378 new_group.append([2 * first, pair[1]])
379 result.append(new_group)
380 return result
383# =============================================================================
384# Difference computation functions
385# =============================================================================
386def compute_dimension_differences(k, X1, X2, n1, n2, rays_X1, rays_X2,
387 derivative_locations_X1, derivative_locations_X2,
388 e_tags_1, e_tags_2, oti_module,
389 X1_np=None, X2_np=None, use_fused=False):
390 """
391 Compute differences for a single dimension k.
392 Only perturbs points at specified derivative_locations with their corresponding rays.
394 Parameters
395 ----------
396 k : int
397 Dimension index.
398 X1, X2 : oti.array or None
399 Input point arrays of shape (n1, d) and (n2, d). Can be None if use_fused=True.
400 n1, n2 : int
401 Number of points in X1, X2.
402 rays_X1 : list of ndarray or None
403 rays_X1[i] has shape (d, len(derivative_locations_X1[i])).
404 rays_X2 : list of ndarray or None
405 rays_X2[i] has shape (d, len(derivative_locations_X2[i])).
406 derivative_locations_X1 : list of list
407 derivative_locations_X1[i] contains indices of X1 points with direction i.
408 derivative_locations_X2 : list of list
409 derivative_locations_X2[i] contains indices of X2 points with direction i.
410 e_tags_1, e_tags_2 : list
411 OTI basis elements for each direction.
412 oti_module : module
413 The PyOTI static module.
414 X1_np, X2_np : ndarray or None
415 Numpy arrays of shape (n1, d) and (n2, d). Required if use_fused=True.
416 use_fused : bool
417 If True, use fused_from_real_with_perturbations C-level function.
419 Returns
420 -------
421 diffs_k : oti.array
422 Differences for dimension k with shape (n1, n2).
423 """
424 # Build perturbation vector for X1
425 perturb_X1_values = [0.0] * n1
426 if rays_X1 is not None:
427 for dir_idx in range(len(rays_X1)):
428 locs = derivative_locations_X1[dir_idx]
429 rays = rays_X1[dir_idx]
430 for j, pt_idx in enumerate(locs):
431 perturb_X1_values[pt_idx] = perturb_X1_values[pt_idx] + e_tags_1[dir_idx] * rays[k, j]
433 # Build perturbation vector for X2
434 perturb_X2_values = [0.0] * n2
435 if rays_X2 is not None:
436 for dir_idx in range(len(rays_X2)):
437 locs = derivative_locations_X2[dir_idx]
438 rays = rays_X2[dir_idx]
439 for j, pt_idx in enumerate(locs):
440 perturb_X2_values[pt_idx] = perturb_X2_values[pt_idx] + e_tags_2[dir_idx] * rays[k, j]
442 # Convert to OTI arrays
443 perturb_X1 = oti_module.array(perturb_X1_values)
444 perturb_X2 = oti_module.array(perturb_X2_values)
446 if use_fused:
447 # --- Fused path ---
448 real_diffs = np.ascontiguousarray(
449 X1_np[:, k:k+1] - X2_np[:, k:k+1].T, dtype=np.float64
450 )
451 # Reshape perturbations from (n,) to (n, 1) for fused function
452 perturb1 = oti_module.zeros((n1, 1)) + perturb_X1
453 perturb2 = oti_module.zeros((n2, 1)) + perturb_X2
454 diffs_k = oti_module.empty((n1, n2))
455 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2)
456 return diffs_k
458 # --- Fallback path ---
459 # Tag coordinates
460 X1_k_tagged = X1[:, k] + perturb_X1
461 X2_k_tagged = X2[:, k] + perturb_X2
463 # Compute differences
464 diffs_k = oti_module.zeros((n1, n2))
465 for i in range(n1):
466 diffs_k[i, :] = X1_k_tagged[i, 0] - oti_module.transpose(X2_k_tagged[:, 0])
468 return diffs_k
471def differences_by_dim_func(X1, X2, rays_X1, rays_X2, derivative_locations_X1, derivative_locations_X2,
472 n_order, oti_module, return_deriv=True):
473 """
474 Compute dimension-wise differences with OTI tagging on both X1 and X2.
476 GDDEGP uses a dual-tag OTI scheme: X1 points are tagged with odd bases
477 (e_1, e_3, e_5, ...) and X2 points with even bases (e_2, e_4, e_6, ...).
478 This requires ``n_bases = 2 * n_direction_types``.
480 The dual-tag approach is necessary because each point can have a unique
481 directional ray, and the kernel matrix requires derivatives with respect to
482 *both* sets of directions simultaneously. In the difference X1 - X2, the
483 OTI coefficient for basis e_i at position (a, b) encodes only the ray of
484 the point that was tagged with e_i. A single-tag scheme (tagging both X1
485 and X2 with the same basis) would conflate the two rays in the difference,
486 making it impossible to recover the correct cross-derivative
487 ``v_i(a)^T H v_j(b)`` needed for K_dd blocks, and producing an asymmetric
488 K_fd block when rays vary per point.
490 Parameters
491 ----------
492 X1 : ndarray of shape (n1, d)
493 First set of input points.
494 X2 : ndarray of shape (n2, d)
495 Second set of input points.
496 rays_X1 : list of ndarray or None
497 List of ray arrays for X1. rays_X1[i] has shape (d, len(derivative_locations_X1[i])).
498 rays_X2 : list of ndarray or None
499 List of ray arrays for X2. rays_X2[i] has shape (d, len(derivative_locations_X2[i])).
500 derivative_locations_X1 : list of list
501 derivative_locations_X1[i] contains indices of X1 points with derivative direction i.
502 derivative_locations_X2 : list of list
503 derivative_locations_X2[i] contains indices of X2 points with derivative direction i.
504 n_order : int
505 Derivative order for OTI tagging.
506 oti_module : module
507 The PyOTI static module (e.g., pyoti.static.onumm4n2).
508 return_deriv : bool, optional
509 If True, use order 2*n_order for derivative-derivative blocks.
511 Returns
512 -------
513 differences_by_dim : list of oti.array
514 List of length d, each element is an (n1, n2) OTI array.
515 """
516 X1_np = np.asarray(X1, dtype=np.float64)
517 X2_np = np.asarray(X2, dtype=np.float64)
518 n1, d = X1_np.shape
519 n2 = X2_np.shape[0]
521 # Check if the fused C-level function is available
522 _use_fused = hasattr(oti_module.zeros((1, 1)), 'fused_from_real_with_perturbations')
524 # Determine number of derivative directions
525 m1 = len(rays_X1) if rays_X1 is not None else 0
526 m2 = len(rays_X2) if rays_X2 is not None else 0
527 m = max(m1, m2)
529 # Pre-compute OTI basis elements
530 e_tags_1 = []
531 e_tags_2 = []
533 if n_order == 0:
534 e_tags_1 = [0] * m
535 e_tags_2 = [0] * m
536 elif not return_deriv:
537 for i in range(m):
538 e_tags_1.append(oti_module.e((2 * i + 1), order=n_order))
539 e_tags_2.append(oti_module.e((2 * i + 2), order=n_order))
540 else:
541 for i in range(m):
542 e_tags_1.append(oti_module.e((2 * i + 1), order=2 * n_order))
543 e_tags_2.append(oti_module.e((2 * i + 2), order=2 * n_order))
545 # Only convert to OTI arrays if using fallback path
546 _use_fused_path = _use_fused and n_order > 0
547 _use_fused_n0 = _use_fused and n_order == 0
548 if _use_fused_path or _use_fused_n0:
549 X1_oti = None
550 X2_oti = None
551 else:
552 X1_oti = oti_module.array(X1_np)
553 X2_oti = oti_module.array(X2_np)
555 # Compute differences for each dimension
556 differences_by_dim = []
558 if _use_fused_n0:
559 # n_order == 0: no perturbation, just real differences via fused path
560 perturb1 = oti_module.zeros((n1, 1))
561 perturb2 = oti_module.zeros((n2, 1))
562 for k in range(d):
563 real_diffs = np.ascontiguousarray(
564 X1_np[:, k:k+1] - X2_np[:, k:k+1].T, dtype=np.float64
565 )
566 diffs_k = oti_module.empty((n1, n2))
567 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2)
568 differences_by_dim.append(diffs_k)
569 return differences_by_dim
571 for k in range(d):
572 diffs_k = compute_dimension_differences(
573 k, X1_oti, X2_oti, n1, n2, rays_X1, rays_X2,
574 derivative_locations_X1, derivative_locations_X2,
575 e_tags_1, e_tags_2, oti_module,
576 X1_np=X1_np, X2_np=X2_np,
577 use_fused=_use_fused_path
578 )
579 differences_by_dim.append(diffs_k)
581 return differences_by_dim
584# =============================================================================
585# Derivative mapping utilities
586# =============================================================================
588def deriv_map(nbases, order):
589 """Create mapping from (order, index) to flattened index."""
590 k = 0
591 map_deriv = []
592 for ordi in range(order + 1):
593 ndir = coti.ndir_order(nbases, ordi)
594 map_deriv_i = [0] * ndir
595 for idx in range(ndir):
596 map_deriv_i[idx] = k
597 k += 1
598 map_deriv.append(map_deriv_i)
599 return map_deriv
602def transform_der_indices(der_indices, der_map):
603 """Transform derivative indices to flattened format."""
604 deriv_ind_transf = []
605 deriv_ind_order = []
606 for deriv in der_indices:
607 imdir = coti.imdir(deriv)
608 idx, order = imdir
609 deriv_ind_transf.append(der_map[order][idx])
610 deriv_ind_order.append(imdir)
611 return deriv_ind_transf, deriv_ind_order
614# =============================================================================
615# RBF Kernel Assembly Functions (Optimized with Numba)
616# =============================================================================
618@profile
619def rbf_kernel(
620 phi,
621 phi_exp,
622 n_order,
623 n_bases,
624 der_indices,
625 index=None
626):
627 """
628 Assembles the full GDDEGP covariance matrix with selective derivative coverage.
630 This version uses Numba-accelerated functions for efficient matrix slicing,
631 replacing expensive np.ix_ operations.
633 Parameters
634 ----------
635 phi : OTI array
636 Base kernel matrix from kernel_func(differences, length_scales).
637 phi_exp : ndarray
638 Expanded derivative array from phi.get_all_derivs().
639 n_order : int
640 Maximum derivative order.
641 n_bases : int
642 Number of OTI bases (must be even).
643 der_indices : list
644 Derivative index specifications.
645 index : list of list
646 index[i] contains indices of points with derivative direction i.
648 Returns
649 -------
650 K : ndarray
651 Kernel matrix with block structure based on derivative locations.
652 """
653 dh = coti.get_dHelp()
655 assert n_bases % 2 == 0, "n_bases must be an even number."
656 PHIrows, PHIcols = phi.shape
657 total_derivs = len(der_indices)
659 # Compute output matrix dimensions
660 n_deriv_rows = sum(len(locs) for locs in index)
661 n_deriv_cols = sum(len(locs) for locs in index)
662 n_output_rows = PHIrows + n_deriv_rows
663 n_output_cols = PHIcols + n_deriv_cols
665 der_map = deriv_map(n_bases, 2 * n_order)
667 # Pre-compute derivative index transformations
668 der_indices_even = make_first_even(der_indices)
669 der_indices_odd = make_first_odd(der_indices)
670 der_indices_tr_even, der_ind_order_even = transform_der_indices(der_indices_even, der_map)
671 der_indices_tr_odd, der_ind_order_odd = transform_der_indices(der_indices_odd, der_map)
673 # Convert index lists to numpy arrays for numba
674 index_arrays = [np.asarray(locs, dtype=np.int64) for locs in index]
676 # Compute block offsets
677 row_offsets = [0, PHIrows]
678 for i in range(total_derivs):
679 row_offsets.append(row_offsets[-1] + len(index[i]))
681 col_offsets = [0, PHIcols]
682 for i in range(total_derivs):
683 col_offsets.append(col_offsets[-1] + len(index[i]))
685 # Allocate output matrix
686 K = np.zeros((n_output_rows, n_output_cols))
688 # Fill blocks
689 for i in range(total_derivs + 1):
690 for j in range(total_derivs + 1):
692 if i == 0 and j == 0:
693 # K_ff: Full function-function block
694 K[0:PHIrows, 0:PHIcols] = phi_exp[0]
696 elif i == 0 and j > 0:
697 # K_fd: Function rows, derivative j columns
698 idx = der_indices_tr_even[j - 1]
699 col_locs = index_arrays[j - 1]
700 col_start = col_offsets[j]
702 # Use numba for efficient column extraction
703 extract_cols_and_assign(phi_exp[idx], col_locs, K,
704 0, col_start, PHIrows)
706 elif i > 0 and j == 0:
707 # K_df: Derivative i rows, function columns
708 idx = der_indices_tr_odd[i - 1]
709 row_locs = index_arrays[i - 1]
710 row_start = row_offsets[i]
712 # Use numba for efficient row extraction
713 extract_rows_and_assign(phi_exp[idx], row_locs, K,
714 row_start, 0, PHIcols)
716 else:
717 # K_dd: Derivative i rows, derivative j columns
718 imdir1 = der_ind_order_even[j - 1]
719 imdir2 = der_ind_order_odd[i - 1]
720 new_idx, new_ord = dh.mult_dir(
721 imdir1[0], imdir1[1], imdir2[0], imdir2[1])
722 idx = der_map[new_ord][new_idx]
724 row_locs = index_arrays[i - 1]
725 col_locs = index_arrays[j - 1]
726 row_start = row_offsets[i]
727 col_start = col_offsets[j]
729 # Use numba for efficient submatrix extraction (replaces np.ix_)
730 extract_and_assign(phi_exp[idx], row_locs, col_locs, K,
731 row_start, col_start)
733 return K
736@numba.jit(nopython=True, cache=True)
737def _assemble_kernel_numba(phi_exp_3d, K, n_rows_func, n_cols_func,
738 fd_flat_indices, df_flat_indices, dd_flat_indices,
739 idx_flat, idx_offsets, idx_sizes,
740 n_deriv_types, row_offsets, col_offsets):
741 """Fused numba kernel for GDDEGP K matrix assembly (no signs, even/odd bases)."""
742 # ff block
743 for r in range(n_rows_func):
744 for c in range(n_cols_func):
745 K[r, c] = phi_exp_3d[0, r, c]
746 # fd block (even indices)
747 for j in range(n_deriv_types):
748 fi = fd_flat_indices[j]
749 co = col_offsets[j]
750 off_j = idx_offsets[j]
751 sz_j = idx_sizes[j]
752 for r in range(n_rows_func):
753 for k in range(sz_j):
754 ci = idx_flat[off_j + k]
755 K[r, co + k] = phi_exp_3d[fi, r, ci]
756 # df block (odd indices)
757 for i in range(n_deriv_types):
758 fi = df_flat_indices[i]
759 ro = row_offsets[i]
760 off_i = idx_offsets[i]
761 sz_i = idx_sizes[i]
762 for k in range(sz_i):
763 ri = idx_flat[off_i + k]
764 for c in range(n_cols_func):
765 K[ro + k, c] = phi_exp_3d[fi, ri, c]
766 # dd block (even × odd)
767 for i in range(n_deriv_types):
768 ro = row_offsets[i]
769 off_i = idx_offsets[i]
770 sz_i = idx_sizes[i]
771 for j in range(n_deriv_types):
772 fi = dd_flat_indices[i, j]
773 co = col_offsets[j]
774 off_j = idx_offsets[j]
775 sz_j = idx_sizes[j]
776 for ki in range(sz_i):
777 ri = idx_flat[off_i + ki]
778 for kj in range(sz_j):
779 ci = idx_flat[off_j + kj]
780 K[ro + ki, co + kj] = phi_exp_3d[fi, ri, ci]
783@numba.jit(nopython=True, cache=True)
784def _project_W_to_phi_space(W, W_proj, n_rows_func, n_cols_func,
785 fd_flat_indices, df_flat_indices, dd_flat_indices,
786 idx_flat, idx_offsets, idx_sizes,
787 n_deriv_types, row_offsets, col_offsets):
788 """
789 Reverse of _assemble_kernel_numba: project W from K-space back into
790 phi_exp-space so that vdot(W, assemble(dphi_exp)) == vdot(W_proj, dphi_exp).
791 No-signs variant for GDDEGP even/odd bases.
792 """
793 for d in range(W_proj.shape[0]):
794 for r in range(W_proj.shape[1]):
795 for c in range(W_proj.shape[2]):
796 W_proj[d, r, c] = 0.0
797 for r in range(n_rows_func):
798 for c in range(n_cols_func):
799 W_proj[0, r, c] += W[r, c]
800 for j in range(n_deriv_types):
801 fi = fd_flat_indices[j]
802 co = col_offsets[j]
803 off_j = idx_offsets[j]
804 sz_j = idx_sizes[j]
805 for r in range(n_rows_func):
806 for k in range(sz_j):
807 ci = idx_flat[off_j + k]
808 W_proj[fi, r, ci] += W[r, co + k]
809 for i in range(n_deriv_types):
810 fi = df_flat_indices[i]
811 ro = row_offsets[i]
812 off_i = idx_offsets[i]
813 sz_i = idx_sizes[i]
814 for k in range(sz_i):
815 ri = idx_flat[off_i + k]
816 for c in range(n_cols_func):
817 W_proj[fi, ri, c] += W[ro + k, c]
818 for i in range(n_deriv_types):
819 ro = row_offsets[i]
820 off_i = idx_offsets[i]
821 sz_i = idx_sizes[i]
822 for j in range(n_deriv_types):
823 fi = dd_flat_indices[i, j]
824 co = col_offsets[j]
825 off_j = idx_offsets[j]
826 sz_j = idx_sizes[j]
827 for ki in range(sz_i):
828 ri = idx_flat[off_i + ki]
829 for kj in range(sz_j):
830 ci = idx_flat[off_j + kj]
831 W_proj[fi, ri, ci] += W[ro + ki, co + kj]
834def precompute_kernel_plan(n_order, n_bases, der_indices, powers, index):
835 """Precompute structural info for rbf_kernel_fast (GDDEGP even/odd variant)."""
836 dh = coti.get_dHelp()
837 assert n_bases % 2 == 0, "n_bases must be an even number."
838 der_map = deriv_map(n_bases, 2 * n_order)
840 n_deriv_types = len(der_indices)
841 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index]
843 index_sizes = np.array([len(idx) for idx in index_arrays], dtype=np.int64)
844 n_pts_with_derivs = int(index_sizes.sum())
846 idx_flat = np.concatenate(index_arrays) if n_deriv_types > 0 else np.array([], dtype=np.int64)
847 idx_offsets = np.zeros(n_deriv_types, dtype=np.int64)
848 for i in range(1, n_deriv_types):
849 idx_offsets[i] = idx_offsets[i - 1] + index_sizes[i - 1]
851 row_offsets = np.zeros(n_deriv_types, dtype=np.int64)
852 col_offsets = np.zeros(n_deriv_types, dtype=np.int64)
853 cumsum = 0
854 for i in range(n_deriv_types):
855 row_offsets[i] = cumsum
856 col_offsets[i] = cumsum
857 cumsum += index_sizes[i]
859 # Even/odd derivative transforms
860 der_indices_even = make_first_even(der_indices)
861 der_indices_odd = make_first_odd(der_indices)
862 der_indices_tr_even, der_ind_order_even = transform_der_indices(der_indices_even, der_map)
863 der_indices_tr_odd, der_ind_order_odd = transform_der_indices(der_indices_odd, der_map)
865 fd_flat_indices = np.array(der_indices_tr_even, dtype=np.int64)
866 df_flat_indices = np.array(der_indices_tr_odd, dtype=np.int64)
868 dd_flat_indices = np.empty((n_deriv_types, n_deriv_types), dtype=np.int64)
869 for i in range(n_deriv_types):
870 for j in range(n_deriv_types):
871 imdir1 = der_ind_order_even[j]
872 imdir2 = der_ind_order_odd[i]
873 new_idx, new_ord = dh.mult_dir(imdir1[0], imdir1[1], imdir2[0], imdir2[1])
874 dd_flat_indices[i, j] = der_map[new_ord][new_idx]
876 return {
877 'signs': np.ones(n_deriv_types + 1, dtype=np.float64), # unused, kept for API
878 'index_arrays': index_arrays,
879 'index_sizes': index_sizes,
880 'n_pts_with_derivs': n_pts_with_derivs,
881 'dd_flat_indices': dd_flat_indices,
882 'n_deriv_types': n_deriv_types,
883 'idx_flat': idx_flat,
884 'idx_offsets': idx_offsets,
885 'row_offsets': row_offsets,
886 'col_offsets': col_offsets,
887 'fd_flat_indices': fd_flat_indices,
888 'df_flat_indices': df_flat_indices,
889 }
892def rbf_kernel_fast(phi_exp_3d, plan, out=None):
893 """Fast kernel assembly using precomputed plan and fused numba kernel."""
894 n_rows_func = phi_exp_3d.shape[1]
895 n_cols_func = phi_exp_3d.shape[2]
896 total = n_rows_func + plan['n_pts_with_derivs']
897 if out is not None:
898 K = out
899 else:
900 K = np.empty((total, total))
902 if 'row_offsets_abs' in plan:
903 row_off = plan['row_offsets_abs']
904 col_off = plan['col_offsets_abs']
905 else:
906 row_off = plan['row_offsets'] + n_rows_func
907 col_off = plan['col_offsets'] + n_cols_func
909 _assemble_kernel_numba(
910 phi_exp_3d, K, n_rows_func, n_cols_func,
911 plan['fd_flat_indices'], plan['df_flat_indices'], plan['dd_flat_indices'],
912 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'],
913 plan['n_deriv_types'], row_off, col_off,
914 )
915 return K
918@profile
919def rbf_kernel_predictions(
920 phi,
921 phi_exp,
922 n_order,
923 n_bases,
924 der_indices,
925 return_deriv,
926 index=None,
927 common_derivs=None,
928 calc_cov=False,
929):
930 """
931 Constructs the RBF kernel matrix for predictions with selective derivative coverage.
933 This version uses Numba-accelerated functions for efficient matrix slicing.
935 Parameters
936 ----------
937 phi : OTI array
938 Base kernel matrix between test and training points.
939 phi_exp : ndarray
940 Expanded derivative array from phi.get_all_derivs().
941 n_order : int
942 Maximum derivative order.
943 n_bases : int
944 Number of OTI bases.
945 der_indices : list
946 Derivative specifications for training data.
947 return_deriv : bool
948 If True, predict derivatives at test points.
949 index : list of list
950 Training point indices for each derivative type.
951 common_derivs : list
952 Common derivative indices to predict.
953 calc_cov : bool
954 If True, computing covariance.
956 Returns
957 -------
958 K : ndarray
959 Prediction kernel matrix.
960 """
961 # Early return for covariance-only case
962 if calc_cov and not return_deriv:
963 return phi.real.T
965 dh = coti.get_dHelp()
967 n_train, n_test = phi.shape
968 n_deriv_types = len(der_indices)
969 n_deriv_types_pred = len(common_derivs) if common_derivs else 0
971 # Handle n_order = 0 case
972 if n_order == 0:
973 return phi.real.T
975 # Convert index lists to numpy arrays for numba
976 index_arrays = [np.asarray(locs, dtype=np.int64) for locs in index]
978 # Determine derivative map
979 if return_deriv:
980 der_map = deriv_map(n_bases, 2 * n_order)
981 derivative_locations_test = [np.arange(n_test, dtype=np.int64)] * n_deriv_types_pred
982 else:
983 der_map = deriv_map(n_bases, n_order)
985 # Create derivative index transformations
986 der_indices_even = make_first_even(der_indices)
987 der_indices_odd = make_first_odd(der_indices)
988 der_indices_tr_odd, der_ind_order_odd = transform_der_indices(der_indices_odd, der_map)
989 der_indices_odd_pred = make_first_odd(common_derivs) if common_derivs else []
990 der_indices_tr_odd_pred, der_ind_order_odd_pred = transform_der_indices(der_indices_odd_pred, der_map) if common_derivs else ([], [])
992 # Compute matrix dimensions
993 n_rows_func = n_test
994 if return_deriv:
995 n_rows_derivs = sum(len(locs) for locs in derivative_locations_test)
996 else:
997 n_rows_derivs = 0
998 total_rows = n_rows_func + n_rows_derivs
1000 if return_deriv and calc_cov:
1001 n_deriv_types = n_deriv_types_pred
1002 n_cols_derivs = sum(len(locs) for locs in derivative_locations_test)
1003 total_cols = n_train + n_cols_derivs
1004 else:
1005 n_cols_derivs = sum(len(locs) for locs in index)
1006 total_cols = n_train + n_cols_derivs
1008 # Compute block offsets
1009 row_offsets = [0, n_test]
1010 if return_deriv:
1011 for i in range(n_deriv_types_pred):
1012 row_offsets.append(row_offsets[-1] + len(derivative_locations_test[i]))
1014 col_offsets = [0, n_train]
1015 if return_deriv and calc_cov:
1016 for i in range(n_deriv_types):
1017 col_offsets.append(col_offsets[-1] + len(derivative_locations_test[i]))
1018 else:
1019 for i in range(n_deriv_types):
1020 col_offsets.append(col_offsets[-1] + len(index[i]))
1022 # Allocate output matrix
1023 K = np.zeros((total_rows, total_cols))
1024 base_shape = (n_train, n_test)
1026 # Block (0,0): Function-Function (K_ff)
1027 content_full = phi_exp[0].reshape(base_shape)
1028 K[:n_test, :n_train] = content_full.T
1030 # First Block-Row: Function-Derivative (K_fd)
1031 for j in range(n_deriv_types):
1032 col_locs = derivative_locations_test[j] if (return_deriv and calc_cov) else index_arrays[j]
1033 col_start = col_offsets[j + 1]
1035 flat_idx = der_indices_tr_odd_pred[j] if calc_cov else der_indices_tr_odd[j]
1036 content_full = phi_exp[flat_idx].reshape(base_shape)
1038 # Use numba for efficient row extraction with transpose
1039 extract_rows_and_assign_transposed(content_full, col_locs, K,
1040 0, col_start, n_test)
1042 if not return_deriv:
1043 return K
1045 # First Block-Column: Derivative-Function (K_df)
1046 der_indices_tr_even, der_ind_order_even = transform_der_indices(der_indices_even, der_map)
1047 der_indices_even_pred = make_first_even(common_derivs)
1048 der_indices_tr_even_pred, der_ind_order_even_pred = transform_der_indices(der_indices_even_pred, der_map)
1050 for i in range(n_deriv_types_pred):
1051 test_locs = derivative_locations_test[i]
1052 row_start = row_offsets[i + 1]
1054 flat_idx = der_indices_tr_even_pred[i]
1055 content_full = phi_exp[flat_idx].reshape(base_shape)
1057 # Use numba for efficient column extraction with transpose
1058 extract_cols_and_assign_transposed(content_full, test_locs, K,
1059 row_start, 0, n_train)
1061 # Inner Blocks: Derivative-Derivative (K_dd)
1062 for i in range(n_deriv_types_pred):
1063 test_locs = derivative_locations_test[i]
1064 row_start = row_offsets[i + 1]
1066 for j in range(n_deriv_types):
1067 col_locs = derivative_locations_test[j] if (return_deriv and calc_cov) else index_arrays[j]
1068 col_start = col_offsets[j + 1]
1070 imdir_train = der_ind_order_odd_pred[j] if calc_cov else der_ind_order_odd[j]
1071 imdir_test = der_ind_order_even_pred[i]
1072 new_idx, new_ord = dh.mult_dir(
1073 imdir_train[0], imdir_train[1],
1074 imdir_test[0], imdir_test[1]
1075 )
1076 flat_idx = der_map[new_ord][new_idx]
1078 content_full = phi_exp[flat_idx].reshape(base_shape)
1080 # Use numba for efficient submatrix extraction with transpose (replaces np.ix_ + .T)
1081 extract_and_assign_transposed(content_full, col_locs, test_locs, K,
1082 row_start, col_start)
1084 return K