TPCScalerSpec.cxx

Go to the documentation of this file.

// 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.
 
 
#include "Framework/Task.h"
#include "Framework/Logger.h"
#include "Framework/DataProcessorSpec.h"
#include "Framework/CCDBParamSpec.h"
#include "Framework/ConfigParamRegistry.h"
#include "TPCBaseRecSim/CDBInterface.h"
#include "DetectorsBase/GRPGeomHelper.h"
#include "TPCCalibration/TPCScaler.h"
#include "TPCCalibration/TPCMShapeCorrection.h"
#include "TTree.h"
#include "TPCCalibration/TPCFastSpaceChargeCorrectionHelper.h"
#include "TPCSpaceCharge/SpaceCharge.h"
#include "CommonUtils/TreeStreamRedirector.h"
#include "TPCCalibration/CorrectionMapsLoader.h"
#include "TPCCalibration/VDriftHelper.h"
 
using namespace o2::framework;
 
namespace o2
{
namespace tpc
{
 
class TPCScalerSpec : public Task
{
 public:
  TPCScalerSpec(std::shared_ptr<o2::base::GRPGeomRequest> req, const o2::tpc::CorrectionMapsGloOpts& sclOpts, bool enableIDCs, bool enableMShape) : mCCDBRequest(req), mEnableIDCs(enableIDCs), mEnableMShape(enableMShape), mGlobOpts(sclOpts)
  {
    mTPCCorrMapsLoader.setLumiScaleType(sclOpts.lumiType);
    mTPCCorrMapsLoader.setLumiScaleMode(sclOpts.lumiMode);
    mTPCCorrMapsLoader.setCheckCTPIDCConsistency(sclOpts.checkCTPIDCconsistency);
  };
 
  void init(framework::InitContext& ic) final
  {
    o2::base::GRPGeomHelper::instance().setRequest(mCCDBRequest);
    mIonDriftTimeMS = ic.options().get<float>("ion-drift-time");
    mMaxTimeWeightsMS = ic.options().get<float>("max-time-for-weights");
    mMShapeScalingFac = ic.options().get<float>("m-shape-scaling-factor");
    mEnableWeights = !(ic.options().get<bool>("disableWeights"));
    const bool enableStreamer = ic.options().get<bool>("enableStreamer");
    const int mshapeThreads = ic.options().get<int>("n-threads");
    mKnotsYMshape = ic.options().get<int>("n-knots-y");
    mKnotsZMshape = ic.options().get<int>("n-knots-z");
    TPCFastSpaceChargeCorrectionHelper::instance()->setNthreads(mshapeThreads);
    o2::tpc::SpaceCharge<double>::setNThreads(mshapeThreads);
 
    if (enableStreamer) {
      mStreamer = std::make_unique<o2::utils::TreeStreamRedirector>("M_Shape.root", "recreate");
    }
    mTPCCorrMapsLoader.init(ic, mEnableIDCs);
  }
 
  void endOfStream(EndOfStreamContext& eos) final
  {
    if (mStreamer) {
      mStreamer->Close();
    }
  }
 
  void run(ProcessingContext& pc) final
  {
    o2::base::GRPGeomHelper::instance().checkUpdates(pc);
    mTPCVDriftHelper.extractCCDBInputs(pc);
    if (mTPCVDriftHelper.isUpdated()) {
      mTPCVDriftHelper.acknowledgeUpdate();
    }
 
    if (mEnableIDCs && pc.inputs().isValid("tpcscaler")) {
      pc.inputs().get<TTree*>("tpcscaler");
    }
 
    if (mEnableWeights && mEnableIDCs) {
      if (pc.inputs().isValid("tpcscalerw")) {
        pc.inputs().get<TPCScalerWeights*>("tpcscalerw");
      }
    }
    if (mEnableMShape && pc.inputs().isValid("mshape")) {
      pc.inputs().get<TTree*>("mshape");
    }
 
    const auto orbitResetTimeMS = o2::base::GRPGeomHelper::instance().getOrbitResetTimeMS();
    const auto firstTFOrbit = pc.services().get<o2::framework::TimingInfo>().firstTForbit;
    const double timestamp = orbitResetTimeMS + firstTFOrbit * o2::constants::lhc::LHCOrbitMUS * 0.001;
    int currRun = pc.services().get<o2::framework::TimingInfo>().runNumber;
    if (mEnableMShape) {
      static int runWarningMS = -1;
      if ((mMShapeTPCScaler.getRun() != -1) && currRun != mMShapeTPCScaler.getRun() && runWarningMS != currRun) {
        LOGP(alarm, "Run number {} of processed data and run number {} of loaded TPC M-shape scaler doesnt match!", pc.services().get<o2::framework::TimingInfo>().runNumber, mMShapeTPCScaler.getRun());
        runWarningMS = currRun;
      }
 
      const auto& boundaryPotential = mMShapeTPCScaler.getBoundaryPotential(timestamp);
      if (!boundaryPotential.mPotential.empty()) {
        LOGP(info, "Calculating M-shape correction from input boundary potential");
        const Side side = Side::A;
        o2::tpc::SpaceCharge<double> sc = o2::tpc::SpaceCharge<double>(mMShapeTPCScaler.getMShapes().mBField, mMShapeTPCScaler.getMShapes().mZ, mMShapeTPCScaler.getMShapes().mR, mMShapeTPCScaler.getMShapes().mPhi);
        for (int iz = 0; iz < mMShapeTPCScaler.getMShapes().mZ; ++iz) {
          for (int iphi = 0; iphi < mMShapeTPCScaler.getMShapes().mPhi; ++iphi) {
            const float pot = mMShapeScalingFac * boundaryPotential.mPotential[iz];
            sc.setPotential(iz, 0, iphi, side, pot);
          }
        }
 
        sc.poissonSolver(side);
        sc.calcEField(side);
        sc.calcGlobalCorrections(sc.getElectricFieldsInterpolator(side));
 
        std::function<void(int, double, double, double, double&, double&, double&)> getCorrections = [&sc = sc](const int roc, double x, double y, double z, double& dx, double& dy, double& dz) {
          Side side = roc < 18 ? Side::A : Side::C;
          if (side == Side::C) {
            dx = 0;
            dy = 0;
            dz = 0;
          } else {
            sc.getCorrections(x, y, z, side, dx, dy, dz);
          }
        };
 
        std::unique_ptr<TPCFastSpaceChargeCorrection> spCorrection = TPCFastSpaceChargeCorrectionHelper::instance()->createFromGlobalCorrection(getCorrections, mKnotsYMshape, mKnotsZMshape);
        std::unique_ptr<TPCFastTransform> fastTransform(TPCFastTransformHelperO2::instance()->create(0, *spCorrection));
        mTPCCorrMapsLoader.setCorrMapMShape(std::move(fastTransform));
      }
 
      if (mStreamer) {
        (*mStreamer) << "treeMShape"
                     << "firstTFOrbit=" << firstTFOrbit
                     << "timestamp=" << timestamp
                     << "boundaryPotential=" << boundaryPotential
                     << "mMShapeScalingFac=" << mMShapeScalingFac
                     << "\n";
      }
    }
 
    float tpcScaler = -1.f;
    if (mEnableIDCs) {
      static int runWarningIDC = -1;
      if (pc.services().get<o2::framework::TimingInfo>().runNumber != mTPCScaler.getRun() && runWarningIDC != currRun) {
        LOGP(alarm, "Run number {} of processed data and run number {} of loaded TPC scaler doesnt match!", pc.services().get<o2::framework::TimingInfo>().runNumber, mTPCScaler.getRun());
        runWarningIDC = currRun;
      }
      float scalerA = mTPCScaler.getMeanScaler(timestamp, o2::tpc::Side::A);
      float scalerC = mTPCScaler.getMeanScaler(timestamp, o2::tpc::Side::C);
      float meanScaler = (scalerA + scalerC) / 2;
      tpcScaler = meanScaler;
      if (mStreamer) {
        (*mStreamer) << "treeIDC"
                     << "scalerA=" << scalerA
                     << "scalerC=" << scalerC
                     << "firstTFOrbit=" << firstTFOrbit
                     << "timestamp=" << timestamp
                     << "\n";
      }
    }
    // check for Maps update
    mTPCCorrMapsLoader.extractCCDBInputs(pc, tpcScaler);
 
    if (mGlobOpts.requestCTPLumi) {
      const float lumiCTP = mTPCCorrMapsLoader.getInstLumiCTP();
      // if CTP lumi was notrequest - defualt of 0 is published, otherwise the value is scaled with the provided factor
      LOGP(info, "Publishing CTP Lumi: {} for timestamp: {}, firstTFOrbit: {}", lumiCTP, timestamp, firstTFOrbit);
      pc.outputs().snapshot(Output{header::gDataOriginCTP, "LUMICTP"}, lumiCTP);
    }
 
    buildMap(pc);
  }
 
  void buildMap(ProcessingContext& pc)
  {
    // reference map
    auto* corrMap = mTPCCorrMapsLoader.getCorrMap();
 
    // // new correction map
    o2::gpu::TPCFastTransform finalMap;
    finalMap.cloneFromObject(*corrMap, nullptr);
    finalMap.setApplyCorrectionOn();
 
    const auto* corrMapRef = mTPCCorrMapsLoader.getCorrMapRef();
    const float lumiScale = mTPCCorrMapsLoader.getLumiScale();
    std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, float>> additionalCorrections;
 
    // if standard scaling is used: map(lumi) = (mean_map - ref_map) * lumiScale + ref_map
    if (mTPCCorrMapsLoader.getLumiScaleMode() == LumiScaleMode::Linear) {
      const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, float>> step0{{&(corrMapRef->getCorrection()), -1.f}};
      // finalMap = (mean_map - finalMap)
      TPCFastSpaceChargeCorrectionHelper::instance()->mergeCorrections(finalMap.getCorrection(), 1, step0, true);
 
      // finalMap = finalMap * lumiScale + ref_map
      const std::vector<std::pair<const o2::gpu::TPCFastSpaceChargeCorrection*, float>> step1{{&(corrMapRef->getCorrection()), 1.f}};
      TPCFastSpaceChargeCorrectionHelper::instance()->mergeCorrections(finalMap.getCorrection(), lumiScale, step1, true);
 
    } else if (mTPCCorrMapsLoader.getLumiScaleMode() == LumiScaleMode::DerivativeMap || mTPCCorrMapsLoader.getLumiScaleMode() == LumiScaleMode::DerivativeMapMC) {
      additionalCorrections.emplace_back(&(corrMapRef->getCorrection()), lumiScale);
    }
 
    // if mshape map valid
    if (!mTPCCorrMapsLoader.isCorrMapMShapeDummy()) {
      LOGP(info, "Adding M-shape correction to the final map with scaling factor {}", mMShapeScalingFac);
      additionalCorrections.emplace_back(&(mTPCCorrMapsLoader.getCorrMapMShape()->getCorrection()), 1.f);
    }
 
    if (!additionalCorrections.empty()) {
      TPCFastSpaceChargeCorrectionHelper::instance()->mergeCorrections(finalMap.getCorrection(), 1, additionalCorrections, true);
    }
 
    Output corrMapOutput{header::gDataOriginTPC, "TPCCORRMAP", 0};
    auto outputBuffer = o2::pmr::vector<char>(pc.outputs().getMemoryResource(corrMapOutput));
    outputBuffer.resize(TPCFastTransformPOD::estimateSize(finalMap.getCorrection()));
    auto* pod = TPCFastTransformPOD::create(outputBuffer.data(), outputBuffer.size(), finalMap);
    const auto& vd = mTPCVDriftHelper.getVDriftObject();
    o2::tpc::TPCFastTransformHelperO2::instance()->updateCalibration(*pod, 0, vd.corrFact, vd.refVDrift, vd.getTimeOffset());
    pc.outputs().adoptContainer(corrMapOutput, std::move(outputBuffer));
  }
 
  void finaliseCCDB(o2::framework::ConcreteDataMatcher& matcher, void* obj) final
  {
    o2::base::GRPGeomHelper::instance().finaliseCCDB(matcher, obj);
    mTPCVDriftHelper.accountCCDBInputs(matcher, obj);
    if (matcher == ConcreteDataMatcher(o2::header::gDataOriginTPC, "TPCSCALERCCDB", 0)) {
      LOGP(info, "Updating TPC scaler");
      mTPCScaler.setFromTree(*((TTree*)obj));
      if (mIonDriftTimeMS > 0) {
        LOGP(info, "Setting ion drift time to: {}", mIonDriftTimeMS);
        mTPCScaler.setIonDriftTimeMS(mIonDriftTimeMS);
      }
      if (mScalerWeights.isValid()) {
        LOGP(info, "Setting TPC scaler weights");
        mTPCScaler.setScalerWeights(mScalerWeights);
        mTPCScaler.useWeights(true);
        if (mIonDriftTimeMS == -1) {
          overWriteIntegrationTime();
        }
      }
    }
    if (matcher == ConcreteDataMatcher(o2::header::gDataOriginTPC, "TPCSCALERWCCDB", 0)) {
      LOGP(info, "Updating TPC scaler weights");
      mScalerWeights = *(TPCScalerWeights*)obj;
      mTPCScaler.setScalerWeights(mScalerWeights);
      mTPCScaler.useWeights(true);
      // in case ion drift time is not specified it is overwritten by the value in the weight object
      if (mIonDriftTimeMS == -1) {
        overWriteIntegrationTime();
      }
    }
    if (matcher == ConcreteDataMatcher(o2::header::gDataOriginTPC, "MSHAPEPOTCCDB", 0)) {
      LOGP(info, "Updating M-shape TPC scaler");
      mMShapeTPCScaler.setFromTree(*((TTree*)obj));
      if (mMShapeTPCScaler.getRun() == -1) {
        LOGP(info, "Loaded default M-Shape correction object from CCDB");
      }
    }
    if (mTPCCorrMapsLoader.accountCCDBInputs(matcher, obj)) {
      return;
    }
  }
 
 private:
  std::shared_ptr<o2::base::GRPGeomRequest> mCCDBRequest;     
  const bool mEnableIDCs{true};                               
  const bool mEnableMShape{false};                            
  const o2::tpc::CorrectionMapsGloOpts mGlobOpts;             
  bool mEnableWeights{false};                                 
  TPCScalerWeights mScalerWeights{};                          
  float mIonDriftTimeMS{-1};                                  
  float mMaxTimeWeightsMS{500};                               
  TPCScaler mTPCScaler;                                       
  float mMShapeScalingFac{0};                                 
  TPCMShapeCorrection mMShapeTPCScaler;                       
  int mKnotsYMshape{4};                                       
  int mKnotsZMshape{4};                                       
  std::unique_ptr<o2::utils::TreeStreamRedirector> mStreamer; 
  o2::tpc::CorrectionMapsLoader mTPCCorrMapsLoader{};
  o2::tpc::VDriftHelper mTPCVDriftHelper{}; 
 
  void overWriteIntegrationTime()
  {
    float integrationTime = std::abs(mScalerWeights.mFirstTimeStampMS);
    if (integrationTime <= 0) {
      return;
    }
    if (integrationTime > mMaxTimeWeightsMS) {
      integrationTime = mMaxTimeWeightsMS;
    }
    LOGP(info, "Setting maximum integration time for weights to: {}", integrationTime);
    mTPCScaler.setIonDriftTimeMS(integrationTime);
  }
};
 
o2::framework::DataProcessorSpec getTPCScalerSpec(bool enableIDCs, bool enableMShape, const o2::tpc::CorrectionMapsGloOpts& sclOpts)
{
  std::vector<InputSpec> inputs;
  if (enableIDCs) {
    LOGP(info, "Publishing IDC scalers for space-charge distortion fluctuation correction");
    inputs.emplace_back("tpcscaler", o2::header::gDataOriginTPC, "TPCSCALERCCDB", 0, Lifetime::Condition, ccdbParamSpec(o2::tpc::CDBTypeMap.at(o2::tpc::CDBType::CalScaler), {}, 1));          // time-dependent
    inputs.emplace_back("tpcscalerw", o2::header::gDataOriginTPC, "TPCSCALERWCCDB", 0, Lifetime::Condition, ccdbParamSpec(o2::tpc::CDBTypeMap.at(o2::tpc::CDBType::CalScalerWeights), {}, 0)); // non time-dependent
  }
  if (enableMShape) {
    LOGP(info, "Publishing M-shape correction map");
    inputs.emplace_back("mshape", o2::header::gDataOriginTPC, "MSHAPEPOTCCDB", 0, Lifetime::Condition, ccdbParamSpec(o2::tpc::CDBTypeMap.at(o2::tpc::CDBType::CalMShape), {}, 1)); // time-dependent
  }
 
  auto ccdbRequest = std::make_shared<o2::base::GRPGeomRequest>(true,                           // orbitResetTime
                                                                false,                          // GRPECS=true for nHBF per TF
                                                                false,                          // GRPLHCIF
                                                                false,                          // GRPMagField
                                                                false,                          // askMatLUT
                                                                o2::base::GRPGeomRequest::None, // geometry
                                                                inputs);
 
  std::vector<OutputSpec> outputs;
  outputs.emplace_back(o2::header::gDataOriginTPC, "TPCCORRMAP", 0, Lifetime::Timeframe);
  if (sclOpts.requestCTPLumi) {
    outputs.emplace_back(o2::header::gDataOriginCTP, "LUMICTP", 0, Lifetime::Timeframe);
  }
  o2::tpc::VDriftHelper::requestCCDBInputs(inputs);
  o2::tpc::CorrectionMapsLoader::requestCCDBInputs(inputs, sclOpts);
 
  return DataProcessorSpec{
    "tpc-scaler",
    inputs,
    outputs,
    AlgorithmSpec{adaptFromTask<TPCScalerSpec>(ccdbRequest, sclOpts, enableIDCs, enableMShape)},
    Options{
      {"ion-drift-time", VariantType::Float, -1.f, {"Overwrite ion drift time if a value >0 is provided"}},
      {"max-time-for-weights", VariantType::Float, 500.f, {"Maximum possible integration time in ms when weights are used"}},
      {"m-shape-scaling-factor", VariantType::Float, 1.f, {"Scale M-shape scaler with this value"}},
      {"disableWeights", VariantType::Bool, false, {"Disable weights for TPC scalers"}},
      {"enableStreamer", VariantType::Bool, false, {"Enable streaming of M-shape scalers"}},
      {"n-threads", VariantType::Int, 4, {"Number of threads used for the M-shape correction"}},
      {"n-knots-y", VariantType::Int, 4, {"Number of knots in y-direction used for the M-shape correction"}},
      {"n-knots-z", VariantType::Int, 4, {"Number of knots in z-direction used for the M-shape correction"}}}};
}
 
} // namespace tpc
} // namespace o2