Note
Click here 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 outdir = os.path.join(
21 "PyFstat_example_data", "PyFstat_example_simple_mcmc_vs_grid_comparison"
22 )
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
116 print("Generating SFTs with injected signal...")
117 writer = pyfstat.Writer(
118 label="simulated_signal",
119 outdir=outdir,
120 tstart=tstart,
121 duration=duration,
122 detectors=detectors,
123 sqrtSX=sqrtSX,
124 Tsft=Tsft,
125 **inj,
126 Band=1, # default band estimation would be too narrow for a wide grid/prior
127 )
128 writer.make_data()
129 print("")
130
131 # set up square search grid with fixed (F0,F1) mismatch
132 # and (optionally) some ad-hoc sky coverage
133 m = 0.001
134 dF0 = np.sqrt(12 * m) / (np.pi * duration)
135 dF1 = np.sqrt(180 * m) / (np.pi * duration**2)
136 DeltaF0 = 500 * dF0
137 DeltaF1 = 200 * dF1
138 if sky:
139 # cover less range to keep runtime down
140 DeltaF0 /= 10
141 DeltaF1 /= 10
142 F0s = [inj["F0"] - DeltaF0 / 2.0, inj["F0"] + DeltaF0 / 2.0, dF0]
143 F1s = [inj["F1"] - DeltaF1 / 2.0, inj["F1"] + DeltaF1 / 2.0, dF1]
144 F2s = [inj["F2"]]
145 search_keys = ["F0", "F1"] # only the ones that aren't 0-width
146 if sky:
147 dSky = 0.01 # rather coarse to keep runtime down
148 DeltaSky = 10 * dSky
149 Alphas = [inj["Alpha"] - DeltaSky / 2.0, inj["Alpha"] + DeltaSky / 2.0, dSky]
150 Deltas = [inj["Delta"] - DeltaSky / 2.0, inj["Delta"] + DeltaSky / 2.0, dSky]
151 search_keys += ["Alpha", "Delta"]
152 else:
153 Alphas = [inj["Alpha"]]
154 Deltas = [inj["Delta"]]
155 search_keys_label = "".join(search_keys)
156
157 print("Performing GridSearch...")
158 gridsearch = pyfstat.GridSearch(
159 label="grid_search_" + search_keys_label,
160 outdir=outdir,
161 sftfilepattern=os.path.join(outdir, "*simulated_signal*sft"),
162 F0s=F0s,
163 F1s=F1s,
164 F2s=F2s,
165 Alphas=Alphas,
166 Deltas=Deltas,
167 tref=inj["tref"],
168 )
169 gridsearch.run()
170 gridsearch.print_max_twoF()
171 gridsearch.generate_loudest()
172
173 # do some plots of the GridSearch results
174 if not sky: # this plotter can't currently deal with too large result arrays
175 print("Plotting 1D 2F distributions...")
176 for key in search_keys:
177 gridsearch.plot_1D(xkey=key, xlabel=labels[key], ylabel=labels["2F"])
178
179 print("Making GridSearch {:s} corner plot...".format("-".join(search_keys)))
180 vals = [np.unique(gridsearch.data[key]) - inj[key] for key in search_keys]
181 twoF = gridsearch.data["twoF"].reshape([len(kval) for kval in vals])
182 corner_labels = [
183 "$f - f_0$ [Hz]",
184 "$\\dot{f} - \\dot{f}_0$ [Hz/s]",
185 ]
186 if sky:
187 corner_labels.append("$\\alpha - \\alpha_0$")
188 corner_labels.append("$\\delta - \\delta_0$")
189 corner_labels.append(labels["2F"])
190 gridcorner_fig, gridcorner_axes = pyfstat.gridcorner(
191 twoF, vals, projection="log_mean", labels=corner_labels, whspace=0.1, factor=1.8
192 )
193 gridcorner_fig.savefig(os.path.join(outdir, gridsearch.label + "_corner.png"))
194 plt.close(gridcorner_fig)
195 print("")
196
197 print("Performing MCMCSearch...")
198 # set up priors in F0 and F1 (over)covering the grid ranges
199 if sky: # MCMC will still be fast in 4D with wider range than grid
200 DeltaF0 *= 50
201 DeltaF1 *= 50
202 theta_prior = {
203 "F0": {
204 "type": "unif",
205 "lower": inj["F0"] - DeltaF0 / 2.0,
206 "upper": inj["F0"] + DeltaF0 / 2.0,
207 },
208 "F1": {
209 "type": "unif",
210 "lower": inj["F1"] - DeltaF1 / 2.0,
211 "upper": inj["F1"] + DeltaF1 / 2.0,
212 },
213 "F2": inj["F2"],
214 }
215 if sky:
216 # also implicitly covering twice the grid range here
217 theta_prior["Alpha"] = {
218 "type": "unif",
219 "lower": inj["Alpha"] - DeltaSky,
220 "upper": inj["Alpha"] + DeltaSky,
221 }
222 theta_prior["Delta"] = {
223 "type": "unif",
224 "lower": inj["Delta"] - DeltaSky,
225 "upper": inj["Delta"] + DeltaSky,
226 }
227 else:
228 theta_prior["Alpha"] = inj["Alpha"]
229 theta_prior["Delta"] = inj["Delta"]
230 # ptemcee sampler settings - in a real application we might want higher values
231 ntemps = 2
232 log10beta_min = -1
233 nwalkers = 100
234 nsteps = [200, 200] # [burnin,production]
235
236 mcmcsearch = pyfstat.MCMCSearch(
237 label="mcmc_search_" + search_keys_label,
238 outdir=outdir,
239 sftfilepattern=os.path.join(outdir, "*simulated_signal*sft"),
240 theta_prior=theta_prior,
241 tref=inj["tref"],
242 nsteps=nsteps,
243 nwalkers=nwalkers,
244 ntemps=ntemps,
245 log10beta_min=log10beta_min,
246 )
247 # walker plot is generated during main run of the search class
248 mcmcsearch.run(
249 walker_plot_args={"plot_det_stat": True, "injection_parameters": inj}
250 )
251 mcmcsearch.print_summary()
252
253 # call some built-in plotting methods
254 # these can all highlight the injection parameters, too
255 print("Making MCMCSearch {:s} corner plot...".format("-".join(search_keys)))
256 mcmcsearch.plot_corner(truths=inj)
257 print("Making MCMCSearch prior-posterior comparison plot...")
258 mcmcsearch.plot_prior_posterior(injection_parameters=inj)
259 print("")
260
261 # NOTE: everything below here is just custom commandline output and plotting
262 # for this particular example, which uses the PyFstat outputs,
263 # but isn't very instructive if you just want to learn the main usage of the package.
264
265 # some informative command-line output comparing search results and injection
266 # get max of GridSearch, contains twoF and all Doppler parameters in the dict
267 max_dict_grid = gridsearch.get_max_twoF()
268 # same for MCMCSearch, here twoF is separate, and non-sampled parameters are not included either
269 max_dict_mcmc, max_2F_mcmc = mcmcsearch.get_max_twoF()
270 print(
271 "max2F={:.4f} from GridSearch, offsets from injection: {:s}.".format(
272 max_dict_grid["twoF"],
273 ", ".join(
274 [
275 "{:.4e} in {:s}".format(max_dict_grid[key] - inj[key], key)
276 for key in search_keys
277 ]
278 ),
279 )
280 )
281 print(
282 "max2F={:.4f} from MCMCSearch, offsets from injection: {:s}.".format(
283 max_2F_mcmc,
284 ", ".join(
285 [
286 "{:.4e} in {:s}".format(max_dict_mcmc[key] - inj[key], key)
287 for key in search_keys
288 ]
289 ),
290 )
291 )
292 # get additional point and interval estimators
293 stats_dict_mcmc = mcmcsearch.get_summary_stats()
294 print(
295 "mean from MCMCSearch: offset from injection by {:s},"
296 " or in fractions of 2sigma intervals: {:s}.".format(
297 ", ".join(
298 [
299 "{:.4e} in {:s}".format(
300 stats_dict_mcmc[key]["mean"] - inj[key], key
301 )
302 for key in search_keys
303 ]
304 ),
305 ", ".join(
306 [
307 "{:.2f}% in {:s}".format(
308 100
309 * np.abs(stats_dict_mcmc[key]["mean"] - inj[key])
310 / (2 * stats_dict_mcmc[key]["std"]),
311 key,
312 )
313 for key in search_keys
314 ]
315 ),
316 )
317 )
318 print(
319 "median from MCMCSearch: offset from injection by {:s},"
320 " or in fractions of 90% confidence intervals: {:s}.".format(
321 ", ".join(
322 [
323 "{:.4e} in {:s}".format(
324 stats_dict_mcmc[key]["median"] - inj[key], key
325 )
326 for key in search_keys
327 ]
328 ),
329 ", ".join(
330 [
331 "{:.2f}% in {:s}".format(
332 100
333 * np.abs(stats_dict_mcmc[key]["median"] - inj[key])
334 / (
335 stats_dict_mcmc[key]["upper90"]
336 - stats_dict_mcmc[key]["lower90"]
337 ),
338 key,
339 )
340 for key in search_keys
341 ]
342 ),
343 )
344 )
345 print()
346
347 # do additional custom plotting
348 print("Loading grid and MCMC search results for custom comparison plots...")
349 gridfile = os.path.join(outdir, gridsearch.label + "_NA_GridSearch.txt")
350 if not os.path.isfile(gridfile):
351 raise RuntimeError(
352 "Failed to load GridSearch results from file '{:s}',"
353 " something must have gone wrong!".format(gridfile)
354 )
355 grid_res = pyfstat.helper_functions.read_txt_file_with_header(gridfile)
356 mcmc_file = os.path.join(outdir, mcmcsearch.label + "_samples.dat")
357 if not os.path.isfile(mcmc_file):
358 raise RuntimeError(
359 "Failed to load MCMCSearch results from file '{:s}',"
360 " something must have gone wrong!".format(mcmc_file)
361 )
362 mcmc_res = pyfstat.helper_functions.read_txt_file_with_header(mcmc_file)
363
364 zoom = {
365 "F0": [inj["F0"] - 10 * dF0, inj["F0"] + 10 * dF0],
366 "F1": [inj["F1"] - 5 * dF1, inj["F1"] + 5 * dF1],
367 }
368
369 # we'll use the two local plotting functions defined above
370 # to avoid code duplication in the sky case
371 print("Creating MCMC-grid comparison plots...")
372 plot_grid_vs_samples(grid_res, mcmc_res, "F0", "F1")
373 plot_2F_scatter(grid_res, "grid", "F0", "F1")
374 plot_2F_scatter(mcmc_res, "mcmc", "F0", "F1")
375 if sky:
376 plot_grid_vs_samples(grid_res, mcmc_res, "Alpha", "Delta")
377 plot_2F_scatter(grid_res, "grid", "Alpha", "Delta")
378 plot_2F_scatter(mcmc_res, "mcmc", "Alpha", "Delta")
Total running time of the script: ( 0 minutes 0.000 seconds)