Python CIC is now more memory and CPU efficient using numexpr

This commit is contained in:
Guilhem Lavaux 2014-09-25 09:55:32 +02:00
parent 0d688fbf7d
commit 4e79e2eec6
2 changed files with 31 additions and 14 deletions

View File

@ -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))

View File

@ -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)