From 88b57faa0e26b7882d7450b2b2a038ab612b64d3 Mon Sep 17 00:00:00 2001
From: Nicolas Wirth <74405269+wirthni@users.noreply.github.com>
Date: Sat, 27 Jun 2026 07:52:28 +0200
Subject: [PATCH 1/2] Refactor collision index handling in diffWakeTreeProducer

Ci is added, because the produced tree is used out of an O2 analysis framework (bare root) and a direct correspondence between collisions and tracks is needed
---
 PWGJE/TableProducer/diffWakeTreeProducer.cxx | 22 +++++++++++++-------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/PWGJE/TableProducer/diffWakeTreeProducer.cxx b/PWGJE/TableProducer/diffWakeTreeProducer.cxx
index 5883df6c432..aeefb423cf8 100644
--- a/PWGJE/TableProducer/diffWakeTreeProducer.cxx
+++ b/PWGJE/TableProducer/diffWakeTreeProducer.cxx
@@ -22,13 +22,11 @@
 #include "Common/DataModel/TrackSelectionTables.h"
 
 #include <Framework/ASoA.h>
-#include <Framework/AnalysisDataModel.h>
 #include <Framework/AnalysisHelpers.h>
 #include <Framework/AnalysisTask.h>
 #include <Framework/HistogramRegistry.h>
 #include <Framework/HistogramSpec.h>
 #include <Framework/InitContext.h>
-#include <Framework/OutputObjHeader.h>
 #include <Framework/runDataProcessing.h>
 
 // Event selection: Only events that contain track above some threshold
@@ -67,7 +65,7 @@ namespace o2::aod
 namespace testcol
 {
 // Event properties
-// DECLARE_SOA_COLUMN(Gi, gi, int64_t);
+DECLARE_SOA_COLUMN(Ci, ci, int64_t);         // Collision index - needed because is used as tree
 DECLARE_SOA_COLUMN(Rn, rn, int32_t);         // run number
 DECLARE_SOA_COLUMN(Cent, cent, float);       // FT0C centrality
 DECLARE_SOA_COLUMN(Mult, mult, int32_t);     // TPC multiplicity
@@ -82,6 +80,7 @@ DECLARE_SOA_COLUMN(Psi3, psi3, int16_t);
 
 DECLARE_SOA_TABLE(TableCols, "AOD", "TABLECOL",
                   o2::soa::Index<>,
+                  testcol::Ci,
                   testcol::Rn,
                   testcol::Cent,
                   testcol::Mult,
@@ -99,6 +98,7 @@ namespace testtrack
 
 // Track properties
 DECLARE_SOA_INDEX_COLUMN(TableCol, tableCol);
+DECLARE_SOA_COLUMN(Ci, ci, int64_t); // is needed to link to collision when used as tree
 DECLARE_SOA_COLUMN(Charge, charge, int16_t);
 DECLARE_SOA_COLUMN(P, p, uint64_t);
 DECLARE_SOA_COLUMN(Dedx, dedx, uint16_t);
@@ -109,6 +109,7 @@ DECLARE_SOA_COLUMN(Dcaz, dcaz, int16_t);
 DECLARE_SOA_TABLE(TableTrack, "AOD", "TABLETRACK",
                   o2::soa::Index<>,
                   testtrack::TableColId,
+                  testtrack::Ci,
                   testtrack::Charge,
                   testtrack::P,
                   testtrack::Dedx,
@@ -196,7 +197,8 @@ struct DiffWakeTreeProducer {
     int16_t substituteEp2 = static_cast<int16_t>(ep2 * 1000);
     int16_t substituteEp3 = static_cast<int16_t>(ep3 * 1000);
 
-    testcol(run,
+    testcol(collisionCounter,
+            run,
             col.centFT0C(),
             col.multTPC(),
             col.trackOccupancyInTimeRange(),
@@ -232,27 +234,30 @@ struct DiffWakeTreeProducer {
 
       int64_t particlePx = (track.px() * 6000);
       if (particlePx < 0)
+      {
         substituteP |= static_cast<uint64_t>(1) << uppermostBit;
-      if (particlePx < 0)
         particlePx = (-1) * particlePx;
+      }
       substituteP |= (particlePx & bitmask20Bits) << lowermostBit;
 
       uppermostBit = 41;
       lowermostBit = 21;
       int64_t particlePy = (track.py() * 6000);
       if (particlePy < 0)
+      {
         substituteP |= static_cast<uint64_t>(1) << uppermostBit;
-      if (particlePy < 0)
         particlePy = (-1) * particlePy;
+      }
       substituteP |= (particlePy & bitmask20Bits) << lowermostBit;
 
       uppermostBit = 62;
       lowermostBit = 42;
       int64_t particlePz = (track.pz() * 6000);
       if (particlePz < 0)
+      {
         substituteP |= static_cast<uint64_t>(1) << uppermostBit;
-      if (particlePz < 0)
         particlePz = (-1) * particlePz;
+      }
       substituteP |= (particlePz & bitmask20Bits) << lowermostBit;
 
       // dEdx
@@ -263,7 +268,8 @@ struct DiffWakeTreeProducer {
       int16_t substituteDCAZ = static_cast<int16_t>(track.dcaZ() * 100);
 
       //--------------- Fill track table ------------------
-      testtrack(collisionCounter,
+      testtrack(testcol.lastIndex(),
+                collisionCounter,
                 track.sign(),
                 substituteP,
                 substituteDEDX,

From 12ea35bab4bdf3baa27cdf1877932d662196c1d9 Mon Sep 17 00:00:00 2001
From: Nicolas Wirth <74405269+wirthni@users.noreply.github.com>
Date: Sat, 27 Jun 2026 07:56:55 +0200
Subject: [PATCH 2/2] Refactor conditional statements for particle momentum

---
 PWGJE/TableProducer/diffWakeTreeProducer.cxx | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/PWGJE/TableProducer/diffWakeTreeProducer.cxx b/PWGJE/TableProducer/diffWakeTreeProducer.cxx
index aeefb423cf8..7e14f949bc9 100644
--- a/PWGJE/TableProducer/diffWakeTreeProducer.cxx
+++ b/PWGJE/TableProducer/diffWakeTreeProducer.cxx
@@ -233,8 +233,7 @@ struct DiffWakeTreeProducer {
       uint64_t bitmask20Bits = 0b11111111111111111111;
 
       int64_t particlePx = (track.px() * 6000);
-      if (particlePx < 0)
-      {
+      if (particlePx < 0) {
         substituteP |= static_cast<uint64_t>(1) << uppermostBit;
         particlePx = (-1) * particlePx;
       }
@@ -243,8 +242,7 @@ struct DiffWakeTreeProducer {
       uppermostBit = 41;
       lowermostBit = 21;
       int64_t particlePy = (track.py() * 6000);
-      if (particlePy < 0)
-      {
+      if (particlePy < 0) {
         substituteP |= static_cast<uint64_t>(1) << uppermostBit;
         particlePy = (-1) * particlePy;
       }
@@ -253,8 +251,7 @@ struct DiffWakeTreeProducer {
       uppermostBit = 62;
       lowermostBit = 42;
       int64_t particlePz = (track.pz() * 6000);
-      if (particlePz < 0)
-      {
+      if (particlePz < 0) {
         substituteP |= static_cast<uint64_t>(1) << uppermostBit;
         particlePz = (-1) * particlePz;
       }