改写NumPy的adstock_geometric函数以支持2D和3D输入的技术求助

你遇到的问题是原函数仅适配1D输入，处理2D/3D的theta数组时会因维度不匹配报错，且直接遍历所有维度的循环效率极低。下面提供两种高效的解决方案，核心思路是利用NumPy的向量化运算，避免逐个元素的慢循环。

方案一：基于时间步循环的向量化实现（推荐） #

这个方案仅遍历时间维度（通常时间步远小于batch规模），每次循环对整个batch的元素进行向量运算，速度非常快，且逻辑和原函数一致，容易理解：

import numpy as np

def adstock_geometric(x: np.ndarray, theta: np.ndarray):
    x = np.asarray(x)
    theta = np.asarray(theta)
    
    # 获取时间序列长度
    T = x.shape[0]
    
    # 处理theta的维度兼容：统一转为最后一维为1的数组
    if theta.ndim == 0:
        theta = theta.reshape((1, 1))
    elif theta.ndim == 1:
        theta = theta.reshape((-1, 1))
    
    # 提取batch维度信息
    batch_dims = theta.shape[:-1]
    
    # 将x扩展为与theta兼容的形状：(*batch_dims, T)
    x_expanded = x.reshape((1,) * len(batch_dims) + (T,))
    x_decayed = np.broadcast_to(x_expanded, batch_dims + (T,))
    
    # 扩展theta以支持广播运算
    theta_expanded = theta.reshape(batch_dims + (1,))
    
    # 仅遍历时间步，而非batch元素，大幅减少循环次数
    for t in range(1, T):
        x_decayed[..., t] += theta_expanded * x_decayed[..., t-1]
    
    # 处理1D输入的返回格式，保持和原函数一致
    if len(batch_dims) == 0:
        return x_decayed.reshape((T,))
    
    return x_decayed

def testrun():
    rand3d = np.random.randint(0, 10, size=(4, 1000, 1)) / 10
    rand2d = np.random.randint(0, 10, size=(1000, 1)) / 10
    x = np.ones(10)
    
    # 1D测试（和原函数结果一致）
    output1d = adstock_geometric(x=x, theta=0.5)
    print("1D输出形状:", output1d.shape)  # (10,)
    
    # 2D测试
    output2d = adstock_geometric(x=x, theta=rand2d)
    print("2D输出形状:", output2d.shape)  # (1000, 10)
    
    # 3D测试
    output3d = adstock_geometric(x=x, theta=rand3d)
    print("3D输出形状:", output3d.shape)  # (4, 1000, 10)

if __name__ == '__main__':
    testrun()

方案优势 #

• 循环次数仅等于时间序列长度（比如示例中的10次），和batch规模无关，即使batch是百万级也不会变慢
• 每次循环都是NumPy的向量运算，效率远高于逐个元素的Python循环
• 自动兼容1D/2D/3D输入，返回形状符合预期

方案二：完全向量化实现（无循环） #

如果你想彻底避免循环，可以利用幂次数组和累积求和实现，逻辑稍复杂但同样高效：

import numpy as np

def adstock_geometric_vectorized(x: np.ndarray, theta: np.ndarray):
    x = np.asarray(x)
    theta = np.asarray(theta)
    
    T = x.shape[0]
    
    # 统一theta维度为(*batch_dims, 1)
    if theta.ndim == 0:
        theta = theta.reshape((1, 1))
    elif theta.ndim == 1:
        theta = theta.reshape((-1, 1))
    
    batch_dims = theta.shape[:-1]
    
    # 扩展x到batch维度
    x_expanded = x.reshape((1,) * len(batch_dims) + (T,))
    x_expanded = np.broadcast_to(x_expanded, batch_dims + (T,))
    
    # 生成theta的幂次数组：theta^0, theta^1, ..., theta^{T-1}
    k = np.arange(T).reshape((1,) * len(batch_dims) + (T,))
    theta_pows = theta.reshape(batch_dims + (1,)) ** k
    
    # 通过反转、累积求和、再反转实现递归衰减的向量化计算
    x_reversed = np.flip(x_expanded, axis=-1)
    theta_pows_reversed = np.flip(theta_pows, axis=-1)
    cumulative = np.cumsum(x_reversed * theta_pows_reversed, axis=-1)
    x_decayed = np.flip(cumulative, axis=-1)
    
    # 适配1D输出格式
    if len(batch_dims) == 0:
        return x_decayed.reshape((T,))
    
    return x_decayed

这个版本通过数学推导将递归运算转化为数组操作，完全没有Python循环，适合对性能要求极高的场景。

备注：内容来源于stack exchange，提问作者richard baws