У меня есть массив 6D NumPy и у меня есть пользовательская функция, которая может быть применена к 2D-массиву, на выходе получается 2D-массив той же формы.Я хотел бы перебрать все двумерные массивы, которые хранятся в индексах 3 и 4, применить пользовательскую функцию и вернуть массив 6D с формой исходного массива 6D и с данными, обработанными, как описано выше, с помощью пользовательской функции.В настоящее время я использую вложенные циклы, но процесс очень медленный.Ниже я делаю то, что я сейчас делаю (размер измерений 2 и 5 равен 1, поэтому не нужно их перебирать).
orig_array = np.reshape(np.zeros(2**4), (2, 2, 1, 2, 2, 1))
def custom_fun(array_2d):
return array_2d + 1
out_array = np.copy(orig_array)
for ind1 in range(orig_array.shape[0]):
for ind2 in range(orig_array.shape[1]):
out_array[ind1, ind2, 0, :, :, 0] = custom_fun(out_array[ind1, ind2, 0, :, :, 0])
Мне было интересно, есть ли способ векторизовать эту процедуру с помощьюцель ускорения вычислений.Я искал документацию NumPy и нашел функцию: np.apply_over_axes, однако, похоже, она не выполняет то, что мне хотелось бы.