ngsolve-webgpu 0.0.1__py3-none-any.whl

ngsolve_webgpu/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.0.1'
+__version_tuple__ = version_tuple = (0, 0, 1)

ngsolve_webgpu/animate.py

@@ -0,0 +1,76 @@
+from webgpu.render_object import RenderObject
+import ngsolve as ngs
+
+
+class Animation(RenderObject):
+    def __init__(self, child):
+        super().__init__()
+        self.child = child
+        self.data = child.data
+        self.time_index = -1
+        self.max_time = -1
+        self.gfs = set()
+        self.parameters = dict()
+        f = self.data.cf
+        self.crawl_function(f)
+        # initial solution
+        self.add_time(initial=True)
+        self.store = True
+
+    def update(self, timestamp):
+        self.child.options = self.options
+        self.child.update(timestamp)
+
+    def get_bounding_box(self):
+        return self.child.get_bounding_box()
+
+    def crawl_function(self, f):
+        if f is None:
+            return
+        if isinstance(f, ngs.GridFunction):
+            self.gfs.add(f)
+        elif isinstance(f, ngs.Parameter) or isinstance(f, ngs.ParameterC):
+            self.parameters[f] = []
+        else:
+            for c in f.data["childs"]:
+                self.crawl_function(c)
+
+    def add_time(self, initial=False):
+        self.max_time += 1
+        self.time_index = self.max_time
+        for gf in self.gfs:
+            gf.AddMultiDimComponent(gf.vec)
+        for par, vals in self.parameters.items():
+            vals.append(par.Get())
+        if not initial:
+            self.slider.max(self.max_time)
+            # set value triggers set_time_index
+            self.slider.setValue(self.time_index)
+
+    def redraw(self, timestamp: float | None = None):
+        if self.store:
+            self.add_time()
+        else:
+            self.child.redraw(timestamp)
+
+    def render(self, encoder):
+        self.child.render(encoder)
+
+    def add_options_to_gui(self, gui):
+        self.slider = gui.slider(
+            0,
+            self.set_time_index,
+            min=0,
+            max=0,
+            step=1,
+            label="animate",
+        )
+        self.child.add_options_to_gui(gui)
+
+    def set_time_index(self, time_index):
+        self.time_index = time_index
+        for gf in self.gfs:
+            gf.vec.data = gf.vecs[time_index + 1]
+        for p, vals in self.parameters.items():
+            p.Set(vals[time_index])
+        self.child.redraw()

ngsolve_webgpu/cf.py

@@ -0,0 +1,369 @@
+import math
+
+import ngsolve as ngs
+import ngsolve.webgui
+import numpy as np
+from webgpu.clipping import Clipping
+from webgpu.colormap import Colormap
+from webgpu.render_object import RenderObject
+from webgpu.utils import (
+    BufferBinding,
+    UniformBinding,
+    buffer_from_array,
+    read_shader_file,
+)
+from webgpu.vectors import BaseVectorRenderObject, VectorRenderer
+from webgpu.webgpu_api import Buffer
+
+from .mesh import Binding as MeshBinding, Mesh2dElementsRenderer
+from .mesh import ElType, MeshData
+
+
+class Binding:
+    FUNCTION_VALUES_2D = 10
+    COMPONENT = 55
+
+_intrules_3d = {}
+
+
+def get_3d_intrules(order):
+    if order in _intrules_3d:
+        return _intrules_3d[order]
+    ref_pts = [
+        [(order - i - j - k) / order, k / order, j / order]
+        for i in range(order + 1)
+        for j in range(order + 1 - i)
+        for k in range(order + 1 - i - j)
+    ]
+    p1_tets = {ngs.ET.TET: [[(1, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 0)]]}
+    p1_tets[ngs.ET.PYRAMID] = [
+        [(1, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 0)],
+        [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0)],
+    ]
+    p1_tets[ngs.ET.PRISM] = [
+        [(1, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 0)],
+        [(0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0)],
+        [(1, 0, 1), (0, 1, 1), (1, 0, 0), (0, 0, 1)],
+    ]
+    p1_tets[ngs.ET.HEX] = [
+        [(1, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 0)],
+        [(0, 1, 1), (1, 1, 1), (1, 1, 0), (1, 0, 1)],
+        [(1, 0, 1), (0, 1, 1), (1, 0, 0), (0, 0, 1)],
+        [(0, 1, 1), (1, 1, 0), (0, 1, 0), (1, 0, 0)],
+        [(0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0)],
+        [(1, 0, 1), (1, 1, 0), (0, 1, 1), (1, 0, 0)],
+    ]
+    rules = {}
+    if order > 1:
+        ho_tets = {}
+        for eltype in p1_tets:
+            for tet in p1_tets[eltype]:
+                ho_tets[eltype] = []
+                for lam in ref_pts:
+                    lami = [*lam, 1 - sum(lam)]
+                    ho_tets[eltype].append(
+                        [sum([lami[j] * tet[j][i] for j in range(4)]) for i in range(3)]
+                    )
+            rules[eltype] = ngs.IntegrationRule(ho_tets[eltype])
+    else:
+        for eltype in p1_tets:
+            rules[eltype] = ngs.IntegrationRule(sum(p1_tets[eltype], []))
+    _intrules_3d[order] = rules
+    return rules
+
+
+def _get_bernstein_matrix_trig(n, intrule):
+    """Create inverse vandermonde matrix for the Bernstein basis functions on a triangle of degree n and given integration points"""
+    ndtrig = int((n + 1) * (n + 2) / 2)
+
+    mat = ngs.Matrix(ndtrig, ndtrig)
+    fac_n = math.factorial(n)
+    for row, ip in enumerate(intrule):
+        col = 0
+        x = 1.0 - ip.point[0] - ip.point[1]
+        y = ip.point[1]
+        z = 1.0 - x - y
+        for i in range(n + 1):
+            factor = fac_n / math.factorial(i) * x**i
+            for j in range(n + 1 - i):
+                k = n - i - j
+                factor2 = 1.0 / (math.factorial(j) * math.factorial(k))
+                mat[row, col] = factor * factor2 * y**j * z**k
+                col += 1
+    return mat
+
+
+def evaluate_cf(cf, mesh, order):
+    """Evaluate a coefficient function on a mesh and returns the values as a flat array, ready to copy to the GPU as storage buffer.
+    The first two entries are the function dimension and the polynomial order of the stored values.
+    """
+    comps = cf.dim
+    int_points = ngsolve.webgui._make_trig(order)
+    intrule = ngs.IntegrationRule(
+        int_points,
+        [
+            0,
+        ]
+        * len(int_points),
+    )
+    ibmat = _get_bernstein_matrix_trig(order, intrule).I
+
+    ndof = ibmat.h
+
+    if isinstance(mesh, ngs.Region):
+        if mesh.VB() == ngs.VOL and mesh.mesh.dim == 3:
+            region = mesh.Boundaries()
+        else:
+            region = mesh
+    else:
+        region = mesh.Materials(".*")
+        if mesh.dim == 3:
+            region = mesh.Boundaries(".*")
+    pts = region.mesh.MapToAllElements(
+        {ngs.ET.TRIG: intrule, ngs.ET.QUAD: intrule}, region
+    )
+    pmat = cf(pts)
+    minval, maxval = (
+        (min(pmat.reshape(-1)), max(pmat.reshape(-1))) if len(pmat) else (0, 1)
+    )
+    pmat = pmat.reshape(-1, ndof, comps)
+
+    values = np.zeros((ndof, pmat.shape[0], comps), dtype=np.float32)
+    for i in range(comps):
+        ngsmat = ngs.Matrix(pmat[:, :, i].transpose())
+        values[:, :, i] = ibmat * ngsmat
+
+    values = values.transpose((1, 0, 2)).flatten()
+    ret = np.concatenate(([np.float32(cf.dim), np.float32(order)], values.reshape(-1)))
+    # print("ret = ", ret)
+    return ret, minval, maxval
+
+
+class FunctionData:
+    mesh_data: MeshData
+    data_2d: np.ndarray | None = None
+    data_3d: np.ndarray | None = None
+    gpu_2d: Buffer | None = None
+    gpu_3d: Buffer | None = None
+    cf: ngs.CoefficientFunction
+    order: int
+    order_3d: int
+    _timestamp: float = -1
+    minval: float = 1e99
+    maxval: float = -1e99
+
+    def __init__(
+        self,
+        mesh_data: MeshData,
+        cf: ngs.CoefficientFunction,
+        order: int,
+        order3d: int = -1,
+    ):
+        self.mesh_data = mesh_data
+        self.cf = cf
+        self.order = order
+        self.order_3d = order if order3d == -1 else order3d
+        self.need_3d = False
+
+    def update(self, timestamp: float):
+        if self._timestamp == timestamp:
+            return
+        self._timestamp = timestamp
+        if self.need_3d:
+            self.mesh_data.need_3d = True
+        self.mesh_data.update(timestamp)
+        self._create_data()
+
+    def _create_data(self):
+        self.gpu_2d = None
+        self.gpu_3d = None
+        self.data_2d, self.minval, self.maxval = evaluate_cf(
+            self.cf, self.mesh_data.ngs_mesh, self.order
+        )
+        if self.need_3d:
+            self.data_3d, minval, maxval = self.evaluate_3d(
+                self.cf, self.mesh_data.ngs_mesh, self.order_3d
+            )
+            self.minval = min(self.minval, minval)
+            self.maxval = max(self.maxval, maxval)
+
+    def get_buffers(self):
+        buffers = self.mesh_data.get_buffers().copy()
+        if self.gpu_2d is None:
+            self.gpu_2d = buffer_from_array(self.data_2d)
+            if self.data_3d is not None:
+                self.gpu_3d = buffer_from_array(self.data_3d)
+        buffers["data_2d"] = self.gpu_2d
+        if self.gpu_3d is not None:
+            buffers["data_3d"] = self.gpu_3d
+        self.data_2d = None
+        self.data_3d = None
+        return buffers
+
+    def get_bounding_box(self):
+        return self.mesh_data.get_bounding_box()
+
+    def evaluate_3d(self, cf, region, order):
+        intrules = get_3d_intrules(order)
+        if not isinstance(region, ngs.Region):
+            region = region.Materials(".*")
+        pts = region.mesh.MapToAllElements(intrules, region)
+        V_inv = vandermonde_3d(order).T
+        vals = cf(pts).reshape(-1, len(intrules[ngs.ET.TET])).dot(V_inv)
+        vmin, vmax = vals.min(), vals.max()
+        ret = np.concatenate(
+            ([np.float32(cf.dim), np.float32(order)], vals.reshape(-1)),
+            dtype=np.float32,
+        )
+        return ret, vmin, vmax
+
+
+_vandermonde_mats = {}
+
+
+def vandermonde_3d(order):
+    if order in _vandermonde_mats:
+        return _vandermonde_mats[order]
+    basis_indices = [
+        (order - i - j - k, k, j, i)
+        for i in range(order + 1)
+        for j in range(order + 1 - i)
+        for k in range(order + 1 - i - j)
+    ]
+    n = len(basis_indices)
+    V = np.zeros((n, n))
+    for r, (i, j, k, l) in enumerate(basis_indices):
+        for c, (a, b, c2, d) in enumerate(basis_indices):
+            multinom_coef = math.factorial(order) / (
+                math.factorial(a)
+                * math.factorial(b)
+                * math.factorial(c2)
+                * math.factorial(d)
+            )
+            V[r, c] = (
+                multinom_coef
+                * (i / order) ** a
+                * (j / order) ** b
+                * (k / order) ** c2
+                * (l / order) ** d
+            )
+    _vandermonde_mats[order] = np.linalg.inv(V)
+    return _vandermonde_mats[order]
+
+
+class CFRenderer(Mesh2dElementsRenderer):
+    """Use "vertices", "index" and "trig_function_values" buffers to render a mesh"""
+    fragment_entry_point = "fragmentTrig"
+
+    def __init__(self, data: FunctionData, component=0, label="CFRenderer"):
+        super().__init__(data=data.mesh_data, label=label)
+        self.data = data
+        self.colormap = Colormap()
+        self.component = component
+
+    def update(self, timestamp):
+        if timestamp == self._timestamp:
+            return
+        self._timestamp = timestamp
+        self.data.update(timestamp)
+        self._buffers = self.data.get_buffers()
+        self.colormap.options = self.options
+
+        self.curvature_subdivision = self.data.mesh_data.curvature_subdivision
+        self.n_vertices = 3 * self.curvature_subdivision**2
+        if self.colormap.autoupdate:
+            self.colormap.set_min_max(
+                self.data.minval, self.data.maxval, set_autoupdate=False
+            )
+        self.colormap.update(timestamp)
+        self.clipping.update(timestamp)
+        self.n_instances = self.data.mesh_data.num_elements[ElType.TRIG]
+        self.component_buffer = buffer_from_array(np.array([self.component], np.int32))
+        self.create_render_pipeline()
+
+    def get_bounding_box(self):
+        return self.data.get_bounding_box()
+
+    def add_options_to_gui(self, gui):
+        if self.data.cf.dim > 1:
+            options = {"Norm": 0}
+            for d in range(self.data.cf.dim):
+                options[str(d)] = d + 1
+            gui.dropdown(func=self.change_cf_dim, label="Component", values=options)
+
+    def change_cf_dim(self, value):
+        self.component = value
+        self.component_buffer = buffer_from_array(np.array([self.component], np.int32))
+        self.options.render_function()
+
+    def get_shader_code(self):
+        shader_code = ""
+
+        for file_name in [
+            "eval.wgsl",
+            "mesh.wgsl",
+            "shader.wgsl",
+            "uniforms.wgsl",
+        ]:
+            shader_code += read_shader_file(file_name, __file__)
+
+        shader_code += self.colormap.get_shader_code()
+        shader_code += self.clipping.get_shader_code()
+        shader_code += self.options.camera.get_shader_code()
+        shader_code += self.options.light.get_shader_code()
+        return shader_code
+
+    def get_bindings(self):
+        return [*super().get_bindings(),
+                *self.colormap.get_bindings(),
+                BufferBinding(Binding.FUNCTION_VALUES_2D, self._buffers["data_2d"]),
+                BufferBinding(Binding.COMPONENT, self.component_buffer)]
+
+class VectorCFRenderer(VectorRenderer):
+    def __init__(
+            self, cf: ngs.CoefficientFunction, mesh: ngs.Mesh, grid_size=20, size=None,
+    ):
+        # calling super-super class to not create points and vectors
+        BaseVectorRenderObject.__init__(self)
+        self.cf = cf
+        self.mesh = mesh
+        # this somehow segfaults in pyodide?
+        self.grid_size = grid_size
+        self.size = size
+
+    def redraw(self, timestamp=None):
+        super().redraw(
+            timestamp=timestamp, cf=self.cf, mesh=self.mesh, grid_size=self.grid_size
+        )
+
+    def update(self, timestamp: float):
+        if self._timestamp == timestamp:
+            return
+        bb = self.mesh.ngmesh.bounding_box
+        self.bounding_box = np.array(
+            [[bb[0][0], bb[0][1], bb[0][2]], [bb[1][0], bb[1][1], bb[1][2]]]
+        )
+        vs = np.linspace(
+            self.bounding_box[0][0],
+            self.bounding_box[1][0],
+            self.grid_size + 1,
+            endpoint=False,
+        )[1:]
+        points = np.meshgrid(vs, vs)
+        xvals = points[0].flatten()
+        yvals = points[1].flatten()
+        self.size = self.size or 1 / 60 * np.linalg.norm(
+            self.bounding_box[1] - self.bounding_box[0]
+        )
+        mpts_ = self.mesh(xvals, yvals, 0.0)
+        pts, mpts = [], []
+        for i in range(len(xvals)):
+            if mpts_[i]["nr"] != -1:
+                mpts.append(mpts_[i])
+                pts.append([xvals[i], yvals[i], 0.0])
+        self.points = np.array(pts, dtype=np.float32).reshape(-1)
+        values = self.cf(mpts)
+        self.vectors = np.array(
+            [values[:, 0], values[:, 1], np.zeros_like(values[:, 0])], dtype=np.float32
+        ).T.reshape(-1)
+        super().update(timestamp)

ngsolve_webgpu/clipping.py

@@ -0,0 +1,155 @@
+from webgpu import create_bind_group, read_shader_file
+from webgpu.utils import buffer_from_array, uniform_from_array
+from webgpu.clipping import Clipping
+from webgpu.colormap import Colormap
+from webgpu.render_object import RenderObject
+from webgpu.utils import BufferBinding, UniformBinding, ReadBuffer
+
+from webgpu.webgpu_api import *
+
+import numpy as np
+
+from .cf import FunctionData
+
+from .mesh import Mesh3dElementsRenderObject, ElType
+from .mesh import Binding as MeshBinding
+
+
+class VolumeCF(Mesh3dElementsRenderObject):
+    fragment_entry_point: str = "cf_fragment_main"
+
+    def __init__(self, data: FunctionData):
+        super().__init__(data=data.mesh_data)
+        self.data = data
+        self.data.need_3d = True
+        self.colormap = Colormap()
+
+    def update(self, timestamp):
+        if self._timestamp == timestamp:
+            return
+        self.colormap.options = self.options
+        self.colormap.update(timestamp)
+        super().update(timestamp)
+
+    def get_bindings(self):
+        return super().get_bindings() + [
+            BufferBinding(10, self._buffers["data_3d"]),
+            *self.colormap.get_bindings(),
+        ]
+
+    def get_shader_code(self):
+        eval_code = read_shader_file("eval.wgsl", __file__)
+        return super().get_shader_code() + self.colormap.get_shader_code() + eval_code
+
+
+class ClippingCF(RenderObject):
+    compute_shader = "clipping/compute.wgsl"
+    n_vertices = 3
+    subdivision = 0
+
+    def __init__(self, data: FunctionData):
+        super().__init__()
+        self.clipping = Clipping()
+        self.colormap = Colormap()
+        self.clipping.callbacks.append(self.build_clip_plane)
+        self.data = data
+        self.data.need_3d = True
+
+    def update(self, timestamp):
+        if timestamp == self._timestamp:
+            return
+        self._timestamp = timestamp
+        self.data.update(timestamp)
+        self.clipping.update(timestamp)
+        self.colormap.options = self.options
+        self.colormap.update(timestamp)
+        self._buffers = self.data.get_buffers()
+        self.build_clip_plane()
+
+    def get_bounding_box(self):
+        return self.data.get_bounding_box()
+
+    def get_shader_code(self, compute=False):
+        shader_code = ""
+        shader_code += self.clipping.get_shader_code()
+        shader_code += self.options.camera.get_shader_code()
+        shader_code += read_shader_file("clipping/common.wgsl", __file__)
+        shader_code += read_shader_file("eval/common.wgsl", __file__)
+        shader_code += read_shader_file("eval/tet.wgsl", __file__)
+        if compute:
+            shader_code += read_shader_file(self.compute_shader, __file__)
+        else:
+            shader_code += read_shader_file("clipping/render.wgsl", __file__)
+            shader_code += self.colormap.get_shader_code()
+            shader_code += self.options.light.get_shader_code()
+        return shader_code
+
+    def get_bindings(self, compute=False):
+        bindings = [
+            *self.options.camera.get_bindings(),
+            BufferBinding(MeshBinding.VERTICES, self._buffers["vertices"]),
+            UniformBinding(22, self.n_tets),
+            UniformBinding(23, self.only_count),
+            BufferBinding(MeshBinding.TET, self._buffers[ElType.TET]),
+            BufferBinding(13, self._buffers["data_3d"]),
+            *self.clipping.get_bindings(),
+        ]
+        if compute:
+            bindings += [
+                BufferBinding(
+                    21,
+                    self.trig_counter,
+                    read_only=False,
+                    visibility=ShaderStage.COMPUTE,
+                ),
+                BufferBinding(24, self.cut_trigs, read_only=False),
+            ]
+        else:
+            bindings += [
+                *self.colormap.get_bindings(),
+                *self.options.light.get_bindings(),
+                BufferBinding(24, self.cut_trigs),
+            ]
+        return bindings
+
+    def build_clip_plane(self):
+        for count in [True, False]:
+            encoder = self.device.createCommandEncoder("build_clip_plane")
+            ntets = self.data.mesh_data.num_elements[ElType.TET] * 4**self.subdivision
+            self.trig_counter = buffer_from_array(
+                np.array([0], dtype=np.uint32),
+                usage=BufferUsage.STORAGE | BufferUsage.COPY_DST | BufferUsage.COPY_SRC,
+            )
+            self.n_tets = uniform_from_array(np.array([ntets], dtype=np.uint32))
+            self.only_count = uniform_from_array(np.array([count], dtype=np.uint32))
+            if count:
+                self.cut_trigs = buffer_from_array(
+                    np.array([0.0] * 64, dtype=np.float32)
+                )
+            else:
+                self.cut_trigs = self.device.createBuffer(
+                    size=64 * self.n_instances, usage=BufferUsage.STORAGE
+                )
+            layout, group = create_bind_group(
+                self.device, self.get_bindings(compute=True), label="create_clip_plane"
+            )
+            shader_module = self.device.createShaderModule(
+                code=self.get_shader_code(compute=True)
+            )
+            pipeline = self.device.createComputePipeline(
+                self.device.createPipelineLayout([layout]),
+                label="create_clip_plane",
+                compute=ComputeState(module=shader_module, entryPoint="main"),
+            )
+            compute_pass = encoder.beginComputePass(label="build_clip_plane")
+            compute_pass.setPipeline(pipeline)
+            compute_pass.setBindGroup(0, group)
+            compute_pass.dispatchWorkgroups(1024)
+            compute_pass.end()
+            if count:
+                read = ReadBuffer(self.trig_counter, encoder)
+            self.device.queue.submit([encoder.finish()])
+            if count:
+                array = read.get_array(dtype=np.uint32)
+                self.n_instances = int(array[0])
+        self.create_render_pipeline()