Parameter tuning

Author: anthoak13

AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h

@@ -70,17 +70,29 @@ class TrackFitterUKFBase {
    float32_t m_weightC0;
    /// Unscented transform weight for the other sigma points
    float32_t m_weighti;
-   /// Kappa parameter for sigma point calculation
-   float32_t m_kappa{3 - DIM_A};
+   /// Lambda parameter for sigma point calculation
+   float32_t m_lambda{0};
 
 public:
-   TrackFitterUKFBase() { updateWeights(); }
+   TrackFitterUKFBase() { setParameters(1, 2, 0); }
    ~TrackFitterUKFBase() = default;
 
-   void setKappa(float32_t kappa)
+   /**
+    * @brief Set the weights used to calculate sigma points.
+    */
+   void setParameters(float alpha, float beta, float kappa)
    {
-      m_kappa = kappa; // Set the kappa parameter for sigma point calculation
-      updateWeights(); // Update the weights based on the new kappa value
+      static_assert(DIM_A > 0, "DIM_A is Zero which leads to numerical issue.");
+
+      m_lambda = alpha * alpha * (DIM_A + kappa) - DIM_A;
+      float32_t denoTerm = m_lambda + static_cast<float32_t>(DIM_A);
+
+      m_weightM0 = m_lambda / denoTerm;
+      m_weightC0 = m_weightM0 + (1.0F - alpha * alpha + beta); // Weight for the mean sigma point in covariance
+      m_weighti = 0.5F / denoTerm;
+      LOG(info) << "Mean weight " << m_weightM0;
+      LOG(info) << "Cov weight: " << m_weightC0;
+      LOG(info) << "Shared weight: " << m_weighti;
    }
 
    /**
@@ -200,24 +212,6 @@ class TrackFitterUKFBase {
    }
 
 protected:
-   /**
-    * @brief Set the weights used to calculate sigma points.
-    */
-   void updateWeights()
-   {
-      static_assert(DIM_A > 0, "DIM_A is Zero which leads to numerical issue.");
-
-      constexpr float alpha = 1.0; // Scaling parameter, set to 1 to match orig
-      constexpr float beta = 2.0;  // Optimal for Gaussian distributions
-
-      float lambda{alpha * alpha * (DIM_A + m_kappa) - DIM_A}; // Lambda parameter for sigma points
-      // lambda = m_kappa
-      const float32_t denoTerm{lambda + static_cast<float32_t>(DIM_A)};
-
-      m_weightM0 = lambda / denoTerm;
-      m_weightC0 = m_weightM0 + (1.0F - alpha * alpha + beta); // Weight for the mean sigma point in covariance
-      m_weighti = 0.5F / denoTerm;
-   }
    /**
     * @brief Add state vector and state covariance matrix to the augmented state vector covariance  matrix.
     */
@@ -328,7 +322,7 @@ class TrackFitterUKFBase {
    Matrix<STATE_DIM, SIGMA_DIM>
    calculateSigmaPoints(const Vector<STATE_DIM> &vecXa, const Matrix<STATE_DIM, STATE_DIM> &matPa)
    {
-      const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_kappa)}; // sqrt(n + \kappa)
+      const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_lambda)}; // sqrt(n + \kappa)
 
       Eigen::LLT<Matrix<STATE_DIM, STATE_DIM>> lltOfPa = calculateCholesky<STATE_DIM>(matPa);
 
@@ -447,6 +441,7 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> {
    using EigenVectorDimX = std::vector<Vector<TF_DIM_X>, Eigen::aligned_allocator<Vector<TF_DIM_X>>>;
    using VectorEigenMatDimX =
       std::vector<Matrix<TF_DIM_X, TF_DIM_X>, Eigen::aligned_allocator<Matrix<TF_DIM_X, TF_DIM_X>>>;
+
    // vectors to hold the information needed for smoothing the UKF
    EigenVectorDimX m_vecXPredHist;    /// @brief History of predicted state vectors at k+1
    VectorEigenMatDimX m_matPPredHist; /// @brief History of predicted state covariances at k+1

AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx

@@ -77,6 +77,7 @@ TEST_F(TrackFitterUKFExampleTest, Prediction)
 
    m_ukf.setCovarianceQ(Q);
    m_ukf.setCovarianceR(R);
+   m_ukf.setParameters(1, 0, 3 - m_ukf.DIM_A); // Set kappa to match the original implementation
 
    m_ukf.predictUKF(funcF, z);
 
@@ -138,6 +139,7 @@ TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection)
 
    m_ukf.setCovarianceQ(Q);
    m_ukf.setCovarianceR(R);
+   m_ukf.setParameters(1, 0, 3 - m_ukf.DIM_A); // Set kappa to match the original implementation
 
    m_ukf.predictUKF(funcF, z);
 
@@ -188,7 +190,7 @@ TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection)
    //  [ 0.00651178 -0.0046378   0.13023154 -0.00210188]
    //  [-0.00465059  0.01344241 -0.00210188  0.1333886 ]]
 
-   ASSERT_NEAR(m_ukf.vecX()[0], 2.47414017F, FLOAT_EPSILON);
+   ASSERT_NEAR(m_ukf.vecX()[0], 2.4758845F, FLOAT_EPSILON);
    ASSERT_NEAR(m_ukf.vecX()[1], 0.53327217F, FLOAT_EPSILON);
    ASSERT_NEAR(m_ukf.vecX()[2], 0.21649734F, FLOAT_EPSILON);
    ASSERT_NEAR(m_ukf.vecX()[3], -0.21214576F, FLOAT_EPSILON);

macro/tests/UKF/AtPropagator.C

@@ -33,7 +33,14 @@ void AtPropagator()
    auto elossModel = std::make_unique<AtTools::AtELossTable>(0);
    elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path
    elossModel->SetDensity(3.553e-5);           // Set density in g/cm^3 for 300 torr H2
-   AtTools::AtPropagator propagator(charge, mass, std::move(elossModel));
+
+   auto elossModel2 = std::make_unique<AtTools::AtELossCATIMA>(3.553e-5);
+   elossModel2->SetProjectile(1, 1, 1);
+   std::vector<std::tuple<int, int, int>> mat;
+   mat.push_back({1, 1, 1});
+   elossModel2->SetMaterial(mat);
+
+   AtTools::AtPropagator propagator(charge, mass, std::move(elossModel2));
    propagator.SetEField({0, 0, 0});    // No electric field
    propagator.SetBField({0, 0, 2.85}); // Magnetic field
    AtTools::AtRK4Stepper stepper;

macro/tests/UKF/UKFSingleTrack.C

@@ -88,6 +88,7 @@ void UKFSingleTrack()
    startPos *= 10;                                         // Convert to mm
    XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c
    startMom *= 1e3;
+   double beginMom = startMom.R(); // Initial momentum in MeV/c
 
    XYZPoint nextPos(x[1], y[1], z[1]);
    startMom = startMom.R() * (nextPos - startPos).Unit(); // Set momentum direction towards the first hit
@@ -98,6 +99,7 @@ void UKFSingleTrack()
    double sigma_theta = 1 * M_PI / 180;    // Angular uncertainty of 1 degree
    double sigma_phi = 1 * M_PI / 180;      // Angular uncertainty of 1 degree
    ukf.fEnableEnStraggling = true;         // Enable energy straggling
+   ukf.setParameters(1e-2, 2, 0);          // Set kappa to match the original implementation
 
    TMatrixD cov(6, 6);
    cov.Zero();
@@ -200,6 +202,8 @@ void UKFSingleTrack()
          eLossSmooth.push_back(0); // First point has no previous state to compare
       }
    }
+   LOG(info) << "Initial smoothed momentum: " << smoothedStates[0][3];
+   LOG(info) << "Starting momentum: " << beginMom;
 
    TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data());
    track->SetTitle("Particle Track;X [mm];Y [mm];Z [mm]");