TPC space-charge: Improve GEM frame charging-up distortions

Adding:
- downsampling of space-charge objects
- simulation of n-sectors only
- some helper functions
- weighted filling of charging-up of GEM frames for smoother potential
- set global distortions from function

Author: matthias-kleiner

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

@@ -186,6 +186,9 @@ struct DataContainer3D {
   /// print the matrix
   void print() const;
 
+  /// convert a data container to a new datacontainer with different grid definition (e.g. different number of vertices)
+  DataContainer3D<DataT> convert(const o2::tpc::RegularGrid3D<DataT>& gridNew, const o2::tpc::RegularGrid3D<DataT>& gridRef, const int threads = 1) const;
+
   /// operator overload
   DataContainer3D<DataT>& operator*=(const DataT value);
   DataContainer3D<DataT>& operator+=(const DataContainer3D<DataT>& other);

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

@@ -19,6 +19,7 @@
 #define ALICEO2_TPC_POISSONSOLVERHELPERS_H_
 
 #include "CommonConstants/MathConstants.h"
+#include "DataFormatsTPC/Defs.h"
 
 namespace o2
 {
@@ -55,7 +56,7 @@ struct MGParameters {                                             ///< Parameter
   inline static int nMGCycle = 200;                               ///< number of multi grid cycle (V type)
   inline static int maxLoop = 7;                                  ///< the number of tree-deep of multi grid
   inline static int gamma = 1;                                    ///< number of iteration at coarsest level !TODO SET TO REASONABLE VALUE!
-  inline static bool normalizeGridToOneSector = false;            ///< the grid in phi direction is squashed from 2 Pi to (2 Pi / SECTORSPERSIDE). This can used to get the potential for phi symmetric sc density or boundary potentials
+  inline static int normalizeGridToNSector = SECTORSPERSIDE;      ///< the grid in phi direction is squashed from 2 Pi to (2 Pi / SECTORSPERSIDE). This can used to get the potential for phi symmetric sc density or boundary potentials
 };
 
 template <typename DataT = double>

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

@@ -204,10 +204,13 @@ class SpaceCharge
 
   /// simulate only one sector instead of 18 per side. This makes currently only sense for the static distortions (ToDo: simplify usage)
   /// phi max will be restricted to 2Pi/18 for this instance and for global instance of poisson solver
-  void setSimOneSector();
+  void setSimOneSector() { setSimNSector(1); }
+
+  /// simulate N sectors
+  void setSimNSector(const int nSectors);
 
   /// unsetting simulation of one sector
-  static void unsetSimOneSector();
+  static void unsetSimNSector();
 
   /// setting default potential (same potential for all GEM frames. The default value of 1000V are matched to distortions observed in laser data without X-Ray etc.
   /// \param side side of the TPC where the potential will be set
@@ -308,10 +311,24 @@ class SpaceCharge
   /// scaling the space-charge density for given stack
   void scaleChargeDensityStack(const float scalingFactor, const Sector sector, const GEMstack stack);
 
+  /// scale the potential by a scaling factor
+  /// \param scalingFactor factor to scale the potential
+  /// \param side side for which the potential will be scaled
+  void scalePotential(const DataT scalingFactor, const Side side) { mPotential[side] *= scalingFactor; }
+
   /// add space charge density from other object (this.mDensity = this.mDensity + other.mDensity)
   /// \param otherSC other space-charge object, which charge will be added to current object
   void addChargeDensity(const SpaceCharge<DataT>& otherSC);
 
+  /// add global corrections from other space charge object
+  void addGlobalCorrections(const SpaceCharge<DataT>& otherSC, const Side side);
+
+  /// convert space-charge object to new definition of number of vertices
+  /// \param nZNew new number of vertices in z direction
+  /// \param nRNew new number of vertices in r direction
+  /// \param nPhiNew new number of vertices in phi direction
+  void downSampleObject(const int nZNew, const int nRNew, const int nPhiNew);
+
   /// step 3: calculate the local distortions and corrections with an electric field
   /// \param type calculate local corrections or local distortions: type = o2::tpc::SpaceCharge<>::Type::Distortions or o2::tpc::SpaceCharge<>::Type::Corrections
   /// \param formulaStruct struct containing a method to evaluate the electric field Er, Ez, Ephi (analytical formula or by TriCubic interpolator)
@@ -415,6 +432,9 @@ class SpaceCharge
   /// \param phi global phi coordinate
   DataT getDensityCyl(const DataT z, const DataT r, const DataT phi, const Side side) const;
 
+  /// get the potential for list of given coordinate
+  std::vector<float> getDensityCyl(const std::vector<DataT>& z, const std::vector<DataT>& r, const std::vector<DataT>& phi, const Side side) const;
+
   /// get the potential for given coordinate
   /// \param z global z coordinate
   /// \param r global r coordinate
@@ -1184,6 +1204,10 @@ class SpaceCharge
   /// \param gCorr function returning global corrections for given global coordinate
   void setGlobalCorrections(const std::function<void(int sector, DataT gx, DataT gy, DataT gz, DataT& gCx, DataT& gCy, DataT& gCz)>& gCorr, const Side side);
 
+  /// setting the global distortions directly from input function provided in global coordinates
+  /// \param gDist function returning global distortions for given global coordinate
+  void setGlobalDistortions(const std::function<void(int sector, DataT gx, DataT gy, DataT gz, DataT& gCx, DataT& gCy, DataT& gCz)>& gDist, const Side side);
+
   /// set misalignment of ROC for shift in z
   /// \param sector sector for which the misalignment in z will be applied (if sector=-1 all sectors are shifted)
   /// \param type 0=IROC, 1=OROC, 2=IROC+OROC
@@ -1229,7 +1253,16 @@ class SpaceCharge
   /// \param tgl tgl of the track
   /// \param nPoints number of points used to calculate the DCAr
   /// \param pcstream if provided debug output is being created
-  float getDCAr(float tgl, const int nPoints, const float phi, o2::utils::TreeStreamRedirector* pcstream = nullptr) const;
+  float getDCAr(float tgl, const int nPoints, const float phi, float rStart = -1, o2::utils::TreeStreamRedirector* pcstream = nullptr) const;
+
+  /// \return returns nearest phi vertex for given phi position
+  size_t getNearestPhiVertex(const DataT phi, const Side side) const { return std::round(phi / getGridSpacingPhi(side)); }
+
+  /// \return returns nearest r vertex for given radius position
+  size_t getNearestRVertex(const DataT r, const Side side) const { return std::round((r - getRMin(side)) / getGridSpacingR(side) + 0.5); }
+
+  /// \return returns number of bins in phi direction for the gap between sectors and for the GEM frame
+  size_t getPhiBinsGapFrame(const Side side) const;
 
  private:
   ParamSpaceCharge mParamGrid{};                                                                          ///< parameters of the grid on which the calculations are performed
@@ -1352,15 +1385,6 @@ class SpaceCharge
   /// dump the created electron tracks with calculateElectronDriftPath function to a tree
   void dumpElectronTracksToTree(const std::vector<std::pair<std::vector<o2::math_utils::Point3D<float>>, std::array<DataT, 3>>>& electronTracks, const int nSamplingPoints, const char* outFile) const;
 
-  /// \return returns nearest phi vertex for given phi position
-  size_t getNearestPhiVertex(const DataT phi, const Side side) const { return std::round(phi / getGridSpacingPhi(side)); }
-
-  /// \return returns nearest r vertex for given radius position
-  size_t getNearestRVertex(const DataT r, const Side side) const { return std::round((r - getRMin(side)) / getGridSpacingR(side) + 0.5); }
-
-  /// \return returns number of bins in phi direction for the gap between sectors and for the GEM frame
-  size_t getPhiBinsGapFrame(const Side side) const;
-
   /// \return setting the boundary potential for given GEM stack
   void setPotentialBoundaryGEMFrameAlongPhi(const std::function<DataT(DataT)>& potentialFunc, const GEMstack stack, const bool bottom, const Side side, const bool outerFrame = false);
 
@@ -1372,23 +1396,25 @@ class SpaceCharge
 
   void initAllBuffers();
 
-  void setBoundaryFromIndices(const std::function<DataT(DataT)>& potentialFunc, const std::vector<size_t>& indices, const Side side);
+  void setBoundaryFromIndices(const std::function<DataT(DataT)>& potentialFunc, const std::vector<std::pair<size_t, float>>& indices, const Side side);
 
   /// get indices of the GEM frame along r
-  std::vector<size_t> getPotentialBoundaryGEMFrameAlongRIndices(const Side side) const;
+  std::vector<std::pair<size_t, float>> getPotentialBoundaryGEMFrameAlongRIndices(const Side side) const;
 
   /// get indices of the GEM frame along phi
-  std::vector<size_t> getPotentialBoundaryGEMFrameAlongPhiIndices(const GEMstack stack, const bool bottom, const Side side, const bool outerFrame, const bool noGap = false) const;
+  std::vector<std::pair<size_t, float>> getPotentialBoundaryGEMFrameAlongPhiIndices(const GEMstack stack, const bool bottom, const Side side, const bool outerFrame, const bool noGap = false) const;
 
   void setROCMisalignment(int stackType, int misalignmentType, int sector, const float potMin, const float potMax);
-  void fillROCMisalignment(const std::vector<size_t>& indicesTop, const std::vector<size_t>& indicesBottom, int sector, int misalignmentType, const std::pair<float, float>& deltaPotPar);
+  void fillROCMisalignment(const std::vector<std::pair<size_t, float>>& indicesTop, const std::vector<std::pair<size_t, float>>& indicesBottom, int sector, int misalignmentType, const std::pair<float, float>& deltaPotPar);
 
   /// 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);
   void calcGlobalDistCorrIterativeLinearCartesian(const DistCorrInterpolator<DataT>& globCorr, const int maxIter, const DataT approachX, const DataT approachY, const DataT approachZ, const DataT diffCorr, const SpaceCharge<DataT>* scSCale, float scale, const Type type);
 
+  void setGlobalDistCorr(const Type type, const std::function<void(int sector, DataT gx, DataT gy, DataT gz, DataT& gCx, DataT& gCy, DataT& gCz)>& gFunc, const Side side);
+
   ClassDefNV(SpaceCharge, 6);
 };
 

Detectors/TPC/spacecharge/src/DataContainer3D.cxx

@@ -331,7 +331,6 @@ void DataContainer3D<DataT>::dumpSlice(std::string_view treename, std::string_vi
   ROOT::RDataFrame dFrame(treename, fileIn);
 
   auto df = dFrame.Define("slice", [rangeiZ, rangeiR, rangeiPhi](const std::pair<long, std::vector<float>>& values, unsigned short nz, unsigned short nr, unsigned short nphi) {
-    const bool simOneSectorOnly = MGParameters::normalizeGridToOneSector;
     std::vector<size_t> ir;
     std::vector<size_t> iphi;
     std::vector<size_t> iz;
@@ -370,12 +369,12 @@ void DataContainer3D<DataT>::dumpSlice(std::string_view treename, std::string_vi
 
       const float rTmp = o2::tpc::GridProperties<float>::getRMin() + o2::tpc::GridProperties<float>::getGridSpacingR(nr) * iRTmp;
       const float zTmp = o2::tpc::GridProperties<float>::getZMin() + o2::tpc::GridProperties<float>::getGridSpacingZ(nz) * iZTmp;
-      const float phiTmp = o2::tpc::GridProperties<float>::getPhiMin() + o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) / (simOneSectorOnly ? SECTORSPERSIDE : 1) * iPhiTmp;
+      const float phiTmp = o2::tpc::GridProperties<float>::getPhiMin() + o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) * (MGParameters::normalizeGridToNSector / double(SECTORSPERSIDE)) * iPhiTmp;
 
       const float x = rTmp * std::cos(phiTmp);
       const float y = rTmp * std::sin(phiTmp);
       const LocalPosition3D pos(x, y, zTmp);
-      unsigned char secNum = simOneSectorOnly ? 0 : std::floor(phiTmp / SECPHIWIDTH);
+      unsigned char secNum = std::floor(phiTmp / SECPHIWIDTH);
       Sector sector(secNum + (pos.Z() < 0) * SECTORSPERSIDE);
       LocalPosition3D lPosTmp = Mapper::GlobalToLocal(pos, sector);
 
@@ -428,10 +427,9 @@ void DataContainer3D<DataT>::dumpInterpolation(std::string_view treename, std::s
 
   // define grid for interpolation
   using GridProp = GridProperties<DataT>;
-  const RegularGrid3D<DataT> mGrid3D(GridProp::ZMIN, GridProp::RMIN, GridProp::PHIMIN, GridProp::getGridSpacingZ(nz), GridProp::getGridSpacingR(nr), o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) / (MGParameters::normalizeGridToOneSector ? SECTORSPERSIDE : 1), ParamSpaceCharge{nr, nz, nphi});
+  const RegularGrid3D<DataT> mGrid3D(GridProp::ZMIN, GridProp::RMIN, GridProp::PHIMIN, GridProp::getGridSpacingZ(nz), GridProp::getGridSpacingR(nr), o2::tpc::GridProperties<float>::getGridSpacingPhi(nphi) * (MGParameters::normalizeGridToNSector / double(SECTORSPERSIDE)), ParamSpaceCharge{nr, nz, nphi});
 
   auto interpolate = [&mGrid3D = std::as_const(mGrid3D), &data = std::as_const(data), rangeR, rangeZ, rangePhi, nR, nZ, nPhi](unsigned int, ULong64_t iPhi) {
-    const bool simOneSectorOnly = MGParameters::normalizeGridToOneSector;
     std::vector<size_t> ir;
     std::vector<size_t> iphi;
     std::vector<size_t> iz;
@@ -473,7 +471,7 @@ void DataContainer3D<DataT>::dumpInterpolation(std::string_view treename, std::s
         const float x = rPos * std::cos(phiPos);
         const float y = rPos * std::sin(phiPos);
         const LocalPosition3D pos(x, y, zPos);
-        unsigned char secNum = simOneSectorOnly ? 0 : std::floor(phiPos / SECPHIWIDTH);
+        unsigned char secNum = std::floor(phiPos / SECPHIWIDTH); // TODO CHECK THIS
         Sector sector(secNum + (pos.Z() < 0) * SECTORSPERSIDE);
         LocalPosition3D lPosTmp = Mapper::GlobalToLocal(pos, sector);
         lPos.emplace_back(lPosTmp);
@@ -512,6 +510,27 @@ bool DataContainer3D<DataT>::getVertices(std::string_view treename, std::string_
   return true;
 }
 
+template <typename DataT>
+DataContainer3D<DataT> DataContainer3D<DataT>::convert(const o2::tpc::RegularGrid3D<DataT>& gridNew, const o2::tpc::RegularGrid3D<DataT>& gridRef, const int threads) const
+{
+  const int nZNew = gridNew.getNZ();
+  const int nRNew = gridNew.getNR();
+  const int nPhiNew = gridNew.getNPhi();
+  DataContainer3D<DataT> contCont(nZNew, nRNew, nPhiNew);
+#pragma omp parallel for num_threads(threads)
+  for (size_t iPhi = 0; iPhi < nPhiNew; ++iPhi) {
+    const DataT phi = gridNew.getPhiVertex(iPhi);
+    for (size_t iR = 0; iR < nRNew; ++iR) {
+      const DataT radius = gridNew.getRVertex(iR);
+      for (size_t iZ = 0; iZ < nZNew; ++iZ) {
+        const DataT z = gridNew.getZVertex(iZ);
+        contCont(iZ, iR, iPhi) = interpolate(z, radius, phi, gridRef);
+      }
+    }
+  }
+  return contCont;
+}
+
 template class o2::tpc::DataContainer3D<float>;
 template class o2::tpc::DataContainer3D<double>;
 

Detectors/TPC/spacecharge/src/PoissonSolver.cxx

@@ -940,7 +940,7 @@ void PoissonSolver<DataT>::residue2D(Vector& residue, const Vector& matricesCurr
     for (int j = 1; j < tnZColumn - 1; ++j) {
       residue(i, j, iPhi) = ih2 * (coefficient1[i] * matricesCurrentV(i + 1, j, iPhi) + coefficient2[i] * matricesCurrentV(i - 1, j, iPhi) + tempRatio * (matricesCurrentV(i, j + 1, iPhi) + matricesCurrentV(i, j - 1, iPhi)) - inverseTempFourth * matricesCurrentV(i, j, iPhi)) + matricesCurrentCharge(i, j, iPhi);
     } // end cols
-  }   // end nRRow
+  } // end nRRow
 
   // Boundary points.
   for (int i = 0; i < tnRRow; ++i) {
@@ -997,7 +997,7 @@ void PoissonSolver<DataT>::residue3D(Vector& residue, const Vector& matricesCurr
                                   coefficient3[i] * (signPlus * matricesCurrentV(i, j, mp1) + signMinus * matricesCurrentV(i, j, mm1)) - inverseCoefficient4[i] * matricesCurrentV(i, j, m)) +
                            matricesCurrentCharge(i, j, m);
       } // end cols
-    }   // end mParamGrid.NRVertices
+    } // end mParamGrid.NRVertices
   }
 }
 
@@ -1263,9 +1263,9 @@ void PoissonSolver<DataT>::relax3D(Vector& matricesCurrentV, const Vector& matri
           for (int i = isw; i < tnRRow - 1; i += 2) {
             (matricesCurrentV)(i, j, m) = (coefficient2[i] * (matricesCurrentV)(i - 1, j, m) + tempRatioZ * ((matricesCurrentV)(i, j - 1, m) + (matricesCurrentV)(i, j + 1, m)) + coefficient1[i] * (matricesCurrentV)(i + 1, j, m) + coefficient3[i] * (signPlus * (matricesCurrentV)(i, j, mp1) + signMinus * (matricesCurrentV)(i, j, mm1)) + (h2 * (matricesCurrentCharge)(i, j, m))) * coefficient4[i];
           } // end cols
-        }   // end mParamGrid.NRVertices
-      }     // end phi
-    }       // end sweep
+        } // end mParamGrid.NRVertices
+      } // end phi
+    } // end sweep
   } else if (MGParameters::relaxType == RelaxType::Jacobi) {
     // for each slice
     for (int m = 0; m < iPhi; ++m) {
@@ -1306,8 +1306,8 @@ void PoissonSolver<DataT>::relax3D(Vector& matricesCurrentV, const Vector& matri
         for (int i = 1; i < tnRRow - 1; ++i) {
           (matricesCurrentV)(i, j, m) = (coefficient2[i] * (matricesCurrentV)(i - 1, j, m) + tempRatioZ * ((matricesCurrentV)(i, j - 1, m) + (matricesCurrentV)(i, j + 1, m)) + coefficient1[i] * (matricesCurrentV)(i + 1, j, m) + coefficient3[i] * (signPlus * (matricesCurrentV)(i, j, mp1) + signMinus * (matricesCurrentV)(i, j, mm1)) + (h2 * (matricesCurrentCharge)(i, j, m))) * coefficient4[i];
         } // end cols
-      }   // end mParamGrid.NRVertices
-    }     // end phi
+      } // end mParamGrid.NRVertices
+    } // end phi
   } else {
     // Case weighted Jacobi
     // TODO
@@ -1329,15 +1329,15 @@ void PoissonSolver<DataT>::relax2D(Vector& matricesCurrentV, const Vector& matri
           matricesCurrentV(i, j, iPhi) = tempFourth * (coefficient1[i] * matricesCurrentV(i + 1, j, iPhi) + coefficient2[i] * matricesCurrentV(i - 1, j, iPhi) +
                                                        tempRatio * (matricesCurrentV(i, j + 1, iPhi) + matricesCurrentV(i, j - 1, iPhi)) + (h2 * matricesCurrentCharge(i, j, iPhi)));
         } // end cols
-      }   // end mParamGrid.NRVertices
-    }     // end pass red-black
+      } // end mParamGrid.NRVertices
+    } // end pass red-black
   } else if (MGParameters::relaxType == RelaxType::Jacobi) {
     for (int j = 1; j < tnZColumn - 1; ++j) {
       for (int i = 1; i < tnRRow - 1; ++i) {
         matricesCurrentV(i, j, iPhi) = tempFourth * (coefficient1[i] * matricesCurrentV(i + 1, j, iPhi) + coefficient2[i] * matricesCurrentV(i - 1, j, iPhi) +
                                                      tempRatio * (matricesCurrentV(i, j + 1, iPhi) + matricesCurrentV(i, j - 1, iPhi)) + (h2 * matricesCurrentCharge(i, j, iPhi)));
       } // end cols
-    }   // end mParamGrid.NRVertices
+    } // end mParamGrid.NRVertices
   } else if (MGParameters::relaxType == RelaxType::WeightedJacobi) {
     // Weighted Jacobi
     // TODO
@@ -1421,7 +1421,7 @@ void PoissonSolver<DataT>::restrict3D(Vector& matricesCurrentCharge, const Vecto
 
           matricesCurrentCharge(i, j, m) = residue(ii, jj, mm) / 8 + s1 / 16 + s2 / 32 + s3 / 64;
         } // end cols
-      }   // end mParamGrid.NRVertices
+      } // end mParamGrid.NRVertices
 
       // for boundary
       for (int j = 0, jj = 0; j < tnZColumn; ++j, jj += 2) {
@@ -1460,7 +1460,7 @@ void PoissonSolver<DataT>::restrict2D(Vector& matricesCurrentCharge, const Vecto
                                             (residue(iip1, jjp1, iphi) + residue(iim1, jjp1, iphi) + residue(iip1, jjm1, iphi) + residue(iim1, jjm1, iphi)) / 16;
       }
     } // end cols
-  }   // end mParamGrid.NRVertices
+  } // end mParamGrid.NRVertices
   // boundary
   // for boundary
   for (int j = 0, jj = 0; j < tnZColumn; ++j, jj += 2) {
@@ -1520,7 +1520,7 @@ void PoissonSolver<DataT>::calcCoefficients2D(unsigned int from, unsigned int to
 template <typename DataT>
 DataT PoissonSolver<DataT>::getGridSizePhiInv()
 {
-  return MGParameters::normalizeGridToOneSector ? (INVTWOPI * SECTORSPERSIDE) : INVTWOPI;
+  return INVTWOPI * SECTORSPERSIDE / MGParameters::normalizeGridToNSector;
 }
 
 template class o2::tpc::PoissonSolver<double>;