Coverage for pesummary/tests/summaryplots_test.py: 100.0%

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

2 

3import os 

4import shutil 

5from glob import glob 

6 

7from pesummary.gw.cli.parser import ArgumentParser 

8from pesummary.gw.cli.inputs import PlottingInput, WebpagePlusPlottingPlusMetaFileInput 

9from pesummary.cli.summaryplots import _GWPlotGeneration as GWPlotGeneration 

10from pesummary.gw.file.meta_file import GWMetaFile 

11from pesummary.cli.summarypages import _GWWebpageGeneration as GWWebpageGeneration 

12from .base import make_result_file, get_list_of_plots, data_dir, testing_dir 

13 

14import pytest 

15import tempfile 

16from pathlib import Path 

17 

18__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

19 

20 

21class TestPlotGeneration(object): 

22 

23 def setup_method(self): 

24 tmpdir = Path(tempfile.TemporaryDirectory(prefix=".", dir=".").name).name 

25 os.makedirs(tmpdir) 

26 self.dir = tmpdir 

27 

28 def teardown_method(self): 

29 """Remove the files created from this class 

30 """ 

31 if os.path.isdir(self.dir): 

32 shutil.rmtree(self.dir) 

33 

34 def test_plot_generation_for_bilby_structure(self): 

35 with open(f"{self.dir}/psd.dat", "w") as f: 

36 f.writelines(["1.00 3.44\n"]) 

37 f.writelines(["100.00 4.00\n"]) 

38 f.writelines(["1000.00 5.00\n"]) 

39 f.writelines(["2000.00 6.00\n"]) 

40 with open(f"{self.dir}/calibration.dat", "w") as f: 

41 f.writelines(["1.0 2.0 3.0 4.0 5.0 6.0 7.0\n"]) 

42 f.writelines(["2000.0 2.0 3.0 4.0 5.0 6.0 7.0"]) 

43 parser = ArgumentParser() 

44 parser.add_all_known_options_to_parser() 

45 make_result_file( 

46 gw=True, extension="hdf5", bilby=True, outdir=self.dir, 

47 n_samples=10 

48 ) 

49 os.rename(f"{self.dir}/test.hdf5", f"{self.dir}/bilby_example.h5") 

50 default_arguments = [ 

51 "--approximant", "IMRPhenomPv2", 

52 "--webdir", self.dir, 

53 "--samples", f"{self.dir}/bilby_example.h5", 

54 "--config", data_dir + "/config_bilby.ini", 

55 "--psd", f"{self.dir}/psd.dat", 

56 "--calibration", f"{self.dir}/calibration.dat", 

57 "--labels", "H10", "--no_ligo_skymap", "--disable_expert", 

58 "--pastro_category_file", f"{testing_dir}/rates.yml", 

59 "--catch_terrestrial_probability_error" 

60 ] 

61 opts = parser.parse_args(default_arguments) 

62 inputs = PlottingInput(opts) 

63 webpage = GWPlotGeneration(inputs) 

64 webpage.generate_plots() 

65 plots = sorted(glob(f"{self.dir}/plots/*.png")) 

66 expected_plots = get_list_of_plots( 

67 gw=True, label="H1", outdir=self.dir, psd=True, 

68 calibration=False, waveform=True 

69 ) 

70 for i, j in zip(expected_plots, plots): 

71 print(i, j) 

72 assert all(i == j for i,j in zip(sorted(expected_plots), sorted(plots))) 

73 

74 def test_plot_generation_for_lalinference_structure(self): 

75 parser = ArgumentParser() 

76 parser.add_all_known_options_to_parser() 

77 make_result_file( 

78 gw=True, extension="hdf5", lalinference=True, 

79 outdir=self.dir, n_samples=10 

80 ) 

81 os.rename( 

82 f"{self.dir}/test.hdf5", f"{self.dir}/lalinference_example.h5" 

83 ) 

84 default_arguments = [ 

85 "--approximant", "IMRPhenomPv2", 

86 "--webdir", self.dir, 

87 "--samples", f"{self.dir}/lalinference_example.h5", 

88 "--config", data_dir + "/config_lalinference.ini", 

89 "--labels", "H10", "--no_ligo_skymap", "--disable_expert", 

90 "--pastro_category_file", f"{testing_dir}/rates.yml", 

91 "--catch_terrestrial_probability_error" 

92 ] 

93 opts = parser.parse_args(default_arguments) 

94 inputs = PlottingInput(opts) 

95 webpage = GWPlotGeneration(inputs) 

96 webpage.generate_plots() 

97 plots = sorted(glob(f"{self.dir}/plots/*.png")) 

98 expected_plots = get_list_of_plots( 

99 gw=True, label="H1", outdir=self.dir, waveform=True 

100 ) 

101 assert all(i == j for i,j in zip(sorted(expected_plots), sorted(plots))) 

102 

103 def test_plot_generation_for_comparison(self): 

104 parser = ArgumentParser() 

105 parser.add_all_known_options_to_parser() 

106 make_result_file( 

107 gw=True, extension="hdf5", lalinference=True, 

108 outdir=self.dir, n_samples=10 

109 ) 

110 os.rename( 

111 f"{self.dir}/test.hdf5", f"{self.dir}/lalinference_example.h5" 

112 ) 

113 make_result_file( 

114 gw=True, extension="hdf5", bilby=True, outdir=self.dir, 

115 n_samples=10 

116 ) 

117 os.rename( 

118 f"{self.dir}/test.hdf5", f"{self.dir}/bilby_example.h5" 

119 ) 

120 default_arguments = [ 

121 "--approximant", "IMRPhenomPv2", "IMRPhenomPv3", 

122 "--webdir", self.dir, 

123 "--samples", f"{self.dir}/bilby_example.h5", 

124 f"{self.dir}/lalinference_example.h5", 

125 "--labels", "H10", "H11", "--no_ligo_skymap", "--disable_expert", 

126 "--pastro_category_file", f"{testing_dir}/rates.yml", 

127 "--catch_terrestrial_probability_error" 

128 ] 

129 opts = parser.parse_args(default_arguments) 

130 inputs = PlottingInput(opts) 

131 webpage = GWPlotGeneration(inputs) 

132 webpage.generate_plots() 

133 plots = sorted(glob(f"{self.dir}/plots/*.png")) 

134 expected_plots = get_list_of_plots( 

135 gw=True, label="H1", number=2, outdir=self.dir, 

136 waveform=True 

137 ) 

138 for i,j in zip(sorted(plots), sorted(expected_plots)): 

139 print(i, j) 

140 assert all(i == j for i,j in zip(sorted(plots), sorted(expected_plots))) 

141 

142 def test_plot_generation_for_add_to_existing(self): 

143 parser = ArgumentParser() 

144 parser.add_all_known_options_to_parser() 

145 make_result_file( 

146 gw=True, extension="hdf5", lalinference=True, 

147 outdir=self.dir, n_samples=10 

148 ) 

149 os.rename( 

150 f"{self.dir}/test.hdf5", f"{self.dir}/lalinference_example.h5" 

151 ) 

152 make_result_file( 

153 gw=True, extension="hdf5", bilby=True, 

154 outdir=self.dir, n_samples=10 

155 ) 

156 os.rename( 

157 f"{self.dir}/test.hdf5", f"{self.dir}/bilby_example.h5" 

158 ) 

159 default_arguments = [ 

160 "--approximant", "IMRPhenomPv2", 

161 "--webdir", self.dir, 

162 "--samples", f"{self.dir}/bilby_example.h5", 

163 "--labels", "H10", "--no_ligo_skymap", "--disable_expert", 

164 "--pastro_category_file", f"{testing_dir}/rates.yml", 

165 "--catch_terrestrial_probability_error" 

166 ] 

167 opts = parser.parse_args(default_arguments) 

168 inputs = WebpagePlusPlottingPlusMetaFileInput(opts) 

169 webpage = GWPlotGeneration(inputs) 

170 webpage.generate_plots() 

171 webpage = GWWebpageGeneration(inputs) 

172 webpage.generate_webpages() 

173 meta_file = GWMetaFile(inputs) 

174 parser = ArgumentParser() 

175 parser.add_all_known_options_to_parser() 

176 default_arguments = [ 

177 "--approximant", "IMRPhenomPv3", 

178 "--existing_webdir", self.dir, 

179 "--samples", f"{self.dir}/lalinference_example.h5", 

180 "--labels", "H11", "--no_ligo_skymap", "--disable_expert", 

181 "--pastro_category_file", f"{testing_dir}/rates.yml", 

182 "--catch_terrestrial_probability_error" 

183 ] 

184 opts = parser.parse_args(default_arguments) 

185 inputs = PlottingInput(opts) 

186 webpage = GWPlotGeneration(inputs) 

187 webpage.generate_plots() 

188 plots = sorted(glob(f"{self.dir}/plots/*.png")) 

189 expected_plots = get_list_of_plots( 

190 gw=True, label="H1", number=2, outdir=self.dir, 

191 waveform=True 

192 ) 

193 assert all(i == j for i, j in zip(sorted(plots), sorted(expected_plots))) 

194 

195 def test_plot_generation_for_multiple_without_comparison(self): 

196 parser = ArgumentParser() 

197 parser.add_all_known_options_to_parser() 

198 make_result_file( 

199 gw=True, extension="hdf5", lalinference=True, 

200 outdir=self.dir, n_samples=10 

201 ) 

202 os.rename( 

203 f"{self.dir}/test.hdf5", f"{self.dir}/lalinference_example.h5" 

204 ) 

205 make_result_file( 

206 gw=True, extension="hdf5", bilby=True, 

207 outdir=self.dir, n_samples=10 

208 ) 

209 os.rename( 

210 f"{self.dir}/test.hdf5", f"{self.dir}/bilby_example.h5" 

211 ) 

212 default_arguments = [ 

213 "--approximant", "IMRPhenomPv2", "IMRPhenomPv3", 

214 "--webdir", self.dir, 

215 "--samples", f"{self.dir}/bilby_example.h5", 

216 f"{self.dir}/lalinference_example.h5", 

217 "--labels", "H10", "H11", "--no_ligo_skymap", 

218 "--disable_comparison", "--disable_expert", 

219 "--pastro_category_file", f"{testing_dir}/rates.yml", 

220 "--catch_terrestrial_probability_error" 

221 ] 

222 opts = parser.parse_args(default_arguments) 

223 inputs = PlottingInput(opts) 

224 webpage = GWPlotGeneration(inputs) 

225 webpage.generate_plots() 

226 plots = sorted(glob(f"{self.dir}/plots/*.png")) 

227 expected_plots = get_list_of_plots( 

228 gw=True, label="H1", number=2, outdir=self.dir, 

229 comparison=False, waveform=True 

230 ) 

231 assert all(i == j for i,j in zip(sorted(plots), sorted(expected_plots))) 

232 

233 def test_plot_generation_for_add_to_existing_without_comparison(self): 

234 parser = ArgumentParser() 

235 parser.add_all_known_options_to_parser() 

236 make_result_file( 

237 gw=True, extension="hdf5", lalinference=True, 

238 outdir=self.dir, n_samples=10 

239 ) 

240 os.rename( 

241 f"{self.dir}/test.hdf5", f"{self.dir}/lalinference_example.h5" 

242 ) 

243 make_result_file( 

244 gw=True, extension="hdf5", bilby=True, 

245 outdir=self.dir, n_samples=10 

246 ) 

247 os.rename( 

248 f"{self.dir}/test.hdf5", f"{self.dir}/bilby_example.h5" 

249 ) 

250 default_arguments = [ 

251 "--approximant", "IMRPhenomPv2", 

252 "--webdir", self.dir, 

253 "--samples", f"{self.dir}/bilby_example.h5", 

254 "--labels", "H10", "--no_ligo_skymap", "--disable_expert", 

255 "--pastro_category_file", f"{testing_dir}/rates.yml", 

256 "--catch_terrestrial_probability_error" 

257 ] 

258 opts = parser.parse_args(default_arguments) 

259 inputs = WebpagePlusPlottingPlusMetaFileInput(opts) 

260 webpage = GWPlotGeneration(inputs) 

261 webpage.generate_plots() 

262 webpage = GWWebpageGeneration(inputs) 

263 webpage.generate_webpages() 

264 meta_file = GWMetaFile(inputs) 

265 parser = ArgumentParser() 

266 parser.add_all_known_options_to_parser() 

267 default_arguments = [ 

268 "--approximant", "IMRPhenomPv3", 

269 "--existing_webdir", self.dir, 

270 "--samples", f"{self.dir}/lalinference_example.h5", 

271 "--labels", "H11", "--no_ligo_skymap", 

272 "--disable_comparison", "--disable_expert", 

273 "--pastro_category_file", f"{testing_dir}/rates.yml", 

274 "--catch_terrestrial_probability_error" 

275 ] 

276 opts = parser.parse_args(default_arguments) 

277 inputs = PlottingInput(opts) 

278 webpage = GWPlotGeneration(inputs) 

279 webpage.generate_plots() 

280 plots = sorted(glob(f"{self.dir}/plots/*.png")) 

281 expected_plots = get_list_of_plots( 

282 gw=True, label="H1", number=2, outdir=self.dir, 

283 comparison=False, waveform=True 

284 ) 

285 assert all(i == j for i, j in zip(sorted(plots), sorted(expected_plots)))