RestOfEvent.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2010 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Anze Zupanc, Matic Lubej, Sviatoslav Bilokin             *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
 
#pragma once
#include <framework/datastore/RelationsObject.h>
#include <mdst/dataobjects/PIDLikelihood.h>
#include <analysis/VariableManager/Utility.h>
#include <analysis/dataobjects/Particle.h>
 
#include <framework/logging/Logger.h>
 
#include <vector>
#include <string>
#include <set>
 
class TLorentzVector;
 
namespace Belle2 {
  // forward declarations
  class Particle;
  class ECLCluster;
  class KLMCluster;
  class Track;
 
  // TODO: Add support for the MdstVee dataobjects when they become available.
 
 
  class RestOfEvent : public RelationsObject {
 
  public:
    struct Mask {
    public:
      Mask(const std::string& name = "", const std::string& origin = "unknown"): m_name(name),
        m_origin(origin)
      {
        B2DEBUG(10, "Mask " << name << " is being initialized by " << origin);
        m_isValid = false;
      };
      std::string getName() const
      {
        return m_name;
      }
      bool isValid() const
      {
        return m_isValid;
      }
      void addParticles(const std::vector<const Particle*>& particles)
      {
        if (isValid()) {
          B2INFO("Mask " + m_name + " originating from "  + m_origin + " is  valid, cannot write to it!");
          return;
        } else {
          for (auto* particle : particles) {
            m_maskedParticleIndices.insert(particle->getArrayIndex());
          }
          m_isValid = true;
        }
      }
      std::set<int> getParticles() const
      {
        return m_maskedParticleIndices;
      }
      std::set<int> getV0s() const
      {
        return m_maskedV0Indices;
      }
      void addV0(const Particle* v0, std::vector<int>& toErase)
      {
        m_maskedV0Indices.insert(v0->getArrayIndex());
        for (int& i : toErase) {
          m_maskedParticleIndices.erase(i);
        }
        m_maskedParticleIndices.insert(v0->getArrayIndex());
      }
      bool hasV0(const Particle* v0) const
      {
        return m_maskedV0Indices.count(v0->getArrayIndex()) > 0;
      }
      void clearParticles()
      {
        m_maskedParticleIndices.clear();
        m_maskedV0Indices.clear();
        m_isValid = false;
      }
      void print()
      {
        B2INFO("Mask name: " + m_name + " originating from " + m_origin);
        if (!m_isValid) {
          B2INFO("\tNot valid!");
        }
        std::string printout =  "\tIndices: ";
        for (const int index : m_maskedParticleIndices) {
          printout += std::to_string(index) +  ", ";
        }
        B2INFO(printout);
      }
    private:
      std::string m_name;                       
      std::string m_origin;                     
      bool m_isValid;                           
      std::set<int> m_maskedParticleIndices;    
      std::set<int> m_maskedV0Indices;          
    };
    explicit RestOfEvent(int pdgCode = 0, bool isNested = false, bool isFromMC = false):
      m_pdgCode(pdgCode), m_isNested(isNested), m_isFromMC(isFromMC) { };
    // setters
    void addParticles(const std::vector<const Particle*>& particle);
    void setPDGCode(int pdgCode)
    {
      m_pdgCode = pdgCode;
    }
    Particle* convertToParticle(const std::string& maskName = "", int pdgCode = 0, bool isSelfConjugated = true);
    int getPDGCode() const
    {
      return m_pdgCode;
    }
    bool hasParticle(const Particle* particle, const std::string& maskName = "") const;
    void initializeMask(const std::string& name, const std::string& origin = "unknown");
    void updateMaskWithCuts(const std::string& name, const std::shared_ptr<Variable::Cut>& trackCut = nullptr,
                            const std::shared_ptr<Variable::Cut>& eclCut = nullptr, const std::shared_ptr<Variable::Cut>& klmCut = nullptr,
                            bool updateExisting = false);
    void excludeParticlesFromMask(const std::string& maskName, const std::vector<const Particle*>& particles,
                                  Particle::EParticleSourceObject listType,
                                  bool discard);
    bool hasMask(const std::string& name) const;
    void updateMaskWithV0(const std::string& name, const Particle* particleV0);
    bool checkCompatibilityOfMaskAndV0(const std::string& name, const Particle* particleV0);
    bool getIsNested() const {return m_isNested;}
    // getters
    std::vector<const Particle*> getParticles(const std::string& maskName = "", bool unpackComposite = true) const;
    std::vector<const Particle*> getPhotons(const std::string& maskName = "", bool unpackComposite = true) const;
    std::vector<const Particle*> getHadrons(const std::string& maskName = "", bool unpackComposite = true) const;
    std::vector<const Particle*> getChargedParticles(const std::string& maskName = "", unsigned int pdg = 0,
                                                     bool unpackComposite = true) const;
 
    std::vector<const Track*> getTracks(const std::string& maskName = "") const;
 
    std::vector<const ECLCluster*> getECLClusters(const std::string& maskName = "") const;
 
    std::vector<const KLMCluster*> getKLMClusters(const std::string& maskName = "") const;
 
    TLorentzVector get4Vector(const std::string& maskName = "") const;
 
    TLorentzVector get4VectorTracks(const std::string& maskName = "") const;
 
    TLorentzVector get4VectorNeutralECLClusters(const std::string& maskName = "") const;
 
    int getNTracks(const std::string& maskName = "") const;
 
    int getNECLClusters(const std::string& maskName = "") const;
 
    int getNKLMClusters(const std::string& maskName = "") const;
 
    std::vector<std::string> getMaskNames() const;
 
    double atcPIDBelleKpiFromPID(const PIDLikelihood* pid) const;
 
    void print() const;
 
  private:
 
    // persistent data members
    std::set<int> m_particleIndices;   
    std::vector<Mask> m_masks;         
    int m_pdgCode;                     
    bool m_isNested;                   
    bool m_isFromMC;                   
    // Private methods
    bool isInParticleList(const Particle* roeParticle, const std::vector<const Particle*>& particlesToUpdate) const;
 
    Mask* findMask(const std::string& name);
    void printIndices(const std::set<int>& indices) const;
    ClassDef(RestOfEvent, 5) 
  };
 
 
} // end namespace Belle2