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