Added parallelization to _project.pyx code

This commit is contained in:
Guilhem Lavaux 2014-06-11 11:14:33 +02:00
parent 5b133a9ac8
commit 401ddc8a8b
9 changed files with 173 additions and 122 deletions

View file

@ -1,5 +1,6 @@
from cpython cimport bool
from cython cimport view
from cython.parallel import prange, parallel
from libc.math cimport sin, cos, abs, floor, sqrt
import numpy as np
cimport numpy as npx
@ -620,12 +621,13 @@ cdef DTYPE_t integrator1(DTYPE_t[:,:,:] density,
# return cube_integral(u, u0, jumper)*d
return cube_integral_trilin(u, u0, jumper, vertex_value)
@cython.boundscheck(False)
def line_of_sight_projection(npx.ndarray[DTYPE_t, ndim=3] density,
npx.ndarray[DTYPE_t] a_u,
cdef DTYPE_t C_line_of_sight_projection(DTYPE_t[:,:,:] density,
DTYPE_t a_u[3],
DTYPE_t min_distance,
DTYPE_t max_distance, DTYPE_t[:] shifter, int integrator_id=0):
DTYPE_t max_distance, DTYPE_t[:] shifter, int integrator_id) nogil except? 0:
cdef DTYPE_t u[3], ifu0[3], u0[3], utot[3]
cdef int u_delta[3]
@ -657,14 +659,17 @@ def line_of_sight_projection(npx.ndarray[DTYPE_t, ndim=3] density,
u0[i] = ifu0[i]-iu0[i]
u_delta[i] = 1 if iu0[i] > 0 else -1
if (not ((iu0[i]>= 0) and (iu0[i] < N))):
raise RuntimeError("iu0[%d] = %d !!" % (i,iu0[i]))
with gil:
raise RuntimeError("iu0[%d] = %d !!" % (i,iu0[i]))
if (not (u0[i]>=0 and u0[i]<=1)):
raise RuntimeError("u0[%d] = %g !" % (i,u0[i]))
with gil:
raise RuntimeError("u0[%d] = %g !" % (i,u0[i]))
completed = 0
if ((int(iu0[0]) >= N) or (int(iu0[0]) <= 0) or
(int(iu0[1]) >= N) or (int(iu0[1]) <= 0) or
(int(iu0[2]) >= N) or (int(iu0[2]) <= 0)):
if ((iu0[0] >= N) or (iu0[0] <= 0) or
(iu0[1] >= N) or (iu0[1] <= 0) or
(iu0[2] >= N) or (iu0[2] <= 0)):
completed = 1
I0 = 0
@ -681,9 +686,9 @@ def line_of_sight_projection(npx.ndarray[DTYPE_t, ndim=3] density,
u0[jumper[0]] = 0
if ((int(iu0[0]) >= N) or (int(iu0[0]) <= 0) or
(int(iu0[1]) >= N) or (int(iu0[1]) <= 0) or
(int(iu0[2]) >= N) or (int(iu0[2]) <= 0)):
if ((iu0[0] >= N) or (iu0[0] <= 0) or
(iu0[1] >= N) or (iu0[1] <= 0) or
(iu0[2] >= N) or (iu0[2] <= 0)):
completed = 1
else:
dist2 = 0
@ -699,35 +704,72 @@ def line_of_sight_projection(npx.ndarray[DTYPE_t, ndim=3] density,
return I0
def line_of_sight_projection(DTYPE_t[:,:,:] density not None,
DTYPE_t[:] a_u not None,
DTYPE_t min_distance,
DTYPE_t max_distance, DTYPE_t[:] shifter not None, int integrator_id=0):
cdef DTYPE_t u[3]
u[0] = a_u[0]
u[1] = a_u[1]
u[2] = a_u[2]
C_line_of_sight_projection(density,
u,
min_distance,
max_distance, shifter, integrator_id)
cdef double _spherical_projloop(double theta, double phi, DTYPE_t[:,:,:] density,
double min_distance, double max_distance,
DTYPE_t[:] shifter, int integrator_id) nogil:
cdef DTYPE_t u0[3]
stheta = sin(theta)
u0[0] = cos(phi)*stheta
u0[1] = sin(phi)*stheta
u0[2] = cos(theta)
return C_line_of_sight_projection(density, u0, min_distance, max_distance, shifter, integrator_id)
@cython.boundscheck(False)
def spherical_projection(int Nside,
npx.ndarray[DTYPE_t, ndim=3] density not None,
DTYPE_t min_distance,
DTYPE_t max_distance, int progress=1, int integrator_id=0, DTYPE_t[:] shifter = None):
DTYPE_t max_distance, int progress=1, int integrator_id=0, DTYPE_t[:] shifter = None, int booster=10):
import healpy as hp
import progressbar as pb
cdef int i
cdef npx.ndarray[DTYPE_t, ndim=1] u
cdef npx.ndarray[DTYPE_t, ndim=1] outm
cdef DTYPE_t[:] theta,phi
cdef DTYPE_t[:,:,:] density_view
cdef DTYPE_t[:] outm
cdef npx.ndarray[DTYPE_t, ndim=1] outm_array
cdef long N
cdef double stheta
if shifter is None:
shifter = view.array(shape=(3,), format=FORMAT_DTYPE, itemsize=sizeof(DTYPE_t))
shifter[:] = 0
outm = np.empty(hp.nside2npix(Nside),dtype=DTYPE)
outm_array = np.empty(hp.nside2npix(Nside),dtype=DTYPE)
outm = outm_array
density_view = density
if progress:
p = pb.ProgressBar(maxval=outm.size).start()
for i in range(outm.size):
if progress:
p.update(i)
u = np.array(hp.pix2vec(Nside, i), dtype=DTYPE)
outm[i] = line_of_sight_projection(density, u, min_distance, max_distance, shifter, integrator_id=integrator_id)
if progress != 0:
p = pb.ProgressBar(maxval=outm.size,widgets=[pb.BouncingBar(), pb.ETA()]).start()
N = outm.size
theta,phi = hp.pix2ang(Nside, np.arange(N))
with nogil, parallel():
for i in prange(N):
if progress != 0 and (i%booster) == 0:
with gil:
p.update(i)
outm[i] = _spherical_projloop(theta[i], phi[i], density_view, min_distance, max_distance, shifter, integrator_id)
if progress:
p.finish()
return outm
return outm_array