python - 生成具有数字范围的条件组合的数据框

Question

我正在尝试计算以下公式:

我写这个是为了计算这个:

%%timeit
df = pd.DataFrame(columns=["h","i","j","k"])
x1=10
y1=8
m=7
P_List = []
m_range = np.arange(0,m+1)

for m in m_range:
        df = pd.DataFrame(columns=["h","i","j","k"])
        h_range = np.arange(0,x1+1)
        for h in h_range:
            if (x1-h)>=0:
                i_range = np.arange(0,x1-h+1)
            else:
                break
            for i in i_range:
                if (y1-(m-i))>=0:
                    k_range = np.arange(0,y1-(m-i)+1)
                else:
                    break
                j = m-i
                if j>=0:
                    for k in k_range:
                        arguments = { "h":h, "i":i, "j":j, "k":k}
                        df = df.append(arguments, ignore_index=True)
        df["P"]= df.apply(lambda x: cal_P(x), axis=1)
        P_List.append(df.P.sum())         

输出:每次循环 5 s ± 184 ms(7 次运行的平均值 ± 标准差，每次 1 次循环)

另一种方法:

%%timeit
df = pd.DataFrame(columns=["h","i","j","k"])
x1=10
y1=8
m=7
P_List = []
m_range = np.arange(0,m+1)

for m in m_range:
    h_range = np.arange(0,x1+1) 
    i_range = np.arange(0,x1+1)
    k_range = np.arange(0,y1+x1+2)
    j_range = np.arange(0,x1+m+1)

    data = np.array(np.meshgrid( h_range, i_range, k_range, j_range)).T.reshape(-1,4)
    df = pd.DataFrame(data=data, columns=[ "h","i","j","k"])

    i_cond = df.i<=(x1-df.h) #(x1-h)
    k_cond = df.k<=(y1-(m-df.i)) #(y1-(m-i))
    j_cond = df.j == (m-df.i)
    df = df.drop(df[~(i_cond&k_cond&j_cond)].index)
    df["P"]= df.apply(lambda x: cal_P(x), axis=1)
    P_List.append(df.P.sum())

输出:每次循环 306 ms ± 5.75 ms(7 次运行的平均值 ± 标准差，每次 1 次循环)

现在的问题是 x1、y1 和 m 的值对于我的数据集来说非常大。对于最大的数据集，x1≈2000，y1≈1800，m≈500。所以第一个方法永远运行，而第二个方法内存不足。

有没有办法在不耗尽内存的情况下更快地完成此操作？或者还有其他更好的方法来计算上述公式吗？

编辑:

P 计算 m 值的超几何分布。我已经更新了上面的代码以反射(reflect)这一点。 在哪里， 3 元素三项函数 (a,b,c)!定义为:

用于计时cal_P是

def cal_P(row):
    l = row.h+row.i+row.j+row.k
    return l

但是计算超几何分布的实际代码是:

from math import factorial as fact
def t_func(a,b,c):
    d=a-b-c
    if d>=0 and a>=0 and b>=0 and c>=0:
        result = fact(a)/(fact(b)*fact(c)*fact(d))     
    else:
        result = 0
    return result 

def hypergeom_XY(h, i, j, k, x1, x2, y1, y2, l):
    pmf = t_func(x1,h,i)*t_func(y1,j,k)*t_func(l-x1-y1,x2-h-j,y2-i-k)/t_func(l,x2,y2)
    return pmf

def cal_P(row):
   P = hypergeom_XY(row.h, row.i, row.j, row.k, x1, x2, y1, y2, l)
   return P
 

Answer 1

我将生成 df 的第二个代码示例转换为以下函数:

def f1(x1, y1, m):
    m_range = np.arange(0, m  + 1)
    h_range = np.arange(0, x1 + 1)
    i_range = np.arange(0, x1 + 1)
    k_range = np.arange(0, y1 + x1 + 2)
    j_range = np.arange(0, x1 + m  + 1)
    data = np.array(np.meshgrid(m_range, h_range, i_range,
        j_range, k_range)).T.reshape(-1, 5)
    df = pd.DataFrame(data=data, columns=['m', 'h', 'i', 'j', 'k'])
    i_cond = df.i <= (x1-df.h)
    k_cond = df.k <= (y1-(df.m-df.i))
    j_cond = df.j == (df.m-df.i)
    rowNo1 = df.index.size
    df = df.drop(df[~(i_cond & k_cond & j_cond)].index)
    rowNo2 = df.index.size
    return df, rowNo1, rowNo2

请注意，它返回 3 个结果:

• df - DataFrame 本身，
• rowNo1 - drop 之前的行数，
• rowNo2 - drop 后的行数。

我对你的数据执行了它:

df, rN1, rN2 = f1(10, 8, 7)
print(f'{rN1}, {rN2}')

获取:

348480, 2010

请注意，rowNo1 相当大。

另请注意，meshgrid 中包含的 j_range 可以省略， 那么行数就会小得多。

名为j的列可以在df生成后计算， 使用df.m - df.i公式。

然后可以消除“不必要的”行，这次使用 df.m >= df.i 公式。

所以我的建议是如何生成源DataFrame:

def f2(x1, y1, m):
    m_range = np.arange(0, m  + 1)
    h_range = np.arange(0, x1 + 1) 
    i_range = np.arange(0, x1 + 1)
    k_range = np.arange(0, y1 + x1 + 2)
    data = np.array(np.meshgrid(m_range, h_range, i_range,
        k_range)).T.reshape(-1, 4)
    df = pd.DataFrame(data=data, columns=["m", "h", "i", "k"])
    df.insert(loc=3, column='j', value=df.m - df.i)
    i_cond = df.i <= (x1-df.h) 
    k_cond = df.k <= (y1-(df.m-df.i))
    j_cond = df.m >= df.i
    rowNo1 = df.index.size
    df = df.drop(df[~(i_cond & k_cond & j_cond)].index)
    rowNo2 = df.index.size
    return df, rowNo1, rowNo2

为了进行比较，我仍然计算drop之前和之后的行号。

然后我执行:

df2, rN1, rN2 = f2(10, 8, 7)
print(f'{rN1}, {rN2}')

获取:

19360, 2010

请注意，这次 rowNo1 明显更小，因此有一个更好的 该函数不会耗尽可用内存的可能性更大 x1、y1 和 m 的值。

为了检查两个结果(df 和 df2)是否相同，我运行了:

df.sort_values(['m', 'h', 'i', 'j', 'k'], ignore_index=True).equals(
    df2.sort_values(['m', 'h', 'i', 'j', 'k'], ignore_index=True))

得到正确。

最后一个重要因素是执行速度。 使用 %timeit 我得到:

• 120 毫秒您的代码 (f1)，
我的代码 (f2)
• 7.69 毫秒 - 快了 15 倍以上。

在您的真实数据上尝试我的函数，而不返回两个行号。