webgpu 0.0.1.dev0__py3-none-any.whl

webgpu/__init__.py

@@ -0,0 +1,12 @@
+from .clipping import Clipping
+from .colormap import Colormap
+from .font import Font
+from .render_object import RenderObject
+from .scene import Scene
+from .utils import (
+    BaseBinding,
+    BufferBinding,
+    UniformBinding,
+    create_bind_group,
+    read_shader_file,
+)

webgpu/camera.py

@@ -0,0 +1,191 @@
+import numpy as np
+
+from .uniforms import BaseBinding, Binding, UniformBase, ct
+from .utils import read_shader_file
+
+
+class CameraUniforms(UniformBase):
+    """Uniforms class, derived from ctypes.Structure to ensure correct memory layout"""
+
+    _binding = Binding.CAMERA
+
+    _fields_ = [
+        ("model_view", ct.c_float * 16),
+        ("model_view_projection", ct.c_float * 16),
+        ("normal_mat", ct.c_float * 16),
+        ("aspect", ct.c_float),
+        ("padding", ct.c_uint32 * 3),
+    ]
+
+
+class Transform:
+    def __init__(self):
+        self._mat = np.identity(4)
+        self._rot_mat = np.identity(4)
+        self._center = (0.5, 0.5, 0)
+        self._scale = 1
+
+    def translate(self, dx=0.0, dy=0.0, dz=0.0):
+        translation = np.array(
+            [[1, 0, 0, dx], [0, 1, 0, dy], [0, 0, 1, dz], [0, 0, 0, 1]]
+        )
+        self._mat = translation @ self._mat
+
+    def scale(self, s):
+        self._scale *= s
+
+    def rotate(self, ang_x, ang_y=0):
+        rx = np.radians(ang_x)
+        cx = np.cos(rx)
+        sx = np.sin(rx)
+
+        rotation_x = np.array(
+            [
+                [1, 0, 0, 0],
+                [0, cx, -sx, 0],
+                [0, sx, cx, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+
+        ry = np.radians(ang_y)
+        cy = np.cos(ry)
+        sy = np.sin(ry)
+        rotation_y = np.array(
+            [
+                [cy, 0, sy, 0],
+                [0, 1, 0, 0],
+                [-sy, 0, cy, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+
+        self._rot_mat = rotation_x @ rotation_y @ self._rot_mat
+
+    @property
+    def mat(self):
+        return self._mat @ self._rot_mat @ self._scale_mat @ self._center_mat
+
+    @property
+    def normal_mat(self):
+        return self._mat @ self._rot_mat @ self._scale_mat @ self._center_mat
+
+    @property
+    def _center_mat(self):
+        cx, cy, cz = self._center
+        return np.array([[1, 0, 0, -cx], [0, 1, 0, -cy], [0, 0, 1, -cz], [0, 0, 0, 1]])
+
+    @property
+    def _scale_mat(self):
+        s = self._scale
+        return np.array([[s, 0, 0, 0], [0, s, 0, 0], [0, 0, s, 0], [0, 0, 0, 1]])
+
+
+class Camera:
+    def __init__(self, canvas):
+        self.canvas = canvas
+        self.uniforms = CameraUniforms(canvas.device)
+        self.transform = Transform()
+        self._render_function = None
+        self._is_moving = False
+        self._is_rotating = False
+
+        canvas.on_resize(self._update_uniforms)
+
+    def get_bindings(self) -> list[BaseBinding]:
+        return self.uniforms.get_bindings()
+
+    def get_shader_code(self):
+        return read_shader_file("camera.wgsl", __file__)
+
+    def __del__(self):
+        del self.uniforms
+
+    def register_callbacks(self, input_handler, redraw_function):
+        self._render_function = redraw_function
+        input_handler.on_mousedown(self._on_mousedown)
+        input_handler.on_mouseup(self._on_mouseup)
+        input_handler.on_mouseout(self._on_mouseup)
+        input_handler.on_mousemove(self._on_mousemove)
+        input_handler.on_wheel(self._on_wheel)
+
+    def unregister_callbacks(self, input_handler):
+        input_handler.unregister("mousedown", self._on_mousedown)
+        input_handler.unregister("mouseup", self._on_mouseup)
+        input_handler.unregister("mouseout", self._on_mouseup)
+        input_handler.unregister("mousemove", self._on_mousemove)
+        input_handler.unregister("wheel", self._on_wheel)
+
+    def _on_mousedown(self, ev):
+        if ev["button"] == 0:
+            self._is_rotating = True
+        if ev["button"] == 1:
+            self._is_moving = True
+
+    def _on_mouseup(self, _):
+        self._is_moving = False
+        self._is_rotating = False
+        self._is_zooming = False
+
+    def _on_wheel(self, ev):
+        self.transform.scale(1 - ev["deltaY"] / 1000)
+        self._render()
+        if hasattr(ev, "preventDefault"):
+            ev.preventDefault()
+
+    def _on_mousemove(self, ev):
+        if self._is_rotating:
+            s = 0.3
+            self.transform.rotate(s * ev["movementY"], s * ev["movementX"])
+            self._render()
+        if self._is_moving:
+            s = 0.01
+            self.transform.translate(s * ev["movementX"], -s * ev["movementY"])
+            self._render()
+
+    def _render(self):
+        self._update_uniforms()
+        if self._render_function:
+            self._render_function()
+
+    def _update_uniforms(self):
+        near = 0.1
+        far = 10
+        fov = 45
+        aspect = self.canvas.width / self.canvas.height
+
+        zoom = 1.0
+        top = near * (np.tan(np.radians(fov) / 2)) * zoom
+        height = 2 * top
+        width = aspect * height
+        left = -0.5 * width
+        right = left + width
+        bottom = top - height
+
+        x = 2 * near / (right - left)
+        y = 2 * near / (top - bottom)
+
+        a = (right + left) / (right - left)
+        b = (top + bottom) / (top - bottom)
+
+        c = -far / (far - near)
+        d = (-far * near) / (far - near)
+
+        proj_mat = np.array(
+            [
+                [x, 0, a, 0],
+                [0, y, b, 0],
+                [0, 0, c, d],
+                [0, 0, -1, 0],
+            ]
+        )
+
+        view_mat = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, -3], [0, 0, 0, 1]])
+        model_view = view_mat @ self.transform.mat
+        model_view_proj = proj_mat @ model_view
+        normal_mat = np.linalg.inv(model_view)
+
+        self.uniforms.model_view[:] = model_view.transpose().flatten()
+        self.uniforms.model_view_projection[:] = model_view_proj.transpose().flatten()
+        self.uniforms.normal_mat[:] = normal_mat.flatten()
+        self.uniforms.update_buffer()

webgpu/canvas.py

@@ -0,0 +1,144 @@
+from typing import Callable
+
+from .input_handler import InputHandler
+from .utils import get_device
+from .webgpu_api import *
+from . import platform
+
+
+def init_webgpu(html_canvas):
+    """Initialize WebGPU, create device and canvas"""
+    device = get_device()
+    return Canvas(device, html_canvas)
+
+
+class Canvas:
+    """Canvas management class, handles "global" state, like webgpu device, canvas, frame and depth buffer"""
+
+    device: Device
+    depth_format: TextureFormat
+    depth_texture: Texture
+    multisample_texture: Texture
+    multisample: MultisampleState
+
+    width: int = 0
+    height: int = 0
+
+    _on_resize_callbacks: list[Callable] = []
+
+    def __init__(self, device, canvas, multisample_count=4):
+
+        self._on_resize_callbacks = []
+
+        self.device = device
+        self.format = platform.js.navigator.gpu.getPreferredCanvasFormat()
+        self.color_target = ColorTargetState(
+            format=self.format,
+            blend=BlendState(
+                color=BlendComponent(
+                    srcFactor=BlendFactor.one,
+                    dstFactor=BlendFactor.one_minus_src_alpha,
+                    operation=BlendOperation.add,
+                ),
+                alpha=BlendComponent(
+                    srcFactor=BlendFactor.one,
+                    dstFactor=BlendFactor.one_minus_src_alpha,
+                    operation=BlendOperation.add,
+                ),
+            ),
+        )
+
+        self.canvas = canvas
+
+        self.context = canvas.getContext("webgpu")
+        self.context.configure(
+            toJS(
+                {
+                    "device": device.handle,
+                    "format": self.format,
+                    "alphaMode": "premultiplied",
+                    "sampleCount": multisample_count,
+                    "usage": TextureUsage.RENDER_ATTACHMENT | TextureUsage.COPY_DST,
+                }
+            )
+        )
+
+        self.multisample = MultisampleState(count=multisample_count)
+        self.depth_format = TextureFormat.depth24plus
+        self.input_handler = InputHandler(canvas)
+
+        self.resize()
+
+        # platform.js.webgpuOnResize(canvas, create_proxy(self.resize, True))
+
+    def on_resize(self, func: Callable):
+        self._on_resize_callbacks.append(func)
+
+    def resize(self, *args):
+        canvas = self.canvas
+        rect = canvas.getBoundingClientRect()
+        width = int(rect.width)
+        height = int(rect.height)
+
+        if width == self.width and height == self.height:
+            return False
+
+        if width == 0 or height == 0:
+            return False
+
+        canvas.width = width
+        canvas.height = height
+
+        self.width = width
+        self.height = height
+
+        device = self.device
+        self.target_texture = device.createTexture(
+            size=[width, height, 1],
+            sampleCount=1,
+            format=self.format,
+            usage=TextureUsage.RENDER_ATTACHMENT | TextureUsage.COPY_SRC,
+            label="target",
+        )
+        self.multisample_texture = device.createTexture(
+            size=[width, height, 1],
+            sampleCount=self.multisample.count,
+            format=self.format,
+            usage=TextureUsage.RENDER_ATTACHMENT,
+            label="multisample",
+        )
+
+        self.depth_texture = device.createTexture(
+            size=[width, height, 1],
+            format=self.depth_format,
+            usage=TextureUsage.RENDER_ATTACHMENT,
+            label="depth_texture",
+            sampleCount=self.multisample.count,
+        )
+
+        self.target_texture_view = self.target_texture.createView()
+
+        for func in self._on_resize_callbacks:
+            func()
+
+    def color_attachments(self, loadOp: LoadOp):
+        return [
+            RenderPassColorAttachment(
+                view=self.multisample_texture.createView(),
+                resolveTarget=self.target_texture_view,
+                # view=self.context.getCurrentTexture().createView(),
+                clearValue=Color(1, 1, 1, 1),
+                loadOp=loadOp,
+            )
+        ]
+
+    def depth_stencil_attachment(self, loadOp: LoadOp):
+        return RenderPassDepthStencilAttachment(
+            self.depth_texture.createView(),
+            depthClearValue=1.0,
+            depthLoadOp=loadOp,
+        )
+
+    def __del__(self):
+        # unregister is needed to remove circular references
+        self.input_handler.unregister_callbacks()

webgpu/clipping.py

@@ -0,0 +1,140 @@
+from .render_object import BaseRenderObject
+from .utils import read_shader_file
+
+from .uniforms import UniformBase, ct
+
+import time
+
+
+class Binding:
+    CLIPPING = 1
+
+
+class ClippingUniforms(UniformBase):
+    _binding = Binding.CLIPPING
+    _fields_ = [
+        ("plane", ct.c_float * 4),
+        ("sphere", ct.c_float * 4),
+        ("mode", ct.c_uint32),
+        ("padding", ct.c_uint32 * 3),
+    ]
+
+    def __init__(self, device, mode=0, **kwargs):
+        super().__init__(device, mode=mode, **kwargs)
+
+
+class Clipping(BaseRenderObject):
+    class Mode:
+        DISABLED = 0
+        PLANE = 1
+        SPHERE = 2
+
+    def __init__(
+        self,
+        mode=Mode.DISABLED,
+        center=[0.0, 0.0, 0.0],
+        normal=[0.0, -1.0, 0.0],
+        radius=1.0,
+    ):
+        self.mode = mode
+        self.center = center
+        self.normal = normal
+        self.radius = radius
+        self.callbacks = []
+
+    def update(self, timestamp):
+        if timestamp == self._timestamp:
+            return
+        self._timestamp = timestamp
+        if not hasattr(self, "uniforms"):
+            self.uniforms = ClippingUniforms(self.device)
+        import numpy as np
+
+        c, n = (
+            np.array(self.center, dtype=np.float32),
+            np.array(self.normal, dtype=np.float32),
+        )
+        if np.linalg.norm(n) == 0:
+            n = np.array([0.0, 0.0, -1.0], dtype=np.float32)
+        else:
+            n = n / np.linalg.norm(n)
+        # convert to normal and distance from origin
+        d = -np.dot(c, n)
+        self.uniforms.mode = self.mode
+        for i in range(4):
+            self.uniforms.plane[i] = [*n, d][i]
+            self.uniforms.sphere[i] = [*c, self.radius][i]
+        self.update_buffer()
+
+    def update_buffer(self):
+        self.uniforms.update_buffer()
+
+    def get_bindings(self):
+        return self.uniforms.get_bindings()
+
+    def get_shader_code(self):
+        return read_shader_file("clipping.wgsl", __file__)
+
+    def get_bounding_box(self) -> tuple[list[float], list[float]] | None:
+        return None
+
+    def __del__(self):
+        if hasattr(self, "uniforms"):
+            self.uniforms._buffer.destroy()
+
+    def add_options_to_gui(self, gui):
+        folder = gui.folder("Clipping", closed=True)
+        folder.checkbox(
+            "enabled", self.mode != self.Mode.DISABLED, self.enable_clipping
+        )
+        folder.value("x", self.center[0], self.set_x_value)
+        folder.value("y", self.center[1], self.set_y_value)
+        folder.value("z", self.center[2], self.set_z_value)
+        folder.value("nx", self.normal[0], self.set_nx_value)
+        folder.value("ny", self.normal[1], self.set_ny_value)
+        folder.value("nz", self.normal[2], self.set_nz_value)
+
+    def render(self, encoder):
+        pass
+
+    def enable_clipping(self, value):
+        self.mode = self.Mode.PLANE if value else self.Mode.DISABLED
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()
+
+    def set_x_value(self, value):
+        self.center[0] = value
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()
+
+    def set_y_value(self, value):
+        self.center[1] = value
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()
+
+    def set_z_value(self, value):
+        self.center[2] = value
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()
+
+    def set_nx_value(self, value):
+        self.normal[0] = value
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()
+
+    def set_ny_value(self, value):
+        self.normal[1] = value
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()
+
+    def set_nz_value(self, value):
+        self.normal[2] = value
+        self.update(time.time())
+        for cb in self.callbacks:
+            cb()