implement cascade processing in v0selector (1-st working version)

Author: lubynets

PWGDQ/Tasks/v0selector.cxx

@@ -49,6 +49,7 @@
 #include <map>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 using namespace o2;
@@ -89,6 +90,18 @@ struct v0selector {
   Configurable<float> cutAPL1{"cutAPL1", 0.107, "cutAPL1"};
   Configurable<float> cutAPL2{"cutAPL2", -0.69, "cutAPL2"};
   Configurable<float> cutAPL3{"cutAPL3", 0.5, "cutAPL3"};
+  // Omega & A-Omega cuts
+  Configurable<float> cutAPOmegaUp1{"cutAPOmegaUp1", 0.25, "cutAPOmegaUp1"};
+  Configurable<float> cutAPOmegaUp2{"cutAPOmegaUp2", 0.358, "cutAPOmegaUp2"};
+  Configurable<float> cutAPOmegaUp3{"cutAPOmegaUp3", 0.35, "cutAPOmegaUp3"};
+  Configurable<float> cutAPOmegaDown1{"cutAPOmegaDown1", 0.15, "cutAPOmegaDown1"};
+  Configurable<float> cutAPOmegaDown2{"cutAPOmegaDown2", 0.358, "cutAPOmegaDown2"};
+  Configurable<float> cutAPOmegaDown3{"cutAPOmegaDown3", 0.16, "cutAPOmegaDown3"};
+  Configurable<float> cutAlphaOmegaHigh{"cutAlphaOmegaHigh", 0.358, "cutAlphaOmegaHigh"};
+  Configurable<float> cutAlphaOmegaLow{"cutAlphaOmegaLow", 0., "cutAlphaOmegaLow"};
+  Configurable<float> cutMassOmegaHigh{"cutMassOmegaHigh", 1.677, "cutMassOmegaHigh"};
+  Configurable<float> cutMassOmegaLow{"cutMassOmegaLow", 1.667, "cutMassOmegaLow"};
+
   Configurable<bool> produceV0ID{"produceV0ID", false, "Produce additional V0ID table"};
   Configurable<bool> produceCascID{"produceCascID", false, "Produce additional CascID table"};
 
@@ -104,7 +117,7 @@ struct v0selector {
 
   Produces<o2::aod::V0Bits> v0bits;
   Produces<o2::aod::V0MapID> v0mapID;
-  Produces<o2::aod::CascMapID> cascMapID;
+  Produces<o2::aod::CascMapID> cascmapID;
 
   // int checkV0(const array<float, 3>& ppos, const array<float, 3>& pneg)
   int checkV0(const float alpha, const float qt)
@@ -161,10 +174,49 @@ struct v0selector {
     return kUndef;
   }
 
-  int checkCascade()
+  std::pair<float, float> evalArmenterosPodolanskiCascade(aod::CascData const& casc)
+  {
+    const float pxv0 = casc.pxpos() + casc.pxneg();
+    const float pyv0 = casc.pypos() + casc.pyneg();
+    const float pzv0 = casc.pzpos() + casc.pzneg();
+
+    const float pxbach = casc.pxbach();
+    const float pybach = casc.pybach();
+    const float pzbach = casc.pzbach();
+
+    const double momTot = RecoDecay::p2(pxv0 + pxbach, pyv0 + pybach, pzv0 + pzbach);
+
+    const float lQlv0 = RecoDecay::dotProd(std::array{pxv0, pyv0, pzv0}, std::array{pxv0 + pxbach, pyv0 + pybach, pzv0 + pzbach}) / momTot;
+    const float lQlbach = RecoDecay::dotProd(std::array{pxbach, pybach, pzbach}, std::array{pxv0 + pxbach, pyv0 + pybach, pzv0 + pzbach}) / momTot;
+    const float dp = RecoDecay::dotProd(std::array{pxbach, pybach, pzbach}, std::array{pxv0 + pxbach, pyv0 + pybach, pzv0 + pzbach});
+
+    float alpha = (lQlv0 - lQlbach) / (lQlv0 + lQlbach);
+    const float qtarm = std::sqrt(RecoDecay::p2(pxbach, pybach, pzbach) - dp * dp / momTot / momTot);
+
+    if (casc.sign() > 0) {
+      alpha = -alpha;
+    }
+
+    return std::make_pair(alpha, qtarm);
+  }
+
+  int checkCascade(float alpha, float qt)
   {
-    // to be implemented;
-    return kOmega;
+    const bool isAlphaPos = alpha > 0;
+    alpha = std::fabs(alpha);
+
+    const float qUp = cutAPOmegaUp1 * std::sqrt(std::abs(1.0f - ((alpha - cutAPOmegaUp2) * (alpha - cutAPOmegaUp2)) / (cutAPOmegaUp3 * cutAPOmegaUp3)));
+    const float qDown = cutAPOmegaDown1 * std::sqrt(std::abs(1.0f - ((alpha - cutAPOmegaDown2) * (alpha - cutAPOmegaDown2)) / (cutAPOmegaDown3 * cutAPOmegaDown3)));
+
+    if (alpha < cutAlphaOmegaLow || alpha > cutAlphaOmegaHigh || qt < qDown || qt > qUp) {
+      return kUndef;
+    }
+
+    if (isAlphaPos) {
+      return kOmega;
+    } else {
+      return kAntiOmega;
+    }
   }
 
   // Configurables
@@ -210,6 +262,10 @@ struct v0selector {
       registry.add("hV0APplot", "hV0APplot", HistType::kTH2F, {{200, -1.0f, +1.0f}, {250, 0.0f, 0.25f}});
       registry.add("hV0APplotSelected", "hV0APplotSelected", HistType::kTH2F, {{200, -1.0f, +1.0f}, {250, 0.0f, 0.25f}});
       registry.add("hV0Psi", "hV0Psi", HistType::kTH2F, {{100, 0, TMath::PiOver2()}, {100, 0, 0.1}});
+      registry.add("hCascAPplot", "hCascAPplot", HistType::kTH2F, {{200, -1.0f, +1.0f}, {500, 0.0f, 0.5f}});
+      registry.add("hCascAPplotSelected", "hCascAPplotSelected", HistType::kTH2F, {{200, -1.0f, +1.0f}, {500, 0.0f, 0.5f}});
+      registry.add("hMassOmega", "hMassOmega", HistType::kTH2F, {{900, 0.0f, 90.0f}, {100, 1.62f, 1.72f}});
+      registry.add("hMassAntiOmega", "hMassAntiOmega", HistType::kTH2F, {{900, 0.0f, 90.0f}, {100, 1.62f, 1.72f}});
     }
   }
 
@@ -395,6 +451,7 @@ struct v0selector {
         v0mapID(v0pidmap[V0.globalIndex()]);
       }
     }
+
     for (const auto& Casc : Cascs) {
       if (fillhisto) {
         registry.fill(HIST("hCascCandidate"), 1);
@@ -406,9 +463,53 @@ struct v0selector {
         registry.fill(HIST("hCascCandidate"), 2);
       }
 
+      const float Cascradius = Casc.cascradius();
+
       const float mOmega = Casc.mOmega();
 
+      auto [alpha, qt] = evalArmenterosPodolanskiCascade(Casc);
+
+      if (fillhisto) {
+        registry.fill(HIST("hCascAPplot"), alpha, qt);
+      }
+
+      const int cascid = checkCascade(alpha, qt);
+      if (cascid < 0) {
+        continue;
+      }
+      if (fillhisto) {
+        registry.fill(HIST("hCascAPplotSelected"), alpha, qt);
+      }
+
+      auto storeCascAddID = [&](auto gix, auto id) {
+        if (produceCascID.value) {
+          cascpidmap[gix] = id;
+        }
+      };
+
+      if (cascid == kOmega) {
+        if (fillhisto) {
+          registry.fill(HIST("hMassOmega"), Cascradius, mOmega);
+        }
+        if (cutMassOmegaLow < mOmega && mOmega < cutMassOmegaHigh) {
+          pidmap[Casc.bachelorId()] |= (uint8_t(1) << kOmega);
+          storeCascAddID(Casc.globalIndex(), kOmega);
+        }
+      } else if (cascid == kAntiOmega) {
+        if (fillhisto) {
+          registry.fill(HIST("hMassAntiOmega"), Cascradius, mOmega);
+        }
+        if (cutMassOmegaLow < mOmega && mOmega < cutMassOmegaHigh) {
+          pidmap[Casc.bachelorId()] |= (uint8_t(1) << kAntiOmega);
+          storeCascAddID(Casc.globalIndex(), kAntiOmega);
+        }
+      }
     } // end of Casc loop
+    if (produceCascID.value) {
+      for (auto& Casc : Cascs) {
+        cascmapID(cascpidmap[Casc.globalIndex()]);
+      }
+    }
     for (auto& track : tracks) {
       // printf("setting pidmap[%lld] = %d\n",track.globalIndex(),pidmap[track.globalIndex()]);
       v0bits(pidmap[track.globalIndex()]);