Python function that handles scalar or arrays -

how best write function can accept either scalar floats or numpy vectors (1-d array), , return scalar, 1-d array, or 2-d array, depending on input?

the function expensive , called often, , don't want place burden on caller special casts arguments or return values. needs treat numbers (not lists or other iterable things).

np.vectorize might slow (broadcasting python function on numpy arrays) , other answers (getting python function cleanly return scalar or list, depending on number of arguments) , np.asarray (a python function accepts argument either scalar or numpy array) not getting dimensions required output array.

this type of code work in matlab, javascript, , other languages:

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

works fine arrays:

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

but not work 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, thought do: ----> 6     xy = np.zeros( ( len(xa), len(ya) ) )       7     j in range(len( ya )):       8         in range(len( xa )):  typeerror: object of type 'float' has no len() 

using np.asarray in 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     j in range(len( ya )):       9         in range(len( xa )):  typeerror: len() of unsized object 

what fast, elegant, , pythonic approach?

all on numpy code base 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      # magic happens here      if scalar_input:         return np.squeeze(ret)     return ret