webgpu 0.0.1__py3-none-any.whl

webgpu/shaders/clipping.wgsl

@@ -0,0 +1,35 @@
+// general uniform structures
+struct ClippingUniforms {
+  plane: vec4<f32>,
+  sphere: vec4<f32>,
+  mode: u32, // 0: disabled, 1: plane, 2: sphere, 3: both
+
+  padding0: u32,
+  padding1: u32,
+  padding2: u32,
+};
+
+@group(0) @binding(1) var<uniform> u_clipping : ClippingUniforms;
+
+fn calcClipping(p: vec3<f32>) -> bool {
+    var result: bool = true;
+    if( (u_clipping.mode & 0x01u) == 1u) {
+        if dot(u_clipping.plane, vec4<f32>(p, 1.0)) > 0. {
+            result = false;
+        }
+    }
+    if ((u_clipping.mode & 0x02u) == 0x02u){
+        let d = distance(p, u_clipping.sphere.xyz);
+        if d > u_clipping.sphere.w {
+            result = false;
+        }
+    }
+    return result;
+}
+
+fn checkClipping(p: vec3<f32>) {
+    if calcClipping(p) == false {
+    discard;
+    }
+}
+

webgpu/shaders/colormap.wgsl

@@ -0,0 +1,60 @@
+@group(0) @binding(6) var u_colormap_texture : texture_1d<f32>;
+@group(0) @binding(7) var u_colormap_sampler : sampler;
+@group(0) @binding(5) var<uniform> u_cmap_uniforms : ColormapUniforms;
+
+struct ColormapUniforms {
+  min: f32,
+  max: f32,
+  position_x: f32,
+  position_y: f32,
+  discrete: u32,
+  n_colors: u32,
+  width: f32,
+  height: f32,
+};
+
+
+fn getColor(value: f32) -> vec4<f32> {
+  var v = (value - u_cmap_uniforms.min) / (u_cmap_uniforms.max - u_cmap_uniforms.min);
+  if(u_cmap_uniforms.discrete == 1u) {
+      v = f32(u32(f32(u_cmap_uniforms.n_colors) * v));
+      v = v / f32(u_cmap_uniforms.n_colors);
+    }
+  return textureSample(u_colormap_texture, u_colormap_sampler, v);
+}
+
+struct VertexOutput {
+  @builtin(position) position: vec4<f32>,
+  @location(0) val: f32,
+};
+
+@vertex
+fn colormap_vertex(@builtin(vertex_index) vertId: u32,
+               @builtin(instance_index) trigId: u32) -> VertexOutput {
+  var posx = trigId / 2u;
+  if(vertId == 2u || (trigId % 2u == 1u && vertId == 1u))
+    {
+      posx = posx + 1u;
+    }
+  var posy = u_cmap_uniforms.position_y;
+  if(vertId == 0u || (trigId % 2u == 1u && vertId == 1u))
+    {
+      posy = posy + u_cmap_uniforms.height;
+    }
+  let position = vec4<f32>(u_cmap_uniforms.position_x +
+                           f32(posx) * u_cmap_uniforms.width / f32(u_cmap_uniforms.n_colors),
+                           posy, 0.0, 1.0);
+  return VertexOutput(position, f32(posx));
+}
+
+@fragment
+fn colormap_fragment(vert: VertexOutput) -> @location(0) vec4<f32> {
+  if (u_cmap_uniforms.height == 0.) {
+    discard;
+  }
+  let min = u_cmap_uniforms.min;
+  let max = u_cmap_uniforms.max;
+  let v = min + (max-min) * vert.val / f32(u_cmap_uniforms.n_colors);
+  let color = getColor(v);
+  return color;
+}

webgpu/shaders/font.wgsl

@@ -0,0 +1,53 @@
+struct FontUniforms {
+  width: u32,
+  height: u32,
+
+  width_normalized: f32,
+  height_normalized: f32,
+};
+
+@group(0) @binding(2) var<uniform> u_font : FontUniforms;
+@group(0) @binding(3) var u_font_texture : texture_2d<f32>;
+
+fn fontGetTexCoord(char: u32, vertexId: u32) -> vec2<f32> {
+    return vec2<f32>(
+        f32((char - 32) * u_font.width),
+        f32(u_font.height)
+    );
+}
+
+struct FontFragmentInput {
+    @builtin(position) fragPosition: vec4<f32>,
+    @location(0) tex_coord: vec2<f32>,
+};
+
+@fragment
+fn fragmentFont(@location(0) tex_coord: vec2<f32>) -> @location(0) vec4<f32> {
+    let alpha: f32 = textureLoad(
+        u_font_texture,
+        vec2i(floor(tex_coord)),
+        0
+    ).x;
+
+    if alpha < 0.01 {
+      discard;
+    }
+
+    return vec4(0., 0., 0., alpha);
+}
+
+fn fontCalc(char: u32, position: vec4<f32>, vertexId: u32) -> FontFragmentInput {
+    var tex_coord = fontGetTexCoord(char, vertexId);
+    var p = position;
+
+    if vertexId == 2 || vertexId == 4 || vertexId == 5 {
+        p.y += u_font.height_normalized * p.w;
+        tex_coord.y = 0.0;
+    }
+
+    if vertexId == 1 || vertexId == 2 || vertexId == 4 {
+        p.x += u_font.width_normalized * p.w;
+        tex_coord.x += f32(u_font.width);
+    }
+  return FontFragmentInput(p, tex_coord);
+}

webgpu/shaders/light.wgsl

@@ -0,0 +1,9 @@
+fn lightCalcBrightness(n: vec3f) -> f32 {
+    let n4 = cameraMapNormal(n);
+    return clamp(dot(normalize(n4.xyz), normalize(vec3<f32>(1., 3., 3.))), .0, 1.) * 0.7 + 0.3;
+}
+
+fn lightCalcColor(n: vec3f, color: vec4f) -> vec4f {
+    let brightness = lightCalcBrightness(n);
+    return vec4f(color.xyz * brightness, color.w);
+}

webgpu/shaders/text.wgsl

@@ -0,0 +1,57 @@
+@group(0) @binding(30) var<storage> u_text : Texts;
+
+struct Texts {
+  n_texts: u32,
+  data: array<u32>, // check in python code of label.py on how the data is stored
+};
+
+struct TextData {
+  pos: vec3f,
+  shift: vec2f,
+  length: u32,
+  ichar: u32,
+  char: u32,
+  apply_camera: u32,
+};
+
+fn textLoadData(i: u32) -> TextData {
+    let offset = u_text.n_texts * 4 + i * 2;
+    let itext = u_text.data[ offset ];
+    let char_data = u_text.data[ offset + 1 ];
+    let ichar = extractBits(char_data, 0u, 16u);
+    let char = extractBits(char_data, 16u, 8u);
+
+    let offset_text = itext * 4;
+    let pos = vec3f(bitcast<f32>(u_text.data[offset_text]), bitcast<f32>(u_text.data[offset_text + 1]), bitcast<f32>(u_text.data[offset_text + 2]));
+    let text_data = u_text.data[offset_text + 3];
+    let length = extractBits(text_data, 0u, 16u);
+    let apply_camera = extractBits(text_data, 16u, 8u);
+
+    let x_align = f32(extractBits(text_data, 24u, 2u));
+    let y_align = f32(extractBits(text_data, 26u, 2u));
+
+    let shift = vec2<f32>(-0.5 * x_align, -0.5 * y_align);
+
+  return TextData(pos, shift, length, ichar, char, apply_camera);
+}
+
+@vertex
+fn vertexText(@builtin(vertex_index) vertexId: u32, @builtin(instance_index) charId: u32) -> FontFragmentInput {
+    let text = textLoadData(charId);
+
+    var position = vec4f(text.pos, 1.0);
+    if (text.apply_camera != 0) {
+        position = cameraMapPoint(text.pos);
+    }
+
+    let w: f32 = u_font.width_normalized * position.w;
+    let h: f32 = u_font.height_normalized * position.w;
+
+    position.x += f32(text.ichar) * w;
+
+    let shift = text.shift;
+    position.x += w * shift.x * f32(text.length);
+    position.y += h * shift.y;
+
+    return fontCalc(text.char, position, vertexId);
+}

webgpu/shaders/triangulation.wgsl

@@ -0,0 +1,34 @@
+
+@group(0) @binding(90) var<storage> u_vertices : array<f32>;
+@group(0) @binding(91) var<storage> u_normals : array<f32>;
+
+struct TriangleFragmentInput {
+  @builtin(position) position: vec4<f32>,
+    @location(0) p: vec3<f32>,
+    @location(1) n: vec3<f32>,
+    @location(2) @interpolate(flat) vertId: u32,
+    @location(3) @interpolate(flat) trigId: u32,
+};
+
+@vertex
+fn vertex_main(@builtin(vertex_index) vertId: u32, @builtin(instance_index) trigId: u32) -> TriangleFragmentInput {
+  let point = vec3<f32>(u_vertices[trigId * 9 + vertId * 3],
+                        u_vertices[trigId * 9 + vertId * 3 + 1],
+                        u_vertices[trigId * 9 + vertId * 3 + 2]);
+  let normal = -vec3<f32>(u_normals[trigId * 9 + vertId * 3],
+                         u_normals[trigId * 9 + vertId * 3 + 1],
+                         u_normals[trigId * 9 + vertId * 3 + 2]);
+
+  let position = cameraMapPoint(point);
+  return TriangleFragmentInput(position,
+                               point,
+                               normal,
+                               vertId,
+                               trigId,
+                               );
+}
+
+@fragment
+fn fragment_main(input: TriangleFragmentInput) -> @location(0) vec4<f32> {
+  return lightCalcColor(input.n, getColor(input.vertId, input.trigId));
+}

webgpu/shaders/vector.wgsl

@@ -0,0 +1,118 @@
+
+@group(0) @binding(81) var<storage, read> vec_points: array<f32>;
+@group(0) @binding(82) var<storage, read> vec_vectors: array<f32>;
+@group(0) @binding(83) var<uniform> vec_options: VectorOptions;
+
+struct VectorOptions {
+ length: f32,
+ scale_with_vector_length: u32,
+ padding2: f32,
+ padding3: f32
+};
+
+
+fn Cross(a: vec3f, b: vec3f) -> vec3f {
+
+    return vec3f(
+        a[1] * b[2] - a[2] * b[1],
+        a[2] * b[0] - a[0] * b[2],
+        a[0] * b[1] - a[1] * b[0]
+    );
+}
+
+
+struct VectorFragmentInput {
+  @builtin(position) fragPosition: vec4<f32>,
+  @location(0) color_val: f32,
+  @location(1) n: vec3<f32>
+};
+
+fn get_point(index: u32) -> vec3f {
+    let i = index * 3u;
+    return vec3f(vec_points[i], vec_points[i + 1u], vec_points[i + 2u]);
+}
+
+fn get_vector(index: u32) -> vec3f {
+    let i = index * 3u;
+    return vec3f(vec_vectors[i], vec_vectors[i + 1u], vec_vectors[i + 2u]);
+}
+
+// Draw vector cone with triangle strip
+
+const cone_points =
+  array<vec3f, 4>(vec3f(-sqrt(0.5), -sqrt(0.5), 0.),
+    vec3f(sqrt(0.5), -sqrt(0.5), 0.),
+    vec3f(0., 1., 0.),
+    vec3f(0., 0., 5.));
+
+const cone_strip = array<u32, 10>(0u, 1u, 2u, 2u, 3u, 1u, 3u, 0u, 3u, 2u);
+const cone_strip_normals = array<u32, 10>(0u, 0u, 0u, 1u, 4u, 2u, 5u, 3u, 6u, 1u);
+
+@vertex
+fn vertex_main(@builtin(vertex_index) index: u32,
+    @builtin(instance_index) instance: u32) -> VectorFragmentInput {
+    let point = get_point(instance);
+    let vector = get_vector(instance);
+    // if length(vector) < 0.0001 {
+    //     return VectorFragmentInput(vec4<f32>(0., 0., 0., 0.),
+    //         0., vec3<f32>(0., 0., 0.));
+    // }
+    let cp = cone_points[cone_strip[index]];
+    let v = normalize(vector);
+    let z_axis = vec3<f32>(0., 0., 1.);
+    var rotation_matrix = mat3x3<f32>(1., 0., 0.,
+        0., 1., 0.,
+        0., 0., 1.);
+    var rotation_axis = -Cross(z_axis, v);
+    let axis_norm = length(rotation_axis);
+    if axis_norm > 0.0001 {
+        rotation_axis = rotation_axis / axis_norm;
+    
+    // Compute the angle between the z-axis and the target vector
+        let angle = acos(clamp(dot(z_axis, v), -1.0, 1.0));
+    
+    // Compute Rodrigues' rotation formula components
+        let K = mat3x3<f32>(0.0, -rotation_axis.z, rotation_axis.y,
+            rotation_axis.z, 0.0, -rotation_axis.x,
+            -rotation_axis.y, rotation_axis.x, 0.0);
+
+        // Final rotation matrix
+        rotation_matrix += sin(angle) * K + (1.0 - cos(angle)) * (K * K);
+    } else {
+        if dot(z_axis, v) < 0. {
+            rotation_matrix = mat3x3<f32>(-1., 0., 0.,
+                0., -1., 0.,
+                0., 0., -1.);
+        }
+    }
+    var rotated = rotation_matrix * cp;
+    rotated = rotated * 0.2 * vec_options.length;
+    if(vec_options.scale_with_vector_length != 0u) {
+      rotated = rotated * length(vector);
+    }
+    let position = point + rotated;
+    let view_position = cameraMapPoint(position);
+    let cone_normals = array<vec3f, 7>(-normalize(Cross(cone_points[1] - cone_points[0],
+        cone_points[2] - cone_points[0])),
+        cone_points[2],
+        cone_points[1],
+        cone_points[0],
+        -normalize(Cross(cone_points[2] - cone_points[3],
+        cone_points[1] - cone_points[3])),
+        -normalize(Cross(cone_points[1] - cone_points[3],
+        cone_points[0] - cone_points[3])),
+        -normalize(Cross(cone_points[0] - cone_points[3],
+        cone_points[2] - cone_points[3])));
+
+    let normal = cone_normals[cone_strip_normals[index]];
+    return VectorFragmentInput(view_position, length(vector),
+        rotation_matrix * normal);
+}
+
+@fragment
+fn fragment_main(input: VectorFragmentInput) -> @location(0) vec4<f32> {
+    if length(input.n) == 0. {
+    discard;
+    }
+    return lightCalcColor(-input.n, getColor(input.color_val));
+}

webgpu/triangles.py

@@ -0,0 +1,66 @@
+import numpy as np
+
+from .render_object import RenderObject
+from .utils import BufferBinding, buffer_from_array, read_shader_file
+
+
+class Binding:
+    VERTICES = 90
+    NORMALS = 91
+    INDICES = 92
+
+
+class TriangulationRenderer(RenderObject):
+    n_vertices: int = 3
+
+    def __init__(self, points, normals=None, color=(0.0, 1.0, 0.0, 1.0), label="Triangulation"):
+        super().__init__(label=label)
+        self.color = color
+        self.points = np.asarray(points, dtype=np.float32).reshape(-1)
+        assert len(self.points) % 9 == 0, "Invalid number of points"
+        if normals is None:
+            ps = self.points.reshape(-1, 3, 3)
+            normals = np.cross((ps[:, 1] - ps[:, 0]), (ps[:, 2] - ps[:, 0]))
+            normals = normals / np.linalg.norm(normals, axis=1)[:, None]
+            self.normals = np.concatenate([normals, normals, normals], axis=1).flatten()
+        else:
+            self.normals = np.asarray(normals, dtype=np.float32).reshape(-1)
+        ps = self.points.reshape(-1, 3)
+        self._bounding_box = ps.min(axis=0), ps.max(axis=0)
+        self.n_instances = len(self.points) // 9
+
+    def update(self, timestamp):
+        if timestamp == self._timestamp:
+            return
+        self._timestamp = timestamp
+
+        self.point_buffer = buffer_from_array(self.points)
+        self.normal_buffer = buffer_from_array(self.normals)
+        self.create_render_pipeline()
+
+    def get_bounding_box(self):
+        return self._bounding_box
+
+    def get_color_shader(self):
+        return """
+fn getColor(vertId: u32, trigId: u32) -> vec4f {{
+  return vec4f{color};
+}}""".format(
+            color=self.color
+        )
+
+    def get_shader_code(self) -> str:
+        return (
+            read_shader_file("triangulation.wgsl", __file__)
+            + self.options.camera.get_shader_code()
+            + self.options.light.get_shader_code()
+            + self.get_color_shader()
+        )
+
+    def get_bindings(self):
+        return [
+            *self.options.camera.get_bindings(),
+            *self.options.light.get_bindings(),
+            BufferBinding(Binding.VERTICES, self.point_buffer),
+            BufferBinding(Binding.NORMALS, self.normal_buffer),
+        ]

webgpu/uniforms.py

@@ -0,0 +1,111 @@
+"""Python equivalents to all uniforms defined in shader code
+
+The UniformBase class are derived from ctypes.Structure to ensure correct memory layout.
+
+CAUTION: 
+- The Binding numbers must match the numbers defined in the shader code. 
+- Uniforms structs must match exactly the memory layout defined in the shader code.
+- The size of each struct must be a multiple of 16 bytes.
+"""
+
+import ctypes as ct
+
+from .utils import BaseBinding, UniformBinding
+from .webgpu_api import BufferUsage, Device
+
+
+class Binding:
+    """Binding numbers for uniforms in shader code in uniforms.wgsl"""
+
+    CAMERA = 0
+    CLIPPING = 1
+    FONT = 2
+    FONT_TEXTURE = 3
+    COLORMAP = 5
+    COLORMAP_TEXTURE = 6
+    COLORMAP_SAMPLER = 7
+
+    EDGES = 8
+    TRIGS = 9
+    TRIG_FUNCTION_VALUES = 10
+    SEG_FUNCTION_VALUES = 11
+    VERTICES = 12
+    TRIGS_INDEX = 13
+    GBUFFERLAM = 14
+
+    MESH = 20
+    EDGE = 21
+    SEG = 22
+    TRIG = 23
+    QUAD = 24
+    TET = 25
+    PYRAMID = 26
+    PRISM = 27
+    HEX = 28
+
+    TEXT = 30
+
+    LINE_INTEGRAL_CONVOLUTION = 40
+    LINE_INTEGRAL_CONVOLUTION_INPUT_TEXTURE = 41
+    LINE_INTEGRAL_CONVOLUTION_OUTPUT_TEXTURE = 42
+
+
+class UniformBase(ct.Structure):
+
+    def __init__(self, device: Device, **kwargs):
+        super().__init__(**kwargs)
+
+        self.device = device
+        self._buffer = device.createBuffer(
+            size=len(bytes(self)),
+            usage=BufferUsage.UNIFORM | BufferUsage.COPY_DST,
+            label=type(self).__name__,
+        )
+
+        size = len(bytes(self))
+        if size % 16:
+            raise ValueError(
+                f"Size of type {type(self)} must be multiple of 16, current size: {size}"
+            )
+
+    def update_buffer(self):
+        self.device.queue.writeBuffer(self._buffer, 0, bytes(self))
+
+    def get_bindings(self) -> list[BaseBinding]:
+        return [UniformBinding(self._binding, self._buffer)]
+
+    def __del__(self):
+        self._buffer.destroy()
+
+
+class MeshUniforms(UniformBase):
+    _binding = Binding.MESH
+    _fields_ = [
+        ("subdivision", ct.c_uint32),
+        ("shrink", ct.c_float),
+        ("padding", ct.c_float * 2),
+    ]
+
+    def __init__(self, device, subdivision=1, shrink=1.0, **kwargs):
+        super().__init__(device, subdivision=subdivision, shrink=shrink, **kwargs)
+
+
+class LineIntegralConvolutionUniforms(UniformBase):
+    _binding = Binding.LINE_INTEGRAL_CONVOLUTION
+    _fields_ = [
+        ("width", ct.c_uint32),
+        ("height", ct.c_uint32),
+        ("kernel_length", ct.c_uint32),
+        ("oriented", ct.c_uint32),
+        ("thickness", ct.c_uint32),
+        ("padding", ct.c_float * 3),
+    ]
+
+    def __init__(self, device, kernel_length=25, oriented=0, thickness=5, **kwargs):
+        super().__init__(
+            device,
+            kernel_length=kernel_length,
+            oriented=oriented,
+            thickness=thickness,
+            **kwargs,
+        )