Если я правильно понимаю, вы ищете это без необходимости в циклах:
idx_max = np.unravel_index(f_5.argmax(), f_5.shape)
idx = idx_max[2]
f_max = f_5[idx_max]
или, что эквивалентно, менее эффективный подход (поскольку он ищет max, а также ищет индекс отдельно):
f_max = np.amax(f_5)
idx = np.argwhere(f_5==f_max)[0,2]