pandas groupby 的加速

Question

我在高性能计算中使用 pandas DataFrame。这个函数是一个重要的时间槽：

def calculate_with_n_electron(self, phi, partition_function,
                              number_density, n_electron):
    ion_populations = pd.DataFrame(data=0.0,
        index=partition_function.index.copy(),
        columns=partition_function.columns.copy(), dtype=np.float64)

    for atomic_number, groups in phi.groupby(level='atomic_number'):
        current_phis = (groups / n_electron).replace(np.nan, 0.0).values
        phis_product = np.cumproduct(current_phis, axis=0)
        neutral_atom_density = (number_density.ix[atomic_number] /
                                (1 + np.sum(phis_product, axis=0)))
        ion_populations.ix[atomic_number, 0] = (
            neutral_atom_density.values)
        ion_populations.ix[atomic_number].values[1:] = (
            neutral_atom_density.values * phis_product)
        ion_populations[ion_populations < self.ion_zero_threshold] = 0.0
    return ion_populations

这是更大范围内的函数：https://github.com/tardis-sn/tardis/blob/master/tardis/plasma/properties/ion_population.py#L151

任何帮助将不胜感激！

Answer 1

在不知道数据是什么样子的情况下，这不太可能完全有效，但希望能给您一些想法 - 基本方法是避免 for 循环并使用矢量化操作进行所有操作。

gb = phi.groupby(level='atomic_number')

# do this outside the groupby, use fillna instead of replace
phi = (phi / n_electron).fillna(0.0)
phi['product'] = gb.cumprod()

# assume number_density has one column named 'density`
phi = phi.join(number_density)
phi['density'] = phi['density'] / (1 + gb['product'].transform('sum'))

# bit of a hack to exclude the first element from each group
# from the multiplication
phi['dummy'] = 1
phi['density'] = df['density'] * np.where(gb['dummy'].cumsum() == 1, 1, df['product'])

phi.loc[phi['density'] < self.ion_zero_threshold] = 0.0