Coverage for pesummary/cli/summarycombine_posteriors.py: 91.4%

1#! /usr/bin/env python 

2 

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

4 

5import os 

6import numpy as np 

7 

8from pesummary.utils.samples_dict import MultiAnalysisSamplesDict 

9from pesummary.utils.utils import logger 

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

11from pesummary.core.cli.inputs import _Input 

12from pesummary.io import read 

13 

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

15__doc__ = """This executable is used to combine posterior samples. This is 

16different from 'summarycombine' as 'summarycombine' combines N files into a single 

17metafile containing N analyses while 'summarycombine_posteriors' combines N 

18posterior samples and creates a single file containing a single analysis""" 

19 

20 

21class ArgumentParser(_ArgumentParser): 

22 def _pesummary_options(self): 

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

24 options.update( 

25 { 

26 "--labels": { 

27 "nargs": "+", 

28 "required": True, 

29 "help": ( 

30 "Analyses you wish to combine. If only a file " 

31 "containing more than one analysis is provided, please " 

32 "pass the labels in that file that you wish to " 

33 "combine. If multiple single analysis files are " 

34 "provided, please pass unique labels to distinguish " 

35 "each analysis. If a file containing more than one " 

36 "analysis is provided alongside a single analysis " 

37 "file, or multiple files containing more than one " 

38 "analysis each, only a single analysis can be combined " 

39 "from each file" 

40 ) 

41 }, 

42 "--weights": { 

43 "nargs": "+", 

44 "type": float, 

45 "help": ( 

46 "Weights to assign to each analysis. Must be same " 

47 "length as labels" 

48 ) 

49 }, 

50 "--use_all": { 

51 "action": "store_true", 

52 "default": False, 

53 "help": "Use all posterior samples (do not weight)" 

54 }, 

55 "--shuffle": { 

56 "action": "store_true", 

57 "default": False, 

58 "help": "Shuffle the combined samples" 

59 }, 

60 "--file_format": { 

61 "type": str, 

62 "default": "dat", 

63 "help": "Format of output file" 

64 }, 

65 "--filename": { 

66 "type": str, 

67 "help": "Name of the output file" 

68 }, 

69 "--outdir": { 

70 "type": str, 

71 "default": "./", 

72 "help": "Directory to save the file", 

73 }, 

74 "--add_to_existing": { 

75 "action": "store_true", 

76 "default": False, 

77 "help": ( 

78 "Add the combined samples to an existing metafile. " 

79 "Only used when a PESummary metafile is provided via " 

80 "the `--samples` option. If this option is provided, " 

81 "the `--file_format` and `--filename` options are " 

82 "ignored" 

83 ) 

84 } 

85 } 

86 ) 

87 return options 

88 

89 

90class Input(_Input): 

91 """Class to handle the core command line arguments 

92 

93 Parameters 

94 ---------- 

95 opts: argparse.Namespace 

96 Namespace object containing the command line options 

97 

98 Attributes 

99 ---------- 

100 result_files: list 

101 list of result files passed 

102 labels: list 

103 list of labels used to distinguish the result files 

104 """ 

105 def __init__(self, opts): 

106 self.opts = opts 

107 self.seed = self.opts.seed 

108 self.result_files = self.opts.samples 

109 self.mcmc_samples = False 

110 self.add_to_existing = False 

111 cond = np.sum([self.is_pesummary_metafile(f) for f in self.result_files]) 

112 if cond > 1: 

113 raise ValueError( 

114 "Can only combine analyses from a single PESummary metafile" 

115 ) 

116 elif cond == 1 and len(self.result_files) > 1: 

117 raise ValueError( 

118 "Can only combine analyses from a single PESummary metafile " 

119 "or multiple non-PESummary metafiles" 

120 ) 

121 self.pesummary = False 

122 if self.is_pesummary_metafile(self.result_files[0]): 

123 self.pesummary = True 

124 self._labels = self.opts.labels 

125 else: 

126 self.labels = self.opts.labels 

127 

128 

129def main(args=None): 

130 """Top level interface for `summarycombine_posteriors` 

131 """ 

132 parser = ArgumentParser() 

133 parser.add_known_options_to_parser( 

134 [ 

135 "--add_to_existing", "--labels", "--weights", "--samples", 

136 "--outdir", "--filename", "--seed", "--shuffle", "--file_format", 

137 "--use_all" 

138 ] 

139 ) 

140 opts, unknown = parser.parse_known_args(args=args) 

141 args = Input(opts) 

142 if not args.pesummary: 

143 samples = { 

144 label: samples for label, samples in 

145 zip(args.labels, args.result_files) 

146 } 

147 mydict = MultiAnalysisSamplesDict.from_files( 

148 samples, disable_prior=True, disable_injection_conversion=True 

149 ) 

150 else: 

151 mydict = read( 

152 args.result_files[0], disable_prior=True, 

153 disable_injection_conversion=True 

154 ).samples_dict 

155 combined = mydict.combine( 

156 use_all=opts.use_all, weights=opts.weights, labels=args.labels, 

157 shuffle=opts.shuffle, logger_level="info" 

158 ) 

159 if opts.add_to_existing and args.pesummary: 

160 from .summarymodify import _Input as _ModifyInput 

161 from pesummary.gw.file.meta_file import _GWMetaFile 

162 from pathlib import Path 

163 

164 class PESummaryInput(_ModifyInput): 

165 def __init__(self, samples): 

166 self.samples = samples 

167 self.data = None 

168 

169 _args = PESummaryInput(args.result_files) 

170 data = _args.data 

171 data["{}_combined".format("_".join(args.labels))] = { 

172 "posterior_samples": combined.to_structured_array(), 

173 "meta_data": { 

174 "sampler": {"nsamples": combined.number_of_samples}, 

175 "meta_data": {"combined": ", ".join(args.labels)} 

176 } 

177 } 

178 input_file = Path(args.result_files[0]) 

179 if opts.filename is None: 

180 logger.warning( 

181 "No filename specified. Saving samples to '{}' in '{}' by " 

182 "default.".format( 

183 input_file.stem + "_combined" + input_file.suffix, 

184 opts.outdir 

185 ) 

186 ) 

187 opts.filename = input_file.stem + "_combined" + input_file.suffix 

188 _metafile = os.path.join(opts.outdir, opts.filename) 

189 logger.info("Saving samples to file '{}'".format(_metafile)) 

190 if os.path.isfile(_metafile): 

191 raise FileExistsError( 

192 "Trying to write samples to '{}' but the file exists. Please " 

193 "specify a different out directory".format(_metafile) 

194 ) 

195 if input_file.suffix in [".h5", ".hdf5"]: 

196 _GWMetaFile.save_to_hdf5( 

197 data, list(data.keys()), None, _metafile, no_convert=True 

198 ) 

199 else: 

200 _GWMetaFile.save_to_json(data, _metafile) 

201 return 

202 elif opts.add_to_existing: 

203 logger.warning( 

204 "Can only use the `--add_to_existing` option when a PESummary " 

205 "metafile is provided via the `--samples` option. Writing " 

206 "combined samples to a `dat` file" 

207 ) 

208 logger.info( 

209 "Saving samples to file '{}'".format( 

210 os.path.join(opts.outdir, opts.filename) 

211 ) 

212 ) 

213 combined.write( 

214 file_format=opts.file_format, filename=opts.filename, outdir=opts.outdir 

215 ) 

216 

217 

218if __name__ == "__main__": 

219 main()