DecayChainStack.h

Go to the documentation of this file.

#include <vector>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <stack>
 
#include "AmpGen/EventType.h"
#include "AmpGen/Particle.h"
#include "AmpGen/MsgService.h"
#include "AmpGen/ParticleProperties.h"
#include "AmpGen/SmallVector.h"
#include "AmpGen/Event.h"
#include <TRandom3.h>
 
namespace AmpGen {
  class DecayChainStackBase {
  public:
    DecayChainStackBase() = default;
    DecayChainStackBase(const Particle &particle);
    virtual ~DecayChainStackBase() {};
    virtual Event makeEvent(TRandom3 *rndm) const = 0;
    virtual double maxWeight() const = 0;
    virtual unsigned NP() const = 0;
    virtual bool operator==(const DecayChainStackBase &other) const = 0;
  };
 
  template <unsigned N> class DecayChainStack : public DecayChainStackBase {
    struct Node {
      enum Type { BW, Flat, Stable, QuasiStable };
      Type type = Type::Flat;
      unsigned l;
      unsigned r;
      std::pair<double, double> range;
      std::pair<double, double> phi_range;
      double bwMass = 0;
      double bwWidth = 0;
      SmallVector<unsigned, N> lfs;
      SmallVector<unsigned, N> rfs;
      Node(Type type = Type::Flat, const unsigned &l = 0, const unsigned &r = 0) : type(type), l(l), r(r) {};
      void print() const {
        std::string type_string = "";
        switch(type) {
        case Type::BW: type_string = "BW         "; break;
        case Type::Flat: type_string = "Flat       "; break;
        case Type::Stable: type_string = "Stable     "; break;
        case Type::QuasiStable: type_string = "QuasiStable"; break;
        };
        INFO(type_string << " decay indices: [" << l << ", " << r << "], range : [" << range.first << ", " << range.second
                         << "] "
                            ", ["
                         << vectorToString(lfs, ",") << "], [" << vectorToString(rfs, ",") << "]");
      }
      template <typename T> T operator()(const T &s) const {
        switch(type) {
        case Type::BW:
          return (bwMass * bwWidth)
                 / ((phi_range.second - phi_range.first) * ((s - bwMass * bwMass) * (s - bwMass * bwMass) + bwMass * bwMass * bwWidth * bwWidth));
 
        case Type::Flat: return 1 / (range.second - range.first);
        case Type::Stable: return 1;
        case Type::QuasiStable: return 1;
        };
        return 0;
      }
 
      double s_nom() const {
        switch(type) {
        case Type::BW: return bwMass * bwMass;
        case Type::Flat: return (range.second + range.first) * 0.5;
        case Type::Stable: return range.first;
        case Type::QuasiStable: return range.first;
        };
        return 0;
      }
      double operator()(TRandom3 *random) const {
        switch(type) {
        case Type::BW: {
          auto y = tan((phi_range.second - phi_range.first) * random->Rndm());
          auto mG = bwMass * bwWidth;
          auto s0 = bwMass * bwMass;
          return s0 + mG * (mG * y + range.first - s0) / (mG - y * (range.first - s0));
        }
        case Type::Flat: return (range.second - range.first) * random->Rndm() + range.first;
        case Type::Stable: return range.first;
        case Type::QuasiStable: return range.first;
        }
        return 0;
      }
      void set(const Particle *particle) {
        auto props = particle->props();
        bwWidth = props->width();
        bwMass = props->mass();
        INFO("Setting particle: " << particle->name() << " " << particle->isStable() << " " << particle->isQuasiStable() << " "
                                  << particle->lineshapeContains({"FormFactor", "None", "TD"}));
        if(particle->isStable())
          type = Type::Stable;
        else if(particle->isQuasiStable())
          type = Type::QuasiStable;
        else if(particle->lineshapeContains({"FormFactor", "None", "TD"}))
          type == Type::Flat;
        else {
          type = Type::BW;
          phi_range.first = atan((range.first - bwMass * bwMass) / (bwMass * bwWidth));
          phi_range.second = atan((range.second - bwMass * bwMass) / (bwMass * bwWidth));
        }
      }
      void setRange(const double &min, const double &max) {
        if(type == Type::BW or type == Type::Flat) {
          range = std::make_pair(min, max);
          if(type == Type::BW) {
            phi_range.first = atan((range.first - bwMass * bwMass) / (bwMass * bwWidth));
            phi_range.second = atan((range.second - bwMass * bwMass) / (bwMass * bwWidth));
          }
        }
      }
      bool operator==(const Node &other) const {
        return type == other.type && range == other.range && phi_range == other.phi_range && bwMass == other.bwMass && bwWidth == other.bwWidth && l == other.l
               && r == other.r && lfs == other.lfs && rfs == other.rfs;
      }
    }; 
    std::array<double, N> m_m0;
    std::array<Node, N - 1> m_nodes;
    double m_rhoMax = 1;
 
  public:
    DecayChainStack(const Particle &particle, bool aggressiveOptimisation) {
      auto fs = particle.getFinalStateParticles();
      double minMass = particle.mass();
      for(int i = 0; i != N; ++i) {
        m_m0[i] = fs[i]->mass();
        minMass -= m_m0[i];
      }
      std::stack<std::pair<std::shared_ptr<Particle>, unsigned>> toDo;
      toDo.emplace(std::make_shared<Particle>(particle), 0);
      unsigned counter = 0;
      while(toDo.size() != 0) {
        auto [current, index] = toDo.top();
        toDo.pop();
        if(current->daughters().size() != 2) {
          std::vector<Particle> particles;
          auto dp = current->daughters();
          for(unsigned int i = 1; i != dp.size(); ++i) particles.push_back(*dp[i]);
          current = std::make_shared<Particle>(current->name(), *dp[0], Particle("NonResS0", particles));
        }
        INFO(index << " " << *current);
        if(current->daughters().size() != 2) FATAL("Should be fixed..");
        auto fs = current->quasiStableTree().daughters();
        auto min_mass = std::accumulate(fs.begin(), fs.end(), 0., [](double acc, auto &p) { return acc + p->mass(); });
        auto G = current->isQuasiStable() ? 0 : current->props()->mass() * current->props()->width() * 10;
        auto s0 = current->mass() * current->mass();
        auto d1_index = current->daughter(0)->index();
        auto d2_index = current->daughter(1)->index();
        if(!current->daughter(0)->isStable()) {
          d1_index = ++counter + N;
          toDo.emplace(current->daughter(0), d1_index - N);
        }
        if(!current->daughter(1)->isStable()) {
          d2_index = ++counter + N;
          toDo.emplace(current->daughter(1), d2_index - N);
        }
        m_nodes[index] = Node(Node::Type::Flat, d1_index, d2_index);
        if(index == 0 or current->isQuasiStable())
          m_nodes[index].range = std::make_pair(s0 - G, s0 + G);
        else
          m_nodes[index].range = std::make_pair(pow(min_mass, 2), pow(minMass + min_mass, 2));
        m_nodes[index].set(&(*current));
      }
      if(aggressiveOptimisation) {
        WARNING("Using aggressive optimisation of phase space, can cause problems for relative normalisation of different decay topologies");
        for(auto &n : m_nodes) {
          if(n.l > N && n.r > N) {
            m_nodes[n.l - N].setRange(m_nodes[n.l - N].range.first,
                                      n.range.second + m_nodes[n.r - N].range.first - 2 * std::sqrt(n.range.second * m_nodes[n.r - N].range.first));
            m_nodes[n.r - N].setRange(m_nodes[n.r - N].range.first,
                                      n.range.second + m_nodes[n.l - N].range.first - 2 * std::sqrt(n.range.second * m_nodes[n.l - N].range.first));
          }
        }
      }
      auto concat = [](auto v1, const auto &v2) {
        v1.insert(v1.end(), v2.begin(), v2.end());
        return v1;
      };
      for(auto n = m_nodes.rbegin(); n != m_nodes.rend(); ++n) {
        auto s1 = n->l < N ? m_m0[n->l] * m_m0[n->l] : m_nodes[n->l - N].range.first;
        auto s2 = n->r < N ? m_m0[n->r] * m_m0[n->r] : m_nodes[n->r - N].range.first;
 
        auto rho = rho_2(n->range.second, s1, s2);
        // INFO(m_nodes.size() - ( n-m_nodes.rbegin()) + N -1  << " " << rho << " " << n->l << " " << n->r << " " << n->range.first << " " << n->range.second
        // << " " << n->type <<
        //     " s(max): " << n->range.second <<
        //     " s1(min): " << s1  << " s2(min): " << s2 );
        m_rhoMax *= rho;
        n->lfs = n->l < N ? SmallVector<unsigned, N>{n->l} : concat(m_nodes[n->l - N].lfs, m_nodes[n->l - N].rfs);
        n->rfs = n->r < N ? SmallVector<unsigned, N>{n->r} : concat(m_nodes[n->r - N].lfs, m_nodes[n->r - N].rfs);
      }
      for(auto &node : m_nodes) node.print();
      if(m_rhoMax <= 0) ERROR("RhoMax^2 < 0! " << m_rhoMax);
      m_rhoMax = std::sqrt(m_rhoMax);
    }
    virtual double maxWeight() const { return m_rhoMax; }
    virtual unsigned NP() const { return N; }
    virtual bool operator==(const DecayChainStackBase &other) const {
      auto op = dynamic_cast<const DecayChainStack<N> *>(&other);
      return op && this->m_m0 == op->m_m0 && this->m_nodes == op->m_nodes;
    }
    /*
       double genPdf(const double* event) const
       {
       double g = 1;
       for( const auto& node : m_nodes ){
    //INFO( "s(" << vectorToString( node.lfs, ",") << "," << vectorToString(node.rfs,",") << ") = " << s(event, node.lfs, node.rfs) );
    g *= node(s(event, node.lfs, node.rfs));
    }
    return g;
    }
    */
    template <typename T> T genPdf(const T *event) const {
      T g = 1;
      for(const auto &node : m_nodes) {
        // INFO( "s(" << vectorToString( node.lfs, ",") << "," << vectorToString(node.rfs,",") << ") = " << s(event, node.lfs, node.rfs) );
        g *= node(s(event, node.lfs, node.rfs));
      }
      return g;
    }
    template <typename T> T s(const T *event, const SmallVector<unsigned, N> &lfs, const SmallVector<unsigned, N> &rfs) const {
      std::array<T, 4> P = {0.};
      for(const auto &it : lfs) {
        P[0] += event[4 * it + 0];
        P[1] += event[4 * it + 1];
        P[2] += event[4 * it + 2];
        P[3] += event[4 * it + 3];
      }
      for(const auto &it : rfs) {
        P[0] += event[4 * it + 0];
        P[1] += event[4 * it + 1];
        P[2] += event[4 * it + 2];
        P[3] += event[4 * it + 3];
      }
      return P[3] * P[3] - P[0] * P[0] - P[1] * P[1] - P[2] * P[2];
    }
 
    double rho_2(const double &s0, const double &s1, const double &s2) const { return 1 - 2 * (s1 + s2) / s0 + (s1 - s2) * (s1 - s2) / (s0 * s0); }
    void boost(double *output, double nx, double ny, double nz, double g, double vg) const {
      double nv = output[0] * nx + output[1] * ny + output[2] * nz;
      output[0] += ((g - 1) * nv + vg * output[3]) * nx;
      output[1] += ((g - 1) * nv + vg * output[3]) * ny;
      output[2] += ((g - 1) * nv + vg * output[3]) * nz;
      output[3] = g * output[3] + vg * nv;
    }
    void set(double *output, double px, double py, double pz, double E) const {
      output[0] = px;
      output[1] = py;
      output[2] = pz;
      output[3] = E;
    }
 
    std::pair<double, std::array<double, 2 * N - 1>> proposal(TRandom3 *rndm) const {
      std::pair<double, std::array<double, 2 * N - 1>> rt;
      auto &[weight, state] = rt;
      for(unsigned int i = 0; i != N; ++i) state[i] = m_m0[i] * m_m0[i];
      weight = 1;
      for(int i = m_nodes.size() - 1; i >= 0; --i) {
        state[i + N] = m_nodes[i](rndm);
        const auto &s = state[i + N];
        const auto &sl = state[m_nodes[i].l];
        const auto &sr = state[m_nodes[i].r];
        double v = rho_2(s, sl, sr) * (s > sl + sr + 2 * std::sqrt(sl * sr));
        weight = v < 0 ? 0 : v * weight;
      }
      weight = std::sqrt(weight);
      return rt;
    }
    void debug(TRandom3 *rndm) const {
      std::pair<double, std::array<double, 2 * N - 1>> rt;
      auto &[weight, state] = rt;
      for(unsigned int i = 0; i != N; ++i) state[i] = m_m0[i] * m_m0[i];
      weight = 1;
      for(int i = m_nodes.size() - 1; i >= 0; --i) {
        state[i + N] = m_nodes[i](rndm);
        const auto &s = state[i + N];
        const auto &sl = state[m_nodes[i].l];
        const auto &sr = state[m_nodes[i].r];
        double v = rho_2(s, sl, sr) * (s > sl + sr + 2 * std::sqrt(sl * sr));
        INFO("Node[" << i << "] w = " << v << "s[" << m_nodes[i].l << " " << m_nodes[i].r << "] = " << s << " [ sl=  " << sl << " sr= " << sr << "]");
        weight = v < 0 ? 0 : v * weight;
      }
      weight = std::sqrt(weight);
    }
 
    void fill_from_state(double *event, const std::array<double, 2 * N - 1> &state, TRandom3 *rndm) const {
      for(int i = 0; i != N; ++i) event[4 * i + 3] = m_m0[i];
      for(auto n = m_nodes.rbegin(); n != m_nodes.rend(); ++n) {
        const auto &sl = state[n->l];
        const auto &sr = state[n->r];
        const auto &s = state[&(*n) - &(*m_nodes.begin()) + N];
        const auto nz = 2 * rndm->Rndm() - 1;
        const auto phi = 2 * M_PI * rndm->Rndm();
        const auto p = 0.5 * std::sqrt(s * rho_2(s, sl, sr));
        const auto sZ = std::sqrt(1 - nz * nz);
        const auto nx = sZ * std::cos(phi);
        const auto ny = sZ * std::sin(phi);
        const auto gl = std::sqrt(1 + p * p / sl);
        const auto vgl = p / std::sqrt(sl);
        const auto gr = std::sqrt(1 + p * p / sr);
        const auto vgr = p / std::sqrt(sr);
 
        if(n->lfs.size == 1)
          set(event + 4 * n->lfs[0], p * nx, p * ny, p * nz, std::sqrt(sl + p * p));
        else
          for(const auto &l : n->lfs) boost(event + 4 * l, nx, ny, nz, gl, vgl);
 
        if(n->rfs.size == 1)
          set(event + 4 * n->rfs[0], -p * nx, -p * ny, -p * nz, std::sqrt(sr + p * p));
        else
          for(const auto &l : n->rfs) boost(event + 4 * l, -nx, -ny, -nz, gr, vgr);
      }
    }
 
    virtual AmpGen::Event makeEvent(TRandom3 *rndm) const {
      std::array<double, 2 * N - 1> state = {0};
      double rho_v = 1;
      do {
        auto rt = proposal(rndm);
        rho_v = std::get<0>(rt);
        state = std::get<1>(rt);
      } while(rho_v < m_rhoMax * rndm->Uniform());
      AmpGen::Event event(N * 4);
      fill_from_state(event, state, rndm);
      return event;
    }
  };
 
  DecayChainStackBase *make_decay_chain_stack(const Particle &particle, bool aggressiveOptimisation = false) {
    auto fs = particle.getFinalStateParticles();
    switch(fs.size()) {
    case(1): return new DecayChainStack<1>(particle, aggressiveOptimisation);
    case(2): return new DecayChainStack<2>(particle, aggressiveOptimisation);
    case(3): return new DecayChainStack<3>(particle, aggressiveOptimisation);
    case(4): return new DecayChainStack<4>(particle, aggressiveOptimisation);
    case(5): return new DecayChainStack<5>(particle, aggressiveOptimisation);
    case(6): return new DecayChainStack<6>(particle, aggressiveOptimisation);
    case(7): return new DecayChainStack<7>(particle, aggressiveOptimisation);
    case(8): return new DecayChainStack<8>(particle, aggressiveOptimisation);
    case(9): return new DecayChainStack<9>(particle, aggressiveOptimisation);
    case(10): return new DecayChainStack<10>(particle, aggressiveOptimisation);
    }
    return nullptr;
  }
 
}