I am trying to extract part of an array from an array.
Let's assume I have an array array1
with shape (M, N, P)
. For my specific case, M = 10
, N = 5
, P = 2000
. I have another array, array2
of shape (M, N, 1)
, which contains the starting points of the interesting data in array1
along the last axis. I want to extract 50 points of this data starting with the indices given by array2
, kind of like this:
array1[:, :, array2:array2 + 50]
I would expect a result of shape (M, N, 50)
. Unfortunatly I get the Error:
TypeError: only integer scalar arrays can be converted to a scalar index
Sure I could also get the result by looping through the array, but I feel that there must be a smarter way, because I needed this quite often.
You can build a mask using a comparison of the values in array2 with an index range of the last dimension:
For example:
import numpy as np
M,N,P,k = 4,2,15,3 # yours would be 10,5,2000,50
A1 = np.arange(M*N*P).reshape((M,N,P))
A2 = np.arange(M*N).reshape((M,N,1)) + 1
rP = np.arange(P)[None,None,:]
A3 = A1[(rP>=A2)&(rP<A2+k)].reshape((M,N,k))
Input:
print(A1)
[[[ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14]
[ 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29]]
[[ 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44]
[ 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59]]
[[ 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74]
[ 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89]]
[[ 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104]
[105 106 107 108 109 110 111 112 113 114 115 116 117 118 119]]]
print(A2)
[[[1]
[2]]
[[3]
[4]]
[[5]
[6]]
[[7]
[8]]]
Output:
print(A3)
[[[ 1 2 3]
[ 17 18 19]]
[[ 33 34 35]
[ 49 50 51]]
[[ 65 66 67]
[ 81 82 83]]
[[ 97 98 99]
[113 114 115]]]