Add training loop for RMSE-loss

app.py

@@ -51,6 +51,15 @@ st.latex(r'''\bf{w}\leftarrow \bf{w}-\eta \frac{\partial\ell(\bf{X},\bf{y}, \bf{
 
 
 # Fitting by the respective cost_function
+w = jnp.array(np.random.random(3))
+learning_rate = 0.05
+NUM_ITER = 1000
+
+fig, ax = plt.subplots(dpi=120)
+ax.set_xlim((0,1))
+ax.set_ylim((-5,26))
+ax.scatter(X[:,1], y, c='#e76254' ,edgecolors='firebrick')
+
 if cost_function == 'RMSE-Loss':
 
      def loss(w):
@@ -60,6 +69,33 @@ if cost_function == 'RMSE-Loss':
      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''')
      st.latex(r'''\ell(X, y, w)= \frac{1}{m}\sum_1^m (\hat{y}_i - y_i)^2''')
+     
+     
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    grad_loss = jax.grad(loss)
+    
+    # Perform gradient descent
+    for i in range(NUM_ITER):
+        # Update progress bar.
+        progress_bar.progress(i + 1)
+
+        # Update parameters.
+        w -= learning_rate * grad_loss(w)
+
+        # Update status text.
+        if (i+1)%100==0:
+            status_text.text(
+                'Trained loss is: %s' % loss(w))
+
+            # add a line to the plot
+            plt.plot(X[jnp.dot(X, w).argsort(), 1], jnp.dot(X, w).sort(), '--')
+      print(f"Trained loss at epoch #{i}:", loss(w))
+    plt.plot(X[jnp.dot(X, w).argsort(), 1], jnp.dot(X, w).sort(), '--', label='Final line')
+
+    status_text.text('Done!')
+
+
 else:
      st.write("You selected the Huber loss function.")
      st.latex(r'''