Add onaxis button

ipi_extract.py

@@ -263,6 +263,7 @@ class Callback(object):
         self.reset_b = None
         self.r_button = None
         self.ind_b = None
+        self.onaxis_b = None
         self.spec_b = None
         self.nfft = 128
         self.ind_select = False
@@ -456,6 +457,15 @@ class Callback(object):
         plt.pause(0.2)
         self.fax.get_figure().set_constrained_layout(False)
 
+    def onaxis(self, event):
+        if self.df[self.offset[self.curr_ind[self.curr], 0]]['onaxis'] == -1:
+            self.df[self.offset[self.curr_ind[self.curr], 0]]['onaxis'] = 1
+        else:
+            self.df[self.offset[self.curr_ind[self.curr], 0]]['onaxis'] ^= 1
+        self.onaxis_b.label.set_text('On-axis'
+                                     if self.df[self.offset[self.curr_ind[self.curr], 0]]['onaxis'] else 'Off-axis')
+        plt.draw()
+
     def increase_freq(self, event):
         self.view_ax[self.curr][1].set_ylim(0, self.view_ax[self.curr][1].get_ylim()[1]+1e3)
         plt.draw()
@@ -504,6 +514,9 @@ class Callback(object):
         self.view_ax[ind][0].set_xlabel(f'Sig man:{dic["ipi_sig"]:.3f}')
         self.view_ax[ind][2].set_xlabel(f'Corr man:{dic["ipi_corr_man"]:.3f} auto:{dic["ipi_corr_auto"]:.3f}')
         self.view_ax[ind][3].set_xlabel(f'Ceps man:{dic["ipi_ceps_man"]:.3f} auto:{dic["ipi_ceps_auto"]:.3f}')
+        self.ind_b.label.set_text(f'Current individual:\n{dic["ind_number"]}')
+        if self.onaxis_b is not None:
+            self.onaxis_b.label.set_text(['Off', 'On', '?'][dic['onaxis']] + '-axis')
 
     def _set_visible(self, ind=None, state=False):
         if ind is None:
@@ -720,9 +733,12 @@ def main(args):
         print(f'Out file {outpath} does not already exist and resume option is set.')
         return 1
 
+    EMLN['onaxis'] = -1
     ref_dict = init(args.input, args.channel)
     if args.resume:
         df = pd.read_hdf(outpath)
+        if 'onaxis' not in df.columns:
+            df['onaxis'] = -1
         ref_dict['callback'].df = df.to_dict(orient='index')
         ref_dict['callback'].offset = np.tile(np.array(list(ref_dict['callback'].df.keys()))[:, np.newaxis], (1,2))
         ref_dict['callback'].offset[:, 1] = ref_dict['callback'].song_resample[(ref_dict['callback'].offset[:, 1] * FSSR).astype(int)]
@@ -732,6 +748,11 @@ def main(args):
             if not np.all(pd.isnull(list(ref_dict['callback'].df[p].values()))):
                 colors[i] = [0, 0, 0, 1]
         ref_dict['callback'].scat.set_color(colors)
+
+    onaxis_b_ax = plt.subplot(ref_dict['gridspec'][-1, 8:10])
+    onaxis_b = Button(onaxis_b_ax, '?-axis')
+    onaxis_b.on_clicked(ref_dict['callback'].onaxis)
+    ref_dict['callback'].onaxis_b = onaxis_b
     plt.draw()
     plt.pause(0.2)
     ref_dict['fig'].set_constrained_layout(False)