diff --git a/python/_project.pyx b/python/_project.pyx index 2de1834..4fea24f 100644 --- a/python/_project.pyx +++ b/python/_project.pyx @@ -1,7 +1,7 @@ from cpython cimport bool from cython cimport view from cython.parallel import prange, parallel -from libc.math cimport sin, cos, abs, floor, sqrt +from libc.math cimport sin, cos, abs, floor, round, sqrt import numpy as np cimport numpy as npx cimport cython @@ -45,13 +45,6 @@ cdef int interp3d_INTERNAL_periodic(DTYPE_t x, DTYPE_t y, ry -= iy rz -= iz - while ix < 0: - ix += Ngrid - while iy < 0: - iy += Ngrid - while iz < 0: - iz += Ngrid - jx = (ix+1)%Ngrid jy = (iy+1)%Ngrid jz = (iz+1)%Ngrid @@ -89,6 +82,69 @@ cdef int interp3d_INTERNAL_periodic(DTYPE_t x, DTYPE_t y, return 0 +@cython.boundscheck(False) +@cython.cdivision(True) +cdef int ngp3d_INTERNAL_periodic(DTYPE_t x, DTYPE_t y, + DTYPE_t z, + DTYPE_t[:,:,:] d, DTYPE_t Lbox, DTYPE_t *retval) nogil: + + cdef int Ngrid = d.shape[0] + cdef DTYPE_t inv_delta = Ngrid/Lbox + cdef int ix, iy, iz + cdef DTYPE_t f[2][2][2] + cdef DTYPE_t rx, ry, rz + cdef int jx, jy, jz + + rx = (inv_delta*x) + ry = (inv_delta*y) + rz = (inv_delta*z) + + ix = int(round(rx)) + iy = int(round(ry)) + iz = int(round(rz)) + + + ix = ix%Ngrid + iy = iy%Ngrid + iz = iz%Ngrid + + retval[0] = d[ix ,iy ,iz ] + + return 0 + + +@cython.boundscheck(False) +@cython.cdivision(True) +cdef int ngp3d_INTERNAL(DTYPE_t x, DTYPE_t y, + DTYPE_t z, + DTYPE_t[:,:,:] d, DTYPE_t Lbox, DTYPE_t *retval) nogil: + + cdef int Ngrid = d.shape[0] + cdef DTYPE_t inv_delta = Ngrid/Lbox + cdef int ix, iy, iz + cdef DTYPE_t f[2][2][2] + cdef DTYPE_t rx, ry, rz + cdef int jx, jy, jz + + rx = (inv_delta*x) + ry = (inv_delta*y) + rz = (inv_delta*z) + + ix = int(round(rx)) + iy = int(round(ry)) + iz = int(round(rz)) + + if (ix < 0 or ix >= Ngrid): + return -1 + if (iy < 0 or iy >= Ngrid): + return -2 + if (iz < 0 or iz >= Ngrid): + return -3 + + retval[0] = d[ix ,iy ,iz ] + + return 0 + @cython.boundscheck(False) @cython.cdivision(True) @@ -152,8 +208,8 @@ cdef int interp3d_INTERNAL(DTYPE_t x, DTYPE_t y, def interp3d(x not None, y not None, z not None, npx.ndarray[DTYPE_t, ndim=3] d not None, DTYPE_t Lbox, - bool periodic=False, bool centered=True): - """ interp3d(x,y,z,d,Lbox,periodic=False) -> interpolated values + bool periodic=False, bool centered=True, bool ngp=False): + """ interp3d(x,y,z,d,Lbox,periodic=False,centered=True,ngp=False) -> interpolated values Compute the tri-linear interpolation of the given field (d) at the given position (x,y,z). It assumes that they are box-centered coordinates. So (x,y,z) == (0,0,0) is equivalent to the pixel at (Nx/2,Ny/2,Nz/2) with Nx,Ny,Nz = d.shape. If periodic is set, it assumes the box is periodic """ @@ -164,10 +220,11 @@ def interp3d(x not None, y not None, cdef DTYPE_t retval cdef long i cdef long Nelt - cdef int myperiodic + cdef int myperiodic, myngp cdef DTYPE_t shifter myperiodic = periodic + myngp = ngp if centered: shifter = Lbox/2 @@ -193,26 +250,46 @@ def interp3d(x not None, y not None, in_slice = d Nelt = ax.size with nogil: - if myperiodic: - for i in xrange(Nelt): - if interp3d_INTERNAL_periodic(shifter+ax[i], shifter+ay[i], shifter+az[i], in_slice, Lbox, &out_slice[i]) < 0: - with gil: - raise ierror + if not myngp: + if myperiodic: + for i in prange(Nelt): + if interp3d_INTERNAL_periodic(shifter+ax[i], shifter+ay[i], shifter+az[i], in_slice, Lbox, &out_slice[i]) < 0: + with gil: + raise ierror + else: + for i in prange(Nelt): + if interp3d_INTERNAL(shifter+ax[i], shifter+ay[i], shifter+az[i], in_slice, Lbox, &out_slice[i]) < 0: + with gil: + raise ierror else: - for i in xrange(Nelt): - if interp3d_INTERNAL(shifter+ax[i], shifter+ay[i], shifter+az[i], in_slice, Lbox, &out_slice[i]) < 0: - with gil: - raise ierror - + if myperiodic: + for i in prange(Nelt): + if ngp3d_INTERNAL_periodic(shifter+ax[i], shifter+ay[i], shifter+az[i], in_slice, Lbox, &out_slice[i]) < 0: + with gil: + raise ierror + else: + for i in prange(Nelt): + if ngp3d_INTERNAL(shifter+ax[i], shifter+ay[i], shifter+az[i], in_slice, Lbox, &out_slice[i]) < 0: + with gil: + raise ierror return out else: - if periodic: - if interp3d_INTERNAL_periodic(shifter+x, shifter+y, shifter+z, d, Lbox, &retval) < 0: - raise ierror + if not myngp: + if periodic: + if interp3d_INTERNAL_periodic(shifter+x, shifter+y, shifter+z, d, Lbox, &retval) < 0: + raise ierror + else: + if interp3d_INTERNAL(shifter+x, shifter+y, shifter+z, d, Lbox, &retval) < 0: + raise ierror else: - if interp3d_INTERNAL(shifter+x, shifter+y, shifter+z, d, Lbox, &retval) < 0: - raise ierror + if periodic: + if ngp3d_INTERNAL_periodic(shifter+x, shifter+y, shifter+z, d, Lbox, &retval) < 0: + raise ierror + else: + if ngp3d_INTERNAL(shifter+x, shifter+y, shifter+z, d, Lbox, &retval) < 0: + raise ierror return retval + @cython.boundscheck(False) @cython.cdivision(True) cdef DTYPE_t interp2d_INTERNAL_periodic(DTYPE_t x, DTYPE_t y, @@ -376,12 +453,17 @@ cdef void INTERNAL_project_cic_no_mass_periodic(npx.ndarray[DTYPE_t, ndim=3] g, cdef double a[3], c[3] cdef int b[3], b1[3] cdef int do_not_put + cdef DTYPE_t[:,:] ax + cdef DTYPE_t[:,:,:] ag + + ax = x + ag = g for i in range(x.shape[0]): do_not_put = 0 for j in range(3): - a[j] = (x[i,j]+half_Box)*delta_Box + a[j] = (ax[i,j]+half_Box)*delta_Box b[j] = int(floor(a[j])) b1[j] = (b[j]+1) % Ngrid @@ -390,18 +472,15 @@ cdef void INTERNAL_project_cic_no_mass_periodic(npx.ndarray[DTYPE_t, ndim=3] g, b[j] %= Ngrid - assert b[j] >= 0 and b[j] < Ngrid - assert b1[j] >= 0 and b1[j] < Ngrid - - g[b[0],b[1],b[2]] += c[0]*c[1]*c[2] - g[b1[0],b[1],b[2]] += a[0]*c[1]*c[2] - g[b[0],b1[1],b[2]] += c[0]*a[1]*c[2] - g[b1[0],b1[1],b[2]] += a[0]*a[1]*c[2] + ag[b[0],b[1],b[2]] += c[0]*c[1]*c[2] + ag[b1[0],b[1],b[2]] += a[0]*c[1]*c[2] + ag[b[0],b1[1],b[2]] += c[0]*a[1]*c[2] + ag[b1[0],b1[1],b[2]] += a[0]*a[1]*c[2] - g[b[0],b[1],b1[2]] += c[0]*c[1]*a[2] - g[b1[0],b[1],b1[2]] += a[0]*c[1]*a[2] - g[b[0],b1[1],b1[2]] += c[0]*a[1]*a[2] - g[b1[0],b1[1],b1[2]] += a[0]*a[1]*a[2] + ag[b[0],b[1],b1[2]] += c[0]*c[1]*a[2] + ag[b1[0],b[1],b1[2]] += a[0]*c[1]*a[2] + ag[b[0],b1[1],b1[2]] += c[0]*a[1]*a[2] + ag[b1[0],b1[1],b1[2]] += a[0]*a[1]*a[2] @cython.boundscheck(False) diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py index 05dbba3..ed3680b 100644 --- a/python/cosmotool/__init__.py +++ b/python/cosmotool/__init__.py @@ -1,7 +1,12 @@ from _cosmotool import * from _project import * -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 time_block, timeit, timeit_quiet +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 time_block, timeit, timeit_quiet + +try: + from .fftw import CubeFT +except ImportError: + print("No FFTW support") diff --git a/python/cosmotool/fftw.py b/python/cosmotool/fftw.py new file mode 100644 index 0000000..a640789 --- /dev/null +++ b/python/cosmotool/fftw.py @@ -0,0 +1,36 @@ +import pyfftw +import multiprocessing +import numpy as np +import numexpr as ne + +class CubeFT(object): + def __init__(self, L, N, max_cpu=-1): + + self.N = N + self.align = pyfftw.simd_alignment + self.L = L + 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._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._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False) + + def rfft(self): + return ne.evaluate('c*a', out=self._dhat, local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3}, casting='unsafe') + + def irfft(self): + return ne.evaluate('c*a', out=self._density, local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3}, casting='unsafe') + + def get_dhat(self): + return self._dhat + def set_dhat(self, in_dhat): + self._dhat[:] = in_dhat + dhat = property(get_dhat, set_dhat, None) + + def get_density(self): + return self._density + def set_density(self, d): + self._density[:] = d + density = property(get_density, set_density, None) + + diff --git a/python_sample/icgen/borgicgen.py b/python_sample/icgen/borgicgen.py index 423d7e8..5e5c029 100644 --- a/python_sample/icgen/borgicgen.py +++ b/python_sample/icgen/borgicgen.py @@ -103,7 +103,7 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, @ct.timeit_quiet -def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'): +def whitify(density, L, cosmo, supergenerate=1, zero_fill=False, func='HU_WIGGLES'): N = density.shape[0] p = ct.CosmologyPower(**cosmo) @@ -141,18 +141,22 @@ def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'): if supergenerate > 1: cond=np.isnan(density_hat_super) - print np.where(np.isnan(density_hat_super))[0].size - Nz = np.count_nonzero(cond) - density_hat_super.real[cond] = np.random.randn(Nz) - density_hat_super.imag[cond] = np.random.randn(Nz) - density_hat_super[cond] /= np.sqrt(2.0) - print np.where(np.isnan(density_hat_super))[0].size + if zero_fill: + density_hat_super[cond] = 0 + else: + + print np.where(np.isnan(density_hat_super))[0].size + Nz = np.count_nonzero(cond) + density_hat_super.real[cond] = np.random.randn(Nz) + density_hat_super.imag[cond] = np.random.randn(Nz) + density_hat_super[cond] /= np.sqrt(2.0) + print np.where(np.isnan(density_hat_super))[0].size # Now we have to fix the Nyquist plane - hNs = Ns/2 - nyquist = density_hat_super[:, :, hNs] - Nplane = nyquist.size - nyquist.flat[:Nplane/2] = np.sqrt(2.0)*nyquist.flat[Nplane:Nplane/2:-1].conj() + hNs = Ns/2 + nyquist = density_hat_super[:, :, hNs] + Nplane = nyquist.size + nyquist.flat[:Nplane/2] = np.sqrt(2.0)*nyquist.flat[Nplane:Nplane/2:-1].conj() print np.where(np.isnan(density_hat_super))[0].size @@ -166,8 +170,8 @@ 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'] + supergenerate=kwargs.get('supergenerate', 1) + zero_fill=kwargs.get('zero_fill', False) posx,vel,density,N,L,a_ic,cosmo = generated_ic ct.simpleWriteGadget("Data/borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) @@ -176,7 +180,7 @@ def write_icfiles(*generated_ic, **kwargs): ct.writeGrafic("Data/ic_deltab", density, L, a_ic, **cosmo) - ct.writeWhitePhase("Data/white.dat", whitify(density, L, cosmo, supergenerate=supergenerate)) + ct.writeWhitePhase("Data/white.dat", whitify(density, L, cosmo, supergenerate=supergenerate,zero_fill=zero_fill)) with file("Data/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'])) diff --git a/python_sample/icgen/gen_ic_from_borg.py b/python_sample/icgen/gen_ic_from_borg.py index b2a6fe5..9249725 100644 --- a/python_sample/icgen/gen_ic_from_borg.py +++ b/python_sample/icgen/gen_ic_from_borg.py @@ -9,10 +9,11 @@ cosmo['omega_B_0']=0.049 cosmo['SIGMA8']=0.8344 cosmo['ns']=0.9624 -supergen=1 +supergen=2 zstart=50 -astart=1/(1.+zstart) +astart=1/69.#1/(1.+zstart) halfPixelShift=False +zero_fill=True if __name__=="__main__": - bic.write_icfiles(*bic.run_generation("initial_density_988.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False), supergenerate=supergen) + bic.write_icfiles(*bic.run_generation("initial_density_1380.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False), supergenerate=supergen, zero_fill=zero_fill)