Last active
October 24, 2019 13:43
-
-
Save anyuzx/e8f8950ed6fcc901a80c65aec28aabba to your computer and use it in GitHub Desktop.
setup.py for parallel Cython codes shown in the blog post
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# cython: language_level=3 | |
import cython | |
import numpy as np | |
from cython.parallel import prange, parallel | |
from libc.math cimport sqrt | |
from libc.math cimport nearbyint | |
@cython.boundscheck(False) | |
@cython.wraparound(False) | |
@cython.cdivision(True) # used to not checking division, can have 20% performance speed up | |
def pdist_cy_omp(double [:,:] positions not None, double l): | |
cdef Py_ssize_t n = positions.shape[0] | |
cdef Py_ssize_t ndim = positions.shape[1] | |
pdistances = np.zeros(n * (n-1) // 2, dtype = np.float64) | |
cdef double [:] pdistances_view = pdistances | |
cdef double d, dd | |
cdef Py_ssize_t i, j, k | |
with nogil, parallel(): | |
for i in prange(n-1, schedule='dynamic'): | |
for j in range(i+1, n): | |
dd = 0.0 | |
for k in range(ndim): | |
d = positions[i,k] - positions[j,k] | |
d = d - nearbyint(d / l) * l | |
dd = dd + d * d | |
pdistances_view[j - 1 + (2 * n - 3 - i) * i // 2] = sqrt(dd) | |
return pdistances |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from distutils.core import setup | |
from distutils.extension import Extension | |
from Cython.Build import cythonize | |
ext_modules = [ | |
Extension( | |
name="pdist_omp", | |
sources=["pdist_omp.pyx"], | |
library_dirs=['/usr/local/lib/gcc/9'], | |
include_dirs=['/usr/local/include'], | |
extra_compile_args=['-O3', '-ffast-math', '-fopenmp'], | |
extra_link_args=['-fopenmp'], | |
) | |
] | |
setup( | |
name='pdist-omp', | |
ext_modules = cythonize(ext_modules), | |
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment