first push

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+    "python.linting.mypyEnabled": true,
+    "python.linting.enabled": true
+}

app.py

@@ -0,0 +1,105 @@
+import streamlit as st
+
+import jax
+import optax
+import jax.numpy as jnp
+import flax.linen as nn
+from flax.linen.initializers import zeros
+from tinygp import kernels, transforms, GaussianProcess
+import numpy as np
+import matplotlib.pyplot as plt
+import regdata as rd
+
+data_dict = {'motorcycle': rd.MotorcycleHelmet()}
+
+data = st.selectbox('Data', [key for key in data_dict.keys()])
+
+# st.markdown(f"{data}")
+
+x, y, t = data_dict[data].get_data()
+x = x.ravel()
+t = t.ravel()
+y = y.ravel()
+
+random = np.random.default_rng(567)
+
+noise = 0.01
+
+# x = np.sort(random.uniform(-1, 1, 100))
+# y = 2 * (x > 0) - 1 + random.normal(0.0, noise, len(x))
+# t = np.linspace(-1.5, 1.5, 500)
+
+# Define a small neural network used to non-linearly transform the input data in our model
+class Transformer(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        x = nn.Dense(features=15)(x)
+        x = nn.relu(x)
+        x = nn.Dense(features=10)(x)
+        x = nn.relu(x)
+        x = nn.Dense(features=1)(x)
+        return x
+
+
+class GPLoss(nn.Module):
+    @nn.compact
+    def __call__(self, x, y, t):
+        # Set up a typical Matern-3/2 kernel
+        log_sigma = self.param("log_sigma", zeros, ())
+        log_rho = self.param("log_rho", zeros, ())
+        # log_jitter = self.param("log_jitter", zeros, ())
+        base_kernel = jnp.exp(2 * log_sigma) * kernels.Matern32(
+            jnp.exp(log_rho)
+        )
+
+        # Define a custom transform to pass the input coordinates through our `Transformer`
+        # network from above
+        transform = Transformer()
+        log_jitter = transform(x.reshape(-1,1)).ravel()
+        kernel = base_kernel
+
+        # Evaluate and return the GP negative log likelihood as usual
+        gp = GaussianProcess(
+            kernel, x[:, None], diag=noise**2 + jnp.exp(log_jitter)
+        )
+        log_prob, gp_cond = gp.condition(y, t[:, None])
+        return -log_prob, (gp_cond.loc, gp_cond.variance), jnp.exp(transform(t[:, None]))
+
+
+# Define and train the model
+def loss(params):
+    return model.apply(params, x, y, t)[0]
+
+
+model = GPLoss()
+params = model.init(jax.random.PRNGKey(1234), x, y, t)
+tx = optax.sgd(learning_rate=1e-3)
+opt_state = tx.init(params)
+loss_grad_fn = jax.jit(jax.value_and_grad(loss))
+for i in range(100):
+    loss_val, grads = loss_grad_fn(params)
+    updates, opt_state = tx.update(grads, opt_state)
+    params = optax.apply_updates(params, updates)
+
+# Plot the results and compare to the true model
+plt.figure()
+_, (mu, var), jitter = model.apply(params, x, y, t)
+var += jitter.ravel()
+# plt.plot(t, 2 * (t > 0) - 1, "k", lw=1, label="truth")
+plt.plot(x, y, ".k", label="data")
+plt.plot(t, mu)
+idx = np.argsort(t)
+plt.plot(t[idx], jitter[idx], label='noise')
+plt.fill_between(
+    t, mu + 2*np.sqrt(var), mu - 2*np.sqrt(var), alpha=0.5, label="95% conf"
+)
+# plt.xlim(-1.5, 1.5)
+# plt.ylim(-1.3, 1.3)
+plt.xlabel("x")
+plt.ylabel("y")
+_ = plt.legend()
+
+col = st.columns(1)[0]
+with col:
+    st.pyplot(plt)
+

requirements.txt

@@ -0,0 +1,6 @@
+jaxlib
+jax
+optax
+flax
+tinygp
+matplotlib