combine_multivariate_gaussians.py

import numpy as np
from scipy.stats import multivariate_normal, chi2
import matplotlib.pyplot as plt

def get_prob(mean_1, cov_1, mean_2, cov_2):
    mean_diff = mean_2 - mean_1
    cov_sum = cov_1 + cov_2
    inv_cov_sum = np.linalg.inv(cov_sum)
    D_squared = mean_diff.T @ inv_cov_sum @ mean_diff
    k = len(mean_1)
    p = 1 - chi2.cdf(D_squared, df=k)
    return p

def product_of_densities(mean_1, cov_1, mean_2, cov_2):
    reusable_inverse = np.linalg.inv(cov_1 + cov_2)
    product_cov = cov_1 @ reusable_inverse @ cov_2
    product_mean = cov_2 @ reusable_inverse @ mean_1 + cov_1 @ reusable_inverse @ mean_2
    return product_mean, product_cov


xg = np.linspace(-3, 5, 100)
yg = np.linspace(-3, 5, 100)
X, Y = np.meshgrid(xg, yg)

def show_density(mean, cov, col='k', cmap='Greys'):
    Z = multivariate_normal.pdf(np.dstack((X, Y)), mean=mean, cov=cov)
    max_density = multivariate_normal.pdf(mean, mean=mean, cov=cov)

    percentiles = np.array([0.99, 0.95, 0.8, 0.6, 0.4, 0.2])
    levels = []
    for percentile in percentiles:
        chi2_val = chi2.ppf(percentile, 2)
        level = max_density * np.exp(-0.5 * chi2_val)
        levels.append(level)
    plt.contour(X, Y, Z, levels=levels, cmap=cmap, alpha=(1-percentiles)**0.3)
    plt.plot(mean[0], mean[1], col + 'x')
    plt.axis('equal')

mean_1 = np.array([0, 0])
mean_2 = np.array([2, 2])
cov_1 = np.eye(2)
cov_2 = np.array([[1, 0.0], [0.0, 0.1]])

print(get_prob(mean_1, cov_1, mean_2, cov_2))
mean_prod = product_of_densities(mean_1, cov_1, mean_2, cov_2)[0]
cov_prod = product_of_densities(mean_1, cov_1, mean_2, cov_2)[1]

print(mean_prod, cov_prod)

show_density(mean_1, cov_1, col='r', cmap='Reds')
show_density(mean_2, cov_2, col='b', cmap='Blues')
show_density(mean_prod, cov_prod, col='g', cmap='Greens')

plt.xlim(-2, 4)
plt.ylim(-2, 3)
plt.show()