test/vertexingTest/main.cc

0001 #include <ConstField.h>
0002 #include <Exception.h>
0003 #include <FieldManager.h>
0004 #include <KalmanFitterRefTrack.h>
0005 #include <StateOnPlane.h>
0006 #include <Track.h>
0007 #include <TrackPoint.h>
0008
0009 #include <MaterialEffects.h>
0010 #include <RKTrackRep.h>
0011 #include <TGeoMaterialInterface.h>
0012
0013 #include <EventDisplay.h>
0014
0015 #include <HelixTrackModel.h>
0016 #include <MeasurementCreator.h>
0017
0018 #include <GFRaveVertexFactory.h>
0019
0020 #include <TDatabasePDG.h>
0021 #include <TEveManager.h>
0022 #include <TGeoManager.h>
0023 #include <TRandom.h>
0024 #include <TVector3.h>
0025 #include <vector>
0026
0027 #include <TROOT.h>
0028 #include <TClonesArray.h>
0029 #include <TFile.h>
0030 #include <TTree.h>
0031 #include <TDatabasePDG.h>
0032 #include <TMath.h>
0033
0034
0035
0036
0037 int main() {
0038
0039   gRandom->SetSeed(14);
0040
0041   // init MeasurementCreator
0042   genfit::MeasurementCreator measurementCreator;
0043
0044
0045   // init geometry and mag. field
0046   new TGeoManager("Geometry", "Geane geometry");
0047   TGeoManager::Import("genfitGeom.root");
0048   genfit::FieldManager::getInstance()->init(new genfit::ConstField(0.,0., 15.)); // 15 kGauss
0049   genfit::MaterialEffects::getInstance()->init(new genfit::TGeoMaterialInterface());
0050
0051
0052   // init event display
0053   genfit::EventDisplay* display = genfit::EventDisplay::getInstance();
0054
0055
0056   // init fitter
0057   genfit::AbsKalmanFitter* fitter = new genfit::KalmanFitterRefTrack();
0058
0059   // init vertex factory
0060   genfit::GFRaveVertexFactory vertexFactory(2);
0061   vertexFactory.setMethod("kalman-smoothing:1");
0062
0063
0064   // init root file
0065   // tracks and vertices are written to two different branches
0066   TFile* trackFile = new TFile("tracks.root", "RECREATE");
0067   trackFile->cd();
0068   TTree* tree = new TTree("tree", "fitted tracks");
0069   TClonesArray trackArray("genfit::Track");
0070   tree->Branch("trackBranch", &trackArray, 32000, -1);
0071
0072   TClonesArray vertexArray("genfit::GFRaveVertex");
0073   tree->Branch("vertexBranch", &vertexArray, 32000, -1);
0074
0075   std::vector<genfit::Track*> tracks;
0076   std::vector<genfit::GFRaveVertex*> vertices;
0077
0078   // main loop
0079   for (unsigned int iEvent=0; iEvent<10; ++iEvent){
0080
0081     // clean up
0082     trackArray.Delete();
0083     vertexArray.Delete();
0084     tracks.clear();
0085     vertices.clear();
0086
0087
0088     unsigned int nTracks = gRandom->Uniform(2, 10);
0089
0090     for (unsigned int iTrack=0; iTrack<nTracks; ++iTrack){
0091
0092       // true start values
0093       TVector3 pos(0, 0, 0);
0094       TVector3 mom(1.,0,0);
0095       mom.SetPhi(gRandom->Uniform(0.,2*TMath::Pi()));
0096       mom.SetTheta(gRandom->Uniform(0.4*TMath::Pi(),0.6*TMath::Pi()));
0097       mom.SetMag(gRandom->Uniform(0.2, 1.));
0098
0099
0100       // helix track model
0101       const int pdg = 13;               // particle pdg code
0102       const double charge = TDatabasePDG::Instance()->GetParticle(pdg)->Charge()/(3.);
0103       genfit::HelixTrackModel* helix = new genfit::HelixTrackModel(pos, mom, charge);
0104       measurementCreator.setTrackModel(helix);
0105
0106
0107       unsigned int nMeasurements = gRandom->Uniform(5, 15);
0108
0109
0110       // smeared start values
0111       const bool smearPosMom = true;     // init the Reps with smeared pos and mom
0112       const double posSmear = 0.1;     // cm
0113       const double momSmear = 3. /180.*TMath::Pi();     // rad
0114       const double momMagSmear = 0.1;   // relative
0115
0116       TVector3 posM(pos);
0117       TVector3 momM(mom);
0118       if (smearPosMom) {
0119         posM.SetX(gRandom->Gaus(posM.X(),posSmear));
0120         posM.SetY(gRandom->Gaus(posM.Y(),posSmear));
0121         posM.SetZ(gRandom->Gaus(posM.Z(),posSmear));
0122
0123         momM.SetPhi(gRandom->Gaus(mom.Phi(),momSmear));
0124         momM.SetTheta(gRandom->Gaus(mom.Theta(),momSmear));
0125         momM.SetMag(gRandom->Gaus(mom.Mag(), momMagSmear*mom.Mag()));
0126       }
0127       // approximate covariance
0128       TMatrixDSym covM(6);
0129       double resolution = 0.01;
0130       for (int i = 0; i < 3; ++i)
0131         covM(i,i) = resolution*resolution;
0132       for (int i = 3; i < 6; ++i)
0133         covM(i,i) = pow(resolution / nMeasurements / sqrt(3), 2);
0134
0135
0136       // trackrep
0137       genfit::AbsTrackRep* rep = new genfit::RKTrackRep(pdg);
0138
0139       // smeared start state
0140       genfit::MeasuredStateOnPlane stateSmeared(rep);
0141       rep->setPosMomCov(stateSmeared, posM, momM, covM);
0142
0143
0144       // create track
0145       TVectorD seedState(6);
0146       TMatrixDSym seedCov(6);
0147       rep->get6DStateCov(stateSmeared, seedState, seedCov);
0148
0149       new(trackArray[iTrack]) genfit::Track(rep, seedState, seedCov);
0150       genfit::Track* trackPtr(static_cast<genfit::Track*>(trackArray.At(iTrack)));
0151       tracks.push_back(trackPtr);
0152
0153       // create random measurement types
0154       std::vector<genfit::eMeasurementType> measurementTypes;
0155       for (unsigned int i = 0; i < nMeasurements; ++i)
0156         measurementTypes.push_back(genfit::eMeasurementType(gRandom->Uniform(8)));
0157
0158
0159       // create smeared measurements and add to track
0160       try{
0161         for (unsigned int i=1; i<measurementTypes.size(); ++i){
0162           std::vector<genfit::AbsMeasurement*> measurements = measurementCreator.create(measurementTypes[i], i*5.);
0163           trackPtr->insertPoint(new genfit::TrackPoint(measurements, trackPtr));
0164         }
0165       }
0166       catch(genfit::Exception& e){
0167         std::cerr<<"Exception, next track"<<std::endl;
0168         std::cerr << e.what();
0169         continue; // here is a memleak!
0170       }
0171
0172       //check
0173       assert(trackPtr->checkConsistency());
0174
0175       // do the fit
0176       try{
0177         fitter->processTrack(trackPtr);
0178       }
0179       catch(genfit::Exception& e){
0180         std::cerr << e.what();
0181         std::cerr << "Exception, next track" << std::endl;
0182         continue;
0183       }
0184
0185       //check
0186       assert(trackPtr->checkConsistency());
0187
0188     } // end loop over tracks
0189
0190
0191
0192     // vertexing
0193     vertexFactory.findVertices(&vertices, tracks);
0194
0195     for (unsigned int i=0; i<vertices.size(); ++i) {
0196       new(vertexArray[i]) genfit::GFRaveVertex(*(vertices[i]));
0197
0198       genfit::GFRaveVertex* vtx = static_cast<genfit::GFRaveVertex*>(vertices[i]);
0199       for (unsigned int j=0; j<vtx->getNTracks(); ++j) {
0200         std::cout << "track parameters uniqueID: " << vtx->getParameters(j)->GetUniqueID() << "\n";
0201       }
0202     }
0203
0204
0205     for (unsigned int i=0; i<tracks.size(); ++i) {
0206       genfit::Track* trk = static_cast<genfit::Track*>(tracks[i]);
0207       std::cout << "track uniqueID: " << trk->GetUniqueID() << "\n";
0208     }
0209
0210
0211     // fill
0212     std::cout << "trackArray nr of entries: " << trackArray.GetEntries() << "\n";
0213     tree->Fill();
0214
0215
0216     if (iEvent < 1000) {
0217       // add tracks to event display
0218       display->addEvent(tracks);
0219     }
0220
0221   } // end loop over events
0222
0223   delete fitter;
0224
0225   // write and close files
0226   trackFile->Write();
0227   trackFile->Close();
0228   /*vertexFile->Write();
0229   vertexFile->Close();*/
0230
0231   // open event display
0232   //display->open();
0233
0234 }
0235
0236