fig3.py

import sys
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import os
import utils

def load_fig3_dat(root):
    dat = {}

    beh_path = os.path.join(root, 'beh')
    # load before learning performance
    beh0 = np.load(os.path.join(beh_path, 'Beh_sup_train2_before_learning.npy'), allow_pickle=1).item()
    dat['mean_beh_bef'] = utils.get_mean_lick_response(beh0, lick_typ='befRew')
    # load after learning performance
    beh1 = np.load(os.path.join(beh_path, 'Beh_sup_train2_after_learning.npy'), allow_pickle=1).item()
    dat['mean_beh_aft'] = utils.get_mean_lick_response(beh1, lick_typ='befRew') 

    fns = ['naive_test1_leaf2_circle1_dprime_leaf2_distribution.npy',
           'sup_train2_before_learning_leaf2_circle1_dprime_leaf2_distribution.npy',
            'unsup_train2_before_learning_leaf2_circle1_dprime_leaf2_distribution.npy',

            'sup_train2_after_learning_leaf2_circle1_dprime_leaf2_distribution.npy',
            'unsup_train2_after_learning_leaf2_circle1_dprime_leaf2_distribution.npy',]
    dat['img'] = [np.load(os.path.join(root, 'process_data', fn), allow_pickle=1).item() for fn in fns]
    dat['outlines'] = np.load(os.path.join(root, 'retinotopy/areas.npz'), allow_pickle = True)['out']
    dat['hotcmp'] = make_hot_cmap() 

    fns = ['sup_train2_before_learning_leaf2_circle1_dprime_frac.npy',
            'sup_train2_after_learning_leaf2_circle1_dprime_frac.npy',
            'unsup_train2_before_learning_leaf2_circle1_dprime_frac.npy',
            'unsup_train2_after_learning_leaf2_circle1_dprime_frac.npy',]
    dat['frac1'] = [np.load(os.path.join(root, 'process_data', fn), allow_pickle=1).item() for fn in fns]

    fns = ['naive_test1_leaf1_leaf2_dprime_distribution.npy',
           'unsup_train2_after_learning_leaf1_leaf2_dprime_distribution.npy',
            'sup_train2_after_learning_leaf1_leaf2_dprime_distribution.npy']
    dat['img1'] = [np.load(os.path.join(root, 'process_data', fn), allow_pickle=1).item() for fn in fns]

    fns = ['naive_test1_leaf1_leaf2_dprime_frac.npy',
            'test1_after_grating_leaf1_leaf2_dprime_frac.npy',
            'unsup_train2_after_learning_leaf1_leaf2_dprime_frac.npy',
            'sup_train2_after_learning_leaf1_leaf2_dprime_frac.npy',]
    dat['frac2'] = [np.load(os.path.join(root, 'process_data', fn), allow_pickle=1).item() for fn in fns]
    return dat

def plot_fig3(dat, root):
    fig = plt.figure(figsize=(7, 7), dpi=500)
    plt.rcParams["font.family"] = "arial"
    plt.rcParams["font.size"] = 5
    
    ax_text = fig.add_axes([0,0.33,1,0.60])
    ax_text.set_facecolor('None')
    ax_text.axis('off')    

    ################## distribution for leaf2  ######################
    x,y, dx,dy, w,h =0.0, 0.67, 0.135, 0.12, 0.12, 0.12

    img_ax1 = [fig.add_axes([x, y+dy*0.6, w, h]), fig.add_axes([x+dx, y+dy, w, h]), 
               fig.add_axes([x+dx, y, w, h]), fig.add_axes([x+2.8*dx, y+dy, w, h]),  
               fig.add_axes([x+2.8*dx, y, w, h])]

    a, b, n =5, 10, 8 # for vmax
    vmax = a/(b**n) # i.e. 5x10e-8
    for i, itn in enumerate(['naive mice', 'task mice\nwhen new', 'unsup. mice\nwhen new', 'task mice\nafter learning', 'unsup. mice\nafter learning']):
        distribution_map(img_ax1[i], dat['img'][i]['img'], dat['outlines'], cmp=dat['hotcmp'], vmax=vmax, scalbar=0)
        img_ax1[i].text(0.35, 0.85, itn, transform=img_ax1[i].transAxes)    
    cbar1 = fig.add_axes([x+0.02,y+0.06,0.05,0.005])
    cbar(cbar1, cmap=dat['hotcmp'], tickLabel=[0, r'$5\times10^{-8}$'], 
              orientation='horizontal', cbarLabelrotation=0,
              cbarLabel='density', ticks=[0, 1], labelpad=-15) 

    ################## change of leaf2 fraction  ######################
    x,y, dx,dy, w,h =0.54,0.69, 0.14,0.14, 0.3,0.215
    ax_frac = fig.add_axes([x,y,w,h])
    plot_frac(ax_frac, dat['frac1'][2]['value'][:, :, 2, 1], dat['frac1'][3]['value'][:, :, 2, 1], col=[0.46,0,0.23])
    plot_frac(ax_frac, dat['frac1'][0]['value'][:, :, 2, 1], dat['frac1'][1]['value'][:, :, 2, 1], col='g')
    for t,txt in enumerate(['V1', 'mHV', 'lHV', 'aHV']):
        ax_frac.text(0.1 + 0.25*t, 0.98, txt, transform=ax_frac.transAxes)
    ax_frac.text(0.75, 0.7, 'task mice', color='g', transform=ax_frac.transAxes)
    ax_frac.text(0.75, 0.65, 'unsupervised', color=[0.46,0,0.23], transform=ax_frac.transAxes)         
    # ################## performance  ######################
    x,y, dx,dy, w,h =0.88, 0.69, 0, 0.08, 0.1, 0.215
    ax_beh = fig.add_axes([x,y,w,h])

    train2_perf_plot(ax_beh, dat['mean_beh_bef'], dat['mean_beh_aft'], title='', yn=1, xlm=[-0.3, 1.3])

    ################## distribution for leaf1 vs leaf2s  ######################
    x,y, dx,dy, w,h =-0.01,0.53, 0.108,0.14, 0.12,0.12
    img_ax2 = [fig.add_axes([x,y,w,h]), fig.add_axes([x+dx,y,w,h]), 
               fig.add_axes([x+2*dx,y,w,h])]
    for i, itn in enumerate(['naive mice', 'unsup. mice\nafter learning', 'task mice\nafter learning']):
        distribution_map(img_ax2[i], dat['img1'][i]['img'], dat['outlines'], cmp=dat['hotcmp'], vmax=vmax, scalbar=0)
        img_ax2[i].text(0.35, 0.85, itn, transform=img_ax2[i].transAxes)     
    cbar2 = fig.add_axes([x+3*dx+0.005, y+0.05, 0.005, 0.06])
    cbar(cbar2, cmap=dat['hotcmp'], tickLabel=[0, r'$5\times10^{-8}$'], cbarLabel='density', ticks=[0, 1], labelpad=-15) 

    ################## leaf1-leaf2 fraction  ######################
    x,y, dx,dy, w,h =0.03,0.36, 0.15,0.13, 0.28,0.14
    ax_frac1 = fig.add_axes([x,y,w,h])

    plot_leaf1_leaf2_frac(ax_frac1, dat['frac2'])

    ################## examples of coding direction  ######################
    x,y, dx,dy, w,h =0.375,0.37, 0.1,0.165, 0.08,0.095

    V1_ax = [fig.add_axes([x,y+dy,w,h]), fig.add_axes([x+dx,y+dy,w,h]), fig.add_axes([x+2*dx,y+dy,w,h])]
    mHV_ax = [fig.add_axes([x,y,w,h]), fig.add_axes([x+dx,y,w,h]), fig.add_axes([x+2*dx,y,w,h])]

    fns = [r'process_data\naive_test1_coding_direction.npy',
          r'process_data\unsup_train2_after_learning_coding_direction.npy',
          r'process_data\sup_train2_after_learning_coding_direction.npy']
    mnames = ['TX109', 'TX119', 'TX108']
    for i in range(3):
        leaf2_coding_direction(V1_ax[i], root, fns[i], mnames[i], 'V1')
        leaf2_coding_direction(mHV_ax[i], root, fns[i], mnames[i], 'mHV')

    ################## Similarity Index  ######################
    x,y, dx,dy, w,h =0.69,0.37, 0.11,0.11, 0.29,0.26
    ax_SI = fig.add_axes([x,y,w,h])
    SI_train2_after_learning(ax_SI, root) 
   
    
    ax_text.text(0.01, .98, r"$\bf{a}$ Distribution of leaf2-selective neurons (leaf2 $vs$ circle1, $d'\geq$0.3)", fontsize=5.5)
    ax_text.text(0.52, 0.98, r"$\bf{b}$ Summary of changes", fontsize=5.5)
    ax_text.text(0.865, 0.98, r"$\bf{c}$ Licking behavior on leaf2", fontsize=5.5)

    ax_text.text(.01, .54, r"$\bf{d}$ Distribution of selective neurons (leaf1 $vs$ leaf2, $|d'|\geq$0.3)", fontsize=5.5)
    ax_text.text(0.01, 0.31, r"$\bf{e}$ Summary of selective neurons (leaf1 $vs$ leaf2, $|d'|\geq$0.3)", fontsize=5.5) 
    

    ax_text.text(.355, .54, r"$\bf{f}$ Coding direction of leaf1-circle1, example mouse (V1)", fontsize=5.5)
    ax_text.text(0.355, 0.26, r"$\bf{g}$ Coding direction of leaf1-circle1, example mouse (medial)", fontsize=5.5)  
    ax_text.text(0.68, 0.54, r"$\bf{h}$ Changes of leaf2-$SI$ (similarity index)", fontsize=5.5)     

def make_hot_cmap():
    new_hot = cm.get_cmap('magma_r', 256)
    newcolors = new_hot(np.linspace(0, 1, 256))
    noCol = np.array([0, 0, 0, 1])
    return ListedColormap(newcolors[:,:])

def distribution_map(ax, img, outlines, scal=10, cmp='', vmax=0.6, hlw = 2, alpha=0.4, scalbar=0):
    sz = img.shape[0]
    ax.imshow(np.flipud(img),cmap=cmp,vmax=vmax,extent=[0, sz*scal, 0, sz*scal], rasterized=True)
    temp_outline=[]
    for j in range(10):
        if j!=7:
            temp = outlines[j].copy()
            temp[:,1] = -(-temp[:,1]+800-2500)+2500
            temp[:,0] = temp[:,0]+800
            ax.plot(temp[:,1],temp[:,0],linewidth=0.5,color='k',alpha=alpha)  
            temp_outline.append(temp)
        else:
            temp_outline.append([])
    if scalbar:
        ax.plot([450,1450],[880,880],'k-',lw=1)      
        
    ax.axis('off')
    utils.fmt(ax, y_invert=0, xlm=[200,4500], ylm=[500,4800],axis_off='off', aspect='equal')
    
def cbar(ax, cmap='gray_r', ticks=[0,1], tickLabel=[], cbarLabel=[], cbarLabelrotation=270, fs_tick=None, fs_label=None, tick_len=1, tick_wid=1, tpad=1, shrink=0.1, orientation='vertical', labelpad=0, outline_color='None'):
    """pos: position of colorbar [x,y,h,w]"""
    cbar = plt.colorbar(cm.ScalarMappable(norm=None, cmap=cmap), cax=ax, ticks=ticks,orientation=orientation,drawedges=False )   
    cbar.ax.tick_params(length=tick_len,width=tick_wid,pad=tpad)
    if any(tickLabel):
        if orientation=='vertical':
            cbar.ax.set_yticklabels(tickLabel)
            for t in cbar.ax.get_yticklabels():
                 t.set_fontsize(fs_tick)  
        elif orientation=='horizontal':
            cbar.ax.set_xticklabels(tickLabel)
            for t in cbar.ax.get_xticklabels():
                 t.set_fontsize(fs_tick)             
    if any(cbarLabel):
        cbar.set_label(cbarLabel,rotation=cbarLabelrotation,fontsize=fs_label, position='bottom', labelpad=labelpad)
    cbar.outline.set_color(outline_color)
    cbar.outline.set_linewidth(0.3)    
    
def plot_frac(ax, frac1, frac2, col='k', alpha=0.3, mk='s',lw0=1, lw1=2.5, elw=2, fs=None, mks=5,ylm=[-0.001,0.46]):
    frac = np.array([frac1, frac2])
    for i in range(4):
        x = np.array([0, 0.5]) + i
        ax.plot(x, frac[:, i, :], color=col, alpha=alpha, lw=0.7)
    u, sem = frac.mean(2), frac.std(2, ddof=1)/np.sqrt(frac.shape[2])
    ax.plot([np.arange(4), np.arange(4)+0.5], u, color=col, lw=1.5)
    ax.errorbar(np.arange(4), u[0, :], yerr=sem[0, :], marker='s', markersize=3, color=col, ls='None', markeredgecolor='k', markeredgewidth=0.5)
    ax.errorbar(np.arange(4)+0.5, u[1, :], yerr=sem[1, :], marker='s', markersize=3, color=col, ls='None', markeredgecolor='k', markeredgewidth=0.5)  
    utils.fmt(ax, ylm=[0, 0.45], xtick=[np.arange(8)/2, ['when\nnew', 'after\nlearning', None, None, None, None, None, None]],
             ylabel="% leaf2-selective neurons ($d'\geq0.3$)", ytick=[[0, 0.1, 0.2, 0.3, 0.4], [0, 10, 20, 30, 40]])     
    
def train2_perf_plot(ax, bef, aft, title='', yn=1, xlm=[-0.3, 1.3]):
    r = np.array([bef['u_sem'][:, 3, 0], aft['u_sem'][:, 2, 0]])
    u, sem = r.mean(1), r.std(1, ddof=1)/np.sqrt(r.shape[1])
    ax.plot([0, 1], r, 'c-', lw=0.5, alpha=0.5)
    ax.plot([0, 1], u, 'c-', lw=2)
    ax.errorbar(0, u[0], yerr=sem[0], marker='s', markersize=3, color='c')
    ax.errorbar(1, u[1], yerr=sem[1], marker='s', markersize=3, color='c')
    yln = 'anticipatory licking (%trials)' if yn else ''
    utils.fmt(ax, xtick=[[0, 1], ['when\nnew', 'after\nlearning']], ytick=[[0, 0.5, 1], [0, 50, 100]],
          ylabel=yln, title=title, xlm=xlm, ylm=[0, 1])    
    
def plot_leaf1_leaf2_frac(ax, frac, col='k', alpha=0.3, mk='s',lw0=1, lw1=2.5, elw=2, fs=None, mks=5,ylm=[-0.001,0.46]):
    xt = np.array([-0.2, 0.04, 0.17, 0.3])
    cols = ['k', '0.5', [0.46,0,0.23], 'g']
    for i in range(len(frac)): # looping across experiments    
        val = frac[i]['value'][:, :, 2, 0] # take fraction at dprime=0.3
        u, sem = val.mean(1), val.std(1, ddof=1)/np.sqrt(val.shape[1])
        ax.scatter(np.repeat(np.arange(4)+xt[i], val.shape[1]), val, marker='o', s=13, color=cols[i], alpha=0.3, edgecolor='None')
        ax.errorbar(np.arange(4)+xt[i], u, yerr=sem, marker='_', markersize=4, color=cols[i], ls='None')  
    utils.fmt(ax, ylm=[-0.01, 0.28], xtick=[xt, ['navie', 'grat.', 'unsup.', 'sup.']], xrot=90,
             ylabel="% selective neurons", ytick=[[0, 0.1, 0.2], [0, 10, 20]])
    xticklabels = ax.get_xticklabels()    
    for label, color in zip(xticklabels, cols):
        label.set_color(color)        
    for t,txt in enumerate(['V1', 'mHV', 'lHV', 'aHV']):
        ax.text(0.1 + 0.25*t, 0.95, txt, transform=ax.transAxes)        
        
def leaf2_coding_direction(ax, root, fn, mname, arn, isyn=1):
    dat = np.load(os.path.join(root, fn), allow_pickle=1).item()
    resp1 = dat['proj_2_stim1'][mname][arn]
    resp2 = dat['proj_2_stim2'][mname][arn]   
    cols = ['r', 'b', 'c']
    ax.axvline(50, linestyle='--', color='k', linewidth=0.3)
    ax.axhline(0, linestyle=':', color='k', linewidth=0.3)    
    for s,sid in enumerate([0, 2, 3]):
        diff = resp1[sid]-resp2[sid]
        ax.plot(diff.T, color=cols[s], lw=0.5, alpha=0.05)
        u,sem = diff.mean(0), diff.std(0, ddof=1)/np.sqrt(diff.shape[0])
        ax.plot(u, lw=1, color=cols[s])
        ax.fill_between(np.arange(len(u)), u-sem, u+sem, color=cols[s], alpha=0.3)  # Shaded STD area 
        yn = 'projection (a.u.)' if isyn else ''
        utils.fmt(ax, xlm=[dat['pos_from_prev'], dat['pos_from_prev']+60], 
                  xtick=[[10, 30, 50, 70], [0, 2, 4, 6]], ytick=[[-1, 0, 1]], ylm=[-1.5, 1.5],
                 xlabel='position (m)', ylabel=yn, ypad=-2)        
        
def SI_train2_after_learning(ax, root):
    fns = ['naive_test1_coding_direction.npy', 
           'test1_after_grating_coding_direction.npy', 
           'unsup_train2_after_learning_coding_direction.npy', 
           'sup_train2_after_learning_coding_direction.npy']
    cols = ['k','0.5', [0.46,0, 0.23], 'g']
    xt = np.array([-0.2, 0.04, 0.17, 0.3])
    for f,fn in enumerate(fns):
        out = utils.load_coding_direction(os.path.join(root, 'process_data'), fn)
        cd_proj_u = out['proj_tr_mean']       
        dx = abs(cd_proj_u[:, :, 3] - cd_proj_u[:, :, 0])
        dy = abs(cd_proj_u[:, :, 3] - cd_proj_u[:, :, 2])
        dxy = abs(cd_proj_u[:, :, 2] - cd_proj_u[:, :, 0])
        SI = (dx-dy) / dxy
        SI = SI.astype(float)
        ax.scatter(np.repeat(np.arange(4)+xt[f], SI.shape[0]), SI.T, marker='o', s=13, color=cols[f], alpha=0.3, edgecolor='None')
        u, sem = np.mean(SI, 0), SI.std(0, ddof=1)/np.sqrt(SI.shape[0])
        ax.plot(np.arange(4)+xt[f], u, marker='_', markersize=4, lw=1.2, color=cols[f], ls='None')
        ax.errorbar(np.arange(4)+xt[f], u, yerr=sem, lw=1, color=cols[f], ls='None')
    yn = 'similarity index ($SI$)'
    utils.fmt(ax, xtick=[xt, ['navie', 'grat.', 'unsup.', 'sup.']], xrot=90, ylabel=yn, ytick=[[-1, 0, 1]], ypad=0)   
    ax.axhline(0, linewidth=0.5, linestyle='--', color='k')
    xticklabels = ax.get_xticklabels()    
    for label, color in zip(xticklabels, cols):
        label.set_color(color)    
    for t,txt in enumerate(['V1', 'mHV', 'lHV', 'aHV']):
        ax.text(0.1 + 0.25*t, 0.99, txt, transform=ax.transAxes)