[u/mrichter/AliRoot.git] / TRD / qaRec / AliTRDpidRefMakerNN.h

#ifndef ALITRDPIDREFMAKERNN_H
#define ALITRDPIDREFMAKERNN_H

//////////////////////////////////////////////////////
//
// Task to build PID reference tree for the training
// of neural networs for the TRD PID
//
// Authors: Alex Wilk    <wilka@uni-muenster.de>
//          Markus Heide <mheide@uni-muenster.de>
//
///////////////////////////////////////////////////////

#ifndef ALITRDRECOTASK_H
#include "AliTRDrecoTask.h"
#endif

#ifndef ALIPID_H
#include "AliPID.h"
#endif

#ifndef ALITRDCALPID_H
#include "../Cal/AliTRDCalPID.h"
#endif

#ifndef ALITRDGEOMETRY_H
#include "../AliTRDgeometry.h"
#endif

class TTree;
class TObjArray;
class TEventList;
class TMultiLayerPerceptron;
class AliPID;
class AliTRDtrackV1;
class AliTRDReconstructor;
class AliTRDpidRefMakerNN : public AliTRDrecoTask
{

public:
  enum  {
    k006  =  0
    ,k008 =  1
    ,k010 =  2
    ,k015 =  3
    ,k020 =  4
    ,k030 =  5
    ,k040 =  6
    ,k050 =  7
    ,k060 =  8
    ,k080 =  9
    ,k100 = 10
    ,kAll = 11
  };

  enum {
    kHistoPDG = 0
    ,kGraphTrain = 1
    ,kGraphTest = 2
  };

  enum {
    kMoniTrain = 50
  };

  AliTRDpidRefMakerNN();

  virtual ~AliTRDpidRefMakerNN();
  
  void    ConnectInputData(Option_t *opt);
  void    CreateOutputObjects();
  void    Exec(Option_t *option);
  Int_t   GetEpochs() {return fEpochs;};
  Int_t   GetMinTrain() {return fMinTrain;};
  Int_t   GetTrainMomBin() {return fTrainMomBin;};

  Bool_t  PostProcess();

  void    SetEpochs(Int_t epochs) {fEpochs = epochs;};
  void    SetMinTrain(Int_t mintrain) {fMinTrain = mintrain;};
  void    SetTrainMomBin(Int_t trainmombin) {fTrainMomBin = trainmombin;};
  void    SetDate(Int_t date) {fDate = date;};
  void    SetDoTraining(Bool_t train) {fDoTraining = train;};
  void    SetContinueTraining(Bool_t continTrain) {fContinueTraining = continTrain;};
  void    SetTrainPath(Int_t path) {fTrainPath = path;};
  void    LoadFile(const Char_t *InFileNN);

  void    Terminate(Option_t *);

  void    MakeTrainingLists();                                 // build the training and the test list
  void    MonitorTraining(Int_t mombin);                       // monitor training process
  void    LoadContainer(const Char_t *InFileCont);
  //void    CreateGraphs();

private:
  AliTRDpidRefMakerNN(const AliTRDpidRefMakerNN&);              // not implemented
  AliTRDpidRefMakerNN& operator=(const AliTRDpidRefMakerNN&);   // not implemented

  void GetV0info(AliTRDtrackV1 *TRDtrack, Float_t *v0pdg);  // get the v0 information
  void TrainNetworks(Int_t mombin);                         // train the neural networks for a given momentum bin

  AliTRDReconstructor *fReconstructor;     //! reconstructor needed for recalculation the PID
  TObjArray     *fV0s;                     //! v0 array
  TTree         *fNN;                      // NN data
  TEventList *fTrain[AliTRDCalPID::kNMom][AliTRDgeometry::kNlayer];          // Training list for each momentum 
  TEventList *fTest[AliTRDCalPID::kNMom][AliTRDgeometry::kNlayer];           // Test list for each momentum 
  TMultiLayerPerceptron *fNet[AliTRDgeometry::kNlayer]; // artificial neural network

  Int_t         fLayer;                    // TRD layer index 
  Int_t         fTrainMomBin;              // momentum bin for the training
  Int_t         fEpochs;                   // Number of epochs for the training of the NNs
  Int_t         fMinTrain;                 // minimum of events needed for training
  Int_t         fDate;                     // date stamp for training of the NNs
  Float_t       fMom;                      // momentum
  Float_t       *fdEdx[10];                // dEdx array
  Float_t       fv0pid[AliPID::kSPECIES];  // pid from v0s
  Bool_t        fDoTraining;               // checks if training will be done
  Bool_t        fContinueTraining;         // checks if training from an older run should be continued
  Int_t         fTrainPath;                // sets the path for continuing the training

  ClassDef(AliTRDpidRefMakerNN, 2); // TRD reference  maker for NN
};

#endif
Commit	Line	Data
72f5f463	1	#ifndef ALITRDPIDREFMAKERNN_H
	2	#define ALITRDPIDREFMAKERNN_H
	3
	4	//////////////////////////////////////////////////////
	5	//
	6	// Task to build PID reference tree for the training
	7	// of neural networs for the TRD PID
	8	//
	9	// Authors: Alex Wilk <wilka@uni-muenster.de>
	10	// Markus Heide <mheide@uni-muenster.de>
	11	//
	12	///////////////////////////////////////////////////////
	13
398f8388	14	#ifndef ALITRDRECOTASK_H
398f8388	15	#include "AliTRDrecoTask.h"
72f5f463	16	#endif
72f5f463	17
398f8388	18	#ifndef ALIPID_H
	19	#include "AliPID.h"
	20	#endif
	21
	22	#ifndef ALITRDCALPID_H
	23	#include "../Cal/AliTRDCalPID.h"
	24	#endif
	25
	26	#ifndef ALITRDGEOMETRY_H
	27	#include "../AliTRDgeometry.h"
	28	#endif
	29
	30	class TTree;
	31	class TObjArray;
	32	class TEventList;
72f5f463	33	class TMultiLayerPerceptron;
398f8388	34	class AliPID;
	35	class AliTRDtrackV1;
	36	class AliTRDReconstructor;
	37	class AliTRDpidRefMakerNN : public AliTRDrecoTask
72f5f463	38	{
398f8388	39
72f5f463	40	public:
398f8388	41	enum {
	42	k006 = 0
	43	,k008 = 1
	44	,k010 = 2
	45	,k015 = 3
	46	,k020 = 4
	47	,k030 = 5
	48	,k040 = 6
	49	,k050 = 7
	50	,k060 = 8
	51	,k080 = 9
	52	,k100 = 10
	53	,kAll = 11
	54	};
	55
	56	enum {
	57	kHistoPDG = 0
	58	,kGraphTrain = 1
	59	,kGraphTest = 2
72f5f463	60	};
72f5f463	61
398f8388	62	enum {
72f5f463	63	kMoniTrain = 50
	64	};
	65
	66	AliTRDpidRefMakerNN();
	67
	68	virtual ~AliTRDpidRefMakerNN();
	69
398f8388	70	void ConnectInputData(Option_t *opt);
72f5f463	71	void CreateOutputObjects();
398f8388	72	void Exec(Option_t *option);
72f5f463	73	Int_t GetEpochs() {return fEpochs;};
	74	Int_t GetMinTrain() {return fMinTrain;};
	75	Int_t GetTrainMomBin() {return fTrainMomBin;};
	76
	77	Bool_t PostProcess();
	78
	79	void SetEpochs(Int_t epochs) {fEpochs = epochs;};
	80	void SetMinTrain(Int_t mintrain) {fMinTrain = mintrain;};
	81	void SetTrainMomBin(Int_t trainmombin) {fTrainMomBin = trainmombin;};
	82	void SetDate(Int_t date) {fDate = date;};
	83	void SetDoTraining(Bool_t train) {fDoTraining = train;};
	84	void SetContinueTraining(Bool_t continTrain) {fContinueTraining = continTrain;};
	85	void SetTrainPath(Int_t path) {fTrainPath = path;};
398f8388	86	void LoadFile(const Char_t *InFileNN);
72f5f463	87
398f8388	88	void Terminate(Option_t *);
72f5f463	89
398f8388	90	void MakeTrainingLists(); // build the training and the test list
72f5f463	91	void MonitorTraining(Int_t mombin); // monitor training process
398f8388	92	void LoadContainer(const Char_t *InFileCont);
398f8388	93	//void CreateGraphs();
72f5f463	94
	95	private:
	96	AliTRDpidRefMakerNN(const AliTRDpidRefMakerNN&); // not implemented
	97	AliTRDpidRefMakerNN& operator=(const AliTRDpidRefMakerNN&); // not implemented
	98
398f8388	99	void GetV0info(AliTRDtrackV1 TRDtrack, Float_t v0pdg); // get the v0 information
398f8388	100	void TrainNetworks(Int_t mombin); // train the neural networks for a given momentum bin
72f5f463	101
398f8388	102	AliTRDReconstructor *fReconstructor; //! reconstructor needed for recalculation the PID
	103	TObjArray *fV0s; //! v0 array
	104	TTree *fNN; // NN data
	105	TEventList *fTrain[AliTRDCalPID::kNMom][AliTRDgeometry::kNlayer]; // Training list for each momentum
	106	TEventList *fTest[AliTRDCalPID::kNMom][AliTRDgeometry::kNlayer]; // Test list for each momentum
72f5f463	107	TMultiLayerPerceptron *fNet[AliTRDgeometry::kNlayer]; // artificial neural network
72f5f463	108
398f8388	109	Int_t fLayer; // TRD layer index
72f5f463	110	Int_t fTrainMomBin; // momentum bin for the training
	111	Int_t fEpochs; // Number of epochs for the training of the NNs
	112	Int_t fMinTrain; // minimum of events needed for training
	113	Int_t fDate; // date stamp for training of the NNs
398f8388	114	Float_t fMom; // momentum
	115	Float_t *fdEdx[10]; // dEdx array
	116	Float_t fv0pid[AliPID::kSPECIES]; // pid from v0s
72f5f463	117	Bool_t fDoTraining; // checks if training will be done
	118	Bool_t fContinueTraining; // checks if training from an older run should be continued
	119	Int_t fTrainPath; // sets the path for continuing the training
	120
398f8388	121	ClassDef(AliTRDpidRefMakerNN, 2); // TRD reference maker for NN
72f5f463	122	};
	123
	124	#endif