complete the traning data part, also plot some figures

python3_samples/generate_sample.py → gen_samples_py3/1-generate_sample.py

@@ -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']])

gen_samples_py3/99-loopplot.py

+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"')

python3_samples/config_files/default.json → gen_samples_py3/config_files/default.json

@@ -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"
 }

python3_samples/config_files/waveform_params.ini → gen_samples_py3/config_files/waveform_params.ini

-; -----------------------------------------------------------------------------
-; 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

gen_samples_py3/plot_sample.py

+"""
+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('')