Jupyter canvas

README.md

@@ -80,6 +80,9 @@ from wgpu.gui.glfw import WgpuCanvas
 # Import PySide2, PyQt5, PySide or PyQt4 before running the line below.
 # The code will detect and use the library that is imported.
 from wgpu.gui.qt import WgpuCanvas
+
+# You can also show wgpu visualizations in Jupyter
+from wgpu.gui.jupyter import WgpuCanvas
 ```
 
 Some functions in the original `wgpu-native` API are async. In the Python API,

docs/reference_gui.rst

@@ -6,7 +6,7 @@ screen is also possible, but we need a *canvas* for that. Since the Python
 ecosystem provides many different GUI toolkits, we need an interface.
 
 For convenience, the wgpu library has builtin support for a few GUI
-toolkits. At the moment these include GLFW and Qt.
+toolkits. At the moment these include GLFW, Qt, and wx (experimental).
 
 
 The canvas interface
@@ -15,34 +15,23 @@ The canvas interface
 To render to a window, an object is needed that implements the few
 functions on the canvas interface, and provide that object to
 :func:`request_adapter() <wgpu.request_adapter>`.
-This interface makes it possible to hook wgpu-py to any GUI that supports GPU rendering.
+This is the minimal interface required to hook wgpu-py to any GUI that supports GPU rendering.
 
 .. autoclass:: wgpu.gui.WgpuCanvasInterface
     :members:
 
 
-The WgpuCanvas classes
-----------------------
+The WgpuCanvas base class
+-------------------------
 
-For each GUI toolkit that wgpu-py has builtin support for, there is a
-``WgpuCanvas`` class, which all derive from the following class. This thus
-provides a single (simple) API to work with windows.
+For each supported GUI toolkit there are specific
+``WgpuCanvas`` classes, which are detailed in the following sections.
+These all derive from the same base class, which defines the common API.
 
 .. autoclass:: wgpu.gui.WgpuCanvasBase
     :members:
 
 
-Offscreen canvases
-------------------
-
-A base class is provided to implement off-screen canvases. Note that you can
-render to a texture without using any canvas object, but in some cases it's
-convenient to do so with a canvas-like API, which is what this class provides.
-
-.. autoclass:: wgpu.gui.WgpuOffscreenCanvas
-    :members:
-
-
 Support for Qt
 --------------
 
@@ -102,3 +91,36 @@ Glfw is a lightweight windowing toolkit. Install it with ``pip install glfw``.
 
 Also see the `GLFW triangle example <https://github.com/pygfx/wgpu-py/blob/main/examples/triangle_glfw.py>`_
 and the `async GLFW example <https://github.com/pygfx/wgpu-py/blob/main/examples/triangle_glfw_asyncio.py>`_.
+
+
+Offscreen canvases
+------------------
+
+A base class is provided to implement off-screen canvases. Note that you can
+render to a texture without using any canvas object, but in some cases it's
+convenient to do so with a canvas-like API.
+
+.. autoclass:: wgpu.gui.WgpuOffscreenCanvas
+    :members:
+
+
+Support for Jupyter lab and notebook
+------------------------------------
+
+WGPU can be used in Jupyter lab and the Jupyter notebook. This canvas
+is based on `jupyter_rfb <https://github.com/vispy/jupyter_rfb>`_ an ipywidget
+subclass implementing a remote frame-buffer. There are also some `wgpu examples <https://jupyter-rfb.readthedocs.io/en/latest/examples/>`_.
+
+To implement interaction, create a subclass and overload the ``handle_event()``
+method (and call ``super().handle_event(event)``).
+
+
+.. code-block:: py
+
+    from wgpu.gui.jupyter import WgpuCanvas
+
+    canvas = WgpuCanvas()
+
+    # ... wgpu code
+
+    canvas  # Use as cell output

examples/triangle.py

@@ -30,7 +30,7 @@
 
 [[stage(vertex)]]
 fn vs_main(in: VertexInput) -> VertexOutput {
-    var positions: array<vec2<f32>, 3> = array<vec2<f32>, 3>(vec2<f32>(0.0, -0.5), vec2<f32>(0.5, 0.5), vec2<f32>(-0.5, 0.7));
+    var positions = array<vec2<f32>, 3>(vec2<f32>(0.0, -0.5), vec2<f32>(0.5, 0.5), vec2<f32>(-0.5, 0.7));
     let index = i32(in.vertex_index);
     let p: vec2<f32> = positions[index];
 

wgpu/gui/base.py

@@ -16,7 +16,7 @@ class WgpuCanvasInterface:
     def __init__(self, *args, **kwargs):
         # The args/kwargs are there because we may be mixed with e.g. a Qt widget
         super().__init__(*args, **kwargs)
-        self._present_context = None
+        self._canvas_context = None
 
     def get_window_id(self):
         """Get the native window id. This is used to obtain a surface id,
@@ -59,13 +59,13 @@ def get_context(self, kind="gpupresent"):
         # Note that this function is analog to HtmlCanvas.get_context(), except
         # here the only valid arg is 'gpupresent', which is also made the default.
         assert kind == "gpupresent"
-        if self._present_context is None:
+        if self._canvas_context is None:
             # Get the active wgpu backend module
             backend_module = sys.modules["wgpu"].GPU.__module__
             # Instantiate the context
             PC = sys.modules[backend_module].GPUCanvasContext  # noqa: N806
-            self._present_context = PC(self)
-        return self._present_context
+            self._canvas_context = PC(self)
+        return self._canvas_context
 
 
 class WgpuCanvasBase(WgpuCanvasInterface):
@@ -102,13 +102,14 @@ def _draw_frame_and_present(self):
         """
         # Perform the user-defined drawing code. When this errors,
         # we should report the error and then continue, otherwise we crash.
+        # Returns the result of the context's present() call or None.
         try:
             self.draw_frame()
         except Exception as err:
             self._log_exception("Draw error", err)
         try:
-            if self._present_context:
-                self._present_context.present()
+            if self._canvas_context:
+                return self._canvas_context.present()
         except Exception as err:
             self._log_exception("Present error", err)
 

wgpu/gui/jupyter.py

@@ -0,0 +1,96 @@
+"""
+Support for rendering in a Jupyter widget. Provides a widget subclass that
+can be used as cell output, or embedded in a ipywidgets gui.
+"""
+
+from .offscreen import WgpuOffscreenCanvas
+
+import numpy as np
+from jupyter_rfb import RemoteFrameBuffer
+
+
+class JupyterWgpuCanvas(WgpuOffscreenCanvas, RemoteFrameBuffer):
+    """An ipywidgets widget providing a wgpu canvas. Needs the jupyter_rfb library."""
+
+    def __init__(self):
+        super().__init__()
+        self._pixel_ratio = 1
+        self._logical_size = 0, 0
+        self._is_closed = False
+
+    # Implementation needed for RemoteFrameBuffer
+
+    def handle_event(self, event):
+        event_type = event.get("event_type", "")
+        if event_type == "close":
+            self._is_closed = True
+        elif event_type == "resize":
+            self._pixel_ratio = event["pixel_ratio"]
+            self._logical_size = event["width"], event["height"]
+
+    def get_frame(self):
+        # The _draw_frame_and_present() does the drawing and then calls
+        # present_context.present(), which calls our present() method.
+        # The resuls is either a numpy array or None, and this matches
+        # with what this method is expected to return.
+        return self._draw_frame_and_present()
+
+    # Implementation needed for WgpuCanvasBase
+
+    def get_pixel_ratio(self):
+        return self._pixel_ratio
+
+    def get_logical_size(self):
+        return self._logical_size
+
+    def get_physical_size(self):
+        return int(self._logical_size[0] * self._pixel_ratio), int(
+            self._logical_size[1] * self._pixel_ratio
+        )
+
+    def set_logical_size(self, width, height):
+        self.css_width = f"{width}px"
+        self.css_height = f"{height}px"
+
+    def close(self):
+        RemoteFrameBuffer.close(self)
+
+    def is_closed(self):
+        return self._is_closed
+
+    def _request_draw(self):
+        RemoteFrameBuffer.request_draw(self)
+
+    # Implementation needed for WgpuOffscreenCanvas
+
+    def present(self, texture_view):
+        # This gets called at the end of a draw pass via GPUCanvasContextOffline
+        device = texture_view._device
+        size = texture_view.size
+        bytes_per_pixel = 4
+        data = device.queue.read_texture(
+            {
+                "texture": texture_view.texture,
+                "mip_level": 0,
+                "origin": (0, 0, 0),
+            },
+            {
+                "offset": 0,
+                "bytes_per_row": bytes_per_pixel * size[0],
+                "rows_per_image": size[1],
+            },
+            size,
+        )
+        return np.frombuffer(data, np.uint8).reshape(size[1], size[0], 4)
+
+    def get_preferred_format(self):
+        # Use a format that maps well to PNG: rgba8norm.
+        # Use srgb for perseptive color mapping. You probably want to
+        # apply this before displaying on screen, but you do not want
+        # to duplicate it. When a PNG is shown on screen in the browser
+        # it's shown as-is (at least it was when I just tried).
+        return "rgba8unorm-srgb"
+
+
+# Make available under a name that is the same for all gui backends
+WgpuCanvas = JupyterWgpuCanvas

wgpu/gui/offscreen.py

@@ -14,17 +14,29 @@ def get_window_id(self):
 
     def get_context(self, kind="gpupresent"):
         """Get the GPUCanvasContext object to obtain a texture to render to."""
+        # Normally this creates a GPUCanvasContext object provided by
+        # the backend (e.g. rs), but here we use our own context.
         assert kind == "gpupresent"
-        if self._present_context is None:
-            self._present_context = GPUCanvasContextOffline(self)
-        return self._present_context
+        if self._canvas_context is None:
+            self._canvas_context = GPUCanvasContextOffline(self)
+        return self._canvas_context
 
     def present(self, texture_view):
         """Method that gets called at the end of each draw event. Subclasses
         should provide the approproate implementation.
         """
         pass
 
+    def get_preferred_format(self):
+        """Get the preferred format for this canvas. This method can
+        be overloaded to control the used texture format. The default
+        is "rgba8unorm" (not including srgb colormapping).
+        """
+        # Use rgba because that order is more common for processing and storage.
+        # Use 8unorm because 8bit is enough and common in most cases.
+        # We DO NOT use srgb colormapping here; we return the "raw" output.
+        return "rgba8unorm"
+
 
 class GPUCanvasContextOffline(base.GPUCanvasContext):
     """Helper class for canvases that render to a texture."""
@@ -41,7 +53,11 @@ def unconfigure(self):
         self._texture = None
 
     def get_preferred_format(self, adapter):
-        return "rgba8unorm"
+        canvas = self._get_canvas()
+        if canvas:
+            return canvas.get_preferred_format()
+        else:
+            return "rgba8unorm"
 
     def get_current_texture(self):
         self._create_new_texture_if_needed()
@@ -51,7 +67,7 @@ def get_current_texture(self):
     def present(self):
         if self._texture_view is not None:
             canvas = self._get_canvas()
-            canvas.present(self._texture_view)
+            return canvas.present(self._texture_view)
 
     def _create_new_texture_if_needed(self):
         canvas = self._get_canvas()