CODEWITHSUNDEEP
pythonarraysnumpy
Daniel
Daniel

Reputation: 27

Returning list of arrays from a function having as argument a vector

I have a function such as:

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

If I use an scalar I will obtain:

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

and if I use an array as an input, I will obtain:

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

Actually I would like to obtain a list of array, namely: A = [A_1,A_2, ..., A_n]

Right now I do not care much about if it is an array of arrays or a list that contain several arrays. I know I can do that using a for loop in x. But I think there is probably another way to do it, and maybe more efficient.

So the output that I would like would be something like:

>>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]])

or something like

>>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]])

Thanks

Upvotes: 0

Views: 51

Answers (2)

FBruzzesi
FBruzzesi

Reputation: 6485

You can do something like:

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

The first line deal with the cases when x is an int or a float, since is not an iterable. Then:

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]]])

Upvotes: 1

Filip Młynarski
Filip Młynarski

Reputation: 3612

In your function we could check type of given parameter. Here, if x is type of np.ndarray we are going to create nested list which we desire, otherwise we'll return output as before.

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))

Output:

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

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

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

Upvotes: 2

Related Questions