From a2bfee6e09a85dd6b3e04bcf7672f07dce527730 Mon Sep 17 00:00:00 2001
From: David Luevano Alvarado <55825613+luevano@users.noreply.github.com>
Date: Thu, 26 Mar 2020 15:49:04 -0700
Subject: Add wasserstein metric

---
 ml_exp/kernels.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/ml_exp/kernels.py b/ml_exp/kernels.py
index d593d83fd..3318fe6cf 100644
--- a/ml_exp/kernels.py
+++ b/ml_exp/kernels.py
@@ -21,6 +21,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
 import numpy as np
+from scipy.stats import wasserstein_distance
 try:
     import tensorflow as tf
     TF_AV = True
@@ -141,3 +142,28 @@ def laplacian_kernel(X1,
             K[i, :] = np.exp(i_sigma * norm)
 
     return K
+
+
+def wasserstein_kernel(X1,
+                       X2,
+                       alpha):
+    """
+    Calculates the Wasserstein Kernel.
+    X1: first representations.
+    X2: second representations.
+    alpha: wasserstein kernel parameter.
+    NOTE: this doesn't work with tensorflow.
+    """
+
+    if X2.ndim == 3:
+        raise TypeError('Representations must be 1D.')
+
+    X1_size = X1.shape[0]
+    X2_size = X2.shape[0]
+
+    K = np.zeros((X1_size, X2_size), dtype=np.float64)
+    for i in range(X1_size):
+        norm = np.array([X2[j] - X1[i] for j in range(X2_size)], dtype=np.float64)
+        K[i, :] = np.exp(- alpha * norm)
+
+    return K
-- 
cgit v1.2.3-70-g09d2