Optimizing an array calculation with numpy

Question

I have the following algorithm which I am trying to optimize in Python:

I define three logical matrices (X, Y, & Z) before hand. Each has N columns. I also define two very large numbers N and A, and an array of arbitrary numbers I call tot_list which has length A. I want to go through X, Y, and Z, find where the elements of each is 1, and then change that corresponding value in a list A_list to a random float. When this operation is done, I update tot_list with A_list. Below is the actual code:

for i in range(0,N):
    A_list = np.ones([1,A])
    
    A_list[0][np.where(X[:,i]==1)[0]] = np.random.random()
    A_list[0][np.where(Y[:,i]==1)[0]] = np.random.random()
    A_list[0][np.where(Z[:,i]==1)[0]] = np.random.random()
    
    tot_list = tot_list+A_list

The code given above does exactly what I want, but it runs terribly slow. I have to deal with gigantic numbers (i.e., with A on the order of 10^6 and N on the order of 10^2), and I have to optimize the above as much as possible. I've tried looking into getting rid of the for-loop, but I am not sure how to implement this. Any suggestions would be welcome.

Answer 1

Ok, without example inputs and outputs it is a bit hard to know the problem, but let me try and answer and I can edit if it isn't right.

As I see it you are looping through every row of X, Y and Z and finding every element that is 1. You can do this for every element of a single array at once by just doing e.g. X == 1.

You then find the places where either X, Y, Z equals one and set that to a random number, otherwise you set it to 1 (from your definition of A_list). So we can find the OR of all the array elements that equal 1 e.g.

equal_one = (X == 1) | (Y == 1) | (Z == 1)

then we can use the three argument version of np.where to set the values equal to 1 to a random number and those not equal to 1 to 1.

out = np.where(equal_one, np.random.random(), 1)

Finally, we sum every row of the array up, along one axis. I think from your example this is equivalent to

tot_list = out.sum(axis=1)

Putting it all together

tot_list = np.where(
    (X == 1) | (Y == 1) | (Z == 1),
    np.random.random(),
    1
).sum(axis=1)