# Python functions taking float and array-like arguments

Published Thursday 26 October 2023 at 12:30 BST

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