Rewriting numpy function to handle 2d and 3d inputs

Question

I am trying to rewrite an numpy function such that it can deal with 2d and 3d inputs. Consider the following code:

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()

As you can see, it works fine for the 1d case but does not work for the 2d and 3d cases. You can imagine that the thetas are stacked in the third dimension. Expected output shape 2d: (1000, 10) Expected output shape 3d: (4, 1000, 10)

The most obvious is just to iterate through all dimensions but this is really slow.

Answer 1

Here is one way to approach the problem:

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

---

Try using the following

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)