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

Moved code. Reorganized more OO like

Browse Source
This commit is contained in:
Guilhem Lavaux 2014-06-01 18:07:44 +02:00
parent 8258bd0fe9
commit 688e4e20de
6 changed files with 192 additions and 140 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
python_sample/cosmogrowth.py
View File

@ -1,37 +0,0 @@
import numpy as np
import cosmolopy as cpy
class CosmoGrowth:
def __init__(self, **cosmo):
self.cosmo = cosmo
def D(self, a):
return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)
def compute_E(self, a):
om = self.cosmo['omega_M_0']
ol = self.cosmo['omega_lambda_0']
ok = self.cosmo['omega_k_0']
E = np.sqrt(om/a**3 + ol + ok/a**2)
H2 = -3*om/a**4 - 2*ok/a**3
Eprime = 0.5*H2/E
return E,Eprime
def Ddot(self, a):
E,Eprime = self.compute_E(a)
D = self.D(a)
Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
Ddot_D *= a
return Ddot_D
def compute_velmul(self, a):
E,_ = self.compute_E(a)
velmul = self.Ddot(a)
velmul *= 100 * a * E
return velmul

103
python_sample/gen_ic_from_borg.py
View File

@ -1,103 +0,0 @@
import cosmotool as ct
import numpy as np
import cosmolopy as cpy
from cosmogrowth import *
cosmo={'omega_M_0':0.3175, 'h':0.6711}
cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
cosmo['omega_k_0'] = 0
a_start=0.001
z_start=1/a_start-1
def fourier_analysis(borg_vol):
L = (borg_vol.ranges[1]-borg_vol.ranges[0])
N = borg_vol.density.shape[0]
return np.fft.rfftn(borg_vol.density)*(L/N)**3, L, N
def half_pixel_shift(borg):
dhat,L,N = fourier_analysis(borg)
ik = np.fft.fftfreq(N,d=L/N)*2*np.pi
phi = 0.5*L/N*(ik[:,None,None]+ik[None,:,None]+ik[None,None,:(N/2+1)])
phase = np.cos(phi)+1j*np.sin(phi)
return dhat*phase, L
def new_shape(N, direction, q=3):
return ((1,)*direction) + (N,) + ((1,)*(q-1-direction))
def build_dir(ik, direction, q=3):
if direction != q-1:
return ik.reshape(new_shape(ik.size, direction, q=q))
else:
N = ik.size/2+1
return ik[:N].reshape(new_shape(N, direction, q=q))
def get_k2(ik, q=3):
N = ik.size
k2 = (ik.reshape(new_shape(N, 0, q=q))**2)
for d in xrange(1,q):
k2 = k2 + build_dir(ik, d, q=q)**2
return k2
def get_displacement(dhat, L, direction=0):
N = dhat.shape[0]
ik = np.fft.fftfreq(N,d=1.0/N)*2*np.pi/L
k2 = get_k2(ik)
k2[0,0,0] = 1
return -build_dir(ik, direction)*1j*dhat / k2
def gen_posgrid(N, L):
ix = np.arange(N)*L/N
x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
z = ix[None,None,:].repeat(N, axis=0).repeat(N, axis=1)
return x.flatten(), y.flatten(), z.flatten()
def run_generation(input_borg, a_ic):
borg_vol = ct.read_borg_vol(input_borg)
N = borg_vol.density.shape[0]
cgrowth = CosmoGrowth(**cosmo)
density_hat, L = half_pixel_shift(borg_vol)
posq = gen_posgrid(N, L)
vel= []
posx = []
velmul = cgrowth.compute_velmul(a_start)
D1 = cgrowth.D(a_ic)
D1_0 = D1/cgrowth.D(a_start)
D2 = 3./7 * D1**2
for j in xrange(3):
psi = D1_0*np.fft.irfftn(get_displacement(density_hat, L, direction=j)).flatten()*(N/L)**3
posx.append(((posq[j] + psi)%L).astype(np.float32))
vel.append((psi*velmul).astype(np.float32))
return posx,vel,N,L,a_ic
def write_icfiles(*generated_ic):
posx,vel,N,L,a_ic = generated_ic
ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic)
for i,c in enumerate(['x','y','z']):
ct.writeGrafic("borg.ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
if __name__=="__main__":
write_icfiles(*run_generation("initial_condition_borg.dat", 1.0))

18
python_sample/icgen/borgadaptor.py Normal file
View File

@ -0,0 +1,18 @@
import numpy as np
def fourier_analysis(borg_vol):
L = (borg_vol.ranges[1]-borg_vol.ranges[0])
N = borg_vol.density.shape[0]
return np.fft.rfftn(borg_vol.density)*(L/N)**3, L, N
def half_pixel_shift(borg):
dhat,L,N = fourier_analysis(borg)
ik = np.fft.fftfreq(N,d=L/N)*2*np.pi
phi = 0.5*L/N*(ik[:,None,None]+ik[None,:,None]+ik[None,None,:(N/2+1)])
phase = np.cos(phi)+1j*np.sin(phi)
return dhat*phase, L

61
python_sample/icgen/borgicgen.py Normal file
View File

@ -0,0 +1,61 @@
import cosmotool as ct
import numpy as np
import cosmolopy as cpy
from cosmogrowth import *
import borgadaptor as ba
def gen_posgrid(N, L):
""" Generate an ordered lagrangian grid"""
ix = np.arange(N)*L/N
x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
z = ix[None,None,:].repeat(N, axis=0).repeat(N, axis=1)
return x.flatten(), y.flatten(), z.flatten()
def run_generation(input_borg, a_borg, a_ic, **cosmo):
""" Generate particles and velocities from a BORG snapshot. Returns a tuple of
(positions,velocities,N,BoxSize,scale_factor)."""
borg_vol = ct.read_borg_vol(input_borg)
N = borg_vol.density.shape[0]
cgrowth = CosmoGrowth(**cosmo)
density_hat, L = ba.half_pixel_shift(borg_vol)
lpt = LagrangianPerturbation(density_hat, L, fourier=True)
# Generate grid
posq = gen_posgrid(N, L)
vel= []
posx = []
# Compute LPT scaling coefficient
D1 = cgrowth.D(a_ic)
D1_0 = D1/cgrowth.D(a_borg)
velmul = cgrowth.compute_velmul(a_ic)*D1_0
D2 = 3./7 * D1**2
for j in xrange(3):
# Generate psi_j (displacement along j)
psi = D1_0*lpt.lpt1(j).flatten()*(N/L)**3
# Generate posx
posx.append(((posq[j] + psi)%L).astype(np.float32))
# Generate vel
vel.append((psi*velmul).astype(np.float32))
return posx,vel,N,L,a_ic
def write_icfiles(*generated_ic, **cosmo):
"""Write the initial conditions from the tuple returned by run_generation"""
posx,vel,N,L,a_ic = generated_ic
ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic)
for i,c in enumerate(['x','y','z']):
ct.writeGrafic("borg.ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)

102
python_sample/icgen/cosmogrowth.py Normal file
View File

@ -0,0 +1,102 @@
import weakref
import numpy as np
import cosmolopy as cpy
class CosmoGrowth(object):
def __init__(self, **cosmo):
self.cosmo = cosmo
def D(self, a):
return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)
def compute_E(self, a):
om = self.cosmo['omega_M_0']
ol = self.cosmo['omega_lambda_0']
ok = self.cosmo['omega_k_0']
E = np.sqrt(om/a**3 + ol + ok/a**2)
H2 = -3*om/a**4 - 2*ok/a**3
Eprime = 0.5*H2/E
return E,Eprime
def Ddot(self, a):
E,Eprime = self.compute_E(a)
D = self.D(a)
Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
Ddot_D *= a
return Ddot_D
def compute_velmul(self, a):
E,_ = self.compute_E(a)
velmul = self.Ddot(a)
velmul *= 100 * a * E
return velmul
class LagrangianPerturbation(object):
def __init__(self,density,L, fourier=False):
self.L = L
self.N = density.shape[0]
self.dhat = np.fft.rfftn(density)*(L/self.N)**3 if not fourier else density
self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
self.cache = weakref.WeakValueDictionary()
def _gradient(self, phi, direction):
return np.fft.irfftn(self._kdir(direction)*1j*phi)*(self.N/self.L)**3
def lpt1(self, direction=0):
k2 = self._get_k2()
k2[0,0,0] = 1
return self._gradient(-self.dhat/k2, direction)
def new_shape(self,direction, q=3, half=False):
N0 = (self.N/2+1) if half else self.N
return ((1,)*direction) + (N0,) + ((1,)*(q-1-direction))
def _kdir(self, direction, q=3):
if direction != q-1:
return self.ik.reshape(self.new_shape(direction, q=q))
else:
return self.ik[:self.N/2+1].reshape(self.new_shape(direction, q=q, half=True))
def _get_k2(self, q=3):
if 'k2' in self.cache:
return self.cache['k2']
k2 = self._kdir(0, q=q)**2
for d in xrange(1,q):
k2 = k2 + self._kdir(d, q=q)**2
self.cache['k2'] = k2
return k2
def lpt2(self, direction=0):
k2 = self._get_k2()
k2[0,0,0] = 1
if 'lpt2_potential' not in self.cache:
div_phi2 = np.zeros((N,N,N), dtype=np.float64)
for j in xrange(3):
q = np.fft.irfftn( build_dir(ik, j)**2*self.dhat / k2 )
for i in xrange(j+1, 3):
div_phi2 += q * np.fft.irfftn( build_dir(ik, i)**2*self.dhat / k2 )
div_phi2 -= (np.fft.irfftn( build_dir(ik, j)*build_dir(ik, i)*self.dhat / k2 ))**2
div_phi2 *= (self.N/self.L)**3
phi2_hat = np.fft.rfftn(div_phi2) * ((L/N)**3) / k2
self.cache['lpt2_potential'] = phi2_hat
del div_phi2
else:
phi2_hat = self.cache['lpt2_potential']
return self._gradient(phi2_hat, direction)

11
python_sample/icgen/gen_ic_from_borg.py Normal file
View File

@ -0,0 +1,11 @@
import borgicgen as bic
cosmo={'omega_M_0':0.3175, 'h':0.6711}
cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
cosmo['omega_k_0'] = 0
zstart=50
astart=1/(1.+zstart)
if __name__=="__main__":
bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, **cosmo), **cosmo)