how can I make a numpy function that accepts a numpy array, an iterable, or a scalar?

Question

Suppose I have this:

def incrementElements(x):
   return x+1

but I want to modify it so that it can take either a numpy array, an iterable, or

Answer 1

Pierre GM's answer is great so long as your function exclusively uses ufuncs (or something similar) to implicitly loop over the input values. If your function needs to iterate over the inputs, then np.asarray doesn't do enough, because you can't iterate over a NumPy scalar:

import numpy as np

x = np.asarray(1)
for xval in x:
    print(np.exp(xval))

Traceback (most recent call last):
  File "Untitled 2.py", line 4, in <module>
    for xval in x:
TypeError: iteration over a 0-d array

If your function needs to iterate over the input, something like the following will work, using np.atleast_1d and np.squeeze (see Array manipulation routines — NumPy Manual). I included an aaout ("Always Array OUT") arg so you can specify whether you want scalar inputs to produce single-element array outputs; it could be dropped if not needed:

def scalar_or_iter_in(x, aaout=False):
    """
    Gather function evaluations over scalar or iterable `x` values.

    aaout :: boolean
        "Always array output" flag:  If True, scalar input produces
        a 1-D, single-element array output.  If false, scalar input
        produces scalar output.
    """
    x = np.asarray(x)
    scalar_in = x.ndim==0

    # Could use np.array instead of np.atleast_1d, as follows:
    # xvals = np.array(x, copy=False, ndmin=1)
    xvals = np.atleast_1d(x)
    y = np.empty_like(xvals, dtype=float)  # dtype in case input is ints
    for i, xx in enumerate(xvals):
        y[i] = np.exp(xx)  # YOUR OPERATIONS HERE!

    if scalar_in and not aaout:
        return np.squeeze(y)
    else:
        return y


print(scalar_or_iter_in(1.))
print(scalar_or_iter_in(1., aaout=True))
print(scalar_or_iter_in([1,2,3]))


2.718281828459045
[2.71828183]
[ 2.71828183  7.3890561  20.08553692]

Of course, for exponentiation you should not explicitly iterate as here, but a more complex operation may not be expressible using NumPy ufuncs. If you do not need to iterate, but want similar control over whether scalar inputs produce single-element array outputs, the middle of the function could be simpler, but the return has to handle the np.atleast_1d:

def scalar_or_iter_in(x, aaout=False):
    """
    Gather function evaluations over scalar or iterable `x` values.

    aaout :: boolean
        "Always array output" flag:  If True, scalar input produces
        a 1-D, single-element array output.  If false, scalar input
        produces scalar output.
    """
    x = np.asarray(x)
    scalar_in = x.ndim==0

    y = np.exp(x)  # YOUR OPERATIONS HERE!

    if scalar_in and not aaout:
        return np.squeeze(y)
    else:
        return np.atleast_1d(y)

I suspect in most cases the aaout flag is not necessary, and that you'd always want scalar outputs with scalar inputs. In such cases, the return should just be:

    if scalar_in:
        return np.squeeze(y)
    else:
        return y

Answer 2

You could try

def incrementElements(x):
    x = np.asarray(x)
    return x+1

np.asarray(x) is the equivalent of np.array(x, copy=False), meaning that a scalar or an iterable will be transformed to a ndarray, but if x is already a ndarray, its data will not be copied.

If you pass a scalar and want a ndarray as output (not a scalar), you can use:

def incrementElements(x):
    x = np.array(x, copy=False, ndmin=1)
    return x

The ndmin=1 argument will force the array to have at least one dimension. Use ndmin=2 for at least 2 dimensions, and so forth. You can also use its equivalent np.atleast_1d (or np.atleast_2d for the 2D version...)

Answer 3

Existing solutions all work. Here's yet another option.

def incrementElements(x):
    try:
        iter(x)
    except TypeError:
        x = [x]
    x = np.array(x)
    # code that operators on x

Answer 4

This is an old question, but here are my two cents.

Although Pierre GM's answer works great, it has the---maybe undesirable---side effect of converting scalars to arrays. If that is what you want/need, then stop reading; otherwise, carry on. While this might be okay (and is probably good for lists and other iterables to return a np.array), it could be argued that for scalars it should return a scalar. If that's the desired behaviour, why not follow python's EAFP philosophy. This is what I usually do (I changed the example to show what could happen when np.asarray returns a "scalar"):

def saturateElements(x):
    x = np.asarray(x)
    try:
        x[x>1] = 1
    except TypeError:
        x = min(x,1)
    return x

I realize it is more verbose than Pierre GM's answer, but as I said, this solution will return a scalar if a scalar is passed, or a np.array is an array or iterable is passed.