Common: DCAFitter add fit status code

f3sch · f3sch · commit ba169e683f8e · 2025-04-01T10:46:56.000+02:00
diff --git a/Common/DCAFitter/include/DCAFitter/DCAFitterN.h b/Common/DCAFitter/include/DCAFitter/DCAFitterN.h
@@ -118,10 +118,26 @@ class DCAFitterN
   using ArrTrPos = o2::gpu::gpustd::array<Vec3D, N>;         // container of Track positions
 
  public:
-  enum BadCovPolicy {   // if encountering non-positive defined cov. matrix, the choice is:
-    Discard = 0,        // stop evaluation
-    Override = 1,       // override correlation coef. to have cov.matrix pos.def and continue
-    OverrideAndFlag = 2 // override correlation coef. to have cov.matrix pos.def, set mPropFailed flag of corresponding candidate to true and continue (up to the user to check the flag)
+  enum BadCovPolicy : uint8_t { // if encountering non-positive defined cov. matrix, the choice is:
+    Discard = 0,                // stop evaluation
+    Override = 1,               // override correlation coef. to have cov.matrix pos.def and continue
+    OverrideAndFlag = 2         // override correlation coef. to have cov.matrix pos.def, set mPropFailed flag of corresponding candidate to true and continue (up to the user to check the flag)
+  };
+
+  enum FitStatus : uint8_t { // fit status of crossing hypothesis
+    None = 0,                // no status set
+    Converged = 1,           // fit converged
+    NoCrossing = 2,          // no reasaonable crossing was found
+    RejRadius = 3,           // radius of crossing was not acceptable
+    RejTrackX = 4,           // one candidate track x was below the mimimum required radius
+    RejTrackRoughZ = 5,      // rejected by rough cut on tracks Z difference
+    RejChi2Max = 6,          // rejected by maximum chi2 cut
+    FailProp = 7,            // propagation of at least prong to PCA failed
+    FailInvCov = 8,          // inversion of cov.-matrix failed
+    FailInvWeight = 9,       // inversion of Ti weight matrix failed
+    FailInv2ndDeriv = 10,    // inversion of 2nd derivatives failed
+    FailCorrTracks = 11,     // correction of tracks to updated x failed
+    FailCloserAlt = 12,      // alternative PCA is closer
   };
 
   static constexpr int getNProngs() { return N; }
@@ -154,7 +170,7 @@ class DCAFitterN
   ///< check if propagation of tracks to candidate vertex was done
   GPUd() bool isPropagateTracksToVertexDone(int cand = 0) const { return mTrPropDone[mOrder[cand]]; }
 
-  ///< check if propagation of tracks to candidate vertex was done
+  ///< check if propagation of tracks to candidate vertex failed
   bool isPropagationFailure(int cand = 0) const { return mPropFailed[mOrder[cand]]; }
 
   ///< track param propagated to V0 candidate (no check for the candidate validity)
@@ -201,6 +217,8 @@ class DCAFitterN
 
   const Track* getOrigTrackPtr(int i) const { return mOrigTrPtr[i]; }
 
+  GPUdi() FitStatus getFitStatus(int cand = 0) const noexcept { return mFitStatus[mOrder[cand]]; }
+
   ///< return number of iterations during minimization (no check for its validity)
   GPUdi() int getNIterations(int cand = 0) const { return mNIters[mOrder[cand]]; }
   GPUdi() void setPropagateToPCA(bool v = true) { mPropagateToPCA = v; }
@@ -315,6 +333,8 @@ class DCAFitterN
   {
     mCurHyp = 0;
     mAllowAltPreference = true;
+    mOrder.fill(0);
+    mFitStatus.fill(FitStatus::None);
   }
 
   GPUdi() static void setTrackPos(Vec3D& pnt, const Track& tr)
@@ -362,12 +382,13 @@ class DCAFitterN
   LogLogThrottler mLoggerBadCov{};
   LogLogThrottler mLoggerBadInv{};
   LogLogThrottler mLoggerBadProp{};
-  MatSym3D mWeightInv;                                   // inverse weight of single track, [sum{M^T E M}]^-1 in EQ.T
+  MatSym3D mWeightInv; // inverse weight of single track, [sum{M^T E M}]^-1 in EQ.T
   o2::gpu::gpustd::array<int, MAXHYP> mOrder{0};
   int mCurHyp = 0;
   int mCrossIDCur = 0;
   int mCrossIDAlt = -1;
   BadCovPolicy mBadCovPolicy{BadCovPolicy::Discard};                                              // what to do in case of non-pos-def. cov. matrix, see BadCovPolicy enum
+  o2::gpu::gpustd::array<FitStatus, MAXHYP> mFitStatus{};                                         // fit status of each hypothesis fit
   bool mAllowAltPreference = true;                                                                // if the fit converges to alternative PCA seed, abandon the current one
   bool mUseAbsDCA = false;                                                                        // use abs. distance minimization rather than chi2
   bool mWeightedFinalPCA = false;                                                                 // recalculate PCA as a cov-matrix weighted mean, even if absDCA method was used
@@ -390,7 +411,7 @@ class DCAFitterN
   float mMaxStep = 2.0;                                                                           // Max step for propagation with Propagator
   int mFitterID = 0;                                                                              // locat fitter ID (mostly for debugging)
   size_t mCallID = 0;
-  ClassDefNV(DCAFitterN, 2);
+  ClassDefNV(DCAFitterN, 3);
 };
 
 ///_________________________________________________________________________
@@ -407,7 +428,8 @@ GPUd() int DCAFitterN<N, Args...>::process(const Tr&... args)
     mTrAux[i].set(*mOrigTrPtr[i], mBz);
   }
   if (!mCrossings.set(mTrAux[0], *mOrigTrPtr[0], mTrAux[1], *mOrigTrPtr[1], mMaxDXYIni, mIsCollinear)) { // even for N>2 it should be enough to test just 1 loop
-    return 0;                                                                                            // no crossing
+    mFitStatus[mCurHyp] = FitStatus::NoCrossing;
+    return 0;
   }
   for (int ih = 0; ih < MAXHYP; ih++) {
     mPropFailed[ih] = false;
@@ -428,6 +450,7 @@ GPUd() int DCAFitterN<N, Args...>::process(const Tr&... args)
   for (int ic = 0; ic < mCrossings.nDCA; ic++) {
     // check if radius is acceptable
     if (mCrossings.xDCA[ic] * mCrossings.xDCA[ic] + mCrossings.yDCA[ic] * mCrossings.yDCA[ic] > mMaxR2) {
+      mFitStatus[mCurHyp] = FitStatus::RejRadius;
       continue;
     }
     mCrossIDCur = ic;
@@ -468,6 +491,7 @@ GPUd() bool DCAFitterN<N, Args...>::calcPCACoefs()
 {
   //< calculate Ti matrices for global vertex decomposition to V = sum_{0<i<N} Ti pi, see EQ.T in the ref
   if (!calcInverseWeight()) {
+    mFitStatus[mCurHyp] = FitStatus::FailInvWeight;
     return false;
   }
   for (int i = N; i--;) { // build Mi*Ei matrix
@@ -558,7 +582,7 @@ GPUd() void DCAFitterN<N, Args...>::calcResidDerivatives()
         dr2[2] += trDx.d2zdx2;
       }
     } // track over which we differentiate
-  } // residual being differentiated
+  }   // residual being differentiated
 }
 
 //__________________________________________________________________________
@@ -608,7 +632,7 @@ GPUd() void DCAFitterN<N, Args...>::calcResidDerivativesNoErr()
       dr2ji[2] = -trDxi.d2zdx2 * NInv;
 
     } // track over which we differentiate
-  } // residual being differentiated
+  }   // residual being differentiated
 }
 
 //__________________________________________________________________________
@@ -727,6 +751,7 @@ GPUd() bool DCAFitterN<N, Args...>::recalculatePCAWithErrors(int cand)
 #endif
           mLoggerBadCov.evCountPrev = mLoggerBadCov.evCount;
         }
+        mFitStatus[mCurHyp] = FitStatus::FailInvCov;
         if (mBadCovPolicy == Discard) {
           return false;
         } else if (mBadCovPolicy == OverrideAndFlag) {
@@ -935,10 +960,14 @@ GPUd() bool DCAFitterN<N, Args...>::minimizeChi2()
   for (int i = N; i--;) {
     mCandTr[mCurHyp][i] = *mOrigTrPtr[i];
     auto x = mTrAux[i].c * mPCA[mCurHyp][0] + mTrAux[i].s * mPCA[mCurHyp][1]; // X of PCA in the track frame
-    if (x < mMinXSeed || !propagateToX(mCandTr[mCurHyp][i], x)) {
+    if (x < mMinXSeed) {
+      mFitStatus[mCurHyp] = FitStatus::RejTrackX;
       return false;
     }
-    setTrackPos(mTrPos[mCurHyp][i], mCandTr[mCurHyp][i]);      // prepare positions
+    if (!propagateToX(mCandTr[mCurHyp][i], x)) {
+      return false;
+    }
+    setTrackPos(mTrPos[mCurHyp][i], mCandTr[mCurHyp][i]);             // prepare positions
     if (!mTrcEInv[mCurHyp][i].set(mCandTr[mCurHyp][i], XerrFactor)) { // prepare inverse cov.matrices at starting point
       if (mLoggerBadCov.needToLog()) {
 #ifndef GPUCA_GPUCODE
@@ -950,6 +979,7 @@ GPUd() bool DCAFitterN<N, Args...>::minimizeChi2()
 #endif
         mLoggerBadCov.evCountPrev = mLoggerBadCov.evCount;
       }
+      mFitStatus[mCurHyp] = FitStatus::FailInvCov;
       if (mBadCovPolicy == Discard) {
         return false;
       } else if (mBadCovPolicy == OverrideAndFlag) {
@@ -959,6 +989,7 @@ GPUd() bool DCAFitterN<N, Args...>::minimizeChi2()
   }
 
   if (mMaxDZIni > 0 && !roughDZCut()) { // apply rough cut on tracks Z difference
+    mFitStatus[mCurHyp] = FitStatus::RejTrackX;
     return false;
   }
 
@@ -979,28 +1010,36 @@ GPUd() bool DCAFitterN<N, Args...>::minimizeChi2()
         printf("fitter %d: error (%ld muted): Inversion failed\n", mFitterID, mLoggerBadCov.getNMuted());
         mLoggerBadInv.evCountPrev = mLoggerBadInv.evCount;
       }
+      mFitStatus[mCurHyp] = FitStatus::FailInv2ndDeriv;
       return false;
     }
     VecND dx = mD2Chi2Dx2 * mDChi2Dx;
     if (!correctTracks(dx)) {
+      mFitStatus[mCurHyp] = FitStatus::FailCorrTracks;
       return false;
     }
     calcPCA(); // updated PCA
     if (mCrossIDAlt >= 0 && closerToAlternative()) {
+      mFitStatus[mCurHyp] = FitStatus::FailCloserAlt;
       mAllowAltPreference = false;
       return false;
     }
     calcTrackResiduals(); // updated residuals
     chi2Upd = calcChi2(); // updated chi2
     if (getAbsMax(dx) < mMinParamChange || chi2Upd > chi2 * mMinRelChi2Change) {
       chi2 = chi2Upd;
+      mFitStatus[mCurHyp] = FitStatus::Converged;
       break; // converged
     }
     chi2 = chi2Upd;
   } while (++mNIters[mCurHyp] < mMaxIter);
   //
   mChi2[mCurHyp] = chi2 * NInv;
-  return mChi2[mCurHyp] < mMaxChi2;
+  if (mChi2[mCurHyp] >= mMaxChi2) {
+    mFitStatus[mCurHyp] = FitStatus::RejChi2Max;
+    return false;
+  }
+  return true;
 }
 
 //___________________________________________________________________
@@ -1012,12 +1051,17 @@ GPUd() bool DCAFitterN<N, Args...>::minimizeChi2NoErr()
   for (int i = N; i--;) {
     mCandTr[mCurHyp][i] = *mOrigTrPtr[i];
     auto x = mTrAux[i].c * mPCA[mCurHyp][0] + mTrAux[i].s * mPCA[mCurHyp][1]; // X of PCA in the track frame
-    if (x < mMinXSeed || !propagateParamToX(mCandTr[mCurHyp][i], x)) {
+    if (x < mMinXSeed) {
+      mFitStatus[mCurHyp] = FitStatus::RejTrackX;
+      return false;
+    }
+    if (!propagateParamToX(mCandTr[mCurHyp][i], x)) {
       return false;
     }
     setTrackPos(mTrPos[mCurHyp][i], mCandTr[mCurHyp][i]); // prepare positions
   }
   if (mMaxDZIni > 0 && !roughDZCut()) { // apply rough cut on tracks Z difference
+    mFitStatus[mCurHyp] = FitStatus::RejTrackX;
     return false;
   }
 
@@ -1035,28 +1079,36 @@ GPUd() bool DCAFitterN<N, Args...>::minimizeChi2NoErr()
         printf("itter %d: error (%ld muted): Inversion failed\n", mFitterID, mLoggerBadCov.getNMuted());
         mLoggerBadInv.evCountPrev = mLoggerBadInv.evCount;
       }
+      mFitStatus[mCurHyp] = FitStatus::FailInv2ndDeriv;
       return false;
     }
     VecND dx = mD2Chi2Dx2 * mDChi2Dx;
     if (!correctTracks(dx)) {
+      mFitStatus[mCurHyp] = FitStatus::FailCorrTracks;
       return false;
     }
     calcPCANoErr(); // updated PCA
     if (mCrossIDAlt >= 0 && closerToAlternative()) {
+      mFitStatus[mCurHyp] = FitStatus::FailCloserAlt;
       mAllowAltPreference = false;
       return false;
     }
     calcTrackResiduals();      // updated residuals
     chi2Upd = calcChi2NoErr(); // updated chi2
     if (getAbsMax(dx) < mMinParamChange || chi2Upd > chi2 * mMinRelChi2Change) {
       chi2 = chi2Upd;
+      mFitStatus[mCurHyp] = FitStatus::Converged;
       break; // converged
     }
     chi2 = chi2Upd;
   } while (++mNIters[mCurHyp] < mMaxIter);
   //
   mChi2[mCurHyp] = chi2 * NInv;
-  return mChi2[mCurHyp] < mMaxChi2;
+  if (mChi2[mCurHyp] >= mMaxChi2) {
+    mFitStatus[mCurHyp] = FitStatus::RejChi2Max;
+    return false;
+  }
+  return true;
 }
 
 //___________________________________________________________________
@@ -1182,6 +1234,7 @@ GPUdi() bool DCAFitterN<N, Args...>::propagateParamToX(o2::track::TrackPar& t, f
     res = t.propagateParamTo(x, mBz);
   }
   if (!res) {
+    mFitStatus[mCurHyp] = FitStatus::FailProp;
     mPropFailed[mCurHyp] = true;
     if (mLoggerBadProp.needToLog()) {
 #ifndef GPUCA_GPUCODE
@@ -1208,6 +1261,7 @@ GPUdi() bool DCAFitterN<N, Args...>::propagateToX(o2::track::TrackParCov& t, flo
     res = t.propagateTo(x, mBz);
   }
   if (!res) {
+    mFitStatus[mCurHyp] = FitStatus::FailProp;
     mPropFailed[mCurHyp] = true;
     if (mLoggerBadProp.needToLog()) {
 #ifndef GPUCA_GPUCODE
