Commit 408522be authored by Gregory Ashton's avatar Gregory Ashton

Clean up and comment of new tutorial

- Removes the old custom_source tutorial
- Adds comments to the create_your_own_source_model
parent ef6e8151
#!/bin/python
"""
A script to demonstrate how to use your own source model
"""
from __future__ import division, print_function
import tupak
import numpy as np
# First set up logging and some output directories and labels
tupak.utils.setup_logger()
outdir = 'outdir'
label = 'create_your_own_source_model'
sampling_frequency = 4096
time_duration = 1
# Here we define out source model - this is the sine-Gaussian model in the
# frequency domain.
def sine_gaussian(f, A, f0, tau, phi0, geocent_time, ra, dec, psi):
arg = -(np.pi * tau * (f-f0))**2 + 1j * phi0
plus = np.sqrt(np.pi) * A * tau * np.exp(arg) / 2.
......@@ -16,33 +23,35 @@ def sine_gaussian(f, A, f0, tau, phi0, geocent_time, ra, dec, psi):
return {'plus': plus, 'cross': cross}
outdir = 'outdir'
label = 'GW150914_sine_gaussian'
time_of_event = 1126259462.422
H1 = tupak.detector.get_interferometer('H1', time_of_event, version=1, outdir=outdir)
L1 = tupak.detector.get_interferometer('L1', time_of_event, version=1, outdir=outdir)
interferometers = [H1, L1]
prior = dict()
prior['A'] = tupak.prior.Uniform(0, 1e-20, 'A')
prior['f0'] = tupak.prior.Uniform(0, 10, 'f')
prior['tau'] = tupak.prior.Uniform(0, 10, 'tau')
prior['geocent_time'] = tupak.prior.Uniform(
time_of_event-0.1, time_of_event+0.1, 'geocent_time')
prior['phi0'] = 0 #tupak.prior.Uniform(0, 2*np.pi, 'phi')
prior['ra'] = 0
prior['dec'] = 0
prior['psi'] = 0
# We now define some parameters that we will inject and then a waveform generator
injection_parameters = dict(A=1e-21, f0=10, tau=1, phi0=0, geocent_time=0,
ra=0, dec=0, psi=0)
waveform_generator = tupak.waveform_generator.WaveformGenerator(
sine_gaussian, H1.sampling_frequency, H1.duration)
frequency_domain_source_model=sine_gaussian,
sampling_frequency=sampling_frequency,
time_duration=time_duration,
parameters=injection_parameters)
hf_signal = waveform_generator.frequency_domain_strain()
# Set up interferometers.
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']]
# Here we define the priors for the search. We use the injection parameters
# except for the amplitude, f0, and geocent_time
prior = injection_parameters.copy()
prior['A'] = tupak.prior.Uniform(0, 1e-20, 'A')
prior['f0'] = tupak.prior.Uniform(0, 20, 'f')
prior['geocent_time'] = tupak.prior.Uniform(-0.01, 0.01, 'geocent_time')
likelihood = tupak.likelihood.Likelihood(interferometers, waveform_generator)
likelihood = tupak.likelihood.Likelihood(IFOs, waveform_generator)
result = tupak.sampler.run_sampler(
likelihood, prior, sampler='pymultinest', outdir=outdir, label=label,
resume=False)
likelihood, prior, sampler='dynesty', outdir=outdir, label=label,
resume=False, sample='unif')
result.plot_walks()
result.plot_distributions()
result.plot_corner()
......
"""
WIP
"""
import numpy as np
import pylab as plt
import dynesty.plotting as dyplot
import corner
import tupak
tupak.utils.setup_logger()
def generate_and_plot_data(waveform_generator, injection_parameters):
hf_signal = waveform_generator.frequency_domain_strain()
# Simulate the data in H1
H1 = tupak.detector.H1
H1_hf_noise, frequencies = H1.power_spectral_density. \
get_noise_realisation(waveform_generator.sampling_frequency,
waveform_generator.time_duration)
H1.set_data(waveform_generator.sampling_frequency,
waveform_generator.time_duration,
frequency_domain_strain=H1_hf_noise)
H1.inject_signal(hf_signal, injection_parameters)
H1.set_spectral_densities()
# Simulate the data in L1
L1 = tupak.detector.L1
L1_hf_noise, frequencies = L1.power_spectral_density. \
get_noise_realisation(waveform_generator.sampling_frequency,
waveform_generator.time_duration)
L1.set_data(waveform_generator.sampling_frequency,
waveform_generator.time_duration,
frequency_domain_strain=L1_hf_noise)
L1.inject_signal(hf_signal, injection_parameters)
L1.set_spectral_densities()
IFOs = [H1, L1]
# Plot the noise and signal
fig, ax = plt.subplots()
plt.loglog(frequencies, np.abs(H1_hf_noise), lw=1.5, label='H1 noise+signal')
plt.loglog(frequencies, np.abs(L1_hf_noise), lw=1.5, label='L1 noise+signal')
plt.loglog(frequencies, np.abs(hf_signal['plus']), lw=0.8, label='signal')
plt.xlim(10, 1000)
plt.legend()
plt.xlabel(r'frequency')
plt.ylabel(r'strain')
fig.savefig('data')
return IFOs
def delta_function_frequency_domain_strain(frequency_array, amplitude,
peak_time, phase, **kwargs):
ht = {'plus': amplitude * np.sin(2 * np.pi * peak_time * frequency_array + phase),
'cross': amplitude * np.cos(2 * np.pi * peak_time * frequency_array + phase)}
return ht
def gaussian_frequency_domain_strain(frequency_array, amplitude, mu, sigma, **kwargs):
ht = {'plus': amplitude * np.exp(-(mu - frequency_array) ** 2 / sigma ** 2 / 2),
'cross': amplitude * np.exp(-(mu - frequency_array) ** 2 / sigma ** 2 / 2)}
return ht
injection_parameters = dict(amplitude=1e-21,
mu=100,
sigma=1,
ra=1.375,
dec=-1.2108,
geocent_time=1126259642.413,
psi=2.659)
sampling_parameters = tupak.prior.parse_floats_to_fixed_priors(injection_parameters)
wg = tupak.waveform_generator.WaveformGenerator(
frequency_domain_source_model=gaussian_frequency_domain_strain,
parameters=injection_parameters
)
IFOs = generate_and_plot_data(wg, injection_parameters)
likelihood = tupak.likelihood.Likelihood(IFOs, wg)
sampling_parameters['amplitude'] = tupak.prior.Uniform(minimum=0.9 * 1e-21, maximum=1.1 * 1e-21)
sampling_parameters['sigma'] = tupak.prior.Uniform(minimum=0, maximum=10)
sampling_parameters['mu'] = tupak.prior.Uniform(minimum=50, maximum=200)
result = tupak.sampler.run_sampler(likelihood, priors=sampling_parameters, verbose=True)
#
# Make some nice plots
#
truths = [injection_parameters[x] for x in result.search_parameter_keys]
corner_plot = corner.corner(result.samples, truths=truths, labels=result.search_parameter_keys)
corner_plot.savefig('corner')
trace_plot, axes = dyplot.traceplot(result['sampler_output'])
trace_plot.savefig('trace')
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment