Fork me on GitHub

source: git/external/TrackCovariance/VertexFit.cc@ 3a105e5

Last change on this file since 3a105e5 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
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;
[7bca620]19 fRstart = -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;
[7bca620]33 fRstart = -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;
[7bca620]57 fRstart = -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
[7bca620]222 std::vector<TMatrixDSym*> Ci; // Position error matrix at fixed phase
[b750b0a]223 std::vector<TVectorD*> wi; // Ci*ni
[7bca620]224 std::vector<Double_t> s_in; // Starting phase
225 //
226 //
[b750b0a]227 //
228 // Track loop
229 for (Int_t i = 0; i < fNtr; i++)
230 {
231 TVectorD par = *fPar[i];
232 TMatrixDSym Cov = *fCov[i];
[7bca620]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)));
[b750b0a]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])));
[7bca620]245 s_in.push_back(s);
[b750b0a]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]);
[7bca620]273 ffi.push_back(si+s_in[i]);
[b750b0a]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();
[7bca620]291 s_in.clear();
[b750b0a]292}
293//
[82db145]294void VertexFit::VertexFitter()
295{
296 //std::cout << "VertexFitter: just in" << std::endl;
[53f4746]297 if (fNtr < 2 && !fVtxCst){
[82db145]298 std::cout << "VertexFit::VertexFitter - Method called with less than 2 tracks - Aborting " << std::endl;
299 std::exit(1);
[537d24f]300 }
301 //
[82db145]302 // Vertex fit
303 //
304 // Initial variable definitions
305 TVectorD x(3);
306 TMatrixDSym covX(3);
307 Double_t Chi2 = 0;
308 //
[b750b0a]309 VtxFitNoSteer(); // Fast vertex finder on first pass (set ffi and fXv)
310 TVectorD x0 = fXv;
[537d24f]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 //
[82db145]319 // Iteration loop
[537d24f]320 while (epsi > eps && Ntry < TryMax) // Iterate until found vertex is stable
321 {
[82db145]322 // Initialize arrays
[537d24f]323 x.Zero();
324 TVectorD cterm(3); TMatrixDSym H(3); TMatrixDSym DW1D(3);
[82db145]325 covX.Zero(); // Reset vertex covariance
[b750b0a]326 cterm.Zero(); // Reset constant term
[537d24f]327 H.Zero(); // Reset H matrix
328 DW1D.Zero();
[127644a]329 //
[82db145]330 // Reset work arrays
331 //
332 ResetWrkArrays();
333 //
334 // Start loop on tracks
335 //
[537d24f]336 for (Int_t i = 0; i < fNtr; i++)
337 {
[127644a]338 // Get track helix parameters and their covariance matrix
[537d24f]339 TVectorD par = *fPar[i];
[127644a]340 TMatrixDSym Cov = *fCov[i];
[82db145]341 //
342 // Update track related arrays
[537d24f]343 //
[82db145]344 UpdateTrkArrays(i);
345 TMatrixDSym Ds = *fDi[i];
346 TMatrixDSym Winv = *fWinvi[i];
[537d24f]347 TMatrixDSym DsW1Ds = Winv.Similarity(Ds); // Service matrix to calculate covX
[82db145]348 //
349 // Update global arrays
[127644a]350 DW1D += DsW1Ds;
[537d24f]351 // Update hessian
352 H += Ds;
353 // update constant term
[53f4746]354 TVectorD xs = *fx0i[i] - *fdi[i];
[b750b0a]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
[537d24f]361 cterm += Ds * xs;
362 } // End loop on tracks
[53f4746]363 // Some additions in case of external constraints
364 if (fVtxCst) {
365 H += fCovCstInv;
366 cterm += fCovCstInv * fxCst;
367 DW1D += fCovCstInv;
368 }
[537d24f]369 //
370 // update vertex position
[82db145]371 TMatrixDSym H1 = RegInv(H);
[91ef0b8]372 x = H1 * cterm;
[82db145]373 //
[537d24f]374 // Update vertex covariance
375 covX = DW1D.Similarity(H1);
[82db145]376 //
[537d24f]377 // Update phases and chi^2
378 Chi2 = 0.0;
379 for (Int_t i = 0; i < fNtr; i++)
380 {
[53f4746]381 TVectorD lambda = (*fDi[i]) * (*fx0i[i] - x - *fdi[i]);
[82db145]382 TMatrixDSym Wm1 = *fWinvi[i];
[537d24f]383 fChi2List(i) = Wm1.Similarity(lambda);
[b750b0a]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 }
[537d24f]388 Chi2 += fChi2List(i);
[82db145]389 TVectorD a = *fai[i];
[b750b0a]390 TVectorD b = (*fWi[i]) * (x - *fx0i[i] + *fdi[i]);
391 ffi[i] += Dot(a, b) / fa2i[i];
[53f4746]392 TVectorD newPar = *fPar[i] - ((*fCov[i]) * (*fAti[i])) * lambda;
393 fParNew[i] = new TVectorD(newPar);
[537d24f]394 }
[53f4746]395 // Add external constraint to Chi2
396 if (fVtxCst) Chi2 += fCovCstInv.Similarity(x - fxCst);
[537d24f]397 //
398 TVectorD dx = x - x0;
399 x0 = x;
400 // update vertex stability
[82db145]401 TMatrixDSym Hess = RegInv(covX);
[537d24f]402 epsi = Hess.Similarity(dx);
403 Ntry++;
404 //
405 // Store result
406 //
[53f4746]407 fXv = x; // Vertex position
[537d24f]408 fcovXv = covX; // Vertex covariance
409 fChi2 = Chi2; // Vertex fit Chi2
[53f4746]410 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2) << std::endl;
[82db145]411 } // end of iteration loop
[537d24f]412 //
[82db145]413 fVtxDone = kTRUE; // Set fit completion flag
[7bca620]414 //fRstart = TMath::Sqrt(fXv(0)*fXv(0) + fXv(1)*fXv(1)); // Store fit
[b750b0a]415 //std::cout << "Found vertex " << fXv(0) << ", " << fXv(1) << ", " << fXv(2)
416 // << ", after "<<Ntry<<" iterations"<<std::endl;
[537d24f]417 //
418}
419//
[82db145]420// Return fit vertex
[537d24f]421TVectorD VertexFit::GetVtx()
422{
[82db145]423 if (!fVtxDone) VertexFitter();
[537d24f]424 return fXv;
425}
[82db145]426//
427// Return fit vertex covariance
[537d24f]428TMatrixDSym VertexFit::GetVtxCov()
429{
[82db145]430 if (!fVtxDone) VertexFitter();
[537d24f]431 return fcovXv;
432}
[82db145]433//
434// Return fit vertex chi2
[537d24f]435Double_t VertexFit::GetVtxChi2()
436{
[82db145]437 if (!fVtxDone) VertexFitter();
[537d24f]438 return fChi2;
439}
[82db145]440//
441// Return array of chi2 contributions from each track
[537d24f]442TVectorD VertexFit::GetVtxChi2List()
443{
[82db145]444 if (!fVtxDone) VertexFitter();
[537d24f]445 return fChi2List;
446}
447//
448// Handle tracks/constraints
449void VertexFit::AddVtxConstraint(TVectorD xv, TMatrixDSym cov) // Add gaussian vertex constraint
450{
[53f4746]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;
[537d24f]461}
462//
[82db145]463// Adding tracks one by one
464void VertexFit::AddTrk(TVectorD *par, TMatrixDSym *Cov) // Add track to input list
[537d24f]465{
[82db145]466 fNtr++;
467 fChi2List.ResizeTo(fNtr); // Resize chi2 array
468 fPar.push_back(par); // add new track
469 fCov.push_back(Cov);
[53f4746]470 fParNew.push_back(par); // add new track
471 fCovNew.push_back(Cov);
[82db145]472 //
473 // Reset previous vertex temp arrays
474 ResetWrkArrays();
475 ffi.clear();
476 fVtxDone = kFALSE; // Reset vertex done flag
[537d24f]477}
[82db145]478//
479// Removing tracks one by one
480void VertexFit::RemoveTrk(Int_t iTrk) // Remove iTrk track
[537d24f]481{
[82db145]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);
[53f4746]487 fParNew.erase(fParNew.begin() + iTrk); // Remove track
488 fCovNew.erase(fCovNew.begin() + iTrk);
[82db145]489 //
490 // Reset previous vertex temp arrays
491 ResetWrkArrays();
492 ffi.clear();
493 fVtxDone = kFALSE; // Reset vertex done flag
[537d24f]494}
Note: See TracBrowser for help on using the repository browser.