import numpy as np
# A dictionary of parameter lists for conversion.
registered_params = {"SB1": ["K", "e", "omega", "P", "T0", "gamma", "amp_f", "l_f"],
"SB2": ["q", "K", "e", "omega", "P", "T0", "gamma", "amp_f", "l_f", "amp_g", "l_g"],
"ST1": ["K_in", "e_in", "omega_in", "P_in", "T0_in", "K_out", "e_out", "omega_out", "P_out", "T0_out", "gamma", "amp_f", "l_f"],
"ST2": ["q_in", "K_in", "e_in", "omega_in", "P_in", "T0_in", "K_out", "e_out", "omega_out", "P_out", "T0_out", "gamma", "amp_f", "l_f"],
"ST3": ["q_in", "K_in", "e_in", "omega_in", "P_in", "T0_in", "q_out", "K_out", "e_out", "omega_out", "P_out", "T0_out", "gamma", "amp_f", "l_f", "amp_g", "l_g", "amp_h", "l_h"]}
registered_models = registered_params.keys()
# Calculate how many orbital parameters there are to each model.
# For now, we'll just calculate this based upon the position of the "gamma" parameter
n_params_orb = {model: (registered_params[model].index("gamma") + 1) for model in registered_params.keys()}
# labels used for plotting
registered_labels = {"SB1": [r"$K$", r"$e$", r"$\omega$", r"$P$", r"$T_0$", r"$\gamma$", r"$a_f$", r"$l_f$"],
"SB2": [r"$q$", r"$K$", r"$e$", r"$\omega$", r"$P$", r"$T_0$", r"$\gamma$", r"$a_f$", r"$l_f$", r"$a_g$", r"$l_g$"],
"ST3": [r"$q_\mathrm{in}$", r"$K_\mathrm{in}$", r"$e_\mathrm{in}$", r"$\omega_\mathrm{in}$", r"$P_\mathrm{in}$", r"$T_{0,\mathrm{in}}$", r"$q_\mathrm{out}$", r"$K_\mathrm{out}$", r"$e_\mathrm{out}$", r"$\omega_\mathrm{out}$", r"$P_\mathrm{out}$", r"$T_{0,\mathrm{out}}$", r"$\gamma$", r"$a_f$", r"$l_f$", r"$a_g$", r"$l_g$", r"$a_h$", r"$l_h$"]}
# For example, if the config.yaml file specifies model as "SB1", then it must list all of these parameters under "SB1".
# It must also specify fix_params, which will always include gamma.
# Then, within the sampling routines, we will create a partial function that takes in model type, list of fixed parameters, and dictionary of default parameter values and upon invocation will convert a vector of only a subset of values into a full vector.
[docs]def convert_vector(p, model, fix_params, **kwargs):
'''Unroll a vector of parameter values into a parameter type, using knowledge about which model we are fitting, the parameters we are fixing, and the default values of those parameters.
Args:
p (np.float) : 1D input array of only a subset of parameter values.
model (str): "SB1", "SB2", etc..
fix_params (list of str): names of parameters that will be fixed
**kwargs: input for ``{param_name: default_value}`` pairs
Returns:
(np.float, np.float) : a 2-tuple of the full vectors for the orbital parameters, and the GP parameters, augmented with the previously missing values.
'''
# Select the registered parameters corresponding to this model
reg_params = registered_params[model]
nparams = len(reg_params)
# fit parameters are the ones in reg_params that are *not* in fix_params
fit_params = [param for param in reg_params if param not in fix_params]
fit_ind = [i for (i,param) in enumerate(reg_params) if param not in fix_params]
# go through each element in fix_params, and find out where it would fit in reg_params
fix_ind = [reg_params.index(param) for param in fix_params]
# make an empty parameter vector of total parameter length
par_vec = np.empty(nparams, dtype=np.float64)
# Stuff it with all the values that we are fitting
par_vec[fit_ind] = p
# Fill in the holes with parameters that we are fixing
# First, convert all fixed parameters to a similar type vector
p_fixed = np.array([kwargs[par_name] for par_name in fix_params])
par_vec[fix_ind] = p_fixed
# Now split it to the orbital parameters and the GP parameters
ind_split = n_params_orb[model]
par_orb = par_vec[:ind_split]
par_GP = par_vec[ind_split:]
return (par_orb, par_GP)
# function convert_dict(p::Dict, model::AbstractString)
[docs]def convert_dict(model, fix_params, **kwargs):
'''Used to turn a dictionary of parameter values (from config.yaml) directly into a parameter type. Generally used for synthesis and plotting command line scripts.'''
# Select the registered parameters corresponding to this model
reg_params = registered_params[model]
nparams = len(reg_params) - len(fix_params)
par_vec = np.empty(nparams, dtype=np.float64)
fit_params = [param for param in reg_params if param not in fix_params]
p_fit = np.array([kwargs[par_name] for par_name in fit_params])
return p_fit
[docs]def get_labels(model, fix_params):
'''
Collect the labels for each model, so that we can plot.
'''
reg_params = registered_params[model]
reg_labels = registered_labels[model]
fit_index = [i for (i,param) in enumerate(reg_params) if param not in fix_params]
labels = [reg_labels[i] for i in fit_index]
return labels
[docs]def gelman_rubin(samplelist):
'''
Given a list of flatchains from separate runs (that already have burn in cut
and have been trimmed, if desired), compute the Gelman-Rubin statistics in
Bayesian Data Analysis 3, pg 284. If you want to compute this for fewer
parameters, then slice the list before feeding it in.
'''
full_iterations = len(samplelist[0])
assert full_iterations % 2 == 0, "Number of iterations must be even. Try cutting off a different number of burn in samples."
shape = samplelist[0].shape
#make sure all the chains have the same number of iterations
for flatchain in samplelist:
assert len(flatchain) == full_iterations, "Not all chains have the same number of iterations!"
assert flatchain.shape == shape, "Not all flatchains have the same shape!"
#make sure all chains have the same number of parameters.
#Following Gelman,
# n = length of split chains
# i = index of iteration in chain
# m = number of split chains
# j = index of which chain
n = full_iterations//2
m = 2 * len(samplelist)
nparams = samplelist[0].shape[-1] #the trailing dimension of a flatchain
#Block the chains up into a 3D array
chains = np.empty((n, m, nparams))
for k, flatchain in enumerate(samplelist):
chains[:,2*k,:] = flatchain[:n] #first half of chain
chains[:,2*k + 1,:] = flatchain[n:] #second half of chain
#Now compute statistics
#average value of each chain
avg_phi_j = np.mean(chains, axis=0, dtype="f8") #average over iterations, now a (m, nparams) array
#average value of all chains
avg_phi = np.mean(chains, axis=(0,1), dtype="f8") #average over iterations and chains, now a (nparams,) array
B = n/(m - 1.0) * np.sum((avg_phi_j - avg_phi)**2, axis=0, dtype="f8") #now a (nparams,) array
s2j = 1./(n - 1.) * np.sum((chains - avg_phi_j)**2, axis=0, dtype="f8") #now a (m, nparams) array
W = 1./m * np.sum(s2j, axis=0, dtype="f8") #now a (nparams,) arary
var_hat = (n - 1.)/n * W + B/n #still a (nparams,) array
std_hat = np.sqrt(var_hat)
R_hat = np.sqrt(var_hat/W) #still a (nparams,) array
data = Table({ "Value": avg_phi,
"Uncertainty": std_hat},
names=["Value", "Uncertainty"])
print(data)
ascii.write(data, sys.stdout, Writer = ascii.Latex, formats={"Value":"%0.3f", "Uncertainty":"%0.3f"}) #
#print("Average parameter value: {}".format(avg_phi))
#print("std_hat: {}".format(np.sqrt(var_hat)))
print("R_hat: {}".format(R_hat))
if np.any(R_hat >= 1.1):
print("You might consider running the chain for longer. Not all R_hats are less than 1.1.")
def estimate_covariance(flatchain, ndim=0):
if ndim == 0:
d = flatchain.shape[1]
else:
d = ndim
import matplotlib.pyplot as plt
#print("Parameters {}".format(flatchain.param_tuple))
#samples = flatchain.samples
cov = np.cov(flatchain, rowvar=0)
#Now try correlation coefficient
cor = np.corrcoef(flatchain, rowvar=0)
# Make a plot of correlation coefficient.
fig, ax = plt.subplots(figsize=(0.5 * d, 0.5 * d), nrows=1, ncols=1)
ext = (0.5, d + 0.5, 0.5, d + 0.5)
ax.imshow(cor, origin="upper", vmin=-1, vmax=1, cmap="bwr", interpolation="none", extent=ext)
fig.savefig("cor_coefficient.png")
print("'Optimal' jumps with covariance (units squared)")
opt_jump = 2.38**2/d * cov
# opt_jump = 1.7**2/d * cov # gives about ??
std_dev = np.sqrt(np.diag(cov))
print("'Optimal' jumps")
print(2.38/np.sqrt(d) * std_dev)
return opt_jump
def main():
# "K", "e", "omega", "P", "T0", "gamma", "amp_f", "l_g"
p = np.array([0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
p_return = convert_vector(p, "SB1", ["gamma", "e"], gamma=1.0, e=-0.4)
print(p_return)
if __name__=="__main__":
main()