How to use the emcee.utils.MPIPool function in emcee

logging.debug("Processing IGS1 map {0}".format(filename))
	conv_map = ConvergenceMap.load(filename)

	#Smooth 1 arcmin
	conv_map.smooth(1.0*arcmin,inplace=True)

	#Measure the moments
	return conv_map.moments(connected=True,dimensionless=True)




logging.basicConfig(level=logging.DEBUG)

try: 
	pool = MPIPool()
except ValueError:
	pool = None

if (pool is not None) and not(pool.is_master()):

	pool.wait()
	sys.exit(0)



map_mock_ids = range(int(sys.argv[1]))

igs1_set = IGS1(root_path="/Users/andreapetri/Documents/Columbia/spurious_shear/convergence_maps")
map_igs1_ids = igs1_set.getNames(z=1.0,realizations=range(1,int(sys.argv[1])+1))

gen = GaussianNoiseGenerator(shape=(2048,2048),side_angle=3.41*deg,label="convergence")

2. We need to figure out which parameters need to be included in the chain,
       which ones are automatically fitted and which ones are automatically
       derived
    3. Initiate the walkers. This can be done in three ways:
        - starting from the previous run
        - starting from posteriors
        - starting from priors
       This is governed by the "init_from" parameter, and the list above is also
       the order of priority that is used when the init_from option is invalid:
       if there is no previous run, it will be checked if we can start from
       posteriors. If there are no posteriors, we'll start from the priors.
    4. Create the sampler and run!    
    """
    # Take care of the MPI pool
    if mpi:
        pool = MPIPool()
        if not pool.is_master():
            pool.wait()
            sys.exit(0)
    else:
        pool = None
    
    # +---------------------------------------+
    # |   STEP 1: load system and parameters  |
    # +---------------------------------------+
    
    # Load the System
    system = universe.load(system_file)
    system.reset_and_clear()
    
    # Load the fit and compute ParameterSets
    compute_params = universe.load(compute_params_file)

#Initialize model and data objects in global namespace
myData = Data(wl, fl)
myModel = LineModelFlag(wl, np.array([1.5, 1.5]))

#Initialize Sampler object
mySamp = Sampler(myData, myModel)

#Create lnprob in global scope
def lnprob(param):
    myModel.set_params(param)
    good = np.sum((myData.y - myModel.y)[myModel.flags]**2)
    bad = np.sum((myData.y - myModel.y)[~myModel.flags]**2)
    return good + bad

pool = MPIPool()

if not pool.is_master():
#   Wait for instructions from the master process.
    pool.wait()
    sys.exit(0) #this is at the very end of the run.

# Initialize the sampler with the chosen specs.
nwalkers = 20
burn_in = 100
ndim = 2
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, pool=pool)

#Declare starting indexes
m = np.random.uniform(low=0, high=3, size=(nwalkers,))
b = np.random.uniform(low=0, high=3, size=(nwalkers,))

def main():

    if config['MPI']:
        from emcee.utils import MPIPool
        # Initialize the MPI-based pool used for parallel run.
        pool = MPIPool()
        print("Running with MPI")

        if not pool.is_master():
        #   Wait for instructions from the master process.
            pool.wait()
            sys.exit(0) #this is at the very end of the run.

        # Initialize the sampler with the chosen specs.
        sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, pool=pool)

    else:
        sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, threads=config['threads'])

    # Create necessary output directories, if they do not already exist
    if not os.path.exists(outdir):
        os.makedirs(outdir, exist_ok=True)

f = open(status_file, "w")
pickle_file = 'initial_pos.pkl'

f.write("#####################################################################################################\n")
f.write("##   loadbalance    runtime_sorting   iteration       mean_time     variance      ideal       actual \n")
f.write("#####################################################################################################\n")

for mean_time in mean_times:
    for variance_fac in xrange(len(variances)):
        first = 0
        variance = variances[variance_fac]*mean_time
        for loadbalance in loadbalancing_options:
            for runtime_sorting_option in runtime_sorting_options:

                # Initialize the MPI-based pool used for parallelization.
                pool = MPIPool(loadbalance=loadbalance)

                if not pool.is_master():
                    # Wait for instructions from the master process.
                    pool.wait()
                    sys.exit(0)

                # Initialize the sampler with the chosen specs.
                sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, pool=pool, runtime_sortingfn=runtime_sorting_option)

                tstart = time.time()
                print("Before running the iterations. loadbalance = {0}".format(loadbalance))
                if first == 0:
                    # Generate random positions (sleep times for lnprob) for the first time
                    p0 = [(mean_time + variance * np.random.randn(ndim)) for col in range(nwalkers)]
                    pos, prob, rstate = sampler.run_mcmc(p0, 1)
                    pkl_file = open(pickle_file, 'wb')

"""

from __future__ import print_function

import sys
import numpy as np
import emcee
from emcee.utils import MPIPool


def lnprob(x):
    return -0.5 * np.sum(x ** 2)

# Initialize the MPI-based pool used for parallelization.
pool = MPIPool(debug=True)

if not pool.is_master():
    # Wait for instructions from the master process.
    pool.wait()
    print("DONE")
    sys.exit(0)

ndim = 2
nwalkers = 6
p0 = [np.random.rand(ndim) for i in range(nwalkers)]

# Initialize the sampler with the chosen specs.
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, pool=pool)

# Run 100 steps as a burn-in.
pos, prob, state = sampler.run_mcmc(p0, 10)

Run this example with:
mpirun -np 2 python example_mpi_pool.py

"""


from __future__ import print_function


import numpy as np
import emcee

from emcee.utils import MPIPool


pool = MPIPool(debug=False)
nwalkers = 50
ndim = 10
p0 = np.random.rand(nwalkers, ndim)


def log_prob(p):
    #A trivial Gaussian
    return -0.5 * (p ** 2).sum()


if pool.is_master():
    sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, pool=pool)
    for sample in sampler.sample(p0, iterations=100):
        print(sample[0])
    pool.close()
else:

smoothed_map = convergence_map.smooth(smoothing_scales[n])
		v,hist_output[descriptor.first:descriptor.last] = smoothed_map.pdf(bin_edges)

	#Return the histograms in array format
	return hist_output

############################################################
########################Main loop###########################
############################################################

if __name__=="__main__":
	
	#Initialize MPI pool
	try: 
		pool = MPIPool()
	except ValueError:
		pool = None

	#Parse command line arguments
	parser = argparse.ArgumentParser()
	parser.add_argument("-v","--verbose",action="store_true",default=False,dest="verbose",help="Display degug info")
	parser.add_argument("-p","--path",action="store",dest="path",default="/Users/andreapetri/Documents/Columbia/spurious_shear/convergence_maps",help="Root path of IGS1 simulations")
	parser.add_argument("-n","--num_realizations",dest="num_realizations",action="store",type=int,default=3,help="How many realizations to process")

	cmd_args = parser.parse_args()

	#Set logging level
	if cmd_args.verbose:
		logging.basicConfig(level=logging.DEBUG)
	else:
		logging.basicConfig(level=logging.INFO)