vxdcdc_alignment.py

#!/usr/bin/env python3
 

 
from glob import glob
from pathlib import Path
import argparse
 
from subprocess import run, CalledProcessError
from shutil import copy2
 
import matplotlib
import matplotlib.pyplot as plt
 
from prompt import ValidationSettings
 
import alignment_validation.dimuon as dimuonval
import alignment_validation.cosmics as cosmicval
 
import ROOT as r
r.PyConfig.IgnoreCommandLineOptions = True
r.gROOT.SetBatch()
 
matplotlib.use('Agg')
plt.style.use("belle2")
 

settings = ValidationSettings(name="Full VXD and CDC Alignment",
                              description=__doc__,
                              download_files=[],
                              expert_config=None)
 
 
def run_validation(job_path, input_data_path=None, **kwargs):
    '''job_path will be replaced with path/to/calibration_results
    input_data_path will be replaced with path/to/data_path used for calibration
    e.g. /group/belle2/dataprod/Data/PromptSkim/'''
 
    collector_output_dir_cosmic = Path(job_path) / 'VXDCDCalignment_validation/0/collector_output/cosmic/'
    collector_output_dir_mumu = Path(job_path) / 'VXDCDCalignment_validation/0/collector_output/mumu/'
 
    output_dir = Path(kwargs.get('output_dir', 'VXDCDCAlignmentValidation_output'))
    # create output directory if it does not exist
    output_dir.mkdir(parents=True, exist_ok=True)
 
    pattern_cosmic = str(collector_output_dir_cosmic) + "/*/cosmic_ana.root"
    pattern_mumu = str(collector_output_dir_mumu) + "/*/dimuon_ana.root"
 
    def hadd_and_get_merged_file(filenames_pattern, input_type):
        root_files = glob(filenames_pattern)
        merged_file = output_dir / f"{input_type}.root"
 
        if len(root_files) > 1:
            try:
                run(
                    ["hadd", "-f", merged_file, *root_files],
                    capture_output=True,
                    text=True,
                )
            except CalledProcessError as e:
                print(f"hadd failed with exit code {e.returncode}")
                print(e.stdout)
                print(e.stderr)
 
        elif len(root_files) == 1:
            copy2(root_files[0], merged_file)
        else:
            raise FileNotFoundError(f"No root files found for pattern: {filenames_pattern}")
 
        return str(merged_file)
 
    cosmic_file = hadd_and_get_merged_file(pattern_cosmic, "cosmics")
    mumu_file = hadd_and_get_merged_file(pattern_mumu, "dimuon")
 
    print(f"Merged ntuples saved in {cosmic_file} and {mumu_file}")
    print("Running validation...")
 
    cosmicval.run_validation([cosmic_file], output_dir=str(output_dir / "cosmics/"))
    dimuonval.run_validation([mumu_file], output_dir=str(output_dir / "dimuon/"))
 
    # Now merge std histograms stored in ColectorOutput.root
    histo_pattern_cosmic = str(collector_output_dir_cosmic) + "/*/CollectorOutput.root"
    histo_pattern_mumu = str(collector_output_dir_mumu) + "/*/CollectorOutput.root"
 
    histo_cosmic = hadd_and_get_merged_file(histo_pattern_cosmic, "cosmic_CollectorOutput")
    histo_mumu = hadd_and_get_merged_file(histo_pattern_mumu, "dimuon_CollectorOutput")
 
    print(f"Merged CollectorOutput histograms saved in {histo_cosmic} and {histo_mumu}")
 
    print("Alignment validation completed.")
 
 
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawTextHelpFormatter)
 
    # b2val-prompt-run wants to pass to the script also input_data_path
    # and requested_iov. As they are not required by this validation I just accept
    # them together with calibration_results_dir and then ignore them
    parser.add_argument('calibration_results_dir',
                        help='The directory that contains the collector outputs',
                        nargs='+')
 
    parser.add_argument('-o', '--output_dir',
                        help='The directory where all the output will be saved',
                        default='VXDCDCAlignmentValidation_output')
    args = parser.parse_args()
 
    run_validation(args.calibration_results_dir[0], output_dir=args.output_dir)