import plotstyle_plt import numpy as np import matplotlib.pyplot as plt from itertools import product from thunderhopper.filetools import search_files from thunderhopper.modeltools import load_data from thunderhopper.filtertools import find_kern_specs from color_functions import load_colors from plot_functions import hide_axis, xlimits, ylimits, xlabel, ylabel, super_ylabel,\ plot_line, plot_barcode, strip_zeros, time_bar from IPython import embed 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 def plot_bi_snippets(axes, time, binary, **kwargs): for ax, binary in zip(axes, binary.T): plot_barcode(ax, time, binary[:, None], **kwargs) return None def side_distributions(axes, snippets, inset_bounds, thresh, ymin=None, ymax=None): bins = np.linspace(snippets.min(), snippets.max(), 50) centers = bins[:-1] + (bins[1] - bins[0]) / 2 for ax, snippet in zip(axes, snippets.T): inset = ax.inset_axes(inset_bounds) inset.axis('off') pdf, _ = np.histogram(snippet, bins, density=True) inset.plot(pdf, centers, c='k', lw=1) inset.fill_betweenx(centers, pdf.min(), pdf, where=(centers > thresh), color=colors['bi'], lw=0) inset.set_xlim(0, pdf.max()) ylimits(centers, inset, minval=ymin, maxval=ymax, pad=0) return None # GENERAL SETTINGS: with_noise = True target = 'Omocestus_rufipes' search_kwargs = dict( incl='subset' if not with_noise else 'subset_noise', dir='../data/inv/thresh_lp/' ) data_paths = search_files(target, **search_kwargs) stages = ['conv', 'bi', 'feat'] load_kwargs = dict( files=stages, keywords=['scales', 'snip', 'measure', 'thresh'] ) save_path = None#'../figures/fig_invariance_thresh_lp_single' if with_noise and save_path is not None: save_path += '_noise' # GRAPH SETTINGS: fig_kwargs = dict( figsize=(32/2.54, 16/2.54), ) super_grid_kwargs = dict( nrows=None, ncols=2, wspace=0, hspace=0, left=0, right=1, bottom=0, top=1 ) snip_grid_kwargs = dict( nrows=len(stages), ncols=None, wspace=0.11, hspace=0.1, left=0.1, right=0.95, bottom=0.01, top=0.85 ) big_grid_kwargs = dict( nrows=1, ncols=1, wspace=0, hspace=0, left=0.15, right=0.96, bottom=0.1, top=0.99 ) inset_bounds = [1, 0, 0.1, 1] # PLOT SETTINGS: colors = load_colors('../data/stage_colors.npz') # lw_snippets = dict( # conv=0.5, # feat=2 # ) # lw_analysis = 3 xlabels = dict( big='scale $\\alpha$', ) xlab_big_kwargs = dict( y=0.01, fontsize=16, ha='center', va='bottom', ) ylabels = dict( conv='$c_i$', bi='$b_i$', feat='$f_i$', big='$\\mu_f$', ) ylab_snip_kwargs = dict( x=0.08, fontsize=20, rotation=0, ha='right', va='center', ) ylab_super_kwargs = dict( x=0.005, fontsize=16, ha='left', va='center', ) ylab_big_kwargs = dict( x=0.02, fontsize=20, ha='center', va='top', ) # xloc = dict( # analysis=10, # ) # letter_snip_kwargs = dict( # x=0.02, # y=1, # ha='left', # va='top', # fontsize=22, # fontweight='bold' # ) # letter_analysis_kwargs = dict( # x=0, # y=1, # ha='left', # va='top', # fontsize=22, # fontweight='bold' # ) # bar_time = 5 # bar_kwargs = dict( # y0=0.7, # y1=0.8, # color='k', # lw=0, # ) kernel = np.array([ [2, 0.008], [4, 0.008], ])[:1] zoom_rel = np.array([0.5, 0.55]) # EXECUTION: for data_path in data_paths: print(f'Processing {data_path}') # Load invariance data: data, config = load_data(data_path, **load_kwargs) t_full = np.arange(data['snip_conv'].shape[0]) / config['env_rate'] zoom_abs = zoom_rel * t_full[-1] zoom_inds = (t_full >= zoom_abs[0]) & (t_full <= zoom_abs[1]) kern_ind = find_kern_specs(config['k_specs'], kerns=kernel)[0] # Reduce to kernel subset and crop time to zoom frame: data['snip_conv'] = data['snip_conv'][zoom_inds, kern_ind, ...] data['snip_bi'] = data['snip_bi'][zoom_inds, kern_ind, ...] data['snip_feat'] = data['snip_feat'][zoom_inds, kern_ind, ...] data['measure_conv'] = data['measure_conv'][:, kern_ind, :] data['measure_feat'] = data['measure_feat'][:, kern_ind, :] data['threshs'] = data['threshs'][:, kern_ind] t_full = np.arange(data['snip_conv'].shape[0]) / config['env_rate'] # Adjust grid parameters: super_grid_kwargs['nrows'] = data['thresh_perc'].size snip_grid_kwargs['ncols'] = data['example_scales'].size # Prepare overall graph: fig = plt.figure(**fig_kwargs) super_grid = fig.add_gridspec(**super_grid_kwargs) # Prepare analysis axis: big_subfig = fig.add_subfigure(super_grid[slice(None), 1]) big_grid = big_subfig.add_gridspec(**big_grid_kwargs) big_ax = big_subfig.add_subplot(big_grid[0, 0]) xlabel(big_ax, xlabels['big'], **xlab_big_kwargs, transform=big_subfig.transSubfigure) ylabel(big_ax, ylabels['big'], **ylab_big_kwargs, transform=big_subfig.transSubfigure) big_ax.set_xlim(data['scales'].min(), data['scales'].max()) ylimits(data['measure_feat'], big_ax, minval=0, pad=0.05) big_ax.set_xscale('symlog', linthresh=data['scales'][1], linscale=0.5) # Prepare snippet axes: snip_axes = {} for i in range(data['thresh_perc'].size): snip_subfig = fig.add_subfigure(super_grid[i, 0]) axes = add_snip_axes(snip_subfig, snip_grid_kwargs) snip_axes[snip_subfig] = axes super_ylabel(f'{data["thresh_perc"][i]}%', snip_subfig, axes[0, 0], axes[-1, 0], **ylab_super_kwargs) for ax, stage in zip(axes[:, 0], stages): ylabel(ax, ylabels[stage], **ylab_snip_kwargs, transform=snip_subfig.transSubfigure) if i == 0: for ax, scale in zip(axes[0, :], data['example_scales']): ax.set_title(f'$\\alpha={strip_zeros(scale)}$') # Plot representation snippets per threshold: for i, (subfig, axes) in enumerate(snip_axes.items()): # Plot kernel response snippets: plot_snippets(axes[0, :], t_full, data['snip_conv'][:, :, i], c=colors['conv'], lw=0.5) # Plot binary snippets: plot_bi_snippets(axes[1, :], t_full, data['snip_bi'][:, :, i], color=colors['bi'], lw=0) # Plot feature snippets: plot_snippets(axes[2, :], t_full, data['snip_feat'][:, :, i], ymin=0, ymax=1, c=colors['feat'], lw=2) # Plot kernel response distribution: side_distributions(axes[0, :], data['snip_conv'][:, :, i], inset_bounds, data['threshs'][i]) # Plot analysis results: big_ax.plot(data['scales'], data['measure_feat'], c=colors['feat'], lw=3) # # 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['conv']], pure_data['example_scales']): # ax.set_title(f'$\\alpha={strip_zeros(scale)}$') # pure_subfig.text(s='a', **letter_snip_kwargs) # # 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) # # Plot pure-song kernel response snippets: # plot_snippets(pure_axes[snip_specs['conv']], t_full, pure_data['conv'], # c=colors['conv'], lw=lw_snippets['conv']) # # Plot pure-song binary snippets: # plot_bi_snippets(pure_axes[snip_specs['bi']], t_full, pure_data['bi'], # color=colors['bi'], lw=0) # # Plot pure-song feature snippets: # plot_snippets(pure_axes[snip_specs['feat']], t_full, pure_data['feat'], # ymin=0, ymax=1, c=colors['feat'], lw=lw_snippets['feat']) # # Indicate time scale: # time_bar(pure_axes[snip_specs['conv']][0], bar_time, **bar_kwargs) # # Plot noise-song kernel response snippets: # plot_snippets(noise_axes[snip_specs['conv']], t_full, noise_data['conv'], # c=colors['conv'], lw=lw_snippets['conv']) # # Plot noise-song binary snippets: # plot_bi_snippets(noise_axes[snip_specs['bi']], t_full, noise_data['bi'], # color=colors['bi'], lw=0) # # Plot noise-song feature snippets: # plot_snippets(noise_axes[snip_specs['feat']], t_full, noise_data['feat'], # ymin=0, ymax=1, c=colors['feat'], lw=lw_snippets['feat']) # # Indicate time scale: # time_bar(noise_axes[snip_specs['conv']][0], bar_time, **bar_kwargs) # # Plot noise-song SD ratios (limited): # analysis_ax.plot(noise_data['scales'], noise_data['measure_conv'], # c=colors['conv'], lw=lw_analysis) # lower, upper = noise_data['spread_conv'] # analysis_ax.fill_between(noise_data['scales'], lower, upper, # color=colors['conv'], **spread_kwargs) # analysis_ax.plot(noise_data['scales'], noise_data['measure_feat'], # c=colors['feat'], lw=lw_analysis) # lower, upper = noise_data['spread_feat'] # analysis_ax.fill_between(noise_data['scales'], lower, upper, # color=colors['feat'], **spread_kwargs) if save_path is not None: fig.savefig(save_path) plt.show() print('Done.') embed()