Using PyTensor.scan to loop over multidimensional array

Questions v5
8

I tried your suggestion, but the result has 4 dimensions (an additional one for i).
Based on your previous suggestions and this example I used the following:

def step(A_diff, A, b, c):
    b_func = k1*A*(A+2*c)
    return A- (A_diff + b_func)

A_sim = A_array[:,:,0]

result, updates = pytensor.scan(fn=step,
                                        outputs_info=A_sim.T,
                                        sequences=A_diff.T,
                                        non_sequences=[b, c])

However, I noticed that result is then missing the first element of A_sim, so I have to create the final array via:

A_new = at.zeros(shape=(at.shape(A_array).eval()))
A_new = at.set_subtensor(A_new [:,:,0], A_array[:,:,0])
A_new = at.set_subtensor(A_new [:,:,1:], result.T)

Maybe there is a smarter way for this last step, but it runs as supposed now and is much faster! It now takes about 0.2 seconds for a simulation.
Thank you very much for your help @jessegrabowski.

1 Like