TrackIsoCalculatorModule.h

/**************************************************************************
 * 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.                  *
 **************************************************************************/
 
#pragma once
 
#include <framework/core/Module.h>
#include <framework/datastore/StoreObjPtr.h>
#include <framework/dataobjects/EventMetaData.h>
#include <framework/logging/LogConfig.h>
#include <framework/datastore/StoreArray.h>
#include <framework/gearbox/Const.h>
#include <framework/database/DBObjPtr.h>
 
#include <analysis/dataobjects/Particle.h>
#include <analysis/dataobjects/ParticleList.h>
#include <analysis/VariableManager/Manager.h>
#include <analysis/VariableManager/Utility.h>
#include <analysis/DecayDescriptor/DecayDescriptor.h>
#include <analysis/dbobjects/PIDDetectorWeights.h>
 
 
namespace Belle2 {
  class DetSurfCylBoundaries {
 
  public:
    DetSurfCylBoundaries() {};
    DetSurfCylBoundaries(float rho, float zfwd, float zbwd, float th_fwd, float th_fwd_brl, float th_bwd_brl, float th_bwd)
    {
      m_rho = rho;
      m_zfwd = zfwd;
      m_zbwd = zbwd;
      m_th_fwd = th_fwd;
      m_th_fwd_brl = th_fwd_brl;
      m_th_bwd_brl = th_bwd_brl;
      m_th_bwd = th_bwd;
    };
    float m_rho; 
    float m_zfwd; 
    float m_zbwd; 
    float m_th_fwd; 
    float m_th_fwd_brl; 
    float m_th_bwd_brl; 
    float m_th_bwd; 
  };
 
  class TrackIsoCalculatorModule : public Module {
 
  public:
 
    TrackIsoCalculatorModule();
 
    ~TrackIsoCalculatorModule() override;
 
    void initialize() override;
 
    void event() override;
 
 
  private:
 
    StoreArray<Particle> m_particles;
 
    std::string m_decayString;
 
    unsigned short m_nSelectedDaughters;
 
    std::string m_pListReferenceName;
 
    std::string m_detName;
 
    unsigned int m_nLayers;
 
    std::unordered_map<std::string, std::string>  m_detLayerToDistVariable;
 
    std::unordered_map<std::string, std::string>  m_detLayerToRefPartIdxVariable;
 
    std::string m_isoScoreVariable;
 
    std::unordered_map<std::string, std::vector<int>> m_detToLayers = {
      {"CDC", {0, 1, 2, 3, 4, 5, 6, 7, 8}},
      {"TOP", {0}},
      {"ARICH", {0}},
      {"ECL", {0, 1}},
      {"KLM", {0}}
    };
 
    std::unordered_map<std::string, DetSurfCylBoundaries> m_detLayerToSurfBoundaries = {
      {"CDC0", DetSurfCylBoundaries(16.8, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC1", DetSurfCylBoundaries(25.7, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC2", DetSurfCylBoundaries(36.52, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC3", DetSurfCylBoundaries(47.69, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC4", DetSurfCylBoundaries(58.41, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC5", DetSurfCylBoundaries(69.53, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC6", DetSurfCylBoundaries(80.25, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC7", DetSurfCylBoundaries(91.37, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"CDC8", DetSurfCylBoundaries(102.09, 150.0, -75.0, 0.0, 0.29, 2.61, 3.14)},
      {"TOP0", DetSurfCylBoundaries(117.8, 193.0, -94.0, 0.24, 0.52, 2.23, 3.14)},
      {"ARICH0", DetSurfCylBoundaries(117.8, 193.0, -94.0, 0.24, 0.52, 2.23, 3.14)},
      {"ECL0", DetSurfCylBoundaries(125.0, 196.0, -102.0, 0.21, 0.56, 2.24, 2.70)},
      {"ECL1", DetSurfCylBoundaries(140.0, 211.0, -117.0, 0.21, 0.56, 2.24, 2.70)},
      {"KLM0", DetSurfCylBoundaries(202.0, 283.9, -189.9, 0.40, 0.82, 2.13, 2.60)},
    };
 
    std::map<std::pair<std::string, int>, double> m_distThreshPerDetLayer = {
      { {Const::parseDetectors(Const::CDC), 0}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 1}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 2}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 3}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 4}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 5}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 6}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 7}, 5.0 },
      { {Const::parseDetectors(Const::CDC), 8}, 5.0 },
      { {Const::parseDetectors(Const::TOP), 0}, 22.0 },
      { {Const::parseDetectors(Const::ARICH), 0}, 10.0 },
      { {Const::parseDetectors(Const::ECL), 0}, 36.0 },
      { {Const::parseDetectors(Const::ECL), 1}, 36.0 },
      { {Const::parseDetectors(Const::KLM), 0}, 20.0 }
    };
 
    StoreObjPtr<EventMetaData> m_event_metadata;
 
    StoreObjPtr<ParticleList> m_pListTarget;
 
    DecayDescriptor m_decaydescriptor;
 
    StoreObjPtr<ParticleList> m_pListReference;
 
    bool m_useHighestProbMassForExt;
 
    bool m_excludePIDDetWeights;
 
    std::string m_payloadName;
 
    std::unique_ptr<DBObjPtr<PIDDetectorWeights>> m_DBWeights;
 
    double getDistAtDetSurface(const Particle* iParticle, const Particle* jParticle, const std::string& detLayerName) const;
 
    double getIsoScore(const Particle* iParticle) const;
 
    double getDistThreshold(Const::EDetector det, int layer) const
    {
      auto detAndLayer = std::make_pair(Const::parseDetectors(det), layer);
      return m_distThreshPerDetLayer.at(detAndLayer);
    };
 
    bool onlySelectedStdChargedInDecay();
 
    Const::EDetector getDetEnum(const std::string& detName) const
    {
 
      if (detName == "CDC") return Const::CDC;
      else if (detName == "TOP") return Const::TOP;
      else if (detName == "ARICH") return Const::ARICH;
      else if (detName == "ECL") return Const::ECL;
      else if (detName == "KLM") return Const::KLM;
      else B2FATAL("Unknown detector component: " << detName);
 
    };
 
  };
}