From 401ddc8a8b407de08568dc5b4c099094cb0b78e4 Mon Sep 17 00:00:00 2001 From: Guilhem Lavaux Date: Wed, 11 Jun 2014 11:14:33 +0200 Subject: [PATCH] Added parallelization to _project.pyx code --- python/_project.pyx | 90 +++++++++++++++++------- python/cosmotool/__init__.py | 5 +- python/cosmotool/grafic.py | 26 ++++++- python/cosmotool/timing.py | 15 ++++ python_sample/build_nbody_skymap.py | 58 +-------------- python_sample/icgen/borgicgen.py | 72 ++++++++++++++++--- python_sample/icgen/gen_ic_from_borg.py | 2 +- python_sample/icgen/test_ic_from_borg.py | 6 +- python_sample/test_spheric_proj.py | 21 +----- 9 files changed, 173 insertions(+), 122 deletions(-) create mode 100644 python/cosmotool/timing.py mode change 100644 => 120000 python_sample/build_nbody_skymap.py mode change 100644 => 120000 python_sample/test_spheric_proj.py diff --git a/python/_project.pyx b/python/_project.pyx index c3d39a0..138eaa2 100644 --- a/python/_project.pyx +++ b/python/_project.pyx @@ -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 diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py index 7a1fad3..3db5697 100644 --- a/python/cosmotool/__init__.py +++ b/python/cosmotool/__init__.py @@ -1,8 +1,7 @@ from _cosmotool import * from _project import * -from grafic import writeGrafic, writeWhitePhase, readGrafic +from grafic import writeGrafic, writeWhitePhase, readGrafic, readWhitePhase from borg import read_borg_vol from cic import cicParticles from simu import loadRamsesAll, simpleWriteGadget, SimulationBare - - +from timing import timeit diff --git a/python/cosmotool/grafic.py b/python/cosmotool/grafic.py index 7eb5f41..65ac6c1 100644 --- a/python/cosmotool/grafic.py +++ b/python/cosmotool/grafic.py @@ -42,9 +42,9 @@ def writeGrafic(filename, field, BoxSize, scalefac, **cosmo): bad, bad, bad, scalefac, cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h'], checkPoint)) - checkPoint = 4*Ny*Nx + checkPoint = 4*Ny*Nz field = field.reshape(field.shape, order='F') - for i in xrange(Nz): + for i in xrange(Nx): f.write(struct.pack("I", checkPoint)) f.write(field[i].astype(np.float32).tostring()) f.write(struct.pack("I", checkPoint)) @@ -58,9 +58,29 @@ def writeWhitePhase(filename, field): checkPoint = 4*4 f.write(struct.pack("IIIIII", checkPoint, Nx, Ny, Nz, 0, checkPoint)) - checkPoint = struct.pack("I", 4*Ny*Nz) + field = field.reshape(field.shape, order='F') + checkPoint = struct.pack("I", 4*Nx*Ny) for i in xrange(Nx): f.write(checkPoint) f.write(field[i].astype(np.float32).tostring()) f.write(checkPoint) + + +def readWhitePhase(filename): + with file(filename, mode="rb") as f: + _, Nx, Ny, Nz, _, _ = struct.unpack("IIIIII", f.read(4*4+2*4)) + + a = np.empty((Nx,Ny,Nz), dtype=np.float32) + + checkPoint_ref = 4*Ny*Nz + + for i in xrange(Nx): + if struct.unpack("I", f.read(4))[0] != checkPoint_ref: + raise ValueError("Invalid unformatted access") + + a[i, :, :] = np.fromfile(f, dtype=np.float32, count=Ny*Nz).reshape((Ny, Nz), order='F') + if struct.unpack("I", f.read(4))[0] != checkPoint_ref: + raise ValueError("Invalid unformatted access") + + return a diff --git a/python/cosmotool/timing.py b/python/cosmotool/timing.py new file mode 100644 index 0000000..58e319d --- /dev/null +++ b/python/cosmotool/timing.py @@ -0,0 +1,15 @@ +import time + +def timeit(method): + + def timed(*args, **kw): + ts = time.time() + result = method(*args, **kw) + te = time.time() + + print '%r (%r, %r) %2.2f sec' % \ + (method.__name__, args, kw, te-ts) + return result + + return timed + diff --git a/python_sample/build_nbody_skymap.py b/python_sample/build_nbody_skymap.py deleted file mode 100644 index 2569a05..0000000 --- a/python_sample/build_nbody_skymap.py +++ /dev/null @@ -1,57 +0,0 @@ -import healpy as hp -import numpy as np -import cosmotool as ct -import h5py as h5 -from matplotlib import pyplot as plt - -L=600. -Nside=128 - -INDATA="/nethome/lavaux/Copy/PlusSimulation/BORG/Input_Data/2m++.npy" -tmpp = np.load(INDATA) - -def build_sky_proj(density, dmax=60.,dmin=0): - - N = density.shape[0] - ix = (np.arange(N)-0.5)*L/N - 0.5 * L - - - dist2 = (ix[:,None,None]**2 + ix[None,:,None]**2 + ix[None,None,:]**2) - - flux = density.transpose().astype(ct.DTYPE) # / dist2 - dmax=N*dmax/L - dmin=N*dmin/L - projsky1 = ct.spherical_projection(Nside, flux, dmin, dmax, integrator_id=1) -# projsky0 = ct.spherical_projection(Nside, flux, 0, 52, integrator_id=0) - - return projsky1*L/N#,projsky0 - -l,b = tmpp['gal_long'],tmpp['gal_lat'] - -l = np.radians(l) -b = np.pi/2 - np.radians(b) - -dcmb = tmpp['velcmb']/100. - -idx = np.where((dcmb>10)*(dcmb<60)) - -plt.figure(1) -plt.clf() -if False: - with h5.File("fields.h5", mode="r") as f: - d = f["density"][:] - d /= np.average(np.average(np.average(d,axis=0),axis=0),axis=0) - proj = build_sky_proj(f["density"][:], dmin=10,dmax=60.) -else: - d = np.load("icgen/dcic0.npy") - proj0 = build_sky_proj(1+d, dmin=10,dmax=60.) - d = np.load("icgen/dcic1.npy") - proj1 = build_sky_proj(1+d, dmin=10,dmax=60.) - -hp.mollview(proj0, fig=1, coord='CG', max=60, cmap=plt.cm.coolwarm) -hp.projscatter(b[idx], l[idx], lw=0, color='g', s=5.0, alpha=0.8) - -plt.figure(2) -plt.clf() -hp.mollview(proj1, fig=2, coord='CG', max=60, cmap=plt.cm.coolwarm) -hp.projscatter(b[idx], l[idx], lw=0, color='g', s=5.0, alpha=0.8) diff --git a/python_sample/build_nbody_skymap.py b/python_sample/build_nbody_skymap.py new file mode 120000 index 0000000..5e7e019 --- /dev/null +++ b/python_sample/build_nbody_skymap.py @@ -0,0 +1 @@ +/nethome/lavaux/wuala/WualaDrive/g_lavaux/PythonCode/BORG/build_nbody_skymap.py \ No newline at end of file diff --git a/python_sample/icgen/borgicgen.py b/python_sample/icgen/borgicgen.py index 8a627ea..111c125 100644 --- a/python_sample/icgen/borgicgen.py +++ b/python_sample/icgen/borgicgen.py @@ -93,17 +93,67 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, density = np.fft.irfftn(lpt.dhat*D1_0)*(supersample*N/L)**3 - return posx,vel,density,N*supersample,L,a_ic + return posx,vel,density,N*supersample,L,a_ic,cosmo -def write_icfiles(*generated_ic, **cosmo): - """Write the initial conditions from the tuple returned by run_generation""" - posx,vel,density,N,L,a_ic = generated_ic +def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'): + N = density.shape[0] + ik = np.fft.fftfreq(N, d=L/N)*2*np.pi - ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) - for i,c in enumerate(["x","y","z"]): - ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo) -# This used to be necessary. However this has been fixed in writeGrafic now -# ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)).transpose(), L, a_ic, **cosmo) + k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2) + p = ct.CosmologyPower(**cosmo) + p.setFunction(func) + p.normalize(cosmo['SIGMA8']) + + Pk = p.compute(k)*cosmo['h']**3*L**3 + Pk[0,0,0]=1 + + density_hat = np.fft.rfftn(density)/N**3/np.sqrt(Pk) + + Ns = N*supergenerate + + density_hat_super = np.zeros((Ns,Ns,Ns/2+1), dtype=np.complex128) + density_hat_super[:] = np.nan + # Copy density hat in place + + hN = N/2 + density_hat_super[:hN, :hN, :hN+1] = density_hat[:hN, :hN, :] + density_hat_super[:hN, (Ns-hN):Ns, :hN+1] = density_hat[:hN, hN:, :] + density_hat_super[(Ns-hN):Ns, (Ns-hN):Ns, :hN+1] = density_hat[hN:, hN:, :] + density_hat_super[(Ns-hN):Ns, :hN, :hN+1] = density_hat[hN:, :hN, :] + + # The moved nyquist place is untouched (so loss of "noise") to keep the structure + # now we just add some noise term + cond=np.isnan(density_hat_super) + x = np.random.randn(np.count_nonzero(cond),2)/np.sqrt(2.0) + density_hat_super[cond] = x[:,0] + 1j * x[:,1] + + # Now we have to fix the Nyquist plane + hNs = Ns/2 + nyquist = density_hat_super[:, :, :hNs] + Nplane = nyquist.size + nyquist.flat[:Nplane/2] = nyquist.flat[Nplane:Nplane/2:-1].conj() + + return np.fft.irfftn(density_hat_super) + +def write_icfiles(*generated_ic, **kwargs): + """Write the initial conditions from the tuple returned by run_generation""" + + supergenerate=1 + if 'supergenerate' in kwargs: + supergenerate=kwargs['supergenerate'] + posx,vel,density,N,L,a_ic,cosmo = generated_ic + + ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) + for i,c in enumerate(["x","y","z"]): + ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo) + + ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo) + + ct.writeWhitePhase("white.dat", whitify(density, L, cosmo, supergenerate=supergenerate)) + + with file("white_params", mode="w") as f: + f.write("4\n%lg, %lg, %lg\n" % (cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h'])) + f.write("%lg\n%lg\n-%lg\n0,0\n" % (cosmo['omega_B_0'],cosmo['ns'],cosmo['SIGMA8'])) + f.write("-%lg\n1\n0\n\n\n\n\n" % L) + f.write("2\n\n0\nwhite.dat\n0\npadding_white.dat\n") - ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo) - ct.writeWhitePhase("white.dat", density) diff --git a/python_sample/icgen/gen_ic_from_borg.py b/python_sample/icgen/gen_ic_from_borg.py index 2fb1d9b..aa6509d 100644 --- a/python_sample/icgen/gen_ic_from_borg.py +++ b/python_sample/icgen/gen_ic_from_borg.py @@ -14,4 +14,4 @@ astart=1/(1.+zstart) halfPixelShift=True if __name__=="__main__": - bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=2, shiftPixel=halfPixelShift, do_lpt2=False), **cosmo) + bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=2, shiftPixel=halfPixelShift, do_lpt2=False)) diff --git a/python_sample/icgen/test_ic_from_borg.py b/python_sample/icgen/test_ic_from_borg.py index e04c26c..dc0265d 100644 --- a/python_sample/icgen/test_ic_from_borg.py +++ b/python_sample/icgen/test_ic_from_borg.py @@ -10,7 +10,7 @@ cosmo['omega_k_0'] = 0 cosmo['omega_B_0']=0.049 cosmo['SIGMA8']=0.8344 cosmo['ns']=0.9624 -N0=256 +N0=128 doSimulation=True @@ -29,7 +29,7 @@ if doSimulation: dsim_hat = np.fft.rfftn(dsim)*(L/N0)**3 Psim, bsim = bic.bin_power(np.abs(dsim_hat)**2/L**3, L, range=(0,1.), bins=150) -pos,_,density,N,L,_ = bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=2, do_lpt2=True) +pos,_,density,N,L,_,_ = bic.run_generation("initial_density_1380.dat", 0.001, astart, cosmo, supersample=2, do_lpt2=True) dcic = ct.cicParticles(pos, L, N0) dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0) @@ -50,5 +50,5 @@ Pref, bref = bic.compute_ref_power(L, N0, cosmo, range=(0,1.), bins=150) Pcic /= D1_0**2 Pdens /= D1_0**2 -borg_evolved = ct.read_borg_vol("final_density_380.dat") +borg_evolved = ct.read_borg_vol("final_density_1380.dat") diff --git a/python_sample/test_spheric_proj.py b/python_sample/test_spheric_proj.py deleted file mode 100644 index 06a1398..0000000 --- a/python_sample/test_spheric_proj.py +++ /dev/null @@ -1,20 +0,0 @@ -import cosmotool as ct -import numpy as np -import healpy as hp - -d = np.zeros((64,64,64), ct.DTYPE) - -d[32,32,32] = 1 - -ii=np.arange(256)*64/256.-32 -xx = ii[:,None,None].repeat(256,axis=1).repeat(256,axis=2).reshape(256**3) -yy = ii[None,:,None].repeat(256,axis=0).repeat(256,axis=2).reshape(256**3) -zz = ii[None,None,:].repeat(256,axis=0).repeat(256,axis=1).reshape(256**3) - -d_high = ct.interp3d(xx, yy, zz, d, 64, periodic=True) -d_high = d_high.reshape((256,256,256)) - -proj0 = ct.spherical_projection(64, d, 0, 20, integrator_id=0, shifter=np.array([0.5,0.5,0.5])) -proj1 = ct.spherical_projection(64, d, 0, 20, integrator_id=1) - -proj0_high = ct.spherical_projection(256, d_high, 0, 30, integrator_id=0, shifter=np.array([0.5,0.5,0.5])) diff --git a/python_sample/test_spheric_proj.py b/python_sample/test_spheric_proj.py new file mode 120000 index 0000000..52aa205 --- /dev/null +++ b/python_sample/test_spheric_proj.py @@ -0,0 +1 @@ +/nethome/lavaux/wuala/WualaDrive/g_lavaux/PythonCode/BORG/test_spheric_proj.py \ No newline at end of file