Возвращение списка массивов из функции, имеющей в качестве аргумента вектор

Question

У меня есть такая функция:

def f(x):
    A =np.array([[0, 1],[0, -1/x]]);
    return A

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

>>x=1
>>f(x)
array([[ 0.,  1.],
       [ 0., -1.]])

, и если я буду использовать массив в качестве входных данных, я получу :

>>x=np.linspace(1,3,3)
>>f(x)
array([[0, 1],
       [0, array([-1.        , -0.5       , -0.33333333])]], dtype=object)

На самом деле я хотел бы получить список массивов, а именно: A = [A_1, A_2, ..., A_n]

Сейчас меня не волнует если это массив массивов или список, содержащий несколько массивов. Я знаю, что могу сделать это, используя для l oop в х. Но я думаю, что, возможно, есть другой способ сделать это, и, возможно, более эффективный.

Поэтому вывод, который я хотел бы, был бы примерно таким:

>>x=np.linspace(1,3,3)
>>r=f(x)
array([[[0, 1],[0,-1]], 
       [[0, 1],[0,-0.5]],
       [[0, 1],[0,-0.33333]]])
>>r[0]
array([[0, 1],[0,-1]])

или что-то вроде

>>x=np.linspace(1,3,3)
>>r=f(x)
[array([[0, 1],[0,-1]]), 
 array([[0, 1],[0,-0.5]]),
 array([[0, 1],[0,-0.33333]])]
>>r[0]
array([[0, 1],[0,-1]])

Спасибо

Answer 1

В вашей функции мы можем проверить тип данного параметра. Здесь, если x является типом np.ndarray, мы собираемся создать вложенный список, который нам нужен, в противном случае мы вернем вывод, как и раньше.

import numpy as np

def f(x):
    if isinstance(x, np.ndarray):
        v = -1/x
        A = np.array([[[0, 1],[0, i]] for i in v])
    else:
        A = np.array([[0, 1],[0, -1/x]])

    return A

x = np.linspace(1,3,3)
print(f(x))

Вывод:

[[[ 0.          1.        ]
  [ 0.         -1.        ]]

 [[ 0.          1.        ]
  [ 0.         -0.5       ]]

 [[ 0.          1.        ]
  [ 0.         -0.33333333]]]

Answer 2

Вы можете сделать что-то вроде:

import numpy as np

def f(x):
    x = np.array([x]) if type(x)==float or type(x)==int else x
    A = np.stack([np.array([[0, 1],[0, -1/i]]) for i in x]);
    return A

В первой строке рассматриваются случаи, когда x является int или float, поскольку не является итеративным. Тогда:

f(1)
array([[[ 0.,  1.],
        [ 0., -1.]]])

f(np.linspace(1,3,3))
array([[[ 0.        ,  1.        ],
        [ 0.        , -1.        ]],

       [[ 0.        ,  1.        ],
        [ 0.        , -0.5       ]],

       [[ 0.        ,  1.        ],
        [ 0.        , -0.33333333]]])