From 81ed6d340c665a7a46d9cd77658f553852ff5099 Mon Sep 17 00:00:00 2001
From: Chengdong Fu <fucd@ihep.ac.cn>
Date: Wed, 19 Aug 2020 15:47:19 +0800
Subject: [PATCH] add TrackSubset
---
.../SiliconTracking/src/TrackSubsetAlg.cpp | 338 ++++++++++++++++++
.../SiliconTracking/src/TrackSubsetAlg.h | 117 ++++++
2 files changed, 455 insertions(+)
create mode 100644 Reconstruction/SiliconTracking/src/TrackSubsetAlg.cpp
create mode 100644 Reconstruction/SiliconTracking/src/TrackSubsetAlg.h
diff --git a/Reconstruction/SiliconTracking/src/TrackSubsetAlg.cpp b/Reconstruction/SiliconTracking/src/TrackSubsetAlg.cpp
new file mode 100644
index 00000000..8961484d
--- /dev/null
+++ b/Reconstruction/SiliconTracking/src/TrackSubsetAlg.cpp
@@ -0,0 +1,338 @@
+#include "TrackSubsetAlg.h"
+
+#include "GearSvc/IGearSvc.h"
+#include "TrackSystemSvc/ITrackSystemSvc.h"
+#include "DataHelper/Navigation.h"
+
+#include <UTIL/ILDConf.h>
+
+#include <gear/BField.h>
+
+#include "KiTrack/SubsetSimple.h"
+#include "KiTrack/SubsetHopfieldNN.h"
+#include "Tools/Fitter.h"
+#include "Tools/KiTrackMarlinTools.h"
+
+#include "TrackSystemSvc/MarlinTrkUtils.h"
+
+using namespace KiTrack;
+
+DECLARE_COMPONENT(TrackSubsetAlg)
+
+TrackSubsetAlg::TrackSubsetAlg(const std::string& name, ISvcLocator* svcLoc)
+ : GaudiAlgorithm(name, svcLoc){
+
+ // modify processor description
+ //_description = "TrackSubsetAlg takes tracks from multiple sources and outputs them (or modified versions, or a subset of them) as one track collection." ;
+
+ //std::vector< std::string > trackInputColNamesDefault;
+ //trackInputColNamesDefault.push_back( "ForwardTracks" );
+ //trackInputColNamesDefault.push_back( "SiTracks" );
+
+ declareProperty("TrackSubsetCollection", _outColHdl, "Handle of the SiTrack output collection");
+}
+
+StatusCode TrackSubsetAlg::initialize() {
+
+ debug() << " init called " << endmsg;
+
+ _nRun = 0 ;
+ _nEvt = 0 ;
+
+ for(unsigned i=0; i<_trackInputColNames.size(); i++){
+ _inTrackColHdls.push_back(new DataHandle<edm4hep::TrackCollection> (_trackInputColNames[i], Gaudi::DataHandle::Reader, this));
+ }
+
+ for(unsigned i=0; i<_trackerHitInputColNames.size(); i++){
+ _inTrackerHitColHdls.push_back(new DataHandle<edm4hep::TrackerHitCollection> (_trackerHitInputColNames[i], Gaudi::DataHandle::Reader, this));
+ }
+ /**********************************************************************************************/
+ /* Initialise the MarlinTrkSystem, needed by the tracks for fitting */
+ /**********************************************************************************************/
+ auto _gear = service<IGearSvc>("GearSvc");
+ if ( !_gear ) {
+ error() << "Failed to find GearSvc ..." << endmsg;
+ return StatusCode::FAILURE;
+ }
+ gear::GearMgr* gearMgr = _gear->getGearMgr();
+ _bField = gearMgr->getBField().at( gear::Vector3D( 0.,0.,0.) ).z() ;
+
+ // set upt the geometry
+ auto _trackSystemSvc = service<ITrackSystemSvc>("TrackSystemSvc");
+ if ( !_trackSystemSvc ) {
+ error() << "Failed to find TrackSystemSvc ..." << endmsg;
+ return StatusCode::FAILURE;
+ }
+ _trkSystem = _trackSystemSvc->getTrackSystem();
+
+ if( _trkSystem == 0 ){
+ error() << "Cannot initialize MarlinTrkSystem of Type: KalTest" <<endmsg;
+ return StatusCode::FAILURE;
+ }
+
+ // set the options
+ _trkSystem->setOption( MarlinTrk::IMarlinTrkSystem::CFG::useQMS, _MSOn ) ; //multiple scattering
+ _trkSystem->setOption( MarlinTrk::IMarlinTrkSystem::CFG::usedEdx, _ElossOn) ; //energy loss
+ _trkSystem->setOption( MarlinTrk::IMarlinTrkSystem::CFG::useSmoothing, _SmoothOn) ; //smoothing
+
+ // initialise the tracking system
+ //_trkSystem->init() ;
+
+ return GaudiAlgorithm::initialize();
+}
+
+StatusCode TrackSubsetAlg::finalize(){
+ for(unsigned i=0; i<_inTrackColHdls.size(); i++){
+ delete _inTrackColHdls[i];
+ }
+ _inTrackColHdls.clear();
+
+ for(unsigned i=0; i<_inTrackerHitColHdls.size(); i++){
+ delete _inTrackerHitColHdls[i];
+ }
+ _inTrackerHitColHdls.clear();
+
+ return GaudiAlgorithm::finalize();
+}
+
+StatusCode TrackSubsetAlg::execute(){
+ std::vector<edm4hep::Track> tracks;
+
+ /**********************************************************************************************/
+ /* Read in the collections */
+ /**********************************************************************************************/
+ debug() << "Try to load " << _trackInputColNames.size() << " input track collections" << endmsg;
+
+ unsigned nTrackLoaded = 0;
+
+ for( unsigned i=0; i < _inTrackColHdls.size(); i++ ){
+ const edm4hep::TrackCollection* trackCol = nullptr;
+ try {
+ trackCol = _inTrackColHdls[i]->get();
+ }
+ catch ( GaudiException &e ) {
+ debug() << "Collection " << _inTrackColHdls[i]->fullKey() << " is unavailable in event " << _nEvt << endmsg;
+ continue;
+ }
+
+ if(trackCol){
+ int nTracks = trackCol->size();
+ debug() << "Load track input collection " << _trackInputColNames[i] << " with " << nTracks << " tracks" << endmsg;
+
+ for(auto track : *trackCol){
+ tracks.push_back( track );
+ nTrackLoaded++;
+ }
+ }
+ else error() << "track input collection " << _trackInputColNames[i] << " could not be found, but not throw exception!!!" << endmsg;
+ }
+
+ debug() << "Loaded all in all " << nTrackLoaded << " tracks, which will now get further processed" << endmsg;
+ if(tracks.size()==0) return StatusCode::SUCCESS;
+
+ Navigation::Instance()->Initialize();
+ for(unsigned i=0; i < _inTrackerHitColHdls.size(); i++){
+ try {
+ auto trackCol = _inTrackerHitColHdls[i]->get();
+ Navigation::Instance()->AddTrackerHitCollection(trackCol);
+ }
+ catch ( GaudiException &e ) {
+ debug() << "Collection " << _inTrackerHitColHdls[i]->fullKey() << " is unavailable in event " << _nEvt << endmsg;
+ continue;
+ }
+ }
+ /**********************************************************************************************/
+ /* Make sure that all tracks are compatible: find the best subset */
+ /**********************************************************************************************/
+ debug() << "Find the best subset of tracks using the Hopfield Neural Network" << endmsg;
+
+ TrackQI trackQI( _trkSystem );
+
+ debug() << "The tracks and their qualities (and their hits ): " << endmsg;
+
+ std::vector<edm4hep::Track*> tracks_p;
+ for( unsigned i=0; i < tracks.size(); i++ ){
+ edm4hep::Track* track = &tracks[i];
+ tracks_p.push_back(track);
+ double qi = trackQI( track );
+ debug() << "Track " << track->id() << " address " << track << "\t" << qi << "( ";
+ std::vector<edm4hep::ConstTrackerHit> hits;
+ std::copy(track->trackerHits_begin(), track->trackerHits_end(), std::back_inserter(hits));
+
+ std::sort( hits.begin(), hits.end(), KiTrackMarlin::compare_TrackerHit_z );
+
+ for( unsigned j=0; j<hits.size(); j++ ){
+ debug() << hits[j].id() << " ";
+ double x = hits[j].getPosition()[0];
+ double y = hits[j].getPosition()[1];
+ double z = hits[j].getPosition()[2];
+ debug() << "[" << x << "," << y << "," << z << "]";
+ }
+ debug() << ")" << endmsg;
+ }
+
+ TrackCompatibility comp;
+
+ SubsetHopfieldNN<edm4hep::Track*> subset;
+ //SubsetSimple<edm4hep::Track* > subset;
+ subset.add( tracks_p );
+ subset.setOmega( _omega );
+ subset.calculateBestSet( comp, trackQI );
+
+ std::vector<edm4hep::Track*> accepted = subset.getAccepted();
+ std::vector<edm4hep::Track*> rejected = subset.getRejected();
+
+ debug() << "\tThe accepted tracks:" << endmsg;
+ for( unsigned i=0; i < accepted.size(); i++ ){
+ debug() << accepted[i]->id() << " address " << accepted[i] << endmsg;
+ }
+
+ debug() << "\tThe rejected tracks:" << endmsg;
+ for( unsigned i=0; i < rejected.size(); i++ ){
+ debug() << rejected[i]->id() << " address " << rejected[i] << endmsg;
+ }
+
+ /**********************************************************************************************/
+ /* Save the tracks to a collection (make new TrackImpls from them) */
+ /**********************************************************************************************/
+ debug() << "Fitting and saving of the tracks" << endmsg;
+
+ auto trkCol = _outColHdl.createAndPut();
+
+ for( unsigned i=0; i < accepted.size(); i++ ){
+ edm4hep::Track trackImpl;
+
+ edm4hep::Track* track = accepted[i];
+
+ std::vector<edm4hep::ConstTrackerHit> trackerHitsObj;
+ std::vector<edm4hep::TrackerHit*> trackerHits;
+ std::copy(track->trackerHits_begin(), track->trackerHits_end(), std::back_inserter(trackerHitsObj));
+
+ for(unsigned i=0; i<trackerHitsObj.size(); i++){
+ //debug() << trackerHitsObj[i].id() << endmsg;
+ trackerHits.push_back(Navigation::Instance()->GetTrackerHit(trackerHitsObj[i].getObjectID()));
+ }
+ // setup initial dummy covariance matrix
+ std::array<float,15> covMatrix;
+ for (unsigned icov = 0; icov<covMatrix.size(); ++icov) {
+ covMatrix[icov] = 0;
+ }
+
+ covMatrix[0] = ( _initialTrackError_d0 ); //sigma_d0^2
+ covMatrix[2] = ( _initialTrackError_phi0 ); //sigma_phi0^2
+ covMatrix[5] = ( _initialTrackError_omega ); //sigma_omega^2
+ covMatrix[9] = ( _initialTrackError_z0 ); //sigma_z0^2
+ covMatrix[14] = ( _initialTrackError_tanL ); //sigma_tanl^2
+
+ std::vector< std::pair<float, edm4hep::TrackerHit*> > r2_values;
+ r2_values.reserve(trackerHits.size());
+
+ for (std::vector<edm4hep::TrackerHit*>::iterator it=trackerHits.begin(); it!=trackerHits.end(); ++it) {
+ edm4hep::TrackerHit* h = *it;
+ float r2 = h->getPosition()[0]*h->getPosition()[0]+h->getPosition()[1]*h->getPosition()[1];
+ r2_values.push_back(std::make_pair(r2, *it));
+ }
+
+ sort(r2_values.begin(),r2_values.end());
+
+ trackerHits.clear();
+ trackerHits.reserve(r2_values.size());
+
+ for (std::vector< std::pair<float, edm4hep::TrackerHit*> >::iterator it=r2_values.begin(); it!=r2_values.end(); ++it) {
+ trackerHits.push_back(it->second);
+ }
+
+ bool fit_backwards = MarlinTrk::IMarlinTrack::backward;
+
+ MarlinTrk::IMarlinTrack* marlinTrk = _trkSystem->createTrack();
+
+ int error = 0;
+
+ try {
+
+ error = MarlinTrk::createFinalisedLCIOTrack(marlinTrk, trackerHits, &trackImpl, fit_backwards, covMatrix, _bField, _maxChi2PerHit);
+
+ } catch (...) {
+
+ // delete Track;
+ // delete marlinTrk;
+
+ throw ;
+
+ }
+
+ // Add hit numbers
+
+ std::vector<std::pair<edm4hep::TrackerHit* , double> > hits_in_fit ;
+ std::vector<std::pair<edm4hep::TrackerHit* , double> > outliers ;
+ std::vector<edm4hep::TrackerHit*> all_hits;
+ all_hits.reserve(300);
+
+ marlinTrk->getHitsInFit(hits_in_fit);
+
+ for ( unsigned ihit = 0; ihit < hits_in_fit.size(); ++ihit) {
+ all_hits.push_back(hits_in_fit[ihit].first);
+ }
+
+ UTIL::BitField64 cellID_encoder( lcio::ILDCellID0::encoder_string ) ;
+
+ MarlinTrk::addHitNumbersToTrack(&trackImpl, all_hits, true, cellID_encoder);
+
+ marlinTrk->getOutliers(outliers);
+
+ for ( unsigned ihit = 0; ihit < outliers.size(); ++ihit) {
+ all_hits.push_back(outliers[ihit].first);
+ }
+
+ MarlinTrk::addHitNumbersToTrack(&trackImpl, all_hits, false, cellID_encoder);
+
+ delete marlinTrk;
+
+ if( error != MarlinTrk::IMarlinTrack::success ) {
+ //delete trackImpl;
+ debug() << "TrackSubsetAlg:: Track fit failed with error code " << error << " track dropped. Number of hits = "<< trackerHits.size() << endmsg;
+ continue ;
+ }
+
+ if( trackImpl.getNdf() < 0) {
+ //delete trackImpl;
+ debug() << "TrackSubsetAlg:: Track fit returns " << trackImpl.getNdf() << " degress of freedom track dropped. Number of hits = "<< trackerHits.size() << endmsg;
+ continue ;
+ }
+
+ trkCol->push_back(trackImpl);
+
+// try{
+//
+// Fitter fitter( trackImpl , _trkSystem );
+//
+// TrackStateImpl* trkStateIP = new TrackStateImpl( fitter.getTrackState( lcio::TrackState::AtIP ) ) ;
+// trackImpl->setChi2( fitter.getChi2( lcio::TrackState::AtIP ) );
+// trackImpl->setNdf( fitter.getNdf( lcio::TrackState::AtIP ) );
+// trkStateIP->setLocation( TrackState::AtIP );
+// trackImpl->addTrackState( trkStateIP );
+//
+// trackVec->addElement(trackImpl);
+//
+// }
+// catch( FitterException e ){
+//
+// delete trackImpl;
+//
+// streamlog_out( ERROR ) << "TrackImpl nr " << i << " rejected, because fit failed: " << e.what() << "\n";
+// continue;
+//
+// }
+ }
+
+ debug() << "Saving " << trkCol->size() << " tracks" << endmsg;
+
+ _nEvt ++ ;
+ return StatusCode::SUCCESS;
+}
+
+
+
+
+
+
diff --git a/Reconstruction/SiliconTracking/src/TrackSubsetAlg.h b/Reconstruction/SiliconTracking/src/TrackSubsetAlg.h
new file mode 100644
index 00000000..5676d977
--- /dev/null
+++ b/Reconstruction/SiliconTracking/src/TrackSubsetAlg.h
@@ -0,0 +1,117 @@
+#ifndef TrackSubsetAlg_h
+#define TrackSubsetAlg_h 1
+
+#include "FWCore/DataHandle.h"
+#include "GaudiAlg/GaudiAlgorithm.h"
+
+#include "edm4hep/TrackCollection.h"
+//#include "edm4hep/Track.h"
+#include "TrackSystemSvc/IMarlinTrkSystem.h"
+
+#include "Math/ProbFunc.h"
+
+/** Processor that takes tracks from multiple sources and outputs them (or modified versions, or a subset of them)
+ * as one track collection.
+ *
+ * <h4>Input - Prerequisites</h4>
+ * Track collections
+ *
+ * <h4>Output</h4>
+ * A single track collection
+ *
+ * @param TrackInputCollections A vector of the input track collections <br>
+ * (default value ForwardTracks SiTracks )
+ *
+ * @param TrackOutputCollection Name of the output track collection <br>
+ * (default value SubsetTracks )
+ *
+ * @param MultipleScatteringOn Whether to take multiple scattering into account when fitting the tracks<br>
+ * (default value true )
+ *
+ * @param EnergyLossOn Whether to take energyloss into account when fitting the tracks<br>
+ * (default value true )
+ *
+ * @param SmoothOn Whether to smooth all measurement sites in fit<br>
+ * (default value false )
+ *
+ * @param Omega The parameter omega for the HNN. Controls the influence of the quality indicator. Between 0 and 1:
+ * 1 means high influence of quality indicator, 0 means no influence.
+ *
+ * @author Robin Glattauer, HEPHY
+ *
+ */
+
+class TrackSubsetAlg : public GaudiAlgorithm {
+ public:
+
+ TrackSubsetAlg(const std::string& name, ISvcLocator* svcLoc);
+
+ virtual StatusCode initialize() ;
+
+ virtual StatusCode execute() ;
+
+ virtual StatusCode finalize() ;
+
+ protected:
+ MarlinTrk::IMarlinTrkSystem* _trkSystem;
+ /* Input collection */
+ std::vector<DataHandle<edm4hep::TrackCollection>* > _inTrackColHdls;
+ std::vector<DataHandle<edm4hep::TrackerHitCollection>* > _inTrackerHitColHdls;
+ /* Output collection */
+ DataHandle<edm4hep::TrackCollection> _outColHdl{"SubsetTracks", Gaudi::DataHandle::Writer, this};
+
+ Gaudi::Property<std::vector<std::string> > _trackInputColNames{this, "TrackInputCollections", {"ForwardTracks", "SiTracks"}};
+ Gaudi::Property<std::vector<std::string> > _trackerHitInputColNames{this, "RawTrackerHitCollections",
+ {"VXDTrackerHits", "SITTrackerHits", "FTDPixelTrackerHits", "FTDStripTrackerHits"}};
+ Gaudi::Property<bool> _MSOn{this, "MultipleScatteringOn", true};
+ Gaudi::Property<bool> _ElossOn{this, "EnergyLossOn", true};
+ Gaudi::Property<bool> _SmoothOn{this, "SmoothOn", false};
+ Gaudi::Property<float> _initialTrackError_d0{this, "InitialTrackErrorD0",1e6};
+ Gaudi::Property<float> _initialTrackError_phi0{this, "InitialTrackErrorPhi0",1e2};
+ Gaudi::Property<float> _initialTrackError_omega{this, "InitialTrackErrorOmega",1e-4};
+ Gaudi::Property<float> _initialTrackError_z0{this, "InitialTrackErrorZ0",1e6};
+ Gaudi::Property<float> _initialTrackError_tanL{this, "InitialTrackErrorTanL",1e2};
+ Gaudi::Property<double> _maxChi2PerHit{this, "MaxChi2PerHit", 1e2};
+ Gaudi::Property<double> _omega{this, "Omega", 0.75};
+
+ float _bField;
+
+ int _nRun ;
+ int _nEvt ;
+};
+
+/** A functor to return whether two tracks are compatible: The criterion is if the share a TrackerHit or more */
+class TrackCompatibility{
+ public:
+ inline bool operator()( edm4hep::Track* trackA, edm4hep::Track* trackB ){
+ unsigned nHitsA = trackA->trackerHits_size();
+ unsigned nHitsB = trackB->trackerHits_size();
+ for( unsigned i=0; i < nHitsA; i++){
+ for( unsigned j=0; j < nHitsB; j++){
+ if ( trackA->getTrackerHits(i) == trackB->getTrackerHits(j) ) return false; // a hit is shared -> incompatible
+ }
+ }
+
+ return true;
+ }
+};
+
+
+/** A functor to return the quality of a track, which is currently the chi2 probability. */
+class TrackQI{
+ public:
+ /** @param trkSystem a pointer to an IMarlinTrkSystem, needed for fitting of tracks */
+ TrackQI( MarlinTrk::IMarlinTrkSystem* trkSystem ): _trkSystem(trkSystem){}
+
+ inline double operator()( edm4hep::Track* track ){
+ return ROOT::Math::chisquared_cdf_c( track->getChi2() , track->getNdf() );
+ }
+
+ protected:
+
+ MarlinTrk::IMarlinTrkSystem* _trkSystem;
+};
+#endif
+
+
+
--
GitLab