vertex finder moved from GlobalESDConverter to a separate component GlobalVertexer
[u/mrichter/AliRoot.git] / HLT / global / AliHLTGlobalVertexerComponent.cxx
1 //**************************************************************************
2 //* This file is property of and copyright by the ALICE HLT Project        * 
3 //* ALICE Experiment at CERN, All rights reserved.                         *
4 //*                                                                        *
5 //* Primary Authors: S.Gorbunov <sergey.gorbunov@kip.uni-heidelberg.de>    *
6 //*                  for The ALICE HLT Project.                            *
7 //*                                                                        *
8 //* Permission to use, copy, modify and distribute this software and its   *
9 //* documentation strictly for non-commercial purposes is hereby granted   *
10 //* without fee, provided that the above copyright notice appears in all   *
11 //* copies and that both the copyright notice and this permission notice   *
12 //* appear in the supporting documentation. The authors make no claims     *
13 //* about the suitability of this software for any purpose. It is          *
14 //* provided "as is" without express or implied warranty.                  *
15 //**************************************************************************
16
17 /** @file   AliHLTGlobalVertexerComponent.cxx
18     @author Sergey Gorbunov
19     @brief  Component for reconstruct primary vertex and V0's
20 */
21
22 #if __GNUC__>= 3
23 using namespace std;
24 #endif
25
26 #include "AliHLTGlobalVertexerComponent.h"
27 #include "AliCDBEntry.h"
28 #include "AliCDBManager.h"
29 #include <TFile.h>
30 #include <TString.h>
31 #include "TObjString.h"
32 #include "TObjArray.h"
33 #include "TH1F.h"
34 #include "TH2F.h"
35 #include "AliESDEvent.h"
36 #include "AliESDtrack.h"
37 #include "AliESDVertex.h"
38 #include "AliESDv0.h"
39 #include "AliHLTMessage.h"
40 #include "AliHLTVertexer.h"
41 #include "TMath.h"
42 #include "AliKFParticle.h"
43 #include "AliKFVertex.h"
44
45
46 /** ROOT macro for the implementation of ROOT specific class methods */
47 ClassImp(AliHLTGlobalVertexerComponent)
48
49 AliHLTGlobalVertexerComponent::AliHLTGlobalVertexerComponent()
50 :
51   fESD(0),
52   fTrackInfos(0),
53   fPrimaryVtx(),  
54   fHistPrimVertexXY(0),
55   fHistPrimVertexZX(0),
56   fHistPrimVertexZY(0),
57   fNEvents(0),
58   fPlotHistograms(1),
59   fFillVtxConstrainedTracks(1),
60   fConstrainedTrackDeviation(4.),
61   fV0DaughterPrimDeviation( 2.5 ),
62   fV0PrimDeviation( 3.5 ),
63   fV0Chi(3.5),
64   fV0DecayLengthInSigmas(3.)
65 {
66   // see header file for class documentation
67   // or
68   // refer to README to build package
69   // or
70   // visit http://web.ift.uib.no/~kjeks/doc/alice-hlt
71
72 }
73
74 AliHLTGlobalVertexerComponent::~AliHLTGlobalVertexerComponent()
75 {
76   // see header file for class documentation
77
78   delete[] fTrackInfos;
79   delete   fHistPrimVertexXY;
80   delete   fHistPrimVertexZX;
81   delete   fHistPrimVertexZY;
82 }
83
84 // Public functions to implement AliHLTComponent's interface.
85 // These functions are required for the registration process
86
87 const char* AliHLTGlobalVertexerComponent::GetComponentID()
88 {
89   // see header file for class documentation
90   
91   return "GlobalVertexer";
92 }
93
94 void AliHLTGlobalVertexerComponent::GetInputDataTypes(AliHLTComponentDataTypeList& list)
95 {
96   // see header file for class documentation
97   list.clear();
98   list.push_back( kAliHLTDataTypeESDObject|kAliHLTDataOriginOut );
99 }
100
101 AliHLTComponentDataType AliHLTGlobalVertexerComponent::GetOutputDataType()
102 {
103   // see header file for class documentation
104   return kAliHLTMultipleDataType;
105 }
106
107 int AliHLTGlobalVertexerComponent::GetOutputDataTypes(AliHLTComponentDataTypeList& tgtList)
108
109 {
110   // see header file for class documentation
111   tgtList.clear();
112   tgtList.push_back(kAliHLTDataTypeESDObject|kAliHLTDataOriginOut);
113   tgtList.push_back(kAliHLTDataTypeHistogram);
114   return tgtList.size();
115 }
116
117 void AliHLTGlobalVertexerComponent::GetOutputDataSize( unsigned long& constBase, double& inputMultiplier )
118 {
119   // see header file for class documentation
120   // XXX TODO: Find more realistic values.
121   constBase = 80000;
122   inputMultiplier = 2.;
123 }
124
125 AliHLTComponent* AliHLTGlobalVertexerComponent::Spawn()
126 {
127   // see header file for class documentation
128   return new AliHLTGlobalVertexerComponent;
129 }
130
131 int AliHLTGlobalVertexerComponent::DoInit( int argc, const char** argv )
132 {
133   // init
134   
135   fHistPrimVertexXY = new TH2F("primVertexXY","HLT:  Primary vertex distribution in XY",60,-5.,5.,60,-5.,5.);
136   fHistPrimVertexXY->SetMarkerStyle(8);
137   fHistPrimVertexXY->SetMarkerSize(0.4);
138   fHistPrimVertexXY->SetYTitle("Y");
139   fHistPrimVertexXY->SetXTitle("X");
140   fHistPrimVertexXY->SetStats(0);
141   fHistPrimVertexXY->SetOption("COLZ");
142
143   fHistPrimVertexZX = new TH2F("primVertexZX","HLT:  Primary vertex distribution in ZX",60,-15.,15.,60,-5.,5.);
144   fHistPrimVertexZX->SetMarkerStyle(8);
145   fHistPrimVertexZX->SetMarkerSize(0.4);
146   fHistPrimVertexZX->SetYTitle("X");
147   fHistPrimVertexZX->SetXTitle("Z");
148   fHistPrimVertexZX->SetStats(0);
149   fHistPrimVertexZX->SetOption("COLZ");
150
151   fHistPrimVertexZY = new TH2F("primVertexZY","HLT:  Primary vertex distribution in ZY",60,-15.,15.,60,-5.,5.);
152   fHistPrimVertexZY->SetMarkerStyle(8);
153   fHistPrimVertexZY->SetMarkerSize(0.4);
154   fHistPrimVertexZY->SetYTitle("Y");
155   fHistPrimVertexZY->SetXTitle("Z");
156   fHistPrimVertexZY->SetStats(0);
157   fHistPrimVertexZY->SetOption("COLZ");
158
159   fNEvents =0;
160
161   int iResult=0;
162   TString configuration="";
163   TString argument="";
164   for (int i=0; i<argc && iResult>=0; i++) {
165     argument=argv[i];
166     if (!configuration.IsNull()) configuration+=" ";
167     configuration+=argument;
168   }
169   
170   if (!configuration.IsNull()) {
171     iResult=Configure(configuration.Data());
172   }  
173
174   return iResult; 
175 }
176   
177 int AliHLTGlobalVertexerComponent::DoDeinit()
178 {
179   // see header file for class documentation
180
181   delete[] fTrackInfos;
182   delete fHistPrimVertexXY;
183   delete fHistPrimVertexZX;
184   delete fHistPrimVertexZY;
185
186   fTrackInfos = 0;
187   fHistPrimVertexXY = 0;
188   fHistPrimVertexZX = 0;
189   fHistPrimVertexZY = 0;
190
191   fPlotHistograms = 1;
192   fFillVtxConstrainedTracks = 1;
193   fConstrainedTrackDeviation = 4.;
194   fV0DaughterPrimDeviation = 2.5 ;
195   fV0PrimDeviation =3.5;
196   fV0Chi = 3.5;
197   fV0DecayLengthInSigmas = 3.;
198   fNEvents = 0;
199   return 0;
200 }
201
202 int AliHLTGlobalVertexerComponent::DoEvent(const AliHLTComponentEventData& /*evtData*/, AliHLTComponentTriggerData& /*trigData*/)
203 {
204   
205   if ( GetFirstInputBlock( kAliHLTDataTypeSOR ) || GetFirstInputBlock( kAliHLTDataTypeEOR ) )
206     return 0;
207
208   fNEvents++;
209
210   for ( const TObject *iter = GetFirstInputObject(kAliHLTDataTypeESDObject); iter != NULL; iter = GetNextInputObject() ) {
211
212     AliESDEvent *event = dynamic_cast<AliESDEvent*>(const_cast<TObject*>( iter ) );
213     event->GetStdContent();
214
215     // primary vertex & V0's 
216       
217     AliHLTVertexer vertexer;
218     vertexer.SetESD( event );
219     vertexer.SetFillVtxConstrainedTracks( fFillVtxConstrainedTracks );
220     vertexer.FindPrimaryVertex();
221     vertexer.FindV0s();
222     const AliESDVertex *vPrim = event->GetPrimaryVertexTracks();
223
224     // plot histograms
225
226     if( fPlotHistograms ){
227       if( vPrim && vPrim->GetNContributors()>=3 ){
228         fHistPrimVertexXY->Fill( vPrim->GetX(),vPrim->GetY()  );
229         fHistPrimVertexZX->Fill( vPrim->GetZ(),vPrim->GetX()  );
230         fHistPrimVertexZY->Fill( vPrim->GetZ(),vPrim->GetY()  );    
231       }  
232       if( fHistPrimVertexXY ) PushBack( (TObject*) fHistPrimVertexXY, kAliHLTDataTypeHistogram,0);
233       if( fHistPrimVertexZX ) PushBack( (TObject*) fHistPrimVertexZX, kAliHLTDataTypeHistogram,0);
234       if( fHistPrimVertexZY ) PushBack( (TObject*) fHistPrimVertexZY, kAliHLTDataTypeHistogram,0);
235     }
236     
237     PushBack( (TObject*) event, kAliHLTDataTypeESDObject, 0);
238   }
239     
240   return 0;
241 }
242
243 int AliHLTGlobalVertexerComponent::Configure(const char* arguments)
244 {
245   
246   int iResult=0;
247   if (!arguments) return iResult;
248   
249   TString allArgs=arguments;
250   TString argument;
251   
252   TObjArray* pTokens=allArgs.Tokenize(" ");
253   int bMissingParam=0;
254
255   if (pTokens) {
256     for (int i=0; i<pTokens->GetEntries() && iResult>=0; i++) {
257       argument=((TObjString*)pTokens->At(i))->GetString();
258       if (argument.IsNull()) continue;
259       
260       if (argument.CompareTo("-fitTracksToVertex")==0) {
261         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
262         HLTInfo("fitTracksToVertex is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
263         fFillVtxConstrainedTracks=((TObjString*)pTokens->At(i))->GetString().Atoi();
264         continue;
265       }
266       else if (argument.CompareTo("-plotHistograms")==0) {
267         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
268         HLTInfo("plotHistograms is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
269         fPlotHistograms=((TObjString*)pTokens->At(i))->GetString().Atoi();
270         continue;
271       }
272       else if (argument.CompareTo("-fillVtxConstrainedTracks")==0) {
273         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
274         HLTInfo("Filling of vertex constrained tracks is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
275         fFillVtxConstrainedTracks=((TObjString*)pTokens->At(i))->GetString().Atoi();
276         continue;
277       }
278       else if (argument.CompareTo("-constrainedTrackDeviation")==0) {
279         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
280         HLTInfo("constrainedTrackDeviation is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
281         fConstrainedTrackDeviation=((TObjString*)pTokens->At(i))->GetString().Atof();
282         continue;
283       }
284       else if (argument.CompareTo("-v0DaughterPrimDeviation")==0) {
285         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
286         HLTInfo("v0DaughterPrimDeviation is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
287         fV0DaughterPrimDeviation=((TObjString*)pTokens->At(i))->GetString().Atof();
288         continue;
289       }
290       else if (argument.CompareTo("-v0PrimDeviation")==0) {
291         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
292         HLTInfo("v0PrimDeviation is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
293         fV0PrimDeviation=((TObjString*)pTokens->At(i))->GetString().Atof();
294         continue;
295       }
296       else if (argument.CompareTo("-v0Chi")==0) {
297         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
298         HLTInfo("v0Chi is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
299         fV0Chi=((TObjString*)pTokens->At(i))->GetString().Atof();
300         continue;
301       }
302       else if (argument.CompareTo("-v0DecayLengthInSigmas")==0) {
303         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
304         HLTInfo("v0DecayLengthInSigmas is set set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
305         fV0DecayLengthInSigmas=((TObjString*)pTokens->At(i))->GetString().Atof();
306         continue;
307       }
308       else {
309         HLTError("unknown argument %s", argument.Data());
310         iResult=-EINVAL;
311         break;
312       }
313     }
314     delete pTokens;
315   }
316   if (bMissingParam) {
317     HLTError("missing parameter for argument %s", argument.Data());
318     iResult=-EINVAL;
319   }  
320   
321   return iResult;
322 }
323
324 int AliHLTGlobalVertexerComponent::Reconfigure(const char* cdbEntry, const char* chainId)
325 {
326   // see header file for class documentation
327   int iResult=0;
328   const char* path="HLT/ConfigTPC/KryptonHistoComponent";
329   const char* defaultNotify="";
330   if (cdbEntry) {
331     path=cdbEntry;
332     defaultNotify=" (default)";
333   }
334   if (path) {
335     HLTInfo("reconfigure from entry %s%s, chain id %s", path, defaultNotify,(chainId!=NULL && chainId[0]!=0)?chainId:"<none>");
336     AliCDBEntry *pEntry = AliCDBManager::Instance()->Get(path/*,GetRunNo()*/);
337     if (pEntry) {
338       TObjString* pString=dynamic_cast<TObjString*>(pEntry->GetObject());
339       if (pString) {
340         HLTInfo("received configuration object string: \'%s\'", pString->GetString().Data());
341         iResult=Configure(pString->GetString().Data());
342       } else {
343         HLTError("configuration object \"%s\" has wrong type, required TObjString", path);
344       }
345     } else {
346       HLTError("can not fetch object \"%s\" from CDB", path);
347     }
348   }
349
350   return iResult;
351 }
352
353
354
355 void AliHLTGlobalVertexerComponent::SetESD( AliESDEvent *event )
356 {
357   //* Fill fTrackInfo array
358
359   delete[] fTrackInfos;
360   fESD = event;
361
362   AliKFParticle::SetField( fESD->GetMagneticField() );
363
364   Int_t nESDTracks=event->GetNumberOfTracks(); 
365   fTrackInfos = new AliESDTrackInfo[ nESDTracks ];
366
367   for (Int_t iTr=0; iTr<nESDTracks; iTr++){ 
368   
369     AliESDTrackInfo &info = fTrackInfos[iTr];
370     info.fOK = 0;
371     info.fPrimUsedFlag = 0;
372     
373     //* track quality check
374
375     AliESDtrack *pTrack = event->GetTrack(iTr);    
376     if( !pTrack  ) continue;
377     if (pTrack->GetKinkIndex(0)>0) continue;
378     if ( !( pTrack->GetStatus()&AliESDtrack::kTPCin ) ) continue;
379     
380     //* Construct KFParticle for the track
381
382     info.fParticle = AliKFParticle( *pTrack->GetInnerParam(), 211 );    
383     info.fOK = 1;
384   }
385 }
386
387
388 void AliHLTGlobalVertexerComponent::FindPrimaryVertex(  )
389 {
390   //* Find event primary vertex
391
392   int nTracks = fESD->GetNumberOfTracks();
393
394   const AliKFParticle **vSelected = new const AliKFParticle*[nTracks]; //* Selected particles for vertex fit
395   Int_t *vIndex = new int [nTracks];                    //* Indices of selected particles
396   Bool_t *vFlag = new bool [nTracks];                    //* Flags returned by the vertex finder
397
398   fPrimaryVtx.Initialize();
399   fPrimaryVtx.SetBeamConstraint(fESD->GetDiamondX(),fESD->GetDiamondY(),0,
400                                 TMath::Sqrt(fESD->GetSigma2DiamondX()),TMath::Sqrt(fESD->GetSigma2DiamondY()),5.3);
401   
402   Int_t nSelected = 0;
403   for( Int_t i = 0; i<nTracks; i++){ 
404     if(!fTrackInfos[i].fOK ) continue;
405     //if( fESD->GetTrack(i)->GetTPCNcls()<60  ) continue;
406     const AliKFParticle &p = fTrackInfos[i].fParticle;
407     Double_t chi = p.GetDeviationFromVertex( fPrimaryVtx );      
408     if( chi > fConstrainedTrackDeviation ) continue;
409     vSelected[nSelected] = &(fTrackInfos[i].fParticle);
410     vIndex[nSelected] = i;
411     nSelected++;  
412   }
413   fPrimaryVtx.ConstructPrimaryVertex( vSelected, nSelected, vFlag, fConstrainedTrackDeviation );
414   for( Int_t i = 0; i<nSelected; i++){ 
415     if( vFlag[i] ) fTrackInfos[vIndex[i]].fPrimUsedFlag = 1;
416   }
417
418   for( Int_t i = 0; i<nTracks; i++ ){
419     AliESDTrackInfo &info = fTrackInfos[i];
420     info.fPrimDeviation = info.fParticle.GetDeviationFromVertex( fPrimaryVtx );   
421   }
422   
423   if( fPrimaryVtx.GetNContributors()>=3 ){
424     AliESDVertex vESD( fPrimaryVtx.Parameters(), fPrimaryVtx.CovarianceMatrix(), fPrimaryVtx.GetChi2(), fPrimaryVtx.GetNContributors() );
425     fESD->SetPrimaryVertexTracks( &vESD );
426
427     // relate the tracks to vertex
428
429     if( fFillVtxConstrainedTracks ){      
430       for( Int_t i = 0; i<nTracks; i++ ){
431         if( !fTrackInfos[i].fPrimUsedFlag ) continue;     
432         if( fTrackInfos[i].fPrimDeviation > fConstrainedTrackDeviation ) continue;
433         fESD->GetTrack(i)->RelateToVertex( &vESD, fESD->GetMagneticField(),100. );
434       }
435     }
436
437   } else {
438     for( Int_t i = 0; i<nTracks; i++)
439       fTrackInfos[i].fPrimUsedFlag = 0;
440   }
441
442
443   delete[] vSelected;
444   delete[] vIndex;
445   delete[] vFlag;
446 }
447
448
449 void AliHLTGlobalVertexerComponent::FindV0s(  )
450 {
451   //* V0 finder
452
453   int nTracks = fESD->GetNumberOfTracks();
454   //AliKFVertex primVtx( *fESD->GetPrimaryVertexTracks() );
455   AliKFVertex &primVtx = fPrimaryVtx;
456   if( primVtx.GetNContributors()<3 ) return;
457
458   bool *constrainedV0   = new bool[nTracks];
459   for( Int_t iTr = 0; iTr<nTracks; iTr++ ){ 
460     constrainedV0[iTr] = 0;
461   }
462   
463
464   for( Int_t iTr = 0; iTr<nTracks; iTr++ ){ //* first daughter
465
466     AliESDTrackInfo &info = fTrackInfos[iTr];
467     if( !info.fOK ) continue;    
468     if( info.fParticle.GetQ() >0 ) continue;    
469     if( info.fPrimDeviation < fV0DaughterPrimDeviation ) continue;
470
471     for( Int_t jTr = 0; jTr<nTracks; jTr++ ){  //* second daughter
472
473       AliESDTrackInfo &jnfo = fTrackInfos[jTr];
474       if( !jnfo.fOK ) continue;
475       if( jnfo.fParticle.GetQ() < 0 ) continue;
476       if( jnfo.fPrimDeviation < fV0DaughterPrimDeviation ) continue;
477    
478       //* construct V0 mother
479
480       AliKFParticle v0( info.fParticle, jnfo.fParticle );     
481
482       //* check V0 Chi^2
483       
484       if( v0.GetChi2()<0 || v0.GetChi2() > fV0Chi*fV0Chi*v0.GetNDF() ) continue;
485
486       //* subtruct daughters from primary vertex 
487
488       AliKFVertex primVtxCopy = primVtx;    
489        
490       if( info.fPrimUsedFlag ){ 
491         if( primVtxCopy.GetNContributors()<=2 ) continue;
492         primVtxCopy -= info.fParticle;
493       }
494       if( jnfo.fPrimUsedFlag ){
495         if( primVtxCopy.GetNContributors()<=2 ) continue;
496         primVtxCopy -= jnfo.fParticle;
497       }
498       //* Check v0 Chi^2 deviation from primary vertex 
499
500       if( v0.GetDeviationFromVertex( primVtxCopy ) > fV0PrimDeviation ) continue;
501
502       //* Add V0 to primary vertex to improve the primary vertex resolution
503
504       primVtxCopy += v0;      
505
506       //* Set production vertex for V0
507
508       v0.SetProductionVertex( primVtxCopy );
509
510       //* Get V0 decay length with estimated error
511
512       Double_t length, sigmaLength;
513       if( v0.GetDecayLength( length, sigmaLength ) ) continue;
514
515       //* Reject V0 if it decays too close[sigma] to the primary vertex
516
517       if( length  < fV0DecayLengthInSigmas*sigmaLength ) continue;
518
519       //* add ESD v0 
520       
521       AliESDv0 v0ESD( *fESD->GetTrack( iTr ), iTr, *fESD->GetTrack( jTr ), jTr );  
522       fESD->AddV0( &v0ESD );
523
524       // relate the tracks to vertex
525
526       if( fFillVtxConstrainedTracks ){
527         if( constrainedV0[iTr] || constrainedV0[jTr]
528             || info.fPrimDeviation < fConstrainedTrackDeviation || jnfo.fPrimDeviation < fConstrainedTrackDeviation ) continue;
529         AliESDVertex vESD(v0.Parameters(), v0.CovarianceMatrix(), v0.GetChi2(), 2);
530         fESD->GetTrack(iTr)->RelateToVertex( &vESD, fESD->GetMagneticField(),100. );
531         fESD->GetTrack(jTr)->RelateToVertex( &vESD, fESD->GetMagneticField(),100. );
532         constrainedV0[iTr] = 1;
533         constrainedV0[jTr] = 1; 
534       }
535     }
536   }
537   delete[] constrainedV0;
538 }
539