KfTrackKalmanFilter.h

Go to the documentation of this file.

/* Copyright (C) 2017-2021 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Sergey Gorbunov [committer], Maksym Zyzak */
 

 
#pragma once  // include this header only once per compilation unit
 
#include "KfFieldRegion.h"
#include "KfMeasurementTime.h"
#include "KfMeasurementU.h"
#include "KfMeasurementXy.h"
#include "KfSimd.h"
#include "KfTrackParam.h"
#include "KfUtils.h"
 
#include <type_traits>
 
namespace cbm::algo::kf
{
  enum class FitDirection
  {
    kUpstream,
    kDownstream
  };
 
  inline FitDirection operator!(FitDirection d)
  {
    return d == FitDirection::kUpstream ? FitDirection::kDownstream : FitDirection::kUpstream;
  }
 
  template<typename T>
  struct LinearizationFull {
    LinearizationFull() = default;
    LinearizationFull(const TrackParam<T> p)
      : x(p.GetX())
      , y(p.GetY())
      , tx(p.GetTx())
      , ty(p.GetTy())
      , qp(p.GetQp())
      , time(p.GetTime())
      , vi(p.GetVi())
    {
    }
 
    T x{T(0.)};     
    T y{T(0.)};     
    T tx{T(0.)};    
    T ty{T(0.)};    
    T qp{T(0.)};    
    T time{T(0.)};  
    T vi{T(0.)};    
  };
 
 
  template<typename T>
  struct LinearizationQp {
    LinearizationQp() = default;
    LinearizationQp(const TrackParam<T> p) : qp(p.GetQp()) {}
 
    T qp;  
  };
 
  template<template<typename T> class LinearizationT_ = LinearizationQp, DoFitTime FlagFitTime = DoFitTime::Y>
  struct FilterSettings {
    template<typename T>
    using LinearizationT             = LinearizationT_<T>;
    static constexpr bool kDoFitTime = (FlagFitTime == DoFitTime::Y);
  };

 
  template<typename DataT, class Settings = FilterSettings<>>
  class TrackKalmanFilter {
 
   public:
    using Linearization_t = typename Settings::template LinearizationT<DataT>;
    using DataTscal       = kf::utils::scaltype<DataT>;
    using DataTmask       = kf::utils::masktype<DataT>;
 
    static constexpr int kNactiveParams = 5 + (Settings::kDoFitTime ? 2 : 0);
 
 
    TrackKalmanFilter() = default;
 
    TrackKalmanFilter(const kf::TrackParam<DataT>& t) { SetTrack(t); }
 
    template<typename T>
    TrackKalmanFilter(const kf::TrackParam<T>& t)
    {
      SetTrack(t);
    }
 
    void SetMask(const DataTmask& m) { fMask = m; }
 
    template<typename T>
    void SetTrack(const kf::TrackParam<T>& t)
    {
      fTr.Set(t);
      fLinearisation = Linearization_t(fTr);
    }
 
    void SetTrack(DataT z, const LinearizationFull<DataT>& lin)
    {
      fTr.SetX(lin.x);
      fTr.SetY(lin.y);
      fTr.SetZ(z);
      fTr.SetTx(lin.tx);
      fTr.SetTy(lin.ty);
      fTr.SetQp(lin.qp);
      if constexpr (Settings::kDoFitTime) {
        fTr.SetTime(lin.time);
        fTr.SetVi(lin.vi);
      }
      fLinearisation = Linearization_t(fTr);
    }
 
    void SetLinearization(const Linearization_t& lin) { fLinearisation = lin; }
 
    kf::TrackParam<DataT>& Tr() { return fTr; }
 
    Linearization_t& Linearization() { return fLinearisation; }
 
    void SetOneEntry(const int i0, const TrackKalmanFilter& T1, const int i1);
 
    std::string ToString(int i = -1);


    void SetParticleMass(DataT mass)
    {
      fMass  = mass;
      fMass2 = mass * mass;
    }

    DataT GetParticleMass() const { return fMass; }

    DataT GetParticleMass2() const { return fMass2; }

    void SetMaxExtrapolationStep(double step) { fMaxExtrapolationStep = DataT(step); }

    DataT GetMaxExtrapolationStep() const { return fMaxExtrapolationStep; }

    void Filter1d(const kf::MeasurementU<DataT>& m);

    void FilterXY(const kf::MeasurementXy<DataT>& m, bool skipUnmeasuredCoordinates = false);

    void FilterTime(DataT t, DataT dt2, const DataTmask& m);

    void FilterTime(kf::MeasurementTime<DataT> mt) { FilterTime(mt.T(), mt.Dt2(), DataTmask(mt.NdfT() > DataT(0.))); }

    void FilterVi(DataT vi);

    void MeasureVelocityWithQp();

    void Extrapolate(DataT z, const kf::FieldRegion<DataT>& F);

    void ExtrapolateInOneStep(DataT z, const kf::FieldRegion<DataT>& Field);

    void ExtrapolateNoField(DataT z);
 

    void ExtrapolateLineInField(DataT z_out, const kf::FieldRegion<DataT>& F);

    void EnergyLossCorrection(DataT radThick, FitDirection direction);

    void EnergyLossCorrection(int atomicZ, DataTscal atomicA, DataTscal rho, DataTscal radLen, DataT radThick,
                              FitDirection direction);
 

    void MultipleScattering(DataT radThick, DataT tx0, DataT ty0, DataT qp0);

    void MultipleScattering(DataT radThick)
    {
      MultipleScattering(radThick, fTr.GetTx(), fTr.GetTy(), fLinearisation.qp);
    }

    void MultipleScatteringInThickMaterial(DataT radThick, DataT thickness, bool fDownstream);


    void GetMeasurementModelAtZline(DataT zm, const kf::FieldRegion<DataT>& Field, std::array<DataT, 5>& Jx,
                                    std::array<DataT, 5>& Jy) const;

    void FilterExtrapolatedXY(const kf::MeasurementXy<DataT>& m, const std::array<DataT, 5>& Jx,
                              const std::array<DataT, 5>& Jy);
 
 
    void FilterExtrapolatedY(const kf::MeasurementXy<DataT>& m, const std::array<DataT, 5>& Jy);
 
    void FilterExtrapolatedY(const kf::MeasurementXy<DataT>& mL, const std::array<DataT, 5>& jL,
                             const kf::MeasurementXy<DataT>& mM, DataT msM, const kf::MeasurementXy<DataT>& mR,
                             const std::array<DataT, 5>& jR);
 
    void FilterExtrapolatedYChi2(const kf::MeasurementXy<DataT>& mL, const std::array<DataT, 5>& jL,
                                 const kf::MeasurementXy<DataT>& mM, DataT msM, const kf::MeasurementXy<DataT>& mR,
                                 const std::array<DataT, 5>& jR);

    std::pair<DataT, DataT> ExtrapolateLineXdX2(DataT z_out) const;

    std::pair<DataT, DataT> ExtrapolateLineYdY2(DataT z_out) const;

    DataT ExtrapolateLineDxy(DataT z_out) const;

    static DataT ApproximateBetheBloch(DataT bg2);

    static DataT ApproximateBetheBloch(DataT bg2, DataT kp0, DataT kp1, DataT kp2, DataT kp3, DataT kp4);

    static std::tuple<DataT, DataT> GetChi2XChi2U(kf::MeasurementXy<DataT> m, DataT x, DataT y, DataT C00, DataT C10,
                                                  DataT C11);

    void GuessTrack(const DataT& trackZ, const DataT hitX[], const DataT hitY[], const DataT hitZ[], const DataT hitT[],
                    const DataT By[], const DataTmask hitW[], const DataTmask hitWtime[], int NHits);
 
   private:
    void CleanNonActiveCovariances();
 
    typedef const DataT cnst;

 
    DataTmask fMask{true};  
 
    kf::TrackParam<DataT> fTr{};       
    Linearization_t fLinearisation{};  
 
    DataT fMass{0.10565800};      
    DataT fMass2{fMass * fMass};  
 
    DataT fMaxExtrapolationStep{50.};  
 
  } _fvecalignment;
 
  // =============================================================================================
 
  template<typename DataT, class Settings>
  inline std::string TrackKalmanFilter<DataT, Settings>::ToString(int i)
  {
    return fTr.ToString(i);
  }
 
 
  template<typename DataT, class Settings>
  inline void TrackKalmanFilter<DataT, Settings>::SetOneEntry(const int i0, const TrackKalmanFilter& T1, const int i1)
  {
    fTr.SetOneEntry(i0, T1.fTr, i1);
    kf::utils::VecCopy<DataT, DataT, false, false>::CopyEntries(fLinearisation.qp, i0, T1.fLinearisation.qp, i1);
  }
 
  template<typename DataT, class Settings>
  inline std::pair<DataT, DataT> TrackKalmanFilter<DataT, Settings>::ExtrapolateLineXdX2(DataT z_out) const
  {
    DataT dz = (z_out - fTr.GetZ());
    return std::pair(fTr.GetX() + fTr.GetTx() * dz, fTr.C00() + dz * (2 * fTr.C20() + dz * fTr.C22()));
  }
 
  template<typename DataT, class Settings>
  inline std::pair<DataT, DataT> TrackKalmanFilter<DataT, Settings>::ExtrapolateLineYdY2(DataT z_out) const
  {
    DataT dz = (z_out - fTr.GetZ());
    return std::pair(fTr.GetY() + fTr.GetTy() * dz, fTr.C11() + dz * (DataT(2.) * fTr.C31() + dz * fTr.C33()));
  }
 
  template<typename DataT, class Settings>
  inline DataT TrackKalmanFilter<DataT, Settings>::ExtrapolateLineDxy(DataT z_out) const
  {
    DataT dz = (z_out - fTr.GetZ());
    return fTr.C10() + dz * (fTr.C21() + fTr.C30() + dz * fTr.C32());
  }
 
 
  template<typename DataT, class Settings>
  inline void TrackKalmanFilter<DataT, Settings>::Extrapolate(DataT z, const kf::FieldRegion<DataT>& F)
  {
    // extrapolates the track to z
    // - makes several steps when the distance is larger than fMaxExtrapolationStep
    if (F.GetFieldType() == kf::EFieldType::Null) {
      ExtrapolateNoField(z);
    }
    else {
      DataT sgn       = kf::utils::iif(fTr.GetZ() < z, DataT(1.), DataT(-1.));
      DataT step      = sgn * fMaxExtrapolationStep;
      DataT zLastLoop = z - step;
      DataT zNext     = kf::utils::iif(fMask, fTr.GetZ() + step, z);
      while (!kf::utils::isFull(sgn * (zLastLoop - zNext) <= DataT(0.))) {
        ExtrapolateInOneStep(zNext, F);
        zNext = kf::utils::iif(fMask, fTr.GetZ() + step, z);
      }
      ExtrapolateInOneStep(z, F);
    }
  }
 
  template<typename DataT, class Settings>
  inline void TrackKalmanFilter<DataT, Settings>::CleanNonActiveCovariances()
  {
    if constexpr (!Settings::kDoFitTime) {
      for (int j = 0; j < 5; j++) {
        fTr.C(5, j) = DataT(0.);
      }
      for (int j = 0; j < 6; j++) {
        fTr.C(6, j) = DataT(0.);
      }
    }
  }
 
 
}  // namespace cbm::algo::kf