цикл для итерации столбцов массива (с размерностью 1D или 2D)

Question

Я пытаюсь написать функцию, которая получает массивный массив INPUT и передает его столбцы одну за другой функции.массив INPUT является 1D или 2D (не более), вторая функция ожидает 1D массивы в качестве параметра.(len (param.shape) == 1)

Я читал аналогичную ветку, в которой ОП хотел суммировать все столбцы и проверить другие условия ... для этого, вероятно, нужен другой ответ.

желаемая операция в псевдокоде:

def func(INPUT,a,b,...)
    for column in INPUT: #whether be a 1D or 2D
        result = another_func(column,...)

попробовал это: вопрос в том, как не проверять размерность массива INPUT внутри func:

if(len(INPUT.shape)==1):
    another_func(INPUT,....)
elif(len(INPUT.shape)==2):
    for c in range(INPUT.shape[1]):
        another_func(INPUT[:,c])

Answer 1

Идея такова: в случае ввода 1d преобразуйте в массив 2d с 1 столбцом, а затем обработайте как ввод 2d.

def func(INPUT, a, b):
    return np.apply_along_axis(
        lambda col: another_func(col, a, b),  # function to apply
        1,  # axis along which to apply; 1 = columns
        np.reshape(np.atleast_2d(H2), (len(H), -1))  # transform 1D->2D, if necessary
    )