Coverage for UFalcon/scripts/construct_shells.py : 0%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (C) 2020 ETH Zurich, Institute for Particle Physics and Astrophysics
2# Author: Raphael Sgier and Jörg Herbel
4# construct_shells.py
6# This file contains example-functions to compute shells (maps with particle counts) using UFalcon
8import argparse
9import numpy as np
10import yaml
11import h5py
12from astropy.cosmology import FlatLambdaCDM
13from astropy import constants as const
14import UFalcon
17def get_redshifts(z_init, z_final, delta_z):
18 """
19 Generates array containing discrete redshift-steps
20 :param z_init: start redshift
21 :param z_final: end redshift
22 :param delta_z: redshift interval between steps
23 :return: array with redshift-steps between z_init and z_final with delta_z-sized steps
24 """
25 to_check = np.array([z_init, z_final, delta_z])
26 to_check = to_check[to_check != 0]
27 n_digits_round = int(np.ceil(np.amax(np.fabs(np.log10(to_check)))))
28 z = np.around(np.arange(z_init, z_final + delta_z, delta_z), decimals=n_digits_round)
29 return z
32def main(path_config, dirpath_in, sim_type, boxsize, nside, path_out):
33 """
34 Computes and stores maps containing the particle counts (shells) from N-Body simulation output
35 :param path_config: path to configuration yaml file
36 :param dirpath_in: path to directory with N-Body simulation output
37 :param sim_type: type of used N-Body simulation: 'l-picola' or 'pkdgrav' (up to now)
38 :param boxsize: size of the box in Gigaparsec
39 :param nside: nside of output maps
40 :param path_out: path where shells will be stored
41 :return: Computes and stores maps containing the particle counts
42 """
44 print('Config: {}'.format(path_config))
45 print('Input directory: {}'.format(dirpath_in))
46 print('Output directory: {}'.format(path_out))
48 # load config
49 with open(path_config, mode='r') as f:
50 config = yaml.load(f)
52 # get redshifts
53 z_low = config['z_low']
54 z = get_redshifts(z_low['min'], z_low['max'], z_low['delta_z'])
56 # get cosmo instance
57 cosmo_params = config.get('cosmology')
58 cosmo = FlatLambdaCDM(H0=cosmo_params['H0'], Om0=cosmo_params['Om0'])
60 # compute shells
61 shells = UFalcon.shells.construct_shells(dirpath=dirpath_in,
62 z_shells=z,
63 boxsize=boxsize,
64 cosmo=cosmo,
65 nside=nside,
66 file_format=sim_type)
68 # store ouput
69 with h5py.File(path_out, mode='w') as fh5:
70 fh5.create_dataset(name='z', data=np.stack((z[:-1], z[1:]), axis=-1))
71 fh5.create_dataset(name='shells', data=shells, compression='lzf')
73 print('Wrote {}'.format(path_out))
76if __name__ == '__main__':
78 parser = argparse.ArgumentParser(description='Construct shells containing particle counts from N-Bodys',
79 add_help=True)
80 parser.add_argument('--path_config', type=str, required=True, help='configuration yaml file')
81 parser.add_argument('--dirpath_in', type=str, required=True, help='path to simulation output')
82 parser.add_argument('--sim_type', type=str, required=True, choices=('l-picola', 'pkdgrav'),
83 help='type of simulation')
84 parser.add_argument('--boxsize', type=float, required=True, help='boxsize in Gpc')
85 parser.add_argument('--nside', type=int, required=True, help='nside of output shells')
86 parser.add_argument('--path_out', type=str, required=True, help='path where shells will be stored')
87 args = parser.parse_args()
89 main(args.path_config, args.dirpath_in, args.sim_type, args.boxsize, args.nside, args.path_out)