TPC: improving simulation of distortions in MC (#13345)

* TPC: improving simulation of distortions in MC

- adding option to recalculate (merging) the distortions in case of large scaling
- accessing the distortions of the derivative map at the distorted
position for better consistency with corrections

* TPC: Renaming variables

(cherry picked from commit 91dd4df)

Author: matthias-kleiner

Detectors/TPC/simulation/include/TPCSimulation/Digitizer.h

@@ -138,6 +138,9 @@ class Digitizer
   void setMeanLumiDistortions(float meanLumi);
   void setMeanLumiDistortionsDerivative(float meanLumi);
 
+  /// in case of scaled distortions, the distortions can be recalculated to ensure consistent distortions and corrections
+  void recalculateDistortions();
+
  private:
   DigitContainer mDigitContainer;      ///< Container for the Digits
   std::unique_ptr<SC> mSpaceCharge;    ///< Handler of full distortions (static + IR dependant)
@@ -151,7 +154,8 @@ class Digitizer
   bool mUseSCDistortions = false;      ///< Flag to switch on the use of space-charge distortions
   int mDistortionScaleType = 0;        ///< type=0: no scaling of distortions, type=1 distortions without any scaling, type=2 distortions scaling with lumi
   float mLumiScaleFactor = 0;          ///< value used to scale the derivative map
-  ClassDefNV(Digitizer, 2);
+  bool mUseScaledDistortions = false;  ///< whether the distortions are already scaled
+  ClassDefNV(Digitizer, 3);
 };
 } // namespace tpc
 } // namespace o2

Detectors/TPC/simulation/src/Digitizer.cxx

@@ -83,7 +83,7 @@ void Digitizer::process(const std::vector<o2::tpc::HitGroup>& hits,
       if (mDistortionScaleType == 1) {
         mSpaceCharge->distortElectron(posEle);
       } else if (mDistortionScaleType == 2) {
-        mSpaceCharge->distortElectron(posEle, mSpaceChargeDer.get(), mLumiScaleFactor);
+        mSpaceCharge->distortElectron(posEle, (mUseScaledDistortions ? nullptr : mSpaceChargeDer.get()), mLumiScaleFactor);
       }
 
       /// Remove electrons that end up more than three sigma of the hit's average diffusion away from the current sector
@@ -237,3 +237,33 @@ void Digitizer::setMeanLumiDistortionsDerivative(float meanLumi)
 {
   mSpaceChargeDer->setMeanLumi(meanLumi);
 }
+
+void Digitizer::recalculateDistortions()
+{
+  if (!mSpaceCharge || !mSpaceChargeDer) {
+    LOGP(info, "Average or derivative distortions not set");
+    return;
+  }
+
+  // recalculate distortions only in case the inst lumi differs from the avg lumi
+  if (mSpaceCharge->getMeanLumi() != CorrMapParam::Instance().lumiInst) {
+    for (int iside = 0; iside < 2; ++iside) {
+      const o2::tpc::Side side = (iside == 0) ? Side::A : Side::C;
+      // this needs to be done only once
+      LOGP(info, "Calculating corrections for average distortions");
+      mSpaceCharge->calcGlobalCorrWithGlobalDistIterative(side, nullptr, 0);
+
+      LOGP(info, "Calculating corrections for derivative distortions");
+      mSpaceChargeDer->calcGlobalCorrWithGlobalDistIterative(side, nullptr, 0);
+
+      LOGP(info, "Calculating scaled distortions with scaling factor {}", mLumiScaleFactor);
+      mSpaceCharge->calcGlobalDistWithGlobalCorrIterative(side, mSpaceChargeDer.get(), mLumiScaleFactor);
+    }
+    // set new lumi of avg map
+    mSpaceCharge->setMeanLumi(CorrMapParam::Instance().lumiInst);
+  } else {
+    LOGP(info, "Inst. lumi {} is same as mean lumi {}. Skip recalculation of distortions", CorrMapParam::Instance().lumiInst, mSpaceCharge->getMeanLumi());
+  }
+
+  mUseScaledDistortions = true;
+}

Detectors/TPC/spacecharge/include/TPCSpaceCharge/SpaceCharge.h

@@ -354,7 +354,20 @@ class SpaceCharge
   /// \param approachR when the difference between the desired r coordinate and the position of the global correction is deltaR, approach the desired r coordinate by deltaR * \p approachR.
   /// \param approachPhi when the difference between the desired phi coordinate and the position of the global correction is deltaPhi, approach the desired phi coordinate by deltaPhi * \p approachPhi.
   /// \param diffCorr if the absolute differences from the interpolated values for the global corrections from the last iteration compared to the current iteration is smaller than this value, set converged to true for current global distortion
-  void calcGlobalDistWithGlobalCorrIterative(const DistCorrInterpolator<DataT>& globCorr, const int maxIter = 100, const DataT approachZ = 0.5, const DataT approachR = 0.5, const DataT approachPhi = 0.5, const DataT diffCorr = 1e-6);
+  /// \param type whether to calculate distortions or corrections
+  void calcGlobalDistWithGlobalCorrIterative(const DistCorrInterpolator<DataT>& globCorr, const int maxIter = 100, const DataT approachZ = 1, const DataT approachR = 1, const DataT approachPhi = 1, const DataT diffCorr = 50e-6, const SpaceCharge<DataT>* scSCale = nullptr, float scale = 0);
+
+  /// step 5: calculate global distortions using the global corrections (FAST)
+  /// \param scSCale possible second sc object
+  /// \param scale scaling for second sc object
+  void calcGlobalDistWithGlobalCorrIterative(const Side side, const SpaceCharge<DataT>* scSCale = nullptr, float scale = 0, const int maxIter = 100, const DataT approachZ = 1, const DataT approachR = 1, const DataT approachPhi = 1, const DataT diffCorr = 50e-6);
+  void calcGlobalDistWithGlobalCorrIterative(const SpaceCharge<DataT>* scSCale = nullptr, float scale = 0, const int maxIter = 100, const DataT approachZ = 1, const DataT approachR = 1, const DataT approachPhi = 1, const DataT diffCorr = 50e-6);
+
+  /// calculate global corrections from global distortions
+  /// \param scSCale possible second sc object
+  /// \param scale scaling for second sc object
+  void calcGlobalCorrWithGlobalDistIterative(const Side side, const SpaceCharge<DataT>* scSCale = nullptr, float scale = 0, const int maxIter = 100, const DataT approachZ = 1, const DataT approachR = 1, const DataT approachPhi = 1, const DataT diffCorr = 50e-6);
+  void calcGlobalCorrWithGlobalDistIterative(const SpaceCharge<DataT>* scSCale = nullptr, float scale = 0, const int maxIter = 100, const DataT approachZ = 1, const DataT approachR = 1, const DataT approachPhi = 1, const DataT diffCorr = 50e-6);
 
   /// \return returns number of vertices in z direction
   unsigned short getNZVertices() const { return mParamGrid.NZVertices; }
@@ -1359,6 +1372,8 @@ class SpaceCharge
   /// set potentialsdue to ROD misalignment
   void initRodAlignmentVoltages(const MisalignmentType misalignmentType, const FCType fcType, const int sector, const Side side, const float deltaPot);
 
+  void calcGlobalDistCorrIterative(const DistCorrInterpolator<DataT>& globCorr, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr, const SpaceCharge<DataT>* scSCale, float scale, const Type type);
+
   ClassDefNV(SpaceCharge, 6);
 };
 

Detectors/TPC/spacecharge/src/SpaceCharge.cxx

@@ -697,12 +697,59 @@ void SpaceCharge<DataT>::calcEField(const Side side)
 }
 
 template <typename DataT>
-void SpaceCharge<DataT>::calcGlobalDistWithGlobalCorrIterative(const DistCorrInterpolator<DataT>& globCorr, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr)
+void SpaceCharge<DataT>::calcGlobalDistWithGlobalCorrIterative(const DistCorrInterpolator<DataT>& globCorr, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr, const SpaceCharge<DataT>* scSCale, float scale)
+{
+  calcGlobalDistCorrIterative(globCorr, maxIter, approachZ, approachR, approachPhi, diffCorr, scSCale, scale, Type::Distortions);
+}
+
+template <typename DataT>
+void SpaceCharge<DataT>::calcGlobalDistWithGlobalCorrIterative(const Side side, const SpaceCharge<DataT>* scSCale, float scale, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr)
+{
+  calcGlobalDistCorrIterative(getGlobalCorrInterpolator(side), maxIter, approachZ, approachR, approachPhi, diffCorr, scSCale, scale, Type::Distortions);
+}
+
+template <typename DataT>
+void SpaceCharge<DataT>::calcGlobalDistWithGlobalCorrIterative(const SpaceCharge<DataT>* scSCale, float scale, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr)
+{
+#pragma omp parallel for num_threads(sNThreads)
+  for (int iside = 0; iside < FNSIDES; ++iside) {
+    const o2::tpc::Side side = (iside == 0) ? Side::A : Side::C;
+    calcGlobalDistWithGlobalCorrIterative(side, scSCale, scale, maxIter, approachZ, approachR, approachPhi, diffCorr);
+  }
+}
+
+template <typename DataT>
+void SpaceCharge<DataT>::calcGlobalCorrWithGlobalDistIterative(const Side side, const SpaceCharge<DataT>* scSCale, float scale, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr)
+{
+  calcGlobalDistCorrIterative(getGlobalDistInterpolator(side), maxIter, approachZ, approachR, approachPhi, diffCorr, scSCale, scale, Type::Corrections);
+}
+
+template <typename DataT>
+void SpaceCharge<DataT>::calcGlobalCorrWithGlobalDistIterative(const SpaceCharge<DataT>* scSCale, float scale, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr)
+{
+#pragma omp parallel for num_threads(sNThreads)
+  for (int iside = 0; iside < FNSIDES; ++iside) {
+    const o2::tpc::Side side = (iside == 0) ? Side::A : Side::C;
+    calcGlobalCorrWithGlobalDistIterative(side, scSCale, scale, maxIter, approachZ, approachR, approachPhi, diffCorr);
+  }
+}
+
+template <typename DataT>
+void SpaceCharge<DataT>::calcGlobalDistCorrIterative(const DistCorrInterpolator<DataT>& globCorr, const int maxIter, const DataT approachZ, const DataT approachR, const DataT approachPhi, const DataT diffCorr, const SpaceCharge<DataT>* scSCale, float scale, const Type type)
 {
   const Side side = globCorr.getSide();
-  initContainer(mGlobalDistdR[side], true);
-  initContainer(mGlobalDistdZ[side], true);
-  initContainer(mGlobalDistdRPhi[side], true);
+  if (type == Type::Distortions) {
+    initContainer(mGlobalDistdR[side], true);
+    initContainer(mGlobalDistdZ[side], true);
+    initContainer(mGlobalDistdRPhi[side], true);
+  } else {
+    initContainer(mGlobalCorrdR[side], true);
+    initContainer(mGlobalCorrdZ[side], true);
+    initContainer(mGlobalCorrdRPhi[side], true);
+  }
+
+  const auto& scSCaleInterpolator = (type == Type::Distortions) ? scSCale->mInterpolatorGlobalCorr[side] : scSCale->mInterpolatorGlobalDist[side];
+
 #pragma omp parallel for num_threads(sNThreads)
   for (unsigned int iPhi = 0; iPhi < mParamGrid.NPhiVertices; ++iPhi) {
     const DataT phi = getPhiVertex(iPhi, side);
@@ -754,13 +801,25 @@ void SpaceCharge<DataT>::calcGlobalDistWithGlobalCorrIterative(const DistCorrInt
 
           // interpolate global correction at new point and calculate position of global correction
           corrdR = globCorr.evaldR(zCurrPos, rCurrPos, phiCurrPos);
+          if (scSCale && scale != 0) {
+            corrdR += scale * scSCaleInterpolator.evaldR(zCurrPos, rCurrPos, phiCurrPos);
+          }
           const DataT rNewPos = rCurrPos + corrdR;
 
-          const DataT corrPhi = globCorr.evaldRPhi(zCurrPos, rCurrPos, phiCurrPos) / rCurrPos;
+          DataT corrPhi = 0;
+          if (scSCale && scale != 0) {
+            corrPhi = scale * scSCaleInterpolator.evaldRPhi(zCurrPos, rCurrPos, phiCurrPos);
+          }
+          corrPhi += globCorr.evaldRPhi(zCurrPos, rCurrPos, phiCurrPos);
+          corrPhi /= rCurrPos;
+
           corrdRPhi = corrPhi * rNewPos; // normalize to new r coordinate
           const DataT phiNewPos = phiCurrPos + corrPhi;
 
           corrdZ = globCorr.evaldZ(zCurrPos, rCurrPos, phiCurrPos);
+          if (scSCale && scale != 0) {
+            corrdZ += scale * scSCaleInterpolator.evaldZ(zCurrPos, rCurrPos, phiCurrPos);
+          }
           const DataT zNewPos = zCurrPos + corrdZ;
 
           // approach desired coordinate
@@ -783,15 +842,27 @@ void SpaceCharge<DataT>::calcGlobalDistWithGlobalCorrIterative(const DistCorrInt
           lastCorrdRPhi = corrdRPhi;
         }
         // set global distortions if algorithm converged or iterations exceed max numbers of iterations
-        mGlobalDistdR[side](iZ, iR, iPhi) = -corrdR;
-        mGlobalDistdRPhi[side](iZ, iR, iPhi) = -corrdRPhi;
-        mGlobalDistdZ[side](iZ, iR, iPhi) = -corrdZ;
+        if (type == Type::Distortions) {
+          mGlobalDistdR[side](iZ, iR, iPhi) = -corrdR;
+          mGlobalDistdRPhi[side](iZ, iR, iPhi) = -corrdRPhi;
+          mGlobalDistdZ[side](iZ, iR, iPhi) = -corrdZ;
+        } else {
+          mGlobalCorrdR[side](iZ, iR, iPhi) = -corrdR;
+          mGlobalCorrdRPhi[side](iZ, iR, iPhi) = -corrdRPhi;
+          mGlobalCorrdZ[side](iZ, iR, iPhi) = -corrdZ;
+        }
       }
     }
     for (unsigned int iR = 0; iR < mParamGrid.NRVertices; ++iR) {
-      mGlobalDistdR[side](0, iR, iPhi) = 3 * (mGlobalDistdR[side](1, iR, iPhi) - mGlobalDistdR[side](2, iR, iPhi)) + mGlobalDistdR[side](3, iR, iPhi);
-      mGlobalDistdRPhi[side](0, iR, iPhi) = 3 * (mGlobalDistdRPhi[side](1, iR, iPhi) - mGlobalDistdRPhi[side](2, iR, iPhi)) + mGlobalDistdRPhi[side](3, iR, iPhi);
-      mGlobalDistdZ[side](0, iR, iPhi) = 3 * (mGlobalDistdZ[side](1, iR, iPhi) - mGlobalDistdZ[side](2, iR, iPhi)) + mGlobalDistdZ[side](3, iR, iPhi);
+      if (type == Type::Distortions) {
+        mGlobalDistdR[side](0, iR, iPhi) = 3 * (mGlobalDistdR[side](1, iR, iPhi) - mGlobalDistdR[side](2, iR, iPhi)) + mGlobalDistdR[side](3, iR, iPhi);
+        mGlobalDistdRPhi[side](0, iR, iPhi) = 3 * (mGlobalDistdRPhi[side](1, iR, iPhi) - mGlobalDistdRPhi[side](2, iR, iPhi)) + mGlobalDistdRPhi[side](3, iR, iPhi);
+        mGlobalDistdZ[side](0, iR, iPhi) = 3 * (mGlobalDistdZ[side](1, iR, iPhi) - mGlobalDistdZ[side](2, iR, iPhi)) + mGlobalDistdZ[side](3, iR, iPhi);
+      } else {
+        mGlobalCorrdR[side](0, iR, iPhi) = 3 * (mGlobalCorrdR[side](1, iR, iPhi) - mGlobalCorrdR[side](2, iR, iPhi)) + mGlobalCorrdR[side](3, iR, iPhi);
+        mGlobalCorrdRPhi[side](0, iR, iPhi) = 3 * (mGlobalCorrdRPhi[side](1, iR, iPhi) - mGlobalCorrdRPhi[side](2, iR, iPhi)) + mGlobalCorrdRPhi[side](3, iR, iPhi);
+        mGlobalCorrdZ[side](0, iR, iPhi) = 3 * (mGlobalCorrdZ[side](1, iR, iPhi) - mGlobalCorrdZ[side](2, iR, iPhi)) + mGlobalCorrdZ[side](3, iR, iPhi);
+      }
     }
   }
 }
@@ -1535,7 +1606,7 @@ void SpaceCharge<DataT>::distortElectron(GlobalPosition3D& point, const SpaceCha
 
   // scale distortions if requested
   if (scSCale && scale != 0) {
-    scSCale->getDistortions(point.X(), point.Y(), point.Z(), side, distXTmp, distYTmp, distZTmp);
+    scSCale->getDistortions(point.X() + distX, point.Y() + distY, point.Z() + distZ, side, distXTmp, distYTmp, distZTmp);
     distX += distXTmp * scale;
     distY += distYTmp * scale;
     distZ += distZTmp * scale;

Steer/DigitizerWorkflow/src/TPCDigitizerSpec.cxx

@@ -120,6 +120,9 @@ class TPCDPLDigitizerTask : public BaseDPLDigitizer
 
     mWithMCTruth = o2::conf::DigiParams::Instance().mctruth;
     auto triggeredMode = ic.options().get<bool>("TPCtriggered");
+    mRecalcDistortions = !(ic.options().get<bool>("do-not-recalculate-distortions"));
+    const int nthreadsDist = ic.options().get<int>("n-threads-distortions");
+    SC::setNThreads(nthreadsDist);
     mUseCalibrationsFromCCDB = ic.options().get<bool>("TPCuseCCDB");
     mMeanLumiDistortions = ic.options().get<float>("meanLumiDistortions");
     mMeanLumiDistortionsDerivative = ic.options().get<float>("meanLumiDistortionsDerivative");
@@ -220,6 +223,9 @@ class TPCDPLDigitizerTask : public BaseDPLDigitizer
       if (mDistortionType == 2) {
         pc.inputs().get<SC*>("tpcdistortionsderiv");
         mDigitizer.setLumiScaleFactor();
+        if (mRecalcDistortions) {
+          mDigitizer.recalculateDistortions();
+        }
       }
     }
 
@@ -475,6 +481,7 @@ class TPCDPLDigitizerTask : public BaseDPLDigitizer
   int mDistortionType = 0;
   float mMeanLumiDistortions = -1;
   float mMeanLumiDistortionsDerivative = -1;
+  bool mRecalcDistortions = false;
 };
 
 o2::framework::DataProcessorSpec getTPCDigitizerSpec(int channel, bool writeGRP, bool mctruth, bool internalwriter, int distortionType)
@@ -513,6 +520,8 @@ o2::framework::DataProcessorSpec getTPCDigitizerSpec(int channel, bool writeGRP,
       {"TPCuseCCDB", VariantType::Bool, false, {"true: load calibrations from CCDB; false: use random calibratoins"}},
       {"meanLumiDistortions", VariantType::Float, -1.f, {"override lumi of distortion object if >=0"}},
       {"meanLumiDistortionsDerivative", VariantType::Float, -1.f, {"override lumi of derivative distortion object if >=0"}},
+      {"do-not-recalculate-distortions", VariantType::Bool, false, {"Do not recalculate the distortions"}},
+      {"n-threads-distortions", VariantType::Int, 4, {"Number of threads used for the calculation of the distortions"}},
     }};
 }