CbmTrdElectronsTrainAnn.h

Go to the documentation of this file.

 
#ifndef CBM_TRD_ELECTRONS_TRAIN_ANN
#define CBM_TRD_ELECTRONS_TRAIN_ANN
 
#include "FairTask.h"
 
#include "TMVA/Factory.h"
#include "TMVA/Reader.h"
#include "TMultiLayerPerceptron.h"
 
#include <string>
#include <vector>
 
class TH1D;
class TTree;
class TGraph;
class TRandom;
//class TMultiLayerPerceptron;
 
enum IdMethod {
  kANN        = 0,
  kBDT        = 1,
  kMEDIAN     = 2,
  kLIKELIHOOD = 3,
  kMeanCut    = 4
};
 
class TrdEloss {
public:
  TrdEloss() : TrdEloss(0., 0., 0.) {}
 
  TrdEloss(double eloss, double dedx = 0., double tr = 0.)
    : fEloss(eloss), fdEdX(dedx), fTR(tr) {}
 
  double fEloss;
  double fdEdX;
  double fTR;
};
 
 
class CbmTrdElectronsTrainAnn : public FairTask {
public:
  CbmTrdElectronsTrainAnn(int nofTrdLayers);
 
  virtual ~CbmTrdElectronsTrainAnn();
 
  virtual InitStatus Init();
 
  virtual void Exec(Option_t* opt);
 
  virtual void Finish();
 
  void Draw(Option_t* = "");
 
private:
  void FillElossVectorSim();
 
  void FillElossVectorReal();
 
  void FillElossHist();
 
  void SortElossAndFillHist();
 
  void Run();
 
  void DoTrain();
 
  void DoTest();
 
  void DoPreTest();
 
  TTree* CreateTree();
 
  std::string CreateAnnString();
 
  TMVA::Factory* CreateFactory(TTree* simu);
 
  TMVA::Reader* CreateTmvaReader();
 
  void Transform();
 
  void Transform1();
 
  void Transform2();
 
  Double_t FindOptimalCut();
 
  void CreateCumProbOutputHist();
 
  TGraph* CreateRocDiagramm();
 
  Double_t Likelihood();
 
  Double_t Median();
 
  Double_t MeanCut();
 
  Double_t Eval(Bool_t isEl);
 
public:
  void SetIsDoTrain(Bool_t doTrain) { fIsDoTrain = doTrain; }
 
  void SetTransformType(Int_t type) { fTransformType = type; }
 
  void SetIdMethod(IdMethod idMethod) { fIdMethod = idMethod; }
 
  void SetOutputDir(const std::string& outputDir) { fOutputDir = outputDir; }
 
  void SetNofAnnEpochs(Int_t nofAnnEpochs) { fNofAnnEpochs = nofAnnEpochs; }
 
  void SetNofTrainSamples(Int_t nofTr) { fNofTrainSamples = nofTr; }
 
  void SetSigmaError(Double_t sigma) { fSigmaError = sigma; }
 
  void SetBeamDataFile(std::string beamDataFile) {
    fBeamDataFile = beamDataFile;
  }
  void SetBeamDataPiHist(std::string beamDataPiHist) {
    fBeamDataPiHist = beamDataPiHist;
  }
  void SetBeamDataElHist(std::string beamDataElHist) {
    fBeamDataElHist = beamDataElHist;
  }
 
  void RunBeamData();
 
private:
  void FillAnnInputHist(Bool_t isEl);
 
  // Input arrays of simulated events
  TClonesArray* fMCTracks;
  TClonesArray* fTrdPoints;
  TClonesArray* fTrdTracks;
  TClonesArray* fTrdTrackMatches;
  TClonesArray* fTrdHits;
 
  // store energy loss information from TrdTrack
  // 1st index -> [0] = electrons, [1] = pions
  // 2nd index -> track index
  // 3rd index -> hit index in track
  std::vector<std::vector<std::vector<TrdEloss>>> fEloss;
 
  std::vector<TH1*> fHists;  //store all pointers to histograms
 
  TH1*
    fhResults;  // histograms for the results storing: pi suppression, el efficiency etc.
 
  // [0] = electrons, [1] = pions
  std::vector<TH1*>
    fhMeanEloss;  // sum of energy losses in all layers divided by number of layers
  std::vector<TH1*> fhEloss;  // energy losses in one layer
 
  // store sorted energy losses and cumulative probability for each TRD layer
  // 1st index -> [0] = electrons, [1] = pions
  // 2nd index -> layer number
  std::vector<std::vector<TH1*>> fhElossSort;
 
  Int_t fEventNum;         // event number
  std::string fOutputDir;  // output directory
  Double_t fSigmaError;  // additional sigma error for energy loss measurements
  Bool_t fIsDoTrain;     // do you want to run training procedure?
  Int_t fTransformType;  // Energy loss transformation type
 
  std::string
    fBeamDataFile;  // path to file with beamtime data for energy losses
  std::string fBeamDataPiHist;  // histogram name with energy losses for pions
  std::string
    fBeamDataElHist;  // histogram name with energy losses for electrons
 
  std::vector<Float_t> fAnnInput;  // input vector for ANN
  Float_t fXOut;                   // output value from ANN
 
  Int_t fNofTrdLayers;  // number of TRD layers
 
  Double_t fMaxEval;
  Double_t fMinEval;
 
  TMultiLayerPerceptron* fNN;
  TMVA::Reader* fReader;
  IdMethod fIdMethod;
  Int_t fNofAnnEpochs;
  Int_t fNofTrainSamples;
  double fElIdEfficiency;
 
  // Histograms for testing
  // [0] = electron, [1] = pion
  std::vector<TH1*> fhOutput;  // algorithm output
  std::vector<TH1*>
    fhCumProbOutput;  // Cumulative probabilities for algorithm's output
  std::vector<std::vector<TH1*>>
    fhInput;  // Input data for algorithm for each input
 
  CbmTrdElectronsTrainAnn(const CbmTrdElectronsTrainAnn&);
  CbmTrdElectronsTrainAnn& operator=(const CbmTrdElectronsTrainAnn&);
 
  ClassDef(CbmTrdElectronsTrainAnn, 1);
};
 
#endif