Coverage for jetgp/full_gddegp/wgddegp_utils.py: 63%
520 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-07 00:57 -0500
1import numpy as np
2from line_profiler import profile
3import pyoti.core as coti
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]]
354# =============================================================================
355# Difference computation functions
356# =============================================================================
358def compute_dimension_differences(k, X1, X2, n1, n2, rays_X1, rays_X2,
359 derivative_locations_X1, derivative_locations_X2,
360 e_tags_1, e_tags_2, oti_module):
361 """
362 Compute differences for a single dimension k.
363 Only perturbs points at specified derivative_locations with their corresponding rays.
365 Parameters
366 ----------
367 k : int
368 Dimension index
369 X1, X2 : oti.array
370 Input point arrays of shape (n1, d) and (n2, d)
371 n1, n2 : int
372 Number of points in X1, X2
373 rays_X1 : list of ndarray or None
374 rays_X1[i] has shape (d, len(derivative_locations_X1[i]))
375 Column j corresponds to point derivative_locations_X1[i][j]
376 rays_X2 : list of ndarray or None
377 rays_X2[i] has shape (d, len(derivative_locations_X2[i]))
378 derivative_locations_X1 : list of list
379 derivative_locations_X1[i] contains indices of X1 points with direction i
380 derivative_locations_X2 : list of list
381 derivative_locations_X2[i] contains indices of X2 points with direction i
382 e_tags_1, e_tags_2 : list
383 OTI basis elements for each direction
384 oti_module : module
385 The PyOTI static module (e.g., pyoti.static.onumm4n2).
387 Returns
388 -------
389 diffs_k : oti.array
390 Differences for dimension k with shape (n1, n2).
391 """
392 # Build perturbation vector for X1
393 perturb_X1_values = [0.0] * n1
394 if rays_X1 is not None:
395 for dir_idx in range(len(rays_X1)):
396 locs = derivative_locations_X1[dir_idx]
397 rays = rays_X1[dir_idx]
398 for j, pt_idx in enumerate(locs):
399 perturb_X1_values[pt_idx] = perturb_X1_values[pt_idx] + e_tags_1[dir_idx] * rays[k, j]
401 # Build perturbation vector for X2
402 perturb_X2_values = [0.0] * n2
403 if rays_X2 is not None:
404 for dir_idx in range(len(rays_X2)):
405 locs = derivative_locations_X2[dir_idx]
406 rays = rays_X2[dir_idx]
407 for j, pt_idx in enumerate(locs):
408 perturb_X2_values[pt_idx] = perturb_X2_values[pt_idx] + e_tags_2[dir_idx] * rays[k, j]
410 # Convert to OTI arrays
411 perturb_X1 = oti_module.array(perturb_X1_values)
412 perturb_X2 = oti_module.array(perturb_X2_values)
414 # Tag coordinates
415 X1_k_tagged = X1[:, k] + perturb_X1
416 X2_k_tagged = X2[:, k] + perturb_X2
418 # Compute differences
419 diffs_k = oti_module.zeros((n1, n2))
420 for i in range(n1):
421 diffs_k[i, :] = X1_k_tagged[i, 0] - oti_module.transpose(X2_k_tagged[:, 0])
423 return diffs_k
426def differences_by_dim_func(X1, X2, rays_X1, rays_X2, derivative_locations_X1, derivative_locations_X2,
427 n_order, oti_module, return_deriv=True):
428 """
429 Compute dimension-wise differences with OTI tagging on both X1 and X2.
430 Only perturbs points at specified derivative_locations with their corresponding rays.
432 Parameters
433 ----------
434 X1 : ndarray of shape (n1, d)
435 First set of input points
436 X2 : ndarray of shape (n2, d)
437 Second set of input points
438 rays_X1 : list of ndarray or None
439 List of ray arrays for X1. rays_X1[i] has shape (d, len(derivative_locations_X1[i]))
440 where column j contains the ray direction for point derivative_locations_X1[i][j]
441 rays_X2 : list of ndarray or None
442 List of ray arrays for X2. rays_X2[i] has shape (d, len(derivative_locations_X2[i]))
443 derivative_locations_X1 : list of list
444 derivative_locations_X1[i] contains indices of X1 points that have derivative direction i
445 derivative_locations_X2 : list of list
446 derivative_locations_X2[i] contains indices of X2 points that have derivative direction i
447 n_order : int
448 Derivative order for OTI tagging
449 oti_module : module
450 The PyOTI static module (e.g., pyoti.static.onumm4n2).
451 return_deriv : bool, optional
452 If True, use order 2*n_order (for training kernel with derivative-derivative blocks)
453 If False, use order n_order (for prediction without derivative outputs)
455 Returns
456 -------
457 differences_by_dim : list of oti.array
458 List of length d, each element is an (n1, n2) OTI array of differences for that dimension
459 """
460 X1_np = np.asarray(X1, dtype=np.float64)
461 X2_np = np.asarray(X2, dtype=np.float64)
462 n1, d = X1_np.shape
463 n2 = X2_np.shape[0]
465 # Check if the fused C-level function is available
466 _use_fused = hasattr(oti_module.zeros((1, 1)), 'fused_from_real_with_perturbations')
468 # Fast path for n_order == 0: no perturbation, just real differences
469 if n_order == 0 and _use_fused:
470 perturb1 = oti_module.zeros((n1, 1))
471 perturb2 = oti_module.zeros((n2, 1))
472 differences_by_dim = []
473 for k in range(d):
474 real_diffs = np.ascontiguousarray(
475 X1_np[:, k:k+1] - X2_np[:, k:k+1].T, dtype=np.float64
476 )
477 diffs_k = oti_module.empty((n1, n2))
478 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2)
479 differences_by_dim.append(diffs_k)
480 return differences_by_dim
482 X1 = oti_module.array(X1_np)
483 X2 = oti_module.array(X2_np)
485 # Determine number of derivative directions from rays arrays
486 m1 = len(rays_X1) if rays_X1 is not None else 0
487 m2 = len(rays_X2) if rays_X2 is not None else 0
488 m = max(m1, m2)
490 # Pre-compute OTI basis elements
491 e_tags_1 = []
492 e_tags_2 = []
494 if n_order == 0:
495 e_tags_1 = [0] * m
496 e_tags_2 = [0] * m
497 elif not return_deriv:
498 for i in range(m):
499 e_tags_1.append(oti_module.e((2 * i + 1), order=n_order))
500 e_tags_2.append(oti_module.e((2 * i + 2), order=n_order))
501 else:
502 for i in range(m):
503 e_tags_1.append(oti_module.e((2 * i + 1), order=2 * n_order))
504 e_tags_2.append(oti_module.e((2 * i + 2), order=2 * n_order))
506 # Compute differences for each dimension
507 differences_by_dim = []
508 for k in range(d):
509 diffs_k = compute_dimension_differences(
510 k, X1, X2, n1, n2, rays_X1, rays_X2,
511 derivative_locations_X1, derivative_locations_X2,
512 e_tags_1, e_tags_2, oti_module
513 )
514 differences_by_dim.append(diffs_k)
516 return differences_by_dim
517# =============================================================================
518# Derivative index transformation utilities
519# =============================================================================
521def make_first_odd(der_indices):
522 """Transform derivative indices to use odd bases (1, 3, 5, ...)."""
523 result = []
524 for group in der_indices:
525 new_group = []
526 for pair in group:
527 first = pair[0]
528 new_group.append([2 * first - 1, pair[1]])
529 result.append(new_group)
530 return result
533def make_first_even(der_indices):
534 """Transform derivative indices to use even bases (2, 4, 6, ...)."""
535 result = []
536 for group in der_indices:
537 new_group = []
538 for pair in group:
539 first = pair[0]
540 new_group.append([2 * first, pair[1]])
541 result.append(new_group)
542 return result
546# =============================================================================
547# Derivative mapping utilities
548# =============================================================================
550def deriv_map(nbases, order):
551 """Create mapping from (order, index) to flattened index."""
552 k = 0
553 map_deriv = []
554 for ordi in range(order + 1):
555 ndir = coti.ndir_order(nbases, ordi)
556 map_deriv_i = [0] * ndir
557 for idx in range(ndir):
558 map_deriv_i[idx] = k
559 k += 1
560 map_deriv.append(map_deriv_i)
561 return map_deriv
564def transform_der_indices(der_indices, der_map):
565 """Transform derivative indices to flattened format."""
566 deriv_ind_transf = []
567 deriv_ind_order = []
568 for deriv in der_indices:
569 imdir = coti.imdir(deriv)
570 idx, order = imdir
571 deriv_ind_transf.append(der_map[order][idx])
572 deriv_ind_order.append(imdir)
573 return deriv_ind_transf, deriv_ind_order
576# =============================================================================
577# RBF Kernel Assembly Functions (Optimized with Numba)
578# =============================================================================
580@profile
581def rbf_kernel(
582 phi,
583 phi_exp,
584 n_order,
585 n_bases,
586 der_indices,
587 powers,
588 index=-1
589):
590 """
591 Assembles the full GDDEGP covariance matrix with support for selective
592 derivative coverage via derivative_locations.
594 This version uses Numba-accelerated functions for efficient matrix slicing,
595 replacing expensive np.ix_ operations.
597 Parameters
598 ----------
599 phi : OTI array
600 Base kernel matrix from kernel_func(differences, length_scales).
601 phi_exp : ndarray
602 Expanded derivative array from phi.get_all_derivs().
603 n_order : int
604 Maximum derivative order.
605 n_bases : int
606 Number of OTI bases (must be even).
607 der_indices : list
608 Derivative index specifications.
609 powers : list of int
610 Powers of (-1) applied to each term (unused but kept for API consistency).
611 index : list of list
612 index[i] contains indices of points with derivative direction i.
614 Returns
615 -------
616 K : ndarray
617 Kernel matrix with block structure based on derivative_locations.
618 """
619 dh = coti.get_dHelp()
621 highest_order = n_order
622 if n_order == 0:
623 n_bases = 0
624 phi_exp = phi.real
625 phi_exp = phi_exp[np.newaxis, :, :]
626 else:
627 n_bases = phi.get_active_bases()[-1]
628 phi_exp = phi.get_all_derivs(n_bases, 2 * highest_order)
629 assert n_bases % 2 == 0, "n_bases must be an even number."
630 PHIrows, PHIcols = phi.shape
631 total_derivs = len(der_indices)
633 # Compute output matrix dimensions
634 n_deriv_rows = sum(len(locs) for locs in index)
635 n_deriv_cols = sum(len(locs) for locs in index)
636 n_output_rows = PHIrows + n_deriv_rows
637 n_output_cols = PHIcols + n_deriv_cols
640 der_map = deriv_map(n_bases, 2 * highest_order)
642 row_iters = total_derivs + 1
643 col_iters = total_derivs + 1
645 # Pre-compute derivative index transformations
646 der_indices_even = make_first_even(der_indices)
647 der_indices_odd = make_first_odd(der_indices)
648 der_indices_tr_even, der_ind_order_even = transform_der_indices(der_indices_even, der_map)
649 der_indices_tr_odd, der_ind_order_odd = transform_der_indices(der_indices_odd, der_map)
651 # Convert index lists to numpy arrays for numba
652 index_arrays = [np.asarray(locs, dtype=np.int64) for locs in index]
654 # Compute block offsets
655 row_offsets = [0, PHIrows]
656 for i in range(total_derivs):
657 row_offsets.append(row_offsets[-1] + len(index[i]))
659 col_offsets = [0, PHIcols]
660 for i in range(total_derivs):
661 col_offsets.append(col_offsets[-1] + len(index[i]))
663 # Allocate output matrix
664 K = np.zeros((n_output_rows, n_output_cols))
666 # Fill blocks
667 for i in range(row_iters):
668 for j in range(col_iters):
670 if i == 0 and j == 0:
671 # K_ff: Full function-function block (all points)
672 idx = 0
673 K[0:PHIrows, 0:PHIcols] = phi_exp[idx]
675 elif i == 0 and j > 0:
676 # K_fd: Function rows (all), derivative j columns (at derivative_locations[j-1])
677 idx = der_indices_tr_even[j - 1]
678 col_locs = index_arrays[j - 1]
679 col_start = col_offsets[j]
681 # Use numba for efficient column extraction
682 extract_cols_and_assign(phi_exp[idx], col_locs, K,
683 0, col_start, PHIrows)
685 elif i > 0 and j == 0:
686 # K_df: Derivative i rows (at derivative_locations[i-1]), function columns (all)
687 idx = der_indices_tr_odd[i - 1]
688 row_locs = index_arrays[i - 1]
689 row_start = row_offsets[i]
691 # Use numba for efficient row extraction
692 extract_rows_and_assign(phi_exp[idx], row_locs, K,
693 row_start, 0, PHIcols)
695 else:
696 # K_dd: Derivative i rows, derivative j columns
697 imdir1 = der_ind_order_even[j - 1]
698 imdir2 = der_ind_order_odd[i - 1]
699 new_idx, new_ord = dh.mult_dir(
700 imdir1[0], imdir1[1], imdir2[0], imdir2[1])
701 idx = der_map[new_ord][new_idx]
703 row_locs = index_arrays[i - 1]
704 col_locs = index_arrays[j - 1]
705 row_start = row_offsets[i]
706 col_start = col_offsets[j]
708 # Use numba for efficient submatrix extraction (replaces np.ix_)
709 extract_and_assign(phi_exp[idx], row_locs, col_locs, K,
710 row_start, col_start)
712 return K
715@numba.jit(nopython=True, cache=True)
716def _assemble_kernel_numba(phi_exp_3d, K, n_rows_func, n_cols_func,
717 fd_flat_indices, df_flat_indices, dd_flat_indices,
718 idx_flat, idx_offsets, idx_sizes,
719 n_deriv_types, row_offsets, col_offsets):
720 """Fused numba kernel for GDDEGP K matrix assembly (no signs, even/odd bases)."""
721 # ff block
722 for r in range(n_rows_func):
723 for c in range(n_cols_func):
724 K[r, c] = phi_exp_3d[0, r, c]
725 # fd block (even indices)
726 for j in range(n_deriv_types):
727 fi = fd_flat_indices[j]
728 co = col_offsets[j]
729 off_j = idx_offsets[j]
730 sz_j = idx_sizes[j]
731 for r in range(n_rows_func):
732 for k in range(sz_j):
733 ci = idx_flat[off_j + k]
734 K[r, co + k] = phi_exp_3d[fi, r, ci]
735 # df block (odd indices)
736 for i in range(n_deriv_types):
737 fi = df_flat_indices[i]
738 ro = row_offsets[i]
739 off_i = idx_offsets[i]
740 sz_i = idx_sizes[i]
741 for k in range(sz_i):
742 ri = idx_flat[off_i + k]
743 for c in range(n_cols_func):
744 K[ro + k, c] = phi_exp_3d[fi, ri, c]
745 # dd block (even × odd)
746 for i in range(n_deriv_types):
747 ro = row_offsets[i]
748 off_i = idx_offsets[i]
749 sz_i = idx_sizes[i]
750 for j in range(n_deriv_types):
751 fi = dd_flat_indices[i, j]
752 co = col_offsets[j]
753 off_j = idx_offsets[j]
754 sz_j = idx_sizes[j]
755 for ki in range(sz_i):
756 ri = idx_flat[off_i + ki]
757 for kj in range(sz_j):
758 ci = idx_flat[off_j + kj]
759 K[ro + ki, co + kj] = phi_exp_3d[fi, ri, ci]
762@numba.jit(nopython=True, cache=True)
763def _project_W_to_phi_space(W, W_proj, n_rows_func, n_cols_func,
764 fd_flat_indices, df_flat_indices, dd_flat_indices,
765 idx_flat, idx_offsets, idx_sizes,
766 n_deriv_types, row_offsets, col_offsets):
767 """
768 Reverse of _assemble_kernel_numba: project W from K-space back into
769 phi_exp-space so that vdot(W, assemble(dphi_exp)) == vdot(W_proj, dphi_exp).
770 No-signs variant for WGDDEGP even/odd bases.
771 """
772 for d in range(W_proj.shape[0]):
773 for r in range(W_proj.shape[1]):
774 for c in range(W_proj.shape[2]):
775 W_proj[d, r, c] = 0.0
776 for r in range(n_rows_func):
777 for c in range(n_cols_func):
778 W_proj[0, r, c] += W[r, c]
779 for j in range(n_deriv_types):
780 fi = fd_flat_indices[j]
781 co = col_offsets[j]
782 off_j = idx_offsets[j]
783 sz_j = idx_sizes[j]
784 for r in range(n_rows_func):
785 for k in range(sz_j):
786 ci = idx_flat[off_j + k]
787 W_proj[fi, r, ci] += W[r, co + k]
788 for i in range(n_deriv_types):
789 fi = df_flat_indices[i]
790 ro = row_offsets[i]
791 off_i = idx_offsets[i]
792 sz_i = idx_sizes[i]
793 for k in range(sz_i):
794 ri = idx_flat[off_i + k]
795 for c in range(n_cols_func):
796 W_proj[fi, ri, c] += W[ro + k, c]
797 for i in range(n_deriv_types):
798 ro = row_offsets[i]
799 off_i = idx_offsets[i]
800 sz_i = idx_sizes[i]
801 for j in range(n_deriv_types):
802 fi = dd_flat_indices[i, j]
803 co = col_offsets[j]
804 off_j = idx_offsets[j]
805 sz_j = idx_sizes[j]
806 for ki in range(sz_i):
807 ri = idx_flat[off_i + ki]
808 for kj in range(sz_j):
809 ci = idx_flat[off_j + kj]
810 W_proj[fi, ri, ci] += W[ro + ki, co + kj]
813@numba.jit(nopython=True, cache=True)
814def _project_W_to_phi_space_accum(W, W_proj, n_rows_func, n_cols_func,
815 fd_flat_indices, df_flat_indices, dd_flat_indices,
816 idx_flat, idx_offsets, idx_sizes,
817 n_deriv_types, row_offsets, col_offsets):
818 """
819 Like _project_W_to_phi_space but accumulates into W_proj without zeroing.
820 Caller must zero W_proj before the first call.
821 No-signs variant for WGDDEGP even/odd bases.
822 """
823 for r in range(n_rows_func):
824 for c in range(n_cols_func):
825 W_proj[0, r, c] += W[r, c]
826 for j in range(n_deriv_types):
827 fi = fd_flat_indices[j]
828 co = col_offsets[j]
829 off_j = idx_offsets[j]
830 sz_j = idx_sizes[j]
831 for r in range(n_rows_func):
832 for k in range(sz_j):
833 ci = idx_flat[off_j + k]
834 W_proj[fi, r, ci] += W[r, co + k]
835 for i in range(n_deriv_types):
836 fi = df_flat_indices[i]
837 ro = row_offsets[i]
838 off_i = idx_offsets[i]
839 sz_i = idx_sizes[i]
840 for k in range(sz_i):
841 ri = idx_flat[off_i + k]
842 for c in range(n_cols_func):
843 W_proj[fi, ri, c] += W[ro + k, c]
844 for i in range(n_deriv_types):
845 ro = row_offsets[i]
846 off_i = idx_offsets[i]
847 sz_i = idx_sizes[i]
848 for j in range(n_deriv_types):
849 fi = dd_flat_indices[i, j]
850 co = col_offsets[j]
851 off_j = idx_offsets[j]
852 sz_j = idx_sizes[j]
853 for ki in range(sz_i):
854 ri = idx_flat[off_i + ki]
855 for kj in range(sz_j):
856 ci = idx_flat[off_j + kj]
857 W_proj[fi, ri, ci] += W[ro + ki, co + kj]
860def precompute_kernel_plan(n_order, n_bases, der_indices, powers, index):
861 """Precompute structural info for rbf_kernel_fast (GDDEGP even/odd variant)."""
862 dh = coti.get_dHelp()
863 assert n_bases % 2 == 0, "n_bases must be an even number."
864 der_map = deriv_map(n_bases, 2 * n_order)
866 n_deriv_types = len(der_indices)
867 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index]
869 index_sizes = np.array([len(idx) for idx in index_arrays], dtype=np.int64)
870 n_pts_with_derivs = int(index_sizes.sum())
872 idx_flat = np.concatenate(index_arrays) if n_deriv_types > 0 else np.array([], dtype=np.int64)
873 idx_offsets = np.zeros(n_deriv_types, dtype=np.int64)
874 for i in range(1, n_deriv_types):
875 idx_offsets[i] = idx_offsets[i - 1] + index_sizes[i - 1]
877 row_offsets = np.zeros(n_deriv_types, dtype=np.int64)
878 col_offsets = np.zeros(n_deriv_types, dtype=np.int64)
879 cumsum = 0
880 for i in range(n_deriv_types):
881 row_offsets[i] = cumsum
882 col_offsets[i] = cumsum
883 cumsum += index_sizes[i]
885 # Even/odd derivative transforms
886 der_indices_even = make_first_even(der_indices)
887 der_indices_odd = make_first_odd(der_indices)
888 der_indices_tr_even, der_ind_order_even = transform_der_indices(der_indices_even, der_map)
889 der_indices_tr_odd, der_ind_order_odd = transform_der_indices(der_indices_odd, der_map)
891 fd_flat_indices = np.array(der_indices_tr_even, dtype=np.int64)
892 df_flat_indices = np.array(der_indices_tr_odd, dtype=np.int64)
894 dd_flat_indices = np.empty((n_deriv_types, n_deriv_types), dtype=np.int64)
895 for i in range(n_deriv_types):
896 for j in range(n_deriv_types):
897 imdir1 = der_ind_order_even[j]
898 imdir2 = der_ind_order_odd[i]
899 new_idx, new_ord = dh.mult_dir(imdir1[0], imdir1[1], imdir2[0], imdir2[1])
900 dd_flat_indices[i, j] = der_map[new_ord][new_idx]
902 return {
903 'signs': np.ones(n_deriv_types + 1, dtype=np.float64), # unused, kept for API
904 'index_arrays': index_arrays,
905 'index_sizes': index_sizes,
906 'n_pts_with_derivs': n_pts_with_derivs,
907 'dd_flat_indices': dd_flat_indices,
908 'n_deriv_types': n_deriv_types,
909 'idx_flat': idx_flat,
910 'idx_offsets': idx_offsets,
911 'row_offsets': row_offsets,
912 'col_offsets': col_offsets,
913 'fd_flat_indices': fd_flat_indices,
914 'df_flat_indices': df_flat_indices,
915 }
918def rbf_kernel_fast(phi_exp_3d, plan, out=None):
919 """Fast kernel assembly using precomputed plan and fused numba kernel."""
920 n_rows_func = phi_exp_3d.shape[1]
921 n_cols_func = phi_exp_3d.shape[2]
922 total = n_rows_func + plan['n_pts_with_derivs']
923 if out is not None:
924 K = out
925 else:
926 K = np.empty((total, total))
928 if 'row_offsets_abs' in plan:
929 row_off = plan['row_offsets_abs']
930 col_off = plan['col_offsets_abs']
931 else:
932 row_off = plan['row_offsets'] + n_rows_func
933 col_off = plan['col_offsets'] + n_cols_func
935 _assemble_kernel_numba(
936 phi_exp_3d, K, n_rows_func, n_cols_func,
937 plan['fd_flat_indices'], plan['df_flat_indices'], plan['dd_flat_indices'],
938 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'],
939 plan['n_deriv_types'], row_off, col_off,
940 )
941 return K
944def rbf_kernel_predictions(
945 phi,
946 phi_exp,
947 n_order,
948 n_bases,
949 der_indices,
950 powers,
951 return_deriv,
952 index=-1,
953 common_derivs=None,
954 calc_cov=False,
955 powers_predict=None
956):
957 """
958 Constructs the RBF kernel matrix for predictions with selective derivative coverage.
960 This version uses Numba-accelerated functions for efficient matrix slicing.
962 Parameters
963 ----------
964 phi : OTI array
965 Base kernel matrix between test and training points.
966 phi_exp : ndarray
967 Expanded derivative array from phi.get_all_derivs().
968 n_order : int
969 Maximum derivative order.
970 n_bases : int
971 Number of OTI bases.
972 der_indices : list
973 Derivative specifications for training data.
974 powers : list of int
975 Sign powers (unused but kept for API consistency).
976 return_deriv : bool
977 If True, predict derivatives at test points.
978 index : list of list
979 Training point indices for each derivative type.
980 common_derivs : list
981 Common derivative indices to predict.
982 calc_cov : bool
983 If True, computing covariance.
984 powers_predict : list of int, optional
985 Sign powers for prediction derivatives (unused but kept for API consistency).
987 Returns
988 -------
989 K : ndarray
990 Prediction kernel matrix.
991 """
992 if calc_cov and not return_deriv:
993 return phi.real
995 dh = coti.get_dHelp()
997 n_train, n_test = phi.shape
998 n_deriv_types = len(der_indices)
999 n_deriv_types_pred = len(common_derivs) if common_derivs else 0
1001 # Handle n_order = 0 case
1002 if n_order == 0:
1003 return phi.real.T
1005 # Convert index lists to numpy arrays for numba
1006 index_arrays = [np.asarray(locs, dtype=np.int64) for locs in index]
1008 # Extract derivative components based on return_deriv
1009 if not return_deriv:
1010 phi_exp = phi.get_all_derivs(n_bases, n_order)
1011 der_map = deriv_map(n_bases, n_order)
1012 else:
1013 phi_exp = phi.get_all_derivs(n_bases, 2 * n_order)
1014 der_map = deriv_map(n_bases, 2 * n_order)
1016 # Create derivative index transformations
1017 der_indices_even = make_first_even(der_indices)
1018 der_indices_odd = make_first_odd(der_indices)
1019 der_indices_tr_odd, der_ind_order_odd = transform_der_indices(der_indices_odd, der_map)
1020 der_indices_odd_pred = make_first_odd(common_derivs) if common_derivs else []
1021 der_indices_tr_odd_pred, der_ind_order_odd_pred = transform_der_indices(der_indices_odd_pred, der_map) if common_derivs else ([], [])
1023 # Compute matrix dimensions
1024 n_rows_func = n_test
1025 if return_deriv:
1026 derivative_locations_test = [np.arange(n_test, dtype=np.int64)] * n_deriv_types_pred
1027 n_rows_derivs = sum(len(locs) for locs in derivative_locations_test)
1028 else:
1029 n_rows_derivs = 0
1030 total_rows = n_rows_func + n_rows_derivs
1032 if return_deriv and calc_cov:
1033 n_cols_func = n_train
1034 n_deriv_types = n_deriv_types_pred
1035 n_cols_derivs = sum(len(locs) for locs in derivative_locations_test)
1036 total_cols = n_cols_func + n_cols_derivs
1037 else:
1038 n_cols_func = n_train
1039 n_cols_derivs = sum(len(locs) for locs in index)
1040 total_cols = n_cols_func + n_cols_derivs
1042 # Compute block offsets
1043 row_offsets = [0, n_test]
1044 if return_deriv:
1045 for i in range(n_deriv_types_pred):
1046 row_offsets.append(row_offsets[-1] + len(derivative_locations_test[i]))
1048 col_offsets = [0, n_train]
1049 for i in range(n_deriv_types):
1050 col_offsets.append(col_offsets[-1] + len(index[i]))
1052 # Allocate output matrix
1053 K = np.zeros((total_rows, total_cols))
1054 base_shape = (n_train, n_test)
1056 # Block (0,0): Function-Function (K_ff)
1057 content_full = phi_exp[0].reshape(base_shape)
1058 K[:n_test, :n_train] = content_full.T
1060 # First Block-Row: Function-Derivative (K_fd)
1061 for j in range(n_deriv_types):
1062 train_locs = index_arrays[j]
1063 col_start = col_offsets[j + 1]
1065 flat_idx = der_indices_tr_odd_pred[j] if calc_cov else der_indices_tr_odd[j]
1066 content_full = phi_exp[flat_idx].reshape(base_shape)
1068 # Use numba for efficient row extraction with transpose
1069 extract_rows_and_assign_transposed(content_full, train_locs, K,
1070 0, col_start, n_test)
1072 if not return_deriv:
1073 return K
1075 # First Block-Column: Derivative-Function (K_df)
1076 der_indices_tr_even, der_ind_order_even = transform_der_indices(der_indices_even, der_map)
1077 der_indices_even_pred = make_first_even(common_derivs)
1078 der_indices_tr_even_pred, der_ind_order_even_pred = transform_der_indices(der_indices_even_pred, der_map)
1080 for i in range(n_deriv_types_pred):
1081 test_locs = derivative_locations_test[i]
1082 row_start = row_offsets[i + 1]
1084 flat_idx = der_indices_tr_even_pred[i]
1085 content_full = phi_exp[flat_idx].reshape(base_shape)
1087 # Use numba for efficient column extraction with transpose
1088 extract_cols_and_assign_transposed(content_full, test_locs, K,
1089 row_start, 0, n_train)
1091 # Inner Blocks: Derivative-Derivative (K_dd)
1092 for i in range(n_deriv_types_pred):
1093 test_locs = derivative_locations_test[i]
1094 row_start = row_offsets[i + 1]
1096 for j in range(n_deriv_types):
1097 train_locs = index_arrays[j]
1098 col_start = col_offsets[j + 1]
1100 imdir_train = der_ind_order_odd_pred[j] if calc_cov else der_ind_order_odd[j]
1101 imdir_test = der_ind_order_even_pred[i]
1102 new_idx, new_ord = dh.mult_dir(
1103 imdir_train[0], imdir_train[1],
1104 imdir_test[0], imdir_test[1]
1105 )
1106 flat_idx = der_map[new_ord][new_idx]
1108 content_full = phi_exp[flat_idx].reshape(base_shape)
1110 # Use numba for efficient submatrix extraction with transpose (replaces np.ix_ + .T)
1111 extract_and_assign_transposed(content_full, train_locs, test_locs, K,
1112 row_start, col_start)
1114 return K
1117# =============================================================================
1118# Utility functions
1119# =============================================================================
1121def determine_weights(diffs_by_dim, diffs_test, length_scales, kernel_func, sigma_n):
1122 """
1123 Vectorized version: compute interpolation weights for multiple test points at once.
1125 Parameters
1126 ----------
1127 diffs_by_dim : list of ndarray
1128 Pairwise differences between training points (by dimension).
1129 diffs_test : list of ndarray
1130 Pairwise differences between test points and training points (by dimension).
1131 Shape: each array is (n_test, n_train) or similar batch dimension.
1132 length_scales : array-like
1133 Kernel hyperparameters.
1134 kernel_func : callable
1135 Kernel function.
1136 sigma_n : float
1137 Noise parameter (if needed).
1139 Returns
1140 -------
1141 weights_matrix : ndarray of shape (n_test, n_train)
1142 Interpolation weights for each test point.
1143 """
1144 # Compute K matrix (training covariance) - same for all test points
1145 K = kernel_func(diffs_by_dim, length_scales).real
1146 n_train = K.shape[0]
1148 # Compute r vectors (test-train covariances) for all test points at once
1149 r_all = kernel_func(diffs_test, length_scales).real
1150 n_test = r_all.shape[0]
1152 # Build augmented system matrix M (same for all test points)
1153 M = np.zeros((n_train + 1, n_train + 1))
1154 M[:n_train, :n_train] = K
1155 M[:n_train, n_train] = 1
1156 M[n_train, :n_train] = 1
1157 M[n_train, n_train] = 0
1159 # Build augmented RHS for all test points
1160 r_augmented = np.zeros((n_test, n_train + 1))
1161 r_augmented[:, :n_train] = r_all
1162 r_augmented[:, n_train] = 1
1164 # Solve for all test points at once
1165 solution = np.linalg.solve(M, r_augmented.T)
1167 # Extract weights (exclude Lagrange multiplier)
1168 weights_matrix = solution[:n_train, :].T
1170 return weights_matrix
1173def to_list(x):
1174 """Convert tuple to list recursively."""
1175 if isinstance(x, tuple):
1176 return [to_list(i) for i in x]
1177 return x
1180def to_tuple(item):
1181 """Convert list to tuple recursively."""
1182 if isinstance(item, list):
1183 return tuple(to_tuple(x) for x in item)
1184 return item
1187def find_common_derivatives(all_indices):
1188 """Find derivative indices common to all submodels."""
1189 sets = [set(to_tuple(elem) for elem in idx_list) for idx_list in all_indices]
1190 return sets[0].intersection(*sets[1:])