fig2.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_fig2_dat(root = r'D:\Results\Zhong-et-al-2025'):
    dat={}
    Beh = np.load(os.path.join(root, 'beh', 'Beh_sup_test1.npy'), allow_pickle=1).item()
    dat['example_lick_raster'] = utils.get_lick_raster(Beh['TX109_2023_04_18_1'])
    # load test1 performance
    beh = np.load(os.path.join(root, r'beh\Beh_sup_test1.npy'), allow_pickle=1).item()
    dat['mean_test1_perf'] = utils.get_mean_lick_response(beh, lick_typ='befRew')
    # load sorted spike
    fn = 'sup_test1_sort_spk.npy'
    dat['sort_spk'] = np.load(os.path.join(root, 'process_data', fn), allow_pickle=1).item() 
    return dat

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

    # example lick raster
    x,y, dx,dy, w,h =0.33,0.75, 0.123,0, 0.1,0.2
    ax_circle1 = fig.add_axes([x,y,w,h])
    lick_raster_plot(ax_circle1, dat['example_lick_raster']['sort_by_cue']['circle1'], 
                          show_reward=1, show_firstLick=0, title='circle1', tcolor='r')

    ax_circle2 = fig.add_axes([x+dx,y,w,h])
    lick_raster_plot(ax_circle2, dat['example_lick_raster']['sort_by_cue']['circle2'], 
                          show_reward=1, show_firstLick=0, title='circle2', tcolor='m')

    ax_leaf2 = fig.add_axes([x+2*dx,y,w,h])
    lick_raster_plot(ax_leaf2, dat['example_lick_raster']['sort_by_cue']['leaf2'], 
                          show_reward=1, show_firstLick=0, title='leaf2', tcolor='c')

    ax_leaf1 = fig.add_axes([x+3*dx,y,w,h])
    lick_raster_plot(ax_leaf1, dat['example_lick_raster']['sort_by_cue']['leaf1'], 
                          show_reward=1, show_firstLick=0, title='leaf1', tcolor='b')

    ################## performance  ######################
    x,y, dx,dy, w,h =0.85,0.725, 0,0.08, 0.15,0.27
    ax_beh = fig.add_axes([x,y,w,h])
    test_perf_plot(ax_beh, dat['mean_test1_perf'], title='', yn=1, xlm=[-0.3, 3.3])

    ################## position sequence in real data  ######################
    x,y, dx,dy, w,h =0.04,0.395, 0.1,0.17, 0.075,0.1

    axes_sequence = [fig.add_axes([x + i*dx, y+dy, w, h]) for i in range(4)]
    test1_peak_pos_scatter_plot(axes_sequence, dat['sort_spk'], mname='VR2', arname='mHV', stim_sort='leaf1', vmax = 1)

    axes_scatter = [fig.add_axes([x + i*dx, y, w, w*10.5/7]) for i in range(4)]
    peak_position_scatter_plot(axes_scatter, dat['sort_spk'], mname='VR2', arname='mHV', stim_sort='leaf1')

    ################## coeficient inside mHV  ######################    
    x,y, dx,dy, w,h =0.475,0.39, 0.08,0.12, 0.15,0.27
    ax_coef = fig.add_axes([x,y,w,h])
    seq_corr_in_mHV(ax_coef, dat['sort_spk'], stim_sort='leaf1')

    ################## coeficent plots  ######################
    x,y, dx,dy, w,h =0.7,0.39, 0.09,0.1, 0.3,0.27
    ax_coef = fig.add_axes([x,y,w,h])
    seq_corr_all_areas(ax_coef, root, stim_sort='leaf1')
    ax_coef.text(0.75, 0.85, 'task mice', color='g', transform=ax_coef.transAxes)
    ax_coef.text(0.75, 0.78, 'unsupervised', color=[0.46,0,0.23], transform=ax_coef.transAxes)
    ax_coef.text(0.75, 0.7, 'naive', color='k', transform=ax_coef.transAxes)

    ################## stimulus responses, average across neurons  ######################
    x,y, dx,dy, w,h =0.02,0.03, 0.07,0.175, 0.045,0.092
    resp_ax = np.array([ [fig.add_axes([x + i*dx, y+dy, w, h]) for i in range(4)],
              [fig.add_axes([x + i*dx, y, w, h]) for i in range(4)] ])
    example_mHV_stim_response(resp_ax, root)

    ######################   arrows   #############################
    x,y, dx,dy, w,h =0.283,0.03, 0.09,0.08, 0.1,0.275
    ax_arrows1 = fig.add_axes([x,y,w,h])
    arrows1(ax_arrows1)

    ################## famililar coding traces  ######################
    x,y, dx,dy, w,h =0.345,0.035, 0.09,0.08, 0.145,0.27
    ax_familiar_CD = fig.add_axes([x,y,w,h])
    example_mHV_coding_direction(ax_familiar_CD, root)

    ################## mean projection across mice for mHV  ######################
    x,y, dx,dy, w,h =0.52,0.035, 0.09,0.08, 0.15,0.27
    ax_cd = fig.add_axes([x,y,w,h])
    mHV_mean_cd_proj(ax_cd, root)

    ################## generalization index  ######################
    x,y, dx,dy, w,h =0.7,0.035, 0.1,0.1, 0.3,0.27
    ax_SI = fig.add_axes([x,y,w,h])
    SI_test1(ax_SI, root) 
    ax_SI.text(0.6, 0.12, 'task mice', color='g', transform=ax_SI.transAxes)
    ax_SI.text(0.6, 0.05, 'unsupervised', color=[0.46,0,0.23], transform=ax_SI.transAxes)
    ax_SI.text(0.6, 0.19, 'naive', color='k', transform=ax_SI.transAxes)    
    
    
    ax_text.text(0.315, 1.01, r'$\bf{a}$ Example lick rasters', fontsize=5.5)
    ax_text.text(0.835, 1.01, r"$\bf{b}$ Licking behavior in $test1$", fontsize=5.5)
    ax_text.text(0.01, 0.689, r"$\bf{c}$ Example sequential responses of, leaf1-selective neurons (medial, task mouse)", fontsize=5.5)

    ax_text.text(.43, .689, r"$\bf{d}$ Sequence similarity ($r$, medial, task mice)", fontsize=5.5)
    ax_text.text(0.66, 0.689, r"$\bf{e}$ Sequence similarity ($r$, all areas)", fontsize=5.5) 
    

    ax_text.text(.01, .325, r"$\bf{f}$ Example leaf1-selective population (medial, task mouse)", fontsize=5.5)
    ax_text.text(0.01, 0.145, r"$\bf{g}$ Example circle1-selective population (medial, task mouse)", fontsize=5.5)  

    ax_text.text(.33, .325, r"$\bf{h}$ Coding direction of leaf1-circle1 (medial, task mouse, test trials)", fontsize=5.5)
    ax_text.text(0.69, 0.325, r"$\bf{i}$ Similarity index ($SI$) on new stimuli", fontsize=5.5)     

def lick_raster_plot(ax, lick, show_reward=1, show_firstLick=0, title='', xlm=[0, 60], tcolor='k'):
    plt.sca(ax)
    SoundPos = lick['SoundPos']
    RewPos = lick['RewPos']
    isRew = lick['isRew']
    LickPos = lick['LickPos']
    LickTr = lick['LickTr']
    fLPos = lick['firstLickPos']
    fLkTr = lick['firstLickTr']
    ntrials = SoundPos.shape[0]

    ax.scatter(LickPos, LickTr, marker='o', s=0.5, color='k', alpha=1, linewidth=0)
    ax.scatter(SoundPos, np.arange(ntrials), marker='o', s=2, color='purple', alpha=1, linewidth=0)
    if show_reward & isRew:
        ax.scatter(RewPos, np.arange(ntrials), marker='o', s=2, color='b', alpha=1, linewidth=0)        
    if show_firstLick:
        ax.scatter(fLPos, fLkTr, marker='o', s=1, color='brown', alpha=1)
    ax.axvline(40, lw=0.5, linestyle='--', color='k')
    utils.fmt(ax, xtick=[[0,20,40,60], [0,2,4,6]], ytick=[[0, ntrials]], tcolor=tcolor,
          ylabel='tirals', xlabel='position (m)', title=title, xlm=xlm, ylm=[0, ntrials], y_invert=1, ypad=-5)
    
def test_perf_plot(ax, perf, title='', yn=1, xlm=[-0.3, 1.3]):
    r = perf['u_sem']
    SID = [0, 1, 3, 2] # xpos for circle1, circle2, leaf2, leaf1
    u, sem = r[:, SID, 0].mean(0), r[:, SID, 0].std(0, ddof=1)/np.sqrt(r.shape[0])
    ax.plot(np.arange(4), r[:, SID, 0].T, 'k-', lw=0.5, alpha=0.5)
    ax.plot(np.arange(4), u, 'k-', lw=2)
    cols = np.array(['r', 'm', 'c', 'b'])
    for i in range(4):
        ax.errorbar(i, u[i], yerr=sem[i], marker='s', markersize=4, color=cols[i], markeredgecolor='k', markeredgewidth=0.5)
    yln = 'anticipatory licking (%trials)' 
    utils.fmt(ax, xtick=[np.arange(len(perf['stimuli'])), perf['stimuli'][SID]], ytick=[[0, 0.5, 1], [0, 50, 100]],
          ylabel=yln, title=title, xlm=xlm, ylm=[0, 1])
    xticklabels = ax.get_xticklabels()    
    for label, color in zip(xticklabels, cols[SID]):
        label.set_color(color)    
        
def test1_peak_pos_scatter_plot(ax, dat, mname='VR2', arname='mHV', stim_sort='leaf1', vmax = 1):
    seq_spk = dat['spk_sort'][mname][arname]['sorted_by_odd_%s'%(stim_sort)]['target']
    cols = ['r', 'm', 'c', 'b']
    tn = ['circle1', 'circle2', 'leaf2', 'leaf1']
    yn = 'neurons (sorted)'
    xn = 'position (m)'
    for i, ic in enumerate([0, 1, 3, 2]):
        ax[i].imshow(seq_spk[ic], vmin=0, vmax=vmax, cmap='gray_r')
        ax[i].axvline(40, linestyle='--', lw=0.5, color='k')
        if i>0:
            xn, yn = '', ''         
        utils.fmt(ax[i], title=tn[i], tcolor=cols[i], y_invert=1, xlabel=xn, ylabel=yn, tpad=0, xpad=-1,
                 xtick=[[0, 20, 40, 60], [0, 2, 4, 6]]) 
        
def positions_scatter(ax, pos1, peak2, mk='s', ms=2, alpha=0.5):
    ax.scatter(pos1, peak2, s=ms, marker=mk, c='k', alpha=alpha, edgecolor='None')
    utils.fmt(ax, xtick=[[0,40],[0,4]], ytick=[[0,40],[0,4]]) 
    
def peak_position_scatter_plot(ax, dat, mname='VR2', arname='mHV', stim_sort='leaf1'):
    cols = ['r', 'm', 'c', 'b']
    tn = ['circle1', 'circle2', 'leaf2', 'leaf1']
    targ_pos = dat['spk_sort'][mname][arname]['sorted_by_odd_%s'%(stim_sort)]['target_maxPos']
    ref_pos = dat['spk_sort'][mname][arname]['sorted_by_odd_%s'%(stim_sort)]['reference_maxPos']
    yn = 'poxition in odd test trial'
    for i, ic in enumerate([0, 1, 3, 2]):
        positions_scatter(ax[i], ref_pos, targ_pos[ic])
        r = np.corrcoef(ref_pos, targ_pos[ic])[0, 1] 
        if i>0:
            yn = ''
        if tn[i]==stim_sort:
            xn = 'position in\neven test trial (m)'
        else:
            xn = 'position (m)'
        utils.fmt(ax[i], title='$r=%.2f$'%(r), tcolor=cols[i], xlabel=xn, ylabel=yn, tpad=0, xpad=-3)   
        
def seq_corr_in_mHV(ax, dat, stim_sort='leaf1'):
    cols = ['r', 'm', 'c', 'b']
    tn = ['circle1', 'circle2', 'leaf2', 'even\nleaf1']
    r = utils.get_seq_corr(dat, stim_sort=stim_sort)[:, 1, :4][:, [0, 1, 3, 2]]
    u, sem = r.mean(0), r.std(0, ddof=1)/np.sqrt(r.shape[0])
    ax.plot(u, lw=1.2, color='k')
    for i, ic in enumerate(cols):
        ax.errorbar(i, u[i], yerr=sem[i], marker='s', markersize=4, color=cols[i], markeredgecolor='k', markeredgewidth=0.5)
    yn = 'correlation with leaf1-sequence (odd trials)'
    ax.plot(r.T, color='k', lw=0.5, alpha=0.5)
    utils.fmt(ax, xtick=[np.arange(4), tn], ylabel=yn)    
    xticklabels = ax.get_xticklabels()    
    for label, color in zip(xticklabels, cols):
        label.set_color(color)           
        
def seq_corr_all_areas(ax, root, stim_sort='leaf1'):
    fn0 = 'sup_test1_sort_spk.npy'
    fn1 = 'unsup_test1_sort_spk.npy'
    fn2 = 'naive_test1_sort_spk.npy'
    cols = ['g', [0.46,0, 0.23], 'k']
    u, sem = np.empty((3, 4, 2)), np.empty((3, 4, 2))
    for f, fn in enumerate([fn0, fn1, fn2]): 
        dat = np.load(os.path.join(root, 'process_data', fn), allow_pickle=1).item()
        r = utils.get_seq_corr(dat, stim_sort=stim_sort)[:, :, 2:4]
        u[f] = r.mean(0)
        sem[f] = r.std(0, ddof=1)/np.sqrt(r.shape[0])
    for f,fcol in enumerate(cols):  
        for a in range(4):
            ax.plot(np.array([0.5, 0])+a, u[f, a], color=fcol, lw=1)
            ax.errorbar(np.array([0.5, 0])+a, u[f, a], yerr=sem[f, a], marker='s', markersize=3, color=fcol, markeredgecolor='k', markeredgewidth=0.5)
    yn = 'correlation with leaf1-sequence (odd trials)'
    utils.fmt(ax, xtick=[[0, 0.5], ['leaf2', 'even\nleaf1']], ylabel=yn, ytick=[[0, 0.5, 1]], ylm=[0, 1])    
    xticklabels = ax.get_xticklabels()    
    for label, color in zip(xticklabels, ['c', 'b']):
        label.set_color(color)          
    for t,txt in enumerate(['V1', 'mHV', 'lHV', 'aHV']):
        ax.text(0.1 + 0.25*t, 0.98, txt, transform=ax.transAxes)        
        
def example_mHV_stim_response(ax, root):
    fn = r'process_data\sup_test1_coding_direction.npy'
    dat = np.load(os.path.join(root, fn), allow_pickle=1).item()   
    stimN = ['circle1', 'circle2', 'leaf2', 'leaf1']
    cols = ['r', 'm', 'c', 'b']
    resp1 = dat['proj_2_stim1']['VR2']['mHV']
    resp2 = dat['proj_2_stim2']['VR2']['mHV']
    sid = [0, 1, 3, 2] # stimulus order ['circe1', 'circel2', 'leaf2', 'leaf1']
    for s in range(4):
        ax[0, s].imshow(resp1[sid[s]], cmap='gray_r', vmin=0, vmax=1)
        ax[0, s].axvline(50, linestyle='--', color='k', linewidth=0.5)
        utils.fmt(ax[0, s], xlm=[dat['pos_from_prev'], dat['pos_from_prev']+60], title=stimN[s], tcolor=cols[s],
                  xtick=[[10, 30, 50, 70], [0, 2, 4, 6]], ytick=[[0, resp1[sid[s]].shape[0]]], tpad=0, boxoff=0)
        ax[1, s].imshow(resp2[sid[s]], cmap='gray_r', vmin=0, vmax=1)
        ax[1, s].axvline(50, linestyle='--', color='k', linewidth=0.5)
        utils.fmt(ax[1, s], xlm=[dat['pos_from_prev'], dat['pos_from_prev']+60], title=stimN[s], tcolor=cols[s],
                  xtick=[[10, 30, 50, 70], [0, 2, 4, 6]], ytick=[[0, resp2[sid[s]].shape[0]]], tpad=0, boxoff=0)        
        
def arrows1(ax):
    ax.plot(np.array([0, 0]),[0, 1.35],'k-', lw=2)
    ax.plot(np.array([0, 0]),[2.55, 3.85],'k-', lw=2)
    ax.plot([0, 0.5, 0.5], [0.75, 0.75, 1.5],'k-', lw=0.5)
    ax.plot([0, 0.5, 0.5], [3.25, 3.25, 2.5],'k-', lw=0.5)
    ax.arrow(0.6, 2, 0.1, 0, head_width =0.3, width=0.1, head_length=0.1, color='k')
    ax.text(0.15, 0.5, 'subtract', rotation=-90, transform=ax.transAxes, verticalalignment='center')
    utils.fmt(ax, axis_off='off', xlm=[0, 2.2], ylm=[0, 4])    
    
def example_mHV_coding_direction(ax, root):
    fn = r'process_data\sup_test1_coding_direction.npy'
    dat = np.load(os.path.join(root, fn), allow_pickle=1).item()   
    stimN = ['circle1', 'circle2', 'leaf2', 'leaf1']
    cols = ['r', 'm', 'b', 'c'] # ['circe1', 'circel2', 'leaf2', 'leaf1']
    resp1 = dat['proj_2_stim1']['VR2']['mHV']
    resp2 = dat['proj_2_stim2']['VR2']['mHV']
    for s in range(4):
        diff = resp1[s]-resp2[s]
        ax.plot(diff.T, color=cols[s], lw=0.5, alpha=0.1)
        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 
        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='projection (a.u.)', ypad=-2)
        ax.axvline(50, linestyle='--', color='k', linewidth=0.5)   
        
def mHV_mean_cd_proj(ax, root):
    fn = 'sup_test1_coding_direction.npy'
    cd_proj = utils.load_coding_direction(os.path.join(root, 'process_data'), fn)['proj_tr_mean'][:, 1, :4] # take medial area
    cd_proj = cd_proj[:, [0, 1, 3, 2]].astype(float)
    print(cd_proj.shape)
    u, sem = np.nanmean(cd_proj, 0), np.nanstd(cd_proj, 0, ddof=1)/np.sqrt(cd_proj.shape[0])
    ax.plot(cd_proj.T, 'k-', lw=0.5, alpha=0.5)
    ax.plot(u, lw=1.2, color='k')
    cols = ['r', 'm', 'c', 'b']
    for i, ic in enumerate(cols):
        ax.errorbar(i, u[i], yerr=sem[i], marker='s', markersize=4, color=cols[i], markeredgecolor='k', markeredgewidth=0.5)
    yn = 'average projection (a.u.)'
    utils.fmt(ax, xtick=[np.arange(4), ['circle1', 'circle2', 'leaf2', 'leaf1']], ylabel=yn, ytick=[[-1, 0, 1]], ypad=0)    
    xticklabels = ax.get_xticklabels()    
    for label, color in zip(xticklabels, cols):
        label.set_color(color)   
        
def SI_test1(ax, root):
    fns = ['naive_test1_coding_direction.npy', 'sup_test1_coding_direction.npy', 'unsup_test1_coding_direction.npy']
    cols = ['k', 'g', [0.46,0, 0.23]]
    for f,fn in enumerate(fns):
        cd_proj = utils.load_coding_direction(os.path.join(root, 'process_data'), fn)['proj_tr_mean'][:, :, :4] # take medial area
        cd_proj = cd_proj[:, :, [0, 1, 3, 2]].astype(float)
        dx = abs(cd_proj[:, :, 1:3]-cd_proj[:, :, 0:1])
        dy = abs(cd_proj[:, :, 3:4]-cd_proj[:, :, 1:3])
        dxy = cd_proj[:, :, 3:4]-cd_proj[:, :, 0:1]
        SI = (dx-dy) / dxy
        u, sem = np.nanmean(SI, 0), np.nanstd(SI, 0, ddof=1)/np.sqrt(SI.shape[0])
        for a in range(4):
            ax.plot(np.array([0, 0.5])+a, u[a], lw=1.2, color=cols[f])
            ax.errorbar(np.array([0, 0.5])+a, u[a], yerr=sem[a], marker='s', markersize=3, color=cols[f], markeredgecolor='k', markeredgewidth=0.5)
    yn = 'similarity index ($SI$)'
    utils.fmt(ax, xtick=[[0, 0.5], ['circle2', 'leaf2']], 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, ['m', 'c']):
        label.set_color(color)    
    for t,txt in enumerate(['V1', 'mHV', 'lHV', 'aHV']):
        ax.text(0.1 + 0.25*t, 0.98, txt, transform=ax.transAxes)