snippets/ElectronID/ElectronID.cc

0001 #include "ElectronID.hh"
0002
0003 #include "edm4hep/utils/vector_utils.h"
0004 #include "edm4hep/utils/kinematics.h"
0005 #include "edm4eic/ClusterCollection.h"
0006
0007 #include <iostream>
0008
0009 ElectronID::ElectronID() {
0010
0011     mEe = 10.;
0012     mEh = 100.;
0013     std::cout << "!!! ElectronID: You have not specified beam energies...defaulting to 10x100 GeV !!!" << std::endl;
0014
0015     mEoP_min = 0.9;
0016     mEoP_max = 1.2;
0017
0018     mDeltaH_min = 0.*mEe;
0019     mDeltaH_min = 5.*mEe;
0020
0021     mIsoR = 0.4;
0022     mIsoE = 0.9;
0023
0024
0025 }
0026
0027 ElectronID::ElectronID(double Ee, double Eh) {
0028
0029     mEe = Ee;
0030     mEh = Eh;
0031
0032     mEoP_min = 0.9;
0033     mEoP_max = 1.2;
0034
0035     mDeltaH_min = 0.*mEe;
0036     mDeltaH_min = 5.*mEe;
0037
0038     mIsoR = 0.4;
0039     mIsoE = 0.9;
0040
0041 }
0042
0043 ElectronID::~ElectronID() {
0044 }
0045
0046
0047 void ElectronID::SetEvent(const podio::Frame* event) {
0048
0049     mEvent = event;
0050
0051 }
0052
0053
0054 edm4eic::ReconstructedParticleCollection ElectronID::FindScatteredElectron() {
0055
0056     // Get all the edm4eic objects needed for electron ID
0057     auto& rcparts = mEvent->get<edm4eic::ReconstructedParticleCollection>("ReconstructedParticles");
0058
0059     // Create collection for storing scattered electron candidates
0060     // (subset collection of ReconstructedParticleCollection)
0061     edm4eic::ReconstructedParticleCollection scatteredElectronCandidates;
0062     scatteredElectronCandidates->setSubsetCollection();
0063
0064     // Loop over all ReconstructedParticles for this event
0065     for (const auto& reconPart : rcparts) {
0066
0067         // Require negative particle
0068         if(reconPart.getCharge() >= 0) continue;
0069
0070         // Require at least one track and one cluster
0071         if(reconPart.getClusters().size() == 0 || reconPart.getTracks().size() == 0) continue;
0072
0073         // Calculate rcpart_ member variables for this event
0074         CalculateParticleValues(reconPart, rcparts);
0075
0076         // Calculate E/p and isolation fraction for this event
0077         // Note that the rcpart_ variables are set in CalculateParticleValues
0078         double recon_EoP = rcpart_sum_cluster_E / edm4hep::utils::magnitude(reconPart.getMomentum());
0079         double recon_isoE = rcpart_sum_cluster_E / rcpart_isolation_E;
0080
0081         // Apply scattered electron ID cuts
0082         if(recon_EoP < mEoP_min || recon_EoP > mEoP_max) continue;
0083         if(recon_isoE < mIsoE) continue;
0084
0085         // If particle passes cuts, add to output collection
0086         scatteredElectronCandidates.push_back(reconPart);
0087
0088     }
0089
0090     return scatteredElectronCandidates;
0091
0092 }
0093
0094 edm4hep::MCParticleCollection ElectronID::GetMCElectron() {
0095
0096     edm4hep::MCParticleCollection meMC;
0097     meMC->setSubsetCollection();
0098
0099     auto& mcparts = mEvent->get<edm4hep::MCParticleCollection>("MCParticles");
0100
0101     std::vector<edm4hep::MCParticle> mc_electrons;
0102
0103     for(const auto& mcp : mcparts) {
0104         if(mcp.getPDG() == 11 && mcp.getGeneratorStatus() == 1) {
0105             mc_electrons.push_back(mcp);
0106         }
0107     }
0108
0109     // For now, just take first electron
0110     if(mc_electrons.size() > 0) {
0111         meMC.push_back(mc_electrons[0]);
0112     }
0113
0114     return meMC;
0115
0116 }
0117
0118 edm4eic::ReconstructedParticleCollection ElectronID::GetTruthReconElectron() {
0119
0120     edm4hep::MCParticleCollection meMC = GetMCElectron();
0121     edm4eic::ReconstructedParticleCollection meRecon;
0122     meRecon->setSubsetCollection();
0123
0124     auto& RecoMC = mEvent->get<edm4eic::MCRecoParticleAssociationCollection>("ReconstructedParticleAssociations");
0125
0126     for(const auto& assoc : RecoMC) {
0127         if(assoc.getSim() == meMC[0]) {
0128             meRecon.push_back(assoc.getRec());
0129             break;
0130         }
0131     }
0132
0133     return meRecon;
0134
0135 }
0136
0137
0138
0139 void ElectronID::CalculateParticleValues(const edm4eic::ReconstructedParticle& rcp,
0140         const edm4eic::ReconstructedParticleCollection& rcparts) {
0141
0142     rcpart_sum_cluster_E = 0.;
0143     rcpart_lead_cluster_E = 0.;
0144     rcpart_isolation_E = 0.;
0145     rcpart_deltaH = 0.;
0146
0147     const edm4eic::Cluster* lead_cluster = nullptr;
0148
0149     for (const auto& cluster : rcp.getClusters()) {
0150         rcpart_sum_cluster_E += cluster.getEnergy();
0151         if(cluster.getEnergy() > rcpart_lead_cluster_E) {
0152             lead_cluster = &cluster;
0153             rcpart_lead_cluster_E = cluster.getEnergy();
0154         }
0155     }
0156
0157     if(!lead_cluster) return;
0158
0159     const auto& lead_pos = lead_cluster->getPosition();
0160     double lead_eta = edm4hep::utils::eta(lead_pos);
0161     double lead_phi = edm4hep::utils::angleAzimuthal(lead_pos);
0162
0163     for (const auto& other_rcp : rcparts) {
0164         if (&other_rcp == &rcp) continue;  // Skip the same particle
0165
0166         for (const auto& other_cluster : other_rcp.getClusters()) {
0167
0168             const auto& other_pos = other_cluster.getPosition();
0169             double other_eta = edm4hep::utils::eta(other_pos);
0170             double other_phi = edm4hep::utils::angleAzimuthal(other_pos);
0171
0172             double d_eta = other_eta - lead_eta;
0173             double d_phi = other_phi - lead_phi;
0174
0175             // Adjust d_phi to be in the range (-pi, pi)
0176             if (d_phi > M_PI) d_phi-=2*M_PI;
0177             if (d_phi < -M_PI) d_phi+=2*M_PI;
0178
0179             double dR = std::sqrt(std::pow(d_eta, 2) + std::pow(d_phi, 2));
0180
0181             // Check if the cluster is within the isolation cone
0182             if (dR < mIsoR) {
0183                 rcpart_isolation_E += other_cluster.getEnergy();
0184             }
0185         }
0186     }
0187
0188     return;
0189 }
0190
0191 edm4eic::ReconstructedParticle ElectronID::SelectHighestPT(const edm4eic::ReconstructedParticleCollection& ecandidates) {
0192
0193     edm4eic::ReconstructedParticle erec;
0194     double max_pT = 0.;
0195
0196     for(auto ecand : ecandidates) {
0197         double e_pT = edm4hep::utils::magnitudeTransverse(ecand.getMomentum());
0198         if(e_pT > max_pT) {
0199             erec = ecand;
0200             max_pT = e_pT;
0201         }
0202     }
0203
0204     return erec;
0205
0206 }
0207
0208 double ElectronID::GetCalorimeterEnergy(const edm4eic::ReconstructedParticle& rcp) {
0209
0210     double sum_cluster_E = 0.;
0211     for (const auto& cluster : rcp.getClusters()) {
0212         sum_cluster_E += cluster.getEnergy();
0213     }
0214     return sum_cluster_E;
0215
0216 }
0217
0218