Merge pull request #400 from JaxGaussianProcesses/fix_noise

Fix obs_noise confusion

Author: henrymoss

.gitignore

@@ -151,4 +151,4 @@ package.json
 package-lock.json
 node_modules/
 
-docs/api
+docs/api

docs/GOVERNANCE.md

@@ -93,4 +93,4 @@ maintainers or reach out over
 -----
 
 This file was adapted from
-[BlackJAX](https://github.com/blackjax-devs/blackjax/blob/main/GOVERNANCE.md).
+[BlackJAX](https://github.com/blackjax-devs/blackjax/blob/main/GOVERNANCE.md).

docs/examples/bayesian_optimisation.py

@@ -204,9 +204,9 @@ def return_optimised_posterior(
     data: gpx.Dataset, prior: gpx.Module, key: Array
 ) -> gpx.Module:
     likelihood = gpx.Gaussian(
-        num_datapoints=data.n, obs_noise=jnp.array(1e-6)
-    )  # Our function is noise-free, so we set the observation noise to a very small value
-    likelihood = likelihood.replace_trainable(obs_noise=False)
+        num_datapoints=data.n, obs_stddev=jnp.array(1e-3)
+    )  # Our function is noise-free, so we set the observation noise's standard deviation to a very small value
+    likelihood = likelihood.replace_trainable(obs_stddev=False)
 
     posterior = prior * likelihood
 

docs/examples/decision_making.py

@@ -153,7 +153,9 @@ def forrester(x: Float[Array, "N 1"]) -> Float[Array, "N 1"]:
 # with the correct number of datapoints:
 
 # %%
-likelihood_builder = lambda n: gpx.Gaussian(num_datapoints=n, obs_noise=jnp.array(1e-6))
+likelihood_builder = lambda n: gpx.Gaussian(
+    num_datapoints=n, obs_stddev=jnp.array(1e-3)
+)
 
 # %% [markdown]
 # Now we have all the components required for constructing our GP posterior. Since we'll

docs/examples/intro_to_kernels.py

@@ -223,9 +223,9 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]:
 prior = gpx.Prior(mean_function=mean, kernel=kernel)
 
 likelihood = gpx.Gaussian(
-    num_datapoints=D.n, obs_noise=jnp.array(1e-6)
-)  # Our function is noise-free, so we set the observation noise to a very small value
-likelihood = likelihood.replace_trainable(obs_noise=False)
+    num_datapoints=D.n, obs_stddev=jnp.array(1e-3)
+)  # Our function is noise-free, so we set the observation noise's standard deviation to a very small value
+likelihood = likelihood.replace_trainable(obs_stddev=False)
 
 no_opt_posterior = prior * likelihood
 

docs/examples/likelihoods_guide.py

@@ -99,12 +99,13 @@
 # Some likelihoods, such as the Gaussian likelihood, contain parameters that we seek
 # to infer. In the case of the Gaussian likelihood, we have a single parameter
 # $\sigma^2$ that determines the observation noise. In GPJax, we can specify the value
-# of this parameter when instantiating the likelihood object. If we do not specify a
+# of $\sigma$ when instantiating the likelihood object. If we do not specify a
 # value, then the likelihood will be initialised with a default value. In the case of
 # the Gaussian likelihood, the default value is $1.0$. If we instead wanted to
-# initialise the likelihood with a value of $0.5$, then we would do this as follows:
+# initialise the likelihood standard deviation with a value of $0.5$, then we would do
+# this as follows:
 
-gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_noise=0.5)
+gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_stddev=0.5)
 
 # To control other properties of the observation noise such as trainability and value
 # constraints, see our [PyTree guide](pytrees.md).
@@ -127,7 +128,7 @@
 meanf = gpx.Zero()
 prior = gpx.Prior(kernel=kernel, mean_function=meanf)
 
-likelihood = gpx.Gaussian(num_datapoints=D.n, obs_noise=0.1)
+likelihood = gpx.Gaussian(num_datapoints=D.n, obs_stddev=0.1)
 
 posterior = prior * likelihood
 
@@ -252,7 +253,7 @@ def q_moments(x):
 
 lquad = gpx.Gaussian(
     num_datapoints=D.n,
-    obs_noise=jnp.array([0.1]),
+    obs_stddev=jnp.array([0.1]),
     integrator=gpx.integrators.GHQuadratureIntegrator(num_points=20),
 )
 

docs/examples/oceanmodelling.py

@@ -226,7 +226,7 @@ def __call__(
 def initialise_gp(kernel, mean, dataset):
     prior = gpx.Prior(mean_function=mean, kernel=kernel)
     likelihood = gpx.Gaussian(
-        num_datapoints=dataset.n, obs_noise=jnp.array([1.0e-6], dtype=jnp.float64)
+        num_datapoints=dataset.n, obs_stddev=jnp.array([1.0e-3], dtype=jnp.float64)
     )
     posterior = prior * likelihood
     return posterior

gpjax/dataset.py

@@ -17,9 +17,9 @@
 import warnings
 
 from beartype.typing import (
+    Literal,
     Optional,
     Union,
-    Literal,
 )
 import jax.numpy as jnp
 from jaxtyping import (

gpjax/distributions.py

@@ -16,9 +16,9 @@
 
 from beartype.typing import (
     Any,
+    Generic,
     Optional,
     Tuple,
-    Generic,
     TypeVar,
     Union,
 )

gpjax/gps.py

@@ -15,7 +15,10 @@
 
 # from __future__ import annotations
 from abc import abstractmethod
-from dataclasses import dataclass, field
+from dataclasses import (
+    dataclass,
+    field,
+)
 from typing import overload
 
 from beartype.typing import (
@@ -25,7 +28,6 @@
 )
 import cola
 from cola.ops import Dense
-
 import jax.numpy as jnp
 from jax.random import (
     PRNGKey,
@@ -47,7 +49,10 @@
     ReshapedDistribution,
     ReshapedGaussianDistribution,
 )
-from gpjax.kernels import RFF, White
+from gpjax.kernels import (
+    RFF,
+    White,
+)
 from gpjax.kernels.base import AbstractKernel
 from gpjax.likelihoods import (
     AbstractLikelihood,
@@ -503,7 +508,7 @@ def predict(
         n_test = len(test_inputs)
 
         # Observation noise o²
-        obs_noise = self.likelihood.obs_noise
+        obs_var = self.likelihood.obs_stddev**2
         mx = self.prior.mean_function(x)
 
         # Precompute Gram matrix, Kxx, at training inputs, x
@@ -512,7 +517,7 @@ def predict(
 
         # Σ = Kxx + Io²
         Sigma = cola.ops.Kronecker(Kxx, Kyy)
-        Sigma += cola.ops.I_like(Sigma) * (obs_noise + self.jitter)
+        Sigma += cola.ops.I_like(Sigma) * (obs_var + self.jitter)
         Sigma = cola.PSD(Sigma)
 
         if mask is not None:
@@ -606,13 +611,10 @@ def sample_approx(
 
         # sample weights v for canonical features
         # v = Σ⁻¹ (y + ε - ɸ⍵) for  Σ = Kxx + Io² and ε ᯈ N(0, o²)
+        obs_var = self.likelihood.obs_stddev**2
         Kxx = self.prior.kernel.gram(train_data.X)  #  [N, N]
-        Sigma = Kxx + cola.ops.I_like(Kxx) * (
-            self.likelihood.obs_noise + self.jitter
-        )  #  [N, N]
-        eps = jnp.sqrt(self.likelihood.obs_noise) * normal(
-            key, [train_data.n, num_samples]
-        )  #  [N, B]
+        Sigma = Kxx + cola.ops.I_like(Kxx) * (obs_var + self.jitter)  #  [N, N]
+        eps = jnp.sqrt(obs_var) * normal(key, [train_data.n, num_samples])  #  [N, B]
         y = train_data.y - self.prior.mean_function(train_data.X)  # account for mean
         Phi = fourier_feature_fn(train_data.X)
         canonical_weights = cola.solve(