Python function that handles scalar or arrays

Question

How best to write a function that can accept either scalar floats or numpy vectors (1-d array), and return a scalar, 1-d array, or 2-d array, depending on the input?

The function is expensive and is called often, and I don't want to place a burden on the caller to do special casts to the arguments or return values. It only needs to treat numbers (not lists or other iterable things).

np.vectorize might be slow (Broadcasting a python function on to numpy arrays) and other answers (Getting a Python function to cleanly return a scalar or list, depending on number of arguments) and np.asarray (A python function that accepts as an argument either a scalar or a numpy array) does not help with getting dimensions required for the output array.

This type of code would work in Matlab, Javascript, and other languages:

import numpy as np  def func( xa, ya ):     # naively, I thought I could do:     xy = np.zeros( ( len(xa), len(ya) ) )     for j in range(len( ya )):         for i in range(len( xa )):             # do something complicated             xy[i,j] = x[i]+y[j]                 return xy 

Works fine for arrays:

x = np.array([1., 2.]) y = np.array([2., 4.]) xy = func(x,y) print xy  [[ 3.  5.]  [ 4.  6.]] 

But does not work for scalar floats:

x = 1. y = 3. xy = func(x,y) print xy  <ipython-input-64-0f85ad330609> in func(xa, ya)       4 def func( xa, ya ):       5     # naively, I thought I could do: ----> 6     xy = np.zeros( ( len(xa), len(ya) ) )       7     for j in range(len( ya )):       8         for i in range(len( xa )):  TypeError: object of type 'float' has no len() 

Using np.asarray in a similar function gives:

<ipython-input-63-9ae8e50196e1> in func(x, y)       5     xa = np.asarray( x );       6     ya = np.asarray( y ); ----> 7     xy = np.zeros( ( len(xa), len(ya) ) )       8     for j in range(len( ya )):       9         for i in range(len( xa )):  TypeError: len() of unsized object 

What is the fast, elegant, and pythonic approach?

Answer 1

All over the numpy code base you find things like:

def func_for_scalars_or_vectors(x):     x = np.asarray(x)     scalar_input = False     if x.ndim == 0:         x = x[None]  # Makes x 1D         scalar_input = True      # The magic happens here      if scalar_input:         return np.squeeze(ret)     return ret 

Answer 2

As a matter of opinion, I would prefer to have a function be flexible on input types, but always return a consistent type; this is what will ultimately prevent callers from having to check return types (the stated goal).

For example, allow scalars and/or arrays as arguments, but always return the array.

def func(x, y):     # allow (x=1,y=2) OR (x=[1,2], y=[3,4]) OR (!) (x=1,y=[2,3])     xn = np.asarray([x]) if np.isscalar(x) else np.asarray(x)     yn = np.asarray([y]) if np.isscalar(y) else np.asarray(y)      # calculations with numpy arrays xn and xy     res = xn + yn  # ..etc...     return res 

(Still, the example can easily be modified to return a scalar, by setting a flag "scalar=True", yada yada yada.. but you'd also have to handle one arg's a scalar, the other is an array, etc.; seems like a lot of YAGNI to me.)