A short example to performing regression with JAXNS

[Cyberface]

The following example can be found in this google colab notebook but I have reproduced it below.

This is a short example how to use the nested sampler JAXNS to perform some simple linear regression.

Thanks to @Joshuaalbert for helping me understand how to use JAXNS!

Preamble

first install JAXNS

!pip install git+http://github.com/Joshuaalbert/jaxns.git

This is using version 0.0.6

Next import the packages we need

from jaxns.nested_sampling import NestedSampler
from jaxns.prior_transforms import PriorChain, UniformPrior, NormalPrior, HalfLaplacePrior
from jaxns.plotting import plot_cornerplot, plot_diagnostics
from jax import random, jit
from jax import numpy as jnp

import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams.update({"font.size": 16})

Data Generation

We will generate some fake data based on a cubic function and use JAXNS to estimate the parameters of this polynomial.

def model(x, a, b, c, d):
    """cubic function"""
    return a*x**3 + b*x**2 + c*x + d

def make_data(a=1, b=0, c=2, d=1, sigma=0.1):
    """generate some data with noise"""
    x = jnp.linspace(-3., 3, 100)[:, None]
    y = model(x, a, b, c, d) + sigma*random.normal(random.PRNGKey(0), shape=(100,1))
    return x, y


"""true parameters"""
a_true = -4
b_true = 3
c_true = 25
d_true = 6
sigma_true = 5

x, y = make_data(a_true, b_true, c_true, d_true, sigma_true)

# note that the shape of x and y will be 2D arrays.
# x.shape = [N, 1]
# y.shape = [N, 1]
# Where N = 100 in this example. i.e. the number of data points.

You can plot the data with

plt.figure()
plt.plot(x, y, 'o')
plt.xlabel('x')
plt.ylabel('y(x)')
plt.show()
plt.close()

Likelihood model

Let's define our likelihood function. Here we use the standard Gaussian likelihood, which assumes the noise is i.i.d with a Normal(loc=0, scale=sigma) distribution.

def lnlike(x, data, a, b, c, d, sigma):
    """
    gaussian likelihood
    """
    
    yhat = model(x, a, b, c, d)
    N = x.size
    
    term1 = -0.5 * N * jnp.log(2. * jnp.pi * sigma**2)
    term2 =  -0.5/sigma**2 * jnp.sum((y - yhat)**2)
    
    return term1 + term2

I believe JAXNS needs a likelihood function where the arguments of the function are the sampling parameters.

So we can use a lambda function to achieve this

log_likelihood = lambda a, b, c, d, sigma, **unused_kwargs: lnlike(x, y, a, b, c, d, sigma)

Setup Priors

We will use uniform priors for the a,b,c,d parameters and a half Laplace prior for the noise terms, because we want the noise term to be positive definate.

prior_a = UniformPrior('a', -20., 20.)
prior_b = UniformPrior('b', -20., 20.)
prior_c = UniformPrior('c', 0, 40)
prior_d = UniformPrior('d', -20., 20.)
prior_sigma = HalfLaplacePrior('sigma', 1)

# finally, in JAXNS you setup the prior with the following line:
prior_chain = PriorChain().push(prior_a).push(prior_b).push(prior_c).push(prior_d).push(prior_sigma)

Setup and run nested sampler

def run_with_n(n):
    ns = NestedSampler(log_likelihood,
                        prior_chain,
                        num_live_points=n,
                        max_samples=1e5,
                        collect_samples=True,
                        sampler_name='slice',
                        sampler_kwargs=dict(depth=7, num_slices=3))
    @jit
    def run(key):
        return ns(key=key,
                    termination_frac=0.01)

    t0 = default_timer()
    results = run(random.PRNGKey(0))
    print("Efficiency", results.efficiency)
    print("Time to run (including compile)", default_timer() - t0)
    t0 = default_timer()
    results = run(random.PRNGKey(1))
    print(results.efficiency)
    print("Time to run (no compile)", default_timer() - t0)
    return results

Run the sampler

# run with options (run twice to see the true speed without compile)
# the first time you run the function `run` it will perform the just-in-time (JIT) compilation.
# the second time you run it, it won't do this step and will be much faster.
# here we run with 500 live points
results = run_with_n(500)

In the google colab notebook the first execution (with compilation) took about 28 seconds.
The second execution just took 1 second!

Plotting results

There are a couple nice built-in plotting function in JAXNS

plot_diagnostics(results)
plot_cornerplot(results)

Getting the posterior samples

To get the posterior samples you need to use the jaxns.utils.resample function

from jaxns.utils import resample

samples = resample(random.PRNGKey(34252), results.samples, results.log_p, S=int(results.ESS))

When we have the independent posterior samples we can compute statistical quantities like the 5th, 50th and 95th percentiles and plot the predictions over the data as a visual check.

def get_percentiles(data):
    return jnp.percentile(data, [5, 50, 95])

a_5, a_50, a_95 = get_percentiles(samples['a'])
b_5, b_50, b_95 = get_percentiles(samples['b'])
c_5, c_50, c_95 = get_percentiles(samples['c'])
d_5, d_50, d_95 = get_percentiles(samples['d'])

plt.figure(figsize=(14, 8))
plt.plot(x, y, 'o', label='data')
plt.plot(x, model(x, a_50, b_50, c_50, d_50), c='k', label='median')

yhat_5 = model(x, a_5, b_5, c_5, d_5)
yhat_95 = model(x, a_95, b_95, c_95, d_95)

plt.fill_between(x[:,0], y1=yhat_5[:,0], y2=yhat_95[:,0], color='k', alpha=0.6, label='95% credible interval')

plt.xlabel('x')
plt.ylabel('y(x)')
plt.legend()
plt.show()
plt.close()

Final result:

[Joshuaalbert]

@Cyberface jaxns 0.0.2 has been released. There's a slightly new API, along with a host of improvements. Can you check out an example for the API and update your example here?

[Cyberface]

Yeah sure no problem.

I've got it working with the new API but I am confused about the following change to prior_chain.

The shape of prior_chain now seems to be the same as the number of keys in the prior_chain dictionary. Is that intentionally? Because now I cannot use prior_chain.U_ndims to get the number of variable parameters.

So right now prior_chain.U_ndims is 14. In this example there are 5 free model parameters. Is there another way to get this number now?

[Joshuaalbert]

Hi @Cyberface U_ndims the dimensionality of the underlying probability sample space, and should be 5 for your case as you expect. Is it showing 14 in the colab?

[Joshuaalbert]

I see it is counting deterministic variables, which it shouldn't be. I'll fix this.

[Joshuaalbert]

Fixed

[Cyberface]

Sorry, took a while to check but yup all working great!

I've updated the example above

[Joshuaalbert]

Hey @Cyberface would you be interested in pull latest version, 1.0 release candidate and running this?