Python functions taking float and array-like arguments

I often find myself wanting to write a Python function that accepts scalars and array-like objects for each argument. In these cases I have to make a decision about that function's return type. Typically the function performs Numpy operations on the arguments passed to it. But the return type of Numpy functions often surprises me. If a scalar or 0d array-like object is passed to such a Numpy function then it returns a numpy.float64 object. Only otherwise does it return a numpy.ndarray object. For example, using numpy.sin I find the following.

>>> import numpy as np
>>> w = np.sin(1); x = np.sin(1.); y = np.sin(np.array(1.)); z = np.sin([1., 1.])
>>> type(w); type(x); type(y); type(z)
<class 'numpy.float64'>
<class 'numpy.float64'>
<class 'numpy.float64'>
<class 'numpy.ndarray'>

Any operation on a numpy.float64 object returns another numpy.float64 object. More surprisingly, any operation on a 0d numpy.ndarray object also returns a numpy.float64 object. Even multiplying by one does this.

>>> x = np.array(1.)
>>> x; type(x); x.ndim; x.shape
array(1.)
<class 'numpy.ndarray'>
0
()
>>> x = 1.*x
>>> x; type(x); x.ndim; x.shape
1.0
<class 'numpy.float64'>
0
()

So what should I do? A home-made function that is just some sequence of Numpy functions will always behave in the same way. Let's imagine that I define such a function, func. It will return the appropriate Numpy type because it is a Numpy function, albeit a composite one.

def func(x):
    return some_function_of(x)

However, there is a good case to be made for a function always returning the same type of object, in this case a numpy.ndarray. To do so I must stop floats and 0d arrays from being returned as numpy.float64 objects. The trick is to cast both of these types as 1d numpy.ndarray objects and then squeeze them to become 0d numpy.ndarray objects. This code does the trick.

def func(x):
    if np.isscalar(x):
        x = np.atleast_1d(x)
        scalar_flag = True
    else:
        scalar_flag = False

    res = some_function_of(x)

    if scalar_flag:
        return res.squeeze()

    return res

In this way, scalar and 0d array-like arguments become 0d Numpy arrays. Other arguments become Numpy arrays while preserving their dimension. Let's say that some_function_of(x) is numpy.log(x). The Numpy function itself works as follows.

>>> x = np.log(1.)
>>> x; type(x); x.ndim; x.shape
0.0
<class 'numpy.float64'>
0
()
>>> x = log([1.])
>>> x; type(x); x.ndim; x.shape
array([1.])
<class 'numpy.ndarray'>
1
(1,)

But func works differently.

>>> x = func(1.)
>>> x; type(x); x.ndim; x.shape
array(0.)
<class 'numpy.ndarray'>
0
()
>>> x = func([1.])
>>> x; type(x); x.ndim; x.shape
array([0.])
<class 'numpy.ndarray'>
1
(1,)

However, I prefer my functions to be consistent with Numpy's, so I need not worry about any of this. But this occasionally leaves me with another problem. Suppose that inside my function I wish to perform some check on the arguments. Let's say that my function takes two arguments and that I want to check that the first is compatible with second. I might coerce both arguments into being Numpy arrays and then compare their shapes.

def func(x, y):
    x = np.asarray(x)
    y = np.asarray(y)
    if not x.shape == y.shape:
        raise ValueError("message.")

    return x

The function works as follows.

>>> func(1., 1.)
array(1.)
>>> func([1.], [1.])
array([1.])
>>> func(1., [1.])
ValueError: message.
>>> func([1.], [1., 1])
ValueError: message.

But if I want my function to be consistent with Numpy then casting scalars to arrays (i.e. promoting the dimension of an object) is exactly what I want to avoid doing. Having promoted the dimensions x and y I am stuck with Numpy arrays. I might try to extract the elements of x and y using x.item() and y.item() but this gives me Python float objects, not Numpy numpy.float64 objects. The trick this time is to treat scalars and array-like objects separately and to coerce the array-like objects into being Numpy arrays instead of coercing the floats as I did before.

def func(x, y):
    if np.isscalar(x) and np.isscalar(y):
        return x
    else:
        x = np.asarray(x)
        y = np.asarray(y)
    if x.shape == y.shape:
        return x
    else:
        raise ValueError("message.")

This behaves how I want it to.

>>> func(1., 1.)
1.
>>> func([1.], [1.])
array([1.])
>>> func(1., [1.])
ValueError: message.
>>> func([1.], [1., 1])
ValueError: message.

Much better.

The users of Stack Overflow helped me get my head straight on this by answering this question of mine.