ENH: Cython nancorr speeds up DataFrame.corr with method='pearson' by > 100x

Author: wesm

pandas/core/frame.py

@@ -3806,22 +3806,26 @@ def corr(self, method='pearson'):
         y : DataFrame
         """
         numeric_df = self._get_numeric_data()
-        mat = numeric_df.values.T
         cols = numeric_df.columns
+        mat = numeric_df.values
 
-        corrf = nanops.get_corr_func(method)
-        K = len(cols)
-        correl = np.empty((K, K), dtype=float)
-        mask = np.isfinite(mat)
-        for i, ac in enumerate(mat):
-            for j, bc  in enumerate(mat):
-                valid = mask[i] & mask[j]
-                if not valid.all():
-                    c = corrf(ac[valid], bc[valid])
-                else:
-                    c = corrf(ac, bc)
-                correl[i, j] = c
-                correl[j, i] = c
+        if method == 'pearson':
+            correl = lib.nancorr(mat)
+        else:
+            mat = mat.T
+            corrf = nanops.get_corr_func(method)
+            K = len(cols)
+            correl = np.empty((K, K), dtype=float)
+            mask = np.isfinite(mat)
+            for i, ac in enumerate(mat):
+                for j, bc  in enumerate(mat):
+                    valid = mask[i] & mask[j]
+                    if not valid.all():
+                        c = corrf(ac[valid], bc[valid])
+                    else:
+                        c = corrf(ac, bc)
+                    correl[i, j] = c
+                    correl[j, i] = c
 
         return self._constructor(correl, index=cols, columns=cols)
 

pandas/src/moments.pyx

@@ -247,7 +247,55 @@ def ewma(ndarray[double_t] input, double_t com):
 
     return output
 
-#-------------------------------------------------------------------------------
+#----------------------------------------------------------------------
+# Pairwise covariance
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def nancorr(ndarray[float64_t, ndim=2] mat):
+    cdef:
+        Py_ssize_t i, j, xi, yi, N, K
+        ndarray[float64_t, ndim=2] result
+        ndarray[uint8_t, ndim=2] mask
+        int64_t nobs = 0
+        float64_t vx, vy, sumx, sumy, sumxx, sumyy, meanx, meany
+
+    N, K = (<object> mat).shape
+
+    result = np.empty((K, K), dtype=np.float64)
+    mask = np.isfinite(mat).view(np.uint8)
+
+    for xi in range(K):
+        for yi in range(xi + 1):
+            nobs = sumxx = sumyy = sumx = sumy = 0
+            for i in range(N):
+                if mask[i, xi] and mask[i, yi]:
+                    vx = mat[i, xi]
+                    vy = mat[i, yi]
+                    nobs += 1
+                    sumx += vx
+                    sumy += vy
+
+            meanx = sumx / nobs
+            meany = sumy / nobs
+
+            # now the cov numerator
+            sumx = 0
+
+            for i in range(N):
+                if mask[i, xi] and mask[i, yi]:
+                    vx = mat[i, xi] - meanx
+                    vy = mat[i, yi] - meany
+
+                    sumx += vx * vy
+                    sumxx += vx * vx
+                    sumyy += vy * vy
+
+            result[xi, yi] = result[yi, xi] = sumx / sqrt(sumxx * sumyy)
+
+    return result
+
+#----------------------------------------------------------------------
 # Rolling variance
 
 def _check_minp(minp, N):