TPC splines: add a possibility to merge specific sectors

Author: cbmsw

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

@@ -41,6 +41,9 @@ using namespace o2::gpu;
 
 class TPCFastSpaceChargeCorrectionHelper
 {
+ public:
+  using SectorScales = std::array<double, TPCFastTransformGeo::getNumberOfSectors()>;
+
  public:
   /// _____________  Constructors / destructors __________________________
 
@@ -115,15 +118,32 @@ class TPCFastSpaceChargeCorrectionHelper
   /// initialise inverse transformation from linear combination of several input corrections
   void initInverse(std::vector<o2::gpu::TPCFastSpaceChargeCorrection*>& corrections, const std::vector<float>& scaling, bool prn);
 
-  /// merge several corrections
+  /// weighted add of several corrections
   /// \param mainCorrection main correction
   /// \param scale scaling factor for the main correction
   /// \param additionalCorrections vector of pairs of additional corrections and their scaling factors
-  /// \param prn printout flag
   /// \return main correction merged with additional corrections
+  void addCorrections(
+    o2::gpu::TPCFastSpaceChargeCorrection& mainCorrection, double scale,
+    const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, double>>& additionalCorrections);
+
+  /// weighted add of several corrections with sector-dependent scaling factors
+  /// \param mainCorrection main correction
+  /// \param scale scaling factor for the main correction
+  /// \param additionalCorrections vector of pairs of additional corrections and their sector-dependent scaling factors
+  /// \return main correction merged with additional corrections
+  void addCorrections(
+    o2::gpu::TPCFastSpaceChargeCorrection& mainCorrection, SectorScales scale,
+    const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, SectorScales>>& additionalCorrections);
+
+  /// merge of two corrections sector-wise
+  /// \param destinationCorrection main correction to which the source correction will be added
+  /// \param sourceCorrection correction to be added to the main correction
+  /// \param sectors vector of sector indices for which the correction will be added
+  /// \return main correction merged with the source correction
   void mergeCorrections(
-    o2::gpu::TPCFastSpaceChargeCorrection& mainCorrection, float scale,
-    const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, float>>& additionalCorrections, bool prn);
+    o2::gpu::TPCFastSpaceChargeCorrection& destinationCorrection, const o2::gpu::TPCFastSpaceChargeCorrection& sourceCorrection,
+    const std::vector<int>& sectors);
 
   /// how far the voxel mean is allowed to be outside of the voxel (1.1 means 10%)
   void setVoxelMeanValidityRange(double range)

Detectors/TPC/calibration/src/TPCFastSpaceChargeCorrectionHelper.cxx

@@ -1013,19 +1013,38 @@ void TPCFastSpaceChargeCorrectionHelper::initInverse(std::vector<o2::gpu::TPCFas
   LOGP(info, "Inverse tooks: {}s", duration);
 }
 
-void TPCFastSpaceChargeCorrectionHelper::mergeCorrections(
-  o2::gpu::TPCFastSpaceChargeCorrection& mainCorrection, float mainScale,
-  const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, float>>& additionalCorrections, bool /*prn*/)
+void TPCFastSpaceChargeCorrectionHelper::addCorrections(
+  o2::gpu::TPCFastSpaceChargeCorrection& mainCorrection, double mainScale,
+  const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, double>>& additionalCorrections)
 {
-  /// merge several corrections
+  /// weighted add of several corrections
+  SectorScales mainSectorScale;
+  mainSectorScale.fill(mainScale);
+  std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, SectorScales>> additionalSectorScales;
+  for (const auto& corr : additionalCorrections) {
+    SectorScales sectorScale;
+    sectorScale.fill(corr.second);
+    additionalSectorScales.emplace_back(corr.first, sectorScale);
+  }
+
+  addCorrections(mainCorrection, mainSectorScale, additionalSectorScales);
+}
+
+void TPCFastSpaceChargeCorrectionHelper::addCorrections(
+  o2::gpu::TPCFastSpaceChargeCorrection& mainCorrection, SectorScales mainScale,
+  const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, SectorScales>>& additionalCorrections)
+{
+  /// weighted add of several corrections
 
   TStopwatch watch;
-  LOG(info) << "fast space charge correction helper: Merge corrections";
+  LOG(info) << "fast space charge correction helper: Add corrections";
 
   const auto& geo = mainCorrection.getGeometry();
 
   for (int sector = 0; sector < geo.getNumberOfSectors(); sector++) {
 
+    float secMainScale = mainScale[sector];
+
     auto myThread = [&](int iThread) {
       for (int row = iThread; row < geo.getNumberOfRows(); row += mNthreads) {
         auto& rowInfo = mainCorrection.getRowInfo(row);
@@ -1040,8 +1059,11 @@ void TPCFastSpaceChargeCorrectionHelper::mergeCorrections(
         constexpr int nKnotPar3d = nKnotPar1d * 3;
 
         { // scale the main correction
-
-          double parscale[4] = {mainScale, mainScale, mainScale, mainScale * mainScale};
+          for (int i = 0; i < 3; i++) {
+            secRowInfo.maxCorr[i] *= secMainScale;
+            secRowInfo.minCorr[i] *= secMainScale;
+          }
+          double parscale[4] = {secMainScale, secMainScale, secMainScale, secMainScale * secMainScale};
           for (int iknot = 0, ind = 0; iknot < spline.getNumberOfKnots(); iknot++) {
             for (int ipar = 0; ipar < nKnotPar1d; ++ipar) {
               for (int idim = 0; idim < 3; idim++, ind++) {
@@ -1074,6 +1096,10 @@ void TPCFastSpaceChargeCorrectionHelper::mergeCorrections(
           const auto& corr = *(additionalCorrections[icorr].first);
           double scale = additionalCorrections[icorr].second;
           auto& linfo = corr.getRowInfo(row);
+          // double scale = additionalCorrections[icorr].second[sector];
+          // auto& linfo = corr.getSectorRowInfo(sector, row);
+          // secRowInfo.updateMaxValues(linfo.getMaxValues(), scale);
+          // secRowInfo.updateMaxValues(linfo.getMinValues(), scale);
 
           double scaleU = rowInfo.gridMeasured.getYscale() / linfo.gridMeasured.getYscale();
           double scaleV = rowInfo.gridMeasured.getZscale() / linfo.gridMeasured.getZscale();
@@ -1150,7 +1176,93 @@ void TPCFastSpaceChargeCorrectionHelper::mergeCorrections(
     }
 
   } // sector
-  float duration = watch.RealTime();
+  double duration = watch.RealTime();
+  LOGP(info, "Merge of corrections tooks: {}s", duration);
+}
+
+void TPCFastSpaceChargeCorrectionHelper::mergeCorrections(o2::gpu::TPCFastSpaceChargeCorrection& destinationCorrection,
+                                                          const o2::gpu::TPCFastSpaceChargeCorrection& sourceCorrection,
+                                                          const std::vector<int>& sectors)
+{
+  /// merge of two corrections sector-wise
+  TStopwatch watch;
+  LOG(info) << "fast space charge correction helper: Merge corrections";
+
+  const auto& geo = destinationCorrection.getGeometry();
+
+  for (int sector : sectors) {
+    if (sector < 0 || sector >= geo.getNumberOfSectors()) {
+      LOGP(fatal, "Invalid sector number {}. Valid range is [0, {})", sector, geo.getNumberOfSectors());
+      continue;
+    }
+    auto myThread = [&](int iThread) {
+      for (int row = iThread; row < geo.getNumberOfRows(); row += mNthreads) {
+
+        { // replace the direct correction
+          const auto& destSpline = destinationCorrection.getSpline(sector, row);
+          float* destSplineParameters = destinationCorrection.getCorrectionData(sector, row);
+          const auto& sourceSpline = sourceCorrection.getSpline(sector, row);
+          const float* sourceSplineParameters = sourceCorrection.getCorrectionData(sector, row);
+
+          // ensure the splines are compatible
+          if (destSpline.getGridX1().getNumberOfKnots() != sourceSpline.getGridX1().getNumberOfKnots() ||
+              destSpline.getGridX2().getNumberOfKnots() != sourceSpline.getGridX2().getNumberOfKnots()) {
+            LOGP(error, "Splines for sector {} row {} are not compatible: number of knots in U or V direction do not match", sector, row);
+            continue;
+          }
+          // replace the destination correction with the source correction for this sector and row
+          memcpy(destSplineParameters, sourceSplineParameters, destSpline.getNumberOfParameters() * sizeof(float));
+        }
+
+        { // replace the inverse correction X
+          const auto& destSpline = destinationCorrection.getSplineInvX(sector, row);
+          float* destSplineParameters = destinationCorrection.getCorrectionDataInvX(sector, row);
+          const auto& sourceSpline = sourceCorrection.getSplineInvX(sector, row);
+          const float* sourceSplineParameters = sourceCorrection.getCorrectionDataInvX(sector, row);
+          // ensure the splines are compatible
+          if (destSpline.getGridX1().getNumberOfKnots() != sourceSpline.getGridX1().getNumberOfKnots() ||
+              destSpline.getGridX2().getNumberOfKnots() != sourceSpline.getGridX2().getNumberOfKnots()) {
+            LOGP(error, "Inverse X splines for sector {} row {} are not compatible: number of knots in U or V direction do not match", sector, row);
+            continue;
+          }
+          memcpy(destSplineParameters, sourceSplineParameters, destSpline.getNumberOfParameters() * sizeof(float));
+        }
+
+        { // replace the inverse correction YZ
+          const auto& destSpline = destinationCorrection.getSplineInvYZ(sector, row);
+          float* destSplineParameters = destinationCorrection.getCorrectionDataInvYZ(sector, row);
+          const auto& sourceSpline = sourceCorrection.getSplineInvYZ(sector, row);
+          const float* sourceSplineParameters = sourceCorrection.getCorrectionDataInvYZ(sector, row);
+          // ensure the splines are compatible
+          if (destSpline.getGridX1().getNumberOfKnots() != sourceSpline.getGridX1().getNumberOfKnots() ||
+              destSpline.getGridX2().getNumberOfKnots() != sourceSpline.getGridX2().getNumberOfKnots()) {
+            LOGP(error, "Inverse YZ splines for sector {} row {} are not compatible: number of knots in U or V direction do not match", sector, row);
+            continue;
+          }
+          memcpy(destSplineParameters, sourceSplineParameters, destSpline.getNumberOfParameters() * sizeof(float));
+        }
+
+        // replace the sector row info
+        auto& destSecRowInfo = destinationCorrection.getSectorRowInfo(sector, row);
+        const auto& sourceSecRowInfo = sourceCorrection.getSectorRowInfo(sector, row);
+        destSecRowInfo = sourceSecRowInfo;
+      } // row
+    }; // thread
+
+    std::vector<std::thread> threads(mNthreads);
+
+    // run n threads
+    for (int i = 0; i < mNthreads; i++) {
+      threads[i] = std::thread(myThread, i);
+    }
+
+    // wait for the threads to finish
+    for (auto& th : threads) {
+      th.join();
+    }
+
+  } // sector
+  double duration = watch.RealTime();
   LOGP(info, "Merge of corrections tooks: {}s", duration);
 }