1"""The time-stepping driver of the full model. A run advances in batches of one
2movie frame (1 / fps s): each batch steps the kernel, records the transverse
3velocity at the audio sample instants into a ring buffer, and at its end forms
4the batch's audio samples from the retarded-time Rayleigh sum over the plate
5and reads back a frame of the transverse displacement. The recording starts at
6the impact, with the plate at rest in its stress-free shape.
8The device is a CUDA GPU when one is available, else Apple's MPS, else the CPU.
9The step is many small tensor operations, which ``torch.compile`` fuses into
10a few kernels (the default; about 6 times faster than the plain step on the
11CPU, where a 16 inch cymbal at 8 mm then runs at about 90 s per second of
12sound, and far more on a GPU, where the plain step would be bound by kernel
13launches).
14"""
16from __future__ import annotations
18import math
19import time
20from typing import Literal
22import numpy as np
23import torch
25from ..constants import CFL
26from ..loss import fit_sigma, in_plane_rate
27from ..params import Listener, Strike
28from ..plate import Plate
29from ..sound import Sound
30from .kernel import Static, static_data, step
31from .reference import max_angular_frequency
34def default_device() -> torch.device:
35 """CUDA when available, else MPS, else the CPU."""
36 if torch.cuda.is_available():
37 return torch.device("cuda")
38 if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
39 return torch.device("mps")
40 return torch.device("cpu")
43class Simulation:
44 """A full-model run that can be advanced batch by batch.
46 Parameters mirror ``simulate``. After construction, ``run`` advances the
47 simulation and ``sound`` returns what has been recorded so far; ``duration``
48 may be raised and ``run`` called again to extend a finished run.
49 """
51 def __init__(
52 self,
53 plate: Plate,
54 strike: Strike = Strike(),
55 listener: Listener = Listener(),
56 *,
57 duration: float = 4.0,
58 fs: int = 48000,
59 fps: float = 120.0,
60 movie: bool = True,
61 device=None,
62 dtype: Literal["float32", "float64"] | torch.dtype = "float32",
63 compile: bool | str = True,
64 loss: tuple[float, float, float] | None = None,
65 ):
66 self.plate, self.strike, self.listener = plate, strike, listener
67 self.fs, self.fps, self.duration, self.movie = fs, fps, duration, movie
68 self.device = torch.device(device) if device is not None else default_device()
69 self.dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype
70 g, m = plate.grid, plate.material
72 # Stable time step, rounded so that an integer number of steps make one audio sample.
73 self.omega_max = max_angular_frequency(g, m, plate.thickness_range[1])
74 dt_max = CFL * 2 / self.omega_max
75 self.steps_per_sample = math.ceil(1 / (fs * dt_max))
76 self.dt = 1 / (fs * self.steps_per_sample)
77 self.samples_per_batch = max(1, round(fs / fps))
78 self.steps_per_batch = self.samples_per_batch * self.steps_per_sample
80 if loss is None:
81 fit = fit_sigma(plate)
82 loss = (fit.sigma0, fit.sigma1, in_plane_rate(plate))
83 self.loss = loss
84 self.static: Static = static_data(plate, strike, self.dt, *loss, self.device, self.dtype)
86 # Listener: weights and delays in samples relative to the nearest node; the ring buffer covers the spread.
87 weight, delay = plate.listener_terms(listener)
88 sel = weight > 0
89 d = np.where(sel, (delay - delay[sel].min()) * fs, 0.0)
90 self.hist_len = self.samples_per_batch + int(math.ceil(d.max())) + 4
91 self._lw = torch.as_tensor(weight.reshape(-1), dtype=self.dtype, device=self.device)
92 self._k0 = torch.as_tensor(np.floor(d).astype(np.int64).reshape(-1), device=self.device)
93 self._fr = torch.as_tensor((d - np.floor(d)).reshape(-1), dtype=self.dtype, device=self.device)
94 self._hist = torch.zeros((self.hist_len, g.n), dtype=self.dtype, device=self.device)
95 self._node_idx = torch.arange(g.n, device=self.device)[None, :]
97 z = lambda: torch.zeros(g.shape, dtype=self.dtype, device=self.device) # noqa: E731
98 self._state = (z(), z(), z(), z(), z(), z(), z())
99 self._stick = torch.tensor([0.0, -strike.velocity, 0.0, 0.0, 0.0, 0.0], dtype=self.dtype, device=self.device)
100 self.total_steps = 0
101 self.samples = 0
102 self.pressure: list[np.ndarray] = []
103 self.frames: list[np.ndarray] = [np.zeros(g.shape, dtype=np.float32)]
104 self.stick = {"x": 0.0, "v": -strike.velocity, "force": 0.0, "max_force": 0.0, "withdrawn": False}
105 self.seconds = 0.0
107 self._step = step
108 self._compile_pending = bool(compile)
109 self._compile_mode = compile if isinstance(compile, str) else "default"
110 # On CUDA the steps of one audio sample are captured into a CUDA graph
111 # (see _try_compile) and replayed on static buffers: self._graph, the
112 # buffer that receives the transverse velocity at the sample instant,
113 # and the stream the graph was captured on.
114 self._graph = None
115 self._sample = None
117 # ---- progress ----
120 def target_samples(self) -> int:
121 return round(self.duration * self.fs)
124 def done(self) -> bool:
125 return self.samples >= self.target_samples
128 def t(self) -> float:
129 """Simulated time since the impact [s]."""
130 return self.total_steps * self.dt
132 def status(self) -> str:
133 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 "")
135 # ---- stepping ----
137 def run_batch(self) -> None:
138 """Advances by one movie frame (or to the end of the recording)."""
139 steps_left = self.target_samples * self.steps_per_sample - self.total_steps
140 steps = min(self.steps_per_batch, steps_left)
141 if steps <= 0:
142 return
143 t0 = time.perf_counter()
144 sps, H = self.steps_per_sample, self.hist_len
145 slots = []
146 state, stick = self._state, self._stick
147 S = self.static
148 if self._compile_pending:
149 self._try_compile()
150 if self._graph is not None:
151 # One graph replay per audio sample; the static buffers advance in place.
152 for local in range(self.total_steps, self.total_steps + steps, sps):
153 self._graph.replay()
154 slot = (local // sps) % H
155 self._hist[slot] = self._sample.reshape(-1)
156 slots.append(slot)
157 state, stick = self._state, self._stick
158 else:
159 for s in range(steps):
160 local = self.total_steps + s
161 state_and_stick = self._step(S, *state, stick)
162 state, stick = state_and_stick[:7], state_and_stick[7]
163 if local % sps == 0:
164 slot = (local // sps) % H
165 self._hist[slot] = state[6].reshape(-1)
166 slots.append(slot)
167 self._state, self._stick = state, stick
168 self.total_steps += steps
169 if slots:
170 self.pressure.append(self._audio(slots))
171 self.samples += len(slots)
172 if self.movie:
173 self.frames.append(state[2].detach().to("cpu", torch.float32).numpy().copy())
174 st = stick.detach().cpu().numpy()
175 self.stick = {"x": float(st[0]), "v": float(st[1]), "force": float(st[2]), "max_force": float(st[3]), "withdrawn": bool(st[5] > 0)}
176 self.seconds += time.perf_counter() - t0
178 def _try_compile(self) -> None:
179 """Wraps the step in ``torch.compile``, which fuses its some two hundred
180 small tensor operations into a few kernels (about 6 times faster than
181 the plain step on the CPU), and on CUDA captures the steps of one audio
182 sample into a CUDA graph, so that a sample costs one graph launch
183 rather than a Python call per kernel (a further 4 to 5 times). Falls
184 back, with a warning, to whatever part of that works."""
185 import warnings
187 self._compile_pending = False
188 try:
189 compiled = torch.compile(step, dynamic=False, mode=self._compile_mode)
190 compiled(self.static, *self._state, self._stick) # compiles; the result is discarded and the state is unchanged
191 self._step = compiled
192 except Exception as e: # noqa: BLE001
193 warnings.warn(f"torch.compile failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the uncompiled step, which is slower", RuntimeWarning, stacklevel=3)
194 self._step = step
195 return
196 if self.device.type != "cuda":
197 return
198 try:
199 self._capture_graph()
200 except Exception as e: # noqa: BLE001
201 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)
202 self._graph = None
204 def _capture_graph(self) -> None:
205 """Captures ``steps_per_sample`` compiled steps into a CUDA graph that
206 reads and writes the state in place. The transverse velocity after the
207 first step of the graph, the sample instant, is copied to ``_sample``."""
208 S, sps = self.static, self.steps_per_sample
209 state = [x.clone() for x in self._state]
210 stick = self._stick.clone()
211 sample = torch.empty_like(state[6])
212 side = torch.cuda.Stream()
213 side.wait_stream(torch.cuda.current_stream())
214 with torch.cuda.stream(side):
215 for _ in range(3): # warm-up on a side stream, as capture requires
216 self._step(S, *state, stick)
217 torch.cuda.current_stream().wait_stream(side)
218 torch.cuda.synchronize()
219 graph = torch.cuda.CUDAGraph()
220 with torch.cuda.graph(graph):
221 cur, st = tuple(state), stick
222 for i in range(sps):
223 out = self._step(S, *cur, st)
224 cur, st = out[:7], out[7]
225 if i == 0:
226 sample.copy_(cur[6])
227 for a, b in zip(state, cur):
228 a.copy_(b)
229 stick.copy_(st)
230 torch.cuda.synchronize()
231 self._state, self._stick = tuple(state), stick
232 self._graph, self._sample = graph, sample
234 def _audio(self, slots: list[int]) -> np.ndarray:
235 """The retarded-time Rayleigh sum for the batch's samples: the acceleration
236 is the first difference of the delayed velocity, the average over one
237 sample interval, which is band-limited and immune to the per-step
238 rounding noise of the force evaluation."""
239 H = self.hist_len
240 slot = torch.as_tensor(slots, device=self.device)[:, None]
241 i0 = (slot - self._k0[None, :]) % H
242 i1 = (i0 - 1) % H
243 i2 = (i0 - 2) % H
244 v0 = torch.gather(self._hist, 0, i0)
245 v1 = torch.gather(self._hist, 0, i1)
246 v2 = torch.gather(self._hist, 0, i2)
247 fr = self._fr[None, :]
248 a = ((v0 + fr * (v1 - v0)) - (v1 + fr * (v2 - v1))) * self.fs
249 return (a @ self._lw).to("cpu", torch.float64).numpy()
251 def run(self, seconds: float | None = None, progress: bool = False) -> "Simulation":
252 """Advances by ``seconds`` of simulated time (to the end by default). With ``progress`` a line is printed per second of sound."""
253 target = self.target_samples if seconds is None else min(self.target_samples, self.samples + round(seconds * self.fs))
254 next_report = self.samples + self.fs
255 while self.samples < target and not self.done:
256 before = self.total_steps
257 self.run_batch()
258 if self.total_steps == before:
259 raise RuntimeError("the run made no progress; the step budget is exhausted before the sample target")
260 if progress and self.samples >= next_report:
261 print(f"{self.t:.2f} s simulated in {self.seconds:.1f} s wall; {self.status()}", flush=True)
262 next_report += self.fs
263 return self
265 def sound(self) -> Sound:
266 """What has been recorded so far."""
267 pressure = np.concatenate(self.pressure) if self.pressure else np.zeros(0)
268 frames = np.stack(self.frames) if self.movie else None
269 info = {
270 "model": "full",
271 "dt": self.dt,
272 "steps_per_sample": self.steps_per_sample,
273 "steps": self.total_steps,
274 "sigma0": self.loss[0],
275 "sigma1": self.loss[1],
276 "sigma_in": self.loss[2],
277 "peak_force": self.stick["max_force"],
278 "device": str(self.device),
279 "dtype": str(self.dtype).replace("torch.", ""),
280 "seconds": self.seconds,
281 }
282 return Sound(pressure.astype(np.float32), self.fs, self.plate, self.strike, self.listener, frames, self.fps, info)
284 # ---- state access, for tests and experiments ----
286 def state(self) -> dict[str, np.ndarray]:
287 """The displacement and velocity fields ``u, v, w, du, dv, dw`` on the grid, as float64 arrays."""
288 u, v, w, _, du, dv, dw = (x.detach().to("cpu", torch.float64).numpy() for x in self._state)
289 return {"u": u, "v": v, "w": w, "du": du, "dv": dv, "dw": dw}
291 def set_state(self, u=None, v=None, w=None, du=None, dv=None, dw=None) -> None:
292 """Overwrites fields of the state (grid arrays); the compensation array of w is reset."""
293 cur = list(self._state)
294 for k, x in ((0, u), (1, v), (2, w), (4, du), (5, dv), (6, dw)):
295 if x is not None:
296 cur[k] = torch.as_tensor(np.asarray(x, dtype=np.float64), dtype=self.dtype, device=self.device) * self.static.ins
297 cur[3] = torch.zeros_like(cur[2])
298 self._assign_state(tuple(cur), self._stick)
300 def _assign_state(self, state, stick) -> None:
301 """Replaces the state; into the graph's static buffers when a graph has been captured."""
302 if self._graph is not None:
303 for a, b in zip(self._state, state):
304 a.copy_(b)
305 self._stick.copy_(stick)
306 else:
307 self._state, self._stick = tuple(state), stick
309 def step_once(self, strike: bool = False) -> None:
310 """Advances the state by one step without recording; without the drumstick unless ``strike``."""
311 S = self.static
312 if not strike and S.strike:
313 from dataclasses import replace
315 S = replace(S, strike=False)
316 out = step(S, *self._state, self._stick)
317 self._assign_state(out[:7], out[7])
318 self.total_steps += 1
321def simulate(
322 plate: Plate,
323 strike: Strike = Strike(),
324 listener: Listener = Listener(),
325 *,
326 duration: float = 4.0,
327 fs: int = 48000,
328 fps: float = 120.0,
329 movie: bool = False,
330 device=None,
331 dtype: Literal["float32", "float64"] | torch.dtype = "float32",
332 compile: bool | str = True,
333 progress: bool = False,
334) -> Sound:
335 """The sound of a strike by the full (geometrically nonlinear) model.
337 ``device`` is a torch device (``"cuda"``, ``"cuda:1"``, ``"mps"``, ``"cpu"``);
338 by default a CUDA GPU when one is available. ``dtype`` is ``"float32"`` (the
339 default, with the rounding safeguards of the kernel) or ``"float64"``.
340 ``compile`` wraps the step in ``torch.compile`` (``True``, the default,
341 with a fallback to the plain step when no compiler is available; ``False``;
342 or a ``torch.compile`` mode string) and, on CUDA, captures the steps of one
343 audio sample into a CUDA graph. ``fps`` is the movie frame rate and
344 also the batch rate of the run. ``loss`` overrides the fitted
345 ``(sigma0, sigma1, sigma_in)`` of the two-parameter loss operator.
346 """
347 sim = Simulation(plate, strike, listener, duration=duration, fs=fs, fps=fps, movie=movie, device=device, dtype=dtype, compile=compile)
348 sim.run(progress=progress)
349 return sim.sound()