Coverage for jetgp/full_degp_sparse/degp_utils.py: 39%

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 

361# @profile 

362def rbf_kernel( 

363 phi, 

364 phi_exp, 

365 n_order, 

366 n_bases, 

367 der_indices, 

368 powers, 

369 index=None 

370): 

371 """ 

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

373 

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

375 replacing expensive np.ix_ operations. 

376 

377 Parameters 

378 ---------- 

379 phi : OTI array 

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

381 phi_exp : ndarray 

382 Expanded derivative array from phi.get_all_derivs(). 

383 n_order : int 

384 Maximum derivative order considered. 

385 n_bases : int 

386 Number of input dimensions. 

387 der_indices : list of lists 

388 Multi-index derivative structures for each derivative component. 

389 powers : list of int 

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

391 index : list of lists or None, optional 

392 If empty list, assumes uniform blocks. 

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

394 

395 Returns 

396 ------- 

397 K : ndarray 

398 Kernel matrix including function values and derivative terms. 

399 """ 

400 dh = coti.get_dHelp() 

401 

402 n_rows_func, n_cols_func = phi.shape 

403 n_deriv_types = len(der_indices) 

404 

405 der_map = deriv_map(n_bases, 2 * n_order) 

406 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

407 

408 # Pre-compute signs (avoid repeated exponentiation) 

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

410 

411 # ========================================================================= 

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

413 # ========================================================================= 

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

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

416 n_cols_func * (n_deriv_types + 1))) 

417 outer_loop_index = n_deriv_types + 1 

418 

419 for j in range(outer_loop_index): 

420 signj = signs[j] 

421 for i in range(n_deriv_types + 1): 

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

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

424 if j == 0 and i == 0: 

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

426 

427 return K 

428 

429 # ========================================================================= 

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

431 # ========================================================================= 

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

433 total_rows = n_rows_func + n_pts_with_derivs_rows 

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

435 total_cols = n_cols_func + n_pts_with_derivs_cols 

436 

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

438 base_shape = (n_rows_func, n_cols_func) 

439 

440 # Convert index lists to numpy arrays for numba 

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

442 

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

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

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

446 

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

448 row_offset = n_rows_func 

449 for i in range(n_deriv_types): 

450 flat_idx = der_indices_tr[i] 

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

452 row_indices = index_arrays[i] 

453 n_pts_this_order = len(row_indices) 

454 

455 # Use numba for efficient row extraction and assignment 

456 extract_rows_and_assign(content_full, row_indices, K, 

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

458 row_offset += n_pts_this_order 

459 

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

461 col_offset = n_cols_func 

462 for j in range(n_deriv_types): 

463 flat_idx = der_indices_tr[j] 

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

465 col_indices = index_arrays[j] 

466 n_pts_this_order = len(col_indices) 

467 

468 # Use numba for efficient column extraction and assignment 

469 extract_cols_and_assign(content_full, col_indices, K, 

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

471 col_offset += n_pts_this_order 

472 

473 # Inner Blocks: Derivative-Derivative (K_dd) 

474 row_offset = n_rows_func 

475 for i in range(n_deriv_types): 

476 col_offset = n_cols_func 

477 row_indices = index_arrays[i] 

478 n_pts_row = len(row_indices) 

479 

480 for j in range(n_deriv_types): 

481 col_indices = index_arrays[j] 

482 n_pts_col = len(col_indices) 

483 

484 imdir1 = der_ind_order[j] 

485 imdir2 = der_ind_order[i] 

486 new_idx, new_ord = dh.mult_dir( 

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

488 flat_idx = der_map[new_ord][new_idx] 

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

490 

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

492 extract_and_assign(content_full, row_indices, col_indices, K, 

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

494 

495 col_offset += n_pts_col 

496 row_offset += n_pts_row 

497 

498 return K 

499 

500 

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

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

503 fd_flat_indices, df_flat_indices, dd_flat_indices, 

504 idx_flat, idx_offsets, idx_sizes, 

505 signs, n_deriv_types, row_offsets, col_offsets): 

506 """ 

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

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

509 """ 

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

511 s0 = signs[0] 

512 for r in range(n_rows_func): 

513 for c in range(n_cols_func): 

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

515 

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

517 for j in range(n_deriv_types): 

518 fi = fd_flat_indices[j] 

519 sj = signs[j + 1] 

520 co = col_offsets[j] 

521 off_j = idx_offsets[j] 

522 sz_j = idx_sizes[j] 

523 for r in range(n_rows_func): 

524 for k in range(sz_j): 

525 ci = idx_flat[off_j + k] 

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

527 

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

529 for i in range(n_deriv_types): 

530 fi = df_flat_indices[i] 

531 ro = row_offsets[i] 

532 off_i = idx_offsets[i] 

533 sz_i = idx_sizes[i] 

534 for k in range(sz_i): 

535 ri = idx_flat[off_i + k] 

536 for c in range(n_cols_func): 

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

538 

539 # Inner Blocks: Derivative-Derivative (dd) 

540 for i in range(n_deriv_types): 

541 ro = row_offsets[i] 

542 off_i = idx_offsets[i] 

543 sz_i = idx_sizes[i] 

544 for j in range(n_deriv_types): 

545 fi = dd_flat_indices[i, j] 

546 sj = signs[j + 1] 

547 co = col_offsets[j] 

548 off_j = idx_offsets[j] 

549 sz_j = idx_sizes[j] 

550 for ki in range(sz_i): 

551 ri = idx_flat[off_i + ki] 

552 for kj in range(sz_j): 

553 ci = idx_flat[off_j + kj] 

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

555 

556 

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

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

559 fd_flat_indices, df_flat_indices, dd_flat_indices, 

560 idx_flat, idx_offsets, idx_sizes, 

561 signs, n_deriv_types, row_offsets, col_offsets): 

562 """ 

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

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

565 

566 This allows computing gradient contributions without materialising the 

567 full dK matrix for each hyperparameter dimension. 

568 """ 

569 # Zero out W_proj 

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

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

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

573 W_proj[d, r, c] = 0.0 

574 

575 s0 = signs[0] 

576 

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

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 

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

583 for j in range(n_deriv_types): 

584 fi = fd_flat_indices[j] 

585 sj = signs[j + 1] 

586 co = col_offsets[j] 

587 off_j = idx_offsets[j] 

588 sz_j = idx_sizes[j] 

589 for r in range(n_rows_func): 

590 for k in range(sz_j): 

591 ci = idx_flat[off_j + k] 

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

593 

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

595 for i in range(n_deriv_types): 

596 fi = df_flat_indices[i] 

597 ro = row_offsets[i] 

598 off_i = idx_offsets[i] 

599 sz_i = idx_sizes[i] 

600 for k in range(sz_i): 

601 ri = idx_flat[off_i + k] 

602 for c in range(n_cols_func): 

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

604 

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

606 for i in range(n_deriv_types): 

607 ro = row_offsets[i] 

608 off_i = idx_offsets[i] 

609 sz_i = idx_sizes[i] 

610 for j in range(n_deriv_types): 

611 fi = dd_flat_indices[i, j] 

612 sj = signs[j + 1] 

613 co = col_offsets[j] 

614 off_j = idx_offsets[j] 

615 sz_j = idx_sizes[j] 

616 for ki in range(sz_i): 

617 ri = idx_flat[off_i + ki] 

618 for kj in range(sz_j): 

619 ci = idx_flat[off_j + kj] 

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

621 

622 

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

624def _project_alpha_to_phi_space(alpha_v, alpha_proj, n_rows_func, n_cols_func, 

625 fd_flat_indices, df_flat_indices, dd_flat_indices, 

626 idx_flat, idx_offsets, idx_sizes, 

627 signs, n_deriv_types, row_offsets, col_offsets): 

628 """ 

629 Project the rank-1 matrix alpha*alpha^T into phi_exp-space. 

630 

631 Same structure as _project_W_to_phi_space but computes 

632 alpha_v[r]*alpha_v[c] on the fly instead of reading W[r,c]. 

633 """ 

634 # Zero out alpha_proj 

635 for d in range(alpha_proj.shape[0]): 

636 for r in range(alpha_proj.shape[1]): 

637 for c in range(alpha_proj.shape[2]): 

638 alpha_proj[d, r, c] = 0.0 

639 

640 s0 = signs[0] 

641 

642 # ff block: alpha_v[r] * alpha_v[c] 

643 for r in range(n_rows_func): 

644 ar = alpha_v[r] 

645 for c in range(n_cols_func): 

646 alpha_proj[0, r, c] += s0 * ar * alpha_v[c] 

647 

648 # fd blocks 

649 for j in range(n_deriv_types): 

650 fi = fd_flat_indices[j] 

651 sj = signs[j + 1] 

652 co = col_offsets[j] 

653 off_j = idx_offsets[j] 

654 sz_j = idx_sizes[j] 

655 for r in range(n_rows_func): 

656 ar = alpha_v[r] 

657 for k in range(sz_j): 

658 ci = idx_flat[off_j + k] 

659 alpha_proj[fi, r, ci] += sj * ar * alpha_v[co + k] 

660 

661 # df blocks 

662 for i in range(n_deriv_types): 

663 fi = df_flat_indices[i] 

664 ro = row_offsets[i] 

665 off_i = idx_offsets[i] 

666 sz_i = idx_sizes[i] 

667 for k in range(sz_i): 

668 ri = idx_flat[off_i + k] 

669 ak = alpha_v[ro + k] 

670 for c in range(n_cols_func): 

671 alpha_proj[fi, ri, c] += s0 * ak * alpha_v[c] 

672 

673 # dd blocks 

674 for i in range(n_deriv_types): 

675 ro = row_offsets[i] 

676 off_i = idx_offsets[i] 

677 sz_i = idx_sizes[i] 

678 for j in range(n_deriv_types): 

679 fi = dd_flat_indices[i, j] 

680 sj = signs[j + 1] 

681 co = col_offsets[j] 

682 off_j = idx_offsets[j] 

683 sz_j = idx_sizes[j] 

684 for ki in range(sz_i): 

685 ri = idx_flat[off_i + ki] 

686 aki = alpha_v[ro + ki] 

687 for kj in range(sz_j): 

688 ci = idx_flat[off_j + kj] 

689 alpha_proj[fi, ri, ci] += sj * aki * alpha_v[co + kj] 

690 

691 

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

693 """ 

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

695 reused across repeated calls with different phi_exp values. 

696 

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

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

699 """ 

700 dh = coti.get_dHelp() 

701 der_map = deriv_map(n_bases, 2 * n_order) 

702 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

703 

704 n_deriv_types = len(der_indices) 

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

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

707 

708 # Precompute sizes and offsets 

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

710 n_pts_with_derivs = int(index_sizes.sum()) 

711 

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

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

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

715 for i in range(1, n_deriv_types): 

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

717 

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

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

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

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

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

723 cumsum = 0 

724 for i in range(n_deriv_types): 

725 row_offsets[i] = cumsum # relative to n_rows_func 

726 col_offsets[i] = cumsum # relative to n_cols_func 

727 cumsum += index_sizes[i] 

728 

729 # Precompute mult_dir results for dd blocks 

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

731 for i in range(n_deriv_types): 

732 for j in range(n_deriv_types): 

733 imdir1 = der_ind_order[j] 

734 imdir2 = der_ind_order[i] 

735 new_idx, new_ord = dh.mult_dir( 

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

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

738 

739 # fd and df flat indices as arrays 

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

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

742 

743 return { 

744 'der_indices_tr': der_indices_tr, 

745 'signs': signs, 

746 'index_arrays': index_arrays, 

747 'index_sizes': index_sizes, 

748 'n_pts_with_derivs': n_pts_with_derivs, 

749 'dd_flat_indices': dd_flat_indices, 

750 'n_deriv_types': n_deriv_types, 

751 # Fused kernel data 

752 'idx_flat': idx_flat, 

753 'idx_offsets': idx_offsets, 

754 'row_offsets': row_offsets, 

755 'col_offsets': col_offsets, 

756 'fd_flat_indices': fd_flat_indices, 

757 'df_flat_indices': df_flat_indices, 

758 } 

759 

760 

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

762 """ 

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

764 

765 Parameters 

766 ---------- 

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

768 Pre-reshaped expanded derivative array. 

769 plan : dict 

770 Precomputed plan from precompute_kernel_plan(). 

771 out : ndarray, optional 

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

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

774 large matrices during optimization loops. 

775 

776 Returns 

777 ------- 

778 K : ndarray 

779 Full kernel matrix. 

780 """ 

781 n_rows_func = phi_exp_3d.shape[1] 

782 n_cols_func = phi_exp_3d.shape[2] 

783 total = n_rows_func + plan['n_pts_with_derivs'] 

784 if out is not None: 

785 K = out 

786 else: 

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

788 

789 # Use cached offsets if available, otherwise compute them 

790 if 'row_offsets_abs' in plan: 

791 row_off = plan['row_offsets_abs'] 

792 col_off = plan['col_offsets_abs'] 

793 else: 

794 row_off = plan['row_offsets'] + n_rows_func 

795 col_off = plan['col_offsets'] + n_cols_func 

796 

797 _assemble_kernel_numba( 

798 phi_exp_3d, K, n_rows_func, n_cols_func, 

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

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

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

802 ) 

803 

804 return K 

805 

806 

807def rbf_kernel_predictions( 

808 phi, 

809 phi_exp, 

810 n_order, 

811 n_bases, 

812 der_indices, 

813 powers, 

814 return_deriv, 

815 index=None, 

816 common_derivs=None, 

817 calc_cov=False, 

818 powers_predict=None 

819): 

820 """ 

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

822 

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

824 

825 Parameters 

826 ---------- 

827 phi : OTI array 

828 Base kernel matrix between test and training points. 

829 phi_exp : ndarray 

830 Expanded derivative array from phi.get_all_derivs(). 

831 n_order : int 

832 Maximum derivative order. 

833 n_bases : int