Reputation: 13
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
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