RootParameterTracker.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
#pragma once
 
// in fw:
#include <tracking/trackFindingVXD/analyzingTools/AnalizerTCInfo.h>
#include <tracking/trackFindingVXD/analyzingTools/AnalyzingAlgorithmFactory.h>
#include <tracking/trackFindingVXD/analyzingTools/KeyValBox.h>
// #include <tracking/trackFindingVXD/analyzingTools/algorithms/AnalyzingAlgorithmClusterBased.h> // TODO
#include <framework/logging/Logger.h>
 
// root:
#include <TTree.h>
#include <TFile.h>
#include <Math/Vector3D.h>
 
// stl:
#include <string>
#include <vector>
#include <utility> // for pair
 
namespace Belle2 {
  class RootParameterTracker {
  protected:
 
    template<class ValueType> using StringKeyBox = KeyValBox<std::string, ValueType>;
 
 
    template<class DataType> using AnalyzingAlgorithm = AnalyzingAlgorithmBase<DataType, AnalizerTCInfo, ROOT::Math::XYZVector>;
 
 
    template<class DataType> using AlgoDataBox =
      StringKeyBox<StringKeyBox<std::pair<AnalyzingAlgorithm<DataType>*, std::vector<DataType>*>>>;
 
 
    StringKeyBox<TTree*> m_treeBox;
 
 
    AlgoDataBox<double> m_algoDataDouble;
 
 
    AlgoDataBox<int> m_algoDataInt;
 
 
    AlgoDataBox<std::vector<double>> m_algoDataVecDouble;
 
 
    TFile* m_file;
 
 
    TTree* prepareTTree(std::string tcTypeName)
    {
      if (m_file == nullptr) {
        B2FATAL("RootParameterTracker::prepareTTree(), Root file not initialized yet! Please call RootParameterTracker::initialize(...) first!");
      }
      m_file->cd();
      //   m_file->ls();
 
      TTree** tree4tcType = m_treeBox.find(tcTypeName);
      if (tree4tcType == nullptr) {
        m_treeBox.push_back({
          tcTypeName,
          new TTree(tcTypeName.c_str(), (std::string("collects data collected to tcType ") + tcTypeName).c_str())
        }
                           );
        tree4tcType = m_treeBox.find(tcTypeName);
        B2WARNING("RootParameterTracker::prepareTTree: new tree for tcType " << tcTypeName << " created, m_treeBox has now " <<
                  m_treeBox.size() << " entries");
        return *tree4tcType;
      }
 
      B2WARNING("RootParameterTracker::prepareTTree: ttree for tcType " << tcTypeName << " is reused, m_treeBox has " << m_treeBox.size()
                << " entries");
      return *tree4tcType;
    }
 
 
    void collectData4DoubleAlgorithms(std::string tcTypeName, const AnalizerTCInfo& aTC);
 
 
    void collectData4IntAlgorithms(std::string tcTypeName, const AnalizerTCInfo& aTC);
 
 
    void collectData4VecDoubleAlgorithms(std::string tcTypeName, const AnalizerTCInfo& aTC);
 
 
  public:
 
    RootParameterTracker() : m_file(nullptr) {}
 
 
    void initialize(std::string fileName, std::string fileTreatment)
    {
      if (fileTreatment != std::string("RECREATE") and fileTreatment != std::string("UPDATE")) {
        B2FATAL("RootParameterTracker::initialize(), specified fileTreatment is " << fileTreatment <<
                ", which is invalid, please read the documentation!");
      }
 
      if (m_file != nullptr) {
        B2FATAL("RootParameterTracker::initialize(), there was a file already linked to this ParameterTracker, which is invalid, please read the documentation!");
      }
      B2DEBUG(20, "RootParameterTracker::initialize(), given parameters are fileName/fileTreatment: " << fileName << "/" <<
              fileTreatment);
      m_file = new TFile(fileName.c_str(), fileTreatment.c_str()); // alternative: UPDATE
      m_file->ls();
    }
 
 
    void addParameters4DoubleAlgorithms(std::string tcTypeName, std::string algorithmName);
 
 
    void addParameters4IntAlgorithms(std::string tcTypeName, std::string algorithmName);
 
 
    void addParameters4VecDoubleAlgorithms(std::string tcTypeName, std::string algorithmName);
 
 
    void collectData(const std::vector<AnalizerTCInfo>& tcVector)
    {
      B2DEBUG(21, "RootParameterTracker::collectData(), size of given tcVector is: " << tcVector.size());
 
      for (const AnalizerTCInfo& tc : tcVector) {
        std::string tcTypeName = TCType::getTypeName(tc.getType());
        B2DEBUG(22, "RootParameterTracker::collectData(), executing TC of type: " << tcTypeName);
 
        collectData4DoubleAlgorithms(tcTypeName, tc);
        collectData4IntAlgorithms(tcTypeName, tc);
        collectData4VecDoubleAlgorithms(tcTypeName, tc);
      } // looping over TCs
    }
 
 
    void fillRoot()
    {
      B2DEBUG(21, "RootParameterTracker::fillRoot(), Executing " << m_treeBox.size() << " ttrees.");
 
      m_file->cd(); //important! without this the famework root I/O (SimpleOutput etc) could mix with the root I/O of this module
//       m_file->ls();
 
      for (auto& boxEntry : m_treeBox) {
        int nBytesWritten = boxEntry.second->Fill();
        B2DEBUG(20, "RootParameterTracker::fillRoot() ttree " << boxEntry.first << " got " << nBytesWritten << " Bytes written");
        // boxEntry.second->Print();
      }
 
      for (auto& algoData2tcType : m_algoDataDouble) {
        for (auto& algoPack : algoData2tcType.second) {
          algoPack.second.second->clear();
        }
      }
 
      for (auto& algoData2tcType : m_algoDataInt) {
        for (auto& algoPack : algoData2tcType.second) {
          algoPack.second.second->clear();
        }
      }
 
      for (auto& algoData2tcType : m_algoDataVecDouble) {
        for (auto& algoPack : algoData2tcType.second) {
          algoPack.second.second->clear();
        }
      }
    }
 
 
    void terminate()
    {
      B2DEBUG(20, "RootParameterTracker::terminate(), Writing results to root-file and clean up heap.");
      if (m_file == nullptr) {
        B2WARNING("RootParameterTracker::terminate(): no rootFile found! skipping writing data into root file!");
        return;
      }
 
      m_file->cd(); //important! without this the famework root I/O (SimpleOutput etc) could mix with the root I/O of this module
      m_file->ls();
 
      for (auto& boxEntry : m_treeBox) {
        boxEntry.second->Write();
 
        if (LogSystem::Instance().isLevelEnabled(LogConfig::c_Debug, 5, PACKAGENAME()) == true) {
          B2DEBUG(20, "RootParameterTracker::terminate(), TTree " << boxEntry.first << " was written:");
          // boxEntry.second->Print();
        }
      }
      m_file->Close();
 
      for (auto& algoData2tcType : m_algoDataDouble) {
        for (auto& algoPack : algoData2tcType.second) {
          delete algoPack.second.first; // algorithm
          delete algoPack.second.second; // vector(data) to algorithm
        }
      }
 
      for (auto& algoData2tcType : m_algoDataInt) {
        for (auto& algoPack : algoData2tcType.second) {
          delete algoPack.second.first; // algorithm
          delete algoPack.second.second; // vector(data) to algorithm
        }
      }
 
      for (auto& algoData2tcType : m_algoDataVecDouble) {
        for (auto& algoPack : algoData2tcType.second) {
          delete algoPack.second.first; // algorithm
          delete algoPack.second.second; // vector(data) to algorithm
        }
      }
    }
 
  };
}