MCInitialParticles.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2015 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Martin Ritter                                            *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
 
#pragma once
 
#include <TLorentzVector.h>
#include <TLorentzRotation.h>
 
namespace Belle2 {
  class MCInitialParticles: public TObject {
 
  public:
 
    enum EGenerationFlags {
      c_generateCMS = 1 << 0,
      c_smearBeamEnergy = 1 << 1,
      c_smearBeamDirection = 1 << 2,
      c_smearBeam = c_smearBeamEnergy | c_smearBeamDirection,
      c_smearVertex = 1 << 3,
      c_smearALL = c_smearVertex | c_smearBeam,
    };
 
    MCInitialParticles(): TObject() {}
 
    MCInitialParticles(const MCInitialParticles& b): TObject(), m_her(b.m_her), m_ler(b.m_ler),
      m_vertex(b.m_vertex), m_validFlag(b.m_validFlag), m_generationFlags(b.m_generationFlags) {}
 
    MCInitialParticles& operator=(const MCInitialParticles& b)
    {
      m_her = b.m_her; m_ler = b.m_ler; m_vertex = b.m_vertex;
      m_validFlag = b.m_validFlag; m_generationFlags = b.m_generationFlags;
      resetBoost();
      return *this;
    }
 
    bool operator==(const MCInitialParticles& b) const
    {
      // FIXME: Ah, bloody hell. sin(x) returns slightly different values on
      // different platforms in some cases so we cannot just do an equality
      // comparison. We need to do this more elegantly, this is just for
      // testing if it solves all problems
#if defined(MCP_DBL_CMP) || defined(MCP_VEC3_CMP) || defined(MCP_VEC4_CMP)
#error Macro already defined, cannot continue
#endif
#define MCP_DBL_CMP(a,b,x) ((a.x()==b.x())||(std::abs(a.x()-b.x())<1e-10))
#define MCP_VEC3_CMP(a,b) (MCP_DBL_CMP(a,b,X) && MCP_DBL_CMP(a,b,Y) && MCP_DBL_CMP(a,b,Z))
#define MCP_VEC4_CMP(a,b) (MCP_VEC3_CMP(a,b) && MCP_DBL_CMP(a,b,E))
      return MCP_VEC4_CMP(m_her, b.m_her) && MCP_VEC4_CMP(m_ler, b.m_ler) && MCP_VEC3_CMP(m_vertex, b.m_vertex)
             && m_validFlag == b.m_validFlag &&
             m_generationFlags == b.m_generationFlags;
#undef MCP_DBL_CMP
#undef MCP_VEC3_CMP
#undef MCP_VEC4_CMP
    }
 
    virtual ~MCInitialParticles() { resetBoost(); }
 
    void set(const TLorentzVector& her, const TLorentzVector& ler, const TVector3& vertex)
    {
      m_her = her;
      m_ler = ler;
      m_vertex = vertex;
      m_validFlag = true;
      resetBoost();
    }
 
    void setHER(const TLorentzVector& her)
    {
      m_her = her;
      resetBoost();
    }
 
    void setLER(const TLorentzVector& ler)
    {
      m_ler = ler;
      resetBoost();
    }
 
    void setVertex(const TVector3& vertex)
    {
      m_vertex = vertex;
    }
 
    void setTime(double time) {m_time = time;}
 
    void setGenerationFlags(int flags) { m_generationFlags = flags; }
 
    const TLorentzVector& getHER() const { return m_her; }
 
    const TLorentzVector& getLER() const { return m_ler; }
 
    const TVector3& getVertex() const { return m_vertex; }
 
    double getTime() const {return m_time;}
 
    double getEnergy() const { return (m_her + m_ler).E(); }
 
    double getMass() const { calculateBoost(); return m_invariantMass; }
 
    const TLorentzRotation& getLabToCMS() const
    {
      calculateBoost(); return *m_labToCMS;
    }
 
    const TLorentzRotation& getCMSToLab() const
    {
      calculateBoost(); return *m_CMSToLab;
    }
 
    bool getValidFlag() const { return m_validFlag; }
 
    int getGenerationFlags() const { return m_generationFlags; }
 
    bool hasGenerationFlags(int flags) const { return (m_generationFlags & flags) == flags; }
 
    std::string getGenerationFlagString(const std::string& separator = " ") const;
 
  private:
 
    void calculateBoost() const;
    void resetBoost();
    TLorentzVector m_her;
    TLorentzVector m_ler;
    TVector3 m_vertex;
    double m_time = 0;
    mutable TLorentzRotation* m_labToCMS{nullptr}; 
 
    mutable TLorentzRotation* m_CMSToLab{nullptr}; 
 
    mutable double m_invariantMass{0.0}; 
 
    bool m_validFlag = false;
    int m_generationFlags{0};
    ClassDef(MCInitialParticles, 3);
  };
 
  inline void MCInitialParticles::calculateBoost() const
  {
    if (m_labToCMS)
      return;
 
    TLorentzVector beam = m_her + m_ler;
    // Save the invariant mass because it's used very often in analysis
    m_invariantMass = beam.M();
 
    // If we generate events in CMS we already are in CMS and there is no
    // transformation so let's use the identity
    if (hasGenerationFlags(c_generateCMS)) {
      m_labToCMS = new TLorentzRotation();
      m_CMSToLab = new TLorentzRotation();
      return;
    }
 
    // Transformation from Lab system to CMS system
    m_labToCMS = new TLorentzRotation(-beam.BoostVector());
    // boost HER e- from Lab system to CMS system
    const TLorentzVector electronCMS = (*m_labToCMS) * m_her;
    // now rotate CMS such that incoming e- is parallel to z-axis
    const TVector3 zaxis(0., 0., 1.);
    TVector3 rotaxis = zaxis.Cross(electronCMS.Vect()) * (1. / electronCMS.Vect().Mag());
    double rotangle = TMath::ASin(rotaxis.Mag());
    m_labToCMS->Rotate(-rotangle, rotaxis);
 
    //cache derived quantities
    m_CMSToLab = new TLorentzRotation(m_labToCMS->Inverse());
  }
 
  inline void MCInitialParticles::resetBoost()
  {
    delete m_labToCMS;
    delete m_CMSToLab;
    m_labToCMS = nullptr;
    m_CMSToLab = nullptr;
    m_invariantMass = 0.0;
  }
 
} //Belle2 namespace