From 8701770d0fb7c17f4ba1cc7ddf93d4d138ee6f52 Mon Sep 17 00:00:00 2001
From: javiber <javier.berneche@gmail.com>
Date: Wed, 3 Apr 2024 10:35:30 -0300
Subject: [PATCH] Added bokeh_wegbl backend

---
 .../implementation/numpy/data/plotter.py      |   9 +-
 .../numpy/data/plotter_bokeh.py               | 107 ++++++++++--------
 2 files changed, 65 insertions(+), 51 deletions(-)

diff --git a/temporian/implementation/numpy/data/plotter.py b/temporian/implementation/numpy/data/plotter.py
index 4ac49d88f..fdcd3776b 100644
--- a/temporian/implementation/numpy/data/plotter.py
+++ b/temporian/implementation/numpy/data/plotter.py
@@ -88,9 +88,16 @@ def bokeh_backend():
     return plotter_bokeh.Plotter
 
 
+def bokeh_webgl_backend():
+    from temporian.implementation.numpy.data import plotter_bokeh
+
+    return plotter_bokeh.PlotterWebGL
+
+
 BACKENDS: Dict[str, Callable] = {
     "matplotlib": matplotlib_backend,
     "bokeh": bokeh_backend,
+    "bokeh_webgl": bokeh_webgl_backend,
 }
 
 
@@ -279,7 +286,7 @@ def plot(
             effectively selects a backend that support interactive plotting.
             Ignored if "backend" is set.
         backend: Plotting library to use. Possible values are: matplotlib,
-            bokeh. If set, overrides the "interactive" argument.
+            bokeh, and bokeh_webgl. If set, overrides the "interactive" argument.
         merge: If true, plots all features in the same plots. If false, plots
             features in separate plots. merge=True on event-sets [e1, e2] is
             equivalent to plotting (e1, e2).
diff --git a/temporian/implementation/numpy/data/plotter_bokeh.py b/temporian/implementation/numpy/data/plotter_bokeh.py
index 1e9ec07c3..4ce297299 100644
--- a/temporian/implementation/numpy/data/plotter_bokeh.py
+++ b/temporian/implementation/numpy/data/plotter_bokeh.py
@@ -21,6 +21,8 @@
 
 
 class Plotter(PlotterBackend):
+    output_backend = "canvas"
+
     def __init__(self, num_plots: int, options: Options):
         super().__init__(num_plots, options)
 
@@ -152,57 +154,62 @@ def finalize(self):
         show(figure_set)
         return figure_set
 
+    def create_fig(
+        self,
+        title: Optional[str],
+        is_unix_timestamp: bool,
+        options: Options,
+        categorical_values: Optional[List[str]],
+    ):
+        tools = [
+            "xpan",
+            "pan",
+            "xwheel_zoom",
+            "ywheel_zoom",
+            "box_zoom",
+            "reset",
+            "undo",
+            "save",
+            "hover",
+        ]
+
+        fig_args = {}
+        if is_unix_timestamp:
+            fig_args["x_axis_type"] = "datetime"
+        if title:
+            fig_args["title"] = title
+
+        if options.min_time is not None or options.max_time is not None:
+            args = {}
+            if options.min_time is not None:
+                args["start"] = (
+                    convert_timestamp_to_datetime(options.min_time)
+                    if is_unix_timestamp
+                    else options.min_time
+                )
+            if options.max_time is not None:
+                args["end"] = (
+                    convert_timestamp_to_datetime(options.max_time)
+                    if is_unix_timestamp
+                    else options.max_time
+                )
+            fig_args["x_range"] = Range1d(**args)
+
+        if categorical_values is not None:
+            fig_args["y_range"] = categorical_values
 
-def create_fig(
-    title: Optional[str],
-    is_unix_timestamp: bool,
-    options: Options,
-    categorical_values: Optional[List[str]],
-):
-    tools = [
-        "xpan",
-        "pan",
-        "xwheel_zoom",
-        "ywheel_zoom",
-        "box_zoom",
-        "reset",
-        "undo",
-        "save",
-        "hover",
-    ]
-
-    fig_args = {}
-    if is_unix_timestamp:
-        fig_args["x_axis_type"] = "datetime"
-    if title:
-        fig_args["title"] = title
-
-    if options.min_time is not None or options.max_time is not None:
-        args = {}
-        if options.min_time is not None:
-            args["start"] = (
-                convert_timestamp_to_datetime(options.min_time)
-                if is_unix_timestamp
-                else options.min_time
-            )
-        if options.max_time is not None:
-            args["end"] = (
-                convert_timestamp_to_datetime(options.max_time)
-                if is_unix_timestamp
-                else options.max_time
-            )
-        fig_args["x_range"] = Range1d(**args)
+        # Note: Ranges cannot be set after the figure is created see:
+        # https://discourse.bokeh.org/t/updating-y-range-to-categorical/2397/3
+        fig = figure(
+            width=options.width_px,
+            height=options.height_per_plot_px,
+            tools=tools,
+            output_backend=self.output_backend,
+            **fig_args,
+        )
 
-    if categorical_values is not None:
-        fig_args["y_range"] = categorical_values
+        return fig
 
-    # Note: Ranges cannot be set after the figure is created see:
-    # https://discourse.bokeh.org/t/updating-y-range-to-categorical/2397/3
-    fig = figure(
-        width=options.width_px,
-        height=options.height_per_plot_px,
-        tools=tools,
-        **fig_args,
-    )
 
-    return fig
+class PlotterWebGL(Plotter):
+    output_backend = "webgl"