/ challenges-old / cymbal2-solutions
Sign in
challenges-old / cymbal2-solutions
main
Branches
main
cymbal2-solutions / solvers / hermite-fem / fem_shell.py
713 lines · 37.7 KBCodeBlameHistory
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"]
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 46 self.rh = p.get("mount_radius", p.get("hole_radius", 0.0)) # inner radius of the shell (the clamped mount)
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)
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 140 inner_cells = set(map(tuple, np.stack([ii.ravel()[rc <= rh], jj.ravel()[rc <= rh]], axis=1).tolist())) # cells inside the mount
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
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 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
38bdd34cymbal2-solutions: Hermite finite element solver (direct and modal) for the cymbal2 challenges, site builder, workflowJeremy Magland 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)
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 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).
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"]
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 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"]
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 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)
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
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 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
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)
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)
f150696Output with the area mean of the acceleration removed; experiment harnessJeremy Magland 432 wmean = T64(np.tile(self.wq, C) / self.area) # (C*G,) quadrature weights over the area: the area mean of a Gauss-point field
434 # ---- strikes
435 nst = len(self.strikes)
436 phi = T64(np.stack([s["phi"] for s in self.strikes])) # (nst, n_dof)
437 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
438 rigM = T64(self.rigid * self.Mdiag[None, :]) # (6, n_dof): <rigid_i, .>_M
439 rig = T64(self.rigid)
441 def project_rigid(): # remove the rigid components of the state (the constraint of MODEL.md, enforced per chunk)
442 for x in (q, v):
443 x.sub_(rig.T @ (rigM @ x))
444 t_on = T64([s["t"] for s in self.strikes])
445 n_on = T64i([int(round(s["t"] / dt)) for s in self.strikes])
446 z = torch.zeros(nst, dtype=torch.float64, device=dev)
447 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
448 active = torch.zeros(nst, dtype=torch.float64, device=dev)
449 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
450 imp = torch.zeros(nst, dtype=torch.float64, device=dev) # contact impulse per stick
451 dur = torch.zeros(nst, dtype=torch.float64, device=dev) # time in contact
452 fmax = torch.zeros(nst, dtype=torch.float64, device=dev)
453 out_direct = torch.zeros(nsamp, dtype=torch.float64, device=dev) if check else None
454 # ---- state
455 q = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
456 v = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
457 acc = torch.zeros(self.n_dof, dtype=torch.float64, device=dev)
458 nstep = torch.zeros((), dtype=torch.int64, device=dev)
459 lag_c = 1.0 / 6.0
461 assert m >= 4, "the retarded-time interpolation assumes at least 4 steps per sample"
463 def rayleigh(n, ag):
464 nf = n.to(torch.float64)
465 kk = torch.ceil((nf - 2 + dsteps) / m)
466 pos = kk * m - dsteps
467 fl = torch.floor(pos)
468 j = nf - fl + 1
469 fr = pos - fl
470 wgt = torch.where(j == 0, -fr * (fr - 1) * (fr - 2) * lag_c,
471 torch.where(j == 1, (fr + 1) * (fr - 1) * (fr - 2) * 0.5,
472 torch.where(j == 2, -(fr + 1) * fr * (fr - 2) * 0.5,
473 torch.where(j == 3, (fr + 1) * fr * (fr - 1) * lag_c, torch.zeros_like(fr)))))
474 valid = (j >= 0) & (j <= 3) & (kk >= 0) & (kk < nsamp)
475 idx = torch.where(valid, kk, torch.zeros_like(kk)).to(torch.int64)
476 return idx, torch.where(valid, rc * wgt * ag, torch.zeros_like(ag))
478 ray_c = torch.compile(rayleigh, dynamic=False) if (compile_ and device.startswith("cuda")) else rayleigh
480 def step():
481 n = nstep
482 # sticks: start the ones whose time has come (at the current surface, with relative speed v0)
483 wbar = phi @ q
484 wbar_dot = phi @ v
485 starting = (n == n_on).to(torch.float64)
486 z.copy_(z * (1 - starting) + wbar * starting)
487 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
488 active.copy_(torch.maximum(active, starting))
489 pen = torch.clamp(wbar - z, min=0.0)
490 F = active * kc * pen * torch.sqrt(pen)
491 Fst.copy_(F)
492 imp.add_(dt * F)
493 dur.add_(dt * (F > 0).to(torch.float64))
494 fmax.copy_(torch.maximum(fmax, F))
495 f = force(q, v)
496 f = f + F @ phi_load # the shell is pushed down (and, with --no-drift, held rigidly): M a = -f
497 acc.copy_(-f * invM - 2 * sigma0 * v)
498 zd.add_(dt * F / ms)
499 z.add_(dt * zd)
500 v.add_(dt * acc)
501 q.add_(dt * v)
502 # Rayleigh integral: scatter the Gauss-point accelerations onto the retarded output samples
503 ag = (acc[cell_dofs_w] @ NwT).reshape(-1) # (C*G,)
f150696Output with the area mean of the acceleration removed; experiment harnessJeremy Magland 504 ag = ag - torch.dot(wmean, ag) # the output uses the acceleration with its area mean removed (MODEL.md)
506 out.index_add_(0, idx, val)
507 if out_direct is not None: # the same integral without retardation, for checking the scatter above
508 ks = torch.div(n, m, rounding_mode="floor")
509 on_sample = ((n % m) == 0) & (ks < nsamp)
510 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))
511 nstep.add_(1)
513 # fields
514 fld = None
515 if field_grid is not None:
516 fx, fy, ft = field_grid
517 idx, wts, inside = self.field_interpolator(fx, fy)
518 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=[])
520 def take_frame():
521 fr = (q[fld["idx"]] * fld["wts"]).sum(1).cpu().numpy().astype(np.float32)
522 fr[~fld["inside"]] = np.nan
523 fld["frames"].append(fr)
525 def reset_state():
526 for tsr in (q, v, acc, out, z, zd, active, Fst, imp, dur, fmax, nstep):
527 tsr.zero_()
528 if out_direct is not None:
529 out_direct.zero_()
531 graph = None
532 if use_graph and device.startswith("cuda"):
533 s = torch.cuda.Stream()
534 s.wait_stream(torch.cuda.current_stream())
535 with torch.cuda.stream(s):
536 for _ in range(2):
537 step()
538 torch.cuda.current_stream().wait_stream(s)
539 # undo the warm-up steps
540 reset_state()
541 graph = torch.cuda.CUDAGraph()
542 with torch.cuda.graph(graph):
543 step()
544 reset_state()
545 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 ''}")
547 t_start = time.perf_counter()
548 last_report = t_start
549 chunk = fld["every"] if fld is not None else 2000
550 if fld is not None:
551 take_frame()
552 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()],
553 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()])
554 contact_log = []
555 Fprev = np.zeros(nst)
556 n_done = 0
557 while n_done < nsteps:
558 k = min(chunk, nsteps - n_done)
559 for _ in range(k):
560 if graph is not None:
561 graph.replay()
562 else:
563 step()
564 n_done += k
565 if self.no_drift:
566 project_rigid()
567 if fld is not None and len(fld["frames"]) < fld["nt"]:
568 take_frame()
569 Fc = Fst.cpu().numpy()
570 for i in range(nst):
571 if (Fc[i] > 0) != (Fprev[i] > 0):
572 contact_log.append((round(n_done * dt, 6), i + 1, "on" if Fc[i] > 0 else "off"))
573 Fprev = Fc
574 now = time.perf_counter()
575 if now - last_report >= report or n_done == nsteps:
576 last_report = now
577 el = now - t_start
578 qn = q.cpu().numpy()
579 if not np.all(np.isfinite(qn)):
580 raise RuntimeError(f"solution diverged at step {n_done} (t = {n_done * dt:.4f} s)")
581 w = qn[0 : 4 * self.n_nodes : 4]
582 sig = out.cpu().numpy()
583 prog = dict(tag=tag, step=n_done, steps=nsteps, t=round(n_done * dt, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
584 eta_s=round(el * (nsteps - n_done) / max(n_done, 1), 1), steps_per_s=round(n_done / el, 1),
585 max_abs_w_m=float(np.abs(w).max()), rms_w_m=float(np.sqrt(np.mean(w**2))),
586 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()],
587 contact_force_N=[round(float(x), 2) for x in Fc], contacts=contact_log[-40:], **stick_stats(),
588 peak_pa=float(np.abs(sig).max()), m=m, dt=dt, device=device)
589 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) "
590 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']}")
591 if out_dir is not None:
592 with open(out_dir / f"{tag}.progress.json", "w") as fh:
593 json.dump(prog, fh, indent=1)
594 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
595 elapsed = time.perf_counter() - t_start
596 sig = out.cpu().numpy()
597 info = dict(m=m, dt=dt, omega_max=omega_max, steps=nsteps, elapsed=elapsed, contacts=contact_log, sticks=stick_stats(),
598 out_direct=out_direct.cpu().numpy() if out_direct is not None else None)
599 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
600 return sig, fields, info
603# ----------------------------------------------------------------------------------------------------------------
605def hardware_string(device: str) -> str:
606 cpu = "unknown CPU"
607 try:
608 for line in open("/proc/cpuinfo"):
609 if line.startswith("model name"):
610 cpu = line.split(":", 1)[1].strip()
611 break
612 except OSError:
613 pass
614 import os
616 s = f"{cpu} ({os.cpu_count()} threads)"
617 if device.startswith("cuda"):
618 import torch
620 s += f"; time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
621 return s
624def main():
625 ch = json.load(open(sys.argv[1]))
626 out_dir = Path(sys.argv[2])
627 args = sys.argv[3:]
629 def take(flag, conv, default=None):
630 nonlocal args
631 if flag in args:
632 i = args.index(flag)
633 val = conv(args[i + 1])
634 args = args[:i] + args[i + 2 :]
635 return val
636 return default
638 spacing = take("--spacing", float, 0.008)
639 seconds = take("--seconds", float)
640 device = take("--device", str, "cpu")
641 report = take("--report", float, 60.0)
642 speed = take("--speed", float)
643 m_over = take("--m", int)
644 tag = take("--tag", str)
645 c_air = take("--c-air", float)
646 check = "--check" in args
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 647 no_drift = "--free" not in args # the rigid constraint of MODEL.md is the default; --free leaves the shell free (diagnostics only)
648 if "--no-drift" in args:
649 args.remove("--no-drift")
650 inplane_scale = take("--inplane-scale", float, 1.0)
652 linear = "--linear" in args
653 fields = "--fields" in args
654 use_graph = "--no-graph" not in args
655 out = ch["output"]
656 fs = out["sample_rate_hz"]
657 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
658 tag = tag or f"fem-g{spacing * 1e3:g}"
659 log = lambda s: print(s, file=sys.stderr, flush=True)
660 out_dir.mkdir(parents=True, exist_ok=True)
661 t0 = time.perf_counter()
662 log(f"{tag}: {ch['challenge_id']}, spacing {spacing * 1e3:g} mm, {nsamp / fs:.2f} s, device {device}{' (linear)' if linear else ''}")
663 model = ShellFEM(ch, spacing, log=log, linear=linear, speed=speed, c_air=c_air, no_drift=no_drift, inplane_scale=inplane_scale)
664 scale = out.get("signal_scale") or 1.0
665 grid = None
666 if fields:
667 fg = ch["fields"]
668 nt, ny, nx = fg["shape"]
669 R = ch["parameters"]["disk_radius"]
670 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
671 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_,
672 field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None)
673 total = time.perf_counter() - t0
674 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
675 meta = {
676 "solution_id": tag,
677 "output": f"{out_dir.name}/{tag}.wav",
678 "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",
679 "parameters": {"grid_spacing_m": spacing, "cells": model.n_cells, "n_dof": model.n_dof, "gauss_points_per_cell": model.G,
680 "oversampling": info["m"], "dt_s": info["dt"], "omega_max_rad_s": info["omega_max"], "inplane_derivative_mass_scale": inplane_scale},
681 "timing_seconds": round(info["elapsed"], 3),
682 "hardware": hardware_string(device),
683 "peak_si": float(np.max(np.abs(sig))),
684 "sticks": info["sticks"],
685 **({"stick_speed_override": speed} if speed is not None else {}),
686 **({"linear": True} if linear else {}),
687 **({"no_drift": True} if no_drift else {}),
688 }
689 if fld is not None:
690 import h5py
692 fg = ch["fields"]
693 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
694 fh.attrs["challenge_id"] = ch["challenge_id"]
695 fh.attrs["solution_id"] = tag
696 fh.attrs["quantity"] = fg["quantity"]
697 fh.attrs["units"] = fg["units"]
698 fh.create_dataset("x", data=grid[0])
699 fh.create_dataset("y", data=grid[1])
700 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
701 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
702 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
703 with open(out_dir / f"{tag}.meta.json", "w") as fh:
704 json.dump(meta, fh, indent=1)
705 if info["out_direct"] is not None:
706 d = info["out_direct"]
707 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)")
708 np.save(out_dir / f"{tag}.direct.npy", d)
709 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
712if __name__ == "__main__":
713 main()
↑↓ move↵openescclose