熊猫Groupby EWM

时间:2019-09-19 01:53:40

标签: python pandas time-series pandas-groupby rolling-computation

我已标记事件(时间序列)数据,其中事件以给定标签的随机间隔发生。我想计算组内ewma并将其作为新列“ X1_EWMA”添加到数据框。到目前为止,这是代码:

import pandas as pd
import numpy as np
import altair as alt

n = 1000
df = pd.DataFrame({
    'T': pd.date_range('20190101', periods=n, freq='H'),
    'C1': np.random.choice(list('PYTHON'), n),
    'C2': np.random.choice(list('FUN'), n),
    'X1': np.random.randn(n),
    'X2': 100 + 10 * np.random.randn(n)
})

ts = df.set_index('T')

display(df.head())
display(ts.head())

感谢SO: Pandas Groupby and apply method with custom function),我能够使用以下公式计算分组的EWMA:

ewm = ts.groupby(['C1']).apply(lambda x: x['X1'].ewm(halflife=10).mean())
ewm.head()

它产生一个由分类变量和日期时间之一索引的序列。系列的长度与原始数据帧和时间系列(df和ts)相同

现在,我想我可以做一些体操运动,以通过连接行索引(假设排序顺序没有变化)将其重新连接到原始数据帧(df),但这似乎并不正确,甚至可能是这是一种冒险的方法,因为groupby仅位于分类标签之一内-我需要小心并进行一些检查/分类/重新索引。

似乎应该有一个更简单的方法,可以将时间序列列直接添加到数据帧(df)或时间序列(ts),而无需创建单独的序列或数据帧并将它们连接在一起。如果我想添加滚动统计信息,例如:

ts.groupby('C1').rolling(10).mean()

在此先感谢您的帮助或投入。

基于可接受答案的结果:

import pandas as pd
import numpy as np
import math
import altair as alt

alt.renderers.enable('notebook')      # for rendering in the notebook
alt.data_transformers.enable('json')  # for plotting data larger than 5000 points

# make a dataframe to test
n = 1000
df = pd.DataFrame({
    'T': pd.date_range('20190101', periods=n, freq='H'),
    'C1': np.random.choice(list('PYTHON'), n),
    'C2': np.random.choice(list('FUN'), n),
    'X1': np.linspace(0, 2*math.pi, n),
    'X2': np.random.randn(n),
})

# add a new variable that is a function of X1, X2 + a random outlier probability
df['X3'] = 0.2 * df['X2'] + np.sin(df['X1']) + np.random.choice(a=[0, 2], size=n, p=[0.98, 0.02])

# make it a time series for later resampling use cases.
ts = df.set_index('T')

#  SOLUTION: Add the ewma line with groupby().transform().
ts['ewm'] = ts.groupby(['C1'])['X3'].transform(lambda x: x.ewm(halflife=1).mean())

# plot the points and ewma using altair faceting and layering
points = alt.Chart().mark_circle(size=20, opacity=0.9).encode(
    x = 'T', 
    y = 'X3',
    color = 'C2',
).properties(width=270, height=170)

lines = alt.Chart().mark_line(size=1, color='red', opacity=1).encode(
    x = 'T', 
    y = 'ewm'
)

alt.layer(points, lines).facet(facet='C1', data=ts.reset_index()).properties(columns=3)

enter image description here

2 个答案:

答案 0 :(得分:1)

让我们使用transform来解决问题:

t['ewm'] = ts.groupby(['C1'])['X1'].transform(lambda x: x.ewm(halflife=10).mean()).values()

答案 1 :(得分:0)

您可以试试吗? 请勿设置ts = df.set_index('T')。然后您可以执行以下操作

ts['ewm']=ts.groupby(['C1'], sort=False).apply(lambda x: x['X1'].ewm(halflife=10).mean()).reset_index(drop=True)