Finished a good part of analysis and figure for Thresh-LP invariance (WIP).

This commit is contained in:
j-hartling
2026-03-06 14:47:22 +01:00
parent 933d28b5f8
commit 0407053c20
15 changed files with 774 additions and 338 deletions

View File

@@ -1,224 +1,281 @@
import plotstyle_plt
import glob
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.transforms import BboxTransformTo
from itertools import product
from thunderhopper.filetools import search_files
from thunderhopper.modeltools import load_data
from color_functions import load_colors
from plot_functions import hide_axis, letter_subplots, ylimits, ylabel, super_xlabel, plot_line
from plot_functions import hide_axis, ylimits, xlabel, ylabel, plot_line, strip_zeros
from IPython import embed
def strip_zeros(num, right_digits=5):
if isinstance(num, int):
return num
num = f'{num:.{right_digits}f}'
left, right = num.split('.')
right = right.rstrip('0')
if right:
return f'{left}.{right}'
return left
def plot_snippets(axes, time, snippets, label, scales=None,
ymin=None, ymax=None, **kwargs):
ymin, ymax = ylimits(snippets, minval=ymin, maxval=ymax, pad=0.05)
ylabel(axes[0], label, x=-0.7, rotation=0, ha='left', va='center')
for i, (ax, snippet) in enumerate(zip(axes, snippets.T)):
plot_line(ax, time, snippet, ymin=ymin, ymax=ymax, **kwargs)
if scales is not None:
ax.set_title(f'$\\alpha={strip_zeros(scales[i])}$')
def time_bar(ax, dur, y0=0.9, y1=0.95, xshift=0.5, parent=None, transform=None, **kwargs):
t0, t1 = ax.get_xlim()
offset = (t1 - t0 - dur) * xshift
x0 = t0 + offset
x1 = x0 + dur
if parent is None:
parent = ax
if transform is None:
transform = BboxTransformTo(parent.bbox)
if transform is not ax.transData:
trans = ax.transData + transform.inverted()
x0 = trans.transform((x0, 0))[0]
x1 = trans.transform((x1, 0))[0]
parent.add_artist(plt.Rectangle((x0, y0), x1 - x0, y1 - y0,
transform=transform, **kwargs))
return None
def add_snip_axes(fig, grid_kwargs):
grid = fig.add_gridspec(**grid_kwargs)
axes = np.zeros((grid.nrows, grid.ncols), dtype=object)
for i, j in product(range(grid.nrows), range(grid.ncols)):
axes[i, j] = fig.add_subplot(grid[i, j])
[hide_axis(ax, 'left') for ax in axes.flatten()]
[hide_axis(ax, 'bottom') for ax in axes.flatten()]
return axes
def plot_snippets(axes, time, snippets, ymin=None, ymax=None, **kwargs):
ymin, ymax = ylimits(snippets, minval=ymin, maxval=ymax, pad=0.05)
for ax, snippet in zip(axes, snippets.T):
plot_line(ax, time, snippet, ymin=ymin, ymax=ymax, **kwargs)
return None
# GENERAL SETTINGS:
target = 'Omocestus_rufipes'
data_paths = glob.glob(f'../data/inv/log_hp/{target}*.npz')
data_paths = search_files(target, excl='noise', dir='../data/inv/log_hp/')
stages = ['env', 'log', 'inv']
files = stages + ['scales', 'plot_scales', 'snr_env', 'snr_log', 'snr_inv']
files = stages + ['scales', 'example_scales', 'limit',
'measure_env', 'measure_log', 'measure_inv']
save_path = '../figures/fig_invariance_log_hp.pdf'
# PLOT SETTINGS:
# GRAPH SETTINGS:
fig_kwargs = dict(
figsize=(32/2.54, 16/2.54),
)
grid1_kwargs = dict(
super_grid_kwargs = dict(
nrows=2,
ncols=2,
wspace=0,
hspace=0,
left=0,
right=1,
bottom=0,
top=1
)
subfig_specs = dict(
pure=(0, 0),
noise=(1, 0),
analysis=(slice(None), 1)
)
pure_grid_kwargs = dict(
nrows=len(stages),
ncols=None,
wspace=0.05,
hspace=0.3,
left=0.1,
right=0.985,
bottom=0.17,
hspace=0.1,
left=0.13,
right=0.95,
bottom=0.15,
top=0.9
)
grid2_kwargs = dict(
noise_grid_kwargs = dict(
nrows=len(stages),
ncols=None,
wspace=0.05,
hspace=0.1,
left=0.13,
right=0.95,
bottom=0.15,
top=0.9
)
analysis_grid_kwargs = dict(
nrows=1,
ncols=3,
wspace=0.35,
ncols=1,
wspace=0,
hspace=0,
left=0.1,
right=0.985,
bottom=0.18,
left=0.15,
right=0.96,
bottom=0.1,
top=0.95
)
snip_specs = dict(
env=(0, slice(None)),
log=(1, slice(None)),
inv=(2, slice(None))
)
# PLOT SETTINGS:
colors = load_colors('../data/stage_colors.npz')
lw_snippets = 0.25
lw_snippets = 0.5
lw_analysis = 3
xlabels = dict(
analysis='scale $\\alpha$',
)
xlab_analysis_kwargs = dict(
y=0.01,
fontsize=16,
ha='center',
va='bottom',
)
ylabels = dict(
env=r'$x_{\text{env}}$',
log=r'$x_{\text{dB}}$',
inv=r'$x_{\text{adapt}}$',
abs='abs. SNR',#'abs. ' + r'$\text{SNR}$',
norm='norm. SNR',#'norm. ' + r'$\text{SNR}$',
gain='rel. SNR gain'
env='$x_{\\text{env}}$',
log='$x_{\\text{dB}}$',
inv='$x_{\\text{adapt}}$',
analysis='ratio $\\text{SD}_{\\alpha}\\,/\\,\\text{SD}_{\\min[\\alpha]}$',
# analysis='ratio $\\sigma_{\\alpha}\\,/\\,\\sigma_{\\min[\\alpha]}$',
)
ylab_snip_kwargs = dict(
x=0.01,
fontsize=20,
rotation=0,
ha='left',
va='center',
)
ylab_analysis_kwargs = dict(
x=0.02,
fontsize=16,
ha='center',
va='top',
)
xloc = dict(
abs=5,
analysis=200,
)
yloc = dict(
env=100,
log=10,
inv=5,
abs=50,
)
letter_kwargs = dict(
x=0.03,
y=0.99,
letter_snip_kwargs = dict(
x=0.02,
y=1,
ha='left',
va='top'
va='top',
fontsize=22,
fontweight='bold'
)
fit_kwargs = dict(
c='darkred',
ls='--',
zorder=1.9
letter_analysis_kwargs = dict(
x=0,
y=1,
ha='left',
va='top',
fontsize=22,
fontweight='bold'
)
fit_lw = dict(
abs=6,
norm=3,
gain=3
)
text_kwargs = dict(
abs=dict(s='$\\alpha^2+1$', fontsize=16, c=fit_kwargs['c'],
x=0.85, y=0.9, ha='right', va='center'),
norm=dict(s='$\\alpha^2+1$', fontsize=16, c=fit_kwargs['c'],
x=0.85, y=0.9, ha='right', va='center'),
gain=dict(s='$\\frac{1}{\\alpha}$', fontsize=24, c=fit_kwargs['c'],
x=0.75, y=0.8, ha='left', va='center')
)
calculated_floor = False
floor_kwargs = dict(
xmin=0,
indicate_unsaturated = False
unsaturated_proportion = 0.85
unsaturated_kwargs = dict(
color=3 * (0.85,),
zorder=0,
lw=0
)
bar_time = 5
bar_kwargs = dict(
y0=0.5,
y1=0.6,
color='k',
lw = 0,
)
# EXECUTION:
for data_path in data_paths:
print(f'Processing {data_path}')
# Load invariance data:
data, config = load_data(data_path, files)
t_full = np.arange(data['env'].shape[0]) / config['env_rate']
nonzero_scales = data['scales'][data['scales'] > 0]
floor_kwargs['xmax'] = data['floor_scale'] if calculated_floor else 1
pure_data, config = load_data(data_path, files)
noise_data, _ = load_data(data_path.replace('.npz', '_noise.npz'), files)
t_full = np.arange(pure_data['env'].shape[0]) / config['env_rate']
# Prepare overall graph:
fig = plt.figure(**fig_kwargs)
super_grid = fig.add_gridspec(2, 1, wspace=0, hspace=1,
left=0, right=1, bottom=0, top=1)
super_grid = fig.add_gridspec(**super_grid_kwargs)
# Prepare snippet axes:
subfig1 = fig.add_subfigure(super_grid[0, :])
grid1_kwargs['ncols'] = data['plot_scales'].size
grid1 = subfig1.add_gridspec(**grid1_kwargs)
axes = np.zeros((grid1.nrows, grid1.ncols), dtype=object)
for i, j in product(range(grid1.nrows), range(grid1.ncols)):
axes[i, j] = subfig1.add_subplot(grid1[i, j])
[hide_axis(ax, 'bottom') for ax in axes[:-1, :].flatten()]
[hide_axis(ax, 'left') for ax in axes[:, 1:].flatten()]
super_xlabel('time [s]', subfig1, axes[-1, 0], axes[-1, -1], y=0)
letter_subplots(axes[:, 0], labels='abc', **letter_kwargs)
# Prepare pure-song snippet axes:
pure_subfig = fig.add_subfigure(super_grid[subfig_specs['pure']])
pure_grid_kwargs['nrows' if pure_grid_kwargs['nrows'] is None else 'ncols'] = pure_data['example_scales'].size
pure_axes = add_snip_axes(pure_subfig, pure_grid_kwargs)
for ax, stage in zip(pure_axes[:, 0], stages):
ylabel(ax, ylabels[stage], **ylab_snip_kwargs,
transform=pure_subfig.transSubfigure)
for ax, scale in zip(pure_axes[snip_specs['env']], pure_data['example_scales']):
ax.set_title(f'$\\alpha={strip_zeros(scale)}$')
pure_subfig.text(s='a', **letter_snip_kwargs)
# Prepare analysis axes:
subfig2 = fig.add_subfigure(super_grid[1, :])
grid2 = subfig2.add_gridspec(**grid2_kwargs)
symlog_kwargs = dict(linthresh=nonzero_scales.min(), linscale=0.2)
# Prepare noise-song snippet axes:
noise_subfig = fig.add_subfigure(super_grid[subfig_specs['noise']])
noise_grid_kwargs['nrows' if noise_grid_kwargs['nrows'] is None else 'ncols'] = noise_data['example_scales'].size
noise_grid = noise_subfig.add_gridspec(**noise_grid_kwargs)
noise_axes = add_snip_axes(noise_subfig, noise_grid_kwargs)
for ax, stage in zip(noise_axes[:, 0], stages):
ylabel(ax, ylabels[stage], **ylab_snip_kwargs,
transform=noise_subfig.transSubfigure)
for ax, scale in zip(noise_axes[snip_specs['env']], noise_data['example_scales']):
ax.set_title(f'$\\alpha={strip_zeros(scale)}$')
noise_subfig.text(s='b', **letter_snip_kwargs)
ax_abs_snr = subfig2.add_subplot(grid2[:, 0])
ax_abs_snr.set_ylabel(ylabels['abs'])
# Prepare analysis axis:
analysis_subfig = fig.add_subfigure(super_grid[subfig_specs['analysis']])
analysis_grid = analysis_subfig.add_gridspec(**analysis_grid_kwargs)
analysis_ax = analysis_subfig.add_subplot(analysis_grid[0, 0])
analysis_ax.set_xlim(noise_data['scales'].min(), noise_data['scales'].max())
analysis_ax.xaxis.set_major_locator(plt.MultipleLocator(xloc['analysis']))
xlabel(analysis_ax, xlabels['analysis'], **xlab_analysis_kwargs,
transform=analysis_subfig.transSubfigure)
analysis_ax.set_yscale('log')
ylabel(analysis_ax, ylabels['analysis'], **ylab_analysis_kwargs,
transform=analysis_subfig.transSubfigure)
analysis_subfig.text(s='c', **letter_analysis_kwargs)
ax_norm_snr = subfig2.add_subplot(grid2[:, 1])
ax_norm_snr.set_ylabel(ylabels['norm'])
ax_norm_snr.set_xscale('symlog', **symlog_kwargs)
ax_norm_snr.set_yscale('log')
# Plot pure-song envelope snippets:
plot_snippets(pure_axes[snip_specs['env']], t_full, pure_data['env'],
ymin=0, c=colors['env'], lw=lw_snippets)
ax_gain = subfig2.add_subplot(grid2[:, 2])
ax_gain.set_ylabel(ylabels['gain'])
ax_gain.set_xscale('symlog', **symlog_kwargs)
ax_gain.set_yscale('log')
super_xlabel('song scale $\\alpha$', subfig2, ax_abs_snr, ax_gain)
letter_subplots([ax_abs_snr, ax_norm_snr, ax_gain], labels='def', **letter_kwargs)
# Plot pure-song logarithmic snippets:
plot_snippets(pure_axes[snip_specs['log']], t_full, pure_data['log'],
ymax=None, c=colors['log'], lw=lw_snippets)
# Plot envelope snippets:
plot_snippets(axes[0, :], t_full, data['env'], ylabels['env'], ymin=0,
scales=data['plot_scales'], c=colors['env'], lw=lw_snippets)
axes[0, 0].yaxis.set_major_locator(plt.MultipleLocator(yloc['env']))
# Plot logarithmic snippets:
plot_snippets(axes[1, :], t_full, data['log'], ylabels['log'], ymax=0,
c=colors['log'], lw=lw_snippets)
axes[1, 0].yaxis.set_major_locator(plt.MultipleLocator(yloc['log']))
# Plot invariant snippets:
plot_snippets(axes[2, :], t_full, data['inv'], ylabels['inv'],
# Plot pure-song invariant snippets:
plot_snippets(pure_axes[snip_specs['inv']], t_full, pure_data['inv'],
c=colors['inv'], lw=lw_snippets)
axes[2, 0].yaxis.set_major_locator(plt.MultipleLocator(yloc['inv']))
# Plot in-representation SNRs (absolute):
ax_abs_snr.plot(data['scales'], data['snr_env'], c=colors['env'], lw=lw_analysis)
ax_abs_snr.plot(data['scales'], data['snr_log'], c=colors['log'], lw=lw_analysis)
ax_abs_snr.plot(data['scales'], data['snr_inv'], c=colors['inv'], lw=lw_analysis)
ax_abs_snr.axvspan(**floor_kwargs)
ax_abs_snr.xaxis.set_major_locator(plt.MultipleLocator(xloc['abs']))
ax_abs_snr.yaxis.set_major_locator(plt.MultipleLocator(yloc['abs']))
ax_abs_snr.set_ylim(0, data['snr_env'].max())
# Indicate time scale:
time_bar(pure_axes[snip_specs['env']][0], bar_time, **bar_kwargs)
# Plot envelope SNR fit:
ax_abs_snr.plot(data['scales'], data['scales'] ** 2 + 1,
lw=fit_lw['abs'], **fit_kwargs)
ax_abs_snr.text(transform=ax_abs_snr.transAxes, **text_kwargs['abs'])
# Plot noise-song envelope snippets:
plot_snippets(noise_axes[snip_specs['env']], t_full, noise_data['env'],
ymin=0, c=colors['env'], lw=lw_snippets)
# Get normalized SNRs:
norm_snr_env = data['snr_env'] / data['snr_env'].max()
norm_snr_log = data['snr_log'] / data['snr_log'].max()
norm_snr_inv = data['snr_inv'] / data['snr_inv'].max()
# Plot noise-song logarithmic snippets:
plot_snippets(noise_axes[snip_specs['log']], t_full, noise_data['log'],
ymax=None, c=colors['log'], lw=lw_snippets)
# Plot in-representation SNRs (normalized):
ax_norm_snr.plot(data['scales'], norm_snr_env, c=colors['env'], lw=lw_analysis)
ax_norm_snr.plot(data['scales'], norm_snr_log, c=colors['log'], lw=lw_analysis)
ax_norm_snr.plot(data['scales'], norm_snr_inv, c=colors['inv'], lw=lw_analysis)
ax_norm_snr.axvspan(**floor_kwargs)
ax_norm_snr.set_ylim(norm_snr_env.min(), 1)
# Plot noise-song invariant snippets:
plot_snippets(noise_axes[snip_specs['inv']], t_full, noise_data['inv'],
c=colors['inv'], lw=lw_snippets)
# # Plot envelope SNR fit:
# ax_norm_snr.plot(nonzero_scales, nonzero_scales / nonzero_scales.max(),
# lw=fit_lw['norm'], **fit_kwargs)
# ax_norm_snr.text(transform=ax_norm_snr.transAxes, **text_kwargs['norm'])
# Indicate time scale:
time_bar(noise_axes[snip_specs['env']][0], bar_time, **bar_kwargs)
# Get relative SNR gain:
gain_log = norm_snr_log / norm_snr_env
gain_inv = norm_snr_inv / norm_snr_env
# Plot pure-song SD ratios (ideal):
base_ind = np.argmin(pure_data['scales'])
# measure_env = pure_data['measure_env'] / pure_data['measure_env'][base_ind]
# measure_log = pure_data['measure_log'] / pure_data['measure_log'][base_ind]
measure_inv = pure_data['measure_inv'] / pure_data['measure_inv'][base_ind]
# analysis_ax.plot(pure_data['scales'], measure_env, c=colors['env'], lw=lw_analysis, ls='--')
# analysis_ax.plot(pure_data['scales'], measure_log, c=colors['log'], lw=lw_analysis, ls='--')
analysis_ax.plot(pure_data['scales'], measure_inv, c=colors['inv'], lw=lw_analysis, ls='--')
# Plot across-representation gain:
ax_gain.plot(data['scales'], gain_log, c=colors['log'], lw=lw_analysis)
ax_gain.plot(data['scales'], gain_inv, c=colors['inv'], lw=lw_analysis)
ax_gain.axvspan(**floor_kwargs)
ax_gain.set_ylim(1, 10)
if indicate_unsaturated:
# Indicate influence of noise floor:
limit = noise_data['limit'] * unsaturated_proportion
thresh_ind = np.nonzero(noise_data['measure_inv'] <= limit)[0][-1]
analysis_ax.axvspan(0, noise_data['scales'][thresh_ind], **unsaturated_kwargs)
# Plot noise-song SD ratios (limited):
base_ind = np.argmin(noise_data['scales'])
measure_env = noise_data['measure_env'] / noise_data['measure_env'][base_ind]
measure_log = noise_data['measure_log'] / noise_data['measure_log'][base_ind]
measure_inv = noise_data['measure_inv'] / noise_data['measure_inv'][base_ind]
analysis_ax.plot(noise_data['scales'], measure_env, c=colors['env'], lw=lw_analysis)
analysis_ax.plot(noise_data['scales'], measure_log, c=colors['log'], lw=lw_analysis)
analysis_ax.plot(noise_data['scales'], measure_inv, c=colors['inv'], lw=lw_analysis)
analysis_ax.set_ylim(0.1, measure_env.max())
# Plot amplification fit:
ax_gain.plot(nonzero_scales, nonzero_scales.max() / nonzero_scales,
lw=fit_lw['gain'], **fit_kwargs)
ax_gain.text(transform=ax_gain.transAxes, **text_kwargs['gain'])
if save_path is not None:
fig.savefig(save_path)
plt.show()