TPCFastSpaceChargeCorrection.h

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.

/// \file  TPCFastSpaceChargeCorrection.h
/// \brief Definition of TPCFastSpaceChargeCorrection class
///
/// \author  Sergey Gorbunov <sergey.gorbunov@cern.ch>

#ifndef ALICEO2_GPUCOMMON_TPCFASTTRANSFORMATION_TPCFASTSPACECHARGECORRECTION_H
#define ALICEO2_GPUCOMMON_TPCFASTTRANSFORMATION_TPCFASTSPACECHARGECORRECTION_H

#include "Spline2D.h"
#include "TPCFastTransformGeo.h"
#include "FlatObject.h"
#include "GPUCommonDef.h"
#include "GPUCommonMath.h"

namespace o2
{
namespace gpu
{

///
/// The TPCFastSpaceChargeCorrection class represents correction of nominal coordinates of TPC clusters
/// using best-fit splines
///
/// Row, U, V -> dX,dU,dV
///
/// The class is flat C structure. No virtual methods, no ROOT types are used.
///
class TPCFastSpaceChargeCorrection : public FlatObject
{
  friend class TPCFastTransformPOD;

 public:
  ///
  /// \brief The struct contains necessary info for TPC padrow
  ///
  struct RowInfo {
    int32_t splineScenarioID{0};  ///< scenario index (which of Spline2D splines to use)
    size_t dataOffsetBytes[3]{0}; ///< offset for the spline data withing a TPC slice
    ClassDefNV(RowInfo, 1);
  };

  struct RowActiveArea {
    float maxDriftLengthCheb[5]{0.f};
    float vMax{0.f};
    float cuMin{0.f};
    float cuMax{0.f};
    float cvMax{0.f};
    ClassDefNV(RowActiveArea, 1);
  };

  struct SliceRowInfo {
    float gridV0{0.f};           ///< V coordinate of the V-grid start
    float gridCorrU0{0.f};       ///< U coordinate of the U-grid start for corrected U
    float gridCorrV0{0.f};       ///< V coordinate of the V-grid start for corrected V
    float scaleCorrUtoGrid{0.f}; ///< scale corrected U to U-grid coordinate
    float scaleCorrVtoGrid{0.f}; ///< scale corrected V to V-grid coordinate
    RowActiveArea activeArea;
    ClassDefNV(SliceRowInfo, 1);
  };

  struct SliceInfo {
    float vMax{0.f}; ///< Max value of V coordinate
    ClassDefNV(SliceInfo, 1);
  };

  typedef Spline2D<float, 3> SplineType;

  /// _____________  Constructors / destructors __________________________

  /// Default constructor: creates an empty uninitialized object
  TPCFastSpaceChargeCorrection();

  /// Copy constructor: disabled to avoid ambiguity. Use cloneTo[In/Ex]ternalBuffer() instead
  TPCFastSpaceChargeCorrection(const TPCFastSpaceChargeCorrection&) = delete;

  /// Assignment operator: disabled to avoid ambiguity. Use cloneTo[In/Ex]ternalBuffer() instead
  TPCFastSpaceChargeCorrection& operator=(const TPCFastSpaceChargeCorrection&) = delete;

  /// Destructor
  ~TPCFastSpaceChargeCorrection();

  /// _____________  FlatObject functionality, see FlatObject class for description  ____________

  /// Memory alignment

  using FlatObject::getBufferAlignmentBytes;
  using FlatObject::getClassAlignmentBytes;

  /// Construction interface

  void cloneFromObject(const TPCFastSpaceChargeCorrection& obj, char* newFlatBufferPtr);
  void destroy();

  /// Making the data buffer external

  using FlatObject::releaseInternalBuffer;
  void moveBufferTo(char* newBufferPtr);

  /// Moving the class with its external buffer to another location

  void setActualBufferAddress(char* actualFlatBufferPtr);
  void setFutureBufferAddress(char* futureFlatBufferPtr);

  /// _______________  Construction interface  ________________________

  /// Starts the construction procedure, reserves temporary memory
  void startConstruction(const TPCFastTransformGeo& geo, int32_t numberOfSplineScenarios);

  /// Initializes a TPC row
  void setRowScenarioID(int32_t iRow, int32_t iScenario);

  /// Sets approximation scenario
  void setSplineScenario(int32_t scenarioIndex, const SplineType& spline);

  /// Finishes construction: puts everything to the flat buffer, releases temporary memory
  void finishConstruction();

  void constructWithNoCorrection(const TPCFastTransformGeo& geo);

  /// _______________  Initialization interface  ________________________

  /// Set no correction
  GPUd() void setNoCorrection();

  /// Sets the time stamp of the current calibaration
  GPUd() void setTimeStamp(int64_t v) { mTimeStamp = v; }

  /// Set safety marging for the interpolation around the TPC row.
  /// Outside of this area the interpolation returns the boundary values.
  GPUd() void setInterpolationSafetyMargin(float val) { fInterpolationSafetyMargin = val; }

  /// Gives const pointer to a spline
  GPUd() const SplineType& getSpline(int32_t slice, int32_t row) const;

  /// Gives pointer to a spline
  GPUd() SplineType& getSpline(int32_t slice, int32_t row);

  /// Gives pointer to spline data
  GPUd() float* getSplineData(int32_t slice, int32_t row, int32_t iSpline = 0);

  /// Gives pointer to spline data
  GPUd() const float* getSplineData(int32_t slice, int32_t row, int32_t iSpline = 0) const;

  /// _______________ The main method: cluster correction  _______________________
  ///
  GPUd() int32_t getCorrection(int32_t slice, int32_t row, float u, float v, float& dx, float& du, float& dv) const;

  /// inverse correction: Corrected U and V -> coorrected X
  GPUd() void getCorrectionInvCorrectedX(int32_t slice, int32_t row, float corrU, float corrV, float& corrX) const;

  /// inverse correction: Corrected U and V -> uncorrected U and V
  GPUd() void getCorrectionInvUV(int32_t slice, int32_t row, float corrU, float corrV, float& nomU, float& nomV) const;

  /// maximal possible drift length of the active area
  GPUd() float getMaxDriftLength(int32_t slice, int32_t row, float pad) const;

  /// maximal possible drift length of the active area
  GPUd() float getMaxDriftLength(int32_t slice, int32_t row) const;

  /// maximal possible drift length of the active area
  GPUd() float getMaxDriftLength(int32_t slice) const;

  /// _______________  Utilities  _______________________________________________

  /// shrink u,v coordinats to the TPC row area +/- fkInterpolationSafetyMargin
  GPUd() void schrinkUV(int32_t slice, int32_t row, float& u, float& v) const;

  /// shrink corrected u,v coordinats to the TPC row area +/- fkInterpolationSafetyMargin
  GPUd() void schrinkCorrectedUV(int32_t slice, int32_t row, float& corrU, float& corrV) const;

  /// convert u,v to internal grid coordinates
  GPUd() void convUVtoGrid(int32_t slice, int32_t row, float u, float v, float& gridU, float& gridV) const;

  /// convert u,v to internal grid coordinates
  GPUd() void convGridToUV(int32_t slice, int32_t row, float gridU, float gridV, float& u, float& v) const;

  /// convert corrected u,v to internal grid coordinates
  GPUd() void convCorrectedUVtoGrid(int32_t slice, int32_t row, float cu, float cv, float& gridU, float& gridV) const;

  /// TPC geometry information
  GPUd() const TPCFastTransformGeo& getGeometry() const
  {
    return mGeo;
  }

  /// Gives the time stamp of the current calibaration parameters
  int64_t getTimeStamp() const { return mTimeStamp; }

  /// Gives the interpolation safety marging  around the TPC row.
  GPUd() float getInterpolationSafetyMargin() const { return fInterpolationSafetyMargin; }

  /// Gives TPC row info
  GPUd() const RowInfo& getRowInfo(int32_t row) const { return mRowInfoPtr[row]; }

  /// Gives TPC slice info
  GPUd() const SliceInfo& getSliceInfo(int32_t slice) const
  {
    return mSliceInfo[slice];
  }

  /// Gives TPC slice info
  GPUd() SliceInfo& getSliceInfo(int32_t slice)
  {
    return mSliceInfo[slice];
  }

  /// Gives TPC slice & row info
  GPUd() const SliceRowInfo& getSliceRowInfo(int32_t slice, int32_t row) const
  {
    return mSliceRowInfoPtr[mGeo.getNumberOfRows() * slice + row];
  }

  /// Gives TPC slice & row info
  GPUd() SliceRowInfo& getSliceRowInfo(int32_t slice, int32_t row)
  {
    return mSliceRowInfoPtr[mGeo.getNumberOfRows() * slice + row];
  }

#if !defined(GPUCA_GPUCODE)
  /// Print method
  void print() const;
  GPUh() double testInverse(bool prn = 0);
#endif

 private:
  /// relocate buffer pointers
  void relocateBufferPointers(const char* oldBuffer, char* newBuffer);
  /// release temporary memory used during construction
  void releaseConstructionMemory();

  /// temporary method with the an way of calculating 2D spline
  GPUd() int32_t getCorrectionOld(int32_t slice, int32_t row, float u, float v, float& dx, float& du, float& dv) const;

  /// _______________  Data members  _______________________________________________

  /// _______________  Construction control  _______________________________________________

  RowInfo* mConstructionRowInfos = nullptr;     //! (transient!!) Temporary container of the row infos during construction
  SplineType* mConstructionScenarios = nullptr; //! (transient!!) Temporary container for spline scenarios

  /// _______________  Geometry  _______________________________________________

  TPCFastTransformGeo mGeo; ///< TPC geometry information

  int32_t mNumberOfScenarios; ///< Number of approximation spline scenarios

  SliceInfo mSliceInfo[TPCFastTransformGeo::getNumberOfSlices()]; ///< SliceInfo array

  SplineType* mScenarioPtr;       //! (transient!!) pointer to spline scenarios
  RowInfo* mRowInfoPtr;           //! (transient!!) pointer to RowInfo array inside the mFlatBufferPtr buffer
  SliceRowInfo* mSliceRowInfoPtr; //! (transient!!) pointer to SliceRowInfo array inside the mFlatBufferPtr

  /// _______________  Calibration data  _______________________________________________

  int64_t mTimeStamp; ///< time stamp of the current calibration

  char* mSplineData[3]; //! (transient!!) pointer to the spline data in the flat buffer

  size_t mSliceDataSizeBytes[3]; ///< size of the data for one slice in the flat buffer

  float fInterpolationSafetyMargin{0.1f}; // 10% area around the TPC row. Outside of this area the interpolation returns the boundary values.

  ClassDefNV(TPCFastSpaceChargeCorrection, 3);
};

/// ====================================================
///       Inline implementations of some methods
/// ====================================================

GPUdi() const TPCFastSpaceChargeCorrection::SplineType& TPCFastSpaceChargeCorrection::getSpline(int32_t slice, int32_t row) const
{
  /// Gives const pointer to spline
  const RowInfo& rowInfo = mRowInfoPtr[row];
  return mScenarioPtr[rowInfo.splineScenarioID];
}

GPUdi() TPCFastSpaceChargeCorrection::SplineType& TPCFastSpaceChargeCorrection::getSpline(int32_t slice, int32_t row)
{
  /// Gives pointer to spline
  const RowInfo& rowInfo = mRowInfoPtr[row];
  return mScenarioPtr[rowInfo.splineScenarioID];
}

GPUdi() float* TPCFastSpaceChargeCorrection::getSplineData(int32_t slice, int32_t row, int32_t iSpline)
{
  /// Gives pointer to spline data
  const RowInfo& rowInfo = mRowInfoPtr[row];
  return reinterpret_cast<float*>(mSplineData[iSpline] + mSliceDataSizeBytes[iSpline] * slice + rowInfo.dataOffsetBytes[iSpline]);
}

GPUdi() const float* TPCFastSpaceChargeCorrection::getSplineData(int32_t slice, int32_t row, int32_t iSpline) const
{
  /// Gives pointer to spline data
  const RowInfo& rowInfo = mRowInfoPtr[row];
  return reinterpret_cast<float*>(mSplineData[iSpline] + mSliceDataSizeBytes[iSpline] * slice + rowInfo.dataOffsetBytes[iSpline]);
}

GPUdi() void TPCFastSpaceChargeCorrection::schrinkUV(int32_t slice, int32_t row, float& u, float& v) const
{
  /// shrink u,v coordinats to the TPC row area +/- fInterpolationSafetyMargin

  const TPCFastTransformGeo::RowInfo& rowInfo = mGeo.getRowInfo(row);

  float uWidth05 = rowInfo.getUwidth() * (0.5f + fInterpolationSafetyMargin);
  float vWidth = mGeo.getTPCzLength(slice);

  if (u < -uWidth05) {
    u = -uWidth05;
  }
  if (u > uWidth05) {
    u = uWidth05;
  }
  if (v < -0.1f * vWidth) {
    v = -0.1f * vWidth;
  }
  if (v > 1.1f * vWidth) {
    v = 1.1f * vWidth;
  }
}

GPUdi() void TPCFastSpaceChargeCorrection::schrinkCorrectedUV(int32_t slice, int32_t row, float& corrU, float& corrV) const
{
  /// shrink corrected u,v coordinats to the TPC row area +/- fInterpolationSafetyMargin

  const TPCFastTransformGeo::RowInfo& rowInfo = mGeo.getRowInfo(row);
  const SliceRowInfo& sliceRowInfo = getSliceRowInfo(slice, row);

  float uMargin = fInterpolationSafetyMargin * rowInfo.getUwidth();
  float vMargin = fInterpolationSafetyMargin * mGeo.getTPCzLength(slice);

  if (corrU < sliceRowInfo.activeArea.cuMin - uMargin) {
    corrU = sliceRowInfo.activeArea.cuMin - uMargin;
  }

  if (corrU > sliceRowInfo.activeArea.cuMax + uMargin) {
    corrU = sliceRowInfo.activeArea.cuMax + uMargin;
  }

  if (corrV < 0.f - vMargin) {
    corrV = 0.f - vMargin;
  }

  if (corrV > sliceRowInfo.activeArea.cvMax + vMargin) {
    corrV = sliceRowInfo.activeArea.cvMax + vMargin;
  }
}

GPUdi() void TPCFastSpaceChargeCorrection::convUVtoGrid(int32_t slice, int32_t row, float u, float v, float& gu, float& gv) const
{
  // TODO optimise !!!
  gu = 0.f;
  gv = 0.f;

  schrinkUV(slice, row, u, v);

  const SliceRowInfo& info = getSliceRowInfo(slice, row);
  const SplineType& spline = getSpline(slice, row);

  float su0 = 0.f, sv0 = 0.f;
  mGeo.convUVtoScaledUV(slice, row, u, info.gridV0, su0, sv0);
  mGeo.convUVtoScaledUV(slice, row, u, v, gu, gv);

  gv = (gv - sv0) / (1.f - sv0);
  gu *= spline.getGridX1().getUmax();
  gv *= spline.getGridX2().getUmax();
}

GPUdi() void TPCFastSpaceChargeCorrection::convGridToUV(int32_t slice, int32_t row, float gridU, float gridV, float& u, float& v) const
{
  // TODO optimise
  /// convert u,v to internal grid coordinates
  float su0 = 0.f, sv0 = 0.f;
  const SliceRowInfo& info = getSliceRowInfo(slice, row);
  const SplineType& spline = getSpline(slice, row);
  mGeo.convUVtoScaledUV(slice, row, 0.f, info.gridV0, su0, sv0);
  float su = gridU / spline.getGridX1().getUmax();
  float sv = sv0 + gridV / spline.getGridX2().getUmax() * (1.f - sv0);
  mGeo.convScaledUVtoUV(slice, row, su, sv, u, v);
}

GPUdi() void TPCFastSpaceChargeCorrection::convCorrectedUVtoGrid(int32_t slice, int32_t row, float corrU, float corrV, float& gridU, float& gridV) const
{
  schrinkCorrectedUV(slice, row, corrU, corrV);

  const SliceRowInfo& sliceRowInfo = getSliceRowInfo(slice, row);

  gridU = (corrU - sliceRowInfo.gridCorrU0) * sliceRowInfo.scaleCorrUtoGrid;
  gridV = (corrV - sliceRowInfo.gridCorrV0) * sliceRowInfo.scaleCorrVtoGrid;
}

GPUdi() int32_t TPCFastSpaceChargeCorrection::getCorrection(int32_t slice, int32_t row, float u, float v, float& dx, float& du, float& dv) const
{
  const SplineType& spline = getSpline(slice, row);
  const float* splineData = getSplineData(slice, row);
  float gridU = 0, gridV = 0;
  convUVtoGrid(slice, row, u, v, gridU, gridV);
  float dxuv[3];
  spline.interpolateU(splineData, gridU, gridV, dxuv);
  if (CAMath::Abs(dxuv[0]) > 100 || CAMath::Abs(dxuv[1]) > 100 || CAMath::Abs(dxuv[2]) > 100) {
    dxuv[0] = dxuv[1] = dxuv[2] = 0;
  }
  dx = dxuv[0];
  du = dxuv[1];
  dv = dxuv[2];
  return 0;
}

GPUdi() int32_t TPCFastSpaceChargeCorrection::getCorrectionOld(int32_t slice, int32_t row, float u, float v, float& dx, float& du, float& dv) const
{
  const SplineType& spline = getSpline(slice, row);
  const float* splineData = getSplineData(slice, row);
  float gridU = 0, gridV = 0;
  convUVtoGrid(slice, row, u, v, gridU, gridV);
  float dxuv[3];
  spline.interpolateUold(splineData, gridU, gridV, dxuv);
  if (CAMath::Abs(dxuv[0]) > 100 || CAMath::Abs(dxuv[1]) > 100 || CAMath::Abs(dxuv[2]) > 100) {
    dxuv[0] = dxuv[1] = dxuv[2] = 0;
  }
  dx = dxuv[0];
  du = dxuv[1];
  dv = dxuv[2];
  return 0;
}

GPUdi() void TPCFastSpaceChargeCorrection::getCorrectionInvCorrectedX(
  int32_t slice, int32_t row, float corrU, float corrV, float& x) const
{
  float gridU, gridV;
  convCorrectedUVtoGrid(slice, row, corrU, corrV, gridU, gridV);

  const Spline2D<float, 1>& spline = reinterpret_cast<const Spline2D<float, 1>&>(getSpline(slice, row));
  const float* splineData = getSplineData(slice, row, 1);
  float dx = 0;
  spline.interpolateU(splineData, gridU, gridV, &dx);
  if (CAMath::Abs(dx) > 100) {
    dx = 0;
  }
  x = mGeo.getRowInfo(row).x + dx;
}

GPUdi() void TPCFastSpaceChargeCorrection::getCorrectionInvUV(
  int32_t slice, int32_t row, float corrU, float corrV, float& nomU, float& nomV) const
{
  float gridU, gridV;
  convCorrectedUVtoGrid(slice, row, corrU, corrV, gridU, gridV);

  const Spline2D<float, 2>& spline = reinterpret_cast<const Spline2D<float, 2>&>(getSpline(slice, row));
  const float* splineData = getSplineData(slice, row, 2);

  float duv[2];
  spline.interpolateU(splineData, gridU, gridV, duv);
  if (CAMath::Abs(duv[0]) > 100 || CAMath::Abs(duv[1]) > 100) {
    duv[0] = duv[1] = 0;
  }
  nomU = corrU - duv[0];
  nomV = corrV - duv[1];
}

GPUdi() float TPCFastSpaceChargeCorrection::getMaxDriftLength(int32_t slice, int32_t row, float pad) const
{
  const RowActiveArea& area = getSliceRowInfo(slice, row).activeArea;
  const float* c = area.maxDriftLengthCheb;
  float x = -1.f + 2.f * pad / mGeo.getRowInfo(row).maxPad;
  float y = c[0] + c[1] * x;
  float f0 = 1.f;
  float f1 = x;
  x *= 2.f;
  for (int32_t i = 2; i < 5; i++) {
    double f = x * f1 - f0;
    y += c[i] * f;
    f0 = f1;
    f1 = f;
  }
  return y;
}

GPUdi() float TPCFastSpaceChargeCorrection::getMaxDriftLength(int32_t slice, int32_t row) const
{
  return getSliceRowInfo(slice, row).activeArea.vMax;
}

GPUdi() float TPCFastSpaceChargeCorrection::getMaxDriftLength(int32_t slice) const
{
  return getSliceInfo(slice).vMax;
}

} // namespace gpu
} // namespace o2

#endif