def __init__(
self,
q,
name="nrsur7dq2",
total_mass=None,
spin_1=None,
spin_2=None,
distance=None,
l_max=4,
modes=None,
times=None,
minimum_frequency=10,
sampling_frequency=4096,
):
"""
Initialise Surrogate MemoryGenerator
Parameters
----------
name: str
Name of the surrogate, default=NRSur7dq2.
l_max: int
Maximum ell value for oscillatory time series.
modes: dict, optional
Modes to load in, default is all ell<=4.
q: float
Binary mass ratio
total_mass: float, optional
Total binary mass in solar units.
distance: float, optional
Distance to the binary in MPC.
spin_1: array-like
Spin vector of more massive black hole.
spin_2: array-like
Spin vector of less massive black hole.
times: array-like
Time array to evaluate the waveforms on, default is
np.linspace(-900, 100, 10001).
"""
self.name = name
if name.lower() == "nrsur7dq2":
try:
from NRSur7dq2 import NRSurrogate7dq2
except ModuleNotFoundError:
print(
"nrsur7sq2 is required for the Surrogate memory generator."
)
raise
self.sur = NRSurrogate7dq2()
if q < 1:
q = 1 / q
if q > 2:
print("WARNING: Surrogate waveform not tested for q>2.")
self.q = q
self.MTot = total_mass
if spin_1 is None:
self.S1 = np.array([0.0, 0.0, 0.0])
else:
self.S1 = np.array(spin_1)
if spin_2 is None:
self.S2 = np.array([0.0, 0.0, 0.0])
else:
self.S2 = np.array(spin_2)
elif name.lower() == "nrhybsur3dq8":
try:
import gwsurrogate
except ModuleNotFoundError:
print(
"gwsurrogate is required for the Surrogate memory generator."
)
raise
self.sur = gwsurrogate.LoadSurrogate("NRHybSur3dq8")
if q < 1:
q = 1 / q
if q > 8:
print(
"WARNING: Hybrdid surrogate waveform not tested for q>8.")
self.q = q
self.MTot = total_mass
if spin_1 is None:
self.chi_1 = 0.0
else:
self.chi_1 = spin_1
if spin_2 is None:
self.chi_2 = 0.0
else:
self.chi_2 = spin_2
self.minimum_frequency = minimum_frequency
self.sampling_frequency = sampling_frequency
else:
raise ValueError(
"Surrogate model {} not implemented.".format(name))
self.distance = distance
self.LMax = l_max
self.modes = modes
if total_mass is None:
self.h_to_geo = 1
self.t_to_geo = 1
else:
self.h_to_geo = self.distance * MPC / self.MTot / SOLAR_MASS / GG * CC**2
self.t_to_geo = 1 / self.MTot / SOLAR_MASS / GG * CC**3
self.h_lm = None
self.times = times
if times is not None and max(times) < 10:
times *= self.t_to_geo
h_lm, times = self.time_domain_oscillatory(modes=modes, times=times)
MemoryGenerator.__init__(self, name=name, h_lm=h_lm, times=times)
import numpy as np
import gwsurrogate
gwsurrogate.catalog.pull('NRSur7dq4')
sur = gwsurrogate.LoadSurrogate('NRSur7dq4')
def _load_NRSur7dq4(self):
import gwsurrogate
from gwsurrogate.new.precessing_surrogate import splinterp_many
self.nrsur = gwsurrogate.LoadSurrogate('NRSur7dq4')
self.splinterp_many = splinterp_many
def __init__(self,
q,
name='nrsur7dq2',
total_mass=None,
spin_1=None,
spin_2=None,
distance=None,
l_max=4,
modes=None,
times=None,
minimum_frequency=10,
sampling_frequency=4096):
"""
Initialise Surrogate MemoryGenerator
Parameters
----------
name: str
Name of the surrogate, default=NRSur7dq2.
l_max: int
Maximum ell value for oscillatory time series.
modes: dict, optional
Modes to load in, default is all ell<=4.
q: float
Binary mass ratio
total_mass: float, optional
Total binary mass in solar units.
distance: float, optional
Distance to the binary in MPC.
spin_1: array-like
Spin vector of more massive black hole.
spin_2: array-like
Spin vector of less massive black hole.
times: array-like
Time array to evaluate the waveforms on, default is
np.linspace(-900, 100, 10001).
"""
self.name = name
# sur(x, dt=dt, f_low=f_low, mode_list=[(2,2), (2,1), (3, 3)], M=M, dist_mpc=dist_mpc,
# inclination=inclination, phi_ref=phi_ref, units='mks')
# h_lm = self.sur(self.q, self.S1, self.S2, MTot=self.MTot,
# distance=self.distance, t=times, LMax=self.LMax)
if name.lower() == "nrsur7dq2":
from NRSur7dq2 import NRSurrogate7dq2
self.sur = NRSurrogate7dq2()
if q < 1:
q = 1 / q
if q > 2:
print('WARNING: Surrogate waveform not tested for q>2.')
self.q = q
self.MTot = total_mass
if spin_1 is None:
self.S1 = np.array([0., 0., 0.])
else:
self.S1 = np.array(spin_1)
if spin_2 is None:
self.S2 = np.array([0., 0., 0.])
else:
self.S2 = np.array(spin_2)
elif name.lower() == "nrsur7dq4":
import gwsurrogate
gwsurrogate.catalog.pull('NRSur7dq4')
self.sur = gwsurrogate.LoadSurrogate('NRSur7dq4')
if q < 1:
q = 1 / q
if q > 4:
print(
'WARNING: Hybrdid surrogate waveform not tested for q>4.')
self.q = q
self.Mtot = total_mass
if spin_1 is None:
self.chi_1 = 0.0
else:
self.chi_1 = spin_1
if spin_2 is None:
self.chi_2 = 0.0
else:
self.chi_2 = spin_2
self.minimum_frequency = minimum_frequency
self.sampling_frequency = sampling_frequency
elif name.lower() == "nrhybsur3dq8":
import gwsurrogate
self.sur = gwsurrogate.LoadSurrogate('NRHybSur3dq8')
if q < 1:
q = 1 / q
if q > 8:
print(
'WARNING: Hybrdid surrogate waveform not tested for q>8.')
self.q = q
self.MTot = total_mass
if spin_1 is None:
self.chi_1 = 0.0
else:
self.chi_1 = spin_1
if spin_2 is None:
self.chi_2 = 0.0
else:
self.chi_2 = spin_2
self.minimum_frequency = minimum_frequency
self.sampling_frequency = sampling_frequency
else:
raise ValueError(
"Surrogate model {} not implemented.".format(name))
self.distance = distance
self.LMax = l_max
self.modes = modes
if total_mass is None:
self.h_to_geo = 1
self.t_to_geo = 1
else:
self.h_to_geo = self.distance * MPC / self.MTot /\
SOLAR_MASS / GG * CC ** 2
self.t_to_geo = 1. / self.MTot / SOLAR_MASS / GG * CC**3
self.h_lm = None
self.times = None
if times is not None and max(times) < 10:
times *= self.t_to_geo
h_lm, times = self.time_domain_oscillatory(modes=modes, times=times)
MemoryGenerator.__init__(self, name=name, h_lm=h_lm, times=times)
def __init__(self):
if not os.path.exists(gwsurrogate.__path__[0]+'/surrogate_downloads/NRHybSur3dq8.h5'):
gwsurrogate.catalog.pull('NRHybSur3dq8')
self.sur = gwsurrogate.LoadSurrogate('NRHybSur3dq8Tidal')
def __init__(self):
if not os.path.exists(gwsurrogate.__path__[0]+'/surrogate_downloads/NRSur7dq4.h5'):
gwsurrogate.catalog.pull('NRSur7dq4')
self.sur = gwsurrogate.LoadSurrogate('NRSur7dq4')
def test_model_regression(generate_regression_data=False):
""" If generate_regression_data = True, this script will generate
an hdf5 file to diff against. No regression will be done.
If generate_regression_data = False, this script will compare
model evaluations to the hdf5 file produced when True. In a typical use case,
this regression file will be downloaded. """
if generate_regression_data:
h5_file = "model_regression_data.h5"
print(
"Generating regression data file... Make sure this step is done BEFORE making any code changes!\n"
)
print(os.path.exists(h5_file))
if os.path.exists(h5_file):
raise RuntimeError("Refusing to overwrite a regression file!")
else:
h5_file = "test/comparison_data.h5" # assumes pytest runs from project-level folder
try: # try importing data. If it doesn't exist, download it
fp_regression = h5py.File("test/model_regression_data.h5", 'r')
except IOError:
print("Downloading regression data...")
os.system(
'wget --directory-prefix=test https://www.dropbox.com/s/vxqsr7fjoffxm5w/model_regression_data.h5'
)
fp_regression = h5py.File("test/model_regression_data.h5", 'r')
# remove models if you don't have them
dont_test = [
"EMRISur1dq1e4", # model data not currently available in public
"NRSur4d2s_TDROM_grid12", # 10 GB file
"NRSur4d2s_FDROM_grid12", # 10 GB file
#"SpEC_q1_10_NoSpin_linear_alt",
#"SpEC_q1_10_NoSpin_linear",
"EOBNRv2", #TODO: this is two surrogates in one. Break up?
#"SpEC_q1_10_NoSpin",
#"EOBNRv2_tutorial",
#"NRHybSur3dq8",
#"NRHybSur3dq8Tidal",
#"NRSur7dq4"
]
# Common directory where all surrogates are assumed to be located
surrogate_path = gws.catalog.download_path()
# repeatability can be useful for regression tests
np.random.seed(0)
# for each model, associate its surrogate data file
models = [model for model in gws.catalog._surrogate_world]
print(models)
models_to_test = {}
for model in models:
surrogate_data = surrogate_path + os.path.basename(
gws.catalog._surrogate_world[model][0])
if os.path.isfile(surrogate_data): # surrogate data file exists
models_to_test[model] = surrogate_data
else: # file missing
msg = "WARNING: Surrogate missing!!!\n"
msg += "Surrogate data assumed to be in the path %s.\n" % surrogate_data
msg += "If the data is somewhere else, change the path or move the file.\n\n"
msg += "To download this surrogate, from ipython do\n\n >>> gws.catalog.pull(%s)\n" % model
print(msg)
time.sleep(1)
# also test the tutorial surrogate
models_to_test["EOBNRv2_tutorial"] = gws.__path__[
0] + "/../tutorial/TutorialSurrogate/EOB_q1_2_NoSpin_Mode22/"
# remove models from testing...
for i in dont_test:
try:
models_to_test.pop(i)
print("model %s removed from testing" % i)
except KeyError:
print("model %s cannot be removed" % i)
fp = h5py.File(h5_file, "w")
# for each model, select three random points to evaluate at
param_samples_tested = []
for model, datafile in models_to_test.items():
print("Generating regression data for model = %s" % model)
print(datafile)
if model in surrogate_old_interface:
sur = gws.EvaluateSurrogate(datafile)
p_mins = sur.param_space.min_vals()
p_maxs = sur.param_space.max_vals()
elif model in surrogate_loader_interface:
# sur = gws.LoadSurrogate(datafile) # tidal and aligned models use the same h5 file (so wrong one loaded)
sur = gws.LoadSurrogate(model)
try:
p_mins = sur._sur_dimless.param_space.min_vals()
p_maxs = sur._sur_dimless.param_space.max_vals()
except AttributeError: # NRSur7dq4 does not have object sur._sur_dimless.param_space
p_mins = None
p_maxs = None
else:
sur = surrogate.FastTensorSplineSurrogate()
sur.load(datafile)
p_mins = sur.param_space.min_vals()
p_maxs = sur.param_space.max_vals()
print("parameter minimum values", p_mins)
print("parameter maximum values", p_maxs)
param_samples = []
if generate_regression_data: # pick new points to compute regression data at
for i in range(3): # sample parameter space 3 times
if model in list(model_sampler.keys()):
custom_sampler = model_sampler[model]
x, tidOpts, pecOpts = custom_sampler(
i) # [q, chiA, chiB], tidOpts
else: # default sampler for spin-aligned BBH models
x = [] # [q, chiAz, chiBz]
for j in range(len(p_mins)):
xj_min = p_mins[j]
xj_max = p_maxs[j]
tmp = float(np.random.uniform(xj_min, xj_max, size=1))
x.append(tmp)
tidOpts = None
pecOpts = None
param_samples.append([x, tidOpts, pecOpts])
else: # get point at which to compute comparison waveform data
for i in range(3):
if model in list(
model_sampler.keys()): # use sample points if provided
custom_sampler = model_sampler[model]
x, tidOpts, pecOpts = custom_sampler(
i) # [q, chiA, chiB], tidOpts
else: # pull regression points from regression data file; none-tidal models only
print(model + "/parameter%i/parameter" % i)
x = []
x.append(
list(fp_regression[model + "/parameter%i/parameter" %
i][:])) # [q, chiAz, chiBz]
x = list(fp_regression[model + "/parameter%i/parameter" %
i][:]) # [q, chiAz, chiBz]
tidOpts = None
pecOpts = None
param_samples.append([x, tidOpts, pecOpts])
param_samples_tested.append(param_samples)
model_grp = fp.create_group(model)
for i, ps in enumerate(param_samples):
x = ps[0]
tidOpts = ps[1]
pecOpts = ps[2]
if model in surrogate_old_interface:
ps_float = x[0] # TODO: generalize interface
modes, t, hp, hc = sur(q=ps_float,
mode_sum=False,
fake_neg_modes=True)
else:
if model in surrogate_loader_interface:
q = x[0]
if type(x[1]) is np.float64 or type(x[1]) is float: # chiz
chiA = np.array([0, 0, x[1]])
chiB = np.array([0, 0, x[2]])
elif len(x[1]) == 3: # spin vector
chiA = np.array(x[1])
chiB = np.array(x[2])
else:
raise ValueError
try:
# Regression samples outside of the training interval.
# Warnings are raised (as they should) but could
# appear bad to a new user
if model in ["NRSur7dq4"]:
with warnings.catch_warnings():
warnings.simplefilter("ignore")
t, h, dyanmics = sur(q,
chiA,
chiB,
f_low=0.0,
tidal_opts=tidOpts,
precessing_opts=pecOpts)
else:
t, h, dyanmics = sur(q,
chiA,
chiB,
f_low=0.0,
tidal_opts=tidOpts,
precessing_opts=pecOpts)
except ValueError: # some models do not allow for f_low=0.0 and require a time step
# step size, Units of M
# initial frequency, Units of cycles/M
t, h, dyanmics = sur(q,
chiA,
chiB,
dt=0.25,
f_low=3.e-3,
tidal_opts=tidOpts,
precessing_opts=pecOpts)
else:
h = sur(x)
try:
h_np = [h[mode] for mode in sur.mode_list]
except AttributeError: # for new interface
h_np = [h[mode] for mode in sur._sur_dimless.mode_list]
h_np = np.vstack(h_np)
hp = np.real(h_np)
hc = np.imag(h_np)
samplei = model_grp.create_group("parameter" + str(i))
if model in list(model_sampler.keys()):
# model samplers return a list of lists, which we flatten into
# a list of numbers for storing in an h5 dataset
x = flatten_params(x)
samplei.create_dataset("parameter", data=x)
samplei.create_dataset("hp", data=hp, dtype='float32')
samplei.create_dataset("hc", data=hc, dtype='float32')
fp.close()
if not generate_regression_data:
fp = h5py.File(h5_file, "r") # reopen comparison data
for model in models_to_test.keys():
print("testing model %s ..." % model)
for i in range(3): # 3 parameter samples
hp_regression = fp_regression[model + "/parameter%i/hp" % i][:]
hc_regression = fp_regression[model + "/parameter%i/hp" % i][:]
hp_comparison = fp[model + "/parameter%i/hp" % i][:]
hc_comparison = fp[model + "/parameter%i/hp" % i][:]
if model == "NRHybSur3dq8":
local_rtol = rtol_gsl
else:
local_rtol = rtol_gsl
np.testing.assert_allclose(hp_regression,
hp_comparison,
rtol=local_rtol,
atol=atol)
np.testing.assert_allclose(hc_regression,
hc_comparison,
rtol=local_rtol,
atol=atol)
# fails due to round-off error differences of different machines
#fp_regression.close()
#process = subprocess.Popen(["h5diff", "test/model_regression_data.h5",h5_file],
# stdin=subprocess.PIPE,
# stdout=subprocess.PIPE,
# stderr=subprocess.PIPE)
#returncode = process.wait()
#stdout, stderr = process.communicate()
#if returncode == 0:
# assert(True)
#else:
# print(stdout)
# print(stderr)
# assert(False)
print("models tested... ")