diff --git a/plotting.py b/plotting.py index 8fba0c09f7ca17be400f8b30853f71a8b8fb94aa..4e4f1ceba92b30e9ed00ceef96422bc0eaa6c950 100644 --- a/plotting.py +++ b/plotting.py @@ -121,6 +121,20 @@ def format_axis_radians(yaxis): yaxis.set_major_formatter(mpl.ticker.FuncFormatter(lambda x, p: '0' if x==0 else '$\\pi$' if x==np.pi else '$-\\pi$' if x==-np.pi else f'${x/np.pi:g}\\pi$')) yaxis.set_minor_formatter(mpl.ticker.FuncFormatter(lambda x, p: '0' if x==0 else '$\\pi$' if x==np.pi else '$-\\pi$' if x==-np.pi else None)) +def format_axis_degrees(yaxis): + yaxis.set_major_locator(mpl.ticker.MaxNLocator('auto', steps=[1, 1.5, 3, 6, 9, 10])) + yaxis.set_minor_locator(mpl.ticker.MultipleLocator(30)) + + +def add_scale(ax, scale, text=None, *, size=0.01, padding=0.1, loc='lower right', color='k', fontsize='x-small'): + """Make a scale or yardstick patch""" + aux = mpl.offsetbox.AuxTransformBox(mpl.transforms.blended_transform_factory(ax.transData, ax.transAxes)) + aux.add_artist(plt.Rectangle((0,0), scale, size, fc=color)) + if text: aux.add_artist(plt.Text(0, size, text, color=color, va='bottom', fontsize=fontsize)) + ab = mpl.offsetbox.AnchoredOffsetbox(loc, borderpad=padding, zorder=100, frameon=False) + ab.set_child(aux) + ax.add_artist(ab) + @@ -181,8 +195,6 @@ def plot_btf(axf, axp, data, *, frev=None, **kwargs): axf.set(ylabel=f'Magnitude / {data.m_unit}', xlim=(np.min(f), np.max(f))) axp.set(ylabel=f'Phase / {data.p_unit}', xlabel='Stimulus tune' if frev else f'Stimulus frequency / {data.f_unit}') - - return fig, (af, ap) def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tune_range=None, fit=False, smoothing=None, return_spectrum=False, **kwargs): @@ -219,6 +231,7 @@ def plot_tune_spectrum(ax, libera_data, xy, turn_range=None, time_range=None, tu ylabel='a.u.') if fit: + q = None try: fitr = (fit if callable(fit) else fit_lorenzian)(freq, mag) q, w = fitr[0][2], fitr[0][3] @@ -350,7 +363,7 @@ def plot_tune_spectrogram(ax, libera_data, xy, *, nperseg=2**12, noverlap=None, ax.get_figure().colorbar(cm, label='FFT magnitude', ax=ax) if show_nperseg and not over_time: - ax.plot(ax.get_xlim()[1]*0.97-np.array((nperseg, 0)), [0.05,0.05], 'k|--', transform=ax.get_xaxis_text1_transform(0)[0]) + add_scale(ax, nperseg) # indicate excitation region