Joop
Joop

Reputation: 8108

Speeding up pandas array calculation

I have working code that achieves the desired calculation result, but I am currently using an algorithm that iterates over the pandas array. this is obviously slower than pure pandas DataFrame calculations. Would like some advice on how i can use pandas functions to speed up this calculation

Code to generate dummy data

df = pd.DataFrame(index=pd.date_range(start='2014-01-01', periods=365))
df['Month'] = df.index.month
df['MTD'] = (df.index.day+0.001)/10000

This is basically a pandas DataFrame with MTD figures for some value. This is purely given so that we have some data to play with.

Needed calculation

what I need is a new DataFrame that has starting (investment) dates as columns - populating them with a few beginning of month values. the index is all possible dates and the values should be the YTD figure. I am using this Dataframe as a lookup/cache for investement dates

pseudocode

YTD = (1+last MTD figure) * ((1+last MTD figure)... for all months to the required date

Working function

def calculate_YTD(df):  # slow takes 3.5s on my machine!!!!!!
    YTD_df = pd.DataFrame(index=df.index)

    for investment_date in [datetime.datetime(2014,x+1,1) for x in range(12)]:  
        YTD_df[investment_date] =1.0  # pre-populate with dummy floats
        for date in df.index:  # iterate over all dates in period
            h = (df[investment_date:date].groupby('Month')['MTD'].max().fillna(0) + 1).product() -1
            YTD_df[investment_date][date] = h
    return YTD_df

I have hardcoded the investment dates list to simplify the problem statement. On my machines this code takes 2.5 to 3.5 seconds. Any suggestions on how i can speed it up?

Upvotes: 0

Views: 392

Answers (1)

chrisb
chrisb

Reputation: 52236

Here's an approach that should be reasonably quick. Quite possible there is something faster/cleaner, but this should be an improvement.

#assuming a fixed number of investments dates, build a list
investment_dates = pd.date_range('2014-1-1', periods=12, freq='MS')

#build a table, by month, which contains the cumulative MTD
#return for each invesment date.  Still have to loop over the investment dates,
#but don't need to loop over each daily value
running_mtd = []
for date in investment_dates:
    curr_mo = (df[df.index >= date].groupby('Month')['MTD'].last() + 1.).cumprod()
    curr_mo.name = date
    running_mtd.append(curr_mo)
running_mtd_df = pd.concat(running_mtd, axis=1)
running_mtd_df = running_mtd_df.shift(1).fillna(1.)

#merge running mtd returns with base dataframe
df = df.merge(running_mtd_df, left_on='Month', right_index=True)


#calculate ytd return for each column / day, by multipling the running 
#monthly return with the current MTD value
for date in investment_dates:
    df[date] = np.where(df.index < date, np.nan, df[date] *  (1. + df['MTD']) - 1.)

Upvotes: 1

Related Questions