如何在numpy数组中同时允许浮点数和和数组

Question

Using numpy and matplotlib it seems quite common that functions allow both a number (float or int) or a numpy array as argument like this:

import numpy as np

print np.sin(0)
# 0

x = np.arange(0,4,0.1)
y = np.sin(x)

In this example I call np.sin once with an integer argument, and once with a numpy array x . I now want to write a function that allows similar treatment, but I don't know how. For example:

def fun(foo, n):
    a = np.zeros(n)
    for i in range(n):
        a[i] = foo
    return a

would allow me to call fun like fun(1, 5) but not like fun(x, 5) . My actual calculation is much more complicated, of course.

How can I initialize a such that I can have both simple numbers or a whole array of numbers as elements?

Thanks a lot for your help!

Answer 1

You need a to inherit the dimensions of foo :

def fun(foo, n):
    a = np.zeros((n,) + np.shape(foo))
    for i in range(n):
        a[i] = foo
    return a

Answer 2

Builtin numpy functions often start with a

 def foo(a, ...):
     a = np.asarray(a)
     ...

That is, they transform the input argument(s) to array (no copy if it is already an array). The allows them to work with scalars and lists.

Once the argument is an array it has a shape and can be broadcasted against other arguments.

In your example, it's unclear what is supposed to happen when foo is an array

def fun(foo, n):
    a = np.zeros(n)
    for i in range(n):
        a[i] = foo
    return a

a is initialized as a dtype float array. That means a[i]=foo works only if foo is a single element number (scalar, possibly a single element array). If foo is an array with more than one value you probably get an error about attempting to set an element with a sequence.

a[i] is short for a[i,...] . That is it indexes on the 1st dimension. So if a was initialed correctly it could accept arrays as inputs (subject to broadcasting rules).

If a was initialed as np.zeros(n, dtype=object) , then a[i]=foo will work with anything, since it a just contains pointers to Python objects.

np.frompyfunc is a way of generating an array from a function. But it returns an array of dtype=object. np.vectorize uses that but gives you more control over the output type. But both work with scalars. An array, if given as argument, is passed element by element to the function.

Answer 3

You can use type identification :

import numpy as np

def double(a):
    if type(a)==int:
        return 2*a
    elif type(a)==float:
        return 2.0*a
    elif type(a)==list:
        return [double(x) for x in a]
    elif type(a)==np.ndarray:
        return 2*a
    else:
        print "bad type"

print double(7)
print double(7.2)
print double([2,9,7])
print double(np.array([[9,8],[2,3]]))

result :

>14
>14.4
>[4, 18, 14]
>[[18 16]
 [ 4  6]]

with eventually a recursive treatement as I did on list