1"""Displacement-based finite element solution of the cymbal2 shallow-shell model on a Cartesian grid.
3The annulus is approximated by the square cells of a uniform grid whose centers lie inside it (a staircase
4outer rim and a staircase hole). The three displacement fields u, v, w are each expanded in bicubic Hermite
5(Bogner-Fox-Schmit) elements, C1 across cells, with the four nodal unknowns f, f_x, f_y, f_xy per field. The
6weak form of MODEL.md is integrated with 3 x 3 Gauss points per cell: the membrane strains of Marguerre theory
7and the bending strains are evaluated at the points, the stress resultants include the Kelvin-Voigt rates, and
8the internal force is the assembled virtual work. Nothing is linearized and no modes are computed: the mass is
9lumped (diagonal scaling of the consistent element mass), the equations are stepped explicitly with the
10leapfrog scheme at a step below the stability limit of the in-plane waves, and the three sticks are point
11masses integrated with the same step through the Hertz contact law. The Rayleigh integral is accumulated at
12every step from the transverse acceleration at the Gauss points with cubic Lagrange interpolation in time at
13the retarded times, so nothing about the acceleration history is stored. Free edges need no treatment: they
14are the natural boundary conditions of the weak form.
16The whole time step runs on the GPU through PyTorch (one dense matmul from element unknowns to Gauss-point
17derivatives, elementwise constitutive evaluation, one matmul back to element forces, a scatter-add), captured
18in a CUDA graph. The CPU path is the same code and serves for tests.
20Usage:
21 python fem_shell.py <challenge.json> <out-dir> [--spacing 0.008] [--seconds S] [--device cuda] [--fields]
22 [--report 60] [--linear] [--speed V] [--tag NAME] [--no-graph]
23Writes <out>/<tag>.wav, <tag>.meta.json, and while running <tag>.progress.json and <tag>.partial.wav (the samples
24computed so far, the rest zero), rewritten every --report seconds of wall time.
25"""
26from __future__ import annotations
28import json
29import math
30import sys
31import time
32from pathlib import Path
34import numpy as np
35from numpy.polynomial import legendre as leg
36from scipy.io import wavfile
38# ----------------------------------------------------------------------------------------------------------------
39# geometry of a challenge
41class Profile:
42 """Mid-surface height Z(r), its slope, and the thickness h(r) of a cymbal2 challenge."""
44 def __init__(self, p: dict):
45 self.R = p["disk_radius"]
46 self.rh = p.get("mount_radius", p.get("hole_radius", 0.0)) # inner radius of the shell (the clamped mount)
47 self.H = p.get("bow_height", 0.0)
48 self.rb = p.get("bell_radius", 0.0)
49 self.hb = p.get("bell_height", 0.0)
50 self.h0 = p.get("thickness")
51 self.hc, self.he = p.get("thickness_center"), p.get("thickness_rim")
53 def Z(self, r):
54 z = self.H * (1.0 - r * r / self.R**2)
55 if self.rb > 0:
56 z = z + np.where(r < self.rb, 0.5 * self.hb * (1.0 + np.cos(np.pi * np.minimum(r, self.rb) / self.rb)), 0.0)
57 return z
59 def dZ(self, r):
60 d = -2.0 * self.H * r / self.R**2
61 if self.rb > 0:
62 d = d + np.where(r < self.rb, -0.5 * self.hb * (np.pi / self.rb) * np.sin(np.pi * np.minimum(r, self.rb) / self.rb), 0.0)
63 return d
65 def slope(self, x, y):
66 r = np.hypot(x, y)
67 d = self.dZ(r)
68 with np.errstate(invalid="ignore", divide="ignore"):
69 zx = np.where(r > 0, d * x / r, 0.0)
70 zy = np.where(r > 0, d * y / r, 0.0)
71 return zx, zy
73 def h(self, r):
74 if self.h0 is not None:
75 return np.full_like(np.asarray(r, dtype=float), self.h0)
76 return np.where(r <= self.rb, self.hc, self.hc + (self.he - self.hc) * (r - self.rb) / (self.R - self.rb))
79# ----------------------------------------------------------------------------------------------------------------
80# bicubic Hermite element on a square cell of side a
82def hermite_1d(s, a):
83 """Cubic Hermite functions on [0, a] at s = x/a: values and x-derivatives up to second order.
84 Returns arrays (4, n) for H, H', H'' with the ordering [f(0), f'(0), f(a), f'(a)]."""
85 s = np.asarray(s, dtype=float)
86 H = np.stack([1 - 3 * s**2 + 2 * s**3, a * (s - 2 * s**2 + s**3), 3 * s**2 - 2 * s**3, a * (-(s**2) + s**3)])
87 dH = np.stack([-6 * s + 6 * s**2, a * (1 - 4 * s + 3 * s**2), 6 * s - 6 * s**2, a * (-2 * s + 3 * s**2)]) / a
88 ddH = np.stack([-6 + 12 * s, a * (-4 + 6 * s), 6 - 12 * s, a * (-2 + 6 * s)]) / a**2
89 return H, dH, ddH
92def element_basis(sx, sy, a):
93 """Shape functions of the 16 element unknowns at points (sx, sy) in [0,1]^2: returns a dict of arrays (n, 16)
94 for N, Nx, Ny, Nxx, Nyy, Nxy. Local unknown k = 4 * node + d with nodes (0,0), (1,0), (0,1), (1,1) and
95 d in [f, f_x, f_y, f_xy]."""
96 Hx, dHx, ddHx = hermite_1d(sx, a)
97 Hy, dHy, ddHy = hermite_1d(sy, a)
98 out = {k: np.zeros((len(np.atleast_1d(sx)), 16)) for k in ("N", "Nx", "Ny", "Nxx", "Nyy", "Nxy")}
99 for node, (i, j) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
100 for d, (p, q) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]): # (x-derivative flag, y-derivative flag)
101 ix, iy = 2 * i + p, 2 * j + q
102 k = 4 * node + d
103 out["N"][:, k] = Hx[ix] * Hy[iy]
104 out["Nx"][:, k] = dHx[ix] * Hy[iy]
105 out["Ny"][:, k] = Hx[ix] * dHy[iy]
106 out["Nxx"][:, k] = ddHx[ix] * Hy[iy]
107 out["Nyy"][:, k] = Hx[ix] * ddHy[iy]
108 out["Nxy"][:, k] = dHx[ix] * dHy[iy]
109 return out
112# ----------------------------------------------------------------------------------------------------------------
114class ShellFEM:
115 NQ = 3 # Gauss points per direction
116 QUANT = ["w_x", "w_y", "w_xx", "w_yy", "w_xy", "u_x", "u_y", "v_x", "v_y"]
118 def __init__(self, ch: dict, spacing: float, log=print, linear: bool = False, speed: float | None = None, c_air: float | None = None, no_drift: bool = False, inplane_scale: float = 1.0):
119 self.ch = ch
120 self.no_drift = no_drift
121 self.inplane_scale = inplane_scale
122 p = dict(ch["parameters"])
123 if c_air is not None:
124 p["c_air"] = c_air
125 self.p = p
126 self.a = a = spacing
127 self.log = log
128 self.linear = linear
129 self.v0 = p["stick_speed"] if speed is None else speed
130 pr = self.profile = Profile(p)
131 R, rh = pr.R, pr.rh
132 # ---- cells: centers inside the annulus; nodes: their corners; the origin is a node
133 n = int(math.ceil(R / a))
134 ii, jj = np.meshgrid(np.arange(-n, n), np.arange(-n, n), indexing="ij")
135 cx, cy = (ii.ravel() + 0.5) * a, (jj.ravel() + 0.5) * a
136 rc = np.hypot(cx, cy)
137 keep = (rc < R) & (rc > rh)
138 self.cells = np.stack([ii.ravel()[keep], jj.ravel()[keep]], axis=1) # lower-left node index (i, j)
139 self.n_cells = len(self.cells)
140 inner_cells = set(map(tuple, np.stack([ii.ravel()[rc <= rh], jj.ravel()[rc <= rh]], axis=1).tolist())) # cells inside the mount
141 # node map
142 node_ids = {}
143 cell_nodes = np.zeros((self.n_cells, 4), dtype=np.int64)
144 for c, (i, j) in enumerate(self.cells):
145 for k, (di, dj) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
146 key = (int(i + di), int(j + dj))
147 if key not in node_ids:
148 node_ids[key] = len(node_ids)
149 cell_nodes[c, k] = node_ids[key]
150 self.n_nodes = len(node_ids)
151 self.node_xy = np.zeros((self.n_nodes, 2))
152 for (i, j), k in node_ids.items():
153 self.node_xy[k] = (i * a, j * a)
154 self.node_ids = node_ids
155 self.n_dof = 12 * self.n_nodes # blocks w, u, v of 4 n_nodes each
156 # the clamped inner edge: every node that touches a cell inside the mount has all its unknowns fixed at zero
157 # (displacements and the slopes of w), which is the clamp of MODEL.md on the staircase inner boundary
158 fixed_nodes = np.array([k for (i, j), k in node_ids.items() if any((i + di, j + dj) in inner_cells for di in (-1, 0) for dj in (-1, 0))], dtype=np.int64)
159 self.fixed_dofs = np.concatenate([b * 4 * self.n_nodes + 4 * fixed_nodes[:, None] + np.arange(4)[None, :] for b in range(3)]).ravel() if len(fixed_nodes) else np.zeros(0, dtype=np.int64)
160 self.free_dofs = np.setdiff1d(np.arange(self.n_dof), self.fixed_dofs)
161 self.clamped = len(fixed_nodes) > 0
162 if self.clamped and no_drift:
163 no_drift = False # the clamp removes the rigid motions; no reaction or projection is needed
164 self.no_drift = False
165 loc = (4 * cell_nodes[:, :, None] + np.arange(4)[None, None, :]).reshape(self.n_cells, 16)
166 self.cell_dofs = np.concatenate([loc, loc + 4 * self.n_nodes, loc + 8 * self.n_nodes], axis=1) # (C, 48): w | u | v
167 self.cell_dofs_w = loc
168 self.area = self.n_cells * a * a
169 # ---- quadrature and the derivative operator B: (C-independent) (NQ*NQ*9, 48)
170 g, wg = leg.leggauss(self.NQ)
171 s = (g + 1) / 2
172 ws = wg / 2
173 SX, SY = np.meshgrid(s, s, indexing="ij")
174 self.sx, self.sy = SX.ravel(), SY.ravel()
175 self.wq = (ws[:, None] * ws[None, :]).ravel() * a * a # (G,)
176 G = len(self.wq)
177 self.G = G
178 eb = element_basis(self.sx, self.sy, a)
179 B = np.zeros((G, 9, 48))
180 B[:, 0, :16], B[:, 1, :16], B[:, 2, :16], B[:, 3, :16], B[:, 4, :16] = eb["Nx"], eb["Ny"], eb["Nxx"], eb["Nyy"], eb["Nxy"]
181 B[:, 5, 16:32], B[:, 6, 16:32] = eb["Nx"], eb["Ny"]
182 B[:, 7, 32:48], B[:, 8, 32:48] = eb["Nx"], eb["Ny"]
183 self.B = B.reshape(G * 9, 48)
184 self.Nw = eb["N"] # (G, 16) transverse displacement at the Gauss points
185 # Gauss point coordinates, slopes, thickness, stiffnesses
186 gx = self.cells[:, 0][:, None] * a + self.sx[None, :] * a # (C, G)
187 gy = self.cells[:, 1][:, None] * a + self.sy[None, :] * a
188 self.gx, self.gy = gx, gy
189 gr = np.hypot(gx, gy)
190 self.Zx, self.Zy = pr.slope(gx, gy)
191 self.hq = pr.h(gr)
192 E, nu = p["E"], p["nu"]
193 self.Cm = E * self.hq / (1 - nu**2) # membrane stiffness per Gauss point (C, G)
194 self.Db = E * self.hq**3 / (12 * (1 - nu**2))
195 self.mass_total = float(np.sum(p["rho"] * self.hq * self.wq[None, :]))
196 # ---- lumped mass (HRZ diagonal scaling per element, per field)
197 Mdiag = np.zeros(self.n_dof)
198 N2 = eb["N"] ** 2 # (G, 16)
199 for blk in range(3):
200 me = (p["rho"] * self.hq * self.wq[None, :]) @ N2 # (C, 16) diagonal of the consistent element mass
201 mcell = np.sum(p["rho"] * self.hq * self.wq[None, :], axis=1) # (C,)
202 trans = me[:, 0::4].sum(axis=1) # translational entries
203 me = me * (mcell / trans)[:, None]
204 np.add.at(Mdiag, self.cell_dofs[:, 16 * blk : 16 * (blk + 1)], me)
205 # optional scaling of the lumped masses of the in-plane derivative unknowns (--inplane-scale). They are tiny and
206 # set the explicit time step through the in-plane block, and scaling them by 16 halves the step count; but it
207 # also pulls the spurious derivative-dominated in-plane modes down from 40 to 60 kHz into the audible band, where
208 # the nonlinear cascade pumps them and they ring as steady lines (seen at 10.7 kHz on the 8 mm grid). The default
209 # is therefore 1 (no scaling).
210 nw = 4 * self.n_nodes
211 for b in (1, 2):
212 for d in (1, 2, 3):
213 Mdiag[nw * b + d::4][: self.n_nodes] *= inplane_scale
214 self.Mdiag = Mdiag
215 # ---- strikes: footprint vectors phi_j = int phi N_j dA (fine quadrature on the cells they touch)
216 self.strikes = []
217 Rs = p["strike_radius"]
218 c_phi = 0.5 - 2.0 / np.pi**2
219 gf, wf = leg.leggauss(12)
220 sf = (gf + 1) / 2
221 wsf = wf / 2
222 FX, FY = np.meshgrid(sf, sf, indexing="ij")
223 ebf = element_basis(FX.ravel(), FY.ravel(), a)["N"] # (144, 16)
224 wff = (wsf[:, None] * wsf[None, :]).ravel() * a * a
225 for st in ch["strikes"]:
226 xs, ys = st["x"], st["y"]
227 phi = np.zeros(self.n_dof)
228 near = np.flatnonzero(np.hypot(cx[keep] - xs, cy[keep] - ys) < Rs + a)
229 for c in near:
230 px = self.cells[c, 0] * a + FX.ravel() * a
231 py = self.cells[c, 1] * a + FY.ravel() * a
232 rr = np.hypot(px - xs, py - ys)
233 bump = np.where(rr <= Rs, 0.5 * (1 + np.cos(np.pi * np.minimum(rr / Rs, 1.0))), 0.0) / (c_phi * np.pi * Rs**2)
234 np.add.at(phi, self.cell_dofs_w[c], (wff * bump) @ ebf)
235 self.strikes.append(dict(t=st["t"], phi=phi, integral=float(phi[0 : 4 * self.n_nodes : 4].sum())))
236 # ---- rigid reaction: for a unit downward force at each strike point, the rigid acceleration field of the shell
237 # (translation along z plus the two tilts) that balances its net force and moments, as a nodal force vector M g.
238 # With --no-drift this reaction is applied together with the stick force, so the shell gains no momentum or
239 # angular momentum and does not drift (the constraint of zero rigid motion; the elastic motion is untouched).
240 nw = 4 * self.n_nodes
241 mw = self.Mdiag[0:nw:4] # translational masses of the w nodes
242 xn, yn = self.node_xy[:, 0], self.node_xy[:, 1]
243 Mt = mw.sum()
244 xc, yc = (mw * xn).sum() / Mt, (mw * yn).sum() / Mt
245 Ixx, Iyy, Ixy = (mw * (yn - yc) ** 2).sum(), (mw * (xn - xc) ** 2).sum(), (mw * (xn - xc) * (yn - yc)).sum()
246 for st in self.strikes:
247 xs, ys = [s_ for s_ in ch["strikes"] if abs(s_["t"] - st["t"]) < 1e-12][0]["x"], [s_ for s_ in ch["strikes"] if abs(s_["t"] - st["t"]) < 1e-12][0]["y"]
248 # rigid field a = c0 + cx (x - xc) + cy (y - yc) with net force F and moments F xs, F ys about the center (per unit F)
249 Minv = np.linalg.inv(np.array([[Iyy, Ixy], [Ixy, Ixx]]))
250 cx, cy = Minv @ np.array([xs - xc, ys - yc])
251 g = np.zeros(self.n_dof)
252 g[0:nw:4] = 1.0 / Mt + cx * (xn - xc) + cy * (yn - yc)
253 g[1:nw:4] = cx # f_x of a linear field
254 g[2:nw:4] = cy
255 st["reaction"] = self.Mdiag * g # nodal force of the rigid acceleration field for unit stick force (upward)
256 # the six rigid fields (M-orthonormalized): used with --no-drift to project the rigid components out of the state
257 rig = np.zeros((6, self.n_dof))
258 rig[0, 0:nw:4] = 1.0 # w translation
259 rig[1, 0:nw:4], rig[1, 1:nw:4] = xn - xc, 1.0 # tilt about y (w = x)
260 rig[2, 0:nw:4], rig[2, 2:nw:4] = yn - yc, 1.0 # tilt about x (w = y)
261 rig[3, nw:2 * nw:4] = 1.0 # u translation
262 rig[4, 2 * nw:3 * nw:4] = 1.0 # v translation
263 rig[5, nw:2 * nw:4], rig[5, nw + 2:2 * nw:4] = -(yn - yc), -1.0 # in-plane rotation: u = -y, v = x
264 rig[5, 2 * nw:3 * nw:4], rig[5, 2 * nw + 1:3 * nw:4] = xn - xc, 1.0
265 for i in range(6): # Gram-Schmidt in the M inner product
266 for j in range(i):
267 rig[i] -= (rig[j] @ (self.Mdiag * rig[i])) * rig[j]
268 rig[i] /= np.sqrt(rig[i] @ (self.Mdiag * rig[i]))
269 self.rigid = rig
270 # ---- Rayleigh integral: weights and delays per Gauss point
271 L = np.array([p["listener_x"], p["listener_y"], p["listener_z"]])
272 gz = pr.Z(gr)
273 d = np.sqrt((L[0] - gx) ** 2 + (L[1] - gy) ** 2 + (L[2] - gz) ** 2) # (C, G)
274 self.ray_c = (p["rho_air"] / (2 * np.pi)) * self.wq[None, :] / d
275 self.ray_tau = d / p["c_air"]
276 self.log(f" mesh: {self.n_cells} cells, {self.n_nodes} nodes, {self.n_dof} unknowns ({len(self.fixed_dofs)} fixed at the mount), spacing {a * 1e3:g} mm, "
277 f"mass {self.mass_total:.4f} kg, footprint integrals {[round(s['integral'], 4) for s in self.strikes]}")
279 # ------------------------------------------------------------------------------------------------------------
280 def assemble_linear(self):
281 """Sparse stiffness matrix of the linearized shell (the Hessian of the strain energy at rest), from the same
282 Gauss-point operators the time stepping uses; symmetric, with the six rigid fields in its null space."""
283 import scipy.sparse as sp
285 p = self.p
286 C, G, nu = self.n_cells, self.G, p["nu"]
287 B = self.B.reshape(G, 9, 48)
288 L = np.zeros((C, G, 6, 9)) # 9 derivative quantities -> [exx, eyy, gxy, kxx, kyy, 2kxy] at the reference state
289 L[:, :, 0, 0] = self.Zx; L[:, :, 0, 5] = 1.0
290 L[:, :, 1, 1] = self.Zy; L[:, :, 1, 8] = 1.0
291 L[:, :, 2, 0] = self.Zy; L[:, :, 2, 1] = self.Zx; L[:, :, 2, 6] = 1.0; L[:, :, 2, 7] = 1.0
292 L[:, :, 3, 2] = 1.0; L[:, :, 4, 3] = 1.0; L[:, :, 5, 4] = 2.0
293 Q = np.array([[1, nu, 0], [nu, 1, 0], [0, 0, (1 - nu) / 2]])
294 Dm = np.zeros((C, G, 6, 6))
295 Dm[:, :, :3, :3] = self.Cm[:, :, None, None] * Q
296 Dm[:, :, 3:, 3:] = self.Db[:, :, None, None] * Q
297 LB = np.einsum("cgij,gjk->cgik", L, B)
298 Ke = np.einsum("g,cgik,cgij,cgjl->ckl", self.wq, LB, Dm, LB, optimize=True)
299 rows = np.repeat(self.cell_dofs, 48, axis=1).ravel()
300 cols = np.tile(self.cell_dofs, (1, 48)).ravel()
301 K = sp.coo_matrix((Ke.ravel(), (rows, cols)), shape=(self.n_dof, self.n_dof)).tocsr()
302 return 0.5 * (K + K.T)
304 def build_force(self, torch, dev, compile_: bool):
305 """The internal force f(q, v) of the assembled shell (Kelvin-Voigt rates included) as a torch function on the
306 device, and the tensors it shares with the time stepping."""
307 T64 = lambda x: torch.as_tensor(np.asarray(x, dtype=np.float64), device=dev)
308 T64i = lambda x: torch.as_tensor(np.asarray(x, dtype=np.int64), device=dev)
309 p = self.p
310 C, G = self.n_cells, self.G
311 cell_dofs = T64i(self.cell_dofs)
312 BT = T64(self.B.T.copy()) # (48, 81)
313 Bw = T64(self.B * np.repeat(self.wq, 9)[:, None]) # (81, 48) with quadrature weights folded in
314 Zx, Zy, Cm, Db = T64(self.Zx), T64(self.Zy), T64(self.Cm), T64(self.Db)
315 nu, eta = p["nu"], p["eta"]
316 sig1rh = T64(2.0 * p.get("sigma1", 0.0) * p["rho"] * self.hq) # 2 sigma1 rho h at the Gauss points (Laplacian damping of w)
317 lin = 0.0 if self.linear else 1.0
318 n_dof = self.n_dof
320 def constitutive(D):
321 Dq, Dv = D[0], D[1]
322 wx, wy = Dq[..., 0], Dq[..., 1]
323 ax, ay = Zx + lin * wx, Zy + lin * wy # Z_x + w_x
324 exx = Dq[..., 5] + Zx * wx + 0.5 * lin * wx * wx
325 eyy = Dq[..., 8] + Zy * wy + 0.5 * lin * wy * wy
326 exy = 0.5 * (Dq[..., 6] + Dq[..., 7] + Zx * wy + Zy * wx + lin * wx * wy)
327 rxx = Dv[..., 5] + ax * Dv[..., 0]
328 ryy = Dv[..., 8] + ay * Dv[..., 1]
329 rxy = 0.5 * (Dv[..., 6] + Dv[..., 7] + ax * Dv[..., 1] + ay * Dv[..., 0])
330 exx = exx + eta * rxx
331 eyy = eyy + eta * ryy
332 exy = exy + eta * rxy
333 Nxx = Cm * (exx + nu * eyy)
334 Nyy = Cm * (eyy + nu * exx)
335 Nxy = Cm * (1 - nu) * exy
336 kxx = Dq[..., 2] + eta * Dv[..., 2]
337 kyy = Dq[..., 3] + eta * Dv[..., 3]
338 kxy = Dq[..., 4] + eta * Dv[..., 4]
339 Mxx = Db * (kxx + nu * kyy)
340 Myy = Db * (kyy + nu * kxx)
341 Mxy = Db * (1 - nu) * kxy
342 # generalized stresses conjugate to [w_x, w_y, w_xx, w_yy, w_xy, u_x, u_y, v_x, v_y]; the sigma1 term is the
343 # virtual work 2 sigma1 rho h grad(w_t) . grad(delta w) of the Laplacian loss
344 return torch.stack([Nxx * ax + Nxy * ay + sig1rh * Dv[..., 0], Nyy * ay + Nxy * ax + sig1rh * Dv[..., 1], Mxx, Myy, 2 * Mxy, Nxx, Nxy, Nxy, Nyy], dim=-1) # (C, G, 9)
346 constitutive_c = torch.compile(constitutive, dynamic=False) if (compile_ and str(dev).startswith("cuda")) else constitutive
348 def force(q, v):
349 qe = q[cell_dofs] # (C, 48)
350 ve = v[cell_dofs]
351 D = (torch.stack([qe, ve]) @ BT).view(2, C, G, 9)
352 S = constitutive_c(D)
353 fe = S.reshape(C, G * 9) @ Bw # (C, 48)
354 f = torch.zeros(n_dof, dtype=torch.float64, device=dev)
355 f.index_add_(0, cell_dofs.view(-1), fe.view(-1))
356 return f
358 invM_np = 1.0 / self.Mdiag
359 invM_np[self.fixed_dofs] = 0.0 # clamped unknowns never accelerate
360 ops = dict(cell_dofs=cell_dofs, cell_dofs_w=T64i(self.cell_dofs_w), NwT=T64(self.Nw.T.copy()), invM=T64(invM_np), T64=T64, T64i=T64i)
361 return force, ops
363 def field_interpolator(self, xs, ys):
364 """Indices (P, 16) and weights (P, 16) giving w at grid points from the w unknowns; mask of points inside."""
365 XX, YY = np.meshgrid(xs, ys, indexing="xy")
366 px, py = XX.ravel(), YY.ravel()
367 a = self.a
368 ci = np.floor(px / a + 1e-9).astype(int)
369 cj = np.floor(py / a + 1e-9).astype(int)
370 lut = {(int(i), int(j)): c for c, (i, j) in enumerate(self.cells)}
371 idx = np.zeros((len(px), 16), dtype=np.int64)
372 wts = np.zeros((len(px), 16))
373 inside = np.zeros(len(px), dtype=bool)
374 for k in range(len(px)):
375 c = lut.get((ci[k], cj[k]))
376 if c is None:
377 # a point on the outer edge of the last cell row belongs to the cell below/left
378 for di, dj in [(-1, 0), (0, -1), (-1, -1)]:
379 c = lut.get((ci[k] + di, cj[k] + dj))
380 if c is not None:
381 ci[k] += di
382 cj[k] += dj
383 break
384 if c is None:
385 continue
386 inside[k] = True
387 sx, sy = px[k] / a - ci[k], py[k] / a - cj[k]
388 idx[k] = self.cell_dofs_w[c]
389 wts[k] = element_basis([sx], [sy], a)["N"][0]
390 return idx, wts, inside
392 # ------------------------------------------------------------------------------------------------------------
393 def run(self, fs: int, nsamp: int, device: str = "cpu", m: int | None = None, field_grid=None, field_rate=None,
394 report: float = 60.0, out_dir: Path | None = None, tag: str = "run", use_graph: bool = True, scale: float = 1.0, check: bool = False, compile_: bool = True):
395 import torch
397 t_setup = time.perf_counter()
398 dev = torch.device(device)
399 force, ops = self.build_force(torch, dev, compile_)
400 T64, T64i = ops["T64"], ops["T64i"]
401 p = self.p
402 C, G = self.n_cells, self.G
403 cell_dofs, cell_dofs_w, NwT, invM = ops["cell_dofs"], ops["cell_dofs_w"], ops["NwT"], ops["invM"]
404 sigma0 = p["sigma0"]
405 ms, kc = p["stick_mass"], p["contact_stiffness"]
406 v0 = self.v0
408 zero = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
409 torch.manual_seed(0)
410 x = torch.randn(self.n_dof, dtype=torch.float64, device=dev)
411 lam = 0.0
412 saved = self.linear
413 self.linear = True
414 force_lin, _ = self.build_force(torch, dev, False)
415 self.linear = saved
416 for _ in range(60):
417 y = invM * force_lin(x, zero)
418 lam = float(torch.dot(x, y) / torch.dot(x, x))
419 x = y / torch.linalg.norm(y)
420 del force_lin
421 omega_max = math.sqrt(max(lam, 1.0)) * 1.05 # power iteration converges from below; margin
422 dt_stab = 2.0 / omega_max
423 if m is None:
424 m = int(math.ceil(1.0 / (fs * 0.85 * dt_stab)))
425 dt = 1.0 / (fs * m)
426 self.log(f" omega_max ~ {omega_max:.3e} rad/s ({omega_max / 2 / math.pi:.0f} Hz), stable dt < {dt_stab:.3e} s; using m = {m}, dt = {dt:.3e} s")
427 nsteps = nsamp * m
429 # ---- Rayleigh integral bookkeeping
430 dsteps = T64(self.ray_tau / dt).view(-1) # (C*G,) delay in steps
431 rc = T64(self.ray_c).view(-1)
432 out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
433 # ---- strikes
434 nst = len(self.strikes)
435 phi = T64(np.stack([s["phi"] for s in self.strikes])) # (nst, n_dof)
436 phi_load = T64(np.stack([s["phi"] - (s["reaction"] if self.no_drift else 0.0) for s in self.strikes])) # force on the shell per unit F, sign as phi
437 rigM = T64(self.rigid * self.Mdiag[None, :]) # (6, n_dof): <rigid_i, .>_M
438 rig = T64(self.rigid)
440 def project_rigid(): # remove the rigid components of the state (the constraint of MODEL.md, enforced per chunk)
441 for x in (q, v):
442 x.sub_(rig.T @ (rigM @ x))
443 t_on = T64([s["t"] for s in self.strikes])
444 n_on = T64i([int(round(s["t"] / dt)) for s in self.strikes])
445 z = torch.zeros(nst, dtype=torch.float64, device=dev)
446 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
447 active = torch.zeros(nst, dtype=torch.float64, device=dev)
448 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
449 imp = torch.zeros(nst, dtype=torch.float64, device=dev) # contact impulse per stick
450 dur = torch.zeros(nst, dtype=torch.float64, device=dev) # time in contact
451 fmax = torch.zeros(nst, dtype=torch.float64, device=dev)
452 out_direct = torch.zeros(nsamp, dtype=torch.float64, device=dev) if check else None
453 # ---- state
454 q = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
455 v = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
456 acc = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
457 nstep = torch.zeros((), dtype=torch.int64, device=dev)
458 lag_c = 1.0 / 6.0
460 assert m >= 4, "the retarded-time interpolation assumes at least 4 steps per sample"
462 def rayleigh(n, ag):
463 nf = n.to(torch.float64)
464 kk = torch.ceil((nf - 2 + dsteps) / m)
465 pos = kk * m - dsteps
466 fl = torch.floor(pos)
467 j = nf - fl + 1
468 fr = pos - fl
469 wgt = torch.where(j == 0, -fr * (fr - 1) * (fr - 2) * lag_c,
470 torch.where(j == 1, (fr + 1) * (fr - 1) * (fr - 2) * 0.5,
471 torch.where(j == 2, -(fr + 1) * fr * (fr - 2) * 0.5,
472 torch.where(j == 3, (fr + 1) * fr * (fr - 1) * lag_c, torch.zeros_like(fr)))))
473 valid = (j >= 0) & (j <= 3) & (kk >= 0) & (kk < nsamp)
474 idx = torch.where(valid, kk, torch.zeros_like(kk)).to(torch.int64)
475 return idx, torch.where(valid, rc * wgt * ag, torch.zeros_like(ag))
477 ray_c = torch.compile(rayleigh, dynamic=False) if (compile_ and device.startswith("cuda")) else rayleigh
479 def step():
480 n = nstep
481 # sticks: start the ones whose time has come (at the current surface, with relative speed v0)
482 wbar = phi @ q
483 wbar_dot = phi @ v
484 starting = (n == n_on).to(torch.float64)
485 z.copy_(z * (1 - starting) + wbar * starting)
486 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
487 active.copy_(torch.maximum(active, starting))
488 pen = torch.clamp(wbar - z, min=0.0)
489 F = active * kc * pen * torch.sqrt(pen)
490 Fst.copy_(F)
491 imp.add_(dt * F)
492 dur.add_(dt * (F > 0).to(torch.float64))
493 fmax.copy_(torch.maximum(fmax, F))
494 f = force(q, v)
495 f = f + F @ phi_load # the shell is pushed down (and, with --no-drift, held rigidly): M a = -f
496 acc.copy_(-f * invM - 2 * sigma0 * v)
497 zd.add_(dt * F / ms)
498 z.add_(dt * zd)
499 v.add_(dt * acc)
500 q.add_(dt * v)
501 # Rayleigh integral: scatter the Gauss-point accelerations onto the retarded output samples
502 ag = (acc[cell_dofs_w] @ NwT).reshape(-1) # (C*G,)
503 idx, val = ray_c(n, ag)
504 out.index_add_(0, idx, val)
505 if out_direct is not None: # the same integral without retardation, for checking the scatter above
506 ks = torch.div(n, m, rounding_mode="floor")
507 on_sample = ((n % m) == 0) & (ks < nsamp)
508 out_direct.index_add_(0, torch.where(on_sample, ks, torch.zeros_like(ks)).view(1), (torch.sum(rc * ag) * on_sample.to(torch.float64)).view(1))
509 nstep.add_(1)
511 # fields
512 fld = None
513 if field_grid is not None:
514 fx, fy, ft = field_grid
515 idx, wts, inside = self.field_interpolator(fx, fy)
516 fld = dict(idx=T64i(idx), wts=T64(wts), inside=inside, every=int(round(fs * m / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])
518 def take_frame():
519 fr = (q[fld["idx"]] * fld["wts"]).sum(1).cpu().numpy().astype(np.float32)
520 fr[~fld["inside"]] = np.nan
521 fld["frames"].append(fr)
523 def reset_state():
524 for tsr in (q, v, acc, out, z, zd, active, Fst, imp, dur, fmax, nstep):
525 tsr.zero_()
526 if out_direct is not None:
527 out_direct.zero_()
529 graph = None
530 if use_graph and device.startswith("cuda"):
531 s = torch.cuda.Stream()
532 s.wait_stream(torch.cuda.current_stream())
533 with torch.cuda.stream(s):
534 for _ in range(2):
535 step()
536 torch.cuda.current_stream().wait_stream(s)
537 # undo the warm-up steps
538 reset_state()
539 graph = torch.cuda.CUDAGraph()
540 with torch.cuda.graph(graph):
541 step()
542 reset_state()
543 self.log(f" setup {time.perf_counter() - t_setup:.1f}s; stepping {nsteps} steps ({nsamp / fs:.2f} s) on {device}{' with a CUDA graph' if graph else ''}")
545 t_start = time.perf_counter()
546 last_report = t_start
547 chunk = fld["every"] if fld is not None else 2000
548 if fld is not None:
549 take_frame()
550 stick_stats = lambda: dict(impulse_Ns=[round(float(x), 5) for x in imp.cpu().numpy()], contact_ms=[round(1e3 * float(x), 4) for x in dur.cpu().numpy()],
551 peak_force_N=[round(float(x), 2) for x in fmax.cpu().numpy()], rebound_mps=[round(float(x), 4) for x in zd.cpu().numpy()])
552 contact_log = []
553 Fprev = np.zeros(nst)
554 n_done = 0
555 while n_done < nsteps:
556 k = min(chunk, nsteps - n_done)
557 for _ in range(k):
558 if graph is not None:
559 graph.replay()
560 else:
561 step()
562 n_done += k
563 if self.no_drift:
564 project_rigid()
565 if fld is not None and len(fld["frames"]) < fld["nt"]:
566 take_frame()
567 Fc = Fst.cpu().numpy()
568 for i in range(nst):
569 if (Fc[i] > 0) != (Fprev[i] > 0):
570 contact_log.append((round(n_done * dt, 6), i + 1, "on" if Fc[i] > 0 else "off"))
571 Fprev = Fc
572 now = time.perf_counter()
573 if now - last_report >= report or n_done == nsteps:
574 last_report = now
575 el = now - t_start
576 qn = q.cpu().numpy()
577 if not np.all(np.isfinite(qn)):
578 raise RuntimeError(f"solution diverged at step {n_done} (t = {n_done * dt:.4f} s)")
579 w = qn[0 : 4 * self.n_nodes : 4]
580 sig = out.cpu().numpy()
581 prog = dict(tag=tag, step=n_done, steps=nsteps, t=round(n_done * dt, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
582 eta_s=round(el * (nsteps - n_done) / max(n_done, 1), 1), steps_per_s=round(n_done / el, 1),
583 max_abs_w_m=float(np.abs(w).max()), rms_w_m=float(np.sqrt(np.mean(w**2))),
584 stick_z_m=[round(float(x), 6) for x in z.cpu().numpy()], stick_v_mps=[round(float(x), 4) for x in zd.cpu().numpy()],
585 contact_force_N=[round(float(x), 2) for x in Fc], contacts=contact_log[-40:], **stick_stats(),
586 peak_pa=float(np.abs(sig).max()), m=m, dt=dt, device=device)
587 self.log(f" t={prog['t']:.3f}s step {n_done}/{nsteps} elapsed {el:.0f}s eta {prog['eta_s']:.0f}s ({prog['steps_per_s']:.0f} steps/s) "
588 f"max|w|={prog['max_abs_w_m']:.3e} peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
589 if out_dir is not None:
590 with open(out_dir / f"{tag}.progress.json", "w") as fh:
591 json.dump(prog, fh, indent=1)
592 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
593 elapsed = time.perf_counter() - t_start
594 sig = out.cpu().numpy()
595 info = dict(m=m, dt=dt, omega_max=omega_max, steps=nsteps, elapsed=elapsed, contacts=contact_log, sticks=stick_stats(),
596 out_direct=out_direct.cpu().numpy() if out_direct is not None else None)
597 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
598 return sig, fields, info
601# ----------------------------------------------------------------------------------------------------------------
603def hardware_string(device: str) -> str:
604 cpu = "unknown CPU"
605 try:
606 for line in open("/proc/cpuinfo"):
607 if line.startswith("model name"):
608 cpu = line.split(":", 1)[1].strip()
609 break
610 except OSError:
611 pass
612 import os
614 s = f"{cpu} ({os.cpu_count()} threads)"
615 if device.startswith("cuda"):
616 import torch
618 s += f"; time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
619 return s
622def main():
623 ch = json.load(open(sys.argv[1]))
624 out_dir = Path(sys.argv[2])
625 args = sys.argv[3:]
627 def take(flag, conv, default=None):
628 nonlocal args
629 if flag in args:
630 i = args.index(flag)
631 val = conv(args[i + 1])
632 args = args[:i] + args[i + 2 :]
633 return val
634 return default
636 spacing = take("--spacing", float, 0.008)
637 seconds = take("--seconds", float)
638 device = take("--device", str, "cpu")
639 report = take("--report", float, 60.0)
640 speed = take("--speed", float)
641 m_over = take("--m", int)
642 tag = take("--tag", str)
643 c_air = take("--c-air", float)
644 check = "--check" in args
645 no_drift = "--free" not in args # the rigid constraint of MODEL.md is the default; --free leaves the shell free (diagnostics only)
646 if "--no-drift" in args:
647 args.remove("--no-drift")
648 inplane_scale = take("--inplane-scale", float, 1.0)
649 compile_ = "--no-compile" not in args
650 linear = "--linear" in args
651 fields = "--fields" in args
652 use_graph = "--no-graph" not in args
653 out = ch["output"]
654 fs = out["sample_rate_hz"]
655 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
656 tag = tag or f"fem-g{spacing * 1e3:g}"
657 log = lambda s: print(s, file=sys.stderr, flush=True)
658 out_dir.mkdir(parents=True, exist_ok=True)
659 t0 = time.perf_counter()
660 log(f"{tag}: {ch['challenge_id']}, spacing {spacing * 1e3:g} mm, {nsamp / fs:.2f} s, device {device}{' (linear)' if linear else ''}")
661 model = ShellFEM(ch, spacing, log=log, linear=linear, speed=speed, c_air=c_air, no_drift=no_drift, inplane_scale=inplane_scale)
662 scale = out.get("signal_scale") or 1.0
663 grid = None
664 if fields:
665 fg = ch["fields"]
666 nt, ny, nx = fg["shape"]
667 R = ch["parameters"]["disk_radius"]
668 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
669 sig, fld, info = model.run(fs, nsamp, device=device, m=m_over, report=report, out_dir=out_dir, tag=tag, use_graph=use_graph, scale=scale, check=check, compile_=compile_,
670 field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None)
671 total = time.perf_counter() - t0
672 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
673 meta = {
674 "solution_id": tag,
675 "output": f"{out_dir.name}/{tag}.wav",
676 "method": "bicubic Hermite (Bogner-Fox-Schmit) finite elements for u, v, w on a Cartesian grid masked to the annulus, lumped mass, explicit leapfrog, Hertz stick, Rayleigh integral accumulated at the retarded times",
677 "parameters": {"grid_spacing_m": spacing, "cells": model.n_cells, "n_dof": model.n_dof, "gauss_points_per_cell": model.G,
678 "oversampling": info["m"], "dt_s": info["dt"], "omega_max_rad_s": info["omega_max"], "inplane_derivative_mass_scale": inplane_scale},
679 "timing_seconds": round(info["elapsed"], 3),
680 "hardware": hardware_string(device),
681 "peak_si": float(np.max(np.abs(sig))),
682 "sticks": info["sticks"],
683 **({"stick_speed_override": speed} if speed is not None else {}),
684 **({"linear": True} if linear else {}),
685 **({"no_drift": True} if no_drift else {}),
686 }
687 if fld is not None:
688 import h5py
690 fg = ch["fields"]
691 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
692 fh.attrs["challenge_id"] = ch["challenge_id"]
693 fh.attrs["solution_id"] = tag
694 fh.attrs["quantity"] = fg["quantity"]
695 fh.attrs["units"] = fg["units"]
696 fh.create_dataset("x", data=grid[0])
697 fh.create_dataset("y", data=grid[1])
698 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
699 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
700 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
701 with open(out_dir / f"{tag}.meta.json", "w") as fh:
702 json.dump(meta, fh, indent=1)
703 if info["out_direct"] is not None:
704 d = info["out_direct"]
705 log(f" check: max|out - direct| = {np.abs(sig - d).max():.3e} Pa against max|direct| = {np.abs(d).max():.3e} Pa (equal only when c_air is huge)")
706 np.save(out_dir / f"{tag}.direct.npy", d)
707 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
710if __name__ == "__main__":
711 main()