FillSliceLID_module.cc

Go to the documentation of this file.

#include "art/Framework/Core/EDProducer.h"
#include "art/Framework/Core/ModuleMacros.h"
#include "art/Framework/Principal/Event.h"
#include "art/Framework/Principal/Handle.h"
#include "art/Framework/Principal/Run.h"
#include "art/Framework/Principal/SubRun.h"
#include "canvas/Persistency/Common/FindManyP.h"
#include "fhiclcpp/ParameterSet.h"
#include "messagefacility/MessageLogger/MessageLogger.h"

#include "Geometry/Geometry.h"
#include "NovaDAQConventions/DAQConventions.h"
#include "RecoBase/Cluster.h"
#include "RecoBase/Vertex.h"
#include "RecoBase/FilterList.h"
#include "SummaryData/SpillData.h"
#include "Utilities/AssociationUtil.h"
#include "Utilities/func/EnvExpand.h"


#include <memory>

#include "SliceLID/art/var_dict_builder/Config.h"
#include "SliceLID/art/var_dict_builder/VarDictBuilder.h"
#include "SliceLID/prediction/Prediction.h"
#include "SliceLID/shifts/TrackLengthCorrection.h"
#include "TensorFlowEvaluator/SliceLID/model/Model.h"
#include "TensorFlowEvaluator/SliceLID/art/producer/utils.h"

namespace SliceLID
{

class FillSliceLID : public art::EDProducer
{

public:
    explicit FillSliceLID(const fhicl::ParameterSet &pset);

    FillSliceLID(FillSliceLID const &) = delete;
    FillSliceLID(FillSliceLID &&) = delete;
    FillSliceLID & operator = (FillSliceLID const &) = delete;
    FillSliceLID & operator = (FillSliceLID &&) = delete;

    void produce(art::Event &evt) override;

    void reconfigure(const fhicl::ParameterSet &pset);

    void fillSliceLIDPreds(
        art::Event &evt,
        std::unique_ptr<std::vector<Prediction>> &sliceLIDPreds,
        std::unique_ptr<art::Assns<Prediction, rb::Cluster>> &sliceLIDAssoc
    );

protected:

    bool getIsRHC(const art::Event &evt);
    Model& selectModel(const art::Event &evt);

    Prediction predict(
        Model &model,
        const art::Event &evt,
        const art::Handle<std::vector<rb::Cluster>> &slices_h,
        size_t sliceIdx,
        const std::vector<art::Ptr<rb::Vertex>> &elastics,
        const std::vector<bool> &sliceIsMC
    );

private:
    Config config;

    std::string generatorLabel;
    std::string beamLabel;

    Model modelFHC;
    Model modelRHC;

    bool fObeyPreselection;
    std::vector<std::string> fPreselectionLabels;

    bool useOppositeHornCurrentNetwork;
};


FillSliceLID::FillSliceLID(const fhicl::ParameterSet &pset)
    : modelFHC(util::EnvExpansion(pset.get<std::string>("modelFHC"))),
      modelRHC(util::EnvExpansion(pset.get<std::string>("modelRHC"))),
      fObeyPreselection  (pset.get<bool>                     ("ObeyPreselection"  )),
      fPreselectionLabels(pset.get<std::vector<std::string>> ("PreselectionLabels")),
      useOppositeHornCurrentNetwork(pset.get<bool>("useOppositeHornCurrentNetwork"))
{
    reconfigure(pset);

    produces<std::vector<Prediction>>();
    produces<art::Assns<Prediction, rb::Cluster>>();
}

void FillSliceLID::reconfigure(const fhicl::ParameterSet &pset)
{
    beamLabel      = pset.get<std::string>("beamLabel");
    generatorLabel = pset.get<std::string>("generatorLabel");

    config.parse(pset);
}

bool FillSliceLID::getIsRHC(const art::Event &evt)
{
    /* TODO: ask if this should be cached per file/run/job/etc ? */
    art::Handle<sumdata::SpillData> spillPot;
    if (! evt.isRealData()) {
        evt.getByLabel(generatorLabel, spillPot);
    }
    else {
        evt.getByLabel(beamLabel, spillPot);
    }

    if (spillPot.failedToGet())
    {
        mf::LogError("FillSliceLID") <<
            "Spill Data not found, aborting without horn current information";
        abort();
    }

    if(useOppositeHornCurrentNetwork) return !spillPot->isRHC;

    // NB - the logic here will cause 0HC to use the FHC network
    return spillPot->isRHC;
}

Model& FillSliceLID::selectModel(const art::Event &evt)
{
    bool isRHC = getIsRHC(evt);

    if (isRHC) {
        return modelRHC;
    }
    else {
        return modelFHC;
    }
}

Prediction FillSliceLID::predict(
    Model &model,
    const art::Event &evt,
    const art::Handle<std::vector<rb::Cluster>> &slices_h,
    size_t sliceIdx,
    const std::vector<art::Ptr<rb::Vertex>> &elastics,
    const std::vector<bool> &sliceIsMC
)
{
    VarDictBuilder builder(evt, slices_h, sliceIdx, elastics, config);
    VarDict varDict = builder.build();

    /* NOTE: this must be removed once geometry is fixed */
    /* TODO: ask if backtracker->HasTruthInfo() is better to get isMC */
    art::ServiceHandle<geo::Geometry> geom;
    TrackLengthCorrection trkLenCorr(
        sliceIsMC[sliceIdx], (geom->DetId() == novadaq::cnv::kFARDET)
    );
    trkLenCorr.shift(varDict);

    return model.predict(varDict);
}

void FillSliceLID::fillSliceLIDPreds(
    art::Event &evt,
    std::unique_ptr<std::vector<Prediction>> &sliceLIDPreds,
    std::unique_ptr<art::Assns<Prediction, rb::Cluster>> &sliceLIDAssoc
)
{
    art::Handle<std::vector<rb::Cluster>> slices_h;
    if (! evt.getByLabel(config.sliceLabel, slices_h))
    {
        mf::LogError("FillSliceLID")
            << "Failed to get slice info" << std::endl;
        return;
    }

    art::FindManyP<rb::Vertex> elasticVxt_fmp(
        slices_h, evt, config.elasticArmsLabel
    );
    if (! elasticVxt_fmp.isValid())
    {
        mf::LogError("FillSliceLID") << "No elastics found" << std::endl;
        return;
    }

    Model &model                = selectModel(evt);
    std::vector<bool> sliceIsMC = getSliceIsMC(*slices_h);

    for (size_t sliceIdx = 0; sliceIdx < slices_h->size(); ++sliceIdx)
    {
        const art::Ptr<rb::Cluster> slicePtr(slices_h, sliceIdx);

        if (slicePtr->IsNoise()) {
            continue;
        }
        if(fObeyPreselection && rb::IsFiltered(evt, slices_h, sliceIdx, fPreselectionLabels)) continue;

        auto elastics = elasticVxt_fmp.at(sliceIdx);

        Prediction pred = predict(
            model, evt, slices_h, sliceIdx, elastics, sliceIsMC
        );

        sliceLIDPreds->emplace_back(std::move(pred));
        util::CreateAssn(
            *this, evt,
            *(sliceLIDPreds.get()), slicePtr, *(sliceLIDAssoc.get())
        );
    }
}

void FillSliceLID::produce(art::Event & evt)
{
    std::unique_ptr<std::vector<Prediction>> sliceLIDPreds(
        new std::vector<Prediction>
    );
    std::unique_ptr<art::Assns<Prediction, rb::Cluster>> sliceLIDAssoc(
        new art::Assns<Prediction, rb::Cluster>
    );

    fillSliceLIDPreds(evt, sliceLIDPreds, sliceLIDAssoc);

    evt.put(std::move(sliceLIDPreds));
    evt.put(std::move(sliceLIDAssoc));
}

DEFINE_ART_MODULE(FillSliceLID)

}