10

I have a 3D numpy array of shape (t, n1, n2):

x = np.random.rand(10, 2, 4)

I need to calculate another 3D array y which is of shape (t, n1, n1) such that: 

y[0] = np.cov(x[0,:,:])

...and so on for all slices along the first axis.

So, a loopy implementation would be:

y = np.zeros((10,2,2))
for i in np.arange(x.shape[0]):
    y[i] = np.cov(x[i, :, :])

Is there any way to vectorize this so I can calculate all covariance matrices in one go? I tried doing:

x1 = x.swapaxes(1, 2)
y = np.dot(x, x1)

But it didn't work.

Matt Hall
  • 7,614
  • 1
  • 23
  • 36
dayum
  • 1,073
  • 15
  • 31

1 Answers1

10

Hacked into numpy.cov source code and tried using the default parameters. As it turns out, np.cov(x[i,:,:]) would be simply :

N = x.shape[2]
m = x[i,:,:]
m -= np.sum(m, axis=1, keepdims=True) / N
cov = np.dot(m, m.T)  /(N - 1)

So, the task was to vectorize this loop that would iterate through i and process all of the data from x in one go. For the same, we could use broadcasting at the third step. For the final step, we are performing sum-reduction there along all slices in first axis. This could be efficiently implemented in a vectorized manner with np.einsum. Thus, the final implementation came to this -

N = x.shape[2]
m1 = x - x.sum(2,keepdims=1)/N
y_out = np.einsum('ijk,ilk->ijl',m1,m1) /(N - 1)

Runtime test

In [155]: def original_app(x):
     ...:     n = x.shape[0]
     ...:     y = np.zeros((n,2,2))
     ...:     for i in np.arange(x.shape[0]):
     ...:         y[i]=np.cov(x[i,:,:])
     ...:     return y
     ...: 
     ...: def proposed_app(x):
     ...:     N = x.shape[2]
     ...:     m1 = x - x.sum(2,keepdims=1)/N
     ...:     out = np.einsum('ijk,ilk->ijl',m1,m1)  / (N - 1)
     ...:     return out
     ...: 

In [156]: # Setup inputs
     ...: n = 10000
     ...: x = np.random.rand(n,2,4)
     ...: 

In [157]: np.allclose(original_app(x),proposed_app(x))
Out[157]: True  # Results verified

In [158]: %timeit original_app(x)
1 loops, best of 3: 610 ms per loop

In [159]: %timeit proposed_app(x)
100 loops, best of 3: 6.32 ms per loop

Huge speedup there!

dayum
  • 1,073
  • 15
  • 31
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Thanks Divakar, looking into it in detail now to understand einsum. – dayum Nov 03 '16 at 23:23
  • I just changed N/(N**2-N) to 1/(N-1) in your code to align with standard formula and ease of understanding – dayum Nov 03 '16 at 23:34
  • @Divakar, could you please explain in more detail this step: `m1 = x - x.sum(2,keepdims=1)/N`? Changing the `.sum` axis, does not seem to affect the final answer... – Newskooler Dec 12 '18 at 20:40
  • Another way, not as fast as the einsum way is: ```out = (m1 @ m1.transpose(0,2,1)) / (N - 1)``` which follows from the definition of the [sample covariance matrix](https://en.wikipedia.org/wiki/Covariance_matrix#Estimation). – Javier TG Sep 05 '21 at 14:45