ProngBPFVars.cxx

Go to the documentation of this file.

#include "LSTME/art/var_dict_builder/ProngBPFVars.h"

#include "RecoBase/Energy.h"
#include "RecoBase/FitSum.h"


namespace LSTME
{

std::array<BPFVars, N_BPF> ProngBPFVars::getEmptyBPFVars()
{
    const BPFVars emptyBPF{ 0.0, 0.0, { 0, 0, 0 } };
    std::array<BPFVars, N_BPF> result;

    for (size_t i = 0; i < N_BPF; i++) {
        result[i] = emptyBPF;
    }

    return result;
}

std::array<BPFVars, N_BPF> ProngBPFVars::getBPFVars(
    const art::Event &evt,
    const art::FindManyP<rb::Track> &prongBPF_fmp,
    size_t pngIdx
)
{
    std::array<BPFVars, N_BPF> result = getEmptyBPFVars();
    if (! prongBPF_fmp.isValid()) {
        return result;
    }

    auto tracks = prongBPF_fmp.at(pngIdx);
    if (tracks.empty()) {
        return result;
    }

    art::FindManyP<rb::FitSum> trackFits_fmp(tracks, evt, prongBPFLabel);
    if (! trackFits_fmp.isValid()) {
        return result;
    }

    art::FindManyP<rb::Energy> trackOverlapE_fmp(
        tracks, evt, art::InputTag(overlapELabel, overlapEBPFLabel)
    );

    if (N_BPF != tracks.size()) {
        return result;
    }

    for (size_t trackIdx = 0; trackIdx < tracks.size(); trackIdx++)
    {
        auto trackFits = trackFits_fmp.at(trackIdx);
        if (trackFits.empty()) {
            continue;
        }

        const rb::FitSum &fit = *trackFits[0];

        result[trackIdx].momentum[0] = fit.FourMom().Px();
        result[trackIdx].momentum[1] = fit.FourMom().Py();
        result[trackIdx].momentum[2] = fit.FourMom().Pz();
        result[trackIdx].energy      = fit.FourMom().E();

        if (trackOverlapE_fmp.isValid())
        {
            auto overlapEnergies = trackOverlapE_fmp.at(trackIdx);
            if (! overlapEnergies.empty()) {
                result[trackIdx].overlapE = overlapEnergies[0]->E();
            }
        }

    }

    return result;
}

void ProngBPFVars::init(const art::Event &evt, const BasicProngVars &pngVars)
{
    if (! pngVars.hasProngs) {
        return;
    }

    art::FindManyP<rb::Track> prongBPF_fmp(pngVars.prongs, evt, prongBPFLabel);
    bpfVars.reserve(pngVars.prongs.size());

    for (size_t pngIdx = 0; pngIdx < pngVars.prongs.size(); pngIdx++) {
        bpfVars.emplace_back(getBPFVars(evt, prongBPF_fmp, pngIdx));
    }
}

void ProngBPFVars::saveToDict(
    std::unordered_map<std::string, std::vector<double>> &dict,
    const std::string &prefix
) const
{
    std::array<std::vector<double>, N_BPF> energy;
    std::array<std::vector<double>, N_BPF> overlapE;
    std::array<std::vector<double>, N_BPF> momentum_x;
    std::array<std::vector<double>, N_BPF> momentum_y;
    std::array<std::vector<double>, N_BPF> momentum_z;

    for (size_t bpfIdx = 0; bpfIdx < N_BPF; bpfIdx++)
    {
        energy[bpfIdx]    .reserve(bpfVars.size());
        overlapE[bpfIdx]  .reserve(bpfVars.size());
        momentum_x[bpfIdx].reserve(bpfVars.size());
        momentum_y[bpfIdx].reserve(bpfVars.size());
        momentum_z[bpfIdx].reserve(bpfVars.size());
    }

    for (auto &bpf : bpfVars)
    {
        for (size_t bpfIdx = 0; bpfIdx < N_BPF; bpfIdx++)
        {
            energy[bpfIdx]    .push_back(bpf[bpfIdx].energy);
            overlapE[bpfIdx]  .push_back(bpf[bpfIdx].overlapE);
            momentum_x[bpfIdx].push_back(bpf[bpfIdx].momentum[0]);
            momentum_y[bpfIdx].push_back(bpf[bpfIdx].momentum[1]);
            momentum_z[bpfIdx].push_back(bpf[bpfIdx].momentum[2]);
        }
    }

    for (size_t bpfIdx = 0; bpfIdx < N_BPF; bpfIdx++)
    {
        const std::string &fullPrefix =
            prefix + "bpf[" + std::to_string(bpfIdx) + "]." ;

        dict.emplace(fullPrefix + "energy",     std::move(energy[bpfIdx]));
        dict.emplace(fullPrefix + "overlapE",   std::move(overlapE[bpfIdx]));
        dict.emplace(fullPrefix + "momentum.x", std::move(momentum_x[bpfIdx]));
        dict.emplace(fullPrefix + "momentum.y", std::move(momentum_y[bpfIdx]));
        dict.emplace(fullPrefix + "momentum.z", std::move(momentum_z[bpfIdx]));
    }
}

ProngBPFVars::ProngBPFVars(
    const art::Event &evt,
    const BasicProngVars &pngVars,
    const Config &config
) : prongBPFLabel(config.prongBPFLabel),
    overlapELabel(config.overlapELabel),
    overlapEBPFLabel(config.overlapEBPFLabel)
{
    init(evt, pngVars);
}

}