dontknowwhatimdoing
dontknowwhatimdoing

Reputation: 13

Accessing array elements within a vectorized function

I am trying to write a function which computes the Simpson integral for given inputs. Parts of these inputs can be in the form of a float or np.array.

Currently, my function is as follows:

s = np.array([0.2,0.3,0.4])

@np.vectorize
def recovery(s,R0,n):

    x = np.linspace(s,1,n)
    y = 1/(x*R0(x))
    ans = simp(y,x)

    return ans

This works fine and returns:

array([a, b, c]) ## Where a, b and c are the correct answers

However if I only want to access the first element in the array within the function, by:

return ans[0]

I get the error message:

IndexError: invalid index to scalar variable

I also get all different kinds of errors when trying to play around with the output "ans", e.g.:

return type(ans)
ValueError: invalid __array_struct__

return len(ans)
TypeError: object of type 'numpy.float64' has no len()

The aim of the function is to return a float if the input s is a float, or an array if the s input is an np.ndarray.

I'm a complete python beginner so any help would be greatly appreciated.

Upvotes: 1

Views: 134

Answers (1)

macroeconomist
macroeconomist

Reputation: 701

To access only the first element of the return array, you need to do so outside the function, not inside, because inside the function ans is still a scalar.

So while

return ans[0]

doesn't work,

recovery(s,R0,n)[0]

or

ans = recovery(s,R0,n)
ans[0]

will work.

The confusing thing here is that the decorator @np.vectorize changes the function from one that maps scalars to scalars to one that maps vectors to vectors. But this manipulation is being done "outside" the function recovery; inside the original source code of the function, everything is as if it's still a function on scalars.

That's why when you wrote return ans[0], Python complained that you were trying to index a scalar; inside the function, ans is a scalar.

Upvotes: 1

Related Questions