CDCSpaceResols.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2016 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: CDC group                                                *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
#pragma once
 
#include <map>
#include <array>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <algorithm>
#include <TObject.h>
 
namespace Belle2 {
  class CDCSpaceResols: public TObject {
 
    typedef std::array<float, 3> array3; 
    typedef unsigned short SigmaID;         
  public:
    enum {c_nSLayers      = 56, 
          c_maxNAlphaBins = 18, 
          c_maxNThetaBins =  7, 
          c_maxNSigmaParams =  9  
         };
 
    CDCSpaceResols() : m_sigmaParamMode(0), m_nSigmaParams(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 setSigmaParamMode(unsigned short mode)
    {
      m_sigmaParamMode = mode;
    }
 
    void setSigmaParams(const SigmaID sigmaID, const std::vector<float>& params)
    {
      unsigned short nSigmaPars = params.size();
 
      if (nSigmaPars <= c_maxNSigmaParams) {
        m_nSigmaParams = nSigmaPars;
        m_sigmas.insert(std::pair<SigmaID, std::vector<float>>(sigmaID, params));
        //  std::cout <<"sigmaID in setSigmaParams= " << sigmaID << std::endl;
      } else {
        B2FATAL("The no. of sigma params. > limit !");
      }
    }
 
    void setSigmaParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta,
                        const std::vector<float>& params)
    {
      const SigmaID sigmaID = getSigmaID(iCLayer, iLR, iAlpha, iTheta);
      setSigmaParams(sigmaID, params);
    }
 
    void addSigmaParams(const SigmaID sigmaID, const std::vector<float>&  delta)
    {
      std::map<SigmaID, std::vector<float>>::iterator it = m_sigmas.find(sigmaID);
 
      if (it != m_sigmas.end()) {
        for (unsigned short i = 0; i < m_nSigmaParams; ++i) {
          (it->second)[i] += delta[i];
        }
      } else {
        B2FATAL("Specified params not found in addSigmaParams !");
      }
    }
 
    void addSigmaParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta,
                        const std::vector<float>& delta)
    {
      const SigmaID sigmaID = getSigmaID(iCLayer, iLR, iAlpha, iTheta);
      addSigmaParams(sigmaID, delta);
    }
 
 
    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 getSigmaParamMode() const
    {
      return m_sigmaParamMode;
    }
 
    SigmaID getSigmaID(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha, unsigned short iTheta) const
    {
      SigmaID id = iCLayer + 64 * iLR + 128 * iAlpha + 4096 * iTheta;
      return id;
    }
 
    SigmaID getSigmaID(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 !");
 
      return getSigmaID(iCLayer, iLR, iAlpha, iTheta);
    }
 
    const std::vector<float>& getSigmaParams(const SigmaID sigmaID) const
    {
      std::map<SigmaID, std::vector<float>>::const_iterator it = m_sigmas.find(sigmaID);
      if (it != m_sigmas.end()) {
        return it->second;
      } else {
        B2FATAL("Specified params. not found in getSigmaParams !");
      }
    }
 
    const std::vector<float>& getSigmaParams(unsigned short iCLayer, unsigned short iLR, unsigned short iAlpha,
                                             unsigned short iTheta) const
    {
      const SigmaID sigmaID = getSigmaID(iCLayer, iLR, iAlpha, iTheta);
      //  std::cout <<"sigmaID in getSigmaParams= " << sigmaID << std::endl;
      return getSigmaParams(sigmaID);
    }
 
    void dump() const
    {
      std::cout << " " << std::endl;
      std::cout << "Contents of sigma 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 sigma" << 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 = getSigmaParams(iCL, iLRp, iA, iT);
              for (unsigned short i = 0; i < m_nSigmaParams; ++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_sigmaParamMode << "  " << m_nSigmaParams << 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);
                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) << iLRp;
                const std::vector<float> params = getSigmaParams(iCL, iLRp, iA, iT);
                for (unsigned short i = 0; i < m_nSigmaParams; ++i) {
                  fout << "  " << std::setw(15) << std::scientific << std::setprecision(8) << params[i];
                }
                fout << std::endl;
              }
            }
          }
        }
        fout.close();
      }
    }
 
  private:
    unsigned short m_sigmaParamMode;    
    unsigned short m_nSigmaParams;      
    std::vector<array3> m_alphaBins; 
    std::vector<array3> m_thetaBins; 
    std::map<SigmaID, std::vector<float>>
                                       m_sigmas;    
    ClassDef(CDCSpaceResols, 1); 
  };
 
} // end namespace Belle2