LorentzVectorBase.hh

Go to the documentation of this file.

#pragma once

#include <climits>

#include "boca/math/Vector3.hh"

namespace boca
{

enum class LorentzDim
{
    x,
    y,
    z,
    t,
    last
};

std::string Name(LorentzDim lorentz_dim);

std::vector<LorentzDim> LorentzDimensions();

template<typename Value_>
class LorentzVectorBase : boost::totally_ordered<LorentzVectorBase<Value_>>
            , boost::additive<LorentzVectorBase<Value_>>
            , Vector3<Value_>
{

    using ValueSquare = boca::ValueSquare<Value_>;

    template<typename Value_2>
    using ValueProduct = boca::ValueProduct<Value_, Value_2>;

    template<typename Value_2>
    using ValueQuotient = ValueQuotient<Value_, Value_2>;

    template <typename> struct IsQuantity : std::false_type {};
    template <typename T> struct IsQuantity<boost::units::quantity<T>> : std::true_type {};

    template<typename Value_2>
    using OnlyIfNotOrSameQuantity = typename std::enable_if < !IsQuantity<Value_2>::value || std::is_same<Value_, Value_2>::value >::type;

protected:

    template<typename Value_2>
    using OnlyIfNotQuantity = typename std::enable_if < !IsQuantity<Value_2>::value >::type;

public:

    constexpr LorentzVectorBase() :
        Vector3<Value_>(),
        scalar_(Value_(0)) {}

    constexpr LorentzVectorBase(Value_ x, Value_ y, Value_ z, Value_ t) :
        Vector3<Value_>(x, y, z),
        scalar_(t) {}

    constexpr LorentzVectorBase(Vector3<Value_> const &vector3, Value_ const &t) :
        Vector3<Value_>(vector3),
        scalar_(t) {}

    template<typename Value_2>
    constexpr LorentzVectorBase(LorentzVectorBase<Value_2> const &lorentz_vector) :
        Vector3<Value_>(lorentz_vector.Spatial()),
        scalar_(lorentz_vector.Scalar()) {}

    void SetVectT(Vector3<Value_> const &spatial, Value_ t)
    {
        Spatial() = spatial;
        Scalar() = t;
    }

    using Vector3<Value_>::SetTheta;

    using Vector3<Value_>::SetPhi;

    using Vector3<Value_>::SetPerp;

    using Vector3<Value_>::SetPerpEtaPhi;

    void SetRho(Value_ rho)
    {
        Spatial().SetMag(rho);
    }

    void SetMag(Value_ magnitude)
    {
        Scalar() = magnitude >= Value_(0)  ?  sqrt(sqr(Spatial()) + sqr(magnitude)) : sqrt(std::max((sqr(Spatial()) - sqr(magnitude)), ValueSquare(0)));
    }

    void SetVectMag(Vector3<Value_> const &spatial, Value_ magnitude)
    {
        Spatial() = spatial;
        SetMag(magnitude);
    }

    void SetPerpEtaPhiMag(Value_ const &perp, boca::Angle const &eta, boca::Angle const &phi, Value_ const &mag)
    {
        SetPerpEtaPhi(perp, eta,  phi);
        SetMag(mag);
    }

    void SetPerpEtaPhiT(Value_ const &perp, boca::Angle const &eta, boca::Angle const &phi, Value_ const &t)
    {
        SetPerpEtaPhi(perp, eta, phi);
        scalar_ = t;
    }


    constexpr Vector3<Value_> const& Spatial() const
    {
        return *this;
    }

    Vector3<Value_> &Spatial()
    {
        return *this;
    }

    using Vector3<Value_>::X;

    using Vector3<Value_>::Y;

    using Vector3<Value_>::Z;

    Value_ &Scalar()
    {
        return scalar_;
    }

    Value_ &T()
    {
        return scalar_;
    }

    constexpr Value_ const& Scalar() const
    {
        return scalar_;
    }

    constexpr Value_ const& T() const
    {
        return scalar_;
    }


    using Vector3<Value_>::Phi;

    using Vector3<Value_>::Theta;

    using Vector3<Value_>::CosTheta;

    using Vector3<Value_>::SinTheta2;

    using Vector3<Value_>::SinTheta;

//     using Vector3<Value_>::Angle;
//
//     /**
//       * @brief Angle wrt. another vector.
//       */
//     boca::Angle Angle(LorentzVectorBase const &vector) const
//     {
//         return Angle(vector.Spatial());
//     }

    using Vector3<Value_>::Eta;

    using Vector3<Value_>::PseudoRapidity;

    boca::Angle Rapidity() const
    {
        return 0.5 * units::log(static_cast<double>(Plus() / Minus()));
    }

    boca::Angle Rap() const
    {
        return Rapidity();
    }

    boca::Angle Rapidity(Vector3<double> const &vector) const
    {
        auto const mag_2 = vector.Mag2();
        if (mag_2 <= 0) return 0_rad;
        auto const vdotu = Spatial().Dot(vector) / sqrt(mag_2);
        if (vdotu == Value_(0)) return 0_rad;
//         if (Scalar() <= Value_(0)) std::cout << "Tried to take rapidity of negative-energy Lorentz vector" << '\n';
//         auto pt = sqrt(units::max(sqr(Scalar() * std::numeric_limits<double>::epsilon()), Perp2(vector) + Mag2()));
        auto const mt = sqrt(Perp2(vector) + Mag2());
        if (mt == Value_(0)) return 0_rad;
        auto const rap = units::log(static_cast<double>((Scalar() + abs(Z())) / mt));
        return Z() > Value_(0) ? rap : -rap;
    }

    using Vector3<Value_>::DeltaPhiTo;

    using Vector3<Value_>::DeltaEtaTo;

    template <typename Value_2_>
    boca::Angle DeltaRapTo(LorentzVectorBase<Value_2_> const &vector) const
    {
        return Rap() - vector.Rap();
    }

    template <typename Value_2_>
    boca::Angle DeltaRTo(LorentzVectorBase<Value_2_> const &vector) const
    {
        return sqrt(sqr(DeltaRapTo(vector)) + sqr(DeltaPhiTo(vector)));
    }


    using Vector3<Value_>::Perp2;

    using Vector3<Value_>::Perp;

    constexpr ValueSquare Rho2() const
    {
        return Spatial().Mag2();
    }

    constexpr Value_ Rho() const
    {
        return Spatial().Mag();
    }

    constexpr ValueSquare Mag2() const
    {
        return sqr(Scalar()) - Rho2();
    }

    constexpr Value_ Mag() const
    {
        auto const mag2 = Mag2();
        return mag2 < ValueSquare(0) ? -sqrt(-mag2) : sqrt(mag2);
    }

    constexpr ValueSquare MagT2() const
    {
        return sqr(Scalar()) - sqr(Z());
    }

    constexpr Value_ MagT() const
    {
        auto const mt2 = MagT2();
        return mt2 < ValueSquare(0) ? -sqrt(-mt2) : sqrt(mt2);
    }

    constexpr ValueSquare ScalarT2() const
    {
        return sqr(Scalar()) * SinTheta2();
    }

    constexpr Value_ ScalarT() const
    {
        return Scalar() < Value_(0) ? -sqrt(ScalarT2()) : sqrt(ScalarT2());
    }

    template<typename Value,  typename = OnlyIfNotQuantity<Value>>
    constexpr ValueSquare ScalarT2(Vector3<Value> const &vector) const
    {
        auto const pt2 = Spatial().Perp2(vector);
        return pt2 == ValueSquare(0) ? ValueSquare(0) : sqr(Scalar()) * pt2 / (pt2 + sqr(Spatial().Dot(vector.Unit())));
    }

    template<typename Value,  typename = OnlyIfNotQuantity<Value>>
    constexpr Value_ ScalarT(Vector3<Value> const &vector) const
    {
        return Scalar() < Value_(0) ? -sqrt(Et2(vector)) : sqrt(Et2(vector));
    }

    constexpr double Beta() const
    {
        return Rho() / Scalar();
    }

    constexpr double Gamma() const
    {
        return 1. / std::sqrt(1. - sqr(Beta()));
    }

    constexpr Value_ Plus() const
    {
        return Scalar() + Z();
    }

    constexpr Value_ Minus() const
    {
        return Scalar() - Z();
    }

    constexpr Vector3<double> BoostVector() const
    {
        return Scalar() == Value_(0) ? Vector3<double> {} :
               Vector3<double> {Spatial() / Scalar()};
    }

    using Vector3<Value_>::Transversal;

    using Vector3<Value_>::EtaPhiVector;

    constexpr Vector3<double> BoostIntoRestFrame() const
    {
        return -BoostVector();
    }

    constexpr LorentzVectorBase<Value_> Boosted(Vector3<double> const &boost) const
    {
        auto lorentz_vector = *this;
        return lorentz_vector.Boost(boost);
    }

    LorentzVectorBase<Value_> &Boost(Vector3<double> const &boost)
    {
        auto const mag_2 = boost.Mag2();
        auto const gamma = 1 / sqrt(1 - mag_2);
        auto const gamma2 = mag_2 > 0 ? (gamma - 1) / mag_2 : 0;
        auto const bp = boost * Spatial();
        auto const scaled = gamma * Scalar();

        Spatial() += boost * (gamma2 * bp + scaled);
        Scalar() = scaled + gamma * bp;
        return *this;
    }

    using Vector3<Value_>::Rotate;

    using Vector3<Value_>::RotateUz;


    template <typename Value_2>
    constexpr LorentzVectorBase<ValueProduct<Value_2>> Scale(Value_2 const &scalar) const
    {
        return {Spatial() *scalar, Scalar() *scalar};
    }

    template <typename Value_2>
    constexpr ValueProduct<Value_2> Dot(LorentzVectorBase<Value_2> const &lorentz_vector) const
    {
        return Scalar() * lorentz_vector.Scalar() - Spatial().Dot(lorentz_vector.Spatial());
    }

    template <typename Value_2>
    constexpr ValueProduct<Value_2> Euclidean(LorentzVectorBase<Value_2> const &lorentz_vector) const
    {
        return Scalar() * lorentz_vector.Scalar() + Spatial().Dot(lorentz_vector.Spatial());
    }


    constexpr bool operator<(LorentzVectorBase const &lorentz_vector) const
    {
        return Mag2() < lorentz_vector.Mag2();
    }

    constexpr bool operator==(LorentzVectorBase const &lorentz_vector) const
    {
        return Spatial() == lorentz_vector.Spatial() && Scalar() == lorentz_vector.Scalar();
    }

    template <typename Value_2, typename = OnlyIfNotOrSameQuantity<Value_2>>
    LorentzVectorBase &operator+=(LorentzVectorBase<Value_2> const &lorentz_vector)
    {
        Spatial() += lorentz_vector.Spatial();
        scalar_ += lorentz_vector.Scalar();
        return *this;
    }

    template <typename Value_2, typename = OnlyIfNotOrSameQuantity<Value_2>>
    LorentzVectorBase &operator-=(LorentzVectorBase<Value_2> const &lorentz_vector)
    {
        Spatial() -= lorentz_vector.Spatial();
        scalar_ -= lorentz_vector.Scalar();
        return *this;
    }

    template <typename Value_2, typename = OnlyIfNotQuantity<Value_2>>
    LorentzVectorBase &operator*=(Value_2 scalar)
    {
        Spatial() *= scalar;
        Scalar() *= scalar;
        return *this;
    }

    template <typename Value_2, typename = OnlyIfNotQuantity<Value_2>>
    LorentzVectorBase &operator/=(Value_2 scalar)
    {
        Spatial() /= scalar;
        Scalar() /= scalar;
        return *this;
    }

    template <typename Value_2>
    constexpr LorentzVectorBase<ValueQuotient<Value_2>> operator/(Value_2 const &scalar) const
    {
        return Scale(1. / scalar);
    }

    constexpr LorentzVectorBase operator-() const
    {
        return { -Spatial(), -Scalar()};
    }

    Value_ const &operator[](LorentzDim lorentz_dim) const
    {
        switch (lorentz_dim) {
        case LorentzDim::x :
            return Spatial()[Dim3::x];
        case LorentzDim::y :
            return Spatial()[Dim3::y];
        case LorentzDim::z :
            return Spatial()[Dim3::z];
        case LorentzDim::t :
            return scalar_;
        default:
            Default("LorentzVector", Name(lorentz_dim));
            return Spatial()[Dim3::x];
        }
    }

    Value_ &operator[](LorentzDim lorentz_dim)
    {
        return const_cast<Value_ &>(static_cast<LorentzVectorBase<Value_> const *>(this)[lorentz_dim]);
    }

    friend auto &operator<<(std::ostream &stream, LorentzVectorBase<Value_> const &lorentz_vector)
    {
        stream << Stream(lorentz_vector.Scalar()) << lorentz_vector.Spatial();
        return stream;
    }


    using Dimension = LorentzDim;

    constexpr ConstIterator<Vector3, Value_> begin() const
    {
        return {this, LorentzDim::x};
    }

    constexpr ConstIterator<Vector3, Value_> end() const
    {
        return {this, LorentzDim::last};
    }

    Iterator<Vector3, Value_> begin()
    {
        return {this, LorentzDim::x};
    }

    Iterator<Vector3, Value_> end()
    {
        return {this, LorentzDim::last};
    }


protected:

    Value_ scalar_;

};

template<typename Value>
using OnlyIfNotLorentzVectorBase = typename std::enable_if<std::is_base_of<LorentzVectorBase<Value>, Value>::value>::type;

template <class Value, class Value_2>
constexpr auto operator*(LorentzVectorBase<Value> const &lorentz_vector_1, LorentzVectorBase<Value_2> const &lorentz_vector_2)
{
    return lorentz_vector_1.Dot(lorentz_vector_2);
}

template < class Value, class Value_2, typename = OnlyIfNotLorentzVectorBase<Value_2> >
constexpr auto operator*(LorentzVectorBase<Value> const &lorentz_vector, Value_2 scalar)
{
    return lorentz_vector.Scale(scalar);
}

template < class Value, class Value_2, typename = OnlyIfNotLorentzVectorBase<Value> >
constexpr auto operator*(Value scalar, LorentzVectorBase<Value_2> const &lorentz_vector)
{
    return lorentz_vector.Scale(scalar);
}

}

namespace boost{

template<typename Value_>
struct range_const_iterator< boca::LorentzVectorBase<Value_> > {
    typedef typename boca::ConstIterator<boca::LorentzVectorBase, Value_> type;
};

template<typename Value_>
struct range_mutable_iterator< boca::LorentzVectorBase<Value_> > {
    typedef typename boca::Iterator<boca::LorentzVectorBase, Value_> type;
};

}