make np.vectorize return scalar value on scalar input

问题

The following code returns an array instead of expected float value.

def f(x):
    return x+1
f = np.vectorize(f, otypes=[np.float])
>>> f(10.5)
array(11.5)

Is there a way to force it return simple scalar value if the input is scalar and not the weird array type?

I find it weird it doesn't do it by default given that all other ufuncs like np.cos, np.sin etc do return regular scalars

Edit: This the the code that works:

import numpy as np
import functools

def as_scalar_if_possible(func):
    @functools.wraps(func) #this is here just to preserve signature
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)[()]

    return wrapper

@as_scalar_if_possible
@np.vectorize
def f(x):
    return x + 1

print(f(11.5)) # prints 12.5

回答1:

The result is technically a scalar as its shape is (). For instance, np.array(11.5)[0] is not a valid operation and will result in an exception. Indeed, the returned results will act as a scalar in most circumstances.

eg.

x = np.array(11.5)
print(x + 1) # prints 12.5
print(x < 12) # prints True, rather than [ True]
x[0] # raises IndexError

If you want to get a "proper" scalar value back then you can just wrap the vectorised function to check the shape of the returned array. This is what numpy ufuncs do behind the scenes.

eg.

import numpy as np

def as_scalar_if_possible(func):
    def wrapper(arr):
        arr = func(arr)
        return arr if arr.shape else np.asscalar(arr)
    return wrapper

@as_scalar_if_possible
@np.vectorize
def f(x):
    return x + 1

print(f(11.5)) # prints 12.5

来源：https://stackoverflow.com/questions/39272465/make-np-vectorize-return-scalar-value-on-scalar-input

标签

python

numpy

numpy-ufunc