Skip to content
Snippets Groups Projects
example_tracking.cxx 8.34 KiB
Newer Older
  • Learn to ignore specific revisions
  • R__LOAD_LIBRARY(libGenFind)
    R__LOAD_LIBRARY(libgenfit2)
    
    R__LOAD_LIBRARY(libDDG4IO.so)
    
    
    #include "ConstField.h"
    #include "Exception.h"
    #include "FieldManager.h"
    #include "KalmanFitterRefTrack.h"
    #include "DAF.h"
    #include "StateOnPlane.h"
    #include "Track.h"
    #include "TrackPoint.h"
    
    #include <MaterialEffects.h>
    #include <RKTrackRep.h>
    #include <TGeoMaterialInterface.h>
    
    #include <EventDisplay.h>
    
    #include <HelixTrackModel.h>
    #include <MeasurementCreator.h>
    #include <WireMeasurement.h>
    #include "PlanarMeasurement.h"
    
    #include <TDatabasePDG.h>
    #include <TEveManager.h>
    #include <TGeoManager.h>
    #include <TRandom.h>
    #include <TVector3.h>
    #include <vector>
    
    
    #include "Math/GenVector/VectorUtil.h"
    
    #include "Math/Vector3D.h"
    #include "Math/Vector4D.h"
    
    #include "ROOT/RDataFrame.hxx"
    
    #include "TCanvas.h"
    
    #include "TChain.h"
    #include "TDatabasePDG.h"
    #include "TEveGeoNode.h"
    #include "TF1.h"
    #include "TFile.h"
    #include "TH1D.h"
    #include "TH1F.h"
    
    #include "TLegend.h"
    #include "TMath.h"
    #include "TRandom3.h"
    #include "TTree.h"
    
    #include <random>
    
    #include <tuple>
    #include <vector>
    
    
    // DD4hep
    #include "DD4hep/Detector.h"
    #include "DDG4/Geant4Data.h"
    #include "DDRec/CellIDPositionConverter.h"
    #include "DDRec/SurfaceManager.h"
    #include "DDRec/Surface.h"
    
    #include "lcio2/TrackerRawDataData.h"
    #include "lcio2/TrackerRawData.h"
    #include "lcio2/TrackerHit.h"
    #include "lcio2/TrackerHitData.h"
    #include "lcio2/TrackCollection.h"
    
    #include "GenFind/GenFindHits.h"
    #include "GenFind/HoughTransform.h"
    #include "DD4hep/DD4hepUnits.h"
    
    #ifdef __CLING__
    #pragma link off all globals;
    #pragma link off all classes;
    #pragma link off all functions;
    #pragma link C++ nestedclass;
    #pragma link C++ class std::vector<std::vector<ROOT::Math::XYZTVector>>+;
    #endif
    
    
    void example_tracking(const char* fname = "gem_tracker_recon.root"){
    
      //ROOT::EnableImplicitMT(0);
      ROOT::DisableImplicitMT(); // Tell ROOT you want to go parallel
    
    
      using namespace ROOT::Math;
      using namespace genfind;
      using namespace lcio2;
    
      double degree = TMath::Pi()/180.0;
    
      std::random_device rd;
      std::mt19937 gen(rd());
    
      TChain* t = new TChain("recon_EVENT");
      t->Add(fname);
    
      ROOT::RDataFrame d0(*t );
    
      //std::cout << t->GetBranch("GEMTrackerHits")->GetClassName() << std::endl;
      //std::vector<dd4hep::sim::Geant4Tracker::Hit*>
      
      auto nhits = [] (const std::vector<dd4hep::sim::Geant4Tracker::Hit*>& hits){ return hits.size(); };
    
      dd4hep::Detector& detector = dd4hep::Detector::getInstance();
    
      detector.fromCompact("gem_tracker.xml");
    
      dd4hep::rec::CellIDPositionConverter cellid_converter(detector);
    
      ROOT::Math::XYZVector position  = {0,0,0}; // position to calculate magnetic field at (the origin in this case)
      //double bField[3]             = {0,0,0};
      auto bField = detector.field().magneticField(position);
      double Bz                    = bField.z()/dd4hep::tesla;
      std::cout << " Magnetic Field Bz             = " << Bz << std::endl;
    
      auto vol_man = detector.volumeManager();
    
      genfit::FieldManager::getInstance()->init(new genfit::ConstField(0.,0., Bz*10.0)); // gentfit uses kilo-Gauss
      genfit::MaterialEffects::getInstance()->init(new genfit::TGeoMaterialInterface());
      genfit::EventDisplay* display = genfit::EventDisplay::getInstance();
      display->setOptions("X");
    
      // Add the geometry to eve display
      TGeoNode*       node1 = gGeoManager->GetTopNode();
      TEveGeoTopNode* its   = new TEveGeoTopNode(gGeoManager, node1);
      gEve->AddGlobalElement(its);
    
      genfit::AbsKalmanFitter* fitter = new genfit::KalmanFitterRefTrack();
      //genfit::AbsKalmanFitter* fitter = new genfit::DAF();
    
      TMatrixDSym covM(6);
      double res1 = 0.0001;
      for (int i = 0; i < 3; ++i)
        covM(i,i) = res1*res1;
      for (int i = 3; i < 6; ++i)
        covM(i,i) = pow(res1/9.0/sqrt(3.0), 2);
    
    
      const int detId(0); // detector ID
      int planeId(0); // detector plane ID
      int hitId(0); // hit ID
    
      double detectorResolution = 0.001; // resolution of planar detectors
      TMatrixDSym hitCov(2);
      hitCov.UnitMatrix();
      hitCov *= detectorResolution*detectorResolution;
      TVectorD hitCoords(2);
    
      dd4hep::rec::SurfaceManager& surfMan = *detector.extension<dd4hep::rec::SurfaceManager>() ;
      //const SurfaceMap& _surfMap = *surfMan.map( "world" ) ;
      const auto& _surfMap = *surfMan.map( "world" ) ;
    
      auto gem_tracks = 
        [&](const std::vector<lcio2::TrackerHitData>& hits) {
          std::vector<ROOT::Math::XYZTVector> vec_hits;
          HoughTransform ht;
          for(const auto& h: hits) {
            std::cout << h.position.at(0) << ", "  << h.position.at(1) << ", " << h.position.at(2) << std::endl;
            vec_hits.push_back({h.position.at(0),h.position.at(1), h.position.at(2), h.time});
          }
          std::vector<std::vector<ROOT::Math::XYZTVector>> res = ht(vec_hits);
          //ROOT::Math::XYZTVector ref_hit = {0.0,0.0,0.0,0.0};
          //std::vector<ConformalHit> chits2 = compute_conformal_hits(vec_hits,ref_hit);
          //std::vector<double> peaks = ht.FindPhiPeaks(chits2);
          //for(const auto& p : chits2){
          //  std::cout << p.fPosition.Z() << std::endl;
          //}
          std::cout << " npeaks = " << res.size() << std::endl;
          if(res.size()>0) std::cout << " nhits(0) = " << res.at(0).size() << std::endl;
          return res;
        };
    
      std::vector<genfit::Track*> all_tracks;
      auto fit_tracks = 
        [&](const std::vector<std::vector<ROOT::Math::XYZTVector>>& possible_tracks) {
          int ihit = 0;
          std::vector<lcio2::TrackData> result_tracks;
          int i_track = 0;
          for(const auto& atrack : possible_tracks) {
            auto first_hit_pos  = atrack.at(0);
            const int pdg = 11; // particle pdg code
    
            // trackrep
            genfit::AbsTrackRep* rep = new genfit::RKTrackRep(pdg);
            // smeared start state
            genfit::MeasuredStateOnPlane stateSmeared(rep);
            // start values for the fit, e.g. from pattern recognition
            TVector3 pos(0.0,0.0,0.0);
            TVector3 mom(0.0, 0.0, 6.0);
            mom.SetMag(6.0);
            stateSmeared.setPosMomCov(pos, mom, covM);
    
            // Create track.
            TVectorD    seedState(6);
            TMatrixDSym seedCov(6);
            stateSmeared.get6DStateCov(seedState, seedCov);
    
            auto aTrack  = new genfit::Track(rep, seedState, seedCov);
            //auto& fitTrack = *aTrack;
            //genfit::Track fitTrack(rep, seedState, seedCov);
    
            for( auto thit : atrack ) {
              ihit++;
    
              TVector3 point      = {thit.X(), thit.Y(), thit.Z()};
              TVector3 u_dir      = {1,0,0};
              TVector3 v_dir      = {0,1,0};
              genfit::SharedPlanePtr plane(new genfit::DetPlane(point, u_dir, v_dir) );
    
              // add some planar hits to track with coordinates
              hitCoords[0] = 0.0;//thit->position.X()*0.1;
              hitCoords[1] = 0.0;//thit->position.Y()*0.1;
    
              genfit::PlanarMeasurement* measurement = new genfit::PlanarMeasurement(hitCoords, hitCov, detId, ++hitId, nullptr);
              static_cast<genfit::PlanarMeasurement*>(measurement)->setPlane(plane, ihit);//thit->cellID);
              (*aTrack).insertPoint(new genfit::TrackPoint(measurement, aTrack));
            }
    
            //check
            aTrack->checkConsistency();
            bool res_is_good = true;
            try {
            // do the fit
            fitter->processTrack(aTrack,true);
    
            // print fit result
            aTrack->getFittedState().Print();
    
            //check
            aTrack->checkConsistency();
    
            } catch (const std::exception& e) {
              res_is_good = false;
            }
    
            if( res_is_good )
              all_tracks.push_back(aTrack);
    
            //genfit::SharedPlanePtr test_plane(new genfit::DetPlane({0,0,0}, {1,0,0} , {0,1,0}) );
            //std::vector<genfit::MeasurementOnPlane*> constructMeasurementsOnPlane  ...
            //delete fitter;
          }
          display->addEvent(all_tracks);
          //std::cout << " |P| = " << (*mc_particles)[0]->psx << ", " << (*mc_particles)[0]->psy << ", " << (*mc_particles)[0]->psz << std::endl ;
    
          //std::string dummy;
          //std::cout << "Enter to continue..." << std::endl;
          //std::getline(std::cin, dummy);
          //display->clear();
          return result_tracks;
        };
    
      auto d1 = d0.Range(10).Define("potential_tracks", gem_tracks, {"TrackerHits"})
                .Define("Tracks", fit_tracks, {"potential_tracks"}) ;
      //d1.Foreach(fit_tracks, {"potential_tracks"});
    
      //TCanvas* c = new TCanvas();
      //h0->DrawClone();
      d1.Snapshot("track_EVENT","test_gem_tracker_track.root");
    
      std::cout << "Display events: " << display->getNEvents() << std::endl;
    
      display->open();
    }