def sampler_emcee(D):
# nlinks: number of iterations for each walker
# please specify as even
D['nlinks'] = 100
# nwalkers: number of interacting chains
D['nwalkers'] = 200
# ntemps:
D['ntemps'] = 5
# ntune: number of initialization steps to discard
D['ntune'] = 500
st = time.time()
# identify starting positions for chains
tmp = start_pos(D['ntemps'] * D['nwalkers'], D)
pos0 = tmp.reshape((D['ntemps'], D['nwalkers'], D['dim']))
likelihood = D['likelihood']
# run MCMC for the model of interest
sampler = PTSampler(ntemps=D['ntemps'],
nwalkers=D['nwalkers'],
dim=len(D['pname']),
logl=likelihood,
logp=prior,
loglargs=(D, ),
logpargs=(D, ))
for pos, lnprob, lnlike in sampler.sample(pos0, iterations=D['ntune']):
pass
burntrace = sampler.chain
sampler.reset()
print(burntrace.shape)
print("burnin completed")
for pos, lnprob, lnlike in sampler.sample(pos, iterations=D['nlinks']):
pass
finaltrace = sampler.chain
print(finaltrace.shape)
D['rawtrace'] = np.concatenate((burntrace[0, ...], finaltrace[0, ...]),
axis=1)
D['sampling_time'] = np.round(time.time() - st, 2)
st = time.time()
D['lnZ'], D['dlnZ'] = sampler.thermodynamic_integration_log_evidence()
elaps = np.round(time.time() - st, 2)
msg = "model evidence evaluation: " + str(elaps) + " seconds"
WP(msg, D['wrt_file'])
return D
covariance = sigma_diag
return covariance
########################################################
ntemps = 2
nwalkers = 50
ndim = 5
sampler = PTSampler(ntemps,
nwalkers,
ndim,
lnprobabmatrix,
lnprior,
loglkwargs={
'obs': obs,
'image_covariance': image_covariance,
'sed_covariance': sed_covariance
})
#sampler=EnsembleSampler(nwalkers,ndim,lnprobabmatrix,lnprior,args)
# Use the MPIPool instead
"""
pool = MPIPool()
if not pool.is_master():
pool.wait()
sys.exit(0)
"""
#sampler = PTSampler(ntemps, nwalkers, ndim, lnprobab, lnprior, threads=15)
#############################################################
#subprocess.call('rm -r data_0.8',shell=True)
os.chdir(olddir)
#STEP 5:
return imageuncertainty, imagechi, seduncertainty.value, sedchi
########################################################
ntemps = 20
nwalkers = 4096
ndim = 7
sampler=PTSampler(ntemps, nwalkers, ndim, lnprobab, lnprior, threads=4)
#############################################################
# Initialize the walkers. The best technique seems to be
# to start in a small ball around the a priori preferred position.
# Dont worry, the walkers quickly branch out and explore the
# rest of the space.
# inclination: w0
# scale height: w1
# dust mass: w2
# max grain size: w3
# alpha: w4
# beta: w5
# weights: w6
#Numnber of iterations (after the burn-in)
niter = 1500000
num_files=50 #I trust YOU, user, to make sure niter / num_files is remainderless
nperfile=niter/num_files
thin=1
#Burn-in for 1/5 of the number of iterations
nburn = 50000
if __name__=='__main__':
print 'Beginning main function'
time.clock()
time.sleep(1)
startTime = datetime.now()
sampler=PTSampler(ntemps, nwalkers, ndim, logl, logp ,loglargs=[eps,x,y,sigx,sigy],logpargs=[eps,x,y,sigx,sigy],threads=32)
sampler.run_mcmc(p0,nburn,thin=thin)
for index in range(num_files):
if index == 0:
p = sampler.chain[:,:,-1,:]
np.savez_compressed(savename+'_burn_'+str(index),af=sampler.acceptance_fraction[0],
chain=sampler.chain[0],lnp=sampler.lnprobability[0],p=sampler.chain[:,:,-1,:])
print 'Burn in complete'
#burnchain = sampler.chain
#pdb.set_trace()
else:
p = sampler.chain[:,:,-1,:]
sampler.reset()
#print 'Burn in complete'
sampler.run_mcmc(p,nperfile,thin=thin)
#print 'orbit fitting complete'
niter = 3000000
num_files = 100 #I trust YOU, user, to make sure niter / num_files is remainderless
nperfile = niter / num_files
thin = 1
#Burn-in for 1/5 of the number of iterations
nburn = 50000
print 'Made it this far, waiting 5 secs'
if __name__ == '__main__':
time.clock()
time.sleep(1)
startTime = datetime.now()
sampler = PTSampler(ntemps,
nwalkers,
ndim,
logl,
logp,
loglargs=[eps, xy, sigs],
threads=16)
sampler.run_mcmc(p0, nburn, thin=thin)
for index in range(num_files):
if index == 0:
p = sampler.chain[:, :, -1, :]
np.savez_compressed(savename + '_burn_' + str(index),
af=sampler.acceptance_fraction[0],
chain=sampler.chain[0],
lnp=sampler.lnprobability[0],
plast=p)
print 'Burn in complete'
#burnchain = sampler.chain
#pdb.set_trace()
def MCMC(self, *args):
if args:
#kplan, mod_lims, ins_lims, acc_lims, sigmas_raw, pos0, logl, logp
kplanets, boundaries, inslims = args[0], args[1], args[2]
acc_lims, sigmas_raw, pos0 = args[3], args[4], args[5]
logl, logp = args[6], args[7]
ndim = 1 + 5 * kplanets + self.nins*2*(self.MOAV+1) + self.totcornum + self.PACC
print(str(self.PM)), 'self.pm!!' # PMPMPM
if kplanets > 0:
if self.PM:
pm_lims = args[8]
ndim += self.lenppm*self.fsig
print('checkpoint 1') # PMPMPM
ndat = len(self.time)
def starinfo():
colors = ['red', 'green', 'blue', 'yellow', 'grey', 'magenta', 'cyan', 'white']
c = sp.random.randint(0,7)
print(colored('\n ###############################################', colors[c]))
print(colored(' # #', colors[c]))
print(colored(' # #', colors[c]))
print(colored(' # E M P E R 0 R #', colors[c]))
print(colored(' # #', colors[c]))
print(colored(' # #', colors[c]))
print(colored(' ###############################################', colors[c]))
print(colored('Exoplanet Mcmc Parallel tEmpering Radial vel0city fitteR', colors[sp.random.randint(0,7)]))
logdat = '\n\nStar Name : '+self.starname
logdat += '\nTemperatures, Walkers, Steps : '+str((self.ntemps, self.nwalkers, self.nsteps))
logdat += '\nN Instruments, K planets, N data : '+str((self.nins, kplanets, self.ndat))
if self.PM:
logdat += '\nN of data for Photometry : '+str(self.ndat_pm)
logdat += '\nN Number of Dimensions : '+str(ndim)
logdat += '\nN Moving Average : '+str(self.MOAV)
logdat += '\nBeta Detail : '+str(self.betas)
logdat += '\n-----------------------------------------------------'
print(logdat)
pass
starinfo()
#'''
#from emperors_library import logp_rv
print(str(self.PM), ndim, 'self.pm y ndim') # PMPMPM
if self.PM:
if kplanets > 0:
logp_params = sp.array([sp.array([self.time, kplanets, self.nins, self.MOAV,
self.totcornum, boundaries, inslims, acc_lims,
sigmas_raw, self.eccprior, self.jittprior,
self.jittmean, self.STARMASS, self.HILL,
self.PACC, self.CHECK]),
sp.array([self.time_pm, self.fsig, self.lenppm,
self.nins_pm, self.MOAV_pm,
self.totcornum_pm, boundaries, sigmas_raw,
self.PACC_pm])])
logl_params = sp.array([sp.array([self.time, self.rv, self.err, self.ins,
self.staract, self.starflag, kplanets, self.nins,
self.MOAV, self.totcornum, self.PACC]),
sp.array([self.time_pm, self.rv_pm, self.err_pm, self.ins_pm,
self.staract_pm, self.starflag_pm, self.fsig,
self.f2k, self.nins_pm, self.MOAV_pm,
self.totcornum_pm, self.PACC_pm, kplanets])])
self.sampler = PTSampler(self.ntemps, self.nwalkers, ndim, logl, logp,
loglargs=[empmir.logl_rvpm, logl_params],
logpargs=[empmir.logp_rvpm, logp_params],
threads=self.cores, betas=self.betas)
#raise ImportError('xd dale al debug mejor')
else:
logp_params = sp.array([self.time, kplanets, self.nins, self.MOAV,
self.totcornum, boundaries, inslims, acc_lims,
sigmas_raw, self.eccprior, self.jittprior,
self.jittmean, self.STARMASS, self.HILL,
self.PACC, self.CHECK])
logl_params = sp.array([self.time, self.rv, self.err, self.ins,
self.staract, self.starflag, kplanets, self.nins,
self.MOAV, self.totcornum, self.PACC])
self.sampler = PTSampler(self.ntemps, self.nwalkers, ndim, logl, logp,
loglargs=[empmir.logl_rv, logl_params],
logpargs=[empmir.logp_rv, logp_params],
threads=self.cores, betas=self.betas)
# raise ImportError
else:
logp_params = sp.array([self.time, kplanets, self.nins, self.MOAV,
self.totcornum, boundaries, inslims, acc_lims,
sigmas_raw, self.eccprior, self.jittprior,
self.jittmean, self.STARMASS, self.HILL,
self.PACC, self.CHECK])
logl_params = sp.array([self.time, self.rv, self.err, self.ins,
self.staract, self.starflag, kplanets, self.nins,
self.MOAV, self.totcornum, self.PACC])
self.sampler = PTSampler(self.ntemps, self.nwalkers, ndim, logl, logp,
loglargs=[empmir.logl_rv, logl_params],
logpargs=[empmir.logp_rv, logp_params],
threads=self.cores, betas=self.betas)
# RVPM THINGY
s0 = chrono.time()
for _ in range(10000):
empmir.logp_rv(pos0[0][0], logp_params)
print('________', chrono.time()-s0)
print('\n --------------------- BURN IN --------------------- \n')
pbar = tqdm(total=self.burn_out)
for p, lnprob, lnlike in self.sampler.sample(pos0, iterations=self.burn_out):
pbar.update(1)
pass
pbar.close()
p0, lnprob0, lnlike0 = p, lnprob, lnlike
print("\nMean acceptance fraction: {0:.3f}".format(sp.mean(self.sampler.acceptance_fraction)))
assert sp.mean(self.sampler.acceptance_fraction) != 0, 'Mean acceptance fraction = 0 ! ! !'
self.sampler.reset()
print('\n ---------------------- CHAIN ---------------------- \n')
pbar = tqdm(total=self.nsteps)
for p, lnprob, lnlike in self.sampler.sample(p0, lnprob0=lnprob0,
lnlike0=lnlike0,
iterations=self.nsteps,
thin=self.thin):
pbar.update(1)
pass
pbar.close()
#'''
assert self.sampler.chain.shape == (self.ntemps, self.nwalkers, self.nsteps/self.thin, ndim), 'something really weird happened'
print("\nMean acceptance fraction: {0:.3f}".format(sp.mean(self.sampler.acceptance_fraction)))
ln_post = self.sampler.lnprobability
posteriors = sp.array([ln_post[i].reshape(-1) for i in range(self.ntemps)])
chains = self.sampler.flatchain
best_post = posteriors[0] == np.max(posteriors[0])
#raise ImportError
thetas_raw = sp.array([chains[i] for i in range(self.ntemps)])
thetas_hen = sp.array([empmir.henshin(chains[i], kplanets) for i in sp.arange(self.ntemps)])
ajuste_hen = thetas_hen[0][best_post][0]
ajuste_raw = thetas_raw[0][best_post][0]
interesting_loc = sp.array([max(posteriors[temp]) - posteriors[temp] < self.bayes_factor for temp in sp.arange(self.ntemps)])
interesting_thetas = sp.array([thetas_hen[temp][interesting_loc[temp]] for temp in sp.arange(self.ntemps)])
thetas_hen = sp.array([thetas_hen[temp] for temp in sp.arange(self.ntemps)])
interesting_thetas_raw = sp.array([thetas_raw[temp][interesting_loc[temp]] for temp in sp.arange(self.ntemps)])
interesting_posts = sp.array([posteriors[temp][interesting_loc[temp]] for temp in range(self.ntemps)])
sigmas = sp.array([ sp.std(interesting_thetas[0][:, i]) for i in range(ndim) ])
sigmas_raw = sp.array([ sp.std(interesting_thetas_raw[0][:, i]) for i in range(ndim) ])
#print('sigmas', sigmas) # for testing
#print('sigmas_raw', sigmas_raw)
#print('mod_lims', boundaries)
print('ALL RIGHT ALL RIGHT ALL RIGHT ALL RIGHT ALL RIGHT ALL RIGHT ALL RIGHT ALL RIGHT ')
return thetas_raw, ajuste_raw, thetas_hen, ajuste_hen, p, lnprob, lnlike, posteriors, self.sampler.betas, interesting_thetas, interesting_posts, sigmas, sigmas_raw
# force close all model images:
model.images.closeimage()
#STEP 5:
return imageuncertainty, imagechi
########################################################
ntemps = 1
nwalkers = 10
ndim = 9
sampler=PTSampler(ntemps, nwalkers, ndim, lnprobab, lnprior)
# Use the MPIPool instead
"""
pool = MPIPool()
if not pool.is_master():
pool.wait()
sys.exit(0)
"""
#sampler = PTSampler(ntemps, nwalkers, ndim, lnprobab, lnprior, threads=15)
#############################################################
# Initialize the walkers. The best technique seems to be
# to start in a small ball around the a priori preferred position.
# Dont worry, the walkers quickly branch out and explore the
# rest of the space.
nwalkers = 100; # Number of iterations. niter = 1000; # Of which burn-in is the first burnin = 500; # Initial spot parameters. # [spotx, spoty, spotradius, spotcontrast] spot = numpy.array([0.204, 0.376, 0.096, 0.524]); # Create 3D matrix for initial state for each temperature and walker. p0 = numpy.repeat(spot[:,numpy.newaxis].T, ntemps*nwalkers, axis=0).reshape(ntemps, nwalkers, ndim); # Randomize the initial states in a small neighborhood. p0 += numpy.random.normal(scale=1e-3, size=p0.shape); # Initialize sampler. sampler = PTSampler(ntemps, nwalkers, ndim, logl, logp); # Run sampler. pos, prob, state = sampler.run_mcmc(p0, niter); # Take a view of the T=0 chain. zerotemp = sampler.chain[0]; # We take iterations at T=0 after burn-in as equilibrium # distribution. With a 100 walkers, this is 1e4 points. eq = zerotemp[:,burnin:,:].reshape([nwalkers*(niter-burnin), ndim]); # Plot distribution of every possible pairs. labels = ["spotx", "spoty", "spotradius", "spotcontrast"]; for ploti in range(ndim-1): for plotj in range(ploti+1,ndim):
def logp(x):
return 0.0
# Training Parameters
sigma = 1e-6
ntemp = 10
nruns = 2000
Temp_i = 0
#initialize the chian with a=0, b=0, c=0.5
pos = np.tile((0,0,0.5),5)/10+1e-4*np.random.randn(ntemp,64, 15)
ntemps, nwalkers, ndim = pos.shape
#first MCMC chain
sampler = PTSampler(ntemps,nwalkers, ndim, log_probability,logp, loglargs=(x, JV_exp, sigma))
sampler.run_mcmc(pos, nruns )
samples = sampler.chain
#%%
#use the values obtained in the first MCMC chain to update the inistal estimate
pos_update = samples[:,:,-1,:]+1e-5*np.random.randn(ntemp,64, 15)
sampler.reset()
#second MCM chain
sampler = PTSampler(ntemps,nwalkers, ndim, log_probability,logp, loglargs=(x, JV_exp, sigma))
sampler.run_mcmc(pos_update, nruns);
flat_samples = sampler.flatchain
zero_flat_samples = flat_samples[Temp_i,:,:]
zero_samples = samples[Temp_i,:,:,:]
#visulize a1
plt.figure()
#sampler = emcee.EnsembleSampler(nwalkers, ndim, lnlike, args=(Geom,wl_obs,f_obs,f_err,gauss_width,logZ,scaling,shift,nspec),threads=nwalkers)
if use_ptsampler:
# use a flat prior
def logp(x):
return 0.0
ntemps = 20
#SAMPLER AND SEEDS ARE DEFINED ACCORDING TO THE MCMC PARAMETERS CHOSEN
if parameters == ['NH', 'Vmax']:
pos = [[NH[i], V[i]] for i in range(nwalkers)]
sampler = PTSampler(ntemps,
nwalkers,
ndim,
lneqnlike,
logp,
loglargs=(Geom, wl_obs, f_obs, f_err, scaling,
shift, nspec)) #,pool=pool)
elif parameters == ['NH', 'Vmax', 'Z']:
pos = [[NH[i], V[i], Z[i]] for i in range(nwalkers)]
sampler = PTSampler(ntemps,
nwalkers,
ndim,
lnlikeZ,
logp,
loglargs=(Geom, wl_obs, f_obs, f_err,
gauss_width, theta, scaling, shift,
nspec),
pool=pool)
elif parameters == ['NH', 'Vmax', 'Z', 'theta']:
def MCMC(self, kplanets, boundaries, inslims, acc_lims, sigmas_raw, pos0,
logl, logp):
ndim = 1 + 4 * kplanets + self.nins * 2 * (
self.MOAV + 1) + self.totcornum + self.PACC
ndat = len(self.time)
def starinfo():
colors = [
'red', 'green', 'blue', 'yellow', 'grey', 'magenta', 'cyan',
'white'
]
c = sp.random.randint(0, 7)
print(
colored(
'\n ###############################################',
colors[c]))
print(
colored(' # #',
colors[c]))
print(
colored(' # #',
colors[c]))
print(
colored(' # E M P E R 0 R #',
colors[c]))
print(
colored(' # #',
colors[c]))
print(
colored(' # #',
colors[c]))
print(
colored(' ###############################################',
colors[c]))
print(
colored(
'Exoplanet Mcmc Parallel tEmpering Radial vel0city fitteR',
colors[sp.random.randint(0, 7)]))
logdat = '\n\nStar Name : ' + self.starname
logdat += '\nTemperatures, Walkers, Steps : ' + str(
(self.ntemps, self.nwalkers, self.nsteps))
logdat += '\nN Instruments, K planets, N data : ' + str(
(self.nins, kplanets, self.ndat))
logdat += '\nN Number of Dimensions : ' + str(ndim)
logdat += '\nN Moving Average : ' + str(self.MOAV)
logdat += '\nBeta Detail : ' + str(
self.betas)
logdat += '\n-----------------------------------------------------'
print(logdat)
pass
starinfo()
#'''
sampler = PTSampler(self.ntemps,
self.nwalkers,
ndim,
logl,
logp,
loglargs=[
self.time, self.rv, self.err, self.ins,
self.staract, self.starflag, kplanets,
self.nins, self.MOAV, self.totcornum, self.PACC
],
logpargs=[
self.time, kplanets, self.nins, self.MOAV,
self.totcornum, boundaries, inslims, acc_lims,
sigmas_raw, self.eccprior, self.jittprior,
self.jittmean, self.STARMASS, self.HILL,
self.PACC, self.CHECK
],
threads=self.cores,
betas=self.betas)
print('\n --------------------- BURN IN --------------------- \n')
pbar = tqdm(total=self.burn_out)
for p, lnprob, lnlike in sampler.sample(pos0,
iterations=self.burn_out):
pbar.update(1)
pass
pbar.close()
p0, lnprob0, lnlike0 = p, lnprob, lnlike
print("\nMean acceptance fraction: {0:.3f}".format(
sp.mean(sampler.acceptance_fraction)))
assert sp.mean(sampler.acceptance_fraction
) != 0, 'Mean acceptance fraction = 0 ! ! !'
sampler.reset()
'''
sampler1 = PTSampler(self.ntemps, self.nwalkers, ndim, logl, logp, loglargs=[self.time, self.rv, self.err, self.ins, self.staract, self.starflag, kplanets, self.nins, self.MOAV, self.totcornum, self.PACC],
logpargs=[self.time, kplanets, self.nins, self.MOAV, self.totcornum, boundaries, inslims, acc_lims, sigmas_raw, self.eccprior, self.jittprior, self.jittmean, self.STARMASS, self.HILL, self.PACC, self.CHECK],
threads=self.cores, betas=self.betas*0.3)
print('\n --------------- EXPERIMENTAL BURN IN --------------- \n')
pbar = tqdm(total=self.burn_out)
for p, lnprob, lnlike in sampler1.sample(pos0, iterations=self.burn_out // 3):
pbar.update(1)
pass
pbar.close()
p0, lnprob0, lnlike0 = p, lnprob, lnlike
print("\nMean acceptance fraction: {0:.3f}".format(sp.mean(sampler1.acceptance_fraction)))
assert sp.mean(sampler1.acceptance_fraction) != 0, 'Mean acceptance fraction = 0 ! ! !'
sampler1.reset()
print('\n ---------------- EXPERIMENTAL CHAIN ---------------- \n')
sampler = PTSampler(self.ntemps, self.nwalkers, ndim, logl, logp, loglargs=[self.time, self.rv, self.err, self.ins, self.staract, self.starflag, kplanets, self.nins, self.MOAV, self.totcornum, self.PACC],
logpargs=[self.time, kplanets, self.nins, self.MOAV, self.totcornum, boundaries, inslims, acc_lims, sigmas_raw, self.eccprior, self.jittprior, self.jittmean, self.STARMASS, self.HILL, self.PACC, self.CHECK],
threads=self.cores, betas=self.betas)
pbar = tqdm(total=self.nsteps)
for p, lnprob, lnlike in sampler.sample(p0, lnprob0=lnprob0,
lnlike0=lnlike0,
iterations=self.nsteps,
thin=self.thin):
pbar.update(1)
pass
pbar.close()
'''
print('\n ---------------------- CHAIN ---------------------- \n')
pbar = tqdm(total=self.nsteps)
for p, lnprob, lnlike in sampler.sample(p0,
lnprob0=lnprob0,
lnlike0=lnlike0,
iterations=self.nsteps,
thin=self.thin):
pbar.update(1)
pass
pbar.close()
#'''
assert sampler.chain.shape == (self.ntemps, self.nwalkers,
self.nsteps / self.thin,
ndim), 'something really weird happened'
print("\nMean acceptance fraction: {0:.3f}".format(
sp.mean(sampler.acceptance_fraction)))
ln_post = sampler.lnprobability
posteriors = sp.array(
[ln_post[i].reshape(-1) for i in range(self.ntemps)])
chains = sampler.flatchain
best_post = posteriors[0] == np.max(posteriors[0])
thetas_raw = sp.array([chains[i] for i in range(self.ntemps)])
thetas_hen = sp.array([
empmir.henshin_PM(chains[i], kplanets) for i in range(self.ntemps)
])
ajuste_hen = thetas_hen[0][best_post][0]
ajuste_raw = thetas_raw[0][best_post][0]
interesting_loc = sp.array([
max(posteriors[temp]) - posteriors[temp] < self.bayes_factor
for temp in sp.arange(self.ntemps)
])
interesting_thetas = sp.array([
thetas_hen[temp][interesting_loc[temp]]
for temp in sp.arange(self.ntemps)
])
interesting_thetas_raw = sp.array([
thetas_raw[temp][interesting_loc[temp]]
for temp in sp.arange(self.ntemps)
])
interesting_posts = sp.array([
posteriors[temp][interesting_loc[temp]]
for temp in range(self.ntemps)
])
sigmas = sp.array(
[sp.std(interesting_thetas[0][:, i]) for i in range(ndim)])
sigmas_raw = sp.array(
[sp.std(interesting_thetas_raw[0][:, i]) for i in range(ndim)])
#print('sigmas', sigmas) # for testing
#print('sigmas_raw', sigmas_raw)
#print('mod_lims', boundaries)
print ajuste_raw, ajuste_hen
return thetas_raw, ajuste_raw, thetas_hen, ajuste_hen, p, lnprob, lnlike, posteriors, sampler.betas, interesting_thetas, interesting_posts, sigmas, sigmas_raw
def make_sampler(self):
return PTSampler(self.ntemps, self.nwalkers, self.ndim, self.lnlike, self.lnprior, threads=self.threads)
def mc_main(s_ident,
ntemps,
nwalkers,
niter,
nburn,
nthin,
nthreads,
mcdata,
mcmod,
partemp=True,
mc_a=2.,
init_samples_fn=None,
write_model=False,
plot=False,
save=False):
start = time.ctime()
time_start_secs = time.time()
print("\nSTART TIME: " + start)
# TEMP!!!
# For now, emcee v3.* offers only an EnsembleSampler, so force that mode
# if such a version is detected.
if (ntemps > 1) and (emcee_version_major >= 3):
ntemps = 1
partemp = False
emcee_v3_msg = "WARNING! FORCED to use EnsembleSampler because emcee v3+ detected. Setting ntemps=1 and partemp=False."
print("\n" + emcee_v3_msg)
print(
"To use a PTSampler, try the ptemcee package (NOT currently compatible with diskmc) or using emcee v2.2.1."
)
else:
emcee_v3_msg = None
data = mcdata.data
uncerts = mcdata.uncerts
data_types = np.array(mcdata.data_types) # need as nd.array for later
stars = mcdata.stars
model_path = os.path.join(
os.path.abspath(os.path.expanduser(mcmod.model_path)), '')
log_path = os.path.join(
os.path.abspath(os.path.expanduser(mcmod.log_path)), '')
lam = mcmod.lam # [microns]
unit_conv = mcmod.unit_conv
# Sort the parameter names.
# NOTE: this must be an array (can't be a list).
pkeys_all = np.array(sorted(mcmod.pkeys))
# Create log file.
mcmc_log_fn = os.path.join(log_path, '%s_mcmc_log.txt' % s_ident)
mcmc_log = open(mcmc_log_fn, 'w')
# FIX ME!!! Need to handle this specific case better.
# Make phi map specifically for conversion of Stokes to radial Stokes.
yy, xx = np.mgrid[:data[0].shape[0], :data[0].shape[1]]
phi_stokes = np.arctan2(yy - stars[0][0], xx - stars[0][1])
# Bin data by factors specified in mcdata.bin_facts list.
# Do nothing if mcdata.bin_facts is None or its elements are 1.
if mcdata.bin_facts is None:
mcdata.bin_facts = len(data) * [1]
data_orig = []
uncerts_orig = []
stars_orig = []
for ii, bin_fact in enumerate(mcdata.bin_facts):
if bin_fact not in [1, None]:
# Store the original data as backup.
data_orig.append(data[ii].copy())
uncerts_orig.append(uncerts[ii].copy())
# Bin data, uncertainties, and mask by interpolation.
datum_binned = zoom(np.nan_to_num(data_orig[ii].data),
1. / bin_fact) * bin_fact
uncert_binned = zoom(np.nan_to_num(uncerts_orig[ii]),
1. / bin_fact,
order=1) * bin_fact
# FIX ME!!! Interpolating the mask may not work perfectly. Linear interpolation (order=1)
# is best so far.
try:
mask_binned = zoom(np.nan_to_num(data_orig[ii].mask),
1. / bin_fact,
order=1)
except:
mask_binned = False
stars_orig.append(stars[ii].copy())
star_binned = stars_orig[ii] / int(bin_fact)
# radii_binned = make_radii(datum_binned, star_binned)
# mask_fit = np.ones(datum_binned.shape).astype(bool)
# mask_fit[star_binned[0]-int(hw_y/bin_fact):star_binned[0]+int(hw_y/bin_fact)+1, star_binned[1]-int(hw_x/bin_fact):star_binned[1]+int(hw_x/bin_fact)+1] = False
# FIX ME!!! Need to specify this inner region mask or happens automatically?
# mask_fit[radii_binned < r_fit/int(bin_fact)] = True
data[ii] = np.ma.masked_array(datum_binned, mask=mask_binned)
uncerts[ii] = uncert_binned
stars[ii] = star_binned
####################################
# ------ INITIALIZE WALKERS ------ #
# This will be done using a uniform distribution drawn from
# plims_lib (first option if not None) or a Gaussian distribution
# drawn from pmeans_lib and psigmas_lib (if plims_lib == None).
ndim = len(pkeys_all)
print(
"\nNtemps = %d, Ndim = %d, Nwalkers = %d, Nstep = %d, Nburn = %d, Nthreads = %d"
% (ntemps, ndim, nwalkers, niter, nburn, nthreads))
# Sort parameters for walker initialization.
if mcmod.plims_lib is not None:
plims_sorted = np.array(
[val for (key, val) in sorted(mcmod.plims_lib.items())])
init_type = 'uniform'
elif mcmod.pmeans_lib is not None:
pmeans_sorted = [
val for (key, val) in sorted(mcmod.pmeans_lib.items())
]
psigmas_sorted = [
val for (key, val) in sorted(mcmod.psigmas_lib.items())
]
init_type = 'gaussian'
# Make the array of initial walker positions p0.
if init_samples_fn is None:
if partemp:
if init_type == 'uniform':
# Uniform initialization between priors.
p0 = np.random.uniform(low=plims_sorted[:, 0],
high=plims_sorted[:, 1],
size=(ntemps, nwalkers, ndim))
elif init_type == 'gaussian':
# Gaussian ball initialization.
p0 = np.random.normal(loc=pmeans_sorted,
scale=psigmas_sorted,
size=(ntemps, nwalkers, ndim))
else:
if init_type == 'uniform':
# Uniform initialization between priors.
p0 = np.random.uniform(low=plims_sorted[:, 0],
high=plims_sorted[:, 1],
size=(nwalkers, ndim))
elif init_type == 'gaussian':
# Gaussian ball initialization.
p0 = np.random.normal(loc=pmeans_sorted,
scale=psigmas_sorted,
size=(nwalkers, ndim))
print("Walker initialization = " + init_type)
else:
init_type == 'sampled'
init_step_ind = -1
init_samples = hickle.load(os.path.join(log_path, init_samples_fn))
if partemp:
p0 = init_samples['_chain'][:, :, init_step_ind, :]
lnprob_init = init_samples['_lnprob'][:, :, init_step_ind]
lnlike_init = init_samples['_lnlikelihood'][:, :, init_step_ind]
else:
p0 = init_samples['_chain'][:, init_step_ind, :]
lnprob_init = init_samples['_lnprob'][:, init_step_ind]
lnlike_init = init_samples['_lnlikelihood'][:, init_step_ind]
print(
"\nLoaded init_samples from %s.\nWalkers will start from those final positions."
% init_samples_fn)
# Try to get the MCFOST version number being used.
try:
output = subprocess.check_output("mcfost -v", shell=True)
mcf_version = output.split('\n')[0].split(' ')[-1]
except:
mcf_version = 'unknown'
print("Running MCFOST version %s" % mcf_version)
print("Running diskmc version %s" % __version__)
log_preamble = [
'|---MCMC LOG---|\n\n',
'%s' % s_ident,
'\nLOG DATE: ' + dt.date.isoformat(dt.date.today()),
'\nJOB START: ' + start,
'\nNPROCESSORS: %d\n' % nthreads,
'\ndiskmc version: %s' % __version__,
'\nMCFOST version: %s' % mcf_version,
'\nMCFOST PARFILE: ',
mcmod.parfile,
'\n\nMCMC PARAMS: Ndim: %d, Nwalkers: %d, Nburn: %d, Niter: %d, Nthin: %d, Nthreads: %d'
% (ndim, nwalkers, nburn, niter, nthin, nthreads),
'\nPARALLEL-TEMPERED?: ' + str(partemp),
' , Ntemps: %d' % ntemps,
'\nINITIALIZATION: ',
init_type,
'\na = %.2f' % mc_a,
'\nWavelength = %s microns' % str(lam),
'\n',
]
mcmc_log.writelines(log_preamble)
if emcee_v3_msg is not None:
mcmc_log.writelines('\n{}\n'.format(emcee_v3_msg))
mcmc_log.close()
# Create emcee sampler instances: parallel-tempered or ensemble.
if partemp:
# Instantiate parallel-tempered sampler.
# Pass data_I, uncertainty_I, and parfile as arguments to emcee sampler.
sampler = PTSampler(
ntemps,
nwalkers,
ndim,
mc_lnlike,
mc_lnprior,
a=mc_a,
loglargs=[
pkeys_all, data, uncerts, data_types, mcmod.mod_bin_factor,
phi_stokes, mcmod.parfile, model_path, unit_conv, mcmod.priors,
mcmod.scatlight, mcmod.fullimg, mcmod.sed, mcmod.dustprops,
lam, partemp, ndim, write_model, s_ident, mcdata.algo_I,
mcmod.modfm
],
logpargs=[pkeys_all, mcmod.priors],
threads=nthreads)
#pool=pool)
else:
# Instantiate ensemble sampler.
if emcee_version_major >= 3:
# Put backend setup here for some future use.
# backend_log_name = os.path.join(log_path, '{}_mcmc_full_sampler.h5'.format(s_ident))
# backend = emcee.backends.HDFBackend(backend_log_name, name='mcmc')
# backend.reset(nwalkers, ndim)
#
# vars(backend)['pkeys_all'] = pkeys_all
backend = None
from multiprocessing import Pool
pool = Pool()
sampler = EnsembleSampler(
nwalkers,
ndim,
mc_lnlike,
a=mc_a,
args=[
pkeys_all, data, uncerts, data_types, mcmod.mod_bin_factor,
phi_stokes, mcmod.parfile, model_path, unit_conv,
mcmod.priors, mcmod.scatlight, mcmod.fullimg, mcmod.sed,
mcmod.dustprops, lam, partemp, ndim, write_model, s_ident,
mcdata.algo_I, mcmod.modfm
],
pool=pool,
backend=backend)
# Add some items to sampler that don't exist in emcee >2.2.1.
vars(sampler)['_chain'] = np.array([])
vars(sampler)['_lnprob'] = np.array([])
else:
sampler = EnsembleSampler(
nwalkers,
ndim,
mc_lnlike,
a=mc_a,
args=[
pkeys_all, data, uncerts, data_types, mcmod.mod_bin_factor,
phi_stokes, mcmod.parfile, model_path, unit_conv,
mcmod.priors, mcmod.scatlight, mcmod.fullimg, mcmod.sed,
mcmod.dustprops, lam, partemp, ndim, write_model, s_ident,
mcdata.algo_I, mcmod.modfm
],
threads=nthreads)
# Insert pkeys and priors into the sampler dict for later use.
# Force '|S' string dtype to avoid unicode (which doesn't hickle well).
vars(sampler)['pkeys_all'] = pkeys_all.astype('S')
vars(sampler)['priors'] = mcmod.priors.copy()
# List of items to delete from the sampler dict that may not hickle well during
# logging. Mix of emcee v2 and v3, Ensemble, and PTSampler items.
sampler_keys_trim = [
'pool', 'lnprobfn', 'log_prob_fn', 'runtime_sortingfn', 'logl', 'logp',
'logpkwargs', 'loglkwargs', 'args', 'kwargs', '_random', '_moves',
'_previous_state', 'backend'
]
###############################
# ------ BURN-IN PHASE ------ #
if nburn > 0:
print("\nBURN-IN START...\n")
for bb, (pburn, lnprob_burn,
lnlike_burn) in enumerate(sampler.sample(p0,
iterations=nburn)):
# Print progress every 25%.
if bb in [nburn // 4, nburn // 2, 3 * nburn // 4]:
print("PROCESSING ITERATION %d; BURN-IN %.1f%% COMPLETE..." %
(bb, 100 * float(bb) / nburn))
pass
# Print burn-in autocorrelation time and acceptance fraction.
try:
max_acl_burn = np.nanmax(
sampler.acor) # fails if too few iterations
except:
max_acl_burn = -1.
print("Largest Burn-in Autocorrelation Time = %.1f" % max_acl_burn)
if partemp:
print("Mean, Median Burn-in Acceptance Fractions: %.2f, %.2f" %
(np.mean(sampler.acceptance_fraction[0]),
np.median(sampler.acceptance_fraction[0])))
else:
print("Mean, Median Burn-in Acceptance Fractions: %.2f, %.2f" %
(np.mean(sampler.acceptance_fraction),
np.median(sampler.acceptance_fraction)))
# Pause interactively between burn-in and main phase.
# Comment this out if you don't want the script to pause here.
# pdb.set_trace()
# Reset the sampler chains and lnprobability after burn-in.
sampler.reset()
print("BURN-IN COMPLETE!")
elif (nburn == 0) & (init_samples_fn is not None):
print(
"\nWalkers initialized from file and no burn-in samples requested."
)
sampler.reset()
pburn = p0
lnprob_burn = None #lnprob_init
lnlike_burn = None #lnlike_init
else:
print("\nNo burn-in samples requested.")
pburn = p0
lnprob_burn = None
lnlike_burn = None
############################
# ------ MAIN PHASE ------ #
print("\nMAIN-PHASE MCMC START...\n")
if partemp:
for nn, (pp, lnprob, lnlike) in enumerate(
sampler.sample(pburn,
lnprob0=lnprob_burn,
lnlike0=lnlike_burn,
iterations=niter)):
# Print progress every 25%.
if nn in [niter // 4, niter // 2, 3 * niter // 4]:
print("Processing iteration %d; MCMC %.1f%% complete..." %
(nn, 100 * float(nn) / niter))
if emcee_version_major < 3:
# Log the full sampler or chain (all temperatures) every so often.
log_message = log_sampler(sampler, sampler_keys_trim,
log_path, s_ident, nn)
else:
for nn, (pp, lnprob, lnlike) in enumerate(
sampler.sample(pburn, lnprob_burn, iterations=niter)):
# Print progress every 25%.
if nn in [niter // 4, niter // 2, 3 * niter // 4]:
print("Processing iteration %d; MCMC %.1f%% complete..." %
(nn, 100 * float(nn) / niter))
# Log the full sampler or chain (all temperatures) every so often.
log_message = log_sampler(sampler, sampler_keys_trim, log_path,
s_ident, nn)
print('\nMCMC RUN COMPLETE!\n')
##################################
# ------ RESULTS HANDLING ------ #
# Log the sampler output and chains. Also get the max and median likelihood
# parameter values and save models for them.
# If possible, save the whole sampler to an HDF5 log file (could be large).
# If that fails because hickle won't handle something in the sampler,
# try to pickle it instead. If that still fails, just log the sampler chains.
# if emcee_version_major < 3:
log_message = log_sampler(sampler, sampler_keys_trim, log_path, s_ident,
'FINAL')
# Chain has shape (ntemps, nwalkers, nsteps/nthin, ndim).
if partemp:
assert sampler.chain.shape == (ntemps, nwalkers, niter / nthin, ndim)
else:
assert sampler.chain.shape == (nwalkers, niter / nthin, ndim)
# Re-open the text log for additional info.
mcmc_log = open(mcmc_log_fn, 'a')
mcmc_log.writelines([
'\n' + log_message,
'\n\n|--- RESULTS FOR ALL ITERATIONS (NO BURN-IN EXCLUDED) ---|', '\n'
])
# If multiple temperatures, take zero temperature walkers because only they
# sample the posterior distribution.
# ch has dimensions [nwalkers, nstep, ndim] and excludes burn-in samples.
if partemp:
ch = sampler.chain[0, :, :, :] # zeroth temperature chain only
samples = ch[:, :, :].reshape((-1, ndim))
lnprob_out = sampler.lnprobability[0] # zero-temp chi-squareds
# Median acceptance fraction of zero temp walkers.
print("\nMean, Median Acceptance Fractions (zeroth temp): %.2f, %.2f" %
(np.mean(sampler.acceptance_fraction[0]),
np.median(sampler.acceptance_fraction[0])))
mcmc_log.writelines(
'\nMean, Median Acceptance Fractions (zeroth temp): %.2f, %.2f' %
(np.mean(sampler.acceptance_fraction[0]),
np.median(sampler.acceptance_fraction[0])))
else:
ch = sampler.chain
samples = ch[:, :, :].reshape((-1, ndim))
lnprob_out = sampler.lnprobability # all chi-squareds
# Median acceptance fraction of all walkers.
print("\nMean, Median Acceptance Fractions: %.2f, %.2f" %
(np.mean(sampler.acceptance_fraction),
np.median(sampler.acceptance_fraction)))
mcmc_log.writelines('\nMean, Median Acceptance Fractions: %.2f, %.2f' %
(np.mean(sampler.acceptance_fraction),
np.median(sampler.acceptance_fraction)))
# Renormalize mass_fraction values so sum to 1.0 (more intuitive).
wh_pops = [ind for ind, key in enumerate(pkeys_all) if 'dust_pop' in key]
samples_orig = samples.copy()
samples[:, wh_pops] /= np.reshape(np.sum(samples_orig[:, wh_pops], axis=1),
(samples_orig.shape[0], 1))
# Haven't implemented blobs handling yet.
blobs = None
# Print zero temp main-phase autocorrelation time and acceptance fraction.
try:
max_acl = np.nanmax(sampler.acor)
if partemp:
acor_T0 = sampler.acor[0]
else:
acor_T0 = sampler.acor
except:
max_acl = -1.
acor_T0 = -1.
print("Largest Main Autocorrelation Time = %.1f" % max_acl)
# Max likelihood params values.
ind_lk_max = np.where(lnprob_out == lnprob_out.max())
lk_max = np.e**lnprob_out.max()
params_ml_mcmc = dict(zip(pkeys_all, ch[ind_lk_max][0]))
params_ml_mcmc_sorted = [
val for (key, val) in sorted(params_ml_mcmc.items())
]
# Get median values (50th percentile) and 1-sigma (68%) confidence intervals
# for each parameter (in order +, -).
params_med_mcmc = list(
map(lambda vv: (vv[1], vv[2] - vv[1], vv[1] - vv[0]),
zip(*np.percentile(samples, [16, 50, 84], axis=0))))
print("\nMax-Likelihood Param Values:")
mcmc_log.writelines('\n\nMAX-LIKELIHOOD PARAM VALUES:')
for kk, key in enumerate(pkeys_all):
print(key + ' = %.3e' % params_ml_mcmc[key])
mcmc_log.writelines('\n%s = %.3e' % (key, params_ml_mcmc[key]))
print(
"\n50%-Likelihood Param Values (50th percentile +/- 1 sigma (i.e., 34%):"
)
mcmc_log.writelines(
'\n\n50%-LIKELIHOOD PARAM VALUES (50th percentile +/- 1 sigma (i.e., 34%):'
)
for kk, key in enumerate(pkeys_all):
print(key + ' = %.3f +/- %.3f/%.3f' %
(params_med_mcmc[kk][0], params_med_mcmc[kk][1],
params_med_mcmc[kk][2]))
mcmc_log.writelines('\n%s = %.3f +/- %.3f/%.3f' %
(key, params_med_mcmc[kk][0],
params_med_mcmc[kk][1], params_med_mcmc[kk][2]))
# Construct max- and med-likelihood models.
print("\nConstructing 'best-fit' models...")
mod_idents = ['maxlk', 'medlk']
params_50th_mcmc = np.array(params_med_mcmc)[:, 0]
for mm, pl in enumerate([params_ml_mcmc_sorted, params_50th_mcmc]):
pl_dict = dict()
pl_dict.update(zip(pkeys_all, pl))
# Name for model and its directory.
try:
fnstring = "%s_mcmc_%s_%s%.3e_%s%.3e_%s%.3e" % \
(s_ident, mod_idents[mm], pkeys_all[0], pl_dict[pkeys_all[0]], pkeys_all[1], pl_dict[pkeys_all[1]],
pkeys_all[2], pl_dict[pkeys_all[2]])
except:
fnstring = "%s_mcmc_%s_%s%.5e" % \
(s_ident, mod_idents[mm], pkeys_all[0], pl_dict[pkeys_all[0]])
# Make the MCFOST model.
make_mcfmod(pkeys_all,
pl_dict,
mcmod.parfile,
model_path,
s_ident,
fnstring,
lam=lam,
scatlight=mcmod.scatlight,
fullimg=mcmod.fullimg)
# Calculate Chi2 for images.
chi2s = chi2_morph(
os.path.join(model_path, fnstring, 'data_%s' % str(lam)), data,
uncerts, data_types, mcmod.mod_bin_factor, phi_stokes, unit_conv)
# Calculate reduced Chi2 for images.
chi2_reds = np.array([
chi2s[ii] /
(np.where(np.isfinite(data[ii].filled(np.nan)))[0].size - ndim)
for ii in range(len(data))
])
chi2_red_total = np.sum(chi2_reds)
if mm == 0:
lk_type = 'Max-Likelihood'
# print('\nMax-Likelihood total chi2_red: %.3e | SED Cushing G: %.3e , I chi2_red: %.3f , Qr chi2_red: %.3f' % (chi2_red_total, G_mm, chi2_red_I, chi2_red_Qr))
# mcmc_log.writelines('\n\nMax-Likelihood total chi2_red: %.3e | SED Cushing G: %.3e , I chi2_red: %.3f , Qr chi2_red: %.3f' % (chi2_red_total, G_mm, chi2_red_I, chi2_red_Qr))
elif mm == 1:
lk_type = '50%-Likelihood'
# print('50%%-Likelihood total chi2_red: %.3e | SED Cushing G: %.3e , I chi2_red: %.3f , Qr chi2_red: %.3f' % (chi2_red_total, G_mm, chi2_red_I, chi2_red_Qr))
# mcmc_log.writelines('\n50%%-Likelihood total chi2_red: %.3e | SED Cushing G: %.3e , I chi2_red: %.3f , Qr chi2_red: %.3f' % (chi2_red_total, G_mm, chi2_red_I, chi2_red_Qr))
# print('%s total chi2_red: %.3e | SED Cushing G: %.3e , I chi2_red: %.3f , Qr chi2_red: %.3f' % (lk_type, chi2_red_total, G_mm, chi2_red_I, chi2_red_Qr))
# mcmc_log.writelines('\n%s total chi2_red: %.3e | SED Cushing G: %.3e , I chi2_red: %.3f , Qr chi2_red: %.3f' % (lk_type, chi2_red_total, G_mm, chi2_red_I, chi2_red_Qr))
print('%s total chi2_red: %.3e' % (lk_type, chi2_red_total))
print("individual chi2_red's: " +
len(chi2_reds) * "%.3e | " % tuple(chi2_reds))
mcmc_log.writelines('\n\n%s total chi2_red: %.3e' %
(lk_type, chi2_red_total))
mcmc_log.writelines("\nindividual chi2_red's: " +
len(chi2_reds) * "%.3e | " % tuple(chi2_reds))
# Make image, sed, and/or dust properties models for maxlk and medlk.
try:
os.chdir(os.path.join(model_path, fnstring))
# Make the dust properties at proper wavelength.
# NOTE: This must come before the image calculation at the same
# wavelength, otherwise MCFOST automatically deletes the image directory!
subprocess.call('rm -rf data_' + str(lam), shell=True)
subprocess.call(
'mcfost ' + fnstring +
'.para -dust_prop -op %s >> dustmcfostout.txt' % lam,
shell=True)
time.sleep(1)
# # Delete the (mostly) empty data_[lam] directory after dust_prop step.
# subprocess.call('rm -rf data_'+str(lam), shell=True)
# Make the SED model.
subprocess.call('mcfost ' + fnstring +
'.para -rt >> sedmcfostout.txt',
shell=True)
# Make the image models (thermal + scattered-light) at demanded wavelength.
subprocess.call('mcfost ' + fnstring + '.para -img ' + str(lam) +
' -rt2 >> imagemcfostout.txt',
shell=True)
time.sleep(1)
print("Saved image and dust properties models.")
except:
print("Failed to save image and dust properties models.")
# Plot and save maxlk and medlk models.
try:
# Placeholder for plotting functions.
pass
print("Max and Median Likelihood models made, plotted, and saved.\n")
except:
print(
"Max and Median Likelihood models made and saved but plotting failed.\n"
)
time_end_secs = time.time()
time_elapsed_secs = time_end_secs - time_start_secs # [seconds]
print("END TIME: " + time.ctime())
print("ELAPSED TIME: %.2f minutes = %.2f hours" %
(time_elapsed_secs / 60., time_elapsed_secs / 3600.))
mcmc_log.writelines(
['\n\nSTART TIME - END TIME: ', start, ' - ',
time.ctime()])
mcmc_log.writelines([
'\nELAPSED TIME: %.2f minutes = %.2f hours' %
(time_elapsed_secs / 60., time_elapsed_secs / 3600.)
])
mcmc_log.writelines('\n\nSUCCESSFUL EXIT')
mcmc_log.close()
# Close MPI pool.
try:
pool.close()
print("\nPool closed")
except:
print("\nNo Pool to close.")
print("\nmc_main function finished\n")
# # Pause interactively before finishing script.
# pdb.set_trace()
return
def start_sampler(data, recovery_prob, burnin, niter, verbose, contact_daylist, max_recovery_time, nsick_param, diagnosis_lag=False, null_comparison=False, **kwargs3):
r"""Sampling performed using emcee """
parameter_estimate=None
##############################################################################
G_raw, health_data, node_health, nodelist, true_value, time_min, time_max, seed_date =data
######################################
### Set number of parameters to estimate
######################################
ndim_base = 2
if recovery_prob: ndim_base += nsick_param
ndim = ndim_base+nsick_param
#######################
##Adjust temperature ladder
#######################
betas = np.linspace(0, -2, 15)
betas = 10**(np.sort(betas)[::-1])
ntemps = 15
###########################################
###set starting positions for the walker
#############################################
nwalkers = max(50, 2*ndim) # number of MCMC walkers
starting_guess = np.zeros((ntemps, nwalkers, ndim))
##starting guess for beta
starting_guess[:, :, 0] = np.random.uniform(low = 0, high = 10, size=(ntemps, nwalkers))
##start epsilon close to zero
epsilons = np.random.power(4, size = (ntemps, nwalkers))
starting_guess[:, :, 1] = 1-epsilons
if diagnosis_lag:
starting_guess[:, :, 2: ] = np.random.uniform(low = 0.001, high = 1,size=(ntemps, nwalkers, ndim-2))
################################################################################
##calculating infection date and infection strength outside loglik to speed up #
##computations
################################################################################
if not diagnosis_lag:
infection_date = [(node, time1) for node in node_health if node_health[node].has_key(1) for (time1,time2) in node_health[node][1]]
infection_date = sorted(infection_date)
infected_strength = {network:{node:{time: calculate_infected_strength(node, time, health_data, G_raw[network]) for time in range(time_min, time_max+1)} for node in nodelist} for network in G_raw}
pool = None
threads = 1
else:
infection_date = None
infected_strength=None
threads = 8
healthy_nodelist = return_healthy_nodelist(node_health)
################################################################################
if threads>1:
sampler = PTSampler(ntemps=ntemps, nwalkers=nwalkers, dim=ndim, betas=betas, logl=log_likelihood, logp=log_prior, a = 1.5, loglargs=(data, infection_date, infected_strength, healthy_nodelist, null_comparison, diagnosis_lag, recovery_prob, nsick_param, contact_daylist, max_recovery_time, parameter_estimate), logpargs=(null_comparison, diagnosis_lag, nsick_param, recovery_prob, parameter_estimate), threads=threads)
if threads<=1:
sampler = PTSampler(ntemps=ntemps, nwalkers=nwalkers, dim=ndim, betas=betas, logl=log_likelihood, logp=log_prior, a = 1.5, loglargs=(data, infection_date, infected_strength, healthy_nodelist, null_comparison, diagnosis_lag, recovery_prob, nsick_param, contact_daylist, max_recovery_time, parameter_estimate), logpargs=(null_comparison, diagnosis_lag, nsick_param, recovery_prob, parameter_estimate))
#Run user-specified burnin
print ("burn in......")
for i, (p, lnprob, lnlike) in enumerate(sampler.sample(starting_guess, iterations = burnin)):
if verbose:print("burnin progress..."), (100 * float(i) / burnin)
else: pass
sampler.reset()
#################################
print ("sampling........")
nthin = 5
for i, (p, lnprob, lnlike) in enumerate(sampler.sample(p, lnprob0 = lnprob, lnlike0= lnlike, iterations= niter, thin= nthin)):
if verbose:print("sampling progress"), (100 * float(i) / niter)
else: pass
##############################
#The resulting samples are stored as the sampler.chain property:
assert sampler.chain.shape == (ntemps, nwalkers, niter/nthin, ndim)
return sampler
def run_emcee(logl,
logp,
p0func,
ntemps=0,
nwalkers=50,
nsamples=2500,
thin=1,
minlogbeta=None,
nupdates=10,
threads=1,
outfilename=None,
saveall=True,
**kwargs):
if minlogbeta is None:
if ntemps == 0:
# use cunning ladder
betas = np.concatenate(
(np.linspace(0, -0.9375, 16), np.linspace(-1, -1.875, 8),
np.linspace(-2, -3.75, 8), np.linspace(-4, -7.5, 8),
np.linspace(-8, -15,
8), np.linspace(-16, -30,
8), np.linspace(-32, -56, 4)))
betas = 10**(np.sort(betas)[::-1])
else:
betas = None # use emcee default
else:
betas = np.logspace(0, minlogbeta, ntemps)
if betas is not None:
ntemps = len(betas)
pos = p0func((ntemps, nwalkers))
if not check_init_pars(logl, logp, pos):
return None
ndim = pos.shape[-1]
if threads > 1:
pool = Pool(threads)
else:
pool = None
sampler = PTSampler(ntemps,
nwalkers,
ndim,
logl,
logp,
betas=betas,
pool=pool)
if nupdates > 0:
start = time.clock()
print('Steps:', end='')
sys.stdout.flush()
nsteps = nsamples // nupdates
for i in range(nupdates):
pos, lnprob, rstate = sampler.run_mcmc(pos, nsteps, thin=thin)
print(' {}'.format((i + 1) * nsteps), end='')
sys.stdout.flush()
nsteps = nsamples - sampler.chain.shape[-2]
if nsteps > 0:
pos, lnprob, rstate = sampler.run_mcmc(pos, nsteps, thin=thin)
if nupdates > 0:
print('\nTime taken = {:.2f} secs'.format(time.clock() - start))
if outfilename is None:
outfilename = 'emcee_sampler.npz'
try:
acor = sampler.acor
except autocorr.AutocorrError:
acor = None
if saveall:
np.savez_compressed(
outfilename,
acceptance_fraction=sampler.acceptance_fraction,
acor=acor,
beta=sampler.betas,
chain=sampler.chain,
lnlikelihood=sampler.lnlikelihood,
lnprobability=sampler.lnprobability,
tswap_acceptance_fraction=sampler.tswap_acceptance_fraction)
else:
# only save lowest temperature and thin samples by factor of ten
np.savez_compressed(
outfilename,
acceptance_fraction=sampler.acceptance_fraction,
acor=acor,
beta=sampler.betas,
chain=sampler.chain[0, :, ::10],
lnlikelihood=sampler.lnlikelihood[0, :, ::10],
lnprobability=sampler.lnprobability[0, :, ::10],
tswap_acceptance_fraction=sampler.tswap_acceptance_fraction)
return sampler
