KalmanVertexFit.cxx

Go to the documentation of this file.

//
//  Kalman Vertex fit 
//  The algoritm is based on a paper by 
//  R. Fruhwirth etc Computer Physics Communications 96(1996) 189-208
//  The formulism for BES3 can be find at BESIII Software Note: No. xx
//  Author: Kanglin He and Min Xu,  Apr. 13, 2007 for created
//
 
 
#include "VertexFit/KalmanVertexFit.h"
#include "VertexFit/HTrackParameter.h"
#include "VertexFit/WTrackParameter.h"
#include "math.h"
 
using namespace std;
 
KalmanVertexFit * KalmanVertexFit::m_pointer = 0;
 
KalmanVertexFit * KalmanVertexFit::instance() {
  if(m_pointer) return m_pointer;
  m_pointer = new KalmanVertexFit();
  return m_pointer;
}
 
KalmanVertexFit::~KalmanVertexFit() {;}
 
KalmanVertexFit::KalmanVertexFit() {;}
 
void KalmanVertexFit::init() {
  m_x = HepVector(3, 0);
  m_C0 = HepSymMatrix(3, 0);
  m_C = HepSymMatrix(3, 0);
  m_chisq = 0;
  m_ndof = -3;
 
  m_flag.clear();
  m_p.clear();
  m_hTrkOrigin.clear();
  m_hTrkInfit.clear(); //xum
  //m_wTrkInfit.clear();
 
  m_G.clear();
  m_A.clear();
  m_B.clear();
  m_c.clear();
  m_W.clear();
  m_GB.clear();
  m_chiF.clear();
 
  //
  // chi-square cut, number of iteration etc
  //
  m_chisqCutforTrack = 50;
  m_maxVertexIteration = 3;
  m_maxTrackIteration = 5;
  m_chisqCutforTrackIteration = 0.1;
}
 
void KalmanVertexFit::initVertex(const VertexParameter vtx) {
  int ifail;
  m_x = vtx.Vx();
  m_C0 = (vtx.Evx()).inverse(ifail);
  m_C = m_C0;
}
 
VertexParameter KalmanVertexFit::vtx() const {
  int ifail;
  VertexParameter v;
  v.setVx(m_x);
 
  std::vector<int> trackid;
    v.setEvx(m_C.inverse(ifail));
  return v;
}
 
HepSymMatrix KalmanVertexFit::Ex() const {
  int ifail;
  return m_C.inverse(ifail);
}
 
void KalmanVertexFit::addTrack(const HTrackParameter htrk) {
 
  int ifail;
 
  HepVector p(3, 0);
  p = htrk.p();
  
  m_p.push_back(p);
  m_flag.push_back(0);
  m_chiF.push_back(999.);
 
  HepMatrix A(5, 3, 0);
  HepMatrix B(5, 3, 0);
  HepSymMatrix G(5, 0); 
 
  G = (htrk.eHel()).inverse(ifail);
 
  m_hTrkOrigin.push_back(htrk);
 
  //xum
  m_hTrkInfit.push_back(htrk); 
  //m_wTrkInfit.push_back(htrk.wTrack());
 
  m_G.push_back(G);
  m_A.push_back(A);
  m_B.push_back(B);
  
  HepSymMatrix W(3, 0);
  HepSymMatrix GB(5, 0);
  HepVector c(5, 0);
 
  m_W.push_back(W);
  m_GB.push_back(GB);
  m_c.push_back(c);
 
}
 
int KalmanVertexFit::numTrack() const {
  int num = 0;
  for(int k = 0; k < m_hTrkOrigin.size(); k++) 
    if(m_flag[k] == 0) num = num +1;
  
  return num;
}
 
std::vector<int> KalmanVertexFit::trackIDList() const {
  std::vector<int> trackid;
  for(int k = 0; k < m_hTrkOrigin.size(); k++)
    if(m_flag[k] == 0) trackid.push_back(trackID(k));
 
  return trackid;
}
 
 
int KalmanVertexFit::filter(const int k) {
  int ifail;
  double chisq[10];
  ////////////////////////////////////////////////
  // start iteration
  ///////////////////////////////////////////////
  HepVector pp(3, 0);
  HepVector xp(3, 0);
  HepSymMatrix CP(3, 0);
  xp = m_x;
  CP = m_C;
  pp = m_p[k];
 
  for(int iter = 0; iter < m_maxTrackIteration; iter++) {
    //
    // calculate derivative matrices for track k
    //
    updateMatrices(k, pp, xp);
    //
    // vertex covariance matrix updated
    //
    HepSymMatrix CK(3, 0);
    CK = CP + (m_GB[k].similarity(m_A[k].T()));
    //
    // state vector x and p updated
    //
    HepVector x(3, 0);
    HepVector p(3, 0);
    x = CK.inverse(ifail) * (CP * xp + m_A[k].T() * m_GB[k] * 
                 (m_hTrkOrigin[k].hel() - m_c[k]));
    p = m_W[k] * m_B[k].T() * m_G[k] *(m_hTrkOrigin[k].hel()- m_c[k]- 
                       m_A[k]*x);
    // chi-square of filter 
    HepSymMatrix chif(1, 0);
    chif = CP.similarity((x-xp).T()) + 
      m_G[k].similarity((m_hTrkOrigin[k].hel()- 
             m_c[k]- m_A[k]*x - m_B[k]*p).T());
    
    //
    // save current 3-momentum and vertex for track k
    //
    // 3-momentum
    pp = p;
    // current vertex and its covaiance matrix
    xp = x;
    CP = CK;
    chisq[iter] = chif[0][0];
    m_chiF[k] = chif[0][0];
    //    std::cout << "chisq, iter = " <<m_chiF[k] << ", " << iter << ", " << k <<std::endl;
    if(iter > 0) {
      double delchi = (chisq[iter]-chisq[iter-1]);
      if(fabs(delchi) < m_chisqCutforTrackIteration) break;
    }
  }
 
  if(m_chiF[k] > m_chisqCutforTrack) {
    //    do not update vertex if failed
    m_flag[k]=1;
    return 1;
  } else { 
    // update vertex if filter success
    m_x = xp;
    m_C = CP;
    m_p[k] = pp;
    return 0;
  }
}
 
void KalmanVertexFit::inverse(const int k) {
  //
  // inverse kalman filter for track k
  //
 
  int ifail;
  HepVector x(3, 0);
  HepVector p(3, 0);
  HepSymMatrix C(3, 0);
  C = m_C - m_GB[k].similarity(m_A[k].T());
  x = C.inverse(ifail)*(m_C*m_x-m_A[k].T() * m_GB[k] * 
            (m_hTrkOrigin[k].hel()-m_c[k]));
  //
  //  update vertex position and its covariance matrix
  //
  m_x = x; 
  m_C = C; 
}
void KalmanVertexFit::filter() {
  
  for(int iter = 0; iter < m_maxVertexIteration; iter++) {
    //
    //  at start of each iteration, 
    //  the covariance matrix take the initial values: m_C0
    //  the vertex position take the update value by the last iteration
    //
    m_C = m_C0;   // covariance matrix of verter --> initial values
    m_chisq = 0;  // total chi-square of smooth
    m_ndof = -3;  // number of degrees of freedom
   
    for(int k = 0; k < m_hTrkOrigin.size(); k++) {
      if(m_flag[k]==1) continue;
      if(filter(k)==0) // filter track k
    m_ndof += 2;
 
    }  
    //
    //  check the chi-square of smooth, make decision to remove bad tracks
    //  user may remove more bad track by tight chi-square cut
    //  through the following interface 
    //  KalmanVertexFit::remove(k)
    //
    for(int k = 0; k < m_hTrkOrigin.size(); k++) {
      if(m_flag[k]==1) continue;
      double chi2 = chiS(k);
      if(chi2 < 0) { // remove track k from vertex
    remove(k);
    continue;
      }
      
      if(chi2 > m_chisqCutforTrack) { // remove track k from vertex
    remove(k);
    continue;
      } 
      m_chisq += chi2;
 
    }
  } // end of iteration
}
 
double KalmanVertexFit::chiS(const int k) const {
  if(m_flag[k]==1) return 999;
  int ifail;
  HepVector x(3, 0);
  HepVector p(3, 0);
  HepSymMatrix C(3, 0);
  C = m_C - m_GB[k].similarity(m_A[k].T());
  x = C.inverse(ifail)*(m_C*m_x-m_A[k].T() * m_GB[k] * 
            (m_hTrkOrigin[k].hel()-m_c[k]));
  p = m_W[k]* m_B[k].T()*m_G[k]*
    (m_hTrkOrigin[k].hel()-m_c[k]-m_A[k]*m_x);
  HepSymMatrix chis(1, 0);
  chis = C.similarity((x-m_x).T())+
    m_G[k].similarity((m_hTrkOrigin[k].hel()-m_c[k]-
               m_A[k]*m_x-m_B[k]*p).T());
  return chis[0][0];
}
 
void KalmanVertexFit::smooth(const int k) {
  // xum
  //
  // update htrk and wtrk parameters
  //
  if(m_flag[k] == 1) return;
  int ifail;
 
  HepVector pp(3, 0);
  HepVector xp(3, 0);
  HepSymMatrix CP(3, 0);
  xp = m_x;
  CP = m_C;
  pp = m_p[k];
 
  updateMatrices(k, pp, xp);
  pp = m_W[k] * m_B[k].T() * m_G[k] * (m_hTrkOrigin[k].hel() - m_c[k] - m_A[k] * xp);
  
  updateMatrices(k, pp, xp);
  //for htrk
  HepVector helix(5, 0);
  helix = m_c[k] + m_A[k] * xp + m_B[k] * pp;
  HepMatrix E(3, 3, 0);
  E = -CP.inverse(ifail) * m_A[k].T() * m_G[k] * m_B[k] * m_W[k];
  HepSymMatrix D(3, 0);
  D = m_W[k] + CP.similarity(E.T());
  
  HepMatrix middle(5, 5, 0);
  HepSymMatrix error(5, 0);
  middle = (CP.inverse(ifail)).similarity(m_A[k]) + m_A[k] * E * m_B[k].T() +
          (m_A[k] * E * m_B[k].T()).T() + D.similarity(m_B[k]);  
  error.assign(middle);
 
  m_hTrkInfit[k].setHel(helix);
  m_hTrkInfit[k].setEHel(error);
  
  m_p[k] = pp;
/*
  //for wtrk
  double mass = m_hTrkOrigin[k].xmass(m_hTrkOrigin[k].partID());
  double Energy = sqrt(pp[0]*pp[0] + pp[1]*pp[1] + pp[2]*pp[2] + mass * mass);
 
  HepMatrix Awtrk(7, 3, 0), Bwtrk(7, 3, 0);
  Awtrk[4][0] = Awtrk[5][1] = Awtrk[6][2] = 1;
  Bwtrk[0][0] = Bwtrk[1][1] = Bwtrk[2][2] = 1;
  Bwtrk[3][0] = pp[0] / Energy;
  Bwtrk[3][1] = pp[1] / Energy;
  Bwtrk[3][2] = pp[2] / Energy;
    
  HepVector w(7, 0);
  HepSymMatrix Ew(7, 0);
  w[0] = pp[0];   w[1] = pp[1];  w[2] = pp[2];  w[3] = Energy; 
  w[4] = xp[0];   w[5] = xp[1];  w[6] = xp[2];
  
  HepSymMatrix Gwtrk(7, 0);
  Gwtrk = m_hTrkOrigin[k].wTrack().Ew().inverse(ifail);
  HepSymMatrix Wwtrk(3, 0);
  Wwtrk = Gwtrk.similarity(Bwtrk.T()).inverse(ifail);
 
  HepMatrix Ewtrk(3, 3, 0);
  Ewtrk = -CP.inverse(ifail) * m_Awtrk[k].T() * m_Gwtrk[k] * m_Bwtrk[k] * m_Wwtrk[k];
  HepSymMatrix Dwtrk(3, 0);
  Dwtrk = m_Wwtrk[k] + CP.similarity(Ewtrk.T());
 
  HepMatrix Ewmiddle(7, 7, 0);
  Ewmiddle = (CP.inverse(ifail)).similarity(m_Awtrk[k]) + m_Awtrk[k] * Ewtrk * m_Bwtrk[k].T() +
          (m_Awtrk[k] * Ewtrk * m_Bwtrk[k].T()).T() + Dwtrk.similarity(m_Bwtrk[k]);         
  Ew.assign(Ewmiddle);
 
  m_wTrkInfit[k].setCharge(m_hTrkOrigin[k].charge());
  m_wTrkInfit[k].setW(w);
  m_wTrkInfit[k].setEw(Ew);*/
}
 
void KalmanVertexFit::remove(const int k) {
  //
  // remove track k from the vertex fit
  //
  /*
    m_hTrkOrigin.erase(m_hTrkOrigin.begin()+k);
    m_p.erase(m_p.begin()+k);
    m_flag.erase(m_flag.begin()+k);
    m_A.erase(m_A.begin()+k);
    m_B.erase(m_B.begin()+k);
    m_c.erase(m_c.begin()+k);
    m_G.erase(m_G.begin()+k);
    m_GB.erase(m_GB.begin()+k);
    m_W.erase(m_W.begin()+k);
    m_chiF.erase(m_chiF.begin()+k);
  */
 
  if(m_flag[k]==1) return; // track k already removed
  inverse(k);
  m_ndof -= 2;
  m_flag[k]=1;
  //xum
  //m_chisq -=chiS(k);
}
 
 
HTrackParameter KalmanVertexFit::hTrack(const int k) const{
  //add xum
  //HTrackParameter htrk;
  //return htrk;
  return m_hTrkInfit[k];
}
 
WTrackParameter KalmanVertexFit::wTrack(const int k, const double mass) const {
  if(m_flag[k] == 1) return m_hTrkOrigin[k].wTrack();
  int ifail;
 
  HepVector pp(3, 0);
  HepVector xp(3, 0);
  HepSymMatrix CP(3, 0);
  xp = m_x;
  CP = m_C;
  pp = m_p[k];
 
  WTrackParameter wtrk;
  double Energy = sqrt(pp[0]*pp[0] + pp[1]*pp[1] + pp[2]*pp[2] + mass * mass);
  
  HepMatrix Awtrk(7, 3, 0), Bwtrk(7, 3, 0);
  Awtrk[4][0] = Awtrk[5][1] = Awtrk[6][2] = 1;
  Bwtrk[0][0] = Bwtrk[1][1] = Bwtrk[2][2] = 1;
  Bwtrk[3][0] = pp[0] / Energy;
  Bwtrk[3][1] = pp[1] / Energy;
  Bwtrk[3][2] = pp[2] / Energy;
    
  HepVector w(7, 0);
  HepSymMatrix Ew(7, 0);
  w[0] = pp[0];   w[1] = pp[1];  w[2] = pp[2];  w[3] = Energy;
  w[4] = xp[0];   w[5] = xp[1];  w[6] = xp[2];  
  
  HepSymMatrix Gwtrk(7, 0);
  Gwtrk = m_hTrkOrigin[k].wTrack().Ew().inverse(ifail);
  HepSymMatrix Wwtrk(3, 0);
  Wwtrk = Gwtrk.similarity(Bwtrk.T()).inverse(ifail);
  
  HepMatrix Ewtrk(3, 3, 0);
  Ewtrk = -CP.inverse(ifail) * Awtrk.T() * Gwtrk * Bwtrk * Wwtrk;
  HepSymMatrix Dwtrk(3, 0);
  Dwtrk = Wwtrk + CP.similarity(Ewtrk.T());
 
  HepMatrix Ewmiddle(7, 7, 0);
  Ewmiddle = (CP.inverse(ifail)).similarity(Awtrk) + Awtrk * Ewtrk * Bwtrk.T() +
          (Awtrk * Ewtrk * Bwtrk.T()).T() + Dwtrk.similarity(Bwtrk);
  Ew.assign(Ewmiddle);
 
  wtrk.setCharge(m_hTrkOrigin[k].charge());
  wtrk.setW(w);
  wtrk.setEw(Ew);
 
  return wtrk;
}
 
//
// private functions, update derivative matrices
//
void KalmanVertexFit::updateMatrices(const int k, const HepVector p, const HepVector x) {
 
  int ifail;
  //
  // expand measurement equation at (x, p)
  //
  HTrackParameter he(m_hTrkOrigin[k].charge(), p, x);
 
  // derivative matrix
  
  m_A[k] = he.dHdx(p, x);
  m_B[k] = he.dHdp(p, x);
 
  // W, GB, c
  m_W[k] = (m_G[k].similarity(m_B[k].T())).inverse(ifail);
  m_GB[k] = m_G[k] - (m_W[k].similarity(m_B[k])).similarity(m_G[k]);
  m_c[k] = he.hel() - m_A[k] * x - m_B[k] * p;  
}
 
void KalmanVertexFit::updateMatrices(const int k) {
 
  int ifail;
  HepVector p(3, 0);
  HepVector x(3, 0);
  p = m_p[k];
  x = m_x;
  updateMatrices(k, p, x);
}
 
HepVector KalmanVertexFit::pull(const int k) {
  HepVector n_pull(5, 0);
  n_pull[0] = n_pull[1] = n_pull[2] = n_pull[3] = n_pull[4] = 999.;
  if(m_flag[k] == 1) return n_pull; 
  for (int i = 0 ; i < 5; i++) {
    double cov = m_hTrkOrigin[k].eHel()[i][i] - m_hTrkInfit[k].eHel()[i][i];
    if (cov == 0.) continue;    
    n_pull[i] = (m_hTrkOrigin[k].hel()[i] - m_hTrkInfit[k].hel()[i]) / sqrt(cov);
  }
  return n_pull;
}
 
double KalmanVertexFit::pullmomentum(const int k) {
  double pull = 999.;
  if(m_flag[k] == 1) return pull;
 
  double kappa_origin = m_hTrkOrigin[k].kappa();
  double kappa_infit = m_hTrkInfit[k].kappa();
  double lamb_origin = m_hTrkOrigin[k].lambda();
  double lamb_infit = m_hTrkInfit[k].lambda();
 
  double Vkappa_origin = m_hTrkOrigin[k].eHel()[2][2];
  double Vkappa_infit = m_hTrkInfit[k].eHel()[2][2];
  double Vlamb_origin = m_hTrkOrigin[k].eHel()[4][4];
  double Vlamb_infit = m_hTrkInfit[k].eHel()[4][4];
  double V_kappa_lamb_origin = m_hTrkOrigin[k].eHel()[4][2];
  double V_kappa_lamb_infit = m_hTrkInfit[k].eHel()[4][2];
 
  double P_origin = calculationP(kappa_origin, lamb_origin);
  //cout << "P_origin = " << P_origin << endl;
  //P_origin = m_hTrkOrigin[k].p()[0] * m_hTrkOrigin[k].p()[0]
  //          + m_hTrkOrigin[k].p()[1] * m_hTrkOrigin[k].p()[1]
  //          + m_hTrkOrigin[k].p()[2] * m_hTrkOrigin[k].p()[2];
  //cout << "P = " << P_origin << endl;
  double P_infit = calculationP(kappa_infit, lamb_infit);
 
  double SigmaP_origin = calculationSigmaP(kappa_origin, lamb_origin, Vkappa_origin,
                        Vlamb_origin, V_kappa_lamb_origin);
  double SigmaP_infit = calculationSigmaP(kappa_infit, lamb_infit, Vkappa_infit,
                        Vlamb_infit, V_kappa_lamb_infit);
  if ((SigmaP_origin - SigmaP_infit) <= 0.)  return pull;
  pull = (P_origin - P_infit) / sqrt (SigmaP_origin - SigmaP_infit);
 
  return pull;
}
 
//====================== sub routine ==================================
double KalmanVertexFit::calculationP(const double kappa, const double lamb) {
  double P = 0.0;
  P = sqrt(1 + lamb * lamb) / kappa;
  //cout << "P in calculationP = " << P << endl;
  return P; 
}
 
double KalmanVertexFit::calculationSigmaP(const double kappa, const double lamb, 
                                  const double Vkappa, const double Vlamb, 
                                  const double Vkappa_lamb) {
  double SigmaP = 0.0;
  double dp_dkappa = -sqrt(1 + lamb*lamb) /kappa/kappa;
  double dp_dlamb = lamb / kappa / sqrt(1 + lamb*lamb);
  SigmaP = dp_dkappa * dp_dkappa * Vkappa + dp_dlamb * dp_dlamb * Vlamb
         + 2 * dp_dkappa * dp_dlamb * Vkappa_lamb;
  return SigmaP;
}