Jerome
Jerome

Reputation: 49

returning elements in bins as arrays in python

I have x,y,v arrays of data points and I am binning v on x-y plane. I am trying to get the x,y,v values back after binning but I want them as arrays corresponding to each bin. My code can get them individually but that will not work for large data sets with many bins. Maybe I need to use loops of some kind but my understanding of loops is weak. Code:

from scipy import stats
import numpy as np


x=np.array([-10,-2,4,12,3,6,8,14,3])
y=np.array([5,5,-6,8,-20,10,2,2,8])
v=np.array([4,-6,-10,40,22,-14,20,8,-10])

ret = stats.binned_statistic_2d(x, 
                                y,
                                values,
                                'count',
                                bins=2,
                                expand_binnumbers=True)

print('counts=',ret.statistic)
print('binnumber=', ret.binnumber)

binnumber = ret.binnumber
statistic = ret.statistic

# get the bin numbers according to some condition
idx_bin_x, idx_bin_y =  np.where(statistic==statistic[1][1])#[0]
print('idx_binx=',idx_bin_x)
print('idx_bin_y=',idx_bin_y)

# A binnumber of i means the corresponding value is 
# between (bin_edges[i-1], bin_edges[i]).
# -> increment the bin indices by one 

idx_bin_x += 1
idx_bin_y += 1

print('idx_binx+1=',idx_bin_x)
print('idx_bin_y+1=',idx_bin_y)

#  get the boolean mask and apply it
is_event_x = np.in1d(binnumber[0], idx_bin_x)
print('eventx=',is_event_x)
is_event_y = np.in1d(binnumber[1], idx_bin_y)
print('eventy=',is_event_y)
is_event_xy = np.logical_and(is_event_x, is_event_y)
print('event_xy=', is_event_xy)


events_x = x[is_event_xy]
events_y = y[is_event_xy]
event_v=v[is_event_xy]

print('x=', events_x)
print('y=', events_y)
print('v=',event_v)

This outputs x,y,v for the bin with count=5 but I want all 4 bins returning 4 arrays for each x,y,v. eg for bin1: x_bin1=[...], y_bin1=[...], v_bin1=[...] and so on for 4 bins.

Also, feel free to suggest if you think there are easier ways to bin 2d planes (x,y) with values (v) like mine and getting binned values. Thank you!

Upvotes: 0

Views: 310

Answers (1)

Chris Seeling
Chris Seeling

Reputation: 656

Using np.array facilitates a compact way to recover the arrays you are after:

from scipy import stats

# coordinates
x = np.array([-10,-2,4,12,3,6,8,14,3])
y = np.array([5,5,-6,8,-20,10,2,2,8])
v = np.array([4,-6,-10,40,22,-14,20,8,-10])

ret = stats.binned_statistic_2d(x, y, None, 'count', bins=2, expand_binnumbers=True)
b = ret.binnumber
for i in [1,2]:
    for j in [1,2]:
        m = (b[0] == i) & (b[1] == j)    # mask
        print((list(x[m]),list(y[m]),list(v[m])))

which gives for each of the four bins a tuple of 3 lists corresponding to x, y and v values:

([], [], [])
([-10, -2], [5, 5], [4, -6])
([4, 3], [-6, -20], [-10, 22])
([12, 6, 8, 14, 3], [8, 10, 2, 2, 8], [40, -14, 20, 8, -10])

Upvotes: 1

Related Questions