Works for me:
>>> import numpy as np
>>> @np.vectorize
... def diff_if_bigger(x, y):
... return y - x if y > x else 0
...
>>> diff_if_bigger(np.array([5.6,7.0]), 8)
array([ 2.4, 1. ])
Note that np.vectorize
isn't really meant as a decorator except for the simplest cases. If you need to specify an explicit otype
, use the usual form new_func = np.vectorize(old_func, otypes=...)
or use functools.partial
to get a decorator.
Note too that np.vectorize
, by default, gets its output type from evaluating the function on the first argument:
The data type of the output of
vectorized
is determined by calling the function with the first element of the input.
So, you should pass float
and return float
if you want to ensure that it infers float
as the output dtype (e.g. use else 0.0
and pass y = 8.0
).