EWP-BplusToKstMuMu-AngAna/Code/Ganga/UsefulScripts/getLFNs.py

import subprocess
from subprocess import Popen, PIPE


decay_number_dict = {
    "SigMC" : { 
        "2011": {
            "E": "3500",
            "nu": "2",
            "trigger": "0x40760037",
            "stripping": "21r1",
            "reco": "Reco14c",
            "simVer": "Sim09a",
            "dst": "DST",
            "merge": ""
        },
        "2012": {
            "E": "4000",
            "nu": "2.5",
            "trigger": "0x409f0045",
            "stripping": "21",
            "reco": "Reco14c",
            "simVer": "Sim09a",
            "dst": "DST",
            "merge": ""
        },
        "2015": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x411400a2",
            "stripping": "24r2",
            "reco":  "Reco15a/Turbo02",
            "simVer": "Sim09i",
            "dst": "DST",
            "merge": ""
        },
        "2016": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x6139160F",
            "stripping": "28r2",
            "reco":  "Reco16/Turbo03a",
            "simVer": "Sim09i",
            "dst": "DST",
            "merge": ""
        },
        "2017": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x62661709",
            "stripping": "29r2",
            "reco":  "Reco17/Turbo04a-WithTurcal",
            "simVer": "Sim09e",
            "dst": "MDST",
            "merge": ""
        },
        "2018": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x617d18a4",
            "stripping": "34r0p1", #34 for 12115102
            "reco":  "Reco18/Turbo05-WithTurcal", 
            "simVer": "Sim09h", #Sim09f for 12115102
            "dst": "DST",
            "merge": ""
        },
        "general":{
            "dec_ID": "12113100", #same for 12115102
            "filter": "NoPrescalingFlagged",
            "stream": "ALLSTREAMS",
            "pythiaVer": "8"
        }
    },
    "RefMC" : { 
        "2011": {
            "E": "3500",
            "nu": "2",
            "trigger": "0x40760037",
            "stripping": "21r1",
            "reco": "Reco14c",
            "simVer": "Sim09a",
            "dst": "DST",
            "merge": ""
        },
        "2012": {
            "E": "4000",
            "nu": "2.5",
            "trigger": "0x409f0045",
            "stripping": "21",
            "reco": "Reco14c",
            "simVer": "Sim09a",
            "dst": "DST",
            "merge": ""
        },
        "2015": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x411400a2",
            "stripping": "24r1",
            "reco":  "Reco15a/Turbo02",
            "simVer": "Sim09e",
            "dst": "DST",
            "merge": ""
        },
        "2016": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x6139160F",
            "stripping": "28r1",
            "reco":  "Reco16/Turbo03",
            "simVer": "Sim09e",
            "dst": "DST",
            "merge": ""
        },
        "general":{
            "dec_ID": "12143401", #same for 12145102
            "filter": "NoPrescalingFlagged",
            "stream": "ALLSTREAMS",
            "pythiaVer": "8"
        }
    },
    "PHSP" : {  #same for 12115179
        "2011": {
            "E": "3500",
            "nu": "2",
            "trigger": "0x40760037",
            "stripping": "21r1p1",
            "reco": "Reco14c",
            "simVer": "Sim09f",
            "dst": "DST",
            "merge": "Merge14/",
            "pythiaVer": "8"
        },
        "2012": {
            "E": "4000",
            "nu": "2.5",
            "trigger": "0x409f0045",
            "stripping": "21r0p1",
            "reco": "Reco14c",
            "simVer": "Sim09f",
            "dst": "DST",
            "merge": "Merge14/",
            "pythiaVer": "8"
        },
        "2015": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x411400a2",
            "stripping": "24r1",
            "reco":  "Reco15a/Turbo02",
            "simVer": "Sim09h",
            "dst": "MDST",
            "merge": "",
            "pythiaVer": "6"
        },
        "2016": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x6139160F",
            "stripping": "28r1",
            "reco":  "Reco16/Turbo03",
            "simVer": "Sim09f",
            "dst": "MDST",
            "merge": "Merge14/",
            "pythiaVer": "8"
        },
        "2017": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x62661709",
            "stripping": "29r2",
            "reco":  "Reco17/Turbo04a-WithTurcal",
            "simVer": "Sim09f",
            "dst": "MDST",
            "merge": "Merge14/",
            "pythiaVer": "8"
        },
        "2018": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "0x617d18a4",
            "stripping": "34",
            "reco":  "Reco18/Turbo05-WithTurcal",
            "simVer": "Sim09f",
            "dst": "MDST",
            "merge": "Merge14/",
            "pythiaVer": "8"
        },
        "general":{
            "dec_ID": "12113446", #same for 12115179
            "filter": "Filtered",
            "stream": "B2XMUMU.STRIP"
        }
    },
    "data"  : { 
         
        "11": {
            "E": "3500",
            "nu": "2",
            "trigger": "",
            "stripping": "21r1",
            "reco": "Reco14",
            "dst": "MDST"
        },
        "12": {
            "E": "4000",
            "nu": "2.5",
            "trigger": "",
            "stripping": "21r0p1a",
            "reco": "Reco14",
            "dst": "MDST"
        },
        "15": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "",
            "stripping": "24r2",
            "reco":  "Reco15a",
            "dst": "MDST"
        },
        "16": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "",
            "stripping": "28r2",
            "reco":  "Reco16",
            "dst": "MDST"
        },
        "17": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "",
            "stripping": "29r2",
            "reco":  "Reco17",
            "dst": "MDST"
        },
        "18": {
            "E": "6500",
            "nu": "1.6-25ns",
            "trigger": "",
            "stripping": "34r0p1", #34 for 12115102
            "reco":  "Reco18", 
            "dst": "MDST"
        },
        "general":{
            "dec_ID": "90000000", 
            "filter": "NoPrescalingFlagged",
            "stream": "ALLSTREAMS",
            "pythiaVer": "8",
            "merge": "",
            "simVer": "", 
            "trigger": ""

        }
   
    }
}

year_dict ={"SigMC": ["2011", "2012","2015", "2016", "2017", "2018"],
            "RefMC": ["2011", "2012","2015", "2016"],
            "PHSP":  ["2011", "2012","2015", "2016", "2017", "2018"],
            "data":  ["11", "12","15", "16", "17", "18"]
            }

def getPath(year, pol, sample):
    tmp_dict = dict(decay_number_dict[sample][year])
    gen_dict = dict(decay_number_dict[sample]["general"])
    tmp_dict.update(gen_dict)

    tmp_dict["year"] = year
    tmp_dict["pol"] = pol
    if (sample == "data"): #data
        PATH ="/LHCb/Collision%(year)s/Beam%(E)sGeV-VeloClosed-Mag%(pol)s/Real Data/%(reco)s/Stripping%(stripping)s/%(dec_ID)s/LEPTONIC.%(dst)s" % tmp_dict               
    else: #MC
        PATH = "/MC/%(year)s/Beam%(E)sGeV-%(year)s-Mag%(pol)s-Nu%(nu)s-Pythia%(pythiaVer)s/%(simVer)s/Trig%(trigger)s/%(reco)s/Stripping%(stripping)s%(filter)s/%(merge)s%(dec_ID)s/%(stream)s.%(dst)s" % tmp_dict
    print PATH
			
    OUTFILE = "%(dec_ID)s_%(year)s%(pol)s.py" % tmp_dict
    #print 'LFN output file: ', OUTFILE

    return PATH, OUTFILE


def getLFNs(nFiles,BKKpath):
  	
    #Call Dirac
    dirac_command = subprocess.Popen('lb-run LHCbDIRAC dirac-bookkeeping-get-files --BKQuery \"{0}\"'.format(BKKpath), shell=True, stdout=subprocess.PIPE)	
    #Read in Dirac output
    stdout = dirac_command.communicate()[0]
    #Split the output per line
    stdout = stdout.split('\n')

    #Create an empty list
    LFNlist= []

    #Check how many files are required
    if (nFiles == -1): max_file = len(stdout)-2
    else:              max_file =  min(nFiles+3,len(stdout)-2)

	#First 3 entries are junk from LHCbDIRAC dirac-bookkeeping-get-files
    for i in range(3, max_file+1):
        tmpLFN = [stdout[i].split('  ')][0]
        tmpLFN = tmpLFN[0]
        LFNlist.append(tmpLFN)
        
    return LFNlist

def writeLFNs(nFiles, year, pol, sample):

    #Get the paths
    BKKpath, LFNfile = getPath(year,pol, sample)

    #Get the LFNs from the paths
    lfn_list =  getLFNs(nFiles,BKKpath)    

    #Open the lfn target file and write all the lfns into it
    address_file = open(LFNfile,'w')
    address_file.write("""from Gaudi.Configuration import *\n""")
    address_file.write("""from GaudiConf import IOHelper\n""")
    address_file.write("""IOHelper('ROOT').inputFiles([\n""")
    for lfn in lfn_list:
        address_file.write("""'LFN:""" + lfn + """'\n""")
    address_file.write("""], clear=True)\n""")
    return

def getProductionID(year, pol, sample):
    #Exctract the production ID from LFN
    path = getPath(year,pol,sample)
    lfn = getLFNs(1,path[0])[0]
    lfn_split = lfn.split('/')
    return lfn_split[5]

def printAllProductionIDs(sample):
    for year in year_dict[sample]:
       print "Production ID for", sample, "year", year, "\t", getProductionID(year,"Down",sample)    
    return


def get_eff(year, pol, sample):
    import os 
    command = "lb-run LHCbDirac/prod dirac-bookkeeping-rejection-stats -B " + str(getPath(year,pol, sample)[0])
    print command
    os.system(command)


def get_tags(year, pol, sample):
    import os 
    #command = "lb-run LHCbDirac/prod dirac-bookkeeping-production-information " + str(getProductionID(year,pol, sample)) + " | grep CONDDB"
    #print command
    #os.system(command)

    command = "lb-run LHCbDirac/prod dirac-bookkeeping-production-information " + str(getProductionID(year,pol, sample)) + " | grep DDB"
    #print command2
    os.system(command)


#getPath("2011","Down", "SigMC")
#getPath("2012","Down", "SigMC")
#getPath("2015","Down", "SigMC")
getPath("2016","Down", "SigMC")
#getPath("2017","Down", "SigMC")
#getPath("2018","Down", "SigMC")

#getPath("2011","Down", "RefMC")
#getPath("2012","Down", "RefMC")
#getPath("2015","Down", "RefMC")
#getPath("2016","Down", "RefMC")

#getPath("2011","Down", "PHSP")
#getPath("2012","Down", "PHSP")
#getPath("2015","Down", "PHSP")
#getPath("2016","Down", "PHSP")
#getPath("2017","Down", "PHSP")
#getPath("2018","Down", "PHSP")
#printAllProductionIDs("RefMC")
#printAllProductionIDs("PHSP")