Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2024-06-01 07:07:12

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023 Christopher Dilks
0003 #ifdef INCLUDE_PODIO
0004 
0005 #include "AnalysisEpicPodio.h"
0006 
0007 AnalysisEpicPodio::AnalysisEpicPodio(TString infileName_, TString outfilePrefix_)
0008   : Analysis(infileName_, outfilePrefix_)
0009   , crossCheckKinematics(false)
0010 {};
0011 
0012 AnalysisEpicPodio::~AnalysisEpicPodio() {};
0013 
0014 void AnalysisEpicPodio::Execute()
0015 {
0016   // setup
0017   Prepare();
0018 
0019   // produce flat list of files from `infiles`
0020   std::vector<std::string> infilesFlat;
0021   for(const auto fileList : infiles)
0022     for(const auto fileName : fileList)
0023       infilesFlat.push_back(fileName);
0024 
0025   // create PODIO event store
0026   podioReader.openFiles(infilesFlat);
0027   evStore.setReader(&podioReader);
0028   ENT = podioReader.getEntries();
0029   if(maxEvents>0) ENT = std::min(maxEvents,ENT);
0030   
0031   // calculate Q2 weights
0032   CalculateEventQ2Weights();
0033 
0034   // upstream reconstruction methods
0035   // - list of upstream methods
0036   const std::vector<std::string> upstreamReconMethodList = {
0037     "Truth",
0038     "Electron",
0039     "DA",
0040     "JB",
0041     "Sigma"
0042   };
0043   // - association of `Kinematics::CalculateDIS` reconstruction method with upstream;
0044   //   for those unavailable upstream, use `"NONE"`
0045   const std::map<TString,std::string> associatedUpstreamMethodMap = {
0046     { "Ele",    "Electron" },
0047     { "DA",     "DA"       },
0048     { "JB",     "JB"       },
0049     { "Sigma",  "Sigma"    },
0050     { "Mixed",  "NONE"     },
0051     { "eSigma", "NONE"     }
0052   };
0053   // - get upstream method associated with `reconMethod`
0054   const auto& associatedUpstreamMethod = associatedUpstreamMethodMap.at(reconMethod);
0055 
0056   // event loop =========================================================
0057   fmt::print("begin event loop...\n");
0058   for(unsigned e=0; e<ENT; e++) {
0059     if(e%10000==0) fmt::print("{} events...\n",e);
0060     if(verbose) fmt::print("\n\n{:=<70}\n",fmt::format("EVENT {} ",e));
0061 
0062     // next event
0063     // FIXME: check that we analyze ALL of the events: do we miss the first or last one?
0064     if(e>0) {
0065       evStore.clear();
0066       podioReader.endOfEvent();
0067     }
0068 
0069     // resets
0070     kin->ResetHFS();
0071     kinTrue->ResetHFS();
0072     double mcPartElectronP   = 0.0;
0073     bool double_counted_beam = false;
0074     int num_ele_beams        = 0;
0075     int num_ion_beams        = 0;
0076     int num_sim_electrons    = 0;
0077     int num_rec_electrons    = 0;
0078     
0079     // read particle collections for this event
0080     const auto& simParts    = evStore.get<edm4hep::MCParticleCollection>("MCParticles");
0081     const auto& recParts    = evStore.get<edm4eic::ReconstructedParticleCollection>("ReconstructedParticles");
0082     const auto& mcRecAssocs = evStore.get<edm4eic::MCRecoParticleAssociationCollection>("ReconstructedParticlesAssoc");
0083 
0084     // data objects
0085     edm4hep::MCParticle mcPartEleBeam;
0086     edm4hep::MCParticle mcPartIonBeam;
0087     edm4hep::MCParticle mcPartElectron;
0088 
0089     // loop over generated particles
0090     if(verbose) fmt::print("\n{:-<60}\n","MCParticles ");
0091     for(auto simPart : simParts) {
0092 
0093       // print out this MCParticle
0094       // if(verbose) PrintParticle(simPart);
0095 
0096       // generated particle properties
0097       auto simPDG = simPart.getPDG();
0098 
0099       // add to Hadronic Final State (HFS) sums
0100       kinTrue->AddToHFS(GetP4(simPart));
0101 
0102       // filter for beam particles
0103       if(simPart.getGeneratorStatus() == constants::statusBeam) {
0104         switch(simPDG) {
0105           case constants::pdgElectron:
0106             if(num_ele_beams>0) double_counted_beam = true;
0107             mcPartEleBeam = simPart;
0108             num_ele_beams++;
0109             break;
0110           case constants::pdgProton:
0111             if(num_ion_beams>0) double_counted_beam = true;
0112             mcPartIonBeam = simPart;
0113             num_ion_beams++;
0114             break;
0115           default:
0116             ErrorPrint(fmt::format("WARNING: Unknown beam particle with PDG={}",simPDG));
0117         }
0118       }
0119 
0120       // filter for scattered electron: select the one with the highest |p|
0121       if(simPart.getGeneratorStatus() == constants::statusFinal) {
0122         if(simPDG == constants::pdgElectron) {
0123           auto eleP = edm4hep::utils::p(simPart);
0124           if(eleP>mcPartElectronP) {
0125             mcPartElectron  = simPart;
0126             mcPartElectronP = eleP;
0127             num_sim_electrons++;
0128           }
0129         }
0130       }
0131 
0132     } // end loop over generated particles
0133 
0134     // check for found generated particles
0135     if(num_ele_beams==0)     { ErrorPrint("WARNING: missing MC electron beam");      continue; };
0136     if(num_ion_beams==0)     { ErrorPrint("WARNING: missing MC ion beam");           continue; };
0137     if(num_sim_electrons==0) { ErrorPrint("WARNING: missing scattered electron");    continue; };
0138     if(double_counted_beam)  { ErrorPrint("WARNING: found multiple beam particles"); continue; };
0139 
0140     // set Kinematics 4-momenta
0141     kinTrue->vecEleBeam  = GetP4(mcPartEleBeam);
0142     kinTrue->vecIonBeam  = GetP4(mcPartIonBeam);
0143     kinTrue->vecElectron = GetP4(mcPartElectron);
0144 
0145     // print beam particles
0146     if(verbose) {
0147       if(verbose) fmt::print("\n{:-<60}\n","GENERATED BEAMS ");
0148       PrintParticle(mcPartEleBeam);
0149       PrintParticle(mcPartIonBeam);
0150       if(verbose) fmt::print("\n{:-<60}\n","GENERATED SCATTERED ELECTRON ");
0151       PrintParticle(mcPartElectron);
0152     }
0153 
0154     // add reconstructed particles to Hadronic Final State (HFS)
0155     /* the following will run loops over Reconstructed Particle <-> MC Particle associations
0156      * - uses high-order function `LoopMCRecoAssocs` for common tasks, such as quality cuts
0157      *   and getting the reconstructed PID (PDG)
0158      * - first define a first-order function (`payload`), then call `LoopMCRecoAssocs`
0159      * - see `LoopMCRecoAssocs` for `payload` signature
0160      */
0161     auto AllRecPartsToHFS = [&] (auto& simPart, auto& recPart, auto recPDG) {
0162       kin->AddToHFS(GetP4(recPart));
0163     };
0164     if(verbose) fmt::print("\n{:-<60}\n","MC<->Reco ASSOCIATIONS ");
0165     LoopMCRecoAssocs(mcRecAssocs, AllRecPartsToHFS, verbose);
0166 
0167     // find reconstructed electron
0168     // ============================================================================
0169     /* FIXME: need realistic electron finder; all of the following options rely
0170      * on MC-truth matching; is there any common upstream realistic electron finder
0171      */
0172 
0173     // find scattered electron by simply matching to truth
0174     // FIXME: not working, until we have truth matching and/or reconstructed PID
0175     // FIXME: does `simPart==mcPartElectron` work as expected?
0176     /*
0177     auto FindRecoEleByTruth = [&] (auto& simPart, auto& recPart, auto recPDG) {
0178       if(recPDG==constants::pdgElectron && simPart==mcPartElectron) {
0179         num_rec_electrons++;
0180         kin->vecElectron = GetP4(recPart);
0181       };
0182     };
0183     LoopMCRecoAssocs(mcRecAssocs, FindRecoEleByTruth);
0184     */
0185 
0186     // use electron finder from upstream algorithm `InclusiveKinematics*`
0187     // FIXME: is the correct upstream electron finder used here? The
0188     // `InclusiveKinematics*` recon algorithms seem to rely on
0189     // `Jug::Base::Beam::find_first_scattered_electron(mcParts)` and matching
0190     // to truth; this guarantees we get the correct reconstructed scattered
0191     // electron
0192     const auto& disCalcs = evStore.get<edm4eic::InclusiveKinematicsCollection>("InclusiveKinematicsElectron");
0193     if(disCalcs.size() != 1) ErrorPrint(fmt::format("WARNING: disCalcs.size = {} != 1 for this event",disCalcs.size()));
0194     for(const auto& calc : disCalcs) {
0195       auto ele = calc.getScat();
0196       if( ! ele.isAvailable()) {
0197         ErrorPrint("WARNING: `disCalcs` scattered electron unavailable");
0198         continue;
0199       }
0200       num_rec_electrons++;
0201       kin->vecElectron = GetP4(ele);
0202     }
0203 
0204     // check for found reconstructed scattered electron
0205     if(num_rec_electrons == 0) { ErrorPrint("WARNING: reconstructed scattered electron not found"); continue; };
0206     if(num_rec_electrons >  1) { ErrorPrint("WARNING: found more than 1 reconstructed scattered electron"); };
0207 
0208     // subtract electron from hadronic final state variables
0209     kin->SubtractElectronFromHFS();
0210     kinTrue->SubtractElectronFromHFS();
0211 
0212     // skip the event if there are no reconstructed particles (other than the
0213     // electron), otherwise hadronic recon methods will fail
0214     if(kin->countHadrons == 0) { ErrorPrint("WARNING: no hadrons"); };
0215   
0216     // calculate DIS kinematics; skip the event if the calculation did not go well
0217     if( ! kin->CalculateDIS(reconMethod)     ) continue; // reconstructed
0218     if( ! kinTrue->CalculateDIS(reconMethod) ) continue; // generated (truth)
0219 
0220     // Get the weight for this event's Q2
0221     //   FIXME: we are in a podio::EventStore event loop, thus we need an
0222     //          alternative to `chain->GetTreeNumber()`; currently disabling weighting
0223     //          for now, by setting `wTrack=1.0`
0224     // auto Q2weightFactor = GetEventQ2Weight(kinTrue->Q2, inLookup[chain->GetTreeNumber()]);
0225     auto Q2weightFactor = 1.0;
0226 
0227     // fill inclusive histograms, if only `inclusive` is included in output
0228     // (otherwise they will be filled in track and jet loops)
0229     if(includeOutputSet["inclusive_only"]) {
0230       auto wInclusive = Q2weightFactor * weightInclusive->GetWeight(*kinTrue);
0231       wInclusiveTotal += wInclusive;
0232       FillHistosInclusive(wInclusive);
0233     }
0234 
0235     // loop over Reco<->MC associations again
0236     /* - calculate SIDIS kinematics
0237      * - fill output data structures
0238      */
0239     auto SidisOutput = [&] (auto& simPart, auto& recPart, auto recPDG) {
0240 
0241       // final state cut
0242       // - check PID, to see if it's a final state we're interested in
0243       auto kv = PIDtoFinalState.find(recPDG);
0244       if(kv!=PIDtoFinalState.end()) finalStateID = kv->second; else return;
0245       if(activeFinalStates.find(finalStateID)==activeFinalStates.end()) return;
0246 
0247       // set SIDIS particle 4-momenta, and calculate their kinematics
0248       kinTrue->vecHadron = GetP4(simPart);
0249       kinTrue->CalculateHadronKinematics();
0250       kin->vecHadron = GetP4(recPart);
0251       kin->CalculateHadronKinematics();
0252 
0253       // weighting
0254       auto wTrack = Q2weightFactor * weightTrack->GetWeight(*kinTrue);
0255       wTrackTotal += wTrack;
0256 
0257       // fill single-hadron histograms in activated bins
0258       FillHistos1h(wTrack);
0259       FillHistosInclusive(wTrack);
0260 
0261       // fill simple tree
0262       // - not binned
0263       // - `IsActiveEvent()` is only true if at least one bin gets filled for this track
0264       if( writeSidisTree && HD->IsActiveEvent() ) ST->FillTree(wTrack);
0265     };
0266     LoopMCRecoAssocs(mcRecAssocs, SidisOutput);
0267 
0268 
0269     // read kinematics calculations from upstream /////////////////////////
0270     // TODO: cross check these with our calculations from `Kinematics`
0271     if(crossCheckKinematics) {
0272       auto PrintRow = [] <typename T> (std::string name, std::vector<T> vals, bool header=false) {
0273         fmt::print("  {:>16}",name);
0274         if(header) { for(auto val : vals) fmt::print(" {:>8}",   val); }
0275         else       { for(auto val : vals) fmt::print(" {:8.4f}", val); }
0276         fmt::print("\n");
0277       };
0278       // upstream calculations
0279       fmt::print("\n{:-<75}\n","KINEMATICS, calculated from upstream: ");
0280       PrintRow("", std::vector<std::string>({ "x", "Q2", "W", "y", "nu" }), true);
0281       for(const auto upstreamReconMethod : upstreamReconMethodList)
0282         for(const auto& calc : evStore.get<edm4eic::InclusiveKinematicsCollection>("InclusiveKinematics"+upstreamReconMethod) )
0283           PrintRow( upstreamReconMethod, std::vector<float>({
0284               calc.getX(),
0285               calc.getQ2(),
0286               calc.getW(),
0287               calc.getY(),
0288               calc.getNu()
0289               }));
0290       // local calculations
0291       fmt::print("{:-<75}\n",fmt::format("KINEMATICS, calculated locally in EPIC-ANALYSIS, with method \"{}\": ",reconMethod));
0292       auto PrintKinematics = [&PrintRow] (std::string name, std::shared_ptr<Kinematics> K) {
0293         PrintRow( name, std::vector<Double_t>({
0294             K->x,
0295             K->Q2,
0296             K->W,
0297             K->y,
0298             K->Nu
0299             }));
0300       };
0301       PrintKinematics("Truth",kinTrue);
0302       PrintKinematics("Reconstructed",kin);
0303       // compare upstream and local
0304       if(associatedUpstreamMethod != "NONE") {
0305         fmt::print("{:-<75}\n",fmt::format("DIFFERENCE: upstream({}) - local({}): ",associatedUpstreamMethod,reconMethod));
0306         for(const auto upstreamMethod : std::vector<std::string>({"Truth",associatedUpstreamMethod})) {
0307           const auto& upstreamCalcs = evStore.get<edm4eic::InclusiveKinematicsCollection>("InclusiveKinematics"+upstreamMethod);
0308           for(const auto& upstreamCalc : upstreamCalcs) {
0309             auto K    = upstreamMethod=="Truth" ? kinTrue : kin;
0310             auto name = upstreamMethod=="Truth" ? "Truth" : "Reconstructed";
0311             PrintRow( name, std::vector<Double_t>({
0312                 upstreamCalc.getX()  - K->x,
0313                 upstreamCalc.getQ2() - K->Q2,
0314                 upstreamCalc.getW()  - K->W,
0315                 upstreamCalc.getY()  - K->y,
0316                 upstreamCalc.getNu() - K->Nu
0317                 }));
0318           }
0319         }
0320       }
0321       else fmt::print("{:-<75}\n  method \"{}\" is not available upstream\n","DIFFERENCE: ",reconMethod);
0322     } // if crossCheckKinematics
0323 
0324   } // event loop
0325   fmt::print("end event loop\n");
0326 
0327   // finish execution
0328   evStore.clear();
0329   podioReader.endOfEvent();
0330   podioReader.closeFile();
0331   Finish();
0332 }
0333 
0334 
0335 // particle printers //////////////////////////////////////////////
0336 
0337 void AnalysisEpicPodio::PrintParticle(const edm4hep::MCParticle& P) { 
0338   fmt::print("\n");
0339   fmt::print("  {:>20}: {}\n", "PDG",          P.getPDG()             );
0340   fmt::print("  {:>20}: {}\n", "Status",       P.getGeneratorStatus() );
0341   fmt::print("  {:>20}: {}\n", "Energy",       P.getEnergy()          );
0342   fmt::print("  {:>20}: {}\n", "p=|Momentum|", edm4hep::utils::p(P)   );
0343   fmt::print("  {:>20}: {}\n", "pT_lab",       edm4hep::utils::pT(P)  );
0344   fmt::print("  {:>20}: ({}, {}, {})\n",
0345       "3-Momentum",
0346       P.getMomentum().x,
0347       P.getMomentum().y,
0348       P.getMomentum().z
0349       );
0350   fmt::print("  {:>20}: ({}, {}, {})\n",
0351       "Vertex",
0352       P.getVertex().x,
0353       P.getVertex().y,
0354       P.getVertex().z
0355       );
0356   fmt::print("  {:>20}:\n", "Parents");
0357   for(const auto& parent : P.getParents())
0358     fmt::print("    {:>20}: {}\n", "PDG", parent.getPDG());
0359   fmt::print("  {:>20}:\n", "Daughters");
0360   for(const auto& daughter : P.getDaughters())
0361     fmt::print("    {:>20}: {}\n", "PDG", daughter.getPDG());
0362 }
0363 
0364 void AnalysisEpicPodio::PrintParticle(const edm4eic::ReconstructedParticle& P) {
0365   fmt::print("\n");
0366   fmt::print("  {:>20}: ", "PDG");
0367   if(P.getParticleIDUsed().isAvailable()) fmt::print("{}\n", P.getParticleIDUsed().getPDG());
0368   else fmt::print("???\n");
0369   fmt::print("  {:>20}: {}\n", "Mass",         P.getMass()           );
0370   fmt::print("  {:>20}: {}\n", "Charge",       P.getCharge()         );
0371   fmt::print("  {:>20}: {}\n", "Energy",       P.getEnergy()         );
0372   fmt::print("  {:>20}: {}\n", "p=|Momentum|", edm4hep::utils::p(P)  );
0373   fmt::print("  {:>20}: {}\n", "pT_lab",       edm4hep::utils::pT(P) );
0374   fmt::print("  {:>20}: ({}, {}, {})\n",
0375       "3-Momentum",
0376       P.getMomentum().x,
0377       P.getMomentum().y,
0378       P.getMomentum().z
0379       );
0380   fmt::print("  {:>20}: {}\n", "# of clusters", P.clusters_size()    );
0381   fmt::print("  {:>20}: {}\n", "# of tracks",   P.tracks_size()      );
0382   fmt::print("  {:>20}: {}\n", "# of PIDs",     P.particleIDs_size() );
0383   fmt::print("  {:>20}: {}\n", "# of recParts", P.particles_size()   );
0384   // for(const auto& track : P.getTracks()) {
0385   //   // ...
0386   // }
0387   // for(const auto& cluster : P.getClusters()) {
0388   //   // ...
0389   // }
0390 }
0391 
0392 
0393 // helper methods /////////////////////////////////////////////////////////////
0394 
0395 // common loop over Reconstructed Particle <-> MC Particle associations
0396 /* - get PID
0397  * - basic quality cuts
0398  * - execute `payload`
0399 //     payload signature: (simPart, recPart, reconstructed PDG)
0400  */
0401 void AnalysisEpicPodio::LoopMCRecoAssocs(
0402     const edm4eic::MCRecoParticleAssociationCollection& mcRecAssocs,
0403     std::function<void(const edm4hep::MCParticle&, const edm4eic::ReconstructedParticle&, int)> payload,
0404     bool printParticles
0405     )
0406 {
0407   for(const auto& assoc : mcRecAssocs ) {
0408 
0409     // get reconstructed and simulated particles, and check for matching
0410     auto recPart = assoc.getRec(); // reconstructed particle
0411     auto simPart = assoc.getSim(); // simulated (truth) particle
0412     // if(!simPart.isAvailable()) continue; // FIXME: consider using this once we have matching
0413 
0414     // print out this reconstructed particle, and its matching truth 
0415     if(printParticles) {
0416       fmt::print("\n   {:->35}\n"," reconstructed particle:");
0417       PrintParticle(recPart);
0418       fmt::print("\n   {:.>35}\n"," truth match:");
0419       if(simPart.isAvailable())
0420         PrintParticle(simPart);
0421       else
0422         fmt::print("     {:>35}\n","NO MATCH");
0423       fmt::print("\n");
0424     }
0425 
0426     // get reconstructed PDG from PID
0427     bool usedTruthPID = false;
0428     auto recPDG = GetReconstructedPDG(simPart, recPart, usedTruthPID);
0429     if(verbose) fmt::print("   GetReconstructedPDG = {}\n",recPDG);
0430     // if(usedTruthPID) continue; // FIXME: consider using this once we have decent PID
0431 
0432     // run payload
0433     payload(simPart, recPart, recPDG);
0434 
0435   } // end loop over Reco<->MC associations
0436 } // end LoopMCRecoAssocs
0437 
0438 
0439 // get PDG from reconstructed particle; resort to true PDG, if
0440 // PID is unavailable (sets `usedTruth` to true)
0441 int AnalysisEpicPodio::GetReconstructedPDG(
0442     const edm4hep::MCParticle& simPart,
0443     const edm4eic::ReconstructedParticle& recPart,
0444     bool& usedTruth
0445     )
0446 {
0447   int pdg = 0;
0448   usedTruth = false;
0449 
0450   // if using edm4hep::ReconstructedParticle:
0451   /*
0452   if(recPart.getParticleIDUsed().isAvailable()) // FIXME: not available
0453     pdg = recPart.getParticleIDUsed().getPDG();
0454   */
0455   
0456   // if using edm4eic::ReconstructedParticle:
0457   // pdg = recPart.getPDG(); // FIXME: not available either
0458 
0459   // if reconstructed PID is unavailable, use MC PDG
0460   if(pdg==0) {
0461     usedTruth = true;
0462     if(simPart.isAvailable())
0463       pdg = simPart.getPDG();
0464   }
0465 
0466   return pdg;
0467 }
0468 
0469 #endif