Parallel OpenMP Cython example

openmp.py

# distutils: extra_compile_args = -openmp
# distutils: extra_link_args = -openmp

from cython cimport boundscheck, wraparound
from cython.parallel cimport prange, threadid
from libc.stdio cimport  stdout, fprintf

import numpy as np

cdef inline double norm2(double complex z) nogil:
	return z.real * z.real + z.imag * z.imag


cdef int escape(double complex z,
				double complex c,
				double z_max,
				int n_max) nogil:

	cdef:
		int i = 0
		double z_max2 = z_max * z_max

	while norm2(z) < z_max2 and i < n_max:
		z = z * z + c
		i += 1

	return i


@boundscheck(False)
@wraparound(False)
def calc_julia(int resolution, double complex c,
			   double bound=1.5, double z_max=4.0, int n_max=1000):

	cdef:
		double step = 2.0 * bound / resolution
		int i, j
		double complex z
		double real, imag
		int[:, ::1] counts

	counts = np.zeros((resolution+1, resolution+1), dtype=np.int32)

	for i in prange(resolution + 1, nogil=True,
					schedule='static', chunksize=1):
		real = -bound + i * step
		for j in range(resolution + 1):
			imag = -bound + j * step
			z = real + imag * 1j
			counts[i,j] = escape(z, c, z_max, n_max)

	return np.asarray(counts)

@boundscheck(False)
@wraparound(False)
def julia_fraction_0(int[:,::1] counts, int maxval=1000):
	cdef:
		int total = 0
		int i, j, N, M
	N = counts.shape[0]; M = counts.shape[1]

	for i in prange(N, nogil=True):
		for j in range(M):
			if counts[i,j] == maxval:
				total += 1
	return total / float(counts.size)




	
@boundscheck(False)
@wraparound(False)
def julia_fraction(int[:,::1] counts, int maxval=1000):
	cdef:
		unsigned int thread_id
		int total = 0
		int i, j, N, M
	N = counts.shape[0]; M = counts.shape[1]

	for i in prange(N, nogil=True):
		thread_id = threadid()
		#fprintf(stdout, "%d\n", <int>thread_id)
		for j in range(M):
			if counts[i,j] == maxval:
				total += 1
	return total / float(counts.size)	

run.py

import julia
import numpy as np
from time import clock
import matplotlib.pyplot as plt

t1 = clock()
jl = julia.calc_julia(1000, (0.322 + 0.05j))
print "time:", clock() - t1

print "julia fraction:", julia.julia_fraction(jl)

plt.imshow(np.log(jl))
plt.show()

setup.py

from distutils.core import setup
from Cython.Build import cythonize

setup(name="julia",
      ext_modules=cythonize("julia.pyx"))