np.apply_along_axis возвращает несколько значений

Question

Я применяю некоторые функции к массиву numpy.

import numpy as np
MATRIX=np.random.rand(5,7)


def F_a(x):
   return 1

def F_b(x):
   return [1,2]
FUNCS=[F_a,F_b]

results=list(map(lambda f:np.apply_along_axis(f,1, MATRIX),FUNCS))

когда все FUNCS возвращают однокалярное значение, тогда оно работает как положено. Однако некоторые моих итераций возврата FUNC (например, списки переменной длины).

Как применить функции, возвращающие несколько значений таким образом (map + apply_along_axis)?

Если я запускаю MVE, я получаю:

[array([1, 1, 1, 1, 1]), array([[1, 2],
        [1, 2],
        [1, 2],
        [1, 2],
        [1, 2]])]

где в качестве желаемого результата будет один массив размером 5x3 вместо двух массивов 5x1 и 5x2

Если я попытаюсь:

results=np.stack(list(map(lambda f:np.apply_along_axis(f,1, MATRIX),FUNCS)))

Я получаю:

ValueError: all input arrays must have the same shape

То же самое с np.vstack, np.hstack

Answer 1

С осью 1 и входом 2d apply_along_axis просто перебирает строки (другие измерения) и возвращает массив, форма которого определяется измерением итерации, и значение функции.

Если fn возвращает скаляр, результат равен (5,), размер строк:

In [291]: np.apply_along_axis(lambda a1: 1, 1,arr)                                                     
Out[291]: array([1, 1, 1, 1, 1])
In [292]: _.shape                                                                                      
Out[292]: (5,)

если он возвращает список или массив из одного элемента, это измерение имеет размер 1:

In [293]: np.apply_along_axis(lambda a1: [1], 1,arr)                                                   
Out[293]: 
array([[1],
       [1],
       [1],
       [1],
       [1]])
In [294]: _.shape                                                                                      
Out[294]: (5, 1)

И так далее для других возвратов:

In [295]: np.apply_along_axis(lambda a1: np.arange(4), 1,arr)                                          
Out[295]: 
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])
In [296]: _.shape                                                                                      
Out[296]: (5, 4)
In [297]: np.apply_along_axis(lambda a1: a1, 1,arr)                                                    
Out[297]: 
array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]])
In [298]: _.shape                                                                                      
Out[298]: (5, 3)

Это то же самое, что создать массив из списка:

In [302]: [(lambda a1:1)(row) for row in arr]                                                          
Out[302]: [1, 1, 1, 1, 1]
In [303]: [(lambda a1:[1,2])(row) for row in arr]                                                      
Out[303]: [[1, 2], [1, 2], [1, 2], [1, 2], [1, 2]]

С такими 2-мерными входами apply_along_axis - это немного больше, чем удобство записи, если это так. Для большего размера это удобство больше. Но это не улучшает скорость по сравнению с более прямой итерацией.

Если вы не понимаете apply_along_axis или у вас есть проблемы с его использованием, не беспокойтесь. Выполните ту итерацию, которую вы понимаете.

====

Если

def F_a(x):
    return [1]

ваша карта списка создаст (5,1) и (5,2) массивы, которые затем можно будет объединить по 2-й оси.

Кроме того, кроме проблем с apply, похоже, вам не нравятся размеры numpy массива и то, как они контролируют использование concatenate (и различных вариантов stack).