Coverage for pesummary/gw/plots/tgr.py: 93.0%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3import os 

4 

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) 

13 

14 

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. 

32 

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 

70 

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" 

105 

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 ] 

116 

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 

124 

125 

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. 

132 

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"]