Refactor to use a backend manager

backend/src/manager.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+import matplotlib.lines as mlines
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.figure import Figure
+from sympy import sympify
+
+from logic import (
+    DataGenerationOptions,
+    Dataset,
+    PlotsData,
+    compute_plot_values,
+    generate_dataset,
+    load_dataset_from_csv,
+)
+
+
+class Manager:
+    def __init__(self, dataset: Dataset | None = None, plots_data: PlotsData | None = None):
+        self.dataset = dataset
+        self.plots_data = plots_data
+
+    def update_dataset(
+        self,
+        dataset_type: str,
+        function: str,
+        x1_range_input: str,
+        x2_range_input: str,
+        x_selection_method: str,
+        sigma: float,
+        nsample: int,
+        csv_file: str,
+        has_header: bool,
+        x1_col: int,
+        x2_col: int,
+        y_col: int,
+    ) -> None:
+        if dataset_type == "Generate":
+            try:
+                parsed_function = sympify(function)
+            except Exception as e:
+                raise ValueError(f"Invalid function {e}")
+
+            x1_range = self._parse_range(x1_range_input)
+            x2_range = self._parse_range(x2_range_input)
+
+            method = x_selection_method.lower()
+            if method not in ("grid", "random"):
+                raise ValueError(f"Invalid x_selection_method: {x_selection_method}")
+
+            self.dataset = generate_dataset(
+                parsed_function,
+                x1_range,
+                x2_range,
+                DataGenerationOptions(method, int(nsample), float(sigma)),
+            )
+
+        elif dataset_type == "CSV":
+            csv_path = self._resolve_csv_path(csv_file)
+            try:
+                self.dataset = load_dataset_from_csv(
+                    csv_path,
+                    bool(has_header),
+                    int(x1_col),
+                    int(x2_col),
+                    int(y_col),
+                )
+            except Exception as e:
+                raise ValueError(f"Failed to load dataset from CSV: {e}")
+
+        else:
+            raise ValueError(f"Invalid dataset_type: {dataset_type}")
+
+
+    def compute_plots_data(
+        self,
+        loss_type: str,
+        regularizer_type: str,
+        reg_levels_input: str,
+        w1_range_input: str,
+        w2_range_input: str,
+        resolution: int,
+    ) -> None:
+        if self.dataset is None:
+            raise ValueError("Dataset is not initialized")
+
+        if loss_type not in ("l1", "l2"):
+            raise ValueError(f"Invalid loss_type: {loss_type}")
+        if regularizer_type not in ("l1", "l2"):
+            raise ValueError(f"Invalid regularizer_type: {regularizer_type}")
+
+        reg_levels = self._parse_levels(reg_levels_input)
+        w1_range = self._parse_range(w1_range_input)
+        w2_range = self._parse_range(w2_range_input)
+
+        self.plots_data = compute_plot_values(
+            self.dataset,
+            loss_type,
+            regularizer_type,
+            reg_levels,
+            w1_range,
+            w2_range,
+            int(resolution),
+        )
+
+    def handle_generate_plots(
+        self,
+        dataset_type: str,
+        function: str,
+        x1_range_input: str,
+        x2_range_input: str,
+        x_selection_method: str,
+        sigma: float,
+        nsample: int,
+        csv_file: str,
+        has_header: bool,
+        x1_col: int,
+        x2_col: int,
+        y_col: int,
+        loss_type: str,
+        regularizer_type: str,
+        reg_levels_input: str,
+        w1_range_input: str,
+        w2_range_input: str,
+        resolution: int,
+    ) -> tuple[Manager, Figure, Figure, Figure]:
+        self.update_dataset(
+            dataset_type,
+            function,
+            x1_range_input,
+            x2_range_input,
+            x_selection_method,
+            sigma,
+            nsample,
+            csv_file,
+            has_header,
+            x1_col,
+            x2_col,
+            y_col,
+        )
+        self.compute_plots_data(
+            loss_type,
+            regularizer_type,
+            reg_levels_input,
+            w1_range_input,
+            w2_range_input,
+            resolution,
+        )
+
+        if self.dataset is None or self.plots_data is None:
+            raise ValueError("Failed to generate plot data")
+
+        contour_plot = self._generate_contour_plot(self.plots_data)
+        data_plot = self._generate_data_plot(self.dataset)
+        strength_plot = self._generate_strength_plot(self.plots_data.path)
+        return self, contour_plot, data_plot, strength_plot
+
+    @staticmethod
+    def _generate_contour_plot(plots_data: PlotsData) -> Figure:
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.set_xlabel("w1")
+        ax.set_ylabel("w2")
+
+        cmap = plt.get_cmap("viridis")
+        n_levels = len(plots_data.reg_levels)
+        if n_levels == 1:
+            colors = [cmap(0.5)]
+        else:
+            colors = [cmap(i / (n_levels - 1)) for i in range(n_levels)]
+
+        cs1 = ax.contour(
+            plots_data.W1,
+            plots_data.W2,
+            plots_data.norms,
+            levels=plots_data.reg_levels,
+            colors=colors,
+            linestyles="dashed",
+        )
+        ax.clabel(cs1, inline=True, fontsize=8)
+
+        cs2 = ax.contour(
+            plots_data.W1,
+            plots_data.W2,
+            plots_data.loss_values,
+            levels=plots_data.loss_levels,
+            colors=colors[::-1],
+        )
+        ax.clabel(cs2, inline=True, fontsize=8)
+
+        if plots_data.unreg_solution.ndim == 1:
+            ax.plot(
+                plots_data.unreg_solution[0],
+                plots_data.unreg_solution[1],
+                "bx",
+                markersize=5,
+                label="unregularized solution",
+            )
+        else:
+            ax.plot(
+                plots_data.unreg_solution[:, 0],
+                plots_data.unreg_solution[:, 1],
+                "b-",
+                label="unregularized solution",
+            )
+
+        ax.plot(plots_data.path[:, 0], plots_data.path[:, 1], "r-", label="regularization path")
+
+        handles = [
+            mlines.Line2D([], [], color="black", linestyle="-", label="loss"),
+            mlines.Line2D([], [], color="black", linestyle="--", label="regularization"),
+            mlines.Line2D([], [], color="red", linestyle="-", label="regularization path"),
+        ]
+        if plots_data.unreg_solution.ndim == 1:
+            handles.append(
+                mlines.Line2D([], [], color="blue", marker="x", linestyle="None", label="unregularized solution")
+            )
+        else:
+            handles.append(mlines.Line2D([], [], color="blue", linestyle="-", label="unregularized solution"))
+
+        ax.legend(handles=handles)
+        ax.grid(True)
+        return fig
+
+    @staticmethod
+    def _generate_data_plot(dataset: Dataset) -> Figure:
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.set_xlabel("x1")
+        ax.set_ylabel("x2")
+
+        scatter = ax.scatter(dataset.x1, dataset.x2, c=dataset.y, cmap="viridis")
+        ax.grid(True)
+        fig.colorbar(scatter, ax=ax)
+        return fig
+
+    @staticmethod
+    def _generate_strength_plot(path: np.ndarray) -> Figure:
+        reg_levels = np.logspace(-4, 4, path.shape[0])
+
+        fig, ax = plt.subplots(figsize=(8, 6))
+        ax.set_xlabel("Regularization Strength")
+        ax.set_ylabel("Weight")
+
+        ax.plot(reg_levels, path[:, 0], "r-", label="w1")
+        ax.plot(reg_levels, path[:, 1], "b-", label="w2")
+        ax.set_xscale("log")
+        ax.legend()
+        ax.grid(True)
+        return fig
+
+    @staticmethod
+    def _parse_range(range_input: str) -> tuple[float, float]:
+        values = tuple(float(x.strip()) for x in range_input.split(","))
+        if len(values) != 2:
+            raise ValueError("Range must contain exactly two comma-separated values")
+        return values
+
+    @staticmethod
+    def _parse_levels(levels_input: str) -> list[float]:
+        values = [float(x.strip()) for x in levels_input.split(",")]
+        if not values:
+            raise ValueError("At least one regularization level is required")
+        return values
+
+    @staticmethod
+    def _resolve_csv_path(csv_file: str) -> str:
+        if csv_file is None:
+            raise ValueError("CSV file is required")
+        if isinstance(csv_file, str):
+            return csv_file
+        if isinstance(csv_file, dict) and "name" in csv_file:
+            return csv_file["name"]
+        if hasattr(csv_file, "name"):
+            return csv_file.name
+        raise ValueError("Unsupported CSV file input")

frontends/gradio/main.py

@@ -1,28 +1,17 @@
-import io
 from typing import Literal
 
 import gradio as gr
 from matplotlib.figure import Figure
-import matplotlib.pyplot as plt
-import matplotlib.lines as mlines
-import numpy as np
-from sympy import sympify
 
 import sys
 from pathlib import Path
+
 root_dir = Path(__file__).resolve().parent.parent.parent
 backend_src = root_dir / "backend" / "src"
 if str(backend_src) not in sys.path:
     sys.path.append(str(backend_src))
 
-from backend import (
-    compute_plot_values, 
-    generate_dataset, 
-    load_dataset_from_csv, 
-    Dataset,
-    DataGenerationOptions,
-    PlotsData,
-)
+from manager import Manager
 
 CSS = """
 .hidden-button {
@@ -31,7 +20,38 @@ CSS = """
 """
 
 
-def get_dataset(
+def handle_dataset_type_change(dataset_type: Literal["Generate", "CSV"]):
+    if dataset_type == "Generate":
+        return (
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    return (
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=True),
+        gr.update(visible=True),
+        gr.update(visible=True),
+        gr.update(visible=True),
+        gr.update(visible=True),
+    )
+
+
+def handle_generate_plots(
+    manager: Manager,
     dataset_type: str,
     function: str,
     x1_range_input: str,
@@ -44,232 +64,35 @@ def get_dataset(
     x1_col: int,
     x2_col: int,
     y_col: int,
-) -> Dataset:
-    if dataset_type == "Generate":
-        try:
-            function = sympify(function)
-        except Exception as e:
-            raise ValueError(f"Invalid function: {e}")
-
-        x1_range = tuple(float(x.strip()) for x in x1_range_input.split(","))
-        x2_range = tuple(float(x.strip()) for x in x2_range_input.split(","))
-
-        if (len(x1_range) != 2 or len(x2_range) != 2):
-            raise ValueError("x1_range and x2_range must be tuples of length 2")
-        
-        x_selection_method = x_selection_method.lower()
-        if x_selection_method not in ("grid", "random"):
-            raise ValueError(f"Invalid x_selection_method: {x_selection_method}")
-
-        dataset = generate_dataset(
-            function,
-            x1_range,
-            x2_range,
-            DataGenerationOptions(
-                x_selection_method,
-                nsample,
-                sigma,
-            )
-        )
-
-    elif dataset_type == "CSV":
-        try:
-            dataset = load_dataset_from_csv(
-                csv_file,
-                has_header,
-                x1_col,
-                x2_col,
-                y_col,
-            )
-        except Exception as e:
-            gr.Info(f"Error loading CSV: {e}")
-            raise e
-
-    else:
-        raise ValueError(f"Invalid dataset_type: {dataset_type}")
-
-    return dataset
-
-
-def parse_plot_settings(
-    dataset: Dataset,
     loss_type: str,
     regularizer_type: str,
     reg_levels_input: str,
     w1_range_input: str,
     w2_range_input: str,
     resolution: int,
-) -> tuple[Dataset, Literal["l1", "l2"], Literal["l1", "l2"], list[float], tuple[float, float], tuple[float, float], int]:
-    reg_levels = [float(x.strip()) for x in reg_levels_input.split(",")]
-    w1_range = tuple(float(x.strip()) for x in w1_range_input.split(","))
-    w2_range = tuple(float(x.strip()) for x in w2_range_input.split(","))
-
-    if loss_type not in ("l1", "l2"):
-        raise ValueError(f"Invalid loss_type: {loss_type}")
-    if regularizer_type not in ("l1", "l2"):
-        raise ValueError(f"Invalid regularizer_type: {regularizer_type}")
-
-    if len(w1_range) != 2 or len(w2_range) != 2:
-        raise ValueError("w1_range and w2_range must be tuples of length 2")
-
-    return (
-        dataset,
+) -> tuple[Manager, Figure, Figure, Figure]:
+    return manager.handle_generate_plots(
+        dataset_type,
+        function,
+        x1_range_input,
+        x2_range_input,
+        x_selection_method,
+        sigma,
+        nsample,
+        csv_file,
+        has_header,
+        x1_col,
+        x2_col,
+        y_col,
         loss_type,
         regularizer_type,
-        reg_levels,
-        w1_range,
-        w2_range,
+        reg_levels_input,
+        w1_range_input,
+        w2_range_input,
         resolution,
     )
 
 
-def generate_contour_plot(
-    W1: np.ndarray,
-    W2: np.ndarray,
-    losses: np.ndarray,
-    norms: np.ndarray,
-    loss_levels: list[float],
-    reg_levels: list[float],
-    unreg_solution: np.ndarray,
-    path: np.ndarray,
-) -> Figure:
-    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, norms, 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)
-
-    # unregularized solution
-    if unreg_solution.ndim == 1:
-        ax.plot(unreg_solution[0], unreg_solution[1], "bx", markersize=5, label="unregularized solution")
-    else:
-        ax.plot(unreg_solution[:, 0], unreg_solution[:, 1], "b-", label="unregularized solution")
-
-    ax.plot(path[:, 0], path[:, 1], "r-", label="regularization path")
-
-    # legend
-    loss_line = mlines.Line2D([], [], color='black', linestyle='-', label='loss')
-    reg_line = mlines.Line2D([], [], color='black', linestyle='--', label='regularization')
-    handles = [loss_line, reg_line]
-    
-    path_line = mlines.Line2D([], [], color='red', linestyle='-', label='regularization path')
-    handles.append(path_line)
-
-    if unreg_solution.ndim == 1:
-        handles.append(
-            mlines.Line2D([], [], color='blue', marker='x', linestyle='None', label='unregularized solution')
-        )
-    else:
-        handles.append(
-            mlines.Line2D([], [], color='blue', linestyle='-', label='unregularized solution')
-        )
-
-    ax.legend(handles=handles)
-
-    ax.grid(True)
-
-    return fig
-
-
-def generate_data_plot(dataset: Dataset) -> Figure:
-    fig, ax = plt.subplots(figsize=(8, 8))
-    ax.set_xlabel("x1")
-    ax.set_ylabel("x2")
-
-    sc = ax.scatter(dataset.x1, dataset.x2, c=dataset.y, cmap='viridis')
-    ax.grid(True)
-    fig.colorbar(sc, ax=ax)
-
-    return fig
-
-
-def generate_strength_plot(
-    path: np.ndarray, 
-    reg_levels: np.ndarray,
-):
-    fig, ax = plt.subplots(figsize=(8, 6))
-    ax.set_xlabel("Regularization Strength")
-    ax.set_ylabel("Weight")
-
-    ax.plot(reg_levels, path[:, 0], 'r-', label='w1')
-    ax.plot(reg_levels, path[:, 1], 'b-', label='w2')
-
-    ax.set_xscale('log')
-    ax.legend()
-    ax.grid(True)
-
-    return fig
-
-
-def generate_all_plots(
-    dataset: Dataset,
-    plots_data: PlotsData,
-) -> tuple[Figure, Figure]:
-    contour_plot = generate_contour_plot(
-        plots_data.W1,
-        plots_data.W2,
-        plots_data.loss_values,
-        plots_data.norms,
-        plots_data.loss_levels,
-        plots_data.reg_levels,
-        plots_data.unreg_solution,
-        plots_data.path,
-    )
-
-    data_plot = generate_data_plot(
-        dataset
-    )
-
-    strength_plot = generate_strength_plot(
-        plots_data.path,
-        np.logspace(-4, 4, 100),
-    )
-
-    return contour_plot, data_plot, strength_plot
-
-
-def handle_dataset_type_change(dataset_type: Literal["Generate", "CSV"]):
-    if dataset_type == "Generate":
-        return (
-            gr.update(visible=True),  # function
-            gr.update(visible=True),  # x1_textbox
-            gr.update(visible=True),  # x2_textbox
-            gr.update(visible=True),  # x_selection_method
-            gr.update(visible=True),  # sigma
-            gr.update(visible=True),  # nsample
-            gr.update(visible=False), # csv_file
-            gr.update(visible=False), # has_header
-            gr.update(visible=False), # x1_col
-            gr.update(visible=False), # x2_col
-            gr.update(visible=False), # y_col
-        )
-    else:  # CSV
-        return (
-            gr.update(visible=False), # function
-            gr.update(visible=False), # x1_textbox
-            gr.update(visible=False), # x2_textbox
-            gr.update(visible=False), # x_selection_method
-            gr.update(visible=False), # sigma
-            gr.update(visible=False), # nsample
-            gr.update(visible=True),  # csv_file
-            gr.update(visible=True),  # has_header
-            gr.update(visible=True),  # x1_col
-            gr.update(visible=True),  # x2_col
-            gr.update(visible=True),  # y_col
-        )
-
-
 def launch():
     default_dataset_type = "Generate"
 
@@ -293,71 +116,41 @@ def launch():
     default_w2_range = "-100, 100"
     default_resolution = 100
 
+    manager = Manager()
+    manager, default_contour_plot, default_data_plot, default_strength_plot = manager.handle_generate_plots(
+        default_dataset_type,
+        default_function,
+        default_x1_range,
+        default_x2_range,
+        default_x_selection_method,
+        default_sigma,
+        default_num_points,
+        default_csv_file,
+        default_has_header,
+        default_x1_col,
+        default_x2_col,
+        default_y_col,
+        default_loss_type,
+        default_regularizer_type,
+        default_reg_levels,
+        default_w1_range,
+        default_w2_range,
+        default_resolution,
+    )
+
     with gr.Blocks() as demo:
         gr.HTML("<div style='text-align:left; font-size:40px; font-weight: bold;'>Regularization visualizer</div>")
 
-        dataset = gr.State(
-            get_dataset(
-                default_dataset_type,
-                default_function,
-                default_x1_range,
-                default_x2_range,
-                default_x_selection_method,
-                default_sigma,
-                default_num_points,
-                default_csv_file,
-                default_has_header,
-                default_x1_col,
-                default_x2_col,
-                default_y_col,
-            )
-        )
-
-        # wrapped PlotsData
-        plots_data = gr.State(
-            compute_plot_values(
-                *parse_plot_settings(
-                    dataset.value,
-                    default_loss_type,
-                    default_regularizer_type,
-                    default_reg_levels,
-                    default_w1_range,
-                    default_w2_range,
-                    default_resolution,
-                )
-            )
-        )
+        manager_state = gr.State(manager)
 
         with gr.Row():
             with gr.Column(scale=2):
                 with gr.Tab("Contours"):
-                    main_plot = gr.Plot(
-                        value=generate_contour_plot(
-                            plots_data.value.W1,
-                            plots_data.value.W2,
-                            plots_data.value.loss_values,
-                            plots_data.value.norms,
-                            plots_data.value.loss_levels,
-                            plots_data.value.reg_levels,
-                            plots_data.value.unreg_solution,
-                            plots_data.value.path,
-                        )
-                    )
+                    main_plot = gr.Plot(value=default_contour_plot)
                 with gr.Tab("Data"):
-                    # todo
-                    data_plot = gr.Plot(
-                        value=generate_data_plot(
-                            dataset.value
-                        )
-                    )
+                    data_plot = gr.Plot(value=default_data_plot)
                 with gr.Tab("Strength"):
-                    # todo
-                    strength_plot = gr.Plot(
-                        value=generate_strength_plot(
-                            plots_data.value.path,
-                            np.logspace(-4, 4, 100),  # todo
-                        )
-                    )
+                    strength_plot = gr.Plot(value=default_strength_plot)
 
             with gr.Column(scale=1):
                 with gr.Tab("Data"):
@@ -371,7 +164,7 @@ def launch():
 
                     with gr.Row():
                         function = gr.Textbox(
-                            label="Function (in terms of x1 and x2)", 
+                            label="Function (in terms of x1 and x2)",
                             value=default_function,
                             interactive=True,
                         )
@@ -401,28 +194,28 @@ def launch():
                             interactive=True,
                         )
                         nsample = gr.Number(
-                            label="Number of points", 
+                            label="Number of points",
                             value=default_num_points,
                             interactive=True,
                         )
 
                     with gr.Row():
                         csv_file = gr.File(
-                            label="Upload CSV file - must have columns: (x1, x2, y)", 
-                            file_types=['.csv'],
-                            visible=False,  # function mode is default
+                            label="Upload CSV file - must have columns: (x1, x2, y)",
+                            file_types=[".csv"],
+                            visible=False,
                         )
 
                     with gr.Row():
                         has_header = gr.Checkbox(
-                            label="CSV has header row", 
+                            label="CSV has header row",
                             value=default_has_header,
                             visible=False,
                         )
-                    
+
                     with gr.Row():
                         x1_col = gr.Number(
-                            label="x1 column index (0-based)", 
+                            label="x1 column index (0-based)",
                             value=default_x1_col,
                             visible=False,
                         )
@@ -436,7 +229,7 @@ def launch():
                             value=default_y_col,
                             visible=False,
                         )
-                    
+
                     dataset_type.change(
                         fn=handle_dataset_type_change,
                         inputs=[dataset_type],
@@ -499,12 +292,9 @@ def launch():
                     )
 
                 gr.Button("Regenerate Plots").click(
-                    fn=lambda: (gr.Plot(), gr.Plot(), gr.Plot()),
-                    inputs=[],
-                    outputs=[main_plot, data_plot, strength_plot],
-                ).then(
-                    fn=get_dataset,
+                    fn=handle_generate_plots,
                     inputs=[
+                        manager_state,
                         dataset_type,
                         function,
                         x1_textbox,
@@ -517,12 +307,6 @@ def launch():
                         x1_col,
                         x2_col,
                         y_col,
-                    ],
-                    outputs=[dataset],
-                ).then(
-                    fn=lambda *args: compute_plot_values(*parse_plot_settings(*args)),
-                    inputs=[
-                        dataset,
                         loss_type_dropdown,
                         regularizer_type_dropdown,
                         regularizer_levels_textbox,
@@ -530,11 +314,7 @@ def launch():
                         w2_range_textbox,
                         resolution_slider,
                     ],
-                    outputs=[plots_data],
-                ).then(
-                    fn=generate_all_plots,
-                    inputs=[dataset, plots_data],
-                    outputs=[main_plot, data_plot, strength_plot],
+                    outputs=[manager_state, main_plot, data_plot, strength_plot],
                 )
 
     demo.launch(css=CSS)