2

I'm trying to take certain data out of a tensor, but I get weird error. Here, I will try to generate the error:

a=np.random.randn(5, 10, 5, 5)
a[:, [1, 6], np.triu_indices(5, 0)[0], np.triu_indices(5, 0)[1]].shape

I get this error

shape mismatch: indexing arrays could not be broadcast together with shapes

I'm not even doing any broadcasting! it is all slicing stuff.

What do I want? leave zeroth axis as is (get everything), get [1] and [6] from first axis, reshape the last two axes from [5, 5] to [15] by taking upper triangle elements only.

Alex
  • 599
  • 1
  • 5
  • 14
  • Only the 1st index is a slice, the rest are arrays that are broadcasted together to create a composite index on 3 dimensions. – hpaulj Aug 02 '19 at 04:32
  • So whats the intended result shape, (5,2,15)? – hpaulj Aug 02 '19 at 04:35
  • @hpaulj Exactly! that's the required shape. I managed to get it but in two stages. I failed to make it in one line – Alex Aug 02 '19 at 04:40
  • Show the 2 steps. It may be simpler. See https://stackoverflow.com/q/57273662/901925 and its duplicate. – hpaulj Aug 02 '19 at 05:53
  • @hpaulj step one: tmp = a[:, :, np.triu_indices(5, 0)[0], np.triu_indices(5, 0)[1]] Step two: a= tmp[:, [1, 6], ...] so basically I avoided mixed indexing – Alex Aug 02 '19 at 08:00

1 Answers1

2

We need to extend the second axis indexing array to 2D, so that it forms an outer-plane against the indices off np.triu_indices. Thus, it give us a 2D grid of mxn array with m being the length of that second axis indexing array and n being the lengths of the np.triu_indices ones. So, essentially, the entire solution would simplify to something like this -

r,c = np.triu_indices(5, 0)
out = a[:, np.array([1, 6])[:,None], r, c]

Or feed in that extended version as a list, i.e. -

out = a[:, [[1],[6]], r, c]

We can also use a masking based one with np.tri/np.triu, which might be faster on larger arrays as we would skip creating all the integer indices, like so -

mask = ~np.tri(5, k=-1, dtype=bool)
out = a[:, np.array([1, 6])[:,None], mask]
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • I'm trying to understand this! so the problem is mixed indexing methods. If I do it in two steps, it works. step one: treat last two axes, get a new matrix, step two: treat axes number 1 simple by a[:, [1, 6], ...] Can you expand on what's outer-plane? – Alex Aug 02 '19 at 08:03
  • @Alex Think of those as two vectors (in matrix notation) that are "outer" with respect to each other. So, `np.array([1, 6])[:,None]` is a column vector and r,c are row vectors. So, when operated together, they form an outer-operation obliging the numpy-broadcasting to extract 2D blocks. You gotta read up on numpy.ix_ and numpy-broadcasting docs and https://stackoverflow.com/a/35608029/ for more info. – Divakar Aug 02 '19 at 08:08