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["hole_radius"]
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 # node map
141 node_ids = {}
142 cell_nodes = np.zeros((self.n_cells, 4), dtype=np.int64)
143 for c, (i, j) in enumerate(self.cells):
144 for k, (di, dj) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
145 key = (int(i + di), int(j + dj))
146 if key not in node_ids:
147 node_ids[key] = len(node_ids)
148 cell_nodes[c, k] = node_ids[key]
149 self.n_nodes = len(node_ids)
150 self.node_xy = np.zeros((self.n_nodes, 2))
151 for (i, j), k in node_ids.items():
152 self.node_xy[k] = (i * a, j * a)
153 self.node_ids = node_ids
154 self.n_dof = 12 * self.n_nodes # blocks w, u, v of 4 n_nodes each
155 loc = (4 * cell_nodes[:, :, None] + np.arange(4)[None, None, :]).reshape(self.n_cells, 16)
156 self.cell_dofs = np.concatenate([loc, loc + 4 * self.n_nodes, loc + 8 * self.n_nodes], axis=1) # (C, 48): w | u | v
157 self.cell_dofs_w = loc
158 self.area = self.n_cells * a * a
159 # ---- quadrature and the derivative operator B: (C-independent) (NQ*NQ*9, 48)
160 g, wg = leg.leggauss(self.NQ)
161 s = (g + 1) / 2
162 ws = wg / 2
163 SX, SY = np.meshgrid(s, s, indexing="ij")
164 self.sx, self.sy = SX.ravel(), SY.ravel()
165 self.wq = (ws[:, None] * ws[None, :]).ravel() * a * a # (G,)
166 G = len(self.wq)
167 self.G = G
168 eb = element_basis(self.sx, self.sy, a)
169 B = np.zeros((G, 9, 48))
170 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"]
171 B[:, 5, 16:32], B[:, 6, 16:32] = eb["Nx"], eb["Ny"]
172 B[:, 7, 32:48], B[:, 8, 32:48] = eb["Nx"], eb["Ny"]
173 self.B = B.reshape(G * 9, 48)
174 self.Nw = eb["N"] # (G, 16) transverse displacement at the Gauss points
175 # Gauss point coordinates, slopes, thickness, stiffnesses
176 gx = self.cells[:, 0][:, None] * a + self.sx[None, :] * a # (C, G)
177 gy = self.cells[:, 1][:, None] * a + self.sy[None, :] * a
178 self.gx, self.gy = gx, gy
179 gr = np.hypot(gx, gy)
180 self.Zx, self.Zy = pr.slope(gx, gy)
181 self.hq = pr.h(gr)
182 E, nu = p["E"], p["nu"]
183 self.Cm = E * self.hq / (1 - nu**2) # membrane stiffness per Gauss point (C, G)
184 self.Db = E * self.hq**3 / (12 * (1 - nu**2))
185 self.mass_total = float(np.sum(p["rho"] * self.hq * self.wq[None, :]))
186 # ---- lumped mass (HRZ diagonal scaling per element, per field)
187 Mdiag = np.zeros(self.n_dof)
188 N2 = eb["N"] ** 2 # (G, 16)
189 for blk in range(3):
190 me = (p["rho"] * self.hq * self.wq[None, :]) @ N2 # (C, 16) diagonal of the consistent element mass
191 mcell = np.sum(p["rho"] * self.hq * self.wq[None, :], axis=1) # (C,)
192 trans = me[:, 0::4].sum(axis=1) # translational entries
193 me = me * (mcell / trans)[:, None]
194 np.add.at(Mdiag, self.cell_dofs[:, 16 * blk : 16 * (blk + 1)], me)
195 # optional scaling of the lumped masses of the in-plane derivative unknowns (--inplane-scale). They are tiny and
196 # set the explicit time step through the in-plane block, and scaling them by 16 halves the step count; but it
197 # also pulls the spurious derivative-dominated in-plane modes down from 40 to 60 kHz into the audible band, where
198 # the nonlinear cascade pumps them and they ring as steady lines (seen at 10.7 kHz on the 8 mm grid). The default
199 # is therefore 1 (no scaling).
200 nw = 4 * self.n_nodes
201 for b in (1, 2):
202 for d in (1, 2, 3):
203 Mdiag[nw * b + d::4][: self.n_nodes] *= inplane_scale
204 self.Mdiag = Mdiag
205 # ---- strikes: footprint vectors phi_j = int phi N_j dA (fine quadrature on the cells they touch)
206 self.strikes = []
207 Rs = p["strike_radius"]
208 c_phi = 0.5 - 2.0 / np.pi**2
209 gf, wf = leg.leggauss(12)
210 sf = (gf + 1) / 2
211 wsf = wf / 2
212 FX, FY = np.meshgrid(sf, sf, indexing="ij")
213 ebf = element_basis(FX.ravel(), FY.ravel(), a)["N"] # (144, 16)
214 wff = (wsf[:, None] * wsf[None, :]).ravel() * a * a
215 for st in ch["strikes"]:
216 xs, ys = st["x"], st["y"]
217 phi = np.zeros(self.n_dof)
218 near = np.flatnonzero(np.hypot(cx[keep] - xs, cy[keep] - ys) < Rs + a)
219 for c in near:
220 px = self.cells[c, 0] * a + FX.ravel() * a
221 py = self.cells[c, 1] * a + FY.ravel() * a
222 rr = np.hypot(px - xs, py - ys)
223 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)
224 np.add.at(phi, self.cell_dofs_w[c], (wff * bump) @ ebf)
225 self.strikes.append(dict(t=st["t"], phi=phi, integral=float(phi[0 : 4 * self.n_nodes : 4].sum())))
226 # ---- rigid reaction: for a unit downward force at each strike point, the rigid acceleration field of the shell
227 # (translation along z plus the two tilts) that balances its net force and moments, as a nodal force vector M g.
228 # With --no-drift this reaction is applied together with the stick force, so the shell gains no momentum or
229 # angular momentum and does not drift (the constraint of zero rigid motion; the elastic motion is untouched).
230 nw = 4 * self.n_nodes
231 mw = self.Mdiag[0:nw:4] # translational masses of the w nodes
232 xn, yn = self.node_xy[:, 0], self.node_xy[:, 1]
233 Mt = mw.sum()
234 xc, yc = (mw * xn).sum() / Mt, (mw * yn).sum() / Mt
235 Ixx, Iyy, Ixy = (mw * (yn - yc) ** 2).sum(), (mw * (xn - xc) ** 2).sum(), (mw * (xn - xc) * (yn - yc)).sum()
236 for st in self.strikes:
237 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"]
238 # 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)
239 Minv = np.linalg.inv(np.array([[Iyy, Ixy], [Ixy, Ixx]]))
240 cx, cy = Minv @ np.array([xs - xc, ys - yc])
241 g = np.zeros(self.n_dof)
242 g[0:nw:4] = 1.0 / Mt + cx * (xn - xc) + cy * (yn - yc)
243 g[1:nw:4] = cx # f_x of a linear field
244 g[2:nw:4] = cy
245 st["reaction"] = self.Mdiag * g # nodal force of the rigid acceleration field for unit stick force (upward)
246 # the six rigid fields (M-orthonormalized): used with --no-drift to project the rigid components out of the state
247 rig = np.zeros((6, self.n_dof))
248 rig[0, 0:nw:4] = 1.0 # w translation
249 rig[1, 0:nw:4], rig[1, 1:nw:4] = xn - xc, 1.0 # tilt about y (w = x)
250 rig[2, 0:nw:4], rig[2, 2:nw:4] = yn - yc, 1.0 # tilt about x (w = y)
251 rig[3, nw:2 * nw:4] = 1.0 # u translation
252 rig[4, 2 * nw:3 * nw:4] = 1.0 # v translation
253 rig[5, nw:2 * nw:4], rig[5, nw + 2:2 * nw:4] = -(yn - yc), -1.0 # in-plane rotation: u = -y, v = x
254 rig[5, 2 * nw:3 * nw:4], rig[5, 2 * nw + 1:3 * nw:4] = xn - xc, 1.0
255 for i in range(6): # Gram-Schmidt in the M inner product
256 for j in range(i):
257 rig[i] -= (rig[j] @ (self.Mdiag * rig[i])) * rig[j]
258 rig[i] /= np.sqrt(rig[i] @ (self.Mdiag * rig[i]))
259 self.rigid = rig
260 # ---- Rayleigh integral: weights and delays per Gauss point
261 L = np.array([p["listener_x"], p["listener_y"], p["listener_z"]])
262 gz = pr.Z(gr)
263 d = np.sqrt((L[0] - gx) ** 2 + (L[1] - gy) ** 2 + (L[2] - gz) ** 2) # (C, G)
264 self.ray_c = (p["rho_air"] / (2 * np.pi)) * self.wq[None, :] / d
265 self.ray_tau = d / p["c_air"]
266 self.log(f" mesh: {self.n_cells} cells, {self.n_nodes} nodes, {self.n_dof} unknowns, spacing {a * 1e3:g} mm, mass {self.mass_total:.4f} kg, "
267 f"footprint integrals {[round(s['integral'], 4) for s in self.strikes]}")
269 # ------------------------------------------------------------------------------------------------------------
270 def assemble_linear(self):
271 """Sparse stiffness matrix of the linearized shell (the Hessian of the strain energy at rest), from the same
272 Gauss-point operators the time stepping uses; symmetric, with the six rigid fields in its null space."""
273 import scipy.sparse as sp
275 p = self.p
276 C, G, nu = self.n_cells, self.G, p["nu"]
277 B = self.B.reshape(G, 9, 48)
278 L = np.zeros((C, G, 6, 9)) # 9 derivative quantities -> [exx, eyy, gxy, kxx, kyy, 2kxy] at the reference state
279 L[:, :, 0, 0] = self.Zx; L[:, :, 0, 5] = 1.0
280 L[:, :, 1, 1] = self.Zy; L[:, :, 1, 8] = 1.0
281 L[:, :, 2, 0] = self.Zy; L[:, :, 2, 1] = self.Zx; L[:, :, 2, 6] = 1.0; L[:, :, 2, 7] = 1.0
282 L[:, :, 3, 2] = 1.0; L[:, :, 4, 3] = 1.0; L[:, :, 5, 4] = 2.0
283 Q = np.array([[1, nu, 0], [nu, 1, 0], [0, 0, (1 - nu) / 2]])
284 Dm = np.zeros((C, G, 6, 6))
285 Dm[:, :, :3, :3] = self.Cm[:, :, None, None] * Q
286 Dm[:, :, 3:, 3:] = self.Db[:, :, None, None] * Q
287 LB = np.einsum("cgij,gjk->cgik", L, B)
288 Ke = np.einsum("g,cgik,cgij,cgjl->ckl", self.wq, LB, Dm, LB, optimize=True)
289 rows = np.repeat(self.cell_dofs, 48, axis=1).ravel()
290 cols = np.tile(self.cell_dofs, (1, 48)).ravel()
291 K = sp.coo_matrix((Ke.ravel(), (rows, cols)), shape=(self.n_dof, self.n_dof)).tocsr()
292 return 0.5 * (K + K.T)
294 def build_force(self, torch, dev, compile_: bool):
295 """The internal force f(q, v) of the assembled shell (Kelvin-Voigt rates included) as a torch function on the
296 device, and the tensors it shares with the time stepping."""
297 T64 = lambda x: torch.as_tensor(np.asarray(x, dtype=np.float64), device=dev)
298 T64i = lambda x: torch.as_tensor(np.asarray(x, dtype=np.int64), device=dev)
299 p = self.p
300 C, G = self.n_cells, self.G
301 cell_dofs = T64i(self.cell_dofs)
302 BT = T64(self.B.T.copy()) # (48, 81)
303 Bw = T64(self.B * np.repeat(self.wq, 9)[:, None]) # (81, 48) with quadrature weights folded in
304 Zx, Zy, Cm, Db = T64(self.Zx), T64(self.Zy), T64(self.Cm), T64(self.Db)
305 nu, eta = p["nu"], p["eta"]
306 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)
307 lin = 0.0 if self.linear else 1.0
308 n_dof = self.n_dof
310 def constitutive(D):
311 Dq, Dv = D[0], D[1]
312 wx, wy = Dq[..., 0], Dq[..., 1]
313 ax, ay = Zx + lin * wx, Zy + lin * wy # Z_x + w_x
314 exx = Dq[..., 5] + Zx * wx + 0.5 * lin * wx * wx
315 eyy = Dq[..., 8] + Zy * wy + 0.5 * lin * wy * wy
316 exy = 0.5 * (Dq[..., 6] + Dq[..., 7] + Zx * wy + Zy * wx + lin * wx * wy)
317 rxx = Dv[..., 5] + ax * Dv[..., 0]
318 ryy = Dv[..., 8] + ay * Dv[..., 1]
319 rxy = 0.5 * (Dv[..., 6] + Dv[..., 7] + ax * Dv[..., 1] + ay * Dv[..., 0])
320 exx = exx + eta * rxx
321 eyy = eyy + eta * ryy
322 exy = exy + eta * rxy
323 Nxx = Cm * (exx + nu * eyy)
324 Nyy = Cm * (eyy + nu * exx)
325 Nxy = Cm * (1 - nu) * exy
326 kxx = Dq[..., 2] + eta * Dv[..., 2]
327 kyy = Dq[..., 3] + eta * Dv[..., 3]
328 kxy = Dq[..., 4] + eta * Dv[..., 4]
329 Mxx = Db * (kxx + nu * kyy)
330 Myy = Db * (kyy + nu * kxx)
331 Mxy = Db * (1 - nu) * kxy
332 # 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
333 # virtual work 2 sigma1 rho h grad(w_t) . grad(delta w) of the Laplacian loss
334 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)
336 constitutive_c = torch.compile(constitutive, dynamic=False) if (compile_ and str(dev).startswith("cuda")) else constitutive
338 def force(q, v):
339 qe = q[cell_dofs] # (C, 48)
340 ve = v[cell_dofs]
341 D = (torch.stack([qe, ve]) @ BT).view(2, C, G, 9)
342 S = constitutive_c(D)
343 fe = S.reshape(C, G * 9) @ Bw # (C, 48)
344 f = torch.zeros(n_dof, dtype=torch.float64, device=dev)
345 f.index_add_(0, cell_dofs.view(-1), fe.view(-1))
346 return f
348 ops = dict(cell_dofs=cell_dofs, cell_dofs_w=T64i(self.cell_dofs_w), NwT=T64(self.Nw.T.copy()), invM=T64(1.0 / self.Mdiag), T64=T64, T64i=T64i)
349 return force, ops
351 def field_interpolator(self, xs, ys):
352 """Indices (P, 16) and weights (P, 16) giving w at grid points from the w unknowns; mask of points inside."""
353 XX, YY = np.meshgrid(xs, ys, indexing="xy")
354 px, py = XX.ravel(), YY.ravel()
355 a = self.a
356 ci = np.floor(px / a + 1e-9).astype(int)
357 cj = np.floor(py / a + 1e-9).astype(int)
358 lut = {(int(i), int(j)): c for c, (i, j) in enumerate(self.cells)}
359 idx = np.zeros((len(px), 16), dtype=np.int64)
360 wts = np.zeros((len(px), 16))
361 inside = np.zeros(len(px), dtype=bool)
362 for k in range(len(px)):
363 c = lut.get((ci[k], cj[k]))
364 if c is None:
365 # a point on the outer edge of the last cell row belongs to the cell below/left
366 for di, dj in [(-1, 0), (0, -1), (-1, -1)]:
367 c = lut.get((ci[k] + di, cj[k] + dj))
368 if c is not None:
369 ci[k] += di
370 cj[k] += dj
371 break
372 if c is None:
373 continue
374 inside[k] = True
375 sx, sy = px[k] / a - ci[k], py[k] / a - cj[k]
376 idx[k] = self.cell_dofs_w[c]
377 wts[k] = element_basis([sx], [sy], a)["N"][0]
378 return idx, wts, inside
380 # ------------------------------------------------------------------------------------------------------------
381 def run(self, fs: int, nsamp: int, device: str = "cpu", m: int | None = None, field_grid=None, field_rate=None,
382 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):
383 import torch
385 t_setup = time.perf_counter()
386 dev = torch.device(device)
387 force, ops = self.build_force(torch, dev, compile_)
388 T64, T64i = ops["T64"], ops["T64i"]
389 p = self.p
390 C, G = self.n_cells, self.G
391 cell_dofs, cell_dofs_w, NwT, invM = ops["cell_dofs"], ops["cell_dofs_w"], ops["NwT"], ops["invM"]
392 sigma0 = p["sigma0"]
393 ms, kc = p["stick_mass"], p["contact_stiffness"]
394 v0 = self.v0
396 zero = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
397 torch.manual_seed(0)
398 x = torch.randn(self.n_dof, dtype=torch.float64, device=dev)
399 lam = 0.0
400 saved = self.linear
401 self.linear = True
402 force_lin, _ = self.build_force(torch, dev, False)
403 self.linear = saved
404 for _ in range(60):
405 y = invM * force_lin(x, zero)
406 lam = float(torch.dot(x, y) / torch.dot(x, x))
407 x = y / torch.linalg.norm(y)
408 del force_lin
409 omega_max = math.sqrt(max(lam, 1.0)) * 1.05 # power iteration converges from below; margin
410 dt_stab = 2.0 / omega_max
411 if m is None:
412 m = int(math.ceil(1.0 / (fs * 0.85 * dt_stab)))
413 dt = 1.0 / (fs * m)
414 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")
415 nsteps = nsamp * m
417 # ---- Rayleigh integral bookkeeping
418 dsteps = T64(self.ray_tau / dt).view(-1) # (C*G,) delay in steps
419 rc = T64(self.ray_c).view(-1)
420 out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
421 # ---- strikes
422 nst = len(self.strikes)
423 phi = T64(np.stack([s["phi"] for s in self.strikes])) # (nst, n_dof)
424 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
425 rigM = T64(self.rigid * self.Mdiag[None, :]) # (6, n_dof): <rigid_i, .>_M
426 rig = T64(self.rigid)
428 def project_rigid(): # remove the rigid components of the state (the constraint of MODEL.md, enforced per chunk)
429 for x in (q, v):
430 x.sub_(rig.T @ (rigM @ x))
431 t_on = T64([s["t"] for s in self.strikes])
432 n_on = T64i([int(round(s["t"] / dt)) for s in self.strikes])
433 z = torch.zeros(nst, dtype=torch.float64, device=dev)
434 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
435 active = torch.zeros(nst, dtype=torch.float64, device=dev)
436 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
437 imp = torch.zeros(nst, dtype=torch.float64, device=dev) # contact impulse per stick
438 dur = torch.zeros(nst, dtype=torch.float64, device=dev) # time in contact
439 fmax = torch.zeros(nst, dtype=torch.float64, device=dev)
440 out_direct = torch.zeros(nsamp, dtype=torch.float64, device=dev) if check else None
441 # ---- state
442 q = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
443 v = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
444 acc = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
445 nstep = torch.zeros((), dtype=torch.int64, device=dev)
446 lag_c = 1.0 / 6.0
448 assert m >= 4, "the retarded-time interpolation assumes at least 4 steps per sample"
450 def rayleigh(n, ag):
451 nf = n.to(torch.float64)
452 kk = torch.ceil((nf - 2 + dsteps) / m)
453 pos = kk * m - dsteps
454 fl = torch.floor(pos)
455 j = nf - fl + 1
456 fr = pos - fl
457 wgt = torch.where(j == 0, -fr * (fr - 1) * (fr - 2) * lag_c,
458 torch.where(j == 1, (fr + 1) * (fr - 1) * (fr - 2) * 0.5,
459 torch.where(j == 2, -(fr + 1) * fr * (fr - 2) * 0.5,
460 torch.where(j == 3, (fr + 1) * fr * (fr - 1) * lag_c, torch.zeros_like(fr)))))
461 valid = (j >= 0) & (j <= 3) & (kk >= 0) & (kk < nsamp)
462 idx = torch.where(valid, kk, torch.zeros_like(kk)).to(torch.int64)
463 return idx, torch.where(valid, rc * wgt * ag, torch.zeros_like(ag))
465 ray_c = torch.compile(rayleigh, dynamic=False) if (compile_ and device.startswith("cuda")) else rayleigh
467 def step():
468 n = nstep
469 # sticks: start the ones whose time has come (at the current surface, with relative speed v0)
470 wbar = phi @ q
471 wbar_dot = phi @ v
472 starting = (n == n_on).to(torch.float64)
473 z.copy_(z * (1 - starting) + wbar * starting)
474 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
475 active.copy_(torch.maximum(active, starting))
476 pen = torch.clamp(wbar - z, min=0.0)
477 F = active * kc * pen * torch.sqrt(pen)
478 Fst.copy_(F)
479 imp.add_(dt * F)
480 dur.add_(dt * (F > 0).to(torch.float64))
481 fmax.copy_(torch.maximum(fmax, F))
482 f = force(q, v)
483 f = f + F @ phi_load # the shell is pushed down (and, with --no-drift, held rigidly): M a = -f
484 acc.copy_(-f * invM - 2 * sigma0 * v)
485 zd.add_(dt * F / ms)
486 z.add_(dt * zd)
487 v.add_(dt * acc)
488 q.add_(dt * v)
489 # Rayleigh integral: scatter the Gauss-point accelerations onto the retarded output samples
490 ag = (acc[cell_dofs_w] @ NwT).reshape(-1) # (C*G,)
491 idx, val = ray_c(n, ag)
492 out.index_add_(0, idx, val)
493 if out_direct is not None: # the same integral without retardation, for checking the scatter above
494 ks = torch.div(n, m, rounding_mode="floor")
495 on_sample = ((n % m) == 0) & (ks < nsamp)
496 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))
497 nstep.add_(1)
499 # fields
500 fld = None
501 if field_grid is not None:
502 fx, fy, ft = field_grid
503 idx, wts, inside = self.field_interpolator(fx, fy)
504 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=[])
506 def take_frame():
507 fr = (q[fld["idx"]] * fld["wts"]).sum(1).cpu().numpy().astype(np.float32)
508 fr[~fld["inside"]] = np.nan
509 fld["frames"].append(fr)
511 def reset_state():
512 for tsr in (q, v, acc, out, z, zd, active, Fst, imp, dur, fmax, nstep):
513 tsr.zero_()
514 if out_direct is not None:
515 out_direct.zero_()
517 graph = None
518 if use_graph and device.startswith("cuda"):
519 s = torch.cuda.Stream()
520 s.wait_stream(torch.cuda.current_stream())
521 with torch.cuda.stream(s):
522 for _ in range(2):
523 step()
524 torch.cuda.current_stream().wait_stream(s)
525 # undo the warm-up steps
526 reset_state()
527 graph = torch.cuda.CUDAGraph()
528 with torch.cuda.graph(graph):
529 step()
530 reset_state()
531 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 ''}")
533 t_start = time.perf_counter()
534 last_report = t_start
535 chunk = fld["every"] if fld is not None else 2000
536 if fld is not None:
537 take_frame()
538 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()],
539 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()])
540 contact_log = []
541 Fprev = np.zeros(nst)
542 n_done = 0
543 while n_done < nsteps:
544 k = min(chunk, nsteps - n_done)
545 for _ in range(k):
546 if graph is not None:
547 graph.replay()
548 else:
549 step()
550 n_done += k
551 if self.no_drift:
552 project_rigid()
553 if fld is not None and len(fld["frames"]) < fld["nt"]:
554 take_frame()
555 Fc = Fst.cpu().numpy()
556 for i in range(nst):
557 if (Fc[i] > 0) != (Fprev[i] > 0):
558 contact_log.append((round(n_done * dt, 6), i + 1, "on" if Fc[i] > 0 else "off"))
559 Fprev = Fc
560 now = time.perf_counter()
561 if now - last_report >= report or n_done == nsteps:
562 last_report = now
563 el = now - t_start
564 qn = q.cpu().numpy()
565 if not np.all(np.isfinite(qn)):
566 raise RuntimeError(f"solution diverged at step {n_done} (t = {n_done * dt:.4f} s)")
567 w = qn[0 : 4 * self.n_nodes : 4]
568 sig = out.cpu().numpy()
569 prog = dict(tag=tag, step=n_done, steps=nsteps, t=round(n_done * dt, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
570 eta_s=round(el * (nsteps - n_done) / max(n_done, 1), 1), steps_per_s=round(n_done / el, 1),
571 max_abs_w_m=float(np.abs(w).max()), rms_w_m=float(np.sqrt(np.mean(w**2))),
572 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()],
573 contact_force_N=[round(float(x), 2) for x in Fc], contacts=contact_log[-40:], **stick_stats(),
574 peak_pa=float(np.abs(sig).max()), m=m, dt=dt, device=device)
575 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) "
576 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']}")
577 if out_dir is not None:
578 with open(out_dir / f"{tag}.progress.json", "w") as fh:
579 json.dump(prog, fh, indent=1)
580 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
581 elapsed = time.perf_counter() - t_start
582 sig = out.cpu().numpy()
583 info = dict(m=m, dt=dt, omega_max=omega_max, steps=nsteps, elapsed=elapsed, contacts=contact_log, sticks=stick_stats(),
584 out_direct=out_direct.cpu().numpy() if out_direct is not None else None)
585 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
586 return sig, fields, info
589# ----------------------------------------------------------------------------------------------------------------
591def hardware_string(device: str) -> str:
592 cpu = "unknown CPU"
593 try:
594 for line in open("/proc/cpuinfo"):
595 if line.startswith("model name"):
596 cpu = line.split(":", 1)[1].strip()
597 break
598 except OSError:
599 pass
600 import os
602 s = f"{cpu} ({os.cpu_count()} threads)"
603 if device.startswith("cuda"):
604 import torch
606 s += f"; time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
607 return s
610def main():
611 ch = json.load(open(sys.argv[1]))
612 out_dir = Path(sys.argv[2])
613 args = sys.argv[3:]
615 def take(flag, conv, default=None):
616 nonlocal args
617 if flag in args:
618 i = args.index(flag)
619 val = conv(args[i + 1])
620 args = args[:i] + args[i + 2 :]
621 return val
622 return default
624 spacing = take("--spacing", float, 0.008)
625 seconds = take("--seconds", float)
626 device = take("--device", str, "cpu")
627 report = take("--report", float, 60.0)
628 speed = take("--speed", float)
629 m_over = take("--m", int)
630 tag = take("--tag", str)
631 c_air = take("--c-air", float)
632 check = "--check" in args
633 no_drift = "--free" not in args # the rigid constraint of MODEL.md is the default; --free leaves the shell free (diagnostics only)
634 if "--no-drift" in args:
635 args.remove("--no-drift")
636 inplane_scale = take("--inplane-scale", float, 1.0)
637 compile_ = "--no-compile" not in args
638 linear = "--linear" in args
639 fields = "--fields" in args
640 use_graph = "--no-graph" not in args
641 out = ch["output"]
642 fs = out["sample_rate_hz"]
643 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
644 tag = tag or f"fem-g{spacing * 1e3:g}"
645 log = lambda s: print(s, file=sys.stderr, flush=True)
646 out_dir.mkdir(parents=True, exist_ok=True)
647 t0 = time.perf_counter()
648 log(f"{tag}: {ch['challenge_id']}, spacing {spacing * 1e3:g} mm, {nsamp / fs:.2f} s, device {device}{' (linear)' if linear else ''}")
649 model = ShellFEM(ch, spacing, log=log, linear=linear, speed=speed, c_air=c_air, no_drift=no_drift, inplane_scale=inplane_scale)
650 scale = out.get("signal_scale") or 1.0
651 grid = None
652 if fields:
653 fg = ch["fields"]
654 nt, ny, nx = fg["shape"]
655 R = ch["parameters"]["disk_radius"]
656 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
657 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_,
658 field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None)
659 total = time.perf_counter() - t0
660 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
661 meta = {
662 "solution_id": tag,
663 "output": f"{out_dir.name}/{tag}.wav",
664 "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",
665 "parameters": {"grid_spacing_m": spacing, "cells": model.n_cells, "n_dof": model.n_dof, "gauss_points_per_cell": model.G,
666 "oversampling": info["m"], "dt_s": info["dt"], "omega_max_rad_s": info["omega_max"], "inplane_derivative_mass_scale": inplane_scale},
667 "timing_seconds": round(info["elapsed"], 3),
668 "hardware": hardware_string(device),
669 "peak_si": float(np.max(np.abs(sig))),
670 "sticks": info["sticks"],
671 **({"stick_speed_override": speed} if speed is not None else {}),
672 **({"linear": True} if linear else {}),
673 **({"no_drift": True} if no_drift else {}),
674 }
675 if fld is not None:
676 import h5py
678 fg = ch["fields"]
679 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
680 fh.attrs["challenge_id"] = ch["challenge_id"]
681 fh.attrs["solution_id"] = tag
682 fh.attrs["quantity"] = fg["quantity"]
683 fh.attrs["units"] = fg["units"]
684 fh.create_dataset("x", data=grid[0])
685 fh.create_dataset("y", data=grid[1])
686 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
687 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
688 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
689 with open(out_dir / f"{tag}.meta.json", "w") as fh:
690 json.dump(meta, fh, indent=1)
691 if info["out_direct"] is not None:
692 d = info["out_direct"]
693 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)")
694 np.save(out_dir / f"{tag}.direct.npy", d)
695 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
698if __name__ == "__main__":
699 main()