Simplifying nested loops using recursion

Question

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.

Asked By: user129994

||

Source

Answer 1

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

Answered By: Michael Cao

Answer 2

You seem to want to take the mean along the first two axes. You can do this simply by x.mean(axis=(0, 1))

Example:

out* is the mean calculated by your approach
out*_np is the mean calculated by np.mean like I show above.

x = np.random.normal(0, 1, [4, 4, 3])
out = np.zeros(x[0, 0, :].shape)
out_np = x.mean(axis=(0,1))
for k in range(3):
    out[k] =  np.mean(x[:, :, k])

print(np.allclose(out, out_np)) # True

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

print(np.allclose(out2, out2_np)) # True

To generalize this to any function that operates on a 2D matrix, you can reshape your original matrix to a 3D matrix, apply your method over the last axis, and reshape the result to the required shape:

def apply_2dfunc(X, func):
    X_r = X.reshape((*X.shape[0:2], -1))
    res = np.zeros((X_r.shape[-1],))
    for i in range(len(res)):
        res[i] = func(X_r[:, :, i])

    return res.reshape(X.shape[2:])

With your example of func=np.mean,

out_g = apply_2dfunc(x, np.mean)
print(np.allclose(out_g, out)) # True


out2_g = apply_2dfunc(x2, np.mean)
print(np.allclose(out2_g, out2)) # True

Answered By: Pranav Hosangadi

Simplifying nested loops using recursion

Question:

Answers: