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

complete the traning data part, also plot some figures

parent 18514b44
...@@ -24,24 +24,8 @@ from utils.waveforms import WaveformParameterGenerator ...@@ -24,24 +24,8 @@ from utils.waveforms import WaveformParameterGenerator
if __name__ == '__main__': 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() 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 # Set up the parser and add arguments
parser = argparse.ArgumentParser(description='Generate a GW data sample.') parser = argparse.ArgumentParser(description='Generate a GW data sample.')
parser.add_argument('--config-file', parser.add_argument('--config-file',
...@@ -50,9 +34,7 @@ if __name__ == '__main__': ...@@ -50,9 +34,7 @@ if __name__ == '__main__':
default='default.json') default='default.json')
# Parse the arguments that were passed when calling this script # Parse the arguments that were passed when calling this script
print('Parsing command line arguments...', end=' ')
command_line_arguments = vars(parser.parse_args()) command_line_arguments = vars(parser.parse_args())
print('Done!')
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
# Read in JSON config file specifying the sample generation process # Read in JSON config file specifying the sample generation process
...@@ -89,8 +71,6 @@ if __name__ == '__main__': ...@@ -89,8 +71,6 @@ if __name__ == '__main__':
# Define some useful shortcuts # Define some useful shortcuts
random_seed = config['random_seed'] random_seed = config['random_seed']
max_runtime = config['max_runtime']
bkg_data_dir = config['background_data_directory']
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
# Construct a generator for sampling waveform parameters # Construct a generator for sampling waveform parameters
...@@ -107,13 +87,6 @@ if __name__ == '__main__': ...@@ -107,13 +87,6 @@ if __name__ == '__main__':
waveform_parameters = \ waveform_parameters = \
(waveform_parameter_generator.draw() for _ in iter(int, 1)) (waveform_parameter_generator.draw() for _ in iter(int, 1))
# -------------------------------------------------------------------------
# Construct a generator for sampling valid noise times
# -------------------------------------------------------------------------
# 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') print('Using synthetic noise! (background_data_directory = None)\n')
# Create a iterator that returns a fake "event time", which we will # Create a iterator that returns a fake "event time", which we will
...@@ -123,31 +96,6 @@ if __name__ == '__main__': ...@@ -123,31 +96,6 @@ if __name__ == '__main__':
# None, so that we know that we need to generate synthetic noise. # None, so that we know that we need to generate synthetic noise.
noise_times = ((1000000000 + _, None) for _ in count()) 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')
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
# Define a convenience function to generate arguments for the simulation # Define a convenience function to generate arguments for the simulation
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
...@@ -170,18 +118,38 @@ if __name__ == '__main__': ...@@ -170,18 +118,38 @@ if __name__ == '__main__':
samples = dict(injection_samples=[], noise_samples=[]) samples = dict(injection_samples=[], noise_samples=[])
injection_parameters = dict(injection_samples=[], noise_samples=[]) injection_parameters = dict(injection_samples=[], noise_samples=[])
for sample_type in ['injection_samples', 'noise_samples']:
# ---------------------------------------------------------------------
# Define some sample_type-specific shortcuts
# ---------------------------------------------------------------------
if sample_type == 'injection_samples':
print('Generating samples containing an injection...') print('Generating samples containing an injection...')
n_samples = config['n_injection_samples'] n_samples = config['n_injection_samples']
arguments_generator = \ arguments_generator = \
(generate_arguments(injection=True) for _ in iter(int, 1)) (generate_arguments(injection=True) for _ in iter(int, 1))
print('Number of samples:',n_samples)
sample_type = 'injection_samples' else:
for i in range(n_samples): print('Generating samples *not* containing an injection...')
print(i) n_samples = config['n_noise_samples']
arguments_generator = \
(generate_arguments(injection=False) for _ in iter(int, 1))
# ---------------------------------------------------------------------
# If we do not need to generate any samples, skip ahead:
# ---------------------------------------------------------------------
if n_samples == 0:
print('Done! (n_samples=0)\n')
continue
results_list = [] results_list = []
# ---------------------------------------------------------------------
# Loop over injections or noise:
# ---------------------------------------------------------------------
for i in range(n_samples):
print('Generating the',i,'th samples\n')
arguments = next(arguments_generator) arguments = next(arguments_generator)
print(arguments) print(arguments)
result = generate_sample(**arguments) result = generate_sample(**arguments)
...@@ -202,7 +170,10 @@ if __name__ == '__main__': ...@@ -202,7 +170,10 @@ if __name__ == '__main__':
injection_parameters[sample_type] = \ injection_parameters[sample_type] = \
list([injection_parameters[sample_type][i] for i in idx]) 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 # Compute the normalization parameters for this file
...@@ -248,8 +219,8 @@ if __name__ == '__main__': ...@@ -248,8 +219,8 @@ if __name__ == '__main__':
static_arguments=static_arguments) static_arguments=static_arguments)
# Collect and add samples (with and without injection) # Collect and add samples (with and without injection)
for sample_type in ('injection_samples', 'noise_samples'): for sample_type in ['injection_samples', 'noise_samples']:
for key in ('event_time', 'h1_strain', 'l1_strain'): for key in ['event_time', 'h1_strain', 'l1_strain']:
if samples[sample_type]: if samples[sample_type]:
value = np.array([_[key] for _ in list(samples[sample_type])]) value = np.array([_[key] for _ in list(samples[sample_type])])
else: else:
...@@ -258,10 +229,10 @@ if __name__ == '__main__': ...@@ -258,10 +229,10 @@ if __name__ == '__main__':
# Collect and add injection_parameters (ignore noise samples here, because # Collect and add injection_parameters (ignore noise samples here, because
# for those, the injection_parameters are always None) # for those, the injection_parameters are always None)
other_keys = ['h1_signal', 'h1_snr', 'l1_signal', 'l1_snr', 'scale_factor'] #other_keys = ['h1_signal', 'h1_snr', 'l1_signal', 'l1_snr', 'scale_factor']
print(injection_parameters['injection_samples']) 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): for key in list(variable_arguments + other_keys):
#print(key)
if injection_parameters['injection_samples']: if injection_parameters['injection_samples']:
value = np.array([_[key] for _ in value = np.array([_[key] for _ in
injection_parameters['injection_samples']]) 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.
...@@ -5,8 +5,8 @@ ...@@ -5,8 +5,8 @@
"inj_bits": [0, 1, 2, 4], "inj_bits": [0, 1, 2, 4],
"waveform_params_file_name": "waveform_params.ini", "waveform_params_file_name": "waveform_params.ini",
"max_runtime": 60, "max_runtime": 60,
"n_injection_samples": 4, "n_injection_samples": 100,
"n_noise_samples": 16, "n_noise_samples": 100,
"n_processes": 4, "n_processes": 4,
"output_file_name": "default.hdf" "output_file_name": "train.hdf"
} }
; -----------------------------------------------------------------------------
; DECLARE ARGUMENTS
; -----------------------------------------------------------------------------
[variable_args] [variable_args]
; Waveform parameters that will vary in MCMC
mass1 = mass1 =
mass2 = mass2 =
spin1z = spin1z =
...@@ -15,9 +10,7 @@ inclination = ...@@ -15,9 +10,7 @@ inclination =
polarization = polarization =
injection_snr = injection_snr =
[static_args] [static_args]
; Waveform parameters that will not change in MCMC
approximant = SEOBNRv4 approximant = SEOBNRv4
domain = time domain = time
f_lower = 18 f_lower = 18
...@@ -57,60 +50,39 @@ seconds_after_event = 2.5 ...@@ -57,60 +50,39 @@ seconds_after_event = 2.5
tukey_alpha = 0.25 tukey_alpha = 0.25
; -----------------------------------------------------------------------------
; DEFINE DISTRIBUTIONS FOR PARAMETERS
; -----------------------------------------------------------------------------
[prior-mass1] [prior-mass1]
; Prior for mass1
name = uniform name = uniform
min-mass1 = 10. min-mass1 = 10.
max-mass1 = 80. max-mass1 = 80.
[prior-mass2] [prior-mass2]
; Prior for mass2
name = uniform name = uniform
min-mass2 = 10. min-mass2 = 10.
max-mass2 = 80. max-mass2 = 80.
[prior-spin1z] [prior-spin1z]
; Prior for spin1z
name = uniform name = uniform
min-spin1z = 0 min-spin1z = 0
max-spin1z = 0.998 max-spin1z = 0.998
[prior-spin2z] [prior-spin2z]
; Prior for spin2z
name = uniform name = uniform
min-spin2z = 0 min-spin2z = 0
max-spin2z = 0.998 max-spin2z = 0.998
[prior-injection_snr] [prior-injection_snr]
; Prior for the injection SNR
name = uniform name = uniform
min-injection_snr = 5 min-injection_snr = 5
max-injection_snr = 20 max-injection_snr = 20
[prior-coa_phase] [prior-coa_phase]
; Coalescence phase prior
name = uniform_angle name = uniform_angle
[prior-inclination] [prior-inclination]
; Inclination prior
name = sin_angle name = sin_angle
[prior-ra+dec] [prior-ra+dec]
; Sky position prior
name = uniform_sky name = uniform_sky
[prior-polarization] [prior-polarization]
; Polarization prior
name = uniform_angle 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