在.groupby()和.apply()之后在DataFrame中添加新列太慢

时间:2020-09-23 06:39:18

标签: python pandas numpy

我正在尝试使用.groupby()、. apply()和自定义函数创建新列。但是,我不知道如何获得良好的性能。

我尝试了下面的代码,该代码不起作用,“ distance_pref”列中充满了NaN值。

def prefVar(results):
    results["distance_pref"] = results\
           .sort_values(by=['Name', 'Date'])\
           .groupby(['Name'])\
           .apply(lambda x: distance_calc(x.Distance.to_numpy().astype(np.float64),
                                          x.Intercept.to_numpy().astype(np.float64),
                                          x.Coef.to_numpy().astype(np.float64)))
    return results

@numba.jit(nopython=True)
def distance_calc(dist, inter, coef):
    assert dist.shape == inter.shape == coef.shape, "Inputs mismatched"
    preference = np.zeros(dist.shape[0])
    for i in numba.prange(1, dist.shape[0]):
        preference[i] = inter[i - 1] + (coef[i - 1] * dist[i])
    return preference

仔细检查结果。groupby()。apply(lambda x:distance_calc(...))返回下面的DataFrame。我不完全确定如何以正确的格式获取它,然后将其合并回原始结果DataFrame中:

NAME
NAME1      [0.0, 0.0, 0.0]
NAME2      [0.0, 0.0, 0.0, 0.0, -0.16666666666666674, -0....]
NAME3      [0.0, 0.0, 0.0, 0.0, -3.4857142857143106, -0.0...]
NAME4      [0.0, 0.0, 0.0, 0.0, -0.15167326376394402, -0....]
NAME5      [0.0, 0.0, 0.0, 0.0, -0.0704500978473579, 0.09...]
                                         ...                        
NAME10000  [0.0, 0.0, 0.0, 0.0, 0.1481644757577409, 0.115...]
NAME10001  [0.0, 0.0, 0.0, 0.0, -0.08333333333333333, 0.0...]
NAME10002  [0.0, 0.0, 0.0, 0.0, -0.08270676691729384, 0.0...]
NAME10003  [0.0, 0.0, 0.0, 0.0, 0.057142857142858605, -0....]

我确实使下面的一些代码正常工作,但是速度很慢。

def prefVar(results):
    return results \
        .sort_values(by=['Name', 'Date']) \
        .groupby(['Name']) \
        .apply(lambda g: distance_pref_on_group(g))


def distance_pref_on_group(group):   
    group["distance_preference"] = distance_calc(group.Distance.to_numpy().astype(np.float64),
                                                 group.Intercept.to_numpy().astype(np.float64),
                                                 group.Coef.to_numpy().astype(np.float64))
    return group

任何有关如何加快上述代码速度或采用更好方法的建议将不胜感激。

0 个答案:

没有答案