Coverage for pesummary/core/file/formats/json.py: 65.4%

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

2 

3import json 

4import numpy as np 

5import inspect 

6from pesummary.utils.dict import load_recursively, paths_to_key 

7 

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

9 

10 

11class PESummaryJsonEncoder(json.JSONEncoder): 

12 """Personalised JSON encoder for PESummary 

13 """ 

14 def default(self, obj): 

15 """Return a json serializable object for 'obj' 

16 

17 Parameters 

18 ---------- 

19 obj: object 

20 object you wish to make json serializable 

21 """ 

22 if isinstance(obj, np.ndarray): 

23 return obj.tolist() 

24 if inspect.isfunction(obj): 

25 return str(obj) 

26 if isinstance(obj, np.integer): 

27 return int(obj) 

28 elif isinstance(obj, np.floating): 

29 return float(obj) 

30 elif isinstance(obj, (bool, np.bool_)): 

31 return str(obj) 

32 elif isinstance(obj, bytes): 

33 return str(obj) 

34 elif isinstance(obj, type): 

35 return str(obj) 

36 return json.JSONEncoder.default(self, obj) 

37 

38 

39def PESummaryJsonDecoder(obj): 

40 if isinstance(obj, dict): 

41 if "__array__" in obj.keys() and "content" in obj.keys(): 

42 return obj["content"] 

43 elif "__complex__" in obj.keys(): 

44 return obj["real"] + obj["imag"] * 1j 

45 return obj 

46 

47 

48def read_json(path, path_to_samples=None, decoder=PESummaryJsonDecoder): 

49 """Grab the parameters and samples in a .json file 

50 

51 Parameters 

52 ---------- 

53 path: str 

54 path to the result file you wish to read in 

55 """ 

56 import json 

57 

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

59 data = json.load(f, object_hook=decoder) 

60 if not path_to_samples: 

61 try: 

62 path_to_samples, = paths_to_key("posterior", data) 

63 path_to_samples = path_to_samples[0] 

64 path_to_samples += "/posterior" 

65 except ValueError: 

66 try: 

67 path_to_samples, = paths_to_key("posterior_samples", data) 

68 path_to_samples = path_to_samples[0] 

69 path_to_samples += "/posterior_samples" 

70 except ValueError: 

71 raise ValueError( 

72 "Unable to find a 'posterior' or 'posterior_samples' group " 

73 "in the file '{}'".format(path_to_samples) 

74 ) 

75 reduced_data, = load_recursively(path_to_samples, data) 

76 if "content" in list(reduced_data.keys()): 

77 reduced_data = reduced_data["content"] 

78 parameters = list(reduced_data.keys()) 

79 reduced_data = { 

80 j: list([reduced_data[j]]) if not isinstance(reduced_data[j], list) else 

81 reduced_data[j] for j in parameters 

82 } 

83 _original_parameters = reduced_data.copy().keys() 

84 _non_numeric = [] 

85 numeric_types = (float, int, np.number) 

86 for key in _original_parameters: 

87 if any(np.iscomplex(reduced_data[key])): 

88 reduced_data[key + "_amp"] = np.abs(reduced_data[key]) 

89 reduced_data[key + "_angle"] = np.angle(reduced_data[key]) 

90 reduced_data[key] = np.real(reduced_data[key]) 

91 elif not all(isinstance(_, numeric_types) for _ in reduced_data[key]): 

92 _non_numeric.append(key) 

93 elif all(isinstance(_, (bool, np.bool_)) for _ in reduced_data[key]): 

94 _non_numeric.append(key) 

95 

96 parameters = list(reduced_data.keys()) 

97 if len(_non_numeric): 

98 from pesummary.utils.utils import logger 

99 logger.info( 

100 "Removing the parameters: '{}' from the posterior table as they " 

101 "are non-numeric".format(", ".join(_non_numeric)) 

102 ) 

103 for key in _non_numeric: 

104 parameters.remove(key) 

105 samples = [ 

106 [reduced_data[j][i] for j in parameters] for i in 

107 range(len(reduced_data[parameters[0]])) 

108 ] 

109 return parameters, samples 

110 

111 

112def _write_json( 

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

114 indent=4, sort_keys=True, dataset_name="posterior_samples", 

115 cls=PESummaryJsonEncoder, **kwargs 

116): 

117 """Write a set of samples to a json file 

118 

119 Parameters 

120 ---------- 

121 parameters: list 

122 list of parameters 

123 samples: 2d list 

124 list of samples. Columns correspond to a given parameter 

125 outdir: str, optional 

126 directory to write the dat file 

127 label: str, optional 

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

129 if not specified 

130 filename: str, optional 

131 The name of the file that you wish to write 

132 overwrite: Bool, optional 

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

134 indent: int, optional 

135 The indentation to use in json.dump. Default 4 

136 sort_keys: Bool, optional 

137 Whether or not to sort the keys in json.dump. Default True 

138 dataset_name: str, optional 

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

140 cls: class, optional 

141 Class to use as the JsonEncoder. Default PESumamryJsonEncoder 

142 """ 

143 from pesummary.utils.utils import check_filename 

144 

145 default_filename = "pesummary_{}.json" 

146 filename = check_filename( 

147 default_filename=default_filename, outdir=outdir, label=label, filename=filename, 

148 overwrite=overwrite 

149 ) 

150 _samples = np.array(samples).T 

151 data = { 

152 dataset_name: { 

153 param: _samples[num] for num, param in enumerate(parameters) 

154 } 

155 } 

156 with open(filename, "w") as f: 

157 json.dump(data, f, indent=indent, sort_keys=sort_keys, cls=cls) 

158 

159 

160def write_json( 

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

162 indent=4, sort_keys=True, cls=PESummaryJsonEncoder, **kwargs 

163): 

164 """Write a set of samples to a json file 

165 

166 Parameters 

167 ---------- 

168 parameters: list 

169 list of parameters 

170 samples: 2d list 

171 list of samples. Columns correspond to a given parameter 

172 outdir: str, optional 

173 directory to write the dat file 

174 label: str, optional 

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

176 if not specified 

177 filename: str, optional 

178 The name of the file that you wish to write 

179 overwrite: Bool, optional 

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

181 indent: int, optional 

182 The indentation to use in json.dump. Default 4 

183 sort_keys: Bool, optional 

184 Whether or not to sort the keys in json.dump. Default True 

185 cls: class, optional 

186 Class to use as the JsonEncoder. Default PESumamryJsonEncoder 

187 """ 

188 from pesummary.io.write import _multi_analysis_write 

189 

190 _multi_analysis_write( 

191 _write_json, parameters, samples, outdir=outdir, label=label, 

192 filename=filename, overwrite=overwrite, indent=indent, 

193 sort_keys=sort_keys, cls=cls, file_format="json", **kwargs 

194 )