CVNEventTF_module.cc

Go to the documentation of this file.

#include<iostream>

#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 "art/Framework/Core/ModuleMacros.h"
#include "canvas/Utilities/InputTag.h"
#include "canvas/Persistency/Common/FindManyP.h"
#include "canvas/Persistency/Common/Assns.h"
#include "fhiclcpp/ParameterSet.h"
#include "messagefacility/MessageLogger/MessageLogger.h"


#include "RecoBase/Cluster.h"
#include "RecoBase/Prong.h"
#include "RecoBase/PID.h"
#include "SummaryData/SpillData.h"
#include "CVN/func/PixelMap.h"
#include "CVN/func/Result.h"
#include "CVN/func/InteractionType.h"
#include "CVN/func/AssignLabels.h"
#include "CVN/func/TrainingData.h"
#include "CVN/func/InteractionType.h"
#include "CVN/func/EventLabeledPMaps.h"
#include "CVN/func/ProngType.h"
#include "CVN/func/ProngTrainingData.h"
#include "Utilities/AssociationUtil.h"
#include "Utilities/func/EnvExpand.h"

#include "TensorFlowHandler/TFHandler.h"

#include <memory>

namespace cvneventtf{
  class CVNEventTF : public art::EDProducer{
  public:
    explicit CVNEventTF(fhicl::ParameterSet const &pset);
    virtual ~CVNEventTF();

    void produce(art::Event& evt);
    bool IsRHC(const art::Event &evt);
    tensorflow::Tensor vector_to_tensor(std::vector<unsigned char>);
    tensorflow::TFHandler* GetModel(const art::Event &evt);

  protected:
    std::string fSliceLabel;
    std::string fPixelMapInput;
    std::string fGeneratorLabel;
    std::string fNuMILabel;
    bool        fUseGeV;
    std::string fLibPath;
    std::string fModelFHCName;
    std::string fModelRHCName;
    unsigned int fNOutput;
    std::string fInputName;
    std::string fOutputName;
    unsigned int fCPUlimit;
    tensorflow::TFHandler* fTFFHC;
    tensorflow::TFHandler* fTFRHC;
  };
}

namespace cvneventtf{
  CVNEventTF::CVNEventTF(fhicl::ParameterSet const &pset):
    fSliceLabel   (pset.get<std::string>("SliceLabel")),
    fPixelMapInput(pset.get<std::string>("PixelMapInput")),
    fGeneratorLabel (pset.get<std::string>("GeneratorLabel")),
    fNuMILabel (pset.get<std::string>("NuMILabel")),
    fUseGeV       (pset.get<bool>       ("UseGeV")),
    fLibPath    (pset.get<std::string>("LibPath")),
    fModelFHCName    (pset.get<std::string>("ModelFHCName")),
    fModelRHCName    (pset.get<std::string>("ModelRHCName")),
    fNOutput      (pset.get<unsigned int>("NOutput")),
    fInputName    (pset.get<std::string>("InputName")),
    fOutputName   (pset.get<std::string>("OutputName")),
    fCPUlimit     (pset.get<unsigned int>("CPUlimit")),
    fTFFHC(0),
    fTFRHC(0)
  {
    fLibPath = fLibPath.empty() ? fLibPath : util::EnvExpansion(fLibPath);
    produces< std::vector<cvn::Result>          >();
    produces< art::Assns<cvn::Result,rb::Cluster> >();
  }

  CVNEventTF::~CVNEventTF()
  { 
    if(fTFFHC) delete fTFFHC;
    if(fTFRHC) delete fTFRHC;
  }
  
  bool CVNEventTF::IsRHC(const art::Event &evt)
  {
      art::Handle<sumdata::SpillData> spillPot;
      if (!evt.isRealData()) 
          evt.getByLabel(fGeneratorLabel, spillPot);
      else 
          evt.getByLabel(fNuMILabel, spillPot);

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

  tensorflow::TFHandler* CVNEventTF::GetModel(const art::Event &evt)
  {
    
    if (IsRHC(evt)) {
      if (!fTFRHC)
        fTFRHC = new tensorflow::TFHandler(fLibPath + fModelRHCName, fCPUlimit);
      return fTFRHC;
    }
    else {
      if (!fTFFHC)
        fTFFHC = new tensorflow::TFHandler(fLibPath + fModelFHCName, fCPUlimit);
      return fTFFHC;
    }
  }

  tensorflow::Tensor CVNEventTF::vector_to_tensor(std::vector<unsigned char> pm)
  {
    const unsigned int vectorSize = pm.size();

    // Initialize the tensors
    tensorflow::Tensor tensor(tensorflow::DT_FLOAT, {1, vectorSize});
    auto rel = tensor.tensor<float,2>();

    // Loop over each element
    for(unsigned int i = 0; i < vectorSize; ++i) rel(0, i) = pm[i];
    
    return tensor;
  }

  void CVNEventTF::produce(art::Event& evt)
  {

    tensorflow::TFHandler* fTF = GetModel(evt);
    //Containers for things we're gonna produce
    std::unique_ptr< std::vector<cvn::Result> > 
                                  resultCol(new std::vector<cvn::Result>);
    std::unique_ptr< art::Assns<cvn::Result, rb::Cluster> >
                                  assocresult(new art::Assns<cvn::Result, rb::Cluster>);

    // Get slices
    art::Handle<std::vector<rb::Cluster> > slicecol;
    evt.getByLabel(fSliceLabel, slicecol);
    art::PtrVector<rb::Cluster> slicelist;
    for(unsigned int i = 0; i < slicecol->size(); ++i){
      slicelist.push_back(art::Ptr<rb::Cluster>(slicecol, i));
    }

    // Get pixel maps
    art::FindManyP<cvn::PixelMap> fmPixelMap(slicecol, evt, fPixelMapInput);

    //loop over slices
    for(size_t iClust = 0; iClust < slicelist.size(); ++iClust) {
      if(!fmPixelMap.isValid()) continue;
      if(slicelist[iClust]->IsNoise()) continue;

      const std::vector<art::Ptr<cvn::PixelMap> > pixelMaps = fmPixelMap.at(iClust);

      if(pixelMaps.empty()) continue;

      std::vector<unsigned char> pmslice = (*pixelMaps[0]).PixelMapToVector(fUseGeV);

      tensorflow::Tensor tensor = vector_to_tensor(pmslice);   

      std::vector<tensorflow::Tensor> result = fTF->Predict({{fInputName,tensor}},
                                                            {fOutputName});
      auto tfoutput = result[0].tensor<float,2>();      

      float resultvec[fNOutput];

      for(unsigned int i = 0; i<fNOutput; i++){  
        resultvec[i] = (float)tfoutput(0,i);
      }
      const float* output = resultvec;
  
      resultCol->emplace_back(output, fNOutput);
      
      util::CreateAssn(*this, evt, *(resultCol.get()),
                       slicelist[iClust], *(assocresult.get()),  UINT_MAX);

    } // slices
   
    evt.put(std::move(resultCol));
    evt.put(std::move(assocresult));

  } // produce
}

DEFINE_ART_MODULE(cvneventtf::CVNEventTF)