Distinguish loss and regularization contours and add legend

regularization.py

@@ -4,6 +4,7 @@ import pickle
 
 import gradio as gr
 import matplotlib.pyplot as plt
+import matplotlib.lines as mlines
 import inspect
 import numpy as np
 import pandas as pd
@@ -175,13 +176,18 @@ class Regularization:
         colors = [cmap(i / (N - 1)) for i in range(N)]
 
         # regularizer contours
-        cs1 = ax.contour(W1, W2, regs, levels=reg_levels, colors=colors)
+        cs1 = ax.contour(W1, W2, regs, levels=reg_levels, colors=colors, linestyles="dashed")
         ax.clabel(cs1, inline=True, fontsize=8) # show contour levels
 
         # loss contours
         cs2 = ax.contour(W1, W2, losses, levels=loss_levels, colors=colors[::-1])
         ax.clabel(cs2, inline=True, fontsize=8)
 
+        # custom legend
+        loss_line = mlines.Line2D([], [], color='black', linestyle='-', label='loss')
+        reg_line = mlines.Line2D([], [], color='black', linestyle='--', label='regularization')
+        ax.legend(handles=[loss_line, reg_line])
+
         # plot solutions
         #for alpha, w, norm, mse in solutions:
             #ax.plot(w[0], w[1], "ro")