Equivalence of numpy.delete() in Asera/Theano?

Questions
1

Hi,
Optimizing my aversive learning script (see here: aversive_learning_simulation/simulating_scrData.ipynb at main · orduek/aversive_learning_simulation · GitHub), I need to perform the following:

  1. calculate the expected value for each trial (per subject).
  2. Remove all values for the shock trials
  3. fit to skin conductance response.

I’m stuck at 2. For example, if I were to use NumPy, I can simply get a vector of indexes for the times’ shock was administered, then use numpy.delete to remove those indexes from the vector.
I can’t find a good way to do it in Theano/Asera.

For convenience, I attach the main portion of the code here:

def update_Q(stim, shock,
             Qs,vec,
             alpha, n_subj):
    """
    This function updates the Q table according to the RL update rule.
    It will be called by theano.scan to do so recursevely, given the observed data and the alpha parameter
    This could have been replaced be the following lamba expression in the theano.scan fn argument:
        fn=lamba action, reward, Qs, alpha: tt.set_subtensor(Qs[action], Qs[action] + alpha * (reward - Qs[action]))
    """
      
    PE = shock - Qs[tt.arange(n_subj), stim]
    Qs = tt.set_subtensor(Qs[tt.arange(n_subj),stim], Qs[tt.arange(n_subj),stim] + alpha * PE)
    
    # in order to get a vector of expected outcome (dependent on the stimulus presentes [CS+, CS-] 
    # we us if statement (switch in theano)
    vec = tt.set_subtensor(vec[tt.arange(n_subj),0], (tt.switch(tt.eq(stim,1), 
                                                                Qs[tt.arange(n_subj),1], Qs[tt.arange(n_subj),0])))
    
   
    return Qs, vec, PE


with pm.Model() as m:
  
      
    alpha = pm.Beta('alpha', 1,1, shape=n_subj)
    beta = pm.Normal('beta',0, 1, shape=n_subj)
    eps = pm.HalfNormal('eps', 5)
    
    Qs = 0.5 * tt.ones((n_subj,2), dtype='float64') # set values for boths stimuli (CS+, CS-)
    vec = 0.5 * tt.ones((n_subj,1), dtype='float64') # vector to save the relevant stimulus's expactation
    
    [Qs,vec, pe], updates = theano.scan(
        fn=update_Q,
        sequences=[stim, shock],
        outputs_info=[Qs, vec, None],
        non_sequences=[alpha, n_subj])
    
    # removing SCR with shocks before fitting
    # transform back to vector
    vec_ = tt.reshape(vec.T, n_subj*n_trials, ndim=1)

    scrV = np.reshape(scrMat.T, n_subj*n_trials)
    # clean scr from shocks
    scrClean = np.delete(scrVec, shIn)
    vec_clean = np.delete(vec_, shIn)
    vec_t = vec_clean * beta[subj]
    
    
    scrs = pm.Normal('scrs', mu = vec_t, sd = eps, observed=scrClean)

Any ideas?

2

I tried running your code to have a look, but it seems there are some missing values (what is scrMat?)

Anyway my first instinct is to just use a boolean mask to select the elements of scrVec that you want to keep. Since (I believe?) shockVec has a one everywhere the trial had a shock, you could do something like: vec_clean = scrVec[shockVec == 0], which would return only elements that correspond to a non-shock trial.

2 Likes
3

Thanks!
Indeed, the code here is partial, so won’t run as is.
Your solution works great! so simple :slight_smile:

1 Like