CODEWITHSUNDEEP
pythonnumpy
varantir
varantir

Reputation: 6854

choose element of ndarray for each element in index array

I am pretty sure this has already been asked and it has to do with numpy.choose, but I cannot figure out how it is accomplished. Consider the following:

N_t = 1000
N_d = 10

X = np.random.random([N_d,N_d,N_t])

jumps = np.random.randint(N_d,size = N_t)
jumps[0] = 0

f = [X[jumps[t],jumps[t-1],t] for t in range(1,N_t)]

Is there a "numpy"-way of constructing f? (Side remark: X is some kind of transition matrix and jumps are the indices of some jump-trajectory in a d-dim state space.

Upvotes: 0

Views: 35

Answers (1)

akuiper
akuiper

Reputation: 215137

Yes. It's called advanced indexing; You do this by constructing indices for each dimension as integer arrays:

X[jumps[1:], jumps[:-1], np.arange(1,N_t)]

np.equal(X[jumps[1:],jumps[:-1],np.arange(1,N_t)], f).all()
# True

Upvotes: 1

Related Questions