Fork me on GitHub

source: git/external/TrackCovariance/VertexFit.cc@ 0983d48

Last change on this file since 0983d48 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
RevLine 
[b750b0a]1/*
2Vertex fitting code
3*/
4#include <TMath.h>
[537d24f]5#include <TVectorD.h>
6#include <TVector3.h>
7#include <TMatrixD.h>
8#include <TMatrixDSym.h>
[b750b0a]9#include <TMatrixDSymEigen.h>
[537d24f]10#include "VertexFit.h"
11//
12// Constructors
13//
[82db145]14//
15// Empty construction (to be used when adding tracks later with AddTrk() )
[537d24f]16VertexFit::VertexFit()
17{
18 fNtr = 0;
[82db145]19 fRold = -1.0;
[537d24f]20 fVtxDone = kFALSE;
21 fVtxCst = kFALSE;
22 fxCst.ResizeTo(3);
[91ef0b8]23 fCovCst.ResizeTo(3, 3);
[53f4746]24 fCovCstInv.ResizeTo(3, 3);
[537d24f]25 fXv.ResizeTo(3);
[91ef0b8]26 fcovXv.ResizeTo(3, 3);
[537d24f]27}
[82db145]28//
29// Build from list of parameters and covariances
[537d24f]30VertexFit::VertexFit(Int_t Ntr, TVectorD** trkPar, TMatrixDSym** trkCov)
31{
32 fNtr = Ntr;
[82db145]33 fRold = -1.0;
[537d24f]34 fVtxDone = kFALSE;
35 fVtxCst = kFALSE;
36 fxCst.ResizeTo(3);
[91ef0b8]37 fCovCst.ResizeTo(3, 3);
[53f4746]38 fCovCstInv.ResizeTo(3, 3);
[537d24f]39 fXv.ResizeTo(3);
[91ef0b8]40 fcovXv.ResizeTo(3, 3);
[537d24f]41 //
[82db145]42 for (Int_t i = 0; i < fNtr; i++)
43 {
[b750b0a]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]));
[82db145]48 }
49 fChi2List.ResizeTo(fNtr);
50 //
[537d24f]51}
[82db145]52//
53// Build from ObsTrk list of tracks
[537d24f]54VertexFit::VertexFit(Int_t Ntr, ObsTrk** track)
55{
56 fNtr = Ntr;
[82db145]57 fRold = -1.0;
[537d24f]58 fVtxDone = kFALSE;
59 fVtxCst = kFALSE;
60 fxCst.ResizeTo(3);
[91ef0b8]61 fCovCst.ResizeTo(3, 3);
[53f4746]62 fCovCstInv.ResizeTo(3, 3);
[537d24f]63 fXv.ResizeTo(3);
[91ef0b8]64 fcovXv.ResizeTo(3, 3);
[537d24f]65 //
[82db145]66 fChi2List.ResizeTo(fNtr);
67 for (Int_t i = 0; i < fNtr; i++)
[537d24f]68 {
[82db145]69 fPar.push_back(new TVectorD(track[i]->GetObsPar()));
[53f4746]70 fParNew.push_back(new TVectorD(track[i]->GetObsPar()));
[82db145]71 fCov.push_back(new TMatrixDSym(track[i]->GetCov()));
[53f4746]72 fCovNew.push_back(new TMatrixDSym(track[i]->GetCov()));
[537d24f]73 }
74}
75//
76// Destructor
[82db145]77//
78void VertexFit::ResetWrkArrays()
79{
[b750b0a]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();
[82db145]100 }
101}
[537d24f]102VertexFit::~VertexFit()
[b750b0a]103{
[82db145]104 fxCst.Clear();
[53f4746]105 fCovCst.Clear();
106 fCovCstInv.Clear();
[82db145]107 fXv.Clear();
108 fcovXv.Clear();
[b750b0a]109 fChi2List.Clear();
[537d24f]110 //
[91ef0b8]111 for (Int_t i = 0; i < fNtr; i++)
[537d24f]112 {
[53f4746]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];
[b750b0a]117 }
[82db145]118 fPar.clear();
[53f4746]119 fParNew.clear();
120 fCov.clear();
121 fCovNew.clear();
[82db145]122 //
123 ResetWrkArrays();
[b750b0a]124 ffi.clear();
[537d24f]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 //
[91ef0b8]145 x0(0) = -D * TMath::Sin(p0);
146 x0(1) = D * TMath::Cos(p0);
[537d24f]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);
[91ef0b8]170 x(2) = x0(2) + ct * phi / (2 * C);
[537d24f]171 //
172 return x;
173}
174//
[82db145]175void VertexFit::UpdateTrkArrays(Int_t i)
[537d24f]176{
177 //
[b750b0a]178 // Get track parameters, covariance and phase
179 Double_t fs = ffi[i]; // Get phase
[53f4746]180 TVectorD par = *fParNew[i];
[82db145]181 TMatrixDSym Cov = *fCov[i];
[537d24f]182 //
[82db145]183 // Fill all track related work arrays arrays
[b750b0a]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
[53f4746]188 TMatrixD At(TMatrixD::kTransposed, A); // A transposed
[b750b0a]189 fAti.push_back(new TMatrixD(At)); // Store A'
[53f4746]190 fWinvi.push_back(new TMatrixDSym(Winv)); // Store W^-1 matrix
[537d24f]191 //
[b750b0a]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 //
[82db145]198 TMatrixDSym W = RegInv(Winv); // W = (A*C*A')^-1
199 fWi.push_back(new TMatrixDSym(W)); // Store W matrix
[537d24f]200 //
[b750b0a]201 TVectorD a = derXds(par, fs); // a = dx/ds = derivatives wrt phase
[82db145]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//
[b750b0a]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//
[82db145]284void VertexFit::VertexFitter()
285{
286 //std::cout << "VertexFitter: just in" << std::endl;
[53f4746]287 if (fNtr < 2 && !fVtxCst){
[82db145]288 std::cout << "VertexFit::VertexFitter - Method called with less than 2 tracks - Aborting " << std::endl;
289 std::exit(1);
[537d24f]290 }
291 //
[82db145]292 // Vertex fit
293 //
294 // Initial variable definitions
295 TVectorD x(3);
296 TMatrixDSym covX(3);
297 Double_t Chi2 = 0;
298 //
[b750b0a]299 VtxFitNoSteer(); // Fast vertex finder on first pass (set ffi and fXv)
300 TVectorD x0 = fXv;
[537d24f]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 //
[82db145]309 // Iteration loop
[537d24f]310 while (epsi > eps && Ntry < TryMax) // Iterate until found vertex is stable
311 {
[82db145]312 // Initialize arrays
[537d24f]313 x.Zero();
314 TVectorD cterm(3); TMatrixDSym H(3); TMatrixDSym DW1D(3);
[82db145]315 covX.Zero(); // Reset vertex covariance
[b750b0a]316 cterm.Zero(); // Reset constant term
[537d24f]317 H.Zero(); // Reset H matrix
318 DW1D.Zero();
[127644a]319 //
[82db145]320 // Reset work arrays
321 //
322 ResetWrkArrays();
323 //
324 // Start loop on tracks
325 //
[537d24f]326 for (Int_t i = 0; i < fNtr; i++)
327 {
[127644a]328 // Get track helix parameters and their covariance matrix
[537d24f]329 TVectorD par = *fPar[i];
[127644a]330 TMatrixDSym Cov = *fCov[i];
[82db145]331 //
332 // Update track related arrays
[537d24f]333 //
[82db145]334 UpdateTrkArrays(i);
335 TMatrixDSym Ds = *fDi[i];
336 TMatrixDSym Winv = *fWinvi[i];
[537d24f]337 TMatrixDSym DsW1Ds = Winv.Similarity(Ds); // Service matrix to calculate covX
[82db145]338 //
339 // Update global arrays
[127644a]340 DW1D += DsW1Ds;
[537d24f]341 // Update hessian
342 H += Ds;
343 // update constant term
[53f4746]344 TVectorD xs = *fx0i[i] - *fdi[i];
[b750b0a]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
[537d24f]351 cterm += Ds * xs;
352 } // End loop on tracks
[53f4746]353 // Some additions in case of external constraints
354 if (fVtxCst) {
355 H += fCovCstInv;
356 cterm += fCovCstInv * fxCst;
357 DW1D += fCovCstInv;
358 }
[537d24f]359 //
360 // update vertex position
[82db145]361 TMatrixDSym H1 = RegInv(H);
[91ef0b8]362 x = H1 * cterm;
[82db145]363 //
[537d24f]364 // Update vertex covariance
365 covX = DW1D.Similarity(H1);
[82db145]366 //
[537d24f]367 // Update phases and chi^2
368 Chi2 = 0.0;
369 for (Int_t i = 0; i < fNtr; i++)
370 {
[53f4746]371 TVectorD lambda = (*fDi[i]) * (*fx0i[i] - x - *fdi[i]);
[82db145]372 TMatrixDSym Wm1 = *fWinvi[i];
[537d24f]373 fChi2List(i) = Wm1.Similarity(lambda);
[b750b0a]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 }
[537d24f]378 Chi2 += fChi2List(i);
[82db145]379 TVectorD a = *fai[i];
[b750b0a]380 TVectorD b = (*fWi[i]) * (x - *fx0i[i] + *fdi[i]);
381 ffi[i] += Dot(a, b) / fa2i[i];
[53f4746]382 TVectorD newPar = *fPar[i] - ((*fCov[i]) * (*fAti[i])) * lambda;
383 fParNew[i] = new TVectorD(newPar);
[537d24f]384 }
[53f4746]385 // Add external constraint to Chi2
386 if (fVtxCst) Chi2 += fCovCstInv.Similarity(x - fxCst);
[537d24f]387 //
388 TVectorD dx = x - x0;
389 x0 = x;
390 // update vertex stability
[82db145]391 TMatrixDSym Hess = RegInv(covX);
[537d24f]392 epsi = Hess.Similarity(dx);
393 Ntry++;
394 //
395 // Store result
396 //
[53f4746]397 fXv = x; // Vertex position
[537d24f]398 fcovXv = covX; // Vertex covariance
399 fChi2 = Chi2; // Vertex fit Chi2
[53f4746]400 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2) << std::endl;
[82db145]401 } // end of iteration loop
[537d24f]402 //
[82db145]403 fVtxDone = kTRUE; // Set fit completion flag
[b750b0a]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;
[537d24f]407 //
408}
409//
[82db145]410// Return fit vertex
[537d24f]411TVectorD VertexFit::GetVtx()
412{
[82db145]413 if (!fVtxDone) VertexFitter();
[537d24f]414 return fXv;
415}
[82db145]416//
417// Return fit vertex covariance
[537d24f]418TMatrixDSym VertexFit::GetVtxCov()
419{
[82db145]420 if (!fVtxDone) VertexFitter();
[537d24f]421 return fcovXv;
422}
[82db145]423//
424// Return fit vertex chi2
[537d24f]425Double_t VertexFit::GetVtxChi2()
426{
[82db145]427 if (!fVtxDone) VertexFitter();
[537d24f]428 return fChi2;
429}
[82db145]430//
431// Return array of chi2 contributions from each track
[537d24f]432TVectorD VertexFit::GetVtxChi2List()
433{
[82db145]434 if (!fVtxDone) VertexFitter();
[537d24f]435 return fChi2List;
436}
437//
438// Handle tracks/constraints
439void VertexFit::AddVtxConstraint(TVectorD xv, TMatrixDSym cov) // Add gaussian vertex constraint
440{
[53f4746]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;
[537d24f]451}
452//
[82db145]453// Adding tracks one by one
454void VertexFit::AddTrk(TVectorD *par, TMatrixDSym *Cov) // Add track to input list
[537d24f]455{
[82db145]456 fNtr++;
457 fChi2List.ResizeTo(fNtr); // Resize chi2 array
458 fPar.push_back(par); // add new track
459 fCov.push_back(Cov);
[53f4746]460 fParNew.push_back(par); // add new track
461 fCovNew.push_back(Cov);
[82db145]462 //
463 // Reset previous vertex temp arrays
464 ResetWrkArrays();
465 ffi.clear();
466 fVtxDone = kFALSE; // Reset vertex done flag
[537d24f]467}
[82db145]468//
469// Removing tracks one by one
470void VertexFit::RemoveTrk(Int_t iTrk) // Remove iTrk track
[537d24f]471{
[82db145]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);
[53f4746]477 fParNew.erase(fParNew.begin() + iTrk); // Remove track
478 fCovNew.erase(fCovNew.begin() + iTrk);
[82db145]479 //
480 // Reset previous vertex temp arrays
481 ResetWrkArrays();
482 ffi.clear();
483 fVtxDone = kFALSE; // Reset vertex done flag
[537d24f]484}
Note: See TracBrowser for help on using the repository browser.