diff --git a/examples/supernova_example/supernova_example.py b/examples/supernova_example/supernova_example.py
index d3ea31cfad66d20db2d52ab3ed600d95214ebd9d..51787a6216a8486a00579f9311b6ca22fcb2d954 100644
--- a/examples/supernova_example/supernova_example.py
+++ b/examples/supernova_example/supernova_example.py
@@ -1,8 +1,10 @@
#!/bin/python
"""
-Tutorial to demonstrate running parameter estimation/model selection on an NR supernova injected signal.
-Signal model is made by applying PCA to a set of supernova waveforms. The first few PCs are then linearly
-combined with a scale factor. (See https://arxiv.org/pdf/1202.3256.pdf)
+
+Tutorial to demonstrate running parameter estimation/model selection on an NR
+supernova injected signal. Signal model is made by applying PCA to a set of
+supernova waveforms. The first few PCs are then linearly combined with a scale
+factor. (See https://arxiv.org/pdf/1202.3256.pdf)
"""
from __future__ import division, print_function
@@ -21,49 +23,51 @@ tupak.utils.setup_logger(outdir=outdir, label=label)
# Set up a random seed for result reproducibility. This is optional!
np.random.seed(170801)
-# We are going to inject a supernova waveform. We first establish a dictionary of parameters that
-# includes all of the different waveform parameters. It will read in a signal to inject from a txt file.
-injection_parameters = dict(file_path = 'MuellerL15_example_inj.txt', luminosity_distance = 10.0, ra = 1.375,
- dec = -1.2108, geocent_time = 1126259642.413, psi= 2.659)
+# We are going to inject a supernova waveform. We first establish a dictionary
+# of parameters that includes all of the different waveform parameters. It will
+# read in a signal to inject from a txt file.
+
+injection_parameters = dict(file_path='MuellerL15_example_inj.txt',
+ luminosity_distance=10.0, ra=1.375,
+ dec=-1.2108, geocent_time=1126259642.413,
+ psi=2.659)
# Create the waveform_generator using a supernova source function
-waveform_generator = tupak.waveform_generator.WaveformGenerator(time_duration=time_duration,
- sampling_frequency=sampling_frequency,
- frequency_domain_source_model=tupak.source.supernova,
- parameters=injection_parameters)
+waveform_generator = tupak.waveform_generator.WaveformGenerator(
+ time_duration=time_duration, sampling_frequency=sampling_frequency,
+ frequency_domain_source_model=tupak.source.supernova,
+ parameters=injection_parameters)
hf_signal = waveform_generator.frequency_domain_strain()
-# Set up interferometers. In this case we'll use three interferometers (LIGO-Hanford (H1), LIGO-Livingston (L1),
-# and Virgo (V1)). These default to their design sensitivity
+# Set up interferometers. In this case we'll use three interferometers
+# (LIGO-Hanford (H1), LIGO-Livingston (L1), and Virgo (V1)). These default to
+# their design sensitivity
IFOs = [tupak.detector.get_interferometer_with_fake_noise_and_injection(
- name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
- sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1', 'V1']]
+ name, injection_polarizations=hf_signal,
+ injection_parameters=injection_parameters, time_duration=time_duration,
+ sampling_frequency=sampling_frequency, outdir=outdir)
+ for name in ['H1', 'L1', 'V1']]
# read in from a file the PCs used to create the signal model.
realPCs = np.loadtxt('SupernovaRealPCs.txt')
imagPCs = np.loadtxt('SupernovaImagPCs.txt')
-# now we have to do the waveform_generator again because the signal model is not the same as the injection in this case.
-simulation_parameters = dict(realPCs=realPCs, imagPCs=imagPCs, pc_coeff1 = 0.1, pc_coeff2 = 0.1,
- pc_coeff3 = 0.1, pc_coeff4 = 0.1, pc_coeff5 = 0.1, luminosity_distance = 10.0,
- ra = 1.375, dec = -1.2108, geocent_time = 1126259642.413, psi=2.659)
+# Now we make another waveform_generator because the signal model is
+# not the same as the injection in this case.
+simulation_parameters = dict(
+ realPCs=realPCs, imagPCs=imagPCs)
-waveform_generator = tupak.waveform_generator.WaveformGenerator(time_duration=time_duration,
- sampling_frequency=sampling_frequency,
- frequency_domain_source_model=tupak.source.supernova_pca_model,
- parameters=simulation_parameters)
+search_waveform_generator = tupak.waveform_generator.WaveformGenerator(
+ time_duration=time_duration, sampling_frequency=sampling_frequency,
+ frequency_domain_source_model=tupak.source.supernova_pca_model,
+ parameters=simulation_parameters)
-# Set up prior, which is a dictionary
+# Set up prior
priors = dict()
-# By default we will sample all terms in the signal models. However, this will take a long time for the calculation,
-# so for this example we will set almost all of the priors to be equall to their injected values. This implies the
-# prior is a delta function at the true, injected value. In reality, the sampler implementation is smart enough to
-# not sample any parameter that has a delta-function prior.
for key in ['psi', 'geocent_time']:
priors[key] = injection_parameters[key]
-
-# don't use default for luminosity distance because we want kpc not Mpc
-priors['luminosity_distance'] = tupak.prior.Uniform(2, 20, 'luminosity_distance')
+priors['luminosity_distance'] = tupak.prior.Uniform(
+ 2, 20, 'luminosity_distance')
priors['pc_coeff1'] = tupak.prior.Uniform(-1, 1, 'pc_coeff1')
priors['pc_coeff2'] = tupak.prior.Uniform(-1, 1, 'pc_coeff2')
priors['pc_coeff3'] = tupak.prior.Uniform(-1, 1, 'pc_coeff3')
@@ -71,13 +75,20 @@ priors['pc_coeff4'] = tupak.prior.Uniform(-1, 1, 'pc_coeff4')
priors['pc_coeff5'] = tupak.prior.Uniform(-1, 1, 'pc_coeff5')
priors['ra'] = tupak.prior.create_default_prior(name='ra')
priors['dec'] = tupak.prior.create_default_prior(name='dec')
-
-# Initialise the likelihood by passing in the interferometer data (IFOs) and the waveoform generator
-likelihood = tupak.likelihood.GravitationalWaveTransient(interferometers=IFOs, waveform_generator=waveform_generator)
-
-# Run sampler. In this case we're going to use the `dynesty` sampler
-result = tupak.sampler.run_sampler(likelihood=likelihood, priors=priors, sampler='dynesty', npoints=100,
- outdir=outdir, label=label)
+priors['geocent_time'] = tupak.prior.Uniform(
+ injection_parameters['geocent_time'] - 1,
+ injection_parameters['geocent_time'] + 1,
+ 'geocent_time')
+
+# Initialise the likelihood by passing in the interferometer data (IFOs) and
+# the waveoform generator
+likelihood = tupak.GravitationalWaveTransient(
+ interferometers=IFOs, waveform_generator=search_waveform_generator)
+
+# Run sampler.
+result = tupak.run_sampler(
+ likelihood=likelihood, priors=priors, sampler='dynesty', npoints=500,
+ outdir=outdir, label=label)
# make some plots of the outputs
result.plot_corner()