jaxns with tinygp models

[zairving]

Hi,

I'm trying to use nested sampling to do model comparison with tinygp models, and I figured this would be the ideal package to use since it's also built on JAX. However, I'm having a hard time getting started.

To show how I'd like to use jaxns, I've adapted a tutorial from the tinygp documentation to show how I'm currently using UltraNest with tinygp. If somebody could adapt my example to use jaxns instead of UltraNest, it should be enough for me to get to grips with using it as a replacement for UltraNest:

import numpy as np
import matplotlib.pyplot as plt
import tinygp
import jax
from ultranest import ReactiveNestedSampler

jax.config.update("jax_enable_x64", True)  # enable double precision in JAX
jax.config.update("jax_platform_name", "cpu")  # force JAX to run on CPU


def kernel(theta):
    return theta[0]**2*tinygp.kernels.ExpSquared(theta[1])*tinygp.kernels.Cosine(theta[2]) + theta[3]**2*tinygp.kernels.Matern32(theta[4])


def build_gp(theta):
    return tinygp.GaussianProcess(kernel(theta), x, diag=yerr**2 + 1e-8, mean=theta[5])


@jax.jit
def loglike(params):
    return build_gp(params).log_probability(y)


def transform(cube):

    priors = cube.copy()

    # uniform priors for all parameters

    lo, hi = 0, max(y) - min(y)
    priors[0] = cube[0]*(hi - lo) + lo  # ExpSquared amplitude
    priors[3] = cube[3]*(hi - lo) + lo  # Matern-3/2 amplitude

    lo, hi = 1e-1, 1e1
    priors[1] = cube[1]*(hi - lo) + lo  # ExpSquared length-scale
    priors[2] = cube[2]*(hi - lo) + lo  # cosine period

    lo, hi = 1, 100
    priors[4] = cube[4]*(hi - lo) + lo  # Matern-3/2 length-scale

    lo, hi = min(y), max(y)
    priors[5] = cube[5]*(hi - lo) + lo  # mean

    return priors


# generate some data
random = np.random.default_rng(42)
x = np.sort(
    np.append(
        random.uniform(0, 3.8, 28),
        random.uniform(5.5, 10, 18),
    )
)
yerr = random.uniform(0.08, 0.22, len(x))
y = (
    0.2 * (x - 5)
    + np.sin(3 * x + 0.1 * (x - 5) ** 2)
    + yerr * random.normal(size=len(x))
)
true_x = np.linspace(0, 10, 100)
true_y = 0.2 * (true_x - 5) + np.sin(3 * true_x + 0.1 * (true_x - 5) ** 2)


# create figure
fig, ax = plt.subplots(tight_layout=True, dpi=200)

#plot data
ax.errorbar(x, y, yerr, marker="x", capsize=2, color="black", linestyle="None", label="data")
ax.plot(true_x, true_y, "k-", alpha=.5, label="true")


# define sampler
sampler = ReactiveNestedSampler(param_names=["ExpSquared amp", "ExpSquared length-scale", "period",
                                             "Matern-3/2 amplitude", "Matern-3/2 length-scale", "mean"],
                                loglike=loglike,
                                transform=transform,
                                log_dir="/home/zac/Downloads/temp",
                                resume="overwrite")

# run sampler
results = sampler.run()
sampler.print_results()

# get medians of posterior distributions
theta = results["posterior"]["median"]

# create GP
gp = build_gp(theta)

# condition GP data
cond_gp = gp.condition(y, true_x).gp
mu, std = cond_gp.mean, np.sqrt(cond_gp.variance)

# plot GP predictive mean and 2 sigma confidence interval
ax.plot(true_x, mu, "r-", label="GP prediction")
ax.fill_between(true_x, mu+2*std, mu-2*std, color="lightgrey", alpha=.5, label="$2 \\sigma$ confidence")


# label plot
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.legend()


plt.show()

The resulting figure can be seen here.

Package versions: tinygp=0.2.3, UltraNest=3.5.7, JAX=0.4.1, Python=3.11.

I really appreciate any help.

[Joshuaalbert]

Here you can find a tutorial on how to do Gaussian processes with JAXNS: https://jaxns.readthedocs.io/en/latest/examples/gaussian_process_marginalisation.html