Update data generation

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import numpy as np
 import jax.numpy as jnp
 import jax
 import matplotlib.pyplot as plt
@@ -21,15 +22,14 @@ noise_standard_deviation = st.sidebar.slider('Standard deviation of the noise',
 cost_function = st.sidebar.radio('What cost function you want to use for the fitting?', options=('RMSE-Loss', 'Huber-Loss'))
 
 # Generate random data
-X = jnp.column_stack((jnp.ones(number_of_observations),
-                     jax.random.uniform(key, shape=(number_of_observations,), minval=0., maxval=1.)))
-w = jnp.array([3.0, -20.0, 32.0])  # coefficients                                    
+X = np.column_stack((np.ones(number_of_observations), 
+                     np.random.random(number_of_observations)))
 X = jnp.column_stack((X, X[:,1] ** 2))   # add x**2 column
-additional_noise = 8 * jax.random.bernoulli(key, p=0.08, shape=[number_of_observations,])
-y = jnp.dot(X, w) + noise_standard_deviation * jax.random.normal(key, shape=[number_of_observations,]) \
-        + additional_noise	
-
 
+additional_noise = 8 * np.random.binomial(1, 0.03, size = number_of_observations)
+y = jnp.array(np.dot(X, w) + noise_standard_deviation * np.random.randn(number_of_observations) \
+        + additional_noise)
+        
 # Plot the data
 fig, ax = plt.subplots(dpi=320)
 ax.set_xlim((0,1))