SecondaryBlockCD.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/ParameterSet.h>
#include <momemta/Module.h>
#include <momemta/Solution.h>
#include <momemta/Math.h>
#include <momemta/InputTag.h>
#include <momemta/Types.h>

#include <Math/GenVector/VectorUtil.h>

class SecondaryBlockCD: public Module {
    public:

        SecondaryBlockCD(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {
                sqrt_s = parameters.globalParameters().get<double>("energy");

                s12 = get<double>(parameters.get<InputTag>("s12"));

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

        virtual Status work() override {

            solutions->clear();

            const double m1 = p1->M();
            const double sq_m1 = SQ(m1);
            const double m2 = p2->M();
            const double sq_m2 = SQ(m2);

            // Don't spend time on unphysical part of phase-space
            if (*s12 >= SQ(sqrt_s) || sq_m1 + sq_m2 >= *s12)
               return Status::NEXT;

            std::vector<double> E1_solutions; // up to two solutions
            const double theta1 = p1->Theta();
            const double phi1 = p1->Phi();

            const double E2 = p2->E();
            const double norm2 = p2->P();

            const double cos_theta12 = ROOT::Math::VectorUtil::CosTheta(*p1, *p2);

            // Equation to be solved for E1 : s12 = (p1+p2)^2 ==> [ 4*SQ(cos_theta12)*SQ(norm2)-4*SQ(E2) ] * SQ(E1) + [ 4*(s12 - SQ(m1) - SQ(m2))*E2 ] * E1 + 4*SQ(m1)*SQ(cos_theta12)*(SQ(m2)-SQ(E2) = 0)
            const double quadraticTerm = 4 * SQ(E2) - 4 * SQ(norm2) * SQ(cos_theta12);
            const double linearTerm = 4 * E2 * (sq_m1 + sq_m2 - *s12);
            const double indepTerm = SQ(sq_m1 + sq_m2 - *s12) + 4 * sq_m1 * SQ(norm2) * SQ(cos_theta12);

            bool foundSolution = solveQuadratic(quadraticTerm, linearTerm, indepTerm, E1_solutions);

            if (!foundSolution) {
                return Status::NEXT;
            }

            for (const double& E1: E1_solutions) {
                // Skip unphysical solutions
                if (E1 <= 0 || E1 <= m1)
                    continue;

                // Avoid introducing spurious solution from the equation s12 = (p1+p2)^2 that has to be squared in the computation
                if ((sq_m1 + sq_m2 + 2 * E1 * E2 - *s12) * cos_theta12 < 0)
                    continue;

                LorentzVector gen_p1_sol;
                double norm1 = std::sqrt(SQ(E1) - sq_m1);
                double gen_pt1_sol = norm1 * std::sin(theta1);
                gen_p1_sol.SetPxPyPzE(
                        gen_pt1_sol * std::cos(phi1),
                        gen_pt1_sol * std::sin(phi1),
                        norm1 * std::cos(theta1),
                        E1);

                if (!ApproxComparison(gen_p1_sol.M() / gen_p1_sol.E(), m1 / gen_p1_sol.E())) {
#ifndef NDEBUG
                    LOG(trace) << "[SecondaryBlockCD] Throwing solution because of invalid mass. " <<
                        "Expected " << m1 << ", got " << gen_p1_sol.M();
#endif
                    continue;
                }


                if (!ApproxComparison((gen_p1_sol + *p2).M2(), *s12)) {
#ifndef NDEBUG
                    LOG(trace) << "[SecondaryBlockCD] Throwing solution because of invalid invariant mass. " <<
                        "Expected " << *s12 << ", got " << (gen_p1_sol + *p2).M2();
#endif
                    continue;
                }
                
                // Compute jacobian
                double jacobian = std::abs( std::sin(theta1) * SQ(norm1) / (32 * CB(M_PI) * (norm1 * E2 - E1 * norm2 * cos_theta12)) );

                Solution s { {gen_p1_sol}, jacobian, true };
                solutions->push_back(s);
            }

            return (solutions->size() > 0) ? Status::OK : Status::NEXT;
        }

    private:
        double sqrt_s;

        // Inputs
        Value<double> s12;
        Value<LorentzVector> p1;
        Value<LorentzVector> p2;

        // Output
        std::shared_ptr<SolutionCollection> solutions = produce<SolutionCollection>("solutions");
};

REGISTER_MODULE(SecondaryBlockCD)
        .Input("s12")
        .Input("p1")
        .Input("p2")
        .Output("solutions")
        .GlobalAttr("energy: double");