Using common HLT track data format for TRD tracks (Theodor)
[u/mrichter/AliRoot.git] / HLT / global / AliHLTGlobalEsdConverterComponent.cxx
1 // $Id$
2
3 //**************************************************************************
4 //* This file is property of and copyright by the ALICE HLT Project        * 
5 //* ALICE Experiment at CERN, All rights reserved.                         *
6 //*                                                                        *
7 //* Primary Authors: Matthias Richter <Matthias.Richter@ift.uib.no>        *
8 //*                  for The ALICE HLT Project.                            *
9 //*                                                                        *
10 //* Permission to use, copy, modify and distribute this software and its   *
11 //* documentation strictly for non-commercial purposes is hereby granted   *
12 //* without fee, provided that the above copyright notice appears in all   *
13 //* copies and that both the copyright notice and this permission notice   *
14 //* appear in the supporting documentation. The authors make no claims     *
15 //* about the suitability of this software for any purpose. It is          *
16 //* provided "as is" without express or implied warranty.                  *
17 //**************************************************************************
18
19 /** @file   AliHLTGlobalEsdConverterComponent.cxx
20     @author Matthias Richter
21     @date   
22     @brief  Global ESD converter component.
23 */
24
25 // see header file for class documentation
26 // or
27 // refer to README to build package
28 // or
29 // visit http://web.ift.uib.no/~kjeks/doc/alice-hlt
30
31 #include <cassert>
32 #include "AliHLTGlobalEsdConverterComponent.h"
33 #include "AliHLTGlobalBarrelTrack.h"
34 #include "AliHLTExternalTrackParam.h"
35 #include "AliHLTTrackMCLabel.h"
36 #include "AliESDEvent.h"
37 #include "AliESDtrack.h"
38 #include "AliCDBEntry.h"
39 #include "AliCDBManager.h"
40 #include "AliPID.h"
41 #include "TTree.h"
42 #include "TList.h"
43
44 /** ROOT macro for the implementation of ROOT specific class methods */
45 ClassImp(AliHLTGlobalEsdConverterComponent)
46
47 AliHLTGlobalEsdConverterComponent::AliHLTGlobalEsdConverterComponent()
48   : AliHLTProcessor()
49   , fESD(NULL)
50   , fSolenoidBz(5)
51   , fWriteTree(0)
52   , fVerbosity(0)
53 {
54   // see header file for class documentation
55   // or
56   // refer to README to build package
57   // or
58   // visit http://web.ift.uib.no/~kjeks/doc/alice-hlt
59 }
60
61 AliHLTGlobalEsdConverterComponent::~AliHLTGlobalEsdConverterComponent()
62 {
63   // see header file for class documentation
64   if (fESD) delete fESD;
65   fESD=NULL;
66 }
67
68 int AliHLTGlobalEsdConverterComponent::Configure(const char* arguments)
69 {
70   // see header file for class documentation
71   int iResult=0;
72   if (!arguments) return iResult;
73
74   TString allArgs=arguments;
75   TString argument;
76   int bMissingParam=0;
77
78   TObjArray* pTokens=allArgs.Tokenize(" ");
79   if (pTokens) {
80     for (int i=0; i<pTokens->GetEntries() && iResult>=0; i++) {
81       argument=((TObjString*)pTokens->At(i))->GetString();
82       if (argument.IsNull()) continue;
83       
84       if (argument.CompareTo("-solenoidBz")==0) {
85         if ((bMissingParam=(++i>=pTokens->GetEntries()))) break;
86         HLTInfo("Magnetic Field set to: %s", ((TObjString*)pTokens->At(i))->GetString().Data());
87         fSolenoidBz=((TObjString*)pTokens->At(i))->GetString().Atof();
88         continue;
89       } else {
90         HLTError("unknown argument %s", argument.Data());
91         iResult=-EINVAL;
92         break;
93       }
94     }
95     delete pTokens;
96   }
97   if (bMissingParam) {
98     HLTError("missing parameter for argument %s", argument.Data());
99     iResult=-EINVAL;
100   }
101
102   return iResult;
103 }
104
105 int AliHLTGlobalEsdConverterComponent::Reconfigure(const char* cdbEntry, const char* chainId)
106 {
107   // see header file for class documentation
108   int iResult=0;
109   const char* path=kAliHLTCDBSolenoidBz;
110   const char* defaultNotify="";
111   if (cdbEntry) {
112     path=cdbEntry;
113     defaultNotify=" (default)";
114   }
115   if (path) {
116     HLTInfo("reconfigure from entry %s%s, chain id %s", path, defaultNotify,(chainId!=NULL && chainId[0]!=0)?chainId:"<none>");
117     AliCDBEntry *pEntry = AliCDBManager::Instance()->Get(path/*,GetRunNo()*/);
118     if (pEntry) {
119       TObjString* pString=dynamic_cast<TObjString*>(pEntry->GetObject());
120       if (pString) {
121         HLTInfo("received configuration object string: \'%s\'", pString->GetString().Data());
122         iResult=Configure(pString->GetString().Data());
123       } else {
124         HLTError("configuration object \"%s\" has wrong type, required TObjString", path);
125       }
126     } else {
127       HLTError("can not fetch object \"%s\" from CDB", path);
128     }
129   }
130   
131   return iResult;
132 }
133
134 void AliHLTGlobalEsdConverterComponent::GetInputDataTypes(AliHLTComponentDataTypeList& list)
135 {
136   // see header file for class documentation
137   list.push_back(kAliHLTDataTypeTrack);
138   list.push_back(kAliHLTDataTypeTrackMC);
139 }
140
141 AliHLTComponentDataType AliHLTGlobalEsdConverterComponent::GetOutputDataType()
142 {
143   // see header file for class documentation
144   return kAliHLTDataTypeESDObject|kAliHLTDataOriginOut;
145 }
146
147 void AliHLTGlobalEsdConverterComponent::GetOutputDataSize(unsigned long& constBase, double& inputMultiplier)
148 {
149   // see header file for class documentation
150   constBase=2000000;
151   inputMultiplier=10.0;
152 }
153
154 int AliHLTGlobalEsdConverterComponent::DoInit(int argc, const char** argv)
155 {
156   // see header file for class documentation
157   int iResult=0;
158   TString argument="";
159   int bMissingParam=0;
160   for (int i=0; i<argc && iResult>=0; i++) {
161     argument=argv[i];
162     if (argument.IsNull()) continue;
163
164     // -notree
165     if (argument.CompareTo("-notree")==0) {
166       fWriteTree=0;
167
168       // -tree
169     } else if (argument.CompareTo("-tree")==0) {
170       fWriteTree=1;
171
172     } else {
173       HLTError("unknown argument %s", argument.Data());
174       break;
175     }
176   }
177   if (bMissingParam) {
178     HLTError("missing parameter for argument %s", argument.Data());
179     iResult=-EINVAL;
180   }
181
182   if (iResult>=0) {
183     iResult=Reconfigure(NULL, NULL);
184   }
185
186   if (iResult>=0) {
187     fESD = new AliESDEvent;
188     if (fESD) {
189       fESD->CreateStdContent();
190     } else {
191       iResult=-ENOMEM;
192     }
193   }
194
195   return iResult;
196 }
197
198 int AliHLTGlobalEsdConverterComponent::DoDeinit()
199 {
200   // see header file for class documentation
201   if (fESD) delete fESD;
202   fESD=NULL;
203
204   return 0;
205 }
206
207 int AliHLTGlobalEsdConverterComponent::DoEvent(const AliHLTComponentEventData& /*evtData*/, 
208                                                AliHLTComponentTriggerData& /*trigData*/)
209 {
210   // see header file for class documentation
211   int iResult=0;
212   if (!fESD) return -ENODEV;
213
214   AliESDEvent* pESD = fESD;
215
216   pESD->Reset(); 
217   pESD->SetMagneticField(fSolenoidBz);
218
219   TTree* pTree = NULL;
220   if (fWriteTree)
221     pTree = new TTree("esdTree", "Tree with HLT ESD objects");
222  
223   if (pTree) {
224     pTree->SetDirectory(0);
225   }
226
227   if ((iResult=ProcessBlocks(pTree, pESD))>0) {
228     // TODO: set the specification correctly
229     if (pTree) {
230       // the esd structure is written to the user info and is
231       // needed in te ReadFromTree method to read all objects correctly
232       pTree->GetUserInfo()->Add(pESD);
233       pESD->WriteToTree(pTree);
234       iResult=PushBack(pTree, kAliHLTDataTypeESDTree|kAliHLTDataOriginOut, 0);
235     } else {
236       iResult=PushBack(pESD, kAliHLTDataTypeESDObject|kAliHLTDataOriginOut, 0);
237     }
238   }
239   if (pTree) {
240     // clear user info list to prevent objects from being deleted
241     pTree->GetUserInfo()->Clear();
242     delete pTree;
243   }
244   return iResult;
245 }
246
247 int AliHLTGlobalEsdConverterComponent::ProcessBlocks(TTree* pTree, AliESDEvent* pESD)
248 {
249   // see header file for class documentation
250
251   int iResult=0;
252   int iAddedDataBlocks=0;
253
254   // Barrel tracking
255
256   // in the first attempt this component reads the TPC tracks and updates in the
257   // second step from the ITS tracks
258
259
260   // first read MC information (if present)
261   std::map<int,int> mcLabels;
262
263   for (const AliHLTComponentBlockData* pBlock=GetFirstInputBlock(kAliHLTDataTypeTrackMC|kAliHLTDataOriginTPC);
264        pBlock!=NULL; pBlock=GetNextInputBlock()) {
265     AliHLTTrackMCData* dataPtr = reinterpret_cast<AliHLTTrackMCData*>( pBlock->fPtr );
266     if (sizeof(AliHLTTrackMCData)+dataPtr->fCount*sizeof(AliHLTTrackMCLabel)==pBlock->fSize) {
267       for( unsigned int il=0; il<dataPtr->fCount; il++ ){
268         AliHLTTrackMCLabel &lab = dataPtr->fLabels[il];
269         mcLabels[lab.fTrackID] = lab.fMCLabel;
270       }
271     } else {
272       HLTWarning("data mismatch in block %s (0x%08x): count %d, size %d -> ignoring track MC information", 
273                  DataType2Text(pBlock->fDataType).c_str(), pBlock->fSpecification, 
274                  dataPtr->fCount, pBlock->fSize);
275     }
276   }
277
278   // convert the TPC tracks to ESD tracks
279   for (const AliHLTComponentBlockData* pBlock=GetFirstInputBlock(kAliHLTDataTypeTrack|kAliHLTDataOriginTPC);
280        pBlock!=NULL; pBlock=GetNextInputBlock()) {
281     vector<AliHLTGlobalBarrelTrack> tracks;
282     if ((iResult=AliHLTGlobalBarrelTrack::ConvertTrackDataArray(reinterpret_cast<const AliHLTTracksData*>(pBlock->fPtr), pBlock->fSize, tracks))>0) {
283       for (vector<AliHLTGlobalBarrelTrack>::iterator element=tracks.begin();
284            element!=tracks.end(); element++) {
285         Float_t points[4] = {
286           element->GetX(),
287           element->GetY(),
288           element->GetLastPointX(),
289           element->GetLastPointY()
290         };
291
292         Int_t mcLabel = -1;
293         if( mcLabels.find(element->TrackID())!=mcLabels.end() )
294           mcLabel = mcLabels[element->TrackID()];
295         element->SetLabel( mcLabel );
296
297         AliESDtrack iotrack;
298         iotrack.UpdateTrackParams(&(*element),AliESDtrack::kTPCin);
299         iotrack.SetTPCPoints(points);
300
301         pESD->AddTrack(&iotrack);
302         if (fVerbosity>0) element->Print();
303       }
304       HLTInfo("converted %d track(s) to AliESDtrack and added to ESD", tracks.size());
305       iAddedDataBlocks++;
306     } else if (iResult<0) {
307       HLTError("can not extract tracks from data block of type %s (specification %08x) of size %d: error %d", 
308                DataType2Text(pBlock->fDataType).c_str(), pBlock->fSpecification, pBlock->fSize, iResult);
309     }
310   }
311
312   // now update ESD tracks with the ITS info
313   for (const AliHLTComponentBlockData* pBlock=GetFirstInputBlock(kAliHLTDataTypeTrack|kAliHLTDataOriginITS);
314        pBlock!=NULL; pBlock=GetNextInputBlock()) {
315     vector<AliHLTGlobalBarrelTrack> tracks;
316     if ((iResult=AliHLTGlobalBarrelTrack::ConvertTrackDataArray(reinterpret_cast<const AliHLTTracksData*>(pBlock->fPtr), pBlock->fSize, tracks))>0) {
317       for (vector<AliHLTGlobalBarrelTrack>::iterator element=tracks.begin();
318            element!=tracks.end(); element++) {
319         int ncl=0;
320         const UInt_t* pointsArray=element->GetPoints();
321         for( unsigned il=0; il<element->GetNumberOfPoints(); il++ ){
322           // TODO: check what needs to be done with the clusters
323           if( pointsArray[il]<~(UInt_t)0 ) {/*tITS.SetClusterIndex(ncl,  tr.fClusterIds[il]);*/}
324           ncl++;
325         }
326         //tITS.SetNumberOfClusters( ncl );
327         int tpcID=element->TrackID();
328         // the ITS tracker assigns the TPC track used as seed for a certain track to
329         // the trackID
330         if( tpcID<0 || tpcID>=pESD->GetNumberOfTracks()) continue;
331
332         AliESDtrack *tESD = pESD->GetTrack( tpcID );
333         if( tESD ) tESD->UpdateTrackParams( &(*element), AliESDtrack::kITSin );
334       }
335     }
336   }
337
338   // convert the HLT TRD tracks to ESD tracks                        
339   for (const AliHLTComponentBlockData* pBlock=GetFirstInputBlock(kAliHLTDataTypeTrack | kAliHLTDataOriginTRD);
340        pBlock!=NULL; pBlock=GetNextInputBlock()) {
341     vector<AliHLTGlobalBarrelTrack> tracks;
342     if ((iResult=AliHLTGlobalBarrelTrack::ConvertTrackDataArray(reinterpret_cast<const AliHLTTracksData*>(pBlock->fPtr), pBlock->fSize, tracks))>0) {
343       for (vector<AliHLTGlobalBarrelTrack>::iterator element=tracks.begin();
344            element!=tracks.end(); element++) {
345         
346         Double_t TRDpid[AliPID::kSPECIES], eProb(0.2), restProb((1-eProb)/(AliPID::kSPECIES-1)); //eprob(element->GetTRDpid...);
347         for(Int_t i=0; i<AliPID::kSPECIES; i++){
348           switch(i){
349           case AliPID::kElectron: TRDpid[AliPID::kElectron]=eProb; break;
350           default: TRDpid[i]=restProb; break;
351           }
352         }
353         
354         AliESDtrack iotrack;
355         iotrack.UpdateTrackParams(&(*element),AliESDtrack::kTRDin);
356         iotrack.SetTRDpid(TRDpid);
357         
358         pESD->AddTrack(&iotrack);
359         if (fVerbosity>0) element->Print();
360       }
361       HLTInfo("converted %d track(s) to AliESDtrack and added to ESD", tracks.size());
362       iAddedDataBlocks++;
363     } else if (iResult<0) {
364       HLTError("can not extract tracks from data block of type %s (specification %08x) of size %d: error %d", 
365                DataType2Text(pBlock->fDataType).c_str(), pBlock->fSpecification, pBlock->fSize, iResult);
366     }
367   }
368       
369       // primary vertex & V0's 
370       /*
371       //AliHLTVertexer vertexer;
372       //vertexer.SetESD( pESD );
373       //vertexer.FindPrimaryVertex();
374       //vertexer.FindV0s();
375       */
376       if (iAddedDataBlocks>0 && pTree) {
377        pTree->Fill();
378       }
379   
380   if (iResult>=0) iResult=iAddedDataBlocks;
381   return iResult;
382 }