Coverage for pesummary/core/file/read.py: 92.9%

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

2 

3from pesummary.core.file.formats.base_read import Read 

4from pesummary.core.file.formats.bilby import Bilby, bilby_version 

5from pesummary.core.file.formats.default import Default 

6from pesummary.core.file.formats.dingo import Dingo 

7from pesummary.core.file.formats.pesummary import PESummary, PESummaryDeprecated 

8from pesummary.utils.utils import logger 

9import os 

10 

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

12 

13 

14def is_dingo_hdf5_file(path): 

15 """Determine if the results file is a dingo hdf5 results file 

16 

17 Parameters 

18 ---------- 

19 path: str 

20 path to the results file 

21 """ 

22 import h5py 

23 

24 try: 

25 with h5py.File(path, "r") as f: 

26 if "version" in f and "dingo" in str(f["version"][()]): 

27 return True 

28 else: 

29 return False 

30 except: 

31 return False 

32 

33 

34def is_bilby_hdf5_file(path): 

35 """Determine if the results file is a bilby hdf5 results file 

36 

37 Parameters 

38 ---------- 

39 path: str 

40 path to the results file 

41 """ 

42 import h5py 

43 

44 f = h5py.File(path, "r") 

45 cond = _check_bilby_version(f) 

46 f.close() 

47 return cond 

48 

49 

50def is_bilby_json_file(path): 

51 """Determine if the results file is a bilby json results file 

52 

53 Parameters 

54 ---------- 

55 path: str 

56 path to the results file 

57 """ 

58 import json 

59 

60 with open(path, "r") as f: 

61 data = json.load(f) 

62 return _check_bilby_version(data) 

63 

64 

65def _check_bilby_version(data): 

66 """Identify if 'bilby' is in the version information extracted from the 

67 file 

68 

69 Parameters 

70 ---------- 

71 data: iterable 

72 contents of the result file, opened with either h5py or json libraries 

73 """ 

74 try: 

75 version = bilby_version(data=data) 

76 if isinstance(version, bytes): 

77 version = version.decode("utf-8") 

78 if "bilby" in version: 

79 return True 

80 return False 

81 except Exception: 

82 return False 

83 

84 

85def _is_pesummary_hdf5_file(path, check_function): 

86 """Determine if the results file is a pesummary hdf5 file 

87 

88 Parameters 

89 ---------- 

90 path: str 

91 path to the results file 

92 check_function: func 

93 function used to check the result file 

94 """ 

95 import h5py 

96 

97 f = h5py.File(path, "r") 

98 outcome = check_function(f) 

99 f.close() 

100 return outcome 

101 

102 

103def is_pesummary_hdf5_file_deprecated(path): 

104 """Determine if the results file is a deprecated pesummary hdf5 file 

105 

106 Parameters 

107 ---------- 

108 path: str 

109 path to the results file 

110 """ 

111 return _is_pesummary_hdf5_file(path, _check_pesummary_file_deprecated) 

112 

113 

114def is_pesummary_hdf5_file(path): 

115 """Determine if the results file is a pesummary hdf5 file 

116 

117 Parameters 

118 ---------- 

119 path: str 

120 path to the results file 

121 """ 

122 return _is_pesummary_hdf5_file(path, _check_pesummary_file) 

123 

124 

125def _is_pesummary_json_file(path, check_function): 

126 """Determine if the result file is a pesummary json file 

127 

128 Parameters 

129 ---------- 

130 path: str 

131 path to the results file 

132 check_function: func 

133 function used to check the result file 

134 """ 

135 import json 

136 

137 with open(path, "r") as f: 

138 data = json.load(f) 

139 return check_function(data) 

140 

141 

142def is_pesummary_json_file(path): 

143 """Determine if the results file is a pesummary json file 

144 

145 Parameters 

146 ---------- 

147 path: str 

148 path to results file 

149 """ 

150 return _is_pesummary_json_file(path, _check_pesummary_file) 

151 

152 

153def is_pesummary_json_file_deprecated(path): 

154 """Determine if the results file is a deprecated pesummary json file 

155 

156 Parameters 

157 ---------- 

158 path: str 

159 path to results file 

160 """ 

161 return _is_pesummary_json_file(path, _check_pesummary_file_deprecated) 

162 

163 

164def _check_pesummary_file_deprecated(f): 

165 """Check the contents of a dictionary to see if it is a deprecated pesummary 

166 dictionary 

167 

168 Parameters 

169 ---------- 

170 f: dict 

171 dictionary of the contents of the file 

172 """ 

173 if "posterior_samples" in f.keys(): 

174 try: 

175 import collections 

176 

177 labels = f["posterior_samples"].keys() 

178 if isinstance(labels, collections.abc.KeysView): 

179 _label = list(labels)[0] 

180 try: 

181 _param = list(f["posterior_samples"][_label].keys())[0] 

182 samples = f["posterior_samples"][_label][_param] 

183 return True 

184 except Exception: 

185 return False 

186 else: 

187 return False 

188 except Exception: 

189 return False 

190 else: 

191 return False 

192 

193 

194def _check_pesummary_file(f): 

195 """Check the contents of a dictionary to see if it is a pesummary dictionary 

196 

197 Parameters 

198 ---------- 

199 f: dict 

200 dictionary of the contents of the file 

201 """ 

202 labels = f.keys() 

203 if "version" not in labels: 

204 return False 

205 try: 

206 if all( 

207 "posterior_samples" in f[label].keys() 

208 for label in labels 

209 if label != "version" and label != "history" and label != "strain" 

210 ): 

211 return True 

212 elif all( 

213 "mcmc_chains" in f[label].keys() 

214 for label in labels 

215 if label != "version" and label != "history" and label != "strain" 

216 ): 

217 return True 

218 else: 

219 return False 

220 except Exception: 

221 return False 

222 

223 

224CORE_HDF5_LOAD = { 

225 is_dingo_hdf5_file: Dingo.load_file, 

226 is_bilby_hdf5_file: Bilby.load_file, 

227 is_pesummary_hdf5_file: PESummary.load_file, 

228 is_pesummary_hdf5_file_deprecated: PESummaryDeprecated.load_file, 

229} 

230 

231CORE_JSON_LOAD = { 

232 is_bilby_json_file: Bilby.load_file, 

233 is_pesummary_json_file: PESummary.load_file, 

234 is_pesummary_json_file_deprecated: PESummaryDeprecated.load_file, 

235} 

236 

237CORE_DEFAULT_LOAD = {"default": Default.load_file} 

238 

239DEFAULT_FORMATS = ["default", "dat", "json", "hdf5", "h5", "txt"] 

240 

241 

242def _read(path, load_options, default=CORE_DEFAULT_LOAD, **load_kwargs): 

243 """Try and load a result file according to multiple options 

244 

245 Parameters 

246 ---------- 

247 path: str 

248 path to results file 

249 load_options: dict 

250 dictionary of checks and loading functions 

251 """ 

252 for check, load in load_options.items(): 

253 if check(path): 

254 try: 

255 return load(path, **load_kwargs) 

256 except ImportError as e: 

257 logger.warning( 

258 "Failed due to import error: {}. Using default load".format(e) 

259 ) 

260 return default["default"](path, **load_kwargs) 

261 except Exception as e: 

262 logger.info( 

263 "Failed to read in {} with the {} class because {}".format( 

264 path, load, e 

265 ) 

266 ) 

267 continue 

268 if len(load_options): 

269 logger.warning( 

270 "Using the default load because {} failed the following checks: {}".format( 

271 path, ", ".join([function.__name__ for function in load_options.keys()]) 

272 ) 

273 ) 

274 return default["default"](path, **load_kwargs) 

275 

276 

277def _file_format(file_format, options): 

278 """Return a dictionary of load options. If a file format is specified, then return 

279 a dictionary containing only that one load option 

280 

281 Parameters 

282 ---------- 

283 file_format: str 

284 the file format you wish to use. If None, all possible load options will be 

285 returned 

286 options: dict 

287 dictionary of possible load options. Key is a function which returns True 

288 or False depending on whether the input file belongs to that class of objects, 

289 value is the load function 

290 """ 

291 if file_format is None: 

292 return options 

293 elif any(file_format.lower() in key for key in DEFAULT_FORMATS): 

294 reduced = {} 

295 elif not any(file_format in key.__name__ for key in options.keys()): 

296 raise ValueError( 

297 "Unrecognised file format. The format must be a substring of any" 

298 "of the following: {}".format( 

299 ", ".join([key.__name__ for key in options.keys()] + DEFAULT_FORMATS) 

300 ) 

301 ) 

302 else: 

303 reduced = { 

304 key: value for key, value in options.items() if file_format in key.__name__ 

305 } 

306 return reduced 

307 

308 

309def read( 

310 path, 

311 HDF5_LOAD=CORE_HDF5_LOAD, 

312 JSON_LOAD=CORE_JSON_LOAD, 

313 DEFAULT=CORE_DEFAULT_LOAD, 

314 file_format=None, 

315 **kwargs, 

316): 

317 """Read in a results file. 

318 

319 Parameters 

320 ---------- 

321 path: str 

322 path to results file 

323 format: str 

324 the file format you wish to use. Default None. If None, the read 

325 function loops through all possible options 

326 **kwargs: dict, optional 

327 all additional kwargs are passed directly to the load_file class method 

328 """ 

329 path = os.path.expanduser(path) 

330 extension = Read.extension_from_path(path) 

331 

332 if extension in ["gz"]: 

333 from pesummary.utils.utils import unzip 

334 

335 path = unzip(path) 

336 if extension in ["hdf5", "h5", "hdf"]: 

337 options = _file_format(file_format, HDF5_LOAD) 

338 return _read(path, options, default=DEFAULT, **kwargs) 

339 elif extension == "json": 

340 options = _file_format(file_format, JSON_LOAD) 

341 return _read(path, options, default=DEFAULT, **kwargs) 

342 else: 

343 return DEFAULT["default"](path, file_format=file_format, **kwargs)