<jupyter_start><jupyter_text>define params globally first before passing to RDM, mess w/ model params and explain what can do. (can turn on and off dale's law). some masking, input output connectivity. train_params variables stuff<jupyter_code>dt = 1 # time-step
tau = 10 # intrinsic time constant of neural state decay
T = 200 # time to run for (number of steps is T/dt)
N_batch = 50 # number of trials per training step
N_rec = 100 # number of recurrent units
name = 'basicModel'
np.random.rand() * T
from psychrnn.tasks.task import Task


"""
Binary Decision task. 
Takes two channels of noisy input.
Binary output with a one hot encoding towards the higher mean channel
"""
class sinewave(Task):
    def __init__(self, dt, tau, T, N_batch):
        super(sinewave,self).__init__(1, 1, dt, tau, T, N_batch)
                                    #N_in, N_out
    def generate_trial_params(self, batch, trial):

        # ----------------------------------
        # Define parameters of a trial
        # ----------------------------------
        params = dict()
        params['onset'] = np.random.rand() * T * 0.8
        params['choice'] = np.random.choice([0,1,2])
        freqs = [0.03,0.03,0.03]
        params['freq'] = freqs[params['choice']]
        params['onset_time'] = np.random.random() * self.T / 8.0
        params['stim_duration'] = np.random.random() * self.T *10
        return params

    def trial_function(self, t, params):

        # ----------------------------------
        # Initialize with noise
        # ----------------------------------
        x_t = np.zeros(self.N_in)
        y_t = np.zeros(self.N_out)
        mask_t = np.ones(self.N_out)

        # ----------------------------------
        # Retrieve parameters
        # ----------------------------------
        freq = params['freq']
        onset = params['onset_time']
        stim_dur = params['stim_duration']

        # ----------------------------------
        # Compute values
        # ----------------------------------
        x_t = np.sin(t * params['freq']) 

        y_t = np.sin(t * params['freq']) 

        if t < onset + stim_dur:
            mask_t = np.zeros(self.N_out)

        return x_t, y_t, mask_t

        return x_t, y_t, mask_t
    
    
class noisy_decode(Task):
    def __init__(self, dt, tau, T, N_batch):
        super(noisy_decode,self).__init__(50, 50, dt, tau, T, N_batch)
                                    #N_in, N_out
    def generate_trial_params(self, batch, trial):

        # ----------------------------------
        # Define parameters of a trial
        # ----------------------------------
        params = dict()
        params['choice'] = np.random.choice([0,1,2,3,4,5])
        rates = [10,11,12,13,14,15]
        params['rate'] = rates[params['choice']]
        return params

    def trial_function(self, t, params):

        # ----------------------------------
        # Initialize with noise
        # ----------------------------------
        x_t = np.zeros(self.N_in)
        y_t = np.zeros(self.N_out)
        mask_t = np.ones(self.N_out)

        # ----------------------------------
        # Retrieve parameters
        # ----------------------------------
        rate = params['rate']
        

        # ----------------------------------
        # Compute values
        # ----------------------------------
        x_t = np.random.randn()*3+ params['rate']

        y_t = params['rate']

        

        return x_t, y_t, mask_t

        return x_t, y_t, mask_t
    
    
class many2one_noisy_decode(Task):
    def __init__(self, dt, tau, T, N_batch):
        super(many2one_noisy_decode,self).__init__(10, 1, dt, tau, T, N_batch)
                                    #N_in, N_out
    def generate_trial_params(self, batch, trial):

        # ----------------------------------
        # Define parameters of a trial
        # ----------------------------------
        params = dict()
        params['rate'] = np.random.randn()*3+10
        params['neuron_noise'] = np.random.choice([0,1,2,3])
        params['rate_noise'] = np.random.choice([0,1,2,3])
        return params

    def trial_function(self, t, params):

        # ----------------------------------
        # Initialize with noise
        # ----------------------------------
        x_t = np.zeros(self.N_in)
        y_t = np.zeros(self.N_out)
        mask_t = np.ones(self.N_out)

        # ----------------------------------
        # Retrieve parameters
        # ----------------------------------
        rate = params['rate']
        

        # ----------------------------------
        # Compute values
        # ----------------------------------

        x_t = np.random.randn(N_in)*params+ params['rate']

        y_t = params['rate']

        

        return x_t, y_t, mask_t

    
def gaussian_input(N_in, mu = 0, sigma = 1):
    y = (1/(sigma * np.sqrt(2 * np.pi)) *np.exp( - (np.arange(N_in) - mu)**2 / (2 * sigma**2)))
    y = y/np.max(y)
    return y

onset = 20
duration = 100
mu = 25
sigma = 3
neuron_noise = 3
num_steps = int(T/dt)
baseline_rate = 10
peak_rate = 25

inputs = np.zeros((N_in,num_steps)) + baseline_rate
outputs = np.zeros((N_in,num_steps)) + baseline_rate
true_rates = gaussian_input(N_in = N_in, mu = mu, sigma = sigma)*peak_rate 

for i,t in enumerate(np.arange(0,T,dt)):
    if (t>onset) and (t<(onset+duration)):
        outputs[:,i] += true_rates
        inputs[:,i] += true_rates 

inputs = inputs + np.random.randn(*inputs.shape)*neuron_noise 
inputs[inputs<0] = 0

swv = sinewave(dt = dt, tau = tau, T = T, N_batch = N_batch)
gen = swv.batch_generator()
tsk = noisy_decode(dt = dt, tau = tau, T = T, N_batch = N_batch)
gen = tsk.batch_generator()
params = tsk.__dict__
print(params)<jupyter_output>{'N_batch': 50, 'N_in': 1, 'N_out': 1, 'dt': 10, 'tau': 100, 'T': 2000, 'alpha': 0.1, 'N_steps': 200}<jupyter_text>Generate *N_batch* trials to be used in training We add in a few params that Basic(RNN) needs but that RDM doesn't generate for us.<jupyter_code>params['name'] = name #Used to scope out a namespace for global variables.
params['N_rec'] = N_rec<jupyter_output><empty_output><jupyter_text>There are some other optional parameters we can add in:<jupyter_code>params['dale_ratio'] = None # Default: None -- when the dale_ratio is set, dale's law is applied
params['rec_noise'] = 0.0 # Default: 0.0 -- how much noise to add to the new_state calculation
params['W_in_train'] = True # Indicates whether W_in is trainable. Default: True
params['W_rec_train'] = True # Indicates whether W_rec is trainable. Default: True
params['W_out_train'] = True # Indicates whether W_out is trainable. Default: True
params['b_rec_train'] = True # Indicates whether b_rec is trainable. Default: True
params['b_out_train'] = True # Indicates whether b_out is trainable. Default: True
params['init_state_train'] = True # Indicates whether init_state is trainable. Default: True
params['load_weights_path'] = None # When given a path, loads weights from file in that path. Default: None
params['L2_out'] = 0<jupyter_output><empty_output><jupyter_text>Instantiate our model<jupyter_code>basicModel.destruct()
basicModel = Basic(params)<jupyter_output><empty_output><jupyter_text>Start a tensorflow session with loss, regularization, predictions, and regularized loss defined.<jupyter_code>basicModel.build()<jupyter_output><empty_output><jupyter_text>Set the training parameters for our model. All of the parameters below are optional.<jupyter_code>train_params = {}
train_params['save_weights_path'] =  '../weights/refactor_weights.npz' # Where to save the model after training. Default: None
train_params['training_iters'] = 10000 # number of iterations to train for Default: 10000
train_params['learning_rate'] = .001 # Sets learning rate if use default optimizer Default: .001
train_params['loss_epoch'] = 10 # Compute and record loss every 'loss_epoch' epochs. Default: 10
train_params['verbosity'] = True # If true, prints information as training progresses. Default: True
train_params['save_training_weights_epoch'] = 100 # save training weights every 'save_training_weights_epoch' epochs. Default: 100
train_params['training_weights_path'] = None # where to save training weights as training progresses. Default: None
train_params['generator_function'] = None # replaces trial_batch_generator with the generator_function when not none. Default: None
train_params['optimizer'] = tf.train.AdamOptimizer(learning_rate=train_params['learning_rate']) # What optimizer to use to compute gradients. Default: tf.train.AdamOptimizer(learning_rate=train_params['learning_rate'])
train_params['clip_grads'] = True # If true, clip gradients by norm 1. Default: True
train_params<jupyter_output><empty_output><jupyter_text>Train our model.<jupyter_code>basicModel.train(gen, train_params)<jupyter_output>Iter 500, Minibatch Loss= 101.306046
Iter 1000, Minibatch Loss= 87.013535
Iter 1500, Minibatch Loss= 26.358490
Iter 2000, Minibatch Loss= 17.759918
Iter 2500, Minibatch Loss= 10.280129
Iter 3000, Minibatch Loss= 9.679141
Iter 3500, Minibatch Loss= 6.142443
Iter 4000, Minibatch Loss= 5.710996
Iter 4500, Minibatch Loss= 5.383883
Iter 5000, Minibatch Loss= 3.765890
Training weights saved in file: None100
Iter 5500, Minibatch Loss= 3.252282
Iter 6000, Minibatch Loss= 4.421055
Iter 6500, Minibatch Loss= 2.657529
Iter 7000, Minibatch Loss= 2.402243
Iter 7500, Minibatch Loss= 3.946777
Iter 8000, Minibatch Loss= 2.094603
Iter 8500, Minibatch Loss= 2.025302
Iter 9000, Minibatch Loss= 2.133550
Iter 9500, Minibatch Loss= 1.691504
Optimization finished!
Model saved in file: ../weights/refactor_weights.npz<jupyter_text>Plot results Extract the next trial from the generator Run the trained model on this trial (not included in the training set).<jupyter_code>x,y,m = next(gen)
plt.plot(range(0, len(x[0,:,:])*dt,dt), x[0,:,:])
plt.plot(range(0, len(y[0,:,:])*dt,dt), y[0,:,:])
results = basicModel.test(x)
output = results[0]
state_var = results[1]
plt.plot(range(0, len(x[0,:,:])*dt,dt),x[0,:,:])
plt.plot(range(0, len(output[0,:,:])*dt,dt),output[0,:,:])
plt.ylabel("Activity of Output Unit")
plt.xlabel("Time (ms)")
plt.title("Output on New Sample")
weights = basicModel.get_weights()
E = np.diag(weights['Dale_rec'])>0
I = np.diag(weights['Dale_rec'])<0

f,ax = plt.subplots(3,1,figsize = (10,10))

ax[0].plot(range(0, len(state_var[0,:,:])*dt,dt),state_var[0,:,:])
plt.ylabel("State Variable Value")
plt.xlabel("Time (ms)")
plt.title("Evolution of State Variables over Time")

for ind in np.where(E)[0]:
    ax[1].plot(range(0, len(state_var[0,:,ind])*dt,dt),state_var[0,:,ind],'r',alpha = 0.3)
plt.ylabel("State Variable Value")
plt.xlabel("Time (ms)")

for ind in np.where(I)[0]:
    ax[2].plot(range(0, len(state_var[0,:,ind])*dt,dt),state_var[0,:,ind],'b',alpha = 0.3)
plt.ylabel("State Variable Value")
plt.xlabel("Time (ms)")
plt.tight_layout()
peak_times_first_cycle = np.argmax(state_var[0,:50,:],axis =0)
state_var.shape
peak_times_first_cycle
np.sort(peak_times_first_cycle)
peak_order = np.argsort(peak_times_first_cycle)
peak_order
sort_state = np.zeros(state_var.shape)
for trial in range(state_var.shape[0]):
    for i,this_neuron in enumerate(peak_order):
        sort_state[trial,:,i] = state_var[trial,:,this_neuron]
plt.imshow(sort_state[1,:,:].T,aspect = 'auto')
plt.imshow(state_var[0,:,peak_order],aspect = 'auto')
plt.plot(range(0, len(x[0,:,:])*dt,dt), x[0,:,:])
plt.ylabel("Input Magnitude")
plt.xlabel("Time (ms)")
plt.title("Input Data")
plt.plot(range(0, len(y[0,:,:])*dt,dt), y[0,:,:])
weights = basicModel.get_weights()
E = np.diag(weights['Dale_rec'])>0
I = np.diag(weights['Dale_rec'])<0
weights['input_connectivity'].shape
weights.keys()
toplot = ['W_in', 'W_rec', 
 'W_out','Dale_rec', 
 'Dale_out', 'input_connectivity', 
 'rec_connectivity', 'output_connectivity']
for key in toplot:
    plt.figure()
    plt.imshow(weights[key])
    plt.title(key)
N_rec = 10
dale_ratio = 0.8
dale_vec = np.ones(N_rec)
        
dale_vec[int(dale_ratio * N_rec):] = -1
dale_rec = np.diag(dale_vec)
dale_vec[int(dale_ratio * N_rec):] = 0
dale_out = np.diag(dale_vec)
dale_rec<jupyter_output><empty_output><jupyter_text>Clean up the model to clear out the tensorflow namespace<jupyter_code>basicModel.destruct()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19300381462543625, 'fraction_numerical': 0.03425046668289911, 'mean_word_length': 2.7785955228457526, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '122996', 'license': 'MIT', 'n_tokens_mistral': 4459, 'n_tokens_neox': 4044, 'n_words': 992, 'path': 'noisy_decoding_task.ipynb', 'repo_name': 'ANaka/PsychRNN', 'size': '12321'}