gpjax 0.11.0__tar.gz → 0.11.1__tar.gz

{gpjax-0.11.0 → gpjax-0.11.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.11.0
+Version: 0.11.1
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues

{gpjax-0.11.0 → gpjax-0.11.1}/examples/constructing_new_kernels.py

@@ -33,7 +33,6 @@ from jaxtyping import (
     install_import_hook,
 )
 import matplotlib.pyplot as plt
-import numpyro.distributions as npd
 from numpyro.distributions import constraints
 import numpyro.distributions.transforms as npt
 
@@ -52,8 +51,6 @@ with install_import_hook("gpjax", "beartype.beartype"):
     import gpjax as gpx
 
 
-tfb = tfp.bijectors
-
 # set the default style for plotting
 use_mpl_style()
 

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/__init__.py

@@ -32,14 +32,15 @@ from gpjax.citation import cite
 from gpjax.dataset import Dataset
 from gpjax.fit import (
     fit,
+    fit_lbfgs,
     fit_scipy,
 )
 
 __license__ = "MIT"
-__description__ = "Didactic Gaussian processes in JAX"
+__description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.11.0"
+__version__ = "0.11.1"
 
 __all__ = [
     "base",
@@ -56,5 +57,6 @@ __all__ = [
     "fit",
     "Module",
     "param_field",
+    "fit_lbfgs",
     "fit_scipy",
 ]

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/fit.py

@@ -15,13 +15,13 @@
 
 import typing as tp
 
-from flax import nnx
 import jax
-from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random as jr
-from numpyro.distributions.transforms import Transform
 import optax as ox
+from flax import nnx
+from jax.flatten_util import ravel_pytree
+from numpyro.distributions.transforms import Transform
 from scipy.optimize import minimize
 
 from gpjax.dataset import Dataset
@@ -127,7 +127,6 @@ def fit(  # noqa: PLR0913
         _check_verbose(verbose)
 
     # Model state filtering
-
     graphdef, params, *static_state = nnx.split(model, Parameter, ...)
 
     # Parameters bijection to unconstrained space
@@ -253,6 +252,110 @@ def fit_scipy(  # noqa: PLR0913
     return model, history
 
 
+def fit_lbfgs(
+    *,
+    model: Model,
+    objective: Objective,
+    train_data: Dataset,
+    params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
+    max_iters: int = 100,
+    safe: bool = True,
+    max_linesearch_steps: int = 32,
+    gtol: float = 1e-5,
+) -> tuple[Model, jax.Array]:
+    r"""Train a Module model with respect to a supplied Objective function.
+
+    Uses Optax's LBFGS implementation and a jax.lax.while loop.
+
+     Args:
+         model: the model Module to be optimised.
+         objective: The objective function that we are optimising with
+             respect to.
+         train_data (Dataset): The training data to be used for the optimisation.
+         max_iters (int): The maximum number of optimisation steps to run. Defaults
+             to 500.
+         safe (bool): Whether to check the types of the inputs.
+         max_linesearch_steps (int): The maximum number of linesearch steps to use
+            for finding the stepsize.
+        gtol (float): Terminate the optimisation if the L2 norm of the gradient is
+            below this threshold.
+
+     Returns:
+         A tuple comprising the optimised model and final loss.
+    """
+    if safe:
+        # Check inputs
+        _check_model(model)
+        _check_train_data(train_data)
+        _check_num_iters(max_iters)
+
+    # Model state filtering
+    graphdef, params, *static_state = nnx.split(model, Parameter, ...)
+
+    # Parameters bijection to unconstrained space
+    if params_bijection is not None:
+        params = transform(params, params_bijection, inverse=True)
+
+    # Loss definition
+    def loss(params: nnx.State) -> ScalarFloat:
+        params = transform(params, params_bijection)
+        model = nnx.merge(graphdef, params, *static_state)
+        return objective(model, train_data)
+
+    # Initialise optimiser
+    optim = ox.lbfgs(
+        linesearch=ox.scale_by_zoom_linesearch(
+            max_linesearch_steps=max_linesearch_steps,
+            initial_guess_strategy="one",
+        )
+    )
+    opt_state = optim.init(params)
+    loss_value_and_grad = ox.value_and_grad_from_state(loss)
+
+    # Optimisation step.
+    def step(carry):
+        params, opt_state = carry
+
+        # Using optax's value_and_grad_from_state is more efficient given LBFGS uses a linesearch
+        # See https://optax.readthedocs.io/en/latest/api/utilities.html#optax.value_and_grad_from_state
+        loss_val, loss_gradient = loss_value_and_grad(params, state=opt_state)
+        updates, opt_state = optim.update(
+            loss_gradient,
+            opt_state,
+            params,
+            value=loss_val,
+            grad=loss_gradient,
+            value_fn=loss,
+        )
+        params = ox.apply_updates(params, updates)
+
+        return params, opt_state
+
+    def continue_fn(carry):
+        _, opt_state = carry
+        n = ox.tree_utils.tree_get(opt_state, "count")
+        g = ox.tree_utils.tree_get(opt_state, "grad")
+        g_l2_norm = ox.tree_utils.tree_l2_norm(g)
+        return (n == 0) | ((n < max_iters) & (g_l2_norm >= gtol))
+
+    # Optimisation loop
+    params, opt_state = jax.lax.while_loop(
+        continue_fn,
+        step,
+        (params, opt_state),
+    )
+    final_loss = ox.tree_utils.tree_get(opt_state, "value")
+
+    # Parameters bijection to constrained space
+    if params_bijection is not None:
+        params = transform(params, params_bijection)
+
+    # Reconstruct model
+    model = nnx.merge(graphdef, params, *static_state)
+
+    return model, final_loss
+
+
 def get_batch(train_data: Dataset, batch_size: int, key: KeyArray) -> Dataset:
     """Batch the data into mini-batches. Sampling is done with replacement.
 

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/kernels/nonstationary/arccosine.py

@@ -23,7 +23,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -91,9 +94,9 @@ class ArcCosine(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
         if isinstance(bias_variance, nnx.Variable):
             self.bias_variance = bias_variance

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/kernels/nonstationary/linear.py

@@ -23,7 +23,7 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import NonNegativeReal
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -64,9 +64,9 @@ class Linear(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
     def __call__(
         self,

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/kernels/nonstationary/polynomial.py

@@ -23,7 +23,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -76,9 +79,9 @@ class Polynomial(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarArray], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
 
         self.name = f"Polynomial (degree {self.degree})"
 

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/kernels/stationary/base.py

@@ -25,7 +25,10 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
-from gpjax.parameters import PositiveReal
+from gpjax.parameters import (
+    NonNegativeReal,
+    PositiveReal,
+)
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -85,11 +88,11 @@ class StationaryKernel(AbstractKernel):
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
-            self.variance = PositiveReal(variance)
+            self.variance = NonNegativeReal(variance)
 
             # static typing
             if tp.TYPE_CHECKING:
-                self.variance = tp.cast(PositiveReal[ScalarFloat], self.variance)
+                self.variance = tp.cast(NonNegativeReal[ScalarFloat], self.variance)
 
     @property
     def spectral_density(self) -> npd.Normal | npd.StudentT:

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/likelihoods.py

@@ -28,7 +28,7 @@ from gpjax.integrators import (
     GHQuadratureIntegrator,
 )
 from gpjax.parameters import (
-    PositiveReal,
+    NonNegativeReal,
     Static,
 )
 from gpjax.typing import (
@@ -134,7 +134,7 @@ class Gaussian(AbstractLikelihood):
         self,
         num_datapoints: int,
         obs_stddev: tp.Union[
-            ScalarFloat, Float[Array, "#N"], PositiveReal, Static
+            ScalarFloat, Float[Array, "#N"], NonNegativeReal, Static
         ] = 1.0,
         integrator: AbstractIntegrator = AnalyticalGaussianIntegrator(),
     ):
@@ -148,8 +148,8 @@ class Gaussian(AbstractLikelihood):
                 likelihoods. Must be an instance of `AbstractIntegrator`. For the Gaussian likelihood, this defaults to
                 the `AnalyticalGaussianIntegrator`, as the expected log likelihood can be computed analytically.
         """
-        if not isinstance(obs_stddev, (PositiveReal, Static)):
-            obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        if not isinstance(obs_stddev, (NonNegativeReal, Static)):
+            obs_stddev = NonNegativeReal(jnp.asarray(obs_stddev))
         self.obs_stddev = obs_stddev
 
         super().__init__(num_datapoints, integrator)

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/mean_functions.py

@@ -207,5 +207,5 @@ SumMeanFunction = ft.partial(
     CombinationMeanFunction, operator=ft.partial(jnp.sum, axis=0)
 )
 ProductMeanFunction = ft.partial(
-    CombinationMeanFunction, operator=ft.partial(jnp.sum, axis=0)
+    CombinationMeanFunction, operator=ft.partial(jnp.prod, axis=0)
 )

{gpjax-0.11.0 → gpjax-0.11.1}/gpjax/parameters.py

@@ -82,6 +82,14 @@ class Parameter(nnx.Variable[T]):
         self._tag = tag
 
 
+class NonNegativeReal(Parameter[T]):
+    """Parameter that is non-negative."""
+
+    def __init__(self, value: T, tag: ParameterTag = "non_negative", **kwargs):
+        super().__init__(value=value, tag=tag, **kwargs)
+        _safe_assert(_check_is_non_negative, self.value)
+
+
 class PositiveReal(Parameter[T]):
     """Parameter that is strictly positive."""
 
@@ -143,6 +151,7 @@ class LowerTriangular(Parameter[T]):
 
 DEFAULT_BIJECTION = {
     "positive": npt.SoftplusTransform(),
+    "non_negative": npt.SoftplusTransform(),
     "real": npt.IdentityTransform(),
     "sigmoid": npt.SigmoidTransform(),
     "lower_triangular": FillTriangularTransform(),
@@ -164,6 +173,13 @@ def _check_is_arraylike(value: T) -> None:
         )
 
 
+@checkify.checkify
+def _check_is_non_negative(value):
+    checkify.check(
+        jnp.all(value >= 0), "value needs to be non-negative, got {value}", value=value
+    )
+
+
 @checkify.checkify
 def _check_is_positive(value):
     checkify.check(

{gpjax-0.11.0 → gpjax-0.11.1}/tests/test_fit.py

@@ -13,20 +13,24 @@
 # limitations under the License.
 # ==============================================================================
 
-from flax import nnx
 import jax.numpy as jnp
 import jax.random as jr
-from jaxtyping import (
-    Float,
-    Num,
-)
 import optax as ox
 import pytest
 import scipy
-
+from beartype.typing import Any
+from flax import nnx
 from gpjax.dataset import Dataset
 from gpjax.fit import (
+    _check_batch_size,
+    _check_log_rate,
+    _check_model,
+    _check_num_iters,
+    _check_optim,
+    _check_train_data,
+    _check_verbose,
     fit,
+    fit_lbfgs,
     fit_scipy,
     get_batch,
 )
@@ -50,6 +54,10 @@ from gpjax.parameters import (
 )
 from gpjax.typing import Array
 from gpjax.variational_families import VariationalGaussian
+from jaxtyping import (
+    Float,
+    Num,
+)
 
 
 def test_fit_simple() -> None:
@@ -141,6 +149,46 @@ def test_fit_scipy_simple():
     assert trained_model.bias.value == 1.0
 
 
+def test_fit_lbfgs_simple():
+    # Create dataset:
+    X = jnp.linspace(0.0, 10.0, 100).reshape(-1, 1)
+    y = 2.0 * X + 1.0 + 10 * jr.normal(jr.PRNGKey(0), X.shape).reshape(-1, 1)
+    D = Dataset(X, y)
+
+    # Define linear model:
+    class LinearModel(nnx.Module):
+        def __init__(self, weight: float, bias: float):
+            self.weight = PositiveReal(weight)
+            self.bias = Static(bias)
+
+        def __call__(self, x):
+            return self.weight.value * x + self.bias.value
+
+    model = LinearModel(weight=1.0, bias=1.0)
+
+    # Define loss function:
+    def mse(model, data):
+        pred = model(data.X)
+        return jnp.mean((pred - data.y) ** 2)
+
+    # Train with bfgs!
+    trained_model, final_loss = fit_lbfgs(
+        model=model,
+        objective=mse,
+        train_data=D,
+        max_iters=10,
+    )
+
+    # Ensure we return a model of the same class
+    assert isinstance(trained_model, LinearModel)
+
+    # Test reduction in loss:
+    assert mse(trained_model, D) < mse(model, D)
+
+    # Test stop_gradient on bias:
+    assert trained_model.bias.value == 1.0
+
+
 @pytest.mark.parametrize("n_data", [20])
 @pytest.mark.parametrize("verbose", [True, False])
 def test_fit_gp_regression(n_data: int, verbose: bool) -> None:
@@ -179,8 +227,7 @@ def test_fit_gp_regression(n_data: int, verbose: bool) -> None:
 
 
 @pytest.mark.parametrize("n_data", [20])
-@pytest.mark.parametrize("verbose", [True, False])
-def test_fit_scipy_gp_regression(n_data: int, verbose: bool) -> None:
+def test_fit_lbfgs_gp_regression(n_data: int) -> None:
     # Create dataset:
     key = jr.PRNGKey(123)
     x = jnp.sort(
@@ -195,20 +242,16 @@ def test_fit_scipy_gp_regression(n_data: int, verbose: bool) -> None:
     posterior = prior * likelihood
 
     # Train with BFGS!
-    trained_model_bfgs, history_bfgs = fit_scipy(
+    trained_model_bfgs, final_loss = fit_lbfgs(
         model=posterior,
         objective=conjugate_mll,
         train_data=D,
         max_iters=40,
-        verbose=verbose,
     )
 
     # Ensure the trained model is a Gaussian process posterior
     assert isinstance(trained_model_bfgs, ConjugatePosterior)
 
-    # Ensure we return a history_bfgs of the correct length
-    assert len(history_bfgs) > 2
-
     # Ensure we reduce the loss
     assert conjugate_mll(trained_model_bfgs, D) < conjugate_mll(posterior, D)
 
@@ -324,3 +367,142 @@ def test_get_batch(n_data: int, n_dim: int, batch_size: int):
     assert New.y.shape[1:] == y.shape[1:]
     assert jnp.sum(New.X == B.X) <= n_dim * batch_size / n_data
     assert jnp.sum(New.y == B.y) <= n_dim * batch_size / n_data
+
+
+@pytest.fixture
+def valid_model() -> nnx.Module:
+    """Return a valid model for testing."""
+
+    class LinearModel(nnx.Module):
+        def __init__(self, weight: float, bias: float) -> None:
+            self.weight = PositiveReal(weight)
+            self.bias = Static(bias)
+
+        def __call__(self, x: Any) -> Any:
+            return self.weight.value * x + self.bias.value
+
+    return LinearModel(weight=1.0, bias=1.0)
+
+
+@pytest.fixture
+def valid_dataset() -> Dataset:
+    """Return a valid dataset for testing."""
+    X = jnp.array([[1.0], [2.0], [3.0]])
+    y = jnp.array([[1.0], [2.0], [3.0]])
+    return Dataset(X=X, y=y)
+
+
+def test_check_model_valid(valid_model: nnx.Module) -> None:
+    """Test that a valid model passes validation."""
+    _check_model(valid_model)
+
+
+def test_check_model_invalid() -> None:
+    """Test that an invalid model raises a TypeError."""
+    model = "not a model"
+    with pytest.raises(
+        TypeError, match="Expected model to be a subclass of nnx.Module"
+    ):
+        _check_model(model)
+
+
+def test_check_train_data_valid(valid_dataset: Dataset) -> None:
+    """Test that valid training data passes validation."""
+    _check_train_data(valid_dataset)
+
+
+def test_check_train_data_invalid() -> None:
+    """Test that invalid training data raises a TypeError."""
+    train_data = "not a dataset"
+    with pytest.raises(
+        TypeError, match="Expected train_data to be of type gpjax.Dataset"
+    ):
+        _check_train_data(train_data)
+
+
+def test_check_optim_valid() -> None:
+    """Test that a valid optimiser passes validation."""
+    optim = ox.sgd(0.1)
+    _check_optim(optim)
+
+
+def test_check_optim_invalid() -> None:
+    """Test that an invalid optimiser raises a TypeError."""
+    optim = "not an optimiser"
+    with pytest.raises(
+        TypeError, match="Expected optim to be of type optax.GradientTransformation"
+    ):
+        _check_optim(optim)
+
+
+@pytest.mark.parametrize("num_iters", [1, 10, 100])
+def test_check_num_iters_valid(num_iters: int) -> None:
+    """Test that valid number of iterations passes validation."""
+    _check_num_iters(num_iters)
+
+
+def test_check_num_iters_invalid_type() -> None:
+    """Test that an invalid num_iters type raises a TypeError."""
+    num_iters = "not an int"
+    with pytest.raises(TypeError, match="Expected num_iters to be of type int"):
+        _check_num_iters(num_iters)
+
+
+@pytest.mark.parametrize("num_iters", [0, -5])
+def test_check_num_iters_invalid_value(num_iters: int) -> None:
+    """Test that an invalid num_iters value raises a ValueError."""
+    with pytest.raises(ValueError, match="Expected num_iters to be positive"):
+        _check_num_iters(num_iters)
+
+
+@pytest.mark.parametrize("log_rate", [1, 10, 100])
+def test_check_log_rate_valid(log_rate: int) -> None:
+    """Test that a valid log rate passes validation."""
+    _check_log_rate(log_rate)
+
+
+def test_check_log_rate_invalid_type() -> None:
+    """Test that an invalid log_rate type raises a TypeError."""
+    log_rate = "not an int"
+    with pytest.raises(TypeError, match="Expected log_rate to be of type int"):
+        _check_log_rate(log_rate)
+
+
+@pytest.mark.parametrize("log_rate", [0, -5])
+def test_check_log_rate_invalid_value(log_rate: int) -> None:
+    """Test that an invalid log_rate value raises a ValueError."""
+    with pytest.raises(ValueError, match="Expected log_rate to be positive"):
+        _check_log_rate(log_rate)
+
+
+@pytest.mark.parametrize("verbose", [True, False])
+def test_check_verbose_valid(verbose: bool) -> None:
+    """Test that valid verbose values pass validation."""
+    _check_verbose(verbose)
+
+
+def test_check_verbose_invalid() -> None:
+    """Test that an invalid verbose value raises a TypeError."""
+    verbose = "not a bool"
+    with pytest.raises(TypeError, match="Expected verbose to be of type bool"):
+        _check_verbose(verbose)
+
+
+@pytest.mark.parametrize("batch_size", [1, 10, 100, -1])
+def test_check_batch_size_valid(batch_size: int) -> None:
+    """Test that valid batch sizes pass validation."""
+    _check_batch_size(batch_size)
+
+
+def test_check_batch_size_invalid_type() -> None:
+    """Test that an invalid batch_size type raises a TypeError."""
+    batch_size = "not an int"
+    with pytest.raises(TypeError, match="Expected batch_size to be of type int"):
+        _check_batch_size(batch_size)
+
+
+@pytest.mark.parametrize("batch_size", [0, -2, -5])
+def test_check_batch_size_invalid_value(batch_size: int) -> None:
+    """Test that invalid batch_size values raise a ValueError."""
+    with pytest.raises(ValueError, match="Expected batch_size to be positive or -1"):
+        _check_batch_size(batch_size)

{gpjax-0.11.0 → gpjax-0.11.1}/tests/test_kernels/test_nonstationary.py

@@ -31,7 +31,7 @@ from gpjax.kernels.nonstationary import (
     Polynomial,
 )
 from gpjax.parameters import (
-    PositiveReal,
+    NonNegativeReal,
     Static,
 )
 
@@ -96,8 +96,8 @@ def test_init_override_paramtype(kernel_request):
             continue
         new_params[param] = Static(value)
 
-    k = kernel(**new_params, variance=PositiveReal(variance))
-    assert isinstance(k.variance, PositiveReal)
+    k = kernel(**new_params, variance=NonNegativeReal(variance))
+    assert isinstance(k.variance, NonNegativeReal)
 
     for param in params.keys():
         if param in ("degree", "order"):
@@ -112,7 +112,7 @@ def test_init_defaults(kernel: type[AbstractKernel]):
 
     # Check that the parameters are set correctly
     assert isinstance(k.compute_engine, type(AbstractKernelComputation()))
-    assert isinstance(k.variance, PositiveReal)
+    assert isinstance(k.variance, NonNegativeReal)
 
 
 @pytest.mark.parametrize("kernel", [k[0] for k in TESTED_KERNELS])
@@ -122,7 +122,7 @@ def test_init_variances(kernel: type[AbstractKernel], variance):
     k = kernel(variance=variance)
 
     # Check that the parameters are set correctly
-    assert isinstance(k.variance, PositiveReal)
+    assert isinstance(k.variance, NonNegativeReal)
     assert jnp.allclose(k.variance.value, jnp.asarray(variance))
 
     # Check that error is raised if variance is not valid