Adding injections and bank corner plots to pyGRB workflows

bin/plotting/pycbc_plot_bank_corner

@@ -124,6 +124,8 @@ parser.add_argument('--fits-file',
 parser.add_argument('--title',
     help="A title for the plot. If not given, the files supplied "
          "and the number of templates will be used")
+parser.add_argument('--no-suptitle', action='store_true',
+                    help="If true no suptitle is shown.")
 
 add_style_opt_to_parser(parser)
 args = parser.parse_args()
@@ -293,7 +295,9 @@ title_text = f"{os.path.basename(args.bank_file)}"
 if args.fits_file is not None:
     title_text += f", {os.path.basename(args.fits_file)}"
 title_text += f" - {banklen}\u00a0templates"
-fig.suptitle('\n'.join(wrap(args.title if args.title is not None else title_text, 60)))
+title_text = '\n'.join(wrap(args.title if args.title is not None else title_text, 60))
+if not args.no_suptitle:
+    fig.suptitle(title_text)
 for k, v in axis_dict.items():
     # Some may be long labels - tilt the label a little to fit together
     xlab = v[0].get_xlabel()
@@ -345,7 +349,7 @@ caption = ("Template bank as a corner plot with "
 metadata.save_fig_with_metadata(
                  fig, args.output_plot_file,
                  cmd=" ".join(sys.argv),
-                 title="Template bank corner plot",
+                 title=title_text,
                  caption=caption,
                  fig_kwds={'bbox_inches': 'tight'})
 

bin/pygrb/pycbc_make_offline_grb_workflow

@@ -490,6 +490,7 @@ if post_proc_method == "PYGRB_OFFLINE":
         logging.info('Entering results module')
         results_files = _workflow.setup_pygrb_results_workflow(wflow, pp_dir,
                                                                clustered_files,
+                                                               inj_files,
                                                                inj_find_files,
                                                                full_bank_file,
                                                                seg_dir,

bin/pygrb/pycbc_pygrb_results_workflow

@@ -143,6 +143,12 @@ parser.add_argument(
     action="store",
     help="The location of the full template bank file",
 )
+parser.add_argument(
+    "--full-inj-files",
+    action="store",
+    nargs="+",
+    help="Location(s) of the generated injections file(s)",
+)
 parser.add_argument(
     "--segment-dir",
     action="store",
@@ -219,6 +225,7 @@ if not set(inj_sets).issubset(eff_secs):
     logging.error(err_msg, inj_sets, eff_secs)
 inj_sets = [i.upper() for i in inj_sets]
 inj_files = labels_in_files_metadata(inj_sets, start_rundir, args.inj_files)
+full_inj_files = labels_in_files_metadata(inj_sets, start_rundir, args.full_inj_files)
 
 # File instance of the veto file
 veto_file = args.veto_file
@@ -273,9 +280,15 @@ plot_node, skygrid_plot = _workflow.make_pygrb_plot(
     wflow, 'pygrb_plot_skygrid', out_dir, trig_file=offsource_file
 )
 plotting_nodes.append(plot_node)
+# Plot the template bank
+plot_node, bank_plot = _workflow.make_pygrb_plot(
+    wflow, 'pycbc_plot_bank_corner', out_dir, bank_file=bank_file
+)
+
 # 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])]
+summary_layout = [(grb_info_table[0],), (seg_plot, skygrid_plot[0]),
+                  (bank_plot[0],)]
 layout.two_column_layout(out_dir, summary_layout)
 
 #
@@ -468,6 +481,25 @@ layout.two_column_layout(coinc_out_dir, coinc_files)
 #
 out_dir = rdir['injections']
 _workflow.makedir(out_dir)
+## Create corner plot
+files = _workflow.FileList([])
+for inj_set in inj_sets:
+    # Retrieve the injections full files
+    inj_file = full_inj_files.find_output_with_tag(
+        inj_set, fail_if_not_single_file=True
+    )
+    # Generate a corner bank plot for each injection set
+    tags = [inj_set]
+    plot_node, output_files = _workflow.make_pygrb_plot(
+                            wflow,
+                            'pycbc_plot_bank_corner',
+                            out_dir,
+                            inj_file=inj_file,
+                            tags=tags,
+                        )
+    plotting_nodes.append(plot_node)
+    files.append(output_files[0])
+layout.single_layout(out_dir, files)
 # Loop over injection plots requested by the user
 inj_plots = wflow.cp.get_subsections('pygrb_plot_injs_results')
 # The command above also picks up the injection set names so we remove them

pycbc/workflow/grb_utils.py

@@ -514,8 +514,14 @@ def make_pygrb_plot(workflow, exec_name, out_dir,
         node.add_opt('--trigger-time', trig_time)
     # Pass the injection file as an input File instance
     if inj_file is not None and exec_name not in \
-            ['pygrb_plot_skygrid', 'pygrb_plot_stats_distribution']:
+            ['pygrb_plot_skygrid', 'pygrb_plot_stats_distribution',
+             'pycbc_plot_bank_corner']:
         node.add_input_opt('--found-missed-file', inj_file)
+    if exec_name == 'pycbc_plot_bank_corner':
+        if inj_file is not None:
+            node.add_input_opt('--bank-file', inj_file)
+        else:
+            node.add_input_opt('--bank-file', bank_file)
     # IFO option
     if ifo:
         node.add_opt('--ifo', ifo)
@@ -540,6 +546,9 @@ def make_pygrb_plot(workflow, exec_name, out_dir,
             node.new_output_file_opt(workflow.analysis_time, '.json',
                                      '--exclusion-dist-output-file',
                                      tags=extra_tags)
+    elif exec_name == 'pycbc_plot_bank_corner':
+        node.new_output_file_opt(workflow.analysis_time, '.png',
+                                 '--output-plot-file', tags=extra_tags+tags)
     else:
         node.new_output_file_opt(workflow.analysis_time, '.png',
                                  '--output-file', tags=extra_tags)
@@ -566,6 +575,15 @@ def make_pygrb_plot(workflow, exec_name, out_dir,
             if workflow.cp.has_option_tags(exec_name, log_flag,
                                            tags=[subsection]):
                 node.add_opt('--'+log_flag)
+    elif exec_name == 'pycbc_plot_bank_corner':
+        node.add_opt('--no-suptitle')
+        if inj_file:
+            node.add_opt('--title', f'\"{tags[0]} injections\"')
+            params = workflow.cp.get_opt_tags(exec_name, 'parameters', ['injs'])
+        else:
+            node.add_opt('--title', f'\"Template bank\"')
+            params = workflow.cp.get_opt_tags(exec_name, 'parameters', ['bank'])
+        node.add_opt('--parameters', params)
 
     # Add job node to workflow
     workflow += node
@@ -759,7 +777,8 @@ def setup_pygrb_minifollowups(workflow, followups_file, trigger_file,
 
 
 def setup_pygrb_results_workflow(workflow, res_dir, trig_files,
-                                 inj_files, bank_file, seg_dir,
+                                 full_injs_files, inj_files,
+                                 bank_file, seg_dir,
                                  veto_file=None, tags=None,
                                  explicit_dependencies=None):
     """Create subworkflow to produce plots, tables,
@@ -772,6 +791,7 @@ def setup_pygrb_results_workflow(workflow, res_dir, trig_files,
     res_dir: The post-processing directory where
         results (plots, etc.) will be stored
     trig_files: FileList of trigger files
+    full_injs_files: Filelist of generated injections
     inj_files: FileList of injection results
     bank_file: The template bank File object
     seg_dir: The directory path with the segments files
@@ -796,6 +816,7 @@ def setup_pygrb_results_workflow(workflow, res_dir, trig_files,
     # Grab and pass all necessary files
     node.add_input_list_opt('--trig-files', trig_files)
     # node.add_input_opt('--config-files', config_file)
+    node.add_input_list_opt('--full-inj-files', full_injs_files)
     node.add_input_list_opt('--inj-files', inj_files)
     node.add_input_opt('--bank-file', bank_file)
     node.add_opt('--segment-dir', seg_dir)