mesh 2d element colors

Author: ChrLackner

ngsolve_webgpu/__init__.py

@@ -3,6 +3,7 @@
 from .clipping import ClippingCF
 from webgpu.colormap import Colorbar, Colormap
 from webgpu.clipping import Clipping
+from .geometry import GeometryRenderer
 
 
 from webgpu.utils import register_shader_directory as _register_shader_directory

ngsolve_webgpu/cf.py

@@ -15,7 +15,7 @@
 from webgpu.vectors import BaseVectorRenderer, VectorRenderer
 from webgpu.webgpu_api import Buffer
 
-from .mesh import Binding as MeshBinding, MeshElements2d
+from .mesh import Binding as MeshBinding, BaseMeshElements2d
 from .mesh import ElType, MeshData
 
 
@@ -264,7 +264,7 @@ def vandermonde_3d(order):
     return _vandermonde_mats[order]
 
 
-class CFRenderer(MeshElements2d):
+class CFRenderer(BaseMeshElements2d):
     """Use "vertices", "index" and "trig_function_values" buffers to render a mesh"""
 
     fragment_entry_point = "fragmentTrig"

ngsolve_webgpu/geometry.py

@@ -20,6 +20,12 @@ class Binding:
 
 
 class BaseGeometryRenderer(Renderer):
+    clipping: Clipping | None = None
+
+    def __init__(self, clipping, *args, **kwargs):
+        self.clipping = clipping
+        super().__init__(*args, **kwargs)
+
     def create_render_pipeline(self, options):
         super().create_render_pipeline(options)
         self.create_pick_pipeline(options)
@@ -71,10 +77,9 @@ def pick_index_render(self, encoder, texture, depth_texture, load_op):
 
 class GeometryFaceRenderer(BaseGeometryRenderer):
     n_vertices: int = 3
-    clipping: Clipping | None = None
 
-    def __init__(self, geo):
-        super().__init__(label="GeometryFaces")
+    def __init__(self, geo, clipping):
+        super().__init__(clipping, label="GeometryFaces")
         self.geo = geo
         self.colors = None
         self.active = True
@@ -119,15 +124,14 @@ def get_shader_code(self):
 class GeometryEdgeRenderer(BaseGeometryRenderer):
     n_vertices: int = 4
     topology: PrimitiveTopology = PrimitiveTopology.triangle_strip
-    clipping: Clipping | None = None
 
     # make sure that edges are rendered on top of faces
     depthBias: int = -5
     depthBiasSlopeScale: int = -5
 
-    def __init__(self, geo):
+    def __init__(self, geo, clipping):
         self.geo = geo
-        super().__init__(label="GeometryEdges")
+        super().__init__(clipping, label="GeometryEdges")
         self.active = True
         self.thickness = 0.005
         self._buffers = {}
@@ -163,11 +167,10 @@ def get_bindings(self):
 class GeometryVertexRenderer(BaseGeometryRenderer):
     n_vertices: int = 4
     topology: PrimitiveTopology = PrimitiveTopology.triangle_strip
-    clipping: Clipping | None = None
 
-    def __init__(self, geo):
+    def __init__(self, geo, clipping):
         self.geo = geo
-        super().__init__(label="GeometryVertices")
+        super().__init__(clipping, label="GeometryVertices")
         self.active = True
         self.thickness = 0.05
         self._buffers = {}
@@ -203,12 +206,12 @@ def get_bindings(self):
 
 
 class GeometryRenderer(MultipleRenderer):
-    def __init__(self, geo, label="Geometry"):
+    def __init__(self, geo, label="Geometry", clipping=None):
         self.geo = geo
-        self.faces = GeometryFaceRenderer(geo)
-        self.edges = GeometryEdgeRenderer(geo)
-        self.vertices = GeometryVertexRenderer(geo)
-        self.clipping = Clipping()
+        self.clipping = clipping or Clipping()
+        self.faces = GeometryFaceRenderer(geo, self.clipping)
+        self.edges = GeometryEdgeRenderer(geo, self.clipping)
+        self.vertices = GeometryVertexRenderer(geo, self.clipping)
         self.faces.clipping = self.clipping
         self.edges.clipping = self.clipping
         self.vertices.clipping = self.clipping

ngsolve_webgpu/mesh.py

@@ -5,6 +5,7 @@
 from webgpu.clipping import Clipping
 from webgpu.font import Font
 from webgpu.renderer import Renderer, RenderOptions, check_timestamp
+from webgpu.colormap import Colormap
 
 # from webgpu.uniforms import Binding
 from webgpu.uniforms import UniformBase, ct
@@ -206,7 +207,7 @@ def _create_data(self):
         self.num_elements[ElType.TRIG] = len(trigs)
         trigs_data = np.zeros((len(trigs), 4), dtype=np.uint32)
         trigs_data[:, :3] = trigs["nodes"][:, :3] - 1
-        trigs_data[:, 3] = trigs["index"]
+        trigs_data[:, 3] = trigs["index"] - 1
         self.elements[ElType.TRIG] = trigs_data
 
         # 3d Elements
@@ -276,7 +277,7 @@ def get_buffers(self):
         return self.gpu_elements
 
 
-class MeshElements2d(Renderer):
+class BaseMeshElements2d(Renderer):
     depthBias: int = 1
     depthBiasSlopeScale: float = 1.0
     vertex_entry_point: str = "vertexTrigP1Indexed"
@@ -328,8 +329,29 @@ def get_bindings(self):
     def get_shader_code(self):
         return read_shader_file("ngsolve/mesh.wgsl")
 
+class MeshElements2d(BaseMeshElements2d):
+    fragment_entry_point = "fragment2dElement"
 
-class MeshWireframe2d(MeshElements2d):
+    def __init__(self, data: MeshData, clipping=None,
+                 colors: list | None = None,
+                 label="MeshElements2d"):
+        super().__init__(data, label=label, clipping=clipping)
+        if colors is None:
+            mesh = data.mesh
+            colors = [[int(ci * 255) for ci in fd.color] for fd in mesh.FaceDescriptors()]
+        self.colormap = Colormap(colormap=colors, minval=-0.5, maxval=len(colors)-0.5)
+        self.colormap.discrete = 0
+        self.colormap.n_colors = 4*len(colors)
+
+    def update(self, options: RenderOptions):
+        super().update(options)
+        self.colormap.update(options)
+
+    def get_bindings(self):
+        return super().get_bindings() + self.colormap.get_bindings()
+
+
+class MeshWireframe2d(BaseMeshElements2d):
     depthBias: int = 0
     depthBiasSlopeScale: float = 0.
     topology: PrimitiveTopology = PrimitiveTopology.line_strip
@@ -374,6 +396,7 @@ def shrink(self, value):
         self._shrink = value
         if self.uniforms is not None:
             self.uniforms.shrink = value
+            self.uniforms.update_buffer()
 
     def get_bounding_box(self) -> tuple[list[float], list[float]] | None:
         return self.data.get_bounding_box()

ngsolve_webgpu/shaders/mesh.wgsl

@@ -247,7 +247,7 @@ fn vertexMeshHex(@builtin(vertex_index) vertId: u32, @builtin(instance_index) el
 @fragment
 fn fragment2dElement(input: VertexOutput2d) -> @location(0) vec4<f32> {
   checkClipping(input.p);
-  return lightCalcColor(input.p, input.n, u_mesh_color);
+  return lightCalcColor(input.p, input.n, getColor(f32(input.index)));
 }
 
 @fragment

ngsolve_webgpu/shaders/shader.wgsl

@@ -13,6 +13,7 @@ struct VertexOutput2d {
   @location(1) lam: vec2<f32>,
   @location(2) @interpolate(flat) id: u32,
   @location(3) n: vec3<f32>,
+  @location(4) @interpolate(flat) index: u32,
 };
 
 struct VertexOutput3d {
@@ -37,7 +38,8 @@ fn vertexEdgeP1(@builtin(vertex_index) vertexId: u32, @builtin(instance_index) e
     return VertexOutput1d(position, p, lam, edgeId);
 }
 
-fn calcTrig(p: array<vec3<f32>, 3>, vertexId: u32, trigId: u32) -> VertexOutput2d {
+fn calcTrig(p: array<vec3<f32>, 3>, vertexId: u32, trigId: u32, index: u32)
+  -> VertexOutput2d {
     let subdivision = u_subdivision;
     let h = 1.0 / f32(subdivision);
 
@@ -83,7 +85,8 @@ fn calcTrig(p: array<vec3<f32>, 3>, vertexId: u32, trigId: u32) -> VertexOutput2
 
     let mapped_position = cameraMapPoint(position);
 
-    return VertexOutput2d(mapped_position, position, lam, trigId, normal);
+    return VertexOutput2d(mapped_position, position, lam, trigId, normal,
+                          index);
 }
 
 
@@ -95,12 +98,13 @@ fn vertexTrigP1Indexed(@builtin(vertex_index) vertexId: u32, @builtin(instance_i
         trigs[4 * trigId + 2]
     );
 
+    let index = trigs[4 * trigId + 3];
     var p = array<vec3<f32>, 3>(
         vec3<f32>(vertices[vid[0] ], vertices[vid[0] + 1], vertices[vid[0] + 2]),
         vec3<f32>(vertices[vid[1] ], vertices[vid[1] + 1], vertices[vid[1] + 2]),
         vec3<f32>(vertices[vid[2] ], vertices[vid[2] + 1], vertices[vid[2] + 2])
     );
-    return calcTrig(p, vertexId, trigId);
+    return calcTrig(p, vertexId, trigId, index);
 }
 
 @vertex
@@ -110,6 +114,7 @@ fn vertexWireframe2d(@builtin(vertex_index) vertexId: u32, @builtin(instance_ind
         trigs[4 * trigId + 1],
         trigs[4 * trigId + 2]
     );
+    let index = trigs[4 * trigId + 3];
 
     var p = array<vec3<f32>, 3>(
         vec3<f32>(vertices[vid[0] ], vertices[vid[0] + 1], vertices[vid[0] + 2]),
@@ -155,7 +160,8 @@ fn vertexWireframe2d(@builtin(vertex_index) vertexId: u32, @builtin(instance_ind
       position += u_deformation_scale * evalTrigVec3(&u_deformation_values_2d, trigId, lam);
     }
     return VertexOutput2d(cameraMapPoint(position), position, lam, trigId,
-                          normalize(cross(p[1] - p[0], p[2] - p[0])));
+                          normalize(cross(p[1] - p[0], p[2] - p[0])),
+                          index);
 }