Condition a multivariate Gaussian on observations

condition.py

def condition(mu, sigma, observations):
    """
    Condition on the observations to produce a new mean and covariance matrix
    Parameters:
    ----------
        mu - the mean vector
        sigma - the covariance matrix
        observations - the observed variables. 
            Variables that are not observed should be Nan in this vector. 
            The resulting distribution will be over these variables in the order they appear
            
    returns:
    -------
        mu_bar - the conditional mean over the unobserved variables
        sigma_bar - the conditional covariance over unobserved variables
    """
    not_obs = np.isnan(observations) # masks the unobserved variables as true
    obs = np.logical_not(not_obs) # masks the observed variables as true
        
    a = observations[obs] # observed variables only
    mu1 = mu[not_obs]
    mu2 = mu[obs]
    sigma_11 = sigma[np.ix_(not_obs, not_obs)]
    sigma_22 = sigma[np.ix_(obs, obs)]
    sigma_21 = sigma[np.ix_(obs, not_obs)]
    sigma_12 = sigma[np.ix_(not_obs, obs)]
        
    mu_bar = mu1 + sigma_12.dot(np.linalg.inv(sigma_22)).dot(a - mu2)
    sigma_bar = sigma_11 - sigma_12.dot(np.linalg.inv(sigma_22)).dot(sigma_21)
    return mu_bar, sigma_bar