Moved code. Reorganized more OO like
This commit is contained in:
Guilhem Lavaux
parent
8258bd0fe9
commit
688e4e20de
@ -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
|
||||
|
||||
@ -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
18
python_sample/icgen/borgadaptor.py
Normal 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
61
python_sample/icgen/borgicgen.py
Normal 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
102
python_sample/icgen/cosmogrowth.py
Normal 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
11
python_sample/icgen/gen_ic_from_borg.py
Normal 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)
|
||||
Loading…
Reference in New Issue
Block a user