add simple PU subtraction

megaClustering.py

@@ -9,7 +9,7 @@
 # import math
 import hgcalHelpers
 # import hgcalHistHelpers
-# import numpy as np
+import numpy as np
 import pandas as pd
 from itertools import repeat
 maxlayer = 52
@@ -24,7 +24,39 @@ def getConeRadius(frontRadius, backRadius, z, maxval=9999.):
     return val
 
 
-def getMegaClusters(genParticles, multiClusters, layerClusters, recHits, gun_type, GEN_engpt, pidSelected, energyRadius=6, frontRadius=3, backRadius=8):
+def pileupSubtraction(matchedMultiCluster, selectedLayerClusters, recHits, layer, energyRadius, frontRadius, backRadius):
+    """For now, the code is the same as in getMegaClusters, but the phi coordinate is changed by pi."""
+
+    pTSum = 0
+    energySum = 0
+    # take first layer cluster z value
+    layer_z = selectedLayerClusters.head(1).z.item()
+    # get multi cluster x and y coordinates
+    matchedMultiCluster_phi = float(matchedMultiCluster.phi) - np.pi
+    if (matchedMultiCluster_phi < -np.pi):
+        matchedMultiCluster_phi += 2*np.pi
+    multiClusPosDF = hgcalHelpers.convertToXY(matchedMultiCluster.eta, matchedMultiCluster_phi, layer_z)
+    # calculate radius based on current layer's z position
+    coneRadius = getConeRadius(frontRadius, backRadius, layer_z)
+    # mind that we need only the first index since there is only one multiCluster
+    layerClusterIndices = hgcalHelpers.getIndicesWithinRadius(multiClusPosDF[['x', 'y']], selectedLayerClusters[['x', 'y']], coneRadius)
+    # now we need to recalculate the layer cluster energies using associated RecHits
+    for layerClusterIndex in layerClusterIndices[0]:
+        associatedRecHits = recHits.iloc[selectedLayerClusters.iloc[layerClusterIndex].rechits]
+        # find maximum energy RecHit
+        maxEnergyRecHitIndex = associatedRecHits['energy'].argmax()
+        # considering only associated RecHits within a radius of energyRadius (6 cm)
+        matchedRecHitIndices = hgcalHelpers.getIndicesWithinRadius(associatedRecHits.loc[[maxEnergyRecHitIndex]][['x', 'y']], associatedRecHits[['x', 'y']], energyRadius)[maxEnergyRecHitIndex]
+        # sum up energies and pT
+        selectedRecHits = associatedRecHits.iloc[matchedRecHitIndices]
+        # correct energy by subdetector weights
+        energySum += selectedRecHits[["energy"]].sum()[0]*energyWeights[layer-1]*1.38
+        pTSum += selectedRecHits[["pt"]].sum()[0]*energyWeights[layer-1]*1.38
+
+    return (energySum, pTSum)
+
+
+def getMegaClusters(genParticles, multiClusters, layerClusters, recHits, gun_type, GEN_engpt, pidSelected, energyRadius=6, frontRadius=3, backRadius=8, doPileupSubtraction=True):
     """
     get the actual mega clusters.
     frontRadius: cone at front of EE
@@ -81,8 +113,12 @@ def getMegaClusters(genParticles, multiClusters, layerClusters, recHits, gun_typ
                 # sum up energies and pT
                 selectedRecHits = associatedRecHits.iloc[matchedRecHitIndices]
                 # correct energy by subdetector weights
-                energySum += selectedRecHits[["energy"]].sum()[0]*energyWeights[layer-1]
-                pTSum += selectedRecHits[["pt"]].sum()[0]*energyWeights[layer-1]
+                energySum += selectedRecHits[["energy"]].sum()[0]*energyWeights[layer-1]*1.38
+                pTSum += selectedRecHits[["pt"]].sum()[0]*energyWeights[layer-1]*1.38
+            if (doPileupSubtraction):
+                (pu_energySum, pu_pTSum) = pileupSubtraction(matchedMultiCluster, selectedLayerClusters, recHits, layer, energyRadius, frontRadius, backRadius)
+                energySum -= pu_energySum
+                pTSum -= pu_pTSum
 
         # use as coordinates eta and phi of matched multi cluster
         megaCluster = [pTSum, matchedMultiCluster['eta'].item(), matchedMultiCluster['phi'].item(), energySum]

resolutionScaleFiller.py

@@ -220,6 +220,8 @@ def fillComparisonHistograms(resolutionScaleObjects, GEN_engpt, histDict):
             # fill the hists
             if len(valueLists) > 0:
                 rangeGeV = GEN_engpt * 1.6
+                if (GEN_engpt < 30):
+                    rangeGeV = GEN_engpt * 5
                 nbins = int(rangeGeV)
                 if (rangeGeV < 50.):
                     nbins = int(10 * rangeGeV)
@@ -233,7 +235,7 @@ def fillComparisonHistograms(resolutionScaleObjects, GEN_engpt, histDict):
                 # object of interest Pt
                 histDict[etaBinName][phiBinName] = hgcalHistHelpers.histValue1D(valueLists[:, 3], histDict[etaBinName][phiBinName], tag="obj_Pt_eta" + etaBinName + "_phi" + phiBinName, title="Pt of object of interest, #eta=" + GEN_eta + ", #phi=" + GEN_phi + ")",   axunit="p_{T} [GeV]", binsBoundariesX=binsBoundariesX_eng, ayunit="N(clusters)")
                 # response
-                binsBoundariesX_relDiff = [[800, -100, 60], [650, -80, 50]]["1p" in etaBinName]
+                binsBoundariesX_relDiff = [800, -100, 60]
                 histDict[etaBinName][phiBinName] = hgcalHistHelpers.histValue1D(valueLists[:, 4]*100, histDict[etaBinName][phiBinName], tag="obj_dEoverE_eta" + etaBinName + "_phi" + phiBinName, title="dEoverE, #eta=" + GEN_eta + ", #phi=" + GEN_phi + ")", axunit="#Delta E_{clust}/E_{clust}[%]", binsBoundariesX=binsBoundariesX_relDiff, ayunit="N(clusters)")
                 histDict[etaBinName][phiBinName] = hgcalHistHelpers.histValue1D(valueLists[:, 5]*100, histDict[etaBinName][phiBinName], tag="obj_dPtoverPt_eta" + etaBinName + "_phi" + phiBinName, title="dPtoverPt, #eta=" + GEN_eta + ", #phi=" + GEN_phi + ")", axunit="#Delta p_{T, clust}/p_{T, clust}[%]", binsBoundariesX=binsBoundariesX_relDiff, ayunit="N(clusters)")
                 # for resolution
@@ -254,10 +256,11 @@ def main():
     parser = optparse.OptionParser(usage)
 
     # input options
-    parser.add_option('', '--files', dest='fileString', type='string',  default='root://eoscms.cern.ch//eos/cms/store/cmst3/group/hgcal/CMG_studies/Production/_RelValSingleGammaPt25Eta1p7_2p7_CMSSW_9_3_2-PU25ns_93X_upgrade2023_realistic_v2_2023D17PU200-v1_GEN-SIM-RECO/NTUP/_RelValSingleGammaPt25Eta1p7_2p7_CMSSW_9_3_2-PU25ns_93X_upgrade2023_realistic_v2_2023D17PU200-v1_GEN-SIM-RECO_NTUP_1.root', help='comma-separated file list')
+    # parser.add_option('', '--files', dest='fileString', type='string',  default='root://eoscms.cern.ch//eos/cms/store/cmst3/group/hgcal/CMG_studies/Production/_SinglePiPt50Eta1p6_2p8_PhaseIITDRFall17DR-noPUFEVT_93X_upgrade2023_realistic_v2-v1_GEN-SIM-RECO/NTUP/_SinglePiPt50Eta1p6_2p8_PhaseIITDRFall17DR-noPUFEVT_93X_upgrade2023_realistic_v2-v1_GEN-SIM-RECO_NTUP_1_0.root', help='comma-separated file list')
+    parser.add_option('', '--files', dest='fileString', type='string',  default='root://eoscms.cern.ch//eos/cms/store/cmst3/group/hgcal/CMG_studies/Production/_SinglePiPt50Eta1p6_2p8_PhaseIITDRFall17DR-PU200FEVT_93X_upgrade2023_realistic_v2-v1_GEN-SIM-RECO/NTUP/_SinglePiPt50Eta1p6_2p8_PhaseIITDRFall17DR-PU200FEVT_93X_upgrade2023_realistic_v2-v1_GEN-SIM-RECO_NTUP_2.root', help='comma-separated file list')
     parser.add_option('', '--gunType', dest='gunType', type='string',  default='pt', help='pt or e')
-    parser.add_option('', '--pid', dest='pid', type='int',  default=22, help='pdgId int')
-    parser.add_option('', '--genValue', dest='genValue', type='int',  default=25, help='generated pT or energy')
+    parser.add_option('', '--pid', dest='pid', type='int',  default=211, help='pdgId int')
+    parser.add_option('', '--genValue', dest='genValue', type='int',  default=50, help='generated pT or energy')
     parser.add_option('', '--tag', dest='tag', type='string',  default='noPU', help='some tag, best used for PU and other info')
     parser.add_option('', '--ref', dest='refName', type='string',  default='genpart', help='reference collection')
     parser.add_option('', '--obj', dest='objName', type='string',  default='pfcluster', help='object of interest collection')