Coverage for jetgp/wdegp/wdegp_utils.py: 58%
451 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-10 23:19 -0500
1import pyoti.core as coti
2import numpy as np
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, parallel=False)
97def extract_and_assign(content_full, row_indices, col_indices, K,
98 row_start, col_start, sign):
99 """
100 Extract submatrix and assign directly to K with sign multiplication.
101 Combines extraction and assignment in one pass for better performance.
103 Parameters
104 ----------
105 content_full : ndarray of shape (n_rows_full, n_cols_full)
106 Source matrix.
107 row_indices : ndarray of int64
108 Row indices to extract.
109 col_indices : ndarray of int64
110 Column indices to extract.
111 K : ndarray
112 Target matrix to fill.
113 row_start : int
114 Starting row index in K.
115 col_start : int
116 Starting column index in K.
117 sign : float
118 Sign multiplier (+1.0 or -1.0).
119 """
120 n_rows = len(row_indices)
121 n_cols = len(col_indices)
122 for i in range(n_rows):
123 ri = row_indices[i]
124 for j in range(n_cols):
125 K[row_start + i, col_start + j] = content_full[ri, col_indices[j]] * sign
128@numba.jit(nopython=True, cache=True)
129def extract_rows_and_assign(content_full, row_indices, K,
130 row_start, col_start, n_cols, sign):
131 """
132 Extract rows and assign directly to K with sign multiplication.
134 Parameters
135 ----------
136 content_full : ndarray of shape (n_rows_full, n_cols)
137 Source matrix.
138 row_indices : ndarray of int64
139 Row indices to extract.
140 K : ndarray
141 Target matrix to fill.
142 row_start : int
143 Starting row index in K.
144 col_start : int
145 Starting column index in K.
146 n_cols : int
147 Number of columns to copy.
148 sign : float
149 Sign multiplier (+1.0 or -1.0).
150 """
151 n_rows = len(row_indices)
152 for i in range(n_rows):
153 ri = row_indices[i]
154 for j in range(n_cols):
155 K[row_start + i, col_start + j] = content_full[ri, j] * sign
158@numba.jit(nopython=True, cache=True)
159def extract_cols_and_assign(content_full, col_indices, K,
160 row_start, col_start, n_rows, sign):
161 """
162 Extract columns and assign directly to K with sign multiplication.
164 Parameters
165 ----------
166 content_full : ndarray of shape (n_rows, n_cols_full)
167 Source matrix.
168 col_indices : ndarray of int64
169 Column indices to extract.
170 K : ndarray
171 Target matrix to fill.
172 row_start : int
173 Starting row index in K.
174 col_start : int
175 Starting column index in K.
176 n_rows : int
177 Number of rows to copy.
178 sign : float
179 Sign multiplier (+1.0 or -1.0).
180 """
181 n_cols = len(col_indices)
182 for i in range(n_rows):
183 for j in range(n_cols):
184 K[row_start + i, col_start + j] = content_full[i, col_indices[j]] * sign
186def differences_by_dim_func(X1, X2, n_order, oti_module, return_deriv=True):
187 """
188 Compute pairwise differences between two input arrays X1 and X2 for each dimension,
189 embedding hypercomplex units along each dimension for automatic differentiation.
191 For each dimension k, this function computes:
192 diff_k[i, j] = X1[i, k] + e_{k+1} - X2[j, k]
193 where e_{k+1} is a hypercomplex unit for the (k+1)-th dimension with order 2 * n_order.
195 Parameters
196 ----------
197 X1 : array_like of shape (n1, d)
198 First set of input points with n1 samples in d dimensions.
199 X2 : array_like of shape (n2, d)
200 Second set of input points with n2 samples in d dimensions.
201 n_order : int
202 The base order used to construct hypercomplex units (e_{k+1}) with order 2 * n_order.
203 oti_module : module
204 The PyOTI static module (e.g., pyoti.static.onumm4n2).
205 return_deriv : bool, optional
206 If True, use 2*n_order for derivative predictions.
208 Returns
209 -------
210 differences_by_dim : list of length d
211 A list where each element is an array of shape (n1, n2), containing the differences
212 between corresponding dimensions of X1 and X2, augmented with hypercomplex units.
213 """
214 # Keep numpy copies for fused path
215 X1_np = np.asarray(X1, dtype=np.float64)
216 X2_np = np.asarray(X2, dtype=np.float64)
217 n1, d = X1_np.shape
218 n2 = X2_np.shape[0]
220 # Check if the fused C-level function is available
221 _use_fused = hasattr(oti_module.zeros((1, 1)), 'fused_from_real_with_perturbations')
223 if _use_fused:
224 # --- Fused path: numpy broadcast for real part, C-level OTI fill ---
225 differences_by_dim = []
226 perturb2 = oti_module.zeros((n2, 1)) # X2 has no perturbation in WDEGP
228 if n_order == 0:
229 # No perturbation needed — just real differences
230 perturb1 = oti_module.zeros((n1, 1))
231 for k in range(d):
232 real_diffs = np.ascontiguousarray(
233 X1_np[:, k:k+1] - X2_np[:, k:k+1].T, dtype=np.float64
234 )
235 diffs_k = oti_module.empty((n1, n2))
236 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2)
237 differences_by_dim.append(diffs_k)
238 else:
239 if return_deriv:
240 hc_order = 2 * n_order
241 else:
242 hc_order = n_order
244 for k in range(d):
245 # Real differences via numpy broadcasting (fast)
246 real_diffs = np.ascontiguousarray(
247 X1_np[:, k:k+1] - X2_np[:, k:k+1].T, dtype=np.float64
248 )
249 # Perturbation: e_{k+1} broadcast to all n1 points
250 perturb1 = oti_module.zeros((n1, 1)) + oti_module.e(k + 1, order=hc_order)
251 # Fused fill: out[i,j].real = real_diffs[i,j], out[i,j].im = perturb1[i] - 0
252 diffs_k = oti_module.empty((n1, n2))
253 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2)
254 differences_by_dim.append(diffs_k)
256 return differences_by_dim
258 # --- Fallback: original Python loop path ---
259 X1 = oti_module.array(X1_np)
260 X2 = oti_module.array(X2_np)
262 differences_by_dim = []
264 if n_order == 0:
265 for k in range(d):
266 diffs_k = oti_module.zeros((n1, n2))
267 for i in range(n1):
268 diffs_k[i, :] = X1[i, k] - (oti_module.transpose(X2[:, k]))
269 differences_by_dim.append(diffs_k)
270 elif not return_deriv:
271 for k in range(d):
272 diffs_k = oti_module.zeros((n1, n2))
273 for i in range(n1):
274 diffs_k[i, :] = (
275 X1[i, k]
276 + oti_module.e(k + 1, order=n_order)
277 - (X2[:, k].T)
278 )
279 differences_by_dim.append(diffs_k)
280 else:
281 for k in range(d):
282 diffs_k = oti_module.zeros((n1, n2))
283 for i in range(n1):
284 diffs_k[i, :] = X1[i, k] - (X2[:, k].T)
285 differences_by_dim.append(diffs_k + oti_module.e(k + 1, order=2 * n_order))
287 return differences_by_dim
290# =============================================================================
291# Derivative mapping utilities
292# =============================================================================
294def deriv_map(nbases, order):
295 """
296 Creates a mapping from (order, index_within_order) to a single
297 flattened index for all derivative components.
298 """
299 k = 0
300 map_deriv = []
301 for ordi in range(order + 1):
302 ndir = coti.ndir_order(nbases, ordi)
303 map_deriv_i = [0] * ndir
304 for idx in range(ndir):
305 map_deriv_i[idx] = k
306 k += 1
307 map_deriv.append(map_deriv_i)
308 return map_deriv
311def transform_der_indices(der_indices, der_map):
312 """
313 Transforms a list of user-facing derivative specifications into the
314 internal (order, index) format and the final flattened index.
315 """
316 deriv_ind_transf = []
317 deriv_ind_order = []
318 for deriv in der_indices:
319 imdir = coti.imdir(deriv)
320 idx, order = imdir
321 deriv_ind_transf.append(der_map[order][idx])
322 deriv_ind_order.append(imdir)
323 return deriv_ind_transf, deriv_ind_order
326# =============================================================================
327# RBF Kernel Assembly Functions (Optimized with Numba)
328# =============================================================================
330def rbf_kernel(
331 phi,
332 phi_exp,
333 n_order,
334 n_bases,
335 der_indices,
336 powers,
337 index=-1,
338):
339 """
340 Constructs the RBF kernel matrix with derivative entries using an
341 efficient pre-allocation strategy combined with a single call to
342 extract all derivative components.
344 This version uses Numba-accelerated functions for efficient matrix slicing,
345 replacing expensive np.ix_ operations.
347 Parameters
348 ----------
349 phi : OTI array
350 Base kernel matrix from kernel_func(differences, length_scales).
351 phi_exp : ndarray
352 Expanded derivative array from phi.get_all_derivs().
353 n_order : int
354 Maximum derivative order.
355 n_bases : int
356 Number of OTI bases.
357 der_indices : list
358 Derivative specifications.
359 powers : list of int
360 Sign powers for each derivative type.
361 index : list of lists
362 Training point indices for each derivative type.
364 Returns
365 -------
366 K : ndarray
367 Full RBF kernel matrix with mixed function and derivative entries.
368 """
369 dh = coti.get_dHelp()
371 # Create maps to translate derivative specifications to flat indices
372 der_map = deriv_map(n_bases, 2 * n_order)
373 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map)
375 # Determine Block Sizes and Pre-allocate Matrix
376 n_rows_func, n_cols_func = phi.shape
377 n_deriv_types = len(der_indices)
378 n_pts_with_derivs_cols = sum(len(order_indices) for order_indices in index)
379 n_pts_with_derivs_rows = sum(len(order_indices) for order_indices in index)
380 total_rows = n_rows_func + n_pts_with_derivs_rows
381 total_cols = n_cols_func + n_pts_with_derivs_cols
383 K = np.zeros((total_rows, total_cols))
384 base_shape = (n_rows_func, n_cols_func)
386 # Pre-compute signs (avoid repeated exponentiation)
387 signs = np.array([(-1.0) ** p for p in powers], dtype=np.float64)
389 # Convert index lists to numpy arrays for numba
390 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index]
392 # Block (0,0): Function-Function (K_ff)
393 content_full = phi_exp[0].reshape(base_shape)
394 K[:n_rows_func, :n_cols_func] = content_full * signs[0]
396 # First Block-Row: Function-Derivative (K_fd)
397 col_offset = n_cols_func
398 for j in range(n_deriv_types):
399 flat_idx = der_indices_tr[j]
400 content_full = phi_exp[flat_idx].reshape(base_shape)
401 current_indices = index_arrays[j]
402 n_pts_this_order = len(current_indices)
404 # Use numba for efficient column extraction and assignment
405 extract_cols_and_assign(content_full, current_indices, K,
406 0, col_offset, n_rows_func, signs[j + 1])
407 col_offset += n_pts_this_order
409 # First Block-Column: Derivative-Function (K_df)
410 row_offset = n_rows_func
411 for i in range(n_deriv_types):
412 flat_idx = der_indices_tr[i]
413 content_full = phi_exp[flat_idx].reshape(base_shape)
414 current_indices = index_arrays[i]
415 n_pts_this_order = len(current_indices)
417 # Use numba for efficient row extraction and assignment
418 extract_rows_and_assign(content_full, current_indices, K,
419 row_offset, 0, n_cols_func, signs[0])
420 row_offset += n_pts_this_order
422 # Inner Blocks: Derivative-Derivative (K_dd)
423 row_offset = n_rows_func
424 for i in range(n_deriv_types):
425 col_offset = n_cols_func
426 row_indices = index_arrays[i]
427 n_pts_row = len(row_indices)
429 for j in range(n_deriv_types):
430 col_indices = index_arrays[j]
431 n_pts_col = len(col_indices)
433 # Multiply the derivative indices to find the correct flat index
434 imdir1 = der_ind_order[j]
435 imdir2 = der_ind_order[i]
436 new_idx, new_ord = dh.mult_dir(
437 imdir1[0], imdir1[1], imdir2[0], imdir2[1])
438 flat_idx = der_map[new_ord][new_idx]
439 content_full = phi_exp[flat_idx].reshape(base_shape)
441 # Use numba for efficient submatrix extraction and assignment (replaces np.ix_)
442 extract_and_assign(content_full, row_indices, col_indices, K,
443 row_offset, col_offset, signs[j + 1])
445 col_offset += n_pts_col
446 row_offset += n_pts_row
448 return K
451@numba.jit(nopython=True, cache=True)
452def _assemble_kernel_numba(phi_exp_3d, K, n_rows_func, n_cols_func,
453 fd_flat_indices, df_flat_indices, dd_flat_indices,
454 idx_flat, idx_offsets, idx_sizes,
455 signs, n_deriv_types, row_offsets, col_offsets):
456 """
457 Fused numba kernel that assembles the entire K matrix in a single call.
458 Handles ff, fd, df, and dd blocks without Python-level loop overhead.
459 """
460 # Block (0,0): Function-Function
461 s0 = signs[0]
462 for r in range(n_rows_func):
463 for c in range(n_cols_func):
464 K[r, c] = phi_exp_3d[0, r, c] * s0
466 # First Block-Row: Function-Derivative (fd)
467 for j in range(n_deriv_types):
468 fi = fd_flat_indices[j]
469 sj = signs[j + 1]
470 co = col_offsets[j]
471 off_j = idx_offsets[j]
472 sz_j = idx_sizes[j]
473 for r in range(n_rows_func):
474 for k in range(sz_j):
475 ci = idx_flat[off_j + k]
476 K[r, co + k] = phi_exp_3d[fi, r, ci] * sj
478 # First Block-Column: Derivative-Function (df)
479 for i in range(n_deriv_types):
480 fi = df_flat_indices[i]
481 ro = row_offsets[i]
482 off_i = idx_offsets[i]
483 sz_i = idx_sizes[i]
484 for k in range(sz_i):
485 ri = idx_flat[off_i + k]
486 for c in range(n_cols_func):
487 K[ro + k, c] = phi_exp_3d[fi, ri, c] * s0
489 # Inner Blocks: Derivative-Derivative (dd)
490 for i in range(n_deriv_types):
491 ro = row_offsets[i]
492 off_i = idx_offsets[i]
493 sz_i = idx_sizes[i]
494 for j in range(n_deriv_types):
495 fi = dd_flat_indices[i, j]
496 sj = signs[j + 1]
497 co = col_offsets[j]
498 off_j = idx_offsets[j]
499 sz_j = idx_sizes[j]
500 for ki in range(sz_i):
501 ri = idx_flat[off_i + ki]
502 for kj in range(sz_j):
503 ci = idx_flat[off_j + kj]
504 K[ro + ki, co + kj] = phi_exp_3d[fi, ri, ci] * sj
507@numba.jit(nopython=True, cache=True)
508def _project_W_to_phi_space(W, W_proj, n_rows_func, n_cols_func,
509 fd_flat_indices, df_flat_indices, dd_flat_indices,
510 idx_flat, idx_offsets, idx_sizes,
511 signs, n_deriv_types, row_offsets, col_offsets):
512 """
513 Reverse of _assemble_kernel_numba: project W from K-space back into
514 phi_exp-space so that vdot(W, assemble(dphi_exp)) == vdot(W_proj, dphi_exp).
515 """
516 for d in range(W_proj.shape[0]):
517 for r in range(W_proj.shape[1]):
518 for c in range(W_proj.shape[2]):
519 W_proj[d, r, c] = 0.0
520 s0 = signs[0]
521 for r in range(n_rows_func):
522 for c in range(n_cols_func):
523 W_proj[0, r, c] += s0 * W[r, c]
524 for j in range(n_deriv_types):
525 fi = fd_flat_indices[j]
526 sj = signs[j + 1]
527 co = col_offsets[j]
528 off_j = idx_offsets[j]
529 sz_j = idx_sizes[j]
530 for r in range(n_rows_func):
531 for k in range(sz_j):
532 ci = idx_flat[off_j + k]
533 W_proj[fi, r, ci] += sj * W[r, co + k]
534 for i in range(n_deriv_types):
535 fi = df_flat_indices[i]
536 ro = row_offsets[i]
537 off_i = idx_offsets[i]
538 sz_i = idx_sizes[i]
539 for k in range(sz_i):
540 ri = idx_flat[off_i + k]
541 for c in range(n_cols_func):
542 W_proj[fi, ri, c] += s0 * W[ro + k, c]
543 for i in range(n_deriv_types):
544 ro = row_offsets[i]
545 off_i = idx_offsets[i]
546 sz_i = idx_sizes[i]
547 for j in range(n_deriv_types):
548 fi = dd_flat_indices[i, j]
549 sj = signs[j + 1]
550 co = col_offsets[j]
551 off_j = idx_offsets[j]
552 sz_j = idx_sizes[j]
553 for ki in range(sz_i):
554 ri = idx_flat[off_i + ki]
555 for kj in range(sz_j):
556 ci = idx_flat[off_j + kj]
557 W_proj[fi, ri, ci] += sj * W[ro + ki, co + kj]
560@numba.jit(nopython=True, cache=True)
561def _project_W_to_phi_space_accum(W, W_proj, n_rows_func, n_cols_func,
562 fd_flat_indices, df_flat_indices, dd_flat_indices,
563 idx_flat, idx_offsets, idx_sizes,
564 signs, n_deriv_types, row_offsets, col_offsets):
565 """
566 Like _project_W_to_phi_space but accumulates into W_proj without zeroing.
567 Caller must zero W_proj before the first call.
568 """
569 s0 = signs[0]
570 for r in range(n_rows_func):
571 for c in range(n_cols_func):
572 W_proj[0, r, c] += s0 * W[r, c]
573 for j in range(n_deriv_types):
574 fi = fd_flat_indices[j]
575 sj = signs[j + 1]
576 co = col_offsets[j]
577 off_j = idx_offsets[j]
578 sz_j = idx_sizes[j]
579 for r in range(n_rows_func):
580 for k in range(sz_j):
581 ci = idx_flat[off_j + k]
582 W_proj[fi, r, ci] += sj * W[r, co + k]
583 for i in range(n_deriv_types):
584 fi = df_flat_indices[i]
585 ro = row_offsets[i]
586 off_i = idx_offsets[i]
587 sz_i = idx_sizes[i]
588 for k in range(sz_i):
589 ri = idx_flat[off_i + k]
590 for c in range(n_cols_func):
591 W_proj[fi, ri, c] += s0 * W[ro + k, c]
592 for i in range(n_deriv_types):
593 ro = row_offsets[i]
594 off_i = idx_offsets[i]
595 sz_i = idx_sizes[i]
596 for j in range(n_deriv_types):
597 fi = dd_flat_indices[i, j]
598 sj = signs[j + 1]
599 co = col_offsets[j]
600 off_j = idx_offsets[j]
601 sz_j = idx_sizes[j]
602 for ki in range(sz_i):
603 ri = idx_flat[off_i + ki]
604 for kj in range(sz_j):
605 ci = idx_flat[off_j + kj]
606 W_proj[fi, ri, ci] += sj * W[ro + ki, co + kj]
609def precompute_kernel_plan(n_order, n_bases, der_indices, powers, index):
610 """
611 Precompute all structural information needed by rbf_kernel so it can be
612 reused across repeated calls with different phi_exp values.
614 Returns a dict containing flat indices, signs, index arrays, precomputed
615 offsets/sizes, and mult_dir results for the dd block.
616 """
617 dh = coti.get_dHelp()
618 der_map = deriv_map(n_bases, 2 * n_order)
619 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map)
621 n_deriv_types = len(der_indices)
622 signs = np.array([(-1.0) ** p for p in powers], dtype=np.float64)
623 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index]
625 # Precompute sizes and offsets
626 index_sizes = np.array([len(idx) for idx in index_arrays], dtype=np.int64)
627 n_pts_with_derivs = int(index_sizes.sum())
629 # Pack all index arrays into a single flat array with offsets
630 idx_flat = np.concatenate(index_arrays) if n_deriv_types > 0 else np.array([], dtype=np.int64)
631 idx_offsets = np.zeros(n_deriv_types, dtype=np.int64)
632 for i in range(1, n_deriv_types):
633 idx_offsets[i] = idx_offsets[i - 1] + index_sizes[i - 1]
635 # Precompute row/col offsets in K for each deriv type
636 row_offsets = np.zeros(n_deriv_types, dtype=np.int64)
637 col_offsets = np.zeros(n_deriv_types, dtype=np.int64)
638 # Note: n_rows_func == n_cols_func for training kernel, but we store
639 # offsets relative to n_rows_func which is added at call time
640 cumsum = 0
641 for i in range(n_deriv_types):
642 row_offsets[i] = cumsum # relative to n_rows_func
643 col_offsets[i] = cumsum # relative to n_cols_func
644 cumsum += index_sizes[i]
646 # Precompute mult_dir results for dd blocks
647 dd_flat_indices = np.empty((n_deriv_types, n_deriv_types), dtype=np.int64)
648 for i in range(n_deriv_types):
649 for j in range(n_deriv_types):
650 imdir1 = der_ind_order[j]
651 imdir2 = der_ind_order[i]
652 new_idx, new_ord = dh.mult_dir(
653 imdir1[0], imdir1[1], imdir2[0], imdir2[1])
654 dd_flat_indices[i, j] = der_map[new_ord][new_idx]
656 # fd and df flat indices as arrays
657 fd_flat_indices = np.array(der_indices_tr, dtype=np.int64)
658 df_flat_indices = np.array(der_indices_tr, dtype=np.int64)
660 return {
661 'der_indices_tr': der_indices_tr,
662 'signs': signs,
663 'index_arrays': index_arrays,
664 'index_sizes': index_sizes,
665 'n_pts_with_derivs': n_pts_with_derivs,
666 'dd_flat_indices': dd_flat_indices,
667 'n_deriv_types': n_deriv_types,
668 # Fused kernel data
669 'idx_flat': idx_flat,
670 'idx_offsets': idx_offsets,
671 'row_offsets': row_offsets,
672 'col_offsets': col_offsets,
673 'fd_flat_indices': fd_flat_indices,
674 'df_flat_indices': df_flat_indices,
675 }
678def rbf_kernel_fast(phi_exp_3d, plan, out=None):
679 """
680 Fast kernel assembly using a precomputed plan and fused numba kernel.
682 Parameters
683 ----------
684 phi_exp_3d : ndarray of shape (n_derivs, n_rows_func, n_cols_func)
685 Pre-reshaped expanded derivative array.
686 plan : dict
687 Precomputed plan from precompute_kernel_plan().
688 out : ndarray, optional
689 Pre-allocated output array. If None, a new array is allocated.
691 Returns
692 -------
693 K : ndarray
694 Full kernel matrix.
695 """
696 n_rows_func = phi_exp_3d.shape[1]
697 n_cols_func = phi_exp_3d.shape[2]
698 total = n_rows_func + plan['n_pts_with_derivs']
699 if out is not None:
700 K = out
701 else:
702 K = np.empty((total, total))
704 if 'row_offsets_abs' in plan:
705 row_off = plan['row_offsets_abs']
706 col_off = plan['col_offsets_abs']
707 else:
708 row_off = plan['row_offsets'] + n_rows_func
709 col_off = plan['col_offsets'] + n_cols_func
711 _assemble_kernel_numba(
712 phi_exp_3d, K, n_rows_func, n_cols_func,
713 plan['fd_flat_indices'], plan['df_flat_indices'], plan['dd_flat_indices'],
714 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'],
715 plan['signs'], plan['n_deriv_types'], row_off, col_off,
716 )
718 return K
721def rbf_kernel_predictions(
722 phi,
723 phi_exp,
724 n_order,
725 n_bases,
726 der_indices,
727 powers,
728 return_deriv,
729 index=-1,
730 common_derivs=None,
731 calc_cov=False,
732 powers_predict=None
733):
734 """
735 Constructs the RBF kernel matrix for predictions with derivative entries.
737 This version uses Numba-accelerated functions for efficient matrix slicing.
739 Parameters
740 ----------
741 phi : OTI array
742 Base kernel matrix between test and training points.
743 phi_exp : ndarray
744 Expanded derivative array from phi.get_all_derivs().
745 n_order : int
746 Maximum derivative order.
747 n_bases : int
748 Number of OTI bases.
749 der_indices : list
750 Derivative specifications for training data.
751 powers : list of int
752 Sign powers for each derivative type.
753 return_deriv : bool
754 If True, predict derivatives at test points.
755 index : list of lists
756 Training point indices for each derivative type.
757 common_derivs : list
758 Common derivative indices to predict.
759 calc_cov : bool
760 If True, computing covariance.
761 powers_predict : list of int, optional
762 Sign powers for prediction derivatives.
764 Returns
765 -------
766 K : ndarray
767 Prediction kernel matrix.
768 """
769 if calc_cov and not return_deriv:
770 return phi.real
772 dh = coti.get_dHelp()
774 n_rows_func, n_cols_func = phi.shape
775 n_deriv_types = len(der_indices)
776 n_deriv_types_pred = len(common_derivs) if common_derivs else 0
778 # Pre-compute signs
779 signs = np.array([(-1.0) ** p for p in powers], dtype=np.float64)
780 if powers_predict is not None:
781 signs_predict = np.array([(-1.0) ** p for p in powers_predict], dtype=np.float64)
782 else:
783 signs_predict = signs
785 if return_deriv:
786 der_map = deriv_map(n_bases, 2 * n_order)
787 index_2 = np.arange(phi_exp.shape[-1], dtype=np.int64)
788 if calc_cov:
789 index_cov = np.arange(phi_exp.shape[-1], dtype=np.int64)
790 n_deriv_types = n_deriv_types_pred
791 n_pts_with_derivs_rows = n_deriv_types * len([i for i in range(n_cols_func) if i < len(index_2)])
792 else:
793 n_pts_with_derivs_rows = sum(len(order_indices) for order_indices in index) if isinstance(index, list) else 0
794 else:
795 der_map = deriv_map(n_bases, n_order)
796 index_2 = np.array([], dtype=np.int64)
797 n_pts_with_derivs_rows = sum(len(order_indices) for order_indices in index) if isinstance(index, list) else 0
799 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map)
800 der_indices_tr_pred, der_ind_order_pred = transform_der_indices(common_derivs, der_map) if common_derivs else ([], [])
801 n_pts_with_derivs_cols = n_deriv_types_pred * len([i for i in range(n_cols_func) if i < len(index_2)])
803 total_rows = n_rows_func + n_pts_with_derivs_rows
804 total_cols = n_cols_func + n_pts_with_derivs_cols
806 K = np.zeros((total_rows, total_cols))
807 base_shape = (n_rows_func, n_cols_func)
809 # Convert index lists to numpy arrays for numba
810 if isinstance(index, list) and len(index) > 0:
811 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index]
812 else:
813 index_arrays = []
815 # Block (0,0): Function-Function (K_ff)
816 content_full = phi_exp[0].reshape(base_shape)
817 K[:n_rows_func, :n_cols_func] = content_full * signs[0]
819 if not return_deriv:
820 # First Block-Column: Derivative-Function (K_df)
821 row_offset = n_rows_func
822 for i in range(n_deriv_types):
823 if not index_arrays:
824 break
825 row_indices = index_arrays[i]
826 n_pts_row = len(row_indices)
828 flat_idx = der_indices_tr[i]
829 content_full = phi_exp[flat_idx].reshape(base_shape)
831 # Use numba for efficient row extraction
832 extract_rows_and_assign(content_full, row_indices, K,
833 row_offset, 0, n_cols_func, signs[0])
834 row_offset += n_pts_row
835 return K
837 # --- return_deriv=True case ---
839 # First Block-Row: Function-Derivative (K_fd)
840 col_offset = n_cols_func
841 for j in range(n_deriv_types_pred):
842 col_indices = index_2
843 n_pts_col = len(col_indices)
845 flat_idx = der_indices_tr_pred[j]
846 content_full = phi_exp[flat_idx].reshape(base_shape)
848 # Use numba for efficient column extraction
849 extract_cols_and_assign(content_full, col_indices, K,
850 0, col_offset, n_rows_func, signs_predict[j + 1])
851 col_offset += n_pts_col
853 # First Block-Column: Derivative-Function (K_df)
854 row_offset = n_rows_func
855 for i in range(n_deriv_types):
856 if calc_cov:
857 row_indices = index_cov
858 flat_idx = der_indices_tr_pred[i]
859 else:
860 if not index_arrays:
861 break
862 row_indices = index_arrays[i]
863 flat_idx = der_indices_tr[i]
864 n_pts_row = len(row_indices)
866 content_full = phi_exp[flat_idx].reshape(base_shape)
868 # Use numba for efficient row extraction
869 extract_rows_and_assign(content_full, row_indices, K,
870 row_offset, 0, n_cols_func, signs[0])
871 row_offset += n_pts_row
873 # Inner Blocks: Derivative-Derivative (K_dd)
874 row_offset = n_rows_func
875 for i in range(n_deriv_types):
876 if calc_cov:
877 row_indices = index_cov
878 else:
879 if not index_arrays:
880 break
881 row_indices = index_arrays[i]
882 n_pts_row = len(row_indices)
884 col_offset = n_cols_func
885 for j in range(n_deriv_types_pred):
886 col_indices = index_2
887 n_pts_col = len(col_indices)
889 # Multiply the derivative indices to find the correct flat index
890 imdir1 = der_ind_order_pred[j]
891 imdir2 = der_ind_order_pred[i] if calc_cov else der_ind_order[i]
892 new_idx, new_ord = dh.mult_dir(
893 imdir1[0], imdir1[1], imdir2[0], imdir2[1])
894 flat_idx = der_map[new_ord][new_idx]
896 content_full = phi_exp[flat_idx].reshape(base_shape)
898 # Use numba for efficient submatrix extraction and assignment (replaces np.ix_)
899 extract_and_assign(content_full, row_indices, col_indices, K,
900 row_offset, col_offset, signs_predict[j + 1])
901 col_offset += n_pts_col
902 row_offset += n_pts_row
904 return K
907# =============================================================================
908# Utility functions
909# =============================================================================
911def determine_weights(diffs_by_dim, diffs_test, length_scales, kernel_func, sigma_n):
912 """
913 Vectorized version: compute interpolation weights for multiple test points at once.
915 Parameters
916 ----------
917 diffs_by_dim : list of ndarray
918 Pairwise differences between training points (by dimension).
919 diffs_test : list of ndarray
920 Pairwise differences between test points and training points (by dimension).
921 Shape: each array is (n_test, n_train) or similar batch dimension.
922 length_scales : array-like
923 Kernel hyperparameters.
924 kernel_func : callable
925 Kernel function.
926 sigma_n : float
927 Noise parameter (if needed).
929 Returns
930 -------
931 weights_matrix : ndarray of shape (n_test, n_train)
932 Interpolation weights for each test point.
933 """
934 # Compute K matrix (training covariance) - same for all test points
935 K = kernel_func(diffs_by_dim, length_scales).real
936 n_train = K.shape[0]
938 # Compute r vectors (test-train covariances) for all test points at once
939 r_all = kernel_func(diffs_test, length_scales).real
940 n_test = r_all.shape[0]
942 # Build augmented system matrix M (same for all test points)
943 M = np.zeros((n_train + 1, n_train + 1))
944 M[:n_train, :n_train] = K
945 M[:n_train, n_train] = 1
946 M[n_train, :n_train] = 1
947 M[n_train, n_train] = 0
949 # Build augmented RHS for all test points
950 r_augmented = np.zeros((n_test, n_train + 1))
951 r_augmented[:, :n_train] = r_all
952 r_augmented[:, n_train] = 1
954 # Solve for all test points at once
955 solution = np.linalg.solve(M, r_augmented.T)
957 # Extract weights (exclude Lagrange multiplier)
958 weights_matrix = solution[:n_train, :].T
960 return weights_matrix
963def to_tuple(item):
964 """Convert list to tuple recursively."""
965 if isinstance(item, list):
966 return tuple(to_tuple(x) for x in item)
967 return item
970def to_list(x):
971 """Convert tuple to list recursively."""
972 if isinstance(x, tuple):
973 return [to_list(i) for i in x]
974 return x
977def find_common_derivatives(all_indices):
978 """Find derivative indices common to all submodels."""
979 sets = [set(to_tuple(elem) for elem in idx_list) for idx_list in all_indices]
980 return sets[0].intersection(*sets[1:])