Nearly finished 1st draft of species-specific Thresh-LP invariance figure (WIP).

python/fig_invariance_thresh-lp_subset.py

@@ -0,0 +1,160 @@
+import plotstyle_plt
+import numpy as np
+import matplotlib.pyplot as plt
+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, shade_colors
+from plot_functions import hide_axis, ylimits, xlabel, ylabel, super_ylabel,\
+                           plot_line, plot_barcode, strip_zeros, time_bar,\
+                           letter_subplot, letter_subplots, hide_ticks,\
+                           super_xlabel, super_ylabel, assign_colors
+from IPython import embed
+
+# GENERAL SETTINGS:
+target = 'Omocestus_rufipes'
+search_kwargs = dict(
+    incl='subset',
+    excl='noise',
+    dir='../data/inv/thresh_lp/'
+)
+pure_paths = search_files(target, **search_kwargs)
+load_kwargs = dict(
+    keywords=['scales', 'measure', 'thresh']
+)
+save_path = None#'../figures/fig_invariance_thresh_lp_subset.pdf'
+
+# GRAPH SETTINGS:
+fig_kwargs = dict(
+    figsize=(32/2.54, 16/2.54),
+)
+super_grid_kwargs = dict(
+    nrows=1,
+    ncols=1,
+    wspace=0,
+    hspace=0,
+    left=0,
+    right=1,
+    bottom=0,
+    top=1
+)
+grid_kwargs = dict(
+    nrows=2,
+    ncols=1,
+    wspace=0,
+    hspace=0.1,
+    left=0.15,
+    right=0.95,
+    bottom=0.1,
+    top=0.85
+)
+inset_bounds = [0.2, 1.01, 0.6, 0.4]
+
+# PLOT SETTINGS:
+colors = load_colors('../data/stage_colors.npz')
+color_factors = [-0.5, 0.5]
+lw = dict(
+    one=3,
+    kern=3,
+    all=1,
+)
+ax_labels = dict(
+    x='scale $\\alpha$',
+    y='$\\mu_f$',
+)
+xlab_kwargs = dict(
+    y=0.005,
+    fontsize=16,
+    ha='center',
+    va='bottom',
+)
+ylab_kwargs = dict(
+    x=0,
+    fontsize=20,
+    ha='left',
+    va='center',
+)
+yloc = 0.2
+
+# EXECUTION:
+for pure_path in pure_paths:
+    print(f'Processing {pure_path}')
+    noise_path = pure_path.replace('.npz', '_noise.npz')
+
+    # Load kernel invariance data:
+    pure_data, config = load_data(pure_path, **load_kwargs)
+    noise_data, _ = load_data(noise_path, **load_kwargs)
+    scales = pure_data['scales']
+
+    # Adjust grid parameters:
+    n_columns = config['k_specs'].shape[0] + 1
+    super_grid_kwargs['ncols'] = n_columns
+
+    # Prepare overall graph:
+    fig = plt.figure(**fig_kwargs)
+    super_grid = fig.add_gridspec(**super_grid_kwargs)
+
+    # Prepare axes:
+    all_axes = np.zeros((grid_kwargs['nrows'], n_columns), dtype=object)
+    subfigs = []
+    for i in range(n_columns):
+        subfig = fig.add_subfigure(super_grid[0, i])
+        grid = subfig.add_gridspec(**grid_kwargs)
+        subfigs.append(subfig)
+        for j in range(grid_kwargs['nrows']):
+            ax = subfig.add_subplot(grid[j, 0])
+            ax.set_xlim(scales[0], scales[-1])
+            ax.set_ylim(0, 1)
+            ax.set_xscale('symlog', linthresh=scales[1], linscale=0.5)
+            ax.yaxis.set_major_locator(plt.MultipleLocator(yloc))
+            if i > 0:
+                hide_ticks(ax, side='left')
+            all_axes[j, i] = ax
+        hide_ticks(all_axes[0, i], side='bottom')
+    super_xlabel(ax_labels['x'], fig, all_axes[-1, 0], all_axes[-1, -1], **xlab_kwargs)
+    super_ylabel(ax_labels['y'], fig, all_axes[0, 0], all_axes[1, 0], **ylab_kwargs)
+
+    # Plot kernel-specific results:
+    in_min, in_high = ylimits(config['kernels'], pad=0.05)
+    for i in range(config['k_specs'].shape[0]):
+        pure_ax, noise_ax = all_axes[:, i]
+        # Plot results of pure-song analysis:
+        pure_ax.plot(scales, pure_data['measure_feat'][:, i, :],
+                     c=colors['feat'], lw=lw['one'])
+        # Plot results of noise-song analysis:
+        noise_ax.plot(scales, noise_data['measure_feat'][:, i, :],
+                      c=colors['feat'], lw=lw['one'])
+        # Indicate kernel waveform:
+        inset = pure_ax.inset_axes(inset_bounds)
+        inset.plot(config['k_times'], config['kernels'][:, i], c='k', lw=lw['kern'])
+        inset.set_xlim(config['k_times'][0], config['k_times'][-1])
+        inset.set_ylim(in_min, in_high)
+        inset.axis('off')
+
+    # Load population invariance data:
+    pure_data, config = load_data(pure_path.replace('_subset', ''), **load_kwargs)
+    noise_data, _ = load_data(noise_path.replace('_subset', ''), **load_kwargs)
+    scales = pure_data['scales']
+
+    # Get kernel type-specific colors:
+    types, ind = np.unique(config['k_specs'][:, 0], return_index=True)
+    types = types[np.argsort(ind)].astype(int)
+    factors = np.linspace(*color_factors, types.size)
+    kern_colors = shade_colors(colors['feat'], factors)
+    kern_colors = dict(zip(types.astype(str), kern_colors))
+
+    # Plot population-wide results:
+    pure_ax, noise_ax = all_axes[:, -1]
+    handles = pure_ax.plot(scales, pure_data['measure_feat'], c='k', lw=lw['all'])
+    assign_colors(handles, config['k_specs'][:, 0], kern_colors)
+
+    handles = noise_ax.plot(scales, noise_data['measure_feat'], c='k', lw=lw['all'])
+    assign_colors(handles, config['k_specs'][:, 0], kern_colors)
+
+
+    if save_path is not None:
+        fig.savefig(save_path)
+    plt.show()
+
+print('Done.')
+embed()