[PWGLF] Add on-the-fly coal to h-d corr

PWGLF/Tasks/Nuspex/hadronnucleicorrelation.cxx

@@ -46,6 +46,7 @@
 #include <TParticlePDG.h>
 #include <TString.h>
 
+#include <algorithm>
 #include <chrono>
 #include <cmath>
 #include <cstddef>
@@ -91,11 +92,20 @@ struct HadronNucleiCorrelation {
   Configurable<bool> doQA{"doQA", true, "save QA histograms"};
   Configurable<bool> doMCQA{"doMCQA", false, "save MC QA histograms"};
   Configurable<bool> isMC{"isMC", false, "is MC"};
-  Configurable<bool> isMCGen{"isMCGen", false, "is isMCGen"};
   Configurable<bool> isPrim{"isPrim", true, "is isPrim"};
   Configurable<bool> doCorrection{"doCorrection", false, "do efficiency correction"};
   Configurable<bool> doQuadraticPID{"doQuadraticPID", false, "do PID with sum in quadrature of TOF and TPC"};
 
+  struct : ConfigurableGroup {
+    std::string prefix = "Coalescence"; // JSON group name
+    Configurable<bool> doMCGenCoalescence{"doMCGenCoalescence", false, "do a simple coalescence on the generated level"};
+    // Coalescence parameters used in Eur. Phys. J. A 59 (2023) 72:
+    //   p_max = 0.148 GeV/c
+    //   r_max = 2 fm
+    Configurable<float> pMax{"pMax", 0.148, "maximum momentum for coalescence"};
+    Configurable<float> rMax{"rMax", 2.0, "maximum radius for coalescence"};
+  } settingsCoalescence;
+
   Configurable<std::string> fCorrectionPath{"fCorrectionPath", "", "Correction path to file"};
   Configurable<std::string> fCorrectionHisto{"fCorrectionHisto", "", "Correction histogram"};
   Configurable<std::string> cfgUrl{"cfgUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
@@ -516,11 +526,11 @@ struct HadronNucleiCorrelation {
         registry.add("hReco_Pt_Proton_TPCEl", "Reco (anti)protons in reco collisions", {HistType::kTH1F, {ptAxisSmall}});
         registry.add("hReco_Pt_Proton_TPCEl_or_TOF", "Reco (anti)protons in reco collisions", {HistType::kTH1F, {ptAxisSmall}});
         registry.add("hReco_Pt_Deuteron_TPCEl", "Reco (anti)deuterons in reco collisions", {HistType::kTH1F, {ptAxisSmall}});
-        registry.add("hReco_Pt_Deuteron_TPCEl_or_TOF", "Reco (anti)protons in reco collisions", {HistType::kTH1F, {ptAxisSmall}});
+        registry.add("hReco_Pt_Deuteron_TPCEl_or_TOF", "Reco (anti)deuterons in reco collisions", {HistType::kTH1F, {ptAxisSmall}});
       }
     }
 
-    if (isMCGen) {
+    if (doprocessMixedEventGen || doprocessSameEventGen) {
       registry.add("Generated/hNEventsMC", "hNEventsMC", {HistType::kTH1D, {{1, 0.f, 1.f}}});
       registry.get<TH1>(HIST("Generated/hNEventsMC"))->GetXaxis()->SetBinLabel(1, "All");
 
@@ -549,6 +559,8 @@ struct HadronNucleiCorrelation {
       registry.add("Generated/hAntiDeuteronsVsPt", "hAntiDeuteronsVsPt", {HistType::kTH1D, {ptAxisGen}});
       registry.add("Generated/hProtonsVsPt", "hProtonsVsPt", {HistType::kTH1D, {ptAxisGen}});
       registry.add("Generated/hAntiProtonsVsPt", "hAntiProtonsVsPt", {HistType::kTH1D, {ptAxisGen}});
+    } else if (settingsCoalescence.doMCGenCoalescence.value) {
+      LOG(fatal) << "This is not a Gen Run but the Gen coalescence is required. Turn off doMCGenCoalescence";
     }
   }
 
@@ -1695,10 +1707,253 @@ struct HadronNucleiCorrelation {
 
     registry.fill(HIST("Generated/hNEventsMC"), 0.5);
 
+    // Local representation of a generated particle.
+    struct GenCoalescenceCandidate {
+      GenCoalescenceCandidate(float pt,
+                              float eta,
+                              float phi,
+                              int pdg,
+                              float vx,
+                              float vy,
+                              float vz,
+                              float vt,
+                              bool physicalPrimary)
+        : mKinematics{pt, eta, phi},
+          mPdg{pdg},
+          mVx{vx},
+          mVy{vy},
+          mVz{vz},
+          mVt{vt},
+          mPhysicalPrimary{physicalPrimary}
+      {
+      }
+      [[nodiscard]] float pt() const { return mKinematics.pt(); }
+      [[nodiscard]] float eta() const { return mKinematics.eta(); }
+      [[nodiscard]] float phi() const { return mKinematics.phi(); }
+
+      [[nodiscard]] float px() const { return pt() * std::cos(phi()); }
+      [[nodiscard]] float py() const { return pt() * std::sin(phi()); }
+      [[nodiscard]] float pz() const { return mKinematics.pz(); }
+
+      [[nodiscard]] float vx() const { return mVx; }
+      [[nodiscard]] float vy() const { return mVy; }
+      [[nodiscard]] float vz() const { return mVz; }
+      [[nodiscard]] float vt() const { return mVt; }
+
+      [[nodiscard]] float mass() const
+      {
+        switch (std::abs(mPdg)) {
+          case PDG_t::kProton:
+            return o2::track::PID::getMass(o2::track::PID::Proton);
+          case PDG_t::kNeutron:
+            return o2::constants::physics::MassNeutron;
+          case o2::constants::physics::Pdg::kDeuteron:
+            return o2::track::PID::getMass(o2::track::PID::Deuteron);
+          default:
+            LOG(fatal) << "Unhandled pdg " << mPdg;
+            return 0.f;
+        }
+      }
+      [[nodiscard]] float energy() const { return std::hypot(mKinematics.p(), mass()); }
+      [[nodiscard]] float rapidity() const { return mKinematics.rapidityForMass(mass()); }
+      [[nodiscard]] float y() const { return rapidity(); }
+      [[nodiscard]] bool isPhysicalPrimary() const { return mPhysicalPrimary; }
+      [[nodiscard]] int pdgCode() const { return mPdg; }
+
+     private:
+      GenCandidate mKinematics;
+
+      int mPdg = 0;
+      float mVx = 0.f;
+      float mVy = 0.f;
+      float mVz = 0.f;
+      float mVt = 0.f;
+
+      bool mPhysicalPrimary = false;
+    };
+
+    std::vector<GenCoalescenceCandidate> particlesToProcess;
+    particlesToProcess.reserve(mcParticles.size());
+
+    for (const auto& particle : mcParticles) {
+      switch (particle.pdgCode()) {
+        case PDG_t::kProton:
+        case PDG_t::kProtonBar:
+        case PDG_t::kNeutron:
+        case PDG_t::kNeutronBar:
+        case o2::constants::physics::Pdg::kDeuteron:
+        case -o2::constants::physics::Pdg::kDeuteron:
+          particlesToProcess.emplace_back(particle.pt(),
+                                          particle.eta(),
+                                          particle.phi(),
+                                          particle.pdgCode(),
+                                          particle.vx(),
+                                          particle.vy(),
+                                          particle.vz(),
+                                          particle.vt(),
+                                          particle.isPhysicalPrimary());
+          break;
+        default:
+          break;
+      }
+    }
+
+    if (settingsCoalescence.doMCGenCoalescence.value) {
+      std::vector<bool> consumed(particlesToProcess.size(), false);
+      std::vector<GenCoalescenceCandidate> deuterons;
+
+      // Try p+n -> d and pbar+nbar -> dbar independently.
+      for (const int sign : {+1, -1}) {
+        const int protonPDG = sign * PDG_t::kProton;
+        const int neutronPDG = sign * PDG_t::kNeutron;
+        const int deuteronPDG = sign * o2::constants::physics::Pdg::kDeuteron;
+
+        for (size_t ip = 0; ip < particlesToProcess.size(); ++ip) {
+          if (consumed[ip] || particlesToProcess[ip].pdgCode() != protonPDG) {
+            continue;
+          }
+
+          const auto& proton = particlesToProcess[ip];
+
+          for (size_t in = 0; in < particlesToProcess.size(); ++in) {
+            if (consumed[in] || particlesToProcess[in].pdgCode() != neutronPDG) {
+              continue;
+            }
+
+            const auto& neutron = particlesToProcess[in];
+
+            const float ep = proton.energy();
+            const float en = neutron.energy();
+
+            // Propagate the particle freezing out first to the freeze-out time of the other particle
+            float xp = proton.vx();
+            float yp = proton.vy();
+            float zp = proton.vz();
+            const float tp = proton.vt();
+
+            float xn = neutron.vx();
+            float yn = neutron.vy();
+            float zn = neutron.vz();
+            const float tn = neutron.vt();
+
+            const float commonTime = std::max(static_cast<float>(tp), static_cast<float>(tn));
+
+            if (tp < commonTime) {
+              const float dt = commonTime - tp;
+              xp += proton.px() / ep * dt;
+              yp += proton.py() / ep * dt;
+              zp += proton.pz() / ep * dt;
+            }
+
+            if (tn < commonTime) {
+              const float dt = commonTime - tn;
+              xn += neutron.px() / en * dt;
+              yn += neutron.py() / en * dt;
+              zn += neutron.pz() / en * dt;
+            }
+
+            // Velocity of the p-n centre-of-mass frame.
+            const float totalE = ep + en;
+            const float totalPx = proton.px() + neutron.px();
+            const float totalPy = proton.py() + neutron.py();
+            const float totalPz = proton.pz() + neutron.pz();
+
+            const float bx = totalPx / totalE;
+            const float by = totalPy / totalE;
+            const float bz = totalPz / totalE;
+
+            const float beta2 = bx * bx + by * by + bz * bz;
+            if (beta2 >= 1.) {
+              continue;
+            }
+
+            const float gamma = 1. / std::sqrt(1. - beta2);
+
+            // Boost the proton momentum into the p-n rest frame.
+            // In that frame p_p* = -p_n*, therefore |p_p*| is the relative momentum entering the coalescence cut.
+            float pxStar = proton.px();
+            float pyStar = proton.py();
+            float pzStar = proton.pz();
+
+            if (beta2 > 0.) {
+              const float betaDotP = bx * proton.px() + by * proton.py() + bz * proton.pz();
+
+              const float factor = ((gamma - 1.) * betaDotP / beta2) - gamma * ep;
+
+              pxStar += factor * bx;
+              pyStar += factor * by;
+              pzStar += factor * bz;
+            }
+
+            const float pRelative = std::sqrt(pxStar * pxStar + pyStar * pyStar + pzStar * pzStar);
+
+            if (pRelative >= settingsCoalescence.pMax.value) {
+              continue;
+            }
+
+            // The two particles have already been propagated to equal time in the lab. Transform their spatial separation to the pair CM frame.
+            float dx = xp - xn;
+            float dy = yp - yn;
+            float dz = zp - zn;
+
+            if (beta2 > 0.) {
+              const float betaDotR = bx * dx + by * dy + bz * dz;
+              const float factor = (gamma - 1.) * betaDotR / beta2;
+
+              dx += factor * bx;
+              dy += factor * by;
+              dz += factor * bz;
+            }
+
+            const float rRelative = std::sqrt(dx * dx + dy * dy + dz * dz);
+
+            if (rRelative >= settingsCoalescence.rMax.value) {
+              continue;
+            }
+
+            // Successful coalescence: remove the constituent proton and neutron and replace them by a deuteron carrying their total three-momentum
+            consumed[ip] = true;
+            consumed[in] = true;
+
+            const float deuteronPt = std::hypot(totalPx, totalPy);
+            const float deuteronPhi = std::atan2(totalPy, totalPx);
+            const float deuteronEta = std::asinh(totalPz / std::max(deuteronPt, 1.e-12f));
+
+            deuterons.emplace_back(deuteronPt,
+                                   deuteronEta,
+                                   deuteronPhi,
+                                   deuteronPDG,
+                                   0.5f * (xp + xn),
+                                   0.5f * (yp + yn),
+                                   0.5f * (zp + zn),
+                                   commonTime,
+                                   proton.isPhysicalPrimary() && neutron.isPhysicalPrimary());
+            break;
+          }
+        }
+      }
+
+      // Construct the post-coalescence particle list.
+      std::vector<GenCoalescenceCandidate> coalescedParticles;
+      coalescedParticles.reserve(particlesToProcess.size() + deuterons.size());
+
+      for (size_t i = 0; i < particlesToProcess.size(); ++i) {
+        if (!consumed[i]) {
+          coalescedParticles.push_back(std::move(particlesToProcess[i]));
+        }
+      }
+
+      for (auto& deuteron : deuterons) {
+        coalescedParticles.push_back(std::move(deuteron));
+      }
+
+      particlesToProcess = std::move(coalescedParticles);
+    }
+
     // Pairing candidates are collected during the QA loop below (which already visits every particle)
     GenCollisionCache genCache;
 
-    for (const auto& particle : mcParticles) {
+    for (const auto& particle : particlesToProcess) {
       auto fillGeneratedQa = [this, &particle](const float binPosition) {
         switch (particle.pdgCode()) {
           case PDG_t::kProton: