updated notebooks

docs/examples/barycentres.py

@@ -26,7 +26,6 @@
 import jax.scipy.linalg as jsl
 from jaxtyping import install_import_hook
 import matplotlib.pyplot as plt
-import optax as ox
 import jaxopt
 import tensorflow_probability.substrates.jax.distributions as tfd
 
@@ -137,13 +136,10 @@ def fit_gp(x: jax.Array, y: jax.Array) -> tfd.MultivariateNormalFullCovariance:
 
     likelihood = gpx.Gaussian(num_datapoints=n)
     posterior = gpx.Prior(mean_function=gpx.Constant(), kernel=gpx.RBF()) * likelihood
-
     opt_posterior, _ = gpx.fit(
         model=posterior,
         train_data=D,
-        solver=jaxopt.OptaxSolver(
-            gpx.ConjugateMLL(negative=True), opt=ox.adam(0.01), maxiter=500
-        ),
+        solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
         key=key,
     )
     latent_dist = opt_posterior.predict(xtest, train_data=D)

docs/examples/bayesian_optimisation.py

@@ -19,7 +19,6 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 from matplotlib import cm
-import optax as ox
 import jaxopt
 import tensorflow_probability.substrates.jax as tfp
 from typing import List, Tuple
@@ -218,9 +217,7 @@ def return_optimised_posterior(
     opt_posterior, history = gpx.fit(
         model=posterior,
         train_data=D,
-        solver=jaxopt.OptaxSolver(
-            gpx.ConjugateMLL(negative=True), opt=ox.adam(0.01), maxiter=1000
-        ),
+        solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
         safe=True,
         key=key,
         verbose=False,

docs/examples/classification.py

@@ -42,7 +42,6 @@
     install_import_hook,
 )
 import matplotlib.pyplot as plt
-import optax as ox
 import jaxopt
 import tensorflow_probability.substrates.jax as tfp
 from tqdm import trange
@@ -122,7 +121,7 @@
 opt_posterior, history = gpx.fit(
     model=posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(negative_lpd, opt=ox.adam(0.01), maxiter=1000),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
     key=key,
 )
 

docs/examples/constructing_new_kernels.py

@@ -272,9 +272,7 @@ def __call__(
 opt_posterior, history = gpx.fit(
     model=circular_posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(
-        gpx.ConjugateMLL(negative=True), opt=ox.adamw(0.05), maxiter=500
-    ),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
     key=key,
 )
 

docs/examples/decision_making.py

@@ -28,7 +28,6 @@
 import jax.random as jr
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-import optax as ox
 
 import gpjax as gpx
 from gpjax.decision_making.utility_functions import (
@@ -164,17 +163,14 @@ def forrester(x: Float[Array, "N 1"]) -> Float[Array, "N 1"]:
 # defined above. We tend to also optimise the hyperparameters of the GP prior when
 # "fitting" our GP, as demonstrated in the [Regression
 # notebook](https://docs.jaxgaussianprocesses.com/examples/regression/). This will be
-# using the GPJax `fit` method under the hood, which requires an `optimization_objective`,
-# `optimizer` and `num_optimization_iters`. Therefore, we also pass these to the
-# `PosteriorHandler` as demonstrated below:
+# using the GPJax `fit` method under the hood, which requires an jaxopt `solver`.
+# Therefore, we also pass this to the `PosteriorHandler` as demonstrated below:
 
 # %%
 posterior_handler = PosteriorHandler(
     prior,
     likelihood_builder=likelihood_builder,
-    optimization_objective=gpx.ConjugateMLL(negative=True),
-    optimizer=ox.adam(learning_rate=0.01),
-    num_optimization_iters=1000,
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
 )
 posterior_handlers = {OBJECTIVE: posterior_handler}
 

docs/examples/graph_kernels.py

@@ -22,7 +22,6 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import networkx as nx
-import optax as ox
 import jaxopt
 
 with install_import_hook("gpjax", "beartype.beartype"):
@@ -157,9 +156,7 @@
 opt_posterior, training_history = gpx.fit(
     model=posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(
-        gpx.ConjugateMLL(negative=True), opt=ox.adamw(0.01), maxiter=1000
-    ),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
     key=key,
 )
 

docs/examples/intro_to_kernels.py

@@ -16,7 +16,6 @@
 from jaxtyping import install_import_hook, Float
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-import optax as ox
 import jaxopt
 import pandas as pd
 from docs.examples.utils import clean_legend
@@ -238,9 +237,7 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:
 opt_posterior, history = gpx.fit(
     model=no_opt_posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(
-        gpx.ConjugateMLL(negative=True), opt=ox.adamw(0.01), maxiter=2000
-    ),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
     safe=True,
     key=key,
 )
@@ -540,9 +537,7 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:
 opt_posterior, history = gpx.fit(
     model=posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(
-        gpx.ConjugateMLL(negative=True), opt=ox.adamw(0.01), maxiter=1000
-    ),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
     safe=True,
     key=key,
 )

docs/examples/oceanmodelling.py

@@ -24,7 +24,6 @@
 from matplotlib import rcParams
 import matplotlib.pyplot as plt
 import jaxopt
-import optax as ox
 import pandas as pd
 import tensorflow_probability as tfp
 
@@ -251,7 +250,7 @@ def optimise_mll(posterior, dataset, NIters=1000, key=key, plot_history=True):
     opt_posterior, history = gpx.fit(
         model=posterior,
         train_data=dataset,
-        solver=jaxopt.OptaxSolver(objective, opt=ox.adam(0.1), maxiter=NIters),
+        solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=NIters),
         safe=True,
         key=key,
     )

docs/examples/regression.py

@@ -31,7 +31,6 @@
 from jaxtyping import install_import_hook
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-import optax as ox
 import jaxopt
 from docs.examples.utils import clean_legend
 
@@ -218,7 +217,7 @@
 opt_posterior, history = gpx.fit(
     model=posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(negative_mll, opt=ox.adamw(0.01), maxiter=500),
+    solver=jaxopt.LBFGS(negative_mll, maxiter=500),
     safe=True,
     key=key,
 )

docs/examples/spatial.py

@@ -54,7 +54,6 @@
 )
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-import optax as ox
 import jaxopt
 import pandas as pd
 import planetary_computer
@@ -191,7 +190,7 @@ def __call__(self, x: Float[Array, "N D"]) -> Float[Array, "N 1"]:
 posterior, history = gpx.fit(
     model=posterior,
     train_data=D,
-    solver=jaxopt.OptaxSolver(negative_mll, opt=optim, maxiter=3000),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=1000),
     safe=True,
     key=key,
 )

docs/examples/yacht.py

@@ -35,7 +35,6 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
-import optax as ox
 import jaxopt
 import pandas as pd
 from sklearn.metrics import (
@@ -195,7 +194,7 @@
 opt_posterior, history = gpx.fit(
     model=posterior,
     train_data=training_data,
-    solver=jaxopt.OptaxSolver(negative_mll, opt=ox.adamw(0.05), maxiter=500),
+    solver=jaxopt.LBFGS(gpx.ConjugateMLL(negative=True), maxiter=500),
     key=key,
 )