LSTMEVar.h

Go to the documentation of this file.

#pragma once

#include "CAFAna/Core/Var.h"
#include "CAFAna/Core/MultiVar.h"
#include "3FlavorAna/Vars/LSTMEHelperVars.h"

#include "SliceLID/shifts/TrackLengthCorrection.h"
#include "TensorFlowEvaluator/LSTME/model/Model.h"
#include "TensorFlowEvaluator/SliceLID/cafana/ResultCacher.hxx"

namespace LSTME
{

using namespace ana;
using CAFAnaModel = ResultCacher<LSTMEnergy>;

static VarDict fillVarDict(
    const caf::SRProxy *sr,
    const std::vector<std::pair<std::string, Var>>      &sliceVarDefs,
    const std::vector<std::pair<std::string, MultiVar>> &png2dVarDefs,
    const std::vector<std::pair<std::string, MultiVar>> &png3dVarDefs
)
{
    VarDict result;

    for (auto &namedVar : sliceVarDefs) {
        result.scalarVarMap[namedVar.first] = namedVar.second(sr);
    }

    for (auto &namedMultiVar : png2dVarDefs) {
        result.vectorVarMap[namedMultiVar.first] = namedMultiVar.second(sr);
    }

    for (auto &namedMultiVar : png3dVarDefs) {
        result.vectorVarMap[namedMultiVar.first] = namedMultiVar.second(sr);
    }

    return result;
}

static std::shared_ptr<CAFAnaModel> initCAFAnaModel(
    const std::string &modelSavedir
)
{
    auto model = std::make_shared<Model>(modelSavedir);

    return std::make_shared<CAFAnaModel>(
        [
          model      = std::move(model),
          trkLenCorr = TrackLengthCorrection()
        ] (const caf::SRProxy *sr) mutable -> LSTMEnergy
        {
            auto varDict = fillVarDict(
                sr, kSliceVarDefs, kPng2dVarDefs, kPng3dVarDefs
            );

            /* TODO: remove air bubble correction once geometry is fixed */
            bool isMC = (sr->mc.nallnus > 0) || (sr->mc.nallcosmics > 0);
            trkLenCorr.setParams(isMC, (sr->hdr.det == caf::kFARDET));
            trkLenCorr.shift(varDict);

            return model->predict(varDict);
        }
    );
}

static Var primaryEnergy(const std::shared_ptr<CAFAnaModel> &model)
{
    return Var(
        [model] (const caf::SRProxy *sr) -> double
        {
            return (*model)(sr).primaryE;
        }
    );
}

static Var totalEnergy(const std::shared_ptr<CAFAnaModel> &model)
{
    return Var(
        [model] (const caf::SRProxy *sr) -> double
        {
            return (*model)(sr).totalE;
        }
    );
}

static Var secondaryEnergy(const std::shared_ptr<CAFAnaModel> &model)
{
    return Var(
        [model] (const caf::SRProxy *sr) -> double
        {
            return (*model)(sr).totalE - (*model)(sr).primaryE;
        }
    );
}

}