[PWGCF] Feature: Derived-to-derived for kinks and pt recalc

PWGCF/Femto/Core/kinkBuilder.h

@@ -307,9 +307,19 @@ class KinkSelection : public BaseSelection<float, o2::aod::femtodatatypes::KinkM
   void computeKinkMotherKinematics(const T& kinkCand)
   {
     std::array<float, 3> momMother = {kinkCand.pxMoth(), kinkCand.pyMoth(), kinkCand.pzMoth()};
-    mKinkMotherPt = RecoDecay::pt(momMother);
+    std::array<float, 3> momDaughter = {kinkCand.pxDaug(), kinkCand.pyDaug(), kinkCand.pzDaug()};
+
     mKinkMotherEta = RecoDecay::eta(momMother);
     mKinkMotherPhi = RecoDecay::phi(momMother);
+
+    // Recalculate pT using kinematic constraints
+    float ptRecalc = utils::calcPtnew(momMother[0], momMother[1], momMother[2],
+                                      momDaughter[0], momDaughter[1], momDaughter[2]);
+    if (ptRecalc > 0.f) {
+      mKinkMotherPt = ptRecalc;
+    } else {
+      mKinkMotherPt = -1.f;
+    }
   }
 
   template <typename T>
@@ -401,7 +411,7 @@ class KinkBuilder
   {
     if constexpr (modes::isEqual(kinkType, modes::Kink::kSigma)) {
       LOG(info) << "Initialize femto Sigma builder...";
-      mProduceSigmas = utils::enableTable("FSigmas_001", table.produceSigmas.value, initContext);
+      mProduceSigmas = utils::enableTable("FSigmas_002", table.produceSigmas.value, initContext);
       mProduceSigmaMasks = utils::enableTable("FSigmaMasks_001", table.produceSigmaMasks.value, initContext);
       mProduceSigmaExtras = utils::enableTable("FSigmaExtras_001", table.produceSigmaExtras.value, initContext);
     }
@@ -510,11 +520,12 @@ class KinkBuilder
   template <typename T1, typename T2, typename T3>
   void fillSigma(T1& collisionProducts, T2& kinkProducts, T3 const& kink, int64_t daughterIndex)
   {
+    // Mass is calculated from original momentum components stored in kink table
     float mass = kink.mSigmaMinus();
 
     if (mProduceSigmas) {
       kinkProducts.producedSigmas(collisionProducts.producedCollision.lastIndex(),
-                                  kink.mothSign() * mKinkSelection.getKinkMotherPt(),
+                                  kink.mothSign() * mKinkSelection.getKinkMotherPt(), // Recalculated pT
                                   mKinkSelection.getKinkMotherEta(),
                                   mKinkSelection.getKinkMotherPhi(),
                                   mass,
@@ -538,11 +549,12 @@ class KinkBuilder
   template <typename T1, typename T2, typename T3>
   void fillSigmaPlus(T1& collisionProducts, T2& kinkProducts, T3 const& kink, int64_t daughterIndex)
   {
+    // Mass is calculated from original momentum components stored in kink table
     float mass = kink.mSigmaPlus();
 
     if (mProduceSigmaPlus) {
       kinkProducts.producedSigmaPlus(collisionProducts.producedCollision.lastIndex(),
-                                     kink.mothSign() * mKinkSelection.getKinkMotherPt(),
+                                     kink.mothSign() * mKinkSelection.getKinkMotherPt(), // Recalculated pT
                                      mKinkSelection.getKinkMotherEta(),
                                      mKinkSelection.getKinkMotherPhi(),
                                      mass,
@@ -576,6 +588,103 @@ class KinkBuilder
   bool mProduceSigmaPlusExtras = false;
 };
 
+struct ConfKinkTablesDerivedToDerived : o2::framework::ConfigurableGroup {
+  std::string prefix = std::string("KinkTables");
+  o2::framework::Configurable<int> limitSigma{"limitSigma", 1, "At least this many sigmas need to be in the collision"};
+  o2::framework::Configurable<int> limitSigmaPlus{"limitSigmaPlus", 0, "At least this many sigma pluses need to be in the collision"};
+};
+
+struct KinkBuilderDerivedToDerivedProducts : o2::framework::ProducesGroup {
+  o2::framework::Produces<o2::aod::StoredFSigmas_002> producedSigmas;
+  o2::framework::Produces<o2::aod::StoredFSigmaMasks_001> producedSigmaMasks;
+  o2::framework::Produces<o2::aod::StoredFSigmaPlus_001> producedSigmaPluses;
+  o2::framework::Produces<o2::aod::StoredFSigmaPlusMasks_001> producedSigmaPlusMasks;
+};
+
+class KinkBuilderDerivedToDerived
+{
+ public:
+  KinkBuilderDerivedToDerived() = default;
+  ~KinkBuilderDerivedToDerived() = default;
+
+  template <typename T>
+  void init(T& config)
+  {
+    mLimitSigma = config.limitSigma.value;
+    mLimitSigmaPlus = config.limitSigmaPlus.value;
+
+    if (mLimitSigma == 0 && mLimitSigmaPlus == 0) {
+      LOG(fatal) << "Both sigma limit and sigmaplus limit are 0. Breaking...";
+    }
+  }
+
+  template <typename T1, typename T2, typename T3, typename T4>
+  bool collisionHasTooFewSigma(T1 const& col, T2 const& /*sigmaTable*/, T3& partitionSigma, T4& cache)
+  {
+    auto sigmaSlice = partitionSigma->sliceByCached(o2::aod::femtobase::stored::fColId, col.globalIndex(), cache);
+    if (sigmaSlice.size() >= mLimitSigma) {
+      return false;
+    }
+    return true;
+  }
+
+  template <typename T1, typename T2, typename T3, typename T4>
+  bool collisionHasTooFewSigmaPlus(T1 const& col, T2 const& /*sigmaPlusTable*/, T3& partitionSigmaPlus, T4& cache)
+  {
+    auto sigmaPlusSlice = partitionSigmaPlus->sliceByCached(o2::aod::femtobase::stored::fColId, col.globalIndex(), cache);
+    if (sigmaPlusSlice.size() >= mLimitSigmaPlus) {
+      return false;
+    }
+    return true;
+  }
+
+  template <typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9>
+  void processSigma(T1 const& col, T2 const& /*sigmaTable*/, T3 const& oldTrackTable, T4& partitionSigma, T5& trackBuilder, T6& cache, T7& newSigmaTable, T8& newTrackTable, T9& newCollisionTable)
+  {
+    auto sigmaSlice = partitionSigma->sliceByCached(o2::aod::femtobase::stored::fColId, col.globalIndex(), cache);
+
+    for (auto const& sigma : sigmaSlice) {
+
+      auto chaDaughter = oldTrackTable.rawIteratorAt(sigma.chaDauId() - oldTrackTable.offset());
+
+      int chaDaughterIndex = trackBuilder.getDaughterIndex(chaDaughter, newTrackTable, newCollisionTable);
+
+      newSigmaTable.producedSigmas(newCollisionTable.producedCollision.lastIndex(),
+                                   sigma.signedPt(),
+                                   sigma.eta(),
+                                   sigma.phi(),
+                                   sigma.mass(),
+                                   chaDaughterIndex);
+      newSigmaTable.producedSigmaMasks(sigma.mask());
+    }
+  }
+
+  template <typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9>
+  void processSigmaPlus(T1 const& col, T2 const& /*sigmaPlusTable*/, T3 const& oldTrackTable, T4& partitionSigmaPlus, T5& trackBuilder, T6& cache, T7& newSigmaPlusTable, T8& newTrackTable, T9& newCollisionTable)
+  {
+    auto sigmaPlusSlice = partitionSigmaPlus->sliceByCached(o2::aod::femtobase::stored::fColId, col.globalIndex(), cache);
+
+    for (auto const& sigmaPlus : sigmaPlusSlice) {
+
+      auto chaDaughter = oldTrackTable.rawIteratorAt(sigmaPlus.chaDauId() - oldTrackTable.offset());
+
+      int chaDaughterIndex = trackBuilder.getDaughterIndex(chaDaughter, newTrackTable, newCollisionTable);
+
+      newSigmaPlusTable.producedSigmaPluses(newCollisionTable.producedCollision.lastIndex(),
+                                            sigmaPlus.signedPt(),
+                                            sigmaPlus.eta(),
+                                            sigmaPlus.phi(),
+                                            sigmaPlus.mass(),
+                                            chaDaughterIndex);
+      newSigmaPlusTable.producedSigmaPlusMasks(sigmaPlus.mask());
+    }
+  }
+
+ private:
+  int mLimitSigma = 0;
+  int mLimitSigmaPlus = 0;
+};
+
 } // namespace kinkbuilder
 } // namespace o2::analysis::femto
 #endif // PWGCF_FEMTO_CORE_KINKBUILDER_H_

PWGCF/Femto/TableProducer/femtoProducerDerivedToDerived.cxx

@@ -14,6 +14,7 @@
 /// \author Anton Riedel, TU München, anton.riedel@tum.de
 
 #include "PWGCF/Femto/Core/collisionBuilder.h"
+#include "PWGCF/Femto/Core/kinkBuilder.h"
 #include "PWGCF/Femto/Core/partitions.h"
 #include "PWGCF/Femto/Core/trackBuilder.h"
 #include "PWGCF/Femto/Core/v0Builder.h"
@@ -41,6 +42,8 @@ struct FemtoProducerDerivedToDerived {
   using Tracks = o2::soa::Join<o2::aod::FTracks, o2::aod::FTrackMasks>;
   using Lambdas = o2::soa::Join<o2::aod::FLambdas, o2::aod::FLambdaMasks>;
   using K0shorts = o2::soa::Join<o2::aod::FK0shorts, o2::aod::FK0shortMasks>;
+  using Sigma = o2::soa::Join<o2::aod::FSigmas, o2::aod::FSigmaMasks>;
+  using SigmaPlus = o2::soa::Join<o2::aod::FSigmaPlus, o2::aod::FSigmaPlusMasks>;
 
   o2::framework::SliceCache cache;
 
@@ -74,12 +77,26 @@ struct FemtoProducerDerivedToDerived {
   o2::framework::Partition<K0shorts> k0shortPartition = MAKE_K0SHORT_PARTITION(k0shortSelection1);
   o2::framework::Preslice<K0shorts> perColK0shorts = o2::aod::femtobase::stored::fColId;
 
+  // kink builder
+  kinkbuilder::KinkBuilderDerivedToDerived kinkBuilder;
+  kinkbuilder::KinkBuilderDerivedToDerivedProducts kinkBuilderProducts;
+  kinkbuilder::ConfKinkTablesDerivedToDerived confKinkBuilder;
+
+  kinkbuilder::ConfSigmaSelection1 sigmaSelection1;
+  o2::framework::Partition<Sigma> sigmaPartition = MAKE_SIGMA_PARTITION(sigmaSelection1);
+  o2::framework::Preslice<Sigma> perColSigma = o2::aod::femtobase::stored::fColId;
+
+  kinkbuilder::ConfSigmaPlusSelection1 sigmaPlusSelection1;
+  o2::framework::Partition<SigmaPlus> sigmaPlusPartition = MAKE_SIGMAPLUS_PARTITION(sigmaPlusSelection1);
+  o2::framework::Preslice<SigmaPlus> perColSigmaPlus = o2::aod::femtobase::stored::fColId;
+
   void init(o2::framework::InitContext& /*context*/)
   {
     trackBuilder.init(confTrackBuilder);
     v0Builder.init(confV0Builder);
+    kinkBuilder.init(confKinkBuilder);
 
-    if ((doprocessTracks + doprocessTracksLambdas + doprocessTracksK0shorts) > 1) {
+    if ((doprocessTracks + doprocessTracksLambdas + doprocessTracksK0shorts + doprocessTracksSigma + doprocessTracksSigmaPlus) > 1) {
       LOG(fatal) << "Only one proccess function can be activated";
     }
   }
@@ -116,6 +133,28 @@ struct FemtoProducerDerivedToDerived {
     v0Builder.processK0shorts(col, k0shorts, tracks, k0shortPartition, trackBuilder, cache, v0BuilderProducts, trackBuilderProducts, collisionBuilderProducts);
   }
   PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksK0shorts, "Process k0short and tracks", false);
+
+  void processTracksSigma(FilteredCollision const& col, Tracks const& tracks, Sigma const& sigma)
+  {
+    if (trackBuilder.collisionHasTooFewTracks(col, tracks, trackPartition1, trackPartition2, cache) || kinkBuilder.collisionHasTooFewSigma(col, sigma, sigmaPartition, cache)) {
+      return;
+    }
+    collisionBuilder.processCollision(col, collisionBuilderProducts);
+    trackBuilder.processTracks(col, tracks, trackPartition1, trackPartition2, cache, trackBuilderProducts, collisionBuilderProducts);
+    kinkBuilder.processSigma(col, sigma, tracks, sigmaPartition, trackBuilder, cache, kinkBuilderProducts, trackBuilderProducts, collisionBuilderProducts);
+  }
+  PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksSigma, "Process sigma and tracks", false);
+
+  void processTracksSigmaPlus(FilteredCollision const& col, Tracks const& tracks, SigmaPlus const& sigmaplus)
+  {
+    if (trackBuilder.collisionHasTooFewTracks(col, tracks, trackPartition1, trackPartition2, cache) || kinkBuilder.collisionHasTooFewSigmaPlus(col, sigmaplus, sigmaPlusPartition, cache)) {
+      return;
+    }
+    collisionBuilder.processCollision(col, collisionBuilderProducts);
+    trackBuilder.processTracks(col, tracks, trackPartition1, trackPartition2, cache, trackBuilderProducts, collisionBuilderProducts);
+    kinkBuilder.processSigmaPlus(col, sigmaplus, tracks, sigmaPlusPartition, trackBuilder, cache, kinkBuilderProducts, trackBuilderProducts, collisionBuilderProducts);
+  }
+  PROCESS_SWITCH(FemtoProducerDerivedToDerived, processTracksSigmaPlus, "Process sigmaPlus and tracks", false);
 };
 
 o2::framework::WorkflowSpec defineDataProcessing(o2::framework::ConfigContext const& cfgc)