Coverage for pesummary/tests/base.py: 85.4%

302 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 

2 

3import os 

4import numpy as np 

5from pathlib import Path 

6from pesummary.gw.cli.inputs import WebpagePlusPlottingPlusMetaFileInput 

7from pesummary.core.cli.inputs import ( 

8 WebpagePlusPlottingPlusMetaFileInput as Input 

9) 

10 

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

12 

13 

14class Namespace(object): 

15 def __init__(self, **kwargs): 

16 self.__dict__.update(kwargs) 

17 

18 

19def namespace(args): 

20 """Generate a namespace for testing purposes 

21 

22 Parameters 

23 ---------- 

24 args: dict 

25 dictionary of arguments 

26 """ 

27 base = Namespace() 

28 for key in list(args.keys()): 

29 setattr(base, key, args[key]) 

30 return base 

31 

32 

33def gw_parameters(): 

34 parameters = [ 

35 'mass_1', 'mass_2', 'a_1', 'a_2', 'tilt_1', 'tilt_2', 'phi_jl', 

36 'phi_12', 'psi', 'theta_jn', 'ra', 'dec', 'luminosity_distance', 

37 'geocent_time', 'log_likelihood', 'mass_ratio', 'total_mass', 

38 'chirp_mass', 'symmetric_mass_ratio', 'iota', 'spin_1x', 'spin_1y', 

39 'spin_1z', 'spin_2x', 'spin_2y', 'spin_2z', 'chi_p', 'chi_eff', 

40 'cos_tilt_1', 'cos_tilt_2', 'redshift', 'comoving_distance', 

41 'mass_1_source', 'mass_2_source', 'total_mass_source', 

42 'chirp_mass_source', 'phi_1', 'phi_2', 'cos_theta_jn', 'cos_iota', 

43 'peak_luminosity_non_evolved', 'final_spin_non_evolved', 

44 'final_mass_non_evolved', 'final_mass_source_non_evolved', 

45 'radiated_energy_non_evolved', 'inverted_mass_ratio', 

46 'viewing_angle', 'chi_p_2spin', 'comoving_volume' 

47 ] 

48 return parameters 

49 

50 

51def get_list_of_files( 

52 gw=False, number=1, existing_plot=False, parameters=[], sections=[], 

53 outdir=".outdir", extra_gw_pages=True 

54): 

55 """Return a list of files that should be generated from a typical workflow 

56 """ 

57 from pesummary.conf import gw_2d_plots 

58 if not gw and not len(parameters): 

59 import string 

60 parameters = list(string.ascii_lowercase)[:17] + ["log_likelihood"] 

61 elif not len(parameters): 

62 parameters = gw_parameters() 

63 if not gw: 

64 label = "core" 

65 else: 

66 label = "gw" 

67 html = [ 

68 "%s/html/error.html" % (outdir), 

69 "%s/html/Version.html" % (outdir), 

70 "%s/html/Logging.html" % (outdir), 

71 "%s/html/About.html" % (outdir), 

72 "%s/html/Downloads.html" % (outdir)] 

73 if gw and not len(sections): 

74 sections = [ 

75 "spins", "spin_angles", "timings", "source", "remnant", "others", 

76 "masses", "location", "inclination", "energy" 

77 ] 

78 elif not len(sections): 

79 sections = ["A-D", "E-F", "I-L", "M-P", "Q-T"] 

80 for num in range(number): 

81 html.append("%s/html/%s%s_%s%s.html" % (outdir, label, num, label, num)) 

82 if gw and extra_gw_pages: 

83 html.append("%s/html/%s%s_%s%s_Classification.html" % (outdir, label, num, label, num)) 

84 html.append("%s/html/%s%s_%s%s_Corner.html" % (outdir, label, num, label, num)) 

85 html.append("%s/html/%s%s_%s%s_Config.html" % (outdir, label, num, label, num)) 

86 html.append("%s/html/%s%s_%s%s_Custom.html" % (outdir, label, num, label, num)) 

87 html.append("%s/html/%s%s_%s%s_All.html" % (outdir, label, num, label, num)) 

88 html.append("%s/html/%s%s_%s%s_Interactive_Corner.html" % ( 

89 outdir, label, num, label, num 

90 )) 

91 for section in sections: 

92 html.append("%s/html/%s%s_%s%s_%s_all.html" % (outdir, label, num, label, num, section)) 

93 for j in parameters: 

94 html.append("%s/html/%s%s_%s%s_%s.html" % (outdir, label, num, label, num, j)) 

95 if gw: 

96 # 2d histogram pages 

97 for section in ["location", "masses", "spins", "timings"]: 

98 if section == "location": 

99 pairs = [p for p in gw_2d_plots if "luminosity_distance" in p] 

100 if not any(all(p in parameters for p in _) for _ in pairs): 

101 continue 

102 elif section == "masses": 

103 pairs = [p for p in gw_2d_plots if "mass_1" in p] 

104 if not any(all(p in parameters for p in _) for _ in pairs): 

105 continue 

106 elif section == "spins": 

107 pairs = [p for p in gw_2d_plots if "chi_p" in p] 

108 if not any(all(p in parameters for p in _) for _ in pairs): 

109 continue 

110 elif section == "timings": 

111 pairs = [p for p in gw_2d_plots if "geocent_time" in p] 

112 if not any(all(p in parameters for p in _) for _ in pairs): 

113 continue 

114 html.append("%s/html/%s%s_%s%s_%s.html" % (outdir, label, num, label, num, section)) 

115 if existing_plot: 

116 html.append("%s/html/%s%s_%s%s_Additional.html" % (outdir, label, num, label, num)) 

117 

118 if number > 1: 

119 html.append("%s/html/Comparison.html" % (outdir)) 

120 html.append("%s/html/Comparison_Custom.html" % (outdir)) 

121 html.append("%s/html/Comparison_All.html" % (outdir)) 

122 html.append("%s/html/Comparison_Interactive_Ridgeline.html" % (outdir)) 

123 for j in parameters: 

124 if j != "classification": 

125 html.append("%s/html/Comparison_%s.html" % (outdir, j)) 

126 for section in sections: 

127 html.append("%s/html/Comparison_%s_all.html" % (outdir, section)) 

128 return sorted(html) 

129 

130 

131def get_list_of_plots( 

132 gw=False, number=1, mcmc=False, label=None, outdir=".outdir", 

133 comparison=True, psd=False, calibration=False, existing_plot=False, 

134 expert=False, waveform=False, parameters=[], extra_gw_plots=True 

135): 

136 """Return a list of plots that should be generated from a typical workflow 

137 """ 

138 from pesummary.conf import gw_2d_plots 

139 if not gw and not len(parameters): 

140 import string 

141 

142 parameters = list(string.ascii_lowercase)[:17] + ["log_likelihood"] 

143 elif not len(parameters): 

144 parameters = gw_parameters() 

145 if not gw and label is None: 

146 label = "core" 

147 elif label is None: 

148 label = "gw" 

149 

150 plots = [] 

151 for num in range(number): 

152 for i in ["sample_evolution", "autocorrelation", "1d_posterior", "cdf"]: 

153 for j in parameters: 

154 plots.append("%s/plots/%s%s_%s_%s.png" % (outdir, label, num, i, j)) 

155 if mcmc: 

156 for j in parameters: 

157 plots.append("%s/plots/%s%s_1d_posterior_%s_combined.png" % (outdir, label, num, j)) 

158 if gw: 

159 for pair in gw_2d_plots: 

160 if not all(p in parameters for p in pair): 

161 continue 

162 plots.append("%s/plots/%s%s_2d_posterior_%s_%s.png" % ( 

163 outdir, label, num, pair[0], pair[1] 

164 )) 

165 if psd: 

166 plots.append("%s/plots/%s%s_psd_plot.png" % (outdir, label, num)) 

167 if calibration: 

168 plots.append("%s/plots/%s%s_calibration_plot.png" % (outdir, label, num)) 

169 if waveform: 

170 plots.append("%s/plots/%s%s_waveform.png" % (outdir, label, num)) 

171 plots.append("%s/plots/%s%s_waveform_time_domain.png" % (outdir, label, num)) 

172 if existing_plot: 

173 plots.append("%s/plots/test.png" % (outdir)) 

174 if expert: 

175 for j in parameters: 

176 plots.append("%s/plots/%s%s_2d_contour_%s_log_likelihood.png" % (outdir, label, num, j)) 

177 plots.append("%s/plots/%s%s_1d_posterior_%s_bootstrap.png" % (outdir, label, num, j)) 

178 plots.append("%s/plots/%s%s_sample_evolution_%s_log_likelihood_colored.png" % (outdir, label, num, j)) 

179 if number > 1 and comparison: 

180 for i in ["1d_posterior", "boxplot", "cdf"]: 

181 for j in parameters: 

182 plots.append("%s/plots/combined_%s_%s.png" % (outdir, i, j)) 

183 if gw and extra_gw_plots: 

184 for num in range(number): 

185 plots.append("%s/plots/%s%s_skymap.png" % (outdir, label, num)) 

186 plots.append("%s/plots/%s%s_default_pepredicates.png" % (outdir, label, num)) 

187 plots.append("%s/plots/%s%s_default_pepredicates_bar.png" % (outdir, label, num)) 

188 plots.append("%s/plots/%s%s_population_pepredicates.png" % (outdir, label, num)) 

189 plots.append("%s/plots/%s%s_population_pepredicates_bar.png" % (outdir, label, num)) 

190 if number > 1 and comparison: 

191 plots.append("%s/plots/combined_skymap.png" % (outdir)) 

192 

193 return sorted(plots) 

194 

195 

196def make_argparse(gw=True, extension="json", bilby=False, lalinference=False, 

197 number=1, existing=False, disable_expert=True, outdir="./.outdir"): 

198 """ 

199 """ 

200 default_args = [] 

201 if gw: 

202 from pesummary.gw.cli.parser import ArgumentParser 

203 default_args.append("--gw") 

204 default_args.append("--nsamples_for_skymap") 

205 default_args.append("10") 

206 else: 

207 from pesummary.core.cli.parser import ArgumentParser 

208 parser = ArgumentParser() 

209 parser.add_all_known_options_to_parser() 

210 params, data = make_result_file( 

211 extension=extension, gw=gw, bilby=bilby, lalinference=lalinference, 

212 outdir=outdir 

213 ) 

214 if not existing: 

215 default_args.append("--webdir") 

216 else: 

217 default_args.append("--existing_webdir") 

218 default_args.append(outdir) 

219 default_args.append("--samples") 

220 for i in range(number): 

221 default_args.append("%s/test.%s" % (outdir, extension)) 

222 default_args.append("--labels") 

223 if not existing: 

224 for i in range(number): 

225 if not gw: 

226 default_args.append("core%s" % (i)) 

227 else: 

228 default_args.append("gw%s" % (i)) 

229 else: 

230 if not gw: 

231 default_args.append("core1") 

232 else: 

233 default_args.append("gw1") 

234 default_args.append("--config") 

235 for i in range(number): 

236 default_args.append(testing_dir + "/example_config.ini") 

237 if disable_expert: 

238 default_args.append("--disable_expert") 

239 opts = parser.parse_args(default_args) 

240 if gw: 

241 func = WebpagePlusPlottingPlusMetaFileInput 

242 else: 

243 func = Input 

244 return opts, func(opts) 

245 

246 

247def read_result_file(outdir="./.outdir", extension="json", bilby=False, 

248 lalinference=False, pesummary=False): 

249 """ 

250 """ 

251 if bilby: 

252 from bilby.core.result import read_in_result 

253 

254 if extension == "json": 

255 f = read_in_result(outdir + "/test.json") 

256 elif extension == "h5" or extension == "hdf5": 

257 f = read_in_result(outdir + "/test.hdf5") 

258 posterior = f.posterior 

259 posterior = posterior.select_dtypes(include=[float, int]) 

260 samples = {key: val for key, val in posterior.items()} 

261 elif lalinference: 

262 import h5py 

263 

264 f = h5py.File(outdir + "/test.hdf5", "r") 

265 

266 posterior = f["lalinference"]["lalinference_nest"]["posterior_samples"] 

267 samples = { 

268 i: [j[num] for j in posterior] for num, i in enumerate( 

269 posterior.dtype.names 

270 ) 

271 } 

272 elif pesummary: 

273 from pesummary.gw.file.read import read 

274 

275 if extension == "json": 

276 f = read(outdir + "/test.json") 

277 else: 

278 f = read(outdir + "/test.h5") 

279 data = f.samples_dict 

280 labels = f.labels 

281 samples = data[labels[0]] 

282 elif extension == "dat": 

283 import numpy as np 

284 

285 data = np.genfromtxt(outdir + "/test.dat", names=True) 

286 samples = { 

287 i: [j[num] for j in data] for num, i in enumerate(data.dtype.names) 

288 } 

289 else: 

290 samples = {} 

291 return samples 

292 

293 

294def make_psd(outdir="./.outdir"): 

295 """Make a psd file 

296 """ 

297 frequencies = np.linspace(0, 1024, 1000) 

298 strains = np.random.uniform(10, 0.1, 1000) 

299 data = np.vstack([frequencies, strains]).T 

300 np.savetxt( 

301 "{}/psd.dat".format(outdir), data, delimiter="\t", 

302 header="\t".join(["frequencies", "strain"]) 

303 ) 

304 

305 

306def make_calibration(outdir="./.outdir"): 

307 """Make a calibration file 

308 """ 

309 frequencies = np.linspace(0, 1024, 1000) 

310 columns = [np.random.uniform(10, 0.1, 1000) for _ in range(6)] 

311 data = np.vstack([frequencies] + columns).T 

312 np.savetxt( 

313 "{}/calibration.dat".format(outdir), data, delimiter="\t", 

314 header="\t".join(["frequencies", "a", "b", "c", "d", "e", "f"]) 

315 ) 

316 

317 

318def make_injection_file( 

319 outdir="./.outdir", extension="json", return_filename=True, 

320 return_injection_dict=True 

321): 

322 """ 

323 """ 

324 from pesummary.io import write 

325 

326 filename = os.path.join(outdir, "injection.{}".format(extension)) 

327 parameters = gw_parameters() 

328 samples = np.array([[np.random.random()] for i in range(len(parameters))]).T 

329 write(parameters, samples, filename=filename, file_format=extension) 

330 args = [] 

331 if return_filename: 

332 args.append(filename) 

333 if return_injection_dict: 

334 args.append({param: samples[0][num] for num, param in enumerate(parameters)}) 

335 return args 

336 

337 

338def make_result_file(outdir="./.outdir/", extension="json", gw=True, bilby=False, 

339 lalinference=False, pesummary=False, pesummary_label="label", 

340 config=None, psd=None, calibration=None, random_seed=None, 

341 n_samples=1000): 

342 """Make a result file that can be read in by PESummary 

343 

344 Parameters 

345 ---------- 

346 outdir: str 

347 directory where you would like to store the result file 

348 extension: str 

349 the file extension of the result file 

350 gw: Bool 

351 if True, gw parameters will be used 

352 """ 

353 if random_seed is not None: 

354 np.random.seed(random_seed) 

355 print(extension, gw, bilby, lalinference, pesummary) 

356 if outdir[-1] != "/": 

357 outdir += "/" 

358 data = np.array([np.random.random(18) for i in range(n_samples)]) 

359 if gw: 

360 parameters = ["mass_1", "mass_2", "a_1", "a_2", "tilt_1", "tilt_2", 

361 "phi_jl", "phi_12", "psi", "theta_jn", "ra", "dec", 

362 "luminosity_distance", "geocent_time", "redshift", 

363 "mass_1_source", "mass_2_source", "log_likelihood"] 

364 distance = np.random.random(n_samples) * 500 

365 mass_1 = np.random.random(n_samples) * 100 

366 q = np.random.random(n_samples) * 100 

367 a_1 = np.random.uniform(0, 0.99, n_samples) 

368 a_2 = np.random.uniform(0, 0.99, n_samples) 

369 for num, i in enumerate(data): 

370 data[num][12] = distance[num] 

371 data[num][0] = mass_1[num] 

372 data[num][1] = mass_1[num] * q[num] 

373 data[num][2] = a_1[num] 

374 data[num][3] = a_2[num] 

375 else: 

376 import string 

377 

378 parameters = list(string.ascii_lowercase)[:17] + ["log_likelihood"] 

379 if extension == "dat": 

380 np.savetxt(outdir + "test.dat", data, delimiter=" ", 

381 header=" ".join(parameters), comments="") 

382 elif extension == "csv": 

383 np.savetxt(outdir + "test.csv", data, delimiter=",", 

384 header=",".join(parameters), comments="") 

385 elif extension == "npy": 

386 from pesummary.utils.samples_dict import SamplesDict 

387 samples = SamplesDict(parameters, np.array(data).T).to_structured_array() 

388 np.save(outdir + "test.npy", samples) 

389 elif extension == "json" and not bilby and not pesummary and not lalinference: 

390 import json 

391 

392 dictionary = {"NameOfCode": {"posterior_samples": {key: 

393 [i[num] for i in data] for num, key in 

394 enumerate(parameters)}}} 

395 with open(outdir + "test.json", "w") as f: 

396 json.dump(dictionary, f, indent=4, sort_keys=True) 

397 elif (extension == "hdf5" or extension == "h5") and not bilby and not pesummary and not lalinference: 

398 import h5py 

399 

400 h5py_data = np.array( 

401 [tuple(i) for i in data], dtype=[tuple([i, 'float64']) for i in parameters]) 

402 f = h5py.File(outdir + "test.h5", "w") 

403 name = f.create_group("NameOfCode") 

404 samples = f.create_dataset("posterior_samples", data=h5py_data) 

405 f.close() 

406 elif bilby and not pesummary and not lalinference: 

407 import bilby 

408 from bilby.core.result import Result 

409 from bilby.core.prior import PriorDict 

410 from pandas import DataFrame 

411 

412 priors = PriorDict() 

413 priors.update({"%s" % (i): bilby.core.prior.Uniform(0.1, 0.5, 0) for i in parameters}) 

414 posterior_data_frame = DataFrame(data, columns=parameters) 

415 injection_parameters = {par: 1. for par in parameters} 

416 bilby_object = Result( 

417 search_parameter_keys=parameters, samples=data, 

418 posterior=posterior_data_frame, label="test", 

419 injection_parameters=injection_parameters, 

420 priors=priors, 

421 log_bayes_factor=0.5, log_evidence_err=0.1, log_noise_evidence=0.1, 

422 log_evidence=0.2, version=["bilby=0.5.3:"], 

423 meta_data={"likelihood": {"time_marginalization": "True"}}) 

424 if extension == "json": 

425 bilby_object.save_to_file( 

426 filename=outdir + "test.json", extension="json") 

427 elif extension == "hdf5" or extension == "h5": 

428 bilby_object.save_to_file( 

429 filename=outdir + "test.hdf5", extension="hdf5") 

430 elif lalinference and not bilby and not pesummary: 

431 import h5py 

432 

433 h5py_data = np.array( 

434 [tuple(i) for i in data], dtype=[tuple([i, 'float64']) for i in parameters]) 

435 f = h5py.File(outdir + "test.hdf5", "w") 

436 lalinference = f.create_group("lalinference") 

437 nest = lalinference.create_group("lalinference_nest") 

438 samples = nest.create_dataset("posterior_samples", data=h5py_data) 

439 f.close() 

440 elif pesummary and not lalinference and not bilby: 

441 dictionary = { 

442 pesummary_label: { 

443 "posterior_samples": { 

444 "parameter_names": parameters, 

445 "samples": [list(i) for i in data] 

446 }, 

447 "injection_data": { 

448 "injection_values": [float("nan") for i in range(len(parameters))] 

449 }, 

450 "version": ["No version information found"], 

451 "meta_data": { 

452 "sampler": {"log_evidence": 0.5}, 

453 "meta_data": {} 

454 } 

455 }, 

456 "version": { 

457 "pesummary": ["v0.1.7"] 

458 } 

459 } 

460 if config is not None: 

461 dictionary[pesummary_label]["config_file"] = config 

462 if psd is not None: 

463 dictionary[pesummary_label]["psds"] = psds 

464 if calibration is not None: 

465 dictionary[pesummary_label]["calibration_envelope"] = calibration 

466 

467 if extension == "json": 

468 import json 

469 

470 with open(outdir + "test.json", "w") as f: 

471 json.dump(dictionary, f, indent=4, sort_keys=True) 

472 elif extension == "hdf5" or extension == "h5": 

473 import h5py 

474 from pesummary.core.file.meta_file import recursively_save_dictionary_to_hdf5_file 

475 

476 f = h5py.File(outdir + "test.h5", "w") 

477 recursively_save_dictionary_to_hdf5_file( 

478 f, dictionary, extra_keys=list(dictionary.keys()) 

479 ) 

480 f.close() 

481 return parameters, data 

482 

483 

484testing_dir = str(Path(__file__).parent) 

485data_dir = str(Path(__file__).parent / "files")