CbmBbaAlignTask.h

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 */

 
#ifndef CbmBbaAlignTask_H
#define CbmBbaAlignTask_H
 
#if __has_include(<omp.h>)
#define HAVE_OMP
#endif
 
#include "CbmBbaAlignmentBody.h"
#include "CbmKfTrackFitter.h"
#include "FairTask.h"
#include "TClonesArray.h"
#include "TFile.h"
#include "TGeoPhysicalNode.h"
#include "TH1F.h"
#include "TString.h"
#include "bba/Parameter.h"
 
#include <TMatrixD.h>
 
#include <cstddef>
#include <string>
 
namespace
{
  constexpr size_t kNdfBody = cbm::bba::AlignmentBody::NofParameters;  // degrees of freedom per alignment body
}
 
class CbmBbaAlignTask : public FairTask {
 public:
  CbmBbaAlignTask(const char* name = "CbmBbaAlignTask", Int_t iVerbose = 0,
                  TString histoFileName = "./CbmBbaAlignmentHisto.root", size_t nMaxTracks = 1000);
 
  InitStatus Init();
  void Exec(Option_t* opt);
  void Finish();
 
  // clang-format off
  enum AlignmentMode
  {
    station,  // one alignment body per tracking station
    sensor    // one alignment body per sensor
  };
  // clang-format on
 
  class Trajectory : public CbmKfTrackFitter<cbm::algo::kf::DoFitTime::N>::Trajectory {
   public:
    using CbmKfTrackFitter<cbm::algo::kf::DoFitTime::N>::Trajectory::Trajectory;
 
    static int GetSensorId(const Node& node) { return node.userReferences[3]; }
    static void SetSensorId(Node& node, int sensorId) { node.userReferences[3] = sensorId; }
    void SetSensorId(size_t inode, int sensorId) { SetSensorId(fNodes[inode], sensorId); }
 
    static int GetAlignmentBodyId(const Node& node) { return node.userReferences[4]; }
    static void SetAlignmentBodyId(Node& node, int bodyId) { node.userReferences[4] = bodyId; }
    void SetAlignmentBodyId(size_t inode, int bodyId) { SetAlignmentBodyId(fNodes[inode], bodyId); }
  };

  struct TrackContainer {
    Trajectory fTrack;                                         // track
    cbm::algo::kf::TrackParamD fFittedTrackParam;              // fitted track parameters
    int fNmvdHits{0};                                          // number of MVD hits
    int fNstsHits{0};                                          // number of STS hits
    int fNmuchHits{0};                                         // number of MUCH hits
    int fNtrd1dHits{0};                                        // number of TRD hits
    int fNtrd2dHits{0};                                        // number of TRD hits
    int fNtofHits{0};                                          // number of TOF hits
    bool fIsActive{1};                                         // is the track active
    std::vector<std::array<double, 3>> fOriginalHitPositions;  // original hit positions
    void StoreOriginalHitPositions();
    void MakeConsistent();
  };
 
  struct Sensor {  // TGeoNode to be aligned
    ECbmModuleId fSystemId{0};
    int fTrackingStation{-1};
    int fAlignmentBody{-1};
    std::string fNodePath{""};  // full path to the node in the root geometry
 
    bool operator<(const Sensor& other) const
    {
      if (fSystemId < other.fSystemId) return true;
      if (fSystemId > other.fSystemId) return false;
      if (fTrackingStation < other.fTrackingStation) return true;
      if (fTrackingStation > other.fTrackingStation) return false;
      return (fNodePath < other.fNodePath);
    };
 
    bool operator==(const Sensor& other) const { return fNodePath == other.fNodePath; };
  };
 
  struct Histograms {
    std::vector<std::unique_ptr<TH1F>> fResidualsX;  // residuals of x-coordinates of the hits
    std::vector<std::unique_ptr<TH1F>> fResidualsY;  // residuals of y-coordinates of the hits
    std::vector<std::unique_ptr<TH1F>> fPullsX;      // pulls of x-coordinates of the hits
    std::vector<std::unique_ptr<TH1F>> fPullsY;      // pulls of y-coordinates of the hits
  };
 
  void SetSimulatedMisalignmentRange(double range) { fSimulatedMisalignmentRange = range; }
  double GetSimulatedMisalignmentRange() const { return fSimulatedMisalignmentRange; }
  void SetRandomSeed(int seed) { fRandomSeed = seed; }
  int GetRandomSeed() const { return fRandomSeed; }
  void SetMaxTracks(size_t nMaxTracks) { fNmaxTracks = nMaxTracks; }
  size_t GetMaxTracks() const { return fNmaxTracks; }
 
 private:
  unsigned SelectTracks(TClonesArray* inputTracks, std::vector<TrackContainer>& tracks);
  std::vector<Sensor> SensorsFromTracks(const std::vector<TrackContainer>& tracks) const;
  std::vector<cbm::bba::AlignmentBody> CreateAndSetAlignmentBodies(AlignmentMode mode,
                                                                   std::vector<Sensor>& sensors) const;
  void SetReferences(const std::vector<Sensor>& sensors, std::vector<TrackContainer>& tracks) const;
 
  std::vector<bba::Parameter> CreateAlignmentParameters() const;
  std::vector<double> InvertParameters(std::vector<double> par) const;
  std::vector<double> DiffParameters(const std::vector<double>& par1, const std::vector<double>& par2) const;
  std::string PrintParameters(const std::vector<double>& par) const;
  std::string PrintRMSE(const std::vector<bba::Parameter>& par, const std::vector<double>& parMC,
                        const std::vector<double>& parValues) const;
 
  std::vector<double> SimulateMisalignment(const std::vector<bba::Parameter>& par, std::vector<TrackContainer>& tracks,
                                           std::vector<cbm::bba::AlignmentBody>& AlignmentBodies) const;
 
  void InitializeHistograms();
  void FillHistograms(const std::vector<TrackContainer>& tracks, Histograms& histo);
  void WriteObjectsCurFile(TObject* obj);
 
  void SetAlignment(const std::vector<double>& par, std::vector<cbm::bba::AlignmentBody>& AlignmentBodies) const;
  void ApplyAlignment(const std::vector<double>& par, std::vector<TrackContainer>& tracks,
                      std::vector<cbm::bba::AlignmentBody>& AlignmentBodies, bool invert) const;
 
  double CostFunction(const std::vector<double>& par, std::vector<TrackContainer>& tracks);
  double CostFunction(const std::vector<double>& par);
 
  // Output
  std::pair<std::string, TGeoHMatrix> CreateAlignmentNode(std::string path, double shiftX, double shiftY, double shiftZ,
                                                          double rotX, double rotY, double rotZ) const;
  std::map<std::string, TGeoHMatrix> CreateAlignmentMatrices(const std::vector<double>& par) const;
  TString fMatrixOutFileName{"AlignmentMatrices.root"};
 
  std::vector<TrackContainer> fTracks;
  std::vector<Sensor> fSensors;
  std::vector<cbm::bba::AlignmentBody> fAlignmentBodies;
 
  TClonesArray* fInputStsTracks{nullptr};
  CbmKfTrackFitter<cbm::algo::kf::DoFitTime::N> fFitter;
 
  int fNtrackingStations{0};  // number of tracking stations
  size_t fNevents{0};         // number of processed events
  size_t fNmaxTracks{0};      // maximum number of tracks in the event
  size_t fNcallsCost{};
 
  // Simulation settings
  double fSimulatedMisalignmentRange{0.};
  int fRandomSeed{1};  // random seed for misalignment
  std::vector<double> fParMC;
  std::vector<double> fParStart;
 
  // paramaters for alignment
  AlignmentMode fAlignmentMode{station};
 
  double fCostIdeal{1.e10};
  double fCostInitial{0.};
 
  // inputs/outputs of the cost function
  int fNthreads{1};
  long fFixedNdf{-1};
 
  // histograms
  Histograms fHistoBeforeAlignment;
  Histograms fHistoAfterAlignment;
 
  TString fHistoFileName{"CbmBbaAlignmentHisto.root"};
  TFile* fHistoFile{nullptr};
  TDirectory* fHistoDir{nullptr};
 
  // make root happy
  ClassDef(CbmBbaAlignTask, 0);
};
#endif