Starting expansion of PyGRB post-processing workflow generator

bin/pygrb/pycbc_pygrb_pp_workflow

@@ -43,6 +43,48 @@ __date__ = pycbc.version.date
 __program__ = "pycbc_pygrb_pp_workflow"
 
 
+# Function used to retrive a File instance from a FileList based on an
+# injection set label
+def fish_label_in_filelist(label, filelist):
+    """Return the File in a FileList with label in its name if the label
+    appears only once in the FileList, otherwise return None.
+    """
+    out_file = None
+    if filelist is not None:
+        out_file = list(filter(lambda x: label.lower()
+                        in x.name.lower(), filelist))
+        out_file = out_file[0] if len(out_file) == 1 else None
+
+    return out_file
+
+
+# Function to retrieve the segments plot (produced in the preprocessing stage)
+# that ensures its copy is saved in the output directory before returning to
+# the original working directory. Appropriate meta data is added to the plot.
+def display_seg_plot(output_dir, segments_dir):
+    """Return the File of the segments plot (which was already produced during
+    the pre-processing) after adding the appropriate metadata to it.
+    """
+    import glob
+    from PIL import Image, PngImagePlugin
+
+    curr_dir = os.getcwd()
+    os.chdir(output_dir)
+    segments_plot = \
+        _workflow.resolve_url_to_file(glob.glob(segments_dir +
+                                                '/*_segments.png')[0])
+    segments_plot_path = segments_plot.storage_path
+    im = Image.open(segments_plot_path)
+    meta = PngImagePlugin.PngInfo()
+    meta.add_text('title', str('Segments and analysis time'))
+    meta.add_text('caption', str('Segments (horizontal bands) available around the trigger time (orange vertical line) and analysis time used (orange box).'))
+    meta.add_text('cmd', str('This plot is generated by the make_grb_segments_plot function during the pre-processing stage of the analysis.'))
+    im.save(segments_plot_path, "png", pnginfo=meta)
+    os.chdir(curr_dir)
+
+    return segments_plot
+
+
 # =============================================================================
 # Main script
 # =============================================================================
@@ -118,8 +160,10 @@ wflow.ifos = ifos
 plotting_nodes = []
 html_nodes = []
 
+# Convert the segments files to a FileList
 seg_filenames = ['bufferSeg.txt', 'offSourceSeg.txt', 'onSourceSeg.txt']
 seg_files = [os.path.join(args.segment_dir, f) for f in seg_filenames]
+seg_files = _workflow.build_segment_filelist(args.segment_dir)
 
 # Logfile of this workflow
 wf_log_file = _workflow.File(wflow.ifos, 'workflow-log', wflow.analysis_time,
@@ -139,25 +183,25 @@ logging.info("Created log file %s", wf_log_file.storage_path)
 
 out_dir = rdir.base
 _workflow.makedir(out_dir)
-files = _workflow.FileList([])
 
 #
-# Create information table about the GRB trigger and plot the search grid
+# Opening page
 #
-info_table_node, grb_info_table = _workflow.make_info_table(wflow, 'pygrb_grb_info_table', out_dir)
+# Create the information table about the GRB trigger
+info_table_node, grb_info_table = \
+    _workflow.make_pygrb_info_table(wflow, 'pygrb_grb_info_table', out_dir)
 html_nodes.append(info_table_node)
-files.append(grb_info_table)
-#
+# Plot the search grid
 plot_node, skygrid_plot = _workflow.make_pygrb_plot(wflow,
                                                     'pygrb_plot_skygrid',
                                                     out_dir,
                                                     trig_file=offsource_file)
 plotting_nodes.append(plot_node)
-files.append(skygrid_plot)
-summary_layout = [(grb_info_table[0],), (skygrid_plot[0],)]
+# Retrieve the segments plot (produced in the preprocessing stage)
+seg_plot = display_seg_plot(out_dir, args.segment_dir)
+summary_layout = [(grb_info_table[0],), (seg_plot, skygrid_plot[0])]
 layout.two_column_layout(out_dir, summary_layout)
 
-
 #
 # Plot SNR timeseries
 #