Coverage for pesummary/cli/summaryclean.py: 80.0%

1#! /usr/bin/env python 

2 

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

4 

5from pesummary.gw.file.read import read as GWRead 

6from pesummary.core.cli.parser import ArgumentParser as _ArgumentParser 

7 

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

9__doc__ = """This executable returns a cleaned data file""" 

10 

11 

12class ArgumentParser(_ArgumentParser): 

13 def _pesummary_options(self): 

14 options = super(ArgumentParser, self)._pesummary_options() 

15 options.update( 

16 { 

17 "--file_format": { 

18 "default": "dat", 

19 "help": "Save the cleaned data in this format", 

20 "choices": [ 

21 "dat", "lalinference", "bilby", "lalinference_dat" 

22 ], 

23 } 

24 } 

25 ) 

26 return options 

27 

28 

29def clean_data_file(path): 

30 """Clean the data file and return a PESummary result file object 

31 

32 Parameters 

33 ---------- 

34 path: str 

35 path to the result file 

36 """ 

37 f = GWRead(path) 

38 f.generate_all_posterior_samples() 

39 return f 

40 

41 

42def save(pesummary_object, file_format, webdir=None, label=None): 

43 """Save the pesummary_object to a given format 

44 

45 Parameters 

46 ---------- 

47 pesummary_object: pesummary.gw.file.formats 

48 pesummary results file object 

49 file_format: str 

50 the file format that you wish to save the file as 

51 webdir: str 

52 directory to save the cleaned data file 

53 """ 

54 if file_format == "dat": 

55 pesummary_object.to_dat(outdir=webdir, label=label) 

56 elif file_format == "lalinference": 

57 pesummary_object.to_lalinference(outdir=webdir, label=label) 

58 elif file_format == "lalinference_dat": 

59 pesummary_object.to_lalinference(outdir=webdir, label=label, dat=True) 

60 elif file_format == "bilby": 

61 pesummary_object.to_bilby() 

62 

63 

64def main(args=None): 

65 """Top level interface for `summaryclean` 

66 """ 

67 parser = ArgumentParser(description=__doc__) 

68 parser.add_known_options_to_parser( 

69 ["--webdir", "--samples", "--labels", "--file_format"] 

70 ) 

71 opts = parser.parse_args(args=args) 

72 if opts.labels: 

73 if len(opts.labels) != len(opts.samples): 

74 raise Exception("Please provide labels for all result files") 

75 for num, i in enumerate(opts.samples): 

76 f = clean_data_file(i) 

77 label = None 

78 if opts.labels: 

79 label = opts.labels[num] 

80 save(f, opts.file_format, webdir=opts.webdir, label=label)