Step.h

Go to the documentation of this file.

#ifndef STEP_HEADER
#define STEP_HEADER
 
#include <cassert>
#include <cmath>
 
#include <array>
#include <deque>
#include <iostream>
#include <limits>
#include <numeric>
#include <vector>
 
#include <TF1.h>
#include <TH1.h>
#include <TH2.h>
#include <TVectorD.h>
#include <TMatrixD.h>
#include <TDecompSVD.h>
#include <TMath.h>
#include <TRandom3.h>
#include <TString.h>
 
#include "Math/VectorUtil.h"
#include "Math/Minimizer.h"
#include "Math/Factory.h"
#include "Math/Functor.h"
 
namespace Step {
 
static const char * green = "\x1B[32m",
                  * red   = "\x1B[31m",
                  * bold  = "\e[1m",
                  * def   = "\e[0m";
 
static const auto eps = std::numeric_limits<double>::epsilon();
static const auto inf = std::numeric_limits<double>::infinity();
static bool verbose = false;
 
struct Result {
    const bool valid; 
    /*const*/ double chi2; 
    /*const*/ unsigned int ndf; //<! number of degrees of freedom
    std::vector<double> X; 
    size_t degree () const { return X.size()-1; } 
    double x2n    () const { return    chi2/ndf ; } 
    double x2nErr () const { return sqrt(2./ndf); } 
    double prob   () const { return TMath::Prob(chi2, ndf); } 
 
    // replicas
    std::vector<double> chi2sT, 
                        chi2sV; 
    double chi2T () const { return std::accumulate(std::begin(chi2sT), std::end(chi2sT), 0.) / chi2sT.size(); }
    double chi2V () const { return std::accumulate(std::begin(chi2sV), std::end(chi2sV), 0.) / chi2sV.size(); }
    double x2nT () const { return chi2T() /  ndf            ; }
    double x2nV () const { return chi2V() / (ndf + X.size()); }
 
    // get TF1
    template<typename POL>
    TF1 * GetTF1 (const char * name, POL pol, double m, double M) const {
        pol.npars = X.size();
        auto f = new TF1(name, pol, m, M, pol.npars);
        f->SetParameters(X.data());
        f->SetTitle(Form("#chi^{2}_{%zu}/ndf = %.1f/%d", degree(), chi2, ndf));
        return f;
    }
};
 
std::ostream& operator<< (std::ostream& stream, const Result& result)
{
    using namespace std;
    if (result.valid) {
        stream << bold << "Chi2/ndf = " << result.chi2 << " / " << result.ndf
                            << " = " << result.x2n() << "\u00B1" << result.x2nErr()
                            << " (prob = " << result.prob() << ")\n" << def;
        size_t nrep = result.chi2sT.size();
        assert(nrep == result.chi2sV.size());
        stream << result.X.size() << " parameters: ";
        for (auto x: result.X) stream << std::setw(12) << x;
        stream << '\n';
        if (nrep > 0)
            stream << nrep <<   " training replicas: Chi2T/" <<  result.ndf                    << " = " << result.x2nT() << '\n'
                   << nrep << " validation replicas: Chi2V/" << (result.ndf + result.X.size()) << " = " << result.x2nV() << '\n';
    }
    else stream << "Not a fit\n";
    return stream;
}
 
static std::deque<Result> chi2s; 
static auto bestResult = chi2s.begin(); 
 
double Ftest (const Result& former, const Result& current)
{
    using namespace std;
    if (!former.valid || !current.valid) throw -1.;
    double RSS1 = former.chi2,
           RSS2 = current.chi2;
    size_t p1 = former.X.size(),
           p2 = current.X.size();
    assert(p2 > p1);
    int ndf = current.ndf; // = nbins - p2
    double x = ((RSS1-RSS2)/(p2-p1))/(RSS2/ndf);
    double pvalue = TMath::FDistI(x, p2-p1, ndf);
    return pvalue;
}
 
double Xval (const Result& current, const Result& former)
{
    if (former.chi2sV.size() != current.chi2sV.size())
        return 0;
 
    size_t nrep = former.chi2sV.size();
    if (nrep == 0) return 0;
 
    int Nf = former .ndf + former .X.size(),
        Nc = current.ndf + current.X.size();
    double nBetter = 0;
    for (size_t i = 0; i < nrep; ++i)
        if (former.chi2sV[i]/Nf < current.chi2sV[i]/Nc)
            ++nBetter;
    nBetter /= nrep;
    return nBetter;
}
 
void PrintAllTests (std::function<double(const Result&, const Result&)> test,
                    std::ostream& stream = std::cout, size_t best = 0)
{
    using namespace std;
 
    int nGoodFits = 0;
    for (auto result = begin(chi2s); result != end(chi2s); ++result)
        if (result->valid) ++nGoodFits;
    if (nGoodFits < 2) {
        stream << red << "Only one valid fit. No test to run.\n" << def;
        return;
    }
    stream << "   " << bold;
    for (auto result = begin(chi2s); result != prev(end(chi2s)); ++result) {
        // loop over all results except the highest one
        size_t deg = result->degree();
        if (deg == best) stream << green;
        if (!result->valid) stream << red;
        stream << setw(10) << deg << def;
    }
    stream << '\n';
    for (auto result2 = next(begin(chi2s)); result2 != end(chi2s); ++result2) {
        // loop over all results
        size_t deg2 = result2->degree();
        if (deg2 == best) stream << green;
        if (!result2->valid) stream << red;
        stream << setw(3) << deg2 << def;
        if (deg2 == best) stream << def;
        stream << setprecision(4);
        for (auto result1 = begin(chi2s); result1 != result2; ++result1) {
            // compare to all results with less parameters
            try            { stream << setw(10) << test(*result1, *result2); }
            catch (double) { stream << red << "         -" << def; }
        }
        stream << bold << '\n';
    }
    stream << def;
}
 
double identity (double x) { return x; }
 
struct FunctionalForm {
 
    unsigned int npars; 
    const double m, 
                 M; 
    std::ostream& stream; 
 
protected:
 
    FunctionalForm (const char * classname, unsigned int Npars, double mi, double Ma, std::ostream& Stream = std::cout) :
        npars(Npars), m(mi), M(Ma), stream(Stream)
    {
        if (verbose)
            stream << "Declaring a `" << classname << "` with " << npars << " parameters between " << mi << " and " << Ma << '\n';
    }
 
    virtual ~FunctionalForm () = default;
 
public:
 
    virtual double operator() (const double *x, const double *p) const = 0;
};
 
namespace Basis {
    double Chebyshev (double x, int i)
    {
        if (std::abs(x) > 1) return 0.;
        switch (i) {
            case 0: return 1.;
            case 1: return x;
            default: return 2*x*Chebyshev(x, i-1)-Chebyshev(x, i-2);
        }
    }
 
    double Legendre (double x, int i)
    {
        if (std::abs(x) > 1) return 0.;
        switch (i) {
            case 0: return 1.;
            case 1: return x;
            default: return ( (2*i-1)*x*Legendre(x,i-1) - (i-1)*Legendre(x, i-2) )/i;
        }
    }
}
 
template<double (*T)(double,int) = Basis::Chebyshev,
         double (*Log)(double) = identity,
         double (*Exp)(double) = identity>
struct Polynomial : public FunctionalForm {
 
    Polynomial (unsigned int d, double mi, double Ma, std::ostream& Stream = std::cout) :
        FunctionalForm(__func__, d+1, mi, Ma, Stream) { }
 
    virtual ~Polynomial () = default;
 
    double operator() (const double *x, const double *p) const override
    {
        double nx = -1 + 2* (Log(*x) - Log(m)) / (Log(M) - Log(m));
        double result = 0;
        for (unsigned int i = 0; i < npars; ++i)
            result += p[i]*T(nx,i);
        return Exp(result);
    }
};
 
template<typename FunctionalForm>
struct Correlator {
 
    const int n; 
    TVectorD c, 
             e; 
    TVectorD v, 
             w, 
             s; 
    TMatrixD cov, 
             rot, 
             invcov; 
 
    FunctionalForm& f; 
    std::ostream& stream; 
 
    Correlator (
            TH1 * h, 
            TH2 * hcov, 
            int im, 
            int iM, 
            FunctionalForm& ff, 
            std::ostream& Stream = std::cout 
            ) :
        n(iM-im+1), c(n), e(n+1), v(n), w(n), s(n), cov(n,n), rot(n,n), invcov(n,n), f(ff), stream(Stream)
    {
        using namespace std;
 
        // sanity checks
 
        if (verbose)
            stream << "n = " << n << '\n';
 
        if (verbose)
            stream << "Checking histogram\n";
        for (int i = im; i <= iM; ++i) {
            auto content = h->GetBinContent(i);
            if (verbose)
                stream << setw(3) << i
                       << setw(12) << h->GetBinCenter(i)
                       << setw(12) << content << '\n';
            assert(isnormal(content));
        }
 
        if (hcov != nullptr) {
            if (verbose)
                stream << "Checking covariance\n";
            for (int i = im; i <= iM; ++i)
            for (int j = im; j <= iM; ++j) {
                auto content = hcov->GetBinContent(i,j);
                if (verbose) {
                    stream << setw(12) << content;
                    if (j == iM) stream << '\n';
                }
                assert(isfinite(content));
            }
        }
 
        // conversion for matrix algebra
 
        int shift = im;
        if (verbose)
            stream << "Converting from histograms to vector and matrix (shift = " << shift << "):";
        for (int i = 0; i < n; ++i) {
            e(i) = h->GetBinLowEdge(i+shift);
            c(i) = h->GetBinCenter (i+shift);
            v(i) = h->GetBinContent(i+shift);
            if (verbose)
                stream << '\n' << setw(5) << i << setw(10) << c(i) << setw(15) << v(i) << setw(15);
            assert(abs(v(i)) > 0);
            for (int j = 0; j < n; ++j) {
                cov(i,j) = hcov != nullptr ? hcov->GetBinContent(i+shift,j+shift)
                                  : i == j ? pow(h->GetBinError(i+shift),2) : 0;
                if (i == j && verbose) stream << (100*sqrt(cov(i,j))/v(i)) << '%';
            }
        }
        e(n) = h->GetBinLowEdge(n+shift);
        if (verbose)
            stream << '\n';
 
        if (!cov.IsSymmetric()) {
            if (verbose)
                stream << "Forcing to be symmetric\n";
 
            auto cov2 = cov;
            cov2.T();
            cov += cov2;
            cov *= 0.5;
            assert(cov.IsSymmetric());
        }
 
        // preparing variables for replicas
        w = v;
        rot = cov.EigenVectors(s); // `s` will be modified here, but will be overwritten later
        if (verbose)
            stream << "Getting eigenvalues, expect different columns only in case of non-zero bin-to-bin correlations:";
        for (int k = 0; k < n; ++k) {
            if (verbose)
                stream << '\n' << setw(5) << k << setw(15) << s[k] << setw(15) << cov(k,k);
            s(k) = sqrt(s(k)); // to be used as input of `TRandom3::Gaus()`
            assert(s[k] > 0);
        }
        if (verbose)
            stream << '\n';
 
        // inverting
        if (verbose)
            stream << "Inverting...\n";
        TDecompSVD svd(cov);
        svd.SetTol(eps);
        bool status = false;
        invcov = svd.Invert(status);
        assert(status);
 
        if (verbose)
            stream << "Ready to fit\n";
    }
 
    double operator() (const double * p)
    {
        auto r = w;
        for (int i = 0; i < n; ++i) {
            auto y0 = f(&e(i), p),
                 y1 = f(&c(i), p),
                 y2 = f(&e(i+1), p);
            r(i) -= (y0+4*y1+y2)/6;
        }
        return r * (invcov * r);
    }
 
    void SetReplica (UInt_t seed)
    {
        TRandom3 r(seed);
 
        TVectorD d(n);
        for (int k = 0; k < n; ++k)
            d(k) = r.Gaus(0, s(k));
 
        w = v + rot*d;
    }
 
    void UnsetReplica ()
    {
        w = v;
    }
};
 
enum EarlyStopping {
    None, 
    Chi2ndf, 
    fTest, 
    xVal 
};
 
template<double (*T)(double, int) = Basis::Chebyshev,
         double (*Log)(double) = identity,
         double (*Exp)(double) = identity>
TF1 * GetSmoothFit (
        TH1 * h, 
        TH2 * cov, 
        int im, 
        int iM, 
        unsigned int maxdegree, 
        EarlyStopping criterion = None, 
        double stopParam = 1, 
        UInt_t nrep = 0, 
        std::ostream& stream = std::cout 
        )
{
    using namespace std;
 
    assert(h != nullptr);
 
    if (verbose)
        stream << "im = " << im << ", iM = " << iM << '\n';
 
    unsigned int nbins = iM-im+1;
    if (verbose)
        stream << "nbins = " << nbins << '\n';
    assert(nbins > 2);
 
    double m = h->GetBinLowEdge(im),
           M = h->GetBinLowEdge(iM+1);
    if (verbose)
        stream << "m = " << m << ", M = " << M << '\n';
 
    Polynomial<T,Log,Exp> pol(maxdegree, m, M, stream);
    Correlator cor(h, cov, im, iM, pol, stream);
    ROOT::Math::Functor functor(cor, pol.npars);
 
    auto minimiser = ROOT::Math::Factory::CreateMinimizer("Minuit2", "Migrad");
    const int Ncalls = 1e6;
    minimiser->SetMaxFunctionCalls(Ncalls);
    minimiser->SetFunction(functor);
    minimiser->SetTolerance(0.001);
    minimiser->SetPrintLevel(verbose);
 
    double fm = Log(h->GetBinContent(im)),
           fM = Log(h->GetBinContent(iM));
    auto step = 0.001;
    double p0 = (fM+fm)/2, p1 = (fM-fm)/2;
    minimiser->SetVariable(0,"p0",p0,step);
    minimiser->SetVariable(1,"p1",p1,step);
    for (unsigned int d = 2; d <= maxdegree; ++d) {
        minimiser->SetVariable(d,Form("p%d", d),0.,step);
        minimiser->FixVariable(d);
    }
 
    // running
    chi2s.clear();
    chi2s.push_back({false, inf, nbins-2, {p0,p1}});
    bestResult = chi2s.begin();
    bool found = false;
    while (chi2s.size() <= maxdegree) {
        unsigned int degree = chi2s.size();
        minimiser->ReleaseVariable(degree);
        bool valid = minimiser->Minimize();
 
        const double * X = minimiser->X();
        assert(X != nullptr);
        double chi2 = cor(X);
 
        assert(isnormal(chi2));
        unsigned int npars = minimiser->NFree(),
                     ndf = nbins - npars;
        assert(npars == degree+1);
 
        const auto& former = chi2s.back();
        chi2s.push_back({valid, chi2, ndf, vector<double>(X,X+npars)});
        auto& current = chi2s.back();
 
        switch (minimiser->Status()) {
            case 0:  stream << current;
                     {
                         // check EarlyStopping::Chi2ndf criterion
                         bool plateau = current.x2n() > former.x2n(),
                              goodchi2 = abs(current.x2n()-1) <= stopParam*current.x2nErr();
                         if (plateau)
                             stream << red << "Chi2/ndf has increased.\n" << def;
                         else if (goodchi2)
                             stream << green << "Chi2 ~ ndf.\n" << def;
 
                         // check EarlyStopping::fTest criterion
                         double ftestRes = 0;
                         if (bestResult != chi2s.begin())
                         try {
                             ftestRes = Ftest(former, current);
                             stream << "F-test result: " << ftestRes << '\n';
                         }
                         catch (double) { stream << red << "No F-test possible.\n" << def; }
                         bool fvalid = ftestRes > stopParam; 
 
                         // check EarlyStopping::xVal criterion
                         double xvalFrac = 0;
                         if (nrep > 0) {
                             current.chi2sT.reserve(nrep);
                             current.chi2sV.reserve(nrep);
                             for (UInt_t seed = 1; seed <= nrep; ++seed) {
                                 cor.SetReplica(seed);
                                 for (unsigned int i = 0; i <= degree; ++i)
                                     minimiser->SetVariableValue(i, current.X[i]);
                                 const double * X = minimiser->X();
                                 assert(X != nullptr);
                                 current.chi2sT.push_back(cor(X));
                                 cor.SetReplica(seed+nrep);
                                 current.chi2sV.push_back(cor(X));
                             }
                             cor.UnsetReplica();
                             for (unsigned int i = 0; i <= degree; ++i)
                                 minimiser->SetVariableValue(i, current.X[i]);
                             X = minimiser->X();
                             xvalFrac = Xval(current, former);
                         }
                         bool xvalid = 1.-xvalFrac > stopParam;
                         if (nrep > 0)
                             stream << "Fraction of validation replicas with better fit Chi2/ndf = " << (xvalid ? green : red) << xvalFrac << def << '\n';
 
                         // apply early stopping criterion
                         if (!found)
                         switch (criterion) {
                             case EarlyStopping::Chi2ndf:
                                 found = plateau || goodchi2;
                                 if (plateau) --bestResult;
                                 break;
                             case EarlyStopping::fTest:
                                 found = !fvalid;
                                 break;
                             case EarlyStopping::xVal:
                                 assert(nrep > 0);
                                 found = !xvalid;
                                 break;
                             case EarlyStopping::None:
                                 found = false;
                         }
                     }
                     break;
            case 1:  stream << red << "Warning: Fit parameter covariance was made positive defined\n" << def;
                     for (unsigned int i = 0; i < npars; ++i) {
                         for (unsigned int j = 0; j < npars; ++j)
                             stream << setw(12) << minimiser->CovMatrix(i,j);
                         stream << '\n';
                     }
                     break;
            case 2:  stream << red << "Warning: Hesse is invalid\n" << def;
                     {
                         double * hessian = nullptr;
                         minimiser->GetHessianMatrix(hessian);
                         assert(hessian != nullptr);
                         for (unsigned int i = 0; i < npars; ++i) {
                             for (unsigned int j = 0; j < npars; ++j)
                                 stream << setw(12) << hessian[i*pol.npars + j];
                             stream << '\n';
                         }
                     }
                     break;
            case 3:  stream << red << "Warning: Expected Distance reached from the Minimum (EDM = " << minimiser->Edm() << ") is above max\n" << def;
                     break;
            case 4:  stream << red << "Warning: Reached call limit (Ncalls = " << Ncalls << ")\n" << def;
                     break;
            case 5:
            default: stream << red << "Warning: Unknown failure\n" << def;
        }
 
        if (!found) ++bestResult;
 
        if (verbose) {
            stream << "Correlations of the fit parameters:\n";
            for (unsigned int i = 0; i < npars; ++i) {
                for (unsigned int j = 0; j < npars; ++j)
                    stream << setw(12) << minimiser->Correlation(i,j);
                stream << '\n';
            }
        }
 
        if (current.ndf == 2) {
            stream << red << "Only 2 d.o.f. left. Stopping.\n" << def;
            break;
        }
 
        stream << "---\n";
    } // end of `while`-loop on number of parameters
    chi2s.pop_front(); // removing first entry, which is not a fit (chi2 = inf)
 
    const char * name = Form("smooth_%s", h->GetName());
    TF1 * f = bestResult->GetTF1(name, pol, m, M);
    stream << "Best result: " << f->GetTitle() << "\n-----\n";
 
    stream << bold << "Running F-test:\n" << def;
    PrintAllTests(Ftest, stream, bestResult->degree());
 
    if (nrep > 0) {
        stream << bold << "Running Cross-validation:\n" << def;
        PrintAllTests(Xval, stream, bestResult->degree());
    }
 
    return f;
}
 
template<double (*T)(double, int) = Basis::Chebyshev,
         double (*Log)(double) = identity,
         double (*Exp)(double) = identity>
TF1 * GetSmoothFit (TH1 * h, TH2 * cov, unsigned int maxdegree,
        EarlyStopping criterion = EarlyStopping::None, double stopParam = 1, UInt_t nrep = 0,
        std::ostream& stream = std::cout)
{
    int im = 1, iM = h->GetNbinsX();
    return GetSmoothFit<T,Log,Exp>(h, cov, im, iM,
            maxdegree, criterion, stopParam, nrep, stream);
}
 
template<double (*T)(double, int) = Basis::Chebyshev,
         double (*Log)(double) = identity,
         double (*Exp)(double) = identity>
TF1 * GetSmoothFit (TH1 * h, int im, int iM, unsigned int maxdegree,
        EarlyStopping criterion = EarlyStopping::None, double stopParam = 1, UInt_t nrep = 0,
        std::ostream& stream = std::cout)
{
    return GetSmoothFit<T,Log,Exp>(h, nullptr, im, iM,
            maxdegree, criterion, stopParam, nrep, stream);
}
 
template<double (*T)(double, int) = Basis::Chebyshev,
         double (*Log)(double) = identity,
         double (*Exp)(double) = identity>
TF1 * GetSmoothFit (TH1 * h, unsigned int maxdegree,
        EarlyStopping criterion = EarlyStopping::None, double stopParam = 1, UInt_t nrep = 0,
        std::ostream& stream = std::cout)
{
    return GetSmoothFit<T,Log,Exp>(h, nullptr,
            maxdegree, criterion, stopParam, nrep, stream);
}
 
 
} // end of Step namespace
#endif