Convert function arguments that are floats to numpy array

Question

I am trying to vectorize a function in python. I want the function's arguments to be able to be passed as either scalars (floats, ints, etc.) or as a numpy array. For arguments that are passed as scalars, they should then be broadcast as a numpy array of a specified length.

The function's arguments will then be used later in the function, so the name of the argument needs to be maintained.

This is the code I currently have, but it's not working:

def f(arg1, arg2, length = 4):
    arguments = locals()
    for name in arguments:
        if isinstance(arguments[name], (float,int)) : 
            eval(name) = np.array([float(arguments[name])] * length)
    return arg1,arg2

The function should return as such:

f(1,2,length = 4)

=> array([1.,1.,1.,1.]) , array([2.,2.,2.,2.])

Thanks!

Answer 1

How about this:

import numpy as np

def f(*args, **kwargs):
  length = kwargs.get("length", 1)
  ret = []
  for arg in args:
    if isinstance(arg, (float, int)):
      ret.append(np.repeat(arg, length))
    else:
      ret.append(arg)
  return tuple(ret)

print f(1, 2, length=4)

The whole approach strikes me as a bit odd. Would love to see a convincing use case for this.