Skip to content

Commit 796d893

Browse files
anthoak13anthoak13
committed
Cleanup UKF doc and headers
1 parent 2b070b9 commit 796d893

2 files changed

Lines changed: 156 additions & 162 deletions

File tree

‎AtReconstruction/AtFitter/OpenKF/kalman_filter/TrackFitterUKF.h‎

Lines changed: 150 additions & 155 deletions
Original file line numberDiff line numberDiff line change
@@ -26,146 +26,98 @@
2626
#include "unscented_kalman_filter.h"
2727
namespace kf {
2828

29-
/// @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm.
30-
/// @tparam DIM_X Dimension of the state vector.
31-
/// @tparam DIM_Z Dimension of the measurement vector.
32-
/// @tparam DIM_V Dimension of the process noise vector.
33-
/// @tparam DIM_N Dimension of the measurement noise vector.
29+
/** @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm.
30+
*
31+
* Serves as a templated base class for UKF calculations. It hold no physics information, just
32+
* the machinery that underlies the UKF formalism. It is a modified version of the UKF provided
33+
* by OpenKF, that has been expanded to allow for more hooks into the method.
34+
*
35+
* Templated because I believe Eigen can do quite a bit of operation for small matrices like we have here if
36+
* the size is known at compile time. Worth checking that though.
37+
*
38+
* @tparam DIM_X Dimension of the state vector.
39+
* @tparam DIM_Z Dimension of the measurement vector.
40+
* @tparam DIM_V Dimension of the process noise vector.
41+
* @tparam DIM_N Dimension of the measurement noise vector.
42+
*/
3443
template <int32_t DIM_X = 6, int32_t DIM_Z = 3, int32_t DIM_V = 2, int32_t DIM_N = 3>
35-
class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
36-
44+
class TrackFitterUKFBase {
3745
public:
3846
// Augmented state vector is just the process noise and state vector. The measurement noise is not included as that
3947
// is independent of the propagation and measurement model and just adds linearly.
4048
static constexpr int32_t DIM_A{DIM_X + DIM_V}; ///< @brief Augmented state dimension
4149
static constexpr int32_t SIGMA_DIM_A{2 * DIM_A + 1}; ///< @brief Sigma points dimension for augmented state
50+
Matrix<DIM_A, SIGMA_DIM_A> m_matSigmaXa{Matrix<DIM_A, SIGMA_DIM_A>::Zero()}; ///< @brief Sigma points matrix
4251

4352
protected:
44-
using KalmanFilter<DIM_X, DIM_Z>::m_vecX; // from Base KalmanFilter class
45-
using KalmanFilter<DIM_X, DIM_Z>::m_matP; // from Base KalmanFilter class
46-
47-
float32_t m_weight0; /// @brief unscented transform weight 0 for mean
48-
float32_t m_weighti; /// @brief unscented transform weight i for none mean samples
49-
float32_t m_kappa{0}; ///< @brief Kappa parameter for finding sigma points
50-
51-
Vector<DIM_A> m_vecXa{Vector<DIM_A>::Zero()}; /// @brief augmented state vector (incl. process
52-
/// and measurement noise means)
53-
Matrix<DIM_A, DIM_A> m_matPa{Matrix<DIM_A, DIM_A>::Zero()}; /// @brief augmented state covariance (incl.
54-
/// process and measurement noise covariances)
55-
56-
// Add variables to track the covariances of the process and measurement noise.
57-
Matrix<DIM_V, DIM_V> m_matQ; // @brief Process noise covariance matrix
58-
Matrix<DIM_N, DIM_N> m_matR; // @brief Measurement noise covariance matrix
53+
Vector<DIM_X> m_vecX{Vector<DIM_X>::Zero()}; /// @brief estimated state vector
54+
Matrix<DIM_X, DIM_X> m_matP{Matrix<DIM_X, DIM_X>::Zero()}; /// @brief state covariance matrix
55+
56+
/// @brief Augmented state vector (incl. process and measurement noise means)
57+
Vector<DIM_A> m_vecXa{Vector<DIM_A>::Zero()};
58+
/// @brief augmented state covariance (incl. process and measurement noise covariances)
59+
Matrix<DIM_A, DIM_A> m_matPa{Matrix<DIM_A, DIM_A>::Zero()};
60+
61+
// Process and measurement noise covariance matrices
62+
/// Process noise covariance matrix (Q)
63+
Matrix<DIM_V, DIM_V> m_matQ{Matrix<DIM_V, DIM_V>::Zero()};
64+
/// Measurement noise covariance matrix (R)
65+
Matrix<DIM_N, DIM_N> m_matR{Matrix<DIM_N, DIM_N>::Zero()};
66+
67+
/// Unscented transform weight for the mean sigma point
68+
float32_t m_weight0;
69+
/// Unscented transform weight for the other sigma points
70+
float32_t m_weighti;
71+
/// Kappa parameter for sigma point calculation
72+
float32_t m_kappa{3 - DIM_A};
5973

6074
public:
61-
Matrix<DIM_A, SIGMA_DIM_A> m_matSigmaXa{Matrix<DIM_A, SIGMA_DIM_A>::Zero()}; ///< @brief Sigma points matrix
62-
63-
TrackFitterUKFBase()
64-
: KalmanFilter<DIM_X, DIM_Z>(), m_kappa(3 - DIM_A), m_matQ(Matrix<DIM_V, DIM_V>::Zero()),
65-
m_matR(Matrix<DIM_N, DIM_N>::Zero())
66-
{
67-
// 1. calculate weights
68-
updateWeights();
69-
}
70-
71-
~TrackFitterUKFBase() {}
75+
TrackFitterUKFBase() { updateWeights(); }
76+
~TrackFitterUKFBase() = default;
7277

7378
void setKappa(float32_t kappa)
7479
{
7580
m_kappa = kappa; // Set the kappa parameter for sigma point calculation
7681
updateWeights(); // Update the weights based on the new kappa value
7782
}
7883

79-
///
80-
/// @brief adding process noise covariance Q to the augmented state covariance
81-
/// matPa in the middle element of the diagonal.
82-
///
83-
void setCovarianceQ(const Matrix<DIM_V, DIM_V> &matQ)
84-
{
85-
m_matQ = matQ; // Store the process noise covariance matrix
86-
}
87-
88-
///
89-
/// @brief set the measurement noise covariance R to be used in the update step
90-
///
91-
void setCovarianceR(const Matrix<DIM_N, DIM_N> &matR)
92-
{
93-
m_matR = matR; // Store the measurement noise covariance matrix
94-
}
95-
96-
/// Add state vector (m_vecX) to the augment state vector (m_vecXa) and also
97-
/// add covariance matrix (m_matP) to the augment covariance (m_matPa).
98-
void updateAugWithState()
99-
{
100-
// Copy state vector to augmented state vector
101-
for (int32_t i{0}; i < DIM_X; ++i) {
102-
m_vecXa[i] = m_vecX[i];
103-
}
104-
105-
// Copy state covariance matrix to augmented covariance matrix
106-
for (int32_t i{0}; i < DIM_X; ++i) {
107-
for (int32_t j{0}; j < DIM_X; ++j) {
108-
m_matPa(i, j) = m_matP(i, j);
109-
}
110-
}
111-
}
112-
113-
virtual std::array<float32_t, DIM_V> calculateProcessNoiseMean()
114-
{
115-
// Calculate the expectation value of the process noise using the current value of the state vector m_vecX
116-
std::array<float32_t, DIM_V> processNoiseMean{0};
117-
118-
// TODO: Set the mean energy loss based on the momentum and particle type. Probably best to track stopping power?
119-
return processNoiseMean;
120-
}
121-
122-
virtual Matrix<DIM_V, DIM_V> calculateProcessNoiseCovariance()
123-
{
124-
// Calculate the process noise covariance matrix
125-
Matrix<DIM_V, DIM_V> matQ{Matrix<DIM_V, DIM_V>::Zero()};
126-
matQ = m_matQ; // Use the stored process noise covariance matrix
127-
128-
// TODO: Set the process noise covariance for angular straggle and energy loss.
129-
return matQ;
130-
}
131-
132-
void updateAugWithProcessNoise()
133-
{
134-
auto processNoiseMean = calculateProcessNoiseMean();
135-
m_matQ = calculateProcessNoiseCovariance();
136-
137-
// Add the mean process noise to the augmented state vector
138-
for (int32_t i{0}; i < DIM_V; ++i) {
139-
m_vecXa[DIM_X + i] = processNoiseMean[i];
140-
}
84+
/**
85+
* @brief Set process noise covariance Q to be used in the prediction step.
86+
*/
87+
void setCovarianceQ(const Matrix<DIM_V, DIM_V> &matQ) { m_matQ = matQ; }
88+
/**
89+
* @brief Set the measurement noise covariance R to be used in the update step.
90+
*/
91+
void setCovarianceR(const Matrix<DIM_N, DIM_N> &matR) { m_matR = matR; }
92+
virtual Vector<DIM_X> &vecX() { return m_vecX; }
93+
virtual const Vector<DIM_X> &vecX() const { return m_vecX; }
14194

142-
// Add process noise covariance to the augmented covariance matrix
143-
const int32_t S_IDX{DIM_X};
144-
const int32_t L_IDX{S_IDX + DIM_V};
95+
virtual Matrix<DIM_X, DIM_X> &matP() { return m_matP; }
96+
virtual const Matrix<DIM_X, DIM_X> &matP() const { return m_matP; }
14597

146-
for (int32_t i{S_IDX}; i < L_IDX; ++i) {
147-
for (int32_t j{S_IDX}; j < L_IDX; ++j) {
148-
m_matPa(i, j) = m_matQ(i - S_IDX, j - S_IDX);
149-
}
150-
}
151-
}
152-
153-
///
154-
/// @brief update the augmented state vector and covariance matrix
155-
/// This functions fully updates the augmented state vector (m_vecXa) and covariance matrix (m_matPa)
156-
/// by setting both the state vector and process noise components.
157-
///
98+
/**
99+
* @brief update the augmented state vector and covariance matrix
100+
*
101+
* This function fully updates the augmented state vector and covariance matrix using
102+
* the state and process noise.
103+
*/
158104
void updateAugmentedStateAndCovariance()
159105
{
160106
updateAugWithState();
161107
updateAugWithProcessNoise();
162108
}
163109

164-
///
165-
/// @brief state prediction step of the unscented Kalman filter (UKF).
166-
/// @param predictionModelFunc callback to the prediction/process model
167-
/// function
168-
///
110+
/**
111+
* @brief state prediction step of the unscented Kalman filter (UKF).
112+
* @param predictionModelFunc callback to the prediction/process model function.
113+
* @param vecZ actual measurement vector.`
114+
*
115+
* A template is used here for performace reasons since there is only really a single prediction
116+
* model used. Honestly it probably does not make a difference compared to an std::function, but
117+
* this was not changed from OpenKF.
118+
*
119+
* This modifies the state vector and state covariance.
120+
*/
169121
template <typename PredictionModelCallback>
170122
void predictUKF(PredictionModelCallback predictionModelFunc, const Vector<DIM_Z> &vecZ)
171123
{
@@ -199,11 +151,11 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
199151
calculateWeightedMeanAndCovariance<DIM_X>(sigmaXx, m_vecX, m_matP);
200152
}
201153

202-
///
203-
/// @brief measurement correction step of the unscented Kalman filter (UKF).
204-
/// @param measurementModelFunc callback to the measurement model function
205-
/// @param vecZ actual measurement vector.
206-
///
154+
/**
155+
* @brief measurement correction step of the unscented Kalman filter (UKF).
156+
* @param measurementModelFunc callback to the measurement model function
157+
* @param vecZ actual measurement vector.
158+
*/
207159
template <typename MeasurementModelCallback>
208160
void correctUKF(MeasurementModelCallback measurementModelFunc, const Vector<DIM_Z> &vecZ)
209161
{
@@ -235,7 +187,6 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
235187
// Add in the measurement noise covariance matrix to the measurement covariance matrix.
236188
matPzz += m_matR; // Add measurement noise covariance
237189

238-
// TODO: calculate cross correlation
239190
const Matrix<DIM_X, DIM_Z> matPxz{calculateCrossCorrelation(sigmaXx, m_vecX, sigmaZ, vecZhat)};
240191

241192
// kalman gain
@@ -246,9 +197,9 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
246197
}
247198

248199
protected:
249-
///
250-
/// @brief algorithm to calculate the weights used to draw the sigma points
251-
///
200+
/**
201+
* @brief Set the weights used to calculate sigma points.
202+
*/
252203
void updateWeights()
253204
{
254205
static_assert(DIM_A > 0, "DIM_A is Zero which leads to numerical issue.");
@@ -258,17 +209,59 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
258209
m_weight0 = m_kappa / denoTerm;
259210
m_weighti = 0.5F / denoTerm;
260211
}
212+
/**
213+
* @brief Add state vector and state covariance matrix to the augmented state vector covariance matrix.
214+
*/
215+
void updateAugWithState()
216+
{
217+
// Copy state vector to augmented state vector
218+
for (int32_t i{0}; i < DIM_X; ++i) {
219+
m_vecXa[i] = m_vecX[i];
220+
}
221+
222+
// Copy state covariance matrix to augmented covariance matrix
223+
for (int32_t i{0}; i < DIM_X; ++i) {
224+
for (int32_t j{0}; j < DIM_X; ++j) {
225+
m_matPa(i, j) = m_matP(i, j);
226+
}
227+
}
228+
}
229+
230+
virtual std::array<float32_t, DIM_V> calculateProcessNoiseMean() { return std::array<float32_t, DIM_V>{0}; }
231+
232+
virtual Matrix<DIM_V, DIM_V> calculateProcessNoiseCovariance() { return m_matQ; }
233+
234+
void updateAugWithProcessNoise()
235+
{
236+
auto processNoiseMean = calculateProcessNoiseMean();
237+
m_matQ = calculateProcessNoiseCovariance();
238+
239+
// Add the mean process noise to the augmented state vector
240+
for (int32_t i{0}; i < DIM_V; ++i) {
241+
m_vecXa[DIM_X + i] = processNoiseMean[i];
242+
}
243+
244+
// Add process noise covariance to the augmented covariance matrix
245+
const int32_t S_IDX{DIM_X};
246+
const int32_t L_IDX{S_IDX + DIM_V};
247+
248+
for (int32_t i{S_IDX}; i < L_IDX; ++i) {
249+
for (int32_t j{S_IDX}; j < L_IDX; ++j) {
250+
m_matPa(i, j) = m_matQ(i - S_IDX, j - S_IDX);
251+
}
252+
}
253+
}
261254

262-
///
263-
/// @brief algorithm to calculate the deterministic sigma points for
264-
/// the unscented transformation
265-
///
266-
/// @param vecX mean of the normally distributed state
267-
/// @param matPxx covariance of the normally distributed state
268-
/// @param STATE_DIM dimension of the vector used to calculate the sigma points
269-
/// @param SIGMA_DIM number of sigma points required (default is 2 * STATE_DIM + 1)
270-
/// @return matrix of sigma points where each column is a sigma point
271-
///
255+
/**
256+
* @brief Algorithm to calculate the deterministic sigma points for
257+
* the unscented transformation.
258+
*
259+
* @param vecX Mean of the normally distributed state.
260+
* @param matPxx Covariance of the normally distributed state.
261+
* @param STATE_DIM Dimension of the vector used to calculate the sigma points.
262+
* @param SIGMA_DIM Number of sigma points required (default is 2 * STATE_DIM + 1).
263+
* @return Matrix of sigma points where each column is a sigma point.
264+
*/
272265
template <int32_t STATE_DIM, int32_t SIGMA_DIM = 2 * STATE_DIM + 1>
273266
Matrix<STATE_DIM, SIGMA_DIM>
274267
calculateSigmaPoints(const Vector<STATE_DIM> &vecXa, const Matrix<STATE_DIM, STATE_DIM> &matPa)
@@ -321,14 +314,12 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
321314
return sigmaXa;
322315
}
323316

324-
///
325-
/// @brief calculate the weighted mean and covariance given a set of sigma
326-
/// points
327-
/// @param[in] sigmaX matrix of (probably posterior) sigma points where each column contain single
328-
/// sigma point.
329-
/// @param[out] vecX output weighted mean of the sigma points
330-
/// @param[out] matPxx output weighted covariance of the sigma points
331-
///
317+
/**
318+
* @brief Calculate the weighted mean and covariance given a set of sigma points.
319+
* @param[in] sigmaX Matrix of (probably posterior) sigma points where each column contains a single sigma point.
320+
* @param[out] vecX Output weighted mean of the sigma points.
321+
* @param[out] matPxx Output weighted covariance of the sigma points.
322+
*/
332323
template <int32_t STATE_DIM, int32_t SIGMA_DIM>
333324
void calculateWeightedMeanAndCovariance(const Matrix<STATE_DIM, SIGMA_DIM> &sigmaX, Vector<STATE_DIM> &vecX,
334325
Matrix<STATE_DIM, STATE_DIM> &matPxx)
@@ -356,17 +347,17 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
356347
}
357348
}
358349

359-
///
360-
/// @brief calculate the cross-correlation given two sets sigma points X and Y
361-
/// and their means x and y
362-
/// @param sigmaX first matrix of sigma points where each column contain
363-
/// single sigma point
364-
/// @param vecX mean of the first set of sigma points
365-
/// @param sigmaY second matrix of sigma points where each column contain
366-
/// single sigma point
367-
/// @param vecY mean of the second set of sigma points
368-
/// @return matPxy, the cross-correlation matrix
369-
///
350+
/**
351+
* @brief calculate the cross-correlation given two sets sigma points X and Y
352+
* and their means x and y
353+
* @param sigmaX first matrix of sigma points where each column contain
354+
* single sigma point
355+
* @param vecX mean of the first set of sigma points
356+
* @param sigmaY second matrix of sigma points where each column contain
357+
* single sigma point
358+
* @param vecY mean of the second set of sigma points
359+
* @return matPxy, the cross-correlation matrix
360+
*/
370361
template <int32_t SIGMA_DIM>
371362
Matrix<DIM_X, DIM_Z> calculateCrossCorrelation(const Matrix<DIM_X, SIGMA_DIM> &sigmaX, const Vector<DIM_X> &vecX,
372363
const Matrix<DIM_Z, SIGMA_DIM> &sigmaY, const Vector<DIM_Z> &vecY)
@@ -389,8 +380,11 @@ class TrackFitterUKFBase : public KalmanFilter<DIM_X, DIM_Z> {
389380
}
390381
};
391382

392-
// Define template dimension variables for clarity and reuse
393-
383+
/**
384+
* @brief Class for fitting tracks using the Unscented Kalman Filter (UKF) algorithm.
385+
*
386+
* UKF specialized for fitting tracks. This class is where all of the physics is
387+
*/
394388
class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> {
395389
protected:
396390
static constexpr int32_t TF_DIM_X = 6;
@@ -402,6 +396,7 @@ class TrackFitterUKF : public TrackFitterUKFBase<6, 3, 1, 3> {
402396
std::unique_ptr<AtTools::AtStepper> fStepper{nullptr};
403397
AtTools::AtPropagator::StepState fMeanStep; /// Holds the step information for POCA propagation of mean state
404398
ROOT::Math::Plane3D fMeasurementPlane; ///< Holds the measurement plane for the track fitter
399+
405400
public:
406401
bool fEnableEnStraggling{true}; ///< @brief Flag to enable/disable energy straggling
407402
double fMaxStragglingFactor{1. / 3.}; ///< @brief Maximum straggling factor for energy loss

0 commit comments

Comments
 (0)