1"""The time loop: leapfrog integration of the shell with contact, damping, rigid-mode projection
2and radiation, block-compiled with JAX.
4One sub-step advances the state (q, v) of the shell, the state (zs, vs) of every exciter and the
5two retarded-time pressure rings:
7 f = f_int(q) - contact(q, zs) - 2 sigma1 rho h Lap(v_w)
8 v' = P_rigid [ (v + dt M^-1 f) / (1 + 2 sigma0 dt) ]
9 q' = q + dt v'
10 vs' = vs + dt F / m, zs' = zs + dt vs'
12where P_rigid removes the six rigid-body modes. The sub-step rate is an integer multiple m of the
13audio rate chosen from the linear stability bound. The acceleration that radiates, and that the
14optional record keeps, is the normal component of (v' - v) / dt, which for this scheme is the
15exact central second difference of q. It is decimated to the audio rate inside the loop by the
16polyphase accumulators of ``decimation.py``, so that no second difference is ever taken on
17stored displacements. At the optional probe points (``output.probes``) the same acceleration is
18interpolated, kept at the sub-step rate for one block, and decimated on the host with the same
19kernel as the listener pressure. See docs/model.md, section "Time integration", and docs/format.md.
20"""
21from __future__ import annotations
23import dataclasses
24import logging
25import math
26import time
27from dataclasses import dataclass, field
29import jax
30import jax.numpy as jnp
31import numpy as np
32from jax import lax
33from scipy.io import wavfile
35from .config import CONTINUATION_FIELDS, SimulationConfig, dict_differences
36from .decimation import DECIMATION_SPAN, StreamingDecimator, decimation_taps
37from .exciter import HitSchedule, build_patch
38from .geometry import CymbalProfile
39from .mesh import Mesh, disk_mesh
40from .radiation import RadiationModel
41from .shell import N_DOF, ShellModel
43log = logging.getLogger("cymsim")
47class HitSummary:
48 time: float
49 kind: str
50 peak_force: float # [N]
51 contact_time: float # total time with F > 0 [s]
52 rebound_velocity: float # exciter velocity at the end of the run [m/s], positive = upward
53 energy_in: float # kinetic energy at release [J]
54 energy_out: float # kinetic energy left in the exciter [J]
58class EnergyRecord:
59 time: float
60 kinetic: float
61 membrane: float
62 bending: float
63 exciters: float
64 max_w_over_h: float # nonlinearity meter: < 1 effectively linear, 3 to 5 clear cascade
67 def total(self) -> float:
68 return self.kinetic + self.membrane + self.bending + self.exciters
72class AccelerationRecord:
73 """Normal acceleration of every node at the audio rate, int16 against a scale per sample.
75 ``data[j, i]`` * ``scale[j]`` is the acceleration of node i along the rest-surface normal
76 (front side positive) at t = j / fs, low-pass filtered by the kernel ``taps`` (see
77 decimation.py). Sample j is final once the run has passed t = (j + W/2) / fs.
78 """
79 fs: int
80 data: np.ndarray # (T, N) int16
81 scale: np.ndarray # (T,) float32
82 taps: np.ndarray # (W, m) float64 polyphase kernel
85 def n_samples(self) -> int:
86 return len(self.data)
89 def time(self) -> np.ndarray:
90 return np.arange(self.n_samples) / self.fs
92 def as_float(self, start: int = 0, stop: int | None = None) -> np.ndarray:
93 """Decoded accelerations [m/s^2] for samples [start, stop), float32 (T', N)."""
94 return self.data[start:stop].astype(np.float32) * self.scale[start:stop, None]
97 def encode(a: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
98 """(int16 data, float32 scale) of a float block (T', N)."""
99 peak = np.abs(a).max(axis=1)
100 scale = np.where(peak > 0, peak / 32767.0, 1.0).astype(np.float32)
101 return np.round(a / scale[:, None]).astype(np.int16), np.where(peak > 0, scale, 0.0).astype(np.float32)
104 def empty(cls, fs: int, n_nodes: int, taps: np.ndarray) -> "AccelerationRecord":
105 return cls(fs, np.zeros((0, n_nodes), np.int16), np.zeros(0, np.float32), taps)
107 def extend(self, data: np.ndarray, scale: np.ndarray) -> None:
108 self.data = np.concatenate([self.data, data])
109 self.scale = np.concatenate([self.scale, scale])
113class DisplacementRecord:
114 """Transverse displacement w of every node [m] at a low frame rate, for display."""
115 rate: float
116 time: np.ndarray # (Td,)
117 w: np.ndarray # (Td, N) in the run's precision
120 def n_frames(self) -> int:
121 return len(self.time)
124 def empty(cls, rate: float, n_nodes: int, dtype=np.float32) -> "DisplacementRecord":
125 return cls(rate, np.zeros(0), np.zeros((0, n_nodes), dtype))
127 def extend(self, time: np.ndarray, w: np.ndarray) -> None:
128 self.time = np.concatenate([self.time, time])
129 self.w = np.concatenate([self.w, w])
133class SimulatorState:
134 """Everything the Python integrator needs to continue a run from step ``n_steps_done``.
136 Written to the ``state`` group of a .cymsim file. It is specific to this implementation
137 (the web companion keeps its own); readers of the format may ignore it.
138 """
139 n_steps_done: int
140 q: np.ndarray # (N, 5)
141 v: np.ndarray # (N, 5)
142 exciter_z: np.ndarray # (K,)
143 exciter_v: np.ndarray # (K,)
144 peak_force: np.ndarray # (K,)
145 contact_steps: np.ndarray # (K,) int32
146 ring_front: np.ndarray # (L,) pending retarded-time pressure, sub-step rate
147 ring_back: np.ndarray # (L,)
148 decimator_tail_front: np.ndarray # sub-steps still needed by non-final audio samples
149 decimator_tail_back: np.ndarray
150 acceleration_ring: np.ndarray | None = None # (W, N) partial audio samples of the record
151 decimator_tail_probes: np.ndarray | None = None # (P, L) the same for the probe signals
152 # Final outputs past the end of the duration: the run overshoots it by the kernel's half-span
153 # and by the rounding to whole blocks. A continued run begins its records with them. None in
154 # a state written before runs could be continued.
155 overhang_front: np.ndarray | None = None # (T',) baffled pressure [Pa]
156 overhang_unbaffled: np.ndarray | None = None # (T',)
157 overhang_probes: np.ndarray | None = None # (T', P)
158 overhang_acceleration_data: np.ndarray | None = None # (T', N) int16
159 overhang_acceleration_scale: np.ndarray | None = None
160 overhang_displacement_time: np.ndarray | None = None # (Td',)
161 overhang_displacement_w: np.ndarray | None = None # (Td', N)
164# After every block of a tangent run. The first four columns: the squared energy norm of the tangent,
165# kinetic_energy(dv) + linear_energy(dq), and the two parts of the same norm of the run. Then the parts:
166# of the tangent its kinetic, bending and linear membrane energy; of the run its bending, linear
167# membrane and (nonlinear, true) membrane energy. At large deflection the linear membrane energy of
168# the run, and presumably of the tangent, is far larger than the true one: the in-plane motion that
169# cancels the quadratic stretch shows up as linear strain.
170NORM_COLUMNS = ("time tangent_norm2 kinetic linear_strain "
171 "tangent_kinetic tangent_bending tangent_membrane_linear bending membrane_linear membrane")
172N_NORM_COLUMNS = len(NORM_COLUMNS.split())
176class TangentRecord:
177 """The tangent of a run (numerics.tangent): the derivative of its outputs and state with
178 respect to delta when every prescribed peak force is multiplied by (1 + delta), the limit of
179 (twin - run) / delta for twin runs. Signals as in SimulationResult; ``norms`` holds, after
180 every block, the time, the squared energy norm of the tangent state, kinetic_energy(dv) +
181 linear_energy(dq), the two parts of the same norm of the run's own state, and the energy parts of
182 both (NORM_COLUMNS; files written before the parts were logged have the first four columns only)."""
183 pressure_rayleigh: np.ndarray
184 pressure_unbaffled: np.ndarray
185 probe_acceleration: np.ndarray | None
186 displacement: DisplacementRecord | None
187 norms: np.ndarray # (B, N_NORM_COLUMNS), columns NORM_COLUMNS
188 state: SimulatorState | None = None
192class SimulationResult:
193 config: SimulationConfig
194 fs: int
195 dt: float
196 substeps: int
197 mesh: Mesh
198 z0_node: np.ndarray # (N,) rest height of the nodes [m]
199 h_node: np.ndarray # (N,) thickness at the nodes [m]
200 normal_node: np.ndarray # (N, 3) unit normal of the rest surface, front side
201 r_front: np.ndarray # (E,) element centroid to listener [m]
202 r_back: np.ndarray # (E,) shortest path over the rim to the listener [m]
203 pressure_rayleigh: np.ndarray # baffled Rayleigh integral at the listener [Pa]
204 pressure_unbaffled: np.ndarray # front minus diffracted back term [Pa]
205 hits: list[HitSummary]
206 energy_log: list[EnergyRecord] = field(default_factory=list)
207 acceleration: AccelerationRecord | None = None
208 displacement: DisplacementRecord | None = None
209 complete: bool = True
210 n_steps_done: int = 0
211 state: SimulatorState | None = None
212 wall_time: float = 0.0
213 creator: str = ""
214 probes: np.ndarray | None = None # (P, 2) probe points [m]
215 probe_acceleration: np.ndarray | None = None # (T, P) normal acceleration at the probes [m/s^2]
216 tangent: TangentRecord | None = None # with numerics.tangent
219 def n_nodes(self) -> int:
220 return self.mesh.n_nodes
223 def n_elements(self) -> int:
224 return self.mesh.n_elements
227 def time(self) -> np.ndarray:
228 return np.arange(len(self.pressure_rayleigh)) / self.fs
231 def duration_done(self) -> float:
232 """Simulated time covered by the run so far [s]."""
233 return self.n_steps_done * self.dt
235 def write_wav(self, prefix: str, normalize: bool = True) -> list[str]:
236 """Write <prefix>_rayleigh.wav and <prefix>_unbaffled.wav as 16-bit PCM. With
237 ``normalize`` the peak is placed at 0.9 full scale; otherwise 1 Pa maps to 0.05 FS."""
238 paths = []
239 for name, p in (("rayleigh", self.pressure_rayleigh), ("unbaffled", self.pressure_unbaffled)):
240 peak = float(np.abs(p).max()) if len(p) else 0.0
241 scale = 0.9 / (peak or 1.0) if normalize else 0.05
242 path = f"{prefix}_{name}.wav"
243 wavfile.write(path, self.fs, (np.clip(p * scale, -1, 1) * 32767).astype(np.int16))
244 paths.append(path)
245 return paths
247 def write_npz(self, prefix: str) -> str:
248 path = f"{prefix}_result.npz"
249 np.savez(path, fs=self.fs, dt=self.dt, pressure_rayleigh=self.pressure_rayleigh,
250 pressure_unbaffled=self.pressure_unbaffled,
251 energy_time=[e.time for e in self.energy_log],
252 energy=[[e.kinetic, e.membrane, e.bending, e.exciters, e.max_w_over_h] for e in self.energy_log],
253 **({} if self.probe_acceleration is None else
254 dict(probes=self.probes, probe_acceleration=self.probe_acceleration)))
255 return path
257 def save(self, path: str) -> str:
258 """Write a .cymsim file (see docs/format.md)."""
259 from .cymsim_file import save_cymsim
260 return save_cymsim(self, path)
263def steps_for(n_out: int, substeps: int, steps_per_block: int) -> int:
264 """Sub-steps of a run of n_out audio samples: the samples, the kernel's look-ahead of W/2
265 samples, rounded up to whole blocks."""
266 return math.ceil((n_out + DECIMATION_SPAN // 2) * substeps / steps_per_block) * steps_per_block
269def _log_until(energy_log: list[EnergyRecord], t_end: float) -> list[EnergyRecord]:
270 """The reports up to t_end (with a margin of rounding)."""
271 return [e for e in energy_log if e.time <= t_end * (1 + 1e-12)]
274def _cut_signals(front, unb, probe, disp, s: SimulatorState, n_out: int, dt: float, substeps: int):
275 """Signals cut to n_out samples (the state's overhang appended first), and the overhang
276 fields of the state past n_out."""
277 cat = lambda a, b: np.concatenate([a, b])
278 front, unb = cat(front, s.overhang_front), cat(unb, s.overhang_unbaffled)
279 over = dict(overhang_front=front[n_out:], overhang_unbaffled=unb[n_out:])
280 if probe is not None:
281 probe = cat(probe, s.overhang_probes)
282 over["overhang_probes"] = probe[n_out:]
283 probe = probe[:n_out]
284 if disp is not None:
285 t, w = cat(disp.time, s.overhang_displacement_time), cat(disp.w, s.overhang_displacement_w)
286 k_out = int(np.count_nonzero(np.round(t / dt) <= n_out * substeps)) # frame times are steps * dt
287 over.update(overhang_displacement_time=t[k_out:], overhang_displacement_w=w[k_out:])
288 disp = DisplacementRecord(disp.rate, t[:k_out], w[:k_out])
289 return front[:n_out], unb[:n_out], probe, disp, over
292def cut_run(prev: SimulationResult, config: SimulationConfig) -> SimulationResult | None:
293 """The run of ``config`` taken from the longer run ``prev`` (whose configuration differs at
294 most in CONTINUATION_FIELDS), or None if ``prev`` does not reach the end of ``config``'s run.
296 This is what ``Simulator(config).run(resume=prev)`` returns in that case, without building the
297 model: the records are cut to the duration, the samples past it join the state's overhang, and
298 the energy log stops where an uninterrupted run would have stopped."""
299 s, num = prev.state, config.numerics
300 n_out = int(round(prev.fs * num.duration))
301 n_steps = steps_for(n_out, prev.substeps, num.steps_per_block)
302 if s is None or s.overhang_front is None or s.n_steps_done < n_steps:
303 return None
304 front, unb, probe, disp, over = _cut_signals(prev.pressure_rayleigh, prev.pressure_unbaffled,
305 prev.probe_acceleration, prev.displacement, s, n_out,
306 prev.dt, prev.substeps)
307 acc = prev.acceleration
308 if acc is not None:
309 cat = lambda a, b: np.concatenate([a, b])
310 data, scale = cat(acc.data, s.overhang_acceleration_data), cat(acc.scale, s.overhang_acceleration_scale)
311 over.update(overhang_acceleration_data=data[n_out:], overhang_acceleration_scale=scale[n_out:])
312 acc = AccelerationRecord(acc.fs, data[:n_out], scale[:n_out], acc.taps)
313 tangent = prev.tangent
314 if tangent is not None:
315 ts = tangent.state
316 tf, tu, tp, td, tover = _cut_signals(tangent.pressure_rayleigh, tangent.pressure_unbaffled,
317 tangent.probe_acceleration, tangent.displacement, ts, n_out,
318 prev.dt, prev.substeps)
319 norms = tangent.norms[tangent.norms[:, 0] <= n_steps * prev.dt * (1 + 1e-12)]
320 tangent = TangentRecord(tf, tu, tp, td, norms, dataclasses.replace(ts, **tover))
321 return dataclasses.replace(
322 prev, config=config, pressure_rayleigh=front, pressure_unbaffled=unb, acceleration=acc,
323 displacement=disp, complete=True, state=dataclasses.replace(s, **over),
324 energy_log=_log_until(prev.energy_log, n_steps * prev.dt), probe_acceleration=probe, tangent=tangent)
327class _Stream:
328 """What a state tuple streams out block by block (the run's, or its tangent's): the listener
329 pressure and the probe signals through their decimators, and the displacement frames."""
331 def __init__(self, sim: "Simulator", prev=None):
332 """``prev``: (pressure_rayleigh, pressure_unbaffled, probe_acceleration, displacement,
333 state) of the run being continued, or None for a run from rest."""
334 m, P = sim.substeps, len(sim.probe_xy)
335 self.sim, self.P = sim, P
336 self.disp = (DisplacementRecord.empty(sim.disp_rate, sim.mesh.n_nodes, np.dtype(sim.shell.dtype))
337 if sim.disp_rate > 0 else None)
338 if prev is None:
339 self.f, self.u, self.p = [np.zeros(0)], [np.zeros(0)], [np.zeros((0, P))]
340 self.dec_f, self.dec_b = StreamingDecimator(m), StreamingDecimator(m)
341 self.dec_p = [StreamingDecimator(m) for _ in range(P)]
342 return
343 front, unb, probe, disp, s = prev
344 n = s.n_steps_done
345 self.f, self.u = [front, s.overhang_front], [unb, s.overhang_unbaffled]
346 self.p = [probe, s.overhang_probes] if P else [np.zeros((0, 0))]
347 self.dec_f = StreamingDecimator(m, tail=s.decimator_tail_front, n_seen=n)
348 self.dec_b = StreamingDecimator(m, tail=s.decimator_tail_back, n_seen=n)
349 self.dec_p = [StreamingDecimator(m, tail=t, n_seen=n) for t in s.decimator_tail_probes[:P]] if P else []
350 if self.disp is not None:
351 self.disp.extend(disp.time, disp.w)
352 self.disp.extend(s.overhang_displacement_time, s.overhang_displacement_w)
354 def take(self, state, n0: int, n1: int, k0: int, k1: int):
355 """Collect the outputs of the block [n0, n1) from ``state``; returns the state with the
356 pressure ring slots it read cleared."""
357 sim = self.sim
358 # the ring slots of this block's sub-steps are final; read, clear, decimate
359 slots = jnp.asarray(sim.radiation.slots(n0, n1))
360 pf, pb = np.asarray(state[4][slots], np.float64), np.asarray(state[5][slots], np.float64)
361 state = state[:4] + (sim._clear(state[4], slots), sim._clear(state[5], slots)) + state[6:]
362 # (samples past the duration are kept too; _result moves them to the state's overhang)
363 new_f, new_b = self.dec_f.push(pf), self.dec_b.push(pb)
364 self.f.append(new_f)
365 self.u.append(new_f - new_b)
366 if self.P:
367 rows = np.asarray(state[11], np.float64)
368 self.p.append(np.stack([d.push(rows[:, i]) for i, d in enumerate(self.dec_p)], 1))
369 if self.disp is not None and k1 > k0:
370 frames = np.asarray(state[10][:k1 - k0], np.dtype(sim.shell.dtype))
371 self.disp.extend(sim.frame_steps[k0:k1] * sim.dt, frames)
372 return state
374 def split(self, n_done: int, state, record_ring=None):
375 """(front, unbaffled, probe, displacement) up to the duration and the SimulatorState
376 after n_done steps, which keeps everything past the duration as its overhang."""
377 sim = self.sim
378 q, v, zs, vs, bf, bb, Fmax, ncont = jax.device_get(state[:8])
379 front, unb = np.concatenate(self.f), np.concatenate(self.u)
380 n = min(len(front), len(unb))
381 probe = np.concatenate(self.p) if self.P else None
382 empty = SimulatorState(0, *(None,) * 10, overhang_front=np.zeros(0), overhang_unbaffled=np.zeros(0),
383 overhang_probes=np.zeros((0, self.P)) if self.P else None,
384 overhang_displacement_time=np.zeros(0) if self.disp is not None else None,
385 overhang_displacement_w=(np.zeros((0, sim.mesh.n_nodes), self.disp.w.dtype)
386 if self.disp is not None else None))
387 front, unb, probe, disp, over = _cut_signals(front[:n], unb[:n], probe, self.disp, empty, sim.n_out,
388 sim.dt, sim.substeps)
389 ncont = np.asarray(ncont) if np.issubdtype(np.asarray(ncont).dtype, np.integer) \
390 else np.zeros(np.shape(ncont), np.int32) # a tangent's float0 counts
391 st = SimulatorState(
392 n_done, np.asarray(q), np.asarray(v), np.asarray(zs), np.asarray(vs), np.asarray(Fmax), ncont,
393 np.asarray(bf), np.asarray(bb), self.dec_f.tail.copy(), self.dec_b.tail.copy(),
394 record_ring, np.stack([d.tail for d in self.dec_p]) if self.dec_p else None, **over)
395 return front, unb, probe, disp, st
398class Simulator:
399 """Build every part of the model from a SimulationConfig; ``run()`` integrates it."""
401 def __init__(self, config: SimulationConfig):
402 config.validate()
403 self.config = config
404 num, mat, out = config.numerics, config.material, config.output
405 t0 = time.time()
407 self.profile = CymbalProfile(config.profile, config.imperfections)
408 self.mesh = disk_mesh(config.profile.radius, config.mesh.element_size, config.mesh.smoothing)
409 log.info("mesh: %d nodes, %d triangles, %d DOFs (nonlinear = %s)",
410 self.mesh.n_nodes, self.mesh.n_elements, N_DOF * self.mesh.n_nodes, num.nonlinear)
412 self.shell = ShellModel(self.mesh, self.profile, mat, num)
413 sh = self.shell
414 rim = sh.z0_node[self.mesh.boundary]
415 log.info("profile: height %.1f mm, rim warp p-p %.2f mm, %d dents, h in [%.2f, %.2f] mm, "
416 "max |grad z0| = %.3f, mass %.3f kg",
417 sh.z0_node.max() * 1e3, (rim.max() - rim.min()) * 1e3, config.imperfections.n_dents,
418 sh.h_e.min() * 1e3, sh.h_e.max() * 1e3, sh.max_slope, sh.mass)
419 if sh.max_slope > config.profile.max_slope:
420 log.warning("max slope %.2f exceeds the shallow-shell bound %.2f", sh.max_slope, config.profile.max_slope)
422 # ---- time step: integer number of sub-steps per audio sample, inside the stability bound.
423 # The run covers W/2 audio samples beyond the duration so that every sample of the
424 # duration is final after the decimation kernel's look-ahead.
425 self.fs = num.fs
426 self.substeps = math.ceil(1.0 / (self.fs * num.cfl_safety * sh.dt_stable))
427 self.dt = 1.0 / (self.fs * self.substeps)
428 self.n_out = int(round(self.fs * num.duration))
429 self.span = DECIMATION_SPAN
430 self.n_steps = steps_for(self.n_out, self.substeps, num.steps_per_block)
431 self.n_blocks = self.n_steps // num.steps_per_block
432 xc, yc = self.mesh.centroid[sh.limiting_element]
433 log.info("omega_max: element bound %.3e rad/s (limiting element at r = %.3f m, h = %.2f mm), "
434 "global %.3e rad/s (%d power iterations); %s bound used: dt = %.3e s, "
435 "%d sub-steps per sample, %d steps",
436 sh.omega_max, math.hypot(xc, yc), sh.h_e[sh.limiting_element] * 1e3, sh.omega_global,
437 sh.omega_global_iterations, "global" if sh.omega_step < sh.omega_max else "element",
438 self.dt, self.substeps, self.n_steps)
439 self.taps = decimation_taps(self.substeps)
441 # ---- hits
442 self.patches = [build_patch(h, self.mesh, self.profile) for h in config.hits]
443 self.schedule = HitSchedule(self.patches, self.dt, sh.dtype)
444 for k, p in enumerate(self.patches):
445 tau = p.free_contact_time()
446 if p.prescribed:
447 log.info("hit %d at t = %.3f s: prescribed raised-cosine force at r = %.3f m, peak %.1f N, "
448 "duration %.2f ms", k, p.time, p.hit.r, p.force_peak, p.force_duration * 1e3)
449 continue
450 log.info("hit %d at t = %.3f s: %s, %d-point patch at r = %.3f m, k = %.3g N/m^%.2f, "
451 "m = %.3f kg, v = %.1f m/s, energy %.2f J, free contact ~ %.2f ms "
452 "(direct excitation below ~%.0f Hz)",
453 k, p.time, p.hit.exciter.kind, p.n_points, p.hit.r, p.k_total, p.alpha,
454 p.mass, p.velocity, p.energy, tau * 1e3, 0.5 / tau)
455 if p.time >= num.duration:
456 log.warning("hit %d is scheduled after the end of the run", k)
458 # ---- radiation
459 self.radiation = RadiationModel(self.mesh, self.profile, sh.z0_node, config.listener,
460 self.dt, num.steps_per_block, sh.dtype)
462 # ---- probes: containing element and barycentric weights of each point
463 self.probe_xy = np.array(out.probes, dtype=np.float64).reshape(-1, 2)
464 loc = [self.mesh.locate(x, y) for x, y in self.probe_xy]
465 self.probe_nodes = jnp.asarray(np.array([l[0] for l in loc]).reshape(-1, 3), dtype=jnp.int32)
466 self.probe_w = jnp.asarray(np.array([l[1] for l in loc]).reshape(-1, 3), dtype=sh.dtype)
468 # ---- records
469 self.record = bool(out.record_acceleration)
470 self.disp_rate = float(out.displacement_rate)
471 if self.record:
472 nbytes = self.n_out * (2 * self.mesh.n_nodes + 4)
473 log.info("acceleration record: %d samples x %d nodes, int16, %.2f GB in memory",
474 self.n_out, self.mesh.n_nodes, nbytes / 1e9)
475 if nbytes > 2e9:
476 log.warning("the acceleration record needs %.1f GB; consider a shorter duration or "
477 "a coarser mesh", nbytes / 1e9)
478 S = self.fs * self.substeps # sub-steps per second
479 if self.disp_rate > 0:
480 # every frame the run passes, including those past the duration (the state's overhang)
481 n_frames = int(math.floor(self.n_steps * self.disp_rate / S)) + 1
482 self.frame_steps = np.round(np.arange(n_frames) * S / self.disp_rate).astype(np.int64)
483 else:
484 self.frame_steps = np.zeros(0, np.int64)
485 self._build_step()
486 self.setup_time = time.time() - t0
487 log.info("setup done in %.1f s; backend %s", self.setup_time, jax.devices()[0].platform)
489 # ---- one sub-step and its block loop
490 def _build_step(self):
491 sh, sched, rad = self.shell, self.schedule, self.radiation
492 dt, sigma0 = self.dt, float(self.config.material.sigma0) # a numpy scalar would promote float32
493 m, W, N = self.substeps, self.span, self.mesh.n_nodes
494 spb = self.config.numerics.steps_per_block
495 tri, normal = sh.tri_j, sh.normal_j
496 record, rate = self.record, self.disp_rate
497 taps_j = jnp.asarray(self.taps, sh.dtype)
498 S = float(self.fs * m)
499 self.acc_rows = spb // m + 2 # audio samples emitted per block, at most
500 self.disp_rows = int(math.ceil(spb * rate / S)) + 2 if rate > 0 else 0
502 n_probes = len(self.probe_xy)
503 p_nodes, p_w = self.probe_nodes, self.probe_w
505 def step(n, k_block0, j_block0, st, f_peak=None):
506 q, v, zs, vs, bf, bb, Fmax, ncont, ring, acc_out, disp_out, pbuf = st
507 if rate > 0:
508 # frame k sits at sub-step round(k S / rate); emit q at step n if n is one of them
509 k_cand = jnp.round(n * (rate / S))
510 is_frame = jnp.round(k_cand * (S / rate)) == n
511 row = jnp.clip(k_cand.astype(jnp.int32) - k_block0, 0, self.disp_rows - 1)
512 disp_out = disp_out.at[row].add(jnp.where(is_frame, q[:, 2], 0.0))
513 Fk, F, zs, vs = sched.contact(n, q, zs, vs, f_peak)
514 fo = sh.internal_force(q)
515 fo = sched.apply_to_plate(fo, Fk)
516 fo = fo.at[:, 2].add(-sh.damping_force_w(v[:, 2]))
517 v1 = (v + dt * fo * sh.Minv_j) / (1 + 2 * sigma0 * dt) # sigma0 implicit
518 v1 = sh.project_rigid(v1)
519 q1 = q + dt * v1
520 zs1, vs1 = sched.advance(zs, vs, F, dt)
521 # normal acceleration at step n: central second difference of q, exact for leapfrog
522 a_node = (((v1 - v) / dt)[:, :3] * normal).sum(1)
523 acc_e = a_node[tri].mean(1)
524 bf, bb = rad.scatter(n, acc_e, bf, bb)
525 if n_probes:
526 pbuf = pbuf.at[n % spb].set((a_node[p_nodes] * p_w).sum(1)) # blocks start at multiples of spb
527 if record:
528 j0, kk = n // m, n % m
529 emit = kk == 0
530 j_e = j0 - W // 2 # the sample that just became final
531 slot_e = j_e % W
532 done = ring[slot_e]
533 row = jnp.clip(j_e - j_block0, 0, self.acc_rows - 1)
534 acc_out = acc_out.at[row].add(jnp.where(emit, done, 0.0))
535 ring = ring.at[slot_e].set(jnp.where(emit, 0.0, done))
536 # slot s holds sample j0 + 1 - W/2 + w with s = (j0 + 1 - W/2 + w) mod W
537 r = (j0 + 1 - W // 2) % W
538 tw = jnp.take(taps_j[:, kk], (jnp.arange(W) - r) % W)
539 ring = ring + tw[:, None] * a_node[None, :]
540 return (q1, v1, zs1, vs1, bf, bb, jnp.maximum(Fmax, F), ncont + (F > 0), ring, acc_out, disp_out, pbuf)
543 def run_block(n0, k_block0, j_block0, st):
544 return lax.fori_loop(0, spb, lambda k, s: step((n0 + k).astype(jnp.int32), k_block0, j_block0, s), st)
547 def tangent_block(n0, k_block0, j_block0, st, dst):
548 """One block of the run together with its tangent dst: the derivative of the state
549 with respect to delta when every prescribed peak force is multiplied by (1 + delta),
550 carried exactly by forward-mode differentiation of the step."""
551 fp = sched.f_peak
553 def body(k, carry):
554 n = (n0 + k).astype(jnp.int32)
555 return jax.jvp(lambda s, f: step(n, k_block0, j_block0, s, f), (carry[0], fp), (carry[1], fp))
556 return lax.fori_loop(0, spb, body, (st, dst))
559 def energies(st):
560 q, v, _, vs = st[0], st[1], st[2], st[3]
561 Em, Eb = sh.energy_parts_per_element(q)
562 return (sh.kinetic_energy(v), Em.sum(), Eb.sum(), sched.energy(vs),
563 jnp.abs(q[:, 2] / sh.h_node_j).max())
566 def clear(buf, idx):
567 return buf.at[idx].set(0.0)
570 def norms(st, dst):
571 """The NORM_COLUMNS after the time: energy norms and energy parts of tangent and run."""
572 tk = sh.kinetic_energy(dst[1])
573 tm, tb = sh.linear_energy_parts(dst[0])
574 k = sh.kinetic_energy(st[1])
575 ml, b = sh.linear_energy_parts(st[0])
576 m, _ = sh.energy_parts_per_element(st[0])
577 return (tk + tm + tb, k, ml + b, tk, tb, tm, b, ml, m.sum())
579 self._step, self._run_block, self._energies, self._clear = step, run_block, energies, clear
580 self._tangent_block, self._norms = tangent_block, norms
583 def zero_tangent(state):
584 """A zero tangent for a device state tuple (float0 for its integer arrays)."""
585 return jax.tree_util.tree_map(
586 lambda a: np.zeros(a.shape, jax.dtypes.float0) if jnp.issubdtype(a.dtype, jnp.integer)
587 else jnp.zeros_like(a), state)
589 def initial_state(self):
590 N, dtype = self.mesh.n_nodes, self.shell.dtype
591 zs, vs = self.schedule.initial_state()
592 bf, bb = self.radiation.initial_buffers()
593 K = self.schedule.n_hits
594 ring = jnp.zeros((self.span, N) if self.record else (0,), dtype)
595 return (jnp.zeros((N, N_DOF), dtype), jnp.zeros((N, N_DOF), dtype), zs, vs, bf, bb,
596 jnp.zeros(K, dtype), jnp.zeros(K, jnp.int32), ring, self._empty_outputs()[0],
597 self._empty_outputs()[1], jnp.zeros((self.config.numerics.steps_per_block, len(self.probe_xy)), dtype))
599 def _state_from(self, s: SimulatorState):
600 """The device state tuple of ``initial_state`` rebuilt from a saved SimulatorState."""
601 dtype = self.shell.dtype
602 f = lambda a: jnp.asarray(a, dtype)
603 ring = f(s.acceleration_ring) if self.record else jnp.zeros((0,), dtype)
604 acc_out, disp_out = self._empty_outputs()
605 return (f(s.q), f(s.v), f(s.exciter_z), f(s.exciter_v), f(s.ring_front), f(s.ring_back),
606 f(s.peak_force), jnp.asarray(s.contact_steps, jnp.int32), ring, acc_out, disp_out,
607 jnp.zeros((self.config.numerics.steps_per_block, len(self.probe_xy)), dtype))
609 def _tangent_from(self, s: SimulatorState):
610 """The device tangent tuple rebuilt from a saved tangent SimulatorState."""
611 st = self._state_from(s)
612 return st[:7] + (np.zeros(st[7].shape, jax.dtypes.float0),) + st[8:]
614 def check_resumable(self, prev: SimulationResult) -> None:
615 """Raise ValueError unless this simulator can continue the run ``prev``."""
616 s = prev.state
617 if s is None:
618 raise ValueError("the earlier run has no integrator state of this program")
619 if s.overhang_front is None:
620 raise ValueError("the earlier run was saved before cymsim could continue runs; run it again")
621 diffs = dict_differences(prev.config.identity_dict(), self.config.identity_dict())
622 if diffs:
623 raise ValueError("the earlier run has a different configuration (" + ", ".join(diffs) + "); only "
624 + ", ".join(".".join(f) for f in CONTINUATION_FIELDS) + " may change")
625 if (prev.fs, prev.substeps, prev.dt, prev.n_nodes) != (self.fs, self.substeps, self.dt, self.mesh.n_nodes):
626 raise ValueError("the earlier run has a different mesh or time step; was cymsim changed since?")
627 if np.dtype(s.q.dtype) != np.dtype(self.shell.dtype):
628 raise ValueError(f"the earlier run's state is {s.q.dtype}, this run is {np.dtype(self.shell.dtype)}")
629 if s.n_steps_done % self.config.numerics.steps_per_block:
630 raise ValueError("the earlier run stopped inside a block")
631 if self.config.numerics.tangent and (prev.tangent is None or prev.tangent.state is None):
632 raise ValueError("the earlier run has no tangent state")
633 if self.config.numerics.tangent and prev.tangent.norms.shape[1] != N_NORM_COLUMNS:
634 raise ValueError("the earlier tangent run logged fewer energy columns (written by an older cymsim)")
636 def _empty_outputs(self):
637 N, dtype = self.mesh.n_nodes, self.shell.dtype
638 return (jnp.zeros((self.acc_rows, N) if self.record else (0,), dtype),
639 jnp.zeros((self.disp_rows, N) if self.disp_rate > 0 else (0,), dtype))
641 def _hit_summaries(self, Fmax, ncont, vs) -> list[HitSummary]:
642 hits = []
643 for k, p in enumerate(self.patches):
644 hits.append(HitSummary(p.time, p.hit.exciter.kind, float(Fmax[k]), float(ncont[k]) * self.dt,
645 float(vs[k]), p.energy, 0.5 * p.mass * float(vs[k]) ** 2))
646 return hits
648 def _result(self, state, energy_log, stream: _Stream, acc, n_done, wall, tangent=None) -> SimulationResult:
649 """The result so far. Every series holds all final samples; those past the duration are
650 moved to the state's overhang. ``tangent``: (tangent state, its _Stream, norm rows)."""
651 n_out = self.n_out
652 Fmax, ncont, vs = jax.device_get((state[6], state[7], state[3]))
653 acc_over = {}
654 if acc is not None:
655 acc_over = dict(overhang_acceleration_data=acc.data[n_out:], overhang_acceleration_scale=acc.scale[n_out:])
656 acc = AccelerationRecord(acc.fs, acc.data[:n_out], acc.scale[:n_out], acc.taps)
657 ring = np.asarray(jax.device_get(state[8])) if self.record else None
658 p_front, p_unb, probe, disp, sim_state = stream.split(n_done, state, ring)
659 sim_state = dataclasses.replace(sim_state, **acc_over)
660 trec = None
661 if tangent is not None:
662 dstate, tstream, norms = tangent
663 tf, tu, tp, td, tst = tstream.split(n_done, dstate)
664 trec = TangentRecord(tf, tu, tp, td, np.array(norms, np.float64).reshape(len(norms), N_NORM_COLUMNS), tst)
665 from . import __version__
666 return SimulationResult(
667 self.config, self.fs, self.dt, self.substeps, self.mesh, self.shell.z0_node, self.shell.h_node,
668 self.shell.normal_node, self.radiation.r_front, self.radiation.r_back, p_front, p_unb,
669 self._hit_summaries(Fmax, ncont, vs), list(energy_log), acc, disp, n_done >= self.n_steps, n_done,
670 sim_state, wall, f"cymsim {__version__}", self.probe_xy if len(self.probe_xy) else None, probe, trec)
672 def run(self, save_path: str | None = None, save_every: float | None = None,
673 resume: SimulationResult | None = None, checkpoint=None,
674 checkpoint_every: float | None = None) -> SimulationResult:
675 """Integrate the run. With ``save_path`` a .cymsim file is written at the end and, if
676 ``save_every`` (seconds of wall time) is given, also periodically during the run, each
677 time as a consistent partial file that can be loaded (see docs/format.md). Likewise
678 ``checkpoint``, a function of a partial SimulationResult, is called every
679 ``checkpoint_every`` seconds of wall time (the run cache uploads partial runs with it).
681 With ``resume``, an earlier run (usually loaded from a .cymsim file) of a configuration
682 that differs from this one at most in CONTINUATION_FIELDS, the run continues from that
683 run's integrator state, and its records begin with the earlier ones. On the same device
684 the result is identical to that of an uninterrupted run. If the earlier run already
685 covers the duration, nothing is integrated and its records are cut to the duration."""
686 num = self.config.numerics
687 m, spb, N = self.substeps, num.steps_per_block, self.mesh.n_nodes
688 acc = AccelerationRecord.empty(self.fs, N, self.taps) if self.record else None
689 tangent = bool(num.tangent)
690 dstate = tstream = norms = None
691 if resume is None:
692 state, n_done, wall0 = self.initial_state(), 0, 0.0
693 energy_log: list[EnergyRecord] = []
694 stream = _Stream(self)
695 if tangent:
696 dstate, tstream, norms = self.zero_tangent(state), _Stream(self), []
697 else:
698 self.check_resumable(resume)
699 s = resume.state
700 state, n_done, wall0 = self._state_from(s), s.n_steps_done, resume.wall_time
701 energy_log = list(resume.energy_log)
702 stream = _Stream(self, (resume.pressure_rayleigh, resume.pressure_unbaffled, resume.probe_acceleration,
703 resume.displacement, s))
704 if tangent:
705 t = resume.tangent
706 dstate = self._tangent_from(t.state)
707 tstream = _Stream(self, (t.pressure_rayleigh, t.pressure_unbaffled, t.probe_acceleration,
708 t.displacement, t.state))
709 norms = [tuple(r) for r in t.norms]
710 if acc is not None:
711 acc.extend(resume.acceleration.data, resume.acceleration.scale)
712 acc.extend(s.overhang_acceleration_data, s.overhang_acceleration_scale)
713 if n_done >= self.n_steps:
714 energy_log = _log_until(energy_log, self.n_steps * self.dt)
715 if tangent:
716 norms = [r for r in norms if r[0] <= self.n_steps * self.dt * (1 + 1e-12)]
717 log.info("continuing the run from t = %.3f s", n_done * self.dt)
718 result_now = lambda n, wall: self._result(state, energy_log, stream, acc, n, wall,
719 (dstate, tstream, norms) if tangent else None)
720 writer = None
721 if save_path is not None:
722 from .cymsim_file import CymsimWriter
723 writer = CymsimWriter(save_path, result_now(n_done, wall0))
724 t_run = time.time()
725 t_saved = t_checkpoint = t_run
726 report_every = max(1, self.n_blocks // max(1, num.n_progress))
727 blk0 = n_done // spb
728 log.info("running %d blocks of %d steps (nonlinearity meter: max|w|/h < 1 effectively "
729 "linear; 3 to 5 clear cascade)", max(0, self.n_blocks - blk0), spb)
730 for blk in range(blk0, self.n_blocks):
731 n0, n1 = blk * spb, (blk + 1) * spb
732 k0 = int(np.searchsorted(self.frame_steps, n0))
733 k1 = int(np.searchsorted(self.frame_steps, n1))
734 j_block0 = -(-n0 // m) - self.span // 2
735 acc_out, disp_out = self._empty_outputs()
736 state = state[:9] + (acc_out, disp_out) + state[11:]
737 if tangent:
738 dstate = dstate[:9] + self._empty_outputs() + dstate[11:]
739 state, dstate = self._tangent_block(jnp.int32(n0), jnp.int32(k0), jnp.int32(j_block0), state, dstate)
740 dstate = tstream.take(dstate, n0, n1, k0, k1)
741 norms.append((n1 * self.dt, *(float(x) for x in self._norms(state, dstate))))
742 else:
743 state = self._run_block(jnp.int32(n0), jnp.int32(k0), jnp.int32(j_block0), state)
744 state = stream.take(state, n0, n1, k0, k1)
745 n_done = n1
747 # ---- the acceleration record
748 if self.record:
749 n_emit = -(-n1 // m) - -(-n0 // m)
750 rows = np.asarray(state[9][:n_emit], np.float32)
751 j = j_block0 + np.arange(n_emit)
752 sel = j >= 0
753 if sel.any():
754 acc.extend(*AccelerationRecord.encode(rows[sel]))
756 if blk % report_every == 0 or blk == self.n_blocks - 1:
757 ke, em, eb, ks, woh = (float(x) for x in self._energies(state))
758 t = n1 * self.dt
759 energy_log.append(EnergyRecord(t, ke, em, eb, ks, woh))
760 log.info(" t = %6.3f s KE %.3e membrane %.3e bending %.3e exciters %.3e J "
761 "max|w|/h = %5.2f [%.0f s wall]", t, ke, em, eb, ks, woh, time.time() - t_run)
762 if not math.isfinite(ke + em + eb):
763 raise FloatingPointError("energy is not finite: the run has blown up "
764 "(lower numerics.cfl_safety)")
765 res = None
766 if writer is not None and save_every is not None and time.time() - t_saved >= save_every \
767 and blk < self.n_blocks - 1:
768 res = result_now(n_done, wall0 + time.time() - t_run)
769 writer.update(res)
770 t_saved = time.time()
771 log.info(" saved partial run (%.3f s) to %s", res.duration_done, save_path)
772 if checkpoint is not None and checkpoint_every is not None \
773 and time.time() - t_checkpoint >= checkpoint_every and blk < self.n_blocks - 1:
774 checkpoint(res or result_now(n_done, wall0 + time.time() - t_run))
775 t_checkpoint = time.time()
777 result = result_now(n_done, wall0 + time.time() - t_run)
778 for h in result.hits:
779 log.info("hit (%s): peak force %.1f N, contact %.2f ms, rebound velocity %+.2f m/s "
780 "(energy left in exciter %.2f of %.2f J)", h.kind, h.peak_force,
781 h.contact_time * 1e3, h.rebound_velocity, h.energy_out, h.energy_in)
782 for name, p in (("rayleigh", result.pressure_rayleigh), ("unbaffled", result.pressure_unbaffled)):
783 peak = float(np.abs(p).max())
784 log.info("%-10s: peak %.3f Pa (%.1f dB SPL peak)", name, peak, 20 * np.log10(max(peak, 1e-12) / 2e-5))
785 if writer is not None:
786 writer.update(result)
787 writer.close()
788 log.info("wrote %s", save_path)
789 return result
792def simulate(config: SimulationConfig, resume: SimulationResult | None = None,
793 cache: str | None = None) -> SimulationResult:
794 """Run ``config`` from rest, from the earlier run ``resume``, or through the cache at the URL
795 ``cache`` (cache.py)."""
796 if cache is not None:
797 if resume is not None:
798 raise ValueError("give either resume or cache, not both")
799 from .cache import cached_run
800 return cached_run(config, cache)
801 return Simulator(config).run(resume=resume)