Add class and function docs.
[u/mrichter/AliRoot.git] / EVE / EveBase / AliEveTrackFitter.cxx
1 // $Id$
2 // Main authors: Matevz Tadel & Alja Mrak-Tadel: 2006, 2007
3
4 /**************************************************************************
5  * Copyright(c) 1998-2008, ALICE Experiment at CERN, all rights reserved. *
6  * See http://aliceinfo.cern.ch/Offline/AliRoot/License.html for          *
7  * full copyright notice.                                                 *
8  **************************************************************************/
9
10 #include "AliEveTrackFitter.h"
11
12 #include "TCanvas.h"
13 #include "TGraph.h"
14 #include "TGraphErrors.h"
15
16 #include "AliCluster.h"
17 #include "AliRieman.h"
18 #include "AliExternalTrackParam.h"
19
20 #include <TEveTrack.h>
21 #include <TEveTrackPropagator.h>
22 #include <TEveVSDStructs.h>
23 #include <TEveManager.h>
24
25
26 //______________________________________________________________________________
27 // AliEveTrackFitter
28 //
29 // AliEveTrackFitter is an interface to TEvePointSet allowing AliRieman fit.
30 // It builds a list of points by listening to selection signal of any object of type
31 // TEvePointSet. After selection the list is feeded to AliRieman fitter,
32 // which returns helix parameters visualized with TEveTrack.
33 //
34
35 ClassImp(AliEveTrackFitter)
36
37 AliEveTrackFitter::AliEveTrackFitter(const Text_t* name, Int_t n_points) :
38     TEvePointSet   (name, n_points),
39
40     fAlpha         (0),
41     fRieman        (0),
42
43     fConnected     (kFALSE),
44     fSPMap         (),
45     fTrackList     (0),
46
47     fGraphPicked   (0),
48     fGraphHelix    (0)
49 {
50   // Constructor.
51
52   SetMarkerColor(3);
53   SetOwnIds(kFALSE);
54
55   fTrackList = new TEveTrackList("Tracks");
56   fTrackList->SetLineWidth(2);
57   fTrackList->SetLineColor(8);
58   fTrackList->IncDenyDestroy();
59   fTrackList->GetPropagator()->SetEditPathMarks(kTRUE);
60   gEve->AddElement(fTrackList, this);
61   UpdateItems();
62
63   fGraphPicked = new TGraph();
64   fGraphPicked->SetName("Selected points");
65   fGraphPicked->SetMarkerColor(4);
66   fGraphPicked->SetMarkerStyle(4);
67   fGraphPicked->SetMarkerSize(2);
68
69   fGraphHelix = new TGraphErrors();
70   fGraphHelix->SetName("Fitted points");
71   fGraphHelix->SetMarkerColor(2);
72 }
73
74 AliEveTrackFitter::~AliEveTrackFitter()
75 {
76   // Destructor.
77
78   if(fRieman) delete fRieman;
79
80   fTrackList->DecDenyDestroy();
81   delete fTrackList;
82 }
83
84 /******************************************************************************/
85 void AliEveTrackFitter::Start()
86 {
87   // Clear existing point selection and maintain connection to the
88   // TEvePointSet signal.
89
90   Reset();
91   if(fConnected == kFALSE)
92   {
93     TQObject::Connect("TEvePointSet", "PointCtrlClicked(TEvePointSet*,Int_t)",
94                       "AliEveTrackFitter", this, "AddFitPoint(TEvePointSet*,Int_t)");
95     fConnected = kTRUE;
96   }
97 }
98
99 void AliEveTrackFitter::Stop()
100 {
101   // Stop adding points for the fit.
102
103   if(fConnected)
104   {
105     TQObject::Disconnect("TEvePointSet", "AddFitPoint(TEvePointSet*,Int_t)");
106     fConnected = kFALSE;
107   }
108 }
109
110 void AliEveTrackFitter::Reset(Int_t n, Int_t ids)
111 {
112   // Reset selection.
113
114   if(fRieman) fRieman->Reset();
115   TEvePointSet::Reset(n, ids);
116   fSPMap.clear();
117 }
118
119 /******************************************************************************/
120
121 void AliEveTrackFitter::AddFitPoint(TEvePointSet* ps, Int_t n)
122 {
123   // Add or remove given point depending if exists in the map.
124
125   Float_t x, y, z;
126
127   PointMap_t::iterator g = fSPMap.find(Point_t(ps, n));
128   if(g != fSPMap.end())
129   {
130     Int_t idx = g->second;
131     if(idx != fLastPoint)
132     {
133       GetPoint(fLastPoint, x, y, z);
134       SetPoint(idx, x, y, z);
135     }
136     fSPMap.erase(g);
137     fLastPoint--;
138   }
139   else
140   {
141     fSPMap[Point_t(ps, n)] = Size();
142     ps->GetPoint(n, x, y, z);
143     SetNextPoint(x, y, z);
144     SetPointId(ps->GetPointId(n));
145   }
146
147   ResetBBox();
148   ElementChanged(kTRUE, kTRUE);
149 }
150
151 /******************************************************************************/
152
153 void AliEveTrackFitter::FitTrack()
154 {
155   // Fit selected points with AliRieman fitter.
156
157   using namespace TMath;
158
159   if(fRieman) delete fRieman;
160   fRieman = new AliRieman(Size());
161
162   Float_t x, y, z;
163   Int_t alphaIdx = 0;
164   GetPoint(alphaIdx, x, y, z);
165   Float_t minR2=x*x + y*y;
166   for (Int_t i=1; i<=fLastPoint; i++)
167   {
168     GetPoint(i, x, y, z);
169     Float_t cR2 = x*x + y*y;
170     if ( minR2 > cR2 )
171     {
172       minR2 = cR2;
173       alphaIdx = i;
174     }
175   }
176   GetPoint(alphaIdx, x, y, z);
177   fAlpha = ATan2(y, x);
178   Float_t sin = Sin(-fAlpha);
179   Float_t cos = Cos(-fAlpha);
180   for (Int_t i=0; i<=fLastPoint; i++) {
181     GetPoint(i, x, y, z);
182     fRieman->AddPoint(cos*x - sin*y, cos*y + sin*x, z, 1, 1);
183   }
184   fRieman->Update();
185
186   Double_t r = Sqrt(minR2);
187   Double_t param[5];
188   Double_t cov[15];
189   fRieman->GetExternalParameters(r, param, cov);
190   // curvature to pt
191   param[4] /= TEveTrackPropagator::fgDefMagField*TEveTrackPropagator::fgkB2C;
192   // sign in tang
193   if(param[4] < 0) param[3] *= -1;
194   AliExternalTrackParam trackParam(r, fAlpha, param, cov);
195   trackParam.Print();
196
197   // make track
198   Double_t AliEveV0[3];
199   trackParam.GetXYZAt(r, TEveTrackPropagator::fgDefMagField, AliEveV0);
200   Double_t P0[3];
201   trackParam.GetPxPyPzAt(r, TEveTrackPropagator::fgDefMagField, P0);
202   TEveRecTrack rc;
203   rc.fV.Set(AliEveV0);
204   rc.fP.Set(P0);
205   rc.fSign = trackParam.Charge();
206
207   TEveTrack* track = new TEveTrack(&rc, fTrackList->GetPropagator());
208   track->SetName(Form("track %f", fAlpha));
209   TEvePathMark* pm = new TEvePathMark(TEvePathMark::kDaughter);
210   for(Int_t i=0; i==fLastPoint; i++)
211   {
212     GetPoint(i, x, y, z);
213     pm->fV.Set(x, y, z);
214     pm->fP.Set(P0);
215     track->AddPathMark(pm);
216   }
217   track->MakeTrack();
218   track->SetAttLineAttMarker(fTrackList);
219   gEve->AddElement(track, fTrackList);
220 }
221
222
223 /******************************************************************************/
224 void AliEveTrackFitter::DestroyElements()
225 {
226   // Virtual method of base class TEveElement.
227   // Preserves TEveTrackPropagator object for fitted helices.
228
229   TEveElement::DestroyElements();
230
231   gEve->AddElement(fTrackList, this);
232   fTrackList->DestroyElements();
233   UpdateItems();
234 }
235
236 /******************************************************************************/
237 void AliEveTrackFitter::DrawDebugGraph()
238 {
239   // Draw graph of picked points and helix points.
240
241   static const TEveException eH("AliEveTrackFitter::DrawRiemanGraph ");
242
243   if(fRieman == 0)
244     throw(eH + "fitter not set.");
245
246   Int_t nR = fRieman->GetN();
247   fGraphPicked->Set(nR);
248   fGraphHelix->Set(nR);
249
250   Double_t* x  =  fRieman->GetX();
251   Double_t* y  =  fRieman->GetY();
252   Double_t* sy =  fRieman->GetSy();
253   for (Int_t i=0; i<nR; i++)
254   {
255     fGraphPicked->SetPoint(i, x[i], y[i]);
256     fGraphHelix->SetPoint (i, x[i], fRieman->GetYat(x[i]));
257     fGraphHelix->SetPointError(i, 0.1, sy[i]); // now faked
258   }
259
260   if (gPad)
261     gPad->Clear();
262
263   fGraphPicked->Draw("AP");
264   fGraphHelix->Draw("SAME P");
265   gPad->GetCanvas()->SetTitle(Form("AliRieman alpha: %f", fAlpha));
266   gPad->Modified();
267   gPad->Update();
268 }