Implement code export for plot reproduction

- note: support for passing model arguments is not yet implemented

decision_boundary.py

@@ -1,9 +1,14 @@
+from collections import deque
+from pathlib import Path
+import pickle
+
 import gradio as gr
 import matplotlib.pyplot as plt
 import inspect
 import numpy as np
 import pandas as pd
 import io
+from jinja2 import Template
 from PIL import Image
 import sklearn
 from sklearn.linear_model import LogisticRegression
@@ -13,9 +18,7 @@ from sklearn.datasets import load_iris
 from sklearn.decomposition import PCA
 from sklearn.metrics import classification_report
 import traceback
-
-from collections import deque
-import pickle
+import yaml
 
 from util import *
 
@@ -60,14 +63,17 @@ class CoordinateProjection:
     def __init__(self, n_components=None, dims=[0, 1]):
         self.dims = dims
 
-    def fit_transform(self, X):
-        print(X)
+    def fit(self, X):
         self.mean = X.mean(axis=0)
-        return X[:, self.dims]
+        return self
 
     def transform(self, X):
         return X[:, self.dims]
 
+    def fit_transform(self, X):
+        self.fit(X)
+        return self.transform(X)
+
     def inverse_transform(self, Z):
         X = np.ones((len(Z), 1)) * self.mean
         X[:, self.dims] = Z
@@ -77,6 +83,9 @@ class InteractiveDecisionBoundary:
     DATASET_FILE = "dataset.csv"
     MODEL_FILE = "model.pkl"
 
+    CODE_FILE = "generated_code.py"
+    EXPORT_CODE_TEMPLATE = "export_code_template.py.j2"
+
     def __init__(self, width, height):
         # initialized in draw_plot
         #self.canvas_width = -1
@@ -97,13 +106,15 @@ class InteractiveDecisionBoundary:
         for cls_name, cls in inspect.getmembers(module, inspect.isclass):
             self.embedders[cls_name] = cls
 
-        self.Embedder = CoordinateProjection
-        #self.Embedder = PCA
+        self.embedder_class = CoordinateProjection
+        #self.embedder_class = PCA
+        self.embedder = self.embedder_class()
 
         # default classifier model
-        #self.model = LogisticRegression
-        self.model = LinearSVC
+        #self.model_class = LogisticRegression
+        self.model_class = LinearSVC
         self.model_args = ""
+        self.model = self.model_class()
 
         # todo: support arbitrary number of classes and user-defined class labels 
         #self.dataset = toydata()
@@ -168,7 +179,7 @@ class InteractiveDecisionBoundary:
                 logger.info("Target:\n" + str(y)) 
 
                 # compute embedding
-                embedder = self.Embedder(n_components=2)
+                embedder = self.embedder_class(n_components=2)
                 self.embedder = embedder
                 logger.info("Embedder = " + str(self.embedder))
 
@@ -192,7 +203,7 @@ class InteractiveDecisionBoundary:
 
                 # plot the decision boundary
                 if decision_boundary:
-                    model = self.model(**parse_param_string(self.model_args))
+                    model = self.model_class(**parse_param_string(self.model_args))
                     model.fit(X, y)
                     self.model = model
 
@@ -290,9 +301,9 @@ class InteractiveDecisionBoundary:
         return self.DATASET_FILE
 
     def update_model(self, classifier_name):
-        self.model = self.classifiers[classifier_name]
+        self.model_class = self.classifiers[classifier_name]
         self.args_textbox.value = ""
-        logger.info(f'Updated model to {self.model}')
+        logger.info(f'Updated model to {self.model_class}')
 
         return ""
 
@@ -302,13 +313,64 @@ class InteractiveDecisionBoundary:
         logger.info(f"{self.MODEL_FILE} updated")
         return self.MODEL_FILE
 
+    def save_code(self):
+        model_class = str(self.model_class.__name__)
+        model_imports = yaml.safe_load(Path("model_imports.yaml").read_text())
+        if model_class not in model_imports:
+            raise ValueError(f"Model {model_class} not found in model_imports.yaml")
+        model_import_stmt = model_imports[model_class]
+
+        embedder_class = str(self.embedder_class.__name__)
+        if embedder_class == "CoordinateProjection":
+            embedder_import_stmt = f"\n{inspect.getsource(CoordinateProjection)}"
+        else:
+            embedder_imports = yaml.safe_load(Path("embedder_imports.yaml").read_text())
+            if embedder_class not in embedder_imports:
+                raise ValueError(f"Embedder {embedder_class} not found in embedder_imports.yaml")
+            embedder_import_stmt = embedder_imports[embedder_class]
+
+        if self.dataset_type == 'Draw2D':
+            x_min = 0
+            x_max = 1
+            y_min = 0
+            y_max = 1
+        else:
+            X = self.dataset.loc[:, self.dataset.columns != 'target'].values
+            Z = self.embedder.fit_transform(X)
+            x_min = Z[:, 0].min()
+            x_max = Z[:, 0].max()
+            y_min = Z[:, 1].min()
+            y_max = Z[:, 1].max()
+
+        template = Template(Path(self.EXPORT_CODE_TEMPLATE).read_text())
+        variables = {
+            'model_import_statement': model_import_stmt,
+            'embedder_import_statement': embedder_import_stmt,
+            'dataset_file': self.DATASET_FILE,
+            'embedder_class': embedder_class,
+            'model_class': model_class,
+            'fig_width': self.canvas_width / 100,
+            'fig_height': self.canvas_height / 100,
+            'dpi': 100,
+            'num_dots': self.num_dots,
+            'x_min': x_min,
+            'x_max': x_max,
+            'y_min': y_min,
+            'y_max': y_max,
+        }
+
+        rendered_code = template.render(variables)
+        Path(self.CODE_FILE).write_text(rendered_code)
+        logger.info(f"{self.CODE_FILE} updated")
+        return self.CODE_FILE
+
     def update_args(self, model_args):
         self.model_args = model_args
         print('updated model_args:', self.model_args)
 
     def update_embedder(self, embedder):
-        self.Embedder = self.embedders[embedder]
-        print('updated embedder:', self.Embedder)
+        self.embedder_class = self.embedders[embedder]
+        print('updated Embedder:', self.embedder_class)
         return self.plot()
 
     def handle_dataset_radio(self, type):
@@ -446,6 +508,7 @@ class InteractiveDecisionBoundary:
                         btn_export_model_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_model_hidden", elem_classes="hidden-button")
 
                         btn_export_code = gr.Button('Code')
+                        btn_export_code_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_code_hidden", elem_classes="hidden-button")
 
                     with gr.Tab("Options"):
                         slider = gr.Slider(minimum=100, maximum=1000, value=100, step=1, label="Resolution (#points)")
@@ -500,6 +563,13 @@ class InteractiveDecisionBoundary:
             ).then(
                 fn=None, inputs=None, outputs=None, js="() => document.querySelector('#btn_export_model_hidden').click()"
             )
+            btn_export_code.click(
+                fn=self.save_code, 
+                inputs=None, 
+                outputs=[btn_export_code_hidden]
+            ).then(
+                fn=None, inputs=None, outputs=None, js="() => document.querySelector('#btn_export_code_hidden').click()"
+            )
 
 
         demo.launch()

embedder_imports.yaml

@@ -0,0 +1,15 @@
+GaussianRandomProjection: from sklearn.random_projection import GaussianRandomProjection
+SparseRandomProjection: from sklearn.random_projection import SparseRandomProjection
+DictionaryLearning: from sklearn.decomposition import DictionaryLearning
+FactorAnalysis: from sklearn.decomposition import FactorAnalysis
+FastICA: from sklearn.decomposition import FastICA
+IncrementalPCA: from sklearn.decomposition import IncrementalPCA
+KernelPCA: from sklearn.decomposition import KernelPCA
+LatentDirichletAllocation: from sklearn.decomposition import LatentDirichletAllocation
+MiniBatchDictionaryLearning: from sklearn.decomposition import MiniBatchDictionaryLearning
+MiniBatchNMF: from sklearn.decomposition import MiniBatchNMF
+MiniBatchSparsePCA: from sklearn.decomposition import MiniBatchSparsePCA
+NMF: from sklearn.decomposition import NMF
+PCA: from sklearn.decomposition import PCA
+SparsePCA: from sklearn.decomposition import SparsePCA
+TruncatedSVD: from sklearn.decomposition import TruncatedSVD

export_code_template.py.j2

@@ -0,0 +1,76 @@
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+{{ model_import_statement }}
+{{ embedder_import_statement }}
+
+
+def label2color(labels):
+    cmap = plt.get_cmap("tab10")
+    named_colors = set(matplotlib.colors.CSS4_COLORS.keys())
+    colors = []
+    for i, label in enumerate(labels):
+        if label.lower() in named_colors:
+            colors.append(label.lower())
+        else:
+            colors.append(cmap(i % 10))
+    return colors
+
+
+def load_dataset(file_path):
+    """Load dataset from a CSV file."""
+    data = pd.read_csv(file_path)
+    X = data.loc[:, data.columns != 'target'].values
+    y = data['target'].values.astype(str)
+
+    if len(X) == 0:
+        raise ValueError("The dataset is empty or not properly formatted.")
+
+    return X, y
+
+
+def create_plot(X, y, model, embedder):
+    # plot data points
+    labels = np.unique(y)
+    colors = label2color(labels)
+    l2c = dict(zip(labels, colors))
+
+    X_embedded = embedder.transform(X)
+
+    for i, label in enumerate(labels):
+        subset = X_embedded[y == label]
+        plt.scatter(subset[:, 0], subset[:, 1], color=colors[i], label=label)
+    plt.legend()
+
+    # plot decision regions
+    xx, yy = np.meshgrid(
+        np.linspace({{ x_min }}, {{ x_max }}, {{ num_dots }}),
+        np.linspace({{ y_min }}, {{ y_max }}, {{ num_dots }}),
+    )
+    xx = xx.ravel()
+    yy = yy.ravel()
+
+    grid = np.c_[xx, yy]
+    predictions = model.predict(embedder.inverse_transform(grid))
+    predictions = predictions.ravel()
+    plt.scatter(xx, yy, c=[l2c[p] for p in predictions], s=1, alpha=0.5)
+
+
+def main():
+    # data loading and preprocessing
+    X, y = load_dataset("{{ dataset_file }}")
+    embedder = {{ embedder_class }}(n_components=2).fit(X)
+
+    # model training
+    model = {{ model_class }}()
+    model.fit(X, y)
+
+    create_plot(X, y, model, embedder)
+    plt.show()
+    # uncomment the line below if you want to save as image
+    # plt.savefig("plot.png")
+
+
+if __name__ == "__main__":
+    main()

model_imports.yaml

@@ -0,0 +1,28 @@
+LinearSVC: from sklearn.svm import LinearSVC
+CalibratedClassifierCV: from sklearn.calibration import CalibratedClassifierCV
+LinearDiscriminantAnalysis: from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
+QuadraticDiscriminantAnalysis: from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
+DummyClassifier: from sklearn.dummy import DummyClassifier
+AdaBoostClassifier: from sklearn.ensemble import AdaBoostClassifier
+BaggingClassifier: from sklearn.ensemble import BaggingClassifier
+ExtraTreesClassifier: from sklearn.ensemble import ExtraTreesClassifier
+GradientBoostingClassifier: from sklearn.ensemble import GradientBoostingClassifier
+HistGradientBoostingClassifier: from sklearn.ensemble import HistGradientBoostingClassifier
+RandomForestClassifier: from sklearn.ensemble import RandomForestClassifier
+GaussianProcessClassifier: from sklearn.gaussian_process import GaussianProcessClassifier
+LogisticRegression: from sklearn.linear_model import LogisticRegression
+PassiveAggressiveClassifier: from sklearn.linear_model import PassiveAggressiveClassifier
+Perceptron: from sklearn.linear_model import Perceptron
+RidgeClassifier: from sklearn.linear_model import RidgeClassifier
+SGDClassifier: from sklearn.linear_model import SGDClassifier
+GaussianNB: from sklearn.naive_bayes import GaussianNB
+KNeighborsClassifier: from sklearn.neighbors import KNeighborsClassifier
+NearestCentroid: from sklearn.neighbors import NearestCentroid
+RadiusNeighborsClassifier: from sklearn.neighbors import RadiusNeighborsClassifier
+MLPClassifier: from sklearn.neural_network import MLPClassifier
+LabelPropagation: from sklearn.semi_supervised import LabelPropagation
+LabelSpreading: from sklearn.semi_supervised import LabelSpreading
+NuSVC: from sklearn.svm import NuSVC
+SVC: from sklearn.svm import SVC
+DecisionTreeClassifier: from sklearn.tree import DecisionTreeClassifier
+ExtraTreeClassifier: from sklearn.tree import ExtraTreeClassifier