JetsAndHF/JetValidation/MakeRunScripts.py

0001 # =============================================================================
0002 #! @file    MakeRunScripts.py
0003 #! @authors Derek Anderson (derek.murphy.anderson@protonmail.com)
0004 # -----------------------------------------------------------------------------
0005 #! @brief Helper script to generate various scripts to hand-off
0006 #!   to eic-shell. This is to get around arguments not being
0007 #!   able to be passed to scripts ran in eic-shell.
0008 #!
0009 #! @usage
0010 #!     MakeRunScripts.py -d <my rucio DID> \
0011 #!                       -c <descriptive tag for campaign> \
0012 #!                       -a <descriptive tag for dataset> \
0013 #!                       -s <number of lists to split files into> \
0014 #!                       -l <padding length for indices>
0015 # =============================================================================
0016
0017 from typing import List
0018 import argparse as ap
0019 import os
0020
0021 class ScriptMaker:
0022     def __init__(self, did: str, camp: str, data: str, splits: int, length: int):
0023         self.did    = did
0024         self.camp   = camp
0025         self.data   = data
0026         self.splits = splits
0027         self.length = length
0028         self.suffix = f"epic{self.camp}_{self.data}"
0029
0030     def _make_split_names(self, pre: str, ext: str) -> List[str]:
0031         return [f"{pre}.split_{split:0{self.length}}.{self.suffix}.{ext}" for split in range(self.splits)]
0032
0033     def make_rucio_script(self) -> None:
0034         script_name = f"run/scripts/do_rucio.{self.suffix}.sh"
0035         list_name   = f"run/lists/files.{self.suffix}.list"
0036         split_pref  = f"run/lists/files.split_"
0037         list_comm   = f"rucio replica list file --protocols root --pfns --rses isopenaccess {self.did} > {list_name}"
0038         split_comm  = f"split --suffix-length={self.length} --additional-suffix=.{self.suffix}.list -d -n l/{self.splits} {list_name} {split_pref}"
0039
0040         # generate script
0041         os.makedirs(os.path.dirname(script_name), exist_ok=True)
0042         os.makedirs(os.path.dirname(list_name), exist_ok=True)
0043         with open(script_name, 'w') as script:
0044             script.write("#!/bin/bash\n\n")
0045             script.write(f"{list_comm}\n")
0046             script.write(f"{split_comm}\n")
0047         os.chmod(script_name, 0o777)
0048
0049     def make_hist_scripts(self) -> None:
0050         drivers   = self._make_split_names("run/scripts/do_hists", "sh")
0051         in_lists  = self._make_split_names("run/lists/files", "list")
0052         out_files = self._make_split_names("run/hists/hists", "root")
0053         for split, driver in enumerate(drivers):
0054
0055             # generate command
0056             hist_comm = f'root -b -q "MakeJetValidationHists.C(\\\"{out_files[split]}\\\", \\\"{in_lists[split]}\\\", 1)"'
0057
0058             # generate script
0059             os.makedirs(os.path.dirname(out_files[split]), exist_ok=True)
0060             with open(driver, 'w') as script:
0061                 script.write("#!/bin/bash\n\n")
0062                 script.write(hist_comm)
0063             os.chmod(driver, 0o777)
0064
0065     def make_plot_scripts(self) -> None:
0066         hist_files  = " ".join(self._make_split_names("run/hists/hists", "root"))
0067         merge_file  = f"run/hists/hist.merged.{self.suffix}.root"
0068         merge_comm  = f"hadd {merge_file} {hist_files}"
0069         plot_comm   = f'root -b -q "MakeJetValidationPlots.C(\\\"run/plots\\\", \\\"{self.suffix}\\\", \\\"{merge_file}\\\")"'
0070         script_name = f"run/scripts/do_plots.{self.suffix}.sh"
0071
0072         # generate script
0073         os.makedirs("run/plots", exist_ok=True)
0074         with open(script_name, 'w') as script:
0075             script.write("#!/bin/bash\n\n")
0076             script.write(f"{merge_comm}\n")
0077             script.write(f"{plot_comm}\n")
0078         os.chmod(script_name, 0o777)
0079
0080
0081 if __name__ == "__main__":
0082
0083     parser = ap.ArgumentParser()
0084     parser.add_argument("-d", "--did")
0085     parser.add_argument("-c", "--camp")
0086     parser.add_argument("-a", "--data")
0087     parser.add_argument("-s", "--splits")
0088     parser.add_argument("-l", "--length")
0089     args = parser.parse_args()
0090
0091     maker = ScriptMaker(args.did, args.camp, args.data, int(args.splits), int(args.length))
0092     maker.make_rucio_script()
0093     maker.make_hist_scripts()
0094     maker.make_plot_scripts()