add gpflow

app.py

@@ -55,21 +55,38 @@ st.title("Heteroscedastic Gaussian Processes")
 
 st.markdown(r"We are learning the noise v/s inputs relationship with a neural net.")
 
-data = st.selectbox("Data", ["Motorcycle", "Olympic"])
+data = st.selectbox("Data", ["Motorcycle", "Olympic", 'GPflow'])
 if data == "Motorcycle":
     data_x, data_y = mcycle_x, mcycle_y
 elif data == "Olympic":
     data_x, data_y = oly_x, oly_y
+elif data == 'GPflow':
+    N = 1001
+
+    # Build inputs X
+    data_x = np.linspace(0, 4 * np.pi, N)
+
+    # Deterministic functions in place of latent ones
+    f1 = np.sin
+    f2 = np.cos
+
+    # Use transform = exp to ensure positive-only scale values
+    transform = np.exp
+
+    # Compute loc and scale as functions of input X
+    loc = f1(data_x)
+    scale = transform(f2(data_x))
+
+    # Sample outputs Y from Gaussian Likelihood
+    data_y = np.random.normal(loc, scale)
     
 x = (data_x - data_x.mean()) / data_x.std()
 y = (data_y - data_y.mean()) / data_y.std()
 
 n_tests = st.number_input(
-    "Number of test points", min_value=10, max_value=200, value=100
+    "Number of test points", min_value=50, max_value=1000, value=100
 )
 
-n_iters = st.number_input("Number of iterations", min_value=1, max_value=100, value=10)
-
 t = np.linspace(x.min(), x.max(), n_tests)
 noise = 0.01
 
@@ -78,12 +95,16 @@ noise = 0.01
 # t = np.linspace(-1.5, 1.5, 500)
 
 # Define a small neural network used to non-linearly transform the input data in our model
+
+fet1 = st.slider("Number of neurons in Layer1", min_value=2, max_value=30, value=15)
+fet2 = st.slider("Number of neurons in Layer1", min_value=2, max_value=30, value=10)
+
 class Transformer(nn.Module):
     @nn.compact
     def __call__(self, x):
-        x = nn.Dense(features=15)(x)
+        x = nn.Dense(features=fet1)(x)
         x = nn.relu(x)
-        x = nn.Dense(features=10)(x)
+        x = nn.Dense(features=fet2)(x)
         x = nn.relu(x)
         x = nn.Dense(features=1)(x)
         return x
@@ -138,15 +159,19 @@ def loss(params):
 
 base_model = BaseGPLoss()
 model = GPLoss()
-base_params = base_model.init(jax.random.PRNGKey(1234), x, y, t)
-params = model.init(jax.random.PRNGKey(1234), x, y, t)
-tx = optax.sgd(learning_rate=1e-3)
+seed = np.random.randint(0,100)
+base_params = base_model.init(jax.random.PRNGKey(seed), x, y, t)
+params = model.init(jax.random.PRNGKey(np.random.randint(seed)), x, y, t)
+n_iters = st.number_input("Number of iterations", min_value=1, max_value=200, value=100)
+lr = st.selectbox("Learning rate", [0.1, 0.01, 0.001, 0.0001], 1)
+tx = optax.sgd(learning_rate=lr)
 base_opt_state = tx.init(base_params)
 opt_state = tx.init(params)
 loss_grad_fn = jax.jit(jax.value_and_grad(loss))
 base_losses = []
 losses = []
-for i in range(200):
+my_bar = st.progress(0)
+for i in range(n_iters):
     m = base_model
     base_loss_val, base_grads = loss_grad_fn(base_params)
     m = model
@@ -157,6 +182,7 @@ for i in range(200):
     params = optax.apply_updates(params, updates)
     losses.append(loss_val)
     base_losses.append(base_loss_val)
+    my_bar.progress((i+1) / n_iters)
 
 # Plot the results and compare to the true model
 fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(10, 4))