Can I use a more pythonic way for my conditional array operations?

In Python using numpy I have to do some iterations over a numpy 2D array. I boiled down the original problem to its essentials:

f = np.array([[0, 0, 0, 0, 0, 0, 0], [0, 10, 22, 30, 40, 50, 0], [0, 11, 22, 33, 44, 55, 0], [0, 0, 0, 0, 0, 0, 0]])
u = np.array([[1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, -1, 1],     [1, 1, -1, -1, -1, 1, 1],   [1, 1, 1, 1, 1, 1, 1]])
x = np.zeros_like(f)

for i in range(1,u.shape[0]-1):
    for j in range(1,u.shape[1]-1):
        if u[i,j] > 0: 
            x[i,j] = u[i,j]*(f[i,j]-f[i,j-1])
        else:
            x[i,j] = -u[i,j]*(f[i,j+1]-f[i,j])

In this iteration I have to iterate over all 2D indices. Is there a more elegant way? If there would be no condition on u I would just write:

y = np.zeros_like(f)    

y[1:-1, 1:-1] = u[1:-1, 1:-1] * (f[1:-1, 1:-1] - f[1:-1, :-2])

which gives the same result, given u>0

I was thinking to use m = ma.masked_array(u, (u > 0)).mask somehow, but by indexing like f[m] I get only a 1D array, so I get the 2D structure lose. Is there a more Pythonic way to carry out this iteration?

Solution

Your two operations can be seen as a single diff, just shifted depending on the sign of u:

x = np.zeros_like(f)
d = np.diff(f, axis=1)
x[:, 1:-1] = np.where(u[:, 1:-1]>0, d[:, :-1], d[:, 1:])

NB. you can replace : on the first axis in all slicings by 1:-1 to really match your loop.

Output:

array([[  0,   0,   0,   0,   0,   0,   0],
       [  0,  10,  12,   8,  10, -50,   0],
       [  0,  11,  11,  11,  11,  11,   0],
       [  0,   0,   0,   0,   0,   0,   0]])

the exact conditional that matches your loop is:

x[1:-1, 1:-1] = np.where(u[1:-1, 1:-1]>0,
                         u[1:-1, 1:-1]*(f[1:-1, 1:-1]-f[1:-1, :-2]),
                        -u[1:-1, 1:-1]*(f[1:-1, 2:]-f[1:-1, 1:-1]))