TrTrackFill.cpp

Go to the documentation of this file.

#include "Containers/TrTrack.h"
#include "Utility/Logging/Logging.h"

#include "NtpMaker/NtpTools.h"

#include "bcorr.h"

#include <iostream>

namespace NAIA {
using namespace TrTrack;

static double RigidityCorrection() {
  float rigidityCorrection = 1;
  float fact;
  int retv(0);
  float mfcor[2] = {0, 0};
  retv = MagnetVarp::btempcor(fact, 0, 1);
  if (retv == 0)
    mfcor[0] = fact;
  retv = MagnetVarp::btempcor(fact, 0, 2);
  if (retv == 0)
    mfcor[1] = fact;
  if (mfcor[0] && mfcor[1])
    rigidityCorrection = (mfcor[0] + mfcor[1]) / 2;
  else if (mfcor[0])
    rigidityCorrection = mfcor[0];
  else if (mfcor[1])
    rigidityCorrection = mfcor[1];

  return rigidityCorrection == 0 ? 1 : rigidityCorrection;
}

static std::array<unsigned long int, numSpans> patterns = {3, 5, 6, 7, 0, 0, 3};
std::array<int, numFits> fitAlgoCode{1, 6};
static bool TrackHasExtHits(TrTrackR *trackPtr, Span span) {
  switch (span) {
  case Span::InnerOnly:
    return true;
  case Span::InnerL1:
    return trackPtr->TestHitLayerJHasXY(1);
  case Span::InnerL9:
    return trackPtr->TestHitLayerJHasXY(9);
  case Span::FullSpan:
    return (trackPtr->TestHitLayerJHasXY(1) && trackPtr->TestHitLayerJHasXY(9));
  case Span::UpperHalfInner:
  case Span::LowerHalfInner:
  case Span::InnerNoMS:
    return true;
  default:
    throw std::runtime_error("Span not supported");
  }
}

static constexpr std::array<double, 9> powersOfTen = {1, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8};

void ComputePatterns(TrTrackR *trackPtr) {
  patterns[Span::UpperHalfInner] = patterns[Span::LowerHalfInner] = 0;

  for (int il = 2; il < 9; il++) {
    if (!trackPtr->TestHitLayerJ(il))
      continue;

    if (il < 7) {
      patterns[Span::UpperHalfInner] += static_cast<unsigned long int>(9 * powersOfTen[il - 1]);
    }

    if (il > 2) {
      patterns[Span::LowerHalfInner] += static_cast<unsigned long int>(9 * powersOfTen[il - 1]);
    }
  }
}

bool TrTrackBaseData::_FitIDExists(Fit fit, Span span) const {
  auto itr = _fit_ID.find(fit);
  if (itr == end(_fit_ID)) {
    return false;
  }

  auto fit_ID_itr = itr->second.find(span);
  if (fit_ID_itr == end(itr->second)) {
    return false;
  }

  return (fit_ID_itr->second > 0);
}

Span TrTrackBaseData::_GetBestSpan(Fit fit) const {
  Span bestspan = _FitIDExists(fit, Span::FullSpan)
                      ? Span::FullSpan
                      : (_FitIDExists(fit, Span::InnerL1) > 0
                             ? Span::InnerL1
                             : (_FitIDExists(fit, Span::InnerL9) > 0 ? Span::InnerL9 : Span::InnerOnly));
  return bestspan;
}

void TrTrackBaseData::Fill(TrTrackR *trackPtr) {
  // some helper default variables
  static const int refit = _isMC ? 3 : 1;

  const float tmpCharge = floor(trackPtr->GetInnerQ_all(_beta).Mean + 0.5);
  const float fitMass = tmpCharge > 1 ? static_cast<float>(0.5 * TrFit::Mhelium * tmpCharge) : TrTrackR::DefaultMass;

  ComputePatterns(trackPtr);

  // Rigidity
  for (auto fit : fitTypes) {
    // Kalman with electron mass hyp. to be done in dedicated fill method
    if (fit == Fit::KalmanElectron)
      continue;

    for (auto span : spanTypes) {
      if (TrackHasExtHits(trackPtr, span)) {
        Logging::MuteOutput(); // this can get quite verbose...
        _fit_ID[fit][span] =
            trackPtr->iTrTrackPar(fitAlgoCode[fit], static_cast<int>(patterns[span]), 20 + refit, fitMass, tmpCharge);
        Logging::UnmuteOutput();
        if (_fit_ID[fit][span] > 0) {
          RigidityCorr[fit][span] =
              static_cast<float>(_isMC ? trackPtr->GetRigidity(_fit_ID[fit][span], 1) // no corr for MC
                                       : trackPtr->GetCorrectedRigidity(_fit_ID[fit][span], 3, 1));
          TrChiSq[fit][span][Side::X] =
              static_cast<float>(trackPtr->GetChisqX(_fit_ID[fit][span]) / trackPtr->GetNdofX(_fit_ID[fit][span]));
          TrChiSq[fit][span][Side::Y] =
              static_cast<float>(trackPtr->GetChisqY(_fit_ID[fit][span]) / trackPtr->GetNdofY(_fit_ID[fit][span]));
        }
      }
    }
  }

  // Charge
  //  - Inner tracker
  Charge[ChargeRecoType::STD] = trackPtr->GetQ_all(_beta).Mean;
  Charge[ChargeRecoType::HL] = trackPtr->GetQH_all(2, _beta).Mean;
  // Charge[ChargeRecoType::YJ] does not support "global" charge;
  InnerCharge[ChargeRecoType::STD] = trackPtr->GetInnerQ_all(_beta).Mean;
  InnerCharge[ChargeRecoType::HL] = trackPtr->GetInnerQH_all(2, _beta).Mean;
  InnerCharge[ChargeRecoType::YJ] = trackPtr->GetQYJ_all(2, _beta).InnerQ;

  //  - single layers
  for (int ijl = 0; ijl < 9; ijl++) {
    Fit fit = Fit::Choutko;
    Span span = Span::NoSpan;
    switch (ijl + 1) {
    case 1:
      span = Span::InnerL1;
      break;
    case 2:
    case 3:
    case 4:
    case 5:
    case 6:
    case 7:
    case 8:
      span = Span::InnerOnly;
      break;
    case 9:
      span = Span::InnerL9;
      break;
    }

    if (trackPtr->GetHitLJ(ijl + 1)) {
      LayerChargeStatus[ijl] = trackPtr->GetHitLJ(ijl + 1)->GetQStatus();

      float rigidity = (_fit_ID[fit][span] > 0)
                           ? RigidityCorr[fit][span]
                           : ((_fit_ID[fit][Span::InnerOnly] > 0) ? RigidityCorr[fit][Span::InnerOnly] : 0);

      LayerChargeXY[ijl][ChargeRecoType::STD] = trackPtr->GetHitLJ(ijl + 1)->GetQ(2, _beta, rigidity);
      LayerChargeXY[ijl][ChargeRecoType::HL] = trackPtr->GetLayerJQH(ijl + 1, 2, _beta, _fit_ID[fit][span]);
      LayerChargeXY[ijl][ChargeRecoType::YJ] = trackPtr->GetLayerQYJ(ijl + 1, 2, _beta, _fit_ID[fit][span]);
    }
  }
}

void TrTrackPlusData::Fill(TrTrackR *trackPtr) {
  const auto &_beta = trkBase->_beta;
  const int refit = trkBase->_isMC ? 3 : 1;
  static unsigned int prescaleCounter = 0;
  static float prescaleRigidityThreshold = -1; // -1 GV

  if (trkBase->InnerCharge.empty())
    return;

  auto *pev = AMSEventR::Head();

  const float tmpCharge = floor(trkBase->InnerCharge.at(ChargeRecoType::STD) + 0.5);
  const float fitMass = tmpCharge > 1 ? static_cast<float>(0.5 * TrFit::Mhelium * tmpCharge) : TrTrackR::DefaultMass;

  auto getPartialPattern = [](TrTrackR *track, unsigned int jLayer) {
    unsigned int pattern = 0;
    for (unsigned int ijl = 2; ijl < 9; ijl++) {
      if (track->TestHitLayerJ(ijl) && ijl != jLayer)
        pattern += 9 * powersOfTen[ijl - 1];
    }
    return pattern;
  };

  for (unsigned int ijl = 2; ijl < 9; ijl++) {
    TrackFeetDistance[ijl - 1] = pev->GetTkFeetDist(ijl, trackPtr);
  }

  // Rigidity
  for (auto fit : fitTypes) {
    // Kalman with electron mass hyp. to be done in dedicated fill method
    if (fit == Fit::KalmanElectron)
      continue;

    for (auto span : spanTypes) {
      if (trkBase->FitIDExists(fit, span)) {
        _fit_ID[fit][span] =
            trackPtr->iTrTrackPar(fitAlgoCode[fit], static_cast<int>(patterns[span]), 20 + refit, fitMass, tmpCharge);
        if (_fit_ID[fit][span] < 0)
          continue;
        Rigidity[fit][span] = static_cast<float>(trackPtr->GetRigidity(_fit_ID[fit][span], 1));
        if (fit == Fit::Kalman) {
          RigidityTOI[fit][span] =
              static_cast<float>(trkBase->_isMC ? trackPtr->GetRigidity(_fit_ID[fit][span], 0)
                                                : trackPtr->GetCorrectedRigidity(_fit_ID[fit][span], 3, 0));
        }

        const TrTrackPar &trackPar = trackPtr->gTrTrackPar(_fit_ID[fit][span]);
        InvRigErr[fit][span] = static_cast<float>(trackPar.ErrRinv / (!trackPar.BcorrFlag ? RigidityCorrection() : 1));
      }
    }

    // fill prescaled ("partial") variables
    // doesn't make sense to try if inner only fit failed...
    if (trkBase->FitIDExists(fit, Span::InnerOnly)) {
      bool bypassPrescaling =
          _beta > 0 && trkBase->RigidityCorr.at(fit).at(Span::InnerOnly) < prescaleRigidityThreshold;

      if (!prescaleCounter || bypassPrescaling) {
        for (unsigned int ijl = 2; ijl < 9; ijl++) {
          if (trackPtr->TestHitLayerJ(ijl))
            continue;
          auto partialPattern = getPartialPattern(trackPtr, ijl);

          int pfit_ID = trackPtr->iTrTrackPar(fitAlgoCode[fit], partialPattern, 20 + refit, fitMass, tmpCharge);
          if (pfit_ID < 0)
            continue;

          const TrTrackPar &partialTrackPar = trackPtr->gTrTrackPar(pfit_ID);
          PartialRigidity[ijl - 1][fit] =
              partialTrackPar.Rigidity * (!partialTrackPar.BcorrFlag ? RigidityCorrection() : 1);
          PartialInvRigErr[ijl - 1][fit] =
              partialTrackPar.ErrRinv / (!partialTrackPar.BcorrFlag ? RigidityCorrection() : 1);
          PartialTrChiSq[ijl - 1][fit][Side::X] = static_cast<float>(partialTrackPar.ChisqX / partialTrackPar.NdofX);
          PartialTrChiSq[ijl - 1][fit][Side::Y] = static_cast<float>(partialTrackPar.ChisqY / partialTrackPar.NdofY);
        }
      }
    }
  }

  prescaleCounter = ++prescaleCounter % 100;

  // Charge
  //  - Inner tracker
  ChargeRMS[ChargeRecoType::STD] = trackPtr->GetQ_all(_beta).RMS;
  ChargeRMS[ChargeRecoType::HL] = trackPtr->GetQH_all(2, _beta).RMS;
  // ChargeRMS[ChargeRecoType::YJ] = trackPtr->GetQYJ_all(2, _beta).InnerQRMS;
  InnerChargeRMS[ChargeRecoType::STD] = trackPtr->GetInnerQ_all(_beta).RMS;
  InnerChargeRMS[ChargeRecoType::HL] = trackPtr->GetInnerQH_all(2, _beta).RMS;
  InnerChargeRMS[ChargeRecoType::YJ] = trackPtr->GetQYJ_all(2, _beta).InnerQRMS;

  //  - single layers
  for (int ijl = 0; ijl < 9; ijl++) {
    Fit fit = Fit::Choutko;
    Span span = Span::NoSpan;
    switch (ijl + 1) {
    case 1:
      span = Span::InnerL1;
      break;
    case 2:
    case 3:
    case 4:
    case 5:
    case 6:
    case 7:
    case 8:
      span = Span::InnerOnly;
      break;
    case 9:
      span = Span::InnerL9;
      break;
    }

    int fit_ID = 0;
    if (trkBase->FitIDExists(fit, span))
      fit_ID = _fit_ID.at(fit).at(span);

    if (trackPtr->GetHitLJ(ijl + 1)) {
      float rigidity =
          (_fit_ID[fit][span] > 0)
              ? trkBase->RigidityCorr.at(fit).at(span)
              : ((_fit_ID[fit][Span::InnerOnly] > 0) ? trkBase->RigidityCorr.at(fit).at(Span::InnerOnly) : 0);

      LayerCharge[ijl][ChargeRecoType::STD][Side::X] = trackPtr->GetHitLJ(ijl + 1)->GetQ(0, _beta, rigidity);
      LayerCharge[ijl][ChargeRecoType::HL][Side::X] = trackPtr->GetLayerJQH(ijl + 1, 0, _beta, fit_ID);
      LayerCharge[ijl][ChargeRecoType::YJ][Side::X] = trackPtr->GetLayerQYJ(ijl + 1, 0, _beta, fit_ID);

      if (trackPtr->GetHitLJ(ijl + 1)->GetXCluster())
        LayerEdep[ijl][Side::X] = trackPtr->GetHitLJ(ijl + 1)->GetXCluster()->GetEdep();
      if (trackPtr->GetHitLJ(ijl + 1)->GetYCluster())
        LayerEdep[ijl][Side::Y] = trackPtr->GetHitLJ(ijl + 1)->GetYCluster()->GetEdep();

      LayerCharge[ijl][ChargeRecoType::STD][Side::Y] = trackPtr->GetHitLJ(ijl + 1)->GetQ(1, _beta, rigidity);
      LayerCharge[ijl][ChargeRecoType::HL][Side::Y] = trackPtr->GetLayerJQH(ijl + 1, 1, _beta, fit_ID);
      LayerCharge[ijl][ChargeRecoType::YJ][Side::Y] = trackPtr->GetLayerQYJ(ijl + 1, 1, _beta, fit_ID);
    }
  }

  // geometry and position
  for (int ijl = 0; ijl < 9; ijl++) {
    if (trackPtr->GetHitLJ(ijl + 1) && trackPtr->GetHitLJ(ijl + 1)->GetXCluster())
      TrTrackHitPos[ijl][Side::X] = static_cast<float>(trackPtr->GetHitCooLJ(ijl + 1).x());
    if (trackPtr->GetHitLJ(ijl + 1) && trackPtr->GetHitLJ(ijl + 1)->GetYCluster())
      TrTrackHitPos[ijl][Side::Y] = static_cast<float>(trackPtr->GetHitCooLJ(ijl + 1).y());
  }

  for (const auto fitPos : fitPositionHeights) {
    // TODO: compute track position at Ecal COG
    for (auto fit : fitTypes) {
      for (auto span : spanTypes) {
        if (trkBase->FitIDExists(fit, span)) {
          int fit_ID = _fit_ID.at(fit).at(span);
          if (fit_ID < 0)
            continue;
          AMSPoint trackPos;
          AMSDir trackDir;
          trackPtr->Interpolate(fitPositionHeightZ[fitPos], trackPos, trackDir, fit_ID);
          TrTrackFitPos[fitPos][fit][span][Side::X] = static_cast<float>(trackPos.x());
          TrTrackFitPos[fitPos][fit][span][Side::Y] = static_cast<float>(trackPos.y());
        }
      }
    }
  }

  Span bestspan = trkBase->GetBestSpan(Fit::Choutko);
  if (_fit_ID.at(Fit::Choutko).at(bestspan) > 0) {
    Theta = static_cast<float>(trackPtr->GetTheta(_fit_ID.at(Fit::Choutko).at(bestspan)));
    Phi = static_cast<float>(trackPtr->GetPhi(_fit_ID.at(Fit::Choutko).at(bestspan)));
  }

  for (auto fit : fitTypes) {
    for (auto span : spanTypes) {
      if (trkBase->FitIDExists(fit, span)) {
        int fit_ID = _fit_ID.at(fit).at(span);
        if (fit_ID < 0)
          continue;
        for (int ijl = 0; ijl < 9; ijl++) {
          if (trackPtr->GetHitLJ(ijl + 1)) {
            AMSPoint point = trackPtr->GetResidualJ(ijl + 1, fit_ID);
            TrTrackResidual[ijl][fit][span][Side::X] = static_cast<float>(point.x());
            TrTrackResidual[ijl][fit][span][Side::Y] = static_cast<float>(point.y());
          }
        }
      }
    }
  }

  static constexpr std::array<float, 5> window = {0.1, 1., 2., 5., 10.};

  // from dbar code
  for (int jLayer = 1; jLayer <= 9; jLayer++) {
    ClusterSignalRatio[jLayer - 1] = pev->GetTrackerRawSignalRatio(jLayer, 10);

    // identify the clusters used in trtrack
    TrClusterR *trackClusterPtrX = trackPtr->GetHitLJ(jLayer) ? trackPtr->GetHitLJ(jLayer)->GetXCluster() : nullptr;
    TrClusterR *trackClusterPtrY = trackPtr->GetHitLJ(jLayer) ? trackPtr->GetHitLJ(jLayer)->GetYCluster() : nullptr;

    // Get track interpolation to jLayer
    AMSPoint fitcoo;
    AMSDir fitdir;
    trackPtr->InterpolateLayerJ(jLayer, fitcoo, fitdir, _fit_ID[Fit::Kalman][Span::InnerOnly]);

    // TODO: Check if zero threshold is correct
    static constexpr float noiseThreshold = 0;

    for (int icl = 0; icl < pev->nTrCluster(); icl++) {
      auto *cluster = pev->pTrCluster(icl);

      if (!cluster || cluster->GetLayerJ() != jLayer)
        continue;

      float edep = 1000 * cluster->GetEdep();
      if (cluster == trackClusterPtrX) {
        NClusters[jLayer - 1][DistanceFromTrack::OnTrack][Side::X]++;
        ClustersEdep[jLayer - 1][DistanceFromTrack::OnTrack][Side::X] += edep;
        continue;
      }
      if (cluster == trackClusterPtrY) {
        NClusters[jLayer - 1][DistanceFromTrack::OnTrack][Side::Y]++;
        ClustersEdep[jLayer - 1][DistanceFromTrack::OnTrack][Side::Y] += edep;
        continue;
      }

      if (edep < noiseThreshold)
        continue; // apply a threshold to account for noise

      // Deal with multiplicities: Choose closest to track
      float clusterDistanceToTrack = std::numeric_limits<float>::max();
      for (int im = 0; im < cluster->GetMultiplicity(); im++) {
        float coo = cluster->GetGCoord(im);

        clusterDistanceToTrack =
            std::min(clusterDistanceToTrack, fabs(static_cast<float>(fitcoo[cluster->GetSide()] - coo)));
      }

      Side side = cluster->GetSide() == 0 ? Side::X : Side::Y;

      NClusters[jLayer - 1][DistanceFromTrack::AllLayer][side]++;
      ClustersEdep[jLayer - 1][DistanceFromTrack::AllLayer][side] += edep;
      if (edep > MaxClusterEdep[jLayer - 1][side]) {
        MaxClusterEdep[jLayer - 1][side] = edep;
        MaxClusterDistance[jLayer - 1][side] = clusterDistanceToTrack;
      }

      // FIXME: I don't really like using an int instead of the right enum
      for (unsigned int iw = 0; iw < window.size(); iw++) {
        if (std::fabs(clusterDistanceToTrack) < window[iw]) {
          NClusters[jLayer - 1][static_cast<DistanceFromTrack>(iw + 1)][side]++;
          ClustersEdep[jLayer - 1][static_cast<DistanceFromTrack>(iw + 1)][side] += edep;
        }
      }
    }
  }
}

void TrTrackBaseData::FillElectronVars(TrTrackR *trackPtr, bool refitKalman) {
  // Rigidity
  auto fit = Fit::KalmanElectron;
  auto span = _GetBestSpan(Fit::Kalman);

  if (TrackHasExtHits(trackPtr, span) && refitKalman) {
    Logging::MuteOutput(); // this can get quite verbose...
    _fit_ID[fit][span] = trackPtr->iTrTrackPar(fitAlgoCode[fit], static_cast<int>(patterns[span]), 22,
                                               static_cast<float>(TrFit::Melectron), 1);
    Logging::UnmuteOutput();
    if (_fit_ID[fit][span] > 0) {
      RigidityCorr[fit][span] =
          static_cast<float>(_isMC ? trackPtr->GetRigidity(_fit_ID[fit][span], 1) // no corr for MC
                                   : trackPtr->GetCorrectedRigidity(_fit_ID[fit][span], 3, 1));
      TrChiSq[fit][span][Side::X] =
          static_cast<float>(trackPtr->GetChisqX(_fit_ID[fit][span]) / trackPtr->GetNdofX(_fit_ID[fit][span]));
      TrChiSq[fit][span][Side::Y] =
          static_cast<float>(trackPtr->GetChisqY(_fit_ID[fit][span]) / trackPtr->GetNdofY(_fit_ID[fit][span]));
    }
  }
}

void TrTrackPlusData::FillElectronVars(TrTrackR *trackPtr, float zEcalCOG, bool refitKalman) {
  const auto &base_fit_ID = trkBase->_fit_ID;

  // Get Electron fit rigidity and chisquares
  {
    auto fit = Fit::KalmanElectron;
    auto span = trkBase->GetBestSpan(Fit::Kalman);
    if (trkBase->FitIDExists(fit, span) && refitKalman) {
      int fit_ID = base_fit_ID.at(fit).at(span);
      Rigidity[fit][span] = static_cast<float>(trackPtr->GetRigidity(fit_ID, 1));
      RigidityTOI[fit][span] = static_cast<float>(trkBase->_isMC ? trackPtr->GetRigidity(fit_ID, 0)
                                                                 : trackPtr->GetCorrectedRigidity(fit_ID, 3, 0));
      const TrTrackPar &trackPar = trackPtr->gTrTrackPar(fit_ID);
      InvRigErr[fit][span] = static_cast<float>(trackPar.ErrRinv / (!trackPar.BcorrFlag ? RigidityCorrection() : 1));
    }
  }

  if (zEcalCOG < std::numeric_limits<float>::max()) {
    for (auto fit : fitTypes) {
      for (auto span : spanTypes) {
        if (trkBase->FitIDExists(fit, span)) {
          int fit_ID = base_fit_ID.at(fit).at(span);
          AMSPoint trackPos;
          AMSDir trackDir;
          trackPtr->Interpolate(zEcalCOG, trackPos, trackDir, fit_ID);
          TrTrackFitPos[FitPositionHeight::EcalCOG][fit][span][Side::X] = static_cast<float>(trackPos.x());
          TrTrackFitPos[FitPositionHeight::EcalCOG][fit][span][Side::Y] = static_cast<float>(trackPos.y());
        }
      }
    }
  }
}

// void TrTrackPlusData::Fill(const AMSEventR *eventPtr) {}
} // namespace NAIA