Refactor code to use gr.State to prevent race conditions with multiple users

regularization.py

@@ -30,10 +30,9 @@ logging.basicConfig(
 )
 logger = logging.getLogger("ELVIS")
 
+
 def min_corresponding_entries(W1, W2, w1, tol=0.1):
-    #mask = np.isclose(W1, w1, atol=tol, rtol=0)
     mask = (W1 <= w1)
-    #print(W1.max(), W1.min(), w1)
 
     values = W2[mask]
 
@@ -80,6 +79,19 @@ def l2_loss_regularization_path(y, X, regularization_type):
 
 
 class Regularization:
+    LOSS_TYPES = ['l1', 'l2']
+    REGULARIZER_TYPES = ['l1', 'l2']
+
+    LOSS_FUNCTIONS = {
+        'l1': l1_loss,
+        'l2': l2_loss,
+    }
+
+    REGULARIZER_FUNCTIONS = {
+        'l1': l1_norm,
+        'l2': l2_norm,
+    }
+
     def __init__(self, width, height):
         # initialized in draw_plot
         #self.canvas_width = -1
@@ -89,64 +101,120 @@ class Regularization:
         self.canvas_height = height
 
         self.css ="""
-#my-button {
-    height: 30px;
-    font-size: 16px;
-}
-
-#rowheight {
-    height: 90px;
-}
-
 .hidden-button {
     display: none;
 }
+"""
+
+    def compute_and_plot_loss_and_reg(
+        self, 
+        X: np.ndarray, 
+        y: np.ndarray, 
+        loss_type: str, 
+        reg_type: str, 
+        reg_levels: list,
+        w1_range: list, 
+        w2_range: list, 
+        num_dots: int, 
+        plot_path: bool,
+    ):
+        W1, W2 = self._build_parameter_grid(
+            w1_range, w2_range, num_dots
+        )
 
-.report-table {
-    border: 0 !important;
-}
-.report-table tr, .report-table th, .report-table td, .report-table tbody, .report-table thead {
-    border: 0 !important;
-    padding: 6px 12px;
-    text-align: center;
-}"""
-
-        # Regularization strengths
-        self.alphas = [0.01, 0.1, 1, 10, 100]
-
-        self.losses = {
-            #'l1': mean_absolute_error, # slow
-            #'l1': lambda y, pred: np.mean(abs(y - pred)),
-            # 'l1': l1_loss,
-            #'l2': mean_squared_error, # slow
-            #'l2': lambda y, pred: np.mean((y - pred)**2)
-            # 'l2': l2_loss
-            "l1": l1_loss,
-            "l2": l2_loss,
-        }
-
-        self.regularizers = {
-            # 'l1': lambda w: sum(abs(w)),
-            # 'l2': np.linalg.norm
-            "l1": l1_norm,
-            "l2": l2_norm,
-        }
-
-        #self.Model = Ridge #l2 loss + l2 reg
-        #self.Model = Lasso #l2 loss + l1 reg
-
-        self.loss_type = 'l2'
-        self.reg_type = 'l2'
-
-        self.loss = self.losses[self.loss_type]
-        self.regularizer = self.regularizers[self.reg_type]
-
-        self.reg_levels = [10, 20, 30]
-        self.w1_range = (-100, 100)
-        self.w2_range = (-100, 100)
-        self.num_dots = 500
-
-        self.plot_regularization_path = False
+        losses = self._compute_losses(
+            X, y, loss_type, W1, W2
+        )
+
+        reg_values = self._compute_reg_values(
+            W1, W2, reg_type
+        )
+
+        loss_levels = [
+            min_corresponding_entries(
+                reg_values, losses, reg_level
+            ) 
+            for reg_level in reg_levels
+        ]
+        loss_levels.reverse()
+
+        if plot_path:
+            if loss_type == "l2":
+                path_w = l2_loss_regularization_path(y, X, regularization_type=reg_type)
+            else:
+                min_loss_reg = reg_values.ravel()[np.argmin(losses)]
+                path_reg_levels = np.linspace(0, min_loss_reg, 20)
+                path_w = []
+                for reg_level in path_reg_levels:
+                    mask = reg_values <= reg_level
+                    if np.sum(mask) == 0:
+                        continue
+                    idx = np.argmin(losses[mask])
+                    path_w.append(
+                        np.stack((W1, W2), axis=-1)[mask][idx]
+                    )
+                
+                path_w = np.array(path_w)
+        else:
+            path_w = None
+
+        return self.plot_loss_and_reg(
+            W1, 
+            W2,
+            losses, 
+            reg_values, 
+            loss_levels, 
+            reg_levels,
+            path_w,
+        )
+        
+    def plot_loss_and_reg(
+        self, 
+        W1: np.ndarray,
+        W2: np.ndarray,
+        losses: np.ndarray, 
+        reg_values: np.ndarray, 
+        loss_levels: list, 
+        reg_levels: list,
+        path_w: np.ndarray | None,
+    ):
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.set_title("")
+        ax.set_xlabel("w1")
+        ax.set_ylabel("w2")
+
+        cmap = plt.get_cmap("viridis")   
+        N = len(reg_levels)
+        colors = [cmap(i / (N - 1)) for i in range(N)]
+
+        # regularizer contours
+        cs1 = ax.contour(W1, W2, reg_values, 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)
+
+        # regularization path
+        if path_w is not None:
+            ax.plot(path_w[:, 0], path_w[:, 1], "r-")
+
+        # legend
+        loss_line = mlines.Line2D([], [], color='black', linestyle='-', label='loss')
+        reg_line = mlines.Line2D([], [], color='black', linestyle='--', label='regularization')
+        handles = [loss_line, reg_line]
+        if path_w is not None:
+            path_line = mlines.Line2D([], [], color='red', linestyle='-', label='regularization path')
+            handles.append(path_line)
+        ax.legend(handles=handles)
+
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+        plt.close(fig)
+        buf.seek(0)
+        img = Image.open(buf)
+
+        return img
 
     def plot_regularization_contour(self):
         '''
@@ -264,25 +332,48 @@ class Regularization:
 
         return img
 
-    def plot_data(self):
-        # make sure the data is the same as the one used in plot_regularization_contour
-        _, _, coef = make_regression(n_samples=200, n_features=2, noise=15, random_state=0, coef=True)
+    def plot_data(self, X_data: np.ndarray, y_data: np.ndarray, coefs: np.ndarray):
+        x1_min = X_data[:, 0].min() - 1
+        x1_max = X_data[:, 0].max() + 1
+        x2_min = X_data[:, 1].min() - 1
+        x2_max = X_data[:, 1].max() + 1
 
-        x1 = np.linspace(-1, 1, 50)
-        x2 = np.linspace(-1, 1, 50)
+        x1 = np.linspace(x1_min, x1_max, 100)
+        x2 = np.linspace(x2_min, x2_max, 100)
         mesh_x1, mesh_x2 = np.meshgrid(x1, x2)
         X = np.stack((mesh_x1.ravel(), mesh_x2.ravel()), axis=-1)
-        y = X @ coef
+        y = X @ coefs
         
         z = y.reshape(mesh_x1.shape)
 
-        fig = go.Figure(data=go.Surface(
-            z=z, 
-            x=mesh_x1, 
-            y=mesh_x2,
-            colorscale='Viridis',
-            opacity=0.8
-        ))
+        fig = go.Figure()
+
+        fig.add_trace(
+            go.Surface(
+                z=z, 
+                x=mesh_x1, 
+                y=mesh_x2,
+                colorscale='Viridis',
+                opacity=0.8,
+                name='True function',
+            )
+        )
+
+        fig.add_trace(
+            go.Scatter3d(
+                x=X_data[:, 0],
+                y=X_data[:, 1],
+                z=y_data,
+                mode='markers',
+                marker=dict(
+                    size=3,
+                    color='red',
+                    opacity=0.8,
+                    symbol='circle',
+                ),
+                name='Data Points',
+            )
+        )
         
         fig.update_layout(
             title="Data",
@@ -297,12 +388,10 @@ class Regularization:
         )
         return fig
 
-    def plot_strength_vs_weight(self):
-        # make sure the data is the same as the one used in plot_regularization_contour
-        X, y = make_regression(n_samples=200, n_features=2, noise=15, random_state=0)
+    def plot_strength_vs_weight(self, X: np.ndarray, y: np.ndarray, loss_type: str, reg_type: str):
         alphas = np.concat([np.zeros(1), np.logspace(-2, 2, 100)])
-        if self.loss_type == "l2":
-            l1_ratio = 1 if self.reg_type == "l1" else 0
+        if loss_type == "l2":
+            l1_ratio = 1 if reg_type == "l1" else 0
             alphas, coefs, *_ = ElasticNet.path(X, y, l1_ratio=l1_ratio, alphas=alphas)
         else:
             coefs = np.random.randn(2, len(alphas))  # temporary
@@ -324,42 +413,68 @@ class Regularization:
 
         return img
 
-    def update_loss(self, loss_type):
-        self.loss_type = loss_type
-        self.loss = self.losses[loss_type]
-        return self.plot_regularization_contour(), self.plot_strength_vs_weight()
+    def update_loss_type(self, loss_type: str):
+        if loss_type not in self.LOSS_TYPES:
+            raise ValueError(f"loss_type must be one of {self.LOSS_TYPES}")
+        return loss_type
 
-    def update_regularizer(self, reg_type):
-        self.reg_type = reg_type
-        self.regularizer = self.regularizers[reg_type]
+    def update_regularizer(self, reg_type: str):
+        if reg_type not in self.REGULARIZER_TYPES:
+            raise ValueError(f"reg_type must be one of {self.REGULARIZER_TYPES}")
+        return reg_type
 
-        return self.plot_regularization_contour(), self.plot_strength_vs_weight()
+    def update_reg_levels(self, reg_levels_input: str):
+        reg_levels = [float(reg_level) for reg_level in reg_levels_input.split(",")]
+        return reg_levels
 
-    def update_reg_levels(self, reg_levels):
-        self.reg_levels = [float(reg_level) for reg_level in reg_levels.split(",")]
+    def update_w1_range(self, w1_range_input: str):
+        w1_range = [float(w1) for w1 in w1_range_input.split(",")]
+        return w1_range
 
-        return self.plot_regularization_contour()
+    def update_w2_range(self, w2_range_input: str):
+        w2_range = [float(w2) for w2 in w2_range_input.split(",")]
+        return w2_range
 
-    def update_w1_range(self, w1_range):
-        self.w1_range = [float(w1) for w1 in w1_range.split(",")]
-        logger.info("Updated w1 range to " + str(self.w1_range))
+    def update_resolution(self, num_dots: int):
+        return num_dots
 
-        return self.plot_regularization_contour()
+    def update_plot_path(self, plot_path: bool):
+        return plot_path
 
-    def update_w2_range(self, w2_range):
-        self.w2_range = [float(w2) for w2 in w2_range.split(",")]
-        logger.info("Updated w2 range to " + str(self.w2_range))
-
-        return self.plot_regularization_contour()
+    def _build_parameter_grid(
+        self, 
+        w1_range: list, 
+        w2_range: list, 
+        num_dots: int,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        # build grid in parameter space
+        w1 = np.linspace(w1_range[0], w1_range[1], num_dots)
+        w2 = np.linspace(w2_range[0], w2_range[1], num_dots)
+        W1, W2 = np.meshgrid(w1, w2)
 
-    def update_resolution(self, num_dots):
-        self.num_dots = num_dots
-        logger.info("updated resolution to " + str(num_dots))
-        return self.plot_regularization_contour()
+        return W1, W2
 
-    def update_plot_path(self, plot_path):
-        self.plot_regularization_path = plot_path
-        return self.plot_regularization_contour()
+    def _compute_losses(
+        self, 
+        X: np.ndarray, 
+        y: np.ndarray, 
+        loss_type: str, 
+        W1: np.ndarray,
+        W2: np.ndarray,
+    ) -> np.ndarray:
+        stacked = np.stack((W1, W2), axis=-1)
+        losses = self.LOSS_FUNCTIONS[loss_type](stacked, y, X)
+        return losses
+
+    def _compute_reg_values(
+        self, 
+        W1: np.ndarray,
+        W2: np.ndarray,
+        reg_type: str,
+    ) -> np.ndarray:
+        stacked = np.stack((W1, W2), axis=-1)
+        regs = self.REGULARIZER_FUNCTIONS[reg_type](stacked)
+        return regs
 
     def launch(self):
         # build the Gradio interface
@@ -367,57 +482,102 @@ class Regularization:
             # app title
             gr.HTML("<div style='text-align:left; font-size:40px; font-weight: bold;'>Regularization visualizer</div>")
 
+            # states
+            loss_type = gr.State("l2")
+            reg_type = gr.State("l2")
+            reg_levels = gr.State([10, 20, 30])
+            w1_range = gr.State([-100, 100])
+            w2_range = gr.State([-100, 100])
+            num_dots = gr.State(500)
+            plot_regularization_path = gr.State(False)
+
+            X, y, coefs = make_regression(
+                n_samples=200, n_features=2, noise=15, random_state=0, coef=True
+            )
+            X = gr.State(X)
+            y = gr.State(y)
+            coefs = gr.State(coefs)
+
             # GUI elements and layout 
             with gr.Row():
                 with gr.Column(scale=2):
-                    with gr.Tab("Regularization contour"):
-                        self.regularization_contour = gr.Image(value=self.plot_regularization_contour(), container=True)
+                    with gr.Tab("Loss and Regularization"):
+                        self.loss_and_regularization_plot = gr.Image(
+                            value=self.compute_and_plot_loss_and_reg(
+                                X.value,
+                                y.value,
+                                loss_type.value,
+                                reg_type.value,
+                                reg_levels.value,
+                                w1_range.value,
+                                w2_range.value,
+                                num_dots.value,
+                                plot_regularization_path.value,
+                            ), 
+                            container=True,
+                        )
                     with gr.Tab("Data"):
-                        self.data_3d_plot = gr.Plot(value=self.plot_data(), container=True)
+                        self.data_3d_plot = gr.Plot(
+                            value=self.plot_data(X.value, y.value, coefs.value), container=True
+                        )
                     with gr.Tab("Strength vs weight"):
-                        self.strength_vs_weight = gr.Image(value=self.plot_strength_vs_weight(), container=True)
+                        self.strength_vs_weight = gr.Image(
+                            value=self.plot_strength_vs_weight(
+                                X.value, y.value, loss_type.value, reg_type.value
+                            ), 
+                            container=True,
+                        )
 
                 with gr.Column(scale=1): 
                     with gr.Tab("Settings"):
-                        dataset_radio = gr.Radio(["make_regression", "Upload"], 
-                                                value="make_regression", label="Dataset type", elem_id="rowheight")
-
-                        # upload data
-                        file_chooser = gr.File(label="Choose a file", visible=False, elem_id="rowheight")
-                        self.file_chooser = file_chooser
-
-                        # loss type 
-                        loss_type = gr.Dropdown(choices=['l1', 'l2'],
-                                                label='Loss type',
-                                                value='l2',
-                                                visible=True)
+                        dataset_radio = gr.Radio(
+                            ["make_regression", "Upload"], 
+                            value="make_regression", 
+                            label="Dataset type", 
+                        )
+
+                        loss_type_selection = gr.Dropdown(
+                            choices=['l1', 'l2'],
+                            label='Loss type',
+                            value='l2',
+                            visible=True,
+                        )
 
                         with gr.Row():
-                            # regularizer type 
-                            regularizer_type = gr.Dropdown(choices=['l1', 'l2'],
-                                                           label='Regularizer type',
-                                                           value='l2',
-                                                           visible=True)
-
-                            # regularization strength
-                            #reg_textbox = gr.Textbox(label="Regularization constants")
-                            reg_textbox = gr.Textbox(label="Regularizer levels",
-                                                     value="10, 20, 30",
-                                                     interactive=True)
+                            regularizer_type_selection = gr.Dropdown(
+                                choices=['l1', 'l2'],
+                                label='Regularizer type',
+                                value='l2',
+                                visible=True,
+                            )
+
+                            reg_textbox = gr.Textbox(
+                                label="Regularizer levels",
+                                value="10, 20, 30",
+                                interactive=True,
+                            )
                             self.reg_textbox = reg_textbox
 
                         with gr.Row():
-                            # parameter value ranges
-                            w1_textbox = gr.Textbox(label="w1 range",
-                                                     value="-100, 100",
-                                                     interactive=True)
-
-                            w2_textbox = gr.Textbox(label="w2 range",
-                                                     value="-100, 100",
-                                                     interactive=True)
-
-                        # resolution
-                        slider = gr.Slider(minimum=100, maximum=1000, value=500, step=1, label="Resolution (#points)")
+                            w1_textbox = gr.Textbox(
+                                label="w1 range",
+                                value="-100, 100",
+                                interactive=True,
+                            )
+
+                            w2_textbox = gr.Textbox(
+                                label="w2 range",
+                                value="-100, 100",
+                                interactive=True,
+                            )
+
+                        resolution_slider = gr.Slider(
+                            minimum=100, 
+                            maximum=1000, 
+                            value=500, 
+                            step=1, 
+                            label="Resolution (#points)",
+                        )
 
                         # plot path
                         path_checkbox = gr.Checkbox(label="Show regularization path", value=False)
@@ -438,34 +598,172 @@ class Regularization:
 
                     with gr.Tab("Usage"):
                         gr.Markdown(''.join(open('usage.md', 'r').readlines()))
-
             
                     # event handlers for GUI elements
-                    loss_type.change(
-                        fn=self.update_loss, 
-                        inputs=loss_type,
-                        outputs=(self.regularization_contour, self.strength_vs_weight),
+                    loss_type_selection.change(
+                        fn=self.update_loss_type, 
+                        inputs=[loss_type_selection],
+                        outputs=[loss_type],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X, 
+                            y, 
+                            loss_type, 
+                            reg_type, 
+                            reg_levels, 
+                            w1_range, 
+                            w2_range, 
+                            num_dots,
+                            plot_regularization_path,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
+                    ).then(
+                        fn=self.plot_strength_vs_weight,
+                        inputs=[
+                            X, 
+                            y, 
+                            loss_type, 
+                            reg_type,
+                        ],
+                        outputs=self.strength_vs_weight,
                     )
 
-                    regularizer_type.change(
+                    regularizer_type_selection.change(
                         fn=self.update_regularizer,
-                        inputs=regularizer_type, 
-                        outputs=(self.regularization_contour, self.strength_vs_weight),
+                        inputs=[regularizer_type_selection], 
+                        outputs=[reg_type],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X, 
+                            y, 
+                            loss_type, 
+                            reg_type, 
+                            reg_levels, 
+                            w1_range, 
+                            w2_range, 
+                            num_dots,
+                            plot_regularization_path,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
+                    ).then(
+                        fn=self.plot_strength_vs_weight,
+                        inputs=[
+                            X, 
+                            y, 
+                            loss_type, 
+                            reg_type,
+                        ],
+                        outputs=self.strength_vs_weight,
                     )
 
-                    reg_textbox.submit(self.update_reg_levels, inputs=reg_textbox,
-                                       outputs=self.regularization_contour)
+                    reg_textbox.submit(
+                        self.update_reg_levels, 
+                        inputs=[reg_textbox],
+                        outputs=[reg_levels],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X, 
+                            y, 
+                            loss_type, 
+                            reg_type, 
+                            reg_levels, 
+                            w1_range, 
+                            w2_range,
+                            num_dots,
+                            plot_regularization_path,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
+                    ).then(
+                        fn=self.plot_strength_vs_weight,
+                        inputs=[
+                            X, 
+                            y, 
+                            loss_type, 
+                            reg_type,
+                        ],
+                        outputs=self.strength_vs_weight,
+                    )
 
-                    w1_textbox.submit(self.update_w1_range, inputs=w1_textbox,
-                                       outputs=self.regularization_contour)
+                    w1_textbox.submit(
+                        self.update_w1_range, 
+                        inputs=[w1_textbox],
+                        outputs=[w1_range],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X,
+                            y,
+                            loss_type,
+                            reg_type,
+                            reg_levels,
+                            w1_range,
+                            w2_range,
+                            num_dots,
+                            plot_regularization_path,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
+                    )
 
-                    w2_textbox.submit(self.update_w2_range, inputs=w2_textbox,
-                                       outputs=self.regularization_contour)
+                    w2_textbox.submit(
+                        self.update_w2_range, 
+                        inputs=[w2_textbox],
+                        outputs=[w2_range],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X,
+                            y,
+                            loss_type,
+                            reg_type,
+                            reg_levels,
+                            w1_range,
+                            w2_range,
+                            num_dots,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
+                    )
 
-                    slider.change(self.update_resolution, inputs=slider, outputs=self.regularization_contour)
+                    resolution_slider.change(
+                        self.update_resolution, 
+                        inputs=[resolution_slider], 
+                        outputs=[num_dots],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X,
+                            y,
+                            loss_type,
+                            reg_type,
+                            reg_levels,
+                            w1_range,
+                            w2_range,
+                            num_dots,
+                            plot_regularization_path,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
+                    )
 
                     path_checkbox.change(
-                        self.update_plot_path, inputs=path_checkbox, outputs=self.regularization_contour
+                        self.update_plot_path, 
+                        inputs=[path_checkbox], 
+                        outputs=[plot_regularization_path],
+                    ).then(
+                        fn=self.compute_and_plot_loss_and_reg,
+                        inputs=[
+                            X,
+                            y,
+                            loss_type,
+                            reg_type,
+                            reg_levels,
+                            w1_range,
+                            w2_range,
+                            num_dots,
+                            plot_regularization_path,
+                        ],
+                        outputs=self.loss_and_regularization_plot,
                     )
 
         demo.launch()