fig4.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
from scipy import stats

def load_fig4_dat(root = r'D:\Results\Zhong-et-al-2025'):
    dat = {}
    fns = ['sup_train1_before_learning_rew_distribution.npy',
            'sup_train1_after_learning_rew_distribution.npy',
            'unsup_train1_before_learning_rew_distribution.npy',
            'unsup_train1_after_learning_rew_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() 

    dat['RewResp_test1'] = np.load(os.path.join(root, 'process_data', 'sup_test1_reward_response.npy'), allow_pickle=True).item()
    dat['RewResp_test2'] = np.load(os.path.join(root, 'process_data', 'sup_test2_reward_response.npy'), allow_pickle=True).item()
    dat['RewResp_test3'] = np.load(os.path.join(root, 'process_data', 'sup_test3_reward_response.npy'), allow_pickle=True).item()  
    return dat

def plot_fig4(dat, root):
    fig = plt.figure(figsize=(7, 7*7/10.5), dpi=500)
    ax_text = fig.add_axes([0,0.05,1,0.94])
    ax_text.set_facecolor('None')
    ax_text.axis('off')
    plt.rcParams["font.family"] = "arial"
    plt.rcParams["font.size"] = 5

    ################## distribution of reward prediction neurons  ######################
    x,y, dx,dy, w,h =-0.00,0.3, 0.17,0.17, 0.18,0.18
    ax_rew_dist = [fig.add_axes([x,y+dy,w,h],rasterized=True), fig.add_axes([x+dx,y+dy,w,h],rasterized=True),
                  fig.add_axes([x,y,w,h],rasterized=True), fig.add_axes([x+dx,y,w,h],rasterized=True)]

    a, b, n =5, 10, 8
    vmax = a/(b**n) # i.e. 5x10e-8
    for i, itn in enumerate(['task mice\nbefore learning', 'task mice\nafter learning', 'unsupervised\nbefore learning', 'unsupervised\nafter learning']):
        distribution_map(ax_rew_dist[i], dat['img'][i]['img'], dat['outlines'], cmp=dat['hotcmp'], vmax=vmax, scalbar=0)
        ax_rew_dist[i].text(0.35, 0.85, itn, transform=ax_rew_dist[i].transAxes)

    ################# distribution summary  ######################
    x,y, dx,dy, w,h =0.42,0.345, 0.1,0.1, 0.14,0.30
    ax_frac=fig.add_axes([x,y,w,h])
    plot_rewPred_neu_frac(ax_frac, root, xlm=[2.9, 3.6])
#     ax_frac.text(0.4, 0.98, 'aHV', transform=ax_frac.transAxes)
    utils.fmt(ax_frac, xtick=[[3, 3.5], ['before\nlearning', 'after\nlearning']])

    x,y, dx,dy, w,h =0.62,0.365, 0.095,0.16, 0.075,0.1
    # example reward prediciton neuron, test1
    axs = [fig.add_axes([x, y+dy, w, h]), fig.add_axes([x+dx, y+dy, w, h]), fig.add_axes([x+dx*2, y+dy, w, h]), fig.add_axes([x+dx*3 ,y+dy, w, h])]
    example_rewPred_resp_test1(axs, root, nneu=9)

    # example mouse, test1
    axs = [fig.add_axes([x,y,w,h]), fig.add_axes([x+dx,y,w,h]), fig.add_axes([x+dx*2,y,w,h]), fig.add_axes([x+dx*3,y,w,h])]
    test1_rewPred_resp(axs, dat['RewResp_test1']['VR2'])

    # example mouse, test2  
    x,y, dx,dy, w,h =0.62,0.03, 0.095,0.16, 0.075,0.1
    axs = [fig.add_axes([x, y+dy, w, h]), fig.add_axes([x+dx, y+dy, w, h]), fig.add_axes([x+dx*2, y+dy, w, h]), fig.add_axes([x+dx*3 ,y+dy, w, h])]
    test2_rewPred_resp(axs, dat['RewResp_test2']['VR2'])

    # example mouse, test3  
    axs = [fig.add_axes([x, y, w, h]), fig.add_axes([x+dx, y, w, h]), fig.add_axes([x+dx*2, y, w, h]), fig.add_axes([x+dx*3 ,y, w, h])]
    test3_rewPred_resp(axs, dat['RewResp_test3']['VR2_swap1'])

    # value aligned to cue, test1 
    x,y, dx,dy, w,h =0.02,0.03, 0.12,0.12, 0.105,0.23
    ax_val2cue=fig.add_axes([x,y,w,h])
    plot_rewResp_2_Cue(ax_val2cue, dat['RewResp_test1'])

    # value aligned to first lick, test1  
    x,y, dx,dy, w,h =0.18,0.03, 0.12,0.12, 0.105,0.23
    ax_spk2FL=fig.add_axes([x, y, w, h])
    plot_rewResp_2_firstLick(ax_spk2FL, dat['RewResp_test1'])

    # value vs beh, test1
    x,y, dx,dy, w,h =0.35,0.03, 0.14,0.12, 0.09,0.23
    ax_val2beh=fig.add_axes([x,y,w,h])
    plot_rewResp_in_leaf2(ax_val2beh, dat['RewResp_test1'])

    # stim vs beh, test1
    ax_stim2beh=fig.add_axes([x+dx,y,w,h],rasterized=True)
    plot_stimResp_in_leaf2(ax_stim2beh, dat['RewResp_test1'])  
    
    ax_text.text(0.01, 0.65, r"$\bf{a}$ Distribution of reward-prediction neurons ($d'_{late\ vs.\ early}$ $\geq$ 0.3)", fontsize=5.5)
    ax_text.text(0.385, 0.65, r"$\bf{b}$ Summary of changes in anterior areas", fontsize=5.5)
    ax_text.text(0.605, 0.65, r"$\bf{c}$ Example reward-prediction activity (anterior, $test1$)", fontsize=5.5)

    ax_text.text(.74, .622, r"activity of example neuron", fontsize=5.5)
    ax_text.text(0.72, 0.452, r"average activity across neurons", fontsize=5.5) 
    

    ax_text.text(0, .25, r"$\bf{d}$ Reward-prediction neurons,", fontsize=5.5)
    ax_text.text(0, .23, r"aligned to sound cue (in leaf1)", fontsize=5.5)
    ax_text.text(.16, .25, r"$\bf{e}$ Reward-prediction neurons,", fontsize=5.5)
    ax_text.text(.16, .23, r"aligned to first lick (in leaf1)", fontsize=5.5)
    ax_text.text(.315, .25, r"$\bf{f}$ Reward-prediction neurons,", fontsize=5.5)
    ax_text.text(.315, .23, r"(anterior in leaf2)", fontsize=5.5)
    ax_text.text(.465, .25, r"$\bf{g}$ leaf1-selective neurons,", fontsize=5.5)
    ax_text.text(.465, .23, r"(medial in leaf2)", fontsize=5.5)
    
    ax_text.text(.605, .265, r"$\bf{g}$ Example average reward-prediction activity (anterior, $test2$)", fontsize=5.5)
    ax_text.text(.605, .095, r"$\bf{g}$ Example average reward-prediction activity (anterior, $test3$)", 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 plot_frac(ax, frac1, frac2, col='k', alpha=0.3, mk='s',lw0=0.7, lw1=2, 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=lw0)
    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=lw1)
    ax.errorbar(np.arange(4), u[0, :], yerr=sem[0, :], marker='s', markersize=3, color=col, ls='None')
    ax.errorbar(np.arange(4)+0.5, u[1, :], yerr=sem[1, :], marker='s', markersize=3, color=col, ls='None')  
    utils.fmt(ax, ylm=[0, 0.25], xtick=[np.arange(8)/2, ['before\nlearning', 'after\nlearning', None, None, None, None, None, None]],
             ylabel=r"% neurons with $d'_{late vs. early} \geq 0.3$", ytick=[[0, 0.1, 0.2], [0, 10, 20]])
    
def plot_rewPred_neu_frac(ax, root, xlm=[]):
    # load reward neurons fractions
    fn0='sup_train1_before_learning_rew_frac.npy'
    fn1='sup_train1_after_learning_rew_frac.npy'
    fn2='unsup_train1_before_learning_rew_frac.npy'
    fn3='unsup_train1_after_learning_rew_frac.npy'    
    sup_rew_bef = np.load(os.path.join(root, 'process_data', fn0), allow_pickle=1).item()
    sup_rew_aft = np.load(os.path.join(root, 'process_data', fn1), allow_pickle=1).item()
    unsup_rew_bef = np.load(os.path.join(root, 'process_data', fn2), allow_pickle=1).item()
    unsup_rew_aft = np.load(os.path.join(root, 'process_data', fn3), allow_pickle=1).item()  
    plot_frac(ax, sup_rew_bef['value'], sup_rew_aft['value'], col='g')
    plot_frac(ax, unsup_rew_bef['value'], unsup_rew_aft['value'], col=[0.46, 0, 0.23])
    utils.fmt(ax, ylm=[0, 0.12], xlm=xlm)   
    
def example_rewPred_resp_test1(ax, root, vmin=0, vmax = 0.5, ms=1.5, nneu=9):
    dat = np.load(os.path.join(root, 'process_data', 'Example_reward_neurons_in_sup_test1.npy'), allow_pickle=1).item()
    resp = dat['Example_reward_neurons_VR2_2021_04_11_1']['resp'][nneu]
    beh = dat['Example_reward_neurons_VR2_2021_04_11_1']['beh']
    CuePos = np.mod(beh['SoundDelPos'], 60)
    RewPos = beh['RewPos']
    uniqW, WallN, stim_id = beh['UniqWalls'], beh['WallName'], beh['stim_id']
    stimN = ['circle1', 'circle2', 'leaf2', 'leaf1']
    cols = ['r', 'm', 'c', 'b']
    for i, sid in enumerate([0, 1, 3, 2]):
        stim  = WallN==uniqW[stim_id==sid]
        s_cue = CuePos[stim]
        sort = np.argsort(s_cue)
        ax[i].imshow(resp[stim][sort], cmap='gray_r', vmin=vmin, vmax=vmax)
        ax[i].plot(s_cue[sort], np.arange(stim.sum()),  marker='.', color='purple', ms=ms, linestyle='None', markeredgewidth=0)
        ax[i].axvline(40, linewidth=0.5, linestyle='--',color='k')
        utils.fmt(ax[i], xtick=[[0,20,40,60],[0,2,4,6]], ytick=[[0, stim.sum()]], xlabel='pos. in %s'%(stimN[i])) 
        
def test1_rewPred_resp(ax, dat, vmin=0, vmax = 0.5, ms=1.5):
    CuePos = np.mod(dat['beh']['SoundDelPos'], 60)
    RewPos = dat['beh']['RewPos']
    uniqW, WallN, stim_id = dat['beh']['UniqWalls'], dat['beh']['WallName'], dat['beh']['stim_id']
    stimN = ['circle1', 'circle2', 'leaf2', 'leaf1']
    cols = ['r', 'm', 'c', 'b']
    for i, sid in enumerate([0, 1, 3, 2]):
        stim  = WallN==uniqW[stim_id==sid]
        s_cue = CuePos[stim]
        sort = np.argsort(s_cue)
        ax[i].imshow(dat['resp'][stim][sort], cmap='gray_r', vmin=vmin, vmax=vmax)
        ax[i].plot(s_cue[sort], np.arange(stim.sum()),  marker='.', color='purple', ms=ms, linestyle='None', markeredgewidth=0)
        ax[i].axvline(40, linewidth=0.5, linestyle='--',color='k')
        utils.fmt(ax[i], xtick=[[0,20,40,60],[0,2,4,6]], ytick=[[0, stim.sum()]], xlabel='pos. in %s'%(stimN[i])) 
        
def test2_rewPred_resp(ax, dat, vmin=0, vmax = 0.5, ms=1.5):
    CuePos = np.mod(dat['beh']['SoundDelPos'], 60)
    RewPos = dat['beh']['RewPos']
    uniqW, WallN, stim_id = dat['beh']['UniqWalls'], dat['beh']['WallName'], dat['beh']['stim_id']
    stimN = ['circle1', 'leaf3', 'leaf2', 'leaf1']
    cols = ['r', [0.27,0.51,0.71], 'c', 'b']
    for i, sid in enumerate([0, 4, 3, 2]):
        stim  = WallN==uniqW[stim_id==sid]
        s_cue = CuePos[stim]
        sort = np.argsort(s_cue)
        ax[i].imshow(dat['resp'][stim][sort], cmap='gray_r', vmin=vmin, vmax=vmax)
        ax[i].plot(s_cue[sort], np.arange(stim.sum()),  marker='.', color='purple', ms=ms, linestyle='None', markeredgewidth=0)
        ax[i].axvline(40, linewidth=0.5, linestyle='--',color='k')
        utils.fmt(ax[i], xtick=[[0,20,40,60],[0,2,4,6]], ytick=[[0, stim.sum()]], xlabel='pos. in %s'%(stimN[i]))  
        
def test3_rewPred_resp(ax, dat, vmin=0, vmax = 0.5, ms=1.5):
    CuePos = np.mod(dat['beh']['SoundDelPos'], 60)
    RewPos = stats.zscore(dat['beh']['RewPos'])
    uniqW, WallN, stim_id = dat['beh']['UniqWalls'], dat['beh']['WallName'], dat['beh']['stim_id']
    stimN = ['circle1', 'leaf1_swap', 'leaf2', 'leaf1']
    cols = ['r', [0,0.47,0.47], 'c', 'b']
    for i, sid in enumerate([0, 5, 3, 2]): # leaf1_swap index: 5 or 6 
        stim  = WallN==uniqW[stim_id==sid]
        s_cue = CuePos[stim]
        sort = np.argsort(s_cue)
        ax[i].imshow(dat['resp'][stim][sort], cmap='gray_r', vmin=vmin, vmax=vmax)
        ax[i].plot(s_cue[sort], np.arange(stim.sum()),  marker='.', color='purple', ms=ms, linestyle='None', markeredgewidth=0)
        ax[i].axvline(40, linewidth=0.5, linestyle='--',color='k')
        utils.fmt(ax[i], xtick=[[0,20,40,60],[0,2,4,6]], ytick=[[0, stim.sum()]], xlabel='pos. in %s'%(stimN[i]))         
        
def plot_rewResp_2_firstLick(ax, dat):
    u_spk, u_lick = [], []
    for kn in dat:
        beh = dat[kn]['beh']
        stim_id, WallN, uniqWN = beh['stim_id'], beh['WallName'], beh['UniqWalls']
        rewStim = WallN==uniqWN[stim_id==2]

        late_FL = dat[kn]['FL_pos']>=20
        resp2FL = dat[kn]['resp2FL'][rewStim & late_FL]
        lick2FL = dat[kn]['lick2FL'][rewStim & late_FL]
        if resp2FL.shape[0]>10:
#             print(resp2FL.shape)
            resp2FL = (resp2FL - resp2FL.mean()) / resp2FL.std()
            u_spk.append(resp2FL.mean(0))
            u_lick.append(lick2FL.mean(0)  * 3)
    u_spk = np.array(u_spk)
    u_lick = np.array(u_lick)
    u0, sem0 = u_spk.mean(0), u_spk.std(0, ddof=1)/np.sqrt(u_spk.shape[0])
    u1, sem1 = u_lick.mean(0), u_lick.std(0, ddof=1)/np.sqrt(u_lick.shape[0])    
    ax.plot(u0, color='b', lw=1)
    ax.fill_between(np.arange(len(u0)), u0-sem0, u0+sem0, color='b', alpha=0.3, edgecolor='None')
    ax2 = ax.twinx()
    ax2.plot(u1, color='0.5', lw=1)
    ax2.fill_between(np.arange(len(u1)), u1-sem1, u1+sem1, color='0.5', alpha=0.3, edgecolor='None')
    utils.fmt(ax, ylabel='average activity (zscore)', boxoff=0, ylm=[-0.8, 1.3], ytick=[[0, 1, 2]], xlabel='time to first lick (s)',
              xtick=[np.linspace(0,30,11)[1::2]-0.5,np.linspace(0,10,11)[1::2].astype(int)-5], xlm=[-0.5, len(u0)-0.5])
    utils.fmt(ax2, ylabel='lick rate (counts/s)', boxoff=0, ylm=[-0.2, 4.5])
    ax.spines['top'].set_visible(False)
    ax2.spines['top'].set_visible(False)    
    
def plot_rewResp_2_Cue(ax, dat):
    u_spk, u_lick = [], []
    for kn in dat:
        beh = dat[kn]['beh']
        stim_id, WallN, uniqWN = beh['stim_id'], beh['WallName'], beh['UniqWalls']
        rewStim = WallN==uniqWN[stim_id==2]
        resp2Cue = dat[kn]['resp2Cue'][rewStim]
        lick2Cue = dat[kn]['lick2Cue'][rewStim]
        
        resp2Cue = (resp2Cue - np.nanmean(resp2Cue)) / np.nanstd(resp2Cue)
        u_spk.append(np.nanmean(resp2Cue, 0))
        u_lick.append(lick2Cue.mean(0) * 3) # recording frame rate is 3 Hz,
        
    u_spk = np.array(u_spk)
    u_lick = np.array(u_lick)
    u0, sem0 = np.nanmean(u_spk, 0), np.nanstd(u_spk, 0, ddof=1)/np.sqrt(u_spk.shape[0])
    u1, sem1 = u_lick.mean(0), u_lick.std(0, ddof=1)/np.sqrt(u_lick.shape[0])    
    ax.plot(u0, color='b', lw=1)
    ax.fill_between(np.arange(len(u0)), u0-sem0, u0+sem0, color='b', alpha=0.3, edgecolor='None')
    ax.axvline(15-0.5, linewidth=0.5, linestyle=':', color='0.5')
    ax2 = ax.twinx()
    ax2.plot(u1, color='0.5', lw=1)
    ax2.fill_between(np.arange(len(u1)), u1-sem1, u1+sem1, color='0.5', alpha=0.3, edgecolor='None')
    utils.fmt(ax, ylabel='average activity (zscore)', boxoff=0, ylm=[-0.7,2], xlabel='time to cue (s)', 
              xtick=[np.linspace(0,30,11)[1::2]-0.5,np.linspace(0,10,11)[1::2].astype(int)-5], xlm=[-0.5, len(u0)-0.5])
    utils.fmt(ax2, ylabel='lick rate (counts/s)', boxoff=0, ylm=[-0.3, 3.7])
    ax.spines['top'].set_visible(False)
    ax2.spines['top'].set_visible(False)        
    
def plot_rewResp_in_leaf2(ax, dat):
    u_spk, u_lick = [], []
    for kn in dat:
        beh = dat[kn]['beh']
        stim_id, WallN, uniqWN = beh['stim_id'], beh['WallName'], beh['UniqWalls']
        rewStim = WallN==uniqWN[stim_id==3]
        islick = utils.lickCount(dat[kn]['beh'], def_range=[0, 40])['inRange'][rewStim]
        if (len(islick)-islick.sum())>=5:
            resp = dat[kn]['resp'][rewStim][:, 0:40].mean(1)
            u_spk.append([resp[~islick].mean(0), resp[islick].mean(0)])
    u_spk = np.array(u_spk)
    u, sem = u_spk.mean(0), u_spk.std(0, ddof=1)/np.sqrt(u_spk.shape[0])
    ax.plot([0, 1], u_spk.T, color='c', lw=0.5, alpha=0.5)
    ax.plot([0, 1], u, color='k', lw=1)
    ax.errorbar([0, 1], u, sem, marker='s', color='k', markersize=3)
    utils.fmt(ax, ylabel='activity (zscore)', ylm=[-0.1, 0.5], ytick=[[0, 0.5]],
              xtick=[[0, 1], ['no lick\ntrials', 'lick\ntrials']], xlm=[-0.2, 1.2])    
    
def plot_stimResp_in_leaf2(ax, dat):
    u_spk, u_lick = [], []
    for kn in dat:
        beh = dat[kn]['beh']
        stim_id, WallN, uniqWN = beh['stim_id'], beh['WallName'], beh['UniqWalls']
        rewStim = WallN==uniqWN[stim_id==3]
        islick = utils.lickCount(dat[kn]['beh'], def_range=[0, 40])['inRange'][rewStim]
        if (len(islick)-islick.sum())>=5:
            resp = dat[kn]['stim_resp'][rewStim][:, 0:40].mean(1)
            u_spk.append([resp[~islick].mean(0), resp[islick].mean(0)])
    u_spk = np.array(u_spk)
    u, sem = u_spk.mean(0), u_spk.std(0, ddof=1)/np.sqrt(u_spk.shape[0])
    ax.plot([0, 1], u_spk.T, color='c', lw=0.5, alpha=0.5)
    ax.plot([0, 1], u, color='k', lw=1)
    ax.errorbar([0, 1], u, sem, marker='s', color='k', markersize=3)
    utils.fmt(ax, ylabel='activity (zscore)', ylm=[-0.1, 0.5], ytick=[[0, 0.5]],
              xtick=[[0, 1], ['no lick\ntrials', 'lick\ntrials']], xlm=[-0.2, 1.2])