Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-08-05 08:18:28

0001 //
0002 // Write fit results to tree, read again, compare.
0003 //
0004 
0005 #include <ConstField.h>
0006 #include <Exception.h>
0007 #include <FieldManager.h>
0008 #include <KalmanFitterRefTrack.h>
0009 #include <StateOnPlane.h>
0010 #include <Track.h>
0011 #include <TrackPoint.h>
0012 
0013 #include <MaterialEffects.h>
0014 #include <RKTrackRep.h>
0015 #include <TGeoMaterialInterface.h>
0016 
0017 #include <EventDisplay.h>
0018 
0019 #include <HelixTrackModel.h>
0020 #include <MeasurementCreator.h>
0021 
0022 #include <TDatabasePDG.h>
0023 #include <TEveManager.h>
0024 #include <TGeoManager.h>
0025 #include <TRandom.h>
0026 #include <TVector3.h>
0027 #include <vector>
0028 
0029 #include "TDatabasePDG.h"
0030 #include <TMath.h>
0031 #include <TFile.h>
0032 #include <TTree.h>
0033 
0034 
0035 #define FILENAME "/tmp/streamerTest.root"
0036 
0037 constexpr bool verbose = false;
0038 
0039 bool emptyTrackTest()
0040 {
0041   TFile *f = TFile::Open(FILENAME, "RECREATE");
0042   f->cd();
0043   genfit::Track *t = new genfit::Track();
0044   try {
0045     t->checkConsistency();
0046   } catch (genfit::Exception&) {
0047     return false;
0048   }
0049 
0050   t->Write("direct");
0051   f->Close();
0052   delete t;
0053   delete f;
0054 
0055   f = TFile::Open(FILENAME, "READ");
0056   t = (genfit::Track*)f->Get("direct");
0057 
0058   bool result = false;
0059   try {
0060     t->checkConsistency();
0061     result = true;
0062   } catch (genfit::Exception&) {
0063     result = false;
0064   }
0065   delete t;
0066   delete f;
0067   return result;
0068 }
0069 
0070 
0071 int main() {
0072   if (!emptyTrackTest()) {
0073     std::cout << "emptyTrackTest failed." << std::endl;
0074     return 1;
0075   }
0076 
0077 
0078   // prepare output tree for Tracks 
0079   // std::unique_ptr<genfit::Track> fitTrack(new genfit::Track());
0080   genfit::Track* fitTrack = new genfit::Track();
0081   TVectorD stateFinal;
0082   TMatrixDSym covFinal;
0083   genfit::DetPlane planeFinal;
0084 
0085   TFile* fOut = new TFile(FILENAME, "RECREATE");
0086   fOut->cd();
0087   TTree* tResults = new TTree("tResults", "results from track fit");
0088   // it is important to set the splitLevel to -1, because some of the genfit classes use custom streamers
0089   tResults->Branch("gfTrack", "genfit::Track", &fitTrack, 32000, -1);
0090   tResults->Branch("stateFinal", &stateFinal);
0091   tResults->Branch("covFinal", &covFinal, 32000, -1);
0092   tResults->Branch("planeFinal", &planeFinal, 32000, -1);
0093 
0094 
0095 
0096 
0097 
0098   gRandom->SetSeed(14);
0099 
0100   // init MeasurementCreator
0101   genfit::MeasurementCreator measurementCreator;
0102 
0103 
0104   // init geometry and mag. field
0105   new TGeoManager("Geometry", "Geane geometry");
0106   TGeoManager::Import("genfitGeom.root");
0107   genfit::FieldManager::getInstance()->init(new genfit::ConstField(0.,0., 15.)); // 15 kGauss
0108   genfit::MaterialEffects::getInstance()->init(new genfit::TGeoMaterialInterface());
0109 
0110 
0111   // init fitter
0112   genfit::AbsKalmanFitter* fitter = new genfit::KalmanFitterRefTrack();
0113 
0114   unsigned int nEvents(100);
0115 
0116   // main loop
0117   for (unsigned int iEvent=0; iEvent<nEvents; ++iEvent){
0118 
0119     // true start values
0120     TVector3 pos(0, 0, 0);
0121     TVector3 mom(1.,0,0);
0122     mom.SetPhi(gRandom->Uniform(0.,2*TMath::Pi()));
0123     mom.SetTheta(gRandom->Uniform(0.4*TMath::Pi(),0.6*TMath::Pi()));
0124     mom.SetMag(gRandom->Uniform(0.2, 1.));
0125 
0126 
0127     // helix track model
0128     const int pdg = 13;               // particle pdg code
0129     const double charge = TDatabasePDG::Instance()->GetParticle(pdg)->Charge()/(3.);
0130     genfit::HelixTrackModel* helix = new genfit::HelixTrackModel(pos, mom, charge);
0131     measurementCreator.setTrackModel(helix);
0132 
0133 
0134     unsigned int nMeasurements = gRandom->Uniform(5, 15);
0135 
0136 
0137     // smeared start values
0138     const bool smearPosMom = true;     // init the Reps with smeared pos and mom
0139     const double posSmear = 0.1;     // cm
0140     const double momSmear = 3. /180.*TMath::Pi();     // rad
0141     const double momMagSmear = 0.1;   // relative
0142 
0143     TVector3 posM(pos);
0144     TVector3 momM(mom);
0145     if (smearPosMom) {
0146       posM.SetX(gRandom->Gaus(posM.X(),posSmear));
0147       posM.SetY(gRandom->Gaus(posM.Y(),posSmear));
0148       posM.SetZ(gRandom->Gaus(posM.Z(),posSmear));
0149 
0150       momM.SetPhi(gRandom->Gaus(mom.Phi(),momSmear));
0151       momM.SetTheta(gRandom->Gaus(mom.Theta(),momSmear));
0152       momM.SetMag(gRandom->Gaus(mom.Mag(), momMagSmear*mom.Mag()));
0153     }
0154     // approximate covariance
0155     TMatrixDSym covM(6);
0156     double resolution = 0.01;
0157     for (int i = 0; i < 3; ++i)
0158       covM(i,i) = resolution*resolution;
0159     for (int i = 3; i < 6; ++i)
0160       covM(i,i) = pow(resolution / nMeasurements / sqrt(3), 2);
0161 
0162 
0163     // trackrep
0164     genfit::AbsTrackRep* rep = new genfit::RKTrackRep(pdg);
0165 
0166     // smeared start state
0167     genfit::MeasuredStateOnPlane stateSmeared(rep);
0168     rep->setPosMomCov(stateSmeared, posM, momM, covM);
0169 
0170 
0171     // create track
0172     TVectorD seedState(6);
0173     TMatrixDSym seedCov(6);
0174     rep->get6DStateCov(stateSmeared, seedState, seedCov);
0175     fitTrack = new genfit::Track(rep, seedState, seedCov);
0176 
0177 
0178     // create random measurement types
0179     std::vector<genfit::eMeasurementType> measurementTypes;
0180     for (unsigned int i = 0; i < nMeasurements; ++i)
0181       measurementTypes.push_back(genfit::eMeasurementType(gRandom->Uniform(8)));
0182 
0183 
0184     // create smeared measurements and add to track
0185     try{
0186       for (unsigned int i=0; i<measurementTypes.size(); ++i){
0187         std::vector<genfit::AbsMeasurement*> measurements = measurementCreator.create(measurementTypes[i], i*5.);
0188         genfit::TrackPoint* tp = new genfit::TrackPoint(measurements, fitTrack);
0189         // test scatterers
0190         genfit::ThinScatterer* sc = new genfit::ThinScatterer(genfit::SharedPlanePtr(new genfit::DetPlane(TVector3(1,1,1), TVector3(1,1,1))),
0191                                                               genfit::Material(1,2,3,4,5));
0192         tp->setScatterer(sc);
0193 
0194         fitTrack->insertPoint(tp);
0195       }
0196     }
0197     catch(genfit::Exception& e){
0198       std::cerr<<"Exception, next track"<<std::endl;
0199       std::cerr << e.what();
0200       delete fitTrack;
0201       fitTrack = 0;
0202       continue;
0203     }
0204 
0205     fitTrack->checkConsistency();
0206 
0207     // do the fit
0208     fitter->processTrack(fitTrack);
0209 
0210     fitTrack->checkConsistency();
0211 
0212 
0213     stateFinal.ResizeTo(fitTrack->getFittedState().getState());
0214     stateFinal = fitTrack->getFittedState().getState();
0215     covFinal.ResizeTo(fitTrack->getFittedState().getCov());
0216     covFinal = fitTrack->getFittedState().getCov();
0217     planeFinal =  *(fitTrack->getFittedState().getPlane());
0218 
0219     switch (iEvent % 5) {
0220     case 0:
0221       break;
0222     case 1:
0223       fitTrack->prune("FL");
0224       break;
0225     case 2:
0226       fitTrack->prune("W");
0227       break;
0228     case 3:
0229       fitTrack->prune("RC");
0230       break;
0231     case 4:
0232       fitTrack->prune("U");
0233       break;
0234     }
0235 
0236     tResults->Fill();
0237 
0238     //fitTrack->Print();
0239 
0240     delete fitTrack;
0241     fitTrack = 0;
0242 
0243   }// end loop over events
0244 
0245   delete fitter;
0246 
0247 
0248 
0249 
0250   fOut->Write();
0251   delete fOut;
0252 
0253   fOut = TFile::Open(FILENAME, "READ");
0254   fOut->GetObject("tResults", tResults);
0255   TVectorD* pState = 0;
0256   tResults->SetBranchAddress("stateFinal", &pState);
0257   TMatrixDSym* pMatrix = 0;
0258   tResults->SetBranchAddress("covFinal", &pMatrix);
0259   genfit::DetPlane* plane = 0;
0260   tResults->SetBranchAddress("planeFinal", &plane);
0261   tResults->SetBranchAddress("gfTrack", &fitTrack);
0262 
0263   int fail(0);
0264 
0265   for (Long_t nEntry = 0; nEntry < tResults->GetEntries(); ++nEntry) {
0266     tResults->GetEntry(nEntry);
0267 
0268     try {
0269       fitTrack->checkConsistency();
0270     } catch (genfit::Exception& e) {
0271       std::cout << e.getExcString() << std::endl;
0272       return 1;
0273     }
0274 
0275     try {
0276       if (*pState == fitTrack->getFittedState().getState() &&
0277           *pMatrix == fitTrack->getFittedState().getCov() &&
0278           *plane ==  *(fitTrack->getFittedState().getPlane())) {
0279         // track ok
0280       }
0281       else {
0282         if (verbose) {
0283           std::cout << "stored track not equal, small differences can occur if some info has been pruned." << std::endl;
0284           pState->Print();
0285           fitTrack->getFittedState().getState().Print();
0286           pMatrix->Print();
0287           fitTrack->getFittedState().getCov().Print();
0288           plane->Print();
0289           fitTrack->getFittedState().getPlane()->Print();
0290           }
0291         ++fail;
0292         //return 1;
0293       }
0294     }
0295     catch (genfit::Exception& e) {
0296         if (verbose) {
0297             std::cerr << e.what();
0298         }
0299       return 1;
0300     }
0301   }
0302   std::cout << nEvents - fail << " stored tracks are identical to fitted tracks, as far as tested." << std::endl;
0303   std::cout << fail << " tracks were not identical to fitted tracks, as far as tested." << std::endl;
0304   delete fitTrack;
0305   std::cout << "deleteing didn't segfault" << std::endl;
0306 
0307   return 0;
0308 }