Numbering vectors instead of loops

Question

I wrote code in Python that works great but is very slow; I think because of the cycles. I hope you can speed up the following operations with numpy commands. Let me set a goal.

Suppose I have an array of 2D numpy all_CMs size row x col . For example, consider a 6 x 11 array (see the figure below).

• I want to calculate the average value for all rows, i.e. sumⱼ aᵢⱼ, which will result in an array. This, of course, can be easily done. (I call this CM_tilde value)

• Now for each row, I want to calculate the average value of some selected values, namely, all values ​​below a certain threshold, calculating their sum and dividing it by the number of all columns ( N ). If the value exceeds the specified threshold, the CM_tilde value is CM_tilde (the average value for the entire row). This value is called CM

• Then the CM value is subtracted from each element in the row.

In addition to this, I want to have a numpy array or list listing all of these CM values.

Picture:

The following code works, but very slow (especially if arrays are getting large)

 CM_tilde = np.mean(data, axis=1) N = data.shape[1] data_cm = np.zeros(( data.shape[0], data.shape[1], data.shape[2] )) all_CMs = np.zeros(( data.shape[0], data.shape[2])) for frame in range(data.shape[2]): for row in range(data.shape[0]): CM=0 for col in range(data.shape[1]): if data[row, col, frame] < (CM_tilde[row, frame]+threshold): CM += data[row, col, frame] else: CM += CM_tilde[row, frame] CM = CM/N all_CMs[row, frame] = CM # calculate CM corrected value for col in range(data.shape[1]): data_cm[row, col, frame] = data[row, col, frame] - CM print "frame: ", frame return data_cm, all_CMs 

Any ideas?

Answer 1

It is very easy to digitize what you do:

 import numpy as np #generate dummy data nrows=6 ncols=11 nframes=3 threshold=0.3 data=np.random.rand(nrows,ncols,nframes) CM_tilde = np.mean(data, axis=1) N = data.shape[1] all_CMs2 = np.mean(np.where(data < (CM_tilde[:,None,:]+threshold),data,CM_tilde[:,None,:]),axis=1) data_cm2 = data - all_CMs2[:,None,:] 

Comparing this to your originals:

 In [684]: (data_cm==data_cm2).all() Out[684]: True In [685]: (all_CMs==all_CMs2).all() Out[685]: True 

The logic is that we are simultaneously working with arrays of size [nrows,ncols,nframes] . The main trick is to use python translation by turning CM_tilde size [nrows,nframes] to CM_tilde[:,None,:] size [nrows,1,nframes] . Python will then use the same values ​​for each column, as this is a one-dimensional dimension of this modified CM_tilde .

Using np.where , we choose (based on threshold ) whether we want to get the corresponding data value or, again, the CM_tilde broadcast CM_tilde . New usage of np.mean allows to calculate all_CMs2 .

At the last stage, we used the broadcast, subtracting this new all_CMs2 from the corresponding data elements.

This can help you vectorize your code this way if you look at the implicit indices of your temporary variables. I mean, your temporary CM variable lives inside the loop over [nrows,nframes] , and its value is reset with each iteration. This means that CM is actually the number of CM[row,frame] (later explicitly assigned to the 2d all_CMs ), and from here it is easy to see that you can build it by summing the corresponding number of CMtmp[row,col,frames] along its column size. If this helps, you can name the np.where(...) part as CMtmp for this purpose, and then calculate np.mean(CMtmp,axis=1) . The same result, obviously, but probably more transparent.

Answer 2

Here is my vectorization of your function. I worked from the inside out and commented on earlier versions when I went. So the first loop that I have vectorized has ### comment marks.

This is not as clean and sound as @Andras's answer, but I hope it is educational, giving an idea of ​​how you can solve this problem gradually.

 def foo2(data, threshold): CM_tilde = np.mean(data, axis=1) N = data.shape[1] #data_cm = np.zeros(( data.shape[0], data.shape[1], data.shape[2] )) ##all_CMs = np.zeros(( data.shape[0], data.shape[2])) bmask = data < (CM_tilde[:,None,:] + threshold) CM = np.zeros_like(data) CM[:] = CM_tilde[:,None,:] CM[bmask] = data[bmask] CM = CM.sum(axis=1) CM = CM/N all_CMs = CM.copy() """ for frame in range(data.shape[2]): for row in range(data.shape[0]): ###print(frame, row) ###mask = data[row, :, frame] < (CM_tilde[row, frame]+threshold) ###print(mask) ##mask = bmask[row,:,frame] ##CM = data[row, mask, frame].sum() ##CM += (CM_tilde[row, frame]*(~mask)).sum() ##CM = CM/N ##all_CMs[row, frame] = CM ## calculate CM corrected value #for col in range(data.shape[1]): # data_cm[row, col, frame] = data[row, col, frame] - CM[row,frame] print "frame: ", frame """ data_cm = data - CM[:,None,:] return data_cm, all_CMs 

The output matches for this small test case, which more than helped me get the right sizes.

 threshold = .1 data = np.arange(4*3*2,dtype=float).reshape(4,3,2)