Как быстро получить равномерно распределенные данные с помощью мультииндекса в pandas

#python #pandas

Вопрос:

У меня есть фрейм данных, индексированный биржевым тикером и датой, который довольно скуден, что-то вроде:

 df = pd.DataFrame({
    'ticker': ['SPY', 'GOOGL', 'GOOGL', 'TSLA', 'TSLA'],
    'date': ['2021-01-01', '2021-09-01', '2021-09-21', '2021-09-21', '2021-09-22'],
    'price': [430.0, 2500.0, 2600.0, 700.0, 710.0],
}).astype({'date': 'datetime64[ns]'}).set_index(['ticker', 'date'])
                    price
ticker  date    
SPY     2021-01-01  430.0
GOOGL   2021-09-01  2500.0
        2021-09-21  2600.0
TSLA    2021-09-21  700.0
        2021-09-22  710.0

 

Я хочу, чтобы в итоге получился кадр данных, содержащий данные за последние три дня, насколько нам известно, например,

 want = pd.DataFrame({
    'ticker': ['SPY', 'SPY', 'SPY', 'GOOGL', 'GOOGL', 'GOOGL', 'TSLA', 'TSLA', 'TSLA'],
    'date': ['2021-09-20', '2021-09-21', '2021-09-22', '2021-09-20', '2021-09-21', '2021-09-22', '2021-09-20', '2021-09-21', '2021-09-22'],
    'price': [430.0, 430.0, 430.0, 2500.0, 2600.0, 2600.0, 0.0, 700.0, 710.0],
}).astype({'date': 'datetime64[ns]'}).set_index(['ticker', 'date'])
                    price
ticker  date    
SPY     2021-09-20  430.0
        2021-09-21  430.0
        2021-09-22  430.0
GOOGL   2021-09-20  2500.0
        2021-09-21  2600.0
        2021-09-22  2600.0
TSLA    2021-09-20  0.0
        2021-09-21  700.0
        2021-09-22  710.0
 

Я придумал пару способов сделать это, до сих пор я думаю, что самый ясный-это группа с пользовательским приложением, т. Е.,

 OUTPUT_DATES = pd.date_range(
    start=pd.Timestamp.today() - pd.DateOffset(days=2),
    end=pd.Timestamp.today(),
    freq='D')

def LastNDays(df):
  return (
      df
      .reset_index(level=0, drop=True)
      .reindex(OUTPUT_DATES, method='ffill')
      .rename_axis('date')
      .fillna(0))
  
df.groupby(level=0).apply(LastNDays)
 

И это работает. Однако это также очень медленно для моего фактического набора данных (несколько сотен тысяч точек данных). Я думаю, что все дело в переиндексировании? Это кажется довольно распространенной задачей для панд (возьмите некоторые странные данные о запасах, приведите их в соответствие с точками данных), поэтому я чувствую, что, вероятно, есть лучший способ сделать это, но я даже не знаю, что искать. Есть какие-нибудь идеи о том, как сделать это быстрее?

Ответ №1:

Вы, вероятно, увидите огромное улучшение производительности при использовании asof слияния. Сначала создайте все необходимые строки из декартова произведения уникальных меток тикера и последних трех дат. Затем выполните слияние, чтобы перенести ближайшее значение (на ту же дату или в прошлом), сконструировать и отсортировать мультииндекс и заполнить отсутствующие значения 0.

 import pandas as pd

dates = pd.date_range(pd.to_datetime('today').normalize(), freq='-1D', periods=3)
#DatetimeIndex(['2021-09-22', '2021-09-21', '2021-09-20'], dtype='datetime64[ns]', freq='-1D')    

df1 = pd.DataFrame(product(dates, df.index.get_level_values('ticker').unique()),
                   columns=['date', 'ticker'])

result = (pd.merge_asof(df1.sort_values('date'), df.reset_index().sort_values('date'),
                        by='ticker', on='date', direction='backward')
            .set_index(['ticker', 'date'])
            .sort_index()
            .fillna(0, downcast='infer')
         )
 

 print(result)

                   price
ticker date             
GOOGL  2021-09-20   2500
       2021-09-21   2600
       2021-09-22   2600
SPY    2021-09-20    430
       2021-09-21    430
       2021-09-22    430
TSLA   2021-09-20      0
       2021-09-21    700
       2021-09-22    710
 

Комментарии:

1. Спасибо! Это занимает ~30 секунд по сравнению с 30 минутами, которые потребовались моей группе!

2. @кристина, да, я так и думал. Это groupby apply будет медленный цикл по группам, поэтому для многих групп может потребоваться довольно много времени, а переиндексация каждой группы-сложная задача для векторизации, поэтому, учитывая, что вам все равно нужна каждая строка в конце, merge это хорошее решение