I maintain some code and I run across something like:
travel_time_vec = np.zeros(...)
for v in some_indexes: # some_indexes is a list of row indexes
traveltimes = traveltime_2d_array[v, list_of_column_indexes]
best_index = np.argmin(traveltimes)
travel_time_vec[v] = traveltimes[best_index]
I would like to drop the for loop and do all the operations below at once - but naively asking for traveltime_2d_array[some_indexes, list_of_column_indexes]
results in:
{IndexError}shape mismatch: indexing arrays could not be broadcast together with shapes (4,) (8,)
Got it - I need to pass the some_indexes
as a list of lists so numpy broadcasts each one to the columns in list_of_column_indexes
. So this:
travel_time_vec = np.zeros(...)
# newaxis below tranforms [1, 2, 3] to [[1], [2], [3]]
traveltimes = traveltime_2d_array[np.array(some_indexes)[:, np.newaxis],
list_of_column_indexes]
# get the index of the min time on each row
best_index = np.argmin(traveltimes, axis=1)
travel_time_vec[some_indexes] = traveltimes[:, best_index]
works as expected, and no more looping