Coverage for jetgp/full_ddegp/ddegp_utils.py: 63%

1import numpy as np 

2import numba 

3import pyoti.core as coti 

4from line_profiler import profile 

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 

191def differences_by_dim_func(X1, X2, rays, n_order, oti_module, return_deriv=True, index=-1): 

192 """ 

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

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

195  

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

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

198  

199 Parameters 

200 ---------- 

201 X1 : ndarray of shape (n1, d) 

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

203 X2 : ndarray of shape (n2, d) 

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

205 rays : ndarray of shape (d, n_rays) 

206 Directional vectors for derivative computation. 

207 n_order : int 

208 The base order used to construct hypercomplex units. 

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

210 When return_deriv=False, uses order n_order. 

211 oti_module : module 

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

213 return_deriv : bool, optional (default=True) 

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

215 derivative-derivative blocks in training kernel). 

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

217 derivative outputs). 

218 index : int, optional 

219 Currently unused. Reserved for future enhancements. 

220  

221 Returns 

222 ------- 

223 differences_by_dim : list of length d 

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

225 the differences between corresponding dimensions of X1 and X2,  

226 augmented with directional hypercomplex perturbations. 

227  

228 Notes 

229 ----- 

230 - The function leverages hypercomplex arithmetic from the pyOTI library. 

231 - The directional perturbation is computed as: perts = rays @ e_bases 

232 where e_bases are the hypercomplex units for each ray direction. 

233 - This routine is typically used in the construction of directional  

234 derivative kernels for Gaussian processes. 

235  

236 Example 

237 ------- 

238 >>> X1 = np.array([[1.0, 2.0], [3.0, 4.0]]) 

239 >>> X2 = np.array([[1.5, 2.5], [3.5, 4.5]]) 

240 >>> rays = np.eye(2) # Standard basis directions 

241 >>> n_order = 1 

242 >>> oti_module = get_oti_module(2, 1) # dim=2, n_order=1 

243 >>> diffs = differences_by_dim_func(X1, X2, rays, n_order, oti_module) 

244 >>> len(diffs) 

245 2 

246 >>> diffs[0].shape 

247 (2, 2) 

248 """ 

249 # Keep numpy copies for fused path 

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

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

252 n1, d = X1_np.shape 

253 n2 = X2_np.shape[0] 

254 n_rays = rays.shape[1] 

255 

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

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

258 

259 differences_by_dim = [] 

260 

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

262 if n_order == 0: 

263 if _use_fused: 

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

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

266 for k in range(d): 

267 real_diffs = np.ascontiguousarray( 

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

269 ) 

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

271 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2) 

272 differences_by_dim.append(diffs_k) 

273 else: 

274 X1_oti = oti_module.array(X1_np) 

275 X2_oti = oti_module.array(X2_np) 

276 for k in range(d): 

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

278 for i in range(n1): 

279 diffs_k[i, :] = X1_oti[i, k] - oti_module.transpose(X2_oti[:, k]) 

280 differences_by_dim.append(diffs_k) 

281 return differences_by_dim 

282 

283 # Determine the order for hypercomplex units based on return_deriv 

284 if return_deriv: 

285 hc_order = 2 * n_order 

286 else: 

287 hc_order = n_order 

288 

289 # Pre-calculate the perturbation vector using directional rays 

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

291 perts = np.dot(rays, e_bases) 

292 

293 if _use_fused: 

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

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

296 

297 for k in range(d): 

298 # Real differences via numpy broadcasting (fast) 

299 real_diffs = np.ascontiguousarray( 

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

301 ) 

302 # Perturbation: perts[k] (ray-projected) broadcast to all n1 points 

303 perturb1 = oti_module.zeros((n1, 1)) + perts[k] 

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

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

306 diffs_k.fused_from_real_with_perturbations(real_diffs, perturb1, perturb2) 

307 differences_by_dim.append(diffs_k) 

308 

309 return differences_by_dim 

310 

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

312 X1 = oti_module.array(X1_np) 

313 X2 = oti_module.array(X2_np) 

314 

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

316 if not return_deriv: 

317 for k in range(d): 

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

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

320 X2_k = X2[:, k] 

321 

322 # Pre-allocate the result matrix for this dimension 

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

324 

325 # Use an efficient single loop for subtraction 

326 for i in range(n1): 

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

328 

329 differences_by_dim.append(diffs_k) 

330 

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

332 else: 

333 for k in range(d): 

334 X2_k = X2[:, k] 

335 

336 # Pre-allocate the result matrix for this dimension 

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

338 

339 # Compute differences without perturbation first 

340 for i in range(n1): 

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

342 

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

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

345 

346 return differences_by_dim 

347 

348# ============================================================================= 

349# Derivative mapping utilities 

350# ============================================================================= 

351 

352def deriv_map(nbases, order): 

353 """ 

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

355 flattened index for all derivative components. 

356  

357 Parameters 

358 ---------- 

359 nbases : int 

360 Number of base dimensions. 

361 order : int 

362 Maximum derivative order. 

363  

364 Returns 

365 ------- 

366 map_deriv : list of lists 

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

368 """ 

369 k = 0 

370 map_deriv = [] 

371 for ordi in range(order + 1): 

372 ndir = coti.ndir_order(nbases, ordi) 

373 map_deriv_i = [0] * ndir 

374 for idx in range(ndir): 

375 map_deriv_i[idx] = k 

376 k += 1 

377 map_deriv.append(map_deriv_i) 

378 return map_deriv 

379 

380 

381def transform_der_indices(der_indices, der_map): 

382 """ 

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

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

385  

386 Parameters 

387 ---------- 

388 der_indices : list 

389 User-facing derivative specifications. 

390 der_map : list of lists 

391 Derivative mapping from deriv_map(). 

392  

393 Returns 

394 ------- 

395 deriv_ind_transf : list 

396 Flattened indices for each derivative. 

397 deriv_ind_order : list 

398 (index, order) tuples for each derivative. 

399 """ 

400 deriv_ind_transf = [] 

401 deriv_ind_order = [] 

402 for deriv in der_indices: 

403 imdir = coti.imdir(deriv) 

404 idx, order = imdir 

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

406 deriv_ind_order.append(imdir) 

407 return deriv_ind_transf, deriv_ind_order 

408 

409 

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

411# RBF Kernel Assembly Functions (Optimized with Numba) 

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

413 

414@profile 

415def rbf_kernel( 

416 phi, 

417 phi_exp, 

418 n_order, 

419 n_bases, 

420 der_indices, 

421 powers, 

422 index=-1 

423): 

424 """ 

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

426 derivative array and block-wise matrix filling. 

427  

428 Supports both uniform blocks (all derivatives at all points) and non-contiguous 

429 indices (different derivatives at different subsets of points). 

430  

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

432 replacing expensive np.ix_ operations. 

433  

434 Parameters 

435 ---------- 

436 phi : OTI array 

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

438 phi_exp : ndarray 

439 Expanded derivative array from phi.get_all_derivs(). 

440 n_order : int 

441 Maximum derivative order considered. 

442 n_bases : int 

443 Number of input dimensions (rays). 

444 der_indices : list of lists 

445 Multi-index derivative structures for each derivative component. 

446 powers : list of int 

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

448 index : list of lists or int, optional (default=-1) 

449 If empty list, assumes all derivative types apply to all training points. 

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

451 allowing non-contiguous index support and variable block sizes. 

452  

453 Returns 

454 ------- 

455 K : ndarray 

456 Full kernel matrix with function values and derivative blocks. 

457 """ 

458 # --- 1. Initial Setup and Efficient Derivative Extraction --- 

459 dh = coti.get_dHelp() 

460 

461 # Create maps to translate derivative specifications to flat indices 

462 der_map = deriv_map(n_bases, 2 * n_order) 

463 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

464 

465 # --- 2. Determine Block Sizes and Pre-allocate Matrix --- 

466 n_rows_func, n_cols_func = phi.shape 

467 n_deriv_types = len(der_indices) 

468 

469 # Pre-compute signs (avoid repeated exponentiation) 

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

471 

472 # Convert index lists to numpy arrays for numba (if provided) 

473 if isinstance(index, list) and len(index) > 0: 

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

475 else: 

476 index_arrays = [] 

477 

478 n_pts_with_derivs_cols = sum(len(idx) for idx in index_arrays) if index_arrays else 0 

479 n_pts_with_derivs_rows = n_pts_with_derivs_cols 

480 total_rows = n_rows_func + n_pts_with_derivs_rows 

481 total_cols = n_cols_func + n_pts_with_derivs_cols 

482 

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

484 base_shape = (n_rows_func, n_cols_func) 

485 

486 # --- 3. Fill the Matrix Block by Block --- 

487 

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

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

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

491 

492 if not index_arrays: 

493 # No derivative indices provided, return early 

494 return K 

495 

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

497 row_offset = n_rows_func 

498 for i in range(n_deriv_types): 

499 flat_idx = der_indices_tr[i] 

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

501 

502 row_indices = index_arrays[i] 

503 n_pts_this_order = len(row_indices) 

504 

505 # Use numba for efficient row extraction and assignment 

506 extract_rows_and_assign(content_full, row_indices, K, 

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

508 row_offset += n_pts_this_order 

509 

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

511 col_offset = n_cols_func 

512 for j in range(n_deriv_types): 

513 flat_idx = der_indices_tr[j] 

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

515 

516 col_indices = index_arrays[j] 

517 n_pts_this_order = len(col_indices) 

518 

519 # Use numba for efficient column extraction and assignment 

520 extract_cols_and_assign(content_full, col_indices, K, 

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

522 col_offset += n_pts_this_order 

523 

524 # Inner Blocks: Derivative-Derivative (K_dd) 

525 row_offset = n_rows_func 

526 for i in range(n_deriv_types): 

527 col_offset = n_cols_func 

528 

529 row_indices = index_arrays[i] 

530 n_pts_row = len(row_indices) 

531 

532 for j in range(n_deriv_types): 

533 col_indices = index_arrays[j] 

534 n_pts_col = len(col_indices) 

535 

536 # Multiply derivative indices to find correct flat index 

537 imdir1 = der_ind_order[j] 

538 imdir2 = der_ind_order[i] 

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

540 flat_idx = der_map[new_ord][new_idx] 

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

542 

543 # Use numba for efficient submatrix extraction and assignment 

544 # This replaces the expensive np.ix_ operation 

545 extract_and_assign(content_full, row_indices, col_indices, K, 

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

547 

548 col_offset += n_pts_col 

549 

550 row_offset += n_pts_row 

551 

552 return K 

553 

554 

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

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

557 fd_flat_indices, df_flat_indices, dd_flat_indices, 

558 idx_flat, idx_offsets, idx_sizes, 

559 signs, n_deriv_types, row_offsets, col_offsets): 

560 """ 

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

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

563 """ 

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

565 s0 = signs[0] 

566 for r in range(n_rows_func): 

567 for c in range(n_cols_func): 

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

569 

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

571 for j in range(n_deriv_types): 

572 fi = fd_flat_indices[j] 

573 sj = signs[j + 1] 

574 co = col_offsets[j] 

575 off_j = idx_offsets[j] 

576 sz_j = idx_sizes[j] 

577 for r in range(n_rows_func): 

578 for k in range(sz_j): 

579 ci = idx_flat[off_j + k] 

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

581 

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

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 K[ro + k, c] = phi_exp_3d[fi, ri, c] * s0 

592 

593 # Inner Blocks: Derivative-Derivative (dd) 

594 for i in range(n_deriv_types): 

595 ro = row_offsets[i] 

596 off_i = idx_offsets[i] 

597 sz_i = idx_sizes[i] 

598 for j in range(n_deriv_types): 

599 fi = dd_flat_indices[i, j] 

600 sj = signs[j + 1] 

601 co = col_offsets[j] 

602 off_j = idx_offsets[j] 

603 sz_j = idx_sizes[j] 

604 for ki in range(sz_i): 

605 ri = idx_flat[off_i + ki] 

606 for kj in range(sz_j): 

607 ci = idx_flat[off_j + kj] 

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

609 

610 

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

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

613 fd_flat_indices, df_flat_indices, dd_flat_indices, 

614 idx_flat, idx_offsets, idx_sizes, 

615 signs, n_deriv_types, row_offsets, col_offsets): 

616 """ 

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

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

619 """ 

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

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

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

623 W_proj[d, r, c] = 0.0 

624 s0 = signs[0] 

625 for r in range(n_rows_func): 

626 for c in range(n_cols_func): 

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

628 for j in range(n_deriv_types): 

629 fi = fd_flat_indices[j] 

630 sj = signs[j + 1] 

631 co = col_offsets[j] 

632 off_j = idx_offsets[j] 

633 sz_j = idx_sizes[j] 

634 for r in range(n_rows_func): 

635 for k in range(sz_j): 

636 ci = idx_flat[off_j + k] 

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

638 for i in range(n_deriv_types): 

639 fi = df_flat_indices[i] 

640 ro = row_offsets[i] 

641 off_i = idx_offsets[i] 

642 sz_i = idx_sizes[i] 

643 for k in range(sz_i): 

644 ri = idx_flat[off_i + k] 

645 for c in range(n_cols_func): 

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

647 for i in range(n_deriv_types): 

648 ro = row_offsets[i] 

649 off_i = idx_offsets[i] 

650 sz_i = idx_sizes[i] 

651 for j in range(n_deriv_types): 

652 fi = dd_flat_indices[i, j] 

653 sj = signs[j + 1] 

654 co = col_offsets[j] 

655 off_j = idx_offsets[j] 

656 sz_j = idx_sizes[j] 

657 for ki in range(sz_i): 

658 ri = idx_flat[off_i + ki] 

659 for kj in range(sz_j): 

660 ci = idx_flat[off_j + kj] 

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

662 

663 

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

665 """ 

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

667 reused across repeated calls with different phi_exp values. 

668 

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

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

671 """ 

672 dh = coti.get_dHelp() 

673 der_map = deriv_map(n_bases, 2 * n_order) 

674 der_indices_tr, der_ind_order = transform_der_indices(der_indices, der_map) 

675 

676 n_deriv_types = len(der_indices) 

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

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

679 

680 # Precompute sizes and offsets 

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

682 n_pts_with_derivs = int(index_sizes.sum()) 

683 

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

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

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

687 for i in range(1, n_deriv_types): 

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

689 

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

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

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

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

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

695 cumsum = 0 

696 for i in range(n_deriv_types): 

697 row_offsets[i] = cumsum # relative to n_rows_func 

698 col_offsets[i] = cumsum # relative to n_cols_func 

699 cumsum += index_sizes[i] 

700 

701 # Precompute mult_dir results for dd blocks 

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

703 for i in range(n_deriv_types): 

704 for j in range(n_deriv_types): 

705 imdir1 = der_ind_order[j] 

706 imdir2 = der_ind_order[i] 

707 new_idx, new_ord = dh.mult_dir( 

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

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

710 

711 # fd and df flat indices as arrays 

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

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

714 

715 return { 

716 'der_indices_tr': der_indices_tr, 

717 'signs': signs, 

718 'index_arrays': index_arrays, 

719 'index_sizes': index_sizes, 

720 'n_pts_with_derivs': n_pts_with_derivs, 

721 'dd_flat_indices': dd_flat_indices, 

722 'n_deriv_types': n_deriv_types, 

723 # Fused kernel data 

724 'idx_flat': idx_flat, 

725 'idx_offsets': idx_offsets, 

726 'row_offsets': row_offsets, 

727 'col_offsets': col_offsets, 

728 'fd_flat_indices': fd_flat_indices, 

729 'df_flat_indices': df_flat_indices, 

730 } 

731 

732 

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

734 """ 

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

736 

737 Parameters 

738 ---------- 

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

740 Pre-reshaped expanded derivative array. 

741 plan : dict 

742 Precomputed plan from precompute_kernel_plan(). 

743 out : ndarray, optional 

744 Pre-allocated output array. If None, a new array is allocated. 

745 

746 Returns 

747 ------- 

748 K : ndarray 

749 Full kernel matrix. 

750 """ 

751 n_rows_func = phi_exp_3d.shape[1] 

752 n_cols_func = phi_exp_3d.shape[2] 

753 total = n_rows_func + plan['n_pts_with_derivs'] 

754 if out is not None: 

755 K = out 

756 else: 

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

758 

759 if 'row_offsets_abs' in plan: 

760 row_off = plan['row_offsets_abs'] 

761 col_off = plan['col_offsets_abs'] 

762 else: 

763 row_off = plan['row_offsets'] + n_rows_func 

764 col_off = plan['col_offsets'] + n_cols_func 

765 

766 _assemble_kernel_numba( 

767 phi_exp_3d, K, n_rows_func, n_cols_func, 

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

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

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

771 ) 

772 

773 return K 

774 

775 

776def rbf_kernel_predictions( 

777 phi, 

778 phi_exp, 

779 n_order, 

780 n_bases, 

781 der_indices, 

782 powers, 

783 return_deriv, 

784 index=-1, 

785 common_derivs=None, 

786 calc_cov=False, 

787 powers_predict=None 

788): 

789 """ 

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

791  

792 This handles the asymmetric case where: 

793 - Rows: Test points (predictions) 

794 - Columns: Training points (with derivative structure from index) 

795  

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

797 

798 Parameters 

799 ---------- 

800 phi : OTI array 

801 Base kernel matrix between test and training points. 

802 phi_exp : ndarray 

803 Expanded derivative array from phi.get_all_derivs(). 

804 n_order : int 

805 Maximum derivative order. 

806 n_bases : int 

807 Number of input dimensions (rays). 

808 der_indices : list 

809 Derivative specifications for training data. 

810 powers : list of int 

811 Sign powers for each derivative type. 

812 return_deriv : bool 

813 If True, predict derivatives at ALL test points. 

814 index : list of lists or int, optional (default=-1) 

815 Training point indices for each derivative type. 

816 common_derivs : list, optional 

817 Common derivative indices to predict (intersection of training and requested). 

818 calc_cov : bool, optional (default=False) 

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

820 powers_predict : list of int, optional 

821 Sign powers for prediction derivatives. 

822 

823 Returns 

824 ------- 

825 K : ndarray 

826 Prediction kernel matrix. 

827 """ 

828 # --- 1. Initial Setup --- 

829 if calc_cov and not return_deriv: 

830 return phi.real 

831 

832 dh = coti.get_dHelp() 

833 

834 # Pre-compute signs 

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

836 if powers_predict is not None: 

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

838 else: 

839 signs_predict = signs 

840 

841 # --- 2. Determine Block Sizes and Pre-allocate Matrix --- 

842 n_rows_func, n_cols_func = phi.shape 

843 n_deriv_types = len(der_indices) 

844 n_deriv_types_pred = len(common_derivs) if common_derivs else 0 

845 

846 # Convert index to numpy arrays 

847 if isinstance(index, list) and len(index) > 0 and isinstance(index[0], (list, np.ndarray)): 

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

849 else: 

850 index_arrays = [] 

851 

852 if return_deriv: 

853 der_map = deriv_map(n_bases, 2 * n_order) 

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

855 if calc_cov: 

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

857 n_deriv_types = n_deriv_types_pred 

858 n_pts_with_derivs_rows = n_deriv_types * n_cols_func 

859 else: 

860 n_pts_with_derivs_rows = sum(len(idx) for idx in index_arrays) if index_arrays else 0 

861 else: 

862 der_map = deriv_map(n_bases, n_order) 

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