python numpy operation instead of for loops
Question: I wrote some lines in python which work fine but is very slow; I think due to
the for loops. I hope one can speed up the following operations using numpy
commands. Let me define the goal.

Let's assume I have (always) a 3d numpy array (first dim: row, second dim:
col, third dim: frame); for simplecity I take a 2d array (since the operation
must be applied for every 2d array in the third dimension). Moreover, I take a
6x11 array, also for simplicity reasons (see drawing below).

  1. I want to calculate the mean for all rows, i.e. sum_j a_ij resulting in an array. This, of course can be easily done. (I call this calue CM_tilde)

  2. Now, for **each row** I want to calculate a selected mean, i.e. adding all values below a certain threshold and dividing them by the number of all columns (=N). If the value is above this defined threshold, the CM_tilde value (mean of the entire row) is added. This value is called CM

  3. Afterwards, the CM value is subtracted from each element in the row

In addition to this I want to have a numpy array or list where all those CM
values are listed.

The figure:

[![figure](http://i.stack.imgur.com/w0V6p.png)](http://i.stack.imgur.com/w0V6p.png)

The following code is working but very slow (especially if the arrays 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: It's quite easy to vectorize what you're doing:

    
    
    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 with 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 work with arrays of size `[nrows,ncols,nframes]`
simultaneously. The main trick is to make use of python's broadcasting, by
turning `CM_tilde` of size `[nrows,nframes]` into `CM_tilde[:,None,:]` of size
`[nrows,1,nframes]`. Python will then use the same values for each column,
since that is a singleton dimension of this modified `CM_tilde`.

By using `np.where` we choose (based on the `threshold`) whether we want to
get the corresponding value of `data`, or, again, the broadcast value of
`CM_tilde`. A new use of `np.mean` allows us to compute `all_CMs2`.

In the final step we made use of broadcasting by directly subtracting this new
`all_CMs2` from the corresponding elements of `data`.

It might help in vectorizing code this way by looking at the implicit indices
of your temporary variables. What I mean is that your temporary variable `CM`
lives inside a loop over `[nrows,nframes]`, and its value is reset with each
iteration. This means that `CM` is in effect a quantity `CM[row,frame]` (later
explicitly assigned to the 2d array `all_CMs`), and from here it's easy to see
that you can construct it by summing up an appropriate `CMtmp[row,col,frames]`
quantity along its column dimension. If it helps, you can name the
`np.where(...)` part as `CMtmp` for this purpose, and then compute
`np.mean(CMtmp,axis=1)` from that. Same result, obviously, but probably more
transparent.

