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 ("reduce-overhead" if self.device.type == "cuda" else "default")
110 # With CUDA graphs ("reduce-overhead") the outputs of one step are the
111 # inputs of the next, and torch must be told where a step begins so that
112 # it may reuse its output buffers after they have been copied.
113 self._mark_step = lambda: None
115 # ---- progress ----
118 def target_samples(self) -> int:
119 return round(self.duration * self.fs)
122 def done(self) -> bool:
123 return self.samples >= self.target_samples
126 def t(self) -> float:
127 """Simulated time since the impact [s]."""
128 return self.total_steps * self.dt
130 def status(self) -> str:
131 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 "")
133 # ---- stepping ----
135 def run_batch(self) -> None:
136 """Advances by one movie frame (or to the end of the recording)."""
137 steps_left = self.target_samples * self.steps_per_sample - self.total_steps
138 steps = min(self.steps_per_batch, steps_left)
139 if steps <= 0:
140 return
141 t0 = time.perf_counter()
142 sps, H = self.steps_per_sample, self.hist_len
143 slots = []
144 state, stick = self._state, self._stick
145 S = self.static
146 if self._compile_pending:
147 self._try_compile()
148 mark = self._mark_step
149 for s in range(steps):
150 local = self.total_steps + s
151 mark()
152 state_and_stick = self._step(S, *state, stick)
153 state, stick = state_and_stick[:7], state_and_stick[7]
154 if local % sps == 0:
155 slot = (local // sps) % H
156 self._hist[slot] = state[6].reshape(-1)
157 slots.append(slot)
158 self._state, self._stick = state, stick
159 self.total_steps += steps
160 if slots:
161 self.pressure.append(self._audio(slots))
162 self.samples += len(slots)
163 if self.movie:
164 self.frames.append(state[2].detach().to("cpu", torch.float32).numpy().copy())
165 st = stick.detach().cpu().numpy()
166 self.stick = {"x": float(st[0]), "v": float(st[1]), "force": float(st[2]), "max_force": float(st[3]), "withdrawn": bool(st[5] > 0)}
167 self.seconds += time.perf_counter() - t0
169 def _try_compile(self) -> None:
170 """Wraps the step in ``torch.compile`` (which fuses its many small tensor
171 operations; about 6 times faster on the CPU and needed on a GPU) and
172 checks it on the current state; falls back to the eager step with a
173 warning when no compiler is available."""
174 import warnings
176 self._compile_pending = False
177 try:
178 compiled = torch.compile(step, dynamic=False, mode=self._compile_mode)
179 if self.device.type == "cuda" and self._compile_mode in ("reduce-overhead", "max-autotune"):
180 self._mark_step = torch.compiler.cudagraph_mark_step_begin
181 self._mark_step()
182 compiled(self.static, *self._state, self._stick) # compiles; the result is discarded and the state is unchanged
183 self._step = compiled
184 except Exception as e: # noqa: BLE001
185 warnings.warn(f"torch.compile failed ({type(e).__name__}: {str(e).splitlines()[0][:200]}); running the uncompiled step, which is slower", RuntimeWarning, stacklevel=3)
186 self._step = step
188 def _audio(self, slots: list[int]) -> np.ndarray:
189 """The retarded-time Rayleigh sum for the batch's samples: the acceleration
190 is the first difference of the delayed velocity, the average over one
191 sample interval, which is band-limited and immune to the per-step
192 rounding noise of the force evaluation."""
193 H = self.hist_len
194 slot = torch.as_tensor(slots, device=self.device)[:, None]
195 i0 = (slot - self._k0[None, :]) % H
196 i1 = (i0 - 1) % H
197 i2 = (i0 - 2) % H
198 v0 = torch.gather(self._hist, 0, i0)
199 v1 = torch.gather(self._hist, 0, i1)
200 v2 = torch.gather(self._hist, 0, i2)
201 fr = self._fr[None, :]
202 a = ((v0 + fr * (v1 - v0)) - (v1 + fr * (v2 - v1))) * self.fs
203 return (a @ self._lw).to("cpu", torch.float64).numpy()
205 def run(self, seconds: float | None = None, progress: bool = False) -> "Simulation":
206 """Advances by ``seconds`` of simulated time (to the end by default). With ``progress`` a line is printed per second of sound."""
207 target = self.target_samples if seconds is None else min(self.target_samples, self.samples + round(seconds * self.fs))
208 next_report = self.samples + self.fs
209 while self.samples < target and not self.done:
210 before = self.total_steps
211 self.run_batch()
212 if self.total_steps == before:
213 raise RuntimeError("the run made no progress; the step budget is exhausted before the sample target")
214 if progress and self.samples >= next_report:
215 print(f"{self.t:.2f} s simulated in {self.seconds:.1f} s wall; {self.status()}", flush=True)
216 next_report += self.fs
217 return self
219 def sound(self) -> Sound:
220 """What has been recorded so far."""
221 pressure = np.concatenate(self.pressure) if self.pressure else np.zeros(0)
222 frames = np.stack(self.frames) if self.movie else None
223 info = {
224 "model": "full",
225 "dt": self.dt,
226 "steps_per_sample": self.steps_per_sample,
227 "steps": self.total_steps,
228 "sigma0": self.loss[0],
229 "sigma1": self.loss[1],
230 "sigma_in": self.loss[2],
231 "peak_force": self.stick["max_force"],
232 "device": str(self.device),
233 "dtype": str(self.dtype).replace("torch.", ""),
234 "seconds": self.seconds,
235 }
236 return Sound(pressure.astype(np.float32), self.fs, self.plate, self.strike, self.listener, frames, self.fps, info)
238 # ---- state access, for tests and experiments ----
240 def state(self) -> dict[str, np.ndarray]:
241 """The displacement and velocity fields ``u, v, w, du, dv, dw`` on the grid, as float64 arrays."""
242 u, v, w, _, du, dv, dw = (x.detach().to("cpu", torch.float64).numpy() for x in self._state)
243 return {"u": u, "v": v, "w": w, "du": du, "dv": dv, "dw": dw}
245 def set_state(self, u=None, v=None, w=None, du=None, dv=None, dw=None) -> None:
246 """Overwrites fields of the state (grid arrays); the compensation array of w is reset."""
247 cur = list(self._state)
248 for k, x in ((0, u), (1, v), (2, w), (4, du), (5, dv), (6, dw)):
249 if x is not None:
250 cur[k] = torch.as_tensor(np.asarray(x, dtype=np.float64), dtype=self.dtype, device=self.device) * self.static.ins
251 cur[3] = torch.zeros_like(cur[2])
252 self._state = tuple(cur)
254 def step_once(self, strike: bool = False) -> None:
255 """Advances the state by one step without recording; without the drumstick unless ``strike``."""
256 S = self.static
257 if not strike and S.strike:
258 from dataclasses import replace
260 S = replace(S, strike=False)
261 out = step(S, *self._state, self._stick)
262 self._state, self._stick = out[:7], out[7]
263 self.total_steps += 1
266def simulate(
267 plate: Plate,
268 strike: Strike = Strike(),
269 listener: Listener = Listener(),
270 *,
271 duration: float = 4.0,
272 fs: int = 48000,
273 fps: float = 120.0,
274 movie: bool = False,
275 device=None,
276 dtype: Literal["float32", "float64"] | torch.dtype = "float32",
277 compile: bool | str = True,
278 progress: bool = False,
279) -> Sound:
280 """The sound of a strike by the full (geometrically nonlinear) model.
282 ``device`` is a torch device (``"cuda"``, ``"cuda:1"``, ``"mps"``, ``"cpu"``);
283 by default a CUDA GPU when one is available. ``dtype`` is ``"float32"`` (the
284 default, with the rounding safeguards of the kernel) or ``"float64"``.
285 ``compile`` wraps the step in ``torch.compile`` (``True``, the default,
286 with a fallback to the plain step when no compiler is available; ``False``;
287 or a ``torch.compile`` mode string). ``fps`` is the movie frame rate and
288 also the batch rate of the run. ``loss`` overrides the fitted
289 ``(sigma0, sigma1, sigma_in)`` of the two-parameter loss operator.
290 """
291 sim = Simulation(plate, strike, listener, duration=duration, fs=fs, fps=fps, movie=movie, device=device, dtype=dtype, compile=compile)
292 sim.run(progress=progress)
293 return sim.sound()