BPFTmvaTrainer_module.cc

Go to the documentation of this file.

////////////////////////////////////////////////////////////////////////
// \file    BPFTmvaTrainer_module.cc
// \brief   Module to create the TTree(s) used to train TMVA for the
//          BPF particle PIDs and the BPF energy estimator. This module
//          is ONLY intended to run over MC files.
// \version 
// \author  $Author: mbaird42 $
////////////////////////////////////////////////////////////////////////

// C/C++ includes
#include <cmath>
#include <iostream>
#include <vector>

// ROOT includes
#include "TH1F.h"
#include "TH2F.h"
#include "TFile.h"
#include "TTree.h"
#include "TMVA/Factory.h"
#include "TMVA/Tools.h"
#include "TMVA/Reader.h"
#include "TMVA/MethodKNN.h"

// Framework includes
#include "art/Framework/Core/EDAnalyzer.h"
#include "canvas/Persistency/Common/FindManyP.h"
#include "art/Framework/Principal/Event.h"
#include "art/Framework/Principal/SubRun.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 "canvas/Persistency/Common/Ptr.h"
#include "canvas/Persistency/Common/PtrVector.h"
#include "fhiclcpp/ParameterSet.h"
#include "messagefacility/MessageLogger/MessageLogger.h"

// NOvASoft includes
#include "RecoBase/CellHit.h"
#include "RecoBase/RecoHit.h"
#include "RecoBase/Cluster.h"
#include "RecoBase/Prong.h"
#include "RecoBase/Track.h"
#include "RecoBase/FitSum.h"

#include "ReMId/ReMId.h"

#include "Utilities/AssociationUtil.h"
#include "Geometry/Geometry.h"
#include "Calibrator/Calibrator.h"

#include "BreakPointFitter/func/dEdxCalculator.h"
#include "MuonRemove/art/TrackCleanUpAlg.h"

#include "MCCheater/BackTracker.h"
#include "nusimdata/SimulationBase/MCNeutrino.h"
#include "nusimdata/SimulationBase/MCParticle.h"

#include "SummaryData/POTSum.h"



// BPFTmvaTrainer header
namespace bpfit {
  class BPFTmvaTrainer : public art::EDAnalyzer {
  public:
    explicit BPFTmvaTrainer(fhicl::ParameterSet const& pset);    
    ~BPFTmvaTrainer();                        
    
    void analyze(const art::Event& evt); 
    void reconfigure(const fhicl::ParameterSet& pset);
    void beginJob();
    void endJob();
    virtual void endSubRun(const art::SubRun& sr);

    void resetVars();
    void setTruthVars(const std::vector<cheat::NeutrinoEffPur>& nuEP);

  private:
    // FCL parameters:
    const art::ProductToken<std::vector<rb::Cluster>> fSlicerToken;
    const art::ProductToken<sumdata::POTSum> fPOTSumToken;
    std::string   fProngLabel;    ///< label for module that made the prongs
    std::string   fProngInstance; ///< instance label for the prongs to be used
    std::string   fTrackLabel;    ///< label for module that made the tracks
    std::string   fHistoFile;     ///< location of dE/dx log-likelihood histogram
    std::string   fXMLFile;       ///< location of the XML file for the muon PID
    double        fdxTOL;         ///< lower limit for dx values to be used in dE/dx calculation
    unsigned int  fTrainMode;     ///< training mode (energy estimator or muon PID?)
    int           fPIDmode;       ///< make the training TTree using the muon, pion, or proton fit

    // Things used for internal processing and helper classes:
    TH2F  *fdEdxVSbg[2][2];                  ///< histo of computed dE/dx for cell hits (vs. beta*gamma)
    bpfit::dEdxCalculator fdEdxCalc;         ///< helper class for computing dEdx values
    TMVA::Reader fMuonPID;                   ///< TMVA reader object for the muon PID
    murem::TrackCleanUpAlg fTrackCleanUpAlg; ///< object to determine hadronic energy on the muon track
    double fTotalPOT;                        ///< total POT counted for all events in this job

    // Output TTrees:
    TTree *fBPFTrackSignalTree;      ///< TTree with track level signal info (for PID training)
    TTree *fBPFTrackBackgroundTree;  ///< TTree with track level background info (for PID training)
    TTree *fBPFSliceTree;            ///< TTree with slice level info (for energy estimator training)


    
    //
    // Variables to go into the Track TTrees.
    //

    // training variables:
    float fLength;      ///< track length
    float fChi2H;       ///< chi^2 hits value per DOF
    float fChi2A;       ///< chi^2 angle value per DOF
    float fChi2T;       ///< total chi^2
    float fdEdxLL;      ///< dEdx LL value
    float fHitRatio;    ///< ratio of track hits to prong hits (used as a track/shower discriminator)
    float fLengthRatio; ///< ratio of track to longest track in the slice
    float fDChi2Tmupi;  ///< delta chi2 total muon - pion
    float fDChi2Tmupr;  ///< delta chi2 total muon - proton
    float fDChi2Tpipr;  ///< delta chi2 total pion - proton

    // track PDG code used to separate signal and background
    float fTrackPDG; ///< PDG code for the primary track particle



    //
    // Variables to go into the Slice TTree. Input variables are mapped to the
    // target variables using the cut variables for event selection.
    //
    // fE and the cut variables are also used for the PID tree
    //

    // target variables:
    float fE;   ///< true energy
    float fEsq; ///< true energy squared (used to estiamte st. dev.)

    // input variables:
    float fPmu;       ///< reconstructed muon momentum from the track selected to be the muon
    float fDirZmu;    ///< muon track dirZ from the track selected to be the muon
    float fN3Dprongs; ///< total number of 3D prongs
    float fE3Dprongs; ///< sum of hit energies in 3D prongs (excluding the "muon" prong) - this does NOT account for dead material
    float fEremain;   ///< sum of remaining slice energy not in 3D prongs PLUS extra energy added back in from TrackCleanUpAlg
    float fSumPE;     ///< sum of the PE for all hits not on the muon track

    // spectator variables:
    float fEhadMu;    ///< Hadronic energy contamination on the muon track as computed by TrackCleanUpAlg

    // cut variables:
    float fMinX;  ///< minimum hit X value in the slice
    float fMinY;  ///< minimum hit Y value in the slice
    float fMinZ;  ///< minimum hit Z value in the slice
    float fMaxX;  ///< maximum hit X value in the slice
    float fMaxY;  ///< maximum hit Y value in the slice
    float fMaxZ;  ///< maximum hit Z value in the slice
    int fNHitX;   ///< number of X-view hits in the slice
    int fNHitY;   ///< number of Y-view hits in the slice
    int fCCNC;    ///< is this CC or NC (from truth)
    int fIntType; ///< neutrino interaction type (from truth)
    int fnuPDG;   ///< PDG code for the best matched nu (from truth)
    float fBestBPFmuPID; ///< best BPF muon PID value
    float fBestRemid;    ///< best remid value
    int fBestBPFmuPIDpdg; ///< pdg code for the particle that contributed the most energy to the track with the best BPF muon PID

    TH1D *fPOT; ///< Total POT accumulated for all events processed in this job.

  };
}



namespace bpfit
{
  //.......................................................................
  BPFTmvaTrainer::BPFTmvaTrainer(fhicl::ParameterSet const& pset) 
    : EDAnalyzer(pset),
      fSlicerToken(consumes<std::vector<rb::Cluster>>(pset.get<std::string>("SlicerLabel"))),
      fPOTSumToken(consumes<sumdata::POTSum, art::InSubRun>("generator")),
      fdEdxCalc(),
      fMuonPID(),
      fTrackCleanUpAlg(pset.get<fhicl::ParameterSet>("TrackCleanUpAlgPSet")),
      fTotalPOT(0.0)
  {
    this->reconfigure(pset);
  }

  //......................................................................
  BPFTmvaTrainer::~BPFTmvaTrainer()
  {
    //======================================================================
    // Clean up any memory allocated by your module... ...if you dare...
    //======================================================================
  }

  //......................................................................
  void BPFTmvaTrainer::reconfigure(const fhicl::ParameterSet& pset)
  {
    fProngLabel    = pset.get<std::string> ("ProngLabel");
    fProngInstance = pset.get<std::string> ("ProngInstance");
    fTrackLabel    = pset.get<std::string> ("TrackLabel");
    fHistoFile     = pset.get<std::string> ("HistoFile");
    fXMLFile       = pset.get<std::string> ("XMLFile");
    fdxTOL         = pset.get<double>      ("dxTOL");
    fTrainMode     = pset.get<unsigned int>("TrainMode");
    fPIDmode       = pset.get<int>         ("PIDmode");
  }
      
  //......................................................................
  void BPFTmvaTrainer::beginJob()
  {

    //
    // Check that the requested PID training mode is supported.
    //
    if(fTrainMode == 0) {
      if(fPIDmode != 13 && fPIDmode != 211 && fPIDmode != 2212) {
        std::cerr << "\n\n\nWARNING: PID training mode for PDG code = "
                  << fPIDmode
                  << " is not currently supported. ABORTING...\n\n\n";
        abort();
      }
    }

    //
    // Read in the log-likelihood histos from the input file
    //
    TFile infile(fHistoFile.c_str());

    fdEdxVSbg[0][0] = (TH2F*)infile.FindObjectAny("fdEdxVSbg_W0_X0");
    fdEdxVSbg[0][1] = (TH2F*)infile.FindObjectAny("fdEdxVSbg_W0_X1");
    fdEdxVSbg[1][0] = (TH2F*)infile.FindObjectAny("fdEdxVSbg_W1_X0");
    fdEdxVSbg[1][1] = (TH2F*)infile.FindObjectAny("fdEdxVSbg_W1_X1");
    
    infile.Close();



    //
    // Set up the TMVA reader for the muon PID.
    //
    if(fTrainMode == 1) {
      fMuonPID.AddVariable("length",&fLength);
      fMuonPID.AddVariable("chi2T",&fChi2T);
      fMuonPID.AddVariable("dEdxLL",&fdEdxLL);
      fMuonPID.AddVariable("hitRatio",&fHitRatio);
      fMuonPID.BookMVA("KNN Method",fXMLFile);
    }



    art::ServiceHandle<art::TFileService> tfs;

    //
    // Book the PID Training TTrees
    //

    // The signal TTree...
    fBPFTrackSignalTree = tfs->make<TTree>("BPFTrackSignalTree","Tree for TMVA training with signal track level BPF info");

    // input variables:
    fBPFTrackSignalTree->Branch("length",&fLength,"length/F");
    fBPFTrackSignalTree->Branch("chi2H",&fChi2H,"chi2H/F");
    fBPFTrackSignalTree->Branch("chi2A",&fChi2A,"chi2A/F");
    fBPFTrackSignalTree->Branch("chi2T",&fChi2T,"chi2T/F");
    fBPFTrackSignalTree->Branch("dEdxLL",&fdEdxLL,"dEdxLL/F");
    fBPFTrackSignalTree->Branch("hitRatio",&fHitRatio,"hitRatio/F");
    fBPFTrackSignalTree->Branch("lengthRatio",&fLengthRatio,"lengthRatio/F");
    fBPFTrackSignalTree->Branch("DChi2Tmupi",&fDChi2Tmupi,"DChi2Tmupi/F");
    fBPFTrackSignalTree->Branch("DChi2Tmupr",&fDChi2Tmupr,"DChi2Tmupr/F");
    fBPFTrackSignalTree->Branch("DChi2Tpipr",&fDChi2Tpipr,"DChi2Tpipr/F");

    // signal/background discriminator
    fBPFTrackSignalTree->Branch("trackPDG",&fTrackPDG,"trackPDG/F");

    // cut/weight variables:
    fBPFTrackSignalTree->Branch("E",&fE,"E/F");
    fBPFTrackSignalTree->Branch("MinX",&fMinX,"MinX/F");
    fBPFTrackSignalTree->Branch("MinY",&fMinY,"MinY/F");
    fBPFTrackSignalTree->Branch("MinZ",&fMinZ,"MinZ/F");
    fBPFTrackSignalTree->Branch("MaxX",&fMaxX,"MaxX/F");
    fBPFTrackSignalTree->Branch("MaxY",&fMaxY,"MaxY/F");
    fBPFTrackSignalTree->Branch("MaxZ",&fMaxZ,"MaxZ/F");
    fBPFTrackSignalTree->Branch("NHitX",&fNHitX,"NHitX/I");
    fBPFTrackSignalTree->Branch("NHitY",&fNHitY,"NHitY/I");
    fBPFTrackSignalTree->Branch("CCNC",&fCCNC,"CCNC/I");
    fBPFTrackSignalTree->Branch("IntType",&fIntType,"IntType/I");
    fBPFTrackSignalTree->Branch("nuPDG",&fnuPDG,"nuPDG/I");



    // The background TTree...
    fBPFTrackBackgroundTree = tfs->make<TTree>("BPFTrackBackgroundTree","Tree for TMVA training with background track level BPF info");

    // input variables:
    fBPFTrackBackgroundTree->Branch("length",&fLength,"length/F");
    fBPFTrackBackgroundTree->Branch("chi2H",&fChi2H,"chi2H/F");
    fBPFTrackBackgroundTree->Branch("chi2A",&fChi2A,"chi2A/F");
    fBPFTrackBackgroundTree->Branch("chi2T",&fChi2T,"chi2T/F");
    fBPFTrackBackgroundTree->Branch("dEdxLL",&fdEdxLL,"dEdxLL/F");
    fBPFTrackBackgroundTree->Branch("hitRatio",&fHitRatio,"hitRatio/F");
    fBPFTrackBackgroundTree->Branch("lengthRatio",&fLengthRatio,"lengthRatio/F");
    fBPFTrackBackgroundTree->Branch("DChi2Tmupi",&fDChi2Tmupi,"DChi2Tmupi/F");
    fBPFTrackBackgroundTree->Branch("DChi2Tmupr",&fDChi2Tmupr,"DChi2Tmupr/F");
    fBPFTrackBackgroundTree->Branch("DChi2Tpipr",&fDChi2Tpipr,"DChi2Tpipr/F");
    

    // signal/background discriminator
    fBPFTrackBackgroundTree->Branch("trackPDG",&fTrackPDG,"trackPDG/F");

    // cut/weight variables:
    fBPFTrackBackgroundTree->Branch("E",&fE,"E/F");
    fBPFTrackBackgroundTree->Branch("MinX",&fMinX,"MinX/F");
    fBPFTrackBackgroundTree->Branch("MinY",&fMinY,"MinY/F");
    fBPFTrackBackgroundTree->Branch("MinZ",&fMinZ,"MinZ/F");
    fBPFTrackBackgroundTree->Branch("MaxX",&fMaxX,"MaxX/F");
    fBPFTrackBackgroundTree->Branch("MaxY",&fMaxY,"MaxY/F");
    fBPFTrackBackgroundTree->Branch("MaxZ",&fMaxZ,"MaxZ/F");
    fBPFTrackBackgroundTree->Branch("NHitX",&fNHitX,"NHitX/I");
    fBPFTrackBackgroundTree->Branch("NHitY",&fNHitY,"NHitY/I");
    fBPFTrackBackgroundTree->Branch("CCNC",&fCCNC,"CCNC/I");
    fBPFTrackBackgroundTree->Branch("IntType",&fIntType,"IntType/I");
    fBPFTrackBackgroundTree->Branch("nuPDG",&fnuPDG,"nuPDG/I");



    //
    // Book the Energy Estimator Training TTree
    //
    fBPFSliceTree = tfs->make<TTree>("BPFSliceTree","Tree for TMVA training with slice level BPF info");

    // target variables:
    fBPFSliceTree->Branch("E",&fE,"E/F");
    fBPFSliceTree->Branch("Esq",&fEsq,"Esq/F");
    
    // input variables:
    fBPFSliceTree->Branch("Pmu",&fPmu,"Pmu/F");
    fBPFSliceTree->Branch("DirZmu",&fDirZmu,"DirZmu/F");
    fBPFSliceTree->Branch("N3Dprongs",&fN3Dprongs,"N3Dprongs/F");
    fBPFSliceTree->Branch("E3Dprongs",&fE3Dprongs,"E3Dprongs/F");
    fBPFSliceTree->Branch("Eremain",&fEremain,"Eremain/F");
    fBPFSliceTree->Branch("SumPE",&fSumPE,"SumPE/F");

    // spectator variables:
    fBPFSliceTree->Branch("EhadMu",&fEhadMu,"EhadMu/F");
    
    // cut variables:
    fBPFSliceTree->Branch("MinX",&fMinX,"MinX/F");
    fBPFSliceTree->Branch("MinY",&fMinY,"MinY/F");
    fBPFSliceTree->Branch("MinZ",&fMinZ,"MinZ/F");
    fBPFSliceTree->Branch("MaxX",&fMaxX,"MaxX/F");
    fBPFSliceTree->Branch("MaxY",&fMaxY,"MaxY/F");
    fBPFSliceTree->Branch("MaxZ",&fMaxZ,"MaxZ/F");
    fBPFSliceTree->Branch("NHitX",&fNHitX,"NHitX/I");
    fBPFSliceTree->Branch("NHitY",&fNHitY,"NHitY/I");
    fBPFSliceTree->Branch("CCNC",&fCCNC,"CCNC/I");
    fBPFSliceTree->Branch("IntType",&fIntType,"IntType/I");
    fBPFSliceTree->Branch("nuPDG",&fnuPDG,"nuPDG/I");
    fBPFSliceTree->Branch("BestBPFmuPID",&fBestBPFmuPID,"BestBPFmuPID/F");
    fBPFSliceTree->Branch("BestRemid",&fBestRemid,"BestRemid/F");
    fBPFSliceTree->Branch("BestBPFmuPIDpdg",&fBestBPFmuPIDpdg,"BestBPFmuPIDpdg/I");



    //
    // Book histograms...
    //
    fPOT = tfs->make<TH1D>("fPOT","Total POT",1,0,1);
    
  }

  //......................................................................
  void BPFTmvaTrainer::endJob()
  {
    fPOT->Fill(0.5,fTotalPOT);
  }

  //......................................................................
  void BPFTmvaTrainer::endSubRun(const art::SubRun& sr)
  {
    art::Handle<sumdata::POTSum> pot;
    sr.getByToken(fPOTSumToken, pot);
    if(!pot.failedToGet()) fTotalPOT += pot->totgoodpot;
  }

  //......................................................................
  void BPFTmvaTrainer::resetVars()
  {
    //
    // PID variables:
    //
    fLength      = -5.0;
    fChi2H       = 1.0e9;
    fChi2A       = 1.0e9;
    fChi2T       = 1.0e9;
    fdEdxLL      = -1.0e9;
    fHitRatio    = -5.0;
    fLengthRatio = -5.0;
    fDChi2Tmupi  = 1.0e9;
    fDChi2Tmupr  = 1.0e9;
    fDChi2Tpipr  = 1.0e9;
    fTrackPDG    = 0;



    //
    // energy estimator variables:
    //
    fE         = -5.0;
    fEsq       = -5.0;
    fPmu       = -5.0;
    fDirZmu    = -5.0;
    fN3Dprongs = -5.0;
    fE3Dprongs = -5.0;
    fEremain   = -5.0;
    fSumPE     = -5.0;
    fEhadMu    = -5.0;
    fMinX      = 1.0e9;
    fMinY      = 1.0e9;
    fMinZ      = 1.0e9;
    fMaxX      = -1.0e9;
    fMaxY      = -1.0e9;
    fMaxZ      = -1.0e9;
    fNHitX     = -5;
    fNHitY     = -5;
    fCCNC      = -5;
    fIntType   = -5;
    fnuPDG     = 0;
    fBestBPFmuPID = -5.0;
    fBestRemid    = -5.0;
    fBestBPFmuPIDpdg = 0;

  }

  //......................................................................
  void BPFTmvaTrainer::setTruthVars(const std::vector<cheat::NeutrinoEffPur>& nuEP)
  {
    //
    // This function sets any variables for the TTree that come from truth.
    // Currently, the most pure neutrino is assumed to be the best
    // match.
    //

    float mostPure = 0.0;
    unsigned int match = 0;
    for(unsigned int i = 0; i < nuEP.size(); ++i) {
      if(nuEP[i].purity > mostPure && nuEP[i].neutrinoInt->NeutrinoSet()) {
        mostPure = nuEP[i].purity;
        match = i;
      }
    }
    
    if(mostPure > 0 && nuEP[match].neutrinoInt->NeutrinoSet()) {
      simb::MCNeutrino nu = nuEP[match].neutrinoInt->GetNeutrino();
      simb::MCParticle Nu = nu.Nu();
      fE       = Nu.E();
      fEsq     = fE*fE;
      fCCNC    = nu.CCNC();
      fIntType = nu.InteractionType();
      fnuPDG   = Nu.PdgCode(); 
    }
    
  }

  //......................................................................
  void BPFTmvaTrainer::analyze(const art::Event& evt)
  {

    //
    // General note on efficiency:
    //
    // This module has two training modes: PID and energy estimator.
    // Currently, all of the information for both training modes is
    // gathered and only the info needed for the specified training mode
    // is stored in the TTree output. This is inefficient but this module
    // is relatively fast anyway...

    //
    // NOTE: Currently for the muon PID training, we only use the variables
    //       fLength, fChi2T, fdEdxLL, and fHitRatio. All other variables were
    //       found to be less effective but further exploration in the future
    //       could be done to make them better.
    //

    // get the slices
    art::Handle< std::vector< rb::Cluster > > slices;
    evt.getByToken(fSlicerToken, slices);

    // create the FindMany object for getting prongs
    art::FindManyP<rb::Prong> prongFMP(slices, evt, art::InputTag(fProngLabel,fProngInstance));

    // create the FindMany object for getting Kalman tracks (used later for getting remid values)
    art::FindManyP<rb::Track> KtrackFMP(slices, evt, "kalmantrackmerge");

    // Make the required service handles...
    art::ServiceHandle<geo::Geometry> geom;
    art::ServiceHandle<calib::Calibrator> cal;
    art::ServiceHandle<cheat::BackTracker> BT;

    // Put the slices into a PtrVector and get the list neutrinos matched
    // to each slice from backtracker.
    art::PtrVector<rb::Cluster> ARTslices;
    for(unsigned int i = 0; i < slices->size(); ++i) {
      ARTslices.push_back(art::Ptr<rb::Cluster>(slices,i));
    }
    std::vector< std::vector< cheat::NeutrinoEffPur > > ALLnuEffPur;
    ALLnuEffPur = BT->SlicesToMCTruthsTable(ARTslices);


    
    // loop over all slices and get the fuzzyk 3D prongs
    for(unsigned int islice = 0; islice < slices->size(); ++islice) {

      // As usual, skip the noise slice...
      if((*slices)[islice].IsNoise()) continue;
      
      // NOTE: Currently, the fundamental object of analysis is the slice.
      // If in the future we change things so that the fundamental object
      // of analysis becomes the vertex, then we will have to get prongs
      // associated with the vertex instead.

      // Reset all TTree variables.
      this->resetVars();

      // Set slice level and truth variables for the TTrees.
      fMinX  = (*slices)[islice].MinX();
      fMinY  = (*slices)[islice].MinY();
      fMinZ  = (*slices)[islice].MinZ();
      fMaxX  = (*slices)[islice].MaxX();
      fMaxY  = (*slices)[islice].MaxY();
      fMaxZ  = (*slices)[islice].MaxZ();
      fNHitX = (*slices)[islice].NXCell();
      fNHitY = (*slices)[islice].NYCell();
      this->setTruthVars(ALLnuEffPur[islice]);



      // get the 3D prongs associated with this slice
      std::vector< art::Ptr< rb::Prong > > prongs;
      if(prongFMP.isValid()) prongs = prongFMP.at(islice);

      // create the FindMany object for getting tracks
      art::FindManyP<rb::Track> trackFMP(prongs, evt, fTrackLabel);



      // Variables used to keep track of the longest track (for the lengthRatio variable)
      // and for the track with the best muon PID.
      float longest   = -1.0; // Length of longest track.
      float bestPID   = -1.0; // best muonPID value
      int   prongPick = -1;   // Prong picked as having the longest track.



      //
      // For the LengthRatio variable, determine first the lenght of the
      // longest track.
      //
      for(unsigned int iprong = 0; iprong < prongs.size(); ++iprong) {

        std::vector< art::Ptr< rb::Track > > tracks;
        if(trackFMP.isValid()) tracks = trackFMP.at(iprong);
        art::FindManyP<rb::FitSum> fitsumFMP(tracks, evt, fTrackLabel);

        for(unsigned int itrack = 0; itrack < tracks.size(); ++itrack) {

          std::vector< art::Ptr< rb::FitSum > > fitsums;
          if(fitsumFMP.isValid()) fitsums = fitsumFMP.at(itrack);
          if(fitsums.size() != 1) abort();

          float len = tracks[itrack]->TotalLength();
          if(len > longest && fitsums[0]->PDG() == fPIDmode) {
            longest = len;
          }

        } // end loop over tracks
      } // end loop over prongs



      //////////
      //
      // Step 1:  Determine which prong is the most muon-like.
      //
      //////////

      // loop over all prongs and get the tracks
      for(unsigned int iprong = 0; iprong < prongs.size(); ++iprong) {

        // get the tracks associated with this prong
        std::vector< art::Ptr< rb::Track > > tracks;
        if(trackFMP.isValid()) tracks = trackFMP.at(iprong);

        // create the FindMany object for getting FitSum objects
        art::FindManyP<rb::FitSum> fitsumFMP(tracks, evt, fTrackLabel);

        // Keep track of the total chi2 values for each of the track fit assupmtions.
        float chi2Tmu = 1.0e9;
        float chi2Tpi = 1.0e9;
        float chi2Tpr = 1.0e9;

        // Loop over all tracks...
        for(unsigned int itrack = 0; itrack < tracks.size(); ++itrack) {

          // get the FitSum object associated with this track
          std::vector< art::Ptr< rb::FitSum > > fitsums;
          if(fitsumFMP.isValid()) fitsums = fitsumFMP.at(itrack);

          // There can be only one!!!
          // (something went really wrong with BPF if there's not one and
          // only one of these per track.)
          if(fitsums.size() != 1) abort();



          float LLtemp = 0.0;

          // fill in chi2 and NDOF values
          float tempH = 0.0;
          float tempA = 0.0;
          float tempN = 0.0;
          tempH = fitsums[0]->Chi2(0)+fitsums[0]->Chi2(1);   // chi2 from the hits
          tempA = fitsums[0]->Chi2(2)+fitsums[0]->Chi2(3);   // chi2 from the angles
          tempN = fitsums[0]->Ndof(0)+fitsums[0]->Ndof(1)-4; // NDOF (NOTE: currently 4 is NOT subtracted from NDOF when the fitsum object is made by BPF)

          // Keep track of the appropriate track fit info...
          if(fitsums[0]->PDG() == fPIDmode && tempN > 0.0) {
            fChi2H       = tempH/tempN;
            fChi2A       = tempA/tempN;
            fChi2T       = fChi2H+fChi2A;
            fLength      = tracks[itrack]->TotalLength();
            fLengthRatio = fLength/longest;
          }

          if(fitsums[0]->PDG() == 13   && tempN > 0.0) chi2Tmu = tempH/tempN+tempA/tempN;
          if(fitsums[0]->PDG() == 211  && tempN > 0.0) chi2Tpi = tempH/tempN+tempA/tempN;
          if(fitsums[0]->PDG() == 2212 && tempN > 0.0) chi2Tpr = tempH/tempN+tempA/tempN;

          //
          // Calculate dE/dx for each cell the track passes through using
          // the helper class.
          //
          fdEdxCalc.computeDEDX(tracks[itrack],fitsums[0]->PDG());
          
          std::vector<double> dE;
          std::vector<double> dx;
          std::vector<double> W;
          std::vector<double> BG;
          std::vector<double> s; // this is a throw away variable...

          fdEdxCalc.getResults(dE,dx,W,BG,s,fdxTOL);

          //
          // Compute the dE/dx likelihood values for all cells.
          //
          for(unsigned int a = 0; a < dE.size(); ++a) {

            unsigned int w = 0;
            unsigned int x = 0;

            if(W[a]  > 0.0) w = 1;
            if(dx[a] > 8.0) x = 1;

            // For the NearDet, we will just the "close to the readout end" FarDet histos.
            if(geom->DetId() == novadaq::cnv::kNEARDET) w = 1;

            double content = fdEdxVSbg[w][x]->GetBinContent(fdEdxVSbg[w][x]->FindBin(log(BG[a]),dE[a]/dx[a]));
            // We can't take the log of zero, so put a minimum on it.
            if(content < 1.0e-9) content = 1.0e-9;
            LLtemp += log(content);

          }

          if(dE.size() > 0) {
            // Keep track of the appropriate track fit info...
            if(fitsums[0]->PDG() == fPIDmode) {
              fdEdxLL = LLtemp/(float)dE.size();
              // NOTE: For this ratio, the number of hits on the track is taken to be the
              //       number of hits for which a value of dE/dx was successfully calculated.
              fHitRatio = (float)dE.size()/(float)prongs[iprong]->NCell();
            }
          }
          // end computing dE/dx LL info...
          


          // Set the variables for the track with the best muonPID
          if(fTrainMode == 1 && fitsums[0]->PDG() == fPIDmode) {
            if(fdEdxLL > -1.0e9 && fChi2T < 1.0e9) {
              // Get the PID value...
              float pid = fMuonPID.EvaluateMVA("KNN Method");
              if(pid > bestPID) {
                bestPID   = pid;
                prongPick = (int)iprong;
                fPmu      = fitsums[0]->FourMom().P();
                fDirZmu   = tracks[itrack]->Dir().Z();
              }
            }
          }
          
        } // end loop over tracks (itrack)



        // Set the delta-chi2 values. We will only do this if BPF succeeded in making all
        // three types of track (muon, pion, proton) otherwise these values won't make any
        // sense.
        //
        // NOTE: If in the future, BPF makes more than three types of track then this section
        //       will need to be adjusted!
        if(tracks.size() == 3) {
          fDChi2Tmupi = chi2Tmu - chi2Tpi;
          fDChi2Tmupr = chi2Tmu - chi2Tpr;
          fDChi2Tpipr = chi2Tpi - chi2Tpr;
        }


        // Determine the true particle type for this prong.
        const std::vector< const sim::Particle * > particles = BT->HitsToParticle(prongs[iprong]->AllCells());
        if(particles.size() > 0) fTrackPDG = particles[0]->PdgCode();

        // Fill the PID training TTree.
        if(fLength > 0.0 && fTrainMode == 0) {
          // WARNING to the user! We are treating particle and anti-particle the same here...
          if(std::abs(fTrackPDG) == fPIDmode) {
            fBPFTrackSignalTree->Fill();
          }
          else {
            fBPFTrackBackgroundTree->Fill();
          }
        }



      } // end loop over prongs (iprong)



      //////////
      //
      // Step 2:  Compute the values to be used for the energy estimator.
      //
      //////////

      //
      // Add up the energy of the non-muon 3D prongs.
      //

      float N3Dpr = 0.0;
      float E3Dpr = 0.0;
      float sumPE = 0.0;
      rb::Cluster  hits3Dprongs; // a container for all hits on 3D prongs

      for(unsigned int p = 0; p < prongs.size(); ++p) {
        
        // Skip 2D prongs.
        if(prongs[p]->Is2D()) continue;
        N3Dpr++;

        // Keep track of all hits in 3D prongs (to be used later when adding up
        // the remaining slice energy.)
        // NOTE: FuzzyK prongs can sometimes share hits, so we have to be careful
        //       not to double count.
        for(unsigned int h = 0; h < prongs[p]->NCell(); ++h) {

          const rb::CellHit *pronghit = prongs[p]->Cell(h).get();

          // Check to see if this hit is already in the list.
          bool onlist = false;
          for(unsigned int hpr = 0; hpr < hits3Dprongs.NCell(); ++hpr) {
            const rb::CellHit *pronghit3D = hits3Dprongs.Cell(hpr).get();
            if((*pronghit) == (*pronghit3D)) {
              onlist = true;
              break;
            }
          } // end loop over hits in 3D prongs (hpr)

          // Add the hit if its not already on the list.
          if(!onlist) hits3Dprongs.Add(prongs[p]->Cell(h));

        } // end loop over hits (h)

        // Skip this prong if it is the identified muon prong.
        if((int)p == prongPick) continue;

        // Loop over all hits in this 3D, non-muon prong and add up their energies.
        //
        // NOTE: For the hit W position, I am currently using the prong W instead of
        //       the track W. Since I haven't determined which of the three BPF fits
        //       is the best one for this prong, this saves me from having to try
        //       to guess at which track to use. My assumption is that prong->W will
        //       be good enough since I am lumping all hits from 3D prongs together
        //       into one 3D prong energy number...
        //
        // WARNING: Since hits can be shared by fuzzyK prongs, there is the potential
        //          below to double (or triple) count energy on shared cell hits. For
        //          now, I am assuming that this is a small effect but something
        //          smarter could/should be put into place in the future.
        for(unsigned int h = 0; h < prongs[p]->NCell(); ++h) {
          
          const rb::CellHit *chit = prongs[p]->Cell(h).get();
          float W = prongs[p]->W(chit);
          rb::RecoHit rhit(cal->MakeRecoHit(*chit,W));

          sumPE += chit->PE();
          if(rhit.IsCalibrated()) E3Dpr += rhit.GeV();

        } // end loop over hits (h)

      } // end loop over prongs (p)

      // Set the prong variables.
      fN3Dprongs = N3Dpr;
      fE3Dprongs = E3Dpr;



      //
      // Add up the remaining energy in the slice.
      //

      rb::Cluster hitsRemain; // a container for all hits NOT on 3D prongs

      // Loop over all hits in this slice.
      for(unsigned int hsl = 0; hsl < (*slices)[islice].NCell(); ++hsl) {

        const rb::CellHit *slicehit = (*slices)[islice].Cell(hsl).get();

        // Check to see if this hit is in the list of hits on 3D prongs.
        bool onlist = false;
        for(unsigned int hpr = 0; hpr < hits3Dprongs.NCell(); ++hpr) {
          const rb::CellHit *pronghit = hits3Dprongs.Cell(hpr).get();
          if((*pronghit) == (*slicehit)) {
            onlist = true;
            break;
          }
        } // end loop over hits in 3D prongs (hpr)

        // If this hit is not in a 3D prong, add it to the list of remaining slice hits.
        if(!onlist) {
          hitsRemain.Add((*slices)[islice].Cell(hsl));
          sumPE += slicehit->PE();
        }

      } // end loop over hits in this slice (hsl)

      // Set the remaining slice energy variables.
      if(hitsRemain.NCell() > 0) fEremain = hitsRemain.TotalGeV();
      else                       fEremain = 0.0;
      
      fSumPE = sumPE;



      //
      // Compute the hadronic contamination on the muon track and add that back in.
      //
      float muonTrackHadE = 0.0;

      // Only do this if we actually picked a track to be the muon.
      if(prongPick >= 0) {

        std::vector< art::Ptr< rb::Track > > tracks;
        if(trackFMP.isValid()) tracks = trackFMP.at(prongPick);
        art::FindManyP<rb::FitSum> fitsumFMP(tracks, evt, fTrackLabel);

        // Figure out which BPF track was made with the muon assumption.
        for(unsigned int itrack = 0; itrack < tracks.size(); ++itrack) {

          std::vector< art::Ptr< rb::FitSum > > fitsums;
          if(fitsumFMP.isValid()) fitsums = fitsumFMP.at(itrack);
          if(fitsums.size() != 1) abort();

          // Skip if not the muon track...
          if(fitsums[0]->PDG() != 13) continue;

          // If it is the muon track, do the extra Ehad calculation
          muonTrackHadE = fTrackCleanUpAlg.ExtraEOnTrackInGeV(*tracks[itrack],*ARTslices[islice]);

        } // end loop over tracks (itrack)

        fEhadMu = muonTrackHadE;

      } // end if(prongPick >= 0)

      fEremain += muonTrackHadE;



      //
      // Set the PID variables for this event:
      //

      // set the best BPF muon PID value
      fBestBPFmuPID = bestPID;

      // loop over Kalman tracks to find the best remid value
      if(KtrackFMP.isValid()) {

        std::vector< art::Ptr<rb::Track> > KalTracks = KtrackFMP.at(islice);
        art::FindManyP<remid::ReMId> remidFMP(KalTracks, evt, "remid");
        float bestRemid = -5.0;

        if(!remidFMP.isValid()) continue;

        for(unsigned int ikt = 0; ikt < KalTracks.size(); ++ikt) {
          std::vector< art::Ptr<remid::ReMId> > remids = remidFMP.at(ikt);
          // should only be one remid value per track
          if(remids.size() != 1) continue;
          if(remids[0]->Value() > bestRemid) bestRemid = remids[0]->Value();
        }

        fBestRemid = bestRemid;

      }
      


      // Record the pdg code for the particle that contributed the most energy
      // to the track with the highest BPF muon PID value.
      if(prongPick >= 0) {
        const std::vector< const sim::Particle * > particles = BT->HitsToParticle(prongs[prongPick]->AllCells());
        int pdg = 0;
        if(particles.size() > 0) pdg = abs(particles[0]->PdgCode());
        fBestBPFmuPIDpdg = pdg;
      }



      // If there is a neutrino match to this slice and we have at least one
      // 3D prong for which BPF made at least one track under the muon assumption,
      // fill the TTree...
      if(fE > 0.0 && fPmu > 0.0 && fTrainMode == 1) fBPFSliceTree->Fill();

    } // end loop over slices (islice)
    
    return;

  }

} // end namespace bpfit
////////////////////////////////////////////////////////////////////////



#include "art/Framework/Core/ModuleMacros.h"

namespace bpfit
{
  DEFINE_ART_MODULE(BPFTmvaTrainer)
}