Introduce helper function to load MClabels from a TTree/TFile

sawenzel · sawenzel · commit af12c27889a8 · 2023-02-06T12:36:29.000+01:00
Function returns a constant MC label container. The persistent
format can be either a MCTruthContainer or IOMCTruthContainer with special
split format.
diff --git a/DataFormats/simulation/CMakeLists.txt b/DataFormats/simulation/CMakeLists.txt
@@ -20,6 +20,7 @@ o2_add_library(SimulationDataFormat
                        src/MCUtils.cxx
                        src/O2DatabasePDG.cxx
                        src/InteractionSampler.cxx
+                       src/ConstMCTruthContainer.cxx
                PUBLIC_LINK_LIBRARIES Microsoft.GSL::GSL
                                      FairRoot::Base
                                      O2::DetectorsCommonDataFormats
diff --git a/DataFormats/simulation/include/SimulationDataFormat/ConstMCTruthContainer.h b/DataFormats/simulation/include/SimulationDataFormat/ConstMCTruthContainer.h
@@ -21,6 +21,8 @@
 #include <Framework/Traits.h>
 #endif
 
+class TTree;
+
 namespace o2
 {
 namespace dataformats
@@ -211,6 +213,14 @@ class ConstMCTruthContainerView
 using ConstMCLabelContainer = o2::dataformats::ConstMCTruthContainer<o2::MCCompLabel>;
 using ConstMCLabelContainerView = o2::dataformats::ConstMCTruthContainerView<o2::MCCompLabel>;
 
+class MCLabelIOHelper
+{
+ public:
+  /// Convenience function to loads MC labels for some entry from a TTree and TBranch.
+  /// Labels can be stored as either MCTruthContainer or IOMCTruthContainer. The caller takes ownership of the returned pointer.
+  static ConstMCTruthContainer<o2::MCCompLabel>* loadFromTTree(TTree* tree, std::string const& brname, int entry);
+};
+
 } // namespace dataformats
 } // namespace o2
 
diff --git a/DataFormats/simulation/src/ConstMCTruthContainer.cxx b/DataFormats/simulation/src/ConstMCTruthContainer.cxx
@@ -0,0 +1,59 @@
+// 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.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// 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.
+
+#include <SimulationDataFormat/ConstMCTruthContainer.h>
+#include <SimulationDataFormat/MCTruthContainer.h>
+#include <SimulationDataFormat/IOMCTruthContainerView.h>
+#include <SimulationDataFormat/MCCompLabel.h>
+#include <TFile.h>
+#include <TTree.h>
+#include <iostream>
+
+using namespace o2::dataformats;
+
+ConstMCTruthContainer<o2::MCCompLabel>* MCLabelIOHelper::loadFromTTree(TTree* tree, std::string const& brname, int entry)
+{
+  o2::dataformats::MCTruthContainer<o2::MCCompLabel>* labels = nullptr;
+  o2::dataformats::IOMCTruthContainerView* labelROOTbuffer = nullptr;
+  o2::dataformats::ConstMCTruthContainer<o2::MCCompLabel>* constlabels = new ConstMCTruthContainer<o2::MCCompLabel>();
+
+  auto branch = tree->GetBranch(brname.c_str());
+  if (!branch) {
+    return nullptr;
+  }
+  auto labelClass = branch->GetClassName();
+  bool oldlabelformat = false;
+
+  // we check how the labels are persisted and react accordingly
+  if (TString(labelClass).Contains("IOMCTruthContainer")) {
+    branch->SetAddress(&labelROOTbuffer);
+    branch->GetEntry(entry);
+    if (labelROOTbuffer) {
+      labelROOTbuffer->copyandflatten(*constlabels);
+      delete labelROOTbuffer;
+    }
+  } else {
+    // case when we directly streamed MCTruthContainer
+    // TODO: support more cases such as ConstMCTruthContainer etc
+    std::string serializedtype("o2::dataformats::MCTruthContainer<o2::MCCompLabel>");
+    if (!TString(labelClass).EqualTo(serializedtype.c_str())) {
+      std::cerr << "Error: expected serialized type " << serializedtype << " but found " << labelClass;
+      return nullptr;
+    }
+    branch->SetAddress(&labels);
+    branch->GetEntry(entry);
+    if (labels) {
+      labels->flatten_to(*constlabels);
+      delete labels;
+    }
+  }
+  return constlabels;
+}
diff --git a/DataFormats/simulation/test/testMCTruthContainer.cxx b/DataFormats/simulation/test/testMCTruthContainer.cxx
@@ -355,6 +355,7 @@ BOOST_AUTO_TEST_CASE(MCTruthContainer_ROOTIO)
     TFile f("tmp2.root", "RECREATE");
     TTree tree("o2sim", "o2sim");
     auto br = tree.Branch("Labels", &io, 32000, 2);
+    tree.Branch("LabelsOriginal", &container, 32000, 2);
     tree.Fill();
     tree.Write();
     f.Close();
@@ -383,6 +384,21 @@ BOOST_AUTO_TEST_CASE(MCTruthContainer_ROOTIO)
   BOOST_CHECK(cc.getLabels(BIGSIZE - 1).size() == 2);
   BOOST_CHECK(cc.getLabels(BIGSIZE - 1)[0] == TruthElement(BIGSIZE - 1, BIGSIZE - 1, BIGSIZE - 1));
   BOOST_CHECK(cc.getLabels(BIGSIZE - 1)[1] == TruthElement(BIGSIZE, BIGSIZE - 1, BIGSIZE - 1));
+
+  // testing convenience API to retrieve a constant label container from a ROOT file, entry 0
+  auto cont = o2::dataformats::MCLabelIOHelper::loadFromTTree(tree2, "Labels", 0);
+  auto cont2 = o2::dataformats::MCLabelIOHelper::loadFromTTree(tree2, "LabelsOriginal", 0);
+
+  BOOST_CHECK(cont);
+  BOOST_CHECK(cont2);
+  BOOST_CHECK(cont->getNElements() == (BIGSIZE - 1) * 2);
+  BOOST_CHECK(cont2->getNElements() == (BIGSIZE - 1) * 2);
+  BOOST_CHECK(cont->getLabels(0).size() == 0);
+  BOOST_CHECK(cont2->getLabels(0).size() == 0);
+  BOOST_CHECK(cont->getLabels(BIGSIZE - 1)[0] == TruthElement(BIGSIZE - 1, BIGSIZE - 1, BIGSIZE - 1));
+  BOOST_CHECK(cont->getLabels(BIGSIZE - 1)[1] == TruthElement(BIGSIZE, BIGSIZE - 1, BIGSIZE - 1));
+  BOOST_CHECK(cont2->getLabels(BIGSIZE - 1)[0] == TruthElement(BIGSIZE - 1, BIGSIZE - 1, BIGSIZE - 1));
+  BOOST_CHECK(cont2->getLabels(BIGSIZE - 1)[1] == TruthElement(BIGSIZE, BIGSIZE - 1, BIGSIZE - 1));
 }
 
 } // namespace o2
