surface vectors, lint

mhochsteger · mhochsteger · commit 79738cfb25a4 · 2025-06-12T13:30:30.000+02:00
diff --git a/ngsolve_webgpu/geometry.py b/ngsolve_webgpu/geometry.py
@@ -227,8 +227,7 @@ def update(self, options: RenderOptions):
 
     def get_bounding_box(self):
         pmin, pmax = self.geo.shape.bounding_box
-        return ([pmin[0], pmin[1], pmin[2]],
-                [pmax[0], pmax[1], pmax[2]])
+        return ([pmin[0], pmin[1], pmin[2]], [pmax[0], pmax[1], pmax[2]])
 
     def render(self, encoder):
         for r in self.render_objects:
diff --git a/ngsolve_webgpu/mesh.py b/ngsolve_webgpu/mesh.py
@@ -255,8 +255,7 @@ def _create_data(self):
 
     def get_bounding_box(self):
         pmin, pmax = self.mesh.bounding_box
-        return ([pmin[0], pmin[1], pmin[2]],
-                [pmax[0], pmax[1], pmax[2]])
+        return ([pmin[0], pmin[1], pmin[2]], [pmax[0], pmax[1], pmax[2]])
 
     def get_buffers(self):
         for eltype in self.elements:
@@ -329,19 +328,20 @@ def get_bindings(self):
     def get_shader_code(self):
         return read_shader_file("ngsolve/mesh.wgsl")
 
+
 class MeshElements2d(BaseMeshElements2d):
     fragment_entry_point = "fragment2dElement"
 
-    def __init__(self, data: MeshData, clipping=None,
-                 colors: list | None = None,
-                 label="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 = Colormap(colormap=colors, minval=-0.5, maxval=len(colors) - 0.5)
         self.colormap.discrete = 0
-        self.colormap.n_colors = 4*len(colors)
+        self.colormap.n_colors = 4 * len(colors)
 
     def update(self, options: RenderOptions):
         super().update(options)
@@ -353,7 +353,7 @@ def get_bindings(self):
 
 class MeshWireframe2d(BaseMeshElements2d):
     depthBias: int = 0
-    depthBiasSlopeScale: float = 0.
+    depthBiasSlopeScale: float = 0.0
     topology: PrimitiveTopology = PrimitiveTopology.line_strip
     color = (0, 0, 0, 1)
     fragment_entry_point: str = "fragmentWireframe2d"
diff --git a/ngsolve_webgpu/shaders/surface_vectors.wgsl b/ngsolve_webgpu/shaders/surface_vectors.wgsl
@@ -0,0 +1,120 @@
+#import ngsolve/eval/trig
+#import ngsolve/uniforms
+#import colormap
+
+@group(0) @binding(21) var<storage, read_write> count_vectors: atomic<u32>;
+@group(0) @binding(22) var<storage, read_write> positions: array<f32>;
+@group(0) @binding(23) var<storage, read_write> directions: array<f32>;
+@group(0) @binding(25) var<storage, read_write> values: array<f32>;
+@group(0) @binding(24) var<uniform> u_ntrigs: u32;
+
+@compute @workgroup_size(256)
+fn compute_surface_vectors(@builtin(global_invocation_id) id: vec3<u32>) {
+  for (var trigId = id.x; trigId<u_ntrigs; trigId+=256*1024) {
+    var vid = 3 * vec3u(
+        trigs[4 * trigId + 2],
+        trigs[4 * trigId + 0],
+        trigs[4 * trigId + 1]
+    );
+
+    let 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])
+    );
+
+    // let pmin = min(p[0], min(p[1], p[2]));
+    let pmin = vec3f(-1, -1, -1);
+    let rad = 1.0;
+    let swap_lam = false;
+
+    var dir: u32 =0;
+    var dir1: u32 =0;
+    var dir2: u32 =0;
+
+    let n = cross (p[1]-p[0], p[2]-p[0]);
+    let na  = abs(n);
+    if (na[0] > na[1] && na[0] > na[2]) {
+      dir = 0;
+    }
+    else if (na[1] > na[2]) {
+      dir = 1;
+    }
+    else  {
+      dir = 2;
+    }
+    
+    dir1 = (dir+1) % 3;
+    dir2 = (dir1+1) % 3;
+
+    var p2d: array<vec2f, 3>;
+
+    for (var k: u32 = 0; k < 3; k++)
+      {
+        p2d[k] = vec2f((p[k][dir1] - pmin[dir1]) / (2*rad),
+                       (p[k][dir2] - pmin[dir2]) / (2*rad));
+      }
+
+    var min2d = min(min(p2d[0], p2d[1]), p2d[2]);
+    var max2d = max(max(p2d[0], p2d[1]), p2d[2]);
+
+    let m = mat2x2f(
+      p2d[1] - p2d[0],
+      p2d[2] - p2d[0]
+    );
+    let mdet = determinant(m);
+
+    let minv = 1.0/mdet * mat2x2f( m[1][1], -m[0][1], -m[1][0], m[0][0] );
+    
+    let gridsize = 0.03;
+
+    let xmin = floor(min2d.x / gridsize) * gridsize;
+
+    for (var s = 0.0; s <= 1.; s += 1.0 * gridsize) {
+      if (s >= min2d.x && s <= max2d.x) 
+      {
+        for (var t = 0.; t <= 1.; t += 1.0 * gridsize) {
+          if (t >= min2d.y && t <= max2d.y)
+            {
+              let lam = minv * (vec2f(s, t) - p2d[0]);
+              
+              if (lam.x >= 0 && lam.y >= 0 && lam.x+lam.y <= 1)
+                {
+                  var cp = p[0] + lam.x * (p[1] - p[0]) + lam.y * (p[2] - p[0]);
+                  
+                  if(@MODE@ == 0) {
+  // just count
+                      atomicAdd(&count_vectors, 1);
+                    }
+                    else {
+                      // write output to buffer
+                      let v = evalTrigVec3(&u_function_values_2d, trigId, lam);
+                      
+                      let val = length(v);
+                      var scale = (val - u_cmap_uniforms.min) / (u_cmap_uniforms.max - u_cmap_uniforms.min);
+                      scale = 2 * gridsize * clamp(scale, 0.5, 1.0);
+                      let direction = scale * normalize(v) ;
+                      let index = atomicAdd(&count_vectors, 1);
+                      if (u_curvature_values_2d[0] != -1.) {
+                        cp = evalTrigVec3(&u_curvature_values_2d, trigId, lam);
+                      }
+                      cp += 0.5 * gridsize * normalize(n);
+
+                      positions[index*3+0] = cp[0];
+                      positions[index*3+1] = cp[1];
+                      positions[index*3+2] = cp[2];
+                      values[index] = val;
+                      directions[index*3+0] = direction[0];
+                      directions[index*3+1] = direction[1];
+                      directions[index*3+2] = direction[2];
+                    }
+
+                }
+
+
+  }
+    }
+      }
+    }
+  }
+}
diff --git a/ngsolve_webgpu/surface_vectors.py b/ngsolve_webgpu/surface_vectors.py
@@ -0,0 +1,142 @@
+import ngsolve as ngs
+import numpy as np
+from webgpu.shapes import ShapeRenderer, generate_cone, generate_cylinder
+from webgpu.utils import (
+    BufferUsage,
+    BufferBinding,
+    ReadBuffer,
+    UniformBinding,
+    read_buffer,
+    read_shader_file,
+    run_compute_shader,
+    uniform_from_array,
+    buffer_from_array,
+    write_array_to_buffer,
+)
+
+from .cf import FunctionData, MeshData, Binding as FunctionBinding
+from .mesh import Binding as MeshBinding
+from .mesh import ElType
+
+
+class SurfaceVectors(ShapeRenderer):
+    def __init__(
+        self,
+        function_data: FunctionData,
+        mesh: MeshData,
+        grid_size: float = 0.02,
+    ):
+        self.function_data = function_data
+        self.mesh = mesh
+
+        bbox = mesh.get_bounding_box()
+        grid_size = np.linalg.norm(np.array(bbox[1]) - np.array(bbox[0])) * grid_size
+
+        cyl = generate_cylinder(8, 0.05, 0.5, bottom_face=True)
+        cone = generate_cone(8, 0.2, 0.5, bottom_face=True)
+        arrow = cyl + cone.move((0, 0, 0.5))
+
+        super().__init__(arrow, None, None)
+        # self.scale_mode = ShapeRenderer.SCALE_Z
+
+    def get_bounding_box(self):
+        return self.mesh.get_bounding_box()
+
+    def get_compute_bindings(self):
+        return []
+
+    def compute_vectors(self):
+        self.u_nvectors = buffer_from_array(
+            np.array([0], dtype=np.uint32),
+            label="n_vectors",
+            usage=BufferUsage.STORAGE | BufferUsage.COPY_DST | BufferUsage.COPY_SRC,
+        )
+
+        mesh_buffers = self.mesh.get_buffers()
+        func_buffers = self.function_data.get_buffers()
+        n_trigs = self.mesh.num_elements[ElType.TRIG]
+        self.u_ntrigs = uniform_from_array(np.array([n_trigs], dtype=np.uint32), label="n_trigs")
+
+        positions = buffer_from_array(np.array([0], dtype=np.float32), label="positions")
+        directions = buffer_from_array(np.array([0], dtype=np.float32), label="positions")
+        values = buffer_from_array(np.array([0], dtype=np.float32), label="positions")
+
+        bindings = [
+            *self.colormap.get_bindings(),
+            BufferBinding(MeshBinding.VERTICES, mesh_buffers["vertices"]),
+            BufferBinding(MeshBinding.TRIGS_INDEX, mesh_buffers[ElType.TRIG]),
+            BufferBinding(22, positions, read_only=False),
+            BufferBinding(23, directions, read_only=False),
+            BufferBinding(25, values, read_only=False),
+            BufferBinding(21, self.u_nvectors, read_only=False),
+            UniformBinding(24, self.u_ntrigs),
+            BufferBinding(MeshBinding.CURVATURE_VALUES_2D, mesh_buffers["curvature_2d"]),
+            # BufferBinding(MeshBinding.DEFORMATION_VALUES, mesh_buffers["deformation_2d"]),
+            UniformBinding(MeshBinding.DEFORMATION_SCALE, mesh_buffers["deformation_scale"]),
+            BufferBinding(FunctionBinding.FUNCTION_VALUES_2D, func_buffers["data_2d"]),
+        ]
+        run_compute_shader(
+            read_shader_file("ngsolve/surface_vectors.wgsl"),
+            bindings,
+            min(n_trigs // 256 + 1, 1024),
+            entry_point="compute_surface_vectors",
+            defines={
+                "MODE": 0,
+                "MAX_EVAL_ORDER": self.function_data.order,
+                "MAX_EVAL_ORDER_VEC3": self.function_data.order,
+            },
+        )
+
+        self.n_vectors = int(read_buffer(self.u_nvectors, np.uint32)[0])
+        write_array_to_buffer(self.u_nvectors, np.array([0], dtype=np.uint32))
+        buffers = {}
+        for name in ["positions", "directions"]:
+            buffers[name] = self.device.createBuffer(
+                size=3 * 4 * self.n_vectors,
+                usage=BufferUsage.STORAGE | BufferUsage.COPY_SRC,
+                label=name,
+            )
+        buffers["values"] = self.device.createBuffer(
+            size=4 * self.n_vectors,
+            usage=BufferUsage.STORAGE | BufferUsage.COPY_SRC,
+            label="values",
+        )
+
+        bindings = [
+            *self.colormap.get_bindings(),
+            BufferBinding(MeshBinding.VERTICES, mesh_buffers["vertices"]),
+            BufferBinding(MeshBinding.TRIGS_INDEX, mesh_buffers[ElType.TRIG]),
+            BufferBinding(22, buffers["positions"], read_only=False),
+            BufferBinding(23, buffers["directions"], read_only=False),
+            BufferBinding(25, buffers["values"], read_only=False),
+            BufferBinding(21, self.u_nvectors, read_only=False),
+            BufferBinding(MeshBinding.CURVATURE_VALUES_2D, mesh_buffers["curvature_2d"]),
+            BufferBinding(FunctionBinding.FUNCTION_VALUES_2D, func_buffers["data_2d"]),
+            # BufferBinding(MeshBinding.DEFORMATION_VALUES, mesh_buffers["deformation_2d"]),
+            UniformBinding(MeshBinding.DEFORMATION_SCALE, mesh_buffers["deformation_scale"]),
+            UniformBinding(24, self.u_ntrigs),
+        ]
+
+        run_compute_shader(
+            read_shader_file("ngsolve/surface_vectors.wgsl"),
+            bindings,
+            min(n_trigs // 256 + 1, 1024),
+            entry_point="compute_surface_vectors",
+            defines={
+                "MODE": 1,
+                "MAX_EVAL_ORDER": self.function_data.order,
+                "MAX_EVAL_ORDER_VEC3": self.function_data.order,
+            },
+        )
+
+        self.positions = read_buffer(buffers["positions"], np.float32).reshape(-1)
+        self.values = read_buffer(buffers["values"], np.float32).reshape(-1)
+        self.directions = read_buffer(buffers["directions"], np.float32).reshape(-1)
+
+    def update(self, options):
+        self.mesh.update(options)
+        self.function_data.update(options)
+        self.colormap.update(options)
+        self.compute_vectors()
+        super().update(options)
+        return
