1// Rigid-body projection: (1) mass-weighted sums of the six rigid modes for
2// displacement and velocity, kinetic energy and max |w| in one workgroup;
3// (2) subtraction of the rigid components from every node.
5@group(0) @binding(0) var<uniform> P: Params;
6@group(0) @binding(1) var<storage, read_write> pos: array<f32>; // [u | v | w | wLo]
7@group(0) @binding(2) var<storage, read_write> vel: array<f32>; // [du | dv | dw]
8@group(0) @binding(3) var<storage, read_write> stats: array<f32>; // 0..11 coefficients, 12 KE, 13 max|w|
10const NS = 14u;
11var<workgroup> sh: array<array<f32, NS>, 256>;
13@compute @workgroup_size(256)
14fn sums(@builtin(local_invocation_id) lid: vec3u) {
15 let t = lid.x;
16 var s: array<f32, NS>;
17 for (var q = 0u; q < NS; q++) { s[q] = 0.0; }
18 for (var k = t; k < P.n; k += 256u) {
19 let i = i32(k % P.nx); let j = i32(k / P.nx);
20 let x = nodeX(i); let y = nodeY(j);
21 let m = P.rhoh * P.hx * P.hy * trapWeight(i, j);
22 let n = P.n;
23 let p = vec3f(pos[k], pos[n + k], pos[2u * n + k]);
24 let v = vec3f(vel[k], vel[n + k], vel[2u * n + k]);
25 s[0] += m * p.x; s[1] += m * p.y; s[2] += m * p.z;
26 s[3] += m * (x * p.y - y * p.x); s[4] += m * y * p.z; s[5] += m * (-x * p.z);
27 s[6] += m * v.x; s[7] += m * v.y; s[8] += m * v.z;
28 s[9] += m * (x * v.y - y * v.x); s[10] += m * y * v.z; s[11] += m * (-x * v.z);
29 s[12] += 0.5 * m * dot(v, v);
30 s[13] = max(s[13], abs(p.z));
31 }
32 for (var q = 0u; q < NS; q++) { sh[t][q] = s[q]; }
33 workgroupBarrier();
34 for (var stride = 128u; stride > 0u; stride >>= 1u) {
35 if (t < stride) {
36 for (var q = 0u; q < 13u; q++) { sh[t][q] += sh[t + stride][q]; }
37 sh[t][13] = max(sh[t][13], sh[t + stride][13]);
38 }
39 workgroupBarrier();
40 }
41 if (t == 0u) {
42 let inv = array<f32, 6>(1.0 / P.mass, 1.0 / P.mass, 1.0 / P.mass, 1.0 / P.Iz, 1.0 / P.Ix, 1.0 / P.Iy);
43 for (var q = 0u; q < 6u; q++) {
44 stats[q] = sh[0][q] * inv[q];
45 stats[q + 6u] = sh[0][q + 6u] * inv[q];
46 }
47 // Kinetic energy of the non-rigid motion (the six rigid modes are mass-orthogonal).
48 var keRigid = 0.0;
49 for (var q = 6u; q < 12u; q++) { keRigid += 0.5 * sh[0][q] * sh[0][q] * inv[q - 6u]; }
50 stats[12] = max(sh[0][12] - keRigid, 0.0);
51 stats[13] = sh[0][13];
52 }
53}
55@compute @workgroup_size(256)
56fn apply(@builtin(global_invocation_id) gid: vec3u) {
57 let k = gid.x;
58 if (k >= P.n) { return; }
59 let i = i32(k % P.nx); let j = i32(k / P.nx);
60 let x = nodeX(i); let y = nodeY(j);
61 let n = P.n;
62 pos[k] -= stats[0] - stats[3] * y;
63 pos[n + k] -= stats[1] + stats[3] * x;
64 pos[2u * n + k] -= stats[2] + stats[4] * y - stats[5] * x;
65 vel[k] -= stats[6] - stats[9] * y;
66 vel[n + k] -= stats[7] + stats[9] * x;
67 vel[2u * n + k] -= stats[8] + stats[10] * y - stats[11] * x;
68}
70// Removes the three rigid components (translation and two rotations) from a
71// per-node acceleration field so that it is self-equilibrated. One workgroup.
72@group(0) @binding(4) var<storage, read_write> field: array<f32>;
73var<workgroup> sh3: array<vec3f, 256>;
75@compute @workgroup_size(256)
76fn projectField(@builtin(local_invocation_id) lid: vec3u) {
77 let t = lid.x;
78 var s = vec3f(0.0);
79 for (var k = t; k < P.n; k += 256u) {
80 let i = i32(k % P.nx); let j = i32(k / P.nx);
81 let m = P.rhoh * P.hx * P.hy * trapWeight(i, j);
82 let f = field[k];
83 s += vec3f(m * f, m * nodeY(j) * f, -m * nodeX(i) * f);
84 }
85 sh3[t] = s;
86 workgroupBarrier();
87 for (var stride = 128u; stride > 0u; stride >>= 1u) {
88 if (t < stride) { sh3[t] += sh3[t + stride]; }
89 workgroupBarrier();
90 }
91 let c = vec3f(sh3[0].x / P.mass, sh3[0].y / P.Ix, sh3[0].z / P.Iy);
92 for (var k = t; k < P.n; k += 256u) {
93 let i = i32(k % P.nx); let j = i32(k / P.nx);
94 field[k] -= c.x + c.y * nodeY(j) - c.z * nodeX(i);
95 }
96}