Use of material map in the track propagator (#3682)

Author: sgorbuno

GPU/GPUTracking/DataCompression/AliHLTTPCClusterStatComponent.cxx

@@ -315,10 +315,8 @@ int AliHLTTPCClusterStatComponent::DoEvent(const AliHLTComponentEventData& evtDa
   const AliHLTUInt8_t* pCurrent = reinterpret_cast<const AliHLTUInt8_t*>(tracks->fTracklets);
   if (fCompressionStudy) {
     GPUTPCGMPropagator prop;
-    const float kRho = 1.025e-3;  // 0.9e-3;
-    const float kRadLen = 29.532; // 28.94;
     prop.SetMaxSinPhi(.999);
-    prop.SetMaterial(kRadLen, kRho);
+    prop.SetMaterialTPC();
     GPUTPCGMPolynomialField field;
     int err = GPUTPCGMPolynomialFieldManager::GetPolynomialField(field);
     if (err != 0) {

GPU/GPUTracking/DataCompression/GPUTPCCompressionTrackModel.cxx

@@ -24,9 +24,7 @@ using namespace GPUCA_NAMESPACE::gpu;
 #ifdef GPUCA_COMPRESSION_TRACK_MODEL_MERGER
 GPUd() void GPUTPCCompressionTrackModel::Init(float x, float y, float z, float alpha, unsigned char qPt, const GPUParam& GPUrestrict() param)
 {
-  static constexpr float kRho = 1.025e-3f;  // 0.9e-3;
-  static constexpr float kRadLen = 29.532f; // 28.94;
-  mProp.SetMaterial(kRadLen, kRho);
+  mProp.SetMaterialTPC();
   mProp.SetMaxSinPhi(GPUCA_MAX_SIN_PHI);
   mProp.SetToyMCEventsFlag(false);
   mProp.SetSeedingErrors(true); // Larger errors for seeds, better since we don't start with good hypothesis

GPU/GPUTracking/Merger/GPUTPCGMMerger.cxx

@@ -342,10 +342,8 @@ GPUd() void GPUTPCGMMerger::ClearTrackLinks(int nBlocks, int nThreads, int iBloc
 
 GPUd() int GPUTPCGMMerger::RefitSliceTrack(GPUTPCGMSliceTrack& sliceTrack, const GPUTPCTrack* inTrack, float alpha, int slice)
 {
-  static constexpr float kRho = 1.025e-3f;  // 0.9e-3;
-  static constexpr float kRadLen = 29.532f; // 28.94;
   GPUTPCGMPropagator prop;
-  prop.SetMaterial(kRadLen, kRho);
+  prop.SetMaterialTPC();
   prop.SetMaxSinPhi(GPUCA_MAX_SIN_PHI);
   prop.SetToyMCEventsFlag(false);
   prop.SetSeedingErrors(true); // Larger errors for seeds, better since we don't start with good hypothesis
@@ -1747,7 +1745,8 @@ void GPUCA_KRNL_BACKEND_CLASS::runKernelBackendInternal<GPUTPCGMMergerSortTracks
 
 GPUd() void GPUTPCGMMerger::SortTracks(int nBlocks, int nThreads, int iBlock, int iThread)
 {
-  auto comp = [cmp = mOutputTracks](const int aa, const int bb) { // Have to duplicate sort comparison: Thrust cannot use the Lambda but OpenCL cannot use the object
+  // Have to duplicate sort comparison: Thrust cannot use the Lambda but OpenCL cannot use the object
+  auto comp = [cmp = mOutputTracks](const int aa, const int bb) {
     const GPUTPCGMMergedTrack& GPUrestrict() a = cmp[aa];
     const GPUTPCGMMergedTrack& GPUrestrict() b = cmp[bb];
     if (a.CCE() != b.CCE()) {
@@ -1765,7 +1764,8 @@ GPUd() void GPUTPCGMMerger::SortTracks(int nBlocks, int nThreads, int iBlock, in
 GPUd() void GPUTPCGMMerger::SortTracksQPt(int nBlocks, int nThreads, int iBlock, int iThread)
 {
   unsigned int* trackSort = (unsigned int*)mTmpMem;
-  auto comp = [cmp = mOutputTracks](const int aa, const int bb) { // Have to duplicate sort comparison: Thrust cannot use the Lambda but OpenCL cannot use the object
+  // Have to duplicate sort comparison: Thrust cannot use the Lambda but OpenCL cannot use the object
+  auto comp = [cmp = mOutputTracks](const int aa, const int bb) {
     const GPUTPCGMMergedTrack& GPUrestrict() a = cmp[aa];
     const GPUTPCGMMergedTrack& GPUrestrict() b = cmp[bb];
     return (CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt()));

GPU/GPUTracking/Merger/GPUTPCGMPropagator.cxx

@@ -29,16 +29,16 @@
 
 using namespace GPUCA_NAMESPACE::gpu;
 
-GPUd() void GPUTPCGMPropagator::GetBxByBz(float Alpha, float X, float Y, float Z, float B[3]) const
+GPUd() void GPUTPCGMPropagator::GetBxByBzBase(float cosAlpha, float sinAlpha, float X, float Y, float Z, float B[3]) const
 {
   // get global coordinates
 
-  float cs = CAMath::Cos(Alpha);
-  float sn = CAMath::Sin(Alpha);
+  float gx = getGlobalX(cosAlpha, sinAlpha, X, Y);
+  float gy = getGlobalY(cosAlpha, sinAlpha, X, Y);
 
 #if defined(GPUCA_GM_USE_FULL_FIELD)
   const double kCLight = 0.000299792458;
-  double r[3] = {X * cs - Y * sn, X * sn + Y * cs, Z};
+  double r[3] = {gx, gy, Z};
   double bb[3];
   AliTracker::GetBxByBz(r, bb);
   bb[0] *= kCLight;
@@ -57,22 +57,22 @@ GPUd() void GPUTPCGMPropagator::GetBxByBz(float Alpha, float X, float Y, float Z
   float bb[3];
   switch (mFieldRegion) {
     case ITS:
-      mField->GetFieldIts(X * cs - Y * sn, X * sn + Y * cs, Z, bb);
+      mField->GetFieldIts(gx, gy, Z, bb);
       break;
     case TRD:
-      mField->GetFieldTrd(X * cs - Y * sn, X * sn + Y * cs, Z, bb);
+      mField->GetFieldTrd(gx, gy, Z, bb);
       break;
     case TPC:
     default:
-      mField->GetField(X * cs - Y * sn, X * sn + Y * cs, Z, bb);
+      mField->GetField(gx, gy, Z, bb);
   }
 
 #endif
 
   // rotate field to local coordinates
 
-  B[0] = bb[0] * cs + bb[1] * sn;
-  B[1] = -bb[0] * sn + bb[1] * cs;
+  B[0] = bb[0] * cosAlpha + bb[1] * sinAlpha;
+  B[1] = -bb[0] * sinAlpha + bb[1] * cosAlpha;
   B[2] = bb[2];
   /*if( mToyMCEvents ){ // special treatment for toy monte carlo
     B[0] = 0;
@@ -81,36 +81,35 @@ GPUd() void GPUTPCGMPropagator::GetBxByBz(float Alpha, float X, float Y, float Z
   }*/
 }
 
-GPUd() float GPUTPCGMPropagator::GetBz(float Alpha, float X, float Y, float Z) const
+GPUd() float GPUTPCGMPropagator::GetBzBase(float cosAlpha, float sinAlpha, float X, float Y, float Z) const
 {
   if (mToyMCEvents) { // special treatment for toy monte carlo
     float B[3];
-    GetBxByBz(Alpha, X, Y, Z, B);
+    GetBxByBzBase(cosAlpha, sinAlpha, X, Y, Z, B);
     return B[2];
   }
 
   // get global coordinates
 
-  float cs = CAMath::Cos(Alpha);
-  float sn = CAMath::Sin(Alpha);
+  float gx = getGlobalX(cosAlpha, sinAlpha, X, Y);
+  float gy = getGlobalY(cosAlpha, sinAlpha, X, Y);
 
 #if defined(GPUCA_GM_USE_FULL_FIELD)
   const double kCLight = 0.000299792458;
-  double r[3] = {X * cs - Y * sn, X * sn + Y * cs, Z};
+  double r[3] = {gx, gy, Z};
   double bb[3];
   AliTracker::GetBxByBz(r, bb);
   return bb[2] * kCLight;
 #else
   switch (mFieldRegion) {
     case ITS:
-      return mField->GetFieldItsBz(X * cs - Y * sn, X * sn + Y * cs, Z);
+      return mField->GetFieldItsBz(gx, gy, Z);
     case TRD:
-      return mField->GetFieldTrdBz(X * cs - Y * sn, X * sn + Y * cs, Z);
+      return mField->GetFieldTrdBz(gx, gy, Z);
     case TPC:
     default:
-      return mField->GetFieldBz(X * cs - Y * sn, X * sn + Y * cs, Z);
+      return mField->GetFieldBz(gx, gy, Z);
   }
-
 #endif
 }
 
@@ -121,8 +120,11 @@ GPUd() int GPUTPCGMPropagator::RotateToAlpha(float newAlpha)
   // return value is error code (0==no error)
   //
 
-  float cc = CAMath::Cos(newAlpha - mAlpha);
-  float ss = CAMath::Sin(newAlpha - mAlpha);
+  float newCosAlpha = CAMath::Cos(newAlpha);
+  float newSinAlpha = CAMath::Sin(newAlpha);
+
+  float cc = newCosAlpha * mCosAlpha + newSinAlpha * mSinAlpha; // cos(newAlpha - mAlpha);
+  float ss = newSinAlpha * mCosAlpha - newCosAlpha * mSinAlpha; //sin(newAlpha - mAlpha);
 
   GPUTPCGMPhysicalTrackModel t0 = mT0;
 
@@ -271,6 +273,8 @@ GPUd() int GPUTPCGMPropagator::RotateToAlpha(float newAlpha)
   }
 
   mAlpha = newAlpha;
+  mCosAlpha = newCosAlpha;
+  mSinAlpha = newSinAlpha;
   mT0 = t0;
 
   return 0;
@@ -288,7 +292,7 @@ GPUd() int GPUTPCGMPropagator::PropagateToXAlpha(float posX, float posAlpha, boo
   }
 
   float B[3];
-  GetBxByBz(mAlpha, mT0.X(), mT0.Y(), mT0.Z(), B);
+  GetBxByBz(mT0.X(), mT0.Y(), mT0.Z(), B);
 
   // propagate mT0 to t0e
 
@@ -313,7 +317,7 @@ GPUd() int GPUTPCGMPropagator::PropagateToXAlphaBz(float posX, float posAlpha, b
     }
   }
 
-  float Bz = GetBz(mAlpha, mT0.X(), mT0.Y(), mT0.Z());
+  float Bz = GetBz(mT0.X(), mT0.Y(), mT0.Z());
 
   // propagate mT0 to t0e
 
@@ -375,6 +379,17 @@ GPUd() int GPUTPCGMPropagator::FollowLinearization(const GPUTPCGMPhysicalTrackMo
     return -4;
   }
 
+  if (mMatLUT) {
+    float xyz1[3] = {getGlobalX(mT0.GetX(), mT0.GetY()), getGlobalY(mT0.GetX(), mT0.GetY()), mT0.GetZ()};
+    float xyz2[3] = {getGlobalX(t0e.GetX(), t0e.GetY()), getGlobalY(t0e.GetX(), t0e.GetY()), t0e.GetZ()};
+    o2::base::MatBudget mat = getMatBudget(xyz1, xyz2);
+    if (mat.meanX2X0 > 1.e-8) {
+      SetMaterial(mat.length / mat.meanX2X0, mat.meanRho);
+    } else {
+      SetMaterialTPC();
+    }
+  }
+
   mT0 = t0e;
   mT->X() = t0e.GetX();
   p[0] = t0e.GetY() + d0 + j02 * d2 + j04 * d4;
@@ -538,7 +553,7 @@ GPUd() int GPUTPCGMPropagator::FollowLinearization(const GPUTPCGMPhysicalTrackMo
 
 GPUd() int GPUTPCGMPropagator::GetPropagatedYZ(float x, float& GPUrestrict() projY, float& GPUrestrict() projZ)
 {
-  float bz = GetBz(mAlpha, mT->X(), mT->Y(), mT->Z());
+  float bz = GetBz(mT->X(), mT->Y(), mT->Z());
   float k = mT0.QPt() * bz;
   float dx = x - mT->X();
   float kdx = k * dx;
@@ -882,7 +897,7 @@ GPUd() void GPUTPCGMPropagator::CalculateMaterialCorrection()
   float p2 = w2 / pti2; // impuls 2
   float betheRho = ApproximateBetheBloch(p2 / mass2) * mMaterial.rho;
   float E = CAMath::Sqrt(p2 + mass2);
-  float theta2 = (14.1f * 14.1f / 1.e6f) / (beta2 * p2) * mMaterial.rhoOverRadLen;
+  float theta2 = (14.1f * 14.1f / 1.e6f) / (beta2 * p2) * mMaterial.radLenInv;
 
   mMaterial.EP2 = E / p2;
 
@@ -923,6 +938,8 @@ GPUd() void GPUTPCGMPropagator::Rotate180()
   while (mAlpha < -M_PI) {
     mAlpha += 2 * M_PI;
   }
+  mCosAlpha = -mCosAlpha;
+  mSinAlpha = -mSinAlpha;
 
   float* c = mT->Cov();
   c[6] = -c[6];
@@ -956,7 +973,7 @@ GPUd() void GPUTPCGMPropagator::Mirror(bool inFlyDirection)
 {
   // mirror the track and the track approximation to the point which has the same X, but located on the other side of trajectory
   float B[3];
-  GetBxByBz(mAlpha, mT0.X(), mT0.Y(), mT0.Z(), B);
+  GetBxByBz(mT0.X(), mT0.Y(), mT0.Z(), B);
   float Bz = B[2];
   if (CAMath::Abs(Bz) < 1.e-8f) {
     Bz = 1.e-8f;
@@ -1036,7 +1053,7 @@ GPUd() void GPUTPCGMPropagator::Mirror(bool inFlyDirection)
   }
 }
 
-GPUd() o2::base::MatBudget GPUTPCGMPropagator::getMatBudget(float* p1, float* p2)
+GPUd() o2::base::MatBudget GPUTPCGMPropagator::getMatBudget(const float* p1, const float* p2)
 {
 #ifdef HAVE_O2HEADERS
   return mMatLUT->getMatBudget(p1[0], p1[1], p1[2], p2[0], p2[1], p2[2]);

GPU/GPUTracking/Merger/GPUTPCGMPropagator.h

@@ -55,13 +55,17 @@ class GPUTPCGMPropagator
   GPUdDefault() GPUTPCGMPropagator() CON_DEFAULT;
 
   struct MaterialCorrection {
-    GPUhd() MaterialCorrection() : radLen(29.532f), rho(1.025e-3f), rhoOverRadLen(rho / radLen), DLMax(0.f), EP2(0.f), sigmadE2(0.f), k22(0.f), k33(0.f), k43(0.f), k44(0.f) {}
+    GPUhd() MaterialCorrection() : radLen(28811.7f), rho(1.025e-3f), radLenInv(1.f / radLen), DLMax(0.f), EP2(0.f), sigmadE2(0.f), k22(0.f), k33(0.f), k43(0.f), k44(0.f) {}
 
-    float radLen, rho, rhoOverRadLen, DLMax, EP2, sigmadE2, k22, k33, k43, k44; // precalculated values for MS and EnergyLoss correction
+    float radLen;                                              // [cm]
+    float rho;                                                 // [g/cm^3]
+    float radLenInv, DLMax, EP2, sigmadE2, k22, k33, k43, k44; // precalculated values for MS and EnergyLoss correction
   };
 
   GPUd() void SetMaterial(float radLen, float rho);
-  GPUd() o2::base::MatBudget getMatBudget(float* p1, float* p2);
+  GPUd() void SetMaterialTPC() { SetMaterial(28811.7f, 1.025e-3f); }
+
+  GPUd() o2::base::MatBudget getMatBudget(const float* p1, const float* p2);
 
   GPUd() void SetPolynomialField(const GPUTPCGMPolynomialField* field) { mField = field; }
 
@@ -109,7 +113,14 @@ class GPUTPCGMPropagator
     return 0;
   }
 
+  /// Bz in local coordinates rotated to mAlpha
+  GPUd() float GetBz(float X, float Y, float Z) const;
+  /// Bx,By,Bz in local coordinates rotated to mAlpha
+  GPUd() void GetBxByBz(float X, float Y, float Z, float B[3]) const;
+
+  /// Bz in local coordinates rotated to Alpha
   GPUd() float GetBz(float Alpha, float X, float Y, float Z) const;
+  /// Bx,By,Bz in local coordinates rotated to Alpha
   GPUd() void GetBxByBz(float Alpha, float X, float Y, float Z, float B[3]) const;
 
   GPUd() void GetErr2(float& err2Y, float& err2Z, const GPUParam& param, float posZ, int iRow, short clusterState) const;
@@ -137,11 +148,26 @@ class GPUTPCGMPropagator
   GPUd() static float ApproximateBetheBloch(float beta2);
   GPUd() int FollowLinearization(const GPUTPCGMPhysicalTrackModel& t0e, float Bz, float dLp, bool inFlyDirection);
 
+  /// Bz in local coordinates rotated to cosAlpha, sinAlpha
+  GPUd() float GetBzBase(float cosAlpha, float sinAlpha, float X, float Y, float Z) const;
+  /// Bx,By,Bz in local coordinates rotated to cosAlpha, sinAlpha
+  GPUd() void GetBxByBzBase(float cosAlpha, float sinAlpha, float X, float Y, float Z, float B[3]) const;
+  // X in global coordinates
+  GPUd() float getGlobalX(float cosAlpha, float sinAlpha, float X, float Y) const;
+  // Y in global coordinates
+  GPUd() float getGlobalY(float cosAlpha, float sinAlpha, float X, float Y) const;
+  // X in global coordinates
+  GPUd() float getGlobalX(float X, float Y) const;
+  // Y in global coordinates
+  GPUd() float getGlobalY(float X, float Y) const;
+
   const GPUTPCGMPolynomialField* mField = nullptr;
   FieldRegion mFieldRegion = TPC;
 
   GPUTPCGMTrackParam* mT = nullptr;
-  float mAlpha = 0; // rotation angle of the track coordinate system
+  float mAlpha = 0.f;    // rotation angle of the track coordinate system
+  float mCosAlpha = 1.f; // cos of the rotation angle
+  float mSinAlpha = 0.f; // sin of the rotation angle
   GPUTPCGMPhysicalTrackModel mT0;
   MaterialCorrection mMaterial;
   bool mSeedingErrors = 0;
@@ -154,11 +180,31 @@ class GPUTPCGMPropagator
   const o2::base::MatLayerCylSet* mMatLUT = nullptr;
 };
 
+GPUdi() void GPUTPCGMPropagator::GetBxByBz(float Alpha, float X, float Y, float Z, float B[3]) const
+{
+  GetBxByBzBase(CAMath::Cos(Alpha), CAMath::Sin(Alpha), X, Y, Z, B);
+}
+
+GPUdi() float GPUTPCGMPropagator::GetBz(float Alpha, float X, float Y, float Z) const
+{
+  return GetBzBase(CAMath::Cos(Alpha), CAMath::Sin(Alpha), X, Y, Z);
+}
+
+GPUdi() void GPUTPCGMPropagator::GetBxByBz(float X, float Y, float Z, float B[3]) const
+{
+  GetBxByBzBase(mCosAlpha, mSinAlpha, X, Y, Z, B);
+}
+
+GPUdi() float GPUTPCGMPropagator::GetBz(float X, float Y, float Z) const
+{
+  return GetBzBase(mCosAlpha, mSinAlpha, X, Y, Z);
+}
+
 GPUdi() void GPUTPCGMPropagator::SetMaterial(float radLen, float rho)
 {
   mMaterial.rho = rho;
   mMaterial.radLen = radLen;
-  mMaterial.rhoOverRadLen = (radLen > 1.e-4f) ? rho / radLen : 0.f;
+  mMaterial.radLenInv = (radLen > 1.e-4f) ? 1.f / radLen : 0.f;
   CalculateMaterialCorrection();
 }
 
@@ -170,12 +216,14 @@ GPUdi() void GPUTPCGMPropagator::SetTrack(GPUTPCGMTrackParam* GPUrestrict() trac
   }
   mT0.Set(*mT);
   mAlpha = Alpha;
+  mCosAlpha = CAMath::Cos(mAlpha);
+  mSinAlpha = CAMath::Sin(mAlpha);
   CalculateMaterialCorrection();
 }
 
 GPUdi() float GPUTPCGMPropagator::GetMirroredYModel() const
 {
-  float Bz = GetBz(mAlpha, mT0.GetX(), mT0.GetY(), mT0.GetZ());
+  float Bz = GetBz(mT0.GetX(), mT0.GetY(), mT0.GetZ());
   return mT0.GetMirroredY(Bz);
 }
 
@@ -184,9 +232,30 @@ GPUdi() float GPUTPCGMPropagator::GetMirroredYTrack() const
   if (!mT) {
     return -1.E10f;
   }
-  float Bz = GetBz(mAlpha, mT->GetX(), mT->GetY(), mT->GetZ());
+  float Bz = GetBz(mT->GetX(), mT->GetY(), mT->GetZ());
   return mT->GetMirroredY(Bz);
 }
+
+GPUdi() float GPUTPCGMPropagator::getGlobalX(float cosAlpha, float sinAlpha, float X, float Y) const
+{
+  return X * cosAlpha - Y * sinAlpha;
+}
+
+GPUdi() float GPUTPCGMPropagator::getGlobalY(float cosAlpha, float sinAlpha, float X, float Y) const
+{
+  return X * sinAlpha + Y * cosAlpha;
+}
+
+GPUdi() float GPUTPCGMPropagator::getGlobalX(float X, float Y) const
+{
+  return getGlobalX(mCosAlpha, mSinAlpha, X, Y);
+}
+
+GPUdi() float GPUTPCGMPropagator::getGlobalY(float X, float Y) const
+{
+  return getGlobalY(mCosAlpha, mSinAlpha, X, Y);
+}
+
 } // namespace gpu
 } // namespace GPUCA_NAMESPACE