Make zero-mean fn. constant a Static class (#500)

* Make zero-mean fn. constant a Static class

* Remove stray comment

* Add missing imports

* Fix zero-mean bug

* Fix zero-mean bug

* Bump version

Author: thomaspinder

gpjax/__init__.py

@@ -39,7 +39,7 @@
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.10.0"
+__version__ = "0.10.1"
 
 __all__ = [
     "base",

gpjax/kernels/base.py

@@ -32,6 +32,7 @@
 from gpjax.parameters import (
     Parameter,
     Real,
+    Static,
 )
 from gpjax.typing import (
     Array,
@@ -220,7 +221,9 @@ class Constant(AbstractKernel):
     def __init__(
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
-        constant: tp.Union[ScalarFloat, Parameter[ScalarFloat]] = jnp.array(0.0),
+        constant: tp.Union[
+            ScalarFloat, Parameter[ScalarFloat], Static[ScalarFloat]
+        ] = jnp.array(0.0),
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),
     ):
         if isinstance(constant, Parameter):

gpjax/mean_functions.py

@@ -28,6 +28,7 @@
 from gpjax.parameters import (
     Parameter,
     Real,
+    Static
 )
 from gpjax.typing import (
     Array,
@@ -130,9 +131,9 @@ class Constant(AbstractMeanFunction):
     """
 
     def __init__(
-        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter] = 0.0
+        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0
     ):
-        if isinstance(constant, Parameter):
+        if isinstance(constant, Parameter) or isinstance(constant, Static):
             self.constant = constant
         else:
             self.constant = Real(jnp.array(constant))
@@ -158,7 +159,7 @@ class Zero(Constant):
     """
 
     def __init__(self):
-        super().__init__(constant=jnp.array(0.0))
+        super().__init__(constant=Static(jnp.array(0.0)))
 
 
 class CombinationMeanFunction(AbstractMeanFunction):

tests/test_mean_functions.py

@@ -5,17 +5,20 @@
 
 
 import jax.numpy as jnp
+import jax.random as jr
 from jaxtyping import (
     Array,
     Float,
 )
 import pytest
 
+import gpjax as gpx
 from gpjax.mean_functions import (
     AbstractMeanFunction,
     Constant,
     Zero,
 )
+from gpjax.parameters import Static
 
 
 def test_abstract() -> None:
@@ -49,38 +52,28 @@ def test_constant(constant: Float[Array, " Q"]) -> None:
     ).all()
 
 
-# TODO: rewrite this test after work on fit
-# def test_zero_mean_remains_zero() -> None:
-#     key = jr.PRNGKey(123)
-
-#     x = jr.uniform(key=key, minval=0, maxval=1, shape=(20, 1))
-#     y = jnp.full((20, 1), 50, dtype=jnp.float64)  # Dataset with non-zero mean
-#     D = gpx.Dataset(X=x, y=y)
-
-#     kernel = gpx.kernels.Constant(constant=jnp.array(0.0))
-#     kernel = kernel.replace_trainable(
-#         constant=False
-#     )  # Prevent kernel from modelling non-zero mean
-#     meanf = Zero()
-#     prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
-#     likelihood = gpx.likelihoods.Gaussian(
-#         num_datapoints=D.n, obs_stddev=jnp.array(1e-3)
-#     )
-#     likelihood = likelihood.replace_trainable(obs_stddev=False)
-#     posterior = prior * likelihood
-
-#     negative_mll = gpx.objectives.ConjugateMLL(negative=True)
-#     opt_posterior, _ = gpx.fit(
-#         model=posterior,
-#         objective=negative_mll,
-#         train_data=D,
-#         optim=ox.adam(learning_rate=0.5),
-#         num_iters=1000,
-#         safe=True,
-#         key=key,
-#     )
-
-#     assert opt_posterior.prior.mean_function.constant == 0.0
+def test_zero_mean_remains_zero() -> None:
+    key = jr.PRNGKey(123)
+
+    x = jr.uniform(key=key, minval=0, maxval=1, shape=(20, 1))
+    y = jnp.full((20, 1), 50, dtype=jnp.float64)  # Dataset with non-zero mean
+    D = gpx.Dataset(X=x, y=y)
+
+    constant = Static(jnp.array(0.0))
+    kernel = gpx.kernels.Constant(constant=constant)
+    meanf = Zero()
+    prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
+    likelihood = gpx.likelihoods.Gaussian(
+        num_datapoints=D.n, obs_stddev=jnp.array(1e-3)
+    )
+    posterior = prior * likelihood
+
+    opt_posterior, _ = gpx.fit_scipy(
+        model=posterior,
+        objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
+        train_data=D,
+    )
+    assert opt_posterior.prior.mean_function.constant.value == 0.0
 
 
 def test_initialising_zero_mean_with_constant_raises_error():