Skip to content

Commit

Permalink
Codeclimate and cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@titodalcanton
titodalcanton committed Sep 19, 2023
1 parent c5d509f commit 0d39fd5
Showing 1 changed file with 64 additions and 51 deletions.
115 changes: 64 additions & 51 deletions pycbc/pnutils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
from scipy.optimize import bisect, brentq, minimize
from pycbc import conversions, libutils

lalsimulation = libutils.import_optional('lalsimulation')
lalsim = libutils.import_optional('lalsimulation')

def nearest_larger_binary_number(input_len):
""" Return the nearest binary number larger than input_len.
Expand Down Expand Up @@ -318,8 +318,7 @@ def f_LRD(m1, m2):
return 1.2 * f_FRD(m1, m2)

def _get_freq(freqfunc, m1, m2, s1z, s2z):
"""
Wrapper of the LALSimulation function returning the frequency
"""Wrapper of the LALSimulation function returning the frequency
for a given frequency function and template parameters.
Parameters
Expand All @@ -340,11 +339,17 @@ def _get_freq(freqfunc, m1, m2, s1z, s2z):
f : float
Frequency in Hz
"""
# Convert to SI units for lalsimulation
m1kg = float(m1) * lal.MSUN_SI
m2kg = float(m2) * lal.MSUN_SI
return lalsimulation.SimInspiralGetFrequency(
m1kg, m2kg, 0, 0, float(s1z), 0, 0, float(s2z), int(freqfunc))
return lalsim.SimInspiralGetFrequency(
solar_mass_to_kg(m1),
solar_mass_to_kg(m2),
0,
0,
float(s1z),
0,
0,
float(s2z),
int(freqfunc)
)

# vectorize to enable calls with numpy arrays
_vec_get_freq = numpy.vectorize(_get_freq)
Expand Down Expand Up @@ -376,8 +381,7 @@ def get_freq(freqfunc, m1, m2, s1z, s2z):
return _vec_get_freq(lalsim_ffunc, m1, m2, s1z, s2z)

def _get_final_freq(approx, m1, m2, s1z, s2z):
"""
Wrapper of the LALSimulation function returning the final (highest)
"""Wrapper of the LALSimulation function returning the final (highest)
frequency for a given approximant an template parameters
Parameters
Expand All @@ -398,11 +402,17 @@ def _get_final_freq(approx, m1, m2, s1z, s2z):
f : float
Frequency in Hz
"""
# Convert to SI units for lalsimulation
m1kg = float(m1) * lal.MSUN_SI
m2kg = float(m2) * lal.MSUN_SI
return lalsimulation.SimInspiralGetFinalFreq(
m1kg, m2kg, 0, 0, float(s1z), 0, 0, float(s2z), int(approx))
return lalsim.SimInspiralGetFinalFreq(
solar_mass_to_kg(m1),
solar_mass_to_kg(m2),
0,
0,
float(s1z),
0,
0,
float(s2z),
int(approx)
)

# vectorize to enable calls with numpy arrays
_vec_get_final_freq = numpy.vectorize(_get_final_freq)
Expand Down Expand Up @@ -432,7 +442,7 @@ def get_final_freq(approx, m1, m2, s1z, s2z):
f : float or numpy.array
Frequency in Hz
"""
lalsim_approx = lalsimulation.GetApproximantFromString(approx)
lalsim_approx = lalsim.GetApproximantFromString(approx)
return _vec_get_final_freq(lalsim_approx, m1, m2, s1z, s2z)

# Dictionary of functions with uniform API taking a
Expand Down Expand Up @@ -511,27 +521,27 @@ def _get_imr_duration(m1, m2, s1z, s2z, f_low, approximant="SEOBNRv4"):
m1, m2, s1z, s2z, f_low = float(m1), float(m2), float(s1z), float(s2z),\
float(f_low)
if approximant == "SEOBNRv2":
chi = lalsimulation.SimIMRPhenomBComputeChi(m1, m2, s1z, s2z)
time_length = lalsimulation.SimIMRSEOBNRv2ChirpTimeSingleSpin(
chi = lalsim.SimIMRPhenomBComputeChi(m1, m2, s1z, s2z)
time_length = lalsim.SimIMRSEOBNRv2ChirpTimeSingleSpin(
m1 * lal.MSUN_SI, m2 * lal.MSUN_SI, chi, f_low)
elif approximant == 'IMRPhenomXAS':
time_length = lalsimulation.SimIMRPhenomXASDuration(
time_length = lalsim.SimIMRPhenomXASDuration(
m1 * lal.MSUN_SI, m2 * lal.MSUN_SI, s1z, s2z, f_low)
elif approximant == "IMRPhenomD":
time_length = lalsimulation.SimIMRPhenomDChirpTime(
time_length = lalsim.SimIMRPhenomDChirpTime(
m1 * lal.MSUN_SI, m2 * lal.MSUN_SI, s1z, s2z, f_low)
elif approximant == "SEOBNRv4":
# NB for no clear reason this function has f_low as first argument
time_length = lalsimulation.SimIMRSEOBNRv4ROMTimeOfFrequency(
time_length = lalsim.SimIMRSEOBNRv4ROMTimeOfFrequency(
f_low, m1 * lal.MSUN_SI, m2 * lal.MSUN_SI, s1z, s2z)
elif approximant == 'SEOBNRv5_ROM':
time_length = lalsimulation.SimIMRSEOBNRv5ROMTimeOfFrequency(
time_length = lalsim.SimIMRSEOBNRv5ROMTimeOfFrequency(
f_low, m1 * lal.MSUN_SI, m2 * lal.MSUN_SI, s1z, s2z)
elif approximant == 'SPAtmplt' or approximant == 'TaylorF2':
chi = lalsimulation.SimInspiralTaylorF2ReducedSpinComputeChi(
chi = lalsim.SimInspiralTaylorF2ReducedSpinComputeChi(
m1, m2, s1z, s2z
)
time_length = lalsimulation.SimInspiralTaylorF2ReducedSpinChirpTime(
time_length = lalsim.SimInspiralTaylorF2ReducedSpinChirpTime(
f_low, m1 * lal.MSUN_SI, m2 * lal.MSUN_SI, chi, -1
)
else:
Expand Down Expand Up @@ -568,42 +578,45 @@ class Params:

# FIXME spins are not taken into account
f_high = f_SchwarzISCO(mass1 + mass2)
tof_func = lambda f: findchirp_chirptime(
float(mass1),
float(mass2),
float(f),
pn_2order
)
def tof_func(f):
return findchirp_chirptime(
float(mass1),
float(mass2),
float(f),
pn_2order
)
elif approximant.startswith('SEOBNRv'):
approximant_prefix = approximant[:len('SEOBNRv*')]
f_high = get_final_freq(approximant_prefix, mass1, mass2, spin1, spin2)
f_high *= 0.999 # avoid errors due to rounding
tof_func_map = {
# use HI function for v2 as it has wider freq range validity
'SEOBNRv2': lalsimulation.SimIMRSEOBNRv2ROMDoubleSpinHITimeOfFrequency,
'SEOBNRv4': lalsimulation.SimIMRSEOBNRv4ROMTimeOfFrequency,
'SEOBNRv5': lalsimulation.SimIMRSEOBNRv5ROMTimeOfFrequency
'SEOBNRv2': lalsim.SimIMRSEOBNRv2ROMDoubleSpinHITimeOfFrequency,
'SEOBNRv4': lalsim.SimIMRSEOBNRv4ROMTimeOfFrequency,
'SEOBNRv5': lalsim.SimIMRSEOBNRv5ROMTimeOfFrequency
}
tof_func = lambda f: tof_func_map[approximant_prefix](
f,
solar_mass_to_kg(mass1),
solar_mass_to_kg(mass2),
float(spin1),
float(spin2)
)
def tof_func(f):
return tof_func_map[approximant_prefix](
f,
solar_mass_to_kg(mass1),
solar_mass_to_kg(mass2),
float(spin1),
float(spin2)
)
elif approximant in ['IMRPhenomD', 'IMRPhenomXAS']:
f_high = get_final_freq(approximant, mass1, mass2, spin1, spin2)
tof_func_map = {
'IMRPhenomD': lalsimulation.SimIMRPhenomDChirpTime,
'IMRPhenomXAS': lalsimulation.SimIMRPhenomXASDuration
'IMRPhenomD': lalsim.SimIMRPhenomDChirpTime,
'IMRPhenomXAS': lalsim.SimIMRPhenomXASDuration
}
tof_func = lambda f: tof_func_map[approximant](
solar_mass_to_kg(mass1),
solar_mass_to_kg(mass2),
float(spin1),
float(spin2),
f
)
def tof_func(f):
return tof_func_map[approximant](
solar_mass_to_kg(mass1),
solar_mass_to_kg(mass2),
float(spin1),
float(spin2),
f
)
else:
raise ValueError(f'Approximant {approximant} not supported')
track_f = numpy.logspace(numpy.log10(f_low), numpy.log10(f_high), n_points)
Expand Down Expand Up @@ -1043,7 +1056,7 @@ def jframe_to_l0frame(mass1, mass2, f_ref, phiref=0., thetajn=0., phijl=0.,
dimensionless spin.
"""
inclination, spin1x, spin1y, spin1z, spin2x, spin2y, spin2z = \
lalsimulation.SimInspiralTransformPrecessingNewInitialConditions(
lalsim.SimInspiralTransformPrecessingNewInitialConditions(
thetajn, phijl, spin1_polar, spin2_polar, spin12_deltaphi,
spin1_a, spin2_a, mass1*lal.MSUN_SI, mass2*lal.MSUN_SI, f_ref,
phiref)
Expand Down Expand Up @@ -1120,7 +1133,7 @@ def l0frame_to_jframe(mass1, mass2, f_ref, phiref=0., inclination=0.,
# Note: unlike other LALSimulation functions, this one takes masses in
# solar masses
thetajn, phijl, s1pol, s2pol, s12_deltaphi, spin1_a, spin2_a = \
lalsimulation.SimInspiralTransformPrecessingWvf2PE(
lalsim.SimInspiralTransformPrecessingWvf2PE(
inclination, spin1x, spin1y, spin1z, spin2x, spin2y, spin2z,
mass1, mass2, f_ref, phiref)
out = {'thetajn': thetajn,
Expand Down

0 comments on commit 0d39fd5

Please sign in to comment.