Hide keyboard shortcuts

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 

3 

4import os 

5import numpy as np 

6import healpy as hp 

7from UFalcon import utils 

8 

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 """ 

17 

18 n_rows_total = 0 

19 

20 with open(path, mode='rb') as fh: 

21 

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) 

25 

26 if header.size == 0: 

27 break 

28 

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 

33 

34 # pre-allocate 

35 data = np.empty((n_rows_total, 3), dtype=np.float32) 

36 n_rows_read = 0 

37 

38 # read out data 

39 fh.seek(0) 

40 while True: 

41 header = np.fromfile(fh, dtype=np.uint32, count=4) 

42 

43 if header.size == 0: 

44 break 

45 

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 

52 

53 fh.seek(4, 1) # skip end marker 

54 

55 # transform to Mpc and subtract origin 

56 origin = boxsize * 500.0 

57 data /= h 

58 data -= origin 

59 

60 return data 

61 

62 

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 """ 

71 

72 # get the total number of rows 

73 n_rows = os.stat(path).st_size // 7 // 4 

74 

75 # initialize output 

76 data = np.empty((n_rows, 3), dtype=np.float32) 

77 

78 # read in blocks 

79 n_block = int(7 * n_rows_per_block) 

80 n_rows_in = 0 

81 

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] 

85 

86 if block.size == 0: 

87 break 

88 

89 data[n_rows_in: n_rows_in + block.shape[0]] = block 

90 n_rows_in += block.shape[0] 

91 

92 # transforms to Mpc 

93 data *= boxsize * 1000 

94 

95 return data 

96 

97 

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 """ 

107 

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"') 

114 

115 return xyz 

116 

117 

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 """ 

124 

125 x = xyz_coord[:, 0] 

126 y = xyz_coord[:, 1] 

127 z = xyz_coord[:, 2] 

128 spherical_coord = np.empty_like(xyz_coord) 

129 

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) 

134 

135 return spherical_coord 

136 

137 

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 

150 

151 

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 """ 

164 

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"') 

173 

174 print('Will process {} files'.format(len(filelist))) 

175 

176 # initialize shells 

177 shells = np.zeros((len(z_shells) - 1, hp.nside2npix(nside)), dtype=np.int32) 

178 

179 # compute comoving distances of the shell boundaries 

180 com_shells = [utils.comoving_distance(0, z, cosmo) for z in z_shells] 

181 

182 print('Processing file ', end='', flush=True) 

183 

184 for i, filename in enumerate(filelist): 

185 

186 print('{} '.format(i + 1), end='', flush=True) 

187 

188 filepath = os.path.join(dirpath, filename) 

189 

190 # read out cartesian coordinates 

191 coord = read_file(filepath, boxsize, cosmo, file_format=file_format) 

192 

193 # transform to spherical coordinates 

194 coord[:] = xyz_to_spherical(coord) 

195 

196 # sort by comoving radius 

197 coord[:] = coord[np.argsort(coord[:, 0])] 

198 

199 # sort particles into shells 

200 ind_shells = np.searchsorted(coord[:, 0], com_shells, side='left') 

201 

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 

207 

208 return shells