combination of numpy+scipy and cython

Question

I am trying to write a function to calculate the correlation matrix of a NumPy array with for loops and the pearsonr and the spearmanr functions from SciPy to calculate the matrix (I know that pandas do the same!!!). I am trying to use Cython to improve the performance of the function, as its current speed is not even comparable to its pandas equivalent (pd.corr()). I am not even sure if this is the way to do it, so, feel free to suggest anything. here is te code for .pyx and setup.py and also for the performance evaluation with three functions: pandas corr(), one custom corr function in python, and the same function in cython with types.

main.pyx

import numpy as np
from scipy.stats import pearsonr
cimport numpy as np


np.import_array()


DTYPE = np.float64


ctypedef np.float64_t DTYPE_t


def cor_custom_c(np.ndarray[DTYPE_t, ndim=2] x):
    cdef int dim = x.shape[1]
    cdef np.ndarray[DTYPE_t, ndim=2] mat = np.ones([dim, dim], dtype=DTYPE)
    cdef int i, j
    cdef double pv
    x = x.T
    for i in range(dim):
        for j in range(dim):
            if i <= j:
                continue
            else:
                pv = pearsonr(x[i], x[j])[0]
            mat[i, j] = pv
            mat[j, i] = pv
    return mat

setup.py

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy

setup(name='Main',
      ext_modules=cythonize([Extension("main", ["main.pyx"], include_dirs=[numpy.get_include()])])
)

check.py

import time
import main
import numpy as np
from scipy.stats import pearsonr
import pandas as pd


def cor_custom(x):
    mat = np.ones((x.shape[1], x.shape[1]))
    dim = x.shape[1]
    x = x.T
    for i in range(dim):
        for j in range(dim):
            if i <= j:
                continue
            else:
                pv = pearsonr(x[i], x[j])[0]
            mat[i, j] = pv
            mat[j, i] = pv
    return mat


data = {}
sample_size = 100
n_feature = 100
print('Sample size: ', sample_size)
print('Number of Features: ', n_feature)
for i in range(n_feature):
    data['col_'+str(i)] = np.random.randint(0, 10, sample_size)
df = pd.DataFrame.from_dict(data)

xx = df.to_numpy(dtype=np.float64, na_value=np.nan, copy=False)
st = time.time()
pd_cr = df.corr(method='pearson')
print('pandas: ', "----%.5f----" % (time.time()-st))

st = time.time()
num_cr = cor_custom(x=xx)
print('numpy only: ', "----%.5f----" % (time.time()-st))

st = time.time()
cython_cr = main.cor_custom_c(xx)
print('cython: ', "----%.5f----" % (time.time()-st))

print('Check pandas corr() equal to cython corr(): ', all(cython_cr == pd_cr))

print('numpy+scipy+cython: ', "----%.5f----" % (time.time()-st))

The performance is as follows:

Sample size: 100
Number of Features: 100
pandas: ----0.00282----
numpy only: ----0.32579----
numpy+scipy+cython: ----0.32199----

combination of numpy+scipy and cython

0 Answers