added selection on existence of extrapolation at calorimeter

MuonColliderSoft · Jul 16, 2024 · e44c841 · e44c841
1 parent 84676e1
commit e44c841
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/source/Utils/include/FilterTracks.h b/source/Utils/include/FilterTracks.h
@@ -61,6 +61,7 @@ class FilterTracks : public marlin::Processor
    std::string _OutputTrackCollection {};
 
    bool _BarrelOnly = false;
+   bool _HasCaloState = false;
 
    //! Cut off for total number of hits
    int _NHitsTotal = 7;

diff --git a/source/Utils/src/FilterTracks.cc b/source/Utils/src/FilterTracks.cc
@@ -25,6 +25,12 @@ FilterTracks::FilterTracks()
 			     _BarrelOnly
 			      );  
 
+  registerProcessorParameter("HasCaloState",
+		  	     "If true, just keep tracks that have a TrackState at the Calorimeter surface",
+			     _HasCaloState,
+			     _HasCaloState
+			      );  
+
   registerProcessorParameter("NHitsTotal",
 		  	     "Minimum number of hits on track",
 			     _NHitsTotal,
@@ -140,6 +146,17 @@ void FilterTracks::processEvent( LCEvent * evt )
 
     int nholes = trk->getNholes();
 
+    bool foundCaloState = false;
+    // Check if a TrackState at the calo surface exists
+    const std::vector<EVENT::TrackState*>& trackStates = trk->getTrackStates();
+    for (const auto& state : trackStates) {
+      if (state->getLocation() == EVENT::TrackState::AtCalorimeter) {
+        foundCaloState = true;
+        break;
+      }
+    }
+    if (_HasCaloState && !foundCaloState) { continue; }
+
     if(_BarrelOnly == true) {
       bool endcaphits = false;
       for(int j=0; j<nhittotal; ++j) {