MCInitialParticles.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/gearbox/Const.h>
#include <framework/utilities/LabToCms.h>
 
#include <TVector3.h>
#include <Math/Boost.h>
#include <Math/LorentzRotation.h>
#include <Math/Vector3D.h>
#include <Math/Vector4D.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: 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 ROOT::Math::PxPyPzEVector& her, const ROOT::Math::PxPyPzEVector& ler, const ROOT::Math::XYZVector& vertex)
    {
      m_her = her;
      m_ler = ler;
      m_vertex = vertex;
      m_validFlag = true;
      resetBoost();
    }
 
    void setByLorentzTransformation(double Ecms, double bX, double bY, double bZ, double angleXZ, double angleYZ,
                                    const ROOT::Math::XYZVector& vertex)
    {
      if (m_labToCMS) delete m_labToCMS;
      if (m_CMSToLab) delete m_CMSToLab;
 
      m_invariantMass = Ecms;
      m_CMSToLab      = new ROOT::Math::LorentzRotation();
      m_labToCMS      = new ROOT::Math::LorentzRotation();
      *m_CMSToLab     = cmsToLab(bX, bY, bZ, angleXZ, angleYZ);
      *m_labToCMS     = m_CMSToLab->Inverse();
 
      const double me = Const::electron.getMass();
      double p = sqrt(Ecms * Ecms / 4 - me * me);
      m_her = (*m_CMSToLab) * ROOT::Math::PxPyPzEVector(0.0, 0.0,  p, Ecms / 2);
      m_ler = (*m_CMSToLab) * ROOT::Math::PxPyPzEVector(0.0, 0.0, -p, Ecms / 2);
 
      m_vertex = vertex;
      m_validFlag = true;
    }
 
    void setHER(const ROOT::Math::PxPyPzEVector& her)
    {
      m_her = her;
      resetBoost();
    }
 
    void setLER(const ROOT::Math::PxPyPzEVector& ler)
    {
      m_ler = ler;
      resetBoost();
    }
 
    void setVertex(const ROOT::Math::XYZVector& vertex)
    {
      m_vertex = vertex;
    }
 
    void setTime(double time) {m_time = time;}
 
    virtual void setGenerationFlags(int flags) { m_generationFlags = flags; }
 
    const ROOT::Math::PxPyPzEVector& getHER() const { return m_her; }
 
    const ROOT::Math::PxPyPzEVector& getLER() const { return m_ler; }
 
    const ROOT::Math::XYZVector& 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; }
 
    ROOT::Math::XYZVector getBoostVector() const { calculateBoost(); return -(m_her + m_ler).BoostToCM(); }
 
    ROOT::Math::PxPyPzEVector getBoostedHER() const { calculateBoost(); return m_boostedHER; }
 
    const ROOT::Math::LorentzRotation& getLabToCMS() const
    {
      calculateBoost(); return *m_labToCMS;
    }
 
    const ROOT::Math::LorentzRotation& 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;
 
 
    static ROOT::Math::LorentzRotation cmsToLab(double bX, double bY, double bZ, double angleXZ, double angleYZ);
 
  private:
 
    void calculateBoost() const;
    void resetBoost();
    ROOT::Math::PxPyPzEVector m_her;
    ROOT::Math::PxPyPzEVector m_ler;
    ROOT::Math::XYZVector m_vertex;
    double m_time = 0;
    mutable ROOT::Math::LorentzRotation* m_labToCMS{nullptr}; 
    mutable ROOT::Math::LorentzRotation* m_CMSToLab{nullptr}; 
    mutable double m_invariantMass{0.0}; 
    mutable ROOT::Math::PxPyPzEVector m_boostedHER; 
    bool m_validFlag = false;
 
  protected:
 
    int m_generationFlags{0};
 
  private:
 
    ClassDef(MCInitialParticles, 5);
  };
 
  inline void MCInitialParticles::calculateBoost() const
  {
    if (m_labToCMS)
      return;
 
    ROOT::Math::PxPyPzEVector 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 ROOT::Math::LorentzRotation();
      m_CMSToLab = new ROOT::Math::LorentzRotation();
      return;
    }
 
    m_boostedHER = ROOT::Math::LorentzRotation(ROOT::Math::Boost(beam.BoostToCM())) * m_her;
    double cmsAngleXZ = atan(m_boostedHER.X() / m_boostedHER.Z());
    double cmsAngleYZ = atan(m_boostedHER.Y() / m_boostedHER.Z());
    m_labToCMS = new ROOT::Math::LorentzRotation(
      LabToCms::rotateLabToCms(-beam.BoostToCM(), cmsAngleXZ, cmsAngleYZ)
    );
 
    //cache derived quantities
    m_CMSToLab = new ROOT::Math::LorentzRotation(m_labToCMS->Inverse());
  }
 
 
 
 
 
  inline void MCInitialParticles::resetBoost()
  {
    delete m_labToCMS;
    delete m_CMSToLab;
    m_labToCMS = nullptr;
    m_CMSToLab = nullptr;
    m_invariantMass = 0.0;
    m_boostedHER = ROOT::Math::PxPyPzEVector();
  }
 
} //Belle2 namespace