[PWGLF] added function to compute 2d weights (#10646)

Author: alcaliva

PWGLF/Tasks/Strangeness/strangenessInJets.cxx

@@ -1881,262 +1881,170 @@ struct StrangenessInJets {
   }
   PROCESS_SWITCH(StrangenessInJets, processMCefficiency, "Process MC Efficiency", false);
 
-  /*
   void processGen(o2::aod::McCollisions const& mcCollisions, aod::McParticles const& mcParticles)
   {
     for (const auto& mccollision : mcCollisions) {
 
-      registryMC.fill(HIST("number_of_events_mc"), 3.5);
-
       // Selection on z_{vertex}
-      if (std::fabs(mccollision.posZ()) > 10)
+      if (std::fabs(mccollision.posZ()) > 10.0)
         continue;
-      registryMC.fill(HIST("number_of_events_mc"), 4.5);
 
       // MC Particles per Collision
       auto mcParticlesPerColl = mcParticles.sliceBy(perMCCollision, mccollision.globalIndex());
 
-      // List of Tracks
-      std::vector<TVector3> trk;
-      std::vector<int> ntrk;
-
+      // loop over generated MC particles
+      std::vector<fastjet::PseudoJet> fjParticles;
       for (const auto& particle : mcParticlesPerColl) {
 
-        // Select Primary Particles
-        double dx = particle.vx() - mccollision.posX();
-        double dy = particle.vy() - mccollision.posY();
-        double dz = particle.vz() - mccollision.posZ();
-        double dcaxy = std::sqrt(dx * dx + dy * dy);
-        double dcaz = std::fabs(dz);
-        if (dcaxy > 0.25)
-          continue;
-        if (dcaz > 2.0)
+        if (!particle.isPhysicalPrimary())
           continue;
         if (std::fabs(particle.eta()) > 0.8)
           continue;
-        if (particle.pt() < 0.15)
+        if (particle.pt() < 0.1)
           continue;
 
-        // PDG Selection
-        int pdg = std::fabs(particle.pdgCode());
-        if ((pdg != 11) && (pdg != 211) && (pdg != 321) && (pdg != 2212))
-          continue;
-
-        TVector3 momentum(particle.px(), particle.py(), particle.pz());
-        trk.push_back(momentum);
-        ntrk.push_back(1);
+        // 4-momentum representation of a particle
+        double energy = std::sqrt(particle.p() * particle.p() + MassPionCharged * MassPionCharged);
+        fastjet::PseudoJet fourMomentum(particle.px(), particle.py(), particle.pz(), energy);
+        fjParticles.emplace_back(fourMomentum);
       }
+      // reject empty events
+      if (fjParticles.size() < 1)
+        continue;
+
+      // cluster particles using the anti-kt algorithm
+      fastjet::JetDefinition jetDef(fastjet::antikt_algorithm, rJet);
+      fastjet::AreaDefinition areaDef(fastjet::active_area, fastjet::GhostedAreaSpec(1.0));
+      fastjet::ClusterSequenceArea cs(fjParticles, jetDef, areaDef);
+      std::vector<fastjet::PseudoJet> jets = fastjet::sorted_by_pt(cs.inclusive_jets());
+      auto [rhoPerp, rhoMPerp] = backgroundSub.estimateRhoPerpCone(fjParticles, jets);
 
-      // Anti-kt Jet Finder
-      int nParticlesRemoved(0);
-      std::vector<TVector3> jet;
+      // jet selection
+      bool isAtLeastOneJetSelected = false;
+      std::vector<TVector3> selectedJet;
       std::vector<TVector3> ue1;
       std::vector<TVector3> ue2;
-      std::vector<int> nParticlesInjet;
-
-      do {
-        double dijMin(1e+06), diBmin(1e+06);
-        int iMin(0), jMin(0), iBmin(0);
-        for (int i = 0; i < static_cast<int>(trk.size()); i++) {
-          if (trk[i].Mag() == 0)
-            continue;
-          double diB = 1.0 / (trk[i].Pt() * trk[i].Pt());
-          if (diB < diBmin) {
-            diBmin = diB;
-            iBmin = i;
-          }
-          for (int j = (i + 1); j < static_cast<int>(trk.size()); j++) {
-            if (trk[j].Mag() == 0)
-              continue;
-            double dij = calculateDij(trk[i], trk[j], rJet);
-            if (dij < dijMin) {
-              dijMin = dij;
-              iMin = i;
-              jMin = j;
-            }
-          }
-        }
-        if (dijMin < diBmin) {
-          trk[iMin] = trk[iMin] + trk[jMin];
-          ntrk[iMin] = ntrk[iMin] + ntrk[jMin];
-          trk[jMin].SetXYZ(0, 0, 0);
-          ntrk[jMin] = 0;
-          nParticlesRemoved++;
-        }
-        if (dijMin > diBmin) {
-          jet.push_back(trk[iBmin]);
-          nParticlesInjet.push_back(ntrk[iBmin]);
-          trk[iBmin].SetXYZ(0, 0, 0);
-          nParticlesRemoved++;
-        }
-      } while (nParticlesRemoved < static_cast<int>(trk.size()));
 
-      // Jet Selection
-      std::vector<int> isSelected;
-      int nJetsSelected(0);
-      for (int i = 0; i < static_cast<int>(jet.size()); i++) {
+      for (auto& jet : jets) { // o2-linter: disable=[const-ref-in-for-loop]
 
-        // Initialization
-        isSelected.push_back(0);
-
-        // Jet fully contained inside acceptance
-        if ((std::fabs(jet[i].Eta()) + rJet) > (etaMax - 0.5))
+        // jet must be fully contained in the acceptance
+        if ((std::fabs(jet.eta()) + rJet) > (etaMax - deltaEtaEdge))
           continue;
-        if (nParticlesInjet[i] < 2)
+
+        // jet pt must be larger than threshold
+        fastjet::PseudoJet jetMinusBkg = backgroundSub.doRhoAreaSub(jet, rhoPerp, rhoMPerp);
+        if (getCorrectedPt(jetMinusBkg.pt()) < minJetPt)
           continue;
+        isAtLeastOneJetSelected = true;
 
-        // Perpendicular cones
+        // perpendicular cone
+        TVector3 jetAxis(jet.px(), jet.py(), jet.pz());
         TVector3 ueAxis1(0, 0, 0);
         TVector3 ueAxis2(0, 0, 0);
-        getPerpendicularAxis(jet[i], ueAxis1, +1);
-        getPerpendicularAxis(jet[i], ueAxis2, -1);
-        ue1.push_back(ueAxis1);
-        ue2.push_back(ueAxis2);
-
-        double ptJetCorr = jet[i].Pt() - 0.5;
-        if (ptJetCorr < minJetPt)
-          continue;
-
-        nJetsSelected++;
-        isSelected[i] = 1;
+        getPerpendicularAxis(jetAxis, ueAxis1, +1);
+        getPerpendicularAxis(jetAxis, ueAxis2, -1);
+        selectedJet.emplace_back(jetAxis);
+        ue1.emplace_back(ueAxis1);
+        ue2.emplace_back(ueAxis2);
       }
-      if (nJetsSelected == 0)
+      if (!isAtLeastOneJetSelected)
         continue;
 
-      for (int i = 0; i < static_cast<int>(jet.size()); i++) {
-
-        if (isSelected[i] == 0)
-          continue;
-
-        // Generated Particles
+      // loop over selected jets / UE
+      for (int i = 0; i < static_cast<int>(selectedJet.size()); i++) {
         for (const auto& particle : mcParticlesPerColl) {
 
           if (!particle.isPhysicalPrimary())
             continue;
+          if (std::fabs(particle.eta()) > 0.8)
+            continue;
+          if (particle.pt() < 0.1)
+            continue;
 
           TVector3 particleDir(particle.px(), particle.py(), particle.pz());
-          const double deltaEtaJet = particleDir.Eta() - jet[i].Eta();
-          const double deltaPhiJet = getDeltaPhi(particleDir.Phi(), jet[i].Phi());
-          const double deltaRjet = std::sqrt(deltaEtaJet * deltaEtaJet + deltaPhiJet * deltaPhiJet);
-          const double deltaEtaUe1 = particleDir.Eta() - ue1[i].Eta();
-          const double deltaPhiUe1 = getDeltaPhi(particleDir.Phi(), ue1[i].Phi());
-          const double deltaRue1 = std::sqrt(deltaEtaUe1 * deltaEtaUe1 + deltaPhiUe1 * deltaPhiUe1);
-          const double deltaEtaUe2 = particleDir.Eta() - ue2[i].Eta();
-          const double deltaPhiUe2 = getDeltaPhi(particleDir.Phi(), ue2[i].Phi());
-          const double deltaRue2 = std::sqrt(deltaEtaUe2 * deltaEtaUe2 + deltaPhiUe2 * deltaPhiUe2);
-
-          switch (particle.pdgCode()) {
-            case 211:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("pi_plus_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
-                registryMC.fill(HIST("pi_plus_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case -211:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("pi_minus_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
-                registryMC.fill(HIST("pi_minus_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case 321:
-              if (deltaRjet < rJet) {
+          float deltaEtaJet = particleDir.Eta() - selectedJet[i].Eta();
+          float deltaPhiJet = getDeltaPhi(particleDir.Phi(), selectedJet[i].Phi());
+          float deltaRjet = std::sqrt(deltaEtaJet * deltaEtaJet + deltaPhiJet * deltaPhiJet);
+          float deltaEtaUe1 = particleDir.Eta() - ue1[i].Eta();
+          float deltaPhiUe1 = getDeltaPhi(particleDir.Phi(), ue1[i].Phi());
+          float deltaRue1 = std::sqrt(deltaEtaUe1 * deltaEtaUe1 + deltaPhiUe1 * deltaPhiUe1);
+          float deltaEtaUe2 = particleDir.Eta() - ue2[i].Eta();
+          float deltaPhiUe2 = getDeltaPhi(particleDir.Phi(), ue2[i].Phi());
+          float deltaRue2 = std::sqrt(deltaEtaUe2 * deltaEtaUe2 + deltaPhiUe2 * deltaPhiUe2);
+
+          // In jet
+          if (deltaRjet < rJet) {
+            switch (particle.pdgCode()) {
+              case 211:
                 registryMC.fill(HIST("pi_plus_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
-                registryMC.fill(HIST("pi_plus_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case -321:
-              if (deltaRjet < rJet) {
+                break;
+              case -211:
                 registryMC.fill(HIST("pi_minus_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
-                registryMC.fill(HIST("pi_minus_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case 2212:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("pi_plus_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case 310:
+                registryMC.fill(HIST("K0s_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              case 3122:
+                registryMC.fill(HIST("Lambda_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              case -3122:
+                registryMC.fill(HIST("AntiLambda_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              case 3312:
+                registryMC.fill(HIST("Xi_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              case -3312:
+                registryMC.fill(HIST("AntiXi_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              case 3334:
+                registryMC.fill(HIST("Omega_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              case -3334:
+                registryMC.fill(HIST("AntiOmega_eta_pt_jet"), particle.pt(), particle.eta());
+                break;
+              default:
+                continue;
+            }
+          }
+
+          if (deltaRue1 < rJet || deltaRue2 < rJet) {
+            switch (particle.pdgCode()) {
+              case 211:
                 registryMC.fill(HIST("pi_plus_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case -2212:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("pi_minus_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case -211:
                 registryMC.fill(HIST("pi_minus_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case 310:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("K0s_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case 310:
                 registryMC.fill(HIST("K0s_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case 3122:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("Lambda_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case 3122:
                 registryMC.fill(HIST("Lambda_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case -3122:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("AntiLambda_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case -3122:
                 registryMC.fill(HIST("AntiLambda_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case 3312:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("Xi_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case 3312:
                 registryMC.fill(HIST("Xi_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case -3312:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("AntiXi_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case -3312:
                 registryMC.fill(HIST("AntiXi_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case 3334:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("Omega_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case 3334:
                 registryMC.fill(HIST("Omega_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
-            case -3334:
-              if (deltaRjet < rJet) {
-                registryMC.fill(HIST("AntiOmega_eta_pt_jet"), particle.pt(), particle.eta());
-              }
-              if (deltaRue1 < rJet || deltaRue2 < rJet) {
+                break;
+              case -3334:
                 registryMC.fill(HIST("AntiOmega_eta_pt_ue"), particle.pt(), particle.eta());
-              }
-              break;
+                break;
+              default:
+                continue;
+            }
           }
         }
       }
     }
   }
   PROCESS_SWITCH(StrangenessInJets, processGen, "Process generated MC", false);
-  */
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)