-
Notifications
You must be signed in to change notification settings - Fork 1
/
EDGES_EMCEE_V2.py
143 lines (125 loc) · 5.37 KB
/
EDGES_EMCEE_V2.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
# Sample EDGES likelihood with EMCEE in yabf
from edges_estimate.likelihoods import LinearFG
from edges_cal.modelling import LinLog
import attr
# yabf: for MCMC
from yabf import Component, Parameter
from yabf.samplers import emcee
import py21cmfast as p21c
import numpy as np
from functools import cached_property
from scipy.interpolate import InterpolatedUnivariateSpline as spline
# print('Do conda activate base first')
@attr.s
class AbsorptionProfile(Component):
provides = ["eor_spectrum"]
# Obviously, put in the parameters of 21cmFAST below. Each can take a min and max.
# The parameters defined here are just the "possible" parameters to fit, they
# don't define the actively fit parameters in the likelihood itself.
# As a pre-flight test, let's fit fR and L_X param first
# CJS: How does yabf know whether fR and L_X are astro (and not cosmo)?
base_parameters = [
Parameter(name="fR", fiducial=2.0, min=1.0, max=7.0, latex="f_{R}"),
Parameter(name="L_X", fiducial=37.0, min=35.0, max=45.0, latex="L_x"),
]
observed_redshifts: np.ndarray = attr.ib(kw_only=True, eq=attr.cmp_using(eq=np.array_equal))
# Add other initialization parameters you require here, using attr.ib(). Eg.
user_params: p21c.UserParams = attr.ib()
cosmo_params: p21c.CosmoParams = attr.ib()
flag_options: p21c.FlagOptions = attr.ib()
astro_params: p21c.AstroParams = attr.ib()
run_lightcone_kwargs: dict = attr.ib(factory=dict)
# cache_loc: str = attr.ib()
# set default in attr.ib()
# cache_loc: str = '/home/dm/watson/21cmFAST-data/'
cache_loc = attr.ib(default="/home/dm/watson/21cmFAST-data/cache/")
# You probably want to run the initial conditions etc and cache them...
@cached_property
def initial_conditions(self):
return p21c.initial_conditions(
user_params=self.user_params,
cosmo_params=self.cosmo_params,
direc=self.cache_loc
)
def calculate(self, ctx, **params):
# This is the thing that actually has to calculate the global signal.
# "params" is a dictionary of {name: value}, where the "name" is one of the
# names defined in your base_parameters.
# So, something like this...
self.astro_params.update(**params)
lc = p21c.run_lightcone(
astro_params=self.astro_params,
init_box = self.initial_conditions,
flag_options=self.flag_options,
direc=self.cache_loc,
redshift=self.observed_redshifts.min(),
max_redshift=self.observed_redshifts.max(),
**self.run_lightcone_kwargs
)
# spline requires z to be in increasing order
z = lc.node_redshifts[-1:0:-1]
T21 = lc.global_brightness_temp[-1:0:-1]
return spline(z, T21)(self.observed_redshifts) / 1000.0
def spectrum(self, ctx, **params):
# Don't change this.
return ctx["eor_spectrum"]
# ---- Let's run it! ----
data = np.genfromtxt("Data_EDGES.txt")
freq = data[:, 0]
wght = data[:, 1]
tsky = data[:, 2]
freq = freq[wght>0]
tsky = tsky[wght>0]
# Let's fix these params around the best-guess settings
user_params = p21c.UserParams(
BOX_LEN = 150,
HII_DIM = 20, # Should be at least 50 for the official run
N_THREADS = 10 # Set to 1 if using MPIRUN
)
astro_params = p21c.AstroParams(
F_STAR10 = -0.8,
F_STAR7_MINI = -2.0,
M_TURN = 7.1,
NU_X_THRESH = 200,
t_STAR = 0.3
)
flag_options = p21c.FlagOptions(
USE_MASS_DEPENDENT_ZETA = True,
USE_TS_FLUCT = True,
INHOMO_RECO = True,
USE_RADIO_ACG = True
)
eor = AbsorptionProfile(
observed_redshifts = 1420/freq - 1,
user_params = user_params,
cosmo_params = p21c.CosmoParams(),
flag_options = flag_options,
astro_params = astro_params,
params = {
'fR': {'min': 2.0, 'max': 7.0},
'L_X':{'min':37.0,'max':45.0}
}, # these are the params that are actually fit. The names have to be in the `base_parameters` above
cache_loc = '/home/dm/watson/21cmFAST-data/cache/',
# run_lightcone_kwargs = {"ZPRIME_STEP_FACTOR": 1.03}
run_lightcone_kwargs = {"ZPRIME_STEP_FACTOR": 1.03, 'write':False} # do not write files
)
fg_model = LinLog(n_terms=5)
my_likelihood = LinearFG(freq=freq, t_sky=tsky, var=0.03**2, fg=fg_model, eor=eor)
# You can call the likelihood like this:
# params here should be fiducials for params you want to fit
# logp is log(posterior)
# a = my_likelihood.partial_linear_model.logp(params=[2, 42.0])
# print(a)
sampler = emcee(
my_likelihood.partial_linear_model, # The actual likelihood to sample from
save_full_config = False, # Otherwise would save a YAML file that is hard to read.
output_dir = "Chains", # Directory in which to save all the output chains.
output_prefix = "PopII_Test_EMCEE", # A prefix for all files output.
# sampler_kwargs = { # Anything that can be passed to PolychordSettings,
# "nlives": 256 # see https://github.com/PolyChord/PolyChordLite/pypolychord/settings.py#L5
# }
)
# Actually run the sampling. You'll get a bunch of files in the output dir, that can
# be read by getdist.
samples = sampler.sample(nsteps=10000,progress = True)
# samples.saveAsText(root)