how can I make a numpy function that accepts a numpy array, an iterable, or a scalar?
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...)
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