Question

这是我的功能：

def calculate_employment_two_digit_industry(df):
    df['intersection'] = df['racEmpProb'] * df['wacEmpProb']
    df['empProb'] = df['intersection'] / df['intersection'].sum()
    df['newEmp'] = df['empProb'] * df['Emp']

    df = df[['h_zcta', 'w_zcta', 'indID', 'newEmp', 'empProb']]
    df.rename(columns = {'newEmp' : 'Emp'}, inplace = True)

    return df

这是我的测试：

def test_calculate_employment_two_digit_industry():
    testDf = pandas.DataFrame({'h_zcta'     : [99163, 99163, 99163, 99163],
                           'w_zcta'     : [83843, 83843, 83843, 83843],
                           'indID'      : [11, 21, 22, 42],
                           'Emp'        : [20, 20, 40, 40],
                           'racEmpProb' : [0.5, 0.5, 0.6, 0.4],
                           'wacEmpProb' : [0.7, 0.3, 0.625, 0.375],
                           '1_digit'    : [1, 1, 2, 2]})

    expectedDf = pandas.DataFrame({'h_zcta'   : [99163, 99163, 99163, 99163],
                             'w_zcta'   : [83843, 83843, 83843, 83843],
                             'indID'    : [11, 21, 22, 42],
                             'Emp'      : [14, 6, 28.5716, 11.4284],
                             'empProb'  : [0.7, 0.3, 0.71429, 0.28571]})

    expectedDf = expectedDf[['h_zcta', 'w_zcta', 'indID', 'Emp', 'empProb']]

    final = testDf.groupby(['h_zcta', 'w_zcta', '1_digit'])\
               .apply(calculate_employment_two_digit_industry).reset_index()

    assert expected.equals(final)

正如你在测试中看到的那样，我有我希望函数返回的内容。除了我可以修复的代码中的潜在数学错误，这里是返回的数据帧，如何让它返回正常的数据帧（如果正常是正确的术语），即没有层只是列和行？

                      h_zcta  w_zcta  indID   Emp  empProb
h_zcta w_zcta 1_digit                                        
99163  83843  1       0   99163   83843     11  14.0      0.7
                      1   99163   83843     21   6.0      0.3
              2       0   99163   83843     22  28.0      0.7
                      1   99163   83843     42  12.0      0.3

提前谢谢。

Answer 1

您需要.reset_index(drop=True)

那是：

final = testDf.groupby(['h_zcta', 'w_zcta', '1_digit']).apply(
    calculate_employment_two_digit_industry).reset_index(drop=True)

>>> final.index
RangeIndex(start=0, stop=4, step=1)

Python groupby应用返回奇数数据帧

1 个答案: