Multidimensional Indexing of an multidemensional matrix in Python

Question

I want to use a multidimensional index matrix to access another multidimensional matrix. My problem is, that methods such as np.newaxis are not working because of broadcasting issues (mismatch of shapes).

My data matrix has a shape of (5001, 3, 240, 16).

import numpy as np

# n_examples, n_channels, n_pictures, n_meta_information
data = np.ones((5001, 3, 240, 16))

# select randomly 32 examples
batch_size = 32
possible_indices = np.arange(5001, dtype=np.int)
random_example_indices = np.random.choice(possible_indices, size=batch_size, replace=True)


# select all three channels
n_channels = 3
channel_indices = np.arange(n_channels)
#channel_indices = np.expand_dims(channel_indices , axis=0)
#channel_indices = np.repeat(channel_indices, batch_size, axis=0)

final_pictures_indices = []
for random_sample_idx in range(batch_size):
    # select a random start index and take 120 successive indices
    # is the same for all three channels
    start_index = np.random.randint(0, max(1, 240 - 120 + 1))
    end_index = start_index + 120
    pictures_indices = np.arange(start_index , end_index , dtype=np.int)
    final_pictures_indices.append(pictures_indices)

# batch_size x n_pictures
final_pictures_indices = np.array(final_pictures_indices)


# should have the shape: (32, 3, 120, 16)
result = data[random_example_indices[:, np.newaxis], channel_indices, final_pictures_indices].shape
print(result)

Unfortunately I got the following error:

result = data[random_example_indices[:, np.newaxis], channel_indices, final_pictures_indices].shape
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (32,1) (3,) (32,120)

I have also tried to fuse all the index information into a single matrix but there I had the problem that stacking matrices with different shapes is not possible.

Thank you for every tip.

Answer 1

your indexing arrays need to be able to broadcast against each other, i.e. from the right, have either your expected final dimension, 1, or not exist. So:

result = data[random_example_indices[:, np.newaxis, np.newaxis],   # (32,   1,    1)
              channel_indices[:, np.newaxis],                      #       (3,    1)
              final_pictures_indices[:, np.newaxis, :].shape       # (32,   1,  120)

That should get you the expected shape.