RestOfEventUpdaterModule.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 <analysis/modules/RestOfEventUpdater/RestOfEventUpdaterModule.h>
 
#include <framework/datastore/StoreArray.h>
#include <framework/datastore/StoreObjPtr.h>
 
#include <framework/logging/Logger.h>
 
using namespace std;
 
using namespace Belle2;
 
//-----------------------------------------------------------------
//                 Register the Module
//-----------------------------------------------------------------
REG_MODULE(RestOfEventUpdater);
 
//-----------------------------------------------------------------
//                 Implementation
//-----------------------------------------------------------------
 
RestOfEventUpdaterModule::RestOfEventUpdaterModule() : Module()
{
  // Set module properties
  setDescription("Updates an existing mask (map of boolean values) for tracks or eclClusters in RestOfEvent with an available property (e.g. after performing training).");
  setPropertyFlags(c_ParallelProcessingCertified);
 
  // Add parameters
  std::vector<std::string> emptyMaskVector;
  std::string emptyCutString;
 
  addParam("particleList", m_inputListName, "Name of the ParticleList which contains information that will be used for updating");
  addParam("updateMasks", m_maskNamesForUpdating, "List of all mask names which will be updated", emptyMaskVector);
  addParam("cutString", m_selection, "Cut string which will be used for updating masks", emptyCutString);
  addParam("discard", m_discard,
           "Update the ROE mask by passing or discarding particles in the provided particle list, default is to pass", false);
}
 
void RestOfEventUpdaterModule::initialize()
{
  StoreArray<Particle>().isRequired();
  m_inputList.isRequired(m_inputListName);
 
  m_cut = Variable::Cut::compile(m_selection);
 
  B2INFO("RestOfEventUpdater updated track/eclCluster ROEMask(s) with infoList: " << m_inputListName << " and cut: " << m_selection);
}
 
void RestOfEventUpdaterModule::event()
{
  if (!m_inputList) {
    B2WARNING("Input list " << m_inputList.getName() << " was not created?");
    return;
  }
  StoreObjPtr<RestOfEvent> roe("RestOfEvent");
  if (!roe.isValid()) {
    B2WARNING("ROE list is not valid somehow, ROE masks are not updated!");
    return;
  }
  std::set<Particle::EParticleSourceObject> encounteredSources;
  // Particle lists can contain Particles from different mdst sources
  // Thus, we split them based on their mdst source
  // Only particles surviving the provided cut are considered
  std::vector<const Particle*> particlesFromTracksToUpdate;
  std::vector<const Particle*> particlesFromECLClustersToUpdate;
  std::vector<const Particle*> particlesFromKLMClustersToUpdate;
  std::vector<const Particle*> compositeParticlesToUpdate;
  for (unsigned j = 0; j < m_inputList->getListSize(); j++) {
    const Particle* partWithInfo = m_inputList->getParticle(j);
    Particle::EParticleSourceObject mdstSource = partWithInfo->getParticleSource();
    encounteredSources.insert(mdstSource);
    if (m_cut->check(partWithInfo)) {
      if (mdstSource == Particle::EParticleSourceObject::c_Track) {
        particlesFromTracksToUpdate.push_back(partWithInfo);
      } else if (mdstSource == Particle::EParticleSourceObject::c_ECLCluster) {
        particlesFromECLClustersToUpdate.push_back(partWithInfo);
      } else if (mdstSource == Particle::EParticleSourceObject::c_KLMCluster) {
        particlesFromKLMClustersToUpdate.push_back(partWithInfo);
      } else if (mdstSource == Particle::EParticleSourceObject::c_Composite or
                 mdstSource == Particle::EParticleSourceObject::c_V0) {
        compositeParticlesToUpdate.push_back(partWithInfo);
      }
    }
  }
  if (encounteredSources.count(Particle::EParticleSourceObject::c_Track) > 0) {
    updateMasksWithParticles(roe, particlesFromTracksToUpdate, Particle::EParticleSourceObject::c_Track);
  } else { // If we have a track-based particle in the particle list there can not be any other mdst source
    if (encounteredSources.count(Particle::EParticleSourceObject::c_ECLCluster) > 0) {
      updateMasksWithParticles(roe, particlesFromECLClustersToUpdate, Particle::EParticleSourceObject::c_ECLCluster);
    }
    if (encounteredSources.count(Particle::EParticleSourceObject::c_KLMCluster) > 0) {
      updateMasksWithParticles(roe, particlesFromKLMClustersToUpdate, Particle::EParticleSourceObject::c_KLMCluster);
    }
    updateMasksWithV0(roe, compositeParticlesToUpdate); // in updateMasksWithV0 it is checked whether the vector is empty
  }
}
 
void RestOfEventUpdaterModule::updateMasksWithV0(const StoreObjPtr<RestOfEvent>& roe,
                                                 const std::vector<const Particle*>& particlesToUpdate)
{
  if (particlesToUpdate.size() == 0) {
    B2DEBUG(10, "No particles in list provided, nothing to do");
    return;
  }
  for (auto& maskToUpdate : m_maskNamesForUpdating) {
    if (maskToUpdate == "") {
      B2FATAL("Cannot update ROE mask with no name!");
    }
    for (auto* particleV0 : particlesToUpdate) {
      if (!roe->checkCompatibilityOfMaskAndV0(maskToUpdate, particleV0)) {
        continue;
      }
      roe->updateMaskWithV0(maskToUpdate, particleV0);
    }
  }
}
 
void RestOfEventUpdaterModule::updateMasksWithParticles(const StoreObjPtr<RestOfEvent>& roe,
                                                        std::vector<const Particle*>& particlesToUpdate, Particle::EParticleSourceObject listType)
{
  for (auto& maskToUpdate : m_maskNamesForUpdating) {
    if (maskToUpdate == "") {
      B2FATAL("Cannot update ROE mask with no name!");
    }
    if (!roe->hasMask(maskToUpdate)) {
      // Change name to get all ROE particles in case of new mask
      roe->initializeMask(maskToUpdate, "ROEUpdaterModule");
    }
    roe->excludeParticlesFromMask(maskToUpdate, particlesToUpdate, listType, m_discard);
 
  }
}