[FIXED] Calculate correlation between 4d numpy arrays in 4th dimension

Issue

I have two 4d numpy arrays and some of the rows will be all zero:

a = np.random.rand(360).reshape(4, 5, 3, 6)
b = np.random.rand(360).reshape(4, 5, 3, 6)
a[0,1,0,:]=0
b[0,1,0,:]=0

I want to calculate the correlation between the two 4d numpy arrays in the forth dimension, so that I can get a 3D numpy array with dimension of (4, 5, 3). Each cell represent the correlation value in the new 3D numpy array.

For example the new 3D array [0, 0, 0] position is calculated from : np.corrcoef(a[0, 0, 0, :], b[0, 0, 0, :][1][0]

Instead of loop through all the positions, is there a effecient way to do this?

Moreover, since some of my 4th dimension are all zero, the np.corrcoef produces nan,I also need to change those "nan" into 0.

Thanks in advance. Any insight is helpful!

Solution

You can just replicate the formula:

corrs = (a*b).mean(axis=-1) - a.mean(axis=-1) * b.mean(axis=-1)
std_prod =  np.sqrt(np.var(a, axis=-1) * np.var(b, axis=-1))

out = corrs / std_prod

out = np.where(np.isnan(out), 0, out)

Answered By - Quang Hoang