Performing complex conditional operation on numpy arrays

Question

import numpy as np

mat_a = np.random.random((5, 5))
mat_b = np.random.random((5, 5))
mat_c = np.random.random((5, 5))

Lets say for a specific grid cell, the values in the same position in the 3 arrays are as follows:
mat_a, A = 0.3
mat_b, B = 0.2
mat_c, C = 0.1

Here, we find the array with the least value, in this case it is C

1. We compute the amount of C that should be allocated to B as 0.1 * (0.2/ (0.2 + 0.3)) i.e. Value of cell in C multiplied by the fraction of B with total being A + B. The newly computed value is stored in a 2D array called C_B

2. Similarly, the amount of C that should be allocated to A is 0.1 * (0.3/(0.2 + 0.3)). The newly computed value is stored in a 2D array called C_A.

3. We repeat this process for cells where least value is in array B, storing the newly computed results in 2D arrays B_C and B_A respectively.

4. We repeat this process for cells where least value is in array A, storing the newly computed results in 2D arrays A_C and A_B respectively.

The only way I can think of doing this is using nested for loops, but that would be prohibitive for larger arrays and not very pythonic. Is there a fast and pythonic solution?

-- edit

C_B should contain 0 where mat_c does not contain smallest value

Answer 1

One solution is to calculate all values, replace unwanted ones with zeros.

mat_a = np.random.random((5, 5))
mat_b = np.random.random((5, 5))
mat_c = np.random.random((5, 5))
bigmat = np.stack((mat_a, mat_b, mat_c)) # this is a 3, 5, 5 array
minima = np.argmin(bigmat, axis=0) # contains a 5x5 array of 0,1,2 for a,b,c respectively
c_a = mat_c * mat_a / (mat_b + mat_c)
c_a[minima != 2] = 0

You can repeat this for the other 5 answer arrays. Or, you could also do:

c_a = np.zeros((5,5))
c_a[minima == 2] = (mat_c * mat_a / (mat_b + mat_c))[minima == 2]