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