For example, I want to change 'a' to 0 if 'a' is less than 5
def foo(a):
return 0 if a < 5 else a
to make it work for numpy array, I it change to:
def foo2(a):
a[a < 5] = 0
return a
The problem is that I want the function work for both scalar and arrays.
The isscalar() function can test if 'a' is a scalar, but it returns false for 0d-arrays, like array(12).
Is there a pythonic way to change scalar and 0d-array to 1d array and remain other ndarray unchanged?
well, I come with a solution that seems to work
def foo3(a):
return a * (a >= 5)
foo3(4)
=> 0
foo3(6)
=> 6
foo3(np.array(3))
=> 0
foo3(np.array(6))
=> 6
foo3(np.array([1, 5]))
=> array([0, 5])
It works fine, but I don't know whether it is safe to do so.
You can use numpy.vectorize, with the original scalar implemenation.
@np.vectorize
def foo(a):
return 0 if a < 5 else a
foo(3)
=> array(0)
foo(7)
=> array(7)
foo([[3,7], [8,-1]])
=> array([[0, 7],
[8, 0]])
Note that when using vectorize, you give up speed (your calculation is no longer vectorized in the numpy/C level) for simplicity (you write your function in its simple form, for scalars).
If you don't mind the function returning an array even if it is supplied with a scalar, I'd be inclined to do it this way:
import numpy as np
def foo(a,k=5):
b = np.array(a)
if not b.ndim:
b = np.array([a])
b[b < k] = 0
return b
print(foo(3))
print(foo(6))
print(foo([1,2,3,4,5,6,7,8,9]))
print(foo(np.array([1,2,3,4,5,6,7,8,9])))
... which produces the results:
[0]
[6]
[0 0 0 0 5 6 7 8 9]
[0 0 0 0 5 6 7 8 9]
As you can see from the test examples, this function works as intended if it is supplied with a regular Python list instead of a numpy array or a scalar.
Creating two arrays in the process may seem wasteful but, first, creating b prevents the function from having unwanted side-effects. Consider that:
def foobad(a,k=5):
a[a < k] = 0
return a
x = np.array([1,2,3,4,5,6,7,8,9])
foobad(x)
print (x)
... prints:
[0 0 0 0 5 6 7 8 9]
... which is probably not what was desired by the user of the function. Second, if the second array creation occurs because the function was supplied with a scalar, it will only be creating an array from a 1-element list, which should be very quick.
Use np.atleast_1d
This will work for any input (scalar or array):
def foo(a):
a = np.atleast_1d(a)
a[a < 5] = 0
return a
Note though, that this will return a 1d array for a scalar input.
np.array or np.asarray to convert the input. Then check if the number of dimensions is 0. If a dimension is added using np.newaxis, None or np.atleast1d, you need to use remove it later, e.g. using np.squeeze or just a[0].This is an answer to the last part of your question. A quick way to change a scalar or a 0d array to a 1d array using np.reshape after checking the dimension with np.ndim.
import numpy as np
a = 1
if np.ndim(a) == 0:
np.reshape(a, (-1))
=> array([1])
Then,
b = np.array(1)
=> array(1) # 0d array
if np.ndim(b) == 0:
np.reshape(b, (-1))
=> array([1]) # 1d array. you can iterate over this.
Try this
def foo(a, b=5):
ix = a < b
if not np.isscalar(ix):
a[ix] = 0
elif ix:
a = 0
return a
print([foo(1), foo(10), foo(np.array(1)), foo(np.array(10)), foo(np.arange(10))])
Outputs
[0, 10, 0, array(10), array([0, 0, 0, 0, 0, 5, 6, 7, 8, 9])]
Note that array(1) > 0 gives bool instead of np.bool_, so safe to use np.isscalar for ix.
Reapeating my answer elsewhere,
The numpy functions naturally handle scalar or nd 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, np.where is your friend.
def foo(a):
return np.where(a<5, a, 0)
foo(6)
> array(0)
foo(np.array([3,4,5,6]))
> array([3, 4, 0, 0])
a.ndim == 0to recognize 0d-array. Take a look at this question stackoverflow.com/questions/773030/…zero_d_array[()]orzero_d_array[...]both work. There does seem to be a weird edge case with boolean advanced indexing, though.ain place? That'll produce strange inconsistencies between the scalar and array cases.