Coverage for jetgp/full_degp/degp_utils.py: 65%

1import numpy as np 

2import pyoti.core as coti 

3from line_profiler import profile 

4import numba 

5 

6 

7# ============================================================================= 

8# Numba-accelerated helper functions for efficient matrix slicing 

9# ============================================================================= 

10 

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. 

15 

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. 

24 

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 

37 

38 

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. 

43 

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. 

52 

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 

64 

65 

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. 

71 

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. 

80 

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 

94 

95 

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. 

102 

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 

126 

127 

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. 

133 

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 

156 

157 

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. 

163 

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 

185 

186 

187# ============================================================================= 

188# Difference computation functions 

189# ============================================================================= 

190 

191 

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

193 """ 

194 Compute pairwise differences between two input arrays X1 and X2 for each dimension, 

195 embedding hypercomplex units along each dimension for automatic differentiation. 

196 

197 For each dimension k, this function computes: 

198 diff_k[i, j] = X1[i, k] + e_{k+1} - X2[j, k] 

199 where e_{k+1} is a hypercomplex unit for the (k+1)-th dimension with order 2 * n_order. 

200 

201 Parameters 

202 ---------- 

203 X1 : array_like of shape (n1, d) 

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

205 X2 : array_like of shape (n2, d) 

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

207 n_order : int 

208 The base order used to construct hypercomplex units (e_{k+1}) with order 2 * n_order. 

209 oti_module : module 

210 The PyOTI static module (e.g., pyoti.static.onumm4n2). 

211 return_deriv : bool, optional 

212 If True, use 2*n_order for derivative predictions. 

213 

214 Returns 

215 ------- 

216 differences_by_dim : list of length d 

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

218 between corresponding dimensions of X1 and X2, augmented with hypercomplex units. 

219 """ 

220 # Keep numpy copies for fused path 

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

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

223 n1, d = X1_np.shape 

224 n2 = X2_np.shape[0] 

225 

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

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

228 

229 if _use_fused: 

230 # --- Fused path: numpy broadcast for real part, C-level OTI fill --- 

231 differences_by_dim = [] 

232 perturb2 = oti_module.zeros((n2, 1)) # X2 has no perturbation in DEGP 

233 

234 if n_order == 0: 

235 # No perturbation needed — just real differences 

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

237 for k in range(d): 

238 real_diffs = np.ascontiguousarray( 

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

240 ) 

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

242 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2) 

243 differences_by_dim.append(diffs_k) 

244 else: 

245 if return_deriv: 

246 hc_order = 2 * n_order 

247 else: 

248 hc_order = n_order 

249 

250 for k in range(d): 

251 # Real differences via numpy broadcasting (fast) 

252 real_diffs = np.ascontiguousarray( 

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

254 ) 

255 # Perturbation: e_{k+1} broadcast to all n1 points 

256 perturb1 = oti_module.zeros((n1, 1)) + oti_module.e(k + 1, order=hc_order) 

257 # Fused fill: out[i,j].real = real_diffs[i,j], out[i,j].im = perturb1[i] - 0 

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

259 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2) 

260 differences_by_dim.append(diffs_k) 

261 

262 return differences_by_dim 

263 

264 # --- Fallback: original Python loop path --- 

265 X1 = oti_module.array(X1_np) 

266 X2 = oti_module.array(X2_np) 

267 

268 differences_by_dim = [] 

269 

270 if n_order == 0: 

271 for k in range(d): 

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

273 for i in range(n1): 

274 diffs_k[i, :] = X1[i, k] - (oti_module.transpose(X2[:, k])) 

275 differences_by_dim.append(diffs_k) 

276 elif not return_deriv: 

277 for k in range(d): 

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

279 for i in range(n1): 

280 diffs_k[i, :] = ( 

281 X1[i, k] 

282 + oti_module.e(k + 1, order=n_order) 

283 - (X2[:, k].T) 

284 ) 

285 differences_by_dim.append(diffs_k) 

286 else: 

287 for k in range(d): 

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

289 for i in range(n1): 

290 diffs_k[i, :] = X1[i, k] - (X2[:, k].T) 

291 differences_by_dim.append( 

292 diffs_k + oti_module.e(k + 1, order=2 * n_order)) 

293 

294 return differences_by_dim 

295 

296 

297# ============================================================================= 

298# Derivative mapping utilities 

299# ============================================================================= 

300 

301def deriv_map(nbases, order): 

302 """ 

303 Create mapping from (order, index) to flattened index. 

304 

305 Parameters 

306 ---------- 

307 nbases : int 

308 Number of base dimensions. 

309 order : int 

310 Maximum derivative order. 

311 

312 Returns 

313 ------- 

314 map_deriv : list of lists 

315 Mapping where map_deriv[order][idx] gives the flattened index. 

316 """ 

317 k = 0 

318 map_deriv = [] 

319 for ordi in range(order + 1): 

320 ndir = coti.ndir_order(nbases, ordi) 

321 map_deriv_i = [0] * ndir 

322 for idx in range(ndir): 

323 map_deriv_i[idx] = k 

324 k += 1 

325 map_deriv.append(map_deriv_i) 

326 return map_deriv 

327 

328 

329def transform_der_indices(der_indices, der_map): 

330 """ 

331 Transform derivative indices to flattened format. 

332 

333 Parameters 

334 ---------- 

335 der_indices : list 

336 User-facing derivative specifications. 

337 der_map : list of lists 

338 Derivative mapping from deriv_map(). 

339 

340 Returns 

341 ------- 

342 deriv_ind_transf : list 

343 Flattened indices for each derivative. 

344 deriv_ind_order : list 

345 (index, order) tuples for each derivative. 

346 """ 

347 deriv_ind_transf = [] 

348 deriv_ind_order = [] 

349 for deriv in der_indices: 

350 imdir = coti.imdir(deriv) 

351 idx, order = imdir 

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

353 deriv_ind_order.append(imdir) 

354 return deriv_ind_transf, deriv_ind_order 

355 

356 

357# ============================================================================= 

358# RBF Kernel Assembly Functions (Optimized with Numba) 

359# ============================================================================= 

360 

361def rbf_kernel( 

362 phi, 

363 phi_exp, 

364 n_order, 

365 n_bases, 

366 der_indices, 

367 powers, 

368 index=None 

369): 

370 """ 

371 Compute the derivative-enhanced RBF kernel matrix (optimized version). 

372 

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

374 replacing expensive np.ix_ operations. 

375 

376 Parameters 

377 ---------- 

378 phi : OTI array 

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

380 phi_exp : ndarray 

381 Expanded derivative array from phi.get_all_derivs(). 

382 n_order : int 

383 Maximum derivative order considered. 

384 n_bases : int 

385 Number of input dimensions. 

386 der_indices : list of lists 

387 Multi-index derivative structures for each derivative component. 

388 powers : list of int 

389 Powers of (-1) applied to each term. 

390 index : list of lists or None, optional 

391 If empty list, assumes uniform blocks. 

392 If provided, specifies which training point indices have each derivative type. 

393 

394 Returns 

395 ------- 

396 K : ndarray 

397 Kernel matrix including function values and derivative terms. 

398 """ 

399 dh = coti.get_dHelp() 

400 

401 n_rows_func, n_cols_func = phi.shape 

402 n_deriv_types = len(der_indices) 

403 

404 der_map = deriv_map(n_bases, 2 * n_order) 

405 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

406 

407 # Pre-compute signs (avoid repeated exponentiation) 

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

409 

410 # ========================================================================= 

411 # CASE 1: Uniform blocks (original behavior) - index is None or empty 

412 # ========================================================================= 

413 if index is None or len(index) == 0: 

414 K = np.zeros((n_rows_func * (n_deriv_types + 1), 

415 n_cols_func * (n_deriv_types + 1))) 

416 outer_loop_index = n_deriv_types + 1 

417 

418 for j in range(outer_loop_index): 

419 signj = signs[j] 

420 for i in range(n_deriv_types + 1): 

421 Klocal = K[i * n_rows_func: (i + 1) * n_rows_func, 

422 j * n_cols_func: (j + 1) * n_cols_func] 

423 if j == 0 and i == 0: 

424 Klocal[:, :] = phi_exp[0] * signj 

425 

426 return K 

427 

428 # ========================================================================= 

429 # CASE 2: Non-contiguous indices - index is provided 

430 # ========================================================================= 

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

432 total_rows = n_rows_func + n_pts_with_derivs_rows 

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

434 total_cols = n_cols_func + n_pts_with_derivs_cols 

435 

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

437 base_shape = (n_rows_func, n_cols_func) 

438 

439 # Convert index lists to numpy arrays for numba 

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

441 

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

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

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

445 

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

447 row_offset = n_rows_func 

448 for i in range(n_deriv_types): 

449 flat_idx = der_indices_tr[i] 

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

451 row_indices = index_arrays[i] 

452 n_pts_this_order = len(row_indices) 

453 

454 # Use numba for efficient row extraction and assignment 

455 extract_rows_and_assign(content_full, row_indices, K, 

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

457 row_offset += n_pts_this_order 

458 

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

460 col_offset = n_cols_func 

461 for j in range(n_deriv_types): 

462 flat_idx = der_indices_tr[j] 

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

464 col_indices = index_arrays[j] 

465 n_pts_this_order = len(col_indices) 

466 

467 # Use numba for efficient column extraction and assignment 

468 extract_cols_and_assign(content_full, col_indices, K, 

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

470 col_offset += n_pts_this_order 

471 

472 # Inner Blocks: Derivative-Derivative (K_dd) 

473 row_offset = n_rows_func 

474 for i in range(n_deriv_types): 

475 col_offset = n_cols_func 

476 row_indices = index_arrays[i] 

477 n_pts_row = len(row_indices) 

478 

479 for j in range(n_deriv_types): 

480 col_indices = index_arrays[j] 

481 n_pts_col = len(col_indices) 

482 

483 imdir1 = der_ind_order[j] 

484 imdir2 = der_ind_order[i] 

485 new_idx, new_ord = dh.mult_dir( 

486 imdir1[0], imdir1[1], imdir2[0], imdir2[1]) 

487 flat_idx = der_map[new_ord][new_idx] 

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

489 

490 # Use numba for direct extraction and assignment (replaces np.ix_) 

491 extract_and_assign(content_full, row_indices, col_indices, K, 

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

493 

494 col_offset += n_pts_col 

495 row_offset += n_pts_row 

496 

497 return K 

498 

499 

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

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

502 fd_flat_indices, df_flat_indices, dd_flat_indices, 

503 idx_flat, idx_offsets, idx_sizes, 

504 signs, n_deriv_types, row_offsets, col_offsets): 

505 """ 

506 Fused numba kernel that assembles the entire K matrix in a single call. 

507 Handles ff, fd, df, and dd blocks without Python-level loop overhead. 

508 """ 

509 # Block (0,0): Function-Function 

510 s0 = signs[0] 

511 for r in range(n_rows_func): 

512 for c in range(n_cols_func): 

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

514 

515 # First Block-Row: Function-Derivative (fd) 

516 for j in range(n_deriv_types): 

517 fi = fd_flat_indices[j] 

518 sj = signs[j + 1] 

519 co = col_offsets[j] 

520 off_j = idx_offsets[j] 

521 sz_j = idx_sizes[j] 

522 for r in range(n_rows_func): 

523 for k in range(sz_j): 

524 ci = idx_flat[off_j + k] 

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

526 

527 # First Block-Column: Derivative-Function (df) 

528 for i in range(n_deriv_types): 

529 fi = df_flat_indices[i] 

530 ro = row_offsets[i] 

531 off_i = idx_offsets[i] 

532 sz_i = idx_sizes[i] 

533 for k in range(sz_i): 

534 ri = idx_flat[off_i + k] 

535 for c in range(n_cols_func): 

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

537 

538 # Inner Blocks: Derivative-Derivative (dd) 

539 for i in range(n_deriv_types): 

540 ro = row_offsets[i] 

541 off_i = idx_offsets[i] 

542 sz_i = idx_sizes[i] 

543 for j in range(n_deriv_types): 

544 fi = dd_flat_indices[i, j] 

545 sj = signs[j + 1] 

546 co = col_offsets[j] 

547 off_j = idx_offsets[j] 

548 sz_j = idx_sizes[j] 

549 for ki in range(sz_i): 

550 ri = idx_flat[off_i + ki] 

551 for kj in range(sz_j): 

552 ci = idx_flat[off_j + kj] 

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

554 

555 

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

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

558 fd_flat_indices, df_flat_indices, dd_flat_indices, 

559 idx_flat, idx_offsets, idx_sizes, 

560 signs, n_deriv_types, row_offsets, col_offsets): 

561 """ 

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

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

564 

565 This allows computing gradient contributions without materialising the 

566 full dK matrix for each hyperparameter dimension. 

567 """ 

568 # Zero out W_proj 

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

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

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

572 W_proj[d, r, c] = 0.0 

573 

574 s0 = signs[0] 

575 

576 # ff block: K[r, c] = phi_exp[0, r, c] * s0 

577 for r in range(n_rows_func): 

578 for c in range(n_cols_func): 

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

580 

581 # fd blocks: K[r, co+k] = phi_exp[fi, r, idx[k]] * sj 

582 for j in range(n_deriv_types): 

583 fi = fd_flat_indices[j] 

584 sj = signs[j + 1] 

585 co = col_offsets[j] 

586 off_j = idx_offsets[j] 

587 sz_j = idx_sizes[j] 

588 for r in range(n_rows_func): 

589 for k in range(sz_j): 

590 ci = idx_flat[off_j + k] 

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

592 

593 # df blocks: K[ro+k, c] = phi_exp[fi, idx[k], c] * s0 

594 for i in range(n_deriv_types): 

595 fi = df_flat_indices[i] 

596 ro = row_offsets[i] 

597 off_i = idx_offsets[i] 

598 sz_i = idx_sizes[i] 

599 for k in range(sz_i): 

600 ri = idx_flat[off_i + k] 

601 for c in range(n_cols_func): 

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

603 

604 # dd blocks: K[ro+ki, co+kj] = phi_exp[dd_fi[i,j], idx_i[ki], idx_j[kj]] * sj 

605 for i in range(n_deriv_types): 

606 ro = row_offsets[i] 

607 off_i = idx_offsets[i] 

608 sz_i = idx_sizes[i] 

609 for j in range(n_deriv_types): 

610 fi = dd_flat_indices[i, j] 

611 sj = signs[j + 1] 

612 co = col_offsets[j] 

613 off_j = idx_offsets[j] 

614 sz_j = idx_sizes[j] 

615 for ki in range(sz_i): 

616 ri = idx_flat[off_i + ki] 

617 for kj in range(sz_j): 

618 ci = idx_flat[off_j + kj] 

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

620 

621 

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

623 """ 

624 Precompute all structural information needed by rbf_kernel so it can be 

625 reused across repeated calls with different phi_exp values. 

626 

627 Returns a dict containing flat indices, signs, index arrays, precomputed 

628 offsets/sizes, and mult_dir results for the dd block. 

629 """ 

630 dh = coti.get_dHelp() 

631 der_map = deriv_map(n_bases, 2 * n_order) 

632 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

633 

634 n_deriv_types = len(der_indices) 

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

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

637 

638 # Precompute sizes and offsets 

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

640 n_pts_with_derivs = int(index_sizes.sum()) 

641 

642 # Pack all index arrays into a single flat array with offsets 

643 idx_flat = np.concatenate(index_arrays) if n_deriv_types > 0 else np.array([], dtype=np.int64) 

644 idx_offsets = np.zeros(n_deriv_types, dtype=np.int64) 

645 for i in range(1, n_deriv_types): 

646 idx_offsets[i] = idx_offsets[i - 1] + index_sizes[i - 1] 

647 

648 # Precompute row/col offsets in K for each deriv type 

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

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

651 # Note: n_rows_func == n_cols_func for training kernel, but we store 

652 # offsets relative to n_rows_func which is added at call time 

653 cumsum = 0 

654 for i in range(n_deriv_types): 

655 row_offsets[i] = cumsum # relative to n_rows_func 

656 col_offsets[i] = cumsum # relative to n_cols_func 

657 cumsum += index_sizes[i] 

658 

659 # Precompute mult_dir results for dd blocks 

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

661 for i in range(n_deriv_types): 

662 for j in range(n_deriv_types): 

663 imdir1 = der_ind_order[j] 

664 imdir2 = der_ind_order[i] 

665 new_idx, new_ord = dh.mult_dir( 

666 imdir1[0], imdir1[1], imdir2[0], imdir2[1]) 

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

668 

669 # fd and df flat indices as arrays 

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

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

672 

673 return { 

674 'der_indices_tr': der_indices_tr, 

675 'signs': signs, 

676 'index_arrays': index_arrays, 

677 'index_sizes': index_sizes, 

678 'n_pts_with_derivs': n_pts_with_derivs, 

679 'dd_flat_indices': dd_flat_indices, 

680 'n_deriv_types': n_deriv_types, 

681 # Fused kernel data 

682 'idx_flat': idx_flat, 

683 'idx_offsets': idx_offsets, 

684 'row_offsets': row_offsets, 

685 'col_offsets': col_offsets, 

686 'fd_flat_indices': fd_flat_indices, 

687 'df_flat_indices': df_flat_indices, 

688 } 

689 

690 

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

692 """ 

693 Fast kernel assembly using a precomputed plan and fused numba kernel. 

694 

695 Parameters 

696 ---------- 

697 phi_exp_3d : ndarray of shape (n_derivs, n_rows_func, n_cols_func) 

698 Pre-reshaped expanded derivative array. 

699 plan : dict 

700 Precomputed plan from precompute_kernel_plan(). 

701 out : ndarray, optional 

702 Pre-allocated output array of shape (total, total). If None, a new 

703 array is allocated. Reusing a buffer avoids repeated allocation of 

704 large matrices during optimization loops. 

705 

706 Returns 

707 ------- 

708 K : ndarray 

709 Full kernel matrix. 

710 """ 

711 n_rows_func = phi_exp_3d.shape[1] 

712 n_cols_func = phi_exp_3d.shape[2] 

713 total = n_rows_func + plan['n_pts_with_derivs'] 

714 if out is not None: 

715 K = out 

716 else: 

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

718 

719 # Use cached offsets if available, otherwise compute them 

720 if 'row_offsets_abs' in plan: 

721 row_off = plan['row_offsets_abs'] 

722 col_off = plan['col_offsets_abs'] 

723 else: 

724 row_off = plan['row_offsets'] + n_rows_func 

725 col_off = plan['col_offsets'] + n_cols_func 

726 

727 _assemble_kernel_numba( 

728 phi_exp_3d, K, n_rows_func, n_cols_func, 

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

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

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

732 ) 

733 

734 return K 

735 

736 

737def rbf_kernel_predictions( 

738 phi, 

739 phi_exp, 

740 n_order, 

741 n_bases, 

742 der_indices, 

743 powers, 

744 return_deriv, 

745 index=None, 

746 common_derivs=None, 

747 calc_cov=False, 

748 powers_predict=None 

749): 

750 """ 

751 Constructs the RBF kernel matrix for predictions with derivative entries. 

752 

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

754 

755 Parameters 

756 ---------- 

757 phi : OTI array 

758 Base kernel matrix between test and training points. 

759 phi_exp : ndarray 

760 Expanded derivative array from phi.get_all_derivs(). 

761 n_order : int 

762 Maximum derivative order. 

763 n_bases : int 

764 Number of input dimensions. 

765 der_indices : list 

766 Derivative specifications for training data. 

767 powers : list of int 

768 Sign powers for each derivative type. 

769 return_deriv : bool 

770 If True, predict derivatives at ALL test points. 

771 index : list of lists or None 

772 Training point indices for each derivative type. 

773 common_derivs : list 

774 Common derivative indices to predict. 

775 calc_cov : bool 

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

777 powers_predict : list of int, optional 

778 Sign powers for prediction derivatives. 

779 

780 Returns 

781 ------- 

782 K : ndarray 

783 Prediction kernel matrix. 

784 """ 

785 # Early return for covariance-only case 

786 if calc_cov and not return_deriv: 

787 return phi.real 

788 

789 dh = coti.get_dHelp() 

790 

791 n_rows_func, n_cols_func = phi.shape 

792 n_deriv_types = len(der_indices) 

793 n_deriv_types_pred = len(common_derivs) if common_derivs else 0 

794 

795 # Pre-compute signs 

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

797 if powers_predict is not None: 

798 signs_predict = np.array( 

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

800 else: 

801 signs_predict = signs 

802 

803 # Determine derivative map and index structures 

804 if return_deriv: 

805 der_map = deriv_map(n_bases, 2 * n_order) 

806 index_2 = np.arange(n_cols_func, dtype=np.int64) 

807 if calc_cov: 

808 index_cov = np.arange(n_cols_func, dtype=np.int64) 

809 n_deriv_types = n_deriv_types_pred 

810 n_pts_with_derivs_rows = n_deriv_types * n_cols_func 

811 else: 

812 n_pts_with_derivs_rows = sum(len(order_indices) 

813 for order_indices in index) if index else 0 

814 else: 

815 der_map = deriv_map(n_bases, n_order) 

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

817 n_pts_with_derivs_rows = sum(len(order_indices) 

818 for order_indices in index) if index else 0 

819 

820 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

821 der_indices_tr_pred, der_ind_order_pred = transform_der_indices( 

822 common_derivs, der_map) if common_derivs else ([], []) 

823 n_pts_with_derivs_cols = n_deriv_types_pred * len(index_2) 

824 

825 total_rows = n_rows_func + n_pts_with_derivs_rows 

826 total_cols = n_cols_func + n_pts_with_derivs_cols 

827 

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

829 base_shape = (n_rows_func, n_cols_func) 

830 

831 # Convert index lists to numpy arrays for numba 

832 if index is not None and len(index) > 0 and isinstance(index[0], (list, np.ndarray)): 

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

834 else: 

835 index_arrays = [] 

836 

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

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

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

840 

841 if not return_deriv: 

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

843 row_offset = n_rows_func