Adding option to recalculate the inverse correction during reconstruc…

…tion

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

@@ -56,6 +56,7 @@ class CorrectionMapsLoader : public o2::gpu::CorrectionMapsHelper
   void updateVDrift(float vdriftCorr, float vdrifRef, float driftTimeOffset = 0);
   void init(o2::framework::InitContext& ic);
   void copySettings(const CorrectionMapsLoader& src);
+  void updateInverse(); /// recalculate inverse correction
 
   static void requestCCDBInputs(std::vector<o2::framework::InputSpec>& inputs, std::vector<o2::framework::ConfigParamSpec>& options, const CorrectionMapsLoaderGloOpts& gloOpts);
   static void addGlobalOptions(std::vector<o2::framework::ConfigParamSpec>& options);
@@ -68,6 +69,7 @@ class CorrectionMapsLoader : public o2::gpu::CorrectionMapsHelper
 
   float mInstLumiFactor = 1.0; // multiplicative factor for inst. lumi
   int mCTPLumiSource = 0;      // 0: main, 1: alternative CTP lumi source
+  int mNthreadsInv = 1;        // number of threads used for calculating the inverse correction
 #endif
 };
 

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

@@ -97,6 +97,12 @@ class TPCFastSpaceChargeCorrectionHelper
 
   void testGeometry(const TPCFastTransformGeo& geo) const;
 
+  /// initialise inverse transformation
+  void initInverse(o2::gpu::TPCFastSpaceChargeCorrection& correction, bool prn);
+
+  /// 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);
+
  private:
   /// geometry initialization
   void initGeometry();
@@ -107,9 +113,6 @@ class TPCFastSpaceChargeCorrectionHelper
   /// initialise max drift length
   void initMaxDriftLength(o2::gpu::TPCFastSpaceChargeCorrection& correction, bool prn);
 
-  /// initialise inverse transformation
-  void initInverse(o2::gpu::TPCFastSpaceChargeCorrection& correction, bool prn);
-
   static TPCFastSpaceChargeCorrectionHelper* sInstance; ///< singleton instance
   bool mIsInitialized = 0;                              ///< initialization flag
   int mNthreads = 1;                                    ///< n of threads to use

Detectors/TPC/calibration/src/CorrectionMapsLoader.cxx

@@ -23,6 +23,7 @@
 #include "Framework/DeviceSpec.h"
 #include "Framework/ConfigParamRegistry.h"
 #include "DataFormatsCTP/LumiInfo.h"
+#include "TPCCalibration/TPCFastSpaceChargeCorrectionHelper.h"
 
 using namespace o2::tpc;
 using namespace o2::framework;
@@ -70,6 +71,10 @@ void CorrectionMapsLoader::extractCCDBInputs(ProcessingContext& pc)
     setVShapeScaler(vshapescaler);
     pc.inputs().get<o2::gpu::TPCFastTransform*>("tpcCorrMapVShape");
   }
+  // update inverse in case it is requested
+  if (!mScaleInverse) {
+    updateInverse();
+  }
 }
 
 //________________________________________________________
@@ -110,6 +115,8 @@ void CorrectionMapsLoader::addOptions(std::vector<ConfigParamSpec>& options)
 {
   // these are options which should be added at the level of device using TPC corrections
   // At the moment - nothing, all options are moved to configurable param CorrMapParam
+  addOption(options, ConfigParamSpec{"recalculate-inverse-correction", o2::framework::VariantType::Bool, false, {"Recalculate the inverse correction in case linear scaling of corrections is used"}});
+  addOption(options, ConfigParamSpec{"nthreads-inverse-correction", o2::framework::VariantType::Int, -1, {"Number of threads used for calculating the inverse correction (-1=all threads)"}});
 }
 
 //________________________________________________________
@@ -220,6 +227,11 @@ void CorrectionMapsLoader::init(o2::framework::InitContext& ic)
       enableVShapeCorrection(true);
     }
   }
+  mScaleInverse = !(ic.options().get<bool>("recalculate-inverse-correction"));
+  mNthreadsInv = (ic.options().get<int>("nthreads-inverse-correction"));
+  if (!mScaleInverse) {
+    LOGP(info, "Recalculating the inverse correction for every TF");
+  }
 }
 
 //________________________________________________________
@@ -238,6 +250,21 @@ void CorrectionMapsLoader::copySettings(const CorrectionMapsLoader& src)
   mInstLumiFactor = src.mInstLumiFactor;
   mCTPLumiSource = src.mCTPLumiSource;
   mLumiScaleMode = src.mLumiScaleMode;
+  mScaleInverse = src.getScaleInverse();
+  mNthreadsInv = src.mNthreadsInv;
+}
+
+void CorrectionMapsLoader::updateInverse()
+{
+  (mNthreadsInv < 0) ? TPCFastSpaceChargeCorrectionHelper::instance()->setNthreadsToMaximum() : TPCFastSpaceChargeCorrectionHelper::instance()->setNthreads(mNthreadsInv);
+  if (mLumiScaleMode == 1 || mLumiScaleMode == 2) {
+    setUpdatedMap();
+    std::vector<float> scaling{1, mLumiScale, mVShapeScaler};
+    std::vector<o2::gpu::TPCFastSpaceChargeCorrection*> corr{&(mCorrMap->getCorrection()), &(mCorrMapRef->getCorrection()), &(mCorrMapVShape->getCorrection())};
+    TPCFastSpaceChargeCorrectionHelper::instance()->initInverse(corr, scaling, false);
+  } else {
+    LOGP(info, "Reinitializing inverse correction with lumi scale mode {} not supported for now", mLumiScaleMode);
+  }
 }
 
 #endif // #ifndef GPUCA_GPUCODE_DEVICE

Detectors/TPC/calibration/src/TPCFastSpaceChargeCorrectionHelper.cxx

@@ -28,6 +28,7 @@
 #include "Riostream.h"
 #include <fairlogger/Logger.h>
 #include <thread>
+#include "TStopwatch.h"
 
 using namespace o2::gpu;
 
@@ -653,10 +654,22 @@ void TPCFastSpaceChargeCorrectionHelper::initMaxDriftLength(o2::gpu::TPCFastSpac
 
 void TPCFastSpaceChargeCorrectionHelper::initInverse(o2::gpu::TPCFastSpaceChargeCorrection& correction, bool prn)
 {
-  /// initialise inverse transformation
+  std::vector<o2::gpu::TPCFastSpaceChargeCorrection*> corr{&correction};
+  initInverse(corr, std::vector<float>{1}, prn);
+}
 
+void TPCFastSpaceChargeCorrectionHelper::initInverse(std::vector<o2::gpu::TPCFastSpaceChargeCorrection*>& corrections, const std::vector<float> scaling, bool prn)
+{
+  /// initialise inverse transformation
+  TStopwatch watch;
   LOG(info) << "fast space charge correction helper: init inverse";
 
+  if (corrections.size() != scaling.size()) {
+    LOGP(error, "Input corrections and scaling values have different size");
+    return;
+  }
+
+  auto& correction = *(corrections.front());
   initMaxDriftLength(correction, prn);
 
   double tpcR2min = mGeo.getRowInfo(0).x - 1.;
@@ -709,6 +722,17 @@ void TPCFastSpaceChargeCorrectionHelper::initInverse(o2::gpu::TPCFastSpaceCharge
           for (double v = v0; v < v1 + stepV; v += stepV) {
             float dx, du, dv;
             correction.getCorrection(slice, row, u, v, dx, du, dv);
+            dx *= scaling[0];
+            du *= scaling[0];
+            dv *= scaling[0];
+            // add remaining corrections
+            for (int i = 1; i < corrections.size(); ++i) {
+              float dxTmp, duTmp, dvTmp;
+              corrections[i]->getCorrection(slice, row, u, v, dxTmp, duTmp, dvTmp);
+              dx += dxTmp * scaling[i];
+              du += duTmp * scaling[i];
+              dv += dvTmp * scaling[i];
+            }
             double cx = x + dx;
             double cu = u + du;
             double cv = v + dv;
@@ -789,6 +813,8 @@ void TPCFastSpaceChargeCorrectionHelper::initInverse(o2::gpu::TPCFastSpaceCharge
     }
 
   } // slice
+  float duration = watch.RealTime();
+  LOGP(info, "Inverse took: {}s", duration);
 }
 
 } // namespace tpc

GPU/TPCFastTransformation/CorrectionMapsHelper.cxx

@@ -30,6 +30,7 @@ void CorrectionMapsHelper::clear()
   mMeanLumi = 0.f;
   mMeanLumiRef = 0.f;
   mVShapeScaler = 0.f;
+  mScaleInverse = false;
 }
 
 void CorrectionMapsHelper::setOwner(bool v)
@@ -98,5 +99,5 @@ void CorrectionMapsHelper::setCorrMapVShape(std::unique_ptr<TPCFastTransform>&&
 //________________________________________________________
 void CorrectionMapsHelper::reportScaling()
 {
-  LOGP(info, "Map scaling update: InstLumiOverride={}, LumiScaleType={} -> instLumi={}, meanLumi={} -> LumiScale={}, lumiScaleMode={}, mVShapeScaler={}", getInstLumiOverride(), getLumiScaleType(), getInstLumi(), getMeanLumi(), getLumiScale(), getLumiScaleMode(), getVShapeScaler());
+  LOGP(info, "Map scaling update: InstLumiOverride={}, LumiScaleType={} -> instLumi={}, meanLumiRef={}, meanLumi={} -> LumiScale={}, lumiScaleMode={}, mVShapeScaler={}", getInstLumiOverride(), getLumiScaleType(), getInstLumi(), getMeanLumiRef(), getMeanLumi(), getLumiScale(), getLumiScaleMode(), getVShapeScaler());
 }

GPU/TPCFastTransformation/CorrectionMapsHelper.h

@@ -48,12 +48,12 @@ class CorrectionMapsHelper
 
   GPUd() void InverseTransformYZtoX(int slice, int row, float y, float z, float& x) const
   {
-    mCorrMap->InverseTransformYZtoX(slice, row, y, z, x, mCorrMapRef, mCorrMapVShape, mLumiScale, mVShapeScaler, mLumiScaleMode);
+    mCorrMap->InverseTransformYZtoX(slice, row, y, z, x, mCorrMapRef, mCorrMapVShape, (mScaleInverse ? mLumiScale : 0), (mScaleInverse ? mVShapeScaler : 0), mLumiScaleMode);
   }
 
   GPUd() void InverseTransformYZtoNominalYZ(int slice, int row, float y, float z, float& ny, float& nz) const
   {
-    mCorrMap->InverseTransformYZtoNominalYZ(slice, row, y, z, ny, nz, mCorrMapRef, mCorrMapVShape, mLumiScale, mVShapeScaler, mLumiScaleMode);
+    mCorrMap->InverseTransformYZtoNominalYZ(slice, row, y, z, ny, nz, mCorrMapRef, mCorrMapVShape, (mScaleInverse ? mLumiScale : 0), (mScaleInverse ? mVShapeScaler : 0), mLumiScaleMode);
   }
 
   GPUd() const GPUCA_NAMESPACE::gpu::TPCFastTransform* getCorrMap() const { return mCorrMap; }
@@ -163,6 +163,8 @@ class CorrectionMapsHelper
 
   int getUpdateFlags() const { return mUpdatedFlags; }
 
+  bool getScaleInverse() const { return mScaleInverse; }
+
  protected:
   enum UpdateFlags { MapBit = 0x1,
                      MapRefBit = 0x2,
@@ -181,6 +183,7 @@ class CorrectionMapsHelper
   float mInstLumiOverride = -1.f;                                  // optional value to override inst lumi
   bool mEnableVShape = false;                                      ///< use v shape correction
   float mVShapeScaler = 0;                                         // scaling value for V-shape distortions
+  bool mScaleInverse{false};                                       // if set to false the inverse correction is already scaled and will not scaled again
   GPUCA_NAMESPACE::gpu::TPCFastTransform* mCorrMap{nullptr};       // current transform
   GPUCA_NAMESPACE::gpu::TPCFastTransform* mCorrMapRef{nullptr};    // reference transform
   GPUCA_NAMESPACE::gpu::TPCFastTransform* mCorrMapVShape{nullptr}; // correction map for v-shape distortions on A-side

GPU/TPCFastTransformation/TPCFastTransform.h

@@ -500,12 +500,19 @@ GPUdi() void TPCFastTransform::TransformInternal(int slice, int row, float& u, f
       float lxT = x + dx;
       getGeometry().convUVtoLocal(slice, uCorr, vCorr, lyT, lzT);
 
+      float invYZtoXScaled;
+      InverseTransformYZtoX(slice, row, lyT, lzT, invYZtoXScaled, ref, ref2, scale, scale2, scaleMode);
+
       float invYZtoX;
-      InverseTransformYZtoX(slice, row, lyT, lzT, invYZtoX, ref, ref2, scale, scale2, scaleMode);
+      InverseTransformYZtoX(slice, row, lyT, lzT, invYZtoX);
 
       float YZtoNominalY;
       float YZtoNominalZ;
-      InverseTransformYZtoNominalYZ(slice, row, lyT, lzT, YZtoNominalY, YZtoNominalZ, ref, ref2, scale, scale2, scaleMode);
+      InverseTransformYZtoNominalYZ(slice, row, lyT, lzT, YZtoNominalY, YZtoNominalZ);
+
+      float YZtoNominalYScaled;
+      float YZtoNominalZScaled;
+      InverseTransformYZtoNominalYZ(slice, row, lyT, lzT, YZtoNominalYScaled, YZtoNominalZScaled, ref, ref2, scale, scale2, scaleMode);
 
       float dxRef, duRef, dvRef;
       if (ref) {
@@ -554,8 +561,11 @@ GPUdi() void TPCFastTransform::TransformInternal(int slice, int row, float& u, f
                                                                                          << "gz=" << gz
                                                                                          // some transformations which are applied
                                                                                          << "invYZtoX=" << invYZtoX
+                                                                                         << "invYZtoXScaled=" << invYZtoXScaled
                                                                                          << "YZtoNominalY=" << YZtoNominalY
+                                                                                         << "YZtoNominalYScaled=" << YZtoNominalYScaled
                                                                                          << "YZtoNominalZ=" << YZtoNominalZ
+                                                                                         << "YZtoNominalZScaled=" << YZtoNominalZScaled
                                                                                          << "scaleMode=" << scaleMode
                                                                                          << "\n";
     })