CbmBbaAlignTask.cxx

Go to the documentation of this file.

/* Copyright (C) 2025 Johann Wolfgang Goethe-Universitaet Frankfurt, Frankfurt am Main
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Nora Bluhme [committer], Sergey Gorbunov */

 
#include "CbmBbaAlignTask.h"
 
#include "CbmBbaAlignmentBody.h"
#include "CbmL1.h"
#include "FairRootManager.h"
#include "FairRunAna.h"
#include "TClonesArray.h"
#include "TFile.h"
#include "TGeoManager.h"
#include "TRandom.h"
#include "bba/BBA.h"
#include "bba/Parameter.h"
 
#include <Logger.h>
 
#include <TGeoManager.h>
#include <TGeoMatrix.h>
 
#include <cmath>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
 
namespace
{
  using namespace cbm::algo;
  constexpr double shiftToRotFactor = 10.;  // factor to convert shifts [cm] to rotations [grad] for boundary settings
  constexpr bool kMultipleScatteringInCostFunction =
    false;  // include multiple scattering in the cost function track refit
  constexpr bool kMultipleScatteringInTrackSelection =
    false;  // include multiple scattering in the track selection refit
}  // namespace
 
 
CbmBbaAlignTask::CbmBbaAlignTask(const char* name, Int_t iVerbose, TString histoFileName, size_t nMaxTracks)
  : FairTask(name, iVerbose)
  , fNmaxTracks(nMaxTracks)
  , fHistoFileName(histoFileName)
{
  TFile* curFile           = gFile;
  TDirectory* curDirectory = gDirectory;
 
  if (!(fHistoFileName == "")) {
    fHistoFile = new TFile(fHistoFileName.Data(), "RECREATE");
  }
  else {
    fHistoFile = gFile;
  }
 
  fHistoFile->cd();
  fHistoDir = fHistoFile->mkdir("Bba");
  fHistoDir->cd();
 
  gFile      = curFile;
  gDirectory = curDirectory;
}
 
 
InitStatus CbmBbaAlignTask::Init()
{
  std::cout << "CbmBbaAlignTask::Init()" << std::endl;
 
#ifdef HAVE_OMP
  {
    const char* env = std::getenv("OMP_NUM_THREADS");
    if (env) {
      fNthreads = TString(env).Atoi();
    }
    // ensure that at least one thread is set
    if (fNthreads < 1) {
      fNthreads = 1;
    }
    LOG(info) << "BBA: running with " << fNthreads << " threads";
  }
#else
  {
    fNthreads = 1;
    LOG(info) << "BBA: OpenMP not available, running with single thread";
  }
#endif
 
  // FairRootManager manages the root input/output
  FairRootManager* ioman = FairRootManager::Instance();
  if (!ioman) {
    LOG(fatal) << "CbmBbaAlignTask: no FairRootManager";
    return kFATAL;
  }
 
  // initialize mvd-sts track input
  fInputStsTracks = static_cast<TClonesArray*>(ioman->GetObject("StsTrack"));
  if (!fInputStsTracks) {
    LOG(error) << "CbmBbaAlignTask::Init: STS track array not found!";
    return kERROR;
  }
  else {
    LOG(info) << "CbmBbaAlignTask::Init: STS track array found.";
  }
 
  // initialize fitter
  fFitter.Init();
  fFitter.SetIgnorePoorCoordinates(true);
 
  //fFitter.SetMaxExtrapolationStep(70.); // less granular extrapolation for performance
  fFitter.SetIgnoreMultipleScattering(!kMultipleScatteringInTrackSelection);
 
  // initialize the number of tracking stations from the CbmStsTrackingInterface
  // (TODO we may need to use only active stations from the L1 setup)
 
  const auto* pStsIfs = cbm::RecoSetupManager::Instance()->GetSetup().GetSts();
  if (!pStsIfs) {
    LOG(fatal) << "CbmBbaAlignTask::Init: STS RecoSetupUnit not found! Initialize "
                  "CbmSetup::Instance()->LoadSetup(setup) before running the alignment task.";
    return kFATAL;
  }
  fNtrackingStations = pStsIfs->GetNofTrackingStations();
 
  LOG(info) << "CbmBbaAlignTask::Init: Number of STS tracking stations from interface: " << fNtrackingStations;
 
  return kSUCCESS;
}
 
 
void CbmBbaAlignTask::Exec(Option_t* /*unused*/)
{
  std::cout << "CbmBbaAlignTask::Exec()" << std::endl;
 
  LOG(info) << "BBA: exec event N " << fNevents;
  fNevents++;
 
  LOG(info) << "CbmBbaAlignTask::Exec: STS track array found with size: " << fInputStsTracks->GetEntries();
 
 
  // select tracks for alignment and store them
  // fFitter.SetDefaultMomentumForMs(0.1); // not needed for STS tracks with magnetic field
 
  SelectTracks(fInputStsTracks, fTracks);
}
 
void CbmBbaAlignTask::Finish()
{
  std::cout << "CbmBbaAlignTask::Finish()" << std::endl;
  LOG(info) << "BBA: start the global alignment procedure with " << fTracks.size() << " tracks ...";
 
  // fix the material radiation thickness to avoid the chi2 rattling at the material map bin boundaries
  fFitter.SetFixMaterial(true);
 
  // Initializations we can only do here, after the tracks are selected
  fSensors = SensorsFromTracks(fTracks);  // create the list of sensors from the tracks
  fAlignmentBodies =
    CreateAndSetAlignmentBodies(fAlignmentMode, fSensors);  // create the list of alignment bodies from the sensors
  SetReferences(fSensors, fTracks);                         // bind the sensors and alignment bodies to the track nodes
  InitializeHistograms();                                   // initialize histograms for the alignment
 
  LOG(info) << "BBA: " << fSensors.size() << " sensors, " << fAlignmentBodies.size() << " alignment bodies";
 
  LOG(debug) << "\n Sensors: ";
  for (unsigned int is = 0; is < fSensors.size(); is++) {
    auto& s = fSensors[is];
    LOG(debug) << "Sensor " << is << " sys " << s.fSystemId << " sta " << s.fTrackingStation << " body "
               << s.fAlignmentBody;
  }
  LOG(info) << "\n Alignment bodies: ";
  for (unsigned int ib = 0; ib < fAlignmentBodies.size(); ib++) {
    auto& b = fAlignmentBodies[ib];
    LOG(info) << "Body: " << ib << " station: " << b.GetTrackingStation();
  }
 
  // "load" the alignment parameters
  std::vector<bba::Parameter> par = CreateAlignmentParameters();
 
  // create artificial misalignement
  if (fSimulatedMisalignmentRange <= 0.) {
    LOG(info) << "no simulated misalignment";
    fParMC = std::vector<double>(kNdfBody * fAlignmentBodies.size(), 0.);
  }
  else {
    LOG(info) << "Simulated misalignment range: " << fSimulatedMisalignmentRange;
    fParMC = SimulateMisalignment(par, fTracks, fAlignmentBodies);  // apply the misalignment to the hits
  }
 
  // create bba object and set parameters
  bba::BBA alignment;
  alignment.setParameters(par);
  alignment.setChi2([this](const std::vector<double>& p) { return CostFunction(p); });
  alignment.printInfo();
 
  // extract the (initial) start parameters from bba to verify the fitter/cost function and calculate the initial cost
  fParStart = alignment.getResult();
 
  {
    // calculate the initial cost function
    // (before alignment)
    LOG(debug) << "calculate initial cost function first time ...";
    auto tracksCopy = fTracks;            // store the initial (misaligned) tracks
    CostFunction(fParStart, tracksCopy);  // run the fitter
 
    // reject bad tracks (caused by artificial misalignement) and recalculate the initial cost function
    unsigned tracks_rejected = 0;
    for (unsigned int itr = 0; itr < tracksCopy.size(); itr++) {
      auto& t     = tracksCopy[itr];
      double ndf  = t.fFittedTrackParam.GetNdf();
      double chi2 = t.fFittedTrackParam.GetChiSq();
      if (ndf < 0. || chi2 < 0. || !std::isfinite(ndf) || !std::isfinite(chi2)) {
        fTracks[itr].fIsActive = false;  // we deactivat the original track
        tracks_rejected++;
      }
    }
    LOG(info) << "reject " << tracks_rejected << " bad tracks and recalculate the initial cost function...";
  }
  {
    auto initialTracks = fTracks;                           // store the remaining initial (misaligned) tracks
    fCostInitial = CostFunction(fParStart, initialTracks);  // recalculate the cost function with the remaining tracks
    // plot the residuals before alignment (after applying the first guessed start parameters)
    FillHistograms(initialTracks, fHistoBeforeAlignment);
  }
 
  fFixedNdf = -1;  //fNdfTotal;  TODO: why here?
 
  LOG(debug) << "calculate ideal cost function second time...";
  fCostIdeal = CostFunction(fParMC);  // we have to do this after rejecting bad tracks
 
  // run the actual alignment ////////////////////
  LOG(info) << "BBA: start the alignment procedure ...";
  LOG(info) << " true (MC) parameters: " << PrintParameters(fParMC);
  LOG(info) << " initial parameters: " << PrintParameters(fParStart);
  LOG(info) << " Cost function for the true parameters: " << fCostIdeal;
  LOG(info) << " Cost function for the initial parameters: " << fCostInitial;
  alignment.align();  // loops until convergence
 
  auto alignedTracks = fTracks;
  double costResult =
    CostFunction(alignment.getResult(),
                 alignedTracks);  // apply the alignment parameters to the hits and calculate the cost function
 
  LOG(info) << "BBA: alignment procedure finished ...";
 
 
  LOG(info) << "BBA: alignment results: ";
  LOG(info) << " aligned parameters: " << PrintParameters(alignment.getResult());
  LOG(info) << " difference to MC parameters: " << PrintParameters(DiffParameters(alignment.getResult(), fParMC));
  LOG(info) << " RMSE to MC parameters: " << PrintRMSE(par, fParMC, alignment.getResult());
  LOG(info) << " Number of cost function calls: " << fNcallsCost;
  LOG(info) << " Cost function for the initial parameters: " << fCostInitial;
  LOG(info) << " Cost function for the true parameters: " << fCostIdeal;
  LOG(info) << " Cost function for the final parameters: " << costResult
            << ", diff to true cost: " << costResult - fCostIdeal
            << ", diff to initial cost: " << costResult - fCostInitial;
 
  // fill histograms after alignment
  FillHistograms(alignedTracks, fHistoAfterAlignment);
 
  // TODO: we may need to run the cost function again several times, at least this is done in the original task
 
 
  // prepare the output alignment parameters
  LOG(info) << "BBA: creating alignment matrices for the output ...";
  std::map<std::string, TGeoHMatrix> alignmentMatrices = CreateAlignmentMatrices(alignment.getResult());
 
  // write the alignment matrices to the output file
  TFile* misalignmentMatrixRootfile = new TFile(fMatrixOutFileName, "RECREATE");
  if (misalignmentMatrixRootfile->IsOpen()) {
    gDirectory->WriteObject(&alignmentMatrices, "AlignmentMatrices");
    misalignmentMatrixRootfile->Write();
    misalignmentMatrixRootfile->Close();
  }
 
  // write histograms to file
  TDirectory* curr   = gDirectory;
  TFile* currentFile = gFile;
  // Open output file and write histograms
 
  fHistoFile->cd();
  WriteObjectsCurFile(fHistoDir);
  if (!(fHistoFileName == "")) {
    fHistoFile->Close();
    fHistoFile->Delete();
  }
  gFile      = currentFile;
  gDirectory = curr;
 
}  // CbmBbaAlignTask::Finish()
 
 
unsigned CbmBbaAlignTask::SelectTracks(TClonesArray* inputTracks, std::vector<TrackContainer>& Tracks)
{
  // debug statistics
  unsigned num_accepted               = 0;
  unsigned num_rejected_create_failed = 0;
  unsigned num_rejected_fit           = 0;
  unsigned num_rejected_hitNumber     = 0;
  unsigned num_rejected_infiniteQp    = 0;
  unsigned num_rejected_momentum      = 0;
  unsigned num_fix_rad_max            = 0;
  unsigned num_fix_rad_min            = 0;
  unsigned num_rejected_fit_node      = 0;
  unsigned num_rejected_traj          = 0;
  int numPosMomentum                  = 0;
  int numNegMomentum                  = 0;
 
  fFitter.SetFixMaterial(false);
 
  // loop over all STS tracks and ...
  for (int iTr = 0; iTr < inputTracks->GetEntriesFast(); iTr++) {
    if (Tracks.size() >= fNmaxTracks) break;
    //const CbmStsTrack* stsTrack = dynamic_cast<CbmStsTrack*>(inputTracks->At(iTr));
    //assert(stsTrack);
    // ... create a track container for the track
    TrackContainer t;  // TODO: refactor track container to class with setters
    if (!fFitter.CreateFromMvdStsTrack(t.fTrack, iTr, true)) {
      num_rejected_create_failed++;
      continue;
    }
    t.MakeConsistent();
 
    // ... fit the track to initialize the track parameters
    if (!fFitter.FitTrajectory(t.fTrack)) {
      num_rejected_fit++;
      continue;
    }
 
    if (!fFitter.FitTrajectoryUpstream(t.fTrack, t.fFittedTrackParam)) {
      num_rejected_fit++;
      continue;
    }
 
    const auto& par = t.fFittedTrackParam;
 
    if (t.fNstsHits < fNtrackingStations) {
      num_rejected_hitNumber++;
      continue;  // cut: all tracking stations must have a hit
    }
    if (!std::isfinite(par.GetQp())) {
      num_rejected_infiniteQp++;
      continue;  // cut: Q/p must be finite
    }
    if (fabs(par.GetQp()) > 3) {
      num_rejected_momentum++;
      continue;  // cut: momentum > 1 GeV
    }
 
    // correct the material radiation thickness for a case of too large/small values
    for (unsigned int i = 0; i < t.fTrack.GetNofNodes(); i++) {
      auto& n = t.fTrack.GetNodeReference(i);
      if (!n.isFitted) {
        t.fIsActive = false;
        num_rejected_fit_node++;
        break;
      }
      if (n.radThick > 0.01) {
        LOG(debug) << "CbmBbaAlignTask: radiation thickness is too large: " << n.radThick;
        num_fix_rad_max++;
        n.radThick = 0.01;
      }
      if (n.radThick < 0.001) {
        LOG(debug) << "CbmBbaAlignTask: radiation thickness is too small: " << n.radThick;
        num_fix_rad_min++;
        n.radThick = 0.001;
      }
    }
    if (!t.fIsActive) continue;
 
 
    // ensure that all the hits are not too much deviated from the trajectory
    // (to be improved)
    const double cutX = 8.;
    const double cutY = 8.;
    for (unsigned int i = 0; i < t.fTrack.GetNofNodes(); i++) {
      // copy the track to avoid modifying the original one
      auto t_copy = t.fTrack;
      auto& n     = t_copy.GetNode(i);
      if (!n.isXySet) {
        continue;
      }
      t_copy.DisableMeasurementAtNode(i);
      if (!fFitter.FitTrajectory(t_copy)) {
        t.fIsActive = false;
        num_rejected_traj++;
        break;
      }
      double x_residual = n.measurementXY.X() - n.paramDn.X();  // this should be the difference vector
      double y_residual = n.measurementXY.Y() - n.paramDn.Y();
      double x_pull     = x_residual / sqrt(n.measurementXY.Dx2() + n.paramDn.GetCovariance(0, 0));
      double y_pull     = y_residual / sqrt(n.measurementXY.Dy2() + n.paramDn.GetCovariance(1, 1));
      if (!n.isFitted || (n.measurementXY.NdfX() > 0. && fabs(x_pull) > cutX)
          || (n.measurementXY.NdfY() > 0. && fabs(y_pull) > cutY)) {
        t.fIsActive = false;
        num_rejected_traj++;
        break;
      }
    }
 
    // copy the accepted track to track vector
    if (t.fIsActive) {
      Tracks.push_back(t);
      Tracks.back().StoreOriginalHitPositions();
      if (par.GetQp() > 0.) {
        numPosMomentum++;
      }
      else {
        numNegMomentum++;
      }
      num_accepted++;
    }
  }
 
  // Debug output
  if (num_rejected_create_failed != 0) {
    LOG(warn) << "failed to create " << num_rejected_create_failed << " tracks";
  }
  if (num_rejected_fit != 0) {
    LOG(warn) << "failed to fit " << num_rejected_fit << " tracks";
  }
  if (num_rejected_hitNumber != 0) {
    LOG(info) << "rejected " << num_rejected_hitNumber << " tracks for insufficient number of hits";
  }
  if (num_rejected_infiniteQp != 0) {
    LOG(warn) << "rejected " << num_rejected_infiniteQp << " tracks for infinite Q/p";
  }
  if (num_rejected_momentum != 0) {
    LOG(info) << "rejected " << num_rejected_momentum << " tracks for low momentum";
  }
  if (num_fix_rad_max != 0) {
    LOG(warn) << "CbmBbaAlignTask: radiation thickness is too large " << num_fix_rad_max << " times";
  }
  if (num_fix_rad_min != 0) {
    LOG(warn) << "CbmBbaAlignTask: radiation thickness is too small " << num_fix_rad_min << " times";
  }
  if (num_rejected_fit_node != 0) {
    LOG(warn) << "Rejected fit node  " << num_rejected_fit_node << " tracks";
  }
  if (num_rejected_traj != 0) {
    LOG(warn) << "Rejected  " << num_rejected_traj << " tracks for hits deviating from the trajectory";
  }
 
  LOG(info) << "BBA: selected " << num_accepted << " (" << Tracks.size() << " total) tracks for alignment, "
            << numPosMomentum << " positive and " << numNegMomentum << " negative momentum tracks";
 
 
  return num_accepted;
}
 
std::vector<CbmBbaAlignTask::Sensor> CbmBbaAlignTask::SensorsFromTracks(const std::vector<TrackContainer>& tracks) const
{
  // create the set of sensors and pack into sorted vector
  // loop over all hits and extract the sensors from the hits
  std::set<Sensor> sensorSet;
  std::shared_ptr<const cbm::GeoNodeMap> geoNodeMap = cbm::RecoSetupManager::Instance()->GetGeoNodeMap();
  for (auto& t : tracks) {
    for (auto& n : t.fTrack.GetNodes()) {
      if (!n.isXySet) {
        continue;
      }
      Sensor s;
      s.fSystemId = Trajectory::GetHitSystemId(n);
      // TODO: get the station index from n.fHitSystemId, n.fHitAddress
      s.fTrackingStation = n.materialLayer;
      if (n.materialLayer < 0) {
        LOG(fatal) << "CbmBbaAlignTask::SensorsFromTracks: negative material layer in track node! for "
                   << Trajectory::GetHitAddress(n) << "in " << Trajectory::GetHitSystemId(n);
      }
      if (s.fSystemId == ECbmModuleId::kSts) {
        s.fNodePath = geoNodeMap->GetNodePath(Trajectory::GetHitAddress(n));  // NOTE: for STS gives a path to a sensor
      }
      sensorSet.insert(s);
    }
  }
 
  // pack sensors into a vector and sort them
  std::vector<Sensor> sensors(sensorSet.begin(), sensorSet.end());
  // TODO: isn't the set sorted correctly already?
  // NOTE: the sensors are automatically sorted inside std::set, no sorting is needed here (SZh)
  std::sort(sensors.begin(), sensors.end(), [](const Sensor& a, const Sensor& b) { return a < b; });
 
  LOG(info) << "BBA: found " << sensors.size() << " sensors used by selected tracks:";
 
  return sensors;
}
 
 
std::vector<cbm::bba::AlignmentBody> CbmBbaAlignTask::CreateAndSetAlignmentBodies(CbmBbaAlignTask::AlignmentMode mode,
                                                                                  std::vector<Sensor>& sensors) const
{
  std::vector<cbm::bba::AlignmentBody> alignmentBodies;
  // bind alignment bodies to sensors or stations
  switch (mode) {
    case station:  // one alignment body per tracking station
      //alignmentBodies.resize(fNtrackingStations);
      for (int i = 0; i < fNtrackingStations; i++) {
        //alignmentBodies[i].SetTrackingStation(i);
        double z = 0;
        switch (i) {
          case 0: z = -14.; break;
          case 1: z = -4.; break;
          case 2: z = 6.; break;
          case 3: z = 16.; break;
          case 4: z = 26.; break;
          case 5: z = 36.; break;
          case 6: z = 46.; break;
          case 7: z = 56.; break;
          default: z = 0.; break;
        }
        alignmentBodies.emplace_back(cbm::bba::AlignmentBody(0., 0., z, 0., 0., 0.));
        alignmentBodies.back().SetTrackingStation(i);
      }
      for (auto& s : sensors) {
        s.fAlignmentBody = s.fTrackingStation;
      }
      break;
    case sensor:  // one alignment body per sensor
      alignmentBodies.clear();
      alignmentBodies.reserve(sensors.size());
 
      for (unsigned int i = 0; i < sensors.size(); i++) {
        auto& s = sensors[i];
        alignmentBodies.emplace_back(cbm::bba::AlignmentBody(s.fNodePath));
        s.fAlignmentBody = i;
        alignmentBodies.back().SetTrackingStation(s.fTrackingStation);
      }
      break;
  }
 
  return alignmentBodies;
}
 
 
void CbmBbaAlignTask::SetReferences(const std::vector<Sensor>& sensors, std::vector<TrackContainer>& tracks) const
{
  // bind the sensors and alignment bodies to the track nodes
  std::shared_ptr<const cbm::GeoNodeMap> geoNodeMap = cbm::RecoSetupManager::Instance()->GetGeoNodeMap();
  for (auto& t : tracks) {
    for (unsigned int i = 0; i < t.fTrack.GetNofNodes(); i++) {
      const auto& n = t.fTrack.GetNode(i);
      if (!n.isXySet) {
        continue;
      }
      // create the sensor we want to find
      Sensor s;
      s.fSystemId = Trajectory::GetHitSystemId(n);
      if (n.materialLayer < 0) {
        LOG(fatal) << "CbmBbaAlignTask::SetReferences: negative material layer in track node! for "
                   << Trajectory::GetHitAddress(n) << "in " << Trajectory::GetHitSystemId(n);
      }
      s.fTrackingStation = n.materialLayer;
      if (s.fSystemId == ECbmModuleId::kSts) {
        s.fNodePath = geoNodeMap->GetNodePath(Trajectory::GetHitAddress(n));
      }
      // find the sensor in the vector and set the reference to sensor index
      auto iter =
        std::lower_bound(sensors.begin(), sensors.end(), s, [](const Sensor& a, const Sensor& b) { return a < b; });
      assert(iter != sensors.end());
      assert(*iter == s);
      int iSensor = std::distance(sensors.begin(), iter);
      t.fTrack.SetSensorId(i, iSensor);  // store the sensor index in the node
      t.fTrack.SetAlignmentBodyId(i,
                                  sensors.at(iSensor).fAlignmentBody);  // store the alignment body index in the node
    }
  }
}
 
std::vector<bba::Parameter> CbmBbaAlignTask::CreateAlignmentParameters() const
{
  int nParameters = kNdfBody * fAlignmentBodies.size();  // x, y, z
  std::vector<bba::Parameter> par(nParameters);
 
  // initialize all parameters
  for (int ip = 0; ip < nParameters; ip++) {
    bba::Parameter& p = par[ip];
    p.SetID(ip);  // ID requires newer BBA version
    p.SetActive(1);
    p.SetValue(0.);
    if (ip % kNdfBody < 3) {
      p.SetBoundaryMin(-fSimulatedMisalignmentRange);
      p.SetBoundaryMax(fSimulatedMisalignmentRange);
    }
    else {
      p.SetBoundaryMin(-fSimulatedMisalignmentRange * shiftToRotFactor);  // in degrees
      p.SetBoundaryMax(fSimulatedMisalignmentRange * shiftToRotFactor);
    }
    if (ip % kNdfBody < 3) {
      p.SetStepMin(1.e-6);
      p.SetStepMax(0.1);
      p.SetStep(1.e-2);
    }
    else {
      p.SetStepMin(1.e-5);
      p.SetStepMax(0.5);
      p.SetStep(1.e-1);
    }
  }
 
  // fix first and last station and x of middle station
  for (unsigned int ib = 0; ib < fAlignmentBodies.size(); ib++) {
    auto& b = fAlignmentBodies[ib];
    if (b.GetTrackingStation() == 0) {  // the first station
      par[kNdfBody * ib + 0].SetActive(0);
      par[kNdfBody * ib + 1].SetActive(0);
      par[kNdfBody * ib + 2].SetActive(0);
      par[kNdfBody * ib + 3].SetActive(0);
      par[kNdfBody * ib + 4].SetActive(0);
      par[kNdfBody * ib + 5].SetActive(0);
      continue;
    }
    if (b.GetTrackingStation() == fNtrackingStations - 1) {  // the last station
      par[kNdfBody * ib + 0].SetActive(0);
      par[kNdfBody * ib + 1].SetActive(0);
      par[kNdfBody * ib + 2].SetActive(0);
      par[kNdfBody * ib + 3].SetActive(0);
      par[kNdfBody * ib + 4].SetActive(0);
      par[kNdfBody * ib + 5].SetActive(0);
      continue;
    }
    if (b.GetTrackingStation() == 4) {
      par[kNdfBody * ib + 0].SetActive(0);
      par[kNdfBody * ib + 1].SetActive(0);
      par[kNdfBody * ib + 2].SetActive(0);
      par[kNdfBody * ib + 3].SetActive(0);
      par[kNdfBody * ib + 4].SetActive(0);
      par[kNdfBody * ib + 5].SetActive(0);
      continue;
    }
    // par[kNdfBody * ib + 5].SetActive(0);  // fix gamma for all stations
  }
  return par;
}
 
std::vector<double> CbmBbaAlignTask::InvertParameters(std::vector<double> par) const
{
  for (unsigned int i = 0; i < par.size(); i++) {
    par[i] = -par[i];
  }
  return par;
}
 
std::vector<double> CbmBbaAlignTask::DiffParameters(const std::vector<double>& par1,
                                                    const std::vector<double>& par2) const
{
  if (par1.size() != par2.size()) {
    LOG(fatal) << "CbmBbaAlignTask::DiffParameters: parameter vectors have different sizes!";
  }
  std::vector<double> diff(par1.size(), 0.);
  for (unsigned int i = 0; i < par1.size(); i++) {
    diff[i] = par1[i] - par2[i];
  }
  return diff;
}
 
std::string CbmBbaAlignTask::PrintParameters(const std::vector<double>& par) const
{
  std::stringstream ss;
  ss << "BBA: alignment parameters:";
  for (unsigned int is = 0; is < fAlignmentBodies.size(); is++) {
    auto& b = fAlignmentBodies[is];
    ss << "\nBody " << is << " sta " << b.GetTrackingStation() << ": x " << std::setprecision(10)
       << par[kNdfBody * is + 0] << " y " << par[kNdfBody * is + 1] << " z " << par[kNdfBody * is + 2]
       << " alpha: " << par[kNdfBody * is + 3] << " beta: " << par[kNdfBody * is + 4]
       << " gamma: " << par[kNdfBody * is + 5];
  }
  return ss.str();
}
 
std::string CbmBbaAlignTask::PrintRMSE(const std::vector<bba::Parameter>& par, const std::vector<double>& parMC,
                                       const std::vector<double>& parValues) const
{
  std::vector<double> rmse(kNdfBody, 0.);  // RMSE for each parameter species (x, y, z)
  std::vector<unsigned int> active(kNdfBody, 0);
  for (unsigned int is = 0; is < fAlignmentBodies.size(); is++) {
    unsigned int indexX = kNdfBody * is + 0;
    unsigned int indexY = kNdfBody * is + 1;
    unsigned int indexZ = kNdfBody * is + 2;
    unsigned int indexA = kNdfBody * is + 3;
    unsigned int indexB = kNdfBody * is + 4;
    unsigned int indexC = kNdfBody * is + 5;
    if (par[indexX].IsActive()) {
      double dx = parValues[indexX] - parMC[indexX];
      rmse[0] += dx * dx;
      active[0]++;
    }
    if (par[indexY].IsActive()) {
      double dy = parValues[indexY] - parMC[indexY];
      rmse[1] += dy * dy;
      active[1]++;
    }
    if (par[indexZ].IsActive()) {
      double dz = parValues[indexZ] - parMC[indexZ];
      rmse[2] += dz * dz;
      active[2]++;
    }
    if (par[indexA].IsActive()) {
      double da = parValues[indexA] - parMC[indexA];
      rmse[3] += da * da;
      active[3]++;
    }
    if (par[indexB].IsActive()) {
      double db = parValues[indexB] - parMC[indexB];
      rmse[4] += db * db;
      active[4]++;
    }
    if (par[indexC].IsActive()) {
      double dc = parValues[indexC] - parMC[indexC];
      rmse[5] += dc * dc;
      active[5]++;
    }
  }
  for (unsigned int i = 0; i < kNdfBody; i++) {
    rmse[i] = sqrt(rmse[i] / active[i]);
  }
  std::stringstream ss;
  ss << "BBA: RMSE of alignment parameters:";
  for (unsigned int i = 0; i < kNdfBody; i++) {
    const char* label = "xyzabc";
    ss << "\n" << label[i] << ": " << std::setprecision(10) << rmse[i] << " (" << active[i] << " active)";
  }
  return ss.str();
}
 
std::vector<double> CbmBbaAlignTask::SimulateMisalignment(const std::vector<bba::Parameter>& par,
                                                          std::vector<TrackContainer>& tracks,
                                                          std::vector<cbm::bba::AlignmentBody>& AlignmentBodies) const
{
  gRandom->SetSeed(fRandomSeed);
  double d    = fSimulatedMisalignmentRange;
  double dRot = d * shiftToRotFactor;
  std::vector<double> par_misaligned(kNdfBody * AlignmentBodies.size(), 0.);
 
  for (unsigned int is = 0; is < AlignmentBodies.size(); is++) {
    const bba::Parameter& px = par[kNdfBody * is + 0];
    const bba::Parameter& py = par[kNdfBody * is + 1];
    const bba::Parameter& pz = par[kNdfBody * is + 2];
    const bba::Parameter& pa = par[kNdfBody * is + 3];
    const bba::Parameter& pb = par[kNdfBody * is + 4];
    const bba::Parameter& pc = par[kNdfBody * is + 5];
 
    // +- d cm misalignment
    if (px.IsActive()) {
      par_misaligned.at(kNdfBody * is + 0) = gRandom->Uniform(2. * d) - d;
    }
    if (py.IsActive()) {
      par_misaligned.at(kNdfBody * is + 1) = gRandom->Uniform(2. * d) - d;
    }
    if (pz.IsActive()) {
      par_misaligned.at(kNdfBody * is + 2) = gRandom->Uniform(2. * d) - d;
    }
    if (pa.IsActive()) {
      par_misaligned.at(kNdfBody * is + 3) = gRandom->Uniform(2. * (dRot)) - (dRot);
    }
    if (pb.IsActive()) {
      par_misaligned.at(kNdfBody * is + 4) = gRandom->Uniform(2. * (dRot)) - (dRot);
    }
    if (pc.IsActive()) {
      par_misaligned.at(kNdfBody * is + 5) = gRandom->Uniform(2. * (dRot)) - (dRot);
    }
  }
  //
  ApplyAlignment(par_misaligned, tracks, AlignmentBodies, true);
  // Store the hit positions after applying the misalignment as original hit positions:
  for (auto& t : tracks) {
    t.StoreOriginalHitPositions();
  }
  return par_misaligned;
}
 
void CbmBbaAlignTask::InitializeHistograms()
{
  TDirectory* curDirectory = gDirectory;
  int nHistos              = fAlignmentBodies.size();
 
  fHistoDir->cd();
  for (int i = -1; i < nHistos; i++) {
    const char* histName  = Form("Station%d", i);
    const char* histTitle = Form(", Station %d", i);
    if (i == -1) {
      histName  = "All";
      histTitle = ", all Stations";
    }
    //fHistoDir->mkdir(n1);
    //fHistoDir->cd(n1);
 
    int nBins     = 301;
    double sizeRX = 0.1;
    double sizeRY = 0.1;
    double sizePX = 10.;
    double sizePY = 10.;
 
    fHistoBeforeAlignment.fResidualsX.push_back(std::make_unique<TH1F>(
      Form("%s_X_Res_BeforeAli", histName), Form("X-Residuals Before Alignment%s", histTitle), nBins, -sizeRX, sizeRX));
    fHistoBeforeAlignment.fResidualsY.push_back(std::make_unique<TH1F>(
      Form("%s_Y_Res_BeforeAli", histName), Form("Y-Residuals Before Alignment%s", histTitle), nBins, -sizeRY, sizeRY));
 
    fHistoBeforeAlignment.fPullsX.push_back(std::make_unique<TH1F>(
      Form("%s_X_Pull_BeforeAli", histName), Form("X-Pulls Before Alignment%s", histTitle), nBins, -sizePX, sizePX));
    fHistoBeforeAlignment.fPullsY.push_back(std::make_unique<TH1F>(
      Form("%s_Y_Pull_BeforeAli", histName), Form("Y-Pulls Before Alignment%s", histTitle), nBins, -sizePY, sizePY));
 
 
    fHistoAfterAlignment.fResidualsX.push_back(std::make_unique<TH1F>(
      Form("%s_X_Res_AfterAli", histName), Form("X-Residuals After Alignment%s", histTitle), nBins, -sizeRX, sizeRX));
    fHistoAfterAlignment.fResidualsY.push_back(std::make_unique<TH1F>(
      Form("%s_Y_Res_AfterAli", histName), Form("Y-Residuals After Alignment%s", histTitle), nBins, -sizeRY, sizeRY));
 
    fHistoAfterAlignment.fPullsX.push_back(std::make_unique<TH1F>(
      Form("%s_X_Pull_AfterAli", histName), Form("X-Pulls After Alignment%s", histTitle), nBins, -sizePX, sizePX));
    fHistoAfterAlignment.fPullsY.push_back(std::make_unique<TH1F>(
      Form("%s_Y_Pull_AfterAli", histName), Form("Y-Pulls After Alignment%s", histTitle), nBins, -sizePY, sizePY));
  }
 
  for (int i = 0; i < nHistos + 1; i++) {
    // set line colors
    fHistoAfterAlignment.fResidualsX[i]->SetLineColor(kRed);
    fHistoAfterAlignment.fResidualsY[i]->SetLineColor(kRed);
    fHistoAfterAlignment.fPullsX[i]->SetLineColor(kRed);
    fHistoAfterAlignment.fPullsY[i]->SetLineColor(kRed);
 
    // set histograms to dynamic binning
    fHistoBeforeAlignment.fResidualsX[i]->SetCanExtend(TH1::kXaxis);
    fHistoBeforeAlignment.fResidualsY[i]->SetCanExtend(TH1::kXaxis);
    fHistoBeforeAlignment.fPullsX[i]->SetCanExtend(TH1::kXaxis);
    fHistoBeforeAlignment.fPullsY[i]->SetCanExtend(TH1::kXaxis);
 
    fHistoAfterAlignment.fResidualsX[i]->SetCanExtend(TH1::kXaxis);
    fHistoAfterAlignment.fResidualsY[i]->SetCanExtend(TH1::kXaxis);
    fHistoAfterAlignment.fPullsX[i]->SetCanExtend(TH1::kXaxis);
    fHistoAfterAlignment.fPullsY[i]->SetCanExtend(TH1::kXaxis);
  }
 
  gDirectory = curDirectory;
}
 
void CbmBbaAlignTask::WriteObjectsCurFile(TObject* obj)
{
  if (!obj->IsFolder())
    obj->Write();
  else {
    TDirectory* cur    = gDirectory;
    TFile* currentFile = gFile;
 
    TDirectory* sub = cur->GetDirectory(obj->GetName());
    sub->cd();
    TList* listSub = (dynamic_cast<TDirectory*>(obj))->GetList();
    TIter it(listSub);
    while (TObject* obj1 = it()) {
      WriteObjectsCurFile(obj1);
    }
    cur->cd();
    gFile      = currentFile;
    gDirectory = cur;
  }
}
 
void CbmBbaAlignTask::FillHistograms(const std::vector<TrackContainer>& tracks, Histograms& histo)
{
  // calculate the residuals for all tracks at each TrajectoryNode
  for (const auto& t : tracks) {
    if (!t.fIsActive) continue;
    for (unsigned int in = 0; in < t.fTrack.GetNofNodes(); in++) {
      // copy the track to avoid modifying the original one
      auto tr          = t.fTrack;
      const auto& node = tr.GetNode(in);
      if (!node.isXySet) {
        continue;
      }
      tr.DisableMeasurementAtNode(in);
 
      fFitter.FitTrajectory(tr);
 
      if (!node.isFitted) {  // TODO understand why this is needed
        continue;
      }
 
      int alignmentBody = Trajectory::GetAlignmentBodyId(node);
 
      double x_residual = node.measurementXY.X() - node.paramDn.X();
      double y_residual = node.measurementXY.Y() - node.paramDn.Y();
      double x_pull     = x_residual / sqrt(node.measurementXY.Dx2() + node.paramDn.GetCovariance(0, 0));
      double y_pull     = y_residual / sqrt(node.measurementXY.Dy2() + node.paramDn.GetCovariance(1, 1));
 
      if (node.measurementXY.NdfX() > 0) {
        histo.fResidualsX[0]->Fill(x_residual);
        histo.fPullsX[0]->Fill(x_pull);
      }
      if (node.measurementXY.NdfY() > 0) {
        histo.fResidualsY[0]->Fill(y_residual);
        histo.fPullsY[0]->Fill(y_pull);
      }
 
      // TODO: why can we drop the check for NdfX/Y > 0 here?
      histo.fResidualsX[alignmentBody + 1]->Fill(x_residual);
      histo.fResidualsY[alignmentBody + 1]->Fill(y_residual);
 
      histo.fPullsX[alignmentBody + 1]->Fill(x_pull);
      histo.fPullsY[alignmentBody + 1]->Fill(y_pull);
    }
  }
}
 
void CbmBbaAlignTask::SetAlignment(const std::vector<double>& par,
                                   std::vector<cbm::bba::AlignmentBody>& AlignmentBodies) const
{
  // apply alignment parameters to the alignment bodies
  for (unsigned int ib = 0; ib < AlignmentBodies.size(); ib++) {
    double dx = par[kNdfBody * ib + 0];
    double dy = par[kNdfBody * ib + 1];
    double dz = par[kNdfBody * ib + 2];
    double da = par[kNdfBody * ib + 3];
    double db = par[kNdfBody * ib + 4];
    double dc = par[kNdfBody * ib + 5];
    AlignmentBodies[ib].SetParameters(dx, dy, dz, da, db, dc);
  }
}
 
void CbmBbaAlignTask::ApplyAlignment(const std::vector<double>& par, std::vector<TrackContainer>& tracks,
                                     std::vector<cbm::bba::AlignmentBody>& AlignmentBodies, bool invert = false) const
{
  SetAlignment(par, AlignmentBodies);
  // apply alignment parameters to the hits
  for (TrackContainer& t : tracks) {
    for (unsigned int in = 0; in < t.fTrack.GetNofNodes(); in++) {
      auto& node = t.fTrack.GetNodeReference(in);
      if (!node.isXySet) {
        continue;
      }
 
      int AlignmentBodyIdx = Trajectory::GetAlignmentBodyId(node);
      if (AlignmentBodyIdx < 0) {
        LOG(fatal) << "BBA: ApplyAlignment: node has no alignment body! ";
        continue;
      }
      std::array<double, 3> alignedHit;
      if (!invert) {
        alignedHit = AlignmentBodies[AlignmentBodyIdx].ApplyAlignmentToHit(t.fOriginalHitPositions[in]);
      }
      else {
        alignedHit = AlignmentBodies[AlignmentBodyIdx].ApplyInverseAlignmentToHit(t.fOriginalHitPositions[in]);
      }
 
      // store the shift in z to shift the track parameters accordingly
      double dz = alignedHit[2] - node.z;
      node.measurementXY.SetX(alignedHit[0]);
      node.measurementXY.SetY(alignedHit[1]);
      node.z = alignedHit[2];
 
      // shift the fitted track to the Z position of the hit
      {
        auto& p = node.paramDn;
        p.SetX(p.X() + p.Tx() * dz);
        p.SetY(p.Y() + p.Ty() * dz);
        p.SetZ(node.z);
      }
      {
        auto& p = node.paramUp;
        p.SetX(p.X() + p.Tx() * dz);
        p.SetY(p.Y() + p.Ty() * dz);
        p.SetZ(node.z);
      }
    }
  }
}
 
double CbmBbaAlignTask::CostFunction(const std::vector<double>& par)
{
  auto tracksCopy = fTracks;
  return CostFunction(par, tracksCopy);
}
 
double CbmBbaAlignTask::CostFunction(const std::vector<double>& par, std::vector<TrackContainer>& tracks)
{
  // apply new parameters to the hits of a copy of tracks
  ApplyAlignment(par, tracks, fAlignmentBodies);
  double chi2Total = 0.;
  long ndfTotal    = 0;
 
  CbmKfTrackFitter fit(fFitter);
  fit.SetIgnoreMultipleScattering(!kMultipleScatteringInCostFunction);
 
#pragma omp parallel for firstprivate(fit) reduction(+ : chi2Total, ndfTotal) num_threads(fNthreads)
  for (unsigned int itr = 0; itr < tracks.size(); itr++) {
    auto& t = tracks[itr];
    if (!t.fIsActive) continue;
 
    // fit.SetDebugInfo(" track " + std::to_string(itr));
 
    for (unsigned int in = 0; in < t.fTrack.GetNofNodes(); in++) {
      const auto& node = t.fTrack.GetNode(in);
      if (!node.isFitted) {
        LOG(fatal) << "BBA: Cost function: aligned node is not fitted! ";
      }
    }
    bool ok     = fit.FitTrajectoryUpstream(t.fTrack, t.fFittedTrackParam);
    double chi2 = t.fFittedTrackParam.GetChiSq();
    double ndf  = t.fFittedTrackParam.GetNdf();
    if (!ok || !std::isfinite(chi2) || chi2 <= 0. || ndf <= 0.) {
      // TODO: deactivate the track and continue
      LOG(fatal) << "BBA: fit failed! ";
      chi2 = 0.;
      ndf  = 0.;
    }
    chi2Total += chi2;
    ndfTotal += (int) ndf;
  }
 
  double cost = chi2Total / (ndfTotal + 1);
  if (fFixedNdf > 0 && ndfTotal != fFixedNdf) {
    cost = 1.e30;
  }
 
  // monitoring
  if (++fNcallsCost % 100 == 0) {
    LOG(info) << "Cost function called " << fNcallsCost << " times, current cost: " << cost
              << ", diff to ideal cost: " << cost - fCostIdeal << ", ndf: " << ndfTotal << ", chi2: " << chi2Total;
    LOG(info) << PrintParameters(DiffParameters(par, fParMC));
  }
 
  return cost;
}
 
// This function creates an alignment node with the given path and alignment parameters.
// It takes the global alignment parameters
std::pair<std::string, TGeoHMatrix> CbmBbaAlignTask::CreateAlignmentNode(std::string path, double shiftX, double shiftY,
                                                                         double shiftZ, double rotX, double rotY,
                                                                         double rotZ) const
{
  if (path.empty()) {
    LOG(fatal) << "Alignment node path is empty!";
  }
  if (path[0] != '/') {
    // LOG(error)
    //   << "Alignment node path provided by the Cbm***TrackingInterface is not global. It must start with \"/\" " << path[0] << " in the path ";
  }
 
  // alignment matrix in global coordinates
  TGeoHMatrix alignmentMatrix;
  alignmentMatrix.SetDx(shiftX);
  alignmentMatrix.SetDy(shiftY);
  alignmentMatrix.SetDz(shiftZ);
  alignmentMatrix.RotateX(rotX);
  alignmentMatrix.RotateY(rotY);
  alignmentMatrix.RotateZ(rotZ);
 
  // Get global transformation matrix (position and rotation) of the TGeoNode at the given path
  if (!gGeoManager->cd(path.c_str())) {
    LOG(fatal) << "Failed to cd to path " << path;
  }
  const TGeoNode* node = gGeoManager->GetCurrentNode();
  if (!node) {
    LOG(fatal) << "Failed to find TGeoNode at path " << path;
  }
  TGeoHMatrix nominalGlobal = *(gGeoManager->GetCurrentMatrix());
  TGeoHMatrix alignedGlobal =
    alignmentMatrix
    * nominalGlobal;  // TODO: make sure that the order of multiplication is correct and nominalTransformation is not altered here!
 
  // get nominal transformation matrix of the mother volume
  gGeoManager->CdUp();
  TGeoHMatrix motherGlobal = *(gGeoManager->GetCurrentMatrix());
  // get aligned local matrix alignedLocal:
  // alignedGlobal = motherGlobal * alignedLocal;
  TGeoHMatrix alignedLocal = motherGlobal.Inverse() * alignedGlobal;
 
  //extract the local alignment matrix
  // get original local matrix of the node
  TGeoHMatrix origLocal = *node->GetMatrix();
  // if the node matrix already contains some alignment, replace it with the original unaligned matrix
  {
    bool found        = false;
    TObjArray* pNodes = gGeoManager->GetListOfPhysicalNodes();
    for (int i = 0; i < pNodes->GetEntriesFast(); i++) {
      const TGeoPhysicalNode* pnode = dynamic_cast<const TGeoPhysicalNode*>(pNodes->At(i));
      if (pnode->GetNode() == node) {
        // there is a physical node for this geo node -> the node already has some alignment
        found     = true;
        origLocal = *pnode->GetOriginalMatrix();
        break;
      }
    }
    if (!found) {
      LOG(debug) << "Failed to find physical node for node " << path;
    }
    else {
      LOG(debug) << "Found physical node for node " << path;
    }
  }
 
  // calculate the local alignment matrix alignmentLocal that includes the existing alignment
  // alignedLocal = alignmentLocal * origLocal
  TGeoHMatrix alignmentLocal = alignedLocal * origLocal.Inverse();  // TODO: check the order of multiplication!!
  alignmentLocal.SetName("aligned with bba (local)");
  origLocal.SetName("original local matrix of sensor");
  alignmentMatrix.SetName("alignment matrix (global)");
 
  if (fair::Logger::GetConsoleSeverity() == fair::Severity::debug) {
    LOG(debug) << "Write alignment matrix for the node " << path << ":";
    alignmentMatrix.Print();
    origLocal.Print();
    alignmentLocal.Print();
  }
 
  return {path, alignmentLocal};  // return the path and the alignment matrix in local coordinates
}
 
// Create a pair of sensor path and alignment matrix in local coordinates for each sensor
std::map<std::string, TGeoHMatrix> CbmBbaAlignTask::CreateAlignmentMatrices(const std::vector<double>& par) const
{
  // create alignment matrices from the parameters
  std::map<std::string, TGeoHMatrix> alignmentMatrices;
  for (auto& s : fSensors) {
    // get the alignment body index
    int i = s.fAlignmentBody;
    assert(i < (int) fAlignmentBodies.size());
    if (i < 0) continue;                // skip sensors without alignment body
    if (s.fNodePath.empty()) continue;  // skip sensors without node path
    alignmentMatrices.insert(CreateAlignmentNode(s.fNodePath, par[i * kNdfBody + 0], par[i * kNdfBody + 1],
                                                 par[i * kNdfBody + 2], 0., 0., 0.));
  }
  return alignmentMatrices;
}
 
 
void CbmBbaAlignTask::TrackContainer::MakeConsistent()
{
  fNmvdHits   = 0;
  fNstsHits   = 0;
  fNmuchHits  = 0;
  fNtrd1dHits = 0;
  fNtrd2dHits = 0;
  fNtofHits   = 0;
  for (const auto& n : fTrack.GetNodes()) {
    switch (Trajectory::GetHitSystemId(n)) {
      case ECbmModuleId::kMvd: fNmvdHits++; break;
      case ECbmModuleId::kSts: fNstsHits++; break;
      case ECbmModuleId::kMuch: fNmuchHits++; break;
      case ECbmModuleId::kTrd: fNtrd1dHits++; break;
      case ECbmModuleId::kTrd2d: fNtrd2dHits++; break;
      case ECbmModuleId::kTof: fNtofHits++; break;
      default: break;
    }
  }
}
 
void CbmBbaAlignTask::TrackContainer::StoreOriginalHitPositions()
{
  fOriginalHitPositions.clear();
  for (const auto& n : fTrack.GetNodes()) {
    if (n.isXySet) {
      fOriginalHitPositions.push_back({n.measurementXY.X(), n.measurementXY.Y(), n.z});
    }
    else {
      fOriginalHitPositions.push_back({0., 0., 0.});
    }
  }
}