fix: faster (threaded) plotting (#207)

* fix: faster (threaded) plotting

Signed-off-by: Henry Schreiner <[email protected]>

* fix: faster (threaded) plotting

Signed-off-by: Henry Schreiner <[email protected]>

* fix: better plotting

Signed-off-by: Henry Schreiner <[email protected]>

---------

Signed-off-by: Henry Schreiner <[email protected]>

Author: henryiii

src/uproot_browser/plot.py

@@ -13,6 +13,8 @@
 import numpy as np
 import plotext as plt
 import uproot
+import uproot.behaviors.TH1
+import uproot.models.RNTuple
 
 from uproot_browser.exceptions import EmptyTreeError
 
@@ -42,7 +44,7 @@ def make_hist_title(item: Any, histogram: hist.Hist) -> str:
 
 
 @functools.singledispatch
-def plot(tree: Any, *, expr: str = "") -> None:  # noqa: ARG001
+def plot(tree: Any, *, width: int = 100, expr: str = "") -> None:  # noqa: ARG001
     """
     Implement this for each type of plottable.
     """
@@ -53,7 +55,10 @@ def plot(tree: Any, *, expr: str = "") -> None:  # noqa: ARG001
 # Simpler in Python 3.11+
 @plot.register(uproot.TBranch)
 def plot_branch(
-    tree: uproot.TBranch | uproot.models.RNTuple.RField, *, expr: str = ""
+    tree: uproot.TBranch | uproot.models.RNTuple.RField,
+    *,
+    width: int = 100,
+    expr: str = "",
 ) -> None:
     """
     Plot a single tree branch.
@@ -64,14 +69,13 @@ def plot_branch(
     if len(finite) < 1:
         msg = f"Branch {tree.name} is empty."
         raise EmptyTreeError(msg)
-    histogram: hist.Hist = hist.numpy.histogram(finite, bins=100, histogram=hist.Hist)
+    histogram: hist.Hist = hist.numpy.histogram(finite, bins=width, histogram=hist.Hist)
     if expr:
         # pylint: disable-next=eval-used
         histogram = eval(expr, {"h": histogram})
     plt.bar(
-        histogram.axes[0].centers,
+        histogram.axes[0].edges,
         histogram.values().astype(float),
-        width=histogram.axes[0].widths,
     )
     plt.ylim(lower=0)
     plt.xticks(np.linspace(histogram.axes[0][0][0], histogram.axes[0][-1][-1], 5))
@@ -83,15 +87,19 @@ def plot_branch(
 
 
 @plot.register
-def plot_hist(tree: uproot.behaviors.TH1.Histogram, expr: str = "") -> None:
+def plot_hist(
+    tree: uproot.behaviors.TH1.Histogram,
+    width: int = 100,  # noqa: ARG001
+    expr: str = "",
+) -> None:
     """
     Plot a 1-D Histogram.
     """
     histogram = hist.Hist(tree.to_hist())
     if expr:
         # pylint: disable-next=eval-used
         histogram = eval(expr, {"h": histogram})
-    plt.bar(histogram.axes[0].centers, histogram.values().astype(float))
+    plt.bar(histogram.axes[0].edges, histogram.values().astype(float))
     plt.ylim(lower=0)
     plt.xticks(np.linspace(histogram.axes[0][0][0], histogram.axes[0][-1][-1], 5))
     plt.xlabel(histogram.axes[0].name)

src/uproot_browser/tui/browser.py

@@ -13,6 +13,7 @@
 import textual.containers
 import textual.events
 import textual.widgets
+import textual.worker
 from textual.reactive import var
 
 with contextlib.suppress(AttributeError):
@@ -36,7 +37,7 @@
 from .header import Header
 from .help import HelpScreen
 from .left_panel import UprootTree
-from .messages import ErrorMessage, UprootSelected
+from .messages import ErrorMessage, RequestPlot, UprootSelected
 from .plot import Plotext
 from .tools import Info, Tools
 from .viewer import ViewWidget
@@ -140,6 +141,16 @@ def on_empty_message(self) -> None:
     def on_error_message(self, message: ErrorMessage) -> None:
         self.view_widget.item = message.err
 
+    def on_request_plot(self, message: RequestPlot) -> None:
+        self.render_plot(message.plot)
+
+    @textual.work(exclusive=True, thread=True)
+    def render_plot(self, plot: Plotext) -> None:
+        worker = textual.worker.get_current_worker()
+        new_plot = plot.make_plot()
+        if new_plot and not worker.is_cancelled:
+            self.call_from_thread(self.view_widget.plot_widget.update, new_plot)
+
 
 if __name__ in {"<run_path>", "__main__"}:
     fname = "../scikit-hep-testdata/src/skhep_testdata/data/uproot-Event.root"

src/uproot_browser/tui/messages.py

@@ -7,6 +7,7 @@
 
 if TYPE_CHECKING:
     from .error import Error
+    from .plot import Plotext
 
 
 @rich.repr.auto
@@ -27,3 +28,10 @@ class ErrorMessage(textual.message.Message, bubble=True):
     def __init__(self, err: Error) -> None:
         self.err = err
         super().__init__()
+
+
+@rich.repr.auto
+class RequestPlot(textual.message.Message, bubble=True):
+    def __init__(self, plot: Plotext) -> None:
+        self.plot = plot
+        super().__init__()

src/uproot_browser/tui/plot.py

@@ -7,13 +7,12 @@
 
 import plotext as plt  # plots in text
 import rich.text
-import textual.widgets
 
 import uproot_browser.plot
 from uproot_browser.exceptions import EmptyTreeError
 
 from .error import Error
-from .messages import EmptyMessage, ErrorMessage
+from .messages import EmptyMessage, ErrorMessage, RequestPlot
 
 if TYPE_CHECKING:
     from .browser import Browser
@@ -32,7 +31,7 @@ def make_plot(item: Any, theme: str, *size: int, expr: str) -> Any:
     plt.clf()
     plt.theme(theme)
     plt.plotsize(*size)
-    uproot_browser.plot.plot(item, expr=expr)
+    uproot_browser.plot.plot(item, width=(size[0] - 5) * 4, expr=expr)
     return plt.build()
 
 
@@ -44,6 +43,24 @@ class Plotext:
     theme: str
     app: Browser
     expr: str = ""
+    size: tuple[int, int] | None = None
+    previous: rich.text.Text | None = None
+    old_expr: str = ""
+
+    def make_plot(self) -> Plotext | None:
+        *_, item = apply_selection(self.upfile, self.selection.split(":"))
+        assert self.size
+        try:
+            canvas = make_plot(item, self.theme, *self.size, expr=self.expr)
+            return dataclasses.replace(self, previous=rich.text.Text.from_ansi(canvas))
+        except EmptyTreeError:
+            self.app.post_message(EmptyMessage())
+            return None
+        except Exception:
+            exc = sys.exc_info()
+            assert exc[1]
+            self.app.post_message(ErrorMessage(Error(exc)))
+            return None
 
     def __rich_console__(
         self, console: rich.console.Console, options: rich.console.ConsoleOptions
@@ -59,13 +76,17 @@ def __rich_console__(
         width = options.max_width or console.width
         height = options.height or console.height
 
-        try:
-            canvas = make_plot(item, self.theme, width, height, expr=self.expr)
-            yield rich.text.Text.from_ansi(canvas)
-        except EmptyTreeError:
-            self.app.post_message(EmptyMessage())
-        except Exception:
-            self.app.query_one("#plot-input", textual.widgets.Input).value = ""
-            exc = sys.exc_info()
-            assert exc[1]
-            self.app.post_message(ErrorMessage(Error(exc)))
+        if (
+            self.size
+            and (width, height) == self.size
+            and self.previous is not None
+            and self.old_expr == self.expr
+        ):
+            yield self.previous
+
+        else:
+            self.size = (width, height)
+            self.previous = rich.text.Text("... plotting ...")
+            self.old_expr = self.expr
+            yield self.previous
+            self.app.post_message(RequestPlot(self))

tests/test_tui.py

@@ -16,6 +16,7 @@ async def test_browse_plot() -> None:
         skhep_testdata.data_path("uproot-Event.root")
     ).run_test() as pilot:
         await pilot.press("down", "down", "down", "enter")
+        await pilot.pause()
         assert isinstance(pilot.app.view_widget.item, Plotext)
 
 
@@ -24,6 +25,7 @@ async def test_browse_empty() -> None:
         skhep_testdata.data_path("uproot-empty.root")
     ).run_test() as pilot:
         await pilot.press("down", "space", "down", "enter")
+        await pilot.pause()
         assert pilot.app.view_widget.item is None
 
 
@@ -32,6 +34,7 @@ async def test_browse_empty_vim() -> None:
         skhep_testdata.data_path("uproot-empty.root")
     ).run_test() as pilot:
         await pilot.press("j", "l", "j", "enter")
+        await pilot.pause()
         assert pilot.app.view_widget.item is None