Python: Binning one coordinate and averaging another based on these bins

Viewed 1936

I have two vectors rev_count and stars. The elements of those form pairs (let's say rev_count is the x coordinate and stars is the y coordinate).

I would like to bin the data by rev_count and then average the stars in a single rev_count bin (I want to bin along the x axis and compute the average y coordinate in that bin).

This is the code that I tried to use (inspired by my matlab background):

import matplotlib.pyplot as plt
import numpy

binwidth = numpy.max(rev_count)/10
revbin = range(0, numpy.max(rev_count), binwidth)
revbinnedstars = [None]*len(revbin)

for i in range(0, len(revbin)-1):
    revbinnedstars[i] = numpy.mean(stars[numpy.argwhere((revbin[i]-binwidth/2) < rev_count < (revbin[i]+binwidth/2))])

print('Plotting binned stars with count')
plt.figure(3)
plt.plot(revbin, revbinnedstars, '.')
plt.show()

However, this seems to be incredibly slow/inefficient. Is there a more natural way to do this in python?

3 Answers

The function I use for binning (x,y) data and determining summary statistics such as mean values in those bins is based upon the scipy.stats.statistic() function. I have written a wrapper for it, because I use it a lot. You may find this useful...

def binXY(x,y,statistic='mean',xbins=10,xrange=None):
    """
    Finds statistical value of x and y values in each x bin. 
    Returns the same type of statistic for both x and y.
    See scipy.stats.binned_statistic() for options.
    
    Parameters
    ----------
    x : array
        x values.
    y : array
        y values.
    statistic : string or callable, optional
        See documentation for scipy.stats.binned_statistic(). Default is mean.
    xbins : int or sequence of scalars, optional
        If xbins is an integer, it is the number of equal bins within xrange.
        If xbins is an array, then it is the location of xbin edges, similar
        to definitions used by np.histogram. Default is 10 bins.
        All but the last (righthand-most) bin is half-open. In other words, if 
        bins is [1, 2, 3, 4], then the first bin is [1, 2) (including 1, but 
        excluding 2) and the second [2, 3). The last bin, however, is [3, 4], 
        which includes 4.    
        
    xrange : (float, float) or [(float, float)], optional
        The lower and upper range of the bins. If not provided, range is 
        simply (x.min(), x.max()). Values outside the range are ignored.
    
    Returns
    -------
    x_stat : array
        The x statistic (e.g. mean) in each bin. 
    y_stat : array
        The y statistic (e.g. mean) in each bin.       
    n : array of dtype int
        The count of y values in each bin.
        """
    x_stat, xbin_edges, binnumber = stats.binned_statistic(x, x, 
                                 statistic=statistic, bins=xbins, range=xrange)
    
    y_stat, xbin_edges, binnumber = stats.binned_statistic(x, y, 
                                 statistic=statistic, bins=xbins, range=xrange)
    
    n, xbin_edges, binnumber = stats.binned_statistic(x, y, 
                                 statistic='count', bins=xbins, range=xrange)
            
    return x_stat, y_stat, n
Related