重写 numpy 函数来处理 2d 和 3d 输入

Question

我正在尝试重写一个 numpy 函数，以便它可以处理 2d 和 3d 输入。 考虑以下代码：

import numpy as np
def adstock_geometric(x: np.array, theta: 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] + theta * x_decayed[xi - 1]

    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)

    output1d = adstock_geometric(x=x, theta=0.5)  # works fine
    output3d = adstock_geometric(x=x, theta=rand3d)
    output2d = adstock_geometric(x=x, theta=rand2d)
    
if __name__ == '__main__':
    testrun()

如您所见，它适用于 1d 情况，但不适用于 2d 和 3d 情况。 您可以想象 θ 是在三维空间中堆叠的。 预期输出形状 2d：(1000, 10) 预期输出形状 3d：(4, 1000, 10)

最明显的就是迭代所有维度，但这真的很慢。

Answer 1

这是解决该问题的一种方法：

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

---

尝试使用以下方法

x = np.ones(10)
rand1d = 0.5
rand2d = np.random.randint(0, 10, size=(5, 1)) / 10
rand3d = np.random.randint(0, 10, size=(4, 2, 1)) / 10
addstock_geom(x, rand1d)
addstock_geom(x, rand2d)
addstock_geom(x, rand3d)