How to write a function to process arrays as input element by element and returns array

Question

I am trying to write a function and I want it to return one element when the input is element and an array of outputs if the input is array such that each element of output array is associated with the same place in input array. I am giving a dummy example:

import numpy as np
def f(a):
    if a<5:
        print a;
f(np.arange(11))

This code returns the error: if a<5:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I expect the output to be:

0
1
2
3
4

How can I make it to work the way I explained as I believe many python functions are working in this way?

Thanks.

Solution 1

When I have had to deal with this kind of thing, I usually start by doing an np.asarray of the input at the beginning, setting a flag if it is 0-dimensional (i.e. a scalar), promoting it to 1-D, running the function on the array, and converting it back to a scalar before returning if the flag was set. With you example:

def f(a):
    a = np.asarray(a)
    is_scalar = False if a.ndim > 0 else True
    a.shape = (1,)*(1-a.ndim) + a.shape
    less_than_5 = a[a < 5]
    return (less_than_5 if not is_scalar else
            (less_than_5[0] if less_than_5 else None))

>>> f(4)
4
>>> f(5)
>>> f([3,4,5,6])
array([3, 4])
>>> f([5,6,7])
array([], dtype=int32)

If you do this very often, you could add all that handling in a function decorator.

Solution 2

if you want the function to react depending upon whether the given input is a list or just an int, use:

def f(a):
    if type(a)==type([]):
        #do stuff
    elif type(a)==type(5):
        #do stuff
    else
        print "Enter an int or list"

by the above step, the function checks if the given input is an array, if the condition is true, the first block is used. next if block checks if the input is an int. Else the else block is executed

Solution 3

import numpy as np
def f(a):
    result = a[a<5]
    return result

def report(arr):
    for elt in arr:
        print(elt)

report(f(np.arange(11)))

Generally speaking I dislike putting print statements in functions (unless the function does nothing but print.) If you keep the I/O separate from the computation, then your functions will be more re-usable.

---

It is also generally a bad idea to write a function that returns different types of output, such as a scalar for some input and an array for other input. If you do that, then subsequent code that uses this function will all have to check if the output is a scalar or an array. Or, the code will have to be written very carefully to control what kind of input is sent to the function. The code can be come very complicated or very buggy if you do this.

Write simple functions -- ones that either always return an array, or always return a scalar.

Comments

[Cupitor]

As I said in my question this is a dummy body only to give an example. In my original code I don't have any printing. The way you solved the problem is problem specific. I am trying to figure out how python or numpy functions take care of this situations when they are unaware of the shape of input.

[Cupitor]

Thank you. This will work, but I am wondering is this the way that such cases are handled by built-in or for example numpy functions?

[Aswin Murugesh]

This is the way to check if we do not know the type of input we receive

[Cupitor]

Well usually numpy arrays are written in this way as they usually deal with matrices! My situation is nearly the same!

[unutbu]

Can you give an example of a numpy function that behaves this way?

[Cupitor]

np.dot(5,3) returns 15 if you give(2,[3,3]) it will return [6,6]

[unutbu]

I suspect np.dot is calling np.asarray on its arguments to promote everything to arrays, and then using broadcasting to promote arrays to equal shapes. Even though np.dot(5,3) returns a scalar, if you are doing that kind of thing in your code, you are misusing NumPy. np.dot(5,3) is 250x slower than 5*3.

[Cupitor]

Well that is what I am wondering. What is the general routine in such cases. Is it as you suspect or any other way?