From 4e79e2eec6872ef967d47cd9afb8c22935635bc6 Mon Sep 17 00:00:00 2001 From: Guilhem Lavaux Date: Thu, 25 Sep 2014 09:55:32 +0200 Subject: [PATCH] Python CIC is now more memory and CPU efficient using numexpr --- python/cosmotool/cic.py | 34 ++++++++++++------- .../build_dipole_ksz_from_galaxies.py | 11 +++++- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/python/cosmotool/cic.py b/python/cosmotool/cic.py index 06553f3..481c2ac 100644 --- a/python/cosmotool/cic.py +++ b/python/cosmotool/cic.py @@ -1,27 +1,35 @@ +import numexpr as ne import numpy as np def cicParticles(particles, L, N): + if type(N) not in [int,long]: + raise TypeError("N must be a numeric type") def shifted(i, t): - return (i[2]+t[2])%N + N*((i[1]+t[1])%N + N*((i[0]+t[0])%N)) + a = np.empty(i[0].size, dtype=np.int64) + return ne.evaluate('(i2+t2)%N + N*((i1+t1)%N + N*((i0+t0)%N) )', local_dict={'i2':i[2], 't2':t[2], 'i1':i[1], 't1':t[1], 'i0':i[0], 't0':t[0], 'N':N}, out=a) i =[] r = [] for d in xrange(3): - q = (particles[d]%L)*N/L - o = np.floor(q).astype(int) + q = ne.evaluate('(p%L)*N/L', local_dict={'p':particles[d], 'L':L, 'N':N }) + o = np.empty(q.size, dtype=np.int64) + ne.evaluate('floor(q)', out=o, casting='unsafe') i.append(o) - r.append(q-o) - - density = np.bincount(shifted(i, (1,1,1)), weights= r[0] * r[1] * r[2], minlength=N*N*N) - density += np.bincount(shifted(i, (1,1,0)), weights= r[0] * r[1] *(1-r[2]), minlength=N*N*N) - density += np.bincount(shifted(i, (1,0,1)), weights= r[0] *(1-r[1])* r[2], minlength=N*N*N) - density += np.bincount(shifted(i, (1,0,0)), weights= r[0] *(1-r[1])*(1-r[2]), minlength=N*N*N) - density += np.bincount(shifted(i, (0,1,1)), weights=(1-r[0])* r[1] * r[2], minlength=N*N*N) - density += np.bincount(shifted(i, (0,1,0)), weights=(1-r[0])* r[1] *(1-r[2]), minlength=N*N*N) - density += np.bincount(shifted(i, (0,0,1)), weights=(1-r[0])*(1-r[1])* r[2], minlength=N*N*N) - density += np.bincount(shifted(i, (0,0,0)), weights=(1-r[0])*(1-r[1])*(1-r[2]), minlength=N*N*N) + r.append(ne.evaluate('q-o')) + + D = {'a':r[0],'b':r[1],'c':r[2]} + N3 = N*N*N + + density = np.bincount(shifted(i, (1,1,1)), weights=ne.evaluate('a*b*c', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (1,1,0)), weights=ne.evaluate('a*b*(1-c)', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (1,0,1)), weights=ne.evaluate('a*(1-b)*c', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (1,0,0)), weights=ne.evaluate('a*(1-b)*(1-c)', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (0,1,1)), weights=ne.evaluate('(1-a)*b*c', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (0,1,0)), weights=ne.evaluate('(1-a)*b*(1-c)', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (0,0,1)), weights=ne.evaluate('(1-a)*(1-b)*c', local_dict=D), minlength=N3) + density += np.bincount(shifted(i, (0,0,0)), weights=ne.evaluate('(1-a)*(1-b)*(1-c)', local_dict=D), minlength=N3) return density.reshape((N,N,N)) diff --git a/python_sample/build_dipole_ksz_from_galaxies.py b/python_sample/build_dipole_ksz_from_galaxies.py index 29358b1..271584e 100644 --- a/python_sample/build_dipole_ksz_from_galaxies.py +++ b/python_sample/build_dipole_ksz_from_galaxies.py @@ -99,6 +99,7 @@ def get_args(): parser.add_argument('--ksz_map', type=str, required=True) parser.add_argument('--base_fig', type=str, default="kszfig.png") parser.add_argument('--build_dipole', type=bool, default=False) + parser.add_argument('--degrade', type=int, default=-1) return parser.parse_args() def main(): @@ -113,11 +114,19 @@ def main(): proj,mask = generate_from_catalog(args.depth_min,args.depth_max,args.Nside) + if args.degrade > 0: + proj *= mask + proj = hp.ud_grade(proj, nside_out=args.degrade) + mask = hp.ud_grade(mask, nside_out=args.degrade) + Nside = args.degrade + else: + Nside = args.Nside + hp.write_map(args.ksz_map + ".fits", proj) hp.write_map(args.ksz_map + "_mask.fits", mask) if args.build_dipole: - x,y,z=hp.pix2vec(args.Nside, np.arange(hp.nside2npix(args.Nside))) + x,y,z=hp.pix2vec(Nside, np.arange(hp.nside2npix(Nside))) hp.write_map(args.ksz_map + "_x.fits", proj*x) hp.write_map(args.ksz_map + "_y.fits", proj*y) hp.write_map(args.ksz_map + "_z.fits", proj*z)