From 7f9db57281056f0f4f92fff40c030da3af0f60c5 Mon Sep 17 00:00:00 2001
From: Fabrizio Grosa <fabrizio.grosa@cern.ch>
Date: Wed, 3 Jun 2026 14:58:14 +0200
Subject: [PATCH] Apply correction for energy (non)conservation in coalescence
 for c-deuteron

---
 .../generator_pythia8_embed_charmnuclei.C     | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/MC/config/PWGHF/external/generator/generator_pythia8_embed_charmnuclei.C b/MC/config/PWGHF/external/generator/generator_pythia8_embed_charmnuclei.C
index 4b6e9a3f7..9e2667e4f 100644
--- a/MC/config/PWGHF/external/generator/generator_pythia8_embed_charmnuclei.C
+++ b/MC/config/PWGHF/external/generator/generator_pythia8_embed_charmnuclei.C
@@ -9,6 +9,9 @@
 #include <string>
 #include <vector>
 
+#include "Math/Vector3D.h"
+#include "Math/Vector4D.h"
+
 R__ADD_INCLUDE_PATH($O2DPG_MC_CONFIG_ROOT)
 #include "MC/config/common/external/generator/CoalescencePythia8.h"
 
@@ -240,6 +243,7 @@ class GeneratorPythia8HFEmbedCharmNuclei : public o2::eventgen::GeneratorPythia8
         // we try the coalescence here, if successful we copy particles in the pythia event and we move to the next charm nucleus
         isCoalSuccess = CoalescencePythia8(mPythiaGun.event, std::vector<unsigned int>{1000010020}, mTrivialCoal, mCoalMomentum, dausToCoal[0], dausToCoal[1], 10.);
         if (isCoalSuccess) {
+          restoreEnergyConservation(mPythiaGun.event, idxCharmNucleus);
           int offset = mPythia.event.size(); // we need to rescale the indices of mothers and daughters, accounting for the particles that are already appended to the event
           for (int iPart{0}; iPart<mPythiaGun.event.size(); ++iPart) {
             auto part = mPythiaGun.event[iPart];
@@ -274,6 +278,32 @@ class GeneratorPythia8HFEmbedCharmNuclei : public o2::eventgen::GeneratorPythia8
 
 
 private:
+
+  void restoreEnergyConservation(Pythia8::Event& event, int idxCharmNucleus, float targetMassTolerance = 1e-5) {
+    /// In the coalescence, the energy is not conserved, we rescale the momentum of the charmed nuclei daughters to restore it to avoid changes in the invariant mass of the charmed nucleus
+
+    float scale = 1.f;
+    float invMass{0.f};
+    while (abs(invMass - mMassCharmNucleus) > targetMassTolerance) {
+      ROOT::Math::PxPyPzMVector fourVecCharmNucleus;
+      for (int iDau{event[idxCharmNucleus].daughter1()}; iDau<=event[idxCharmNucleus].daughter2(); ++iDau) {
+        auto dau = event[iDau];
+        fourVecCharmNucleus += ROOT::Math::PxPyPzMVector(dau.px() * scale, dau.py() * scale, dau.pz() * scale, dau.m());
+      }
+      invMass = fourVecCharmNucleus.M();
+      scale *= mMassCharmNucleus / invMass;
+    }
+
+      for (int iDau{event[idxCharmNucleus].daughter1()}; iDau<=event[idxCharmNucleus].daughter2(); ++iDau) {
+      event[iDau].px(event[iDau].px() * scale);
+      event[iDau].py(event[iDau].py() * scale);
+      event[iDau].pz(event[iDau].pz() * scale);
+    }
+    event[idxCharmNucleus].px(event[idxCharmNucleus].px() * scale);
+    event[idxCharmNucleus].py(event[idxCharmNucleus].py() * scale);
+    event[idxCharmNucleus].pz(event[idxCharmNucleus].pz() * scale);
+  }
+
   // Properties of selection
   float mMassCharmNucleus;           /// mass of the charmed nucleus
   int mPdgCharmNucleus;              /// pdg code of the charmed nucleus