VXDTFTrainingDataCollectorModule.cc

/**************************************************************************
 * 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.                  *
 **************************************************************************/
 
#include <tracking/trackFindingVXD/filterMap/map/FiltersContainer.h>
#include <tracking/trackFindingVXD/sectorMapTools/SecMapTrainer.h>
 
#include <tracking/modules/vxdtfRedesign/VXDTFTrainingDataCollectorModule.h>
 
#include <framework/logging/Logger.h>
#include <framework/datastore/StoreObjPtr.h>
 
#include <TRandom.h>
 
using namespace Belle2;
 
REG_MODULE(VXDTFTrainingDataCollector);
 
VXDTFTrainingDataCollectorModule::VXDTFTrainingDataCollectorModule() :
  Module(),
  m_eventData("EventMetaData", DataStore::c_Event)
{
 
  //Set module properties
  setDescription("This module collect the data needed to train the VXDTF sector map. The data are saved on a root file that eventually will be chained and used by the training module.");
  setPropertyFlags(c_ParallelProcessingCertified |
                   c_TerminateInAllProcesses);
 
  addParam("SpacePointTrackCandsName", m_PARAMSpacePointTrackCandsName,
           "the name of the storeArray containing the SpacePointTrackCands used for extracting and collecting the training data.",
           std::string(""));
 
  addParam("NameTag", m_PARAMNameTag, "A name tag that will be attached to the name of the output file. If left empty (\"\") a "
           "random number will be attached!", std::string(""));
 
  addParam("outputDir", m_PARAMoutputDir,
           "Name of the output directory. The output file created by this module will be written into that directory.", m_PARAMoutputDir);
}
 
void VXDTFTrainingDataCollectorModule::initialize()
{
 
  FiltersContainer<SpacePoint>& filtersContainer = Belle2::FiltersContainer<SpacePoint>::getInstance();
 
  m_spacePointTrackCands.isRequired(m_PARAMSpacePointTrackCandsName);
 
  for (auto setup : filtersContainer.getAllSetups()) {
    auto config = setup.second->getConfig();
 
    std::string nameAppendix = m_PARAMNameTag;
    // if the name tag was not set a random number will be attached!
    if (nameAppendix == std::string("")) {
      int randomInt = gRandom->Integer(std::numeric_limits<int>::max());
      nameAppendix = std::to_string(randomInt);
    }
 
    SecMapTrainer<SelectionVariableFactory<SecMapTrainerHit> >
    newMap(setup.first, nameAppendix, m_PARAMoutputDir);
 
    m_secMapTrainers.push_back(std::move(newMap));
 
  }
 
  for (auto& trainer : m_secMapTrainers) {
    trainer.initialize();
  }
 
 
}
 
 
void VXDTFTrainingDataCollectorModule::event()
{
  //get the data
  int thisExperiment = m_eventData->getExperiment();
  int thisRun = m_eventData->getRun();
  int thisEvent = m_eventData->getEvent();
 
  for (auto& collector : m_secMapTrainers)
    collector.initializeEvent(thisExperiment, thisRun, thisEvent);
 
 
  // number of space point track candidates from the MC track finder
  unsigned nSPTCs = m_spacePointTrackCands.getEntries();
 
 
  for (unsigned iTC = 0; iTC < nSPTCs; ++ iTC) {
    const SpacePointTrackCand* currentTC = m_spacePointTrackCands[iTC];
    for (auto& dataCollector : m_secMapTrainers)
      dataCollector.storeTC(*currentTC, iTC);
  }
 
  // process raw data:
  for (auto& dataCollector : m_secMapTrainers) {
    unsigned nTCsProcessed = dataCollector.processTracks();
 
    B2DEBUG(20,
            "VXDTFTrainingDataCollectorModule, event " << thisEvent <<
            " with mapTrainer " << dataCollector.getConfig().secMapName <<
            ": number of TCs processed: " << nTCsProcessed <<
            ", calculations done!");
  }
}
 
void VXDTFTrainingDataCollectorModule::terminate()
{
  for (auto& dataCollector : m_secMapTrainers)
    dataCollector.terminate();
}