Working prediction plot

Author: MilesCranmer

gui/app.py

@@ -222,7 +222,8 @@ def main():
                     "batch_size",
                 ]
             ],
-            outputs=blocks["df"],
+            outputs=[blocks["df"], blocks["predictions_plot"]],
+            show_progress=True,
         )
 
         # Any update to the equation choice will trigger a plot_example_data:

gui/plots.py

@@ -61,6 +61,20 @@ def plot_example_data(test_equation, num_points, noise_level, data_seed):
     return fig
 
 
+def plot_predictions(y, ypred):
+    fig, ax = plt.subplots(figsize=(6, 6), dpi=100)
+
+    ax.scatter(y, ypred, alpha=0.7, edgecolors="w", s=50)
+
+    stylize_axis(ax)
+
+    ax.set_xlabel("true")
+    ax.set_ylabel("prediction")
+    fig.tight_layout(pad=2)
+
+    return fig
+
+
 def stylize_axis(ax):
     ax.grid(True, which="both", ls="--", linewidth=0.5, color="gray", alpha=0.5)
     ax.spines["top"].set_visible(False)

gui/processing.py

@@ -7,6 +7,7 @@
 
 import pandas as pd
 from data import generate_data, read_csv
+from plots import plot_predictions
 
 EMPTY_DF = lambda: pd.DataFrame(
     {
@@ -188,7 +189,10 @@ def processing(
                 )
             )
             out = PERSISTENT_READER.out_queue.get()
+            predictions = out["ypred"]
             equations = out["equations"]
-            yield equations[["Complexity", "Loss", "Equation"]]
+            yield equations[["Complexity", "Loss", "Equation"]], plot_predictions(
+                y, predictions
+            )
 
         time.sleep(0.1)