BuildInitialState.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 <functional>

#include <Math/Vector3D.h>
#include <Math/Boost.h>

#include <momemta/Module.h>
#include <momemta/ParameterSet.h>
#include <momemta/Types.h>

class BuildInitialState: public Module {
    public:

        BuildInitialState(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {
            if (parameters.get<bool>("do_transverse_boost", false)) {
                do_compute_initials = compute_initials_boost;
            } else {
                do_compute_initials = compute_initials_trivial;
            }

            halved_sqrt_s = parameters.globalParameters().get<double>("energy") / 2;

            std::vector<InputTag> input_particles_tags = parameters.get<std::vector<InputTag>>("particles");
            for (auto& t: input_particles_tags)
                input_particles.push_back(get<LorentzVector>(t));
        };

        virtual Status work() override {

            partons->clear();

            LorentzVectorRefCollection particles;
            for (auto& p: input_particles) {
                particles.push_back(std::ref(*p));
            }

            do_compute_initials(particles);

            // Check if solutions are physical
            if ((*partons)[0].E() > halved_sqrt_s || (*partons)[1].E() > halved_sqrt_s )
                return Status::NEXT;

            return Status::OK;
        }

    private:
        using compute_initials_signature = std::function<void(const LorentzVectorRefCollection&)>;

        compute_initials_signature do_compute_initials;

        compute_initials_signature compute_initials_trivial =
            [&](const LorentzVectorRefCollection& particles) {
                LorentzVector tot;
                for (const auto& p: particles)
                    tot += p.get();

                double q1Pz = (tot.Pz() + tot.E()) / 2.;
                double q2Pz = (tot.Pz() - tot.E()) / 2.;

                partons->push_back(LorentzVector(0., 0., q1Pz, std::abs(q1Pz)));
                partons->push_back(LorentzVector(0., 0., q2Pz, std::abs(q2Pz)));
            };

        compute_initials_signature compute_initials_boost =
            [&](const LorentzVectorRefCollection& particles) {
                LorentzVector tot;
                for (const auto& p: particles)
                    tot += p.get();

                // Define boost that puts the transverse total momentum vector in its CoM frame
                LorentzVector transverse_tot = tot;
                transverse_tot.SetPz(0);
                ROOT::Math::XYZVector isr_deBoost_vector( transverse_tot.BoostToCM() );

                // In the "transverse" CoM frame, use total Pz and E to define initial longitudinal quark momenta
                const ROOT::Math::Boost isr_deBoost( isr_deBoost_vector );
                const ROOT::Math::PxPyPzEVector new_tot( isr_deBoost * tot );

                double q1Pz = (new_tot.Pz() + new_tot.E()) / 2.;
                double q2Pz = (new_tot.Pz() - new_tot.E()) / 2.;

                partons->push_back(LorentzVector(0., 0., q1Pz, std::abs(q1Pz)));
                partons->push_back(LorentzVector(0., 0., q2Pz, std::abs(q2Pz)));

                // Boost initial parton momenta by the opposite of the transverse boost needed to put the whole system in its CoM
                const ROOT::Math::Boost isr_boost( -isr_deBoost_vector );
                (*partons)[0] = isr_boost * (*partons)[0];
                (*partons)[1] = isr_boost * (*partons)[1];
            };

        double halved_sqrt_s;

        std::vector<Value<LorentzVector>> input_particles;

        std::shared_ptr<std::vector<LorentzVector>> partons = produce<std::vector<LorentzVector>>("partons");
};

REGISTER_MODULE(BuildInitialState)
        .Inputs("particles")
        .Output("partons")
        .GlobalAttr("energy:double")
        .Attr("do_transverse_boost:bool=false");