Sum up over np.array or np.float

Question

We have a numpy-based algorithm that is supposed to handle data of different type.

def my_fancy_algo(a):
    b = np.sum(a, axis=1)
    # Do something b
    return b

If we pass a=np.array[1.0, 2.0, 3.0] then b evaluates to [6.0].

If we pass a=6.0 then we get

*** ValueError: 'axis' entry is out of bounds

The desired behavior would be that we get same return value 6.0 not ([6.0]) for both inputs.

What is the correct pythonic and safe way to handle this? type? shape?

Answer 1

Your example array actually gives the same problem as a scalar:

>>> a = np.array([1.0,2.0,3.0])
>>> np.sum(a, axis=1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.4/site-packages/numpy/core/fromnumeric.py", line 1724, in sum
    out=out, keepdims=keepdims)
  File "/usr/lib/python3.4/site-packages/numpy/core/_methods.py", line 32, in _sum
    return umr_sum(a, axis, dtype, out, keepdims)
ValueError: 'axis' entry is out of bounds

The good news is that there's a numpy function exactly for ensuring making numpy calls with axis=1 will work - it's called np.atleast_2d:

>>> np.sum(np.atleast_2d(a), axis=1)
array([ 6.])
>>> np.sum(np.atleast_2d(6.0), axis=1)
array([ 6.])

But since you apparently want a scalar answer, you could instead just drop the axis argument entirely:

>>> np.sum(a)
6.0
>>> np.sum(6.0)
6.0

Answer 2

np.sum(np.array([1.0, 2.0, 3.0]), axis=1) produces ValueError: 'axis' entry is out of bounds for me.

Did you mean to put axis=0 in line 2? Then it works for arrays as well as scalars:

>>> np.sum(np.array([1.0, 2.0, 3.0]), axis=0)
6
>>> np.sum(3, axis=0)
3