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 self.omega, self.Phi = sliced_modes(K, fem.Mdiag, fem.rigid, f_max, log=log)
107 self.n_modes = len(self.omega)
108 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")
109 p = fem.p
110 self.alpha = p["sigma0"] + 0.5 * p["eta"] * self.omega**2
112 def run(self, fs: int, nsamp: int, m: int, device: str = "cpu", field_grid=None, field_rate=None, report: float = 60.0,
113 out_dir: Path | None = None, tag: str = "modal", scale: float = 1.0, compile_: bool = True, use_graph: bool = True):
114 import torch
116 fem, p = self.fem, self.fem.p
117 t_setup = time.perf_counter()
118 dev = torch.device(device)
119 force, ops = fem.build_force(torch, dev, compile_)
120 T64, T64i = ops["T64"], ops["T64i"]
121 N = self.n_modes
122 dt = 1.0 / (fs * m)
123 nsteps = nsamp * m
124 Phi = T64(self.Phi) # (n_dof, N)
125 PhiT = Phi.T.contiguous()
126 w2 = T64(self.omega**2)
127 al = T64(self.alpha)
128 eta, sigma0 = p["eta"], p["sigma0"]
129 ms, kc, v0 = p["stick_mass"], p["contact_stiffness"], fem.v0
130 # exact two-step recurrence of the damped linear oscillators
131 disc = np.sqrt(self.alpha**2 - self.omega**2 + 0j)
132 E1, E2 = np.exp((-self.alpha + disc) * dt), np.exp((-self.alpha - disc) * dt)
133 a1 = T64((E1 + E2).real)
134 a2 = T64(-(E1 * E2).real)
135 b = T64((1.0 - (E1 + E2).real + (E1 * E2).real) / self.omega**2)
136 # Rayleigh integral: per-step weight table G[d, k] from the modes at the Gauss points
137 tau = fem.ray_tau.reshape(-1)
138 cq = fem.ray_c.reshape(-1)
139 tau_steps = tau / dt
140 d0 = np.floor(tau_steps).astype(np.int64)
141 fr = tau_steps - d0
142 lag = np.stack([-fr * (fr - 1.0) * (fr - 2.0) / 6.0, (fr + 1.0) * (fr - 1.0) * (fr - 2.0) / 2.0,
143 -(fr + 1.0) * fr * (fr - 2.0) / 2.0, (fr + 1.0) * fr * (fr - 1.0) / 6.0], axis=1)
144 dmax = int(d0.max()) + 2
145 L = dmax + 1
146 Nw = T64(fem.Nw) # (G, 16)
147 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
148 Gt = torch.zeros((L, N), dtype=torch.float64, device=dev)
149 for j, off in enumerate([-1, 0, 1, 2]):
150 Gt.index_add_(0, T64i(d0 + off), T64(cq * lag[:, j])[:, None] * Psi)
151 del Psi
152 # sticks
153 nst = len(fem.strikes)
154 gains = T64(np.stack([s["phi"] for s in fem.strikes])) @ Phi # (nst, N): footprint average of the modes
155 n_on = T64i([int(round(s["t"] / dt)) for s in fem.strikes])
156 z = torch.zeros(nst, dtype=torch.float64, device=dev)
157 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
158 active = torch.zeros(nst, dtype=torch.float64, device=dev)
159 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
160 imp = torch.zeros(nst, dtype=torch.float64, device=dev)
161 dur = torch.zeros(nst, dtype=torch.float64, device=dev)
162 fmax_t = torch.zeros(nst, dtype=torch.float64, device=dev)
163 # state
164 a = torch.zeros(N, dtype=torch.float64, device=dev)
165 a_prev = torch.zeros(N, dtype=torch.float64, device=dev)
166 q = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
167 q_prev = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
168 ring = torch.zeros((L, N), dtype=torch.float64, device=dev)
169 out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
170 nstep = torch.zeros((), dtype=torch.int64, device=dev)
171 ar = torch.arange(L, device=dev)
172 fld = None
173 if field_grid is not None:
174 fx, fy, ft = field_grid
175 idx, wts, inside = fem.field_interpolator(fx, fy)
176 Wg = (T64(wts)[:, :, None] * Phi[T64i(idx)]).sum(1) # (P, N)
177 fld = dict(Wg=Wg, inside=inside, every=int(round(fs / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])
179 def step():
180 n = nstep
181 q.copy_(Phi @ a)
182 v = (q - q_prev) / dt
183 adot = (a - a_prev) / dt
184 # sticks
185 wbar = gains @ a
186 wbar_dot = gains @ adot
187 starting = (n == n_on).to(torch.float64)
188 z.copy_(z * (1 - starting) + wbar * starting)
189 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
190 active.copy_(torch.maximum(active, starting))
191 pen = torch.clamp(wbar - z, min=0.0)
192 F = active * kc * pen * torch.sqrt(pen)
193 Fst.copy_(F)
194 imp.add_(dt * F)
195 dur.add_(dt * (F > 0).to(torch.float64))
196 fmax_t.copy_(torch.maximum(fmax_t, F))
197 # modal force: nonlinear remainder of the internal force plus the sticks
198 f = force(q, v)
199 Fm = -(PhiT @ f) + w2 * a + eta * w2 * adot - F @ gains
200 a_next = a1 * a + a2 * a_prev + b * Fm
201 add = Fm - w2 * a - al * ((a_next - a_prev) / dt)
202 ring[n % L] = add
203 zd.add_(dt * F / ms)
204 z.add_(dt * zd)
205 q_prev.copy_(q)
206 a_prev.copy_(a)
207 a.copy_(a_next)
208 nstep.add_(1)
210 def sample(k):
211 # sample k is at step k m: the first step below computes and stores the acceleration there, the output sample
212 # folds the table with the ring (delays counted back from that step), then the remaining steps of the sample
213 step()
214 idx = (nstep - 1 - ar) % L
215 out[k] = torch.sum(Gt * ring[idx])
216 for _ in range(m - 1):
217 step()
219 graph = None
220 k_t = torch.zeros((), dtype=torch.int64, device=dev)
222 def sample_graph():
223 sample(k_t)
225 def reset():
226 for t in (z, zd, active, Fst, imp, dur, fmax_t, a, a_prev, q, q_prev, ring, out, nstep, k_t):
227 t.zero_()
229 if use_graph and device.startswith("cuda"):
230 s = torch.cuda.Stream()
231 s.wait_stream(torch.cuda.current_stream())
232 with torch.cuda.stream(s):
233 for _ in range(2):
234 sample_graph()
235 torch.cuda.current_stream().wait_stream(s)
236 reset()
237 graph = torch.cuda.CUDAGraph()
238 with torch.cuda.graph(graph):
239 sample_graph()
240 reset()
241 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 ''}")
243 def take_frame():
244 fr_ = (fld["Wg"] @ a).cpu().numpy().astype(np.float32)
245 fr_[~fld["inside"]] = np.nan
246 fld["frames"].append(fr_)
248 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()],
249 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()])
250 t_start = time.perf_counter()
251 last_report = t_start
252 if fld is not None:
253 take_frame()
254 for k in range(nsamp):
255 if graph is not None:
256 k_t.fill_(k)
257 graph.replay()
258 else:
259 sample(k)
260 if fld is not None and (k + 1) % fld["every"] == 0 and len(fld["frames"]) < fld["nt"]:
261 take_frame()
262 now = time.perf_counter()
263 if (now - last_report >= report or k == nsamp - 1) and k > 0:
264 last_report = now
265 el = now - t_start
266 an = a.cpu().numpy()
267 if not np.all(np.isfinite(an)):
268 raise RuntimeError(f"solution diverged at sample {k} (t = {k / fs:.4f} s); increase m")
269 sig = out.cpu().numpy()
270 prog = dict(tag=tag, step=(k + 1) * m, steps=nsteps, t=round((k + 1) / fs, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
271 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()),
272 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()],
273 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)
274 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} "
275 f"peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
276 if out_dir is not None:
277 with open(out_dir / f"{tag}.progress.json", "w") as fh:
278 json.dump(prog, fh, indent=1)
279 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
280 elapsed = time.perf_counter() - t_start
281 sig = out.cpu().numpy()
282 if not np.all(np.isfinite(sig)):
283 raise RuntimeError("output is not finite; increase m")
284 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
285 return sig, fields, dict(m=m, dt=dt, elapsed=elapsed, sticks=stick_stats())
288def main():
289 ch = json.load(open(sys.argv[1]))
290 out_dir = Path(sys.argv[2])
291 args = sys.argv[3:]
293 def take(flag, conv, default=None):
294 nonlocal args
295 if flag in args:
296 i = args.index(flag)
297 val = conv(args[i + 1])
298 args = args[:i] + args[i + 2 :]
299 return val
300 return default
302 spacing = take("--spacing", float, 0.008)
303 f_max = take("--fmax", float, 20000.0)
304 m = take("--m", int, 4)
305 seconds = take("--seconds", float)
306 device = take("--device", str, "cpu")
307 report = take("--report", float, 60.0)
308 speed = take("--speed", float)
309 tag = take("--tag", str)
310 inplane_scale = take("--inplane-scale", float, 16.0)
311 fields = "--fields" in args
312 compile_ = "--no-compile" not in args
313 use_graph = "--no-graph" not in args
314 out = ch["output"]
315 fs = out["sample_rate_hz"]
316 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
317 tag = tag or f"modal-g{spacing * 1e3:g}-f{f_max / 1e3:g}k-m{m}"
318 log = lambda s: print(s, file=sys.stderr, flush=True)
319 out_dir.mkdir(parents=True, exist_ok=True)
320 t0 = time.perf_counter()
321 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}")
322 fem = ShellFEM(ch, spacing, log=log, speed=speed, no_drift=True, inplane_scale=inplane_scale)
323 model = ModalShell(fem, f_max, log=log)
324 t1 = time.perf_counter()
325 scale = out.get("signal_scale") or 1.0
326 grid = None
327 if fields:
328 fg = ch["fields"]
329 nt, ny, nx = fg["shape"]
330 R = ch["parameters"]["disk_radius"]
331 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
332 sig, fld, info = model.run(fs, nsamp, m, device=device, field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None,
333 report=report, out_dir=out_dir, tag=tag, scale=scale, compile_=compile_, use_graph=use_graph)
334 total = time.perf_counter() - t0
335 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
336 meta = {
337 "solution_id": tag,
338 "output": f"{out_dir.name}/{tag}.wav",
339 "method": "modes of the bicubic Hermite finite element system on a Cartesian grid masked to the annulus (lumped mass), truncated at a cutoff frequency, "
340 "exact linear propagation with the nonlinear force evaluated at the Gauss points, Hertz stick, Rayleigh integral by a retarded weight table",
341 "parameters": {"grid_spacing_m": spacing, "cells": fem.n_cells, "n_dof": fem.n_dof, "n_modes": model.n_modes, "mode_cutoff_hz": f_max,
342 "highest_mode_hz": round(float(model.omega.max() / 2 / np.pi), 1), "oversampling": m, "dt_s": info["dt"], "inplane_derivative_mass_scale": inplane_scale},
343 "timing_seconds": round(info["elapsed"], 3),
344 "setup_seconds": round(t1 - t0, 3),
345 "hardware": hardware_string(device),
346 "peak_si": float(np.max(np.abs(sig))),
347 "sticks": info["sticks"],
348 **({"stick_speed_override": speed} if speed is not None else {}),
349 }
350 if fld is not None:
351 import h5py
353 fg = ch["fields"]
354 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
355 fh.attrs["challenge_id"] = ch["challenge_id"]
356 fh.attrs["solution_id"] = tag
357 fh.attrs["quantity"] = fg["quantity"]
358 fh.attrs["units"] = fg["units"]
359 fh.create_dataset("x", data=grid[0])
360 fh.create_dataset("y", data=grid[1])
361 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
362 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
363 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
364 with open(out_dir / f"{tag}.meta.json", "w") as fh:
365 json.dump(meta, fh, indent=1)
366 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
369if __name__ == "__main__":
370 main()