熊猫群加速由

Question

我在高性能计算中使用熊猫数据帧。这个函数是一个显著时间片：

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