1"""Modal solution of the cymbal2 shallow-shell model on the same discretization as fem_shell.py.
3The semi-discrete system is the one fem_shell.py builds: bicubic Hermite elements for u, v, w on the Cartesian grid,
43 x 3 Gauss points, lumped mass (with the same in-plane derivative-mass scaling). Here the linear part of that system
5is diagonalized once, K phi = omega^2 M phi, and only the modes below a cutoff frequency f_max are kept; the six rigid
6fields are excluded exactly, by a penalty in the eigenproblem that lifts them out of range (the rigid constraint of
7MODEL.md). The state is the vector of modal amplitudes a(t), and the equations are
9 a_k'' + 2 alpha_k a_k' + omega_k^2 a_k = F_k(a, a'), alpha_k = sigma0 + eta omega_k^2 / 2,
11where F is the modal projection of the nonlinear remainder of the internal force, evaluated at the Gauss points of the
12mesh by the same routine as the finite element solver (so the nonlinear terms are not truncated to any tensor of modal
13couplings), plus the stick loads. The time stepping is the two-step scheme of the v1 modal solver: the damped linear
14oscillators are propagated exactly over each step and the force is applied with a static-exact weight, which is
15second-order accurate and free of dispersion in the linear limit, so the step can be a small multiple of the output
16sample. The sticks obey the Hertz law with the footprint-averaged displacement as in fem_shell.py. The Rayleigh
17integral is a per-step weight table folded with a ring buffer of modal accelerations (cubic Lagrange interpolation in
18time at the retarded time of every Gauss point).
20The two solvers therefore solve the same discrete equations and differ only in the modal truncation and the time
21integration; their difference measures the effect of the modes above f_max and of the step.
23Usage:
24 python modal_shell.py <challenge.json> <out-dir> [--spacing 0.008] [--fmax 20000] [--m 4] [--seconds S] [--device cuda]
25 [--fields] [--report 60] [--tag NAME] [--inplane-scale 16]
26"""
27from __future__ import annotations
29import json
30import math
31import sys
32import time
33from pathlib import Path
35import numpy as np
36import scipy.sparse as sp
37import scipy.sparse.linalg as spla
38from scipy.io import wavfile
40from fem_shell import ShellFEM, hardware_string
43def sliced_modes(K, Mdiag, rigid, f_max: float, k_slice: int = 250, log=print, penalty: float = 1e12):
44 """Every eigenpair of K phi = lambda M phi with lambda <= (2 pi f_max)^2, except the six rigid fields, which are
45 lifted to the eigenvalue `penalty` by the rank-6 term M R R^T M (R the M-orthonormal rigid fields) and thereby
46 made exactly orthogonal to every returned mode. Shift-invert ARPACK in slices of k_slice eigenvalues, each slice
47 factorizing K - sigma M once (the low-rank term is handled with the Woodbury identity)."""
48 n = K.shape[0]
49 M = sp.diags(Mdiag).tocsc()
50 U = (Mdiag[:, None] * rigid.T) * math.sqrt(penalty) # (n, 6): penalty term is U U^T
51 Kop = spla.LinearOperator((n, n), matvec=lambda x: K @ x + U @ (U.T @ x), dtype=np.float64)
52 target = (2 * math.pi * f_max) ** 2
53 vals_all, vecs_all = [], []
54 lo = -np.inf
55 sigma = -1e3
56 width = None
57 tries = 0
58 while True:
59 t0 = time.perf_counter()
60 lu = spla.splu((K - sigma * M).tocsc())
61 AiU = lu.solve(U) # (n, 6)
62 S = np.linalg.inv(np.eye(6) + U.T @ AiU)
64 def opinv(x, lu=lu, AiU=AiU, S=S):
65 y = lu.solve(x)
66 return y - AiU @ (S @ (U.T @ y))
68 OPinv = spla.LinearOperator((n, n), matvec=opinv, dtype=np.float64)
69 vals, vecs = spla.eigsh(Kop, k=k_slice, M=M, sigma=sigma, which="LM", OPinv=OPinv, tol=1e-9)
70 order = np.argsort(vals)
71 vals, vecs = vals[order], vecs[:, order]
72 tries += 1
73 if np.isfinite(lo) and vals[0] > lo * (1 + 1e-9) + 1e-3:
74 if tries > 40:
75 raise RuntimeError("mode slicing did not close the gap")
76 sigma = 0.5 * (sigma + lo) # the slice did not reach back to the covered range: step back
77 log(f" slice at {math.sqrt(max(sigma, 0)) / 2 / math.pi:.0f} Hz left a gap, retrying")
78 continue
79 keep = vals > lo if np.isfinite(lo) else np.ones(len(vals), dtype=bool)
80 vals_all.append(vals[keep])
81 vecs_all.append(vecs[:, keep])
82 lo = vals[-1]
83 width = vals[-1] - vals[0]
84 log(f" slice: {int(keep.sum())} new modes up to {math.sqrt(max(lo, 0)) / 2 / math.pi:.0f} Hz in {time.perf_counter() - t0:.1f}s")
85 if lo >= target:
86 break
87 sigma = lo + 0.45 * width
88 vals = np.concatenate(vals_all)
89 vecs = np.concatenate(vecs_all, axis=1)
90 sel = (vals <= target) & (vals < 0.5 * penalty)
91 vals, vecs = vals[sel], vecs[:, sel]
92 # rigid content check and clean-up of the near-null modes of the flat geometry (exact zeros the penalty removed)
93 rc = np.abs(rigid @ (Mdiag[:, None] * vecs)).max()
94 log(f" {len(vals)} modes below {f_max:g} Hz; largest rigid component {rc:.1e}")
95 return np.sqrt(np.maximum(vals, 0.0)), vecs
98class ModalShell:
99 def __init__(self, fem: ShellFEM, f_max: float, log=print):
100 self.fem = fem
101 self.log = log
102 t0 = time.perf_counter()
103 K = fem.assemble_linear()
104 self.log(f" linear stiffness assembled ({K.nnz} nonzeros) in {time.perf_counter() - t0:.1f}s")
105 t0 = time.perf_counter()
106 cache_dir = Path(__file__).resolve().parent / "cache"
107 cache_dir.mkdir(exist_ok=True)
108 import hashlib
110 geom_key = hashlib.sha1(json.dumps({k: v for k, v in fem.ch["parameters"].items() if k in ("disk_radius", "hole_radius", "bow_height", "bell_radius", "bell_height", "thickness", "thickness_center", "thickness_rim", "E", "nu", "rho")}, sort_keys=True).encode()).hexdigest()[:8]
111 cache = cache_dir / f"modes-{fem.ch['challenge_id']}-g{fem.a * 1e3:g}-f{f_max / 1e3:g}k-s{fem.inplane_scale:g}-{geom_key}.npz"
112 if cache.exists():
113 z = np.load(cache)
114 self.omega, self.Phi = z["omega"], z["Phi"]
115 self.log(f" modes loaded from {cache.name}")
116 else:
117 self.omega, self.Phi = sliced_modes(K, fem.Mdiag, fem.rigid, f_max, log=log)
118 np.savez(cache, omega=self.omega, Phi=self.Phi)
119 self.n_modes = len(self.omega)
120 self.log(f" {self.n_modes} modes in {time.perf_counter() - t0:.1f}s: f = {np.round(self.omega[:5] / 2 / np.pi, 2)} ... {self.omega[-1] / 2 / np.pi:.0f} Hz")
121 p = fem.p
122 self.alpha = p["sigma0"] + 0.5 * p["eta"] * self.omega**2
124 def run(self, fs: int, nsamp: int, m: int, device: str = "cpu", field_grid=None, field_rate=None, report: float = 60.0,
125 out_dir: Path | None = None, tag: str = "modal", scale: float = 1.0, compile_: bool = True, use_graph: bool = True):
126 import torch
128 fem, p = self.fem, self.fem.p
129 t_setup = time.perf_counter()
130 dev = torch.device(device)
131 force, ops = fem.build_force(torch, dev, compile_)
132 T64, T64i = ops["T64"], ops["T64i"]
133 N = self.n_modes
134 dt = 1.0 / (fs * m)
135 nsteps = nsamp * m
136 Phi = T64(self.Phi) # (n_dof, N)
137 PhiT = Phi.T.contiguous()
138 w2 = T64(self.omega**2)
139 al = T64(self.alpha)
140 eta, sigma0 = p["eta"], p["sigma0"]
141 ms, kc, v0 = p["stick_mass"], p["contact_stiffness"], fem.v0
142 # exact two-step recurrence of the damped linear oscillators
143 disc = np.sqrt(self.alpha**2 - self.omega**2 + 0j)
144 E1, E2 = np.exp((-self.alpha + disc) * dt), np.exp((-self.alpha - disc) * dt)
145 a1 = T64((E1 + E2).real)
146 a2 = T64(-(E1 * E2).real)
147 b = T64((1.0 - (E1 + E2).real + (E1 * E2).real) / self.omega**2)
148 # Rayleigh integral: per-step weight table G[d, k] from the modes at the Gauss points
149 tau = fem.ray_tau.reshape(-1)
150 cq = fem.ray_c.reshape(-1)
151 tau_steps = tau / dt
152 d0 = np.floor(tau_steps).astype(np.int64)
153 fr = tau_steps - d0
154 lag = np.stack([-fr * (fr - 1.0) * (fr - 2.0) / 6.0, (fr + 1.0) * (fr - 1.0) * (fr - 2.0) / 2.0,
155 -(fr + 1.0) * fr * (fr - 2.0) / 2.0, (fr + 1.0) * fr * (fr - 1.0) / 6.0], axis=1)
156 dmax = int(d0.max()) + 2
157 L = dmax + 1
158 Nw = T64(fem.Nw) # (G, 16)
159 Psi = torch.einsum("gi,cin->cgn", Nw, Phi[ops["cell_dofs_w"]]).reshape(-1, N) # (C*G, N) w of the modes at the Gauss points
160 Gt = torch.zeros((L, N), dtype=torch.float64, device=dev)
161 for j, off in enumerate([-1, 0, 1, 2]):
162 Gt.index_add_(0, T64i(d0 + off), T64(cq * lag[:, j])[:, None] * Psi)
163 del Psi
164 # sticks
165 nst = len(fem.strikes)
166 gains = T64(np.stack([s["phi"] for s in fem.strikes])) @ Phi # (nst, N): footprint average of the modes
167 n_on = T64i([int(round(s["t"] / dt)) for s in fem.strikes])
168 z = torch.zeros(nst, dtype=torch.float64, device=dev)
169 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
170 active = torch.zeros(nst, dtype=torch.float64, device=dev)
171 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
172 imp = torch.zeros(nst, dtype=torch.float64, device=dev)
173 dur = torch.zeros(nst, dtype=torch.float64, device=dev)
174 fmax_t = torch.zeros(nst, dtype=torch.float64, device=dev)
175 # state
176 a = torch.zeros(N, dtype=torch.float64, device=dev)
177 a_prev = torch.zeros(N, dtype=torch.float64, device=dev)
178 q = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
179 q_prev = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
180 ring = torch.zeros((L, N), dtype=torch.float64, device=dev)
181 out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
182 nstep = torch.zeros((), dtype=torch.int64, device=dev)
183 ar = torch.arange(L, device=dev)
184 fld = None
185 if field_grid is not None:
186 fx, fy, ft = field_grid
187 idx, wts, inside = fem.field_interpolator(fx, fy)
188 Wg = (T64(wts)[:, :, None] * Phi[T64i(idx)]).sum(1) # (P, N)
189 fld = dict(Wg=Wg, inside=inside, every=int(round(fs / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])
191 def step():
192 n = nstep
193 q.copy_(Phi @ a)
194 v = (q - q_prev) / dt
195 adot = (a - a_prev) / dt
196 # sticks
197 wbar = gains @ a
198 wbar_dot = gains @ adot
199 starting = (n == n_on).to(torch.float64)
200 z.copy_(z * (1 - starting) + wbar * starting)
201 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
202 active.copy_(torch.maximum(active, starting))
203 pen = torch.clamp(wbar - z, min=0.0)
204 F = active * kc * pen * torch.sqrt(pen)
205 Fst.copy_(F)
206 imp.add_(dt * F)
207 dur.add_(dt * (F > 0).to(torch.float64))
208 fmax_t.copy_(torch.maximum(fmax_t, F))
209 # modal force: nonlinear remainder of the internal force plus the sticks
210 f = force(q, v)
211 Fm = -(PhiT @ f) + w2 * a + eta * w2 * adot - F @ gains
212 a_next = a1 * a + a2 * a_prev + b * Fm
213 add = Fm - w2 * a - al * ((a_next - a_prev) / dt)
214 ring.index_copy_(0, (n % L).view(1), add.view(1, -1))
215 zd.add_(dt * F / ms)
216 z.add_(dt * zd)
217 q_prev.copy_(q)
218 a_prev.copy_(a)
219 a.copy_(a_next)
220 nstep.add_(1)
222 def sample(k):
223 # sample k is at step k m: the first step below computes and stores the acceleration there, the output sample
224 # folds the table with the ring (delays counted back from that step), then the remaining steps of the sample
225 step()
226 idx = (nstep - 1 - ar) % L
227 out.index_copy_(0, k.view(1), torch.sum(Gt * ring[idx]).view(1))
228 for _ in range(m - 1):
229 step()
231 graph = None
232 k_t = torch.zeros((), dtype=torch.int64, device=dev)
234 def sample_graph():
235 sample(k_t)
237 def reset():
238 for t in (z, zd, active, Fst, imp, dur, fmax_t, a, a_prev, q, q_prev, ring, out, nstep, k_t):
239 t.zero_()
241 if use_graph and device.startswith("cuda"):
242 s = torch.cuda.Stream()
243 s.wait_stream(torch.cuda.current_stream())
244 with torch.cuda.stream(s):
245 for _ in range(2):
246 sample_graph()
247 torch.cuda.current_stream().wait_stream(s)
248 reset()
249 graph = torch.cuda.CUDAGraph()
250 with torch.cuda.graph(graph):
251 sample_graph()
252 reset()
253 self.log(f" setup {time.perf_counter() - t_setup:.1f}s; stepping {nsteps} steps ({nsamp / fs:.2f} s, m = {m}) on {device}{' with a CUDA graph' if graph else ''}")
255 def take_frame():
256 fr_ = (fld["Wg"] @ a).cpu().numpy().astype(np.float32)
257 fr_[~fld["inside"]] = np.nan
258 fld["frames"].append(fr_)
260 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()],
261 peak_force_N=[round(float(x), 2) for x in fmax_t.cpu().numpy()], rebound_mps=[round(float(x), 4) for x in zd.cpu().numpy()])
262 t_start = time.perf_counter()
263 last_report = t_start
264 if fld is not None:
265 take_frame()
266 for k in range(nsamp):
267 k_t.fill_(k)
268 if graph is not None:
269 graph.replay()
270 else:
271 sample(k_t)
272 if fld is not None and (k + 1) % fld["every"] == 0 and len(fld["frames"]) < fld["nt"]:
273 take_frame()
274 now = time.perf_counter()
275 if (now - last_report >= report or k == nsamp - 1) and k > 0:
276 last_report = now
277 el = now - t_start
278 an = a.cpu().numpy()
279 if not np.all(np.isfinite(an)):
280 raise RuntimeError(f"solution diverged at sample {k} (t = {k / fs:.4f} s); increase m")
281 sig = out.cpu().numpy()
282 prog = dict(tag=tag, step=(k + 1) * m, steps=nsteps, t=round((k + 1) / fs, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
283 eta_s=round(el * (nsamp - k - 1) / (k + 1), 1), steps_per_s=round((k + 1) * m / el, 1), max_abs_a=float(np.abs(an).max()),
284 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()],
285 contact_force_N=[round(float(x), 2) for x in Fst.cpu().numpy()], **stick_stats(), peak_pa=float(np.abs(sig).max()), m=m, dt=dt, device=device, n_modes=N)
286 self.log(f" t={prog['t']:.3f}s elapsed {el:.0f}s eta {prog['eta_s']:.0f}s ({prog['steps_per_s']:.0f} steps/s) max|a|={prog['max_abs_a']:.3e} "
287 f"peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
288 if out_dir is not None:
289 with open(out_dir / f"{tag}.progress.json", "w") as fh:
290 json.dump(prog, fh, indent=1)
291 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
292 elapsed = time.perf_counter() - t_start
293 sig = out.cpu().numpy()
294 if not np.all(np.isfinite(sig)):
295 raise RuntimeError("output is not finite; increase m")
296 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
297 return sig, fields, dict(m=m, dt=dt, elapsed=elapsed, sticks=stick_stats())
300def main():
301 ch = json.load(open(sys.argv[1]))
302 out_dir = Path(sys.argv[2])
303 args = sys.argv[3:]
305 def take(flag, conv, default=None):
306 nonlocal args
307 if flag in args:
308 i = args.index(flag)
309 val = conv(args[i + 1])
310 args = args[:i] + args[i + 2 :]
311 return val
312 return default
314 spacing = take("--spacing", float, 0.008)
315 f_max = take("--fmax", float, 20000.0)
316 m = take("--m", int, 4)
317 seconds = take("--seconds", float)
318 device = take("--device", str, "cpu")
319 report = take("--report", float, 60.0)
320 speed = take("--speed", float)
321 tag = take("--tag", str)
322 inplane_scale = take("--inplane-scale", float, 1.0)
323 fields = "--fields" in args
324 compile_ = "--no-compile" not in args
325 use_graph = "--no-graph" not in args
326 out = ch["output"]
327 fs = out["sample_rate_hz"]
328 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
329 tag = tag or f"modal-g{spacing * 1e3:g}-f{f_max / 1e3:g}k-m{m}"
330 log = lambda s: print(s, file=sys.stderr, flush=True)
331 out_dir.mkdir(parents=True, exist_ok=True)
332 t0 = time.perf_counter()
333 log(f"{tag}: {ch['challenge_id']}, spacing {spacing * 1e3:g} mm, modes to {f_max:g} Hz, m = {m}, {nsamp / fs:.2f} s, device {device}")
334 fem = ShellFEM(ch, spacing, log=log, speed=speed, no_drift=True, inplane_scale=inplane_scale)
335 model = ModalShell(fem, f_max, log=log)
336 t1 = time.perf_counter()
337 scale = out.get("signal_scale") or 1.0
338 grid = None
339 if fields:
340 fg = ch["fields"]
341 nt, ny, nx = fg["shape"]
342 R = ch["parameters"]["disk_radius"]
343 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
344 sig, fld, info = model.run(fs, nsamp, m, device=device, field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None,
345 report=report, out_dir=out_dir, tag=tag, scale=scale, compile_=compile_, use_graph=use_graph)
346 total = time.perf_counter() - t0
347 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
348 meta = {
349 "solution_id": tag,
350 "output": f"{out_dir.name}/{tag}.wav",
351 "method": "modes of the bicubic Hermite finite element system on a Cartesian grid masked to the annulus (lumped mass), truncated at a cutoff frequency, "
352 "exact linear propagation with the nonlinear force evaluated at the Gauss points, Hertz stick, Rayleigh integral by a retarded weight table",
353 "parameters": {"grid_spacing_m": spacing, "cells": fem.n_cells, "n_dof": fem.n_dof, "n_modes": model.n_modes, "mode_cutoff_hz": f_max,
354 "highest_mode_hz": round(float(model.omega.max() / 2 / np.pi), 1), "oversampling": m, "dt_s": info["dt"], "inplane_derivative_mass_scale": inplane_scale},
355 "timing_seconds": round(info["elapsed"], 3),
356 "setup_seconds": round(t1 - t0, 3),
357 "hardware": hardware_string(device),
358 "peak_si": float(np.max(np.abs(sig))),
359 "sticks": info["sticks"],
360 **({"stick_speed_override": speed} if speed is not None else {}),
361 }
362 if fld is not None:
363 import h5py
365 fg = ch["fields"]
366 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
367 fh.attrs["challenge_id"] = ch["challenge_id"]
368 fh.attrs["solution_id"] = tag
369 fh.attrs["quantity"] = fg["quantity"]
370 fh.attrs["units"] = fg["units"]
371 fh.create_dataset("x", data=grid[0])
372 fh.create_dataset("y", data=grid[1])
373 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
374 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
375 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
376 with open(out_dir / f"{tag}.meta.json", "w") as fh:
377 json.dump(meta, fh, indent=1)
378 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
381if __name__ == "__main__":
382 main()