How to get better performance when taking the mean of a specific subset of a dataframe in pandas?

Question

I have a performance problem. The following code took 3 hours to loop through 5000 items out of 50 000.

I have a dataframe df, and a list of dictionary keys to loop through key_list. Each key corresponds to a single index of the dataframe. At each index, I want to get the mean of the columns mean_cols a few rows before and a few rows after the index, and then create a new dictionary with the before and after columns.

mean_cols = ['A', 'B', 'C']
rows_list = []
key_list = list(some_dict.keys()) # around 50k items

for key in key_list:

    means_after = df[mean_cols].iloc[key:key+5].mean()
    means_before = df[mean_cols].iloc[key-5:key].mean()

    for col in mean_cols:
        row_dict[str(col+'_after')] = round(means_after[col], 2)
        row_dict[str(col+'_before')] = round(means_before[col], 2)

    rows_list.append(row_dict)

I am pretty sure its the two lines,

means_after = df[mean_cols].iloc[key:key+5].mean()
means_before = df[mean_cols].iloc[key-5:key].mean()

however I can't think of a faster way to do it. Anyone have any ideas?

Answer 1

Have a look at the pandas rolling functionality. The docs list various windowing strategies and options to attempt to tackle your problem. You'll use the rolling method of your DataFrame with options suitable for your desired windows and then chain a mean call to get a resulting rolling mean DataFrame.

EDIT: Here is a simplified example to get you started. You'll want to review the output and make sure that's what you are going for, and this also has NaN values until the window size is full.

If I increase the columns to each have 50k items this takes 7 milliseconds on my laptop.

import pandas as pd

mean_cols = ["A", "B"]

df = pd.DataFrame({
    "A": range(0, 50000),
    "B": range(0, 50000)
})

rm_before = df[mean_cols].rolling(5).mean()
rm_after = rm_before.shift(-4)

for col in mean_cols:
    df[col+"_before"] = rm_before[col]
    df[col+"_after"] = rm_after[col]

print(df)

Answer 2

Pandas .mean() seems to have reputation to be slow.

An idea I would have is to use numpy by converting using pandas' built-in .to_numpy(). However, then if you want to have column-wise mean calculation, numpy's .mean() needs and axis specification - otherwise it will calculate mean of all values in the numpy array.

import pandas as pd
import numpy as np
import random

# from @totalhack
mean_cols = ["A", "B"]

df = pd.DataFrame({
    "A": range(0, 50000),
    "B": range(0, 50000)
})

key_list = random.sample(range(50000), k=50000)
# in case that key_list are rownames (indexes), convert them into
# row_indexes, because numpy array won't have names. E.g. by:
# my_rownames = [x for x in your_df_with_rownames.indexes]
# key_list = [my_rownames.index(k) for k in your_old_keylist]

df_mc = np.array(df[mean_cols])

rows_list = []

for key in keys_list:

    means_after = df_mc[key:key+5].mean(axis=0)
    means_before = df_mc[key-5:key].mean(axis=0)
    row_dict = {}

    for col in mean_cols:
        row_dict[str(col+'_after')] = round(means_after[mean_cols.index(col)], 2)
        row_dict[str(col+'_before')] = round(means_before[mean_cols.index(col)], 2)

    rows_list.append(row_dict)

If the data frame has only numeric values, it would accelerate calculations a lot more to convert it as early as possible to np.arrays. However, I guess there are text or date data too in the data frame. So the earliest time point to convert to numpy array is I guess directly after subsetting the mean_cols - so that is where I put .to_numpy().

Or using parallelization (using more cpus in parallel)

df_mc = np.array(df[mean_cols])

def mean_after(key, np_array=df_mc):
    return list(np.round(np_array[key: key+5].mean(axis=0), 2))

def mean_before(key, np_array=df_mc):
    return list(np.round(np_array[key-5:key].mean(axis=0), 2))


import multiprocessing as mp

pool = mp.Pool()

afters = pool.map(mean_after, keys_list)
befores = pool.map(mean_before, keys_list)

# for what do you need rows_list with dictionaires for each column value?
# why not accessing like this the afters or befores?

afters[row_idx][mean_cols.index(col)]
befores[row_idx][mean_cols.index(col)]