SingleTreeChain.cpp

Go to the documentation of this file.

#include "Chain/SingleTreeChain.h"
#include "Utility/Logging/Logging.h"

#include "TTreeIndex.h"

#include <functional>
#include <unordered_map>

namespace NAIA {

SingleTreeChain::SingleTreeChain(AccessMode mode) : m_accessMode{mode} {
  switch (m_accessMode) {
  case AccessMode::Read:
    m_data.chain = new TChain("NAIAChain");
    m_rti.chain = new TChain("RTI");
    m_file.chain = new TChain("FileInfo");
    break;
  case AccessMode::Write:
    m_data.tree = new TTree("NAIAChain", "AMS-Italy ntuple tree");
    m_rti.tree = new TTree("RTI", "RTIInfo tree");
    m_file.tree = new TTree("FileInfo", "FileInfo tree");
    break;
  default:
    throw std::runtime_error("SingleTreeChain::GetEntries - Unrecognized access mode");
  }
}

Event &SingleTreeChain::GetEvent(unsigned long long iEv) {
  static int currentTree = -1;

  // when the tree changes we need to reload the branches...
  auto treeReadEntry = m_data.chain->LoadTree(iEv);
  if (m_data.chain->GetTreeNumber() != currentTree) {
    spdlog::trace("Event {} -> CHANGING TREE", iEv);
    m_event.SetAllBranchAddress(this->m_data.chain);
    currentTree = m_data.chain->GetTreeNumber();
  }

  // do stuff to propagate iEv to the Event object (containers need to know which event to load)
  m_event.SetEventNumber(treeReadEntry);
  return m_event;
}

Event &SingleTreeChain::GetEventWithIndex(unsigned int run, unsigned int eventno) {
  return GetEvent(m_data.chain->GetEntryNumberWithIndex(run, eventno));
}

const RTIInfo &SingleTreeChain::GetEventRTIInfo() {
  static unsigned int lastSec = 0;

  if (IsMC())
    throw(std::runtime_error("GetEventRTIInfo called on a MC chain"));

  if (m_event.header->UTCTime != lastSec) {
    lastSec = m_event.header->UTCTime;
    m_rti.chain->GetEntryWithIndex(m_event.header->UTCTime);
  }

  return m_rtiInfo;
}

const FileInfo &SingleTreeChain::GetEventFileInfo() {
  static unsigned int lastRun = 0;

  if (m_event.header->Run != lastRun) {
    lastRun = m_event.header->Run;
    m_file.chain->GetEntryWithIndex(m_event.header->Run);
  }

  return m_fileInfo;
}

const MCFileInfo &SingleTreeChain::GetEventMCFileInfo() {
  static unsigned int lastRun = 0;

  if (!IsMC())
    throw(std::runtime_error("GetEventMCFileInfo called on a non-MC chain"));

  if (m_event.header->Run != lastRun) {
    lastRun = m_event.header->Run;
    m_file.chain->GetEntryWithIndex(m_event.header->Run);
  }

  return m_fileInfoMC;
}

TChain *SingleTreeChain::GetRTITree() {
  static const std::string &fnName = "SingleTreeChain::GetRTITree";
  auto logger = getLogger(fnName);

  switch (m_accessMode) {
  case AccessMode::Read:
    return m_rti.chain;
  case AccessMode::Write:
    logger->error("Not supported in Write mode");
    return nullptr;
  default:
    throw std::runtime_error("SingleTreeChain::GetRTITree - Unrecognized access mode");
  }
}

TChain *SingleTreeChain::GetFileInfoTree() {
  static const std::string &fnName = "SingleTreeChain::GetFileInfoTree";
  auto logger = getLogger(fnName);

  switch (m_accessMode) {
  case AccessMode::Read:
    return m_file.chain;
  case AccessMode::Write:
    logger->error("Not supported in Write mode");
    return nullptr;
  default:
    throw std::runtime_error("SingleTreeChain::GetFileInfoTree - Unrecognized access mode");
  }
}

void SingleTreeChain::SetDirectory(TDirectory *directory) {
  m_data.tree->SetDirectory(directory);
  m_rti.tree->SetDirectory(directory);
  m_file.tree->SetDirectory(directory);
}

int SingleTreeChain::Add(const std::string &filePath) {
  switch (m_accessMode) {
  case AccessMode::Read:
    m_file.chain->Add(filePath.c_str());
    m_rti.chain->Add(filePath.c_str());

    return m_data.chain->Add(filePath.c_str());
  case AccessMode::Write:
    return -1;
  default:
    throw std::runtime_error("SingleTreeChain::Add - Unrecognized access mode");
  }
}

int SingleTreeChain::Fill() {
  switch (m_accessMode) {
  case AccessMode::Read:
    return -1;
  case AccessMode::Write:
    return m_data.tree->Fill();
  default:
    throw std::runtime_error("SingleTreeChain::Fill - Unrecognized access mode");
  }
}

int SingleTreeChain::FillRTI() {
  switch (m_accessMode) {
  case AccessMode::Read:
    return -1;
  case AccessMode::Write:
    return m_rti.tree->Fill();
  default:
    throw std::runtime_error("SingleTreeChain::FillRTI - Unrecognized access mode");
  }
}

int SingleTreeChain::FillFileInfo() {
  switch (m_accessMode) {
  case AccessMode::Read:
    return -1;
  case AccessMode::Write:
    return m_file.tree->Fill();
  default:
    throw std::runtime_error("SingleTreeChain::FillFileInfo - Unrecognized access mode");
  }
}

unsigned long int SingleTreeChain::GetEntries() {
  switch (m_accessMode) {
  case AccessMode::Read:
    return m_data.chain->GetEntries();
  case AccessMode::Write:
    return m_data.tree->GetEntries();
  default:
    throw std::runtime_error("SingleTreeChain::GetEntries - Unrecognized access mode");
  }
}

void SingleTreeChain::Print() {
  switch (m_accessMode) {
  case AccessMode::Read:
    m_file.chain->Print();
    m_rti.chain->Print();
    return m_data.chain->Print();
  case AccessMode::Write:
    m_file.tree->Print();
    m_rti.tree->Print();
    return m_data.tree->Print();
  default:
    throw std::runtime_error("SingleTreeChain::Print - Unrecognized access mode");
  }
}

int SingleTreeChain::Write() {
  switch (m_accessMode) {
  case AccessMode::Read:
    m_file.chain->Write();
    m_rti.chain->Write();
    return m_data.chain->Write();
  case AccessMode::Write:
    m_file.tree->BuildIndex("FileInfo.Run");
    m_file.tree->Write();

    m_rti.tree->BuildIndex("RTIInfo.UTCTime");
    m_rti.tree->Write();

    m_data.tree->BuildIndex("HeaderData.Run", "HeaderData.EventNo");
    return m_data.tree->Write();
  default:
    throw std::runtime_error("SingleTreeChain::Write - Unrecognized access mode");
  }
}

void SingleTreeChain::SetupBranches(bool isMC) {
  switch (m_accessMode) {
  case AccessMode::Read:
    if (m_data.chain->GetBranch("MCTruthBaseData")) {
      m_event.SetMC(true);
    }

    m_event.SetAllBranchAddress(this->m_data.chain);
    m_file.chain->SetBranchAddress("FileInfo", &m_fileInfoPtr);
    if (m_file.chain->GetBranch("MCFileInfo"))
      m_file.chain->SetBranchAddress("MCFileInfo", &m_fileInfoMCPtr);
    m_rti.chain->SetBranchAddress("RTIInfo", &m_rtiInfoPtr);

    m_file.chain->BuildIndex("FileInfo.Run");
    m_rti.chain->BuildIndex("RTIInfo.UTCTime");
    break;
  case AccessMode::Write:
    m_event.BranchAll(this->m_data.tree);
    m_file.tree->Branch("FileInfo", &m_fileInfo);
    if (isMC)
      m_file.tree->Branch("MCFileInfo", &m_fileInfoMC);
    m_rti.tree->Branch("RTIInfo", &m_rtiInfo);
    break;
  default:
    throw std::runtime_error("SingleTreeChain::SetupBranches - Unrecognized access mode");
  }
}

static const std::unordered_map<std::string, std::function<void(Event &)>> evDisablers = {
    {"EventSummary", [](Event &event) { event.evSummary.DisableIO(); }},
    {"DAQ", [](Event &event) { event.daq.DisableIO(); }},
    {"TofBase", [](Event &event) { event.tofBase.DisableIO(); }},
    {"TofPlus", [](Event &event) { event.tofPlus.DisableIO(); }},
    {"TofBaseStandalone", [](Event &event) { event.tofBaseSt.DisableIO(); }},
    {"TofPlusStandalone", [](Event &event) { event.tofPlusSt.DisableIO(); }},
    {"EcalBase", [](Event &event) { event.ecalBase.DisableIO(); }},
    {"EcalPlus", [](Event &event) { event.ecalPlus.DisableIO(); }},
    {"TtrTrackBase", [](Event &event) { event.trTrackBase.DisableIO(); }},
    {"TtrTrackPlus", [](Event &event) { event.trTrackPlus.DisableIO(); }},
    {"TrdKBase", [](Event &event) { event.trdKBase.DisableIO(); }},
    {"TrdKBaseStandalone", [](Event &event) { event.trdKBaseSt.DisableIO(); }},
    {"RichBase", [](Event &event) { event.richBase.DisableIO(); }},
    {"RichPlus", [](Event &event) { event.richPlus.DisableIO(); }},
    {"UnbExtHitBase", [](Event &event) { event.extHitBase.DisableIO(); }},
    {"MCTruthBase", [](Event &event) { event.mcTruthBase.DisableIO(); }},
    {"MCTruthPlus", [](Event &event) { event.mcTruthPlus.DisableIO(); }},
};

SkimTreeHandle<SingleTreeChain> SingleTreeChain::CreateSkimTree(const std::string &filename,
                                                                const std::string &exclBranches) {
  static const std::string &fnName = "SingleTreeChain::CreateSkimTree";
  auto logger = getLogger(fnName);

  auto outFile = std::make_unique<TFile>(filename.c_str(), "recreate");
  auto newChain = std::make_unique<SingleTreeChain>(SingleTreeChain::AccessMode::Write);

  if (!exclBranches.empty()) {
    for (const std::string &container : TokenizeString(exclBranches, ';')) {
      if (evDisablers.find(container) != std::end(evDisablers)) {
        logger->info("Disabling container {}", container);
        evDisablers.at(container)(newChain->m_event);
      } else {
        logger->error("Unknown container {}", container);
      }
    }
  }

  // link existing branches to new tree branches
  newChain->m_event.MirrorBranches(newChain->m_data.tree, m_event);
  newChain->m_rti.tree->Branch("RTIInfo", &m_rtiInfo);
  newChain->m_file.tree->Branch("FileInfo", &m_fileInfo);
  if (m_file.tree->GetBranch("MCFileInfo"))
    newChain->m_file.tree->Branch("MCFileInfo", &m_fileInfoMC);

  return {std::move(outFile), std::move(newChain), this};
}

void SingleTreeChain::SetEntryList(TEntryList *entryList, Option_t *option) {
  static const std::string &fnName = "SingleTreeChain::SetEntryList";
  auto logger = getLogger(fnName);

  switch (m_accessMode) {
  case AccessMode::Write:
    logger->error("TEntryList not supported in Write mode");
    break;
  case AccessMode::Read:
    m_data.chain->SetEntryList(entryList, option);
    break;
  default:
    break;
  }
}

} // namespace NAIA