fit_scipy#

gpjax.fit.fit_scipy(*, model, objective, train_data, max_iters=500, verbose=True, safe=True)[source]#

Train a Module model with respect to a supplied Objective function using SciPy’s L-BFGS-B optimiser.

Parameters are transformed to unconstrained space, flattened into a single vector, and passed to scipy.optimize.minimize. Gradients are computed via JAX’s value_and_grad.

Parameters:

  • model (Module) – The model to be optimised.

  • objective (Objective) – The objective function to minimise with respect to the model parameters.

  • train_data (Dataset) – The training data used to evaluate the objective.

  • max_iters (int) – Maximum number of L-BFGS-B iterations. Defaults to 500.

  • verbose (bool) – Whether to print optimisation progress. Defaults to True.

  • safe (bool) – Whether to validate inputs before optimisation. Defaults to True.

Returns:

A tuple of the optimised model and an array of

objective values recorded at each iteration.

Return type:

tuple[Module, Array]

Example

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import gpjax as gpx
>>> import jax.numpy as jnp

>>> xtrain = jnp.linspace(0, 1).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)

>>> meanf = gpx.mean_functions.Constant()
>>> kernel = gpx.kernels.RBF()
>>> likelihood = gpx.likelihoods.Gaussian()
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
>>> posterior = prior * likelihood

>>> nmll = lambda p, d: -gpx.objectives.conjugate_mll(p, d)
>>> trained_model, history = gpx.fit_scipy(
...     model=posterior, objective=nmll, train_data=D
... )