fa-xr-ufunc

fa-xr-ufunc.py

Y = xr.DataArray(Y_raw, dims=('observed_columns','rows'))
Y = Y.assign_coords({'observed_columns':np.arange(d),'rows':np.arange(n)})

def post_hoc_fun(W, psi, Y, rng=None):
    '''
    Recovers latent variables given observation Y and single draw of W, psi
    from the posterior distribution.
    '''
    k, d = W.shape[-2:]
    _, n = Y.shape
    WW = np.linalg.inv(np.einsum("...ji,...jk", W, W))
    WW_Wt = np.einsum("...ij,...kj", WW, W)
    F_mu = np.einsum("...ij,...jk", 1 / np.sqrt(psi) * WW_Wt, Y)
    WW_chol = np.linalg.cholesky(WW)
    return F_mu + np.einsum("...ij,...jk", WW_chol, rng.standard_normal(size=(k, n)))

xr.apply_ufunc(
    post_hoc_fun,
    trace.posterior['W'],
    trace.posterior['psi'],
    Y,
    input_core_dims=[['observed_columns', 'latent_columns'],
                     [],
                     ["observed_columns","rows"]],
    output_core_dims=[["latent_columns", "rows"]],
    kwargs={"rng": rng}
)