Python实现二维分箱后按箱返回x、y、v数组的方法

问题描述 #

我有x、y、v三个数组，要在x-y平面对v进行分箱操作，希望分箱后能按每个箱分别提取对应的x、y、v数组，但当前代码只能单独获取单个箱的数据，没法适配多箱的大数据集，循环知识薄弱不知道怎么实现。原代码如下：

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])

# 注意：原代码里的values应该是v，这里先修正
ret = stats.binned_statistic_2d(x, 
                                y,
                                v,  # 修正原代码的values为v
                                '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)

当前代码只能输出计数为5的单个箱数据，需要获取全部4个箱的对应数组，比如bin1对应x_bin1=[...]、y_bin1=[...]、v_bin1=[...]，以此类推。同时希望有更简便的分箱方法。

解决方案：循环遍历所有箱 #

利用binnumber中每个数据点对应的x/y箱号，遍历所有可能的箱组合，批量提取每个箱的数据：

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])

# 执行分箱，修正values为v
ret = stats.binned_statistic_2d(x, 
                                y,
                                v,
                                'count',
                                bins=2,
                                expand_binnumbers=True)

binnumber = ret.binnumber  # shape (2, n_samples)，第一行是x的箱号，第二行是y的箱号
x_bins_count = ret.statistic.shape[0]
y_bins_count = ret.statistic.shape[1]

# 存储所有箱的数据，用字典方便索引
bins_data = {}

# 遍历所有x箱号和y箱号（箱号从1开始）
for x_bin in range(1, x_bins_count + 1):
    for y_bin in range(1, y_bins_count + 1):
        # 找到同时属于当前x箱和y箱的数据点索引
        mask = np.logical_and(binnumber[0] == x_bin, binnumber[1] == y_bin)
        # 提取对应数据
        bins_data[f'bin_{x_bin}_{y_bin}'] = {
            'x': x[mask],
            'y': y[mask],
            'v': v[mask]
        }

# 打印所有箱的数据
for bin_name, data in bins_data.items():
    print(f'\n{bin_name}:')
    print(f'x: {data["x"]}')
    print(f'y: {data["y"]}')
    print(f'v: {data["v"]}')

这段代码会自动遍历所有箱，把每个箱的x/y/v数据存在字典里，键名是bin_x_y（x和y是对应的箱号），直接通过键就能获取对应箱的数据。

更简便的方法：用Pandas分组分箱 #

如果用Pandas处理会更直观，代码更简洁：

import pandas as pd
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])

# 转为DataFrame
df = pd.DataFrame({'x': x, 'y': y, 'v': v})

# 对x和y分箱，bins=2表示分成2个箱
df['x_bin'] = pd.cut(df['x'], bins=2, labels=False) + 1  # 箱号从1开始
df['y_bin'] = pd.cut(df['y'], bins=2, labels=False) + 1

# 按x_bin和y_bin分组
grouped = df.groupby(['x_bin', 'y_bin'])

# 遍历每个分组，获取数据
for (x_bin, y_bin), group in grouped:
    print(f'\nbin_{x_bin}_{y_bin}:')
    print(f'x: {group["x"].values}')
    print(f'y: {group["y"].values}')
    print(f'v: {group["v"].values}')

Pandas的cut函数直接完成分箱，groupby之后就能批量处理每个箱的数据，逻辑更清晰，代码量更少，适合大数据集处理。

内容的提问来源于stack exchange，提问作者Jerome