diff --git a/pandas/_libs/algos.pyx b/pandas/_libs/algos.pyx index b3c519ab99b6e..2de8573b13b91 100644 --- a/pandas/_libs/algos.pyx +++ b/pandas/_libs/algos.pyx @@ -2,10 +2,12 @@ import cython from cython import Py_ssize_t +from cython.parallel import prange from libc.stdlib cimport malloc, free from libc.string cimport memmove from libc.math cimport fabs, sqrt +from cpython cimport bool import numpy as np cimport numpy as cnp @@ -230,14 +232,15 @@ def kth_smallest(numeric[:] a, Py_ssize_t k) -> numeric: @cython.boundscheck(False) @cython.wraparound(False) -def nancorr(ndarray[float64_t, ndim=2] mat, bint cov=0, minp=None): +def nancorr(float64_t[:, :] mat, bint cov=0, minp=None, bool parallel=False): cdef: Py_ssize_t i, j, xi, yi, N, K bint minpv - ndarray[float64_t, ndim=2] result - ndarray[uint8_t, ndim=2] mask + float64_t[:, :] result + uint8_t[:, :] mask int64_t nobs = 0 float64_t vx, vy, sumx, sumy, sumxx, sumyy, meanx, meany, divisor + int64_t blah = 0 N, K = (mat).shape @@ -249,44 +252,82 @@ def nancorr(ndarray[float64_t, ndim=2] mat, bint cov=0, minp=None): result = np.empty((K, K), dtype=np.float64) mask = np.isfinite(mat).view(np.uint8) - with nogil: - 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 + if parallel: + with nogil: + for xi in prange(K, schedule='dynamic'): + nancorr_single_row(mat, N, K, result, xi, mask, minpv, cov) + else: + with nogil: + for xi in range(K): + nancorr_single_row(mat, N, K, result, xi, mask, minpv, cov) - if nobs < minpv: - result[xi, yi] = result[yi, xi] = NaN - else: - meanx = sumx / nobs - meany = sumy / nobs + return np.asarray(result) - # 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 +@cython.boundscheck(False) +@cython.wraparound(False) +cdef void nancorr_single_row(float64_t[:, :] mat, + Py_ssize_t N, + Py_ssize_t K, + float64_t[:, :] result, + Py_ssize_t xi, + uint8_t[:, :] mask, + bint minpv, + bint cov=0) nogil: + for yi in range(xi + 1): + nancorr_single(mat, N, K, result, xi, yi, mask, minpv, cov) - sumx += vx * vy - sumxx += vx * vx - sumyy += vy * vy - divisor = (nobs - 1.0) if cov else sqrt(sumxx * sumyy) +@cython.boundscheck(False) +@cython.wraparound(False) +cdef void nancorr_single(float64_t[:, :] mat, + Py_ssize_t N, + Py_ssize_t K, + float64_t[:, :] result, + Py_ssize_t xi, + Py_ssize_t yi, + uint8_t[:, :] mask, + bint minpv, + bint cov=0) nogil: + cdef: + Py_ssize_t i, j + int64_t nobs = 0 + float64_t vx, vy, sumx, sumy, sumxx, sumyy, meanx, meany, divisor - if divisor != 0: - result[xi, yi] = result[yi, xi] = sumx / divisor - else: - result[xi, yi] = result[yi, xi] = NaN + 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 + + if nobs < minpv: + result[xi, yi] = result[yi, xi] = NaN + else: + 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 + + divisor = (nobs - 1.0) if cov else sqrt(sumxx * sumyy) + + if divisor != 0: + result[xi, yi] = result[yi, xi] = sumx / divisor + else: + result[xi, yi] = result[yi, xi] = NaN - return result # ---------------------------------------------------------------------- # Pairwise Spearman correlation diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 7bbbdd70e062e..feef6cc5eac3e 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -6996,7 +6996,8 @@ def corr(self, method='pearson', min_periods=1): mat = numeric_df.values if method == 'pearson': - correl = libalgos.nancorr(ensure_float64(mat), minp=min_periods) + correl = libalgos.nancorr(ensure_float64(mat), minp=min_periods, + parallel=True) elif method == 'spearman': correl = libalgos.nancorr_spearman(ensure_float64(mat), minp=min_periods) diff --git a/setup.py b/setup.py index ed2d905f4358b..8e10d1f6e5181 100755 --- a/setup.py +++ b/setup.py @@ -9,6 +9,7 @@ import os from os.path import join as pjoin +import numpy import pkg_resources import platform from distutils.sysconfig import get_config_var @@ -677,10 +678,11 @@ def srcpath(name=None, suffix='.pyx', subdir='src'): obj = Extension('pandas.{name}'.format(name=name), sources=sources, depends=data.get('depends', []), - include_dirs=include, + include_dirs=include + [numpy.get_include()], language=data.get('language', 'c'), define_macros=data.get('macros', macros), - extra_compile_args=extra_compile_args) + extra_compile_args=['-fopenmp'] + extra_compile_args, + extra_link_args=['-fopenmp']) extensions.append(obj) @@ -704,12 +706,13 @@ def srcpath(name=None, suffix='.pyx', subdir='src'): np_datetime_sources), include_dirs=['pandas/_libs/src/ujson/python', 'pandas/_libs/src/ujson/lib', - 'pandas/_libs/src/datetime'], - extra_compile_args=(['-D_GNU_SOURCE'] + + 'pandas/_libs/src/datetime', + numpy.get_include()], + extra_compile_args=(['-D_GNU_SOURCE', '-fopenmp'] + extra_compile_args), + extra_link_args=['-fopenmp'], define_macros=macros) - extensions.append(ujson_ext) # ----------------------------------------------------------------------