Commit f9952731 authored by Yifan Wang's avatar Yifan Wang
Browse files

complete the traning data part, also plot some figures

parent 18514b44
This diff is collapsed.
This diff is collapsed.
......@@ -24,24 +24,8 @@ from utils.waveforms import WaveformParameterGenerator
if __name__ == '__main__':
# -------------------------------------------------------------------------
# Preliminaries
# -------------------------------------------------------------------------
# Disable output buffering ('flush' option is not available for Python 2)
#sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
# Start the stopwatch
script_start = time.time()
print('')
print('GENERATE A GW DATA SAMPLE FILE')
print('')
# -------------------------------------------------------------------------
# Parse the command line arguments
# -------------------------------------------------------------------------
# Set up the parser and add arguments
parser = argparse.ArgumentParser(description='Generate a GW data sample.')
parser.add_argument('--config-file',
......@@ -50,9 +34,7 @@ if __name__ == '__main__':
default='default.json')
# Parse the arguments that were passed when calling this script
print('Parsing command line arguments...', end=' ')
command_line_arguments = vars(parser.parse_args())
print('Done!')
# -------------------------------------------------------------------------
# Read in JSON config file specifying the sample generation process
......@@ -89,8 +71,6 @@ if __name__ == '__main__':
# Define some useful shortcuts
random_seed = config['random_seed']
max_runtime = config['max_runtime']
bkg_data_dir = config['background_data_directory']
# -------------------------------------------------------------------------
# Construct a generator for sampling waveform parameters
......@@ -107,46 +87,14 @@ if __name__ == '__main__':
waveform_parameters = \
(waveform_parameter_generator.draw() for _ in iter(int, 1))
# -------------------------------------------------------------------------
# Construct a generator for sampling valid noise times
# -------------------------------------------------------------------------
print('Using synthetic noise! (background_data_directory = None)\n')
# If the 'background_data_directory' is None, we will use synthetic noise
if config['background_data_directory'] is None:
print('Using synthetic noise! (background_data_directory = None)\n')
# Create a iterator that returns a fake "event time", which we will
# use as a seed for the RNG to ensure the reproducibility of the
# generated synthetic noise.
# For the HDF file path that contains that time, we always yield
# None, so that we know that we need to generate synthetic noise.
noise_times = ((1000000000 + _, None) for _ in count())
# Otherwise, we set up a timeline object for the background noise, that
# is, we read in all HDF files in the raw_data_directory and figure out
# which parts of it are useable (i.e., have the right data quality and
# injection bits set as specified in the config file).
else:
print('Using real noise from LIGO recordings! '
'(background_data_directory = {})'.format(bkg_data_dir))
print('Reading in raw data. This may take several minutes...', end=' ')
# Create a timeline object by running over all HDF files once
noise_timeline = NoiseTimeline(background_data_directory=bkg_data_dir,
random_seed=random_seed)
# Create a noise time generator so that can sample valid noise times
# simply by calling next(noise_time_generator)
delta_t = int(static_arguments['noise_interval_width'] / 2)
noise_times = (noise_timeline.sample(delta_t=delta_t,
dq_bits=config['dq_bits'],
inj_bits=config['inj_bits'],
return_paths=True)
for _ in iter(int, 1))
print('Done!\n')
# Create a iterator that returns a fake "event time", which we will
# use as a seed for the RNG to ensure the reproducibility of the
# generated synthetic noise.
# For the HDF file path that contains that time, we always yield
# None, so that we know that we need to generate synthetic noise.
noise_times = ((1000000000 + _, None) for _ in count())
# -------------------------------------------------------------------------
# Define a convenience function to generate arguments for the simulation
......@@ -170,23 +118,43 @@ if __name__ == '__main__':
samples = dict(injection_samples=[], noise_samples=[])
injection_parameters = dict(injection_samples=[], noise_samples=[])
for sample_type in ['injection_samples', 'noise_samples']:
print('Generating samples containing an injection...')
n_samples = config['n_injection_samples']
arguments_generator = \
# ---------------------------------------------------------------------
# Define some sample_type-specific shortcuts
# ---------------------------------------------------------------------
if sample_type == 'injection_samples':
print('Generating samples containing an injection...')
n_samples = config['n_injection_samples']
arguments_generator = \
(generate_arguments(injection=True) for _ in iter(int, 1))
print('Number of samples:',n_samples)
else:
print('Generating samples *not* containing an injection...')
n_samples = config['n_noise_samples']
arguments_generator = \
(generate_arguments(injection=False) for _ in iter(int, 1))
sample_type = 'injection_samples'
for i in range(n_samples):
print(i)
# ---------------------------------------------------------------------
# If we do not need to generate any samples, skip ahead:
# ---------------------------------------------------------------------
if n_samples == 0:
print('Done! (n_samples=0)\n')
continue
results_list = []
arguments = next(arguments_generator)
print(arguments)
result = generate_sample(**arguments)
results_list.append(result)
# ---------------------------------------------------------------------
# Loop over injections or noise:
# ---------------------------------------------------------------------
for i in range(n_samples):
print('Generating the',i,'th samples\n')
arguments = next(arguments_generator)
print(arguments)
result = generate_sample(**arguments)
results_list.append(result)
# ---------------------------------------------------------------------
# Process results in the results_list
# ---------------------------------------------------------------------
......@@ -202,7 +170,10 @@ if __name__ == '__main__':
injection_parameters[sample_type] = \
list([injection_parameters[sample_type][i] for i in idx])
print('Sample generation completed!\n')
if sample_type == 'injection_samples':
print('Signal+noise generation completed!\n')
else:
print('Noise generation completed!\n')
# -------------------------------------------------------------------------
# Compute the normalization parameters for this file
......@@ -248,8 +219,8 @@ if __name__ == '__main__':
static_arguments=static_arguments)
# Collect and add samples (with and without injection)
for sample_type in ('injection_samples', 'noise_samples'):
for key in ('event_time', 'h1_strain', 'l1_strain'):
for sample_type in ['injection_samples', 'noise_samples']:
for key in ['event_time', 'h1_strain', 'l1_strain']:
if samples[sample_type]:
value = np.array([_[key] for _ in list(samples[sample_type])])
else:
......@@ -258,10 +229,10 @@ if __name__ == '__main__':
# Collect and add injection_parameters (ignore noise samples here, because
# for those, the injection_parameters are always None)
other_keys = ['h1_signal', 'h1_snr', 'l1_signal', 'l1_snr', 'scale_factor']
print(injection_parameters['injection_samples'])
#other_keys = ['h1_signal', 'h1_snr', 'l1_signal', 'l1_snr', 'scale_factor']
other_keys = ['h1_signal', 'h1_output_signal','h1_snr', 'l1_signal','l1_output_signal', 'l1_snr', 'scale_factor']
for key in list(variable_arguments + other_keys):
#print(key)
if injection_parameters['injection_samples']:
value = np.array([_[key] for _ in
injection_parameters['injection_samples']])
......
import os
for i in range(10):
print(i)
os.system('python plot_sample.py --hdf-file-path output/train.hdf --sample-id '+str(i)+' --plot-path "plots/'+str(i)+'.png"')
for i in range(100,110):
print(i)
os.system('python plot_sample.py --hdf-file-path output/train.hdf --sample-id '+str(i)+' --plot-path "plots/'+str(i)+'.png"')
This diff is collapsed.
This diff is collapsed.
......@@ -5,8 +5,8 @@
"inj_bits": [0, 1, 2, 4],
"waveform_params_file_name": "waveform_params.ini",
"max_runtime": 60,
"n_injection_samples": 4,
"n_noise_samples": 16,
"n_injection_samples": 100,
"n_noise_samples": 100,
"n_processes": 4,
"output_file_name": "default.hdf"
"output_file_name": "train.hdf"
}
; -----------------------------------------------------------------------------
; DECLARE ARGUMENTS
; -----------------------------------------------------------------------------
[variable_args]
; Waveform parameters that will vary in MCMC
mass1 =
mass2 =
spin1z =
......@@ -15,9 +10,7 @@ inclination =
polarization =
injection_snr =
[static_args]
; Waveform parameters that will not change in MCMC
approximant = SEOBNRv4
domain = time
f_lower = 18
......@@ -57,60 +50,39 @@ seconds_after_event = 2.5
tukey_alpha = 0.25
; -----------------------------------------------------------------------------
; DEFINE DISTRIBUTIONS FOR PARAMETERS
; -----------------------------------------------------------------------------
[prior-mass1]
; Prior for mass1
name = uniform
min-mass1 = 10.
max-mass1 = 80.
[prior-mass2]
; Prior for mass2
name = uniform
min-mass2 = 10.
max-mass2 = 80.
[prior-spin1z]
; Prior for spin1z
name = uniform
min-spin1z = 0
max-spin1z = 0.998
[prior-spin2z]
; Prior for spin2z
name = uniform
min-spin2z = 0
max-spin2z = 0.998
[prior-injection_snr]
; Prior for the injection SNR
name = uniform
min-injection_snr = 5
max-injection_snr = 20
[prior-coa_phase]
; Coalescence phase prior
name = uniform_angle
[prior-inclination]
; Inclination prior
name = sin_angle
[prior-ra+dec]
; Sky position prior
name = uniform_sky
[prior-polarization]
; Polarization prior
name = uniform_angle
"""
Plot the results produced by the generate_sample.py script.
"""
import argparse
import numpy as np
import os
import sys
import time
from utils.samplefiles import SampleFile
# We need to load a different backend for matplotlib before import plt to
# avoid problems on environments where the $DISPLAY variable is not set.
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt # noqa
# -----------------------------------------------------------------------------
# MAIN CODE
# -----------------------------------------------------------------------------
if __name__ == '__main__':
# -------------------------------------------------------------------------
# Preliminaries
# -------------------------------------------------------------------------
# Start the stopwatch
script_start_time = time.time()
print('')
print('PLOT A GENERATED SAMPLE (WITH / WITHOUT AN INJECTION)')
print('')
# -------------------------------------------------------------------------
# Parse the command line arguments
# -------------------------------------------------------------------------
# Set up the parser
parser = argparse.ArgumentParser(description='Plot a generated sample.')
# Add arguments (and set default values where applicable)
parser.add_argument('--hdf-file-path',
help='Path to the HDF sample file (generated with '
'generate_sample.py) to be used. '
'Default: ./output/default.hdf.',
default='./output/default.hdf')
parser.add_argument('--sample-id',
help='ID of the sample to be viewed (an integer '
'between 0 and n_injection_samples + '
'n_noise_samples). Default: 0.',
default=0)
parser.add_argument('--seconds-before',
help='Seconds to plot before the event_time. '
'Default: 0.15.',
default=0.15)
parser.add_argument('--seconds-after',
help='Seconds to plot after the event_time. '
'Default: 0.05.',
default=0.05)
parser.add_argument('--plot-path',
help='Where to save the plot of the sample. '
'Default: ./sample.pdf.',
default='sample.pdf')
# Parse the arguments that were passed when calling this script
print('Parsing command line arguments...', end=' ')
arguments = vars(parser.parse_args())
print('Done!')
# Set up shortcuts for the command line arguments
hdf_file_path = str(arguments['hdf_file_path'])
sample_id = int(arguments['sample_id'])
seconds_before = float(arguments['seconds_before'])
seconds_after = float(arguments['seconds_after'])
plot_path = str(arguments['plot_path'])
# -------------------------------------------------------------------------
# Read in the sample file
# -------------------------------------------------------------------------
print('Reading in HDF file...', end=' ')
data = SampleFile()
data.read_hdf(hdf_file_path)
df = data.as_dataframe(injection_parameters=True,
static_arguments=True)
print('Done!')
# -------------------------------------------------------------------------
# Plot the desired sample
# -------------------------------------------------------------------------
print('Plotting sample...', end=' ')
# Select the sample (i.e., the row from the data frame of samples)
try:
sample = df.loc[sample_id]
except KeyError:
raise KeyError('Given sample_id is too big! Maximum value = {}'.
format(len(df) - 1))
# Check if the sample we have received contains an injection or not
if 'h1_signal' in sample.keys():
has_injection = isinstance(sample['h1_signal'], np.ndarray)
else:
has_injection = False
# Read out and construct some necessary values for plotting
seconds_before_event = float(sample['seconds_before_event'])
seconds_after_event = float(sample['seconds_after_event'])
target_sampling_rate = float(sample['target_sampling_rate'])
sample_length = float(sample['sample_length'])
# Create a grid on which the sample can be plotted so that the
# event_time is at position 0
grid = np.linspace(0 - seconds_before_event, 0 + seconds_after_event,
int(target_sampling_rate * sample_length))
# Create subplots for H1 and L1
fig, axes1 = plt.subplots(nrows=2)
# If the sample has an injection, we need a second y-axis to plot the
# pure (i.e., unwhitened) detector signals
if has_injection:
axes2 = [ax.twinx() for ax in axes1]
else:
axes2 = None
# Plot the strains for H1 and L1
for i, (det_name, det_string) in enumerate([('H1', 'h1_strain'),
('L1', 'l1_strain')]):
axes1[i].plot(grid, sample[det_string], color='C0')
axes1[i].set_xlim(-seconds_before, seconds_after)
axes1[i].set_ylim(-150, 150)
axes1[i].tick_params('y', colors='C0', labelsize=8)
axes1[i].set_ylabel('Amplitude of Whitened Strain ({})'
.format(det_name), color='C0', fontsize=8)
# If applicable, also plot the detector signals for H1 and L1
if has_injection:
# Get the maximum value of the detector signal (for norming them)
maximum = max(np.max(sample['h1_signal']), np.max(sample['l1_signal']))
for i, (det_name, det_string) in enumerate([('H1', 'h1_signal'),
('L1', 'l1_signal')]):
axes2[i].plot(grid, sample[det_string] / maximum, color='C1')
axes2[i].set_xlim(-seconds_before, seconds_after)
axes2[i].set_ylim(-1.2, 1.2)
axes2[i].tick_params('y', colors='C1', labelsize=8)
axes2[i].set_ylabel('Rescaled Amplitude of Simulated\n'
'Detector Signal ({})'.format(det_name),
color='C1', fontsize=8)
# Also add the injection parameters
if has_injection:
keys = ('mass1', 'mass2', 'spin1z', 'spin2z', 'ra', 'dec',
'coa_phase', 'inclination', 'polarization', 'injection_snr')
string = ', '.join(['{} = {:.2f}'.format(_, float(sample[_]))
for _ in keys])
else:
string = '(sample does not contain an injection)'
plt.figtext(0.5, 0.9, 'Injection Parameters:\n' + string,
fontsize=8, ha='center')
# Add a vertical line at the position of the event (x=0)
axes1[0].axvline(x=0, color='black', ls='--', lw=1)
axes1[1].axvline(x=0, color='black', ls='--', lw=1)
# Set x-labels
axes1[0].set_xticklabels([])
axes1[1].set_xlabel('Time from event time (in seconds)')
# Adjust the size and spacing of the subplots
plt.gcf().set_size_inches(12, 6, forward=True)
plt.tight_layout(rect=[0, 0, 1, 0.9])
plt.subplots_adjust(wspace=0, hspace=0)
# Add a title
plt.suptitle('Sample #{}'.format(sample_id), y=0.975)
# Save the plot at the given location
plt.savefig(plot_path, bbox_inches='tight', pad_inches=0)
print('Done!')
# -------------------------------------------------------------------------
# Postliminaries
# -------------------------------------------------------------------------
# Print the total run time
print('')
print('Total runtime: {:.1f} seconds!'
.format(time.time() - script_start_time))
print('')
Supports Markdown
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