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…)