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.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx new file mode 100644 index 000000000..6ee181b29 --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.cxx @@ -0,0 +1,276 @@ +#include "TrackFitterUKF.h" + +#include + +#include +#include + +#include + +namespace kf { + +void TrackFitterUKF::Reset() +{ + // Reset the state vector and covariance matrix + TrackFitterUKFBase::Reset(); + m_vecXSmooth.clear(); + m_matPSmooth.clear(); + m_matCPredHist.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 + 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); + } + } + + 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); + 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 + + 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 +{ + 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) { + 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 + 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 (!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(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"); + } + // 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 +} + +Vector TrackFitterUKF::funcH(const Vector &x) +{ + Vector vecZ; + // Calculate the measurement vector based on the position and momentum + vecZ[0] = x[0]; // X coordinate + vecZ[1] = x[1]; // Y coordinate + vecZ[2] = x[2]; // Z coordinate + + return vecZ; // Return the measurement vector +} + +void TrackFitterUKF::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(debug) << "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(debug) << "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); + + // 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_vecXFiltHist.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); + + // logEigen("State PCorr", m_matP, m_matPCorrHist.size()); +} + +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_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; + + // 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_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} + 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} + + // 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 + logEigen("State P+", m_matPSmooth[i - 1], i - 1); + } +} + +} // 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 new file mode 100644 index 000000000..b48b90852 --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h @@ -0,0 +1,592 @@ +/// +/// 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 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. + * + * 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. + * + * 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. + * + * @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: + // 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 + + // 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: + 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 + 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 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()}; + + /// 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; + /// 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 filtered (after correction) state vectors at k + VectorEigenVecDimX m_vecXFiltHist; + /// History of filtered (after correction) state covariances at k + 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()}; + +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_vecXFiltHist.clear(); + m_matPFiltHist.clear(); + } + + /** + * @brief Set the weights used to calculate sigma points. + */ + void setParameters(float alpha, float beta, float kappa) + { + 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; + } + + /** + * @brief Set process noise covariance Q to be used in the prediction step. + */ + 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 SetMeasurementNoise(const Matrix &matR) { m_matR = matR; } + virtual Vector &vecX() { return m_vecX; } + virtual const Vector &vecX() const { return m_vecX; } + + 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 + * + * This function fully updates the augmented state vector and covariance matrix using + * the state and process noise. + */ + void updateAugmentedStateAndCovariance() + { + updateAugWithState(); + updateAugWithProcessNoise(); + // ensurePD(m_matPa); // Ensure the augmented covariance matrix is positive definite + } + + /** + * @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) + { + 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 + 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{ + 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, vecZ)}; // 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); + logEigen("P-", m_matP, 0); // Log the eigenvalues of the covariance matrix + 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 + m_matPPredHist.push_back(m_matP); // Store the predicted covariance matrix + + logEigen("P-Post", m_matP, 0); + } + + /** + * @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; // 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 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)}; + + // 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 + + 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 + + // 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 + 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(); + 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. + */ + 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_matQaug; } + + void updateAugWithProcessNoise() + { + auto processNoiseMean = calculateProcessNoiseMean(); + m_matQaug = 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_matQaug(i - S_IDX, j - S_IDX); + } + } + } + + /** + * @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) { + 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) + { + symmetrize(matP); + Eigen::LLT> lltOfP(matP); + if (lltOfP.info() != Eigen::Success) { + 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) { + 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; + nTouch++; + // Check if the regularized matrix is now positive definite + } + if (lltOfP.info() != Eigen::Success) { + 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. + * + * @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) + { + const float32_t scalarMultiplier{std::sqrt(STATE_DIM + m_lambda)}; // sqrt(n + \kappa) + + Eigen::LLT> lltOfPa = calculateCholesky(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 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) + { + // 1. calculate mean of the sigma points + 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] + } + + // 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_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} + + 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 + } + // ensurePD(matPxx); // Ensure the covariance matrix is positive definite + } + + /** + * @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_weightC0 * (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; + } +}; + +/** + * @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; + 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 + 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 + + /// Smoothed state vector and covariance + 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 + 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. + * @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, std::unique_ptr &&stepper) + : TrackFitterUKFBase(), fPropagator(std::move(propagator)), fStepper(std::move(stepper)) + { + } + virtual void Reset() override; + void SetInitialState(const ROOT::Math::XYZPoint &initialPosition, const ROOT::Math::XYZVector &initialMomentum, + const TMatrixD &initialCovariance); + + void SetMeasCov(const TMatrixD &measCov); + + TMatrixD GetStateCovariance() const; + TMatrixD GetAugStateCovariance() const; + std::array GetAugStateVector() const; + const VectorEigenVecDimX &GetSmoothedStates() const { return m_vecXSmooth; }; + const VectorEigenMatDimX &GetSmoothedCovariances() const { return m_matPSmooth; }; + + void predictUKF(const ROOT::Math::XYZPoint &z); + void correctUKF(const ROOT::Math::XYZPoint &z); + void smoothUKF(); + +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 + +#endif // TRACKFITTERUKF_H diff --git a/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx new file mode 100644 index 000000000..393cdca24 --- /dev/null +++ b/AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKFTest.cxx @@ -0,0 +1,539 @@ +#include "kalman_filter/TrackFitterUKF.h" + +#include "AtELossTable.h" +#include "AtKinematics.h" +#include "AtPropagator.h" + +#include + +#include "TMatrixD.h" + +#include "Math/Vector3D.h" +#include "gtest/gtest.h" + +class TrackFitterUKFExampleTest : 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::TrackFitterUKFBase 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, const kf::Vector &z) + { + 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(TrackFitterUKFExampleTest, 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.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); + + // 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(TrackFitterUKFExampleTest, 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.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); + + // 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.4758845 0.53327217 0.21649734 -0.21214576] + // P = + // [[ 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); +} + +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 + +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; } + Vector getStateVector() { return m_vecX; } + + 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; + + virtual void TearDown() override {} + + static constexpr float FLOAT_EPSILON{0.001F}; + + 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 + + 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}; + 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); + 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 + /// @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, const kf::Vector &z) + { + return {}; + } +}; + +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) +{ + 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 + + // 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 + + 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); + LOG(info) << "Finished setting initial state"; + + 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(); +} + +TEST_F(TrackFitterUKFFixture, TestPredictionAndCorrectStep) +{ + 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 + + // 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/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 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/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; 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/AtReconstruction/CMakeLists.txt b/AtReconstruction/CMakeLists.txt index 3b5eac823..585f5ce2a 100755 --- a/AtReconstruction/CMakeLists.txt +++ b/AtReconstruction/CMakeLists.txt @@ -120,6 +120,8 @@ 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}") @@ -138,6 +140,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/TrackFitterUKFTest.cxx ) attpcroot_generate_tests(OpenKFTests 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/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..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); @@ -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/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) diff --git a/AtTools/AtKinematics.cxx b/AtTools/AtKinematics.cxx index 2431d34f2..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) @@ -303,5 +308,22 @@ 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; +} +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 b0d343fd2..ce09b7a44 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); @@ -73,6 +76,26 @@ 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. + * @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); +double KE(double mom, double mass); + +double GetRelMomFromKE(double KE, 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..5838ac706 --- /dev/null +++ b/AtTools/AtPropagator.cxx @@ -0,0 +1,489 @@ +#include "AtPropagator.h" + +#include "AtKinematics.h" + +#include + +// 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}, + {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) +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 +{ + auto v = Kinematics::GetVel(mom, fState.fMass); + + 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) + + auto drag = -dedx_si * mom.Unit(); + LOG(debug) << "drag: " << drag << " mom " << mom << " dedx " << dedx_si; + + return F_lorentz + drag; // Force in N +} + +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(), fState.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::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(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); + 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(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); + if (!result) { + LOG(error) << "Integration step failed, aborting propagation."; + return; // Abort propagation if step failed + } + fState = result; // Update the internal state + + 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(debug) << "------ Reached measurement surface ------"; + double finalH = (fState.fLastPos - fState.fPos).R(); // Distance traveled in the last step + double approach = surface.Distance(fState.fLastPos); + + 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 + fState.fPos = fState.fLastPos; // Set position to last position + fState.fMom = fState.fLastMom; // Set momentum to last momentum + result = stepper.Step(fState); + if (!result) { + LOG(error) << "Failed to propagate to measurement point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fState.h; // Save original step size + fState = result; // Update the internal state + fState.h = 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(fState.fLastMom, fState.fMass); + double deltaE = KE_last - fStopTol; + deltaE = std::max(deltaE, 0.0); // Ensure we don't have negative energy loss + + 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(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 + fState.fMom = fState.fLastMom; // Set momentum to last momentum + result = stepper.Step(fState); + if (!result) { + LOG(error) << "Failed to propagate to stopping point, aborting."; + return; // Abort propagation if step failed + } + auto origH = fState.h; // Save original step size + fState = result; // Update the internal state + fState.h = origH; // Restore original step size + 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. + + // If we still haven't intersected the surface, we need to adjust the step size + double h = surface.Distance(fState.fPos); // Reduce step size so we hit the surface + if (h <= fDistTol || surface.PassedSurface(result)) { + reachedMeasurementPoint = true; + break; + } + 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) { + LOG(error) << "Failed to propagate to surface after stopping, aborting."; + return; // Abort propagation if step failed + } + fState = result; // Update the internal state + 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 + } + + if (reachedMeasurementPoint || particleStopped || momentumReversed) { + double distanceToSurface = surface.Distance(fState.fPos); + + double KE_final = Kinematics::KE(fState.fMom, fState.fMass); + 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(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(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 +} + +void AtPropagator::PropagateToMeasurementSurface(const AtMeasurementSurface &surface, double eLoss, AtStepper &stepper) +{ + LOG(debug) << "Propagating to surface with eLoss: " << eLoss; + + if (eLoss == 0) { + LOG(warn) << "No energy loss specified, propagating without energy loss adjustment."; + PropagateToMeasurementSurface(surface, stepper); + return; + } + + int iterations = 0; + double calc_eLoss = 0; + 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) { + 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; + + 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."); + } + + iterations++; + PropagateToMeasurementSurface(surface, stepper); // Propagate without energy loss adjustment + + double KE_final = Kinematics::KE(fState.fMom, fState.fMass); + calc_eLoss = KE_initial - KE_final; // Energy loss in MeV + fScalingFactor *= eLoss / calc_eLoss; + 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(debug) << "Energy loss converged after " << iterations << " iterations."; + + fScalingFactor = 1; // Reset scaling factor after convergence +} + +AtPropagator::StepState AtRK4Stepper::Step(const AtPropagator::StepState &state) const +{ + + auto result = state; + result.fLastPos = state.fPos; + result.fLastMom = state.fMom; + result.hUsed = state.h; + result.status = AtPropagator::StepStateStatus::kSuccess; + + auto h = state.h; // Step size in m + auto fPos = state.fPos; + auto fMom = state.fMom; + + 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. + + 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 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.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.fPos = pos_SI * 1e3; // Convert back to mm + + return result; +} + +AtPropagator::StepState AtRK4AdaptiveStepper::Step(const AtPropagator::StepState &state) const +{ + + // Take h to be the step size in m. + auto result = state; + result.fLastPos = state.fPos; + result.fLastMom = state.fMom; + result.status = AtPropagator::StepStateStatus::kSuccess; + + 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. + + auto x0_mm = fPos; + auto p0 = fMom; + 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) + 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(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) + XYZVector p_err = (p_5_MeV - p_4_MeV); // Error in momentum (MeV/c) + + // Calculate the overall error + 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() / (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); + + 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.fPos = x_5_mm; // Update position in mm + result.fMom = p_5_MeV; // Update momentum in MeV/c + 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(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.status = AtPropagator::StepStateStatus::kInvalidStepSize; + result.hUsed = h; + return result; // Abort propagation if step size is out of bounds + } + } + } +} + +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()); + auto currDeriv = (fPoint - result.fPos).Dot(result.fMom.Unit()); + LOG(debug) << "Last Derivative: " << lastDeriv << ", Current Derivative: " << currDeriv; + return lastDeriv * currDeriv <= 0; +} + +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; + auto currSign = fPlane.Distance(result.fPos) > 0 ? 1 : -1; + 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 new file mode 100644 index 000000000..6713e81fd --- /dev/null +++ b/AtTools/AtPropagator.h @@ -0,0 +1,274 @@ +#ifndef ATPROPAGATOR_H +#define ATPROPAGATOR_H + +#include "AtELossModel.h" + +#include + +#include "Math/Plane3D.h" +#include "Math/Point3D.h" +#include "Math/Vector3D.h" + +namespace AtTools { + +class AtMeasurementSurface; +class AtStepper; + +/** + * @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 { +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; + 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 + std::unique_ptr fELossModel; // Energy loss model + + // Internal state variables for the propagator + 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) + + 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. + * @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) + : fELossModel(std::move(elossModel)) + { + fState.fMass = mass; + fState.fQ = charge; + } + AtPropagator(AtPropagator &&) = default; + /** + * @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 XYZPoint &pos, const XYZVector &mom) + { + fState.fPos = pos; + 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; } + + /** + * @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 by adjusting fScalingFactor. + */ + void PropagateToMeasurementSurface(const AtMeasurementSurface &point, double eLoss, AtStepper &stepper); + + 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. + * + * @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(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; + + 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. + } + + std::pair + Derivatives(const ROOT::Math::XYZPoint &pos, const ROOT::Math::XYZVector &mom) const + { + return {dxds(pos, mom), dpds(pos, mom)}; + } + +protected: + /** + * @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 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 + + virtual ~AtStepper() = default; + +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: + 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 + 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 + +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; } +}; + +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(AtPropagator::StepState &result) const override; + ROOT::Math::XYZPoint ProjectToSurface(const ROOT::Math::XYZPoint &pos) const override; +}; +} // namespace AtTools +#endif // #ifndef ATPROPAGATOR_H diff --git a/AtTools/AtPropagatorTest.cxx b/AtTools/AtPropagatorTest.cxx new file mode 100644 index 000000000..dbcd93c95 --- /dev/null +++ b/AtTools/AtPropagatorTest.cxx @@ -0,0 +1,381 @@ +#include "AtPropagator.h" + +#include "AtELossTable.h" +#include "AtKinematics.h" + +#include + +#include +#include +#include +using ROOT::Math::Plane3D; +using ROOT::Math::XYZPoint; +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 +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 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; } + 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) +{ + 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 + 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) +{ + 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 + + // 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) +{ + 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 + + // 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); +} + +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(getEnergyPath()); // Use the function to get the path + 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; + 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); + + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); + + 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.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.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 + 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"); + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path + 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; + 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 (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); + + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); + + 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"); + elossModel->LoadSrimTable(getEnergyPath()); // Use the function to get the path + AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4Stepper stepper; + + 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 + AtMeasurementPlane measurementPlane(plane); + propagator.PropagateToMeasurementSurface(measurementPlane, stepper); + + 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, 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(getEnergyPath()); // Use the function to get the path + AtPropagator propagator(charge, mass, std::move(elossModel)); + AtRK4Stepper stepper; + + 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 + AtMeasurementPlane measurementPlane(plane); + propagator.PropagateToMeasurementSurface(measurementPlane, stepper); + + 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 + 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(getEnergyPath()); // Use the function to get the path + 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; + 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 + stepper.fInitialStep = 1; // Set initial step size to 1 m + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); + + 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); + + propagator.SetState(startPos, startMom); + propagator.SetEField({0, 0, 0}); // No electric field + propagator.SetBField({0, 0, 0}); // No magnetic field + stepper.fInitialStep = 1e-6; // Set initial step size to 1e-6 m + + ASSERT_NEAR(propagator.GetMomentum().X(), 43.331, 1e-1); + + propagator.PropagateToMeasurementSurface(measurementPoint, stepper); + + 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); +} + +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(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 + 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; +} 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} diff --git a/CMakeLists.txt b/CMakeLists.txt index 55a74e19e..3d9fff690 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -126,20 +126,40 @@ 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_MakeAvailable(catima) + +FetchContent_GetProperties(catima) + if (NOT catima_POPULATED) + message(STATUS "Populating CATIMA") + 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/AtPropagator.C b/macro/tests/UKF/AtPropagator.C new file mode 100644 index 000000000..a20fef05a --- /dev/null +++ b/macro/tests/UKF/AtPropagator.C @@ -0,0 +1,83 @@ +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 + + 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; + + 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/TestManyTracks.C b/macro/tests/UKF/TestManyTracks.C new file mode 100644 index 000000000..7f7cf6760 --- /dev/null +++ b/macro/tests/UKF/TestManyTracks.C @@ -0,0 +1,207 @@ +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 *hMomSampled = 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.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 +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_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); + + 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_sample * fTrueMom.R()); + pSampled += bias; + hMomSampled->Fill(pSampled); + + // pSampled = fTrueMom.R(); + + ROOT::Math::Polar3DVector sampledMom(pSampled, fTrueMom.Theta(), fTrueMom.Phi()); + SingleUKF(fTruePos, XYZVector(sampledMom), CalculateInitialCov(pSampled)); + + 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"); + 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"); +} + +/*********** 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 new file mode 100644 index 000000000..99886f9ac --- /dev/null +++ b/macro/tests/UKF/UKFSingleTrack.C @@ -0,0 +1,345 @@ +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; +int pointsToCluster = 5; +void LoadHits() +{ + 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.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 % pointsToCluster != 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; + std::cout << x[0] << " " << x[1] << " " << x[2] << 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, p2, sigmap2, lambda2, sigmalambda2, residual; + std::vector xSmooth, ySmooth, zSmooth, pSmooth, sigmapSmooth, residualSmooth, eLossSmooth; + + // 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 + + 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 + + // 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; + 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 + + // Initial uncertainties + 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 + ukf.fEnableEnStraggling = true; // Enable energy straggling + ukf.setParameters(1e-3, 2, 0); // alpha, beta, kappa + + 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()); + 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()); + + // Skip the first point since it is the initial state. + // 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 + + 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.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]); + 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; + + 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 + } + + 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(); + 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]; + 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 + 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 + } + } + 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]"); + 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); + + 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) { + 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); + 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]); + pGraph->SetPointError(i, 0, + sigmap2[i] * 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 [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"); + 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); + 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 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