Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
main.cc
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file main.cc
1 #include <ConstField.h>
2 #include <Exception.h>
3 #include <FieldManager.h>
4 #include <KalmanFitterRefTrack.h>
5 #include <StateOnPlane.h>
6 #include <Track.h>
7 #include <TrackPoint.h>
8 
9 #include <MaterialEffects.h>
10 #include <RKTrackRep.h>
11 #include <TGeoMaterialInterface.h>
12 
13 #include <EventDisplay.h>
14 
15 #include <HelixTrackModel.h>
16 #include <MeasurementCreator.h>
17 
18 #include <GFRaveVertexFactory.h>
19 
20 #include <TDatabasePDG.h>
21 #include <TEveManager.h>
22 #include <TGeoManager.h>
23 #include <TRandom.h>
24 #include <TVector3.h>
25 #include <vector>
26 
27 #include <TROOT.h>
28 #include <TClonesArray.h>
29 #include <TFile.h>
30 #include <TTree.h>
31 #include <TDatabasePDG.h>
32 #include <TMath.h>
33 
34 
35 
36 
37 int main() {
38 
39  gRandom->SetSeed(14);
40 
41  // init MeasurementCreator
42  genfit::MeasurementCreator measurementCreator;
43 
44 
45  // init geometry and mag. field
46  new TGeoManager("Geometry", "Geane geometry");
47  TGeoManager::Import("genfitGeom.root");
48  genfit::FieldManager::getInstance()->init(new genfit::ConstField(0.,0., 15.)); // 15 kGauss
50 
51 
52  // init event display
54 
55 
56  // init fitter
58 
59  // init vertex factory
60  genfit::GFRaveVertexFactory vertexFactory(2);
61  vertexFactory.setMethod("kalman-smoothing:1");
62 
63 
64  // init root file
65  // tracks and vertices are written to two different branches
66  TFile* trackFile = new TFile("tracks.root", "RECREATE");
67  trackFile->cd();
68  TTree* tree = new TTree("tree", "fitted tracks");
69  TClonesArray trackArray("genfit::Track");
70  tree->Branch("trackBranch", &trackArray, 32000, -1);
71 
72  TClonesArray vertexArray("genfit::GFRaveVertex");
73  tree->Branch("vertexBranch", &vertexArray, 32000, -1);
74 
75  std::vector<genfit::Track*> tracks;
76  std::vector<genfit::GFRaveVertex*> vertices;
77 
78  // main loop
79  for (unsigned int iEvent=0; iEvent<10; ++iEvent){
80 
81  // clean up
82  trackArray.Delete();
83  vertexArray.Delete();
84  tracks.clear();
85  vertices.clear();
86 
87 
88  unsigned int nTracks = gRandom->Uniform(2, 10);
89 
90  for (unsigned int iTrack=0; iTrack<nTracks; ++iTrack){
91 
92  // true start values
93  TVector3 pos(0, 0, 0);
94  TVector3 mom(1.,0,0);
95  mom.SetPhi(gRandom->Uniform(0.,2*TMath::Pi()));
96  mom.SetTheta(gRandom->Uniform(0.4*TMath::Pi(),0.6*TMath::Pi()));
97  mom.SetMag(gRandom->Uniform(0.2, 1.));
98 
99 
100  // helix track model
101  const int pdg = 13; // particle pdg code
102  const double charge = TDatabasePDG::Instance()->GetParticle(pdg)->Charge()/(3.);
103  genfit::HelixTrackModel* helix = new genfit::HelixTrackModel(pos, mom, charge);
104  measurementCreator.setTrackModel(helix);
105 
106 
107  unsigned int nMeasurements = gRandom->Uniform(5, 15);
108 
109 
110  // smeared start values
111  const bool smearPosMom = true; // init the Reps with smeared pos and mom
112  const double posSmear = 0.1; // cm
113  const double momSmear = 3. /180.*TMath::Pi(); // rad
114  const double momMagSmear = 0.1; // relative
115 
116  TVector3 posM(pos);
117  TVector3 momM(mom);
118  if (smearPosMom) {
119  posM.SetX(gRandom->Gaus(posM.X(),posSmear));
120  posM.SetY(gRandom->Gaus(posM.Y(),posSmear));
121  posM.SetZ(gRandom->Gaus(posM.Z(),posSmear));
122 
123  momM.SetPhi(gRandom->Gaus(mom.Phi(),momSmear));
124  momM.SetTheta(gRandom->Gaus(mom.Theta(),momSmear));
125  momM.SetMag(gRandom->Gaus(mom.Mag(), momMagSmear*mom.Mag()));
126  }
127  // approximate covariance
128  TMatrixDSym covM(6);
129  double resolution = 0.01;
130  for (int i = 0; i < 3; ++i)
131  covM(i,i) = resolution*resolution;
132  for (int i = 3; i < 6; ++i)
133  covM(i,i) = pow(resolution / nMeasurements / sqrt(3), 2);
134 
135 
136  // trackrep
137  genfit::AbsTrackRep* rep = new genfit::RKTrackRep(pdg);
138 
139  // smeared start state
140  genfit::MeasuredStateOnPlane stateSmeared(rep);
141  rep->setPosMomCov(stateSmeared, posM, momM, covM);
142 
143 
144  // create track
145  TVectorD seedState(6);
146  TMatrixDSym seedCov(6);
147  rep->get6DStateCov(stateSmeared, seedState, seedCov);
148 
149  new(trackArray[iTrack]) genfit::Track(rep, seedState, seedCov);
150  genfit::Track* trackPtr(static_cast<genfit::Track*>(trackArray.At(iTrack)));
151  tracks.push_back(trackPtr);
152 
153  // create random measurement types
154  std::vector<genfit::eMeasurementType> measurementTypes;
155  for (unsigned int i = 0; i < nMeasurements; ++i)
156  measurementTypes.push_back(genfit::eMeasurementType(gRandom->Uniform(8)));
157 
158 
159  // create smeared measurements and add to track
160  try{
161  for (unsigned int i=1; i<measurementTypes.size(); ++i){
162  std::vector<genfit::AbsMeasurement*> measurements = measurementCreator.create(measurementTypes[i], i*5.);
163  trackPtr->insertPoint(new genfit::TrackPoint(measurements, trackPtr));
164  }
165  }
166  catch(genfit::Exception& e){
167  std::cerr<<"Exception, next track"<<std::endl;
168  std::cerr << e.what();
169  continue; // here is a memleak!
170  }
171 
172  //check
173  assert(trackPtr->checkConsistency());
174 
175  // do the fit
176  try{
177  fitter->processTrack(trackPtr);
178  }
179  catch(genfit::Exception& e){
180  std::cerr << e.what();
181  std::cerr << "Exception, next track" << std::endl;
182  continue;
183  }
184 
185  //check
186  assert(trackPtr->checkConsistency());
187 
188  } // end loop over tracks
189 
190 
191 
192  // vertexing
193  vertexFactory.findVertices(&vertices, tracks);
194 
195  for (unsigned int i=0; i<vertices.size(); ++i) {
196  new(vertexArray[i]) genfit::GFRaveVertex(*(vertices[i]));
197 
198  genfit::GFRaveVertex* vtx = static_cast<genfit::GFRaveVertex*>(vertices[i]);
199  for (unsigned int j=0; j<vtx->getNTracks(); ++j) {
200  std::cout << "track parameters uniqueID: " << vtx->getParameters(j)->GetUniqueID() << "\n";
201  }
202  }
203 
204 
205  for (unsigned int i=0; i<tracks.size(); ++i) {
206  genfit::Track* trk = static_cast<genfit::Track*>(tracks[i]);
207  std::cout << "track uniqueID: " << trk->GetUniqueID() << "\n";
208  }
209 
210 
211  // fill
212  std::cout << "trackArray nr of entries: " << trackArray.GetEntries() << "\n";
213  tree->Fill();
214 
215 
216  if (iEvent < 1000) {
217  // add tracks to event display
218  display->addEvent(tracks);
219  }
220 
221  } // end loop over events
222 
223  delete fitter;
224 
225  // write and close files
226  trackFile->Write();
227  trackFile->Close();
228  /*vertexFile->Write();
229  vertexFile->Close();*/
230 
231  // open event display
232  //display->open();
233 
234 }
235 
236