加快 pandas 数组计算

Question

我有实现所需计算结果的工作代码，但我目前正在使用一种迭代 pandas 数组的算法。这显然比纯 pandas DataFrame 计算要慢。想要一些关于如何使用 pandas 函数来加快计算速度的建议

生成虚拟数据的代码

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

这基本上是一个带有 MTD 数据的 pandas DataFrame，具有一定的价值。这纯粹是为了让我们有一些数据可以使用。

需要计算

我需要的是一个新的 DataFrame，它将开始（投资）日期作为列 - 用一些月初值填充它们。索引是所有可能的日期，值应该是 YTD 数字。我正在使用此 Dataframe 作为投资日期的查找/缓存

伪代码

YTD = (1+最后一个 MTD 数字) * ((1+最后一个 MTD 数字)... 直到所需日期的所有月份

工作功能

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

我已硬编码投资日期列表以简化问题陈述。在我的机器上，这段代码需要 2.5 到 3.5 秒。关于如何加快速度的任何建议？

Answer 1

这是一种应该相当快的方法。很可能有更快/更清洁的东西，但这应该是一种改进。

#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.)