Coverage for pesummary/gw/conversions/cosmology.py: 78.5%
93 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3import numpy as np
4from pesummary.gw.cosmology import get_cosmology
5from pesummary.utils.utils import logger
6from pesummary.utils.decorators import array_input
8__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
10try:
11 from astropy.cosmology import z_at_value
12 import astropy.units as u
13except ImportError:
14 pass
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)
27def _z_from_quantity_exact(samples, quantity, unit, cosmology):
28 """Return the redshift given samples for a cosmological quantity, assuming
29 a given cosmology
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
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
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
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
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
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 )
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 )
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 )
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 )
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
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)
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)
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)
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)
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)
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)
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)
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)
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)
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)
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)
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)
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)
260class Redshift(object):
261 class Distance(object):
262 exact = z_from_dL_exact
263 approx = z_from_dL_approx
265 class ComovingVolume(object):
266 exact = z_from_comoving_volume_exact
267 approx = z_from_comoving_volume_exact