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 

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.h5", 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 opts = parser.parse_args(default_arguments) 

59 inputs = PlottingInput(opts) 

60 webpage = GWPlotGeneration(inputs) 

61 webpage.generate_plots() 

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

63 expected_plots = get_list_of_plots( 

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

65 calibration=False 

66 ) 

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

68 print(i, j) 

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

70 

71 def test_plot_generation_for_lalinference_structure(self): 

72 parser = ArgumentParser() 

73 parser.add_all_known_options_to_parser() 

74 make_result_file( 

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

76 outdir=self.dir, n_samples=10 

77 ) 

78 os.rename( 

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

80 ) 

81 default_arguments = [ 

82 "--approximant", "IMRPhenomPv2", 

83 "--webdir", self.dir, 

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

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

86 "--labels", "H10", "--no_ligo_skymap", "--disable_expert"] 

87 opts = parser.parse_args(default_arguments) 

88 inputs = PlottingInput(opts) 

89 webpage = GWPlotGeneration(inputs) 

90 webpage.generate_plots() 

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

92 expected_plots = get_list_of_plots( 

93 gw=True, label="H1", outdir=self.dir, 

94 ) 

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

96 

97 def test_plot_generation_for_comparison(self): 

98 parser = ArgumentParser() 

99 parser.add_all_known_options_to_parser() 

100 make_result_file( 

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

102 outdir=self.dir, n_samples=10 

103 ) 

104 os.rename( 

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

106 ) 

107 make_result_file( 

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

109 n_samples=10 

110 ) 

111 os.rename( 

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

113 ) 

114 default_arguments = [ 

115 "--approximant", "IMRPhenomPv2", "IMRPhenomP", 

116 "--webdir", self.dir, 

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

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

119 "--labels", "H10", "H11", "--no_ligo_skymap", "--disable_expert"] 

120 opts = parser.parse_args(default_arguments) 

121 inputs = PlottingInput(opts) 

122 webpage = GWPlotGeneration(inputs) 

123 webpage.generate_plots() 

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

125 expected_plots = get_list_of_plots( 

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

127 ) 

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

129 print(i, j) 

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

131 

132 def test_plot_generation_for_add_to_existing(self): 

133 parser = ArgumentParser() 

134 parser.add_all_known_options_to_parser() 

135 make_result_file( 

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

137 outdir=self.dir, n_samples=10 

138 ) 

139 os.rename( 

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

141 ) 

142 make_result_file( 

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

144 outdir=self.dir, n_samples=10 

145 ) 

146 os.rename( 

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

148 ) 

149 default_arguments = [ 

150 "--approximant", "IMRPhenomPv2", 

151 "--webdir", self.dir, 

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

153 "--labels", "H10", "--no_ligo_skymap", "--disable_expert"] 

154 opts = parser.parse_args(default_arguments) 

155 inputs = WebpagePlusPlottingPlusMetaFileInput(opts) 

156 webpage = GWPlotGeneration(inputs) 

157 webpage.generate_plots() 

158 webpage = GWWebpageGeneration(inputs) 

159 webpage.generate_webpages() 

160 meta_file = GWMetaFile(inputs) 

161 parser = ArgumentParser() 

162 parser.add_all_known_options_to_parser() 

163 default_arguments = [ 

164 "--approximant", "IMRPhenomP", 

165 "--existing_webdir", self.dir, 

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

167 "--labels", "H11", "--no_ligo_skymap", "--disable_expert"] 

168 opts = parser.parse_args(default_arguments) 

169 inputs = PlottingInput(opts) 

170 webpage = GWPlotGeneration(inputs) 

171 webpage.generate_plots() 

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

173 expected_plots = get_list_of_plots( 

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

175 ) 

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

177 

178 def test_plot_generation_for_multiple_without_comparison(self): 

179 parser = ArgumentParser() 

180 parser.add_all_known_options_to_parser() 

181 make_result_file( 

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

183 outdir=self.dir, n_samples=10 

184 ) 

185 os.rename( 

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

187 ) 

188 make_result_file( 

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

190 outdir=self.dir, n_samples=10 

191 ) 

192 os.rename( 

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

194 ) 

195 default_arguments = [ 

196 "--approximant", "IMRPhenomPv2", "IMRPhenomP", 

197 "--webdir", self.dir, 

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

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

200 "--labels", "H10", "H11", "--no_ligo_skymap", 

201 "--disable_comparison", "--disable_expert" 

202 ] 

203 opts = parser.parse_args(default_arguments) 

204 inputs = PlottingInput(opts) 

205 webpage = GWPlotGeneration(inputs) 

206 webpage.generate_plots() 

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

208 expected_plots = get_list_of_plots( 

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

210 comparison=False 

211 ) 

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

213 print(i, j) 

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

215 

216 def test_plot_generation_for_add_to_existing_without_comparison(self): 

217 parser = ArgumentParser() 

218 parser.add_all_known_options_to_parser() 

219 make_result_file( 

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

221 outdir=self.dir, n_samples=10 

222 ) 

223 os.rename( 

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

225 ) 

226 make_result_file( 

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

228 outdir=self.dir, n_samples=10 

229 ) 

230 os.rename( 

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

232 ) 

233 default_arguments = [ 

234 "--approximant", "IMRPhenomPv2", 

235 "--webdir", self.dir, 

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

237 "--labels", "H10", "--no_ligo_skymap", "--disable_expert"] 

238 opts = parser.parse_args(default_arguments) 

239 inputs = WebpagePlusPlottingPlusMetaFileInput(opts) 

240 webpage = GWPlotGeneration(inputs) 

241 webpage.generate_plots() 

242 webpage = GWWebpageGeneration(inputs) 

243 webpage.generate_webpages() 

244 meta_file = GWMetaFile(inputs) 

245 parser = ArgumentParser() 

246 parser.add_all_known_options_to_parser() 

247 default_arguments = [ 

248 "--approximant", "IMRPhenomP", 

249 "--existing_webdir", self.dir, 

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

251 "--labels", "H11", "--no_ligo_skymap", 

252 "--disable_comparison", "--disable_expert" 

253 ] 

254 opts = parser.parse_args(default_arguments) 

255 inputs = PlottingInput(opts) 

256 webpage = GWPlotGeneration(inputs) 

257 webpage.generate_plots() 

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

259 expected_plots = get_list_of_plots( 

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

261 comparison=False 

262 ) 

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