add visu to prediction regression

frontend/pages/prediction_regression.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from sklearn.linear_model import LinearRegression
 import pandas as pd
+import matplotlib.pyplot as plt
+import numpy as np
 
 st.header("Prediction: Regression")
 
@@ -25,5 +27,34 @@ if "data" in st.session_state:
         prediction = model.predict(pd.DataFrame([pred_values], columns=data_name))
 
         st.write("Prediction:", prediction[0])
+
+        fig = plt.figure()
+        dataframe_sorted = pd.concat([X, y], axis=1).sort_values(by=data_name)
+
+        if len(data_name) == 1:
+            X = dataframe_sorted[data_name[0]]
+            y = dataframe_sorted[target_name]
+
+            prediction_array_y = [
+                model.predict(pd.DataFrame([[dataframe_sorted[data_name[0]].iloc[i]]], columns=data_name))[0]
+                for i in range(dataframe_sorted.shape[0])
+            ]
+
+            plt.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[target_name], color='b')
+            plt.scatter(dataframe_sorted[data_name[0]], prediction_array_y, color='r')
+        else:
+            ax = fig.add_subplot(111, projection='3d')
+
+            prediction_array_y = [
+                model.predict(pd.DataFrame([[dataframe_sorted[data_name[0]].iloc[i], dataframe_sorted[data_name[1]].iloc[i]]], columns=data_name))[0]
+                for i in range(dataframe_sorted.shape[0])
+            ]
+
+            ax.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[data_name[1]], dataframe_sorted[target_name], color='b')
+            ax.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[data_name[1]], prediction_array_y, color='r')
+
+        st.pyplot(fig)
+
+
 else:
     st.error("File not loaded")