3

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?

cncggvg
  • 657
  • 6
  • 12
  • Scalars can't be indexed, you'll need to handle the scalar case separately, possibly by converting it to a 0d-array first then back. – simonzack Aug 23 '14 at 04:58
  • You can use `a.ndim == 0` to recognize 0d-array. Take a look at this question http://stackoverflow.com/questions/773030/why-are-0d-arrays-in-numpy-not-considered-scalar – Barmar Aug 23 '14 at 05:14
  • 0d-array cannot be indexed either, so I want to change them to 1d-arrays, but I just cannot find a good way to do this. – cncggvg Aug 23 '14 at 05:16
  • The error message is misleading; 0-dimensional arrays can be indexed. `zero_d_array[()]` or `zero_d_array[...]` both work. There does seem to be a weird edge case with boolean advanced indexing, though. – user2357112 Aug 23 '14 at 06:06
  • Are you sure you want your function to modify `a` in place? That'll produce strange inconsistencies between the scalar and array cases. – user2357112 Aug 23 '14 at 06:11

7 Answers7

4

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.

cncggvg
  • 657
  • 6
  • 12
2

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

shx2
  • 61,779
  • 13
  • 130
  • 153
1

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.

Simon
  • 10,679
  • 1
  • 30
  • 44
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.
Taro Kiritani
  • 724
  • 6
  • 24
0

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.

Syrtis Major
  • 3,791
  • 1
  • 30
  • 40
0

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.

nivniv
  • 3,421
  • 5
  • 33
  • 40
  • This needs to be slightly modified to return scalars. The best example I found so far is https://stackoverflow.com/a/29319864/7919597. You need use `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]. – Joe Dec 16 '17 at 08:05
0

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])
jChoi
  • 151
  • 6