0

Basically is there a Numpy or PyTorch function that does this:

vp_sa_s=mdp_data['sa_s'].detach().clone()
dims = vp_sa_s.size()
for i in range(dims[0]):
    for j in range(dims[1]):
        for k in range(dims[2]):
            # to mimic matlab functionality: vp(mdp_data.sa_s)
            try:
                vp_sa_s[i,j,k] = vp[mdp_data['sa_s'][i,j,k]]
            except:
                pass

Given that vp_sa_s is size (10,5,5) and each value is a valid index vp i.e in range 0-9. vp is size (10,1) with a bunch of random values.

Matlab do it elegantly and quickly with vp(mdp_data.sa_s) which will form a new (10,5,5) matrix. If all values in mdp_data.sa_s are 1, the result would be a (10,5,5) tensor with each value being the 1st value in vp.

Does a function or method that exists that can achieve this in less than O(N^3) time as the above code is terribly inefficient.

Thanks!

Antoine Delia
  • 1,728
  • 5
  • 26
  • 42

1 Answers1

1

What is wrong with

result = vp[vp_sa_s, 0]

note that since your vp is of shape (10, 1) (it has a trailing singleton dimension) you need to add the , 0] index in the assignment to get rid of this extra dimension.

Shai
  • 111,146
  • 38
  • 238
  • 371
  • Man I feel like an idiot. Thank you so much! I was overcomplicating it so much. Would you be able to explain how Numpy interprets that as "fill a result array size == vp_sa_s with values result[I] == vp[vp_sa_s[I]]. As my understanding is that vp[vp_sa_s] would be interpreted as Index VP with the dimensions stored in vp_sa_s? –  Jan 26 '21 at 09:51
  • 1
    @joekadi [Fancy Indexing](https://numpy.org/doc/stable/reference/arrays.indexing.html) in `numpy` is designed specifically to mirror `MATLAB` functionality in this case. – Daniel F Jan 26 '21 at 09:53
  • @joekadi you should not feel bad about your question. – Shai Jan 26 '21 at 09:55
  • 1
    [Here](https://stackoverflow.com/questions/44574086/how-is-numpys-fancy-indexing-implemented) is some more information about how it's implemented if you're interested. – Daniel F Jan 26 '21 at 09:55