Операция с массивом Numpy на основе пользовательского ввода

Question

У меня есть двумерный массив. Скажем:

abc = np.arange(40).reshape(10,4)

Теперь я принимаю строковую переменную пользовательского ввода 'mm'. Эта переменная определяет, что я хочу сделать с массивом. Имеет следующие 4 возможных значения:

'min': возвращает столбец со всеми минимальными значениями

'max': возвращает столбец со всеми максимальными значениями

«среднее»: возвращает столбец со всеми средними значениями

'time': возвращает столбец со значением в определенном столбце. Это также нуждается в дополнительном входе от пользователя; переменная 'idx' с номером столбца.

Я включил первые три возможности, определив словарную переменную, которая выбирает соответствующую функцию.

dict = {'mean':np.mean, 'max':np.max, 'min':np.min}
fn = dict[mm]

Поэтому моя выходная переменная будет:

op = fn(abc,axis=1)

Но я застреваю, когда пытаюсь включить четвертую возможность нарезки определенного столбца на основе пользовательского ввода. Можно ли определить любую словесную функцию в словаре dict, которая также может включать эту операцию?

Я бы предпочел не использовать условие if-else, так как это делает код слишком длинным. Фактические переменные намного больше и сложнее в обращении.

Ожидаемый результат, когда пользовательская переменная 'mm' имеет разные значения:

op_min = np.min(abc,axis=1)
op_max = np.max(abc,axis=1)
op_mean = np.mean(abc,axis=1)
op_time = abc[:,2]  # assuming the user input for idx is 2

Answer 1

Если вы можете транспонировать матрицу, эта функция может выбрать строку:

abc = np.arange(40).reshape(10,4)
row_num = 2
f = operator.itemgetter(row_num)
f(abc)

Также вы можете определить одну строковую функцию по лямбде для выбора столбца:

f = lambda x: x[:,column_num]
f(abc)

Answer 2

Хорошо, я нашел решение.Я транспонировал матрицу, как предложено пользователем 3371603, и использовал функцию operator.getitem.Затем я использовал индексную переменную (idx) со значением по умолчанию 0. Итак,

import operator
dict = {'mean':np.mean, 'max':np.max, 'min':np.min, 'time':operator.getitem}
fn = dict[mm]

idx = 0
if mm == 'time': idx=2     # or user input

cba = np.swapaxes(abc,0,1)
op = fn(cba, idx)

Отдельное спасибо user3371603!