Lots of stuff. Syncing to home.

python/fig_invariance_thresh-lp_single.py

@@ -6,17 +6,18 @@ 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,\
+from plot_functions import hide_axis, ylimits, xlabel, ylabel, super_ylabel, super_xlabel,\
                            plot_line, plot_barcode, strip_zeros, time_bar,\
-                           letter_subplot, letter_subplots
+                           letter_subplot, letter_subplots, title_subplot
 from IPython import embed
 
-def add_snip_axes(fig, grid_kwargs):
+def add_snip_axes(fig, grid_kwargs, snip_specs):
     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()]
+    axes = np.zeros((len(snip_specs), grid.ncols), dtype=object)
+    for i, specs in enumerate(snip_specs.values()):
+        for j in range(grid.ncols):
+            axes[i, j] = fig.add_subplot(grid[specs, j])
+    [hide_axis(ax, 'left') for ax in axes[:, 1:].flatten()]
     [hide_axis(ax, 'bottom') for ax in axes.flatten()]
     return axes
 
@@ -88,14 +89,14 @@ subfig_specs = dict(
     big=(slice(None), -1),
 )
 snip_grid_kwargs = dict(
-    nrows=len(stages),
+    nrows=len(stages) + 1,
     ncols=None,
     wspace=0.3,
-    hspace=0.1,
-    left=0.1,
+    hspace=0,
+    left=0.17,
     right=0.93,
-    bottom=0.05,
-    top=0.85
+    bottom=0.15,
+    top=0.9
 )
 big_grid_kwargs = dict(
     nrows=1,
@@ -107,6 +108,11 @@ big_grid_kwargs = dict(
     bottom=0.1,
     top=0.99
 )
+snip_specs = dict(
+    conv=slice(0, -2),
+    bi=-2,
+    feat=-1
+)
 inset_bounds = [1.02, 0, 0.2, 1]
 
 # PLOT SETTINGS:
@@ -119,20 +125,27 @@ lw = dict(
     big=4,
 )
 xlabels = dict(
+    snip='time [s]',
     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$',
 )
+xlab_snip_kwargs = dict(
+    y=0,
+    fontsize=16,
+    ha='center',
+    va='bottom',
+)
+xlab_big_kwargs = dict(
+    y=0,
+    fontsize=16,
+    ha='center',
+    va='bottom',
+)
 ylab_snip_kwargs = dict(
     x=0.08,
     fontsize=20,
@@ -141,7 +154,7 @@ ylab_snip_kwargs = dict(
     va='center',
 )
 ylab_super_kwargs = dict(
-    x=0.005,
+    x=0,
     fontsize=16,
     ha='left',
     va='center',
@@ -153,11 +166,21 @@ ylab_big_kwargs = dict(
     va='top',
 )
 yloc = dict(
+    conv=100,
+    bi=1,
+    feat=1,
     big=0.2,
 )
+title_kwargs = dict(
+    x=0.5,
+    yref=1,
+    ha='center',
+    va='top',
+    fontsize=16,
+)
 letter_snip_kwargs = dict(
-    x=0.01,
-    y=0.9,
+    x=0,
+    y=1,
     ha='left',
     va='top',
     fontsize=22,
@@ -178,7 +201,7 @@ dist_fill_kwargs = dict(
     color=colors['bi'],
     lw=0.1,
 )
-bar_time = 0.5
+bar_time = 0.1
 bar_kwargs = dict(
     y0=0.3,
     y1=0.4,
@@ -190,7 +213,7 @@ kernel = np.array([
     [2, 0.004],
     [3, 0.002],
 ])[np.array([1])]
-zoom_rel = np.array([0.5, 0.525])
+zoom_rel = np.array([0.5, 0.515])
 
 
 # EXECUTION:
@@ -232,17 +255,22 @@ for data_path in data_paths:
     for i in range(data['threshs'].size):
         subfig_specs['snip'] = (i, subfig_specs['snip'][1])
         snip_subfig = fig.add_subfigure(super_grid[subfig_specs['snip']])
-        axes = add_snip_axes(snip_subfig, snip_grid_kwargs)
+        axes = add_snip_axes(snip_subfig, snip_grid_kwargs, snip_specs)
         snip_axes[snip_subfig] = axes
         super_ylabel(f'{strip_zeros(100 * data["thresh_perc"][i])}%',
                      snip_subfig, axes[-1, 0], axes[0, 0], **ylab_super_kwargs)
         for ax, stage in zip(axes[:, 0], stages):
+            ax.yaxis.set_major_locator(plt.MultipleLocator(yloc[stage]))
             ylabel(ax, ylabels[stage], **ylab_snip_kwargs,
                    transform=snip_subfig.transSubfigure)
         if i == 0:
+            axes[0, 0].set_xlim(t_full[0], t_full[-1])
             time_bar(axes[0, 0], bar_time, **bar_kwargs)
             for ax, scale in zip(axes[0, :], data['example_scales']):
-                ax.set_title(f'$\\alpha={strip_zeros(scale)}$')
+                title = f'$\\alpha={strip_zeros(scale)}$'
+                title_subplot(ax, title, **title_kwargs, ref=fig)
+        elif i == data['threshs'].size - 1: 
+            super_xlabel(xlabels['snip'], snip_subfig, axes[-1, 0], axes[-1, -1], **xlab_snip_kwargs)
     letter_subplots(snip_axes.keys(), **letter_snip_kwargs)
 
     # Prepare analysis axis: