<jupyter_start><jupyter_text>The usual CO2 example<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
from statsmodels.datasets import co2

data = co2.load_pandas().data
t = 2000 + (np.array(data.index.to_julian_date()) - 2451545.0) / 365.25
y = np.array(data.co2)
m = np.isfinite(t) & np.isfinite(y) & (t < 1996)
t, y = t[m][::4], y[m][::4]

plt.plot(t, y, ".k")
plt.xlim(t.min(), t.max())
plt.xlabel("year")
_ = plt.ylabel("CO$_2$ in ppm")
from jax.config import config

config.update("jax_enable_x64", True)

import jax
import jax.numpy as jnp

from tinygp import kernels, GaussianProcess


def build_kernel(theta):
    k1 = theta[0] ** 2 * kernels.ExpSquared(theta[1])
    k2 = (
        theta[2] ** 2
        * kernels.ExpSquared(theta[3])
        * kernels.ExpSineSquared(theta[4], gamma=theta[5])
    )
    k3 = theta[6] ** 2 * kernels.RationalQuadratic(theta[7], alpha=theta[8])
    k4 = theta[9] ** 2 * kernels.ExpSquared(theta[10])
    return k1 + k2 + k3 + k4


def build_gp(theta, X):
    return GaussianProcess(
        build_kernel(jnp.exp(theta)),
        X,
        diag=jnp.exp(2 * theta[11]),
        mean=theta[12],
    )


def neg_log_likelihood(theta, X, y):
    gp = build_gp(theta, X)
    return -gp.condition(y)


theta_init = np.append(
    np.log(
        [66.0, 67.0, 2.4, 90.0, 1.0, 4.3, 0.66, 1.2, 0.78, 0.18, 1.6, 0.19]
    ),
    340.0,
)
obj = jax.jit(jax.value_and_grad(neg_log_likelihood))
print(f"Initial negative log likelihood: {obj(theta_init, t, y)[0]}")
from scipy.optimize import minimize

soln = minimize(obj, theta_init, jac=True, args=(t, y))
print(f"Final negative log likelihood: {soln.fun}")
x = np.linspace(max(t), 2025, 2000)
gp = build_gp(soln.x, t)
mu, var = gp.predict(y, x, return_var=True)

plt.plot(t, y, ".k")
plt.fill_between(
    x, mu + np.sqrt(var), mu - np.sqrt(var), color="C0", alpha=0.5
)
plt.plot(x, mu, color="C0", lw=2)

plt.xlim(t.min(), 2025)
plt.xlabel("year")
_ = plt.ylabel("CO$_2$ in ppm")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1805486284289277, 'fraction_numerical': 0.056359102244389024, 'mean_word_length': 3.6869158878504673, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '327397', 'license': 'MIT', 'n_tokens_mistral': 940, 'n_tokens_neox': 849, 'n_words': 177, 'path': 'docs/tutorials/quickstart.ipynb', 'repo_name': 'CBS2606/tinygp', 'size': '2005'}