[PWGLF,Trigger] Nuclei trigger: fix wrong nested loop and V0 CosPA calculation (#11120)

Co-authored-by: Francesco Mazzaschi <fmazzasc@alipap1.cern.ch>

Author: fmazzasc

EventFiltering/PWGLF/nucleiFilter.cxx

@@ -431,73 +431,86 @@ struct nucleiFilter {
           h2TPCsignal[iN]->Fill(track.sign() * track.tpcInnerParam() * fixTPCrigidity, track.tpcSignal());
         }
       }
+      //
+      // fill QA histograms
+      //
+      qaHists.fill(HIST("fTPCsignal"), track.sign() * track.tpcInnerParam() * fixTPCrigidity, track.tpcSignal());
 
-      for (const auto& track : tracks) {
-        if (track.itsNCls() < cfgCutNclusITS ||
-            track.tpcNClsFound() < cfgCutNclusTPC ||
-            std::abs(track.dcaXY()) > cfgCutDCAxy ||
-            std::abs(track.dcaZ()) > cfgCutDCAz ||
-            std::abs(track.eta()) > cfgCutEta) {
+    } // end loop over tracks
+
+    for (const auto& track : tracks) {
+      if (track.itsNCls() < cfgCutNclusITS ||
+          track.tpcNClsFound() < cfgCutNclusTPC ||
+          std::abs(track.dcaXY()) > cfgCutDCAxy ||
+          std::abs(track.dcaZ()) > cfgCutDCAz ||
+          std::abs(track.eta()) > cfgCutEta) {
+        continue;
+      }
+      const ROOT::Math::PtEtaPhiMVector trackVector(track.pt(), track.eta(), track.phi(), constants::physics::MassPiMinus);
+      for (size_t iH3{0}; iH3 < h3vectors.size(); ++iH3) {
+        if (h3indices[iH3] == track.globalIndex()) {
           continue;
         }
-        const ROOT::Math::PtEtaPhiMVector trackVector(track.pt(), track.eta(), track.phi(), constants::physics::MassPiMinus);
-        for (size_t iH3{0}; iH3 < h3vectors.size(); ++iH3) {
-          if (h3indices[iH3] == track.globalIndex()) {
-            continue;
-          }
-          const auto& h3vector = h3vectors[iH3];
-          auto pivector = trackVector;
-          auto cm = h3vector + trackVector;
-          const ROOT::Math::Boost boost(cm.BoostToCM());
-          boost(pivector);
-          if (pivector.P() < cfgCutKstar) {
-            keepEvent[kTritonFemto] = true;
-            break;
-          }
+        const auto& h3vector = h3vectors[iH3];
+        auto pivector = trackVector;
+        auto cm = h3vector + trackVector;
+        const ROOT::Math::Boost boost(cm.BoostToCM());
+        boost(pivector);
+        if (pivector.P() < cfgCutKstar) {
+          keepEvent[kTritonFemto] = true;
+          break;
         }
       }
+    }
 
-      for (const auto& v0 : v0s) {
-        const auto& posTrack = v0.posTrack_as<TrackCandidates>();
-        const auto& negTrack = v0.negTrack_as<TrackCandidates>();
-        if ((posTrack.itsNCls() < cfgCutNclusITS || posTrack.tpcNClsFound() < cfgCutNclusTPC) &&
-            (negTrack.itsNCls() < cfgCutNclusITS || negTrack.tpcNClsFound() < cfgCutNclusTPC)) {
-          continue;
-        }
-        float nSigmas[2]{
-          cfgBetheBlochParams->get(2, 5u) > 0.f ? getNsigma(posTrack, 2, 0) : posTrack.tpcNSigmaHe(),
-          cfgBetheBlochParams->get(2, 5u) > 0.f ? getNsigma(negTrack, 2, 1) : negTrack.tpcNSigmaHe()};
-        if ((nSigmas[0] < cfgCutsPID->get(2, 0u) || nSigmas[0] > cfgCutsPID->get(2, 1u)) &&
-            (nSigmas[1] < cfgCutsPID->get(2, 0u) || nSigmas[1] > cfgCutsPID->get(2, 1u))) {
-          continue;
-        }
-        int n2bodyVtx = fitter2body.process(getTrackParCov(posTrack), getTrackParCov(negTrack));
-        if (n2bodyVtx == 0) {
-          continue;
-        }
-        auto vtxXYZ = fitter2body.getPCACandidate();
-        o2::gpu::gpustd::array<float, 3> mom = {0.};
-        vtxXYZ[0] -= collision.posX();
-        vtxXYZ[1] -= collision.posY();
-        vtxXYZ[2] -= collision.posZ();
-        auto momTrackParCov = fitter2body.createParentTrackPar();
-        momTrackParCov.getPxPyPzGlo(mom);
-        double cosPA = (vtxXYZ[0] * mom[0] + vtxXYZ[1] * mom[1] + vtxXYZ[2] * mom[2]) /
-                       std::sqrt((vtxXYZ[0] * vtxXYZ[0] + vtxXYZ[1] * vtxXYZ[1] + vtxXYZ[2] * vtxXYZ[2]) *
-                                 (mom[0] * mom[0] + mom[1] * mom[1] + mom[2] * mom[2]));
-        if (cosPA < cfgCutCosPAheV0) {
-          continue;
-        }
-        keepEvent[kHeV0] = true;
-        break;
+    for (const auto& v0 : v0s) {
+      const auto& posTrack = v0.posTrack_as<TrackCandidates>();
+      const auto& negTrack = v0.negTrack_as<TrackCandidates>();
+      if ((posTrack.itsNCls() < cfgCutNclusITS || posTrack.tpcNClsFound() < cfgCutNclusTPC) &&
+          (negTrack.itsNCls() < cfgCutNclusITS || negTrack.tpcNClsFound() < cfgCutNclusTPC)) {
+        continue;
       }
+      float nSigmas[2]{
+        cfgBetheBlochParams->get(2, 5u) > 0.f ? getNsigma(posTrack, 2, 0) : posTrack.tpcNSigmaHe(),
+        cfgBetheBlochParams->get(2, 5u) > 0.f ? getNsigma(negTrack, 2, 1) : negTrack.tpcNSigmaHe()};
 
-      //
-      // fill QA histograms
-      //
-      qaHists.fill(HIST("fTPCsignal"), track.sign() * track.tpcInnerParam() * fixTPCrigidity, track.tpcSignal());
+      bool isHe3 = nSigmas[0] > cfgCutsPID->get(2, 0u) && nSigmas[0] < cfgCutsPID->get(2, 1u);
+      bool isAntiHe3 = nSigmas[1] > cfgCutsPID->get(2, 0u) && nSigmas[1] < cfgCutsPID->get(2, 1u);
+      if (!isHe3 && !isAntiHe3) {
+        continue;
+      }
+      auto& he3Track = isHe3 ? posTrack : negTrack;
+      auto& piTrack = isHe3 ? negTrack : posTrack;
 
-    } // end loop over tracks
+      int n2bodyVtx = fitter2body.process(getTrackParCov(he3Track), getTrackParCov(piTrack));
+      if (n2bodyVtx == 0) {
+        continue;
+      }
+      auto vtxXYZ = fitter2body.getPCACandidate();
+      vtxXYZ[0] -= collision.posX();
+      vtxXYZ[1] -= collision.posY();
+      vtxXYZ[2] -= collision.posZ();
+
+      o2::gpu::gpustd::array<float, 3> momHe3 = {0.};
+      o2::gpu::gpustd::array<float, 3> momPi = {0.};
+      o2::gpu::gpustd::array<float, 3> momTot = {0.};
+      auto& hePropTrack = fitter2body.getTrack(0);
+      auto& piPropTrack = fitter2body.getTrack(1);
+      hePropTrack.getPxPyPzGlo(momHe3);
+      piPropTrack.getPxPyPzGlo(momPi);
+      for (int i = 0; i < 3; ++i) {
+        momHe3[i] *= 2;
+        momTot[i] = momHe3[i] + momPi[i];
+      }
+      double cosPA = (vtxXYZ[0] * momTot[0] + vtxXYZ[1] * momTot[1] + vtxXYZ[2] * momTot[2]) /
+                     std::sqrt((vtxXYZ[0] * vtxXYZ[0] + vtxXYZ[1] * vtxXYZ[1] + vtxXYZ[2] * vtxXYZ[2]) *
+                               (momTot[0] * momTot[0] + momTot[1] * momTot[1] + momTot[2] * momTot[2]));
+      if (cosPA < cfgCutCosPAheV0) {
+        continue;
+      }
+      keepEvent[kHeV0] = true;
+      break;
+    }
 
     // fH3L3Body trigger
     auto bc = collision.bc_as<aod::BCsWithTimestamps>();