A python function that accepts as an argument either a scalar or a numpy array

Here's one solution:

import numpy as np

def sign(x):
    y = np.ones_like(x)
    y[np.asarray(x) < 0] = -1

    if isinstance(x, np.ndarray):
        return y
    else:
        return type(x)(y)

This should return a value of the same type as the input. For example sign(42) gives 1, sign(42.0) gives 1.0. If you give it an ndarray, it will work like np.sign.

In general, you might proceed with the assumption that your input is an ndarray. If you try to access an attribute or method that an ndarray has, but your input does not, then you fall back to operating on a scalar type. Use exceptions to implement this. For example:

def foo_on_scalars(x):
    # do scalar things

def foo(x):
    try:
        # assume x is an ndarray
    except AttributeError:
        foo_on_scalars(x)

The numpy functions naturally handle scalar or array inputs and preserve the shape in the output. So, it's always best to find the numpy functions doing the job. In this case, the function should be np.sign as suggested above. For different logics, you can use np.where(x>0, 1, -1), which works for scalar and array values of x.


np.vectorize can be used to achieve that, but would be slow because all it does, when your decorated function is called with an array, is looping through the array elements and apply the scalar function to each, i.e. not leveraging numpy's speed.

A method I find useful for vectorizing functions involving if-else is using np.choose:

def sign_non_zero(x):
    return np.choose(
        x > 0,  # bool values, used as indices to the array
        [
            -1, # index=0=False, i.e. x<=0
            1,  # index=1=True, i.e. x>0
        ])

This works when x is either scalar or an array, and is faster than looping in python-space.

The only disadvantage of using np.choose is that it is not intuitive to write if-else logic in that manner, and the code is less readable. Whenver I use it, I include comments like the ones above, to make it easier on the reader to understand what is going on.


i wonder if it's a vectorized function that you want:

>>> import numpy as NP

>>> def fnx(a):
        if a > 0:
            return 1
        else:
            return -1

>>> vfnx = NP.vectorize(fnx)

>>> a = NP.random.randint(1, 10, 5)
array([4, 9, 7, 9, 2])

>>> a0 = 7

>>> vfnx(a)
array([1, 1, 1, 1])

>>> vfnx(a0)
array(1)