From 688e4e20de0141bf276d33beb7cbcb32110effa2 Mon Sep 17 00:00:00 2001 From: Guilhem Lavaux Date: Sun, 1 Jun 2014 18:07:44 +0200 Subject: [PATCH] Moved code. Reorganized more OO like --- python_sample/cosmogrowth.py | 37 --------- python_sample/gen_ic_from_borg.py | 103 ------------------------ python_sample/icgen/borgadaptor.py | 18 +++++ python_sample/icgen/borgicgen.py | 61 ++++++++++++++ python_sample/icgen/cosmogrowth.py | 102 +++++++++++++++++++++++ python_sample/icgen/gen_ic_from_borg.py | 11 +++ 6 files changed, 192 insertions(+), 140 deletions(-) delete mode 100644 python_sample/cosmogrowth.py delete mode 100644 python_sample/gen_ic_from_borg.py create mode 100644 python_sample/icgen/borgadaptor.py create mode 100644 python_sample/icgen/borgicgen.py create mode 100644 python_sample/icgen/cosmogrowth.py create mode 100644 python_sample/icgen/gen_ic_from_borg.py diff --git a/python_sample/cosmogrowth.py b/python_sample/cosmogrowth.py deleted file mode 100644 index 8e378e5..0000000 --- a/python_sample/cosmogrowth.py +++ /dev/null @@ -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 - diff --git a/python_sample/gen_ic_from_borg.py b/python_sample/gen_ic_from_borg.py deleted file mode 100644 index 56b4330..0000000 --- a/python_sample/gen_ic_from_borg.py +++ /dev/null @@ -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)) diff --git a/python_sample/icgen/borgadaptor.py b/python_sample/icgen/borgadaptor.py new file mode 100644 index 0000000..93fcd9b --- /dev/null +++ b/python_sample/icgen/borgadaptor.py @@ -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 + diff --git a/python_sample/icgen/borgicgen.py b/python_sample/icgen/borgicgen.py new file mode 100644 index 0000000..193f3c5 --- /dev/null +++ b/python_sample/icgen/borgicgen.py @@ -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) + diff --git a/python_sample/icgen/cosmogrowth.py b/python_sample/icgen/cosmogrowth.py new file mode 100644 index 0000000..65279c9 --- /dev/null +++ b/python_sample/icgen/cosmogrowth.py @@ -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) diff --git a/python_sample/icgen/gen_ic_from_borg.py b/python_sample/icgen/gen_ic_from_borg.py new file mode 100644 index 0000000..0c08630 --- /dev/null +++ b/python_sample/icgen/gen_ic_from_borg.py @@ -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)