修复支持多维输入的Numpy几何衰减函数实现

以下是完全向量化的实现方案，兼容多维输入（如(N, T)形状的x），支持L截断参数与归一化选项，效率远高于循环版本：

import numpy as np

def adstock_geometric_vectorized(x, theta, L=7, normalize=False):
    # 统一theta格式：支持标量或(N,)数组输入
    theta = np.asarray(theta)
    if theta.ndim == 0:
        theta = np.full(x.shape[0], theta)
    
    N, T = x.shape
    # 在时间维度开头补L-1个0，用于生成完整滑动窗口
    x_pad = np.pad(x, ((0, 0), (L-1, 0)), mode='constant')
    # 提取每个时间步的前L个元素窗口，形状为(N, T, L)
    windows = np.lib.stride_tricks.sliding_window_view(x_pad, window_shape=L, axis=1)
    # 反转窗口内元素顺序，对应[x[t], x[t-1], ..., x[t-L+1]]
    windows_reversed = windows[:, :, ::-1]
    
    # 构造几何权重矩阵：theta[i]^0, theta[i]^1, ..., theta[i]^(L-1)
    l_indices = np.arange(L)
    weights = theta[:, np.newaxis] ** l_indices
    
    # 计算加权累加和
    cumsum = (windows_reversed * weights[:, np.newaxis, :]).sum(axis=2)
    
    if normalize:
        # 计算每个时间步的权重和（模拟输入全1的衰减结果）
        ones_pad = np.pad(np.ones_like(x), ((0, 0), (L-1, 0)), mode='constant')
        ones_windows = np.lib.stride_tricks.sliding_window_view(ones_pad, window_shape=L, axis=1)
        ones_windows_reversed = ones_windows[:, :, ::-1]
        sum_weights = (ones_windows_reversed * weights[:, np.newaxis, :]).sum(axis=2)
        cumsum = cumsum / sum_weights
    
    return cumsum

核心思路 #

• 滑动窗口：利用np.lib.stride_tricks.sliding_window_view生成每个时间步的前L个元素窗口，避免手动循环遍历时间步
• 广播权重：通过广播机制生成每个样本对应的几何权重数组，实现逐样本的加权计算
• 归一化处理：通过模拟全1输入的衰减结果，得到每个时间步的权重和，实现归一化

正确性验证 #

可以和原循环版本对比结果，确保功能一致：

# 生成测试数据
np.random.seed(42)
N = 10
T = 20
x = np.random.randn(N, T)
theta = np.random.uniform(0, 1, size=N)
L = 5

# 原循环版本代码
def adstock_geometric_loop(x, theta, L=7, normalize=False):
    result = []
    for i in range(x.shape[1]):
        cumsum = 0
        wcumsum = 0
        for l in range(min(i + 1, L)):
            w = theta ** l
            cumsum += w * x[:, i - l]
            wcumsum += w
        if not normalize:
            result.append(cumsum)
        else:
            result.append(cumsum / wcumsum)
    return np.array(result).T

# 对比结果
output_loop = adstock_geometric_loop(x, theta, L=L)
output_vectorized = adstock_geometric_vectorized(x, theta, L=L)

print(f"最大绝对误差：{np.max(np.abs(output_loop - output_vectorized)):.10f}")
print(f"结果是否一致：{np.allclose(output_loop, output_vectorized)}")

# 验证归一化场景
output_loop_norm = adstock_geometric_loop(x, theta, L=L, normalize=True)
output_vectorized_norm = adstock_geometric_vectorized(x, theta, L=L, normalize=True)

print(f"归一化后最大绝对误差：{np.max(np.abs(output_loop_norm - output_vectorized_norm)):.10f}")
print(f"归一化结果是否一致：{np.allclose(output_loop_norm, output_vectorized_norm)}")

优势说明 #

• 无显式Python循环，完全基于Numpy底层优化，处理大规模数据时效率提升显著
• 支持(N, T)形状输入与(N,)形状的衰减系数theta，也兼容theta为标量的情况
• 完整保留原循环版本的L截断与归一化功能
• 可轻松扩展到更高维输入（如(N, M, T)），只需调整pad和sliding_window_view的axis参数

内容的提问来源于stack exchange，提问作者rutkov