python - 获取下n个交易日功能太慢

Question

我创建了一个函数来生成下一个/前一个 n 个交易日，但对于我的目的来说太慢了。谁能建议一种加速该功能的方法？

def next_trading_day(start_day, num_trading_days, direction):
    '''returns the next/previous trading day. Business_days determines how many days
    back or into the future, direction determines whether back (-1) or forward (1)'''
    for i in range(0, num_trading_days, direction):
        next_day = start_day +datetime.timedelta(days=direction)
        while next_day.weekday() in [5,6] or next_day not in mcal.get_calendar('NYSE').valid_days(start_date='2000-12-20', end_date='2020-01-10'):
            next_day +=  datetime.timedelta(days=direction)
        start_day = next_day
    return start_day

这是我使用该功能的方式:

import pandas as pd
dict1 = [
        {'date': '2016-11-27'},
        {'date': '2016-11-28'},
{'date': '2016-11-27'},
]
df1= pd.DataFrame(dict1)
df1['date']      = pd.to_datetime(df1['date'])

df['Date-1']=df['date'].dt.date.apply(next_business_day, args=[-1,-1,])

Answer 1

此检查 next_day not in mcal.get_calendar('NYSE').valid_days(start_date='2000-12-20', end_date='2020-01-10') 非常及时 -消耗，因为它需要从超过 7000 天的数组中查找。而且你需要对每一个操作都执行此操作，因此我认为这是效率低下的主要原因。

您可以通过将 mcal.get_calendar('NYSE').valid_days(start_date='2000-12-20', end_date='2020-01-10') 转换为设置，这会将查找时间从 O(N) 降低到 O(log N)。

但我会选择另一种策略:

1. 创建一个将每个交易日与其下一个和/或最后一个交易日相匹配的表格
2. 将上表与您的数据中的日期合并
3. 估算缺失值
4. 将新创建的表与原始数据合并

编辑:允许任意数量的滞后和超前

import pandas as pd
import pandas_market_calendars as mcal

def get_next_trading_day(df1, n):
    trading_days = pd.DataFrame({"date": mcal.get_calendar('NYSE').valid_days(start_date='2016-11-10', end_date='2016-12-01')})
    trading_days['date'] = trading_days['date'].dt.tz_convert(None)
    trading_days = trading_days[~trading_days.date.dt.weekday.isin([5,6])]
    trading_days['next_trading_day'] = trading_days.date.shift(-n)
    # extract unique  date from df1
    df2 = pd.DataFrame({"date": pd.unique(df1['date'])})

    # merge with the trading days data (non-trading day will have NA fields)
    df2 = df2.merge(trading_days, on='date', how='outer')

    # impute NA values
    df2.sort_values(by='date', inplace=True)

    df2['next_trading_day'].fillna(method=  'ffill' if n>0 else 'bfill', inplace=True)

    return df1.merge(df2, on='date', how='left')

dict1 = [
        {'date': '2016-11-27'},
        {'date': '2016-11-28'},
        {'date': '2016-11-27'},
        ]
df1= pd.DataFrame(dict1)
df1['date']      = pd.to_datetime(df1['date'])


print("Next trading day")
print(get_next_trading_day(df1, 1))
print()

print("Previous trading day")
print(get_next_trading_day(df1, -1))
print()

print("Next next trading day")
print(get_next_trading_day(df1, 2))
print()

print("Previous previous trading day")
print(get_next_trading_day(df1, -2))
print()

输出

Next trading day
        date next_trading_day
0 2016-11-27       2016-11-28
1 2016-11-28       2016-11-29
2 2016-11-27       2016-11-28

Previous trading day
        date next_trading_day
0 2016-11-27       2016-11-25
1 2016-11-28       2016-11-25
2 2016-11-27       2016-11-25

Next next trading day
        date next_trading_day
0 2016-11-27       2016-11-29
1 2016-11-28       2016-11-30
2 2016-11-27       2016-11-29

Previous previous trading day
        date next_trading_day
0 2016-11-27       2016-11-23
1 2016-11-28       2016-11-23
2 2016-11-27       2016-11-23