Trainers

class ife_surrogate.gp.trainers.Trainer(sample_parameters=True, verbose=False, save_history=False)

Bases: ABC

Parameters:

  • sample_parameters (bool)

  • verbose (bool)

  • save_history (bool)

property model
loss_fn(updated_params)
Parameters:

updated_params (Dict)

Return type:

float

abstract train(*args, **kwargs)

Run optimization and return (best_run, history).

Return type:

Tuple[Dict, Dict | None]

class ife_surrogate.gp.trainers.OptaxTrainer(key=Array((), dtype=key<fry>) overlaying: [0 0], optimizer=(<function chain.<locals>.init_fn>, <function chain.<locals>.update_fn>), number_iterations=100, number_restarts=1, tolerance=0.01, patience=20, **kwargs)

Bases: Trainer

Trainer for Gaussian Process models using gradient-based optimization with Optax[1].

This class leverages JAX and Optax to optimize model parameters using gradient descent, supporting multiple restarts and early stopping.

Parameters:

  • key (Key)

  • number_iterations (int)

  • number_restarts (int)

  • tolerance (Float)

  • patience (Float)

key

JAX random key for parameter sampling and reproducibility.

Type:

Key

optimizer

An Optax optimizer instance (e.g., optax.adam).[2]

number_iterations

Maximum number of iterations per training run.

Type:

int

number_restarts

Number of independent training restarts.

Type:

int

tolerance

Minimum improvement threshold to reset early stopping.

Type:

float

patience

Number of iterations to wait without improvement before stopping.

Type:

int

verbose

Whether to print iteration progress.

Type:

bool

save_history

Whether to return the full optimization history.

Type:

bool

sample_parameters

Whether to sample kernel hyperparameters at each restart.

Type:

bool

train(model

GPModel) -> Tuple[Dict, Optional[Dict]]: Trains the given GP model and returns the best parameters and optional history.

info()

Prints the current trainer settings for easy reference.

[1] “https://optax.readthedocs.io/en/latest/” [2] “https://optax.readthedocs.io/en/latest/api/optimizers.html”

info()

Prints current settings of the OptaxTrainer.

train(model)

Run optimization and return (best_run, history).

Parameters:

model (GPModel)

Return type:

Tuple[Dict, Dict | None]

class ife_surrogate.gp.trainers.SwarmTrainer(swarm_settings={'c1': 1.5, 'c2': 1.5, 'w': 0.5}, number_iterations=100, number_particles=20, number_restarts=1, bounds=None, **kwargs)

Bases: Trainer

Trainer class that optimizes model parameters using Particle Swarm Optimization (PSO) via pyswarms[1]

This class extends a generic Trainer and implements a global optimization method for training models, particularly useful when gradients are unavailable or the loss landscape is highly non-convex.

Parameters:

  • swarm_settings (dict, optional) – Dictionary of PSO hyperparameters: - c1 (float): Cognitive parameter (default: 0.5) - c2 (float): Social parameter (default: 0.3) - w (float): Inertia weight (default: 0.9)

  • number_iterations (int, optional) – Number of iterations for the PSO algorithm per restart (default: 100).

  • number_particles (int, optional) – Number of particles in the swarm (default: 20).

  • number_restarts (int, optional) – Number of independent PSO runs to avoid local minima (default: 1).

  • (dict[str (bounds) – (float, float)], optional): Lower and upper bounds for each parameter. Not defined parameter bounds are automatically inferred from kernel priors.

  • **kwargs – Additional keyword arguments passed to the base Trainer class.

  • bounds (dict)

swarm_settings

Hyperparameters controlling particle behavior in PSO.

Type:

dict

number_iterations

Iterations per PSO run.

Type:

int

number_particles

Number of particles in the swarm.

Type:

int

number_restarts

Number of independent PSO runs.

Type:

int

bounds (dict[str

(float, float)]): Parameter bounds.

model

The model to be optimized.

verbose

Flag to control printing of optimization progress (inherited from Trainer).

Type:

bool

info()

Prints current PSO settings and bounds.

train(model)

Optimizes the provided model’s kernel parameters using PSO and returns the best parameters and optionally the optimization history.

Example

>>> trainer = SwarmTrainer(number_particles=30, number_iterations=200)
>>> best_run, history = trainer.train(model)

[1] ‘https://github.com/ljvmiranda921/pyswarms’

info()
train(model)

Run optimization and return (best_run, history).