diff --git a/README.md b/README.md index ae221c47..fcb8745f 100644 --- a/README.md +++ b/README.md @@ -697,7 +697,7 @@ python runner.py --workflow emctag_ttdilep_sf --json metadata/test_bta_run3.json - `ci:skip syst` : remove `--isSyst all` option - `ci:JERC_split` : change systematic option to split JERC uncertainty sources `--isSyst JERC_split` - `ci:weight_only` : change systematic option to weight only variations `--isSyst weight_only` - + ### Running jupyter remotely 1. On your local machine, edit `.ssh/config`: ``` diff --git a/scripts/dohadd.py b/scripts/dohadd.py new file mode 100644 index 00000000..da42168b --- /dev/null +++ b/scripts/dohadd.py @@ -0,0 +1,25 @@ +import os, sys +from glob import glob + +indir = sys.argv[1] + +systs = os.listdir(indir) + +outfile = open("hadd.sh", "w") + +for syst in systs: + roots = glob(f"{indir}/{syst}/*/*.root") + if len(roots) == 0: + print(f"Skipping {indir}/{syst}. Not the right directory structure.") + continue + samps = os.listdir(f"{indir}/{syst}") + for samp in samps: + if len(glob(f"{indir}/{syst}/{samp}/*.root")) == 0: + continue + outfile.write( + f"hadd -v 0 {indir}/{syst}/{samp}.root {indir}/{syst}/{samp}/*.root\n" + ) + +print( + "Now run `parallel :::: hadd.sh` from an environment with ROOT installed. E.g. \nconda activate rootenv\nparallel :::: hadd.sh\nconda activate btv_coffea" +) diff --git a/scripts/dump_prescale.py b/scripts/dump_prescale.py index 43b49c13..cb7c0768 100644 --- a/scripts/dump_prescale.py +++ b/scripts/dump_prescale.py @@ -25,9 +25,7 @@ def get_prescale(HLT, lumimask, verbose=False): if not os.path.exists( f"src/BTVNanoCommissioning/data/Prescales/HLTinfo_{HLT}_run{runs[0]}_{runs[-1]}.csv" ): - for run in runs: - os.system( f"brilcalc trg --prescale --hltpath 'HLT_{HLT}*' -r {run} --output-style csv &>tmp.csv" ) diff --git a/scripts/fetch.py b/scripts/fetch.py index cb6899ce..a77ecd16 100644 --- a/scripts/fetch.py +++ b/scripts/fetch.py @@ -62,6 +62,23 @@ action="store_true", default=False, ) +parser.add_argument( + "-r", + "--redirector", + help="xrootd ridirector in case sites are not found", + choices=["infn", "fnal", "cern"], + default="infn", +) +parser.add_argument( + "-j", "--ncpus", help="Number of CPUs to use for validation", default="4" +) +parser.add_argument( + "--skipvalidation", + action="store_true", + help="If true, the readability of files will not be validated.", + default=False, +) + parser.add_argument("--campaign", help="campaign info", default=None, type=str) @@ -210,9 +227,14 @@ def getFilesFromDas(args): if xrd is None: print( - f"No SITE available in the whitelist for file {dsname}, change to global redirector" + f"No SITE available in the whitelist for file {dsname}, change to global redirector: {args.redirector}" ) - xrd = "root://xrootd-cms.infn.it//" + redirector = { + "infn": "root://xrootd-cms.infn.it//", + "fnal": "root://cmsxrootd.fnal.gov/", + "cern": "root://cms-xrd-global.cern.ch/", + } + xrd = redirector[args.redirector] if args.limit is not None: flist = flist[: args.limit] if dsname not in fdict: @@ -335,7 +357,7 @@ def remove_bad_files(sample_dict, outname, remove_bad=True): _rmap = p_map( validate, sample_dict[sample], - num_cpus=4, + num_cpus=int(args.ncpus), desc=f"Validating {sample[:20]}...", ) @@ -419,7 +441,8 @@ def main(args): empty = False assert empty, "you have empty lists" output_file = "./%s" % (args.output) - # fdict = remove_bad_files(fdict, args.output, True) # remove bad files + if not args.skipvalidation: + fdict = remove_bad_files(fdict, args.output, True) # remove bad files with open(output_file, "w") as fp: json.dump(fdict, fp, indent=4) print("The file is saved at: ", output_file) diff --git a/scripts/plotdataMC.py b/scripts/plotdataMC.py index 59ef7d3c..a616f03c 100644 --- a/scripts/plotdataMC.py +++ b/scripts/plotdataMC.py @@ -1,5 +1,5 @@ import numpy as np -import argparse, os, arrow, glob, re +import argparse, os, arrow, glob, re, sys from coffea.util import load import matplotlib.pyplot as plt from matplotlib.offsetbox import AnchoredText @@ -195,8 +195,12 @@ else: var_set = args.variable.split(",") for index, discr in enumerate(var_set): - if not isinstance(collated["mc"][discr], hist.hist.Hist): - continue + try: + if not isinstance(collated["mc"][discr], hist.hist.Hist): + continue + except: + print(f"{discr} not found. Variable must be in", collated["mc"].keys()) + sys.exit(1) ## remove empty if ( discr not in collated["mc"].keys() diff --git a/src/BTVNanoCommissioning/data/lumiMasks/Cert_Collisions2022_eraCD_355862_357900_Golden.json b/src/BTVNanoCommissioning/data/lumiMasks/Cert_Collisions2022_eraCD_355862_357900_Golden.json new file mode 100644 index 00000000..f7e30033 --- /dev/null +++ b/src/BTVNanoCommissioning/data/lumiMasks/Cert_Collisions2022_eraCD_355862_357900_Golden.json @@ -0,0 +1,178 @@ +{ + "355862": [[121, 133]], + "355863": [[1, 14]], + "355870": [[31, 67]], + "355871": [[1, 5]], + "355872": [ + [1, 738], + [758, 995], + [997, 1217] + ], + "355892": [[14, 197]], + "355912": [[43, 200]], + "355913": [[1, 106]], + "355921": [[38, 442]], + "355933": [[75, 448]], + "355942": [[24, 189], [193, 213]], + "355988": [[43, 80], [85, 90]], + "355989": [[1, 24]], + "355998": [[1, 35]], + "355999": [[1, 9]], + "356004": [[1, 19]], + "356005": [[1, 187]], + "356043": [[1, 65]], + "356071": [[37, 191]], + "356074": [[1, 26]], + "356075": [[1, 125]], + "356076": [[1, 153]], + "356077": [[1, 472]], + "356135": [[46, 71]], + "356309": [[61, 184]], + "356316": [[45, 185]], + "356322": [[1, 19]], + "356323": [[1, 67], [69, 650]], + "356371": [[41, 50], [67, 72]], + "356375": [[35, 77], [101, 125]], + "356378": [ + [8, 208], + [210, 219], + [221, 304] + ], + "356381": [[1, 1193]], + "356383": [[1, 33]], + "356385": [[1, 30]], + "356386": [[1, 122]], + "356426": [[39, 60]], + "356428": [[1, 300]], + "356433": [[1, 310]], + "356434": [[1, 13]], + "356435": [[1, 3], [8, 8]], + "356446": [[10, 623]], + "356523": [[32, 410], [412, 898]], + "356531": [[1, 56]], + "356563": [ + [36, 113], + [117, 164], + [168, 177], + [181, 191], + [193, 194], + [199, 343] + ], + "356568": [[42, 64]], + "356569": [[1, 251]], + "356570": [[1, 98]], + "356576": [[58, 240]], + "356578": [[1, 865]], + "356580": [[1, 51]], + "356582": [[7, 104]], + "356614": [[1, 10], [16, 19], [27, 62]], + "356615": [[1, 1297]], + "356619": [[1, 173]], + "356810": [[44, 163]], + "356811": [[1, 44]], + "356812": [[1, 107]], + "356813": [[1, 54]], + "356814": [ + [1, 305], + [307, 309], + [311, 366], + [368, 672] + ], + "356815": [[1, 54], [176, 219]], + "356824": [[1, 66]], + "356908": [[1, 26]], + "356919": [[29, 116]], + "356937": [[20, 138]], + "356946": [[1, 129]], + "356947": [[1, 350]], + "356948": [[1, 88]], + "356949": [[1, 94]], + "356951": [[1, 274]], + "356954": [[1, 364]], + "356955": [[1, 380]], + "356956": [[1, 109]], + "356968": [[81, 252]], + "356969": [[1, 236]], + "356970": [[1, 366]], + "356998": [[1, 5]], + "356999": [[1, 58]], + "357000": [[1, 50]], + "357001": [[1, 183]], + "357079": [[1, 22]], + "357080": [[1, 616]], + "357081": [[1, 759]], + "357101": [[54, 103]], + "357102": [[1, 13], [43, 134]], + "357104": [[1, 4]], + "357106": [[1, 60]], + "357112": [[1, 519]], + "357268": [[70, 143]], + "357271": [[1, 20], [22, 1570]], + "357328": [[44, 105]], + "357329": [[1, 668]], + "357330": [[1, 157]], + "357331": [[1, 23]], + "357332": [[1, 430]], + "357333": [[1, 207]], + "357401": [[48, 664]], + "357406": [[50, 174]], + "357438": [[35, 230]], + "357440": [[1, 354]], + "357441": [[1, 83]], + "357442": [[1, 1373]], + "357447": [[40, 50]], + "357472": [[34, 60]], + "357478": [[43, 50]], + "357479": [[1, 1046]], + "357482": [[1, 5], [21, 220]], + "357538": [[39, 63]], + "357542": [[1, 11], [13, 252]], + "357550": [[1, 36]], + "357610": [[63, 253]], + "357611": [[1, 412]], + "357612": [[1, 736]], + "357613": [[1, 256]], + "357688": [[1, 380]], + "357696": [[31, 319], [341, 410]], + "357697": [[1, 39]], + "357698": [[1, 63]], + "357699": [[1, 30]], + "357700": [[1, 757]], + "357701": [[1, 310]], + "357705": [[1, 202]], + "357706": [[1, 161]], + "357720": [[32, 92]], + "357732": [[30, 157]], + "357734": [[1, 300]], + "357735": [ + [1, 387], + [389, 821], + [823, 947], + [949, 1126] + ], + "357754": [[29, 145]], + "357756": [[1, 425]], + "357757": [[1, 9]], + "357758": [[1, 85]], + "357759": [[1, 70]], + "357766": [[10, 124]], + "357777": [[1, 85]], + "357778": [[1, 359]], + "357779": [[1, 74]], + "357781": [[1, 10]], + "357802": [[43, 206]], + "357803": [[1, 153]], + "357804": [[1, 23]], + "357805": [[1, 88]], + "357806": [[1, 56]], + "357807": [[1, 244]], + "357808": [[1, 17]], + "357809": [[1, 41]], + "357812": [[1, 50]], + "357813": [[1, 293]], + "357814": [[1, 212]], + "357815": [[1, 944]], + "357898": [[1, 313]], + "357899": [[1, 637]], + "357900": [[1, 516]] +} \ No newline at end of file diff --git a/src/BTVNanoCommissioning/data/lumiMasks/Cert_Collisions2022_eraEFG_359022_362760_Golden.json b/src/BTVNanoCommissioning/data/lumiMasks/Cert_Collisions2022_eraEFG_359022_362760_Golden.json new file mode 100644 index 00000000..6a14fd64 --- /dev/null +++ b/src/BTVNanoCommissioning/data/lumiMasks/Cert_Collisions2022_eraEFG_359022_362760_Golden.json @@ -0,0 +1,258 @@ +{ + "359569": [[1, 83]], + "359571": [ + [1, 4], + [7, 28], + [30, 39], + [41, 50], + [52, 61], + [63, 72], + [74, 83], + [85, 94], + [96, 105], + [107, 116], + [118, 127], + [129, 138], + [140, 149], + [151, 160], + [162, 169] + ], + "359575": [ + [1, 13], + [15, 24], + [26, 35], + [37, 46], + [48, 57], + [59, 68], + [70, 89] + ], + "359595": [[17, 289]], + "359597": [[1, 94]], + "359602": [[1, 74]], + "359609": [[1, 25]], + "359612": [[1, 12]], + "359661": [[42, 233]], + "359685": [[1, 132]], + "359686": [[1, 934]], + "359691": [[19, 139]], + "359693": [[1, 604]], + "359694": [[1, 1080]], + "359699": [[28, 1883]], + "359718": [[1, 199], [202, 368]], + "359750": [[52, 186]], + "359751": [[15, 1007]], + "359762": [[44, 266]], + "359763": [[1, 410], [412, 459]], + "359764": [[28, 1111]], + "359776": [[104, 432]], + "359806": [[42, 410]], + "359808": [[1, 46]], + "359809": [[1, 43]], + "359810": [[1, 183]], + "359812": [[1, 227]], + "359814": [[1, 366]], + "359870": [[2, 7]], + "359871": [[1, 228]], + "359887": [[30, 502]], + "359899": [[19, 305]], + "359903": [[1, 12]], + "359998": [ + [24, 757], + [759, 759], + [761, 785], + [787, 1918] + ], + "360017": [[42, 81]], + "360019": [[1, 153], [155, 1489]], + "360075": [[20, 429], [462, 809]], + "360090": [[1, 794]], + "360116": [[77, 432]], + "360125": [ + [132, 195], + [200, 262], + [264, 324] + ], + "360126": [[1, 35], [37, 370]], + "360127": [[1, 201]], + "360128": [[1, 163]], + "360130": [[87, 91]], + "360131": [[1, 378]], + "360141": [[25, 775]], + "360224": [[43, 259]], + "360225": [[1, 169], [171, 356]], + "360295": [[35, 771], [774, 1475]], + "360296": [[1, 519]], + "360327": [[32, 383]], + "360390": [[5, 36], [38, 76]], + "360391": [[1, 4]], + "360392": [[1, 61]], + "360393": [[1, 124]], + "360400": [[40, 119]], + "360413": [[37, 126]], + "360428": [[50, 123]], + "360432": [[1, 28]], + "360435": [[1, 37], [43, 53]], + "360437": [[1, 40]], + "360438": [[1, 6]], + "360458": [[45, 215], [234, 246]], + "360459": [[1, 1332]], + "360460": [[1, 1077]], + "360486": [[45, 322]], + "360490": [[8, 936], [940, 1187]], + "360491": [[3, 157], [159, 179]], + "360737": [[10, 27]], + "360761": [[34, 190]], + "360794": [[37, 713]], + "360795": [[1, 395]], + "360797": [[1, 5]], + "360820": [[1, 885]], + "360825": [[1, 358]], + "360826": [[1, 91], [93, 119]], + "360856": [[42, 97]], + "360874": [[31, 82]], + "360876": [[1, 341]], + "360887": [[12, 134]], + "360888": [[1, 289]], + "360889": [[1, 36]], + "360890": [[1, 216], [219, 506]], + "360892": [[1, 296]], + "360895": [[1, 700]], + "360919": [[39, 836]], + "360920": [[1, 4]], + "360921": [[1, 218]], + "360927": [ + [37, 331], + [333, 581], + [583, 1930] + ], + "360941": [[1, 72]], + "360942": [[1, 344]], + "360945": [[1, 91]], + "360946": [[1, 364]], + "360948": [[1, 119]], + "360950": [ + [1, 179], + [181, 468], + [470, 718] + ], + "360951": [[1, 190]], + "360991": [[41, 270]], + "360992": [[1, 99]], + "361020": [[41, 439]], + "361044": [[42, 166]], + "361045": [[1, 322], [325, 947]], + "361050": [[1, 10]], + "361052": [[1, 65]], + "361054": [[1, 461]], + "361083": [[6, 59]], + "361091": [[38, 418]], + "361105": [[41, 448]], + "361106": [[1, 31]], + "361107": [[1, 156]], + "361108": [[1, 7]], + "361110": [[1, 197]], + "361188": [[46, 186]], + "361193": [[1, 29]], + "361195": [[1, 20]], + "361197": [[1, 960], [991, 2177]], + "361223": [[36, 730]], + "361239": [[35, 1090]], + "361240": [[1, 1015]], + "361272": [[1, 125]], + "361280": [[1, 267]], + "361283": [[1, 26]], + "361284": [[1, 36]], + "361297": [[38, 970]], + "361303": [[36, 2275]], + "361318": [[34, 490]], + "361319": [[1, 162]], + "361320": [[1, 382]], + "361333": [[36, 337]], + "361361": [[60, 105]], + "361362": [[1, 31]], + "361363": [[1, 27]], + "361365": [[1, 489]], + "361366": [[1, 201]], + "361400": [[41, 496]], + "361417": [[38, 1659]], + "361443": [[40, 2313]], + "361468": [[3, 1883]], + "361475": [ + [37, 614], + [617, 836], + [838, 838], + [840, 1121], + [1124, 1273] + ], + "361512": [[36, 849], [851, 1573]], + "361569": [[35, 598]], + "361573": [[1, 305]], + "361579": [[43, 905]], + "361580": [[1, 422]], + "361957": [[1, 1163]], + "361971": [[38, 731], [733, 2381]], + "361989": [[37, 112], [114, 151]], + "361990": [[1, 204]], + "361994": [[1, 116]], + "362058": [[39, 209]], + "362059": [[1, 63]], + "362060": [[1, 83], [85, 85]], + "362061": [[1, 314]], + "362062": [[1, 207]], + "362063": [[1, 70], [77, 92]], + "362064": [ + [93, 218], + [220, 508], + [510, 1067] + ], + "362087": [[1, 13], [15, 146]], + "362091": [[1, 1418]], + "362104": [[33, 147]], + "362105": [[1, 179]], + "362106": [[1, 78]], + "362107": [[1, 785]], + "362148": [[39, 758]], + "362153": [[1, 140]], + "362154": [[1, 784]], + "362159": [[29, 35]], + "362161": [[1, 120]], + "362163": [[1, 65]], + "362166": [[1, 331]], + "362167": [[1, 646]], + "362437": [[1, 11], [73, 92]], + "362439": [[8, 34], [249, 1595]], + "362597": [[1, 139]], + "362614": [[1, 153]], + "362615": [[1, 149]], + "362616": [[1, 351]], + "362617": [[1, 244]], + "362618": [[1, 274]], + "362653": [[40, 305]], + "362654": [[1, 26]], + "362655": [ + [1, 140], + [143, 144], + [148, 338] + ], + "362657": [[1, 198]], + "362695": [[39, 372]], + "362696": [ + [1, 1112], + [1114, 1120], + [1122, 1124], + [1126, 1128], + [1130, 1132], + [1134, 1140], + [1142, 1144], + [1146, 1148], + [1150, 1152], + [1154, 1720] + ], + "362698": [[1, 374]], + "362720": [[52, 949]], + "362727": [[4, 7]], + "362728": [[1, 388]], + "362757": [[32, 93]], + "362758": [[1, 205]], + "362760": [[1, 728]] +} diff --git a/src/BTVNanoCommissioning/helpers/definitions.py b/src/BTVNanoCommissioning/helpers/definitions.py index a0f6e93d..f55074fa 100644 --- a/src/BTVNanoCommissioning/helpers/definitions.py +++ b/src/BTVNanoCommissioning/helpers/definitions.py @@ -6216,7 +6216,6 @@ def SV_definitions(): def axes_name(var): - output = {} unit = "" diff --git a/src/BTVNanoCommissioning/helpers/func.py b/src/BTVNanoCommissioning/helpers/func.py index 509abdb6..9041afe6 100644 --- a/src/BTVNanoCommissioning/helpers/func.py +++ b/src/BTVNanoCommissioning/helpers/func.py @@ -2,6 +2,7 @@ import numpy as np from coffea import processor import psutil, os +import uproot def memory_usage_psutil(): @@ -103,9 +104,12 @@ def uproot_writeable(events, include=["events", "run", "luminosityBlock"]): and "Flavor" not in n ): continue - if not _is_rootcompat(events[bname][n]) and ak.num( - events[bname][n], axis=0 - ) != len(flatten(events[bname][n])): + evnums = ak.num(events[bname][n], axis=0) + if not isinstance(evnums, int): + continue + if not _is_rootcompat(events[bname][n]) and evnums != len( + flatten(events[bname][n]) + ): continue # skip IdxG if "IdxG" in n: diff --git a/src/BTVNanoCommissioning/helpers/update_branch.py b/src/BTVNanoCommissioning/helpers/update_branch.py index 7702fcbf..9476a52f 100644 --- a/src/BTVNanoCommissioning/helpers/update_branch.py +++ b/src/BTVNanoCommissioning/helpers/update_branch.py @@ -79,7 +79,9 @@ def missing_branch(events): ) if hasattr(events, "METFixEE2017"): events.MET = events.METFixEE2017 - if not hasattr(events.PuppiMET, "MetUnclustEnUpDeltaX"): + if hasattr(events.PuppiMET, "ptUnclusteredUp") and not hasattr( + events.PuppiMET, "MetUnclustEnUpDeltaX" + ): met = events.PuppiMET met["MetUnclustEnUpDeltaX"] = met.ptUnclusteredUp * np.cos(met.phiUnclusteredUp) met["MetUnclustEnUpDeltaY"] = met.ptUnclusteredUp * np.sin(met.phiUnclusteredUp) diff --git a/src/BTVNanoCommissioning/utils/AK4_parameters.py b/src/BTVNanoCommissioning/utils/AK4_parameters.py index 301c3973..156824e4 100644 --- a/src/BTVNanoCommissioning/utils/AK4_parameters.py +++ b/src/BTVNanoCommissioning/utils/AK4_parameters.py @@ -99,7 +99,7 @@ # "Run2023C-22Sep2023_v4": "Summer23Prompt23_RunCv4_V1", # }, "JME": "jec_compiled.pkl.gz", - # "jetveto": {"Run2023BC jetvetomap_all": "Summer23Prompt23_RunC_v1.histo.root"}, + "jetveto": {"Run2023BC jetvetomap_all": "Summer23Prompt23_RunC_v1.histo.root"}, "JPCalib": { "Run2023C-22Sep2023_v1": "calibeHistoWrite_Data2023C-22Sep2023_v1.root", "Run2023C-22Sep2023_v2": "calibeHistoWrite_Data2023C-22Sep2023_v2.root", @@ -111,14 +111,14 @@ "Summer23BPix": { "lumiMask": "Cert_Collisions2023_366442_370790_Golden.json", "PU": "puwei_Summer23BPix.histo.root", - # "JME": { - # "MC": "Summer23BPixPrompt23_V1", - # "Run2023D": "Summer23BPixPrompt23_RunD_V1", - # }, - "JME": "jec_compiled.pkl.gz", - # "jetveto": { - # "Run2023D jetvetomap_all": "Summer23BPixPrompt23_RunD_v1.histo.root" - # }, # this is from Mikko https://indico.cern.ch/event/1315421/contributions/5532963/attachments/2697975/4683826/2023_08_16_jetvetomaps_v3.pdf + "JME": { + "MC": "Summer23BPixPrompt23_V1", + "Run2023D": "Summer23BPixPrompt23_RunD_V1", + }, + # "JME": "jec_compiled.pkl.gz", + "jetveto": { + "Run2023D jetvetomap_all": "Summer23BPixPrompt23_RunD_v1.histo.root" + }, # this is from Mikko https://indico.cern.ch/event/1315421/contributions/5532963/attachments/2697975/4683826/2023_08_16_jetvetomaps_v3.pdf "JPCalib": { "Run2023D-22Sep2023_v1": "calibeHistoWrite_Data2023D-22Sep2023_v1.root", "Run2023D-22Sep2023_v2": "calibeHistoWrite_Data2023D-22Sep2023_v2.root", diff --git a/src/BTVNanoCommissioning/utils/array_writer.py b/src/BTVNanoCommissioning/utils/array_writer.py new file mode 100644 index 00000000..144db85b --- /dev/null +++ b/src/BTVNanoCommissioning/utils/array_writer.py @@ -0,0 +1,95 @@ +from BTVNanoCommissioning.helpers.func import uproot_writeable +import numpy as np +import awkward as ak +import os, uproot + + +def array_writer( + processor_class, # the NanoProcessor class ("self") + pruned_event, # the event with specific calculated variables stored + nano_event, # entire NanoAOD/PFNano event with many variables + systname, # name of systematic shift + dataset, # dataset name + isRealData, # boolean + remove=[ + "SoftMuon", + "MuonJet", + "dilep", + "OtherJets", + "Jet", + ], # remove from variable list + kinOnly=[ + "Muon", + "Jet", + "SoftMuon", + "dilep", + "charge", + "MET", + ], # variables for which only kinematic properties are kept + kins=[ + "pt", + "eta", + "phi", + "mass", + "pfRelIso04_all", + "pfRelIso03_all", + "dxy", + "dz", + ], # kinematic propoerties for the above variables + othersData=[ + "PFCands_*", + "MuonJet_*", + "SV_*", + "PV_npvs", + "PV_npvsGood", + "Rho_*", + "SoftMuon_dxySig", + "Muon_sip3d", + ], # other fields, for Data and MC + othersMC=["Pileup_nTrueInt", "Pileup_nPU"], # other fields, for MC only + empty=False, +): + if empty: + print("WARNING: No events selected. Writing blank file.") + out_branch = [] + else: + # Get only the variables that were added newly + out_branch = np.setdiff1d( + np.array(pruned_event.fields), np.array(nano_event.fields) + ) + + # Handle kinOnly vars + for v in remove: + out_branch = np.delete(out_branch, np.where((out_branch == v))) + + for kin in kins: + for obj in kinOnly: + if "MET" in obj and ("pt" != kin or "phi" != kin): + continue + if (obj != "Muon" and obj != "SoftMuon") and ( + "pfRelIso04_all" == kin or "d" in kin + ): + continue + out_branch = np.append(out_branch, [f"{obj}_{kin}"]) + + # Handle data vars + out_branch = np.append(out_branch, othersData) + + if not isRealData: + out_branch = np.append(out_branch, othersMC) + + # Write to root files + # print("Branches to write:", out_branch) + outdir = f"{processor_class.name}/{systname}/{dataset}/" + os.system(f"mkdir -p {outdir}") + + with uproot.recreate( + f"{outdir}/{nano_event.metadata['filename'].split('/')[-1].replace('.root','')}_{int(nano_event.metadata['entrystop']/processor_class.chunksize)}.root" + ) as fout: + if not empty: + fout["Events"] = uproot_writeable(pruned_event, include=out_branch) + fout["TotalEventCount"] = ak.Array( + [nano_event.metadata["entrystop"] - nano_event.metadata["entrystart"]] + ) + if not isRealData: + fout["TotalEventWeight"] = ak.Array([ak.sum(nano_event.genWeight)]) diff --git a/src/BTVNanoCommissioning/utils/compile_jec.py b/src/BTVNanoCommissioning/utils/compile_jec.py index 662a4654..8f14aa42 100644 --- a/src/BTVNanoCommissioning/utils/compile_jec.py +++ b/src/BTVNanoCommissioning/utils/compile_jec.py @@ -282,7 +282,6 @@ def jet_factory_factory(files): "src/BTVNanoCommissioning/data/JME/Summer23/Summer23Prompt23_V1_MC_Uncertainty_AK4PFPuppi.junc.txt", "src/BTVNanoCommissioning/data/JME/Summer23/Summer23Prompt23_JRV1_MC_SF_AK4PFPuppi.jersf.txt", "src/BTVNanoCommissioning/data/JME/Summer23/Summer23Prompt23_JRV1_MC_PtResolution_AK4PFPuppi.jr.txt", - ], "dataCv123": [ "src/BTVNanoCommissioning/data/JME/Summer23/Summer23Prompt23_RunCv123_V1_DATA_L1FastJet_AK4PFPuppi.jec.txt", @@ -308,6 +307,14 @@ def jet_factory_factory(files): "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_JRV1_MC_PtResolution_AK4PFPuppi.jr.txt", "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_JRV1_MC_SF_AK4PFPuppi.jersf.txt", ], + # "dataD":[ + # "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_RunD_V1_DATA_L1FastJet_AK4PFPuppi.jec.txt", + # "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_RunD_V1_DATA_L2Relative_AK4PFPuppi.jec.txt", + # "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_RunD_V1_DATA_L2Residual_AK4PFPuppi.jec.txt", + # "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_RunD_V1_DATA_L3Absolute_AK4PFPuppi.jec.txt", + # "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_RunD_V1_DATA_Uncertainty_AK4PFPuppi.junc.txt", + # "src/BTVNanoCommissioning/data/JME/Summer23BPix/Summer23BPixPrompt23_RunD_V1_DATA_UncertaintySources_AK4PFPuppi.junc.txt" + # ] }, } diff --git a/src/BTVNanoCommissioning/utils/histogrammer.py b/src/BTVNanoCommissioning/utils/histogrammer.py index 0d204cd0..34d3cd06 100644 --- a/src/BTVNanoCommissioning/utils/histogrammer.py +++ b/src/BTVNanoCommissioning/utils/histogrammer.py @@ -181,9 +181,10 @@ def histogrammer(events, workflow): syst_axis, flav_axis, dz_axis, Hist.storage.Weight() ) else: - _hist_dict[f"{i}_pfRelIso04_all"] = Hist.Hist( - syst_axis, iso_axis, Hist.storage.Weight() - ) + if "m" in workflow: + _hist_dict[f"{i}_pfRelIso04_all"] = Hist.Hist( + syst_axis, iso_axis, Hist.storage.Weight() + ) _hist_dict[f"{i}_dxy"] = Hist.Hist( syst_axis, qcddxy_axis, Hist.storage.Weight() ) @@ -309,9 +310,10 @@ def histogrammer(events, workflow): ) _hist_dict["dr_mumu"] = Hist.Hist(syst_axis, dr_axis, Hist.storage.Weight()) for i in ["posl", "negl"]: - _hist_dict[f"{i}_pfRelIso04_all"] = Hist.Hist( - syst_axis, iso_axis, Hist.storage.Weight() - ) + if "m" in workflow: + _hist_dict[f"{i}_pfRelIso04_all"] = Hist.Hist( + syst_axis, iso_axis, Hist.storage.Weight() + ) _hist_dict[f"{i}_dxy"] = Hist.Hist( syst_axis, dxy_axis, Hist.storage.Weight() ) diff --git a/src/BTVNanoCommissioning/utils/selection.py b/src/BTVNanoCommissioning/utils/selection.py index cc45cc1f..dc4e36c5 100644 --- a/src/BTVNanoCommissioning/utils/selection.py +++ b/src/BTVNanoCommissioning/utils/selection.py @@ -37,12 +37,13 @@ def ele_mvatightid(events, campaign): return elemask -def softmu_mask(events, campaign): +def softmu_mask(events, campaign, dxySigCut=0): softmumask = ( (events.Muon.pt < 25) & (abs(events.Muon.eta) < 2.4) & (events.Muon.tightId > 0.5) & (events.Muon.pfRelIso04_all > 0.2) + & (abs(events.Muon.dxy / events.Muon.dxyErr) > dxySigCut) & (events.Muon.jetIdx != -1) ) diff --git a/src/BTVNanoCommissioning/workflows/BTA_producer.py b/src/BTVNanoCommissioning/workflows/BTA_producer.py index 987f13da..f1a3b0d7 100644 --- a/src/BTVNanoCommissioning/workflows/BTA_producer.py +++ b/src/BTVNanoCommissioning/workflows/BTA_producer.py @@ -63,6 +63,8 @@ def process(self, events): dirname += "_addAllTracks" if self.addPFMuons: dirname += "_addPFMuons" + if self.isSyst: + fname = "systematic/" + fname checkf = os.popen( f"gfal-ls root://eoscms.cern.ch//eos/cms/store/group/phys_btag/milee/{dirname}/{self._campaign.replace('Run3','')}/systematic/{fname}" ).read() @@ -104,8 +106,6 @@ def process_shift(self, events, shift_name): events = events[events.HLT.PFJet80] if len(events) == 0: return {dataset: len(events)} - # if "JME" in self.SF_map.keys() or "jetveto" in self.SF_map.keys(): - # events.Jet = update(events.Jet, {"veto": jetveto(events, self.SF_map)}) # basic variables basic_vars = { "Run": events.run, @@ -1115,6 +1115,8 @@ def process_shift(self, events, shift_name): output["GenV0"] = GenV0 os.system(f"mkdir -p {dataset}") fname = f"{dataset}/{events.metadata['filename'].split('/')[-1].replace('.root','')}_{int(events.metadata['entrystop']/self.chunksize)}.root" + if self.isSyst: + fname = "systematic/" + fname dirname = "BTA" if self.addAllTracks: dirname += "_addAllTracks" diff --git a/src/BTVNanoCommissioning/workflows/QCD_validation.py b/src/BTVNanoCommissioning/workflows/QCD_validation.py index 58f9cffe..3ca714e6 100644 --- a/src/BTVNanoCommissioning/workflows/QCD_validation.py +++ b/src/BTVNanoCommissioning/workflows/QCD_validation.py @@ -4,6 +4,7 @@ from BTVNanoCommissioning.utils.selection import jet_cut from BTVNanoCommissioning.helpers.func import flatten, update, dump_lumi from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.helpers.update_branch import missing_branch from BTVNanoCommissioning.utils.correction import ( load_SF, diff --git a/src/BTVNanoCommissioning/workflows/__init__.py b/src/BTVNanoCommissioning/workflows/__init__.py index 33ccaab8..51f8208f 100644 --- a/src/BTVNanoCommissioning/workflows/__init__.py +++ b/src/BTVNanoCommissioning/workflows/__init__.py @@ -20,41 +20,30 @@ from BTVNanoCommissioning.workflows.ctag_dileptt_valid_sf import ( NanoProcessor as CTAGDilepttValidSFProcessor, ) -from BTVNanoCommissioning.workflows.ctag_eDY_valid_sf import ( - NanoProcessor as CTAGeDYValidSFProcessor, -) -from BTVNanoCommissioning.workflows.ctag_eWc_valid_sf import ( - NanoProcessor as CTAGeWcValidSFProcessor, -) -from BTVNanoCommissioning.workflows.ctag_Wc_valid_sf import ( - NanoProcessor as CTAGWcValidSFProcessor, +from BTVNanoCommissioning.workflows.ctag_Wctt_valid_sf import ( + NanoProcessor as CTAGWcTTValidSFProcessor, ) from BTVNanoCommissioning.workflows.ctag_DY_valid_sf import ( NanoProcessor as CTAGDYValidSFProcessor, ) -from BTVNanoCommissioning.workflows.ctag_edileptt_valid_sf import ( - NanoProcessor as CTAGEDilepttValidSFProcessor, -) -from BTVNanoCommissioning.workflows.ctag_ettsemilep_valid_sf import ( - NanoProcessor as CTAGETTSemilepValidSFProcessor, -) -from BTVNanoCommissioning.workflows.ctag_semileptt_valid_sf import ( - NanoProcessor as CTAGSemilepttValidSFProcessor, + +##QCD +from BTVNanoCommissioning.workflows.QCD_validation import ( + NanoProcessor as QCDValidProcessor, ) + +## BTA - for SFs from BTVNanoCommissioning.workflows.BTA_producer import ( NanoProcessor as BTA_processor, ) from BTVNanoCommissioning.workflows.BTA_ttbar_producer import ( NanoProcessor as BTA_ttbar_processor, -) +) # ttbar -kinFit # from BTVNanoCommissioning.workflows.example import ( # NanoProcessor as ExampleProcessor, # ) -##QCD -from BTVNanoCommissioning.workflows.QCD_validation import ( - NanoProcessor as QCDValidProcessor, -) + # FIXME - make names more systematic? workflows = {} @@ -63,22 +52,31 @@ # TTBar workflows["ttdilep_sf"] = TTdilepValidSFProcessor workflows["ttsemilep_sf"] = TTsemilepValidSFProcessor + workflows["emctag_ttdilep_sf"] = CTAGEMDilepttValidSFProcessor -workflows["ctag_ttdilep_sf"] = CTAGDilepttValidSFProcessor -workflows["ectag_ttdilep_sf"] = CTAGEDilepttValidSFProcessor -workflows["ctag_ttsemilep_sf"] = CTAGSemilepttValidSFProcessor -workflows["ectag_ttsemilep_sf"] = CTAGETTSemilepValidSFProcessor +workflows["ctag_ttdilep_sf"] = partial( + CTAGDilepttValidSFProcessor, selectionModifier="dilepttM" +) +workflows["ectag_ttdilep_sf"] = partial( + CTAGDilepttValidSFProcessor, selectionModifier="dilepttE" +) +workflows["ctag_ttsemilep_sf"] = partial( + CTAGWcTTValidSFProcessor, selectionModifier="semittM" +) +workflows["ectag_ttsemilep_sf"] = partial( + CTAGWcTTValidSFProcessor, selectionModifier="semittE" +) ##QCD workflows["QCD_sf"] = QCDValidProcessor # W+c -workflows["ctag_Wc_sf"] = CTAGWcValidSFProcessor -workflows["ectag_Wc_sf"] = CTAGeWcValidSFProcessor +workflows["ctag_Wc_sf"] = partial(CTAGWcTTValidSFProcessor, selectionModifier="WcM") +workflows["ectag_Wc_sf"] = partial(CTAGWcTTValidSFProcessor, selectionModifier="WcE") # DY -workflows["ctag_DY_sf"] = CTAGDYValidSFProcessor -workflows["ectag_DY_sf"] = CTAGeDYValidSFProcessor +workflows["ctag_DY_sf"] = partial(CTAGDYValidSFProcessor, selectionModifier="DYM") +workflows["ectag_DY_sf"] = partial(CTAGDYValidSFProcessor, selectionModifier="DYE") # Tutorial # workflows["example"] = ExampleProcessor diff --git a/src/BTVNanoCommissioning/workflows/ctag_DY_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_DY_valid_sf.py index 9daeb6e1..2bd2b791 100644 --- a/src/BTVNanoCommissioning/workflows/ctag_DY_valid_sf.py +++ b/src/BTVNanoCommissioning/workflows/ctag_DY_valid_sf.py @@ -8,6 +8,7 @@ load_lumi, load_SF, muSFs, + eleSFs, puwei, btagSFs, JME_shifts, @@ -17,12 +18,11 @@ from BTVNanoCommissioning.helpers.func import ( flatten, update, - uproot_writeable, dump_lumi, ) from BTVNanoCommissioning.helpers.update_branch import missing_branch - from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import ( jet_id, mu_idiso, @@ -41,6 +41,7 @@ def __init__( isArray=False, noHist=False, chunksize=75000, + selectionModifier="DYM", ): self._year = year self._campaign = campaign @@ -50,6 +51,7 @@ def __init__( self.noHist = noHist self.lumiMask = load_lumi(self._campaign) self.chunksize = chunksize + self.selMod = selectionModifier ## Load corrections self.SF_map = load_SF(self._campaign) @@ -98,19 +100,31 @@ def process(self, events): def process_shift(self, events, shift_name): dataset = events.metadata["dataset"] isRealData = not hasattr(events, "genWeight") + + isMu = False + isEle = False + if "DYM" in self.selMod: + triggers = ["Mu17_TrkIsoVVL_Mu8_TrkIsoVVL_DZ_Mass8"] + isMu = True + elif "DYE" in self.selMod: + triggers = ["Ele23_Ele12_CaloIdL_TrackIdL_IsoVL"] + isEle = True + else: + raise ValueError(self.selMod, "is not a valid selection modifier.") + + histname = {"DYM": "ctag_DY_sf", "DYE": "ectag_DY_sf"} _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ctag_DY_sf") + {"": None} if self.noHist else histogrammer(events, histname[self.selMod]) ) output = { "sumw": processor.defaultdict_accumulator(float), **_hist_event_dict, } - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) + if isRealData: + output["sumw"] = len(events) + else: + output["sumw"] = ak.sum(events.genWeight) #################### # Selections # @@ -124,7 +138,6 @@ def process_shift(self, events, shift_name): output = dump_lumi(events[req_lumi], output) ## HLT - triggers = ["Mu17_TrkIsoVVL_Mu8_TrkIsoVVL_DZ_Mass8"] checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) if ak.all(checkHLT == False): raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) @@ -145,19 +158,46 @@ def process_shift(self, events, shift_name): dilep_ele = events.Electron[ (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) ] + if isMu: + thisdilep = dilep_mu + otherdilep = dilep_ele + else: + thisdilep = dilep_ele + otherdilep = dilep_mu ## dilepton - pos_dilep = dilep_mu[dilep_mu.charge > 0] - neg_dilep = dilep_mu[dilep_mu.charge < 0] + pos_dilep = thisdilep[thisdilep.charge > 0] + neg_dilep = thisdilep[thisdilep.charge < 0] req_dilep = ak.fill_none( ( (ak.num(pos_dilep.pt) >= 1) & (ak.num(neg_dilep.pt) >= 1) - & (ak.num(dilep_mu.charge) >= 2) - & (ak.num(dilep_ele.charge) == 0) + & (ak.num(thisdilep.charge) >= 2) + & (ak.num(otherdilep.charge) == 0) + ), + False, + axis=-1, + ) + + jet_sel = ak.fill_none( + jet_id(events, self._campaign) + & ( + ak.all( + events.Jet.metric_table(pos_dilep) > 0.4, + axis=2, + mask_identity=True, + ) + ) + & ( + ak.all( + events.Jet.metric_table(neg_dilep) > 0.4, + axis=2, + mask_identity=True, + ) ), False, axis=-1, ) + pos_dilep = ak.pad_none(pos_dilep, 1, axis=1) neg_dilep = ak.pad_none(neg_dilep, 1, axis=1) @@ -188,30 +228,13 @@ def process_shift(self, events, shift_name): axis=-1, ) ] - req_jets = ak.num(event_jet.pt) >= 1 - event_jet = ak.pad_none(event_jet, 1, axis=1) + req_jets = ak.count(event_jet.pt, axis=1) >= 1 + # event_jet = ak.pad_none(event_jet, 1, axis=1) ## store jet index for PFCands, create mask on the jet index jetindx = ak.mask( ak.local_index(events.Jet.pt), - ( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(pos_dilep[:, 0]) > 0.4, - axis=2, - mask_identity=True, - ) - ) - & ( - ak.all( - events.Jet.metric_table(neg_dilep[:, 0]) > 0.4, - axis=2, - mask_identity=True, - ) - ) - ) - == 1, + jet_sel == 1, ) jetindx = ak.pad_none(jetindx, 1) jetindx = jetindx[:, 0] @@ -221,6 +244,16 @@ def process_shift(self, events, shift_name): False, ) if len(events[event_level]) == 0: + if self.isArray: + array_writer( + self, + events[event_level], + events, + "nominal", + dataset, + isRealData, + empty=True, + ) return {dataset: output} #################### # Selected objects # @@ -268,8 +301,10 @@ def process_shift(self, events, shift_name): weights, syst_wei, ) - if "MUO" in self.SF_map.keys(): + if isMu and "MUO" in self.SF_map.keys(): muSFs(smu, self.SF_map, weights, syst_wei, False) + if isEle and "EGM" in self.SF_map.keys(): + eleSFs(smu, self.SF_map, weights, syst_wei, False) if "BTV" in self.SF_map.keys(): btagSFs(sel_jet, self.SF_map, weights, "DeepJetC", syst_wei) btagSFs(sel_jet, self.SF_map, weights, "DeepJetB", syst_wei) @@ -406,13 +441,22 @@ def process_shift(self, events, shift_name): if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] - pruned_ev.Jet = sel_jet - pruned_ev.Muon = smu + pruned_ev["SelJet"] = sjets + pruned_ev["LeadJet"] = sel_jet + if isMu: + pruned_ev["MuonPlus"] = sposmu + pruned_ev["MuonMinus"] = snegmu + kinOnly = ["Muon", "MuonPlus", "MuonMinus", "Jet"] + else: + pruned_ev["ElectronPlus"] = sposmu + pruned_ev["ElectronMinus"] = snegmu + kinOnly = ["Electron", "ElectronPlus", "ElectronMinus", "Jet"] pruned_ev["dilep"] = sposmu + snegmu pruned_ev["dilep", "pt"] = pruned_ev.dilep.pt pruned_ev["dilep", "eta"] = pruned_ev.dilep.eta pruned_ev["dilep", "phi"] = pruned_ev.dilep.phi pruned_ev["dilep", "mass"] = pruned_ev.dilep.mass + pruned_ev["njet"] = njet if "PFCands" in events.fields: pruned_ev.PFCands = spfcands # Add custom variables @@ -426,29 +470,17 @@ def process_shift(self, events, shift_name): pruned_ev["dr_mu1jet"] = sposmu.delta_r(sel_jet) pruned_ev["dr_mu2jet"] = snegmu.delta_r(sel_jet) - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where((out_branch == "dilep")), + array_writer( + self, + pruned_ev, + events, + systematics[0], + dataset, + isRealData, + kinOnly=kinOnly, + remove=kinOnly, ) - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in ["Muon", "Jet", "dilep"]: - if (obj != "Muon") and ("pfRelIso04_all" == kin or "d" in kin): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) return {dataset: output} def postprocess(self, accumulator): diff --git a/src/BTVNanoCommissioning/workflows/ctag_Wc_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_Wctt_valid_sf.py similarity index 71% rename from src/BTVNanoCommissioning/workflows/ctag_Wc_valid_sf.py rename to src/BTVNanoCommissioning/workflows/ctag_Wctt_valid_sf.py index d924cce7..d813fac9 100644 --- a/src/BTVNanoCommissioning/workflows/ctag_Wc_valid_sf.py +++ b/src/BTVNanoCommissioning/workflows/ctag_Wctt_valid_sf.py @@ -1,6 +1,5 @@ import os import collections, awkward as ak, numpy as np -import os import uproot from coffea import processor from coffea.analysis_tools import Weights @@ -9,6 +8,7 @@ load_lumi, load_SF, muSFs, + eleSFs, puwei, btagSFs, JME_shifts, @@ -17,11 +17,11 @@ from BTVNanoCommissioning.helpers.func import ( flatten, update, - uproot_writeable, dump_lumi, ) from BTVNanoCommissioning.helpers.update_branch import missing_branch from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import ( jet_id, mu_idiso, @@ -40,6 +40,7 @@ def __init__( isArray=False, noHist=False, chunksize=75000, + selectionModifier="WcM", ): self._year = year self._campaign = campaign @@ -51,6 +52,7 @@ def __init__( self.chunksize = chunksize ## Load corrections self.SF_map = load_SF(self._campaign) + self.selMod = selectionModifier @property def accumulator(self): @@ -97,8 +99,34 @@ def process(self, events): def process_shift(self, events, shift_name): dataset = events.metadata["dataset"] isRealData = not hasattr(events, "genWeight") + + isMu = False + isEle = False + if "WcM" in self.selMod or "semittM" in self.selMod: + triggers = ["IsoMu27", "IsoMu24"] + isMu = True + dxySigcut = 0 # 1 + muNeEmSum = 1 # 0.7 + muonpTratioCut = 1 # 0.4 + isolepdz, isolepdxy, isolepsip3d = 0.01, 0.002, 2 + elif "WcE" in self.selMod or "semittE" in self.selMod: + triggers = ["Ele32_WPTight_Gsf_L1DoubleEG"] + isEle = True + dxySigcut = 0 + muNeEmSum = 1 + muonpTratioCut = 1 # 0.6 + isolepdz, isolepdxy, isolepsip3d = 0.02, 0.01, 2.5 + else: + raise ValueError(self.selMod, "is not a valid selection modifier.") + + histoname = { + "WcM": "ctag_Wc_sf", + "WcE": "ectag_Wc_sf", + "semittM": "ctag_Wc_sf", # same histogram representation as W+c + "semittE": "ectag_Wc_sf", # same histogram representation as W+c + } _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ctag_Wc_sf") + {"": None} if self.noHist else histogrammer(events, histoname[self.selMod]) ) output = { @@ -106,11 +134,10 @@ def process_shift(self, events, shift_name): **_hist_event_dict, } - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) + if isRealData: + output["sumw"] = len(events) + else: + output["sumw"] = ak.sum(events.genWeight) #################### # Selections # #################### @@ -123,7 +150,6 @@ def process_shift(self, events, shift_name): output = dump_lumi(events[req_lumi], output) ## HLT - triggers = ["IsoMu24"] checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) if ak.all(checkHLT == False): raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) @@ -136,49 +162,69 @@ def process_shift(self, events, shift_name): for t in trig_arrs: req_trig = req_trig | t - ## Muon cuts - # muon twiki: https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideMuonIdRun2 - iso_muon = events.Muon[(events.Muon.pt > 30) & mu_idiso(events, self._campaign)] - req_muon = ak.count(iso_muon.pt, axis=1) == 1 + ## Lepton cuts + if isMu: + # muon twiki: https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideMuonIdRun2 + iso_lep = events.Muon[ + (events.Muon.pt > 30) & mu_idiso(events, self._campaign) + ] + elif isEle: + iso_lep = events.Electron[ + (events.Electron.pt > 34) & ele_mvatightid(events, self._campaign) + ] + req_lep = ak.count(iso_lep.pt, axis=1) == 1 jet_sel = ak.fill_none( jet_id(events, self._campaign) - & (ak.all(events.Jet.metric_table(iso_muon) > 0.5, axis=2)), + & (ak.all(events.Jet.metric_table(iso_lep) > 0.5, axis=2)), False, axis=-1, ) - iso_muon = ak.pad_none(iso_muon, 1, axis=1) - iso_muon = iso_muon[:, 0] - iso_muindx = ak.mask( - ak.local_index(events.Muon.pt), - ((events.Muon.pt > 30) & mu_idiso(events, self._campaign)) == 1, - ) - iso_muindx = ak.pad_none(iso_muindx, 1) - iso_muindx = iso_muindx[:, 0] + iso_lep = ak.pad_none(iso_lep, 1, axis=1) + iso_lep = iso_lep[:, 0] + if isMu: + iso_lepindx = ak.mask( + ak.local_index(events.Muon.pt), + ((events.Muon.pt > 30) & mu_idiso(events, self._campaign)) == 1, + ) + elif isEle: + iso_lepindx = ak.mask( + ak.local_index(events.Electron.pt), + ((events.Electron.pt > 34) & ele_mvatightid(events, self._campaign)) + == 1, + ) + iso_lepindx = ak.pad_none(iso_lepindx, 1) + iso_lepindx = iso_lepindx[:, 0] ## Jet cuts if "DeepJet_nsv" in events.Jet.fields: jet_sel = jet_sel & (events.Jet.DeepJet_nsv > 0) event_jet = events.Jet[jet_sel] - req_jets = (ak.num(event_jet.pt) >= 1) & (ak.num(event_jet.pt) <= 3) + nseljet = ak.count(event_jet.pt, axis=1) + if "Wc" in self.selMod: + req_jets = (nseljet >= 1) & (nseljet <= 3) + else: + req_jets = nseljet >= 4 ## Soft Muon cuts soft_muon = events.Muon[ softmu_mask(events, self._campaign) - & (abs(events.Muon.dxy / events.Muon.dxyErr) > 1.0) + & (abs(events.Muon.dxy / events.Muon.dxyErr) > dxySigcut) ] req_softmu = ak.count(soft_muon.pt, axis=1) >= 1 mujetsel = ak.fill_none( ( (ak.all(event_jet.metric_table(soft_muon) <= 0.4, axis=2)) & ((event_jet.muonIdx1 != -1) | (event_jet.muonIdx2 != -1)) - & ((event_jet.muEF + event_jet.neEmEF) < 0.7) + & ((event_jet.muEF + event_jet.neEmEF) < muNeEmSum) + & (event_jet.pt > 20) + & ((event_jet.pt / event_jet.E) > 0.03) ), False, axis=-1, ) mujetsel2 = ak.fill_none( ( - ((events.Jet.muEF + events.Jet.neEmEF) < 0.7) + ((events.Jet.muEF + events.Jet.neEmEF) < muNeEmSum) & ( ak.all( events.Jet.metric_table(soft_muon) <= 0.4, @@ -192,8 +238,10 @@ def process_shift(self, events, shift_name): axis=-1, ) soft_muon = ak.pad_none(soft_muon, 1, axis=1) + soft_muon["dxySig"] = soft_muon.dxy / soft_muon.dxyErr ## Muon-jet cuts + event_jet["isMuonJet"] = mujetsel mu_jet = event_jet[mujetsel] otherjets = event_jet[~mujetsel] req_mujet = ak.num(mu_jet.pt, axis=1) >= 1 @@ -208,24 +256,24 @@ def process_shift(self, events, shift_name): jetindx = jetindx[:, 0] # Other cuts - req_pTratio = (soft_muon[:, 0].pt / mu_jet[:, 0].pt) < 0.4 - idx = np.where(iso_muon.jetIdx == -1, 0, iso_muon.jetIdx) - req_QCDveto = ( - (iso_muon.pfRelIso04_all < 0.05) - & (abs(iso_muon.dz) < 0.01) - & (abs(iso_muon.dxy) < 0.002) - & (iso_muon.sip3d < 2) - & ( - iso_muon.pt - / ak.firsts( - events.Jet[ - (events.Jet.muonIdx1 == iso_muindx) - | ((events.Jet.muonIdx2 == iso_muindx)) - ].pt - ) - > 0.75 - ) - ) + req_pTratio = (soft_muon[:, 0].pt / mu_jet[:, 0].pt) < muonpTratioCut + idx = np.where(iso_lep.jetIdx == -1, 0, iso_lep.jetIdx) + # req_QCDveto = ( + # (iso_lep.pfRelIso04_all < 0.05) + # & (abs(iso_lep.dz) < isolepdz) + # & (abs(iso_lep.dxy) < isolepdxy) + # & (iso_lep.sip3d < isolepsip3d) + # & ( + # iso_lep.pt + # / ak.firsts( + # events.Jet[ + # (events.Jet.muonIdx1 == iso_lepindx) + # | ((events.Jet.muonIdx2 == iso_lepindx)) + # ].pt + # ) + # > 0.75 + # ) + # ) dilep_mu = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] dilep_ele = events.Electron[ @@ -235,31 +283,63 @@ def process_shift(self, events, shift_name): ak.count(dilep_mu.pt, axis=1) + ak.count(dilep_ele.pt, axis=1) != 2 ) - dilep_mass = iso_muon + soft_muon[:, 0] - req_dilepmass = (dilep_mass.mass > 12.0) & ( - (dilep_mass.mass < 80) | (dilep_mass.mass > 100) - ) + dilep_mass = iso_lep + soft_muon[:, 0] + if isMu: + req_dilepmass = dilep_mass.mass > 12.0 + # & ( + # (dilep_mass.mass < 80) | (dilep_mass.mass > 100) + # ) + elif isEle: + req_dilepmass = iso_lep.pt > 0 - MET = ak.zip( + iso_lep_trans = ak.zip( { - "pt": events.MET.pt, - "eta": ak.zeros_like(events.MET.pt), - "phi": events.MET.phi, - "mass": ak.zeros_like(events.MET.pt), + "pt": iso_lep.pt, + "eta": ak.zeros_like(iso_lep.pt), + "phi": iso_lep.phi, + "mass": iso_lep.mass, }, with_name="PtEtaPhiMLorentzVector", ) - Wmass = MET + iso_muon - # modified to transverse mass - req_mtw = ( - np.sqrt(2 * iso_muon.pt * MET.pt * (1 - np.cos(iso_muon.delta_phi(MET)))) - > 55 - ) + if "Run3" not in self._campaign: + MET = ak.zip( + { + "pt": events.MET.pt, + "eta": ak.zeros_like(events.MET.pt), + "phi": events.MET.phi, + "mass": ak.zeros_like(events.MET.pt), + }, + with_name="PtEtaPhiMLorentzVector", + ) + else: + MET = ak.zip( + { + "pt": events.PuppiMET.pt, + "eta": ak.zeros_like(events.PuppiMET.pt), + "phi": events.PuppiMET.phi, + "mass": ak.zeros_like(events.PuppiMET.pt), + }, + with_name="PtEtaPhiMLorentzVector", + ) + + wmasscut = 55 + if "semitt" in self.selMod: + wmasscut = 0 + Wcand = MET + iso_lep_trans # transverse mass + Wmass = Wcand.mass + Wpt = Wcand.pt + req_mtw = Wmass > wmasscut + + # ==This is the manual calculation for transverse mass== + # dphi = iso_lep.phi-events.PuppiMET.phi + # dphi = np.where(dphinp.pi,dphi-2*np.pi,dphi) + # trans = np.sqrt(2*iso_lep.pt*events.PuppiMET.pt*(1-np.cos(dphi))) event_level = ( req_trig & req_lumi - & req_muon + & req_lep & req_jets & req_softmu & req_dilepmass @@ -270,24 +350,42 @@ def process_shift(self, events, shift_name): ) event_level = ak.fill_none(event_level, False) if len(events[event_level]) == 0: + if self.isArray: + array_writer( + self, + events[event_level], + events, + "nominal", + dataset, + isRealData, + empty=True, + ) return {dataset: output} #################### # Selected objects # #################### - shmu = iso_muon[event_level] + shmu = iso_lep[event_level] + wm = Wmass[event_level] + wp = Wpt[event_level] sjets = event_jet[event_level] ssmu = soft_muon[event_level] smet = MET[event_level] smuon_jet = mu_jet[event_level] sotherjets = otherjets[event_level] + sdilep = dilep_mass[event_level] nsoftmu = ak.count(ssmu.pt, axis=1) nmujet = ak.count(smuon_jet.pt, axis=1) smuon_jet = smuon_jet[:, 0] ssmu = ssmu[:, 0] sz = shmu + ssmu sw = shmu + smet - osss = ak.values_astype(shmu.charge * ssmu.charge * -1, int) + + osss = 1 + ossswrite = shmu.charge * ssmu.charge * -1 + if "Wc" in self.selMod: + osss = shmu.charge * ssmu.charge * -1 + njet = ak.count(sjets.pt, axis=1) # Find the PFCands associate with selected jets. Search from jetindex->JetPFCands->PFCand if "PFCands" in events.fields: @@ -320,8 +418,10 @@ def process_shift(self, events, shift_name): weights, syst_wei, ) - if "MUO" in self.SF_map.keys(): + if isMu and "MUO" in self.SF_map.keys(): muSFs(shmu, self.SF_map, weights, syst_wei, False) + if isEle and "EGM" in self.SF_map.keys(): + eleSFs(shmu, self.SF_map, weights, syst_wei, False) if "BTV" in self.SF_map.keys(): btagSFs(smuon_jet, self.SF_map, weights, "DeepJetC", syst_wei) btagSFs(smuon_jet, self.SF_map, weights, "DeepJetB", syst_wei) @@ -436,7 +536,11 @@ def process_shift(self, events, shift_name): syst="noSF", flav=smflav, osss=osss, - discr=smuon_jet[histname.replace(f"_{i}", "")], + discr=np.where( + smuon_jet[histname.replace(f"_{i}", "")] < 0, + -0.2, + smuon_jet[histname.replace(f"_{i}", "")], + ), weight=weights.partial_weight(exclude=exclude_btv), ) if not isRealData and "btag" in self.SF_map.keys(): @@ -444,10 +548,16 @@ def process_shift(self, events, shift_name): syst=syst, flav=smflav, osss=osss, - discr=smuon_jet[histname.replace(f"_{i}", "")], + discr=np.where( + smuon_jet[histname.replace(f"_{i}", "")] < 0, + -0.2, + smuon_jet[histname.replace(f"_{i}", "")], + ), weight=weight, ) elif "btag" in histname and "Trans" in histname: + if histname not in smuon_jet: + continue for i in range(2): histname = histname.replace("Trans", "").replace(f"_{i}", "") h.fill( @@ -522,12 +632,20 @@ def process_shift(self, events, shift_name): if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] - pruned_ev["Jet"] = sjets + pruned_ev["SelJet"] = sjets pruned_ev["Muon"] = shmu pruned_ev["MuonJet"] = smuon_jet pruned_ev["SoftMuon"] = ssmu pruned_ev["OtherJets"] = sotherjets - pruned_ev["osss"] = osss + if "Wc" in self.selMod: + pruned_ev["osss"] = osss + else: + pruned_ev["osss"] = ossswrite + pruned_ev["njet"] = njet + pruned_ev["W_transmass"] = wm + pruned_ev["W_pt"] = wp + pruned_ev["dilep_mass"] = sdilep.mass + pruned_ev["dilep_pt"] = sdilep.pt if "PFCands" in events.fields: pruned_ev.PFCands = spfcands # Add custom variables @@ -543,47 +661,10 @@ def process_shift(self, events, shift_name): pruned_ev["dr_lep1_softmu"] = shmu.delta_r(ssmu) pruned_ev["soft_l_ptratio"] = ssmu.pt / smuon_jet.pt pruned_ev["l1_ptratio"] = shmu.pt / smuon_jet.pt + pruned_ev["MuonJet_beta"] = smuon_jet.pt / smuon_jet.E + pruned_ev["MuonJet_muneuEF"] = smuon_jet.muEF + smuon_jet.neEmEF - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where( - (out_branch == "SoftMuon") - | (out_branch == "MuonJet") - | (out_branch == "dilep") - | (out_branch == "OtherJets") - ), - ) - - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in [ - "Muon", - "Jet", - "SoftMuon", - "MuonJet", - "dilep", - "charge", - "MET", - ]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon" and obj != "SoftMuon") and ( - "pfRelIso04_all" == kin or "d" in kin - ): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*", "SV_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/{systematics[0]}_{events.metadata['filename'].split('/')[-1].replace('.root','')}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) + array_writer(self, pruned_ev, events, systematics[0], dataset, isRealData) return {dataset: output} diff --git a/src/BTVNanoCommissioning/workflows/ctag_dileptt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_dileptt_valid_sf.py index 3e4b5d9d..b9d99a47 100644 --- a/src/BTVNanoCommissioning/workflows/ctag_dileptt_valid_sf.py +++ b/src/BTVNanoCommissioning/workflows/ctag_dileptt_valid_sf.py @@ -8,6 +8,7 @@ load_lumi, load_SF, muSFs, + eleSFs, puwei, btagSFs, JME_shifts, @@ -17,12 +18,11 @@ from BTVNanoCommissioning.helpers.func import ( flatten, update, - uproot_writeable, - _is_rootcompat, dump_lumi, ) from BTVNanoCommissioning.helpers.update_branch import missing_branch from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import ( jet_id, mu_idiso, @@ -41,6 +41,7 @@ def __init__( isArray=False, noHist=False, chunksize=75000, + selectionModifier="DilepTTEM", ): self._year = year self._campaign = campaign @@ -50,6 +51,7 @@ def __init__( self.noHist = noHist self.lumiMask = load_lumi(self._campaign) self.chunksize = chunksize + self.selMod = selectionModifier ## Load corrections self.SF_map = load_SF(self._campaign) @@ -98,20 +100,6 @@ def process(self, events): def process_shift(self, events, shift_name): dataset = events.metadata["dataset"] isRealData = not hasattr(events, "genWeight") - _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ctag_ttdilep_sf") - ) - - output = { - "sumw": processor.defaultdict_accumulator(float), - **_hist_event_dict, - } - - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) #################### # Selections # @@ -120,12 +108,12 @@ def process_shift(self, events, shift_name): req_lumi = np.ones(len(events), dtype="bool") if isRealData: req_lumi = self.lumiMask(events.run, events.luminosityBlock) - # only dump for nominal case - if shift_name is None: - output = dump_lumi(events[req_lumi], output) ## HLT - triggers = ["Mu17_TrkIsoVVL_Mu8_TrkIsoVVL_DZ_Mass8"] + if self.selMod == "dilepttM": + triggers = ["Mu17_TrkIsoVVL_Mu8_TrkIsoVVL_DZ_Mass8"] + elif self.selMod == "dilepttE": + triggers = ["Ele23_Ele12_CaloIdL_TrackIdL_IsoVL"] checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) if ak.all(checkHLT == False): raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) @@ -138,18 +126,43 @@ def process_shift(self, events, shift_name): for t in trig_arrs: req_trig = req_trig | t - ## Muon cuts - # muon twiki: https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideMuonIdRun2 - iso_muon = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] - iso_muon = ak.pad_none(iso_muon, 2) - req_muon = (ak.count(iso_muon.pt, axis=1) == 2) & (iso_muon[:, 0].pt > 20) + # Lepton selections + if self.selMod == "dilepttM": + # dilepton selections + iso_muon = events.Muon[ + (events.Muon.pt > 12) & mu_idiso(events, self._campaign) + ] + iso_lep = ak.pad_none(iso_muon, 2) + req_lep = (ak.count(iso_lep.pt, axis=1) == 2) & (iso_lep[:, 0].pt > 20) + # veto other flavors + dilep_ele = events.Electron[ + (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) + ] + req_dilepveto = ak.count(dilep_ele.pt, axis=1) == 0 + elif self.selMod == "dilepttE": + # dilepton selections + iso_ele = events.Electron[ + (events.Electron.pt > 25) & ele_mvatightid(events, self._campaign) + ] + iso_lep = ak.pad_none(iso_ele, 2) + req_lep = (ak.count(iso_lep.pt, axis=1) == 2) & (iso_lep[:, 0].pt > 27) + # veto other flavors + dilep_mu = events.Muon[ + (events.Muon.pt > 12) & mu_idiso(events, self._campaign) + ] + req_dilepveto = ak.count(dilep_mu.pt, axis=1) == 0 + # veto Z events + dilep_mass = iso_lep[:, 0] + iso_lep[:, 1] + req_dilepmass = (dilep_mass.mass > 12.0) & ( + (dilep_mass.mass < 75) | (dilep_mass.mass > 105) + ) ## Jet cuts event_jet = events.Jet[ ak.fill_none( jet_id(events, self._campaign) & ak.all( - events.Jet.metric_table(iso_muon) > 0.4, axis=2, mask_identity=True + events.Jet.metric_table(iso_lep) > 0.4, axis=2, mask_identity=True ), False, axis=-1, @@ -198,17 +211,6 @@ def process_shift(self, events, shift_name): jetindx = ak.pad_none(jetindx, 1) jetindx = jetindx[:, 0] - ## Other cuts - dilep_ele = events.Electron[ - (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) - ] - req_dilepveto = ak.count(dilep_ele.pt, axis=1) == 0 - - dilep_mass = iso_muon[:, 0] + iso_muon[:, 1] - req_dilepmass = (dilep_mass.mass > 12.0) & ( - (dilep_mass.mass < 75) | (dilep_mass.mass > 105) - ) - MET = ak.zip( { "pt": events.MET.pt, @@ -223,7 +225,7 @@ def process_shift(self, events, shift_name): event_level = ( req_trig & req_lumi - & req_muon + & req_lep & req_dilepveto & req_dilepmass & req_MET @@ -232,20 +234,50 @@ def process_shift(self, events, shift_name): & req_mujet ) event_level = ak.fill_none(event_level, False) + histname = { + "dilepttM": "ctag_ttdilep_sf", + "dilepttE": "ectag_ttdilep_sf", + } + _hist_event_dict = ( + {"": None} if self.noHist else histogrammer(events, histname[self.selMod]) + ) + + output = { + "sumw": processor.defaultdict_accumulator(float), + **_hist_event_dict, + } + if shift_name is None: + if isRealData: + output["sumw"] = len(events) + else: + output["sumw"] = ak.sum(events.genWeight) + # only dump for nominal case + if shift_name is None: + output = dump_lumi(events[req_lumi], output) if len(events[event_level]) == 0: + if self.isArray: + array_writer( + self, + events[event_level], + events, + "nominal", + dataset, + isRealData, + empty=True, + ) return {dataset: output} #################### # Selected objects # #################### - shmu = iso_muon[event_level] + shlep = iso_lep[event_level] ssmu = soft_muon[event_level] nsoftmu = ak.count(ssmu.pt, axis=1) softmu0 = ssmu softmu0 = ssmu[:, 0] - sz = shmu[:, 0] + shmu[:, 1] - isomu0 = shmu[:, 0] - isomu1 = shmu[:, 1] + sz = shlep[:, 0] + shlep[:, 1] + isomu0 = shlep[:, 0] + isomu1 = shlep[:, 1] sjets = event_jet[event_level] smuon_jet = mu_jet[event_level] nmujet = ak.count(smuon_jet.pt, axis=1) @@ -282,8 +314,10 @@ def process_shift(self, events, shift_name): weights, syst_wei, ) - if "MUO" in self.SF_map.keys(): - muSFs(shmu, self.SF_map, weights, syst_wei, False) + if "MUO" in self.SF_map.keys() and self.selMod == "dilepttM": + muSFs(shlep, self.SF_map, weights, syst_wei, False) + if "EGM" in self.SF_map.keys() and self.selMod == "dilepttE": + eleSFs(shlep, self.SF_map, weights, syst_wei, False) if "BTV" in self.SF_map.keys(): btagSFs(smuon_jet, self.SF_map, weights, "DeepJetC", syst_wei) btagSFs(smuon_jet, self.SF_map, weights, "DeepJetB", syst_wei) @@ -459,21 +493,27 @@ def process_shift(self, events, shift_name): ####################### # Create root files # ####################### + if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] - pruned_ev.Jet = sjets - pruned_ev.Muon = shmu - pruned_ev["dilep"] = shmu[:, 0] + shmu[:, 1] + # pruned_ev["SelJet"] = sjets + if self.selMod == "dilepttM": + pruned_ev["Muon1"] = shlep[:, 0] + pruned_ev["Muon2"] = shlep[:, 1] + if self.selMod == "dilepttE": + pruned_ev["Electron1"] = shlep[:, 0] + pruned_ev["Electron2"] = shlep[:, 1] + + pruned_ev["MuonJet"] = smuon_jet + pruned_ev["SoftMuon"] = ssmu + pruned_ev["dilep"] = shlep[:, 0] + shlep[:, 1] pruned_ev["dilep", "pt"] = pruned_ev.dilep.pt pruned_ev["dilep", "eta"] = pruned_ev.dilep.eta pruned_ev["dilep", "phi"] = pruned_ev.dilep.phi pruned_ev["dilep", "mass"] = pruned_ev.dilep.mass if "PFCands" in events.fields: pruned_ev.PFCands = spfcands - pruned_ev["MuonJet"] = smuon_jet - pruned_ev["SoftMuon"] = ssmu - # Add custom variables if not isRealData: pruned_ev["weight"] = weights.weight() @@ -481,46 +521,11 @@ def process_shift(self, events, shift_name): pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( include=[ind_wei] ) - pruned_ev["dr_mujet_softmu"] = ssmu[:, 0].delta_r(smuon_jet) - pruned_ev["dr_mujet_lep1"] = shmu[:, 0].delta_r(smuon_jet) - pruned_ev["dr_mujet_lep2"] = shmu[:, 1].delta_r(smuon_jet) - pruned_ev["dr_lep1_softmu"] = shmu[:, 0].delta_r(ssmu[:, 0]) pruned_ev["soft_l_ptratio"] = ssmu[:, 0].pt / smuon_jet.pt - pruned_ev["l1_ptratio"] = shmu[:, 0].pt / smuon_jet.pt - pruned_ev["l2_ptratio"] = shmu[:, 1].pt / smuon_jet.pt - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where( - (out_branch == "SoftMuon") - | (out_branch == "MuonJet") - | (out_branch == "dilep") - ), - ) + array_writer(self, pruned_ev, events, systematics[0], dataset, isRealData) - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in ["Muon", "Jet", "SoftMuon", "MuonJet", "dilep", "MET"]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon" and obj != "SoftMuon") and ( - "pfRelIso04_all" == kin or "d" in kin - ): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) return {dataset: output} def postprocess(self, accumulator): diff --git a/src/BTVNanoCommissioning/workflows/ctag_eDY_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_eDY_valid_sf.py deleted file mode 100644 index 80bca99c..00000000 --- a/src/BTVNanoCommissioning/workflows/ctag_eDY_valid_sf.py +++ /dev/null @@ -1,448 +0,0 @@ -import collections, awkward as ak, numpy as np -import os -import uproot -from coffea import processor -from coffea.analysis_tools import Weights - -from BTVNanoCommissioning.utils.correction import ( - load_lumi, - load_SF, - eleSFs, - puwei, - btagSFs, - JME_shifts, - Roccor_shifts, -) -from BTVNanoCommissioning.helpers.func import ( - flatten, - update, - uproot_writeable, - dump_lumi, -) -from BTVNanoCommissioning.helpers.update_branch import missing_branch -from BTVNanoCommissioning.utils.histogrammer import histogrammer -from BTVNanoCommissioning.utils.selection import jet_id, mu_idiso, ele_mvatightid - - -class NanoProcessor(processor.ProcessorABC): - def __init__( - self, - year="2022", - campaign="Summer22Run3", - name="", - isSyst=False, - isArray=False, - noHist=False, - chunksize=75000, - ): - self._year = year - self._campaign = campaign - self.name = name - self.isSyst = isSyst - self.isArray = isArray - self.noHist = noHist - self.lumiMask = load_lumi(self._campaign) - self.chunksize = chunksize - ## Load corrections - self.SF_map = load_SF(self._campaign) - - @property - def accumulator(self): - return self._accumulator - - def process(self, events): - isRealData = not hasattr(events, "genWeight") - dataset = events.metadata["dataset"] - events = missing_branch(events) - shifts = [] - if "JME" in self.SF_map.keys(): - syst_JERC = self.isSyst - if self.isSyst == "JERC_split": - syst_JERC = "split" - shifts = JME_shifts( - shifts, self.SF_map, events, self._campaign, isRealData, syst_JERC - ) - else: - if int(self._year) < 2020: - shifts = [ - ({"Jet": events.Jet, "MET": events.MET, "Muon": events.Muon}, None) - ] - else: - shifts = [ - ( - { - "Jet": events.Jet, - "MET": events.PuppiMET, - "Muon": events.Muon, - }, - None, - ) - ] - if "roccor" in self.SF_map.keys(): - shifts = Roccor_shifts(shifts, self.SF_map, events, isRealData, False) - else: - shifts[0][0]["Muon"] = events.Muon - - return processor.accumulate( - self.process_shift(update(events, collections), name) - for collections, name in shifts - ) - - def process_shift(self, events, shift_name): - dataset = events.metadata["dataset"] - isRealData = not hasattr(events, "genWeight") - _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ectag_DY_sf") - ) - - output = { - "sumw": processor.defaultdict_accumulator(float), - **_hist_event_dict, - } - - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) - #################### - # Selections # - #################### - ## Lumimask - req_lumi = np.ones(len(events), dtype="bool") - if isRealData: - req_lumi = self.lumiMask(events.run, events.luminosityBlock) - # only dump for nominal case - if shift_name is None: - output = dump_lumi(events[req_lumi], output) - - ## HLT - triggers = ["Ele23_Ele12_CaloIdL_TrackIdL_IsoVL"] - checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) - if ak.all(checkHLT == False): - raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) - elif ak.any(checkHLT == False): - print(np.array(triggers)[~checkHLT], " not exist in", dataset) - trig_arrs = [ - events.HLT[_trig] for _trig in triggers if hasattr(events.HLT, _trig) - ] - req_trig = np.zeros(len(events), dtype="bool") - for t in trig_arrs: - req_trig = req_trig | t - - ## Muon cuts - dilep_mu = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] - ## Electron cuts - dilep_ele = events.Electron[ - (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) - ] - - ## dilepton - pos_dilep = dilep_ele[dilep_ele.charge > 0] - neg_dilep = dilep_ele[dilep_ele.charge < 0] - req_dilep = ( - (ak.num(pos_dilep.pt) >= 1) - & (ak.num(neg_dilep.pt) >= 1) - & (ak.num(dilep_ele.charge) >= 2) - & (ak.num(dilep_mu.charge) == 0) - ) - pos_dilep = ak.pad_none(pos_dilep, 1, axis=1) - neg_dilep = ak.pad_none(neg_dilep, 1, axis=1) - - dilep_mass = pos_dilep[:, 0] + neg_dilep[:, 0] - req_dilepmass = ak.fill_none( - ( - (dilep_mass.mass > 81) - & (dilep_mass.mass < 101) - & (dilep_mass.pt > 15) - & ((pos_dilep[:, 0].pt > 27) | (neg_dilep[:, 0].pt > 27)) - ), - False, - axis=-1, - ) - - ## Jet cuts - event_jet = events.Jet[ - ak.fill_none( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(pos_dilep[:, 0]) > 0.4, - axis=2, - mask_identity=True, - ) - ) - & ( - ak.all( - events.Jet.metric_table(neg_dilep[:, 0]) > 0.4, - axis=2, - mask_identity=True, - ) - ), - False, - axis=-1, - ) - ] - req_jets = ak.num(event_jet.pt) >= 1 - - ## store jet index for PFCands, create mask on the jet index - jetindx = ak.mask( - ak.local_index(events.Jet.pt), - ( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(pos_dilep[:, 0]) > 0.4, - axis=2, - mask_identity=True, - ) - ) - & ( - ak.all( - events.Jet.metric_table(neg_dilep[:, 0]) > 0.4, - axis=2, - mask_identity=True, - ) - ) - ) - == 1, - ) - jetindx = ak.pad_none(jetindx, 1) - jetindx = jetindx[:, 0] - - event_level = ak.fill_none( - req_lumi & req_trig & req_dilep & req_dilepmass & req_jets, False - ) - if len(events[event_level]) == 0: - return {dataset: output} - - #################### - # Selected objects # - #################### - sposmu = pos_dilep[event_level] - sposmu = sposmu[:, 0] - snegmu = neg_dilep[event_level] - snegmu = snegmu[:, 0] - sz = sposmu + snegmu - sjets = event_jet[event_level] - njet = ak.count(sjets.pt, axis=1) - sel_mu = ak.concatenate([sposmu, snegmu]) - sel_jet = sjets[:, 0] - sele = ak.zip( - { - b: ak.Array(np.reshape(sel_mu[b].to_numpy(), (len(sposmu[b]), 2))) - for b in sposmu.fields - } - ) - # Find the PFCands associate with selected jets. Search from jetindex->JetPFCands->PFCand - if "PFCands" in events.fields: - spfcands = events[event_level].PFCands[ - events[event_level] - .JetPFCands[ - events[event_level].JetPFCands.jetIdx == jetindx[event_level] - ] - .pFCandsIdx - ] - - #################### - # Weight & Geninfo # - #################### - weights = Weights(len(events[event_level]), storeIndividual=True) - if not isRealData: - weights.add("genweight", events[event_level].genWeight) - par_flav = (sel_jet.partonFlavour == 0) & (sel_jet.hadronFlavour == 0) - genflavor = sel_jet.hadronFlavour + 1 * par_flav - if len(self.SF_map.keys()) > 0: - syst_wei = True if self.isSyst != False else False - if "PU" in self.SF_map.keys(): - puwei( - events[event_level].Pileup.nTrueInt, - self.SF_map, - weights, - syst_wei, - ) - if "EGM" in self.SF_map.keys(): - eleSFs(sele, self.SF_map, weights, syst_wei, False) - if "BTV" in self.SF_map.keys(): - btagSFs(sel_jet, self.SF_map, weights, "DeepJetC", syst_wei) - btagSFs(sel_jet, self.SF_map, weights, "DeepJetB", syst_wei) - btagSFs(sel_jet, self.SF_map, weights, "DeepCSVB", syst_wei) - btagSFs(sel_jet, self.SF_map, weights, "DeepCSVC", syst_wei) - else: - genflavor = ak.zeros_like(sel_jet.pt, dtype=int) - - # Systematics information - if shift_name is None: - systematics = ["nominal"] + list(weights.variations) - else: - systematics = [shift_name] - exclude_btv = [ - "DeepCSVC", - "DeepCSVB", - "DeepJetB", - "DeepJetC", - ] # exclude b-tag SFs for btag inputs - - #################### - # Fill histogram # - #################### - for syst in systematics: - if self.isSyst == False and syst != "nominal": - break - if self.noHist: - break - weight = ( - weights.weight() - if syst == "nominal" or syst == shift_name - else weights.weight(modifier=syst) - ) - for histname, h in output.items(): - if ( - "Deep" in histname - and "btag" not in histname - and histname in events.Jet.fields - ): - h.fill( - syst, - genflavor, - sel_jet[histname], - weight=weights.partial_weight(exclude=exclude_btv), - ) - elif ( - "PFCands" in events.fields - and "PFCands" in histname - and histname.split("_")[1] in events.PFCands.fields - ): - h.fill( - syst, - flatten(ak.broadcast_arrays(genflavor, spfcands["pt"])[0]), - flatten(spfcands[histname.replace("PFCands_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), - spfcands["pt"], - )[0] - ), - ) - elif ( - "posl_" in histname - and histname.replace("posl_", "") in sposmu.fields - ): - h.fill( - syst, - flatten(sposmu[histname.replace("posl_", "")]), - weight=weight, - ) - elif ( - "negl_" in histname - and histname.replace("negl_", "") in snegmu.fields - ): - h.fill( - syst, - flatten(snegmu[histname.replace("negl_", "")]), - weight=weight, - ) - - elif "jet_" in histname: - h.fill( - syst, - genflavor, - sel_jet[histname.replace("jet_", "")], - weight=weight, - ) - elif ( - "btag" in histname - and "0" in histname - and histname.replace("_0", "") in events.Jet.fields - ): - h.fill( - syst="noSF", - flav=genflavor, - discr=np.where( - sel_jet[histname.replace("_0", "")] < 0, - -0.2, - sel_jet[histname.replace("_0", "")], - ), - weight=weights.partial_weight(exclude=exclude_btv), - ) - if not isRealData and "btag" in self.SF_map.keys(): - h.fill( - syst=syst, - flav=genflavor, - discr=np.where( - sel_jet[histname.replace("_0", "")] < 0, - -0.2, - sel_jet[histname.replace("_0", "")], - ), - weight=weight, - ) - output["njet"].fill(syst, njet, weight=weight) - output["dr_mumu"].fill(syst, snegmu.delta_r(sposmu), weight=weight) - output["z_pt"].fill(syst, flatten(sz.pt), weight=weight) - output["z_eta"].fill(syst, flatten(sz.eta), weight=weight) - output["z_phi"].fill(syst, flatten(sz.phi), weight=weight) - output["z_mass"].fill(syst, flatten(sz.mass), weight=weight) - output["npvs"].fill( - syst, - events[event_level].PV.npvs, - weight=weight, - ) - if not isRealData: - output["pu"].fill( - syst, - events[event_level].Pileup.nTrueInt, - weight=weight, - ) - ####################### - # Create root files # - ####################### - if self.isArray: - # Keep the structure of events and pruned the object size - pruned_ev = events[event_level] - pruned_ev.Jet = sel_jet - pruned_ev.Muon = sele - pruned_ev["dilep"] = sposmu + snegmu - pruned_ev["dilep", "pt"] = pruned_ev.dilep.pt - pruned_ev["dilep", "eta"] = pruned_ev.dilep.eta - pruned_ev["dilep", "phi"] = pruned_ev.dilep.phi - pruned_ev["dilep", "mass"] = pruned_ev.dilep.mass - if "PFCands" in events.fields: - pruned_ev.PFCands = spfcands - # Add custom variables - if not isRealData: - pruned_ev["weight"] = weights.weight() - for ind_wei in weights.weightStatistics.keys(): - pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( - include=[ind_wei] - ) - - pruned_ev["dr_mu1jet"] = sposmu.delta_r(sel_jet) - pruned_ev["dr_mu2jet"] = snegmu.delta_r(sel_jet) - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where((out_branch == "dilep")), - ) - - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in ["Muon", "Jet", "dilep"]: - if (obj != "Muon") and ("pfRelIso04_all" == kin or "d" in kin): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) - return {dataset: output} - - def postprocess(self, accumulator): - return accumulator diff --git a/src/BTVNanoCommissioning/workflows/ctag_eWc_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_eWc_valid_sf.py deleted file mode 100644 index 966dfef8..00000000 --- a/src/BTVNanoCommissioning/workflows/ctag_eWc_valid_sf.py +++ /dev/null @@ -1,565 +0,0 @@ -import collections, awkward as ak, numpy as np -import os -import uproot - -from coffea import processor -from coffea.analysis_tools import Weights - -from BTVNanoCommissioning.utils.correction import ( - load_lumi, - load_SF, - eleSFs, - puwei, - btagSFs, - JME_shifts, - Roccor_shifts, -) -from BTVNanoCommissioning.helpers.func import ( - flatten, - update, - uproot_writeable, - dump_lumi, -) -from BTVNanoCommissioning.helpers.update_branch import missing_branch -from BTVNanoCommissioning.utils.histogrammer import histogrammer -from BTVNanoCommissioning.utils.selection import ( - jet_id, - mu_idiso, - ele_mvatightid, - softmu_mask, -) - - -class NanoProcessor(processor.ProcessorABC): - def __init__( - self, - year="2022", - campaign="Summer22Run3", - name="", - isSyst=False, - isArray=False, - noHist=False, - chunksize=75000, - ): - self._year = year - self._campaign = campaign - self.name = name - self.isSyst = isSyst - self.isArray = isArray - self.noHist = noHist - self.lumiMask = load_lumi(self._campaign) - self.chunksize = chunksize - ## Load corrections - self.SF_map = load_SF(self._campaign) - - @property - def accumulator(self): - return self._accumulator - - def process(self, events): - isRealData = not hasattr(events, "genWeight") - dataset = events.metadata["dataset"] - events = missing_branch(events) - shifts = [] - if "JME" in self.SF_map.keys(): - syst_JERC = self.isSyst - if self.isSyst == "JERC_split": - syst_JERC = "split" - shifts = JME_shifts( - shifts, self.SF_map, events, self._campaign, isRealData, syst_JERC - ) - else: - if int(self._year) > 2020: - shifts = [ - ({"Jet": events.Jet, "MET": events.MET, "Muon": events.Muon}, None) - ] - else: - shifts = [ - ( - { - "Jet": events.Jet, - "MET": events.PuppiMET, - "Muon": events.Muon, - }, - None, - ) - ] - if "roccor" in self.SF_map.keys(): - shifts = Roccor_shifts(shifts, self.SF_map, events, isRealData, False) - else: - shifts[0][0]["Muon"] = events.Muon - - return processor.accumulate( - self.process_shift(update(events, collections), name) - for collections, name in shifts - ) - - def process_shift(self, events, shift_name): - dataset = events.metadata["dataset"] - isRealData = not hasattr(events, "genWeight") - _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ectag_Wc_sf") - ) - - output = { - "sumw": processor.defaultdict_accumulator(float), - **_hist_event_dict, - } - - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) - #################### - # Selections # - #################### - ## Lumimask - req_lumi = np.ones(len(events), dtype="bool") - if isRealData: - req_lumi = self.lumiMask(events.run, events.luminosityBlock) - # only dump for nominal case - if shift_name is None: - output = dump_lumi(events[req_lumi], output) - - ## HLT - triggers = ["Ele32_WPTight_Gsf_L1DoubleEG"] - checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) - if ak.all(checkHLT == False): - raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) - elif ak.any(checkHLT == False): - print(np.array(triggers)[~checkHLT], " not exist in", dataset) - trig_arrs = [ - events.HLT[_trig] for _trig in triggers if hasattr(events.HLT, _trig) - ] - req_trig = np.zeros(len(events), dtype="bool") - for t in trig_arrs: - req_trig = req_trig | t - - ## Electron cuts - iso_ele = events.Electron[ - (events.Electron.pt > 34) & ele_mvatightid(events, self._campaign) - ] - req_ele = ak.count(iso_ele.pt, axis=1) == 1 - jet_sel = ak.fill_none( - jet_id(events, self._campaign) - & (ak.all(events.Jet.metric_table(iso_ele) > 0.5, axis=2)), - False, - axis=-1, - ) - iso_ele = ak.pad_none(iso_ele, 1, axis=1) - iso_ele = iso_ele[:, 0] - iso_eindx = ak.mask( - ak.local_index(events.Electron.pt), - ((events.Electron.pt > 34) & ele_mvatightid(events, self._campaign)) == 1, - ) - iso_eindx = ak.pad_none(iso_eindx, 1) - iso_eindx = iso_eindx[:, 0] - - ## Jet cuts - - if "DeepJet_nsv" in events.Jet.fields: - jet_sel = jet_sel & (events.Jet.DeepJet_nsv > 0) - event_jet = events.Jet[jet_sel] - req_jets = (ak.num(event_jet.pt) >= 1) & (ak.num(event_jet.pt) <= 3) - - ## Soft Muon cuts - soft_muon = events.Muon[softmu_Ï€mask(events, self._campaign)] - req_softmu = ak.count(soft_muon.pt, axis=1) >= 1 - mujetsel = ak.fill_none( - (ak.all(event_jet.metric_table(soft_muon) <= 0.4, axis=2)) - & ((event_jet.muonIdx1 != -1) | (event_jet.muonIdx2 != -1)), - False, - axis=-1, - ) - mujetsel2 = ak.fill_none( - ( - ak.all( - events.Jet.metric_table(soft_muon) <= 0.4, - axis=2, - ) - ) - & ((events.Jet.muonIdx1 != -1) | (events.Jet.muonIdx2 != -1)), - False, - axis=-1, - ) - - ## Muon-jet cuts - soft_muon = ak.pad_none(soft_muon, 1, axis=1) - - ## Muon-jet cuts - mu_jet = event_jet[mujetsel] - otherjets = event_jet[~mujetsel] - req_mujet = ak.num(mu_jet.pt, axis=1) >= 1 - mu_jet = ak.pad_none(mu_jet, 1, axis=1) - - ## store jet index for PFCands, create mask on the jet index - jet_selpf = (jet_sel) & (mujetsel2) - if "DeepJet_nsv" in events.Jet.fields: - jet_selpf = jet_selpf & (events.Jet.DeepJet_nsv > 0) - jetindx = ak.mask(ak.local_index(events.Jet.pt), jet_selpf == True) - jetindx = ak.pad_none(jetindx, 1) - jetindx = jetindx[:, 0] - - # Other cuts - req_pTratio = (soft_muon[:, 0].pt / mu_jet[:, 0].pt) < 0.6 - - req_QCDveto = ( - (iso_ele.pfRelIso03_all < 0.05) - & (abs(iso_ele.dz) < 0.02) - & (abs(iso_ele.dxy) < 0.01) - & (iso_ele.sip3d < 2.5) - & ( - iso_ele.pt - / ak.firsts( - events.Jet[ - (events.Jet.electronIdx1 == iso_eindx) - | ((events.Jet.electronIdx2 == iso_eindx)) - ].pt - ) - > 0.75 - ) - ) - - dilep_mu = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] - dilep_ele = events.Electron[ - (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) - ] - req_dilepveto = ( - ak.count(dilep_mu.pt, axis=1) + ak.count(dilep_ele.pt, axis=1) != 2 - ) - - MET = ak.zip( - { - "pt": events.MET.pt, - "eta": ak.zeros_like(events.MET.pt), - "phi": events.MET.phi, - "mass": ak.zeros_like(events.MET.pt), - }, - with_name="PtEtaPhiMLorentzVector", - ) - Wmass = MET + iso_ele - req_Wmass = Wmass.mass > 55 - - event_level = ( - req_trig - & req_lumi - & req_ele - & req_jets - & req_softmu - & req_mujet - & req_Wmass - & req_dilepveto - & req_QCDveto - & req_pTratio - ) - event_level = ak.fill_none(event_level, False) - if len(events[event_level]) == 0: - return {dataset: output} - - #################### - # Selected objects # - #################### - shmu = iso_ele[event_level] - sjets = event_jet[event_level] - ssmu = soft_muon[event_level] - smet = MET[event_level] - smuon_jet = mu_jet[event_level] - sotherjets = otherjets[event_level] - nsoftmu = ak.count(ssmu.pt, axis=1) - nmujet = ak.count(smuon_jet.pt, axis=1) - smuon_jet = smuon_jet[:, 0] - ssmu = ssmu[:, 0] - sz = shmu + ssmu - sw = shmu + smet - osss = ak.values_astype(shmu.charge * ssmu.charge * -1, int) - njet = ak.count(sjets.pt, axis=1) - # Find the PFCands associate with selected jets. Search from jetindex->JetPFCands->PFCand - if "PFCands" in events.fields: - spfcands = events[event_level].PFCands[ - events[event_level] - .JetPFCands[ - events[event_level].JetPFCands.jetIdx == jetindx[event_level] - ] - .pFCandsIdx - ] - - #################### - # Weight & Geninfo # - #################### - weights = Weights(len(events[event_level]), storeIndividual=True) - if not isRealData: - weights.add("genweight", events[event_level].genWeight) - genflavor = sjets.hadronFlavour + 1 * ( - (sjets.partonFlavour == 0) & (sjets.hadronFlavour == 0) - ) - smflav = smuon_jet.hadronFlavour + 1 * ( - (smuon_jet.partonFlavour == 0) & (smuon_jet.hadronFlavour == 0) - ) - if len(self.SF_map.keys()) > 0: - syst_wei = True if self.isSyst != False else False - if "PU" in self.SF_map.keys(): - puwei( - events[event_level].Pileup.nTrueInt, - self.SF_map, - weights, - syst_wei, - ) - if "EGM" in self.SF_map.keys(): - eleSFs(shmu, self.SF_map, weights, syst_wei, False) - if "BTV" in self.SF_map.keys(): - btagSFs(smuon_jet, self.SF_map, weights, "DeepJetC", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepJetB", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepCSVB", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepCSVC", syst_wei) - - else: - genflavor = ak.zeros_like(sjets.pt, dtype=int) - smflav = ak.zeros_like(smuon_jet.pt, dtype=int) - - # Systematics information - if shift_name is None: - systematics = ["nominal"] + list(weights.variations) - else: - systematics = [shift_name] - exclude_btv = [ - "DeepCSVC", - "DeepCSVB", - "DeepJetB", - "DeepJetC", - ] # exclude b-tag SFs for btag inputs - - #################### - # Fill histogram # - #################### - for syst in systematics: - if self.isSyst == False and syst != "nominal": - break - if self.noHist: - break - weight = ( - weights.weight() - if syst == "nominal" or syst == shift_name - else weights.weight(modifier=syst) - ) - for histname, h in output.items(): - if ( - "Deep" in histname - and "btag" not in histname - and histname in events.Jet.fields - ): - h.fill( - syst, - flatten(genflavor), - flatten(ak.broadcast_arrays(osss, sjets["pt"])[0]), - flatten(sjets[histname]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), sjets["pt"] - )[0] - ), - ) - elif ( - "PFCands" in events.fields - and "PFCands" in histname - and histname.split("_")[1] in events.PFCands.fields - ): - h.fill( - syst, - flatten(ak.broadcast_arrays(smflav, spfcands["pt"])[0]), - flatten(ak.broadcast_arrays(osss, spfcands["pt"])[0]), - flatten(spfcands[histname.replace("PFCands_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), - spfcands["pt"], - )[0] - ), - ) - elif "jet_" in histname and "mu" not in histname: - h.fill( - syst, - flatten(genflavor), - flatten(ak.broadcast_arrays(osss, sjets["pt"])[0]), - flatten(sjets[histname.replace("jet_", "")]), - weight=flatten(ak.broadcast_arrays(weight, sjets["pt"])[0]), - ) - elif "hl_" in histname and histname.replace("hl_", "") in shmu.fields: - h.fill( - syst, - osss, - flatten(shmu[histname.replace("hl_", "")]), - weight=weight, - ) - elif ( - "soft_l" in histname - and histname.replace("soft_l_", "") in ssmu.fields - ): - h.fill( - syst, - smflav, - osss, - flatten(ssmu[histname.replace("soft_l_", "")]), - weight=weight, - ) - elif "mujet_" in histname: - h.fill( - syst, - smflav, - osss, - flatten(smuon_jet[histname.replace("mujet_", "")]), - weight=weight, - ) - elif "btag" in histname: - for i in range(2): - if ( - str(i) not in histname - or histname.replace(f"_{i}", "") not in events.Jet.fields - ): - continue - h.fill( - syst="noSF", - flav=smflav, - osss=osss, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weights.partial_weight(exclude=exclude_btv), - ) - if not isRealData and "btag" in self.SF_map.keys(): - h.fill( - syst=syst, - flav=smflav, - osss=osss, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weight, - ) - - output["njet"].fill(syst, osss, njet, weight=weight) - output["nmujet"].fill(syst, osss, nmujet, weight=weight) - output["nsoftmu"].fill(syst, osss, nsoftmu, weight=weight) - output["hl_ptratio"].fill( - syst, - genflavor[:, 0], - osss=osss, - ratio=shmu.pt / sjets[:, 0].pt, - weight=weight, - ) - output["soft_l_ptratio"].fill( - syst, - flav=smflav, - osss=osss, - ratio=ssmu.pt / smuon_jet.pt, - weight=weight, - ) - output["dr_lmujetsmu"].fill( - syst, - flav=smflav, - osss=osss, - dr=smuon_jet.delta_r(ssmu), - weight=weight, - ) - output["dr_lmujethmu"].fill( - syst, - flav=smflav, - osss=osss, - dr=smuon_jet.delta_r(shmu), - weight=weight, - ) - output["dr_lmusmu"].fill( - syst, - osss=osss, - dr=shmu.delta_r(ssmu), - weight=weight, - ) - output["z_pt"].fill(syst, osss, flatten(sz.pt), weight=weight) - output["z_eta"].fill(syst, osss, flatten(sz.eta), weight=weight) - output["z_phi"].fill(syst, osss, flatten(sz.phi), weight=weight) - output["z_mass"].fill(syst, osss, flatten(sz.mass), weight=weight) - output["w_pt"].fill(syst, osss, flatten(sw.pt), weight=weight) - output["w_eta"].fill(syst, osss, flatten(sw.eta), weight=weight) - output["w_phi"].fill(syst, osss, flatten(sw.phi), weight=weight) - output["w_mass"].fill(syst, osss, flatten(sw.mass), weight=weight) - output["MET_pt"].fill(syst, osss, flatten(smet.pt), weight=weight) - output["MET_phi"].fill(syst, osss, flatten(smet.phi), weight=weight) - output["npvs"].fill( - syst, - events[event_level].PV.npvs, - weight=weight, - ) - if not isRealData: - output["pu"].fill( - syst, - events[event_level].Pileup.nTrueInt, - weight=weight, - ) - ####################### - # Create root files # - ####################### - if self.isArray: - # Keep the structure of events and pruned the object size - pruned_ev = events[event_level] - pruned_ev["Jet"] = sjets - pruned_ev["Muon"] = shmu - pruned_ev["MuonJet"] = smuon_jet - pruned_ev["OtherJets"] = sotherjets - pruned_ev["SoftMuon"] = ssmu - pruned_ev["osss"] = osss - if "PFCands" in events.fields: - pruned_ev.PFCands = spfcands - # Add custom variables - if not isRealData: - pruned_ev["weight"] = weights.weight() - for ind_wei in weights.weightStatistics.keys(): - pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( - include=[ind_wei] - ) - - pruned_ev["dr_mujet_softmu"] = ssmu.delta_r(smuon_jet) - pruned_ev["dr_mujet_lep1"] = shmu.delta_r(smuon_jet) - pruned_ev["dr_lep1_softmu"] = shmu.delta_r(ssmu) - pruned_ev["soft_l_ptratio"] = ssmu.pt / smuon_jet.pt - pruned_ev["l1_ptratio"] = shmu.pt / smuon_jet.pt - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where( - (out_branch == "SoftMuon") - # | (out_branch == "MuonJet") - | (out_branch == "dilep") - ), - ) - - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in [ - "Muon", - "Jet", - "SoftMuon", - # "MuonJet", - "dilep", - "charge", - "MET", - ]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon" and obj != "SoftMuon") and ( - "pfRelIso04_all" == kin or "d" in kin - ): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*", "SV_*"] - ) - # # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - - with uproot.recreate( - f"{self.name}/{dataset}/{events.metadata['filename'][events.metadata['filename'].rfind('/')+1:-5]}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) - return {dataset: output} - - def postprocess(self, accumulator): - return accumulator diff --git a/src/BTVNanoCommissioning/workflows/ctag_edileptt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_edileptt_valid_sf.py deleted file mode 100644 index 5c7e07ea..00000000 --- a/src/BTVNanoCommissioning/workflows/ctag_edileptt_valid_sf.py +++ /dev/null @@ -1,523 +0,0 @@ -import collections, awkward as ak, numpy as np -import os -import uproot -from coffea import processor -from coffea.analysis_tools import Weights - -from BTVNanoCommissioning.utils.correction import ( - load_lumi, - load_SF, - eleSFs, - puwei, - btagSFs, - JME_shifts, - Roccor_shifts, -) -from BTVNanoCommissioning.helpers.func import ( - flatten, - update, - uproot_writeable, - dump_lumi, -) -from BTVNanoCommissioning.helpers.update_branch import missing_branch -from BTVNanoCommissioning.utils.histogrammer import histogrammer -from BTVNanoCommissioning.utils.selection import ( - jet_id, - mu_idiso, - ele_mvatightid, - softmu_mask, -) - - -class NanoProcessor(processor.ProcessorABC): - def __init__( - self, - year="2022", - campaign="Summer22Run3", - name="", - isSyst=False, - isArray=False, - noHist=False, - chunksize=75000, - ): - self._year = year - self._campaign = campaign - self.name = name - self.isSyst = isSyst - self.isArray = isArray - self.noHist = noHist - self.lumiMask = load_lumi(self._campaign) - self.chunksize = chunksize - ## Load corrections - self.SF_map = load_SF(self._campaign) - - @property - def accumulator(self): - return self._accumulator - - def process(self, events): - isRealData = not hasattr(events, "genWeight") - dataset = events.metadata["dataset"] - events = missing_branch(events) - shifts = [] - if "JME" in self.SF_map.keys(): - syst_JERC = False - - if self.isSyst == "JERC_split": - syst_JERC = "split" - shifts = JME_shifts( - shifts, self.SF_map, events, self._campaign, isRealData, syst_JERC - ) - else: - if int(self._year) < 2020: - shifts = [ - ({"Jet": events.Jet, "MET": events.MET, "Muon": events.Muon}, None) - ] - else: - shifts = [ - ( - { - "Jet": events.Jet, - "MET": events.PuppiMET, - "Muon": events.Muon, - }, - None, - ) - ] - if "roccor" in self.SF_map.keys(): - shifts = Roccor_shifts(shifts, self.SF_map, events, isRealData, False) - else: - shifts[0][0]["Muon"] = events.Muon - - return processor.accumulate( - self.process_shift(update(events, collections), name) - for collections, name in shifts - ) - - def process_shift(self, events, shift_name): - dataset = events.metadata["dataset"] - isRealData = not hasattr(events, "genWeight") - _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ectag_ttdilep_sf") - ) - - output = { - "sumw": processor.defaultdict_accumulator(float), - **_hist_event_dict, - } - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) - #################### - # Selections # - #################### - ## Lumimask - req_lumi = np.ones(len(events), dtype="bool") - if isRealData: - req_lumi = self.lumiMask(events.run, events.luminosityBlock) - # only dump for nominal case - if shift_name is None: - output = dump_lumi(events[req_lumi], output) - - ## HLT - triggers = ["Ele23_Ele12_CaloIdL_TrackIdL_IsoVL"] - checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) - if ak.all(checkHLT == False): - raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) - elif ak.any(checkHLT == False): - print(np.array(triggers)[~checkHLT], " not exist in", dataset) - trig_arrs = [ - events.HLT[_trig] for _trig in triggers if hasattr(events.HLT, _trig) - ] - req_trig = np.zeros(len(events), dtype="bool") - for t in trig_arrs: - req_trig = req_trig | t - - ## Electron cuts - iso_ele = events.Electron[ - (events.Electron.pt > 25) & ele_mvatightid(events, self._campaign) - ] - iso_ele = ak.pad_none(iso_ele, 2) - req_ele = (ak.count(iso_ele.pt, axis=1) == 2) & (iso_ele[:, 0].pt > 27) - - ## Jet cuts - event_jet = events.Jet[ - ak.fill_none( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(iso_ele) > 0.4, - axis=2, - mask_identity=True, - ) - ), - False, - axis=-1, - ) - ] - req_jets = ak.count(event_jet.pt, axis=1) >= 2 - - ## Soft Muon cuts - soft_muon = events.Muon[softmu_mask(events, self._campaign)] - req_softmu = ak.count(soft_muon.pt, axis=1) >= 1 - - ## Muon jet cuts - mu_jet = event_jet[ - ak.fill_none( - ( - ak.all( - event_jet.metric_table(soft_muon) <= 0.4, - axis=2, - mask_identity=True, - ) - ) - & ((event_jet.muonIdx1 != -1) | (event_jet.muonIdx2 != -1)), - False, - axis=-1, - ) - ] - req_mujet = ak.count(mu_jet.pt, axis=1) >= 1 - - ## store jet index for PFCands, create mask on the jet index - jetindx = ak.mask( - ak.local_index(events.Jet.pt), - ( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(soft_muon) <= 0.4, - axis=2, - mask_identity=True, - ) - ) - & ((events.Jet.muonIdx1 != -1) | (events.Jet.muonIdx2 != -1)) - ) - == 1, - ) - jetindx = ak.pad_none(jetindx, 1) - jetindx = jetindx[:, 0] - - # Other cuts - dilep_mu = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] - req_dilepveto = ak.count(dilep_mu.pt, axis=1) == 0 - - dilep_mass = iso_ele[:, 0] + iso_ele[:, 1] - req_dilepmass = (dilep_mass.mass > 12.0) & ( - (dilep_mass.mass < 75) | (dilep_mass.mass > 105) - ) - - MET = ak.zip( - { - "pt": events.MET.pt, - "eta": ak.zeros_like(events.MET.pt), - "phi": events.MET.phi, - "mass": ak.zeros_like(events.MET.pt), - }, - with_name="PtEtaPhiMLorentzVector", - ) - req_MET = MET.pt > 40 - - event_level = ( - req_trig - & req_lumi - & req_ele - & req_dilepveto - & req_dilepmass - & req_MET - & req_jets - & req_softmu - & req_mujet - ) - event_level = ak.fill_none(event_level, False) - if len(events[event_level]) == 0: - return {dataset: output} - - #################### - # Selected objects # - #################### - shmu = iso_ele[event_level] - ssmu = soft_muon[event_level] - nsoftmu = ak.count(ssmu.pt, axis=1) - softmu0 = ssmu[:, 0] - sz = shmu[:, 0] + shmu[:, 1] - isomu0 = shmu[:, 0] - isomu1 = shmu[:, 1] - sjets = event_jet[event_level] - smuon_jet = mu_jet[event_level] - nmujet = ak.count(smuon_jet.pt, axis=1) - smuon_jet = smuon_jet[:, 0] - smet = MET[event_level] - njet = ak.count(sjets.pt, axis=1) - # Find the PFCands associate with selected jets. Search from jetindex->JetPFCands->PFCand - if "PFCands" in events.fields: - spfcands = events[event_level].PFCands[ - events[event_level] - .JetPFCands[ - events[event_level].JetPFCands.jetIdx == jetindx[event_level] - ] - .pFCandsIdx - ] - - #################### - # Weight & Geninfo # - #################### - weights = Weights(len(events[event_level]), storeIndividual=True) - if not isRealData: - weights.add("genweight", events[event_level].genWeight) - par_flav = (sjets.partonFlavour == 0) & (sjets.hadronFlavour == 0) - genflavor = ak.values_astype(sjets.hadronFlavour + 1 * par_flav, int) - smpu = (smuon_jet.partonFlavour == 0) & (smuon_jet.hadronFlavour == 0) - smflav = ak.values_astype(1 * smpu + smuon_jet.hadronFlavour, int) - if len(self.SF_map.keys()) > 0: - syst_wei = True if self.isSyst != False else False - if "PU" in self.SF_map.keys(): - puwei( - events[event_level].Pileup.nTrueInt, - self.SF_map, - weights, - syst_wei, - ) - if "EGM" in self.SF_map.keys(): - eleSFs(shmu, self.SF_map, weights, syst_wei, False) - if "BTV" in self.SF_map.keys(): - btagSFs(sjets, self.SF_map, weights, "DeepJetC", syst_wei) - btagSFs(sjets, self.SF_map, weights, "DeepJetB", syst_wei) - btagSFs(sjets, self.SF_map, weights, "DeepCSVB", syst_wei) - btagSFs(sjets, self.SF_map, weights, "DeepCSVC", syst_wei) - else: - genflavor = ak.zeros_like(sjets.pt, dtype=int) - smflav = ak.zeros_like(smuon_jet.pt, dtype=int) - - # Systematics information - if shift_name is None: - systematics = ["nominal"] + list(weights.variations) - else: - systematics = [shift_name] - exclude_btv = [ - "DeepCSVC", - "DeepCSVB", - "DeepJetB", - "DeepJetC", - ] # exclude b-tag SFs for btag inputs - - #################### - # Fill histogram # - #################### - for syst in systematics: - if self.isSyst == False and syst != "nominal": - break - if self.noHist: - break - weight = ( - weights.weight() - if syst == "nominal" or syst == shift_name - else weights.weight(modifier=syst) - ) - for histname, h in output.items(): - if ( - "Deep" in histname - and "btag" not in histname - and histname in events.Jet.fields - ): - h.fill( - syst, - flatten(genflavor), - flatten(sjets[histname]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), sjets["pt"] - )[0] - ), - ) - elif ( - "PFCands" in events.fields - and "PFCands" in histname - and histname.split("_")[1] in events.PFCands.fields - ): - h.fill( - syst, - flatten(ak.broadcast_arrays(smflav, spfcands["pt"])[0]), - flatten(spfcands[histname.replace("PFCands_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), - spfcands["pt"], - )[0] - ), - ) - elif "jet_" in histname and "mu" not in histname: - h.fill( - syst, - flatten(genflavor), - flatten(sjets[histname.replace("jet_", "")]), - weight=flatten(ak.broadcast_arrays(weight, sjets["pt"])[0]), - ) - elif "hl_" in histname and histname.replace("hl_", "") in isomu0.fields: - h.fill( - syst, - flatten(isomu0[histname.replace("hl_", "")]), - weight=weight, - ) - elif "sl_" in histname and histname.replace("sl_", "") in isomu1.fields: - h.fill( - syst, - flatten(isomu1[histname.replace("sl_", "")]), - weight=weight, - ) - elif "soft_l" in histname and not "ptratio" in histname: - h.fill( - syst, - smflav, - flatten(softmu0[histname.replace("soft_l_", "")]), - weight=weight, - ) - elif "lmujet_" in histname: - h.fill( - syst, - smflav, - flatten(smuon_jet[histname.replace("lmujet_", "")]), - weight=weight, - ) - elif "btag" in histname: - for i in range(2): - if ( - str(i) not in histname - or histname.replace(f"_{i}", "") not in events.Jet.fields - ): - continue - h.fill( - syst="noSF", - flav=smflav, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weights.partial_weight(exclude=exclude_btv), - ) - if not isRealData and "btag" in self.SF_map.keys(): - h.fill( - syst=syst, - flav=smflav, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weight, - ) - - output["njet"].fill(syst, njet, weight=weight) - output["nmujet"].fill(syst, nmujet, weight=weight) - output["nsoftmu"].fill(syst, nsoftmu, weight=weight) - output["hl_ptratio"].fill( - syst, - genflavor[:, 0], - ratio=isomu0.pt / sjets[:, 0].pt, - weight=weight, - ) - output["sl_ptratio"].fill( - syst, - genflavor[:, 0], - ratio=isomu1.pt / sjets[:, 0].pt, - weight=weight, - ) - output["soft_l_ptratio"].fill( - syst, - flav=smflav, - ratio=softmu0.pt / smuon_jet.pt, - weight=weight, - ) - output["dr_lmujetsmu"].fill( - syst, - flav=smflav, - dr=smuon_jet.delta_r(softmu0), - weight=weight, - ) - output["dr_lmujethmu"].fill( - syst, - flav=smflav, - dr=smuon_jet.delta_r(isomu0), - weight=weight, - ) - output["dr_lmusmu"].fill(syst, isomu0.delta_r(softmu0), weight=weight) - output["z_pt"].fill(syst, flatten(sz.pt), weight=weight) - output["z_eta"].fill(syst, flatten(sz.eta), weight=weight) - output["z_phi"].fill(syst, flatten(sz.phi), weight=weight) - output["z_mass"].fill(syst, flatten(sz.mass), weight=weight) - output["MET_pt"].fill(syst, flatten(smet.pt), weight=weight) - output["MET_phi"].fill(syst, flatten(smet.phi), weight=weight) - output["npvs"].fill( - syst, - events[event_level].PV.npvs, - weight=weight, - ) - if not isRealData: - output["pu"].fill( - syst, - events[event_level].Pileup.nTrueInt, - weight=weight, - ) - ####################### - # Create root files # - ####################### - if self.isArray: - # Keep the structure of events and pruned the object size - pruned_ev = events[event_level] - pruned_ev.Jet = sjets - pruned_ev.Muon = shmu - pruned_ev["dilep"] = shmu[:, 0] + shmu[:, 1] - pruned_ev["dilep", "pt"] = pruned_ev.dilep.pt - pruned_ev["dilep", "eta"] = pruned_ev.dilep.eta - pruned_ev["dilep", "phi"] = pruned_ev.dilep.phi - pruned_ev["dilep", "mass"] = pruned_ev.dilep.mass - if "PFCands" in events.fields: - pruned_ev.PFCands = spfcands - pruned_ev["MuonJet"] = smuon_jet - pruned_ev["SoftMuon"] = ssmu - - # Add custom variables - if not isRealData: - pruned_ev["weight"] = weights.weight() - for ind_wei in weights.weightStatistics.keys(): - pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( - include=[ind_wei] - ) - - pruned_ev["dr_mujet_softmu"] = ssmu[:, 0].delta_r(smuon_jet) - pruned_ev["dr_mujet_lep1"] = shmu[:, 0].delta_r(smuon_jet) - pruned_ev["dr_mujet_lep2"] = shmu[:, 1].delta_r(smuon_jet) - pruned_ev["dr_lep1_softmu"] = shmu[:, 0].delta_r(ssmu[:, 0]) - pruned_ev["soft_l_ptratio"] = ssmu[:, 0].pt / smuon_jet.pt - pruned_ev["l1_ptratio"] = shmu[:, 0].pt / smuon_jet.pt - pruned_ev["l2_ptratio"] = shmu[:, 1].pt / smuon_jet.pt - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where( - (out_branch == "SoftMuon") - | (out_branch == "MuonJet") - | (out_branch == "dilep") - ), - ) - - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in ["Muon", "Jet", "SoftMuon", "MuonJet", "MET"]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon" and obj != "SoftMuon") and ( - "pfRelIso04_all" == kin or "d" in kin - ): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) - return {dataset: output} - - def postprocess(self, accumulator): - return accumulator diff --git a/src/BTVNanoCommissioning/workflows/ctag_emdileptt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_emdileptt_valid_sf.py index 04215a4e..c56e86a1 100644 --- a/src/BTVNanoCommissioning/workflows/ctag_emdileptt_valid_sf.py +++ b/src/BTVNanoCommissioning/workflows/ctag_emdileptt_valid_sf.py @@ -17,11 +17,11 @@ from BTVNanoCommissioning.helpers.func import ( flatten, update, - uproot_writeable, dump_lumi, ) from BTVNanoCommissioning.helpers.update_branch import missing_branch from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import ( jet_id, mu_idiso, diff --git a/src/BTVNanoCommissioning/workflows/ctag_ettsemilep_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_ettsemilep_valid_sf.py deleted file mode 100644 index aa788623..00000000 --- a/src/BTVNanoCommissioning/workflows/ctag_ettsemilep_valid_sf.py +++ /dev/null @@ -1,563 +0,0 @@ -import gc, collections, pickle, os, sys, numpy as np, awkward as ak -import os -import uproot - -from coffea import processor -from coffea.analysis_tools import Weights - -from BTVNanoCommissioning.utils.correction import ( - load_lumi, - load_SF, - eleSFs, - puwei, - btagSFs, - JME_shifts, - Roccor_shifts, -) -from BTVNanoCommissioning.helpers.func import ( - flatten, - update, - uproot_writeable, - dump_lumi, -) -from BTVNanoCommissioning.helpers.update_branch import missing_branch -from BTVNanoCommissioning.utils.histogrammer import histogrammer -from BTVNanoCommissioning.utils.selection import ( - jet_id, - mu_idiso, - ele_mvatightid, - softmu_mask, -) - - -class NanoProcessor(processor.ProcessorABC): - # Define histograms - def __init__( - self, - year="2022", - campaign="Summer22Run3", - name="", - isSyst=False, - isArray=False, - noHist=False, - chunksize=75000, - ): - self._year = year - self._campaign = campaign - self.name = name - self.isSyst = isSyst - self.isArray = isArray - self.noHist = noHist - self.lumiMask = load_lumi(self._campaign) - self.chunksize = chunksize - ## Load corrections - self.SF_map = load_SF(self._campaign) - - @property - def accumulator(self): - return self._accumulator - - def process(self, events): - isRealData = not hasattr(events, "genWeight") - dataset = events.metadata["dataset"] - events = missing_branch(events) - shifts = [] - if "JME" in self.SF_map.keys(): - syst_JERC = self.isSyst - if self.isSyst == "JERC_split": - syst_JERC = "split" - shifts = JME_shifts( - shifts, self.SF_map, events, self._campaign, isRealData, syst_JERC - ) - else: - if int(self._year) < 2020: - shifts = [ - ({"Jet": events.Jet, "MET": events.MET, "Muon": events.Muon}, None) - ] - else: - shifts = [ - ( - { - "Jet": events.Jet, - "MET": events.PuppiMET, - "Muon": events.Muon, - }, - None, - ) - ] - if "roccor" in self.SF_map.keys(): - shifts = Roccor_shifts(shifts, self.SF_map, events, isRealData, False) - else: - shifts[0][0]["Muon"] = events.Muon - - return processor.accumulate( - self.process_shift(update(events, collections), name) - for collections, name in shifts - ) - - def process_shift(self, events, shift_name): - dataset = events.metadata["dataset"] - isRealData = not hasattr(events, "genWeight") - _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ectag_ttsemilep_sf") - ) - if _hist_event_dict == None: - _hist_event_dict[""]: None - output = { - "sumw": processor.defaultdict_accumulator(float), - **_hist_event_dict, - } - - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) - - #################### - # Selections # - #################### - ## Lumimask - req_lumi = np.ones(len(events), dtype="bool") - if isRealData: - req_lumi = self.lumiMask(events.run, events.luminosityBlock) - # only dump for nominal case - if shift_name is None: - output = dump_lumi(events[req_lumi], output) - - ## HLT - triggers = ["Ele32_WPTight_Gsf_L1DoubleEG"] - checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) - if ak.all(checkHLT == False): - raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) - elif ak.any(checkHLT == False): - print(np.array(triggers)[~checkHLT], " not exist in", dataset) - trig_arrs = [ - events.HLT[_trig] for _trig in triggers if hasattr(events.HLT, _trig) - ] - req_trig = np.zeros(len(events), dtype="bool") - for t in trig_arrs: - req_trig = req_trig | t - - ## Electron cuts - iso_ele = events.Electron[ - (events.Electron.pt > 34) & ele_mvatightid(events, self._campaign) - ] - req_ele = ak.count(iso_ele.pt, axis=1) == 1 - iso_ele = ak.pad_none(iso_ele, 1, axis=1) - iso_ele = iso_ele[:, 0] - iso_eindx = ak.mask( - ak.local_index(events.Electron.pt), - ((events.Electron.pt > 34) & ele_mvatightid(events, self._campaign)) == 1, - ) - iso_eindx = ak.pad_none(iso_eindx, 1) - iso_eindx = iso_eindx[:, 0] - ## Jet cuts - event_jet = events.Jet[ - ak.fill_none( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(iso_ele) > 0.5, - axis=2, - mask_identity=True, - ) - ), - False, - axis=-1, - ) - ] - req_jets = ak.num(event_jet.pt) >= 4 - - ## Soft Muon cuts - soft_muon = events.Muon[softmu_mask(events, self._campaign)] - req_softmu = ak.count(soft_muon.pt, axis=1) >= 1 - soft_muon = ak.pad_none(soft_muon, 1, axis=1) - - ## Muon-jet cuts - mu_jet = event_jet[ - ak.fill_none( - ( - ak.all( - event_jet.metric_table(soft_muon) <= 0.4, - axis=2, - mask_identity=True, - ) - ) - & ((event_jet.muonIdx1 != -1) | (event_jet.muonIdx2 != -1)), - False, - axis=-1, - ) - ] - req_mujet = ak.num(mu_jet.pt, axis=1) >= 1 - mu_jet = ak.pad_none(mu_jet, 1, axis=1) - - ## store jet index for PFCands, create mask on the jet index - jetindx = ak.mask( - ak.local_index(events.Jet.pt), - ( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(iso_ele) > 0.5, - axis=2, - mask_identity=True, - ) - ) - & ( - ak.all( - events.Jet.metric_table(soft_muon) <= 0.4, - axis=2, - mask_identity=True, - ) - ) - & ((events.Jet.muonIdx1 != -1) | (events.Jet.muonIdx2 != -1)) - ) - == 1, - ) - jetindx = ak.pad_none(jetindx, 1) - jetindx = jetindx[:, 0] - - # Other cuts - req_pTratio = (soft_muon[:, 0].pt / mu_jet[:, 0].pt) < 0.6 - - req_QCDveto = ( - (iso_ele.pfRelIso03_all < 0.05) - & (abs(iso_ele.dz) < 0.02) - & (abs(iso_ele.dxy) < 0.01) - & (iso_ele.sip3d < 2.5) - & ( - iso_ele.pt - / ak.firsts( - events.Jet[ - (events.Jet.electronIdx1 == iso_eindx) - | ((events.Jet.electronIdx2 == iso_eindx)) - ].pt - ) - > 0.75 - ) - ) - - dilep_mu = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] - dilep_ele = events.Electron[ - (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) - ] - req_dilepveto = ( - ak.count(dilep_mu.pt, axis=1) + ak.count(dilep_ele.pt, axis=1) != 2 - ) - - MET = ak.zip( - { - "pt": events.MET.pt, - "eta": ak.zeros_like(events.MET.pt), - "phi": events.MET.phi, - "mass": ak.zeros_like(events.MET.pt), - }, - with_name="PtEtaPhiMLorentzVector", - ) - Wmass = MET + iso_ele - req_Wmass = Wmass.mass > 55 - - event_level = ( - req_trig - & req_lumi - & req_ele - & req_jets - & req_softmu - & req_mujet - & req_Wmass - & req_dilepveto - & req_QCDveto - & req_pTratio - ) - event_level = ak.fill_none(event_level, False) - if len(events[event_level]) == 0: - return {dataset: output} - - #################### - # Selected objects # - #################### - shmu = iso_ele[event_level] - sjets = event_jet[event_level] - ssmu = soft_muon[event_level] - smet = MET[event_level] - smuon_jet = mu_jet[event_level] - nsoftmu = ak.count(ssmu.pt, axis=1) - nmujet = ak.count(smuon_jet.pt, axis=1) - smuon_jet = smuon_jet[:, 0] - ssmu = ssmu[:, 0] - sz = shmu + ssmu - sw = shmu + smet - njet = ak.count(sjets.pt, axis=1) - if "PFCands" in events.fields: - spfcands = events[event_level].PFCands[ - events[event_level] - .JetPFCands[ - events[event_level].JetPFCands.jetIdx == jetindx[event_level] - ] - .pFCandsIdx - ] - - #################### - # Weight & Geninfo # - #################### - weights = Weights(len(events[event_level]), storeIndividual=True) - if not isRealData: - weights.add("genweight", events[event_level].genWeight) - genflavor = sjets.hadronFlavour + 1 * ( - (sjets.partonFlavour == 0) & (sjets.hadronFlavour == 0) - ) - smflav = ( - 1 * ((smuon_jet.partonFlavour == 0) & (smuon_jet.hadronFlavour == 0)) - + smuon_jet.hadronFlavour - ) - if len(self.SF_map.keys()) > 0: - syst_wei = True if self.isSyst != False else False - if "PU" in self.SF_map.keys(): - puwei( - events[event_level].Pileup.nTrueInt, - self.SF_map, - weights, - syst_wei, - ) - if "EGM" in self.SF_map.keys(): - eleSFs(shmu, self.SF_map, weights, syst_wei, False) - if "BTV" in self.SF_map.keys(): - btagSFs(smuon_jet, self.SF_map, weights, "DeepJetC", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepJetB", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepCSVB", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepCSVC", syst_wei) - - else: - genflavor = ak.zeros_like(sjets.pt, dtype=int) - smflav = ak.zeros_like(smuon_jet.pt, dtype=int) - - # Systematics information - if shift_name is None: - systematics = ["nominal"] + list(weights.variations) - else: - systematics = [shift_name] - exclude_btv = [ - "DeepCSVC", - "DeepCSVB", - "DeepJetB", - "DeepJetC", - ] # exclude b-tag SFs for btag inputs - - #################### - # Fill histogram # - #################### - for syst in systematics: - if self.isSyst == False and syst != "nominal": - break - if self.noHist: - break - weight = ( - weights.weight() - if syst == "nominal" or syst == shift_name - else weights.weight(modifier=syst) - ) - for histname, h in output.items(): - if ( - "Deep" in histname - and "btag" not in histname - and histname in events.Jet.fields - ): - h.fill( - syst, - flatten(genflavor), - flatten(sjets[histname]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), sjets["pt"] - )[0] - ), - ) - elif "jet_" in histname and "mu" not in histname: - h.fill( - syst, - flatten(genflavor), - flatten(sjets[histname.replace("jet_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), sjets["pt"] - )[0] - ), - ) - elif "hl_" in histname and histname.replace("hl_", "") in shmu.fields: - h.fill( - syst, - flatten(shmu[histname.replace("hl_", "")]), - weight=weight, - ) - elif ( - "soft_l" in histname - and histname.replace("soft_l_", "") in ssmu.fields - ): - h.fill( - syst, - smflav, - flatten(ssmu[histname.replace("soft_l_", "")]), - weight=weight, - ) - elif "mujet_" in histname: - h.fill( - syst, - smflav, - flatten(smuon_jet[histname.replace("mujet_", "")]), - weight=weight, - ) - elif ( - "PFCands" in events.fields - and "PFCands" in histname - and histname.split("_")[1] in events.PFCands.fields - ): - h.fill( - syst, - flatten(ak.broadcast_arrays(smflav, spfcands["pt"])[0]), - flatten(spfcands[histname.replace("PFCands_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), - spfcands["pt"], - )[0] - ), - ) - elif "btag" in histname: - for i in range(2): - if ( - str(i) not in histname - or histname.replace(f"_{i}", "") not in events.Jet.fields - ): - continue - h.fill( - syst="noSF", - flav=smflav, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weights.partial_weight(exclude=exclude_btv), - ) - if not isRealData and "btag" in self.SF_map.keys(): - h.fill( - syst=syst, - flav=smflav, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weight, - ) - output["njet"].fill(syst, njet, weight=weight) - output["nmujet"].fill(syst, nmujet, weight=weight) - output["nsoftmu"].fill(syst, nsoftmu, weight=weight) - output["hl_ptratio"].fill( - syst, - genflavor[:, 0], - ratio=shmu.pt / sjets[:, 0].pt, - weight=weight, - ) - output["soft_l_ptratio"].fill( - syst, - flav=smflav, - ratio=ssmu.pt / smuon_jet.pt, - weight=weight, - ) - output["dr_lmujetsmu"].fill( - syst, - flav=smflav, - dr=smuon_jet.delta_r(ssmu), - weight=weight, - ) - output["dr_lmujethmu"].fill( - syst, - flav=smflav, - dr=smuon_jet.delta_r(shmu), - weight=weight, - ) - output["dr_lmusmu"].fill(syst, dr=shmu.delta_r(ssmu), weight=weight) - output["z_pt"].fill(syst, flatten(sz.pt), weight=weight) - output["z_eta"].fill(syst, flatten(sz.eta), weight=weight) - output["z_phi"].fill(syst, flatten(sz.phi), weight=weight) - output["z_mass"].fill(syst, flatten(sz.mass), weight=weight) - output["w_pt"].fill(syst, flatten(sw.pt), weight=weight) - output["w_eta"].fill(syst, flatten(sw.eta), weight=weight) - output["w_phi"].fill(syst, flatten(sw.phi), weight=weight) - output["w_mass"].fill(syst, flatten(sw.mass), weight=weight) - output["MET_pt"].fill(syst, flatten(smet.pt), weight=weight) - output["MET_phi"].fill(syst, flatten(smet.phi), weight=weight) - output["npvs"].fill( - syst, - events[event_level].PV.npvs, - weight=weight, - ) - if not isRealData: - output["pu"].fill( - syst, - events[event_level].Pileup.nTrueInt, - weight=weight, - ) - ####################### - # Create root files # - ####################### - if self.isArray: - # Keep the structure of events and pruned the object size - pruned_ev = events[event_level] - pruned_ev.Jet = sjets - pruned_ev.Muon = shmu - pruned_ev["MuonJet"] = smuon_jet - pruned_ev["SoftMuon"] = ssmu - if "PFCands" in events.fields: - pruned_ev.PFCands = spfcands - # Add custom variables - if not isRealData: - pruned_ev["weight"] = weights.weight() - for ind_wei in weights.weightStatistics.keys(): - pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( - include=[ind_wei] - ) - - pruned_ev["dr_mujet_softmu"] = ssmu.delta_r(smuon_jet) - pruned_ev["dr_mujet_lep1"] = shmu.delta_r(smuon_jet) - pruned_ev["dr_lep1_softmu"] = shmu.delta_r(ssmu) - pruned_ev["soft_l_ptratio"] = ssmu.pt / smuon_jet.pt - pruned_ev["l1_ptratio"] = shmu.pt / smuon_jet.pt - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where( - (out_branch == "SoftMuon") - | (out_branch == "MuonJet") - | (out_branch == "dilep") - ), - ) - - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in [ - "Muon", - "Jet", - "SoftMuon", - "MuonJet", - "dilep", - "charge", - "MET", - ]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon" and obj != "SoftMuon") and ( - "pfRelIso04_all" == kin or "d" in kin - ): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*", "SV_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) - return {dataset: output} - - def postprocess(self, accumulator): - return accumulator diff --git a/src/BTVNanoCommissioning/workflows/ctag_semileptt_valid_sf.py b/src/BTVNanoCommissioning/workflows/ctag_semileptt_valid_sf.py deleted file mode 100644 index 361bffae..00000000 --- a/src/BTVNanoCommissioning/workflows/ctag_semileptt_valid_sf.py +++ /dev/null @@ -1,580 +0,0 @@ -import collections, awkward as ak, numpy as np -import os -import uproot - -from coffea import processor -from coffea.analysis_tools import Weights - -from BTVNanoCommissioning.utils.correction import ( - load_lumi, - load_SF, - muSFs, - puwei, - btagSFs, - JME_shifts, - Roccor_shifts, -) -from BTVNanoCommissioning.helpers.func import ( - flatten, - update, - uproot_writeable, - dump_lumi, -) -from BTVNanoCommissioning.helpers.update_branch import missing_branch -from BTVNanoCommissioning.utils.histogrammer import histogrammer -from BTVNanoCommissioning.utils.selection import ( - jet_id, - mu_idiso, - ele_mvatightid, - softmu_mask, -) - - -class NanoProcessor(processor.ProcessorABC): - def __init__( - self, - year="2022", - campaign="Summer22Run3", - name="", - isSyst=False, - isArray=False, - noHist=False, - chunksize=75000, - ): - self._year = year - self._campaign = campaign - self.name = name - self.isSyst = isSyst - self.isArray = isArray - self.noHist = noHist - self.lumiMask = load_lumi(self._campaign) - self.chunksize = chunksize - ## Load corrections - self.SF_map = load_SF(self._campaign) - - @property - def accumulator(self): - return self._accumulator - - def process(self, events): - isRealData = not hasattr(events, "genWeight") - dataset = events.metadata["dataset"] - events = missing_branch(events) - shifts = [] - if "JME" in self.SF_map.keys(): - syst_JERC = self.isSyst - if self.isSyst == "JERC_split": - syst_JERC = "split" - shifts = JME_shifts( - shifts, self.SF_map, events, self._campaign, isRealData, syst_JERC - ) - else: - if int(self._year) < 2020: - shifts = [ - ({"Jet": events.Jet, "MET": events.MET, "Muon": events.Muon}, None) - ] - else: - shifts = [ - ( - { - "Jet": events.Jet, - "MET": events.PuppiMET, - "Muon": events.Muon, - }, - None, - ) - ] - if "roccor" in self.SF_map.keys(): - shifts = Roccor_shifts(shifts, self.SF_map, events, isRealData, False) - else: - shifts[0][0]["Muon"] = events.Muon - - return processor.accumulate( - self.process_shift(update(events, collections), name) - for collections, name in shifts - ) - - def process_shift(self, events, shift_name): - dataset = events.metadata["dataset"] - isRealData = not hasattr(events, "genWeight") - _hist_event_dict = ( - {"": None} if self.noHist else histogrammer(events, "ctag_ttsemilep_sf") - ) - - output = { - "sumw": processor.defaultdict_accumulator(float), - **_hist_event_dict, - } - - if shift_name is None: - if isRealData: - output["sumw"] = len(events) - else: - output["sumw"] = ak.sum(events.genWeight) - - ############### - # Event level # - ############### - ## Lumimask - req_lumi = np.ones(len(events), dtype="bool") - if isRealData: - req_lumi = self.lumiMask(events.run, events.luminosityBlock) - # only dump for nominal case - if shift_name is None: - output = dump_lumi(events[req_lumi], output) - - ## HLT - triggers = ["IsoMu24", "IsoMu27"] - checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers]) - if ak.all(checkHLT == False): - raise ValueError("HLT paths:", triggers, " are all invalid in", dataset) - elif ak.any(checkHLT == False): - print(np.array(triggers)[~checkHLT], " not exist in", dataset) - trig_arrs = [ - events.HLT[_trig] for _trig in triggers if hasattr(events.HLT, _trig) - ] - req_trig = np.zeros(len(events), dtype="bool") - for t in trig_arrs: - req_trig = req_trig | t - - ## Muon cuts - # muon twiki: https://twiki.cern.ch/twiki/bin/view/CMS/SWGuideMuonIdRun2 - iso_muon = events.Muon[(events.Muon.pt > 30) & mu_idiso(events, self._campaign)] - req_muon = ak.count(iso_muon.pt, axis=1) == 1 - iso_muon = ak.pad_none(iso_muon, 1, axis=1) - iso_muon = iso_muon[:, 0] - iso_muindx = ak.mask( - ak.local_index(events.Muon.pt), - ((events.Muon.pt > 30) & mu_idiso(events, self._campaign)) == 1, - ) - iso_muindx = ak.pad_none(iso_muindx, 1) - iso_muindx = iso_muindx[:, 0] - - ## Jet cuts - event_jet = events.Jet[ - ak.fill_none( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(iso_muon) > 0.5, - axis=2, - mask_identity=True, - ) - ), - False, - axis=-1, - ) - ] - req_jets = ak.num(event_jet.pt) >= 4 - - ## Soft Muon cuts - - soft_muon = events.Muon[softmu_mask(events, self._campaign)] - req_softmu = ak.count(soft_muon.pt, axis=1) >= 1 - soft_muon = ak.pad_none(soft_muon, 1, axis=1) - - ## Muon-jet cuts - mu_jet = event_jet[ - ak.fill_none( - ( - ak.all( - event_jet.metric_table(soft_muon) <= 0.4, - axis=2, - mask_identity=True, - ) - ) - & ((event_jet.muonIdx1 != -1) | (event_jet.muonIdx2 != -1)) - & ((event_jet.muEF + event_jet.neEmEF) < 0.7), - False, - axis=-1, - ) - ] - req_mujet = ak.num(mu_jet.pt, axis=1) >= 1 - mu_jet = ak.pad_none(mu_jet, 1, axis=1) - - ## store jet index for PFCands, create mask on the jet index - jetindx = ak.mask( - ak.local_index(events.Jet.pt), - ( - jet_id(events, self._campaign) - & ( - ak.all( - events.Jet.metric_table(iso_muon) > 0.5, - axis=2, - mask_identity=True, - ) - ) - & ((events.Jet.muEF + events.Jet.neEmEF) < 0.7) - & ( - ak.all( - events.Jet.metric_table(iso_muon) > 0.5, - axis=2, - mask_identity=True, - ) - ) - & ( - ak.all( - events.Jet.metric_table(soft_muon) <= 0.4, - axis=2, - mask_identity=True, - ) - ) - & ((events.Jet.muonIdx1 != -1) | (events.Jet.muonIdx2 != -1)) - ) - == 1, - ) - jetindx = ak.pad_none(jetindx, 1) - jetindx = jetindx[:, 0] - - ## Other cuts - req_pTratio = (soft_muon[:, 0].pt / mu_jet[:, 0].pt) < 0.4 - - req_QCDveto = ( - (iso_muon.pfRelIso04_all < 0.05) - & (abs(iso_muon.dz) < 0.01) - & (abs(iso_muon.dxy) < 0.002) - & (iso_muon.sip3d < 2) - & ( - iso_muon.pt - / ak.firsts( - events.Jet[ - (events.Jet.muonIdx1 == iso_muindx) - | ((events.Jet.muonIdx2 == iso_muindx)) - ].pt - ) - > 0.75 - ) - ) - - dilep_mu = events.Muon[(events.Muon.pt > 12) & mu_idiso(events, self._campaign)] - dilep_ele = events.Electron[ - (events.Electron.pt > 15) & ele_mvatightid(events, self._campaign) - ] - req_dilepveto = ( - ak.count(dilep_mu.pt, axis=1) + ak.count(dilep_ele.pt, axis=1) != 2 - ) - - dilep_mass = iso_muon + soft_muon[:, 0] - req_dilepmass = (dilep_mass.mass > 12.0) & ( - (dilep_mass.mass < 80) | (dilep_mass.mass > 100) - ) - MET = ak.zip( - { - "pt": events.MET.pt, - "eta": ak.zeros_like(events.MET.pt), - "phi": events.MET.phi, - "mass": ak.zeros_like(events.MET.pt), - }, - with_name="PtEtaPhiMLorentzVector", - ) - Wmass = MET + iso_muon - req_Wmass = Wmass.mass > 55 - - event_level = ( - req_trig - & req_lumi - & req_muon - & req_jets - & req_softmu - & req_dilepmass - & req_mujet - & req_Wmass - & req_dilepveto - & req_QCDveto - & req_pTratio - ) - - event_level = ak.fill_none(event_level, False) - if len(events[event_level]) == 0: - return {dataset: output} - - #################### - # Selected objects # - #################### - shmu = iso_muon[event_level] - sjets = event_jet[event_level] - ssmu = soft_muon[event_level] - smet = MET[event_level] - smuon_jet = mu_jet[event_level] - nsoftmu = ak.count(ssmu.pt, axis=1) - nmujet = ak.count(smuon_jet.pt, axis=1) - smuon_jet = smuon_jet[:, 0] - ssmu = ssmu[:, 0] - sz = shmu + ssmu - sw = shmu + smet - osss = ak.values_astype(shmu.charge * ssmu.charge * -1, int) - njet = ak.count(sjets.pt, axis=1) - # Find the PFCands associate with selected jets. Search from jetindex->JetPFCands->PFCand - if "PFCands" in events.fields: - spfcands = events[event_level].PFCands[ - events[event_level] - .JetPFCands[ - events[event_level].JetPFCands.jetIdx == jetindx[event_level] - ] - .pFCandsIdx - ] - - #################### - # Weight & Geninfo # - #################### - weights = Weights(len(events[event_level]), storeIndividual=True) - if not isRealData: - weights.add("genweight", events[event_level].genWeight) - genflavor = sjets.hadronFlavour + 1 * ( - (sjets.partonFlavour == 0) & (sjets.hadronFlavour == 0) - ) - smflav = smuon_jet.hadronFlavour + 1 * ( - (smuon_jet.partonFlavour == 0) & (smuon_jet.hadronFlavour == 0) - ) - if len(self.SF_map.keys()) > 0: - syst_wei = True if self.isSyst != False else False - if "PU" in self.SF_map.keys(): - puwei( - events[event_level].Pileup.nTrueInt, - self.SF_map, - weights, - syst_wei, - ) - if "MUO" in self.SF_map.keys(): - muSFs(shmu, self.SF_map, weights, syst_wei, False) - if "BTV" in self.SF_map.keys(): - btagSFs(smuon_jet, self.SF_map, weights, "DeepJetC", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepJetB", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepCSVB", syst_wei) - btagSFs(smuon_jet, self.SF_map, weights, "DeepCSVC", syst_wei) - - else: - genflavor = ak.zeros_like(sjets.pt, dtype=int) - smflav = ak.zeros_like(smuon_jet.pt, dtype=int) - - # Systematics information - if shift_name is None: - systematics = ["nominal"] + list(weights.variations) - else: - systematics = [shift_name] - exclude_btv = [ - "DeepCSVC", - "DeepCSVB", - "DeepJetB", - "DeepJetC", - ] # exclude b-tag SFs for btag inputs - - #################### - # Fill histogram # - #################### - for syst in systematics: - if self.isSyst == False and syst != "nominal": - break - if self.noHist: - break - weight = ( - weights.weight() - if syst == "nominal" or syst == shift_name - else weights.weight(modifier=syst) - ) - for histname, h in output.items(): - if ( - "Deep" in histname - and "btag" not in histname - and histname in events.Jet.fields - ): - h.fill( - syst, - flatten(genflavor), - flatten(sjets[histname]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), sjets["pt"] - )[0] - ), - ) - elif "jet_" in histname and "mu" not in histname: - h.fill( - syst, - flatten(genflavor), - flatten(sjets[histname.replace("jet_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), sjets["pt"] - )[0] - ), - ) - elif "hl_" in histname and histname.replace("hl_", "") in shmu.fields: - h.fill( - syst, - flatten(shmu[histname.replace("hl_", "")]), - weight=weight, - ) - elif ( - "soft_l" in histname - and histname.replace("soft_l_", "") in ssmu.fields - ): - print(histname) - h.fill( - syst, - smflav, - flatten(ssmu[histname.replace("soft_l_", "")]), - weight=weight, - ) - elif "mujet_" in histname: - h.fill( - syst, - smflav, - flatten(smuon_jet[histname.replace("mujet_", "")]), - weight=weight, - ) - elif ( - "PFCands" in events.fields - and "PFCands" in histname - and histname.split("_")[1] in events.PFCands.fields - ): - h.fill( - syst, - flatten(ak.broadcast_arrays(smflav, spfcands["pt"])[0]), - flatten(spfcands[histname.replace("PFCands_", "")]), - weight=flatten( - ak.broadcast_arrays( - weights.partial_weight(exclude=exclude_btv), - spfcands["pt"], - )[0] - ), - ) - elif "btag" in histname: - for i in range(2): - if ( - str(i) not in histname - or histname.replace(f"_{i}", "") not in events.Jet.fields - ): - continue - h.fill( - syst="noSF", - flav=smflav, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weights.partial_weight(exclude=exclude_btv), - ) - if not isRealData and "btag" in self.SF_map.keys(): - h.fill( - syst=syst, - flav=smflav, - discr=smuon_jet[histname.replace(f"_{i}", "")], - weight=weight, - ) - output["njet"].fill(syst, njet, weight=weight) - output["nmujet"].fill(syst, nmujet, weight=weight) - output["nsoftmu"].fill(syst, nsoftmu, weight=weight) - output["hl_ptratio"].fill( - syst, - genflavor[:, 0], - ratio=shmu.pt / sjets[:, 0].pt, - weight=weight, - ) - output["soft_l_ptratio"].fill( - syst, - flav=smflav, - ratio=ssmu.pt / smuon_jet.pt, - weight=weight, - ) - output["dr_lmujetsmu"].fill( - syst, - flav=smflav, - dr=smuon_jet.delta_r(ssmu), - weight=weight, - ) - output["dr_lmujethmu"].fill( - syst, - flav=smflav, - dr=smuon_jet.delta_r(shmu), - weight=weight, - ) - output["dr_lmusmu"].fill(syst, dr=shmu.delta_r(ssmu), weight=weight) - output["z_pt"].fill(syst, flatten(sz.pt), weight=weight) - output["z_eta"].fill(syst, flatten(sz.eta), weight=weight) - output["z_phi"].fill(syst, flatten(sz.phi), weight=weight) - output["z_mass"].fill(syst, flatten(sz.mass), weight=weight) - output["w_pt"].fill(syst, flatten(sw.pt), weight=weight) - output["w_eta"].fill(syst, flatten(sw.eta), weight=weight) - output["w_phi"].fill(syst, flatten(sw.phi), weight=weight) - output["w_mass"].fill(syst, flatten(sw.mass), weight=weight) - output["MET_pt"].fill(syst, flatten(smet.pt), weight=weight) - output["MET_phi"].fill(syst, flatten(smet.phi), weight=weight) - output["npvs"].fill( - syst, - events[event_level].PV.npvs, - weight=weight, - ) - if not isRealData: - output["pu"].fill( - syst, - events[event_level].Pileup.nTrueInt, - weight=weight, - ) - ####################### - # Create root files # - ####################### - if self.isArray: - # Keep the structure of events and pruned the object size - pruned_ev = events[event_level] - pruned_ev.Jet = sjets - pruned_ev.Muon = shmu - pruned_ev["MuonJet"] = smuon_jet - pruned_ev["SoftMuon"] = ssmu - - if "PFCands" in events.fields: - pruned_ev.PFCands = spfcands - # Add custom variables - if not isRealData: - pruned_ev["weight"] = weights.weight() - for ind_wei in weights.weightStatistics.keys(): - pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( - include=[ind_wei] - ) - - pruned_ev["dr_mujet_softmu"] = ssmu.delta_r(smuon_jet) - pruned_ev["dr_mujet_lep1"] = shmu.delta_r(smuon_jet) - pruned_ev["dr_lep1_softmu"] = shmu.delta_r(ssmu) - pruned_ev["soft_l_ptratio"] = ssmu.pt / smuon_jet.pt - pruned_ev["l1_ptratio"] = shmu.pt / smuon_jet.pt - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - out_branch = np.delete( - out_branch, - np.where( - (out_branch == "SoftMuon") - | (out_branch == "MuonJet") - | (out_branch == "dilep") - ), - ) - - for kin in ["pt", "eta", "phi", "mass", "pfRelIso04_all", "dxy", "dz"]: - for obj in [ - "Muon", - "Jet", - "SoftMuon", - "MuonJet", - "dilep", - "charge", - "MET", - ]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon" and obj != "SoftMuon") and ( - "pfRelIso04_all" == kin or "d" in kin - ): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*", "SV_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) - return {dataset: output} - - def postprocess(self, accumulator): - return accumulator diff --git a/src/BTVNanoCommissioning/workflows/example.py b/src/BTVNanoCommissioning/workflows/example.py index 3ca91318..d50b2873 100644 --- a/src/BTVNanoCommissioning/workflows/example.py +++ b/src/BTVNanoCommissioning/workflows/example.py @@ -27,6 +27,7 @@ ## load histograms & selctions for this workflow from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import jet_id, mu_idiso, ele_cuttightid @@ -264,7 +265,14 @@ def process_shift(self, events, shift_name): if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] # pruned events + pruned_ev.Muon = smu # replace muon collections with selected muon + if self.isArray: + # Keep the structure of events and pruned the object size + pruned_ev = events[event_level] + pruned_ev["SelJet"] = sjets + pruned_ev["Muon"] = smu + # Add custom variables if not isRealData: pruned_ev["weight"] = weights.weight() @@ -272,17 +280,8 @@ def process_shift(self, events, shift_name): pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( include=[ind_wei] ) - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) # stored customed variables - out_branch = np.append(out_branch, ["Jet_btagDeep*", "Muon_pt"]) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) + array_writer(self, pruned_ev, events, systematics[0], dataset, isRealData) + return {dataset: output} ## post process, return the accumulator, compressed diff --git a/src/BTVNanoCommissioning/workflows/ttdilep_valid_sf.py b/src/BTVNanoCommissioning/workflows/ttdilep_valid_sf.py index fff6a1b4..10bdc3d0 100644 --- a/src/BTVNanoCommissioning/workflows/ttdilep_valid_sf.py +++ b/src/BTVNanoCommissioning/workflows/ttdilep_valid_sf.py @@ -20,13 +20,13 @@ from BTVNanoCommissioning.helpers.func import ( flatten, update, - uproot_writeable, dump_lumi, ) from BTVNanoCommissioning.helpers.update_branch import missing_branch ## load histograms & selctions for this workflow from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import ( jet_id, mu_idiso, @@ -223,6 +223,16 @@ def process_shift(self, events, shift_name): ) event_level = ak.fill_none(event_level, False) if len(events[event_level]) == 0: + if self.isArray: + array_writer( + self, + events[event_level], + events, + "nominal", + dataset, + isRealData, + empty=True, + ) return {dataset: output} #################### @@ -426,9 +436,9 @@ def process_shift(self, events, shift_name): if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] - pruned_ev.Jet = sjets - pruned_ev.Electron = sel - pruned_ev.Muon = smu + pruned_ev["SelJet"] = sjets + pruned_ev["Muon"] = smu + pruned_ev["Electron"] = sel if "PFCands" in events.fields: pruned_ev.PFCands = spfcands # Add custom variables @@ -441,32 +451,8 @@ def process_shift(self, events, shift_name): pruned_ev["dr_mujet0"] = smu.delta_r(sjets[:, 0]) pruned_ev["dr_mujet1"] = smu.delta_r(sjets[:, 1]) + array_writer(self, pruned_ev, events, systematics[0], dataset, isRealData) - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - for kin in ["pt", "eta", "phi", "mass", "dz", "dxy"]: - for obj in ["Jet", "Electron", "Muon"]: - if obj == "Jet" and "d" in kin: - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, - [ - "Jet_btagDeep*", - "Jet_DeepJet*", - "PFCands_*", - "Electron_pfRelIso03_all", - "Muon_pfRelIso03_all", - ], - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) return {dataset: output} def postprocess(self, accumulator): diff --git a/src/BTVNanoCommissioning/workflows/ttsemilep_valid_sf.py b/src/BTVNanoCommissioning/workflows/ttsemilep_valid_sf.py index e2868a45..eaee6736 100644 --- a/src/BTVNanoCommissioning/workflows/ttsemilep_valid_sf.py +++ b/src/BTVNanoCommissioning/workflows/ttsemilep_valid_sf.py @@ -23,6 +23,7 @@ ) from BTVNanoCommissioning.helpers.update_branch import missing_branch from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import jet_id, btag_mu_idiso, MET_filters import hist @@ -196,6 +197,16 @@ def process_shift(self, events, shift_name): req_trig & req_jets & req_muon & req_MET & req_lumi & req_metfilter, False ) if len(events[event_level]) == 0: + if self.isArray: + array_writer( + self, + events[event_level], + events, + "nominal", + dataset, + isRealData, + empty=True, + ) return {dataset: output} #################### # Selected objects # @@ -382,8 +393,8 @@ def process_shift(self, events, shift_name): if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] - pruned_ev.Jet = sjets - pruned_ev.Muon = smu + pruned_ev["SelJet"] = sjets + pruned_ev["Muon"] = smu if "PFCands" in events.fields: pruned_ev.PFCands = spfcands # Add custom variables @@ -396,28 +407,7 @@ def process_shift(self, events, shift_name): for i in range(4): pruned_ev[f"dr_mujet{i}"] = smu.delta_r(sjets[:, i]) - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - for kin in ["pt", "eta", "phi", "mass"]: # ,"pfRelIso04_all","dxy", "dz"]: - # for obj in ["Jet", "Muon", "MET","PuppiMET"]: - for obj in ["Muon"]: - if "MET" in obj and ("pt" != kin or "phi" != kin): - continue - if (obj != "Muon") and ("pfRelIso04_all" == kin or "d" in kin): - continue - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) + array_writer(self, pruned_ev, events, systematics[0], dataset, isRealData) return {dataset: output} def postprocess(self, accumulator): diff --git a/src/BTVNanoCommissioning/workflows/validation.py b/src/BTVNanoCommissioning/workflows/validation.py index 65b31842..42a16c7a 100644 --- a/src/BTVNanoCommissioning/workflows/validation.py +++ b/src/BTVNanoCommissioning/workflows/validation.py @@ -21,12 +21,12 @@ from BTVNanoCommissioning.helpers.func import ( flatten, update, - uproot_writeable, dump_lumi, ) from BTVNanoCommissioning.helpers.update_branch import missing_branch from BTVNanoCommissioning.utils.histogrammer import histogrammer +from BTVNanoCommissioning.utils.array_writer import array_writer from BTVNanoCommissioning.utils.selection import ( jet_id, mu_idiso, @@ -371,7 +371,8 @@ def process_shift(self, events, shift_name): if self.isArray: # Keep the structure of events and pruned the object size pruned_ev = events[event_level] - pruned_ev.Jet = sjets + pruned_ev["SelJet"] = sjets + # Add custom variables if not isRealData: pruned_ev["weight"] = weights.weight() @@ -379,23 +380,7 @@ def process_shift(self, events, shift_name): pruned_ev[f"{ind_wei}_weight"] = weights.partial_weight( include=[ind_wei] ) - - # Create a list of variables want to store. For objects from the PFNano file, specify as {object}_{variable}, wildcard option only accepted at the end of the string - out_branch = np.setdiff1d( - np.array(pruned_ev.fields), np.array(events.fields) - ) - for kin in ["pt", "eta", "phi", "mass"]: - for obj in ["Jet"]: - out_branch = np.append(out_branch, [f"{obj}_{kin}"]) - out_branch = np.append( - out_branch, ["Jet_btagDeep*", "Jet_DeepJet*", "PFCands_*"] - ) - # write to root files - os.system(f"mkdir -p {self.name}/{dataset}") - with uproot.recreate( - f"{self.name}/{dataset}/f{events.metadata['filename'].split('_')[-1].replace('.root','')}_{systematics[0]}_{int(events.metadata['entrystop']/self.chunksize)}.root" - ) as fout: - fout["Events"] = uproot_writeable(pruned_ev, include=out_branch) + array_writer(self, pruned_ev, events, systematics[0], dataset, isRealData) return {dataset: output} def postprocess(self, accumulator):