CODEWITHSUNDEEP
arrayspython-3.xnumpy
I Khalil
I Khalil

Reputation: 31

NumPy: Replace all elements along third dimension with averages in 3D array

I have a 3D array of the dimensions 400*800*3. I want to replace all the elements in the third dimension with the average of that dimension. Right now I have been able to accomplish this by using loops.

test_data=np.random.randint(0,256,size=(400,800,3))
for i in range(400):
    for j in range(800):
        mn = np.mean(test_data[i,j])
        test_data[i,j]=mn

I want to know if there is a, efficient and less verbose way of achieving the same objective. Thanks.

Also, what if instead of a mean I would like a weighted average of the three elements i.e the last line should be replaced by

test_data[i,j,0]=test_data[i,j,0]*0.2
test_data[i,j,1]=test_data[i,j,1]*0.5
test_data[i,j,2]=test_data[i,j,2]*0.3

Upvotes: 3

Views: 2693

Answers (3)

Divakar
Divakar

Reputation: 221574

Get the mean values along the last axis and assign into all places with [:] -

test_data[:] = test_data.mean(axis=-1,keepdims=1)

Alternatively, we can create a new array with replication along the last axis -

mean_vals = test_data.mean(axis=-1,keepdims=1).astype(test_data.dtype)
test_data_out = np.repeat(mean_vals,3,axis=-1)

For a read-only version and much faster way, use np.broadcast_to -

test_data_out = np.broadcast_to(mean_vals, test_data.shape)

For the weighted average part, if you meant :

test_data[i,j]= test_data[i,j,0]*0.2 + test_data[i,j,1]*0.5 + test_data[i,j,2]*0.3

For that, we could use multi-dim tensor reduction : np.tensordot -

W = [0.2,0.5,0.3] # weights
W_mean = np.tensordot(test_data,W, axes=((-1,-1)))[...,None]
test_data[:] = W_mean.astype(test_data.dtype)

Instead, if you meant :

test_data[i,j,0]=test_data[i,j,0]*0.2
test_data[i,j,1]=test_data[i,j,1]*0.5
test_data[i,j,2]=test_data[i,j,2]*0.3

For that, there's no sum-reduction, so we could simply make use of broadcasting -

test_data[:] = (test_data*[0.2, 0.5, 0.3]).astype(test_data.dtype)

Upvotes: 3

hpaulj
hpaulj

Reputation: 231385

In [556]: test_data=np.arange(24.).reshape(2,4,3)
In [557]: test_data
Out[557]: 
array([[[  0.,   1.,   2.],
        [  3.,   4.,   5.],
        [  6.,   7.,   8.],
        [  9.,  10.,  11.]],

       [[ 12.,  13.,  14.],
        [ 15.,  16.,  17.],
        [ 18.,  19.,  20.],
        [ 21.,  22.,  23.]]])

Scaling the values along the last dimension is just a multiplication:

In [558]: test_data*np.array([0.2, 0.5, 0.3])
Out[558]: 
array([[[  0. ,   0.5,   0.6],
        [  0.6,   2. ,   1.5],
        [  1.2,   3.5,   2.4],
        [  1.8,   5. ,   3.3]],

       [[  2.4,   6.5,   4.2],
        [  3. ,   8. ,   5.1],
        [  3.6,   9.5,   6. ],
        [  4.2,  11. ,   6.9]]])

It could be done in place with test_data *= np.array([....])

Taking the mean on the last dimension is just:

In [559]: test_data.mean(axis=-1)
Out[559]: 
array([[  1.,   4.,   7.,  10.],
       [ 13.,  16.,  19.,  22.]])

This is 2d. @divakar has shown how to keep it 3d (keep_dims) or even full size (2,4,1) or even (2,4,3).

In [561]: test_data.mean(axis=-1,keepdims=True)
Out[561]: 
array([[[  1.],
        [  4.],
        [  7.],
        [ 10.]],

       [[ 13.],
        [ 16.],
        [ 19.],
        [ 22.]]])

np.average lets you specify weights (weighted average):

In [568]: np.average(test_data, axis=-1, weights=[0.2,0.5,0.3])
Out[568]: 
array([[  1.1,   4.1,   7.1,  10.1],
       [ 13.1,  16.1,  19.1,  22.1]])

This is the same as taking a dot product on the last dimension

In [569]: np.dot(test_data, [0.2,0.5,0.3])
Out[569]: 
array([[  1.1,   4.1,   7.1,  10.1],
       [ 13.1,  16.1,  19.1,  22.1]])

Upvotes: 0

6502
6502

Reputation: 114521

This should work:

test_data[:,:,0] += test_data[:,:,1]
test_data[:,:,0] += test_data[:,:,2]
test_data[:,:,0] *= 1.0/3
test_data[:,:,1] = test_data[:,:,0]*0.5
test_data[:,:,2] = test_data[:,:,0]*0.3
test_data[:,:,0] *= 0.2

This approach will work "inplace" without requiring extra memory

Upvotes: 0

Related Questions