Fork me on GitHub

Context Navigation

source: git/external/TrackCovariance/VertexFit.cc

Visit:

Last change on this file 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

Rev	Line
[b750b0a]	1	/*
	2	Vertex 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]	16	VertexFit::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]	30	VertexFit::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]	54	VertexFit::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	//
	78	void 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]	102	VertexFit::~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	//
	130	TVectorD 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	//
	152	TVectorD 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]	175	void 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 = ACA'
	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 = (ACA')^-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]	215	void 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; // Cini
[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((fRstartfRstart-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]	294	void 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]	421	TVectorD VertexFit::GetVtx()
	422	{
[82db145]	423	if (!fVtxDone) VertexFitter();
[537d24f]	424	return fXv;
	425	}
[82db145]	426	//
	427	// Return fit vertex covariance
[537d24f]	428	TMatrixDSym VertexFit::GetVtxCov()
	429	{
[82db145]	430	if (!fVtxDone) VertexFitter();
[537d24f]	431	return fcovXv;
	432	}
[82db145]	433	//
	434	// Return fit vertex chi2
[537d24f]	435	Double_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]	442	TVectorD VertexFit::GetVtxChi2List()
	443	{
[82db145]	444	if (!fVtxDone) VertexFitter();
[537d24f]	445	return fChi2List;
	446	}
	447	//
	448	// Handle tracks/constraints
	449	void 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
	464	void 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
	480	void 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.

Download in other formats: