wasserstein.py

import matplotlib.pyplot as plt
import numpy as np
from scipy import stats
from scipy.stats import wasserstein_distance,beta


wasserstein_distance([1,2,3,4],[1,2,3,4,4])
x = np.linspace(0, 1, 100)
dist1 = stats.beta.pdf(x,5,5)
dist2 = stats.beta.pdf(x,8,5)
ws_distance = wasserstein_distance(dist1,dist2)


fig, ax = plt.subplots(1, figsize=(8, 6))
ax.fill_between(x, dist1, alpha=0.5) 
ax.fill_between(x, dist2, alpha=0.5) 
ax.plot(0, 0, label="Wasserstein Distance \n = {:3.2f}".format(ws_distance), alpha=0)
ax.legend(loc=2, fontsize=15).get_frame().set_alpha(0)