From b1acf6e2c5c969cd7143e65c6a82a6bd69693c23 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 9 Jun 2025 15:17:47 -0400 Subject: [PATCH 01/71] OpenKF for our case is building and predicting Our prediction step matches the results of the OpenKF example and is passing that portion of the test. It is failing the update step, but this might not be a problem as we are not adding the measurement noise to the augmented state vector, as the OpenKF example does. Instead we are treating the measurement noise analytically. This might cause the small differences in the results, but I need to run through the math to confirm that. --- .../OpenKF/kalman_filter/AtTrackFitterUKF.h | 374 ++++++++++++++++++ .../kalman_filter/AtTrackFitterUKFTest.cxx | 206 ++++++++++ AtReconstruction/AtFitter/OpenKF/kf_util.h | 8 +- AtReconstruction/CMakeLists.txt | 2 + 4 files changed, 587 insertions(+), 3 deletions(-) create mode 100644 AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h create mode 100644 AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h new file mode 100644 index 000000000..70ae0d9dd --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h @@ -0,0 +1,374 @@ +/// +/// Copyright 2022 Mohanad Youssef (Al-khwarizmi) +/// +/// Use of this source code is governed by an GPL-3.0 - style +/// license that can be found in the LICENSE file or at +/// https://opensource.org/licenses/GPL-3.0 +/// +/// @author Mohanad Youssef +/// @file unscented_kalman_filter.h +/// + +#ifndef UNSCENTED_KALMAN_FILTER_LIB_H +#define UNSCENTED_KALMAN_FILTER_LIB_H + +#include "kalman_filter.h" +#include "kf_util.h" +#include "unscented_kalman_filter.h" + +namespace kf { + +/// @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. +/// @tparam DIM_X Dimension of the state vector. +/// @tparam DIM_Z Dimension of the measurement vector. +/// @tparam DIM_V Dimension of the process noise vector. +/// @tparam DIM_N Dimension of the measurement noise vector. +template +class TrackFitterUKF : public KalmanFilter { +public: + // Augmented state vector is just the process noise and state vector. The measurement noise is not included as that + // is independent of the propagation and measurement model and just adds linearly. + static constexpr int32_t DIM_A{DIM_X + DIM_V}; ///< @brief Augmented state dimension + static constexpr int32_t SIGMA_DIM_A{2 * DIM_A + 1}; ///< @brief Sigma points dimension for augmented state + float32_t m_kappa{0}; ///< @brief Kappa parameter for finding sigma points + + // Add variables to track the covariances of the process and measurement noise. + Matrix m_matQ; // @brief Process noise covariance matrix + Matrix m_matR; // @brief Measurement noise covariance matrix + + Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix + + TrackFitterUKF() + : KalmanFilter(), m_kappa(3 - DIM_A), m_matQ(Matrix::Zero()), + m_matR(Matrix::Zero()) + { + // 1. calculate weights + updateWeights(); + } + + ~TrackFitterUKF() {} + + void setKappa(float32_t kappa) + { + m_kappa = kappa; // Set the kappa parameter for sigma point calculation + updateWeights(); // Update the weights based on the new kappa value + } + + // This code uses two different conventions for managing noise. + // The state vector noise is set in the updateAugmentedStateAndCovariance() method, while + // the noise vectors for the process and measurement models are set in the setCovarianceQ() and + // setCovarianceR() methods. This is an odd choice. We will be moving everything into a common + // structure where updateAugmentedStateAndCovariance() handle all covariance updates that are actually + // part of the augmented state vector. + + /// + /// @brief adding process noise covariance Q to the augmented state covariance + /// matPa in the middle element of the diagonal. + /// + void setCovarianceQ(const Matrix &matQ) + { + m_matQ = matQ; // Store the process noise covariance matrix + } + + /// + /// @brief adding measurement noise covariance R to the augmented state + /// covariance matPa in the third element of the diagonal. + /// + void setCovarianceR(const Matrix &matR) + { + m_matR = matR; // Store the measurement noise covariance matrix + } + + /// Add state vector (m_vecX) to the augment state vector (m_vecXa) and also + /// add covariance matrix (m_matP) to the augment covariance (m_matPa). + void updateAugWithState() + { + // Copy state vector to augmented state vector + for (int32_t i{0}; i < DIM_X; ++i) { + m_vecXa[i] = m_vecX[i]; + } + + // Copy state covariance matrix to augmented covariance matrix + for (int32_t i{0}; i < DIM_X; ++i) { + for (int32_t j{0}; j < DIM_X; ++j) { + m_matPa(i, j) = m_matP(i, j); + } + } + } + + std::array calculateProcessNoiseMean() + { + // Calculate the expectation value of the process noise using the current value of the state vector m_vecX + std::array processNoiseMean{0}; + + // TODO: Set the mean energy loss based on the momentum and particle type. Probably best to track stopping power? + return processNoiseMean; + } + + Matrix calculateProcessNoiseCovariance() + { + // Calculate the process noise covariance matrix + Matrix matQ{Matrix::Zero()}; + + // TODO: Set the process noise covariance for angular straggle and energy loss. + return matQ; + } + + void updateAugWithProcessNoise() + { + auto processNoiseMean = calculateProcessNoiseMean(); + m_matQ; // = calculateProcessNoiseCovariance(); + + // Add the mean process noise to the augmented state vector + for (int32_t i{0}; i < DIM_V; ++i) { + m_vecXa[DIM_X + i] = processNoiseMean[i]; + } + + // Add process noise covariance to the augmented covariance matrix + const int32_t S_IDX{DIM_X}; + const int32_t L_IDX{S_IDX + DIM_V}; + + for (int32_t i{S_IDX}; i < L_IDX; ++i) { + for (int32_t j{S_IDX}; j < L_IDX; ++j) { + m_matPa(i, j) = m_matQ(i - S_IDX, j - S_IDX); + } + } + } + + /// + /// @brief update the augmented state vector and covariance matrix + /// This functions fully updates the augmented state vector (m_vecXa) and covariance matrix (m_matPa) + /// by setting both the state vector and process noise components. + /// + void updateAugmentedStateAndCovariance() + { + updateAugWithState(); + updateAugWithProcessNoise(); + } + + /// + /// @brief state prediction step of the unscented Kalman filter (UKF). + /// @param predictionModelFunc callback to the prediction/process model + /// function + /// + template + void predictUKF(PredictionModelCallback predictionModelFunc) + { + setKappa(3 - DIM_A); // Set kappa for the augmented state vector and update the weights. + updateAugmentedStateAndCovariance(); + + // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma + // point. + m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); + + // Pull out the sigma points for the state vector and process noise in two different matrices. + Matrix sigmaXx{m_matSigmaXa.block(0, 0, DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector + Matrix sigmaXv{ + m_matSigmaXa.block(DIM_X, 0, DIM_V, SIGMA_DIM_A)}; // Sigma points for process noise + + // Get each sigma point, apply the prediction model function, and store the results + // back into the sigmaXx matrix (safe since each sigma point is independent). + for (int32_t i{0}; i < SIGMA_DIM_A; ++i) { + const Vector sigmaXxi{util::getColumnAt(i, sigmaXx)}; + const Vector sigmaXvi{util::getColumnAt(i, sigmaXv)}; + + const Vector Yi{predictionModelFunc(sigmaXxi, sigmaXvi)}; // y = f(x) + + // Copy the predicted state vector back into the sigmaXx matrix. + util::copyToColumn(i, sigmaXx, Yi); + // Copy the predicted state vector back into the augmented state vector (for use in future functions). + util::copyToColumn(i, m_matSigmaXa, Yi); + } + + // Calculate the weighted mean and covariance of the sigma points for the state vector. + // This will be the new state vector and covariance matrix. + calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP); + } + + /// + /// @brief measurement correction step of the unscented Kalman filter (UKF). + /// @param measurementModelFunc callback to the measurement model function + /// @param vecZ actual measurement vector. + /// + template + void correctUKF(MeasurementModelCallback measurementModelFunc, const Vector &vecZ) + { + // The state vector used here is an unaugmented state vector (m_vecX) and the covariance matrix is + // an unaugmented covariance matrix (m_matP). This is because we are assuming the measurement noise + // is independent of the state vector and process noise, so we can just use the unaugmented state vector + // and covariance matrix for the measurement correction step, then add the measurement noise covariance. + + // Pull out the sigma points for the state vector after prediction. + Matrix sigmaXx{m_matSigmaXa.block(0, 0, DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector + + // Get each sigma point, apply the prediction model function, and store the results + // in the sigmaZ matrix. + Matrix sigmaZ; + for (int32_t i{0}; i < SIGMA_DIM_A; ++i) { + const Vector sigmaXxi{util::getColumnAt(i, sigmaXx)}; + + const Vector Zi{measurementModelFunc(sigmaXxi)}; // z = h(x) + + util::copyToColumn(i, sigmaZ, Zi); + } + + // calculate the mean measurement vector and covariance matrix + // from the sigma points. + Vector vecZhat; + Matrix matPzz; + calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz); + + // Add in the measurement noise covariance matrix to the measurement covariance matrix. + matPzz += m_matR; // Add measurement noise covariance + + // TODO: calculate cross correlation + const Matrix matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)}; + + // kalman gain + const Matrix matK{matPxz * matPzz.inverse()}; + + m_vecX += matK * (vecZ - vecZhat); + m_matP -= matK * matPzz * matK.transpose(); + } + +private: + using KalmanFilter::m_vecX; // from Base KalmanFilter class + using KalmanFilter::m_matP; // from Base KalmanFilter class + + float32_t m_weight0; /// @brief unscented transform weight 0 for mean + float32_t m_weighti; /// @brief unscented transform weight i for none mean samples + + Vector m_vecXa{Vector::Zero()}; /// @brief augmented state vector (incl. process + /// and measurement noise means) + Matrix m_matPa{Matrix::Zero()}; /// @brief augmented state covariance (incl. + /// process and measurement noise covariances) + + /// + /// @brief algorithm to calculate the weights used to draw the sigma points + /// + void updateWeights() + { + static_assert(DIM_A > 0, "DIM_A is Zero which leads to numerical issue."); + + const float32_t denoTerm{m_kappa + static_cast(DIM_A)}; + + m_weight0 = m_kappa / denoTerm; + m_weighti = 0.5F / denoTerm; + } + + /// + /// @brief algorithm to calculate the deterministic sigma points for + /// the unscented transformation + /// + /// @param vecX mean of the normally distributed state + /// @param matPxx covariance of the normally distributed state + /// @param STATE_DIM dimension of the vector used to calculate the sigma points + /// @param SIGMA_DIM number of sigma points required (default is 2 * STATE_DIM + 1) + /// @return matrix of sigma points where each column is a sigma point + /// + template + Matrix + calculateSigmaPoints(const Vector &vecXa, const Matrix &matPa) + { + setKappa(3 - STATE_DIM); // Set kappa for the sigma points calculation + const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_kappa)}; // sqrt(n + \kappa) + + // cholesky factorization to get matrix Pxx square-root + Eigen::LLT> lltOfPa(matPa); + Matrix matSa{lltOfPa.matrixL()}; // sqrt(P_{a}) + + matSa *= scalarMultiplier; // sqrt( (n + \kappa) * P_{a} ) + + Matrix sigmaXa; + + // X_0 = \bar{xa} + util::copyToColumn(0, sigmaXa, vecXa); + + for (int32_t i{0}; i < STATE_DIM; ++i) { + const int32_t IDX_1{i + 1}; + const int32_t IDX_2{i + STATE_DIM + 1}; + + util::copyToColumn(IDX_1, sigmaXa, vecXa); + util::copyToColumn(IDX_2, sigmaXa, vecXa); + + const Vector vecShiftTerm{util::getColumnAt(i, matSa)}; + + util::addColumnFrom(IDX_1, sigmaXa, vecShiftTerm); // X_i^a = \bar{xa} + sqrt( (n^a + + // \kappa) * P^{a} ) + util::subColumnFrom(IDX_2, sigmaXa, vecShiftTerm); // X_{i+n}^a = \bar{xa} - sqrt( (n^a + + // \kappa) * P^{a} ) + } + + return sigmaXa; + } + + /// + /// @brief calculate the weighted mean and covariance given a set of sigma + /// points + /// @param[in] sigmaX matrix of (probably posterior) sigma points where each column contain single + /// sigma point. + /// @param[out] vecX output weighted mean of the sigma points + /// @param[out] matPxx output weighted covariance of the sigma points + /// + template + void calculateWeightedMeanAndCovariance(const Matrix &sigmaX, Vector &vecX, + Matrix &matPxx) + { + // 1. calculate mean of the sigma points + vecX = m_weight0 * util::getColumnAt(0, sigmaX); + for (int32_t i{1}; i < SIGMA_DIM; ++i) { + vecX += m_weighti * util::getColumnAt(i, sigmaX); // y += W[0, i] Y[:, i] + } + + // 2. calculate covariance: P_{yy} = \sum_{i_0}^{2n} W[0, i] (Y[:, i] - + // \bar{y}) (Y[:, i] - \bar{y})^T + Vector devXi{util::getColumnAt(0, sigmaX) - vecX}; // Y[:, 0] - \bar{ y } + + matPxx = m_weight0 * devXi * devXi.transpose(); // P_0 = W[0, 0] (Y[:, 0] - \bar{y}) (Y[:, 0] - + // \bar{y})^T + + for (int32_t i{1}; i < SIGMA_DIM; ++i) { + devXi = util::getColumnAt(i, sigmaX) - vecX; // Y[:, i] - \bar{y} + + const Matrix Pi{m_weighti * devXi * devXi.transpose()}; // P_i = W[0, i] (Y[:, i] - + // \bar{y}) (Y[:, i] - \bar{y})^T + + matPxx += Pi; // y += W[0, i] (Y[:, i] - \bar{y}) (Y[:, i] - \bar{y})^T + } + } + + /// + /// @brief calculate the cross-correlation given two sets sigma points X and Y + /// and their means x and y + /// @param sigmaX first matrix of sigma points where each column contain + /// single sigma point + /// @param vecX mean of the first set of sigma points + /// @param sigmaY second matrix of sigma points where each column contain + /// single sigma point + /// @param vecY mean of the second set of sigma points + /// @return matPxy, the cross-correlation matrix + /// + template + Matrix calculateCrossCorrelation(const Matrix &sigmaX, const Vector &vecX, + const Matrix &sigmaY, const Vector &vecY) + { + Vector devXi{util::getColumnAt(0, sigmaX) - vecX}; // X[:, 0] - \bar{ x } + Vector devYi{util::getColumnAt(0, sigmaY) - vecY}; // Y[:, 0] - \bar{ y } + + // P_0 = W[0, 0] (X[:, 0] - \bar{x}) (Y[:, 0] - \bar{y})^T + Matrix matPxy{m_weight0 * (devXi * devYi.transpose())}; + + for (int32_t i{1}; i < SIGMA_DIM; ++i) { + devXi = util::getColumnAt(i, sigmaX) - vecX; // X[:, i] - \bar{x} + devYi = util::getColumnAt(i, sigmaY) - vecY; // Y[:, i] - \bar{y} + + matPxy += m_weighti * (devXi * devYi.transpose()); // y += W[0, i] (Y[:, i] - + // \bar{y}) (Y[:, i] - \bar{y})^T + } + + return matPxy; + } +}; +} // namespace kf + +#endif // UNSCENTED_KALMAN_FILTER_LIB_H diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx new file mode 100644 index 000000000..f4e018c16 --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx @@ -0,0 +1,206 @@ +#include "kalman_filter/AtTrackFitterUKF.h" + +#include "gtest/gtest.h" + +class AtTrackFitterUKFTest : public testing::Test { +public: + virtual void SetUp() override {} + virtual void TearDown() override {} + + static constexpr float FLOAT_EPSILON{0.001F}; + + static constexpr size_t DIM_X{4}; + static constexpr size_t DIM_V{4}; + static constexpr size_t DIM_Z{2}; + static constexpr size_t DIM_N{2}; + + kf::TrackFitterUKF m_ukf; + + /// @brief to propagate the state vector using the process model + /// @param x state vector + /// @param v process noise vector + /// @return propagated (unaugmented) state vector + static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v) + { + kf::Vector y; + y[0] = x[0] + x[2] + v[0]; + y[1] = x[1] + x[3] + v[1]; + y[2] = x[2] + v[2]; + y[3] = x[3] + v[3]; + return y; + } + + /// @brief to apply the measurement model to the state vector + /// @param x the state vector of the system + /// @return the measurement vector + static kf::Vector funcH(const kf::Vector &x) + { + kf::Vector y; + + kf::float32_t px{x[0]}; + kf::float32_t py{x[1]}; + + y[0] = std::sqrt((px * px) + (py * py)); + y[1] = std::atan(py / (px + std::numeric_limits::epsilon())); + return y; + } +}; + +TEST_F(AtTrackFitterUKFTest, test_UKF_Prediction) +{ + kf::Vector x; + x << 2.0F, 1.0F, 0.0F, 0.0F; + + kf::Matrix P; + P << 0.01F, 0.0F, 0.0F, 0.0F, 0.0F, 0.01F, 0.0F, 0.0F, 0.0F, 0.0F, 0.05F, 0.0F, 0.0F, 0.0F, 0.0F, 0.05F; + + kf::Matrix Q; + Q << 0.05F, 0.0F, 0.0F, 0.0F, 0.0F, 0.05F, 0.0F, 0.0F, 0.0F, 0.0F, 0.1F, 0.0F, 0.0F, 0.0F, 0.0F, 0.1F; + + kf::Matrix R; + R << 0.01F, 0.0F, 0.0F, 0.01F; + + kf::Vector z; + z << 2.5F, 0.05F; + + m_ukf.vecX() = x; + m_ukf.matP() = P; + + m_ukf.setCovarianceQ(Q); + m_ukf.setCovarianceR(R); + + m_ukf.predictUKF(funcF); + + // Expectation from the python results: + // ===================================== + // x = + // [2.0 1.0 0.0 0.0] + // P = + // [[0.11 0.00 0.05 0.00] + // [0.00 0.11 0.00 0.05] + // [0.05 0.00 0.15 0.00] + // [0.00 0.05 0.00 0.15]] + + ASSERT_NEAR(m_ukf.vecX()[0], 2.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[1], 1.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[2], 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[3], 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(0, 0), 0.11F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 1), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 2), 0.05F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 3), 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(1, 0), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 1), 0.11F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 2), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 3), 0.05F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(2, 0), 0.05F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 1), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 2), 0.15F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 3), 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(3, 0), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 1), 0.05F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 2), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 3), 0.15F, FLOAT_EPSILON); +} + +TEST_F(AtTrackFitterUKFTest, test_UKF_PredictionAndCorrection) +{ + kf::Vector x; + x << 2.0F, 1.0F, 0.0F, 0.0F; + + kf::Matrix P; + P << 0.01F, 0.0F, 0.0F, 0.0F, 0.0F, 0.01F, 0.0F, 0.0F, 0.0F, 0.0F, 0.05F, 0.0F, 0.0F, 0.0F, 0.0F, 0.05F; + + kf::Matrix Q; + Q << 0.05F, 0.0F, 0.0F, 0.0F, 0.0F, 0.05F, 0.0F, 0.0F, 0.0F, 0.0F, 0.1F, 0.0F, 0.0F, 0.0F, 0.0F, 0.1F; + + kf::Matrix R; + R << 0.01F, 0.0F, 0.0F, 0.01F; + + kf::Vector z; + z << 2.5F, 0.05F; + + m_ukf.vecX() = x; + m_ukf.matP() = P; + + m_ukf.setCovarianceQ(Q); + m_ukf.setCovarianceR(R); + + m_ukf.predictUKF(funcF); + + // Expectation from the python results: + // ===================================== + // x = + // [2.0 1.0 0.0 0.0] + // P = + // [[0.11 0.00 0.05 0.00] + // [0.00 0.11 0.00 0.05] + // [0.05 0.00 0.15 0.00] + // [0.00 0.05 0.00 0.15]] + + ASSERT_NEAR(m_ukf.vecX()[0], 2.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[1], 1.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[2], 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[3], 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(0, 0), 0.11F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 1), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 2), 0.05F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 3), 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(1, 0), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 1), 0.11F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 2), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 3), 0.05F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(2, 0), 0.05F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 1), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 2), 0.15F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 3), 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(3, 0), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 1), 0.05F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 2), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 3), 0.15F, FLOAT_EPSILON); + + m_ukf.correctUKF(funcH, z); + + // Expectations from the python results: + // ====================================== + // x = + // [ 2.554 0.356 0.252 -0.293] + // P = + // [[ 0.01 -0.001 0.005 -0. ] + // [-0.001 0.01 - 0. 0.005 ] + // [ 0.005 - 0. 0.129 - 0. ] + // [-0. 0.005 - 0. 0.129]] + + ASSERT_NEAR(m_ukf.vecX()[0], 2.554F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[1], 0.356F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[2], 0.252F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[3], -0.293F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(0, 0), 0.01F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 1), -0.001F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 2), 0.005F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 3), 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(1, 0), -0.001F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 1), 0.01F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 2), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 3), 0.005F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(2, 0), 0.005F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 1), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 2), 0.129F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 3), 0.0F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(3, 0), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 1), 0.005F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 2), 0.0F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 3), 0.129F, FLOAT_EPSILON); +} diff --git a/AtReconstruction/AtFitter/OpenKF/kf_util.h b/AtReconstruction/AtFitter/OpenKF/kf_util.h index 4a01d3c19..311ed25eb 100644 --- a/AtReconstruction/AtFitter/OpenKF/kf_util.h +++ b/AtReconstruction/AtFitter/OpenKF/kf_util.h @@ -18,10 +18,12 @@ namespace kf { namespace util { -template -void copyToColumn(const int32_t colIdx, Matrix &lhsSigmaX, const Vector &rhsVecX) +template +void copyToColumn(const int32_t colIdx, Matrix &lhsSigmaX, const Vector &rhsVecX) { - for (int32_t i{0}; i < ROWS; ++i) { // rows + assert(colIdx < COLS); // assert if colIdx is out of boundary + assert(ROWS_COPY <= ROWS); // assert if rhsVecX is larger than lhsSigmaX + for (int32_t i{0}; i < ROWS_COPY; ++i) { // rows lhsSigmaX(i, colIdx) = rhsVecX[i]; } } diff --git a/AtReconstruction/CMakeLists.txt b/AtReconstruction/CMakeLists.txt index 3b5eac823..bb3024824 100755 --- a/AtReconstruction/CMakeLists.txt +++ b/AtReconstruction/CMakeLists.txt @@ -120,6 +120,7 @@ if(TARGET Eigen3::Eigen) set(SRCS ${SRCS} AtPatternRecognition/triplclust/src/orthogonallsq.cxx AtFitter/AtOpenKFTest.cxx + ) message(STATUS "Current sources: ${SRCS}") @@ -138,6 +139,7 @@ if(TARGET Eigen3::Eigen) set(TEST_SRCS_OPENKF AtFitter/OpenKF/kalman_filter/kalman_filter_test.cxx AtFitter/OpenKF/kalman_filter/unscented_kalman_filter_test.cxx + AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx ) attpcroot_generate_tests(OpenKFTests From c507853decd72bc9206d579f4b7f8999c596b782 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 9 Jun 2025 16:33:09 -0400 Subject: [PATCH 02/71] Implement update portion of UKF test --- .../OpenKF/kalman_filter/AtTrackFitterUKF.h | 26 +++++++ .../kalman_filter/AtTrackFitterUKFTest.cxx | 69 +++++++++++-------- 2 files changed, 65 insertions(+), 30 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h index 70ae0d9dd..010518a3d 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h @@ -16,6 +16,8 @@ #include "kf_util.h" #include "unscented_kalman_filter.h" +#include + namespace kf { /// @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. @@ -212,20 +214,44 @@ class TrackFitterUKF : public KalmanFilter { util::copyToColumn(i, sigmaZ, Zi); } + std::cout << "Sigma Points for Measurement:" << std::endl; + // Print sigmaZ in CSV format + for (int32_t row = 0; row < DIM_Z; ++row) { + for (int32_t col = 0; col < SIGMA_DIM_A; ++col) { + std::cout << sigmaZ(row, col); + if (col < SIGMA_DIM_A - 1) + std::cout << ","; + } + std::cout << std::endl; + } + // calculate the mean measurement vector and covariance matrix // from the sigma points. Vector vecZhat; Matrix matPzz; calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz); + std::cout << "Mean Measurement Vector (vecZhat):" << std::endl; + std::cout << vecZhat.transpose() << std::endl; + std::cout << "Measurement Covariance Matrix (matPzz):" << std::endl; + std::cout << matPzz << std::endl; + // Add in the measurement noise covariance matrix to the measurement covariance matrix. matPzz += m_matR; // Add measurement noise covariance + std::cout << "S Matrix (matPzz):" << std::endl; + std::cout << matPzz << std::endl; + // TODO: calculate cross correlation const Matrix matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)}; + std::cout << "Cross Correlation Matrix (matPxz):" << std::endl; + std::cout << matPxz << std::endl; + // kalman gain const Matrix matK{matPxz * matPzz.inverse()}; + std::cout << "Kalman Gain (matK):" << std::endl; + std::cout << matK << std::endl; m_vecX += matK * (vecZ - vecZhat); m_matP -= matK * matPzz * matK.transpose(); diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx index f4e018c16..9a29320a2 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx @@ -167,40 +167,49 @@ TEST_F(AtTrackFitterUKFTest, test_UKF_PredictionAndCorrection) ASSERT_NEAR(m_ukf.matP()(3, 2), 0.0F, FLOAT_EPSILON); ASSERT_NEAR(m_ukf.matP()(3, 3), 0.15F, FLOAT_EPSILON); + std::cout << "Sigma points after prediction:\n"; + for (int32_t j{0}; j < kf::TrackFitterUKF::DIM_A; ++j) { + for (int32_t i{0}; i < kf::TrackFitterUKF::DIM_A * 2 + 1; ++i) { + + std::cout << m_ukf.m_matSigmaXa(j, i) << ","; + } + std::cout << std::endl; + } + m_ukf.correctUKF(funcH, z); // Expectations from the python results: // ====================================== // x = - // [ 2.554 0.356 0.252 -0.293] + // [ 2.4758845 0.53327217 0.21649734 -0.21214576] // P = - // [[ 0.01 -0.001 0.005 -0. ] - // [-0.001 0.01 - 0. 0.005 ] - // [ 0.005 - 0. 0.129 - 0. ] - // [-0. 0.005 - 0. 0.129]] - - ASSERT_NEAR(m_ukf.vecX()[0], 2.554F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.vecX()[1], 0.356F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.vecX()[2], 0.252F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.vecX()[3], -0.293F, FLOAT_EPSILON); - - ASSERT_NEAR(m_ukf.matP()(0, 0), 0.01F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(0, 1), -0.001F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(0, 2), 0.005F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(0, 3), 0.0F, FLOAT_EPSILON); - - ASSERT_NEAR(m_ukf.matP()(1, 0), -0.001F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(1, 1), 0.01F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(1, 2), 0.0F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(1, 3), 0.005F, FLOAT_EPSILON); - - ASSERT_NEAR(m_ukf.matP()(2, 0), 0.005F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(2, 1), 0.0F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(2, 2), 0.129F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(2, 3), 0.0F, FLOAT_EPSILON); - - ASSERT_NEAR(m_ukf.matP()(3, 0), 0.0F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(3, 1), 0.005F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(3, 2), 0.0F, FLOAT_EPSILON); - ASSERT_NEAR(m_ukf.matP()(3, 3), 0.129F, FLOAT_EPSILON); + // [[ 0.01433114 -0.01026142 0.00651178 -0.00465059] + // [-0.01026142 0.0295458 -0.0046378 0.01344241] + // [ 0.00651178 -0.0046378 0.13023154 -0.00210188] + // [-0.00465059 0.01344241 -0.00210188 0.1333886 ]] + + 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); + + ASSERT_NEAR(m_ukf.matP()(0, 0), 0.01433114F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 1), -0.01026142F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 2), 0.00651178F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(0, 3), -0.00465059F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(1, 0), -0.01026142F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 1), 0.0295458F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 2), -0.0046378F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(1, 3), 0.01344241F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(2, 0), 0.00651178F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 1), -0.0046378F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 2), 0.13023154F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(2, 3), -0.00210188F, FLOAT_EPSILON); + + ASSERT_NEAR(m_ukf.matP()(3, 0), -0.00465059F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 1), 0.01344241F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 2), -0.00210188F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.matP()(3, 3), 0.1333886F, FLOAT_EPSILON); } From f0f9ee273f0d84c109f0f89157fcc341d9bfbba3 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 9 Jun 2025 16:43:02 -0400 Subject: [PATCH 03/71] Remove debug statements --- .../OpenKF/kalman_filter/AtTrackFitterUKF.h | 30 ++----------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h index 010518a3d..9b3a8c0ae 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h @@ -16,8 +16,6 @@ #include "kf_util.h" #include "unscented_kalman_filter.h" -#include - namespace kf { /// @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. @@ -111,6 +109,7 @@ class TrackFitterUKF : public KalmanFilter { { // Calculate the process noise covariance matrix Matrix matQ{Matrix::Zero()}; + matQ = m_matQ; // Use the stored process noise covariance matrix // TODO: Set the process noise covariance for angular straggle and energy loss. return matQ; @@ -119,7 +118,7 @@ class TrackFitterUKF : public KalmanFilter { void updateAugWithProcessNoise() { auto processNoiseMean = calculateProcessNoiseMean(); - m_matQ; // = calculateProcessNoiseCovariance(); + m_matQ = calculateProcessNoiseCovariance(); // Add the mean process noise to the augmented state vector for (int32_t i{0}; i < DIM_V; ++i) { @@ -156,7 +155,6 @@ class TrackFitterUKF : public KalmanFilter { template void predictUKF(PredictionModelCallback predictionModelFunc) { - setKappa(3 - DIM_A); // Set kappa for the augmented state vector and update the weights. updateAugmentedStateAndCovariance(); // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma @@ -214,44 +212,20 @@ class TrackFitterUKF : public KalmanFilter { util::copyToColumn(i, sigmaZ, Zi); } - std::cout << "Sigma Points for Measurement:" << std::endl; - // Print sigmaZ in CSV format - for (int32_t row = 0; row < DIM_Z; ++row) { - for (int32_t col = 0; col < SIGMA_DIM_A; ++col) { - std::cout << sigmaZ(row, col); - if (col < SIGMA_DIM_A - 1) - std::cout << ","; - } - std::cout << std::endl; - } - // calculate the mean measurement vector and covariance matrix // from the sigma points. Vector vecZhat; Matrix matPzz; calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz); - std::cout << "Mean Measurement Vector (vecZhat):" << std::endl; - std::cout << vecZhat.transpose() << std::endl; - std::cout << "Measurement Covariance Matrix (matPzz):" << std::endl; - std::cout << matPzz << std::endl; - // Add in the measurement noise covariance matrix to the measurement covariance matrix. matPzz += m_matR; // Add measurement noise covariance - std::cout << "S Matrix (matPzz):" << std::endl; - std::cout << matPzz << std::endl; - // TODO: calculate cross correlation const Matrix matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)}; - std::cout << "Cross Correlation Matrix (matPxz):" << std::endl; - std::cout << matPxz << std::endl; - // kalman gain const Matrix matK{matPxz * matPzz.inverse()}; - std::cout << "Kalman Gain (matK):" << std::endl; - std::cout << matK << std::endl; m_vecX += matK * (vecZ - vecZhat); m_matP -= matK * matPzz * matK.transpose(); From 090951930cbb3de3ff27c6d2deea80ee9a9b5226 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 9 Jun 2025 17:00:14 -0400 Subject: [PATCH 04/71] Add check for positive definite cov matrix Removed some unnecessary updates to kappa Fixed missleading comment --- .../AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h index 9b3a8c0ae..a29c1a51b 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h @@ -71,8 +71,7 @@ class TrackFitterUKF : public KalmanFilter { } /// - /// @brief adding measurement noise covariance R to the augmented state - /// covariance matPa in the third element of the diagonal. + /// @brief set the measurement noise covariance R to be used in the update step /// void setCovarianceR(const Matrix &matR) { @@ -270,11 +269,13 @@ class TrackFitterUKF : public KalmanFilter { Matrix calculateSigmaPoints(const Vector &vecXa, const Matrix &matPa) { - setKappa(3 - STATE_DIM); // Set kappa for the sigma points calculation const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_kappa)}; // sqrt(n + \kappa) // cholesky factorization to get matrix Pxx square-root Eigen::LLT> lltOfPa(matPa); + if (lltOfPa.info() != Eigen::Success) { + throw std::runtime_error("Cholesky decomposition failed, matrix is not positive definite."); + } Matrix matSa{lltOfPa.matrixL()}; // sqrt(P_{a}) matSa *= scalarMultiplier; // sqrt( (n + \kappa) * P_{a} ) From 87587b2f927711d630678372489455b490ba8059 Mon Sep 17 00:00:00 2001 From: anthoak13 Date: Mon, 9 Jun 2025 20:58:05 -0400 Subject: [PATCH 05/71] Add physics test cases to be filled --- .../{AtTrackFitterUKF.h => TrackFitterUKF.h} | 0 ...tterUKFTest.cxx => TrackFitterUKFTest.cxx} | 81 ++++++++++++++++--- AtReconstruction/CMakeLists.txt | 2 +- 3 files changed, 69 insertions(+), 14 deletions(-) rename AtReconstruction/AtFitter/OpenKF/kalman_filter/{AtTrackFitterUKF.h => TrackFitterUKF.h} (100%) rename AtReconstruction/AtFitter/OpenKF/kalman_filter/{AtTrackFitterUKFTest.cxx => TrackFitterUKFTest.cxx} (77%) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h similarity index 100% rename from AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKF.h rename to AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx similarity index 77% rename from AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx rename to AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 9a29320a2..5767d80e0 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -1,8 +1,8 @@ -#include "kalman_filter/AtTrackFitterUKF.h" +#include "kalman_filter/TrackFitterUKF.h" #include "gtest/gtest.h" -class AtTrackFitterUKFTest : public testing::Test { +class TrackFitterUKFExampleTest : public testing::Test { public: virtual void SetUp() override {} virtual void TearDown() override {} @@ -46,7 +46,7 @@ class AtTrackFitterUKFTest : public testing::Test { } }; -TEST_F(AtTrackFitterUKFTest, test_UKF_Prediction) +TEST_F(TrackFitterUKFExampleTest, Prediction) { kf::Vector x; x << 2.0F, 1.0F, 0.0F, 0.0F; @@ -107,7 +107,7 @@ TEST_F(AtTrackFitterUKFTest, test_UKF_Prediction) ASSERT_NEAR(m_ukf.matP()(3, 3), 0.15F, FLOAT_EPSILON); } -TEST_F(AtTrackFitterUKFTest, test_UKF_PredictionAndCorrection) +TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection) { kf::Vector x; x << 2.0F, 1.0F, 0.0F, 0.0F; @@ -167,15 +167,6 @@ TEST_F(AtTrackFitterUKFTest, test_UKF_PredictionAndCorrection) ASSERT_NEAR(m_ukf.matP()(3, 2), 0.0F, FLOAT_EPSILON); ASSERT_NEAR(m_ukf.matP()(3, 3), 0.15F, FLOAT_EPSILON); - std::cout << "Sigma points after prediction:\n"; - for (int32_t j{0}; j < kf::TrackFitterUKF::DIM_A; ++j) { - for (int32_t i{0}; i < kf::TrackFitterUKF::DIM_A * 2 + 1; ++i) { - - std::cout << m_ukf.m_matSigmaXa(j, i) << ","; - } - std::cout << std::endl; - } - m_ukf.correctUKF(funcH, z); // Expectations from the python results: @@ -213,3 +204,67 @@ TEST_F(AtTrackFitterUKFTest, test_UKF_PredictionAndCorrection) ASSERT_NEAR(m_ukf.matP()(3, 2), -0.00210188F, FLOAT_EPSILON); ASSERT_NEAR(m_ukf.matP()(3, 3), 0.1333886F, FLOAT_EPSILON); } + +class TrackFitterUKFPhysicsTest : public testing::Test { +public: + virtual void SetUp() override {} + virtual void TearDown() override {} + + static constexpr float FLOAT_EPSILON{0.001F}; + + static constexpr size_t DIM_X{6}; + static constexpr size_t DIM_V{2}; + static constexpr size_t DIM_Z{3}; + static constexpr size_t DIM_N{3}; + + kf::TrackFitterUKF m_ukf; + + /// @brief to propagate the state vector using the process model + /// @param x state vector + /// @param v process noise vector + /// @return propagated (unaugmented) state vector + static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v) + { + //TODO: This needs to be filled with an RK4 solver for the physics model + kf::Vector y{x}; + + // For now, we just return the state vector as is + return y; + } + + /// @brief to apply the measurement model to the state vector + /// @param x the state vector of the system + /// @return the measurement vector + static kf::Vector funcH(const kf::Vector &x) + { + kf::Vector y; + y[0] = x[0]; + y[1] = x[1]; + y[2] = x[2]; + + return y; + } +}; + +TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) +{ + // TODO: This needs to be filled with a proper physics test. + kf::Vector x; // Initial state vector + + kf::Matrix P; // Initial state vector covariance matrix + + // Note: process noise is defined in the UKF class, so we don't need to set it here. + + kf::Vector z; // Measurement vector to be used in the correction step + z << std::sqrt(5), std::atan2(1.f ,2.f); + + kf::Matrix R; // Covariance matrix for the measurement noise + + m_ukf.vecX() = x; + m_ukf.matP() = P; + + m_ukf.setCovarianceR(R); + m_ukf.predictUKF(funcF); + + +} diff --git a/AtReconstruction/CMakeLists.txt b/AtReconstruction/CMakeLists.txt index bb3024824..bc0de8a62 100755 --- a/AtReconstruction/CMakeLists.txt +++ b/AtReconstruction/CMakeLists.txt @@ -139,7 +139,7 @@ if(TARGET Eigen3::Eigen) set(TEST_SRCS_OPENKF AtFitter/OpenKF/kalman_filter/kalman_filter_test.cxx AtFitter/OpenKF/kalman_filter/unscented_kalman_filter_test.cxx - AtFitter/OpenKF/kalman_filter/AtTrackFitterUKFTest.cxx + AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx ) attpcroot_generate_tests(OpenKFTests From 6b8698bef83092a1a9fcfa0def11970d1d1edcb9 Mon Sep 17 00:00:00 2001 From: anthoak13 Date: Mon, 9 Jun 2025 21:00:40 -0400 Subject: [PATCH 06/71] Run clang format --- .../OpenKF/kalman_filter/TrackFitterUKFTest.cxx | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 5767d80e0..1751ea443 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -225,9 +225,9 @@ class TrackFitterUKFPhysicsTest : public testing::Test { /// @return propagated (unaugmented) state vector static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v) { - //TODO: This needs to be filled with an RK4 solver for the physics model + // TODO: This needs to be filled with an RK4 solver for the physics model kf::Vector y{x}; - + // For now, we just return the state vector as is return y; } @@ -248,7 +248,7 @@ class TrackFitterUKFPhysicsTest : public testing::Test { TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) { - // TODO: This needs to be filled with a proper physics test. + // TODO: This needs to be filled with a proper physics test. kf::Vector x; // Initial state vector kf::Matrix P; // Initial state vector covariance matrix @@ -256,7 +256,7 @@ TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) // Note: process noise is defined in the UKF class, so we don't need to set it here. kf::Vector z; // Measurement vector to be used in the correction step - z << std::sqrt(5), std::atan2(1.f ,2.f); + z << std::sqrt(5), std::atan2(1.f, 2.f); kf::Matrix R; // Covariance matrix for the measurement noise @@ -265,6 +265,4 @@ TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) m_ukf.setCovarianceR(R); m_ukf.predictUKF(funcF); - - } From 9efdcfcdcb3e6936c1da551157ce6226c8f52bff Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 17 Jul 2025 15:55:37 +0200 Subject: [PATCH 07/71] Force is working and propogator with no fields is working --- .../AtFitter/OpenKF/kalman_filter/HinH.txt | 143 +++++++++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 4 +- .../kalman_filter/TrackFitterUKFTest.cxx | 297 +++++++++++++++++- AtTools/AtELossTable.h | 2 + 4 files changed, 437 insertions(+), 9 deletions(-) create mode 100644 AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt b/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt new file mode 100644 index 000000000..b436b8c8c --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt @@ -0,0 +1,143 @@ + ================================================================== + SRIM version ---> SRIM-2013.00 + Calc. date ---> julio 16. 2025 + ================================================================== + + Disk File Name = SRIM Outputs\Hydrogen in Hydrogen (gas).txt + + Ion = Hydrogen [1] . Mass = 1.008 amu + + Target Density = 4.1906E-05 g/cm3 = 2.5036E+19 atoms/cm3 + Target is a GAS + ======= Target Composition ======== + Atom Atom Atomic Mass + Name Numb Percent Percent + ---- ---- ------- ------- + H 1 100.00 100.00 + ==================================== + Bragg Correction = 0.00% + Stopping Units = MeV / (mg/cm2) + See bottom of Table for other Stopping units + + Ion dE/dx dE/dx Projected Longitudinal Lateral + Energy Elec. Nuclear Range Straggling Straggling + -------------- ---------- ---------- ---------- ---------- ---------- +999.999 eV 7.504E-01 1.692E-01 267.71 um 61.83 um 55.80 um + 1.10 keV 7.871E-01 1.600E-01 291.11 um 64.85 um 59.45 um + 1.20 keV 8.221E-01 1.519E-01 313.96 um 67.62 um 62.90 um + 1.30 keV 8.556E-01 1.447E-01 336.31 um 70.18 um 66.16 um + 1.40 keV 8.879E-01 1.383E-01 358.18 um 72.56 um 69.25 um + 1.50 keV 9.191E-01 1.325E-01 379.59 um 74.77 um 72.19 um + 1.60 keV 9.492E-01 1.272E-01 400.57 um 76.83 um 74.99 um + 1.70 keV 9.785E-01 1.224E-01 421.14 um 78.76 um 77.65 um + 1.80 keV 1.007E+00 1.180E-01 441.32 um 80.57 um 80.20 um + 2.00 keV 1.061E+00 1.101E-01 480.60 um 83.97 um 84.97 um + 2.25 keV 1.124E+00 1.019E-01 527.87 um 87.75 um 90.40 um + 2.50 keV 1.183E+00 9.491E-02 573.34 um 91.09 um 95.34 um + 2.75 keV 1.239E+00 8.894E-02 617.19 um 94.06 um 99.86 um + 3.00 keV 1.293E+00 8.377E-02 659.59 um 96.72 um 104.03 um + 3.25 keV 1.344E+00 7.923E-02 700.68 um 99.13 um 107.88 um + 3.50 keV 1.393E+00 7.521E-02 740.56 um 101.32 um 111.46 um + 3.75 keV 1.440E+00 7.162E-02 779.35 um 103.33 um 114.80 um + 4.00 keV 1.486E+00 6.840E-02 817.12 um 105.18 um 117.93 um + 4.50 keV 1.572E+00 6.285E-02 889.92 um 108.70 um 123.65 um + 5.00 keV 1.654E+00 5.821E-02 959.49 um 111.75 um 128.76 um + 5.50 keV 1.731E+00 5.428E-02 1.03 mm 114.44 um 133.36 um + 6.00 keV 1.804E+00 5.090E-02 1.09 mm 116.83 um 137.55 um + 6.50 keV 1.873E+00 4.795E-02 1.15 mm 118.97 um 141.38 um + 7.00 keV 1.939E+00 4.536E-02 1.21 mm 120.92 um 144.91 um + 8.00 keV 2.064E+00 4.101E-02 1.33 mm 124.82 um 151.22 um + 9.00 keV 2.178E+00 3.748E-02 1.44 mm 128.13 um 156.72 um + 10.00 keV 2.285E+00 3.457E-02 1.54 mm 131.01 um 161.59 um + 11.00 keV 2.384E+00 3.211E-02 1.64 mm 133.54 um 165.95 um + 12.00 keV 2.477E+00 3.001E-02 1.74 mm 135.79 um 169.89 um + 13.00 keV 2.564E+00 2.819E-02 1.83 mm 137.82 um 173.48 um + 14.00 keV 2.646E+00 2.660E-02 1.92 mm 139.67 um 176.77 um + 15.00 keV 2.723E+00 2.519E-02 2.01 mm 141.36 um 179.82 um + 16.00 keV 2.796E+00 2.393E-02 2.09 mm 142.92 um 182.65 um + 17.00 keV 2.864E+00 2.281E-02 2.17 mm 144.37 um 185.29 um + 18.00 keV 2.929E+00 2.179E-02 2.26 mm 145.72 um 187.77 um + 20.00 keV 3.048E+00 2.003E-02 2.41 mm 149.01 um 192.31 um + 22.50 keV 3.179E+00 1.822E-02 2.60 mm 153.11 um 197.34 um + 25.00 keV 3.292E+00 1.673E-02 2.79 mm 156.75 um 201.82 um + 27.50 keV 3.389E+00 1.549E-02 2.96 mm 160.04 um 205.87 um + 30.00 keV 3.473E+00 1.443E-02 3.14 mm 163.06 um 209.56 um + 32.50 keV 3.544E+00 1.351E-02 3.30 mm 165.87 um 212.98 um + 35.00 keV 3.604E+00 1.272E-02 3.47 mm 168.50 um 216.15 um + 37.50 keV 3.654E+00 1.201E-02 3.63 mm 170.98 um 219.13 um + 40.00 keV 3.695E+00 1.139E-02 3.79 mm 173.35 um 221.95 um + 45.00 keV 3.753E+00 1.033E-02 4.11 mm 180.64 um 227.18 um + 50.00 keV 3.784E+00 9.466E-03 4.43 mm 187.42 um 231.99 um + 55.00 keV 3.794E+00 8.743E-03 4.74 mm 193.85 um 236.49 um + 60.00 keV 3.785E+00 8.129E-03 5.06 mm 200.03 um 240.76 um + 65.00 keV 3.763E+00 7.601E-03 5.37 mm 206.04 um 244.84 um + 70.00 keV 3.730E+00 7.142E-03 5.69 mm 211.93 um 248.78 um + 80.00 keV 3.638E+00 6.382E-03 6.33 mm 232.69 um 256.39 um + 90.00 keV 3.526E+00 5.776E-03 7.00 mm 252.74 um 263.79 um + 100.00 keV 3.403E+00 5.282E-03 7.69 mm 272.49 um 271.10 um + 110.00 keV 3.276E+00 4.870E-03 8.40 mm 292.22 um 278.42 um + 120.00 keV 3.150E+00 4.522E-03 9.14 mm 312.12 um 285.84 um + 130.00 keV 3.026E+00 4.222E-03 9.91 mm 332.30 um 293.41 um + 140.00 keV 2.907E+00 3.963E-03 10.72 mm 352.84 um 301.18 um + 150.00 keV 2.794E+00 3.735E-03 11.55 mm 373.81 um 309.18 um + 160.00 keV 2.687E+00 3.533E-03 12.42 mm 395.24 um 317.44 um + 170.00 keV 2.587E+00 3.353E-03 13.33 mm 417.16 um 326.00 um + 180.00 keV 2.492E+00 3.192E-03 14.26 mm 439.58 um 334.87 um + 200.00 keV 2.320E+00 2.915E-03 16.25 mm 524.78 um 353.64 um + 225.00 keV 2.134E+00 2.632E-03 18.92 mm 651.26 um 379.16 um + 250.00 keV 1.975E+00 2.402E-03 21.83 mm 773.93 um 407.12 um + 275.00 keV 1.840E+00 2.211E-03 24.95 mm 895.52 um 437.62 um + 300.00 keV 1.722E+00 2.050E-03 28.30 mm 1.02 mm 470.70 um + 325.00 keV 1.620E+00 1.911E-03 31.87 mm 1.14 mm 506.38 um + 350.00 keV 1.530E+00 1.791E-03 35.65 mm 1.26 mm 544.65 um + 375.00 keV 1.451E+00 1.686E-03 39.65 mm 1.39 mm 585.49 um + 400.00 keV 1.381E+00 1.594E-03 43.86 mm 1.52 mm 628.85 um + 450.00 keV 1.260E+00 1.437E-03 52.90 mm 1.99 mm 723.02 um + 500.00 keV 1.162E+00 1.310E-03 62.75 mm 2.44 mm 826.75 um + 550.00 keV 1.079E+00 1.205E-03 73.40 mm 2.88 mm 939.68 um + 600.00 keV 1.008E+00 1.116E-03 84.83 mm 3.31 mm 1.06 mm + 650.00 keV 9.477E-01 1.039E-03 97.02 mm 3.75 mm 1.19 mm + 700.00 keV 8.947E-01 9.735E-04 109.96 mm 4.18 mm 1.33 mm + 800.00 keV 8.064E-01 8.650E-04 138.03 mm 5.79 mm 1.63 mm + 900.00 keV 7.355E-01 7.791E-04 169.00 mm 7.29 mm 1.96 mm + 1.00 MeV 6.772E-01 7.095E-04 202.78 mm 8.74 mm 2.32 mm + 1.10 MeV 6.322E-01 6.518E-04 239.22 mm 10.17 mm 2.71 mm + 1.20 MeV 5.879E-01 6.031E-04 278.33 mm 11.61 mm 3.13 mm + 1.30 MeV 5.513E-01 5.616E-04 320.21 mm 13.07 mm 3.57 mm + 1.40 MeV 5.197E-01 5.256E-04 364.75 mm 14.54 mm 4.05 mm + 1.50 MeV 4.918E-01 4.942E-04 411.91 mm 16.04 mm 4.54 mm + 1.60 MeV 4.670E-01 4.665E-04 461.66 mm 17.55 mm 5.07 mm + 1.70 MeV 4.448E-01 4.418E-04 513.97 mm 19.09 mm 5.62 mm + 1.80 MeV 4.247E-01 4.197E-04 568.83 mm 20.65 mm 6.19 mm + 2.00 MeV 3.899E-01 3.819E-04 686.03 mm 26.57 mm 7.41 mm + 2.25 MeV 3.542E-01 3.436E-04 846.45 mm 35.05 mm 9.09 mm + 2.50 MeV 3.250E-01 3.125E-04 1.02 m 43.08 mm 10.91 mm + 2.75 MeV 3.005E-01 2.868E-04 1.21 m 50.96 mm 12.88 mm + 3.00 MeV 2.797E-01 2.652E-04 1.42 m 58.81 mm 14.99 mm + 3.25 MeV 2.618E-01 2.467E-04 1.64 m 66.70 mm 17.25 mm + 3.50 MeV 2.462E-01 2.307E-04 1.87 m 74.66 mm 19.66 mm + 3.75 MeV 2.325E-01 2.168E-04 2.12 m 82.73 mm 22.20 mm + 4.00 MeV 2.203E-01 2.045E-04 2.39 m 90.91 mm 24.88 mm + 4.50 MeV 1.997E-01 1.838E-04 2.95 m 121.70 mm 30.66 mm + 5.00 MeV 1.829E-01 1.671E-04 3.58 m 150.69 mm 36.98 mm + 5.50 MeV 1.689E-01 1.533E-04 4.26 m 179.04 mm 43.83 mm + 6.00 MeV 1.570E-01 1.416E-04 4.99 m 207.27 mm 51.20 mm + 6.50 MeV 1.468E-01 1.317E-04 5.78 m 235.63 mm 59.08 mm + 7.00 MeV 1.379E-01 1.231E-04 6.61 m 264.26 mm 67.48 mm + 8.00 MeV 1.232E-01 1.090E-04 8.44 m 370.83 mm 85.77 mm + 9.00 MeV 1.116E-01 9.794E-05 10.48 m 470.37 mm 106.02 mm + 10.00 MeV 1.021E-01 8.897E-05 12.71 m 567.83 mm 128.19 mm + 11.00 MeV 9.416E-02 8.155E-05 15.15 m 665.14 mm 152.25 mm +----------------------------------------------------------- + Multiply Stopping by for Stopping Units + ------------------- ------------------ + 4.1905E-04 eV / Angstrom + 4.1905E-03 keV / micron + 4.1905E-03 MeV / mm + 1.0000E+00 keV / (ug/cm2) + 1.0000E+00 MeV / (mg/cm2) + 1.0000E+03 keV / (mg/cm2) + 1.6738E+00 eV / (1E15 atoms/cm2) + 5.5947E-01 L.S.S. reduced units + ================================================================== + (C) 1984.1989.1992.1998.2008 by J.P. Biersack and J.F. Ziegler diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index a29c1a51b..3481cb98f 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -152,7 +152,7 @@ class TrackFitterUKF : public KalmanFilter { /// function /// template - void predictUKF(PredictionModelCallback predictionModelFunc) + void predictUKF(PredictionModelCallback predictionModelFunc, const Vector &vecZ) { updateAugmentedStateAndCovariance(); @@ -171,7 +171,7 @@ class TrackFitterUKF : public KalmanFilter { const Vector sigmaXxi{util::getColumnAt(i, sigmaXx)}; const Vector sigmaXvi{util::getColumnAt(i, sigmaXv)}; - const Vector Yi{predictionModelFunc(sigmaXxi, sigmaXvi)}; // y = f(x) + const Vector Yi{predictionModelFunc(sigmaXxi, sigmaXvi, vecZ)}; // y = f(x) // Copy the predicted state vector back into the sigmaXx matrix. util::copyToColumn(i, sigmaXx, Yi); diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 1751ea443..e4b635ad3 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -1,5 +1,9 @@ #include "kalman_filter/TrackFitterUKF.h" +#include "AtELossTable.h" +#include "AtKinematics.h" + +#include "Math/Vector3D.h" #include "gtest/gtest.h" class TrackFitterUKFExampleTest : public testing::Test { @@ -20,7 +24,7 @@ class TrackFitterUKFExampleTest : public testing::Test { /// @param x state vector /// @param v process noise vector /// @return propagated (unaugmented) state vector - static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v) + static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z) { kf::Vector y; y[0] = x[0] + x[2] + v[0]; @@ -69,7 +73,7 @@ TEST_F(TrackFitterUKFExampleTest, Prediction) m_ukf.setCovarianceQ(Q); m_ukf.setCovarianceR(R); - m_ukf.predictUKF(funcF); + m_ukf.predictUKF(funcF, z); // Expectation from the python results: // ===================================== @@ -130,7 +134,7 @@ TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection) m_ukf.setCovarianceQ(Q); m_ukf.setCovarianceR(R); - m_ukf.predictUKF(funcF); + m_ukf.predictUKF(funcF, z); // Expectation from the python results: // ===================================== @@ -207,6 +211,8 @@ TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection) class TrackFitterUKFPhysicsTest : public testing::Test { public: + using XYZVector = ROOT::Math::XYZVector; + virtual void SetUp() override {} virtual void TearDown() override {} @@ -217,18 +223,192 @@ class TrackFitterUKFPhysicsTest : public testing::Test { static constexpr size_t DIM_Z{3}; static constexpr size_t DIM_N{3}; + static XYZVector fBField; // B-field in tesla + static XYZVector fEField; // E-field in V/m + static constexpr double fC = 299792458; // m/s + kf::TrackFitterUKF m_ukf; + /** + * @param pos Position of particle in mm + * @param mom Momentum of particle in MeV/c + * @param charge charge of the particle in Coulombs + * @param mass mass of particle in MeV/c^2 + * @param dedx Stopping power in MeV/mm + * + * @returns Force in N + */ + static XYZVector Force(XYZVector pos, XYZVector mom, double charge, double mass, double dedx) + { + + // auto fourMom = AtTools::Kinematics::Get4Vector(mom, mass); + // auto v = mom / fourMom.E() * c; // m/s + auto v = GetVel(mom, mass); + + auto F_lorentz = charge * (fEField + v.Cross(fBField)); + // std::cout << "F_lorentz: " << F_lorentz << std::endl; + auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) + + auto drag = -dedx_si * mom.Unit(); + // std::cout << "drag: " << drag << " mom " << mom << " dedx " << dedx_si << std::endl; + + return F_lorentz + drag; + } + + static XYZVector GetVel(XYZVector mom, double mass) + { + auto fourMom = AtTools::Kinematics::Get4Vector(mom, mass); + const double c = 299792458; // m/s + return mom / fourMom.E() * c; // m/s + } + + static double dist(const XYZVector &x, const XYZVector &z) { return std::sqrt((x - z).Mag2()); } + /// @brief to propagate the state vector using the process model /// @param x state vector /// @param v process noise vector + /// @param vecZ The next measurement point used to stop the propagation. /// @return propagated (unaugmented) state vector - static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v) + static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z) { + std::cout << "Staring to run funcF" << std::endl; // TODO: This needs to be filled with an RK4 solver for the physics model kf::Vector y{x}; + XYZVector measurement(z[0], z[1], z[2]); // Measurement point in mm + + double charge = x[6]; + double eLoss = v[0]; + double mass = 938.272; // Mass in MeV/c^2 + + double mat_density = 0; // Density of the material in g/cm^3 + AtTools::AtELossTable dedxModel(mat_density); + dedxModel.LoadSrimTable( + "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); // Load the + // SRIM table + // for energy + // loss + double scalingFactor = 1.0; + int iterations = 0; + double calc_eLoss = 0; + + while (std::abs(calc_eLoss - eLoss) > 1e-3) { + std::cout << "Running iteration " << iterations << " with scaling factor: " << scalingFactor + << " and energy loss: " << calc_eLoss << std::endl; + + if (iterations > 100) { + // If we are not converging, we should probably throw an error. + throw std::runtime_error("Energy loss did not converge after 100 iterations."); + } + + // Variables needed in a single run of the RK4 solver. This section needs to be repeated + // until the energy loss converges to the correct value. + double h = 1e-10; // Timestep in s (100 ns to start) + double lastApproach = std::numeric_limits::max(); + bool approaching = true; + iterations++; + + XYZVector pos(x[0], x[1], x[2]); + XYZVector mom(x[3], x[4], x[5]); + double KE_initial = std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV + + while (true) { + std::cout << "Position: " << pos.X() << ", " << pos.Y() << ", " << pos.Z() << std::endl; + std::cout << "Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; + + // Using timestep, propagate state forward one step. + double KE = std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV + auto dedx = scalingFactor * dedxModel.GetdEdx(KE); // Get the stopping power in MeV/mm + // std::cout << "KE: " << KE << " dedx: " << dedx << std::endl; + + auto spline = dedxModel.GetSpline(); + // std::cout << "Spline: " << spline.get_x_min() << " to " << spline.get_x_max() << std::endl; + // std::cout << "dxde " << spline(KE) << " dxde " << dedx << std::endl; + + auto x_k1 = GetVel(mom, mass); + auto p_k1 = Force(pos, mom, charge, mass, dedx); + // std::cout << "vel: " << x_k1 << " speed " << x_k1.R() << std::endl; + // std::cout << "Force: " << p_k1 << std::endl; + + auto x_k2 = GetVel(mom + p_k1 * h / 2, mass); + auto p_k2 = Force(pos + x_k1 * h / 2, mom + p_k1 * h / 2, charge, mass, dedx); + // std::cout << "vel: " << x_k2 << " speed " << x_k2.R() << std::endl; + // std::cout << "Force: " << p_k2 << std::endl; + + auto x_k3 = GetVel(mom + p_k2 * h / 2, mass); + auto p_k3 = Force(pos + x_k2 * h / 2, mom + p_k2 * h / 2, charge, mass, dedx); + // std::cout << "vel: " << x_k3 << " speed " << x_k3.R() << std::endl; + // std::cout << "Force: " << p_k3 << std::endl; + + auto x_k4 = GetVel(mom + p_k3 * h, mass); + auto p_k4 = Force(pos + x_k3 * h, mom + p_k3 * h, charge, mass, dedx); + // std::cout << "vel: " << x_k4 << " speed " << x_k4.R() << std::endl; + // std::cout << "Force: " << p_k4 << std::endl; + + auto mom_SItoMeV = 1.60218e-13 / 299792458; // Factor to convert momentum to MeV/c from kg*m/s + auto F_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; + + // Convert momentum to SI, update, and convert back to MeV/c + auto mom_SI = mom * mom_SItoMeV; + mom_SI += F_SI * h; + mom = mom_SI / mom_SItoMeV; // Convert back to MeV/c + + // Convert position to SI, update, and convert back to mm + auto pos_SI = pos / 1e3; // Convert mm to m + pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; + pos = pos_SI * 1e3; // Convert back to mm + + std::cout << "Average force: " << F_SI << " N" << std::endl; + std::cout << "Momentum: " << mom * mom_SItoMeV << " kg m/s" << std::endl; + std::cout << "Delta x: " << (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6 << " m/s" << std::endl; + + auto approach = dist(pos, measurement); + std::cout << "pos: " << pos << " measurement" << measurement << std::endl; + std::cout << "Approach: " << approach << " last approach: " << lastApproach << std::endl; + if (approach < lastApproach) { + // We are still approaching the measurement point + approaching = true; + lastApproach = approach; + + continue; + } + + bool reachedMeasurementPoint = (approaching && approach > lastApproach); + bool particleStopped = std::sqrt(mom.Mag2() + mass * mass) - mass < 0.01; + if (reachedMeasurementPoint || particleStopped) { + // Last iteration we were still approaching the measurement point. Now we are further away + // then before. We have probably reached the measurement point if things are well behaved. + // I can think of cases where this will not be true. A better solution might be to run + // tracking the point of closest approach until the distance between the current state and + // the measurement point is larger than the distance between the last state and the measurement point. + + // Undo the last update, and break out of the integration loop + // pos = pos - (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; + // mom = mom - (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) * h / 6; + y[0] = pos.X(); + y[1] = pos.Y(); + y[2] = pos.Z(); + y[3] = mom.X(); + y[4] = mom.Y(); + y[5] = mom.Z(); + + // Update the scaling factor + double KE_final = std::sqrt(mom.Mag2() + mass * mass) - mass; + calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + scalingFactor *= calc_eLoss / eLoss; + std::cout << "------- End of RK4 interation ---------" << std::endl; + std::cout << "Particle stopped: " << particleStopped << std::endl; + std::cout << "Reached measurement point: " << reachedMeasurementPoint << std::endl; + std::cout << "Last approach: " << lastApproach << " Current approach: " << approach << std::endl; + std::cout << "Desired energy loss: " << eLoss << " MeV" << std::endl; + std::cout << "Calculated energy loss: " << calc_eLoss << " MeV" << std::endl; + std::cout << "New scaling factor: " << scalingFactor << std::endl; + std::cout << "Final Position: " << pos.X() << ", " << pos.Y() << ", " << pos.Z() << std::endl; + std::cout << "Final Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; + return y; + } + } // End of loop over RK4 integration + } // End loop over energy loss convergence - // For now, we just return the state vector as is return y; } @@ -244,8 +424,15 @@ class TrackFitterUKFPhysicsTest : public testing::Test { return y; } + + const double mass_p = 938.272; // Mass of proton in MeV/c^2 + const double charge_p = 1.602176634e-19; // Charge of proton }; +// Definition of static member variables +TrackFitterUKFPhysicsTest::XYZVector TrackFitterUKFPhysicsTest::fBField; +TrackFitterUKFPhysicsTest::XYZVector TrackFitterUKFPhysicsTest::fEField; + TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) { // TODO: This needs to be filled with a proper physics test. @@ -263,6 +450,102 @@ TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) m_ukf.vecX() = x; m_ukf.matP() = P; - m_ukf.setCovarianceR(R); - m_ukf.predictUKF(funcF); + // m_ukf.setCovarianceR(R); + // m_ukf.predictUKF(funcF, z); + ASSERT_EQ(true, true); +} + +TEST_F(TrackFitterUKFPhysicsTest, TestForceNoFields) +{ + XYZVector pos(0, 0, 0); // Position in mm + XYZVector mom(100, 0, 0); // Momentum in MeV/c + fBField = XYZVector(0, 0, 0); // B-field in tesla + fEField = XYZVector(0, 0, 0); // E-field + + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + double dedx = 1; // Stopping power in MeV/mm + + auto force = Force(pos, mom, charge, mass, dedx); + + ASSERT_NEAR(force.X(), -1.602e-10, FLOAT_EPSILON); + ASSERT_NEAR(force.Y(), 0, FLOAT_EPSILON); + ASSERT_NEAR(force.Z(), 0, FLOAT_EPSILON); + + mom = XYZVector(100, 0, 100); // Reset momentum + force = Force(pos, mom, charge, mass, dedx); + ASSERT_NEAR(force.X(), -1.602e-10 / std::sqrt(2), FLOAT_EPSILON); + ASSERT_NEAR(force.Y(), 0, FLOAT_EPSILON); + ASSERT_NEAR(force.Z(), -1.602e-10 / std::sqrt(2), FLOAT_EPSILON); +} + +TEST_F(TrackFitterUKFPhysicsTest, TestForceEField) +{ + XYZVector pos(0, 0, 0); // Position in mm + XYZVector mom(100, 0, 0); // Momentum in MeV/c + fBField = XYZVector(0, 0, 1); // B-field in tesla + fEField = XYZVector(0, 0, 0); // E-field in V/m + + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + double dedx = 0; // Stopping power in MeV/mm + + auto force = Force(pos, mom, charge, mass, dedx); + + ASSERT_NEAR(force.X(), 0, FLOAT_EPSILON); + ASSERT_NEAR(force.Y(), 0, FLOAT_EPSILON); + ASSERT_NEAR(force.Z(), 1.121e-14, FLOAT_EPSILON); +} + +TEST_F(TrackFitterUKFPhysicsTest, TestForceBField) +{ + XYZVector pos(0, 0, 0); // Position in mm + XYZVector mom(100, 0, 0); // Momentum in MeV/c + fBField = XYZVector(0, 0, 1); // B-field in tesla + fEField = XYZVector(0, 0, 0); // E-field + + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + double dedx = 0; // Stopping power in MeV/mm + + auto force = Force(pos, mom, charge, mass, dedx); + + ASSERT_NEAR(force.X(), 0, FLOAT_EPSILON); + ASSERT_NEAR(force.Y(), -5.09e-12, FLOAT_EPSILON); + ASSERT_NEAR(force.Z(), 0, FLOAT_EPSILON); +} + +TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) +{ + + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + fBField = XYZVector(0, 0, 0); // B-field in tesla + fEField = XYZVector(0, 0, 0); // E-field + + kf::Vector x; // Initial state vector + x[0] = 0; + x[1] = 0; + x[2] = 0; + x[3] = p; // p_x + x[4] = 0; + x[5] = 0; + + ASSERT_NEAR(x[3], 43.331, 1e-1); // Make sure momentum is calculated correctly + + kf::Vector v; // Process noise vector + v[0] = 1; // Energy loss in MeV + v[1] = 0.0; // No process noise in this example + + kf::Vector z; // Measurement vector + z[0] = 1e3; + z[1] = 0; + z[2] = 0; + + auto final = funcF(x, v, z); // Propagate the state vector using the process model + + // Check the final position is close to the stopping point from LISE + ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 + ASSERT_NEAR(final[0], 210, 10); // Final position in x-direction should be close to 210 mm } diff --git a/AtTools/AtELossTable.h b/AtTools/AtELossTable.h index de3c29657..683bdf57d 100644 --- a/AtTools/AtELossTable.h +++ b/AtTools/AtELossTable.h @@ -51,6 +51,8 @@ class AtELossTable : public AtELossModel { [[deprecated]] double GetEnergyOld(double energyIni, double distance) const; + const tk::spline &GetSpline() const { return fdXdE; } + private: void LoadTable(const std::vector &energy, const std::vector &dEdX); void LoadRangeVariance(const std::vector &energy, const std::vector &rangeVariance) From abae9bc06958a953abd1a79c5127426ff9a56ecd Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 17 Jul 2025 15:58:56 +0200 Subject: [PATCH 08/71] Add additional test --- .../OpenKF/kalman_filter/TrackFitterUKFTest.cxx | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index e4b635ad3..a7dbb9c68 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -381,9 +381,6 @@ class TrackFitterUKFPhysicsTest : public testing::Test { // tracking the point of closest approach until the distance between the current state and // the measurement point is larger than the distance between the last state and the measurement point. - // Undo the last update, and break out of the integration loop - // pos = pos - (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; - // mom = mom - (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) * h / 6; y[0] = pos.X(); y[1] = pos.Y(); y[2] = pos.Z(); @@ -548,4 +545,13 @@ TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) // Check the final position is close to the stopping point from LISE ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 ASSERT_NEAR(final[0], 210, 10); // Final position in x-direction should be close to 210 mm + + KE = 0.5; + E = KE + mass_p; + p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + x[3] = p; // Reset momentum + final = funcF(x, v, z); // Propagate the state vector using the + + ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 + ASSERT_NEAR(final[0], 68.6, 5); // Final position in x } From a7036b6624dbf08efce476ea38c4505b6b41c8c1 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 17 Jul 2025 17:01:24 +0200 Subject: [PATCH 09/71] Add test for fixing energy loss --- .../kalman_filter/TrackFitterUKFTest.cxx | 84 +++++++++++++++---- 1 file changed, 67 insertions(+), 17 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index a7dbb9c68..903e5fb7d 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -291,7 +291,7 @@ class TrackFitterUKFPhysicsTest : public testing::Test { int iterations = 0; double calc_eLoss = 0; - while (std::abs(calc_eLoss - eLoss) > 1e-3) { + while (std::abs(calc_eLoss - eLoss) > 1e-4) { std::cout << "Running iteration " << iterations << " with scaling factor: " << scalingFactor << " and energy loss: " << calc_eLoss << std::endl; @@ -309,9 +309,12 @@ class TrackFitterUKFPhysicsTest : public testing::Test { XYZVector pos(x[0], x[1], x[2]); XYZVector mom(x[3], x[4], x[5]); + double KE_initial = std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV while (true) { + XYZVector lastPos = pos; + XYZVector lastMom = mom; std::cout << "Position: " << pos.X() << ", " << pos.Y() << ", " << pos.Z() << std::endl; std::cout << "Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; @@ -357,13 +360,13 @@ class TrackFitterUKFPhysicsTest : public testing::Test { pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; pos = pos_SI * 1e3; // Convert back to mm - std::cout << "Average force: " << F_SI << " N" << std::endl; - std::cout << "Momentum: " << mom * mom_SItoMeV << " kg m/s" << std::endl; - std::cout << "Delta x: " << (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6 << " m/s" << std::endl; + // std::cout << "Average force: " << F_SI << " N" << std::endl; + // std::cout << "Momentum: " << mom * mom_SItoMeV << " kg m/s" << std::endl; + // std::cout << "Delta x: " << (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6 << " m/s" << std::endl; auto approach = dist(pos, measurement); - std::cout << "pos: " << pos << " measurement" << measurement << std::endl; - std::cout << "Approach: " << approach << " last approach: " << lastApproach << std::endl; + // std::cout << "pos: " << pos << " measurement" << measurement << std::endl; + // std::cout << "Approach: " << approach << " last approach: " << lastApproach << std::endl; if (approach < lastApproach) { // We are still approaching the measurement point approaching = true; @@ -381,27 +384,30 @@ class TrackFitterUKFPhysicsTest : public testing::Test { // tracking the point of closest approach until the distance between the current state and // the measurement point is larger than the distance between the last state and the measurement point. - y[0] = pos.X(); - y[1] = pos.Y(); - y[2] = pos.Z(); - y[3] = mom.X(); - y[4] = mom.Y(); - y[5] = mom.Z(); + // Undo the last step since we were closer last time. + y[0] = lastPos.X(); + y[1] = lastPos.Y(); + y[2] = lastPos.Z(); + y[3] = lastMom.X(); + y[4] = lastMom.Y(); + y[5] = lastMom.Z(); // Update the scaling factor - double KE_final = std::sqrt(mom.Mag2() + mass * mass) - mass; + double KE_final = std::sqrt(lastMom.Mag2() + mass * mass) - mass; calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - scalingFactor *= calc_eLoss / eLoss; - std::cout << "------- End of RK4 interation ---------" << std::endl; + scalingFactor *= eLoss / calc_eLoss; + std::cout << "------- End of RK4 interation " << iterations << " ---------" << std::endl; std::cout << "Particle stopped: " << particleStopped << std::endl; std::cout << "Reached measurement point: " << reachedMeasurementPoint << std::endl; std::cout << "Last approach: " << lastApproach << " Current approach: " << approach << std::endl; std::cout << "Desired energy loss: " << eLoss << " MeV" << std::endl; std::cout << "Calculated energy loss: " << calc_eLoss << " MeV" << std::endl; + std::cout << "Difference: " << calc_eLoss - eLoss << " MeV" << std::endl; std::cout << "New scaling factor: " << scalingFactor << std::endl; std::cout << "Final Position: " << pos.X() << ", " << pos.Y() << ", " << pos.Z() << std::endl; std::cout << "Final Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; - return y; + break; + // return y; } } // End of loop over RK4 integration } // End loop over energy loss convergence @@ -512,7 +518,7 @@ TEST_F(TrackFitterUKFPhysicsTest, TestForceBField) ASSERT_NEAR(force.Z(), 0, FLOAT_EPSILON); } -TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) +TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorStoppingNoField) { double KE = 1; // Kinetic energy in MeV @@ -555,3 +561,47 @@ TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 ASSERT_NEAR(final[0], 68.6, 5); // Final position in x } + +TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) +{ + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + fBField = XYZVector(0, 0, 0); // B-field in tesla + fEField = XYZVector(0, 0, 0); // E-field + + kf::Vector x; // Initial state vector + x[0] = 0; + x[1] = 0; + x[2] = 0; + x[3] = p; // p_x + x[4] = 0; + x[5] = 0; + + ASSERT_NEAR(x[3], 43.331, 1e-1); // Make sure momentum is calculated correctly + + kf::Vector v; // Process noise vector + v[0] = 0.0285; // Energy loss in MeV + v[1] = 0.0; // No process noise in this example + + kf::Vector z; // Measurement vector + z[0] = 10; // Measure after 10 mm + z[1] = 0; + z[2] = 0; + + auto final = funcF(x, v, z); // Propagate the state vector using the process model + + // Check the final position is close to the stopping point from LISE + double E_fin = KE - v[0] + mass_p; + double p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); + ASSERT_NEAR(final[3], p_fin, 0.1); // Final momentum in x-direction + ASSERT_NEAR(final[0], 10, .5); // Final position in x-direction should be close to 10 mm + + v[0] = 0.3237; // Energy loss in MeV + z[0] = 100; // Measure after 100 mm + final = funcF(x, v, z); // Propagate the state vector using the process model + E_fin = KE - v[0] + mass_p; + p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); + ASSERT_NEAR(final[3], p_fin, 0.1); // Final momentum + ASSERT_NEAR(final[0], 100, .5); // Final position +} From 37fc28bfa376c76c5ac0f27e5e15f1154ea947d7 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 17 Jul 2025 18:13:37 +0200 Subject: [PATCH 10/71] Move some code from tests to AtTools --- AtTools/AtKinematics.cxx | 9 +++- AtTools/AtKinematics.h | 19 ++++++++ AtTools/AtPropagator.cxx | 20 ++++++++ AtTools/AtPropagator.h | 93 +++++++++++++++++++++++++++++++++++ AtTools/AtPropagatorTest.cxx | 95 ++++++++++++++++++++++++++++++++++++ AtTools/AtToolsLinkDef.h | 1 + AtTools/CMakeLists.txt | 2 + 7 files changed, 238 insertions(+), 1 deletion(-) create mode 100644 AtTools/AtPropagator.cxx create mode 100644 AtTools/AtPropagator.h create mode 100644 AtTools/AtPropagatorTest.cxx diff --git a/AtTools/AtKinematics.cxx b/AtTools/AtKinematics.cxx index 2431d34f2..755baf91c 100644 --- a/AtTools/AtKinematics.cxx +++ b/AtTools/AtKinematics.cxx @@ -303,5 +303,12 @@ double EtoA(double mass) { return mass / 931.5; } - +ROOT::Math::XYZVector GetVel(ROOT::Math::XYZVector mom, double mass) +{ + return mom / Get4Vector(mom, mass).E() * fC; +} +double KE(ROOT::Math::XYZVector mom, double mass) +{ + return std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV +} } // namespace AtTools::Kinematics diff --git a/AtTools/AtKinematics.h b/AtTools/AtKinematics.h index b0d343fd2..8be5c0111 100644 --- a/AtTools/AtKinematics.h +++ b/AtTools/AtKinematics.h @@ -7,6 +7,7 @@ #ifndef ATKINEMATICS_H #define ATKINEMATICS_H +#include #include #include // for PxPyPzEVector #include // for Double_t, THashConsistencyHolder, Int_t, ClassDef @@ -64,6 +65,8 @@ class AtKinematics : public TObject { namespace Kinematics { +static constexpr double fC = 299792458.0; // Speed of light in m/s + double GetGamma(double KE, double m1, double m2); double GetGamma(double beta); double GetVelocity(double gamma); @@ -74,6 +77,22 @@ double GetRelMom(double gamma, double mass); double AtoE(double Amu); double EtoA(double mass); +/** + * Calculate the kinetic energy of a particle given its momentum and mass. + * @param mom Momentum vector (in MeV/c) + * @param mass Mass of the particle (in MeV/c^2) + * @returns Kinetic energy in MeV + */ +double KE(ROOT::Math::XYZVector mom, double mass); + +/** + * Calculate the velocity vector from momentum and mass + * @param mom Momentum vector (in MeV/c) + * @param mass Mass of the particle (in MeV/c^2) + * @returns Velocity vector in m/s + */ +ROOT::Math::XYZVector GetVel(ROOT::Math::XYZVector mom, double mass); + template ROOT::Math::PxPyPzEVector Get4Vector(Vector mom, double m) { diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx new file mode 100644 index 000000000..c7a9b8858 --- /dev/null +++ b/AtTools/AtPropagator.cxx @@ -0,0 +1,20 @@ +#include "AtPropagator.h" + +#include "AtKinematics.h" +namespace AtTools { + +AtPropagator::XYZVector AtPropagator::Force(XYZVector pos, XYZVector mom) const +{ + auto v = Kinematics::GetVel(mom, fMass); + + auto F_lorentz = fQ * (fEField + v.Cross(fBField)); + // std::cout << "F_lorentz: " << F_lorentz << std::endl; + auto dedx = fScalingFactor * fELossModel->GetdEdx(Kinematics::KE(mom, fMass)); // Stopping power in MeV/mm + auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) + + auto drag = -dedx_si * mom.Unit(); + // std::cout << "drag: " << drag << " mom " << mom << " dedx " << dedx_si << std::endl; + + return F_lorentz + drag; +} +} // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h new file mode 100644 index 000000000..58a2df0d3 --- /dev/null +++ b/AtTools/AtPropagator.h @@ -0,0 +1,93 @@ +#ifndef ATPROPAGATOR_H +#define ATPROPAGATOR_H + +#include "AtELossModel.h" + +#include "Math/Vector3D.h" + +namespace AtTools { + +/** + * @brief Class for propagating particles through a medium. + * + * This class is responsible for simulating the propagation of particles + * through a medium, taking into account energy loss and other effects. + * Uses an AtELossModel to calculate the energy loss and propagates particles + * in the presence of electric and magnetic fields. + * + * Class is designed to be used with a single particle type. Create a new instance of the + * class if the material or particle type changes. + */ +class AtPropagator { +protected: + using XYZVector = ROOT::Math::XYZVector; + XYZVector fEField{0, 0, 0}; // Electric field vector + XYZVector fBField{0, 0, 0}; // Magnetic field vector + + const double fQ; // Charge of the particle in Coulombs + const double fMass; // Mass of the particle in MeV/c^2 + const std::unique_ptr fELossModel; // Energy loss model + + // Internal state variables for the propagator + double fH = 1e-10; /// Step size for propagation in s + double fETolerance = 1e-4; /// Energy tolerance for convergence when fixing energy loss + double fScalingFactor = 1.0; /// Scaling factor for energy loss + + XYZVector fPos; // Current position of the particle in mm + XYZVector fMom; // Current momentum of the particle in MeV/c + + static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) + +public: + AtPropagator(double charge, double mass, std::unique_ptr elossModel) + : fQ(charge), fMass(mass), fELossModel(std::move(elossModel)) + { + } + /** + * @brief Set the electric field (V/m) + */ + void SetEField(const XYZVector &eField) { fEField = eField; } + void SetEField(double ex, double ey, double ez) { fEField.SetXYZ(ex, ey, ez); } + /** + * @brief Set the magnetic field (T) + */ + void SetBField(const XYZVector &bField) { fBField = bField; } + void SetBField(double bx, double by, double bz) { fBField.SetXYZ(bx, by, bz); } + + /** + * @brief Set the state of the particle. + * + * @param pos Position of the particle in mm. + * @param mom Momentum of the particle in MeV/c. + */ + void SetState(const XYZVector &pos, const XYZVector &mom) + { + fPos = pos; + fMom = mom; + } + XYZVector GetPosition() const { return fPos; } + XYZVector GetMomentum() const { return fMom; } + + /** + * @brief Propagate the particle to the point of closest approach to the given point. + * + * Propagate to a given point in space, adjusting the magnitude of the stopping power + * to ensure that a specific about of energy is lost during the propagation. + * + * @param point The point to approach. + * @param eLoss If not 0, constrain the energy loss to this value (adjusting the stopping power). + */ + void PropagateToPoint(const XYZVector &point, double eLoss = 0); + + /** + * @brief Calculate the force acting on the particle. + * + * @param pos Position of the particle in mm. + * @param mom Momentum of the particle in MeV/c. + * @return The force acting on the particle in N. + */ + XYZVector Force(XYZVector pos, XYZVector mom) const; +}; + +} // namespace AtTools +#endif // #ifndef ATPROPAGATOR_H diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx new file mode 100644 index 000000000..56a6ba77f --- /dev/null +++ b/AtTools/AtPropagatorTest.cxx @@ -0,0 +1,95 @@ +#include "AtPropagator.h" + +#include "AtKinematics.h" + +#include + +#include +#include +using ROOT::Math::XYZVector; + +using namespace AtTools; + +const double mass_p = 938.272; // Mass of proton in MeV/c^2 +const double charge_p = 1.602176634e-19; // Charge of proton + +class DummyELossModel : public AtELossModel { +public: + double eLoss = 1; + DummyELossModel() : AtELossModel(0) {} + + double GetdEdx(double /*KE*/) const override { return eLoss; } + double GetRange(double /*energyIni*/, double /*energyFin = 0*/) const override { return 1.0; } + double GetEnergyLoss(double /*energyIni*/, double /*distance*/) const override { return 1.0; } + double GetEnergy(double /*energyIni*/, double /*distance*/) const override { return 1.0; } +}; + +TEST(AtPropagatorTest, ForceNoField) +{ + XYZVector pos(0, 0, 0); // Position in mm + XYZVector mom(100, 0, 0); // Momentum in MeV/c + + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + double dedx = 1; // Stopping power in MeV/mm + + // Create a dummy energy loss model + auto elossModel = std::make_unique(); + AtPropagator propagator(charge, mass, std::move(elossModel)); + propagator.SetEField({0, 0, 0}); + propagator.SetBField({0, 0, 0}); + + auto force = propagator.Force(pos, mom); + + ASSERT_NEAR(force.X(), -1.602e-10, 1e-12); + ASSERT_NEAR(force.Y(), 0, 1e-12); + ASSERT_NEAR(force.Z(), 0, 1e-12); + + mom = XYZVector(100, 0, 100); // Reset momentum + force = propagator.Force(pos, mom); + ASSERT_NEAR(force.X(), -1.602e-10 / std::sqrt(2), 1e-12); + ASSERT_NEAR(force.Y(), 0, 1e-12); + ASSERT_NEAR(force.Z(), -1.602e-10 / std::sqrt(2), 1e-12); +} + +TEST(AtPropagatorTest, ForceEField) +{ + XYZVector pos(0, 0, 0); // Position in mm + XYZVector mom(100, 0, 0); // Momentum in MeV/c + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + + // Create a dummy energy loss model + auto elossModel = std::make_unique(); + elossModel->eLoss = 0; // No energy loss for this test + AtPropagator propagator(charge, mass, std::move(elossModel)); + propagator.SetEField({0, 0, 70000}); + propagator.SetBField({0, 0, 0}); + + auto force = propagator.Force(pos, mom); + + ASSERT_NEAR(force.X(), 0, 1e-12); + ASSERT_NEAR(force.Y(), 0, 1e-12); + ASSERT_NEAR(force.Z(), 1.121e-14, 1e-15); +} + +TEST(AtPropagatorTest, ForceBField) +{ + XYZVector pos(0, 0, 0); // Position in mm + XYZVector mom(100, 0, 0); // Momentum in MeV/c + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + + // Create a dummy energy loss model + auto elossModel = std::make_unique(); + elossModel->eLoss = 0; // No energy loss for this test + AtPropagator propagator(charge, mass, std::move(elossModel)); + propagator.SetEField({0, 0, 0}); + propagator.SetBField({0, 0, 1}); + + auto force = propagator.Force(pos, mom); + + ASSERT_NEAR(force.X(), 0, 1e-12); + ASSERT_NEAR(force.Y(), -5.09e-12, 1e-13); + ASSERT_NEAR(force.Z(), 0, 1e-12); +} \ No newline at end of file diff --git a/AtTools/AtToolsLinkDef.h b/AtTools/AtToolsLinkDef.h index 729423c95..c8f178919 100644 --- a/AtTools/AtToolsLinkDef.h +++ b/AtTools/AtToolsLinkDef.h @@ -27,6 +27,7 @@ #pragma link C++ class AtEDistortionModel - !; #pragma link C++ class AtTools::AtKinematics + ; +#pragma link C++ class AtTools::AtPropagator - !; #pragma link C++ class AtTools::AtVirtualTerminal + ; #pragma link C++ class RandomSample::AtSample - !; diff --git a/AtTools/CMakeLists.txt b/AtTools/CMakeLists.txt index eae5925fc..82b04ad96 100644 --- a/AtTools/CMakeLists.txt +++ b/AtTools/CMakeLists.txt @@ -19,6 +19,7 @@ set(SRCS AtEDistortionModel.cxx AtVirtualTerminal.cxx AtKinematics.cxx + AtPropagator.cxx AtFormat.cxx AtSpline.cxx @@ -77,6 +78,7 @@ Set(INCLUDE_DIR set(TEST_SRCS DataCleaning/AtkNNTest.cxx AtELossTableTest.cxx + AtPropagatorTest.cxx ) if(CATIMA_FOUND) set(TEST_SRCS ${TEST_SRCS} From b0bed98062bc83f75fd3d3714ce38fa0b3665502 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 21 Jul 2025 11:42:12 +0200 Subject: [PATCH 11/71] Add Propagator class for charged particles in mat Also adds tests. --- .../kalman_filter/TrackFitterUKFTest.cxx | 27 ++--- AtTools/AtPropagator.cxx | 108 +++++++++++++++++- AtTools/AtPropagator.h | 12 +- AtTools/AtPropagatorTest.cxx | 84 ++++++++++++++ 4 files changed, 212 insertions(+), 19 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 903e5fb7d..183e96c0c 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -243,7 +243,7 @@ class TrackFitterUKFPhysicsTest : public testing::Test { // auto fourMom = AtTools::Kinematics::Get4Vector(mom, mass); // auto v = mom / fourMom.E() * c; // m/s - auto v = GetVel(mom, mass); + auto v = AtTools::Kinematics::GetVel(mom, mass); auto F_lorentz = charge * (fEField + v.Cross(fBField)); // std::cout << "F_lorentz: " << F_lorentz << std::endl; @@ -255,13 +255,6 @@ class TrackFitterUKFPhysicsTest : public testing::Test { return F_lorentz + drag; } - static XYZVector GetVel(XYZVector mom, double mass) - { - auto fourMom = AtTools::Kinematics::Get4Vector(mom, mass); - const double c = 299792458; // m/s - return mom / fourMom.E() * c; // m/s - } - static double dist(const XYZVector &x, const XYZVector &z) { return std::sqrt((x - z).Mag2()); } /// @brief to propagate the state vector using the process model @@ -310,7 +303,8 @@ class TrackFitterUKFPhysicsTest : public testing::Test { XYZVector pos(x[0], x[1], x[2]); XYZVector mom(x[3], x[4], x[5]); - double KE_initial = std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV + double KE_initial = + AtTools::Kinematics::KE(mom, mass); // std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV while (true) { XYZVector lastPos = pos; @@ -319,30 +313,30 @@ class TrackFitterUKFPhysicsTest : public testing::Test { std::cout << "Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; // Using timestep, propagate state forward one step. - double KE = std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV - auto dedx = scalingFactor * dedxModel.GetdEdx(KE); // Get the stopping power in MeV/mm + double KE = AtTools::Kinematics::KE(mom, mass); // Kinetic energy in MeV + auto dedx = scalingFactor * dedxModel.GetdEdx(KE); // Get the stopping power in MeV/mm // std::cout << "KE: " << KE << " dedx: " << dedx << std::endl; auto spline = dedxModel.GetSpline(); // std::cout << "Spline: " << spline.get_x_min() << " to " << spline.get_x_max() << std::endl; // std::cout << "dxde " << spline(KE) << " dxde " << dedx << std::endl; - auto x_k1 = GetVel(mom, mass); + auto x_k1 = AtTools::Kinematics::GetVel(mom, mass); auto p_k1 = Force(pos, mom, charge, mass, dedx); // std::cout << "vel: " << x_k1 << " speed " << x_k1.R() << std::endl; // std::cout << "Force: " << p_k1 << std::endl; - auto x_k2 = GetVel(mom + p_k1 * h / 2, mass); + auto x_k2 = AtTools::Kinematics::GetVel(mom + p_k1 * h / 2, mass); auto p_k2 = Force(pos + x_k1 * h / 2, mom + p_k1 * h / 2, charge, mass, dedx); // std::cout << "vel: " << x_k2 << " speed " << x_k2.R() << std::endl; // std::cout << "Force: " << p_k2 << std::endl; - auto x_k3 = GetVel(mom + p_k2 * h / 2, mass); + auto x_k3 = AtTools::Kinematics::GetVel(mom + p_k2 * h / 2, mass); auto p_k3 = Force(pos + x_k2 * h / 2, mom + p_k2 * h / 2, charge, mass, dedx); // std::cout << "vel: " << x_k3 << " speed " << x_k3.R() << std::endl; // std::cout << "Force: " << p_k3 << std::endl; - auto x_k4 = GetVel(mom + p_k3 * h, mass); + auto x_k4 = AtTools::Kinematics::GetVel(mom + p_k3 * h, mass); auto p_k4 = Force(pos + x_k3 * h, mom + p_k3 * h, charge, mass, dedx); // std::cout << "vel: " << x_k4 << " speed " << x_k4.R() << std::endl; // std::cout << "Force: " << p_k4 << std::endl; @@ -457,6 +451,7 @@ TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) // m_ukf.predictUKF(funcF, z); ASSERT_EQ(true, true); } +/* TEST_F(TrackFitterUKFPhysicsTest, TestForceNoFields) { @@ -553,6 +548,7 @@ TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorStoppingNoField) ASSERT_NEAR(final[0], 210, 10); // Final position in x-direction should be close to 210 mm KE = 0.5; + v[0] = 0.5; // Energy loss in MeV E = KE + mass_p; p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c x[3] = p; // Reset momentum @@ -561,6 +557,7 @@ TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorStoppingNoField) ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 ASSERT_NEAR(final[0], 68.6, 5); // Final position in x } +*/ TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) { diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index c7a9b8858..3728335af 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -1,6 +1,8 @@ #include "AtPropagator.h" #include "AtKinematics.h" + +#include namespace AtTools { AtPropagator::XYZVector AtPropagator::Force(XYZVector pos, XYZVector mom) const @@ -8,13 +10,113 @@ AtPropagator::XYZVector AtPropagator::Force(XYZVector pos, XYZVector mom) const auto v = Kinematics::GetVel(mom, fMass); auto F_lorentz = fQ * (fEField + v.Cross(fBField)); - // std::cout << "F_lorentz: " << F_lorentz << std::endl; + LOG(debug) << "F_lorentz: " << F_lorentz; auto dedx = fScalingFactor * fELossModel->GetdEdx(Kinematics::KE(mom, fMass)); // Stopping power in MeV/mm auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) auto drag = -dedx_si * mom.Unit(); - // std::cout << "drag: " << drag << " mom " << mom << " dedx " << dedx_si << std::endl; + LOG(debug) << "drag: " << drag << " mom " << mom << " dedx " << dedx_si; + + return F_lorentz + drag; // Force in N +} + +void AtPropagator::RK4Step() +{ + double h = fH; // Step size in seconds + + auto x_k1 = Kinematics::GetVel(fMom, fMass); + auto p_k1 = Force(fPos, fMom); + + auto x_k2 = Kinematics::GetVel(fMom + p_k1 * h / 2, fMass); + auto p_k2 = Force(fPos + x_k1 * h / 2, fMom + p_k1 * h / 2); + + auto x_k3 = Kinematics::GetVel(fMom + p_k2 * h / 2, fMass); + auto p_k3 = Force(fPos + x_k2 * h / 2, fMom + p_k2 * h / 2); + + auto x_k4 = Kinematics::GetVel(fMom + p_k3 * h, fMass); + auto p_k4 = Force(fPos + x_k3 * h, fMom + p_k3 * h); + + auto F_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // Force in SI units (N) + + auto mom_SI = fReltoSImom * fMom; + mom_SI += F_SI * h; // Update momentum in SI units (kg m/s) + fMom = mom_SI / fReltoSImom; // Convert back to + + auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) + pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; // Update position in SI units (m + fPos = pos_SI * 1e3; // Convert back to mm +} + +void AtPropagator::PropagateToPoint(const XYZVector &point, double eLoss) +{ + LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; + + double scalingFactor = 1.0; + int iterations = 0; + double calc_eLoss = 0; + + while (std::abs(calc_eLoss - eLoss) > 1e-4 || eLoss == 0) { + LOG(debug) << "Running iteration " << iterations << " with scaling factor: " << scalingFactor + << " and energy loss: " << calc_eLoss; + + if (iterations > 100) { + // If we are not converging, we should probably throw an error. + throw std::runtime_error("Energy loss did not converge after 100 iterations."); + } + + double lastApproach = std::numeric_limits::max(); + bool approaching = true; + iterations++; + auto KE_initial = Kinematics::KE(fMom, fMass); + + while (true) { + XYZVector lastPos = fPos; + XYZVector lastMom = fMom; + LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + + RK4Step(); + + auto approach = (fPos - point).R(); + if (approach < lastApproach) { + // We are still approaching the measurement point + approaching = true; + lastApproach = approach; + continue; + } + + bool reachedMeasurementPoint = (approaching && approach > lastApproach); + bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; + if (reachedMeasurementPoint || particleStopped) { + // Last iteration we were still approaching the measurement point. Now we are further away + // then before. We have probably reached the measurement point if things are well behaved. + // I can think of cases where this will not be true. A better solution might be to run + // tracking the point of closest approach until the distance between the current state and + // the measurement point is larger than the distance between the last state and the measurement point. + + // Undo the last step since we were closer last time. + fPos = lastPos; + fMom = lastMom; + + double KE_final = Kinematics::KE(fMom, fMass); + calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + scalingFactor *= eLoss / calc_eLoss; + LOG(info) << "------- End of RK4 interation " << iterations << " ---------"; + LOG(info) << "Particle stopped: " << particleStopped; + LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; + LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; + LOG(info) << "Desired energy loss: " << eLoss << " MeV"; + LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(info) << "Difference: " << calc_eLoss - eLoss << " MeV"; + LOG(info) << "New scaling factor: " << scalingFactor; + LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - return F_lorentz + drag; + if (eLoss == 0) + return; // If no energy loss is specified, we are done. + break; // Else rerun with adjusted scaling factor + } + } // End of loop over RK4 integration + } // End loop over energy loss convergence } } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 58a2df0d3..066d894f4 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -30,7 +30,8 @@ class AtPropagator { // Internal state variables for the propagator double fH = 1e-10; /// Step size for propagation in s - double fETolerance = 1e-4; /// Energy tolerance for convergence when fixing energy loss + double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss + double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped double fScalingFactor = 1.0; /// Scaling factor for energy loss XYZVector fPos; // Current position of the particle in mm @@ -87,6 +88,15 @@ class AtPropagator { * @return The force acting on the particle in N. */ XYZVector Force(XYZVector pos, XYZVector mom) const; + +protected: + /** + * @brief Perform a single RK4 step for propagation. + * + * This method performs a single Runge-Kutta 4th order step to propagate the particle's state. + * Updates fPos and fMom. + */ + void RK4Step(); }; } // namespace AtTools diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 56a6ba77f..f3c130187 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -1,10 +1,12 @@ #include "AtPropagator.h" +#include "AtELossTable.h" #include "AtKinematics.h" #include #include +#include #include using ROOT::Math::XYZVector; @@ -92,4 +94,86 @@ TEST(AtPropagatorTest, ForceBField) ASSERT_NEAR(force.X(), 0, 1e-12); ASSERT_NEAR(force.Y(), -5.09e-12, 1e-13); ASSERT_NEAR(force.Z(), 0, 1e-12); +} + +TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) +{ + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable( + "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + AtPropagator propagator(charge, mass, std::move(elossModel)); + + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + XYZVector startPos(0, 0, 0); // Start position in mm + XYZVector startMom(p, 0, 0); // Start momentum in MeV/c + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + XYZVector targetPoint(1e3, 0, 0); // Target point to propagate to (1 m) + propagator.PropagateToPoint(targetPoint, KE); + + auto finalPos = propagator.GetPosition(); + auto finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), 210, 10); // Final position in x-direction should be close to 210 mm + ASSERT_NEAR(finalMom.X(), 0, 0.1); + + KE = 0.5; + E = KE + mass_p; + p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + startMom.SetXYZ(p, 0, 0); // Reset momentum + propagator.SetState(startPos, startMom); + + propagator.PropagateToPoint(targetPoint, KE); // Propagate to range + finalPos = propagator.GetPosition(); + finalMom = propagator.GetMomentum(); + ASSERT_NEAR(finalPos.X(), 68.6, 5); // Final position in x-direction should be close to 68.6 mm + ASSERT_NEAR(finalMom.X(), 0, 0.1); // Final momentum in x-direction should be close to 0 +} + +TEST(AtPropagatorTest, PropagateToPoint_NoField) +{ + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable( + "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + AtPropagator propagator(charge, mass, std::move(elossModel)); + + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + XYZVector startPos(0, 0, 0); // Start position in mm + XYZVector startMom(p, 0, 0); // Start momentum in MeV/c + + double eLoss = 0.0285; // Expected energy loss in MeV (LISE) + double E_fin = KE - eLoss + mass_p; // Expected final energy after loss + double p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); // Expected final momentum in MeV/c + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + std::cout << "STARTING PROPAGATION " << std::endl << std::endl; + + XYZVector targetPoint(10, 0, 0); // Target point to propagate to 10 mm + propagator.PropagateToPoint(targetPoint); + + std::cout << "FINISHING PROPAGATION " << std::endl << std::endl; + + auto finalPos = propagator.GetPosition(); + auto finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), 10, 1); // Final position in x-direction should be close to 10 mm + ASSERT_NEAR(finalMom.X(), p_fin, 0.1); } \ No newline at end of file From 72e80879771cebe28fa8ff7ac51c4e1ddd5dcedc Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 21 Jul 2025 12:50:00 +0200 Subject: [PATCH 12/71] Update to using XYZPoint over XYZVector --- AtTools/AtPropagator.cxx | 10 ++++++---- AtTools/AtPropagator.h | 27 ++++++++++++++++++++++----- AtTools/AtPropagatorTest.cxx | 16 +++++++++------- 3 files changed, 37 insertions(+), 16 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 3728335af..1d5f639e2 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -5,7 +5,7 @@ #include namespace AtTools { -AtPropagator::XYZVector AtPropagator::Force(XYZVector pos, XYZVector mom) const +AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const { auto v = Kinematics::GetVel(mom, fMass); @@ -47,7 +47,9 @@ void AtPropagator::RK4Step() fPos = pos_SI * 1e3; // Convert back to mm } -void AtPropagator::PropagateToPoint(const XYZVector &point, double eLoss) +void AtPropagator::PropagateTo() {} + +void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) { LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; @@ -70,8 +72,8 @@ void AtPropagator::PropagateToPoint(const XYZVector &point, double eLoss) auto KE_initial = Kinematics::KE(fMom, fMass); while (true) { - XYZVector lastPos = fPos; - XYZVector lastMom = fMom; + auto lastPos = fPos; + auto lastMom = fMom; LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 066d894f4..be30a50cd 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -3,6 +3,10 @@ #include "AtELossModel.h" +#include + +#include "Math/Plane3D.h" +#include "Math/Point3D.h" #include "Math/Vector3D.h" namespace AtTools { @@ -21,6 +25,9 @@ namespace AtTools { class AtPropagator { protected: using XYZVector = ROOT::Math::XYZVector; + using XYZPoint = ROOT::Math::XYZPoint; + using Plane3D = ROOT::Math::Plane3D; + using DistanceFunc = std::function; XYZVector fEField{0, 0, 0}; // Electric field vector XYZVector fBField{0, 0, 0}; // Magnetic field vector @@ -34,7 +41,7 @@ class AtPropagator { double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped double fScalingFactor = 1.0; /// Scaling factor for energy loss - XYZVector fPos; // Current position of the particle in mm + XYZPoint fPos; // Current position of the particle in mm XYZVector fMom; // Current momentum of the particle in MeV/c static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) @@ -61,12 +68,12 @@ class AtPropagator { * @param pos Position of the particle in mm. * @param mom Momentum of the particle in MeV/c. */ - void SetState(const XYZVector &pos, const XYZVector &mom) + void SetState(const XYZPoint &pos, const XYZVector &mom) { fPos = pos; fMom = mom; } - XYZVector GetPosition() const { return fPos; } + XYZPoint GetPosition() const { return fPos; } XYZVector GetMomentum() const { return fMom; } /** @@ -78,7 +85,15 @@ class AtPropagator { * @param point The point to approach. * @param eLoss If not 0, constrain the energy loss to this value (adjusting the stopping power). */ - void PropagateToPoint(const XYZVector &point, double eLoss = 0); + void PropagateToPoint(const XYZPoint &point, double eLoss = 0); + + /** + * @brief Propagate the particle to the given plane. + * + * Propagate the particle until it reaches the specified plane, adjusting the magnitude + * of the stopping power to ensure that a specific amount of energy is lost during the propagation. + */ + void PropagateToPlane(const Plane3D &plane, double eLoss = 0); /** * @brief Calculate the force acting on the particle. @@ -87,7 +102,7 @@ class AtPropagator { * @param mom Momentum of the particle in MeV/c. * @return The force acting on the particle in N. */ - XYZVector Force(XYZVector pos, XYZVector mom) const; + XYZVector Force(XYZPoint pos, XYZVector mom) const; protected: /** @@ -97,6 +112,8 @@ class AtPropagator { * Updates fPos and fMom. */ void RK4Step(); + + void PropagateTo(); }; } // namespace AtTools diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index f3c130187..6a48aedd8 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -8,6 +8,8 @@ #include #include #include +using ROOT::Math::Plane3D; +using ROOT::Math::XYZPoint; using ROOT::Math::XYZVector; using namespace AtTools; @@ -28,7 +30,7 @@ class DummyELossModel : public AtELossModel { TEST(AtPropagatorTest, ForceNoField) { - XYZVector pos(0, 0, 0); // Position in mm + XYZPoint pos(0, 0, 0); // Position in mm XYZVector mom(100, 0, 0); // Momentum in MeV/c double charge = charge_p; // Charge in Coulombs @@ -56,7 +58,7 @@ TEST(AtPropagatorTest, ForceNoField) TEST(AtPropagatorTest, ForceEField) { - XYZVector pos(0, 0, 0); // Position in mm + XYZPoint pos(0, 0, 0); // Position in mm XYZVector mom(100, 0, 0); // Momentum in MeV/c double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 @@ -77,7 +79,7 @@ TEST(AtPropagatorTest, ForceEField) TEST(AtPropagatorTest, ForceBField) { - XYZVector pos(0, 0, 0); // Position in mm + XYZPoint pos(0, 0, 0); // Position in mm XYZVector mom(100, 0, 0); // Momentum in MeV/c double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 @@ -108,7 +110,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c - XYZVector startPos(0, 0, 0); // Start position in mm + XYZPoint startPos(0, 0, 0); // Start position in mm XYZVector startMom(p, 0, 0); // Start momentum in MeV/c propagator.SetState(startPos, startMom); @@ -117,7 +119,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - XYZVector targetPoint(1e3, 0, 0); // Target point to propagate to (1 m) + XYZPoint targetPoint(1e3, 0, 0); // Target point to propagate to (1 m) propagator.PropagateToPoint(targetPoint, KE); auto finalPos = propagator.GetPosition(); @@ -151,7 +153,7 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c - XYZVector startPos(0, 0, 0); // Start position in mm + XYZPoint startPos(0, 0, 0); // Start position in mm XYZVector startMom(p, 0, 0); // Start momentum in MeV/c double eLoss = 0.0285; // Expected energy loss in MeV (LISE) @@ -166,7 +168,7 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) std::cout << "STARTING PROPAGATION " << std::endl << std::endl; - XYZVector targetPoint(10, 0, 0); // Target point to propagate to 10 mm + XYZPoint targetPoint(10, 0, 0); // Target point to propagate to 10 mm propagator.PropagateToPoint(targetPoint); std::cout << "FINISHING PROPAGATION " << std::endl << std::endl; From 27ca18468a803fa5e881ef055c8d88372859643e Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 21 Jul 2025 12:53:41 +0200 Subject: [PATCH 13/71] Use internal scaling factor --- AtTools/AtPropagator.cxx | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 1d5f639e2..0d9190f8e 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -53,12 +53,11 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) { LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; - double scalingFactor = 1.0; int iterations = 0; double calc_eLoss = 0; while (std::abs(calc_eLoss - eLoss) > 1e-4 || eLoss == 0) { - LOG(debug) << "Running iteration " << iterations << " with scaling factor: " << scalingFactor + LOG(debug) << "Running iteration " << iterations << " with scaling factor: " << fScalingFactor << " and energy loss: " << calc_eLoss; if (iterations > 100) { @@ -102,7 +101,7 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) double KE_final = Kinematics::KE(fMom, fMass); calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - scalingFactor *= eLoss / calc_eLoss; + fScalingFactor *= eLoss / calc_eLoss; LOG(info) << "------- End of RK4 interation " << iterations << " ---------"; LOG(info) << "Particle stopped: " << particleStopped; LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; @@ -110,15 +109,19 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) LOG(info) << "Desired energy loss: " << eLoss << " MeV"; LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; LOG(info) << "Difference: " << calc_eLoss - eLoss << " MeV"; - LOG(info) << "New scaling factor: " << scalingFactor; + LOG(info) << "New scaling factor: " << fScalingFactor; LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - if (eLoss == 0) - return; // If no energy loss is specified, we are done. - break; // Else rerun with adjusted scaling factor + if (eLoss == 0) { + fScalingFactor = 1; // Reset scaling factor after convergence + return; // If no energy loss is specified, we are done. + } + break; // Else rerun with adjusted scaling factor } } // End of loop over RK4 integration } // End loop over energy loss convergence + + fScalingFactor = 1; // Reset scaling factor after convergence } } // namespace AtTools From ef1dec0b342fd1e3e4f7574a0668720355d197b0 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 00:13:31 +0200 Subject: [PATCH 14/71] Fixed step propagator stopping in correct place --- AtTools/AtKinematics.cxx | 6 + AtTools/AtKinematics.h | 1 + AtTools/AtPropagator.cxx | 498 ++++++++++++++++++++++++++++++----- AtTools/AtPropagator.h | 104 +++++++- AtTools/AtPropagatorTest.cxx | 84 +++++- 5 files changed, 609 insertions(+), 84 deletions(-) diff --git a/AtTools/AtKinematics.cxx b/AtTools/AtKinematics.cxx index 755baf91c..3f9d416f7 100644 --- a/AtTools/AtKinematics.cxx +++ b/AtTools/AtKinematics.cxx @@ -303,6 +303,12 @@ double EtoA(double mass) { return mass / 931.5; } +double GetSpeed(double p, double mass) +{ + // Calculate the speed of a particle given its momentum and mass in m/s + double beta = GetBeta(p, mass); + return beta * fC; // Speed in m/s +} ROOT::Math::XYZVector GetVel(ROOT::Math::XYZVector mom, double mass) { return mom / Get4Vector(mom, mass).E() * fC; diff --git a/AtTools/AtKinematics.h b/AtTools/AtKinematics.h index 8be5c0111..622f347d9 100644 --- a/AtTools/AtKinematics.h +++ b/AtTools/AtKinematics.h @@ -76,6 +76,7 @@ double GetBeta(double p, double mass); double GetRelMom(double gamma, double mass); double AtoE(double Amu); double EtoA(double mass); +double GetSpeed(double p, double mass); /// Calculate the speed of a particle given its momentum and mass in m/s /** * Calculate the kinetic energy of a particle given its momentum and mass. diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 0d9190f8e..0d6bef446 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -3,6 +3,29 @@ #include "AtKinematics.h" #include + +// Butcher tableau (c, a_ij) and the two b vectors for Dormand–Prince 5(4) +// c1 = 0 +// Butcher tableau coefficients for Dormand–Prince 5(4) method + +static constexpr double c[7] = {0.0, 1.0 / 5.0, 3.0 / 10.0, 4.0 / 5.0, 8.0 / 9.0, 1.0, 1.0}; + +static constexpr double a[7][6] = { + {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, + {1.0 / 5.0, 0.0, 0.0, 0.0, 0.0, 0.0}, + {3.0 / 40.0, 9.0 / 40.0, 0.0, 0.0, 0.0, 0.0}, + {44.0 / 45.0, -56.0 / 15.0, 32.0 / 9.0, 0.0, 0.0, 0.0}, + {19372.0 / 6561.0, -25360.0 / 2187.0, 64448.0 / 6561.0, -212.0 / 729.0, 0.0, 0.0}, + {9017.0 / 3168.0, -355.0 / 33.0, 46732.0 / 5247.0, 49.0 / 176.0, -5103.0 / 18656.0, 0.0}, + {35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0, 11.0 / 84.0}}; + +// b (5th-order) +static constexpr double b[7] = {35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0, 11.0 / 84.0, 0.0}; + +// b* (4th-order, “star”) +static constexpr double bs[7] = {5179.0 / 57600.0, 0.0, 7571.0 / 16695.0, 393.0 / 640.0, -92097.0 / 339200.0, + 187.0 / 2100.0, 1.0 / 40.0}; + namespace AtTools { AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const @@ -20,26 +43,241 @@ AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const return F_lorentz + drag; // Force in N } -void AtPropagator::RK4Step() +AtPropagator::XYZVector AtPropagator::dpds(const XYZPoint &pos, const XYZVector &mom) const +{ + // Calculate the force acting on the particle at the given position and momentum + auto speed = Kinematics::GetSpeed(mom.R(), fMass); // Speed in m/s + return Force(pos, mom) / speed; +} +AtPropagator::XYZVector AtPropagator::d2xds2(const XYZPoint &pos, const XYZVector &mom) const +{ + auto phat = mom.Unit(); // Unit vector in the direction of momentum + auto p = mom.R(); // Magnitude of the momentum + auto dpds_vec = dpds(pos, mom); // Derivative of momentum w.r.t. arc length + + return 1 / p * (dpds_vec - phat * (phat.Dot(dpds_vec))); // Second derivative of position w.r.t. arc length +} + +void AtPropagator::PropagateToPointAdaptive(const XYZPoint &point) { - double h = fH; // Step size in seconds + LOG(info) << "Propagating to point: " << point; + + auto KE_initial = Kinematics::KE(fMom, fMass); + while (true) { + fLastPos = fPos; + fLastMom = fMom; + LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + + auto acceptedStep = RK4StepAdaptive(fH); + if (!acceptedStep) { + LOG(error) << "RK4 step failed, aborting propagation."; + return; // Abort propagation if step failed + } + + if (!acceptedStep || !ReachedPOCA(point)) + continue; + + bool reachedMeasurementPoint = ReachedPOCA(point); + bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; + if (reachedMeasurementPoint || particleStopped) { + // Last iteration we were still approaching the measurement point. Now we are further away + // then before. We have probably reached the measurement point if things are well behaved. + // I can think of cases where this will not be true. A better solution might be to run + // tracking the point of closest approach until the distance between the current state and + // the measurement point is larger than the distance between the last state and the measurement point. + + // Undo the last step since we were closer last time. + double lastApproach = (fLastPos - point).R(); + double approach = (fPos - point).R(); + fPos = fLastPos; + fMom = fLastMom; + + double KE_final = Kinematics::KE(fMom, fMass); + auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + LOG(info) << "------- End of RK4 interation ---------"; + LOG(info) << "Particle stopped: " << particleStopped; + LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; + LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; + LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(info) << "Scaling factor: " << fScalingFactor; + LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + return; + } + } // End of loop over RK4 integration +} + +bool AtPropagator::RK4StepAdaptive(double &h) +{ + // Take h to be the step size in m. + // Use DP5(4) method for adaptive step size control. + + double atol_pos = 1e-2; // Absolute tolerance for position (mm) + double atol_mom = 1e-2; // Absolute tolerance for momentum (MeV/c) + double rtol = 1e-4; // Relative tolerance for both position and momentum - auto x_k1 = Kinematics::GetVel(fMom, fMass); - auto p_k1 = Force(fPos, fMom); + auto x0_mm = fPos; + auto p0 = fMom; + LOG(info) << "Starting RK4 step with initial position: " << x0_mm.X() << ", " << x0_mm.Y() << ", " << x0_mm.Z(); + LOG(info) << "Initial momentum: " << p0.X() << ", " << p0.Y() << ", " << p0.Z(); - auto x_k2 = Kinematics::GetVel(fMom + p_k1 * h / 2, fMass); - auto p_k2 = Force(fPos + x_k1 * h / 2, fMom + p_k1 * h / 2); + while (true) { + auto x_SI = fPos * 1e-3; // Convert position to SI units (m) + auto p_SI = fReltoSImom * fMom; // Convert momentum to SI units (kg m/s) + XYZVector kx[7]; // kx[i] will hold the position derivatives (unitless) + XYZVector kp[7]; // kp[i] will hold the momentum derivatives (SI units) - auto x_k3 = Kinematics::GetVel(fMom + p_k2 * h / 2, fMass); - auto p_k3 = Force(fPos + x_k2 * h / 2, fMom + p_k2 * h / 2); + // anonymous lambda to calculate and store the kx and kp values. Input is SI units. + auto calc_k = [&](const XYZPoint &x, const XYZVector &p, int i) { + kx[i] = p.Unit(); // The derivative of the position is then just the unit vector of the momentum. + kp[i] = dpds(x * 1e-3, p / fReltoSImom); // The derivative of the momentum is dpds. + }; - auto x_k4 = Kinematics::GetVel(fMom + p_k3 * h, fMass); - auto p_k4 = Force(fPos + x_k3 * h, fMom + p_k3 * h); + // anonymous lambda to calculate the position and momentum at the i-th stage + auto calc_xp = [&](int i) { + XYZVector dx(0, 0, 0); + XYZVector dp(0, 0, 0); + for (int j = 0; j < i; ++j) { + dx = dx + kx[j] * a[i][j]; + dp = dp + kp[j] * a[i][j]; + } + XYZPoint x = x_SI + dx * h; + XYZVector p = p_SI + dp * h; + return std::make_pair(x, p); + }; + + // Calculate kx and kp for each stage + // build stage 0 + calc_k(x_SI, p0, 0); + + // build stage 1 + auto [x1, p1] = calc_xp(1); + calc_k(x1, p1, 1); // k1 + + // build stage 2 + auto [x2, p2] = calc_xp(2); + calc_k(x2, p2, 2); // k2 + + // build stage 3 + auto [x3, p3] = calc_xp(3); + calc_k(x3, p3, 3); // k3 + + // build stage 4 + auto [x4, p4] = calc_xp(4); + calc_k(x4, p4, 4); // k4 + + // build stage 5 + auto [x5, p5] = calc_xp(5); + calc_k(x5, p5, 5); // k5 + + // build stage 6 + auto [x6, p6] = calc_xp(6); + calc_k(x6, p6, 6); // k6 + + // Calculate the new position and momentum using the 5th-order method + XYZVector dx(0, 0, 0); + XYZVector dp(0, 0, 0); + for (int i = 0; i < 7; ++i) { + dx = dx + kx[i] * b[i]; + dp = dp + kp[i] * b[i]; + } + XYZPoint x_new_5 = x_SI + dx * h; // New position in SI units (m) + XYZVector p_new_5 = p_SI + dp * h; // New momentum in SI units (kg m/s) + + // Calculate the new position and momentum using the 4th-order method + dx = XYZVector(0, 0, 0); + dp = XYZVector(0, 0, 0); + for (int i = 0; i < 7; ++i) { + dx = dx + kx[i] * bs[i]; + dp = dp + kp[i] * bs[i]; + } + XYZPoint x_new_4 = x_SI + dx * h; // New position in SI units (m) + XYZVector p_new_4 = p_SI + dp * h; // New momentum in SI units (kg m/s) - auto F_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // Force in SI units (N) + auto x_4_mm = x_new_4 * 1e3; // Convert back to mm + auto p_4_MeV = p_new_4 / fReltoSImom; // Convert back to MeV/c + auto x_5_mm = x_new_5 * 1e3; // Convert back to mm + auto p_5_MeV = p_new_5 / fReltoSImom; // Convert back to MeV/c + LOG(info) << "New position (5th order): " << x_5_mm.X() << ", " << x_5_mm.Y() << ", " << x_5_mm.Z(); + LOG(info) << "New momentum (5th order): " << p_5_MeV.X() << ", " << p_5_MeV.Y() << ", " << p_5_MeV.Z(); + LOG(info) << "New position (4th order): " << x_4_mm.X() << ", " << x_4_mm.Y() << ", " << x_4_mm.Z(); + LOG(info) << "New momentum (4th order): " << p_4_MeV.X() << ", " << p_4_MeV.Y() << ", " << p_4_MeV.Z(); + + // Convert back to mm and MeV/c + XYZVector x_err = (x_5_mm - x_4_mm); // Error in position (mm) + XYZVector p_err = (p_5_MeV - p_4_MeV); // Error in momentum (MeV/c) + + // Calculate the overall error + double ex = x_err.X() / (atol_pos + rtol * std::abs(x_5_mm.X())); + double ey = x_err.Y() / (atol_pos + rtol * std::abs(x_5_mm.Y())); + double ez = x_err.Z() / (atol_pos + rtol * std::abs(x_5_mm.Z())); + + double ep_x = p_err.X() / (atol_mom + rtol * std::abs(p_5_MeV.X())); + double ep_y = p_err.Y() / (atol_mom + rtol * std::abs(p_5_MeV.Y())); + double ep_z = p_err.Z() / (atol_mom + rtol * std::abs(p_5_MeV.Z())); + + // Combine errors (norm) + double err = std::sqrt(ex * ex + ey * ey + ez * ez + ep_x * ep_x + ep_y * ep_y + ep_z * ep_z); + + double factor = std::pow(err, -1.0 / 5.0); // Adjust step size based on error + factor = std::clamp(factor, 0.25, 4.0); // Clamp factor to reasonable limits + double hNew = h * factor; + // We now know the local error at this point. Now we need to decide to accept the point or not. + if (err <= 1.0) { + // Accept the step + fPos = x_5_mm; // Update position in mm + fMom = p_5_MeV; // Update momentum in MeV/c + LOG(info) << "Accepted step with error: " << err; + LOG(info) << "Step size: " << h << " m"; + LOG(info) << "New step size: " << hNew << " m"; + LOG(info) << "New Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "New Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + + // Adjust the step size for the next iteration + h = hNew; + return true; // Step accepted + } else { + // Reject the step and reduce the step size + LOG(info) << "Rejected step with error: " << err; + LOG(info) << "Step size: " << h << " m"; + LOG(info) << "Reducing step size to: " << hNew << " m"; + + h = hNew; // Reduce step size + if (h < 1e-6) { + LOG(error) << "Step size too small, aborting propagation."; + return false; // Abort propagation if step size is too small + } + } + } +} + +void AtPropagator::RK4Step(double h) +{ + // Take h to be the step size in m. + + auto x_k1 = fMom.Unit(); // The derivative of the position is then just the unit vector of the momentum. + auto p_k1 = dpds(fPos, fMom); // The derivative of the momentum is dpds. + + auto x_2 = fPos + x_k1 * h / 2; // Position at the midpoint + auto p_2 = fMom + p_k1 * h / 2; // Momentum at the midpoint + auto x_k2 = p_2.Unit(); + auto p_k2 = dpds(x_2, p_2); + + auto x_3 = fPos + x_k2 * h / 2; // Position at the second midpoint + auto p_3 = fMom + p_k2 * h / 2; // Momentum at the second midpoint + auto x_k3 = p_3.Unit(); + auto p_k3 = dpds(x_3, p_3); + + auto x_4 = fPos + x_k3 * h; // Position at the end of the step + auto p_4 = fMom + p_k3 * h; // Momentum at the end of the step + auto x_k4 = p_4.Unit(); + auto p_k4 = dpds(x_4, p_4); + + auto dpds_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // "Force" in SI units (N) auto mom_SI = fReltoSImom * fMom; - mom_SI += F_SI * h; // Update momentum in SI units (kg m/s) + mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) fMom = mom_SI / fReltoSImom; // Convert back to auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) @@ -47,16 +285,181 @@ void AtPropagator::RK4Step() fPos = pos_SI * 1e3; // Convert back to mm } -void AtPropagator::PropagateTo() {} +bool AtPropagator::ReachedPOCA(const XYZPoint &point) +{ + // Here we need to check if we are getting closer or further away from the POCA. + // We may walk right past it so need to look for a change in the sign of the derivative or + // something like that. + auto lastDeriv = (fLastPos - point).Dot(fLastMom.Unit()); // proportional missing constants + auto currDeriv = (fPos - point).Dot(fMom.Unit()); + LOG(debug) << "Last Derivative: " << lastDeriv << ", Current Derivative: " << currDeriv; + return lastDeriv * currDeriv < 0; + + auto lastApproach = (fLastPos - point).R(); + auto approach = (fPos - point).R(); + return (approach > lastApproach); +} + +bool AtPropagator::IntersectedPlane(const Plane3D &plane) +{ + // Check if the particle has crossed the plane this step. + auto prevSign = plane.Distance(fLastPos) > 0 ? 1 : -1; + auto currSign = plane.Distance(fPos) > 0 ? 1 : -1; + return (prevSign != currSign); +} + +void AtPropagator::PropagateToPlane(const Plane3D &plane) +{ + LOG(info) << "Propagating to plane: " << plane; + + auto KE_initial = Kinematics::KE(fMom, fMass); + while (true) { + fLastPos = fPos; + fLastMom = fMom; + LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + + RK4Step(fH); + + if (!IntersectedPlane(plane)) + continue; + + bool reachedMeasurementPoint = IntersectedPlane(plane); + bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; + if (reachedMeasurementPoint || particleStopped) { + // Last iteration we were still approaching the measurement point. Now we are further away + // then before. We have probably reached the measurement point if things are well behaved. + // I can think of cases where this will not be true. A better solution might be to run + // tracking the point of closest approach until the distance between the current state and + // the measurement point is larger than the distance between the last state and the measurement point. + + // Undo the last step since we were closer last time. + double lastApproach = plane.Distance(fLastPos); + double approach = plane.Distance(fPos); + fPos = fLastPos; + fMom = fLastMom; + + double KE_final = Kinematics::KE(fMom, fMass); + auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + LOG(info) << "------- End of RK4 interation ---------"; + LOG(info) << "Particle stopped: " << particleStopped; + LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; + LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; + LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(info) << "Scaling factor: " << fScalingFactor; + LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + return; + } + } // End of loop over RK4 integration +} + +void AtPropagator::PropagateToPoint(const XYZPoint &point) +{ + LOG(info) << "Propagating to point: " << point; + + auto KE_initial = Kinematics::KE(fMom, fMass); + while (true) { + fLastPos = fPos; + fLastMom = fMom; + LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + + RK4Step(fH); + + if (!ReachedPOCA(point)) + continue; + + bool reachedMeasurementPoint = ReachedPOCA(point); + bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; + bool momentumReversed = (fLastMom.Dot(fMom) < 0); + + // If we stopped, then we should figure out about where we stopped assuming linear de/dx over this last step + if (particleStopped || momentumReversed) { + LOG(info) << "------ Particle stopped ------"; + + double finalH = (fPos - fLastPos).R(); // Distance traveled in the last step + double E_loss = + Kinematics::KE(fLastMom, fMass) + Kinematics::KE(fMom, fMass); // Energy loss in MeV in the last step + double dedx = E_loss / finalH; // Stopping power in MeV/mm + + LOG(info) << "Particle stopped with final step size: " << finalH << " mm"; + LOG(info) << "Energy loss in last step: " << E_loss << " MeV"; + LOG(info) << "Stopping power (dE/dx): " << dedx << " MeV/mm"; + LOG(info) << "Energy before stopping: " << Kinematics::KE(fLastMom, fMass) << " MeV"; + finalH = Kinematics::KE(fLastMom, fMass) / dedx; // Distance to stop in mm + LOG(info) << "Estimated distance to stop: " << finalH << " mm"; + + fPos = fLastPos; + fMom = fLastMom; // Reset to last position and momentum + RK4Step(finalH); // Propagate to the point where we stopped + fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped + fLastMom = fMom; + + LOG(info) << "Final Position after stopping: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + } + + if (reachedMeasurementPoint) { + // We reached the measurement point, so we should figure out how far we are from the measurement point + // and update that remaining amount + LOG(info) << "------ Reached measurement point------"; + double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step + double approach = (fLastPos - point).R(); // Distance to the measurement point + LOG(info) << "Distance to measurement point: " << approach << " mm"; + LOG(info) << "Final step size: " << finalH << " mm"; + + finalH = approach * 1e-3; // Convert to meters for the RK4 step + fPos = fLastPos; + fMom = fLastMom; + RK4Step(finalH); // Propagate to the measurement point + } + + if (reachedMeasurementPoint || particleStopped) { + // Undo the last step since we were closer last time. + double lastApproach = (fLastPos - point).R(); + double approach = (fPos - point).R(); + + double KE_final = Kinematics::KE(fMom, fMass); + auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + LOG(info) << "------- End of RK4 interation ---------"; + LOG(info) << "Particle stopped: " << particleStopped; + LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; + LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; + LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(info) << "Scaling factor: " << fScalingFactor; + LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(info) << "Last Position: " << fLastPos.X() << ", " << fLastPos.Y() << ", " << fLastPos.Z(); + LOG(info) << "Last Momentum: " << fLastMom.X() << ", " << fLastMom.Y() << ", " << fLastMom.Z(); + + // fPos = fLastPos; + // fMom = fLastMom; + return; + } + } // End of loop over RK4 integration +} void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) { LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; + if (eLoss == 0) { + LOG(warn) << "No energy loss specified, propagating without energy loss adjustment."; + PropagateToPoint(point); + return; + } + int iterations = 0; double calc_eLoss = 0; + double KE_initial = Kinematics::KE(fMom, fMass); + auto initialMom = fMom; // Save initial momentum for energy loss calculation + auto initialPos = fPos; // Save initial position for energy loss calculation + + while (std::abs(calc_eLoss - eLoss) > 1e-4) { + fMom = initialMom; // Reset position and momentum to initial values for the next iteration + fPos = initialPos; - while (std::abs(calc_eLoss - eLoss) > 1e-4 || eLoss == 0) { LOG(debug) << "Running iteration " << iterations << " with scaling factor: " << fScalingFactor << " and energy loss: " << calc_eLoss; @@ -65,62 +468,21 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) throw std::runtime_error("Energy loss did not converge after 100 iterations."); } - double lastApproach = std::numeric_limits::max(); - bool approaching = true; iterations++; - auto KE_initial = Kinematics::KE(fMom, fMass); - - while (true) { - auto lastPos = fPos; - auto lastMom = fMom; - LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - - RK4Step(); - - auto approach = (fPos - point).R(); - if (approach < lastApproach) { - // We are still approaching the measurement point - approaching = true; - lastApproach = approach; - continue; - } + PropagateToPoint(point); // Propagate without energy loss adjustment - bool reachedMeasurementPoint = (approaching && approach > lastApproach); - bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - if (reachedMeasurementPoint || particleStopped) { - // Last iteration we were still approaching the measurement point. Now we are further away - // then before. We have probably reached the measurement point if things are well behaved. - // I can think of cases where this will not be true. A better solution might be to run - // tracking the point of closest approach until the distance between the current state and - // the measurement point is larger than the distance between the last state and the measurement point. - - // Undo the last step since we were closer last time. - fPos = lastPos; - fMom = lastMom; - - double KE_final = Kinematics::KE(fMom, fMass); - calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - fScalingFactor *= eLoss / calc_eLoss; - LOG(info) << "------- End of RK4 interation " << iterations << " ---------"; - LOG(info) << "Particle stopped: " << particleStopped; - LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; - LOG(info) << "Desired energy loss: " << eLoss << " MeV"; - LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; - LOG(info) << "Difference: " << calc_eLoss - eLoss << " MeV"; - LOG(info) << "New scaling factor: " << fScalingFactor; - LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - - if (eLoss == 0) { - fScalingFactor = 1; // Reset scaling factor after convergence - return; // If no energy loss is specified, we are done. - } - break; // Else rerun with adjusted scaling factor - } - } // End of loop over RK4 integration - } // End loop over energy loss convergence + double KE_final = Kinematics::KE(fMom, fMass); + calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + fScalingFactor *= eLoss / calc_eLoss; + LOG(info) << "Desired energy loss: " << eLoss << " MeV"; + LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(info) << "Difference: " << calc_eLoss - eLoss << " MeV"; + LOG(info) << "New scaling factor: " << fScalingFactor; + LOG(info) << "Condition: " << (std::abs(calc_eLoss - eLoss) > 1e-4); + + } // End loop over energy loss convergence + + LOG(info) << "Energy loss converged after " << iterations << " iterations."; fScalingFactor = 1; // Reset scaling factor after convergence } diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index be30a50cd..01b916fa1 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -36,14 +36,19 @@ class AtPropagator { const std::unique_ptr fELossModel; // Energy loss model // Internal state variables for the propagator - double fH = 1e-10; /// Step size for propagation in s + double fH = 1e-4; /// Step size for propagation in m + double fDelta = 1e-3; /// Relative error tolerance for adaptive step size. 10^-3 means each 1m of propagation + /// introduces at most 1mm of error. double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss - double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped + double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped (MeV) double fScalingFactor = 1.0; /// Scaling factor for energy loss XYZPoint fPos; // Current position of the particle in mm XYZVector fMom; // Current momentum of the particle in MeV/c + XYZPoint fLastPos; + XYZVector fLastMom; + static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) public: @@ -56,6 +61,7 @@ class AtPropagator { */ void SetEField(const XYZVector &eField) { fEField = eField; } void SetEField(double ex, double ey, double ez) { fEField.SetXYZ(ex, ey, ez); } + /** * @brief Set the magnetic field (T) */ @@ -73,6 +79,10 @@ class AtPropagator { fPos = pos; fMom = mom; } + + void SetDelta(double delta) { fDelta = delta; } + void SetH(double h) { fH = h; } + XYZPoint GetPosition() const { return fPos; } XYZVector GetMomentum() const { return fMom; } @@ -83,17 +93,24 @@ class AtPropagator { * to ensure that a specific about of energy is lost during the propagation. * * @param point The point to approach. - * @param eLoss If not 0, constrain the energy loss to this value (adjusting the stopping power). + * @param eLoss If not 0, constrain the energy loss to this value by adjusting fScalingFactor. + */ + void PropagateToPoint(const XYZPoint &point, double eLoss); + + /** + * @brief Propagate the particle to the point of closest approach to the given point. + * + * @param point The point to approach. */ - void PropagateToPoint(const XYZPoint &point, double eLoss = 0); + void PropagateToPoint(const XYZPoint &point); + void PropagateToPointAdaptive(const XYZPoint &point); /** * @brief Propagate the particle to the given plane. * - * Propagate the particle until it reaches the specified plane, adjusting the magnitude - * of the stopping power to ensure that a specific amount of energy is lost during the propagation. + * @param plane The plane to approach. */ - void PropagateToPlane(const Plane3D &plane, double eLoss = 0); + void PropagateToPlane(const Plane3D &plane); /** * @brief Calculate the force acting on the particle. @@ -104,17 +121,86 @@ class AtPropagator { */ XYZVector Force(XYZPoint pos, XYZVector mom) const; + /** + * @brief Calculate the derivate of the momentum w.r.t. arc length. + * + * @param pos Position of the particle in mm. + * @param mom Momentum of the particle in MeV/c. + * @return The derivative of the momentum w.r.t. arc length in N/m. + */ + XYZVector dpds(const XYZPoint &pos, const XYZVector &mom) const; + + /** + * @brief Calculate the second derivative of the position w.r.t. arc length. + * + * \frac{d^2\vec{x}}{ds^2} = \frac{1}{p} \left( \frac{d\vec{p}}{ds} - \hat{p} (\hat{p} \cdot \frac{d\vec{p}}{ds}) + * \right) + * + * @param pos Position of the particle in mm. + * @param mom Momentum of the particle in MeV/c. + * @return The second derivative of the position w.r.t. arc length in m/m^2. + */ + XYZVector d2xds2(const XYZPoint &pos, const XYZVector &mom) const; + + XYZVector dxds(const XYZPoint &pos, const XYZVector &mom) const + { + return mom.Unit(); // The derivative of the position is just the unit vector of the momentum. + } + protected: /** * @brief Perform a single RK4 step for propagation. * * This method performs a single Runge-Kutta 4th order step to propagate the particle's state. * Updates fPos and fMom. + * @param h Step size for the RK4 step in meters. + */ + void RK4Step(double h); + + /** + * @brief Perform a single RK4 step using the Nystrom method. + * This method performs a single Runge-Kutta 4th order step using the Nystrom method + * to propagate the particle's state. Updates fPos and fMom. + * @param h Step size for the RK4 step in meters. */ - void RK4Step(); + void RK4StepNystrom(double h); - void PropagateTo(); + /** + * @brief Perform an adaptive RK4 step for propagation. + * This method performs an adaptive Runge-Kutta 4th order step to propagate the particle's state. + * Updates fPos and fMom. + * + * The error is based on the difference in the positions at the end of the step. i.e: + * \eps = \sqrt{\eps_x^2 + \eps_y^2 + \eps_z^2}, where + * \eps_x = 1/30*|x_1 - x_2|, where x_1 is using h/2 and x_2 is using h. + * + * Step size is adjust to ensure the local error is less than fDelta. + * + * @param h Step size for the RK4 step in seconds. Modified in place to reflect the new step size. + * @return True if the step was accepted, false otherwise. + */ + bool RK4StepAdaptive(double &h); + + void PropagateTo(DistanceFunc distanceFunc); + + bool ReachedPOCA(const XYZPoint &point); + bool IntersectedPlane(const Plane3D &plane); }; +static constexpr double a21 = 1.0 / 5.0; +static constexpr double a31 = 3.0 / 40.0, a32 = 9.0 / 40.0; +static constexpr double a41 = 44.0 / 45.0, a42 = -56.0 / 15.0, a43 = 32.0 / 9.0; +static constexpr double a51 = 19372.0 / 6561.0, a52 = -25360.0 / 2187.0, a53 = 64448.0 / 6561.0, a54 = -212.0 / 729.0; +static constexpr double a61 = 9017.0 / 3168.0, a62 = -355.0 / 33.0, a63 = 46732.0 / 5247.0, a64 = 49.0 / 176.0, + a65 = -5103.0 / 18656.0; +static constexpr double a71 = 35.0 / 384.0, a72 = 0.0, a73 = 500.0 / 1113.0, a74 = 125.0 / 192.0, + a75 = -2187.0 / 6784.0, a76 = 11.0 / 84.0; +// b (5th-order) +static constexpr double b1 = 35.0 / 384.0, b3 = 500.0 / 1113.0, b4 = 125.0 / 192.0, b5 = -2187.0 / 6784.0, + b6 = 11.0 / 84.0; +// b* (4th-order, “star”) +static constexpr double bs1 = 5179.0 / 57600.0, bs3 = 7571.0 / 16695.0, bs4 = 393.0 / 640.0, bs5 = -92097.0 / 339200.0, + bs6 = 187.0 / 2100.0, bs7 = 1.0 / 40.0; + } // namespace AtTools #endif // #ifndef ATPROPAGATOR_H diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 6a48aedd8..2e831836c 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -120,7 +120,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); XYZPoint targetPoint(1e3, 0, 0); // Target point to propagate to (1 m) - propagator.PropagateToPoint(targetPoint, KE); + propagator.PropagateToPoint(targetPoint); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -128,16 +128,16 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) ASSERT_NEAR(finalPos.X(), 210, 10); // Final position in x-direction should be close to 210 mm ASSERT_NEAR(finalMom.X(), 0, 0.1); - KE = 0.5; + KE = 0.75; E = KE + mass_p; p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c startMom.SetXYZ(p, 0, 0); // Reset momentum propagator.SetState(startPos, startMom); - propagator.PropagateToPoint(targetPoint, KE); // Propagate to range + propagator.PropagateToPoint(targetPoint); // Propagate to range finalPos = propagator.GetPosition(); finalMom = propagator.GetMomentum(); - ASSERT_NEAR(finalPos.X(), 68.6, 5); // Final position in x-direction should be close to 68.6 mm + ASSERT_NEAR(finalPos.X(), 130, 10); // Final position in x-direction should be close to 130 mm ASSERT_NEAR(finalMom.X(), 0, 0.1); // Final momentum in x-direction should be close to 0 } @@ -166,16 +166,86 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - std::cout << "STARTING PROPAGATION " << std::endl << std::endl; - XYZPoint targetPoint(10, 0, 0); // Target point to propagate to 10 mm propagator.PropagateToPoint(targetPoint); - std::cout << "FINISHING PROPAGATION " << std::endl << std::endl; + auto finalPos = propagator.GetPosition(); + auto finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), 10, 1); // Final position in x-direction should be close to 10 mm + ASSERT_NEAR(finalMom.X(), p_fin, 0.01); +} + +TEST(AtPropagatorTest, PropagateToPlane_NoField) +{ + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable( + "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + AtPropagator propagator(charge, mass, std::move(elossModel)); + + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + XYZPoint startPos(0, 0, 0); // Start position in mm + XYZVector startMom(p, 0, 0); // Start momentum in MeV/c + + double eLoss = 0.0285; // Expected energy loss in MeV in 10 mm (LISE) + double E_fin = KE - eLoss + mass_p; // Expected final energy after loss + double p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); // Expected final momentum in MeV/c + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + XYZPoint planePoint(10, 10, 10); // Target point to propagate to 10 mm + XYZVector planeNormal(1, 0, 0); // Normal vector of the plane in x-direction + Plane3D plane(planeNormal, planePoint); // Create the plane + propagator.PropagateToPlane(plane); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); ASSERT_NEAR(finalPos.X(), 10, 1); // Final position in x-direction should be close to 10 mm ASSERT_NEAR(finalMom.X(), p_fin, 0.1); +} + +TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) +{ + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable( + "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + AtPropagator propagator(charge, mass, std::move(elossModel)); + + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + XYZPoint startPos(0, 0, 0); // Start position in mm + XYZVector startMom(p, 0, 0); // Start momentum in MeV/c + + double eLoss = 0.0285; // Expected energy loss in MeV in 10 mm (LISE) + double E_fin = KE - eLoss + mass_p; // Expected final energy after loss + double p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); // Expected final momentum in MeV/c + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + propagator.SetDelta(1e-3); // Set relative error tolerance. Traveling 10 mm means at most 10 microns of error. + propagator.SetH(1); // Set initial step size to 1 s + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + XYZPoint targetPoint(10, 0, 0); // Target point to propagate to (10 mm) + propagator.PropagateToPointAdaptive(targetPoint); + + auto finalPos = propagator.GetPosition(); + auto finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), 10, 10 * 1e-3); // Final position in x-direction should be close to 10 mm + ASSERT_NEAR(finalMom.X(), p_fin, 0.1); } \ No newline at end of file From 57ea952cbe149b33874c32511ffa136a337118af Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 01:22:10 +0200 Subject: [PATCH 15/71] Non-Adaptive tests working --- AtTools/AtPropagator.cxx | 107 +++++++++++++++++++++++++---------- AtTools/AtPropagator.h | 1 + AtTools/AtPropagatorTest.cxx | 33 +++++++++++ 3 files changed, 110 insertions(+), 31 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 0d6bef446..376588409 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -316,38 +316,83 @@ void AtPropagator::PropagateToPlane(const Plane3D &plane) while (true) { fLastPos = fPos; fLastMom = fMom; - LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(info) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); RK4Step(fH); - if (!IntersectedPlane(plane)) - continue; - bool reachedMeasurementPoint = IntersectedPlane(plane); bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - if (reachedMeasurementPoint || particleStopped) { - // Last iteration we were still approaching the measurement point. Now we are further away - // then before. We have probably reached the measurement point if things are well behaved. - // I can think of cases where this will not be true. A better solution might be to run - // tracking the point of closest approach until the distance between the current state and - // the measurement point is larger than the distance between the last state and the measurement point. + bool momentumReversed = (fLastMom.Dot(fMom) < 0); - // Undo the last step since we were closer last time. - double lastApproach = plane.Distance(fLastPos); - double approach = plane.Distance(fPos); + if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { + // We reached the measurement point, so we should figure out how far we are from the measurement point + LOG(info) << "------ Reached measurement point ------"; + double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step + double approach = std::abs(plane.Distance(fLastPos)); + + LOG(info) << "Distance to plane: " << approach << " mm"; + LOG(info) << "Final step size: " << finalH << " mm"; + + finalH = approach * 1e-3; // Convert to meters for the RK4 step fPos = fLastPos; fMom = fLastMom; + RK4Step(finalH); // Propagate to the measurement point + } + + if (particleStopped || momentumReversed) { + // In this case the particle stopped before hitting the plane + // we should throw a warning to let the user know that there wasn't + // enough energy to reach the plane. + LOG(warning) << "------ Particle stopped before intersecting plane ------"; + + // Calculate how far to travel before stopping + double KE_last = Kinematics::KE(fLastMom, fMass); + double deltaE = KE_last - fStopTol; + deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss + + LOG(info) << "Last KE: " << KE_last << " MeV"; + LOG(info) << "Energy to loose to stop: " << deltaE << " MeV"; + + double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm + RK4Step(h_Stop); + LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; + + while (!IntersectedPlane(plane)) { + fScalingFactor = 0; // Turn off enregy loss. + + // If we still haven't intersected the plane, we need to adjust the step size + double h = std::abs(plane.Distance(fPos)); // Reduce step size so we hit the plane + if (h <= fDistTol) + break; + LOG(info) << "Propagating to plane after stopping with step size: " << h << " mm"; + RK4Step(h * 1e-3); // Convert mm to m for RK4 step + LOG(info) << "New position after adjusting step size: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + } + fLastMom = fMom; + fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped + reachedMeasurementPoint = true; + } + + if (reachedMeasurementPoint || particleStopped || momentumReversed) { + double distanceToPlane = std::abs(plane.Distance(fPos)); double KE_final = Kinematics::KE(fMom, fMass); auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV LOG(info) << "------- End of RK4 interation ---------"; LOG(info) << "Particle stopped: " << particleStopped; LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; + LOG(info) << "Distance to plane: " << distanceToPlane << " mm"; LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; LOG(info) << "Scaling factor: " << fScalingFactor; LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + + // Project the position onto the plane. Cannot use ProjectOnPlane since it is templated in such + // a way that it can't separate our XYZPoint and its internal XYZPoint. + double d = plane.Distance(fPos); // Distance from the point to the plane + fPos = XYZPoint(fPos.X() - plane.A() * d, fPos.Y() - plane.B() * d, fPos.Z() - plane.C() * d); + LOG(info) << "Projected Position on plane: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); return; } @@ -374,6 +419,21 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; bool momentumReversed = (fLastMom.Dot(fMom) < 0); + if (reachedMeasurementPoint) { + // We reached the measurement point, so we should figure out how far we are from the measurement point + // and update that remaining amount + LOG(info) << "------ Reached measurement point------"; + double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step + double approach = (fLastPos - point).R(); // Distance to the measurement point + LOG(info) << "Distance to measurement point: " << approach << " mm"; + LOG(info) << "Final step size: " << finalH << " mm"; + + finalH = approach * 1e-3; // Convert to meters for the RK4 step + fPos = fLastPos; + fMom = fLastMom; + RK4Step(finalH); // Propagate to the measurement point + } + // If we stopped, then we should figure out about where we stopped assuming linear de/dx over this last step if (particleStopped || momentumReversed) { LOG(info) << "------ Particle stopped ------"; @@ -400,22 +460,7 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); } - if (reachedMeasurementPoint) { - // We reached the measurement point, so we should figure out how far we are from the measurement point - // and update that remaining amount - LOG(info) << "------ Reached measurement point------"; - double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step - double approach = (fLastPos - point).R(); // Distance to the measurement point - LOG(info) << "Distance to measurement point: " << approach << " mm"; - LOG(info) << "Final step size: " << finalH << " mm"; - - finalH = approach * 1e-3; // Convert to meters for the RK4 step - fPos = fLastPos; - fMom = fLastMom; - RK4Step(finalH); // Propagate to the measurement point - } - - if (reachedMeasurementPoint || particleStopped) { + if (reachedMeasurementPoint || particleStopped || momentumReversed) { // Undo the last step since we were closer last time. double lastApproach = (fLastPos - point).R(); double approach = (fPos - point).R(); diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 01b916fa1..0f38e8f75 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -41,6 +41,7 @@ class AtPropagator { /// introduces at most 1mm of error. double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped (MeV) + double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) double fScalingFactor = 1.0; /// Scaling factor for energy loss XYZPoint fPos; // Current position of the particle in mm diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 2e831836c..10311e97a 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -213,6 +213,39 @@ TEST(AtPropagatorTest, PropagateToPlane_NoField) ASSERT_NEAR(finalMom.X(), p_fin, 0.1); } +TEST(AtPropagatorTest, PropagateToPlane_StoppingNoField) +{ + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable( + "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + AtPropagator propagator(charge, mass, std::move(elossModel)); + + double KE = 1; // Kinetic energy in MeV + double E = KE + mass_p; + double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c + XYZPoint startPos(0, 0, 0); // Start position in mm + XYZVector startMom(p, 0, 0); // Start momentum in MeV/c + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + XYZPoint planePoint(220, 0, 0); // Target point to propagate to 215 mm + XYZVector planeNormal(1, 0, 0); // Normal vector of the plane in x-direction + Plane3D plane(planeNormal, planePoint); // Create the plane + propagator.PropagateToPlane(plane); + + auto finalPos = propagator.GetPosition(); + auto finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), 220, 1); // Final position in x-direction should be close to 215 mm + ASSERT_NEAR(finalMom.X(), 0, 0.1); +} + TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) { double charge = charge_p; // Charge in Coulombs From fdee72ac5055506b3231548fe4b2ce025bbb34e6 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 01:31:02 +0200 Subject: [PATCH 16/71] Adaptive working --- AtTools/AtPropagator.cxx | 58 +++++++++++++++++++++++++++++++++------- 1 file changed, 49 insertions(+), 9 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 376588409..78e795a13 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -80,18 +80,53 @@ void AtPropagator::PropagateToPointAdaptive(const XYZPoint &point) bool reachedMeasurementPoint = ReachedPOCA(point); bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - if (reachedMeasurementPoint || particleStopped) { - // Last iteration we were still approaching the measurement point. Now we are further away - // then before. We have probably reached the measurement point if things are well behaved. - // I can think of cases where this will not be true. A better solution might be to run - // tracking the point of closest approach until the distance between the current state and - // the measurement point is larger than the distance between the last state and the measurement point. + bool momentumReversed = (fLastMom.Dot(fMom) < 0); + if (reachedMeasurementPoint) { + // We reached the measurement point, so we should figure out how far we are from the measurement point + // and update that remaining amount + LOG(info) << "------ Reached measurement point------"; + double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step + double approach = (fLastPos - point).R(); // Distance to the measurement point + LOG(info) << "Distance to measurement point: " << approach << " mm"; + LOG(info) << "Final step size: " << finalH << " mm"; + + finalH = approach * 1e-3; // Convert to meters for the RK4 step + fPos = fLastPos; + fMom = fLastMom; + RK4StepAdaptive(finalH); // Propagate to the measurement point + } + + // If we stopped, then we should figure out about where we stopped assuming linear de/dx over this last step + if (particleStopped || momentumReversed) { + LOG(info) << "------ Particle stopped ------"; + + double finalH = (fPos - fLastPos).R(); // Distance traveled in the last step + double E_loss = + Kinematics::KE(fLastMom, fMass) + Kinematics::KE(fMom, fMass); // Energy loss in MeV in the last step + double dedx = E_loss / finalH; // Stopping power in MeV/mm + + LOG(info) << "Particle stopped with final step size: " << finalH << " mm"; + LOG(info) << "Energy loss in last step: " << E_loss << " MeV"; + LOG(info) << "Stopping power (dE/dx): " << dedx << " MeV/mm"; + LOG(info) << "Energy before stopping: " << Kinematics::KE(fLastMom, fMass) << " MeV"; + finalH = Kinematics::KE(fLastMom, fMass) / dedx; // Distance to stop in mm + LOG(info) << "Estimated distance to stop: " << finalH << " mm"; + + fPos = fLastPos; + fMom = fLastMom; // Reset to last position and momentum + RK4Step(finalH); // Propagate to the point where we stopped + fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped + fLastMom = fMom; + + LOG(info) << "Final Position after stopping: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + } + + if (reachedMeasurementPoint || particleStopped || momentumReversed) { // Undo the last step since we were closer last time. double lastApproach = (fLastPos - point).R(); double approach = (fPos - point).R(); - fPos = fLastPos; - fMom = fLastMom; double KE_final = Kinematics::KE(fMom, fMass); auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV @@ -103,6 +138,11 @@ void AtPropagator::PropagateToPointAdaptive(const XYZPoint &point) LOG(info) << "Scaling factor: " << fScalingFactor; LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(info) << "Last Position: " << fLastPos.X() << ", " << fLastPos.Y() << ", " << fLastPos.Z(); + LOG(info) << "Last Momentum: " << fLastMom.X() << ", " << fLastMom.Y() << ", " << fLastMom.Z(); + + // fPos = fLastPos; + // fMom = fLastMom; return; } } // End of loop over RK4 integration @@ -115,7 +155,7 @@ bool AtPropagator::RK4StepAdaptive(double &h) double atol_pos = 1e-2; // Absolute tolerance for position (mm) double atol_mom = 1e-2; // Absolute tolerance for momentum (MeV/c) - double rtol = 1e-4; // Relative tolerance for both position and momentum + double rtol = 1e-6; // Relative tolerance for both position and momentum auto x0_mm = fPos; auto p0 = fMom; From 8b38a50a1a2fe96f7085f69b145002511bd04f79 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 01:34:52 +0200 Subject: [PATCH 17/71] Add test for small initial step size --- AtTools/AtPropagatorTest.cxx | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 10311e97a..813426116 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -268,7 +268,6 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) propagator.SetState(startPos, startMom); propagator.SetEField({0, 0, 0}); // No electric field propagator.SetBField({0, 0, 0}); // No magnetic field - propagator.SetDelta(1e-3); // Set relative error tolerance. Traveling 10 mm means at most 10 microns of error. propagator.SetH(1); // Set initial step size to 1 s ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); @@ -281,4 +280,19 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(finalPos.X(), 10, 10 * 1e-3); // Final position in x-direction should be close to 10 mm ASSERT_NEAR(finalMom.X(), p_fin, 0.1); + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + propagator.SetH(1e-6); // Set initial step size to 1e-6 m + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + propagator.PropagateToPointAdaptive(targetPoint); + + finalPos = propagator.GetPosition(); + finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), 10, 10 * 1e-3); // Final position in x-direction should be close to 10 mm + ASSERT_NEAR(finalMom.X(), p_fin, 0.1); } \ No newline at end of file From 6cda9957ad0b05bef18a2ca56f8a0b985737b688 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 11:21:42 +0200 Subject: [PATCH 18/71] Fix unit issue in stopping plane calculation. --- AtTools/AtPropagator.cxx | 28 ++++++++++++++------ AtTools/AtPropagator.h | 55 ++++++++++++++++++++++++++++++---------- 2 files changed, 61 insertions(+), 22 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 78e795a13..bf4b3b431 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -26,6 +26,9 @@ static constexpr double b[7] = {35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0 static constexpr double bs[7] = {5179.0 / 57600.0, 0.0, 7571.0 / 16695.0, 393.0 / 640.0, -92097.0 / 339200.0, 187.0 / 2100.0, 1.0 / 40.0}; +using ROOT::Math::Plane3D; +using ROOT::Math::XYZPoint; +using ROOT::Math::XYZVector; namespace AtTools { AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const @@ -299,18 +302,18 @@ void AtPropagator::RK4Step(double h) auto x_k1 = fMom.Unit(); // The derivative of the position is then just the unit vector of the momentum. auto p_k1 = dpds(fPos, fMom); // The derivative of the momentum is dpds. - auto x_2 = fPos + x_k1 * h / 2; // Position at the midpoint - auto p_2 = fMom + p_k1 * h / 2; // Momentum at the midpoint + auto x_2 = fPos + x_k1 * h / 2; // Position at the midpoint + auto p_2 = fMom + p_k1 * h / 2 / fReltoSImom; // Momentum at the midpoint auto x_k2 = p_2.Unit(); auto p_k2 = dpds(x_2, p_2); - auto x_3 = fPos + x_k2 * h / 2; // Position at the second midpoint - auto p_3 = fMom + p_k2 * h / 2; // Momentum at the second midpoint + auto x_3 = fPos + x_k2 * h / 2; // Position at the second midpoint + auto p_3 = fMom + p_k2 * h / 2 / fReltoSImom; // Momentum at the second midpoint auto x_k3 = p_3.Unit(); auto p_k3 = dpds(x_3, p_3); - auto x_4 = fPos + x_k3 * h; // Position at the end of the step - auto p_4 = fMom + p_k3 * h; // Momentum at the end of the step + auto x_4 = fPos + x_k3 * h; // Position at the end of the step + auto p_4 = fMom + p_k3 * h / fReltoSImom; // Momentum at the end of the step auto x_k4 = p_4.Unit(); auto p_k4 = dpds(x_4, p_4); @@ -393,9 +396,10 @@ void AtPropagator::PropagateToPlane(const Plane3D &plane) LOG(info) << "Last KE: " << KE_last << " MeV"; LOG(info) << "Energy to loose to stop: " << deltaE << " MeV"; - double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm - RK4Step(h_Stop); + LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; + + RK4Step(h_Stop * 1e-3); LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; @@ -571,4 +575,12 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) fScalingFactor = 1; // Reset scaling factor after convergence } + +AtStepper::StepResult +AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom, DerivFunc derivFunc) const +{ + // Take h to be the step size in m. + + return {}; +} } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 0f38e8f75..1b1a29247 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -188,20 +188,47 @@ class AtPropagator { bool IntersectedPlane(const Plane3D &plane); }; -static constexpr double a21 = 1.0 / 5.0; -static constexpr double a31 = 3.0 / 40.0, a32 = 9.0 / 40.0; -static constexpr double a41 = 44.0 / 45.0, a42 = -56.0 / 15.0, a43 = 32.0 / 9.0; -static constexpr double a51 = 19372.0 / 6561.0, a52 = -25360.0 / 2187.0, a53 = 64448.0 / 6561.0, a54 = -212.0 / 729.0; -static constexpr double a61 = 9017.0 / 3168.0, a62 = -355.0 / 33.0, a63 = 46732.0 / 5247.0, a64 = 49.0 / 176.0, - a65 = -5103.0 / 18656.0; -static constexpr double a71 = 35.0 / 384.0, a72 = 0.0, a73 = 500.0 / 1113.0, a74 = 125.0 / 192.0, - a75 = -2187.0 / 6784.0, a76 = 11.0 / 84.0; -// b (5th-order) -static constexpr double b1 = 35.0 / 384.0, b3 = 500.0 / 1113.0, b4 = 125.0 / 192.0, b5 = -2187.0 / 6784.0, - b6 = 11.0 / 84.0; -// b* (4th-order, “star”) -static constexpr double bs1 = 5179.0 / 57600.0, bs3 = 7571.0 / 16695.0, bs4 = 393.0 / 640.0, bs5 = -92097.0 / 339200.0, - bs6 = 187.0 / 2100.0, bs7 = 1.0 / 40.0; +class AtStepper { +public: + struct StepResult { + ROOT::Math::XYZPoint pos; // Position of the particle in mm + ROOT::Math::XYZVector mom; // Momentum of the particle in MeV/c + ROOT::Math::XYZPoint lastPos; // Last position of the particle in mm + ROOT::Math::XYZVector lastMom; // Last momentum of the particle in MeV/c + double h; // Step size for the step in m + bool success; // Whether the step was successful + }; + /** + * @brief Function type defining the derivative of the position and momentum w.r.t. distance. + * + * This function takes the current position and momentum and returns the derivate of the position and momentum. + * + * @param pos Current position of the particle in mm. + * @param mom Current momentum of the particle in MeV/c. + * @return A pair containing the derivatives of the position and momentum in SI units (m and kg m/s). + * The first element is the derivative of the position, and the second element is the derivative + * of the momentum. + */ + using DerivFunc = std::function( + const ROOT::Math::XYZPoint &, const ROOT::Math::XYZVector &)>; + + virtual StepResult + Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom, DerivFunc derivFunc) const = 0; + +protected: + static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) +}; + +class AtRK4Stepper : public AtStepper { +public: + StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom, + DerivFunc derivFunc) const override; +}; +class AtRK4AdaptiveStepper : public AtStepper { +public: + StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom, + DerivFunc derivFunc) const override; +}; } // namespace AtTools #endif // #ifndef ATPROPAGATOR_H From 511e04fb6070249223a86bae1176162f63f93d2d Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 11:50:47 +0200 Subject: [PATCH 19/71] PropagateToPoint uses AtStepper interface --- AtTools/AtPropagator.cxx | 86 ++++++++++++++++++++++++++++++++----- AtTools/AtPropagator.h | 92 +++++++++++++++++++++++++--------------- 2 files changed, 132 insertions(+), 46 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index bf4b3b431..b4c04a483 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -448,13 +448,19 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) LOG(info) << "Propagating to point: " << point; auto KE_initial = Kinematics::KE(fMom, fMass); + AtRK4Stepper stepper; + stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; + while (true) { - fLastPos = fPos; - fLastMom = fMom; LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - RK4Step(fH); + auto result = stepper.Step(fH, fPos, fMom); + if (!result.success) { + LOG(error) << "Integration step failed, aborting propagation."; + return; // Abort propagation if step failed + } + CopyFromState(result); // Copy the new state from the stepper if (!ReachedPOCA(point)) continue; @@ -463,7 +469,7 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; bool momentumReversed = (fLastMom.Dot(fMom) < 0); - if (reachedMeasurementPoint) { + if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { // We reached the measurement point, so we should figure out how far we are from the measurement point // and update that remaining amount LOG(info) << "------ Reached measurement point------"; @@ -473,15 +479,42 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) LOG(info) << "Final step size: " << finalH << " mm"; finalH = approach * 1e-3; // Convert to meters for the RK4 step - fPos = fLastPos; - fMom = fLastMom; - RK4Step(finalH); // Propagate to the measurement point + result = stepper.Step(finalH, fLastPos, fLastMom); + if (!result.success) { + LOG(error) << "Failed to propagate to measurement point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fH; // Save original step size + CopyFromState(result); // Update position and momentum to the new state + fH = origH; // Restore original step size } // If we stopped, then we should figure out about where we stopped assuming linear de/dx over this last step if (particleStopped || momentumReversed) { - LOG(info) << "------ Particle stopped ------"; + LOG(info) << "------ Particle stopped before measurement point/surface------"; + + // Calculate how far to travel before stopping + double KE_last = Kinematics::KE(fLastMom, fMass); + double deltaE = KE_last - fStopTol; + deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss + double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm + + LOG(info) << "KE at last point: " << KE_last << " MeV"; + LOG(info) << "Energy to lose to stop: " << deltaE << " MeV"; + LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; + LOG(info) << "dE/dx at last point: " << fELossModel->GetdEdx(KE_last) << " MeV/mm"; + result = stepper.Step(h_Stop * 1e-3, fLastPos, fLastMom); + if (!result.success) { + LOG(error) << "Failed to propagate to stopping point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fH; // Save original step size + CopyFromState(result); // Update position and momentum to the new state + fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped + fLastMom = fMom; // Update last momentum to zero + fH = origH; // Restore original step size + /* double finalH = (fPos - fLastPos).R(); // Distance traveled in the last step double E_loss = Kinematics::KE(fLastMom, fMass) + Kinematics::KE(fMom, fMass); // Energy loss in MeV in the last step @@ -499,6 +532,7 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) RK4Step(finalH); // Propagate to the point where we stopped fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped fLastMom = fMom; + */ LOG(info) << "Final Position after stopping: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); @@ -576,11 +610,41 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) fScalingFactor = 1; // Reset scaling factor after convergence } -AtStepper::StepResult -AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom, DerivFunc derivFunc) const +AtStepper::StepResult AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const { // Take h to be the step size in m. + StepResult result; + result.lastPos = fPos; + result.lastMom = fMom; + result.h = h; + result.success = true; + + auto [x_k1, p_k1] = + fDeriv(fPos, fMom); // The derivative of the position is then just the unit vector of the momentum. + + auto x_2 = fPos + x_k1 * h / 2; // Position at the midpoint + auto p_2 = fMom + p_k1 * h / 2 / fReltoSImom; // Momentum at the midpoint + + auto [x_k2, p_k2] = fDeriv(x_2, p_2); // Derivative at the midpoint + + auto x_3 = fPos + x_k2 * h / 2; // Position at the second midpoint + auto p_3 = fMom + p_k2 * h / 2 / fReltoSImom; // Momentum at the second midpoint + auto [x_k3, p_k3] = fDeriv(x_3, p_3); + + auto x_4 = fPos + x_k3 * h; // Position at the end of the step + auto p_4 = fMom + p_k3 * h / fReltoSImom; // Momentum at the end of the step + auto [x_k4, p_k4] = fDeriv(x_4, p_4); + + auto dpds_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // "Force" in SI units (N) + + auto mom_SI = fReltoSImom * fMom; + mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) + result.mom = mom_SI / fReltoSImom; // Convert back to + + auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) + pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; // Update position in SI units (m + result.pos = pos_SI * 1e3; // Convert back to mm - return {}; + return result; } } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 1b1a29247..0cdbea65e 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -11,6 +11,38 @@ namespace AtTools { +class AtStepper { +public: + struct StepResult { + ROOT::Math::XYZPoint pos; // Position of the particle in mm + ROOT::Math::XYZVector mom; // Momentum of the particle in MeV/c + ROOT::Math::XYZPoint lastPos; // Last position of the particle in mm + ROOT::Math::XYZVector lastMom; // Last momentum of the particle in MeV/c + double h; // Step size for the step in m + bool success; // Whether the step was successful + }; + /** + * @brief Function type defining the derivative of the position and momentum w.r.t. distance. + * + * This function takes the current position and momentum and returns the derivate of the position and momentum. + * + * @param pos Current position of the particle in mm. + * @param mom Current momentum of the particle in MeV/c. + * @return A pair containing the derivatives of the position and momentum in SI units (m and kg m/s). + * The first element is the derivative of the position, and the second element is the derivative + * of the momentum. + */ + using DerivFunc = std::function( + const ROOT::Math::XYZPoint &, const ROOT::Math::XYZVector &)>; + + DerivFunc fDeriv; + + virtual StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const = 0; + +protected: + static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) +}; + /** * @brief Class for propagating particles through a medium. * @@ -148,6 +180,12 @@ class AtPropagator { return mom.Unit(); // The derivative of the position is just the unit vector of the momentum. } + std::pair + Derivatives(const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const + { + return {dxds(pos, mom), dpds(pos, mom)}; + } + protected: /** * @brief Perform a single RK4 step for propagation. @@ -186,48 +224,32 @@ class AtPropagator { bool ReachedPOCA(const XYZPoint &point); bool IntersectedPlane(const Plane3D &plane); -}; - -class AtStepper { -public: - struct StepResult { - ROOT::Math::XYZPoint pos; // Position of the particle in mm - ROOT::Math::XYZVector mom; // Momentum of the particle in MeV/c - ROOT::Math::XYZPoint lastPos; // Last position of the particle in mm - ROOT::Math::XYZVector lastMom; // Last momentum of the particle in MeV/c - double h; // Step size for the step in m - bool success; // Whether the step was successful - }; - /** - * @brief Function type defining the derivative of the position and momentum w.r.t. distance. - * - * This function takes the current position and momentum and returns the derivate of the position and momentum. - * - * @param pos Current position of the particle in mm. - * @param mom Current momentum of the particle in MeV/c. - * @return A pair containing the derivatives of the position and momentum in SI units (m and kg m/s). - * The first element is the derivative of the position, and the second element is the derivative - * of the momentum. - */ - using DerivFunc = std::function( - const ROOT::Math::XYZPoint &, const ROOT::Math::XYZVector &)>; - - virtual StepResult - Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom, DerivFunc derivFunc) const = 0; - -protected: - static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) + void CopyFromState(const AtStepper::StepResult &result) + { + fPos = result.pos; + fMom = result.mom; + fLastPos = result.lastPos; + fLastMom = result.lastMom; + fH = result.h; + } + AtStepper::StepResult CopyToState() const + { + AtStepper::StepResult result; + result.pos = fPos; + result.mom = fMom; + result.lastPos = fLastPos; + result.lastMom = fLastMom; + return result; + } }; class AtRK4Stepper : public AtStepper { public: - StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom, - DerivFunc derivFunc) const override; + StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; }; class AtRK4AdaptiveStepper : public AtStepper { public: - StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom, - DerivFunc derivFunc) const override; + StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; }; } // namespace AtTools From a848c2b6c3b24468622a6cf235093e3c8f4197cd Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 12:07:41 +0200 Subject: [PATCH 20/71] All RK4 to point tests use stepper interface --- AtTools/AtPropagator.cxx | 182 ++++++++++++++++++++++++++++++----- AtTools/AtPropagator.h | 5 +- AtTools/AtPropagatorTest.cxx | 8 +- 3 files changed, 166 insertions(+), 29 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index b4c04a483..4a21f0919 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -443,12 +443,11 @@ void AtPropagator::PropagateToPlane(const Plane3D &plane) } // End of loop over RK4 integration } -void AtPropagator::PropagateToPoint(const XYZPoint &point) +void AtPropagator::PropagateToPoint(const XYZPoint &point, AtStepper &stepper) { LOG(info) << "Propagating to point: " << point; auto KE_initial = Kinematics::KE(fMom, fMass); - AtRK4Stepper stepper; stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; while (true) { @@ -514,25 +513,6 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped fLastMom = fMom; // Update last momentum to zero fH = origH; // Restore original step size - /* - double finalH = (fPos - fLastPos).R(); // Distance traveled in the last step - double E_loss = - Kinematics::KE(fLastMom, fMass) + Kinematics::KE(fMom, fMass); // Energy loss in MeV in the last step - double dedx = E_loss / finalH; // Stopping power in MeV/mm - - LOG(info) << "Particle stopped with final step size: " << finalH << " mm"; - LOG(info) << "Energy loss in last step: " << E_loss << " MeV"; - LOG(info) << "Stopping power (dE/dx): " << dedx << " MeV/mm"; - LOG(info) << "Energy before stopping: " << Kinematics::KE(fLastMom, fMass) << " MeV"; - finalH = Kinematics::KE(fLastMom, fMass) / dedx; // Distance to stop in mm - LOG(info) << "Estimated distance to stop: " << finalH << " mm"; - - fPos = fLastPos; - fMom = fLastMom; // Reset to last position and momentum - RK4Step(finalH); // Propagate to the point where we stopped - fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped - fLastMom = fMom; - */ LOG(info) << "Final Position after stopping: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); @@ -563,13 +543,13 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point) } // End of loop over RK4 integration } -void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) +void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss, AtStepper &stepper) { LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; if (eLoss == 0) { LOG(warn) << "No energy loss specified, propagating without energy loss adjustment."; - PropagateToPoint(point); + PropagateToPoint(point, stepper); return; } @@ -592,7 +572,7 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss) } iterations++; - PropagateToPoint(point); // Propagate without energy loss adjustment + PropagateToPoint(point, stepper); // Propagate without energy loss adjustment double KE_final = Kinematics::KE(fMom, fMass); calc_eLoss = KE_initial - KE_final; // Energy loss in MeV @@ -647,4 +627,158 @@ AtStepper::StepResult AtRK4Stepper::Step(double h, const XYZPoint &fPos, const X return result; } + +AtStepper::StepResult AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const +{ + // Take h to be the step size in m. + StepResult result; + result.lastPos = fPos; + result.lastMom = fMom; + result.h = h; + result.success = true; + + // Take h to be the step size in m. + // Use DP5(4) method for adaptive step size control. + + double atol_pos = 1e-2; // Absolute tolerance for position (mm) + double atol_mom = 1e-2; // Absolute tolerance for momentum (MeV/c) + double rtol = 1e-6; // Relative tolerance for both position and momentum + + auto x0_mm = fPos; + auto p0 = fMom; + LOG(info) << "Starting RK4 step with initial position: " << x0_mm.X() << ", " << x0_mm.Y() << ", " << x0_mm.Z(); + LOG(info) << "Initial momentum: " << p0.X() << ", " << p0.Y() << ", " << p0.Z(); + + while (true) { + auto x_SI = fPos * 1e-3; // Convert position to SI units (m) + auto p_SI = fReltoSImom * fMom; // Convert momentum to SI units (kg m/s) + XYZVector kx[7]; // kx[i] will hold the position derivatives (unitless) + XYZVector kp[7]; // kp[i] will hold the momentum derivatives (SI units) + + // anonymous lambda to calculate and store the kx and kp values. Input is SI units. + auto calc_k = [&](const XYZPoint &x, const XYZVector &p, int i) { + auto [k_x, k_p] = fDeriv(x * 1e-3, p / fReltoSImom); + kx[i] = k_x; // Store the position derivative (unitless) + kp[i] = k_p; // Store the momentum derivative (SI units) + }; + + // anonymous lambda to calculate the position and momentum at the i-th stage + auto calc_xp = [&](int i) { + XYZVector dx(0, 0, 0); + XYZVector dp(0, 0, 0); + for (int j = 0; j < i; ++j) { + dx = dx + kx[j] * a[i][j]; + dp = dp + kp[j] * a[i][j]; + } + XYZPoint x = x_SI + dx * h; + XYZVector p = p_SI + dp * h; + return std::make_pair(x, p); + }; + + // Calculate kx and kp for each stage + // build stage 0 + calc_k(x_SI, p0, 0); + + // build stage 1 + auto [x1, p1] = calc_xp(1); + calc_k(x1, p1, 1); // k1 + + // build stage 2 + auto [x2, p2] = calc_xp(2); + calc_k(x2, p2, 2); // k2 + + // build stage 3 + auto [x3, p3] = calc_xp(3); + calc_k(x3, p3, 3); // k3 + + // build stage 4 + auto [x4, p4] = calc_xp(4); + calc_k(x4, p4, 4); // k4 + + // build stage 5 + auto [x5, p5] = calc_xp(5); + calc_k(x5, p5, 5); // k5 + + // build stage 6 + auto [x6, p6] = calc_xp(6); + calc_k(x6, p6, 6); // k6 + + // Calculate the new position and momentum using the 5th-order method + XYZVector dx(0, 0, 0); + XYZVector dp(0, 0, 0); + for (int i = 0; i < 7; ++i) { + dx = dx + kx[i] * b[i]; + dp = dp + kp[i] * b[i]; + } + XYZPoint x_new_5 = x_SI + dx * h; // New position in SI units (m) + XYZVector p_new_5 = p_SI + dp * h; // New momentum in SI units (kg m/s) + + // Calculate the new position and momentum using the 4th-order method + dx = XYZVector(0, 0, 0); + dp = XYZVector(0, 0, 0); + for (int i = 0; i < 7; ++i) { + dx = dx + kx[i] * bs[i]; + dp = dp + kp[i] * bs[i]; + } + XYZPoint x_new_4 = x_SI + dx * h; // New position in SI units (m) + XYZVector p_new_4 = p_SI + dp * h; // New momentum in SI units (kg m/s) + + auto x_4_mm = x_new_4 * 1e3; // Convert back to mm + auto p_4_MeV = p_new_4 / fReltoSImom; // Convert back to MeV/c + auto x_5_mm = x_new_5 * 1e3; // Convert back to mm + auto p_5_MeV = p_new_5 / fReltoSImom; // Convert back to MeV/c + LOG(info) << "New position (5th order): " << x_5_mm.X() << ", " << x_5_mm.Y() << ", " << x_5_mm.Z(); + LOG(info) << "New momentum (5th order): " << p_5_MeV.X() << ", " << p_5_MeV.Y() << ", " << p_5_MeV.Z(); + LOG(info) << "New position (4th order): " << x_4_mm.X() << ", " << x_4_mm.Y() << ", " << x_4_mm.Z(); + LOG(info) << "New momentum (4th order): " << p_4_MeV.X() << ", " << p_4_MeV.Y() << ", " << p_4_MeV.Z(); + + // Convert back to mm and MeV/c + XYZVector x_err = (x_5_mm - x_4_mm); // Error in position (mm) + XYZVector p_err = (p_5_MeV - p_4_MeV); // Error in momentum (MeV/c) + + // Calculate the overall error + double ex = x_err.X() / (atol_pos + rtol * std::abs(x_5_mm.X())); + double ey = x_err.Y() / (atol_pos + rtol * std::abs(x_5_mm.Y())); + double ez = x_err.Z() / (atol_pos + rtol * std::abs(x_5_mm.Z())); + + double ep_x = p_err.X() / (atol_mom + rtol * std::abs(p_5_MeV.X())); + double ep_y = p_err.Y() / (atol_mom + rtol * std::abs(p_5_MeV.Y())); + double ep_z = p_err.Z() / (atol_mom + rtol * std::abs(p_5_MeV.Z())); + + // Combine errors (norm) + double err = std::sqrt(ex * ex + ey * ey + ez * ez + ep_x * ep_x + ep_y * ep_y + ep_z * ep_z); + + double factor = std::pow(err, -1.0 / 5.0); // Adjust step size based on error + factor = std::clamp(factor, 0.25, 4.0); // Clamp factor to reasonable limits + double hNew = h * factor; + // We now know the local error at this point. Now we need to decide to accept the point or not. + if (err <= 1.0) { + // Accept the step + result.pos = x_5_mm; // Update position in mm + result.mom = p_5_MeV; // Update momentum in MeV/c + LOG(info) << "Accepted step with error: " << err; + LOG(info) << "Step size: " << h << " m"; + LOG(info) << "New step size: " << hNew << " m"; + LOG(info) << "New Position: " << result.pos.X() << ", " << result.pos.Y() << ", " << result.pos.Z(); + LOG(info) << "New Momentum: " << result.mom.X() << ", " << result.mom.Y() << ", " << result.mom.Z(); + + // Adjust the step size for the next iteration + result.h = hNew; + result.success = true; // Step accepted + return result; + } else { + // Reject the step and reduce the step size + LOG(info) << "Rejected step with error: " << err; + LOG(info) << "Step size: " << h << " m"; + LOG(info) << "Reducing step size to: " << hNew << " m"; + + result.h = hNew; // Reduce step size + if (result.h < 1e-6) { + LOG(error) << "Step size too small, aborting propagation."; + result.success = false; + return result; // Abort propagation if step size is too small + } + } + } +} } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 0cdbea65e..78e1f2694 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -128,14 +128,15 @@ class AtPropagator { * @param point The point to approach. * @param eLoss If not 0, constrain the energy loss to this value by adjusting fScalingFactor. */ - void PropagateToPoint(const XYZPoint &point, double eLoss); + void PropagateToPoint(const XYZPoint &point, double eLoss, AtStepper &stepper); /** * @brief Propagate the particle to the point of closest approach to the given point. * * @param point The point to approach. + * @param stepper The stepper to use for propagation. */ - void PropagateToPoint(const XYZPoint &point); + void PropagateToPoint(const XYZPoint &point, AtStepper &stepper); void PropagateToPointAdaptive(const XYZPoint &point); /** diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 813426116..5cf21a84a 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -106,6 +106,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) elossModel->LoadSrimTable( "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4Stepper stepper; double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -120,7 +121,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); XYZPoint targetPoint(1e3, 0, 0); // Target point to propagate to (1 m) - propagator.PropagateToPoint(targetPoint); + propagator.PropagateToPoint(targetPoint, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -134,7 +135,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) startMom.SetXYZ(p, 0, 0); // Reset momentum propagator.SetState(startPos, startMom); - propagator.PropagateToPoint(targetPoint); // Propagate to range + propagator.PropagateToPoint(targetPoint, stepper); // Propagate to range finalPos = propagator.GetPosition(); finalMom = propagator.GetMomentum(); ASSERT_NEAR(finalPos.X(), 130, 10); // Final position in x-direction should be close to 130 mm @@ -149,6 +150,7 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) elossModel->LoadSrimTable( "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4Stepper stepper; double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -167,7 +169,7 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); XYZPoint targetPoint(10, 0, 0); // Target point to propagate to 10 mm - propagator.PropagateToPoint(targetPoint); + propagator.PropagateToPoint(targetPoint, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); From 08c7ccba2b4f86515ab3211c018cb28b0938dd75 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 12:12:16 +0200 Subject: [PATCH 21/71] Adaptive test working with stepper --- AtTools/AtPropagator.cxx | 1 + AtTools/AtPropagatorTest.cxx | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 4a21f0919..7a070a64a 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -773,6 +773,7 @@ AtStepper::StepResult AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, LOG(info) << "Reducing step size to: " << hNew << " m"; result.h = hNew; // Reduce step size + h = hNew; // Update the step size for the next iteration if (result.h < 1e-6) { LOG(error) << "Step size too small, aborting propagation."; result.success = false; diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 5cf21a84a..55b1321a7 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -256,6 +256,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) elossModel->LoadSrimTable( "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4AdaptiveStepper stepper; double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -275,7 +276,8 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); XYZPoint targetPoint(10, 0, 0); // Target point to propagate to (10 mm) - propagator.PropagateToPointAdaptive(targetPoint); + // propagator.PropagateToPointAdaptive(targetPoint); + propagator.PropagateToPoint(targetPoint, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -290,7 +292,8 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - propagator.PropagateToPointAdaptive(targetPoint); + // propagator.PropagateToPointAdaptive(targetPoint); + propagator.PropagateToPoint(targetPoint, stepper); finalPos = propagator.GetPosition(); finalMom = propagator.GetMomentum(); From 365e2114c5b066ac17f95db474697748f534f017 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 12:29:21 +0200 Subject: [PATCH 22/71] Plane propagation uses AtStepper interface --- AtTools/AtPropagator.cxx | 327 +++++------------------------------ AtTools/AtPropagator.h | 39 +---- AtTools/AtPropagatorTest.cxx | 8 +- 3 files changed, 50 insertions(+), 324 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 7a070a64a..53a7c792b 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -61,273 +61,6 @@ AtPropagator::XYZVector AtPropagator::d2xds2(const XYZPoint &pos, const XYZVecto return 1 / p * (dpds_vec - phat * (phat.Dot(dpds_vec))); // Second derivative of position w.r.t. arc length } -void AtPropagator::PropagateToPointAdaptive(const XYZPoint &point) -{ - LOG(info) << "Propagating to point: " << point; - - auto KE_initial = Kinematics::KE(fMom, fMass); - while (true) { - fLastPos = fPos; - fLastMom = fMom; - LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - - auto acceptedStep = RK4StepAdaptive(fH); - if (!acceptedStep) { - LOG(error) << "RK4 step failed, aborting propagation."; - return; // Abort propagation if step failed - } - - if (!acceptedStep || !ReachedPOCA(point)) - continue; - - bool reachedMeasurementPoint = ReachedPOCA(point); - bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - bool momentumReversed = (fLastMom.Dot(fMom) < 0); - - if (reachedMeasurementPoint) { - // We reached the measurement point, so we should figure out how far we are from the measurement point - // and update that remaining amount - LOG(info) << "------ Reached measurement point------"; - double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step - double approach = (fLastPos - point).R(); // Distance to the measurement point - LOG(info) << "Distance to measurement point: " << approach << " mm"; - LOG(info) << "Final step size: " << finalH << " mm"; - - finalH = approach * 1e-3; // Convert to meters for the RK4 step - fPos = fLastPos; - fMom = fLastMom; - RK4StepAdaptive(finalH); // Propagate to the measurement point - } - - // If we stopped, then we should figure out about where we stopped assuming linear de/dx over this last step - if (particleStopped || momentumReversed) { - LOG(info) << "------ Particle stopped ------"; - - double finalH = (fPos - fLastPos).R(); // Distance traveled in the last step - double E_loss = - Kinematics::KE(fLastMom, fMass) + Kinematics::KE(fMom, fMass); // Energy loss in MeV in the last step - double dedx = E_loss / finalH; // Stopping power in MeV/mm - - LOG(info) << "Particle stopped with final step size: " << finalH << " mm"; - LOG(info) << "Energy loss in last step: " << E_loss << " MeV"; - LOG(info) << "Stopping power (dE/dx): " << dedx << " MeV/mm"; - LOG(info) << "Energy before stopping: " << Kinematics::KE(fLastMom, fMass) << " MeV"; - finalH = Kinematics::KE(fLastMom, fMass) / dedx; // Distance to stop in mm - LOG(info) << "Estimated distance to stop: " << finalH << " mm"; - - fPos = fLastPos; - fMom = fLastMom; // Reset to last position and momentum - RK4Step(finalH); // Propagate to the point where we stopped - fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped - fLastMom = fMom; - - LOG(info) << "Final Position after stopping: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - } - - if (reachedMeasurementPoint || particleStopped || momentumReversed) { - // Undo the last step since we were closer last time. - double lastApproach = (fLastPos - point).R(); - double approach = (fPos - point).R(); - - double KE_final = Kinematics::KE(fMom, fMass); - auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - LOG(info) << "------- End of RK4 interation ---------"; - LOG(info) << "Particle stopped: " << particleStopped; - LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; - LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; - LOG(info) << "Scaling factor: " << fScalingFactor; - LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - LOG(info) << "Last Position: " << fLastPos.X() << ", " << fLastPos.Y() << ", " << fLastPos.Z(); - LOG(info) << "Last Momentum: " << fLastMom.X() << ", " << fLastMom.Y() << ", " << fLastMom.Z(); - - // fPos = fLastPos; - // fMom = fLastMom; - return; - } - } // End of loop over RK4 integration -} - -bool AtPropagator::RK4StepAdaptive(double &h) -{ - // Take h to be the step size in m. - // Use DP5(4) method for adaptive step size control. - - double atol_pos = 1e-2; // Absolute tolerance for position (mm) - double atol_mom = 1e-2; // Absolute tolerance for momentum (MeV/c) - double rtol = 1e-6; // Relative tolerance for both position and momentum - - auto x0_mm = fPos; - auto p0 = fMom; - LOG(info) << "Starting RK4 step with initial position: " << x0_mm.X() << ", " << x0_mm.Y() << ", " << x0_mm.Z(); - LOG(info) << "Initial momentum: " << p0.X() << ", " << p0.Y() << ", " << p0.Z(); - - while (true) { - auto x_SI = fPos * 1e-3; // Convert position to SI units (m) - auto p_SI = fReltoSImom * fMom; // Convert momentum to SI units (kg m/s) - XYZVector kx[7]; // kx[i] will hold the position derivatives (unitless) - XYZVector kp[7]; // kp[i] will hold the momentum derivatives (SI units) - - // anonymous lambda to calculate and store the kx and kp values. Input is SI units. - auto calc_k = [&](const XYZPoint &x, const XYZVector &p, int i) { - kx[i] = p.Unit(); // The derivative of the position is then just the unit vector of the momentum. - kp[i] = dpds(x * 1e-3, p / fReltoSImom); // The derivative of the momentum is dpds. - }; - - // anonymous lambda to calculate the position and momentum at the i-th stage - auto calc_xp = [&](int i) { - XYZVector dx(0, 0, 0); - XYZVector dp(0, 0, 0); - for (int j = 0; j < i; ++j) { - dx = dx + kx[j] * a[i][j]; - dp = dp + kp[j] * a[i][j]; - } - XYZPoint x = x_SI + dx * h; - XYZVector p = p_SI + dp * h; - return std::make_pair(x, p); - }; - - // Calculate kx and kp for each stage - // build stage 0 - calc_k(x_SI, p0, 0); - - // build stage 1 - auto [x1, p1] = calc_xp(1); - calc_k(x1, p1, 1); // k1 - - // build stage 2 - auto [x2, p2] = calc_xp(2); - calc_k(x2, p2, 2); // k2 - - // build stage 3 - auto [x3, p3] = calc_xp(3); - calc_k(x3, p3, 3); // k3 - - // build stage 4 - auto [x4, p4] = calc_xp(4); - calc_k(x4, p4, 4); // k4 - - // build stage 5 - auto [x5, p5] = calc_xp(5); - calc_k(x5, p5, 5); // k5 - - // build stage 6 - auto [x6, p6] = calc_xp(6); - calc_k(x6, p6, 6); // k6 - - // Calculate the new position and momentum using the 5th-order method - XYZVector dx(0, 0, 0); - XYZVector dp(0, 0, 0); - for (int i = 0; i < 7; ++i) { - dx = dx + kx[i] * b[i]; - dp = dp + kp[i] * b[i]; - } - XYZPoint x_new_5 = x_SI + dx * h; // New position in SI units (m) - XYZVector p_new_5 = p_SI + dp * h; // New momentum in SI units (kg m/s) - - // Calculate the new position and momentum using the 4th-order method - dx = XYZVector(0, 0, 0); - dp = XYZVector(0, 0, 0); - for (int i = 0; i < 7; ++i) { - dx = dx + kx[i] * bs[i]; - dp = dp + kp[i] * bs[i]; - } - XYZPoint x_new_4 = x_SI + dx * h; // New position in SI units (m) - XYZVector p_new_4 = p_SI + dp * h; // New momentum in SI units (kg m/s) - - auto x_4_mm = x_new_4 * 1e3; // Convert back to mm - auto p_4_MeV = p_new_4 / fReltoSImom; // Convert back to MeV/c - auto x_5_mm = x_new_5 * 1e3; // Convert back to mm - auto p_5_MeV = p_new_5 / fReltoSImom; // Convert back to MeV/c - LOG(info) << "New position (5th order): " << x_5_mm.X() << ", " << x_5_mm.Y() << ", " << x_5_mm.Z(); - LOG(info) << "New momentum (5th order): " << p_5_MeV.X() << ", " << p_5_MeV.Y() << ", " << p_5_MeV.Z(); - LOG(info) << "New position (4th order): " << x_4_mm.X() << ", " << x_4_mm.Y() << ", " << x_4_mm.Z(); - LOG(info) << "New momentum (4th order): " << p_4_MeV.X() << ", " << p_4_MeV.Y() << ", " << p_4_MeV.Z(); - - // Convert back to mm and MeV/c - XYZVector x_err = (x_5_mm - x_4_mm); // Error in position (mm) - XYZVector p_err = (p_5_MeV - p_4_MeV); // Error in momentum (MeV/c) - - // Calculate the overall error - double ex = x_err.X() / (atol_pos + rtol * std::abs(x_5_mm.X())); - double ey = x_err.Y() / (atol_pos + rtol * std::abs(x_5_mm.Y())); - double ez = x_err.Z() / (atol_pos + rtol * std::abs(x_5_mm.Z())); - - double ep_x = p_err.X() / (atol_mom + rtol * std::abs(p_5_MeV.X())); - double ep_y = p_err.Y() / (atol_mom + rtol * std::abs(p_5_MeV.Y())); - double ep_z = p_err.Z() / (atol_mom + rtol * std::abs(p_5_MeV.Z())); - - // Combine errors (norm) - double err = std::sqrt(ex * ex + ey * ey + ez * ez + ep_x * ep_x + ep_y * ep_y + ep_z * ep_z); - - double factor = std::pow(err, -1.0 / 5.0); // Adjust step size based on error - factor = std::clamp(factor, 0.25, 4.0); // Clamp factor to reasonable limits - double hNew = h * factor; - // We now know the local error at this point. Now we need to decide to accept the point or not. - if (err <= 1.0) { - // Accept the step - fPos = x_5_mm; // Update position in mm - fMom = p_5_MeV; // Update momentum in MeV/c - LOG(info) << "Accepted step with error: " << err; - LOG(info) << "Step size: " << h << " m"; - LOG(info) << "New step size: " << hNew << " m"; - LOG(info) << "New Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "New Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - - // Adjust the step size for the next iteration - h = hNew; - return true; // Step accepted - } else { - // Reject the step and reduce the step size - LOG(info) << "Rejected step with error: " << err; - LOG(info) << "Step size: " << h << " m"; - LOG(info) << "Reducing step size to: " << hNew << " m"; - - h = hNew; // Reduce step size - if (h < 1e-6) { - LOG(error) << "Step size too small, aborting propagation."; - return false; // Abort propagation if step size is too small - } - } - } -} - -void AtPropagator::RK4Step(double h) -{ - // Take h to be the step size in m. - - auto x_k1 = fMom.Unit(); // The derivative of the position is then just the unit vector of the momentum. - auto p_k1 = dpds(fPos, fMom); // The derivative of the momentum is dpds. - - auto x_2 = fPos + x_k1 * h / 2; // Position at the midpoint - auto p_2 = fMom + p_k1 * h / 2 / fReltoSImom; // Momentum at the midpoint - auto x_k2 = p_2.Unit(); - auto p_k2 = dpds(x_2, p_2); - - auto x_3 = fPos + x_k2 * h / 2; // Position at the second midpoint - auto p_3 = fMom + p_k2 * h / 2 / fReltoSImom; // Momentum at the second midpoint - auto x_k3 = p_3.Unit(); - auto p_k3 = dpds(x_3, p_3); - - auto x_4 = fPos + x_k3 * h; // Position at the end of the step - auto p_4 = fMom + p_k3 * h / fReltoSImom; // Momentum at the end of the step - auto x_k4 = p_4.Unit(); - auto p_k4 = dpds(x_4, p_4); - - auto dpds_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // "Force" in SI units (N) - - auto mom_SI = fReltoSImom * fMom; - mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) - fMom = mom_SI / fReltoSImom; // Convert back to - - auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) - pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; // Update position in SI units (m - fPos = pos_SI * 1e3; // Convert back to mm -} - bool AtPropagator::ReachedPOCA(const XYZPoint &point) { // Here we need to check if we are getting closer or further away from the POCA. @@ -351,18 +84,23 @@ bool AtPropagator::IntersectedPlane(const Plane3D &plane) return (prevSign != currSign); } -void AtPropagator::PropagateToPlane(const Plane3D &plane) +void AtPropagator::PropagateToPlane(const Plane3D &plane, AtStepper &stepper) { LOG(info) << "Propagating to plane: " << plane; auto KE_initial = Kinematics::KE(fMom, fMass); + stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; + while (true) { - fLastPos = fPos; - fLastMom = fMom; - LOG(info) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - RK4Step(fH); + auto result = stepper.Step(fH, fPos, fMom); + if (!result.success) { + LOG(error) << "Integration step failed, aborting propagation."; + return; // Abort propagation if step failed + } + CopyFromState(result); // Copy the new state from the stepper bool reachedMeasurementPoint = IntersectedPlane(plane); bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; @@ -378,9 +116,14 @@ void AtPropagator::PropagateToPlane(const Plane3D &plane) LOG(info) << "Final step size: " << finalH << " mm"; finalH = approach * 1e-3; // Convert to meters for the RK4 step - fPos = fLastPos; - fMom = fLastMom; - RK4Step(finalH); // Propagate to the measurement point + result = stepper.Step(finalH, fLastPos, fLastMom); + if (!result.success) { + LOG(error) << "Failed to propagate to measurement point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fH; // Save original step size + CopyFromState(result); // Update position and momentum to the new state + fH = origH; // Restore original step size } if (particleStopped || momentumReversed) { @@ -399,19 +142,31 @@ void AtPropagator::PropagateToPlane(const Plane3D &plane) double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; - RK4Step(h_Stop * 1e-3); + result = stepper.Step(h_Stop * 1e-3, fLastPos, fLastMom); + if (!result.success) { + LOG(error) << "Failed to propagate to stopping point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fH; // Save original step size + CopyFromState(result); // Update position and momentum to the new state + fH = origH; // Restore original step size LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; while (!IntersectedPlane(plane)) { - fScalingFactor = 0; // Turn off enregy loss. + fScalingFactor = 0; // Turn off energy loss. // If we still haven't intersected the plane, we need to adjust the step size double h = std::abs(plane.Distance(fPos)); // Reduce step size so we hit the plane if (h <= fDistTol) break; LOG(info) << "Propagating to plane after stopping with step size: " << h << " mm"; - RK4Step(h * 1e-3); // Convert mm to m for RK4 step + result = stepper.Step(h * 1e-3, fPos, fMom); + if (!result.success) { + LOG(error) << "Failed to propagate to plane after stopping, aborting."; + return; // Abort propagation if step failed + } + CopyFromState(result); // Update position and momentum to the new state LOG(info) << "New position after adjusting step size: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); } fLastMom = fMom; @@ -599,6 +354,10 @@ AtStepper::StepResult AtRK4Stepper::Step(double h, const XYZPoint &fPos, const X result.h = h; result.success = true; + LOG(debug) << "Starting RK4 step with initial position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Initial momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(debug) << "Step size (h): " << h << " m"; + auto [x_k1, p_k1] = fDeriv(fPos, fMom); // The derivative of the position is then just the unit vector of the momentum. @@ -616,14 +375,18 @@ AtStepper::StepResult AtRK4Stepper::Step(double h, const XYZPoint &fPos, const X auto [x_k4, p_k4] = fDeriv(x_4, p_4); auto dpds_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // "Force" in SI units (N) + auto dxds_SI = (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6; // Position derivative in SI units (m) + + LOG(debug) << "dp/ds (SI units): " << dpds_SI.X() << ", " << dpds_SI.Y() << ", " << dpds_SI.Z(); + LOG(debug) << "dx/ds (SI units): " << dxds_SI.X() << ", " << dxds_SI.Y() << ", " << dxds_SI.Z(); auto mom_SI = fReltoSImom * fMom; mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) result.mom = mom_SI / fReltoSImom; // Convert back to - auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) - pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; // Update position in SI units (m - result.pos = pos_SI * 1e3; // Convert back to mm + auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) + pos_SI += dxds_SI * h; // Update position in SI units (m + result.pos = pos_SI * 1e3; // Convert back to mm return result; } diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 78e1f2694..13dc2d71e 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -59,7 +59,6 @@ class AtPropagator { using XYZVector = ROOT::Math::XYZVector; using XYZPoint = ROOT::Math::XYZPoint; using Plane3D = ROOT::Math::Plane3D; - using DistanceFunc = std::function; XYZVector fEField{0, 0, 0}; // Electric field vector XYZVector fBField{0, 0, 0}; // Magnetic field vector @@ -137,14 +136,13 @@ class AtPropagator { * @param stepper The stepper to use for propagation. */ void PropagateToPoint(const XYZPoint &point, AtStepper &stepper); - void PropagateToPointAdaptive(const XYZPoint &point); /** * @brief Propagate the particle to the given plane. * * @param plane The plane to approach. */ - void PropagateToPlane(const Plane3D &plane); + void PropagateToPlane(const Plane3D &plane, AtStepper &stepper); /** * @brief Calculate the force acting on the particle. @@ -188,41 +186,6 @@ class AtPropagator { } protected: - /** - * @brief Perform a single RK4 step for propagation. - * - * This method performs a single Runge-Kutta 4th order step to propagate the particle's state. - * Updates fPos and fMom. - * @param h Step size for the RK4 step in meters. - */ - void RK4Step(double h); - - /** - * @brief Perform a single RK4 step using the Nystrom method. - * This method performs a single Runge-Kutta 4th order step using the Nystrom method - * to propagate the particle's state. Updates fPos and fMom. - * @param h Step size for the RK4 step in meters. - */ - void RK4StepNystrom(double h); - - /** - * @brief Perform an adaptive RK4 step for propagation. - * This method performs an adaptive Runge-Kutta 4th order step to propagate the particle's state. - * Updates fPos and fMom. - * - * The error is based on the difference in the positions at the end of the step. i.e: - * \eps = \sqrt{\eps_x^2 + \eps_y^2 + \eps_z^2}, where - * \eps_x = 1/30*|x_1 - x_2|, where x_1 is using h/2 and x_2 is using h. - * - * Step size is adjust to ensure the local error is less than fDelta. - * - * @param h Step size for the RK4 step in seconds. Modified in place to reflect the new step size. - * @return True if the step was accepted, false otherwise. - */ - bool RK4StepAdaptive(double &h); - - void PropagateTo(DistanceFunc distanceFunc); - bool ReachedPOCA(const XYZPoint &point); bool IntersectedPlane(const Plane3D &plane); void CopyFromState(const AtStepper::StepResult &result) diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 55b1321a7..249f69487 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -186,6 +186,7 @@ TEST(AtPropagatorTest, PropagateToPlane_NoField) elossModel->LoadSrimTable( "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4Stepper stepper; double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -206,7 +207,7 @@ TEST(AtPropagatorTest, PropagateToPlane_NoField) XYZPoint planePoint(10, 10, 10); // Target point to propagate to 10 mm XYZVector planeNormal(1, 0, 0); // Normal vector of the plane in x-direction Plane3D plane(planeNormal, planePoint); // Create the plane - propagator.PropagateToPlane(plane); + propagator.PropagateToPlane(plane, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -223,6 +224,7 @@ TEST(AtPropagatorTest, PropagateToPlane_StoppingNoField) elossModel->LoadSrimTable( "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4Stepper stepper; double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -239,7 +241,7 @@ TEST(AtPropagatorTest, PropagateToPlane_StoppingNoField) XYZPoint planePoint(220, 0, 0); // Target point to propagate to 215 mm XYZVector planeNormal(1, 0, 0); // Normal vector of the plane in x-direction Plane3D plane(planeNormal, planePoint); // Create the plane - propagator.PropagateToPlane(plane); + propagator.PropagateToPlane(plane, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -276,7 +278,6 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); XYZPoint targetPoint(10, 0, 0); // Target point to propagate to (10 mm) - // propagator.PropagateToPointAdaptive(targetPoint); propagator.PropagateToPoint(targetPoint, stepper); auto finalPos = propagator.GetPosition(); @@ -292,7 +293,6 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - // propagator.PropagateToPointAdaptive(targetPoint); propagator.PropagateToPoint(targetPoint, stepper); finalPos = propagator.GetPosition(); From 1d563a79d558fa7978f4ee43a5d0d37af08f7b54 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 13:57:50 +0200 Subject: [PATCH 23/71] Propogate to point using common surface formalism --- AtTools/AtPropagator.cxx | 129 +++++++++++++++++++++++++++++++++-- AtTools/AtPropagator.h | 44 ++++++++++++ AtTools/AtPropagatorTest.cxx | 16 +++-- 3 files changed, 177 insertions(+), 12 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 53a7c792b..25249f206 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -69,11 +69,7 @@ bool AtPropagator::ReachedPOCA(const XYZPoint &point) auto lastDeriv = (fLastPos - point).Dot(fLastMom.Unit()); // proportional missing constants auto currDeriv = (fPos - point).Dot(fMom.Unit()); LOG(debug) << "Last Derivative: " << lastDeriv << ", Current Derivative: " << currDeriv; - return lastDeriv * currDeriv < 0; - - auto lastApproach = (fLastPos - point).R(); - auto approach = (fPos - point).R(); - return (approach > lastApproach); + return lastDeriv * currDeriv <= 0; } bool AtPropagator::IntersectedPlane(const Plane3D &plane) @@ -247,6 +243,8 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, AtStepper &stepper) if (particleStopped || momentumReversed) { LOG(info) << "------ Particle stopped before measurement point/surface------"; + result.mass = fMass; // Ensure mass is set in the result + // Calculate how far to travel before stopping double KE_last = Kinematics::KE(fLastMom, fMass); double deltaE = KE_last - fStopTol; @@ -298,6 +296,117 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, AtStepper &stepper) } // End of loop over RK4 integration } +void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) +{ + LOG(info) << "Propagating to measurement surface"; + + auto KE_initial = Kinematics::KE(fMom, fMass); + stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; + + while (true) { + LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + + auto result = stepper.Step(fH, fPos, fMom); + if (!result.success) { + LOG(error) << "Integration step failed, aborting propagation."; + return; // Abort propagation if step failed + } + CopyFromState(result); // Copy the new state from the stepper + + bool reachedMeasurementPoint = surface.PassedSurface(result); + bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; + bool momentumReversed = (fLastMom.Dot(fMom) < 0); + + if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { + // We reached the measurement surface, so we should figure out how far we are from the measurement point + LOG(info) << "------ Reached measurement surface ------"; + double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step + double approach = surface.Distance(fLastPos); + + LOG(info) << "Distance to plane: " << approach << " mm"; + LOG(info) << "Final step size: " << finalH << " mm"; + + finalH = approach * 1e-3; // Convert to meters for the RK4 step + result = stepper.Step(finalH, fLastPos, fLastMom); + if (!result.success) { + LOG(error) << "Failed to propagate to measurement point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fH; // Save original step size + CopyFromState(result); // Update position and momentum to the new state + fH = origH; // Restore original step size + } + + if (particleStopped || momentumReversed) { + // In this case the particle stopped before hitting the plane + // we should throw a warning to let the user know that there wasn't + // enough energy to reach the surface. + LOG(warning) << "------ Particle stopped before reaching measurement surface ------"; + + // Calculate how far to travel before stopping + double KE_last = Kinematics::KE(fLastMom, fMass); + double deltaE = KE_last - fStopTol; + deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss + + LOG(info) << "Last KE: " << KE_last << " MeV"; + LOG(info) << "Energy to loose to stop: " << deltaE << " MeV"; + double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm + LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; + + result = stepper.Step(h_Stop * 1e-3, fLastPos, fLastMom); + if (!result.success) { + LOG(error) << "Failed to propagate to stopping point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fH; // Save original step size + CopyFromState(result); // Update position and momentum to the new state + fH = origH; // Restore original step size + LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; + + while (surface.fClipToSurface && !surface.PassedSurface(result)) { + fScalingFactor = 0; // Turn off energy loss. + + // If we still haven't intersected the surface, we need to adjust the step size + double h = surface.Distance(fPos); // Reduce step size so we hit the surface + if (h <= fDistTol) + break; + LOG(info) << "Propagating to surface after stopping with step size: " << h << " mm"; + result = stepper.Step(h * 1e-3, fPos, fMom); + if (!result.success) { + LOG(error) << "Failed to propagate to surface after stopping, aborting."; + return; // Abort propagation if step failed + } + CopyFromState(result); // Update position and momentum to the new state + LOG(info) << "New position after adjusting step size: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + } + fLastMom = fMom; + fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped + reachedMeasurementPoint = true; + } + + if (reachedMeasurementPoint || particleStopped || momentumReversed) { + double distanceToSurface = surface.Distance(fPos); + + double KE_final = Kinematics::KE(fMom, fMass); + auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + LOG(info) << "------- End of RK4 interation ---------"; + LOG(info) << "Particle stopped: " << particleStopped; + LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; + LOG(info) << "Distance to plane: " << distanceToSurface << " mm"; + LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(info) << "Scaling factor: " << fScalingFactor; + LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + + fPos = surface.ProjectToSurface(fPos); + LOG(info) << "Projected Position on plane: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + return; + } + } // End of loop over RK4 integration +} + void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss, AtStepper &stepper) { LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; @@ -545,4 +654,14 @@ AtStepper::StepResult AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, } } } + +bool AtMeasurementPoint::PassedSurface(AtStepper::StepResult &result) const +{ + // Check if the particle has passed the measurement point + auto lastDeriv = (fPoint - result.lastPos).Dot(result.lastMom.Unit()); + auto currDeriv = (fPoint - result.pos).Dot(result.mom.Unit()); + LOG(debug) << "Last Derivative: " << lastDeriv << ", Current Derivative: " << currDeriv; + return lastDeriv * currDeriv <= 0; +} + } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 13dc2d71e..07b3fc490 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -11,6 +11,7 @@ namespace AtTools { +class AtMeasurementSurface; class AtStepper { public: struct StepResult { @@ -18,6 +19,7 @@ class AtStepper { ROOT::Math::XYZVector mom; // Momentum of the particle in MeV/c ROOT::Math::XYZPoint lastPos; // Last position of the particle in mm ROOT::Math::XYZVector lastMom; // Last momentum of the particle in MeV/c + double mass; // Mass of the particle in MeV/c^2 double h; // Step size for the step in m bool success; // Whether the step was successful }; @@ -144,6 +146,8 @@ class AtPropagator { */ void PropagateToPlane(const Plane3D &plane, AtStepper &stepper); + void PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper); + /** * @brief Calculate the force acting on the particle. * @@ -203,6 +207,9 @@ class AtPropagator { result.mom = fMom; result.lastPos = fLastPos; result.lastMom = fLastMom; + result.mass = fMass; + result.h = fH; + result.success = true; // Assume success unless proven otherwise return result; } }; @@ -216,5 +223,42 @@ class AtRK4AdaptiveStepper : public AtStepper { StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; }; +/** + * @brief Class for measurement surface in the AT-TPC. + * + * This class represents a measurement surface or point in the AT-TPC. It's used to define the stopping + * point and behavior of the propagator. + */ +class AtMeasurementSurface { +public: + bool fClipToSurface = false; // Whether to clip to the surface + + /** + * @brief Calculate the distance from the position to the surface. + */ + virtual double Distance(const ROOT::Math::XYZPoint &pos) const = 0; + + /** + * @brief Check if we have passed the surface between the last position and the current position. + */ + virtual bool PassedSurface(AtStepper::StepResult &result) const = 0; + + virtual ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const + { + return pos; // Default implementation returns the position as is + } +}; + +class AtMeasurementPoint : public AtMeasurementSurface { +protected: + ROOT::Math::XYZPoint fPoint; // The measurement point in mm + +public: + AtMeasurementPoint(const ROOT::Math::XYZPoint &point) : fPoint(point) {} + + double Distance(const ROOT::Math::XYZPoint &pos) const override { return (fPoint - pos).R(); } + + bool PassedSurface(AtStepper::StepResult &result) const override; +}; } // namespace AtTools #endif // #ifndef ATPROPAGATOR_H diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 249f69487..c945993bb 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -107,6 +107,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; + AtMeasurementPoint measurementPoint({1e3, 0, 0}); double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -120,8 +121,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - XYZPoint targetPoint(1e3, 0, 0); // Target point to propagate to (1 m) - propagator.PropagateToPoint(targetPoint, stepper); + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -135,7 +135,7 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) startMom.SetXYZ(p, 0, 0); // Reset momentum propagator.SetState(startPos, startMom); - propagator.PropagateToPoint(targetPoint, stepper); // Propagate to range + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); // Propagate to range finalPos = propagator.GetPosition(); finalMom = propagator.GetMomentum(); ASSERT_NEAR(finalPos.X(), 130, 10); // Final position in x-direction should be close to 130 mm @@ -151,6 +151,7 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; + AtMeasurementPoint measurementPoint({10, 0, 0}); double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -169,7 +170,8 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); XYZPoint targetPoint(10, 0, 0); // Target point to propagate to 10 mm - propagator.PropagateToPoint(targetPoint, stepper); + // propagator.PropagateToPoint(targetPoint, stepper); + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -259,6 +261,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4AdaptiveStepper stepper; + AtMeasurementPoint measurementPoint({10, 0, 0}); double KE = 1; // Kinetic energy in MeV double E = KE + mass_p; @@ -277,8 +280,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - XYZPoint targetPoint(10, 0, 0); // Target point to propagate to (10 mm) - propagator.PropagateToPoint(targetPoint, stepper); + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -293,7 +295,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - propagator.PropagateToPoint(targetPoint, stepper); + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); finalPos = propagator.GetPosition(); finalMom = propagator.GetMomentum(); From dcccdfdc89c34d87d6a4f56d4920542448c05095 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 14:01:43 +0200 Subject: [PATCH 24/71] Migrate fixed eloss to new interfaces --- AtTools/AtPropagator.cxx | 110 ++------------------------------------- AtTools/AtPropagator.h | 10 +--- 2 files changed, 5 insertions(+), 115 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 25249f206..17d734816 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -194,108 +194,6 @@ void AtPropagator::PropagateToPlane(const Plane3D &plane, AtStepper &stepper) } // End of loop over RK4 integration } -void AtPropagator::PropagateToPoint(const XYZPoint &point, AtStepper &stepper) -{ - LOG(info) << "Propagating to point: " << point; - - auto KE_initial = Kinematics::KE(fMom, fMass); - stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; - - while (true) { - LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - - auto result = stepper.Step(fH, fPos, fMom); - if (!result.success) { - LOG(error) << "Integration step failed, aborting propagation."; - return; // Abort propagation if step failed - } - CopyFromState(result); // Copy the new state from the stepper - - if (!ReachedPOCA(point)) - continue; - - bool reachedMeasurementPoint = ReachedPOCA(point); - bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - bool momentumReversed = (fLastMom.Dot(fMom) < 0); - - if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { - // We reached the measurement point, so we should figure out how far we are from the measurement point - // and update that remaining amount - LOG(info) << "------ Reached measurement point------"; - double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step - double approach = (fLastPos - point).R(); // Distance to the measurement point - LOG(info) << "Distance to measurement point: " << approach << " mm"; - LOG(info) << "Final step size: " << finalH << " mm"; - - finalH = approach * 1e-3; // Convert to meters for the RK4 step - result = stepper.Step(finalH, fLastPos, fLastMom); - if (!result.success) { - LOG(error) << "Failed to propagate to measurement point, aborting."; - return; // Abort propagation if step failed - } - auto origH = fH; // Save original step size - CopyFromState(result); // Update position and momentum to the new state - fH = origH; // Restore original step size - } - - // If we stopped, then we should figure out about where we stopped assuming linear de/dx over this last step - if (particleStopped || momentumReversed) { - LOG(info) << "------ Particle stopped before measurement point/surface------"; - - result.mass = fMass; // Ensure mass is set in the result - - // Calculate how far to travel before stopping - double KE_last = Kinematics::KE(fLastMom, fMass); - double deltaE = KE_last - fStopTol; - deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss - double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm - - LOG(info) << "KE at last point: " << KE_last << " MeV"; - LOG(info) << "Energy to lose to stop: " << deltaE << " MeV"; - LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; - LOG(info) << "dE/dx at last point: " << fELossModel->GetdEdx(KE_last) << " MeV/mm"; - - result = stepper.Step(h_Stop * 1e-3, fLastPos, fLastMom); - if (!result.success) { - LOG(error) << "Failed to propagate to stopping point, aborting."; - return; // Abort propagation if step failed - } - auto origH = fH; // Save original step size - CopyFromState(result); // Update position and momentum to the new state - fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped - fLastMom = fMom; // Update last momentum to zero - fH = origH; // Restore original step size - - LOG(info) << "Final Position after stopping: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum after stopping: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - } - - if (reachedMeasurementPoint || particleStopped || momentumReversed) { - // Undo the last step since we were closer last time. - double lastApproach = (fLastPos - point).R(); - double approach = (fPos - point).R(); - - double KE_final = Kinematics::KE(fMom, fMass); - auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - LOG(info) << "------- End of RK4 interation ---------"; - LOG(info) << "Particle stopped: " << particleStopped; - LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Last approach: " << lastApproach << " Current approach: " << approach; - LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; - LOG(info) << "Scaling factor: " << fScalingFactor; - LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - LOG(info) << "Last Position: " << fLastPos.X() << ", " << fLastPos.Y() << ", " << fLastPos.Z(); - LOG(info) << "Last Momentum: " << fLastMom.X() << ", " << fLastMom.Y() << ", " << fLastMom.Z(); - - // fPos = fLastPos; - // fMom = fLastMom; - return; - } - } // End of loop over RK4 integration -} - void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) { LOG(info) << "Propagating to measurement surface"; @@ -407,13 +305,13 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur } // End of loop over RK4 integration } -void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss, AtStepper &stepper) +void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, double eLoss, AtStepper &stepper) { - LOG(info) << "Propagating to point: " << point << " with eLoss: " << eLoss; + LOG(info) << "Propagating to surface with eLoss: " << eLoss; if (eLoss == 0) { LOG(warn) << "No energy loss specified, propagating without energy loss adjustment."; - PropagateToPoint(point, stepper); + PropagateToMeasurementSurface(surface, stepper); return; } @@ -436,7 +334,7 @@ void AtPropagator::PropagateToPoint(const XYZPoint &point, double eLoss, AtStepp } iterations++; - PropagateToPoint(point, stepper); // Propagate without energy loss adjustment + PropagateToMeasurementSurface(surface, stepper); // Propagate without energy loss adjustment double KE_final = Kinematics::KE(fMom, fMass); calc_eLoss = KE_initial - KE_final; // Energy loss in MeV diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 07b3fc490..8e8e33c31 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -129,15 +129,7 @@ class AtPropagator { * @param point The point to approach. * @param eLoss If not 0, constrain the energy loss to this value by adjusting fScalingFactor. */ - void PropagateToPoint(const XYZPoint &point, double eLoss, AtStepper &stepper); - - /** - * @brief Propagate the particle to the point of closest approach to the given point. - * - * @param point The point to approach. - * @param stepper The stepper to use for propagation. - */ - void PropagateToPoint(const XYZPoint &point, AtStepper &stepper); + void PropagateToMeasurementSurface(const AtMeasurementSurface &point, double eLoss, AtStepper &stepper); /** * @brief Propagate the particle to the given plane. From df4ec309bfd8f2eba755b2365672e0173f54da26 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 14:08:23 +0200 Subject: [PATCH 25/71] Plans using interface --- AtTools/AtPropagator.cxx | 8 ++++++++ AtTools/AtPropagator.h | 10 ++++++++++ AtTools/AtPropagatorTest.cxx | 6 ++++-- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 17d734816..4d87820b0 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -562,4 +562,12 @@ bool AtMeasurementPoint::PassedSurface(AtStepper::StepResult &result) const return lastDeriv * currDeriv <= 0; } +bool AtMeasurementPlane::PassedSurface(AtStepper::StepResult &result) const +{ + // Check if the particle has crossed the plane this step. + auto prevSign = fPlane.Distance(result.lastPos) > 0 ? 1 : -1; + auto currSign = fPlane.Distance(result.pos) > 0 ? 1 : -1; + return (prevSign != currSign); +} + } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 8e8e33c31..5d52acd1e 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -252,5 +252,15 @@ class AtMeasurementPoint : public AtMeasurementSurface { bool PassedSurface(AtStepper::StepResult &result) const override; }; + +class AtMeasurementPlane : public AtMeasurementSurface { +protected: + ROOT::Math::Plane3D fPlane; // The measurement plane +public: + AtMeasurementPlane(const ROOT::Math::Plane3D &plane) : fPlane(plane) { fClipToSurface = true; } + + double Distance(const ROOT::Math::XYZPoint &pos) const override { return std::abs(fPlane.Distance(pos)); } + bool PassedSurface(AtStepper::StepResult &result) const override; +}; } // namespace AtTools #endif // #ifndef ATPROPAGATOR_H diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index c945993bb..75467531b 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -209,7 +209,8 @@ TEST(AtPropagatorTest, PropagateToPlane_NoField) XYZPoint planePoint(10, 10, 10); // Target point to propagate to 10 mm XYZVector planeNormal(1, 0, 0); // Normal vector of the plane in x-direction Plane3D plane(planeNormal, planePoint); // Create the plane - propagator.PropagateToPlane(plane, stepper); + AtMeasurementPlane measurementPlane(plane); + propagator.PropagateToMeasurementSurface(measurementPlane, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); @@ -243,7 +244,8 @@ TEST(AtPropagatorTest, PropagateToPlane_StoppingNoField) XYZPoint planePoint(220, 0, 0); // Target point to propagate to 215 mm XYZVector planeNormal(1, 0, 0); // Normal vector of the plane in x-direction Plane3D plane(planeNormal, planePoint); // Create the plane - propagator.PropagateToPlane(plane, stepper); + AtMeasurementPlane measurementPlane(plane); + propagator.PropagateToMeasurementSurface(measurementPlane, stepper); auto finalPos = propagator.GetPosition(); auto finalMom = propagator.GetMomentum(); From c847df19c2021a79334649b48e49372a8e47ef82 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 14:28:14 +0200 Subject: [PATCH 26/71] Refactor and only project if at surface --- AtTools/AtPropagator.cxx | 163 +++++---------------------------------- AtTools/AtPropagator.h | 97 ++++++++++------------- 2 files changed, 62 insertions(+), 198 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 4d87820b0..25bc5751a 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -61,139 +61,6 @@ AtPropagator::XYZVector AtPropagator::d2xds2(const XYZPoint &pos, const XYZVecto return 1 / p * (dpds_vec - phat * (phat.Dot(dpds_vec))); // Second derivative of position w.r.t. arc length } -bool AtPropagator::ReachedPOCA(const XYZPoint &point) -{ - // Here we need to check if we are getting closer or further away from the POCA. - // We may walk right past it so need to look for a change in the sign of the derivative or - // something like that. - auto lastDeriv = (fLastPos - point).Dot(fLastMom.Unit()); // proportional missing constants - auto currDeriv = (fPos - point).Dot(fMom.Unit()); - LOG(debug) << "Last Derivative: " << lastDeriv << ", Current Derivative: " << currDeriv; - return lastDeriv * currDeriv <= 0; -} - -bool AtPropagator::IntersectedPlane(const Plane3D &plane) -{ - // Check if the particle has crossed the plane this step. - auto prevSign = plane.Distance(fLastPos) > 0 ? 1 : -1; - auto currSign = plane.Distance(fPos) > 0 ? 1 : -1; - return (prevSign != currSign); -} - -void AtPropagator::PropagateToPlane(const Plane3D &plane, AtStepper &stepper) -{ - LOG(info) << "Propagating to plane: " << plane; - - auto KE_initial = Kinematics::KE(fMom, fMass); - stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; - - while (true) { - LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - - auto result = stepper.Step(fH, fPos, fMom); - if (!result.success) { - LOG(error) << "Integration step failed, aborting propagation."; - return; // Abort propagation if step failed - } - CopyFromState(result); // Copy the new state from the stepper - - bool reachedMeasurementPoint = IntersectedPlane(plane); - bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - bool momentumReversed = (fLastMom.Dot(fMom) < 0); - - if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { - // We reached the measurement point, so we should figure out how far we are from the measurement point - LOG(info) << "------ Reached measurement point ------"; - double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step - double approach = std::abs(plane.Distance(fLastPos)); - - LOG(info) << "Distance to plane: " << approach << " mm"; - LOG(info) << "Final step size: " << finalH << " mm"; - - finalH = approach * 1e-3; // Convert to meters for the RK4 step - result = stepper.Step(finalH, fLastPos, fLastMom); - if (!result.success) { - LOG(error) << "Failed to propagate to measurement point, aborting."; - return; // Abort propagation if step failed - } - auto origH = fH; // Save original step size - CopyFromState(result); // Update position and momentum to the new state - fH = origH; // Restore original step size - } - - if (particleStopped || momentumReversed) { - // In this case the particle stopped before hitting the plane - // we should throw a warning to let the user know that there wasn't - // enough energy to reach the plane. - LOG(warning) << "------ Particle stopped before intersecting plane ------"; - - // Calculate how far to travel before stopping - double KE_last = Kinematics::KE(fLastMom, fMass); - double deltaE = KE_last - fStopTol; - deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss - - LOG(info) << "Last KE: " << KE_last << " MeV"; - LOG(info) << "Energy to loose to stop: " << deltaE << " MeV"; - double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm - LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; - - result = stepper.Step(h_Stop * 1e-3, fLastPos, fLastMom); - if (!result.success) { - LOG(error) << "Failed to propagate to stopping point, aborting."; - return; // Abort propagation if step failed - } - auto origH = fH; // Save original step size - CopyFromState(result); // Update position and momentum to the new state - fH = origH; // Restore original step size - LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; - - while (!IntersectedPlane(plane)) { - fScalingFactor = 0; // Turn off energy loss. - - // If we still haven't intersected the plane, we need to adjust the step size - double h = std::abs(plane.Distance(fPos)); // Reduce step size so we hit the plane - if (h <= fDistTol) - break; - LOG(info) << "Propagating to plane after stopping with step size: " << h << " mm"; - result = stepper.Step(h * 1e-3, fPos, fMom); - if (!result.success) { - LOG(error) << "Failed to propagate to plane after stopping, aborting."; - return; // Abort propagation if step failed - } - CopyFromState(result); // Update position and momentum to the new state - LOG(info) << "New position after adjusting step size: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - } - fLastMom = fMom; - fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped - reachedMeasurementPoint = true; - } - - if (reachedMeasurementPoint || particleStopped || momentumReversed) { - double distanceToPlane = std::abs(plane.Distance(fPos)); - - double KE_final = Kinematics::KE(fMom, fMass); - auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - LOG(info) << "------- End of RK4 interation ---------"; - LOG(info) << "Particle stopped: " << particleStopped; - LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Distance to plane: " << distanceToPlane << " mm"; - LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; - LOG(info) << "Scaling factor: " << fScalingFactor; - LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - - // Project the position onto the plane. Cannot use ProjectOnPlane since it is templated in such - // a way that it can't separate our XYZPoint and its internal XYZPoint. - double d = plane.Distance(fPos); // Distance from the point to the plane - fPos = XYZPoint(fPos.X() - plane.A() * d, fPos.Y() - plane.B() * d, fPos.Z() - plane.C() * d); - LOG(info) << "Projected Position on plane: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - return; - } - } // End of loop over RK4 integration -} - void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) { LOG(info) << "Propagating to measurement surface"; @@ -263,13 +130,15 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; - while (surface.fClipToSurface && !surface.PassedSurface(result)) { + while (surface.fClipToSurface) { fScalingFactor = 0; // Turn off energy loss. // If we still haven't intersected the surface, we need to adjust the step size double h = surface.Distance(fPos); // Reduce step size so we hit the surface - if (h <= fDistTol) + if (h <= fDistTol || surface.PassedSurface(result)) { + reachedMeasurementPoint = true; break; + } LOG(info) << "Propagating to surface after stopping with step size: " << h << " mm"; result = stepper.Step(h * 1e-3, fPos, fMom); if (!result.success) { @@ -281,7 +150,6 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur } fLastMom = fMom; fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped - reachedMeasurementPoint = true; } if (reachedMeasurementPoint || particleStopped || momentumReversed) { @@ -297,7 +165,10 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(info) << "Scaling factor: " << fScalingFactor; LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - fPos = surface.ProjectToSurface(fPos); + // If we reached the measurement surface, we should project the position onto the surface + if (reachedMeasurementPoint) { + fPos = surface.ProjectToSurface(fPos); + } LOG(info) << "Projected Position on plane: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); return; @@ -352,10 +223,10 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur fScalingFactor = 1; // Reset scaling factor after convergence } -AtStepper::StepResult AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const +AtStepper::StepState AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const { // Take h to be the step size in m. - StepResult result; + StepState result; result.lastPos = fPos; result.lastMom = fMom; result.h = h; @@ -398,10 +269,10 @@ AtStepper::StepResult AtRK4Stepper::Step(double h, const XYZPoint &fPos, const X return result; } -AtStepper::StepResult AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const +AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const { // Take h to be the step size in m. - StepResult result; + StepState result; result.lastPos = fPos; result.lastMom = fMom; result.h = h; @@ -553,7 +424,7 @@ AtStepper::StepResult AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, } } -bool AtMeasurementPoint::PassedSurface(AtStepper::StepResult &result) const +bool AtMeasurementPoint::PassedSurface(AtStepper::StepState &result) const { // Check if the particle has passed the measurement point auto lastDeriv = (fPoint - result.lastPos).Dot(result.lastMom.Unit()); @@ -562,7 +433,7 @@ bool AtMeasurementPoint::PassedSurface(AtStepper::StepResult &result) const return lastDeriv * currDeriv <= 0; } -bool AtMeasurementPlane::PassedSurface(AtStepper::StepResult &result) const +bool AtMeasurementPlane::PassedSurface(AtStepper::StepState &result) const { // Check if the particle has crossed the plane this step. auto prevSign = fPlane.Distance(result.lastPos) > 0 ? 1 : -1; @@ -570,4 +441,10 @@ bool AtMeasurementPlane::PassedSurface(AtStepper::StepResult &result) const return (prevSign != currSign); } +ROOT::Math::XYZPoint AtMeasurementPlane::ProjectToSurface(const ROOT::Math::XYZPoint &pos) const +{ + // Project the position onto the measurement plane + auto dist = fPlane.Distance(pos); + return pos - dist * fPlane.Normal(); +} } // namespace AtTools diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 5d52acd1e..cef8486d5 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -14,14 +14,15 @@ namespace AtTools { class AtMeasurementSurface; class AtStepper { public: - struct StepResult { - ROOT::Math::XYZPoint pos; // Position of the particle in mm - ROOT::Math::XYZVector mom; // Momentum of the particle in MeV/c - ROOT::Math::XYZPoint lastPos; // Last position of the particle in mm - ROOT::Math::XYZVector lastMom; // Last momentum of the particle in MeV/c - double mass; // Mass of the particle in MeV/c^2 - double h; // Step size for the step in m - bool success; // Whether the step was successful + struct StepState { + ROOT::Math::XYZPoint pos; /// Position of the particle in mm + ROOT::Math::XYZVector mom; /// Momentum of the particle in MeV/c + ROOT::Math::XYZPoint lastPos; /// Last position of the particle in mm + ROOT::Math::XYZVector lastMom; /// Last momentum of the particle in MeV/c + double mass; /// Mass of the particle in MeV/c^2 + double h; /// Step size to use in m + double hUsed; /// Step size used in this step in m + bool success; /// Whether the step was successful }; /** * @brief Function type defining the derivative of the position and momentum w.r.t. distance. @@ -39,12 +40,35 @@ class AtStepper { DerivFunc fDeriv; - virtual StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const = 0; + virtual StepState Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const = 0; protected: static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) }; +/** + * @brief Class for measurement surface in the AT-TPC. + * + * This class represents a measurement surface or point in the AT-TPC. It's used to define the stopping + * point and behavior of the propagator. + */ +class AtMeasurementSurface { +public: + bool fClipToSurface = false; // Whether to clip to the surface + + /** + * @brief Calculate the distance from the position to the surface. + */ + virtual double Distance(const ROOT::Math::XYZPoint &pos) const = 0; + + /** + * @brief Check if we have passed the surface between the last position and the current position. + */ + virtual bool PassedSurface(AtStepper::StepState &result) const = 0; + + virtual ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const = 0; +}; + /** * @brief Class for propagating particles through a medium. * @@ -70,8 +94,6 @@ class AtPropagator { // Internal state variables for the propagator double fH = 1e-4; /// Step size for propagation in m - double fDelta = 1e-3; /// Relative error tolerance for adaptive step size. 10^-3 means each 1m of propagation - /// introduces at most 1mm of error. double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped (MeV) double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) @@ -114,7 +136,6 @@ class AtPropagator { fMom = mom; } - void SetDelta(double delta) { fDelta = delta; } void SetH(double h) { fH = h; } XYZPoint GetPosition() const { return fPos; } @@ -131,13 +152,6 @@ class AtPropagator { */ void PropagateToMeasurementSurface(const AtMeasurementSurface &point, double eLoss, AtStepper &stepper); - /** - * @brief Propagate the particle to the given plane. - * - * @param plane The plane to approach. - */ - void PropagateToPlane(const Plane3D &plane, AtStepper &stepper); - void PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper); /** @@ -182,9 +196,7 @@ class AtPropagator { } protected: - bool ReachedPOCA(const XYZPoint &point); - bool IntersectedPlane(const Plane3D &plane); - void CopyFromState(const AtStepper::StepResult &result) + void CopyFromState(const AtStepper::StepState &result) { fPos = result.pos; fMom = result.mom; @@ -192,9 +204,9 @@ class AtPropagator { fLastMom = result.lastMom; fH = result.h; } - AtStepper::StepResult CopyToState() const + AtStepper::StepState CopyToState() const { - AtStepper::StepResult result; + AtStepper::StepState result; result.pos = fPos; result.mom = fMom; result.lastPos = fLastPos; @@ -208,37 +220,11 @@ class AtPropagator { class AtRK4Stepper : public AtStepper { public: - StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; + StepState Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; }; class AtRK4AdaptiveStepper : public AtStepper { public: - StepResult Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; -}; - -/** - * @brief Class for measurement surface in the AT-TPC. - * - * This class represents a measurement surface or point in the AT-TPC. It's used to define the stopping - * point and behavior of the propagator. - */ -class AtMeasurementSurface { -public: - bool fClipToSurface = false; // Whether to clip to the surface - - /** - * @brief Calculate the distance from the position to the surface. - */ - virtual double Distance(const ROOT::Math::XYZPoint &pos) const = 0; - - /** - * @brief Check if we have passed the surface between the last position and the current position. - */ - virtual bool PassedSurface(AtStepper::StepResult &result) const = 0; - - virtual ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const - { - return pos; // Default implementation returns the position as is - } + StepState Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; }; class AtMeasurementPoint : public AtMeasurementSurface { @@ -249,8 +235,8 @@ class AtMeasurementPoint : public AtMeasurementSurface { AtMeasurementPoint(const ROOT::Math::XYZPoint &point) : fPoint(point) {} double Distance(const ROOT::Math::XYZPoint &pos) const override { return (fPoint - pos).R(); } - - bool PassedSurface(AtStepper::StepResult &result) const override; + bool PassedSurface(AtStepper::StepState &result) const override; + ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override { return fPoint; } }; class AtMeasurementPlane : public AtMeasurementSurface { @@ -260,7 +246,8 @@ class AtMeasurementPlane : public AtMeasurementSurface { AtMeasurementPlane(const ROOT::Math::Plane3D &plane) : fPlane(plane) { fClipToSurface = true; } double Distance(const ROOT::Math::XYZPoint &pos) const override { return std::abs(fPlane.Distance(pos)); } - bool PassedSurface(AtStepper::StepResult &result) const override; + bool PassedSurface(AtStepper::StepState &result) const override; + ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override; }; } // namespace AtTools #endif // #ifndef ATPROPAGATOR_H From 1a8bdcc0533ec9d2ee117edd333afc43d7e95f31 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 14:34:13 +0200 Subject: [PATCH 27/71] Refactor state struct --- AtTools/AtPropagator.cxx | 34 +++++++++++++++++----------------- AtTools/AtPropagator.h | 36 +++++++++++++++++++----------------- 2 files changed, 36 insertions(+), 34 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 25bc5751a..b99be9a6f 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -227,8 +227,8 @@ AtStepper::StepState AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XY { // Take h to be the step size in m. StepState result; - result.lastPos = fPos; - result.lastMom = fMom; + result.fLastPos = fPos; + result.fLastMom = fMom; result.h = h; result.success = true; @@ -259,12 +259,12 @@ AtStepper::StepState AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XY LOG(debug) << "dx/ds (SI units): " << dxds_SI.X() << ", " << dxds_SI.Y() << ", " << dxds_SI.Z(); auto mom_SI = fReltoSImom * fMom; - mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) - result.mom = mom_SI / fReltoSImom; // Convert back to + mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) + result.fMom = mom_SI / fReltoSImom; // Convert back to - auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) - pos_SI += dxds_SI * h; // Update position in SI units (m - result.pos = pos_SI * 1e3; // Convert back to mm + auto pos_SI = fPos * 1e-3; // Convert position to SI units (m) + pos_SI += dxds_SI * h; // Update position in SI units (m + result.fPos = pos_SI * 1e3; // Convert back to mm return result; } @@ -273,8 +273,8 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, { // Take h to be the step size in m. StepState result; - result.lastPos = fPos; - result.lastMom = fMom; + result.fLastPos = fPos; + result.fLastMom = fMom; result.h = h; result.success = true; @@ -395,13 +395,13 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, // We now know the local error at this point. Now we need to decide to accept the point or not. if (err <= 1.0) { // Accept the step - result.pos = x_5_mm; // Update position in mm - result.mom = p_5_MeV; // Update momentum in MeV/c + result.fPos = x_5_mm; // Update position in mm + result.fMom = p_5_MeV; // Update momentum in MeV/c LOG(info) << "Accepted step with error: " << err; LOG(info) << "Step size: " << h << " m"; LOG(info) << "New step size: " << hNew << " m"; - LOG(info) << "New Position: " << result.pos.X() << ", " << result.pos.Y() << ", " << result.pos.Z(); - LOG(info) << "New Momentum: " << result.mom.X() << ", " << result.mom.Y() << ", " << result.mom.Z(); + LOG(info) << "New Position: " << result.fPos.X() << ", " << result.fPos.Y() << ", " << result.fPos.Z(); + LOG(info) << "New Momentum: " << result.fMom.X() << ", " << result.fMom.Y() << ", " << result.fMom.Z(); // Adjust the step size for the next iteration result.h = hNew; @@ -427,8 +427,8 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, bool AtMeasurementPoint::PassedSurface(AtStepper::StepState &result) const { // Check if the particle has passed the measurement point - auto lastDeriv = (fPoint - result.lastPos).Dot(result.lastMom.Unit()); - auto currDeriv = (fPoint - result.pos).Dot(result.mom.Unit()); + auto lastDeriv = (fPoint - result.fLastPos).Dot(result.fLastMom.Unit()); + auto currDeriv = (fPoint - result.fPos).Dot(result.fMom.Unit()); LOG(debug) << "Last Derivative: " << lastDeriv << ", Current Derivative: " << currDeriv; return lastDeriv * currDeriv <= 0; } @@ -436,8 +436,8 @@ bool AtMeasurementPoint::PassedSurface(AtStepper::StepState &result) const bool AtMeasurementPlane::PassedSurface(AtStepper::StepState &result) const { // Check if the particle has crossed the plane this step. - auto prevSign = fPlane.Distance(result.lastPos) > 0 ? 1 : -1; - auto currSign = fPlane.Distance(result.pos) > 0 ? 1 : -1; + auto prevSign = fPlane.Distance(result.fLastPos) > 0 ? 1 : -1; + auto currSign = fPlane.Distance(result.fPos) > 0 ? 1 : -1; return (prevSign != currSign); } diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index cef8486d5..2c7dafb6d 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -15,14 +15,14 @@ class AtMeasurementSurface; class AtStepper { public: struct StepState { - ROOT::Math::XYZPoint pos; /// Position of the particle in mm - ROOT::Math::XYZVector mom; /// Momentum of the particle in MeV/c - ROOT::Math::XYZPoint lastPos; /// Last position of the particle in mm - ROOT::Math::XYZVector lastMom; /// Last momentum of the particle in MeV/c - double mass; /// Mass of the particle in MeV/c^2 - double h; /// Step size to use in m - double hUsed; /// Step size used in this step in m - bool success; /// Whether the step was successful + ROOT::Math::XYZPoint fPos; /// Position of the particle in mm + ROOT::Math::XYZVector fMom; /// Momentum of the particle in MeV/c + ROOT::Math::XYZPoint fLastPos; /// Last position of the particle in mm + ROOT::Math::XYZVector fLastMom; /// Last momentum of the particle in MeV/c + double fMass; /// Mass of the particle in MeV/c^2 + double h; /// Step size to use in m + double hUsed; /// Step size used in this step in m + bool success; /// Whether the step was successful }; /** * @brief Function type defining the derivative of the position and momentum w.r.t. distance. @@ -99,6 +99,8 @@ class AtPropagator { double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) double fScalingFactor = 1.0; /// Scaling factor for energy loss + AtStepper::StepState fState; // Current state of the particle + XYZPoint fPos; // Current position of the particle in mm XYZVector fMom; // Current momentum of the particle in MeV/c @@ -198,20 +200,20 @@ class AtPropagator { protected: void CopyFromState(const AtStepper::StepState &result) { - fPos = result.pos; - fMom = result.mom; - fLastPos = result.lastPos; - fLastMom = result.lastMom; + fPos = result.fPos; + fMom = result.fMom; + fLastPos = result.fLastPos; + fLastMom = result.fLastMom; fH = result.h; } AtStepper::StepState CopyToState() const { AtStepper::StepState result; - result.pos = fPos; - result.mom = fMom; - result.lastPos = fLastPos; - result.lastMom = fLastMom; - result.mass = fMass; + result.fPos = fPos; + result.fMom = fMom; + result.fLastPos = fLastPos; + result.fLastMom = fLastMom; + result.fMass = fMass; result.h = fH; result.success = true; // Assume success unless proven otherwise return result; From ec81eb11b81e7eb4980cb54a06e5d5fede3c1a67 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 15:08:48 +0200 Subject: [PATCH 28/71] Refactor to passing around state struct Old variables need removed along with debug statements. --- AtTools/AtPropagator.cxx | 127 ++++++++++++++++++++++++--------------- AtTools/AtPropagator.h | 60 +++++++++--------- 2 files changed, 106 insertions(+), 81 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index b99be9a6f..4339d0e01 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -65,41 +65,51 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur { LOG(info) << "Propagating to measurement surface"; - auto KE_initial = Kinematics::KE(fMom, fMass); + auto KE_initial = Kinematics::KE(fState.fMom, fState.fMass); stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; while (true) { - LOG(debug) << "Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(info) << "Position: " << GetPosition().X() << ", " << GetPosition().Y() << ", " << GetPosition().Z(); + LOG(info) << "Momentum: " << GetMomentum().X() << ", " << GetMomentum().Y() << ", " << GetMomentum().Z(); - auto result = stepper.Step(fH, fPos, fMom); + auto result = stepper.Step(fState); if (!result.success) { LOG(error) << "Integration step failed, aborting propagation."; return; // Abort propagation if step failed } CopyFromState(result); // Copy the new state from the stepper + fState = result; // Update the internal state - bool reachedMeasurementPoint = surface.PassedSurface(result); - bool particleStopped = Kinematics::KE(fMom, fMass) < fStopTol; - bool momentumReversed = (fLastMom.Dot(fMom) < 0); + bool reachedMeasurementPoint = surface.PassedSurface(fState); + bool particleStopped = Kinematics::KE(fState.fMom, fState.fMass) < fStopTol; + bool momentumReversed = (fState.fLastMom.Dot(fState.fMom) < 0); if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { // We reached the measurement surface, so we should figure out how far we are from the measurement point LOG(info) << "------ Reached measurement surface ------"; - double finalH = (fLastPos - fPos).R(); // Distance traveled in the last step - double approach = surface.Distance(fLastPos); + double finalH = (fState.fLastPos - fState.fPos).R(); // Distance traveled in the last step + double approach = surface.Distance(fState.fLastPos); LOG(info) << "Distance to plane: " << approach << " mm"; LOG(info) << "Final step size: " << finalH << " mm"; - - finalH = approach * 1e-3; // Convert to meters for the RK4 step - result = stepper.Step(finalH, fLastPos, fLastMom); + LOG(info) << "Current position: " << fState.fLastPos.X() << ", " << fState.fLastPos.Y() << ", " + << fState.fLastPos.Z(); + LOG(info) << "Current momentum: " << fState.fLastMom.X() << ", " << fState.fLastMom.Y() << ", " + << fState.fLastMom.Z(); + + finalH = approach * 1e-3; // Convert to meters for the RK4 step + fState.h = finalH; // Set the step size to the distance to the surface + fState.fPos = fState.fLastPos; // Set position to last position + fState.fMom = fState.fLastMom; // Set momentum to last momentum + result = stepper.Step(fState); if (!result.success) { LOG(error) << "Failed to propagate to measurement point, aborting."; return; // Abort propagation if step failed } auto origH = fH; // Save original step size CopyFromState(result); // Update position and momentum to the new state + fState = result; // Update the internal state + fState.h = origH; // Restore original step size fH = origH; // Restore original step size } @@ -110,7 +120,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(warning) << "------ Particle stopped before reaching measurement surface ------"; // Calculate how far to travel before stopping - double KE_last = Kinematics::KE(fLastMom, fMass); + double KE_last = Kinematics::KE(fState.fLastMom, fState.fMass); double deltaE = KE_last - fStopTol; deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss @@ -119,58 +129,69 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; - result = stepper.Step(h_Stop * 1e-3, fLastPos, fLastMom); + fState.h = h_Stop * 1e-3; // Convert to meters for the RK4 step + fState.fPos = fState.fLastPos; // Set position to last position + fState.fMom = fState.fLastMom; // Set momentum to last momentum + result = stepper.Step(fState); if (!result.success) { LOG(error) << "Failed to propagate to stopping point, aborting."; return; // Abort propagation if step failed } auto origH = fH; // Save original step size CopyFromState(result); // Update position and momentum to the new state + fState = result; // Update the internal state + fState.h = origH; // Restore original step size fH = origH; // Restore original step size - LOG(info) << "Propagated to stopping point: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Energy after stopping: " << Kinematics::KE(fMom, fMass) << " MeV"; + LOG(info) << "Propagated to stopping point: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " + << fState.fPos.Z(); + LOG(info) << "Energy after stopping: " << Kinematics::KE(fState.fMom, fState.fMass) << " MeV"; while (surface.fClipToSurface) { fScalingFactor = 0; // Turn off energy loss. // If we still haven't intersected the surface, we need to adjust the step size - double h = surface.Distance(fPos); // Reduce step size so we hit the surface + double h = surface.Distance(fState.fPos); // Reduce step size so we hit the surface if (h <= fDistTol || surface.PassedSurface(result)) { reachedMeasurementPoint = true; break; } LOG(info) << "Propagating to surface after stopping with step size: " << h << " mm"; - result = stepper.Step(h * 1e-3, fPos, fMom); + fState.h = h * 1e-3; // Convert to meters for the RK4 step + result = stepper.Step(fState); if (!result.success) { LOG(error) << "Failed to propagate to surface after stopping, aborting."; return; // Abort propagation if step failed } CopyFromState(result); // Update position and momentum to the new state - LOG(info) << "New position after adjusting step size: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + fState = result; // Update the internal state + LOG(info) << "New position after adjusting step size: " << fState.fPos.X() << ", " << fState.fPos.Y() + << ", " << fState.fPos.Z(); } - fLastMom = fMom; - fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped + fState.fLastMom = fState.fMom; + fState.fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped } if (reachedMeasurementPoint || particleStopped || momentumReversed) { - double distanceToSurface = surface.Distance(fPos); + double distanceToSurface = surface.Distance(fState.fPos); - double KE_final = Kinematics::KE(fMom, fMass); + double KE_final = Kinematics::KE(fState.fMom, fState.fMass); + LOG(info) << "Initial KE" << KE_initial << " MeV"; + LOG(info) << "Final KE: " << KE_final << " MeV"; auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV LOG(info) << "------- End of RK4 interation ---------"; LOG(info) << "Particle stopped: " << particleStopped; LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Distance to plane: " << distanceToSurface << " mm"; + LOG(info) << "Distance to surface: " << distanceToSurface << " mm"; LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; LOG(info) << "Scaling factor: " << fScalingFactor; - LOG(info) << "Final Position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Final Position: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); // If we reached the measurement surface, we should project the position onto the surface if (reachedMeasurementPoint) { - fPos = surface.ProjectToSurface(fPos); + fState.fPos = surface.ProjectToSurface(fState.fPos); } - LOG(info) << "Projected Position on plane: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Final Momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(info) << "Projected on surface: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); + LOG(info) << "Final Momentum: " << fState.fMom.X() << ", " << fState.fMom.Y() << ", " << fState.fMom.Z(); return; } } // End of loop over RK4 integration @@ -188,13 +209,13 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur int iterations = 0; double calc_eLoss = 0; - double KE_initial = Kinematics::KE(fMom, fMass); - auto initialMom = fMom; // Save initial momentum for energy loss calculation - auto initialPos = fPos; // Save initial position for energy loss calculation + double KE_initial = Kinematics::KE(fState.fMom, fState.fMass); + auto initialMom = fState.fMom; // Save initial momentum for energy loss calculation + auto initialPos = fState.fPos; // Save initial position for energy loss calculation while (std::abs(calc_eLoss - eLoss) > 1e-4) { - fMom = initialMom; // Reset position and momentum to initial values for the next iteration - fPos = initialPos; + fState.fMom = initialMom; // Reset position and momentum to initial values for the next iteration + fState.fPos = initialPos; LOG(debug) << "Running iteration " << iterations << " with scaling factor: " << fScalingFactor << " and energy loss: " << calc_eLoss; @@ -207,7 +228,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur iterations++; PropagateToMeasurementSurface(surface, stepper); // Propagate without energy loss adjustment - double KE_final = Kinematics::KE(fMom, fMass); + double KE_final = Kinematics::KE(fState.fMom, fState.fMass); calc_eLoss = KE_initial - KE_final; // Energy loss in MeV fScalingFactor *= eLoss / calc_eLoss; LOG(info) << "Desired energy loss: " << eLoss << " MeV"; @@ -223,18 +244,22 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur fScalingFactor = 1; // Reset scaling factor after convergence } -AtStepper::StepState AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const +AtStepper::StepState AtRK4Stepper::Step(const StepState &state) const { // Take h to be the step size in m. - StepState result; - result.fLastPos = fPos; - result.fLastMom = fMom; - result.h = h; + StepState result = state; + result.fLastPos = state.fPos; + result.fLastMom = state.fMom; + result.hUsed = state.h; result.success = true; - LOG(debug) << "Starting RK4 step with initial position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(debug) << "Initial momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - LOG(debug) << "Step size (h): " << h << " m"; + auto h = state.h; // Step size in m + auto fPos = state.fPos; + auto fMom = state.fMom; + + LOG(info) << "Starting RK4 step with initial position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(info) << "Initial momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(info) << "Step size (h): " << h << " m"; auto [x_k1, p_k1] = fDeriv(fPos, fMom); // The derivative of the position is then just the unit vector of the momentum. @@ -255,8 +280,8 @@ AtStepper::StepState AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XY auto dpds_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // "Force" in SI units (N) auto dxds_SI = (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6; // Position derivative in SI units (m) - LOG(debug) << "dp/ds (SI units): " << dpds_SI.X() << ", " << dpds_SI.Y() << ", " << dpds_SI.Z(); - LOG(debug) << "dx/ds (SI units): " << dxds_SI.X() << ", " << dxds_SI.Y() << ", " << dxds_SI.Z(); + LOG(info) << "dp/ds (SI units): " << dpds_SI.X() << ", " << dpds_SI.Y() << ", " << dpds_SI.Z(); + LOG(info) << "dx/ds (SI units): " << dxds_SI.X() << ", " << dxds_SI.Y() << ", " << dxds_SI.Z(); auto mom_SI = fReltoSImom * fMom; mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) @@ -269,14 +294,16 @@ AtStepper::StepState AtRK4Stepper::Step(double h, const XYZPoint &fPos, const XY return result; } -AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, const XYZVector &fMom) const +AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const { // Take h to be the step size in m. - StepState result; - result.fLastPos = fPos; - result.fLastMom = fMom; - result.h = h; + StepState result = state; + result.fLastPos = state.fPos; + result.fLastMom = state.fMom; result.success = true; + auto h = state.h; // Step size in m + auto fPos = state.fPos; + auto fMom = state.fMom; // Take h to be the step size in m. // Use DP5(4) method for adaptive step size control. @@ -405,6 +432,7 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, // Adjust the step size for the next iteration result.h = hNew; + result.hUsed = h; // Store the step size used result.success = true; // Step accepted return result; } else { @@ -418,6 +446,7 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(double h, const XYZPoint &fPos, if (result.h < 1e-6) { LOG(error) << "Step size too small, aborting propagation."; result.success = false; + result.hUsed = h; return result; // Abort propagation if step size is too small } } diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 2c7dafb6d..c1f6e2f6a 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -20,6 +20,7 @@ class AtStepper { ROOT::Math::XYZPoint fLastPos; /// Last position of the particle in mm ROOT::Math::XYZVector fLastMom; /// Last momentum of the particle in MeV/c double fMass; /// Mass of the particle in MeV/c^2 + double fQ; /// Charge of the particle in Coulombs double h; /// Step size to use in m double hUsed; /// Step size used in this step in m bool success; /// Whether the step was successful @@ -40,7 +41,7 @@ class AtStepper { DerivFunc fDeriv; - virtual StepState Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const = 0; + virtual StepState Step(const StepState &state) const = 0; protected: static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) @@ -99,20 +100,22 @@ class AtPropagator { double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) double fScalingFactor = 1.0; /// Scaling factor for energy loss - AtStepper::StepState fState; // Current state of the particle - - XYZPoint fPos; // Current position of the particle in mm - XYZVector fMom; // Current momentum of the particle in MeV/c - - XYZPoint fLastPos; - XYZVector fLastMom; - + AtStepper::StepState fState; // Current state of the particle static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) public: + /** + * @brief Constructor for AtPropagator. + * @param charge Charge of the particle in Coulombs. + * @param mass Mass of the particle in MeV/c^2. + * @param elossModel Energy loss model to use for the particle. + */ AtPropagator(double charge, double mass, std::unique_ptr elossModel) : fQ(charge), fMass(mass), fELossModel(std::move(elossModel)) { + fState.fMass = mass; + fState.fQ = charge; + fState.h = fH; // Initialize step size } /** * @brief Set the electric field (V/m) @@ -134,14 +137,18 @@ class AtPropagator { */ void SetState(const XYZPoint &pos, const XYZVector &mom) { - fPos = pos; - fMom = mom; + fState.fPos = pos; + fState.fMom = mom; } - void SetH(double h) { fH = h; } + void SetH(double h) + { + fH = h; + fState.h = h; + } // Set the step size in m - XYZPoint GetPosition() const { return fPos; } - XYZVector GetMomentum() const { return fMom; } + XYZPoint GetPosition() const { return fState.fPos; } + XYZVector GetMomentum() const { return fState.fMom; } /** * @brief Propagate the particle to the point of closest approach to the given point. @@ -200,33 +207,22 @@ class AtPropagator { protected: void CopyFromState(const AtStepper::StepState &result) { - fPos = result.fPos; - fMom = result.fMom; - fLastPos = result.fLastPos; - fLastMom = result.fLastMom; + fState = result; + // fPos = result.fPos; + // fMom = result.fMom; + // fLastPos = result.fLastPos; + // fLastMom = result.fLastMom; fH = result.h; } - AtStepper::StepState CopyToState() const - { - AtStepper::StepState result; - result.fPos = fPos; - result.fMom = fMom; - result.fLastPos = fLastPos; - result.fLastMom = fLastMom; - result.fMass = fMass; - result.h = fH; - result.success = true; // Assume success unless proven otherwise - return result; - } }; class AtRK4Stepper : public AtStepper { public: - StepState Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; + StepState Step(const StepState &state) const override; }; class AtRK4AdaptiveStepper : public AtStepper { public: - StepState Step(double h, const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const override; + StepState Step(const StepState &state) const override; }; class AtMeasurementPoint : public AtMeasurementSurface { From 7332716f71cb7b62b27bfa91bb609209af5dc322 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 15:55:45 +0200 Subject: [PATCH 29/71] Remove double counted variables from state --- AtTools/AtPropagator.cxx | 10 ++++------ AtTools/AtPropagator.h | 11 +++-------- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 4339d0e01..99966deae 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -69,8 +69,8 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; while (true) { - LOG(info) << "Position: " << GetPosition().X() << ", " << GetPosition().Y() << ", " << GetPosition().Z(); - LOG(info) << "Momentum: " << GetMomentum().X() << ", " << GetMomentum().Y() << ", " << GetMomentum().Z(); + LOG(debug) << "Position: " << GetPosition().X() << ", " << GetPosition().Y() << ", " << GetPosition().Z(); + LOG(debug) << "Momentum: " << GetMomentum().X() << ", " << GetMomentum().Y() << ", " << GetMomentum().Z(); auto result = stepper.Step(fState); if (!result.success) { @@ -106,11 +106,10 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(error) << "Failed to propagate to measurement point, aborting."; return; // Abort propagation if step failed } - auto origH = fH; // Save original step size + auto origH = fState.h; // Save original step size CopyFromState(result); // Update position and momentum to the new state fState = result; // Update the internal state fState.h = origH; // Restore original step size - fH = origH; // Restore original step size } if (particleStopped || momentumReversed) { @@ -137,11 +136,10 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(error) << "Failed to propagate to stopping point, aborting."; return; // Abort propagation if step failed } - auto origH = fH; // Save original step size + auto origH = fState.h; // Save original step size CopyFromState(result); // Update position and momentum to the new state fState = result; // Update the internal state fState.h = origH; // Restore original step size - fH = origH; // Restore original step size LOG(info) << "Propagated to stopping point: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); LOG(info) << "Energy after stopping: " << Kinematics::KE(fState.fMom, fState.fMass) << " MeV"; diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index c1f6e2f6a..3c37cb39c 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -94,7 +94,6 @@ class AtPropagator { const std::unique_ptr fELossModel; // Energy loss model // Internal state variables for the propagator - double fH = 1e-4; /// Step size for propagation in m double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped (MeV) double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) @@ -115,7 +114,7 @@ class AtPropagator { { fState.fMass = mass; fState.fQ = charge; - fState.h = fH; // Initialize step size + fState.h = 1e-4; // Initial step size in m } /** * @brief Set the electric field (V/m) @@ -141,11 +140,7 @@ class AtPropagator { fState.fMom = mom; } - void SetH(double h) - { - fH = h; - fState.h = h; - } // Set the step size in m + void SetH(double h) { fState.h = h; } // Set the step size in m XYZPoint GetPosition() const { return fState.fPos; } XYZVector GetMomentum() const { return fState.fMom; } @@ -212,7 +207,7 @@ class AtPropagator { // fMom = result.fMom; // fLastPos = result.fLastPos; // fLastMom = result.fLastMom; - fH = result.h; + // fH = result.h; } }; From 03e6d77720207c6a10913ed3cd4430c65b7f252b Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 16:07:40 +0200 Subject: [PATCH 30/71] Remove debug and more refactoring --- AtTools/AtPropagator.cxx | 52 ++++++++++++++++-------------------- AtTools/AtPropagator.h | 57 +++++++++++++++++++--------------------- 2 files changed, 49 insertions(+), 60 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 99966deae..ef78b2177 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -33,12 +33,12 @@ namespace AtTools { AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const { - auto v = Kinematics::GetVel(mom, fMass); + auto v = Kinematics::GetVel(mom, fState.fMass); - auto F_lorentz = fQ * (fEField + v.Cross(fBField)); + auto F_lorentz = fState.fQ * (fEField + v.Cross(fBField)); LOG(debug) << "F_lorentz: " << F_lorentz; - auto dedx = fScalingFactor * fELossModel->GetdEdx(Kinematics::KE(mom, fMass)); // Stopping power in MeV/mm - auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) + auto dedx = fScalingFactor * fELossModel->GetdEdx(Kinematics::KE(mom, fState.fMass)); // Stopping power in MeV/mm + auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) auto drag = -dedx_si * mom.Unit(); LOG(debug) << "drag: " << drag << " mom " << mom << " dedx " << dedx_si; @@ -49,7 +49,7 @@ AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const AtPropagator::XYZVector AtPropagator::dpds(const XYZPoint &pos, const XYZVector &mom) const { // Calculate the force acting on the particle at the given position and momentum - auto speed = Kinematics::GetSpeed(mom.R(), fMass); // Speed in m/s + auto speed = Kinematics::GetSpeed(mom.R(), fState.fMass); // Speed in m/s return Force(pos, mom) / speed; } AtPropagator::XYZVector AtPropagator::d2xds2(const XYZPoint &pos, const XYZVector &mom) const @@ -77,8 +77,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(error) << "Integration step failed, aborting propagation."; return; // Abort propagation if step failed } - CopyFromState(result); // Copy the new state from the stepper - fState = result; // Update the internal state + fState = result; // Update the internal state bool reachedMeasurementPoint = surface.PassedSurface(fState); bool particleStopped = Kinematics::KE(fState.fMom, fState.fMass) < fStopTol; @@ -107,7 +106,6 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur return; // Abort propagation if step failed } auto origH = fState.h; // Save original step size - CopyFromState(result); // Update position and momentum to the new state fState = result; // Update the internal state fState.h = origH; // Restore original step size } @@ -137,7 +135,6 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur return; // Abort propagation if step failed } auto origH = fState.h; // Save original step size - CopyFromState(result); // Update position and momentum to the new state fState = result; // Update the internal state fState.h = origH; // Restore original step size LOG(info) << "Propagated to stopping point: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " @@ -160,8 +157,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(error) << "Failed to propagate to surface after stopping, aborting."; return; // Abort propagation if step failed } - CopyFromState(result); // Update position and momentum to the new state - fState = result; // Update the internal state + fState = result; // Update the internal state LOG(info) << "New position after adjusting step size: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); } @@ -299,6 +295,7 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const result.fLastPos = state.fPos; result.fLastMom = state.fMom; result.success = true; + auto h = state.h; // Step size in m auto fPos = state.fPos; auto fMom = state.fMom; @@ -306,14 +303,10 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const // Take h to be the step size in m. // Use DP5(4) method for adaptive step size control. - double atol_pos = 1e-2; // Absolute tolerance for position (mm) - double atol_mom = 1e-2; // Absolute tolerance for momentum (MeV/c) - double rtol = 1e-6; // Relative tolerance for both position and momentum - auto x0_mm = fPos; auto p0 = fMom; - LOG(info) << "Starting RK4 step with initial position: " << x0_mm.X() << ", " << x0_mm.Y() << ", " << x0_mm.Z(); - LOG(info) << "Initial momentum: " << p0.X() << ", " << p0.Y() << ", " << p0.Z(); + LOG(debug) << "Starting RK4 step with initial position: " << x0_mm.X() << ", " << x0_mm.Y() << ", " << x0_mm.Z(); + LOG(debug) << "Initial momentum: " << p0.X() << ", " << p0.Y() << ", " << p0.Z(); while (true) { auto x_SI = fPos * 1e-3; // Convert position to SI units (m) @@ -403,13 +396,13 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const XYZVector p_err = (p_5_MeV - p_4_MeV); // Error in momentum (MeV/c) // Calculate the overall error - double ex = x_err.X() / (atol_pos + rtol * std::abs(x_5_mm.X())); - double ey = x_err.Y() / (atol_pos + rtol * std::abs(x_5_mm.Y())); - double ez = x_err.Z() / (atol_pos + rtol * std::abs(x_5_mm.Z())); + double ex = x_err.X() / (fAtolPos + fRtol * std::abs(x_5_mm.X())); + double ey = x_err.Y() / (fAtolPos + fRtol * std::abs(x_5_mm.Y())); + double ez = x_err.Z() / (fAtolPos + fRtol * std::abs(x_5_mm.Z())); - double ep_x = p_err.X() / (atol_mom + rtol * std::abs(p_5_MeV.X())); - double ep_y = p_err.Y() / (atol_mom + rtol * std::abs(p_5_MeV.Y())); - double ep_z = p_err.Z() / (atol_mom + rtol * std::abs(p_5_MeV.Z())); + double ep_x = p_err.X() / (fAtolMom + fRtol * std::abs(p_5_MeV.X())); + double ep_y = p_err.Y() / (fAtolMom + fRtol * std::abs(p_5_MeV.Y())); + double ep_z = p_err.Z() / (fAtolMom + fRtol * std::abs(p_5_MeV.Z())); // Combine errors (norm) double err = std::sqrt(ex * ex + ey * ey + ez * ez + ep_x * ep_x + ep_y * ep_y + ep_z * ep_z); @@ -428,8 +421,7 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const LOG(info) << "New Position: " << result.fPos.X() << ", " << result.fPos.Y() << ", " << result.fPos.Z(); LOG(info) << "New Momentum: " << result.fMom.X() << ", " << result.fMom.Y() << ", " << result.fMom.Z(); - // Adjust the step size for the next iteration - result.h = hNew; + result.h = hNew; // Adjust the step size for the next iteration result.hUsed = h; // Store the step size used result.success = true; // Step accepted return result; @@ -439,13 +431,13 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const LOG(info) << "Step size: " << h << " m"; LOG(info) << "Reducing step size to: " << hNew << " m"; - result.h = hNew; // Reduce step size - h = hNew; // Update the step size for the next iteration - if (result.h < 1e-6) { - LOG(error) << "Step size too small, aborting propagation."; + result.h = hNew; // Reduce step size for next iteration + h = hNew; // Update h for the next iteration + if (result.h < fMinStep || result.h > fMaxStep) { + LOG(error) << "Step size out of bounds, aborting propagation."; result.success = false; result.hUsed = h; - return result; // Abort propagation if step size is too small + return result; // Abort propagation if step size is out of bounds } } } diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 3c37cb39c..e1c357dfa 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -86,20 +86,21 @@ class AtPropagator { using XYZVector = ROOT::Math::XYZVector; using XYZPoint = ROOT::Math::XYZPoint; using Plane3D = ROOT::Math::Plane3D; - XYZVector fEField{0, 0, 0}; // Electric field vector - XYZVector fBField{0, 0, 0}; // Magnetic field vector - const double fQ; // Charge of the particle in Coulombs - const double fMass; // Mass of the particle in MeV/c^2 + // Variables used for the force + XYZVector fEField{0, 0, 0}; // Electric field vector + XYZVector fBField{0, 0, 0}; // Magnetic field vector const std::unique_ptr fELossModel; // Energy loss model // Internal state variables for the propagator - double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss - double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped (MeV) - double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) double fScalingFactor = 1.0; /// Scaling factor for energy loss + AtStepper::StepState fState; /// Current state of the particle + + // Tolerances and limits + double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss + double fStopTol = 0.01; /// Maximum kinetic energy to consider the particle stopped (MeV) + double fDistTol = 1e-2; /// Distance tolerance when considering positions equal. (mm) - AtStepper::StepState fState; // Current state of the particle static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) public: @@ -110,7 +111,7 @@ class AtPropagator { * @param elossModel Energy loss model to use for the particle. */ AtPropagator(double charge, double mass, std::unique_ptr elossModel) - : fQ(charge), fMass(mass), fELossModel(std::move(elossModel)) + : fELossModel(std::move(elossModel)) { fState.fMass = mass; fState.fQ = charge; @@ -176,18 +177,6 @@ class AtPropagator { */ XYZVector dpds(const XYZPoint &pos, const XYZVector &mom) const; - /** - * @brief Calculate the second derivative of the position w.r.t. arc length. - * - * \frac{d^2\vec{x}}{ds^2} = \frac{1}{p} \left( \frac{d\vec{p}}{ds} - \hat{p} (\hat{p} \cdot \frac{d\vec{p}}{ds}) - * \right) - * - * @param pos Position of the particle in mm. - * @param mom Momentum of the particle in MeV/c. - * @return The second derivative of the position w.r.t. arc length in m/m^2. - */ - XYZVector d2xds2(const XYZPoint &pos, const XYZVector &mom) const; - XYZVector dxds(const XYZPoint &pos, const XYZVector &mom) const { return mom.Unit(); // The derivative of the position is just the unit vector of the momentum. @@ -200,15 +189,17 @@ class AtPropagator { } protected: - void CopyFromState(const AtStepper::StepState &result) - { - fState = result; - // fPos = result.fPos; - // fMom = result.fMom; - // fLastPos = result.fLastPos; - // fLastMom = result.fLastMom; - // fH = result.h; - } + /** + * @brief Calculate the second derivative of the position w.r.t. arc length. + * + * \frac{d^2\vec{x}}{ds^2} = \frac{1}{p} \left( \frac{d\vec{p}}{ds} - \hat{p} (\hat{p} \cdot \frac{d\vec{p}}{ds}) + * \right) + * + * @param pos Position of the particle in mm. + * @param mom Momentum of the particle in MeV/c. + * @return The second derivative of the position w.r.t. arc length in m/m^2. + */ + XYZVector d2xds2(const XYZPoint &pos, const XYZVector &mom) const; }; class AtRK4Stepper : public AtStepper { @@ -217,6 +208,12 @@ class AtRK4Stepper : public AtStepper { }; class AtRK4AdaptiveStepper : public AtStepper { public: + double fAtolPos = 1e-2; /// Absolute tolerance for position in mm + double fAtolMom = 1e-2; /// Absolute tolerance for momentum in MeV/c + double fRtol = 1e-6; /// Relative tolerance for position and momentum + double fMinStep = 1e-6; /// Minimum step size in m + double fMaxStep = 10.0; /// Maximum step size in m + StepState Step(const StepState &state) const override; }; From 4801616b620f34fd01403ed638e5ac4752bdc8f9 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 18:36:02 +0200 Subject: [PATCH 31/71] Refactor step size defaults --- AtTools/AtPropagator.cxx | 80 +++++++++--------- AtTools/AtPropagator.h | 154 +++++++++++++++++++---------------- AtTools/AtPropagatorTest.cxx | 6 +- 3 files changed, 124 insertions(+), 116 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index ef78b2177..fd526d77f 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -4,12 +4,10 @@ #include -// Butcher tableau (c, a_ij) and the two b vectors for Dormand–Prince 5(4) -// c1 = 0 // Butcher tableau coefficients for Dormand–Prince 5(4) method +// https://en.wikipedia.org/wiki/Dormand%E2%80%93Prince_method static constexpr double c[7] = {0.0, 1.0 / 5.0, 3.0 / 10.0, 4.0 / 5.0, 8.0 / 9.0, 1.0, 1.0}; - static constexpr double a[7][6] = { {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, {1.0 / 5.0, 0.0, 0.0, 0.0, 0.0, 0.0}, @@ -18,11 +16,9 @@ static constexpr double a[7][6] = { {19372.0 / 6561.0, -25360.0 / 2187.0, 64448.0 / 6561.0, -212.0 / 729.0, 0.0, 0.0}, {9017.0 / 3168.0, -355.0 / 33.0, 46732.0 / 5247.0, 49.0 / 176.0, -5103.0 / 18656.0, 0.0}, {35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0, 11.0 / 84.0}}; - // b (5th-order) static constexpr double b[7] = {35.0 / 384.0, 0.0, 500.0 / 1113.0, 125.0 / 192.0, -2187.0 / 6784.0, 11.0 / 84.0, 0.0}; - -// b* (4th-order, “star”) +// b* (4th-order) static constexpr double bs[7] = {5179.0 / 57600.0, 0.0, 7571.0 / 16695.0, 393.0 / 640.0, -92097.0 / 339200.0, 187.0 / 2100.0, 1.0 / 40.0}; @@ -64,6 +60,7 @@ AtPropagator::XYZVector AtPropagator::d2xds2(const XYZPoint &pos, const XYZVecto void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) { LOG(info) << "Propagating to measurement surface"; + fState.h = stepper.GetInitialStep(); // Set the initial step size auto KE_initial = Kinematics::KE(fState.fMom, fState.fMass); stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; @@ -73,7 +70,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(debug) << "Momentum: " << GetMomentum().X() << ", " << GetMomentum().Y() << ", " << GetMomentum().Z(); auto result = stepper.Step(fState); - if (!result.success) { + if (!result) { LOG(error) << "Integration step failed, aborting propagation."; return; // Abort propagation if step failed } @@ -101,7 +98,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur fState.fPos = fState.fLastPos; // Set position to last position fState.fMom = fState.fLastMom; // Set momentum to last momentum result = stepper.Step(fState); - if (!result.success) { + if (!result) { LOG(error) << "Failed to propagate to measurement point, aborting."; return; // Abort propagation if step failed } @@ -130,7 +127,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur fState.fPos = fState.fLastPos; // Set position to last position fState.fMom = fState.fLastMom; // Set momentum to last momentum result = stepper.Step(fState); - if (!result.success) { + if (!result) { LOG(error) << "Failed to propagate to stopping point, aborting."; return; // Abort propagation if step failed } @@ -153,7 +150,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(info) << "Propagating to surface after stopping with step size: " << h << " mm"; fState.h = h * 1e-3; // Convert to meters for the RK4 step result = stepper.Step(fState); - if (!result.success) { + if (!result) { LOG(error) << "Failed to propagate to surface after stopping, aborting."; return; // Abort propagation if step failed } @@ -238,22 +235,22 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur fScalingFactor = 1; // Reset scaling factor after convergence } -AtStepper::StepState AtRK4Stepper::Step(const StepState &state) const +AtPropagator::StepState AtRK4Stepper::Step(const AtPropagator::StepState &state) const { - // Take h to be the step size in m. - StepState result = state; + + auto result = state; result.fLastPos = state.fPos; result.fLastMom = state.fMom; result.hUsed = state.h; - result.success = true; + result.status = AtPropagator::StepStateStatus::kSuccess; auto h = state.h; // Step size in m auto fPos = state.fPos; auto fMom = state.fMom; - LOG(info) << "Starting RK4 step with initial position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); - LOG(info) << "Initial momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); - LOG(info) << "Step size (h): " << h << " m"; + LOG(debug) << "Starting RK4 step with initial position: " << fPos.X() << ", " << fPos.Y() << ", " << fPos.Z(); + LOG(debug) << "Initial momentum: " << fMom.X() << ", " << fMom.Y() << ", " << fMom.Z(); + LOG(debug) << "Step size (h): " << h << " m"; auto [x_k1, p_k1] = fDeriv(fPos, fMom); // The derivative of the position is then just the unit vector of the momentum. @@ -274,8 +271,8 @@ AtStepper::StepState AtRK4Stepper::Step(const StepState &state) const auto dpds_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; // "Force" in SI units (N) auto dxds_SI = (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6; // Position derivative in SI units (m) - LOG(info) << "dp/ds (SI units): " << dpds_SI.X() << ", " << dpds_SI.Y() << ", " << dpds_SI.Z(); - LOG(info) << "dx/ds (SI units): " << dxds_SI.X() << ", " << dxds_SI.Y() << ", " << dxds_SI.Z(); + LOG(debug) << "dp/ds (SI units): " << dpds_SI.X() << ", " << dpds_SI.Y() << ", " << dpds_SI.Z(); + LOG(debug) << "dx/ds (SI units): " << dxds_SI.X() << ", " << dxds_SI.Y() << ", " << dxds_SI.Z(); auto mom_SI = fReltoSImom * fMom; mom_SI += dpds_SI * h; // Update momentum in SI units (kg m/s) @@ -288,13 +285,14 @@ AtStepper::StepState AtRK4Stepper::Step(const StepState &state) const return result; } -AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const +AtPropagator::StepState AtRK4AdaptiveStepper::Step(const AtPropagator::StepState &state) const { + // Take h to be the step size in m. - StepState result = state; + auto result = state; result.fLastPos = state.fPos; result.fLastMom = state.fMom; - result.success = true; + result.status = AtPropagator::StepStateStatus::kSuccess; auto h = state.h; // Step size in m auto fPos = state.fPos; @@ -386,10 +384,10 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const auto p_4_MeV = p_new_4 / fReltoSImom; // Convert back to MeV/c auto x_5_mm = x_new_5 * 1e3; // Convert back to mm auto p_5_MeV = p_new_5 / fReltoSImom; // Convert back to MeV/c - LOG(info) << "New position (5th order): " << x_5_mm.X() << ", " << x_5_mm.Y() << ", " << x_5_mm.Z(); - LOG(info) << "New momentum (5th order): " << p_5_MeV.X() << ", " << p_5_MeV.Y() << ", " << p_5_MeV.Z(); - LOG(info) << "New position (4th order): " << x_4_mm.X() << ", " << x_4_mm.Y() << ", " << x_4_mm.Z(); - LOG(info) << "New momentum (4th order): " << p_4_MeV.X() << ", " << p_4_MeV.Y() << ", " << p_4_MeV.Z(); + LOG(debug) << "New position (5th order): " << x_5_mm.X() << ", " << x_5_mm.Y() << ", " << x_5_mm.Z(); + LOG(debug) << "New momentum (5th order): " << p_5_MeV.X() << ", " << p_5_MeV.Y() << ", " << p_5_MeV.Z(); + LOG(debug) << "New position (4th order): " << x_4_mm.X() << ", " << x_4_mm.Y() << ", " << x_4_mm.Z(); + LOG(debug) << "New momentum (4th order): " << p_4_MeV.X() << ", " << p_4_MeV.Y() << ", " << p_4_MeV.Z(); // Convert back to mm and MeV/c XYZVector x_err = (x_5_mm - x_4_mm); // Error in position (mm) @@ -415,27 +413,27 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const // Accept the step result.fPos = x_5_mm; // Update position in mm result.fMom = p_5_MeV; // Update momentum in MeV/c - LOG(info) << "Accepted step with error: " << err; - LOG(info) << "Step size: " << h << " m"; - LOG(info) << "New step size: " << hNew << " m"; - LOG(info) << "New Position: " << result.fPos.X() << ", " << result.fPos.Y() << ", " << result.fPos.Z(); - LOG(info) << "New Momentum: " << result.fMom.X() << ", " << result.fMom.Y() << ", " << result.fMom.Z(); - - result.h = hNew; // Adjust the step size for the next iteration - result.hUsed = h; // Store the step size used - result.success = true; // Step accepted + LOG(debug) << "Accepted step with error: " << err; + LOG(debug) << "Step size: " << h << " m"; + LOG(debug) << "New step size: " << hNew << " m"; + LOG(debug) << "New Position: " << result.fPos.X() << ", " << result.fPos.Y() << ", " << result.fPos.Z(); + LOG(debug) << "New Momentum: " << result.fMom.X() << ", " << result.fMom.Y() << ", " << result.fMom.Z(); + + result.h = hNew; // Adjust the step size for the next iteration + result.hUsed = h; // Store the step size used + result.status = AtPropagator::StepStateStatus::kSuccess; // Step accepted return result; } else { // Reject the step and reduce the step size - LOG(info) << "Rejected step with error: " << err; - LOG(info) << "Step size: " << h << " m"; - LOG(info) << "Reducing step size to: " << hNew << " m"; + LOG(debug) << "Rejected step with error: " << err; + LOG(debug) << "Step size: " << h << " m"; + LOG(debug) << "Reducing step size to: " << hNew << " m"; result.h = hNew; // Reduce step size for next iteration h = hNew; // Update h for the next iteration if (result.h < fMinStep || result.h > fMaxStep) { LOG(error) << "Step size out of bounds, aborting propagation."; - result.success = false; + result.status = AtPropagator::StepStateStatus::kInvalidStepSize; result.hUsed = h; return result; // Abort propagation if step size is out of bounds } @@ -443,7 +441,7 @@ AtStepper::StepState AtRK4AdaptiveStepper::Step(const StepState &state) const } } -bool AtMeasurementPoint::PassedSurface(AtStepper::StepState &result) const +bool AtMeasurementPoint::PassedSurface(AtPropagator::StepState &result) const { // Check if the particle has passed the measurement point auto lastDeriv = (fPoint - result.fLastPos).Dot(result.fLastMom.Unit()); @@ -452,7 +450,7 @@ bool AtMeasurementPoint::PassedSurface(AtStepper::StepState &result) const return lastDeriv * currDeriv <= 0; } -bool AtMeasurementPlane::PassedSurface(AtStepper::StepState &result) const +bool AtMeasurementPlane::PassedSurface(AtPropagator::StepState &result) const { // Check if the particle has crossed the plane this step. auto prevSign = fPlane.Distance(result.fLastPos) > 0 ? 1 : -1; diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index e1c357dfa..f5e4b13b9 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -12,63 +12,7 @@ namespace AtTools { class AtMeasurementSurface; -class AtStepper { -public: - struct StepState { - ROOT::Math::XYZPoint fPos; /// Position of the particle in mm - ROOT::Math::XYZVector fMom; /// Momentum of the particle in MeV/c - ROOT::Math::XYZPoint fLastPos; /// Last position of the particle in mm - ROOT::Math::XYZVector fLastMom; /// Last momentum of the particle in MeV/c - double fMass; /// Mass of the particle in MeV/c^2 - double fQ; /// Charge of the particle in Coulombs - double h; /// Step size to use in m - double hUsed; /// Step size used in this step in m - bool success; /// Whether the step was successful - }; - /** - * @brief Function type defining the derivative of the position and momentum w.r.t. distance. - * - * This function takes the current position and momentum and returns the derivate of the position and momentum. - * - * @param pos Current position of the particle in mm. - * @param mom Current momentum of the particle in MeV/c. - * @return A pair containing the derivatives of the position and momentum in SI units (m and kg m/s). - * The first element is the derivative of the position, and the second element is the derivative - * of the momentum. - */ - using DerivFunc = std::function( - const ROOT::Math::XYZPoint &, const ROOT::Math::XYZVector &)>; - - DerivFunc fDeriv; - - virtual StepState Step(const StepState &state) const = 0; - -protected: - static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) -}; - -/** - * @brief Class for measurement surface in the AT-TPC. - * - * This class represents a measurement surface or point in the AT-TPC. It's used to define the stopping - * point and behavior of the propagator. - */ -class AtMeasurementSurface { -public: - bool fClipToSurface = false; // Whether to clip to the surface - - /** - * @brief Calculate the distance from the position to the surface. - */ - virtual double Distance(const ROOT::Math::XYZPoint &pos) const = 0; - - /** - * @brief Check if we have passed the surface between the last position and the current position. - */ - virtual bool PassedSurface(AtStepper::StepState &result) const = 0; - - virtual ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const = 0; -}; +class AtStepper; /** * @brief Class for propagating particles through a medium. @@ -82,6 +26,25 @@ class AtMeasurementSurface { * class if the material or particle type changes. */ class AtPropagator { +public: + enum class StepStateStatus { + kSuccess, /// Step was successful + kInvalidStepSize /// Step failed + }; + struct StepState { + ROOT::Math::XYZPoint fPos; /// Position of the particle in mm + ROOT::Math::XYZVector fMom; /// Momentum of the particle in MeV/c + ROOT::Math::XYZPoint fLastPos; /// Last position of the particle in mm + ROOT::Math::XYZVector fLastMom; /// Last momentum of the particle in MeV/c + double fMass = 0; /// Mass of the particle in MeV/c^2 + double fQ = 0; /// Charge of the particle in Coulombs + double h = 0; /// Step size to use in m + double hUsed = 0; /// Step size used in this step in m + StepStateStatus status = StepStateStatus::kSuccess; /// Whether the step was successful + + operator bool() const { return status == StepStateStatus::kSuccess; } + }; + protected: using XYZVector = ROOT::Math::XYZVector; using XYZPoint = ROOT::Math::XYZPoint; @@ -94,7 +57,7 @@ class AtPropagator { // Internal state variables for the propagator double fScalingFactor = 1.0; /// Scaling factor for energy loss - AtStepper::StepState fState; /// Current state of the particle + StepState fState; /// Current state of the particle // Tolerances and limits double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss @@ -115,8 +78,8 @@ class AtPropagator { { fState.fMass = mass; fState.fQ = charge; - fState.h = 1e-4; // Initial step size in m } + /** * @brief Set the electric field (V/m) */ @@ -141,8 +104,6 @@ class AtPropagator { fState.fMom = mom; } - void SetH(double h) { fState.h = h; } // Set the step size in m - XYZPoint GetPosition() const { return fState.fPos; } XYZVector GetMomentum() const { return fState.fMom; } @@ -201,22 +162,73 @@ class AtPropagator { */ XYZVector d2xds2(const XYZPoint &pos, const XYZVector &mom) const; }; +class AtStepper { +public: + /** + * @brief Function type defining the derivative of the position and momentum w.r.t. distance. + * + * This function takes the current position and momentum and returns the derivate of the position and momentum. + * + * @param pos Current position of the particle in mm. + * @param mom Current momentum of the particle in MeV/c. + * @return A pair containing the derivatives of the position and momentum in SI units (m and kg m/s). + * The first element is the derivative of the position, and the second element is the derivative + * of the momentum. + */ + using DerivFunc = std::function( + const ROOT::Math::XYZPoint &, const ROOT::Math::XYZVector &)>; + + DerivFunc fDeriv; + + virtual AtPropagator::StepState Step(const AtPropagator::StepState &state) const = 0; + virtual double GetInitialStep() const { return 1e-4; } /// Default initial step size in m + +protected: + static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) +}; class AtRK4Stepper : public AtStepper { + double fStepSize = 1e-4; + public: - StepState Step(const StepState &state) const override; + AtPropagator::StepState Step(const AtPropagator::StepState &state) const override; + double GetInitialStep() const override { return fStepSize; } /// Default initial step size in m }; class AtRK4AdaptiveStepper : public AtStepper { public: - double fAtolPos = 1e-2; /// Absolute tolerance for position in mm - double fAtolMom = 1e-2; /// Absolute tolerance for momentum in MeV/c - double fRtol = 1e-6; /// Relative tolerance for position and momentum - double fMinStep = 1e-6; /// Minimum step size in m - double fMaxStep = 10.0; /// Maximum step size in m - - StepState Step(const StepState &state) const override; + double fAtolPos = 1e-2; /// Absolute tolerance for position in mm + double fAtolMom = 1e-2; /// Absolute tolerance for momentum in MeV/c + double fRtol = 1e-6; /// Relative tolerance for position and momentum + double fMinStep = 1e-6; /// Minimum step size in m + double fMaxStep = 10.0; /// Maximum step size in m + double fInitialStep = 1e-4; /// Initial step size in m + + AtPropagator::StepState Step(const AtPropagator::StepState &state) const override; + double GetInitialStep() const override { return fInitialStep; } /// Default initial step size in m }; +/** + * @brief Class for measurement surface in the AT-TPC. + * + * This class represents a measurement surface or point in the AT-TPC. It's used to define the stopping + * point and behavior of the propagator. + */ +class AtMeasurementSurface { +public: + bool fClipToSurface = false; // Whether to clip to the surface + + /** + * @brief Calculate the distance from the position to the surface. + */ + virtual double Distance(const ROOT::Math::XYZPoint &pos) const = 0; + + /** + * @brief Check if we have passed the surface between the last position and the current position. + */ + virtual bool PassedSurface(AtPropagator::StepState &result) const = 0; + + virtual ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const = 0; +}; class AtMeasurementPoint : public AtMeasurementSurface { protected: ROOT::Math::XYZPoint fPoint; // The measurement point in mm @@ -225,7 +237,7 @@ class AtMeasurementPoint : public AtMeasurementSurface { AtMeasurementPoint(const ROOT::Math::XYZPoint &point) : fPoint(point) {} double Distance(const ROOT::Math::XYZPoint &pos) const override { return (fPoint - pos).R(); } - bool PassedSurface(AtStepper::StepState &result) const override; + bool PassedSurface(AtPropagator::StepState &result) const override; ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override { return fPoint; } }; @@ -236,7 +248,7 @@ class AtMeasurementPlane : public AtMeasurementSurface { AtMeasurementPlane(const ROOT::Math::Plane3D &plane) : fPlane(plane) { fClipToSurface = true; } double Distance(const ROOT::Math::XYZPoint &pos) const override { return std::abs(fPlane.Distance(pos)); } - bool PassedSurface(AtStepper::StepState &result) const override; + bool PassedSurface(AtPropagator::StepState &result) const override; ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override; }; } // namespace AtTools diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 75467531b..608cb220a 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -169,8 +169,6 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); - XYZPoint targetPoint(10, 0, 0); // Target point to propagate to 10 mm - // propagator.PropagateToPoint(targetPoint, stepper); propagator.PropagateToMeasurementSurface(measurementPoint, stepper); auto finalPos = propagator.GetPosition(); @@ -278,7 +276,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) propagator.SetState(startPos, startMom); propagator.SetEField({0, 0, 0}); // No electric field propagator.SetBField({0, 0, 0}); // No magnetic field - propagator.SetH(1); // Set initial step size to 1 s + stepper.fInitialStep = 1; // Set initial step size to 1 m ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); @@ -293,7 +291,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) propagator.SetState(startPos, startMom); propagator.SetEField({0, 0, 0}); // No electric field propagator.SetBField({0, 0, 0}); // No magnetic field - propagator.SetH(1e-6); // Set initial step size to 1e-6 m + stepper.fInitialStep = 1e-6; // Set initial step size to 1e-6 m ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); From 9d06ade3f2e4940053753e3a644cf51d7a5ae857 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 22 Jul 2025 20:36:14 +0200 Subject: [PATCH 32/71] Fix comments that implied AtEnergyLoss density was mg/cm^3 --- AtTools/AtELossModel.cxx | 4 +++- AtTools/AtELossTable.cxx | 9 +++++++-- AtTools/AtPropagator.cxx | 4 +++- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/AtTools/AtELossModel.cxx b/AtTools/AtELossModel.cxx index 7ad30bdca..bda0f6966 100644 --- a/AtTools/AtELossModel.cxx +++ b/AtTools/AtELossModel.cxx @@ -1,10 +1,12 @@ #include "AtELossModel.h" +#include "FairLogger.h" + #include namespace AtTools { /** - * Set the density of the material we are calculating energy losses for in mg/cm^3. + * Set the density of the material we are calculating energy losses for in g/cm^3. * Likely not fully tested, but I want to keep it around to remind myself of it. */ void AtELossModel::SetDensity(double density) diff --git a/AtTools/AtELossTable.cxx b/AtTools/AtELossTable.cxx index c31ba26ed..8f093a118 100644 --- a/AtTools/AtELossTable.cxx +++ b/AtTools/AtELossTable.cxx @@ -198,7 +198,7 @@ void AtELossTable::LoadSrimTable(std::string fileName) try { if (atConversion && tokens.at(1) == "MeV" && tokens.at(3) == "mm") { conversion = std::stod(tokens.at(0)); - LOG(info) << "Using conversion factor of " << conversion; + LOG(info) << "Using conversion factor of " << conversion << " from " << tokens.at(0); break; } @@ -221,8 +221,13 @@ void AtELossTable::LoadSrimTable(std::string fileName) } } // end loop over file - for (auto &dedx : dEdX) + for (auto &dedx : dEdX) { dedx *= conversion; + } + + for (int i = 0; i < dEdX.size(); ++i) { + LOG(debug) << "Energy: " << energy[i] << " MeV dEdX: " << dEdX[i] << " MeV/mm"; + } LoadTable(energy, dEdX); LoadRangeVariance(energy, rangeVar); diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index fd526d77f..ee1a0cb6c 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -33,6 +33,7 @@ AtPropagator::XYZVector AtPropagator::Force(XYZPoint pos, XYZVector mom) const auto F_lorentz = fState.fQ * (fEField + v.Cross(fBField)); LOG(debug) << "F_lorentz: " << F_lorentz; + auto dedx = fScalingFactor * fELossModel->GetdEdx(Kinematics::KE(mom, fState.fMass)); // Stopping power in MeV/mm auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) @@ -66,7 +67,8 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; while (true) { - LOG(debug) << "Position: " << GetPosition().X() << ", " << GetPosition().Y() << ", " << GetPosition().Z(); + LOG(info) << "Position: " << GetPosition().X() / 10 << ", " << GetPosition().Y() / 10 << ", " + << GetPosition().Z() / 10; LOG(debug) << "Momentum: " << GetMomentum().X() << ", " << GetMomentum().Y() << ", " << GetMomentum().Z(); auto result = stepper.Step(fState); From caffb8b6fe7bae471ecdaaeff07c345c6d0ea1df Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 23 Jul 2025 11:41:36 +0200 Subject: [PATCH 33/71] Possibly make tests portable. --- AtTools/AtPropagator.h | 7 +-- AtTools/AtPropagatorTest.cxx | 84 +++++++++++++++++++++++++++++++----- 2 files changed, 77 insertions(+), 14 deletions(-) diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index f5e4b13b9..0fd07435c 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -56,8 +56,7 @@ class AtPropagator { const std::unique_ptr fELossModel; // Energy loss model // Internal state variables for the propagator - double fScalingFactor = 1.0; /// Scaling factor for energy loss - StepState fState; /// Current state of the particle + StepState fState; /// Current state of the particle // Tolerances and limits double fETol = 1e-4; /// Energy tolerance for convergence when fixing energy loss @@ -67,6 +66,8 @@ class AtPropagator { static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) public: + double fScalingFactor = 1.0; /// Scaling factor for energy loss + /** * @brief Constructor for AtPropagator. * @param charge Charge of the particle in Coulombs. @@ -238,7 +239,7 @@ class AtMeasurementPoint : public AtMeasurementSurface { double Distance(const ROOT::Math::XYZPoint &pos) const override { return (fPoint - pos).R(); } bool PassedSurface(AtPropagator::StepState &result) const override; - ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override { return fPoint; } + ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override { return pos; } }; class AtMeasurementPlane : public AtMeasurementSurface { diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 608cb220a..81cd1ad5b 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -103,8 +103,9 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 auto elossModel = std::make_unique(0); - elossModel->LoadSrimTable( - "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + // elossModel->LoadSrimTable( + // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; AtMeasurementPoint measurementPoint({1e3, 0, 0}); @@ -147,8 +148,9 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 auto elossModel = std::make_unique(0); - elossModel->LoadSrimTable( - "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + // elossModel->LoadSrimTable( + // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; AtMeasurementPoint measurementPoint({10, 0, 0}); @@ -183,8 +185,9 @@ TEST(AtPropagatorTest, PropagateToPlane_NoField) double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 auto elossModel = std::make_unique(0); - elossModel->LoadSrimTable( - "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + // elossModel->LoadSrimTable( + // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; @@ -222,8 +225,9 @@ TEST(AtPropagatorTest, PropagateToPlane_StoppingNoField) double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 auto elossModel = std::make_unique(0); - elossModel->LoadSrimTable( - "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + // elossModel->LoadSrimTable( + // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; @@ -257,8 +261,9 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) double charge = charge_p; // Charge in Coulombs double mass = mass_p; // Mass in MeV/c^2 auto elossModel = std::make_unique(0); - elossModel->LoadSrimTable( - "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + // elossModel->LoadSrimTable( + // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4AdaptiveStepper stepper; AtMeasurementPoint measurementPoint({10, 0, 0}); @@ -302,4 +307,61 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) ASSERT_NEAR(finalPos.X(), 10, 10 * 1e-3); // Final position in x-direction should be close to 10 mm ASSERT_NEAR(finalMom.X(), p_fin, 0.1); -} \ No newline at end of file +} + +TEST(AtPropagatorTest, PropagateToPoint_Field) +{ + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(0); + // elossModel->LoadSrimTable( + // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); + elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools + elossModel->SetDensity(3.3084e-05); // Set density in g/cm^3 for 300 torr H2 + AtPropagator propagator(charge, mass, std::move(elossModel)); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 2.85}); // Magnetic field + AtRK4Stepper stepper; + + XYZPoint startPos(-3.40046e-05, -1.49863e-05, 0.10018); // Start position in cm + startPos *= 10; // Convert to mm + XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c + startMom *= 1e3; // Convert to MeV/c + + auto KE = Kinematics::KE(startMom, mass); // Convert momentum to kinetic energy + std::cout << "Propagating proton with KE: " << KE << " MeV" << std::endl; + std::cout << "Initial position: " << startPos.X() << ", " << startPos.Y() << ", " << startPos.Z() << std::endl; + + propagator.SetState(startPos, startMom); + + XYZPoint point({-1.4895, -4.8787, 1.01217}); // measurement point in cm + point *= 10; // Convert to mm + AtMeasurementPoint measurementPoint(point); + + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); + + auto finalPos = propagator.GetPosition(); + auto finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), point.X(), 1); // Check final position is within 1 mm of the measurement point + ASSERT_NEAR(finalPos.Y(), point.Y(), 1); + ASSERT_NEAR(finalPos.Z(), point.Z(), 1); + std::cout << "Difference in position: " << measurementPoint.Distance(finalPos) << " mm" << std::endl; + + /*** Propagate to new measurement point ****/ + propagator.SetState(startPos, startMom); + + point = XYZPoint({-3.6942, -6.13106, 1.45025}); // measurement point in cm + point *= 10; // Convert to mm + measurementPoint = AtMeasurementPoint(point); + + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); + + finalPos = propagator.GetPosition(); + finalMom = propagator.GetMomentum(); + + ASSERT_NEAR(finalPos.X(), point.X(), 1); // Check final position is within 1 mm of the measurement point + ASSERT_NEAR(finalPos.Y(), point.Y(), 1); + ASSERT_NEAR(finalPos.Z(), point.Z(), 1); + std::cout << "Difference in position: " << measurementPoint.Distance(finalPos) << " mm" << std::endl; +} From 2d24c7aaae6e67962d3e60042e32c71e498bba35 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 23 Jul 2025 11:47:15 +0200 Subject: [PATCH 34/71] Try to setup env in github runner --- AtTools/AtPropagatorTest.cxx | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 81cd1ad5b..8f1e67a5c 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -17,6 +17,14 @@ using namespace AtTools; const double mass_p = 938.272; // Mass of proton in MeV/c^2 const double charge_p = 1.602176634e-19; // Charge of proton +std::string getEnergyPath() +{ + auto env = std::getenv("VMCWORKDIR"); + if (env == nullptr) { + return "../../resources/energy_loss/HinH.txt"; // Default path assuming cwd is build/AtTools + } + return std::string(env) + "/resources/energy_loss/HinH.txt"; // Use environment variable +} class DummyELossModel : public AtELossModel { public: double eLoss = 1; @@ -105,7 +113,8 @@ TEST(AtPropagatorTest, PropagateToPoint_StoppingNoField) auto elossModel = std::make_unique(0); // elossModel->LoadSrimTable( // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); - elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools + + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; AtMeasurementPoint measurementPoint({1e3, 0, 0}); @@ -150,7 +159,7 @@ TEST(AtPropagatorTest, PropagateToPoint_NoField) auto elossModel = std::make_unique(0); // elossModel->LoadSrimTable( // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); - elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; AtMeasurementPoint measurementPoint({10, 0, 0}); @@ -187,7 +196,7 @@ TEST(AtPropagatorTest, PropagateToPlane_NoField) auto elossModel = std::make_unique(0); // elossModel->LoadSrimTable( // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); - elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; @@ -227,7 +236,7 @@ TEST(AtPropagatorTest, PropagateToPlane_StoppingNoField) auto elossModel = std::make_unique(0); // elossModel->LoadSrimTable( // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); - elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4Stepper stepper; @@ -263,7 +272,7 @@ TEST(AtPropagatorTest, PropagateToPointAdaptive_NoField) auto elossModel = std::make_unique(0); // elossModel->LoadSrimTable( // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); - elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtPropagator propagator(charge, mass, std::move(elossModel)); AtRK4AdaptiveStepper stepper; AtMeasurementPoint measurementPoint({10, 0, 0}); @@ -316,8 +325,8 @@ TEST(AtPropagatorTest, PropagateToPoint_Field) auto elossModel = std::make_unique(0); // elossModel->LoadSrimTable( // "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); - elossModel->LoadSrimTable("../../resources/energy_loss/HinH.txt"); // Assumes cwd is build/AtTools - elossModel->SetDensity(3.3084e-05); // Set density in g/cm^3 for 300 torr H2 + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path + elossModel->SetDensity(3.3084e-05); // Set density in g/cm^3 for 300 torr H2 AtPropagator propagator(charge, mass, std::move(elossModel)); propagator.SetEField({0, 0, 0}); // No electric field propagator.SetBField({0, 0, 2.85}); // Magnetic field From 0f009dcf0cd813e9658009fe479512e244d9427d Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 23 Jul 2025 11:59:13 +0200 Subject: [PATCH 35/71] Second attempt at setting CI env --- .github/workflows/CI-tests.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/CI-tests.yml b/.github/workflows/CI-tests.yml index 810be15de..1f4652135 100644 --- a/.github/workflows/CI-tests.yml +++ b/.github/workflows/CI-tests.yml @@ -25,6 +25,9 @@ jobs: # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type run: cmake -B ${{github.workspace}}/build -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} + - name: Configure Environment + run: source ${{github.workspace}}/build/config.sh + - name: Build # Build your program with the given configuration run: cmake --build ${{github.workspace}}/build --config ${{env.BUILD_TYPE}} -j $(nproc) From 8a29bd05f3281409ad4333899f5eb33534126829 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 23 Jul 2025 12:04:15 +0200 Subject: [PATCH 36/71] Try 3 - didn't run at all --- .github/workflows/CI-tests.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/CI-tests.yml b/.github/workflows/CI-tests.yml index 1f4652135..fd8ccb341 100644 --- a/.github/workflows/CI-tests.yml +++ b/.github/workflows/CI-tests.yml @@ -27,6 +27,7 @@ jobs: - name: Configure Environment run: source ${{github.workspace}}/build/config.sh + shell: bash - name: Build # Build your program with the given configuration From aacb5273aba9549ef6d59b678a0182f41d75dcc9 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 23 Jul 2025 12:06:54 +0200 Subject: [PATCH 37/71] Probably really works --- .github/workflows/CI-tests.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/workflows/CI-tests.yml b/.github/workflows/CI-tests.yml index fd8ccb341..810be15de 100644 --- a/.github/workflows/CI-tests.yml +++ b/.github/workflows/CI-tests.yml @@ -25,10 +25,6 @@ jobs: # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type run: cmake -B ${{github.workspace}}/build -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} - - name: Configure Environment - run: source ${{github.workspace}}/build/config.sh - shell: bash - - name: Build # Build your program with the given configuration run: cmake --build ${{github.workspace}}/build --config ${{env.BUILD_TYPE}} -j $(nproc) From 9b41edb9141f1e0070761875af5dc2f33f5cd8fa Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 23 Jul 2025 12:27:55 +0200 Subject: [PATCH 38/71] Add integration test for drawing figures --- AtTools/AtPropagator.cxx | 15 ++ AtTools/AtPropagator.h | 9 + macro/tests/UKF/AtPropagator.C | 76 ++++++++ macro/tests/UKF/hits.txt | 314 +++++++++++++++++++++++++++++++++ 4 files changed, 414 insertions(+) create mode 100644 macro/tests/UKF/AtPropagator.C create mode 100644 macro/tests/UKF/hits.txt diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index ee1a0cb6c..da36cca07 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -58,6 +58,21 @@ AtPropagator::XYZVector AtPropagator::d2xds2(const XYZPoint &pos, const XYZVecto return 1 / p * (dpds_vec - phat * (phat.Dot(dpds_vec))); // Second derivative of position w.r.t. arc length } +void AtPropagator::PropagateOneStep(AtStepper &stepper) +{ + if (fState.h == 0) + fState.h = stepper.GetInitialStep(); // Set the initial step size + + stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; + + auto result = stepper.Step(fState); + if (!result) { + LOG(error) << "Integration step failed, aborting propagation."; + return; // Abort propagation if step failed + } + fState = result; // Update the internal state +} + void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) { LOG(info) << "Propagating to measurement surface"; diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 0fd07435c..2e199b497 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -104,6 +104,7 @@ class AtPropagator { fState.fPos = pos; fState.fMom = mom; } + const StepState &GetState() const { return fState; } XYZPoint GetPosition() const { return fState.fPos; } XYZVector GetMomentum() const { return fState.fMom; } @@ -121,6 +122,14 @@ class AtPropagator { void PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper); + /** + * @brief Propagate the particle using the given stepper. + * Propagates one step using the provided stepper. + * + * @param stepper The stepper to use for propagation. + */ + void PropagateOneStep(AtStepper &stepper); + /** * @brief Calculate the force acting on the particle. * diff --git a/macro/tests/UKF/AtPropagator.C b/macro/tests/UKF/AtPropagator.C new file mode 100644 index 000000000..80774a23c --- /dev/null +++ b/macro/tests/UKF/AtPropagator.C @@ -0,0 +1,76 @@ +std::string getEnergyPath() +{ + auto env = std::getenv("VMCWORKDIR"); + if (env == nullptr) { + return "../../resources/energy_loss/HinH.txt"; // Default path assuming cwd is build/AtTools + } + return std::string(env) + "/resources/energy_loss/HinH.txt"; // Use environment variable +} + +const double mass_p = 938.272; // Mass of proton in MeV/c^2 +const double charge_p = 1.602176634e-19; // Charge of proton + +// This test should plot the trajectory of a particle in a magnetic field using +// the output from GEANT and the AtPropagator class. +void AtPropagator() +{ + using namespace AtTools; + + std::vector x, y, z; + std::vector x2, y2, z2; + + std::ifstream infile("hits.txt"); + double xi, yi, zi, Ei; + while (infile >> xi >> yi >> zi >> Ei) { + x.push_back(xi * 10); + y.push_back(yi * 10); + z.push_back(zi * 10); + } + + // Our propagator setup + double charge = charge_p; // Charge in Coulombs + double mass = mass_p; // Mass in MeV/c^2 + auto elossModel = std::make_unique(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)); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 2.85}); // Magnetic field + AtTools::AtRK4Stepper stepper; + + XYZPoint startPos(-3.40046e-05, -1.49863e-05, 0.10018); // Start position in cm + startPos *= 10; // Convert to mm + XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c + startMom *= 1e3; // Convert to MeV/c + + propagator.SetState(startPos, startMom); + + // Loop through until the particle is stopped + while (Kinematics::KE(propagator.GetState().fMom, propagator.GetState().fMass) > 0.1) { + // Propagate to the next point + propagator.PropagateOneStep(stepper); + + // Get the current position and momentum + auto pos = propagator.GetPosition(); + + // Store the position for plotting + x2.push_back(pos.X()); + y2.push_back(pos.Y()); + z2.push_back(pos.Z()); + } + + TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); + track->SetTitle("Particle Track;X [mm];Y [mm];Z [mm]"); + track->SetMarkerStyle(20); + track->SetMarkerSize(0.8); + + TGraph2D *track2 = new TGraph2D(x2.size(), x2.data(), y2.data(), z2.data()); + track2->SetTitle("Propagated Particle Track;X [mm];Y [mm];Z [mm]"); + track2->SetMarkerStyle(21); + track2->SetMarkerSize(0.8); + track2->SetMarkerColor(kRed); + + TCanvas *c1 = new TCanvas("c1", "Particle Track", 800, 600); + track->Draw("P"); + track2->Draw("PSAME"); +} \ No newline at end of file diff --git a/macro/tests/UKF/hits.txt b/macro/tests/UKF/hits.txt new file mode 100644 index 000000000..5f76f895d --- /dev/null +++ b/macro/tests/UKF/hits.txt @@ -0,0 +1,314 @@ +-3.40046e-05 -1.49863e-05 0.10018 0 +0.0189471 -0.0966173 0.117714 0.00258037 +0.036285 -0.193528 0.135248 0.0010415 +0.0519068 -0.290728 0.152794 0.000541655 +0.0656858 -0.388206 0.170338 0.00194783 +0.0777186 -0.485911 0.187907 0.00412452 +0.0878053 -0.583817 0.205582 0.00169653 +0.0960944 -0.681881 0.223318 0.00205692 +0.102599 -0.780065 0.241129 0.000151766 +0.107452 -0.878339 0.258972 0.00118051 +0.110508 -0.976677 0.276864 0.00104222 +0.111671 -1.07506 0.294729 0.0035253 +0.111012 -1.17344 0.312619 0.00308229 +0.108501 -1.27179 0.330529 0.00160089 +0.104209 -1.37008 0.348443 0.000583303 +0.0982273 -1.46827 0.366392 0.00753745 +0.0903449 -1.56634 0.384264 0.00259151 +0.0806727 -1.66425 0.402186 0.00123866 +0.0691679 -1.76195 0.420098 0.00378331 +0.0559072 -1.85943 0.438033 0.00155726 +0.0403906 -1.95681 0.454628 0.000931355 +0.0230265 -2.05397 0.470691 0.00707037 +0.00386458 -2.1508 0.486689 0.00451979 +-0.0172429 -2.24723 0.502676 0.000756807 +-0.040074 -2.34325 0.518767 0.00152287 +-0.0646638 -2.43883 0.53488 0.00209089 +-0.0909762 -2.53395 0.550959 0.00253983 +-0.119025 -2.62858 0.567065 0.000634231 +-0.148834 -2.72266 0.583183 0.00193713 +-0.180375 -2.81617 0.599325 0.000131761 +-0.213376 -2.90915 0.615642 0.00343676 +-0.248073 -3.00151 0.631919 0.0029619 +-0.284436 -3.0932 0.648338 0.00100025 +-0.322544 -3.1842 0.664698 0.00209536 +-0.362341 -3.27446 0.681099 0.00147181 +-0.403855 -3.36395 0.697465 0.00263272 +-0.446964 -3.45269 0.713791 0.00307259 +-0.49172 -3.54062 0.73008 0.000251718 +-0.538022 -3.62773 0.74641 0.000878999 +-0.585925 -3.71398 0.762729 0.00357119 +-0.63545 -3.79932 0.778955 0.00164382 +-0.686629 -3.88363 0.795469 0.0012815 +-0.739481 -3.9669 0.811974 0.00430081 +-0.793872 -4.04916 0.828547 0.000666631 +-0.849819 -4.13036 0.845154 0.00341725 +-0.907271 -4.2105 0.861818 0.00646584 +-0.966224 -4.28955 0.878398 0.0016157 +-1.02671 -4.36742 0.895043 0.00470272 +-1.08867 -4.44412 0.911731 0.000640672 +-1.15203 -4.51964 0.928504 0.00518449 +-1.21675 -4.59399 0.945318 0.00193881 +-1.28287 -4.66712 0.962047 0.000694573 +-1.35037 -4.73898 0.978776 0.00125284 +-1.41928 -4.80951 0.995432 0.000281865 +-1.4895 -4.8787 1.01217 0.00203059 +-1.56108 -4.94649 1.02892 0.0011867 +-1.63397 -5.01288 1.04565 0.00125706 +-1.70809 -5.07786 1.06248 0.00812622 +-1.78342 -5.14141 1.0794 0.00211417 +-1.85991 -5.20359 1.09625 0.00246438 +-1.93756 -5.26433 1.113 0.00164815 +-2.01636 -5.32357 1.12976 0.000725691 +-2.09635 -5.38118 1.14655 0.00244181 +-2.17745 -5.43721 1.16334 0.00103049 +-2.25963 -5.49165 1.18015 0.00192944 +-2.34288 -5.54443 1.19699 0.00149925 +-2.42712 -5.59561 1.21388 0.00251826 +-2.51229 -5.64519 1.23083 0.000926565 +-2.59843 -5.69307 1.24776 0.0020667 +-2.68552 -5.73922 1.26464 0.000874783 +-2.77346 -5.78369 1.28164 0.00165935 +-2.86226 -5.82643 1.29861 0.000426512 +-2.95185 -5.86746 1.31565 0.00106954 +-3.0422 -5.90683 1.33254 0.00296486 +-3.13335 -5.94435 1.34937 0.00219723 +-3.22523 -5.98001 1.36629 0.000638691 +-3.31777 -6.01396 1.38312 0.00169432 +-3.41098 -6.04601 1.4 0.000985825 +-3.5048 -6.07624 1.41687 0.00221741 +-3.59923 -6.10458 1.43358 0.000573825 +-3.6942 -6.13106 1.45025 0.00238183 +-3.78972 -6.15558 1.46684 0.00207704 +-3.88566 -6.17838 1.48345 0.000456725 +-3.98202 -6.19931 1.50007 0.000939948 +-4.07874 -6.21849 1.5167 0.000472584 +-4.17584 -6.23571 1.53327 0.00105617 +-4.27324 -6.25105 1.5499 0.0010503 +-4.37094 -6.26453 1.56645 0.00100485 +-4.46887 -6.27607 1.58306 0.00503887 +-4.56699 -6.2857 1.59976 0.00703127 +-4.66531 -6.29321 1.61638 0.00334247 +-4.76374 -6.29885 1.63314 0.00331044 +-4.86224 -6.30255 1.64995 0.00073014 +-4.9608 -6.30434 1.66675 0.0033623 +-5.05941 -6.30428 1.68339 0.00535778 +-5.158 -6.30231 1.69999 0.00125925 +-5.25655 -6.29836 1.71646 0.00204625 +-5.355 -6.29256 1.73302 0.00118473 +-5.45332 -6.28474 1.74952 0.000353303 +-5.55145 -6.2749 1.76601 0.00177079 +-5.64938 -6.2631 1.78248 0.00324638 +-5.74706 -6.24947 1.79897 0.00210638 +-5.84443 -6.2338 1.81553 0.00294219 +-5.94147 -6.21615 1.83199 0.00309201 +-6.03816 -6.19668 1.84845 0.00203747 +-6.13449 -6.17531 1.8647 0.00511082 +-6.23034 -6.15191 1.88094 0.00501013 +-6.32571 -6.12657 1.89718 0.00261361 +-6.42049 -6.09919 1.91348 0.00235029 +-6.51472 -6.06993 1.92974 0.00343822 +-6.60831 -6.0387 1.946 0.00569468 +-6.70125 -6.00557 1.96227 0.00298844 +-6.79347 -5.97049 1.97856 0.00283765 +-6.88492 -5.93348 1.99486 0.00251219 +-6.9756 -5.89462 2.0112 0.00333844 +-7.06541 -5.85385 2.02769 0.00139711 +-7.15432 -5.81119 2.04427 0.00457198 +-7.24233 -5.76668 2.06081 0.00183194 +-7.32941 -5.72039 2.07733 0.00210649 +-7.41551 -5.67229 2.09385 0.00146382 +-7.50064 -5.6225 2.11039 0.00213869 +-7.58468 -5.57094 2.12706 0.0044235 +-7.66763 -5.51765 2.1438 0.000694616 +-7.74947 -5.46267 2.16051 0.00246119 +-7.83012 -5.40594 2.17712 0.0023732 +-7.9096 -5.34756 2.19371 0.00286895 +-7.98777 -5.28744 2.21029 0.000563268 +-8.06466 -5.22571 2.22691 0.00198034 +-8.14019 -5.16232 2.24354 0.00232113 +-8.21438 -5.09739 2.26026 0.00149609 +-8.28724 -5.03093 2.27682 0.00112195 +-8.35863 -4.96291 2.29342 0.00461114 +-8.42857 -4.8934 2.31008 0.00188045 +-8.49702 -4.8224 2.32661 0.00237145 +-8.56394 -4.74996 2.34315 0.000600861 +-8.6293 -4.67611 2.35968 0.00217227 +-8.69297 -4.60078 2.37615 0.00469567 +-8.75504 -4.5241 2.39252 0.00190909 +-8.81544 -4.44612 2.40896 0.00436395 +-8.87398 -4.36673 2.42536 0.00375093 +-8.93089 -4.28614 2.44172 0.00169305 +-8.98603 -4.20437 2.45822 0.00205361 +-9.03946 -4.12146 2.47469 0.00651101 +-9.09112 -4.03747 2.4913 0.00196265 +-9.14097 -3.95242 2.50803 0.000512267 +-9.18891 -3.86627 2.52481 0.00193968 +-9.23497 -3.77914 2.54172 0.00264933 +-9.27901 -3.69096 2.55858 0.00512826 +-9.32122 -3.60189 2.57545 0.00234888 +-9.36127 -3.51181 2.59219 0.0041278 +-9.39927 -3.42083 2.60889 0.00394476 +-9.43523 -3.32903 2.62559 0.00626612 +-9.4692 -3.23646 2.64222 0.000954117 +-9.50126 -3.14323 2.65897 0.00068102 +-9.53154 -3.04941 2.67572 0.00117896 +-9.55987 -2.95499 2.69251 0.00440072 +-9.58636 -2.86011 2.70969 0.00454022 +-9.61094 -2.76471 2.72688 0.00324835 +-9.63324 -2.66874 2.74393 0.00130486 +-9.65335 -2.57229 2.76103 0.00461174 +-9.67143 -2.47544 2.77816 0.00257374 +-9.68746 -2.37823 2.79528 0.00149974 +-9.70129 -2.28067 2.81234 0.00222129 +-9.71289 -2.1828 2.82926 0.00402998 +-9.72231 -2.08472 2.84631 0.00510242 +-9.72954 -1.98643 2.86322 0.00427729 +-9.73469 -1.888 2.88007 0.00285612 +-9.73773 -1.78947 2.89691 0.00215436 +-9.73855 -1.69093 2.91392 0.00178831 +-9.73715 -1.5924 2.93087 0.00886926 +-9.7336 -1.4939 2.94776 0.00359489 +-9.72787 -1.39549 2.96459 0.00169257 +-9.71993 -1.29722 2.9813 0.00271018 +-9.70987 -1.19916 2.99809 0.00207105 +-9.69759 -1.10135 3.01489 0.00110787 +-9.68311 -1.00385 3.03175 0.00443925 +-9.66639 -0.906729 3.04869 0.00466506 +-9.64755 -0.809974 3.06552 0.00218618 +-9.62655 -0.713686 3.08246 0.00540695 +-9.60321 -0.617922 3.09933 0.00178606 +-9.57794 -0.522661 3.11624 0.00490372 +-9.55052 -0.427988 3.13312 0.0020954 +-9.5209 -0.334002 3.15012 0.00149254 +-9.48924 -0.240682 3.16712 0.00538447 +-9.45551 -0.148121 3.18427 0.00312085 +-9.41973 -0.0563445 3.20148 0.00155854 +-9.38187 0.0346355 3.21849 0.00141054 +-9.34195 0.124743 3.23541 0.00628454 +-9.29985 0.213875 3.25223 0.00376027 +-9.25561 0.301932 3.26921 0.00207038 +-9.20936 0.388958 3.28615 0.00546088 +-9.16116 0.474953 3.3029 0.00611755 +-9.11083 0.559761 3.31945 0.00548134 +-9.05847 0.643318 3.33608 0.00600543 +-9.00422 0.725669 3.35267 0.00233206 +-8.94817 0.806832 3.36911 0.00304746 +-8.89026 0.886668 3.38559 0.000929448 +-8.83035 0.96505 3.40192 0.00375323 +-8.76866 1.04206 3.41817 0.00100678 +-8.70529 1.11771 3.43432 0.00250021 +-8.64014 1.19182 3.45052 0.00498279 +-8.57323 1.26439 3.46653 0.0045065 +-8.50462 1.33533 3.48265 0.00271847 +-8.43438 1.40462 3.4989 0.00125119 +-8.36242 1.47215 3.51502 0.00592709 +-8.28883 1.53789 3.53123 0.00554028 +-8.21372 1.60193 3.54727 0.00357857 +-8.13701 1.66402 3.5634 0.00670072 +-8.05872 1.72404 3.57977 0.0056806 +-7.97891 1.78203 3.59609 0.00242805 +-7.89768 1.83806 3.61225 0.00477683 +-7.81498 1.89188 3.62853 0.00434346 +-7.73107 1.94389 3.64443 0.00370326 +-7.64573 1.99355 3.66029 0.00179178 +-7.55918 2.0411 3.67601 0.00395621 +-7.47151 2.0865 3.69186 0.00552891 +-7.38269 2.12963 3.70769 0.00401387 +-7.29272 2.17023 3.72372 0.00357124 +-7.20175 2.20856 3.73968 0.0023659 +-7.10988 2.24472 3.75555 0.000686133 +-7.01708 2.27832 3.77159 0.00335168 +-6.92342 2.30947 3.78762 0.00215076 +-6.82897 2.33807 3.8038 0.0038759 +-6.73384 2.36436 3.81987 0.00601634 +-6.63806 2.38822 3.83587 0.00322579 +-6.54173 2.40968 3.85197 0.00534104 +-6.44488 2.42855 3.86824 0.00193304 +-6.34768 2.44507 3.88488 0.000624437 +-6.25007 2.4593 3.90132 0.00186719 +-6.15216 2.47101 3.9179 0.00585969 +-6.05396 2.48007 3.9345 0.00482671 +-5.95555 2.48643 3.95102 0.00396111 +-5.85699 2.49014 3.96753 0.00325424 +-5.75836 2.49127 3.98395 0.00156398 +-5.65969 2.48966 4.00013 0.00376579 +-5.56111 2.48531 4.0163 0.00445 +-5.46269 2.47853 4.03264 0.00726225 +-5.36453 2.46889 4.0491 0.00526993 +-5.26666 2.45677 4.06564 0.00606185 +-5.16916 2.44191 4.08217 0.00703524 +-5.07198 2.42499 4.0986 0.00514467 +-4.97514 2.40634 4.11514 0.00376785 +-4.8789 2.38496 4.13185 0.00565096 +-4.78333 2.36073 4.14855 0.0037355 +-4.68843 2.3339 4.1651 0.00553483 +-4.59438 2.30437 4.18186 0.00654986 +-4.50109 2.27238 4.19839 0.00664659 +-4.40862 2.23804 4.21478 0.00382731 +-4.31694 2.20158 4.23109 0.00598498 +-4.22605 2.16309 4.2471 0.0048853 +-4.13638 2.12183 4.26309 0.00251647 +-4.04781 2.07827 4.27914 0.00667167 +-3.96054 2.03206 4.29488 0.00526986 +-3.87461 1.98342 4.31068 0.00493183 +-3.79018 1.93218 4.32638 0.0049293 +-3.70734 1.87843 4.34207 0.00384186 +-3.62593 1.82247 4.35762 0.00262308 +-3.54625 1.76407 4.37309 0.00171779 +-3.46825 1.70342 4.3885 0.00421162 +-3.39201 1.64052 4.40366 0.00472597 +-3.31764 1.57537 4.41864 0.00473131 +-3.24538 1.50787 4.43347 0.00393331 +-3.17543 1.43802 4.44858 0.00539462 +-3.10784 1.3659 4.4637 0.00439951 +-3.04248 1.29179 4.47902 0.00379672 +-2.97943 1.21573 4.4945 0.00658802 +-2.91898 1.13763 4.51017 0.00606156 +-2.86095 1.05768 4.52567 0.00384396 +-2.80568 0.975826 4.54127 0.00531629 +-2.75292 0.892352 4.55699 0.00849088 +-2.70273 0.807312 4.57274 0.00465122 +-2.6553 0.720703 4.58852 0.00597102 +-2.61083 0.63255 4.60435 0.00444418 +-2.56834 0.543396 4.62001 0.00679056 +-2.52841 0.453078 4.63572 0.00762534 +-2.49192 0.361235 4.65097 0.00403043 +-2.45884 0.268045 4.66581 0.00671071 +-2.42906 0.173723 4.68048 0.00419242 +-2.40246 0.0784021 4.6948 0.00486919 +-2.37922 -0.0177892 4.70915 0.00907706 +-2.35945 -0.114742 4.72358 0.00737428 +-2.34287 -0.212286 4.73804 0.00748826 +-2.32996 -0.310375 4.75256 0.00449589 +-2.32048 -0.408811 4.76736 0.00571873 +-2.31453 -0.50755 4.78199 0.00665646 +-2.31252 -0.606374 4.7971 0.00376963 +-2.31409 -0.705117 4.81279 0.00617805 +-2.31923 -0.803835 4.82785 0.00547875 +-2.32858 -0.902283 4.84266 0.00787066 +-2.34165 -1.00028 4.85765 0.00620081 +-2.35864 -1.09758 4.87318 0.00880822 +-2.37961 -1.19409 4.88882 0.00935348 +-2.40467 -1.28948 4.9053 0.00793634 +-2.43373 -1.38378 4.92148 0.00829069 +-2.46697 -1.47665 4.93789 0.0098699 +-2.50301 -1.56848 4.95432 0.0101575 +-2.54236 -1.65894 4.97067 0.00987816 +-2.58544 -1.74788 4.98578 0.00680355 +-2.63341 -1.83437 5.00048 0.0109609 +-2.68466 -1.91879 5.01616 0.0104423 +-2.74016 -2.0004 5.03219 0.0145097 +-2.80018 -2.07879 5.04797 0.0122098 +-2.86443 -2.15356 5.06462 0.0123009 +-2.93286 -2.22468 5.08063 0.0087081 +-3.0053 -2.29149 5.09751 0.0125225 +-3.08243 -2.35284 5.1143 0.0133418 +-3.16441 -2.40731 5.13179 0.0142218 +-3.25044 -2.45546 5.14835 0.0142074 +-3.34096 -2.49508 5.16332 0.0134456 +-3.43483 -2.52575 5.17852 0.0119897 +-3.52011 -2.5456 5.19268 0.0120695 +-3.58404 -2.5518 5.20013 0.00473408 +-3.63536 -2.54465 5.19907 0.00380166 +-3.66951 -2.53396 5.20118 0.00210337 From b3efecb50503bd7543bfc390b0e3f9b3d6a0a904 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 12:06:25 +0200 Subject: [PATCH 39/71] Remove old UKF tests in prep for real first test --- .../kalman_filter/TrackFitterUKFTest.cxx | 380 +----------------- 1 file changed, 8 insertions(+), 372 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 183e96c0c..ca0671954 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -218,45 +218,17 @@ class TrackFitterUKFPhysicsTest : public testing::Test { static constexpr float FLOAT_EPSILON{0.001F}; - static constexpr size_t DIM_X{6}; - static constexpr size_t DIM_V{2}; - static constexpr size_t DIM_Z{3}; - static constexpr size_t DIM_N{3}; + static constexpr size_t DIM_X{6}; // Set state vector to be (x,y,z,p,theta,phi) + static constexpr size_t DIM_V{1}; // Noise vector is energy straggling + static constexpr size_t DIM_Z{3}; // Measurement vector is (x,y,z) + static constexpr size_t DIM_N{3}; // Measurement noise vector - static XYZVector fBField; // B-field in tesla - static XYZVector fEField; // E-field in V/m + static XYZVector fBField{0, 0, 2.85}; // B-field in tesla + static XYZVector fEField{0, 0, 0}; // E-field in V/m static constexpr double fC = 299792458; // m/s kf::TrackFitterUKF m_ukf; - /** - * @param pos Position of particle in mm - * @param mom Momentum of particle in MeV/c - * @param charge charge of the particle in Coulombs - * @param mass mass of particle in MeV/c^2 - * @param dedx Stopping power in MeV/mm - * - * @returns Force in N - */ - static XYZVector Force(XYZVector pos, XYZVector mom, double charge, double mass, double dedx) - { - - // auto fourMom = AtTools::Kinematics::Get4Vector(mom, mass); - // auto v = mom / fourMom.E() * c; // m/s - auto v = AtTools::Kinematics::GetVel(mom, mass); - - auto F_lorentz = charge * (fEField + v.Cross(fBField)); - // std::cout << "F_lorentz: " << F_lorentz << std::endl; - auto dedx_si = dedx * 1.60218e-10; // de_dx in SI units (J/m) - - auto drag = -dedx_si * mom.Unit(); - // std::cout << "drag: " << drag << " mom " << mom << " dedx " << dedx_si << std::endl; - - return F_lorentz + drag; - } - - static double dist(const XYZVector &x, const XYZVector &z) { return std::sqrt((x - z).Mag2()); } - /// @brief to propagate the state vector using the process model /// @param x state vector /// @param v process noise vector @@ -264,341 +236,5 @@ class TrackFitterUKFPhysicsTest : public testing::Test { /// @return propagated (unaugmented) state vector static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z) { - std::cout << "Staring to run funcF" << std::endl; - // TODO: This needs to be filled with an RK4 solver for the physics model - kf::Vector y{x}; - XYZVector measurement(z[0], z[1], z[2]); // Measurement point in mm - - double charge = x[6]; - double eLoss = v[0]; - double mass = 938.272; // Mass in MeV/c^2 - - double mat_density = 0; // Density of the material in g/cm^3 - AtTools::AtELossTable dedxModel(mat_density); - dedxModel.LoadSrimTable( - "/home/adam/fair_install/ATTPCROOTv2/AtReconstruction/AtFitter/OpenKF/kalman_filter/HinH.txt"); // Load the - // SRIM table - // for energy - // loss - double scalingFactor = 1.0; - int iterations = 0; - double calc_eLoss = 0; - - while (std::abs(calc_eLoss - eLoss) > 1e-4) { - std::cout << "Running iteration " << iterations << " with scaling factor: " << scalingFactor - << " and energy loss: " << calc_eLoss << std::endl; - - if (iterations > 100) { - // If we are not converging, we should probably throw an error. - throw std::runtime_error("Energy loss did not converge after 100 iterations."); - } - - // Variables needed in a single run of the RK4 solver. This section needs to be repeated - // until the energy loss converges to the correct value. - double h = 1e-10; // Timestep in s (100 ns to start) - double lastApproach = std::numeric_limits::max(); - bool approaching = true; - iterations++; - - XYZVector pos(x[0], x[1], x[2]); - XYZVector mom(x[3], x[4], x[5]); - - double KE_initial = - AtTools::Kinematics::KE(mom, mass); // std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV - - while (true) { - XYZVector lastPos = pos; - XYZVector lastMom = mom; - std::cout << "Position: " << pos.X() << ", " << pos.Y() << ", " << pos.Z() << std::endl; - std::cout << "Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; - - // Using timestep, propagate state forward one step. - double KE = AtTools::Kinematics::KE(mom, mass); // Kinetic energy in MeV - auto dedx = scalingFactor * dedxModel.GetdEdx(KE); // Get the stopping power in MeV/mm - // std::cout << "KE: " << KE << " dedx: " << dedx << std::endl; - - auto spline = dedxModel.GetSpline(); - // std::cout << "Spline: " << spline.get_x_min() << " to " << spline.get_x_max() << std::endl; - // std::cout << "dxde " << spline(KE) << " dxde " << dedx << std::endl; - - auto x_k1 = AtTools::Kinematics::GetVel(mom, mass); - auto p_k1 = Force(pos, mom, charge, mass, dedx); - // std::cout << "vel: " << x_k1 << " speed " << x_k1.R() << std::endl; - // std::cout << "Force: " << p_k1 << std::endl; - - auto x_k2 = AtTools::Kinematics::GetVel(mom + p_k1 * h / 2, mass); - auto p_k2 = Force(pos + x_k1 * h / 2, mom + p_k1 * h / 2, charge, mass, dedx); - // std::cout << "vel: " << x_k2 << " speed " << x_k2.R() << std::endl; - // std::cout << "Force: " << p_k2 << std::endl; - - auto x_k3 = AtTools::Kinematics::GetVel(mom + p_k2 * h / 2, mass); - auto p_k3 = Force(pos + x_k2 * h / 2, mom + p_k2 * h / 2, charge, mass, dedx); - // std::cout << "vel: " << x_k3 << " speed " << x_k3.R() << std::endl; - // std::cout << "Force: " << p_k3 << std::endl; - - auto x_k4 = AtTools::Kinematics::GetVel(mom + p_k3 * h, mass); - auto p_k4 = Force(pos + x_k3 * h, mom + p_k3 * h, charge, mass, dedx); - // std::cout << "vel: " << x_k4 << " speed " << x_k4.R() << std::endl; - // std::cout << "Force: " << p_k4 << std::endl; - - auto mom_SItoMeV = 1.60218e-13 / 299792458; // Factor to convert momentum to MeV/c from kg*m/s - auto F_SI = (p_k1 + 2 * p_k2 + 2 * p_k3 + p_k4) / 6; - - // Convert momentum to SI, update, and convert back to MeV/c - auto mom_SI = mom * mom_SItoMeV; - mom_SI += F_SI * h; - mom = mom_SI / mom_SItoMeV; // Convert back to MeV/c - - // Convert position to SI, update, and convert back to mm - auto pos_SI = pos / 1e3; // Convert mm to m - pos_SI += (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) * h / 6; - pos = pos_SI * 1e3; // Convert back to mm - - // std::cout << "Average force: " << F_SI << " N" << std::endl; - // std::cout << "Momentum: " << mom * mom_SItoMeV << " kg m/s" << std::endl; - // std::cout << "Delta x: " << (x_k1 + 2 * x_k2 + 2 * x_k3 + x_k4) / 6 << " m/s" << std::endl; - - auto approach = dist(pos, measurement); - // std::cout << "pos: " << pos << " measurement" << measurement << std::endl; - // std::cout << "Approach: " << approach << " last approach: " << lastApproach << std::endl; - if (approach < lastApproach) { - // We are still approaching the measurement point - approaching = true; - lastApproach = approach; - - continue; - } - - bool reachedMeasurementPoint = (approaching && approach > lastApproach); - bool particleStopped = std::sqrt(mom.Mag2() + mass * mass) - mass < 0.01; - if (reachedMeasurementPoint || particleStopped) { - // Last iteration we were still approaching the measurement point. Now we are further away - // then before. We have probably reached the measurement point if things are well behaved. - // I can think of cases where this will not be true. A better solution might be to run - // tracking the point of closest approach until the distance between the current state and - // the measurement point is larger than the distance between the last state and the measurement point. - - // Undo the last step since we were closer last time. - y[0] = lastPos.X(); - y[1] = lastPos.Y(); - y[2] = lastPos.Z(); - y[3] = lastMom.X(); - y[4] = lastMom.Y(); - y[5] = lastMom.Z(); - - // Update the scaling factor - double KE_final = std::sqrt(lastMom.Mag2() + mass * mass) - mass; - calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - scalingFactor *= eLoss / calc_eLoss; - std::cout << "------- End of RK4 interation " << iterations << " ---------" << std::endl; - std::cout << "Particle stopped: " << particleStopped << std::endl; - std::cout << "Reached measurement point: " << reachedMeasurementPoint << std::endl; - std::cout << "Last approach: " << lastApproach << " Current approach: " << approach << std::endl; - std::cout << "Desired energy loss: " << eLoss << " MeV" << std::endl; - std::cout << "Calculated energy loss: " << calc_eLoss << " MeV" << std::endl; - std::cout << "Difference: " << calc_eLoss - eLoss << " MeV" << std::endl; - std::cout << "New scaling factor: " << scalingFactor << std::endl; - std::cout << "Final Position: " << pos.X() << ", " << pos.Y() << ", " << pos.Z() << std::endl; - std::cout << "Final Momentum: " << mom.X() << ", " << mom.Y() << ", " << mom.Z() << std::endl; - break; - // return y; - } - } // End of loop over RK4 integration - } // End loop over energy loss convergence - - return y; - } - - /// @brief to apply the measurement model to the state vector - /// @param x the state vector of the system - /// @return the measurement vector - static kf::Vector funcH(const kf::Vector &x) - { - kf::Vector y; - y[0] = x[0]; - y[1] = x[1]; - y[2] = x[2]; - - return y; - } - - const double mass_p = 938.272; // Mass of proton in MeV/c^2 - const double charge_p = 1.602176634e-19; // Charge of proton -}; - -// Definition of static member variables -TrackFitterUKFPhysicsTest::XYZVector TrackFitterUKFPhysicsTest::fBField; -TrackFitterUKFPhysicsTest::XYZVector TrackFitterUKFPhysicsTest::fEField; - -TEST_F(TrackFitterUKFPhysicsTest, PhysicsPrediction) -{ - // TODO: This needs to be filled with a proper physics test. - kf::Vector x; // Initial state vector - - kf::Matrix P; // Initial state vector covariance matrix - - // Note: process noise is defined in the UKF class, so we don't need to set it here. - - kf::Vector z; // Measurement vector to be used in the correction step - z << std::sqrt(5), std::atan2(1.f, 2.f); - - kf::Matrix R; // Covariance matrix for the measurement noise - - m_ukf.vecX() = x; - m_ukf.matP() = P; - - // m_ukf.setCovarianceR(R); - // m_ukf.predictUKF(funcF, z); - ASSERT_EQ(true, true); -} -/* - -TEST_F(TrackFitterUKFPhysicsTest, TestForceNoFields) -{ - XYZVector pos(0, 0, 0); // Position in mm - XYZVector mom(100, 0, 0); // Momentum in MeV/c - fBField = XYZVector(0, 0, 0); // B-field in tesla - fEField = XYZVector(0, 0, 0); // E-field - - double charge = charge_p; // Charge in Coulombs - double mass = mass_p; // Mass in MeV/c^2 - double dedx = 1; // Stopping power in MeV/mm - - auto force = Force(pos, mom, charge, mass, dedx); - - ASSERT_NEAR(force.X(), -1.602e-10, FLOAT_EPSILON); - ASSERT_NEAR(force.Y(), 0, FLOAT_EPSILON); - ASSERT_NEAR(force.Z(), 0, FLOAT_EPSILON); - - mom = XYZVector(100, 0, 100); // Reset momentum - force = Force(pos, mom, charge, mass, dedx); - ASSERT_NEAR(force.X(), -1.602e-10 / std::sqrt(2), FLOAT_EPSILON); - ASSERT_NEAR(force.Y(), 0, FLOAT_EPSILON); - ASSERT_NEAR(force.Z(), -1.602e-10 / std::sqrt(2), FLOAT_EPSILON); -} - -TEST_F(TrackFitterUKFPhysicsTest, TestForceEField) -{ - XYZVector pos(0, 0, 0); // Position in mm - XYZVector mom(100, 0, 0); // Momentum in MeV/c - fBField = XYZVector(0, 0, 1); // B-field in tesla - fEField = XYZVector(0, 0, 0); // E-field in V/m - - double charge = charge_p; // Charge in Coulombs - double mass = mass_p; // Mass in MeV/c^2 - double dedx = 0; // Stopping power in MeV/mm - - auto force = Force(pos, mom, charge, mass, dedx); - - ASSERT_NEAR(force.X(), 0, FLOAT_EPSILON); - ASSERT_NEAR(force.Y(), 0, FLOAT_EPSILON); - ASSERT_NEAR(force.Z(), 1.121e-14, FLOAT_EPSILON); -} - -TEST_F(TrackFitterUKFPhysicsTest, TestForceBField) -{ - XYZVector pos(0, 0, 0); // Position in mm - XYZVector mom(100, 0, 0); // Momentum in MeV/c - fBField = XYZVector(0, 0, 1); // B-field in tesla - fEField = XYZVector(0, 0, 0); // E-field - - double charge = charge_p; // Charge in Coulombs - double mass = mass_p; // Mass in MeV/c^2 - double dedx = 0; // Stopping power in MeV/mm - - auto force = Force(pos, mom, charge, mass, dedx); - - ASSERT_NEAR(force.X(), 0, FLOAT_EPSILON); - ASSERT_NEAR(force.Y(), -5.09e-12, FLOAT_EPSILON); - ASSERT_NEAR(force.Z(), 0, FLOAT_EPSILON); -} - -TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorStoppingNoField) -{ - - double KE = 1; // Kinetic energy in MeV - double E = KE + mass_p; - double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c - fBField = XYZVector(0, 0, 0); // B-field in tesla - fEField = XYZVector(0, 0, 0); // E-field - - kf::Vector x; // Initial state vector - x[0] = 0; - x[1] = 0; - x[2] = 0; - x[3] = p; // p_x - x[4] = 0; - x[5] = 0; - - ASSERT_NEAR(x[3], 43.331, 1e-1); // Make sure momentum is calculated correctly - - kf::Vector v; // Process noise vector - v[0] = 1; // Energy loss in MeV - v[1] = 0.0; // No process noise in this example - - kf::Vector z; // Measurement vector - z[0] = 1e3; - z[1] = 0; - z[2] = 0; - - auto final = funcF(x, v, z); // Propagate the state vector using the process model - - // Check the final position is close to the stopping point from LISE - ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 - ASSERT_NEAR(final[0], 210, 10); // Final position in x-direction should be close to 210 mm - - KE = 0.5; - v[0] = 0.5; // Energy loss in MeV - E = KE + mass_p; - p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c - x[3] = p; // Reset momentum - final = funcF(x, v, z); // Propagate the state vector using the - - ASSERT_NEAR(final[3], 0, 0.1); // Final momentum in x-direction should be close to 0 - ASSERT_NEAR(final[0], 68.6, 5); // Final position in x -} -*/ - -TEST_F(TrackFitterUKFPhysicsTest, TestPropagatorNoField) -{ - double KE = 1; // Kinetic energy in MeV - double E = KE + mass_p; - double p = std::sqrt(E * E - mass_p * mass_p); // Momentum in MeV/c - fBField = XYZVector(0, 0, 0); // B-field in tesla - fEField = XYZVector(0, 0, 0); // E-field - - kf::Vector x; // Initial state vector - x[0] = 0; - x[1] = 0; - x[2] = 0; - x[3] = p; // p_x - x[4] = 0; - x[5] = 0; - - ASSERT_NEAR(x[3], 43.331, 1e-1); // Make sure momentum is calculated correctly - - kf::Vector v; // Process noise vector - v[0] = 0.0285; // Energy loss in MeV - v[1] = 0.0; // No process noise in this example - - kf::Vector z; // Measurement vector - z[0] = 10; // Measure after 10 mm - z[1] = 0; - z[2] = 0; - - auto final = funcF(x, v, z); // Propagate the state vector using the process model - - // Check the final position is close to the stopping point from LISE - double E_fin = KE - v[0] + mass_p; - double p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); - ASSERT_NEAR(final[3], p_fin, 0.1); // Final momentum in x-direction - ASSERT_NEAR(final[0], 10, .5); // Final position in x-direction should be close to 10 mm - - v[0] = 0.3237; // Energy loss in MeV - z[0] = 100; // Measure after 100 mm - final = funcF(x, v, z); // Propagate the state vector using the process model - E_fin = KE - v[0] + mass_p; - p_fin = std::sqrt(E_fin * E_fin - mass_p * mass_p); - ASSERT_NEAR(final[3], p_fin, 0.1); // Final momentum - ASSERT_NEAR(final[0], 100, .5); // Final position -} + ASSERT_EQ(true, true); + } \ No newline at end of file From a711861b80c14bcb41388976941a0d0a03e8607e Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 14:01:37 +0200 Subject: [PATCH 40/71] Reset KF changes --- AtTools/AtPropagator.h | 15 +++++++++++---- AtTools/AtPropagatorTest.cxx | 5 +++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 2e199b497..153d4fa88 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -51,9 +51,9 @@ class AtPropagator { using Plane3D = ROOT::Math::Plane3D; // Variables used for the force - XYZVector fEField{0, 0, 0}; // Electric field vector - XYZVector fBField{0, 0, 0}; // Magnetic field vector - const std::unique_ptr fELossModel; // Energy loss model + XYZVector fEField{0, 0, 0}; // Electric field vector + XYZVector fBField{0, 0, 0}; // Magnetic field vector + std::unique_ptr fELossModel; // Energy loss model // Internal state variables for the propagator StepState fState; /// Current state of the particle @@ -80,7 +80,7 @@ class AtPropagator { fState.fMass = mass; fState.fQ = charge; } - + AtPropagator(AtPropagator &&) = default; /** * @brief Set the electric field (V/m) */ @@ -105,6 +105,7 @@ class AtPropagator { fState.fMom = mom; } const StepState &GetState() const { return fState; } + const AtELossModel *GetELossModel() const { return fELossModel.get(); } XYZPoint GetPosition() const { return fState.fPos; } XYZVector GetMomentum() const { return fState.fMom; } @@ -246,6 +247,12 @@ class AtMeasurementPoint : public AtMeasurementSurface { public: AtMeasurementPoint(const ROOT::Math::XYZPoint &point) : fPoint(point) {} + template + AtMeasurementPoint(const T &point) + { + fPoint = ROOT::Math::XYZPoint(point[0], point[1], point[2]); + } + double Distance(const ROOT::Math::XYZPoint &pos) const override { return (fPoint - pos).R(); } bool PassedSurface(AtPropagator::StepState &result) const override; ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override { return pos; } diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx index 8f1e67a5c..dbcd93c95 100644 --- a/AtTools/AtPropagatorTest.cxx +++ b/AtTools/AtPropagatorTest.cxx @@ -16,6 +16,7 @@ using namespace AtTools; const double mass_p = 938.272; // Mass of proton in MeV/c^2 const double charge_p = 1.602176634e-19; // Charge of proton +namespace { std::string getEnergyPath() { @@ -25,6 +26,7 @@ std::string getEnergyPath() } return std::string(env) + "/resources/energy_loss/HinH.txt"; // Use environment variable } +} // namespace class DummyELossModel : public AtELossModel { public: double eLoss = 1; @@ -34,6 +36,9 @@ class DummyELossModel : public AtELossModel { double GetRange(double /*energyIni*/, double /*energyFin = 0*/) const override { return 1.0; } double GetEnergyLoss(double /*energyIni*/, double /*distance*/) const override { return 1.0; } double GetEnergy(double /*energyIni*/, double /*distance*/) const override { return 1.0; } + double GetElossStraggling(double /*energyIni*/, double /*energyFin*/) const override { return 0.0; } + double GetdEdxStraggling(double /*energyIni*/, double /*energyFin*/) const override { return 0.0; } + double GetRangeVariance(double /*energy*/) const override { return 0.0; } }; TEST(AtPropagatorTest, ForceNoField) From 6f74ffda5cd02c1bffad4753d82f899be0270f36 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 14:03:38 +0200 Subject: [PATCH 41/71] Original tests building again. Need fresh slate --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index ca0671954..c2979bb18 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -223,8 +223,8 @@ class TrackFitterUKFPhysicsTest : public testing::Test { static constexpr size_t DIM_Z{3}; // Measurement vector is (x,y,z) static constexpr size_t DIM_N{3}; // Measurement noise vector - static XYZVector fBField{0, 0, 2.85}; // B-field in tesla - static XYZVector fEField{0, 0, 0}; // E-field in V/m + XYZVector fBField{0, 0, 2.85}; // B-field in tesla + XYZVector fEField{0, 0, 0}; // E-field in V/m static constexpr double fC = 299792458; // m/s kf::TrackFitterUKF m_ukf; @@ -236,5 +236,6 @@ class TrackFitterUKFPhysicsTest : public testing::Test { /// @return propagated (unaugmented) state vector static kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z) { - ASSERT_EQ(true, true); - } \ No newline at end of file + return {}; + } +}; \ No newline at end of file From b288edd702f53b8616bc4cad2fa152d82ac54643 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 14:06:09 +0200 Subject: [PATCH 42/71] Base class with virtual tests pass --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h | 10 +++++----- .../OpenKF/kalman_filter/TrackFitterUKFTest.cxx | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 3481cb98f..a6f062941 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -24,7 +24,7 @@ namespace kf { /// @tparam DIM_V Dimension of the process noise vector. /// @tparam DIM_N Dimension of the measurement noise vector. template -class TrackFitterUKF : public KalmanFilter { +class TrackFitterUKFBase : public KalmanFilter { public: // Augmented state vector is just the process noise and state vector. The measurement noise is not included as that // is independent of the propagation and measurement model and just adds linearly. @@ -38,7 +38,7 @@ class TrackFitterUKF : public KalmanFilter { Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix - TrackFitterUKF() + TrackFitterUKFBase() : KalmanFilter(), m_kappa(3 - DIM_A), m_matQ(Matrix::Zero()), m_matR(Matrix::Zero()) { @@ -46,7 +46,7 @@ class TrackFitterUKF : public KalmanFilter { updateWeights(); } - ~TrackFitterUKF() {} + ~TrackFitterUKFBase() {} void setKappa(float32_t kappa) { @@ -95,7 +95,7 @@ class TrackFitterUKF : public KalmanFilter { } } - std::array calculateProcessNoiseMean() + virtual std::array calculateProcessNoiseMean() { // Calculate the expectation value of the process noise using the current value of the state vector m_vecX std::array processNoiseMean{0}; @@ -104,7 +104,7 @@ class TrackFitterUKF : public KalmanFilter { return processNoiseMean; } - Matrix calculateProcessNoiseCovariance() + virtual Matrix calculateProcessNoiseCovariance() { // Calculate the process noise covariance matrix Matrix matQ{Matrix::Zero()}; diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index c2979bb18..d9714f0da 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -18,7 +18,7 @@ class TrackFitterUKFExampleTest : public testing::Test { static constexpr size_t DIM_Z{2}; static constexpr size_t DIM_N{2}; - kf::TrackFitterUKF m_ukf; + kf::TrackFitterUKFBase m_ukf; /// @brief to propagate the state vector using the process model /// @param x state vector @@ -227,7 +227,7 @@ class TrackFitterUKFPhysicsTest : public testing::Test { XYZVector fEField{0, 0, 0}; // E-field in V/m static constexpr double fC = 299792458; // m/s - kf::TrackFitterUKF m_ukf; + kf::TrackFitterUKFBase m_ukf; /// @brief to propagate the state vector using the process model /// @param x state vector From 73f294b1bca8893e62690a1fe5eb0a20d6032279 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 14:19:02 +0200 Subject: [PATCH 43/71] UKF w/ prop is instantiating --- .../OpenKF/kalman_filter/TrackFitterUKF.h | 73 ++++++++++++++++++- .../kalman_filter/TrackFitterUKFTest.cxx | 34 ++++++++- 2 files changed, 100 insertions(+), 7 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index a6f062941..c0d2ad4cf 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -9,8 +9,10 @@ /// @file unscented_kalman_filter.h /// -#ifndef UNSCENTED_KALMAN_FILTER_LIB_H -#define UNSCENTED_KALMAN_FILTER_LIB_H +#ifndef TRACKFITTERUKF_H +#define TRACKFITTERUKF_H + +#include "AtPropagator.h" #include "kalman_filter.h" #include "kf_util.h" @@ -370,6 +372,71 @@ class TrackFitterUKFBase : public KalmanFilter { return matPxy; } }; + +// Define template dimension variables for clarity and reuse + +class TrackFitterUKF : public TrackFitterUKFBase<6, 1, 3, 3> { +protected: + static constexpr int32_t TRACKFITTER_DIM_X = 6; + static constexpr int32_t TRACKFITTER_DIM_Z = 1; + static constexpr int32_t TRACKFITTER_DIM_V = 3; + static constexpr int32_t TRACKFITTER_DIM_N = 3; + + AtTools::AtPropagator fPropagator; ///< @brief Propagator for the track fitter + AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state + +public: + /** + * @brief Constructor for the TrackFitterUKF class. + * @param propagator The propagator to be used for the track fitting, must be passed as an rvalue reference. + * + * Example usage: + * ``` + * AtTools::AtPropagator propagator; + * kf::TrackFitterUKF trackFitterUKF(std::move(propagator)); + * ``` + */ + TrackFitterUKF(AtTools::AtPropagator &&propagator) : TrackFitterUKFBase(), fPropagator(std::move(propagator)) {} + + template + void predictUKF(PredictionModelCallback predictionModelFunc, const Vector &vecZ) + { + // First we need to propagate the mean state vector to the next measurement point. + AtTools::AtRK4Stepper stepper; // Use RK4 stepper for propagation + fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(vecZ), stepper); + fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator + } + +protected: + std::array calculateProcessNoiseMean() override + { + // The process noise is the scaling factor for dedx. By definition the mean should be 1 + std::array processNoiseMean{1}; + return processNoiseMean; + } + + Matrix calculateProcessNoiseCovariance() override + { + assert(TRACKFITTER_DIM_V == 1 && "Process noise covariance is only implemented for DIM_V = 1"); + // Calculate the process noise covariance matrix + Matrix matQ{Matrix::Zero()}; + + // We need to know what the energy of the particle before/after transport. + double eIn = 0; + double eOut = 0; + + if (const auto *elossModel = fPropagator.GetELossModel()) { + double dedx_straggle = elossModel->GetdEdxStraggling(eIn, eOut); + double factor = dedx_straggle / elossModel->GetdEdx(eIn); + matQ(0, 0) = factor * factor; // Variance for the dedx straggling. + + } else { + throw std::runtime_error("Cannot calculate process noise covariance without an energy loss model"); + } + // TODO: Add multiple scattering + return matQ; + } +}; } // namespace kf -#endif // UNSCENTED_KALMAN_FILTER_LIB_H +#endif // TRACKFITTERUKF_H diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index d9714f0da..b43064d0a 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -209,11 +209,22 @@ TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection) ASSERT_NEAR(m_ukf.matP()(3, 3), 0.1333886F, FLOAT_EPSILON); } -class TrackFitterUKFPhysicsTest : public testing::Test { +namespace { + +std::string getEnergyPath() +{ + auto env = std::getenv("VMCWORKDIR"); + if (env == nullptr) { + return "../../resources/energy_loss/HinH.txt"; // Default path assuming cwd is build/AtTools + } + return std::string(env) + "/resources/energy_loss/HinH.txt"; // Use environment variable +} +} // namespace + +class TrackFitterUKFFixture : public testing::Test { public: using XYZVector = ROOT::Math::XYZVector; - virtual void SetUp() override {} virtual void TearDown() override {} static constexpr float FLOAT_EPSILON{0.001F}; @@ -227,7 +238,17 @@ class TrackFitterUKFPhysicsTest : public testing::Test { XYZVector fEField{0, 0, 0}; // E-field in V/m static constexpr double fC = 299792458; // m/s - kf::TrackFitterUKFBase m_ukf; + std::unique_ptr m_ukf{nullptr}; + + void SetUp() override + { + const double mass_p = 938.272; // Mass of proton in MeV/c^2 + const double charge_p = 1.602176634e-19; // Charge of protonble + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path + AtTools::AtPropagator propagator(charge_p, mass_p, std::move(elossModel)); + m_ukf = std::make_unique(std::move(propagator)); + } /// @brief to propagate the state vector using the process model /// @param x state vector @@ -238,4 +259,9 @@ class TrackFitterUKFPhysicsTest : public testing::Test { { return {}; } -}; \ No newline at end of file +}; + +TEST_F(TrackFitterUKFFixture, TestInstantiation) +{ + assert(true); +} \ No newline at end of file From 8d670562939730516182a3fc7e37be413e232640 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 19:00:06 +0200 Subject: [PATCH 44/71] Prediction code runs --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 86 +++++++++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 118 ++++++------ .../kalman_filter/TrackFitterUKFTest.cxx | 172 +++++++++++++++++- AtReconstruction/CMakeLists.txt | 1 + AtTools/AtKinematics.cxx | 9 + AtTools/AtKinematics.h | 3 + AtTools/AtPropagator.h | 2 + 7 files changed, 325 insertions(+), 66 deletions(-) create mode 100644 AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx new file mode 100644 index 000000000..4534689b7 --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -0,0 +1,86 @@ +#include "TrackFitterUKF.h" + +#include + +#include +#include + +#include + +namespace kf { + +void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition, + const ROOT::Math::XYZVector &initialMomentum, const TMatrixD &initialCovariance) +{ + m_vecX[0] = initialPosition.X(); // X position + m_vecX[1] = initialPosition.Y(); // Y position + m_vecX[2] = initialPosition.Z(); // Z position + m_vecX[3] = initialMomentum.R(); // Momentum magnitude + m_vecX[4] = initialMomentum.Theta(); // Polar angle + m_vecX[5] = initialMomentum.Phi(); // Azimuthal angle + + // Copy elements from initialCovariance to m_matP + for (int i = 0; i < m_matP.rows(); ++i) { + for (int j = 0; j < m_matP.cols(); ++j) { + m_matP(i, j) = initialCovariance(i, j); + } + } +} + +std::array TrackFitterUKF::calculateProcessNoiseMean() +{ + // The process noise is the scaling factor for dedx. By definition the mean should be 1 + std::array processNoiseMean{1}; + return processNoiseMean; +} + +Matrix TrackFitterUKF::calculateProcessNoiseCovariance() +{ + assert(TF_DIM_V == 1 && "Process noise covariance is only implemented for DIM_V = 1"); + // Calculate the process noise covariance matrix + Matrix matQ{Matrix::Zero()}; + + // We need to know what the energy of the particle before/after transport. + double eIn = AtTools::Kinematics::KE(fMeanStep.fLastMom, fMeanStep.fMass); + double eOut = AtTools::Kinematics::KE(fMeanStep.fMom, fMeanStep.fMass); + + if (const auto *elossModel = fPropagator.GetELossModel()) { + double dedx_straggle = elossModel->GetdEdxStraggling(eIn, eOut); + double factor = dedx_straggle / elossModel->GetdEdx(eIn); + matQ(0, 0) = factor * factor; // Variance for the dedx straggling. + + } else { + throw std::runtime_error("Cannot calculate process noise covariance without an energy loss model"); + } + // TODO: Add multiple scattering + return matQ; +} + +Vector TrackFitterUKF::funcF(const Vector &x, + const Vector &v, + const Vector &z) +{ + // Set the state of the propagator to the current state vector + using namespace ROOT::Math; + XYZPoint fPos(x[0], x[1], x[2]); // Position from state vector + Polar3DVector fMom(x[3], x[4], x[5]); // Momentum from state vector + + fPropagator.SetState(fPos, XYZVector(fMom)); + + fPropagator.fScalingFactor = v[0]; // Set the scaling factor for energy loss + fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPlane(fMeasurementPlane), *fStepper); + fPropagator.fScalingFactor = 1.0; // Reset the scaling factor after propagation + + auto fState = fPropagator.GetState(); // Get the propagated state + Vector vecX{Vector::Zero()}; + vecX[0] = fState.fPos.X(); // X position + vecX[1] = fState.fPos.Y(); // Y position + vecX[2] = fState.fPos.Z(); // Z position + vecX[3] = fState.fMom.R(); // Momentum magnitude + vecX[4] = fState.fMom.Theta(); // Polar angle + vecX[5] = fState.fMom.Phi(); // Azimuthal + + return vecX; // Return the propagated state vector +} + +} // namespace kf \ No newline at end of file diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index c0d2ad4cf..42459991c 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -12,12 +12,18 @@ #ifndef TRACKFITTERUKF_H #define TRACKFITTERUKF_H +#include "AtKinematics.h" #include "AtPropagator.h" +#include + +#include +#include +#include + #include "kalman_filter.h" #include "kf_util.h" #include "unscented_kalman_filter.h" - namespace kf { /// @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. @@ -27,17 +33,31 @@ namespace kf { /// @tparam DIM_N Dimension of the measurement noise vector. template class TrackFitterUKFBase : public KalmanFilter { + public: // Augmented state vector is just the process noise and state vector. The measurement noise is not included as that // is independent of the propagation and measurement model and just adds linearly. static constexpr int32_t DIM_A{DIM_X + DIM_V}; ///< @brief Augmented state dimension static constexpr int32_t SIGMA_DIM_A{2 * DIM_A + 1}; ///< @brief Sigma points dimension for augmented state - float32_t m_kappa{0}; ///< @brief Kappa parameter for finding sigma points + +protected: + using KalmanFilter::m_vecX; // from Base KalmanFilter class + using KalmanFilter::m_matP; // from Base KalmanFilter class + + float32_t m_weight0; /// @brief unscented transform weight 0 for mean + float32_t m_weighti; /// @brief unscented transform weight i for none mean samples + float32_t m_kappa{0}; ///< @brief Kappa parameter for finding sigma points + + Vector m_vecXa{Vector::Zero()}; /// @brief augmented state vector (incl. process + /// and measurement noise means) + Matrix m_matPa{Matrix::Zero()}; /// @brief augmented state covariance (incl. + /// process and measurement noise covariances) // Add variables to track the covariances of the process and measurement noise. Matrix m_matQ; // @brief Process noise covariance matrix Matrix m_matR; // @brief Measurement noise covariance matrix +public: Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix TrackFitterUKFBase() @@ -56,13 +76,6 @@ class TrackFitterUKFBase : public KalmanFilter { updateWeights(); // Update the weights based on the new kappa value } - // This code uses two different conventions for managing noise. - // The state vector noise is set in the updateAugmentedStateAndCovariance() method, while - // the noise vectors for the process and measurement models are set in the setCovarianceQ() and - // setCovarianceR() methods. This is an odd choice. We will be moving everything into a common - // structure where updateAugmentedStateAndCovariance() handle all covariance updates that are actually - // part of the augmented state vector. - /// /// @brief adding process noise covariance Q to the augmented state covariance /// matPa in the middle element of the diagonal. @@ -232,18 +245,7 @@ class TrackFitterUKFBase : public KalmanFilter { m_matP -= matK * matPzz * matK.transpose(); } -private: - using KalmanFilter::m_vecX; // from Base KalmanFilter class - using KalmanFilter::m_matP; // from Base KalmanFilter class - - float32_t m_weight0; /// @brief unscented transform weight 0 for mean - float32_t m_weighti; /// @brief unscented transform weight i for none mean samples - - Vector m_vecXa{Vector::Zero()}; /// @brief augmented state vector (incl. process - /// and measurement noise means) - Matrix m_matPa{Matrix::Zero()}; /// @brief augmented state covariance (incl. - /// process and measurement noise covariances) - +protected: /// /// @brief algorithm to calculate the weights used to draw the sigma points /// @@ -276,6 +278,10 @@ class TrackFitterUKFBase : public KalmanFilter { // cholesky factorization to get matrix Pxx square-root Eigen::LLT> lltOfPa(matPa); if (lltOfPa.info() != Eigen::Success) { + LOG(error) << "Cholesky decomposition failed, matrix is not positive definite."; + for (int32_t i{0}; i < STATE_DIM; ++i) { + LOG(error) << "Pxx[" << i << "]: " << matPa(i, i); + } throw std::runtime_error("Cholesky decomposition failed, matrix is not positive definite."); } Matrix matSa{lltOfPa.matrixL()}; // sqrt(P_{a}) @@ -375,16 +381,17 @@ class TrackFitterUKFBase : public KalmanFilter { // Define template dimension variables for clarity and reuse -class TrackFitterUKF : public TrackFitterUKFBase<6, 1, 3, 3> { +class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { protected: - static constexpr int32_t TRACKFITTER_DIM_X = 6; - static constexpr int32_t TRACKFITTER_DIM_Z = 1; - static constexpr int32_t TRACKFITTER_DIM_V = 3; - static constexpr int32_t TRACKFITTER_DIM_N = 3; + static constexpr int32_t TF_DIM_X = 6; + static constexpr int32_t TF_DIM_Z = 3; + static constexpr int32_t TF_DIM_V = 1; + static constexpr int32_t TF_DIM_N = 3; - AtTools::AtPropagator fPropagator; ///< @brief Propagator for the track fitter + AtTools::AtPropagator fPropagator; ///< @brief Propagator for the track fitter + std::unique_ptr fStepper{nullptr}; AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state - + ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter public: /** * @brief Constructor for the TrackFitterUKF class. @@ -396,46 +403,37 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 1, 3, 3> { * kf::TrackFitterUKF trackFitterUKF(std::move(propagator)); * ``` */ - TrackFitterUKF(AtTools::AtPropagator &&propagator) : TrackFitterUKFBase(), fPropagator(std::move(propagator)) {} - - template - void predictUKF(PredictionModelCallback predictionModelFunc, const Vector &vecZ) + TrackFitterUKF(AtTools::AtPropagator &&propagator, std::unique_ptr &&stepper) + : TrackFitterUKFBase(), fPropagator(std::move(propagator)), fStepper(std::move(stepper)) { - // First we need to propagate the mean state vector to the next measurement point. - AtTools::AtRK4Stepper stepper; // Use RK4 stepper for propagation - fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(vecZ), stepper); - fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator } -protected: - std::array calculateProcessNoiseMean() override - { - // The process noise is the scaling factor for dedx. By definition the mean should be 1 - std::array processNoiseMean{1}; - return processNoiseMean; - } + void SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, + const TMatrixD &initialCovariance); - Matrix calculateProcessNoiseCovariance() override - { - assert(TRACKFITTER_DIM_V == 1 && "Process noise covariance is only implemented for DIM_V = 1"); - // Calculate the process noise covariance matrix - Matrix matQ{Matrix::Zero()}; + kf::Vector + funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z); - // We need to know what the energy of the particle before/after transport. - double eIn = 0; - double eOut = 0; + void predictUKF(const Vector &z) + { + // First we need to propagate the mean state vector to the next measurement point. + fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(z), *fStepper); + fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator - if (const auto *elossModel = fPropagator.GetELossModel()) { - double dedx_straggle = elossModel->GetdEdxStraggling(eIn, eOut); - double factor = dedx_straggle / elossModel->GetdEdx(eIn); - matQ(0, 0) = factor * factor; // Variance for the dedx straggling. + // Now we can construct the reference plane. + using namespace ROOT::Math; + fMeasurementPlane = Plane3D(fMeanStep.fMom.Unit(), + XYZPoint(z)); // Create a plane using the momentum direction and position - } else { - throw std::runtime_error("Cannot calculate process noise covariance without an energy loss model"); - } - // TODO: Add multiple scattering - return matQ; + auto callback = [this](const kf::Vector &x_, const kf::Vector &v_, + const kf::Vector &z_) { return funcF(x_, v_, z_); }; + TrackFitterUKFBase::predictUKF(callback, z); } + +protected: + std::array calculateProcessNoiseMean() override; + + Matrix calculateProcessNoiseCovariance() override; }; } // namespace kf diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index b43064d0a..4ca0f983c 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -2,6 +2,11 @@ #include "AtELossTable.h" #include "AtKinematics.h" +#include "AtPropagator.h" + +#include + +#include "TMatrixD.h" #include "Math/Vector3D.h" #include "gtest/gtest.h" @@ -221,6 +226,28 @@ std::string getEnergyPath() } } // namespace +namespace kf { +class TrackFitterUKFTest : public TrackFitterUKF { +public: + Vector vecXa() { return m_vecXa; } + Matrix matPa() { return m_matPa; } + Matrix matSigmaXa() { return m_matSigmaXa; } + AtTools::AtPropagator &getPropagator() { return fPropagator; } + AtTools::AtPropagator::StepState &getState() { return fMeanStep; } + + void calcSigmaPoints() { m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); } + void calcMeanAndCovFromSigma(const Matrix &sigmaXa, Vector &meanXa, + Matrix &covXa) + { + calculateWeightedMeanAndCovariance(sigmaXa, meanXa, covXa); + } + + TrackFitterUKFTest(AtTools::AtPropagator &&propagator, std::unique_ptr &&stepper) + : TrackFitterUKF(std::move(propagator), std::move(stepper)) + { + } +}; +} // namespace kf class TrackFitterUKFFixture : public testing::Test { public: using XYZVector = ROOT::Math::XYZVector; @@ -234,20 +261,23 @@ class TrackFitterUKFFixture : public testing::Test { static constexpr size_t DIM_Z{3}; // Measurement vector is (x,y,z) static constexpr size_t DIM_N{3}; // Measurement noise vector - XYZVector fBField{0, 0, 2.85}; // B-field in tesla - XYZVector fEField{0, 0, 0}; // E-field in V/m - static constexpr double fC = 299792458; // m/s + XYZVector fBField{0, 0, 2.85}; // B-field in tesla + XYZVector fEField{0, 0, 0}; // E-field in V/m + static constexpr double fC = 299792458; // m/s + const double mass_p = 938.272; // Mass of proton in MeV/c^2 + const double charge_p = 1.602176634e-19; // Charge of protonble std::unique_ptr m_ukf{nullptr}; void SetUp() override { - const double mass_p = 938.272; // Mass of proton in MeV/c^2 - const double charge_p = 1.602176634e-19; // Charge of protonble auto elossModel = std::make_unique(0); elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtTools::AtPropagator propagator(charge_p, mass_p, std::move(elossModel)); - m_ukf = std::make_unique(std::move(propagator)); + propagator.SetBField(fBField); + propagator.SetEField(fEField); + auto stepper = std::make_unique(); // Use RK4 stepper for propagation + m_ukf = std::make_unique(std::move(propagator), std::move(stepper)); } /// @brief to propagate the state vector using the process model @@ -264,4 +294,134 @@ class TrackFitterUKFFixture : public testing::Test { TEST_F(TrackFitterUKFFixture, TestInstantiation) { assert(true); +} + +TEST_F(TrackFitterUKFFixture, TestAugmentedStateUpdate) +{ + auto testPtr = dynamic_cast(m_ukf.get()); + + // Create a state change where a proton has lost 1 MeV of energy (to rest) + auto p = AtTools::Kinematics::GetRelMomFromKE(1.0, mass_p); + AtTools::AtPropagator::StepState fInitialState; + fInitialState.fMass = mass_p; + fInitialState.fQ = charge_p; + + // Current state + fInitialState.fLastPos = {0, 0, 0}; + fInitialState.fLastMom = {p, 0, 0}; + // State at next measurement point + fInitialState.fPos = {202.78, 0, 0}; // Pulled from SRIM table + fInitialState.fMom = {0, 0, 0}; + double var = 0.01; + TMatrixD cov(6, 6); + cov.Zero(); + for (int i = 0; i < 6; ++i) { + cov(i, i) = var; // Set diagonal covariance to some small number + } + + m_ukf->SetInitialState(fInitialState.fLastPos, fInitialState.fLastMom, cov); + testPtr->getState() = fInitialState; // Set the mean state in the test class (rather than propagate) + + // Fill m_vecXa and m_matPa with the provided state + m_ukf->updateAugmentedStateAndCovariance(); + + assert(testPtr != nullptr); + + // Check the augmented state vector + double eps = 1e-4; + ASSERT_NEAR(testPtr->vecXa()[0], fInitialState.fLastPos.X(), eps); + ASSERT_NEAR(testPtr->vecXa()[1], fInitialState.fLastPos.Y(), eps); + ASSERT_NEAR(testPtr->vecXa()[2], fInitialState.fLastPos.Z(), eps); + ASSERT_NEAR(testPtr->vecXa()[3], fInitialState.fLastMom.R(), eps); + ASSERT_NEAR(testPtr->vecXa()[4], fInitialState.fPos.Theta(), eps); + ASSERT_NEAR(testPtr->vecXa()[5], fInitialState.fPos.Phi(), eps); + ASSERT_NEAR(testPtr->vecXa()[6], 1, eps); // Average energy straggling factor is 1 + + // Check the augmented covariance matrix. Nothing set but energy straggling + ASSERT_NEAR(testPtr->matPa()(0, 0), var, eps); + ASSERT_NEAR(testPtr->matPa()(1, 1), var, eps); + ASSERT_NEAR(testPtr->matPa()(2, 2), var, eps); + ASSERT_NEAR(testPtr->matPa()(3, 3), var, eps); + ASSERT_NEAR(testPtr->matPa()(4, 4), var, eps); + ASSERT_NEAR(testPtr->matPa()(5, 5), var, eps); + + double sigma = 8.74 / 202.78; // Energy straggling factor (from range straggling table) + ASSERT_NEAR(testPtr->matPa()(6, 6), sigma * sigma, eps); +} + +TEST_F(TrackFitterUKFFixture, TestSigmaPoints) +{ + auto testPtr = dynamic_cast(m_ukf.get()); + + // Create a state change where a proton has lost 1 MeV of energy (to rest) + auto p = AtTools::Kinematics::GetRelMomFromKE(1.0, mass_p); + AtTools::AtPropagator::StepState fInitialState; + fInitialState.fMass = mass_p; + fInitialState.fQ = charge_p; + + // Current state + fInitialState.fLastPos = {0, 0, 0}; + fInitialState.fLastMom = {p, 0, 0}; + // State at next measurement point + fInitialState.fPos = {202.78, 0, 0}; // Pulled from SRIM table + fInitialState.fMom = {0, 0, 0}; + double var = 0.01; + TMatrixD cov(6, 6); + cov.Zero(); + for (int i = 0; i < 6; ++i) { + cov(i, i) = var; // Set diagonal covariance to some small number + } + + m_ukf->SetInitialState(fInitialState.fLastPos, fInitialState.fLastMom, cov); + testPtr->getState() = fInitialState; // Set the mean state in the test class (rather than propagate) + + // Fill m_vecXa and m_matPa with the provided state + m_ukf->updateAugmentedStateAndCovariance(); + + testPtr->calcSigmaPoints(); + auto sigmaPoints = testPtr->matSigmaXa(); + + kf::Matrix covXa; + kf::Vector meanXa; + testPtr->calcMeanAndCovFromSigma(sigmaPoints, meanXa, covXa); + + // Verify the mean and cov match what we put in + for (int i = 0; i < kf::TrackFitterUKF::DIM_A; ++i) { + ASSERT_NEAR(meanXa[i], testPtr->vecXa()[i], FLOAT_EPSILON); + for (int j = 0; j < kf::TrackFitterUKF::DIM_A; ++j) { + ASSERT_NEAR(covXa(i, j), testPtr->matPa()(i, j), FLOAT_EPSILON); + } + } + + LOG(debug) << "Sigma points:\n" << sigmaPoints; +} + +TEST_F(TrackFitterUKFFixture, TestPredictionStep) +{ + auto testPtr = dynamic_cast(m_ukf.get()); + + // Create a state change where a proton has lost 1 MeV of energy (to rest) + auto p = AtTools::Kinematics::GetRelMomFromKE(1.0, mass_p); + AtTools::AtPropagator::StepState fInitialState; + fInitialState.fMass = mass_p; + fInitialState.fQ = charge_p; + + // Current state + fInitialState.fLastPos = {0, 0, 0}; + fInitialState.fLastMom = {p, 0, 0}; + // State at next measurement point + fInitialState.fPos = {202.78, 0, 0}; // Pulled from SRIM table + fInitialState.fMom = {0, 0, 0}; + double var = 0.01; + TMatrixD cov(6, 6); + cov.Zero(); + for (int i = 0; i < 6; ++i) { + cov(i, i) = var; // Set diagonal covariance to some small number + } + + m_ukf->SetInitialState(fInitialState.fLastPos, fInitialState.fLastMom, cov); + testPtr->getState() = fInitialState; // Set the mean state in the test class (rather than propagate) + + // Fill m_vecXa and m_matPa with the provided state + m_ukf->updateAugmentedStateAndCovariance(); } \ No newline at end of file diff --git a/AtReconstruction/CMakeLists.txt b/AtReconstruction/CMakeLists.txt index bc0de8a62..585f5ce2a 100755 --- a/AtReconstruction/CMakeLists.txt +++ b/AtReconstruction/CMakeLists.txt @@ -120,6 +120,7 @@ if(TARGET Eigen3::Eigen) set(SRCS ${SRCS} AtPatternRecognition/triplclust/src/orthogonallsq.cxx AtFitter/AtOpenKFTest.cxx + AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx ) message(STATUS "Current sources: ${SRCS}") diff --git a/AtTools/AtKinematics.cxx b/AtTools/AtKinematics.cxx index 3f9d416f7..7470effa9 100644 --- a/AtTools/AtKinematics.cxx +++ b/AtTools/AtKinematics.cxx @@ -291,6 +291,11 @@ double GetRelMom(double gamma, double mass) { return std::sqrt(gamma * gamma - 1) * mass; } +double GetRelMomFromKE(double KE, double mass) +{ + + return std::sqrt((KE + mass) * (KE + mass) - mass * mass); +} /** * Get the mass in MeV of a fragment of mass in amu (or A) @@ -317,4 +322,8 @@ double KE(ROOT::Math::XYZVector mom, double mass) { return std::sqrt(mom.Mag2() + mass * mass) - mass; // Kinetic energy in MeV } +double KE(double mom, double mass) +{ + return std::sqrt(mom * mom + mass * mass) - mass; +} } // namespace AtTools::Kinematics diff --git a/AtTools/AtKinematics.h b/AtTools/AtKinematics.h index 622f347d9..ce09b7a44 100644 --- a/AtTools/AtKinematics.h +++ b/AtTools/AtKinematics.h @@ -85,6 +85,9 @@ double GetSpeed(double p, double mass); /// Calculate the speed of a particle gi * @returns Kinetic energy in MeV */ double KE(ROOT::Math::XYZVector mom, double mass); +double KE(double mom, double mass); + +double GetRelMomFromKE(double KE, double mass); /** * Calculate the velocity vector from momentum and mass diff --git a/AtTools/AtPropagator.h b/AtTools/AtPropagator.h index 153d4fa88..6713e81fd 100644 --- a/AtTools/AtPropagator.h +++ b/AtTools/AtPropagator.h @@ -194,6 +194,8 @@ class AtStepper { virtual AtPropagator::StepState Step(const AtPropagator::StepState &state) const = 0; virtual double GetInitialStep() const { return 1e-4; } /// Default initial step size in m + virtual ~AtStepper() = default; + protected: static constexpr double fReltoSImom = 1.60218e-13 / 299792458; // Conversion factor from MeV/c to kg m/s (SI units) }; From 9573b645907bbcf11b8e634d019faa53ae503f02 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Thu, 24 Jul 2025 19:59:22 +0200 Subject: [PATCH 45/71] Seemingly working forward step --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 41 ++++- .../OpenKF/kalman_filter/TrackFitterUKF.h | 23 ++- .../kalman_filter/TrackFitterUKFTest.cxx | 145 +++++++++++++++--- AtTools/AtELossTable.cxx | 2 +- AtTools/AtPropagator.cxx | 8 +- 5 files changed, 189 insertions(+), 30 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 4534689b7..8c3824eca 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -12,12 +12,13 @@ namespace kf { void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, const TMatrixD &initialCovariance) { - m_vecX[0] = initialPosition.X(); // X position - m_vecX[1] = initialPosition.Y(); // Y position - m_vecX[2] = initialPosition.Z(); // Z position - m_vecX[3] = initialMomentum.R(); // Momentum magnitude - m_vecX[4] = initialMomentum.Theta(); // Polar angle - m_vecX[5] = initialMomentum.Phi(); // Azimuthal angle + fPropagator.SetState(initialPosition, initialMomentum); // Set the initial state in the propagator + m_vecX[0] = initialPosition.X(); // X position + m_vecX[1] = initialPosition.Y(); // Y position + m_vecX[2] = initialPosition.Z(); // Z position + m_vecX[3] = initialMomentum.R(); // Momentum magnitude + m_vecX[4] = initialMomentum.Theta(); // Polar angle + m_vecX[5] = initialMomentum.Phi(); // Azimuthal angle // Copy elements from initialCovariance to m_matP for (int i = 0; i < m_matP.rows(); ++i) { @@ -27,6 +28,19 @@ void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition } } +void TrackFitterUKF::SetMeasCov(const TMatrixD &measCov) +{ + if (measCov.GetNrows() != TF_DIM_Z || measCov.GetNcols() != TF_DIM_Z) { + throw std::runtime_error("Measurement covariance matrix must be of size " + std::to_string(TF_DIM_Z) + "x" + + std::to_string(TF_DIM_Z)); + } + for (int i = 0; i < TF_DIM_Z; ++i) { + for (int j = 0; j < TF_DIM_Z; ++j) { + m_matR(i, j) = measCov(i, j); // Copy measurement covariance to m_matR + } + } +} + std::array TrackFitterUKF::calculateProcessNoiseMean() { // The process noise is the scaling factor for dedx. By definition the mean should be 1 @@ -83,4 +97,19 @@ Vector TrackFitterUKF::funcF(const Vector TrackFitterUKF::funcH(const Vector &x) +{ + // Measurement function to convert state vector to measurement vector + using namespace ROOT::Math; + Vector vecZ; + XYZPoint fPos(x[0], x[1], x[2]); // Position from state vector + + // Calculate the measurement vector based on the position and momentum + vecZ[0] = fPos.X(); // X coordinate + vecZ[1] = fPos.Y(); // Y coordinate + vecZ[2] = fPos.Z(); // Z coordinate + + return vecZ; // Return the measurement vector +} + } // namespace kf \ No newline at end of file diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 42459991c..a5227a2da 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -411,11 +411,15 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { void SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, const TMatrixD &initialCovariance); + void SetMeasCov(const TMatrixD &measCov); + kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z); + kf::Vector funcH(const kf::Vector &x); - void predictUKF(const Vector &z) + void predictUKF(const ROOT::Math::XYZPoint &z) { + // First we need to propagate the mean state vector to the next measurement point. fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(z), *fStepper); fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator @@ -424,10 +428,23 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { using namespace ROOT::Math; fMeasurementPlane = Plane3D(fMeanStep.fMom.Unit(), XYZPoint(z)); // Create a plane using the momentum direction and position - + Vector zVec; // Initialize the measurement vector + zVec[0] = z.X(); + zVec[1] = z.Y(); + zVec[2] = z.Z(); auto callback = [this](const kf::Vector &x_, const kf::Vector &v_, const kf::Vector &z_) { return funcF(x_, v_, z_); }; - TrackFitterUKFBase::predictUKF(callback, z); + TrackFitterUKFBase::predictUKF(callback, zVec); + } + + void correctUKF(const ROOT::Math::XYZPoint &z) + { + Vector zVec; // Initialize the measurement vector + zVec[0] = z.X(); + zVec[1] = z.Y(); + zVec[2] = z.Z(); + auto callback = [this](const kf::Vector &x_) { return funcH(x_); }; + TrackFitterUKFBase::correctUKF(callback, zVec); } protected: diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 4ca0f983c..9709f4797 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -234,6 +234,7 @@ class TrackFitterUKFTest : public TrackFitterUKF { Matrix matSigmaXa() { return m_matSigmaXa; } AtTools::AtPropagator &getPropagator() { return fPropagator; } AtTools::AtPropagator::StepState &getState() { return fMeanStep; } + Vector getStateVector() { return m_vecX; } void calcSigmaPoints() { m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); } void calcMeanAndCovFromSigma(const Matrix &sigmaXa, Vector &meanXa, @@ -268,10 +269,12 @@ class TrackFitterUKFFixture : public testing::Test { const double charge_p = 1.602176634e-19; // Charge of protonble std::unique_ptr m_ukf{nullptr}; + AtTools::AtELossModel *fElossModel{nullptr}; void SetUp() override { auto elossModel = std::make_unique(0); + fElossModel = elossModel.get(); // Store the model for later use elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path AtTools::AtPropagator propagator(charge_p, mass_p, std::move(elossModel)); propagator.SetBField(fBField); @@ -398,30 +401,136 @@ TEST_F(TrackFitterUKFFixture, TestSigmaPoints) TEST_F(TrackFitterUKFFixture, TestPredictionStep) { + using namespace AtTools; + using namespace ROOT::Math; + + // Set the conditions for simulated data auto testPtr = dynamic_cast(m_ukf.get()); + fElossModel->SetDensity(3.3084e-05); + testPtr->getPropagator().SetEField({0, 0, 0}); // No electric field + testPtr->getPropagator().SetBField({0, 0, 2.85}); // Magnetic field - // Create a state change where a proton has lost 1 MeV of energy (to rest) - auto p = AtTools::Kinematics::GetRelMomFromKE(1.0, mass_p); - AtTools::AtPropagator::StepState fInitialState; - fInitialState.fMass = mass_p; - fInitialState.fQ = charge_p; + // Attempt to predict the state of a proton from simulation + XYZPoint startPos(-3.40046e-05, -1.49863e-05, 0.10018); // Start position in cm + startPos *= 10; // Convert to mm + XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c + startMom *= 1e3; // Convert to MeV/c + + double sigma_pos = 10; // Position uncertainty of 10 mm + double sigma_mom = 0.1 * startMom.R(); // Momentum uncertainty of 10% MeV/c + double sigma_theta = 1 * M_PI / 180; // Angular uncertainty of 1 degree + double sigma_phi = 1 * M_PI / 180; // Angular uncertainty of 1 degree - // Current state - fInitialState.fLastPos = {0, 0, 0}; - fInitialState.fLastMom = {p, 0, 0}; - // State at next measurement point - fInitialState.fPos = {202.78, 0, 0}; // Pulled from SRIM table - fInitialState.fMom = {0, 0, 0}; - double var = 0.01; TMatrixD cov(6, 6); cov.Zero(); - for (int i = 0; i < 6; ++i) { - cov(i, i) = var; // Set diagonal covariance to some small number + for (int i = 0; i < 3; ++i) { + + cov(i, i) = sigma_pos * sigma_pos; // Set diagonal covariance to some small number } + cov(3, 3) = sigma_mom * sigma_mom; // Momentum uncertainty + cov(4, 4) = sigma_theta * sigma_theta; // Angular uncertainty + cov(5, 5) = sigma_phi * sigma_phi; // Angular uncertainty + + m_ukf->SetInitialState(startPos, startMom, cov); + + XYZPoint point({-1.4895, -4.8787, 1.01217}); // measurement point in cm + point *= 10; // Convert to mm + m_ukf->predictUKF(point); + + // Check the predicted state vector before correction + auto stateVec = testPtr->getStateVector(); + XYZPoint predictedPos(stateVec[0], stateVec[1], stateVec[2]); + Polar3DVector predictedMom(stateVec[3], stateVec[4], stateVec[5]); + LOG(info) << "Measurement point: " << point; + LOG(info) << "Mean position: " << testPtr->getState().fLastPos; + LOG(info) << "Predicted position (sigma): " << predictedPos; + + LOG(info) << "Mean momentum: " << testPtr->getState().fLastMom; + LOG(info) << "Predicted momentum (sigma): " << XYZPoint(predictedMom); + + LOG(info) << "Initial covariance matrix:\n"; + cov.Print(); + LOG(info) << "Predicted covariance matrix:\n" << testPtr->matP(); +} - m_ukf->SetInitialState(fInitialState.fLastPos, fInitialState.fLastMom, cov); - testPtr->getState() = fInitialState; // Set the mean state in the test class (rather than propagate) +TEST_F(TrackFitterUKFFixture, TestPredictionAndCorrectStep) +{ + using namespace AtTools; + using namespace ROOT::Math; - // Fill m_vecXa and m_matPa with the provided state - m_ukf->updateAugmentedStateAndCovariance(); + // Set the conditions for simulated data + auto testPtr = dynamic_cast(m_ukf.get()); + fElossModel->SetDensity(3.3084e-05); + testPtr->getPropagator().SetEField({0, 0, 0}); // No electric field + testPtr->getPropagator().SetBField({0, 0, 2.85}); // Magnetic field + + // Attempt to predict the state of a proton from simulation + XYZPoint startPos(-3.40046e-05, -1.49863e-05, 0.10018); // Start position in cm + startPos *= 10; // Convert to mm + XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c + startMom *= 1e3; // Convert to MeV/c + + double sigma_pos = 1; // Position uncertainty of 10 mm + double sigma_mom = 0.01 * startMom.R(); // Momentum uncertainty of 10% MeV/c + double sigma_theta = 5 * M_PI / 180; // Angular uncertainty of 1 degree + double sigma_phi = 5 * M_PI / 180; // Angular uncertainty of 1 degree + + TMatrixD cov(6, 6); + cov.Zero(); + for (int i = 0; i < 3; ++i) { + + cov(i, i) = sigma_pos * sigma_pos; // Set diagonal covariance to some small number + } + cov(3, 3) = sigma_mom * sigma_mom; // Momentum uncertainty + cov(4, 4) = sigma_theta * sigma_theta; // Angular uncertainty + cov(5, 5) = sigma_phi * sigma_phi; // Angular uncertainty + + m_ukf->SetInitialState(startPos, startMom, cov); + + XYZPoint point({-1.4895, -4.8787, 1.01217}); // measurement point in cm + point *= 10; // Convert to mm + TMatrixD cov_meas(3, 3); + cov_meas.Zero(); + for (int i = 0; i < 3; ++i) { + cov_meas(i, i) = sigma_pos * sigma_pos; + } + m_ukf->SetMeasCov(cov_meas); // Set measurement noise covariance + + /** Make prediction */ + m_ukf->predictUKF(point); + + // Check the predicted state vector before correction + auto stateVec = testPtr->getStateVector(); + XYZPoint predictedPos(stateVec[0], stateVec[1], stateVec[2]); + Polar3DVector predictedMom(stateVec[3], stateVec[4], stateVec[5]); + + LOG(info) << "Prediction step complete: " << std::endl; + + LOG(info) << "Measurement point: " << point; + LOG(info) << "Mean position: " << testPtr->getState().fLastPos; + LOG(info) << "Predicted position (sigma): " << predictedPos; + + LOG(info) << "Mean momentum: " << testPtr->getState().fLastMom; + LOG(info) << "Predicted momentum (sigma): " << XYZPoint(predictedMom); + + LOG(info) << "Initial covariance matrix:\n"; + cov.Print(); + LOG(info) << "Predicted covariance matrix:\n" << testPtr->matP(); + + // Now correct the state with the measurement + m_ukf->correctUKF(point); + auto correctedStateVec = testPtr->getStateVector(); + XYZPoint correctedPos(correctedStateVec[0], correctedStateVec[1], correctedStateVec[2]); + Polar3DVector correctedMom(correctedStateVec[3], correctedStateVec[4], correctedStateVec[5]); + + LOG(info) << "Correction complete: " << std::endl; + + LOG(info) << "Measurement point: " << point; + LOG(info) << "Predicted position (sigma): " << predictedPos; + LOG(info) << "Corrected position: " << correctedPos; + + LOG(info) << "Predicted momentum (sigma): " << XYZPoint(predictedMom); + LOG(info) << "Corrected momentum: " << XYZPoint(correctedMom); + + LOG(info) << "Corrected covariance matrix:\n" << testPtr->matP(); } \ No newline at end of file diff --git a/AtTools/AtELossTable.cxx b/AtTools/AtELossTable.cxx index 8f093a118..33159c92d 100644 --- a/AtTools/AtELossTable.cxx +++ b/AtTools/AtELossTable.cxx @@ -26,7 +26,7 @@ void AtELossTable::SetDensity(double density) elem = fScaling / elem; // scale dx/de by the density } for (auto &elem : r_var) { - LOG(info) << "Scaling range variance by " << fScaling * fScaling; + LOG(debug) << "Scaling range variance by " << fScaling * fScaling; elem = elem / (fScaling * fScaling); // scale range variance by the density } LoadTable(x, y); diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index da36cca07..8b458ed6d 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -79,11 +79,15 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur fState.h = stepper.GetInitialStep(); // Set the initial step size auto KE_initial = Kinematics::KE(fState.fMom, fState.fMass); + if (KE_initial < fStopTol) { + LOG(warning) << "Initial kinetic energy is below stopping threshold, cannot propagate."; + return; // Cannot propagate if the initial kinetic energy is below the stopping threshold + } stepper.fDeriv = [this](const XYZPoint &pos, const XYZVector &mom) { return this->Derivatives(pos, mom); }; while (true) { - LOG(info) << "Position: " << GetPosition().X() / 10 << ", " << GetPosition().Y() / 10 << ", " - << GetPosition().Z() / 10; + LOG(debug) << "Position: " << GetPosition().X() / 10 << ", " << GetPosition().Y() / 10 << ", " + << GetPosition().Z() / 10; LOG(debug) << "Momentum: " << GetMomentum().X() << ", " << GetMomentum().Y() << ", " << GetMomentum().Z(); auto result = stepper.Step(fState); From 5e8471c2108f7355d8d52254cf2aef9c99fef56a Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Sat, 26 Jul 2025 22:35:05 +0200 Subject: [PATCH 46/71] Add full test track. Decomp failing --- .../OpenKF/kalman_filter/TrackFitterUKF.h | 21 ++- AtReconstruction/AtReconstructionLinkDef.h | 2 + macro/tests/UKF/UKFSingleTrack.C | 175 ++++++++++++++++++ 3 files changed, 193 insertions(+), 5 deletions(-) create mode 100644 macro/tests/UKF/UKFSingleTrack.C diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index a5227a2da..d8011f8d3 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -275,14 +275,21 @@ class TrackFitterUKFBase : public KalmanFilter { { const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_kappa)}; // sqrt(n + \kappa) - // cholesky factorization to get matrix Pxx square-root Eigen::LLT> lltOfPa(matPa); if (lltOfPa.info() != Eigen::Success) { - LOG(error) << "Cholesky decomposition failed, matrix is not positive definite."; - for (int32_t i{0}; i < STATE_DIM; ++i) { - LOG(error) << "Pxx[" << i << "]: " << matPa(i, i); + LOG(error) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; + // Add a small value to the diagonal to regularize the matrix + Matrix matPaReg = matPa; + matPaReg = (matPaReg + matPaReg.transpose()) * 0.5; // Ensure symmetry + matPaReg += Matrix::Identity() * 1e-6; // Regularization value + lltOfPa.compute(matPaReg); + if (lltOfPa.info() != Eigen::Success) { + for (int32_t i{0}; i < STATE_DIM; ++i) { + LOG(error) << "\n" << matPaReg; + } + throw std::runtime_error( + "Cholesky decomposition failed, matrix is not positive definite even after regularization."); } - throw std::runtime_error("Cholesky decomposition failed, matrix is not positive definite."); } Matrix matSa{lltOfPa.matrixL()}; // sqrt(P_{a}) @@ -412,6 +419,10 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { const TMatrixD &initialCovariance); void SetMeasCov(const TMatrixD &measCov); + std::array GetStateVector() const + { + return {m_vecX[0], m_vecX[1], m_vecX[2], m_vecX[3], m_vecX[4], m_vecX[5]}; + } kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z); diff --git a/AtReconstruction/AtReconstructionLinkDef.h b/AtReconstruction/AtReconstructionLinkDef.h index f9f287f94..7b1780c62 100755 --- a/AtReconstruction/AtReconstructionLinkDef.h +++ b/AtReconstruction/AtReconstructionLinkDef.h @@ -41,6 +41,8 @@ #pragma link C++ namespace kf; #pragma link C++ namespace kf::util; +#pragma link C++ class kf::TrackFitterUKFBase - !; +#pragma link C++ class kf::TrackFitterUKF - !; /* Classes that depend on Genfit2 */ #pragma link C++ class genfit::AtSpacepointMeasurement + ; diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C new file mode 100644 index 000000000..402643318 --- /dev/null +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -0,0 +1,175 @@ +std::string getEnergyPath() +{ + auto env = std::getenv("VMCWORKDIR"); + if (env == nullptr) { + return "../../resources/energy_loss/HinH.txt"; // Default path assuming cwd is build/AtTools + } + return std::string(env) + "/resources/energy_loss/HinH.txt"; // Use environment variable +} + +const double mass_p = 938.272; // Mass of proton in MeV/c^2 +const double charge_p = 1.602176634e-19; // Charge of proton + +// Simulated (measurement) hits +std::vector x, y, z, Eloss; + +void LoadHits() +{ + std::ifstream infile("hits.txt"); + double xi, yi, zi, Ei; + int i = 0; + double eLoss = 0; + + // Save first point. + eLoss = Ei * 1e3; // Initialize energy loss + x.push_back(xi * 10); // Convert to mm + y.push_back(yi * 10); // Convert to mm + z.push_back(zi * 10); // Convert to mm + + while (infile >> xi >> yi >> zi >> Ei) { + Ei *= 1e3; // Convert to MeV + + if (i++ % 5 != 0) { + eLoss += Ei; + continue; // Skip every 5th point + } + + double dx = x.back() - xi; + double dy = y.back() - yi; + double dz = z.back() - zi; + double distance = std::sqrt(dx * dx + dy * dy + dz * dz); + + x.push_back(xi * 10); + y.push_back(yi * 10); + z.push_back(zi * 10); + Eloss.push_back(eLoss); + eLoss = 0; // Reset energy loss for the next segment + } + + std::cout << "Finished loading hits. Total points: " << x.size() << std::endl; +} + +// This test should plot the trajectory of a particle in a magnetic field using +// the output from GEANT and the AtPropagator class. +void UKFSingleTrack() +{ + LoadHits(); // Load hits from file + + std::cout << " Creating the UKF class" << std::endl; + using namespace AtTools; + + std::vector x2, y2, z2, Eloss2; + + // Setup the Propagator for UKF + auto elossModel = std::make_unique(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_p, mass_p, std::move(elossModel)); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 2.85}); // Magnetic field + + // Setup stepper for UKF + auto stepper = std::make_unique(); + + // Setup UKF + kf::TrackFitterUKF ukf(std::move(propagator), std::move(stepper)); + + XYZPoint startPos(-3.40046e-05, -1.49863e-05, 0.10018); // Start position in cm + startPos *= 10; // Convert to mm + XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c + startMom *= 1e3; // Convert to MeV/c + + // Initial uncertainties + double sigma_pos = 5; // Position uncertainty of 10 mm + double sigma_mom = 0.01 * startMom.R(); // Momentum uncertainty of 10% MeV/c + double sigma_theta = 5 * M_PI / 180; // Angular uncertainty of 1 degree + double sigma_phi = 5 * M_PI / 180; // Angular uncertainty of 1 degree + + TMatrixD cov(6, 6); + cov.Zero(); + for (int i = 0; i < 3; ++i) { + + cov(i, i) = sigma_pos * sigma_pos; // Set diagonal covariance to some small number + } + cov(3, 3) = sigma_mom * sigma_mom; // Momentum uncertainty + cov(4, 4) = sigma_theta * sigma_theta; // Angular uncertainty + cov(5, 5) = sigma_phi * sigma_phi; // Angular uncertainty + + // Set the initial state + std::cout << "Setting initial state" << std::endl; + + ukf.SetInitialState(startPos, startMom, cov); + TMatrixD cov_meas(3, 3); + cov_meas.Zero(); + for (int i = 0; i < 3; ++i) { + cov_meas(i, i) = sigma_pos * sigma_pos; + } + // Create the covariance for measurement points. Assume constant + + x2.push_back(startPos.X()); + y2.push_back(startPos.Y()); + z2.push_back(startPos.Z()); + + ROOT::Math::XYZVector lastMom = ROOT::Math::XYZVector(startMom.X(), startMom.Y(), startMom.Z()); + + for (size_t i = 0; i < x.size(); ++i) { + std::cout << "Processing hit " << i << " of " << x.size() << std::endl; + XYZPoint point(x[i], y[i], z[i]); // measurement point in mm + ukf.SetMeasCov(cov_meas); // Set measurement noise covariance + + ukf.predictUKF(point); + ukf.correctUKF(point); + + auto state = ukf.GetStateVector(); + ROOT::Math::XYZPoint pos(state[0], state[1], state[2]); + ROOT::Math::Polar3DVector momPolar(state[3], state[4], state[5]); + ROOT::Math::XYZVector mom(momPolar); + + std::cout << "Predicted position: " << pos << std::endl; + std::cout << "Predicted momentum: " << mom << std::endl; + std::cout << "Measurement point: " << point << std::endl; + + auto KE_in = Kinematics::KE(lastMom, mass_p); + auto KE_out = Kinematics::KE(mom, mass_p); + lastMom = mom; + + x2.push_back(pos.X()); + y2.push_back(pos.Y()); + z2.push_back(pos.Z()); + Eloss2.push_back((KE_in - KE_out)); + } + + TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); + track->SetTitle("Particle Track;X [mm];Y [mm];Z [mm]"); + track->SetMarkerStyle(20); + track->SetMarkerSize(0.8); + + TGraph2D *track2 = new TGraph2D(x2.size(), x2.data(), y2.data(), z2.data()); + track2->SetTitle("Propagated Particle Track;X [mm];Y [mm];Z [mm]"); + track2->SetMarkerStyle(21); + track2->SetMarkerSize(0.8); + track2->SetMarkerColor(kRed); + + TCanvas *c1 = new TCanvas("c1", "Particle Track", 800, 600); + track->Draw("P"); + track2->Draw("PSAME"); + + TGraph *elossGraph = new TGraph(Eloss.size()); + for (size_t i = 0; i < Eloss.size(); ++i) { + elossGraph->SetPoint(i, i, Eloss[i]); + } + elossGraph->SetTitle("Energy Loss per Hit;Hit Number;Energy Loss [MeV]"); + elossGraph->SetMarkerStyle(20); + + TGraph *eloss2Graph = new TGraph(Eloss2.size()); + for (size_t i = 0; i < Eloss2.size(); ++i) { + eloss2Graph->SetPoint(i, i, Eloss2[i]); + } + eloss2Graph->SetTitle("Propagated Energy Loss per Hit;Hit Number;Energy Loss [MeV]"); + eloss2Graph->SetMarkerStyle(21); + eloss2Graph->SetMarkerColor(kRed); + + TCanvas *c2 = new TCanvas("c2", "Energy Loss per Hit", 800, 600); + elossGraph->Draw("AP"); + // eloss2Graph->Draw("PSAME"); +} \ No newline at end of file From 479bb8471957be8423681cf98887be4dfd6ddb5a Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Sat, 26 Jul 2025 23:01:07 +0200 Subject: [PATCH 47/71] Fix decomposition, but UKF seems suspicious. It seems to be propagating initially in the wrong direction. It gets "dragged" back to following the points, but that seems to blow up the momentum uncertainty (which makes sense if the initial momentum is somehow in the wrong direction, and the measurement points need to force it to change its mind). I think that causes it to fail somewhere around point 23, becuase of sigma points are sampling such a large distribution due to the momentum uncertainty that they start stopping in the gas (or have a negative momentum). Checks for physicality need to be added (step size, momentum values, etc.) --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h | 9 ++++++--- macro/tests/UKF/UKFSingleTrack.C | 13 ++++++++++--- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index d8011f8d3..7714bc4dc 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -284,9 +284,12 @@ class TrackFitterUKFBase : public KalmanFilter { matPaReg += Matrix::Identity() * 1e-6; // Regularization value lltOfPa.compute(matPaReg); if (lltOfPa.info() != Eigen::Success) { - for (int32_t i{0}; i < STATE_DIM; ++i) { - LOG(error) << "\n" << matPaReg; - } + LOG(error) << "Cholesky decomposition failed even after regularization. Attempting again"; + matPaReg += Matrix::Identity() * 1e-3; // Increase regularization value + lltOfPa.compute(matPaReg); + } + if (lltOfPa.info() != Eigen::Success) { + LOG(error) << "\n" << matPaReg; throw std::runtime_error( "Cholesky decomposition failed, matrix is not positive definite even after regularization."); } diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 402643318..faa7071b8 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -21,6 +21,7 @@ void LoadHits() double eLoss = 0; // Save first point. + infile >> xi >> yi >> zi >> Ei; eLoss = Ei * 1e3; // Initialize energy loss x.push_back(xi * 10); // Convert to mm y.push_back(yi * 10); // Convert to mm @@ -47,6 +48,7 @@ void LoadHits() } std::cout << "Finished loading hits. Total points: " << x.size() << std::endl; + std::cout << x[0] << " " << x[1] << " " << x[2] << std::endl; } // This test should plot the trajectory of a particle in a magnetic field using @@ -82,8 +84,8 @@ void UKFSingleTrack() // Initial uncertainties double sigma_pos = 5; // Position uncertainty of 10 mm double sigma_mom = 0.01 * startMom.R(); // Momentum uncertainty of 10% MeV/c - double sigma_theta = 5 * M_PI / 180; // Angular uncertainty of 1 degree - double sigma_phi = 5 * M_PI / 180; // Angular uncertainty of 1 degree + double sigma_theta = 1 * M_PI / 180; // Angular uncertainty of 1 degree + double sigma_phi = 1 * M_PI / 180; // Angular uncertainty of 1 degree TMatrixD cov(6, 6); cov.Zero(); @@ -112,13 +114,17 @@ void UKFSingleTrack() ROOT::Math::XYZVector lastMom = ROOT::Math::XYZVector(startMom.X(), startMom.Y(), startMom.Z()); - for (size_t i = 0; i < x.size(); ++i) { + // Skip the first point since it is the initial state. + // Stop when things break. + for (size_t i = 1; i < 21; ++i) { std::cout << "Processing hit " << i << " of " << x.size() << std::endl; XYZPoint point(x[i], y[i], z[i]); // measurement point in mm ukf.SetMeasCov(cov_meas); // Set measurement noise covariance ukf.predictUKF(point); + std::cout << std::endl << "Prediction step complete." << std::endl; ukf.correctUKF(point); + std::cout << std::endl << "Correction step complete." << std::endl; auto state = ukf.GetStateVector(); ROOT::Math::XYZPoint pos(state[0], state[1], state[2]); @@ -127,6 +133,7 @@ void UKFSingleTrack() std::cout << "Predicted position: " << pos << std::endl; std::cout << "Predicted momentum: " << mom << std::endl; + std::cout << "Measurement point: " << point << std::endl; auto KE_in = Kinematics::KE(lastMom, mass_p); From ab3b6425af9a0311945d7df21bd9f06a9c0ff579 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Sat, 26 Jul 2025 23:23:04 +0200 Subject: [PATCH 48/71] Fix issue with initial COV being too high. Filter not follows measurement points. --- AtTools/AtPropagator.cxx | 10 +++++----- macro/tests/UKF/UKFSingleTrack.C | 11 +++++++---- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index 8b458ed6d..d3ee01989 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -75,7 +75,7 @@ void AtPropagator::PropagateOneStep(AtStepper &stepper) void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) { - LOG(info) << "Propagating to measurement surface"; + LOG(info) << "Propagating to measurement surface from: " << fState.fPos; fState.h = stepper.GetInitialStep(); // Set the initial step size auto KE_initial = Kinematics::KE(fState.fMom, fState.fMass); @@ -109,9 +109,9 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur LOG(info) << "Distance to plane: " << approach << " mm"; LOG(info) << "Final step size: " << finalH << " mm"; - LOG(info) << "Current position: " << fState.fLastPos.X() << ", " << fState.fLastPos.Y() << ", " + LOG(info) << "Position before surface: " << fState.fLastPos.X() << ", " << fState.fLastPos.Y() << ", " << fState.fLastPos.Z(); - LOG(info) << "Current momentum: " << fState.fLastMom.X() << ", " << fState.fLastMom.Y() << ", " + LOG(info) << "Momentum before surface: " << fState.fLastMom.X() << ", " << fState.fLastMom.Y() << ", " << fState.fLastMom.Z(); finalH = approach * 1e-3; // Convert to meters for the RK4 step @@ -187,16 +187,16 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur double distanceToSurface = surface.Distance(fState.fPos); double KE_final = Kinematics::KE(fState.fMom, fState.fMass); - LOG(info) << "Initial KE" << KE_initial << " MeV"; + LOG(info) << "Initial KE: " << KE_initial << " MeV"; LOG(info) << "Final KE: " << KE_final << " MeV"; auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - LOG(info) << "------- End of RK4 interation ---------"; LOG(info) << "Particle stopped: " << particleStopped; LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; LOG(info) << "Distance to surface: " << distanceToSurface << " mm"; LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; LOG(info) << "Scaling factor: " << fScalingFactor; LOG(info) << "Final Position: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); + LOG(info) << "------- End of RK4 interation ---------" << std::endl; // If we reached the measurement surface, we should project the position onto the surface if (reachedMeasurementPoint) { diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index faa7071b8..71d37742e 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -79,10 +79,13 @@ void UKFSingleTrack() XYZPoint startPos(-3.40046e-05, -1.49863e-05, 0.10018); // Start position in cm startPos *= 10; // Convert to mm XYZVector startMom(0.00935463, -0.0454279, 0.00826042); // Start momentum in GeV/c - startMom *= 1e3; // Convert to MeV/c + startMom *= 1e3; + + XYZPoint nextPos(x[1], y[1], z[1]); + startMom = startMom.R() * (nextPos - startPos).Unit(); // Set momentum direction towards the first hit // Initial uncertainties - double sigma_pos = 5; // Position uncertainty of 10 mm + double sigma_pos = 1; // Position uncertainty of 10 mm double sigma_mom = 0.01 * startMom.R(); // Momentum uncertainty of 10% MeV/c double sigma_theta = 1 * M_PI / 180; // Angular uncertainty of 1 degree double sigma_phi = 1 * M_PI / 180; // Angular uncertainty of 1 degree @@ -115,8 +118,8 @@ void UKFSingleTrack() ROOT::Math::XYZVector lastMom = ROOT::Math::XYZVector(startMom.X(), startMom.Y(), startMom.Z()); // Skip the first point since it is the initial state. - // Stop when things break. - for (size_t i = 1; i < 21; ++i) { + // Stop when things break (point 21). + for (size_t i = 1; i < x.size() && i < 100; ++i) { std::cout << "Processing hit " << i << " of " << x.size() << std::endl; XYZPoint point(x[i], y[i], z[i]); // measurement point in mm ukf.SetMeasCov(cov_meas); // Set measurement noise covariance From 3c516c265574aaf74838325270588f84b4b1a903 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 28 Jul 2025 23:49:01 +0200 Subject: [PATCH 49/71] Move most propagator logging to debug --- AtTools/AtPropagator.cxx | 71 ++++++++++++++++++++-------------------- 1 file changed, 36 insertions(+), 35 deletions(-) diff --git a/AtTools/AtPropagator.cxx b/AtTools/AtPropagator.cxx index d3ee01989..5838ac706 100644 --- a/AtTools/AtPropagator.cxx +++ b/AtTools/AtPropagator.cxx @@ -75,7 +75,7 @@ void AtPropagator::PropagateOneStep(AtStepper &stepper) void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, AtStepper &stepper) { - LOG(info) << "Propagating to measurement surface from: " << fState.fPos; + LOG(debug) << "Propagating to measurement surface from: " << fState.fPos; fState.h = stepper.GetInitialStep(); // Set the initial step size auto KE_initial = Kinematics::KE(fState.fMom, fState.fMass); @@ -103,16 +103,16 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur if (reachedMeasurementPoint && !particleStopped && !momentumReversed) { // We reached the measurement surface, so we should figure out how far we are from the measurement point - LOG(info) << "------ Reached measurement surface ------"; + LOG(debug) << "------ Reached measurement surface ------"; double finalH = (fState.fLastPos - fState.fPos).R(); // Distance traveled in the last step double approach = surface.Distance(fState.fLastPos); - LOG(info) << "Distance to plane: " << approach << " mm"; - LOG(info) << "Final step size: " << finalH << " mm"; - LOG(info) << "Position before surface: " << fState.fLastPos.X() << ", " << fState.fLastPos.Y() << ", " - << fState.fLastPos.Z(); - LOG(info) << "Momentum before surface: " << fState.fLastMom.X() << ", " << fState.fLastMom.Y() << ", " - << fState.fLastMom.Z(); + LOG(debug) << "Distance to plane: " << approach << " mm"; + LOG(debug) << "Final step size: " << finalH << " mm"; + LOG(debug) << "Position before surface: " << fState.fLastPos.X() << ", " << fState.fLastPos.Y() << ", " + << fState.fLastPos.Z(); + LOG(debug) << "Momentum before surface: " << fState.fLastMom.X() << ", " << fState.fLastMom.Y() << ", " + << fState.fLastMom.Z(); finalH = approach * 1e-3; // Convert to meters for the RK4 step fState.h = finalH; // Set the step size to the distance to the surface @@ -139,10 +139,10 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur double deltaE = KE_last - fStopTol; deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss - LOG(info) << "Last KE: " << KE_last << " MeV"; - LOG(info) << "Energy to loose to stop: " << deltaE << " MeV"; + LOG(debug) << "Last KE: " << KE_last << " MeV"; + LOG(debug) << "Energy to loose to stop: " << deltaE << " MeV"; double h_Stop = deltaE / fELossModel->GetdEdx(KE_last); // Distance to stop in mm - LOG(info) << "Estimated distance to stop: " << h_Stop << " mm"; + LOG(debug) << "Estimated distance to stop: " << h_Stop << " mm"; fState.h = h_Stop * 1e-3; // Convert to meters for the RK4 step fState.fPos = fState.fLastPos; // Set position to last position @@ -155,9 +155,9 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur auto origH = fState.h; // Save original step size fState = result; // Update the internal state fState.h = origH; // Restore original step size - LOG(info) << "Propagated to stopping point: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " - << fState.fPos.Z(); - LOG(info) << "Energy after stopping: " << Kinematics::KE(fState.fMom, fState.fMass) << " MeV"; + LOG(debug) << "Propagated to stopping point: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " + << fState.fPos.Z(); + LOG(debug) << "Energy after stopping: " << Kinematics::KE(fState.fMom, fState.fMass) << " MeV"; while (surface.fClipToSurface) { fScalingFactor = 0; // Turn off energy loss. @@ -168,7 +168,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur reachedMeasurementPoint = true; break; } - LOG(info) << "Propagating to surface after stopping with step size: " << h << " mm"; + LOG(debug) << "Propagating to surface after stopping with step size: " << h << " mm"; fState.h = h * 1e-3; // Convert to meters for the RK4 step result = stepper.Step(fState); if (!result) { @@ -176,8 +176,8 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur return; // Abort propagation if step failed } fState = result; // Update the internal state - LOG(info) << "New position after adjusting step size: " << fState.fPos.X() << ", " << fState.fPos.Y() - << ", " << fState.fPos.Z(); + LOG(debug) << "New position after adjusting step size: " << fState.fPos.X() << ", " << fState.fPos.Y() + << ", " << fState.fPos.Z(); } fState.fLastMom = fState.fMom; fState.fMom = XYZVector(0, 0, 0); // Set momentum to zero since we stopped @@ -187,23 +187,24 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur double distanceToSurface = surface.Distance(fState.fPos); double KE_final = Kinematics::KE(fState.fMom, fState.fMass); - LOG(info) << "Initial KE: " << KE_initial << " MeV"; - LOG(info) << "Final KE: " << KE_final << " MeV"; + LOG(debug) << "Initial KE: " << KE_initial << " MeV"; + LOG(debug) << "Final KE: " << KE_final << " MeV"; auto calc_eLoss = KE_initial - KE_final; // Energy loss in MeV - LOG(info) << "Particle stopped: " << particleStopped; - LOG(info) << "Reached measurement point: " << reachedMeasurementPoint; - LOG(info) << "Distance to surface: " << distanceToSurface << " mm"; - LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; - LOG(info) << "Scaling factor: " << fScalingFactor; - LOG(info) << "Final Position: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); - LOG(info) << "------- End of RK4 interation ---------" << std::endl; + LOG(debug) << "Particle stopped: " << particleStopped; + LOG(debug) << "Reached measurement point: " << reachedMeasurementPoint; + LOG(debug) << "Distance to surface: " << distanceToSurface << " mm"; + LOG(debug) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(debug) << "Scaling factor: " << fScalingFactor; + LOG(debug) << "Final Position: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); + LOG(debug) << "------- End of RK4 interation ---------" << std::endl; // If we reached the measurement surface, we should project the position onto the surface if (reachedMeasurementPoint) { fState.fPos = surface.ProjectToSurface(fState.fPos); } - LOG(info) << "Projected on surface: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " << fState.fPos.Z(); - LOG(info) << "Final Momentum: " << fState.fMom.X() << ", " << fState.fMom.Y() << ", " << fState.fMom.Z(); + LOG(debug) << "Projected on surface: " << fState.fPos.X() << ", " << fState.fPos.Y() << ", " + << fState.fPos.Z(); + LOG(debug) << "Final Momentum: " << fState.fMom.X() << ", " << fState.fMom.Y() << ", " << fState.fMom.Z(); return; } } // End of loop over RK4 integration @@ -211,7 +212,7 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, double eLoss, AtStepper &stepper) { - LOG(info) << "Propagating to surface with eLoss: " << eLoss; + LOG(debug) << "Propagating to surface with eLoss: " << eLoss; if (eLoss == 0) { LOG(warn) << "No energy loss specified, propagating without energy loss adjustment."; @@ -243,15 +244,15 @@ void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &sur double KE_final = Kinematics::KE(fState.fMom, fState.fMass); calc_eLoss = KE_initial - KE_final; // Energy loss in MeV fScalingFactor *= eLoss / calc_eLoss; - LOG(info) << "Desired energy loss: " << eLoss << " MeV"; - LOG(info) << "Calculated energy loss: " << calc_eLoss << " MeV"; - LOG(info) << "Difference: " << calc_eLoss - eLoss << " MeV"; - LOG(info) << "New scaling factor: " << fScalingFactor; - LOG(info) << "Condition: " << (std::abs(calc_eLoss - eLoss) > 1e-4); + LOG(debug) << "Desired energy loss: " << eLoss << " MeV"; + LOG(debug) << "Calculated energy loss: " << calc_eLoss << " MeV"; + LOG(debug) << "Difference: " << calc_eLoss - eLoss << " MeV"; + LOG(debug) << "New scaling factor: " << fScalingFactor; + LOG(debug) << "Condition: " << (std::abs(calc_eLoss - eLoss) > 1e-4); } // End loop over energy loss convergence - LOG(info) << "Energy loss converged after " << iterations << " iterations."; + LOG(debug) << "Energy loss converged after " << iterations << " iterations."; fScalingFactor = 1; // Reset scaling factor after convergence } From 2b070b938e4cdcc3c2df7d129e8087dbf3294a4e Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Mon, 28 Jul 2025 23:49:15 +0200 Subject: [PATCH 50/71] Fix bug in straggling. Fix issue with stragling being calculated over last timestep rather then the whole propagation from previous point to next measurement POCA. Add additional functions to UKF to disable straggling. Add functions to pull intermediate results from UKF. Add some additional CATIMA straggling functions (pull from calculation, not just range tables that are pre-calculated). --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 44 +++++++++++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 22 +++++- AtTools/AtELossCATIMA.cxx | 21 ++++++ AtTools/AtELossCATIMA.h | 1 + AtTools/AtELossCATIMATest.cxx | 17 +++++ CMakeLists.txt | 7 +- macro/tests/UKF/UKFSingleTrack.C | 73 +++++++++++++++++-- 7 files changed, 174 insertions(+), 11 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 8c3824eca..1c8012494 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -28,6 +28,33 @@ void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition } } +TMatrixD TrackFitterUKF::GetStateCovariance() const +{ + TMatrixD cov(m_matP.rows(), m_matP.cols()); + for (int i = 0; i < m_matP.rows(); ++i) { + for (int j = 0; j < m_matP.cols(); ++j) { + cov(i, j) = m_matP(i, j); // Copy covariance matrix to TMatrixD + } + } + return cov; +} + +std::array TrackFitterUKF::GetAugStateVector() const +{ + return {m_vecXa[0], m_vecXa[1], m_vecXa[2], m_vecXa[3], + m_vecXa[4], m_vecXa[5], m_vecXa[6]}; // Return the augmented state vector as an array +} +TMatrixD TrackFitterUKF::GetAugStateCovariance() const +{ + TMatrixD cov(m_matPa.rows(), m_matPa.cols()); + for (int i = 0; i < m_matPa.rows(); ++i) { + for (int j = 0; j < m_matPa.cols(); ++j) { + cov(i, j) = m_matPa(i, j); // Copy augmented covariance matrix to TMatrixD + } + } + return cov; +} + void TrackFitterUKF::SetMeasCov(const TMatrixD &measCov) { if (measCov.GetNrows() != TF_DIM_Z || measCov.GetNcols() != TF_DIM_Z) { @@ -58,10 +85,27 @@ Matrix TrackFitterUKF::calcu double eIn = AtTools::Kinematics::KE(fMeanStep.fLastMom, fMeanStep.fMass); double eOut = AtTools::Kinematics::KE(fMeanStep.fMom, fMeanStep.fMass); + if (!fEnableEnStraggling) { + // If energy straggling is disabled, we set the process noise to zero. + matQ(0, 0) = 0.0; + return matQ; + } + if (const auto *elossModel = fPropagator.GetELossModel()) { double dedx_straggle = elossModel->GetdEdxStraggling(eIn, eOut); double factor = dedx_straggle / elossModel->GetdEdx(eIn); + if (factor > fMaxStragglingFactor) { + LOG(warn) << "Process noise factor for energy straggling is greater than " << fMaxStragglingFactor + << ". To maintain stability, we will " + "use a factor of " + << fMaxStragglingFactor << "."; + factor = fMaxStragglingFactor; + } matQ(0, 0) = factor * factor; // Variance for the dedx straggling. + LOG(info) << "Calculating process noise for straggling between " << eIn << " MeV and " << eOut << " MeV over " + << elossModel->GetRange(eIn, eOut) << " mm."; + LOG(info) << "Process noise covariance for energy straggling: " << matQ(0, 0) << " (factor: " << factor + << ", dedx_straggle: " << dedx_straggle << ", dEdx: " << elossModel->GetdEdx(eIn) << ")"; } else { throw std::runtime_error("Cannot calculate process noise covariance without an energy loss model"); diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 7714bc4dc..1e38f505e 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -403,6 +403,9 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter public: + bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling + double fMaxStragglingFactor{1. / 3.}; ///< @brief Maximum straggling factor for energy loss + /** * @brief Constructor for the TrackFitterUKF class. * @param propagator The propagator to be used for the track fitting, must be passed as an rvalue reference. @@ -426,6 +429,9 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { { return {m_vecX[0], m_vecX[1], m_vecX[2], m_vecX[3], m_vecX[4], m_vecX[5]}; } + TMatrixD GetStateCovariance() const; + TMatrixD GetAugStateCovariance() const; + std::array GetAugStateVector() const; kf::Vector funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z); @@ -433,13 +439,25 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { void predictUKF(const ROOT::Math::XYZPoint &z) { + using namespace ROOT::Math; // First we need to propagate the mean state vector to the next measurement point. + XYZPoint startingPosition{m_vecX[0], m_vecX[1], m_vecX[2]}; // Get the starting position from the state vector + Polar3DVector startingMomentum{m_vecX[3], m_vecX[4], + m_vecX[5]}; // Get the starting momentum from the state vector + + LOG(info) << "Propagating reference state from position: " << startingPosition + << " with momentum: " << XYZVector(startingMomentum); + + fPropagator.SetState(startingPosition, XYZVector(startingMomentum)); fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(z), *fStepper); - fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator + fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator + fMeanStep.fLastPos = startingPosition; // Store the last position + fMeanStep.fLastMom = startingMomentum; // Store the last momentum + + LOG(info) << "Propagated to position: " << fMeanStep.fPos << " with momentum: " << fMeanStep.fMom; // Now we can construct the reference plane. - using namespace ROOT::Math; fMeasurementPlane = Plane3D(fMeanStep.fMom.Unit(), XYZPoint(z)); // Create a plane using the momentum direction and position Vector zVec; // Initialize the measurement vector diff --git a/AtTools/AtELossCATIMA.cxx b/AtTools/AtELossCATIMA.cxx index 3f85a981d..2518d8402 100644 --- a/AtTools/AtELossCATIMA.cxx +++ b/AtTools/AtELossCATIMA.cxx @@ -129,6 +129,27 @@ double AtELossCATIMA::GetdEdxStraggling(double energyIni, double energyFin) cons auto factor = dE_st / (energyIni - energyFin); return factor * dedx_min; } +double AtELossCATIMA::GetdEdxStragglingCATIMA(double energy, double intDistance) const +{ + if (fProjectile == nullptr || fMaterial == nullptr) { + LOG(error) << "Projectile or material not set. dEdx straggling is 0."; + return 0; + } + catima::Result result = catima::calculate(*fProjectile, *fMaterial, energy / fProjectileMassAmu); + double dEdxi = result.dEdxi; // MeV/(g/cm^2) + + intDistance *= 10 * fDensity; // Convert to g/cm^2. + + auto oldT = fProjectile->T; + fProjectile->T = energy / fProjectileMassAmu; // Set the projectile's + auto dE_st = catima::domega2dx(*fProjectile, *fMaterial) / dEdxi; // domega in MeV/(g/cm^2) + fProjectile->T = oldT; // Restore the projectile's T + + dE_st /= intDistance; // Get the variance in stopping power (MeV/(g/cm^2)) + dE_st *= fDensity * fDensity; // Convert to MeV^2/cm^2 + dE_st *= 100; // Convert to MeV^2/mm^2 + return std::sqrt(dE_st); // Returns the factor for dEdx straggling +} std::vector> AtELossCATIMA::GetBraggCurve(double energy, double rangeStepSize, double totalFractionELoss) const diff --git a/AtTools/AtELossCATIMA.h b/AtTools/AtELossCATIMA.h index bf63862e7..33f4c61c7 100644 --- a/AtTools/AtELossCATIMA.h +++ b/AtTools/AtELossCATIMA.h @@ -45,6 +45,7 @@ class AtELossCATIMA : public AtELossModel { virtual double GetRangeVariance(double energy) const override; virtual double GetElossStraggling(double energyIni, double energyFin) const override; virtual double GetdEdxStraggling(double energyIni, double energyFin) const override; + double GetdEdxStragglingCATIMA(double energy, double distance) const; virtual std::vector> GetBraggCurve(double energy, double rangeStepSize = 0, double totalFractionELoss = 0.001) const override; diff --git a/AtTools/AtELossCATIMATest.cxx b/AtTools/AtELossCATIMATest.cxx index 9a438ae60..d7d89dfae 100644 --- a/AtTools/AtELossCATIMATest.cxx +++ b/AtTools/AtELossCATIMATest.cxx @@ -70,6 +70,7 @@ TEST_F(AtELossCATIMATestFixture, TestEnergyLossStragglingDistance) { double expectedSigma = 0.0084 * mass; // Expected sigma from LISE double eloss_straggling = model.GetElossStragglingDistance(1.0, 50.0); // 1 MeV over 10 mm + LOG(info) << "Straggling over 5 cm: " << eloss_straggling; ASSERT_NEAR(eloss_straggling, expectedSigma, 0.1 * expectedSigma); expectedSigma = 0.0376 * mass; @@ -94,6 +95,22 @@ TEST_F(AtELossCATIMATestFixture, TestdEdxStraggling) ASSERT_NEAR(e_min, e_min_dedx, 0.01 * e_min); // Check minimum energy loss } +TEST_F(AtELossCATIMATestFixture, TestdEdxStragglingCalc) +{ + double intDistance = 10; // 1 mm integration distance. + auto Eout = model.GetEnergy(5.0, intDistance); // Initialize the model with a known energy + double dedx_straggling = model.GetdEdxStraggling(5.0, Eout); // Get over small range + double dedx_st_calc = model.GetdEdxStragglingCATIMA(5.0, intDistance); + + LOG(info) << "Energy after " << intDistance << " mm: " << Eout << " MeV (" << Eout / mass << " MeV/u)"; + LOG(info) << "Energy loss: " << 5.0 - Eout << " MeV"; + LOG(info) << "Straggling factor: " << dedx_straggling * intDistance / (5.0 - Eout); + LOG(info) << "Straggling factor: " << dedx_straggling / model.GetdEdx(5.0); + LOG(info) << "Energy straggling calc: " << dedx_st_calc * intDistance << " MeV"; + LOG(info) << "Energy straggling range: " << dedx_straggling * intDistance << " MeV"; + ASSERT_NEAR(dedx_straggling, dedx_st_calc, 1e-4); // Check dEdx straggling +} + TEST(AtELossCATIMATest, LISE_Match) { // Create the vector with the gas components for H2. diff --git a/CMakeLists.txt b/CMakeLists.txt index 55a74e19e..1ae91b0e3 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -138,8 +138,11 @@ if(NOT CATIMA_FOUND) set(APPS OFF CACHE BOOL "Disable CATIMA applications" FORCE) set(GENERATE_DATA OFF CACHE BOOL "Disable CATIMA data generator" FORCE) - FetchContent_MakeAvailable(catima) - +FetchContent_GetProperties(catima) + if (NOT catima_POPULATED) + FetchContent_Populate(catima) + add_subdirectory(${catima_SOURCE_DIR} ${catima_BINARY_DIR}) + endif() # Suppress shadow warnings for CATIMA target if(TARGET catima) get_target_property(catima_compile_options catima COMPILE_OPTIONS) diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 71d37742e..a6a82cd5a 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -12,7 +12,7 @@ const double charge_p = 1.602176634e-19; // Charge of proton // Simulated (measurement) hits std::vector x, y, z, Eloss; - +int pointsToCluster = 5; void LoadHits() { std::ifstream infile("hits.txt"); @@ -22,15 +22,15 @@ void LoadHits() // Save first point. infile >> xi >> yi >> zi >> Ei; - eLoss = Ei * 1e3; // Initialize energy loss + eLoss = Ei; // Initialize energy loss x.push_back(xi * 10); // Convert to mm y.push_back(yi * 10); // Convert to mm z.push_back(zi * 10); // Convert to mm while (infile >> xi >> yi >> zi >> Ei) { - Ei *= 1e3; // Convert to MeV + // Ei *= 1e3; // Convert to MeV - if (i++ % 5 != 0) { + if (++i % pointsToCluster != 0) { eLoss += Ei; continue; // Skip every 5th point } @@ -60,13 +60,20 @@ void UKFSingleTrack() std::cout << " Creating the UKF class" << std::endl; using namespace AtTools; - std::vector x2, y2, z2, Eloss2; + std::vector x2, y2, z2, Eloss2, p2, sigmap2, lambda2, sigmalambda2, residual; // Setup the Propagator for UKF auto elossModel = std::make_unique(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_p, mass_p, std::move(elossModel)); + + auto elossModel2 = std::make_unique(3.553e-5); + elossModel2->SetProjectile(1, 1, 1); + std::vector> mat; + mat.push_back({1, 1, 1}); + elossModel2->SetMaterial(mat); + + AtTools::AtPropagator propagator(charge_p, mass_p, std::move(elossModel2)); propagator.SetEField({0, 0, 0}); // No electric field propagator.SetBField({0, 0, 2.85}); // Magnetic field @@ -89,6 +96,7 @@ void UKFSingleTrack() double sigma_mom = 0.01 * startMom.R(); // Momentum uncertainty of 10% MeV/c 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 TMatrixD cov(6, 6); cov.Zero(); @@ -114,6 +122,9 @@ void UKFSingleTrack() x2.push_back(startPos.X()); y2.push_back(startPos.Y()); z2.push_back(startPos.Z()); + p2.push_back(startMom.R()); + sigmap2.push_back(sigma_mom); + residual.push_back(0); // Initial residual is zero ROOT::Math::XYZVector lastMom = ROOT::Math::XYZVector(startMom.X(), startMom.Y(), startMom.Z()); @@ -125,11 +136,15 @@ void UKFSingleTrack() ukf.SetMeasCov(cov_meas); // Set measurement noise covariance ukf.predictUKF(point); + auto augState = ukf.GetAugStateVector(); + auto augCov = ukf.GetAugStateCovariance(); std::cout << std::endl << "Prediction step complete." << std::endl; ukf.correctUKF(point); std::cout << std::endl << "Correction step complete." << std::endl; auto state = ukf.GetStateVector(); + auto cov = ukf.GetStateCovariance(); + ROOT::Math::XYZPoint pos(state[0], state[1], state[2]); ROOT::Math::Polar3DVector momPolar(state[3], state[4], state[5]); ROOT::Math::XYZVector mom(momPolar); @@ -143,10 +158,17 @@ void UKFSingleTrack() auto KE_out = Kinematics::KE(mom, mass_p); lastMom = mom; + double residualValue = (point - pos).R(); + x2.push_back(pos.X()); y2.push_back(pos.Y()); z2.push_back(pos.Z()); Eloss2.push_back((KE_in - KE_out)); + p2.push_back(mom.R()); + sigmap2.push_back(std::sqrt(cov(3, 3))); // Propagate momentum uncertainty + lambda2.push_back(augState[6]); // Energy straggling factor + sigmalambda2.push_back(std::sqrt(augCov(6, 6))); // Propagate energy straggling uncertainty + residual.push_back(residualValue); // Store the residual for this hit } TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); @@ -179,7 +201,44 @@ void UKFSingleTrack() eloss2Graph->SetMarkerStyle(21); eloss2Graph->SetMarkerColor(kRed); + TGraphErrors *pGraph = new TGraphErrors(p2.size()); + for (size_t i = 0; i < p2.size(); ++i) { + pGraph->SetPoint(i, i, p2[i] * 1e-3 * pointsToCluster / 5.); + pGraph->SetPointError(i, 0, + sigmap2[i] * 5 * 1e-3 * pointsToCluster / + 5.); // Error bars from sigmap2, converted to GeV/c + } + pGraph->SetTitle("Momentum per Hit;Hit Number;Momentum [MeV/c]"); + pGraph->SetMarkerStyle(20); + pGraph->SetMarkerColor(kBlue); + + TGraphErrors *lambdaGraph = new TGraphErrors(lambda2.size()); + for (size_t i = 0; i < lambda2.size(); ++i) { + lambdaGraph->SetPoint(i, i, lambda2[i] * 0.03 * pointsToCluster / 5.); + lambdaGraph->SetPointError(i, 0, sigmalambda2[i] * 0.03 * pointsToCluster / 5.); + std::cout << "Lambda: " << lambda2[i] << ", Error: " << sigmalambda2[i] << std::endl; + } + lambdaGraph->SetTitle("Lambda per Hit (scaled);Hit Number;Lambda [scaled]"); + lambdaGraph->SetMarkerStyle(22); + lambdaGraph->SetMarkerColor(kGreen + 2); + + TGraph *residualGraph = new TGraph(residual.size()); + for (size_t i = 0; i < residual.size(); ++i) { + residualGraph->SetPoint(i, i, residual[i] * .1); + } + residualGraph->SetTitle("Residual per Hit;Hit Number;Residual [mm]"); + residualGraph->SetMarkerStyle(23); + residualGraph->SetMarkerColor(kMagenta); + TCanvas *c2 = new TCanvas("c2", "Energy Loss per Hit", 800, 600); elossGraph->Draw("AP"); - // eloss2Graph->Draw("PSAME"); + eloss2Graph->Draw("PSAME"); + pGraph->Draw("PSAME"); + // lambdaGraph->Draw("PSAME"); + residualGraph->Draw("PSAME"); + + double sumEloss = std::accumulate(Eloss.begin(), Eloss.end(), 0.0); + double sumEloss2 = std::accumulate(Eloss2.begin(), Eloss2.end(), 0.0); + std::cout << "Sum of Eloss: " << sumEloss << std::endl; + std::cout << "Sum of Eloss2: " << sumEloss2 << std::endl; } \ No newline at end of file From 796d89365b8e42faf78e181a34748f14b82ab740 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 00:33:41 +0200 Subject: [PATCH 51/71] Cleanup UKF doc and headers --- .../OpenKF/kalman_filter/TrackFitterUKF.h | 305 +++++++++--------- .../OpenKF/kalman_filter/kalman_filter.h | 13 +- 2 files changed, 156 insertions(+), 162 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 1e38f505e..f525661ec 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -26,49 +26,54 @@ #include "unscented_kalman_filter.h" namespace kf { -/// @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. -/// @tparam DIM_X Dimension of the state vector. -/// @tparam DIM_Z Dimension of the measurement vector. -/// @tparam DIM_V Dimension of the process noise vector. -/// @tparam DIM_N Dimension of the measurement noise vector. +/** @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. + * + * Serves as a templated base class for UKF calculations. It hold no physics information, just + * the machinery that underlies the UKF formalism. It is a modified version of the UKF provided + * by OpenKF, that has been expanded to allow for more hooks into the method. + * + * Templated because I believe Eigen can do quite a bit of operation for small matrices like we have here if + * the size is known at compile time. Worth checking that though. + * + * @tparam DIM_X Dimension of the state vector. + * @tparam DIM_Z Dimension of the measurement vector. + * @tparam DIM_V Dimension of the process noise vector. + * @tparam DIM_N Dimension of the measurement noise vector. + */ template -class TrackFitterUKFBase : public KalmanFilter { - +class TrackFitterUKFBase { public: // Augmented state vector is just the process noise and state vector. The measurement noise is not included as that // is independent of the propagation and measurement model and just adds linearly. static constexpr int32_t DIM_A{DIM_X + DIM_V}; ///< @brief Augmented state dimension static constexpr int32_t SIGMA_DIM_A{2 * DIM_A + 1}; ///< @brief Sigma points dimension for augmented state + Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix protected: - using KalmanFilter::m_vecX; // from Base KalmanFilter class - using KalmanFilter::m_matP; // from Base KalmanFilter class - - float32_t m_weight0; /// @brief unscented transform weight 0 for mean - float32_t m_weighti; /// @brief unscented transform weight i for none mean samples - float32_t m_kappa{0}; ///< @brief Kappa parameter for finding sigma points - - Vector m_vecXa{Vector::Zero()}; /// @brief augmented state vector (incl. process - /// and measurement noise means) - Matrix m_matPa{Matrix::Zero()}; /// @brief augmented state covariance (incl. - /// process and measurement noise covariances) - - // Add variables to track the covariances of the process and measurement noise. - Matrix m_matQ; // @brief Process noise covariance matrix - Matrix m_matR; // @brief Measurement noise covariance matrix + Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector + Matrix m_matP{Matrix::Zero()}; /// @brief state covariance matrix + + /// @brief Augmented state vector (incl. process and measurement noise means) + Vector m_vecXa{Vector::Zero()}; + /// @brief augmented state covariance (incl. process and measurement noise covariances) + Matrix m_matPa{Matrix::Zero()}; + + // Process and measurement noise covariance matrices + /// Process noise covariance matrix (Q) + Matrix m_matQ{Matrix::Zero()}; + /// Measurement noise covariance matrix (R) + Matrix m_matR{Matrix::Zero()}; + + /// Unscented transform weight for the mean sigma point + float32_t m_weight0; + /// 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}; public: - Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix - - TrackFitterUKFBase() - : KalmanFilter(), m_kappa(3 - DIM_A), m_matQ(Matrix::Zero()), - m_matR(Matrix::Zero()) - { - // 1. calculate weights - updateWeights(); - } - - ~TrackFitterUKFBase() {} + TrackFitterUKFBase() { updateWeights(); } + ~TrackFitterUKFBase() = default; void setKappa(float32_t kappa) { @@ -76,96 +81,43 @@ class TrackFitterUKFBase : public KalmanFilter { updateWeights(); // Update the weights based on the new kappa value } - /// - /// @brief adding process noise covariance Q to the augmented state covariance - /// matPa in the middle element of the diagonal. - /// - void setCovarianceQ(const Matrix &matQ) - { - m_matQ = matQ; // Store the process noise covariance matrix - } - - /// - /// @brief set the measurement noise covariance R to be used in the update step - /// - void setCovarianceR(const Matrix &matR) - { - m_matR = matR; // Store the measurement noise covariance matrix - } - - /// Add state vector (m_vecX) to the augment state vector (m_vecXa) and also - /// add covariance matrix (m_matP) to the augment covariance (m_matPa). - void updateAugWithState() - { - // Copy state vector to augmented state vector - for (int32_t i{0}; i < DIM_X; ++i) { - m_vecXa[i] = m_vecX[i]; - } - - // Copy state covariance matrix to augmented covariance matrix - for (int32_t i{0}; i < DIM_X; ++i) { - for (int32_t j{0}; j < DIM_X; ++j) { - m_matPa(i, j) = m_matP(i, j); - } - } - } - - virtual std::array calculateProcessNoiseMean() - { - // Calculate the expectation value of the process noise using the current value of the state vector m_vecX - std::array processNoiseMean{0}; - - // TODO: Set the mean energy loss based on the momentum and particle type. Probably best to track stopping power? - return processNoiseMean; - } - - virtual Matrix calculateProcessNoiseCovariance() - { - // Calculate the process noise covariance matrix - Matrix matQ{Matrix::Zero()}; - matQ = m_matQ; // Use the stored process noise covariance matrix - - // TODO: Set the process noise covariance for angular straggle and energy loss. - return matQ; - } - - void updateAugWithProcessNoise() - { - auto processNoiseMean = calculateProcessNoiseMean(); - m_matQ = calculateProcessNoiseCovariance(); - - // Add the mean process noise to the augmented state vector - for (int32_t i{0}; i < DIM_V; ++i) { - m_vecXa[DIM_X + i] = processNoiseMean[i]; - } + /** + * @brief Set process noise covariance Q to be used in the prediction step. + */ + void setCovarianceQ(const Matrix &matQ) { m_matQ = matQ; } + /** + * @brief Set the measurement noise covariance R to be used in the update step. + */ + void setCovarianceR(const Matrix &matR) { m_matR = matR; } + virtual Vector &vecX() { return m_vecX; } + virtual const Vector &vecX() const { return m_vecX; } - // Add process noise covariance to the augmented covariance matrix - const int32_t S_IDX{DIM_X}; - const int32_t L_IDX{S_IDX + DIM_V}; + virtual Matrix &matP() { return m_matP; } + virtual const Matrix &matP() const { return m_matP; } - for (int32_t i{S_IDX}; i < L_IDX; ++i) { - for (int32_t j{S_IDX}; j < L_IDX; ++j) { - m_matPa(i, j) = m_matQ(i - S_IDX, j - S_IDX); - } - } - } - - /// - /// @brief update the augmented state vector and covariance matrix - /// This functions fully updates the augmented state vector (m_vecXa) and covariance matrix (m_matPa) - /// by setting both the state vector and process noise components. - /// + /** + * @brief update the augmented state vector and covariance matrix + * + * This function fully updates the augmented state vector and covariance matrix using + * the state and process noise. + */ void updateAugmentedStateAndCovariance() { updateAugWithState(); updateAugWithProcessNoise(); } - /// - /// @brief state prediction step of the unscented Kalman filter (UKF). - /// @param predictionModelFunc callback to the prediction/process model - /// function - /// + /** + * @brief state prediction step of the unscented Kalman filter (UKF). + * @param predictionModelFunc callback to the prediction/process model function. + * @param vecZ actual measurement vector.` + * + * A template is used here for performace reasons since there is only really a single prediction + * model used. Honestly it probably does not make a difference compared to an std::function, but + * this was not changed from OpenKF. + * + * This modifies the state vector and state covariance. + */ template void predictUKF(PredictionModelCallback predictionModelFunc, const Vector &vecZ) { @@ -199,11 +151,11 @@ class TrackFitterUKFBase : public KalmanFilter { calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP); } - /// - /// @brief measurement correction step of the unscented Kalman filter (UKF). - /// @param measurementModelFunc callback to the measurement model function - /// @param vecZ actual measurement vector. - /// + /** + * @brief measurement correction step of the unscented Kalman filter (UKF). + * @param measurementModelFunc callback to the measurement model function + * @param vecZ actual measurement vector. + */ template void correctUKF(MeasurementModelCallback measurementModelFunc, const Vector &vecZ) { @@ -235,7 +187,6 @@ class TrackFitterUKFBase : public KalmanFilter { // Add in the measurement noise covariance matrix to the measurement covariance matrix. matPzz += m_matR; // Add measurement noise covariance - // TODO: calculate cross correlation const Matrix matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)}; // kalman gain @@ -246,9 +197,9 @@ class TrackFitterUKFBase : public KalmanFilter { } protected: - /// - /// @brief algorithm to calculate the weights used to draw the sigma points - /// + /** + * @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."); @@ -258,17 +209,59 @@ class TrackFitterUKFBase : public KalmanFilter { m_weight0 = m_kappa / denoTerm; m_weighti = 0.5F / denoTerm; } + /** + * @brief Add state vector and state covariance matrix to the augmented state vector covariance matrix. + */ + void updateAugWithState() + { + // Copy state vector to augmented state vector + for (int32_t i{0}; i < DIM_X; ++i) { + m_vecXa[i] = m_vecX[i]; + } + + // Copy state covariance matrix to augmented covariance matrix + for (int32_t i{0}; i < DIM_X; ++i) { + for (int32_t j{0}; j < DIM_X; ++j) { + m_matPa(i, j) = m_matP(i, j); + } + } + } + + virtual std::array calculateProcessNoiseMean() { return std::array{0}; } + + virtual Matrix calculateProcessNoiseCovariance() { return m_matQ; } + + void updateAugWithProcessNoise() + { + auto processNoiseMean = calculateProcessNoiseMean(); + m_matQ = calculateProcessNoiseCovariance(); + + // Add the mean process noise to the augmented state vector + for (int32_t i{0}; i < DIM_V; ++i) { + m_vecXa[DIM_X + i] = processNoiseMean[i]; + } + + // Add process noise covariance to the augmented covariance matrix + const int32_t S_IDX{DIM_X}; + const int32_t L_IDX{S_IDX + DIM_V}; + + for (int32_t i{S_IDX}; i < L_IDX; ++i) { + for (int32_t j{S_IDX}; j < L_IDX; ++j) { + m_matPa(i, j) = m_matQ(i - S_IDX, j - S_IDX); + } + } + } - /// - /// @brief algorithm to calculate the deterministic sigma points for - /// the unscented transformation - /// - /// @param vecX mean of the normally distributed state - /// @param matPxx covariance of the normally distributed state - /// @param STATE_DIM dimension of the vector used to calculate the sigma points - /// @param SIGMA_DIM number of sigma points required (default is 2 * STATE_DIM + 1) - /// @return matrix of sigma points where each column is a sigma point - /// + /** + * @brief Algorithm to calculate the deterministic sigma points for + * the unscented transformation. + * + * @param vecX Mean of the normally distributed state. + * @param matPxx Covariance of the normally distributed state. + * @param STATE_DIM Dimension of the vector used to calculate the sigma points. + * @param SIGMA_DIM Number of sigma points required (default is 2 * STATE_DIM + 1). + * @return Matrix of sigma points where each column is a sigma point. + */ template Matrix calculateSigmaPoints(const Vector &vecXa, const Matrix &matPa) @@ -321,14 +314,12 @@ class TrackFitterUKFBase : public KalmanFilter { return sigmaXa; } - /// - /// @brief calculate the weighted mean and covariance given a set of sigma - /// points - /// @param[in] sigmaX matrix of (probably posterior) sigma points where each column contain single - /// sigma point. - /// @param[out] vecX output weighted mean of the sigma points - /// @param[out] matPxx output weighted covariance of the sigma points - /// + /** + * @brief Calculate the weighted mean and covariance given a set of sigma points. + * @param[in] sigmaX Matrix of (probably posterior) sigma points where each column contains a single sigma point. + * @param[out] vecX Output weighted mean of the sigma points. + * @param[out] matPxx Output weighted covariance of the sigma points. + */ template void calculateWeightedMeanAndCovariance(const Matrix &sigmaX, Vector &vecX, Matrix &matPxx) @@ -356,17 +347,17 @@ class TrackFitterUKFBase : public KalmanFilter { } } - /// - /// @brief calculate the cross-correlation given two sets sigma points X and Y - /// and their means x and y - /// @param sigmaX first matrix of sigma points where each column contain - /// single sigma point - /// @param vecX mean of the first set of sigma points - /// @param sigmaY second matrix of sigma points where each column contain - /// single sigma point - /// @param vecY mean of the second set of sigma points - /// @return matPxy, the cross-correlation matrix - /// + /** + * @brief calculate the cross-correlation given two sets sigma points X and Y + * and their means x and y + * @param sigmaX first matrix of sigma points where each column contain + * single sigma point + * @param vecX mean of the first set of sigma points + * @param sigmaY second matrix of sigma points where each column contain + * single sigma point + * @param vecY mean of the second set of sigma points + * @return matPxy, the cross-correlation matrix + */ template Matrix calculateCrossCorrelation(const Matrix &sigmaX, const Vector &vecX, const Matrix &sigmaY, const Vector &vecY) @@ -389,8 +380,11 @@ class TrackFitterUKFBase : public KalmanFilter { } }; -// Define template dimension variables for clarity and reuse - +/** + * @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm. + * + * UKF specialized for fitting tracks. This class is where all of the physics is + */ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { protected: static constexpr int32_t TF_DIM_X = 6; @@ -402,6 +396,7 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { std::unique_ptr fStepper{nullptr}; AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter + public: bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling double fMaxStragglingFactor{1. / 3.}; ///< @brief Maximum straggling factor for energy loss diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/kalman_filter.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/kalman_filter.h index 90651eb99..66694183e 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/kalman_filter.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/kalman_filter.h @@ -18,10 +18,13 @@ namespace kf { template class KalmanFilter { -public: - KalmanFilter() {} - ~KalmanFilter() {} +protected: + Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector + Matrix m_matP{Matrix::Zero()}; /// @brief state covariance matrix +public: + KalmanFilter() = default; + ~KalmanFilter() = default; virtual Vector &vecX() { return m_vecX; } virtual const Vector &vecX() const { return m_vecX; } @@ -88,10 +91,6 @@ class KalmanFilter { m_vecX = m_vecX + matKk * (vecZ - measurementModelFunc(m_vecX)); m_matP = (matI - matKk * matJcobH) * m_matP; } - -protected: - Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector - Matrix m_matP{Matrix::Zero()}; /// @brief state covariance matrix }; } // namespace kf From bd62e903f7bc3de98047de2d44b927fb5aebb335 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 02:02:33 +0200 Subject: [PATCH 52/71] Save the information needed for smoothing --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 91 +++++++++++++++++++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 84 ++++++----------- 2 files changed, 120 insertions(+), 55 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 1c8012494..407a09edb 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -8,10 +8,31 @@ #include namespace kf { +void TrackFitterUKF::Reset() +{ + // Reset the state vector and covariance matrix + m_vecX.setZero(); + m_matP.setZero(); + m_vecXa.setZero(); + m_matPa.setZero(); + m_matQ.setZero(); + m_matR.setZero(); + m_matSigmaXa.setZero(); + + // Clear the history vectors + m_vecXPredHist.clear(); + m_matPPredHist.clear(); + m_matCPredHist.clear(); + m_vecXHist.clear(); + m_matPHist.clear(); + fMeanStep = AtTools::AtPropagator::StepState(); // Reset the step state +} void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, const TMatrixD &initialCovariance) { + // If we are setting the initial state, then we should clear the history. + Reset(); fPropagator.SetState(initialPosition, initialMomentum); // Set the initial state in the propagator m_vecX[0] = initialPosition.X(); // X position m_vecX[1] = initialPosition.Y(); // Y position @@ -26,6 +47,19 @@ void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition m_matP(i, j) = initialCovariance(i, j); } } + + // Save the initial state in our history vectors + m_vecXHist.push_back(m_vecX); + m_matPHist.push_back(m_matP); + m_vecXPredHist.push_back(m_vecX); + m_matPPredHist.push_back(m_matP); + m_matCPredHist.push_back(Matrix::Zero()); // Cross-correlation is not defined for the first point + + // We need to calculate the sigma points for the initial state + updateAugmentedStateAndCovariance(); // Update the augmented state vector and covariance matrix + m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); // Calculate the sigma points for the initial state + // Now we grab the sigma points only for the state. + m_matSigmaXPred = m_matSigmaXa.block(0, 0, TF_DIM_X, SIGMA_DIM_A); // Extract the state sigma points } TMatrixD TrackFitterUKF::GetStateCovariance() const @@ -156,4 +190,61 @@ Vector TrackFitterUKF::funcH(const Vector zVec; // Initialize the measurement vector + zVec[0] = z.X(); + zVec[1] = z.Y(); + zVec[2] = z.Z(); + auto callback = [this](const kf::Vector &x_, const kf::Vector &v_, + const kf::Vector &z_) { return funcF(x_, v_, z_); }; + TrackFitterUKFBase::predictUKF(callback, zVec); + + // Now we need to store the predicted state and covariance for smoothing later. + m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector + m_matPPredHist.push_back(m_matP); // Store the predicted covariance matrix + + // Get the sigma points belonging to the predicted state + Matrix sigmaXx{m_matSigmaXa.block(0, 0, TF_DIM_X, SIGMA_DIM_A)}; + + // Calculate the cross-corelation between the filtered state at k and predicted state at k+1 + auto matCPred = + calculateCrossCorrelation(m_matSigmaXPred, m_vecXHist.back(), sigmaXx, m_vecXPredHist.back()); + m_matCPredHist.push_back(matCPred); // Store the cross-correlation matrix +} + +void TrackFitterUKF::correctUKF(const ROOT::Math::XYZPoint &z) +{ + Vector zVec; // Initialize the measurement vector + zVec[0] = z.X(); + zVec[1] = z.Y(); + zVec[2] = z.Z(); + auto callback = [this](const kf::Vector &x_) { return funcH(x_); }; + TrackFitterUKFBase::correctUKF(callback, zVec); + + // After correction we need to save the filtered state + m_vecXHist.push_back(m_vecX); // Store the filtered state vector + m_matPHist.push_back(m_matP); // Store the filtered covariance matrix +} + } // namespace kf \ No newline at end of file diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index f525661ec..3a6237688 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -358,19 +358,20 @@ class TrackFitterUKFBase { * @param vecY mean of the second set of sigma points * @return matPxy, the cross-correlation matrix */ - template - Matrix calculateCrossCorrelation(const Matrix &sigmaX, const Vector &vecX, - const Matrix &sigmaY, const Vector &vecY) + template + Matrix + calculateCrossCorrelation(const Matrix &sigmaX, const Vector &vecX, + const Matrix &sigmaY, const Vector &vecY) { - Vector devXi{util::getColumnAt(0, sigmaX) - vecX}; // X[:, 0] - \bar{ x } - Vector devYi{util::getColumnAt(0, sigmaY) - vecY}; // Y[:, 0] - \bar{ y } + Vector devXi{util::getColumnAt(0, sigmaX) - vecX}; // X[:, 0] - \bar{ x } + Vector devYi{util::getColumnAt(0, sigmaY) - vecY}; // Y[:, 0] - \bar{ y } // P_0 = W[0, 0] (X[:, 0] - \bar{x}) (Y[:, 0] - \bar{y})^T - Matrix matPxy{m_weight0 * (devXi * devYi.transpose())}; + Matrix matPxy{m_weight0 * (devXi * devYi.transpose())}; for (int32_t i{1}; i < SIGMA_DIM; ++i) { - devXi = util::getColumnAt(i, sigmaX) - vecX; // X[:, i] - \bar{x} - devYi = util::getColumnAt(i, sigmaY) - vecY; // Y[:, i] - \bar{y} + devXi = util::getColumnAt(i, sigmaX) - vecX; // X[:, i] - \bar{x} + devYi = util::getColumnAt(i, sigmaY) - vecY; // Y[:, i] - \bar{y} matPxy += m_weighti * (devXi * devYi.transpose()); // y += W[0, i] (Y[:, i] - // \bar{y}) (Y[:, i] - \bar{y})^T @@ -397,6 +398,19 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter + // vectors to hold the information needed for smoothing the UKF + std::vector> m_vecXPredHist; /// @brief History of predicted state vectors at k+1 + std::vector> m_matPPredHist; /// @brief History of predicted state covariances at k+1 + /// History of cross correlation between filtered state at k and predicted at k+1 + std::vector> m_matCPredHist; + /// History of filtered (after correction) state vectors at k + std::vector> m_vecXHist; + /// History of filtered (after correction) state covariances at k + std::vector> m_matPHist; + + /// The sigma points after propagation for the last prediction step. + Matrix m_matSigmaXPred{Matrix::Zero()}; + public: bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling double fMaxStragglingFactor{1. / 3.}; ///< @brief Maximum straggling factor for energy loss @@ -415,7 +429,7 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { : TrackFitterUKFBase(), fPropagator(std::move(propagator)), fStepper(std::move(stepper)) { } - + void Reset(); void SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, const TMatrixD &initialCovariance); @@ -428,56 +442,16 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { TMatrixD GetAugStateCovariance() const; std::array GetAugStateVector() const; - kf::Vector - funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z); - kf::Vector funcH(const kf::Vector &x); - - void predictUKF(const ROOT::Math::XYZPoint &z) - { - using namespace ROOT::Math; - - // First we need to propagate the mean state vector to the next measurement point. - XYZPoint startingPosition{m_vecX[0], m_vecX[1], m_vecX[2]}; // Get the starting position from the state vector - Polar3DVector startingMomentum{m_vecX[3], m_vecX[4], - m_vecX[5]}; // Get the starting momentum from the state vector - - LOG(info) << "Propagating reference state from position: " << startingPosition - << " with momentum: " << XYZVector(startingMomentum); - - fPropagator.SetState(startingPosition, XYZVector(startingMomentum)); - fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(z), *fStepper); - fMeanStep = fPropagator.GetState(); // Get the mean step information from the propagator - fMeanStep.fLastPos = startingPosition; // Store the last position - fMeanStep.fLastMom = startingMomentum; // Store the last momentum - - LOG(info) << "Propagated to position: " << fMeanStep.fPos << " with momentum: " << fMeanStep.fMom; - - // Now we can construct the reference plane. - fMeasurementPlane = Plane3D(fMeanStep.fMom.Unit(), - XYZPoint(z)); // Create a plane using the momentum direction and position - Vector zVec; // Initialize the measurement vector - zVec[0] = z.X(); - zVec[1] = z.Y(); - zVec[2] = z.Z(); - auto callback = [this](const kf::Vector &x_, const kf::Vector &v_, - const kf::Vector &z_) { return funcF(x_, v_, z_); }; - TrackFitterUKFBase::predictUKF(callback, zVec); - } - - void correctUKF(const ROOT::Math::XYZPoint &z) - { - Vector zVec; // Initialize the measurement vector - zVec[0] = z.X(); - zVec[1] = z.Y(); - zVec[2] = z.Z(); - auto callback = [this](const kf::Vector &x_) { return funcH(x_); }; - TrackFitterUKFBase::correctUKF(callback, zVec); - } + void predictUKF(const ROOT::Math::XYZPoint &z); + void correctUKF(const ROOT::Math::XYZPoint &z); protected: std::array calculateProcessNoiseMean() override; - Matrix calculateProcessNoiseCovariance() override; + + kf::Vector + funcF(const kf::Vector &x, const kf::Vector &v, const kf::Vector &z); + kf::Vector funcH(const kf::Vector &x); }; } // namespace kf From 322e738f3bb3ab43170c2f987ec660ebcbec7b7d Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 11:30:44 +0200 Subject: [PATCH 53/71] Refactor LLT decomp into function --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 49 ++++++++++++- .../OpenKF/kalman_filter/TrackFitterUKF.h | 69 ++++++++++++------- macro/tests/UKF/UKFSingleTrack.C | 13 ++++ 3 files changed, 105 insertions(+), 26 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 407a09edb..36981c980 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -218,6 +218,18 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) zVec[2] = z.Z(); auto callback = [this](const kf::Vector &x_, const kf::Vector &v_, const kf::Vector &z_) { return funcF(x_, v_, z_); }; + + updateAugmentedStateAndCovariance(); + + // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma + // point. + auto sigmaPoints = calculateSigmaPoints(m_vecXa, m_matPa); + + // Pull out the sigma points for the state vector and process noise in two different matrices. + Matrix sigmaXxPrior{ + sigmaPoints.block(0, 0, TF_DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector + m_matSigmaXPred = sigmaXxPrior; // Store the sigma points for the state vector + TrackFitterUKFBase::predictUKF(callback, zVec); // Now we need to store the predicted state and covariance for smoothing later. @@ -228,8 +240,7 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) Matrix sigmaXx{m_matSigmaXa.block(0, 0, TF_DIM_X, SIGMA_DIM_A)}; // Calculate the cross-corelation between the filtered state at k and predicted state at k+1 - auto matCPred = - calculateCrossCorrelation(m_matSigmaXPred, m_vecXHist.back(), sigmaXx, m_vecXPredHist.back()); + auto matCPred = calculateCrossCorrelation(sigmaXxPrior, m_vecXHist.back(), sigmaXx, m_vecXPredHist.back()); m_matCPredHist.push_back(matCPred); // Store the cross-correlation matrix } @@ -247,4 +258,38 @@ void TrackFitterUKF::correctUKF(const ROOT::Math::XYZPoint &z) m_matPHist.push_back(m_matP); // Store the filtered covariance matrix } +void TrackFitterUKF::smoothUKF() +{ + + // Smoothing is done by iterating backwards over the history of predicted states and covariances. + // Here i = k+1 + m_vecXSmooth.resize(m_vecXPredHist.size()); + m_matPSmooth.resize(m_matPPredHist.size()); + m_vecXSmooth.back() = m_vecXHist.back(); // The last smoothed state is the last corrected state + m_matPSmooth.back() = m_matPHist.back(); // The last smoothed covariance is the last corrected covariance + for (size_t i = m_vecXPredHist.size() - 1; i > 0; --i) { + + // Get the predicted state and covariance at step i + const auto &xPred = m_vecXPredHist[i]; // m_{k+1}^- + const auto &pPred = m_matPPredHist[i]; // P_{k+1}^- + const auto &ccor = m_matCPredHist[i]; // C_{k+1} + + // Get the filtered state and covariance at step i-1 + const auto &xFilt = m_vecXHist[i - 1]; // m_{k} + const auto &pFilt = m_matPHist[i - 1]; // P_{k} + + // Get the smoothed state and covariance at step i + auto &xSmooth = m_vecXSmooth[i]; // m^s_{k+1} + auto &pSmooth = m_matPSmooth[i]; // P^s_{k+1} + + auto llt = calculateCholesky(pPred); // Perform Cholesky decomposition + auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} + + std::cout << "D matrix at step " << i << ":\n" << D << "\n"; + m_vecXSmooth[i - 1] = xFilt + D * (xSmooth - xPred); // m^s_{k} = m_{k} + D * (m^s_{k+1} - m_{k+1}^-) + m_matPSmooth[i - 1] = + pFilt + D * (pSmooth - pPred) * D.transpose(); // P^s_{k} = P_{k} + D * (P^s_{k+1} - P_{k+1}^-) * D^T + } +} + } // namespace kf \ No newline at end of file diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 3a6237688..4617e7000 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -252,6 +252,31 @@ class TrackFitterUKFBase { } } + template + Eigen::LLT> calculateCholesky(const Matrix &matP) + { + Eigen::LLT> lltOfP(matP); + if (lltOfP.info() != Eigen::Success) { + LOG(error) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; + // Add a small value to the diagonal to regularize the matrix + Matrix matPReg = matP; + matPReg = (matPReg + matPReg.transpose()) * 0.5; // Ensure symmetry + matPReg += Matrix::Identity() * 1e-6; // Regularization value + lltOfP.compute(matPReg); + if (lltOfP.info() != Eigen::Success) { + LOG(error) << "Cholesky decomposition failed even after regularization. Attempting again"; + matPReg += Matrix::Identity() * 1e-3; // Increase regularization value + lltOfP.compute(matPReg); + } + if (lltOfP.info() != Eigen::Success) { + LOG(error) << "\n" << matPReg; + throw std::runtime_error( + "Cholesky decomposition failed, matrix is not positive definite even after regularization."); + } + } + return lltOfP; // Return the Cholesky decomposition of the covariance matrix + } + /** * @brief Algorithm to calculate the deterministic sigma points for * the unscented transformation. @@ -268,25 +293,8 @@ class TrackFitterUKFBase { { const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_kappa)}; // sqrt(n + \kappa) - Eigen::LLT> lltOfPa(matPa); - if (lltOfPa.info() != Eigen::Success) { - LOG(error) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; - // Add a small value to the diagonal to regularize the matrix - Matrix matPaReg = matPa; - matPaReg = (matPaReg + matPaReg.transpose()) * 0.5; // Ensure symmetry - matPaReg += Matrix::Identity() * 1e-6; // Regularization value - lltOfPa.compute(matPaReg); - if (lltOfPa.info() != Eigen::Success) { - LOG(error) << "Cholesky decomposition failed even after regularization. Attempting again"; - matPaReg += Matrix::Identity() * 1e-3; // Increase regularization value - lltOfPa.compute(matPaReg); - } - if (lltOfPa.info() != Eigen::Success) { - LOG(error) << "\n" << matPaReg; - throw std::runtime_error( - "Cholesky decomposition failed, matrix is not positive definite even after regularization."); - } - } + Eigen::LLT> lltOfPa = calculateCholesky(matPa); + Matrix matSa{lltOfPa.matrixL()}; // sqrt(P_{a}) matSa *= scalarMultiplier; // sqrt( (n + \kappa) * P_{a} ) @@ -398,15 +406,23 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter + using EigenVectorDimX = std::vector, Eigen::aligned_allocator>>; + using VectorEigenMatDimX = + std::vector, Eigen::aligned_allocator>>; // vectors to hold the information needed for smoothing the UKF - std::vector> m_vecXPredHist; /// @brief History of predicted state vectors at k+1 - std::vector> m_matPPredHist; /// @brief History of predicted state covariances at k+1 + EigenVectorDimX m_vecXPredHist; /// @brief History of predicted state vectors at k+1 + VectorEigenMatDimX m_matPPredHist; /// @brief History of predicted state covariances at k+1 /// History of cross correlation between filtered state at k and predicted at k+1 - std::vector> m_matCPredHist; + VectorEigenMatDimX m_matCPredHist; /// History of filtered (after correction) state vectors at k - std::vector> m_vecXHist; + EigenVectorDimX m_vecXHist; /// History of filtered (after correction) state covariances at k - std::vector> m_matPHist; + VectorEigenMatDimX m_matPHist; + + /// Smoothed state vector and covariance + EigenVectorDimX m_vecXSmooth; + /// Smoothed state covariance + VectorEigenMatDimX m_matPSmooth; /// The sigma points after propagation for the last prediction step. Matrix m_matSigmaXPred{Matrix::Zero()}; @@ -441,9 +457,14 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { TMatrixD GetStateCovariance() const; TMatrixD GetAugStateCovariance() const; std::array GetAugStateVector() const; + const EigenVectorDimX &GetSmoothedStates() const { return m_vecXSmooth; }; + const VectorEigenMatDimX &GetSmoothedCovariances() const { return m_matPSmooth; }; + const EigenVectorDimX &GetFilteredStates() const { return m_vecXHist; }; + const VectorEigenMatDimX &GetFilteredCovariances() const { return m_matPHist; }; void predictUKF(const ROOT::Math::XYZPoint &z); void correctUKF(const ROOT::Math::XYZPoint &z); + void smoothUKF(); protected: std::array calculateProcessNoiseMean() override; diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index a6a82cd5a..565f72492 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -171,6 +171,19 @@ void UKFSingleTrack() residual.push_back(residualValue); // Store the residual for this hit } + // At this point we have the full trajectory of the particle + ukf.smoothUKF(); // Perform smoothing + auto smoothedStates = ukf.GetSmoothedStates(); + auto smoothedCovariances = ukf.GetSmoothedCovariances(); + auto filteredStates = ukf.GetFilteredStates(); + auto filteredCovariances = ukf.GetFilteredCovariances(); + for (int i = 0; i < smoothedStates.size(); ++i) { + auto &state = smoothedStates[i]; + auto &filteredState = filteredStates[i]; + std::cout << "Smoothed state " << i << ": " << state.transpose() << std::endl; + std::cout << "Filtered state " << i << ": " << filteredState.transpose() << std::endl; + } + TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); track->SetTitle("Particle Track;X [mm];Y [mm];Z [mm]"); track->SetMarkerStyle(20); From 841a7cdb6b3b3932a66d4e4e846eb7a791b18b0a Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 12:32:22 +0200 Subject: [PATCH 54/71] Adjust weights to use method that should protect PD of COV better --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 6 +++-- .../OpenKF/kalman_filter/TrackFitterUKF.h | 24 +++++++++++------ .../kalman_filter/TrackFitterUKFTest.cxx | 2 +- macro/tests/UKF/UKFSingleTrack.C | 26 +++++++++++++++++++ 4 files changed, 47 insertions(+), 11 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 36981c980..bdb68edc6 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -268,6 +268,7 @@ void TrackFitterUKF::smoothUKF() m_vecXSmooth.back() = m_vecXHist.back(); // The last smoothed state is the last corrected state m_matPSmooth.back() = m_matPHist.back(); // The last smoothed covariance is the last corrected covariance for (size_t i = m_vecXPredHist.size() - 1; i > 0; --i) { + LOG(info) << "Smoothing step " << i << " of " << m_vecXPredHist.size() - 1; // Get the predicted state and covariance at step i const auto &xPred = m_vecXPredHist[i]; // m_{k+1}^- @@ -283,9 +284,10 @@ void TrackFitterUKF::smoothUKF() auto &pSmooth = m_matPSmooth[i]; // P^s_{k+1} auto llt = calculateCholesky(pPred); // Perform Cholesky decomposition - auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} + auto D = ccor * llt.solve(Matrix::Identity()); + // auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} - std::cout << "D matrix at step " << i << ":\n" << D << "\n"; + // std::cout << "D matrix at step " << i << ":\n" << D << "\n"; m_vecXSmooth[i - 1] = xFilt + D * (xSmooth - xPred); // m^s_{k} = m_{k} + D * (m^s_{k+1} - m_{k+1}^-) m_matPSmooth[i - 1] = pFilt + D * (pSmooth - pPred) * D.transpose(); // P^s_{k} = P_{k} + D * (P^s_{k+1} - P_{k+1}^-) * D^T diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 4617e7000..e0f032614 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -64,8 +64,10 @@ class TrackFitterUKFBase { /// Measurement noise covariance matrix (R) Matrix m_matR{Matrix::Zero()}; - /// Unscented transform weight for the mean sigma point - float32_t m_weight0; + /// Unscented transform weight for the mean sigma point in mean + float32_t m_weightM0; + /// Unscented transform weight for the mean sigma point in covariance + float32_t m_weightC0; /// Unscented transform weight for the other sigma points float32_t m_weighti; /// Kappa parameter for sigma point calculation @@ -204,9 +206,15 @@ class TrackFitterUKFBase { { static_assert(DIM_A > 0, "DIM_A is Zero which leads to numerical issue."); - const float32_t denoTerm{m_kappa + static_cast(DIM_A)}; + constexpr float alpha = 1.0; // Scaling parameter, set to 1 to match orig + constexpr float beta = 2.0; // Optimal for Gaussian distributions - m_weight0 = m_kappa / denoTerm; + 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(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; } /** @@ -333,7 +341,7 @@ class TrackFitterUKFBase { Matrix &matPxx) { // 1. calculate mean of the sigma points - vecX = m_weight0 * util::getColumnAt(0, sigmaX); + vecX = m_weightM0 * util::getColumnAt(0, sigmaX); for (int32_t i{1}; i < SIGMA_DIM; ++i) { vecX += m_weighti * util::getColumnAt(i, sigmaX); // y += W[0, i] Y[:, i] } @@ -342,8 +350,8 @@ class TrackFitterUKFBase { // \bar{y}) (Y[:, i] - \bar{y})^T Vector devXi{util::getColumnAt(0, sigmaX) - vecX}; // Y[:, 0] - \bar{ y } - matPxx = m_weight0 * devXi * devXi.transpose(); // P_0 = W[0, 0] (Y[:, 0] - \bar{y}) (Y[:, 0] - - // \bar{y})^T + matPxx = m_weightC0 * devXi * devXi.transpose(); // P_0 = W[0, 0] (Y[:, 0] - \bar{y}) (Y[:, 0] - + // \bar{y})^T for (int32_t i{1}; i < SIGMA_DIM; ++i) { devXi = util::getColumnAt(i, sigmaX) - vecX; // Y[:, i] - \bar{y} @@ -375,7 +383,7 @@ class TrackFitterUKFBase { Vector devYi{util::getColumnAt(0, sigmaY) - vecY}; // Y[:, 0] - \bar{ y } // P_0 = W[0, 0] (X[:, 0] - \bar{x}) (Y[:, 0] - \bar{y})^T - Matrix matPxy{m_weight0 * (devXi * devYi.transpose())}; + Matrix matPxy{m_weightC0 * (devXi * devYi.transpose())}; for (int32_t i{1}; i < SIGMA_DIM; ++i) { devXi = util::getColumnAt(i, sigmaX) - vecX; // X[:, i] - \bar{x} diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 9709f4797..4f502d9f1 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -188,7 +188,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.4758845F, FLOAT_EPSILON); + ASSERT_NEAR(m_ukf.vecX()[0], 2.47414017F, 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); diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 565f72492..92a852ef7 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -61,6 +61,7 @@ void UKFSingleTrack() using namespace AtTools; std::vector x2, y2, z2, Eloss2, p2, sigmap2, lambda2, sigmalambda2, residual; + std::vector xSmooth, ySmooth, zSmooth, pSmooth, sigmapSmooth; // Setup the Propagator for UKF auto elossModel = std::make_unique(0); @@ -182,6 +183,11 @@ void UKFSingleTrack() auto &filteredState = filteredStates[i]; std::cout << "Smoothed state " << i << ": " << state.transpose() << std::endl; std::cout << "Filtered state " << i << ": " << filteredState.transpose() << std::endl; + xSmooth.push_back(state[0]); + ySmooth.push_back(state[1]); + zSmooth.push_back(state[2]); + pSmooth.push_back(state[3]); + sigmapSmooth.push_back(std::sqrt(smoothedCovariances[i](3, 3))); // Momentum uncertainty } TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); @@ -195,9 +201,29 @@ void UKFSingleTrack() track2->SetMarkerSize(0.8); track2->SetMarkerColor(kRed); + TGraph2D *smoothedTrack = new TGraph2D(xSmooth.size(), xSmooth.data(), ySmooth.data(), zSmooth.data()); + smoothedTrack->SetTitle("Smoothed Particle Track;X [mm];Y [mm];Z [mm]"); + smoothedTrack->SetMarkerStyle(22); + smoothedTrack->SetMarkerSize(0.8); + smoothedTrack->SetMarkerColor(kGreen + 2); + TCanvas *c1 = new TCanvas("c1", "Particle Track", 800, 600); + + // Set axis ranges based on track and track2 + double xmin = std::min(*std::min_element(x.begin(), x.end()), *std::min_element(x2.begin(), x2.end())); + double xmax = std::max(*std::max_element(x.begin(), x.end()), *std::max_element(x2.begin(), x2.end())); + double ymin = std::min(*std::min_element(y.begin(), y.end()), *std::min_element(y2.begin(), y2.end())); + double ymax = std::max(*std::max_element(y.begin(), y.end()), *std::max_element(y2.begin(), y2.end())); + double zmin = std::min(*std::min_element(z.begin(), z.end()), *std::min_element(z2.begin(), z2.end())); + double zmax = std::max(*std::max_element(z.begin(), z.end()), *std::max_element(z2.begin(), z2.end())); + + track->GetXaxis()->SetLimits(xmin, xmax); + track->GetYaxis()->SetLimits(ymin, ymax); + track->GetZaxis()->SetLimits(zmin, zmax); + track->Draw("P"); track2->Draw("PSAME"); + smoothedTrack->Draw("PSAME"); TGraph *elossGraph = new TGraph(Eloss.size()); for (size_t i = 0; i < Eloss.size(); ++i) { From d73785f90097dc648a7c22d11b1df0c7b55e24f6 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 13:05:37 +0200 Subject: [PATCH 55/71] Add function to ensure PD and only save PD matrices Smoother is more stable, but it still fail part way through, so we are still poorly conditioned somewhere. --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 3 ++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 53 +++++++++++++++---- .../kalman_filter/TrackFitterUKFTest.cxx | 1 + 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index bdb68edc6..eee2d4288 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -223,7 +223,9 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma // point. + LOG(info) << "Calculating sigma points for prediction step."; auto sigmaPoints = calculateSigmaPoints(m_vecXa, m_matPa); + LOG(info) << "Finished calculating sigma points for prediction step."; // Pull out the sigma points for the state vector and process noise in two different matrices. Matrix sigmaXxPrior{ @@ -231,6 +233,7 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) m_matSigmaXPred = sigmaXxPrior; // Store the sigma points for the state vector TrackFitterUKFBase::predictUKF(callback, zVec); + LOG(info) << "Finished prediction step."; // Now we need to store the predicted state and covariance for smoothing later. m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index e0f032614..afd8a23e3 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -107,6 +107,7 @@ class TrackFitterUKFBase { { updateAugWithState(); updateAugWithProcessNoise(); + ensurePD(m_matPa); // Ensure the augmented covariance matrix is positive definite } /** @@ -127,7 +128,9 @@ class TrackFitterUKFBase { // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma // point. + LOG(info) << "Calculating sigma points for prediction step2."; m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); + LOG(info) << "Finished calculating sigma points for prediction step2."; // Pull out the sigma points for the state vector and process noise in two different matrices. Matrix sigmaXx{m_matSigmaXa.block(0, 0, DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector @@ -260,31 +263,58 @@ class TrackFitterUKFBase { } } + /** + * @brief Calculate Cholesky decomposition of a covariance matrix and update the matrix so it is PD. + * + * Modifies the input matrix to ensure it is symmetric and positive definite. + * If the decomposition fails, it attempts to regularize the matrix by adding a small value + * to the diagonal and retrying the decomposition. + */ template Eigen::LLT> calculateCholesky(const Matrix &matP) { Eigen::LLT> lltOfP(matP); if (lltOfP.info() != Eigen::Success) { - LOG(error) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; + throw std::runtime_error("Cholesky decomposition failed, matrix is not positive definite."); + } + + return lltOfP; // Return the Cholesky decomposition of the covariance matrix + } + + template + Eigen::LLT> ensurePD(Matrix &matP) + { + Eigen::LLT> lltOfP(matP); + if (lltOfP.info() != Eigen::Success) { + LOG(warn) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; // Add a small value to the diagonal to regularize the matrix - Matrix matPReg = matP; - matPReg = (matPReg + matPReg.transpose()) * 0.5; // Ensure symmetry - matPReg += Matrix::Identity() * 1e-6; // Regularization value - lltOfP.compute(matPReg); - if (lltOfP.info() != Eigen::Success) { - LOG(error) << "Cholesky decomposition failed even after regularization. Attempting again"; - matPReg += Matrix::Identity() * 1e-3; // Increase regularization value - lltOfP.compute(matPReg); + int i = 0; + while (lltOfP.info() != Eigen::Success && i < 3) { + LOG(debug) << "Attempting to regularize covariance matrix, iteration: " << i; + + symmetrize(matP); // Ensure symmetry before regularization + matP += Matrix::Identity() * std::pow(10, -6 + i); // Regularization value + lltOfP.compute(matP); + i = i + 1; + // Check if the regularized matrix is now positive definite } if (lltOfP.info() != Eigen::Success) { - LOG(error) << "\n" << matPReg; + LOG(error) << "\n" << matP; throw std::runtime_error( "Cholesky decomposition failed, matrix is not positive definite even after regularization."); - } + } else + LOG(warn) << "Cholesky decomposition succeeded after regularization of order " << i - 1; } return lltOfP; // Return the Cholesky decomposition of the covariance matrix } + template + void symmetrize(Matrix &matP) + { + // Ensure the matrix is symmetric + matP = (matP + matP.transpose()) * 0.5; + } + /** * @brief Algorithm to calculate the deterministic sigma points for * the unscented transformation. @@ -361,6 +391,7 @@ class TrackFitterUKFBase { matPxx += Pi; // y += W[0, i] (Y[:, i] - \bar{y}) (Y[:, i] - \bar{y})^T } + ensurePD(matPxx); // Ensure the covariance matrix is positive definite } /** diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 4f502d9f1..50ad94688 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -432,6 +432,7 @@ TEST_F(TrackFitterUKFFixture, TestPredictionStep) cov(5, 5) = sigma_phi * sigma_phi; // Angular uncertainty m_ukf->SetInitialState(startPos, startMom, cov); + LOG(info) << "Finished setting initial state"; XYZPoint point({-1.4895, -4.8787, 1.01217}); // measurement point in cm point *= 10; // Convert to mm From 22c624e38b356521281b802e9b7a25d3732006f3 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 14:10:31 +0200 Subject: [PATCH 56/71] Remove debug. Add diagnostic plots --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 3 - .../OpenKF/kalman_filter/TrackFitterUKF.h | 3 +- macro/tests/UKF/UKFSingleTrack.C | 62 ++++++++++++++++--- 3 files changed, 55 insertions(+), 13 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index eee2d4288..bdb68edc6 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -223,9 +223,7 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma // point. - LOG(info) << "Calculating sigma points for prediction step."; auto sigmaPoints = calculateSigmaPoints(m_vecXa, m_matPa); - LOG(info) << "Finished calculating sigma points for prediction step."; // Pull out the sigma points for the state vector and process noise in two different matrices. Matrix sigmaXxPrior{ @@ -233,7 +231,6 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) m_matSigmaXPred = sigmaXxPrior; // Store the sigma points for the state vector TrackFitterUKFBase::predictUKF(callback, zVec); - LOG(info) << "Finished prediction step."; // Now we need to store the predicted state and covariance for smoothing later. m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index afd8a23e3..c12f95062 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -128,9 +128,7 @@ class TrackFitterUKFBase { // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma // point. - LOG(info) << "Calculating sigma points for prediction step2."; m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); - LOG(info) << "Finished calculating sigma points for prediction step2."; // Pull out the sigma points for the state vector and process noise in two different matrices. Matrix sigmaXx{m_matSigmaXa.block(0, 0, DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector @@ -284,6 +282,7 @@ class TrackFitterUKFBase { template Eigen::LLT> ensurePD(Matrix &matP) { + symmetrize(matP); Eigen::LLT> lltOfP(matP); if (lltOfP.info() != Eigen::Success) { LOG(warn) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 92a852ef7..4d7007cd4 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -61,7 +61,7 @@ void UKFSingleTrack() using namespace AtTools; std::vector x2, y2, z2, Eloss2, p2, sigmap2, lambda2, sigmalambda2, residual; - std::vector xSmooth, ySmooth, zSmooth, pSmooth, sigmapSmooth; + std::vector xSmooth, ySmooth, zSmooth, pSmooth, sigmapSmooth, residualSmooth, eLossSmooth; // Setup the Propagator for UKF auto elossModel = std::make_unique(0); @@ -188,6 +188,17 @@ void UKFSingleTrack() zSmooth.push_back(state[2]); pSmooth.push_back(state[3]); sigmapSmooth.push_back(std::sqrt(smoothedCovariances[i](3, 3))); // Momentum uncertainty + residualSmooth.push_back( + (XYZPoint(state[0], state[1], state[2]) - XYZPoint(filteredState[0], filteredState[1], filteredState[2])).R()); + if (i > 0) { + auto lastMom = smoothedStates[i - 1][3]; + auto mom = smoothedStates[i][3]; + auto KE_in = Kinematics::KE(lastMom, mass_p); + auto KE_out = Kinematics::KE(mom, mass_p); + eLossSmooth.push_back(KE_in - KE_out); // Energy loss between smoothed states + } else { + eLossSmooth.push_back(0); // First point has no previous state to compare + } } TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); @@ -239,13 +250,19 @@ void UKFSingleTrack() eloss2Graph->SetTitle("Propagated Energy Loss per Hit;Hit Number;Energy Loss [MeV]"); eloss2Graph->SetMarkerStyle(21); eloss2Graph->SetMarkerColor(kRed); + TGraph *elossSmoothGraph = new TGraph(eLossSmooth.size()); + for (size_t i = 0; i < eLossSmooth.size(); ++i) { + elossSmoothGraph->SetPoint(i, i, eLossSmooth[i]); + } + elossSmoothGraph->SetTitle("Smoothed Energy Loss per Hit;Hit Number;Energy Loss [MeV]"); + elossSmoothGraph->SetMarkerStyle(22); + elossSmoothGraph->SetMarkerColor(kGreen + 2); TGraphErrors *pGraph = new TGraphErrors(p2.size()); for (size_t i = 0; i < p2.size(); ++i) { - pGraph->SetPoint(i, i, p2[i] * 1e-3 * pointsToCluster / 5.); + pGraph->SetPoint(i, i, p2[i]); pGraph->SetPointError(i, 0, - sigmap2[i] * 5 * 1e-3 * pointsToCluster / - 5.); // Error bars from sigmap2, converted to GeV/c + sigmap2[i] * 5); // Error bars from sigmap2, converted to GeV/c } pGraph->SetTitle("Momentum per Hit;Hit Number;Momentum [MeV/c]"); pGraph->SetMarkerStyle(20); @@ -265,16 +282,45 @@ void UKFSingleTrack() for (size_t i = 0; i < residual.size(); ++i) { residualGraph->SetPoint(i, i, residual[i] * .1); } - residualGraph->SetTitle("Residual per Hit;Hit Number;Residual [mm]"); + residualGraph->SetTitle("Residual per Hit;Hit Number;Residual [cm]"); residualGraph->SetMarkerStyle(23); residualGraph->SetMarkerColor(kMagenta); + TGraphErrors *pSmoothGraph = new TGraphErrors(pSmooth.size()); + for (size_t i = 0; i < pSmooth.size(); ++i) { + pSmoothGraph->SetPoint(i, i, pSmooth[i]); + pSmoothGraph->SetPointError(i, 0, + sigmapSmooth[i] * 5); // Error bars from sigmapSmooth, converted to GeV/c + } + pSmoothGraph->SetTitle("Smoothed Momentum per Hit;Hit Number;Momentum [MeV/c]"); + pSmoothGraph->SetMarkerStyle(22); + pSmoothGraph->SetMarkerColor(kGreen + 2); + + TGraph *residualSmoothGraph = new TGraph(residualSmooth.size()); + for (size_t i = 0; i < residualSmooth.size(); ++i) { + residualSmoothGraph->SetPoint(i, i, residualSmooth[i] * .1); + } + residualSmoothGraph->SetTitle("Smoothed Residual per Hit;Hit Number;Residual [cm]"); + residualSmoothGraph->SetMarkerStyle(24); + residualSmoothGraph->SetMarkerColor(kOrange + 7); + TCanvas *c2 = new TCanvas("c2", "Energy Loss per Hit", 800, 600); elossGraph->Draw("AP"); eloss2Graph->Draw("PSAME"); - pGraph->Draw("PSAME"); - // lambdaGraph->Draw("PSAME"); - residualGraph->Draw("PSAME"); + elossSmoothGraph->Draw("PSAME"); + // pGraph->Draw("PSAME"); + // lambdaGraph->Draw("PSAME"); + // residualGraph->Draw("PSAME"); + // pSmoothGraph->Draw("PSAME"); + // residualSmoothGraph->Draw("PSAME"); + + TCanvas *c3 = new TCanvas("c3", "Momentum at Hit (error bars 5X)", 800, 600); + pGraph->Draw("AP"); + pSmoothGraph->Draw("PSAME"); + + TCanvas *c4 = new TCanvas("c4", "Residual at Hit", 800, 600); + residualGraph->Draw("AP"); + residualSmoothGraph->Draw("PSAME"); double sumEloss = std::accumulate(Eloss.begin(), Eloss.end(), 0.0); double sumEloss2 = std::accumulate(Eloss2.begin(), Eloss2.end(), 0.0); From 56bc6218b5e3663712d448d3269c9a4a3629974a Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 15:40:57 +0200 Subject: [PATCH 57/71] Parameter tuning --- .../OpenKF/kalman_filter/TrackFitterUKF.h | 45 +++++++++---------- .../kalman_filter/TrackFitterUKFTest.cxx | 4 +- macro/tests/UKF/AtPropagator.C | 9 +++- macro/tests/UKF/UKFSingleTrack.C | 8 +++- 4 files changed, 38 insertions(+), 28 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index c12f95062..fc35231f8 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/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(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(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 calculateSigmaPoints(const Vector &vecXa, const Matrix &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> lltOfPa = calculateCholesky(matPa); @@ -447,6 +441,7 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { using EigenVectorDimX = std::vector, Eigen::aligned_allocator>>; using VectorEigenMatDimX = std::vector, Eigen::aligned_allocator>>; + // 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 diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 50ad94688..8ce7dfc3d 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/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); diff --git a/macro/tests/UKF/AtPropagator.C b/macro/tests/UKF/AtPropagator.C index 80774a23c..a20fef05a 100644 --- a/macro/tests/UKF/AtPropagator.C +++ b/macro/tests/UKF/AtPropagator.C @@ -33,7 +33,14 @@ void AtPropagator() auto elossModel = std::make_unique(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(3.553e-5); + elossModel2->SetProjectile(1, 1, 1); + std::vector> 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; diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 4d7007cd4..d6639004e 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/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,9 @@ 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; + LOG(info) << "Error in momentum reconstruction: " << (smoothedStates[0][3] - beginMom) / beginMom * 100 << "%"; TGraph2D *track = new TGraph2D(x.size(), x.data(), y.data(), z.data()); track->SetTitle("Particle Track;X [mm];Y [mm];Z [mm]"); @@ -272,7 +277,7 @@ void UKFSingleTrack() for (size_t i = 0; i < lambda2.size(); ++i) { lambdaGraph->SetPoint(i, i, lambda2[i] * 0.03 * pointsToCluster / 5.); lambdaGraph->SetPointError(i, 0, sigmalambda2[i] * 0.03 * pointsToCluster / 5.); - std::cout << "Lambda: " << lambda2[i] << ", Error: " << sigmalambda2[i] << std::endl; + // std::cout << "Lambda: " << lambda2[i] << ", Error: " << sigmalambda2[i] << std::endl; } lambdaGraph->SetTitle("Lambda per Hit (scaled);Hit Number;Lambda [scaled]"); lambdaGraph->SetMarkerStyle(22); @@ -326,4 +331,5 @@ void UKFSingleTrack() double sumEloss2 = std::accumulate(Eloss2.begin(), Eloss2.end(), 0.0); std::cout << "Sum of Eloss: " << sumEloss << std::endl; std::cout << "Sum of Eloss2: " << sumEloss2 << std::endl; + std::cout << "Initial energy: " << Kinematics::KE(beginMom, mass_p) << " MeV" << std::endl; } \ No newline at end of file From 06919b8a4b54f96259bf12af8c3fff0337e2b161 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 18:03:10 +0200 Subject: [PATCH 58/71] Add ability to build CATIMA offline if source is already downloaded. --- CMakeLists.txt | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 1ae91b0e3..3d9fff690 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -126,20 +126,37 @@ find_package2(PUBLIC CATIMA) find_package2(PUBLIC Eigen3 3.3 NO_MODULE) # If catima was not found, then download and build it locally +set(CATIMA_LOCAL_SOURCE "${CMAKE_BINARY_DIR}/_deps/catima-src") +set(CATIMA_USE_LOCAL ON) +if(NOT EXISTS "${CATIMA_LOCAL_SOURCE}/CMakeLists.txt" AND ${CATIMA_USE_LOCAL}) + message(WARNING "Local cmake requested, but the source could not be found at ${CATIMA_LOCAL_SOURCE}.") + message(INFO EXISTS "${CATIMA_LOCAL_SOURCE}/CMakeLists.txt" ) + SET(CATIMA_USE_LOCAL OFF) +endif() + if(NOT CATIMA_FOUND) message(STATUS "CATIMA not found, downloading and building locally") - FetchContent_Declare(catima - GIT_REPOSITORY https://github.com/hrosiak/catima.git - GIT_TAG master) + if(${CATIMA_USE_LOCAL}) + FetchContent_Declare(catima + SOURCE_DIR ${CATIMA_LOCAL_DIR}) + else() + FetchContent_Declare(catima + GIT_REPOSITORY https://github.com/hrosiak/catima.git + GIT_TAG master + SOURCE_DIR ${CATIMA_LOCAL_DIR}) + endif() + # Set CATIMA build options before making it available set(TESTS OFF CACHE BOOL "Disable CATIMA tests" FORCE) set(EXAMPLES OFF CACHE BOOL "Disable CATIMA examples" FORCE) set(APPS OFF CACHE BOOL "Disable CATIMA applications" FORCE) set(GENERATE_DATA OFF CACHE BOOL "Disable CATIMA data generator" FORCE) + FetchContent_GetProperties(catima) if (NOT catima_POPULATED) + message(STATUS "Populating CATIMA") FetchContent_Populate(catima) add_subdirectory(${catima_SOURCE_DIR} ${catima_BINARY_DIR}) endif() From 520786fc1ceb67027b4df104fff6ba42142d62b9 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 19:03:17 +0200 Subject: [PATCH 59/71] Increase stability of filter Ensures each update to a covariance matrix maintains PD. --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 3 ++- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h | 14 +++++++++----- macro/tests/UKF/UKFSingleTrack.C | 2 +- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index bdb68edc6..0752750e7 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -285,7 +285,8 @@ void TrackFitterUKF::smoothUKF() auto llt = calculateCholesky(pPred); // Perform Cholesky decomposition auto D = ccor * llt.solve(Matrix::Identity()); - // auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} + // auto D = llt.solve(ccor.transpose()).transpose(); + // auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} // std::cout << "D matrix at step " << i << ":\n" << D << "\n"; m_vecXSmooth[i - 1] = xFilt + D * (xSmooth - xPred); // m^s_{k} = m_{k} + D * (m^s_{k+1} - m_{k+1}^-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index fc35231f8..43388f379 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -195,20 +195,24 @@ class TrackFitterUKFBase { // calculate the mean measurement vector and covariance matrix // from the sigma points. - Vector vecZhat; - Matrix matPzz; + Vector vecZhat; // Predicted measurement vector + Matrix matPzz; // Measurement covariance matrix calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz); // Add in the measurement noise covariance matrix to the measurement covariance matrix. - matPzz += m_matR; // Add measurement noise covariance + matPzz += m_matR; // Add measurement noise covariance so we gen the innovation covariance matrix. + ensurePD(matPzz); // Ensure the covariance matrix is positive definite const Matrix matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)}; // kalman gain - const Matrix matK{matPxz * matPzz.inverse()}; + auto llt = calculateCholesky(matPzz); + const Matrix matK = llt.solve(matPxz.transpose()).transpose(); + // Matrix matK = {matPxz * llt.solve(Matrix::Identity())}; m_vecX += matK * (vecZ - vecZhat); m_matP -= matK * matPzz * matK.transpose(); + ensurePD(m_matP); // Ensure the covariance matrix is positive definite } protected: @@ -279,7 +283,7 @@ class TrackFitterUKFBase { symmetrize(matP); Eigen::LLT> lltOfP(matP); if (lltOfP.info() != Eigen::Success) { - LOG(warn) << "Cholesky decomposition failed, matrix is not positive definite. Attempting recovery..."; + LOG(warn) << "Cholesky decomposition failed while ensuring PD. Attempting recovery..."; // Add a small value to the diagonal to regularize the matrix int i = 0; while (lltOfP.info() != Eigen::Success && i < 3) { diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index d6639004e..b524138ed 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -99,7 +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 + ukf.setParameters(1e-1, 2, 0); // Set kappa to match the original implementation TMatrixD cov(6, 6); cov.Zero(); From d9247ca0958eeb71c67195b476828a989453a51e Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 19:21:01 +0200 Subject: [PATCH 60/71] Optimize Kalman gain calculation --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 43388f379..e5cacd4fb 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -207,11 +207,13 @@ class TrackFitterUKFBase { // kalman gain auto llt = calculateCholesky(matPzz); - const Matrix matK = llt.solve(matPxz.transpose()).transpose(); - // Matrix matK = {matPxz * llt.solve(Matrix::Identity())}; + // const Matrix matK = llt.solve(matPxz.transpose()).transpose(); + Matrix matK = {matPxz * llt.solve(Matrix::Identity())}; + // Matrix matK = matPxz * matPzz.inverse(); By far the worst method for filter stability m_vecX += matK * (vecZ - vecZhat); m_matP -= matK * matPzz * matK.transpose(); + // m_matP -= matPxz * matK.transpose(); ensurePD(m_matP); // Ensure the covariance matrix is positive definite } From 42994ef75d92a4d449670f437fa3cff94ec6968a Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Tue, 29 Jul 2025 19:30:29 +0200 Subject: [PATCH 61/71] Finish optimizing of smoothing gain --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 0752750e7..0f754ebcd 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -283,10 +283,9 @@ void TrackFitterUKF::smoothUKF() auto &xSmooth = m_vecXSmooth[i]; // m^s_{k+1} auto &pSmooth = m_matPSmooth[i]; // P^s_{k+1} - auto llt = calculateCholesky(pPred); // Perform Cholesky decomposition - auto D = ccor * llt.solve(Matrix::Identity()); - // auto D = llt.solve(ccor.transpose()).transpose(); - // auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} + auto llt = calculateCholesky(pPred); + // auto D = ccor * llt.solve(Matrix::Identity()); + auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} // std::cout << "D matrix at step " << i << ":\n" << D << "\n"; m_vecXSmooth[i - 1] = xFilt + D * (xSmooth - xPred); // m^s_{k} = m_{k} + D * (m^s_{k+1} - m_{k+1}^-) From 075ab927ed3672fd3bc767bd6f0a27c76451644a Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 12:03:17 +0200 Subject: [PATCH 62/71] Zero missing parts of filter --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 0f754ebcd..3fbb31e27 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -18,6 +18,7 @@ void TrackFitterUKF::Reset() m_matQ.setZero(); m_matR.setZero(); m_matSigmaXa.setZero(); + m_matSigmaXPred.setZero(); // Clear the history vectors m_vecXPredHist.clear(); @@ -25,6 +26,8 @@ void TrackFitterUKF::Reset() m_matCPredHist.clear(); m_vecXHist.clear(); m_matPHist.clear(); + m_vecXSmooth.clear(); + m_matPSmooth.clear(); fMeanStep = AtTools::AtPropagator::StepState(); // Reset the step state } @@ -136,10 +139,10 @@ Matrix TrackFitterUKF::calcu factor = fMaxStragglingFactor; } matQ(0, 0) = factor * factor; // Variance for the dedx straggling. - LOG(info) << "Calculating process noise for straggling between " << eIn << " MeV and " << eOut << " MeV over " - << elossModel->GetRange(eIn, eOut) << " mm."; - LOG(info) << "Process noise covariance for energy straggling: " << matQ(0, 0) << " (factor: " << factor - << ", dedx_straggle: " << dedx_straggle << ", dEdx: " << elossModel->GetdEdx(eIn) << ")"; + LOG(debug) << "Calculating process noise for straggling between " << eIn << " MeV and " << eOut << " MeV over " + << elossModel->GetRange(eIn, eOut) << " mm."; + LOG(debug) << "Process noise covariance for energy straggling: " << matQ(0, 0) << " (factor: " << factor + << ", dedx_straggle: " << dedx_straggle << ", dEdx: " << elossModel->GetdEdx(eIn) << ")"; } else { throw std::runtime_error("Cannot calculate process noise covariance without an energy loss model"); @@ -198,8 +201,8 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) XYZPoint startingPosition{m_vecX[0], m_vecX[1], m_vecX[2]}; // Get the starting position from the state vector Polar3DVector startingMomentum{m_vecX[3], m_vecX[4], m_vecX[5]}; // Get the starting momentum from the state vector - LOG(info) << "Propagating reference state from position: " << startingPosition - << " with momentum: " << XYZVector(startingMomentum); + LOG(debug) << "Propagating reference state from position: " << startingPosition + << " with momentum: " << XYZVector(startingMomentum); fPropagator.SetState(startingPosition, XYZVector(startingMomentum)); fPropagator.PropagateToMeasurementSurface(AtTools::AtMeasurementPoint(z), *fStepper); @@ -207,7 +210,7 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) fMeanStep.fLastPos = startingPosition; // Store the last position fMeanStep.fLastMom = startingMomentum; // Store the last momentum - LOG(info) << "Propagated to position: " << fMeanStep.fPos << " with momentum: " << fMeanStep.fMom; + LOG(debug) << "Propagated to position: " << fMeanStep.fPos << " with momentum: " << fMeanStep.fMom; // Now we can construct the reference plane. fMeasurementPlane = Plane3D(fMeanStep.fMom.Unit(), @@ -268,7 +271,7 @@ void TrackFitterUKF::smoothUKF() m_vecXSmooth.back() = m_vecXHist.back(); // The last smoothed state is the last corrected state m_matPSmooth.back() = m_matPHist.back(); // The last smoothed covariance is the last corrected covariance for (size_t i = m_vecXPredHist.size() - 1; i > 0; --i) { - LOG(info) << "Smoothing step " << i << " of " << m_vecXPredHist.size() - 1; + LOG(debug) << "Smoothing step " << i << " of " << m_vecXPredHist.size() - 1; // Get the predicted state and covariance at step i const auto &xPred = m_vecXPredHist[i]; // m_{k+1}^- From 534bfe10ea7feb3384e109e52f9e673f79a428f3 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 12:04:04 +0200 Subject: [PATCH 63/71] Add test for many tracks --- macro/tests/UKF/TestManyTracks.C | 200 +++++++++++++++++++++++++++++++ macro/tests/UKF/UKFSingleTrack.C | 9 +- 2 files changed, 202 insertions(+), 7 deletions(-) create mode 100644 macro/tests/UKF/TestManyTracks.C diff --git a/macro/tests/UKF/TestManyTracks.C b/macro/tests/UKF/TestManyTracks.C new file mode 100644 index 000000000..f025b74b1 --- /dev/null +++ b/macro/tests/UKF/TestManyTracks.C @@ -0,0 +1,200 @@ +std::string getEnergyPath() +{ + auto env = std::getenv("VMCWORKDIR"); + if (env == nullptr) { + return "../../resources/energy_loss/HinH.txt"; // Default path assuming cwd is build/AtTools + } + return std::string(env) + "/resources/energy_loss/HinH.txt"; // Use environment variable +} +using ROOT::Math::XYZPoint; +using ROOT::Math::XYZVector; + +const double mass_p = 938.272; // Mass of proton in MeV/c^2 +const double charge_p = 1.602176634e-19; // Charge of proton + +// Vectors to store the simulated points to compare to +std::vector x_sim, y_sim, z_sim, Eloss_sim; +TH1F *hMom = nullptr; +TH1F *hMomError = nullptr; +TCanvas *c1 = new TCanvas(); +TCanvas *c2 = new TCanvas(); + +// Parameters for model +const int pointsToCluster = 5; +const double sigma_pos = 1; // Position uncertainty of 10 mm +const double sigma_mom = 0.01; // Momentum uncertainty in percentage +const double sigma_theta = 1 * M_PI / 180; // Angular uncertainty of 1 degree +const double sigma_phi = 1 * M_PI / 180; // Angular uncertainty of 1 degree +const double gasDensity = 3.553e-5; // g/cm^3 +const double fAlpha = 1e-1; +const double fBeta = 2; +const double fKappa = 0; + +// Global variables +kf::TrackFitterUKF *ukf = nullptr; + +// Functions in file +/// Loads hits from an input file +void LoadHits(); +/// Run a single UKF for the passed initial state +void SingleUKF(XYZPoint initialPos, XYZVector initialMom, TMatrixD initialCov); +/// @brief Create the ukf pointer to reuse. +void CreateUKF(); + +TMatrixD CalculateInitialCov(double p); +TMatrixD CalculatePosCov(); + +/// @brief Test many tracks, saving stat properties +/// @param n +/// @param bias +void TestManyTracks(int n, double bias = 0) +{ + LoadHits(); + CreateUKF(); + + XYZPoint fTruePos(-3.40046e-04, -1.49863e-04, 1.0018); + XYZVector fTrueMom(0.00935463, -0.0454279, 0.00826042); + fTrueMom *= 1e3; + double fSigmaMom = fTrueMom.R() * sigma_mom; + + hMom = new TH1F("hMom", "Reconstructed Momentum (MeV/c)", 100, fTrueMom.R() - 4 * fSigmaMom, + fTrueMom.R() + 4 * fSigmaMom); + hMomError = + new TH1F("hMomError", "Error (%)", 100, -4 * fSigmaMom / fTrueMom.R() * 100, 4 * fSigmaMom / fTrueMom.R() * 100); + + for (int i = 0; i < n; ++i) { + + if (i % 100 == 0) + std::cout << "On iteration " << i << std::endl; + + double pSampled = gRandom->Gaus(fTrueMom.R(), sigma_mom * fTrueMom.R()); + pSampled += bias; + + // pSampled = fTrueMom.R(); + + ROOT::Math::Polar3DVector sampledMom(pSampled, fTrueMom.Theta(), fTrueMom.Phi()); + // try { + SingleUKF(fTruePos, XYZVector(sampledMom), CalculateInitialCov(pSampled)); + //} catch (...) { + // std::cerr << "Failed to propagate iteration " << i << " with seed momentum " << pSampled << std::endl; + // continue; + //} + + auto filtState = ukf->GetFilteredStates()[0]; + auto smoothState = ukf->GetSmoothedStates()[0]; + + double pReco = smoothState[3]; + hMom->Fill(pReco); + double error = (pReco - fTrueMom.R()) / fTrueMom.R() * 100; + hMomError->Fill(error); + + std::cout << std::endl + << std::endl + << "With initial momentum " << pSampled << " reconstructed " << pReco << " Error: " << error << "%" + << std::endl + << std::endl; + } + + // Draw results + c1->cd(); + hMom->Draw("hist"); + c2->cd(); + hMomError->Draw("hist"); +} + +/*********** Function implementations **************/ +void LoadHits() +{ + if (x_sim.size() != 0) + return; + Eloss_sim.clear(); + std::ifstream infile("hits.txt"); + double xi, yi, zi, Ei; + int i = 0; + double eLoss = 0; + + // Save first point. + infile >> xi >> yi >> zi >> Ei; + eLoss = Ei; // Initialize energy loss + x_sim.push_back(xi * 10); // Convert to mm + y_sim.push_back(yi * 10); // Convert to mm + z_sim.push_back(zi * 10); // Convert to mm + + while (infile >> xi >> yi >> zi >> Ei) { + // Ei *= 1e3; // Convert to MeV + + if (++i % pointsToCluster != 0) { + eLoss += Ei; + continue; // Skip every 5th point + } + x_sim.push_back(xi * 10); + y_sim.push_back(yi * 10); + z_sim.push_back(zi * 10); + Eloss_sim.push_back(eLoss); + eLoss = 0; // Reset energy loss for the next segment + } +} + +void CreateUKF() +{ + auto elossModel2 = std::make_unique(gasDensity); + elossModel2->SetProjectile(1, 1, 1); + std::vector> mat; + mat.push_back({1, 1, 1}); + elossModel2->SetMaterial(mat); + + auto elossModel = std::make_unique(0); + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path + elossModel->SetDensity(gasDensity); + + AtTools::AtPropagator propagator(charge_p, mass_p, std::move(elossModel)); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 2.85}); // Magnetic field + + // Setup stepper for UKF + auto stepper = std::make_unique(); + + // Setup UKF + ukf = new kf::TrackFitterUKF(std::move(propagator), std::move(stepper)); + ukf->setParameters(fAlpha, fBeta, fKappa); +} + +TMatrixD CalculateInitialCov(double p) +{ + TMatrixD cov(6, 6); + cov.Zero(); + for (int i = 0; i < 3; ++i) { + + cov(i, i) = sigma_pos * sigma_pos; // Set diagonal covariance to some small number + } + cov(3, 3) = p * p * sigma_mom * sigma_mom; // Momentum uncertainty + cov(4, 4) = sigma_theta * sigma_theta; // Angular uncertainty + cov(5, 5) = sigma_phi * sigma_phi; // Angular uncertainty + + return cov; +} +TMatrixD CalculatePosCov() +{ + TMatrixD cov(3, 3); + cov.Zero(); + for (int i = 0; i < 3; ++i) { + + cov(i, i) = sigma_pos * sigma_pos; // Set diagonal covariance to some small number + } + return cov; +} + +void SingleUKF(XYZPoint intitialPos, XYZVector initialMom, TMatrixD intitialCov) +{ + ukf->SetInitialState(intitialPos, initialMom, intitialCov); + + for (int i = 1; i < x_sim.size(); ++i) { + ukf->SetMeasCov(CalculatePosCov()); + XYZPoint point(x_sim[i], y_sim[i], z_sim[i]); + + ukf->predictUKF(point); + ukf->correctUKF(point); + } + + ukf->smoothUKF(); +} \ No newline at end of file diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index b524138ed..1c0ebf82e 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -12,7 +12,7 @@ const double charge_p = 1.602176634e-19; // Charge of proton // Simulated (measurement) hits std::vector x, y, z, Eloss; -int pointsToCluster = 5; +int pointsToCluster = 20; void LoadHits() { std::ifstream infile("hits.txt"); @@ -99,7 +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-1, 2, 0); // Set kappa to match the original implementation + ukf.setParameters(1e-1, 2, 0); // alpha, beta, kappa TMatrixD cov(6, 6); cov.Zero(); @@ -141,9 +141,7 @@ void UKFSingleTrack() ukf.predictUKF(point); auto augState = ukf.GetAugStateVector(); auto augCov = ukf.GetAugStateCovariance(); - std::cout << std::endl << "Prediction step complete." << std::endl; ukf.correctUKF(point); - std::cout << std::endl << "Correction step complete." << std::endl; auto state = ukf.GetStateVector(); auto cov = ukf.GetStateCovariance(); @@ -154,7 +152,6 @@ void UKFSingleTrack() std::cout << "Predicted position: " << pos << std::endl; std::cout << "Predicted momentum: " << mom << std::endl; - std::cout << "Measurement point: " << point << std::endl; auto KE_in = Kinematics::KE(lastMom, mass_p); @@ -183,8 +180,6 @@ void UKFSingleTrack() for (int i = 0; i < smoothedStates.size(); ++i) { auto &state = smoothedStates[i]; auto &filteredState = filteredStates[i]; - std::cout << "Smoothed state " << i << ": " << state.transpose() << std::endl; - std::cout << "Filtered state " << i << ": " << filteredState.transpose() << std::endl; xSmooth.push_back(state[0]); ySmooth.push_back(state[1]); zSmooth.push_back(state[2]); From 59f4cb40c9b044297d36d6e142ac80da144d5f57 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 08:33:28 -0400 Subject: [PATCH 64/71] Add diagnostic and add COV correction --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 8 +++++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 29 ++++++++++++++++--- macro/tests/UKF/UKFSingleTrack.C | 19 ++++++++++-- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 3fbb31e27..62fe27ba2 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -63,6 +63,9 @@ void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); // Calculate the sigma points for the initial state // Now we grab the sigma points only for the state. m_matSigmaXPred = m_matSigmaXa.block(0, 0, TF_DIM_X, SIGMA_DIM_A); // Extract the state sigma points + + logEigen("Initial cov", m_matP, 0); // Log the eigenvalues of the initial covariance matrix + LOG(info) << "Initial COV:" << std::endl << m_matP; } TMatrixD TrackFitterUKF::GetStateCovariance() const @@ -235,6 +238,8 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) TrackFitterUKFBase::predictUKF(callback, zVec); + // logEigen("State P-", m_matP, m_matPPredHist.size()); + // Now we need to store the predicted state and covariance for smoothing later. m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector m_matPPredHist.push_back(m_matP); // Store the predicted covariance matrix @@ -256,6 +261,8 @@ void TrackFitterUKF::correctUKF(const ROOT::Math::XYZPoint &z) auto callback = [this](const kf::Vector &x_) { return funcH(x_); }; TrackFitterUKFBase::correctUKF(callback, zVec); + // logEigen("State PCorr", m_matP, m_matPHist.size()); + // After correction we need to save the filtered state m_vecXHist.push_back(m_vecX); // Store the filtered state vector m_matPHist.push_back(m_matP); // Store the filtered covariance matrix @@ -294,6 +301,7 @@ void TrackFitterUKF::smoothUKF() m_vecXSmooth[i - 1] = xFilt + D * (xSmooth - xPred); // m^s_{k} = m_{k} + D * (m^s_{k+1} - m_{k+1}^-) m_matPSmooth[i - 1] = pFilt + D * (pSmooth - pPred) * D.transpose(); // P^s_{k} = P_{k} + D * (P^s_{k+1} - P_{k+1}^-) * D^T + logEigen("State P+", m_matPSmooth[i - 1], i - 1); } } diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index e5cacd4fb..ba728cde3 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -49,6 +49,7 @@ class TrackFitterUKFBase { static constexpr int32_t SIGMA_DIM_A{2 * DIM_A + 1}; ///< @brief Sigma points dimension for augmented state Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix + int nTouch = 0; // Variable to track the number of times a matrix has a floor added. protected: Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector Matrix m_matP{Matrix::Zero()}; /// @brief state covariance matrix @@ -163,7 +164,12 @@ class TrackFitterUKFBase { // Calculate the weighted mean and covariance of the sigma points for the state vector. // This will be the new state vector and covariance matrix. - calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP); + calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP, "P-"); + m_matP(0, 0) += 1e-4; + m_matP(1, 1) += 1e-4; + m_matP(2, 2) += 1e-4; + ensurePD(m_matP); // Ensure the covariance matrix is positive definite + logEigen("P-Post", m_matP, 0); } /** @@ -197,10 +203,11 @@ class TrackFitterUKFBase { // from the sigma points. Vector vecZhat; // Predicted measurement vector Matrix matPzz; // Measurement covariance matrix - calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz); + calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz, "SnoR"); // Add in the measurement noise covariance matrix to the measurement covariance matrix. matPzz += m_matR; // Add measurement noise covariance so we gen the innovation covariance matrix. + logEigen("S", matPzz, 0); ensurePD(matPzz); // Ensure the covariance matrix is positive definite const Matrix matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)}; @@ -217,6 +224,18 @@ class TrackFitterUKFBase { ensurePD(m_matP); // Ensure the covariance matrix is positive definite } + template + void logEigen(std::string tag, const T &P, int k) + { + Eigen::SelfAdjointEigenSolver es(P); + double lmin = es.eigenvalues().minCoeff(); + double lmax = es.eigenvalues().maxCoeff(); + double cond = lmax / lmin; + + LOG(info) << "k: " << k << " " << tag << " Eval: min = " << lmin << ", max = " << lmax + << ", condition number = " << cond; + } + protected: /** * @brief Add state vector and state covariance matrix to the augmented state vector covariance matrix. @@ -295,6 +314,7 @@ class TrackFitterUKFBase { matP += Matrix::Identity() * std::pow(10, -6 + i); // Regularization value lltOfP.compute(matP); i = i + 1; + nTouch++; // Check if the regularized matrix is now positive definite } if (lltOfP.info() != Eigen::Success) { @@ -367,7 +387,7 @@ class TrackFitterUKFBase { */ template void calculateWeightedMeanAndCovariance(const Matrix &sigmaX, Vector &vecX, - Matrix &matPxx) + Matrix &matPxx, std::string tag = "") { // 1. calculate mean of the sigma points vecX = m_weightM0 * util::getColumnAt(0, sigmaX); @@ -390,7 +410,8 @@ class TrackFitterUKFBase { matPxx += Pi; // y += W[0, i] (Y[:, i] - \bar{y}) (Y[:, i] - \bar{y})^T } - ensurePD(matPxx); // Ensure the covariance matrix is positive definite + logEigen(tag, matPxx, 0); // Log the eigenvalues of the covariance matrix + // ensurePD(matPxx); // Ensure the covariance matrix is positive definite } /** diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 1c0ebf82e..32c09d472 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -12,7 +12,7 @@ const double charge_p = 1.602176634e-19; // Charge of proton // Simulated (measurement) hits std::vector x, y, z, Eloss; -int pointsToCluster = 20; +int pointsToCluster = 5; void LoadHits() { std::ifstream infile("hits.txt"); @@ -99,7 +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-1, 2, 0); // alpha, beta, kappa + ukf.setParameters(1e-2, 2, 0); // alpha, beta, kappa TMatrixD cov(6, 6); cov.Zero(); @@ -141,6 +141,19 @@ void UKFSingleTrack() ukf.predictUKF(point); auto augState = ukf.GetAugStateVector(); auto augCov = ukf.GetAugStateCovariance(); + auto covP = ukf.matP(); + + if (i == 1) { + LOG(info) << "P-: " << std::endl << covP; + Eigen::SelfAdjointEigenSolver es(covP); // float version + float λmin = es.eigenvalues()(0); + auto vmin = es.eigenvectors().col(0); // length-6 + + std::printf("λmin = %.3e eigvec = [", λmin); + for (int i = 0; i < 6; ++i) + std::printf(" %.2e", vmin(i)); + std::printf(" ]\n"); + } ukf.correctUKF(point); auto state = ukf.GetStateVector(); @@ -171,6 +184,8 @@ void UKFSingleTrack() residual.push_back(residualValue); // Store the residual for this hit } + LOG(info) << "After forward pass " << ukf.nTouch << " touches."; + // At this point we have the full trajectory of the particle ukf.smoothUKF(); // Perform smoothing auto smoothedStates = ukf.GetSmoothedStates(); From 596d81b435fd8021c64b38bd7e71b1e922e3e609 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 08:43:57 -0400 Subject: [PATCH 65/71] Add controllable diagnostics --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index ba728cde3..fc59e26dc 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -49,7 +49,10 @@ class TrackFitterUKFBase { static constexpr int32_t SIGMA_DIM_A{2 * DIM_A + 1}; ///< @brief Sigma points dimension for augmented state Matrix m_matSigmaXa{Matrix::Zero()}; ///< @brief Sigma points matrix - int nTouch = 0; // Variable to track the number of times a matrix has a floor added. + // Controls and variables for running numerical diagnostics. + int nTouch{0}; // Variable to track the number of times a matrix has a floor added. + bool kLogEigen{false}; + protected: Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector Matrix m_matP{Matrix::Zero()}; /// @brief state covariance matrix @@ -227,6 +230,9 @@ class TrackFitterUKFBase { template void logEigen(std::string tag, const T &P, int k) { + if (!kLogEigen) { + return; // If logging is disabled, do not log the eigenvalues. + } Eigen::SelfAdjointEigenSolver es(P); double lmin = es.eigenvalues().minCoeff(); double lmax = es.eigenvalues().maxCoeff(); From f14be203b9aa0bb8f132110ad51ff396a26c1823 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 09:24:02 -0400 Subject: [PATCH 66/71] Rename noises --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 2 +- .../OpenKF/kalman_filter/TrackFitterUKF.h | 27 ++++++++++++++----- .../kalman_filter/TrackFitterUKFTest.cxx | 8 +++--- AtReconstruction/AtFitter/OpenKF/types.h | 2 +- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 62fe27ba2..ffa1bde15 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -15,7 +15,7 @@ void TrackFitterUKF::Reset() m_matP.setZero(); m_vecXa.setZero(); m_matPa.setZero(); - m_matQ.setZero(); + m_matQaug.setZero(); m_matR.setZero(); m_matSigmaXa.setZero(); m_matSigmaXPred.setZero(); diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index fc59e26dc..a4508e876 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -32,6 +32,11 @@ namespace kf { * the machinery that underlies the UKF formalism. It is a modified version of the UKF provided * by OpenKF, that has been expanded to allow for more hooks into the method. * + * Introduces two different types of process noise. We have augmented process noise and + * model process noise. Augmented process noise acts during the propagator, model process + * noise can be used to describe noise in the propagation itself and allows for accounting + * for uncertainty that may arise from unmodeled dynamics. + * * Templated because I believe Eigen can do quite a bit of operation for small matrices like we have here if * the size is known at compile time. Worth checking that though. * @@ -63,8 +68,11 @@ class TrackFitterUKFBase { Matrix m_matPa{Matrix::Zero()}; // Process and measurement noise covariance matrices - /// Process noise covariance matrix (Q) - Matrix m_matQ{Matrix::Zero()}; + /// Process noise covariance matrix incorporated in propagator (Q_aug) + Matrix m_matQaug{Matrix::Zero()}; + /// Process noise covariance matrix for model noise (Q_mod) + Matrix m_matQmod{Matrix::Zero()}; + /// Measurement noise covariance matrix (R) Matrix m_matR{Matrix::Zero()}; @@ -102,11 +110,16 @@ class TrackFitterUKFBase { /** * @brief Set process noise covariance Q to be used in the prediction step. */ - void setCovarianceQ(const Matrix &matQ) { m_matQ = matQ; } + void setAugmentNoise(const Matrix &matQ) { m_matQaug = matQ; } + /** + * @brief Set the process noise covariance Q_mod to be used in the prediction step. + */ + void setModelNoise(const Matrix &matQmod) { m_matQmod = matQmod; } + /** * @brief Set the measurement noise covariance R to be used in the update step. */ - void setCovarianceR(const Matrix &matR) { m_matR = matR; } + void SetMeasurementNoise(const Matrix &matR) { m_matR = matR; } virtual Vector &vecX() { return m_vecX; } virtual const Vector &vecX() const { return m_vecX; } @@ -263,12 +276,12 @@ class TrackFitterUKFBase { virtual std::array calculateProcessNoiseMean() { return std::array{0}; } - virtual Matrix calculateProcessNoiseCovariance() { return m_matQ; } + virtual Matrix calculateProcessNoiseCovariance() { return m_matQaug; } void updateAugWithProcessNoise() { auto processNoiseMean = calculateProcessNoiseMean(); - m_matQ = calculateProcessNoiseCovariance(); + m_matQaug = calculateProcessNoiseCovariance(); // Add the mean process noise to the augmented state vector for (int32_t i{0}; i < DIM_V; ++i) { @@ -281,7 +294,7 @@ class TrackFitterUKFBase { for (int32_t i{S_IDX}; i < L_IDX; ++i) { for (int32_t j{S_IDX}; j < L_IDX; ++j) { - m_matPa(i, j) = m_matQ(i - S_IDX, j - S_IDX); + m_matPa(i, j) = m_matQaug(i - S_IDX, j - S_IDX); } } } diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx index 8ce7dfc3d..393cdca24 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -75,8 +75,8 @@ TEST_F(TrackFitterUKFExampleTest, Prediction) m_ukf.vecX() = x; m_ukf.matP() = P; - m_ukf.setCovarianceQ(Q); - m_ukf.setCovarianceR(R); + m_ukf.setAugmentNoise(Q); + m_ukf.SetMeasurementNoise(R); m_ukf.setParameters(1, 0, 3 - m_ukf.DIM_A); // Set kappa to match the original implementation m_ukf.predictUKF(funcF, z); @@ -137,8 +137,8 @@ TEST_F(TrackFitterUKFExampleTest, PredictionAndCorrection) m_ukf.vecX() = x; m_ukf.matP() = P; - m_ukf.setCovarianceQ(Q); - m_ukf.setCovarianceR(R); + m_ukf.setAugmentNoise(Q); + m_ukf.SetMeasurementNoise(R); m_ukf.setParameters(1, 0, 3 - m_ukf.DIM_A); // Set kappa to match the original implementation m_ukf.predictUKF(funcF, z); diff --git a/AtReconstruction/AtFitter/OpenKF/types.h b/AtReconstruction/AtFitter/OpenKF/types.h index 215f8fc79..4111fb8dc 100644 --- a/AtReconstruction/AtFitter/OpenKF/types.h +++ b/AtReconstruction/AtFitter/OpenKF/types.h @@ -17,7 +17,7 @@ #include namespace kf { -using float32_t = float; +using float32_t = double; template using Matrix = Eigen::Matrix; From e94e3f0e292498c6978f1888d0157a50f4ed338b Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 09:55:33 -0400 Subject: [PATCH 67/71] More refactoring --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 17 +---- .../OpenKF/kalman_filter/TrackFitterUKF.h | 70 +++++++++++-------- macro/tests/UKF/UKFSingleTrack.C | 2 +- 3 files changed, 45 insertions(+), 44 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index ffa1bde15..cd5444f73 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -8,24 +8,11 @@ #include namespace kf { + void TrackFitterUKF::Reset() { // Reset the state vector and covariance matrix - m_vecX.setZero(); - m_matP.setZero(); - m_vecXa.setZero(); - m_matPa.setZero(); - m_matQaug.setZero(); - m_matR.setZero(); - m_matSigmaXa.setZero(); - m_matSigmaXPred.setZero(); - - // Clear the history vectors - m_vecXPredHist.clear(); - m_matPPredHist.clear(); - m_matCPredHist.clear(); - m_vecXHist.clear(); - m_matPHist.clear(); + TrackFitterUKFBase::Reset(); m_vecXSmooth.clear(); m_matPSmooth.clear(); fMeanStep = AtTools::AtPropagator::StepState(); // Reset the step state diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index a4508e876..a86f89b06 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -59,6 +59,9 @@ class TrackFitterUKFBase { bool kLogEigen{false}; protected: + using VectorEigenVecDimX = std::vector, Eigen::aligned_allocator>>; + using VectorEigenMatDimX = std::vector, Eigen::aligned_allocator>>; + Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector Matrix m_matP{Matrix::Zero()}; /// @brief state covariance matrix @@ -67,12 +70,10 @@ class TrackFitterUKFBase { /// @brief augmented state covariance (incl. process and measurement noise covariances) Matrix m_matPa{Matrix::Zero()}; - // Process and measurement noise covariance matrices /// Process noise covariance matrix incorporated in propagator (Q_aug) Matrix m_matQaug{Matrix::Zero()}; /// Process noise covariance matrix for model noise (Q_mod) Matrix m_matQmod{Matrix::Zero()}; - /// Measurement noise covariance matrix (R) Matrix m_matR{Matrix::Zero()}; @@ -85,10 +86,42 @@ class TrackFitterUKFBase { /// Lambda parameter for sigma point calculation float32_t m_lambda{0}; + // Track state over time (can be used for smoothing or diagnostics) + VectorEigenVecDimX m_vecXPredHist; /// @brief History of predicted state vectors at k+1 + VectorEigenMatDimX m_matPPredHist; /// @brief History of predicted state covariances at k+1 + /// History of cross correlation between filtered state at k and predicted at k+1 + VectorEigenMatDimX m_matCPredHist; + /// History of filtered (after correction) state vectors at k + VectorEigenVecDimX m_vecXHist; + /// History of filtered (after correction) state covariances at k + VectorEigenMatDimX m_matPHist; + /// The sigma points after propagation for the last prediction step. + Matrix m_matSigmaXPred{Matrix::Zero()}; + public: TrackFitterUKFBase() { setParameters(1, 2, 0); } ~TrackFitterUKFBase() = default; + virtual void Reset() + { + m_vecX.setZero(); + m_matP.setZero(); + m_vecXa.setZero(); + m_matPa.setZero(); + m_matQaug.setZero(); + m_matQmod.setZero(); + m_matR.setZero(); + m_matSigmaXa.setZero(); + m_matSigmaXPred.setZero(); + + // Clear the history vectors + m_vecXPredHist.clear(); + m_matPPredHist.clear(); + m_matCPredHist.clear(); + m_vecXHist.clear(); + m_matPHist.clear(); + } + /** * @brief Set the weights used to calculate sigma points. */ @@ -136,7 +169,7 @@ class TrackFitterUKFBase { { updateAugWithState(); updateAugWithProcessNoise(); - ensurePD(m_matPa); // Ensure the augmented covariance matrix is positive definite + // ensurePD(m_matPa); // Ensure the augmented covariance matrix is positive definite } /** @@ -478,34 +511,18 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { static constexpr int32_t TF_DIM_Z = 3; static constexpr int32_t TF_DIM_V = 1; static constexpr int32_t TF_DIM_N = 3; + // using DIM_X = TrackFitterUKFBase::DIM_X; AtTools::AtPropagator fPropagator; ///< @brief Propagator for the track fitter std::unique_ptr fStepper{nullptr}; AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter - using EigenVectorDimX = std::vector, Eigen::aligned_allocator>>; - using VectorEigenMatDimX = - std::vector, Eigen::aligned_allocator>>; - - // 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 - /// History of cross correlation between filtered state at k and predicted at k+1 - VectorEigenMatDimX m_matCPredHist; - /// History of filtered (after correction) state vectors at k - EigenVectorDimX m_vecXHist; - /// History of filtered (after correction) state covariances at k - VectorEigenMatDimX m_matPHist; - /// Smoothed state vector and covariance - EigenVectorDimX m_vecXSmooth; + VectorEigenVecDimX m_vecXSmooth; /// Smoothed state covariance VectorEigenMatDimX m_matPSmooth; - /// The sigma points after propagation for the last prediction step. - Matrix m_matSigmaXPred{Matrix::Zero()}; - public: bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling double fMaxStragglingFactor{1. / 3.}; ///< @brief Maximum straggling factor for energy loss @@ -524,21 +541,18 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { : TrackFitterUKFBase(), fPropagator(std::move(propagator)), fStepper(std::move(stepper)) { } - void Reset(); + virtual void Reset() override; void SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, const TMatrixD &initialCovariance); void SetMeasCov(const TMatrixD &measCov); - std::array GetStateVector() const - { - return {m_vecX[0], m_vecX[1], m_vecX[2], m_vecX[3], m_vecX[4], m_vecX[5]}; - } + TMatrixD GetStateCovariance() const; TMatrixD GetAugStateCovariance() const; std::array GetAugStateVector() const; - const EigenVectorDimX &GetSmoothedStates() const { return m_vecXSmooth; }; + const VectorEigenVecDimX &GetSmoothedStates() const { return m_vecXSmooth; }; const VectorEigenMatDimX &GetSmoothedCovariances() const { return m_matPSmooth; }; - const EigenVectorDimX &GetFilteredStates() const { return m_vecXHist; }; + const VectorEigenVecDimX &GetFilteredStates() const { return m_vecXHist; }; const VectorEigenMatDimX &GetFilteredCovariances() const { return m_matPHist; }; void predictUKF(const ROOT::Math::XYZPoint &z); diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 32c09d472..5be40ab4c 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -99,7 +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); // alpha, beta, kappa + ukf.setParameters(1e-3, 2, 0); // alpha, beta, kappa TMatrixD cov(6, 6); cov.Zero(); From ce22ed31d776e9348f9f12d3fa003234537b676c Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 10:12:59 -0400 Subject: [PATCH 68/71] More refactoring --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 49 +++++-------------- .../OpenKF/kalman_filter/TrackFitterUKF.h | 45 ++++++++++++----- macro/tests/UKF/UKFSingleTrack.C | 4 +- 3 files changed, 47 insertions(+), 51 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index cd5444f73..13dddea8e 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -15,6 +15,7 @@ void TrackFitterUKF::Reset() TrackFitterUKFBase::Reset(); m_vecXSmooth.clear(); m_matPSmooth.clear(); + m_matCPredHist.clear(); fMeanStep = AtTools::AtPropagator::StepState(); // Reset the step state } @@ -39,8 +40,8 @@ void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition } // Save the initial state in our history vectors - m_vecXHist.push_back(m_vecX); - m_matPHist.push_back(m_matP); + m_vecXFiltHist.push_back(m_vecX); + m_matPFiltHist.push_back(m_matP); m_vecXPredHist.push_back(m_vecX); m_matPPredHist.push_back(m_matP); m_matCPredHist.push_back(Matrix::Zero()); // Cross-correlation is not defined for the first point @@ -170,15 +171,11 @@ Vector TrackFitterUKF::funcF(const Vector TrackFitterUKF::funcH(const Vector &x) { - // Measurement function to convert state vector to measurement vector - using namespace ROOT::Math; Vector vecZ; - XYZPoint fPos(x[0], x[1], x[2]); // Position from state vector - // Calculate the measurement vector based on the position and momentum - vecZ[0] = fPos.X(); // X coordinate - vecZ[1] = fPos.Y(); // Y coordinate - vecZ[2] = fPos.Z(); // Z coordinate + vecZ[0] = x[0]; // X coordinate + vecZ[1] = x[1]; // Y coordinate + vecZ[2] = x[2]; // Z coordinate return vecZ; // Return the measurement vector } @@ -212,30 +209,14 @@ void TrackFitterUKF::predictUKF(const ROOT::Math::XYZPoint &z) auto callback = [this](const kf::Vector &x_, const kf::Vector &v_, const kf::Vector &z_) { return funcF(x_, v_, z_); }; - updateAugmentedStateAndCovariance(); - - // Calculate the sigma points for the augmented state vector and save in a matrix where each column is a sigma - // point. - auto sigmaPoints = calculateSigmaPoints(m_vecXa, m_matPa); - - // Pull out the sigma points for the state vector and process noise in two different matrices. - Matrix sigmaXxPrior{ - sigmaPoints.block(0, 0, TF_DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector - m_matSigmaXPred = sigmaXxPrior; // Store the sigma points for the state vector - TrackFitterUKFBase::predictUKF(callback, zVec); - // logEigen("State P-", m_matP, m_matPPredHist.size()); - - // Now we need to store the predicted state and covariance for smoothing later. - m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector - m_matPPredHist.push_back(m_matP); // Store the predicted covariance matrix - // Get the sigma points belonging to the predicted state Matrix sigmaXx{m_matSigmaXa.block(0, 0, TF_DIM_X, SIGMA_DIM_A)}; // Calculate the cross-corelation between the filtered state at k and predicted state at k+1 - auto matCPred = calculateCrossCorrelation(sigmaXxPrior, m_vecXHist.back(), sigmaXx, m_vecXPredHist.back()); + auto matCPred = + calculateCrossCorrelation(m_matSigmaXPred, m_vecXFiltHist.back(), sigmaXx, m_vecXPredHist.back()); m_matCPredHist.push_back(matCPred); // Store the cross-correlation matrix } @@ -248,11 +229,7 @@ void TrackFitterUKF::correctUKF(const ROOT::Math::XYZPoint &z) auto callback = [this](const kf::Vector &x_) { return funcH(x_); }; TrackFitterUKFBase::correctUKF(callback, zVec); - // logEigen("State PCorr", m_matP, m_matPHist.size()); - - // After correction we need to save the filtered state - m_vecXHist.push_back(m_vecX); // Store the filtered state vector - m_matPHist.push_back(m_matP); // Store the filtered covariance matrix + // logEigen("State PCorr", m_matP, m_matPCorrHist.size()); } void TrackFitterUKF::smoothUKF() @@ -262,8 +239,8 @@ void TrackFitterUKF::smoothUKF() // Here i = k+1 m_vecXSmooth.resize(m_vecXPredHist.size()); m_matPSmooth.resize(m_matPPredHist.size()); - m_vecXSmooth.back() = m_vecXHist.back(); // The last smoothed state is the last corrected state - m_matPSmooth.back() = m_matPHist.back(); // The last smoothed covariance is the last corrected covariance + m_vecXSmooth.back() = m_vecXFiltHist.back(); // The last smoothed state is the last corrected state + m_matPSmooth.back() = m_matPFiltHist.back(); // The last smoothed covariance is the last corrected covariance for (size_t i = m_vecXPredHist.size() - 1; i > 0; --i) { LOG(debug) << "Smoothing step " << i << " of " << m_vecXPredHist.size() - 1; @@ -273,8 +250,8 @@ void TrackFitterUKF::smoothUKF() const auto &ccor = m_matCPredHist[i]; // C_{k+1} // Get the filtered state and covariance at step i-1 - const auto &xFilt = m_vecXHist[i - 1]; // m_{k} - const auto &pFilt = m_matPHist[i - 1]; // P_{k} + const auto &xFilt = m_vecXFiltHist[i - 1]; // m_{k} + const auto &pFilt = m_matPFiltHist[i - 1]; // P_{k} // Get the smoothed state and covariance at step i auto &xSmooth = m_vecXSmooth[i]; // m^s_{k+1} diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index a86f89b06..f544f44f8 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -60,6 +60,7 @@ class TrackFitterUKFBase { protected: using VectorEigenVecDimX = std::vector, Eigen::aligned_allocator>>; + using VectorEigenVecDimZ = std::vector, Eigen::aligned_allocator>>; using VectorEigenMatDimX = std::vector, Eigen::aligned_allocator>>; Vector m_vecX{Vector::Zero()}; /// @brief estimated state vector @@ -89,12 +90,12 @@ class TrackFitterUKFBase { // Track state over time (can be used for smoothing or diagnostics) VectorEigenVecDimX m_vecXPredHist; /// @brief History of predicted state vectors at k+1 VectorEigenMatDimX m_matPPredHist; /// @brief History of predicted state covariances at k+1 - /// History of cross correlation between filtered state at k and predicted at k+1 - VectorEigenMatDimX m_matCPredHist; /// History of filtered (after correction) state vectors at k - VectorEigenVecDimX m_vecXHist; + VectorEigenVecDimX m_vecXFiltHist; /// History of filtered (after correction) state covariances at k - VectorEigenMatDimX m_matPHist; + VectorEigenMatDimX m_matPFiltHist; + /// History of measurement points + VectorEigenVecDimZ m_vecZHist; /// The sigma points after propagation for the last prediction step. Matrix m_matSigmaXPred{Matrix::Zero()}; @@ -117,9 +118,8 @@ class TrackFitterUKFBase { // Clear the history vectors m_vecXPredHist.clear(); m_matPPredHist.clear(); - m_matCPredHist.clear(); - m_vecXHist.clear(); - m_matPHist.clear(); + m_vecXFiltHist.clear(); + m_matPFiltHist.clear(); } /** @@ -158,6 +158,10 @@ class TrackFitterUKFBase { virtual Matrix &matP() { return m_matP; } virtual const Matrix &matP() const { return m_matP; } + const VectorEigenVecDimX &GetFilteredStates() const { return m_vecXFiltHist; } + const VectorEigenMatDimX &GetFilteredCovariances() const { return m_matPFiltHist; } + const VectorEigenVecDimX &GetPredictedStates() const { return m_vecXPredHist; } + const VectorEigenMatDimX &GetPredictedCovariances() const { return m_matPPredHist; } /** * @brief update the augmented state vector and covariance matrix @@ -192,6 +196,9 @@ class TrackFitterUKFBase { // point. m_matSigmaXa = calculateSigmaPoints(m_vecXa, m_matPa); + // Pull out the sigma points for the state vector + m_matSigmaXPred = m_matSigmaXa.block(0, 0, DIM_X, SIGMA_DIM_A); + // Pull out the sigma points for the state vector and process noise in two different matrices. Matrix sigmaXx{m_matSigmaXa.block(0, 0, DIM_X, SIGMA_DIM_A)}; // Sigma points for state vector Matrix sigmaXv{ @@ -213,11 +220,20 @@ class TrackFitterUKFBase { // Calculate the weighted mean and covariance of the sigma points for the state vector. // This will be the new state vector and covariance matrix. - calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP, "P-"); + calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP); + logEigen("P-", m_matP, 0); // Log the eigenvalues of the covariance matrix + + // Here we add process noise m_matP(0, 0) += 1e-4; m_matP(1, 1) += 1e-4; m_matP(2, 2) += 1e-4; + ensurePD(m_matP); // Ensure the covariance matrix is positive definite + + // Now we need to store the predicted state and covariance for smoothing later. + m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector + m_matPPredHist.push_back(m_matP); // Store the predicted covariance matrix + logEigen("P-Post", m_matP, 0); } @@ -252,7 +268,7 @@ class TrackFitterUKFBase { // from the sigma points. Vector vecZhat; // Predicted measurement vector Matrix matPzz; // Measurement covariance matrix - calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz, "SnoR"); + calculateWeightedMeanAndCovariance(sigmaZ, vecZhat, matPzz); // Add in the measurement noise covariance matrix to the measurement covariance matrix. matPzz += m_matR; // Add measurement noise covariance so we gen the innovation covariance matrix. @@ -271,6 +287,10 @@ class TrackFitterUKFBase { m_matP -= matK * matPzz * matK.transpose(); // m_matP -= matPxz * matK.transpose(); ensurePD(m_matP); // Ensure the covariance matrix is positive definite + + // After correction we need to save the filtered state + m_vecXFiltHist.push_back(m_vecX); // Store the filtered state vector + m_matPFiltHist.push_back(m_matP); // Store the filtered covariance matrix } template @@ -439,7 +459,7 @@ class TrackFitterUKFBase { */ template void calculateWeightedMeanAndCovariance(const Matrix &sigmaX, Vector &vecX, - Matrix &matPxx, std::string tag = "") + Matrix &matPxx) { // 1. calculate mean of the sigma points vecX = m_weightM0 * util::getColumnAt(0, sigmaX); @@ -462,7 +482,6 @@ class TrackFitterUKFBase { matPxx += Pi; // y += W[0, i] (Y[:, i] - \bar{y}) (Y[:, i] - \bar{y})^T } - logEigen(tag, matPxx, 0); // Log the eigenvalues of the covariance matrix // ensurePD(matPxx); // Ensure the covariance matrix is positive definite } @@ -522,6 +541,8 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { VectorEigenVecDimX m_vecXSmooth; /// Smoothed state covariance VectorEigenMatDimX m_matPSmooth; + /// History of cross correlation between filtered state at k and predicted at k+1 + VectorEigenMatDimX m_matCPredHist; public: bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling @@ -552,8 +573,6 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { std::array GetAugStateVector() const; const VectorEigenVecDimX &GetSmoothedStates() const { return m_vecXSmooth; }; const VectorEigenMatDimX &GetSmoothedCovariances() const { return m_matPSmooth; }; - const VectorEigenVecDimX &GetFilteredStates() const { return m_vecXHist; }; - const VectorEigenMatDimX &GetFilteredCovariances() const { return m_matPHist; }; void predictUKF(const ROOT::Math::XYZPoint &z); void correctUKF(const ROOT::Math::XYZPoint &z); diff --git a/macro/tests/UKF/UKFSingleTrack.C b/macro/tests/UKF/UKFSingleTrack.C index 5be40ab4c..99886f9ac 100644 --- a/macro/tests/UKF/UKFSingleTrack.C +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -156,8 +156,8 @@ void UKFSingleTrack() } ukf.correctUKF(point); - auto state = ukf.GetStateVector(); - auto cov = ukf.GetStateCovariance(); + auto state = ukf.vecX(); + auto cov = ukf.matP(); ROOT::Math::XYZPoint pos(state[0], state[1], state[2]); ROOT::Math::Polar3DVector momPolar(state[3], state[4], state[5]); From 75ab27232e0494b0e8ffdfa7ee4d83a06477e7b9 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 11:02:18 -0400 Subject: [PATCH 69/71] Remove inverse calls --- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx | 4 ++-- .../AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h | 7 +++---- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index 13dddea8e..f006798af 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -258,8 +258,8 @@ void TrackFitterUKF::smoothUKF() auto &pSmooth = m_matPSmooth[i]; // P^s_{k+1} auto llt = calculateCholesky(pPred); - // auto D = ccor * llt.solve(Matrix::Identity()); - auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} + auto D = ccor * llt.solve(Matrix::Identity()); + // auto D = ccor * pPred.inverse(); // D = C_{k+1} * (P_{k+1}^-)^{-1} // std::cout << "D matrix at step " << i << ":\n" << D << "\n"; m_vecXSmooth[i - 1] = xFilt + D * (xSmooth - xPred); // m^s_{k} = m_{k} + D * (m^s_{k+1} - m_{k+1}^-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index f544f44f8..05ce1b95a 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -279,9 +279,9 @@ class TrackFitterUKFBase { // kalman gain auto llt = calculateCholesky(matPzz); - // const Matrix matK = llt.solve(matPxz.transpose()).transpose(); - Matrix matK = {matPxz * llt.solve(Matrix::Identity())}; - // Matrix matK = matPxz * matPzz.inverse(); By far the worst method for filter stability + const Matrix matK = llt.solve(matPxz.transpose()).transpose(); + // Matrix matK = {matPxz * llt.solve(Matrix::Identity())}; + // Matrix matK = matPxz * matPzz.inverse(); By far the worst method for filter stability m_vecX += matK * (vecZ - vecZhat); m_matP -= matK * matPzz * matK.transpose(); @@ -530,7 +530,6 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { static constexpr int32_t TF_DIM_Z = 3; static constexpr int32_t TF_DIM_V = 1; static constexpr int32_t TF_DIM_N = 3; - // using DIM_X = TrackFitterUKFBase::DIM_X; AtTools::AtPropagator fPropagator; ///< @brief Propagator for the track fitter std::unique_ptr fStepper{nullptr}; From 4caa6010ff6e87ff845b749abd29d375ff609712 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 11:19:28 -0400 Subject: [PATCH 70/71] Incorporate model noise --- .../OpenKF/kalman_filter/TrackFitterUKF.cxx | 4 ++++ .../OpenKF/kalman_filter/TrackFitterUKF.h | 16 +++++++++------- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx index f006798af..6ee181b29 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -39,6 +39,10 @@ void TrackFitterUKF::SetInitialState(const ROOT::Math::XYZPoint &initialPosition } } + for (int i = 0; i < m_matQmod.rows(); ++i) { + m_matQmod(i, i) = fPosModelNoise; // Initialize model noise covariance to zero + } + // Save the initial state in our history vectors m_vecXFiltHist.push_back(m_vecX); m_matPFiltHist.push_back(m_matP); diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h index 05ce1b95a..b48b90852 100644 --- a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -222,13 +222,8 @@ class TrackFitterUKFBase { // This will be the new state vector and covariance matrix. calculateWeightedMeanAndCovariance(sigmaXx, m_vecX, m_matP); logEigen("P-", m_matP, 0); // Log the eigenvalues of the covariance matrix - - // Here we add process noise - m_matP(0, 0) += 1e-4; - m_matP(1, 1) += 1e-4; - m_matP(2, 2) += 1e-4; - - ensurePD(m_matP); // Ensure the covariance matrix is positive definite + m_matP += m_matQmod; // Add the model process noise covariance to the state covariance matrix. + ensurePD(m_matP); // Ensure the covariance matrix is positive definite // Now we need to store the predicted state and covariance for smoothing later. m_vecXPredHist.push_back(m_vecX); // Store the predicted state vector @@ -546,6 +541,13 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> { public: bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling double fMaxStragglingFactor{1. / 3.}; ///< @brief Maximum straggling factor for energy loss + /** + * Uncertainty in the position of the propagated track. Used to set a floor + * in the model error covariance. Physically this represents the unmodeled + * properties (error in RK4, etc.) Numerically is breaks the strong correlation that + * exists between x/y/z in the propagator when the track is basically straight. + */ + double fPosModelNoise{1e-4}; /** * @brief Constructor for the TrackFitterUKF class. From a0789379f2f2dbc684d52ddd7dbad10ca9df8b64 Mon Sep 17 00:00:00 2001 From: Adam Anthony Date: Wed, 30 Jul 2025 11:50:33 -0400 Subject: [PATCH 71/71] Add more output to many track test macro --- macro/tests/UKF/TestManyTracks.C | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/macro/tests/UKF/TestManyTracks.C b/macro/tests/UKF/TestManyTracks.C index f025b74b1..7f7cf6760 100644 --- a/macro/tests/UKF/TestManyTracks.C +++ b/macro/tests/UKF/TestManyTracks.C @@ -15,6 +15,7 @@ const double charge_p = 1.602176634e-19; // Charge of proton // Vectors to store the simulated points to compare to std::vector x_sim, y_sim, z_sim, Eloss_sim; TH1F *hMom = nullptr; +TH1F *hMomSampled = nullptr; TH1F *hMomError = nullptr; TCanvas *c1 = new TCanvas(); TCanvas *c2 = new TCanvas(); @@ -22,7 +23,8 @@ TCanvas *c2 = new TCanvas(); // Parameters for model const int pointsToCluster = 5; const double sigma_pos = 1; // Position uncertainty of 10 mm -const double sigma_mom = 0.01; // Momentum uncertainty in percentage +const double sigma_mom = 0.1; // Momentum uncertainty in percentage +const double sigma_mom_sample = 0.05; // Sampled momentum uncertainty in percentage const double sigma_theta = 1 * M_PI / 180; // Angular uncertainty of 1 degree const double sigma_phi = 1 * M_PI / 180; // Angular uncertainty of 1 degree const double gasDensity = 3.553e-5; // g/cm^3 @@ -55,10 +57,12 @@ void TestManyTracks(int n, double bias = 0) XYZPoint fTruePos(-3.40046e-04, -1.49863e-04, 1.0018); XYZVector fTrueMom(0.00935463, -0.0454279, 0.00826042); fTrueMom *= 1e3; - double fSigmaMom = fTrueMom.R() * sigma_mom; + double fSigmaMom = fTrueMom.R() * sigma_mom_sample; hMom = new TH1F("hMom", "Reconstructed Momentum (MeV/c)", 100, fTrueMom.R() - 4 * fSigmaMom, fTrueMom.R() + 4 * fSigmaMom); + hMomSampled = new TH1F("hMom", "Reconstructed Momentum (MeV/c)", 100, fTrueMom.R() - 4 * fSigmaMom, + fTrueMom.R() + 4 * fSigmaMom); hMomError = new TH1F("hMomError", "Error (%)", 100, -4 * fSigmaMom / fTrueMom.R() * 100, 4 * fSigmaMom / fTrueMom.R() * 100); @@ -67,18 +71,14 @@ void TestManyTracks(int n, double bias = 0) if (i % 100 == 0) std::cout << "On iteration " << i << std::endl; - double pSampled = gRandom->Gaus(fTrueMom.R(), sigma_mom * fTrueMom.R()); + double pSampled = gRandom->Gaus(fTrueMom.R(), sigma_mom_sample * fTrueMom.R()); pSampled += bias; + hMomSampled->Fill(pSampled); // pSampled = fTrueMom.R(); ROOT::Math::Polar3DVector sampledMom(pSampled, fTrueMom.Theta(), fTrueMom.Phi()); - // try { SingleUKF(fTruePos, XYZVector(sampledMom), CalculateInitialCov(pSampled)); - //} catch (...) { - // std::cerr << "Failed to propagate iteration " << i << " with seed momentum " << pSampled << std::endl; - // continue; - //} auto filtState = ukf->GetFilteredStates()[0]; auto smoothState = ukf->GetSmoothedStates()[0]; @@ -98,6 +98,13 @@ void TestManyTracks(int n, double bias = 0) // Draw results c1->cd(); hMom->Draw("hist"); + hMomSampled->SetLineColor(kRed); + hMomSampled->Draw("same hist"); + + auto legend = new TLegend(0.65, 0.75, 0.88, 0.88); + legend->AddEntry(hMom, "Reconstructed", "l"); + legend->AddEntry(hMomSampled, "Sampled", "l"); + legend->Draw(); c2->cd(); hMomError->Draw("hist"); }