1

I am attempting to extract out and perform math on the subset of one numpy array that is of shape (3,32), the subset of data I am attempting to extract out is that of shape (3,9) and the ranges of this data originate at the indices contained in another array size (3). As an example, I have a data set of values from three channels operating in the time domain and I extract the index of the max value of each channel into an array

a = np.random.randint(20,size = (3,32))
a
array([[18,  3, 10,  6, 12,  1, 10,  8,  4, 11, 13, 14,  9,  9, 10,  2,
         9,  0,  0, 16, 14, 19,  1, 19, 14, 19, 19,  2, 14,  0,  4, 18],
       [ 9, 19,  2, 12,  0, 14, 18,  7,  3,  0,  7,  3, 12, 19,  4,  2,
         5,  9,  2, 11, 15, 19, 16, 17,  3,  4, 17,  5,  6,  1,  2, 17],
       [ 0, 11, 18,  8,  9,  2,  9, 15,  9,  6,  0,  8,  9, 16,  9,  6,
         1, 19,  1,  9, 12,  8,  0,  0,  7, 15,  3, 14, 15,  8, 10, 19]])
b = np.argmax(a,1)
b
array([21,  1, 17], dtype=int64)

my goal at this point is to derive a new array consisting of the three values from each of the indexes specified. For instance I would be looking to extract out :

[21,22,23] from a[0]
[1,2,3] from a[1]
[17,18,19] from a[2]

all into a new array of size [3,3]

I've been able to achieve this using loops already but I suspect that there is a more efficient way of producing this result without loops (speed is a bit of an issue with this application). I have been able to effect a similar result by manually populating a smaller matrix manually ...

c = np.asarray([[1,2,3],[2,3,4],[3,4,5]])
a[np.arange(3)[:,None],c]
array([[ 3, 10,  6],
       [ 2, 12,  0],
       [ 8,  9,  2]])

However, given the dynamic nature of this application I would like to write this such that it can be dynamically scaled (range of indeces out to 9 values beyond the root index, etc). I just don't know if there is such a way to do this. I have used syntax similar to the following in an effort to slice the array ...

a[np.arange(3)[:,None],b[:]:(b[:] + 2)]

resulting in error messages in the nature of ...

builtins.TypeError: only integer scalar arrays can be converted to a scalar index
Matthew Connolly
  • 45
  • 5
  • What if there aren't 3 values remaining in the row in `a`? For example, what if the max was at index `31`? – user3483203 May 28 '19 at 22:18
  • Yeah I considered that I should stipulate that, the range will never exceed the bounds of the array given the nature of the application (in the actual application the max indices always are 10 +/- 1) – Matthew Connolly May 28 '19 at 22:24
  • In the accepted solution, use `b` as `idx`. – Divakar May 29 '19 at 04:28

1 Answers1

1

Since you say you cannot overflow, this gets much less tricky. In general, since you have your starting indices, you can use basic broadcasting to create an (n, 3) shape array with your indices, and use take_along_axis to pull those elements from the original array.

np.take_along_axis(a, b[:, None] + np.arange(3), axis=1)


array([[19,  1, 19],
       [19,  2, 12],
       [19,  1,  9]])
user3483203
  • 50,081
  • 9
  • 65
  • 94