CVNProngEvaluatorTF_module.cc

Go to the documentation of this file.

////////////////////////////////////////////////////////////////////////
// \file    CVNProngEvaluatorTF_module.cc
// \brief   Producer module creating 2D prong CVN neural net results
////////////////////////////////////////////////////////////////////////

// C/C++ includes
#include <iostream>
#include <sstream>

// ROOT includes
#include "TTree.h"
#include "TH2F.h"

// Framework includes
#include "art/Framework/Core/EDProducer.h"
#include "art/Framework/Principal/Event.h"
#include "art/Framework/Principal/Handle.h"
#include "art/Framework/Services/Optional/TFileDirectory.h"
#include "art/Framework/Services/Optional/TFileService.h"
#include "art/Framework/Services/Registry/ServiceHandle.h"
#include "fhiclcpp/ParameterSet.h"
#include "messagefacility/MessageLogger/MessageLogger.h"
#include "canvas/Persistency/Common/FindManyP.h"
#include "art/Framework/Core/ModuleMacros.h"
#include "canvas/Persistency/Common/Assns.h"

// NOvASoft includes
#include "RecoBase/PID.h"
#include "RecoBase/Cluster.h"
#include "RecoBase/FilterList.h"
#include "Utilities/AssociationUtil.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 "TensorFlowHandler/TFHandler.h"
#include "Utilities/func/EnvExpand.h"

using namespace cvn;

namespace cvntf {

  class CVNProngEvaluatorTF : public art::EDProducer {
  public:
    explicit CVNProngEvaluatorTF(fhicl::ParameterSet const& pset);
    ~CVNProngEvaluatorTF();

    bool filter(std::vector<unsigned char> pm, int view);
    std::vector<tensorflow::Tensor> vector_to_tensor(std::vector<unsigned char>, std::vector<unsigned char>);

    void produce(art::Event& evt);
    void beginJob();
    void endJob();



  private:
    /// Module lablel for input clusters
    const art::ProductToken<std::vector<rb::Cluster>> fClusterToken;

    /// Module label for input pixel maps
    std::string fPixelMapInput;
    std::string fSlicePMLabel;
    std::string fProngPMLabel;

    /// Module lablel for input prongs
    std::string    fProngModLabel;
    /// Instance lablel for input 3D prongs
    // use 2d view
    bool        fIs2DProng;
    std::string fProng3DLabel;
    std::string fProng2DLabel;
    mutable std::string fProngLabel;

    ///  Check rb::IsFiltered?
    bool        fObeyPreselection;
    bool        fSkipLongProngs;
    bool        fProngOnly;
    double      fProngLengthCut;

    //

    /// Number of outputs fron neural net
    unsigned int fNOutput;
    bool        fUseGeV;
    std::vector<std::string> fPreselectionLabels;
    std::string fLibPath;
    std::string fModelName;
    tensorflow::TFHandler* fTF;
  };



  //.......................................................................
  CVNProngEvaluatorTF::CVNProngEvaluatorTF(fhicl::ParameterSet const& pset):

    fClusterToken(consumes<std::vector<rb::Cluster>>(pset.get<std::string>("ClusterLabel"))),
    fPixelMapInput (pset.get<std::string>         ("PixelMapInput")),
    fSlicePMLabel  (pset.get<std::string>         ("SlicePMLabel")),
    fProngPMLabel  (pset.get<std::string>         ("ProngPMLabel")),
    fProngModLabel (pset.get<std::string>         ("ProngModLabel")),
    fIs2DProng     (pset.get<bool>                ("Is2DProng")),
    fProng3DLabel  (pset.get<std::string>         ("Prong3DLabel")),
    fProng2DLabel  (pset.get<std::string>         ("Prong2DLabel")),
    fProngLabel  (pset.get<std::string>           ("ProngLabel")),//(fIs2Dorong?"Prong2DLabel": "Prong3DLabel")),

    fObeyPreselection (pset.get<bool>             ("ObeyPreselection")),
    fSkipLongProngs(pset.get<bool>                ("SkipLongProngs")),
    fProngOnly     (pset.get<bool>                ("ProngOnly")),
    fProngLengthCut (pset.get<double>             ("ProngLengthCut")),
    fNOutput(8),
    fUseGeV(true),
    fPreselectionLabels (pset.get<std::vector<std::string>> ("PreselectionLabels")),
    fLibPath(pset.get<std::string>("LibPath")),
    fModelName(pset.get<std::string>("ModelName")),
    fTF(0)
  {
    fLibPath = fLibPath.empty() ? fLibPath : util::EnvExpansion(fLibPath);
    fTF = new tensorflow::TFHandler(fLibPath + fModelName);
    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> >();
  }

  //......................................................................
  CVNProngEvaluatorTF::~CVNProngEvaluatorTF()
  {
    if(fTF) delete fTF;

    //======================================================================
    // Clean up any memory allocated by your module
    //======================================================================
  }

  bool CVNProngEvaluatorTF::filter(std::vector<unsigned char> pm, int view)
  {
    float x(0), y(0), tmpx(0), tmpy(0);
    for(int i = 0; i<8000; i++){
      tmpx=pm[i];
      if(tmpx>0){
        x++; 
      }
    }
    for(int i = 8000; i<16000; i++){
      tmpy = pm[i];
      if(tmpy>0) y++;
    }
    if(view && y<=3) return false; 
    if((!view) && x<=3) return false;
    else return true;
  }

  std::vector<tensorflow::Tensor> CVNProngEvaluatorTF::vector_to_tensor(std::vector<unsigned char> prongpm, std::vector<unsigned char> slicepm)
  {
    tensorflow::Tensor tensorx(tensorflow::DT_FLOAT, {1,16000});
    tensorflow::Tensor tensory(tensorflow::DT_FLOAT, {1,16000});

    auto relx = tensorx.tensor<float,2>();
    auto rely = tensory.tensor<float,2>();
    float x(0),y(0);

    for(int i = 0; i<8000; i++){
      //odd-view
      relx(0, i) = slicepm[i];
      relx(0, i+8000) = prongpm[i];
      x += prongpm[i];
    }
    for(int i = 8000; i< 16000; i++){
      rely(0, i-8000) = slicepm[i]; 
      rely(0, i) = prongpm[i];
      y += prongpm[i];
    }
    return {tensorx, tensory};
  }

  //......................................................................
  void CVNProngEvaluatorTF::beginJob()
  {  }

  //......................................................................
  void CVNProngEvaluatorTF::endJob()
  {
  }

  //......................................................................
  void CVNProngEvaluatorTF::produce(art::Event& evt)
  {

    /// Define containers for the things we're going to produce
    std::unique_ptr< std::vector<Result> >
                                  resultCol(new std::vector<Result>);
    std::unique_ptr< std::vector<rb::PID> >
                                  pidCol(new std::vector<rb::PID>);
    std::unique_ptr< art::Assns<Result, rb::Cluster> >
                                  assoc(new art::Assns<Result, rb::Cluster>);
    std::unique_ptr< art::Assns<Result, rb::Prong> >
                                  assocp(new art::Assns<Result, rb::Prong>);
    std::unique_ptr< art::Assns<rb::PID, rb::Prong> >
                                  assocpid(new art::Assns<rb::PID, rb::Prong>);

    // get the slices
    art::Handle< std::vector< rb::Cluster > > clusterHandle;
    evt.getByToken(fClusterToken, clusterHandle);
    const std::vector<rb::Cluster> & clusters = *clusterHandle;

    art::PtrVector<rb::Cluster>  sliceCol;
    for(unsigned int i = 0; i<clusterHandle->size(); i++){
      art::Ptr<rb::Cluster> clust(clusterHandle,i);
      sliceCol.push_back(clust);
    }

    art::FindManyP<PixelMap> fmSlicePM(clusterHandle, evt, fPixelMapInput);//art::InputTag(fPixelMapInput, fSlicePMLabel));
    if(!fmSlicePM.isValid()) return;

    for(size_t iClust = 0; iClust < clusters.size(); iClust++) {
      if(fObeyPreselection && rb::IsFiltered(evt,clusterHandle,iClust,fPreselectionLabels))
        continue;
      if(clusters[iClust].IsNoise()) continue;

      const std::vector<art::Ptr<PixelMap> > slicePMs = fmSlicePM.at(iClust);

      art::FindManyP<rb::Prong> fmProng(clusterHandle, evt, art::InputTag(fProngModLabel, fProng2DLabel));
      if(!fmProng.isValid()) continue;
      //for 3d's x,y 2ds
      std::vector<art::Ptr<rb::Prong>> prongs;

      prongs = fmProng.at(iClust);      

      if(prongs.empty()) continue;
      art::FindManyP<PixelMap> fmProngPM(prongs, evt, fPixelMapInput);
      if( !fmProngPM.isValid()) continue;
      // prong first
      for( unsigned int iProng = 0; iProng < prongs.size(); iProng++ ){ 
        art::Ptr<rb::Prong> prong = prongs[iProng];
        int view = 2;
        
        if(fIs2DProng){
          if(prong->NXCell()) view = 0;
          else if(prong->NYCell()) view = 1;
          else std::cout<<"How the hell do we get to no cells at all..."<<std::endl;
        }
        else{
          if(!prong->NXCell() || !prong->NYCell()){
           std::cout<<" 3d prong with an empty view!!!!!!"<<std::endl;
           continue;
          }
          abort(); //abort here for now
        }
        if ( fSkipLongProngs && prong->TotalLength() > fProngLengthCut ) continue;
        const std::vector<art::Ptr<PixelMap>> pm = fmProngPM.at(iProng);
        if (pm.empty()) continue;
        if (slicePMs.empty()) continue;
        std::vector<unsigned char> prongpm = (*pm[0]).PixelMapToVector(fUseGeV);
        std::vector<unsigned char> slicepm = (*slicePMs[0]).PixelMapToVector(fUseGeV);

        if(!filter(prongpm, view)) continue;
        std::vector<tensorflow::Tensor> tensor = vector_to_tensor(prongpm, slicepm);
        std::vector<tensorflow::Tensor> result = fTF->Predict({std::make_pair("input", tensor[view])}, {"output/Softmax"});

        float output[fNOutput];

        auto tfoutput = result[0].tensor<float,2>();

        std::vector<std::pair<int, double>> outputwithpdg;

        for ( unsigned int j = 0; j < fNOutput; j++ ){
          outputwithpdg.push_back(std::make_pair(GetPDGByPType((cvn::PType)j), (double)tfoutput(0,j)));
          pidCol->emplace_back(outputwithpdg[j].first, outputwithpdg[j].second);
          util::CreateAssn(*this, evt, *(pidCol.get()),
                           prongs[iProng], *(assocpid.get()),  UINT_MAX);
          output[j] = (float)tfoutput(0,j);
        }
        const float* outputptr = output; 
        resultCol->emplace_back(outputptr, fNOutput);
//        PixelMap tmp = *pm[0];
//        (tmp).Print();
//        for (unsigned int i = 0; i < fNOutput; i++) {std::cout<<*outputptr<<std::endl;outputptr++;}

          // Output object with PDGs rather than version dependednt cvn ID
        util::CreateAssn(*this, evt, *(resultCol.get()),
                           sliceCol[iClust], *(assoc.get()),  UINT_MAX);
        util::CreateAssn(*this, evt, *(resultCol.get()),
                           prongs[iProng], *(assocp.get()),  UINT_MAX);
      }//iProng
    }//iClust

    evt.put(std::move(pidCol));
    evt.put(std::move(resultCol));
    evt.put(std::move(assoc));
    evt.put(std::move(assocp));
    evt.put(std::move(assocpid));

  }

  //----------------------------------------------------------------------


}
  DEFINE_ART_MODULE(cvntf::CVNProngEvaluatorTF)
////////////////////////////////////////////////////////////////////////