[PWGCF] Femto: Addition of SigmaPlus variables in output tables

PWGCF/Femto/Core/kinkBuilder.h

@@ -554,7 +554,7 @@ class KinkBuilder
       mProduceSigmaPlus = utils::enableTable("FSigmaPlus_001", table.produceSigmaPlus.value, initContext);
       mProduceLiteSigmaPlus = utils::enableTable("FLiteSigmaPlus_001", table.produceLiteSigmaPlus.value, initContext);
       mProduceSigmaPlusMasks = utils::enableTable("FSigmaPlusMasks_001", table.produceSigmaPlusMasks.value, initContext);
-      mProduceSigmaPlusExtras = utils::enableTable("FSigmaPlusExtras_001", table.produceSigmaPlusExtras.value, initContext);
+      mProduceSigmaPlusExtras = utils::enableTable("FSigmaPlusExtras_002", table.produceSigmaPlusExtras.value, initContext);
 
       if (mProduceSigmaPlus && mProduceLiteSigmaPlus) {
         LOG(fatal) << "FSigmaPlus and FLiteSigmaPlus are mutually exclusive -- enable only one. "
@@ -615,7 +615,8 @@ class KinkBuilder
         fillSigma(collisionBuilder, kinkProducts, kink, daughterIndex);
       }
       if constexpr (modes::isEqual(kinkType, modes::Kink::kSigmaPlus)) {
-        fillSigmaPlus(collisionBuilder, kinkProducts, kink, daughterIndex);
+        auto mother = kink.template trackMoth_as<T7>();
+        fillSigmaPlus(collisionBuilder, kinkProducts, kink, daughter, mother, daughterIndex);
       }
     }
   }
@@ -651,7 +652,8 @@ class KinkBuilder
         mcBuilder.template fillMcSigmaWithLabel<system>(daughter, mcParticles, mcCols, mcProducts);
       }
       if constexpr (modes::isEqual(kinkType, modes::Kink::kSigmaPlus)) {
-        fillSigmaPlus(collisionBuilder, kinkProducts, kink, daughterIndex);
+        auto mother = kink.template trackMoth_as<T8>();
+        fillSigmaPlus(collisionBuilder, kinkProducts, kink, daughter, mother, daughterIndex);
         mcBuilder.template fillMcSigmaPlusWithLabel<system>(daughter, mcParticles, mcCols, mcProducts);
       }
     }
@@ -695,8 +697,8 @@ class KinkBuilder
     }
   }
 
-  template <typename T1, typename T2, typename T3>
-  void fillSigmaPlus(T1& collisionBuilder, T2& kinkProducts, T3 const& kink, int64_t daughterIndex)
+  template <typename T1, typename T2, typename T3, typename T4, typename T5>
+  void fillSigmaPlus(T1& collisionBuilder, T2& kinkProducts, T3 const& kink, T4 const& daughter, T5 const& mother, int64_t daughterIndex)
   {
     // Mass is calculated from original momentum components stored in kink table
     float mass = kink.mSigmaPlus();
@@ -728,7 +730,13 @@ class KinkBuilder
         kink.xDecVtx(),
         kink.yDecVtx(),
         kink.zDecVtx(),
-        mKinkSelection.getKinkTransRadius());
+        mKinkSelection.getKinkTransRadius(),
+        kink.dcaKinkTopo(),
+        daughter.tpcNSigmaPr(),
+        daughter.tpcNSigmaPi(),
+        daughter.tofNSigmaPr(),
+        kink.ptMoth(),
+        mother.itsClusterSizes());
     }
   }
 
@@ -759,6 +767,7 @@ struct KinkBuilderDerivedToDerivedProducts : o2::framework::ProducesGroup {
   o2::framework::Produces<o2::aod::StoredFSigmaMasks> producedSigmaMasks;
   o2::framework::Produces<o2::aod::StoredFSigmaPlus> producedSigmaPluses;
   o2::framework::Produces<o2::aod::StoredFSigmaPlusMasks> producedSigmaPlusMasks;
+  o2::framework::Produces<o2::aod::StoredFSigmaPlusExtras> producedSigmaPlusExtras;
 };
 
 class KinkBuilderDerivedToDerived
@@ -843,6 +852,7 @@ class KinkBuilderDerivedToDerived
                                             sigmaPlus.mass(),
                                             chaDaughterIndex);
       newSigmaPlusTable.producedSigmaPlusMasks(sigmaPlus.mask());
+      fillSigmaPlusExtensions(sigmaPlus, newSigmaPlusTable);
     }
   }
 
@@ -894,11 +904,32 @@ class KinkBuilderDerivedToDerived
                                             sigmaPlus.mass(),
                                             chaDaughterIndex);
       newSigmaPlusTable.producedSigmaPlusMasks(sigmaPlus.mask());
+      fillSigmaPlusExtensions(sigmaPlus, newSigmaPlusTable);
       mcBuilder.fillSigmaPlusWithLabel(sigmaPlus, mcCols, mcParticles, mcMothers, mcPartonicMothers, mcProducts);
     }
   }
 
  private:
+  template <typename T1, typename T2>
+  void fillSigmaPlusExtensions(T1 const& sigmaPlus, T2& newSigmaPlusTable)
+  {
+    if constexpr (requires(T1 t) { t.transRadius(); }) {
+      newSigmaPlusTable.producedSigmaPlusExtras(sigmaPlus.kinkAngle(),
+                                                sigmaPlus.dcaDaugToPV(),
+                                                sigmaPlus.dcaMothToPV(),
+                                                sigmaPlus.decayVtxX(),
+                                                sigmaPlus.decayVtxY(),
+                                                sigmaPlus.decayVtxZ(),
+                                                sigmaPlus.transRadius(),
+                                                sigmaPlus.dcaKinkTopo(),
+                                                sigmaPlus.dauTpcNSigmaPr(),
+                                                sigmaPlus.dauTpcNSigmaPi(),
+                                                sigmaPlus.dauTofNSigmaPr(),
+                                                sigmaPlus.ptOriginal(),
+                                                sigmaPlus.itsClusterSizes());
+    }
+  }
+
   int mLimitSigma = 0;
   int mLimitSigmaPlus = 0;
 };

PWGCF/Femto/DataModel/FemtoTables.h

@@ -1014,6 +1014,11 @@ DECLARE_SOA_COLUMN(DecayVtxX, decayVtxX, float);     //! x coordinate of decay v
 DECLARE_SOA_COLUMN(DecayVtxY, decayVtxY, float);     //! y coordinate of decay vertex (relative to PV)
 DECLARE_SOA_COLUMN(DecayVtxZ, decayVtxZ, float);     //! z coordinate of decay vertex (relative to PV)
 DECLARE_SOA_COLUMN(TransRadius, transRadius, float); //! Transverse decay radius from PV
+DECLARE_SOA_COLUMN(DcaKinkTopo, dcaKinkTopo, float);
+DECLARE_SOA_COLUMN(DauTpcNSigmaPr, dauTpcNSigmaPr, float);
+DECLARE_SOA_COLUMN(DauTpcNSigmaPi, dauTpcNSigmaPi, float);
+DECLARE_SOA_COLUMN(DauTofNSigmaPr, dauTofNSigmaPr, float);
+DECLARE_SOA_COLUMN(PtOriginal, ptOriginal, float);
 
 // id column for charged daughter track
 DECLARE_SOA_INDEX_COLUMN_FULL(ChaDau, chaDau, int32_t, FTracks, "_ChaDau"); //!
@@ -1167,7 +1172,23 @@ DECLARE_SOA_TABLE_STAGED_VERSIONED(FSigmaPlusExtras_001, "FSIGMAPLUSEXTRAS", 1,
                                    femtokinks::DecayVtxZ,
                                    femtokinks::TransRadius);
 
-using FSigmaPlusExtras = FSigmaPlusExtras_001;
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FSigmaPlusExtras_002, "FSIGMAPLUSEXTRAS", 2,
+                                   femtokinks::KinkAngle,
+                                   femtokinks::DcaDaugToPV,
+                                   femtokinks::DcaMothToPV,
+                                   femtokinks::DecayVtxX,
+                                   femtokinks::DecayVtxY,
+                                   femtokinks::DecayVtxZ,
+                                   femtokinks::TransRadius,
+                                   femtokinks::DcaKinkTopo,
+                                   femtokinks::DauTpcNSigmaPr,
+                                   femtokinks::DauTpcNSigmaPi,
+                                   femtokinks::DauTofNSigmaPr,
+                                   femtokinks::PtOriginal,
+                                   femtotracks::ItsClusterSizes);
+
+using FSigmaPlusExtras = FSigmaPlusExtras_002;
+using StoredFSigmaPlusExtras = StoredFSigmaPlusExtras_002;
 
 namespace femtocascades
 {

PWGCF/Femto/TableProducer/femtoProducerDerivedToDerived.cxx

@@ -54,10 +54,12 @@ struct FemtoProducerDerivedToDerived {
   using FemtoOmegas = o2::soa::Join<o2::aod::FOmegas, o2::aod::FOmegaMasks>;
   using FemtoSigma = o2::soa::Join<o2::aod::FSigmas, o2::aod::FSigmaMasks>;
   using FemtoSigmaPlus = o2::soa::Join<o2::aod::FSigmaPlus, o2::aod::FSigmaPlusMasks>;
+  using FemtoSigmaPlusExtended = o2::soa::Join<FemtoSigmaPlus, o2::aod::FSigmaPlusExtras>;
 
   using FemtoMcTracks = o2::soa::Join<FemtoTracks, o2::aod::FTrackLabels>;
   using FemtoMcSigma = o2::soa::Join<FemtoSigma, o2::aod::FSigmaLabels>;
   using FemtoMcSigmaPlus = o2::soa::Join<FemtoSigmaPlus, o2::aod::FSigmaPlusLabels>;
+  using FemtoMcSigmaPlusExtended = o2::soa::Join<FemtoSigmaPlusExtended, o2::aod::FSigmaPlusLabels>;
   using FemtoMcParticles = o2::soa::Join<o2::aod::FMcParticles, o2::aod::FMcMotherLabels>;
 
   o2::framework::SliceCache cache;
@@ -121,6 +123,8 @@ struct FemtoProducerDerivedToDerived {
   kinkbuilder::ConfSigmaPlusSelection1 sigmaPlusSelection1;
   o2::framework::Partition<FemtoSigmaPlus> sigmaPlusPartition = MAKE_SIGMAPLUS_PARTITION(sigmaPlusSelection1);
   o2::framework::Preslice<FemtoSigmaPlus> perColSigmaPlus = o2::aod::femtobase::stored::fColId;
+  o2::framework::Partition<FemtoSigmaPlusExtended> sigmaPlusExtendedPartition = MAKE_SIGMAPLUS_PARTITION(sigmaPlusSelection1);
+  o2::framework::Preslice<FemtoSigmaPlusExtended> perColSigmaPlusExtended = o2::aod::femtobase::stored::fColId;
 
   // cascade builder
   cascadebuilder::CascadeBuilderDerivedToDerived cascadeBuilder;
@@ -149,6 +153,8 @@ struct FemtoProducerDerivedToDerived {
 
   o2::framework::Partition<FemtoMcSigmaPlus> mcSigmaPlusPartition = MAKE_SIGMAPLUS_PARTITION(sigmaPlusSelection1);
   o2::framework::Preslice<FemtoMcSigmaPlus> perColMcSigmaPlus = o2::aod::femtobase::stored::fColId;
+  o2::framework::Partition<FemtoMcSigmaPlusExtended> mcSigmaPlusExtendedPartition = MAKE_SIGMAPLUS_PARTITION(sigmaPlusSelection1);
+  o2::framework::Preslice<FemtoMcSigmaPlusExtended> perColMcSigmaPlusExtended = o2::aod::femtobase::stored::fColId;
 
   void init(o2::framework::InitContext& /*context*/)
   {
@@ -158,10 +164,10 @@ struct FemtoProducerDerivedToDerived {
       static_cast<int>(doprocessLambdas) +
       static_cast<int>(doprocessTracksXis) + static_cast<int>(doprocessTracksOmegas) +
       static_cast<int>(doprocessTracksK0shorts) + static_cast<int>(doprocessTracksSigma) +
-      static_cast<int>(doprocessTracksSigmaPlus) +
+      static_cast<int>(doprocessTracksSigmaPlus) + static_cast<int>(doprocessTracksSigmaPlusExtended) +
       static_cast<int>(doprocessTracksRho0) + static_cast<int>(doprocessTracksPhi) +
       static_cast<int>(doprocessTracksMc) + static_cast<int>(doprocessTracksSigmaMc) +
-      static_cast<int>(doprocessTracksSigmaPlusMc);
+      static_cast<int>(doprocessTracksSigmaPlusMc) + static_cast<int>(doprocessTracksSigmaPlusExtendedMc);
     if (activeProcesses != 1) {
       LOG(fatal) << "Exactly one process function must be activated (got " << activeProcesses << ").";
     }
@@ -292,6 +298,18 @@ struct FemtoProducerDerivedToDerived {
   }
   PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksSigmaPlus, "Process sigmaPlus and tracks", false);
 
+  void processTracksSigmaPlusExtended(FilteredFemtoCollision const& col, FemtoTracks const& tracks, FemtoSigmaPlusExtended const& sigmaplus)
+  {
+    if (trackBuilder.collisionHasTooFewTracks(col, tracks, trackPartition1, trackPartition2, cache) || kinkBuilder.collisionHasTooFewSigmaPlus(col, sigmaplus, sigmaPlusExtendedPartition, cache)) {
+      return;
+    }
+    trackBuilder.reset(tracks);
+    collisionBuilder.processCollision(col, collisionBuilderProducts);
+    trackBuilder.processTracks(col, tracks, trackPartition1, trackPartition2, cache, trackBuilderProducts, collisionBuilderProducts);
+    kinkBuilder.processSigmaPlus(col, sigmaplus, tracks, sigmaPlusExtendedPartition, trackBuilder, cache, kinkBuilderProducts, trackBuilderProducts, collisionBuilderProducts);
+  }
+  PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksSigmaPlusExtended, "Process sigmaPlus (with SigmaPlusExtras) and tracks", false);
+
   void processTracksRho0(FilteredFemtoCollision const& col, FemtoTracks const& tracks, FemtoRho0s const& rho0s)
   {
     if (trackBuilder.collisionHasTooFewTracks(col, tracks, trackPartition1, trackPartition2, cache) || resonanceBuilder.collisionHasTooFewRho0s(col, rho0s, rho0Partition, cache)) {
@@ -396,6 +414,32 @@ struct FemtoProducerDerivedToDerived {
     }
   }
   PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksSigmaPlusMc, "Process sigmaPlus and tracks with MC information", false);
+
+  void processTracksSigmaPlusExtendedMc(FilteredFemtoMcCollisions const& cols,
+                                        FemtoMcTracks const& tracks,
+                                        FemtoMcSigmaPlusExtended const& sigmaplus,
+                                        o2::aod::FMcCols const& mcCols,
+                                        FemtoMcParticles const& mcParticles,
+                                        o2::aod::FMcMothers const& mcMothers,
+                                        o2::aod::FMcPartMoths const& mcPartonicMothers)
+  {
+    mcBuilder.reset(mcCols, mcParticles);
+
+    for (auto const& col : cols) {
+      if (trackBuilder.collisionHasTooFewTracks(col, tracks, mcTrackPartition1, mcTrackPartition2, cache) || kinkBuilder.collisionHasTooFewSigmaPlus(col, sigmaplus, mcSigmaPlusExtendedPartition, cache)) {
+        continue;
+      }
+      trackBuilder.reset(tracks);
+      collisionBuilder.processCollision(col, collisionBuilderProducts);
+      mcBuilder.fillCollisionWithLabel(col, mcCols, mcBuilderProducts);
+      trackBuilder.processTracksWithMc(col, tracks, mcTrackPartition1, mcTrackPartition2, cache,
+                                       trackBuilderProducts, collisionBuilderProducts,
+                                       mcBuilder, mcCols, mcParticles, mcMothers, mcPartonicMothers, mcBuilderProducts);
+      kinkBuilder.processSigmaPlusWithMc(col, sigmaplus, tracks, mcSigmaPlusExtendedPartition, trackBuilder, cache, kinkBuilderProducts, trackBuilderProducts, collisionBuilderProducts,
+                                         mcBuilder, mcCols, mcParticles, mcMothers, mcPartonicMothers, mcBuilderProducts);
+    }
+  }
+  PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksSigmaPlusExtendedMc, "Process sigmaPlus (with SigmaPlusExtras) and tracks with MC information", false);
 };
 
 o2::framework::WorkflowSpec defineDataProcessing(o2::framework::ConfigContext const& context)