1"""Matplotlib views of a plate, its modes, and a sound. Every function draws
2into ``ax`` (a new figure by default) and returns the axes. Requires
3``matplotlib`` (``pip install cymsim[plot]``)."""
5from __future__ import annotations
7import numpy as np
9from .plate import Plate
10from .sound import Sound
11from .spectrogram import Spectrogram
13__all__ = ["plate", "mode", "frame", "signal", "spectrogram", "decay_times", "force", "animate"]
16def _ax(ax, figsize=(6, 5)):
17 import matplotlib.pyplot as plt
19 if ax is None:
20 _, ax = plt.subplots(figsize=figsize)
21 return ax
24def _extent(pl: Plate):
25 g = pl.grid
26 return [g.x[0] * 1e3 - g.hx * 500, g.x[-1] * 1e3 + g.hx * 500, g.y[0] * 1e3 - g.hy * 500, g.y[-1] * 1e3 + g.hy * 500]
29def field(pl: Plate, values: np.ndarray, ax=None, *, cmap="viridis", label: str = "", symmetric: bool = False, strike=None, listener=None):
30 """A color map of a grid array over the plate (outside masked), in mm coordinates."""
31 ax = _ax(ax)
32 v = np.ma.masked_where(~pl.grid.mask, values)
33 kw = {}
34 if symmetric:
35 m = float(np.abs(v).max()) or 1.0
36 kw = {"vmin": -m, "vmax": m}
37 im = ax.imshow(v, origin="lower", extent=_extent(pl), cmap=cmap, interpolation="nearest", **kw)
38 ax.figure.colorbar(im, ax=ax, label=label)
39 if strike is not None:
40 x, y = pl.strike_point(strike)
41 ax.plot(x * 1e3, y * 1e3, "o", mfc="none", mec="yellow", mew=2, ms=10)
42 if listener is not None:
43 lx, ly, _ = listener.position
44 ax.plot(lx * 1e3, ly * 1e3, "v", color="white", ms=8, clip_on=False)
45 ax.set_xlabel("x [mm]")
46 ax.set_ylabel("y [mm]")
47 ax.set_aspect("equal")
48 return ax
51def plate(pl: Plate, ax=None, what: str = "auto", strike=None, listener=None):
52 """The plate: its profile (when it has one) or its thickness, with the strike point and the listener when given."""
53 if what == "auto":
54 what = "thickness" if pl.is_flat else "profile"
55 if what == "profile":
56 ax = field(pl, pl.w0 * 1e3, ax, cmap="cividis", label="stress-free height [mm]", strike=strike, listener=listener)
57 elif what == "thickness":
58 ax = field(pl, pl.h * 1e3, ax, cmap="magma", label="thickness [mm]", strike=strike, listener=listener)
59 else:
60 raise ValueError("what must be 'auto', 'profile' or 'thickness'")
61 ax.set_title(pl.name)
62 return ax
65def mode(modes, m: int, ax=None, amplitude: float = 1e-3):
66 """The transverse shape of mode ``m``, scaled to the given peak displacement [m]."""
67 w = modes.w[m]
68 w = w / (np.abs(w).max() or 1.0) * amplitude
69 ax = field(modes.plate, w * 1e3, ax, cmap="RdBu_r", label="w [mm]", symmetric=True)
70 ax.set_title(f"mode {m}: {modes.freqs[m]:.1f} Hz, T60 {modes.t60[m]:.2f} s")
71 return ax
74def frame(sound: Sound, index: int, ax=None, scale: float | None = None):
75 """Movie frame ``index`` of the transverse displacement; ``scale`` fixes the color range [m]."""
76 if sound.frames is None:
77 raise ValueError("this sound has no movie frames; run with movie=True")
78 w = sound.frames[index]
79 ax = _ax(ax)
80 v = np.ma.masked_where(~sound.plate.grid.mask, w * 1e3)
81 m = (scale or float(np.abs(sound.frames).max()) or 1e-9) * 1e3
82 im = ax.imshow(v, origin="lower", extent=_extent(sound.plate), cmap="RdBu_r", vmin=-m, vmax=m, interpolation="nearest")
83 ax.figure.colorbar(im, ax=ax, label="w [mm]")
84 ax.set_title(f"t = {sound.frame_time(index) * 1e3:.1f} ms")
85 ax.set_xlabel("x [mm]")
86 ax.set_ylabel("y [mm]")
87 ax.set_aspect("equal")
88 return ax
91def signal(sound: Sound, ax=None, tmax: float | None = None):
92 """The listener pressure against time."""
93 ax = _ax(ax, (8, 3))
94 t = sound.time
95 sel = slice(None) if tmax is None else t <= tmax
96 ax.plot(t[sel], sound.pressure[sel], lw=0.5)
97 ax.set_xlabel("time [s]")
98 ax.set_ylabel("pressure [Pa]")
99 ax.set_title(f"{sound.plate.name}, {sound.info.get('model', '')} model, peak {sound.peak_spl:.0f} dB SPL")
100 return ax
103def spectrogram(s: Sound | Spectrogram, ax=None, *, db_range: float = 90.0, fmin: float = 20.0, fmax: float | None = None, log_freq: bool = True, modes=None, window: int = 2048, hop: int = 256):
104 """The spectrogram of a sound (or a precomputed ``Spectrogram``), with the
105 modal frequencies drawn as ticks on the left edge when ``modes`` is given."""
106 sp = s.spectrogram(window, hop) if isinstance(s, Sound) else s
107 ax = _ax(ax, (8, 4))
108 fmax = fmax or sp.freqs[-1]
109 sel = (sp.freqs >= (fmin if log_freq else 0)) & (sp.freqs <= fmax)
110 ax.pcolormesh(sp.times, sp.freqs[sel], sp.db[sel], cmap="inferno", vmin=-db_range, vmax=0, shading="nearest", rasterized=True)
111 if log_freq:
112 ax.set_yscale("log")
113 ax.set_ylim(fmin, fmax)
114 if modes is not None:
115 f = modes.freqs[(modes.freqs >= fmin) & (modes.freqs <= fmax)]
116 ax.plot(np.full(len(f), sp.times[0]), f, "_", color="cyan", ms=6, mew=1)
117 ax.set_xlabel("time [s]")
118 ax.set_ylabel("frequency [Hz]")
119 return ax
122def decay_times(modes, ax=None):
123 """Decay time of every mode against frequency, total and by mechanism."""
124 ax = _ax(ax, (8, 4))
125 by = modes.t60_by_mechanism()
126 ax.plot(modes.freqs, modes.t60, "k.", ms=4, label="total")
127 for name, marker in (("radiation", "b."), ("material", "g."), ("thermal", "r.")):
128 ax.plot(modes.freqs, by[name], marker, ms=2, alpha=0.6, label=name)
129 ax.set_xscale("log")
130 ax.set_yscale("log")
131 ax.set_ylim(0.01, 1e4)
132 ax.set_xlabel("frequency [Hz]")
133 ax.set_ylabel("T60 [s]")
134 ax.legend()
135 return ax
138def force(sound: Sound, ax=None):
139 """The stick force during the contact (linear model only, which records it)."""
140 if "force" not in sound.info:
141 raise ValueError("the force history is recorded by the linear model only")
142 ax = _ax(ax, (6, 3))
143 f = sound.info["force"]
144 ax.plot(np.arange(len(f)) * sound.info["force_dt"] * 1e3, f)
145 ax.set_xlabel("time [ms]")
146 ax.set_ylabel("force [N]")
147 ax.set_title(f"peak {sound.info['peak_force']:.0f} N, contact {sound.info['contact_time'] * 1e3:.2f} ms")
148 return ax
151def animate(sound: Sound, interval_ms: float | None = None, scale: float | None = None, stride: int = 1):
152 """A ``matplotlib.animation.FuncAnimation`` of the movie frames (display with ``HTML(anim.to_jshtml())`` in a notebook)."""
153 import matplotlib.pyplot as plt
154 from matplotlib.animation import FuncAnimation
156 if sound.frames is None:
157 raise ValueError("this sound has no movie frames; run with movie=True")
158 fig, ax = plt.subplots(figsize=(6, 5))
159 m = (scale or float(np.abs(sound.frames).max()) or 1e-9) * 1e3
160 mask = ~sound.plate.grid.mask
161 im = ax.imshow(np.ma.masked_where(mask, sound.frames[0] * 1e3), origin="lower", extent=_extent(sound.plate), cmap="RdBu_r", vmin=-m, vmax=m, interpolation="nearest")
162 fig.colorbar(im, ax=ax, label="w [mm]")
163 ax.set_aspect("equal")
164 title = ax.set_title("")
166 def update(i):
167 im.set_data(np.ma.masked_where(mask, sound.frames[i] * 1e3))
168 title.set_text(f"t = {sound.frame_time(i) * 1e3:.1f} ms")
169 return im, title
171 frames = range(0, len(sound.frames), stride)
172 return FuncAnimation(fig, update, frames=frames, interval=interval_ms or 1000 * stride / sound.fps, blit=False)