CDCXtRelations.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
#pragma once
 
#include <framework/logging/Logger.h>
 
#include <map>
#include <array>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <algorithm>
#include <TObject.h>
 
namespace Belle2 {
  class CDCXtRelations: public TObject {
 
    typedef std::array<float, 3> array3; 
    typedef unsigned short XtID;         
  public:
    enum {c_nSLayers      = 56, 
          c_maxNAlphaBins = 18, 
          c_maxNThetaBins =  7, 
          c_maxNXtParams  =  8  
         };
 
    CDCXtRelations() : m_xtParamMode(0), m_nXtParams(0)
    {}
 
    void setAlphaBin(const array3& alpha)
    {
      if (m_alphaBins.size() <= c_maxNAlphaBins) {
        m_alphaBins.push_back(alpha);
        sort(m_alphaBins.begin(), m_alphaBins.end(), comp);
      } else {
        //  std::cout<< m_alphaBins.size() <<" "<< c_maxNAlphaBins <<std::endl;
        B2FATAL("The no. of alpha bins > limit !");
      }
    }
 
    void setThetaBin(const array3& theta)
    {
      if (m_thetaBins.size() <= c_maxNThetaBins) {
        m_thetaBins.push_back(theta);
        sort(m_thetaBins.begin(), m_thetaBins.end(), comp);
      } else {
        B2FATAL("The no. of theta bins > limit !");
      }
    }
 
    static bool comp(const array3& lhs, const array3& rhs)
    {
      return lhs[0] < rhs[0];
    }
 
    void setXtParamMode(unsigned short mode)
    {
      m_xtParamMode = mode;
    }
 
    void setXtParams(const XtID xtID, const std::vector<float>& params)
    {
      unsigned short nXtParams = params.size();
 
      if (nXtParams <= c_maxNXtParams) {
        m_nXtParams = nXtParams;
        m_xts.insert(std::pair<XtID, std::vector<float>>(xtID, params));
        //  std::cout <<"xtID in setXtParams= " << xtID << std::endl;
      } else {
        B2FATAL("The no. of xt params. > limit !");
      }
    }
 
    void setXtParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta,
                     const std::vector<float>& params)
    {
      const XtID xtID = getXtID(iCLayer, iLR, iAlpha, iTheta);
      setXtParams(xtID, params);
    }
 
    void addXTParams(const XtID xtID, const std::vector<float>&  delta)
    {
      std::map<XtID, std::vector<float>>::iterator it = m_xts.find(xtID);
 
      if (it != m_xts.end()) {
        for (unsigned short i = 0; i < m_nXtParams; ++i) {
          (it->second)[i] += delta[i];
        }
      } else {
        B2FATAL("Specified xt params not found in addXTParams !");
      }
    }
 
    void addXTParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta,
                     const std::vector<float>& delta)
    {
      const XtID xtID = getXtID(iCLayer, iLR, iAlpha, iTheta);
      addXTParams(xtID, delta);
    }
 
    void replaceXTParams(const XtID xtID, const std::vector<float>&  param)
    {
      std::map<XtID, std::vector<float>>::iterator it = m_xts.find(xtID);
 
      if (it != m_xts.end()) {
        for (unsigned short i = 0; i < m_nXtParams; ++i) {
          (it->second)[i] = param[i];
        }
      } else {
        B2FATAL("Specified xt params not found in replaceXTParams !");
      }
    }
 
    void replaceXTParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta,
                         const std::vector<float>& param)
    {
      const XtID xtID = getXtID(iCLayer, iLR, iAlpha, iTheta);
      replaceXTParams(xtID, param);
    }
 
 
    unsigned short getNoOfAlphaBins() const
    {
      return m_alphaBins.size();
    }
 
    unsigned short getNoOfThetaBins() const
    {
      return m_thetaBins.size();
    }
 
    const array3& getAlphaBin(unsigned short i) const
    {
      return m_alphaBins[i];
    }
 
    float getAlphaPoint(unsigned short i) const
    {
      return m_alphaBins[i][2];
    }
 
    const array3& getThetaBin(unsigned short i) const
    {
      return m_thetaBins[i];
    }
 
    float getThetaPoint(unsigned short i) const
    {
      return m_thetaBins[i][2];
    }
 
    unsigned short getXtParamMode() const
    {
      return m_xtParamMode;
    }
 
    XtID getXtID(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta) const
    {
      XtID id = iCLayer + 64 * iLR + 128 * iAlpha + 4096 * iTheta;
      return id;
    }
 
    XtID getXtID(unsigned short iCLayer, unsigned short iLR, float alpha, float theta) const
    {
      /*
      unsigned short iTheta = 999;
      unsigned short ibin = 0;
      for (std::vector<array3>::const_iterator it = m_thetaBins.begin(); it != m_thetaBins.end(); ++it) {
      if ((*it)[0] <= theta && theta <= (*it)[1]) {
      iTheta = ibin;
      break;
      }
      ++ibin;
      }
      */
      unsigned short iTheta = 999;
      unsigned short ibin = 0;
      for (auto const& it : m_thetaBins) {
        if (it[0] <= theta && theta <= it[1]) {
          iTheta = ibin;
          break;
        }
        ++ibin;
      }
      if (iTheta == 999) B2FATAL("Theta bin not found !");
 
      /*
      unsigned short iAlpha = 999;
      ibin = 0;
      for (std::vector<array3>::const_iterator it = m_alphaBins.begin(); it != m_alphaBins.end(); ++it) {
      if ((*it)[0] <= alpha && alpha <= (*it)[1]) {
      iAlpha = ibin;
      break;
      }
      ++ibin;
      }
      */
      unsigned short iAlpha = 999;
      ibin = 0;
      for (auto const& it : m_alphaBins) {
        if (it[0] <= alpha && alpha <= it[1]) {
          iAlpha = ibin;
          break;
        }
        ++ibin;
      }
      if (iAlpha == 999) B2FATAL("Alpha bin not found ! " << alpha);
 
      return getXtID(iCLayer, iLR, iAlpha, iTheta);
    }
 
    const std::vector<float>& getXtParams(const XtID xtID) const
    {
      std::map<XtID, std::vector<float>>::const_iterator it = m_xts.find(xtID);
      if (it != m_xts.end()) {
        return it->second;
      } else {
        B2FATAL("Specified xt params. not found in getXtParams !");
      }
    }
 
    const std::vector<float>& getXtParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha,
                                          unsigned short iTheta) const
    {
      const XtID xtID = getXtID(iCLayer, iLR, iAlpha, iTheta);
      //  std::cout <<"xtID in getXtParams= " << xtID << std::endl;
      return getXtParams(xtID);
    }
 
    void dump() const
    {
      std::cout << " " << std::endl;
      std::cout << "Contents of xt db" << std::endl;
      std::cout << "alpha bins" << std::endl;
 
      const double deg = 180. / M_PI;
 
      unsigned short nAlphaBins = m_alphaBins.size();
      for (unsigned short i = 0; i < nAlphaBins; ++i) {
        std::cout << " " << deg* m_alphaBins[i][0] << " " << deg* m_alphaBins[i][1] << " " << deg* m_alphaBins[i][2] << " " << std::endl;
      }
 
      std::cout << " " << std::endl;
      std::cout << "theta bins" << std::endl;
 
      unsigned short nThetaBins = m_thetaBins.size();
      for (unsigned short i = 0; i < nThetaBins; ++i) {
        std::cout << " " << deg* m_thetaBins[i][0] << " " << deg* m_thetaBins[i][1] << " " << deg* m_thetaBins[i][2] << " " << std::endl;
      }
 
      std::cout << " " << std::endl;
      std::cout << "coefficients for xt" << std::endl;
 
      for (unsigned short iT = 0; iT < nThetaBins; ++iT) {
        for (unsigned short iA = 0; iA < nAlphaBins; ++iA) {
          for (unsigned short iCL = 0; iCL < c_nSLayers; ++iCL) {
            for (unsigned short iLR = 0; iLR < 2; ++iLR) {
              unsigned short iLRp = abs(iLR - 1);
              std::cout << iCL << "  " << deg* m_thetaBins[iT][2] << "  " << deg* m_alphaBins[iA][2] << "  " << iLRp;
              const std::vector<float> params = getXtParams(iCL, iLRp, iA, iT);
              for (unsigned short i = 0; i < m_nXtParams; ++i) {
                std::cout << "  " << params[i];
              }
              std::cout << "  " << std::endl;
            }
          }
        }
      }
    }
 
    void outputToFile(std::string fileName) const
    {
      std::ofstream fout(fileName);
 
      if (fout.bad()) {
        B2ERROR("Specified output file could not be opened!");
      } else {
        const double deg = 180. / M_PI;
 
        unsigned short nAlphaBins = m_alphaBins.size();
        fout << nAlphaBins << std::endl;
 
        for (unsigned short i = 0; i < nAlphaBins; ++i) {
          fout << deg* m_alphaBins[i][0] << "  " << deg* m_alphaBins[i][1] << "  " << deg* m_alphaBins[i][2] << std::endl;
        }
 
        unsigned short nThetaBins = m_thetaBins.size();
        fout << nThetaBins << std::endl;
 
        for (unsigned short i = 0; i < nThetaBins; ++i) {
          fout << deg* m_thetaBins[i][0] << "  " << deg* m_thetaBins[i][1] << "  " << deg* m_thetaBins[i][2] << std::endl;
        }
 
        fout << m_xtParamMode << "  " << m_nXtParams << std::endl;
 
        signed short phiAngle = 0.;
        for (unsigned short iT = 0; iT < nThetaBins; ++iT) {
          for (unsigned short iA = 0; iA < nAlphaBins; ++iA) {
            for (unsigned short iCL = 0; iCL < c_nSLayers; ++iCL) {
              for (unsigned short iLR = 0; iLR < 2; ++iLR) {
                unsigned short iLRp = abs(iLR - 1);
                fout << std::setw(2) << std::right << std::fixed << iCL << "  " << std::setw(5) << std::setprecision(
                       1) << deg* m_thetaBins[iT][2] << "  " << std::setw(5) << std::right << deg* m_alphaBins[iA][2] << "  " << std::setw(
                       1) << std::setprecision(1) << phiAngle << "  " << std::setw(1) << iLRp;
                const std::vector<float> params = getXtParams(iCL, iLRp, iA, iT);
                for (unsigned short i = 0; i < m_nXtParams; ++i) {
                  fout << "  " << std::setw(15) << std::scientific << std::setprecision(8) << params[i];
                }
                fout << std::endl;
              }
            }
          }
        }
        fout.close();
      }
    }
 
    // ------------- Interface to global Millepede calibration ----------------
    static unsigned short getGlobalUniqueID() {return 29;}
    double getGlobalParam(unsigned short xtId, unsigned short i) const
    {
      return getXtParams(xtId).at(i);
    }
    void setGlobalParam(double value, unsigned short xtId, unsigned short i)
    {
 
      std::map<XtID, std::vector<float>>::const_iterator it = m_xts.find(xtId);
      if (it != m_xts.end()) {
        std::vector<float> allParams =  it->second;
        allParams.at(i) = value;
        setXtParams(xtId, allParams);
      } else {
        B2INFO("Specified xt params. not found in getXtParams.");
        std::vector<float> allParams {0., 0., 0., 0., 0., 0., 0., 0.};
        allParams.at(i) = value;
        setXtParams(xtId, allParams);
      }
    }
    std::vector<std::pair<unsigned short, unsigned short>> listGlobalParams() const
    {
      std::vector<std::pair<unsigned short, unsigned short>> result;
      for (auto ixt : m_xts) {
        for (int i = 0; i < 8; ++i) {
          result.push_back({ixt.first, i});
        }
      }
      return result;
    }
  private:
    unsigned short m_xtParamMode;    
    unsigned short m_nXtParams;      
    std::vector<array3> m_alphaBins; 
    std::vector<array3> m_thetaBins; 
    std::map<XtID, std::vector<float>>
                                    m_xts;    
    ClassDef(CDCXtRelations, 2); 
  };
 
} // end namespace Belle2