TPC: Improve standalone dEdx calculation class

Author: tubagundem

Detectors/TPC/calibration/CMakeLists.txt

@@ -165,6 +165,10 @@ o2_add_test_root_macro(macro/drawCMV.C
                        COMPILE_ONLY
                        PUBLIC_LINK_LIBRARIES O2::TPCCalibration O2::TPCBase
                        LABELS tpc)
+o2_add_test_root_macro(macro/calculatedEdx.C
+                       COMPILE_ONLY
+                       PUBLIC_LINK_LIBRARIES O2::TPCCalibration O2::TPCBase
+                       LABELS tpc)
 
 o2_add_test(IDCFourierTransform
             COMPONENT_NAME calibration

Detectors/TPC/calibration/include/TPCCalibration/CalculatedEdx.h

@@ -20,16 +20,32 @@
 // o2 includes
 #include "DataFormatsTPC/TrackTPC.h"
 #include "DataFormatsTPC/dEdxInfo.h"
+#include "TPCBase/Mapper.h"
 #include "GPUO2InterfaceRefit.h"
 #include "CalibdEdxContainer.h"
+#include "CorrectionMapsHelper.h"
 #include "CommonUtils/TreeStreamRedirector.h"
 #include "TPCCalibration/CorrectdEdxDistortions.h"
 #include "TPCFastTransformPOD.h"
+#include "GPUCommonRtypes.h"
 #include <vector>
+#include <map>
+#include <unordered_map>
+#include <string>
+#include <utility>
 
 namespace o2::tpc
 {
 
+/// \brief average cluster occupancy of a track, per TPC region
+struct AverageOccupancy {
+  double IROC = 0.;
+  double OROC1 = 0.;
+  double OROC2 = 0.;
+  double OROC3 = 0.;
+  ClassDefNV(AverageOccupancy, 1);
+};
+
 /// \brief dEdx calculation class
 ///
 /// This class is used to calculate dEdx of reconstructed tracks.
@@ -46,7 +62,7 @@ namespace o2::tpc
 /// c.setMembers(tpcTrackClIdxVecInput, clusterIndex, tpcTracks); // set the member variables: TrackTPC, TPCClRefElem, o2::tpc::ClusterNativeAccess
 /// c.setRefit(); // set the refit pointer to perform refitting of tracks, otherwise setPropagateTrack to true
 /// start looping over the tracks
-/// c.calculatedEdx(track, output, 0.015, 0.60, CorrectionFlags::TopologyPol | CorrectionFlags::dEdxResidual, ClusterFlags::ExcludeEdgeCl) // this will fill the dEdxInfo output for given track
+/// c.calculatedEdx(track, output, averageOcc, 0.015, 0.60, CorrectionFlags::TopologyPol | CorrectionFlags::dEdxResidual, ClusterFlags::ExcludeEdgeCl) // this will fill the dEdxInfo output and per-region average track occupancy averageOcc for given track
 
 enum class CorrectionFlags : unsigned short {
   None = 0,
@@ -61,11 +77,14 @@ enum class CorrectionFlags : unsigned short {
 enum class ClusterFlags : unsigned short {
   None = 0,
   ExcludeSingleCl = 1 << 0,         ///< flag to exclude single clusters in dEdx calculation
-  ExcludeSplitCl = 1 << 1,          ///< flag to exclude split clusters in dEdx calculation
-  ExcludeEdgeCl = 1 << 2,           ///< flag to exclude sector edge clusters in dEdx calculation
-  ExcludeSubthresholdCl = 1 << 3,   ///< flag to exclude subthreshold clusters in dEdx calculation
-  ExcludeSectorBoundaries = 1 << 4, ///< flag to exclude sector boundary clusters in subthreshold cluster treatment
-  ExcludeSharedCl = 1 << 5,         ///< flag to exclude clusters shared between tracks
+  ExcludeSplitPadCl = 1 << 1,       ///< flag to exclude split pad clusters in dEdx calculation
+  ExcludeSplitTimeCl = 1 << 2,      ///< flag to exclude split time clusters in dEdx calculation
+  ExcludeSplitCl = 1 << 3,          ///< flag to exclude split pad or time clusters in dEdx calculation
+  ExcludeEdgeCl = 1 << 4,           ///< flag to exclude sector edge clusters in dEdx calculation
+  ExcludeSubthresholdCl = 1 << 5,   ///< flag to exclude subthreshold clusters in dEdx calculation
+  ExcludeSectorBoundaries = 1 << 6, ///< flag to exclude sector boundary clusters in subthreshold cluster treatment
+  ExcludeSharedCl = 1 << 7,         ///< flag to exclude clusters shared between tracks in dEdx calculation
+  ExcludeSamePadRowCl = 1 << 8,     ///< flag to exclude clusters in the same pad row in dEdx calculation
 };
 
 inline CorrectionFlags operator&(CorrectionFlags a, CorrectionFlags b) { return static_cast<CorrectionFlags>(static_cast<unsigned short>(a) & static_cast<unsigned short>(b)); }
@@ -95,6 +114,9 @@ class CalculatedEdx
   /// \param propagate propagate the tracks to extract the track parameters instead of performing a refit
   void setPropagateTrack(const bool propagate) { mPropagateTrack = propagate; }
 
+  /// \param propagate propagate the tracks to extract the track parameters instead of performing a refit
+  void setPropagateParams(const bool propagate) { mPropagateParams = propagate; }
+
   /// \param debug use debug streamer and set debug vectors
   void setDebug(const bool debug) { mDebug = debug; }
 
@@ -110,12 +132,22 @@ class CalculatedEdx
   /// \param minChargeMaxThreshold upper limit for the possible minimum charge max in subthreshold treatment
   void setMinChargeMaxThreshold(float minChargeMaxThreshold) { mMinChargeMaxThreshold = minChargeMaxThreshold; }
 
-  /// set the debug streamer
-  void setStreamer(const char* debugRootFile) { mStreamer = std::make_unique<o2::utils::TreeStreamRedirector>(debugRootFile, "recreate"); };
+  /// set the debug streamer for a given output file; a new streamer is only created the first time a given debugRootFile is seen,
+  /// so different calculatedEdx() calls using different debugRootFile names each get their own independent debug file
+  void setStreamer(const char* debugRootFile)
+  {
+    auto& streamer = mStreamers[debugRootFile];
+    if (!streamer) {
+      streamer = std::make_unique<o2::utils::TreeStreamRedirector>(debugRootFile, "recreate");
+    }
+  };
 
   /// set the debug streamer of the space-charge dedx correction
   void setSCStreamer(const char* debugRootFile = "debug_sc_corrections.root") { mSCdEdxCorrection.setStreamer(debugRootFile); }
 
+  /// \param lumi set luminosity for space-charge correction map scaling
+  void setLumi(const float lumi) { mSCdEdxCorrection.setLumi(lumi); }
+
   /// \return returns magnetic field in kG
   float getFieldNominalGPUBz() { return mFieldNominalGPUBz; }
 
@@ -128,18 +160,21 @@ class CalculatedEdx
   /// \return returns the upper limit for the possible minimum charge max in subthreshold treatment
   float getMinChargeMaxThreshold() { return mMinChargeMaxThreshold; }
 
-  /// fill missing clusters with minimum charge (method=0) or minimum charge/2 (method=1) or Landau (method=2)
+  /// fill missing clusters with minimum charge (method=0) or minimum charge/2 (method=1)
   void fillMissingClusters(int missingClusters[4], float minChargeTot, float minChargeMax, int method, std::array<std::vector<float>, 5>& chargeTotROC, std::array<std::vector<float>, 5>& chargeMaxROC);
 
+  void handleSameRowClusters(o2::tpc::TrackTPC& track, std::map<std::pair<unsigned char, unsigned char>, std::vector<int>>& clustersByRow, std::map<std::pair<unsigned char, unsigned char>, o2::tpc::ClusterNative>& combinedClustersByRow, std::map<int, std::tuple<unsigned char, unsigned char, unsigned int>>& clusterReferencesByIndex);
+
   /// get the truncated mean for the input track with the truncation range, charge type, region and corrections
   /// the cluster charge is normalized by effective length*gain, you can turn off the normalization by setting all corrections to false
   /// \param track input track
   /// \param output output dEdxInfo
+  /// \param averageOcc output average cluster occupancy of the track, per TPC region
   /// \param low lower cluster cut
   /// \param high higher cluster cut
-  /// \param mask to apply different corrections: TopologySimple = simple analytical topology correction, TopologyPol = topology correction from polynomials, GainFull = full gain map from calibration container,
+  /// \param correctionMask to apply different corrections: TopologySimple = simple analytical topology correction, TopologyPol = topology correction from polynomials, GainFull = full gain map from calibration container,
   ///                                                      GainResidual = residuals gain map from calibration container, dEdxResidual = residual dEdx correction
-  void calculatedEdx(TrackTPC& track, dEdxInfo& output, float low = 0.015f, float high = 0.6f, CorrectionFlags correctionMask = CorrectionFlags::TopologyPol | CorrectionFlags::dEdxResidual, ClusterFlags clusterMask = ClusterFlags::None, int subthresholdMethod = 0, const char* debugRootFile = "dEdxDebug.root");
+  void calculatedEdx(TrackTPC& track, dEdxInfo& output, AverageOccupancy& averageOcc, float low = 0.015f, float high = 0.6f, CorrectionFlags correctionMask = CorrectionFlags::TopologyPol | CorrectionFlags::dEdxResidual, ClusterFlags clusterMask = ClusterFlags::None, int subthresholdMethod = 0, int stackBoundaryMethod = 0, const char* debugRootFile = "dEdxDebug.root");
 
   /// get the truncated mean for the input charge vector and the truncation range low*nCl<nCl<high*nCl
   /// \param charge input vector
@@ -155,14 +190,24 @@ class CalculatedEdx
 
   /// get effective track length using topology correction from polynomials
   /// \param track input track
-  /// \param cl cluster
+  /// \param clPad cluster pad
+  /// \param clTime cluster time
   /// \param region pad region
   /// \param charge total or maximum charge of the cluster, cl
   /// \param chargeType total or maximum
   /// \param threshold zero supression threshold
   /// \return returns topology correction from polynomials
   float getTrackTopologyCorrectionPol(const o2::tpc::TrackTPC& track, const o2::tpc::ClusterNative& cl, const unsigned int region, const float charge, ChargeType chargeType, const float threshold) const;
 
+  /// \return returns space-charge dedx correctin
+  auto& getSCCorrection() { return mSCdEdxCorrection; }
+
+  /// \return returns cluster occupancy for given cluster time; only valid (non-sentinel) when the refit method is used, since the occupancy map is only filled by setRefit()
+  unsigned int getOccupancy(float clTime) const;
+
+  /// \return returns true if given row index is in a stack boundary
+  bool isInStackBoundaries(int stackNumber, unsigned char rowIndex, int stackBoundaryMethod);
+
   /// load calibration objects from CCDB
   /// \param runNumberOrTimeStamp run number or time stamp
   /// \param isMC set if dEdx space-charge corrections will be loaded for MC or real data
@@ -214,15 +259,6 @@ class CalculatedEdx
   /// \param object name of the object to load
   void setPropagatorFromFile(const char* folder, const char* file, const char* object);
 
-  /// \param lumi set luminosity for space-charge correction map scaling
-  void setLumi(const float lumi) { mSCdEdxCorrection.setLumi(lumi); }
-
-  /// \return returns space-charge dedx correctin
-  auto& getSCCorrection() { return mSCdEdxCorrection; }
-
-  /// \return returns cluster occupancy for given cluster
-  unsigned int getOccupancy(const o2::tpc::ClusterNative& cl) const;
-
  private:
   std::vector<TrackTPC>* mTracks{nullptr};                    ///< vector containing the tpc tracks which will be processed
   std::vector<TPCClRefElem>* mTPCTrackClIdxVecInput{nullptr}; ///< input vector with TPC tracks cluster indicies
@@ -233,18 +269,22 @@ class CalculatedEdx
   std::vector<unsigned int> mTPCRefitterOccMap;                  ///< externally set TPC clusters occupancy map
   std::unique_ptr<o2::gpu::GPUO2InterfaceRefit> mRefit{nullptr}; ///< TPC refitter used for TPC tracks refit during the reconstruction
 
-  int mMaxMissingCl{1};                                                ///< maximum number of missing clusters for subthreshold check
-  float mMinChargeTotThreshold{50};                                    ///< upper limit for minimum charge tot value in subthreshold treatment, i.e for a high dEdx track adding a minimum value of 500 to track as a virtual charge doesn't make sense
-  float mMinChargeMaxThreshold{50};                                    ///< upper limit for minimum charge max value in subthreshold treatment, i.e for a high dEdx track adding a minimum value of 500 to track as a virtual charge doesn't make sense
-  float mFieldNominalGPUBz{5};                                         ///< magnetic field in kG, used for track propagation
-  bool mPropagateTrack{false};                                         ///< propagating the track instead of performing a refit
-  bool mDebug{false};                                                  ///< use the debug streamer
-  CalibdEdxContainer mCalibCont;                                       ///< calibration container
-  std::unique_ptr<o2::utils::TreeStreamRedirector> mStreamer{nullptr}; ///< debug streamer
+  int mMaxMissingCl{1};                                                                         ///< maximum number of missing clusters for subthreshold check
+  float mMinChargeTotThreshold{50};                                                             ///< upper limit for minimum charge tot value in subthreshold treatment, i.e for a high dEdx track adding a minimum value of 500 to track as a virtual charge doesn't make sense
+  float mMinChargeMaxThreshold{50};                                                             ///< upper limit for minimum charge max value in subthreshold treatment, i.e for a high dEdx track adding a minimum value of 500 to track as a virtual charge doesn't make sense
+  float mFieldNominalGPUBz{5};                                                                  ///< magnetic field in kG, used for track propagation
+  bool mPropagateTrack{false};                                                                  ///< propagating the track instead of performing a refit (faster than refit)
+  bool mPropagateParams{false};                                                                 ///< propagating the parameters instead of full propagation (faster than track propagation)
+  bool mDebug{false};                                                                           ///< use the debug streamer
+  CalibdEdxContainer mCalibCont;                                                                ///< calibration container
+  std::unordered_map<std::string, std::unique_ptr<o2::utils::TreeStreamRedirector>> mStreamers; ///< debug streamers, keyed by output file name so each debugRootFile gets its own tree
+  long mDebugTrackIndex{-1};                                                                    ///< running index of the track being processed, written to the debug trees so per-cluster rows can be grouped back into tracks
 
   CorrectdEdxDistortions mSCdEdxCorrection; ///< for space-charge correction of dE/dx
+
+  std::array<std::vector<unsigned char>, 4> mStackBoundaries = {{{0, 62}, {63, 96}, {97, 126}, {127, 151}}}; // for excluding stack boundaries in dEdx calculation
 };
 
 } // namespace o2::tpc
 
-#endif
+#endif