Add option to use a custom function

gp_visualizer.py

@@ -9,6 +9,7 @@ from jinja2 import Template
 import matplotlib.pyplot as plt
 import matplotlib.lines as mlines
 import numpy as np
+import numexpr
 import pandas as pd
 from PIL import Image
 import plotly.graph_objects as go
@@ -35,6 +36,14 @@ logging.basicConfig(
 )
 logger = logging.getLogger("ELVIS")
 
+
+NUMEXPR_CONSTANTS = {
+    'pi': np.pi,
+    'PI': np.pi,
+    'e': np.e,
+}
+
+
 def eval_kernel(kernel_str):
     # List of allowed kernel constructors
     allowed_names = {
@@ -64,6 +73,25 @@ def eval_kernel(kernel_str):
 
     return result
 
+
+def get_function(function, xlim=(-1, 1), nsample=100):
+    pi = np.pi
+    PI = np.pi
+
+    x = np.linspace(xlim[0], xlim[1], nsample)
+    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS})
+    x = x.reshape(-1, 1)
+    return x, y
+
+
+def get_data_points(function, xlim=(-1, 1), nsample=10, sigma=0):
+    x = xlim[0] + (xlim[1] - xlim[0]) * np.random.rand(nsample)
+    x = np.sort(x)
+    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS}) + sigma * np.random.randn(nsample)
+    x = x.reshape(-1, 1)
+    return x, y
+
+
 def make_sine(xlim=(0,1), nsample=20, sigma=0.1, uniform=False, sort=True):
     np.random.seed(42)
     if uniform:
@@ -79,15 +107,21 @@ def make_sine(xlim=(0,1), nsample=20, sigma=0.1, uniform=False, sort=True):
 
     return X, y
 
+
 class GPVisualizer:
     DEFAULT_KERNEL = "RBF() + WhiteKernel()"
+    DEFAULT_FUNCTION = "sin(2 * pi * x)"
 
     def __init__(self, width, height):
         self.canvas_width = width
         self.canvas_height = height
 
+        self.function = self.DEFAULT_FUNCTION
         self.kernel = eval_kernel(self.DEFAULT_KERNEL)
 
+        self.x_train, self.y_train = self.generate_data(self.function)
+        self.model = self.train_model(self.kernel, self.x_train, self.y_train)
+
         self.plot_options = {
             "show_training_data": True,
             "show_confidence_interval": True,
@@ -102,6 +136,15 @@ class GPVisualizer:
     display: none;
 }"""
 
+    def generate_data(self, function):
+        return get_data_points(function, xlim=(-1, 1), nsample=30, sigma=0.1)
+
+    def train_model(self, kernel, x_train, y_train):
+        gpr = GaussianProcessRegressor(kernel=kernel, random_state=0)
+        logger.info('fitting ' + str(gpr))
+        gpr.fit(x_train, y_train)
+        return gpr
+
     def plot(self):
         '''
         '''
@@ -113,17 +156,8 @@ class GPVisualizer:
         ax = fig.add_axes([0., 0., 1., 1.]) # 
         ax.margins(x=0, y=0) # no padding in both directions
 
-        # make a synthetic dataset 
-        logger.info("Generating data")
-        X_tr, y_tr = make_sine(xlim=(0, 1), nsample=20)
-        print(X_tr.shape, y_tr.shape)
-        X_ts, y_ts = make_sine(xlim=(-1, 2), uniform=True, sort=True, nsample=100)
-
-        # fit GP
-        gpr = GaussianProcessRegressor(kernel=self.kernel, random_state=0)
-        logger.info('fitting ' + str(gpr))
-        gpr.fit(X_tr, y_tr)
-        y_pred, y_std = gpr.predict(X_ts, return_std=True)
+        x_test, y_test = get_function(self.function, xlim=(-2, 2), nsample=100)
+        y_pred, y_std = self.model.predict(x_test, return_std=True)
 
         # plot
         fig, ax = plt.subplots(figsize=(8, 8))
@@ -131,28 +165,25 @@ class GPVisualizer:
         ax.set_xlabel("x")
         ax.set_ylabel("y")
 
-        #plt.scatter(X_tr.flatten(), y_tr)
-        #plt.errorbar(X_ts.flatten(), y_pred, yerr=1.95*y_std, fmt='+', color='orange')
-
-        R2 = gpr.score(X_tr, y_tr)
+        R2 = self.model.score(self.x_train, self.y_train)
 
         if self.plot_options["show_training_data"]:
-            if len(X_tr) > 1:
-                plt.scatter(X_tr.flatten(), y_tr, label='training data (R2=%.2f)' % (R2), color=self.plot_cmap(0))
+            if len(self.x_train) > 1:
+                plt.scatter(self.x_train.flatten(), self.y_train, label='training data (R2=%.2f)' % (R2), color=self.plot_cmap(0))
             else:
-                plt.scatter(X_tr.flatten(), y_tr, label='training data', color=self.plot_cmap(0))
+                plt.scatter(self.x_train.flatten(), self.y_train, label='training data', color=self.plot_cmap(0))
 
         if self.plot_options["show_true_function"]:
-            plt.plot(X_ts.flatten(), np.sin(2*np.pi*X_ts.flatten()), label='true function', color=self.plot_cmap(1))
+            plt.plot(x_test.flatten(), y_test, label='true function', color=self.plot_cmap(1))
 
         if self.plot_options["show_predictions"]:
-            plt.scatter(X_ts.flatten(), y_pred, marker='+', label='predictions', color=self.plot_cmap(2))
+            plt.plot(x_test.flatten(), y_pred, linestyle="--", label='mean prediction', color=self.plot_cmap(2))
 
         if self.plot_options["show_confidence_interval"]:
             plt.fill_between(
-                X_ts.flatten(), 
-                y_pred - 1.96*y_std, 
-                y_pred + 1.96*y_std, 
+                x_test.flatten(), 
+                y_pred - 1.96 * y_std, 
+                y_pred + 1.96 * y_std, 
                 alpha=0.2,
                 label='95% confidence interval', 
                 color=self.plot_cmap(3)
@@ -201,10 +232,21 @@ class GPVisualizer:
         )
         return fig
 
-    #def update_kernel_type(self, kernel_type):
-        #self.kernel_type = kernel_type
-        #self.kernel = self.kernels[kernel_type]
-        #return self.plot()
+    def update_function(self, function):
+        # test if function is valid
+        try:
+            x = np.linspace(-1, 1, 10)
+            y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS})
+        except Exception as e:
+            raise ValueError(f"Invalid function: {e}")
+
+        self.function = function
+
+        # reset data and model
+        self.x_train, self.y_train = self.generate_data(self.function)
+        self.model = self.train_model(self.kernel, self.x_train, self.y_train)
+
+        return self.plot()
 
     def update_kernel_spec(self, kernel_spec):
         self.kernel = eval_kernel(kernel_spec)
@@ -231,8 +273,18 @@ class GPVisualizer:
 
                 with gr.Column(scale=1): 
                     with gr.Tab("Settings"):
-                        dataset_radio = gr.Radio(["sine", "Upload"], 
-                                                value="sine", label="Dataset", elem_id="rowheight")
+                        dataset_radio = gr.Radio(
+                            ["Generate", "Upload"], 
+                            value="Generate", 
+                            label="Dataset", 
+                        )
+
+                        function_box = gr.Textbox(
+                            label="Function",
+                            placeholder="function of x",
+                            value=self.DEFAULT_FUNCTION,
+                            interactive=True,
+                        )
 
                         # upload data
                         file_chooser = gr.File(label="Choose a file", visible=False, elem_id="rowheight")
@@ -272,14 +324,12 @@ class GPVisualizer:
                     with gr.Tab("Usage"):
                         gr.Markdown(''.join(open('usage.md', 'r').readlines()))
 
+                    function_box.submit(
+                        fn=self.update_function,
+                        inputs=function_box,
+                        outputs=[self.canvas],
+                    )
             
-                    # event handlers for GUI elements
-                    #kernel_type.change(
-                        #fn=self.update_kernel_type, 
-                        #inputs=kernel_type,
-                        #outputs=(self.canvas)
-                    #)
-
                     kernel_spec.submit(
                         fn=self.update_kernel_spec,
                         inputs=kernel_spec,