Coverage for pesummary/gw/conversions/evolve.py: 76.5%
85 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 08:42 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3import numpy as np
4import multiprocessing
6from pesummary.gw.conversions import (
7 tilt_angles_and_phi_12_from_spin_vectors_and_L
8)
9from pesummary.gw.waveform import _get_start_freq_from_approximant
10from pesummary.utils.utils import iterator, logger
11from pesummary.utils.decorators import array_input
13__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
15try:
16 from lal import MTSUN_SI, MSUN_SI
17 import lalsimulation
18 from lalsimulation import SimInspiralSpinTaylorPNEvolveOrbit
19except ImportError:
20 pass
23def evolve_spins(*args, evolve_limit="ISCO", **kwargs):
24 """Evolve spins to a given limit.
26 Parameters
27 ----------
28 *args: tuple
29 all arguments passed to either evolve_angles_forwards or
30 evolve_angles_backwards
31 evolve_limit: str/float, optional
32 limit to evolve frequencies. If evolve_limit=='infinite_separation' or
33 evolve_limit==0, evolve spins to infinite separation. if
34 evolve_limit=='ISCO', evolve spins to ISCO frequency. If any other
35 float, evolve spins to that frequency.
36 **kwargs: dict, optional
37 all kwargs passed to either evolve_angles_forwards or evolve_angles_backwards
38 """
39 _infinite_string = "infinite_separation"
40 cond1 = isinstance(evolve_limit, str) and evolve_limit.lower() == _infinite_string
41 if cond1 or evolve_limit == 0:
42 return evolve_angles_backwards(*args, **kwargs)
43 else:
44 return evolve_angles_forwards(*args, **kwargs)
47@array_input(
48 ignore_kwargs=[
49 "final_velocity", "tolerance", "dt", "multi_process",
50 "evolution_approximant"
51 ], force_return_array=True
52)
53def evolve_angles_forwards(
54 mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, phi_12, f_low, f_ref,
55 approximant, final_velocity="ISCO", tolerance=1e-3,
56 dt=0.1, multi_process=1, evolution_approximant="SpinTaylorT5"
57):
58 """Evolve the BBH spin angles forwards to a specified value using
59 lalsimulation.SimInspiralSpinTaylorPNEvolveOrbit. By default this is
60 the Schwarzchild ISCO velocity.
62 Parameters
63 ----------
64 mass_1: float/np.ndarray
65 float/array of primary mass samples of the binary
66 mass_2: float/np.ndarray
67 float/array of secondary mass samples of the binary
68 a_1: float/np.ndarray
69 float/array of primary spin magnitudes
70 a_2: float/np.ndarray
71 float/array of secondary spin magnitudes
72 tilt_1: float/np.ndarray
73 float/array of primary spin tilt angle from the orbital angular momentum
74 tilt_2: float/np.ndarray
75 float/array of secondary spin tilt angle from the orbital angular momentum
76 phi_12: float/np.ndarray
77 float/array of samples for the angle between the in-plane spin
78 components
79 f_low: float
80 low frequency cutoff used in the analysis
81 f_ref: float
82 reference frequency where spins are defined
83 approximant: str
84 Approximant used to generate the posterior samples
85 final_velocity: str, float
86 final orbital velocity for the evolution. This can either be the
87 Schwarzschild ISCO velocity 6**-0.5 ~= 0.408 ('ISCO') or a
88 fraction of the speed of light
89 tolerance: float
90 Only evolve spins if at least one spin's magnitude is greater than
91 tolerance
92 dt: float
93 steps in time for the integration, in terms of the mass of the binary
94 multi_process: int, optional
95 number of cores to run on when evolving the spins. Default: 1
96 evolution_approximant: str
97 name of the approximant you wish to use to evolve the spins. Default
98 is SpinTaylorT5. Other choices are SpinTaylorT1 or SpinTaylorT4
99 """
100 if isinstance(final_velocity, str) and final_velocity.lower() == "isco":
101 final_velocity = 6. ** -0.5
102 else:
103 final_velocity = float(final_velocity)
105 f_start = _get_start_freq_from_approximant(approximant, f_low, f_ref)
106 with multiprocessing.Pool(multi_process) as pool:
107 args = np.array([
108 mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, phi_12,
109 [f_start] * len(mass_1), [final_velocity] * len(mass_1),
110 [tolerance] * len(mass_1), [dt] * len(mass_1),
111 [evolution_approximant] * len(mass_1)
112 ], dtype="object").T
113 data = np.array(
114 list(
115 iterator(
116 pool.imap(_wrapper_for_evolve_angles_forwards, args),
117 tqdm=True, logger=logger, total=len(mass_1),
118 desc="Evolving spins forward for remnant fits evaluation"
119 )
120 )
121 )
122 tilt_1_evol, tilt_2_evol, phi_12_evol = data.T
123 return tilt_1_evol, tilt_2_evol, phi_12_evol
126def _wrapper_for_evolve_angles_forwards(args):
127 """Wrapper function for _evolve_angles_forwards for a pool of workers
129 Parameters
130 ----------
131 args: tuple
132 All args passed to _evolve_angles_forwards
133 """
134 return _evolve_angles_forwards(*args)
137def _evolve_angles_forwards(
138 mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, phi_12, f_start, final_velocity,
139 tolerance, dt, evolution_approximant
140):
141 """Wrapper function for the SimInspiralSpinTaylorPNEvolveOrbit function
143 Parameters
144 ----------
145 mass_1: float
146 primary mass of the binary
147 mass_2: float
148 secondary mass of the binary
149 a_1: float
150 primary spin magnitude
151 a_2: float
152 secondary spin magnitude
153 tilt_1: float
154 primary spin tilt angle from the orbital angular momentum
155 tilt_2: float
156 secondary spin tilt angle from the orbital angular momentum
157 phi_12: float
158 the angle between the in-plane spin components
159 f_start: float
160 frequency to start the evolution from
161 final_velocity: float
162 Final velocity to evolve the spins up to
163 tolerance: float
164 Only evolve spins if at least one spins magnitude is greater than
165 tolerance
166 dt: float
167 steps in time for the integration, in terms of the mass of the binary
168 evolution_approximant: str
169 name of the approximant you wish to use to evolve the spins.
170 """
171 from packaging import version
172 if np.logical_or(a_1 > tolerance, a_2 > tolerance):
173 # Total mass in seconds
174 total_mass = (mass_1 + mass_2) * MTSUN_SI
175 f_final = final_velocity ** 3 / (total_mass * np.pi)
176 _approx = getattr(lalsimulation, evolution_approximant)
177 if version.parse(lalsimulation.__version__) >= version.parse("2.5.2"):
178 spinO = 6
179 else:
180 spinO = 7
181 data = SimInspiralSpinTaylorPNEvolveOrbit(
182 deltaT=dt * total_mass, m1=mass_1 * MSUN_SI,
183 m2=mass_2 * MSUN_SI, fStart=f_start, fEnd=f_final,
184 s1x=a_1 * np.sin(tilt_1), s1y=0.,
185 s1z=a_1 * np.cos(tilt_1),
186 s2x=a_2 * np.sin(tilt_2) * np.cos(phi_12),
187 s2y=a_2 * np.sin(tilt_2) * np.sin(phi_12),
188 s2z=a_2 * np.cos(tilt_2), lnhatx=0., lnhaty=0., lnhatz=1.,
189 e1x=1., e1y=0., e1z=0., lambda1=0., lambda2=0., quadparam1=1.,
190 quadparam2=1., spinO=spinO, tideO=0, phaseO=7, lscorr=0,
191 approx=_approx
192 )
193 # Set index to take from array output by SimInspiralSpinTaylorPNEvolveOrbit:
194 # -1 for evolving forward in time and 0 for evolving backward in time
195 if f_start <= f_final:
196 idx_use = -1
197 else:
198 idx_use = 0
199 a_1_evolve = np.array(
200 [
201 data[2].data.data[idx_use], data[3].data.data[idx_use],
202 data[4].data.data[idx_use]
203 ]
204 )
205 a_2_evolve = np.array(
206 [
207 data[5].data.data[idx_use], data[6].data.data[idx_use],
208 data[7].data.data[idx_use]
209 ]
210 )
211 Ln_evolve = np.array(
212 [
213 data[8].data.data[idx_use], data[9].data.data[idx_use],
214 data[10].data.data[idx_use]
215 ]
216 )
217 tilt_1_evol, tilt_2_evol, phi_12_evol = \
218 tilt_angles_and_phi_12_from_spin_vectors_and_L(
219 a_1_evolve, a_2_evolve, Ln_evolve
220 )
221 else:
222 tilt_1_evol, tilt_2_evol, phi_12_evol = tilt_1, tilt_2, phi_12
223 return tilt_1_evol, tilt_2_evol, phi_12_evol
226def _wrapper_for_evolve_angles_backwards(args):
227 """Wrapper function for evolving tilts backwards for a pool of workers
229 Parameters
230 ----------
231 args: tuple
232 Zeroth arg is the function you wish to use when evolving the tilts.
233 1st to 8th args are arguments passed to function. All other arguments
234 are treated as kwargs passed to function
235 """
236 _function = args[0]
237 _args = args[1:9]
238 _kwargs = args[9:]
239 return _function(*_args, **_kwargs[0])
242@array_input(
243 ignore_kwargs=[
244 "method", "multi_process", "return_fits_used", "version"
245 ], force_return_array=True
246)
247def evolve_angles_backwards(
248 mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, phi_12, f_low, f_ref,
249 approximant, method="precession_averaged", multi_process=1,
250 return_fits_used=False, version="v2", evolution_approximant="SpinTaylorT5",
251 **kwargs
252):
253 """Evolve BBH tilt angles backwards to infinite separation
255 Parameters
256 ----------
257 mass_1: float/np.ndarray
258 float/array of primary mass samples of the binary
259 mass_2: float/np.ndarray
260 float/array of secondary mass samples of the binary
261 a_1: float/np.ndarray
262 float/array of primary spin magnitudes
263 a_2: float/np.ndarray
264 float/array of secondary spin magnitudes
265 tilt_1: float/np.ndarray
266 float/array of primary spin tilt angle from the orbital angular momentum
267 tilt_2: float/np.ndarray
268 float/array of secondary spin tilt angle from the orbital angular momentum
269 phi_12: float/np.ndarray
270 float/array of samples for the angle between the in-plane spin
271 components
272 f_ref: float
273 reference frequency where spins are defined
274 approximant: str
275 Approximant used to generate the posterior samples
276 method: str
277 Method to use when evolving tilts to infinity. Possible options are
278 'precession_averaged' and 'hybrid_orbit_averaged'. 'precession_averaged'
279 computes tilt angles at infinite separation assuming that precession
280 averaged spin evolution from Gerosa et al. is valid starting from f_ref.
281 'hybrid_orbit_averaged' combines orbit-averaged evolution and
282 'precession_averaged' evolution as in Johnson-McDaniel et al. This is more
283 accurate but slower than the 'precession_averaged' method.
284 multi_process: int, optional
285 number of cores to run on when evolving the spins. Default: 1
286 return_fits_used: Bool, optional
287 return a dictionary of fits used. Default False
288 version: str, optional
289 version of the
290 tilts_at_infinity.hybrid_spin_evolution.calc_tilts_at_infty_hybrid_evolve
291 function to use within the lalsimulation library. Default 'v2'. If
292 an old version of lalsimulation is installed where 'v2' is not
293 available, fall back to 'v1'.
294 evolution_approximant: str, optional
295 the approximant to use when evolving the spins. For allowed
296 approximants see
297 tilts_at_infinity.hybrid_spin_evolution.calc_tilts_at_infty_hybrid_evolve
298 Default 'SpinTaylorT5'
299 **kwargs: dict, optional
300 all kwargs passed to the
301 tilts_at_infinity.hybrid_spin_evolution.calc_tilts_at_infty_hybrid_evolve
302 function in the lalsimulation library
303 """
304 from lalsimulation.tilts_at_infinity import hybrid_spin_evolution
305 import warnings
306 _mds = ["precession_averaged", "hybrid_orbit_averaged"]
307 if method.lower() not in _mds:
308 raise ValueError(
309 "Invalid method. Please choose either {}".format(", ".join(_mds))
310 )
311 # check to see if the provided version is available in lalsimulation. If not
312 # fall back to 'v1'
313 try:
314 with warnings.catch_warnings():
315 warnings.simplefilter("ignore")
316 _ = hybrid_spin_evolution.calc_tilts_at_infty_hybrid_evolve(
317 2., 1., 1., 1., 1., 1., 1., 1., version=version
318 )
319 except ValueError as e:
320 if "Only version" not in str(e):
321 raise
322 logger.warning(
323 "Unknown version '{}'. Falling back to 'v1'".format(version)
324 )
325 version = "v1"
327 kwargs.update(
328 {
329 "prec_only": method.lower() == "precession_averaged",
330 "version": version, "approx": evolution_approximant
331 }
332 )
333 f_start = _get_start_freq_from_approximant(approximant, f_low, f_ref)
334 with multiprocessing.Pool(multi_process) as pool:
335 args = np.array(
336 [
337 [hybrid_spin_evolution.calc_tilts_at_infty_hybrid_evolve] * len(mass_1),
338 mass_1 * MSUN_SI, mass_2 * MSUN_SI, a_1, a_2, tilt_1, tilt_2, phi_12,
339 [f_start] * len(mass_1), [kwargs] * len(mass_1)
340 ], dtype=object
341 ).T
342 data = np.array(
343 list(
344 iterator(
345 pool.imap(_wrapper_for_evolve_angles_backwards, args),
346 tqdm=True, desc="Evolving spins backwards to infinite separation",
347 logger=logger, total=len(mass_1)
348 )
349 )
350 )
351 tilt_1_inf = np.array([l["tilt1_inf"] for l in data])
352 tilt_2_inf = np.array([l["tilt2_inf"] for l in data])
353 if return_fits_used:
354 fits_used = [
355 method.lower(), (
356 "lalsimulation.tilts_at_infinity.hybrid_spin_evolution."
357 "calc_tilts_at_infty_hybrid_evolve=={}".format(version)
358 )
359 ]
360 if method.lower() == "hybrid_orbit_averaged":
361 fits_used.append("approx={}".format(evolution_approximant))
362 return [tilt_1_inf, tilt_2_inf, phi_12], fits_used
363 return tilt_1_inf, tilt_2_inf, phi_12