LumiDB/scripts/makePileupJSON.py

0001 #!/usr/bin/env python3
0002 from __future__ import print_function
0003 import argparse
0004 import RecoLuminosity.LumiDB.LumiConstants as LumiConstants
0005 import re
0006 from math import sqrt
0007
0008 ##############################
0009 ## ######################## ##
0010 ## ## ################## ## ##
0011 ## ## ## Main Program ## ## ##
0012 ## ## ################## ## ##
0013 ## ######################## ##
0014 ##############################
0015
0016 # This script takes a .csv file containing the per-BX luminosity values from brilcalc and processes them to
0017 # produce an output file containing the average and RMS pileup for each lumi section. This can then be fed
0018 # into pileupCalc.py to calculate a final pileup histogram. For more documentation see:
0019 # https://twiki.cern.ch/twiki/bin/viewauth/CMS/PileupJSONFileforData
0020
0021 # modified from the estimatePileup.py script in RecoLuminosity/LumiDB
0022 # originally 5 Jan, 2012  Mike Hildreth
0023
0024 if __name__ == '__main__':
0025     parameters = LumiConstants.ParametersObject()
0026
0027     parser = argparse.ArgumentParser(description="Script to estimate average and RMS pileup using the per-bunch luminosity information provided by brilcalc. The output is a JSON file containing a dictionary by runs with one entry per lumi section.")
0028     parser.add_argument('inputFile', help='CSV input file as produced from brilcalc')
0029     parser.add_argument('outputFile', help='Name of JSON output file')
0030     parser.add_argument('-b', '--selBX', metavar='BXLIST', help='Comma-separated list of BXs to use (will use all by default)')
0031     parser.add_argument('-n', '--no-threshold', action='store_true', help='By default, to avoid including spurious luminosity from afterglow in the pileup calculation, bunches with luminosity below a given threshold are excluded. This threshold is 8.0/ub/LS for HFOC at nominal energy, 2.0/ub/LS for HFOC at low energy, and 1.2/ub/LS for other luminometers. If the input data has already been preprocessed (e.g. by using the --xingTr argument in brilcalc) to exclude these bunches, or if you are running on special data with low overall luminosity, then use this flag to disable the application of the threshold.')
0032     args = parser.parse_args()
0033
0034     output = args.outputFile
0035
0036     sel_bx = set()
0037     if args.selBX:
0038         for ibx in args.selBX.split(","):
0039             try:
0040                 bx = int(ibx)
0041                 sel_bx.add(bx)
0042             except:
0043                 print(ibx,"is not an int")
0044         print("Processing",args.inputFile,"with selected BXs:", sorted(sel_bx))
0045     else:
0046         print("Processing",args.inputFile,"with all BX")
0047
0048     # The "CSV" file actually is a little complicated, since we also want to split on the colons separating
0049     # run/fill as well as the spaces separating the per-BX information.
0050     sepRE = re.compile(r'[\]\[\s,;:]+')
0051     csv_input = open(args.inputFile, 'r')
0052     last_run = -1
0053
0054     last_valid_lumi = []
0055
0056     output_line = '{'
0057     for line in csv_input:
0058         if line[0] == '#':
0059             continue # skip comment lines
0060
0061         pieces = sepRE.split(line.strip())
0062
0063         if len(pieces) < 15:
0064             # The most likely cause of this is that we're using a csv file without the bunch data, so might as well
0065             # just give up now.
0066             raise RuntimeError("Not enough fields in input line; maybe you forgot to include --xing in your brilcalc command?\n"+line)
0067         try:
0068             run = int(pieces[0])
0069             lumi_section = int(pieces[2])
0070             beam_energy = float(pieces[10])
0071             #tot_del_lumi = float(pieces[11])
0072             #tot_rec_lumi = float(pieces[12])
0073             luminometer = pieces[14]
0074
0075             if luminometer == "HFOC":
0076                 if beam_energy > 3000:
0077                     # Use higher threshold for nominal runs otherwise we'll get a lot of junk.
0078                     threshold = 8.0
0079                 else:
0080                     # Use lower threshold for 5.02 GeV runs since the overall lumi is quite low.
0081                     threshold = 2.0
0082             else:
0083                 threshold = 1.2
0084
0085             xing_lumi_array = []
0086             for bxid, bunch_del_lumi, bunch_rec_lumi in zip(pieces[15::3], pieces[16::3], pieces[17::3]):
0087                 if sel_bx and int(bxid) not in sel_bx:
0088                     continue
0089                 if args.no_threshold or float(bunch_del_lumi) > threshold:
0090                     xing_lumi_array.append([int(bxid), float(bunch_del_lumi), float(bunch_rec_lumi)])
0091         except:
0092             print("Failed to parse line: check if the input format has changed")
0093             print(pieces[0],pieces[1],pieces[2],pieces[3],pieces[4],pieces[5],pieces[6],pieces[7],pieces[8],pieces[9])
0094             continue
0095
0096         # In principle we could also have a check for if len(pieces) == 15 (i.e. no BX data is present) but
0097         # luminosity is present, which implies we're using a luminometer without BX data. In this case we
0098         # could just extrapolate from the previous good LS (just scaling the pileup by the ratio of
0099         # luminosity). In practice this is an extremely small number of lumi sections, and in 2018 the 14
0100         # lumisections in normtag_PHYSICS without BX data (all RAMSES) are all in periods with zero recorded
0101         # lumi, so they don't affect the resulting pileup at all. So for run 2 I haven't bothered to implement
0102         # this.
0103
0104         if run != last_run:
0105             # the script also used to add a dummy LS at the end of runs but this is not necessary in run 2
0106             if last_run > 0:
0107                 output_line = output_line[:-1] + '], '
0108             last_run = run
0109             output_line += ('\n"%d":' % run )
0110             output_line += ' ['
0111
0112         # Now do the actual parsing.
0113         total_lumi = 0
0114         total_int = 0
0115         total_int2 = 0
0116         total_weight2 = 0
0117
0118         # first loop to get sum for (weighted) mean
0119         for bxid, bunch_del_lumi, bunch_rec_lumi in xing_lumi_array:
0120             total_lumi += bunch_rec_lumi
0121             # this will eventually be_pileup*bunch_rec_lumi but it's quicker to apply the factor once outside the loop
0122             total_int += bunch_del_lumi*bunch_rec_lumi
0123         # filled_xings = len(xing_lumi_array)
0124
0125         # convert sum to pileup and get the mean
0126         total_int *= parameters.orbitLength / parameters.lumiSectionLength
0127         if total_lumi > 0:
0128             mean_int = total_int/total_lumi
0129         else:
0130             mean_int = 0
0131
0132         # second loop to get (weighted) RMS
0133         for bxid, bunch_del_lumi, bunch_rec_lumi in xing_lumi_array:
0134             mean_pileup = bunch_del_lumi * parameters.orbitLength / parameters.lumiSectionLength
0135             if mean_pileup > 100:
0136                 print("mean number of pileup events > 100 for run %d, lum %d : m %f l %f" % \
0137                       (runNumber, lumi_section, mean_pileup, bunch_del_lumi))
0138                 #print "mean number of pileup events for lum %d: m %f idx %d l %f" % (lumi_section, mean_pileup, bxid, bunch_rec_lumi)
0139
0140             total_int2 += bunch_rec_lumi*(mean_pileup-mean_int)*(mean_pileup-mean_int)
0141             total_weight2 += bunch_rec_lumi*bunch_rec_lumi
0142
0143         # compute final RMS and write it out
0144         #print " LS, Total lumi, filled xings %d, %f, %d" %(lumi_section,total_lumi,filled_xings)
0145         bunch_rms_lumi = 0
0146         denom = total_lumi*total_lumi-total_weight2
0147         if total_lumi > 0 and denom > 0:
0148             bunch_rms_lumi = sqrt(total_lumi*total_int2/denom)
0149
0150         output_line += "[%d,%2.4e,%2.4e,%2.4e]," % (lumi_section, total_lumi, bunch_rms_lumi, mean_int)
0151         last_valid_lumi = [lumi_section, total_lumi, bunch_rms_lumi, mean_int]
0152
0153     output_line = output_line[:-1] + ']}'
0154     csv_input.close()
0155
0156     outputfile = open(output,'w')
0157     if not outputfile:
0158         raise RuntimeError("Could not open '%s' as an output JSON file" % output)
0159
0160     outputfile.write(output_line)
0161     outputfile.close()
0162     print("Output written to", output)