DRAFT: Music10 tev

CMakeLists.txt

@@ -68,8 +68,8 @@ find_package( DD4hep REQUIRED COMPONENTS DDRec)
 set(CMAKE_MODULE_PATH  ${CMAKE_MODULE_PATH}  ${DD4hep_ROOT}/cmake )
 include( DD4hep )
 
-find_package( ROOT REQUIRED )
-set( ROOT_COMPONENT_LIBRARIES Geom Reflex)
+#find_package( ROOT REQUIRED )
+#set( ROOT_COMPONENT_LIBRARIES Geom Reflex TMVA)
 
 if(DD4HEP_USE_XERCESC)
   find_package(XercesC)
@@ -78,8 +78,7 @@ include(DD4hep_XML_setup)
 INCLUDE_DIRECTORIES( BEFORE SYSTEM ${DD4hep_INCLUDE_DIRS} )
 LINK_LIBRARIES( ${DD4hep_LIBRARIES} ${DD4hep_COMPONENT_LIBRARIES}  )
 
-#FIND_PACKAGE( ROOT REQUIRED ) 
-
+FIND_PACKAGE( ROOT REQUIRED COMPONENTS Geom TMVA) 
 INCLUDE_DIRECTORIES( SYSTEM ${ROOT_INCLUDE_DIRS} )
 LINK_LIBRARIES( ${ROOT_LIBRARIES} )
 ADD_DEFINITIONS( ${ROOT_DEFINITIONS} )

source/Utils/include/FilterTracks.h

@@ -16,7 +16,13 @@ namespace TrackPerf
  *  @parameter NHitsVertex Minimum number of hits on vertex detector
  *  @parameter NHitsInner Minimum number of hits on inner tracker
  *  @parameter NHitsOuter Minimum number of hits on outer tracker
+ *  @parameter MaxOutliers Maximum number of outliers hits on track
+ *  @parameter MaxHoles Maximum number of holes on track
+ *  @parameter MinNdf Minimum value for ndf
  *  @parameter MinPt Minimum transverse momentum
+ *  @parameter MaxPt Max transverse momentum
+ *  @parameter MinTheta Minimum theta
+ *  @parameter MaxTheta Max theta
  *  @parameter Chi2Spatial Spatial chi squared
  *
  * @author N. Bruhwiler
@@ -70,14 +76,33 @@ class FilterTracks : public marlin::Processor
    int _NHitsInner = 2;
    //! Cut off for number of hits in outer tracker (barrel and endcap combined)
    int _NHitsOuter = 1;
+   //! Cut off for number of holes
+   int _MaxHoles = 10;
 
    //! Cut off for momentum (GeV)
-   float _MinPt = 1.0;   //units GeV
+   float _MinPt = 0.5;    //units GeV
+   float _MaxPt = 1000.0; //units GeV
+
+   //! Cut off for theta (rad)
+   float _MinTheta = 0;    
+   float _MaxTheta = 3.14; 
+
+   //! Cut off for the value ndf
+   int _MinNdf = 1;
+
+   //! Cut off for outliers number   
+   int _MaxOutl = 10;  
 
    //! Cut off for spatial and temporal chi squared values
    float _Chi2Spatial = 0;
 
+   //! NN parameters
+   std::string _NNmethod = ""; // if defined apply the NN
+   std::string _NNweights = ""; // xml file with weights
+   std::vector<std::string> _NNvars; // sorted list of variables used by NN
+   float _NNthr = 0; // NN threshold   
+
    //! Default magnetic field value (Tesla)
-   float _Bz = 3.57;
+   float _Bz = 5.0;
 };
 

source/Utils/src/FilterTracks.cc

@@ -7,11 +7,18 @@
 #include <EVENT/Track.h>
 #include <EVENT/TrackerHit.h>
 #include <IMPL/LCCollectionVec.h>
+#include <IMPL/LCFlagImpl.h>
+#include <IMPL/TrackImpl.h>
 #include <UTIL/CellIDDecoder.h>
 #include <UTIL/LCTrackerConf.h>
 
+#include "TMVA/Reader.h"
+
+#include <iostream> 
+
 FilterTracks aFilterTracks ;
 
+
 FilterTracks::FilterTracks()
   : Processor("FilterTracks")
 {
@@ -55,12 +62,79 @@ FilterTracks::FilterTracks()
 			     _MinPt
 		 	      );
 
+  registerProcessorParameter("MaxPt",
+             "Max transverse momentum",
+           _MaxPt,
+           _MaxPt
+            );
+
+  registerProcessorParameter("MinTheta",
+             "Minimum theta",
+           _MinTheta,
+           _MinTheta
+            );
+
+  registerProcessorParameter("MaxTheta",
+             "Max theta",
+           _MaxTheta,
+           _MaxTheta
+            );
+
   registerProcessorParameter("Chi2Spatial",
-		  	     "Spatial chi squared",
+		  	     "iMinimum value for Spatial chi squared",
 			     _Chi2Spatial,
 			     _Chi2Spatial
 		 	      );
 
+  registerProcessorParameter("MinNdf",
+             "Minimum value for ndf",
+           _MinNdf,
+           _MinNdf
+            );
+
+  registerProcessorParameter("MaxHoles",
+             "Max number of holes",
+           _MaxHoles,
+           _MaxHoles
+            );
+
+  registerProcessorParameter("MaxOutliers",
+             "Max number of outliers",
+           _MaxOutl,
+           _MaxOutl
+          );
+
+  registerProcessorParameter("Bz",
+             "Magnetic field in Tesla (default: 5)",
+           _Bz,
+           _Bz
+            );
+
+  registerProcessorParameter("NNmethod",
+           "Name of the NN method, if empty uses standard cuts",
+           _NNmethod,
+           std::string("")
+            );
+
+  registerProcessorParameter("NNweights",
+             "File xml with the weights of the NN",
+           _NNweights,
+           std::string("")
+            );
+
+  registerProcessorParameter("NNvars",
+             "Sorted list with the names of the variables used in the training"
+              "Supported variables are: trtvhn, trtihn, trtohn, trthn, trtnh, trch2, trndf",
+           _NNvars,
+           _NNvars
+            );
+
+  registerProcessorParameter("NNthr",
+             "NN threshold (Takes tracks with prediction > NNthr)",
+           _NNthr,
+           _NNthr
+            );
+
   registerInputCollection( LCIO::TRACK,
 		  	   "InputTrackCollectionName" ,
 			   "Name of the input collection",
@@ -81,7 +155,8 @@ void FilterTracks::init()
 {
   // Print the initial parameters
   printParameters() ;
-  buildBfield() ;
+  // it requires that geometry has been already initialized
+  // buildBfield() ;
 }
 
 void FilterTracks::processRunHeader( LCRunHeader* /*run*/)
@@ -106,57 +181,110 @@ void FilterTracks::processEvent( LCEvent * evt )
 {
   // Make the output track collection
   LCCollectionVec *OutputTrackCollection = new LCCollectionVec(LCIO::TRACK);
-  OutputTrackCollection->setSubset(true);
+  // do not store pointers but real tracks
+  OutputTrackCollection->setSubset(false);
 
-  // Get input collection
-  LCCollection* InputTrackCollection  =evt->getCollection(_InputTrackCollection);
+  // if we want to point back to the hits we need to set the flag
+  LCFlagImpl trkFlag(0);
+  trkFlag.setBit(LCIO::TRBIT_HITS);
+  OutputTrackCollection->setFlag(trkFlag.getFlag());
 
+  TMVA::Reader* reader = new TMVA::Reader();
+
+  // Get input collection
+  LCCollection* InputTrackCollection = evt->getCollection(_InputTrackCollection);
+
   if( InputTrackCollection->getTypeName() != lcio::LCIO::TRACK )
     { throw EVENT::Exception( "Invalid collection type: " + InputTrackCollection->getTypeName() ) ; }
 
   // Filter
   std::string encoderString = lcio::LCTrackerCellID::encoding_string();
   UTIL::CellIDDecoder<lcio::TrackerHit> decoder(encoderString);
 
+  std::unordered_map<std::string, float> vars = {
+    {"trtvhn", 0},
+    {"trtihn", 0},
+    {"trtohn", 0},
+    {"trthn", 0},
+    {"trtnh", 0},
+    {"trch2", 0},
+    {"trndf", 0},
+  };
+
+  if ( not _NNmethod.empty() ) {
+
+    for (unsigned i=0,N=_NNvars.size(); i<N; ++i) {
+      std::string name = _NNvars[i];
+      if (vars.find(name) == vars.end())
+        throw EVENT::Exception( "NN variable not supported: " + name ) ;  
+      reader->AddVariable( name, &vars[name] );
+    }
+    reader->BookMVA(_NNmethod, _NNweights);
+  }
+
   for(int i=0; i<InputTrackCollection->getNumberOfElements(); i++) {
     EVENT::Track *trk=dynamic_cast<EVENT::Track*>(InputTrackCollection->getElementAt(i));
 
-    int nhittotal  = trk->getTrackerHits().size();
+    vars["trtnh"]  = trk->getTrackerHits().size();
 
     const EVENT::IntVec& subdetectorHits = trk->getSubdetectorHitNumbers();
-    int nhitvertex = subdetectorHits[1]+subdetectorHits[2];
-    int nhitinner = subdetectorHits[3]+subdetectorHits[4];
-    int nhitouter = subdetectorHits[5]+subdetectorHits[6];
+    vars["trtvhn"] = subdetectorHits[1]+subdetectorHits[2];
+    vars["trtihn"] = subdetectorHits[3]+subdetectorHits[4];
+    vars["trtohn"] = subdetectorHits[5]+subdetectorHits[6];
 
+    vars["trthn"] = trk->getNholes();
     float pt = fabs(0.3*_Bz/trk->getOmega()/1000);
+    float theta = M_PI_2-atan(trk->getTanLambda());
+    vars["trch2"] = trk->getChi2();
 
-    float chi2spatial = trk->getChi2();
+    vars["trndf"] = trk->getNdf();
 
     if(_BarrelOnly == true) {
       bool endcaphits = false;
-      for(int j=0; j<nhittotal; ++j) {
-	//Find what subdetector the hit is on 
-	uint32_t systemID = decoder(trk->getTrackerHits()[j])["system"];
-	if(systemID == 2 || systemID == 4 || systemID == 6) {
-	  endcaphits = true;
-	  break;
-	}
+      for(int j=0; j<vars["trtnh"]; ++j) {
+	      //Find what subdetector the hit is on 
+	      uint32_t systemID = decoder(trk->getTrackerHits()[j])["system"];
+	      if(systemID == 2 || systemID == 4 || systemID == 6) {
+	        endcaphits = true;
+	        break;
+	      }
+      }
+      if (endcaphits == false) { 
+        auto itrk = dynamic_cast<IMPL::TrackImpl*>(trk);
+        OutputTrackCollection->addElement(new IMPL::TrackImpl(*itrk)); 
       }
-      if(endcaphits == false) { OutputTrackCollection->addElement(trk); }
     } else { // track property cuts
-      if(nhittotal    > _NHitsTotal  &&
-	 nhitvertex   > _NHitsVertex &&
-	 nhitinner    > _NHitsInner  &&
-	 nhitouter    > _NHitsOuter  &&
-	 pt           > _MinPt       &&
-	 chi2spatial  > _Chi2Spatial)
-	{ OutputTrackCollection->addElement(trk); }
+      if ( not _NNmethod.empty() ) { // NN cuts
+        auto prediction = reader->EvaluateMVA(_NNmethod);
+        if ( prediction > _NNthr ) {
+          auto itrk = dynamic_cast<IMPL::TrackImpl*>(trk);
+          OutputTrackCollection->addElement(new IMPL::TrackImpl(*itrk)); 
+        }
+      } else { // user cuts
+        if (vars["trtnh"]  > _NHitsTotal  &&
+	          vars["trtvhn"] > _NHitsVertex &&
+	          vars["trtihn"] > _NHitsInner  &&
+	          vars["trtohn"] > _NHitsOuter  &&
+	          pt             > _MinPt       &&
+            pt             < _MaxPt       &&
+            theta          > _MinTheta    &&
+            theta          < _MaxTheta    &&
+	          vars["trch2"]  > _Chi2Spatial &&
+            vars["trndf"]  > _MinNdf      &&
+            vars["trtnh"]-vars["trndf"]/2 < _MaxOutl &&
+            vars["trthn"]  < _MaxHoles)
+	          { 
+              auto itrk = dynamic_cast<IMPL::TrackImpl*>(trk);
+              OutputTrackCollection->addElement(new IMPL::TrackImpl(*itrk)); 
+            }
+      }
     }
   }
 
   // Save output track collection
   evt->addCollection(OutputTrackCollection, _OutputTrackCollection);  
+  delete reader;
 }
 
 void FilterTracks::end()
-{ }
+{  }