KalmanStepper.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2017 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Nils Braun                                               *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
#pragma once
 
#include <tracking/ckf/general/utilities/EigenHelper.h>
 
#include <genfit/MeasuredStateOnPlane.h>
#include <genfit/MeasurementOnPlane.h>
 
#include <Eigen/Dense>
 
 
namespace Belle2 {
  template <unsigned int Dimension>
  class KalmanStepper {
    using MeasurementState = Eigen::Matrix<double, Dimension, 1>;
    using MeasurementCovariance = Eigen::Matrix<double, Dimension, Dimension>;
    using HMatrix = Eigen::Matrix<double, Dimension, 5>;
    using TrackState = Eigen::Matrix<double, 5, 1>;
    using TrackCovariance = Eigen::Matrix<double, 5, 5>;
 
  public:
    double kalmanStep(genfit::MeasuredStateOnPlane& measuredStateOnPlane,
                      const genfit::MeasurementOnPlane& measurementOnPlane) const
    {
      // Extract matrices from the state
      TrackState x_k = convertToEigen<5>(measuredStateOnPlane.getState());
      TrackCovariance C_k = convertToEigen<5, 5>(measuredStateOnPlane.getCov());
 
      // extract matrices from the measurement
      const MeasurementState& m_k = convertToEigen<Dimension>(measurementOnPlane.getState());
      const HMatrix& H_k = convertToEigen<Dimension, 5>(measurementOnPlane.getHMatrix()->getMatrix());
      const MeasurementCovariance& V_k = convertToEigen<Dimension, Dimension>(measurementOnPlane.getCov());
 
      const double chi2 = kalmanStep(x_k, C_k, m_k, V_k, H_k);
 
      // set the state back
      measuredStateOnPlane.setState(TVectorD(5, x_k.data()));
      measuredStateOnPlane.setCov(TMatrixDSym(5, C_k.data()));
 
      return chi2;
    }
 
    double calculateResidual(genfit::MeasuredStateOnPlane& measuredStateOnPlane,
                             const genfit::MeasurementOnPlane& measurementOnPlane) const
    {
      // Extract matrices from the state
      const TrackState& x_k = convertToEigen<5>(measuredStateOnPlane.getState());
 
      // extract matrices from the measurement
      const MeasurementState& m_k = convertToEigen<Dimension>(measurementOnPlane.getState());
      const HMatrix& H_k = convertToEigen<Dimension, 5>(measurementOnPlane.getHMatrix()->getMatrix());
 
      const MeasurementState& residual = m_k - H_k * x_k;
      return residual.norm();
    }
 
  private:
    double kalmanStep(TrackState& x_k, TrackCovariance& C_k,
                      const MeasurementState& m_k,
                      const MeasurementCovariance& V_k,
                      const HMatrix& H_k) const
    {
      const Eigen::Matrix<double, 5, Dimension>& K_k = C_k * H_k.transpose() * (V_k + H_k * C_k * H_k.transpose()).inverse();
      C_k -= K_k * H_k * C_k;
      x_k += K_k * (m_k - H_k * x_k);
 
      const MeasurementState& residual = m_k - H_k * x_k;
      const double chi2 = (residual.transpose() * (V_k - H_k * C_k * H_k.transpose()).inverse() * residual).value();
      return chi2;
    }
  };
}