MatrixElement.cc

/*
 *  MoMEMta: a modular implementation of the Matrix Element Method
 *  Copyright (C) 2016  Universite catholique de Louvain (UCL), Belgium
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

// Must be loaded before `momemta/Logging.h`, otherwise there's conflict between usage
// of `log()` and `namespace log`
#include <numeric>
#include <LHAPDF/LHAPDF.h>

#include <momemta/Logging.h>
#include <momemta/MatrixElement.h>
#include <momemta/MatrixElementFactory.h>
#include <momemta/ParameterSet.h>
#include <momemta/Math.h>
#include <momemta/Module.h>
#include <momemta/Types.h>
#include <momemta/Utils.h>

class MatrixElement: public Module {
    struct ParticleId {
        int64_t pdg_id;
        int64_t me_index;
    };

    public:

        MatrixElement(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {

            sqrt_s = parameters.globalParameters().get<double>("energy");

            use_pdf = parameters.get<bool>("use_pdf", true);

            m_partons = get<std::vector<LorentzVector>>(parameters.get<InputTag>("initialState"));

            const auto& particles_set = parameters.get<ParameterSet>("particles");

            auto particle_tags = particles_set.get<std::vector<InputTag>>("inputs");
            for (auto& tag: particle_tags)
                m_particles.push_back(get<LorentzVector>(tag));

            const auto& particles_ids_set = particles_set.get<std::vector<ParameterSet>>("ids");
            for (const auto& s: particles_ids_set) {
                ParticleId id;
                id.pdg_id = s.get<int64_t>("pdg_id");
                id.me_index = s.get<int64_t>("me_index");
                m_particles_ids.push_back(id);
            }

            if (m_particles.size() != m_particles_ids.size()) {
                LOG(fatal) << "The number of particles ids is not consistent with the number of particles. Did you"
                            " forget some ids?";

                throw Module::invalid_configuration("Inconsistent number of ids and number of particles");
            }

            const auto& jacobians_tags = parameters.get<std::vector<InputTag>>("jacobians", std::vector<InputTag>());
            for (const auto& tag: jacobians_tags) {
                m_jacobians.push_back(get<double>(tag));
            }

            std::string matrix_element = parameters.get<std::string>("matrix_element");
            const ParameterSet& matrix_element_configuration = parameters.get<ParameterSet>("matrix_element_parameters");
            m_ME = MatrixElementFactory::get().create(matrix_element, matrix_element_configuration);

            if (parameters.exists("override_parameters")) {
                const ParameterSet& matrix_element_params = parameters.get<ParameterSet>("override_parameters");
                auto p = m_ME->getParameters();

                for (const auto& name: matrix_element_params.getNames()) {
                    double value = matrix_element_params.get<double>(name);
                    p->setParameter(name, value);
                }

                p->cacheParameters();
                p->cacheCouplings();
            }

            // PDF, if asked
            if (use_pdf) {
                // Silence LHAPDF
                LHAPDF::setVerbosity(0);

                std::string pdf = parameters.get<std::string>("pdf");
                m_pdf.reset(LHAPDF::mkPDF(pdf, 0));

                double pdf_scale = parameters.get<double>("pdf_scale");
                pdf_scale_squared = SQ(pdf_scale);
            }

            // Prepare arrays for sorting particles
            for (size_t i = 0; i < m_particles_ids.size(); i++) {
                indexing.push_back(m_particles_ids[i].me_index - 1);
            }

            // Pre-allocate memory for the finalState array
            finalState.resize(m_particles_ids.size());

            // Sort the array taking into account the indexing in the configuration
            std::vector<int64_t> suite(indexing.size());
            std::iota(suite.begin(), suite.end(), 0);

            permutations = get_permutations(suite, indexing);
        };

        virtual void beginIntegration() {
            // Don't assume the non-zero helicities will be the same for each event
            // In principle they are, but this protects against buggy calls to the ME (e.g. returning NaN or inf)
            m_ME->resetHelicities();
        }

        virtual Status work() override {
            static std::vector<LorentzVector> empty_vector;

            *m_integrand = 0;
            const std::vector<LorentzVector>& partons = *m_partons;

            LorentzVectorRefCollection particles;
            for (const auto& p: m_particles)
                particles.push_back(std::ref(*p));

            for (size_t i = 0; i < m_particles_ids.size(); i++) {
                finalState[i] = std::make_pair(m_particles_ids[i].pdg_id, toVector(particles[i].get()));
            }

            // Sort the array taking into account the indexing in the configuration
            apply_permutations(finalState, permutations);

            std::pair<std::vector<double>, std::vector<double>> initialState { toVector(partons[0]),
                                                                               toVector(partons[1]) };

            auto result = m_ME->compute(initialState, finalState);

            double x1 = std::abs(partons[0].Pz() / (sqrt_s / 2.));
            double x2 = std::abs(partons[1].Pz() / (sqrt_s / 2.));

            // Compute flux factor 1/(2*x1*x2*s)
            double phaseSpaceIn = 1. / (2. * x1 * x2 * SQ(sqrt_s));

            double integrand = phaseSpaceIn;
            for (const auto& jacobian: m_jacobians) {
                integrand *= (*jacobian);
            }

            // PDF
            double final_integrand = 0;
            for (const auto& me: result) {
                double pdf1 = use_pdf ? m_pdf->xfxQ2(me.first.first, x1, pdf_scale_squared) / x1 : 1;
                double pdf2 = use_pdf ? m_pdf->xfxQ2(me.first.second, x2, pdf_scale_squared) / x2 : 1;

                final_integrand += me.second * pdf1 * pdf2;
            }

            final_integrand *= integrand;
            *m_integrand = final_integrand;

            return Status::OK;
        }

    private:
        double sqrt_s;
        bool use_pdf;
        double pdf_scale_squared = 0;
        std::shared_ptr<momemta::MatrixElement> m_ME;
        std::shared_ptr<LHAPDF::PDF> m_pdf;

        std::vector<int64_t> indexing;
        std::vector<size_t> permutations;
        std::vector<std::pair<int, std::vector<double>>> finalState;

        // Inputs
        Value<std::vector<LorentzVector>> m_partons;

        std::vector<Value<LorentzVector>> m_particles;
        std::vector<ParticleId> m_particles_ids;

        std::vector<Value<double>> m_jacobians;

        // Outputs
        std::shared_ptr<double> m_integrand = produce<double>("output");
};

REGISTER_MODULE(MatrixElement)
        .Input("initialState")
        .OptionalInputs("jacobians")
        .Inputs("particles/inputs")
        .Output("output")
        .GlobalAttr("energy:double")
        .Attr("matrix_element:string")
        .Attr("matrix_element_parameters:pset")
        .OptionalAttr("override_parameters:pset")
        .Attr("particles:pset")
        .Attr("use_pdf:bool=true")
        .OptionalAttr("pdf:string")
        .OptionalAttr("pdf_scale:double");