<jupyter_start><jupyter_text>User-defined Parameters<jupyter_code>##################################################################################################################
## PARAMETERS THAT YOU SHOULD CHANGE
##################################################################################################################
# Path to data and atlas
mouse_name = 'IJ1'
session_name = 'Mar3'
root_dir = '/media/cat/4TBSSD/yuki/yongxu/data/'
data_folder = root_dir+mouse_name+'/'+session_name+'/'


# spatial_data_filename = "IA1_spatial.npy" #
spatial_data_filename =mouse_name+'pm_'+session_name+'_30Hz_code_04_trial_ROItimeCourses_30sec_pca_0.95_spatial.npy'
# temporal_data_filename = "IA1_temporal.npy" # 
temporal_data_filename = mouse_name+'pm_'+session_name+'_30Hz_code_04_trial_ROItimeCourses_30sec_pca_0.95.npy'

# atlas_filename = "maskwarp_1.npy" # contains 'atlas'
random_spatial_data_filename =mouse_name+'pm_'+session_name+'_30Hz_code_04_random_ROItimeCourses_30sec_pca_0.95_spatial.npy'
# temporal_data_filename = "IA1_temporal.npy" # 
random_temporal_data_filename = mouse_name+'pm_'+session_name+'_30Hz_code_04_random_ROItimeCourses_30sec_pca_0.95.npy'

atlas_filename = "atlas_split.npy"

# maxrank = how many max components per brain region. Set maxrank to around 4 for regular dataset.
maxrank = 1
  
# min_pixels = minimum number of pixels in Allen map for it to be considered a brain region
# default min_pixels = 100
min_pixels = 200

# loc_thresh = Localization threshold, i.e. percentage of area restricted to be inside the 'Allen boundary'
# default loc_thresh = 80
loc_thresh = 75

# r2_thresh = Fraction of variance in the data to capture with LocaNMF
# default r2_thresh = 0.99
r2_thresh = 0.96

# Do you want nonnegative temporal components? The data itself should also be nonnegative in this case.
# default nonnegative_temporal = False
nonnegative_temporal = False 

##################################################################################################################
## PARAMETERS THAT YOU SHOULD PROBABLY NOT CHANGE (unless you know what you're doing)
##################################################################################################################
# maxiter_hals = Number of iterations in innermost loop (HALS). Keeping this low provides a sort of regularization.
# default maxiter_hals = 20
maxiter_hals = 20

# maxiter_lambda = Number of iterations for the lambda loop. Keep this high for finding a good solution.
# default maxiter_lambda = 100
maxiter_lambda = 150

# lambda_step = Amount to multiply lambda after every lambda iteration. 
# lambda_init = initial value of lambda. Keep this low. default lambda_init = 0.000001
# lambda_{i+1}=lambda_i*lambda_step. lambda_0=lambda_init. default lambda_step = 1.35
lambda_step = 1.25
lambda_init = 1e-4

# # spatial_data_filename.shape
# a=np.load(data_folder+spatial_data_filename)
# print(a.shape)
# b=np.load(data_folder+temporal_data_filename)
# print(b.shape)<jupyter_output><empty_output><jupyter_text>Load & Format Data<jupyter_code>spatial = np.load(data_folder+spatial_data_filename)
# spatial_random = np.load(data_folder+random_spatial_data_filename)
# spatial=np.concatenate(spatial_trial,spatial_random,axis=0)
spatial = np.transpose(spatial,[1,0])
denoised_spatial_name = np.reshape(spatial,[128,128,-1])

temporal_trial = np.load(data_folder+temporal_data_filename)
temporal_random = np.load(data_folder+random_temporal_data_filename)
temporal=np.concatenate((temporal_trial,temporal_random),axis=0)
temporal = np.transpose(temporal,[1,0,2])
# denoised_temporal_name = np.reshape(temporal,[-1,42*601])
denoised_temporal_name = np.reshape(temporal,[-1,temporal.shape[1]*temporal.shape[2]])
print('loaded data\n',flush=True)

atlas = np.load('/home/cat/code/widefieldPredict/locanmf/atlas_fixed_pixel.npy')#['atlas'].astype(float)
#areanames = sio.loadmat(data_folder+atlas_filename)['areanames']
# atlas.shape
fig = plt.figure()
plt.imshow(atlas)
# fig=plt.figure(figsize=(20,15))
# for it in np.unique(atlas):
#     plotmap=np.zeros((atlas.shape)); plotmap.fill(np.nan); plotmap[atlas==it]=atlas[atlas==it]
# #     plt.subplot(5,6,it+1)
#     plt.imshow(plotmap,cmap='Spectral'); plt.axis('off'); plt.title('Allen region map'); plt.show(); 
#     plt.show()

print(denoised_temporal_name.shape)
print(denoised_spatial_name.shape)
print(atlas.shape)

fig = plt.figure()
plt.plot(denoised_temporal_name[:,:601].T); plt.show()

fig=plt.figure()
for i in np.arange(7):
    plt.imshow(denoised_spatial_name[:,:,i]); plt.show()
# Get data in the correct format
V=denoised_temporal_name
U=denoised_spatial_name

#
brainmask = np.ones(U.shape[:2],dtype=bool)

# Load true areas if simulated data
simulation=0

# Include nan values of U in brainmask, and put those values to 0 in U 
brainmask[np.isnan(np.sum(U,axis=2))]=False
U[np.isnan(U)]=0

# Preprocess V: flatten and remove nans
dimsV=V.shape
keepinds=np.nonzero(np.sum(np.isfinite(V),axis=0))[0]
V=V[:,keepinds]

# del arrays
# U.shape
# Check that data has the correct shapes. V [K_d x T], U [X x Y x K_d], brainmask [X x Y]
if V.shape[0]!=U.shape[-1]: 
    print('Wrong dimensions of U and V!')
    
print("Rank of video : %d" % V.shape[0]); print("Number of timepoints : %d" % V.shape[1]); 

# Plot the maximum of U
# plotmap=np.zeros((atlas.shape)); plotmap.fill(np.nan); plotmap[brainmask]=atlas[brainmask]
# fig=plt.figure()
# plt.imshow(plotmap,cmap='Spectral'); plt.axis('off'); plt.title('Allen region map'); plt.show(); 
# fig=plt.figure()
# plt.imshow(np.max(U,axis=2)); plt.axis('off'); plt.title('Max U'); plt.show()

# Perform the LQ decomposition. Time everything.
t0_global = time.time()
t0 = time.time()
if nonnegative_temporal:
    r = V.T
else:
    q, r = np.linalg.qr(V.T)
time_ests={'qr_decomp':time.time() - t0}<jupyter_output><empty_output><jupyter_text>Initialize LocaNMF<jupyter_code># Put in data structure for LocaNMF
video_mats = (np.copy(U[brainmask]), r.T)
rank_range = (1, maxrank, 1)
del U

# region_mats[0] = [unique regions x pixels] the mask of each region
# region_mats[1] = [unique regions x pixels] the distance penalty of each region
# region_mats[2] = [unique regions] area code
region_mats = LocaNMF.extract_region_metadata(brainmask, atlas, min_size=min_pixels)

region_metadata = LocaNMF.RegionMetadata(region_mats[0].shape[0],
                                         region_mats[0].shape[1:],
                                         device=device)

region_metadata.set(torch.from_numpy(region_mats[0].astype(np.uint8)),
                    torch.from_numpy(region_mats[1]),
                    torch.from_numpy(region_mats[2].astype(np.int64)))

# print (region_mats[1].shape)

# print (region_mats[2])

# print (region_mats[2].shape)
# grab region names
rois=np.load('/home/cat/code/widefieldPredict/locanmf/rois_50.npz')
rois_name=rois['names']
# rois_name

rois_ids=rois['ids']
# rois_ids
# Do SVD as initialization
if device=='cuda': 
    torch.cuda.synchronize()

# 
print('v SVD Initialization')
t0 = time.time()
region_videos = LocaNMF.factor_region_videos(video_mats,
                                               region_mats[0],
                                               rank_range[1],
                                               device=device)
#
if device=='cuda': 
    torch.cuda.synchronize()
print("\'-total : %f" % (time.time() - t0))
time_ests['svd_init'] = time.time() - t0
# 
low_rank_video = LocaNMF.LowRankVideo(
    (int(np.sum(brainmask)),) + video_mats[1].shape, device=device
)
low_rank_video.set(torch.from_numpy(video_mats[0].T),
                   torch.from_numpy(video_mats[1]))<jupyter_output><empty_output><jupyter_text>LocaNMF<jupyter_code># 
if device=='cuda': 
    torch.cuda.synchronize()

# 
print('v Rank Line Search')
t0 = time.time()

#locanmf_comps,loc_save 

(nmf_factors,
 loc_save,
 save_lam,
 save_scale,
 save_per,
 save_spa,
 save_scratch) = LocaNMF.rank_linesearch(low_rank_video,
                                             region_metadata,
                                             region_videos,
                                             maxiter_rank=maxrank,
                                             maxiter_lambda=maxiter_lambda,      # main param to tweak
                                             maxiter_hals=maxiter_hals,
                                             lambda_step=lambda_step,
                                             lambda_init=lambda_init, 
                                             loc_thresh=loc_thresh,
                                             r2_thresh=r2_thresh,
                                             rank_range=rank_range,
#                                             nnt=nonnegative_temporal,
                                             verbose=[True, False, False],
                                             sample_prop=(1,1),
                                             device=device
                                            )
#
if device=='cuda': 
    torch.cuda.synchronize()
    
# 
print("\'-total : %f" % (time.time() - t0))
time_ests['rank_linesearch'] = time.time() - t0
# print (loc_save.scale.data)

# print (loc_save.spatial.data.shape)

# print (loc_save.spatial.scratch.shape)

# how much of spatial components is in each region
ratio=torch.norm(loc_save.spatial.scratch,p=2,dim=-1)/torch.norm(loc_save.spatial.data,p=2,dim=-1)
print ("ratio: ", ratio)

# how much of the spatial component is inside the ROI
per=100*(ratio**2)
print ('percentage; ', per)

#
print ("The threshold is 75%.  all copmonents should be above, otherwise increase lambda iterations ")
# # 
# print("Number of components : %d" % len(locanmf_comps))

# computing the variance in each component;
print ("Compute the variance explained by each component ")
mov = torch.matmul(low_rank_video.spatial.data.t(),
                   low_rank_video.temporal.data)

# 
var = torch.mean(torch.var(mov, dim=1, unbiased=False))  # TODO: Precompute this
var_ests=np.zeros((len(nmf_factors)))
for i in np.arange(len(nmf_factors)):
    mov = torch.matmul(torch.index_select(nmf_factors.spatial.data,0,torch.tensor([i])).t(),
                       torch.index_select(nmf_factors.temporal.data,0,torch.tensor([i])))
    var_i = torch.mean(torch.var(mov, dim=1, unbiased=False))  
    
    # mean(var(dataest))/mean(var(data))
    var_ests[i] = var_i.item() / var.item()
# to return to this; should sum to 1, perhaps
plt.plot(var_ests); 
plt.show()

# 
print (np.argsort(var_ests))
print (var_ests[np.argsort(var_ests)])
print (var_ests.sum())

print ("TODO: verify that locaNMF svm decoding is similar to pca")

# compute the r-squared again
mov = torch.matmul(low_rank_video.spatial.data.t(),low_rank_video.temporal.data)
var = torch.mean(torch.var(mov, dim=1, unbiased=False))  # TODO: Precompute this

torch.addmm(beta=1,
            input=mov,
            alpha=-1,
            mat1=nmf_factors.spatial.data.t(),
            mat2=nmf_factors.temporal.data,
            out=mov)

r2_est = 1 - (torch.mean(mov.pow_(2)).item() / var.item())

print(r2_est); 

print ("TODO: Save rsquared components")

# # for each component what is the rsquared - OPTIONAL - for checks
# r2_ests=np.zeros((len(locanmf_comps)))
# for i in np.arange(len(locanmf_comps)):
#     mov = torch.matmul(low_rank_video.spatial.data.t(),low_rank_video.temporal.data)
#     var = torch.mean(torch.var(mov, dim=1, unbiased=False))  # TODO: Precompute this
    
#     # mov = data-dataest = data-mat1*mat2
#     torch.addmm(beta=1,
#                 input=mov,
#                 alpha=-1,
#                 mat1=torch.index_select(locanmf_comps.spatial.data,
#                                                 0,
#                                                 torch.tensor([i])).t(),
#                 mat2=torch.index_select(locanmf_comps.temporal.data,
#                                                 0,
#                                                 torch.tensor([i])),
#                 out=mov)
    
#     # r2_ests = 1-mean(var(data-dataest))/mean(var(data))
#     r2_ests[i] = 1 - (torch.mean(mov.pow_(2)).item() / var.item())
# # for each component what is the rsquared - OPTIONAL - for checks
# mov = torch.matmul(low_rank_video.spatial.data.t(),low_rank_video.temporal.data)
# var = torch.mean(torch.var(mov, dim=1, unbiased=False))  # TODO: Precompute this

# r2_ests_2=np.zeros((len(locanmf_comps)))
# for i in np.arange(len(locanmf_comps)):
#     mov_i = torch.matmul(torch.index_select(locanmf_comps.spatial.data,0,torch.tensor([i])).t(),
#                        torch.index_select(locanmf_comps.temporal.data,0,torch.tensor([i])))
#     var_i = torch.mean(torch.var(mov-mov_i, dim=1, unbiased=False))  
    
#     # 1 - mean(var(data-dataest))/mean(var(data))
#     r2_ests_2[i] = 1 - (var_i.item() / var.item())
# r2_ests_loo=np.zeros((len(locanmf_comps)))
# for i in np.arange(len(locanmf_comps)):
#     mov = torch.matmul(low_rank_video.spatial.data.t(),low_rank_video.temporal.data)
#     var = torch.mean(torch.var(mov, dim=1, unbiased=False))  # TODO: Precompute this

#     torch.addmm(beta=1,
#                 input=mov,
#                 alpha=-1,
#                 mat1=torch.index_select(locanmf_comps.spatial.data,
#                                                 0,
#                                                 torch.tensor(np.concatenate((np.arange(0,i),np.arange(i+1,len(locanmf_comps)))))).t(),
#                 mat2=torch.index_select(locanmf_comps.temporal.data,
#                                                 0,
#                                                 torch.tensor(np.concatenate((np.arange(0,i),np.arange(i+1,len(locanmf_comps)))))),
#                 out=mov)

#     r2_ests_loo[i] = 1 - (torch.mean(mov.pow_(2)).item() / var.item())
# _,r2_fit=LocaNMF.evaluate_fit_to_region(low_rank_video,
#                                            locanmf_comps,
#                                            region_metadata.support.data.sum(0),
#                                            sample_prop=(1, 1))


# Evaluate R^2
_,r2_fit=LocaNMF.evaluate_fit_to_region(low_rank_video,
                                        nmf_factors,
                                        region_metadata.support.data.sum(0),
                                        sample_prop=(1, 1)
                                       )
# 
print("R^2 fit on all data : %f" % r2_fit)

time_ests['global_time'] = time.time()-t0_global
# C is the temporal components
C = np.matmul(q,nmf_factors.temporal.data.cpu().numpy().T).T
print ("n_comps, n_time pts x n_trials: ", C.shape)
qc, rc = np.linalg.qr(C.T)
# back to visualizing variance
print(np.sum(var_ests))
plt.bar(np.arange(len(nmf_factors)),var_ests); plt.show()
# 
var_est=np.zeros((len(nmf_factors)))
for i in np.arange(len(nmf_factors)):
    var_est[i]=np.var(C[i,:])/np.var(C)
    
# locanmf_comps.regions.data<jupyter_output><empty_output><jupyter_text>Reformat spatial and temporal matrices, and save<jupyter_code># Assigning regions to components
region_ranks = []; region_idx = []
locanmf_comps = nmf_factors

for rdx in torch.unique(locanmf_comps.regions.data, sorted=True):
    region_ranks.append(torch.sum(rdx == locanmf_comps.regions.data).item())
    region_idx.append(rdx.item())

areas=region_metadata.labels.data[locanmf_comps.regions.data].cpu().numpy()
# Get LocaNMF spatial and temporal components
A=locanmf_comps.spatial.data.cpu().numpy().T
A_reshape=np.zeros((brainmask.shape[0],brainmask.shape[1],A.shape[1])); 
A_reshape.fill(np.nan)
A_reshape[brainmask,:]=A

# C is already computed above delete above
if nonnegative_temporal:
    C=locanmf_comps.temporal.data.cpu().numpy()
else:
    C=np.matmul(q,locanmf_comps.temporal.data.cpu().numpy().T).T
    
# Add back removed columns from C as nans
C_reshape=np.full((C.shape[0],dimsV[1]),np.nan)
C_reshape[:,keepinds]=C
C_reshape=np.reshape(C_reshape,[C.shape[0],dimsV[1]])

# Get lambdas
lambdas=np.squeeze(locanmf_comps.lambdas.data.cpu().numpy())
# A_reshape.shape

# c_p is the trial sturcutre
c_p=C_reshape.reshape(A_reshape.shape[2],int(C_reshape.shape[1]/1801),1801)

# 
c_plot=c_p.transpose((1,0,2))
c_plot.shape
# save LocaNMF data
data_folder = root_dir+mouse_name+'/'+session_name+'/'

areas_saved = []
for area in areas:
    idx = np.where(rois_ids==np.abs(area))[0]
    temp_name = str(rois_name[idx].squeeze())
    if area <0:
        temp_name += " - right"
    else:
        temp_name += " - left"
    

    areas_saved.append(temp_name)
    
np.savez(os.path.join(root_dir,mouse_name,session_name,'locanmf_trial.npz'),
        temporal = c_plot[:int(c_plot.shape[0]/2),:,:],
        areas = areas,
        names = areas_saved)

np.savez(os.path.join(root_dir,mouse_name,session_name,'locanmf_random.npz'),
        temporal = c_plot[int(c_plot.shape[0]/2):,:,:],
        areas = areas,
        names = areas_saved)


d = np.load('/media/cat/4TBSSD/yuki/yongxu/data/IJ1/Mar3/locanmf_trial.npz')
temporal = d['temporal']
areas = d['areas']
names = d['names']
print (temporal.shape)
print (areas)
print (names)
# from scipy.signal import savgol_filter

# t = np.arange(temporal.shape[2])/30 - 30
# max_vals = []
# for k in range(temporal.shape[1]):
#     temp = temporal[:,k].mean(0)
#     #plt.plot(temp)
#     temp2 = savgol_filter(temp, 15, 2)
    
#     plt.subplot(121)
#     plt.xlim(-15,0)
#     plt.plot(t,temp2)
    
#     m = np.max(temp2[:temp2.shape[0]//2])
#     max_vals.append(m)
    
#     plt.subplot(122)
#     plt.xlim(-15,0)
#     temp2 = temp2/np.max(temp2)
#     plt.plot(t,temp2)
    
#     #break
# plt.show()

# max_vals = np.array(max_vals)
# args = np.argsort(max_vals)[::-1]
# print (max_vals[args])
# print (names[args])


# temp componetns of behavior + control data vstacked
t_plot=temporal.transpose((1,0,2))

print (temporal.shape)
# Plot the distribution of lambdas. OPTIONAL
# If lots of values close to the minimum, decrease lambda_init.
# If lots of values close to the maximum, increase maxiter_lambda or lambda_step.
plt.hist(locanmf_comps.lambdas.data.cpu(), 
         bins=torch.unique(locanmf_comps.lambdas.data).shape[0])

plt.show()
print(locanmf_comps.lambdas.data.cpu())
region_name=region_mats[2]
region_name.shape

region_name

def parse_areanames_new(region_name,rois_name):
    areainds=[]; areanames=[];
    for i,area in enumerate(region_name):
        areainds.append(area)
        areanames.append(rois_name[np.where(rois_ids==np.abs(area))][0])
    sortvec=np.argsort(np.abs(areainds))
    areanames=[areanames[i] for i in sortvec]
    areainds=[areainds[i] for i in sortvec]
    return areainds,areanames

# 
region_name=region_mats[2]


# Get area names for all components
areainds,areanames_all = parse_areanames_new(region_name,rois_name)
areanames_area=[]
for i,area in enumerate(areas):
    areanames_area.append(areanames_all[areainds.index(area)])
    
# # Get area names for all components
areainds,areanames_all =parse_areanames_new(region_name,rois_name)
areanames_area=[]
for i,area in enumerate(areas):
    areanames_area.append(areanames_all[areainds.index(area)])
    
# Save results! - USE .NPZ File
print("LocaNMF completed successfully in "+ str(time.time()-t0_global) + "\n")
print("Results saved in "+data_folder+'locanmf_decomp_loc'+str(loc_thresh)+'.mat')

# Prefer to save c_p which is already converted
sio.savemat(data_folder+'locanmf_decomp_loc'+str(loc_thresh)+'.mat',
            {'C':C_reshape,
             'A':A_reshape,
             'lambdas':lambdas,
             'areas':areas,
             'r2_fit':r2_fit,
             'time_ests':time_ests,
             'areanames':areanames_area
            })


torch.cuda.empty_cache()
atlas_split=atlas
itt=0
fig=plt.figure() #figsize=(10,10))
b_=[]
for it in np.unique(atlas_split):
    if np.abs(it) !=0:
        plotmap=np.zeros((atlas_split.shape)); plotmap.fill(np.nan); plotmap[atlas_split==it]=atlas_split[atlas_split==it]
        plt.subplot(5,8,itt+1)
        plt.imshow(plotmap,cmap='Spectral'); plt.axis('off'); 
        plt.title(rois_name[np.where(rois_ids==np.abs(it))][0],fontsize=6); 
        b_.append(plotmap)
        #     plt.show()
        itt=itt+1
plt.tight_layout(h_pad=0.5,w_pad=0.5)
# how much of the spatial component is inside the ROI
# per
######################################################
##### PLACE TO LOOK FOR large spatial components #####
######################################################
fig=plt.figure()
for i in range(A_reshape.shape[2]):
    plt.subplot(5,8,i+1)
    plt.imshow(A_reshape[:,:,i])
    plt.title(areanames_area[i],fontsize=6)
plt.tight_layout(h_pad=0.5,w_pad=0.5)
plt.show()
# calculate ROI data
roi_spatial=np.zeros((A_reshape.shape[2],denoised_spatial_name.shape[2]))
for i in range(denoised_spatial_name.shape[2]):
    for j in range(A_reshape.shape[2]):
        A_masking=np.zeros((A_reshape[:,:,j].shape))
        A_masking[A_reshape[:,:,j]!=0]=1
        A_multiply=A_masking*denoised_spatial_name[:,:,i]
        roi_spatial[j,i]=np.sum(A_multiply)/np.sum(A_masking)
        
roi_data=[]
for s in range(temporal_trial.shape[0]):
    roi_each=roi_spatial@temporal_trial[s]
    roi_data.append(roi_each)
roi_save_trial=np.array(roi_data)
roi_save_trial.shape
# roi_spatial=np.zeros((A_reshape.shape[2],denoised_spatial_name.shape[2]))
# for i in range(denoised_spatial_name.shape[2]):
#     for j in range(A_reshape.shape[2]):
#         A_masking=np.zeros((A_reshape[:,:,j].shape))
#         A_masking[A_reshape[:,:,j]!=0]=1
#         A_multiply=A_masking*denoised_spatial_name[:,:,i]
#         roi_spatial[j,i]=np.sum(A_multiply)/np.sum(A_masking)
        
# roi_data=[]
# for s in range(temporal_random.shape[0]):
#     roi_each=roi_spatial@temporal_random[s]
#     roi_data.append(roi_each)
# roi_save_random=np.array(roi_data)
# roi_save_random.shape
# save ROI data
# np.save(save_folder+save_name+'trial.npy',roi_save_trial)
# np.save(save_folder+save_name+'random.npy',roi_save_random)<jupyter_output><empty_output><jupyter_text>Visualization of components<jupyter_code># CAT: Use actual averages not random data
C_area_rois = c_plot[:int(c_plot.shape[0]/2),:,:].mean(0)
print (C_area_rois.shape)
t = np.arange(C_area_rois.shape[1])/30.-30

clrs_local = ['magenta','brown','pink','lightblue','darkblue']

    

fig = plt.figure(figsize=(12,16))
############### SHREYA"S CODE ####################
# Spatial and Temporal Components: Summary
atlascolor=np.zeros((atlas.shape[0],atlas.shape[1],4))
A_color=np.zeros((A_reshape.shape[0],A_reshape.shape[1],4))
cmap=plt.cm.get_cmap('jet')
colors=cmap(np.arange(len(areainds))/len(areainds))

# 
for i,area_i in enumerate(areainds):
    if area_i not in areas:
        continue
    atlascolor[atlas==area_i,:]=colors[i,:]
    C_area=C[np.where(areas==area_i)[0],:]
    
    # 
    for z in range(len(names_plot)):
        if names_plot[z] in names_plot[z]:
            clr = clrs_local[z]
            break
      
    # 
    for j in np.arange(colors.shape[1]):
        A_color[:,:,j]=A_color[:,:,j]+colors[i,j]*A_reshape[:,:,np.where(areas==area_i)[0][0]]
        
        
        
#fig=plt.figure(figsize=(15,8))
ax1=fig.add_subplot(2,2,1)
ax1.imshow(atlascolor)
#ax1.set_title('Atlas Regions')
ax1.axis('off')
ax2=fig.add_subplot(2,2,3)
ax2.imshow(A_color)
#ax2.set_title('Spatial Components (One per region)')
ax2.axis('off')
ax3=fig.add_subplot(1,2,2)
axvar=0
print (areas)

names_plot = ['Retrosplenial','barrel','limb','visual','motor']

for i,area_i in enumerate(areainds):
    if area_i not in areas:
        continue
    # C_area=C[np.where(areas==area_i)[0][0],:min(1000,C.shape[1])]
    C_area=C_area_rois[np.where(areas==area_i)[0][0]]
    #ax3.plot(1.5*axvar+C_area/np.nanmax(np.abs(C_area)),color=colors[i,:])
    
    for z in range(len(names_plot)):
        if names_plot[z] in names_plot[z]:
            clr = clrs_local[z]
            break
        
        
    ax3.plot(t, 1.5*axvar+C_area*23,
             #color=colors[i,:],
             color = clr,
             linewidth=3)
    
    print (C_area.shape)
    axvar+=1

ax3.set_xlim(-15,0)
#ax3.set_title('Temporal Components (One per region)')
ax3.axis('off') 

if True:
    plt.savefig('/home/cat/locanmf.svg',dpi=300)
    plt.close()
else:
    plt.show()
data1 = np.load('/media/cat/4TBSSD/yuki/AQ2/tif_files/AQ2am_Dec17_30Hz/AQ2am_Dec17_30Hz_locanmf.npz', 
              allow_pickle=True)
names_rois = data1['names']
print (names_rois)
ids = data1['areas']
print (areas)
# Plotting all the regions' components
for i,area in enumerate(areas):
    try:
        fig=plt.figure(figsize=(20,4))
        ax1 = fig.add_subplot(1,3,1)
        plotmap_area = np.zeros((atlas.shape));
        plotmap_area.fill(np.nan);
        plotmap_area[brainmask] = atlas[brainmask]==area
        ax1.imshow(plotmap_area);
        ax1.set_title('Atlas '+areanames_area[i])
        ax1.axis('off')

        ax2 = fig.add_subplot(1,3,2)
        ax2.imshow(A_reshape[:,:,i])
        ax2.set_title('LocaNMF A [%s]'%(i+1))
        ax2.axis('off')

        ax3 = fig.add_subplot(1,3,3)
        ax3.plot(C[i,:min(1000,C.shape[1])],'k')
        if simulation:
            ax3.plot(V[np.where(area==trueareas)[0][0],:min(1000,V.shape[1])],'r');
            if i==0: ax3.legend(('LocaNMF','True'))
        ax3.set_title('LocaNMF C [%s]'%(i+1))
        ax3.axis('off')
        plt.show()
    except:
        pass

# Calculate Canonical Correlations between components in each pair of regions
corrmat=np.zeros((len(areainds),len(areainds)))
skipinds=[]
for i,area_i in enumerate(areainds):
    for j,area_j in enumerate(areainds):
        if i==0 and area_j not in areas:
            skipinds.append(j)
        C_i=C[np.where(areas==area_i)[0],:].T
        C_j=C[np.where(areas==area_j)[0],:].T
        if i not in skipinds and j not in skipinds:
            cca=CCA(n_components=1)
            cca.fit(C_i,C_j)
            C_i_cca,C_j_cca=cca.transform(C_i,C_j)
            try: C_i_cca=C_i_cca[:,0]
            except: pass
            try: C_j_cca=C_j_cca[:,0]
            except: pass               
            corrmat[i,j]=np.corrcoef(C_i_cca,C_j_cca)[0,1]
corrmat=np.delete(corrmat,skipinds,axis=0); 
corrmat=np.delete(corrmat,skipinds,axis=1);
corr_areanames=np.delete(areanames_all,skipinds)
# Plot correlations
fig=plt.figure()
plt.imshow(corrmat,cmap=plt.cm.get_cmap('jet')); plt.clim(-1,1); plt.colorbar(shrink=0.8)
plt.get_cmap('jet')
plt.xticks(ticks=np.arange(len(areainds)-len(skipinds)),labels=corr_areanames,rotation=90); 
plt.yticks(ticks=np.arange(len(areainds)-len(skipinds)),labels=corr_areanames); 
plt.title('CCA between all regions',fontsize=12)
plt.xlabel('Region i',fontsize=12)
plt.ylabel('Region j',fontsize=12)
plt.show()
# Save visualized components and correlations
# print('Saving postprocessing results!')
# postprocess.plot_components(A_reshape,C,areas,atlas,areanames,data_folder)
# postprocess.plot_correlations(A_reshape,C,areas,atlas,areanames,data_folder)
import numpy as np
data = np.load('/media/cat/4TBSSD/yuki/IA1/tif_files/IA1pm_Feb2_30Hz/IA1pm_Feb2_30Hz_locanmf.npz',
              allow_pickle=True)

trials = data['temporal_trial']
random = data['temporal_random']
print (trials.shape, random.shape)


trials = np.load('/media/cat/4TBSSD/yuki/IA1/tif_files/IA1am_Mar4_30Hz/IA1am_Mar4_30Hz_code_04_trial_ROItimeCourses_30sec_pca_0.95.npy')
print (trials.shape)
data = np.load('/media/cat/4TBSSD/yuki/IA1/tif_files/IA1am_Mar4_30Hz/IA1am_Mar4_30Hz_whole_stack_trial_ROItimeCourses_15sec_pca30components.npy')
print (data.shape)

spatial = np.load('/media/cat/4TBSSD/yuki/IA1/tif_files/IA1am_Mar4_30Hz/IA1am_Mar4_30Hz_whole_stack_trial_ROItimeCourses_15sec_pca30components_spatial.npy')
print (spatial.shape)<jupyter_output>(40000, 30)
(30, 16384)
{'fraction_non_alphanumeric': 0.1622770671951, 'fraction_numerical': 0.020933165195460277, 'mean_word_length': 3.4905355120530657, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 66, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '131543', 'license': 'MIT', 'n_tokens_mistral': 10718, 'n_tokens_neox': 10203, 'n_words': 1552, 'path': 'run_locaNMF_notebooks/run_locaNMF_and_yongxu_original_code.ipynb', 'repo_name': 'catubc/widefield', 'size': '27755'}