Coverage for UFalcon/shells.py : 99%
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
4import os
5import numpy as np
6import healpy as hp
7from UFalcon import utils
9def read_lpicola(path, h, boxsize):
10 """
11 Reads in a binary data file produced by L-Picola.
12 :param path: path to file
13 :param h: Dimensionless Hubble parameter
14 :param boxsize: size of the box in Gigaparsec
15 :return: 3-tuple containing (x, y, z) particle positions in Megaparsec
16 """
18 n_rows_total = 0
20 with open(path, mode='rb') as fh:
22 # first get the total number of blocks, such that we can pre-allocate memory
23 while True:
24 header = np.fromfile(fh, dtype=np.uint32, count=4)
26 if header.size == 0:
27 break
29 n_rows_current = header[1]
30 n_rows_total += n_rows_current
31 block_size = n_rows_current * 7
32 fh.seek(block_size * 4 + 4, 1) # data + endmarker
34 # pre-allocate
35 data = np.empty((n_rows_total, 3), dtype=np.float32)
36 n_rows_read = 0
38 # read out data
39 fh.seek(0)
40 while True:
41 header = np.fromfile(fh, dtype=np.uint32, count=4)
43 if header.size == 0:
44 break
46 n_rows_current = header[1]
47 block_size = n_rows_current * 7
48 data[n_rows_read: n_rows_read + n_rows_current] = np.fromfile(fh,
49 dtype=np.float32,
50 count=block_size).reshape(-1, 7)[:, :3]
51 n_rows_read += n_rows_current
53 fh.seek(4, 1) # skip end marker
55 # transform to Mpc and subtract origin
56 origin = boxsize * 500.0
57 data /= h
58 data -= origin
60 return data
63def read_pkdgrav(path, boxsize, n_rows_per_block=int(1e6)):
64 """
65 Reads in a binary data file produced by PKDGRAV.
66 :param path: path to file
67 :param boxsize: size of the box in Gigaparsec
68 :param n_rows_per_block: number of rows to read in one block, allows to limit memory consumption for large files
69 :return: 3-tuple containing (x, y, z) particle positions in Megaparsec
70 """
72 # get the total number of rows
73 n_rows = os.stat(path).st_size // 7 // 4
75 # initialize output
76 data = np.empty((n_rows, 3), dtype=np.float32)
78 # read in blocks
79 n_block = int(7 * n_rows_per_block)
80 n_rows_in = 0
82 with open(path, mode='rb') as f:
83 while True:
84 block = np.fromfile(f, dtype=np.float32, count=n_block).reshape(-1, 7)[:, :3]
86 if block.size == 0:
87 break
89 data[n_rows_in: n_rows_in + block.shape[0]] = block
90 n_rows_in += block.shape[0]
92 # transforms to Mpc
93 data *= boxsize * 1000
95 return data
98def read_file(path, boxsize, cosmo, file_format='l-picola'):
99 """
100 Reads in particle positions stored in a binary file produced by either L-PICOLA or PKDGRAV.
101 :param path: path to binary file holding particle positions
102 :param boxsize: size of the box in Gigaparsec
103 :param cosmo: Astropy.Cosmo instance, controls the cosmology used
104 :param file_format: data format, either l-picola or pkdgrav
105 :return: theta- and phi-coordinates of particles inside the shell
106 """
108 if file_format == 'l-picola':
109 xyz = read_lpicola(path, cosmo.h, boxsize)
110 elif file_format == 'pkdgrav':
111 xyz = read_pkdgrav(path, boxsize)
112 else:
113 raise ValueError('Data format {} is not supported, choose either "l-picola" or "pkdgrav"')
115 return xyz
118def xyz_to_spherical(xyz_coord):
119 """
120 Transform from comoving cartesian (x, y, z)- to spherical coordinates (comoving radius, healpix theta, healpix phi).
121 :param xyz_coord: cartesian coordinates, shape: (number of particles, 3)
122 :return: comoving radius, theta, phi
123 """
125 x = xyz_coord[:, 0]
126 y = xyz_coord[:, 1]
127 z = xyz_coord[:, 2]
128 spherical_coord = np.empty_like(xyz_coord)
130 # comoving radius
131 spherical_coord[:, 0] = np.sqrt(x ** 2 + y ** 2 + z ** 2)
132 # theta, phi
133 spherical_coord[:, 1], spherical_coord[:, 2] = hp.vec2ang(xyz_coord)
135 return spherical_coord
138def thetaphi_to_pixelcounts(theta, phi, nside):
139 """
140 Transforms angular particle positions to counts in healpix pixels. The size of the output array equals the index of
141 the last non-empty pixel (i.e. the largest healpix index with at least one count).
142 :param theta: healpix theta-coordinate
143 :param phi: healpix phi-coordinate
144 :param nside: nside of the healpix map
145 :return: counts in each pixel, maximum size: nside - 1
146 """
147 pix_ind = hp.ang2pix(nside, theta, phi, nest=False)
148 counts = np.bincount(pix_ind)
149 return counts
152def construct_shells(dirpath, z_shells, boxsize, cosmo, nside, file_format='l-picola'):
153 """
154 Reads in particle positions stored in all the binary file produced by either L-PICOLA or PKDGRAV and transforms
155 their angular positions to counts in healpix pixels corresponding to shells at different redshifts.
156 :param dirpath: path to the directory holding the binary files with particle positions
157 :param z_shells: array containing the discrete redshifts steps, over which the lightcone is constructed
158 :param boxsize: size of the box in Gigaparsec
159 :param cosmo: Astropy.Cosmo instance, controls the cosmology used
160 :param nside: nside of the healpix map
161 :param file_format: data format, either l-picola or pkdgrav
162 :return: matrix with dimension (len(z_shells) - 1, Npix) containing the number counts for each shell and pixel-index
163 """
165 # find all files to process
166 if file_format == 'l-picola':
167 filelist = list(filter(lambda fn: '_lightcone.' in fn and os.path.splitext(fn)[1] != '.info',
168 os.listdir(dirpath)))
169 elif file_format == 'pkdgrav':
170 filelist = list(filter(lambda fn: '.lcp.' in fn, os.listdir(dirpath)))
171 else:
172 raise ValueError('Data format {} is not supported, choose either "l-picola" or "pkdgrav"')
174 print('Will process {} files'.format(len(filelist)))
176 # initialize shells
177 shells = np.zeros((len(z_shells) - 1, hp.nside2npix(nside)), dtype=np.int32)
179 # compute comoving distances of the shell boundaries
180 com_shells = [utils.comoving_distance(0, z, cosmo) for z in z_shells]
182 print('Processing file ', end='', flush=True)
184 for i, filename in enumerate(filelist):
186 print('{} '.format(i + 1), end='', flush=True)
188 filepath = os.path.join(dirpath, filename)
190 # read out cartesian coordinates
191 coord = read_file(filepath, boxsize, cosmo, file_format=file_format)
193 # transform to spherical coordinates
194 coord[:] = xyz_to_spherical(coord)
196 # sort by comoving radius
197 coord[:] = coord[np.argsort(coord[:, 0])]
199 # sort particles into shells
200 ind_shells = np.searchsorted(coord[:, 0], com_shells, side='left')
202 for i_shell in range(shells.shape[0]):
203 i_low = ind_shells[i_shell]
204 i_up = ind_shells[i_shell + 1]
205 counts_shell = thetaphi_to_pixelcounts(coord[i_low: i_up, 1], coord[i_low: i_up, 2], nside)
206 shells[i_shell, :counts_shell.size] += counts_shell
208 return shells