Update to convergence criteria

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])

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()

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
+