A function works fine with one number, but not an array. What am I missing?-CodePudding

Pardon my ignorance, but I'm very new to coding in python. I have a pretty simple function; it just needs to make a calculation based on b's relative location to a and c:

a = 6
b = 3
c = 2

def function(a, b, c):
    if ((a >= b) & (b >= c)):
        return b - c 
    elif ((a <= b) & (b >= c)):
        return a - c
    else:
        return 0 
t = function(a, b, c)
print(t)

When I run it with simple numbers like above, it gives me the right answer no matter what I make b. (In this case 1)

But When I run it with a,b, and c as Numpy Arrays, it only returns b - c across the entire "t" array.

It's not too much different, but here's what I'm using for the array version:

def function(a, b, c):
    if ((a >= b) & (b >= c)).any():
        return b - c 
    elif ((a <= b) & (b >= c)).any():
        return a - c
    else:
        return 0 
t = function(a, b, c[i>1])
print(t)

(The [i>1] is there because there is a variable amount of array input, and another function will be used for when [i = 0])

I've also tried this: t = np.where(((prev2 >= Head_ELV) & (Head_ELV >= Bottom_ELV)).any, Head_ELV - Bottom_ELV, 0) but ran into the same result.

Would a while-loop work better?

CodePudding user response：

I don't think you need looping here as the problem can be solved using array operations. You could try the below, assuming the arrays are of the same length.

# import numpy to be able to work with arrays
import numpy as np

def function(a, b, c):
    # declare array t with only zeros
    t = np.zeros_like(a)

    # declare filters
    mask_1 = (a >= b) * (b >= c)
    mask_2 = (a <= b) * (b >= c)

    # modifying t based on the filters above
    t[mask_1] = (b - c)[mask_1]
    t[mask_2] = (a - c)[mask_2]
    return t

# example 2d arrays
a = np.array([[1800,5], [5,5]])
b = np.array([[3416,2], [3,4]])
c = np.array([[1714,2], [3,4]])

# run function
function(a, b, c)