legend added

app.py

@@ -55,9 +55,9 @@ def mean_and_bounds_for_pnf(x,y,test_xs, choice):
     bounds = model.criterion.quantile(logits,center_prob=.682).squeeze(1)
     return model.criterion.mean(logits).squeeze(1), bounds[:,0], bounds[:,1]
 
-def plot_w_conf_interval(ax_or_plt, x, m, lb, ub, color):
-    ax_or_plt.plot(x.squeeze(-1),m, color=color)
-    ax_or_plt.fill_between(x.squeeze(-1), lb, ub, alpha=.1, color=color)
+def plot_w_conf_interval(ax_or_plt, x, m, lb, ub, color, label_prefix):
+    ax_or_plt.plot(x.squeeze(-1),m, color=color, label=label_prefix+' mean')
+    ax_or_plt.fill_between(x.squeeze(-1), lb, ub, alpha=.1, color=color, label=label_prefix+' conf. interval')
 
 
 
@@ -74,22 +74,24 @@ def infer(table, choice):
         return excuse, None
     x = torch.tensor(table[:,0]).unsqueeze(1)
     y = torch.tensor(table[:,1])
-    fig = plt.figure(figsize=(4,2),dpi=1000)
+    fig = plt.figure(figsize=(8,4),dpi=1000)
 
     if len(x) > 4:
         return excuse_max_examples, None
     if (x<0.).any() or (x>1.).any() or (y<-1).any() or (y>1).any():
         return excuse, None
 
-    plt.scatter(x,y)
+    plt.scatter(x,y, color='black', label='Examples in given dataset')
 
 
     
     test_xs = torch.linspace(0,1,100).unsqueeze(1)
     
-    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_gp(x,y,test_xs), 'green')
-    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_pnf(x,y,test_xs, choice), 'blue')
-
+    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_gp(x,y,test_xs), 'green', 'GP')
+    plot_w_conf_interval(plt, test_xs, *mean_and_bounds_for_pnf(x,y,test_xs, choice), 'blue', 'PFN')
+    
+    plt.legend(ncol=2,bbox_to_anchor=[0.5,-.08],loc="upper center")
+    plt.tight_layout()
 
     
     return '', plt.gcf()