如何在Python/NumPy中无循环实现滚动加权求和？

尝试用np.convolve计算滚动成交量加权平均价（VWAP）时，速度远快于标准for循环，但结果不正确，还遗漏了数组最后一项。求无需for循环的正确滚动加权求和实现方法。

已尝试方案 #

标准for循环（速度慢，结果正确） #

def calc_vwap_1(price, volume, period_lookback):
    """
    Calculates the volume-weighted average price (VWAP) for a given period of time.
    The VWAP is calculated by taking the sum of the product of each price and volume over a given period, 
    and dividing by the sum of the volume over that period.
    
    Parameters:
        price (numpy.ndarray): A list or array of prices.
        volume (numpy.ndarray): A list or array of volumes, corresponding to the prices.
        period_lookback (int): The number of days to look back when calculating VWAP.
        
    Returns:
        numpy.ndarray: An array of VWAP values, one for each day in the input period.
    """
    vwap = np.zeros(len(price))
    for i in range(period_lookback, len(price)):
        lb = i - period_lookback  # lower bound
        ub = i + 1  # upper bound
        volume_sum = volume[lb:ub].sum()
        if volume_sum > 0:
            vwap[i] = (price[lb:ub] * volume[lb:ub]).sum() / volume_sum
        else:
            vwap[i] = np.nan
    return vwap

使用np.convolve的same模式 #

def calc_vwap_2(price, volume, period_lookback):
    price_volume = price * volume
    # Use convolve to get the rolling sum of product of price and volume
    price_volume_conv = np.convolve(price_volume, np.ones(period_lookback), mode='same')[period_lookback-1:]
    # Use convolve to get the rolling sum of volume
    volume_conv = np.convolve(volume, np.ones(period_lookback), mode='same')[period_lookback-1:]
    # Create a mask to check if the volume sum is greater than 0
    mask = volume_conv > 0
    # Initialize the vwap array
    vwap = np.zeros(len(price))
    # Use the mask to check if volume sum is greater than zero, if it is, proceed with the division and store the result in vwap array, otherwise store NaN
    vwap[period_lookback-1:] = np.where(mask, price_volume_conv / volume_conv, np.nan)
    return vwap

使用np.convolve的valid模式 #

def calc_vwap_3(price, volume, period_lookback):
    # Calculate product of price and volume
    price_volume = price * volume
    # Use convolve to get the rolling sum of product of price and volume and volume array
    price_volume_conv = np.convolve(price_volume, np.ones(period_lookback), mode='valid')
    # Use convolve to get the rolling sum of volume
    volume_conv = np.convolve(volume, np.ones(period_lookback), mode='valid')
    # Create a mask to check if the volume sum is greater than 0
    mask = volume_conv > 0
    # Initialize the vwap array
    vwap = np.zeros(len(price))
    # Use the mask to check if volume sum is greater than zero, if it is, proceed with the division and store the result in vwap array, otherwise store NaN
    vwap[period_lookback-1:] = np.where(mask, price_volume_conv / volume_conv, np.nan)
    return vwap

使用np.cumsum结合切片 #

def calc_vwap_4(price, volume, period_lookback):
    price_volume = price * volume
    # Use cumsum to get the rolling sum of product of price and volume
    price_volume_cumsum = np.cumsum(price_volume)[period_lookback-1:]
    # Use cumsum to get the rolling sum of volume
    volume_cumsum = np.cumsum(volume)[period_lookback-1:]
    # Create a mask to check if the volume sum is greater than 0
    mask = volume_cumsum > 0
    # Initialize the vwap array
    vwap = np.zeros(len(price))
    # Use the mask to check if volume sum is greater than zero, if it is, proceed with the division and store the result in vwap array, otherwise store NaN
    vwap[period_lookback-1:] = np.where(mask, price_volume_cumsum / volume_cumsum, np.nan)
    return vwap

使用np.reduceat #

def calc_vwap_5(price, volume, period_lookback):
    price_volume = price * volume
    # Use reduceat to get the rolling sum of product of price and volume
    price_volume_cumsum = np.add.reduceat(price_volume, np.arange(0, len(price), period_lookback))[period_lookback-1:]
    # Use reduceat to get the rolling sum of volume
    volume_cumsum = np.add.reduceat(volume, np.arange(0, len(price), period_lookback))[period_lookback-1:]
    # Create a mask to check if the volume sum is greater than 0
    mask = volume_cumsum > 0
    # Initialize the vwap array
    vwap = np.zeros(len(price))
    # Use the mask to check if volume sum is greater than zero, if it is, proceed with the division and store the result in vwap array, otherwise store NaN
    vwap[period_lookback-1:] = np.where(mask, price_volume_cumsum / volume_cumsum, np.nan)
    return vwap

使用np.lib.stride_tricks.as_strided #

def calc_vwap_6(price, volume, period_lookback):
    price_volume = price * volume
    price_volume_strided = np.lib.stride_tricks.as_strided(price_volume, shape=(len(price)-period_lookback+1, period_lookback), strides=(price_volume.strides[0], price_volume.strides[0]))
    volume_strided = np.lib.stride_tricks.as_strided(volume, shape=(len(price)-period_lookback+1, period_lookback), strides=(volume.strides[0], volume.strides[0]))
    price_volume_sum = price_volume_strided.sum(axis=1)
    volume_sum = volume_strided.sum(axis=1)
    mask = volume_sum > 0
    vwap = np.zeros(len(price))
    vwap[period_lookback-1:] = np.where(mask, price_volume_sum / volume_sum, np.nan)
    return vwap

测试数据 #

import numpy as np

price = np.random.random(10000)
volume = np.random.random(10000)
print(calc_vwap_1(price, volume, 100))
print()
print(calc_vwap_2(price, volume, 100))
print()
print(calc_vwap_3(price, volume, 100))
print()
print(calc_vwap_4(price, volume, 100))
print()
print(calc_vwap_5(price, volume, 100))
print()
print(calc_vwap_6(price, volume, 100))
print()

计算结果 #

vwap_1 -> [0.         0.         0.         ... 0.47375965 0.47762679 0.48448903] # 正确结果

vwap_2 -> [0.         0.         0.         ... 0.53108759 0.51933363 0.51360848]

vwap_3 -> [0.         0.         0.         ... 0.49834202 0.4984141  0.49845759]

vwap_4 -> [0.         0.         0.         ... 0.49834202 0.4984141  0.49845759]

vwap_5 -> [0.         0.         0.         ... 0.48040529 0.48040529 0.48040529]

vwap_6 -> [0.         0.         0.         ... 0.47027032 0.48009596 0.48040529]

内容的提问来源于Stack Exchange，提问作者Hofbr