guilhem_lavaux/cosmotool
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Python CIC is now more memory and CPU efficient using numexpr

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

34
python/cosmotool/cic.py
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))