Coverage for pesummary/gw/plots/tgr.py: 93.0%
43 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3import os
5__author__ = [
6 "Charlie Hoy <charlie.hoy@ligo.org>",
7 "Aditya Vijaykumar <aditya.vijaykumar@ligo.org>",
8]
9DEFAULT_PLOT_KWARGS = dict(
10 grid=True, smooth=4, type="triangle", fontsize=dict(label=20, legend=14),
11 fig_kwargs=dict(wspace=0.2, hspace=0.2),
12)
15def imrct_plot(
16 imrct_deviations,
17 samples=None,
18 evolve_spins=False,
19 make_diagnostic_plots=False,
20 inspiral_string="inspiral",
21 postinspiral_string="postinspiral",
22 cmap="YlOrBr",
23 levels=[0.68, 0.95],
24 level_kwargs={"colors": ["k", "k"]},
25 xlabel=r"$\Delta M_{\mathrm{f}} / \bar{M}_{\mathrm{f}}$",
26 ylabel=r"$\Delta a_{\mathrm{f}} / \bar{a}_{\mathrm{f}}$",
27 _default_plot_kwargs=DEFAULT_PLOT_KWARGS,
28 **plot_kwargs
29):
30 """Generate a triangle plot showing the IMR deviation parameters and
31 diagnostic plots if specified.
33 Parameters
34 ----------
35 imrct_deviations: ProbabilityDict2D
36 Output of imrct_deviation_parameters_from_final_mass_final_spin
37 samples: MultiAnalysisSamplesDict, optional
38 Dictionary containing inspiral and postinspiral samples. Default None
39 evolve_spins: Bool, optional
40 if True, use the evolved spin remnant properties for the diagnostic
41 plots. Default False
42 make_diagnostic_plots: Bool, optional
43 if True, generate 3 diagnostic triangle plots, one showing the
44 a_1-a_2 posterior distribution, another showing the
45 mass_1-mass_2 posterior distribution and another showing the
46 remnant properties. Default False
47 inspiral_string: str, optional
48 The label of the inspiral analysis in the MultiAnalysisSamplesDict.
49 Default 'inspiral'
50 postinspiral_string: str, optional
51 The label of the postinspiral analysis in the MultiAnalysisSamplesDict.
52 Default 'postinspiral'
53 cmap: str, optional
54 The cmap to use when generating the IMRCT deviation triangle plot.
55 Default 'YlOrBr'
56 levels: tuple, optional
57 The levels to plot in the IMRCT deviation triangle plot. Default
58 [0.68, 0.95]
59 level_kwargs: dict, optional
60 Level kwargs to use in the IMRCT deviation triangle plot. Default
61 {'colors': ['k', 'k']}
62 xlabel: str, optional
63 the xlabel to use in the IMRCT deviation triangle plot. Default
64 r'$\Delta M_{\mathrm{f}} / \bar{M}_{\mathrm{f}}$'
65 ylabel: str, optional
66 the ylabel to use in the IMRCT deviation triangle plot. Default
67 r'$\Delta a_{\mathrm{f}} / \bar{a}_{\mathrm{f}}$'
68 plot_kwargs: dict, optional
69 all additional kwargs passed to the IMRCT deviation triangle plot
71 Returns
72 -------
73 figs: dict
74 dictionary of figures. The IMRCT deviation plot has key
75 'imrct_deviation' and diagnostic plots have keys
76 'diagnostic_{}_{}'
77 """
78 figs = {}
79 _plot_kwargs = _default_plot_kwargs.copy()
80 _plot_kwargs.update(
81 {
82 "cmap": cmap,
83 "levels": levels,
84 "level_kwargs": level_kwargs,
85 "xlabel": xlabel,
86 "ylabel": ylabel,
87 }
88 )
89 _plot_kwargs.update(plot_kwargs)
90 fig, _ax1, ax_2d, _ax3 = imrct_deviations.plot(
91 "final_mass_final_spin_deviations",
92 **_plot_kwargs,
93 )
94 ax_2d.plot(0, 0, "k+", ms=12, mew=2)
95 figs["imrct_deviation"] = [fig, _ax1, ax_2d, _ax3]
96 if make_diagnostic_plots and samples is None:
97 raise ValueError(
98 "Please provide a MultiAnalysisSamplesDict object containing the "
99 "posterior samples for the inspiral and postinspiral analysis"
100 )
101 elif make_diagnostic_plots:
102 evolve_spins_string = ""
103 if not evolve_spins:
104 evolve_spins_string = "_non_evolved"
106 samples_string = "final_{}" + evolve_spins_string
107 plot_kwargs = _default_plot_kwargs.copy()
108 plot_kwargs.update(
109 {"fill_alpha": 0.2, "labels": [inspiral_string, postinspiral_string]}
110 )
111 parameters_to_plot = [
112 [samples_string.format("mass"), samples_string.format("spin")],
113 ["mass_1", "mass_2"],
114 ["a_1", "a_2"],
115 ]
117 for parameters in parameters_to_plot:
118 fig, ax1, ax2, ax3 = samples.plot(
119 parameters,
120 **plot_kwargs,
121 )
122 figs["diagnostic_{}_{}".format(*parameters)] = [fig, ax1, ax2, ax3]
123 return figs
126def make_and_save_imrct_plots(
127 *args, webdir="./", plot_label=None, return_fig=False,
128 save=True, **kwargs
129):
130 """Generate and save a triangle plot showing the IMR deviation parameters
131 and diagnostic plots if specified.
133 Parameters
134 ----------
135 *args: tuple
136 all args passed to the imrct_plot function
137 webdir: str, optional
138 the directory to save the plots. Default "./"
139 plot_label: str, optional
140 label to prepend the
141 return_fig: Bool, optional
142 if True, return the figure and axes associated with the imrct_deviation
143 plot
144 save: Bool, optional
145 if True, save the figures
146 **kwargs: dict, optional
147 all additional kwargs passed to the imrct_plot function
148 """
149 plotdir = os.path.join(webdir, "plots")
150 if plot_label is not None:
151 base_string = os.path.join(plotdir, "%s_imrct_{}.png" % (plot_label))
152 else:
153 base_string = os.path.join(plotdir, "imrct_{}.png")
154 figs = imrct_plot(*args, **kwargs)
155 if save:
156 fig = figs["imrct_deviation"][0]
157 fig.savefig(
158 base_string.format("deviations_triangle_plot"), bbox_inches="tight"
159 )
160 fig.close()
161 diagnostic_keys = [key for key in figs.keys() if "diagnostic" in key]
162 for diag in diagnostic_keys:
163 fig = figs[diag][0]
164 save_string = "_".join(diag.split("_")[1:])
165 fig.savefig(base_string.format(save_string))
166 fig.close()
167 if return_fig:
168 return figs["imrct_deviation"]