BlockA.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/>.
*/

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

#include <stdexcept>

class BlockA: public Module {
    public:

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

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

            p1 = get<LorentzVector>(parameters.get<InputTag>("p1"));
            p2 = get<LorentzVector>(parameters.get<InputTag>("p2"));

            auto branches_tags = parameters.get<std::vector<InputTag>>("branches");
            for (auto& t: branches_tags)
                m_branches.push_back(get<LorentzVector>(t));
            if (m_branches.empty()) {
                auto exception = std::invalid_argument("BlockA is not valid without at least a third particle in the event.");
                LOG(fatal) << exception.what();
                throw exception;
            }
        };

        virtual Status work() override {

            solutions->clear();

            LorentzVector pb;
            for (size_t i = 0; i < m_branches.size(); i++) {
                pb += *m_branches[i];
            }

            double pbx = pb.Px();
            double pby = pb.Py();
            const double theta1 = p1->Theta();
            const double phi1 = p1->Phi();
            const double theta2 = p2->Theta();
            const double phi2 = p2->Phi();
            const double m1 = p1->M();
            const double m2 = p2->M();

            // pT = p1+p2+pb = 0. Equivalent to the following system:
            // p1x+p2x = -pbx
            // p1y+p2y = -pby,
            // where: p1x=modp1*sin(theta1)*cos(phi1), p1y=modp1*sin(theta1)*sin(phi1),
            //        p2x=modp2*sin(theta2)*cos(phi2), p2y=modp2*sin(theta2)*sin(phi2)
            // Get modp1, modp2 as solutions of this system

            std::vector<double> modp1;
            std::vector<double> modp2;

            const double sin_theta1 = std::sin(theta1);
            const double cos_phi1 = std::cos(phi1);
            const double sin_theta2 = std::sin(theta2);
            const double cos_phi2 = std::cos(phi2);
            const double sin_phi1 = std::sin(phi1);
            const double sin_phi2 = std::sin(phi2);
            const double cos_theta1 = std::cos(theta1);
            const double cos_theta2 = std::cos(theta2);

            const double a10 = sin_theta1 * cos_phi1;
            const double a01 = sin_theta2 * cos_phi2;
            const double a00 = pbx;
            const double b10 = sin_theta1 * sin_phi1;
            const double b01 = sin_theta2 * sin_phi2;
            const double b00 = pby;

            bool foundSolution = solve2Linear(a10, a01, a00, b10, b01, b00, modp1, modp2, false);

            if (!foundSolution)
               return Status::NEXT;

            double mod_p1 = modp1[0];
            double mod_p2 = modp2[0];
            if (mod_p1 < 0 || mod_p2 < 0)
               return Status::NEXT;

            double E1 = sqrt(SQ(mod_p1) + SQ(m1));
            double E2 = sqrt(SQ(mod_p2) + SQ(m2));

            LorentzVector gen_p1(mod_p1 * sin_theta1 * cos_phi1, mod_p1 * sin_theta1 * sin_phi1, mod_p1 * cos_theta1, E1);
            LorentzVector gen_p2(mod_p2 * sin_theta2 * cos_phi2, mod_p2 * sin_theta2 * sin_phi2, mod_p2 * cos_theta2, E2);

            // Check if solutions are physical
            LorentzVector tot = gen_p1 + gen_p2 + pb;
            double q1Pz = std::abs(tot.Pz() + tot.E()) / 2.;
            double q2Pz = std::abs(tot.Pz() - tot.E()) / 2.;
            if (q1Pz > sqrt_s / 2 || q2Pz > sqrt_s / 2)
                return Status::NEXT;

            // Pt should be 0
            if (!ApproxComparison(tot.Pt(), 0)) {
#ifndef NDEBUG
                LOG(trace) << "[BlockA] Throwing solution because total Pt is not null: " << tot.Pt();
#endif
                return Status::NEXT;
            }

            double jacobian = (SQ(mod_p1) * SQ(mod_p2)) / (8 * SQ(M_PI * sqrt_s) * E1 * E2);
            jacobian *= 1. / std::abs(cos_phi1 * sin_phi2 - sin_phi1 * cos_phi2);

            Solution s = { {gen_p1, gen_p2}, jacobian, true };
            solutions->push_back(s);

            return Status::OK;
        }


    private:
        double sqrt_s;

        // Inputs
        Value<LorentzVector> p1, p2;
        std::vector<Value<LorentzVector>> m_branches;
        // Outputs
        std::shared_ptr<SolutionCollection> solutions = produce<SolutionCollection>("solutions");
};

REGISTER_MODULE(BlockA)
    .Input("p1")
    .Input("p2")
    .OptionalInputs("branches")
    .Output("solutions")
    .GlobalAttr("energy:double");