Fork me on GitHub

source: git/external/TrackCovariance/VertexFit.cc@ cc8716b

Last change on this file since cc8716b was 7bca620, checked in by Franco BEDESCHI <bed@…>, 3 years ago

Add feature to force starting radius in vertex fit
  • Property mode set to 100644
File size: 12.5 KB
Line 
1/*
2Vertex fitting code
3*/
4#include <TMath.h>
5#include <TVectorD.h>
6#include <TVector3.h>
7#include <TMatrixD.h>
8#include <TMatrixDSym.h>
9#include <TMatrixDSymEigen.h>
10#include "VertexFit.h"
11//
12// Constructors
13//
14//
15// Empty construction (to be used when adding tracks later with AddTrk() )
16VertexFit::VertexFit()
17{
18 fNtr = 0;
19 fRstart = -1.0;
20 fVtxDone = kFALSE;
21 fVtxCst = kFALSE;
22 fxCst.ResizeTo(3);
23 fCovCst.ResizeTo(3, 3);
24 fCovCstInv.ResizeTo(3, 3);
25 fXv.ResizeTo(3);
26 fcovXv.ResizeTo(3, 3);
27}
28//
29// Build from list of parameters and covariances
30VertexFit::VertexFit(Int_t Ntr, TVectorD** trkPar, TMatrixDSym** trkCov)
31{
32 fNtr = Ntr;
33 fRstart = -1.0;
34 fVtxDone = kFALSE;
35 fVtxCst = kFALSE;
36 fxCst.ResizeTo(3);
37 fCovCst.ResizeTo(3, 3);
38 fCovCstInv.ResizeTo(3, 3);
39 fXv.ResizeTo(3);
40 fcovXv.ResizeTo(3, 3);
41 //
42 for (Int_t i = 0; i < fNtr; i++)
43 {
44 fPar.push_back(new TVectorD(*trkPar[i]));
45 fParNew.push_back(new TVectorD(*trkPar[i]));
46 fCov.push_back(new TMatrixDSym(*trkCov[i]));
47 fCovNew.push_back(new TMatrixDSym(*trkCov[i]));
48 }
49 fChi2List.ResizeTo(fNtr);
50 //
51}
52//
53// Build from ObsTrk list of tracks
54VertexFit::VertexFit(Int_t Ntr, ObsTrk** track)
55{
56 fNtr = Ntr;
57 fRstart = -1.0;
58 fVtxDone = kFALSE;
59 fVtxCst = kFALSE;
60 fxCst.ResizeTo(3);
61 fCovCst.ResizeTo(3, 3);
62 fCovCstInv.ResizeTo(3, 3);
63 fXv.ResizeTo(3);
64 fcovXv.ResizeTo(3, 3);
65 //
66 fChi2List.ResizeTo(fNtr);
67 for (Int_t i = 0; i < fNtr; i++)
68 {
69 fPar.push_back(new TVectorD(track[i]->GetObsPar()));
70 fParNew.push_back(new TVectorD(track[i]->GetObsPar()));
71 fCov.push_back(new TMatrixDSym(track[i]->GetCov()));
72 fCovNew.push_back(new TMatrixDSym(track[i]->GetCov()));
73 }
74}
75//
76// Destructor
77//
78void VertexFit::ResetWrkArrays()
79{
80 Int_t N = (Int_t)fdi.size();
81 if(N > 0){
82 for (Int_t i = 0; i < N; i++)
83 {
84 if (fx0i[i]) { fx0i[i]->Clear(); delete fx0i[i]; }
85 if (fai[i]) { fai[i]->Clear(); delete fai[i]; }
86 if (fdi[i]) { fdi[i]->Clear(); delete fdi[i]; }
87 if (fAti[i]) { fAti[i]->Clear(); delete fAti[i]; }
88 if (fDi[i]) { fDi[i]->Clear(); delete fDi[i]; }
89 if (fWi[i]) { fWi[i]->Clear(); delete fWi[i]; }
90 if (fWinvi[i]){ fWinvi[i]->Clear(); delete fWinvi[i]; }
91 }
92 fa2i.clear();
93 fx0i.clear();
94 fai.clear();
95 fdi.clear();
96 fAti.clear();
97 fDi.clear();
98 fWi.clear();
99 fWinvi.clear();
100 }
101}
102VertexFit::~VertexFit()
103{
104 fxCst.Clear();
105 fCovCst.Clear();
106 fCovCstInv.Clear();
107 fXv.Clear();
108 fcovXv.Clear();
109 fChi2List.Clear();
110 //
111 for (Int_t i = 0; i < fNtr; i++)
112 {
113 fPar[i]->Clear(); delete fPar[i];
114 fParNew[i]->Clear(); delete fParNew[i];
115 fCov[i]->Clear(); delete fCov[i];
116 fCovNew[i]->Clear(); delete fCovNew[i];
117 }
118 fPar.clear();
119 fParNew.clear();
120 fCov.clear();
121 fCovNew.clear();
122 //
123 ResetWrkArrays();
124 ffi.clear();
125 fNtr = 0;
126}
127//
128//
129//
130TVectorD VertexFit::Fill_x0(TVectorD par)
131{
132 //
133 // Calculate track 3D position at R = |D| (minimum approach to z-axis)
134 //
135 TVectorD x0(3);
136 //
137 // Decode input arrays
138 //
139 Double_t D = par(0);
140 Double_t p0 = par(1);
141 Double_t C = par(2);
142 Double_t z0 = par(3);
143 Double_t ct = par(4);
144 //
145 x0(0) = -D * TMath::Sin(p0);
146 x0(1) = D * TMath::Cos(p0);
147 x0(2) = z0;
148 //
149 return x0;
150}
151//
152TVectorD VertexFit::Fill_x(TVectorD par, Double_t phi)
153{
154 //
155 // Calculate track 3D position for a given phase, phi
156 //
157 TVectorD x(3);
158 //
159 // Decode input arrays
160 //
161 Double_t D = par(0);
162 Double_t p0 = par(1);
163 Double_t C = par(2);
164 Double_t z0 = par(3);
165 Double_t ct = par(4);
166 //
167 TVectorD x0 = Fill_x0(par);
168 x(0) = x0(0) + (TMath::Sin(phi + p0) - TMath::Sin(p0)) / (2 * C);
169 x(1) = x0(1) - (TMath::Cos(phi + p0) - TMath::Cos(p0)) / (2 * C);
170 x(2) = x0(2) + ct * phi / (2 * C);
171 //
172 return x;
173}
174//
175void VertexFit::UpdateTrkArrays(Int_t i)
176{
177 //
178 // Get track parameters, covariance and phase
179 Double_t fs = ffi[i]; // Get phase
180 TVectorD par = *fParNew[i];
181 TMatrixDSym Cov = *fCov[i];
182 //
183 // Fill all track related work arrays arrays
184 TMatrixD A = derXdPar(par, fs); // A = dx/da = derivatives wrt track parameters
185 TMatrixDSym Winv = Cov;
186 Winv.Similarity(A); // W^-1 = A*C*A'
187
188 TMatrixD At(TMatrixD::kTransposed, A); // A transposed
189 fAti.push_back(new TMatrixD(At)); // Store A'
190 fWinvi.push_back(new TMatrixDSym(Winv)); // Store W^-1 matrix
191 //
192 TVectorD xs = Fill_x(par, fs);
193 fx0i.push_back(new TVectorD(xs)); // Start helix position
194 //
195 TVectorD di = A * (par - *fPar[i]); // x-shift
196 fdi.push_back(new TVectorD(di)); // Store x-shift
197 //
198 TMatrixDSym W = RegInv(Winv); // W = (A*C*A')^-1
199 fWi.push_back(new TMatrixDSym(W)); // Store W matrix
200 //
201 TVectorD a = derXds(par, fs); // a = dx/ds = derivatives wrt phase
202 fai.push_back(new TVectorD(a)); // Store a
203 //
204 Double_t a2 = W.Similarity(a);
205 fa2i.push_back(a2); // Store a2
206 //
207 // Build D matrix
208 TMatrixDSym B(3);
209 B.Rank1Update(a, -1. / a2);
210 B.Similarity(W);
211 TMatrixDSym Ds = W + B; // D matrix
212 fDi.push_back(new TMatrixDSym(Ds)); // Store D matrix
213}
214//
215void VertexFit::VtxFitNoSteer()
216{
217 //
218 // Initialize
219 //
220 std::vector<TVectorD*> x0i; // Tracks at ma
221 std::vector<TVectorD*> ni; // Track derivative wrt phase
222 std::vector<TMatrixDSym*> Ci; // Position error matrix at fixed phase
223 std::vector<TVectorD*> wi; // Ci*ni
224 std::vector<Double_t> s_in; // Starting phase
225 //
226 //
227 //
228 // Track loop
229 for (Int_t i = 0; i < fNtr; i++)
230 {
231 TVectorD par = *fPar[i];
232 TMatrixDSym Cov = *fCov[i];
233 Double_t s = 0.;
234 // Case when starting radius is provided
235 if(fRstart > TMath::Abs(par(0))){
236 s = 2.*TMath::ASin(par(2)*TMath::Sqrt((fRstart*fRstart-par(0)*par(0))/(1.+2.*par(2)*par(0))));
237 }
238 //
239 x0i.push_back(new TVectorD(Fill_x(par, s)));
240 ni.push_back(new TVectorD(derXds(par, s)));
241 TMatrixD A = derXdPar(par, s);
242 Ci.push_back(new TMatrixDSym(Cov.Similarity(A)));
243 TMatrixDSym Cinv = RegInv(*Ci[i]);
244 wi.push_back(new TVectorD(Cinv * (*ni[i])));
245 s_in.push_back(s);
246 }
247 //std::cout << "Vtx init completed. fNtr = "<<fNtr << std::endl;
248 //
249 // Get fit vertex
250 //
251 TMatrixDSym D(3); D.Zero();
252 TVectorD Dx(3); Dx.Zero();
253 for (Int_t i = 0; i < fNtr; i++)
254 {
255 TMatrixDSym Cinv = RegInv(*Ci[i]);
256 TMatrixDSym W(3);
257 W.Rank1Update(*wi[i], 1. / Ci[i]->Similarity(*wi[i]));
258 TMatrixDSym Dd = Cinv - W;
259 D += Dd;
260 Dx += Dd * (*x0i[i]);
261 }
262 if(fVtxCst){
263 D += fCovCstInv;
264 Dx += fCovCstInv*fxCst;
265 }
266 fXv = RegInv(D) * Dx;
267 //std::cout << "Fast vertex (x, y, z) = "<<fXv(0)<<", "<<fXv(1)<<", "<<fXv(2) << std::endl;
268 //
269 // Get fit phases
270 //
271 for (Int_t i = 0; i < fNtr; i++){
272 Double_t si = Dot(*wi[i], fXv - (*x0i[i])) / Ci[i]->Similarity(*wi[i]);
273 ffi.push_back(si+s_in[i]);
274 //TVectorD xvi = Fill_x(*fPar[i],si);
275 //std::cout << "Fast vertex "<<i<<": xvi = "<<xvi(0)<<", "<<xvi(1)<<", "<<xvi(2)
276 // <<", si = "<<si<< std::endl;
277 }
278 //
279 // Cleanup
280 for (Int_t i = 0; i < fNtr; i++)
281 {
282 x0i[i]->Clear(); delete x0i[i];
283 ni[i]->Clear(); delete ni[i];
284 Ci[i]->Clear(); delete Ci[i];
285 wi[i]->Clear(); delete wi[i];
286 }
287 x0i.clear();;
288 ni.clear();
289 Ci.clear();
290 wi.clear();
291 s_in.clear();
292}
293//
294void VertexFit::VertexFitter()
295{
296 //std::cout << "VertexFitter: just in" << std::endl;
297 if (fNtr < 2 && !fVtxCst){
298 std::cout << "VertexFit::VertexFitter - Method called with less than 2 tracks - Aborting " << std::endl;
299 std::exit(1);
300 }
301 //
302 // Vertex fit
303 //
304 // Initial variable definitions
305 TVectorD x(3);
306 TMatrixDSym covX(3);
307 Double_t Chi2 = 0;
308 //
309 VtxFitNoSteer(); // Fast vertex finder on first pass (set ffi and fXv)
310 TVectorD x0 = fXv;
311 //
312 // Iteration properties
313 //
314 Int_t Ntry = 0;
315 Int_t TryMax = 100;
316 Double_t eps = 1.0e-9; // vertex stability
317 Double_t epsi = 1000.;
318 //
319 // Iteration loop
320 while (epsi > eps && Ntry < TryMax) // Iterate until found vertex is stable
321 {
322 // Initialize arrays
323 x.Zero();
324 TVectorD cterm(3); TMatrixDSym H(3); TMatrixDSym DW1D(3);
325 covX.Zero(); // Reset vertex covariance
326 cterm.Zero(); // Reset constant term
327 H.Zero(); // Reset H matrix
328 DW1D.Zero();
329 //
330 // Reset work arrays
331 //
332 ResetWrkArrays();
333 //
334 // Start loop on tracks
335 //
336 for (Int_t i = 0; i < fNtr; i++)
337 {
338 // Get track helix parameters and their covariance matrix
339 TVectorD par = *fPar[i];
340 TMatrixDSym Cov = *fCov[i];
341 //
342 // Update track related arrays
343 //
344 UpdateTrkArrays(i);
345 TMatrixDSym Ds = *fDi[i];
346 TMatrixDSym Winv = *fWinvi[i];
347 TMatrixDSym DsW1Ds = Winv.Similarity(Ds); // Service matrix to calculate covX
348 //
349 // Update global arrays
350 DW1D += DsW1Ds;
351 // Update hessian
352 H += Ds;
353 // update constant term
354 TVectorD xs = *fx0i[i] - *fdi[i];
355 //TVectorD xx0 = *fx0i[i];
356
357 //std::cout << "Iter. " << Ntry << ", trk " << i << ", xs= "
358 // << xs(0) << ", " << xs(1) << ", " << xs(2)<<
359 // ", ph0= "<<par(1)<< std::endl;
360
361 cterm += Ds * xs;
362 } // End loop on tracks
363 // Some additions in case of external constraints
364 if (fVtxCst) {
365 H += fCovCstInv;
366 cterm += fCovCstInv * fxCst;
367 DW1D += fCovCstInv;
368 }
369 //
370 // update vertex position
371 TMatrixDSym H1 = RegInv(H);
372 x = H1 * cterm;
373 //
374 // Update vertex covariance
375 covX = DW1D.Similarity(H1);
376 //
377 // Update phases and chi^2
378 Chi2 = 0.0;
379 for (Int_t i = 0; i < fNtr; i++)
380 {
381 TVectorD lambda = (*fDi[i]) * (*fx0i[i] - x - *fdi[i]);
382 TMatrixDSym Wm1 = *fWinvi[i];
383 fChi2List(i) = Wm1.Similarity(lambda);
384 if(fChi2List(i) < 0.0){
385 //std::cout<<"# "<<i<<", Chi2= "<<fChi2List(i)<<", Wm1:"<<std::endl; Wm1.Print();
386 //std::cout<<"Lambda= "<<std::endl; lambda.Print();
387 }
388 Chi2 += fChi2List(i);
389 TVectorD a = *fai[i];
390 TVectorD b = (*fWi[i]) * (x - *fx0i[i] + *fdi[i]);
391 ffi[i] += Dot(a, b) / fa2i[i];
392 TVectorD newPar = *fPar[i] - ((*fCov[i]) * (*fAti[i])) * lambda;
393 fParNew[i] = new TVectorD(newPar);
394 }
395 // Add external constraint to Chi2
396 if (fVtxCst) Chi2 += fCovCstInv.Similarity(x - fxCst);
397 //
398 TVectorD dx = x - x0;
399 x0 = x;
400 // update vertex stability
401 TMatrixDSym Hess = RegInv(covX);
402 epsi = Hess.Similarity(dx);
403 Ntry++;
404 //
405 // Store result
406 //
407 fXv = x; // Vertex position
408 fcovXv = covX; // Vertex covariance
409 fChi2 = Chi2; // Vertex fit Chi2
410 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2) << std::endl;
411 } // end of iteration loop
412 //
413 fVtxDone = kTRUE; // Set fit completion flag
414 //fRstart = TMath::Sqrt(fXv(0)*fXv(0) + fXv(1)*fXv(1)); // Store fit
415 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2)
416 // << ", after "<<Ntry<<" iterations"<<std::endl;
417 //
418}
419//
420// Return fit vertex
421TVectorD VertexFit::GetVtx()
422{
423 if (!fVtxDone) VertexFitter();
424 return fXv;
425}
426//
427// Return fit vertex covariance
428TMatrixDSym VertexFit::GetVtxCov()
429{
430 if (!fVtxDone) VertexFitter();
431 return fcovXv;
432}
433//
434// Return fit vertex chi2
435Double_t VertexFit::GetVtxChi2()
436{
437 if (!fVtxDone) VertexFitter();
438 return fChi2;
439}
440//
441// Return array of chi2 contributions from each track
442TVectorD VertexFit::GetVtxChi2List()
443{
444 if (!fVtxDone) VertexFitter();
445 return fChi2List;
446}
447//
448// Handle tracks/constraints
449void VertexFit::AddVtxConstraint(TVectorD xv, TMatrixDSym cov) // Add gaussian vertex constraint
450{
451 //std::cout << "VertexFit::AddVtxConstraint: Not implemented yet" << std::endl;
452 fVtxCst = kTRUE; // Vertex constraint flag
453 fxCst = xv; // Constraint value
454 fCovCst = cov; // Constraint covariance
455 fCovCstInv = cov;
456 fCovCstInv.Invert(); // Its inverse
457 //
458 // Set starting vertex as external constraint
459 fXv = fxCst;
460 fcovXv = fCovCst;
461}
462//
463// Adding tracks one by one
464void VertexFit::AddTrk(TVectorD *par, TMatrixDSym *Cov) // Add track to input list
465{
466 fNtr++;
467 fChi2List.ResizeTo(fNtr); // Resize chi2 array
468 fPar.push_back(par); // add new track
469 fCov.push_back(Cov);
470 fParNew.push_back(par); // add new track
471 fCovNew.push_back(Cov);
472 //
473 // Reset previous vertex temp arrays
474 ResetWrkArrays();
475 ffi.clear();
476 fVtxDone = kFALSE; // Reset vertex done flag
477}
478//
479// Removing tracks one by one
480void VertexFit::RemoveTrk(Int_t iTrk) // Remove iTrk track
481{
482 fNtr--;
483 fChi2List.Clear();
484 fChi2List.ResizeTo(fNtr); // Resize chi2 array
485 fPar.erase(fPar.begin() + iTrk); // Remove track
486 fCov.erase(fCov.begin() + iTrk);
487 fParNew.erase(fParNew.begin() + iTrk); // Remove track
488 fCovNew.erase(fCovNew.begin() + iTrk);
489 //
490 // Reset previous vertex temp arrays
491 ResetWrkArrays();
492 ffi.clear();
493 fVtxDone = kFALSE; // Reset vertex done flag
494}
Note: See TracBrowser for help on using the repository browser.