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 if rigid is None: # a clamped shell has no rigid fields: plain shift-invert
51 rigid = np.zeros((0, n))
52 U = (Mdiag[:, None] * rigid.T) * math.sqrt(penalty) # (n, r): penalty term is U U^T
53 nr = U.shape[1]
54 Kop = spla.LinearOperator((n, n), matvec=lambda x: K @ x + U @ (U.T @ x), dtype=np.float64)
55 target = (2 * math.pi * f_max) ** 2
56 vals_all, vecs_all = [], []
57 lo = -np.inf
58 sigma = -1e3
59 width = None
60 tries = 0
61 while True:
62 t0 = time.perf_counter()
63 lu = spla.splu((K - sigma * M).tocsc())
64 AiU = lu.solve(U) if nr else np.zeros((n, 0)) # (n, r)
65 S = np.linalg.inv(np.eye(nr) + U.T @ AiU) if nr else np.zeros((0, 0))
67 def opinv(x, lu=lu, AiU=AiU, S=S):
68 y = lu.solve(x)
69 return y - AiU @ (S @ (U.T @ y))
71 OPinv = spla.LinearOperator((n, n), matvec=opinv, dtype=np.float64)
72 vals, vecs = spla.eigsh(Kop, k=k_slice, M=M, sigma=sigma, which="LM", OPinv=OPinv, tol=1e-9)
73 order = np.argsort(vals)
74 vals, vecs = vals[order], vecs[:, order]
75 tries += 1
76 if np.isfinite(lo) and vals[0] > lo * (1 + 1e-9) + 1e-3:
77 if tries > 40:
78 raise RuntimeError("mode slicing did not close the gap")
79 sigma = 0.5 * (sigma + lo) # the slice did not reach back to the covered range: step back
80 log(f" slice at {math.sqrt(max(sigma, 0)) / 2 / math.pi:.0f} Hz left a gap, retrying")
81 continue
82 keep = vals > lo if np.isfinite(lo) else np.ones(len(vals), dtype=bool)
83 vals_all.append(vals[keep])
84 vecs_all.append(vecs[:, keep])
85 lo = vals[-1]
86 width = vals[-1] - vals[0]
87 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")
88 if lo >= target:
89 break
90 sigma = lo + 0.45 * width
91 vals = np.concatenate(vals_all)
92 vecs = np.concatenate(vecs_all, axis=1)
93 sel = (vals <= target) & (vals < 0.5 * penalty)
94 vals, vecs = vals[sel], vecs[:, sel]
95 # rigid content check and clean-up of the near-null modes of the flat geometry (exact zeros the penalty removed)
96 rc = np.abs(rigid @ (Mdiag[:, None] * vecs)).max() if nr else 0.0
97 log(f" {len(vals)} modes below {f_max:g} Hz" + (f"; largest rigid component {rc:.1e}" if nr else " (clamped: no rigid fields)"))
98 return np.sqrt(np.maximum(vals, 0.0)), vecs
101class ModalShell:
102 def __init__(self, fem: ShellFEM, f_max: float, log=print):
103 self.fem = fem
104 self.log = log
105 t0 = time.perf_counter()
106 K = fem.assemble_linear()
107 self.log(f" linear stiffness assembled ({K.nnz} nonzeros) in {time.perf_counter() - t0:.1f}s")
108 t0 = time.perf_counter()
109 cache_dir = Path(__file__).resolve().parent / "cache"
110 cache_dir.mkdir(exist_ok=True)
111 import hashlib
113 geom_key = hashlib.sha1(json.dumps({k: v for k, v in fem.ch["parameters"].items() if k in ("disk_radius", "hole_radius", "mount_radius", "bow_height", "bell_radius", "bell_height", "thickness", "thickness_center", "thickness_rim", "E", "nu", "rho")}, sort_keys=True).encode()).hexdigest()[:8]
114 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"
115 if cache.exists():
116 z = np.load(cache)
117 self.omega, self.Phi = z["omega"], z["Phi"]
118 self.log(f" modes loaded from {cache.name}")
119 else:
120 if fem.clamped: # eigenproblem on the free unknowns only; the modes are zero at the mount
121 fr = fem.free_dofs
122 om, Pf = sliced_modes(K[fr][:, fr].tocsr(), fem.Mdiag[fr], None, f_max, log=log)
123 Phi = np.zeros((fem.n_dof, len(om)))
124 Phi[fr] = Pf
125 self.omega, self.Phi = om, Phi
126 else:
127 self.omega, self.Phi = sliced_modes(K, fem.Mdiag, fem.rigid, f_max, log=log)
128 np.savez(cache, omega=self.omega, Phi=self.Phi)
129 self.n_modes = len(self.omega)
130 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")
131 p = fem.p
132 self.alpha = p["sigma0"] + 0.5 * p["eta"] * self.omega**2
134 def run(self, fs: int, nsamp: int, m: int, device: str = "cpu", field_grid=None, field_rate=None, report: float = 60.0,
135 out_dir: Path | None = None, tag: str = "modal", scale: float = 1.0, compile_: bool = True, use_graph: bool = True):
136 import torch
138 fem, p = self.fem, self.fem.p
139 t_setup = time.perf_counter()
140 dev = torch.device(device)
141 force, ops = fem.build_force(torch, dev, compile_)
142 T64, T64i = ops["T64"], ops["T64i"]
143 N = self.n_modes
144 dt = 1.0 / (fs * m)
145 nsteps = nsamp * m
146 Phi = T64(self.Phi) # (n_dof, N)
147 PhiT = Phi.T.contiguous()
148 w2 = T64(self.omega**2)
149 al = T64(self.alpha)
150 eta, sigma0 = p["eta"], p["sigma0"]
151 ms, kc, v0 = p["stick_mass"], p["contact_stiffness"], fem.v0
152 # exact two-step recurrence of the damped linear oscillators
153 disc = np.sqrt(self.alpha**2 - self.omega**2 + 0j)
154 E1, E2 = np.exp((-self.alpha + disc) * dt), np.exp((-self.alpha - disc) * dt)
155 a1 = T64((E1 + E2).real)
156 a2 = T64(-(E1 * E2).real)
157 b = T64((1.0 - (E1 + E2).real + (E1 * E2).real) / self.omega**2)
158 # Rayleigh integral: per-step weight table G[d, k] from the modes at the Gauss points
159 tau = fem.ray_tau.reshape(-1)
160 cq = fem.ray_c.reshape(-1)
161 tau_steps = tau / dt
162 d0 = np.floor(tau_steps).astype(np.int64)
163 fr = tau_steps - d0
164 lag = np.stack([-fr * (fr - 1.0) * (fr - 2.0) / 6.0, (fr + 1.0) * (fr - 1.0) * (fr - 2.0) / 2.0,
165 -(fr + 1.0) * fr * (fr - 2.0) / 2.0, (fr + 1.0) * fr * (fr - 1.0) / 6.0], axis=1)
166 dmax = int(d0.max()) + 2
167 L = dmax + 1
168 Nw = T64(fem.Nw) # (G, 16)
169 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
170 Gt = torch.zeros((L, N), dtype=torch.float64, device=dev)
171 for j, off in enumerate([-1, 0, 1, 2]):
172 Gt.index_add_(0, T64i(d0 + off), T64(cq * lag[:, j])[:, None] * Psi)
173 del Psi
174 # sticks
175 nst = len(fem.strikes)
176 gains = T64(np.stack([s["phi"] for s in fem.strikes])) @ Phi # (nst, N): footprint average of the modes
177 n_on = T64i([int(round(s["t"] / dt)) for s in fem.strikes])
178 z = torch.zeros(nst, dtype=torch.float64, device=dev)
179 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
180 active = torch.zeros(nst, dtype=torch.float64, device=dev)
181 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
182 imp = torch.zeros(nst, dtype=torch.float64, device=dev)
183 dur = torch.zeros(nst, dtype=torch.float64, device=dev)
184 fmax_t = torch.zeros(nst, dtype=torch.float64, device=dev)
185 # state
186 a = torch.zeros(N, dtype=torch.float64, device=dev)
187 a_prev = torch.zeros(N, dtype=torch.float64, device=dev)
188 q = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
189 q_prev = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
190 ring = torch.zeros((L, N), dtype=torch.float64, device=dev)
191 out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
192 nstep = torch.zeros((), dtype=torch.int64, device=dev)
193 ar = torch.arange(L, device=dev)
194 fld = None
195 if field_grid is not None:
196 fx, fy, ft = field_grid
197 idx, wts, inside = fem.field_interpolator(fx, fy)
198 Wg = (T64(wts)[:, :, None] * Phi[T64i(idx)]).sum(1) # (P, N)
199 fld = dict(Wg=Wg, inside=inside, every=int(round(fs / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])
201 def step():
202 n = nstep
203 q.copy_(Phi @ a)
204 v = (q - q_prev) / dt
205 adot = (a - a_prev) / dt
206 # sticks
207 wbar = gains @ a
208 wbar_dot = gains @ adot
209 starting = (n == n_on).to(torch.float64)
210 z.copy_(z * (1 - starting) + wbar * starting)
211 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
212 active.copy_(torch.maximum(active, starting))
213 pen = torch.clamp(wbar - z, min=0.0)
214 F = active * kc * pen * torch.sqrt(pen)
215 Fst.copy_(F)
216 imp.add_(dt * F)
217 dur.add_(dt * (F > 0).to(torch.float64))
218 fmax_t.copy_(torch.maximum(fmax_t, F))
219 # modal force: nonlinear remainder of the internal force plus the sticks
220 f = force(q, v)
221 Fm = -(PhiT @ f) + w2 * a + eta * w2 * adot - F @ gains
222 a_next = a1 * a + a2 * a_prev + b * Fm
223 add = Fm - w2 * a - al * ((a_next - a_prev) / dt)
224 ring.index_copy_(0, (n % L).view(1), add.view(1, -1))
225 zd.add_(dt * F / ms)
226 z.add_(dt * zd)
227 q_prev.copy_(q)
228 a_prev.copy_(a)
229 a.copy_(a_next)
230 nstep.add_(1)
232 def sample(k):
233 # sample k is at step k m: the first step below computes and stores the acceleration there, the output sample
234 # folds the table with the ring (delays counted back from that step), then the remaining steps of the sample
235 step()
236 idx = (nstep - 1 - ar) % L
237 out.index_copy_(0, k.view(1), torch.sum(Gt * ring[idx]).view(1))
238 for _ in range(m - 1):
239 step()
241 graph = None
242 k_t = torch.zeros((), dtype=torch.int64, device=dev)
244 def sample_graph():
245 sample(k_t)
247 def reset():
248 for t in (z, zd, active, Fst, imp, dur, fmax_t, a, a_prev, q, q_prev, ring, out, nstep, k_t):
249 t.zero_()
251 if use_graph and device.startswith("cuda"):
252 s = torch.cuda.Stream()
253 s.wait_stream(torch.cuda.current_stream())
254 with torch.cuda.stream(s):
255 for _ in range(2):
256 sample_graph()
257 torch.cuda.current_stream().wait_stream(s)
258 reset()
259 graph = torch.cuda.CUDAGraph()
260 with torch.cuda.graph(graph):
261 sample_graph()
262 reset()
263 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 ''}")
265 def take_frame():
266 fr_ = (fld["Wg"] @ a).cpu().numpy().astype(np.float32)
267 fr_[~fld["inside"]] = np.nan
268 fld["frames"].append(fr_)
270 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()],
271 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()])
272 t_start = time.perf_counter()
273 last_report = t_start
274 if fld is not None:
275 take_frame()
276 for k in range(nsamp):
277 k_t.fill_(k)
278 if graph is not None:
279 graph.replay()
280 else:
281 sample(k_t)
282 if fld is not None and (k + 1) % fld["every"] == 0 and len(fld["frames"]) < fld["nt"]:
283 take_frame()
284 now = time.perf_counter()
285 if (now - last_report >= report or k == nsamp - 1) and k > 0:
286 last_report = now
287 el = now - t_start
288 an = a.cpu().numpy()
289 if not np.all(np.isfinite(an)):
290 raise RuntimeError(f"solution diverged at sample {k} (t = {k / fs:.4f} s); increase m")
291 sig = out.cpu().numpy()
292 prog = dict(tag=tag, step=(k + 1) * m, steps=nsteps, t=round((k + 1) / fs, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
293 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()),
294 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()],
295 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)
296 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} "
297 f"peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
298 if out_dir is not None:
299 with open(out_dir / f"{tag}.progress.json", "w") as fh:
300 json.dump(prog, fh, indent=1)
301 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
302 elapsed = time.perf_counter() - t_start
303 sig = out.cpu().numpy()
304 if not np.all(np.isfinite(sig)):
305 raise RuntimeError("output is not finite; increase m")
306 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
307 return sig, fields, dict(m=m, dt=dt, elapsed=elapsed, sticks=stick_stats())
310def main():
311 ch = json.load(open(sys.argv[1]))
312 out_dir = Path(sys.argv[2])
313 args = sys.argv[3:]
315 def take(flag, conv, default=None):
316 nonlocal args
317 if flag in args:
318 i = args.index(flag)
319 val = conv(args[i + 1])
320 args = args[:i] + args[i + 2 :]
321 return val
322 return default
324 spacing = take("--spacing", float, 0.008)
325 f_max = take("--fmax", float, 20000.0)
326 m = take("--m", int, 4)
327 seconds = take("--seconds", float)
328 device = take("--device", str, "cpu")
329 report = take("--report", float, 60.0)
330 speed = take("--speed", float)
331 tag = take("--tag", str)
332 inplane_scale = take("--inplane-scale", float, 1.0)
333 fields = "--fields" in args
334 compile_ = "--no-compile" not in args
335 use_graph = "--no-graph" not in args
336 out = ch["output"]
337 fs = out["sample_rate_hz"]
338 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
339 tag = tag or f"modal-g{spacing * 1e3:g}-f{f_max / 1e3:g}k-m{m}"
340 log = lambda s: print(s, file=sys.stderr, flush=True)
341 out_dir.mkdir(parents=True, exist_ok=True)
342 t0 = time.perf_counter()
343 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}")
344 fem = ShellFEM(ch, spacing, log=log, speed=speed, no_drift=True, inplane_scale=inplane_scale)
345 model = ModalShell(fem, f_max, log=log)
346 t1 = time.perf_counter()
347 scale = out.get("signal_scale") or 1.0
348 grid = None
349 if fields:
350 fg = ch["fields"]
351 nt, ny, nx = fg["shape"]
352 R = ch["parameters"]["disk_radius"]
353 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
354 sig, fld, info = model.run(fs, nsamp, m, device=device, field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None,
355 report=report, out_dir=out_dir, tag=tag, scale=scale, compile_=compile_, use_graph=use_graph)
356 total = time.perf_counter() - t0
357 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
358 meta = {
359 "solution_id": tag,
360 "output": f"{out_dir.name}/{tag}.wav",
361 "method": "modes of the bicubic Hermite finite element system on a Cartesian grid masked to the annulus (lumped mass), truncated at a cutoff frequency, "
362 "exact linear propagation with the nonlinear force evaluated at the Gauss points, Hertz stick, Rayleigh integral by a retarded weight table",
363 "parameters": {"grid_spacing_m": spacing, "cells": fem.n_cells, "n_dof": fem.n_dof, "n_modes": model.n_modes, "mode_cutoff_hz": f_max,
364 "highest_mode_hz": round(float(model.omega.max() / 2 / np.pi), 1), "oversampling": m, "dt_s": info["dt"], "inplane_derivative_mass_scale": inplane_scale},
365 "timing_seconds": round(info["elapsed"], 3),
366 "setup_seconds": round(t1 - t0, 3),
367 "hardware": hardware_string(device),
368 "peak_si": float(np.max(np.abs(sig))),
369 "sticks": info["sticks"],
370 **({"stick_speed_override": speed} if speed is not None else {}),
371 }
372 if fld is not None:
373 import h5py
375 fg = ch["fields"]
376 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
377 fh.attrs["challenge_id"] = ch["challenge_id"]
378 fh.attrs["solution_id"] = tag
379 fh.attrs["quantity"] = fg["quantity"]
380 fh.attrs["units"] = fg["units"]
381 fh.create_dataset("x", data=grid[0])
382 fh.create_dataset("y", data=grid[1])
383 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
384 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
385 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
386 with open(out_dir / f"{tag}.meta.json", "w") as fh:
387 json.dump(meta, fh, indent=1)
388 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
391if __name__ == "__main__":
392 main()