GGSTRootReader.cpp

Go to the documentation of this file.

/*
 * GGSTRootReader.cpp
 *
 *  Created on: 23 Aug 2011
 *      Author: Nicola Mori
 */

// GGS headers
#include "montecarlo/readers/GGSTRootReader.h"
#include "utils/GGSSmartLog.h"

// Root headers
#include "TString.h"

// C++ headers
#include <algorithm>
#include <filesystem>
#include <fstream>
#include <regex>

GGSTRootReader::GGSTRootReader() : _currentEv(-1) {}

GGSTRootReader::~GGSTRootReader() {
  for (unsigned int iHandler = 0; iHandler < _filesHandlers.size(); iHandler++) {
    delete _filesHandlers[iHandler];
  }
  for (unsigned int iReader = 0; iReader < _readers.size(); iReader++) {
    delete _readers[iReader];
  }
}

GGSTFilesHandler *GGSTRootReader::Open(const std::string &fileName, GGSTFilesHandler *filesHandler, bool force) {
  static const std::string routineName("GGSTRootReader::Open");

  GGSTFilesHandler *handler = nullptr;
  if (filesHandler == nullptr)
    handler = new GGSTFilesHandler;
  else {
    handler = filesHandler;
  }

  if (fileName.find('?') != std::string::npos || fileName.find('*') != std::string::npos) {
    // ******* Name with wildcards *******

    // ++++ Convert fileName into a regular expression ++++
    std::string regExpFiles(fileName);
    // 1. escape all .
    std::string::size_type pos = 0;
    while ((pos = regExpFiles.find('.', pos)) != std::string::npos) {
      regExpFiles.replace(pos, 1, "\\.");
      pos += 2;
    }
    // 2. convert all * to .*
    pos = 0;
    while ((pos = regExpFiles.find('*', pos)) != std::string::npos) {
      regExpFiles.replace(pos, 1, ".*");
      pos += 2;
    }
    // 3. convert all ? to .
    std::replace(regExpFiles.begin(), regExpFiles.end(), '?', '.');

    // ++++ Search files ++++
    std::string basePathName(fileName);
    // Extract path
    size_t posOfLastSlash = basePathName.find_last_of('/');
    if (posOfLastSlash != std::string::npos)
      basePathName.erase(posOfLastSlash);
    else
      basePathName = std::filesystem::current_path().string();

    std::filesystem::path basePath(std::filesystem::absolute(basePathName));

    std::filesystem::directory_iterator endItr; // Default constructor yields past-the-end
    std::filesystem::directory_iterator dirIter(basePath);
    if (dirIter == endItr) {
      if (filesHandler == nullptr) {
        delete handler;
      }
      return nullptr;
    }

    for (; dirIter != endItr; ++dirIter) {
      // Skip if not a file
      if (!std::filesystem::is_regular_file(dirIter->status()))
        continue;
      std::smatch what;
      // Skip if no match
      auto pathStr = dirIter->path().string();
      if (!std::regex_match(pathStr, what, std::regex(regExpFiles)))
        continue;
      // Skip if not a Root file
      if (pathStr.substr(pathStr.size() - 5) != ".root")
        continue;
      // File matches, store it
      handler->_AddFile(pathStr, force);
    }
    if (handler->GetNFiles() == 0) {
      GGSCOUT(WARNING) << fileName << ": no file found." << GGSENDL;
      delete handler;
      return nullptr;
    }
    handler->_SortFiles();

  } else if (fileName.substr(fileName.size() - 4) == ".txt") {
    // *******  Text file with file names *******
    if (std::filesystem::exists(fileName)) {
      std::ifstream file(fileName.data(), std::ios::in);
      std::string currFile;
      while (file.good()) {
        getline(file, currFile);
        if (std::filesystem::exists(currFile) && currFile.substr(currFile.size() - 5) == ".root")
          handler->_AddFile(currFile, force);
      }
    } else {
      if (filesHandler == nullptr) {
        delete handler;
      }
      return nullptr;
    }
  } else if (fileName.substr(fileName.size() - 5) == ".root") {
    // *******  Single root file *******
    if (!handler->_AddFile(fileName, force)) {
      if (filesHandler == nullptr) {
        delete handler;
      }
      return nullptr;
    }
  } else {
    if (filesHandler == nullptr) {
      delete handler;
    }
    return nullptr;
  }

  if (find(_filesHandlers.begin(), _filesHandlers.end(), handler) == _filesHandlers.end())
    _filesHandlers.push_back(handler);
  return handler;
}

TChain *GGSTRootReader::_FindChain(const TString &treeName, GGSTFilesHandler *filesHandler) {

  // Search in files handler (chain already opened by another reader)
  for (unsigned int iChain = 0; iChain < filesHandler->_chains.size(); iChain++) {
    if (filesHandler->_chains[iChain]->GetName() == treeName) {
      return filesHandler->_chains[iChain];
    }
  }

  // Search in files handled by the handler
  return _CreateChain(treeName, filesHandler);
}

TChain *GGSTRootReader::_CreateChain(const TString &treeName, GGSTFilesHandler *filesHandler) {

  static const std::string routineName("GGSTRootReader::_CreateChain");
  std::vector<TChain *> chains(1);
  chains[0] = new TChain(treeName);
  unsigned int iFile = 0;

  // Find the first file containing the desired tree
  while (iFile < filesHandler->_files.size()) {
    TFile *file = TFile::Open(filesHandler->_files[iFile].data());
    TTree *tree = dynamic_cast<TTree *>(file->Get(treeName));
    if (tree) {
      chains[0]->Add(filesHandler->_files[iFile].data(), 0);
      chains[0]->SetTitle(tree->GetTitle());
      delete file;
      break;
    }
    iFile++;
    delete file;
  }

  if (iFile == filesHandler->_files.size()) {
    // Tree not found in any file
    delete chains[0];
    return nullptr;
  }

  // First tree found.
  for (iFile += 1; iFile < filesHandler->_files.size(); iFile++) {
    TFile *file = TFile::Open(filesHandler->_files[iFile].data());
    if (!(file->IsZombie())) {
      TTree *tree = dynamic_cast<TTree *>(file->Get(treeName));
      if (tree) {
        // Check if branches match with at least one chain
        bool chainNotFound = true;
        for (unsigned int iChain = 0; iChain < chains.size(); iChain++) {
          if (_CheckBranches(chains[iChain], tree)) {
            // Number of branches and branch names match. Add this file to the current chain.
            chains[iChain]->Add(filesHandler->_files[iFile].data(), 0);
            chainNotFound = false;
            delete file;
            file = nullptr;
            break;
          }
        }
        if (chainNotFound) {
          chains.push_back(new TChain(treeName));
          chains.back()->SetTitle(tree->GetTitle());
          chains.back()->Add(filesHandler->_files[iFile].data(), 0);
        }
      }
    }
    delete file;
  }

  // Set friends if more than one chain has been created
  for (unsigned int iChain = 1; iChain < chains.size(); iChain++) {
    if (chains[iChain]->GetEntries() != chains[0]->GetEntries()) {
      GGSCOUT(ERROR) << "Mismatch in number of elements:" << GGSENDL;
      GGSCCOUT(ERROR) << "  Chain 0: " << chains[0]->GetEntries() << " entries" << GGSENDL;
      GGSCCOUT(ERROR) << "    Title: \"" << chains[0]->GetTitle() << "\"" << GGSENDL;
      GGSCCOUT(ERROR) << "  Chain " << iChain << ": " << chains[iChain]->GetEntries() << " entries" << GGSENDL;
      GGSCCOUT(ERROR) << "    Title \": " << chains[iChain]->GetTitle() << "\"" << GGSENDL;
      for (unsigned int iChain1 = 0; iChain1 < chains.size(); iChain1++) {
        delete chains[iChain1];
        return nullptr;
      }
    } else {
      chains[0]->AddFriend(chains[iChain]);
    }
  }

  // Pass the ownership of all the chains to the file handler
  for (unsigned int iChain = 0; iChain < chains.size(); iChain++) {
    filesHandler->_chains.push_back(chains[iChain]);
  }

  return chains[0];
}

bool GGSTRootReader::_CheckBranches(TTree *tree1, TTree *tree2) {

  TObjArray *branchesInTree1 = tree1->GetListOfBranches();
  // Retrieve the list of branches in the current tree
  TObjArray *branchesInTree2 = tree2->GetListOfBranches();
  // Check if the two lists have the same number of branches with the same names
  if (branchesInTree1->GetEntries() == branchesInTree2->GetEntries()) {
    for (int iBranch = 0; iBranch < branchesInTree1->GetEntries(); iBranch++) {
      if (strcmp(((TBranch *)(branchesInTree1->At(iBranch)))->GetName(),
                 ((TBranch *)(branchesInTree2->At(iBranch)))->GetName())) {
        return false;
      }
    }
  } else
    return false;

  return true;
}

Long64_t GGSTRootReader::_GetEntriesBlind() {

  if (_filesHandlers.size() == 0)
    return 0;

  if (_filesHandlers[0]->_files.size() == 0)
    return 0;

  TFile *file = TFile::Open(_filesHandlers[0]->_files[0].data());
  TString treeName("");
  TList *keysList = file->GetListOfKeys();
  for (int iKey = 0; iKey < keysList->GetEntries(); iKey++) {
    TString keyName(keysList->At(iKey)->GetName());
    TTree *key = dynamic_cast<TTree *>(file->Get(keyName));
    if (key != nullptr) {
      treeName = ((TTree *)key)->GetName();
    }
  }

  delete file;
  file = nullptr;
  if (treeName == "")
    return 0;

  TChain *chain = _CreateChain(treeName, _filesHandlers[0]);
  Long64_t nEntries = 0;
  if (chain)
    nEntries = chain->GetEntries();
  delete chain;

  return nEntries;
}

void GGSTRootReader::GetEntry(Long64_t entry) {

  for (unsigned int i = 0; i < _readers.size(); i++) {
    _readers[i]->reader->GetEntry(entry);
  }
  _currentEv = entry;
}

Long64_t GGSTRootReader::GetEntries() {
  static const std::string routineName("GGSTRootReader::GetEntries");

  if (_readers.size() > 0)
    return _readers[0]->chain->GetEntries();
  else {
    GGSCOUT(WARNING) << "GetEntries may return a wrong number of entries when called before GetReader." << GGSENDL;
    return _GetEntriesBlind();
  }
}

const GGSTSimInfo *GGSTRootReader::GetSimInfo(const GGSTFilesHandler *filesHandler) {

  if (filesHandler == nullptr) {
    if (_filesHandlers.size() > 0)
      filesHandler = _filesHandlers[0];
    else
      return nullptr;
  }

  return &(filesHandler->_simInfo);
}

const GGSTGeoParams *GGSTRootReader::GetGeoParams(const GGSTFilesHandler *filesHandler) {

  if (filesHandler == nullptr) {
    if (_filesHandlers.size() > 0)
      filesHandler = _filesHandlers[0];
    else
      return nullptr;
  }

  return &(filesHandler->_geoParams);
}

const GGSTParameters *GGSTRootReader::GetGenParams(const GGSTFilesHandler *filesHandler) {

  if (filesHandler == nullptr) {
    if (_filesHandlers.size() > 0)
      filesHandler = _filesHandlers[0];
    else
      return nullptr;
  }

  return &(filesHandler->_genParams);
}

const TGeoManager *GGSTRootReader::GetROOTGeometry(GGSTFilesHandler *filesHandler) {
  if (filesHandler == nullptr) {
    if (_filesHandlers.size() > 0)
      filesHandler = _filesHandlers[0];
    else
      return nullptr;
  }
  if (!(filesHandler->_geoManager)) {
    if (filesHandler->_chains.size() != 0) {
      auto readEntry = filesHandler->_chains[0]->GetReadEntry();
      filesHandler->_chains[0]->GetEntry(0);
      if (GGSSmartLog::verboseLevel < GGSSmartLog::DEBUG) {
        GGSSmartLog::MuteOutput();
      }
      filesHandler->_geoManager = (TGeoManager *)(filesHandler->_chains[0]->GetCurrentFile()->Get("GGSGeometry"));
      GGSSmartLog::UnmuteOutput();
      filesHandler->_chains[0]->GetEntry(readEntry);
    } else if (filesHandler->_files.size() != 0) {
      TFile *tGeoFile = TFile::Open(filesHandler->_files[0].c_str());
      filesHandler->_geoManager = (TGeoManager *)(tGeoFile->Get("GGSGeometry"));
      tGeoFile->Close();
      delete tGeoFile;
    }
  }

  return filesHandler->_geoManager;
}

bool GGSTRootReader::HasROOTGeometry(GGSTFilesHandler *filesHandler) {
  if (filesHandler == nullptr) {
    if (_filesHandlers.size() > 0)
      filesHandler = _filesHandlers[0];
    else
      return false;
  }
  if (filesHandler->_geoManager) {
    return true;
  }

  TKey *geoKey = nullptr;
  if (filesHandler->_chains.size() != 0) {
    auto readEntry = filesHandler->_chains[0]->GetReadEntry();
    filesHandler->_chains[0]->GetEntry(0);
    if (GGSSmartLog::verboseLevel < GGSSmartLog::DEBUG) {
      GGSSmartLog::MuteOutput();
    }
    geoKey = (filesHandler->_chains[0]->GetCurrentFile()->GetKey("GGSGeometry"));
    GGSSmartLog::UnmuteOutput();
    filesHandler->_chains[0]->GetEntry(readEntry);
  } else if (filesHandler->_files.size() != 0) {
    TFile *tGeoFile = TFile::Open(filesHandler->_files[0].c_str());
    geoKey = tGeoFile->GetKey("GGSGeometry");
    tGeoFile->Close();
    delete tGeoFile;
  }

  return geoKey;
}