Coverage for jetgp/full_ddegp/wddegp_utils.py: 61%

1import numpy as np 

2import pyoti.sparse as oti 

3import pyoti.core as coti 

4from line_profiler import profile 

5import numba 

6 

7 

8# ============================================================================= 

9# Numba-accelerated helper functions for efficient matrix slicing 

10# ============================================================================= 

11 

12@numba.jit(nopython=True, cache=True) 

13def extract_rows(content_full, row_indices, n_cols): 

14 """ 

15 Extract rows from content_full at specified indices. 

16  

17 Parameters 

18 ---------- 

19 content_full : ndarray of shape (n_rows_full, n_cols) 

20 Source matrix. 

21 row_indices : ndarray of int64 

22 Row indices to extract. 

23 n_cols : int 

24 Number of columns. 

25  

26 Returns 

27 ------- 

28 result : ndarray of shape (len(row_indices), n_cols) 

29 Extracted rows. 

30 """ 

31 n_rows = len(row_indices) 

32 result = np.empty((n_rows, n_cols)) 

33 for i in range(n_rows): 

34 ri = row_indices[i] 

35 for j in range(n_cols): 

36 result[i, j] = content_full[ri, j] 

37 return result 

38 

39 

40@numba.jit(nopython=True, cache=True) 

41def extract_cols(content_full, col_indices, n_rows): 

42 """ 

43 Extract columns from content_full at specified indices. 

44  

45 Parameters 

46 ---------- 

47 content_full : ndarray of shape (n_rows, n_cols_full) 

48 Source matrix. 

49 col_indices : ndarray of int64 

50 Column indices to extract. 

51 n_rows : int 

52 Number of rows. 

53  

54 Returns 

55 ------- 

56 result : ndarray of shape (n_rows, len(col_indices)) 

57 Extracted columns. 

58 """ 

59 n_cols = len(col_indices) 

60 result = np.empty((n_rows, n_cols)) 

61 for i in range(n_rows): 

62 for j in range(n_cols): 

63 result[i, j] = content_full[i, col_indices[j]] 

64 return result 

65 

66 

67@numba.jit(nopython=True, cache=True) 

68def extract_submatrix(content_full, row_indices, col_indices): 

69 """ 

70 Extract submatrix from content_full at specified row and column indices. 

71 Replaces the expensive np.ix_ operation. 

72  

73 Parameters 

74 ---------- 

75 content_full : ndarray of shape (n_rows_full, n_cols_full) 

76 Source matrix. 

77 row_indices : ndarray of int64 

78 Row indices to extract. 

79 col_indices : ndarray of int64 

80 Column indices to extract. 

81  

82 Returns 

83 ------- 

84 result : ndarray of shape (len(row_indices), len(col_indices)) 

85 Extracted submatrix. 

86 """ 

87 n_rows = len(row_indices) 

88 n_cols = len(col_indices) 

89 result = np.empty((n_rows, n_cols)) 

90 for i in range(n_rows): 

91 ri = row_indices[i] 

92 for j in range(n_cols): 

93 result[i, j] = content_full[ri, col_indices[j]] 

94 return result 

95 

96 

97@numba.jit(nopython=True, cache=True, parallel=False) 

98def extract_and_assign(content_full, row_indices, col_indices, K, 

99 row_start, col_start, sign): 

100 """ 

101 Extract submatrix and assign directly to K with sign multiplication. 

102 Combines extraction and assignment in one pass for better performance. 

103  

104 Parameters 

105 ---------- 

106 content_full : ndarray of shape (n_rows_full, n_cols_full) 

107 Source matrix. 

108 row_indices : ndarray of int64 

109 Row indices to extract. 

110 col_indices : ndarray of int64 

111 Column indices to extract. 

112 K : ndarray 

113 Target matrix to fill. 

114 row_start : int 

115 Starting row index in K. 

116 col_start : int 

117 Starting column index in K. 

118 sign : float 

119 Sign multiplier (+1.0 or -1.0). 

120 """ 

121 n_rows = len(row_indices) 

122 n_cols = len(col_indices) 

123 for i in range(n_rows): 

124 ri = row_indices[i] 

125 for j in range(n_cols): 

126 K[row_start + i, col_start + j] = content_full[ri, col_indices[j]] * sign 

127 

128 

129@numba.jit(nopython=True, cache=True) 

130def extract_rows_and_assign(content_full, row_indices, K, 

131 row_start, col_start, n_cols, sign): 

132 """ 

133 Extract rows and assign directly to K with sign multiplication. 

134  

135 Parameters 

136 ---------- 

137 content_full : ndarray of shape (n_rows_full, n_cols) 

138 Source matrix. 

139 row_indices : ndarray of int64 

140 Row indices to extract. 

141 K : ndarray 

142 Target matrix to fill. 

143 row_start : int 

144 Starting row index in K. 

145 col_start : int 

146 Starting column index in K. 

147 n_cols : int 

148 Number of columns to copy. 

149 sign : float 

150 Sign multiplier (+1.0 or -1.0). 

151 """ 

152 n_rows = len(row_indices) 

153 for i in range(n_rows): 

154 ri = row_indices[i] 

155 for j in range(n_cols): 

156 K[row_start + i, col_start + j] = content_full[ri, j] * sign 

157 

158 

159@numba.jit(nopython=True, cache=True) 

160def extract_cols_and_assign(content_full, col_indices, K, 

161 row_start, col_start, n_rows, sign): 

162 """ 

163 Extract columns and assign directly to K with sign multiplication. 

164  

165 Parameters 

166 ---------- 

167 content_full : ndarray of shape (n_rows, n_cols_full) 

168 Source matrix. 

169 col_indices : ndarray of int64 

170 Column indices to extract. 

171 K : ndarray 

172 Target matrix to fill. 

173 row_start : int 

174 Starting row index in K. 

175 col_start : int 

176 Starting column index in K. 

177 n_rows : int 

178 Number of rows to copy. 

179 sign : float 

180 Sign multiplier (+1.0 or -1.0). 

181 """ 

182 n_cols = len(col_indices) 

183 for i in range(n_rows): 

184 for j in range(n_cols): 

185 K[row_start + i, col_start + j] = content_full[i, col_indices[j]] * sign 

186 

187 

188# ============================================================================= 

189# Difference computation functions 

190# ============================================================================= 

191 

192def differences_by_dim_func(X1, X2, rays, n_order,oti_module, return_deriv=True): 

193 """ 

194 Compute dimension-wise pairwise differences between X1 and X2, 

195 including hypercomplex perturbations in the directions specified by `rays`. 

196  

197 This optimized version pre-calculates the perturbation and uses a single 

198 efficient loop for subtraction, avoiding broadcasting issues with OTI arrays. 

199  

200 Parameters 

201 ---------- 

202 X1 : ndarray of shape (n1, d) 

203 First set of input points with n1 samples in d dimensions. 

204 X2 : ndarray of shape (n2, d) 

205 Second set of input points with n2 samples in d dimensions. 

206 rays : ndarray of shape (d, n_rays) 

207 Directional vectors for derivative computation. 

208 n_order : int 

209 The base order used to construct hypercomplex units. 

210 When return_deriv=True, uses order 2*n_order. 

211 When return_deriv=False, uses order n_order. 

212 return_deriv : bool, optional (default=True) 

213 If True, use order 2*n_order for hypercomplex units (needed for  

214 derivative-derivative blocks in training kernel). 

215 If False, use order n_order (sufficient for prediction without  

216 derivative outputs). 

217  

218 Returns 

219 ------- 

220 differences_by_dim : list of length d 

221 A list where each element is an array of shape (n1, n2), containing  

222 the differences between corresponding dimensions of X1 and X2,  

223 augmented with directional hypercomplex perturbations. 

224 """ 

225 X1_np = np.asarray(X1, dtype=np.float64) 

226 X2_np = np.asarray(X2, dtype=np.float64) 

227 n1, d = X1_np.shape 

228 n2 = X2_np.shape[0] 

229 n_rays = rays.shape[1] 

230 

231 # Check if the fused C-level function is available 

232 _use_fused = hasattr(oti_module.zeros((1, 1)), 'fused_from_real_with_perturbations') 

233 

234 differences_by_dim = [] 

235 

236 # Case 1: n_order == 0 (no hypercomplex perturbation) 

237 if n_order == 0: 

238 if _use_fused: 

239 perturb1 = oti_module.zeros((n1, 1)) 

240 perturb2 = oti_module.zeros((n2, 1)) 

241 for k in range(d): 

242 real_diffs = np.ascontiguousarray( 

243 X1_np[:, k:k+1] - X2_np[:, k:k+1].T, dtype=np.float64 

244 ) 

245 diffs_k = oti_module.empty((n1, n2)) 

246 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2) 

247 differences_by_dim.append(diffs_k) 

248 else: 

249 X1 = oti_module.array(X1_np) 

250 X2 = oti_module.array(X2_np) 

251 for k in range(d): 

252 diffs_k = oti_module.zeros((n1, n2)) 

253 for i in range(n1): 

254 diffs_k[i, :] = X1[i, k] - X2[:, k].T 

255 differences_by_dim.append(diffs_k) 

256 return differences_by_dim 

257 

258 # Convert to OTI arrays for non-fused paths 

259 X1 = oti_module.array(X1_np) 

260 X2 = oti_module.array(X2_np) 

261 

262 # Determine the order for hypercomplex units based on return_deriv 

263 if return_deriv: 

264 hc_order = 2 * n_order 

265 else: 

266 hc_order = n_order 

267 

268 # Pre-calculate the perturbation vector using directional rays 

269 e_bases = [oti_module.e(i + 1, order=hc_order) for i in range(n_rays)] 

270 perts = np.dot(rays, e_bases) 

271 

272 # Case 2: return_deriv=False (prediction without derivative outputs) 

273 if not return_deriv: 

274 for k in range(d): 

275 # Add the pre-calculated perturbation for the current dimension to all points in X1 

276 X1_k_tagged = X1[:, k] + perts[k] 

277 X2_k = X2[:, k] 

278 

279 # Pre-allocate the result matrix for this dimension 

280 diffs_k = oti_module.zeros((n1, n2)) 

281 

282 # Use an efficient single loop for subtraction 

283 for i in range(n1): 

284 diffs_k[i, :] = X1_k_tagged[i, 0] - X2_k[:, 0].T 

285 

286 differences_by_dim.append(diffs_k) 

287 

288 # Case 3: return_deriv=True (training kernel with derivative-derivative blocks) 

289 else: 

290 for k in range(d): 

291 X2_k = X2[:, k] 

292 

293 # Pre-allocate the result matrix for this dimension 

294 diffs_k = oti_module.zeros((n1, n2)) 

295 

296 # Compute differences without perturbation first 

297 for i in range(n1): 

298 diffs_k[i, :] = X1[i, k] - X2_k[:, 0].T 

299 

300 # Add perturbation to the entire matrix (more efficient) 

301 differences_by_dim.append(diffs_k + perts[k]) 

302 

303 return differences_by_dim 

304 

305 

306# ============================================================================= 

307# Derivative mapping utilities 

308# ============================================================================= 

309 

310def deriv_map(nbases, order): 

311 """ 

312 Creates a mapping from (order, index_within_order) to a single 

313 flattened index for all derivative components. 

314 """ 

315 k = 0 

316 map_deriv = [] 

317 for ordi in range(order + 1): 

318 ndir = coti.ndir_order(nbases, ordi) 

319 map_deriv_i = [0] * ndir 

320 for idx in range(ndir): 

321 map_deriv_i[idx] = k 

322 k += 1 

323 map_deriv.append(map_deriv_i) 

324 return map_deriv 

325 

326 

327def transform_der_indices(der_indices, der_map): 

328 """ 

329 Transforms a list of user-facing derivative specifications into the 

330 internal (order, index) format and the final flattened index. 

331 """ 

332 deriv_ind_transf = [] 

333 deriv_ind_order = [] 

334 for deriv in der_indices: 

335 imdir = coti.imdir(deriv) 

336 idx, order = imdir 

337 deriv_ind_transf.append(der_map[order][idx]) 

338 deriv_ind_order.append(imdir) 

339 return deriv_ind_transf, deriv_ind_order 

340 

341 

342# ============================================================================= 

343# RBF Kernel Assembly Functions (Optimized with Numba) 

344# ============================================================================= 

345 

346def rbf_kernel( 

347 phi, 

348 phi_exp, 

349 n_order, 

350 n_bases, 

351 der_indices, 

352 powers, 

353 index=-1 

354): 

355 """ 

356 Assembles the full DD-GP covariance matrix using an efficient, pre-computed 

357 derivative array and block-wise matrix filling. 

358  

359 This version uses Numba-accelerated functions for efficient matrix slicing, 

360 replacing expensive np.ix_ operations. 

361  

362 Parameters 

363 ---------- 

364 phi : OTI array 

365 Base kernel matrix from kernel_func(differences, length_scales). 

366 phi_exp : ndarray 

367 Expanded derivative array from phi.get_all_derivs(). 

368 n_order : int 

369 Maximum derivative order considered. 

370 n_bases : int 

371 Total number of bases (function value + derivative terms). 

372 der_indices : list of lists 

373 Multi-index derivative structures for each derivative component. 

374 powers : list of int 

375 Powers of (-1) applied to each term (for symmetry or sign conventions). 

376 index : list of lists 

377 Specifies which training point indices have each derivative type. 

378  

379 Returns 

380 ------- 

381 K : ndarray 

382 Full kernel matrix with function values and derivative blocks. 

383 """ 

384 dh = coti.get_dHelp() 

385 

386 # Create maps to translate derivative specifications to flat indices 

387 der_map = deriv_map(n_bases, 2 * n_order) 

388 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

389 

390 # Determine Block Sizes and Pre-allocate Matrix 

391 n_rows_func, n_cols_func = phi.shape 

392 n_deriv_types = len(der_indices) 

393 n_pts_with_derivs_cols = sum(len(order_indices) for order_indices in index) 

394 n_pts_with_derivs_rows = sum(len(order_indices) for order_indices in index) 

395 total_rows = n_rows_func + n_pts_with_derivs_rows 

396 total_cols = n_cols_func + n_pts_with_derivs_cols 

397 

398 K = np.zeros((total_rows, total_cols)) 

399 base_shape = (n_rows_func, n_cols_func) 

400 

401 # Pre-compute signs (avoid repeated exponentiation) 

402 signs = np.array([(-1.0) ** p for p in powers], dtype=np.float64) 

403 

404 # Convert index lists to numpy arrays for numba 

405 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index] 

406 

407 # Block (0,0): Function-Function (K_ff) 

408 content_full = phi_exp[0].reshape(base_shape) 

409 K[:n_rows_func, :n_cols_func] = content_full * signs[0] 

410 

411 # First Block-Column: Derivative-Function (K_df) 

412 row_offset = n_rows_func 

413 for i in range(n_deriv_types): 

414 flat_idx = der_indices_tr[i] 

415 content_full = phi_exp[flat_idx].reshape(base_shape) 

416 

417 current_indices = index_arrays[i] 

418 n_pts_this_order = len(current_indices) 

419 

420 # Use numba for efficient row extraction and assignment 

421 extract_rows_and_assign(content_full, current_indices, K, 

422 row_offset, 0, n_cols_func, signs[0]) 

423 row_offset += n_pts_this_order 

424 

425 # First Block-Row: Function-Derivative (K_fd) 

426 col_offset = n_cols_func 

427 for j in range(n_deriv_types): 

428 flat_idx = der_indices_tr[j] 

429 content_full = phi_exp[flat_idx].reshape(base_shape) 

430 

431 current_indices = index_arrays[j] 

432 n_pts_this_order = len(current_indices) 

433 

434 # Use numba for efficient column extraction and assignment 

435 extract_cols_and_assign(content_full, current_indices, K, 

436 0, col_offset, n_rows_func, signs[j + 1]) 

437 col_offset += n_pts_this_order 

438 

439 # Inner Blocks: Derivative-Derivative (K_dd) 

440 row_offset = n_rows_func 

441 for i in range(n_deriv_types): 

442 col_offset = n_cols_func 

443 

444 row_indices = index_arrays[i] 

445 n_pts_row = len(row_indices) 

446 

447 for j in range(n_deriv_types): 

448 col_indices = index_arrays[j] 

449 n_pts_col = len(col_indices) 

450 

451 # Multiply derivative indices to find correct flat index 

452 imdir1 = der_ind_order[j] 

453 imdir2 = der_ind_order[i] 

454 new_idx, new_ord = dh.mult_dir(imdir1[0], imdir1[1], imdir2[0], imdir2[1]) 

455 flat_idx = der_map[new_ord][new_idx] 

456 content_full = phi_exp[flat_idx].reshape(base_shape) 

457 

458 # Use numba for efficient submatrix extraction and assignment (replaces np.ix_) 

459 extract_and_assign(content_full, row_indices, col_indices, K, 

460 row_offset, col_offset, signs[j + 1]) 

461 

462 col_offset += n_pts_col 

463 

464 row_offset += n_pts_row 

465 

466 return K 

467 

468 

469@numba.jit(nopython=True, cache=True) 

470def _assemble_kernel_numba(phi_exp_3d, K, n_rows_func, n_cols_func, 

471 fd_flat_indices, df_flat_indices, dd_flat_indices, 

472 idx_flat, idx_offsets, idx_sizes, 

473 signs, n_deriv_types, row_offsets, col_offsets): 

474 """Fused numba kernel for entire K matrix assembly.""" 

475 s0 = signs[0] 

476 for r in range(n_rows_func): 

477 for c in range(n_cols_func): 

478 K[r, c] = phi_exp_3d[0, r, c] * s0 

479 for j in range(n_deriv_types): 

480 fi = fd_flat_indices[j] 

481 sj = signs[j + 1] 

482 co = col_offsets[j] 

483 off_j = idx_offsets[j] 

484 sz_j = idx_sizes[j] 

485 for r in range(n_rows_func): 

486 for k in range(sz_j): 

487 ci = idx_flat[off_j + k] 

488 K[r, co + k] = phi_exp_3d[fi, r, ci] * sj 

489 for i in range(n_deriv_types): 

490 fi = df_flat_indices[i] 

491 ro = row_offsets[i] 

492 off_i = idx_offsets[i] 

493 sz_i = idx_sizes[i] 

494 for k in range(sz_i): 

495 ri = idx_flat[off_i + k] 

496 for c in range(n_cols_func): 

497 K[ro + k, c] = phi_exp_3d[fi, ri, c] * s0 

498 for i in range(n_deriv_types): 

499 ro = row_offsets[i] 

500 off_i = idx_offsets[i] 

501 sz_i = idx_sizes[i] 

502 for j in range(n_deriv_types): 

503 fi = dd_flat_indices[i, j] 

504 sj = signs[j + 1] 

505 co = col_offsets[j] 

506 off_j = idx_offsets[j] 

507 sz_j = idx_sizes[j] 

508 for ki in range(sz_i): 

509 ri = idx_flat[off_i + ki] 

510 for kj in range(sz_j): 

511 ci = idx_flat[off_j + kj] 

512 K[ro + ki, co + kj] = phi_exp_3d[fi, ri, ci] * sj 

513 

514 

515@numba.jit(nopython=True, cache=True) 

516def _project_W_to_phi_space(W, W_proj, n_rows_func, n_cols_func, 

517 fd_flat_indices, df_flat_indices, dd_flat_indices, 

518 idx_flat, idx_offsets, idx_sizes, 

519 signs, n_deriv_types, row_offsets, col_offsets): 

520 """ 

521 Reverse of _assemble_kernel_numba: project W from K-space back into 

522 phi_exp-space so that vdot(W, assemble(dphi_exp)) == vdot(W_proj, dphi_exp). 

523 """ 

524 for d in range(W_proj.shape[0]): 

525 for r in range(W_proj.shape[1]): 

526 for c in range(W_proj.shape[2]): 

527 W_proj[d, r, c] = 0.0 

528 s0 = signs[0] 

529 for r in range(n_rows_func): 

530 for c in range(n_cols_func): 

531 W_proj[0, r, c] += s0 * W[r, c] 

532 for j in range(n_deriv_types): 

533 fi = fd_flat_indices[j] 

534 sj = signs[j + 1] 

535 co = col_offsets[j] 

536 off_j = idx_offsets[j] 

537 sz_j = idx_sizes[j] 

538 for r in range(n_rows_func): 

539 for k in range(sz_j): 

540 ci = idx_flat[off_j + k] 

541 W_proj[fi, r, ci] += sj * W[r, co + k] 

542 for i in range(n_deriv_types): 

543 fi = df_flat_indices[i] 

544 ro = row_offsets[i] 

545 off_i = idx_offsets[i] 

546 sz_i = idx_sizes[i] 

547 for k in range(sz_i): 

548 ri = idx_flat[off_i + k] 

549 for c in range(n_cols_func): 

550 W_proj[fi, ri, c] += s0 * W[ro + k, c] 

551 for i in range(n_deriv_types): 

552 ro = row_offsets[i] 

553 off_i = idx_offsets[i] 

554 sz_i = idx_sizes[i] 

555 for j in range(n_deriv_types): 

556 fi = dd_flat_indices[i, j] 

557 sj = signs[j + 1] 

558 co = col_offsets[j] 

559 off_j = idx_offsets[j] 

560 sz_j = idx_sizes[j] 

561 for ki in range(sz_i): 

562 ri = idx_flat[off_i + ki] 

563 for kj in range(sz_j): 

564 ci = idx_flat[off_j + kj] 

565 W_proj[fi, ri, ci] += sj * W[ro + ki, co + kj] 

566 

567 

568@numba.jit(nopython=True, cache=True) 

569def _project_W_to_phi_space_accum(W, W_proj, n_rows_func, n_cols_func, 

570 fd_flat_indices, df_flat_indices, dd_flat_indices, 

571 idx_flat, idx_offsets, idx_sizes, 

572 signs, n_deriv_types, row_offsets, col_offsets): 

573 """ 

574 Like _project_W_to_phi_space but accumulates into W_proj without zeroing. 

575 Caller must zero W_proj before the first call. 

576 """ 

577 s0 = signs[0] 

578 for r in range(n_rows_func): 

579 for c in range(n_cols_func): 

580 W_proj[0, r, c] += s0 * W[r, c] 

581 for j in range(n_deriv_types): 

582 fi = fd_flat_indices[j] 

583 sj = signs[j + 1] 

584 co = col_offsets[j] 

585 off_j = idx_offsets[j] 

586 sz_j = idx_sizes[j] 

587 for r in range(n_rows_func): 

588 for k in range(sz_j): 

589 ci = idx_flat[off_j + k] 

590 W_proj[fi, r, ci] += sj * W[r, co + k] 

591 for i in range(n_deriv_types): 

592 fi = df_flat_indices[i] 

593 ro = row_offsets[i] 

594 off_i = idx_offsets[i] 

595 sz_i = idx_sizes[i] 

596 for k in range(sz_i): 

597 ri = idx_flat[off_i + k] 

598 for c in range(n_cols_func): 

599 W_proj[fi, ri, c] += s0 * W[ro + k, c] 

600 for i in range(n_deriv_types): 

601 ro = row_offsets[i] 

602 off_i = idx_offsets[i] 

603 sz_i = idx_sizes[i] 

604 for j in range(n_deriv_types): 

605 fi = dd_flat_indices[i, j] 

606 sj = signs[j + 1] 

607 co = col_offsets[j] 

608 off_j = idx_offsets[j] 

609 sz_j = idx_sizes[j] 

610 for ki in range(sz_i): 

611 ri = idx_flat[off_i + ki] 

612 for kj in range(sz_j): 

613 ci = idx_flat[off_j + kj] 

614 W_proj[fi, ri, ci] += sj * W[ro + ki, co + kj] 

615 

616 

617def precompute_kernel_plan(n_order, n_bases, der_indices, powers, index): 

618 """Precompute structural info for rbf_kernel_fast.""" 

619 dh = coti.get_dHelp() 

620 der_map = deriv_map(n_bases, 2 * n_order) 

621 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

622 

623 n_deriv_types = len(der_indices) 

624 signs = np.array([(-1.0) ** p for p in powers], dtype=np.float64) 

625 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index] 

626 

627 index_sizes = np.array([len(idx) for idx in index_arrays], dtype=np.int64) 

628 n_pts_with_derivs = int(index_sizes.sum()) 

629 

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] 

634 

635 row_offsets = np.zeros(n_deriv_types, dtype=np.int64) 

636 col_offsets = np.zeros(n_deriv_types, dtype=np.int64) 

637 cumsum = 0 

638 for i in range(n_deriv_types): 

639 row_offsets[i] = cumsum 

640 col_offsets[i] = cumsum 

641 cumsum += index_sizes[i] 

642 

643 dd_flat_indices = np.empty((n_deriv_types, n_deriv_types), dtype=np.int64) 

644 for i in range(n_deriv_types): 

645 for j in range(n_deriv_types): 

646 imdir1 = der_ind_order[j] 

647 imdir2 = der_ind_order[i] 

648 new_idx, new_ord = dh.mult_dir(imdir1[0], imdir1[1], imdir2[0], imdir2[1]) 

649 dd_flat_indices[i, j] = der_map[new_ord][new_idx] 

650 

651 fd_flat_indices = np.array(der_indices_tr, dtype=np.int64) 

652 df_flat_indices = np.array(der_indices_tr, dtype=np.int64) 

653 

654 return { 

655 'der_indices_tr': der_indices_tr, 

656 'signs': signs, 

657 'index_arrays': index_arrays, 

658 'index_sizes': index_sizes, 

659 'n_pts_with_derivs': n_pts_with_derivs, 

660 'dd_flat_indices': dd_flat_indices, 

661 'n_deriv_types': n_deriv_types, 

662 'idx_flat': idx_flat, 

663 'idx_offsets': idx_offsets, 

664 'row_offsets': row_offsets, 

665 'col_offsets': col_offsets, 

666 'fd_flat_indices': fd_flat_indices, 

667 'df_flat_indices': df_flat_indices, 

668 } 

669 

670 

671def rbf_kernel_fast(phi_exp_3d, plan, out=None): 

672 """Fast kernel assembly using precomputed plan and fused numba kernel.""" 

673 n_rows_func = phi_exp_3d.shape[1] 

674 n_cols_func = phi_exp_3d.shape[2] 

675 total = n_rows_func + plan['n_pts_with_derivs'] 

676 if out is not None: 

677 K = out 

678 else: 

679 K = np.empty((total, total)) 

680 

681 if 'row_offsets_abs' in plan: 

682 row_off = plan['row_offsets_abs'] 

683 col_off = plan['col_offsets_abs'] 

684 else: 

685 row_off = plan['row_offsets'] + n_rows_func 

686 col_off = plan['col_offsets'] + n_cols_func 

687 

688 _assemble_kernel_numba( 

689 phi_exp_3d, K, n_rows_func, n_cols_func, 

690 plan['fd_flat_indices'], plan['df_flat_indices'], plan['dd_flat_indices'], 

691 plan['idx_flat'], plan['idx_offsets'], plan['index_sizes'], 

692 plan['signs'], plan['n_deriv_types'], row_off, col_off, 

693 ) 

694 return K 

695 

696 

697def rbf_kernel_predictions( 

698 phi, 

699 phi_exp, 

700 n_order, 

701 n_bases, 

702 der_indices, 

703 powers, 

704 return_deriv, 

705 index=-1, 

706 common_derivs=None, 

707 calc_cov=False, 

708 powers_predict=None 

709): 

710 """ 

711 Constructs the RBF kernel matrix for predictions with directional derivative entries. 

712  

713 This version uses Numba-accelerated functions for efficient matrix slicing. 

714 

715 Parameters 

716 ---------- 

717 phi : OTI array 

718 Base kernel matrix between test and training points. 

719 phi_exp : ndarray 

720 Expanded derivative array from phi.get_all_derivs(). 

721 n_order : int 

722 Maximum derivative order. 

723 n_bases : int 

724 Number of input dimensions. 

725 der_indices : list 

726 Derivative specifications. 

727 powers : list of int 

728 Sign powers for each derivative type. 

729 return_deriv : bool 

730 If True, predict derivatives at ALL test points. 

731 index : list of lists 

732 Training point indices for each derivative type. 

733 common_derivs : list 

734 Common derivative indices to predict. 

735 calc_cov : bool 

736 If True, computing covariance (use all indices for rows). 

737 powers_predict : list of int, optional 

738 Sign powers for prediction derivatives. 

739 

740 Returns 

741 ------- 

742 K : ndarray 

743 Prediction kernel matrix. 

744 """ 

745 if calc_cov and not return_deriv: 

746 return phi.real 

747 

748 dh = coti.get_dHelp() 

749 

750 n_rows_func, n_cols_func = phi.shape 

751 n_deriv_types = len(der_indices) 

752 n_deriv_types_pred = len(common_derivs) if common_derivs else 0 

753 

754 # Pre-compute signs 

755 signs = np.array([(-1.0) ** p for p in powers], dtype=np.float64) 

756 if powers_predict is not None: 

757 signs_predict = np.array([(-1.0) ** p for p in powers_predict], dtype=np.float64) 

758 else: 

759 signs_predict = signs 

760 

761 if return_deriv: 

762 der_map = deriv_map(n_bases, 2 * n_order) 

763 index_2 = np.arange(phi_exp.shape[-1], dtype=np.int64) 

764 if calc_cov: 

765 index_cov = np.arange(phi_exp.shape[-1], dtype=np.int64) 

766 n_deriv_types = n_deriv_types_pred 

767 n_pts_with_derivs_rows = n_deriv_types * len([i for i in range(n_cols_func) if i < len(index_2)]) 

768 else: 

769 n_pts_with_derivs_rows = sum(len(order_indices) for order_indices in index) 

770 else: 

771 der_map = deriv_map(n_bases, n_order) 

772 index_2 = np.array([], dtype=np.int64) 

773 n_pts_with_derivs_rows = sum(len(order_indices) for order_indices in index) 

774 

775 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

776 der_indices_tr_pred, der_ind_order_pred = transform_der_indices(common_derivs, der_map) if common_derivs else ([], []) 

777 n_pts_with_derivs_cols = n_deriv_types_pred * len([i for i in range(n_cols_func) if i < len(index_2)]) 

778 

779 total_rows = n_rows_func + n_pts_with_derivs_rows 

780 total_cols = n_cols_func + n_pts_with_derivs_cols 

781 

782 K = np.zeros((total_rows, total_cols)) 

783 base_shape = (n_rows_func, n_cols_func) 

784 

785 # Convert index lists to numpy arrays for numba 

786 if index != -1 and isinstance(index, list) and len(index) > 0: 

787 index_arrays = [np.asarray(idx, dtype=np.int64) for idx in index] 

788 else: 

789 index_arrays = [] 

790 

791 # Block (0,0): Function-Function (K_ff) 

792 content_full = phi_exp[0].reshape(base_shape) 

793 K[:n_rows_func, :n_cols_func] = content_full * signs[0] 

794 

795 if not return_deriv: 

796 # First Block-Column: Derivative-Function (K_df) 

797 row_offset = n_rows_func 

798 for i in range(n_deriv_types): 

799 if calc_cov: 

800 row_indices = index_cov 

801 else: 

802 if not index_arrays: