大熊猫在多个条件下的Bin平均

Question

bounty到期在7天内。该问题的答案有资格获得+ 250声望赏金。Fred S想要引起更多关注对此问题。

问题

我有一个目标变量x以及一些其他变量A和B。当要满足x和A的某些条件时，我想计算B的平均值（以及其他统计信息）。一个现实的例子是，当太阳辐射（x）和风速（A）落入某些预定义的bin范围内时，通过一连串的测量值来计算平均气温（B）。

潜在的解决方案

我已经能够通过循环来完成此操作（请参见下面的示例），但是我了解到我应该避免在数据帧上循环。从我对该站点的研究中，我觉得可能存在使用pd.cut或np.select的更为优雅的矢量化解决方案，但坦率地说，我不知道该怎么做。

示例

生成样本数据

import pandas as pd
import numpy as np

n = 100
df = pd.DataFrame(
    {
        "x": np.random.randn(n),
        "A": np.random.randn(n)+5,
        "B": np.random.randn(n)+10
    }
)

[df.head()输出：

    x           A           B
0   -0.585313   6.038620    9.909762
1   0.412323    3.991826    8.836848
2   0.211713    5.019520    9.667349
3   0.710699    5.353677    9.757903
4   0.681418    4.452754    10.647738

计算箱平均数

# define bin ranges
bins_A = np.arange(3, 8)
bins_B = np.arange(8, 13)

# prepare output lists
A_mins= []
A_maxs= []
B_mins= []
B_maxs= []
x_means= []
x_stds= []
x_counts= []

# loop over bins
for i_A in range(0, len(bins_A)-1):
    A_min = bins_A[i_A]
    A_max = bins_A[i_A+1]
    for i_B in range(0, len(bins_B)-1):
        B_min = bins_B[i_B]
        B_max = bins_B[i_B+1]

        # binning conditions for current step
        conditions = np.logical_and.reduce(
            [
                df["A"] > A_min,
                df["A"] < A_max,
                df["B"] > B_min,
                df["B"] < B_max,
            ]
        )

        # calculate statistics for x and store values in lists
        x_values = df.loc[conditions, "x"]
        x_means.append(x_values.mean())
        x_stds.append(x_values.std())
        x_counts.append(x_values.count())

        A_mins.append(A_min)
        A_maxs.append(A_max)
        B_mins.append(B_min)
        B_maxs.append(B_max)

将结果存储在新的数据框中

binned = pd.DataFrame(
    data={
        "A_min": A_mins,
        "A_max": A_maxs,
        "B_min": B_mins,
        "B_max": B_maxs,
        "x_mean": x_means,
        "x_std": x_stds,
        "x_count": x_counts 
        }
)

[binned.head()输出：

    A_min   A_max   B_min   B_max   x_mean      x_std       x_count
0   3       4       8       9       0.971624    0.790972    2
1   3       4       9       10      0.302795    0.380102    3
2   3       4       10      11      0.447398    1.787659    5
3   3       4       11      12      0.462149    1.195844    2
4   4       5       8       9       0.379431    0.983965    4

Answer 1

如果您担心的是关于[[性能，则可以使用for循环，如果您使用numba，则可以进行较小的更改

这里有一个执行计算的函数。关键是calculate使用numba，因此速度非常快。其余仅用于创建熊猫数据框：

from numba import njit

def calc_numba(df, bins_A, bins_B):
    """ wrapper for the timeit. It only creates a dataframe """

    @njit
    def calculate(A, B, x, bins_A, bins_B):

        size = (len(bins_A) - 1)*(len(bins_B) - 1)
        out = np.empty((size, 7))

        index = 0
        for i_A, A_min in enumerate(bins_A[:-1]):
            A_max = bins_A[i_A + 1]

            for i_B, B_min in enumerate(bins_B[:-1]):
                B_max = bins_B[i_B + 1]

                mfilter = (A_min < A)*(A < A_max)*(B_min < B)*(B < B_max)
                x_values = x[mfilter]

                out[index, :] = [
                    A_min,
                    A_max,
                    B_min,
                    B_max,
                    x_values.mean(),
                    x_values.std(),
                    len(x_values)
                ]

                index += 1

        return out

    columns = ["A_min", "A_max", "B_min", "B_max", "mean", "std", "count"]
    out = calculate(df["A"].values, df["B"].values, df["x"].values, bins_A, bins_B)
    return pd.DataFrame(out, columns=columns)

性能测试
使用n = 1_000_000以及相同的bins_A和bins_B，我们得到：
%timeit code_question(df, bins_A, bins_B)
15.7 s ± 428 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit calc_numba(df, bins_A, bins_B)
507 ms ± 12.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

比问题代码快[<30

由于pandas内置方法使用了类似的增强功能，因此很难击败numba的性能。