Averaging out sections of a multiple row array in Python
Question: I've got a 2-row array called C like this:

    
    
    from numpy import *
    A = [1,2,3,4,5]
    B = [50,40,30,20,10]
    C = vstack((A,B))
    

I want to take all the columns in C where the value in the first row falls
between i and i+2, and average them. I can do this with just A no problem:

    
    
    i = 0
    A_avg = []
    
    while(i<6):
        selection = A[logical_and(A >= i, A < i+2)] 
        A_avg.append(mean(selection))
        i += 2
    

then A_avg is:

    
    
    [1.0,2.5,4.5]
    

I want to carry out the same process with my two-row array C, but I want to
take the average of each row separately, while doing it in a way that's
dictated by the first row. For example, for C, I want to end up with a 2 x 3
array that looks like:

    
    
    [[1.0,2.5,4.5],
     [50,35,15]]
    

Where the first row is A averaged in blocks between i and i+2 as before, and
the second row is B averaged in the same blocks as A, regardless of the values
it has. So the first entry is unchanged, the next two get averaged together,
and the next two get averaged together, for each row separately. Anyone know
of a clever way to do this? Many thanks!


Answer: I hope this is not too clever. TIL boolean indexing does not broadcast, so I
had to manually do the broadcasting. Let me know if anything is unclear.

    
    
    import numpy as np
    A = [1,2,3,4,5]
    B = [50,40,30,20,10]
    C = np.vstack((A,B)) # float so that I can use np.nan
    
    i = np.arange(0, 6, 2)[:, None]
    selections = np.logical_and(A >= i, A < i+2)[None]
    
    D, selections = np.broadcast_arrays(C[:, None], selections)
    D = D.astype(float)     # allows use of nan, and makes a copy to prevent repeated behavior
    D[~selections] = np.nan # exclude these elements from mean
    
    D = np.nanmean(D, axis=-1)
    

Then,

    
    
    >>> D
    array([[  1. ,   2.5,   4.5],
           [ 50. ,  35. ,  15. ]])
    

Another way, using `np.histogram` to bin your data. This may be faster for
large arrays, but is only useful for few rows, since a hist must be done with
different weights for each row:

    
    
    bins = np.arange(0, 7, 2)     # include the end
    n = np.histogram(A, bins)[0]  # number of columns in each bin
    a_mean = np.histogram(A, bins, weights=A)[0]/n
    b_mean = np.histogram(A, bins, weights=B)[0]/n
    D = np.vstack([a_mean, b_mean])
    

