forked from allisony/lyapy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathGJ176_MCMC_d2h_fixed.py
239 lines (178 loc) · 9.18 KB
/
GJ176_MCMC_d2h_fixed.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
# -*- coding: utf-8 -*-
import emcee
import numpy as np
import matplotlib.pyplot as plt
import astropy.io.fits as pyfits
import lyapy
import lya_plot
plt.ion()
## Change this value around sometimes
np.random.seed(82)
## Fixing the D/H ratio at 1.5e-5. 1.56e-5 might be a better value to use though (Wood+ 2004 is the reference, I believe)
d2h_true = 1.5e-5
descrip = '_d2h_fixed' ## appended to saved files throughout
## MCMC parameters
ndim = 9 #ndim is number of fitted parameters
nwalkers = 30 # must be even, not odd
nsteps = 500
burnin = 300
## Read in fits file ##
## I've used my lyapy.geocoronal_subtraction function to put the FITS file into the
## format I want... hence the "modAY" suffix to the file below.
input_filename = 'p_msl_pan_-----_gj176_panspec_native_resolution_waverange1100.0-1300.0_modAY.fits'
spec_hdu = pyfits.open(input_filename)
spec = spec_hdu[1].data
spec_header = spec_hdu[1].header
## Define wave, flux, and error variables ##
wave_to_fit = spec['wave']
flux_to_fit = spec['flux']
error_to_fit = spec['error']
resolution = float(spec_header['RES_POW'])
## This part is just making sure the error bars in the low-flux wings aren't smaller than the RMS
## in the wings
rms_wing = np.std(flux_to_fit[np.where((wave_to_fit > np.min(wave_to_fit)) & (wave_to_fit < 1214.5))])
error_to_fit[np.where(error_to_fit < rms_wing)] = rms_wing
## Masking the core of the HI absorption because it may be contaminated by geocoronal emission
## 18 Jan 2018 - I don't think this part is actually necessary. The convolution should fit this inner region well.
mask = lyapy.mask_spectrum(flux_to_fit,interactive=False,mask_lower_limit=36.,mask_upper_limit=42.)
flux_masked = np.ma.masked_array(spec['flux'],mask=mask)
wave_masked = np.ma.masked_array(spec['wave'],mask=mask)
error_masked = np.ma.masked_array(spec['error'],mask=mask)
## Defining parameter ranges. Below I use uniform priors for most of the parameters -- as long
## as they fit inside these ranges.
## 18 Jan 2018 -- To Do: There should be a way here to fix certain parameters.
vs_n_min = 0.
vs_n_max = 75.
am_n_min = -13.1
am_n_max = -11.8
fw_n_min = 125.
fw_n_max = 275.
vs_b_min = -5.
vs_b_max = 100.
am_b_min = -15.4
am_b_max = -13.4
fw_b_min = 300.
fw_b_max = 700.
h1_col_min = 17.0
h1_col_max = 19.5
h1_b_min = 1.
h1_b_max = 20.
h1_vel_min = 20.
h1_vel_max = 100.
# Define the probability function as likelihood * prior.
# I use a flat/uniform prior for everything but h1_b.
def lnprior(theta):
vs_n, am_n, fw_n, vs_b, am_b, fw_b, h1_col, h1_b, h1_vel = theta
if (vs_n_min < vs_n < vs_n_max) and (am_n_min < am_n < am_n_max) and (fw_n_min < fw_n < fw_n_max) and (vs_b_min < vs_b < vs_b_max) and (am_b_min < am_b < am_b_max) and (fw_b_min < fw_b < fw_b_max) and (h1_col_min < h1_col < h1_col_max) and (h1_b_min < h1_b < h1_b_max) and (h1_vel_min < h1_vel < h1_vel_max):
return np.log(h1_b)
return -np.inf
def lnlike(theta, x, y, yerr):
vs_n, am_n, fw_n, vs_b, am_b, fw_b, h1_col, h1_b, h1_vel = theta
y_model = lyapy.damped_lya_profile(x,vs_n,10**am_n,fw_n,vs_b,10**am_b,fw_b,h1_col,
h1_b,h1_vel,d2h=d2h_true,resolution=resolution,
single_component_flux=False)/1e14
return -0.5 * np.sum(np.log(2 * np.pi * yerr**2) + (y - y_model) ** 2 / yerr**2)
def lnprob(theta, x, y, yerr):
lp = lnprior(theta)
if not np.isfinite(lp):
return -np.inf
return lp + lnlike(theta, x, y, yerr)
# Set up the sampler. There are multiple ways to initialize the walkers,
# and I chose uniform sampling of the parameter ranges.
pos = [np.array([np.random.uniform(low=vs_n_min,high=vs_n_max,size=1)[0],
np.random.uniform(low=am_n_min,high=am_n_max,size=1)[0],
np.random.uniform(low=fw_n_min,high=fw_n_max,size=1)[0],
np.random.uniform(low=vs_b_min,high=vs_b_max,size=1)[0],
np.random.uniform(low=am_b_min,high=am_b_max,size=1)[0],
np.random.uniform(low=fw_b_min,high=fw_b_max,size=1)[0],
np.random.uniform(low=h1_col_min,high=h1_col_max,size=1)[0],
np.random.uniform(low=h1_b_min,high=h1_b_max,size=1)[0],
np.random.uniform(low=h1_vel_min,high=h1_vel_max,size=1)[0]]) for i in range(nwalkers)]
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(wave_masked,flux_masked,error_masked))
vs_n_pos = np.zeros(len(pos))
am_n_pos = np.zeros(len(pos))
fw_n_pos = np.zeros(len(pos))
vs_b_pos = np.zeros(len(pos))
am_b_pos = np.zeros(len(pos))
fw_b_pos = np.zeros(len(pos))
h1_col_pos = np.zeros(len(pos))
h1_b_pos = np.zeros(len(pos))
h1_vel_pos = np.zeros(len(pos))
for i in range(len(pos)):
vs_n_pos[i] = pos[i][0]
am_n_pos[i] = pos[i][1]
fw_n_pos[i] = pos[i][2]
vs_b_pos[i] = pos[i][3]
am_b_pos[i] = pos[i][4]
fw_b_pos[i] = pos[i][5]
h1_col_pos[i] = pos[i][6]
h1_b_pos[i] = pos[i][7]
h1_vel_pos[i] = pos[i][8]
# Clear and run the production chain.
print("Running MCMC...")
sampler.run_mcmc(pos, nsteps, rstate0=np.random.get_state())
print("Done.")
## remove the burn-in period from the sampler
samples = sampler.chain[:, burnin:, :].reshape((-1, ndim))
## print best fit parameters + uncertainties
vs_n_mcmc, am_n_mcmc, fw_n_mcmc, vs_b_mcmc, am_b_mcmc, fw_b_mcmc, h1_col_mcmc, h1_b_mcmc, \
h1_vel_mcmc = map(lambda v: [v[1], v[2]-v[1], v[1]-v[0]], \
zip(*np.percentile(samples, [16, 50, 84], axis=0)))
print("""MCMC result:
vs_n = {0[0]} +{0[1]} -{0[2]}
am_n = {1[0]} +{1[1]} -{1[2]}
fw_n = {2[0]} +{2[1]} -{2[2]}
vs_b = {3[0]} +{3[1]} -{3[2]}
am_b = {4[0]} +{4[1]} -{4[2]}
fw_b = {5[0]} +{5[1]} -{5[2]}
h1_col = {6[0]} +{6[1]} -{6[2]}
h1_b = {7[0]} +{7[1]} -{7[2]}
h1_vel = {8[0]} +{8[1]} -{8[2]}
""".format(vs_n_mcmc, am_n_mcmc, fw_n_mcmc, vs_b_mcmc, am_b_mcmc, fw_b_mcmc, h1_col_mcmc,
h1_b_mcmc, h1_vel_mcmc))
print("Mean acceptance fraction: {0:.3f}"
.format(np.mean(sampler.acceptance_fraction)))
print("should be between 0.25 and 0.5")
## best fit intrinsic profile
lya_intrinsic_profile_mcmc,lya_intrinsic_flux_mcmc = lyapy.lya_intrinsic_profile_func(wave_to_fit,
vs_n_mcmc[0],10**am_n_mcmc[0],fw_n_mcmc[0],vs_b_mcmc[0],10**am_b_mcmc[0],fw_b_mcmc[0],
return_flux=True)
## best fit attenuated profile
model_best_fit = lyapy.damped_lya_profile(wave_to_fit,vs_n_mcmc[0],10**am_n_mcmc[0],fw_n_mcmc[0],
vs_b_mcmc[0],10**am_b_mcmc[0],fw_b_mcmc[0],h1_col_mcmc[0],
h1_b_mcmc[0],h1_vel_mcmc[0],d2h_true,resolution,
single_component_flux=False)/1e14
## Here's the big messy part where I determine the 1-sigma error bars on the
## reconstructed, intrinsic LyA flux. From my paper: "For each of the 9 parameters,
## the best-fit values are taken as the 50th percentile (the median) of the marginalized
## distributions, and 1-σ error bars as the 16th and 84th percentiles (shown as dashed
## vertical lines in Figures 3 and 4). The best-fit reconstructed Lyα fluxes are determined
## from the best-fit amplitude, FWHM, and velocity centroid parameters, and the 1-σ error bars
## of the reconstructed Lyα flux are taken by varying these parameters individually between
## their 1-σ error bars and keeping all others fixed at their best-fit value.
## The resulting minimum and maximum Lyα fluxes become the 1-σ error bars (Table 2)."
lya_intrinsic_profile_limits = [] ## this is for showing the gray-shaded regions in my Figure 1
lya_intrinsic_flux_limits = [] ## this is for finding the 1-σ LyA flux error bars
vs_n_limits = [vs_n_mcmc[0] + vs_n_mcmc[1], vs_n_mcmc[0] - vs_n_mcmc[2]]
am_n_limits = [am_n_mcmc[0] + am_n_mcmc[1], am_n_mcmc[0] - am_n_mcmc[2]]
fw_n_limits = [fw_n_mcmc[0] + fw_n_mcmc[1], fw_n_mcmc[0] - fw_n_mcmc[2]]
vs_b_limits = [vs_b_mcmc[0] + vs_b_mcmc[1], vs_b_mcmc[0] - vs_b_mcmc[2]]
am_b_limits = [am_b_mcmc[0] + am_b_mcmc[1], am_b_mcmc[0] - am_b_mcmc[2]]
fw_b_limits = [fw_b_mcmc[0] + fw_b_mcmc[1], fw_b_mcmc[0] - fw_b_mcmc[2]]
h1_col_limits = [h1_col_mcmc[0] + h1_col_mcmc[1], h1_col_mcmc[0] - h1_col_mcmc[2]]
h1_b_limits = [h1_b_mcmc[0] + h1_b_mcmc[1], h1_b_mcmc[0] - h1_b_mcmc[2]]
h1_vel_limits = [h1_vel_mcmc[0] + h1_vel_mcmc[1], h1_vel_mcmc[0] - h1_vel_mcmc[2]]
## Whenever I say "damped" I mean "attenuated"
lya_plot.walkers(sampler)
lya_plot.corner(samples)
lya_plot.profile(samples, wave_to_fit, flux_to_fit, error_to_fit, resolution,
model_best_fit, lya_intrinsic_profile_mcmc, d2h_true = 1.5e-5)
## Estimating EUV spectrum -- 9.27 pc is the distance to GJ 176
lyapy.EUV_spectrum(lya_intrinsic_flux_mcmc,9.27,spec_header['STAR'],Mstar=True,savefile=True,doplot=True)
## Saving intrinsic profile
outfile_str = spec_header['STAR'] + '_LyA_intrinsic_profile.txt'
dx = wave_to_fit[1] - wave_to_fit[0]
wave_to_fit_extended = np.arange(1209.5,1222.0+dx,dx)
lya_intrinsic_profile_mcmc_extended = lyapy.lya_intrinsic_profile_func(wave_to_fit_extended,
vs_n_mcmc[0],10**am_n_mcmc[0],fw_n_mcmc[0],vs_b_mcmc[0],10**am_b_mcmc[0],fw_b_mcmc[0])
np.savetxt(outfile_str,np.transpose([wave_to_fit_extended,lya_intrinsic_profile_mcmc_extended]))