B2Vector3.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/logging/Logger.h>
 
#include <TVector3.h>
#include <string>
#include <iostream>     // std::cout, std::fixed
#include <iomanip>      // std::setprecision
#include <typeinfo>
#include <cmath>
 
 
namespace Belle2 {
  template<typename DataType>
  class B2Vector3 {
  protected:
    static_assert(std::is_floating_point<DataType>::value, "B2Vector3 only works with floating point types");
    DataType m_coordinates[3];
  public:
    typedef DataType value_type;
 
    B2Vector3(void) : m_coordinates {static_cast<DataType>(0), static_cast<DataType>(0), static_cast<DataType>(0)} {};
    B2Vector3(const DataType xVal, const DataType yVal, const DataType zVal): m_coordinates {xVal, yVal, zVal} {};
    explicit B2Vector3(const DataType(& coords)[3]): m_coordinates {coords[0], coords[1], coords[2]} {};
    explicit B2Vector3(const DataType(* coords)[3]): m_coordinates {(*coords)[0], (*coords)[1], (*coords)[2]} {};
    // cppcheck-suppress noExplicitConstructor
    B2Vector3(const TVector3& tVec3): m_coordinates {static_cast<DataType>(tVec3.X()), static_cast<DataType>(tVec3.Y()), static_cast<DataType>(tVec3.Z())} {};
    // cppcheck-suppress noExplicitConstructor
    B2Vector3(const TVector3* tVec3): m_coordinates {static_cast<DataType>(tVec3->X()), static_cast<DataType>(tVec3->Y()), static_cast<DataType>(tVec3->Z())} {};
    explicit B2Vector3(const B2Vector3<DataType>& b2Vec3): m_coordinates {b2Vec3.X(), b2Vec3.Y(), b2Vec3.Z()} {};
    explicit B2Vector3(const B2Vector3<DataType>* b2Vec3): m_coordinates {b2Vec3->X(), b2Vec3->Y(), b2Vec3->Z()} {};
    // cppcheck-suppress noExplicitConstructor
    template <typename OtherType> B2Vector3(const B2Vector3<OtherType>& b2Vec3):
      m_coordinates {static_cast<DataType>(b2Vec3.X()), static_cast<DataType>(b2Vec3.Y()), static_cast<DataType>(b2Vec3.Z())} {};
    template <typename OtherType> explicit B2Vector3(const B2Vector3<OtherType>* b2Vec3):
      m_coordinates {static_cast<DataType>(b2Vec3->X()), static_cast<DataType>(b2Vec3->Y()), static_cast<DataType>(b2Vec3->Z())} {};
 
    DataType operator()(unsigned i) const { return m_coordinates[i]; }
    DataType operator[](unsigned i) const { return m_coordinates[i]; }
    DataType& operator()(unsigned i) { return m_coordinates[i]; }
    DataType& operator[](unsigned i) { return m_coordinates[i]; }
 
    B2Vector3<DataType>& operator = (const B2Vector3<DataType>& b);
    B2Vector3<DataType>& operator = (const TVector3& b);
 
    operator TVector3() const { return GetTVector3(); }
 
    bool operator == (const B2Vector3<DataType>& b) const { return X() == b.X() && Y() == b.Y() && Z() == b.Z(); }
    bool operator == (const TVector3& b) const { return X() == b.X() && Y() == b.Y() && Z() == b.Z(); }
    bool operator != (const B2Vector3<DataType>& b) const { return !(*this == b); }
    bool operator != (const TVector3& b) const { return !(*this == b); }
 
    B2Vector3<DataType>& operator += (const B2Vector3<DataType>& b);
    B2Vector3<DataType>& operator -= (const B2Vector3<DataType>& b);
    B2Vector3<DataType>& operator *= (DataType a);
    B2Vector3<DataType> operator - () const { return B2Vector3<DataType>(-X(), -Y(), -Z()); }
    B2Vector3<DataType> operator + (const B2Vector3<DataType>& b) const
    {
      return B2Vector3<DataType>(X() + b.X(), Y() + b.Y(), Z() + b.Z());
    }
    B2Vector3<DataType> operator - (const B2Vector3<DataType>& b) const
    {
      return B2Vector3<DataType>(X() - b.X(), Y() - b.Y(), Z() - b.Z());
    }
    B2Vector3<DataType> operator * (DataType a) const
    {
      return B2Vector3<DataType>(a * X(), a * Y(), a * Z());
    }
    B2Vector3<DataType> operator / (DataType a) const
    {
      return B2Vector3<DataType>(X() / a, Y() / a, Z() / a);
    }
    DataType operator * (const B2Vector3<DataType>& b) const { return Dot(b); }
 
 
    DataType Phi() const { return X() == 0 && Y() == 0 ? 0 : atan2(Y(), X()); }
    DataType Theta() const { return X() == 0 && Y() == 0 && Z() == 0 ? 0 : atan2(Perp(), Z()); }
    DataType CosTheta() const { const double pTot = Mag(); return pTot == 0 ? 1 : Z() / pTot; }
    DataType Mag2() const { return X() * X() + Y() * Y() + Z() * Z(); }
    DataType Mag() const { return std::hypot((double)Perp(), (double)Z()); }
 
    void SetPhi(DataType phi)
    {
      const double perp   = Perp();
      SetX(perp * cos((double)phi));
      SetY(perp * sin((double)phi));
    }
 
    void SetTheta(DataType theta)
    {
      const double ma   = Mag();
      const double ph   = Phi();
      const double ctheta = std::cos((double) theta);
      const double stheta = std::sin((double) theta);
      SetX(ma * stheta * std::cos(ph));
      SetY(ma * stheta * std::cos(ph));
      SetZ(ma * ctheta);
    }
 
    void SetMag(DataType mag)
    {
      double factor = Mag();
      if (factor == 0) {
        B2WARNING(name() << "::SetMag: zero vector can't be stretched");
      } else {
        factor = mag / factor;
        SetX(X()*factor);
        SetY(Y()*factor);
        SetZ(Z()*factor);
      }
    }
 
    DataType Perp2() const { return X() * X() + Y() * Y(); }
    DataType Pt() const { return Perp(); }
    DataType Perp() const { return std::hypot((double)X(), (double)Y()); }
 
    void SetPerp(DataType r)
    {
      const double p = Perp();
      if (p != 0.0) {
        m_coordinates[0] *= r / p;
        m_coordinates[1] *= r / p;
      }
    }
 
    DataType Perp2(const B2Vector3<DataType>& axis)  const
    {
      const double tot = axis.Mag2();
      const double ss  = Dot(axis);
      double per = Mag2();
      if (tot > 0.0) per -= ss * ss / tot;
      if (per < 0)   per = 0;
      return per;
    }
 
    DataType Pt(const B2Vector3<DataType>& axis) const { return Perp(axis); }
    DataType Perp(const B2Vector3<DataType>& axis) const { return std::sqrt(Perp2(axis)); }
    DataType DeltaPhi(const B2Vector3<DataType>& v) const { return Mpi_pi(Phi() - v.Phi()); }
 
 
    static DataType Mpi_pi(DataType angle)
    {
      if (std::isnan(angle)) {
        B2ERROR(name() << "::Mpi_pi: function called with NaN");
        return angle;
      }
      angle = std::remainder(angle, 2 * M_PI);
      //for compatibility with ROOT we flip the sign for exactly pi
      if (angle == M_PI) angle = -M_PI;
      return angle;
    }
 
    DataType DeltaR(const B2Vector3<DataType>& v) const
    {
      const double deta = Eta() - v.Eta();
      const double dphi = DeltaPhi(v);
      return std::hypot(deta, dphi);
    }
 
    DataType DrEtaPhi(const B2Vector3<DataType>& v) const
    {
      return DeltaR(v);
    }
 
    void SetMagThetaPhi(DataType mag, DataType theta, DataType phi)
    {
      const double amag = std::abs(mag);
      const double sinTheta = std::sin((double)theta);
      m_coordinates[0] = amag * sinTheta * std::cos((double)phi);
      m_coordinates[1] = amag * sinTheta * std::sin((double)phi);
      m_coordinates[2] = amag * std::cos((double)theta);
    }
 
    B2Vector3<DataType> Unit() const
    {
      const double  tot = Mag2();
      B2Vector3<DataType> p(X(), Y(), Z());
      return tot > 0.0 ? p *= (1.0 / std::sqrt(tot)) : p;
    }
 
    B2Vector3<DataType> Orthogonal() const
    {
      const double xVal = std::abs((double)X());
      const double yVal = std::abs((double)Y());
      const double zVal = std::abs((double)Z());
      if (xVal < yVal) {
        return xVal < zVal ? B2Vector3<DataType>(0, Z(), -Y()) : B2Vector3<DataType>(Y(), -X(), 0);
      } else {
        return yVal < zVal ? B2Vector3<DataType>(-Z(), 0, X()) : B2Vector3<DataType>(Y(), -X(), 0);
      }
    }
 
    DataType Dot(const B2Vector3<DataType>& p) const
    {
      return X() * p.X() + Y() * p.Y() + Z() * p.Z();
    }
 
    B2Vector3<DataType> Cross(const B2Vector3<DataType>& p) const
    {
      return B2Vector3<DataType>(Y() * p.Z() - p.Y() * Z(), Z() * p.X() - p.Z() * X(), X() * p.Y() - p.X() * Y());
    }
 
    DataType Angle(const B2Vector3<DataType>& q) const
    {
      const double ptot2 = Mag2() * q.Mag2();
      if (ptot2 <= 0) {
        return 0.0;
      } else {
        double arg = Dot(q) / std::sqrt(ptot2);
        if (arg >  1.0) arg =  1.0;
        if (arg < -1.0) arg = -1.0;
        return std::acos(arg);
      }
    }
 
    DataType PseudoRapidity() const
    {
      const double cosTheta = CosTheta();
      if (std::abs(cosTheta) < 1) return -0.5 * std::log((1.0 - cosTheta) / (1.0 + cosTheta));
      if (Z()  == 0) return 0;
      //B2WARNING(name() << "::PseudoRapidity: transverse momentum = 0! return +/- 10e10");
      if (Z() > 0) return 10e10;
      else        return -10e10;
    }
 
 
    DataType Eta() const { return PseudoRapidity(); }
 
 
    void RotateX(DataType angle)
    {
      //rotate vector around X
      const double s = std::sin((double)angle);
      const double c = std::cos((double)angle);
      const double yOld = Y();
      m_coordinates[1] = c * yOld - s * Z();
      m_coordinates[2] = s * yOld + c * Z();
    }
 
 
    void RotateY(DataType angle)
    {
      //rotate vector around Y
      const double s = std::sin((double)angle);
      const double c = std::cos((double)angle);
      const double zOld = Z();
      m_coordinates[0] = s * zOld + c * X();
      m_coordinates[2] = c * zOld - s * X();
    }
 
 
    void RotateZ(DataType angle)
    {
      //rotate vector around Z
      const double s = std::sin((double)angle);
      const double c = std::cos((double)angle);
      const double xOld = X();
      m_coordinates[0] = c * xOld - s * Y();
      m_coordinates[1] = s * xOld + c * Y();
    }
 
    void RotateUz(const B2Vector3<DataType>& NewUzVector)
    {
      // NewUzVector must be normalized !
 
      const double u1 = NewUzVector.X();
      const double u2 = NewUzVector.Y();
      const double u3 = NewUzVector.Z();
      double up = u1 * u1 + u2 * u2;
 
      if (up) {
        up = std::sqrt(up);
        DataType px = X(),  py = Y(),  pz = Z();
        m_coordinates[0] = (u1 * u3 * px - u2 * py + u1 * up * pz) / up;
        m_coordinates[1] = (u2 * u3 * px + u1 * py + u2 * up * pz) / up;
        m_coordinates[2] = (u3 * u3 * px -      px + u3 * up * pz) / up;
      } else if (u3 < 0.) {
        m_coordinates[0] = -m_coordinates[0];
        m_coordinates[2] = -m_coordinates[2];
      }
    }
 
    void Rotate(DataType alpha, const B2Vector3<DataType>& v)
    {
      B2Vector3<DataType> n = v.Unit();
      *this = (n * (n.Dot(*this)) + cos(alpha) * ((n.Cross(*this)).Cross(n)) + sin(alpha) * (n.Cross(*this)));
    }
 
    void Abs()
    {
      m_coordinates[0] = std::abs(m_coordinates[0]);
      m_coordinates[1] = std::abs(m_coordinates[1]);
      m_coordinates[2] = std::abs(m_coordinates[2]);
    }
 
    void Sqrt()
    {
      Abs();
      m_coordinates[0] = std::sqrt(m_coordinates[0]);
      m_coordinates[1] = std::sqrt(m_coordinates[1]);
      m_coordinates[2] = std::sqrt(m_coordinates[2]);
    }
 
    DataType at(unsigned i) const;
    DataType x() const { return m_coordinates[0]; }
    DataType y() const { return m_coordinates[1]; }
    DataType z() const { return m_coordinates[2]; }
    DataType X() const { return x(); }
    DataType Y() const { return y(); }
    DataType Z() const { return z(); }
    DataType Px() const { return x(); }
    DataType Py() const { return y(); }
    DataType Pz() const { return z(); }
 
    void GetXYZ(Double_t* carray) const;
    void GetXYZ(Float_t* carray) const;
    void GetXYZ(TVector3* tVec) const;
    TVector3 GetTVector3() const;
 
    void SetX(DataType x) { m_coordinates[0] = x; }
    void SetY(DataType y) { m_coordinates[1] = y; }
    void SetZ(DataType z) { m_coordinates[2] = z; }
 
    void SetXYZ(DataType x, DataType y, DataType z)
    {
      SetX(x); SetY(y); SetZ(z);
    }
    void SetXYZ(const TVector3& tVec);
    void SetXYZ(const TVector3* tVec);
 
    static std::string name();
 
    std::string PrintString(unsigned precision = 4) const
    {
      return name() + " " + PrintStringXYZ(precision) + " " + PrintStringCyl(precision);
    }
 
    std::string PrintStringXYZ(unsigned precision = 4) const
    {
      std::ostringstream output;
      output  << "(x,y,z)=("
              << std::fixed << std::setprecision(precision)
              << X() << "," << Y() << "," << Z() << ")";
      return output.str();
    }
 
    std::string PrintStringCyl(unsigned precision = 4) const
    {
      std::ostringstream output;
      output  << "(rho, theta, phi)=("
              << std::fixed << std::setprecision(precision)
              << Mag() << "," << Theta() * 180. / M_PI << "," << Phi() * 180. / M_PI << ")";
      return output.str();
    }
 
    void Print()
    {
      //print vector parameters
      Print(PrintString().c_str());
    }
 
  };
 
  typedef B2Vector3<double> B2Vector3D;
 
  typedef B2Vector3<float> B2Vector3F;
 
  template <typename DataType>
  Bool_t operator == (const TVector3& a, const B2Vector3<DataType>& b)
  {
    return (a.X() == b.X() && a.Y() == b.Y() && a.Z() == b.Z());
  }
 
  template < typename DataType>
  Bool_t operator != (const TVector3& a, const B2Vector3<DataType>& b)
  {
    return !(a == b);
  }
 
  template < typename DataType>
  B2Vector3<DataType> operator * (DataType a, const B2Vector3<DataType>& p)
  {
    return B2Vector3<DataType>(a * p.X(), a * p.Y(), a * p.Z());
  }
 
  template < typename DataType>
  B2Vector3<DataType> operator + (const TVector3& a, const B2Vector3<DataType>& b)
  {
    return B2Vector3<DataType>(a.X() + b.X(), a.Y() + b.Y(), a.Z() + b.Z());
  }
 
  template < typename DataType>
  B2Vector3<DataType> operator - (const TVector3& a, const B2Vector3<DataType>& b)
  {
    return B2Vector3<DataType>(a.X() - b.X(), a.Y() - b.Y(), a.Z() - b.Z());
  }
 
  template < typename DataType>
  B2Vector3<DataType> operator + (const B2Vector3<DataType>& a, const TVector3& b)
  {
    return B2Vector3<DataType>(a.X() + b.X(), a.Y() + b.Y(), a.Z() + b.Z());
  }
 
  template < typename DataType>
  B2Vector3<DataType> operator - (const B2Vector3<DataType>& a, const TVector3& b)
  {
    return B2Vector3<DataType>(a.X() - b.X(), a.Y() - b.Y(), a.Z() - b.Z());
  }
 
 
  template< typename DataType >
  B2Vector3<DataType>& B2Vector3<DataType>::operator = (const B2Vector3<DataType>& b)
  {
    m_coordinates[0] = b.X();
    m_coordinates[1] = b.Y();
    m_coordinates[2] = b.Z();
    return *this;
  }
 
  template< typename DataType >
  B2Vector3<DataType>& B2Vector3<DataType>::operator = (const TVector3& b)
  {
    m_coordinates[0] = b.X();
    m_coordinates[1] = b.Y();
    m_coordinates[2] = b.Z();
    return *this;
  }
 
  template< typename DataType >
  B2Vector3<DataType>& B2Vector3<DataType>::operator += (const B2Vector3<DataType>& b)
  {
    m_coordinates[0] += b.X();
    m_coordinates[1] += b.Y();
    m_coordinates[2] += b.Z();
    return *this;
  }
 
 
  template< typename DataType >
  B2Vector3<DataType>& B2Vector3<DataType>::operator -= (const B2Vector3<DataType>& b)
  {
    m_coordinates[0] -= b.X();
    m_coordinates[1] -= b.Y();
    m_coordinates[2] -= b.Z();
    return *this;
  }
 
  template< typename DataType >
  B2Vector3<DataType>& B2Vector3<DataType>::operator *= (DataType a)
  {
    m_coordinates[0] *= a;
    m_coordinates[1] *= a;
    m_coordinates[2] *= a;
    return *this;
  }
 
  template< typename DataType >
  void B2Vector3<DataType>::SetXYZ(const TVector3& tVec)
  {
    m_coordinates[0] = static_cast<Double_t>(tVec.X());
    m_coordinates[1] = static_cast<Double_t>(tVec.Y());
    m_coordinates[2] = static_cast<Double_t>(tVec.Z());
  }
 
  template< typename DataType >
  void B2Vector3<DataType>::SetXYZ(const TVector3* tVec)
  {
    m_coordinates[0] = static_cast<Double_t>(tVec->X());
    m_coordinates[1] = static_cast<Double_t>(tVec->Y());
    m_coordinates[2] = static_cast<Double_t>(tVec->Z());
  }
 
  template< typename DataType >
  void B2Vector3<DataType>::GetXYZ(double* carray) const
  {
    carray[0] = X();
    carray[1] = Y();
    carray[2] = Z();
  }
 
  template< typename DataType >
  void B2Vector3<DataType>::GetXYZ(TVector3* tVec) const
  {
    tVec->SetXYZ(static_cast<Double_t>(X()),
                 static_cast<Double_t>(Y()),
                 static_cast<Double_t>(Z()));
  }
 
 
  template< typename DataType >
  TVector3 B2Vector3<DataType>::GetTVector3() const
  {
    return
      TVector3(
        static_cast<Double_t>(X()),
        static_cast<Double_t>(Y()),
        static_cast<Double_t>(Z())
      );
  }
 
 
  template < typename DataType>
  DataType B2Vector3<DataType>::at(unsigned i) const
  {
    switch (i) {
      case 0:
        return B2Vector3<DataType>::m_coordinates[0];
      case 1:
        return B2Vector3<DataType>::m_coordinates[1];
      case 2:
        return B2Vector3<DataType>::m_coordinates[2];
    }
    B2FATAL(this->name() << "::access operator: given index (i=" << i << ") is out of bounds!");
    return 0.;
  }
 
  template < typename DataType>
  std::string B2Vector3<DataType>::name()
  {
    return std::string("B2Vector3<") + typeid(DataType).name() + std::string(">");
  }
 
} // end namespace Belle2