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