Starting expansion of PyGRB post-processing workflow generator (gwast…

…ro#4891)

* Adding 2 functions to the pygrb PP workflow generator and setting up a better landing page

* Avoid glob.glob

* Replacing the fish_label_in_filelist function with labels_in_files_metadata

bin/pygrb/pycbc_pygrb_pp_workflow

@@ -43,6 +43,59 @@ __date__ = pycbc.version.date
 __program__ = "pycbc_pygrb_pp_workflow"
 
 
+# Function that converts a list of file paths to a FileList in which each File
+# is provided with a label
+def labels_in_files_metadata(labels, file_paths):
+    """Return a FileList based on the list of file paths.  Each File in it
+    is provided with its tag from labels: the correctess of the tag is
+    based on the file name.
+    """
+    if len(labels) != len(file_paths):
+        logging.error("The two arguments must have the same length.")
+
+    # Sort the file paths according to labels
+    up_labels = [l.upper() for l in labels]
+    sorted_file_paths = [list(filter(lambda x: l in x.upper(), file_paths))[0]
+                         for l in up_labels]
+
+    # Convert sorted_file_paths to a FileList where each File has its label
+    # in the tags
+    labels_paths = zip(up_labels, sorted_file_paths)
+    out = \
+        _workflow.FileList([_workflow.resolve_url_to_file(p, attrs={'tags': [l]})
+                            for (l, p) in labels_paths])
+
+    return out
+
+
+# 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.
+    """
+    from PIL import Image, PngImagePlugin
+
+    curr_dir = os.getcwd()
+    os.chdir(output_dir)
+    segments_plot_name = 'GRB' + \
+        wflow.cp.get('workflow', 'trigger-name') + '_segments.png'
+    segments_plot = \
+        _workflow.resolve_url_to_file(os.path.join(args.segment_dir,
+                                                   segments_plot_name))
+    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 +171,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 +194,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
 #