is there a pythonic way to change scalar and 0d-array to 1d array?

Question

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?

Answer 1

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

Answer 2

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.

Answer 3

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.

Answer 4

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.

Answer 5

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.

Answer 6

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.

Answer 7

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