-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathSG_MCMC_FLAT.py
120 lines (80 loc) · 3.5 KB
/
SG_MCMC_FLAT.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from __future__ import division
import numpy as np
import emcee
import scipy.integrate as integrate
from scipy.special import betainc, gamma
import math
from time import strftime
import sys
c1 = 2.9979E5 #km / s
# In[2]:
import os
data_file_name = sys.argv[1]
datadir='/global/u2/k/kap146/carver/GMM_Likelihood/Mock_Data' #/global/scratch2/sd/kap146/Input_Data'
#out_datadir = '/global/scratch2/sd/kap146/Output_Data'
datafile_1G =os.path.join(datadir, data_file_name)
M, z, dm = np.loadtxt(datafile_1G, unpack = True, usecols=[0, 1, 2])
# In[5]:
def distmod(Om,w,z):
Ol = 1.0 - Om
x0 = Ol/(Om+Ol)
x = Ol*(1.+z)**(3.*w)/(Om + Ol*(1.+z)**(3.*w))
m = -1./(6.*w)
bx0 = betainc(m, 0.5-m, x0)*gamma(m)*gamma(0.5-m)/gamma(0.5)
bx = betainc(m, 0.5-m, x)*gamma(m)*gamma(0.5-m)/gamma(0.5)
r = 2.*m/np.sqrt(Om) * (Om/Ol)**m *(bx0 - bx)
b = (1.0+z)*r
return 5.0*np.log10(b)
# Define log likelihood function.
def lnlike(theta, x, y, yerr):
z = x
Omega_M, w, sig_int, script_M = theta
var = yerr*yerr + sig_int*sig_int
model = script_M + distmod(Omega_M, w, z) + 5.0*np.log10(c1)
population = np.exp(-0.5* (y-model)**2./var) / (np.sqrt(2.0*np.pi*var))
return np.sum(np.log(population))
# In[6]:
def lnprior(theta):
Omega_M, w, sig_int, script_M = theta
if -15.0 < script_M < 5.0 and 0.001 < Omega_M < 1.0 and -3.0 < w < 1.0 and 0.0 < sig_int < 0.3:
return 0.0
return -np.inf
# In[7]:
def lnprob(theta, x, y, yerr):
lp = lnprior(theta)
if not np.isfinite(lp):
return -np.inf
return lp + lnlike(theta, x, y, yerr)
# In[8]:
sig = 0.10 # assuming know SN very well
script_M_true, Omega_M_true, w_true, sig_int_true = -3.6, 0.30, -0.98, 0.14 #-3.704, 0.30, 0.70, -1.0, 0.15
# Set starting position for parameters. Different options for fitting different parameters
initial = Omega_M_true, w_true, sig_int_true, script_M_true
# In[13]:
ndim, nwalkers = len(initial), 50
step = 100
thin_f = 1
pos0 = [initial + 1e-2*np.random.randn(ndim) for i in range(nwalkers)]
# In[14]:
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, a = 2.5,threads = 8, args = (z, dm, sig))
# In[15]:
## outputs position, probability, and random number seed
pos, prob, state = sampler.run_mcmc(pos0, step/10, thin = 2) ## (initial position, number of steps)
sampler.reset() ## reset cha
# In[16]:
pos, prob, state = sampler.run_mcmc(pos0, step, rstate0=state, thin = thin_f)
# Calculates percent accepted. and tells you how to change. If it is around 0.0, you have a bigger problem.
af = sampler.acceptance_fraction
#np.savetxt('/global/scratch2/sd/kap146/Output_Data/%s_accepatance_fraction_1G_%s_st_%s_d_%s_w_%s_data.dat'%(strftime('%Y%m%d'),step, ndim, nwalkers, len(M)), af)
af_arr = np.mean(af, dtype = np.float64)
print af_arr
f = open('/global/scratch2/sd/kap146/Output_Data/%s_accepatance_fraction_1G_%s_st_%s_d_%s_w_%s_data_0.dat'%(strftime('%Y%m%d'),step,ndim, nwalkers, len(M)), 'w')
f.write('%.2f' % af_arr)
f.close()
## Print the estimated values from the MCMC
fit = np.mean(pos, axis = 0, dtype=np.float64)
np.savetxt('/global/scratch2/sd/kap146/Output_Data/%s_est_param_1G_%s_st_%s_d_%s_w_%s_data_0.dat'%(strftime('%Y%m%d'),step, ndim, nwalkers, len(M)), (fit))
######################### Plotting #########################
# Flatten samples
samples = sampler.chain[:, :, :].reshape((-1, ndim))
np.savetxt("/global/scratch2/sd/kap146/Output_Data/%s_cosmofit_1G_%s_st_%s_d_%s_w_%s_data_samples_0.dat"%(strftime('%Y%m%d'),step, ndim, nwalkers, len(M)), samples)