def main(settings_in,advanced_settings_in, priors_in):
"""
A simple example of how to instantiate the objects used by ExoSOFTmodel
with their required input parameters. This example uses the 5% Jupiter
analogue discussed in the release paper and calculates the chi squared
of the expected parameters for the simulated data.
"""
## load up default settings dictionary
sd = tools.startup(settings_in,advanced_settings_in,priors_in,rePlot=True)
log = KMlogger.getLogger('main',lvl=sd['logLevel'],addFH=False)
## Instantiate main objects/classes:
# ExoSOFTpriors, ExoSOFTdata and ExoSOFTparams.
# These are all instantiated as member variables of ExoSOFTmodel class.
Model = tools.ExoSOFTmodel(sd)
## define a set of starting parameters
# The user can use any reasonable guess here. For the 5% Jupiter analogue
# used in this example, we will use expected values for simplicity.
m2 = const.M_jup.value/const.M_sun.value
sqrte_sinomega = np.sqrt(0.048)*np.sin((np.pi/180.0)*14.8)
sqrte_cosomega = np.sqrt(0.048)*np.cos((np.pi/180.0)*14.8)
# pars: [m1,m2,parallax,long_an, e/sqrte_sinomega,to/tc,p,inc,arg_peri/sqrte_cosomega,v1,v2...]
start_params = [1.0,m2,50,100.6,sqrte_sinomega,2450639.0,11.9,45.0,sqrte_cosomega,0.01]
## NOTE: this array must be a numpy array with dtype=np.dtype('d').
start_params = np.array(start_params,dtype=np.dtype('d'))
## calculate the log posterior for the provided data and input parameters
ln_post = tools.ln_posterior(start_params, Model)
## print a few basic results of the fit
log.importantinfo('\nchi_squared_3d '+str(Model.chi_squared_3d))
log.importantinfo('reduced chi_squared_3d '+str(Model.chi_squared_3d/35.0))
log.importantinfo('prior '+str(Model.prior))
log.importantinfo('ln_post '+str(ln_post)+'\n')
def main(settings_in, advanced_settings_in, priors_in):
"""
An example of using ExoSOFTmodel with emcee (https://github.com/dfm/emcee).
This will show:
1. How to intantiate the necessary objects with their required input
parameters
2. How prepare the inputs to emcee
3. Then run emcee and finally make some simple plots of the results.
"""
## Simple boolean flags to control the basic steps of this script.
# Just do a quick test run?
# Else, run a longer one to give smooth posteriors
quick = True
# What plots should be made?
show_burnin = False
show_chains = True
show_posteriors = True
show_corner = False
sd = tools.startup(settings_in,
advanced_settings_in,
priors_in,
rePlot=True)
log = KMlogger.getLogger('main', lvl=sd['logLevel'], addFH=False)
## Instantiate main objects/classes:
# ExoSOFTpriors, ExoSOFTdata and ExoSOFTparams.
# These are all instantiated as member variables of ExoSOFTmodel class.
Model = tools.ExoSOFTmodel(sd)
## define a set of starting parameters
# The user can use any reasonable guess here. For the 5% Jupiter analogue
# used in this example, we will use expected values for simplicity.
m2 = const.M_jup.value / const.M_sun.value
sqrte_sinomega = np.sqrt(0.048) * np.sin((np.pi / 180.0) * 14.8)
sqrte_cosomega = np.sqrt(0.048) * np.cos((np.pi / 180.0) * 14.8)
# pars: [m1,m2,parallax,long_an, e/sqrte_sinomega,to/tc,p,inc,arg_peri/sqrte_cosomega,v1,v2...]
start_params = [
1.0, m2, 50, 100.6, sqrte_sinomega, 2450639.0, 11.9, 45.0,
sqrte_cosomega, 0.1
]
ncpu = multiprocessing.cpu_count()
ndim = len(sd['range_maxs']) # number of parameters in the model
if quick:
nwalkers = 50 # number of MCMC walkers
nburn = 100 # "burn-in" to stabilize chains
nsteps = 1000 # number of MCMC steps to take
else:
nwalkers = 500
nburn = 500
nsteps = 50000
## NOTE: starting_guesses array must be a numpy array with dtype=np.dtype('d').
starting_guesses = tools.make_starting_params(start_params,
nwalkers,
scale=0.01)
log.importantinfo("\nObjects and inputs prepared, now calling emcee.")
## Call emcee to explore the parameter space
sampler = emcee.EnsembleSampler(nwalkers,
ndim,
tools.ln_posterior,
args=[Model],
threads=ncpu)
sampler.run_mcmc(starting_guesses, nsteps)
# chain is of shape (nwalkers, nsteps, ndim)
# discard burn-in points and reshape
trace = sampler.chain[:, nburn:, :]
trace = trace.reshape(-1, ndim)
labels = ['m2', 'period', 'inclination']
## Show walkers during burn-in
if show_burnin:
fig = plt.figure(figsize=(10, 5))
j = 0
for i, chain in enumerate(sampler.chain[:, :nburn, :].T):
if i in [1, 6, 7]:
plt.subplot(3, 1, j + 1)
if i == 1:
# convert m2 units to Mjup instead of Msun
chain *= const.M_sun.value / const.M_jup.value
plt.plot(chain, drawstyle='steps', color='k', alpha=0.2)
plt.ylabel(labels[j])
j += 1
plt.show()
## Show walkers after burn-in
if show_chains:
fig = plt.figure(figsize=(10, 5))
j = 0
for i, chain in enumerate(sampler.chain[:, nburn:, :].T):
if i in [1, 6, 7]:
plt.subplot(3, 1, j + 1)
if i == 1:
# convert m2 units to Mjup instead of Msun
chain *= const.M_sun.value / const.M_jup.value
plt.plot(chain, drawstyle='steps', color='k', alpha=0.2)
log.info('\nparameter # ' + str(i))
log.info('mean = ' + str(np.mean(chain)))
log.info('median = ' + str(np.median(chain)))
log.info('variance = ' + str(np.var(chain)))
#print('chain ',repr(chain))
plt.ylabel(labels[j])
j += 1
plt.show()
## inspect posteriors
if show_posteriors:
fig = plt.figure(figsize=(12, 3))
j = 0
for i in range(ndim):
if i in [1, 6, 7]:
plt.subplot(1, 3, j + 1)
true_val = start_params[i]
trace_use = trace[:, i]
if i == 1:
# convert m2 units to Mjup instead of Msun
trace_use *= const.M_sun.value / const.M_jup.value
true_val *= const.M_sun.value / const.M_jup.value
plt.hist(trace_use, 100, color="k", histtype="step")
yl = plt.ylim()
plt.vlines(true_val,
yl[0],
yl[1],
color='blue',
lw=3,
alpha=0.25,
label='true')
plt.title("{}".format(labels[j]))
plt.legend()
j += 1
plt.show()
## make a corner plot
if show_corner:
fig = corner.corner(trace,
labels=labels,
quantiles=[0.16, 0.5, 0.84],
truths=start_params)
plt.show()
log.critical("\nFinished emcee example :D \n")
def custom_post(settings_in, advanced_settings_in, priors_in):
## load up settings that were passed in
tic = timeit.default_timer()
settings = tools.startup(settings_in,
advanced_settings_in,
priors_in,
rePlot=True)
log = KMlogger.getLogger('main',
dr=settings['finalFolder'],
lvl=settings['logLevel'])
skipBurnInStrip = True
# try to bring back pickeled fleshed out settings file and final results strings
# if they exists.
[
allFname, burnInStr, clStr, grStr, effPtsStr, durationStrings,
postTime, allTime
] = ['', '', '', '', '', '', '', '']
[MCMCmpoRO, SAmpoRO, STmpoRO, MCmpoRO, emcee_mpoRO,
FINALmpoRO] = [None, None, None, None, None, None]
outFiles = []
if False:
try:
settings = tools.unPklIt(settings, 'settings')
finalSummaryStrs = tools.unPklIt(settings, 'finalSummaryStrs')
[
allFname, outFiles, stageList, clStr, burnInStr, bestFit,
grStr, effPtsStr, allTime, postTime, durationStrings
] = finalSummaryStrs
FINALmpoRO = tools.unPklIt(settings, 'FINALmpoRO')
# mpoRO contains members: bestRedChiSqrs,avgAcceptRates,acceptStrs,stage,retStr,latestRetStr
[MCmpoRO, SAmpoRO, STmpoRO,
MCMCmpoRO] = tools.reloadMpoROs(settings)
'For MCMC stage: ' + str(np.mean(MCMCmpoRO.avgAcceptRates))
except:
print('AN ERROR OCCURRED WHILE TRYING TO LOAD PKLS BACK IN')
## make list of output files
outFiles = []
if settings['stageList'][-1] == 'MCMC':
outFiles = np.sort(glob.glob(os.path.join(settings['finalFolder'],\
"outputDataMCMC*_BIstripped.fits")))
if len(outFiles) == 0:
outFiles = np.sort(glob.glob(os.path.join(settings['finalFolder'],\
"outputDataMCMC*.fits")))
else:
outFiles = np.sort(glob.glob(os.path.join(settings['finalFolder'],\
"outputData"+\
settings['stageList'][-1]+"*.fits")))
if len(outFiles) == 0:
log.error("No appropriate list of individual chain outputs to "+\
"perform burn-in stripping on, or combine.")
## get name for combined output data file
allFname = ''
if os.path.exists(os.path.join(settings['finalFolder'],\
"combined-BIstripped-MCMCdata.fits")):
allFname = os.path.join(settings['finalFolder'],\
"combined-BIstripped-MCMCdata.fits")
skipBurnInStrip = True
elif os.path.exists(os.path.join(settings['finalFolder'],\
"combined"+settings['stageList'][-1]+"data.fits")):
allFname = os.path.join(settings['finalFolder'],\
"combined"+settings['stageList'][-1]+"data.fits")
skipBurnInStrip = False
else:
skipBurnInStrip = True
log.critical("No combined MCMC file available. Please check and/or "+\
"modify the customPost.py script as needed.")
## combine the individual chain data files needed
if (len(outFiles) > 0) and (allFname == ''):
log.importantinfo('about to combine data files together')
if 'BIstripped' in outFiles[0]:
allFname = os.path.join(os.path.dirname(outFiles[0]),\
"combined-BIstripped-MCMCdata.fits")
else:
allFname = os.path.join(os.path.dirname(outFiles[0]),\
"combined"+settings['stageList'][-1]+"data.fits")
tools.combineFits(outFiles, allFname)
## calc and strip burn-in?
burnInStr = ''
if True:
if skipBurnInStrip == False:
log.importantinfo('about to strip burn-in')
if (len(outFiles)>1)and(settings['CalcBurn'] and\
(settings['stageList'][-1]=='MCMC')):
(burnInStr,
burnInLengths) = tools.burnInCalc(outFiles, allFname)
if settings['rmBurn'][0]:
strippedFnames = tools.burnInStripper(
outFiles, burnInLengths)
outFiles = strippedFnames
## combine stripped files to make final file?
if len(strippedFnames) > 0:
strippedAllFname = os.path.join(os.path.dirname(strippedFnames[0]),\
"combined-BIstripped-MCMCdata.fits")
tools.combineFits(strippedFnames, strippedAllFname)
## replace final combined filename with new stripped version
allFname = strippedAllFname
## find best fit
if True:
log.importantinfo('about to find best orbit')
bestFit = tools.findBestOrbit(allFname,
bestToFile=False,
findAgain=False)
else:
# for manual input of best values
bestFit = np.array([
0.9892581, 0.0009696, 49.4824152, 101.82462, 0.05706358,
2450656.61, 2450656.61, 12.012219, 44.5819933, 0.1945267, 5.227630,
46.908327, 8.916892, -0.04136662
])
## If using the 'expected' values for synthetic data, they are in 'stored_pars' format.
# stored_pars: [m1,m2,parallax,long_an,e,to,tc,p,inc,arg_peri,a_tot_au,chi_sqr,K,v1,v2...]
expectedValues = np.array([
1.0, 0.0009543, 50.0, 101.6, 0.048, 2450639.0, 2450639.0, 11.9, 45.0,
14.8, 5.21, 1.0, 8.9, 0.0
])
if False:
log.importantinfo('about to make orbit plots')
plotFnameBase = os.path.join(settings['finalFolder'], 'orbPlot-Manual')
##for reference: DIlims=[[[xMin,xMax],[yMin,yMax]],[[xCropMin,xCropMax],[yCropMin,yCropMax]]] [[[,],[,]],[[,],[]]]
## RVlims=[[yMin,yMax],[yResidMin,yResidMax],[xMin,xMax]]
tools.orbitPlotter(bestFit,
settings,
plotFnameBase,
fl_format='eps',
DIlims=[],
RVlims=[])
## Simulated jupiter
#tools.orbitPlotter(bestFit,settings,plotFnameBase,fl_format='eps',DIlims=[],RVlims=[[-9.95,9.9],[-1.5,1.5],[-0.55,0.55]],diErrMult=1,diLnThk=2.5)
clStr = ''
if True:
log.importantinfo('about to plot posteriors')
tic2 = timeit.default_timer()
plotFilename = os.path.join(settings['finalFolder'], 'posteriors')
clStr = tools.summaryPlotter(allFname,
plotFilename,
paramsToPlot=[],
xLims=[],
bestVals=bestFit,
stage=settings['stageList'][-1],
shadeConfLevels=True,
forceRecalc=False,
plotALLpars=True)
#clStr = tools.summaryPlotter(allFname, plotFilename,paramsToPlot=[0,1,8,10],xLims=[[0.5,1.6],[0.5,1.5],[35,52],[4.0,6.3]],bestVals=expectedValues,stage=settings['stageList'][-1], shadeConfLevels=True,forceRecalc=False)
log.importantinfo("It took a total of " +
tools.timeStrMaker(timeit.default_timer() - tic2))
## make corner plot?
if False:
log.importantinfo('about to make corner plot')
tic2 = timeit.default_timer()
plotFilename = os.path.join(settings['finalFolder'],
'cornerPlot_expectedVals')
label_kwargs = {'fontsize': 28, 'weight': 'heavy', 'labelpad': 1}
other_kwargs = {'labelsize': 18, 'xy_factor': 1.5}
tools.cornerPlotter(allFname,
plotFilename,
paramsToPlot=[0, 1, 8, 10],
bestVals=expectedValues,
smooth=True,
label_kwargs=label_kwargs,
latex=True,
fl_format='png',
other_kwargs=other_kwargs)
log.importantinfo("It took a total of " +
tools.timeStrMaker(timeit.default_timer() - tic2))
if False:
log.importantinfo('about to make 2D density plot')
plotFilename = os.path.join(settings['finalFolder'], 'm1m2-densPlot')
#ranges=[[xMin,xMax],[yMin,yMax]]
tools.densityPlotter2D(allFname,
plotFilename,
paramsToPlot=[0, 1],
bestVals=[bestFit[0], bestFit[1]])
if False:
log.importantinfo('about to make progress plots')
##make progress plot
for stgStr in ['SA', 'ST', 'MCMC']:
emcee_stage = False
if stgStr == 'emcee':
emcee_stage = True
for procNum in range(0, 2):
dataFname = os.path.join(
settings['finalFolder'],
'outputData' + stgStr + str(procNum) + '.fits')
if os.path.exists(dataFname):
log.info("About to make progress plots for file: " +
dataFname)
for parNum in [1, 4, 8]:
plotFilename = os.path.join(
settings['finalFolder'], 'progressPlot-' + stgStr +
str(procNum) + '-' + str(parNum))
try:
#outputDataFilename,plotFilename,paramToPlot,yLims=[],xLims = [],expectedVal=None,emcee_stage=False,downSample=True
tools.progressPlotter(dataFname,
plotFilename,
parNum,
yLims=[],
xLims=[],
expectedVal=None,
emcee_stage=emcee_stage,
downSample=False)
except:
log.importantinfo(
'could not make plot for proc# ' +
str(procNum) + ', and par# ' + str(parNum))
##calc R?
grStr = ''
if False:
log.importantinfo('about to calc GR')
if (len(outFiles) > 1) and (settings['CalcGR'] and
(settings['stageList'][-1] == 'MCMC')):
tic2 = timeit.default_timer()
grStr = tools.gelmanRubinCalc(outFiles, settings['nSamples'])
log.info(grStr)
log.importantinfo("It took a total of " +
tools.timeStrMaker(timeit.default_timer() -
tic2))
## custom re check of the orbit fit
if False:
log.importantinfo(
'about to calculate predicted location for custom date')
orbParams = bestFit
epochs = [2457327.500]
tools.predictLocation(orbParams, settings, epochs)
## calc correlation length & number effective points?
effPtsStr = ''
if False:
log.importantinfo('about to calc # effective points')
if ((len(outFiles) > 1) and
(settings['stageList'][-1]
== 'MCMC')) and (settings['calcCL'] and os.path.exists(allFname)):
tic2 = timeit.default_timer()
effPtsStr = tools.mcmcEffPtsCalc(allFname)
log.info(effPtsStr)
log.importantinfo("It took a total of " +
tools.timeStrMaker(timeit.default_timer() -
tic2))
## calc auto-correlation time using tools from the emcee package
if False:
log.importantinfo('about to calc integrated autocorr time')
if ((settings['stageList'][-1] == 'emcee') or
(settings['stageList'][-1] == 'MCMC')) and (
settings['calcIAC'] and os.path.exists(allFname)):
tic2 = timeit.default_timer()
iacStr = tools.autocorr(allFname)
log.info(iacStr)
log.importantinfo("It took a total of " +
tools.timeStrMaker(timeit.default_timer() -
tic2))
postTime = timeit.default_timer() - tic
## following post-processing stages can take a long time, so write the current
## summary information to the summary file and add the rest later
if False:
if os.path.exists(allFname):
log.importantinfo('about to make summary file')
tools.summaryFile(settings, settings['stageList'], allFname, clStr,
burnInStr, bestFit, grStr, effPtsStr, iacStr,
allTime, postTime, durationStrings, MCmpoRO,
SAmpoRO, STmpoRO, MCMCmpoRO, emcee_mpoRO)
## CAUTION!! This way of cleaning is permenant!!
## Double check your settings and that this is what you want first!!
## clean up files (move to folders or delete them)
if False:
log.importantinfo('about to clean up output directory')
tools.cleanUp(settings, settings['stageList'], allFname)
## copy to dropbox is a legacy function that will be repurposed or killed later
if False and ('copyToDB' in settings):
if settings['copyToDB']:
log.importantinfo('about to copy files to dropbox')
tools.copyToDB(settings)
log.importantinfo("It took a total of " + tools.timeStrMaker(postTime))
### DONE customPost
print("\ncustomPost script complete :-D\n")