PDFConstructor.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 <top/reconstruction_cpp/TOPTrack.h>
#include <top/reconstruction_cpp/InverseRaytracer.h>
#include <top/reconstruction_cpp/FastRaytracer.h>
#include <top/reconstruction_cpp/YScanner.h>
#include <top/reconstruction_cpp/SignalPDF.h>
#include <top/reconstruction_cpp/BackgroundPDF.h>
#include <top/reconstruction_cpp/DeltaRayPDF.h>
#include <top/geometry/TOPGeometryPar.h>
#include <vector>
#include <map>
#include <set>
#include <limits>
 
namespace Belle2 {
  namespace TOP {
 
    class PDFConstructor {
 
    public:
 
      enum EPDFOption {
        c_Rough = 0,   
        c_Fine = 1,    
        c_Optimal = 2  
      };
 
      enum EStoreOption {
        c_Reduced = 0, 
        c_Full = 1     
      };
 
      struct LogL {
        double logL = 0; 
        double expPhotons = 0; 
        unsigned numPhotons = 0; 
        explicit LogL(double phot): logL(-phot), expPhotons(phot)
        {}
      };
 
      struct Pull {
        int pixelID = 0;    
        double time = 0;    
        double peakT0 = 0;  
        double ttsT0 = 0;   
        double sigma = 0;   
        double phiCer = 0;  
        double wt = 0;      
        Pull(int pix, double t, double t0, double tts0, double sig, double phi, double w):
          pixelID(pix), time(t), peakT0(t0), ttsT0(tts0), sigma(sig), phiCer(phi), wt(w)
        {}
      };
 
      PDFConstructor(const TOPTrack& track, const Const::ChargedStable& hypothesis,
                     EPDFOption PDFOption = c_Optimal, EStoreOption storeOption = c_Reduced, double overrideMass = 0);
 
      bool isValid() const {return m_valid;}
 
      void switchOffDeltaRayPDF() const {m_deltaPDFOn = false;}
 
      void switchOnDeltaRayPDF() const {m_deltaPDFOn = true;}
 
      void switchDeltaRayPDF(bool deltaPDFOn) const {m_deltaPDFOn = deltaPDFOn;}
 
      int getModuleID() const {return m_moduleID;}
 
      const Const::ChargedStable& getHypothesis() const {return m_hypothesis;}
 
      const std::vector<TOPTrack::SelectedHit>& getSelectedHits() const {return m_selectedHits;}
 
      double getBkgRate() const {return m_bkgRate;}
 
      double getCosTotal() const {return m_cosTotal;}
 
      double getCosCerenkovAngle(double E) const;
 
      const std::vector<SignalPDF>& getSignalPDF() const {return m_signalPDFs;}
 
      const BackgroundPDF* getBackgroundPDF() const {return m_backgroundPDF;}
 
      const DeltaRayPDF& getDeltaRayPDF() const {return m_deltaRayPDF;}
 
      double getExpectedSignalPhotons() const {return m_signalPhotons;}
 
      double getExpectedDeltaPhotons() const {return m_deltaPDFOn ? m_deltaPhotons : 0;}
 
      double getExpectedBkgPhotons() const
      {
        double photons = m_bkgPhotons;
        for (const auto* other : m_pdfOtherTracks) {
          photons += other->getExpectedSignalPhotons() + other->getExpectedDeltaPhotons();
        }
        return photons;
      }
 
      double getExpectedPhotons() const
      {
        return getExpectedSignalPhotons() + getExpectedDeltaPhotons() + getExpectedBkgPhotons();
      }
 
      double getExpectedSignalPhotons(double minTime, double maxTime) const
      {
        double ps = 0;
        for (const auto& signalPDF : m_signalPDFs) {
          ps += signalPDF.getIntegral(minTime, maxTime);
        }
        return ps * m_signalPhotons;
      }
 
      double getExpectedDeltaPhotons(double minTime, double maxTime) const
      {
        return m_deltaPDFOn ? m_deltaRayPDF.getIntegral(minTime, maxTime) * m_deltaPhotons : 0;
      }
 
      double getExpectedBkgPhotons(double minTime, double maxTime) const
      {
        double photons = (maxTime - minTime) / (m_maxTime - m_minTime) * m_bkgPhotons;
        for (const auto* other : m_pdfOtherTracks) {
          photons += other->getExpectedSignalPhotons(minTime, maxTime) + other->getExpectedDeltaPhotons(minTime, maxTime);
        }
        return photons;
      }
 
      double getExpectedPhotons(double minTime, double maxTime) const
      {
        return getExpectedSignalPhotons(minTime, maxTime) + getExpectedDeltaPhotons(minTime, maxTime) +
               getExpectedBkgPhotons(minTime, maxTime);
      }
 
      double getPDFValue(int pixelID, double time, double timeErr, double sigt = 0) const
      {
        return pdfValue(pixelID, time, timeErr, sigt) / getExpectedPhotons();
      }
 
      LogL getLogL() const;
 
      LogL getLogL(double t0, double sigt = 0) const {return getLogL(t0, m_minTime, m_maxTime, sigt);}
 
      LogL getLogL(double t0, double minTime, double maxTime, double sigt = 0) const;
 
      LogL getBackgroundLogL() const {return getBackgroundLogL(m_minTime, m_maxTime);}
 
      LogL getBackgroundLogL(double minTime, double maxTime) const;
 
      const std::vector<LogL>& getPixelLogLs(double t0, double sigt = 0) const
      {return getPixelLogLs(t0, m_minTime, m_maxTime, sigt);}
 
      const std::vector<LogL>& getPixelLogLs(double t0, double minTime, double maxTime, double sigt = 0) const;
 
      const std::vector<Pull>& getPulls() const;
 
      int getNCalls_setPDF(SignalPDF::EPeakType type) const {return m_ncallsSetPDF[type];}
 
      int getNCalls_expandPDF(SignalPDF::EPeakType type) const {return m_ncallsExpandPDF[type];}
 
      const std::map <double, YScanner::Derivatives>& getDerivatives() const {return m_derivatives;}
 
      void appendPDFOther(const PDFConstructor* pdfOther)
      {
        if (not pdfOther) return;
        m_pdfOtherTracks.push_back(pdfOther);
      }
 
      void clearPDFOther() {m_pdfOtherTracks.clear();}
 
 
    private:
 
      struct PrismSolution {
        double len = 0;    
        double L = 0;      
        double cosFic = 0; 
        double sinFic = 0; 
      };
 
      struct InverseRaytracerDirect {
        const InverseRaytracer* inverseRaytracer = 0; 
        const SignalPDF::EPeakType type = SignalPDF::c_Direct; 
        int solve(double xD, double zD, int, double, double,
                  const TOPTrack::AssumedEmission& assumedEmission,
                  const InverseRaytracer::CerenkovAngle& cer, double step = 0) const
        {
          return inverseRaytracer->solveDirect(xD, zD, assumedEmission, cer, step);
        }
      };
 
      struct InverseRaytracerReflected {
        const InverseRaytracer* inverseRaytracer = 0; 
        const SignalPDF::EPeakType type = SignalPDF::c_Reflected; 
        int solve(double xD, double zD, int Nxm, double xmMin, double xmMax,
                  const TOPTrack::AssumedEmission& assumedEmission,
                  const InverseRaytracer::CerenkovAngle& cer, double step = 0) const
        {
          return inverseRaytracer->solveReflected(xD, zD, Nxm, xmMin, xmMax, assumedEmission, cer, step);
        }
      };
 
      template<class T>
      void setSignalPDF(T& t, unsigned col, double xD, double zD, int Nxm = 0, double xmMin = 0, double xmMax = 0);
 
      const InverseRaytracer::CerenkovAngle& cerenkovAngle(double dE = 0);
 
      void setSignalPDF();
 
      void setSignalPDF_direct();
 
      void setSignalPDF_reflected();
 
      void setSignalPDF_prism();
 
      void setSignalPDF_reflected(int Nxm, double xmMin, double xmMax);
 
      bool detectionPositionX(double xM, int Nxm, std::vector<double>& xDs, double& minLen);
 
      bool doRaytracingCorrections(const InverseRaytracer::Solution& sol, double dFic_dx, double xD);
 
      double deltaXD(double dFic, const InverseRaytracer::Solution& sol, double xD);
 
      bool setDerivatives(YScanner::Derivatives& D, double dL, double de, double dFic);
 
      bool raytrace(const FastRaytracer& rayTracer, double dL = 0, double de = 0, double dFic = 0);
 
      double propagationLosses(double E, double propLen, int nx, int ny, SignalPDF::EPeakType type) const;
 
      void expandSignalPDF(unsigned col, const YScanner::Derivatives& D, SignalPDF::EPeakType type);
 
      bool rangeOfX(double z, double& xmi, double& xma);
 
      double findReflectionExtreme(double xE, double zE, double zD, int Nxm, double A,
                                   const RaytracerBase::Mirror& mirror) const;
 
      double derivativeOfReflectedX(double x, double xe, double ze, double zd) const;
 
      bool prismRaytrace(const PrismSolution& sol, double dL = 0, double dFic = 0, double de = 0);
 
      PrismSolution prismSolution(const PixelPositions::PixelData& pixel, unsigned k, int nx);
 
      PrismSolution prismSolution(const ROOT::Math::XYZPoint& rD, double L);
 
      double pdfValueSignalDelta(int pixelID, double time, double timeErr, double sigt = 0) const;
 
      double pdfValue(int pixelID, double time, double timeErr, double sigt = 0) const;
 
      void initializePixelLogLs(double minTime, double maxTime) const;
 
      void appendPulls(const TOPTrack::SelectedHit& hit) const;
 
      int m_moduleID = 0; 
      const TOPTrack& m_track;   
      const Const::ChargedStable m_hypothesis; 
      const InverseRaytracer* m_inverseRaytracer = 0; 
      const FastRaytracer* m_fastRaytracer = 0; 
      std::vector<FastRaytracer> m_rayTracers; 
      const YScanner* m_yScanner = 0; 
      const BackgroundPDF* m_backgroundPDF = 0; 
      DeltaRayPDF m_deltaRayPDF; 
      EPDFOption m_PDFOption = c_Optimal; 
      EStoreOption m_storeOption = c_Reduced; 
      bool m_valid = false; 
      double m_beta = 0; 
      double m_tof = 0; 
      double m_groupIndex = 0; 
      double m_groupIndexDerivative = 0; 
      double m_cosTotal = 0; 
      double m_minTime = 0; 
      double m_maxTime = 0; 
      std::vector<TOPTrack::SelectedHit> m_selectedHits; 
      double m_bkgRate = 0; 
      std::vector<SignalPDF> m_signalPDFs; 
      double m_signalPhotons = 0; 
      double m_bkgPhotons = 0; 
      double m_deltaPhotons = 0; 
      std::map<double, InverseRaytracer::CerenkovAngle> m_cerenkovAngles; 
      double m_dFic = 0; 
      double m_Fic = 0;  
      mutable std::map <SignalPDF::EPeakType, int> m_ncallsSetPDF; 
      mutable std::map <SignalPDF::EPeakType, int> m_ncallsExpandPDF; 
      mutable std::vector<LogL> m_pixelLLs; 
      mutable std::vector<Pull> m_pulls; 
      mutable bool m_deltaPDFOn = true; 
      std::map <double, YScanner::Derivatives> m_derivatives; 
      mutable std::set<int> m_zeroPixels; 
      std::vector<const PDFConstructor*> m_pdfOtherTracks; 
    };
 
    //--- inline functions ------------------------------------------------------------
 
    inline double PDFConstructor::pdfValueSignalDelta(int pixelID, double time, double timeErr, double sigt) const
    {
      unsigned k = pixelID - 1;
      if (k < m_signalPDFs.size() and m_valid) {
        double f = m_signalPhotons * m_signalPDFs[k].getPDFValue(time, timeErr, sigt);
        if (m_deltaPDFOn) f += m_deltaPhotons * m_deltaRayPDF.getPDFValue(pixelID, time);
        return f;
      }
      return 0;
    }
 
    inline double PDFConstructor::pdfValue(int pixelID, double time, double timeErr, double sigt) const
    {
 
      if (not m_valid) return 0;
 
      double f = pdfValueSignalDelta(pixelID, time, timeErr, sigt); // signal
      f += m_bkgPhotons * m_backgroundPDF->getPDFValue(pixelID); // uniform background
      for (const auto* other : m_pdfOtherTracks) f += other->pdfValueSignalDelta(pixelID, time, timeErr, sigt); // other tracks
 
      return f;
    }
 
    inline double PDFConstructor::deltaXD(double dFic, const InverseRaytracer::Solution& sol, double xD)
    {
      m_dFic = dFic;
      m_fastRaytracer->propagate(m_inverseRaytracer->getReconstructedPhoton(sol, dFic), true);
      if (not m_fastRaytracer->getPropagationStatus()) return std::numeric_limits<double>::quiet_NaN();
      return m_fastRaytracer->getXD() - xD;
    }
 
    inline double PDFConstructor::getCosCerenkovAngle(double E) const
    {
      double refind = TOPGeometryPar::Instance()->getPhaseIndex(E);
      return std::min(1 / m_beta / refind, 1.0);
    }
 
    inline const InverseRaytracer::CerenkovAngle& PDFConstructor::cerenkovAngle(double dE)
    {
      auto& cer = m_cerenkovAngles[dE];
      if (cer.cosThc == 0 and cer.sinThc == 0) {
        double meanE = m_yScanner->getMeanEnergy();
        double cosThc = getCosCerenkovAngle(meanE + dE);
        cer = InverseRaytracer::CerenkovAngle(cosThc);
      }
      return cer;
    }
 
 
    template<class T>
    void PDFConstructor::setSignalPDF(T& t, unsigned col, double xD, double zD, int Nxm, double xmMin, double xmMax)
    {
      m_ncallsSetPDF[t.type]++;
 
      m_inverseRaytracer->clear();
      t.inverseRaytracer = m_inverseRaytracer;
 
      // central solutions
 
      int i0 = t.solve(xD, zD, Nxm, xmMin, xmMax, m_track.getEmissionPoint(), cerenkovAngle());
      if (i0 < 0 or not m_inverseRaytracer->getStatus()) return;
      int n = 0;
      for (unsigned i = 0; i < 2; i++) {
        if (not m_inverseRaytracer->getStatus(i)) continue;
        const auto& solutions = m_inverseRaytracer->getSolutions(i);
        const auto& sol = solutions[i0];
        double time = m_tof + sol.len * m_groupIndex / Const::speedOfLight;
        if (time > m_maxTime + 1.0) continue;
        n++;
      }
      if (n == 0) return;
 
      // solutions with xD displaced by dx
 
      double dx = 0.1; // cm
      int i_dx = t.solve(xD + dx, zD, Nxm, xmMin, xmMax, m_track.getEmissionPoint(), cerenkovAngle(), dx);
      if (i_dx < 0) return;
      int k = 0;
      while (m_inverseRaytracer->isNymDifferent()) { // get rid of discontinuities
        if (k > 8) {
          B2DEBUG(20, "TOP::PDFConstructor::setSignalPDF: failed to find the same Nym (dx)");
          return;
        }
        dx = - dx / 2;
        i_dx = t.solve(xD + dx, zD, Nxm, xmMin, xmMax, m_track.getEmissionPoint(), cerenkovAngle(), dx);
        if (i_dx < 0) return;
        k++;
      }
 
      // loop over the two solutions, do ray-tracing corrections, compute the derivatives and expand PDF in y
 
      for (unsigned i = 0; i < 2; i++) {
        if (not m_inverseRaytracer->getStatus(i)) continue;
        const auto& solutions = m_inverseRaytracer->getSolutions(i);
        const auto& sol = solutions[i0];
        const auto& sol_dx = solutions[i_dx];
 
        // compute dFic/dx needed in raytracing corrections
 
        double dFic_dx = YScanner::Derivatives::dFic_d(sol, sol_dx);
 
        // do raytracing corrections
 
        m_dFic = 0;
        bool ok = doRaytracingCorrections(sol, dFic_dx, xD);
        if (not ok) continue;
        m_Fic = sol.getFic() + m_dFic;
 
        // check if time is still within the reconstruction range
 
        double time = m_tof + m_fastRaytracer->getPropagationLen() * m_groupIndex / Const::speedOfLight;
        if (time > m_maxTime) continue;
 
        // compute the derivatives
 
        YScanner::Derivatives D;
        ok = setDerivatives(D, 0.1, 0.1, 0.001);
        if (not ok) {
          B2DEBUG(20, "TOP::PDFConstructor::setSignalPDF: failed to determine derivatives: "
                  << LogVar("track momentum", m_track.getMomentumMag())
                  << LogVar("impact local z", m_track.getEmissionPoint().position.Z())
                  << LogVar("xD", xD)
                  << LogVar("Nxm", Nxm)
                  << LogVar("time", time)
                  << LogVar("peak type", t.type));
          continue;
        }
        if (m_storeOption == c_Full) m_derivatives[xD] = D;
 
        // expand PDF in y
 
        expandSignalPDF(col, D, t.type);
      }
    }
 
  } // namespace TOP
} // namespace Belle2