webgpu 0.0.1__py3-none-any.whl

webgpu/jupyter.py

@@ -0,0 +1,159 @@
+import base64
+import itertools
+import os
+import pickle
+import time
+
+from . import platform, utils
+from .canvas import Canvas
+from .lilgui import LilGUI
+from .link import js_code as _link_js_code
+from .render_object import *
+from .scene import Scene
+from .triangles import *
+from .utils import init_device_sync
+from .webgpu_api import *
+
+
+def create_package_zip(module_name="webgpu"):
+    """
+    Creates a zip file containing all files in the specified Python package.
+    """
+    import importlib.util
+    import os
+    import tempfile
+    import zipfile
+
+    spec = importlib.util.find_spec(module_name)
+    if spec is None or spec.origin is None:
+        raise ValueError(f"Package {module_name} not found.")
+
+    package_dir = os.path.dirname(spec.origin)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        output_filename = os.path.join(temp_dir, f"{module_name}.zip")
+        with zipfile.ZipFile(output_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
+            for root, _, files in os.walk(package_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, start=os.path.dirname(package_dir))
+                    zipf.write(file_path, arcname)
+
+        return open(output_filename, "rb").read()
+
+
+_id_counter = itertools.count()
+
+
+def _init_html(scene, width, height):
+    from IPython.display import HTML, display
+
+    if isinstance(scene, RenderObject):
+        scene = [scene]
+    if isinstance(scene, list):
+        scene = Scene(scene)
+
+    id_ = f"__webgpu_{next(_id_counter)}_"
+
+    display(
+        HTML(
+            f"""
+            <div id='{id_}root'
+            style="display: flex; justify-content: space-between;"
+            >
+                <canvas 
+                    id='{id_}canvas'
+                    style='background-color: #d0d0d0; flex: 3; width: {width}px; height: {height}px;'
+                >
+                </canvas>
+                <div id='{id_}lilgui'
+                    style='flex: 1;'
+
+                ></div>
+            </div>
+            """
+        )
+    )
+
+    return scene, id_
+
+
+def _draw_scene(scene: Scene, width, height, id_):
+    html_canvas = platform.js.document.getElementById(f"{id_}canvas")
+
+    while html_canvas is None:
+        html_canvas = platform.js.document.getElementById(f"{id_}canvas")
+    html_canvas.width = width
+    html_canvas.height = height
+    gui_element = platform.js.document.getElementById(f"{id_}lilgui")
+
+    canvas = Canvas(utils.get_device(), html_canvas)
+    scene.gui = LilGUI(gui_element, scene)
+    scene.init(canvas)
+    scene.render()
+
+
+def _DrawPyodide(b64_data: str):
+    data = base64.b64decode(b64_data.encode("utf-8"))
+    id_, scene, width, height = pickle.loads(data)
+
+    _draw_scene(scene, width, height, id_)
+    return scene
+
+
+def _DrawHTML(
+    scene: Scene | list[RenderObject] | RenderObject,
+    width=640,
+    height=640,
+):
+    """Draw a scene using display(Javascrip()) with all information in the HTML
+    This way, data is kept in the converted html when running nbconvert
+    The scene object is unpickled and drawn within a pyodide instance in the browser when the html is opened
+    """
+    from IPython.display import Javascript, display
+
+    scene, id_ = _init_html(scene, width, height)
+
+    data = pickle.dumps((id_, scene, width, height))
+    b64_data = base64.b64encode(data).decode("utf-8")
+
+    display(Javascript(f"window.draw_scene('{b64_data}');"))
+    return scene
+
+
+def Draw(
+    scene: Scene | list[RenderObject] | RenderObject,
+    width=640,
+    height=640,
+):
+    scene, id_ = _init_html(scene, width, height)
+    _draw_scene(scene, width, height, id_)
+    return scene
+
+
+if not platform.is_pyodide:
+    from IPython.display import Javascript, display
+
+    is_exporting = "WEBGPU_EXPORTING" in os.environ
+
+    if is_exporting:
+        Draw = _DrawHTML
+        webgpu_module = create_package_zip("webgpu")
+        webgpu_module_b64 = base64.b64encode(webgpu_module).decode("utf-8")
+        js_code = _link_js_code
+        js_code += f"\nconst _webgpu_code = '{webgpu_module_b64}';"
+        js_code += f"\nwindow.pyodide_ready = init_pyodide(_webgpu_code);"
+        display(Javascript(js_code))
+    else:
+        # Not exporting and not running in pyodide -> Start a websocket server and wait for the client to connect
+        while not platform.websocket_server:
+            time.sleep(0.1)
+        port = platform.websocket_server.port
+        host = f"ws://localhost:{port}"
+        js_code = _link_js_code + f"WebsocketLink('{host}');"
+
+        display(Javascript(js_code))
+
+        platform.init()
+
+        device = init_device_sync()

webgpu/jupyter_pyodide.py

@@ -0,0 +1,363 @@
+import base64
+import pickle
+
+from .draw import Draw as DrawPyodide
+from .lilgui import LilGUI
+from .render_object import RenderObject, _render_objects
+from .scene import Scene
+from .utils import _is_pyodide, reload_package
+
+
+def create_package_zip(module_name="webgpu"):
+    """
+    Creates a zip file containing all files in the specified Python package.
+    """
+    import importlib.util
+    import os
+    import tempfile
+    import zipfile
+
+    spec = importlib.util.find_spec(module_name)
+    if spec is None or spec.origin is None:
+        raise ValueError(f"Package {module_name} not found.")
+
+    package_dir = os.path.dirname(spec.origin)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        output_filename = os.path.join(temp_dir, f"{module_name}.zip")
+        with zipfile.ZipFile(output_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
+            for root, _, files in os.walk(package_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, start=os.path.dirname(package_dir))
+                    zipf.write(file_path, arcname)
+
+        return open(output_filename, "rb").read()
+
+
+_package_b64 = base64.b64encode(create_package_zip()).decode("utf-8")
+
+_init_js_code = (
+    r"""
+const SNAPSHOT_URL = 'https://cdn.jsdelivr.net/gh/mhochsteger/ngsolve_pyodide@webgpu1/snapshot.bin.gz';
+
+function decodeB64(base64String) {
+    const binaryString = atob(base64String);
+    const len = binaryString.length;
+    const bytes = new Uint8Array(len);
+    for (let i = 0; i < len; i++) {
+        bytes[i] = binaryString.charCodeAt(i);
+    }
+    return bytes;
+}
+
+async function fetchSnapshot() {
+  const blob = await (await fetch(SNAPSHOT_URL)).blob();
+  const decompressor = new DecompressionStream('gzip');
+  const stream = blob.stream().pipeThrough(decompressor);
+  const response = new Response(stream);
+  return await response.arrayBuffer();
+};
+
+function initLilGUI() {
+    // In generated html files, requirejs is imported before lil-gui is loaded.
+    // Thus, we must load lil-gui using require, use import otherwise.
+    const lil_url = "https://cdn.jsdelivr.net/npm/lil-gui@0.20";
+    if(window.define === undefined){
+        import(lil_url);
+    } else {
+        require([lil_url], (module) => {
+            window.lil = module;
+        });
+    }
+}
+
+async function main() {
+  if(window.webgpu_ready === undefined) {
+      initLilGUI();
+      const pyodide_module = await import("https://cdn.jsdelivr.net/pyodide/v0.27.2/full/pyodide.mjs");
+      window.pyodide = await pyodide_module.loadPyodide( {
+        // _loadSnapshot: await fetchSnapshot(),
+        // lockFileURL: 'https://cdn.jsdelivr.net/gh/mhochsteger/ngsolve_pyodide@webgpu2/pyodide-lock.json',
+        // indexURL: "https://cdn.jsdelivr.net/pyodide/v0.26.2/full/",
+        }
+      );
+      pyodide.setDebug(true);
+      window.pyodide_ready = pyodide.loadPackage(['micropip', 'numpy', 'packaging']);
+      await window.pyodide_ready;
+  }
+  else {
+      await webgpu_ready;
+  }
+  const webgpu_b64 = `"""
+    + _package_b64
+    + r"""`;
+  const webpgu_zip = decodeB64(webgpu_b64);
+  await pyodide.unpackArchive(webpgu_zip, 'zip');
+  await pyodide.runPythonAsync('import webgpu.utils');
+  await pyodide.runPythonAsync('await webgpu.utils.init_device()');
+}
+window.webgpu_ready = main();
+
+"""
+)
+
+
+def _encode_data(data):
+    binary_chunk = pickle.dumps(data)
+    return base64.b64encode(binary_chunk).decode("utf-8")
+
+
+def _decode_data(data):
+    binary_chunk = base64.b64decode(data.encode("utf-8"))
+    return pickle.loads(binary_chunk)
+
+
+def _encode_function(func):
+    import inspect
+
+    return [func.__name__, inspect.getsource(func)]
+
+
+def _decode_function(encoded_func):
+    import __main__
+
+    func_name, func_str = encoded_func
+    symbols = __main__.__dict__
+    exec(func_str, symbols, symbols)
+    return symbols[func_name]
+
+
+def _init(canvas_id="canvas"):
+    import js
+
+    from webgpu.canvas import init_webgpu
+
+    return init_webgpu(js.document.getElementById(canvas_id))
+
+
+def get_render_canvas(canvas_id):
+    return _render_canvases[canvas_id]
+
+
+def _draw_client(canvas_id, scene, assets, globs):
+    from pathlib import Path
+
+    import js
+    import pyodide.ffi
+
+    from webgpu.jupyter import _decode_data, _decode_function
+
+    assets = _decode_data(assets)
+
+    for module_data in assets.get("modules", {}).values():
+        # extract zipfile from binary chunk
+        import io
+        import zipfile
+
+        zipf = zipfile.ZipFile(io.BytesIO(module_data))
+        zipf.extractall()
+
+    for file_name, file_data in assets.get("files", {}).items():
+        with open(file_name, "wb") as f:
+            f.write(file_data)
+
+    for module_name in assets.get("modules", {}):
+        reload_package(module_name)
+
+    canvas = _init(canvas_id)
+    scene = _decode_data(scene)
+
+    if "init_function" in assets:
+        func = _decode_function(assets["init_function"])
+        func(canvas, **scene)
+    elif "init_function_name" in assets:
+        func = globs[assets["init_function_name"]]
+        func(canvas, **scene)
+    else:
+        scene.init(canvas)
+        DrawPyodide(scene, canvas)
+
+
+_draw_js_code_template = r"""
+async function draw() {{
+    var canvas = document.createElement('canvas');
+    var canvas_id = "{canvas_id}";
+    canvas.id = canvas_id;
+    canvas.width = {width};
+    canvas.height = {height};
+    canvas.style = "background-color: #d0d0d0";
+    element.appendChild(canvas);
+    await window.webgpu_ready;
+    await window.pyodide.runPythonAsync('import webgpu.jupyter; webgpu.jupyter._draw_client("{canvas_id}", "{data}", "{assets}", globals())');
+}}
+draw();
+    """
+
+if not _is_pyodide:
+    from IPython.core.magic import register_cell_magic
+    from IPython.display import HTML, Javascript, display
+
+    display(Javascript(_init_js_code))
+
+    _call_counter = 0
+
+    def _get_canvas_id():
+        global _call_counter
+        _call_counter += 1
+        return f"canvas_{_call_counter}"
+
+    def _run_js_code(data, assets, width, height):
+        display(
+            Javascript(
+                _draw_js_code_template.format(
+                    canvas_id=_get_canvas_id(),
+                    data=_encode_data(data),
+                    assets=_encode_data(assets),
+                    width=width,
+                    height=height,
+                )
+            )
+        )
+
+    html_code = r"""
+<div id="{canvas_id}_row" style="display: flex; justify-content: space-between;">
+    <canvas id="{canvas_id}" style="flex: 3; margin-right: 10px; padding: 10px; height: {height}px; width: {width}px; background-color: #d0d0d0;"></canvas>
+    <div id="{canvas_id}_gui" style="flex: 1; margin-left: 10px; padding: 10px;"></div>
+</div>
+"""
+    js_code = r"""
+async function draw() {{
+    await window.webgpu_ready;
+    var gui_element = document.getElementById('{canvas_id}' + '_gui');
+    console.log('gui_element =', gui_element);
+    if(window.lil_guis === undefined) {{
+      window.lil_guis = new Object();
+    }}
+    window.lil_guis['{canvas_id}'] = new lil.GUI({{container: gui_element}});
+    // var canvas2 = document.createElement('canvas');
+    // console.log("canvas2 =", canvas2);
+    var canvas = document.getElementById("{canvas_id}");
+    console.log('canvas size', canvas.clientWidth, canvas.clientHeight);
+    console.log(canvas);
+    canvas.width = Math.floor(canvas.clientWidth/32)*32;
+    canvas.height = Math.floor(canvas.clientHeight/32)*32;
+    canvas.style = "background-color: #d0d0d0; max-width: {width}px; max-height: {height}px;";
+    await window.pyodide.runPythonAsync('import webgpu.jupyter; webgpu.jupyter._draw_client("{canvas_id}", "{scene}", "{assets}", globals())');
+}}
+draw();
+    """
+
+    def Draw(
+        scene: Scene | list[RenderObject] | RenderObject,
+        width=608,
+        height=608,
+        modules=[],
+    ):
+        if isinstance(scene, RenderObject):
+            scene = [scene]
+        if isinstance(scene, list):
+            scene = Scene(scene)
+        canvas_id = _get_canvas_id()
+        scene.gui = LilGUI(canvas_id, scene._id)
+        assets = {"modules": {module: create_package_zip(module) for module in modules}}
+        display(
+            HTML(html_code.format(canvas_id=canvas_id, width=width, height=height)),
+            Javascript(
+                js_code.format(
+                    canvas_id=canvas_id,
+                    scene=_encode_data(scene),
+                    assets=_encode_data(assets),
+                    width=width,
+                    height=height,
+                )
+            ),
+        )
+
+        return scene
+
+    def DrawCustom(
+        client_function,
+        kwargs={},
+        modules: list[str] = [],
+        files: list[str] = [],
+        width=608,
+        height=608,
+    ):
+        assets = {
+            "modules": {module: create_package_zip(module) for module in modules},
+            "files": {f: open(f, "rb").read() for f in files},
+        }
+        if isinstance(client_function, str):
+            assets["init_function_name"] = client_function
+        else:
+            assets["init_function"] = _encode_function(client_function)
+        canvas_id = _get_canvas_id()
+        display(
+            HTML(html_code.format(canvas_id=canvas_id, height=height, width=width)),
+            Javascript(
+                js_code.format(
+                    canvas_id=canvas_id,
+                    scene=_encode_data(kwargs),
+                    assets=_encode_data(assets),
+                    width=width,
+                    height=height,
+                )
+            ),
+        )
+
+    def run_code_in_pyodide(code: str):
+        display(
+            Javascript(
+                f"window.webgpu_ready.then(() => {{ window.pyodide.runPythonAsync(`{code}`) }});"
+            )
+        )
+
+    @register_cell_magic
+    def pyodide(line, cell):
+        run_code_in_pyodide(str(cell))
+
+    @register_cell_magic
+    def pyodide_and_kernel(line, cell):
+        run_code_in_pyodide(str(cell))
+        ip = get_ipython()
+        exec(cell, ip.user_global_ns)
+
+    del pyodide
+
+    class Pyodide:
+        def __setattr__(self, key, value):
+            data = _encode_data(value)
+            display(
+                Javascript(
+                    f"window.webgpu_ready.then(() => {{ window.pyodide.runPythonAsync(`import webgpu.jupyter; {key} = webgpu.jupyter._decode_data('{data}')`) }});"
+                )
+            )
+
+    pyodide = Pyodide()
+
+
+def pyodide_install_packages(packages):
+    if not _is_pyodide:
+        display(
+            Javascript(
+                f"window.webgpu_ready = window.webgpu_ready.then(() => {{ return window.pyodide.loadPackage({packages}); }})"
+            )
+        )
+
+
+def update_render_object(id, **kwargs):
+    if _is_pyodide:
+        obj = _render_objects[id]
+        obj.update(**kwargs)
+    else:
+        kwargs = _encode_data(kwargs)
+        run_code_in_pyodide(
+            f"import webgpu.jupyter; webgpu.jupyter.update_render_object('{id}', **webgpu.jupyter._decode_data('{kwargs}'));"
+        )
+
+
+def redraw_canvas(canvas_id: str):
+    run_code_in_pyodide(
+        f"import webgpu.draw, js; js.requestAnimationFrame(webgpu.draw._canvas_id_to_gpu['{canvas_id}'].input_handler.render_function)"
+    )

webgpu/labels.py

@@ -0,0 +1,132 @@
+import numpy as np
+
+from .font import Font
+from .render_object import RenderObject
+from .uniforms import Binding
+from .utils import BufferBinding, read_shader_file
+from .webgpu_api import *
+
+
+class Labels(RenderObject):
+    vertex_entry_point: str = "vertexText"
+    fragment_entry_point: str = "fragmentFont"
+    n_vertices: int = 6
+
+    """Render a list of strings on screen
+    @param labels: list of strings to render
+    @param positions: list of positions to render the labels at
+    @param apply_camera: whether to apply the camera transformation to the labels
+    @param h_align: horizontal alignment of the labels. Can be one of: left, l, center, c, right, r
+    @param v_align: horizontal alignment of the labels. Can be one of: bottom, b, center, c, top, t
+    @param font_size: font size
+
+    If any of apply_camera, h_align, or v_align is a list, it must have the same length as labels.
+    """
+
+    def __init__(
+        self,
+        labels: list[str],
+        positions: list[tuple],
+        apply_camera: bool | list[bool] = False,
+        h_align: str | list[str] = "left",
+        v_align: str | list[str] = "bottom",
+        font_size=20,
+    ):
+        self.labels = labels
+        self.positions = positions
+        self.font_size = font_size
+        self.apply_camera = apply_camera
+        self.h_align = h_align
+        self.v_align = v_align
+
+    def update(self, timestamp):
+        if timestamp == self._timestamp:
+            return
+        self._timestamp = timestamp
+        n_chars = sum(len(label) for label in self.labels)
+        n_labels = len(self.labels)
+        self.n_vertices = 6
+        self.n_instances = n_chars
+        char_t = np.dtype(
+            [
+                ("itext", np.uint32),
+                ("ichar", np.uint16),
+                ("char", np.uint8),
+                ("padding", np.uint8),
+            ]
+        )
+        char_data = np.zeros(n_chars, dtype=char_t)
+        text_t = np.dtype(
+            [
+                ("pos", np.float32, 3),
+                ("length", np.uint16),
+                ("apply_camera", np.uint8),
+                ("alignment", np.uint8),
+            ]
+        )
+        text_data = np.zeros(n_labels, dtype=text_t)
+
+        align_map = {
+            "c": 1,
+            "center": 1,
+            "r": 2,
+            "right": 2,
+            "t": 2,
+            "top": 2,
+            "b": 0,
+            "bottom": 0,
+            "l": 0,
+            "left": 0,
+        }
+
+        ichar = 0
+        for i, label, pos in zip(range(len(self.labels)), self.labels, self.positions):
+            h_align = self.h_align if isinstance(self.h_align, str) else self.h_align[i]
+            v_align = self.v_align if isinstance(self.v_align, str) else self.v_align[i]
+            align = align_map[h_align] + 4 * align_map[v_align]
+            apply_camera = (
+                self.apply_camera if isinstance(self.apply_camera, bool) else self.apply_camera[i]
+            )
+
+            if len(pos) == 2:
+                pos = (*pos, 0)
+
+            text_data[i]["pos"] = pos
+            text_data[i]["length"] = len(label)
+            text_data[i]["apply_camera"] = apply_camera
+            text_data[i]["alignment"] = align
+
+            i0 = ichar
+            for c in label:
+                char_data[ichar]["itext"] = i
+                char_data[ichar]["ichar"] = ichar - i0
+                char_data[ichar]["char"] = ord(c)
+                ichar += 1
+
+        self.font = Font(self.canvas, self.font_size)
+
+        data = (
+            np.array([len(self.labels)], dtype=np.uint32).tobytes()
+            + text_data.tobytes()
+            + char_data.tobytes()
+        )
+
+        self.buffer = self.device.createBuffer(
+            len(data),
+            usage=BufferUsage.STORAGE | BufferUsage.COPY_DST,
+        )
+        self.device.queue.writeBuffer(self.buffer, 0, data)
+        self.create_render_pipeline()
+
+    def get_shader_code(self):
+        shader_code = read_shader_file("text.wgsl", __file__)
+        shader_code += self.font.get_shader_code()
+        shader_code += self.options.camera.get_shader_code()
+        return shader_code
+
+    def get_bindings(self):
+        return [
+            *self.font.get_bindings(),
+            *self.options.camera.get_bindings(),
+            BufferBinding(Binding.TEXT, self.buffer),
+        ]

webgpu/light.py

@@ -0,0 +1,12 @@
+from .utils import read_shader_file
+
+
+class Light:
+    def __init__(self, device):
+        self.device = device
+
+    def get_bindings(self):
+        return []
+
+    def get_shader_code(self):
+        return read_shader_file("light.wgsl", __file__)