Fork me on GitHub

source: git/external/TrackCovariance/VertexFit.cc@ 78ce8d1

Last change on this file since 78ce8d1 was b750b0a, checked in by Franco BEDESCHI <bed@…>, 3 years ago

Fix CovarianceMatrix scaling error. Vertexing improvements.
  • Property mode set to 100644
File size: 12.2 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 fRold = -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 fRold = -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 fRold = -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 //
225 // Track loop
226 for (Int_t i = 0; i < fNtr; i++)
227 {
228 Double_t s = 0.;
229 TVectorD par = *fPar[i];
230 TMatrixDSym Cov = *fCov[i];
231 x0i.push_back(new TVectorD(Fill_x0(par)));
232 ni.push_back(new TVectorD(derXds(par, s)));
233 TMatrixD A = derXdPar(par, s);
234 Ci.push_back(new TMatrixDSym(Cov.Similarity(A)));
235 TMatrixDSym Cinv = RegInv(*Ci[i]);
236 wi.push_back(new TVectorD(Cinv * (*ni[i])));
237 }
238 //std::cout << "Vtx init completed. fNtr = "<<fNtr << std::endl;
239 //
240 // Get fit vertex
241 //
242 TMatrixDSym D(3); D.Zero();
243 TVectorD Dx(3); Dx.Zero();
244 for (Int_t i = 0; i < fNtr; i++)
245 {
246 TMatrixDSym Cinv = RegInv(*Ci[i]);
247 TMatrixDSym W(3);
248 W.Rank1Update(*wi[i], 1. / Ci[i]->Similarity(*wi[i]));
249 TMatrixDSym Dd = Cinv - W;
250 D += Dd;
251 Dx += Dd * (*x0i[i]);
252 }
253 if(fVtxCst){
254 D += fCovCstInv;
255 Dx += fCovCstInv*fxCst;
256 }
257 fXv = RegInv(D) * Dx;
258 //std::cout << "Fast vertex (x, y, z) = "<<fXv(0)<<", "<<fXv(1)<<", "<<fXv(2) << std::endl;
259 //
260 // Get fit phases
261 //
262 for (Int_t i = 0; i < fNtr; i++){
263 Double_t si = Dot(*wi[i], fXv - (*x0i[i])) / Ci[i]->Similarity(*wi[i]);
264 ffi.push_back(si);
265 //TVectorD xvi = Fill_x(*fPar[i],si);
266 //std::cout << "Fast vertex "<<i<<": xvi = "<<xvi(0)<<", "<<xvi(1)<<", "<<xvi(2)
267 // <<", si = "<<si<< std::endl;
268 }
269 //
270 // Cleanup
271 for (Int_t i = 0; i < fNtr; i++)
272 {
273 x0i[i]->Clear(); delete x0i[i];
274 ni[i]->Clear(); delete ni[i];
275 Ci[i]->Clear(); delete Ci[i];
276 wi[i]->Clear(); delete wi[i];
277 }
278 x0i.clear();;
279 ni.clear();
280 Ci.clear();
281 wi.clear();
282}
283//
284void VertexFit::VertexFitter()
285{
286 //std::cout << "VertexFitter: just in" << std::endl;
287 if (fNtr < 2 && !fVtxCst){
288 std::cout << "VertexFit::VertexFitter - Method called with less than 2 tracks - Aborting " << std::endl;
289 std::exit(1);
290 }
291 //
292 // Vertex fit
293 //
294 // Initial variable definitions
295 TVectorD x(3);
296 TMatrixDSym covX(3);
297 Double_t Chi2 = 0;
298 //
299 VtxFitNoSteer(); // Fast vertex finder on first pass (set ffi and fXv)
300 TVectorD x0 = fXv;
301 //
302 // Iteration properties
303 //
304 Int_t Ntry = 0;
305 Int_t TryMax = 100;
306 Double_t eps = 1.0e-9; // vertex stability
307 Double_t epsi = 1000.;
308 //
309 // Iteration loop
310 while (epsi > eps && Ntry < TryMax) // Iterate until found vertex is stable
311 {
312 // Initialize arrays
313 x.Zero();
314 TVectorD cterm(3); TMatrixDSym H(3); TMatrixDSym DW1D(3);
315 covX.Zero(); // Reset vertex covariance
316 cterm.Zero(); // Reset constant term
317 H.Zero(); // Reset H matrix
318 DW1D.Zero();
319 //
320 // Reset work arrays
321 //
322 ResetWrkArrays();
323 //
324 // Start loop on tracks
325 //
326 for (Int_t i = 0; i < fNtr; i++)
327 {
328 // Get track helix parameters and their covariance matrix
329 TVectorD par = *fPar[i];
330 TMatrixDSym Cov = *fCov[i];
331 //
332 // Update track related arrays
333 //
334 UpdateTrkArrays(i);
335 TMatrixDSym Ds = *fDi[i];
336 TMatrixDSym Winv = *fWinvi[i];
337 TMatrixDSym DsW1Ds = Winv.Similarity(Ds); // Service matrix to calculate covX
338 //
339 // Update global arrays
340 DW1D += DsW1Ds;
341 // Update hessian
342 H += Ds;
343 // update constant term
344 TVectorD xs = *fx0i[i] - *fdi[i];
345 //TVectorD xx0 = *fx0i[i];
346
347 //std::cout << "Iter. " << Ntry << ", trk " << i << ", xs= "
348 // << xs(0) << ", " << xs(1) << ", " << xs(2)<<
349 // ", ph0= "<<par(1)<< std::endl;
350
351 cterm += Ds * xs;
352 } // End loop on tracks
353 // Some additions in case of external constraints
354 if (fVtxCst) {
355 H += fCovCstInv;
356 cterm += fCovCstInv * fxCst;
357 DW1D += fCovCstInv;
358 }
359 //
360 // update vertex position
361 TMatrixDSym H1 = RegInv(H);
362 x = H1 * cterm;
363 //
364 // Update vertex covariance
365 covX = DW1D.Similarity(H1);
366 //
367 // Update phases and chi^2
368 Chi2 = 0.0;
369 for (Int_t i = 0; i < fNtr; i++)
370 {
371 TVectorD lambda = (*fDi[i]) * (*fx0i[i] - x - *fdi[i]);
372 TMatrixDSym Wm1 = *fWinvi[i];
373 fChi2List(i) = Wm1.Similarity(lambda);
374 if(fChi2List(i) < 0.0){
375 //std::cout<<"# "<<i<<", Chi2= "<<fChi2List(i)<<", Wm1:"<<std::endl; Wm1.Print();
376 //std::cout<<"Lambda= "<<std::endl; lambda.Print();
377 }
378 Chi2 += fChi2List(i);
379 TVectorD a = *fai[i];
380 TVectorD b = (*fWi[i]) * (x - *fx0i[i] + *fdi[i]);
381 ffi[i] += Dot(a, b) / fa2i[i];
382 TVectorD newPar = *fPar[i] - ((*fCov[i]) * (*fAti[i])) * lambda;
383 fParNew[i] = new TVectorD(newPar);
384 }
385 // Add external constraint to Chi2
386 if (fVtxCst) Chi2 += fCovCstInv.Similarity(x - fxCst);
387 //
388 TVectorD dx = x - x0;
389 x0 = x;
390 // update vertex stability
391 TMatrixDSym Hess = RegInv(covX);
392 epsi = Hess.Similarity(dx);
393 Ntry++;
394 //
395 // Store result
396 //
397 fXv = x; // Vertex position
398 fcovXv = covX; // Vertex covariance
399 fChi2 = Chi2; // Vertex fit Chi2
400 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2) << std::endl;
401 } // end of iteration loop
402 //
403 fVtxDone = kTRUE; // Set fit completion flag
404 fRold = TMath::Sqrt(fXv(0)*fXv(0) + fXv(1)*fXv(1)); // Store fit
405 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2)
406 // << ", after "<<Ntry<<" iterations"<<std::endl;
407 //
408}
409//
410// Return fit vertex
411TVectorD VertexFit::GetVtx()
412{
413 if (!fVtxDone) VertexFitter();
414 return fXv;
415}
416//
417// Return fit vertex covariance
418TMatrixDSym VertexFit::GetVtxCov()
419{
420 if (!fVtxDone) VertexFitter();
421 return fcovXv;
422}
423//
424// Return fit vertex chi2
425Double_t VertexFit::GetVtxChi2()
426{
427 if (!fVtxDone) VertexFitter();
428 return fChi2;
429}
430//
431// Return array of chi2 contributions from each track
432TVectorD VertexFit::GetVtxChi2List()
433{
434 if (!fVtxDone) VertexFitter();
435 return fChi2List;
436}
437//
438// Handle tracks/constraints
439void VertexFit::AddVtxConstraint(TVectorD xv, TMatrixDSym cov) // Add gaussian vertex constraint
440{
441 //std::cout << "VertexFit::AddVtxConstraint: Not implemented yet" << std::endl;
442 fVtxCst = kTRUE; // Vertex constraint flag
443 fxCst = xv; // Constraint value
444 fCovCst = cov; // Constraint covariance
445 fCovCstInv = cov;
446 fCovCstInv.Invert(); // Its inverse
447 //
448 // Set starting vertex as external constraint
449 fXv = fxCst;
450 fcovXv = fCovCst;
451}
452//
453// Adding tracks one by one
454void VertexFit::AddTrk(TVectorD *par, TMatrixDSym *Cov) // Add track to input list
455{
456 fNtr++;
457 fChi2List.ResizeTo(fNtr); // Resize chi2 array
458 fPar.push_back(par); // add new track
459 fCov.push_back(Cov);
460 fParNew.push_back(par); // add new track
461 fCovNew.push_back(Cov);
462 //
463 // Reset previous vertex temp arrays
464 ResetWrkArrays();
465 ffi.clear();
466 fVtxDone = kFALSE; // Reset vertex done flag
467}
468//
469// Removing tracks one by one
470void VertexFit::RemoveTrk(Int_t iTrk) // Remove iTrk track
471{
472 fNtr--;
473 fChi2List.Clear();
474 fChi2List.ResizeTo(fNtr); // Resize chi2 array
475 fPar.erase(fPar.begin() + iTrk); // Remove track
476 fCov.erase(fCov.begin() + iTrk);
477 fParNew.erase(fParNew.begin() + iTrk); // Remove track
478 fCovNew.erase(fCovNew.begin() + iTrk);
479 //
480 // Reset previous vertex temp arrays
481 ResetWrkArrays();
482 ffi.clear();
483 fVtxDone = kFALSE; // Reset vertex done flag
484}
Note: See TracBrowser for help on using the repository browser.