display_2 center cut time evolution

epileptor/display_2d.py

@@ -7,7 +7,7 @@ from mpl_toolkits.mplot3d import Axes3D
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
-from matplotlib import animation
+from matplotlib import animation, ticker, patches
 
 from epileptor import display_0d
 from epileptor import util
@@ -194,15 +194,36 @@ def center_cut_t_evolution(dt, state, data_filename):
     values = state.values
     # Those values should be set manually. Now they are set for 2019-04-04_19.16
     # 32800 - 33100, 51380 - 52000
-    dt_begin = int(51400 / dt)
-    dt_end = int(51700 / dt)
+    # dt_begin = int(63300 / dt)
+    # dt_end = int(63500 / dt)
+    dt_begin = int(0 / dt)
+    dt_end = int(400 / dt)
     _, y_nh, x_nh = values.shape
     x = x_nh // 2
-    center_cut = values[dt_begin:dt_end, :, x:(x + 2)]
+    center_cut = values[dt_begin:dt_end, :, x:(x + 5)]
     center_cut = center_cut.transpose(0, 2, 1).reshape(-1, y_nh)
-    ax = plt.gca()
-    img = ax.imshow(center_cut.T, cmap='jet')
-    plt.colorbar(mappable=img, ax=ax)
+    # center_cut = center_cut.reshape(-1, y_nh)
+    fig = plt.gcf()
+    ax = fig.add_subplot(111, xticks=[], yticks=[])
+    img = ax.imshow(center_cut.T, cmap='jet', aspect=4.0)
+
+    rect = patches.Rectangle((0, 90), 500, 5, clip_on=False, linewidth=.5, edgecolor='000', facecolor='000')
+    ax.add_patch(rect)
+    ax.text(50, 115, '0.1s', color='000', weight='bold', fontsize=14)
+
+    rect = patches.Rectangle((-35, 65), 15, 15, clip_on=False, linewidth=.5, edgecolor='000', facecolor='000')
+    ax.add_patch(rect)
+    ax.text(-80, 65, '0.1mm', color='000', weight='bold', fontsize=8, rotation=90)
+
+    cbar_ax = fig.add_axes([0.5, 0.25, .4, .1], xticks=[], yticks=[], frameon=False)
+    cbar = plt.colorbar(
+        mappable=img, ax=cbar_ax, orientation='horizontal',
+        fraction=1.0, pad=0.01
+    )
+    tick_locator = ticker.MaxNLocator(nbins=3)
+    cbar.locator = tick_locator
+    cbar.update_ticks()
+
 
     title = 'cut_{}_at{}s_{}'.format(state.name, int(dt_begin / 1000), data_filename)
     fig_fname = '{}.{}'.format(title, img_ext)