Exemplo n.º 1
0
arrayHH4W = tools.appendTreeArray(arrayHH4W)
imgArrayHH4W = img.makeBoostCandFourVector(arrayHH4W)

arrayHH4B = tree2array(treeHH4B, treeVars, sel)
arrayHH4B = tools.appendTreeArray(arrayHH4B)
imgArrayHH4B = img.makeBoostCandFourVector(arrayHH4B)

print "Made candidate 4 vector arrays from the datasets"

#==================================================================================
# Make Jet Images /////////////////////////////////////////////////////////////////
#==================================================================================

jetImagesDF = {}
print "Creating boosted Jet Images for QCD"
jetImagesDF['QCD'] = img.prepareBoostedImages(imgArrayJJ, arrayJJ, 10,
                                              boostAxis)
print "Creating boosted Jet Images for HH->bbbb"
jetImagesDF['HH4B'] = img.prepareBoostedImages(imgArrayHH4B, arrayHH4B, 10,
                                               boostAxis)
print "Creating boosted Jet Images for HH->WWWW"
jetImagesDF['HH4W'] = img.prepareBoostedImages(imgArrayHH4W, arrayHH4W, 10,
                                               boostAxis)

print "Made jet image data frames"

h5f = h5py.File("data/phiCosThetaBoostedJetImages10by10.h5", "w")
h5f.create_dataset('QCD', data=jetImagesDF['QCD'], compression='lzf')
h5f.create_dataset('HH4B', data=jetImagesDF['HH4B'], compression='lzf')
h5f.create_dataset('HH4W', data=jetImagesDF['HH4W'], compression='lzf')

print "Saved Boosted Jet Images"
T_image = []
W_image = []
Z_image = []
BES_vars = []
ET = []
Mass = []
numPass = 0
for index, jet in enumerate(treeCandidates):
    #Selection criteria here
    if index % 1000 == 1: print "Imaging jet", index
    if (
            jet.jetAK8_pt > 500 and jet.jetAK8_pt < 1800
            and jet.jetAK8_SoftDropMass > 10 and jet.nSubjets_Higgs > 3
    ):  #nSubjets cut vetoes only ~0.1% of events, but many variables poorly defined if nSubjets < 4
        H_image.append(
            img.prepareBoostedImages(jet, 'H', 31, smearImage, smearWidth,
                                     smearPoints))
        T_image.append(
            img.prepareBoostedImages(jet, 'T', 31, smearImage, smearWidth,
                                     smearPoints))
        W_image.append(
            img.prepareBoostedImages(jet, 'W', 31, smearImage, smearWidth,
                                     smearPoints))
        Z_image.append(
            img.prepareBoostedImages(jet, 'Z', 31, smearImage, smearWidth,
                                     smearPoints))
        BES_vars.append(tools.GetBESVars(jet, bestVars))
        ET.append(jet.jetAK8_pt)
        Mass.append(jet.jetAK8_mass)
        numPass += 1
        if numPass > 330000: break
print "Boosted Event Shape Variables: ", bestVars

# make arrays from the trees
start, stop, step = 0, 167262, 1
bestArrayHH4B = tree2array(treeHH4B, bestVars, sel, None, start, stop, step)
bestArrayHH4B = tools.appendTreeArray(bestArrayHH4B)

print "Made array with the Boosted Event Shape Variables"

#==================================================================================
# Make Jet Images /////////////////////////////////////////////////////////////////
#==================================================================================

jetImagesDF = {}
print "Creating boosted Jet Images for HH->bbbb"
jetImagesDF['HH4B_images'] = img.prepareBoostedImages(imgArrayHH4B, arrayHH4B,
                                                      31, boostAxis)

print "Made jet image data frames"

#==================================================================================
# Store BEST Variables in DataFrame ///////////////////////////////////////////////
#==================================================================================

jetImagesDF['HH4B_BES_vars'] = bestArrayHH4B
print "Stored BES variables"

#==================================================================================
# Store Data in h5 file ///////////////////////////////////////////////////////////
#==================================================================================

h5f = h5py.File("images/HH4BphiCosThetaBoostedJetImagesX10.h5", "w")
print "Boosted Event Shape Variables: ", bestVars

# make arrays from the trees
start, stop, step = 0, 167262, 1
bestArrayQCD = tree2array(treeQCD, bestVars, sel, None, start, stop, step)
bestArrayQCD = tools.appendTreeArray(bestArrayQCD)

print "Made array with the Boosted Event Shape Variables"

#==================================================================================
# Make Jet Images /////////////////////////////////////////////////////////////////
#==================================================================================

jetImagesDF = {}
print "Creating boosted Jet Images for QCD"
jetImagesDF['QCD_images'] = img.prepareBoostedImages(imgArrayQCD, arrayQCD, 31, boostAxis)

print "Made jet image data frames"

#==================================================================================
# Store BEST Variables in DataFrame ///////////////////////////////////////////////
#==================================================================================

jetImagesDF['QCD_BES_vars'] = bestArrayQCD
print "Stored BES variables"

#==================================================================================
# Store Data in h5 file ///////////////////////////////////////////////////////////
#==================================================================================

h5f = h5py.File("images/OpenDataQCDphiCosThetaBoostedJetImagesX10.h5","w")