Covariance Estimation including uncertainties

Questions
4

I believe you could get what you want by raveling X into a (200, 1) vector and using a single 200 x 200 covariance matrix with 2x2 covariance matrices on the (block) diagonal.

Actually setting this up in PyMC is slightly complex because aesara doesn’t have a dedicated block_diag function like scipy.linalg.block_diag. I use a scan here, but in my experience scans are awful and should be avoided. I’ve also done it with aesara.sparse.CSR pretty painlessly, but I don’t know where I put that code. No idea which is better, beyond my generally bad experiences with large scans.

import aesara
import aesara.tensor as at

def allocate_block(A, out, r_start, c_start, r_stride, c_stride):
    row_slice = slice(r_start, r_start + r_stride)
    col_slice = slice(c_start, c_start + c_stride)

    next_r = r_start + r_stride
    next_c = c_start + c_stride

    return at.set_subtensor(out[row_slice, col_slice], A), next_r, next_c

def block_diag(arr: at.tensor3):
    n, rows, cols = arr.shape
    out = at.zeros((n * rows, n * cols))
    
    result, _ = aesara.scan(allocate_block,
                            sequences=[arr],
                            outputs_info=[out, at.zeros(1, dtype='int64'), at.zeros(1, dtype='int64')],
                            non_sequences=[rows.astype('int64'), cols.astype('int64')])
    return result[0][-1]

A quick test:

arr = at.tensor3('cov_matrics')
with np.printoptions(linewidth=1000, precision=3, suppress=True):
    print(block_diag(arr).eval({arr:np.random.normal(size=(5,2,2))}))

[[ 1.518 -0.084  0.     0.     0.     0.     0.     0.     0.     0.   ]
 [-0.108 -0.433  0.     0.     0.     0.     0.     0.     0.     0.   ]
 [ 0.     0.    -0.08   1.759  0.     0.     0.     0.     0.     0.   ]
 [ 0.     0.     0.599  2.     0.     0.     0.     0.     0.     0.   ]
 [ 0.     0.     0.     0.    -1.505 -1.356  0.     0.     0.     0.   ]
 [ 0.     0.     0.     0.     0.706 -1.205  0.     0.     0.     0.   ]
 [ 0.     0.     0.     0.     0.     0.    -1.046  0.054  0.     0.   ]
 [ 0.     0.     0.     0.     0.     0.    -1.686 -0.323  0.     0.   ]
 [ 0.     0.     0.     0.     0.     0.     0.     0.     1.776  0.464]
 [ 0.     0.     0.     0.     0.     0.     0.     0.    -0.598  0.487]]

So as you can see, you get this banded matrix where every pair of observations is correlated but independent of the rest. I believe this corresponds to the situation you want. The PyMC code is straightforward as this point: somehow sample 100 covariance matrices, stack them into a (100, 2, 2) tensor, feed it into this block function to get a single (200 x 200) covariance matrix, ravel X so that it’s (200, 1) with the correct structure, then feed it all into an MvNormal.

You could then reshape the output of the MvNormal to be (100, 2) if you so desired.

There might also be a covariance function from the GP library you could use to get something like this with less hassle, but I’m not very well versed in all that. The main idea will be the same though: ravel and use a single structured covariance matrix in MvNormal.

2 Likes