Rivet analyses

// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/HeavyHadrons.hh"

namespace Rivet {


  class MC_HFDECAYS : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(MC_HFDECAYS);

    const string whoDis(const int pid) const {
      switch (pid) {
        case PID::B0: return "B0";
        case PID::BPLUS: return "BPLUS";
        case PID::B0S: return "B0S";
        case PID::LAMBDAB: return "LAMBDAB";
        case PID::D0: return "D0";
        case PID::DPLUS: return "DPLUS";
        case PID::DSPLUS: return "DSPLUS";
        case PID::LAMBDACPLUS: return "LAMBDACPLUS";
        default: return "";
      }
    }

    /// Book histograms and initialise projections before the run
    void init() {

      declare(HeavyHadrons(), "HA");

      FastJets jetpro(FinalState(), JetAlg::ANTIKT, 0.4, JetMuons::DECAY, JetInvisibles::DECAY);
      declare(jetpro, "Jets");

      // bar charts
      book(_h["bar_b_jet_width"], "width_B_jet", 7, 0., 0.3);
      book(_h["bar_c_jet_width"], "width_C_jet", 7, 0., 0.3);
      // profiles
      book(_p["b_jet_rho"], "avg_rho_B_jet", 10, 0., 0.4);
      book(_p["c_jet_rho"], "avg_rho_C_jet", 10, 0., 0.4);
      book(_p["b_jet_psi"], "avg_psi_B_jet", 10, 0., 0.4);
      book(_p["c_jet_psi"], "avg_psi_C_jet", 10, 0., 0.4);
      // histograms
      book(_h["b_frac"], "prod_frac_B", 11, 0.5, 11.5);
      book(_h["c_frac"], "prod_frac_C", 7, 0.5, 7.5);
      book(_h["b_jet_frag"], "frag_B_jet", 50, 0., 1.1);
      book(_h["c_jet_frag"], "frag_C_jet", 50, 0., 1.1);
      book(_h["b_jet_pT"], "pT_B_jet", 25, 25., 1025.);
      book(_h["c_jet_pT"], "pT_C_jet", 25, 25., 1025.);
      book(_h["b_jet_pThad"], "pThad_B_jet", 10, 0., 200.);
      book(_h["c_jet_pThad"], "pThad_C_jet", 10, 0., 200.);
      book(_h["b_jet_high_pT"], "high_pT_frag_B_jet", 46, 0., 1.);
      book(_h["c_jet_high_pT"], "high_pT_frag_C_jet", 46, 0., 1.);
      book(_h["b_jet_high_pT_1H"], "high_pT_frag_B_jet_1H", 46, 0., 1.);
      book(_h["c_jet_high_pT_1H"], "high_pT_frag_C_jet_1H", 46, 0., 1.);


      book(_h["ch_B0"], "B0_charged_mult", 30, 0.5, 30.5);
      book(_h["ch_BPLUS"], "BPLUS_charged_mult", 30, 0.5, 30.5);
      book(_h["ch_B0S"], "B0S_charged_mult", 30, 0.5, 30.5);
      book(_h["ch_LAMBDAB"], "LAMBDAB_charged_mult", 30, 0.5, 30.5);
      book(_h["ch_D0"], "D0_charged_mult", 16, 0.5, 16.5);
      book(_h["ch_DPLUS"], "DPLUS_charged_multh", 16, 0.5, 16.5);
      book(_h["ch_DSPLUS"], "DSPLUS_charged_mult", 16, 0.5, 16.5);
      book(_h["ch_LAMBDACPLUS"], "LAMBDACPLUS_charged_mult", 16, 0.5, 16.5);

      book(_h["st_B0"], "B0_stable_mult", 40, 0.5, 40.5);
      book(_h["st_BPLUS"], "BPLUS_stable_mult", 40, 0.5, 40.5);
      book(_h["st_B0S"], "B0S_stable_mult", 40, 0.5, 40.5);
      book(_h["st_LAMBDAB"], "LAMBDAB_stable_mult", 40, 0.5, 40.5);
      book(_h["st_D0"], "D0_stable_mult", 20, 0.5, 20.5);
      book(_h["st_DPLUS"], "DPLUS_stable_mult", 20, 0.5, 20.5);
      book(_h["st_DSPLUS"], "DSPLUS_stable_mult", 20, 0.5, 20.5);
      book(_h["st_LAMBDACPLUS"], "LAMBDACPLUS_stable_mult", 20, 0.5, 20.5);

      book(_h["pt_B0"], "B0_pT", 10, 25., 425.);
      book(_h["pt_BPLUS"], "BPLUS_pT", 10, 25., 425.);
      book(_h["pt_B0S"], "B0S_pT", 10, 25., 425.);
      book(_h["pt_LAMBDAB"], "LAMBDAB_pT", 10, 25., 425.);
      book(_h["pt_D0"], "D0_pT", 10, 25., 425.);
      book(_h["pt_DPLUS"], "DPLUS_pT", 10, 25., 425.);
      book(_h["pt_DSPLUS"], "DSPLUS_pT", 10, 25., 425.);
      book(_h["pt_LAMBDACPLUS"], "LAMBDACPLUS_pT", 10, 25., 425.);

      book(_h["b_jet_ch_mult"], "charged_mult_B_jets", 40, 0.5, 40.5);
      book(_h["c_jet_ch_mult"], "charged_mult_C_jets", 40, 0.5, 40.5);

      book(_h["b_jet_l_pTrel"], "lepton_pTrel_B_jets", 8, 0., 15.);
      book(_h["c_jet_l_pTrel"], "lepton_pTrel_C_jets", 8, 0., 10.);

      book(_h["b_jet_l_pT"], "lepton_pT_B_jets", 10, 0., 100.);
      book(_h["c_jet_l_pT"], "lepton_pT_C_jets", 10, 0., 100.);

      // double-differentials
      vector<double> groupEdges = ptaxis.edges(); // includes +/- inf
      groupEdges.erase(groupEdges.begin());
      groupEdges.pop_back(); // removes +/- inf
      book(_g["avg_B_jet_ch_mult"], groupEdges);
      book(_g["avg_C_jet_ch_mult"], groupEdges);
      book(_g["avg_B_jet_l_pTrel"], groupEdges);
      book(_g["avg_C_jet_l_pTrel"], groupEdges);
      for (size_t i = 0; i < ptaxis.numBins(); ++i) {
        const std::pair<double, double> ptbin = std::make_pair(ptaxis.min(i + 1), ptaxis.max(i + 1));
        const string suff = to_string(int(ptbin.first)) + "_" + to_string(int(ptbin.second));
        book(_p["avg_B_jet_rho_" + to_string(i)], "avg_B_jet_rho_" + suff, 10, 0., 0.4);
        book(_p["avg_C_jet_rho_" + to_string(i)], "avg_C_jet_rho_" + suff, 10, 0., 0.4);

        book(_g["avg_B_jet_ch_mult"]->bin(i + 1), "avg_B_jet_ch_mult_" + suff, 40, 0.5, 40.5);
        book(_g["avg_C_jet_ch_mult"]->bin(i + 1), "avg_C_jet_ch_mult_" + suff, 40, 0.5, 40.5);
        if (i == 0) {
          book(_g["avg_B_jet_l_pTrel"]->bin(i + 1), "avg_B_jet_l_pTrel_" + suff, 2, 0., 4.);
          book(_g["avg_C_jet_l_pTrel"]->bin(i + 1), "avg_C_jet_l_pTrel_" + suff, 2, 0., 3.);
        }
        else if (i <= 2) {
          book(_g["avg_B_jet_l_pTrel"]->bin(i + 1), "avg_B_jet_l_pTrel_" + suff, 4, 0., 8.);
          book(_g["avg_C_jet_l_pTrel"]->bin(i + 1), "avg_C_jet_l_pTrel_" + suff, 5, 0., 6.);
        }
        else {
          book(_g["avg_B_jet_l_pTrel"]->bin(i + 1), "avg_B_jet_l_pTrel_" + suff, 8, 0., 15.);
          book(_g["avg_C_jet_l_pTrel"]->bin(i + 1), "avg_C_jet_l_pTrel_" + suff, 8, 0., 10.);
        }
      }
    }

    double p_annulus(const Jet& jet, const double a, const double b) const {
      // calculate the total momentum inside an annulus with a <= R < b
      return sum(select(jet.particles(),
                        [&](const Particle& p) {
                          const double dr = deltaR(p, jet);
                          return (dr < b && dr >= a);
                        }),
                 Kin::pT, 0.)
          / GeV;
    }

    void count_mult(const Particle& p) {
      unsigned int nst = p.stableDescendants().size() + 0.5;
      unsigned int nch = p.stableDescendants(Cuts::charge != 0).size() + 0.5;
      _h["st_" + whoDis(p.abspid())]->fill(nst);
      _h["ch_" + whoDis(p.abspid())]->fill(nch);
    }

    double pTrel(const Jet& jet, const Particle& p) const {
      return (p.p3().cross(jet.p3())).mod() / (jet.p3().mod());
    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {

      const HeavyHadrons& ha = apply<HeavyHadrons>(event, "HA");

      if (ha.bHadrons().empty() && ha.cHadrons().empty()) vetoEvent;

      for (const Particle& p : ha.bHadrons()) {
        const string name = "pt_" + whoDis(p.abspid());
        if (p.abspid() == PID::B0) { //Take into consideration both particles and anti-particles with abs()
          count_mult(p);
          _h["b_frac"]->fill(1);
          _h[name]->fill(p.pT() / GeV);
        }
        else if (p.abspid() == PID::BPLUS) {
          count_mult(p);
          _h["b_frac"]->fill(2);
          _h[name]->fill(p.pT() / GeV);
        }
        else if (p.abspid() == PID::B0S) {
          count_mult(p);
          _h["b_frac"]->fill(3);
          _h[name]->fill(p.pT() / GeV);
        }
        else if (p.abspid() == PID::BCPLUS)
          _h["b_frac"]->fill(4);
        else if (p.abspid() == PID::LAMBDAB) {
          count_mult(p);
          _h["b_frac"]->fill(5);
          _h[name]->fill(p.pT() / GeV);
        }
        else if (p.abspid() == PID::XIBMINUS)
          _h["b_frac"]->fill(6);
        else if (p.abspid() == PID::XI0B)
          _h["b_frac"]->fill(7);
        else if (p.abspid() == PID::OMEGABMINUS)
          _h["b_frac"]->fill(8);
        else if (p.abspid() == PID::SIGMABMINUS)
          _h["b_frac"]->fill(9);
        else if (p.abspid() == PID::SIGMAB)
          _h["b_frac"]->fill(10);
        else if (p.abspid() == PID::SIGMABPLUS)
          _h["b_frac"]->fill(11);
      }

      for (const Particle& p : ha.cHadrons()) {
        if (!p.fromBottom()) { //take into account only c-hadrons that don't come from a b-hadron decay
                               //Avoid double-counting b-hadron fractions
          const string name = "pt_" + whoDis(p.abspid());
          if (p.abspid() == PID::DPLUS) {
            count_mult(p);
            _h["c_frac"]->fill(1);
            _h[name]->fill(p.pT() / GeV);
          }
          else if (p.abspid() == PID::D0) {
            count_mult(p);
            _h["c_frac"]->fill(2);
            _h[name]->fill(p.pT() / GeV);
          }
          else if (p.abspid() == PID::DSPLUS) {
            count_mult(p);
            _h["c_frac"]->fill(3);
            _h[name]->fill(p.pT() / GeV);
          }
          else if (p.abspid() == PID::LAMBDACPLUS) {
            count_mult(p);
            _h["c_frac"]->fill(4);
            _h[name]->fill(p.pT() / GeV);
          }
          else if (p.abspid() == PID::XICPLUS)
            _h["c_frac"]->fill(5);
          else if (p.abspid() == PID::XI0C)
            _h["c_frac"]->fill(6);
          else if (p.abspid() == PID::OMEGA0C)
            _h["c_frac"]->fill(7);
        }
      }
      Jets jets = apply<FastJets>(event, "Jets").jetsByPt(Cuts::pT > 25. * GeV && Cuts::absrap < 2.5);
      if (jets.empty()) vetoEvent;

      auto r_bins = _p["b_jet_rho"]->xEdges();
      const double dr = r_bins[1] - r_bins[0]; //dr is equal for all bins
      for (const Jet& thisJet : jets) {
        double p_0_R = p_annulus(thisJet, 0., 0.4);
        if (fuzzyEquals(p_0_R, 0., 1e-5)) continue;
        Particles bjets = thisJet.bTags(Cuts::pT > 5 * GeV);
        iselect(bjets, deltaRLess(thisJet, 0.3));
        for (const Particle& thisB : bjets) {
          _h["b_jet_pThad"]->fill(thisB.pT() / GeV);
          double z = thisJet.p3().dot(thisB.p3()) / thisJet.p2();
          _h["b_jet_frag"]->fill(z);
          if (inRange(thisJet.pT(), 500 * GeV, 1000 * GeV)) {
            _h["b_jet_high_pT"]->fill(z);
            if (bjets.size() == 1) _h["b_jet_high_pT_1H"]->fill(z);
          }
          for (size_t i = 0; i < r_bins.size() - 1; ++i) {
            double r = 0.5 * (r_bins[i] + r_bins[i + 1]);
            double rho = p_annulus(thisJet, r_bins[i], r_bins[i + 1]) / p_0_R / dr;
            double psi = p_annulus(thisJet, 0, r_bins[i + 1]) / p_0_R;
            _p["b_jet_rho"]->fill(r, rho);
            _p["b_jet_psi"]->fill(r, psi);
            size_t ptb = ptaxis.index(thisJet.pT() / GeV) - 1; // index in {1, nBins}
            if (ptb < ptaxis.numBins()) _p["avg_B_jet_rho_" + to_string(ptb)]->fill(r, rho);
          }
        }

        Particles cjets = thisJet.cTags(Cuts::pT > 5 * GeV);
        iselect(cjets, deltaRLess(thisJet, 0.3));
        if (bjets.empty()) {
          for (const Particle& thisC : cjets) {
            _h["c_jet_pThad"]->fill(thisC.pT() / GeV);
            double z = thisJet.p3().dot(thisC.p3()) / thisJet.p2();
            _h["c_jet_frag"]->fill(z);
            if (inRange(thisJet.pT(), 500 * GeV, 1000 * GeV)) {
              _h["c_jet_high_pT"]->fill(z);
              if (cjets.size() == 1) {
                _h["c_jet_high_pT_1H"]->fill(z);
              }
            }
            for (size_t i = 0; i < r_bins.size() - 1; ++i) {
              double r = 0.5 * (r_bins[i] + r_bins[i + 1]);
              double rho = p_annulus(thisJet, r_bins[i], r_bins[i + 1]) / p_0_R / dr;
              double psi = p_annulus(thisJet, 0, r_bins[i + 1]) / p_0_R;
              _p["c_jet_rho"]->fill(r, rho);
              _p["c_jet_psi"]->fill(r, psi);
              size_t ptb = ptaxis.index(thisJet.pT() / GeV) - 1; // index in {0, nBins - 1}
              if (ptb < ptaxis.numBins()) _p["avg_C_jet_rho_" + to_string(ptb)]->fill(r, rho);
            }
          }
        }

        if (bjets.size()) {
          double W_num = 0., W_den = 0.;
          long N_charged = 0;
          for (const Particle& pp : thisJet.particles()) {
            if (pp.isStable()) {
              W_num += pp.pT() * deltaR(thisJet, pp);
              W_den += pp.pT();
              if (pp.isCharged()) ++N_charged;
            }
            if (pp.isLepton()) {
              _h["b_jet_l_pT"]->fill(pp.pT() / GeV);
              _h["b_jet_l_pTrel"]->fill(pTrel(thisJet, pp) / GeV);
              _g["avg_B_jet_l_pTrel"]->fill(thisJet.pT() / GeV, pTrel(thisJet, pp) / GeV);
            }
          }
          if (W_den) _h["bar_b_jet_width"]->fill(W_num / W_den);
          _h["b_jet_pT"]->fill(thisJet.pT() / GeV);
          if (N_charged) {
            _h["b_jet_ch_mult"]->fill(N_charged);
            _g["avg_B_jet_ch_mult"]->fill(thisJet.pT() / GeV, (double)N_charged);
          }
        }
        else if (cjets.size()) {
          double W_num = 0., W_den = 0.;
          long N_charged = 0;
          for (const Particle& pp : thisJet.particles()) {
            if (pp.isStable()) {
              W_num += pp.pT() * deltaR(thisJet, pp);
              W_den += pp.pT();
              if (pp.isCharged()) ++N_charged;
            }
            if (pp.isLepton()) {
              _h["c_jet_l_pT"]->fill(pp.pT() / GeV);
              _h["c_jet_l_pTrel"]->fill(pTrel(thisJet, pp) / GeV);
              _g["avg_C_jet_l_pTrel"]->fill(thisJet.pT() / GeV, pTrel(thisJet, pp) / GeV);
            }
          }
          if (W_den) _h["bar_c_jet_width"]->fill(W_num / W_den);
          _h["c_jet_pT"]->fill(thisJet.pT() / GeV);
          if (N_charged) {
            _h["c_jet_ch_mult"]->fill(N_charged);
            _g["avg_C_jet_ch_mult"]->fill(thisJet.pT() / GeV, (double)N_charged);
          }
        }
      }
    }

    /// Normalise histograms etc., after the run
    void finalize() {
      for (const auto& hit : _h) {
        double sf = 1.0;
        if (hit.first.find("bar_") != string::npos) {
          sf = (hit.second->xMax() - hit.second->xMin()) / hit.second->numBins();
        }
        normalize(hit.second, sf);
      }
      normalize(_g);
    }

  private:

    map<string, Histo1DPtr> _h;
    map<string, Profile1DPtr> _p;
    map<string, Histo1DGroupPtr> _g;
    YODA::Axis<double> ptaxis{25, 30, 50, 70, 100, 150, 300, 500, 1000};
  };

  RIVET_DECLARE_PLUGIN(MC_HFDECAYS);
}