BremsFinderModule.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.                  *
 **************************************************************************/
 
// Own header.
#include <analysis/modules/BremsCorrection/BremsFinderModule.h>
 
// framework aux
#include <framework/gearbox/Const.h>
#include <framework/logging/Logger.h>
#include <framework/datastore/RelationArray.h>
#include <framework/datastore/RelationVector.h>
 
// dataobjects
#include <mdst/dataobjects/Track.h>
#include <mdst/dataobjects/ECLCluster.h>
 
// utilities
#include <analysis/DecayDescriptor/ParticleListName.h>
#include <analysis/DecayDescriptor/DecayDescriptor.h>
 
// variables
#include <analysis/variables/ECLVariables.h>
 
#include <cmath>
#include <algorithm>
#include <TMatrixFSym.h>
 
using namespace std;
using namespace Belle2;
 
//-----------------------------------------------------------------
//                 Register module
//-----------------------------------------------------------------
 
REG_MODULE(BremsFinder);
 
//-----------------------------------------------------------------
//                 Implementation
//-----------------------------------------------------------------
 
BremsFinderModule::BremsFinderModule() :
  Module(), m_pdgCode(0)
{
  // set module description (e.g. insert text)
  setDescription(R"DOC(
  This module copies each particle in the ``inputList`` to the ``outputList`` and uses
  the results of the **eclTrackBremFinder** module to look for possible bremsstrahlung photons; if these 
  photons exist, it adds their four momentum to the particle in the ``outputList``.
  It also adds the original particle and these photons as daughters of the new, corrected particle. 
  Track and PID information of the original particle are copied onto the new one to facilitate their access 
  in the analysis scripts.
 
  The **eclTrackBremFinder** module uses the lepton track PXD and SVD hits and extrapolates them to the ECL; 
  then looks for ECL clusters with energies between 0.02 and 1 times the track energy and without associated 
  tracks, and checks if the normalized distance between each of these clusters 
  and the extrapolated hit is smaller than 0.05. If it is, a *Bremsstrahlung* weighted relation 
  between said cluster and the track is established. The weight is determined as
 
  .. math:: 
                 
     \text{max}\left(\frac{\left|\phi_{\text{cluster}}-\phi_{\text{hit}}\right|}{\Delta\phi_{\text{cluster}}+\Delta\phi_{\text{hit}}}, \, \frac{\left|\theta_{\text{cluster}}-\theta_{\text{hit}}\right|}{\Delta\theta_{\text{cluster}}+\Delta\theta_{\text{hit}}}\right)
 
  where :math:`\phi_i` and :math:`\theta_i` are the azimuthal and polar angles of the ECL cluster and the
  extrapolated hit, and :math:`\Delta x` represents the uncertainty of the value :math:`x`. The details of the
  calculation of these quantities are `here`_. By default, only relations with a weight smaller than 3.0 are stored.
  The user can further reduce the maximally allowed value of this weight to remove unwanted photons from the
  bremsstrahlung correction.
 
  This module looks for photons in the ``gammaList`` whose clusters have a *Bremsstrahlung* relation with the track 
  of one of the particles in the ``inputList``, and adds their 4-momentum to the particle's one. It also stores the value
  of each relation weight as ``extraInfo`` of the corrected particle, under the name ``"bremsWeightWithPhotonN"``, where
  N is the index of the photon as daughter of the corrected particle; thus ``"bremsWeightWithPhoton0"`` gives the weight
  of the Bremsstrahlung relation between the new, corrected particle, and the first photon daughter.
 
  Warning:
    Even in the event of no bremsstrahlung photons found, a new particle is still created, and the original one is still 
    added as its daughter.
 
  Warning:
    Studies have shown that the requirements that the energy of the photon must be between 0.2 and 1 times the track energy and
    that only track-photon relations with a weight below three are considered, are too tight. Until these are relaxed and a new
    processing is done (MC15 and proc 13) it might be better to use the alternative `BelleBremRecovery` module.
                
  See also:
    `eclTrackBremFinder module`_
                 
  .. _eclTrackBremFinder module: https://gitlab.desy.de/belle2/software/basf2/-/tree/main/ecl/modules/eclTrackBremFinder
  .. _here: https://gitlab.desy.de/belle2/software/basf2/-/tree/main/ecl/modules/eclTrackBremFinder/src/BremFindingMatchCompute.cc)DOC");
  setPropertyFlags(c_ParallelProcessingCertified);
 
  // Add parameters
  addParam("outputList", m_outputListName, "The output particle list name.");
  addParam("inputList", m_inputListName,
           R"DOC(The initial particle list name containing the particles to correct. *It should already exist* and *the particles must have an associated track.*)DOC");
  addParam("gammaList", m_gammaListName,
           R"DOC(The photon list containing the preselected bremsstrahlung candidates. *It should already exist* and *the particles in the list must be photons*)DOC");
  addParam("maximumAcceptance", m_maximumAcceptance,
           "The maximum value of the relation weight between a bremsstrahlung cluster and a particle track",
           m_maximumAcceptance);
  addParam("multiplePhotons", m_addMultiplePhotons, "If true, use all possible photons to correct the particle's 4-momentum",
           m_addMultiplePhotons);
  addParam("usePhotonOnlyOnce", m_usePhotonOnlyOnce,
           R"DOC(If true, each brems candidate is used to correct maximum 1 particle (the one with the lowest relation weight among all in the ``inputList``).)DOC",
           m_usePhotonOnlyOnce);
  addParam("writeOut", m_writeOut,
           R"DOC(If true, the output ``ParticleList`` will be saved by `RootOutput`. If false, it will be ignored when writing the file.)DOC",
           m_writeOut);
}
 
void BremsFinderModule::initialize()
{
  //Get input lepton particle list
  m_inputList.isRequired(m_inputListName);
 
  DecayDescriptor decDes;
  decDes.init(m_inputListName);
 
  m_pdgCode  = decDes.getMother()->getPDGCode();
  //Check for the particles in the lepton list to have a track associated
  if (Const::chargedStableSet.find(abs(m_pdgCode)) == Const::invalidParticle)
    B2ERROR("[BremsFinderModule] Invalid particle list. the particles in " << m_inputListName << " should have an associated track.");
 
  //Get input photon particle list
  m_gammaList.isRequired(m_gammaListName);
 
  decDes.init(m_gammaListName);
  int temp_pdg = decDes.getMother()->getPDGCode();
 
  //Check that this is a gamma list
  if (temp_pdg != Const::photon.getPDGCode())
    B2ERROR("[BremsFinderModule] Invalid particle list. the particles in " << m_gammaListName << " should be photons!");
 
  decDes.init(m_outputListName);
  temp_pdg = decDes.getMother()->getPDGCode();
  // check the validity of output ParticleList name
  if (m_inputListName == m_outputListName)
    B2ERROR("[BremsFinderModule] Input and output particle list names are the same: " << m_inputListName);
  else if (temp_pdg != m_pdgCode) {
    B2ERROR("[BremsFinderModule] The input and output particle list correspond to different particles: " << m_inputListName << " --> "
            << m_outputListName);
  }
 
  // output particle
  m_outputAntiListName = ParticleListName::antiParticleListName(m_outputListName);
 
  // make output lists
  DataStore::EStoreFlags flags = m_writeOut ? DataStore::c_WriteOut : DataStore::c_DontWriteOut;
  m_outputList.registerInDataStore(m_outputListName, flags);
  m_outputAntiList.registerInDataStore(m_outputAntiListName, flags);
 
  m_particles.registerRelationTo(m_pidlikelihoods);
}
 
void BremsFinderModule::event()
{
  RelationArray particlesToMCParticles(m_particles, m_mcParticles);
 
  // new output particle list
  m_outputList.create();
  m_outputList->initialize(m_pdgCode, m_outputListName);
 
  m_outputAntiList.create();
  m_outputAntiList->initialize(-1 * m_pdgCode, m_outputAntiListName);
  m_outputAntiList->bindAntiParticleList(*(m_outputList));
 
  // Number of photons (calculate it here only once)
  const unsigned int nGamma = m_gammaList->getListSize();
 
  // Number of leptons (calculate it here only once)
  const unsigned int nLep = m_inputList->getListSize();
 
  const std::string relationName = "Bremsstrahlung";
 
  //In the case of only one track per photon
  if (m_usePhotonOnlyOnce) {
    for (unsigned n = 0; n < nGamma; n++) {
      Particle* gamma = m_gammaList->getParticle(n);
      //Skip this if it the best match has already been assigned (pathological case: happens only if you use the same gamma list
      //to correct more than once. Performance studies, in which the same list is used with different options, are an example
      if (gamma->hasExtraInfo("bestMatchIndex")) continue;
 
      auto cluster = gamma->getECLCluster();
      //Get the tracks related to each photon...
      RelationVector<Track> relatedTracks = cluster->getRelationsTo<Track>("", relationName);
      double bestWeight = m_maximumAcceptance;
      unsigned bestMatchIndex = 0;
      unsigned trkIndex = 0;
      //Loop over the related tracks...
      for (auto trk = relatedTracks.begin(); trk != relatedTracks.end(); trk++, trkIndex++) {
        //... and over the input particles' tracks...
        for (unsigned i = 0; i < nLep; i++) {
          const Particle* lepton = m_inputList->getParticle(i);
          auto leptonTrack = lepton->getTrack();
          //... check that the particle track corresponds to the related track....
          if (leptonTrack->getArrayIndex() == trk->getArrayIndex()) {
            double weight = relatedTracks.weight(trkIndex);
            if (weight < bestWeight) {
              bestWeight = weight;
              //... and only select the best match among the tracks in the input list
              bestMatchIndex = trk->getArrayIndex();
            }
            break; //If the particle corresponding to the related track is found, break the loop over the particles and go for the next related track
          }
        }
      }
      //... finally, add the best match index as an extra info for the photon
      gamma->addExtraInfo("bestMatchIndex", bestMatchIndex);
    }
  }
 
  // loop over charged particles, correct them and add them to the output list
 
  for (unsigned i = 0; i < nLep; i++) {
    const Particle* lepton = m_inputList->getParticle(i);
 
    //Get the track of this lepton...
    auto track = lepton->getTrack();
 
    //... and get the bremsstrahlung clusters related to this track
    RelationVector<ECLCluster> bremClusters = track->getRelationsFrom<ECLCluster>("", relationName);
 
    std::vector<std::pair <double, Particle*> > selectedGammas;
 
    unsigned j = 0;
    for (auto bremCluster = bremClusters.begin(); bremCluster != bremClusters.end(); bremCluster++, j++) {
      double weight = bremClusters.weight(j);
 
      if (weight > m_maximumAcceptance) continue;
 
      for (unsigned k = 0; k < nGamma; k++) {
 
        Particle* gamma = m_gammaList->getParticle(k);
        auto cluster = gamma->getECLCluster();
 
        if (bremCluster->getClusterId() == cluster->getClusterId()) {
          if (m_usePhotonOnlyOnce) { //If only one track per photon should be used...
            if (track->getArrayIndex() == static_cast<int>
                (gamma->getExtraInfo("bestMatchIndex")))      //... check if this track is the best match ...
              selectedGammas.push_back(std::make_pair(weight, gamma)); //... and if it is, add it to the selected gammas
          } else {
            selectedGammas.push_back(std::make_pair(weight, gamma));
          }
        }
 
      } // Closes for loop on gammas
 
    } // Closes for loop on brem clusters
 
    //The 4-momentum of the new lepton in the output particle list
    ROOT::Math::PxPyPzEVector new4Vec = lepton->get4Vector();
 
    //Sort weight-particle pairs by weight. Smaller weights go first
    std::sort(selectedGammas.begin(), selectedGammas.end());
 
    //Add to this 4-momentum those of the selected photon(s)
    for (auto const& bremsPair : selectedGammas) {
      Particle* g = bremsPair.second;
      new4Vec += g->get4Vector();
      if (! m_addMultiplePhotons) break; //stop after adding the first photon
    }
 
    //Create the new particle with the 4-momentum calculated before
    Particle correctedLepton(new4Vec, lepton->getPDGCode(), Particle::EFlavorType::c_Flavored, Particle::c_Track,
                             track->getArrayIndex());
 
    //And add the original lepton as its daughter
    correctedLepton.appendDaughter(lepton, false);
 
    const TMatrixFSym& lepErrorMatrix = lepton->getMomentumVertexErrorMatrix();
    TMatrixFSym corLepMatrix(lepErrorMatrix);
 
    double bremsGammaEnergySum = 0.0;
    //Now, if there are any, add the brems photons as daughters as well. As before, we distinguish between the multiple and only one brems photon cases
    int photonIndex = 0;
    for (auto const& bremsPair : selectedGammas) {
      //Add the weights as extra info of the mother
      Particle* bremsGamma = bremsPair.second;
      std::string extraInfoName = "bremsWeightWithPhoton" + std::to_string(photonIndex);
      correctedLepton.addExtraInfo(extraInfoName, bremsPair.first);
      photonIndex++;
      bremsGammaEnergySum += Variable::eclClusterE(bremsGamma);
 
      const TMatrixFSym& gammaErrorMatrix = bremsGamma->getMomentumVertexErrorMatrix();
      for (int irow = 0; irow <= 3; irow++) {
        for (int icol = irow; icol <= 3; icol++) corLepMatrix(irow, icol) += gammaErrorMatrix(irow, icol);
      }
      correctedLepton.appendDaughter(bremsGamma, false);
      B2DEBUG(10, "[BremsFinderModule] Found a bremsstrahlung gamma and added its 4-vector to the charged particle");
      if (! m_addMultiplePhotons) break; //stop after adding the first photon
    }
 
    correctedLepton.setMomentumVertexErrorMatrix(corLepMatrix);
 
    // add the info from original lepton to the new lepton
    correctedLepton.setVertex(lepton->getVertex());
    correctedLepton.setPValue(lepton->getPValue());
    correctedLepton.addExtraInfo("bremsCorrected", float(selectedGammas.size() > 0));
    correctedLepton.addExtraInfo("bremsCorrectedPhotonEnergy", bremsGammaEnergySum);
 
    // add the mc relation
    Particle* newLepton = m_particles.appendNew(correctedLepton);
    const MCParticle* mcLepton = lepton->getRelated<MCParticle>();
    const PIDLikelihood* pid = lepton->getPIDLikelihood();
 
    if (pid) newLepton->addRelationTo(pid);
 
    if (mcLepton != nullptr) newLepton->addRelationTo(mcLepton);
 
    m_outputList->addParticle(newLepton);
 
  } //Closes for loop on leptons
 
} //Close event()