Coverage for pesummary/core/file/formats/hdf5.py: 82.8%

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

2 

3import h5py 

4import numpy as np 

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

6from pesummary.utils.dict import load_recursively, paths_to_key 

7 

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

9 

10 

11def read_hdf5(path, **kwargs): 

12 """Grab the parameters and samples in a .hdf5 file 

13 

14 Parameters 

15 ---------- 

16 path: str 

17 path to the result file you wish to read in 

18 kwargs: dict 

19 all kwargs passed to _read_hdf5_with_h5py functions 

20 """ 

21 return _read_hdf5_with_h5py(path, **kwargs) 

22 

23 

24def _read_hdf5_with_h5py( 

25 path, remove_params=None, path_to_samples=None, 

26 return_posterior_dataset=False 

27): 

28 """Grab the parameters and samples in a .hdf5 file with h5py 

29 

30 Parameters 

31 ---------- 

32 path: str 

33 path to the result file you wish to read in 

34 remove_params: list, optional 

35 parameters you wish to remove from the posterior table 

36 """ 

37 import h5py 

38 import copy 

39 

40 if path_to_samples is None: 

41 path_to_samples = Read.guess_path_to_samples(path) 

42 

43 f = h5py.File(path, 'r') 

44 c1 = isinstance(f[path_to_samples], h5py._hl.group.Group) 

45 if c1 and "parameter_names" not in f[path_to_samples].keys(): 

46 original_parameters = [i for i in f[path_to_samples].keys()] 

47 if remove_params is not None: 

48 parameters = [ 

49 i for i in original_parameters if i not in remove_params 

50 ] 

51 else: 

52 parameters = copy.deepcopy(original_parameters) 

53 n_samples = len(f[path_to_samples][parameters[0]]) 

54 try: 

55 samples = np.array([ 

56 f[path_to_samples][original_parameters.index(i)] for i in 

57 parameters 

58 ]).T 

59 except (AttributeError, KeyError, TypeError): 

60 samples = np.array([f[path_to_samples][i] for i in parameters]).T 

61 cond1 = "loglr" not in parameters or "log_likelihood" not in \ 

62 parameters 

63 cond2 = "likelihood_stats" in f.keys() and "loglr" in \ 

64 f["likelihood_stats"] 

65 if cond1 and cond2: 

66 parameters.append("log_likelihood") 

67 for num, i in enumerate(samples): 

68 samples[num].append(float(f["likelihood_stats/loglr"][num])) 

69 elif c1: 

70 original_parameters = [ 

71 i.decode("utf-8") if isinstance(i, bytes) else i for i in 

72 f[path_to_samples]["parameter_names"] 

73 ] 

74 if remove_params is not None: 

75 parameters = [ 

76 i for i in original_parameters if i not in remove_params 

77 ] 

78 else: 

79 parameters = copy.deepcopy(original_parameters) 

80 samples = np.array(f[path_to_samples]["samples"]) 

81 elif isinstance(f[path_to_samples], h5py._hl.dataset.Dataset): 

82 parameters = list(f[path_to_samples].dtype.names) 

83 samples = np.array(f[path_to_samples]).view((float, len(parameters))).tolist() 

84 for num, par in enumerate(parameters): 

85 if par == "logL": 

86 parameters[num] = "log_likelihood" 

87 if return_posterior_dataset: 

88 return parameters, samples, f[path_to_samples] 

89 f.close() 

90 return parameters, samples 

91 

92 

93def _write_hdf5( 

94 parameters, samples, outdir="./", label=None, filename=None, overwrite=False, 

95 dataset_name="posterior_samples", **kwargs 

96): 

97 """Write a set of samples to a hdf5 file 

98 

99 Parameters 

100 ---------- 

101 parameters: list 

102 list of parameters 

103 samples: 2d list 

104 list of samples. Columns correspond to a given parameter 

105 outdir: str, optional 

106 directory to write the dat file 

107 label: str, optional 

108 The label of the analysis. This is used in the filename if a filename 

109 if not specified 

110 filename: str, optional 

111 The name of the file that you wish to write 

112 overwrite: Bool, optional 

113 If True, an existing file of the same name will be overwritten 

114 dataset_name: str, optional 

115 name of the dataset to store a set of samples. Default posterior_samples 

116 """ 

117 from pesummary.utils.samples_dict import SamplesDict 

118 from pesummary.utils.utils import check_filename 

119 

120 default_filename = "pesummary_{}.h5" 

121 filename = check_filename( 

122 default_filename=default_filename, outdir=outdir, label=label, filename=filename, 

123 overwrite=overwrite 

124 ) 

125 samples = SamplesDict(parameters, np.array(samples).T) 

126 _samples = samples.to_structured_array() 

127 with h5py.File(filename, "w") as f: 

128 f.create_dataset(dataset_name, data=_samples) 

129 

130 

131def write_hdf5( 

132 parameters, samples, outdir="./", label=None, filename=None, overwrite=False, 

133 **kwargs 

134): 

135 """Write a set of samples to a hdf5 file 

136 

137 Parameters 

138 ---------- 

139 parameters: list 

140 list of parameters 

141 samples: 2d list 

142 list of samples. Columns correspond to a given parameter 

143 outdir: str, optional 

144 directory to write the dat file 

145 label: str, optional 

146 The label of the analysis. This is used in the filename if a filename 

147 if not specified 

148 filename: str, optional 

149 The name of the file that you wish to write 

150 overwrite: Bool, optional 

151 If True, an existing file of the same name will be overwritten 

152 """ 

153 from pesummary.io.write import _multi_analysis_write 

154 

155 _multi_analysis_write( 

156 _write_hdf5, parameters, samples, outdir=outdir, label=label, 

157 filename=filename, overwrite=overwrite, file_format="hdf5", **kwargs 

158 )