Cythonify slow for loops.

mblondel · mblondel · commit c59858d2fc7b · 2017-06-05T22:05:42.000+09:00
diff --git a/sdtw/distance.py b/sdtw/distance.py
@@ -2,6 +2,7 @@
 
 from sklearn.metrics.pairwise import euclidean_distances
 
+from .soft_dtw_fast import _jacobian_product_sq_euc
 
 class SquaredEuclidean(object):
 
@@ -45,15 +46,8 @@ def jacobian_product(self, E):
             Product with Jacobian
             ([m x d, m x n] * [m x n] = [m x d]).
         """
-        m = self.X.shape[0]
-        n = self.Y.shape[0]
-        d = self.X.shape[1]
-
         G = np.zeros_like(self.X)
 
-        for i in range(m):
-            for j in range(n):
-                for k in range(d):
-                    G[i, k] += E[i,j] * 2 * (self.X[i, k] - self.Y[j, k])
+        _jacobian_product_sq_euc(self.X, self.Y, E, G)
 
         return G
diff --git a/sdtw/soft_dtw_fast.pyx b/sdtw/soft_dtw_fast.pyx
@@ -99,3 +99,17 @@ def _soft_dtw_grad(np.ndarray[double, ndim=2] D,
             b = exp((R[i, j+1] - R[i, j] - D[i-1, j]) / gamma)
             c = exp((R[i+1, j+1] - R[i, j] - D[i, j]) / gamma)
             E[i, j] = E[i+1, j] * a + E[i, j+1] * b + E[i+1,j+1] * c
+
+
+def _jacobian_product_sq_euc(np.ndarray[double, ndim=2] X,
+                             np.ndarray[double, ndim=2] Y,
+                             np.ndarray[double, ndim=2] E,
+                             np.ndarray[double, ndim=2] G):
+    cdef int m = X.shape[0]
+    cdef int n = Y.shape[0]
+    cdef int d = X.shape[1]
+
+    for i in range(m):
+        for j in range(n):
+            for k in range(d):
+                G[i, k] += E[i,j] * 2 * (X[i, k] - Y[j, k])
