[PWGLF] Update V0 analysis task

PWGLF/TableProducer/Strangeness/cascademlselection.cxx

@@ -1,4 +1,4 @@
 // Copyright 2019-2020 CERN and copyright holders of ALICE O2.
 // See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
 // All rights not expressly granted are reserved.
 //
@@ -8,7 +8,7 @@
 // In applying this license CERN does not waive the privileges and immunities
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
 //
 //  *+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*
 //  Lambdakzero ML selection task
 //  *+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+*
@@ -46,9 +46,9 @@
 #include <TFile.h>
 #include <TH2F.h>
 #include <TProfile.h>
 #include <TLorentzVector.h>
 #include <TPDGCode.h>
 #include <TDatabasePDG.h>
 #include "Tools/ML/MlResponse.h"
 #include "Tools/ML/model.h"
 
@@ -67,12 +67,15 @@
 // For derived data analysis
 using CascDerivedDatas = soa::Join<aod::CascCores, aod::CascExtras, aod::CascCollRefs>;
 
 struct cascademlselection {
   o2::ml::OnnxModel mlModelXiMinus;
   o2::ml::OnnxModel mlModelXiPlus;
   o2::ml::OnnxModel mlModelOmegaMinus;
   o2::ml::OnnxModel mlModelOmegaPlus;
 
+  // Custom grouping
+  std::vector<std::vector<int>> cascadesGrouped;
+
   std::map<std::string, std::string> metadata;
 
   Produces<aod::CascXiMLScores> xiMLSelections;    // optionally aggregate information from ML output for posterior analysis (derived data)
@@ -159,7 +162,7 @@
   }
 
   // function to load models for ML-based classifiers
   void LoadMachines(int64_t timeStampML)
   {
     if (mlConfigurations.loadModelsFromCCDB) {
       ccdbApi.init(ccdbConfigurations.ccdburl);
@@ -261,30 +264,52 @@
     }
   }
 
-  void processDerivedData(soa::Join<aod::StraCollisions, aod::StraStamps>::iterator const& collision, CascDerivedDatas const& cascades)
+  void processDerivedData(soa::Join<aod::StraCollisions, aod::StraStamps> const& collisions, CascDerivedDatas const& cascades)
   {
-    initCCDB(collision);
+    // Custom grouping
+    cascadesGrouped.clear();
+    cascadesGrouped.resize(collisions.size());
+
+    for (const auto& cascade : cascades) {
+      cascadesGrouped[cascade.straCollisionId()].push_back(cascade.globalIndex());
+    }
 
-    histos.fill(HIST("hEventVertexZ"), collision.posZ());
-    for (auto& casc : cascades) {
-      nCandidates++;
-      if (nCandidates % 50000 == 0) {
-        LOG(info) << "Candidates processed: " << nCandidates;
+    for (const auto& collision : collisions) {
+      initCCDB(collision);
+
+      histos.fill(HIST("hEventVertexZ"), collision.posZ());
+      for (std::size_t i = 0; i < cascadesGrouped[collision.globalIndex()].size(); i++) {
+        auto casc = cascades.rawIteratorAt(cascadesGrouped[collision.globalIndex()][i]);
+        nCandidates++;
+        if (nCandidates % 50000 == 0) {
+          LOG(info) << "Candidates processed: " << nCandidates;
+        }
+        processCandidate(casc);
       }
-      processCandidate(casc);
     }
   }
-  void processStandardData(aod::Collision const& collision, CascOriginalDatas const& cascades)
+  void processStandardData(aod::Collisions const& collisions, CascOriginalDatas const& cascades)
   {
-    initCCDB(collision);
+    // Custom grouping
+    cascadesGrouped.clear();
+    cascadesGrouped.resize(collisions.size());
 
-    histos.fill(HIST("hEventVertexZ"), collision.posZ());
-    for (auto& casc : cascades) {
-      nCandidates++;
-      if (nCandidates % 50000 == 0) {
-        LOG(info) << "Candidates processed: " << nCandidates;
+    for (const auto& cascade : cascades) {
+      cascadesGrouped[cascade.collisionId()].push_back(cascade.globalIndex());
+    }
+
+    for (const auto& collision : collisions) {
+      initCCDB(collision);
+
+      histos.fill(HIST("hEventVertexZ"), collision.posZ());
+      for (std::size_t i = 0; i < cascadesGrouped[collision.globalIndex()].size(); i++) {
+        auto casc = cascades.rawIteratorAt(cascadesGrouped[collision.globalIndex()][i]);
+        nCandidates++;
+        if (nCandidates % 50000 == 0) {
+          LOG(info) << "Candidates processed: " << nCandidates;
+        }
+        processCandidate(casc);
       }
-      processCandidate(casc);
     }
   }