LTF.h

/*
  Copyright 2021, D. Britzger, Max-Planck-Institute for Physics, Munich, Germany

  Permission is hereby granted, free of charge, to any person obtaining
  a copy of this software and associated documentation files (the
  "Software"), to deal in the Software without restriction, including
  without limitation the rights to use, copy, modify, merge, publish,
  distribute, sublicense, and/or sell copies of the Software, and to
  permit persons to whom the Software is furnished to do so, subject to
  the following conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  The Software is provided "as is", without warranty of any kind,
  express or implied, including but not limited to the warranties of
  merchantability, fitness for a particular purpose and
  noninfringement. In no event shall the authors or copyright holders be
  liable for any claim, damages or other liability, whether in an action
  of contract, tort or otherwise, arising from, out of or in connection
  with the Software or the use or other dealings in the Software.
*/
#ifndef LTF_h
#define LTF_h

// --------------------------------------------------------------- //
/*
 *
 * LTF
 *
 * The Linear Template Fit implemented using the Eigen library
 *
 */
// --------------------------------------------------------------- //

#include "Eigen/Sparse"
#include "Eigen/Dense"
#include <string>
#include <map>
#include <unordered_map>
#include <iostream>
#include <memory>

using namespace std;

class LTF {

public:
   enum class Uncertainty { Constrained=1, Unconstrained=0, External=-1};
   
public:

   // ------------------------------------------- //
   class LiTeFit {
   public:

      LiTeFit() {}
      ~LiTeFit(){}
      double DoLiTeFit();                                                        
      double DoLiTeFit(int mPolN, int mOrdInfrc, const Eigen::VectorXd& ahat);   
      double DoIterativeFitNewton(int nIter=3, int nPol=2, int nInfrc=1);        
      double DoIterativeFitTaylor(int nIter=4, double StepSize=0.6, int nPol=2, int nInfrc=1);   
      void ApplyLogNormalDistribution();                                         
      void RescaleInputErrors(std::vector<double> values ){                      
         RescaleInputErrors(Eigen::Map<Eigen::VectorXd>(&values[0],values.size()));
      }
      void RescaleInputErrors(const Eigen::VectorXd& values );                   

      // --- input, as used in LiTeFit
      Eigen::VectorXd Dt;                                                        
      Eigen::MatrixXd M ;                                                        
      Eigen::MatrixXd Y ;                                                        
      Eigen::MatrixXd X ;                                                        
      Eigen::MatrixXd A ;                                                        
      std::vector<std::pair<std::string,Eigen::MatrixXd > > Vs    ;              
      std::vector<std::pair<std::string,Eigen::VectorXd > > Sys   ;              
      std::map<std::string,Eigen::MatrixXd >                VsExt ;              
      std::map<std::string,Eigen::VectorXd >                SysExt;              
      std::map<std::string,LTF::Uncertainty>                SysIsConstr;         
      std::map<std::string,Eigen::MatrixXd >                SysY;                
      std::map<std::string,Eigen::MatrixXd >                SysA;                
      std::map<std::string,double>                          corrSys;             
      vector<double> Gamma;                                                      

      // --- output
      double a0 = 0;                                                             
      double a0_errorFit = 0;                                                    
      double a0_errorExt  = 0;                                                   
      double a0_errorTmpl = 0;                                                   
      double a0_errorRespMat = 0;                                                
      double chisq = 0;                                                          
      double chisq_error = 0;                                                    
      Eigen::MatrixXd F;                                                         
      Eigen::VectorXd chisq_y;                                                   
      std::map<std::string,double> chisq_part;                                   
      Eigen::VectorXd ahat;                                                      
      Eigen::VectorXd ahat_errorFit;                                             
      Eigen::VectorXd ahat_errorExt;                                             
      Eigen::VectorXd ahat_errorTmpl;                                            
      Eigen::VectorXd ahat_errorRespMat;                                         
      Eigen::VectorXd TheoFit;                                                   
      std::map<std::string,Eigen::MatrixXd> Vsource;                             
      std::map<std::string,Eigen::VectorXd> DeltaSys;                            
      std::map<std::string,Eigen::MatrixXd> VsourceExt;                          
      std::map<std::string,Eigen::VectorXd> DeltaSysExt;                         
      std::map<std::string,Eigen::VectorXd> DeltaSysY;                           
      std::map<std::string,Eigen::VectorXd> DeltaSysA;                           

      // --- result from alternative pol2-fit to chi2-values of the templates
      Eigen::VectorXd achk;                                                      
      Eigen::VectorXd achk_errorFit;                                             
      double  achk_chisq;                                                        
      
      // --- result from linearized second-order model approximation
      Eigen::VectorXd ahat21;                                                    
      Eigen::VectorXd EDM21;                                                     
      double          EDM21_chisq;                                               
      Eigen::VectorXd delta_m21;                                                 
      
      // ---- printout, and useful getters
      void PrintShort() const;                                                   
      void PrintFull()  const;                                                   
      Eigen::MatrixXd Mc(int nPow=1, int nInfrc=0) const;                        
      Eigen::VectorXd a0pow(int nPow, int nInfrc, const Eigen::VectorXd& ahat) const; 
      Eigen::MatrixXd McApprox(int powmax, int intfrc, const Eigen::VectorXd& ahat) const;  
      Eigen::MatrixXd W() const;                                                 
      Eigen::MatrixXd VFit() const { return LTF::VSum(Vsource,DeltaSys); }       
      Eigen::MatrixXd VExt() const { return LTF::VSum(VsourceExt,DeltaSysExt); } 
      bool GetLogNormal() const { return LogNormal; }                            
      Eigen::VectorXd CalculatePrediction(double v ) const {                     
         return CalculatePrediction(std::vector<double>{v});}
      Eigen::VectorXd CalculatePrediction(std::vector<double> v ) const {        
         return CalculatePrediction(Eigen::Map<Eigen::VectorXd>(&v[0],v.size()));}
      Eigen::VectorXd CalculatePrediction(Eigen::VectorXd v ) const;             

      Eigen::VectorXd ComputeLinearTemplateFit(int mPolN = 1, int mOrdInfrc = 0, const Eigen::VectorXd& ahat = Eigen::VectorXd() ) const;  
      Eigen::VectorXd ComputeNewtonEstimator(int mPolN = 1, int mOrdInfrc = 0, const Eigen::VectorXd& ahat = Eigen::VectorXd() ) /*const*/;  

   protected:
      void ApplyGamma(std::map<std::string,Eigen::VectorXd>& Delta) const;       
      void ApplyGamma(std::map<std::string,Eigen::MatrixXd>& VMat) const;        
      bool LogNormal  = false;                                                   
      std::vector<std::pair<std::string,Eigen::MatrixXd> >     SysAX;            
      Eigen::VectorXd ErrorScale;                                                
      Eigen::VectorXd ybar;
      Eigen::MatrixXd Ytil;

   };  // end class LTF::LiTeFitResult
   // ------------------------------------------- //

public:
   LTF() {}
   ~LTF() {}

   // ------------------------------------------- //
   // --- set data
   // ------------------------------------------- //
   void SetData(const vector<double>& data) {  
      std::vector<double> tmp = data; // copy to no-const for eigen...
      fDt = Eigen::Map<Eigen::VectorXd>(&tmp[0], tmp.size());
   }
   void SetData(size_t n, const double* data) {
      SetData(vector<double>(data,data+n)); 
   }


   // ------------------------------------------- //
   // --- add a new template
   // ------------------------------------------- //
   void AddTemplate(double mval, const vector<double>& dist) {
      AddTemplate(vector<double>{mval},dist);
   }
   void AddTemplate(double mval, size_t n , const double* dist) {
      AddTemplate(mval,std::vector<double>(dist, dist+n));
   }
   void AddTemplate(const vector<double> mvals, size_t n , const double* dist) {
      AddTemplate(mvals,std::vector<double>(dist, dist+n));
   }
   void AddTemplate(const vector<double>& mvals, const vector<double>& dist);


   // ------------------------------------------- //
   // --- set error
   // ------------------------------------------- //
   void AddError(std::string name, size_t n, const double* error, double corr=0., LTF::Uncertainty constr = LTF::Uncertainty::Constrained); 

   void AddError(const std::string& name, const std::vector<double>& error, double corr=0., LTF::Uncertainty constr = LTF::Uncertainty::Constrained) {
      AddError(name,error.size(),&error[0],corr,constr);
   }

   void AddError(const std::string& name, const std::vector<std::vector<double> >& cov, LTF::Uncertainty constr = LTF::Uncertainty::Constrained);

   void AddErrorRelative(const std::string& name, const std::vector<std::vector<double> >& cov_rel, LTF::Uncertainty constr = LTF::Uncertainty::Constrained);

   void AddErrorRelative(std::string name, size_t n, const double* error, double corr=0., LTF::Uncertainty constr = LTF::Uncertainty::Constrained) {
      AddErrorRelative(name,std::vector<double>(error,error+n),corr,constr);
   }



   void AddErrorRelative(const std::string& name, std::vector<double> error, double corr=0., LTF::Uncertainty constr = LTF::Uncertainty::Constrained) {
      if ( fDt.rows() == 0 || size_t(fDt.rows()) != error.size() ) {
         std::cerr<<"Error!  Data distribution not set or its size is inconsistent with input uncertainty vector."<<std::endl;
         exit(1);
      }
      for ( size_t i=0 ; i<error.size( ); i++ ) error[i] *= fDt(i);
      AddError(name,error,corr,constr);
   }

   void AddErrorPercent(std::string name, size_t n, const double* error, double corr=0., LTF::Uncertainty constr = LTF::Uncertainty::Constrained) {
      AddErrorPercent(name,std::vector<double>(error,error+n),corr,constr);
   }

   void AddErrorPercent(const std::string& name, std::vector<double> error, double corr=0., LTF::Uncertainty constr = LTF::Uncertainty::Constrained) {
      if ( fDt.rows() == 0 || size_t(fDt.rows()) != error.size() ) {
         std::cerr<<"Error!  Data distribution not set or its size is inconsistent with input uncertainty vector."<<std::endl;
         exit(1);
      }
      for ( size_t i=0 ; i<error.size( ); i++ ) error[i] *= fDt(i) / 100.;
      AddError(name,error,corr,constr);
   }

   void AddUncorrelatedErrorSquared(const std::string& name, size_t n, const double* error2, LTF::Uncertainty constr = LTF::Uncertainty::Constrained) {
      if ( constr == LTF::Uncertainty::Unconstrained ) {
         std::cerr<<"Error!  An uncorrelated error cannot be 'unconstrained'. Error name: "<<name<<std::endl; 
         exit(1);
      }
      std::vector<std::vector<double> > cov(n);
      for ( size_t i = 0 ; i<n ;i++ ) {
         cov[i].resize(n);
         cov[i][i] = error2[i];
      }
      AddError(name, cov, constr);
   }

   void AddCorrelatedError(const std::string& name, vector<double> error, LTF::Uncertainty constr = LTF::Uncertainty::Constrained);
   
   // ------------------------------------------- //
   void AddTemplateError(const std::string& name, const std::vector<double>& mvals, size_t n, const double* error, double corr=0.);

   void AddTemplateError(const std::string& name, double mval, size_t n, const double* error, double corr=0.) {
      AddTemplateError(name, std::vector<double>{mval}, n, error, corr);
   }

   void AddTemplateError(const std::string& name, const std::vector<double>& mvals, const std::vector<double>& error, double corr=0.) {
      AddTemplateError(name,mvals,error.size(),&error[0],corr);
   }
   void AddTemplateError(const std::string& name, double mval, const std::vector<double>& error, double corr=0.) {
      AddTemplateError(name,std::vector<double>{mval},error.size(),&error[0],corr);
   }

   void AddTemplateErrorSquared(const std::string& name, const std::vector<double>& mvals, size_t n, const double* error2, double corr=0.) {
      AddTemplateErrorSquared(name, mvals, std::vector<double>(error2, error2+n), corr);
   }
   void AddTemplateErrorSquared(const std::string& name, double mval, size_t n, const double* error2, double corr=0.) {
      AddTemplateErrorSquared(name, std::vector<double>{mval}, std::vector<double>(error2, error2+n), corr);
   }

   void AddTemplateErrorSquared(const std::string& name, const std::vector<double>& mvals, const std::vector<double>& error2, double corr=0.) {
      std::vector<double> vect = error2;
      for ( double& e : vect ) e = std::sqrt(e);
      AddTemplateError(name, mvals, vect.size(), &vect[0], corr);
   }

   
   // ------------------------------------------- //
   // --- set a (detector) response matrix [optional]
   // ------------------------------------------- //
   void SetResponseMatrix(const std::vector<std::vector<double> >& A);

   void SetResponseMatrix(const Eigen::MatrixXd& A) { fA = A; }


   // ------------------------------------------- //
   void AddResponseMatrixError(const std::string& name, const std::vector<std::vector<double> >& errorA, double corr=0.) {
      AddResponseMatrixError(name, LTF::Std_to_Mat(errorA), corr);
   }
   void AddResponseMatrixError(const std::string& name, const Eigen::MatrixXd& errorA, double corr=0.) {
      SetCorrSys(name,corr);
      fSysA[name] = errorA;
   }
   void AddResponseMatrixErrorSquared(const std::string& name, std::vector<std::vector<double> > error2A, double corr=0.) {
      for ( auto& v : error2A )
         for ( double& e2 : v ) e2 = sqrt(e2);
      AddResponseMatrixError(name, LTF::Std_to_Mat(error2A), corr);
   }
   void AddResponseMatrixErrorSquared(const std::string& name, const Eigen::MatrixXd& error2A, double corr=0.) {
      AddResponseMatrixError(name, error2A.cwiseSqrt(), corr);
   }

   // ------------------------------------------- //
   // --- set exponential factor(s) gamma
   // ------------------------------------------- //
   void SetGamma(double gamma = 1) {
      if ( fTmplDistN.empty() ) { 
         std::cerr<<"Warning in SetGamma. Templates should be set first!"<<std::endl;
         fGamma = vector<double>(50,gamma);
      }
      else  {
         if ( fTmplDistN.begin()->first.size() > 1 ) 
            std::cout<<"Warning."
                     <<" SetGamma(double) is ambigous when performing a multivariate LTF."
                     <<" Better use SetGamma(vector<double>)."<<std::endl;
         
         fGamma = vector<double>(fTmplDistN.begin()->first.size(),gamma);
      }
   }
   void SetGamma(const vector<double>& gamma) {
      fGamma = gamma;
   }

   // ------------------------------------------- //
   // --- fit range
   // ------------------------------------------- //
   void SetFitRange(int binmin, int binmax) {
      fBinMin = binmin;
      fBinNN  = binmax -binmin+1;
      if ( fBinNN<=0 ) {
         cout<<"Error. No bins left for fit."<<endl;
         exit(1);
      }
   }

   // ------------------------------------------- //
   // --- perform linear template fit
   // ------------------------------------------- //
   const LiTeFit& DoLiTeFit();
   const LiTeFit& DoIterativeFitNewton(int nIter=3, int nPol=2, int nInfrc=1);

   // ------------------------------------------- //
   // --- Choose normal or log-normal uncertainties 
   // ------------------------------------------- //
   void UseLogNormalUncertainties(bool LogNormal = true) {
      fUseLog = LogNormal;
   }
   bool GetLogNormalUncertainties() const {
      return fUseLog;
   }

   // ------------------------------------------- //
   // --- use nuisance parameter or covariance matrices
   // ------------------------------------------- //
   void UseNuisanceParameters(bool UseNuisance=true) {
      fUseNuisance = UseNuisance;
   }


   // ------------------------------------------- //
   // --- rescale (absolute) uncertainties
   // ------------------------------------------- //
   void RescaleInputErrors(std::vector<double> values ) {
      if ( values.size() != size_t(fDt.rows()) ) {
         std::cerr<<"Error in LTF::RescaleInputErrors. Size of input vector not compatible with size of data vector."<<std::endl;
         exit(1);
      }
      fErrorScale = Eigen::Map<Eigen::VectorXd>(&values[0],values.size());
   }


// ------------------------------------------- //
// --- members
protected:
   Eigen::VectorXd fDt;                                                       
   std::vector<std::pair<vector<double>,Eigen::VectorXd> > fTmplDistN;        
   Eigen::MatrixXd                                       fA;                  
   std::vector<std::pair<std::string,Eigen::MatrixXd > > fVs;                 
   std::vector<std::pair<std::string,Eigen::VectorXd > > fSys;                
   std::map<std::string,LTF::Uncertainty>                fSysIsConstr;        
   std::map<std::string,Eigen::MatrixXd >                fVsExt;              
   std::map<std::string,Eigen::VectorXd >                fSysExt;             
   std::map<std::string,std::map<vector<double>,Eigen::VectorXd> > fSysY;     
   std::map<std::string,Eigen::MatrixXd>                 fSysA;               
   std::map<std::string,double>                          fcorrSys;            
                                                                           
   Eigen::VectorXd fErrorScale;                                               
   std::vector<double> fGamma;                                                
   bool   fUseLog = false;                                                    
   bool fUseNuisance = true;                                                  
   LiTeFit fLTF;                                                              
                                                                              
   int fBinMin = 0;                                                           
   int fBinNN  = 0;                                                           

// ------------------------------------------- //
// --- protected methods
protected:
   size_t N() const;                                                          
   Eigen::MatrixXd Calc_M() const;                                            
   Eigen::MatrixXd Calc_Y() const;                                            
   Eigen::MatrixXd Calc_X() const;                                            
   std::map<std::string,Eigen::MatrixXd > Calc_dY(
      std::map<std::string,std::map<vector<double>,Eigen::VectorXd> >& fSysY) const; 
   void SetLiTeFitInput();                                                    
   void SetCorrSys(const std::string& name, double c);                        
   

// ------------------------------------------- //
// ---- static helper functions
public:

   static Eigen::VectorXd GetDelta_V(const Eigen::MatrixXd& V) { 
      return V.diagonal().cwiseSqrt();
   }

   static Eigen::MatrixXd GetV_Delta(const Eigen::VectorXd& d, double corr);

   template< class T >
   static std::map<std::string,Eigen::MatrixXd> GetV_DeltaSys(const T& DeltaSys) {     
      std::map<std::string,Eigen::MatrixXd> ret;
      for ( auto& s : DeltaSys ) ret[s.first] = LTF::GetV_Delta(s.second,1);
      return ret; 
   }

   template< class T >
   static std::map<std::string,Eigen::VectorXd> GetDelta_Vsource(const T& Vsource) {  
      std::map<std::string,Eigen::VectorXd> ret;
      for ( auto& s : Vsource ) ret[s.first] = LTF::GetDelta_V(s.second);
      return ret; 
   }

   static Eigen::MatrixXd VSum( const std::map<std::string,Eigen::MatrixXd>& Vs, const std::map<std::string,Eigen::VectorXd>& DeltaSys );

   static Eigen::MatrixXd VSum( const std::vector<std::pair<std::string,Eigen::MatrixXd > >& Vs );

   static Eigen::MatrixXd Cov_to_Cor(const Eigen::MatrixXd& V);

   static Eigen::MatrixXd Std_to_Mat(const std::vector<std::vector<double> >& A);
   
};

#endif