Search code examples
pythonnumpyrecursionnested-loops

Simplifying nested loops using recursion


Let's assume we have a three-dimensional array x. For each k, I want to have the mean of x[:, :, k] stored as out[k]. The task is simple if it is a 3D or 4D matrix, with the for loop like this:

x = np.random.normal(0, 1, [4, 4, 3])
out = np.zeros(x[0, 0, :].shape)
for k in range(3):
    out[k] =  np.mean(x[:, :, k])
    
x2 = np.random.normal(0, 1, [5, 5, 4, 6])
out2 = np.zeros(x2[0, 0, :, :].shape)
for k in range(4):
    for l in range(6):
        out2[k, l] = np.mean(x2[:, :, k, l])

But the code is getting ugly if I want to go for higher dimension (let's say I want to cover up to 100 dimension), since the number of nested loops increase.

I know that I have to do it recursively with 2D matrices serve as my base case, but I don't know how to implement it.


Solution

  • Don't need recursion for this case. Turns out my previous answer was pretty close, just need a few transposes.

    x2 = np.random.normal(0, 1, [2, 2, 7, 5, 5, 4, 6])
    
    dims = len(x2.shape)
    
    out = np.mean(x2.T, axis = tuple(range(-2, -1))).T
    
    
    print(out[0,0,0,0,0], np.mean(x2[:, :, 0, 0, 0, 0, 0]))
    

    Notably, the values between out and np.mean are sometimes slightly off on the last three digits, so there's some weird precision loss happening somehow.

    If you must do it recursively...

    x2 = np.random.normal(0, 1, [2, 2, 7, 5, 5, 4, 6])
    
    def recursive_mean(x, out = None):
        if out is None:
            out = np.zeros(x.shape[2:])
            
        if len(x.shape) > 3:
            for i in range(x.shape[2]):
                out[i] = recursive_mean(x[:, :, i], out[i])
                
        else:
            for i in range(x.shape[2]):
                out[i] = np.mean(x[:, :, i])
            
            
        return out
    
    out2 = recursive_mean(x2)
    
    print(out2[0,0,0,0,0], np.mean(x2[:, :, 0, 0, 0, 0, 0]))
    print(out2[3, 1, 2, 0, 4], np.mean(x2[:, :, 3, 1, 2, 0, 4]))