Support priors

Author: thomaspinder

gpjax/objectives/mlls.py

@@ -7,7 +7,7 @@
 from ..gps import ConjugatePosterior, NonConjugatePosterior
 from ..kernels import gram
 from ..likelihoods import link_function
-from ..parameters.prior_densities import log_density
+from ..parameters.priors import evaluate_prior, prior_checks
 from ..parameters.transforms import (SoftplusTransformation, Transformation,
                                      untransform)
 from ..types import Array
@@ -28,17 +28,17 @@ def marginal_ll(
     Returns: A multivariate normal distribution
     """
 
-    def mll(params: dict, x: Array, y: Array):
+    def mll(params: dict, x: Array, y: Array, priors: dict = None):
         params = untransform(params, transformation)
         mu = gp.prior.mean_function(x)
         gram_matrix = params["variance"] * gram(gp.prior.kernel, x / params["lengthscale"])
         gram_matrix += params["obs_noise"] * I(x.shape[0])
         L = jnp.linalg.cholesky(gram_matrix)
         random_variable = tfd.MultivariateNormalTriL(mu, L)
-        # TODO: Attach log-prior density sum here
-        constant = jnp.array(-1.0) if negative else jnp.array(1.0)
-        return constant * random_variable.log_prob(y.squeeze()).mean()
 
+        log_prior_density = evaluate_prior(params, priors)
+        constant = jnp.array(-1.0) if negative else jnp.array(1.0)
+        return constant * (random_variable.log_prob(y.squeeze()).mean() + log_prior_density)
     return mll
 
 
@@ -49,7 +49,7 @@ def marginal_ll(
     negative: bool = False,
     jitter: float = 1e-6,
 ) -> Callable:
-    def mll(params: dict, x: Array, y: Array):
+    def mll(params: dict, x: Array, y: Array, priors: dict = {'latent': tfd.Normal(loc=0., scale=1.)}):
         params = untransform(params, transformation)
         n = x.shape[0]
         link = link_function(gp.likelihood)
@@ -59,9 +59,10 @@ def mll(params: dict, x: Array, y: Array):
         F = jnp.matmul(L, params["latent"])
         rv = link(F)
         ll = jnp.sum(rv.log_prob(y))
-        # TODO: Attach full log-prior density sum here
-        latent_prior = jnp.sum(log_density(params["latent"], tfd.Normal(loc=0.0, scale=1.0)))
+
+        priors = prior_checks(gp, priors)
+        log_prior_density = evaluate_prior(params, priors)
         constant = jnp.array(-1.0) if negative else jnp.array(1.0)
-        return constant * (ll + latent_prior)
+        return constant * (ll + log_prior_density)
 
     return mll

gpjax/parameters/__init__.py

@@ -1,4 +1,4 @@
 from .base import complete, initialise
-from .prior_densities import log_density
+from .priors import log_density
 from .transforms import (IdentityTransformation, SoftplusTransformation,
                          transform, untransform)

gpjax/parameters/prior_densities.py

@@ -0,0 +1,39 @@
+import jax.numpy as jnp
+from jax.interpreters.ad import JVPTracer
+from jax.interpreters.partial_eval import DynamicJaxprTracer
+from multipledispatch import dispatch
+from tensorflow_probability.substrates.jax import distributions as tfd
+from ..gps import NonConjugatePosterior
+import warnings
+
+from ..types import Array, NoneType
+
+
+@dispatch((jnp.DeviceArray, JVPTracer, DynamicJaxprTracer), tfd.Distribution)
+def log_density(param: jnp.DeviceArray, density: tfd.Distribution) -> Array:
+    return density.log_prob(param)
+
+
+@dispatch(dict, NoneType)
+def evaluate_prior(params: dict, priors: dict) -> Array:
+    return jnp.array(0.)
+
+
+@dispatch(dict, dict)
+def evaluate_prior(params: dict, priors: dict) -> Array:
+    lpd = jnp.array(0)
+    for param, val in priors.items():
+        lpd+=jnp.sum(log_density(params[param], priors[param]))
+    return lpd
+
+
+@dispatch(NonConjugatePosterior, dict)
+def prior_checks(gp: NonConjugatePosterior, priors: dict) -> dict:
+    if 'latent' in priors.keys():
+        latent_prior = priors['latent']
+        if latent_prior.name != 'Normal':
+            warnings.warn(f'A {latent_prior.name} distribution prior has been placed on the latent function. It is strongly afvised that a unit-Gaussian prior is used.')
+        return priors
+    else:
+        priors['latent'] = tfd.Normal(loc=0., scale=1.)
+        return priors

gpjax/parameters/priors.py

@@ -6,3 +6,4 @@
 
 # Array = Union[jnp.ndarray, ShardedDeviceArray, jnp.DeviceArray]  # Cannot currently dispatch on a Union type
 # Data = Tuple[Array, Array]
+NoneType = type(None)

gpjax/types.py

@@ -4,8 +4,10 @@
 from gpjax.kernels import RBF
 from gpjax.parameters import transform, SoftplusTransformation, initialise
 import jax.numpy as jnp
+import jax.random as jr
 import pytest
 from typing import Callable
+from tensorflow_probability.substrates.jax import distributions as tfd
 
 
 def test_conjugate():
@@ -29,3 +31,27 @@ def test_non_conjugate():
     y = jnp.sin(x)
     params = transform(params=initialise(posterior, n), transformation=SoftplusTransformation)
     assert neg_mll(params, x, y) == jnp.array(-1.)*mll(params, x, y)
+
+
+def test_prior_mll():
+    """
+    Test that the MLL evaluation works with priors attached to the parameter values.
+    """
+    key = jr.PRNGKey(123)
+    x = jnp.sort(jr.uniform(key, minval=-5.0, maxval=5.0, shape=(100, 1)), axis=0)
+    f = lambda x: jnp.sin(jnp.pi * x) / (jnp.pi * x)
+    y = f(x) + jr.normal(key, shape=x.shape) * 0.1
+    posterior = Prior(kernel=RBF()) * Gaussian()
+    mll = marginal_ll(posterior)
+
+    params = initialise(posterior)
+    priors = {
+        "lengthscale": tfd.Gamma(1.0, 1.0),
+        "variance": tfd.Gamma(2.0, 2.0),
+        "obs_noise": tfd.Gamma(2.0, 2.0),
+    }
+    mll_eval = mll(params, x, y)
+    mll_eval_priors = mll(params, x, y, priors)
+
+    assert pytest.approx(mll_eval) == jnp.array(-115.72332969)
+    assert pytest.approx(mll_eval_priors) == jnp.array(-118.97202259)

tests/objectives/test_mlls.py

@@ -0,0 +1,85 @@
+from gpjax.parameters import log_density
+from gpjax.parameters.priors import evaluate_prior, prior_checks
+from gpjax.gps import Prior
+from gpjax.kernels import RBF
+from gpjax.likelihoods import Bernoulli
+from tensorflow_probability.substrates.jax import distributions as tfd
+import pytest
+import jax.numpy as jnp
+
+
+@pytest.mark.parametrize('x', [-1., 0., 1.])
+def test_lpd(x):
+    val = jnp.array(x)
+    dist = tfd.Normal(loc=0., scale=1.)
+    lpd = log_density(val, dist)
+    assert lpd is not None
+
+
+def test_prior_evaluation():
+    """
+    Test the regular setup that every parameter has a corresponding prior distribution attached to its unconstrained
+    value.
+    """
+    params = {
+        "lengthscale": jnp.array([1.]),
+        "variance": jnp.array([1.]),
+        "obs_noise": jnp.array([1.]),
+    }
+    priors = {
+        "lengthscale": tfd.Gamma(1.0, 1.0),
+        "variance": tfd.Gamma(2.0, 2.0),
+        "obs_noise": tfd.Gamma(3.0, 3.0),
+    }
+    lpd = evaluate_prior(params, priors)
+    assert pytest.approx(lpd) == -2.0110168
+
+
+def test_none_prior():
+    """
+    Test that multiple dispatch is working in the case of no priors.
+    """
+    params = {
+        "lengthscale": jnp.array([1.]),
+        "variance": jnp.array([1.]),
+        "obs_noise": jnp.array([1.]),
+    }
+    lpd = evaluate_prior(params, None)
+    assert lpd == 0.
+
+
+def test_incomplete_priors():
+    """
+    Test the case where a user specifies priors for some, but not all, parameters.
+    """
+    params = {
+        "lengthscale": jnp.array([1.]),
+        "variance": jnp.array([1.]),
+        "obs_noise": jnp.array([1.]),
+    }
+    priors = {
+        "lengthscale": tfd.Gamma(1.0, 1.0),
+        "variance": tfd.Gamma(2.0, 2.0),
+    }
+    lpd = evaluate_prior(params, priors)
+    assert pytest.approx(lpd) == -1.6137061
+
+
+def test_checks():
+    incomplete_priors = {'lengthscale': jnp.array([1.])}
+    posterior = Prior(kernel=RBF()) * Bernoulli()
+    priors = prior_checks(posterior, incomplete_priors)
+    assert 'latent' in priors.keys()
+    assert 'variance' not in priors.keys()
+
+
+def test_check_needless():
+    complete_prior = {
+        "lengthscale": tfd.Gamma(1.0, 1.0),
+        "variance": tfd.Gamma(2.0, 2.0),
+        "obs_noise": tfd.Gamma(3.0, 3.0),
+        "latent": tfd.Normal(loc=0., scale=1.)
+    }
+    posterior = Prior(kernel=RBF()) * Bernoulli()
+    priors = prior_checks(posterior, complete_prior)
+    assert priors == complete_prior

tests/parameters/test_prior_densities.py

@@ -0,0 +1,5 @@
+from gpjax.types import NoneType
+
+
+def test_nonetype():
+    assert isinstance(None, NoneType)

tests/parameters/test_priors.py

@@ -1,5 +1,6 @@
-from gpjax.utils import concat_dictionaries, I, merge_dictionaries
+from gpjax.utils import concat_dictionaries, I, merge_dictionaries, standardise, unstandardise
 import jax.numpy as jnp
+import jax.random as jr
 import pytest
 
 
@@ -24,3 +25,17 @@ def test_merge_dicts():
     d = merge_dictionaries(d1, d2)
     assert list(d.keys()) == ['a', 'b']
     assert list(d.values()) == [1, 3]
+
+
+def test_standardise():
+    key = jr.PRNGKey(123)
+    x = jr.uniform(key, shape=(100, 1))
+    xtr, xmean, xstd = standardise(x)
+    assert pytest.approx(jnp.mean(xtr), rel=4) == 0.
+
+    xtr2 = standardise(x, xmean, xstd)
+    assert pytest.approx(jnp.mean(xtr2), rel=4) == 0.
+
+    xuntr = unstandardise(xtr, xmean, xstd)
+    diff = jnp.sum(jnp.abs(xuntr - x))
+    assert pytest.approx(diff) == 0