CVNNeutronProngTF_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 "RecoBase/FilterList.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 cvnneutronprongtf{
  class CVNNeutronProngTF : public art::EDProducer{
  public:
    explicit CVNNeutronProngTF(fhicl::ParameterSet const &pset);
    virtual ~CVNNeutronProngTF();

    void produce(art::Event& evt);
    tensorflow::Tensor vector_to_tensor(std::vector<unsigned char>);
    tensorflow::Tensor vector_to_tensor(std::vector<unsigned char>, std::vector<unsigned char>);
    tensorflow::TFHandler* GetModel2D(const art::Event &evt);
    tensorflow::TFHandler* GetModel3D(const art::Event &evt);
    void CalcResult( std::vector< art::Ptr<cvn::PixelMap> > pixelMaps, tensorflow::TFHandler* fTF, unsigned int NOutput,
                     std::vector<std::pair<int, double>> &outputwithpdg, std::vector<float> &result , std::vector<unsigned char> pmslc);

  protected:
    std::string fSliceLabel;
    std::string fPixelMapInput;
    std::string fGeneratorLabel;
    bool        fObeyPreselection;
    std::vector<std::string> fPreselectionLabels;
    bool        fUseGeV;
    std::string fLibPath;

    std::string fProngModLabel;
    std::string fProng2DLabel;
    std::string fModel2DName;
    unsigned int fNOutput2D;

    std::string fProng3DLabel;
    std::string fModel3DName;
    unsigned int fNOutput3D;

    std::string fInputNameSlc;
    std::string fInputName;
    std::string fOutputName;
    unsigned int fCPUlimit;
    tensorflow::TFHandler* fTF2D;
    tensorflow::TFHandler* fTF3D;

  };

  CVNNeutronProngTF::CVNNeutronProngTF(fhicl::ParameterSet const &pset):
    EDProducer(pset),
    fSliceLabel   (pset.get<std::string>("SliceLabel")),
    fPixelMapInput(pset.get<std::string>("PixelMapInput")),
    fGeneratorLabel (pset.get<std::string>("GeneratorLabel")),
    fObeyPreselection  (pset.get<bool>                     ("ObeyPreselection"  )),
    fPreselectionLabels(pset.get<std::vector<std::string>> ("PreselectionLabels")),
    fUseGeV       (pset.get<bool>       ("UseGeV")),
    fLibPath    (pset.get<std::string>("LibPath")),

    fProngModLabel(pset.get<std::string>("ProngModLabel")),
    fProng2DLabel (pset.get<std::string>("Prong2DLabel")),
    fModel2DName  (pset.get<std::string>("Model2DName")),
    fNOutput2D    (pset.get<unsigned int>("NOutput2D")),

    fProng3DLabel (pset.get<std::string>("Prong3DLabel")),
    fModel3DName  (pset.get<std::string>("Model3DName")),
    fNOutput3D    (pset.get<unsigned int>("NOutput3D")),

    fInputNameSlc    (pset.get<std::string>("InputNameSlice")),
    fInputName    (pset.get<std::string>("InputName")),
    fOutputName   (pset.get<std::string>("OutputName")),
    fCPUlimit     (pset.get<unsigned int>("CPUlimit")),
    fTF2D(0),
    fTF3D(0)
  {
    fLibPath = fLibPath.empty() ? fLibPath : util::EnvExpansion(fLibPath);
    produces< std::vector<cvn::Result>   >();
    produces< art::Assns<cvn::Result, rb::Cluster> >();
    produces< art::Assns<cvn::Result, rb::Prong> >();
    produces< std::vector<rb::PID>          >();
    produces< art::Assns<rb::PID,rb::Prong> >();
  }

  CVNNeutronProngTF::~CVNNeutronProngTF()
  { 
    if(fTF2D) delete fTF2D;
    if(fTF3D) delete fTF3D;

  }

  tensorflow::TFHandler* CVNNeutronProngTF::GetModel2D(const art::Event &evt)
  {
    if (!fTF2D)
      fTF2D = new tensorflow::TFHandler(fLibPath + fModel2DName, fCPUlimit);
    return fTF2D;
  }

  tensorflow::TFHandler* CVNNeutronProngTF::GetModel3D(const art::Event &evt)
  {
    if (!fTF3D)
      fTF3D = new tensorflow::TFHandler(fLibPath + fModel3DName, fCPUlimit);
    return fTF3D;
  }

  tensorflow::Tensor CVNNeutronProngTF::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;
  }

  tensorflow::Tensor CVNNeutronProngTF::vector_to_tensor(std::vector<unsigned char> pmslice, std::vector<unsigned char> pmprong)
  {
    const unsigned int vectorSizeSlice = pmslice.size();
    const unsigned int vectorSizeProng = pmprong.size();
    const unsigned int vectorSizeTotal = vectorSizeSlice + vectorSizeProng;

    // Initialize the tensors
    tensorflow::Tensor tensor(tensorflow::DT_FLOAT, {1, vectorSizeTotal});
    auto rel = tensor.tensor<float,2>();
    
    // Loop over each element
    for(unsigned int i = 0; i < vectorSizeSlice; ++i) rel(0, i) = pmslice[i];
    for(unsigned int j = vectorSizeSlice; j < vectorSizeTotal; ++j) rel(0, j) = pmprong[j-vectorSizeSlice];

    return tensor;
  }

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

    tensorflow::TFHandler* fTF2D = GetModel2D(evt);
    tensorflow::TFHandler* fTF3D = GetModel3D(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> >
                                  assocresultslice(new art::Assns<cvn::Result, rb::Cluster>);
    std::unique_ptr< art::Assns<cvn::Result, rb::Prong> >
                                  assocresultprong(new art::Assns<cvn::Result, rb::Prong>);    
    std::unique_ptr< std::vector<rb::PID> > 
                                  pidCol(new std::vector<rb::PID>);
    std::unique_ptr< art::Assns<rb::PID, rb::Prong> >
                                  assocpid(new art::Assns<rb::PID, rb::Prong>);

    // 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;
      if(fObeyPreselection && rb::IsFiltered(evt, slicecol, iClust, fPreselectionLabels)) continue;

      //get slice pixel maps
      const std::vector<art::Ptr<cvn::PixelMap> > pixelMapsSl = fmPixelMap.at(iClust);
      if(pixelMapsSl.empty()) continue;
      std::vector<unsigned char> pmslice = (*pixelMapsSl[0]).PixelMapToVector(fUseGeV);


      // --- 2D Prongs.
      art::FindManyP<rb::Prong> fmProng2D(slicecol, evt, art::InputTag(fProngModLabel, fProng2DLabel));
      if( fmProng2D.isValid() ) {
        // Get everything out of the event.
        std::vector<art::Ptr<rb::Prong>> prongs2D = fmProng2D.at(iClust);
        art::FindManyP<cvn::PixelMap> fmPixelMap2D(prongs2D, evt, fPixelMapInput);
        if ( fmPixelMap2D.isValid() ) {

          for( unsigned int iProng = 0; iProng < prongs2D.size(); ++iProng ){
            // Calculate result
            std::vector<std::pair<int, double>> outputwithpdg;
            std::vector<float> result;
            if ( fmPixelMap2D.at(iProng).empty() ) continue;
            CalcResult( fmPixelMap2D.at(iProng), fTF2D, fNOutput2D, outputwithpdg, result , pmslice );
            // Ensure that result.size() equals fNOutput
            assert( result.size() == fNOutput2D);
            // Make Assn between result and this prong.
            float resultvec[fNOutput2D];
            for(unsigned int i = 0; i < fNOutput2D; i++){
              resultvec[i] = result[i];
              // Make Assn.
              pidCol->emplace_back(outputwithpdg[i].first, outputwithpdg[i].second);
              util::CreateAssn(evt, *(pidCol.get()),
                               prongs2D[iProng], *(assocpid.get()),  UINT_MAX);
            }
            // Make Assn between full result output and prong.
            const float* resultoutput = resultvec;
            resultCol->emplace_back(resultoutput, fNOutput2D);
            util::CreateAssn(evt, *(resultCol.get()),
                             slicelist[iClust], *(assocresultslice.get()), UINT_MAX);
            util::CreateAssn(evt, *(resultCol.get()),
                             prongs2D[iClust], *(assocresultprong.get()), UINT_MAX);
          }
        }
      }

      // --- 3D Prongs.
      art::FindManyP<rb::Prong> fmProng3D(slicecol, evt, art::InputTag(fProngModLabel, fProng3DLabel) );
      std::vector<art::Ptr<rb::Prong>> prongs3D;
      if( fmProng3D.isValid() ) {
        // Get everything out of the event.
        std::vector<art::Ptr<rb::Prong>> prongs3D = fmProng3D.at(iClust);
        art::FindManyP<cvn::PixelMap> fmPixelMap3D(prongs3D, evt, fPixelMapInput);
        if ( fmPixelMap3D.isValid() ) {
          for( unsigned int iProng = 0; iProng < prongs3D.size(); ++iProng ){
            // Calculate result
            std::vector<std::pair<int, double>> outputwithpdg;
            std::vector<float> result;
            if ( fmPixelMap3D.at(iProng).empty() ) continue;
            CalcResult( fmPixelMap3D.at(iProng), fTF3D, fNOutput3D, outputwithpdg, result , pmslice );
            // Ensure that result.size() equals fNOutput
            //assert( result.size() == fNOutput3D);
            // Make Assn between result and this prong.
            float resultvec[fNOutput3D];
            for(unsigned int i = 0; i < result.size(); i++){
              resultvec[i] = result[i];
              // Make Assn.
              pidCol->emplace_back(outputwithpdg[i].first, outputwithpdg[i].second);
              util::CreateAssn(evt, *(pidCol.get()),
                               prongs3D[iProng], *(assocpid.get()),  UINT_MAX);
            }
            // Make Assn between full result output and prong.
            const float* resultoutput = resultvec;
            resultCol->emplace_back(resultoutput, fNOutput3D);
            util::CreateAssn(evt, *(resultCol.get()),
                             slicelist[iClust], *(assocresultslice.get()), UINT_MAX);
            util::CreateAssn(evt, *(resultCol.get()),
                             prongs3D[iClust], *(assocresultprong.get()), UINT_MAX);
          }
        }
      }

    } // slices
    evt.put(std::move(resultCol));
    evt.put(std::move(assocresultslice));
    evt.put(std::move(assocresultprong));
    evt.put(std::move(pidCol));
    evt.put(std::move(assocpid));
  } // produce

  // --- Function to calculate the CVN inference.
  // --- Need a function so can do both 2D and 3D evaluation.
  void CVNNeutronProngTF::CalcResult( std::vector< art::Ptr<cvn::PixelMap> > pixelMaps, tensorflow::TFHandler* fTF, unsigned int NOutput,
                                      std::vector<std::pair<int, double>> &outputwithpdg, std::vector<float> &result , std::vector<unsigned char> pmslc) {

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

    tensorflow::Tensor tensorslc;
    tensorflow::Tensor tensor;
    tensor = vector_to_tensor(pmprong);
    tensorslc = vector_to_tensor(pmslc);


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

      for(unsigned int i = 0; i < NOutput; i++){
        //Fill and create PID assn
        if(i==0){
          outputwithpdg.push_back(std::make_pair(2212, (double)tfoutput(0,i)));
        }
        else if (i==1){
          outputwithpdg.push_back(std::make_pair(22  , (double)tfoutput(0,i)));
        }
        else if (i==2){
          outputwithpdg.push_back(std::make_pair(0   , (double)tfoutput(0,i)));
        }
        //Fill Result column
        result.push_back( (float)tfoutput(0,i) );
      }

    return;
  }
}
DEFINE_ART_MODULE(cvnneutronprongtf::CVNNeutronProngTF)