/ challenges-old / cymbal2-solutions
Sign in
challenges-old / cymbal2-solutions
cymbal2-solutions / solvers / hermite-fem / modal_shell.py
393 lines · 20.2 KBCodeBlameHistory
38bdd34cymbal2-solutions: Hermite finite element solver (direct and modal) for the cymbal2 challenges, site builder, workflowJeremy Magland 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()
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 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]
38bdd34cymbal2-solutions: Hermite finite element solver (direct and modal) for the cymbal2 challenges, site builder, workflowJeremy Magland 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())
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)
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 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)"))
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()
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 109 cache_dir = Path(__file__).resolve().parent / "cache"
110 cache_dir.mkdir(exist_ok=True)
111 import hashlib
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 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]
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 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:
7044128Clamped mount: fixed unknowns at the inner staircase edge in the FEM, free-unknown eigenproblem in the modal solverJeremy Magland 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)
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 128 np.savez(cache, omega=self.omega, Phi=self.Phi)
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
f150696Output with the area mean of the acceleration removed; experiment harnessJeremy Magland 170 Psi = Psi - (T64(np.tile(fem.wq, fem.n_cells) / fem.area) @ Psi)[None, :] # remove each mode's area mean: the output of MODEL.md
172 for j, off in enumerate([-1, 0, 1, 2]):
173 Gt.index_add_(0, T64i(d0 + off), T64(cq * lag[:, j])[:, None] * Psi)
174 del Psi
175 # sticks
176 nst = len(fem.strikes)
177 gains = T64(np.stack([s["phi"] for s in fem.strikes])) @ Phi # (nst, N): footprint average of the modes
178 n_on = T64i([int(round(s["t"] / dt)) for s in fem.strikes])
179 z = torch.zeros(nst, dtype=torch.float64, device=dev)
180 zd = torch.zeros(nst, dtype=torch.float64, device=dev)
181 active = torch.zeros(nst, dtype=torch.float64, device=dev)
182 Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
183 imp = torch.zeros(nst, dtype=torch.float64, device=dev)
184 dur = torch.zeros(nst, dtype=torch.float64, device=dev)
185 fmax_t = torch.zeros(nst, dtype=torch.float64, device=dev)
186 # state
187 a = torch.zeros(N, dtype=torch.float64, device=dev)
188 a_prev = torch.zeros(N, dtype=torch.float64, device=dev)
189 q = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
190 q_prev = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
191 ring = torch.zeros((L, N), dtype=torch.float64, device=dev)
192 out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
193 nstep = torch.zeros((), dtype=torch.int64, device=dev)
194 ar = torch.arange(L, device=dev)
195 fld = None
196 if field_grid is not None:
197 fx, fy, ft = field_grid
198 idx, wts, inside = fem.field_interpolator(fx, fy)
199 Wg = (T64(wts)[:, :, None] * Phi[T64i(idx)]).sum(1) # (P, N)
200 fld = dict(Wg=Wg, inside=inside, every=int(round(fs / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])
202 def step():
203 n = nstep
204 q.copy_(Phi @ a)
205 v = (q - q_prev) / dt
206 adot = (a - a_prev) / dt
207 # sticks
208 wbar = gains @ a
209 wbar_dot = gains @ adot
210 starting = (n == n_on).to(torch.float64)
211 z.copy_(z * (1 - starting) + wbar * starting)
212 zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
213 active.copy_(torch.maximum(active, starting))
214 pen = torch.clamp(wbar - z, min=0.0)
215 F = active * kc * pen * torch.sqrt(pen)
216 Fst.copy_(F)
217 imp.add_(dt * F)
218 dur.add_(dt * (F > 0).to(torch.float64))
219 fmax_t.copy_(torch.maximum(fmax_t, F))
220 # modal force: nonlinear remainder of the internal force plus the sticks
221 f = force(q, v)
222 Fm = -(PhiT @ f) + w2 * a + eta * w2 * adot - F @ gains
223 a_next = a1 * a + a2 * a_prev + b * Fm
224 add = Fm - w2 * a - al * ((a_next - a_prev) / dt)
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 225 ring.index_copy_(0, (n % L).view(1), add.view(1, -1))
227 z.add_(dt * zd)
228 q_prev.copy_(q)
229 a_prev.copy_(a)
230 a.copy_(a_next)
231 nstep.add_(1)
233 def sample(k):
234 # sample k is at step k m: the first step below computes and stores the acceleration there, the output sample
235 # folds the table with the ring (delays counted back from that step), then the remaining steps of the sample
236 step()
237 idx = (nstep - 1 - ar) % L
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 238 out.index_copy_(0, k.view(1), torch.sum(Gt * ring[idx]).view(1))
240 step()
242 graph = None
243 k_t = torch.zeros((), dtype=torch.int64, device=dev)
245 def sample_graph():
246 sample(k_t)
248 def reset():
249 for t in (z, zd, active, Fst, imp, dur, fmax_t, a, a_prev, q, q_prev, ring, out, nstep, k_t):
250 t.zero_()
252 if use_graph and device.startswith("cuda"):
253 s = torch.cuda.Stream()
254 s.wait_stream(torch.cuda.current_stream())
255 with torch.cuda.stream(s):
256 for _ in range(2):
257 sample_graph()
258 torch.cuda.current_stream().wait_stream(s)
259 reset()
260 graph = torch.cuda.CUDAGraph()
261 with torch.cuda.graph(graph):
262 sample_graph()
263 reset()
264 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 ''}")
266 def take_frame():
267 fr_ = (fld["Wg"] @ a).cpu().numpy().astype(np.float32)
268 fr_[~fld["inside"]] = np.nan
269 fld["frames"].append(fr_)
271 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()],
272 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()])
273 t_start = time.perf_counter()
274 last_report = t_start
275 if fld is not None:
276 take_frame()
277 for k in range(nsamp):
280 graph.replay()
281 else:
38bdd34cymbal2-solutions: Hermite finite element solver (direct and modal) for the cymbal2 challenges, site builder, workflowJeremy Magland 283 if fld is not None and (k + 1) % fld["every"] == 0 and len(fld["frames"]) < fld["nt"]:
284 take_frame()
285 now = time.perf_counter()
286 if (now - last_report >= report or k == nsamp - 1) and k > 0:
287 last_report = now
288 el = now - t_start
289 an = a.cpu().numpy()
290 if not np.all(np.isfinite(an)):
291 raise RuntimeError(f"solution diverged at sample {k} (t = {k / fs:.4f} s); increase m")
292 sig = out.cpu().numpy()
293 prog = dict(tag=tag, step=(k + 1) * m, steps=nsteps, t=round((k + 1) / fs, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
294 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()),
295 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()],
296 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)
297 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} "
298 f"peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
299 if out_dir is not None:
300 with open(out_dir / f"{tag}.progress.json", "w") as fh:
301 json.dump(prog, fh, indent=1)
302 wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
303 elapsed = time.perf_counter() - t_start
304 sig = out.cpu().numpy()
305 if not np.all(np.isfinite(sig)):
306 raise RuntimeError("output is not finite; increase m")
307 fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
308 return sig, fields, dict(m=m, dt=dt, elapsed=elapsed, sticks=stick_stats())
311def main():
312 ch = json.load(open(sys.argv[1]))
313 out_dir = Path(sys.argv[2])
314 args = sys.argv[3:]
316 def take(flag, conv, default=None):
317 nonlocal args
318 if flag in args:
319 i = args.index(flag)
320 val = conv(args[i + 1])
321 args = args[:i] + args[i + 2 :]
322 return val
323 return default
325 spacing = take("--spacing", float, 0.008)
326 f_max = take("--fmax", float, 20000.0)
327 m = take("--m", int, 4)
328 seconds = take("--seconds", float)
329 device = take("--device", str, "cpu")
330 report = take("--report", float, 60.0)
331 speed = take("--speed", float)
332 tag = take("--tag", str)
50128a6Laplacian loss term, rigid constraint by default, geometry-keyed mode cache, README on both solversJeremy Magland 333 inplane_scale = take("--inplane-scale", float, 1.0)
335 compile_ = "--no-compile" not in args
336 use_graph = "--no-graph" not in args
337 out = ch["output"]
338 fs = out["sample_rate_hz"]
339 nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
340 tag = tag or f"modal-g{spacing * 1e3:g}-f{f_max / 1e3:g}k-m{m}"
341 log = lambda s: print(s, file=sys.stderr, flush=True)
342 out_dir.mkdir(parents=True, exist_ok=True)
343 t0 = time.perf_counter()
344 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}")
345 fem = ShellFEM(ch, spacing, log=log, speed=speed, no_drift=True, inplane_scale=inplane_scale)
346 model = ModalShell(fem, f_max, log=log)
347 t1 = time.perf_counter()
348 scale = out.get("signal_scale") or 1.0
349 grid = None
350 if fields:
351 fg = ch["fields"]
352 nt, ny, nx = fg["shape"]
353 R = ch["parameters"]["disk_radius"]
354 grid = (np.linspace(-R, R, nx), np.linspace(-R, R, ny), np.arange(nt) / fg["frame_rate_hz"])
355 sig, fld, info = model.run(fs, nsamp, m, device=device, field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None,
356 report=report, out_dir=out_dir, tag=tag, scale=scale, compile_=compile_, use_graph=use_graph)
357 total = time.perf_counter() - t0
358 wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
359 meta = {
360 "solution_id": tag,
361 "output": f"{out_dir.name}/{tag}.wav",
362 "method": "modes of the bicubic Hermite finite element system on a Cartesian grid masked to the annulus (lumped mass), truncated at a cutoff frequency, "
363 "exact linear propagation with the nonlinear force evaluated at the Gauss points, Hertz stick, Rayleigh integral by a retarded weight table",
364 "parameters": {"grid_spacing_m": spacing, "cells": fem.n_cells, "n_dof": fem.n_dof, "n_modes": model.n_modes, "mode_cutoff_hz": f_max,
365 "highest_mode_hz": round(float(model.omega.max() / 2 / np.pi), 1), "oversampling": m, "dt_s": info["dt"], "inplane_derivative_mass_scale": inplane_scale},
366 "timing_seconds": round(info["elapsed"], 3),
367 "setup_seconds": round(t1 - t0, 3),
368 "hardware": hardware_string(device),
369 "peak_si": float(np.max(np.abs(sig))),
370 "sticks": info["sticks"],
371 **({"stick_speed_override": speed} if speed is not None else {}),
372 }
373 if fld is not None:
374 import h5py
376 fg = ch["fields"]
377 with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
378 fh.attrs["challenge_id"] = ch["challenge_id"]
379 fh.attrs["solution_id"] = tag
380 fh.attrs["quantity"] = fg["quantity"]
381 fh.attrs["units"] = fg["units"]
382 fh.create_dataset("x", data=grid[0])
383 fh.create_dataset("y", data=grid[1])
384 fh.create_dataset("t", data=grid[2][: fld.shape[0]])
385 fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
386 meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
387 with open(out_dir / f"{tag}.meta.json", "w") as fh:
388 json.dump(meta, fh, indent=1)
389 log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")
392if __name__ == "__main__":
393 main()
↑↓ move↵openescclose