```
# default_exp maxDist
```


```
#hide
#all_slow
```

# maxDist
> In this section we create a function to compute an approximation of the 2-Wasserstein distance between two univariate data sets


```
#hide
from dolfin import *
import numpy as np
import ot
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
# import matplotlib.cm as cm
# plt.rcParams['figure.figsize'] = (10,6)
# import required functions from oneDim
from statFEM.oneDim import mean_assembler, kernMat, cov_assembler, sample_gp
from scipy.stats import linregress
from scipy import integrate
from scipy.linalg import sqrtm
from tqdm.notebook import tqdm
import sympy; sympy.init_printing()
# code for displaying matrices nicely
def display_matrix(m):
    display(sympy.Matrix(m))
```

# Computing the 2-Wasserstein distance between two data-sets

We start by creating a function to estimate the 2-Wasserstein distance between two data-sets `a` and `b`, using the Python package [POT](https://github.com/PythonOT/POT).


```
#hide
def wass_old(a,b):
    # get the range of the data
    a_range = (a.min(), a.max())
    b_range = (b.min(), b.max())
    
    # set up bins for computation of 2-wasserstein distance
    x_min = min(a_range[0], b_range[0])
    x_max = max(a_range[1], b_range[1])
    x = np.arange(x_min - 2, x_max + 2)
    
    # get cost matrix
    n = len(x) - 1
    M = ot.dist(x[:-1].reshape((n,1)),x[:-1].reshape((n,1)))
    
    # get histograms for a and b
    a_h = np.histogram(a, bins=x, density=True)[0]
    b_h = np.histogram(b, bins=x, density=True)[0]
    
    # compute the 2-wasserstein distance using POT and return
    return np.sqrt(ot.emd2(a_h,b_h,M))
```


```
#export

import numpy as np
import ot

def wass(a,b,n_bins):
    "This function computes an approximation of the 2-Wasserstein distance between two datasets."
    # get the range of the data
    a_range = (a.min(), a.max())
    b_range = (b.min(), b.max())
    
    # get range for union of a and b
    x_min = min(a_range[0], b_range[0])
    x_max = max(a_range[1], b_range[1])
    
    # get histograms and bins
    a_h, bins_a = np.histogram(a, range=(x_min,x_max),bins=n_bins, density=True)
    b_h, bins_b = np.histogram(b, range=(x_min,x_max),bins=n_bins, density=True)
    
    # ensure bins_a and bins_b are equal
    assert (bins_a == bins_b).all()
    
    # get bin width
    width = bins_a[1] - bins_a[0]
    
    # normalise histograms so they sum to 1
    a_h *= width
    b_h *= width
    
    # get cost matrix
    n = len(bins_a) - 1
    M = ot.dist(bins_a[:-1].reshape((n,1)),bins_a[:-1].reshape((n,1)))
    
    # compute the 2-wasserstein distance using POT and return
    return np.sqrt(ot.emd2(a_h,b_h,M))
```

`wass` takes in the two datasets `a` and `b` as well as an argument `n_bin` which controls how many bins are used to create the histograms for the datasets.

Let's test this function out. First we make sure it gives $\operatorname{wass}(a,a) = 0$ for any dataset $a$.


```
# standard normal
N = 1000 # number of samples
n_bins = 10 # number of bins
np.random.seed(134)
a = np.random.normal(size=N)
assert wass(a,a,n_bins) == 0
```

We also test it on samples from 2 different Gaussians, $a\sim\mathcal{N}(m_0,s_0^{2})$ and $b\sim\mathcal{N}(m_1,s_1^{2})$. We expect, theoretically, that $\operatorname{wass}(a,b)=\sqrt{|m_0-m_1|^{2}+|s_0-s_1|^{2}}$.


```
# set up means and standard deviations
m_0 = 7
m_1 = 58
s_0 = 1.63
s_1 = 0.7

# draw the samples
N = 1000
#####################################
n_bins = 50 # number of bins
#####################################
np.random.seed(2321)
a = np.random.normal(loc = m_0, scale = s_0,size=N)
b = np.random.normal(loc = m_1, scale = s_1,size=N)

# tolerance for the comparison
tol = 1e-1

# compute the 2-wasserstein with our function and also the true theoretical value
W = wass(a,b,n_bins)
W_true = np.sqrt(np.abs(m_0-m_1)**2 + np.abs(s_0-s_1)**2)
# compare
assert np.abs(W - W_true) <= tol
```

Let's take the previous example and compute the distance for a range of different means and standard deviations.


```
# set up range for means and standard deviations
n = 40
m_range = np.linspace(m_0 - 2, m_0 + 2, n)
s_range = np.linspace(s_0/4, 2*s_0, n)

# set up arrays to hold results with our function, the theoretical results, 
# and theoretical results using estimated means and standard deviations
W = np.zeros((n, n))
W_0 = np.zeros((n, n))
W_est = np.zeros((n,n))

N = 10000 # number of samples
################################################
n_bins = 100 # number of bins
################################################
np.random.seed(2321)
a = np.random.normal(loc = m_0, scale = s_0,size=N)
m_a_est = np.mean(a)
s_a_est = np.std(a)

# sample for each m,s in the ranges and compute the results
for i, m in enumerate(m_range):
    for j, s in enumerate(s_range):
        b = np.random.normal(loc = m, scale = s, size = N)
        m_est = np.mean(b)
        s_est = np.std(b)
        
        W[i,j] = wass(a,b,n_bins)
        W_0[i,j] = np.sqrt(np.abs(m - m_0)**2 + np.abs(s - s_0)**2)
        W_est[i,j] = np.sqrt(np.abs(m_est - m_a_est)**2 + np.abs(s_est - s_a_est)**2)
```

Let's visualize the results:


```
#hide_input
M, S = np.meshgrid(m_range, s_range,indexing='ij')
plt.rcParams['figure.figsize'] = (16,5)
fig, axs = plt.subplots(ncols=4, gridspec_kw=dict(width_ratios=[4,4,4,0.2]))
axs[0].contourf(M, S, W)
axs[0].scatter([m_0],[s_0],marker='X',c='red')
axs[0].set_xlabel('$m$')
axs[0].set_ylabel('$s$')
axs[0].set_title('POT')

axs[1].contourf(M, S, W_est)
axs[1].scatter([m_0],[s_0],marker='X',c='red')
axs[1].set_xlabel('$m$')
axs[1].set_ylabel('$s$')
axs[1].set_title('Estimated truth')

axs[2].contourf(M, S, W_0)
axs[2].scatter([m_0],[s_0],marker='X',c='red')
axs[2].set_xlabel('$m$')
axs[2].set_ylabel('$s$')
axs[2].set_title('True')
fig.colorbar(axs[np.argmax([W.max(), W_est.max(),W_0.max()])].collections[0], cax=axs[3])
plt.tight_layout()
plt.show()
```

335cefcd7b440f2ac7a3ee6a532f5c2985692339
ipynb
Jupyter Notebook
06_maxDist.ipynb
YanniPapandreou/statFEM
189ddbb9c2f5a363d6e7e2f62a893cb3706e45bb
Apache-2.0
2022-02-04T09:26:33.000Z
2022-02-04T09:26:33.000Z
06_maxDist.ipynb
YanniPapandreou/statFEM
189ddbb9c2f5a363d6e7e2f62a893cb3706e45bb
Apache-2.0
06_maxDist.ipynb
YanniPapandreou/statFEM
189ddbb9c2f5a363d6e7e2f62a893cb3706e45bb
Apache-2.0
Qwen/Qwen-72B
1. YES
2. YES