Coverage for jetgp/kernel_funcs/kernel_funcs.py: 71%

1import numpy as np 

2import jetgp.utils 

3from line_profiler import profile 

4import importlib 

5import subprocess 

6import sys 

7import os 

8import warnings 

9 

10 

11def get_oti_module(n_bases, n_order, auto_compile=True, otilib_path=None, use_sparse=False): 

12 """ 

13 Dynamically import the correct PyOTI static library. 

14 If the module doesn't exist and auto_compile=True, attempts to compile it. 

15 Falls back to pyoti.sparse if compilation fails or is disabled. 

16 

17 Parameters 

18 ---------- 

19 n_bases : int 

20 Number of bases (dimension of the input space). 

21 n_order : int 

22 Derivative order for the GP. The OTI order will be 2*n_order. 

23 auto_compile : bool, optional (default=False) 

24 If True, attempt to compile missing modules automatically. 

25 Requires jetgp.cmod_writer and jetgp.build_static to be available. 

26 otilib_path : str, optional 

27 Path to otilib-master directory. If None, attempts auto-detection. 

28 

29 ------- 

30 module 

31 The appropriate pyoti.static.onummXnY module, or pyoti.sparse as fallback. 

32 """ 

33 if n_order == 0: 

34 module_name = "pyoti.real" 

35 return importlib.import_module(module_name) 

36 

37 oti_order = 2 * n_order 

38 module_name = f"pyoti.static.onumm{n_bases}n{oti_order}" 

39 if use_sparse: 

40 return importlib.import_module("pyoti.sparse") 

41 try: 

42 return importlib.import_module(module_name) 

43 except ModuleNotFoundError: 

44 if not auto_compile: 

45 warnings.warn( 

46 f"PyOTI static module '{module_name}' not found. " 

47 f"Falling back to pyoti.sparse which is significantly slower.\n" 

48 f"For better performance, compile the static module manually:\n" 

49 f" 1. cd /path/to/otilib-master/build\n" 

50 f" 2. cmake ..\n" 

51 f" 3. make m{n_bases}n{oti_order} -j8\n" 

52 f" 4. python build_static.py m{n_bases}n{oti_order}", 

53 UserWarning 

54 ) 

55 return importlib.import_module("pyoti.sparse") 

56 

57 # Check if auto-compile tools are available 

58 try: 

59 from jetgp.cmod_writer import writer 

60 from jetgp.build_static import build_module 

61 except ImportError: 

62 warnings.warn( 

63 f"PyOTI static module '{module_name}' not found and auto-compile " 

64 f"tools (jetgp.cmod_writer, jetgp.build_static) are not available.\n" 

65 f"Falling back to pyoti.sparse which is significantly slower.\n" 

66 f"For better performance, compile the static module manually:\n" 

67 f" 1. cd /path/to/otilib-master/build\n" 

68 f" 2. cmake ..\n" 

69 f" 3. make m{n_bases}n{oti_order} -j8\n" 

70 f" 4. python build_static.py m{n_bases}n{oti_order}", 

71 UserWarning 

72 ) 

73 return importlib.import_module("pyoti.sparse") 

74 

75 print(f"Module '{module_name}' not found. Attempting to compile...") 

76 

77 try: 

78 _compile_oti_module(n_bases, oti_order, otilib_path) 

79 

80 # Clear import caches and retry 

81 importlib.invalidate_caches() 

82 

83 return importlib.import_module(module_name) 

84 except Exception as e: 

85 warnings.warn( 

86 f"Failed to compile PyOTI static module '{module_name}': {e}\n" 

87 f"Falling back to pyoti.sparse which is significantly slower.\n" 

88 f"For better performance, compile the static module manually:\n" 

89 f" 1. cd /path/to/otilib-master/build\n" 

90 f" 2. cmake ..\n" 

91 f" 3. make m{n_bases}n{oti_order} -j8\n" 

92 f" 4. python build_static.py m{n_bases}n{oti_order}", 

93 UserWarning 

94 ) 

95 return importlib.import_module("pyoti.sparse") 

96 

97 

98def _get_otilib_path(otilib_path=None): 

99 """ 

100 Auto-detect otilib path from the installed pyoti package location. 

101 

102 Parameters 

103 ---------- 

104 otilib_path : str, optional 

105 Override path to otilib-master directory. 

106 

107 Returns 

108 ------- 

109 otilib_path : str 

110 Path to otilib-master directory. 

111 """ 

112 # Use explicit argument if provided 

113 if otilib_path is not None: 

114 if not os.path.isdir(otilib_path): 

115 raise RuntimeError(f"otilib path does not exist: {otilib_path}") 

116 return otilib_path 

117 

118 # Check environment variable 

119 otilib_path = os.environ.get('OTILIB_PATH') 

120 if otilib_path is not None: 

121 if not os.path.isdir(otilib_path): 

122 raise RuntimeError(f"OTILIB_PATH does not exist: {otilib_path}") 

123 return otilib_path 

124 

125 # Auto-detect from installed pyoti 

126 try: 

127 import pyoti 

128 

129 # Get the pyoti package location 

130 if hasattr(pyoti, '__path__'): 

131 pyoti_install_path = pyoti.__path__[0] 

132 elif hasattr(pyoti, '__file__'): 

133 pyoti_install_path = os.path.dirname(pyoti.__file__) 

134 else: 

135 raise AttributeError("Cannot determine pyoti installation path") 

136 

137 # Navigate up from pyoti to find otilib root 

138 # Typical structure: otilib-master/src/python/pyoti/pyoti/__init__.py 

139 current = pyoti_install_path 

140 for _ in range(6): # Navigate up to 6 levels 

141 parent = os.path.dirname(current) 

142 

143 # Check if this looks like otilib root 

144 # Must have BOTH CMakeLists.txt AND src/ directory (not just build dir) 

145 potential_cmake = os.path.join(parent, 'CMakeLists.txt') 

146 potential_src = os.path.join(parent, 'src') 

147 potential_include = os.path.join(parent, 'include') 

148 

149 if (os.path.isfile(potential_cmake) and 

150 os.path.isdir(potential_src) and 

151 os.path.isdir(potential_include)): 

152 otilib_path = parent 

153 break 

154 

155 current = parent 

156 

157 except ImportError: 

158 pass 

159 

160 # Final validation 

161 if otilib_path is None: 

162 raise RuntimeError( 

163 "Could not auto-detect otilib path. Please either:\n" 

164 " 1. Set the OTILIB_PATH environment variable\n" 

165 " 2. Pass otilib_path explicitly to get_oti_module()\n" 

166 " 3. Ensure pyoti is installed from the otilib source tree" 

167 ) 

168 

169 if not os.path.isdir(otilib_path): 

170 raise RuntimeError(f"otilib path does not exist: {otilib_path}") 

171 

172 return otilib_path 

173 

174 

175def _compile_oti_module(n_bases, oti_order, otilib_path=None): 

176 """ 

177 Compile a PyOTI static module. 

178 

179 Parameters 

180 ---------- 

181 n_bases : int 

182 Number of bases. 

183 oti_order : int 

184 OTI order (already multiplied by 2). 

185 otilib_path : str, optional 

186 Path to otilib-master directory. 

187 """ 

188 # Auto-detect path 

189 otilib_path = _get_otilib_path(otilib_path) 

190 

191 build_dir = os.path.join(otilib_path, 'build') 

192 module_target = f"m{n_bases}n{oti_order}" 

193 

194 print(f"Compiling OTI module: {module_target}") 

195 print(f" otilib_path: {otilib_path}") 

196 print(f" build_dir: {build_dir}") 

197 

198 # Step 1: Generate C code using cmod_writer (from jetgp) 

199 print(f"Step 1/4: Generating C code for m={n_bases}, n={oti_order}...") 

200 _run_cmod_writer(n_bases, oti_order, otilib_path) 

201 

202 # Step 2: Run cmake (if needed) 

203 print("Step 2/4: Running cmake...") 

204 _run_cmake(build_dir) 

205 

206 # Step 3: Run make 

207 print(f"Step 3/4: Compiling {module_target}...") 

208 _run_make(build_dir, module_target) 

209 

210 # Step 4: Build and install Python module 

211 print(f"Step 4/4: Building Python module...") 

212 _run_build_static(build_dir, module_target, otilib_path) 

213 

214 print(f"Successfully compiled {module_target}") 

215 

216 

217def _run_cmod_writer(n_bases, oti_order, otilib_path): 

218 """Generate C code using cmod_writer from jetgp.""" 

219 from jetgp.cmod_writer import writer 

220 

221 w = writer(nbases=n_bases, order=oti_order) 

222 w.write_files(base_dir=otilib_path) 

223 

224 

225def _run_cmake(build_dir): 

226 """Run cmake in the build directory.""" 

227 os.makedirs(build_dir, exist_ok=True) 

228 

229 result = subprocess.run( 

230 ['cmake', '..'], 

231 cwd=build_dir, 

232 capture_output=True, 

233 text=True 

234 ) 

235 

236 if result.returncode != 0: 

237 raise RuntimeError( 

238 f"cmake failed with return code {result.returncode}.\n" 

239 f"stdout: {result.stdout}\n" 

240 f"stderr: {result.stderr}" 

241 ) 

242 

243 

244def _run_make(build_dir, module_target, n_jobs=8): 

245 """Run make for the specific module target.""" 

246 result = subprocess.run( 

247 ['make', module_target, f'-j{n_jobs}'], 

248 cwd=build_dir, 

249 capture_output=True, 

250 text=True 

251 ) 

252 

253 if result.returncode != 0: 

254 raise RuntimeError( 

255 f"make failed with return code {result.returncode}.\n" 

256 f"stdout: {result.stdout}\n" 

257 f"stderr: {result.stderr}" 

258 ) 

259 

260 

261def _run_build_static(build_dir, module_target, otilib_path): 

262 """Build and install the Python module using jetgp's build_static.""" 

263 from jetgp.build_static import build_module 

264 

265 build_module(module_target, otilib_path=otilib_path, build_dir=build_dir) 

266 

267 

268class KernelFactory: 

269 """ 

270 Factory for generating different kernel functions (SE, RQ, SineExp, Matérn) 

271 in isotropic and anisotropic forms with caching for improved performance. 

272 

273 Attributes 

274 ---------- 

275 dim : int 

276 Dimensionality of the input space. 

277 normalize : bool 

278 Whether to normalize inputs (scaling differences to [-3, 3]). 

279 differences_by_dim : list of arrays 

280 Pairwise differences between input points, by dimension. 

281 true_noise_std : float, optional 

282 Known noise standard deviation (for adjusting noise bounds). 

283 bounds : list of tuples 

284 Hyperparameter bounds (log10 space). 

285 nu : float 

286 Smoothness parameter for the Matérn kernel. 

287 n_order : int 

288 Order of derivatives for kernel smoothness. 

289 """ 

290 

291 def __init__(self, dim, normalize, differences_by_dim, n_order, 

292 true_noise_std=None, smoothness_parameter=None, oti_module=None): 

293 self.dim = dim 

294 self.normalize = normalize 

295 self.differences_by_dim = differences_by_dim 

296 self.true_noise_std = true_noise_std 

297 self.bounds = [] 

298 self.oti = oti_module 

299 if smoothness_parameter is not None: 

300 self.alpha = smoothness_parameter 

301 self.nu = smoothness_parameter + 0.5 

302 else: 

303 self.alpha = 1 

304 self.nu = 1.5 

305 self.n_order = n_order 

306 # Dynamic OTI import 

307 

308 # Initialize caching infrastructure 

309 self._init_caches() 

310 

311 # ------------------------------------------------------------------- 

312 # Caching Infrastructure 

313 # ------------------------------------------------------------------- 

314 

315 def _init_caches(self): 

316 """Initialize all cache variables.""" 

317 # Temporary array cache 

318 self._cached_shape = None 

319 self._tmp1 = None 

320 self._tmp2 = None 

321 self._sqdist = None 

322 

323 # Hyperparameter cache 

324 self._cached_length_scales = None 

325 self._cached_ell = None 

326 self._cached_sigma_f_sq = None 

327 self._cached_alpha = None 

328 self._cached_p = None 

329 self._cached_pi_over_p = None 

330 

331 # Fused sqdist cache 

332 self._cached_ell_sq = None 

333 self._has_fused_sqdist = None 

334 

335 def clear_caches(self): 

336 """Clear all caches. Call when training data changes.""" 

337 self._init_caches() 

338 

339 def _ensure_temp_arrays(self, shape): 

340 """ 

341 Ensure temporary arrays exist with correct shape. 

342 

343 Parameters 

344 ---------- 

345 shape : tuple 

346 Required shape for temporary arrays. 

347 

348 Returns 

349 ------- 

350 tuple 

351 (tmp1, tmp2, sqdist) temporary arrays. 

352 """ 

353 if self._cached_shape != shape: 

354 self._cached_shape = shape 

355 self._tmp1 = self.oti.zeros(shape) 

356 self._tmp2 = self.oti.zeros(shape) 

357 self._sqdist = self.oti.zeros(shape) 

358 return self._tmp1, self._tmp2, self._sqdist 

359 

360 def _reset_sqdist(self): 

361 """Reset sqdist accumulator to zero.""" 

362 if self._sqdist is None: 

363 return 

364 # Use the most efficient method available in oti 

365 if hasattr(self._sqdist, 'fill'): 

366 self._sqdist.fill(0) 

367 elif hasattr(self.oti, 'set_zero'): 

368 self.oti.set_zero(self._sqdist) 

369 else: 

370 # Fallback: multiply by zero 

371 self.oti.mul(0.0, self._sqdist, out=self._sqdist) 

372 

373 def _compute_sqdist_aniso(self, differences_by_dim, ell, sqdist, tmp1, tmp2): 

374 """Compute sqdist = Σ ell[i]² * diff[i]², using fused C kernel when available.""" 

375 if self._has_fused_sqdist is None: 

376 self._has_fused_sqdist = hasattr(sqdist, 'fused_sqdist') 

377 if self._has_fused_sqdist: 

378 ell_sq = np.ascontiguousarray(ell ** 2, dtype=np.float64) 

379 sqdist.fused_sqdist(differences_by_dim, ell_sq) 

380 else: 

381 self._reset_sqdist() 

382 for i in range(self.dim): 

383 self.oti.mul(ell[i], differences_by_dim[i], out=tmp1) 

384 self.oti.mul(tmp1, tmp1, out=tmp2) 

385 self.oti.sum(sqdist, tmp2, out=sqdist) 

386 

387 def _compute_sqdist_iso(self, differences_by_dim, ell, sqdist, tmp1, tmp2): 

388 """Compute sqdist = Σ ell² * diff[i]², using fused C kernel when available.""" 

389 if self._has_fused_sqdist is None: 

390 self._has_fused_sqdist = hasattr(sqdist, 'fused_sqdist') 

391 if self._has_fused_sqdist: 

392 ell_sq_val = float(ell) ** 2 

393 ell_sq = np.full(self.dim, ell_sq_val, dtype=np.float64) 

394 sqdist.fused_sqdist(differences_by_dim, ell_sq) 

395 else: 

396 self._reset_sqdist() 

397 for i in range(self.dim): 

398 self.oti.mul(ell, differences_by_dim[i], out=tmp1) 

399 self.oti.mul(tmp1, tmp1, out=tmp2) 

400 self.oti.sum(sqdist, tmp2, out=sqdist) 

401 

402 # ------------------------------------------------------------------- 

403 # Hyperparameter Caching Methods 

404 # ------------------------------------------------------------------- 

405 

406 def _cache_se_params_aniso(self, length_scales): 

407 """Cache anisotropic SE kernel hyperparameters.""" 

408 ls_tuple = tuple(float(x) for x in length_scales) 

409 if self._cached_length_scales != ('se_aniso', ls_tuple): 

410 self._cached_length_scales = ('se_aniso', ls_tuple) 

411 self._cached_ell = 10 ** np.array(length_scales[:-1]) 

412 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

413 return self._cached_ell, self._cached_sigma_f_sq 

414 

415 def _cache_se_params_iso(self, length_scales): 

416 """Cache isotropic SE kernel hyperparameters.""" 

417 ls_tuple = tuple(float(x) for x in length_scales) 

418 if self._cached_length_scales != ('se_iso', ls_tuple): 

419 self._cached_length_scales = ('se_iso', ls_tuple) 

420 self._cached_ell = 10 ** length_scales[0] 

421 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

422 return self._cached_ell, self._cached_sigma_f_sq 

423 

424 def _cache_rq_params_aniso(self, length_scales): 

425 """Cache anisotropic RQ kernel hyperparameters.""" 

426 ls_tuple = tuple(float(x) for x in length_scales) 

427 if self._cached_length_scales != ('rq_aniso', ls_tuple): 

428 self._cached_length_scales = ('rq_aniso', ls_tuple) 

429 self._cached_ell = 10 ** np.array(length_scales[:self.dim]) 

430 self._cached_alpha = 10 ** length_scales[self.dim] 

431 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

432 return self._cached_ell, self._cached_alpha, self._cached_sigma_f_sq 

433 

434 def _cache_rq_params_iso(self, length_scales): 

435 """Cache isotropic RQ kernel hyperparameters.""" 

436 ls_tuple = tuple(float(x) for x in length_scales) 

437 if self._cached_length_scales != ('rq_iso', ls_tuple): 

438 self._cached_length_scales = ('rq_iso', ls_tuple) 

439 self._cached_ell = 10 ** length_scales[0] 

440 self._cached_alpha = np.exp(length_scales[1]) 

441 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

442 return self._cached_ell, self._cached_alpha, self._cached_sigma_f_sq 

443 

444 def _cache_sine_exp_params_aniso(self, length_scales): 

445 """Cache anisotropic Sine-Exponential kernel hyperparameters.""" 

446 ls_tuple = tuple(float(x) for x in length_scales) 

447 if self._cached_length_scales != ('sine_exp_aniso', ls_tuple): 

448 self._cached_length_scales = ('sine_exp_aniso', ls_tuple) 

449 self._cached_ell = 10 ** np.array(length_scales[:self.dim]) 

450 self._cached_p = 10 ** np.array(length_scales[self.dim:2*self.dim]) 

451 self._cached_pi_over_p = np.pi / self._cached_p 

452 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

453 return self._cached_ell, self._cached_pi_over_p, self._cached_sigma_f_sq 

454 

455 def _cache_sine_exp_params_iso(self, length_scales): 

456 """Cache isotropic Sine-Exponential kernel hyperparameters.""" 

457 ls_tuple = tuple(float(x) for x in length_scales) 

458 if self._cached_length_scales != ('sine_exp_iso', ls_tuple): 

459 self._cached_length_scales = ('sine_exp_iso', ls_tuple) 

460 self._cached_ell = 10 ** length_scales[0] 

461 self._cached_p = 10 ** length_scales[1] 

462 self._cached_pi_over_p = np.pi / self._cached_p 

463 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

464 return self._cached_ell, self._cached_pi_over_p, self._cached_sigma_f_sq 

465 

466 def _cache_matern_params_aniso(self, length_scales): 

467 """Cache anisotropic Matern kernel hyperparameters.""" 

468 ls_tuple = tuple(float(x) for x in length_scales) 

469 if self._cached_length_scales != ('matern_aniso', ls_tuple): 

470 self._cached_length_scales = ('matern_aniso', ls_tuple) 

471 self._cached_ell = 10 ** np.array(length_scales[:-1]) 

472 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

473 return self._cached_ell, self._cached_sigma_f_sq 

474 

475 def _cache_matern_params_iso(self, length_scales): 

476 """Cache isotropic Matern kernel hyperparameters.""" 

477 ls_tuple = tuple(float(x) for x in length_scales) 

478 if self._cached_length_scales != ('matern_iso', ls_tuple): 

479 self._cached_length_scales = ('matern_iso', ls_tuple) 

480 self._cached_ell = 10 ** length_scales[0] 

481 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

482 return self._cached_ell, self._cached_sigma_f_sq 

483 

484 def _cache_si_params_aniso(self, length_scales): 

485 """Cache anisotropic SI kernel hyperparameters.""" 

486 ls_tuple = tuple(float(x) for x in length_scales) 

487 if self._cached_length_scales != ('si_aniso', ls_tuple): 

488 self._cached_length_scales = ('si_aniso', ls_tuple) 

489 self._cached_ell = 10 ** np.array(length_scales[:-1]) 

490 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

491 return self._cached_ell, self._cached_sigma_f_sq 

492 

493 def _cache_si_params_iso(self, length_scales): 

494 """Cache isotropic SI kernel hyperparameters.""" 

495 ls_tuple = tuple(float(x) for x in length_scales) 

496 if self._cached_length_scales != ('si_iso', ls_tuple): 

497 self._cached_length_scales = ('si_iso', ls_tuple) 

498 self._cached_ell = 10 ** length_scales[0] 

499 self._cached_sigma_f_sq = (10 ** length_scales[-1]) ** 2 

500 return self._cached_ell, self._cached_sigma_f_sq 

501 

502 # ------------------------------------------------------------------- 

503 # Bounds and Factory Methods 

504 # ------------------------------------------------------------------- 

505 

506 def get_bounds_from_data(self): 

507 """ 

508 Computes bounds for hyperparameters based on the observed data range. 

509 """ 

510 self.bounds = [] 

511 for diffs in self.differences_by_dim: 

512 min_val = float(diffs.real.min()) 

513 max_val = float(diffs.real.max()) 

514 self.bounds.append((-3, np.log(max_val))) 

515 

516 def create_kernel(self, kernel_name, kernel_type): 

517 """ 

518 Returns a kernel function based on specified name and type. 

519 

520 Parameters 

521 ---------- 

522 kernel_name : str 

523 Name of the kernel ('SE', 'RQ', 'SineExp', 'Matern'). 

524 kernel_type : str 

525 Type of kernel ('anisotropic' or 'isotropic'). 

526 

527 Returns 

528 ------- 

529 callable 

530 The selected kernel function. 

531 """ 

532 # Clear caches when creating a new kernel 

533 self.clear_caches() 

534 

535 if not self.normalize: 

536 self.get_bounds_from_data() 

537 

538 if kernel_type == "anisotropic": 

539 return self._create_anisotropic(kernel_name) 

540 elif kernel_type == "isotropic": 

541 return self._create_isotropic(kernel_name) 

542 else: 

543 raise ValueError("Invalid kernel_type") 

544 

545 def _create_anisotropic(self, kernel): 

546 """ 

547 Sets bounds and returns the anisotropic kernel function. 

548 

549 Returns 

550 ------- 

551 callable 

552 The anisotropic kernel function. 

553 """ 

554 sigma_n_bound = (-16, -3) 

555 

556 if kernel == "SE": 

557 self._add_bounds([(-1, 5), sigma_n_bound]) 

558 return self.se_kernel_anisotropic 

559 elif kernel == "RQ": 

560 self._add_bounds([(-1, 5), (-1, 5), sigma_n_bound]) 

561 return self.rq_kernel_anisotropic 

562 elif kernel == "SineExp": 

563 self._add_bounds([(0.0, 5)] * self.dim + [(-1, 5), sigma_n_bound]) 

564 return self.sine_exp_kernel_anisotropic 

565 elif kernel == "Matern": 

566 self._add_bounds([(-1, 5), sigma_n_bound]) 

567 self.matern_kernel_prebuild = jetgp.utils.matern_kernel_builder( 

568 self.nu, oti_module=self.oti) 

569 return self.matern_kernel_anisotropic 

570 elif kernel == "SI": 

571 self._add_bounds([(-5, 5), sigma_n_bound]) 

572 self.SI_kernel_prebuild = jetgp.utils.generate_bernoulli_lambda( 

573 self.alpha) 

574 return self.SI_kernel_anisotropic 

575 else: 

576 raise NotImplementedError("Anisotropic kernel not implemented") 

577 

578 def _create_isotropic(self, kernel): 

579 """ 

580 Sets bounds and returns the isotropic kernel function. 

581 

582 Returns 

583 ------- 

584 callable 

585 The isotropic kernel function. 

586 """ 

587 sigma_n_bound = (-16, -3) 

588 

589 if self.normalize: 

590 core_bounds = [(-3, 3)] 

591 else: 

592 self.get_bounds_from_data() 

593 core_bounds = [( 

594 float(min([d.real.min() for d in self.differences_by_dim])), 

595 float(max([d.real.max() for d in self.differences_by_dim])) 

596 )] 

597 

598 if kernel == "SE": 

599 self.bounds = core_bounds + [(-1, 5), sigma_n_bound] 

600 return self.se_kernel_isotropic 

601 elif kernel == "RQ": 

602 self.bounds = core_bounds + [(-1, 5), (-1, 5), sigma_n_bound] 

603 return self.rq_kernel_isotropic 

604 elif kernel == "SineExp": 

605 self.bounds = core_bounds + [(0.0, 3.0), (-1, 5), sigma_n_bound] 

606 return self.sine_exp_kernel_isotropic 

607 elif kernel == "Matern": 

608 self.bounds = core_bounds + [(-1, 5), sigma_n_bound] 

609 self.matern_kernel_prebuild = jetgp.utils.matern_kernel_builder( 

610 self.nu, oti_module=self.oti) 

611 return self.matern_kernel_isotropic 

612 elif kernel == "SI": 

613 self.bounds = core_bounds + [(-5, 5), sigma_n_bound] 

614 self.SI_kernel_prebuild = jetgp.utils.generate_bernoulli_lambda( 

615 self.alpha) 

616 return self.SI_kernel_isotropic 

617 else: 

618 raise NotImplementedError("Isotropic kernel not implemented") 

619 

620 def _add_bounds(self, extra_bounds): 

621 """ 

622 Append additional hyperparameter bounds to the kernel's configuration. 

623 

624 Parameters 

625 ---------- 

626 extra_bounds : list of tuple 

627 Bounds to append, where each tuple is a (min, max) pair in log10 scale. 

628 """ 

629 if self.normalize: 

630 self.bounds = [(-2, 1)] * self.dim + extra_bounds 

631 else: 

632 self.bounds += extra_bounds 

633 

634 # ------------------------------------------------------------------- 

635 # Anisotropic Kernel Implementations with Caching 

636 # ------------------------------------------------------------------- 

637 

638 @profile 

639 def se_kernel_anisotropic(self, differences_by_dim, length_scales): 

640 """ 

641 Anisotropic Squared Exponential (SE) kernel with caching. 

642 

643 Parameters 

644 ---------- 

645 differences_by_dim : list of ndarray 

646 Pairwise differences by dimension. 

647 length_scales : list 

648 Hyperparameters: [ell_1, ..., ell_dim, sigma_f] 

649 

650 Returns 

651 ------- 

652 ndarray 

653 Kernel matrix values. 

654 """ 

655 ell, sigma_f_sq = self._cache_se_params_aniso(length_scales) 

656 tmp1, tmp2, sqdist = self._ensure_temp_arrays( 

657 differences_by_dim[0].shape) 

658 self._compute_sqdist_aniso(differences_by_dim, ell, sqdist, tmp1, tmp2) 

659 

660 self.oti.exp((-0.5) * sqdist, out=tmp1) 

661 self.oti.mul(sigma_f_sq, tmp1, out=tmp2) 

662 return tmp2 

663 

664 def rq_kernel_anisotropic(self, differences_by_dim, length_scales): 

665 """ 

666 Anisotropic Rational Quadratic (RQ) kernel with caching. 

667 

668 Parameters 

669 ---------- 

670 differences_by_dim : list of ndarray 

671 Pairwise differences by dimension. 

672 length_scales : list 

673 Hyperparameters: [ell_1, ..., ell_dim, alpha, sigma_f] 

674 

675 Returns 

676 ------- 

677 ndarray 

678 Kernel matrix values. 

679 """ 

680 ell, alpha, sigma_f_sq = self._cache_rq_params_aniso(length_scales) 

681 tmp1, tmp2, sqdist = self._ensure_temp_arrays( 

682 differences_by_dim[0].shape) 

683 self._compute_sqdist_aniso(differences_by_dim, ell, sqdist, tmp1, tmp2) 

684 

685 # (1 + sqdist / (2 * alpha))^(-alpha) 

686 inv_2alpha = 1.0 / (2 * alpha) 

687 self.oti.mul(sqdist, inv_2alpha, out=tmp1) 

688 self.oti.sum(1.0, tmp1, out=tmp2) 

689 self.oti.pow(tmp2, -alpha, out=tmp1) 

690 self.oti.mul(sigma_f_sq, tmp1, out=tmp2) 

691 return tmp2 

692 

693 def sine_exp_kernel_anisotropic(self, differences_by_dim, length_scales): 

694 """ 

695 Anisotropic Sine-Exponential (Periodic) kernel with caching. 

696 

697 Parameters 

698 ---------- 

699 differences_by_dim : list of ndarray 

700 Pairwise differences by dimension. 

701 length_scales : list 

702 Hyperparameters: [ell_1, ..., ell_dim, p_1, ..., p_dim, sigma_f] 

703 

704 Returns 

705 ------- 

706 ndarray 

707 Kernel matrix values. 

708 """ 

709 ell, pi_over_p, sigma_f_sq = self._cache_sine_exp_params_aniso( 

710 length_scales) 

711 tmp1, tmp2, sqdist = self._ensure_temp_arrays( 

712 differences_by_dim[0].shape) 

713 self._reset_sqdist() 

714 

715 for i in range(self.dim): 

716 self.oti.mul(pi_over_p[i], differences_by_dim[i], out=tmp1) 

717 self.oti.sin(tmp1, out=tmp2) 

718 self.oti.mul(ell[i], tmp2, out=tmp1) 

719 self.oti.mul(tmp1, tmp1, out=tmp2) 

720 self.oti.sum(sqdist, tmp2, out=sqdist) 

721 

722 self.oti.mul(-2.0, sqdist, out=tmp1) 

723 self.oti.exp(tmp1, out=tmp2) 

724 self.oti.mul(sigma_f_sq, tmp2, out=tmp1) 

725 return tmp1 

726 

727 def matern_kernel_anisotropic(self, differences_by_dim, length_scales): 

728 """ 

729 Anisotropic Matérn kernel (half-integer ν) with caching. 

730 

731 Parameters 

732 ---------- 

733 differences_by_dim : list of ndarray 

734 Pairwise differences by dimension. 

735 length_scales : list 

736 Hyperparameters: [ell_1, ..., ell_dim, sigma_f] 

737 

738 Returns 

739 ------- 

740 ndarray 

741 Kernel matrix values. 

742 """ 

743 ell, sigma_f_sq = self._cache_matern_params_aniso(length_scales) 

744 tmp1, tmp2, sqdist = self._ensure_temp_arrays( 

745 differences_by_dim[0].shape) 

746 

747 # Compute scaled squared distance using fused op if available 

748 self._compute_sqdist_aniso(differences_by_dim, ell, sqdist, tmp1, tmp2) 

749 

750 # r = sqrt(sqdist + eps²) — regularise r directly (not each diff) 

751 # so that r.e([d]) = 0 at training-point diagonals, preserving 

752 # correct OTI derivative structure for the covariance blocks. 

753 _eps = 1e-10 

754 r = self.oti.sqrt(self.oti.sum(sqdist, _eps ** 2)) 

755 return sigma_f_sq * self.matern_kernel_prebuild(r) 

756 

757 def SI_kernel_anisotropic(self, differences_by_dim, length_scales): 

758 """ 

759 Anisotropic SI kernel with caching. 

760 

761 Parameters 

762 ---------- 

763 differences_by_dim : list of ndarray 

764 Pairwise differences by dimension. 

765 length_scales : list 

766 Hyperparameters: [ell_1, ..., ell_dim, sigma_f] 

767 

768 Returns 

769 ------- 

770 ndarray 

771 Kernel matrix values. 

772 """ 

773 ell, sigma_f_sq = self._cache_si_params_aniso(length_scales) 

774 

775 val = 1 

776 for i in range(self.dim): 

777 val = val * \ 

778 (1 + ell[i] * self.SI_kernel_prebuild(differences_by_dim[i])) 

779 return sigma_f_sq * val 

780 

781 # ------------------------------------------------------------------- 

782 # Isotropic Kernel Implementations with Caching 

783 # ------------------------------------------------------------------- 

784 

785 def se_kernel_isotropic(self, differences_by_dim, length_scales): 

786 """ 

787 Isotropic Squared Exponential (SE) kernel with caching. 

788 

789 Parameters 

790 ---------- 

791 differences_by_dim : list of ndarray 

792 Pairwise differences by dimension. 

793 length_scales : list 

794 Hyperparameters: [ell, sigma_f] 

795 

796 Returns 

797 ------- 

798 ndarray 

799 Kernel matrix values. 

800 """ 

801 ell, sigma_f_sq = self._cache_se_params_iso(length_scales) 

802 tmp1, tmp2, sqdist = self._ensure_temp_arrays( 

803 differences_by_dim[0].shape) 

804 self._compute_sqdist_iso(differences_by_dim, ell, sqdist, tmp1, tmp2) 

805 

806 self.oti.exp((-0.5) * sqdist, out=tmp1) 

807 self.oti.mul(sigma_f_sq, tmp1, out=tmp2) 

808 return tmp2 

809 

810 def rq_kernel_isotropic(self, differences_by_dim, length_scales): 

811 """ 

812 Isotropic Rational Quadratic (RQ) kernel with caching. 

813 

814 Parameters 

815 ---------- 

816 differences_by_dim : list of ndarray 

817 Pairwise differences by dimension. 

818 length_scales : list 

819 Hyperparameters: [ell, alpha, sigma_f] 

820 

821 Returns 

822 ------- 

823 ndarray 

824 Kernel matrix values. 

825 """ 

826 ell, alpha, sigma_f_sq = self._cache_rq_params_iso(length_scales) 

827 tmp1, tmp2, sqdist = self._ensure_temp_arrays( 

828 differences_by_dim[0].shape) 

829 self._compute_sqdist_iso(differences_by_dim, ell, sqdist, tmp1, tmp2) 

830 

831 # (1 + sqdist / (2 * alpha))^(-alpha) 

832 inv_2alpha = 1.0 / (2 * alpha) 

833 self.oti.mul(sqdist, inv_2alpha, out=tmp1) 

834 self.oti.sum(1.0, tmp1, out=tmp2) 

835 self.oti.pow(tmp2, -alpha, out=tmp1) 

836 self.oti.mul(sigma_f_sq, tmp1, out=tmp2) 

837 return tmp2 

838 

839 def sine_exp_kernel_isotropic(self, differences_by_dim, length_scales): 

840 """ 

841 Isotropic Sine-Exponential (Periodic) kernel with caching. 

842 

843 Parameters