diff --git a/MCcubed/mc/mcmc.py b/MCcubed/mc/mcmc.py index 19a15ab..f614444 100755 --- a/MCcubed/mc/mcmc.py +++ b/MCcubed/mc/mcmc.py @@ -33,7 +33,8 @@ def mcmc(data, uncert=None, func=None, indparams=[], leastsq=True, chisqscale=False, grtest=True, grexit=False, burnin=0, thinning=1, fgamma=1.0, fepsilon=0.0, plots=False, savefile=None, savemodel=None, comm=None, - resume=False, log=None, rms=False, hsize=1): + resume=False, log=None, rms=False, hsize=1, + griter=1000, confreg=0.95, confacc=0.02): """ This beautiful piece of code runs a Markov-chain Monte Carlo algoritm. @@ -116,6 +117,12 @@ def mcmc(data, uncert=None, func=None, indparams=[], File object to write log into. hsize: Int Initial samples for snooker walk. + griter: Int + Number of iterations between GR tests + confreg: Float + Confidence region for accuracy requirement + confacc: Float + Relative accuracy requirement for confreg Returns ------- @@ -218,8 +225,11 @@ def mcmc(data, uncert=None, func=None, indparams=[], hsize = nchains + 1 # Intermediate steps to run GR test and print progress report: - intsteps = chainsize / 10 - + if grtest: + intsteps = griter + else: + intsteps = chainsize / 10 + # Allocate arrays with variables: numaccept = np.zeros(nchains) # Number of accepted proposal jumps outbounds = np.zeros((nchains, nfree), np.int) # Out of bounds proposals @@ -432,7 +442,7 @@ def mcmc(data, uncert=None, func=None, indparams=[], nextp = np.copy(params) # Proposed parameters nextchisq = np.zeros(nchains) # Chi square of nextp - # Gelman-Rubin exit flag: + # Gelman-Rubin calculation flag: grflag = False mrfactor = np.zeros(nchains) @@ -567,22 +577,50 @@ def mcmc(data, uncert=None, func=None, indparams=[], format(bestchisq, str(bestp)), log) # Gelman-Rubin statistic: - if grtest and (i+nold) > burnin: - psrf = gr.convergetest(allparams[:, :, burnin:i+nold+1:thinning]) - mu.msg(1, "Gelman-Rubin statistic for free parameters:\n{:s}". - format(psrf), log) - if np.all(psrf < 1.01): + if (i+nold) > burnin: + if grtest: + psrf = gr.convergetest(allparams[:, :, burnin:i+nold+1:thinning]) + mu.msg(1, "Gelman-Rubin statistic for free parameters:\n{:s}". + format(psrf), log) + + if np.all(psrf < 1.01) and not grflag: mu.msg(1, "All parameters have converged to within 1% of unity.", log) - # End the MCMC if all parameters satisfy GR two consecutive times: - if grexit and grflag: - # Let the workers know that the MCMC is stopping: - if mpi: - endflag = np.tile(np.inf, nchains*mpars) - mu.comm_scatter(comm, endflag, MPI.DOUBLE) - break - grflag = True - else: - grflag = False + # Calculate ESS + allpac = np.zeros((len(ifree), i+nold)) # allparams autocorrelation + nisamp = np.zeros( len(ifree)) # iterations per sample + k = 0 + for p in ifree: + meanapi = np.mean(allparams[0,p,:i+nold]) + autocorr = np.correlate(allparams[0,p,:i+nold]-meanapi, + allparams[0,p,:i+nold]-meanapi, mode='full') + allpac[k] = autocorr[np.size(autocorr)//2:] / np.max(autocorr) + + if np.min(allpac[k]) < 0.01: + cutoff = np.where(allpac[k] < 0.01)[0][0] + else: + cutoff = -1 + + nisamp[k] = 1 + 2 * np.sum(allpac[k,:cutoff]) + k += 1 + + mu.msg(1, "Iterations per sample: {:f}".format(np.max(nisamp)), log) + ess = mu.cred2ess(confreg, confacc) + nineed = np.round(i+nold + ess * np.max(nisamp)) + mu.msg(1, "Total iter for {:.2f}% confidence" + " in a {:.2f}% region: {:.0f}". + format(confacc*(1-confreg)*100, confreg*100, nineed), log) + + grflag = True # Successfully passed GR test + grtest = False # Stop doing GR test for efficiency + + # End the MCMC if enough iterations have been run: + if grexit and grflag and (i+nold)*nchains > nineed: + # Let the workers know that the MCMC is stopping: + if mpi: + endflag = np.tile(np.inf, nchains*mpars) + mu.comm_scatter(comm, endflag, MPI.DOUBLE) + break + # Save current results: if savefile is not None: np.save(savefile, allparams[:,:,0:i+nold]) diff --git a/MCcubed/mccubed.py b/MCcubed/mccubed.py index 5f00845..9e17ecc 100755 --- a/MCcubed/mccubed.py +++ b/MCcubed/mccubed.py @@ -94,8 +94,24 @@ def main(): group.add_argument( "--grexit", dest="grexit", help="Exit the MCMC loop if the MCMC satisfies the GR " - "test two consecutive times [default: %(default)s]", + "test and achieves the desired confidence region " + "accuracy [default: %(default)s]", type=eval, action="store", default=False) + group.add_argument( "--griter", + dest="griter", + help="Number of iterations per GR test evaluation " + "[default: %(default)s]", + type=eval, action="store", default=1000) + group.add_argument( "--confreg", + dest="confreg", + help="Desired confidence region for accuracy test " + "[default: %(default)s]", + type=float, action="store", default=0.95) + group.add_argument( "--confacc", + dest="confacc", + help="Desired relative accuracy on the specified " + "confidence region [default: %(default)s]", + type=float, action="store", default=0.02) group.add_argument("-b", "--burnin", help="Number of burn-in iterations (per chain) " "[default: %(default)s]", @@ -232,6 +248,9 @@ def main(): chisqscale = args2.chisqscale grtest = args2.grtest grexit = args2.grexit + griter = args2.griter + confreg = args2.confreg + confacc = args2.confacc burnin = args2.burnin thinning = args2.thinning fgamma = args2.fgamma @@ -391,7 +410,7 @@ def main(): numit, nchains, walk, wlike, leastsq, chisqscale, grtest, grexit, burnin, thinning, fgamma, fepsilon, plots, savefile, savemodel, - comm, resume, log, rms, hsize) + comm, resume, log, rms, hsize, griter, confreg, confacc) if tracktime: stop = timeit.default_timer() diff --git a/MCcubed/utils/mcutils.py b/MCcubed/utils/mcutils.py index 6045e2f..2c0b355 100644 --- a/MCcubed/utils/mcutils.py +++ b/MCcubed/utils/mcutils.py @@ -3,10 +3,11 @@ __all__ = ["parray", "saveascii", "loadascii", "savebin", "loadbin", "comm_scatter", "comm_gather", "comm_bcast", "comm_disconnect", - "msg", "warning", "error", "progressbar", "sep"] + "msg", "warning", "error", "progressbar", "sep", "cred2ess"] import os, sys, time, traceback, textwrap, struct import numpy as np +import scipy.stats as ss try: from mpi4py import MPI @@ -410,3 +411,26 @@ def progressbar(frac, file=None): if file is not None: file.write(text + "\n") +def cred2ess(p, eps): + """ + Compute the Effective Sample Size (ESS) needed to compute + a credible region with size p (in range 0 to 1), targeting + a relative accuracy in 1-p of eps. + + Parameters + ---------- + p: Float + Credible region size, from 0 to 1. E.g., 0.95 would indicate + a 95% credible region. + eps: Float + Desired relative accuracy in 1-p. E.g., for p = 0.95, an + eps of 0.02 would equate to 0.02 * 0.05 = 0.1% accuracy. + + Returns + ------- + ess: Float + ESS needed to meet eps accuracy on 1-p. + """ + ess = 2*(ss.norm.ppf(0.5 * (1 - p)) / eps)**2 + return ess +