9249a0b047263f6c984bbf575f006783cb6abc83
[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 helix fit. It creates a set of
30 // points, listening to signal PointCtrlClicked() of any
31 // TEvePointSet. Via editor it fits selected points and creates a
32 // reconstructed track.
33 //
34
35 ClassImp(AliEveTrackFitter)
36
37 AliEveTrackFitter::AliEveTrackFitter(const Text_t* name, Int_t n_points) :
38     TEvePointSet   (name, n_points),
39
40     fGraphSelected (0),
41     fGraphFitted   (0),
42     fAlpha         (0),
43     fRieman        (0),
44     fConnected     (kFALSE),
45     fTrackList     (0),
46     fMapPS         ()
47 {
48   // Constructor.
49
50   SetMarkerColor(3);
51   SetOwnIds(kFALSE);
52
53   fGraphSelected = new TGraph();
54   fGraphSelected->SetName("Selected points");
55   fGraphSelected->SetMarkerColor(4);
56   fGraphSelected->SetMarkerStyle(4);
57   fGraphSelected->SetMarkerSize(2);
58
59   fGraphFitted = new TGraphErrors();
60   fGraphFitted->SetName("Fitted points");
61   fGraphFitted->SetMarkerColor(2);
62
63   fTrackList = new TEveTrackList("Tracks");
64   fTrackList->SetLineWidth(2);
65   fTrackList->SetLineColor(8);
66   fTrackList->IncDenyDestroy();
67   fTrackList->GetPropagator()->SetEditPathMarks(kTRUE);
68   gEve->AddElement(fTrackList, this);
69   UpdateItems();
70 }
71
72 AliEveTrackFitter::~AliEveTrackFitter()
73 {
74   // Destructor.
75
76   if(fRieman) delete fRieman;
77   fTrackList->DecDenyDestroy();
78   delete fTrackList;
79 }
80
81 /******************************************************************************/
82 void AliEveTrackFitter::DestroyElements()
83 {
84   // Virtual method of base class TEveElement.
85   // It preserves track list to have coomon track propagator attributes.
86
87   TEveElement::DestroyElements();
88   gEve->AddElement(fTrackList, this);
89   fTrackList->DestroyElements();
90   UpdateItems();
91 }
92
93 /******************************************************************************/
94 void AliEveTrackFitter::Start()
95 {
96   // Start selection of points.
97
98   Reset();
99   if(fConnected == kFALSE)
100   {
101     TQObject::Connect("TEvePointSet", "PointCtrlClicked(TEvePointSet*,Int_t)",
102                       "AliEveTrackFitter", this, "AddFitPoint(TEvePointSet*,Int_t)");
103
104     fConnected = kTRUE;
105   }
106 }
107
108 void AliEveTrackFitter::Stop()
109 {
110   // Stop selection of points.
111
112   if(fConnected)
113   {
114     TQObject::Disconnect("TEvePointSet", "AddFitPoint(TEvePointSet*,Int_t)");
115     fConnected = kFALSE;
116   }
117 }
118
119 /******************************************************************************/
120
121 void AliEveTrackFitter::AddFitPoint(TEvePointSet* ps, Int_t n)
122 {
123   // Add/remove given point depending if exists in the fMapPS.
124
125   Float_t x, y, z;
126
127   std::map<Point_t, Int_t>::iterator g = fMapPS.find(Point_t(ps, n));
128   if(g != fMapPS.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     fMapPS.erase(g);
137     fLastPoint--;
138   }
139   else
140   {
141     fMapPS[Point_t(ps, n)] = Size();
142     ps->GetPoint(n, x, y, z);
143     SetNextPoint(x, y, z);
144     SetPointId(ps->GetPointId(n));
145   }
146   ResetBBox();
147   ElementChanged(kTRUE, kTRUE);
148 }
149
150 /******************************************************************************/
151
152 void AliEveTrackFitter::FitTrack()
153 {
154   // Fit selected points with AliRieman fitter.
155
156   using namespace TMath;
157
158   if(fRieman) delete fRieman;
159   fRieman = new AliRieman(Size());
160
161   Float_t x, y, z;
162   Int_t alphaIdx = 0;
163   GetPoint(alphaIdx, x, y, z);
164   Float_t minR2=x*x + y*y;
165   for (Int_t i=1; i<=fLastPoint; i++)
166   {
167     GetPoint(i, x, y, z);
168     Float_t cR2 = x*x + y*y;
169     if ( minR2 > cR2 )
170     {
171       minR2 = cR2;
172       alphaIdx = i;
173     }
174   }
175   GetPoint(alphaIdx, x, y, z);
176   fAlpha = ATan2(y, x);
177   Float_t sin = Sin(-fAlpha);
178   Float_t cos = Cos(-fAlpha);
179   for (Int_t i=0; i<=fLastPoint; i++) {
180     GetPoint(i, x, y, z);
181     fRieman->AddPoint(cos*x - sin*y, cos*y + sin*x, z, 1, 1);
182   }
183   fRieman->Update();
184
185   Double_t r = Sqrt(minR2);
186   Double_t param[5];
187   Double_t cov[15];
188   fRieman->GetExternalParameters(r, param, cov);
189   // curvature to pt
190   param[4] /= TEveTrackPropagator::fgDefMagField*TEveTrackPropagator::fgkB2C;
191   // sign in tang
192   if(param[4] < 0) param[3] *= -1;
193   AliExternalTrackParam trackParam(r, fAlpha, param, cov);
194   trackParam.Print();
195
196   // make track
197   Double_t AliEveV0[3];
198   trackParam.GetXYZAt(r, TEveTrackPropagator::fgDefMagField, AliEveV0);
199   Double_t P0[3];
200   trackParam.GetPxPyPzAt(r, TEveTrackPropagator::fgDefMagField, P0);
201   TEveRecTrack rc;
202   rc.fV.Set(AliEveV0);
203   rc.fP.Set(P0);
204   rc.fSign = trackParam.Charge();
205
206   TEveTrack* track = new TEveTrack(&rc, fTrackList->GetPropagator());
207   track->SetName(Form("track %f", fAlpha));
208   TEvePathMark* pm = new TEvePathMark(TEvePathMark::kDaughter);
209   for(Int_t i=0; i==fLastPoint; i++)
210   {
211     GetPoint(i, x, y, z);
212     pm->fV.Set(x, y, z);
213     pm->fP.Set(P0);
214     track->AddPathMark(pm);
215   }
216   track->MakeTrack();
217   track->SetAttLineAttMarker(fTrackList);
218   gEve->AddElement(track, fTrackList);
219 }
220
221
222 void AliEveTrackFitter::Reset(Int_t n, Int_t ids)
223 {
224   // Reset selection.
225
226   if(fRieman) fRieman->Reset();
227   TEvePointSet::Reset(n, ids);
228   fMapPS.clear();
229 }
230
231 /******************************************************************************/
232 void AliEveTrackFitter::DrawRiemanGraph()
233 {
234   // Draw graph of rieman fit.
235
236    static const TEveException eH("AliEveTrackFitter::DrawRiemanGraph ");
237
238   if(fRieman == 0)
239     throw(eH + "fitter not set.");
240
241   Int_t nR = fRieman->GetN();
242   fGraphSelected->Set(nR);
243   fGraphFitted->Set(nR);
244
245   Double_t* x =  fRieman->GetX();
246   Double_t* y =  fRieman->GetY();
247   Double_t* sy =  fRieman->GetSy();
248   for (Int_t i=0; i<nR; i++)
249   {
250     fGraphSelected->SetPoint(i, x[i], y[i]);
251     fGraphFitted->SetPoint(i, x[i], fRieman->GetYat(x[i]));
252     fGraphFitted->SetPointError(i, 0.1, sy[i]);
253   }
254
255   if (gPad) gPad->Clear();
256   fGraphSelected->Draw("AP");
257   fGraphFitted->Draw("SAME P");
258   gPad->GetCanvas()->SetTitle(Form("AliRieman alpha: %f", fAlpha));
259   gPad->Modified();
260   gPad->Update();
261 }