Numpy Divide Arrays With Multiple Out Conditions

Question

I have two, two-dimensional arrays (say arrayA & arrayB) that are exactly the same size (2500X, 1500Y). I am interested in dividing array A by array B, but have three conditions that I would like to be excluded from the division and instead replaced with a specific value. These conditions are:

1. If arrayB contains zero at point (Bx,By), replace output (Cx,Cy) with (arrayA*arrayA)
2. If arrayA contains zero at point (Ax,Ay), replace output (Cx,Cy) with 0.50
3. If both arrayA & B at overlapping points (Ax,Ay & Bx,By) contain 0, replace output (Cx,Cy) with 1

I've found that numpy.divide parameters out and where allow me to define each of these individually, so I've taken the first condition and arranged it as follows:

arrayC = np.divide(arrayA, arrayB, out=(arrayA*arrayA), where=arrayB!=0)

My question is how can I combine the other two conditions and their desired outputs within this operation?

Answer 1

One solution, not sure it is the fastest

za=A==0
zb=B==0
case0=(~za)&~zb
case1=zb&~za
case2=za&~zb
case3=za&zb
C=case3*1 + case2*0.5 + case1*A*A # Case 3,2, 1
C[case0]=(A[case0]/B[case0])

Could be more compact with less intermediate values, but I've chosen clarity.

You could also use a cascade of np.where

zb=B==0
C=np.where(A==0, np.where(zb,1,0.5), np.where(zb, A*A, A/B))

Edit: better version (but still not perfect)

zb=B==0
za=A==0
C=np.where(za, np.where(zb,1,0.5), A*A)
np.divide(A, B, out=C, where=(~zb)&~za)

It combines np.where and your np.divide where= It is as fast as the previous solution. And does not complain about 0-division, since division occurs only for the cases where it is needed. Nevertheless, it computes the first version of C (the one before np.divide), and particularly A*A, everywhere, even where it is not needed, since it will be overwritten. So, it could probably be better.