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 a scalar, and promote the argument to a numpy array and add 1 to each element.

How could I do that? I suppose I could test argument class but that seems like a bad idea. If I do this:

def incrementElements(x):
   return numpy.array(x)+1

it works properly on arrays or iterables but not scalars. The problem here is that numpy.array(x) for scalar x produces some weird object that is contained by a numpy array but isn't a "real" array; if I add a scalar to it, the result is demoted to a scalar.

Answer 1

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 2

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