Note
Go to the end to download the full example code.
MCMC search v.s. grid search
An example to compare MCMCSearch and GridSearch on the same data.
8 import os
9
10 import matplotlib.pyplot as plt
11 import numpy as np
12
13 import pyfstat
14
15 # flip this switch for a more expensive 4D (F0,F1,Alpha,Delta) run
16 # instead of just (F0,F1)
17 # (still only a few minutes on current laptops)
18 sky = False
19
20 label = "PyFstatExampleSimpleMCMCvsGridComparison"
21 outdir = os.path.join("PyFstat_example_data", label)
22 logger = pyfstat.set_up_logger(label=label, outdir=outdir)
23 if sky:
24 outdir += "AlphaDelta"
25
26 # parameters for the data set to generate
27 tstart = 1000000000
28 duration = 30 * 86400
29 Tsft = 1800
30 detectors = "H1,L1"
31 sqrtSX = 1e-22
32
33 # parameters for injected signals
34 inj = {
35 "tref": tstart,
36 "F0": 30.0,
37 "F1": -1e-10,
38 "F2": 0,
39 "Alpha": 0.5,
40 "Delta": 1,
41 "h0": 0.05 * sqrtSX,
42 "cosi": 1.0,
43 }
44
45 # latex-formatted plotting labels
46 labels = {
47 "F0": "$f$ [Hz]",
48 "F1": "$\\dot{f}$ [Hz/s]",
49 "2F": "$2\\mathcal{F}$",
50 "Alpha": "$\\alpha$",
51 "Delta": "$\\delta$",
52 }
53 labels["max2F"] = "$\\max\\,$" + labels["2F"]
54
55
56 def plot_grid_vs_samples(grid_res, mcmc_res, xkey, ykey):
57 """local plotting function to avoid code duplication in the 4D case"""
58 plt.plot(grid_res[xkey], grid_res[ykey], ".", label="grid")
59 plt.plot(mcmc_res[xkey], mcmc_res[ykey], ".", label="mcmc")
60 plt.plot(inj[xkey], inj[ykey], "*k", label="injection")
61 grid_maxidx = np.argmax(grid_res["twoF"])
62 mcmc_maxidx = np.argmax(mcmc_res["twoF"])
63 plt.plot(
64 grid_res[xkey][grid_maxidx],
65 grid_res[ykey][grid_maxidx],
66 "+g",
67 label=labels["max2F"] + "(grid)",
68 )
69 plt.plot(
70 mcmc_res[xkey][mcmc_maxidx],
71 mcmc_res[ykey][mcmc_maxidx],
72 "xm",
73 label=labels["max2F"] + "(mcmc)",
74 )
75 plt.xlabel(labels[xkey])
76 plt.ylabel(labels[ykey])
77 plt.legend()
78 plotfilename_base = os.path.join(outdir, "grid_vs_mcmc_{:s}{:s}".format(xkey, ykey))
79 plt.savefig(plotfilename_base + ".png")
80 if xkey == "F0" and ykey == "F1":
81 plt.xlim(zoom[xkey])
82 plt.ylim(zoom[ykey])
83 plt.savefig(plotfilename_base + "_zoom.png")
84 plt.close()
85
86
87 def plot_2F_scatter(res, label, xkey, ykey):
88 """local plotting function to avoid code duplication in the 4D case"""
89 markersize = 3 if label == "grid" else 1
90 sc = plt.scatter(res[xkey], res[ykey], c=res["twoF"], s=markersize)
91 cb = plt.colorbar(sc)
92 plt.xlabel(labels[xkey])
93 plt.ylabel(labels[ykey])
94 cb.set_label(labels["2F"])
95 plt.title(label)
96 plt.plot(inj[xkey], inj[ykey], "*k", label="injection")
97 maxidx = np.argmax(res["twoF"])
98 plt.plot(
99 res[xkey][maxidx],
100 res[ykey][maxidx],
101 "+r",
102 label=labels["max2F"],
103 )
104 plt.legend()
105 plotfilename_base = os.path.join(
106 outdir, "{:s}_{:s}{:s}_2F".format(label, xkey, ykey)
107 )
108 plt.xlim([min(res[xkey]), max(res[xkey])])
109 plt.ylim([min(res[ykey]), max(res[ykey])])
110 plt.savefig(plotfilename_base + ".png")
111 plt.close()
112
113
114 if __name__ == "__main__":
115 logger.info("Generating SFTs with injected signal...")
116 writer = pyfstat.Writer(
117 label=label + "SimulatedSignal",
118 outdir=outdir,
119 tstart=tstart,
120 duration=duration,
121 detectors=detectors,
122 sqrtSX=sqrtSX,
123 Tsft=Tsft,
124 **inj,
125 Band=1, # default band estimation would be too narrow for a wide grid/prior
126 )
127 writer.make_data()
128
129 # set up square search grid with fixed (F0,F1) mismatch
130 # and (optionally) some ad-hoc sky coverage
131 m = 0.001
132 dF0 = np.sqrt(12 * m) / (np.pi * duration)
133 dF1 = np.sqrt(180 * m) / (np.pi * duration**2)
134 DeltaF0 = 500 * dF0
135 DeltaF1 = 200 * dF1
136 if sky:
137 # cover less range to keep runtime down
138 DeltaF0 /= 10
139 DeltaF1 /= 10
140 F0s = [inj["F0"] - DeltaF0 / 2.0, inj["F0"] + DeltaF0 / 2.0, dF0]
141 F1s = [inj["F1"] - DeltaF1 / 2.0, inj["F1"] + DeltaF1 / 2.0, dF1]
142 F2s = [inj["F2"]]
143 search_keys = ["F0", "F1"] # only the ones that aren't 0-width
144 if sky:
145 dSky = 0.01 # rather coarse to keep runtime down
146 DeltaSky = 10 * dSky
147 Alphas = [inj["Alpha"] - DeltaSky / 2.0, inj["Alpha"] + DeltaSky / 2.0, dSky]
148 Deltas = [inj["Delta"] - DeltaSky / 2.0, inj["Delta"] + DeltaSky / 2.0, dSky]
149 search_keys += ["Alpha", "Delta"]
150 else:
151 Alphas = [inj["Alpha"]]
152 Deltas = [inj["Delta"]]
153 search_keys_label = "".join(search_keys)
154
155 logger.info("Performing GridSearch...")
156 gridsearch = pyfstat.GridSearch(
157 label="GridSearch" + search_keys_label,
158 outdir=outdir,
159 sftfilepattern=writer.sftfilepath,
160 F0s=F0s,
161 F1s=F1s,
162 F2s=F2s,
163 Alphas=Alphas,
164 Deltas=Deltas,
165 tref=inj["tref"],
166 )
167 gridsearch.run()
168 gridsearch.print_max_twoF()
169 gridsearch.generate_loudest()
170
171 # do some plots of the GridSearch results
172 if not sky: # this plotter can't currently deal with too large result arrays
173 logger.info("Plotting 1D 2F distributions...")
174 for key in search_keys:
175 gridsearch.plot_1D(xkey=key, xlabel=labels[key], ylabel=labels["2F"])
176
177 logger.info("Making GridSearch {:s} corner plot...".format("-".join(search_keys)))
178 vals = [np.unique(gridsearch.data[key]) - inj[key] for key in search_keys]
179 twoF = gridsearch.data["twoF"].reshape([len(kval) for kval in vals])
180 corner_labels = [
181 "$f - f_0$ [Hz]",
182 "$\\dot{f} - \\dot{f}_0$ [Hz/s]",
183 ]
184 if sky:
185 corner_labels.append("$\\alpha - \\alpha_0$")
186 corner_labels.append("$\\delta - \\delta_0$")
187 corner_labels.append(labels["2F"])
188 gridcorner_fig, gridcorner_axes = pyfstat.gridcorner(
189 twoF, vals, projection="log_mean", labels=corner_labels, whspace=0.1, factor=1.8
190 )
191 gridcorner_fig.savefig(os.path.join(outdir, gridsearch.label + "_corner.png"))
192 plt.close(gridcorner_fig)
193
194 logger.info("Performing MCMCSearch...")
195 # set up priors in F0 and F1 (over)covering the grid ranges
196 if sky: # MCMC will still be fast in 4D with wider range than grid
197 DeltaF0 *= 50
198 DeltaF1 *= 50
199 theta_prior = {
200 "F0": {
201 "type": "unif",
202 "lower": inj["F0"] - DeltaF0 / 2.0,
203 "upper": inj["F0"] + DeltaF0 / 2.0,
204 },
205 "F1": {
206 "type": "unif",
207 "lower": inj["F1"] - DeltaF1 / 2.0,
208 "upper": inj["F1"] + DeltaF1 / 2.0,
209 },
210 "F2": inj["F2"],
211 }
212 if sky:
213 # also implicitly covering twice the grid range here
214 theta_prior["Alpha"] = {
215 "type": "unif",
216 "lower": inj["Alpha"] - DeltaSky,
217 "upper": inj["Alpha"] + DeltaSky,
218 }
219 theta_prior["Delta"] = {
220 "type": "unif",
221 "lower": inj["Delta"] - DeltaSky,
222 "upper": inj["Delta"] + DeltaSky,
223 }
224 else:
225 theta_prior["Alpha"] = inj["Alpha"]
226 theta_prior["Delta"] = inj["Delta"]
227 # ptemcee sampler settings - in a real application we might want higher values
228 ntemps = 2
229 log10beta_min = -1
230 nwalkers = 100
231 nsteps = [200, 200] # [burn-in,production]
232
233 mcmcsearch = pyfstat.MCMCSearch(
234 label="MCMCSearch" + search_keys_label,
235 outdir=outdir,
236 sftfilepattern=writer.sftfilepath,
237 theta_prior=theta_prior,
238 tref=inj["tref"],
239 nsteps=nsteps,
240 nwalkers=nwalkers,
241 ntemps=ntemps,
242 log10beta_min=log10beta_min,
243 )
244 # walker plot is generated during main run of the search class
245 mcmcsearch.run(
246 walker_plot_args={"plot_det_stat": True, "injection_parameters": inj}
247 )
248 mcmcsearch.print_summary()
249
250 # call some built-in plotting methods
251 # these can all highlight the injection parameters, too
252 logger.info("Making MCMCSearch {:s} corner plot...".format("-".join(search_keys)))
253 mcmcsearch.plot_corner(truths=inj)
254 logger.info("Making MCMCSearch prior-posterior comparison plot...")
255 mcmcsearch.plot_prior_posterior(injection_parameters=inj)
256
257 # NOTE: everything below here is just custom commandline output and plotting
258 # for this particular example, which uses the PyFstat outputs,
259 # but isn't very instructive if you just want to learn the main usage of the package.
260
261 # some informative command-line output comparing search results and injection
262 # get max of GridSearch, contains twoF and all Doppler parameters in the dict
263 max_dict_grid = gridsearch.get_max_twoF()
264 # same for MCMCSearch, here twoF is separate, and non-sampled parameters are not included either
265 max_dict_mcmc, max_2F_mcmc = mcmcsearch.get_max_twoF()
266 logger.info(
267 "max2F={:.4f} from GridSearch, offsets from injection: {:s}.".format(
268 max_dict_grid["twoF"],
269 ", ".join(
270 [
271 "{:.4e} in {:s}".format(max_dict_grid[key] - inj[key], key)
272 for key in search_keys
273 ]
274 ),
275 )
276 )
277 logger.info(
278 "max2F={:.4f} from MCMCSearch, offsets from injection: {:s}.".format(
279 max_2F_mcmc,
280 ", ".join(
281 [
282 "{:.4e} in {:s}".format(max_dict_mcmc[key] - inj[key], key)
283 for key in search_keys
284 ]
285 ),
286 )
287 )
288 # get additional point and interval estimators
289 stats_dict_mcmc = mcmcsearch.get_summary_stats()
290 logger.info(
291 "mean from MCMCSearch: offset from injection by {:s},"
292 " or in fractions of 2sigma intervals: {:s}.".format(
293 ", ".join(
294 [
295 "{:.4e} in {:s}".format(
296 stats_dict_mcmc[key]["mean"] - inj[key], key
297 )
298 for key in search_keys
299 ]
300 ),
301 ", ".join(
302 [
303 "{:.2f}% in {:s}".format(
304 100
305 * np.abs(stats_dict_mcmc[key]["mean"] - inj[key])
306 / (2 * stats_dict_mcmc[key]["std"]),
307 key,
308 )
309 for key in search_keys
310 ]
311 ),
312 )
313 )
314 logger.info(
315 "median from MCMCSearch: offset from injection by {:s},"
316 " or in fractions of 90% confidence intervals: {:s}.".format(
317 ", ".join(
318 [
319 "{:.4e} in {:s}".format(
320 stats_dict_mcmc[key]["median"] - inj[key], key
321 )
322 for key in search_keys
323 ]
324 ),
325 ", ".join(
326 [
327 "{:.2f}% in {:s}".format(
328 100
329 * np.abs(stats_dict_mcmc[key]["median"] - inj[key])
330 / (
331 stats_dict_mcmc[key]["upper90"]
332 - stats_dict_mcmc[key]["lower90"]
333 ),
334 key,
335 )
336 for key in search_keys
337 ]
338 ),
339 )
340 )
341
342 # do additional custom plotting
343 logger.info("Loading grid and MCMC search results for custom comparison plots...")
344 gridfile = os.path.join(outdir, gridsearch.label + "_NA_GridSearch.txt")
345 if not os.path.isfile(gridfile):
346 raise RuntimeError(
347 "Failed to load GridSearch results from file '{:s}',"
348 " something must have gone wrong!".format(gridfile)
349 )
350 grid_res = pyfstat.utils.read_txt_file_with_header(gridfile)
351 mcmc_file = os.path.join(outdir, mcmcsearch.label + "_samples.dat")
352 if not os.path.isfile(mcmc_file):
353 raise RuntimeError(
354 "Failed to load MCMCSearch results from file '{:s}',"
355 " something must have gone wrong!".format(mcmc_file)
356 )
357 mcmc_res = pyfstat.utils.read_txt_file_with_header(mcmc_file)
358
359 zoom = {
360 "F0": [inj["F0"] - 10 * dF0, inj["F0"] + 10 * dF0],
361 "F1": [inj["F1"] - 5 * dF1, inj["F1"] + 5 * dF1],
362 }
363
364 # we'll use the two local plotting functions defined above
365 # to avoid code duplication in the sky case
366 logger.info("Creating MCMC-grid comparison plots...")
367 plot_grid_vs_samples(grid_res, mcmc_res, "F0", "F1")
368 plot_2F_scatter(grid_res, "grid", "F0", "F1")
369 plot_2F_scatter(mcmc_res, "mcmc", "F0", "F1")
370 if sky:
371 plot_grid_vs_samples(grid_res, mcmc_res, "Alpha", "Delta")
372 plot_2F_scatter(grid_res, "grid", "Alpha", "Delta")
373 plot_2F_scatter(mcmc_res, "mcmc", "Alpha", "Delta")