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 <TLorentzVector.h>
#include <Math/AxisAngle.h>
#include <Math/Boost.h>
#include <Math/LorentzRotation.h>
#include <Math/Vector3D.h>
#include <Math/Vector4D.h>
 
#include <framework/geometry/B2Vector3.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 B2Vector3D& vertex)
    {
      m_her = her;
      m_ler = ler;
      m_vertex = vertex;
      m_validFlag = true;
      resetBoost();
    }
 
    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 B2Vector3D& 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 B2Vector3D& 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 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;
 
  private:
 
    void calculateBoost() const;
    void resetBoost();
    ROOT::Math::PxPyPzEVector m_her;
    ROOT::Math::PxPyPzEVector m_ler;
    B2Vector3D 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}; 
    bool m_validFlag = false;
 
  protected:
 
    int m_generationFlags{0};
 
  private:
 
    ClassDef(MCInitialParticles, 4);
  };
 
  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;
    }
 
    // Transformation from Lab system to CMS system
    m_labToCMS = new ROOT::Math::LorentzRotation(ROOT::Math::Boost(beam.BoostToCM()));
    // boost HER e- from Lab system to CMS system
    const ROOT::Math::PxPyPzEVector electronCMS = (*m_labToCMS) * m_her;
    // now rotate CMS such that incoming e- is parallel to z-axis
    const ROOT::Math::XYZVector zaxis(0., 0., 1.);
    ROOT::Math::XYZVector rotaxis = zaxis.Cross(electronCMS.Vect()) / electronCMS.P();
    double rotangle = TMath::ASin(rotaxis.R());
    const ROOT::Math::LorentzRotation rotation(ROOT::Math::AxisAngle(rotaxis, -rotangle));
    *m_labToCMS = rotation * (*m_labToCMS);
 
    //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;
  }
 
} //Belle2 namespace