Coverage for pesummary/gw/conversions/cosmology.py: 78.5%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3import numpy as np 

4from pesummary.gw.cosmology import get_cosmology 

5from pesummary.utils.utils import logger 

6from pesummary.utils.decorators import array_input 

7 

8__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

9 

10try: 

11 from astropy.cosmology import z_at_value 

12 import astropy.units as u 

13except ImportError: 

14 pass 

15 

16 

17def _wrapper_for_z_from_quantity_exact(args): 

18 """Wrapper function for _z_from_quantity_exact for a pool of workers 

19Parameters 

20 ---------- 

21 args: tuple 

22 All args passed to _z_from_quantity_exact 

23 """ 

24 return _z_from_quantity_exact(*args) 

25 

26 

27def _z_from_quantity_exact(samples, quantity, unit, cosmology): 

28 """Return the redshift given samples for a cosmological quantity, assuming 

29 a given cosmology 

30 

31 Parameters 

32 ---------- 

33 samples: float/np.array 

34 samples for a cosmological quantity, e.g. luminosity_distance 

35 quantity: str 

36 name of the quantity you wish to compute. This must be an attribute of 

37 cosmology 

38 unit: astropy.units 

39 unit of the samples 

40 cosmology: astropy.cosmology.LambdaCDM 

41 the cosmology to use for conversions 

42 """ 

43 _z = z_at_value(getattr(cosmology, quantity), samples * unit) 

44 return _z.value 

45 

46 

47@array_input(ignore_kwargs=["cosmology", "multi_process"]) 

48def z_from_quantity_exact( 

49 samples, quantity, unit, cosmology="Planck15", multi_process=1 

50): 

51 """Return the redshift given samples for a cosmological quantity 

52 """ 

53 import multiprocessing 

54 from pesummary.utils.utils import iterator 

55 

56 logger.warning( 

57 "Estimating the exact redshift for every {}. This may take a few " 

58 "minutes".format(quantity) 

59 ) 

60 cosmo = get_cosmology(cosmology) 

61 args = np.array( 

62 [ 

63 samples, [quantity] * len(samples), [unit] * len(samples), 

64 [cosmo] * len(samples) 

65 ], dtype=object 

66 ).T 

67 with multiprocessing.Pool(multi_process) as pool: 

68 z = np.array( 

69 list( 

70 iterator( 

71 pool.imap(_wrapper_for_z_from_quantity_exact, args), 

72 tqdm=True, desc="Calculating redshift", logger=logger, 

73 total=len(samples) 

74 ) 

75 ) 

76 ) 

77 return z 

78 

79 

80def _z_from_quantity_approx(samples, quantity, unit, cosmology, N=100, **kwargs): 

81 """Return the redshift given samples for a cosmological quantity, assuming 

82 a given cosmology. This technique uses interpolation to estimate the 

83 redshift 

84 

85 Parameters 

86 ---------- 

87 samples: float/np.array 

88 samples for a cosmological quantity, e.g. luminosity_distance 

89 quantity: str 

90 name of the quantity you wish to compute. This must be an attribute of 

91 cosmology 

92 unit: astropy.units 

93 unit of the samples 

94 cosmology: astropy.cosmology.LambdaCDM 

95 the cosmology to use for conversions 

96 """ 

97 logger.warning("The redshift is being approximated using interpolation. " 

98 "Bear in mind that this does introduce a small error.") 

99 cosmo = get_cosmology(cosmology) 

100 zmin = _z_from_quantity_exact(np.min(samples), quantity, unit, cosmo) 

101 zmax = _z_from_quantity_exact(np.max(samples), quantity, unit, cosmo) 

102 zgrid = np.logspace(np.log10(zmin), np.log10(zmax), N) 

103 grid = _quantity_from_z(zgrid, quantity, cosmology) 

104 zvals = np.interp(samples, grid, zgrid) 

105 return zvals 

106 

107 

108@array_input(ignore_kwargs=["cosmology", "multi_process"]) 

109def z_from_dL_exact(luminosity_distance, cosmology="Planck15", multi_process=1): 

110 """Return the redshift given samples for the luminosity distance 

111 """ 

112 return z_from_quantity_exact( 

113 luminosity_distance, "luminosity_distance", u.Mpc, cosmology=cosmology, 

114 multi_process=multi_process 

115 ) 

116 

117 

118@array_input(ignore_kwargs=["N", "cosmology"]) 

119def z_from_dL_approx( 

120 luminosity_distance, N=100, cosmology="Planck15", **kwargs 

121): 

122 """Return the approximate redshift given samples for the luminosity 

123 distance. This technique uses interpolation to estimate the redshift 

124 """ 

125 return _z_from_quantity_approx( 

126 luminosity_distance, "luminosity_distance", u.Mpc, cosmology, N=N, 

127 **kwargs 

128 ) 

129 

130 

131@array_input(ignore_kwargs=["cosmology", "multi_process"]) 

132def z_from_comoving_volume_exact( 

133 comoving_volume, cosmology="Planck15", multi_process=1): 

134 """Return the redshift given samples for the comoving volume 

135 """ 

136 return z_from_quantity_exact( 

137 comoving_volume, "comoving_volume", u.Mpc**3, cosmology=cosmology, 

138 multi_process=multi_process 

139 ) 

140 

141 

142@array_input(ignore_kwargs=["N", "cosmology"]) 

143def z_from_comoving_volume_approx( 

144 comoving_volume, N=100, cosmology="Planck15", **kwargs 

145): 

146 """Return the approximate redshift given samples for the comoving volume. 

147 This technique uses interpolation to estimate the redshift 

148 """ 

149 return _z_from_quantity_approx( 

150 comoving_volume, "comoving_volume", u.Mpc**3, cosmology, N=N, **kwargs 

151 ) 

152 

153 

154def _quantity_from_z(redshift, quantity, cosmology="Planck15"): 

155 """Return a cosmological quantify given samples for the redshift 

156 """ 

157 cosmo = get_cosmology(cosmology) 

158 return getattr(cosmo, quantity)(redshift).value 

159 

160 

161@array_input(ignore_kwargs=["cosmology"]) 

162def dL_from_z(redshift, cosmology="Planck15"): 

163 """Return the luminosity distance given samples for the redshift 

164 """ 

165 return _quantity_from_z(redshift, "luminosity_distance", cosmology=cosmology) 

166 

167 

168@array_input(ignore_kwargs=["cosmology"]) 

169def comoving_distance_from_z(redshift, cosmology="Planck15"): 

170 """Return the comoving distance given samples for the redshift 

171 """ 

172 return _quantity_from_z(redshift, "comoving_distance", cosmology=cosmology) 

173 

174 

175@array_input(ignore_kwargs=["cosmology"]) 

176def comoving_volume_from_z(redshift, cosmology="Planck15"): 

177 """Return the comoving volume given samples for the redshift 

178 """ 

179 return _quantity_from_z(redshift, "comoving_volume", cosmology=cosmology) 

180 

181 

182def _source_from_detector(parameter, z): 

183 """Return the source-frame parameter given samples for the detector-frame parameter 

184 and the redshift 

185 """ 

186 return parameter / (1. + z) 

187 

188 

189def _detector_from_source(parameter, z): 

190 """Return the detector-frame parameter given samples for the source-frame parameter 

191 and the redshift 

192 """ 

193 return parameter * (1. + z) 

194 

195 

196@array_input() 

197def m1_source_from_m1_z(mass_1, z): 

198 """Return the source-frame primary mass given samples for the 

199 detector-frame primary mass and the redshift 

200 """ 

201 return _source_from_detector(mass_1, z) 

202 

203 

204@array_input() 

205def m1_from_m1_source_z(mass_1_source, z): 

206 """Return the detector-frame primary mass given samples for the 

207 source-frame primary mass and the redshift 

208 """ 

209 return _detector_from_source(mass_1_source, z) 

210 

211 

212@array_input() 

213def m2_source_from_m2_z(mass_2, z): 

214 """Return the source-frame secondary mass given samples for the 

215 detector-frame secondary mass and the redshift 

216 """ 

217 return _source_from_detector(mass_2, z) 

218 

219 

220@array_input() 

221def m2_from_m2_source_z(mass_2_source, z): 

222 """Return the detector-frame secondary mass given samples for the 

223 source-frame secondary mass and the redshift 

224 """ 

225 return _detector_from_source(mass_2_source, z) 

226 

227 

228@array_input() 

229def m_total_source_from_mtotal_z(total_mass, z): 

230 """Return the source-frame total mass of the binary given samples for 

231 the detector-frame total mass and redshift 

232 """ 

233 return _source_from_detector(total_mass, z) 

234 

235 

236@array_input() 

237def mtotal_from_mtotal_source_z(total_mass_source, z): 

238 """Return the detector-frame total mass of the binary given samples for 

239 the source-frame total mass and redshift 

240 """ 

241 return _detector_from_source(total_mass_source, z) 

242 

243 

244@array_input() 

245def mchirp_source_from_mchirp_z(mchirp, z): 

246 """Return the source-frame chirp mass of the binary given samples for 

247 detector-frame chirp mass and redshift 

248 """ 

249 return _source_from_detector(mchirp, z) 

250 

251 

252@array_input() 

253def mchirp_from_mchirp_source_z(mchirp_source, z): 

254 """Return the detector-frame chirp mass of the binary given samples for 

255 the source-frame chirp mass and redshift 

256 """ 

257 return _detector_from_source(mchirp_source, z) 

258 

259 

260class Redshift(object): 

261 class Distance(object): 

262 exact = z_from_dL_exact 

263 approx = z_from_dL_approx 

264 

265 class ComovingVolume(object): 

266 exact = z_from_comoving_volume_exact 

267 approx = z_from_comoving_volume_exact