update display_2d make_animation_board

epileptor/display_2d.py

@@ -74,11 +74,12 @@ def make_animation_board(dt, state, data_filename, points):
         ax.set_axis_off()
         ax.set_title('{}'.format(util.dt_to_sec(shot_i, dt)))
         state_img = ax.imshow(state.values[shot_i, :, :], cmap='jet')
-        for point in points:
+        for point in [s1, s2]:
             ax.text(point['x'] - 2, point['y'] + 1, 'x', {'color': 'w', 'fontsize': 5, 'weight': 'bold'})
         state_img.set_clim(state.min, state.max)
         return state_img
 
+    s1, s2 = points.get(0), points.get(1)
     peak_times = util.get_peak_times(state.values, util.sec_to_dt(20, dt))
     nrows = len(peak_times)
     ncols = 8
@@ -96,9 +97,7 @@ def make_animation_board(dt, state, data_filename, points):
         im = add_to_board(t_shots_dt[i], ax)
         if i == 0:
             ax.set_title(ax.get_title() + 's')
-            s1 = points.get(0)
-            ax.text(s1['x'] - 18, s1['y'] - 2, 'S1', {'color': 'w', 'fontsize': 10})
-            s2 = points.get(1)
+            ax.text(s1['x'] - 15, s1['y'] - 2, 'S1', {'color': 'w', 'fontsize': 10})
             ax.text(s2['x'] - 15, s2['y'] + 2, 'S2', {'color': 'w', 'fontsize': 10})
     fig.subplots_adjust(wspace=0.05)
     fig.colorbar(