IFE Surrogate GP

A flexible and extensible library for Gaussian Processes, built for performance and modularity.
Models can be trained via maximum likelihood estimation or by baysian inference.

Features

• High-performance kernels (JAX-compatible)
• Composable API for building custom models
• Multiple optimizers (Optax, Scipy, etc.)
• Built in Baysian inference of the models with NumPyro
• Automatic hyperparameter handling
• Built-in training workflows

Installation

pip install ife_surrogate

Usage

Quickstart

  from ife_surrogate.gp.kernels import Kriging
  from ife_surrogate.gp.models import WidebandGP
  from ife_surrogate.gp.trainers import SwarmTrainer
  
  from ife_surrogate.utils.datasets import load_lc_filter_dataset

  from numpyro.distributions import Uniform
  import jax.numpy as jnp
  import matplotlib.pyplot as plt


  data = load_lc_filter_dataset()
  x_train_scaled = data["x_train_scaled"]
  x_test_scaled = data["x_test_scaled"]
  y_train_scaled = data["y_train_scaled"]
  y_test_scaled = data["y_test_scaled"]

  freqs = data["freqs"]


  d = x_train_scaled.shape[1]
  priors = {"lengthscale": Uniform(1e-0, 1e1), "power": Uniform(1, 2)}
  kernel = Kriging(lengthscale=jnp.ones(d), power=jnp.ones(d), priors=priors)

  model = WidebandGP(x_train_scaled, y_train_scaled, kernel, freqs)

  trainer = SwarmTrainer(number_iterations=100, number_particles=10)
  best_run, history = trainer.train(model)

  predictions, variances = model.predict(x_test_scaled)


  fig, ax = plt.subplots(figsize=(10, 5))
  plt.xscale("log")
  plt.scatter(freqs, predictions[0], label="predicted")
  plt.scatter(freqs, y_test_scaled[0], label="true")
  plt.legend()
  plt.grid()
  plt.show()

Documentation

• API Reference
• [Tutorials] (Coming soon)

Key Components

• Kernels

  • Kriging
  • RBF
  • Matern
  • SumKernel
  • ProductKernel
  • Scale
  • RQ
  • Noise

• Models

  • WidebandGP
  • ScalerGP
  • WidebandTP
  • ScalerTP

• Trainers

  • OptaxTrainer
  • SwarmTrainer

License

Distributed under the MIT License. See LICENSE for more information.

References

The mathematical background and implementation of these models are based on the following publications:

Gaussian Process & Student-t Theory

• Gaussian Process Basics: Rasmussen, C. E., & Williams, C. K. I. (2006). Gaussian Processes for Machine Learning. MIT Press.
• Student-t Process (TP) Foundations: Shah, A., Wilson, A. G., & Ghahramani, Z. (2014). Student-t Processes as Alternatives to Gaussian Processes. Proceedings of the 17th International Conference on Artificial Intelligence and Statistics (AISTATS).

Wideband & Multi-output Modeling

• Wideband Architecture: Rezende, R. S., Hansen, J., Piwonski, A., & Schuhmann, R. (2024). Wideband Kriging for Multiobjective Optimization of a High-Voltage EMI Filter. IEEE Transactions on Electromagnetic Compatibility, 66(4), 1116–1124.
• Multi-output Separable GPs: Bilionis, I., Zabaras, N., Konomi, B. A., & Lin, G. (2013). Multi-output separable Gaussian process: Towards an efficient, fully Bayesian paradigm for uncertainty quantification. Journal of Computational Physics, 241, 212–239.
• Vector-Valued Kernels: Alvarez, M. A., Rosasco, L., & Lawrence, N. D. (2012). Kernels for Vector-Valued Functions: A Review. [cite_start]Foundations and Trends in Machine Learning.

Inference & Software Stack

• JAX: Bradbury, J., et al. (2018). JAX: Composable transformations of Python+NumPy programs.
• NumPyro: Phan, D., Pradhan, N., & Jankowiak, M. (2019). Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro. arXiv preprint.
• Flax: Heek, J., et al. (2020). Flax: A neural network library and ecosystem for JAX.
• Optax: Babuschkin, I., et al. (2020). Optax: Composable gradient descent optimization for JAX.
• SciPy: Virtanen, P., et al. (2020). SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods.
• NUTS Sampler: Hoffman, M. D., & Gelman, A. (2014). The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research.