1// One explicit time step (symplectic Euler) of the von Karman plate with a
2// thickness field, a stress-free profile, the drumstick, and the pads of the
3// stand.
4//
5// State layout (structure of arrays, n = nx * ny nodes):
6// pos: [u | v | w | wLo] wLo carries the rounding error of the w update
7// vel: [du | dv | dw]
8// A workgroup of 16x16 threads advances a 12x12 tile of nodes in three
9// phases separated by barriers. (1) Every thread loads one entry of the tile
10// plus its 2-node halo (w, u, v, w0, thickness, node weights) into shared
11// memory. (2) The threads compute, once each and again into shared memory,
12// the membrane forces of the 13x13 cells around the tile and the weighted
13// bending moments and mixed differences that the bending stencil gathers.
14// (3) Each of the 144 tile nodes evaluates its update from shared memory.
15// On grids of a few thousand nodes the kernel is bound by the latency of one
16// workgroup rather than by throughput, so the shared phases, which replace a
17// tenfold per-node recomputation of the stencil ingredients, are what make a
18// step cheap.
19//
20// Nodes outside the plate (weight 0) are skipped: a second difference exists
21// only where the node and both neighbors belong to the plate, and a cell
22// contributes only when its four corners do. That gives the free edge of the
23// rectangle and a staircase free edge for any other outline, hole included.
25@group(0) @binding(0) var<uniform> P: Params;
26@group(0) @binding(1) var<uniform> SP: StepParams;
27@group(0) @binding(2) var<storage, read> posIn: array<f32>;
28@group(0) @binding(3) var<storage, read> velIn: array<f32>;
29@group(0) @binding(4) var<storage, read_write> posOut: array<f32>;
30@group(0) @binding(5) var<storage, read_write> velOut: array<f32>;
31@group(0) @binding(6) var<storage, read_write> hist: array<f32>; // [slot * n + node]: transverse velocity at the audio sample instants
32// Two 8-float slots of drumstick state: position, velocity, force, max force,
33// time since the contact released, withdrawn flag, and two spare.
34const STICK_RELEASE: f32 = 5e-3; // [s] a release longer than this ends the hit
35@group(0) @binding(7) var<storage, read_write> stick: array<f32>;
37const TILE = 12; // nodes per workgroup along each axis (must match simulation.ts)
38const HALO = 2;
39const TW = 16; // TILE + 2 * HALO: the loaded tile, one thread per entry
40const CT = 13; // TILE + 1: cells adjacent to the tile's nodes
41const MT = 14; // TILE + 2: nodes at which the bending moments are needed
42var<workgroup> tw: array<f32, 256>; // w
43var<workgroup> tu: array<f32, 256>; // u
44var<workgroup> tv: array<f32, 256>; // v
45var<workgroup> tw0: array<f32, 256>; // stress-free shape
46var<workgroup> th: array<f32, 256>; // thickness
47var<workgroup> tn: array<u32, 64>; // node weights in quarters (0..4), four to a word
48var<workgroup> cellF: array<f32, 2028>; // membrane forces of the CT x CT cells, 4 nodes x 3 components
49var<workgroup> ma: array<f32, 196>; // weighted bending moments at the MT x MT inner nodes
50var<workgroup> mb: array<f32, 196>;
51var<workgroup> cc: array<f32, 169>; // weighted mixed differences at the CT x CT cells
53var<private> ox: i32;
54var<private> oy: i32;
55var<private> ihx2: f32;
56var<private> ihy2: f32;
57var<private> ihxy: f32;
59fn id(i: i32, j: i32) -> u32 { return u32(j) * P.nx + u32(i); }
60fn lidx(i: i32, j: i32) -> u32 { return u32(j - oy) * u32(TW) + u32(i - ox); }
61fn midx(i: i32, j: i32) -> u32 { return u32(j - oy - 1) * u32(MT) + u32(i - ox - 1); }
62fn cidx(ci: i32, cj: i32) -> u32 { return u32(cj - oy - 1) * u32(CT) + u32(ci - ox - 1); }
63fn W(i: i32, j: i32) -> f32 { return tw[lidx(i, j)]; }
64fn U(i: i32, j: i32) -> f32 { return tu[lidx(i, j)]; }
65fn V(i: i32, j: i32) -> f32 { return tv[lidx(i, j)]; }
66fn W0(i: i32, j: i32) -> f32 { return tw0[lidx(i, j)]; }
67fn DW(i: i32, j: i32) -> f32 { return velIn[2u * P.n + id(i, j)]; }
68fn TH(i: i32, j: i32) -> f32 { return th[lidx(i, j)]; }
69fn NWq(i: i32, j: i32) -> u32 { let l = lidx(i, j); return (tn[l >> 2u] >> ((l & 3u) * 8u)) & 0xFFu; }
70fn NW(i: i32, j: i32) -> f32 { return 0.25 * f32(NWq(i, j)); }
71fn INS(i: i32, j: i32) -> bool { return NWq(i, j) != 0u; }
72fn insCell(ci: i32, cj: i32) -> bool { return INS(ci, cj) && INS(ci + 1, cj) && INS(ci + 1, cj + 1) && INS(ci, cj + 1); }
73fn Dn(i: i32, j: i32) -> f32 { let h = TH(i, j); return P.Db * h * h * h; }
74fn cellH(ci: i32, cj: i32) -> f32 { return 0.25 * (TH(ci, cj) + TH(ci + 1, cj) + TH(ci + 1, cj + 1) + TH(ci, cj + 1)); }
76// Graph Laplacian of the transverse velocity at node (i, j), over neighbors in the plate.
77fn lapDW(i: i32, j: i32) -> f32 {
78 let c = DW(i, j);
79 var lap = 0.0;
80 if (INS(i - 1, j)) { lap += (DW(i - 1, j) - c) * ihx2; }
81 if (INS(i + 1, j)) { lap += (DW(i + 1, j) - c) * ihx2; }
82 if (INS(i, j - 1)) { lap += (DW(i, j - 1) - c) * ihy2; }
83 if (INS(i, j + 1)) { lap += (DW(i, j + 1) - c) * ihy2; }
84 return lap;
85}
87// Weighted bending moments at nodes; zero where the second difference is undefined.
88fn Ma(i: i32, j: i32) -> f32 {
89 if (!(INS(i - 1, j) && INS(i, j) && INS(i + 1, j))) { return 0.0; }
90 let a = (W(i + 1, j) - 2.0 * W(i, j) + W(i - 1, j)) * ihx2;
91 var b = 0.0;
92 if (INS(i, j - 1) && INS(i, j + 1)) { b = (W(i, j + 1) - 2.0 * W(i, j) + W(i, j - 1)) * ihy2; }
93 return NW(i, j) * Dn(i, j) * (a + P.nu * b);
94}
95fn Mb(i: i32, j: i32) -> f32 {
96 if (!(INS(i, j - 1) && INS(i, j) && INS(i, j + 1))) { return 0.0; }
97 let b = (W(i, j + 1) - 2.0 * W(i, j) + W(i, j - 1)) * ihy2;
98 var a = 0.0;
99 if (INS(i - 1, j) && INS(i + 1, j)) { a = (W(i + 1, j) - 2.0 * W(i, j) + W(i - 1, j)) * ihx2; }
100 return NW(i, j) * Dn(i, j) * (b + P.nu * a);
101}
102fn Cc(ci: i32, cj: i32) -> f32 {
103 if (!insCell(ci, cj)) { return 0.0; }
104 let h = cellH(ci, cj);
105 return P.Db * h * h * h * (W(ci + 1, cj + 1) - W(ci + 1, cj) - W(ci, cj + 1) + W(ci, cj)) * ihxy;
106}
108// Membrane forces of the bilinear cell (ci, cj) on its four nodes (2x2 Gauss), written to cellF at local cell index lc.
109fn cellForces(ci: i32, cj: i32, lc: u32) {
110 let base = lc * 12u;
111 if (!insCell(ci, cj)) {
112 for (var q = 0u; q < 12u; q++) { cellF[base + q] = 0.0; }
113 return;
114 }
115 let hc = cellH(ci, cj);
116 let C = P.Cm * hc;
117 let G = P.Gm * hc;
118 let Un = vec4f(U(ci, cj), U(ci + 1, cj), U(ci + 1, cj + 1), U(ci, cj + 1));
119 let Vn = vec4f(V(ci, cj), V(ci + 1, cj), V(ci + 1, cj + 1), V(ci, cj + 1));
120 let W0n = vec4f(W0(ci, cj), W0(ci + 1, cj), W0(ci + 1, cj + 1), W0(ci, cj + 1));
121 let Wdn = vec4f(W(ci, cj), W(ci + 1, cj), W(ci + 1, cj + 1), W(ci, cj + 1)); // deviation from the stress-free shape
122 let q = 0.25 * P.hx * P.hy;
123 let ihx = 1.0 / P.hx; let ihy = 1.0 / P.hy;
124 let g3 = 0.57735026918962576;
125 var Fu = vec4f(0.0); var Fv = vec4f(0.0); var Fw = vec4f(0.0);
126 for (var gp = 0; gp < 4; gp++) {
127 var xi = -g3; if ((gp & 1) != 0) { xi = g3; }
128 var eta = -g3; if ((gp & 2) != 0) { eta = g3; }
129 let dNdx = vec4f(-(1.0 - eta), 1.0 - eta, 1.0 + eta, -(1.0 + eta)) * (0.5 * ihx);
130 let dNdy = vec4f(-(1.0 - xi), -(1.0 + xi), 1.0 + xi, 1.0 - xi) * (0.5 * ihy);
131 let ux = dot(dNdx, Un); let uy = dot(dNdy, Un);
132 let vx = dot(dNdx, Vn); let vy = dot(dNdy, Vn);
133 let dwx = dot(dNdx, Wdn); let dwy = dot(dNdy, Wdn);
134 let w0x = dot(dNdx, W0n); let w0y = dot(dNdy, W0n);
135 let wx = w0x + dwx; let wy = w0y + dwy; // total slope
136 // Marguerre strains: the total slope minus the stress-free slope, written
137 // so that every term carries the deviation. The direct form
138 // 0.5 (wx^2 - w0x^2) subtracts two nearly equal numbers, and in single
139 // precision its rounding error is proportional to w0x^2 and so does not
140 // fall with the vibration: on a curved design it acts as a constant
141 // broadband forcing that leaves a noise floor which never decays.
142 let exx = ux + dwx * (w0x + 0.5 * dwx);
143 let eyy = vy + dwy * (w0y + 0.5 * dwy);
144 let gxy = uy + vx + w0x * dwy + w0y * dwx + dwx * dwy;
145 let Nxx = C * (exx + P.nu * eyy);
146 let Nyy = C * (eyy + P.nu * exx);
147 let Nxy = G * gxy;
148 Fu -= q * (Nxx * dNdx + Nxy * dNdy);
149 Fv -= q * (Nyy * dNdy + Nxy * dNdx);
150 Fw -= q * ((Nxx * wx + Nxy * wy) * dNdx + (Nxy * wx + Nyy * wy) * dNdy);
151 }
152 for (var a = 0u; a < 4u; a++) {
153 cellF[base + a * 3u] = Fu[a];
154 cellF[base + a * 3u + 1u] = Fv[a];
155 cellF[base + a * 3u + 2u] = Fw[a];
156 }
157}
159// True for the plate nodes of the drumstick contact patch (same rule as patchAround in params.ts).
160fn inContactPatch(i: i32, j: i32) -> bool {
161 if (nodeWeight(i, j) == 0.0) { return false; }
162 let dxs = nodeX(i) - P.strikeX; let dys = nodeY(j) - P.strikeY;
163 return dxs * dxs + dys * dys <= P.strikeR2 || id(i, j) == P.strikeNode;
164}
166fn cellForce(ci: i32, cj: i32, a: u32) -> vec3f {
167 let base = cidx(ci, cj) * 12u + a * 3u;
168 return vec3f(cellF[base], cellF[base + 1u], cellF[base + 2u]);
169}
171// Relaxed force of the pads on a node with bottom-pad rate kb (at the bounce
172// frequency) and transverse displacement w (same rule as padForce in
173// reference.ts); the rest of the felt's force is in the arms of its chain.
174// The constant part that balances the weight of the plate is subtracted,
175// so w = 0 is rest.
176fn padForce(kb0: f32, w: f32) -> f32 {
177 let kb = P.kInf * kb0;
178 if (P.liftoff == 0u) {
179 if (P.hasTop != 0u) { return -2.0 * kb * w; }
180 return -kb * w;
181 }
182 var f = kb * max(P.preloadBottom - w, 0.0) - kb * P.preloadBottom;
183 if (P.hasTop != 0u) { f -= kb * max(P.preloadTop + w, 0.0) - kb * P.preloadTop; }
184 return f;
185}
186// Number of pads in contact with a node at transverse displacement w.
187fn padsInContact(w: f32) -> f32 {
188 if (P.liftoff == 0u) { return select(1.0, 2.0, P.hasTop != 0u); }
189 var nc = select(0.0, 1.0, P.preloadBottom - w > 0.0);
190 if (P.hasTop != 0u && P.preloadTop + w > 0.0) { nc += 1.0; }
191 return nc;
192}
194@compute @workgroup_size(16, 16)
195fn main(@builtin(local_invocation_id) lid: vec3u, @builtin(workgroup_id) wid: vec3u) {
196 let n = P.n;
197 ox = i32(wid.x) * TILE - HALO;
198 oy = i32(wid.y) * TILE - HALO;
199 ihx2 = 1.0 / (P.hx * P.hx); ihy2 = 1.0 / (P.hy * P.hy); ihxy = 1.0 / (P.hx * P.hy);
200 let lt = lid.y * 16u + lid.x;
202 // Phase 1: one entry of the tile with halo per thread.
203 {
204 let gi = ox + i32(lid.x); let gj = oy + i32(lid.y);
205 var ww = 0.0; var uu = 0.0; var vv = 0.0; var z0 = 0.0; var hh = 0.0;
206 if (gi >= 0 && gi < i32(P.nx) && gj >= 0 && gj < i32(P.ny)) {
207 let k = id(gi, gj);
208 let st = nd[k];
209 ww = posIn[2u * n + k]; uu = posIn[k]; vv = posIn[n + k]; z0 = st.x; hh = st.z;
210 }
211 tw[lt] = ww; tu[lt] = uu; tv[lt] = vv; tw0[lt] = z0; th[lt] = hh;
212 if (lt < 64u) {
213 var word = 0u;
214 for (var q = 0u; q < 4u; q++) {
215 let e = lt * 4u + q;
216 let ei = ox + i32(e % u32(TW)); let ej = oy + i32(e / u32(TW));
217 var wq = 0u;
218 if (ei >= 0 && ei < i32(P.nx) && ej >= 0 && ej < i32(P.ny)) { wq = u32(round(4.0 * nd[id(ei, ej)].y)); }
219 word |= wq << (q * 8u);
220 }
221 tn[lt] = word;
222 }
223 }
224 workgroupBarrier();
226 // Phase 2: membrane forces of the cells, bending moments and mixed differences.
227 if (lt < 169u) {
228 let ci = ox + 1 + i32(lt % u32(CT)); let cj = oy + 1 + i32(lt / u32(CT));
229 cellForces(ci, cj, lt);
230 cc[lt] = Cc(ci, cj);
231 }
232 if (lt < 196u) {
233 let mi = ox + 1 + i32(lt % u32(MT)); let mj = oy + 1 + i32(lt / u32(MT));
234 ma[lt] = Ma(mi, mj);
235 mb[lt] = Mb(mi, mj);
236 }
237 workgroupBarrier();
239 // Phase 3: the update of the TILE x TILE nodes (no barriers past this point).
240 if (lt >= 144u) { return; }
241 let i = ox + HALO + i32(lt % u32(TILE)); let j = oy + HALO + i32(lt / u32(TILE));
242 if (i >= i32(P.nx) || j >= i32(P.ny)) { return; }
243 let k = id(i, j);
244 let weight = NW(i, j);
245 if (weight == 0.0) {
246 posOut[k] = 0.0; posOut[n + k] = 0.0; posOut[2u * n + k] = 0.0; posOut[3u * n + k] = 0.0;
247 velOut[k] = 0.0; velOut[n + k] = 0.0; velOut[2u * n + k] = 0.0;
248 return;
249 }
251 // Elastic forces: bending from the discrete Kirchhoff energy, membrane from the cells.
252 let bend = (ma[midx(i - 1, j)] - 2.0 * ma[midx(i, j)] + ma[midx(i + 1, j)]) * ihx2
253 + (mb[midx(i, j - 1)] - 2.0 * mb[midx(i, j)] + mb[midx(i, j + 1)]) * ihy2
254 + 2.0 * (1.0 - P.nu) * (cc[cidx(i, j)] - cc[cidx(i - 1, j)] - cc[cidx(i, j - 1)] + cc[cidx(i - 1, j - 1)]) * ihxy;
255 var F = vec3f(0.0, 0.0, -P.hx * P.hy * bend);
256 F += cellForce(i, j, 0u);
257 F += cellForce(i - 1, j, 1u);
258 F += cellForce(i - 1, j - 1, 2u);
259 F += cellForce(i, j - 1, 3u);
260 let pOld = vec3f(posIn[k], posIn[n + k], posIn[2u * n + k]);
261 let vel = vec3f(velIn[k], velIn[n + k], velIn[2u * n + k]);
262 // The pads: the relaxed spring to ground on the transverse motion
263 // (compression-only with lift-off) and the felt's shear on the in-plane
264 // motion, plus the arms of the felt's chain, whose extensions padQ carry
265 // the memory of the felt (felt.ts, addPadMemory in reference.ts).
266 let kb = nd[k].w;
267 let onPad = P.standOn != 0u && kb > 0.0;
268 var nc = 0.0;
269 if (onPad) {
270 var nPads = 1.0; if (P.hasTop != 0u) { nPads = 2.0; }
271 let ku = nPads * kb * 0.3846153846; // shear: G / E = 1 / (2 (1 + 0.3))
272 F.x -= P.kInf * ku * pOld.x;
273 F.y -= P.kInf * ku * pOld.y;
274 F.z += padForce(kb, pOld.z);
275 nc = padsInContact(pOld.z);
276 var arm = vec3f(0.0);
277 for (var a = 0u; a < P.nArms; a++) {
278 let base = a * n + k;
279 arm += armK(a) * vec3f(padQ[base], padQ[P.nArms * n + base], padQ[2u * P.nArms * n + base]);
280 }
281 F -= vec3f(ku, ku, nc * kb) * arm;
282 }
283 let mass = P.rho * TH(i, j) * P.hx * P.hy * weight;
284 var acc = F * (1.0 / mass);
286 // Losses: frequency-independent on all components, the graph Laplacian of the
287 // transverse velocity (~k^2), and an extra frequency-independent in-plane loss.
288 acc -= 2.0 * SP.sigma0 * vel;
289 acc.z += 2.0 * SP.sigma1 * lapDW(i, j);
290 acc.x -= 2.0 * SP.sigmaIn * vel.x;
291 acc.y -= 2.0 * SP.sigmaIn * vel.y;
293 // Drumstick: Hertz contact between a free mass and the contact patch (the
294 // plate nodes within the tip radius, plus the nearest one). The contact
295 // coordinate is the mass-weighted mean displacement of the patch, and the
296 // force is spread over the patch in proportion to the node masses, so the
297 // pair is conjugate and the contact has a potential energy. Once the
298 // contact has been released for STICK_RELEASE the stick is withdrawn.
299 if ((SP.flags & 1u) != 0u) {
300 let sIn = 8u * SP.parity;
301 if (P.standOn == 0u) {
302 // Free plate: inertia relief. The force of the previous step is applied
303 // to the whole plate as the rigid acceleration field that balances the
304 // patch load and its moments about the center, so the plate gains no net
305 // momentum (the thread takes it up in reality). A plate on a stand takes
306 // the reaction through the pads.
307 let Fprev = stick[sIn + 2u];
308 acc.z += Fprev * (1.0 / P.mass + nodeY(j) * P.patchCY / P.Ix + nodeX(i) * P.patchCX / P.Iy);
309 }
310 if (inContactPatch(i, j)) {
311 let sx = stick[sIn]; let sv = stick[sIn + 1u]; let sMax = stick[sIn + 3u];
312 var released = stick[sIn + 4u];
313 var withdrawn = stick[sIn + 5u];
314 var Fc = 0.0;
315 if (withdrawn == 0.0) {
316 var wc = 0.0;
317 for (var pj = i32(P.patchJ0); pj <= i32(P.patchJ1); pj++) {
318 for (var pi = i32(P.patchI0); pi <= i32(P.patchI1); pi++) {
319 if (inContactPatch(pi, pj)) { wc += nodeMass(pi, pj) * posIn[2u * n + id(pi, pj)]; }
320 }
321 }
322 wc /= P.patchMass;
323 let pen = wc - sx; // the stick sits on the +z side and moves toward -z
324 if (pen > 0.0) { Fc = P.K * pow(pen, P.alpha); released = 0.0; }
325 else if (sMax > 0.0) {
326 released += P.dt;
327 if (released > STICK_RELEASE) { withdrawn = 1.0; }
328 }
329 }
330 acc.z -= Fc / P.patchMass;
331 if (k == P.strikeNode) {
332 let sOut = 8u * (1u - SP.parity);
333 let svn = sv + P.dt * Fc / P.stickMass;
334 stick[sOut] = sx + P.dt * svn;
335 stick[sOut + 1u] = svn;
336 stick[sOut + 2u] = Fc;
337 stick[sOut + 3u] = max(sMax, Fc);
338 stick[sOut + 4u] = released;
339 stick[sOut + 5u] = withdrawn;
340 }
341 }
342 }
344 let vNew = vel + P.dt * acc;
345 if (SP.audioSlot >= 0) { hist[u32(SP.audioSlot) * n + k] = vNew.z; }
346 var pNew = pOld + P.dt * vNew;
347 // The arms of the felt relax over the step while the node moves by dt vNew
348 // (exact for a linear motion within the step); a node that has left every
349 // pad relaxes its transverse arms at once.
350 if (onPad) {
351 let dw = P.dt * vNew;
352 for (var a = 0u; a < P.nArms; a++) {
353 let base = a * n + k;
354 let iu = base; let iv = P.nArms * n + base; let iw = 2u * P.nArms * n + base;
355 padQ[iu] = padQ[iu] * armDecay(a) + dw.x * armGrowth(a);
356 padQ[iv] = padQ[iv] * armDecay(a) + dw.y * armGrowth(a);
357 padQ[iw] = select(0.0, padQ[iw] * armDecay(a) + dw.z * armGrowth(a), nc > 0.0);
358 }
359 }
360 // Compensated update of w: the increment dt * v can be far below the
361 // float32 resolution of w, so the rounding error is carried separately.
362 let y = P.dt * vNew.z + posIn[3u * n + k];
363 let t = pOld.z + y;
364 let lo = (pOld.z - t) + y;
365 pNew.z = t;
366 posOut[k] = pNew.x; posOut[n + k] = pNew.y; posOut[2u * n + k] = pNew.z; posOut[3u * n + k] = lo;
367 velOut[k] = vNew.x; velOut[n + k] = vNew.y; velOut[2u * n + k] = vNew.z;
368}