spiketrain chapter work

spike_trains/code/firing_rates.py

@@ -228,7 +228,7 @@ def plot_comparison(spike_times, bin_width, sigma, max_t=30., dt=1e-4):
     ax3.set_xticklabels([])
 
     ax4.plot(time, conv_rate, label="convolved rate")
-    ax4.set_xlabel("times [s]", fontsize=10)
+    ax4.set_xlabel("time [s]", fontsize=10)
     ax4.set_ylabel("firing rate [Hz]", fontsize=10)
     ax4.legend(fontsize=10)
     ax4.set_xlim([1.5, 3.5])

spike_trains/code/sta.py

@@ -0,0 +1,81 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.io as scio
+import seaborn as sb
+from IPython import embed
+sb.set_context("paper")
+
+
+def plot_sta(times, stim, dt, t_min=-0.1, t_max=.1):
+    count = 0
+    sta = np.zeros((abs(t_min) + abs(t_max))/dt)
+    time = np.arange(t_min, t_max, dt)
+    if len(stim.shape) > 1 and stim.shape[1] > 1:
+        stim = stim[:,1]
+    for i in range(len(times[0])):
+        times = np.squeeze(spike_times[0][i])
+        for t in times:
+            if (int((t + t_min)/dt) < 0) or ((t + t_max)/dt > len(stim)):
+                continue;
+            
+            min_index = int(np.round((t+t_min)/dt))
+            max_index = int(np.round((t+t_max)/dt))
+            snippet = np.squeeze(stim[ min_index : max_index])            
+            sta += snippet
+            count += 1
+    sta /= count
+    
+    fig = plt.figure()
+    fig.set_size_inches(5, 5)
+    fig.subplots_adjust(left=0.15, bottom=0.125, top=0.95, right=0.95, )
+    fig.set_facecolor("white")
+    
+    ax = fig.add_subplot(111)
+    ax.plot(time, sta, color="darkblue")
+    ax.set_xlabel("time [s]")
+    ax.set_ylabel("stimulus")
+    ax.xaxis.grid('off')
+    ylim = ax.get_ylim()
+    xlim = ax.get_xlim()
+    ax.plot(list(xlim), [0., 0.], zorder=1, color='darkgray', ls='--')
+    ax.plot([0., 0.], list(ylim), zorder=1, color='darkgray', ls='--')
+    ax.set_xlim(list(xlim))
+    ax.set_ylim(list(ylim))
+    fig.savefig("../lecture/images/sta.pdf")
+    plt.close()
+    return sta
+
+
+def reconstruct_stimulus(spike_times, sta, stimulus, t_max=30., dt=1e-4):
+    s_est = np.zeros((spike_times.shape[1], len(stimulus)))
+    for i in range(10):
+        times = np.squeeze(spike_times[0][i])
+        indices = np.asarray((np.round(times/dt)), dtype=int)
+        y = np.zeros(len(stimulus))
+        y[indices] = 1
+        s_est[i, :] = np.convolve(y, sta, mode='same')
+
+    plt.plot(np.arange(0, t_max, dt), stimulus[:,1], label='stimulus', color='darkblue', lw=2.)
+    plt.plot(np.arange(0, t_max, dt), np.mean(s_est, axis=0), label='reconstruction', color='silver', lw=1.5)
+    plt.xlabel('time[s]')
+    plt.ylabel('stimulus')
+    plt.xlim([0.0, 0.25])
+    plt.ylim([-1., 1.])
+    plt.legend()
+    plt.plot([0.0, 0.25], [0., 0.], color="darkgray", lw=1, ls='--', zorder=1)
+
+    fig = plt.gcf()
+    fig.set_size_inches(7.5, 5)
+    fig.subplots_adjust(left=0.15, bottom=0.125, top=0.95, right=0.95, )
+    fig.set_facecolor("white")
+    fig.savefig('../lecture/images/reconstruction.pdf')
+    plt.close()
+
+
+if __name__ == "__main__":
+    punit_data =  scio.loadmat('../../programming/exercises/p-unit_spike_times.mat')
+    punit_stim = scio.loadmat('../../programming/exercises/p-unit_stimulus.mat')
+    spike_times = punit_data["spike_times"]
+    stimulus = punit_stim["stimulus"]
+    sta = plot_sta(spike_times, stimulus, 5e-5, -0.05, 0.05)
+    reconstruct_stimulus(spike_times, sta, stimulus, 10, 5e-5)