CODEWITHSUNDEEP
pandas
Wolfgang Kerzendorf
Wolfgang Kerzendorf

Reputation: 734

Speedup of pandas groupby

I'm using the pandas DataFrame in high performance calculations. This function is a significant time sink:

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

here is the function within the bigger context: https://github.com/tardis-sn/tardis/blob/master/tardis/plasma/properties/ion_population.py#L151

Any help would be appreciated!

Upvotes: 1

Views: 1059

Answers (1)

chrisb
chrisb

Reputation: 52246

Without knowing what the data looks like it's very unlikely this works exactly, but hopefully gives you some ideas - basic approach is to avoid the for loop and do everything with vectorized operations.

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

Upvotes: 2

Related Questions