Connection to tree in Notify() seems to be safer.
[u/mrichter/AliRoot.git] / FMD / AliFMDBackgroundCorrection.cxx
1 /**************************************************************************
2  * Copyright(c) 2004, ALICE Experiment at CERN, All rights reserved. *
3  *                                                                        *
4  * Author: The ALICE Off-line Project.                                    *
5  * Contributors are mentioned in the code where appropriate.              *
6  *                                                                        *
7  * Permission to use, copy, modify and distribute this software and its   *
8  * documentation strictly for non-commercial purposes is hereby granted   *
9  * without fee, provided that the above copyright notice appears in all   *
10  * copies and that both the copyright notice and this permission notice   *
11  * appear in the supporting documentation. The authors make no claims     *
12  * about the suitability of this software for any purpose. It is          *
13  * provided "as is" without express or implied warranty.                  *
14  **************************************************************************/
15
16 // Thil class computes background corrections for the FMD. The
17 // correction is computed in eta,phi cells and the objects stored can
18 // be put into alien to use with analysis.  It is based on the
19 // AliFMDInput class that is used to loop over hits and primaries.
20 //
21 // Author: Hans Hjersing Dalsgaard, NBI, hans.dalsgaard@cern.ch
22 //
23 //
24
25 #include "AliSimulation.h"
26 #include "TStopwatch.h"
27 #include "iostream"
28 #include "TGrid.h"
29 #include "AliRunLoader.h"
30 #include "AliGeomManager.h"
31 #include "AliFMDGeometry.h"
32 #include "AliStack.h"
33 #include "TParticle.h"
34 #include "TDatabasePDG.h"
35 #include "TParticlePDG.h"
36 #include "AliRun.h"
37 #include "AliFMDBackgroundCorrection.h"
38 #include "TSystem.h"
39 #include "AliCDBManager.h"
40 #include "TTree.h"
41 #include "TClonesArray.h"
42 #include "TBranch.h"
43 #include "AliFMDHit.h"
44 #include "AliLoader.h" 
45 #include "AliFMD.h"
46 #include "AliGenEventHeader.h"
47 #include "AliHeader.h"
48 #include "TFile.h"
49 #include "TAxis.h"
50 #include "AliCDBId.h"
51 #include "AliCDBMetaData.h"
52 #include "TROOT.h"
53 #include "AliFMDParameters.h"
54 #include "AliLog.h"
55 #include "TList.h"
56 #include "AliFMDAnaParameters.h"
57 #include "AliFMDAnaCalibBackgroundCorrection.h"
58 #include "AliTrackReference.h"
59 #include "AliFMDStripIndex.h"
60
61 ClassImp(AliFMDBackgroundCorrection)
62 //_____________________________________________________________________
63 AliFMDBackgroundCorrection::AliFMDBackgroundCorrection() : 
64   TNamed(),
65   fCorrectionArray(),
66   fPrimaryList()
67 {} 
68
69 //_____________________________________________________________________
70 AliFMDBackgroundCorrection::AliFMDInputBG::AliFMDInputBG(Bool_t hits_not_trackref) 
71   : AliFMDInput("galice.root"),
72     fPrimaryArray(),
73     fHitArray(),
74     fPrimaryMapInner(),
75     fPrimaryMapOuter(),
76     fHitMap(0),           // nDetector=0 means 51200 slots
77     fLastTrackByStrip(0), // nDetector=0 means 51200 slots
78     fPrim(0),
79     fHits(0),
80     fZvtxCut(0),
81     fNvtxBins(0),
82     fPrevTrack(-1),
83     fPrevDetector(-1),
84     fPrevRing('Q'),
85     fPrevSec(-1),
86     fNbinsEta(100)
87 {
88   if(hits_not_trackref) {
89     AddLoad(kHits);
90     AddLoad(kTracks); 
91   }
92   else
93     AddLoad(kTrackRefs);
94   AddLoad(kKinematics); 
95   AddLoad(kHeader);
96 }
97
98 //_____________________________________________________________________
99
100 void 
101 AliFMDBackgroundCorrection::GenerateBackgroundCorrection(Bool_t from_hits,
102                                                          const Int_t nvtxbins,
103                                                          Float_t zvtxcut, 
104                                                          const Int_t nBinsEta, 
105                                                          Bool_t storeInAlien, 
106                                                          Int_t runNo,
107                                                          Int_t endRunNo, 
108                                                          const Char_t* filename,
109                                                          Bool_t simulate,
110                                                          Int_t nEvents,
111                                                          Bool_t inFile,
112                                                          const Char_t* infilename) 
113 {
114   //TGrid::Connect("alien:",0,0,"t");
115   if(simulate)
116     Simulate(nEvents);
117   else {
118     //AliCDBManager::Instance()->SetDefaultStorage("alien://Folder=/alice/data/2008/LHC08d/OCDB/");
119     AliCDBManager::Instance()->SetDefaultStorage("local://$ALICE_ROOT/OCDB");
120     AliCDBManager::Instance()->SetRun(runNo);
121     
122 #if defined(__CINT__)
123     gSystem->Load("liblhapdf");
124     gSystem->Load("libEGPythia6");
125     gSystem->Load("libpythia6");
126     gSystem->Load("libAliPythia6");
127     gSystem->Load("libgeant321");
128 #endif
129       // 
130   }  
131   
132   //Setting up the geometry
133   //-----------------------------------------------
134   if (AliGeomManager::GetGeometry() == NULL)
135     AliGeomManager::LoadGeometry();
136   
137   AliFMDGeometry* geo = AliFMDGeometry::Instance();
138   geo->Init();
139   geo->InitTransformations();
140     
141   AliInfo("Processing hits and primaries ");
142   
143   AliFMDInputBG input(from_hits);
144   
145   if(!inFile) {
146   
147     input.SetVtxCutZ(zvtxcut);
148     input.SetNvtxBins(nvtxbins);
149     input.SetNbinsEta(nBinsEta);
150     input.Run();
151   }
152   
153   AliInfo(Form("Found %d primaries and %d hits.", 
154                input.GetNprim(),input.GetNhits()));
155   TObjArray* hitArray ;
156   TObjArray* primaryArray;
157   if(inFile) {
158     TFile* infile = TFile::Open(infilename);
159     hitArray     = new TObjArray();
160     primaryArray = new TObjArray();
161     
162     for(Int_t det =1; det<=3;det++) {
163       TObjArray* detArrayHits = new TObjArray();
164       detArrayHits->SetName(Form("FMD%d",det));
165       hitArray->AddAtAndExpand(detArrayHits,det);
166       Int_t nRings = (det==1 ? 1 : 2);
167       for(Int_t ring = 0;ring<nRings;ring++) {
168         Char_t ringChar = (ring == 0 ? 'I' : 'O');
169         TObjArray* vtxArrayHits = new TObjArray();
170         vtxArrayHits->SetName(Form("FMD%d%c",det,ringChar));
171         detArrayHits->AddAtAndExpand(vtxArrayHits,ring);
172         for(Int_t v=0; v<nvtxbins;v++) {
173           
174           TH2F* hHits = 
175             static_cast<TH2F*>(infile->Get(Form("hHits_FMD%d%c_vtx%d",
176                                                 det,ringChar,v)));
177                 
178                 
179           vtxArrayHits->AddAtAndExpand(hHits,v);
180           
181         } 
182       }
183     }
184     
185     for(Int_t iring = 0; iring<2;iring++) {
186       Char_t     ringChar = (iring == 0 ? 'I' : 'O');
187       TObjArray* ringArray = new TObjArray();
188       ringArray->SetName(Form("FMD_%c",ringChar));
189       primaryArray->AddAtAndExpand(ringArray,iring);
190       for(Int_t v=0; v<nvtxbins;v++) {
191         
192         TH2F* hPrimary = 
193           static_cast<TH2F*>(infile->Get(Form("hPrimary_FMD_%c_vtx%d",
194                                               ringChar,v)));
195         ringArray->AddAtAndExpand(hPrimary,v);
196       }
197     }
198     
199     
200   }
201   else {
202     hitArray     = input.GetHits();
203     primaryArray = input.GetPrimaries();
204   }
205   fCorrectionArray.SetName("FMD_bg_correction");
206   fCorrectionArray.SetOwner();
207    
208   TList* primaryList     = new TList();
209   primaryList->SetName("primaries");
210   
211   TList* hitList     = new TList();
212   hitList->SetName("hits");
213   TList* corrList    = new TList();
214   corrList->SetName("corrections");
215   
216   AliFMDAnaCalibBackgroundCorrection* background = 
217     new AliFMDAnaCalibBackgroundCorrection();
218   
219   for(Int_t det= 1; det <=3; det++) {
220     Int_t nRings = (det==1 ? 1 : 2);
221     
222     TObjArray* detArrayCorrection = new TObjArray();
223     detArrayCorrection->SetName(Form("FMD%d",det));
224     fCorrectionArray.AddAtAndExpand(detArrayCorrection,det);
225     
226     
227     for(Int_t iring = 0; iring<nRings; iring++) {
228       TObjArray* primRingArray      = 
229         static_cast<TObjArray*>(primaryArray->At(iring));
230       Char_t     ringChar           = (iring == 0 ? 'I' : 'O');
231       TObjArray* vtxArrayCorrection = new TObjArray();
232       vtxArrayCorrection->SetName(Form("FMD%d%c",det,ringChar));
233       detArrayCorrection->AddAtAndExpand(vtxArrayCorrection,iring);
234       
235       for(Int_t vertexBin=0;vertexBin<nvtxbins;vertexBin++) {
236         TObjArray* detArray  = (TObjArray*)hitArray->At(det);
237         TObjArray* vtxArray  = (TObjArray*)detArray->At(iring);
238         TH2F* hHits          = (TH2F*)vtxArray->At(vertexBin);
239         hitList->Add(hHits);
240         TH2F* hPrimary  = static_cast<TH2F*>(primRingArray->At(vertexBin));
241         primaryList->Add(hPrimary);
242         TH2F* hCorrection = 
243           static_cast<TH2F*>(hHits->Clone(Form("FMD%d%c_vtxbin_%d_correction",
244                                                det,ringChar,vertexBin)));
245         hCorrection->Divide(hPrimary);
246         corrList->Add(hCorrection);
247         hCorrection->SetTitle(hCorrection->GetName());
248         vtxArrayCorrection->AddAtAndExpand(hCorrection,vertexBin);
249         background->SetBgCorrection(det,ringChar,vertexBin,hCorrection);
250       }
251     }
252   }
253   
254   TAxis refAxis(nvtxbins,-1*zvtxcut,zvtxcut);
255   if(inFile) {
256     TFile* infile = TFile::Open(infilename);
257     TAxis* refaxis = (TAxis*)infile->Get("vertexbins");
258     background->SetRefAxis(refaxis);
259       
260   }
261   else background->SetRefAxis(&refAxis);
262   
263   TFile*  fout = new TFile(filename,"RECREATE");
264   refAxis.Write("vertexbins");
265   
266   hitList->Write();
267   primaryList->Write();
268   corrList->Write();
269    
270   TObjArray* container = new TObjArray();
271   container->SetOwner();
272   container->AddAtAndExpand(&refAxis,0);
273   container->AddAtAndExpand(&fCorrectionArray,1);
274   container->AddAtAndExpand(hitArray,2);
275   container->AddAtAndExpand(primaryArray,3);
276   
277   
278   if(storeInAlien) {
279     AliCDBManager* cdb = AliCDBManager::Instance();
280     cdb->SetDefaultStorage("local://$ALICE_ROOT/OCDB");
281     AliCDBId      id(AliFMDAnaParameters::GetBackgroundPath(),runNo,endRunNo);
282     
283     AliCDBMetaData* meta = new AliCDBMetaData;                          
284     meta->SetResponsible(gSystem->GetUserInfo()->fRealName.Data());     
285     meta->SetAliRootVersion(gROOT->GetVersion());                       
286     meta->SetBeamPeriod(1);                                             
287     meta->SetComment("Background Correction for FMD");
288     meta->SetProperty("key1", background );
289     cdb->Put(background, id, meta);
290     
291   }
292   
293   fout->Close();
294  
295   
296   }
297
298 //_____________________________________________________________________
299 void AliFMDBackgroundCorrection::Simulate(Int_t nEvents) {
300   
301   AliSimulation sim ; 
302   sim.SetRunNumber(0);
303   TGrid::Connect("alien:",0,0,"t");
304   // FIXME: Do not use a hard-coded path!
305   sim.SetDefaultStorage("alien://Folder=/alice/data/2008/LHC08d/OCDB/");
306   sim.SetConfigFile("Config.C");
307   sim.SetRunQA("FMD:");
308   TStopwatch timer;
309   timer.Start();
310   sim.RunSimulation(nEvents);    
311   timer.Stop();
312   timer.Print();
313   
314 }
315
316 //_____________________________________________________________________
317 Bool_t 
318 AliFMDBackgroundCorrection::AliFMDInputBG::ProcessHit(AliFMDHit* h, 
319                                                       TParticle* /*p*/) 
320 {
321   
322   if(!h)
323     return kTRUE;
324   Bool_t retval = ProcessEvent(h->Detector(),
325                                h->Ring(),
326                                h->Sector(),
327                                h->Strip(),
328                                h->Track(),
329                                h->Q());
330   
331   // std::cout<<"Length: "<<h->Length()<<std::endl;
332   // std::cout<<"Is stopped "<<h->IsStop()<<"    "<<kTRUE<<std::endl;
333   return retval;
334 }
335 //_____________________________________________________________________
336 Bool_t 
337 AliFMDBackgroundCorrection::AliFMDInputBG::ProcessTrackRef(AliTrackReference* tr, TParticle* p) 
338 {
339   if(!tr)
340     return kTRUE;
341   UShort_t det,sec,strip;
342   Char_t   ring;
343   AliFMDStripIndex::Unpack(tr->UserId(),det,ring,sec,strip);
344   Int_t         nTrack  = tr->GetTrack();
345   TDatabasePDG* pdgDB   = TDatabasePDG::Instance();
346   TParticlePDG* pdgPart = pdgDB->GetParticle(p->GetPdgCode());
347   Float_t       charge  = (pdgPart ? pdgPart->Charge() : 0);
348   Bool_t        retval  = ProcessEvent(det,ring,sec,strip,nTrack,charge);
349   return retval;
350   
351 }
352 //_____________________________________________________________________
353 Bool_t 
354 AliFMDBackgroundCorrection::AliFMDInputBG::ProcessEvent(UShort_t det,
355                                                         Char_t   ring, 
356                                                         UShort_t sec, 
357                                                         UShort_t strip,
358                                                         Int_t    nTrack,
359                                                         Float_t  charge)
360 {
361   
362   if (charge == 0) return kTRUE;
363   
364   /*  if(charge !=  0 && 
365      ((nTrack != fPrevTrack) || 
366       (det != fPrevDetector))){ || 
367       (ring != fPrevRing))){    ||
368                                   (sec != fPrevSec))) {*/
369   /*Float_t nstrips = (ring =='O' ? 256 : 512);
370   
371   
372   Float_t prevStripTrack = -1;
373   if(strip >0)
374     prevStripTrack = fLastTrackByStrip.operator()(det,ring,sec,strip-1);
375   
376   Float_t nextStripTrack = -1;
377   if(strip < (nstrips - 1))
378     nextStripTrack = fLastTrackByStrip.operator()(det,ring,sec,strip+1);
379   */
380   
381   Float_t thisStripTrack = fLastTrackByStrip(det,ring,sec,strip);
382   // if(nTrack == fLastTrackByStrip.operator()(det,ring,sec,strip))
383   //  std::cout<<"Track # "<<nTrack<<"  failed the cut in "<<det<<"   "<<ring<<"   "<<sec<<"  "<<strip<<std::endl;
384   // else
385   //  std::cout<<"Track # "<<nTrack<<"  passed the cut in "<<det<<"   "<<ring<<"   "<<sec<<"  "<<strip<<std::endl;
386   if(nTrack == thisStripTrack) return kTRUE;
387     
388   fHitMap(det,ring,sec,strip) += 1;
389   fHits++;
390   Float_t nstrips = (ring =='O' ? 256 : 512);
391   fLastTrackByStrip(det,ring,sec,strip)     = Float_t(nTrack);
392   if(strip >0)
393     fLastTrackByStrip(det,ring,sec,strip-1) = Float_t(nTrack);
394   if(strip < (nstrips - 1))
395     fLastTrackByStrip(det,ring,sec,strip+1) = Float_t(nTrack);
396   
397   fPrevDetector = det;
398   fPrevRing     = ring;
399   fPrevSec      = sec;
400   fPrevTrack    = nTrack;
401   
402   return kTRUE;
403
404 }
405
406 //_____________________________________________________________________
407 Bool_t AliFMDBackgroundCorrection::AliFMDInputBG::Init() 
408 {
409   fLastTrackByStrip.Reset(-1);
410   
411   fPrimaryArray.SetOwner();
412   fPrimaryArray.SetName("FMD_primary");
413   
414   fPrimaryMapInner.SetBins(fNbinsEta, -6,6, 20, 0, 2*TMath::Pi());
415   fPrimaryMapOuter.SetBins(fNbinsEta, -6,6, 40, 0, 2*TMath::Pi());
416   fPrimaryMapInner.SetName("fPrimaryMapInner");
417   fPrimaryMapInner.SetName("fPrimaryMapOuter");
418   
419   fPrimaryMapInner.Sumw2();
420   fPrimaryMapOuter.Sumw2();
421   for(Int_t iring = 0; iring<2;iring++) {
422     Char_t ringChar = (iring == 0 ? 'I' : 'O');
423     TObjArray* ringArray = new TObjArray();
424     ringArray->SetName(Form("FMD_%c",ringChar));
425     fPrimaryArray.AddAtAndExpand(ringArray,iring);
426     Int_t nSec = (iring == 1 ? 40 : 20);
427     for(Int_t v=0; v<fNvtxBins;v++) {
428
429       TH2F* hPrimary       = new TH2F(Form("hPrimary_FMD_%c_vtx%d",ringChar,v),
430                                       Form("hPrimary_FMD_%c_vtx%d",ringChar,v),
431                                       fNbinsEta, -6,6, nSec, 0,2*TMath::Pi());
432       hPrimary->Sumw2();
433       ringArray->AddAtAndExpand(hPrimary,v);
434     }
435   }
436   
437   
438   fHitArray.SetOwner();
439   fHitArray.SetName("FMD_hits");
440    
441   for(Int_t det =1; det<=3;det++) {
442     TObjArray* detArrayHits = new TObjArray();
443     detArrayHits->SetName(Form("FMD%d",det));
444     fHitArray.AddAtAndExpand(detArrayHits,det);
445     Int_t nRings = (det==1 ? 1 : 2);
446     for(Int_t ring = 0;ring<nRings;ring++) {
447       Int_t nSec = (ring == 1 ? 40 : 20);
448       Char_t ringChar = (ring == 0 ? 'I' : 'O');
449       TObjArray* vtxArrayHits = new TObjArray();
450       vtxArrayHits->SetName(Form("FMD%d%c",det,ringChar));
451       detArrayHits->AddAtAndExpand(vtxArrayHits,ring);
452       for(Int_t v=0; v<fNvtxBins;v++) {
453         TH2F* hHits = new TH2F(Form("hHits_FMD%d%c_vtx%d", det,ringChar,v),
454                                Form("hHits_FMD%d%c_vtx%d", det,ringChar,v),
455                                fNbinsEta, -6,6, nSec, 0, 2*TMath::Pi());
456         hHits->Sumw2();
457         vtxArrayHits->AddAtAndExpand(hHits,v);
458         
459       } 
460     }
461   }
462
463   AliFMDInput::Init();
464   
465   return kTRUE;
466 }
467 //_____________________________________________________________________
468
469 Bool_t AliFMDBackgroundCorrection::AliFMDInputBG::Begin(Int_t event ) 
470 {
471
472   Bool_t             retVal    = AliFMDInput::Begin(event); 
473   AliStack*          partStack = fLoader->Stack();
474   Int_t              nTracks   = partStack->GetNtrack();
475   
476   
477   for(Int_t j=0;j<nTracks;j++) {
478     TParticle* p           = partStack->Particle(j);
479     TDatabasePDG* pdgDB    = TDatabasePDG::Instance();
480     TParticlePDG* pdgPart  = pdgDB->GetParticle(p->GetPdgCode());
481     Float_t       charge   = (pdgPart ? pdgPart->Charge() : 0);
482     if (charge == 0) continue;
483
484     Float_t   phi = TMath::ATan2(p->Py(),p->Px());
485     if(phi<0) phi = phi+2*TMath::Pi();
486     Float_t   eta = p->Eta();   
487     
488     // std::cout<<-1*TMath::Log(TMath::Tan(0.5*p->Theta()))<<std::endl;
489     
490     Bool_t primary = partStack->IsPhysicalPrimary(j);
491     //(charge!=0)&&(TMath::Abs(p->Vx() - vertex.At(0))<0.01)&&(TMath::Abs(p->Vy() - vertex.At(1))<0.01)&&(TMath::Abs(p->Vz() - vertex.At(2))<0.01);
492     if(!primary) continue;
493       
494     fPrim++;
495     
496     fPrimaryMapInner.Fill(eta,phi);
497     fPrimaryMapOuter.Fill(eta,phi);
498   }
499   
500   return retVal;
501 }
502 //_____________________________________________________________________
503 Bool_t AliFMDBackgroundCorrection::AliFMDInputBG::End()  {
504   
505   Bool_t retval = AliFMDInput::End();
506   
507   AliGenEventHeader* genHeader = fLoader->GetHeader()->GenEventHeader();
508   TArrayF vertex;
509   genHeader->PrimaryVertex(vertex);
510   
511   if(TMath::Abs(vertex.At(2)) > fZvtxCut) 
512     return kTRUE;
513   
514   Double_t delta           = 2*fZvtxCut/fNvtxBins;
515   Double_t vertexBinDouble = (vertex.At(2) + fZvtxCut) / delta;
516   Int_t    vertexBin       = Int_t(vertexBinDouble);
517   //Primaries
518   TObjArray* innerArray = static_cast<TObjArray*>(fPrimaryArray.At(0));
519   TObjArray* outerArray = static_cast<TObjArray*>(fPrimaryArray.At(1));
520   
521   TH2F* hPrimaryInner  = static_cast<TH2F*>(innerArray->At(vertexBin));
522   TH2F* hPrimaryOuter  = static_cast<TH2F*>(outerArray->At(vertexBin));
523   
524   hPrimaryInner->Add(&fPrimaryMapInner);
525   hPrimaryOuter->Add(&fPrimaryMapOuter);
526   
527   //Hits
528   for(UShort_t det=1;det<=3;det++) {
529     Int_t nRings = (det==1 ? 1 : 2);
530     for (UShort_t ir = 0; ir < nRings; ir++) {
531       Char_t   ring = (ir == 0 ? 'I' : 'O');
532       UShort_t nsec = (ir == 0 ? 20  : 40);
533       UShort_t nstr = (ir == 0 ? 512 : 256);
534       
535       for(UShort_t sec =0; sec < nsec;  sec++) {
536         
537         for(UShort_t strip = 0; strip < nstr; strip++) {
538           
539           if(fHitMap(det,ring,sec,strip) > 0) {
540             //std::cout<<fHitMap.operator()(det,ring,sec,strip)<<std::endl;
541             Double_t x,y,z;
542             AliFMDGeometry* fmdgeo = AliFMDGeometry::Instance();
543             fmdgeo->Detector2XYZ(det,ring,sec,strip,x,y,z);
544             
545             Int_t iring = (ring == 'I' ? 0 : 1);
546             
547             TObjArray* detArray  = static_cast<TObjArray*>(fHitArray.At(det));
548             TObjArray* vtxArray  = static_cast<TObjArray*>(detArray->At(iring));
549             TH2F* hHits          = static_cast<TH2F*>(vtxArray->At(vertexBin));
550             
551             Float_t   phi   = TMath::ATan2(y,x);
552             if(phi<0) phi   = phi+2*TMath::Pi();
553             Float_t   r     = TMath::Sqrt(TMath::Power(x,2)+TMath::Power(y,2));
554             Float_t   theta = TMath::ATan2(r,z-vertex.At(2));
555             Float_t   eta   = -1*TMath::Log(TMath::Tan(0.5*theta));
556             hHits->Fill(eta,phi,fHitMap.operator()(det,ring,sec,strip));
557           }
558           
559         }
560       }
561     }
562   }
563   
564   fPrimaryMapInner.Reset();
565   fPrimaryMapOuter.Reset();
566   fHitMap.Reset(0);
567   fLastTrackByStrip.Reset(-1);
568   
569   return retval;
570 }
571
572 //
573 // EOF
574 //
575
576