Vegas.h

Go to the documentation of this file.

// Honeycomb, Copyright (C) 2015 NewGamePlus Inc.  Distributed under the Boost Software License v1.0.
#pragma once

#include "Honey/Math/Random/Random.h"

namespace honey
{


template<sdt Dim = 1, sdt DimRes = 1, class Real__ = Real, sdt BinCount = 100>
class Vegas
{
public:
    typedef Real__                          Real;
private:
    typedef typename Numeral<Real>::Real_   Real_;
    typedef typename Real_::DoubleType      Double_;
    typedef typename Double_::Real          Double;
    typedef Alge_<Real>                     Alge;
    typedef Uniform_<Real>                  Uniform;

public:
    static const sdt dim                    = Dim;
    static const sdt dimRes                 = DimRes;
    typedef Vec<dimRes,Real>                VecRes;
    typedef Vec<dim,Real>                   Vec;
    typedef function<VecRes (const Vec&)>   Func;


    Vegas(const Func& func, RandomGen& gen, const Vec& lower, const Vec& upper, sdt sampleCount = 1000, Real warmUp = 0.1, sdt iterCount = 5, Real alpha = 1.5) :

        func(func), gen(gen), lower(lower), upper(upper), sampleCount(sampleCount), warmUp(Alge::max(warmUp,0)), iterCount(iterCount), alpha(alpha),
        tgral(0), chi2a(0), sd(0), _progress(0), sampleTotal(0) {}

    ~Vegas() {}


    const VecRes& integrate(Real progressDelta = 1);

    Real progress() const                                                           { return _progress; }
    const VecRes& result() const                                                    { return tgral; }
    const VecRes& chiSqr() const                                                    { return chi2a; }
    const VecRes& stdDev() const                                                    { return sd; }

private:
    void init(sdt init, sdt sampleCount);
    void rebin(Real rc, sdt nd, Real r[], Real xin[], Real xi[]);
    void integrate_priv(sdt samplesMax);

    // Input
    Func func;
    RandomGen& gen;
    Vec lower;
    Vec upper;
    sdt sampleCount;
    Real warmUp;
    sdt iterCount;
    Real alpha;

    // Output
    VecRes tgral;
    VecRes chi2a;
    VecRes sd;

    // Progress
    Real _progress;
    sdt sampleTotal;
    sdt iterCur;
    sdt npgCur;

    // Init 0 and 1
    sdt ndo;
    sdt ittot;                  // iter count across init > 1
    sdt mds;                    // sampling mode

    // Init 2
    sdt nd;                     // slices in grid (c.f. BinCount)
    sdt ng;  
    sdt npg;                    // number of calls within bin
    Real calls;                 // real total number of calls to fxn
    Real dv2g;
    Real dxg;
    Real xnd;
    Real xJac;                  // Jacobian of integration
    Real dx[dim];               // width of integration region

    // Integrate locals
    // accumulator for stuff at end of iteration...
    struct iterAccu
    {
        Real Wgt;               // weight
        Real sWgt;              // cumulative sum for weights
        Real sChi;              // cumulative sum for chi^2
        Real sInt;              // cumulative sum for integral
    };                          
    iterAccu Ai[dimRes];        // ...one for each integrand

    Real d[BinCount][dim];
    Real di[BinCount][dim];     // delta i
    Real r[BinCount];
    Real xi[dim][BinCount];
    Real xin[BinCount];         // aux. variable for rebinning

    sdt ia[dim];
    sdt kg[dim];

    // accumulator over bins / hypercubes...
    struct binAccu
    {
        Real ti;                // sum for f over bins
        Real tsi;               // sum for variances over bins
    };                 
    binAccu Ab[dimRes];         // ...one for each integrand

    // accumulator over points x within bins...
    struct pointAccu
    {
        Real f2;                // f squared
        Real fb;                // sum for f within bi
        Real f2b;               // sum for f2 within bin
        sdt npg;                // number of calls within bin f != 0
    };               
    pointAccu Ax[dimRes];       // ...one for each integrand
};

template<sdt Dim, sdt DimRes, class Real, sdt BinCount>
const typename Vegas<Dim,DimRes,Real,BinCount>::VecRes&
    Vegas<Dim,DimRes,Real,BinCount>::integrate(Real progressDelta)
{
    sdt sampleWarmUp = warmUp*sampleCount;
    if (_progress == 0)
    {
        //First run, init with warm up or main
        if (sampleWarmUp > 0)
            init(0, sampleWarmUp);
        else
            init(0, sampleCount);
    }
    
    _progress = Alge::min(_progress + progressDelta, 1);
    sdt sampleAcc = _progress*(sampleWarmUp+sampleCount) - sampleTotal;
    sampleTotal += sampleAcc;

    if (sampleTotal - sampleAcc < sampleWarmUp)
    {
        //Warm up grid
        integrate_priv(sampleTotal < sampleWarmUp ? sampleAcc : -1);
        if (sampleTotal < sampleWarmUp)
            return result();
        //Done warm up, discard results and init main
        sampleAcc = sampleTotal - sampleWarmUp;
        init(1, sampleCount);
    }

    //Piecewise sample count is not accurate, so -1 means run until completion
    integrate_priv(_progress < 1 ? sampleAcc : -1);

    if (iterCur == iterCount)
        //Calculation is complete
        _progress = 1;

    return result();
}

template<sdt Dim, sdt DimRes, class Real, sdt BinCount>
void Vegas<Dim,DimRes,Real,BinCount>::init(sdt init, sdt sampleCount)
{
    iterCur = 0;
    npgCur = -2;

    if (init <= 0)      //entry for cold start
    {
        mds = 1;        //1 == use stratified sampling
        ndo = 1;
        for (sdt j=0; j<dim; ++j) xi[j][0] = 1;
    }

    if (init <= 1)      //inherit the previous grid
    {
        for (sdt j=0; j<dimRes; ++j)
        {
            Ai[j].sInt = 0;
            Ai[j].sWgt = 0;
            Ai[j].sChi = 0;
        }
        ittot = 1;
    }

    if (init <= 2)      //inherit grid and results
    {
        nd = BinCount;
        ng = 1;
        if (mds)
        {
            ng = Alge::pow(Real(sampleCount)/2+0.25,1/Real(dim));
            mds = 1;
            if ((2*ng-BinCount) >= 0)
            {
                mds = -1;
                npg = ng/BinCount+1;
                nd = ng/npg;
                ng = npg*nd;
            }
        }
        if (ng <= 0)
        {
            //No samples
            iterCur = iterCount;
            return;
        }
        sdt k = 1;
        for (sdt i=0; i<dim; ++i) k *= ng;
        npg = sampleCount/k > 2 ? sampleCount/k : 2;
        calls = npg * k;
        dxg = 1/Real(ng);
        dv2g = 1;
        for (sdt i=0; i<dim; ++i) dv2g *= dxg;
        dv2g = calls*calls*dv2g*dv2g/npg/npg/(npg-1.0);
        xnd = nd;
        dxg *= xnd;
        xJac = 1/Real(calls);
        for (sdt j=0; j<dim; ++j)
        {
            dx[j] = upper[j]-lower[j];
            xJac *= dx[j];
        }
        if (nd != ndo)
        {
            for (sdt i=0; i<(nd>ndo?nd:ndo); ++i) r[i] = 1;
            for (sdt j=0; j<dim; ++j) rebin(ndo/xnd,nd,r,xin,xi[j]);
            ndo = nd;
        }
    }
}

template<sdt Dim, sdt DimRes, class Real, sdt BinCount>
void Vegas<Dim,DimRes,Real,BinCount>::integrate_priv(sdt samplesMax)
{
    sdt samples = 0;
    samplesMax *= iterCount;

    for (; iterCur < iterCount; ++iterCur, ++ittot)
    {
        if (npgCur == -2)
        {
            //Init loop
            npgCur = -1;
            for (sdt j=0; j<dimRes; ++j) Ab[j].ti = Ab[j].tsi = 0;
            for (sdt j=0; j<dim; ++j)
            {
                kg[j] = 1;
                for (sdt i=0; i<nd; ++i) d[i][j] = di[i][j] = 0;
            }
        }
        for (;;)
        {
            if (npgCur == -1)
            {
                //Init loop
                npgCur = 0;
                for (sdt j=0; j<dimRes; ++j)
                {
                    Ax[j].fb = 0;
                    Ax[j].f2b = 0;
                    Ax[j].npg = 0;
                }
            }
            for (sdt k=npgCur; k<npg; ++k, ++samples)
            {
                if (samplesMax >= 0 && samples >= samplesMax)
                {
                    npgCur = k;
                    return;
                }

                Real wgt = xJac;
                Vec x;
                for (sdt j=0; j<dim; ++j)
                {
                    Real xrand = Uniform::nextStd(gen);
                    Real xn = (kg[j]-xrand)*dxg+1;
                    ia[j] = (xn<BinCount) ? xn : BinCount;
                    ia[j] = (ia[j]>1) ? ia[j] : 1;
                    Real xo, rc;
                    if (ia[j] > 1)
                    {
                        xo = xi[j][ia[j]-1]-xi[j][ia[j]-2];
                        rc = xi[j][ia[j]-2]+(xn-ia[j])*xo;
                    }
                    else
                    {
                        xo = xi[j][ia[j]-1];
                        rc = (xn-ia[j])*xo;
                    }
                    x[j] = lower[j]+rc*dx[j];
                    wgt *= xo*xnd;
                }
                VecRes f = func(const_cast<const Vec&>(x));   // call integrand at point x
                for (sdt j=0; j<dimRes; ++j)
                {
                    if (f[j] != 0) ++Ax[j].npg;
                    f[j] *= wgt;
                    Ax[j].f2 = f[j]*f[j];
                    Ax[j].fb += f[j];
                    Ax[j].f2b += Ax[j].f2;
                }
                for (sdt j=0; j<dim; ++j)
                {
                    di[ia[j]-1][j] += f[0];
                    if (mds >= 0) d[ia[j]-1][j] += Ax[0].f2;
                }
            } // end of loop within hypercube
            npgCur = -1;

            for (sdt j=0; j<dimRes; ++j)
            {
                Ax[j].f2b = Alge::sqrt(Ax[j].f2b*Ax[j].npg);
                Ax[j].f2b = (Ax[j].f2b-Ax[j].fb)*(Ax[j].f2b+Ax[j].fb);
                if (Ax[j].f2b <= 0) Ax[j].f2b = Real_::smallest;
                Ab[j].ti += Ax[j].fb;
                Ab[j].tsi += Ax[j].f2b;
            }
            if (mds < 0)
            {
                for (sdt j=0; j<dim; ++j) d[ia[j]-1][j] += Ax[0].f2b;
            }
            sdt k;
            for (k=dim-1; k>=0; --k)
            {
                kg[k] %= ng;
                if (++kg[k] != 1) break;
            }
            if (k < 0) break;
        }  // end of loop over hypercubes
        npgCur = -2;

        for (sdt j=0; j<dimRes; ++j)
        {
            Ab[j].tsi *= dv2g;
            Ai[j].Wgt = 1/Ab[j].tsi;
            Ai[j].sInt += Ai[j].Wgt*Ab[j].ti;
            Ai[j].sChi += Ai[j].Wgt*Ab[j].ti*Ab[j].ti;
            Ai[j].sWgt += Ai[j].Wgt;
            tgral[j] = Ai[j].sInt/Ai[j].sWgt;
            chi2a[j] = (Ai[j].sChi-Ai[j].sInt*tgral[j])/(ittot-0.9999);
            if (chi2a[j] < 0) chi2a[j] = 0;
            sd[j] = Alge::sqrt(1/Ai[j].sWgt);
            Ab[j].tsi = Alge::sqrt(Ab[j].tsi);
        }

        Real dt[dim];
        for (sdt j=0; j<dim; ++j)
        {
            Real xo = d[0][j];
            Real xn = d[1][j];
            d[0][j] = (xo+xn)/2;
            dt[j] = d[0][j];
            for (sdt i=1; i<nd-1; ++i)
            {
                Real rc = xo+xn;
                xo = xn;
                xn = d[i+1][j];
                d[i][j] = (rc+xn)/3;
                dt[j] += d[i][j];
            }
            d[nd-1][j] = (xo+xn)/2;
            dt[j] += d[nd-1][j];
        }
        for (sdt j=0; j<dim; ++j)
        {
            Real rc = 0;
            for (sdt i=0; i<nd; ++i)
            {
                if (d[i][j] < Real_::smallest) d[i][j] = Real_::smallest;
                r[i] =  Alge::pow((1.0-d[i][j]/dt[j])/
                        (Alge::log(dt[j])-Alge::log(d[i][j])),alpha);
                rc += r[i];
            }
            rebin(rc/xnd,nd,r,xin,xi[j]);
        }
    }
}

template<sdt Dim, sdt DimRes, class Real, sdt BinCount>
void Vegas<Dim, DimRes, Real, BinCount>::rebin(Real rc, sdt nd, Real r[], Real xin[], Real xi[])
{
    sdt i;
    sdt k = 0;
    Real dr = 0;
    Real xn = 0;
    Real xo = 0;

    for (i=0; i<nd-1; i++)
    {
        while (rc > dr)
            dr += r[k++];
        if (k > 1) xo = xi[k-2];
        xn = xi[k-1];
        dr -= rc;
        xin[i] = xn-(xn-xo)*dr/r[k-1];
    }
    for (i=0; i<nd-1; i++) xi[i] = xin[i];
    xi[nd-1] = 1;
}

}