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.
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.
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.
