修复此 numpy 函数，使其接受多维输入

Question

我正在尝试计算以下函数：

地理衰变函数：

在 numpy 中，必须进行向量化。这种情况下的问题是它应该适用于多维输入。例如，x 的维度为 (999, 20)，其中第二个维度表示上述函数中的 t。 h 将具有与 x 的第一个维度匹配的维度，并且应映射到 x 的第一个维度，在上述情况下 h 的维度将为 (999,)，这种情况下的输出应为维度 (999, 20 ).

如果 x 为 1d，则以下代码适用于 2d 和 3d 的 theta（不考虑 L）：

def adstock_geometric(x, h):
    axis = n if (n := len(np.array(h).shape) - 1) > 0 else None
    return ((h**np.arange(x.size)) * x).cumsum(axis)

但是，如果我们有一个 2d x ，其中第一个维度应该映射到 h 的第一个维度，我不知道除了循环每个维度之外我们还能怎么做（因此在上面的情况下是一个 range(999) 循环）..可以有人在这件事上帮助我吗？

另一个一维解是：

def adstock_geometric(x: np.array, h: np.array):
    x_decayed = np.zeros_like(x)
    x_decayed[0] = x[0]

    for xi in range(1, len(x_decayed)):
        x_decayed[xi] = x[xi] + h * x_decayed[xi - 1]

    return x_decayed

但是速度确实很慢..并且不适用于 2d x。此外，我不确定如何合并 L..

更新 下面的似乎有效，但它没有矢量化：

def adstock_geometric(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

Answer 1

一旦您意识到可以通过沿最后一个轴的离散卷积来计算分子和分母，矢量化就成为可能：

import numpy as np
from scipy.signal import oaconvolve

def adstock_vectorized(x, theta, L=7, normalize=False):  # Proposed
    assert theta.shape == x.shape[:-1]  # Check shapes
    l = x.shape[-1]  # Same as `len_samples`
    w = theta[:, np.newaxis] ** np.arange(min(L, l))  # Create `w`
    # Numerator: convolve with `x`; denominator: convolve with ones; both along last axis
    num = oaconvolve(x, w, mode="full", axes=-1)
    denom = oaconvolve(np.ones_like(x), w, mode="full", axes=-1)
    return num[..., :l] / denom[..., :l] if normalize else num[..., :l]

def adstock_geometric(x, theta, L=7, normalize=False):  # From the question, condensed
    result = []
    for i in range(x.shape[1]):
        cumsum, wcumsum = 0, 0
        for l in range(min(i + 1, L)):
            w = theta ** l
            cumsum += w * x[:, i - l]
            wcumsum += w
        result.append((cumsum / wcumsum) if normalize else cumsum)
    return np.array(result).T

if __name__ == "__main__":   
    # Create some demo data
    np.random.seed(42)
    num_samples, len_samples = 30, 10
    x = np.random.normal(size=(num_samples, len_samples))
    h = np.random.normal(size=num_samples)
    
    for normalize in False, True:
        for L in range(1, len_samples):
            a_v = adstock_vectorized(x, h, L=L, normalize=normalize)
            a_g = adstock_geometric(x, h, L=L, normalize=normalize)
            assert np.allclose(a_v, a_g)

建议的解决方案，

adstock_vectorized()

，应该适用于任意数量的维度，只要您的数据位于最后一个维度。

还有两个注意事项：

• 您也可以尝试使用

  fftconvolve()

  代替

  oaconvolve()

  。对于您的实际数据来说，两者都可能更快。不幸的是，您不能直接使用

  convolve()

  ，因为它错过了

  axes

  参数，这在这里至关重要。
• 选择不同的随机种子和不同的数组大小，您可能会发现最后一行的断言偶尔会失败。在我所有的检查中，我总是只能发现这是由于数字错误而不是实际不同的结果。