[PWGLF] Fix indexing and add invariant mass plot

PWGLF/TableProducer/Nuspex/he3LambdaAnalysis.cxx

@@ -1,4 +1,4 @@
 // Copyright 2019-2020 CERN and copyright holders of ALICE O2.
 // See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
 // All rights not expressly granted are reserved.
 //
@@ -8,7 +8,7 @@
 // In applying this license CERN does not waive the privileges and immunities
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
 
 #include "Common/Core/trackUtilities.h"
 #include "Common/DataModel/Centrality.h"
 #include "Common/DataModel/EventSelection.h"
@@ -57,15 +57,15 @@
 
 namespace lfv0he3
 {
 DECLARE_SOA_INDEX_COLUMN(LFEvent, lfEvent); // Collision ID for the event
 DECLARE_SOA_COLUMN(Pt, pt, float);
 DECLARE_SOA_COLUMN(Eta, eta, float);
 DECLARE_SOA_COLUMN(Phi, phi, float);
 DECLARE_SOA_COLUMN(Mass, mass, float);
 DECLARE_SOA_COLUMN(CosPA, cosPA, float);
 DECLARE_SOA_COLUMN(DCAxy, dcaXY, float);
 DECLARE_SOA_COLUMN(DCAz, dcaZ, float);
 DECLARE_SOA_COLUMN(TPCnCls, tpcNCls, int);
 DECLARE_SOA_COLUMN(ITSClusterSizes, itsClusterSizes, uint32_t);
 DECLARE_SOA_COLUMN(NsigmaTPC, nSigmaTPC, float);
 DECLARE_SOA_COLUMN(DCAdaughters, dcaDaughters, float);
@@ -80,7 +80,7 @@
 
 namespace
 {
 constexpr double betheBlochDefault[1][6]{{-1.e32, -1.e32, -1.e32, -1.e32, -1.e32, -1.e32}};
 static const std::vector<std::string> betheBlochParNames{"p0", "p1", "p2", "p3", "p4", "resolution"};
 static const std::vector<std::string> particleName{"He3"};
 o2::base::MatLayerCylSet* matLUT = nullptr;
@@ -106,6 +106,7 @@
 std::shared_ptr<TH2> hTPCnSigmaHe3;
 std::shared_ptr<TH2> hArmenterosPodolanskiAll;
 std::shared_ptr<TH2> hArmenterosPodolanskiSelected;
+std::shared_ptr<TH2> hInvariantMass;
 
 }; // namespace
 
@@ -203,8 +204,7 @@
   std::array<double, 6> mBBparamsHe;
   float mBz = 0.0f; // Magnetic field in T
   HistogramRegistry mRegistry{"He3LambdaAnalysis"};
-  int mRunNumber = 0;  // Current run number
-  int mEventIndex = 0; // Current event index
+  int mRunNumber = 0; // Current run number
 
   void init(InitContext const&)
   {
@@ -213,7 +213,7 @@
     ccdb->setLocalObjectValidityChecking();
     ccdb->setFatalWhenNull(true);
 
     for (int i = 0; i < 5; i++) {
       mBBparamsHe[i] = cfgHe3.betheBlochParams->get("He3", Form("p%i", i));
     }
     mBBparamsHe[5] = cfgHe3.betheBlochParams->get("He3", "resolution");
@@ -248,6 +248,9 @@
     hArmenterosPodolanskiAll = mRegistry.add<TH2>("hArmenterosPodolanskiAll", "Armenteros-Podolanski All", {HistType::kTH2D, {{100, -1., 1.}, {100, 0., 0.5}}});
     hArmenterosPodolanskiSelected = mRegistry.add<TH2>("hArmenterosPodolanskiSelected", "Armenteros-Podolanski Selected", {HistType::kTH2D, {{100, -1., 1.}, {100, 0., 0.5}}});
 
+    constexpr double ConstituentsMass = o2::constants::physics::MassProton + o2::constants::physics::MassNeutron * 2 + o2::constants::physics::MassSigmaPlus;
+    hInvariantMass = mRegistry.add<TH2>("hInvariantMass", "Invariant Mass", {HistType::kTH2D, {{45, 1., 10}, {100, ConstituentsMass - 0.05, ConstituentsMass + 0.05}}});
+
     LOGF(info, "He3-Lambda analysis initialized");
   }
 
@@ -426,14 +429,20 @@
     // Fill output tables
     lfHe3V0Collision(collision.posZ(), collision.centFT0C());
     for (const auto& he3 : he3Candidates) {
-      lfHe3(mEventIndex, he3.momentum.Pt(), he3.momentum.Eta(), he3.momentum.Phi(),
+      lfHe3(lfHe3V0Collision.lastIndex(), he3.momentum.Pt(), he3.momentum.Eta(), he3.momentum.Phi(),
             he3.dcaXY, he3.dcaZ, he3.tpcNClsFound, he3.itsClusterSizes, he3.nSigmaTPC, he3.sign);
     }
     for (const auto& lambda : lambdaCandidates) {
-      lfLambda(mEventIndex, lambda.momentum.Pt(), lambda.momentum.Eta(), lambda.momentum.Phi(),
+      lfLambda(lfHe3V0Collision.lastIndex(), lambda.momentum.Pt(), lambda.momentum.Eta(), lambda.momentum.Phi(),
                lambda.mass, lambda.cosPA, lambda.dcaV0Daughters, lambda.dcaProtonToPV, lambda.dcaPionToPV, lambda.v0Radius, lambda.sign);
     }
-    mEventIndex++;
+
+    for (const auto& he3 : he3Candidates) {
+      for (const auto& lambda : lambdaCandidates) {
+        auto pairMomentum = lambda.momentum + he3.momentum; // Calculate invariant mass
+        hInvariantMass->Fill(pairMomentum.Pt(), pairMomentum.M());
+      }
+    }
   }
   PROCESS_SWITCH(he3LambdaAnalysis, processData, "Process data", true);
 };