Computing a variance/covariance matrix with groupby.map()

[GWeindel]

Hello,

I am trying to make use of xarray.DataArray.groupby().map() to compute a variance-covariance matrix but this reveals to be tricky I guess because xarray handles the dimension name, hence avoiding me using simple matrix multiplication.

A reproducible example:

import xarray as xr 
def vcov(x):
    return x@x.T

x = xr.DataArray([[[0,1,2],[1,0,1]],[[0,1,2],[1,0,1]]], dims=['a','b','c'])

With xarray I try:

print(x.groupby('c').map(vcov))

<xarray.DataArray (c: 3)>
array([ 2, 2, 10])

While calling to_numpy before gives me the expected shape but I have to use a for loop

x = x.to_numpy()
print(x @ x.T)

[[1 1]
[1 1]]
[[1 1]
[1 1]]
[[5 5]
[5 5]]

Is there a way to go from dims=['a','b','c'] to dims = ['a','a',c'] by using matrix multiplication?

[joshdorrington]

I thought perhaps using dummy coords would be a good approach, as in the below:

def vcov_da(x):
    y=x.rename({d:'dummy_'+d for d in x.dims})
    z=(x@y.T)
    z=z.sum(z.dims[1:-1])
    return z.rename({v1:v2 for v1,v2 in zip(z.dims,x.dims)})

Frustratingly it doesn't quite work, and I'm not totally sure why:

x.groupby('c').map(vcov_da):

xarray.DataArray a: 2 b: 2 c: 3
array([[[1, 1, 9],
        [1, 1, 9]],

       [[1, 1, 9],
        [1, 1, 9]]])

But perhaps the approach is in the right line, even if there is a linalg muck up in there somewhere.

[GWeindel]

That does look like a promising approach, I will fiddle around with it and keep you posted, thanks a lot! Le mar. 22 nov. 2022 à 16:48, Josh Dorrington ***@***.***> a écrit :

…

I thought perhaps using dummy coords would be a good approach, as in the below: def vcov_da(x): y=x.rename({d:'dummy_'+d for d in x.dims}) ***@***.***) z=z.sum(z.dims[1:-1]) return z.rename({v1:v2 for v1,v2 in zip(z.dims,x.dims)}) Frustratingly it doesn't quite work, and I'm not totally sure why: x.groupby('c').map(vcov_da): xarray.DataArray a: 2 b: 2 c: 3 array([[[1, 1, 9], [1, 1, 9]], [[1, 1, 9], [1, 1, 9]]]) But perhaps the approach is in the right line, even if there is a linalg muck up in there somewhere. — Reply to this email directly, view it on GitHub <#7290 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AG3LE52UBMKM3WLLKIRPNBTWJTTL3ANCNFSM6AAAAAASAYLNNA> . You are receiving this because you authored the thread.Message ID: ***@***.***>