Coverage for jetgp/kernel_funcs/kernel

1import numpy as np

2import jetgp.utils

3from line_profiler import profile

4import importlib

5import subprocess

6import sys

7import os

8import warnings

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.

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.

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)

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")

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")

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

77 try:

78 _compile_oti_module(n_bases, oti_order, otilib_path)

80 # Clear import caches and retry

81 importlib.invalidate_caches()

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")

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 # Check config file written by setup_otilib

126 from pathlib import Path

127 config_file = Path.home() / ".config" / "jetgp" / "otilib_path"

128 if config_file.exists():

129 candidate = config_file.read_text().strip()

130 if os.path.isdir(candidate):

131 return candidate

132

133 # Auto-detect from installed pyoti

134 try:

135 import pyoti

136

137 # Get the pyoti package location

138 if hasattr(pyoti, '__path__'):

139 pyoti_install_path = pyoti.__path__[0]

140 elif hasattr(pyoti, '__file__'):

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

142 else:

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

144

145 # Navigate up from pyoti to find otilib root

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

147 current = pyoti_install_path

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

149 parent = os.path.dirname(current)

150

151 # Check if this looks like otilib root

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

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

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

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

156

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

158 os.path.isdir(potential_src) and

159 os.path.isdir(potential_include)):

160 otilib_path = parent

161 break

162

163 current = parent

164

165 except ImportError:

166 pass

167

168 # Final validation

169 if otilib_path is None:

170 raise RuntimeError(

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

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

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

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

175 )

176

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

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

179

180 return otilib_path

181

182

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

184 """

185 Compile a PyOTI static module.

186

187 Parameters

188 ----------

189 n_bases : int

190 Number of bases.

191 oti_order : int

192 OTI order (already multiplied by 2).

193 otilib_path : str, optional

194 Path to otilib-master directory.

195 """

196 # Auto-detect path

197 otilib_path = _get_otilib_path(otilib_path)

198

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

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

201

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

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

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

205

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

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

208 _run_cmod_writer(n_bases, oti_order, otilib_path)

209

210 # Step 2: Run cmake (if needed)

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

212 _run_cmake(build_dir)

213

214 # Step 3: Run make

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

216 _run_make(build_dir, module_target)

217

218 # Step 4: Build and install Python module

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

220 _run_build_static(build_dir, module_target, otilib_path)

221

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

223

224

225def _run_cmod_writer(n_bases, oti_order, otilib_path):

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

227 from jetgp.cmod_writer import writer

228

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

230 w.write_files(base_dir=otilib_path)

231

232

233def _run_cmake(build_dir):

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

235 os.makedirs(build_dir, exist_ok=True)

236

237 result = subprocess.run(

238 ['cmake', '..'],

239 cwd=build_dir,

240 capture_output=True,

241 text=True

242 )

243

244 if result.returncode != 0:

245 raise RuntimeError(

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

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

248 f"stderr: {result.stderr}"

249 )

250

251

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

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

254 result = subprocess.run(

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

256 cwd=build_dir,

257 capture_output=True,

258 text=True

259 )

260

261 if result.returncode != 0:

262 raise RuntimeError(

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

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

265 f"stderr: {result.stderr}"

266 )

267

268

269def _run_build_static(build_dir, module_target, otilib_path):

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

271 from jetgp.build_static import build_module

272

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

274

275

276class KernelFactory:

277 """

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

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

280

281 Attributes

282 ----------

283 dim : int

284 Dimensionality of the input space.

285 normalize : bool

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

287 differences_by_dim : list of arrays

288 Pairwise differences between input points, by dimension.

289 true_noise_std : float, optional

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

291 bounds : list of tuples

292 Hyperparameter bounds (log10 space).

293 nu : float

294 Smoothness parameter for the Matérn kernel.

295 n_order : int

296 Order of derivatives for kernel smoothness.

297 """

298

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

300 true_noise_std=None, smoothness_parameter=None, oti_module=None,

301 sparse_diffs=True):

302 self.dim = dim

303 self.normalize = normalize

304 self.differences_by_dim = differences_by_dim

305 self.true_noise_std = true_noise_std

306 self.bounds = []

307 self.oti = oti_module

308 # These cached buffers are reused across kernel evaluations. They must

309 # start zeroed because some fused OTI kernels only write the components

310 # they touch, leaving any uninitialized entries to leak NaNs/Infs.

311 self._alloc = oti_module.zeros

312 self.sparse_diffs = sparse_diffs

313 if smoothness_parameter is not None:

314 self.alpha = smoothness_parameter

315 self.nu = smoothness_parameter + 0.5

316 else:

317 self.alpha = 1

318 self.nu = 1.5

319 self.n_order = n_order

320 # Dynamic OTI import

321

322 # Initialize caching infrastructure

323 self._init_caches()

324

325 # -------------------------------------------------------------------

326 # Caching Infrastructure

327 # -------------------------------------------------------------------

328

329 def _init_caches(self):

330 """Initialize all cache variables."""

331 # Temporary array cache

332 self._cached_shape = None

333 self._tmp1 = None

334 self._tmp2 = None

335 self._sqdist = None

336

337 # Hyperparameter cache

338 self._cached_length_scales = None

339 self._cached_ell = None

340 self._cached_sigma_f_sq = None

341 self._cached_alpha = None

342 self._cached_p = None

343 self._cached_pi_over_p = None

344

345 # Fused sqdist cache

346 self._cached_ell_sq = None

347 self._has_fused_sqdist_sparse = None

348 self._has_fused_sqdist_linear = None

349 self._has_fused_sqdist = None

350

351 def clear_caches(self):

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

353 self._init_caches()

354

355 def _ensure_temp_arrays(self, shape):

356 """

357 Ensure temporary arrays exist with correct shape.

358

359 Parameters

360 ----------

361 shape : tuple

362 Required shape for temporary arrays.

363

364 Returns

365 -------

366 tuple

367 (tmp1, tmp2, sqdist) temporary arrays.

368 """

369 if shape != self._cached_shape:

370 self._cached_shape = shape

371 self._tmp1 = self._alloc(shape)

372 self._tmp2 = self._alloc(shape)

373 self._sqdist = self._alloc(shape)

374 return self._tmp1, self._tmp2, self._sqdist

375

376 def _reset_sqdist(self):

377 """Reset sqdist accumulator to zero."""

378 if self._sqdist is None:

379 return

380 # Use the most efficient method available in oti

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

382 self._sqdist.fill(0)

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

384 self.oti.set_zero(self._sqdist)

385 else:

386 # Fallback: multiply by zero

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

388

389 def _init_fused_sqdist_caps(self, sqdist):

390 """Lazily detect which fused squared-distance kernels are available."""

391 if self._has_fused_sqdist_sparse is None:

392 self._has_fused_sqdist_sparse = (

393 self.sparse_diffs and hasattr(sqdist, 'fused_sqdist_sparse')

394 )

395 if self._has_fused_sqdist_linear is None:

396 self._has_fused_sqdist_linear = hasattr(sqdist, 'fused_sqdist_linear')

397 if self._has_fused_sqdist is None:

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

399

400 @profile

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

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

403 self._init_fused_sqdist_caps(sqdist)

404 if self._has_fused_sqdist_sparse:

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

406 sqdist.fused_sqdist_sparse(differences_by_dim, ell_sq)

407 elif self._has_fused_sqdist_linear:

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

409 sqdist.fused_sqdist_linear(differences_by_dim, ell_sq)

410 elif self._has_fused_sqdist:

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

412 sqdist.fused_sqdist(differences_by_dim, ell_sq)

413 else:

414 self._reset_sqdist()

415 for i in range(self.dim):

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

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

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

419

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

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

422 self._init_fused_sqdist_caps(sqdist)

423 if self._has_fused_sqdist_sparse:

424 ell_sq_val = float(ell) ** 2

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

426 sqdist.fused_sqdist_sparse(differences_by_dim, ell_sq)

427 elif self._has_fused_sqdist_linear:

428 ell_sq_val = float(ell) ** 2

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

430 sqdist.fused_sqdist_linear(differences_by_dim, ell_sq)

431 elif self._has_fused_sqdist:

432 ell_sq_val = float(ell) ** 2

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

434 sqdist.fused_sqdist(differences_by_dim, ell_sq)

435 else:

436 self._reset_sqdist()

437 for i in range(self.dim):

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

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

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

441 @profile

442 def _compute_neg_half_sqdist_aniso(self, differences_by_dim, ell, sqdist, tmp1, tmp2):

443 """Compute -0.5 * Σ ell[i]² * diff[i]² directly, absorbing the -0.5 into ell_sq

444 on the fused path to avoid a separate mul step."""

445 self._init_fused_sqdist_caps(sqdist)

446 if self._has_fused_sqdist_sparse:

447 ell_sq = np.ascontiguousarray(-0.5 * ell ** 2, dtype=np.float64)

448 sqdist.fused_sqdist_sparse(differences_by_dim, ell_sq)

449 elif self._has_fused_sqdist_linear:

450 ell_sq = np.ascontiguousarray(-0.5 * ell ** 2, dtype=np.float64)

451 sqdist.fused_sqdist_linear(differences_by_dim, ell_sq)

452 elif self._has_fused_sqdist:

453 ell_sq = np.ascontiguousarray(-0.5 * ell ** 2, dtype=np.float64)

454 sqdist.fused_sqdist(differences_by_dim, ell_sq)

455 else:

456 self._reset_sqdist()

457 for i in range(self.dim):

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

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

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

461 self.oti.mul(-0.5, sqdist, out=sqdist)

462

463 def _compute_neg_half_sqdist_iso(self, differences_by_dim, ell, sqdist, tmp1, tmp2):

464 """Compute -0.5 * Σ ell² * diff[i]² directly, absorbing the -0.5 into ell_sq

465 on the fused path to avoid a separate mul step."""

466 self._init_fused_sqdist_caps(sqdist)

467 if self._has_fused_sqdist_sparse:

468 ell_sq_val = -0.5 * float(ell) ** 2

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

470 sqdist.fused_sqdist_sparse(differences_by_dim, ell_sq)

471 elif self._has_fused_sqdist_linear:

472 ell_sq_val = -0.5 * float(ell) ** 2

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

474 sqdist.fused_sqdist_linear(differences_by_dim, ell_sq)

475 elif self._has_fused_sqdist:

476 ell_sq_val = -0.5 * float(ell) ** 2

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

478 sqdist.fused_sqdist(differences_by_dim, ell_sq)

479 else:

480 self._reset_sqdist()

481 for i in range(self.dim):

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

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

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

485 self.oti.mul(-0.5, sqdist, out=sqdist)

486

487 def _compute_scaled_sqdist_aniso(self, differences_by_dim, ell, scale, sqdist, tmp1, tmp2):

488 """Compute scale * Σ ell[i]² * diff[i]², absorbing scale into ell_sq on the fused

489 path to avoid a separate mul step after sqdist construction."""

490 self._init_fused_sqdist_caps(sqdist)

491 if self._has_fused_sqdist_sparse:

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

493 sqdist.fused_sqdist_sparse(differences_by_dim, ell_sq)

494 elif self._has_fused_sqdist_linear:

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

496 sqdist.fused_sqdist_linear(differences_by_dim, ell_sq)

497 elif self._has_fused_sqdist:

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

499 sqdist.fused_sqdist(differences_by_dim, ell_sq)

500 else:

501 self._reset_sqdist()

502 for i in range(self.dim):

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

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

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

506 self.oti.mul(scale, sqdist, out=sqdist)

507

508 def _compute_scaled_sqdist_iso(self, differences_by_dim, ell, scale, sqdist, tmp1, tmp2):

509 """Compute scale * Σ ell² * diff[i]², absorbing scale into ell_sq on the fused

510 path to avoid a separate mul step after sqdist construction."""

511 self._init_fused_sqdist_caps(sqdist)

512 if self._has_fused_sqdist_sparse:

513 ell_sq_val = scale * float(ell) ** 2

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

515 sqdist.fused_sqdist_sparse(differences_by_dim, ell_sq)

516 elif self._has_fused_sqdist_linear:

517 ell_sq_val = scale * float(ell) ** 2

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

519 sqdist.fused_sqdist_linear(differences_by_dim, ell_sq)

520 elif self._has_fused_sqdist:

521 ell_sq_val = scale * float(ell) ** 2

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

523 sqdist.fused_sqdist(differences_by_dim, ell_sq)

524 else:

525 self._reset_sqdist()

526 for i in range(self.dim):

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

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

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

530 self.oti.mul(scale, sqdist, out=sqdist)

531

532 # -------------------------------------------------------------------

533 # Hyperparameter Caching Methods

534 # -------------------------------------------------------------------

535

536 def _cache_se_params_aniso(self, length_scales):

537 """Cache anisotropic SE kernel hyperparameters."""

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

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

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

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

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

543 return self._cached_ell, self._cached_sigma_f_sq

544

545 def _cache_se_params_iso(self, length_scales):

546 """Cache isotropic SE kernel hyperparameters."""

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

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

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

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

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

552 return self._cached_ell, self._cached_sigma_f_sq

553

554 def _cache_rq_params_aniso(self, length_scales):

555 """Cache anisotropic RQ kernel hyperparameters."""

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

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

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

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

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

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

562 return self._cached_ell, self._cached_alpha, self._cached_sigma_f_sq

563

564 def _cache_rq_params_iso(self, length_scales):

565 """Cache isotropic RQ kernel hyperparameters."""

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

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

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

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

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

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

572 return self._cached_ell, self._cached_alpha, self._cached_sigma_f_sq

573

574 def _cache_sine_exp_params_aniso(self, length_scales):

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

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

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

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

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

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

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

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

583 return self._cached_ell, self._cached_pi_over_p, self._cached_sigma_f_sq

584

585 def _cache_sine_exp_params_iso(self, length_scales):

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

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

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

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

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

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

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

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

594 return self._cached_ell, self._cached_pi_over_p, self._cached_sigma_f_sq

595

596 def _cache_matern_params_aniso(self, length_scales):

597 """Cache anisotropic Matern kernel hyperparameters."""

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

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

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

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

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

603 return self._cached_ell, self._cached_sigma_f_sq

604

605 def _cache_matern_params_iso(self, length_scales):

606 """Cache isotropic Matern kernel hyperparameters."""

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

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

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

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

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

612 return self._cached_ell, self._cached_sigma_f_sq

613

614 def _cache_si_params_aniso(self, length_scales):

615 """Cache anisotropic SI kernel hyperparameters."""

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

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

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

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

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

621 return self._cached_ell, self._cached_sigma_f_sq

622

623 def _cache_si_params_iso(self, length_scales):

624 """Cache isotropic SI kernel hyperparameters."""

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

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

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

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

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

630 return self._cached_ell, self._cached_sigma_f_sq

631

632 # -------------------------------------------------------------------

633 # Bounds and Factory Methods

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

635

636 def get_bounds_from_data(self):

637 """

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

639 """

640 self.bounds = []

641 for diffs in self.differences_by_dim:

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

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

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

645

646 def create_kernel(self, kernel_name, kernel_type):

647 """

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

649

650 Parameters

651 ----------

652 kernel_name : str

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

654 kernel_type : str

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

656

657 Returns

658 -------

659 callable

660 The selected kernel function.

661 """

662 # Clear caches when creating a new kernel

663 self.clear_caches()

664

665 if not self.normalize:

666 self.get_bounds_from_data()

667

668 if kernel_type == "anisotropic":

669 kernel_func = self._create_anisotropic(kernel_name)

670 elif kernel_type == "isotropic":

671 kernel_func = self._create_isotropic(kernel_name)

672 else:

673 raise ValueError("Invalid kernel_type")

674

675 # Pre-warm the temp array cache for the training shape so the first

676 # optimizer call hits immediately instead of allocating on the fly.

677 if self.differences_by_dim is not None and len(self.differences_by_dim) > 0:

678 self._ensure_temp_arrays(self.differences_by_dim[0].shape)

679

680 return kernel_func

681

682 def _create_anisotropic(self, kernel):

683 """

684 Sets bounds and returns the anisotropic kernel function.

685

686 Returns

687 -------

688 callable

689 The anisotropic kernel function.

690 """

691 sigma_n_bound = (-16, -3)

692

693 if kernel == "SE":

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

695 return self.se_kernel_anisotropic

696 elif kernel == "RQ":

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

698 return self.rq_kernel_anisotropic

699 elif kernel == "SineExp":

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

701 return self.sine_exp_kernel_anisotropic

702 elif kernel == "Matern":

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

704 self.matern_kernel_prebuild = jetgp.utils.matern_kernel_builder(

705 self.nu, oti_module=self.oti)

706 return self.matern_kernel_anisotropic

707 elif kernel == "SI":

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

709 self.SI_kernel_prebuild = jetgp.utils.generate_bernoulli_lambda(

710 self.alpha)

711 return self.SI_kernel_anisotropic

712 else:

713 raise NotImplementedError("Anisotropic kernel not implemented")

714

715 def _create_isotropic(self, kernel):

716 """

717 Sets bounds and returns the isotropic kernel function.

718

719 Returns

720 -------

721 callable

722 The isotropic kernel function.

723 """

724 sigma_n_bound = (-16, -3)

725

726 if self.normalize:

727 core_bounds = [(-3, 3)]

728 else:

729 self.get_bounds_from_data()

730 core_bounds = [(

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

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

733 )]

734

735 if kernel == "SE":

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

737 return self.se_kernel_isotropic

738 elif kernel == "RQ":

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

740 return self.rq_kernel_isotropic

741 elif kernel == "SineExp":

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

743 return self.sine_exp_kernel_isotropic

744 elif kernel == "Matern":

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

746 self.matern_kernel_prebuild = jetgp.utils.matern_kernel_builder(

747 self.nu, oti_module=self.oti)

748 return self.matern_kernel_isotropic

749 elif kernel == "SI":

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

751 self.SI_kernel_prebuild = jetgp.utils.generate_bernoulli_lambda(

752 self.alpha)

753 return self.SI_kernel_isotropic

754 else:

755 raise NotImplementedError("Isotropic kernel not implemented")

756

757 def _add_bounds(self, extra_bounds):

758 """

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

760

761 Parameters

762 ----------

763 extra_bounds : list of tuple

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

765 """

766 if self.normalize:

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

768 else:

769 self.bounds += extra_bounds

770

771 # -------------------------------------------------------------------

772 # Anisotropic Kernel Implementations with Caching

773 # -------------------------------------------------------------------

774

775 @profile

776 def se_kernel_anisotropic(self, differences_by_dim, length_scales):

777 """

778 Anisotropic Squared Exponential (SE) kernel with caching.

779

780 Parameters

781 ----------

782 differences_by_dim : list of ndarray

783 Pairwise differences by dimension.

784 length_scales : list

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

786

787 Returns

788 -------

789 ndarray

790 Kernel matrix values.

791 """

792 ell, sigma_f_sq = self._cache_se_params_aniso(length_scales)

793 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

794 self._compute_neg_half_sqdist_aniso(differences_by_dim, ell, sqdist, tmp1, tmp2)

795

796 # TODO (performance): oti.exp dominates kernel cost (~48% of kernel time).

797 # The full N×N OTI phi array is computed, but the final kernel matrix K is

798 # symmetric — only the upper triangle is needed. For DEGP, phi[j,i] equals

799 # phi[i,j] with first-order OTI epsilon components negated (second-order

800 # unchanged), so a C-level `fused_mirror` op could reconstruct the lower

801 # triangle from the upper without re-evaluating oti.exp, halving its cost.

802 # For GDDEGP, odd/even OTI basis tags are fully swapped between phi[i,j] and

803 # phi[j,i], requiring an odd↔even component remapping — more complex but

804 # follows the same principle. Both require otilib C/Cython changes.

805 self.oti.exp(sqdist, out=tmp1)

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

807 return tmp2

808

809 def rq_kernel_anisotropic(self, differences_by_dim, length_scales):

810 """

811 Anisotropic Rational Quadratic (RQ) kernel with caching.

812

813 Parameters

814 ----------

815 differences_by_dim : list of ndarray

816 Pairwise differences by dimension.

817 length_scales : list

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

819

820 Returns

821 -------

822 ndarray

823 Kernel matrix values.

824 """

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

826 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

827

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

829 inv_2alpha = 1.0 / (2 * alpha)

830 self._compute_scaled_sqdist_aniso(differences_by_dim, ell, inv_2alpha, sqdist, tmp1, tmp2)

831 self.oti.sum(1.0, sqdist, out=tmp1)

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

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

834 return tmp1

835

836 def sine_exp_kernel_anisotropic(self, differences_by_dim, length_scales):

837 """

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

839

840 Parameters

841 ----------

842 differences_by_dim : list of ndarray

843 Pairwise differences by dimension.

844 length_scales : list

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

846

847 Returns

848 -------

849 ndarray

850 Kernel matrix values.

851 """

852 ell, pi_over_p, sigma_f_sq = self._cache_sine_exp_params_aniso(

853 length_scales)

854 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

855 self._reset_sqdist()

856

857 for i in range(self.dim):

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

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

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

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

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

863

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

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

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

867 return tmp1

868

869 def matern_kernel_anisotropic(self, differences_by_dim, length_scales):

870 """

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

872

873 Parameters

874 ----------

875 differences_by_dim : list of ndarray

876 Pairwise differences by dimension.

877 length_scales : list

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

879

880 Returns

881 -------

882 ndarray

883 Kernel matrix values.

884 """

885 ell, sigma_f_sq = self._cache_matern_params_aniso(length_scales)

886 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

887

888 # Compute scaled squared distance using fused op if available

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

890

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

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

893 # correct OTI derivative structure for the covariance blocks.

894 _eps = 1e-10

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

896 return sigma_f_sq * self.matern_kernel_prebuild(r)

897

898 def SI_kernel_anisotropic(self, differences_by_dim, length_scales):

899 """

900 Anisotropic SI kernel with caching.

901

902 Parameters

903 ----------

904 differences_by_dim : list of ndarray

905 Pairwise differences by dimension.

906 length_scales : list

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

908

909 Returns

910 -------

911 ndarray

912 Kernel matrix values.

913 """

914 ell, sigma_f_sq = self._cache_si_params_aniso(length_scales)

915

916 val = 1

917 for i in range(self.dim):

918 val = val * \

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

920 return sigma_f_sq * val

921

922 # -------------------------------------------------------------------

923 # Isotropic Kernel Implementations with Caching

924 # -------------------------------------------------------------------

925

926 def se_kernel_isotropic(self, differences_by_dim, length_scales):

927 """

928 Isotropic Squared Exponential (SE) kernel with caching.

929

930 Parameters

931 ----------

932 differences_by_dim : list of ndarray

933 Pairwise differences by dimension.

934 length_scales : list

935 Hyperparameters: [ell, sigma_f]

936

937 Returns

938 -------

939 ndarray

940 Kernel matrix values.

941 """

942 ell, sigma_f_sq = self._cache_se_params_iso(length_scales)

943 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

944 self._compute_neg_half_sqdist_iso(differences_by_dim, ell, sqdist, tmp1, tmp2)

945

946 self.oti.exp(sqdist, out=tmp1)

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

948 return tmp2

949

950 def rq_kernel_isotropic(self, differences_by_dim, length_scales):

951 """

952 Isotropic Rational Quadratic (RQ) kernel with caching.

953

954 Parameters

955 ----------

956 differences_by_dim : list of ndarray

957 Pairwise differences by dimension.

958 length_scales : list

959 Hyperparameters: [ell, alpha, sigma_f]

960

961 Returns

962 -------

963 ndarray

964 Kernel matrix values.

965 """

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

967 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

968

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

970 inv_2alpha = 1.0 / (2 * alpha)

971 self._compute_scaled_sqdist_iso(differences_by_dim, ell, inv_2alpha, sqdist, tmp1, tmp2)

972 self.oti.sum(1.0, sqdist, out=tmp1)

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

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

975 return tmp1

976

977 def sine_exp_kernel_isotropic(self, differences_by_dim, length_scales):

978 """

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

980

981 Parameters

982 ----------

983 differences_by_dim : list of ndarray

984 Pairwise differences by dimension.

985 length_scales : list

986 Hyperparameters: [ell, p, sigma_f]

987

988 Returns

989 -------

990 ndarray

991 Kernel matrix values.

992 """

993 ell, pi_over_p, sigma_f_sq = self._cache_sine_exp_params_iso(

994 length_scales)

995 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

996 self._reset_sqdist()

997

998 for i in range(self.dim):

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

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

1001 self.oti.mul(ell, tmp2, out=tmp1)

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

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

1004

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

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

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

1008 return tmp1

1009

1010 def matern_kernel_isotropic(self, differences_by_dim, length_scales):

1011 """

1012 Isotropic Matérn kernel (half-integer ν) with caching.

1013

1014 Parameters

1015 ----------

1016 differences_by_dim : list of ndarray

1017 Pairwise differences by dimension.

1018 length_scales : list

1019 Hyperparameters: [ell, sigma_f]

1020

1021 Returns

1022 -------

1023 ndarray

1024 Kernel matrix values.

1025 """

1026 ell, sigma_f_sq = self._cache_matern_params_iso(length_scales)

1027 tmp1, tmp2, sqdist = self._ensure_temp_arrays(differences_by_dim[0].shape)

1028

1029 # Compute scaled squared distance using fused op if available

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

1031

1032 # r = sqrt(sqdist + eps²) — regularise r directly

1033 _eps = 1e-10

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

1035 return sigma_f_sq * self.matern_kernel_prebuild(r)

1036

1037 def SI_kernel_isotropic(self, differences_by_dim, length_scales):

1038 """

1039 Isotropic SI kernel with caching.

1040

1041 Parameters

1042 ----------

1043 differences_by_dim : list of ndarray

1044 Pairwise differences by dimension.

1045 length_scales : list

1046 Hyperparameters: [ell, sigma_f]

1047

1048 Returns

1049 -------

1050 ndarray

1051 Kernel matrix values.

1052 """

1053 ell, sigma_f_sq = self._cache_si_params_iso(length_scales)

1054

1055 val = 1

1056 for i in range(self.dim):

1057 val = val * \

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

1059 return sigma_f_sq * val

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

504 statements