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

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_dL_exact(args): 

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

19 

20 Parameters 

21 ---------- 

22 args: tuple 

23 All args passed to _z_from_dL_exact 

24 """ 

25 return _z_from_dL_exact(*args) 

26 

27 

28def _z_from_dL_exact(luminosity_distance, cosmology): 

29 """Return the redshift given samples for the luminosity distance for a 

30 given cosmology 

31 

32 Parameters 

33 ---------- 

34 luminosity_distance: float/np.array 

35 luminosity distance samples 

36 cosmology: astropy.cosmology.LambdaCDM 

37 the cosmology to use for conversions 

38 """ 

39 _z = z_at_value( 

40 cosmology.luminosity_distance, luminosity_distance * u.Mpc 

41 ) 

42 try: 

43 return _z.value 

44 except AttributeError: 

45 # astropy < 5.0 

46 return _z 

47 

48 

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

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

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

52 """ 

53 import multiprocessing 

54 from pesummary.utils.utils import iterator 

55 

56 logger.warning("Estimating the exact redshift for every luminosity " 

57 "distance. This may take a few minutes.") 

58 cosmo = get_cosmology(cosmology) 

59 args = np.array( 

60 [luminosity_distance, [cosmo] * len(luminosity_distance)], 

61 dtype=object 

62 ).T 

63 with multiprocessing.Pool(multi_process) as pool: 

64 z = np.array( 

65 list( 

66 iterator( 

67 pool.imap(_wrapper_for_z_from_dL_exact, args), tqdm=True, 

68 desc="Calculating redshift", logger=logger, 

69 total=len(luminosity_distance) 

70 ) 

71 ) 

72 ) 

73 return z 

74 

75 

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

77def z_from_dL_approx( 

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

79): 

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

81 distance. This technique uses interpolation to estimate the redshift 

82 """ 

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

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

85 cosmo = get_cosmology(cosmology) 

86 d_min = np.min(luminosity_distance) 

87 d_max = np.max(luminosity_distance) 

88 zmin = _z_from_dL_exact(d_min, cosmo) 

89 zmax = _z_from_dL_exact(d_max, cosmo) 

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

91 Dgrid = cosmo.luminosity_distance(zgrid).value 

92 zvals = np.interp(luminosity_distance, Dgrid, zgrid) 

93 return zvals 

94 

95 

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

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

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

99 """ 

100 cosmo = get_cosmology(cosmology) 

101 return cosmo.luminosity_distance(redshift).value 

102 

103 

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

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

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

107 """ 

108 cosmo = get_cosmology(cosmology) 

109 return cosmo.comoving_distance(redshift).value 

110 

111 

112def _source_from_detector(parameter, z): 

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

114 and the redshift 

115 """ 

116 return parameter / (1. + z) 

117 

118 

119def _detector_from_source(parameter, z): 

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

121 and the redshift 

122 """ 

123 return parameter * (1. + z) 

124 

125 

126@array_input() 

127def m1_source_from_m1_z(mass_1, z): 

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

129 detector-frame primary mass and the redshift 

130 """ 

131 return _source_from_detector(mass_1, z) 

132 

133 

134@array_input() 

135def m1_from_m1_source_z(mass_1_source, z): 

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

137 source-frame primary mass and the redshift 

138 """ 

139 return _detector_from_source(mass_1_source, z) 

140 

141 

142@array_input() 

143def m2_source_from_m2_z(mass_2, z): 

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

145 detector-frame secondary mass and the redshift 

146 """ 

147 return _source_from_detector(mass_2, z) 

148 

149 

150@array_input() 

151def m2_from_m2_source_z(mass_2_source, z): 

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

153 source-frame secondary mass and the redshift 

154 """ 

155 return _detector_from_source(mass_2_source, z) 

156 

157 

158@array_input() 

159def m_total_source_from_mtotal_z(total_mass, z): 

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

161 the detector-frame total mass and redshift 

162 """ 

163 return _source_from_detector(total_mass, z) 

164 

165 

166@array_input() 

167def mtotal_from_mtotal_source_z(total_mass_source, z): 

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

169 the source-frame total mass and redshift 

170 """ 

171 return _detector_from_source(total_mass_source, z) 

172 

173 

174@array_input() 

175def mchirp_source_from_mchirp_z(mchirp, z): 

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

177 detector-frame chirp mass and redshift 

178 """ 

179 return _source_from_detector(mchirp, z) 

180 

181 

182@array_input() 

183def mchirp_from_mchirp_source_z(mchirp_source, z): 

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

185 the source-frame chirp mass and redshift 

186 """ 

187 return _detector_from_source(mchirp_source, z) 

188 

189 

190class Redshift(object): 

191 exact = z_from_dL_exact 

192 approx = z_from_dL_approx