<jupyter_start><jupyter_text>Import packages<jupyter_code>import numpy as np
import numpy.random as npr
from matplotlib import pyplot as plt
%matplotlib inline
from scipy.linalg import orth
import time

from ssa.models import fit_ssa, weighted_pca, weighted_rrr
from ssa.util import get_sample_weights
# %load_ext autoreload
# %autoreload 2<jupyter_output><empty_output><jupyter_text>Simulate data<jupyter_code>#Function that creates a sine wave for a given amount of time (T),
#where the number of cycles (c) occurs during that time
def create_sine(T,c):
    tau=T/(2*np.pi)/c
    return np.sin(np.arange(0,T)/tau)<jupyter_output><empty_output><jupyter_text>Generate the simulated data X0 (size [Time x Num_neurons1]) and Y0 (size [Time x Num_neurons2])<jupyter_code>np.random.seed(0) #To get the same simulated data

T=1200 #Time
N1=50 #Number of neurons in population 1
N2=60 #Number of neurons in population 2
R_shared=3 #Number of shared dimensions in lowD representations
R1=3 #Number of unshared dimensions in lowD representation of pop 1
R2=3 #Number of unshared dimensions in lowD representation of pop 2


#Orthogonal matrix that projects low dimensional space to full neural space of pop1
U_tmp=orth(npr.randn(R_shared+R1,N1).T).T 

#Orthogonal matrix that projects low dimensional space to full neural space of pop2
V_tmp=orth(npr.randn(R_shared+R2,N2).T).T 

#Create shared low dimensional space
Z_shared=np.zeros([T,R_shared])
for i in range(R_shared):
    Z_shared[300*i:300*i+500,i]=create_sine(500,i+1)

#Create low dimensional spaces unique to each population
Zx=npr.randn(T,R1)
Zy=npr.randn(T,R2)
    
#Create high-dimensional neural activity    
b1=npr.randn(N1) #Offset of neurons
b2=npr.randn(N2) #Offset of neurons

X0=Z_shared@U_tmp[:R_shared]+Zx@U_tmp[R_shared:]+b1 #Project into high-dimensional space and add offset
Y0=Z_shared@V_tmp[:R_shared]+Zy@V_tmp[R_shared:]+b2 #Project into high-dimensional space and add offset

# X0=X0+.1*npr.randn(X0.shape[0],X0.shape[1]) #Add noise
# Y0=Y0+.1*npr.randn(Y0.shape[0],Y0.shape[1]) #Add noise<jupyter_output><empty_output><jupyter_text>Preprocess data (optional) I have found that the method usually works better when zero-centering the data.In this specific example, if you don't zero-center the data, it will take ~5000 iterations to converge to the ground truth, rather than ~2000.<jupyter_code># X=np.copy(X0)
# Y=np.copy(Y0)


X=np.copy(X0-np.mean(X0,axis=0)[None,:])
Y=np.copy(Y0-np.mean(Y0,axis=0)[None,:])<jupyter_output><empty_output><jupyter_text>Set required model parameters<jupyter_code>#Number of shared dimensions in the low-D model you're fitting
R_est=4<jupyter_output><empty_output><jupyter_text>Set optional model parameters All of these have default values, so it's not essential to set them. The values listed below are the default values.<jupyter_code>#Strength of the sparsity penalty
lam=.01

#Number of epochs of model fitting
n_epochs=3000

#Learning rate of model fitting
lr=.001

#Whether to print the model error while fitting
verbose=True
#How much to weight each data point in time
#(this can be helpful for making sure dimensions still aim to explain time points with low activity)

sample_weights=np.ones([X.shape[0],1]) #Weight equally

# sample_weights=get_sample_weights(Y) #Weight inversely to norm of activity at each time point<jupyter_output><empty_output><jupyter_text>Fit Reduced Rank Regression Model (for comparison)<jupyter_code>#Fit weighted RRR
#Note that this function does not automatically subtract the mean from the data
__, U_est_rrr,V_est_rrr,b_rrr = weighted_rrr(X,Y,R_est,sample_weights)

#Get the low dimensional representation (the principal components)
rrr_latent = X@U_est_rrr<jupyter_output><empty_output><jupyter_text>Fit SSA Model<jupyter_code>#Fit SSA
model, latent, y_pred = fit_ssa(X=X,Y=Y,R=R_est,sample_weight=sample_weights,lam=lam,lr=lr,n_epochs=n_epochs,verbose=verbose)

#Fit SSA without all the optional parameters
# model, latent, y_pred = fit_ssa(X=X,Y=Y,R=R_est)

#Get the low dimensional representation
ssa_latent=latent.detach().numpy()<jupyter_output>Training loss before training 3550.931640625
Epoch 0: train loss: 3550.931640625
Epoch 10: train loss: 3536.5634765625
Epoch 20: train loss: 3533.59033203125
Epoch 30: train loss: 3533.450439453125
Epoch 40: train loss: 3533.2470703125
Epoch 50: train loss: 3533.028564453125
Epoch 60: train loss: 3532.96533203125
Epoch 70: train loss: 3532.91015625
Epoch 80: train loss: 3532.85888671875
Epoch 90: train loss: 3532.81591796875
Epoch 100: train loss: 3532.769775390625
Epoch 110: train loss: 3532.72314453125
Epoch 120: train loss: 3532.67333984375
Epoch 130: train loss: 3532.620849609375
Epoch 140: train loss: 3532.564453125
Epoch 150: train loss: 3532.5029296875
Epoch 160: train loss: 3532.436767578125
Epoch 170: train loss: 3532.36572265625
Epoch 180: train loss: 3532.28955078125
Epoch 190: train loss: 3532.20751953125
Epoch 200: train loss: 3532.119873046875
Epoch 210: train loss: 3532.026611328125
Epoch 220: train loss: 3531.927978515625
Epoch 230: train loss: 3531.82421875
Epoch 240: [...]<jupyter_text>Plot results  Plot unordered lowD representations<jupyter_code>#Ground truth
Z_extra=np.zeros([T,R_est])
Z_extra[:,:R_shared]=Z_shared


plt.figure(figsize=(15,5))
for i in range(R_est):
    
    #Plot ground truth
    plt.subplot(R_est,3,3*i+1)
    plt.plot((Z_extra)[:,i]) 
    
    plt.ylim([-1.1, 1.1])
    plt.yticks([])
    if i<R_est-1:
        plt.xticks([])
    else:
        plt.xlabel('Time')
    
    # Plot SSA results
    plt.subplot(R_est,3,3*i+2)
    plt.plot(ssa_latent[:,i])
    
    plt.ylim([-1.1, 1.1])
    plt.yticks([])    
    if i<R_est-1:
        plt.xticks([])
    else:
        plt.xlabel('Time')

    # Plot PCA results
    plt.subplot(R_est,3,3*i+3)
    plt.plot(rrr_latent[:,i])
    
    plt.ylim([-1.1, 1.1])
    plt.yticks([])
    if i<R_est-1:
        plt.xticks([])
    else:
        plt.xlabel('Time')        

#Titles
plt.subplot(R_est,3,1)
plt.title('True LowD Projections')

plt.subplot(R_est,3,2)
plt.title('SSA LowD Projections')

plt.subplot(R_est,3,3)
plt.title('RRR LowD Projections')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1088259109311741, 'fraction_numerical': 0.09117408906882592, 'mean_word_length': 5.413291796469367, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '53370', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2694, 'n_tokens_neox': 2315, 'n_words': 609, 'path': 'notebooks/Example_2pops.ipynb', 'repo_name': 'jglaser2/ssa', 'size': '6175'}