Rewrite kernel, gotta fix adj mat

1 files changed, 18 insertions, 7 deletions

diff --git a/ml_exp/kernels.py b/ml_exp/kernels.py
index 3914ffc20..e6855f97e 100644
--- a/ml_exp/kernels.py
+++ b/ml_exp/kernels.py
@@ -26,20 +26,31 @@ import numpy as np
 
 def gaussian_kernel(X1,
                     X2,
-                    sigma):
+                    sigma,
+                    opt=True):
     """
     Calculates the Gaussian Kernel.
     X1: first representations.
     X2: second representations.
     sigma: kernel width.
+    opt: if the optimized algorithm should be used. For benchmarking purposes.
     """
-    inv_sigma = -0.5 / (sigma*sigma)
+    i_sigma = -0.5 / (sigma*sigma)
 
     K = np.zeros((X1.shape[0], X2.shape[0]), dtype=float)
-    for i, x1 in enumerate(X1):
-        for j, x2 in enumerate(X2):
-            f_norm = np.linalg.norm(x1 - x2)
-            # print(f_norm)
-            K[i, j] = math.exp(inv_sigma * f_norm)
+    if opt:
+        # Faster way of calculating the kernel.
+        for i, x1 in enumerate(X1):
+            if X2.ndim == 3:
+                norm = np.linalg.norm(X2 - x1, axis=(1, 2))
+            else:
+                norm = np.linalg.norm(X2 - x1, axis=-1)
+            K[i, :] = np.exp(i_sigma * np.square(norm))
+    else:
+        for i, x1 in enumerate(X1):
+            for j, x2 in enumerate(X2):
+                f_norm = np.linalg.norm(x2 - x1)
+                # print(f_norm)
+                K[i, j] = math.exp(i_sigma * f_norm**2)
 
     return K