3

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!

Improve this question
5
  • Is there a reason why a can't be a list? Commented Apr 22, 2016 at 12:53
  • @chepner the documentation show otherwise, docs.scipy.org/doc/numpy-1.10.1/user/basics.rec.html The columns need to be of one dtype, but not the whole array. Commented Apr 22, 2016 at 13:28
  • What do you want to get from fun(x,5)? 5 copies of x? What shape of an array? Commented Apr 22, 2016 at 15:09
  • @CactusWoman: later on I rely on a being a numpy array. I suppose I could create a as a list, then fill the list and finally write a = np.array(a). But isn't there a "cleaner" solution? Commented Apr 22, 2016 at 15:23
  • @hpaulj I used this extremely simplified function just as an example. I my actual problem I'm calculating quite complicated probabilities. Commented Apr 22, 2016 at 15:26

3 Answers 3

Reset to default
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.

Improve this answer
Sign up to request clarification or add additional context in comments.

Comments

0

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
Improve this answer

3 Comments

I tried something similar, but this way I cannot do fun(1, 5) anymore: Ints or floats don't have a "shape".
@Daniel: That's why I'm using np.shape(foo) not foo.shape, which works just fine on scalars (returning ())
oh, I missed the difference! Indeed I had only tried foo.shape. I will try that and come back to you!
-1

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

Improve this answer

4 Comments

isinstance(var, type) is a better way to write this. Also, your entire function can be rewritten as np.asarray(a) * 2
@Eric, well isinstance is needed if inheritance of type is used which is not the goal of the question. Then np.asarray transforms list to array which makes a difference with my function... By the way, of course such method has no interest, it was just to give a solution to identify int and arrays.
Inheritance is the goal, because this fails: a = np.zeros(2); double(a[0]) when using type(...) = .... So does double(2L)
Ah yeah you are right on this point, I didn t expect inheritance to be so present.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.