average numpy array but retain shape

Question

I have a Numpy 3 axis array whose elements are 3 dimensional. I'd like to average them and return the same shape of the array. The normal average function removes the 3 dimensions and replace it with the average (as expected):

a = np.array([[[0.1, 0.2, 0.3], [0.2, 0.3, 0.4]],
              [[0.4, 0.4, 0.4], [0.7, 0.6, 0.8]]], np.float32)

b = np.average(a, axis=2)
# b = [[0.2, 0.3],
#      [0.4, 0.7]]

Result required:

# b = [[[0.2, 0.2, 0.2], [0.3, 0.3, 0.3]],
#      [[0.4, 0.4, 0.4], [0.7, 0.7, 0.7]]]

Can you do this elegantly or do I just have to iterate over the array in Python (which will be a lot slower compared to a powerful Numpy function).

Can you set the Dtype argument, for the np.mean function, to a 1D array perhaps?

Thanks.

Ozonosphere

Big fan of what you want to see in the question. – Amblyopia

In most cases i can imagine, broadcasting will do it without requiring 3d array. – Noonan

Answer 1

>>> import numpy as np
>>> a = np.array([[[0.1, 0.2, 0.3], [0.2, 0.3, 0.4]],
...               [[0.4, 0.4, 0.4], [0.7, 0.6, 0.8]]], np.float32)
>>> b = np.average(a, axis=2)
>>> b
array([[ 0.2       ,  0.29999998],
       [ 0.40000001,  0.69999999]], dtype=float32)
>>> c = np.dstack((b, b, b))
>>> c
array([[[ 0.2       ,  0.2       ,  0.2       ],
        [ 0.29999998,  0.29999998,  0.29999998]],

       [[ 0.40000001,  0.40000001,  0.40000001],
        [ 0.69999999,  0.69999999,  0.69999999]]], dtype=float32)

Terefah

Answer 2

Ok, CAUTION I don't have my masters in numpyology yet, but just playing around, I came up with:

>>> np.average(a,axis=-1).repeat(a.shape[-1]).reshape(a.shape)
array([[[ 0.2       ,  0.2       ,  0.2       ],
        [ 0.29999998,  0.29999998,  0.29999998]],

       [[ 0.40000001,  0.40000001,  0.40000001],
        [ 0.69999999,  0.69999999,  0.69999999]]], dtype=float32)

Organogenesis

I like it, that's better than what I came up with. – Amblyopia

It's one line and it works for all array shapes and sizes as long as you are taking the average along the last axis. – Organogenesis

Answer 3

Have you considered using broadcasting? Here is more info about broadcasting if you're new to the concept.

Here is an example using broadcast_arrays, keep in mind that the b produced here by broadcast_arrays should be treated as read only, you should make a copy if you want to write to it:

>>> b = np.average(a, axis=2)[:, :, np.newaxis]
>>> b, _ = np.broadcast_arrays(b, a)
>>> b
array([[[ 0.2       ,  0.2       ,  0.2       ],
        [ 0.29999998,  0.29999998,  0.29999998]],

       [[ 0.40000001,  0.40000001,  0.40000001],
        [ 0.69999999,  0.69999999,  0.69999999]]], dtype=float32)

Cloots

Answer 4

Here is a method that avoids making copies:

a = a.T
a[:] = a.mean(axis=0)
a = a.T

Or if you don't want to overwrite a:

b = np.empty_like(a)
b = b.T
b[:] = a.mean(axis=-1).T
b = b.T

Formative

And another great answer. How do I choose?! Thanks! :) – Ozonosphere

Answer 5

This is for an arbitrary axis:

array is the ndimentional array and axis is the axis to average

np.repeat( np.expand_dims( np.mean( array, axis ), axis ), array.shape[axis], axis )

Sal

Answer 6

You can use keepdims=True in numpy.mean to retain the original dimensions:

keepdims : bool, optional

If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.

If the default value is passed, then keepdims will not be passed through to the mean method of sub-classes of ndarray, however any non-default value will be. If the sub-class’ method does not implement keepdims any exceptions will be raised.

This already solves many broadcasting situations. However, if you need the same shape you can then use np.broadcast_to:

import numpy as np
a = np.array([[[0.1, 0.2, 0.3], [0.2, 0.3, 0.4]],
            [[0.4, 0.4, 0.4], [0.7, 0.6, 0.8]]], np.float32)
m1 = np.mean(a, axis=2, keepdims=True)
m2 = np.broadcast_to(m1, a.shape)
print('a.shape: %s\nm1.shape: %s\nm2.shape: %s' %(a.shape, m1.shape, m2.shape))

Output:

a.shape: (2, 2, 3)
m1.shape: (2, 2, 1)
m2.shape: (2, 2, 3)

Absent

I just tried this with numpy 1.26.2 and it does not work. What version of numpy are you using? – Ozonosphere

Thanks @Absent , so your code does what you expect it to but in terms of this question it would need to return a z.shape of (2, 2, 3). – Ozonosphere

@Ozonosphere My bad, I understood you only required to keep the original dimensions (which was the problem I was looking for when I came across this post). I updated my answer to keep the same shape as you require. – Absent

Thanks @Puco4! That makes sense. I had got excited there was such as simple solution but broadcasting seems like a good solution I'll try out. – Ozonosphere