NAIAChain.cpp

Go to the documentation of this file.

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

#include "TTreeIndex.h"

#include <functional>
#include <unordered_map>

namespace NAIA {

NAIAChain::NAIAChain(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("NAIAChain::GetEntries - Unrecognized access mode");
  }
}

Event &NAIAChain::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 &NAIAChain::GetEventWithIndex(unsigned int run, unsigned int eventno) {
  return GetEvent(m_data.chain->GetEntryNumberWithIndex(run, eventno));
}

const RTIInfo &NAIAChain::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;
    auto rc = m_rti.chain->GetEntryWithIndex(m_event.header->UTCTime);
    if (rc < 0) {
      throw std::runtime_error(
          fmt::format("Failed to read entry for UTC second {} from RTI tree. (ec: {})", m_event.header->UTCTime, rc));
    }
  }

  return m_rtiInfo;
}

const FileInfo &NAIAChain::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 &NAIAChain::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 *NAIAChain::GetRTITree() {
  static const std::string &fnName = "NAIAChain::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("NAIAChain::GetRTITree - Unrecognized access mode");
  }
}

TChain *NAIAChain::GetFileInfoTree() {
  static const std::string &fnName = "NAIAChain::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("NAIAChain::GetFileInfoTree - Unrecognized access mode");
  }
}

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

int NAIAChain::Add(const std::string &filePath) {
  switch (m_accessMode) {
  case AccessMode::Read:
    if (!CheckVersion(filePath)) {
      return -1;
    }

    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("NAIAChain::Add - Unrecognized access mode");
  }
}

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

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

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

unsigned long int NAIAChain::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("NAIAChain::GetEntries - Unrecognized access mode");
  }
}

void NAIAChain::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("NAIAChain::Print - Unrecognized access mode");
  }
}

int NAIAChain::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();

    if (!m_isMC) {
      spdlog::debug("Building RTIInfo index");
      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("NAIAChain::Write - Unrecognized access mode");
  }
}

void NAIAChain::SetupBranches(bool create_MC_branches) {
  switch (m_accessMode) {
  case AccessMode::Read:
    if (m_data.chain->GetBranch("MCTruthBaseData")) {
      m_isMC = true;
      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 (create_MC_branches)
      m_file.tree->Branch("MCFileInfo", &m_fileInfoMC);
    m_rti.tree->Branch("RTIInfo", &m_rtiInfo);
    break;
  default:
    throw std::runtime_error("NAIAChain::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(); }},
    {"TrTrackBase", [](Event &event) { event.trTrackBase.DisableIO(); }},
    {"TrTrackPlus", [](Event &event) { event.trTrackPlus.DisableIO(); }},
    {"SecondTrTrackBase", [](Event &event) { event.secondTrTrackBase.DisableIO(); }},
    {"TrTrackBaseStandalone", [](Event &event) { event.trTrackBaseSt.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<NAIAChain> NAIAChain::CreateSkimTree(const std::string &filename, const std::string &exclBranches) {
  static const std::string &fnName = "NAIAChain::CreateSkimTree";
  auto logger = getLogger(fnName);

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

  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);

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

void NAIAChain::SetEntryList(TEntryList *entryList, Option_t *option) {
  static const std::string &fnName = "NAIAChain::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;
  }
}

bool NAIAChain::CheckVersion(const std::string &filePath) {
  static const std::string &fnName = "NAIAChain::CheckVersion";
  auto logger = getLogger(fnName);

  NAIA::Version this_version{};

  auto file = std::unique_ptr<TFile>(TFile::Open(filePath.c_str()));
  NAIA::Version *file_version = file->Get<NAIA::Version>("VersionHeader");
  if (!file_version) {
    logger->error("NAIA version header not found in file {}. This file will be skipped.", filePath);
    return false;
  }

  if (file_version->major != this_version.major) {
    logger->error("Major version mismatch in file {} ({} on file vs {} in this library). This file will be skipped.",
                  filePath, file_version->major, this_version.major);
    return false;
  }

  if (file_version->minor != this_version.minor) {
    logger->warn("Minor version mismatch in file {} ({} on file vs {} in this library). This should be fine, however "
                 "please report any issue to the NAIA development team.",
                 filePath, file_version->minor, this_version.minor);
  }

  return true;
}

} // namespace NAIA