使用 pandas 数据框矢量化复杂切片

Question

出于速度目的，我希望能够对这段代码进行矢量化处理。目的是计算一个函数，在这种情况下是标准偏差，从一对包含在两个单独数组中的日期元组。

import pandas as pd
import numpy as np

asd_1 = pd.Series(0.01 * np.random.randn(252), index=pd.date_range('2011-1-1', periods=252))

index_1 = pd.to_datetime(['2011-2-2', '2011-4-3', '2011-5-1',])
index_2 = pd.to_datetime(['2011-2-15', '2011-4-16', '2011-5-17',])

index_tot = list(zip(index_1,index_2))

aux_learning_std = pd.DataFrame([np.nanstd(asd_1.loc[i:j]) for i, j in index_tot], index=index_1)

可行的解决方案是通过循环执行的，但我宁愿能够通过 numpy/pandas 对其进行矢量化，这样会更快。最初我想使用类似的东西：

df_aux = pd.concat([asd_1 for _ in range(len(index_1))], axis=1)
results = df_aux.apply(lambda x: np.nanstd(x.loc[i,j]), axis = 0)

但在这里我没有将向量放在一个操作中。

欢迎提出任何建议。

p.s.：下面有一张图片用于说明目的

Answer 1

数组中跨范围的向量化标准差

def get_ranges_arr(starts,ends):
    # Taken from http://stackoverflow.com/a/37626057/3293881
    counts = ends - starts
    counts_csum = counts.cumsum()
    id_arr = np.ones(counts_csum[-1],dtype=int)
    id_arr[0] = starts[0]
    id_arr[counts_csum[:-1]] = starts[1:] - ends[:-1] + 1
    return id_arr.cumsum()

def ranged_std(arr,starts,ends):
    # Get all indices and the IDs corresponding to same groups
    idx = get_ranges_arr(starts,ends)
    id_arr = np.repeat(np.arange(starts.size),ends-starts)
    
    # Extract relevant data
    slice_arr = arr[idx]
    
    # Simulate standard deviation implementation for a number of groups
    # using id_arr as the basis to perform various mathematical operations
    # within each group. Since, std. deviation performs sum/mean reduction,
    # we can simply use np.bincount for an efficient implementation.
    # Std. deviation formula used :
    #https://github.com/numpy/numpy/blob/v1.11.0/numpy/core/fromnumeric.py#L2939
    grp_counts = np.bincount(id_arr)
    mean_vals = np.bincount(id_arr,slice_arr)/grp_counts
    abs_vals = np.abs(slice_arr - mean_vals[id_arr])**2
    return np.sqrt(np.bincount(id_arr,abs_vals)/grp_counts)

示例运行（针对循环版本进行验证）

In [173]: arr = np.random.randint(0,9,(20))

In [174]: starts = np.array([2,6,11])

In [175]: ends = np.array([8,9,15])

In [176]: [np.std(arr[i:j]) for i,j in zip(starts,ends)]
Out[176]: [1.9720265943665387, 0.81649658092772603, 0.82915619758884995]

In [177]: ranged_std(arr,starts,ends)
Out[177]: array([ 1.97202659,  0.81649658,  0.8291562 ])    

运行时测试

案例#1：极少数范围3

In [21]: arr = np.random.randint(0,9,(20))

In [22]: starts = np.array([2,6,11])

In [23]: ends = np.array([8,9,15])

In [24]: %timeit [np.std(arr[i:j]) for i,j in zip(starts,ends)]
10000 loops, best of 3: 146 µs per loop

In [25]: %timeit ranged_std(arr,starts,ends)
10000 loops, best of 3: 45 µs per loop

案例#2：相当数量的范围1000

In [32]: arr = np.random.randint(0,9,(1010))

In [33]: starts = np.random.randint(0,9,(1000))

In [34]: ends = starts + np.random.randint(0,9,(1000))

In [35]: %timeit [np.std(arr[i:j]) for i,j in zip(starts,ends)]
10 loops, best of 3: 47.5 ms per loop

In [36]: %timeit ranged_std(arr,starts,ends)
1000 loops, best of 3: 217 µs per loop

案例#3：大量范围10000

In [60]: arr = np.random.randint(0,9,(1010))

In [61]: arr = np.random.randint(0,9,(10010))

In [62]: starts = np.random.randint(0,9,(10000))

In [63]: ends = starts + np.random.randint(0,9,(10000))

In [64]: %timeit [np.std(arr[i:j]) for i,j in zip(starts,ends)]
1 loops, best of 3: 474 ms per loop

In [65]: %timeit ranged_std(arr,starts,ends)
100 loops, best of 3: 2.17 ms per loop

200x+ 的加速速度真的很惊人！

---

使用ranged_std解决我们的案例

# Get start, stop numeric indices as needed for getting ranges array later on
starts = asd_1.index.searchsorted(index_1)
ends = asd_1.index.searchsorted(index_2)

# Create final dataframe output using ranged_std func
df = pd.DataFrame(ranged_std(asd_1.values,starts,ends+1),index=index_1)

用于验证的示例运行 -

In [17]: asd_1 = pd.Series(0.01 * np.random.randn(252), index=\
    ...:                   pd.date_range('2011-1-1', periods=252))
    ...: 
    ...: index_1 = pd.to_datetime(['2011-2-2', '2011-4-3', '2011-5-1',])
    ...: index_2 = pd.to_datetime(['2011-2-15', '2011-4-16', '2011-5-17',])
    ...: 
    ...: index_tot = list(zip(index_1,index_2))
    ...: aux_learning_std = pd.DataFrame([np.nanstd(asd_1.loc[i:j]) for i, j in \
    ...:                                                index_tot], index=index_1)
    ...: 

In [18]: starts = asd_1.index.searchsorted(index_1)
    ...: ends = asd_1.index.searchsorted(index_2)
    ...: df = pd.DataFrame(ranged_std(asd_1.values,starts,ends+1),index=index_1)
    ...: 

In [19]: aux_learning_std
Out[19]: 
                   0
2011-02-02  0.007244
2011-04-03  0.012862
2011-05-01  0.010155

In [20]: df
Out[20]: 
                   0
2011-02-02  0.007244
2011-04-03  0.012862
2011-05-01  0.010155