Compare autocorrelation times for parallel tempered and normal MCMC

mcmc.py

#!/bin/bash

'''
Script exported by refnx for analysing NR/XRR data.

To get help:

    python mcmc.py -h

You will need to install the following packages in a Python interpreter:
  - refnx
  - numpy
  - cython
  - matplotlib
  - scipy

If you wish to run the script on a cluster making use of MPI you'll need to
install:
  - mpi4py
  - schwimmbad

The usage then is:

    mpiexec -n 4 python mcmc.py <other options>
'''

import sys
import time
import argparse
from multiprocessing import Pool
import operator
import numpy as np
from numpy import array
try:
    import tqdm
    import schwimmbad
except ImportError:
    pass

from refnx.analysis import Objective, GlobalObjective, Transform, CurveFitter
from refnx.analysis import Parameter, Parameters, Interval, process_chain
from refnx.analysis import load_chain
from refnx.analysis.parameter import Constant, build_constraint_from_tree
from refnx.analysis.parameter import _BinaryOp, _UnaryOp
from refnx.dataset import ReflectDataset, Data1D

from refnx.reflect import Slab, SLD, Structure, Stack
from refnx.reflect import ReflectModel, LipidLeaflet, MixedReflectModel, Spline
from refnx._lib import flatten

import refnx
# Script created by refnx version: 0.1.7.dev0+4138220

def objective():
    objective_1 = Objective(ReflectModel(Structure(components=[Slab(Parameter(value=0.0, name='fronting - thick', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), SLD([Parameter(value=2.07, name='fronting - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='fronting - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='fronting'), Parameter(value=0.0, name='fronting - rough', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), name='fronting', vfsolv=Parameter(value=0.0, name='fronting - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)), Slab(Parameter(value=9.857694062804677, name='1 - thick', vary=True, bounds=Interval(lb=0, ub=30.0), constraint=None), SLD([Parameter(value=3.47, name='1 - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='1 - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='1'), Parameter(value=2.046970081753427, name='1 - rough', vary=True, bounds=Interval(lb=1.0, ub=3.0), constraint=None), name='sio2', vfsolv=Parameter(value=0.0, name='1 - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)), Slab(Parameter(value=211.90602151931475, name='thick', vary=True, bounds=Interval(lb=0, ub=400.0), constraint=None), SLD([Parameter(value=0.2022933000077, name='sld', vary=True, bounds=Interval(lb=0, ub=2.0), constraint=None), Parameter(value=0.0, name='isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name=''), Parameter(value=3.9322526932829645, name='rough', vary=True, bounds=Interval(lb=0, ub=4.0), constraint=None), name='polymer', vfsolv=Parameter(value=0.036066832341294784, name='vfsolv', vary=True, bounds=Interval(lb=0, ub=1.0), constraint=None)), Slab(Parameter(value=0.0, name='backing - thick', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), SLD([Parameter(value=6.36, name='backing - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='backing - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='backing'), Parameter(value=5.9840990714481395, name='backing - rough', vary=True, bounds=Interval(lb=0, ub=6.0), constraint=None), name='backing', vfsolv=Parameter(value=0.0, name='backing - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None))], name='', solvent=None, reverse_structure=False, contract=0), name='e361r', scale=Parameter(value=1.0210569516720702, name='scale', vary=True, bounds=Interval(lb=0, ub=2.0), constraint=None), bkg=Parameter(value=1.8379213279395567e-05, name='bkg', vary=True, bounds=Interval(lb=0, ub=4e-05), constraint=None), dq=Parameter(value=5.0, name='dq - resolution', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), threads=-1, quad_order=17), Data1D(data='/Users/andrew/Documents/Andy/programming/refnx/refnx/analysis/test/e361r.txt', mask=None), lnsigma=None, use_weights=True, transform=Transform('logY'), logp_extra=None, name='e361r')
    objective_2 = Objective(ReflectModel(Structure(components=[Slab(Parameter(value=0.0, name='fronting - thick', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), SLD([Parameter(value=2.07, name='fronting - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='fronting - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='fronting'), Parameter(value=0.0, name='fronting - rough', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), name='fronting', vfsolv=Parameter(value=0.0, name='fronting - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)), Slab(Parameter(value=9.857694062804677, name='1 - thick', vary=False, bounds=Interval(lb=0, ub=30.0), constraint=Parameter(value=9.857694062804677, name='1 - thick', vary=True, bounds=Interval(lb=0, ub=30.0), constraint=None)), SLD([Parameter(value=3.47, name='1 - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='1 - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='1'), Parameter(value=2.046970081753427, name='1 - rough', vary=False, bounds=Interval(lb=0, ub=6.0), constraint=Parameter(value=2.046970081753427, name='1 - rough', vary=True, bounds=Interval(lb=1.0, ub=3.0), constraint=None)), name='sio2', vfsolv=Parameter(value=0.0, name='1 - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)), Slab(Parameter(value=211.90602151931475, name='thick', vary=False, bounds=Interval(lb=0, ub=400.0), constraint=Parameter(value=211.90602151931475, name='thick', vary=True, bounds=Interval(lb=0, ub=400.0), constraint=None)), SLD([Parameter(value=0.2022933000077, name='sld', vary=False, bounds=Interval(lb=0, ub=2.0), constraint=Parameter(value=0.2022933000077, name='sld', vary=True, bounds=Interval(lb=0, ub=2.0), constraint=None)), Parameter(value=0.0, name='isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name=''), Parameter(value=3.9322526932829645, name='rough', vary=False, bounds=Interval(lb=0, ub=6.0), constraint=Parameter(value=3.9322526932829645, name='rough', vary=True, bounds=Interval(lb=0, ub=4.0), constraint=None)), name='polymer', vfsolv=Parameter(value=0.036066832341294784, name='vfsolv', vary=False, bounds=Interval(lb=0, ub=1.0), constraint=Parameter(value=0.036066832341294784, name='vfsolv', vary=True, bounds=Interval(lb=0, ub=1.0), constraint=None))), Slab(Parameter(value=0.0, name='backing - thick', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), SLD([Parameter(value=3.47, name='backing - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='backing - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='backing'), Parameter(value=5.9840990714481395, name='backing - rough', vary=False, bounds=Interval(lb=0, ub=6.0), constraint=Parameter(value=5.9840990714481395, name='backing - rough', vary=True, bounds=Interval(lb=0, ub=6.0), constraint=None)), name='backing', vfsolv=Parameter(value=0.0, name='backing - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None))], name='', solvent=None, reverse_structure=False, contract=0), name='e365r', scale=Parameter(value=0.9644593548185179, name='scale', vary=True, bounds=Interval(lb=0, ub=2.0), constraint=None), bkg=Parameter(value=2.174885298686226e-05, name='bkg', vary=True, bounds=Interval(lb=0, ub=4e-05), constraint=None), dq=Parameter(value=5.0, name='dq - resolution', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), threads=-1, quad_order=17), Data1D(data='/Users/andrew/Documents/Andy/programming/refnx/refnx/analysis/test/e365r.txt', mask=None), lnsigma=None, use_weights=True, transform=Transform('logY'), logp_extra=None, name='e365r')
    objective_3 = Objective(ReflectModel(Structure(components=[Slab(Parameter(value=0.0, name='fronting - thick', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), SLD([Parameter(value=2.07, name='fronting - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='fronting - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='fronting'), Parameter(value=0.0, name='fronting - rough', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), name='fronting', vfsolv=Parameter(value=0.0, name='fronting - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)), Slab(Parameter(value=9.857694062804677, name='1 - thick', vary=False, bounds=Interval(lb=0, ub=30.0), constraint=Parameter(value=9.857694062804677, name='1 - thick', vary=True, bounds=Interval(lb=0, ub=30.0), constraint=None)), SLD([Parameter(value=3.47, name='1 - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='1 - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='1'), Parameter(value=2.046970081753427, name='1 - rough', vary=False, bounds=Interval(lb=0, ub=6.0), constraint=Parameter(value=2.046970081753427, name='1 - rough', vary=True, bounds=Interval(lb=1.0, ub=3.0), constraint=None)), name='sio2', vfsolv=Parameter(value=0.0, name='1 - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)), Slab(Parameter(value=211.90602151931475, name='thick', vary=False, bounds=Interval(lb=0, ub=400.0), constraint=Parameter(value=211.90602151931475, name='thick', vary=True, bounds=Interval(lb=0, ub=400.0), constraint=None)), SLD([Parameter(value=0.2022933000077, name='sld', vary=False, bounds=Interval(lb=0, ub=2.0), constraint=Parameter(value=0.2022933000077, name='sld', vary=True, bounds=Interval(lb=0, ub=2.0), constraint=None)), Parameter(value=0.0, name='isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name=''), Parameter(value=3.9322526932829645, name='rough', vary=False, bounds=Interval(lb=0, ub=6.0), constraint=Parameter(value=3.9322526932829645, name='rough', vary=True, bounds=Interval(lb=0, ub=4.0), constraint=None)), name='polymer', vfsolv=Parameter(value=0.036066832341294784, name='vfsolv', vary=False, bounds=Interval(lb=0, ub=1.0), constraint=Parameter(value=0.036066832341294784, name='vfsolv', vary=True, bounds=Interval(lb=0, ub=1.0), constraint=None))), Slab(Parameter(value=0.0, name='backing - thick', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), SLD([Parameter(value=-0.56, name='backing - sld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), Parameter(value=0.0, name='backing - isld', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None)], name='backing'), Parameter(value=5.9840990714481395, name='backing - rough', vary=False, bounds=Interval(lb=0, ub=6.0), constraint=Parameter(value=5.9840990714481395, name='backing - rough', vary=True, bounds=Interval(lb=0, ub=6.0), constraint=None)), name='backing', vfsolv=Parameter(value=0.0, name='backing - volfrac solvent', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None))], name='', solvent=None, reverse_structure=False, contract=0), name='e366r', scale=Parameter(value=0.9049357307467163, name='scale', vary=True, bounds=Interval(lb=0, ub=2.0), constraint=None), bkg=Parameter(value=2.1704045071153445e-05, name='bkg', vary=True, bounds=Interval(lb=0, ub=4e-05), constraint=None), dq=Parameter(value=5.0, name='dq - resolution', vary=False, bounds=Interval(lb=-np.inf, ub=np.inf), constraint=None), threads=-1, quad_order=17), Data1D(data='/Users/andrew/Documents/Andy/programming/refnx/refnx/analysis/test/e366r.txt', mask=None), lnsigma=None, use_weights=True, transform=Transform('logY'), logp_extra=None, name='e366r')
    objective_0 = GlobalObjective([objective_1, objective_2, objective_3, ])
    parameters = list(flatten(objective_0.parameters))
    parameters[31].constraint = build_constraint_from_tree([parameters[8]])
    parameters[34].constraint = build_constraint_from_tree([parameters[11]])
    parameters[36].constraint = build_constraint_from_tree([parameters[13]])
    parameters[37].constraint = build_constraint_from_tree([parameters[14]])
    parameters[39].constraint = build_constraint_from_tree([parameters[16]])
    parameters[40].constraint = build_constraint_from_tree([parameters[17]])
    parameters[44].constraint = build_constraint_from_tree([parameters[21]])
    parameters[54].constraint = build_constraint_from_tree([parameters[8]])
    parameters[57].constraint = build_constraint_from_tree([parameters[11]])
    parameters[59].constraint = build_constraint_from_tree([parameters[13]])
    parameters[60].constraint = build_constraint_from_tree([parameters[14]])
    parameters[62].constraint = build_constraint_from_tree([parameters[16]])
    parameters[63].constraint = build_constraint_from_tree([parameters[17]])
    parameters[67].constraint = build_constraint_from_tree([parameters[21]])
    return objective_0


def structure_plot(obj, samples=0):
    # plot sld profiles
    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(111)

    if isinstance(obj, GlobalObjective):
        if samples > 0:
            savedparams = np.array(obj.parameters)
            for pvec in obj.parameters.pgen(ngen=samples):
                obj.setp(pvec)
                for o in obj.objectives:
                    if hasattr(o.model, 'structure'):
                        ax.plot(*o.model.structure.sld_profile(),
                                color="k", alpha=0.01)

            # put back saved_params
            obj.setp(savedparams)

        for o in obj.objectives:
            if hasattr(o.model, 'structure'):
                ax.plot(*o.model.structure.sld_profile(), zorder=20)

        ax.set_ylabel('SLD / $10^{-6}\\AA^{-2}$')
        ax.set_xlabel("z / $\\AA$")

    elif isinstance(obj, Objective) and hasattr(obj.model, 'structure'):
        fig, ax = obj.model.structure.plot(samples=samples)

    fig.savefig('steps_sld.png', dpi=1000)


def main(args):
    nwalkers = args.walkers
    nthin = args.thin
    nsteps = args.steps
    ntemps = args.temps
    nplot = args.plot
    nburn = args.burn
    cfile = args.chain
    verbose = args.verbose
    pargs = ()

    try:
        # we may want to use MPI to parallelise
        import schwimmbad
        if args.mpi:
            pool_klass = schwimmbad.MPIPool
    except ImportError:
        pass

    # not wanting MPI, so just use multiprocessing.Pool to parallelise
    if not args.mpi:
        pool_klass = Pool
        pargs = (args.n_cores,)

    with pool_klass(*pargs) as workers:
        # necessary when using MPI.
        if args.mpi and (not workers.is_master()):
            workers.wait()
            sys.exit(0)

        start_time = time.time()

        obj = objective()

        # turn off pthread'ing of reflectivity calculation if MPI. Otherwise
        # the reflectivity calculation will want to spread out over all the
        # available processors.
        if args.mpi:
            _objectives = [obj]
            if isinstance(obj, GlobalObjective):
                _objectives = obj.objectives

            for o in _objectives:
                o.model.threads = 1

        # Create the fitter and fit
        fitter = CurveFitter(obj, nwalkers=nwalkers, ntemps=ntemps)

        if nsteps:
            if cfile:
                # Initialise the walkers with a pre-existing chain
                chain = load_chain(cfile)
                fitter.initialise(chain)
            else:
                # Initialise the walkers by doing a fit, then using the
                # covariance. The workers kwd is only present in scipy >1.2
                fitter.fit('differential_evolution', workers=workers.map)
                fitter.initialise('covar')

            # Buffering is there so the chain file is not written to
            # continuously
            with open('steps.chain', 'w', buffering=500000) as f:
                res = fitter.sample(nsteps, pool=workers.map, f=f,
                                    verbose=verbose, nthin=nthin);
                f.flush()
            process_chain(obj, fitter.chain, nburn=nburn)
        else:
            # the workers kwd is only present in scipy >1.2
            fitter.fit('differential_evolution', workers=workers.map)

        print(str(obj))
        print('\n')
        print('Duration (s): {}'.format(time.time() - start_time))

        try:
            # create graphs of reflectivity and SLD profiles
            import matplotlib
            import matplotlib.pyplot as plt
            matplotlib.use('agg')

            fig, ax = obj.plot(samples=nplot)
            ax.set_ylabel('R')
            ax.set_xlabel("Q / $\\AA$")
            fig.savefig('steps.png', dpi=1000)

            structure_plot(obj, samples=nplot)

            # corner plot
            fig = obj.corner()
            fig.savefig('steps_corner.png')

            # plot the Autocorrelation function of the chain
            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.plot(fitter.acf())
            ax.set_ylabel('autocorrelation')
            ax.set_xlabel('step')
            fig.savefig('steps-autocorrelation.png')
        except ImportError:
            pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-w', '--walkers', help='number of emcee walkers',
                        type=int, default=200)
    parser.add_argument('-t', '--thin', help='factor to thin chain by',
                        type=int, default=1)
    parser.add_argument('-s', '--steps', help=("number of thinned MCMC steps"
                                               " to save"),
                        type=int, default=1000)
    parser.add_argument('-b', '--burn', help=("number of initial MCMC steps"
                                               " to discard"),
                        type=int, default=0)
    parser.add_argument('-n', '--temps', help=("number of parallel tempering"
                                               " temperatures (requires the"
                                               " ptemcee package)"),
                        type=int, default=-1)
    parser.add_argument('-p', '--plot', help=("create plots of the MCMC"
                                              " using 'plot' samples"),
                        type=int, default=0),
    parser.add_argument('-c', '--chain', help='initialise chain from file',
                        type=str, default='')
    parser.add_argument('-v', '--verbose', help="Displays a progress bar while"
                        " sampling",
                        dest='verbose', action='store_true',  default=False)

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--ncores", dest="n_cores", default=-1,
                       type=int, help=("Number of processes (uses"
                       "multiprocessing)."))
    group.add_argument("--mpi", dest="mpi", default=False,
                       action="store_true", help="Run with MPI.")

    args = parser.parse_args()

    if args.n_cores == -1:
        args.n_cores = None
    main(args)

PTMCMCvsMCMC.ipynb