[Python pandas]複数のcolumnを参照して、rollingでapplyしたい。

はじめに

Pythonライブラリpandasで、移動平均などを計算するときには、窓関数を計算するrollingメソッドを利用することが多いです。rollingメソッドを使って、meanなどのデフォルトで用意されている計算以外を行う場合は、applyメソッドで、独自定義した関数を指定します。その時に複数のcolumnが関連した計算をしたいことがあるのですが、rollingメソッドは、columnごとで計算をするので、単純には複数のcolumnを参照した計算を行うことができません。

この記事は、rollingメソッドを使って、複数のcolumnが関連する計算方法についてのメモです。

環境

• Windows 10 Home
• Python(3.7.6)
• Pandas(1.2.1)

複数のcolumnを参照する方法

pandasの説明では、独自関数の引数では、ndarray(if raw=True)かSeries(if raw=False)を受けることになっているので、複数のcolumnを指定することができません。そこで、argsを使って、独自関数の引数にもとのDataFrameを指定して、複数のcolumnを参照します。

import pandas as pd
import numpy as np
df = pd.DataFrame({'date':pd.date_range('2021/01/20', '2021/01/23'),
                   'A':np.arange(0, 4),
                   'B':np.arange(11, 15),
                   'C':np.arange(34, 38)})
df = df.set_index('date')
#print("Original DataFrame")
#print(df)

#applyしたい関数
def user_define(x, df_u):
    """
    Args:
    	x: applyで与えられるSeries(indexを利用するため）
    	df_u: argsで与える元のDataFrame
    """
    df_A = df_u.loc[x.index, 'A']
    df_C = df_u.loc[x.index, 'C']
    return np.sum(df_A + df_C)

#print("Applied Series ")
df['A'].rolling(2).apply(user_define, args=(df,), raw=False)#argsはタプル指定

もとのDataFrame

　　

rolling.applyを適用したSeries

追記

上記の例だと、method='table'でもできます。

import pandas as pd
import numpy as np
import numba
df = pd.DataFrame({'date':pd.date_range('2021/01/20', '2021/01/23'),
                   'A':np.arange(0, 4),
                   'B':np.arange(11, 15),
                   'C':np.arange(34, 38)})
df = df.set_index('date')

@numba.jit
def user_define(w):
    return np.sum(w[:,0] + w[:,2])

df.rolling(2, method='table').apply(user_define, raw=True, engine='numba').loc[:,'A']

参考