Ctpdev: irframes

Detectors/CTP/workflowIO/src/DigitReaderSpec.cxx

@@ -17,7 +17,10 @@
 #include "DataFormatsCTP/LumiInfo.h"
 #include "Headers/DataHeader.h"
 #include "DetectorsCommonDataFormats/DetID.h"
+#include "SimulationDataFormat/MCCompLabel.h"
+#include "SimulationDataFormat/ConstMCTruthContainer.h"
 #include "CommonUtils/NameConf.h"
+#include "CommonUtils/IRFrameSelector.h"
 #include "Framework/DataProcessorSpec.h"
 #include "Framework/Task.h"
 #include "Framework/ControlService.h"
@@ -50,6 +53,7 @@ class DigitReader : public Task
   std::unique_ptr<TTree> mTree;
 
   bool mUseMC = false; // use MC truth
+  bool mUseIRFrames = false; // selected IRFrames mode
   std::string mDigTreeName = "o2sim";
   std::string mDigitBranchName = "CTPDigits";
   std::string mLumiBranchName = "CTPLumi";
@@ -58,29 +62,78 @@ class DigitReader : public Task
 DigitReader::DigitReader(bool useMC)
 {
   if (useMC) {
-    LOG(info) << "CTP does not support MC truth at the moment";
+    LOG(info) << "CTP : truth = data as CTP inputs are already digital";
   }
 }
 
 void DigitReader::init(InitContext& ic)
 {
   auto filename = o2::utils::Str::concat_string(o2::utils::Str::rectifyDirectory(ic.options().get<std::string>("input-dir")),
                                                 ic.options().get<std::string>("ctp-digit-infile"));
+  if (ic.options().hasOption("ignore-irframes") && !ic.options().get<bool>("ignore-irframes")) {
+    mUseIRFrames = true;
+  }
   connectTree(filename);
 }
 
 void DigitReader::run(ProcessingContext& pc)
 {
-  auto ent = mTree->GetReadEntry() + 1;
-  assert(ent < mTree->GetEntries()); // this should not happen
-
-  mTree->GetEntry(ent);
-  LOG(info) << "DigitReader pushes " << mDigits.size() << " digits at entry " << ent;
-  pc.outputs().snapshot(Output{"CTP", "DIGITS", 0}, mDigits);
-  pc.outputs().snapshot(Output{"CTP", "LUMI", 0}, mLumi);
-  if (mTree->GetReadEntry() + 1 >= mTree->GetEntries()) {
-    pc.services().get<ControlService>().endOfStream();
-    pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
+  gsl::span<const o2::dataformats::IRFrame> irFrames{};
+  // LOG(info) << "Using IRs:" << mUseIRFrames;
+  if (mUseIRFrames) {
+    irFrames = pc.inputs().get<gsl::span<o2::dataformats::IRFrame>>("driverInfo");
+  }
+  auto ent = mTree->GetReadEntry();
+  if (!mUseIRFrames) {
+    ent++;
+    assert(ent < mTree->GetEntries()); // this should not happen
+    mTree->GetEntry(ent);
+    LOG(info) << "DigitReader pushes " << mDigits.size() << " digits at entry " << ent;
+    pc.outputs().snapshot(Output{"CTP", "DIGITS", 0}, mDigits);
+    pc.outputs().snapshot(Output{"CTP", "LUMI", 0}, mLumi);
+    if (mTree->GetReadEntry() + 1 >= mTree->GetEntries()) {
+      pc.services().get<ControlService>().endOfStream();
+      pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
+    }
+  } else {
+    std::vector<o2::ctp::CTPDigit> digitSel;
+    if (irFrames.size()) { // we assume the IRFrames are in the increasing order
+      if (ent < 0) {
+        ent++;
+      }
+      o2::utils::IRFrameSelector irfSel;
+      // MC  digits are already aligned
+      irfSel.setSelectedIRFrames(irFrames, 0, 0, 0, true);
+      const auto irMin = irfSel.getIRFrames().front().getMin(); // use processed IRframes for rough comparisons (possible shift!)
+      const auto irMax = irfSel.getIRFrames().back().getMax();
+      LOGP(info, "Selecting IRFrame {}-{}", irMin.asString(), irMax.asString());
+      while (ent < mTree->GetEntries()) {
+        if (ent > mTree->GetReadEntry()) {
+          mTree->GetEntry(ent);
+        }
+        if (mDigits.front().intRecord <= irMax && mDigits.back().intRecord >= irMin) { // THere is overlap
+          for (int i = 0; i < (int)mDigits.size(); i++) {
+            const auto& dig = mDigits[i];
+            // if(irfSel.check(dig.intRecord)) { // adding selected digit
+            if (dig.intRecord >= irMin && dig.intRecord <= irMax) {
+              digitSel.push_back(dig);
+              LOG(info) << "adding:" << dig.intRecord << " ent:" << ent;
+            }
+          }
+        }
+        if (mDigits.back().intRecord < irMax) { // need to check the next entry
+          ent++;
+          continue;
+        }
+        break; // push collected data
+      }
+    }
+    pc.outputs().snapshot(Output{"CTP", "DIGITS", 0}, digitSel);
+    pc.outputs().snapshot(Output{"CTP", "LUMI", 0}, mLumi); // add full lumi for this TF
+    if (!irFrames.size() || irFrames.back().isLast()) {
+      pc.services().get<ControlService>().endOfStream();
+      pc.services().get<ControlService>().readyToQuit(QuitRequest::Me);
+    }
   }
 }