Update of DGCandidateProducer

EventFiltering/PWGUD/DGHelpers.h

@@ -71,7 +71,7 @@ struct DGSelector {
   template <typename CC, typename BC, typename BCs, typename TCs, typename FWs>
   int IsSelected(DGCutparHolder diffCuts, CC const& collision, BC& bc, BCs& bcRange, TCs& tracks, FWs& fwdtracks)
   {
-    LOGF(debug, "Collision %f BC %i", collision.collisionTime(), bc.globalBC());
+    LOGF(debug, "Collision %f BC %u", collision.collisionTime(), bc.globalBC());
     LOGF(debug, "Number of close BCs: %i", bcRange.size());
 
     // check that there are no FIT signals in any of the compatible BCs

PWGUD/Core/DGPIDSelector.h

@@ -21,7 +21,7 @@
 using namespace o2;
 
 using UDTracksFull = soa::Join<aod::UDTracks, aod::UDTrackCollisionIDs, aod::UDTracksPID>;
-using UDTrackFull = soa::Join<aod::UDTracks, aod::UDTrackCollisionIDs, aod::UDTracksPID>::iterator;
+using UDTrackFull = UDTracksFull::iterator;
 
 float particleMass(TDatabasePDG* pdg, int pid);
 

PWGUD/TableProducer/DGCandProducer.cxx

@@ -52,6 +52,7 @@
 #include "Framework/AnalysisTask.h"
 
 #include "EventFiltering/PWGUD/DGHelpers.h"
+#include "PWGUD/Core/DGMCHelpers.h"
 #include "PWGUD/Core/UDHelperFunctions.h"
 #include "PWGUD/DataModel/UDTables.h"
 
@@ -73,20 +74,19 @@ struct DGCandProducer {
     diffCuts = (DGCutparHolder)DGCuts;
   }
 
+  // data tables
   Produces<aod::UDCollisions> outputCollisions;
   Produces<aod::UDTracks> outputTracks;
   Produces<aod::UDTracksPID> outputTracksPID;
-  Produces<aod::UDTrackCollisionIDs> outputTrackColIDs;
-
-  // helper struct
-  struct FT0Info {
-    float amplitudeA = -1;
-    float amplitudeC = -1;
-    float timeA = -999.;
-    float timeC = -999.;
-    uint8_t triggerMask = 0;
-  };
+  Produces<aod::UDTracksExtra> outputTracksExtra;
+  Produces<aod::UDTrackCollisionIDs> outputTracksCollisionsId;
+
+  // MC tables
+  Produces<aod::UDMcCollisions> outputMcCollisions;
+  Produces<aod::UDMcParticles> outputMcParticles;
+  Produces<aod::UDMcTrackLabels> outputMcTrackLabels;
 
+  // data inputs
   using CCs = soa::Join<aod::Collisions, aod::EvSels>;
   using CC = CCs::iterator;
   using BCs = soa::Join<aod::BCsWithTimestamps, aod::BcSels, aod::Run3MatchedToBCSparse>;
@@ -96,14 +96,122 @@ struct DGCandProducer {
                         aod::TOFSignal, aod::pidTOFFullEl, aod::pidTOFFullMu, aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTOFFullPr>;
   using FWs = aod::FwdTracks;
 
-  void process(CC const& collision,
-               BCs const& bcs,
-               TCs& tracks,
-               FWs& fwdtracks,
-               aod::Zdcs& zdcs,
-               aod::FT0s& ft0s,
-               aod::FV0As& fv0as,
-               aod::FDDs& fdds)
+  // MC inputs
+  using MCCCs = soa::Join<aod::Collisions, aod::EvSels, aod::McCollisionLabels>;
+  using MCCC = MCCCs::iterator;
+  using MCTCs = soa::Join<aod::Tracks, aod::TracksExtra, aod::TrackSelection, aod::McTrackLabels, aod::TOFSignal,
+                          aod::pidTPCFullEl, aod::pidTPCFullMu, aod::pidTPCFullPi, aod::pidTPCFullKa, aod::pidTPCFullPr,
+
+                          aod::pidTOFFullEl, aod::pidTOFFullMu, aod::pidTOFFullPi, aod::pidTOFFullKa, aod::pidTOFFullPr>;
+  using MCTC = MCTCs::iterator;
+
+  // function to update UDTracks, UDTracksPID, and UDTracksExtra
+  template <typename TTrack, typename TBC>
+  void updateUDTrackTables(TTrack const& track, TBC const& bc)
+  {
+    outputTracks(track.px(), track.py(), track.pz(), track.sign(),
+                 bc.globalBC(), track.trackTime(), track.trackTimeRes());
+    outputTracksPID(track.tpcNSigmaEl(),
+                    track.tpcNSigmaMu(),
+                    track.tpcNSigmaPi(),
+                    track.tpcNSigmaKa(),
+                    track.tpcNSigmaPr(),
+                    track.tofNSigmaEl(),
+                    track.tofNSigmaMu(),
+                    track.tofNSigmaPi(),
+                    track.tofNSigmaKa(),
+                    track.tofNSigmaPr());
+    outputTracksExtra(track.itsClusterMap(),
+                      track.tpcNClsFindable(),
+                      track.tpcNClsFindableMinusFound(),
+                      track.tpcNClsFindableMinusCrossedRows(),
+                      track.tpcNClsShared(),
+                      track.trdPattern(),
+                      track.itsChi2NCl(),
+                      track.tpcChi2NCl(),
+                      track.trdChi2(),
+                      track.tofChi2(),
+                      track.tpcSignal(),
+                      track.trdSignal(),
+                      track.length(),
+                      track.tofExpMom(),
+                      track.detectorMap());
+    outputTracksCollisionsId(outputCollisions.lastIndex());
+  }
+
+  // this function properly updates UDMcCollisions and UDMcParticles and returns the value
+  // deltaIndex, which is needed to correct the McParticles indices
+  // For a given McCollision all associated McParticles are saved
+  template <typename TMcCollision, typename TMcParticles, typename TBC>
+  void updateMcUDTables(TMcCollision const& McCol,
+                        TMcParticles const& McParts,
+                        TBC const& mcbc,
+                        int64_t& deltaIndex)
+  {
+    // save McCol
+    outputMcCollisions(mcbc.globalBC(),
+                       McCol.generatorsID(),
+                       McCol.posX(),
+                       McCol.posY(),
+                       McCol.posZ(),
+                       McCol.t(),
+                       McCol.weight(),
+                       McCol.impactParameter());
+
+    // save McParts
+    // calculate conversion from old indices to new indices
+    // old = mcpart.globalIndex()
+    // new = old + deltaIndex
+    // deltaIndex = [outputMcParticles.lastIndex() - McParts.iteratorAt(0).globalIndex() + 1]
+    deltaIndex = outputMcParticles.lastIndex() - McParts.iteratorAt(0).globalIndex() + 1;
+    LOGF(debug, " deltaIndex %i", deltaIndex);
+
+    // new mother and daughter ids
+    std::vector<int32_t> newmids;
+    int32_t newdids[2] = {-1, -1};
+
+    // all particles of the McCollision are saved
+    for (auto mcpart : McParts) {
+      // correct mother and daughter IDs
+      newmids.clear();
+      auto oldmids = mcpart.mothersIds();
+      for (uint ii = 0; ii < oldmids.size(); ii++) {
+        auto newval = oldmids[ii] < 0 ? oldmids[ii] : oldmids[ii] + deltaIndex;
+        LOGF(debug, " mid %i / %i", oldmids[ii], newval);
+        newmids.push_back(newval);
+      }
+      auto olddids = mcpart.daughtersIds();
+      for (uint ii = 0; ii < olddids.size(); ii++) {
+        auto newval = olddids[ii] < 0 ? olddids[ii] : olddids[ii] + deltaIndex;
+        LOGF(debug, " did %i / %i", olddids[ii], newval);
+        newdids[ii] = newval;
+      }
+      LOGF(debug, " ms %i ds %i", oldmids.size(), olddids.size());
+
+      // update UDMcParticles
+      outputMcParticles(outputMcCollisions.lastIndex(),
+                        mcpart.pdgCode(),
+                        mcpart.statusCode(),
+                        mcpart.flags(),
+                        newmids,
+                        newdids,
+                        mcpart.weight(),
+                        mcpart.px(),
+                        mcpart.py(),
+                        mcpart.pz(),
+                        mcpart.e());
+    }
+  }
+
+  // process function for real data
+  void processData(CC const& collision,
+                   BCs const& bcs,
+                   TCs& tracks,
+                   FWs& fwdtracks,
+                   aod::Zdcs& zdcs,
+                   aod::FT0s& ft0s,
+                   aod::FV0As& fv0as,
+                   aod::FDDs& fdds)
   {
     // nominal BC
     auto bc = collision.bc_as<BCs>();
@@ -116,25 +224,121 @@ struct DGCandProducer {
 
     // save DG candidates
     if (isDGEvent == 0) {
-      FT0Info ft0Info;
+      LOGF(info, "  Data: good collision!");
 
-      // update DGCandidates tables
-      outputCollisions(bc.runNumber(), bc.timestamp(),
+      // update DG candidates tables
+      outputCollisions(bc.globalBC(), bc.runNumber(),
                        collision.posX(), collision.posY(), collision.posZ(),
-                       collision.numContrib(), netCharge(tracks), rPVtrwTOF(tracks, collision.numContrib()),
-                       ft0Info.amplitudeA, ft0Info.amplitudeC, ft0Info.timeA, ft0Info.timeC, ft0Info.triggerMask);
+                       collision.numContrib(), netCharge(tracks),
+                       rPVtrwTOF(tracks, collision.numContrib()),
+                       0., 0., 0., 0., 0);
 
       // update DGTracks tables
       for (auto& track : tracks) {
         if (track.isPVContributor()) {
-          outputTracks(track.px(), track.py(), track.pz(), track.sign(), bc.globalBC(), track.trackTime(), track.trackTimeRes());
-          outputTrackColIDs(outputCollisions.lastIndex());
-          outputTracksPID(track.tpcNSigmaEl(), track.tpcNSigmaMu(), track.tpcNSigmaPi(), track.tpcNSigmaKa(), track.tpcNSigmaPr(),
-                          track.tofNSigmaEl(), track.tofNSigmaMu(), track.tofNSigmaPi(), track.tofNSigmaKa(), track.tofNSigmaPr());
+          updateUDTrackTables(track, bc);
         }
       }
     }
   }
+
+  PROCESS_SWITCH(DGCandProducer, processData, "Process real data", false);
+
+  Preslice<MCTCs> tracksPerCollision = aod::track::collisionId;
+  Preslice<FWs> fwdTracksPerCollision = aod::fwdtrack::collisionId;
+
+  // process function for MC data
+  void processMc(aod::McCollision const& McCol,
+                 aod::McParticles const& McParts,
+                 MCCCs const& collisions,
+                 BCs const& bcs,
+                 MCTCs const& tracks,
+                 FWs const& fwdtracks,
+                 aod::Zdcs const& zdcs,
+                 aod::FT0s const& ft0s,
+                 aod::FV0As const& fv0as,
+                 aod::FDDs const& fdds) //)
+  {
+    for (auto McPart : McParts) {
+      LOGF(debug, "McCol %i McPart %i", McCol.globalIndex(), McPart.globalIndex());
+    }
+
+    // is this a central diffractive event?
+    // by default it is assumed to be a MB event
+    bool isPythiaDiff = isPythiaCDE(McParts);
+    bool isGraniittiDiff = isGraniittiCDE(McParts);
+    LOGF(debug, "mcCol %i type %i / %i / %i", (int)McCol.globalIndex(), !isPythiaDiff && !isGraniittiDiff, isPythiaDiff, isGraniittiDiff);
+
+    // MC BC
+    auto mcbc = McCol.bc_as<BCs>();
+
+    // save MCTruth of all diffractive events
+    bool mcColIsSaved = false;
+    int64_t deltaIndex = 0;
+    auto nMcParts0 = outputMcParticles.lastIndex();
+    if (isPythiaDiff || isGraniittiDiff) {
+      // update tables UDMcCollisions and UDMcParticles
+      updateMcUDTables(McCol, McParts, mcbc, deltaIndex);
+      mcColIsSaved = true;
+    }
+
+    // loop over all reconstructed collisions associated with McCol
+    for (auto collision : collisions) {
+      // get the tracks belonging to collision
+      auto collisionTracks = tracks.sliceBy(tracksPerCollision, collision.globalIndex());
+      auto collisionFwdTracks = fwdtracks.sliceBy(fwdTracksPerCollision, collision.globalIndex());
+      LOGF(debug, "  tracks %i / %i", (int)collisionTracks.size(), collisionFwdTracks.size());
+
+      auto bc = collision.bc_as<BCs>();
+      LOGF(debug, "  BC mc %i reco %i", mcbc.globalBC(), bc.globalBC());
+
+      // is this a collision to be saved?
+      // obtain slice of compatible BCs
+      auto bcRange = MCcompatibleBCs(collision, diffCuts.NDtcoll(), bcs, diffCuts.minNBCs());
+
+      // apply DG selection
+      auto isDGEvent = dgSelector.IsSelected(diffCuts, collision, bc, bcRange, collisionTracks, collisionFwdTracks);
+      LOGF(debug, "  isDG %i", (int)isDGEvent);
+
+      // save information of DG events
+      if (isDGEvent == 0) {
+        LOGF(info, "  MC: good collision!");
+
+        // update UDMcCollisions and UDMcParticles if not already done
+        if (!mcColIsSaved) {
+          updateMcUDTables(McCol, McParts, mcbc, deltaIndex);
+          mcColIsSaved = true;
+        }
+
+        // UDCollisions
+        outputCollisions(bc.globalBC(), bc.runNumber(),
+                         collision.posX(), collision.posY(), collision.posZ(),
+                         collision.numContrib(), netCharge(tracks),
+                         rPVtrwTOF(collisionTracks, collision.numContrib()),
+                         0., 0., 0., 0., 0);
+
+        // UDTracks, UDTrackCollisionID, UDTracksExtras, UDMcTrackLabels
+        for (auto& track : collisionTracks) {
+          // but save only the Primary Vertex tracks
+          if (track.isPVContributor()) {
+            updateUDTrackTables(track, bc);
+
+            // properly correct the index into the UDMcParticles tables with deltaIndex
+            auto newval = track.mcParticleId() < 0 ? track.mcParticleId() : track.mcParticleId() + deltaIndex;
+            // only associations with McParticles belonging to the actual McCollision are supported
+            if ((newval < nMcParts0) || (newval > outputMcParticles.lastIndex())) {
+              LOGF(info, "<ATTENTION> UDMcParticles index out of range %i (%i - %i)", newval, nMcParts0 + 1, outputMcParticles.lastIndex());
+              newval = -1;
+            }
+            outputMcTrackLabels(newval,
+                                track.mcMask());
+          }
+        }
+      }
+    }
+  }
+
+  PROCESS_SWITCH(DGCandProducer, processMc, "Process MC data", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)

PWGUD/Tasks/DGCandAnalyzer.cxx

@@ -72,8 +72,6 @@ struct DGCandAnalyzer {
     pidsel.init(anaPars);
   }
 
-  using UDTracksFull = soa::Join<aod::UDTracks, aod::UDTrackCollisionIDs, aod::UDTracksPID>;
-
   void process(aod::UDCollision const& dgcand, UDTracksFull const& dgtracks)
   {