Change numpy code by jax.numpy

app.py

@@ -1,7 +1,12 @@
 import streamlit as st
-import numpy as np
+import jax.numpy as jnp
+import jax
 import matplotlib.pyplot as plt
 
+# Set random key
+seed=321
+key = jax.random.PRNGKey(seed)
+
 st.title('Fitting simple models with JAX')
 st.header('A quadratric regression example')
 
@@ -15,19 +20,18 @@ number_of_observations = st.sidebar.slider('Number of observations', min_value=5
 noise_standard_deviation = st.sidebar.slider('Standard deviation of the noise', min_value = 0.0, max_value=2.0, value=1.0)
 cost_function = st.sidebar.radio('What cost function you want to use for the fitting?', options=('RMSE-Loss', 'Huber-Loss'))
 
-np.random.seed(2)
-
+# Generate random data
 X = np.column_stack((np.ones(number_of_observations), 
                      np.random.random(number_of_observations)))      
 
 w = np.array([3.0, -20.0, 32.0])  # coefficients                                    
-
 X = np.column_stack((X, X[:,1] ** 2))   # add x**2 column
-additional_noise = 8 * np.random.binomial(1, 0.03, size = number_of_observations)
-y = np.dot(X, w) + noise_standard_deviation * np.random.randn(number_of_observations) \
+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	
 
 
+# Plot the data
 fig, ax = plt.subplots(dpi=320)
 ax.set_xlim((0,1))
 ax.set_ylim((-5,26))
@@ -46,6 +50,10 @@ st.latex(r'''\bf{w}\leftarrow \bf{w}-\eta \frac{\partial\ell(\bf{X},\bf{y}, \bf{
 
 # Fitting by the respective cost_function
 if cost_function == 'RMSE-Loss':
+
+     def loss(w):
+         return 1/X.shape[0] * jax.numpy.linalg.norm(jnp.dot(X, w) - y)**2
+         
      st.write('You selected the RMSE loss function.')
      st.latex(r'''\ell(X, y, w)=\frac{1}{m}||Xw - y||_{2}^2''')
      st.latex(r'''\ell(X, y, w)=\frac{1}{m}\big(\sqrt{(Xw - y)\cdot(Xw - y)}\big)^2''')