"""The time-stepping driver of the full model. A run advances in batches of one movie frame (1 / fps s): each batch steps the kernel, records the transverse velocity at the audio sample instants into a ring buffer, and at its end forms the batch's audio samples from the retarded-time Rayleigh sum over the plate and reads back a frame of the transverse displacement. The recording starts at the impact, with the plate at rest in its stress-free shape. The device is a CUDA GPU when one is available, else Apple's MPS, else the CPU. The step is many small tensor operations, which ``torch.compile`` fuses into a few kernels (the default; about 6 times faster than the plain step on the CPU, where a 16 inch cymbal at 8 mm then runs at about 90 s per second of sound, and far more on a GPU, where the plain step would be bound by kernel launches). """ from __future__ import annotations import math import time from typing import Literal import numpy as np import torch from ..constants import CFL from ..loss import fit_sigma, in_plane_rate from ..params import Listener, Strike from ..plate import Plate from ..sound import Sound from .kernel import Static, static_data, step from .reference import max_angular_frequency def default_device() -> torch.device: """CUDA when available, else MPS, else the CPU.""" if torch.cuda.is_available(): return torch.device("cuda") if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") class Simulation: """A full-model run that can be advanced batch by batch. Parameters mirror ``simulate``. After construction, ``run`` advances the simulation and ``sound`` returns what has been recorded so far; ``duration`` may be raised and ``run`` called again to extend a finished run. """ def __init__( self, plate: Plate, strike: Strike = Strike(), listener: Listener = Listener(), *, duration: float = 4.0, fs: int = 48000, fps: float = 120.0, movie: bool = True, device=None, dtype: Literal["float32", "float64"] | torch.dtype = "float32", compile: bool | str = True, loss: tuple[float, float, float] | None = None, remove_rigid: bool = True, ): self.plate, self.strike, self.listener = plate, strike, listener self.remove_rigid = remove_rigid self.fs, self.fps, self.duration, self.movie = fs, fps, duration, movie self.device = torch.device(device) if device is not None else default_device() self.dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype g, m = plate.grid, plate.material # Stable time step, rounded so that an integer number of steps make one audio sample. self.omega_max = max_angular_frequency(g, m, plate.thickness_range[1]) dt_max = CFL * 2 / self.omega_max self.steps_per_sample = math.ceil(1 / (fs * dt_max)) self.dt = 1 / (fs * self.steps_per_sample) self.samples_per_batch = max(1, round(fs / fps)) self.steps_per_batch = self.samples_per_batch * self.steps_per_sample if loss is None: fit = fit_sigma(plate) loss = (fit.sigma0, fit.sigma1, in_plane_rate(plate)) self.loss = loss self.static: Static = static_data(plate, strike, self.dt, *loss, self.device, self.dtype) # Listener: weights and delays in samples relative to the nearest node; the ring buffer covers the spread. weight, delay = plate.listener_terms(listener) sel = weight > 0 d = np.where(sel, (delay - delay[sel].min()) * fs, 0.0) self.hist_len = self.samples_per_batch + int(math.ceil(d.max())) + 4 self._lw = torch.as_tensor(weight.reshape(-1), dtype=self.dtype, device=self.device) self._k0 = torch.as_tensor(np.floor(d).astype(np.int64).reshape(-1), device=self.device) self._fr = torch.as_tensor((d - np.floor(d)).reshape(-1), dtype=self.dtype, device=self.device) self._hist = torch.zeros((self.hist_len, g.n), dtype=self.dtype, device=self.device) self._node_idx = torch.arange(g.n, device=self.device)[None, :] z = lambda: torch.zeros(g.shape, dtype=self.dtype, device=self.device) # noqa: E731 self._state = (z(), z(), z(), z(), z(), z(), z()) self._stick = torch.tensor([0.0, -strike.velocity, 0.0, 0.0, 0.0, 0.0], dtype=self.dtype, device=self.device) self.total_steps = 0 self.samples = 0 self.pressure: list[np.ndarray] = [] self.frames: list[np.ndarray] = [np.zeros(g.shape, dtype=np.float32)] self.stick = {"x": 0.0, "v": -strike.velocity, "force": 0.0, "max_force": 0.0, "withdrawn": False} self.seconds = 0.0 # Mass-orthonormal rigid-body modes of the free plate (three transverse: # translation and two tilts; three in-plane: two translations and the # rotation about the normal), removed from the state after every batch. self._rigid_w, self._rigid_uv = self._rigid_bases() self._mass = torch.as_tensor(plate.node_mass, dtype=self.dtype, device=self.device) self._step = step self._compile_pending = bool(compile) self._compile_mode = compile if isinstance(compile, str) else "default" # On CUDA the steps of one audio sample are captured into a CUDA graph # (see _try_compile) and replayed on static buffers: self._graph, the # buffer that receives the transverse velocity at the sample instant, # and the stream the graph was captured on. self._graph = None self._sample = None # ---- progress ---- @property def target_samples(self) -> int: return round(self.duration * self.fs) @property def done(self) -> bool: return self.samples >= self.target_samples @property def t(self) -> float: """Simulated time since the impact [s].""" return self.total_steps * self.dt def status(self) -> str: return f"t = {self.t:.3f} s, {self.total_steps} steps, {self.samples} samples, stick force {self.stick['force']:.1f} N (max {self.stick['max_force']:.1f})" + (", withdrawn" if self.stick["withdrawn"] else "") # ---- stepping ---- def run_batch(self) -> None: """Advances by one movie frame (or to the end of the recording).""" steps_left = self.target_samples * self.steps_per_sample - self.total_steps steps = min(self.steps_per_batch, steps_left) if steps <= 0: return t0 = time.perf_counter() sps, H = self.steps_per_sample, self.hist_len slots = [] state, stick = self._state, self._stick S = self.static if self._compile_pending: self._try_compile() if self._graph is not None: # One graph replay per audio sample; the static buffers advance in place. for local in range(self.total_steps, self.total_steps + steps, sps): self._graph.replay() slot = (local // sps) % H self._hist[slot] = self._sample.reshape(-1) slots.append(slot) state, stick = self._state, self._stick else: for s in range(steps): local = self.total_steps + s state_and_stick = self._step(S, *state, stick) state, stick = state_and_stick[:7], state_and_stick[7] if local % sps == 0: slot = (local // sps) % H self._hist[slot] = state[6].reshape(-1) slots.append(slot) self._state, self._stick = state, stick self.total_steps += steps if self.remove_rigid: self._remove_rigid() if slots: self.pressure.append(self._audio(slots)) self.samples += len(slots) if self.movie: self.frames.append(state[2].detach().to("cpu", torch.float32).numpy().copy()) st = stick.detach().cpu().numpy() self.stick = {"x": float(st[0]), "v": float(st[1]), "force": float(st[2]), "max_force": float(st[3]), "withdrawn": bool(st[5] > 0)} self.seconds += time.perf_counter() - t0 def _rigid_bases(self): """Rigid-body modes orthonormalized under the mass inner product, as tensors ``(3, ny, nx)`` for w and ``(3, 2, ny, nx)`` for (u, v).""" g = self.plate.grid m = self.plate.node_mass cx, cy = self.plate.center X, Y, one = g.X - cx, g.Y - cy, np.ones(g.shape) z = np.zeros(g.shape) def orthonormal(modes): out = [] for v in modes: v = [np.where(g.mask, c, 0.0) for c in v] for b in out: c = sum((m * bi * vi).sum() for bi, vi in zip(b, v)) v = [vi - c * bi for bi, vi in zip(b, v)] n = math.sqrt(sum((m * vi * vi).sum() for vi in v)) out.append([vi / n for vi in v]) return torch.as_tensor(np.asarray(out), dtype=self.dtype, device=self.device) return orthonormal([[one], [X], [Y]])[:, 0], orthonormal([[one, z], [z, one], [-Y, X]]) def _remove_rigid(self) -> None: """Projects the rigid-body motion of the free plate out of the state, in place. The von Karman strains are not invariant under a finite rotation, and the model conserves neither the angular momentum about the in-plane axes nor, strictly, the in-plane momentum, so a hard strike lets the free plate slowly tilt and spin. A tilt of a few degrees then acts as a shell with that slope and pumps energy into the in-plane grid modes. The rigid modes carry no vibration and, for a free plate, no sound, so they are removed after every batch.""" u, v, w, wlo, du, dv, dw = self._state mass = self._mass for f in (w, dw): c = (mass * f * self._rigid_w).sum(dim=(1, 2)) f.sub_((c[:, None, None] * self._rigid_w).sum(0)) for a, b in ((u, v), (du, dv)): c = (mass * (a * self._rigid_uv[:, 0] + b * self._rigid_uv[:, 1])).sum(dim=(1, 2)) a.sub_((c[:, None, None] * self._rigid_uv[:, 0]).sum(0)) b.sub_((c[:, None, None] * self._rigid_uv[:, 1]).sum(0)) def _try_compile(self) -> None: """Wraps the step in ``torch.compile``, which fuses its some two hundred small tensor operations into a few kernels (about 6 times faster than the plain step on the CPU), and on CUDA captures the steps of one audio sample into a CUDA graph, so that a sample costs one graph launch rather than a Python call per kernel (a further 4 to 5 times). Falls back, with a warning, to whatever part of that works.""" import warnings self._compile_pending = False try: compiled = torch.compile(step, dynamic=False, mode=self._compile_mode) compiled(self.static, *self._state, self._stick) # compiles; the result is discarded and the state is unchanged self._step = compiled except Exception as e: # noqa: BLE001 warnings.warn(f"torch.compile failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the uncompiled step, which is slower", RuntimeWarning, stacklevel=3) self._step = step return if self.device.type != "cuda": return try: self._capture_graph() except Exception as e: # noqa: BLE001 warnings.warn(f"CUDA graph capture failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the compiled step without a graph, which is slower", RuntimeWarning, stacklevel=3) self._graph = None def _capture_graph(self) -> None: """Captures ``steps_per_sample`` compiled steps into a CUDA graph that reads and writes the state in place. The transverse velocity after the first step of the graph, the sample instant, is copied to ``_sample``.""" S, sps = self.static, self.steps_per_sample state = [x.clone() for x in self._state] stick = self._stick.clone() sample = torch.empty_like(state[6]) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): for _ in range(3): # warm-up on a side stream, as capture requires self._step(S, *state, stick) torch.cuda.current_stream().wait_stream(side) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): cur, st = tuple(state), stick for i in range(sps): out = self._step(S, *cur, st) cur, st = out[:7], out[7] if i == 0: sample.copy_(cur[6]) for a, b in zip(state, cur): a.copy_(b) stick.copy_(st) torch.cuda.synchronize() self._state, self._stick = tuple(state), stick self._graph, self._sample = graph, sample def _audio(self, slots: list[int]) -> np.ndarray: """The retarded-time Rayleigh sum for the batch's samples: the acceleration is the first difference of the delayed velocity, the average over one sample interval, which is band-limited and immune to the per-step rounding noise of the force evaluation.""" H = self.hist_len slot = torch.as_tensor(slots, device=self.device)[:, None] i0 = (slot - self._k0[None, :]) % H i1 = (i0 - 1) % H i2 = (i0 - 2) % H v0 = torch.gather(self._hist, 0, i0) v1 = torch.gather(self._hist, 0, i1) v2 = torch.gather(self._hist, 0, i2) fr = self._fr[None, :] a = ((v0 + fr * (v1 - v0)) - (v1 + fr * (v2 - v1))) * self.fs return (a @ self._lw).to("cpu", torch.float64).numpy() def run(self, seconds: float | None = None, progress: bool = False) -> "Simulation": """Advances by ``seconds`` of simulated time (to the end by default). With ``progress`` a line is printed per second of sound.""" target = self.target_samples if seconds is None else min(self.target_samples, self.samples + round(seconds * self.fs)) next_report = self.samples + self.fs while self.samples < target and not self.done: before = self.total_steps self.run_batch() if self.total_steps == before: raise RuntimeError("the run made no progress; the step budget is exhausted before the sample target") if progress and self.samples >= next_report: print(f"{self.t:.2f} s simulated in {self.seconds:.1f} s wall; {self.status()}", flush=True) next_report += self.fs return self def sound(self) -> Sound: """What has been recorded so far.""" pressure = np.concatenate(self.pressure) if self.pressure else np.zeros(0) frames = np.stack(self.frames) if self.movie else None info = { "model": "full", "dt": self.dt, "steps_per_sample": self.steps_per_sample, "steps": self.total_steps, "sigma0": self.loss[0], "sigma1": self.loss[1], "sigma_in": self.loss[2], "peak_force": self.stick["max_force"], "device": str(self.device), "dtype": str(self.dtype).replace("torch.", ""), "seconds": self.seconds, } return Sound(pressure.astype(np.float32), self.fs, self.plate, self.strike, self.listener, frames, self.fps, info) # ---- state access, for tests and experiments ---- def state(self) -> dict[str, np.ndarray]: """The displacement and velocity fields ``u, v, w, du, dv, dw`` on the grid, as float64 arrays.""" u, v, w, _, du, dv, dw = (x.detach().to("cpu", torch.float64).numpy() for x in self._state) return {"u": u, "v": v, "w": w, "du": du, "dv": dv, "dw": dw} def set_state(self, u=None, v=None, w=None, du=None, dv=None, dw=None) -> None: """Overwrites fields of the state (grid arrays); the compensation array of w is reset.""" cur = list(self._state) for k, x in ((0, u), (1, v), (2, w), (4, du), (5, dv), (6, dw)): if x is not None: cur[k] = torch.as_tensor(np.asarray(x, dtype=np.float64), dtype=self.dtype, device=self.device) * self.static.ins cur[3] = torch.zeros_like(cur[2]) self._assign_state(tuple(cur), self._stick) def _assign_state(self, state, stick) -> None: """Replaces the state; into the graph's static buffers when a graph has been captured.""" if self._graph is not None: for a, b in zip(self._state, state): a.copy_(b) self._stick.copy_(stick) else: self._state, self._stick = tuple(state), stick def step_once(self, strike: bool = False) -> None: """Advances the state by one step without recording; without the drumstick unless ``strike``.""" S = self.static if not strike and S.strike: from dataclasses import replace S = replace(S, strike=False) out = step(S, *self._state, self._stick) self._assign_state(out[:7], out[7]) self.total_steps += 1 def simulate( plate: Plate, strike: Strike = Strike(), listener: Listener = Listener(), *, duration: float = 4.0, fs: int = 48000, fps: float = 120.0, movie: bool = False, device=None, dtype: Literal["float32", "float64"] | torch.dtype = "float32", compile: bool | str = True, progress: bool = False, remove_rigid: bool = True, ) -> Sound: """The sound of a strike by the full (geometrically nonlinear) model. ``device`` is a torch device (``"cuda"``, ``"cuda:1"``, ``"mps"``, ``"cpu"``); by default a CUDA GPU when one is available. ``dtype`` is ``"float32"`` (the default, with the rounding safeguards of the kernel) or ``"float64"``. ``compile`` wraps the step in ``torch.compile`` (``True``, the default, with a fallback to the plain step when no compiler is available; ``False``; or a ``torch.compile`` mode string) and, on CUDA, captures the steps of one audio sample into a CUDA graph. ``fps`` is the movie frame rate and also the batch rate of the run. ``loss`` overrides the fitted ``(sigma0, sigma1, sigma_in)`` of the two-parameter loss operator. ``remove_rigid`` projects the rigid-body motion of the free plate out of the state after every batch (see ``Simulation._remove_rigid``); leave it on unless you study that motion itself. """ sim = Simulation(plate, strike, listener, duration=duration, fs=fs, fps=fps, movie=movie, device=device, dtype=dtype, compile=compile, remove_rigid=remove_rigid) sim.run(progress=progress) return sim.sound()