analysis.py

#%%
import glob
import numpy as np
import matplotlib.pyplot as plt

#%%
fnames = glob.glob('results/*.npz')

# helper functions
def get_key(fname):
    fname = fname.split('/')[-1].split('.npz')[0]
    aux = fname.split('_')
    DA = np.round(float(aux[1]), decimals=2) # convert to percentage
    E_to_I_D2_ratio = np.round(float(aux[2]), decimals=2) # convert to percentage
    E_D1_to_D2_ratio = np.round(float(aux[3]), decimals=2) # convert to percentage
    return DA, E_to_I_D2_ratio, E_D1_to_D2_ratio

def get_unique_keys(fnames):
    DAs, E_to_I_D2_ratios, E_D1_to_D2_ratios = set(), set(), set()
    for fname in fnames:
        DA, E_to_I_D2_ratio, E_D1_to_D2_ratio = get_key(fname)
        DAs.add(DA)
        E_to_I_D2_ratios.add(E_to_I_D2_ratio)
        E_D1_to_D2_ratios.add(E_D1_to_D2_ratio)

    DAs = np.array(sorted(list(DAs)))
    E_to_I_D2_ratios = np.array(sorted(list(E_to_I_D2_ratios)))
    E_D1_to_D2_ratios = np.array(sorted(list(E_D1_to_D2_ratios)))
    return DAs, E_to_I_D2_ratios, E_D1_to_D2_ratios

# load data
DAs, E_to_I_D2_ratios, E_D1_to_D2_ratios = get_unique_keys(fnames)
results = {get_key(fname): np.load(fname) for fname in fnames}
base_case = (0.0, E_to_I_D2_ratios.min(), E_D1_to_D2_ratios.min())

#%%
# shared parameters
r = list(results.values())[0]
NE, NI = r['xi'].shape[1], r['xi'].shape[1] // 4
dt = 20 # in ms
TIME_STEPS = np.arange(0, r['runtime'], step=dt) # in ms

#%% 
# spike count matrices with different dVTs indexed with 'ij' order
var_shapes = (len(E_to_I_D2_ratios), len(E_D1_to_D2_ratios))
E_spike_count_mat_ctrl = np.zeros((len(TIME_STEPS), NE))
I_spike_count_mat_ctrl = np.zeros((len(TIME_STEPS), NI))
E_spike_count_mat_DA = np.zeros((*var_shapes, len(TIME_STEPS), NE))
I_spike_count_mat_DA = np.zeros((*var_shapes, len(TIME_STEPS), NI))

print(E_to_I_D2_ratios, E_D1_to_D2_ratios)

r = results[base_case]
for t, l in zip(r['E_spikes'][0], r['E_spikes'][1]):
    E_spike_count_mat_ctrl[int(t/dt),r['idx_perm'][int(l)]] += 1
for t, l in zip(r['I_spikes'][0], r['I_spikes'][1]):
    I_spike_count_mat_ctrl[int(t/dt),int(l)] += 1

DA = DAs.max() # only analyze the highest DA level
for j, E_to_I_D2_ratio in enumerate(E_to_I_D2_ratios):
    for k, E_D1_to_D2_ratio in enumerate(E_D1_to_D2_ratios):
        r = results[(DA, E_to_I_D2_ratio, E_D1_to_D2_ratio)]

        for t, l in zip(r['E_spikes'][0], r['E_spikes'][1]):
            E_spike_count_mat_DA[j,k,int(t/dt),r['idx_perm'][int(l)]] += 1
            # E_spike_count_mat[j,k,int(t/dt),int(l)] += 1

        for t, l in zip(r['I_spikes'][0], r['I_spikes'][1]):
            I_spike_count_mat_DA[j,k,int(t/dt),int(l)] += 1

# %%
# neuronal rate within the time window of interest
t0, t1 = 500, 1000 # in ms
t0_idx, t1_idx = int(t0/dt), int(t1/dt)

# population rates
E_pop_rate = E_spike_count_mat_DA.sum(axis=-1) / NE / (dt * 1e-3)
I_pop_rate = I_spike_count_mat_DA.sum(axis=-1) / NI / (dt * 1e-3)

# within the time window of interest
E_pop_rate_mean = E_pop_rate[:,:,t0_idx:t1_idx].mean(axis=-1)
I_pop_rate_mean = I_pop_rate[:,:,t0_idx:t1_idx].mean(axis=-1)

# %%
# count bursts
DA = DAs.max() # only analyze the highest DA level
E_burst_count_mat = np.zeros((*var_shapes, NE))
for i, E_to_I_D2_ratio in enumerate(E_to_I_D2_ratios):
    for j, E_D1_to_D2_ratio in enumerate(E_D1_to_D2_ratios):
        r = results[(DA, E_to_I_D2_ratio, E_D1_to_D2_ratio)]
        for l in range(NE):
            spike_times = r['E_spikes'][0][r['E_spikes'][1] == l]
            spike_times = spike_times[(spike_times >= t0) & (spike_times < t1)]
            if len(spike_times) == 0:
                continue
            isi = np.diff(spike_times)
            burst_starts = np.where(isi > 20)[0] + 1
            n_bursts = len(burst_starts) + 1 # +1 for the first burst
            E_burst_count_mat[i,j,l] = n_bursts

# %% 
# variables for plotting
dd2, de = (E_to_I_D2_ratios[1] - E_to_I_D2_ratios[0]) / 2, (E_D1_to_D2_ratios[1] - E_D1_to_D2_ratios[0]) / 2
extent = [E_D1_to_D2_ratios[0]-de, E_D1_to_D2_ratios[-1]+de, 
          E_to_I_D2_ratios[0]-dd2, E_to_I_D2_ratios[-1]+dd2]
ylabel = '${HVC_{RA}}$ / ${HVC_{Int}}$ D2 ratio\n($r_{D2, E/I}$)'
xlabel = '${HVC_{RA}}$ D1 / D2 ratio\n($r_{E, D1/D2}$)'

# %%
# plot some example raster plots

def plot_example(DA, examples, example_colors):
    fig, ax = plt.subplots(3, len(examples), sharex='all', sharey='row', 
                           figsize=(0.3+1*len(examples),2.5), height_ratios=[4,1,1])
    if len(examples) == 1:
        ax = ax[:,None]

    tmin, tmax = 450, 1000

    NI_sel = 100
    i_sel_I = np.random.choice(NI, size=NI_sel, replace=False)

    for i in range(ax.shape[1]):
        ax[0,i].set(yticks=[NI_sel * 3], yticklabels=[''], ylim=[0,NE+NI_sel*3])
        ax[2,i].set(xlabel='Time (ms)', xticks=[500, 900], xticklabels=[0, 400], xlim=[tmin, tmax])

    for i, (E_to_I_D2_ratio, E_D1_to_D2_ratio) in enumerate(examples):
        r = results[(DA, E_to_I_D2_ratio, E_D1_to_D2_ratio)]
        ax[0,i].plot(r['E_spikes'][0], r['idx_perm'][r['E_spikes'][1].astype(int)]+NI_sel*3, 
                     ',', ms=1, c=example_colors[i], rasterized=True)
        
        for y, j in enumerate(i_sel_I):
            ts = r['I_spikes'][0][r['I_spikes'][1] == j]
            ax[0,i].plot(ts, y * 3 + np.zeros_like(ts), 
                         ',', ms=1, c=example_colors[i], alpha=0.4, 
                         rasterized=True)
        
        p = np.where(E_D1_to_D2_ratios == E_D1_to_D2_ratio)[0][0]
        q = np.where(E_to_I_D2_ratios == E_to_I_D2_ratio)[0][0]
        
        ax[1,i].plot(TIME_STEPS, E_pop_rate[p,q], c=example_colors[i])
        ax[2,i].plot(TIME_STEPS, I_pop_rate[p,q], c=example_colors[i], alpha=0.5)
    
    ax[0,0].text(tmin-40, NI_sel*3+NE/2, r'$HVC_{RA}$', rotation=90, 
                 ha='right', va='center', fontsize=10)
    ax[0,0].text(tmin-40, NI_sel, r'$HVC_{Int}$', rotation=90, 
                 ha='right', va='center', fontsize=10)
    return fig, ax

def mark_example_points(ax): # for marking points on heatmaps
    for example, c in zip(examples, example_colors):
        ax.scatter(example[1], example[0], marker='*', color=c)


# %%
# Plot statistics
def plot_stats():
    fig, ax = plt.subplots(1, 3, sharex='all', sharey='all', figsize=(4.5,2.5))

    imshow_args = dict(extent=extent, origin='lower', aspect='auto', cmap='binary')

    im = ax[0].imshow(E_burst_count_mat.mean(axis=2), **imshow_args)
    mark_example_points(ax[0])
    ax[0].set(ylabel=ylabel, xlabel=xlabel)
    fig.colorbar(im, ax=ax[0], orientation='horizontal', location='top',
                label=f'{r"$HVC_{RA}$"} num. bursts', shrink=0.8)

    im = ax[1].imshow(E_pop_rate_mean, **imshow_args)
    mark_example_points(ax[1])
    ax[1].set(ylabel=' ')
    fig.colorbar(im, ax=ax[1], orientation='horizontal', location='top',
                    label='$HVC_{RA}$ rate (Hz)', shrink=0.8)

    im = ax[2].imshow(I_pop_rate_mean, **imshow_args)
    mark_example_points(ax[2])
    ax[2].set(ylabel=' ')
    fig.colorbar(im, ax=ax[2], orientation='horizontal', location='top',
                    label='$HVC_{Int}$ rate (Hz)', shrink=0.8)
    ax[2].invert_xaxis()
    return fig, ax

# %%
# Actual plotting for control condition
fig, ax = plot_example(0.0, [base_case[1:]], ['k'])
ax[1,0].set(yticks=[0, 3], ylim=[0, 3])
ax[2,0].set(ylabel='rate (Hz)', yticks=[0, 50], ylim=[0, 50])
fig.savefig('./raster_control.svg', dpi=100)

# %%
# Actual plotting for DA conditions
DA, i = DAs.max(), 0

examples = [(0.65, 1.45), (0.40, 1.40), (0.30, 1.35)]
example_colors = [f'C{i}' for i in range(len(examples))]
fig, ax = plot_example(DA, examples, example_colors)
ax[0,0].set_title('tutor-naive')
ax[0,1].set_title('intermediate')
ax[0,2].set_title('tutored')
ax[1,0].set(yticks=[0, 15], ylim=[0, 15])
ax[2,0].set(ylabel='rate (Hz)', yticks=[20, 50], ylim=[20, 50])
for i in range(len(examples)):
    ax[2,i].axhline(30, color='k', linestyle=':', alpha=0.5, zorder=-10, lw=1)
fig.savefig('./rasters.svg', dpi=100)

fig, ax = plot_stats()
fig.savefig('./pop_stats_heatmaps.svg')

# %%