Coverage for pesummary/tests/summaryplots_test.py: 64.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.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 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, waveform=True 

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, waveform=True 

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.hdf5", f"{self.dir}/bilby_example.h5" 

113 ) 

114 default_arguments = [ 

115 "--approximant", "IMRPhenomPv2", "IMRPhenomPv3", 

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 waveform=True 

128 ) 

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

130 print(i, j) 

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

132 

133 def test_plot_generation_for_add_to_existing(self): 

134 parser = ArgumentParser() 

135 parser.add_all_known_options_to_parser() 

136 make_result_file( 

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

138 outdir=self.dir, n_samples=10 

139 ) 

140 os.rename( 

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

142 ) 

143 make_result_file( 

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

145 outdir=self.dir, n_samples=10 

146 ) 

147 os.rename( 

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

149 ) 

150 default_arguments = [ 

151 "--approximant", "IMRPhenomPv2", 

152 "--webdir", self.dir, 

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

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

155 opts = parser.parse_args(default_arguments) 

156 inputs = WebpagePlusPlottingPlusMetaFileInput(opts) 

157 webpage = GWPlotGeneration(inputs) 

158 webpage.generate_plots() 

159 webpage = GWWebpageGeneration(inputs) 

160 webpage.generate_webpages() 

161 meta_file = GWMetaFile(inputs) 

162 parser = ArgumentParser() 

163 parser.add_all_known_options_to_parser() 

164 default_arguments = [ 

165 "--approximant", "IMRPhenomPv3", 

166 "--existing_webdir", self.dir, 

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

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

169 opts = parser.parse_args(default_arguments) 

170 inputs = PlottingInput(opts) 

171 webpage = GWPlotGeneration(inputs) 

172 webpage.generate_plots() 

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

174 expected_plots = get_list_of_plots( 

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

176 waveform=True 

177 ) 

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

179 

180 def test_plot_generation_for_multiple_without_comparison(self): 

181 parser = ArgumentParser() 

182 parser.add_all_known_options_to_parser() 

183 make_result_file( 

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

185 outdir=self.dir, n_samples=10 

186 ) 

187 os.rename( 

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

189 ) 

190 make_result_file( 

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

192 outdir=self.dir, n_samples=10 

193 ) 

194 os.rename( 

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

196 ) 

197 default_arguments = [ 

198 "--approximant", "IMRPhenomPv2", "IMRPhenomPv3", 

199 "--webdir", self.dir, 

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

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

202 "--labels", "H10", "H11", "--no_ligo_skymap", 

203 "--disable_comparison", "--disable_expert" 

204 ] 

205 opts = parser.parse_args(default_arguments) 

206 inputs = PlottingInput(opts) 

207 webpage = GWPlotGeneration(inputs) 

208 webpage.generate_plots() 

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

210 expected_plots = get_list_of_plots( 

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

212 comparison=False, waveform=True 

213 ) 

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.hdf5", 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", "IMRPhenomPv3", 

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, waveform=True 

262 ) 

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