python学习——解决在pandas和numpy上的rolling apply操作

今天和隔壁某宇宙券商小哥交流时候，发现了这么一个需求，代码看起来一点问题都没有，但是运行不起来

import numpy as np
import pandas as pd

a = np.arange(0,10).reshape(5,2)
b = pd.DataFrame(a, columns=['X', 'Y'])
c = b.rolling(window=2).apply(lambda x: np.linalg.det(x))

其实从上面的代码看起来很简单，想实现的就是对b进行滑动计算自定义函数 （window不一定非得是2，只是这里用了行列式函数，所以是2)
上面的代码肯定是运行不起来的，不然也不会费工夫写这篇博客了，其实解决问题的思路应该分成两步，第一个的话，获得rolling的subset，然后对各个subset依次进行操作，这里要说一点的是，对于dataframe直接做操作的apply，里面的x对应的是df的一行或者一列，那直接rolling后获得的是一个rolling对象，这个时候的x其实是一个index
因此我们有如下的处理方案（直接运行最后一行）：

def dataframe_roll(df):
    def my_fn(window_series):
        window_df = df[(df.index >= window_series.index[0]) & (df.index <= window_series.index[-1])]
        return  np.linalg.det(window_df)
    return my_fn
c = b.rolling(window=2).apply(dataframe_roll(b))

然后小哥说，其实对numpy处理更好，针对numpy，就有一些比较好的工具函数可以用了，比如：
虽然这个名字看起来比较丑，numpy.lib.stride_tricks.sliding_window_view，但是确实是比较好用的
处理方案如下所示：

# ---------------------- Numpy array rolling window操作  ---------------
# 简单操作如果是pandas有的 就转为pandas再处理比较快
def rolling_window(a, window, axis=0):
    """
    返回2D array的滑窗array的array
    """
    if axis == 0:
        shape = (a.shape[0] - window +1, window, a.shape[-1])
        strides = (a.strides[0],) + a.strides
        a_rolling = np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
    elif axis==1:
        shape = (a.shape[-1] - window +1,) + (a.shape[0], window) 
        strides = (a.strides[-1],) + a.strides
        a_rolling = np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
    return a_rolling

# 对于pandas DataFrame自带的函数，对于大维度，还是pandas rolling method快
# 小数组的话，此方法快
# 上述的快，都是能十倍的速度比。
def rolling_mean(A, window=None):
    ret = np.full(A.shape, np.nan)
    A_rolling = rolling_window(A, window=window, axis=0)
    Atmp = np.stack(map(lambda x:np.mean(x, axis=0), A_rolling))
    ret[window-1:,:] = Atmp
    return ret

def rolling_nanmean(A, window=None):
    ret = np.full(A.shape, np.nan)
    A_rolling = rolling_window(A, window=window, axis=0)
    Atmp = np.stack(map(lambda x:np.nanmean(x, axis=0), A_rolling))
    ret[window-1:,:] = Atmp
    return ret

参考文献：
Apply function on a rolling slice of a pandas DataFrame
numpy ndarray的rolling window操作
numpy.lib.stride_tricks.sliding_window_view