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

Fix PYFFTW usage. Use numexpr

Browse Source
This commit is contained in:
Guilhem Lavaux 2014-07-04 16:53:17 +02:00
parent 815c0b616a
commit de1bda366a
3 changed files with 15 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
python_sample/icgen/cosmogrowth.py
View File

@ -1,3 +1,4 @@
import numexpr as ne
import multiprocessing import multiprocessing
import pyfftw import pyfftw
import weakref import weakref
@ -13,14 +14,14 @@ class CubeFT(object):
self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64') self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32') self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
self.irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False) self._irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
self.rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False) self._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
def rfft(self): def rfft(self):
return self.rfft()*(self.L/self.N)**3 return ne.evaluate('c*a', local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3})
def irfft(self): def irfft(self):
return self.irfft()/self.L**3 return ne.evaluate('c*a', local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3})
def get_dhat(self): def get_dhat(self):
return self._dhat return self._dhat
@ -152,16 +153,18 @@ class LagrangianPerturbation(object):
k2 = self._get_k2() k2 = self._get_k2()
k2[0,0,0] = 1 k2[0,0,0] = 1
potgen0 = lambda i: ne.evaluate('kdir**2*d/k2',local_dict={'kdir':self._kdir(i),'d':self.dhat,'k2':k2} )
potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d/k2',local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'k2':k2} )
if 'lpt2_potential' not in self.cache: if 'lpt2_potential' not in self.cache:
print("Rebuilding potential...") print("Rebuilding potential...")
div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64) div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64)
for j in xrange(3): for j in xrange(3):
q = self._do_irfft( self._kdir(j)**2*self.dhat / k2 ).copy() q = self._do_irfft( potgen0(j) ).copy()
for i in xrange(j+1, 3): for i in xrange(j+1, 3):
div_phi2 += q * self._do_irfft( self._kdir(i)**2*self.dhat / k2 ) div_phi2 += q * self._do_irfft( potgen0(i) )
div_phi2 -= (self._do_irfft(self._kdir(j)*self._kdir(i)*self.dhat / k2 ) )**2 div_phi2 -= self._do_irfft(potgen(i,j))**2
div_phi2 *= 1/self.L**6
phi2_hat = -self._do_rfft(div_phi2) / k2 phi2_hat = -self._do_rfft(div_phi2) / k2
#self.cache['lpt2_potential'] = phi2_hat #self.cache['lpt2_potential'] = phi2_hat
del div_phi2 del div_phi2

6
python_sample/icgen/test_ic_from_borg.py
View File

@ -12,8 +12,8 @@ cosmo['SIGMA8']=0.8344
cosmo['ns']=0.9624 cosmo['ns']=0.9624
N0=256 N0=256
doSimulation=True doSimulation=False
simShift=True simShift=False
snap_id=int(sys.argv[1]) snap_id=int(sys.argv[1])
astart=1/100. astart=1/100.
@ -37,7 +37,7 @@ if doSimulation:
dsim_hat = np.fft.rfftn(dsim)*(L/N0)**3 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) 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_density_2588.dat", 0.001, astart, cosmo, supersample=2, do_lpt2=True) pos,_,density,N,L,_,_ = bic.run_generation("initial_density_988.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=True)
dcic = ct.cicParticles(pos, L, N0) dcic = ct.cicParticles(pos, L, N0)
dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0) dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0)